metrale_model_arch/weight_loader/nemotron.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Nemotron-H weight loader: Mamba-2 SSM, MoE and full-attention layers, one per
4//! entry of `config.layer_types`.
5//!
6//! Owner: model-arch weight loader (Nemotron-H).
7//! Invariants:
8//! - Only the full-attention projections are TP-sharded; SSM and MoE layers are not.
9//! - Full-attention layers take KV dtypes from `layer_kv_dtypes` in order of appearance.
10
11use anyhow::Result;
12use metrale_cache::kv_cache::KvCacheDtype;
13use metrale_config::ModelConfig;
14use metrale_gpu_runtime::gpu::GpuBackend;
15use metrale_model_weights::weights::WeightStore;
16
17mod ssm_layer;
18
19use super::ModelWeightLoader;
20use crate::nemotron_moe::NemotronMoeLayer;
21use crate::tp_shard::{TpAttentionDims, TpShardKind, shard_dense_bf16, shard_quantized_nvfp4};
22use metrale_model_layers::layer::TransformerLayer;
23use metrale_model_layers::layers::{FfnComponent, Qwen3AttentionLayer};
24use metrale_model_layers::weight_map::{
25 DenseWeight, MtpWeights, dense, load_nemotron_attention, load_nemotron_moe, quantize_to_nvfp4,
26};
27
28pub struct NemotronHWeightLoader;
29
30impl ModelWeightLoader for NemotronHWeightLoader {
31 fn supports_tp(&self) -> bool {
32 // 2026-09-25: Full-attention layers are TP-sharded in both the
33 // NVFP4-from-disk and the dense arm; SSM and MoE layers are not.
34 true
35 }
36
37 fn load_layers(
38 &self,
39 store: &WeightStore,
40 config: &ModelConfig,
41 gpu: &dyn GpuBackend,
42 layer_kv_dtypes: &[KvCacheDtype],
43 ) -> Result<Vec<Box<dyn TransformerLayer>>> {
44 let layer_types = &config.layer_types;
45 let mut layers: Vec<Box<dyn TransformerLayer>> =
46 Vec::with_capacity(config.num_hidden_layers);
47 let mut attn_idx = 0usize;
48 let h = config.hidden_size;
49
50 // 2026-09-25: Kernels for quantizing BF16 weights to NVFP4 at load.
51 let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
52 let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
53 let stream = gpu.default_stream();
54
55 // 2026-09-25: One reusable BF16 scratch buffer for the SSM and MoE
56 // loaders' dequant intermediates, sized to the largest SSM projection,
57 // shared-expert projection or MoE expert projection.
58 let moe_input = config.moe_input_size();
59 // 2026-09-25: The MoE intermediate size can vary per layer; use the largest.
60 let max_moe_inter = config.max_moe_intermediate_size();
61 let scratch_elems = (config.mamba2_in_proj_size() * h)
62 .max(h * config.mamba2_d_inner())
63 .max(config.shared_expert_intermediate_size * h)
64 .max(h * config.shared_expert_intermediate_size)
65 .max(max_moe_inter * moe_input)
66 .max(moe_input * max_moe_inter);
67 let scratch_bytes = scratch_elems * 2;
68 let scratch = gpu.alloc(scratch_bytes)?;
69
70 for (i, lt) in layer_types.iter().enumerate() {
71 let lp = config.layer_prefix(i);
72 let norm = dense(store, &format!("{lp}.norm.weight"))?;
73
74 match lt {
75 metrale_config::LayerType::LinearAttention => {
76 let layer = Self::build_ssm_layer(
77 gpu, store, config, i, h, &lp, norm, quantize_k, absmax_k, scratch, stream,
78 )?;
79 layers.push(Box::new(layer));
80 }
81 metrale_config::LayerType::SlidingAttention => {
82 unreachable!("unexpected SlidingAttention in this loader")
83 }
84 // 2026-09-25: Sparse attention has no Nemotron path: an error, not the
85 // dense-attention arm.
86 metrale_config::LayerType::SparseAttention => anyhow::bail!(
87 "layer {i}: SparseAttention (deepseek_sparse_attention) has no Nemotron loader"
88 ),
89 metrale_config::LayerType::Moe => {
90 // 2026-09-25: Standalone MoE layer, with this layer's intermediate
91 // size and top-k.
92 let moe_inter = config.moe_intermediate_size_for(i);
93 let top_k = config.num_experts_per_tok_for(i);
94 let moe = load_nemotron_moe(
95 store,
96 i,
97 config.num_experts,
98 gpu,
99 config,
100 Some(absmax_k),
101 Some(quantize_k),
102 stream,
103 Some(scratch),
104 &lp,
105 )?;
106 if i < 4 || moe_inter != config.moe_intermediate_size {
107 tracing::info!(
108 "L{i} MoE: inter={moe_inter} top_k={top_k} latent={} has_fc1={} has_fc2={} shared_up_s2={:.6e} experts[0].up_s2={:.6e}",
109 config.moe_latent_size,
110 moe.fc1_latent_proj.is_some(),
111 moe.fc2_latent_proj.is_some(),
112 moe.shared_up.weight_scale_2,
113 moe.experts
114 .first()
115 .map(|e| e.up_proj.weight_scale_2)
116 .unwrap_or(0.0),
117 );
118 }
119 let mut moe_layer =
120 NemotronMoeLayer::new(moe, norm, config, gpu, moe_inter, top_k)?;
121 // 2026-09-25: Transposed prefill copies; routed experts are
122 // transposed only when `moe_latent_size == 0`
123 // (`nemotron_moe/prefill_weights.rs`).
124 moe_layer.prepare_prefill_weights(
125 gpu,
126 config,
127 crate::nemotron_moe::shared_expert_e4m3(),
128 );
129 layers.push(Box::new(moe_layer));
130 }
131 metrale_config::LayerType::FullAttention => {
132 // 2026-09-25: Attention: `load_nemotron_attention` returns NVFP4
133 // from disk or dense weights (`is_nvfp4`).
134 let (mut attn, mut q_nvfp4, mut k_nvfp4, mut v_nvfp4, mut o_dense, is_nvfp4) =
135 load_nemotron_attention(store, i, gpu, &lp)?;
136 let tp_rank = config.tp_rank;
137 let tp_size = config.tp_world_size.max(1);
138 let dims = TpAttentionDims::from_config(config);
139 if is_nvfp4 && tp_size > 1 {
140 // 2026-09-25: NVFP4 from disk: shard the packed weights and block
141 // scales, then free the unsharded ones.
142 let group_size = 16usize;
143 if let Some(q) = q_nvfp4.as_ref() {
144 let s = shard_quantized_nvfp4(
145 q,
146 dims.full_q_n,
147 dims.h,
148 TpShardKind::ColumnParallel,
149 tp_rank,
150 tp_size,
151 group_size,
152 gpu,
153 )?;
154 gpu.free(q.weight)?;
155 gpu.free(q.weight_scale)?;
156 q_nvfp4 = Some(s);
157 }
158 if let Some(k) = k_nvfp4.as_ref() {
159 let s = shard_quantized_nvfp4(
160 k,
161 dims.full_kv_n,
162 dims.h,
163 TpShardKind::ColumnParallel,
164 tp_rank,
165 tp_size,
166 group_size,
167 gpu,
168 )?;
169 gpu.free(k.weight)?;
170 gpu.free(k.weight_scale)?;
171 k_nvfp4 = Some(s);
172 }
173 if let Some(v) = v_nvfp4.as_ref() {
174 let s = shard_quantized_nvfp4(
175 v,
176 dims.full_kv_n,
177 dims.h,
178 TpShardKind::ColumnParallel,
179 tp_rank,
180 tp_size,
181 group_size,
182 gpu,
183 )?;
184 gpu.free(v.weight)?;
185 gpu.free(v.weight_scale)?;
186 v_nvfp4 = Some(s);
187 }
188 // 2026-09-25: In the NVFP4 arm O is `attn.o_proj`.
189 let o_old = attn.o_proj;
190 let o_sharded = shard_quantized_nvfp4(
191 &o_old,
192 dims.h,
193 dims.full_o_in,
194 TpShardKind::RowParallel,
195 tp_rank,
196 tp_size,
197 group_size,
198 gpu,
199 )?;
200 gpu.free(o_old.weight)?;
201 gpu.free(o_old.weight_scale)?;
202 attn.o_proj = o_sharded;
203 }
204 let mut bf16_o_dense: Option<DenseWeight> = None;
205 let (q_nv, k_nv, v_nv) = if is_nvfp4 {
206 (q_nvfp4, k_nvfp4, v_nvfp4)
207 } else {
208 let num_heads = config.num_attention_heads;
209 let kv_heads = config.num_key_value_heads;
210 let hd = config.head_dim;
211 // 2026-09-25: Dense arm: shard the BF16 weights before any
212 // quantization.
213 if tp_size > 1 {
214 let (qp, _, _) = shard_dense_bf16(
215 attn.q_proj.weight,
216 dims.full_q_n,
217 dims.h,
218 TpShardKind::ColumnParallel,
219 tp_rank,
220 tp_size,
221 gpu,
222 )?;
223 if qp != attn.q_proj.weight {
224 gpu.free(attn.q_proj.weight)?;
225 }
226 attn.q_proj.weight = qp;
227 let (kp, _, _) = shard_dense_bf16(
228 attn.k_proj.weight,
229 dims.full_kv_n,
230 dims.h,
231 TpShardKind::ColumnParallel,
232 tp_rank,
233 tp_size,
234 gpu,
235 )?;
236 if kp != attn.k_proj.weight {
237 gpu.free(attn.k_proj.weight)?;
238 }
239 attn.k_proj.weight = kp;
240 let (vp, _, _) = shard_dense_bf16(
241 attn.v_proj.weight,
242 dims.full_kv_n,
243 dims.h,
244 TpShardKind::ColumnParallel,
245 tp_rank,
246 tp_size,
247 gpu,
248 )?;
249 if vp != attn.v_proj.weight {
250 gpu.free(attn.v_proj.weight)?;
251 }
252 attn.v_proj.weight = vp;
253 let (op, _, _) = shard_dense_bf16(
254 o_dense.weight,
255 dims.h,
256 dims.full_o_in,
257 TpShardKind::RowParallel,
258 tp_rank,
259 tp_size,
260 gpu,
261 )?;
262 if op != o_dense.weight {
263 gpu.free(o_dense.weight)?;
264 }
265 o_dense.weight = op;
266 }
267 // 2026-09-25: Dense attention stays BF16 unless
268 // `METRALE_NEMOTRON_BF16_ATTN=0`, which quantizes Q/K/V/O to
269 // NVFP4.
270 let keep_bf16_attn =
271 std::env::var("METRALE_NEMOTRON_BF16_ATTN").as_deref() != Ok("0");
272 if keep_bf16_attn {
273 tracing::info!(
274 "L{i} attention: keeping checkpoint BF16 Q/K/V/O (no NVFP4 requant)"
275 );
276 bf16_o_dense = Some(o_dense);
277 (None, None, None)
278 } else {
279 let q = quantize_to_nvfp4(
280 &attn.q_proj,
281 num_heads * hd,
282 h,
283 gpu,
284 absmax_k,
285 quantize_k,
286 stream,
287 )?;
288 let k = quantize_to_nvfp4(
289 &attn.k_proj,
290 kv_heads * hd,
291 h,
292 gpu,
293 absmax_k,
294 quantize_k,
295 stream,
296 )?;
297 let v = quantize_to_nvfp4(
298 &attn.v_proj,
299 kv_heads * hd,
300 h,
301 gpu,
302 absmax_k,
303 quantize_k,
304 stream,
305 )?;
306 let o = quantize_to_nvfp4(
307 &o_dense,
308 h,
309 num_heads * hd,
310 gpu,
311 absmax_k,
312 quantize_k,
313 stream,
314 )?;
315 attn.o_proj = o;
316 (Some(q), Some(k), Some(v))
317 }
318 };
319 // 2026-09-25: Transposed NVFP4 Q/K/V/O copies for prefill; a
320 // failed transpose leaves that copy `None`.
321 let q_dim = config.num_attention_heads * config.head_dim;
322 let kv_dim = config.num_key_value_heads * config.head_dim;
323 let qt = q_nv
324 .as_ref()
325 .and_then(|w| w.transpose_for_gemm(gpu, q_dim, h).ok());
326 let kt = k_nv
327 .as_ref()
328 .and_then(|w| w.transpose_for_gemm(gpu, kv_dim, h).ok());
329 let vt = v_nv
330 .as_ref()
331 .and_then(|w| w.transpose_for_gemm(gpu, kv_dim, h).ok());
332 // 2026-09-25: In the BF16 arm `attn.o_proj` is null (the weight
333 // is `o_dense`); only a real quantized o_proj is transposed.
334 let ot = if attn.o_proj.is_null() {
335 None
336 } else {
337 attn.o_proj.transpose_for_gemm(gpu, h, q_dim).ok()
338 };
339
340 let mut attn_layer = Qwen3AttentionLayer::new_ungated(
341 norm,
342 attn,
343 DenseWeight {
344 weight: metrale_gpu_runtime::gpu::DevicePtr::NULL,
345 },
346 FfnComponent::None,
347 attn_idx,
348 q_nv,
349 k_nv,
350 v_nv,
351 gpu,
352 layer_kv_dtypes[attn_idx],
353 config.fp8_kv_calibration_tokens,
354 config,
355 )?;
356 if let Some(od) = bf16_o_dense {
357 attn_layer.set_o_dense_bf16(od);
358 }
359 attn_layer.set_prefill_weights(qt, kt, vt, ot);
360 layers.push(Box::new(attn_layer));
361 attn_idx += 1;
362 }
363 }
364
365 if (i + 1) % 10 == 0 {
366 tracing::info!("Loaded layers 0..{}", i + 1);
367 }
368 }
369
370 tracing::info!(
371 "Nemotron-H weight loader: {} layers ({} SSM, {} MoE, {} attention)",
372 layers.len(),
373 config.num_ssm_layers(),
374 config.num_moe_layers(),
375 attn_idx,
376 );
377
378 Ok(layers)
379 }
380
381 fn load_embedding(
382 &self,
383 store: &WeightStore,
384 config: &ModelConfig,
385 _gpu: &dyn GpuBackend,
386 ) -> Result<DenseWeight> {
387 dense(
388 store,
389 &format!("{}.embeddings.weight", config.weight_prefix),
390 )
391 }
392
393 fn load_final_norm(
394 &self,
395 store: &WeightStore,
396 config: &ModelConfig,
397 _gpu: &dyn GpuBackend,
398 ) -> Result<DenseWeight> {
399 dense(store, &format!("{}.norm_f.weight", config.weight_prefix))
400 }
401
402 fn load_lm_head(
403 &self,
404 store: &WeightStore,
405 config: &ModelConfig,
406 _gpu: &dyn GpuBackend,
407 ) -> Result<DenseWeight> {
408 if store.contains("lm_head.weight") {
409 dense(store, "lm_head.weight")
410 } else {
411 dense(
412 store,
413 &format!("{}.embeddings.weight", config.weight_prefix),
414 )
415 }
416 }
417
418 fn load_mtp_weights(
419 &self,
420 _store: &WeightStore,
421 _config: &ModelConfig,
422 _gpu: &dyn GpuBackend,
423 ) -> Result<Option<MtpWeights>> {
424 Ok(None)
425 }
426}
427
428#[cfg(test)]
429#[path = "nemotron_tests.rs"]
430mod tests;