spark_model/weight_loader/
glm5_next.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GLM-5.3-Flash (`glm5_next`) tensor accounting.
4//!
5//! Slice 1 scope: **classify every tensor in the checkpoint, deliberately.**
6//! No loading, no device work, no forward pass.
7//!
8//! Reference checkpoint: `LibertAIDAI/GLM-5.3-Flash-NVFP4` snapshot
9//! `9e0d74e3cef17f634e84fb8e2223707e02616290` — 120 shards, **113,074 tensors**,
10//! 407 distinct name patterns. Every claim here is scoped to that checkpoint.
11//!
12//! The point of this module is that "we loaded the model" and "we accounted for
13//! the checkpoint" are different statements. A loader that silently ignores an
14//! unrecognised tensor will happily produce a model that is quietly wrong — the
15//! failure mode that cost the DS4F campaign weeks. So the contract is:
16//! `classify` returns `None` for anything it has not been taught, and the test
17//! suite fails if a single tensor in the reference checkpoint lands there.
18
19use std::collections::BTreeMap;
20
21/// What a tensor is *for*. Deliberately coarse — Slice 1 proves coverage, not
22/// placement.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
24pub enum TensorRole {
25    /// Token embedding / final norm / lm_head.
26    Embedding,
27    LmHead,
28    FinalNorm,
29    /// KDA linear-attention mixer (34 of the 45 text layers).
30    KdaProjection,
31    KdaConv,
32    KdaDecay,
33    KdaGate,
34    KdaNorm,
35    /// NoPE sparse MLA (11 text layers + the MTP layer).
36    MlaProjection,
37    MlaNorm,
38    /// DSA top-k indexer that fronts each sparse-MLA layer.
39    Indexer,
40    IndexerNorm,
41    /// MoE router + experts + shared expert.
42    MoeRouter,
43    MoeExpert,
44    MoeShared,
45    /// Dense FFN (layers 0..2, `first_k_dense_replace = 3`).
46    DenseFfn,
47    /// Per-layer norms and mHC (hyper-connection) parameters.
48    LayerNorm,
49    HyperConnection,
50    /// MTP head at layer 45. 🪤 GLM does NOT use `mtp.0.*`.
51    MtpProjection,
52    MtpNorm,
53    /// Vision tower — present in the checkpoint, out of scope for the text port.
54    Vision,
55}
56
57impl TensorRole {
58    /// Norm-family roles. Used by the RMSNorm sanity pass: a tensor whose name
59    /// says "norm" must land in one of these, or the classifier is lying about
60    /// something. (Hazard carried from `notavault-atlas`: RMSNorm weights have
61    /// silently corrupted Atlas numbers before.)
62    pub fn is_norm(self) -> bool {
63        matches!(
64            self,
65            TensorRole::FinalNorm
66                | TensorRole::KdaNorm
67                | TensorRole::MlaNorm
68                | TensorRole::IndexerNorm
69                | TensorRole::LayerNorm
70                | TensorRole::MtpNorm
71        )
72    }
73
74    /// Text-model roles Atlas must eventually implement. Vision is excluded.
75    pub fn is_text_model(self) -> bool {
76        !matches!(self, TensorRole::Vision)
77    }
78}
79
80/// Replace every all-numeric path segment with `#`, so `layers.7.` and
81/// `layers.N.` (our census canonicalisation) collapse to the same key.
82fn normalize(name: &str) -> String {
83    name.split('.')
84        .map(|seg| {
85            if seg.is_empty() {
86                seg
87            } else if seg.chars().all(|c| c.is_ascii_digit()) || seg == "N" || seg == "E" {
88                "#"
89            } else {
90                seg
91            }
92        })
93        .collect::<Vec<_>>()
94        .join(".")
95}
96
97/// Classify one checkpoint tensor name.
98///
99/// Returns `None` for anything unrecognised — callers MUST treat that as a
100/// hard error, never as "skip it".
101pub fn classify(name: &str) -> Option<TensorRole> {
102    let n = normalize(name);
103    let s = n.as_str();
104
105    // ---- non-layer -------------------------------------------------------
106    if s == "lm_head.weight" {
107        return Some(TensorRole::LmHead);
108    }
109    if s == "model.language_model.embed_tokens.weight" {
110        return Some(TensorRole::Embedding);
111    }
112    if s == "model.language_model.norm.weight" {
113        return Some(TensorRole::FinalNorm);
114    }
115    if s.starts_with("model.visual.") || s.starts_with("model.vision") {
116        return Some(TensorRole::Vision);
117    }
118
119    let rest = s.strip_prefix("model.language_model.layers.#.")?;
120
121    // ---- MTP head (layer 45) --------------------------------------------
122    // These names appear ONLY on the MTP layer. Layer-index checking is the
123    // caller's job (see `is_mtp_only_name`); here we classify by role.
124    match rest {
125        "eh_proj.weight" => return Some(TensorRole::MtpProjection),
126        "enorm.weight" | "hnorm.weight" | "shared_head.norm.weight" => {
127            return Some(TensorRole::MtpNorm);
128        }
129        _ => {}
130    }
131
132    // ---- KDA linear attention -------------------------------------------
133    match rest {
134        "self_attn.q_proj.weight"
135        | "self_attn.k_proj.weight"
136        | "self_attn.v_proj.weight"
137        | "self_attn.b_proj.weight"
138        | "self_attn.f_a_proj.weight"
139        | "self_attn.f_b_proj.weight"
140        | "self_attn.g_a_proj.weight"
141        | "self_attn.g_b_proj.weight" => return Some(TensorRole::KdaProjection),
142        "self_attn.q_conv1d.weight" | "self_attn.k_conv1d.weight" | "self_attn.v_conv1d.weight" => {
143            return Some(TensorRole::KdaConv);
144        }
145        "self_attn.A_log" | "self_attn.dt_bias" => return Some(TensorRole::KdaDecay),
146        "self_attn.o_norm.weight" => return Some(TensorRole::KdaNorm),
147        _ => {}
148    }
149
150    // ---- NoPE sparse MLA -------------------------------------------------
151    match rest {
152        "self_attn.q_a_proj.weight"
153        | "self_attn.q_b_proj.weight"
154        | "self_attn.kv_a_proj_with_mqa.weight"
155        | "self_attn.kv_b_proj.weight" => return Some(TensorRole::MlaProjection),
156        "self_attn.q_a_layernorm.weight" | "self_attn.kv_a_layernorm.weight" => {
157            return Some(TensorRole::MlaNorm);
158        }
159        // o_proj is shared by both mixer families (46 occurrences = 34 KDA +
160        // 11 DSA + 1 MTP), so it cannot discriminate; it is an output
161        // projection either way.
162        "self_attn.o_proj.weight" => return Some(TensorRole::MlaProjection),
163        _ => {}
164    }
165
166    // ---- DSA indexer -----------------------------------------------------
167    if let Some(idx) = rest.strip_prefix("self_attn.indexer.") {
168        return Some(match idx {
169            "k_norm.weight" | "k_norm.bias" => TensorRole::IndexerNorm,
170            _ => TensorRole::Indexer,
171        });
172    }
173
174    // ---- mHC hyper-connections ------------------------------------------
175    if rest.starts_with("hc_") {
176        return Some(TensorRole::HyperConnection);
177    }
178
179    // ---- norms -----------------------------------------------------------
180    if rest == "input_layernorm.weight" || rest == "post_attention_layernorm.weight" {
181        return Some(TensorRole::LayerNorm);
182    }
183
184    // ---- MoE / dense FFN -------------------------------------------------
185    if let Some(mlp) = rest.strip_prefix("mlp.") {
186        if mlp.starts_with("gate.") || mlp == "gate.weight" {
187            return Some(TensorRole::MoeRouter);
188        }
189        if mlp.starts_with("experts.#.") {
190            return Some(TensorRole::MoeExpert);
191        }
192        if mlp.starts_with("shared_experts.") {
193            return Some(TensorRole::MoeShared);
194        }
195        // Bare gate/up/down on a layer with no expert dimension = dense FFN.
196        if mlp.starts_with("gate_proj.")
197            || mlp.starts_with("up_proj.")
198            || mlp.starts_with("down_proj.")
199        {
200            return Some(TensorRole::DenseFfn);
201        }
202    }
203
204    None
205}
206
207/// Names that exist ONLY on the MTP layer. Used to locate the MTP layer index
208/// from the checkpoint itself rather than assuming one.
209///
210/// 🪤 Do not look for `mtp.0.*`: GLM-5.3 has **zero** such tensors. Verified by
211/// scanning all 120 shard headers of the reference checkpoint.
212pub fn is_mtp_only_name(name: &str) -> bool {
213    let n = normalize(name);
214    let Some(rest) = n.strip_prefix("model.language_model.layers.#.") else {
215        return false;
216    };
217    matches!(
218        rest,
219        "eh_proj.weight" | "enorm.weight" | "hnorm.weight" | "shared_head.norm.weight"
220    )
221}
222
223/// Result of accounting a whole checkpoint's tensor-name list.
224#[derive(Debug, Default)]
225pub struct Accounting {
226    pub total: usize,
227    pub by_role: BTreeMap<String, usize>,
228    /// Anything `classify` refused. MUST be empty.
229    pub unknown: Vec<String>,
230    /// Layer indices that carry MTP-only tensors.
231    pub mtp_layers: Vec<usize>,
232}
233
234/// Account for a list of `(tensor_name, count)` pairs.
235pub fn account<'a, I>(names: I) -> Accounting
236where
237    I: IntoIterator<Item = (&'a str, usize)>,
238{
239    let mut acc = Accounting::default();
240    let mut mtp = std::collections::BTreeSet::new();
241    for (name, count) in names {
242        acc.total += count;
243        match classify(name) {
244            Some(role) => {
245                *acc.by_role.entry(format!("{role:?}")).or_insert(0) += count;
246            }
247            None => acc.unknown.push(name.to_string()),
248        }
249        if is_mtp_only_name(name)
250            && let Some(i) = layer_index(name)
251        {
252            mtp.insert(i);
253        }
254    }
255    acc.mtp_layers = mtp.into_iter().collect();
256    acc
257}
258
259/// Extract the numeric layer index from a real tensor name (`None` for the
260/// canonicalised `layers.N.` form, which carries no index).
261pub fn layer_index(name: &str) -> Option<usize> {
262    let mut it = name.split('.');
263    while let Some(seg) = it.next() {
264        if seg == "layers" {
265            return it.next().and_then(|s| s.parse().ok());
266        }
267    }
268    None
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    #[test]
276    fn normalizes_numeric_and_canonical_segments_alike() {
277        assert_eq!(
278            normalize("model.language_model.layers.7.self_attn.o_proj.weight"),
279            normalize("model.language_model.layers.N.self_attn.o_proj.weight")
280        );
281        assert_eq!(
282            normalize("model.language_model.layers.45.mlp.experts.12.down_proj.weight"),
283            normalize("model.language_model.layers.N.mlp.experts.E.down_proj.weight")
284        );
285    }
286
287    #[test]
288    fn mtp_is_found_by_layer_name_not_by_mtp_prefix() {
289        assert!(is_mtp_only_name(
290            "model.language_model.layers.45.eh_proj.weight"
291        ));
292        assert!(!is_mtp_only_name("model.language_model.layers.3.eh_proj"));
293        // The DeepSeek convention must NOT be what we key on.
294        assert!(!is_mtp_only_name("model.layers.mtp.0.eh_proj.weight"));
295        assert_eq!(
296            layer_index("model.language_model.layers.45.eh_proj.weight"),
297            Some(45)
298        );
299    }
300
301    #[test]
302    fn unknown_tensor_is_refused_not_skipped() {
303        assert_eq!(
304            classify("model.language_model.layers.4.self_attn.wat"),
305            None
306        );
307        let acc = account([("model.language_model.layers.4.self_attn.wat", 1)]);
308        assert_eq!(acc.unknown.len(), 1);
309    }
310
311    #[test]
312    fn norm_tensors_land_in_norm_roles() {
313        for n in [
314            "model.language_model.norm.weight",
315            "model.language_model.layers.0.self_attn.o_norm.weight",
316            "model.language_model.layers.3.self_attn.q_a_layernorm.weight",
317            "model.language_model.layers.3.self_attn.indexer.k_norm.weight",
318            "model.language_model.layers.5.input_layernorm.weight",
319            "model.language_model.layers.45.enorm.weight",
320        ] {
321            let r = classify(n).unwrap_or_else(|| panic!("unclassified: {n}"));
322            assert!(r.is_norm(), "{n} classified as non-norm {r:?}");
323        }
324    }
325}
spark_model/weight_loader/glm5_next.rs

spark_model/weight_loader/
glm5_next.rs
