metrale_model_arch/weight_loader/glm5_next_load/
loader.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `impl ModelWeightLoader for Glm5NextWeightLoader`.
4//!
5//! Owner: model-arch weight loader.
6//! Invariants: none beyond the types.
7
8use super::*;
9
10impl ModelWeightLoader for Glm5NextWeightLoader {
11    /// 2026-09-25: True only when `metrale_config::glm_vision_enabled()`
12    /// (`METRALE_GLM_VISION`). The GLM-5.3 config parser reads the same function
13    /// before it parses `vision_config`; this method takes no config, so the
14    /// gate is the environment.
15    fn binds_vision_encoder(&self) -> bool {
16        metrale_config::glm_vision_enabled()
17    }
18
19    /// 2026-09-25: Bind the `model.visual.*` tower; see
20    /// [`crate::weight_loader::glm5_next_vision`].
21    fn load_vision_encoder(
22        &self,
23        store: &WeightStore,
24        config: &ModelConfig,
25        gpu: &dyn GpuBackend,
26    ) -> Result<Option<metrale_model_layers::layers::VisionTower>> {
27        crate::weight_loader::glm5_next_vision::load_glm5_next_vision(store, config, gpu)
28    }
29
30    /// 2026-09-25: Keep the MTP layer's BF16 routed-expert weights off the
31    /// device (`is_full_width_mtp_expert`); `bind_expert` quantises them from
32    /// disk. The layer is `num_hidden_layers`. A checkpoint whose MTP experts
33    /// are U8 defers nothing.
34    fn defer_predicate(
35        &self,
36        config: &ModelConfig,
37    ) -> Option<metrale_model_weights::weights::DeferHook> {
38        let num_layers = config.num_hidden_layers;
39        Some(std::sync::Arc::new(
40            move |name: &str, dtype: WeightDtype| is_full_width_mtp_expert(name, dtype, num_layers),
41        ))
42    }
43
44    /// 2026-09-25: DSA, KDA and the MLP all shard under TP (see `glm5_next_load.rs`);
45    /// routed experts are also split by EP (`local_expert_range`).
46    fn supports_tp(&self) -> bool {
47        true
48    }
49
50    fn load_layers(
51        &self,
52        store: &WeightStore,
53        config: &ModelConfig,
54        gpu: &dyn GpuBackend,
55        _layer_kv_dtypes: &[KvCacheDtype],
56    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
57        let skeleton = Glm5NextTextSkeleton::from_config(config)?;
58        // 2026-09-25: `l2_eps` and `chunk` are fixed here; every other field is
59        // read from the config.
60        let kda_cfg = Glm5NextKdaConfig {
61            hidden: config.hidden_size,
62            heads: config.linear_num_value_heads,
63            head_dim: config.linear_value_head_dim,
64            conv_kernel: config.linear_conv_kernel_dim,
65            gate_lower_bound: config.linear_gate_lower_bound,
66            rms_norm_eps: config.rms_norm_eps as f32,
67            l2_eps: 1e-6,
68            chunk: 32,
69        };
70        kda_cfg.validate()?;
71        // 2026-09-25: `gate_rank` is the row count of layer 0's `f_a_proj`; the
72        // config has no such key.
73        let gate_rank = {
74            let n = qualify(0, "self_attn.f_a_proj.weight");
75            let t = store.get(&n).with_context(|| {
76                format!("glm5_next: {n} is needed to size the KDA gate bottleneck")
77            })?;
78            *t.shape.first().context("f_a_proj has no rows")?
79        };
80        let kda_plan = KdaTpPlan::from_config(config, gate_rank)?;
81        let dsa_cfg = Glm5NextDsaConfig::from_config(config)?;
82        let mlp_cfg = Glm5NextMlpConfig::from_config(config)?;
83
84        let kda_kernels = Glm5NextKdaKernels::resolve(gpu)?;
85        let dsa_kernels = Glm5NextDsaKernels::resolve(gpu)?;
86        let dsa_layer_kernels = Glm5NextDsaLayerKernels::resolve(gpu)?;
87        let mlp_kernels = Glm5NextMlpKernels::resolve(gpu)?;
88        let mhc_kernels_probe = Glm5NextMhcKernels::resolve(gpu)?;
89        let rms_norm_k = gpu.kernel("rms_norm_vanilla", "rms_norm_vanilla")?;
90        // 2026-09-25: Only a layer without mHC (the MTP layer) uses `add_k`, and it
91        // fails there if the kernel is missing; `try_kernel` lets a target without
92        // it still build the text stack.
93        let add_k = metrale_model_layers::layers::try_kernel(gpu, "bf16_add", "bf16_add_inplace");
94
95        // 2026-09-25: Every KDA layer shares one workspace (one `kda_cfg`). The
96        // workspaces are sized for `verify_k` rows, the largest of the batched
97        // GEMV's `DENSE_GEMV_BATCHM_MAX_M`, `PREFILL_ROWS` and `prefill_rows()`
98        // (`METRALE_GLM_PREFILL_ROWS`).
99        let verify_k = (metrale_model_layers::layers::ops::DENSE_GEMV_BATCHM_MAX_M as usize)
100            .max(crate::glm5next_layer::PREFILL_ROWS)
101            .max(crate::glm5next_layer::prefill_rows());
102        let kda_ws = std::sync::Arc::new(crate::glm5next_kda::Glm5NextKdaWorkspace::new(
103            gpu, &kda_cfg, verify_k,
104        )?);
105
106        // 2026-09-25: Unless `METRALE_GLM_MLP_WS_SHARED=0`, one MLP workspace serves
107        // every layer; otherwise each layer allocates its own. Either way it is
108        // allocated here, at load, before the KV pool is sized.
109        let mlp_ws_bytes = crate::glm5next_mlp::forward::mlp_ws_total_bytes(&mlp_cfg, verify_k);
110        let shared_mlp_ws = if crate::glm5next_mlp::forward::mlp_ws_shared() {
111            tracing::info!(
112                "GLM MLP workspace: SHARED, 1 x {:.1} MB for {} layers at {verify_k} rows \
113                 (per-layer would be {:.1} MB)",
114                mlp_ws_bytes as f64 / 1e6,
115                skeleton.layers.len(),
116                (mlp_ws_bytes * skeleton.layers.len()) as f64 / 1e6,
117            );
118            Some(std::sync::Arc::new(
119                crate::glm5next_mlp::forward::Glm5NextMlpWorkspace::new(gpu, &mlp_cfg, verify_k)?,
120            ))
121        } else {
122            tracing::warn!(
123                "GLM MLP workspace: PER-LAYER, {} x {:.1} MB at {verify_k} rows",
124                skeleton.layers.len(),
125                mlp_ws_bytes as f64 / 1e6,
126            );
127            None
128        };
129
130        let dsa_plan = crate::glm5next_dsa::tp::DsaTpPlan::new(
131            config.tp_rank,
132            config.tp_world_size.max(1),
133            &dsa_cfg,
134        )?;
135        let last = skeleton.layers.len() - 1;
136        let mut out: Vec<Box<dyn TransformerLayer>> = Vec::with_capacity(skeleton.layers.len());
137
138        // 2026-09-25: The KV pool has `num_attention_layers()` slots, which counts
139        // the sparse-attention layers only, so a DSA layer addresses it by its
140        // ordinal among DSA layers, not by its model index.
141        let mut attn_layer_idx = 0usize;
142
143        for sl in &skeleton.layers {
144            let idx = sl.index;
145            let t_layer = std::time::Instant::now();
146            let src = LayerSource::collect(gpu, store, idx)
147                .with_context(|| format!("glm5_next: collecting layer {idx}"))?;
148            let t_collect = t_layer.elapsed();
149
150            let mixer = match sl.mixer {
151                Mixer::Kda => {
152                    let sharded = KdaShardedSource::new(&src, &kda_plan)?;
153                    let (w, _report) = bind_kda_weights(gpu, &kda_cfg, idx, &sharded)?;
154                    Glm5NextMixer::Kda {
155                        layer: Box::new(Glm5NextKdaLayer::new(idx, kda_cfg, w, kda_kernels)?),
156                        ws: kda_ws.clone(),
157                        cfg: kda_cfg,
158                    }
159                }
160                Mixer::Dsa => {
161                    let load = |n: &str| src.f32(n);
162                    let w = build_dsa_weights(gpu, &dsa_cfg, &dsa_plan, &load)?;
163                    Glm5NextMixer::Dsa(Box::new(Glm5NextDsaLayer {
164                        persist_bt: std::env::var("METRALE_GLM_DSA_ALLOC_PER_STEP").as_deref()
165                            != Ok("1"),
166                        cfg: dsa_cfg,
167                        weights: w,
168                        kernels: dsa_layer_kernels,
169                        select_kernels: dsa_kernels,
170                        decode_kernel:
171                            crate::glm5next_dsa::attend::Glm5NextDsaDecodeKernel::resolve(gpu)?,
172                        workspace: crate::glm5next_dsa::layer::Glm5NextDsaWorkspace::new(
173                            gpu, &dsa_cfg, verify_k,
174                        )?,
175                        layer_idx: idx,
176                        attn_layer_idx: {
177                            let a = attn_layer_idx;
178                            attn_layer_idx += 1;
179                            a
180                        },
181                        rms_eps: config.rms_norm_eps as f32,
182                        kv_scale: 1.0,
183                    }))
184                }
185            };
186
187            let t_mixer = t_layer.elapsed();
188
189            let load = |n: &str| src.f32(n);
190            let mlp = match sl.mlp {
191                Mlp::Dense => Glm5NextMlpSite::Dense(mlp_build::build_dense_mlp(
192                    gpu,
193                    &mlp_cfg,
194                    config.tp_rank,
195                    config.intermediate_size,
196                    "mlp",
197                    &load,
198                )?),
199                Mlp::RoutedMoe => {
200                    let expert = |id: usize| bind_expert(gpu, store, idx, id);
201                    Glm5NextMlpSite::Moe(Box::new(mlp_build::build_moe(
202                        gpu,
203                        &mlp_cfg,
204                        config.tp_rank,
205                        config.shared_expert_intermediate_size,
206                        &load,
207                        &expert,
208                    )?))
209                }
210            };
211
212            let t_mlp = t_layer.elapsed();
213            tracing::info!(
214                "glm5_next layer {idx} built: collect {:.2}s mixer {:.2}s mlp {:.2}s \
215                 (mixer={:?} mlp={:?})",
216                t_collect.as_secs_f64(),
217                (t_mixer - t_collect).as_secs_f64(),
218                (t_mlp - t_mixer).as_secs_f64(),
219                sl.mixer,
220                sl.mlp,
221            );
222
223            let mhc = if sl.hyper_connection {
224                Some(Glm5NextMhc {
225                    kernels: mhc_kernels_probe,
226                    attn: bind_mhc_site(gpu, &src, "attn", config.hc_mult, config.hidden_size)?,
227                    ffn: bind_mhc_site(gpu, &src, "ffn", config.hc_mult, config.hidden_size)?,
228                    hc_mult: config.hc_mult,
229                    sinkhorn_iters: config.hc_sinkhorn_iters,
230                    hc_eps: config.hc_eps,
231                })
232            } else {
233                None
234            };
235
236            out.push(Box::new(Glm5NextLayer {
237                layer_idx: idx,
238                mixer,
239                mlp,
240                mlp_cfg,
241                mlp_kernels,
242                mlp_ws: match &shared_mlp_ws {
243                    Some(ws) => ws.clone(),
244                    None => std::sync::Arc::new(
245                        crate::glm5next_mlp::forward::Glm5NextMlpWorkspace::new(
246                            gpu, &mlp_cfg, verify_k,
247                        )?,
248                    ),
249                },
250                mhc,
251                input_norm: upload_f32_as_bf16(gpu, &src.f32("input_layernorm.weight")?)?,
252                post_attn_norm: upload_f32_as_bf16(
253                    gpu,
254                    &src.f32("post_attention_layernorm.weight")?,
255                )?,
256                rms_norm_k,
257                add_k,
258                rms_eps: config.rms_norm_eps as f32,
259                hidden: config.hidden_size,
260                mixer_all_reduce: match sl.mixer {
261                    Mixer::Kda => kda_plan.needs_output_all_reduce(),
262                    Mixer::Dsa => dsa_plan.needs_output_all_reduce(),
263                },
264                is_first: idx == 0,
265                is_last: idx == last,
266            }));
267        }
268        Ok(out)
269    }
270
271    fn load_embedding(
272        &self,
273        store: &WeightStore,
274        _config: &ModelConfig,
275        _gpu: &dyn GpuBackend,
276    ) -> Result<DenseWeight> {
277        dense(store, "model.language_model.embed_tokens.weight")
278    }
279
280    fn load_final_norm(
281        &self,
282        store: &WeightStore,
283        _config: &ModelConfig,
284        _gpu: &dyn GpuBackend,
285    ) -> Result<DenseWeight> {
286        dense(store, "model.language_model.norm.weight")
287    }
288
289    fn load_lm_head(
290        &self,
291        store: &WeightStore,
292        _config: &ModelConfig,
293        _gpu: &dyn GpuBackend,
294    ) -> Result<DenseWeight> {
295        dense(store, "lm_head.weight")
296    }
297
298    /// 2026-09-25: `None`: the GLM-5.3 MTP layer is loaded by
299    /// `glm5_next_mtp::load_glm5next_mtp_module`, not as `MtpWeights`.
300    fn load_mtp_weights(
301        &self,
302        _store: &WeightStore,
303        _config: &ModelConfig,
304        _gpu: &dyn GpuBackend,
305    ) -> Result<Option<crate::weight_loader::MtpWeights>> {
306        Ok(None)
307    }
308
309    /// 2026-09-25: Free the store tensors `is_reuploaded` matches, then those
310    /// `is_quantized_expert_weight` matches.
311    fn prune_after_load(
312        &self,
313        store: &mut WeightStore,
314        config: &ModelConfig,
315        gpu: &dyn GpuBackend,
316    ) -> Result<()> {
317        let n = config.num_hidden_layers;
318        let (count, bytes) = store.free_matching(gpu, |name| is_reuploaded(name, n))?;
319        tracing::info!(
320            "glm5_next: released {count} store tensors ({:.2} GB) already re-uploaded by the \
321             binders; routed experts and the MTP block kept",
322            bytes as f64 / 1e9,
323        );
324        let quantized: std::collections::BTreeSet<String> = store
325            .names()
326            .filter(|n| {
327                store
328                    .get(n)
329                    .is_ok_and(|t| is_quantized_expert_weight(n, t.dtype))
330            })
331            .map(str::to_string)
332            .collect();
333        if !quantized.is_empty() {
334            let (qcount, qbytes) = store.free_matching(gpu, |name| quantized.contains(name))?;
335            tracing::info!(
336                "glm5_next: released {qcount} full-width BF16 routed-expert tensors \
337                 ({:.2} GB) quantised to NVFP4 at bind time",
338                qbytes as f64 / 1e9,
339            );
340        }
341        Ok(())
342    }
343}
metrale_model_arch/weight_loader/glm5_next_load/loader.rs

metrale_model_arch/weight_loader/glm5_next_load/
loader.rs
