metrale_model_engine/kimi_k3_loader/
dry_run.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Dry run of a Kimi K3 safetensors `index.json` weight map, loading no tensors:
4//! lists the distinct shard files, and fails when a class of `TEXT_CLASSES` is missing or a key
5//! matches no known class.
6//!
7//! Owner: model-engine Kimi K3 loader.
8//! Invariants: none beyond the types.
9
10use std::collections::{BTreeMap, BTreeSet};
11
12use anyhow::{Context, Result, bail};
13use serde_json::Value;
14
15use super::classes::{ClassKind, TEXT_CLASSES, classify};
16
17/// 2026-09-25: Result of walking a K3 `weight_map`. In a returned report `missing_text_classes`
18/// and `unknown_keys` are always empty: either one being non-empty is an error instead.
19#[derive(Debug, Clone)]
20pub struct KimiK3DryRun {
21    pub shards: Vec<String>,
22    pub present_text_classes: Vec<String>,
23    pub missing_text_classes: Vec<String>,
24    /// 2026-09-25: Vision and mm_projector classes seen; filled only when `language_model_only`
25    /// is set.
26    pub ignored_vision_classes: Vec<String>,
27    pub unknown_keys: Vec<String>,
28}
29
30/// 2026-09-25: Parse a safetensors `index.json` and dry-run its `weight_map`.
31pub fn dry_run_index_json(index_json: &str, language_model_only: bool) -> Result<KimiK3DryRun> {
32    let raw: Value = serde_json::from_str(index_json).context("invalid safetensors index.json")?;
33    let map = raw
34        .get("weight_map")
35        .and_then(Value::as_object)
36        .context("safetensors index.json missing weight_map")?;
37    let weight_map: BTreeMap<String, String> = map
38        .iter()
39        .map(|(k, v)| {
40            let shard = v
41                .as_str()
42                .with_context(|| format!("weight_map[{k}] is not a string"))?
43                .to_string();
44            Ok((k.clone(), shard))
45        })
46        .collect::<Result<_>>()?;
47    dry_run_weight_map(&weight_map, language_model_only)
48}
49
50pub fn dry_run_weight_map(
51    weight_map: &BTreeMap<String, String>,
52    language_model_only: bool,
53) -> Result<KimiK3DryRun> {
54    let mut shards = BTreeSet::new();
55    let mut present_text = BTreeSet::new();
56    let mut ignored_vision = BTreeSet::new();
57    let mut unknown_keys = Vec::new();
58
59    for (name, shard) in weight_map {
60        shards.insert(shard.clone());
61        match classify(name) {
62            Some((class, ClassKind::Text)) => {
63                present_text.insert(class.to_string());
64            }
65            Some((class, ClassKind::Vision)) => {
66                if language_model_only {
67                    ignored_vision.insert(class.to_string());
68                }
69            }
70            None => unknown_keys.push(name.clone()),
71        }
72    }
73
74    let missing_text_classes: Vec<String> = TEXT_CLASSES
75        .iter()
76        .filter(|c| !present_text.contains(**c))
77        .map(|c| (*c).to_string())
78        .collect();
79
80    if !missing_text_classes.is_empty() {
81        bail!(
82            "K3 weight map missing required text classes ({}): {}",
83            missing_text_classes.len(),
84            missing_text_classes.join(", ")
85        );
86    }
87    if !unknown_keys.is_empty() {
88        bail!(
89            "K3 weight map has unrecognised keys ({}): {}",
90            unknown_keys.len(),
91            unknown_keys
92                .iter()
93                .take(8)
94                .cloned()
95                .collect::<Vec<_>>()
96                .join(", ")
97        );
98    }
99
100    Ok(KimiK3DryRun {
101        shards: shards.into_iter().collect(),
102        present_text_classes: present_text.into_iter().collect(),
103        missing_text_classes,
104        ignored_vision_classes: ignored_vision.into_iter().collect(),
105        unknown_keys,
106    })
107}
108
109/// 2026-09-25: Fail unless the map names exactly `n` shard files. The TSV in `docs/k3` records
110/// 96 shards for the full checkpoint.
111#[cfg(test)]
112pub fn require_shard_count(report: &KimiK3DryRun, n: usize) -> Result<()> {
113    if report.shards.len() != n {
114        bail!(
115            "K3 weight map expected {n} shards, got {}",
116            report.shards.len()
117        );
118    }
119    Ok(())
120}
121
122#[cfg(test)]
123mod tests {
124    use super::super::classes::example_key;
125    use super::*;
126    use std::collections::BTreeMap;
127
128    /// 2026-09-25: Synthetic 96-shard map: one key per text class, plus keys of the first class
129    /// until every shard file appears.
130    fn synthetic_96_shard_map() -> BTreeMap<String, String> {
131        let mut map = BTreeMap::new();
132        let n_shards = 96usize;
133        for (i, class) in TEXT_CLASSES.iter().enumerate() {
134            let shard = format!("model-{:05}-of-00096.safetensors", (i % n_shards) + 1);
135            map.insert(example_key(class), shard);
136        }
137        let pad_class = TEXT_CLASSES[0];
138        for i in TEXT_CLASSES.len()..n_shards {
139            let shard = format!("model-{:05}-of-00096.safetensors", i + 1);
140            let key = pad_class
141                .replacen('*', "0", 1)
142                .replacen('*', &i.to_string(), 1);
143            map.insert(key, shard);
144        }
145        map
146    }
147
148    #[test]
149    fn kimi_k3_weight_map_dry_run_lists_96_shards() {
150        let map = synthetic_96_shard_map();
151        let report = dry_run_weight_map(&map, true).expect("complete text map");
152        assert_eq!(report.shards.len(), 96);
153        assert_eq!(report.present_text_classes.len(), TEXT_CLASSES.len());
154        assert!(report.missing_text_classes.is_empty());
155        assert!(report.unknown_keys.is_empty());
156        assert_eq!(report.shards[0], "model-00001-of-00096.safetensors");
157        assert_eq!(report.shards[95], "model-00096-of-00096.safetensors");
158    }
159
160    #[test]
161    fn kimi_k3_weight_map_dry_run_refuses_missing_text_class() {
162        let mut map = synthetic_96_shard_map();
163        let drop = example_key("language_model.model.output_attn_res_proj.weight");
164        map.remove(&drop);
165        let err = dry_run_weight_map(&map, true).unwrap_err().to_string();
166        assert!(
167            err.contains("missing required text classes"),
168            "unexpected error: {err}"
169        );
170        assert!(err.contains("output_attn_res_proj"), "{err}");
171    }
172
173    #[test]
174    fn kimi_k3_weight_map_dry_run_ignores_vision_when_language_model_only() {
175        let mut map = synthetic_96_shard_map();
176        map.insert(
177            "vision_tower.patch_embed.proj.weight".into(),
178            "model-00001-of-00096.safetensors".into(),
179        );
180        let report = dry_run_weight_map(&map, true).expect("vision ignored");
181        assert!(
182            report
183                .ignored_vision_classes
184                .iter()
185                .any(|c| c == "vision_tower.patch_embed.proj.weight")
186        );
187    }
188
189    #[test]
190    fn kimi_k3_weight_map_dry_run_index_json() {
191        let map = synthetic_96_shard_map();
192        let index = serde_json::json!({ "weight_map": map });
193        let report = dry_run_index_json(&index.to_string(), true).unwrap();
194        assert_eq!(report.shards.len(), 96);
195        require_shard_count(&report, 96).unwrap();
196    }
197
198    #[test]
199    fn kimi_k3_weight_map_dry_run_refuses_95_shards() {
200        let mut map = synthetic_96_shard_map();
201        map.retain(|_, shard| !shard.starts_with("model-00096-of-"));
202        let report = dry_run_weight_map(&map, true).expect("text classes still present");
203        let err = require_shard_count(&report, 96).unwrap_err().to_string();
204        assert!(err.contains("expected 96 shards"), "{err}");
205        assert!(err.contains("got 95"), "{err}");
206    }
207}
metrale_model_engine/kimi_k3_loader/dry_run.rs

metrale_model_engine/kimi_k3_loader/
dry_run.rs
