metrale_model_arch/weight_loader/
glm5_next_vision.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Bind the GLM-5.3 vision tower (`*.visual.*`) as BF16 device pointers.
4//!
5//! Owner: model-arch weight loader (GLM-5.3).
6//! Invariants:
7//! - Every bound tensor is BF16; any other dtype fails the load (`bf16`).
8//! - No tensor is copied: each field is the store's own pointer.
9//! - The patch-embed conv weight is read as `[hidden_size, 3 * temporal_patch_size *
10//!   patch_size^2]` and the downsample conv weight as `[out_hidden_size, hidden_size *
11//!   spatial_merge_size^2]`, over the same bytes; `check_shape` compares the leading dim
12//!   and the product of the trailing dims.
13
14use anyhow::{Context, Result};
15use metrale_config::{ModelConfig, VisionConfig};
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
17use metrale_model_weights::weights::{WeightDtype, WeightStore};
18
19use metrale_model_layers::layers::glm_vit::init::GlmVitGeometry;
20use metrale_model_layers::layers::{GlmVit, GlmVitBlock, GlmVitMerger, VisionTower};
21
22/// 2026-09-25: Bind one BF16 tensor by name, refusing any other dtype rather
23/// than reinterpreting its bytes.
24fn bf16(store: &WeightStore, name: &str) -> Result<DevicePtr> {
25    let w = store
26        .get(name)
27        .with_context(|| format!("glm5_next vision: missing tensor {name}"))?;
28    anyhow::ensure!(
29        w.dtype == WeightDtype::BF16,
30        "glm5_next vision: {name} is {:?}, but the GLM tower is BF16 throughout (347/347 \
31         tensors in the official NVFP4 export). A quantised tower needs a dequant path, not \
32         a cast of this pointer.",
33        w.dtype
34    );
35    Ok(w.ptr)
36}
37
38/// 2026-09-25: Check a bound tensor's on-disk shape against what the encoder
39/// will read.
40///
41/// Shapes are compared on the PRODUCT of the trailing dims, because the two
42/// tensors that matter here are stored with more axes than they are read with
43/// (the 5-D patch-embed conv and the 4-D downsample conv). Comparing the flat
44/// widths is the claim actually being made — "these bytes are a `[rows, cols]`
45/// GEMM weight" — and it catches a checkpoint whose patch or merge geometry
46/// moved, which is the case a name lookup cannot.
47fn check_shape(store: &WeightStore, name: &str, rows: usize, cols: usize) -> Result<()> {
48    let shape = store.get(name)?.shape.clone();
49    let got_rows = shape.first().copied().unwrap_or(0);
50    let got_cols: usize = shape.iter().skip(1).product::<usize>().max(1);
51    anyhow::ensure!(
52        got_rows == rows && (shape.len() == 1 || got_cols == cols),
53        "glm5_next vision: {name} has shape {shape:?}, which flattens to \
54         [{got_rows}, {got_cols}] — the encoder was built for [{rows}, {cols}] from \
55         vision_config. The checkpoint's tower geometry does not match its config."
56    );
57    Ok(())
58}
59
60fn load_block(store: &WeightStore, vp: &str, i: usize, v: &VisionConfig) -> Result<GlmVitBlock> {
61    let bp = format!("{vp}.blocks.{i}");
62    let h = v.hidden_size;
63    let head_dim = h / v.num_heads.max(1);
64    check_shape(store, &format!("{bp}.attn.qkv.weight"), 3 * h, h)?;
65    check_shape(store, &format!("{bp}.attn.q_norm.weight"), head_dim, 1)?;
66    check_shape(
67        store,
68        &format!("{bp}.mlp.gate_proj.weight"),
69        v.intermediate_size,
70        h,
71    )?;
72    Ok(GlmVitBlock {
73        norm1_w: bf16(store, &format!("{bp}.norm1.weight"))?,
74        qkv_w: bf16(store, &format!("{bp}.attn.qkv.weight"))?,
75        qkv_b: bf16(store, &format!("{bp}.attn.qkv.bias"))?,
76        q_norm_w: bf16(store, &format!("{bp}.attn.q_norm.weight"))?,
77        k_norm_w: bf16(store, &format!("{bp}.attn.k_norm.weight"))?,
78        proj_w: bf16(store, &format!("{bp}.attn.proj.weight"))?,
79        proj_b: bf16(store, &format!("{bp}.attn.proj.bias"))?,
80        norm2_w: bf16(store, &format!("{bp}.norm2.weight"))?,
81        gate_w: bf16(store, &format!("{bp}.mlp.gate_proj.weight"))?,
82        gate_b: bf16(store, &format!("{bp}.mlp.gate_proj.bias"))?,
83        up_w: bf16(store, &format!("{bp}.mlp.up_proj.weight"))?,
84        up_b: bf16(store, &format!("{bp}.mlp.up_proj.bias"))?,
85        down_w: bf16(store, &format!("{bp}.mlp.down_proj.weight"))?,
86        down_b: bf16(store, &format!("{bp}.mlp.down_proj.bias"))?,
87    })
88}
89
90fn load_merger(store: &WeightStore, vp: &str, v: &VisionConfig) -> Result<GlmVitMerger> {
91    let out = v.out_hidden_size;
92    let sms2 = v.spatial_merge_size * v.spatial_merge_size;
93    check_shape(
94        store,
95        &format!("{vp}.downsample.weight"),
96        out,
97        v.hidden_size * sms2,
98    )?;
99    check_shape(store, &format!("{vp}.merger.proj.weight"), out, out)?;
100    check_shape(
101        store,
102        &format!("{vp}.merger.gate_proj.weight"),
103        v.projection_intermediate_size,
104        out,
105    )?;
106    check_shape(
107        store,
108        &format!("{vp}.merger.down_proj.weight"),
109        out,
110        v.projection_intermediate_size,
111    )?;
112    Ok(GlmVitMerger {
113        post_layernorm_w: bf16(store, &format!("{vp}.post_layernorm.weight"))?,
114        downsample_w: bf16(store, &format!("{vp}.downsample.weight"))?,
115        downsample_b: bf16(store, &format!("{vp}.downsample.bias"))?,
116        proj_w: bf16(store, &format!("{vp}.merger.proj.weight"))?,
117        norm_w: bf16(store, &format!("{vp}.merger.post_projection_norm.weight"))?,
118        // 2026-09-25: The merger's post-projection norm is a LayerNorm with
119        // bias (`glm_vit/merge.rs` passes `norm_b`), so a missing bias fails
120        // the load here.
121        norm_b: bf16(store, &format!("{vp}.merger.post_projection_norm.bias"))?,
122        gate_w: bf16(store, &format!("{vp}.merger.gate_proj.weight"))?,
123        up_w: bf16(store, &format!("{vp}.merger.up_proj.weight"))?,
124        down_w: bf16(store, &format!("{vp}.merger.down_proj.weight"))?,
125    })
126}
127
128/// 2026-09-25: The body of `Glm5NextWeightLoader::load_vision_encoder`.
129///
130/// `Ok(None)` when the config has no `vision_config`, and, with a warning,
131/// when it has one but no tower prefix holds `patch_embed.proj.weight`.
132pub fn load_glm5_next_vision(
133    store: &WeightStore,
134    config: &ModelConfig,
135    gpu: &dyn GpuBackend,
136) -> Result<Option<VisionTower>> {
137    let Some(v) = config.vision.clone() else {
138        return Ok(None);
139    };
140    // 2026-09-25: The tower may sit under any of three prefixes; the first
141    // that holds `patch_embed.proj.weight` is used.
142    let vp = ["model.visual", "model.language_model.visual", "visual"]
143        .into_iter()
144        .find(|p| store.contains(&format!("{p}.patch_embed.proj.weight")));
145    let Some(vp) = vp else {
146        tracing::warn!(
147            "glm5_next declares a vision_config but no `*.visual.patch_embed.proj.weight` is \
148             present in the checkpoint; serving TEXT-ONLY (image input will be refused)."
149        );
150        return Ok(None);
151    };
152
153    let patch_dim = 3 * v.temporal_patch_size * v.patch_size * v.patch_size;
154    check_shape(
155        store,
156        &format!("{vp}.patch_embed.proj.weight"),
157        v.hidden_size,
158        patch_dim,
159    )?;
160    let patch_embed_w = bf16(store, &format!("{vp}.patch_embed.proj.weight"))?;
161    let patch_embed_b = bf16(store, &format!("{vp}.patch_embed.proj.bias"))?;
162
163    let mut blocks = Vec::with_capacity(v.depth);
164    for i in 0..v.depth {
165        blocks.push(load_block(store, vp, i, &v)?);
166    }
167    let merger = load_merger(store, vp, &v)?;
168
169    let geo = GlmVitGeometry {
170        hidden_size: v.hidden_size,
171        num_heads: v.num_heads,
172        intermediate_size: v.intermediate_size,
173        spatial_merge_size: v.spatial_merge_size,
174        out_hidden_size: v.out_hidden_size,
175        projection_intermediate_size: v.projection_intermediate_size,
176        patch_size: v.patch_size,
177        temporal_patch_size: v.temporal_patch_size,
178        rms_norm_eps: v.rms_norm_eps,
179        swiglu_limit: v.swiglu_limit,
180        rope_theta: v.rope_theta,
181        max_pixels: v.max_pixels,
182    };
183    let encoder = GlmVit::new(patch_embed_w, patch_embed_b, blocks, merger, &geo, gpu)?;
184    tracing::info!(
185        "GLM-5.3 vision tower bound: depth={} hidden={} heads={} patch={} merge={} → {} \
186         (BF16, {} tensors)",
187        v.depth,
188        v.hidden_size,
189        v.num_heads,
190        v.patch_size,
191        v.spatial_merge_size,
192        v.out_hidden_size,
193        v.depth * 14 + 11,
194    );
195    Ok(Some(VisionTower::glm(encoder)))
196}
197
198#[cfg(test)]
199mod tests {
200    /// 2026-09-25: The bound tensor count: 14 per block (the `GlmVitBlock`
201    /// fields in `load_block`) and 11 outside the blocks (patch-embed weight
202    /// and bias plus the nine `GlmVitMerger` fields), the same count the load
203    /// log prints as `v.depth * 14 + 11`.
204    #[test]
205    fn the_bound_tensor_count_is_the_checkpoints_347() {
206        let per_block = 14;
207        let non_block = 11;
208        assert_eq!(24 * per_block + non_block, 347);
209    }
210}
metrale_model_arch/weight_loader/glm5_next_vision.rs

metrale_model_arch/weight_loader/
glm5_next_vision.rs
