feat: support Nemotron-H / Nemotron-Cascade-2 (#1711)

[michael-rabe]

Adds initial support for hybrid Mamba2 + Attention + MoE models via the unfused_moe registry. A bare AutoRound("nvidia/...") call now produces a coherent INT4 checkpoint without launcher-side workarounds.

Core additions:

• unfused_moe/nemotron_h.py: linear-discoverable MoE block
• unfused_moe/nemotron_h_setup.py: post-load fixups (Zamba2 group_size, SSM/router FP32 restore)
• utils/source_tensor_overrides.py: generic source-checkpoint tensor reload utility
• MODEL_CONFIG["nemotron_h"]: dispatch + upstream-rename preservation
• compressors/base.py: apply_post_load_fixups hook (non-NH no-op)
• Export: norm_dtype and scale_dtype per-layer overrides

Documentation:

• New .claude/skills/adapt-unfused-moe skill covering the pipeline
• adapt-new-llm slimmed to point at the dedicated skill
• README / README_CN / step_by_step notes updated

Tests: 89 CPU tests across registration, post-load, export dtype controls, source-tensor overrides, and missing-tensors symmetry.

Closes #1711

Description

Please briefly describe your main changes, the motivation.

Type of Change

• Bug fix
• New feature
• Documentation update
• Performance improvement
• Code refactoring
• Other (please specify):

Related Issues

Fixes or relates to #

Checklist Before Submitting

• My code has been tested locally.
• Documentation has been updated as needed.
• New or updated tests are included where applicable.

[wenhuach21]

Thanks for the PR. Are all these file changes necessary? It seems to modify many unrelated files.

[michael-rabe]

Thanks for the PR. Are all these file changes necessary? It seems to modify many unrelated files.

I'm afraid it is, but haven't optimized (reduced) much after I got coherent and stable output.
Looking forward to your feedback ;).

Very welcome if we can reduce the number of changes, nemotron models seem to behave way different to others and are more sensitive to low dtypes.
I had coherence problems without these changes, the post apply "patches" became necessary to bypass the "standard" autoround behaviours without changing too much of the autoround core.

[wenhuach21]

Thanks for the PR. Are all these file changes necessary? It seems to modify many unrelated files.

I'm afraid it is, but haven't optimized (reduced) much after I got coherent and stable output. Looking forward to your feedback ;).

Very welcome if we can reduce the number of changes, nemotron models seem to behave way different to others and are more sensitive to low dtypes. I had coherence problems without these changes, the post apply "patches" became necessary to bypass the "standard" autoround behaviours without changing too much of the autoround core.

Thanks for the prompt reply. Could you let me know what indicator led you to change the norm dtype to FP32? I checked the model and noticed that the original norm weights are in BF16. I’m wondering whether modifying the norm dtype in the weights might cause issues with vLLM or SGLang, and whether they would still behave as expected.

For MoE models, I would generally suggest using higher precision for certain parts, for example, 8-bit for critical components (non-moe modules) and 4-bit for others (experts). This is a more common approach when pure 4-bit quantization leads to a significant accuracy drop and avoids the API change.

Please feel free to correct me if I’m mistaken, as I’m not very familiar with this model.

Thanks again for your pr.

[wenhuach21]

for mixed-bits, you could try scheme="int4_mixed"

[michael-rabe]

Thanks for the PR. Are all these file changes necessary? It seems to modify many unrelated files.

I'm afraid it is, but haven't optimized (reduced) much after I got coherent and stable output. Looking forward to your feedback ;).
Very welcome if we can reduce the number of changes, nemotron models seem to behave way different to others and are more sensitive to low dtypes. I had coherence problems without these changes, the post apply "patches" became necessary to bypass the "standard" autoround behaviours without changing too much of the autoround core.

Thanks for the prompt reply. Could you let me know what indicator led you to change the norm dtype to FP32? I checked the model and noticed that the original norm weights are in BF16. I’m wondering whether modifying the norm dtype in the weights might cause issues with vLLM or SGLang, and whether they would still behave as expected.

For MoE models, I would generally suggest using higher precision for certain parts, for example, 8-bit for critical components (non-moe modules) and 4-bit for others (experts). This is a more common approach when pure 4-bit quantization leads to a significant accuracy drop and avoids the API change.

Please feel free to correct me if I’m mistaken, as I’m not very familiar with this model.

Thanks again for your pr.

