Support NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 FP8/NVFP4 PTQ in example

examples/post_training/modelopt/README.md

@@ -32,6 +32,7 @@ knowledge distillation, pruning, speculative decoding, and more.
 | `meta-llama/Llama-4-{Scout,Maverick}-17B-{16,128}E-Instruct` | ✅ | ✅ | - | - |
 | `moonshotai/Kimi-K2-Instruct` | ✅ | ✅ | - | - |
 | `nvidia/NVIDIA-Nemotron-Nano-9B-v2` | ✅ | - | ✅ | ✅ |
+| `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` | ✅ | - | ✅ | ✅ |
 | `openai/gpt-oss-{20b, 120b}` | ✅ | **Online** | ✅ | ✅ |
 | `Qwen/Qwen3-{0.6B, 8B}` | ✅ | ✅ | ✅ | ✅ |
 | `Qwen/Qwen3-{30B-A3B, 235B-A22B}` | **WAR** | ✅ | ✅ | ✅ |
@@ -57,18 +58,19 @@ Provide the pretrained checkpoint path through variable `${HF_MODEL_CKPT}` and p
 Hugging Face-Like quantized checkpoint for TensorRT-LLM, vLLM, or SGLang deployement,
 provide `${EXPORT_DIR}` to `export.sh`.
 
