metrale_model_arch/
precision_schedule.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3#![allow(clippy::doc_lazy_continuation)]
4#![allow(clippy::doc_overindented_list_items)]
5
6//! 2026-09-25: A per-(layer, role) weight precision schedule: [`Role`] tags a
7//! tensor, [`Dtype`] is a requested precision, and [`PrecisionSchedule`]
8//! resolves a (layer, role) pair to a dtype.
9//!
10//! The only schedule in use is `PrecisionSchedule::default()`, which
11//! `ModelWeightLoader::precision_schedule` returns for every loader
12//! (`weight_loader/mod.rs`); it resolves every lookup to `Dtype::Inherit`. The
13//! qwen35 loader only logs a schedule that has an override
14//! (`weight_loader/qwen35/load_layers.rs`), and no loader calls `dtype_for`.
15//!
16//! Owner: model-arch (weight loading).
17//! Invariants:
18//! - `dtype_for` resolves in the order: sensitive-layer dtype (Attention,
19//!   Expert and SharedExpert tensors only), then the role's dtype, then the
20//!   global default; `Inherit` at a step falls through to the next.
21
22use std::collections::BTreeSet;
23
24/// 2026-09-25: Semantic role of a tensor, used for precision lookups. Adding
25/// a role means extending the matches in `Role::name`, `Role::from_str` and
26/// `PrecisionSchedule::dtype_for`.
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
28pub enum Role {
29    /// 2026-09-25: MoE router gate (hidden × num_experts).
30    Router,
31    /// 2026-09-25: Final unembedding (hidden × vocab).
32    LmHead,
33    /// 2026-09-25: Token embedding (vocab × hidden).
34    Embedding,
35    /// 2026-09-25: Attention Q/K/V/O projection.
36    Attention,
37    /// 2026-09-25: Expert FFN (gate / up / down per expert).
38    Expert,
39    /// 2026-09-25: Shared-expert FFN, when the model has one.
40    SharedExpert,
41    /// 2026-09-25: Layer norm scales (RMSNorm `weight`).
42    Norm,
43}
44
45impl Role {
46    pub fn name(&self) -> &'static str {
47        match self {
48            Role::Router => "router",
49            Role::LmHead => "lm_head",
50            Role::Embedding => "embedding",
51            Role::Attention => "attention",
52            Role::Expert => "expert",
53            Role::SharedExpert => "shared_expert",
54            Role::Norm => "norm",
55        }
56    }
57
58    /// 2026-09-25: Parse a role tag (the inverse of `name`). Returns `None`
59    /// for an unknown name.
60    #[allow(clippy::should_implement_trait)]
61    pub fn from_str(s: &str) -> Option<Self> {
62        match s {
63            "router" => Some(Role::Router),
64            "lm_head" => Some(Role::LmHead),
65            "embedding" => Some(Role::Embedding),
66            "attention" => Some(Role::Attention),
67            "expert" => Some(Role::Expert),
68            "shared_expert" => Some(Role::SharedExpert),
69            "norm" => Some(Role::Norm),
70            _ => None,
71        }
72    }
73}
74
75/// 2026-09-25: Target precision for a tensor. `Inherit` means no override;
76/// the other variants request that precision.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum Dtype {
79    /// 2026-09-25: No override: the loader keeps its own format detection.
80    Inherit,
81    Bf16,
82    Fp8,
83    Nvfp4,
84}
85
86impl Dtype {
87    #[allow(clippy::should_implement_trait)]
88    pub fn from_str(s: &str) -> Option<Self> {
89        match s {
90            "inherit" => Some(Dtype::Inherit),
91            "bf16" => Some(Dtype::Bf16),
92            "fp8" => Some(Dtype::Fp8),
93            "nvfp4" => Some(Dtype::Nvfp4),
94            _ => None,
95        }
96    }
97
98    /// 2026-09-25: True for every variant but `Inherit`.
99    pub fn is_override(&self) -> bool {
100        !matches!(self, Dtype::Inherit)
101    }
102}
103
104/// 2026-09-25: Per-(layer, role) precision schedule: one dtype per role, a
105/// global default, and a set of sensitive layers with their own dtype.
106#[derive(Debug, Clone)]
107pub struct PrecisionSchedule {
108    /// 2026-09-25: Dtype for any (layer, role) that nothing more specific
109    /// overrides.
110    default: Dtype,
111    /// 2026-09-25: Per-role dtypes, one field per `Role`.
112    router_dtype: Dtype,
113    lm_head_dtype: Dtype,
114    embedding_dtype: Dtype,
115    attention_dtype: Dtype,
116    expert_dtype: Dtype,
117    shared_expert_dtype: Dtype,
118    norm_dtype: Dtype,
119    /// 2026-09-25: Layer indices marked sensitive. `Attention`, `Expert` and
120    /// `SharedExpert` tensors in these layers get `sensitive_dtype` when it is
121    /// not `Inherit`.
122    sensitive_layers: BTreeSet<u16>,
123    sensitive_dtype: Dtype,
124}
125
126impl Default for PrecisionSchedule {
127    /// 2026-09-25: Empty schedule: every lookup returns `Inherit`.
128    fn default() -> Self {
129        Self {
130            default: Dtype::Inherit,
131            router_dtype: Dtype::Inherit,
132            lm_head_dtype: Dtype::Inherit,
133            embedding_dtype: Dtype::Inherit,
134            attention_dtype: Dtype::Inherit,
135            expert_dtype: Dtype::Inherit,
136            shared_expert_dtype: Dtype::Inherit,
137            norm_dtype: Dtype::Inherit,
138            sensitive_layers: BTreeSet::new(),
139            sensitive_dtype: Dtype::Inherit,
140        }
141    }
142}
143
144impl PrecisionSchedule {
145    /// 2026-09-25: Build a schedule from:
146    ///   - `router_dtype`: dtype for the MoE gate
147    ///   - `lm_head_dtype`: dtype for the final unembedding
148    ///   - `sensitive_block_dtype` + `sensitive_block_indices`: the dtype
149    ///     for the weight tensors of the listed layers
150    ///   - `default_dtype`: the global default
151    ///
152    /// The other roles' dtypes are `Inherit`.
153    pub fn build(
154        router_dtype: Dtype,
155        lm_head_dtype: Dtype,
156        sensitive_block_indices: &[u16],
157        sensitive_block_dtype: Dtype,
158        default_dtype: Dtype,
159    ) -> Self {
160        Self {
161            default: default_dtype,
162            router_dtype,
163            lm_head_dtype,
164            embedding_dtype: Dtype::Inherit,
165            attention_dtype: Dtype::Inherit,
166            expert_dtype: Dtype::Inherit,
167            shared_expert_dtype: Dtype::Inherit,
168            norm_dtype: Dtype::Inherit,
169            sensitive_layers: sensitive_block_indices.iter().copied().collect(),
170            sensitive_dtype: sensitive_block_dtype,
171        }
172    }
173
174    /// 2026-09-25: Resolve the target dtype for a tensor. `layer_idx = None`
175    /// is for non-layer tensors (embedding, lm_head, final norm).
176    /// Lookup order, skipping any step that yields `Inherit`:
177    ///   1. Sensitive-layer dtype (Attention, Expert and SharedExpert only)
178    ///   2. Per-role dtype
179    ///   3. Global default
180    pub fn dtype_for(&self, layer_idx: Option<u16>, role: Role) -> Dtype {
181        if let Some(li) = layer_idx
182            && matches!(role, Role::Attention | Role::Expert | Role::SharedExpert)
183            && self.sensitive_layers.contains(&li)
184            && self.sensitive_dtype.is_override()
185        {
186            return self.sensitive_dtype;
187        }
188        let role_dtype = match role {
189            Role::Router => self.router_dtype,
190            Role::LmHead => self.lm_head_dtype,
191            Role::Embedding => self.embedding_dtype,
192            Role::Attention => self.attention_dtype,
193            Role::Expert => self.expert_dtype,
194            Role::SharedExpert => self.shared_expert_dtype,
195            Role::Norm => self.norm_dtype,
196        };
197        if role_dtype.is_override() {
198            role_dtype
199        } else {
200            self.default
201        }
202    }
203
204    /// 2026-09-25: True iff some lookup can return a dtype other than
205    /// `Inherit`. The qwen35 loader logs the schedule when it is.
206    pub fn has_any_override(&self) -> bool {
207        self.default.is_override()
208            || self.router_dtype.is_override()
209            || self.lm_head_dtype.is_override()
210            || self.embedding_dtype.is_override()
211            || self.attention_dtype.is_override()
212            || self.expert_dtype.is_override()
213            || self.shared_expert_dtype.is_override()
214            || self.norm_dtype.is_override()
215            || (self.sensitive_dtype.is_override() && !self.sensitive_layers.is_empty())
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn default_schedule_is_all_inherit() {
225        let s = PrecisionSchedule::default();
226        assert!(!s.has_any_override());
227        assert_eq!(s.dtype_for(None, Role::LmHead), Dtype::Inherit);
228        assert_eq!(s.dtype_for(Some(0), Role::Router), Dtype::Inherit);
229        assert_eq!(s.dtype_for(Some(38), Role::Expert), Dtype::Inherit);
230    }
231
232    #[test]
233    fn role_override_wins_over_default() {
234        let s = PrecisionSchedule::build(
235            Dtype::Bf16,
236            Dtype::Bf16,
237            &[], // 2026-09-25: no sensitive layers
238            Dtype::Inherit,
239            Dtype::Nvfp4,
240        );
241        assert!(s.has_any_override());
242        assert_eq!(s.dtype_for(None, Role::Router), Dtype::Bf16);
243        assert_eq!(s.dtype_for(None, Role::LmHead), Dtype::Bf16);
244        assert_eq!(s.dtype_for(Some(15), Role::Expert), Dtype::Nvfp4);
245    }
246
247    #[test]
248    fn sensitive_layer_overrides_bulk_default_for_weights() {
249        let s = PrecisionSchedule::build(
250            Dtype::Bf16,
251            Dtype::Bf16,
252            &[0, 1, 38, 39],
253            Dtype::Fp8,
254            Dtype::Nvfp4,
255        );
256        assert_eq!(s.dtype_for(Some(0), Role::Expert), Dtype::Fp8);
257        assert_eq!(s.dtype_for(Some(38), Role::Attention), Dtype::Fp8);
258        assert_eq!(s.dtype_for(Some(1), Role::SharedExpert), Dtype::Fp8);
259        assert_eq!(s.dtype_for(Some(5), Role::Expert), Dtype::Nvfp4);
260
261        let sensitive_only = PrecisionSchedule::build(
262            Dtype::Inherit,
263            Dtype::Inherit,
264            &[7],
265            Dtype::Fp8,
266            Dtype::Inherit,
267        );
268        assert!(sensitive_only.has_any_override());
269        assert_eq!(
270            sensitive_only.dtype_for(Some(7), Role::Attention),
271            Dtype::Fp8
272        );
273    }
274
275    #[test]
276    fn sensitive_layer_does_not_override_router_or_lm_head() {
277        // 2026-09-25: Only Attention, Expert and SharedExpert tensors take the
278        // sensitive dtype; the other roles resolve to their role dtype or the
279        // default.
280        let s = PrecisionSchedule::build(Dtype::Bf16, Dtype::Bf16, &[0], Dtype::Fp8, Dtype::Nvfp4);
281        assert_eq!(s.dtype_for(Some(0), Role::Router), Dtype::Bf16);
282        assert_eq!(s.dtype_for(Some(0), Role::LmHead), Dtype::Bf16);
283        assert_eq!(s.dtype_for(Some(0), Role::Embedding), Dtype::Nvfp4);
284        assert_eq!(s.dtype_for(Some(0), Role::Norm), Dtype::Nvfp4);
285    }
286
287    #[test]
288    fn role_str_round_trips() {
289        for r in [
290            Role::Router,
291            Role::LmHead,
292            Role::Embedding,
293            Role::Attention,
294            Role::Expert,
295            Role::SharedExpert,
296            Role::Norm,
297        ] {
298            assert_eq!(Role::from_str(r.name()), Some(r));
299        }
300        assert_eq!(Role::from_str("nonsense"), None);
301    }
302
303    #[test]
304    fn dtype_str_parsing() {
305        assert_eq!(Dtype::from_str("bf16"), Some(Dtype::Bf16));
306        assert_eq!(Dtype::from_str("fp8"), Some(Dtype::Fp8));
307        assert_eq!(Dtype::from_str("nvfp4"), Some(Dtype::Nvfp4));
308        assert_eq!(Dtype::from_str("inherit"), Some(Dtype::Inherit));
309        assert_eq!(Dtype::from_str("bogus"), None);
310        assert!(!Dtype::Inherit.is_override());
311        assert!(Dtype::Bf16.is_override());
312    }
313}
metrale_model_arch/precision_schedule.rs

metrale_model_arch/
precision_schedule.rs
