diff --git a/README.md b/README.md index 4852bfa7884..487cb9e535e 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ Model Optimizer is also integrated with [NVIDIA Megatron-Bridge](https://github. ## Latest News -- [2026/05/27] [**End-to-end optimization tutorial for Nemotron-3-Nano-30B-A3B**](./examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16): Pruning + distillation (with long context extension) + FP8 quantization achieving 2.6× vLLM throughput and 2.6× memory reduction. +- [2026/05/27] [**End-to-end Optimization tutorial for Nemotron-3-Nano-30B-A3B**](./examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16): Pruning + two-phase distillation + FP8 quantization achieving 2.6× vLLM throughput and 2.6× memory reduction. - [2026/05/13] [**Puzzletron**](./examples/puzzletron): A new algorithm for heterogeneous pruning & NAS of LLM and VLM models. - [2026/04/15] Customer story: [Domyn compresses Colosseum-355B → 260B using ModelOpt's Minitron pruning + distillation](https://www.domyn.com/blog/domyn-large-the-journey-of-a-european-sovereign-ai-model-for-regulated-industries) - [2026/03/17] Customer story: [Bielik.AI builds Bielik Minitron 7B (33% smaller, 50% faster, 90% quality retained) using ModelOpt's Minitron pruning + distillation](https://bielik.ai/en/nvidia-gtc-bielik-minitron-premiere/) @@ -102,12 +102,12 @@ more fine-grained control on installed dependencies or for alternative docker im | **Technique** | **Description** | **Examples** | **Docs** | | :------------: | :------------: | :------------: | :------------: | -| Post Training Quantization | Compress model size by 2x-4x, speeding up inference while preserving model quality! | \[[LLMs](./examples/llm_ptq/)\] \[[diffusers](./examples/diffusers/)\] \[[VLMs](./examples/vlm_ptq/)\] \[[onnx](./examples/onnx_ptq/)\] \[[windows](./examples/windows/)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | -| Quantization Aware Training | Refine accuracy even further with a few training steps! | \[[Hugging Face](./examples/llm_qat/)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | -| Pruning | Reduce your model size and accelerate inference by removing unnecessary weights! | \[[General](./examples/pruning/)\] \[[Megatron-Bridge](./examples/megatron_bridge/README.md#pruning)\] | | -| Distillation | Reduce deployment model size by teaching small models to behave like larger models! | \[[Megatron-Bridge](./examples/llm_distill/README.md#knowledge-distillation-kd-in-nvidia-megatron-bridge-framework)\] \[[Megatron-LM](./examples/llm_distill/README.md#knowledge-distillation-kd-in-nvidia-megatron-lm-framework)\] \[[Hugging Face](./examples/llm_distill/)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/4_distillation.html)\] | -| Speculative Decoding | Train draft modules to predict extra tokens during inference! | \[[Megatron](./examples/speculative_decoding#mlm-example)\] \[[Hugging Face](./examples/speculative_decoding/)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/5_speculative_decoding.html)\] | -| Sparsity | Efficiently compress your model by storing only its non-zero parameter values and their locations | \[[PyTorch](./examples/llm_sparsity/)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/6_sparsity.html)\] | +| Post Training Quantization | Compress model size by 2x-4x, speeding up inference while preserving model quality! | \[[HF LLMs / VLMs](./examples/llm_ptq/)\] \[[Megatron-Bridge LLMs / VLMs](./examples/megatron_bridge/)\] \[[Diffusers](./examples/diffusers/)\] \[[ONNX](./examples/onnx_ptq/)\] \[[Windows](./examples/windows/)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | +| Quantization Aware Training / Distillation | Refine accuracy of quantized models even further with a few training steps! | \[[Hugging Face](./examples/llm_qat/)\] \[[Megatron-Bridge](./examples/megatron_bridge)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | +| Pruning | Reduce your model parameters or memory footprint and accelerate inference by removing unnecessary weights! | \[[General](./examples/pruning/)\] \[[Megatron-Bridge](./examples/megatron_bridge/)\] | | +| Distillation | Reduce deployment model size by teaching small models to behave like larger models! | \[[Hugging Face](./examples/llm_distill/)\] \[[Megatron-Bridge](./examples/megatron_bridge/)\] \[[Megatron-LM](./examples/llm_distill/README.md#knowledge-distillation-kd-in-nvidia-megatron-lm-framework)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/4_distillation.html)\] | +| Speculative Decoding | Train draft modules to predict extra tokens during inference! | \[[Hugging Face](./examples/speculative_decoding/)\] \[[Megatron-LM](./examples/speculative_decoding#mlm-example)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/5_speculative_decoding.html)\] | +| Sparsity | Efficiently compress your model by storing only its non-zero parameter values and their locations | \[[Hugging Face](./examples/llm_sparsity/)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/6_sparsity.html)\] | @@ -131,8 +131,8 @@ more fine-grained control on installed dependencies or for alternative docker im | Model Type | Support Matrix | |------------|----------------| | LLM Quantization | [View Support Matrix](./examples/llm_ptq/README.md#support-matrix) | -| Diffusers Quantization | [View Support Matrix](./examples/diffusers/README.md#support-matrix) | | VLM Quantization | [View Support Matrix](./examples/vlm_ptq/README.md#support-matrix) | +| Diffusers Quantization | [View Support Matrix](./examples/diffusers/README.md#support-matrix) | | ONNX Quantization | [View Support Matrix](./examples/torch_onnx/README.md#onnx-export-supported-llm-models) | | Windows Quantization | [View Support Matrix](./examples/windows/README.md#support-matrix) | | Quantization Aware Training | [View Support Matrix](./examples/llm_qat/README.md#support-matrix) | diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 6bb1f19be5c..98954c4caf9 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -301,6 +301,9 @@ torchrun --nproc_per_node 1 prune_minitron.py --help > uneven PP by setting `--num_layers_in_first_pipeline_stage` and `--num_layers_in_last_pipeline_stage`. > E.g. for Qwen3-8B with 36 layers and 8 GPUs, you can set both to 3 to get 3-5-5-5-5-5-5-3 layers per GPU. +> [!NOTE] +> NAS-based pruning requires ~2x the GPU memory of Manual pruning because it needs to simultaneously hold original model while evaluating each pruned candidate. + > [!NOTE] > If pruning a Nemotron model and you want to save the pruned model back in HF format, please downgrade to `transformers<5` via `python -m pip install "transformers<5"` before pruning. diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 2f5fe777e91..55bd572a01b 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -256,6 +256,9 @@ def get_args() -> argparse.Namespace: raise ValueError("--prune_export_config must parse to a dictionary.") args.prune_export_config = prune_export_config + if args.inference_batch_size is None: + args.inference_batch_size = args.calib_batch_size + print_args(args) return args @@ -374,11 +377,7 @@ def score_func(m): pruning_config["hparams_to_skip"] = args.hparams_to_skip pruning_config["top_k"] = args.top_k # memory_mb constraint requires batch_size and seq_length - pruning_config["batch_size"] = ( - args.inference_batch_size - if args.inference_batch_size is not None - else args.calib_batch_size - ) + pruning_config["batch_size"] = args.inference_batch_size pruning_config["seq_length"] = args.seq_length print_rank_0(f"Pruning constraints: {pruning_constraints}") diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index f3a4db3676b..ee7ea0c3b56 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -298,9 +298,8 @@ def main(args: argparse.Namespace): num_samples=args.calib_num_samples, seq_length=args.seq_length, batch_size=args.calib_batch_size, - # Calibrate on unpacked sequences. pack=True is Megatron pretraining-style global-stream - # document packing, which changes the per-sample calibration statistics. - pack=False, + # pack=True uses Megatron pretraining-style global-stream document packing + pack=True, ) else: warn_rank_0("Dynamic or weight-only quantization detected; skipping calibration.") @@ -327,7 +326,6 @@ def main(args: argparse.Namespace): if dist.is_master(): mtq.print_quant_summary(unwrapped_model, args.export_megatron_path) - print_rank_0(f"\nSaving quantized model to {args.export_megatron_path} in Megatron format...") bridge.save_megatron_model( model, args.export_megatron_path, diff --git a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/ABLATIONS.md b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/ABLATIONS.md index a240aa4940c..15d9c121cc5 100644 --- a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/ABLATIONS.md +++ b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/ABLATIONS.md @@ -1,5 +1,8 @@ # Ablations: Nemotron-3-Nano-30B-A3B-BF16 +> [!NOTE] +> Every benchmark table in this document is from evaluations run **without** eval-time tool calling (no Python sandbox). The [main README](README.md#results) reports both with-tools and no-tools numbers for GPQA and AIME. Consequently, the gains attributed to the `tool_calling` training data in the [data-blend ablation](#effect-of-data-blend-tool_calling) reflect **training-time** exposure to agentic/function-call reasoning, not tools invoked during evaluation. + ## Pruning > [!NOTE] @@ -255,7 +258,7 @@ Per-benchmark difference (Δ = new − old), at matching iter counts during the **Summary — new blend is preferred:** -- **GPQA Diamond** is the most consistent winner — positive Δ at every single checkpoint (+1.2/+0.7/+2.7/+0.6/+0.7). The tool_calling allocation helps throughout training, not just at the end. Calculator-tool use is common in GPQA quantitative items. +- **GPQA Diamond** is the most consistent winner — positive Δ at every single checkpoint (+1.2/+0.7/+2.7/+0.6/+0.7). The tool_calling allocation helps throughout training, not just at the end. Since these evals run **without** eval-time tools, the gain comes from training-time exposure to agentic/function-call reasoning traces transferring to GPQA — not from the model invoking a tool during evaluation. - **SciCode** trends positive early (+2.5/+3.6 at 2.5B/20B) but neutral-to-slightly-negative mid-training (-1.7/-2.4 at 40B/60B) and positive again at 80B (+0.3). Even averaged-of-8, SciCode has the largest residual checkpoint-to-checkpoint noise of any benchmark. - **AIME 2025** trends positive from 40B onward (+3.0/+1.3/+0.5). The early dip is consistent with the math share dropping from 30%→27%; the model catches up once enough tokens have been seen. - **IFBench** is small-positive early (+0.8/+1.8) then dips mid-training (-1.4/-1.1) before recovering to neutral at 80B (0.0). Net roughly flat over training. @@ -264,7 +267,7 @@ Per-benchmark difference (Δ = new − old), at matching iter counts during the ### Effect of long context training -After the 8K-seq-length phase of the old-blend run, training was continued with the same blend but with `seq_length` increased from 8192 to 32768. The longer-context phase is short (200–1000 additional iters) but disproportionately impactful. Numbers below use the old-blend run because it has the longer LC sweep (1000 iters / +25B tokens). The new-blend run's shorter LC phase (+800 iters / +20B tokens, ending at 100B) is in the [main README](README.md) and shows the same qualitative finding (large AIME jump immediately at the start of the LC phase). +After the 8K-seq-length phase of the old-blend run, training was continued with the same blend but with `seq_length` increased from 8192 to 32768. The longer-context phase is short (200–1000 additional iters) but disproportionately impactful. Numbers below use the old-blend run because it has the longer LC sweep (1000 iters / +25B tokens). The new-blend run's shorter LC phase (+800 iters / +20B tokens, ending at 100B) is in the [main README](README.md) and shows the same qualitative finding (large AIME jump immediately at the start of the LC phase) — though at higher absolute AIME scores there, since the README evals enable tools. Average is over the 6 reasoning benchmarks (excluding MMLU), matching the main README: diff --git a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md index 1c7729da51a..22a31ff9db7 100644 --- a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md +++ b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md @@ -7,24 +7,40 @@ End-to-end optimization of [NVIDIA-Nemotron-3-Nano-30B-A3B-BF16](https://hugging 3. **[Distillation](#3-distillation)** — recovering accuracy via Megatron-Bridge knowledge distillation 4. **[Evaluation](#4-evaluation)** — benchmarking with NeMo Evaluator across MMLU Pro, GPQA Diamond, AIME, and more 5. **[Quantization](#5-quantization)** — FP8 PTQ on the distilled checkpoint using ModelOpt's `examples/megatron_bridge/quantize.py` script -6. **[vLLM Inference Benchmarking](#6-vllm-inference-benchmarking)** — throughput comparison of BF16 vs FP8 on a single H100 +6. **[vLLM Inference Benchmarking](#6-vllm-inference-benchmarking)** — throughput comparison across BF16 and FP8 on a single H100 ## Results -![Benchmark Recovery During Knowledge Distillation](figures/learning_curves.png) - -| Model | MMLU Pro | GPQA Diamond | LiveCodeBench v6 | AIME 2025 | IFBench | SciCode (Subtask) | Average | -| --- | --- | --- | --- | --- | --- | --- | --- | -| Pruned 22B/A3.0B (no distillation) | 47.1 | 33.5 | 27.4 | 15.5 | 36.9 | 12.1 | 28.8 | -| Distill @ 2.5B tokens (100 iters at 8K SeqLen) | 73.3 | 63.7 | 55.3 | 77.6 | 59.1 | 25.1 | 59.0 | -| Distill @ 20B tokens (800 iters at 8K SeqLen) | 74.8 | 66.0 | 62.3 | 79.6 | 65.4 | 26.1 | 62.4 | -| Distill @ 40B tokens (1600 iters at 8K SeqLen) | 76.4 | 67.2 | 62.3 | 79.8 | 66.0 | 26.6 | 63.1 | -| Distill @ 60B tokens (2400 iters at 8K SeqLen) | 76.1 | 68.1 | 63.6 | 78.8 | 67.3 | 27.0 | 63.5 | -| Distill @ 80B tokens (3200 iters at 8K SeqLen) | 76.5 | 69.1 | 63.9 | 80.7 | 66.5 | 29.0 | 64.3 | -| Distill @ 82.5B tokens (+100 iters at 32K SeqLen) | 76.2 | 69.8 | 64.8 | 87.0 | 68.2 | 27.0 | 65.5 | -| Distill @ 100B tokens (+800 iters at 32K SeqLen) - **BF16** | 76.6 | 69.6 | 66.1 | 87.3 | 68.9 | 28.4 | 66.2 | -| Distill @ 100B tokens + **FP8 Quantize** | 76.7 | 70.7 | 65.5 | 87.3 | 69.0 | 28.5 | 66.3 | -| Nemotron-3-Nano-30B-A3B-BF16 (official, 31.6B/A3.6B) | 78.0 | 70.3 | 67.9 | 87.1 | 69.1 | 31.8 | 67.4 | +![Benchmark Recovery (BF16) During Knowledge Distillation](figures/learning_curves.png) + +Main results — all models evaluated with the same [setup](#4-evaluation). Values are `mean ± std_dev` across repeats: + +| Model | MMLU Pro | GPQA Diamond | GPQA Diamond (w. tools) | LiveCodeBench v6 | AIME 2025 | AIME 2025 (w. tools) | IFBench | SciCode (Subtask) | Average | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | +| **Pruned 22B/A3.0B + Distilled — 100B tokens (BF16)** | 76.7 | 68.9 ± 2.5 | 71.4 ± 2.1 | 65.1 ± 1.0 | 86.4 ± 3.7 | 97.2 ± 3.3 | 69.2 | 28.8 ± 1.7 | 70.5 | +|   ↳ **FP8** (quantized from BF16) | 75.9 | 71.1 ± 1.4 | 70.4 ± 1.1 | 64.8 ± 0.9 | 87.0 ± 4.2 | 95.1 ± 4.8 | 68.2 | 28.7 ± 2.5 | 70.2 | +| **Official Nemotron-3-Nano-30B-A3B-BF16 (31.6B/A3.6B)** | 78.2 | 70.3 ± 1.7 | 74.2 ± 1.9 | 68.9 ± 0.9 | 86.8 ± 4.4 | 97.7 ± 3.3 | 69.2 | 31.8 ± 1.2 | 72.1 | + +
+Full results — pruning baseline and full distillation trajectory (click to expand) + +| Model | MMLU Pro | GPQA Diamond | GPQA Diamond (w. tools) | LiveCodeBench v6 | AIME 2025 | AIME 2025 (w. tools) | IFBench | SciCode (Subtask) | Average | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | +| Pruned 22B/A3.0B (no distillation) | 47.8 | 33.5 ± 2.3 | 34.2 ± 3.0 | 26.7 ± 1.3 | 15.1 ± 3.8 | 17.8 ± 5.5 | 39.9 | 13.1 ± 1.8 | 28.5 | +| Distill @ 2.5B tokens (100 iters at 8K SeqLen) | 73.0 | 62.2 ± 1.4 | 62.0 ± 2.1 | 58.4 ± 1.9 | 79.7 ± 5.4 | 90.0 ± 5.8 | 59.9 | 21.3 ± 2.5 | 63.3 | +| Distill @ 20B tokens (800 iters at 8K SeqLen) | 74.9 | 66.0 ± 2.1 | 68.0 ± 1.8 | 62.5 ± 0.9 | 78.6 ± 3.8 | 88.3 ± 5.2 | 67.0 | 25.2 ± 2.3 | 66.3 | +| Distill @ 40B tokens (1600 iters at 8K SeqLen) | 75.6 | 67.2 ± 2.2 | 69.1 ± 2.2 | 62.4 ± 1.0 | 79.0 ± 3.9 | 92.5 ± 4.3 | 66.8 | 27.4 ± 1.3 | 67.5 | +| Distill @ 60B tokens (2400 iters at 8K SeqLen) | 76.0 | 68.4 ± 2.0 | 70.1 ± 2.0 | 64.3 ± 1.5 | 79.6 ± 4.5 | 92.0 ± 4.2 | 67.6 | 28.4 ± 2.2 | 68.3 | +| Distill @ 80B tokens (3200 iters at 8K SeqLen) | 76.7 | 68.7 ± 3.1 | 68.2 ± 1.8 | 63.9 ± 0.9 | 81.6 ± 6.0 | 93.4 ± 4.6 | 69.0 | 28.5 ± 2.7 | 68.8 | +| Distill @ 82.5B tokens (+100 iters at 32K SeqLen) | 76.6 | 70.0 ± 0.9 | 70.1 ± 1.6 | 65.4 ± 1.0 | 86.8 ± 4.5 | 96.7 ± 3.5 | 68.9 | 27.9 ± 2.7 | 70.3 | +| Distill @ 100B tokens (+800 iters at 32K SeqLen) - **BF16** | 76.7 | 68.9 ± 2.5 | 71.4 ± 2.1 | 65.1 ± 1.0 | 86.4 ± 3.7 | 97.2 ± 3.3 | 69.2 | 28.8 ± 1.7 | 70.5 | +| Distill @ 100B tokens + **FP8 Quantize** | 75.9 | 71.1 ± 1.4 | 70.4 ± 1.1 | 64.8 ± 0.9 | 87.0 ± 4.2 | 95.1 ± 4.8 | 68.2 | 28.7 ± 2.5 | 70.2 | +| Nemotron-3-Nano-30B-A3B-BF16 (official, 31.6B/A3.6B) | 78.2 | 70.3 ± 1.7 | 74.2 ± 1.9 | 68.9 ± 0.9 | 86.8 ± 4.4 | 97.7 ± 3.3 | 69.2 | 31.8 ± 1.2 | 72.1 | + +
+ +> [!NOTE] +> Some of these benchmarks are very noisy with a large per-run spread (e.g. AIME, which has only 30 problems, and SciCode), so small differences are not always meaningful. For a more reliable comparison, use a larger `num_repeats` in practice. ### vLLM Throughput (single H100, ISL=32768, OSL=1024) @@ -102,13 +118,13 @@ DATA_BLEND=" \ #### General Guidelines -The optimal blend is 30% pretraining and 70% post-training data. Exact proportions may vary depending on the benchmarks you care about. The blend above was designed to maximize recovery on popular General Knowledge, Reasoning, Instruction Following, and Tool Calling benchmarks. The key design decisions were: +The optimal blend is 30% pretraining and 70% post-training data. Exact proportions may vary depending on the benchmarks you care about. The blend above was designed to maximize recovery on popular General Knowledge, Reasoning, Coding, and Instruction Following benchmarks. The key design decisions were: - **30% pretraining data** closes the MMLU gap that arises from training exclusively on reasoning-heavy post-training data. The General split (20%) is upweighted specifically to recover general knowledge recall. - **Math (27%)** is the largest post-training category because AIME and MMLU Pro respond strongly to more math reasoning tokens. We use a mix of `Nemotron-Math-v2` and `Nemotron-SFT-Math-v3` for higher quality math reasoning signal with full reasoning traces. - **Science (13%)** uses `Nemotron-Post-Training-Dataset-v1 / stem` as the primary source for volume and GPQA stability, with small allocations to `Nemotron-Science-v1` MCQ/RQA subsets for format alignment with GPQA's multiple-choice structure. - **Instruction following (5%)** saturates quickly so a small allocation is sufficient. -- **Tool calling (5%)** uses `Nemotron-Agentic-v1 / tool_calling`. Our evals run with `--enable-auto-tool-choice`, so the student needs explicit exposure to function-call schemas; this helps SciCode (heavy Python tool use) and GPQA Diamond (which can benefit from calculator tools). +- **Tool calling (5%)** uses `Nemotron-Agentic-v1 / tool_calling`. Our GPQA and AIME evals enable a Python sandbox tool (`--enable-auto-tool-choice`), so the student needs explicit exposure to function-call schemas to use it; this helps GPQA Diamond and AIME 2025 (which can offload quantitative steps to the tool). This blend intentionally omits capabilities not targeted in this experiment (e.g. multilingual, SWE). Depending on what benchmarks matter for your use case, you can substitute or add datasets from the [Nemotron Post-Training v3 collection](https://huggingface.co/collections/nvidia/nemotron-post-training-v3), for example: @@ -126,40 +142,41 @@ When adding new datasets, reduce weights of lower-priority categories proportion Here we prune the [NVIDIA-Nemotron-3-Nano-30B-A3B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16) HuggingFace checkpoint from 31.6B/A3.6B to 3.0B active parameters. The output is a pruned HuggingFace checkpoint that feeds into the distillation step. -Run on **1 node with 8x H100** (~1 hour) +Run on **1 node with 8x H100** (~30 mins)
Pruning command (click to expand) ```bash torchrun --nproc_per_node 8 /opt/Model-Optimizer/examples/megatron_bridge/prune_minitron.py \ - --pp_size 8 \ --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \ --trust_remote_code \ - --prune_target_params 28e9 \ - --prune_target_active_params 3e9 \ - --hparams_to_skip num_attention_heads \ + --pp_size 8 \ + --num_layers_in_first_pipeline_stage 5 \ + --num_layers_in_last_pipeline_stage 5 \ + --calib_batch_size 8 \ --seq_length 8192 \ - --output_hf_path /path/to/Nemotron-3-Nano-30B-A3B-Pruned-A3.0B \ - --top_k 20 \ - --max_depth_pruning 0.15 \ - --max_width_pruning 0.30 \ + --prune_target_active_params 3e9 \ + --prune_target_params 24e9 \ --prune_score_func mmlu_10pct_bs32 \ - --num_layers_in_first_pipeline_stage 5 \ - --num_layers_in_last_pipeline_stage 5 + --max_width_pruning 0.30 \ + --max_depth_pruning 0.15 \ + --hparams_to_skip num_attention_heads \ + --top_k 10 \ + --output_hf_path /path/to/Nemotron-3-Nano-30B-A3B-Pruned-A3.0B ``` Non-default arguments: -- `--hparams_to_skip num_attention_heads` (default: none) — attention heads pruning is harder to recover, hence skipped -- `--seq_length 8192` (default: 4096) — dataset has longer sequences +- `--num_layers_in_first_pipeline_stage 5 --num_layers_in_last_pipeline_stage 5` — Uneven pipeline parallelism since 52 layers is not divisible by 8 GPUs +- `--calib_batch_size 8` (default: 1) — faster calibration with larger batch size using more memory. +- `--seq_length 8192` (default: 4096) — more tokens for better MoE calibration - `--prune_target_active_params 3e9` — MoE-specific; the **primary** pruning constraint — targets active params rather than total params, which is what matters for MoE inference cost -- `--prune_target_params 28e9` — upper bound on total params only; the actual pruned model total can range anywhere from ~20B to 28B depending on which architecture wins — see pruning logs below for the top 20 candidates. You may also skip this argument all together for simplicity. -- `--top_k 20` (default: 10) — larger candidate pool for better architecture search -- `--max_depth_pruning 0.15` (default: 0.20) — tighter constraint since candidates with 42–46 layers universally fail for this model -- `--max_width_pruning 0.30` (default: 0.40) — tighter constraint to prevent head_dim≤48 and hidden=2048 dead zones +- `--prune_target_params 24e9` — upper bound on total params only; the actual pruned model total can range anywhere from ~19B to 24B depending on which architecture wins — see pruning logs below for the top 10 candidates. You may also skip this argument all together for simplicity. - `--prune_score_func mmlu_10pct_bs32` (default: `mmlu_10pct_bs1`) — batch_size=32 for ~3–4× faster candidate scoring -- `--num_layers_in_first_pipeline_stage 5 --num_layers_in_last_pipeline_stage 5` — Uneven pipeline parallelism since 52 layers is not divisible by 8 GPUs +- `--max_width_pruning 0.30` (default: 0.40) — tighter constraint to prevent head_dim≤48 and hidden=2048 dead zones +- `--max_depth_pruning 0.15` (default: 0.20) — tighter constraint since candidates with 42–46 layers universally fail for this model +- `--hparams_to_skip num_attention_heads` (default: none) — attention heads pruning is harder to recover, hence skipped **NOTE**: The tighter search space constraints here (`--max_depth_pruning`, `--max_width_pruning`) are specific to Nemotron hybrid models (Mamba + Attention + MoE). Unlike standard transformers which expose only layers/hidden/attention/FFN dimensions, these models add Mamba-specific dimensions (`mamba_num_heads`, `mamba_head_dim`) and MoE dimensions (`num_moe_experts`, `moe_ffn_hidden_size`, `moe_shared_expert_intermediate_size`), making the combined search space much larger. The default 40%/20% bounds cast too wide a net and waste compute on dead-zone architectures. @@ -167,14 +184,14 @@ See [ABLATIONS.md](ABLATIONS.md#pruning) for the full architecture search analys
-Pruning logs (top 20 candidates, best subnet, layer patterns) (click to expand) +Pruning logs (top 10 candidates, best subnet, layer patterns) (click to expand) ```text -╭──────────────────────────────────────────────────── Original Model Stats ─────────────────────────────────────────────────────╮ -│ Total Parameters 31.58B │ -│ Active Parameters 3.58B │ -│ Memory (BF16, seq_length=8192, batch_size=1) weights: 60230.1 MB, kv_cache: 48.0 MB, mamba_state: 23.8 MB, Total: 60301.9 MB │ -╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +╭───────────────────────────────────────────────────── Original Model Stats ──────────────────────────────────────────────────────╮ +│ Total Parameters 31.58B │ +│ Active Parameters 3.58B │ +│ Memory (BF16, seq_length=8192, batch_size=8) weights: 60230.1 MB, kv_cache: 384.0 MB, mamba_state: 190.5 MB, Total: 60804.6 MB │ +╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ Search Space (≤30% width / ≤15% depth pruning) @@ -192,50 +209,30 @@ See [ABLATIONS.md](ABLATIONS.md#pruning) for the full architecture search analys │ Search space size │ 10800 │ └─────────────────────────────────────┴────────────────────────────────┘ -Top 20 Candidates with Scores + Top 10 Candidates with Scores ┏━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━┓ ┃ # ┃ export_config ┃ active_params ┃ params ┃ score ┃ ┡━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━┩ -│ 1 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 64, 'num_moe_experts': 120, │ 3.00B │ 27.06B │ 0.3399 │ -│ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 2 │ {'num_layers': 48, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 112, │ 3.00B │ 25.37B │ 0.4650 │ -│ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 3 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 64, 'mamba_head_dim': 56, 'num_moe_experts': 112, │ 3.00B │ 25.37B │ 0.2343 │ -│ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 4 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 56, 'mamba_head_dim': 48, 'num_moe_experts': 96, │ 3.00B │ 20.09B │ 0.2552 │ +│ 1 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 56, 'mamba_head_dim': 48, 'num_moe_experts': 96, │ 3.00B │ 20.09B │ 0.2811 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 5 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 21.61B │ 0.2601 │ +│ 2 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 21.61B │ 0.2622 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 6 │ {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 64, 'num_moe_experts': 96, │ 3.00B │ 19.28B │ 0.3762 │ +│ 3 │ {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 64, 'num_moe_experts': 96, │ 3.00B │ 19.28B │ 0.4098 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3712} │ │ │ │ -│ 7 │ {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 104, │ 3.00B │ 22.28B │ 0.4783 │ +│ 4 │ {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 104, │ 3.00B │ 22.28B │ 0.4993 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 8 │ {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 96, │ 3.00B │ 21.99B │ 0.2420 │ +│ 5 │ {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 96, │ 3.00B │ 21.99B │ 0.2559 │ │ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3328} │ │ │ │ -│ 9 │ {'num_layers': 50, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 112, │ 3.00B │ 25.37B │ 0.2399 │ -│ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3712} │ │ │ │ -│ 10 │ {'num_layers': 50, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 112, │ 3.00B │ 26.17B │ 0.2601 │ -│ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3328} │ │ │ │ -│ 11 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 64, 'num_moe_experts': 112, │ 3.00B │ 25.37B │ 0.2503 │ +│ 6 │ {'num_layers': 48, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.4566 │ │ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 12 │ {'num_layers': 48, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.4329 │ +│ 7 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 64, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.2371 │ │ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 13 │ {'num_layers': 46, 'hidden_size': 2688, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 128, │ 3.00B │ 26.17B │ 0.2587 │ -│ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 2816} │ │ │ │ -│ 14 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 64, 'mamba_head_dim': 56, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.2336 │ -│ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 15 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'mamba_head_dim': 56, 'num_moe_experts': 96, │ 3.00B │ 20.09B │ 0.2559 │ +│ 8 │ {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'mamba_head_dim': 56, 'num_moe_experts': 96, │ 3.00B │ 20.09B │ 0.2601 │ │ │ 'moe_ffn_hidden_size': 1536, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 16 │ {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 96, │ 3.00B │ 20.70B │ 0.4608 │ +│ 9 │ {'num_layers': 52, 'hidden_size': 2304, 'mamba_num_heads': 64, 'mamba_head_dim': 64, 'num_moe_experts': 96, │ 3.00B │ 20.70B │ 0.4734 │ │ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ -│ 17 │ {'num_layers': 50, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.2455 │ +│ 10 │ {'num_layers': 50, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.2699 │ │ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3712} │ │ │ │ -│ 18 │ {'num_layers': 50, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 104, │ 3.00B │ 24.42B │ 0.2503 │ -│ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3328} │ │ │ │ -│ 19 │ {'num_layers': 48, 'hidden_size': 2560, 'mamba_num_heads': 48, 'mamba_head_dim': 48, 'num_moe_experts': 120, │ 3.00B │ 27.92B │ 0.2587 │ -│ │ 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3712} │ │ │ │ -│ 20 │ {'num_layers': 46, 'hidden_size': 2560, 'mamba_num_heads': 56, 'mamba_head_dim': 64, 'num_moe_experts': 104, │ 3.00B │ 23.68B │ 0.2469 │ -│ │ 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3072} │ │ │ │ └────┴───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┴───────────────┴────────┴────────┘ ╭──────────────────────────────────────────────────────────────────────── Best Subnet ─────────────────────────────────────────────────────────────────────────╮ @@ -243,17 +240,17 @@ Top 20 Candidates with Scores │ 'moe_shared_expert_intermediate_size': 3072} │ │ active_params 3.00B │ │ params 22.28B │ -│ score 0.4783 │ +│ score 0.4993 │ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ Original hybrid_layer_pattern: MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME Pruned hybrid_layer_pattern: MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME -╭───────────────────────────────────────────────────── Pruned Model Stats ──────────────────────────────────────────────────────╮ -│ Total Parameters 22.28B │ -│ Active Parameters 3.00B │ -│ Memory (BF16, seq_length=8192, batch_size=1) weights: 42489.7 MB, kv_cache: 48.0 MB, mamba_state: 23.8 MB, Total: 42561.6 MB │ -╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ +╭────────────────────────────────────────────────────── Pruned Model Stats ───────────────────────────────────────────────────────╮ +│ Total Parameters 22.28B │ +│ Active Parameters 3.00B │ +│ Memory (BF16, seq_length=8192, batch_size=8) weights: 42489.7 MB, kv_cache: 384.0 MB, mamba_state: 190.5 MB, Total: 43064.2 MB │ +╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ```
@@ -328,7 +325,8 @@ Phase 2 starts as a separate run from a fresh HuggingFace student checkpoint, so python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \ --hf-model /path/to/Nemotron-3-Nano-30B-A3B-Pruned-A3.0B \ --megatron-path /path/to/distill_output_phase1_8k/checkpoints/iter_0003200 \ - --hf-path /path/to/distill_output_phase1_8k/checkpoints/hf_iter_0003200 + --hf-path /path/to/distill_output_phase1_8k/checkpoints/hf_iter_0003200 \ + --trust-remote-code ``` @@ -385,12 +383,42 @@ For multi-node Slurm runs, see the [Megatron-Bridge README](../../README.md#slur > [!NOTE] > This is pure SFT-style distillation — no RL or online reward signal is used. Adding an RL-based post-training step after distillation is a natural next step that could further improve some of these benchmarks. +#### 3d. Convert Phase 2 final checkpoint to HuggingFace format + +We use the same conversion script to convert the Phase 2 final checkpoint to HuggingFace format. + +
+Checkpoint conversion command (click to expand) + +> NOTE: Below command only works for non-quantized checkpoints. For quantized checkpoints, we use the `export.py` script in Section 5 to directly export the quantized checkpoint to Unified HF format for deployment. + +```bash +python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \ + --hf-model /path/to/Nemotron-3-Nano-30B-A3B-Pruned-A3.0B \ + --megatron-path /path/to/distill_output_phase2_32k/checkpoints/iter_0000800 \ + --hf-path /path/to/distill_output_phase2_32k/checkpoints/hf_iter_0000800 \ + --trust-remote-code +``` + +
+ --- ### 4. Evaluation The eval config in [nemo_evaluator.yaml](nemo_evaluator.yaml) is for Slurm-based evaluation — it submits a vLLM serving job (with tool calling enabled via `--enable-auto-tool-choice --tool-call-parser qwen3_coder`) and runs evals against it. For local model execution and evaluation, refer to the [NeMo Evaluator documentation](https://docs.nvidia.com/nemo/evaluator/latest/) or this [blog](https://huggingface.co/blog/nvidia/nemotron-3-nano-evaluation-recipe). +**Tasks and exact metric names reported in the results table:** + +| Benchmark | Library | num_repeats | Metric name | +| --- | --- | --- | --- | +| MMLU Pro | NeMo Evaluator | 1 | `mmlu-pro_pass_at_1_symbolic_correct` | +| GPQA Diamond | NeMo Evaluator | 8 | `gpqa_pass_at_1_avg-of-8_symbolic_correct` | +| LiveCodeBench v6 | NeMo Evaluator | 4 | `livecodebench_pass_at_1_avg-of-4_accuracy` | +| AIME 2025 | NeMo Evaluator | 32 | `aime25_pass_at_1_avg-of-32_symbolic_correct` | +| IFBench | NeMo Evaluator | 8 | `ifbench_pass_at_1_avg-of-8_average_score` | +| SciCode (Subtask) | NeMo Evaluator | 8 | `scicode_pass_at_1_avg-of-8_subtask_accuracy` | +
Evaluation launch steps (click to expand) @@ -400,9 +428,9 @@ Before running, update the following fields in the `nemo_evaluator.yaml` file or - `execution.account` — your Slurm account - `deployment.checkpoint_path` — Hugging Face checkpoint path (original, pruned, or quantized) -The yaml is set up for a **BF16** checkpoint. For **FP8** checkpoints, also apply the quantization-specific vLLM deployment settings documented at the top of `nemo_evaluator.yaml`: +The yaml is set up for a **BF16** checkpoint. For **FP8** or **NVFP4** checkpoints, also apply the quantization-specific vLLM deployment settings documented at the top of `nemo_evaluator.yaml`: - append `--kv-cache-dtype fp8` to `deployment.extra_args` -- set the matching FlashInfer MoE env vars in `deployment.env_vars` (`VLLM_USE_FLASHINFER_MOE_FP8` plus `VLLM_FLASHINFER_MOE_BACKEND: throughput`) +- set the matching FlashInfer MoE env vars in `deployment.env_vars` (`VLLM_USE_FLASHINFER_MOE_FP8` for FP8 / `VLLM_USE_FLASHINFER_MOE_FP4` for NVFP4, plus `VLLM_FLASHINFER_MOE_BACKEND: throughput`) ```bash pip install "nemo-evaluator-launcher[all]==0.1.82" @@ -427,17 +455,6 @@ nemo-evaluator-launcher run --config nemo_evaluator.yaml
-**Tasks and exact metric names reported in the results table:** - -| Benchmark | Library | num_repeats | Metric name | -| --- | --- | --- | --- | -| MMLU Pro | NeMo Evaluator | 1 | `mmlu-pro_pass_at_1_symbolic_correct` | -| GPQA Diamond | NeMo Evaluator | 8 | `gpqa_pass_at_1_avg-of-8_symbolic_correct` | -| LiveCodeBench v6 | NeMo Evaluator | 4 | `livecodebench_pass_at_1_avg-of-4_accuracy` | -| AIME 2025 | NeMo Evaluator | 32 | `aime25_pass_at_1_avg-of-32_symbolic_correct` | -| IFBench | NeMo Evaluator | 8 | `ifbench_pass_at_1_avg-of-8_average_score` | -| SciCode (Subtask) | NeMo Evaluator | 8 | `scicode_pass_at_1_avg-of-8_subtask_accuracy` | - For more details on NeMo Evaluator, see the [GitHub repo](https://github.com/NVIDIA-NeMo/evaluator) and [documentation](https://docs.nvidia.com/nemo/evaluator/latest/). --- @@ -451,7 +468,7 @@ Similar to the official [Nemotron-3-Nano-30B-A3B-FP8](https://huggingface.co/nvi This is done with the `MAMBA_MOE_FP8_CONSERVATIVE_CFG` config defined in [`modelopt/torch/quantization/config.py`](../../../../modelopt/torch/quantization/config.py), which you select by passing `--quant_cfg MAMBA_MOE_FP8_CONSERVATIVE_CFG` below. For a faster model at the cost of a larger accuracy drop, you can use `MAMBA_MOE_FP8_AGGRESSIVE_CFG` instead. > [!NOTE] -> You can also quantize to NVFP4 using `--quant_cfg MAMBA_MOE_NVFP4_CONSERVATIVE_CFG` or `MAMBA_MOE_NVFP4_AGGRESSIVE_CFG` (faster, more accuracy drop), which may require further distillation (QAD) to recover accuracy and a Blackwell GPU for deployment. +> You can also quantize to NVFP4 using `--quant_cfg MAMBA_MOE_NVFP4_CONSERVATIVE_CFG` or `MAMBA_MOE_NVFP4_AGGRESSIVE_CFG` (faster, more accuracy drop). NVFP4 typically needs further [Quantization Aware Distillation (QAD)](../../README.md#quantization-aware-distillation-qad) to recover accuracy, plus a Blackwell GPU for deployment. Quantization is a two-step flow: `quantize.py` calibrates and saves a Megatron checkpoint, then `export.py` converts it to a deployable HuggingFace checkpoint (the unified HF exporter loads at TP=1, so pipeline parallelism is used to shard across GPUs). Both steps take a few minutes on 8x H100. @@ -466,8 +483,8 @@ torchrun --nproc_per_node 8 /opt/Model-Optimizer/examples/megatron_bridge/quanti --trust_remote_code \ --tp_size 8 \ --quant_cfg MAMBA_MOE_FP8_CONSERVATIVE_CFG \ - --calib_batch_size 32 \ - --seq_length 512 \ + --calib_batch_size 4 \ + --seq_length 8192 \ --export_megatron_path /path/to/distill_output_phase2_32k/checkpoints/iter_0000800_fp8_megatron \ --skip_generate ``` @@ -502,7 +519,7 @@ The exported HuggingFace checkpoint is directly deployable with [vLLM](https://g > ``` > [!TIP] -> You can run the evaluation using the same `nemo_evaluator.yaml` file for the quantized checkpoint also — just apply the FP8 deployment tweaks documented at the top of the yaml. +> You can run the evaluation using the same `nemo_evaluator.yaml` file for the quantized checkpoint also — just apply the FP8/NVFP4 deployment tweaks documented at the top of the yaml. To evaluate an NVFP4 checkpoint, you need a Blackwell GPU. See FP8 vs BF16 results in the [Results](#results) section above. diff --git a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/figures/learning_curves.png b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/figures/learning_curves.png index dbf4f359bc6..15a7ff785c7 100644 Binary files a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/figures/learning_curves.png and b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/figures/learning_curves.png differ diff --git a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/nemo_evaluator.yaml b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/nemo_evaluator.yaml index b2f63f17e86..5ea1def6879 100644 --- a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/nemo_evaluator.yaml +++ b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/nemo_evaluator.yaml @@ -5,10 +5,14 @@ # - `execution.account` — your Slurm account # - `deployment.checkpoint_path` — Hugging Face checkpoint path (original, pruned, or quantized) # -# This config is set up for a BF16 checkpoint. For an FP8 checkpoint, also apply the deployment changes below: -# - FP8: add `--kv-cache-dtype fp8` to `deployment.extra_args`, and set in `deployment.env_vars`: -# VLLM_USE_FLASHINFER_MOE_FP8: "1" -# VLLM_FLASHINFER_MOE_BACKEND: throughput +# This config is set up for a BF16 checkpoint. For quantized checkpoints, also apply the deployment changes below: +# - FP8: add `--kv-cache-dtype fp8` to `deployment.extra_args`, and set in `deployment.env_vars`: +# VLLM_USE_FLASHINFER_MOE_FP8: "1" +# VLLM_FLASHINFER_MOE_BACKEND: throughput +# - NVFP4: add `--kv-cache-dtype fp8` to `deployment.extra_args`, and set in `deployment.env_vars`: +# VLLM_USE_FLASHINFER_MOE_FP4: "1" +# VLLM_FLASHINFER_MOE_BACKEND: throughput +# (NVFP4 deployment requires a Blackwell GPU) # # Usage: # pip install "nemo-evaluator-launcher[all]==0.1.82" @@ -65,14 +69,17 @@ deployment: pipeline_parallel_size: 1 data_parallel_size: 8 gpu_memory_utilization: 0.90 - # extra_args is for the BF16 checkpoint. For an FP8 checkpoint, append `--kv-cache-dtype fp8`. + # extra_args is for the BF16 checkpoint. For FP8/NVFP4 checkpoints, append `--kv-cache-dtype fp8`. extra_args: "--max-model-len 262144 --max-num-seqs 8 --enable-log-requests --no-enable-prefix-caching --trust-remote-code --mamba_ssm_cache_dtype float32\ \ --enable-auto-tool-choice --tool-call-parser qwen3_coder --reasoning-parser-plugin /checkpoint/nano_v3_reasoning_parser.py --reasoning-parser nano_v3" - # env_vars is for the BF16 checkpoint (no MoE backend flags needed). For an FP8 - # checkpoint, replace `{}` with the block below (see notes at top of file): + # env_vars is for the BF16 checkpoint (no MoE backend flags needed). For a quantized + # checkpoint, replace `{}` with the block matching your format (see notes at top of file): # FP8: # VLLM_USE_FLASHINFER_MOE_FP8: "1" # VLLM_FLASHINFER_MOE_BACKEND: throughput + # NVFP4: + # VLLM_USE_FLASHINFER_MOE_FP4: "1" + # VLLM_FLASHINFER_MOE_BACKEND: throughput env_vars: {} endpoints: chat: /v1/chat/completions @@ -87,6 +94,10 @@ evaluation: adapter_config: use_system_prompt: true use_reasoning: false + # vLLM rejects both max_tokens and max_completion_tokens for these tasks + params_to_remove: + - max_tokens + - max_completion_tokens params_to_add: chat_template_kwargs: enable_thinking: true @@ -102,8 +113,8 @@ evaluation: params: parallelism: 64 max_new_tokens: 131072 - temperature: 0.99999 - top_p: 0.99999 + temperature: 1.0 + top_p: 1.0 request_timeout: 3600 max_retries: 10 extra: @@ -119,6 +130,7 @@ evaluation: OPENAI_CLIENT_ID: OPENAI_CLIENT_ID OPENAI_CLIENT_SECRET: OPENAI_CLIENT_SECRET + # For no-tools results: in ns_gpqa and ns_aime2025 below, remove `use_sandbox: true` and the `++tool_modules=[...]` token from `args` tasks: # 1. MMLU Pro - name: ns_mmlu_pro @@ -129,7 +141,7 @@ evaluation: params: extra: num_repeats: 1 - args: "++prompt_config=eval/aai/mcq-10choices-boxed" + args: "++prompt_config=eval/aai/mcq-10choices-boxed ++inference.tokens_to_generate=null" # 2. GPQA Diamond - name: ns_gpqa @@ -140,7 +152,9 @@ evaluation: params: extra: num_repeats: 8 - args: "++prompt_config=eval/aai/mcq-4choices" + # Tool calling: run the Python sandbox tool during generation + use_sandbox: true + args: "++prompt_config=eval/aai/mcq-4choices ++inference.tokens_to_generate=null ++tool_modules=[nemo_skills.mcp.servers.python_tool::PythonTool]" # 3. LiveCodeBench - name: ns_livecodebench @@ -162,6 +176,9 @@ evaluation: params: extra: num_repeats: 32 + # Tool calling: run the Python sandbox tool during generation + use_sandbox: true + args: "++prompt_config=/nemo_run/code/eval_factory_prompts/math-oai.yaml ++inference.tokens_to_generate=null ++tool_modules=[nemo_skills.mcp.servers.python_tool::PythonTool]" # 5. IFBench - name: ns_ifbench @@ -182,3 +199,4 @@ evaluation: params: extra: num_repeats: 8 + args: "++eval_config.num_parallel_requests=20" diff --git a/examples/pruning/README.md b/examples/pruning/README.md index 025c61a6712..db6f4fa1c23 100644 --- a/examples/pruning/README.md +++ b/examples/pruning/README.md @@ -12,15 +12,15 @@ This section focuses on applying Model Optimizer's state-of-the-art complementar
-| **Section** | **Description** | **Link** | **Docs** | -| :------------: | :------------: | :------------: | :------------: | -| Pre-Requisites | Required & optional packages to use this technique | \[[Link](#pre-requisites)\] | | -| Getting Started | Learn how to use the pruning API | \[[Link](#getting-started)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/3_pruning.html)\] | -| Support Matrix | View the support matrix to see available pruning algorithms and their compatibility with different models and frameworks | \[[Link](#support-matrix)\] | | -| Examples | Examples of different pruning methods | \[[Link](#examples)\] | | -| Pruning Guidelines | Guidelines for choosing how and how much to prune for best results | \[[Link](#pruning-guidelines)\] | | -| Tutorials / Results | End-to-end tutorials for Minitron and Puzzletron pruning | \[[Link](#tutorials--results)\] | | -| Resources | Extra links to relevant resources | \[[Link](#resources)\] | | +| **Section** | **Description** | **Link** | +| :------------: | :------------: | :------------: | +| Pre-Requisites | Required & optional packages to use this technique | \[[Link](#pre-requisites)\] | +| Getting Started | Learn how to use the pruning API | \[[Link](#getting-started)\] | +| Support Matrix | View the support matrix to see available pruning algorithms and their compatibility with different models and frameworks | \[[Link](#support-matrix)\] | +| Examples | Examples of different pruning methods | \[[Link](#examples)\] | +| Pruning Guidelines | Guidelines for choosing how and how much to prune for best results | \[[Link](#pruning-guidelines)\] | +| Tutorials / Results | End-to-end tutorials for Minitron and Puzzletron pruning | \[[Link](#tutorials--results)\] | +| Resources | Extra links to relevant resources | \[[Link](#resources)\] |
@@ -174,10 +174,17 @@ If your model parameters are already sorted and you just want to prune the weigh | **Algorithm** | **Model** | **Pruning Constraints** | | :---: | :---: | :---: | -| Minitron | Megatron-core (M-LM, M-Bridge) based GPT / Mamba / MoE / Hybrid LLM Models1 | **Manual:** `export_config` with width (`hidden_size`, `ffn_hidden_size`, `num_attention_heads`, `mamba_num_heads`, `mamba_head_dim`, `num_moe_experts`, `moe_ffn_hidden_size`, `moe_shared_expert_intermediate_size`) and/or depth (`num_layers`) pruned values
**Auto:** one or more of `params`, `active_params`, `memory_mb` (requires `score_func` in config) | +| Minitron | Megatron-core1 (M-Bridge, M-LM) based dense / MoE / hybrid Mamba-Transformer LLMs2 | **Auto:** one or more of `params`, `active_params`, `memory_mb`
**Manual:** `export_config` with width (`hidden_size`, `ffn_hidden_size`, `num_attention_heads`, `mamba_num_heads`, `mamba_head_dim`, `num_moe_experts`, `moe_ffn_hidden_size`, `moe_shared_expert_intermediate_size`) and/or depth (`num_layers`) pruned values | +| Puzzletron | Hugging Face based dense / MoE / hybrid Mamba-Transformer LLMs & VLMs3 | **Target:** one or more of `target_memory`, `num_params`, `target_latency_seconds`
**Heterogeneous (per-layer) search dimensions:**4 FFN `intermediate_size` (different sizes per layer), attention `op`/`no_op` (selective attention-layer removal) and KV heads (GQA grouping), `hidden_size`, and MoE `num_experts` (expert removal) | | FastNAS | Computer Vision models | `flops`, `params` | -> *1.Only models in Pipeline Parallelism (PP) are supported. Hugging Face models can be imported into M-Bridge/M-LM format as long as they are [supported](https://docs.nvidia.com/nemo/megatron-bridge/latest/index.html#supported-models) by the framework.* +> *1.Hugging Face models can be imported into M-Bridge/M-LM format as long as they are [supported](https://docs.nvidia.com/nemo/megatron-bridge/latest/index.html#supported-models) by the framework.* + +> *2.Language model part of VLMs can be pruned as well.* + +> *3.Puzzletron operates on Hugging Face checkpoints via its `AnyModel` abstraction. New architectures can be added by writing a model descriptor + converter — see the [AnyModel Guide](../../modelopt/torch/puzzletron/anymodel/README.md). Available configs are in the Puzzletron [configs](../puzzletron/configs/) directory.* + +> *4.The MIP search produces a heterogeneous architecture (dimensions can differ per layer). Which dimensions are searched is model- and config-dependent.* ## Examples diff --git a/modelopt/torch/puzzletron/plugins/mbridge/base.py b/modelopt/torch/puzzletron/plugins/mbridge/base.py index 688a32fbb71..39c7050738e 100644 --- a/modelopt/torch/puzzletron/plugins/mbridge/base.py +++ b/modelopt/torch/puzzletron/plugins/mbridge/base.py @@ -32,7 +32,6 @@ HeterogeneousTransformerConfig, TransformerConfig, ) -from megatron.bridge.utils.instantiate_utils import register_allowed_target_prefix from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import ( get_gpt_heterogeneous_layer_spec, ) @@ -43,9 +42,14 @@ if not hasattr(TransformerConfig, "get_config_for_layer"): TransformerConfig.get_config_for_layer = lambda self, layer_number: self -__all__ = ["heterogeneous_layer_spec", "GenericHeterogeneousProvider", "HeterogeneousBridgeMixin"] +try: # nemo:26.04.01 onwards needs this fix + from megatron.bridge.utils.instantiate_utils import register_allowed_target_prefix + + register_allowed_target_prefix("modelopt.") +except ImportError: + pass -register_allowed_target_prefix("modelopt.") +__all__ = ["heterogeneous_layer_spec", "GenericHeterogeneousProvider", "HeterogeneousBridgeMixin"] def heterogeneous_layer_spec(config) -> ModuleSpec: diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/nemo_evaluator.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/nemo_evaluator.yaml new file mode 120000 index 00000000000..53e53a6540e --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/nemo_evaluator.yaml @@ -0,0 +1 @@ +../../../../../examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/nemo_evaluator.yaml \ No newline at end of file