-> **📙 NOTE:** ModelOpt supports different quantization formats. By default, we simulate the
-> low-precision numerical behavior (fake-quant) which can be run on GPUs with compute > 80.
+> **📙 NOTE:** ModelOpt supports different quantization formats which are listed in the [ModelOpt quant configs](https://github.com/NVIDIA/Model-Optimizer/blob/7971fff05882da7eae16eae6bc927d1481dcd63f/modelopt/torch/quantization/config.py#L626).
+> The quant config is specified by the full config name in all-caps, e.g. NVFP4_DEFAULT_CFG.
+> By default, we simulate the low-precision numerical behavior (fake-quant) which can be run on GPUs with compute > 80.
 > Real low-precision paramters (e.g. `E4M3` or `E2M1`)
 > and low-precision compute (e.g. `FP8Linear`) are also supported depending on GPU compute capability.
-> **See [Adanvanced Topics](./ADVANCED.md) for details**.
+> **See [Advanced Topics](./ADVANCED.md) for details**.
 
 ```sh
 \
     TP=1 \
     HF_MODEL_CKPT=<pretrained_model_name_or_path> \
     MLM_MODEL_SAVE=/tmp/Llama-3.2-1B-Instruct_quant \
-    ./quantize.sh meta-llama/Llama-3.2-1B-Instruct nvfp4
+    ./quantize.sh meta-llama/Llama-3.2-1B-Instruct NVFP4_DEFAULT_CFG 
 
 \
     PP=1 \
@@ -78,6 +80,8 @@ provide `${EXPORT_DIR}` to `export.sh`.
     ./export.sh meta-llama/Llama-3.2-1B-Instruct
 ```
 
+For KV cache quantization, add a flag like `MLM_EXTRA_ARGS="--export-kv-cache-quant fp8"` while specifying your desired KV cache precision (see `KV_QUANT_CFG_CHOICES` in `quantize.py`).
+
 ### ⭐ Online BF16 EAGLE3 Training
 
 Online EAGLE3 training has both the target (frozen) and draft models in the memory where the `hidden_states`
@@ -100,7 +104,7 @@ deployment.
     ./export.sh meta-llama/Llama-3.2-1B-Instruct
 ```
 
-See [Adanvanced Topics](./ADVANCED.md) for a `moonshotai/Kimi-K2-Instruct` EAGLE3 training example using `slurm`.
+See [Advanced Topics](./ADVANCED.md) for a `moonshotai/Kimi-K2-Instruct` EAGLE3 training example using `slurm`.
 
 ### ⭐ Offline BF16 EAGLE3 Training
 Unlike online EAGLE3 training, offline workflow precomputes target model `hidden_states` and dumps to disk.

examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+
+if [ -z ${HF_MODEL_CKPT} ]; then
+    HF_MODEL_CKPT=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+    TOKENIZER_MODEL=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+else
+    TOKENIZER_MODEL=${HF_MODEL_CKPT}
+fi
+
+MODEL_ARGS=" \
+    --trust-remote-code \
+    --save-interval 100000 \
+    --micro-batch-size 1 \
+    --moe-token-dispatcher-type allgather \
+    --enable-experimental \
+    --moe-permute-fusion \
+    --use-fused-weighted-squared-relu \
+    --cross-entropy-loss-fusion \
+    --cross-entropy-fusion-impl native \
+    --moe-router-score-function sigmoid \
+    --moe-grouped-gemm \
+    --num-experts 128 \
+    --moe-router-topk 6 \
+    --moe-aux-loss-coeff 1e-4 \
+    --moe-router-topk-scaling-factor 2.5 \
+    --moe-router-enable-expert-bias \
+    --moe-router-dtype fp32 \
+    --moe-router-load-balancing-type seq_aux_loss \
+    --moe-shared-expert-intermediate-size 3712 \
+    \
+    --attention-backend flash \
+    --disable-gloo-process-groups \
+    --is-hybrid-model \
+    --mamba-num-heads 64 \
+    --mamba-head-dim 64 \
+    --hybrid-override-pattern MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME \
+    --use-mcore-models \
+    --untie-embeddings-and-output-weights \
+    --disable-bias-linear \
+    --init-method-std 0.0173 \
+    --position-embedding-type none \
+    --squared-relu \
+    --num-layers 52 \
+    --hidden-size 2688 \
+    --num-attention-heads 32 \
+    --group-query-attention \
+    --num-query-groups 2 \
+    --ffn-hidden-size 1856 \
+    --kv-channels 128 \
+    --normalization RMSNorm \
+    \
+    --tokenizer-type HuggingFaceTokenizer \
+    --bf16 \
+    --seq-length 8192 \
+    --max-position-embeddings 8192 \
+    --export-model-type MambaModel \
+    "

examples/post_training/modelopt/convert_model.py

@@ -136,7 +136,10 @@ def check_arguments():
         print_rank_0(
             "Import model from Hugging Face checkpoint in dtype {}.".format(str(import_dtype))
         )
-        import_kwargs = {"dtype": import_dtype}
+        import_kwargs = {
+            "dtype": import_dtype,
+            "moe_router_dtype": args.moe_router_dtype,
+        }
         if modelopt_version_at_least("0.41.0"):
             import_kwargs.update({"trust_remote_code": args.trust_remote_code})
         import_mcore_gpt_from_hf(

examples/post_training/modelopt/export.py

@@ -5,6 +5,7 @@
 import os
 import sys
 import warnings
+from pathlib import Path
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../")))
 
@@ -74,8 +75,11 @@ def add_modelopt_export_args(parser):
     unwrapped_model = unwrap_model(model)[0]
     unwrapped_model.to_empty(device="cpu")
 
-    if args.load is not None:
+    if args.load is not None and Path(args.load).is_dir():
         _ = load_modelopt_checkpoint(model)
+    else:
+        raise ValueError(f"Invalid load checkpoint directory: {args.load}")
+
 
     # Decide whether we are exporting only the extra_modules (e.g. EAGLE3).
     # Only the last pp stage may have extra_modules, hence broadcast from the last rank.
@@ -90,6 +94,7 @@ def add_modelopt_export_args(parser):
         "export_extra_modules": export_extra_modules,
         "dtype": torch.bfloat16,
         "export_dir": args.export_dir,
+        "moe_router_dtype": unwrapped_model.config.moe_router_dtype,
     }
     if modelopt_version_at_least("0.41.0"):
         export_kwargs.update({"trust_remote_code": args.trust_remote_code})

examples/post_training/modelopt/quantize.py

@@ -52,15 +52,9 @@
 
 warnings.filterwarnings("ignore")
 
-# TODO deprecate these aliases in the next release
-QUANT_CFG_CHOICES = {
-    "int8_sq": mtq.INT8_SMOOTHQUANT_CFG,
-    "fp8": mtq.FP8_DEFAULT_CFG,
-    "fp8_blockwise": mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG,
-    "int4_awq": mtq.INT4_AWQ_CFG,
-    "w4a8_awq": mtq.W4A8_AWQ_BETA_CFG,
-    "nvfp4": mtq.NVFP4_DEFAULT_CFG,
-}
+QUANT_CFG_CHOICES = {}
+
+# Auto-load all quant configs by full name
 for k in mtq.config.choices:
     QUANT_CFG_CHOICES[k] = getattr(mtq, k)
 

megatron/core/post_training/modelopt/mamba/model_specs.py

@@ -2,6 +2,7 @@
 
 from megatron.core.extensions.transformer_engine import TEDotProductAttention
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.post_training.modelopt.layers import Norm
 from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules
 from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules
@@ -83,9 +84,23 @@ def get_mamba_stack_modelopt_spec(
         ),
     )
 
+    moe_layer = ModuleSpec(
+        module=TransformerLayer,
+        submodules=TransformerLayerSubmodules(
+            pre_mlp_layernorm=Norm,
+            mlp=get_moe_module_spec(
+                use_te=False, num_experts=8, moe_grouped_gemm=False  # Can be anything non None
+            ),
+            mlp_bda=get_bias_dropout_add,
+        ),
+    )
+
     return ModuleSpec(
         module=MambaStack,
         submodules=MambaStackSubmodules(
-            mamba_layer=mamba_layer, attention_layer=attention_layer, mlp_layer=mlp_layer
+            mamba_layer=mamba_layer,
+            attention_layer=attention_layer,
+            mlp_layer=mlp_layer,
+            moe_layer=moe_layer,
         ),
     )