spark_model/weight_loader/nemotron.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3use anyhow::Result;
4use atlas_core::config::ModelConfig;
5use spark_runtime::gpu::GpuBackend;
6use spark_runtime::kv_cache::KvCacheDtype;
7use spark_runtime::weights::WeightStore;
8
9mod ssm_layer;
10
11use super::ModelWeightLoader;
12use crate::layer::TransformerLayer;
13use crate::layers::{FfnComponent, NemotronMoeLayer, Qwen3AttentionLayer};
14use crate::tp_shard::{TpAttentionDims, TpShardKind, shard_dense_bf16, shard_quantized_nvfp4};
15use crate::weight_map::{
16 DenseWeight, MtpWeights, dense, load_nemotron_attention, load_nemotron_moe, quantize_to_nvfp4,
17};
18
19pub struct NemotronHWeightLoader;
20
21impl ModelWeightLoader for NemotronHWeightLoader {
22 fn supports_tp(&self) -> bool {
23 // FullAttention layers TP-sharded across both quant paths
24 // (NVFP4-from-disk and BF16/FP8 → NVFP4). LinearAttention
25 // (Mamba-2 SSM) and MoE layers run full-replica per rank —
26 // SSM stays correct because hidden in/out is the same on
27 // every rank; MoE under EP+TP composition only uses EP.
28 true
29 }
30
31 fn load_layers(
32 &self,
33 store: &WeightStore,
34 config: &ModelConfig,
35 gpu: &dyn GpuBackend,
36 layer_kv_dtypes: &[KvCacheDtype],
37 ) -> Result<Vec<Box<dyn TransformerLayer>>> {
38 let layer_types = &config.layer_types;
39 let mut layers: Vec<Box<dyn TransformerLayer>> =
40 Vec::with_capacity(config.num_hidden_layers);
41 let mut attn_idx = 0usize;
42 let h = config.hidden_size;
43
44 // Runtime quantization kernels for BF16→NVFP4 conversion of unquantized layers.
45 let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
46 let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
47 let stream = gpu.default_stream();
48
49 // Pre-allocate a reusable scratch buffer for FP8→BF16 dequant intermediates.
50 // On GB10 UVM, gpu.free() posts in-band TLB invalidations that corrupt
51 // nearby allocations (BUG #29). Using a scratch buffer avoids all frees
52 // during loading. Size = max(in_proj, out_proj, shared_up, shared_down) in BF16 bytes.
53 let moe_input = config.moe_input_size();
54 // Puzzle: intermediate size varies per MoE layer — size scratch to max.
55 let max_moe_inter = config.max_moe_intermediate_size();
56 let scratch_elems = (config.mamba2_in_proj_size() * h)
57 .max(h * config.mamba2_d_inner())
58 .max(config.shared_expert_intermediate_size * h)
59 .max(h * config.shared_expert_intermediate_size)
60 .max(max_moe_inter * moe_input)
61 .max(moe_input * max_moe_inter);
62 let scratch_bytes = scratch_elems * 2; // BF16 = 2 bytes
63 let scratch = gpu.alloc(scratch_bytes)?;
64
65 for (i, lt) in layer_types.iter().enumerate() {
66 let lp = config.layer_prefix(i);
67 let norm = dense(store, &format!("{lp}.norm.weight"))?;
68
69 match lt {
70 atlas_core::config::LayerType::LinearAttention => {
71 let layer = Self::build_ssm_layer(
72 gpu, store, config, i, h, &lp, norm, quantize_k, absmax_k, scratch, stream,
73 )?;
74 layers.push(Box::new(layer));
75 }
76 atlas_core::config::LayerType::SlidingAttention => {
77 unreachable!("unexpected SlidingAttention in this loader")
78 }
79 // GLM-5.3's `deepseek_sparse_attention`. Nemotron has no indexer and no
80 // sparse-selection path, so this is a hard error rather than a silent
81 // fallthrough into the dense-attention arm.
82 atlas_core::config::LayerType::SparseAttention => anyhow::bail!(
83 "layer {i}: SparseAttention (deepseek_sparse_attention) has no Nemotron loader"
84 ),
85 atlas_core::config::LayerType::Moe => {
86 // Standalone MoE FFN layer (uniform Super/Nano or Puzzle per-block)
87 let moe_inter = config.moe_intermediate_size_for(i);
88 let top_k = config.num_experts_per_tok_for(i);
89 let moe = load_nemotron_moe(
90 store,
91 i,
92 config.num_experts,
93 gpu,
94 config,
95 Some(absmax_k),
96 Some(quantize_k),
97 stream,
98 Some(scratch),
99 &lp,
100 )?;
101 if i < 4 || moe_inter != config.moe_intermediate_size {
102 tracing::info!(
103 "L{i} MoE: inter={moe_inter} top_k={top_k} latent={} has_fc1={} has_fc2={} shared_up_s2={:.6e} experts[0].up_s2={:.6e}",
104 config.moe_latent_size,
105 moe.fc1_latent_proj.is_some(),
106 moe.fc2_latent_proj.is_some(),
107 moe.shared_up.weight_scale_2,
108 moe.experts
109 .first()
110 .map(|e| e.up_proj.weight_scale_2)
111 .unwrap_or(0.0),
112 );
113 }
114 let mut moe_layer =
115 NemotronMoeLayer::new(moe, norm, config, gpu, moe_inter, top_k)?;
116 // Builds the transposed shared-expert weights (2 small matrices
117 // per layer) so prefill can use `w4a16_gemm_t` instead of the base
118 // `w4a16_gemm`. Routed-expert transposition stays disabled inside
119 // for LatentMoE — 512 experts x 40 layers would not fit.
120 moe_layer.prepare_prefill_weights(gpu, config);
121 layers.push(Box::new(moe_layer));
122 }
123 atlas_core::config::LayerType::FullAttention => {
124 // Attention layer — quantize BF16 Q/K/V/O directly from
125 // WeightStore pointers (no intermediate alloc/free needed).
126 let (mut attn, mut q_nvfp4, mut k_nvfp4, mut v_nvfp4, mut o_dense, is_nvfp4) =
127 load_nemotron_attention(store, i, gpu, &lp)?;
128 let tp_rank = config.tp_rank;
129 let tp_size = config.tp_world_size.max(1);
130 let dims = TpAttentionDims::from_config(config);
131 if is_nvfp4 && tp_size > 1 {
132 // NVFP4-from-disk: shard packed weight + FP8 scales.
133 let group_size = 16usize;
134 if let Some(q) = q_nvfp4.as_ref() {
135 let s = shard_quantized_nvfp4(
136 q,
137 dims.full_q_n,
138 dims.h,
139 TpShardKind::ColumnParallel,
140 tp_rank,
141 tp_size,
142 group_size,
143 gpu,
144 )?;
145 gpu.free(q.weight)?;
146 gpu.free(q.weight_scale)?;
147 q_nvfp4 = Some(s);
148 }
149 if let Some(k) = k_nvfp4.as_ref() {
150 let s = shard_quantized_nvfp4(
151 k,
152 dims.full_kv_n,
153 dims.h,
154 TpShardKind::ColumnParallel,
155 tp_rank,
156 tp_size,
157 group_size,
158 gpu,
159 )?;
160 gpu.free(k.weight)?;
161 gpu.free(k.weight_scale)?;
162 k_nvfp4 = Some(s);
163 }
164 if let Some(v) = v_nvfp4.as_ref() {
165 let s = shard_quantized_nvfp4(
166 v,
167 dims.full_kv_n,
168 dims.h,
169 TpShardKind::ColumnParallel,
170 tp_rank,
171 tp_size,
172 group_size,
173 gpu,
174 )?;
175 gpu.free(v.weight)?;
176 gpu.free(v.weight_scale)?;
177 v_nvfp4 = Some(s);
178 }
179 // O proj is stored on attn.o_proj as QuantizedWeight in NVFP4-disk path.
180 let o_old = attn.o_proj;
181 let o_sharded = shard_quantized_nvfp4(
182 &o_old,
183 dims.h,
184 dims.full_o_in,
185 TpShardKind::RowParallel,
186 tp_rank,
187 tp_size,
188 group_size,
189 gpu,
190 )?;
191 gpu.free(o_old.weight)?;
192 gpu.free(o_old.weight_scale)?;
193 attn.o_proj = o_sharded;
194 }
195 let mut bf16_o_dense: Option<DenseWeight> = None;
196 let (q_nv, k_nv, v_nv) = if is_nvfp4 {
197 (q_nvfp4, k_nvfp4, v_nvfp4)
198 } else {
199 let num_heads = config.num_attention_heads;
200 let kv_heads = config.num_key_value_heads;
201 let hd = config.head_dim;
202 // BF16 / FP8-dequant fallback: shard the dense BF16
203 // before quantization. Dims here are TP-LOCAL after
204 // sharding (config head counts already TP-divided).
205 if tp_size > 1 {
206 let (qp, _, _) = shard_dense_bf16(
207 attn.q_proj.weight,
208 dims.full_q_n,
209 dims.h,
210 TpShardKind::ColumnParallel,
211 tp_rank,
212 tp_size,
213 gpu,
214 )?;
215 if qp != attn.q_proj.weight {
216 gpu.free(attn.q_proj.weight)?;
217 }
218 attn.q_proj.weight = qp;
219 let (kp, _, _) = shard_dense_bf16(
220 attn.k_proj.weight,
221 dims.full_kv_n,
222 dims.h,
223 TpShardKind::ColumnParallel,
224 tp_rank,
225 tp_size,
226 gpu,
227 )?;
228 if kp != attn.k_proj.weight {
229 gpu.free(attn.k_proj.weight)?;
230 }
231 attn.k_proj.weight = kp;
232 let (vp, _, _) = shard_dense_bf16(
233 attn.v_proj.weight,
234 dims.full_kv_n,
235 dims.h,
236 TpShardKind::ColumnParallel,
237 tp_rank,
238 tp_size,
239 gpu,
240 )?;
241 if vp != attn.v_proj.weight {
242 gpu.free(attn.v_proj.weight)?;
243 }
244 attn.v_proj.weight = vp;
245 let (op, _, _) = shard_dense_bf16(
246 o_dense.weight,
247 dims.h,
248 dims.full_o_in,
249 TpShardKind::RowParallel,
250 tp_rank,
251 tp_size,
252 gpu,
253 )?;
254 if op != o_dense.weight {
255 gpu.free(o_dense.weight)?;
256 }
257 o_dense.weight = op;
258 }
259 // Keep attention in BF16 when the checkpoint ships it that
260 // way. ModelOpt left Q/K/V/O unquantized ON PURPOSE here:
261 // Puzzle is Mamba-dominant with only 9 full-attention
262 // layers, so those layers carry the long-range retrieval
263 // and are the last place to spend precision — and at
264 // ~1.2 GB BF16 they are cheap to keep. Crushing them
265 // 16-bit -> 4-bit saved ~0.9 GB and degraded exactly what
266 // they exist for. `ATLAS_NEMOTRON_BF16_ATTN=0` restores
267 // the old quantize-everything behaviour for an A/B.
268 let keep_bf16_attn =
269 std::env::var("ATLAS_NEMOTRON_BF16_ATTN").as_deref() != Ok("0");
270 if keep_bf16_attn {
271 tracing::info!(
272 "L{i} attention: keeping checkpoint BF16 Q/K/V/O (no NVFP4 requant)"
273 );
274 bf16_o_dense = Some(o_dense);
275 (None, None, None)
276 } else {
277 let q = quantize_to_nvfp4(
278 &attn.q_proj,
279 num_heads * hd,
280 h,
281 gpu,
282 absmax_k,
283 quantize_k,
284 stream,
285 )?;
286 let k = quantize_to_nvfp4(
287 &attn.k_proj,
288 kv_heads * hd,
289 h,
290 gpu,
291 absmax_k,
292 quantize_k,
293 stream,
294 )?;
295 let v = quantize_to_nvfp4(
296 &attn.v_proj,
297 kv_heads * hd,
298 h,
299 gpu,
300 absmax_k,
301 quantize_k,
302 stream,
303 )?;
304 let o = quantize_to_nvfp4(
305 &o_dense,
306 h,
307 num_heads * hd,
308 gpu,
309 absmax_k,
310 quantize_k,
311 stream,
312 )?;
313 attn.o_proj = o;
314 (Some(q), Some(k), Some(v))
315 }
316 };
317 // Transposed Q/K/V/O so prefill uses `w4a16_gemm_t` (FP8 MMA,
318 // N128/K32, cp.async) instead of the base `w4a16_gemm`. Same
319 // gap as the SSM/MoE layers: the setter existed but was never
320 // called for Nemotron. Q/K/V/O are small (h=4096, kv=2 heads),
321 // so the extra copies cost ~0.3 GB.
322 let q_dim = config.num_attention_heads * config.head_dim;
323 let kv_dim = config.num_key_value_heads * config.head_dim;
324 let qt = q_nv
325 .as_ref()
326 .and_then(|w| w.transpose_for_gemm(gpu, q_dim, h).ok());
327 let kt = k_nv
328 .as_ref()
329 .and_then(|w| w.transpose_for_gemm(gpu, kv_dim, h).ok());
330 let vt = v_nv
331 .as_ref()
332 .and_then(|w| w.transpose_for_gemm(gpu, kv_dim, h).ok());
333 // Keep-BF16 attention leaves `attn.o_proj` as the NULL
334 // placeholder (`load_nemotron_attention` returns the real
335 // weight via `o_dense`); transposing it would launch
336 // `transpose_u8` on a NULL pointer, and the swallowed `.ok()`
337 // error left the CUDA context poisoned (700 at the next
338 // module load). Only transpose a real quantized o_proj.
339 let ot = if attn.o_proj.is_null() {
340 None
341 } else {
342 attn.o_proj.transpose_for_gemm(gpu, h, q_dim).ok()
343 };
344
345 let mut attn_layer = Qwen3AttentionLayer::new_ungated(
346 norm,
347 attn,
348 DenseWeight {
349 weight: spark_runtime::gpu::DevicePtr::NULL,
350 },
351 FfnComponent::None,
352 attn_idx,
353 q_nv,
354 k_nv,
355 v_nv,
356 gpu,
357 layer_kv_dtypes[attn_idx],
358 config.fp8_kv_calibration_tokens,
359 config,
360 )?;
361 if let Some(od) = bf16_o_dense {
362 // Dispatch checks `o_dense_bf16` first (see gemma4 loader).
363 attn_layer.set_o_dense_bf16(od);
364 }
365 attn_layer.set_prefill_weights(qt, kt, vt, ot);
366 layers.push(Box::new(attn_layer));
367 attn_idx += 1;
368 }
369 }
370
371 if (i + 1) % 10 == 0 {
372 tracing::info!("Loaded layers 0..{}", i + 1);
373 }
374 }
375
376 tracing::info!(
377 "Nemotron-H weight loader: {} layers ({} SSM, {} MoE, {} attention)",
378 layers.len(),
379 config.num_ssm_layers(),
380 config.num_moe_layers(),
381 attn_idx,
382 );
383
384 Ok(layers)
385 }
386
387 fn load_embedding(
388 &self,
389 store: &WeightStore,
390 config: &ModelConfig,
391 _gpu: &dyn GpuBackend,
392 ) -> Result<DenseWeight> {
393 dense(
394 store,
395 &format!("{}.embeddings.weight", config.weight_prefix),
396 )
397 }
398
399 fn load_final_norm(
400 &self,
401 store: &WeightStore,
402 config: &ModelConfig,
403 _gpu: &dyn GpuBackend,
404 ) -> Result<DenseWeight> {
405 dense(store, &format!("{}.norm_f.weight", config.weight_prefix))
406 }
407
408 fn load_lm_head(
409 &self,
410 store: &WeightStore,
411 config: &ModelConfig,
412 _gpu: &dyn GpuBackend,
413 ) -> Result<DenseWeight> {
414 if store.contains("lm_head.weight") {
415 dense(store, "lm_head.weight")
416 } else {
417 dense(
418 store,
419 &format!("{}.embeddings.weight", config.weight_prefix),
420 )
421 }
422 }
423
424 fn load_mtp_weights(
425 &self,
426 _store: &WeightStore,
427 _config: &ModelConfig,
428 _gpu: &dyn GpuBackend,
429 ) -> Result<Option<MtpWeights>> {
430 Ok(None) // Nemotron-H has no MTP
431 }
432}
433
434#[cfg(test)]
435mod tests {
436 use std::collections::HashMap;
437
438 use spark_runtime::gpu::{DevicePtr, mock::MockGpuBackend};
439 use spark_runtime::weights::{WeightDtype, WeightTensor};
440
441 use super::*;
442
443 fn tensor(ptr: u64) -> WeightTensor {
444 WeightTensor {
445 ptr: DevicePtr(ptr),
446 shape: vec![4],
447 dtype: WeightDtype::BF16,
448 }
449 }
450
451 fn config() -> ModelConfig {
452 let mut config = ModelConfig::qwen3_next_80b_nvfp4();
453 config.weight_prefix = "backbone".to_string();
454 config
455 }
456
457 #[test]
458 fn nemotron_layout_routes_embedding_norm_and_tied_head() {
459 let store = WeightStore::from_map(HashMap::from([
460 ("backbone.embeddings.weight".to_string(), tensor(11)),
461 ("backbone.norm_f.weight".to_string(), tensor(12)),
462 ]));
463 let gpu = MockGpuBackend::new();
464 let loader = NemotronHWeightLoader;
465 let config = config();
466
467 assert_eq!(
468 loader.load_embedding(&store, &config, &gpu).unwrap().weight,
469 DevicePtr(11)
470 );
471 assert_eq!(
472 loader
473 .load_final_norm(&store, &config, &gpu)
474 .unwrap()
475 .weight,
476 DevicePtr(12)
477 );
478 assert_eq!(
479 loader.load_lm_head(&store, &config, &gpu).unwrap().weight,
480 DevicePtr(11)
481 );
482 assert!(loader.supports_tp());
483 assert!(
484 loader
485 .load_mtp_weights(&store, &config, &gpu)
486 .unwrap()
487 .is_none()
488 );
489 }
490
491 #[test]
492 fn explicit_lm_head_takes_precedence_over_tied_embedding() {
493 let store = WeightStore::from_map(HashMap::from([
494 ("backbone.embeddings.weight".to_string(), tensor(11)),
495 ("lm_head.weight".to_string(), tensor(13)),
496 ]));
497
498 let actual = NemotronHWeightLoader
499 .load_lm_head(&store, &config(), &MockGpuBackend::new())
500 .unwrap();
501 assert_eq!(actual.weight, DevicePtr(13));
502 }
503}