inferencelayer 0.2.9

//! `precision-search` — derive a per-layer precision policy for ONE model, judged by the KLD
//! ruler: the engine's own version of what Unsloth hand-tunes per model family ("Dynamic 2.0").
//!
//! ```text
//! precision-search --model <dir> [--calib bench/calib_v1.txt] [--band 2] [--rounds 4]
//!                  [--max-docs 10] [--max-tokens 64] [--out precision.json]
//! ```
//!
//! Greedy sensitivity search: candidates are `head=f16` and whole-layer bands `blk:a-b=f16`;
//! each round scores every remaining candidate ON TOP of the accepted set and accepts the best
//! **ΔKLD per Δmegabyte** — so a cheap band that removes little damage loses to an expensive one
//! that removes a lot, and vice versa, on measured numbers rather than heuristics.
//!
//! Cost control: the f16 REFERENCE arm runs the calibration subset once and its per-position
//! log-softmax is cached (f32), so each candidate evaluation is one arm load plus one subset
//! pass. The search subset is deliberately small (`--max-docs`/`--max-tokens`); the emitted
//! policy should be RE-VERIFIED on the full corpus with `kld-eval --arm gpu-prec:<spec>` — the
//! search ranks, the ruler judges.

use anyhow::{Context, Result};
use inferencelayer::weights::{Lfm2Config, Precision, WDtype};

struct Args {
    model: String,
    calib: String,
    band: usize,
    rounds: usize,
    max_docs: usize,
    max_tokens: usize,
    out: Option<String>,
}

fn parse_args() -> Result<Args> {
    let mut a = Args {
        model: String::new(),
        calib: "bench/calib_v1.txt".into(),
        band: 2,
        rounds: 4,
        max_docs: 10,
        max_tokens: 64,
        out: None,
    };
    let mut it = std::env::args().skip(1);
    while let Some(k) = it.next() {
        match k.as_str() {
            "--model" => a.model = it.next().context("--model value")?,
            "--calib" => a.calib = it.next().context("--calib value")?,
            "--band" => a.band = it.next().context("--band value")?.parse()?,
            "--rounds" => a.rounds = it.next().context("--rounds value")?.parse()?,
            "--max-docs" => a.max_docs = it.next().context("--max-docs value")?.parse()?,
            "--max-tokens" => a.max_tokens = it.next().context("--max-tokens value")?.parse()?,
            "--out" => a.out = it.next(),
            other => anyhow::bail!("unknown argument {other}"),
        }
    }
    anyhow::ensure!(!a.model.is_empty(), "--model <dir> required");
    Ok(a)
}

/// One upgrade candidate: the policy term, and its byte cost (MB) going Q4_0 → f16.
#[derive(Clone)]
struct Cand {
    term: String,
    mb: f64,
}

/// Teacher-forced token streams for the calibration subset.
fn calib_ids(args: &Args, dir: &std::path::Path) -> Result<Vec<Vec<u32>>> {
    let text = std::fs::read_to_string(&args.calib)
        .with_context(|| format!("calibration corpus {}", args.calib))?;
    let tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
        .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;
    let mut docs = Vec::new();
    for line in text
        .lines()
        .map(str::trim)
        .filter(|l| !l.is_empty() && !l.starts_with('#'))
        .take(args.max_docs)
    {
        let ids: Vec<u32> = tok
            .encode(line, true)
            .map_err(|e| anyhow::anyhow!("encode: {e}"))?
            .get_ids()
            .iter()
            .copied()
            .take(args.max_tokens)
            .collect();
        if ids.len() >= 2 {
            docs.push(ids);
        }
    }
    anyhow::ensure!(!docs.is_empty(), "no calibration documents");
    Ok(docs)
}

/// Run one arm over the subset. With `cache = None` returns the per-position log-softmax rows
/// (the reference pass); with `cache = Some(rows)` returns the mean KLD(ref ‖ arm) instead.
fn run_arm(
    ctx: &inferencelayer::GpuCtx,
    dir: &std::path::Path,
    spec: &str,
    docs: &[Vec<u32>],
    cache: Option<&[Vec<f32>]>,
) -> Result<(f64, Vec<Vec<f32>>)> {
    let w = inferencelayer::Weights::load_with_precision(ctx, dir, Precision::parse(spec)?)?;
    let g = inferencelayer::Lfm2Gpu::new(ctx, w);
    let mut rows: Vec<Vec<f32>> = Vec::new();
    let (mut kld_sum, mut n) = (0f64, 0usize);
    let mut pos_idx = 0usize;
    for ids in docs {
        g.reset(ctx);
        for pos in 0..ids.len() - 1 {
            let logits = g.forward(ctx, ids[pos], pos)?;
            let lq = inferencelayer::kld::log_softmax_f64(&logits);
            match cache {
                None => rows.push(lq.iter().map(|x| *x as f32).collect()),
                Some(refrows) => {
                    // KLD(ref ‖ arm) = Σ p·(lp − lq), ref log-probs from the f32 cache.
                    let lp = &refrows[pos_idx];
                    let mut k = 0f64;
                    for (a, b) in lp.iter().zip(&lq) {
                        let a = f64::from(*a);
                        let p = a.exp();
                        if p > 0.0 {
                            k += p * (a - b);
                        }
                    }
                    kld_sum += k;
                    n += 1;
                }
            }
            pos_idx += 1;
        }
    }
    Ok((if n > 0 { kld_sum / n as f64 } else { 0.0 }, rows))
}

