zsfm_moment/infer/
mod.rs

1//! MOMENT-1-large Burn-only inference engine (Burn backend, Flex CPU).
2//!
3//! RevIN, patch embedding with mask tokens, T5 relative-bias encoder
4//! blocks, gated-GELU FFN, reconstruction head, with weights loaded
5//! straight from GGUF via [`zsfm_burn::GgufWeights`].
6
7use std::path::Path;
8
9use anyhow::{Context, Result};
10use burn::tensor::{activation, Device, Tensor, TensorData};
11use zsfm_burn::{open_gguf_weights, GgufWeights};
12
13use crate::config::MomentConfig;
14
15fn device() -> Device {
16    Device::flex()
17}
18
19fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
20    weights.load_2d(name)
21}
22
23fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
24    weights.load_1d(name)
25}
26
27fn load_flat(weights: &GgufWeights, name: &str) -> Result<Vec<f32>> {
28    weights.load_flat(name)
29}
30
31fn linear(x: Tensor<2>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<2> {
32    match b {
33        Some(bias) => zsfm_burn::linear::burn_linear_bias(x, w, bias),
34        None => zsfm_burn::linear::burn_linear_nobias(x, w),
35    }
36}
37
38fn rms(x: Tensor<2>, w: Tensor<1>, eps: f32) -> Tensor<2> {
39    zsfm_burn::norm::burn_rms_norm(x, w, eps)
40}
41
42struct BurnEncoderBlock {
43    attn_qkv_w: Tensor<2>,
44    attn_o_w: Tensor<2>,
45    attn_norm_w: Tensor<1>,
46    ffn_wi0_w: Tensor<2>,
47    ffn_wi1_w: Tensor<2>,
48    ffn_wo_w: Tensor<2>,
49    ffn_norm_w: Tensor<1>,
50}
51
52pub struct BurnMomentModel {
53    config: MomentConfig,
54    patch_embed_w: Tensor<2>,
55    pos_embed: Tensor<3>,
56    mask_embed: Tensor<1>,
57    rel_bias_data: Vec<f32>,
58    blocks: Vec<BurnEncoderBlock>,
59    norm_f_w: Tensor<1>,
60    head_w: Tensor<2>,
61    head_b: Tensor<1>,
62}
63
64impl BurnMomentModel {
65    pub fn load(gguf_path: &Path, config: MomentConfig) -> Result<Self> {
66        let weights = open_gguf_weights(gguf_path)
67            .with_context(|| format!("open {}", gguf_path.display()))?;
68
69        let patch_embed_w = load_t(&weights, "patch_embed.weight")?;
70        let pos_embed: Tensor<3> = weights.load_3d("pos_embed.pe")?;
71        let mask_embed = load_v(&weights, "mask_embed")?;
72        let rel_bias_data: Vec<f32> = load_flat(&weights, "blk.0.attn_rel_bias.weight")?;
73
74        let mut blocks = Vec::with_capacity(config.n_layers);
75        for n in 0..config.n_layers {
76            let p = |s: &str| format!("blk.{n}.{s}");
77            let qw = load_t(&weights, &p("attn_q.weight"))?;
78            let kw = load_t(&weights, &p("attn_k.weight"))?;
79            let vw = load_t(&weights, &p("attn_v.weight"))?;
80            let attn_qkv_w = Tensor::cat(vec![qw, kw, vw], 0);
81            blocks.push(BurnEncoderBlock {
82                attn_qkv_w,
83                attn_o_w: load_t(&weights, &p("attn_o.weight"))?,
84                attn_norm_w: load_v(&weights, &p("attn_norm.weight"))?,
85                ffn_wi0_w: load_t(&weights, &p("ffn_wi0.weight"))?,
86                ffn_wi1_w: load_t(&weights, &p("ffn_wi1.weight"))?,
87                ffn_wo_w: load_t(&weights, &p("ffn_wo.weight"))?,
88                ffn_norm_w: load_v(&weights, &p("ffn_norm.weight"))?,
89            });
90        }
91
92        let norm_f_w = load_v(&weights, "norm_f.weight")?;
93        let head_w = load_t(&weights, "head.weight")?;
94        let head_b = load_v(&weights, "head.bias")?;
95
96        Ok(Self {
97            config,
98            patch_embed_w,
99            pos_embed,
100            mask_embed,
101            rel_bias_data,
102            blocks,
103            norm_f_w,
104            head_w,
105            head_b,
106        })
107    }
108
109    pub fn forecast(&self, context: &[f32], horizon: usize) -> Result<Vec<f32>> {
110        let cfg = &self.config;
111        let (loc, scale) = revin_stats(context);
112        let scale = scale.max(1e-8);
113        let mut ctx_scaled: Vec<f32> = context.iter().map(|&v| (v - loc) / scale).collect();
114        if ctx_scaled.len() < cfg.seq_len {
115            let pad = cfg.seq_len - ctx_scaled.len();
116            let mut padded = vec![0.0f32; pad];
117            padded.extend_from_slice(&ctx_scaled);
118            ctx_scaled = padded;
119        } else if ctx_scaled.len() > cfg.seq_len {
120            let start = ctx_scaled.len() - cfg.seq_len;
121            ctx_scaled = ctx_scaled[start..].to_vec();
122        }
123
124        let ctx_patches = cfg.seq_len / cfg.patch_stride;
125        let n_fc = horizon.div_ceil(cfg.patch_len);
126        let total_patches = ctx_patches + n_fc;
127        let rel_bias = self.compute_rel_bias(total_patches);
128
129        let ctx_data = patchify(&ctx_scaled, cfg.patch_len, cfg.patch_stride, ctx_patches);
130        let mut h = self.embed_patches_with_mask(&ctx_data, n_fc, total_patches);
131        for blk in &self.blocks {
132            h = self.forward_block(h, blk, rel_bias.clone(), total_patches);
133        }
134        h = rms(h, self.norm_f_w.clone(), cfg.layer_norm_eps as f32);
135        let future_h: Tensor<2> = h.narrow(0, ctx_patches, n_fc);
136        let pred = linear(future_h, self.head_w.clone(), Some(self.head_b.clone()));
137        let pred_flat: Vec<f32> = pred
138            .to_data()
139            .try_to_vec()
140            .map_err(|e| anyhow::anyhow!("{e:?}"))?;
141        Ok(pred_flat
142            .iter()
143            .take(horizon)
144            .map(|&v| v * scale + loc)
145            .collect())
146    }
147
148    fn embed_patches_with_mask(
149        &self,
150        ctx_patches: &[Vec<f32>],
151        n_fc: usize,
152        total_patches: usize,
153    ) -> Tensor<2> {
154        let cfg = &self.config;
155        let dev = device();
156        let flat: Vec<f32> = ctx_patches.iter().flatten().copied().collect();
157        let x = Tensor::<2>::from_data(
158            TensorData::new(flat, [ctx_patches.len(), cfg.patch_len]),
159            &dev,
160        );
161        let h_ctx = x.matmul(self.patch_embed_w.clone().transpose());
162        let h_fc = self
163            .mask_embed
164            .clone()
165            .unsqueeze_dim::<2>(0)
166            .expand([n_fc, cfg.d_model]);
167        let h = Tensor::cat(vec![h_ctx, h_fc], 0);
168        let pos: Tensor<2> = self
169            .pos_embed
170            .clone()
171            .squeeze_dim::<2>(0)
172            .narrow(0, 0, total_patches);
173        h + pos
174    }
175
176    fn forward_block(
177        &self,
178        hidden: Tensor<2>,
179        blk: &BurnEncoderBlock,
180        rel_bias: Tensor<4>,
181        seq_len: usize,
182    ) -> Tensor<2> {
183        let h = rms(
184            hidden.clone(),
185            blk.attn_norm_w.clone(),
186            self.config.layer_norm_eps as f32,
187        );
188        let h = self.t5_self_attn(h, blk, rel_bias, seq_len);
189        let h = h + hidden;
190        let h2 = rms(
191            h.clone(),
192            blk.ffn_norm_w.clone(),
193            self.config.layer_norm_eps as f32,
194        );
195        let h2 = gated_gelu_ffn(h2, &blk.ffn_wi0_w, &blk.ffn_wi1_w, &blk.ffn_wo_w);
196        h2 + h
197    }
198
199    fn t5_self_attn(
200        &self,
201        hidden: Tensor<2>,
202        blk: &BurnEncoderBlock,
203        rel_bias: Tensor<4>,
204        seq_len: usize,
205    ) -> Tensor<2> {
206        let cfg = &self.config;
207        let qkv = linear(hidden, blk.attn_qkv_w.clone(), None);
208        let q: Tensor<2> = qkv.clone().narrow(1, 0, cfg.d_model);
209        let k: Tensor<2> = qkv.clone().narrow(1, cfg.d_model, cfg.d_model);
210        let v: Tensor<2> = qkv.narrow(1, 2 * cfg.d_model, cfg.d_model);
211        let split = |t: Tensor<2>| {
212            t.reshape([seq_len, cfg.n_heads, cfg.head_dim])
213                .permute([1, 0, 2])
214        };
215        let (q, k, v) = (split(q), split(k), split(v));
216        let scores = q.matmul(k.permute([0, 2, 1]));
217        let rel: Tensor<3> = rel_bias.squeeze_dim(0);
218        let scores = scores + rel;
219        let attn = activation::softmax(scores, 2);
220        let out = attn.matmul(v);
221        let out: Tensor<2> = out.permute([1, 0, 2]).reshape([seq_len, cfg.d_model]);
222        linear(out, blk.attn_o_w.clone(), None)
223    }
224
225    fn compute_rel_bias(&self, seq_len: usize) -> Tensor<4> {
226        let cfg = &self.config;
227        let mut out = vec![0.0f32; cfg.n_heads * seq_len * seq_len];
228        for query_pos in 0..seq_len {
229            for key_pos in 0..seq_len {
230                let rel = key_pos as i64 - query_pos as i64;
231                let bucket = t5_relative_bucket(
232                    rel,
233                    true,
234                    cfg.rel_attn_num_buckets,
235                    cfg.rel_attn_max_distance,
236                );
237                for head in 0..cfg.n_heads {
238                    out[head * seq_len * seq_len + query_pos * seq_len + key_pos] =
239                        self.rel_bias_data[bucket * cfg.n_heads + head];
240                }
241            }
242        }
243        Tensor::<4>::from_data(
244            TensorData::new(out, [1, cfg.n_heads, seq_len, seq_len]),
245            &device(),
246        )
247    }
248}
249
250fn gated_gelu_ffn(
251    x: Tensor<2>,
252    wi0_w: &Tensor<2>,
253    wi1_w: &Tensor<2>,
254    wo_w: &Tensor<2>,
255) -> Tensor<2> {
256    let gate = activation::gelu(linear(x.clone(), wi0_w.clone(), None));
257    let value = linear(x, wi1_w.clone(), None);
258    linear(gate * value, wo_w.clone(), None)
259}
260
261fn t5_relative_bucket(
262    relative_position: i64,
263    bidirectional: bool,
264    num_buckets: usize,
265    max_distance: usize,
266) -> usize {
267    let mut ret = 0usize;
268    let mut num_buckets = num_buckets;
269    let n: usize = if bidirectional {
270        num_buckets /= 2;
271        if relative_position > 0 {
272            ret += num_buckets;
273        }
274        relative_position.unsigned_abs() as usize
275    } else {
276        (-relative_position).max(0) as usize
277    };
278    let max_exact = num_buckets / 2;
279    if n < max_exact {
280        ret += n;
281    } else {
282        let val = max_exact
283            + ((n as f32 / max_exact as f32).ln() / (max_distance as f32 / max_exact as f32).ln()
284                * (num_buckets - max_exact) as f32) as usize;
285        ret += val.min(num_buckets - 1);
286    }
287    ret
288}
289
290fn revin_stats(x: &[f32]) -> (f32, f32) {
291    let n = x.len() as f64;
292    if n == 0.0 {
293        return (0.0, 1.0);
294    }
295    let mean = x.iter().map(|&v| v as f64).sum::<f64>() / n;
296    let var = x.iter().map(|&v| (v as f64 - mean).powi(2)).sum::<f64>() / n;
297    (mean as f32, var.sqrt() as f32)
298}
299
300fn patchify(x: &[f32], patch_len: usize, stride: usize, num_patches: usize) -> Vec<Vec<f32>> {
301    (0..num_patches)
302        .map(|i| {
303            let start = i * stride;
304            let end = (start + patch_len).min(x.len());
305            let mut patch = x[start..end].to_vec();
306            patch.resize(patch_len, 0.0);
307            patch
308        })
309        .collect()
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
316
317    fn pseudo(n: usize, seed: u64) -> Vec<f32> {
318        let mut x = seed;
319        (0..n)
320            .map(|_| {
321                x = x
322                    .wrapping_mul(6364136223846793005)
323                    .wrapping_add(1442695040888963407);
324                ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
325            })
326            .collect()
327    }
328
329    fn tiny_config() -> MomentConfig {
330        MomentConfig {
331            d_model: 16,
332            n_layers: 1,
333            n_heads: 2,
334            head_dim: 8,
335            d_ff: 32,
336            seq_len: 8,
337            patch_len: 2,
338            patch_stride: 2,
339            num_patches: 4,
340            rel_attn_num_buckets: 8,
341            rel_attn_max_distance: 16,
342            layer_norm_eps: 1e-6,
343        }
344    }
345
346    fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
347        *seed += 1;
348        let data: Vec<f32> = pseudo(rows * cols, *seed)
349            .iter()
350            .map(|v| v * scale)
351            .collect();
352        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
353        w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
354    }
355
356    fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
357        *seed += 1;
358        let bytes: Vec<u8> = pseudo(n, *seed)
359            .iter()
360            .flat_map(|v| v.to_le_bytes())
361            .collect();
362        w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
363    }
364
365    fn put3(w: &mut GGUFWriter, name: &str, d0: usize, d1: usize, d2: usize, seed: &mut u64) {
366        *seed += 1;
367        let bytes: Vec<u8> = pseudo(d0 * d1 * d2, *seed)
368            .iter()
369            .flat_map(|v| v.to_le_bytes())
370            .collect();
371        w.add_tensor(
372            name,
373            vec![d2 as u64, d1 as u64, d0 as u64],
374            GGMLType::F32,
375            bytes,
376        );
377    }
378
379    #[test]
380    fn burn_forecast_golden() {
381        let dir = std::env::temp_dir().join(format!("moment-burn-{}", std::process::id()));
382        let _ = std::fs::remove_dir_all(&dir);
383        std::fs::create_dir_all(&dir).unwrap();
384        let path = dir.join("model.gguf");
385        let mut w = GGUFWriter::new();
386        w.add_metadata(
387            "general.architecture",
388            GGUFMetaValue::String("moment".into()),
389        );
390        let mut seed = 500u64;
391        put2(&mut w, "patch_embed.weight", 16, 2, &mut seed, 0.3);
392        put3(&mut w, "pos_embed.pe", 1, 8, 16, &mut seed);
393        put1(&mut w, "mask_embed", 16, &mut seed);
394        put1(&mut w, "blk.0.attn_rel_bias.weight", 16, &mut seed);
395        for n in ["attn_q", "attn_k", "attn_v", "attn_o"] {
396            put2(&mut w, &format!("blk.0.{n}.weight"), 16, 16, &mut seed, 0.3);
397        }
398        put1(&mut w, "blk.0.attn_norm.weight", 16, &mut seed);
399        put2(&mut w, "blk.0.ffn_wi0.weight", 32, 16, &mut seed, 0.3);
400        put2(&mut w, "blk.0.ffn_wi1.weight", 32, 16, &mut seed, 0.3);
401        put2(&mut w, "blk.0.ffn_wo.weight", 16, 32, &mut seed, 0.3);
402        put1(&mut w, "blk.0.ffn_norm.weight", 16, &mut seed);
403        put1(&mut w, "norm_f.weight", 16, &mut seed);
404        put2(&mut w, "head.weight", 2, 16, &mut seed, 0.3);
405        put1(&mut w, "head.bias", 2, &mut seed);
406        let f = std::fs::File::create(&path).unwrap();
407        w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
408
409        let cfg = tiny_config();
410        let burn = BurnMomentModel::load(&path, cfg).unwrap();
411        let ctx: Vec<f32> = (0..8).map(|i| 10.0 + 0.5 * i as f32).collect();
412        let t0 = std::time::Instant::now();
413        let a = burn.forecast(&ctx, 4).unwrap();
414        let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
415        let b = burn.forecast(&ctx, 4).unwrap();
416        assert_eq!(a.len(), 4);
417        assert!(a.iter().all(|v| v.is_finite()));
418        assert!(a.iter().any(|&v| v != 0.0));
419        assert_eq!(a, b, "forecast must be deterministic");
420        println!("moment synthetic: burn {burn_ms:.3}ms out={a:?}");
421        // Golden values recorded from the verified Burn port.
422        let golden: [f32; 4] = [9.266289, 10.292548, 9.2582655, 10.252379];
423        assert_eq!(a.len(), golden.len());
424        for (x, y) in a.iter().zip(golden.iter()) {
425            assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
426        }
427        let _ = std::fs::remove_dir_all(&dir);
428    }
429}
zsfm_moment/infer/mod.rs

zsfm_moment/infer/
mod.rs
