diff --git a/examples/post_training/modelopt/README.md b/examples/post_training/modelopt/README.md index 48e679e4e31..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** | ✅ | ✅ | ✅ | @@ -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= \ 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. 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 new file mode 100644 index 00000000000..c294e03235c --- /dev/null +++ b/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 \ + " 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 old mode 100644 new mode 100755 index 0aa625b875d..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. @@ -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}) diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index 635c18ee545..ceedce606a5 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/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) 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, ), )