1use std::path::Path;
9
10use anyhow::{Context, Result};
11use burn::tensor::{activation, Device, Tensor, TensorData};
12use zsfm_burn::{open_gguf_weights, GgufWeights};
13
14use zsfm_burn::linear::{burn_linear_bias, burn_linear_nobias};
15use zsfm_burn::norm::burn_rms_norm_nd;
16
17use crate::config::MoiraiConfig;
18
19const PATCH_SIZE: usize = 32;
20const PATCH_IDX: usize = 2;
21
22fn device() -> Device {
23 Device::flex()
24}
25
26fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
27 weights.load_2d(name)
28}
29
30fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
31 weights.load_1d(name)
32}
33
34fn load_3d(weights: &GgufWeights, name: &str) -> Result<Tensor<3>> {
35 weights.load_3d(name)
36}
37
38fn load_flat(weights: &GgufWeights, name: &str) -> Result<Vec<f32>> {
39 weights.load_flat(name)
40}
41
42fn linear(x: Tensor<2>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<2> {
43 match b {
44 Some(bias) => burn_linear_bias(x, w, bias),
45 None => burn_linear_nobias(x, w),
46 }
47}
48
49fn linear_nd<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
50 zsfm_burn::linear::burn_linear_nd(x, w, b)
51}
52
53fn rms<const D: usize>(x: Tensor<D>, w: Tensor<1>) -> Tensor<D> {
54 burn_rms_norm_nd(x, w, 1e-6)
55}
56
57fn swiglu_ffn(x: Tensor<2>, fc1_w: Tensor<2>, fc2_w: Tensor<2>, gate_w: Tensor<2>) -> Tensor<2> {
58 let content = activation::silu(linear(x.clone(), fc1_w, None));
59 let gate = linear(x, gate_w, None);
60 linear(content * gate, fc2_w, None)
61}
62
63struct BurnEncoderBlock {
64 norm1_w: Tensor<1>,
65 norm2_w: Tensor<1>,
66 attn_qkv_w: Tensor<2>,
67 attn_o_w: Tensor<2>,
68 attn_qn_w: Tensor<1>,
69 attn_kn_w: Tensor<1>,
70 vbias_obs: Tensor<3>,
71 vbias_mask: Tensor<3>,
72 ffn_fc1_w: Tensor<2>,
73 ffn_fc2_w: Tensor<2>,
74 ffn_gate_w: Tensor<2>,
75}
76
77pub struct BurnMoiraiModel {
78 config: MoiraiConfig,
79 in_proj_w: Tensor<3>,
80 in_proj_b: Tensor<2>,
81 mask_embed: Tensor<2>,
82 blocks: Vec<BurnEncoderBlock>,
83 norm_f_w: Tensor<1>,
84 head_st_loc_w: Tensor<3>,
85 head_st_loc_b: Tensor<2>,
86 rope_inv_freq: Vec<f32>,
87}
88
89impl BurnMoiraiModel {
90 pub fn load(gguf_path: &Path, config: MoiraiConfig) -> Result<Self> {
91 let weights = open_gguf_weights(gguf_path)
92 .with_context(|| format!("open {}", gguf_path.display()))?;
93 let dev = device();
94
95 let in_proj_w = load_3d(&weights, "in_proj.weight")?;
96 let in_proj_b = {
97 let d: Vec<f32> = load_flat(&weights, "in_proj.bias")?;
98 anyhow::ensure!(d.len().is_multiple_of(5), "in_proj.bias len");
99 let d_model = d.len() / 5;
100 Tensor::<2>::from_data(TensorData::new(d, [5, d_model]), &dev)
101 };
102 let mask_embed = load_t(&weights, "mask_embed.weight")?;
103
104 let n_heads = config.n_heads;
105 let mut blocks = Vec::with_capacity(config.n_layers);
106 for n in 0..config.n_layers {
107 let p = |s: &str| format!("blk.{n}.{s}");
108 let qw = load_t(&weights, &p("attn_q.weight"))?;
109 let kw = load_t(&weights, &p("attn_k.weight"))?;
110 let vw = load_t(&weights, &p("attn_v.weight"))?;
111 let attn_qkv_w = Tensor::cat(vec![qw, kw, vw], 0);
112 let vbias_raw: Vec<f32> = load_flat(&weights, &p("attn_vbias.weight"))?;
113 let vbias_obs = Tensor::<3>::from_data(
114 TensorData::new(vbias_raw[0..n_heads].to_vec(), [n_heads, 1, 1]),
115 &dev,
116 );
117 let vbias_mask = Tensor::<3>::from_data(
118 TensorData::new(vbias_raw[n_heads..2 * n_heads].to_vec(), [n_heads, 1, 1]),
119 &dev,
120 );
121 blocks.push(BurnEncoderBlock {
122 norm1_w: load_v(&weights, &p("norm1.weight"))?,
123 norm2_w: load_v(&weights, &p("norm2.weight"))?,
124 attn_qkv_w,
125 attn_o_w: load_t(&weights, &p("attn_o.weight"))?,
126 attn_qn_w: load_v(&weights, &p("attn_qn.weight"))?,
127 attn_kn_w: load_v(&weights, &p("attn_kn.weight"))?,
128 vbias_obs,
129 vbias_mask,
130 ffn_fc1_w: load_t(&weights, &p("ffn_fc1.weight"))?,
131 ffn_fc2_w: load_t(&weights, &p("ffn_fc2.weight"))?,
132 ffn_gate_w: load_t(&weights, &p("ffn_gate.weight"))?,
133 });
134 }
135
136 let norm_f_w = load_v(&weights, "norm_f.weight")?;
137 let head_st_loc_w = load_3d(&weights, "head.st_loc.weight")?;
138 let head_st_loc_b = {
139 let d: Vec<f32> = load_flat(&weights, "head.st_loc.bias")?;
140 anyhow::ensure!(d.len().is_multiple_of(5), "head bias len");
141 let max_ps = d.len() / 5;
142 Tensor::<2>::from_data(TensorData::new(d, [5, max_ps]), &dev)
143 };
144
145 let half = config.head_dim / 2;
146 let rope_inv_freq: Vec<f32> = (0..half)
147 .map(|i| 1.0_f32 / 10000_f32.powf(2.0 * i as f32 / config.head_dim as f32))
148 .collect();
149
150 Ok(Self {
151 config,
152 in_proj_w,
153 in_proj_b,
154 mask_embed,
155 blocks,
156 norm_f_w,
157 head_st_loc_w,
158 head_st_loc_b,
159 rope_inv_freq,
160 })
161 }
162
163 pub fn forecast(&self, context: &[f32], horizon: usize) -> Result<Vec<f32>> {
164 let cfg = &self.config;
165 let patch_size = PATCH_SIZE;
166 let patch_idx = PATCH_IDX;
167 let loc = context.iter().map(|&v| v as f64).sum::<f64>() / context.len() as f64;
168 let scale =
169 context.iter().map(|&v| (v as f64 - loc).abs()).sum::<f64>() / context.len() as f64;
170 let scale = (scale.max(1e-8)) as f32;
171 let loc = loc as f32;
172 let max_ts = cfg.max_seq_len;
173 let mut ctx_scaled: Vec<f32> = context.iter().map(|&v| (v - loc) / scale).collect();
174 if ctx_scaled.len() > max_ts {
175 let start = ctx_scaled.len() - max_ts;
176 ctx_scaled = ctx_scaled[start..].to_vec();
177 }
178 let ctx_len = ctx_scaled.len();
179 let ctx_padded_len = ctx_len.div_ceil(patch_size) * patch_size;
180 if ctx_padded_len > ctx_len {
181 let mut padded = vec![0.0f32; ctx_padded_len - ctx_len];
182 padded.extend_from_slice(&ctx_scaled);
183 ctx_scaled = padded;
184 }
185 let n_ctx_patches = ctx_scaled.len() / patch_size;
186 let n_fc_patches = horizon.div_ceil(patch_size);
187 let total_patches = n_ctx_patches + n_fc_patches;
188 let d_model = cfg.d_model;
189 let max_ps = cfg.max_patch_size;
190 let dev = device();
191
192 let proj_w_full: Tensor<2> = self
193 .in_proj_w
194 .clone()
195 .narrow(0, patch_idx, 1)
196 .squeeze_dim::<2>(0);
197 let proj_w: Tensor<2> = proj_w_full.narrow(1, 0, patch_size);
198 let proj_b: Tensor<1> = self
199 .in_proj_b
200 .clone()
201 .narrow(0, patch_idx, 1)
202 .squeeze_dim::<1>(0);
203 let mut patch_flat = vec![0.0f32; n_ctx_patches * patch_size];
204 for i in 0..n_ctx_patches {
205 let src = &ctx_scaled[i * patch_size..(i + 1) * patch_size];
206 patch_flat[i * patch_size..(i + 1) * patch_size].copy_from_slice(src);
207 }
208 let ctx_patches_t = Tensor::<2>::from_data(
209 TensorData::new(patch_flat, [n_ctx_patches, patch_size]),
210 &dev,
211 );
212 let ctx_emb = linear(ctx_patches_t, proj_w, Some(proj_b));
213 let mask_embed: Tensor<2> = self.mask_embed.clone().reshape([1, d_model]);
214 let fc_emb = mask_embed.expand([n_fc_patches, d_model]);
215 let mut h = Tensor::cat(vec![ctx_emb, fc_emb], 0);
216
217 let mut is_masked = vec![0u8; total_patches];
218 is_masked[n_ctx_patches..total_patches].fill(1);
219 for blk in &self.blocks {
220 h = self.forward_block(h, blk, &is_masked, total_patches)?;
221 }
222 h = rms(h, self.norm_f_w.clone());
223
224 let loc_w_3d: Tensor<2> = self
225 .head_st_loc_w
226 .clone()
227 .narrow(0, patch_idx, 1)
228 .squeeze_dim::<2>(0);
229 let loc_w: Tensor<2> = loc_w_3d.narrow(1, 0, patch_size);
230 let loc_b: Tensor<1> = self
231 .head_st_loc_b
232 .clone()
233 .narrow(0, patch_idx, 1)
234 .squeeze_dim::<1>(0)
235 .narrow(0, 0, patch_size);
236 let future_h: Tensor<2> = h.narrow(0, n_ctx_patches, n_fc_patches);
237 let pred = linear(future_h, loc_w, Some(loc_b));
238 let pred_flat: Vec<f32> = pred
239 .to_data()
240 .try_to_vec()
241 .map_err(|e| anyhow::anyhow!("{e:?}"))?;
242 let _ = max_ps;
243 Ok(pred_flat
244 .iter()
245 .take(horizon)
246 .map(|&v| v * scale + loc)
247 .collect())
248 }
249
250 fn forward_block(
251 &self,
252 hidden: Tensor<2>,
253 blk: &BurnEncoderBlock,
254 is_masked: &[u8],
255 seq_len: usize,
256 ) -> Result<Tensor<2>> {
257 let h = rms(hidden.clone(), blk.norm1_w.clone());
258 let h = self.qk_attn(h, blk, is_masked, seq_len)?;
259 let h = h + hidden;
260 let h2 = rms(h.clone(), blk.norm2_w.clone());
261 let h2 = swiglu_ffn(
262 h2,
263 blk.ffn_fc1_w.clone(),
264 blk.ffn_fc2_w.clone(),
265 blk.ffn_gate_w.clone(),
266 );
267 Ok(h2 + h)
268 }
269
270 fn qk_attn(
271 &self,
272 hidden: Tensor<2>,
273 blk: &BurnEncoderBlock,
274 is_masked: &[u8],
275 seq_len: usize,
276 ) -> Result<Tensor<2>> {
277 let cfg = &self.config;
278 let n_heads = cfg.n_heads;
279 let head_dim = cfg.head_dim;
280 let d_model = cfg.d_model;
281 let qkv = linear_nd(hidden, blk.attn_qkv_w.clone(), None);
282 let q: Tensor<2> = qkv.clone().narrow(1, 0, d_model);
283 let k: Tensor<2> = qkv.clone().narrow(1, d_model, d_model);
284 let v: Tensor<2> = qkv.narrow(1, 2 * d_model, d_model);
285 let q: Tensor<3> = q.reshape([seq_len, n_heads, head_dim]);
286 let k: Tensor<3> = k.reshape([seq_len, n_heads, head_dim]);
287 let q = qk_norm_heads(q, blk.attn_qn_w.clone(), seq_len, n_heads, head_dim);
288 let k = qk_norm_heads(k, blk.attn_kn_w.clone(), seq_len, n_heads, head_dim);
289 let q: Tensor<3> = q.permute([1, 0, 2]);
290 let k: Tensor<3> = k.permute([1, 0, 2]);
291 let (cos_t, sin_t) = self.rope_tables(seq_len);
292 let q = apply_rope_with_tables(q, cos_t.clone(), sin_t.clone(), head_dim);
293 let k = apply_rope_with_tables(k, cos_t, sin_t, head_dim);
294 let v: Tensor<3> = v.reshape([seq_len, n_heads, head_dim]).permute([1, 0, 2]);
295 let scores = q
296 .matmul(k.permute([0, 2, 1]))
297 .mul_scalar(1.0 / (head_dim as f32).sqrt());
298 let scores = apply_var_attn_bias(
299 scores,
300 is_masked,
301 seq_len,
302 blk.vbias_obs.clone(),
303 blk.vbias_mask.clone(),
304 );
305 let attn = activation::softmax(scores, 2);
306 let out = attn.matmul(v);
307 let out: Tensor<2> = out.permute([1, 0, 2]).reshape([seq_len, d_model]);
308 Ok(linear_nd(out, blk.attn_o_w.clone(), None))
309 }
310
311 fn rope_tables(&self, seq_len: usize) -> (Tensor<3>, Tensor<3>) {
312 let dev = device();
313 let half = self.rope_inv_freq.len();
314 let mut cos_v = vec![0.0f32; seq_len * half];
315 let mut sin_v = vec![0.0f32; seq_len * half];
316 for pos in 0..seq_len {
317 for i in 0..half {
318 let theta = pos as f32 * self.rope_inv_freq[i];
319 cos_v[pos * half + i] = theta.cos();
320 sin_v[pos * half + i] = theta.sin();
321 }
322 }
323 let cos_t = Tensor::<2>::from_data(TensorData::new(cos_v, [seq_len, half]), &dev)
324 .unsqueeze_dim::<3>(0);
325 let sin_t = Tensor::<2>::from_data(TensorData::new(sin_v, [seq_len, half]), &dev)
326 .unsqueeze_dim::<3>(0);
327 (cos_t, sin_t)
328 }
329}
330
331fn qk_norm_heads(
332 x: Tensor<3>,
333 weight: Tensor<1>,
334 seq_len: usize,
335 n_heads: usize,
336 head_dim: usize,
337) -> Tensor<3> {
338 let x_flat: Tensor<2> = x.reshape([seq_len * n_heads, head_dim]);
339 burn_rms_norm_nd(x_flat, weight, 1e-6).reshape([seq_len, n_heads, head_dim])
340}
341
342fn apply_var_attn_bias(
343 scores: Tensor<3>,
344 is_masked: &[u8],
345 seq_len: usize,
346 vbias_obs: Tensor<3>,
347 vbias_mask: Tensor<3>,
348) -> Tensor<3> {
349 let dev = device();
350 let mask_f: Vec<f32> = is_masked.iter().map(|&m| m as f32).collect();
351 let mask_t = Tensor::<3>::from_data(TensorData::new(mask_f, [1, 1, seq_len]), &dev);
352 let delta = vbias_mask - vbias_obs.clone();
353 let bias = vbias_obs + delta * mask_t;
354 scores + bias
355}
356
357fn apply_rope_with_tables(
358 x: Tensor<3>,
359 cos_t: Tensor<3>,
360 sin_t: Tensor<3>,
361 head_dim: usize,
362) -> Tensor<3> {
363 let half = head_dim / 2;
364 let x1 = x.clone().narrow(2, 0, half);
365 let x2 = x.clone().narrow(2, half, half);
366 let rot1 = x1.clone() * cos_t.clone() - x2.clone() * sin_t.clone();
367 let rot2 = x1 * sin_t + x2 * cos_t;
368 Tensor::cat(vec![rot1, rot2], 2)
369}
370
371#[cfg(test)]
372mod tests {
373 use super::*;
374 use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
375
376 fn pseudo(n: usize, seed: u64) -> Vec<f32> {
377 let mut x = seed;
378 (0..n)
379 .map(|_| {
380 x = x
381 .wrapping_mul(6364136223846793005)
382 .wrapping_add(1442695040888963407);
383 ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
384 })
385 .collect()
386 }
387
388 fn tiny_config() -> MoiraiConfig {
389 MoiraiConfig {
390 d_model: 32,
391 n_layers: 1,
392 n_heads: 4,
393 head_dim: 8,
394 d_ff: 64,
395 max_seq_len: 64,
396 patch_sizes: vec![8, 16, 32, 64, 128],
397 max_patch_size: 32,
398 }
399 }
400
401 fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
402 *seed += 1;
403 let data: Vec<f32> = pseudo(rows * cols, *seed)
404 .iter()
405 .map(|v| v * scale)
406 .collect();
407 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
408 w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
409 }
410
411 fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
412 *seed += 1;
413 let bytes: Vec<u8> = pseudo(n, *seed)
414 .iter()
415 .flat_map(|v| v.to_le_bytes())
416 .collect();
417 w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
418 }
419
420 fn put3(
421 w: &mut GGUFWriter,
422 name: &str,
423 d0: usize,
424 d1: usize,
425 d2: usize,
426 seed: &mut u64,
427 scale: f32,
428 ) {
429 *seed += 1;
430 let data: Vec<f32> = pseudo(d0 * d1 * d2, *seed)
431 .iter()
432 .map(|v| v * scale)
433 .collect();
434 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
435 w.add_tensor(
436 name,
437 vec![d2 as u64, d1 as u64, d0 as u64],
438 GGMLType::F32,
439 bytes,
440 );
441 }
442
443 #[test]
444 fn burn_forecast_golden() {
445 let dir = std::env::temp_dir().join(format!("moirai-burn-{}", std::process::id()));
447 let _ = std::fs::remove_dir_all(&dir);
448 std::fs::create_dir_all(&dir).unwrap();
449 let path = dir.join("model.gguf");
450 let mut w = GGUFWriter::new();
451 w.add_metadata(
452 "general.architecture",
453 GGUFMetaValue::String("moirai".into()),
454 );
455 let mut seed = 1300u64;
456 let s = 0.1;
457 put3(&mut w, "in_proj.weight", 5, 32, 32, &mut seed, s);
458 put2(&mut w, "in_proj.bias", 5, 32, &mut seed, 1.0);
459 put2(&mut w, "mask_embed.weight", 1, 32, &mut seed, 1.0);
460 put2(&mut w, "blk.0.attn_q.weight", 32, 32, &mut seed, s);
461 put2(&mut w, "blk.0.attn_k.weight", 32, 32, &mut seed, s);
462 put2(&mut w, "blk.0.attn_v.weight", 32, 32, &mut seed, s);
463 put1(&mut w, "blk.0.norm1.weight", 32, &mut seed);
464 put1(&mut w, "blk.0.norm2.weight", 32, &mut seed);
465 put2(&mut w, "blk.0.attn_o.weight", 32, 32, &mut seed, s);
466 put1(&mut w, "blk.0.attn_qn.weight", 8, &mut seed);
467 put1(&mut w, "blk.0.attn_kn.weight", 8, &mut seed);
468 put1(&mut w, "blk.0.attn_vbias.weight", 8, &mut seed);
469 put2(&mut w, "blk.0.ffn_fc1.weight", 64, 32, &mut seed, s);
470 put2(&mut w, "blk.0.ffn_fc2.weight", 32, 64, &mut seed, s);
471 put2(&mut w, "blk.0.ffn_gate.weight", 64, 32, &mut seed, s);
472 put1(&mut w, "norm_f.weight", 32, &mut seed);
473 put3(&mut w, "head.st_loc.weight", 5, 32, 32, &mut seed, s);
474 put2(&mut w, "head.st_loc.bias", 5, 32, &mut seed, 1.0);
475 let f = std::fs::File::create(&path).unwrap();
476 w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
477
478 let cfg = tiny_config();
479 let burn = BurnMoiraiModel::load(&path, cfg).unwrap();
480 let ctx: Vec<f32> = (0..32).map(|i| 10.0 + 0.5 * i as f32).collect();
481 let t0 = std::time::Instant::now();
482 let a = burn.forecast(&ctx, 4).unwrap();
483 let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
484 let b = burn.forecast(&ctx, 4).unwrap();
485 assert_eq!(a.len(), 4);
486 assert!(a.iter().all(|v| v.is_finite()));
487 assert!(a.iter().any(|&v| v != 0.0));
488 assert_eq!(a, b, "forecast must be deterministic");
489 println!("moirai synthetic: burn {burn_ms:.1}ms out={a:?}");
490 let golden: [f32; 4] = [10.968519, 12.493711, 11.747968, 11.247049];
492 assert_eq!(a.len(), golden.len());
493 for (x, y) in a.iter().zip(golden.iter()) {
494 assert!((x - y).abs() < 1e-3, "golden drift: {x} vs {y}");
495 }
496 let _ = std::fs::remove_dir_all(&dir);
497 }
498}