spark_model/
mtp_layout.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Checkpoint-level detection of MTP / next-token-prediction weights.
4//!
5//! Atlas binds MTP through three unrelated loader paths — the Qwen-shaped
6//! `MtpWeights` vec, DeepSeek-V4's `mtp.0.*` module, and GLM-5.3's
7//! `layers.{num_hidden_layers}` block — but the "did the user get what they
8//! asked for" check read only the first of them. On GLM-5.3 that produced
9//!
10//! ```text
11//! GLM-5.3 MTP draft module loaded (layers.45)
12//! `--speculative` was requested but no MTP weights were loaded for this model
13//! ```
14//!
15//! in the same startup log: the module HAD loaded, from
16//! `model.language_model.layers.45.*`. This module is the layout-aware
17//! predicate the check should have been asking, kept separate from any one
18//! loader so a new architecture only has to be named here once.
19
20use atlas_core::config::ModelConfig;
21use spark_runtime::weights::WeightStore;
22
23/// The MTP weight layout a checkpoint ships, if any.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum MtpLayout {
26    /// `mtp.*` — Qwen3.5 (`mtp.fc.weight`, `mtp.layers.0.*`) and the
27    /// DeepSeek `mtp.0.*` multi-module spelling.
28    MtpPrefix,
29    /// Transformer layer(s) one past the skeleton: GLM-5.3's
30    /// `model.language_model.layers.45.*` at `num_hidden_layers = 45`, and
31    /// the DeepSeek-V3 `model.layers.61.*` nextn block.
32    ExtraLayer { first: usize, count: usize },
33}
34
35/// Layer index of a `*.layers.N.*` / `layers.N.*` tensor name.
36///
37/// Same acceptance as `preflight::validate_layer_coverage`: any prefix
38/// (`model.`, `model.language_model.`, `backbone.`) or none at all
39/// (Mistral consolidated checkpoints).
40fn layer_index(name: &str) -> Option<usize> {
41    let tail = if let Some(pos) = name.find(".layers.") {
42        &name[pos + ".layers.".len()..]
43    } else {
44        name.strip_prefix("layers.")?
45    };
46    let end = tail.find('.')?;
47    tail[..end].parse().ok()
48}
49
50/// Detect the MTP layout from raw tensor names.
51///
52/// `mtp.*` wins when both are present — that is the layout the generic
53/// `load_mtp_weights_multi` path binds.
54pub fn detect<'a>(
55    names: impl Iterator<Item = &'a str>,
56    num_hidden_layers: usize,
57) -> Option<MtpLayout> {
58    let mut extras: Vec<usize> = Vec::new();
59    let mut has_prefix = false;
60    for n in names {
61        if n.starts_with("mtp.") {
62            has_prefix = true;
63        } else if let Some(idx) = layer_index(n)
64            && idx >= num_hidden_layers
65        {
66            extras.push(idx);
67        }
68    }
69    if has_prefix {
70        return Some(MtpLayout::MtpPrefix);
71    }
72    if extras.is_empty() {
73        return None;
74    }
75    extras.sort_unstable();
76    extras.dedup();
77    Some(MtpLayout::ExtraLayer {
78        first: extras[0],
79        count: extras.len(),
80    })
81}
82
83/// [`detect`] over a loaded [`WeightStore`].
84///
85/// 🪤 On an EP worker the store has already been filtered by expert index,
86/// but never by layer — the MTP block's non-expert tensors are present on
87/// every rank, so this answers the same on rank 0 and rank 1.
88pub fn detect_in_store(store: &WeightStore, config: &ModelConfig) -> Option<MtpLayout> {
89    detect(store.names(), config.num_hidden_layers)
90}
91
92#[cfg(test)]
93mod tests {
94    use super::{MtpLayout, detect, layer_index};
95
96    #[test]
97    fn glm5_next_layer_45_block_is_detected() {
98        // The regression this module exists for: GLM-5.3 nests the text stack
99        // under `model.language_model.` and puts MTP at layers.45.
100        let names = [
101            "model.language_model.layers.44.self_attn.q_proj.weight",
102            "model.language_model.layers.45.eh_proj.weight",
103            "model.language_model.layers.45.shared_head.norm.weight",
104            "model.language_model.layers.45.mlp.experts.7.down_proj.weight",
105            "lm_head.weight",
106        ];
107        assert_eq!(
108            detect(names.iter().copied(), 45),
109            Some(MtpLayout::ExtraLayer {
110                first: 45,
111                count: 1
112            }),
113        );
114    }
115
116    #[test]
117    fn qwen_mtp_prefix_still_detected() {
118        let names = ["model.layers.0.self_attn.q_proj.weight", "mtp.fc.weight"];
119        assert_eq!(
120            detect(names.iter().copied(), 48),
121            Some(MtpLayout::MtpPrefix)
122        );
123    }
124
125    #[test]
126    fn deepseek_multi_module_prefix_still_detected() {
127        let names = ["mtp.0.self_attn.q_proj.weight", "mtp.1.eh_proj.weight"];
128        assert_eq!(
129            detect(names.iter().copied(), 61),
130            Some(MtpLayout::MtpPrefix)
131        );
132    }
133
134    #[test]
135    fn deepseek_v3_nextn_extra_layer_detected() {
136        let names = [
137            "model.layers.60.mlp.gate.weight",
138            "model.layers.61.eh_proj.weight",
139        ];
140        assert_eq!(
141            detect(names.iter().copied(), 61),
142            Some(MtpLayout::ExtraLayer {
143                first: 61,
144                count: 1
145            }),
146        );
147    }
148
149    #[test]
150    fn plain_checkpoint_has_no_mtp() {
151        let names = [
152            "model.layers.0.self_attn.q_proj.weight",
153            "model.layers.47.mlp.down_proj.weight",
154            "model.embed_tokens.weight",
155            "lm_head.weight",
156        ];
157        assert_eq!(detect(names.iter().copied(), 48), None);
158    }
159
160    #[test]
161    fn vision_tower_blocks_are_not_mistaken_for_mtp() {
162        // `model.visual.blocks.N.*` has no `.layers.` segment; a checkpoint
163        // whose vision depth exceeds num_hidden_layers must NOT read as MTP.
164        let names = [
165            "model.language_model.layers.0.self_attn.q_proj.weight",
166            "model.visual.blocks.23.attn.proj.weight",
167            "model.visual.merger.proj.weight",
168        ];
169        assert_eq!(detect(names.iter().copied(), 45), None);
170    }
171
172    #[test]
173    fn multi_extra_layers_are_counted() {
174        let names = [
175            "model.layers.61.eh_proj.weight",
176            "model.layers.62.eh_proj.weight",
177        ];
178        assert_eq!(
179            detect(names.iter().copied(), 61),
180            Some(MtpLayout::ExtraLayer {
181                first: 61,
182                count: 2
183            }),
184        );
185    }
186
187    #[test]
188    fn unprefixed_mistral_layer_names_parse() {
189        assert_eq!(layer_index("layers.3.attention.wq.weight"), Some(3));
190        assert_eq!(
191            layer_index("model.language_model.layers.45.eh_proj.weight"),
192            Some(45)
193        );
194        assert_eq!(layer_index("lm_head.weight"), None);
195        assert_eq!(layer_index("model.visual.blocks.2.attn.proj.weight"), None);
196    }
197}
spark_model/mtp_layout.rs

spark_model/
mtp_layout.rs