Ah, I missed the proper documentation (I'll fix that tonight).
Thanks for the valuable suggestions, I'm following up on that and set the PR into "draft" status.

You're right that the original norm weights are BF16, and BF16 is numerically fine for the norm weights themselves — upcasting them to FP32 is lossless but gains nothing at the weight level.

The purpose of norm_dtype="fp32" is different: it's a lever for the residual stream, not for the norm weights. In HuggingFace-style RMSNorm the norm output's dtype follows the weight's dtype, so storing norm weights as FP32 pulls the post-norm tensor — and via the residual add, the residual stream itself — into FP32 at runtime. Over 50+ layers the residual is a sum of many block outputs, and keeping that accumulation in FP32 reduces BF16 rounding drift.

Two caveats: this only helps on engines that honor the stored norm-weight dtype (HF transformers does, fused-kernel runtimes like vLLM/SGLang often don't). And the Nemotron-H coherence fix itself was not norm_dtype — it was the always-on post-load restore of the SSM core tensors and the router correction bias to FP32, which are the tensors with genuine FP32 requirements. norm_dtype is orthogonal and optional.

So: FP32 norm weights aren't about upgrading the norm weights themselves — they're an export-side lever for residual-stream precision on compliant inference stacks.

Please correct me, if I'm wrong. I'm new to the club and still in the learning phase ;).
I appreciate your questions.
Sorry for the wall of text.

[michael-rabe]

Progressing, the comments were helpful and can reduce the number of changes.
Still investigating as I'm not yet happy with the results (accuracy) and want to make sure it's not in the quantisation itself.

[michael-rabe]

I tried to remove everything not absolutely necessary (some fallbacks added), I ran a couple of autoround runs to verify the results (BF16 wasn't the actual issue, ignore that comment, the differences are neglectable).
Test-Results: LAMBADA-5153: acc=0.624, tok_acc=0.704.

[wenhuach21]

@copilot resolve the merge conflicts in this pull request

[wenhuach21]

with this pr #2303 and some other optimizations in main branch, we no need to monkey patch moes, please feel free to reopen this pr if there are still some issues

(autoround) wenhuach@mlp-dgx-01:~/auto-round$ CUDA_VISIBLE_DEVICES=6 python3 -m auto_round /models/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/  --output_dir /data2/wenhuach
W0909 09:41:43.345000 4162872 site-packages/torch/utils/_pytree.py:630] <enum 'KernelPreference'> is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.
W0909 09:41:43.394000 4162872 site-packages/torch/utils/_pytree.py:630] <enum 'ScaleCalculationMode'> is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.
2026-09-09 09:41:44 INFO main.py L309: start to quantize /models/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16
2026-09-09 09:41:45 INFO model.py L332: Fused-MoE checkpoint detected: building a meta skeleton and materializing weights per block (set `AR_DISABLE_AUTO_META_LOAD=1` to load the whole model on CPU instead).
[transformers] The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)` is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1d
2026-09-09 09:41:45 INFO device.py L1558: Before applying custom replacements 'peak_ram': 0.94GB
2026-09-09 09:41:46 INFO moe_experts_interface.py L793: [MoE Prep] Unfused 23 MOE experts modules
2026-09-09 09:41:47 INFO device.py L1558: After applying custom replacements 'peak_ram': 1.06GB
2026-09-09 09:41:47 INFO replace_modules.py L101: Prepared 23 MOE modules for quantization
2026-09-09 09:41:47 INFO replace_modules.py L160: Experts (before unfuse) [model.layers.1.mixer.experts] (NemotronHExperts):
NemotronHExperts(
  (act_fn): ReLUSquaredActivation()
)
2026-09-09 09:41:47 INFO replace_modules.py L161: Experts (after unfuse) [model.layers.1.mixer.experts] (NemotronHExperts):
NemotronHExperts(
  (act_fn): ReLUSquaredActivation()
  (0-127): 128 x _ExpertContainer(
    (up_proj): Linear(in_features=2688, out_features=1856, bias=False)
    (down_proj): Linear(in_features=1856, out_features=2688, bias=False)
  )
)
2026-09-09 09:41:49 INFO base.py L515: Using default calibration dataset NeelNanda/pile-10k.
2026-09-09 09:41:49 INFO base.py L2138: format is not set, using default auto_round format.
2026-09-09 09:41:49 INFO base.py L999: Using predefined ignore_layers: model.layers.[1,3,6,8,10,13,15,17,20,22,24,27,29,31,34,36,38,40,43,45,47,49,51].mixer.gate
2026-09-09 09:41:49 INFO layer_config_resolver.py L401: Ignored layers: model.layers.[1,3,6,8,10,13,15,17,20,22,24,27,29,31,34,36,38,40,43,45,47,49,51].mixer.gate
2026-09-09 09:41:50 INFO layer_config_resolver.py L401: Ignored layers: model.layers.[1,3,6,8,10,13,15,17,20,22,24,27,29,31,34,36,38,40,43,45,47,49,51].mixer.gate
2026-09-09 09:41:50 INFO base.py L1340: `torch.compile` is enabled
2026-09-09 09:41:52 INFO orchestrator.py L587: start to cache block inputs
2026-09-09 09:41:52 INFO calib_dataset.py L1113: Preprocessing calibration dataset in a subprocess to avoid memory leaks...
Map: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:13<00:00, 752.57 examples/s]
Filter: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10000/10000 [00:08<00:00, 1214.37 examples/s]
Casting the dataset: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1232/1232 [00:04<00:00, 261.63 examples/s]
2026-09-09 09:42:32 INFO device.py L1560: 'peak_ram': 2.96GB
2026-09-09 09:42:32 INFO orchestrator.py L619: caching done
Quantizing model.layers.0:   0%|                                                                                                                                                                                                         | 0/52 [00:04<?, ?it/s]2026-09-09 09:42:37 INFO logging.py L355: AR_NGRAM_DEVICE=auto is set, but no ngram modules were found under this block.
quantized 2/2 layers in the block, loss iter 0: 1.155e-10 -> iter 175: 3.893e-11
2026-09-09 09:44:44 INFO device.py L1560: 'peak_ram': 18.39GB, 'peak_vram': 9.79GB
Quantizing model.layers.1:   2%|███▋                                                                                                                                                                                          | 1/52 [02:12<1:52:15, 132.07s/it]quantized 258/258 layers in the block, loss iter 0: 3.177e-10 -> iter 128: 1.398e-10
2026-09-09 09:45:56 INFO device.py L1560: 'peak_ram': 21.13GB, 'peak_vram': 36.0GB
Quantizing model.layers.2:   4%|███████▎                                                                                                                                                                                      | 2/52 [03:39<1:28:09, 105.79s/it]quantized 2/2 layers in the block, loss iter 0: 1.140e-09 -> iter 167: 3.979e-10
2026-09-09 09:47:58 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 36.0GB
Quantizing model.layers.3:   6%|██████████▉                                                                                                                                                                                   | 3/52 [05:26<1:26:41, 106.16s/it]quantized 258/258 layers in the block, loss iter 0: 2.176e-09 -> iter 186: 8.946e-10
2026-09-09 09:48:49 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.4:   8%|██████████████▋                                                                                                                                                                                | 4/52 [06:31<1:11:58, 89.98s/it]quantized 2/2 layers in the block, loss iter 0: 4.170e-09 -> iter 69: 1.979e-09
2026-09-09 09:50:29 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.5:  10%|██████████████████▎                                                                                                                                                                            | 5/52 [07:56<1:09:15, 88.41s/it]quantized 4/4 layers in the block, loss iter 0: 5.805e-09 -> iter 150: 3.085e-09
2026-09-09 09:50:52 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.6:  12%|██████████████████████▎                                                                                                                                                                          | 6/52 [08:19<50:43, 66.17s/it]quantized 258/258 layers in the block, loss iter 0: 1.027e-08 -> iter 123: 4.712e-09
2026-09-09 09:51:42 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.7:  13%|█████████████████████████▉                                                                                                                                                                       | 7/52 [09:25<49:25, 65.90s/it]quantized 2/2 layers in the block, loss iter 0: 9.458e-09 -> iter 164: 6.305e-09
2026-09-09 09:53:12 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.8:  15%|█████████████████████████████▋                                                                                                                                                                   | 8/52 [10:40<50:29, 68.86s/it]quantized 258/258 layers in the block, loss iter 0: 1.650e-08 -> iter 191: 1.017e-08
2026-09-09 09:53:53 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.9:  17%|█████████████████████████████████▍                                                                                                                                                               | 9/52 [11:35<46:13, 64.50s/it]quantized 2/2 layers in the block, loss iter 0: 2.522e-08 -> iter 131: 1.652e-08
2026-09-09 09:55:22 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.10:  19%|████████████████████████████████████▋                                                                                                                                                          | 10/52 [12:49<47:20, 67.63s/it]quantized 258/258 layers in the block, loss iter 0: 4.556e-08 -> iter 187: 2.576e-08
2026-09-09 09:56:03 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.11:  21%|████████████████████████████████████████▍                                                                                                                                                      | 11/52 [13:54<45:29, 66.57s/it]quantized 2/2 layers in the block, loss iter 0: 8.021e-08 -> iter 132: 4.383e-08
2026-09-09 09:57:40 INFO device.py L1560: 'peak_ram': 21.14GB, 'peak_vram': 38.17GB
Quantizing model.layers.12:  23%|████████████████████████████████████████████                                                                                                                                                   | 12/52 [15:08<45:54, 68.87s/it]

[wenhuach21]

still has inference issue #2324