fn main() -> Result<()> {
    let args = parse_args()?;
    let dir = std::path::Path::new(&args.model);
    let cfg = Lfm2Config::from_json(&std::fs::read(dir.join("config.json"))?)?;
    let ctx = inferencelayer::GpuCtx::new()?;
    eprintln!(
        "model: {:?}, {} layers, hidden {}, vocab {}",
        cfg.arch, cfg.n_layers, cfg.hidden, cfg.vocab
    );

    // ── Byte model: Δbytes for upgrading a group Q4_0 (4.5 bpw) → f16 (16 bpw) ──────────────
    let dbpw = (16.0 - 4.5) / 8.0; // Δbytes per weight
    let (h, hd) = (cfg.hidden as f64, cfg.head_dim as f64);
    let layer_params = (cfg.n_heads as f64 * hd + 2.0 * cfg.n_kv_heads as f64 * hd) * h // qkv
        + cfg.n_heads as f64 * hd * h                                                   // o
        + 3.0 * h * cfg.intermediate as f64; // gate+up+down
    let head_params = cfg.vocab as f64 * h;
    let band_mb = |nl: usize| nl as f64 * layer_params * dbpw / 1e6;

    let mut cands: Vec<Cand> = vec![Cand {
        term: "head=f16".into(),
        mb: head_params * dbpw / 1e6,
    }];
    let mut a = 0usize;
    while a < cfg.n_layers {
        let b = (a + args.band - 1).min(cfg.n_layers - 1);
        cands.push(Cand {
            term: format!("blk:{a}-{b}=f16"),
            mb: band_mb(b - a + 1),
        });
        a = b + 1;
    }

    // ── Reference pass (cached) + base score ────────────────────────────────────────────────
    let docs = calib_ids(&args, dir)?;
    let npos: usize = docs.iter().map(|d| d.len() - 1).sum();
    eprintln!("subset: {} docs, {npos} positions; caching f16 reference…", docs.len());
    let (_, refrows) = run_arm(&ctx, dir, "f16", &docs, None)?;
    let (mut cur_kld, _) = run_arm(&ctx, dir, "q4", &docs, Some(&refrows))?;
    eprintln!("base q4: KLD {cur_kld:.4} nats (subset)");

    // ── Greedy rounds: accept the best ΔKLD/MB ──────────────────────────────────────────────
    let mut accepted: Vec<Cand> = Vec::new();
    let mut history = Vec::new();
    for round in 1..=args.rounds {
        let mut best: Option<(usize, f64, f64)> = None; // (idx, kld_with, ratio)
        for (i, c) in cands.iter().enumerate() {
            let spec = build_spec(&accepted, Some(c));
            let (kld, _) = run_arm(&ctx, dir, &spec, &docs, Some(&refrows))?;
            let ratio = (cur_kld - kld) / c.mb;
            eprintln!(
                "  round {round}: +{:<14} KLD {kld:.4}  Δ {:+.4}  ({:.1} MB, ratio {ratio:+.5})",
                c.term,
                cur_kld - kld,
                c.mb
            );
            if ratio > 0.0 && best.map_or(true, |(_, _, r)| ratio > r) {
                best = Some((i, kld, ratio));
            }
        }
        let Some((i, kld, ratio)) = best else {
            eprintln!("round {round}: no candidate improves — stopping");
            break;
        };
        let c = cands.remove(i);
        eprintln!(
            "round {round}: ACCEPT {} (KLD {cur_kld:.4} → {kld:.4}, {:.1} MB, ratio {ratio:.5})",
            c.term, c.mb
        );
        history.push(serde_json::json!({
            "round": round, "term": c.term, "mb": c.mb, "kld_before": cur_kld, "kld_after": kld,
        }));
        cur_kld = kld;
        accepted.push(c);
        if cands.is_empty() {
            break;
        }
    }

    let spec = build_spec(&accepted, None);
    let total_mb: f64 = accepted.iter().map(|c| c.mb).sum();
    println!("\n── precision-search result ──────────────────────────────");
    println!("model   {}", args.model);
    println!("policy  {spec}");
    println!("cost    +{total_mb:.1} MB over all-Q4");
    println!("KLD     {cur_kld:.4} nats on the search subset (re-verify on the full corpus!)");
    if let Some(out) = &args.out {
        let doc = serde_json::json!({
            "model": args.model,
            "spec": spec,
            "delta_mb": total_mb,
            "search_subset_kld": cur_kld,
            "history": history,
            "note": "search ranks on a subset; re-verify with kld-eval --arm gpu-prec:<spec>",
        });
        std::fs::write(out, serde_json::to_string_pretty(&doc)?)?;
        println!("wrote   {out}");
    }
    Ok(())
}

/// Assemble `q4,<accepted…>[,<extra>]` — base Q4 plus the chosen upgrades.
fn build_spec(accepted: &[Cand], extra: Option<&Cand>) -> String {
    let mut s = String::from("q4");
    for c in accepted.iter().chain(extra) {
        s.push(',');
        s.push_str(&c.term);
    }
    s
}