spark_model/weight_loader/kimi_k3/
dry_run.rs

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

spark_model/weight_loader/kimi_k3/
dry_run.rs
