metrale_circuit/venn/
checkpoint.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: Check a target instance against the checkpoint it names: the layer kinds its
4//! `config.json` lists, and the precision every bound module declares (ModelOpt
5//! `hf_quant_config.json` or `config.json` `quantization_config`), resolved through the engine's
6//! own policy at the `declared` tier with every W8A8 kernel capability present, so the answer
7//! is the declared weight and activation formats.
8//!
9//! Owner: metrale-circuit (venn).
10//! Invariants: pure (the caller reads the files); every disagreement is listed in one error.
11
12use metrale_config::{DeclaredPrecisionPlan, WeightQuantPolicy};
13
14use super::VennError;
15use super::repo::config_layer_kinds;
16use crate::Loaded;
17use crate::instances::Instance;
18use crate::precision::EdgePrecision;
19use crate::precision_policy::{PolicyPrecision, caps_named, tier_named};
20
21/// 2026-09-29: Every KernelCaps bit: the declared formats, whatever kernels are built today.
22const ALL_CAPS: [&str; 4] = [
23    "w8a8_decode",
24    "w8a8_moe_decode",
25    "w8a8_block_scaled_decode",
26    "fp8_lm_head_batched",
27];
28
29/// 2026-09-29: `Ok` when `config` (and `hf_quant`, if the checkpoint ships one) agree with the
30/// instance and its loaded circuit.
31pub fn check(
32    inst: &Instance,
33    loaded: &Loaded,
34    config: &str,
35    hf_quant: Option<&str>,
36) -> Result<(), VennError> {
37    let bad = |d: String| VennError::Checkpoint(d);
38    let cfg: serde_json::Value =
39        serde_json::from_str(config).map_err(|e| bad(format!("config.json: {e}")))?;
40    let mut problems = Vec::new();
41    let kinds = config_layer_kinds(&cfg)?;
42    if kinds != inst.shape.layer_kinds {
43        problems.push(format!(
44            "layer kinds: config.json has {} layers, INSTANCES.toml {} (or the kinds differ)",
45            kinds.len(),
46            inst.shape.layer_kinds.len()
47        ));
48    }
49    let qc = match hf_quant {
50        Some(text) => serde_json::from_str::<serde_json::Value>(text)
51            .map_err(|e| bad(format!("hf_quant_config.json: {e}")))?,
52        None => cfg
53            .get("quantization_config")
54            .cloned()
55            .unwrap_or(serde_json::Value::Null),
56    };
57    let plan = if qc.is_null() {
58        DeclaredPrecisionPlan::UNDECLARED
59    } else {
60        DeclaredPrecisionPlan::from_quantization_config(&qc)
61            .map_err(|e| bad(format!("quantization: {e:#}")))?
62    };
63    let caps: Vec<String> = ALL_CAPS.iter().map(|c| c.to_string()).collect();
64    let policy = WeightQuantPolicy::new(
65        tier_named("declared").map_err(|e| bad(e.to_string()))?,
66        &plan,
67        caps_named(&caps).map_err(|e| bad(e.to_string()))?,
68    );
69    let declared = PolicyPrecision::new(policy, &[]);
70    let c = &loaded.circuit;
71    for n in c.nodes.iter().filter(|n| n.weight.is_some()) {
72        let Some(module) = n.binding.first() else {
73            continue;
74        };
75        let want = declared.linear(&module.replace('*', "0"));
76        let act = n.inputs.first().map(|&e| c.edges[e].format);
77        if n.weight != Some(want.weight) || act != Some(want.activation) {
78            problems.push(format!(
79                "{module}: the circuit runs {}/{}, the checkpoint declares {}/{}",
80                n.weight.map(|w| w.name()).unwrap_or_default(),
81                act.map(|a| a.name()).unwrap_or_default(),
82                want.weight,
83                want.activation
84            ));
85        }
86    }
87    problems.dedup();
88    if problems.is_empty() {
89        Ok(())
90    } else {
91        Err(bad(format!(
92            "{} disagrees with the checkpoint:\n  {}",
93            inst.recipe,
94            problems.join("\n  ")
95        )))
96    }
97}
metrale_circuit/venn/checkpoint.rs

metrale_circuit/venn/
checkpoint.rs
