open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/nemotron_decode.rs

//! Wave 49 end-to-end bring-up: drive the real
//! NVIDIA-Nemotron-3-Nano-4B-BF16 checkpoint through the
//! full NemotronH forward pass, token by token.
//!
//! ```bash
//! cargo run --release --example nemotron_decode -- \
//!     /path/to/snapshot-dir [n_tokens] [prompt text or id,id,id]
//! ```
//!
//! The snapshot dir must hold `config.json` +
//! `model.safetensors`; with `tokenizer.json` present the
//! third argument is treated as TEXT (encoded by the wave 50
//! Rust BPE tokenizer) and the generated ids are decoded
//! back to text — fully self-contained, no Python anywhere.
//! A comma-separated all-digit argument is still accepted as
//! raw ids. Pass criteria:
//!
//!   1. every step's logits are finite
//!   2. the logit distribution is non-degenerate (std > 0)
//!   3. greedy continuations are stable across re-runs
//!      (decode is deterministic)

use jouleclaw_loader_gguf::bpe_tokenizer::BpeTokenizer;
use jouleclaw_loader_gguf::nemotron_h::{NemotronHConfig, NemotronHModel};
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
use std::path::PathBuf;
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: nemotron_decode <snapshot-dir> [n_tokens]");
        std::process::exit(1);
    }
    let dir = PathBuf::from(&args[1]);
    let n_tokens: usize = args.get(2).map(|s| s.parse()).transpose()?.unwrap_or(8);

    // Config.
    let cfg_bytes = std::fs::read(dir.join("config.json"))?;
    let config: NemotronHConfig = serde_json::from_slice(&cfg_bytes)?;
    println!("config: {} layers, pattern {}",
        config.num_hidden_layers, config.hybrid_override_pattern);
    println!("        hidden {}, mamba {}h×{}d state {}, attn {}q/{}kv×{}, vocab {}",
        config.hidden_size, config.mamba_num_heads, config.mamba_head_dim,
        config.ssm_state_size, config.num_attention_heads,
        config.num_key_value_heads, config.head_dim, config.vocab_size);

    // Checkpoint.
    let t0 = Instant::now();
    let st = read_safetensors_file(dir.join("model.safetensors"))
        .map_err(|e| format!("safetensors: {e:?}"))?;
    println!("checkpoint mapped in {:.2}s ({} tensors)",
        t0.elapsed().as_secs_f64(), st.tensors.len());

    let t1 = Instant::now();
    let model = NemotronHModel::bind(config, st)?;
    println!("bound in {:.2}s (small tensors widened to f32)\n",
        t1.elapsed().as_secs_f64());

    // Tokenizer (optional — needed for text prompts + decode).
    let tokenizer = {
        let tk_path = dir.join("tokenizer.json");
        if tk_path.is_file() {
            let t = Instant::now();
            let tk = BpeTokenizer::from_file(&tk_path)?;
            println!("tokenizer loaded in {:.2}s ({} vocab)",
                t.elapsed().as_secs_f64(), tk.vocab_size());
            Some(tk)
        } else {
            None
        }
    };

    // Optional prompt: raw ids if all-digits-and-commas, else
    // text through the Rust tokenizer.
    let prompt: Vec<u32> = match args.get(3) {
        None => Vec::new(),
        Some(s) if s.chars().all(|c| c.is_ascii_digit() || c == ',') => {
            s.split(',').map(|t| t.trim().parse().unwrap()).collect()
        }
        Some(text) => {
            let tk = tokenizer.as_ref()
                .expect("text prompt needs tokenizer.json in the snapshot dir");
            let ids = tk.encode(text);
            println!("prompt {text:?} → {ids:?}");
            ids
        }
    };

    let bos = 1u32;
    let mut state = model.new_state();
    let mut generated: Vec<u32> = Vec::new();
    // Batched prefill (wave 54): BOS + prompt in one pass;
    // the returned logits are the first decode step.
    let full_prompt: Vec<u32> =
        std::iter::once(bos).chain(prompt.iter().copied()).collect();
    let t_pre = Instant::now();
    let mut logits_next = model.forward_prompt(&full_prompt, &mut state);
    let dt_pre = t_pre.elapsed();
    println!(
        "prefill: {} tokens in {:.0} ms ({:.1} tok/s)",
        full_prompt.len(),
        dt_pre.as_secs_f64() * 1e3,
        full_prompt.len() as f64 / dt_pre.as_secs_f64(),
    );
    println!("greedy decode:");
    for step in 0..n_tokens {
        let t = Instant::now();
        let logits = if step == 0 {
            std::mem::take(&mut logits_next)
        } else {
            model.forward_token(*generated.last().unwrap(), &mut state)
        };
        let dt = t.elapsed();
        // Statistics.
        let finite = logits.iter().all(|v| v.is_finite());
        let mean = logits.iter().sum::<f32>() / logits.len() as f32;
        let var = logits.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>()
            / logits.len() as f32;
        let (argmax, max) = logits.iter().enumerate()
            .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| {
                if v > bv { (i, v) } else { (bi, bv) }
            });
        println!(
            "  step {step:>2}: out={argmax:>6}  max={max:>8.3}  \
             mean={mean:>7.3}  std={:>6.3}  finite={finite}  {:.0} ms",
            var.sqrt(), dt.as_secs_f64() * 1e3,
        );
        if !finite {
            eprintln!("FAILED: non-finite logits at step {step}");
            std::process::exit(1);
        }
        if var.sqrt() < 1e-3 {
            eprintln!("FAILED: degenerate logit distribution at step {step}");
            std::process::exit(1);
        }
        generated.push(argmax as u32);
    }
    println!("\ngenerated token ids: {generated:?}");
    if let Some(tk) = &tokenizer {
        println!("generated text: {:?}", tk.decode(&generated));
    }
    #[cfg(feature = "fabric")]
    {
        let entries = jouleclaw_loader_gguf::fabric_dispatch::ledger_snapshot();
        if !entries.is_empty() {
            let total_gf: f64 = entries.iter().map(|e| e.gflops).sum();
            println!("\nfabric attribution (whole run):");
            for e in &entries {
                let uj = e
                    .projected_uj
                    .map(|u| format!("{:.1} J projected", u / 1e6))
                    .unwrap_or_else(|| "no calibration".into());
                println!(
                    "  {:>18}: {:>6} ops, {:>9.1} GFLOP ({:>4.1}%), {}",
                    e.member_id, e.ops, e.gflops,
                    100.0 * e.gflops / total_gf, uj,
                );
            }
        }
    }
    println!("PASS: {n_tokens} steps, all logits finite and non-degenerate.");
    Ok(())
}