zsfm_moment/infer/
mod.rs

1//! MOMENT-1-large inference engine.
2//!
3//! Architecture: T5 encoder-only transformer with patch embeddings.
4//! - RevIN normalization → patch embedding (value + position) →
5//!   24 T5 encoder blocks (RMSNorm + relative-bias MHA + gated-GELU FFN) →
6//!   final RMSNorm → per-patch reconstruction head
7//! - Zero-shot forecasting: iteratively predict 8 steps at a time using
8//!   the reconstruction head applied to the last patch's encoder output.
9
10use std::collections::HashMap;
11use std::sync::Mutex;
12use std::io::{BufReader, Read, Seek};
13use std::path::Path;
14
15use anyhow::{Context, Result};
16use candle_core::quantized::gguf_file;
17use candle_core::{DType, Device, Tensor, D};
18
19use crate::config::MomentConfig;
20
21// ---------------------------------------------------------------------------
22// Weight structs
23// ---------------------------------------------------------------------------
24
25struct EncoderBlock {
26    attn_qkv_w: Tensor, // fused [3*d_model, d_model]
27    attn_o_w: Tensor,
28    attn_norm_w: Tensor,
29    ffn_wi0_w: Tensor,
30    ffn_wi1_w: Tensor,
31    ffn_wo_w: Tensor,
32    ffn_norm_w: Tensor,
33}
34
35pub struct MomentModel {
36    device: Device,
37    config: MomentConfig,
38    patch_embed_w: Tensor,    // [1024, 8]  value embedding (no bias)
39    pos_embed: Tensor,        // [1, 5000, 1024]
40    mask_embed: Tensor,       // [1024]  learned token for masked (future) patches
41    rel_bias_data: Vec<f32>,  // [32 * 16] flat, precomputed from rel_bias_w
42    blocks: Vec<EncoderBlock>,
43    norm_f_w: Tensor,
44    head_w: Tensor,           // [8, 1024]
45    head_b: Tensor,           // [8]
46    rel_bias_cache: Mutex<HashMap<usize, Tensor>>,
47}
48
49// ---------------------------------------------------------------------------
50// GGUF loading
51// ---------------------------------------------------------------------------
52
53fn load_t(
54    content: &gguf_file::Content,
55    reader: &mut (impl Read + Seek),
56    name: &str,
57    device: &Device,
58) -> Result<Tensor> {
59    zsfm_nn::load_tensor(content, reader, name, device, DType::F32)
60}
61
62impl MomentModel {
63    pub fn load(gguf_path: &Path, config: MomentConfig) -> Result<Self> {
64        let device = Device::Cpu;
65        let file = std::fs::File::open(gguf_path)
66            .with_context(|| format!("open {}", gguf_path.display()))?;
67        let mut reader = BufReader::with_capacity(zsfm_gguf::READ_BUF_CAPACITY, file);
68        let content = gguf_file::Content::read(&mut reader).context("parse GGUF header")?;
69
70        let patch_embed_w = load_t(&content, &mut reader, "patch_embed.weight", &device)?;
71        let pos_embed     = load_t(&content, &mut reader, "pos_embed.pe", &device)?;
72        let mask_embed    = load_t(&content, &mut reader, "mask_embed", &device)?;
73        let rel_bias_data = load_t(&content, &mut reader, "blk.0.attn_rel_bias.weight", &device)?
74            .flatten_all()?.to_vec1::<f32>()?;
75
76        let mut blocks = Vec::with_capacity(config.n_layers);
77        for n in 0..config.n_layers {
78            let p = |s: &str| format!("blk.{n}.{s}");
79            let q_w = load_t(&content, &mut reader, &p("attn_q.weight"), &device)?;
80            let k_w = load_t(&content, &mut reader, &p("attn_k.weight"), &device)?;
81            let v_w = load_t(&content, &mut reader, &p("attn_v.weight"), &device)?;
82            let attn_qkv_w = Tensor::cat(&[&q_w, &k_w, &v_w], 0)
83                .with_context(|| format!("qkv cat blk.{n}"))?;
84            blocks.push(EncoderBlock {
85                attn_qkv_w,
86                attn_o_w:   load_t(&content, &mut reader, &p("attn_o.weight"), &device)?,
87                attn_norm_w: load_t(&content, &mut reader, &p("attn_norm.weight"), &device)?,
88                ffn_wi0_w:  load_t(&content, &mut reader, &p("ffn_wi0.weight"), &device)?,
89                ffn_wi1_w:  load_t(&content, &mut reader, &p("ffn_wi1.weight"), &device)?,
90                ffn_wo_w:   load_t(&content, &mut reader, &p("ffn_wo.weight"), &device)?,
91                ffn_norm_w: load_t(&content, &mut reader, &p("ffn_norm.weight"), &device)?,
92            });
93        }
94
95        let norm_f_w = load_t(&content, &mut reader, "norm_f.weight", &device)?;
96        let head_w   = load_t(&content, &mut reader, "head.weight", &device)?;
97        let head_b   = load_t(&content, &mut reader, "head.bias", &device)?;
98
99        Ok(Self {
100            device, config,
101            patch_embed_w, pos_embed, mask_embed, rel_bias_data,
102            blocks, norm_f_w, head_w, head_b,
103            rel_bias_cache: Mutex::new(HashMap::new()),
104        })
105    }
106
107    // -----------------------------------------------------------------------
108    // Forecasting
109    // -----------------------------------------------------------------------
110
111    /// Zero-shot forecasting via single-pass masked inpainting.
112    ///
113    /// Appends `n_forecast_patches` zero-filled future patches after the 64
114    /// context patches and runs ONE T5 encoder pass. The bidirectional encoder
115    /// fills the masked future positions using the real context, mirroring
116    /// the masking pretraining task. Head applied to future positions gives
117    /// the forecast.
118    pub fn forecast(&self, context: &[f32], horizon: usize) -> Result<Vec<f32>> {
119        let cfg = &self.config;
120        let patch_len = cfg.patch_len;
121        let seq_len = cfg.seq_len;
122        let ctx_patches = seq_len / cfg.patch_stride; // 64
123
124        // RevIN normalization
125        let (loc, scale) = revin_stats(context);
126        let scale = scale.max(1e-8);
127
128        // Scale and pad/truncate context to seq_len
129        let mut ctx_scaled: Vec<f32> = context.iter().map(|&v| (v - loc) / scale).collect();
130        if ctx_scaled.len() < seq_len {
131            let pad = seq_len - ctx_scaled.len();
132            let mut padded = vec![0.0f32; pad];
133            padded.extend_from_slice(&ctx_scaled);
134            ctx_scaled = padded;
135        } else if ctx_scaled.len() > seq_len {
136            let start = ctx_scaled.len() - seq_len;
137            ctx_scaled = ctx_scaled[start..].to_vec();
138        }
139
140        // Number of future patches to generate
141        let n_fc = (horizon + patch_len - 1) / patch_len;
142        let total_patches = ctx_patches + n_fc;
143
144        let rel_bias = {
145            let mut cache = self.rel_bias_cache.lock().unwrap();
146            if !cache.contains_key(&total_patches) {
147                cache.insert(total_patches, self.compute_rel_bias(total_patches)?);
148            }
149            cache[&total_patches].clone()
150        };
151
152        // Single encoder pass: context patches + mask_embed for future positions
153        let ctx_patches_data = patchify(&ctx_scaled, patch_len, cfg.patch_stride, ctx_patches);
154        let mut h = self.embed_patches_with_mask(&ctx_patches_data, n_fc, total_patches)?;
155        for blk in &self.blocks {
156            h = self.forward_block(&h, blk, &rel_bias, total_patches)?;
157        }
158        h = zsfm_nn::rms_norm(&h, Some(&self.norm_f_w), cfg.layer_norm_eps)?;
159
160        // Apply head to future positions
161        let future_h = h.narrow(0, ctx_patches, n_fc)?; // [n_fc, 1024]
162        let pred = zsfm_nn::linear_bias(&future_h, &self.head_w, &self.head_b)?; // [n_fc, 8]
163
164        let pred_flat: Vec<f32> = pred.flatten_all()?.to_vec1()?;
165        let result: Vec<f32> = pred_flat
166            .iter()
167            .take(horizon)
168            .map(|&v| v * scale + loc)
169            .collect();
170
171        Ok(result)
172    }
173
174    /// Embed context patches via value_embedding, then append n_fc mask tokens.
175    fn embed_patches_with_mask(
176        &self,
177        ctx_patches: &[Vec<f32>],
178        n_fc: usize,
179        total_patches: usize,
180    ) -> Result<Tensor> {
181        let cfg = &self.config;
182        let patch_len = cfg.patch_len;
183        let d_model = cfg.d_model;
184        let ctx_n = ctx_patches.len();
185
186        // Context: value_embedding(patch) → [ctx_n, d_model]
187        let flat: Vec<f32> = ctx_patches.iter().flatten().copied().collect();
188        let x = Tensor::from_vec(flat, (ctx_n, patch_len), &self.device)?;
189        let h_ctx = x.matmul(&self.patch_embed_w.t()?)?;
190
191        // Future: expand mask_embed [d_model] → [n_fc, d_model]
192        let h_fc = self.mask_embed.unsqueeze(0)?.broadcast_as((n_fc, d_model))?;
193
194        // Concatenate: [total_patches, d_model]
195        let h = Tensor::cat(&[h_ctx, h_fc], 0)?;
196
197        // Position embedding: pe is [1, 5000, d_model]; take [:total_patches]
198        let pos = self.pos_embed.squeeze(0)?.narrow(0, 0, total_patches)?;
199
200        Ok((h + pos)?)
201    }
202
203    fn forward_block(
204        &self,
205        hidden: &Tensor,
206        blk: &EncoderBlock,
207        rel_bias: &Tensor, // [1, n_heads, seq, seq]
208        seq_len: usize,
209    ) -> Result<Tensor> {
210        // Pre-norm attention (with residual)
211        let res = hidden;
212        let h = zsfm_nn::rms_norm(hidden, Some(&blk.attn_norm_w), self.config.layer_norm_eps)?;
213        let h = self.t5_self_attn(&h, blk, rel_bias, seq_len)?;
214        let h = (h + res)?;
215
216        // Pre-norm FFN (with residual)
217        let res2 = h.clone();
218        let h2 = zsfm_nn::rms_norm(&h, Some(&blk.ffn_norm_w), self.config.layer_norm_eps)?;
219        let h2 = gated_gelu_ffn(&h2, &blk.ffn_wi0_w, &blk.ffn_wi1_w, &blk.ffn_wo_w)?;
220        Ok((h2 + res2)?)
221    }
222
223    fn t5_self_attn(
224        &self,
225        hidden: &Tensor,
226        blk: &EncoderBlock,
227        rel_bias: &Tensor, // [1, n_heads, seq, seq]
228        seq_len: usize,
229    ) -> Result<Tensor> {
230        let cfg = &self.config;
231        let n_heads = cfg.n_heads;
232        let head_dim = cfg.head_dim;
233        let d_model = cfg.d_model;
234
235        // Project: single fused matmul → [seq, 3*d_model]
236        let qkv = zsfm_nn::linear_nobias(hidden, &blk.attn_qkv_w)?;
237        let q = qkv.narrow(D::Minus1, 0, d_model)?;
238        let k = qkv.narrow(D::Minus1, d_model, d_model)?;
239        let v = qkv.narrow(D::Minus1, 2 * d_model, d_model)?;
240
241        // Reshape to [n_heads, seq, head_dim]
242        let q = q.reshape((seq_len, n_heads, head_dim))?.permute((1, 0, 2))?.contiguous()?;
243        let k = k.reshape((seq_len, n_heads, head_dim))?.permute((1, 0, 2))?.contiguous()?;
244        let v = v.reshape((seq_len, n_heads, head_dim))?.permute((1, 0, 2))?.contiguous()?;
245
246        // T5 does NOT scale by sqrt(head_dim) — scaling is absorbed into weight init
247        let scores = q.matmul(&k.permute((0, 2, 1))?)?; // [n_heads, seq, seq]
248        // Add relative position bias (broadcast over batch dim)
249        let rel_bias_squeezed = rel_bias.squeeze(0)?; // [n_heads, seq, seq]
250        let scores = (scores + rel_bias_squeezed)?;
251        let attn = candle_nn::ops::softmax_last_dim(&scores)?;
252
253        // Weighted sum
254        let out = attn.matmul(&v)?; // [n_heads, seq, head_dim]
255        let out = out.permute((1, 0, 2))?.contiguous()?.reshape((seq_len, d_model))?;
256
257        zsfm_nn::linear_nobias(&out, &blk.attn_o_w)
258    }
259
260    /// Compute T5 relative position bias for a given sequence length.
261    fn compute_rel_bias(&self, seq_len: usize) -> Result<Tensor> {
262        let cfg = &self.config;
263        let n_heads = cfg.n_heads;
264        let num_buckets = cfg.rel_attn_num_buckets; // 32
265        let max_distance = cfg.rel_attn_max_distance; // 128
266
267        let mut out = vec![0.0f32; n_heads * seq_len * seq_len];
268        for query_pos in 0..seq_len {
269            for key_pos in 0..seq_len {
270                let rel = key_pos as i64 - query_pos as i64;
271                let bucket = t5_relative_bucket(rel, true, num_buckets, max_distance);
272                for head in 0..n_heads {
273                    // bias_data: row=bucket, col=head → index = bucket*n_heads + head
274                    let bias_val = self.rel_bias_data[bucket * n_heads + head];
275                    out[head * seq_len * seq_len + query_pos * seq_len + key_pos] = bias_val;
276                }
277            }
278        }
279
280        // Shape: [1, n_heads, seq, seq]
281        Ok(Tensor::from_vec(out, (1, n_heads, seq_len, seq_len), &self.device)?)
282    }
283}
284
285// ---------------------------------------------------------------------------
286// T5 relative position bucket
287// ---------------------------------------------------------------------------
288
289fn t5_relative_bucket(
290    relative_position: i64,
291    bidirectional: bool,
292    num_buckets: usize,
293    max_distance: usize,
294) -> usize {
295    let mut ret = 0usize;
296    let mut num_buckets = num_buckets;
297
298    let n: usize = if bidirectional {
299        num_buckets /= 2;  // each direction gets half the buckets
300        if relative_position > 0 {
301            ret += num_buckets;  // positive offset for future positions
302        }
303        relative_position.unsigned_abs() as usize
304    } else {
305        (-relative_position).max(0) as usize
306    };
307
308    let max_exact = num_buckets / 2;
309
310    if n < max_exact {
311        ret += n;
312    } else {
313        let val = max_exact
314            + ((n as f32 / max_exact as f32).ln()
315                / (max_distance as f32 / max_exact as f32).ln()
316                * (num_buckets - max_exact) as f32) as usize;
317        ret += val.min(num_buckets - 1);
318    }
319
320    ret
321}
322
323// ---------------------------------------------------------------------------
324// Ops
325// ---------------------------------------------------------------------------
326
327/// T5 gated-GELU FFN: out = wo(gelu(wi_0(x)) * wi_1(x))
328fn gated_gelu_ffn(
329    x: &Tensor,
330    wi0_w: &Tensor,
331    wi1_w: &Tensor,
332    wo_w: &Tensor,
333) -> Result<Tensor> {
334    let gate  = zsfm_nn::linear_nobias(x, wi0_w)?.gelu_erf()?;
335    let value = zsfm_nn::linear_nobias(x, wi1_w)?;
336    let h = (gate * value)?;
337    zsfm_nn::linear_nobias(&h, wo_w)
338}
339
340// ---------------------------------------------------------------------------
341// Preprocessing
342// ---------------------------------------------------------------------------
343
344/// Compute mean and std of x for RevIN normalization.
345fn revin_stats(x: &[f32]) -> (f32, f32) {
346    let n = x.len() as f64;
347    if n == 0.0 { return (0.0, 1.0); }
348    let mean = x.iter().map(|&v| v as f64).sum::<f64>() / n;
349    let var  = x.iter().map(|&v| (v as f64 - mean).powi(2)).sum::<f64>() / n;
350    (mean as f32, var.sqrt() as f32)
351}
352
353/// Extract non-overlapping patches from a sequence.
354fn patchify(x: &[f32], patch_len: usize, stride: usize, num_patches: usize) -> Vec<Vec<f32>> {
355    (0..num_patches)
356        .map(|i| {
357            let start = i * stride;
358            let end = (start + patch_len).min(x.len());
359            let mut patch = x[start..end].to_vec();
360            patch.resize(patch_len, 0.0);
361            patch
362        })
363        .collect()
364}
365
366// ---------------------------------------------------------------------------
367// zsfm-core::Forecaster
368// ---------------------------------------------------------------------------
369
370impl zsfm_core::Forecaster for MomentModel {
371    type Config = MomentConfig;
372
373    fn load(gguf_path: &Path, config: MomentConfig) -> Result<Self> {
374        MomentModel::load(gguf_path, config)
375    }
376
377    /// MOMENT is univariate-only and point-forecast-only; `mask` is unused.
378    fn forecast(
379        &self,
380        context: &[Vec<f32>],
381        _mask: &[Vec<bool>],
382        horizon: usize,
383    ) -> Result<zsfm_core::QuantileMatrix> {
384        anyhow::ensure!(context.len() == 1, "MomentModel only supports univariate forecasting (1 variate)");
385        let point = MomentModel::forecast(self, &context[0], horizon)?;
386        Ok(vec![vec![point]])
387    }
388}
zsfm_moment/infer/mod.rs

zsfm_moment/infer/
mod.rs
