metrale_model_arch/weight_loader/
glm5_next_load.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Glm5NextWeightLoader`: builds the GLM-5.3 text stack from a
4//! `WeightStore`, plus the helpers the MTP loader shares.
5//!
6//! `Glm5NextTextSkeleton::from_config` fixes each layer's mixer (KDA or DSA)
7//! and MLP (dense or routed MoE) from the config's `layer_types` and
8//! `mlp_only_layers`; the loader builds what it lists.
9//!
10//! Under TP, DSA shards through `DsaTpPlan`, the MLP through the rank passed to
11//! `glm5next_mlp::build`, and KDA through [`KdaShardedSource`], which slices
12//! the host bytes before `bind_kda_weights` sees them. At TP > 1 both mixer
13//! plans report `needs_output_all_reduce`, and the layer all-reduces the mixer
14//! output (`Glm5NextLayer::mixer_all_reduce`).
15//!
16//! Owner: model-arch weight loader.
17//! Invariants: none beyond the types.
18
19use anyhow::{Context, Result, bail};
20use metrale_cache::kv_cache::KvCacheDtype;
21use metrale_config::ModelConfig;
22use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
23use metrale_model_weights::weights::{WeightDtype, WeightStore, WeightTensor};
24
25use super::ModelWeightLoader;
26use crate::glm5next_dsa::build::build_dsa_weights;
27use crate::glm5next_dsa::layer::{Glm5NextDsaLayer, Glm5NextDsaLayerKernels};
28use crate::glm5next_dsa::{Glm5NextDsaConfig, Glm5NextDsaKernels};
29use crate::glm5next_kda::binding::{KdaDtype, KdaTensorSource, RawTensor, bind_kda_weights};
30use crate::glm5next_kda::tp::KdaTpPlan;
31use crate::glm5next_kda::tp_bind::KdaShardedSource;
32use crate::glm5next_kda::{Glm5NextKdaConfig, Glm5NextKdaKernels, Glm5NextKdaLayer};
33use crate::glm5next_layer::{Glm5NextLayer, Glm5NextMhc, Glm5NextMixer, Glm5NextMlpSite};
34use crate::glm5next_mhc::{Glm5NextMhcKernels, Glm5NextMhcSiteWeights, mhc_mix_max_tokens, mix_hc};
35use crate::glm5next_mlp::weights::{Glm5NextExpertWeights, Nvfp4Proj};
36use crate::glm5next_mlp::{Glm5NextMlpConfig, Glm5NextMlpKernels, build as mlp_build};
37use crate::glm5next_skeleton::{Glm5NextTextSkeleton, Mixer, Mlp};
38use metrale_model_layers::layer::TransformerLayer;
39use metrale_model_layers::weight_map::DenseWeight;
40
41#[cfg(test)]
42mod defer_hook_tests;
43mod expert_quant;
44#[cfg(test)]
45mod export_layout_tests;
46mod loader;
47mod nvfp4_dequant;
48mod nvfp4_quant;
49#[cfg(test)]
50mod plan_cast_tests;
51mod plan_dtype;
52
53use expert_quant::{quantize_deferred_expert_proj, quantize_expert_proj};
54
55pub struct Glm5NextWeightLoader;
56
57/// 2026-09-25: The full name of layer `layer`'s tensor `leaf`, under
58/// `model.language_model.layers.`.
59fn qualify(layer: usize, leaf: &str) -> String {
60    format!("model.language_model.layers.{layer}.{leaf}")
61}
62
63/// 2026-09-25: Whether a store tensor belongs to a text layer (index below
64/// `num_layers`) and is not a routed expert. `load_layers` copies those to the
65/// host (`LayerSource::collect`) and the binders upload their own buffers, so
66/// `prune_after_load` frees the store's copies.
67///
68/// Two families are never matched:
69/// - `mlp.experts.*`: `bind_expert` binds U8 experts zero-copy from the store's
70///   pointers.
71/// - layer `num_layers`, the MTP layer: `glm5_next_mtp::load_glm5next_mtp_module`
72///   reads it.
73fn is_reuploaded(name: &str, num_layers: usize) -> bool {
74    let Some(rest) = name.strip_prefix("model.language_model.layers.") else {
75        return false;
76    };
77    let Some((idx, rel)) = rest.split_once('.') else {
78        return false;
79    };
80    let Ok(idx) = idx.parse::<usize>() else {
81        return false;
82    };
83    idx < num_layers && !rel.starts_with("mlp.experts.")
84}
85
86/// 2026-09-25: Whether a store tensor is a BF16 routed-expert projection
87/// weight, which [`bind_expert`] quantised to NVFP4 rather than binding. The
88/// dtype is the whole test: a U8 expert is the kernel's operand and is never
89/// matched. `prune_after_load` frees the matches; `factory::build` calls it
90/// after the MTP module is bound. With a store that honoured
91/// [`Glm5NextWeightLoader::defer_predicate`], the MTP layer's BF16 experts are
92/// never resident.
93fn is_quantized_expert_weight(name: &str, dtype: WeightDtype) -> bool {
94    dtype == WeightDtype::BF16
95        && name.starts_with("model.language_model.layers.")
96        && name.contains(".mlp.experts.")
97        && name.ends_with("_proj.weight")
98}
99
100/// 2026-09-25: Whether a store tensor is a BF16 `mlp.experts.*_proj.weight` of
101/// layer `num_layers`, the MTP layer: the family
102/// [`Glm5NextWeightLoader::defer_predicate`] keeps off the device. The shared
103/// expert (`mlp.shared_experts.*`) does not match; it is read through
104/// [`LayerSource`].
105fn is_full_width_mtp_expert(name: &str, dtype: WeightDtype, num_layers: usize) -> bool {
106    if dtype != WeightDtype::BF16 || !name.ends_with("_proj.weight") {
107        return false;
108    }
109    let Some(rest) = name.strip_prefix("model.language_model.layers.") else {
110        return false;
111    };
112    let Some((idx, rel)) = rest.split_once('.') else {
113        return false;
114    };
115    idx.parse::<usize>().ok() == Some(num_layers) && rel.starts_with("mlp.experts.")
116}
117
118/// 2026-09-25: Read a device tensor back as host bytes.
119fn host_bytes(gpu: &dyn GpuBackend, t: &WeightTensor) -> Result<Vec<u8>> {
120    let mut b = vec![0u8; t.byte_size()];
121    gpu.copy_d2h(t.ptr, &mut b)?;
122    Ok(b)
123}
124
125/// 2026-09-25: Read a BF16 or FP32 device tensor back as host `f32`, by the
126/// tensor's own dtype; any other dtype is an error.
127fn host_f32(gpu: &dyn GpuBackend, t: &WeightTensor, what: &str) -> Result<Vec<f32>> {
128    let b = host_bytes(gpu, t)?;
129    match t.dtype {
130        WeightDtype::BF16 => Ok(b
131            .chunks_exact(2)
132            .map(|c| half::bf16::from_bits(u16::from_le_bytes([c[0], c[1]])).to_f32())
133            .collect()),
134        WeightDtype::FP32 => Ok(b
135            .chunks_exact(4)
136            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
137            .collect()),
138        other => bail!(
139            "{what}: dtype {other:?} cannot be read as f32 without a conversion this loader refuses to guess"
140        ),
141    }
142}
143
144fn upload_f32(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
145    let b: Vec<u8> = v.iter().flat_map(|x| x.to_le_bytes()).collect();
146    let p = gpu.alloc(b.len().max(1))?;
147    gpu.copy_h2d(&b, p)?;
148    Ok(p)
149}
150
151/// 2026-09-25: One layer's tensors under layer-relative names, copied to the
152/// host (routed experts by name only).
153pub(super) struct LayerSource {
154    names: Vec<String>,
155    tensors: std::collections::BTreeMap<String, (WeightDtype, Vec<usize>, Vec<u8>)>,
156    /// 2026-09-25: Copies, at the width the binders bind at, of the tensors
157    /// stored at another width (see [`plan_dtype`]). Only
158    /// [`KdaTensorSource::get`] reads them; [`Self::f32`] reads the stored bytes.
159    plan_cast: std::collections::BTreeMap<String, (WeightDtype, Vec<u8>)>,
160}
161
162impl LayerSource {
163    pub(super) fn collect(gpu: &dyn GpuBackend, store: &WeightStore, layer: usize) -> Result<Self> {
164        let prefix = format!("model.language_model.layers.{layer}.");
165        let mut names = Vec::new();
166        let mut tensors = std::collections::BTreeMap::new();
167        let mut plan_cast = std::collections::BTreeMap::new();
168        let rels: Vec<String> = store
169            .names()
170            .filter_map(|n| n.strip_prefix(&prefix).map(|r| r.to_string()))
171            .collect();
172        for rel in rels {
173            let rel = rel.as_str();
174            // 2026-09-25: Routed experts are listed but not copied: `bind_expert` reads
175            // them from the store.
176            if rel.starts_with("mlp.experts.") {
177                names.push(rel.to_string());
178                continue;
179            }
180            names.push(rel.to_string());
181            let t = store.get(&format!("{prefix}{rel}"))?;
182            let bytes = host_bytes(gpu, t)?;
183            if let Some(cast) = plan_dtype::cast_to_plan_dtype(rel, t.dtype, &bytes)? {
184                plan_cast.insert(rel.to_string(), cast);
185            }
186            tensors.insert(rel.to_string(), (t.dtype, t.shape.clone(), bytes));
187        }
188        Ok(Self {
189            names,
190            tensors,
191            plan_cast,
192        })
193    }
194
195    pub(super) fn f32(&self, name: &str) -> Result<Vec<f32>> {
196        let (dtype, shape, bytes) = self
197            .tensors
198            .get(name)
199            .with_context(|| format!("missing tensor {name}"))?;
200        match dtype {
201            WeightDtype::BF16 => Ok(bytes
202                .chunks_exact(2)
203                .map(|c| half::bf16::from_bits(u16::from_le_bytes([c[0], c[1]])).to_f32())
204                .collect()),
205            WeightDtype::FP32 => Ok(bytes
206                .chunks_exact(4)
207                .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
208                .collect()),
209            // 2026-09-25: A packed-NVFP4 tensor is dequantized, so `f32` accepts
210            // it as well as BF16 and FP32.
211            WeightDtype::UInt8 => self.dequant_packed_nvfp4(name, shape, bytes),
212            other => bail!("{name}: dtype {other:?} is not a plain float tensor"),
213        }
214    }
215
216    /// 2026-09-25: One packed-NVFP4 `.weight` of this layer as `f32`, scaled by
217    /// its `.weight_scale` (F8_E4M3 blocks) and scalar `.weight_scale_2`
218    /// siblings. A missing or differently typed sibling is an error. Reached
219    /// only from the `UInt8` arm of [`Self::f32`].
220    fn dequant_packed_nvfp4(&self, name: &str, shape: &[usize], bytes: &[u8]) -> Result<Vec<f32>> {
221        let base = name.strip_suffix(".weight").with_context(|| {
222            format!("{name}: packed NVFP4 must be a `.weight`, with scale siblings beside it")
223        })?;
224        let (scale_dtype, _, scale_bytes) = self
225            .tensors
226            .get(&format!("{base}.weight_scale"))
227            .with_context(|| format!("{name} is packed NVFP4 but {base}.weight_scale is absent"))?;
228        if *scale_dtype != WeightDtype::FP8E4M3 {
229            bail!("{base}.weight_scale is {scale_dtype:?}, expected F8_E4M3 block scales");
230        }
231        let s2 = self.f32(&format!("{base}.weight_scale_2"))?;
232        let [scale_2] = s2[..] else {
233            bail!("{base}.weight_scale_2 is not a scalar");
234        };
235        nvfp4_dequant::dequant_nvfp4_to_f32(name, bytes, shape, scale_bytes, scale_2)
236    }
237}
238
239impl KdaTensorSource for LayerSource {
240    fn get(&self, name: &str) -> Option<RawTensor<'_>> {
241        let (dtype, shape, bytes) = match self.plan_cast.get(name) {
242            // 2026-09-25: A plan-width copy, when one exists, replaces the stored bytes.
243            Some((dtype, bytes)) => (dtype, &self.tensors.get(name)?.1, bytes),
244            None => {
245                let (dtype, shape, bytes) = self.tensors.get(name)?;
246                (dtype, shape, bytes)
247            }
248        };
249        // 2026-09-25: Only BF16 and FP32 tensors are offered; any other dtype reads
250        // as absent rather than cast.
251        let dtype = match dtype {
252            WeightDtype::BF16 => KdaDtype::Bf16,
253            WeightDtype::FP32 => KdaDtype::F32,
254            _ => return None,
255        };
256        Some(RawTensor {
257            dtype,
258            shape: shape.clone(),
259            bytes,
260        })
261    }
262    fn names(&self) -> Vec<String> {
263        self.names.clone()
264    }
265}
266
267/// 2026-09-25: Bind one mHC site (`attn` or `ffn`) of a layer. `hc_{site}_fn`
268/// is uploaded as BF16 with `hc_fn_bf16` set; `base` and `scale` as F32. Each
269/// tensor's length is checked first.
270fn bind_mhc_site(
271    gpu: &dyn GpuBackend,
272    src: &LayerSource,
273    site: &str,
274    hc_mult: usize,
275    hidden: usize,
276) -> Result<Glm5NextMhcSiteWeights> {
277    let f = src.f32(&format!("hc_{site}_fn"))?;
278    let want = mix_hc(hc_mult) * hc_mult * hidden;
279    if f.len() != want {
280        bail!(
281            "hc_{site}_fn has {} elements, expected mix_hc({hc_mult}) * {hc_mult} * {hidden} = {want}",
282            f.len()
283        );
284    }
285    let base = src.f32(&format!("hc_{site}_base"))?;
286    if base.len() != mix_hc(hc_mult) {
287        bail!(
288            "hc_{site}_base has {} entries, expected mix_hc({hc_mult}) = {}",
289            base.len(),
290            mix_hc(hc_mult)
291        );
292    }
293    let scale = src.f32(&format!("hc_{site}_scale"))?;
294    if scale.len() != 3 {
295        bail!(
296            "hc_{site}_scale has {} entries, expected 3 (pre, post, comb)",
297            scale.len()
298        );
299    }
300    Ok(Glm5NextMhcSiteWeights {
301        hc_fn: upload_f32_as_bf16(gpu, &f)?,
302        hc_fn_bf16: true,
303        hc_scale: upload_f32(gpu, &scale)?,
304        hc_base: upload_f32(gpu, &base)?,
305        // 2026-09-25: The `hc_mix` to `hc_finish` scratch; one per site, so a
306        // layer's two sites never share it.
307        mix: gpu.alloc(mhc_mix_max_tokens() * mix_hc(hc_mult) * 4)?,
308    })
309}
310
311/// 2026-09-25: One routed expert's gate/up/down as NVFP4 triples. Each
312/// projection takes the first arm that applies:
313///
314/// 1. deferred by [`Glm5NextWeightLoader::defer_predicate`]: read from disk,
315///    quantised, uploaded;
316/// 2. resident U8: bound zero-copy with its `weight_scale` and scalar
317///    `weight_scale_2`;
318/// 3. resident BF16: read back and quantised; `prune_after_load` frees the
319///    source.
320///
321/// Any other dtype is an error.
322pub(super) fn bind_expert(
323    gpu: &dyn GpuBackend,
324    store: &WeightStore,
325    layer: usize,
326    id: usize,
327) -> Result<Glm5NextExpertWeights> {
328    let proj = |p: &str| -> Result<Nvfp4Proj> {
329        let base = format!("mlp.experts.{id}.{p}");
330        let wname = qualify(layer, &format!("{base}.weight"));
331        // 2026-09-25: Checked first: a deferred tensor has no store entry.
332        if let Some(d) = store.deferred(&wname) {
333            return quantize_deferred_expert_proj(gpu, store, d, &base);
334        }
335        let w = store.get(&wname)?;
336        match w.dtype {
337            WeightDtype::UInt8 => {
338                let scale = store.get(&qualify(layer, &format!("{base}.weight_scale")))?;
339                let s2 = store.get(&qualify(layer, &format!("{base}.weight_scale_2")))?;
340                let s2 = host_f32(gpu, s2, &format!("{base}.weight_scale_2"))?;
341                let [s2] = s2[..] else {
342                    bail!("{base}.weight_scale_2 is not a scalar");
343                };
344                Ok(Nvfp4Proj {
345                    packed: w.ptr,
346                    scale: scale.ptr,
347                    scale_2: s2,
348                })
349            }
350            WeightDtype::BF16 => quantize_expert_proj(gpu, store, w, &base),
351            other => bail!("{base}.weight is {other:?}, expected packed U8 NVFP4 or BF16"),
352        }
353    };
354    Ok(Glm5NextExpertWeights {
355        gate_proj: proj("gate_proj")?,
356        up_proj: proj("up_proj")?,
357        down_proj: proj("down_proj")?,
358    })
359}
360
361fn dense(store: &WeightStore, name: &str) -> Result<DenseWeight> {
362    Ok(DenseWeight {
363        weight: store.get(name)?.ptr,
364    })
365}
366
367pub(super) fn upload_f32_as_bf16(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
368    let b: Vec<u8> = v
369        .iter()
370        .flat_map(|x| half::bf16::from_f32(*x).to_le_bytes())
371        .collect();
372    let p = gpu.alloc(b.len().max(1))?;
373    gpu.copy_h2d(&b, p)?;
374    Ok(p)
375}
376
377#[cfg(test)]
378mod prune_tests;
379
380/// 2026-09-25: [`LayerSource::collect`], for the MTP loader.
381pub(super) fn layer_source(
382    gpu: &dyn GpuBackend,
383    store: &WeightStore,
384    layer: usize,
385) -> Result<LayerSource> {
386    LayerSource::collect(gpu, store, layer)
387}
388
389/// 2026-09-25: [`bind_expert`], for the MTP loader.
390pub(super) fn bind_expert_at(
391    gpu: &dyn GpuBackend,
392    store: &WeightStore,
393    layer: usize,
394    id: usize,
395) -> Result<Glm5NextExpertWeights> {
396    bind_expert(gpu, store, layer, id)
397}
398
399/// 2026-09-25: `upload_f32_as_bf16`, for the MTP loader.
400pub fn upload_bf16(gpu: &dyn GpuBackend, v: &[f32]) -> Result<DevicePtr> {
401    upload_f32_as_bf16(gpu, v)
402}
403
404#[cfg(test)]
405mod vision_capability_tests;
metrale_model_arch/weight_loader/glm5_next_load.rs

metrale_model_arch/weight_loader/
glm5_next_load.rs
