From b8a494afcf63e662d1ef228921e07d096d0c5941 Mon Sep 17 00:00:00 2001 From: jenchen13 Date: Mon, 5 Jan 2026 15:42:31 -0800 Subject: [PATCH 01/11] nano3 config for modelopt Signed-off-by: jenchen13 --- .../nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh new file mode 100644 index 00000000000..8932244af93 --- /dev/null +++ b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh @@ -0,0 +1,58 @@ +#!/bin/bash + +if [ -z ${HF_MODEL_CKPT} ]; then + HF_MODEL_CKPT=nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3 + TOKENIZER_MODEL=nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3 +else + TOKENIZER_MODEL=${HF_MODEL_CKPT} +fi + +#TOKENIZER_MODEL="/lustre/fsw/portfolios/llmservice/projects/llmservice_nlp_fm/nemotron6/tokenizers/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json" + +MODEL_ARGS=" \ + --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 \ + " From da36e4c4df0b90ca7c048f27cab623e2d987b7e9 Mon Sep 17 00:00:00 2001 From: jenchen13 Date: Mon, 5 Jan 2026 15:46:50 -0800 Subject: [PATCH 02/11] remove comment Signed-off-by: jenchen13 --- .../modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh | 2 -- 1 file changed, 2 deletions(-) diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh index 8932244af93..57260f3522e 100644 --- a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh +++ b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh @@ -7,8 +7,6 @@ else TOKENIZER_MODEL=${HF_MODEL_CKPT} fi -#TOKENIZER_MODEL="/lustre/fsw/portfolios/llmservice/projects/llmservice_nlp_fm/nemotron6/tokenizers/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json" - MODEL_ARGS=" \ --save-interval 100000 \ --micro-batch-size 1 \ From 4a5947f756b2ce85d2b6bc38784b9ca21c140cf4 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 9 Jan 2026 18:09:17 +0000 Subject: [PATCH 03/11] deprecate quant aliases in quantize Signed-off-by: Jennifer Chen --- examples/post_training/modelopt/ADVANCED.md | 2 +- examples/post_training/modelopt/README.md | 13 ++++++++----- examples/post_training/modelopt/quantize.py | 10 +--------- 3 files changed, 10 insertions(+), 15 deletions(-) diff --git a/examples/post_training/modelopt/ADVANCED.md b/examples/post_training/modelopt/ADVANCED.md index 28aad7d7964..5a7b834563d 100644 --- a/examples/post_training/modelopt/ADVANCED.md +++ b/examples/post_training/modelopt/ADVANCED.md @@ -19,7 +19,7 @@ quantization. ```sh \ HF_MODEL_CKPT= \ - bash quantize.sh [model_conf] [qformat] + ./quantize.sh [model_conf] [qformat] ``` > **❗ IMPORTANT:** `model_conf` is used to get the corresponding Megatron-LM `${MODEL_ARGS}`. For example, diff --git a/examples/post_training/modelopt/README.md b/examples/post_training/modelopt/README.md index 48e679e4e31..1a2d3edd04f 100644 --- a/examples/post_training/modelopt/README.md +++ b/examples/post_training/modelopt/README.md @@ -57,18 +57,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= \ 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 +79,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 `quantize.py` `KV_QUANT_CFG_CHOICES`). + ### ⭐ Online BF16 EAGLE3 Training Online EAGLE3 training has both the target (frozen) and draft models in the memory where the `hidden_states` @@ -100,7 +103,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. diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index 635c18ee545..c9cfb1cdf38 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -52,15 +52,7 @@ 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, -} +# Auto-load all quant configs by full name for k in mtq.config.choices: QUANT_CFG_CHOICES[k] = getattr(mtq, k) From 74d1be09c0b850427f5da51f6238fdb53c72286e Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 9 Jan 2026 18:38:05 +0000 Subject: [PATCH 04/11] hybrid moe quantization fixes Signed-off-by: Jennifer Chen --- .../post_training/modelopt/convert_model.py | 5 ++++- examples/post_training/modelopt/export.py | 1 + .../post_training/modelopt/mamba/model_specs.py | 17 ++++++++++++++++- 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/examples/post_training/modelopt/convert_model.py b/examples/post_training/modelopt/convert_model.py index 53ae25f8d92..cf5f6e5bbbb 100644 --- a/examples/post_training/modelopt/convert_model.py +++ b/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( diff --git a/examples/post_training/modelopt/export.py b/examples/post_training/modelopt/export.py index 0aa625b875d..c7540693ed6 100644 --- a/examples/post_training/modelopt/export.py +++ b/examples/post_training/modelopt/export.py @@ -90,6 +90,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}) diff --git a/megatron/core/post_training/modelopt/mamba/model_specs.py b/megatron/core/post_training/modelopt/mamba/model_specs.py index e8a14212bc3..0a38d05b980 100755 --- a/megatron/core/post_training/modelopt/mamba/model_specs.py +++ b/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, ), ) From 54aea08fdd5111bb48e734468c57542125ad4783 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 9 Jan 2026 19:10:50 +0000 Subject: [PATCH 05/11] rename Signed-off-by: Jennifer Chen --- ...n-Nano-31B-A3-v3.sh => NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename examples/post_training/modelopt/conf/nvidia/{NVIDIA-Nemotron-Nano-31B-A3-v3.sh => NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh} (100%) diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh similarity index 100% rename from examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3.sh rename to examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh From ebcf12045fcb91d6ca384cc8350312ab7b6b7288 Mon Sep 17 00:00:00 2001 From: Jenny Chen Date: Wed, 14 Jan 2026 10:57:06 -0800 Subject: [PATCH 06/11] Update examples/post_training/modelopt/README.md Co-authored-by: Asha Anoosheh --- examples/post_training/modelopt/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/post_training/modelopt/README.md b/examples/post_training/modelopt/README.md index 1a2d3edd04f..b6103bdfc85 100644 --- a/examples/post_training/modelopt/README.md +++ b/examples/post_training/modelopt/README.md @@ -79,7 +79,7 @@ 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 `quantize.py` `KV_QUANT_CFG_CHOICES`). +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 From 10ac001e62a947ea997b19602a8c1b304d13ab3d Mon Sep 17 00:00:00 2001 From: jenchen13 Date: Thu, 15 Jan 2026 10:44:19 -0800 Subject: [PATCH 07/11] raise error if invalid export dir Signed-off-by: jenchen13 --- examples/post_training/modelopt/export.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) mode change 100644 => 100755 examples/post_training/modelopt/export.py diff --git a/examples/post_training/modelopt/export.py b/examples/post_training/modelopt/export.py old mode 100644 new mode 100755 index c7540693ed6..9dc66eecb6d --- a/examples/post_training/modelopt/export.py +++ b/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. From e78758ff4295389cec0fb32a81080c5a50e832e5 Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Fri, 23 Jan 2026 14:07:20 -0800 Subject: [PATCH 08/11] chg: update HF Hub reference --- .../conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh index 57260f3522e..18364a2d4e4 100644 --- a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh +++ b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh @@ -1,8 +1,8 @@ #!/bin/bash if [ -z ${HF_MODEL_CKPT} ]; then - HF_MODEL_CKPT=nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3 - TOKENIZER_MODEL=nvidia/NVIDIA-Nemotron-Nano-31B-A3-v3 + 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 From ff17d53d81c47c6b0e5371133ad42091d7bdeeda Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Fri, 23 Jan 2026 16:28:59 -0800 Subject: [PATCH 09/11] chg: update --- examples/post_training/modelopt/quantize.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index c9cfb1cdf38..ceedce606a5 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -52,6 +52,8 @@ warnings.filterwarnings("ignore") +QUANT_CFG_CHOICES = {} + # Auto-load all quant configs by full name for k in mtq.config.choices: QUANT_CFG_CHOICES[k] = getattr(mtq, k) From 18bb514b4fa527b9da415ef0260a9cbfb88b59ee Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Fri, 23 Jan 2026 21:45:25 -0800 Subject: [PATCH 10/11] chg: update --- .../modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh index 18364a2d4e4..c294e03235c 100644 --- a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh +++ b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh @@ -8,6 +8,7 @@ else fi MODEL_ARGS=" \ + --trust-remote-code \ --save-interval 100000 \ --micro-batch-size 1 \ --moe-token-dispatcher-type allgather \ From a956892b755db17b3473de814cd4ea6b3464a0a3 Mon Sep 17 00:00:00 2001 From: Chenhan Yu Date: Mon, 26 Jan 2026 09:23:57 -0800 Subject: [PATCH 11/11] chg: update README support matrix --- examples/post_training/modelopt/ADVANCED.md | 2 +- examples/post_training/modelopt/README.md | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/post_training/modelopt/ADVANCED.md b/examples/post_training/modelopt/ADVANCED.md index 5a7b834563d..28aad7d7964 100644 --- a/examples/post_training/modelopt/ADVANCED.md +++ b/examples/post_training/modelopt/ADVANCED.md @@ -19,7 +19,7 @@ quantization. ```sh \ HF_MODEL_CKPT= \ - ./quantize.sh [model_conf] [qformat] + bash quantize.sh [model_conf] [qformat] ``` > **❗ IMPORTANT:** `model_conf` is used to get the corresponding Megatron-LM `${MODEL_ARGS}`. For example, diff --git a/examples/post_training/modelopt/README.md b/examples/post_training/modelopt/README.md index b6103bdfc85..6ebdd8ac5d6 100644 --- a/examples/post_training/modelopt/README.md +++ b/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** | ✅ | ✅ | ✅ |