metrale_model_arch/weight_loader/
glm5_next.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: GLM-5.3 (`glm5_next`) tensor accounting: `classify` gives the
4//! role of a checkpoint tensor name, and `account` totals a name list by role.
5//! Nothing is loaded.
6//!
7//! `classify` returns `None` for a name it does not know, so an unrecognised
8//! tensor is reported rather than skipped. The counts below are from the
9//! reference checkpoint's name fixture
10//! (`crates/model-engine/tests/fixtures/glm53-nvfp4-9e0d74e3-patterns.tsv`),
11//! whose test fails if any name in it is unknown.
12//!
13//! Owner: model-arch weight loader.
14//! Invariants: none beyond the types.
15
16use std::collections::BTreeMap;
17
18/// 2026-09-25: What a tensor is for.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
20pub enum TensorRole {
21    Embedding,
22    LmHead,
23    FinalNorm,
24    /// 2026-09-25: KDA linear-attention projections. The reference checkpoint
25    /// has 34 KDA layers.
26    KdaProjection,
27    KdaConv,
28    KdaDecay,
29    KdaGate,
30    KdaNorm,
31    /// 2026-09-25: Sparse-MLA projections, plus every mixer's `o_proj`. The
32    /// reference checkpoint has 12 sparse-MLA layers.
33    MlaProjection,
34    MlaNorm,
35    /// 2026-09-25: The DSA top-k indexer of a sparse-MLA layer
36    /// (`self_attn.indexer.*`, its `k_norm` excepted).
37    Indexer,
38    IndexerNorm,
39    MoeRouter,
40    MoeExpert,
41    MoeShared,
42    /// 2026-09-25: Dense FFN: `mlp.gate_proj/up_proj/down_proj` with no expert
43    /// index. The reference checkpoint has 3 dense-FFN layers.
44    DenseFfn,
45    /// 2026-09-25: `input_layernorm` and `post_attention_layernorm`.
46    LayerNorm,
47    HyperConnection,
48    /// 2026-09-25: The MTP layer's `eh_proj`. The MTP layer is a numbered
49    /// `layers.#` layer; the reference checkpoint has no `mtp.0.*` tensor.
50    MtpProjection,
51    MtpNorm,
52    /// 2026-09-25: The vision tower (`model.visual.*`, `model.vision*`), not part
53    /// of the text model.
54    Vision,
55}
56
57impl TensorRole {
58    /// 2026-09-25: Norm-family roles. `norm_sanity_pass` in
59    /// `glm5next_tensor_accounting.rs` requires every text tensor whose name
60    /// contains `norm` to land in one of them.
61    pub fn is_norm(self) -> bool {
62        matches!(
63            self,
64            TensorRole::FinalNorm
65                | TensorRole::KdaNorm
66                | TensorRole::MlaNorm
67                | TensorRole::IndexerNorm
68                | TensorRole::LayerNorm
69                | TensorRole::MtpNorm
70        )
71    }
72
73    /// 2026-09-25: Every role except `Vision`.
74    pub fn is_text_model(self) -> bool {
75        !matches!(self, TensorRole::Vision)
76    }
77}
78
79/// 2026-09-25: Replace every all-numeric path segment, and the fixture's
80/// placeholders `N` and `E`, with `#`, so real and fixture names give one key.
81fn normalize(name: &str) -> String {
82    name.split('.')
83        .map(|seg| {
84            if seg.is_empty() {
85                seg
86            } else if seg.chars().all(|c| c.is_ascii_digit()) || seg == "N" || seg == "E" {
87                "#"
88            } else {
89                seg
90            }
91        })
92        .collect::<Vec<_>>()
93        .join(".")
94}
95
96/// 2026-09-25: Classify one checkpoint tensor name. `None` means unrecognised;
97/// callers must treat it as an error, not as a tensor to skip.
98pub fn classify(name: &str) -> Option<TensorRole> {
99    let n = normalize(name);
100    let s = n.as_str();
101
102    if s == "lm_head.weight" {
103        return Some(TensorRole::LmHead);
104    }
105    if s == "model.language_model.embed_tokens.weight" {
106        return Some(TensorRole::Embedding);
107    }
108    if s == "model.language_model.norm.weight" {
109        return Some(TensorRole::FinalNorm);
110    }
111    if s.starts_with("model.visual.") || s.starts_with("model.vision") {
112        return Some(TensorRole::Vision);
113    }
114
115    let rest = s.strip_prefix("model.language_model.layers.#.")?;
116
117    // 2026-09-25: MTP-layer names are classified by name alone; `account` finds
118    // the MTP layer's index through `is_mtp_only_name`.
119    match rest {
120        "eh_proj.weight" => return Some(TensorRole::MtpProjection),
121        "enorm.weight" | "hnorm.weight" | "shared_head.norm.weight" => {
122            return Some(TensorRole::MtpNorm);
123        }
124        _ => {}
125    }
126
127    match rest {
128        "self_attn.q_proj.weight"
129        | "self_attn.k_proj.weight"
130        | "self_attn.v_proj.weight"
131        | "self_attn.b_proj.weight"
132        | "self_attn.f_a_proj.weight"
133        | "self_attn.f_b_proj.weight"
134        | "self_attn.g_a_proj.weight"
135        | "self_attn.g_b_proj.weight" => return Some(TensorRole::KdaProjection),
136        "self_attn.q_conv1d.weight" | "self_attn.k_conv1d.weight" | "self_attn.v_conv1d.weight" => {
137            return Some(TensorRole::KdaConv);
138        }
139        "self_attn.A_log" | "self_attn.dt_bias" => return Some(TensorRole::KdaDecay),
140        "self_attn.o_norm.weight" => return Some(TensorRole::KdaNorm),
141        _ => {}
142    }
143
144    match rest {
145        "self_attn.q_a_proj.weight"
146        | "self_attn.q_b_proj.weight"
147        | "self_attn.kv_a_proj_with_mqa.weight"
148        | "self_attn.kv_b_proj.weight" => return Some(TensorRole::MlaProjection),
149        "self_attn.q_a_layernorm.weight" | "self_attn.kv_a_layernorm.weight" => {
150            return Some(TensorRole::MlaNorm);
151        }
152        // 2026-09-25: Every mixer has an `o_proj` (46 in the reference
153        // checkpoint: 34 KDA, 12 sparse-MLA), so it does not identify the family.
154        "self_attn.o_proj.weight" => return Some(TensorRole::MlaProjection),
155        _ => {}
156    }
157
158    if let Some(idx) = rest.strip_prefix("self_attn.indexer.") {
159        return Some(match idx {
160            "k_norm.weight" | "k_norm.bias" => TensorRole::IndexerNorm,
161            _ => TensorRole::Indexer,
162        });
163    }
164
165    if rest.starts_with("hc_") {
166        return Some(TensorRole::HyperConnection);
167    }
168
169    if rest == "input_layernorm.weight" || rest == "post_attention_layernorm.weight" {
170        return Some(TensorRole::LayerNorm);
171    }
172
173    if let Some(mlp) = rest.strip_prefix("mlp.") {
174        if mlp.starts_with("gate.") || mlp == "gate.weight" {
175            return Some(TensorRole::MoeRouter);
176        }
177        if mlp.starts_with("experts.#.") {
178            return Some(TensorRole::MoeExpert);
179        }
180        if mlp.starts_with("shared_experts.") {
181            return Some(TensorRole::MoeShared);
182        }
183        if mlp.starts_with("gate_proj.")
184            || mlp.starts_with("up_proj.")
185            || mlp.starts_with("down_proj.")
186        {
187            return Some(TensorRole::DenseFfn);
188        }
189    }
190
191    None
192}
193
194/// 2026-09-25: Whether `name` is an MTP-layer tensor of a numbered layer
195/// (`eh_proj`, `enorm`, `hnorm`, `shared_head.norm`). `account` uses it to find
196/// the MTP layer's index. The reference checkpoint has no `mtp.0.*` tensor.
197pub fn is_mtp_only_name(name: &str) -> bool {
198    let n = normalize(name);
199    let Some(rest) = n.strip_prefix("model.language_model.layers.#.") else {
200        return false;
201    };
202    matches!(
203        rest,
204        "eh_proj.weight" | "enorm.weight" | "hnorm.weight" | "shared_head.norm.weight"
205    )
206}
207
208/// 2026-09-25: The result of [`account`].
209#[derive(Debug, Default)]
210pub struct Accounting {
211    pub total: usize,
212    pub by_role: BTreeMap<String, usize>,
213    /// 2026-09-25: Names `classify` refused.
214    pub unknown: Vec<String>,
215    /// 2026-09-25: Layer indices that carry MTP-layer tensors, ascending.
216    pub mtp_layers: Vec<usize>,
217}
218
219/// 2026-09-25: Total a list of `(tensor_name, count)` pairs by role.
220pub fn account<'a, I>(names: I) -> Accounting
221where
222    I: IntoIterator<Item = (&'a str, usize)>,
223{
224    let mut acc = Accounting::default();
225    let mut mtp = std::collections::BTreeSet::new();
226    for (name, count) in names {
227        acc.total += count;
228        match classify(name) {
229            Some(role) => {
230                *acc.by_role.entry(format!("{role:?}")).or_insert(0) += count;
231            }
232            None => acc.unknown.push(name.to_string()),
233        }
234        if is_mtp_only_name(name)
235            && let Some(i) = layer_index(name)
236        {
237            mtp.insert(i);
238        }
239    }
240    acc.mtp_layers = mtp.into_iter().collect();
241    acc
242}
243
244/// 2026-09-25: The number after the first `layers` segment, or `None` when it
245/// is absent or not a number (as in the fixture's `layers.N.`).
246pub fn layer_index(name: &str) -> Option<usize> {
247    let mut it = name.split('.');
248    while let Some(seg) = it.next() {
249        if seg == "layers" {
250            return it.next().and_then(|s| s.parse().ok());
251        }
252    }
253    None
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn normalizes_numeric_and_canonical_segments_alike() {
262        assert_eq!(
263            normalize("model.language_model.layers.7.self_attn.o_proj.weight"),
264            normalize("model.language_model.layers.N.self_attn.o_proj.weight")
265        );
266        assert_eq!(
267            normalize("model.language_model.layers.45.mlp.experts.12.down_proj.weight"),
268            normalize("model.language_model.layers.N.mlp.experts.E.down_proj.weight")
269        );
270    }
271
272    #[test]
273    fn mtp_is_found_by_layer_name_not_by_mtp_prefix() {
274        assert!(is_mtp_only_name(
275            "model.language_model.layers.45.eh_proj.weight"
276        ));
277        assert!(!is_mtp_only_name("model.language_model.layers.3.eh_proj"));
278        // 2026-09-25: An `mtp.0.*` name is not an MTP-layer name.
279        assert!(!is_mtp_only_name("model.layers.mtp.0.eh_proj.weight"));
280        assert_eq!(
281            layer_index("model.language_model.layers.45.eh_proj.weight"),
282            Some(45)
283        );
284    }
285
286    #[test]
287    fn unknown_tensor_is_refused_not_skipped() {
288        assert_eq!(
289            classify("model.language_model.layers.4.self_attn.wat"),
290            None
291        );
292        let acc = account([("model.language_model.layers.4.self_attn.wat", 1)]);
293        assert_eq!(acc.unknown.len(), 1);
294    }
295
296    #[test]
297    fn norm_tensors_land_in_norm_roles() {
298        for n in [
299            "model.language_model.norm.weight",
300            "model.language_model.layers.0.self_attn.o_norm.weight",
301            "model.language_model.layers.3.self_attn.q_a_layernorm.weight",
302            "model.language_model.layers.3.self_attn.indexer.k_norm.weight",
303            "model.language_model.layers.5.input_layernorm.weight",
304            "model.language_model.layers.45.enorm.weight",
305        ] {
306            let r = classify(n).unwrap_or_else(|| panic!("unclassified: {n}"));
307            assert!(r.is_norm(), "{n} classified as non-norm {r:?}");
308        }
309    }
310}
metrale_model_arch/weight_loader/glm5_next.rs

metrale_model_arch/weight_loader/
glm5_next.rs
