diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 494dcd724e2..6646359c7c3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,6 +25,7 @@ Changelog **New Features** - Support full Transformer Engine spec for Minitron pruning (``mcore_minitron``). Now we no longer need to use custom ModelOpt spec. Note that this does not affect the usage of the pruning workflow but makes pruning slightly faster and may result in slightly different pruned model because of different kernel and numerics. +- Add end-to-end tutorial for Minitron pruning + distillation + quantization + evaluation + vLLM deployment for Nemotron-Nano-9B-v2 → Pruned 7B along with data blend preparation steps (and ablation study). See `examples/pruning/minitron/README.md `_ for details. - Add Puzzletron - a new algorithm for heterogeneous pruning of LLM and VLM models. See `examples/puzzletron/README.md `_ for more details. - Added iterator interface using CalibrationDataReader in ONNX quantization workflow. - Add N:M sparse softmax support to the Triton flash attention kernel (``modelopt.torch.kernels.common.attention.triton_fa``). See `examples/llm_sparsity/attention_sparsity/README.md `_ for usage. diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md new file mode 100644 index 00000000000..c3904d2a0fd --- /dev/null +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -0,0 +1,242 @@ +# Tokenizing for Megatron Frameworks + +| **Section** | **Description** | **Link** | +| :---: | :---: | :---: | +| From JSONL files | Tokenize local JSONL files | \[[Link](#from-jsonl-files)\] | +| From Hugging Face Hub | Stream or download HF datasets and tokenize | \[[Link](#from-hugging-face-hub)\] | +| `reasoning_content` for Post-Training v3 | Control how chain-of-thought traces are handled | \[[Link](#reasoning_content-for-post-training-v3-datasets)\] | +| Nemotron Pre/Post-Training Datasets | Ready-to-run commands for all Nemotron datasets | \[[Link](#ready-to-run-tokenization-commands)\] | + +The distillation and pre-training scripts in Megatron-Bridge or Megatron-LM expect data pre-tokenized in Megatron's binary indexed format (`.bin` / `.idx`). +Use the `megatron_preprocess_data` utility to tokenize any JSONL or Hugging Face dataset. +The tokenization scripts below print the list of output prefixes (e.g. `tokenized_qwen3/data1_text`) that you can use for the `data_paths` argument (with relative weights on different files) in Megatron training scripts. + +**Important Notes:** + +- For Pretraining / raw-text data (`text` key) — use `--append_eod` so Megatron can tell where documents end when concatenating them into long sequences. +- For Post-training chat data (`messages` key) — omit `--append_eod`; the chat template already appends EOS at the end of each conversation. +- Set `--max_sequence_length 256_000` to avoid rare OOM errors if some text is very long. + +## From JSONL files + +```bash +python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --jsonl_paths /path/to/data1.jsonl /path/to/data2.jsonl ... \ + --json_keys text \ + --tokenizer Qwen/Qwen3-0.6B \ + --output_dir tokenized_qwen3 \ + --workers 32 \ + --append_eod +``` + +```bash +python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --jsonl_paths /path/to/sft_data.jsonl \ + --json_keys messages \ + --tokenizer Qwen/Qwen3-0.6B \ + --output_dir tokenized_qwen3 \ + --workers 32 +``` + +Instead of `--jsonl_paths`, pass `--input_dir /path/to/dir` to tokenize all JSONL files in a directory (`.jsonl` and `.jsonl.gz` are both supported). + +## From Hugging Face Hub + +To tokenize a dataset directly from Hugging Face Hub: + +```bash +python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --hf_dataset nvidia/Nemotron-Pretraining-SFT-v1 \ + --hf_name Nemotron-SFT-Code \ + --hf_split train \ + --hf_max_samples_per_split 10_000_000 \ + --json_keys text \ + --tokenizer Qwen/Qwen3-0.6B \ + --output_dir tokenized_qwen3 \ + --workers 32 \ + --append_eod +``` + +Omit `--hf_name` to process all subsets, `--hf_split` for all splits, or `--hf_max_samples_per_split` for all samples. +To quickly test, use [nvidia/Nemotron-Pretraining-Dataset-sample](https://huggingface.co/datasets/nvidia/Nemotron-Pretraining-Dataset-sample). + +For very large datasets (tens of millions of documents), or datasets with complex nested message schemas (e.g. `tool_calls`, `function_call` fields) that cause Arrow type-cast errors in non-streaming mode, add `--hf_streaming` to avoid downloading the full dataset — only the rows actually consumed are fetched. Optionally pair with `--hf_max_samples_per_split ` to cap the row count; without it streaming still works but re-downloads on every run with no disk cache. + +> **Performance note:** Non-streaming mode downloads all Parquet shards once and caches them as Arrow files on disk. +> Re-runs read from cache and are much faster. +> Streaming re-downloads on every run with no cache, so it is slower for full-dataset processing. + +## `reasoning_content` for Post-Training v3 Datasets + +v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from +the final answer). Use `--reasoning_content` to control how it is handled: + +| Value | Behaviour | +| --- | --- | +| `strip` (default) | Field is discarded before `apply_chat_template`. Safe for any tokenizer. | +| `inline` | Wrapped as `…` and prepended to `content`. Preserves reasoning in a tokenizer-agnostic way. | +| `native` | Passed unchanged. Requires the tokenizer's chat template to handle the field (e.g. Qwen3). | + +```bash +python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --hf_dataset nvidia/Nemotron-Math-v2 \ + --hf_split high_part00 \ + --json_keys messages \ + --tokenizer nvidia/NVIDIA-Nemotron-Nano-9B-v2 \ + --output_dir tokenized_nemotron_v2 \ + --workers 32 \ + --reasoning_content inline +``` + +--- + +## Ready-to-run tokenization commands + +Tokenization commands for all Nemotron Pre-Training and Post-Training datasets used in Megatron-Bridge distillation experiments. + +Two parameters vary by model — set them before running the commands below: + +```bash +TOKENIZER=nvidia/NVIDIA-Nemotron-Nano-9B-v2 # HuggingFace tokenizer (or local path) +OUTPUT_DIR=tokenized_nemotron_v2 # Output directory for tokenized files +``` + +> [!TIP] +> Token count for a `.bin` file = file size in bytes ÷ 4. This is also printed by the tokenization script on completion. + +> [!NOTE] +> Tokenizing each of the datasets below will take anywhere between 10 minutes to few hours. You can tokenize all in parallel to speed up the process. +> +> You may tokenize more datasets or skip some datasets depending on your needs. + +### Nemotron Pretraining dataset + +**[nvidia/Nemotron-Pretraining-SFT-v1](https://huggingface.co/datasets/nvidia/Nemotron-Pretraining-SFT-v1)** — raw text; omitting `--hf_name` tokenizes all 3 subsets (Code, General, MATH) in one command, producing a separate output file per subset named after each: + +```bash +python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --hf_dataset nvidia/Nemotron-Pretraining-SFT-v1 \ + --hf_split train \ + --hf_streaming \ + --hf_max_samples_per_split 10_000_000 \ + --json_keys text \ + --tokenizer ${TOKENIZER} \ + --output_dir ${OUTPUT_DIR} \ + --workers 96 \ + --max_sequence_length 256_000 \ + --append_eod \ + --strip_newlines +``` + +--- + +### Nemotron Post-training v1 dataset + +**[nvidia/Nemotron-Post-Training-Dataset-v1](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v1)** — STEM subset, capped at 5M samples. v1 data does not contain reasoning traces: + +```bash +python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --hf_dataset nvidia/Nemotron-Post-Training-Dataset-v1 \ + --hf_name default \ + --hf_split stem \ + --hf_streaming \ + --hf_max_samples_per_split 5_000_000 \ + --json_keys messages \ + --tokenizer ${TOKENIZER} \ + --output_dir ${OUTPUT_DIR} \ + --workers 96 \ + --max_sequence_length 256_000 +``` + +--- + +### Nemotron Post-training v3 collection + +Datasets below are from the [Nemotron Post-Training v3 collection](https://huggingface.co/collections/nvidia/nemotron-post-training-v3). All use `--reasoning_content inline` to preserve `…` traces. The collection contains many more datasets — if you care about benchmarks not covered here (e.g. multilingual, agentic/tool use, SWE, safety), pick the relevant datasets from the collection and tokenize them the same way. + +**[nvidia/Nemotron-Math-v2](https://huggingface.co/datasets/nvidia/Nemotron-Math-v2)** — tokenize `high_part00` and `high_part01` separately: + +```bash +for SPLIT in high_part00 high_part01; do + python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --hf_dataset nvidia/Nemotron-Math-v2 \ + --hf_split ${SPLIT} \ + --json_keys messages \ + --tokenizer ${TOKENIZER} \ + --output_dir ${OUTPUT_DIR} \ + --workers 96 \ + --max_sequence_length 256_000 \ + --reasoning_content inline +done +``` + +**[nvidia/Nemotron-SFT-Competitive-Programming-v2](https://huggingface.co/datasets/nvidia/Nemotron-SFT-Competitive-Programming-v2)** — stored as raw JSONL on HuggingFace, download before tokenizing: + +```bash +hf download nvidia/Nemotron-SFT-Competitive-Programming-v2 \ + --repo-type dataset \ + --local-dir datasets/Nemotron-SFT-Competitive-Programming-v2/ +for FILE in competitive_programming_python_00 competitive_programming_cpp_00; do + python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --jsonl_paths datasets/Nemotron-SFT-Competitive-Programming-v2/data/${FILE}.jsonl \ + --json_keys messages \ + --tokenizer ${TOKENIZER} \ + --output_dir ${OUTPUT_DIR} \ + --workers 96 \ + --max_sequence_length 256_000 \ + --reasoning_content inline +done +``` + +**[nvidia/Nemotron-Science-v1](https://huggingface.co/datasets/nvidia/Nemotron-Science-v1)** — stored as raw JSONL on HuggingFace, download before tokenizing: + +```bash +hf download nvidia/Nemotron-Science-v1 \ + --repo-type dataset \ + --local-dir datasets/Nemotron-Science-v1/ +python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --input_dir datasets/Nemotron-Science-v1/data/ \ + --json_keys messages \ + --tokenizer ${TOKENIZER} \ + --output_dir ${OUTPUT_DIR} \ + --workers 96 \ + --max_sequence_length 256_000 \ + --reasoning_content inline +``` + +**[nvidia/Nemotron-SFT-Instruction-Following-Chat-v2](https://huggingface.co/datasets/nvidia/Nemotron-SFT-Instruction-Following-Chat-v2)** — stored as raw JSONL on HuggingFace, download before tokenizing: + +```bash +hf download nvidia/Nemotron-SFT-Instruction-Following-Chat-v2 \ + --repo-type dataset \ + --local-dir datasets/Nemotron-SFT-Instruction-Following-Chat-v2/ +python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ + --input_dir datasets/Nemotron-SFT-Instruction-Following-Chat-v2/data/ \ + --json_keys messages \ + --tokenizer ${TOKENIZER} \ + --output_dir ${OUTPUT_DIR} \ + --workers 96 \ + --max_sequence_length 256_000 \ + --reasoning_content inline +``` + +--- + +### Expected output + +After running all commands above, `${OUTPUT_DIR}/` should contain the following `.bin` / `.idx` file pairs: + +```text +nvidia--Nemotron-Pretraining-SFT-v1_Nemotron-SFT-Code_train_text_max10000000.{bin,idx} +nvidia--Nemotron-Pretraining-SFT-v1_Nemotron-SFT-General_train_text_max10000000.{bin,idx} +nvidia--Nemotron-Pretraining-SFT-v1_Nemotron-SFT-MATH_train_text_max10000000.{bin,idx} +nvidia--Nemotron-Post-Training-Dataset-v1_default_stem_messages_max5000000.{bin,idx} +nvidia--Nemotron-Math-v2_default_high_part00_messages.{bin,idx} +nvidia--Nemotron-Math-v2_default_high_part01_messages.{bin,idx} +competitive_programming_python_00_messages.{bin,idx} +competitive_programming_cpp_00_messages.{bin,idx} +MCQ_messages.{bin,idx} +RQA_messages.{bin,idx} +reasoning_off_messages.{bin,idx} +reasoning_on_messages.{bin,idx} +``` diff --git a/examples/dataset/README.md b/examples/dataset/README.md index 15cb21613c2..d073237cf6a 100644 --- a/examples/dataset/README.md +++ b/examples/dataset/README.md @@ -5,7 +5,7 @@ | **Section** | **Description** | **Link** | | :------------: | :------------: | :------------: | | Building Chat Datasets | Scripts to build conversation datasets from Nemotron and other HuggingFace sources | \[[Link](#building-chat-datasets)\] | -| Tokenizing for Megatron Frameworks | Convert JSONL or HF datasets to Megatron binary format for distillation and pre-training | \[[Link](#tokenizing-for-megatron-frameworks)\] | +| Tokenizing for Megatron Frameworks | Convert JSONL or HF datasets to Megatron binary format for distillation and pre-training | \[[Link](MEGATRON_DATA_PREP.md)\] | @@ -140,85 +140,7 @@ In `generate` mode, assistant turns are stripped so the row ends with a user tur ## Tokenizing for Megatron Frameworks -The distillation and pre-training scripts in Megatron-Bridge or Megatron-LM expect data pre-tokenized in Megatron's binary indexed format (`.bin` / `.idx`). -Use the `megatron_preprocess_data` utility to tokenize any JSONL or Hugging Face dataset. -The tokenization scripts below prints the list of output prefixes (e.g. `tokenized_qwen3/data1_text`) that you can use for the `data_paths` argument (with relative weights on different files) in Megatron training scripts. - -**Important Notes:** - -- For Pretraining / raw-text data (`text` key) — use `--append_eod` so Megatron can tell where documents end when concatenating them into long sequences. -- For Post-training chat data (`messages` key) — omit `--append_eod`; the chat template already appends EOS at the end of each conversation. -- Set `--max_sequence_length 256_000` to avoid rare OOM errors if some text is very long. - -### From JSONL files - -```bash -python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ - --jsonl_paths /path/to/data1.jsonl /path/to/data2.jsonl ... \ - --json_keys text \ - --tokenizer Qwen/Qwen3-0.6B \ - --output_dir tokenized_qwen3 \ - --workers 32 \ - --append_eod -``` - -```bash -python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ - --jsonl_paths /path/to/sft_data.jsonl \ - --json_keys messages \ - --tokenizer Qwen/Qwen3-0.6B \ - --output_dir tokenized_qwen3 \ - --workers 32 -``` - -Instead of `--jsonl_paths`, pass `--input_dir /path/to/dir` to tokenize all JSONL files in a directory (`.jsonl` and `.jsonl.gz` are both supported). - -### From Hugging Face Hub - -To tokenize a dataset directly from Hugging Face Hub: - -```bash -python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ - --hf_dataset nvidia/Nemotron-Pretraining-SFT-v1 \ - --hf_name Nemotron-SFT-Code \ - --hf_split train \ - --hf_max_samples_per_split 10_000_000 \ - --json_keys text \ - --tokenizer Qwen/Qwen3-0.6B \ - --output_dir tokenized_qwen3 \ - --workers 32 \ - --append_eod -``` - -Omit `--hf_name` to process all subsets, `--hf_split` for all splits, or `--hf_max_samples_per_split` for all samples. -To quickly test, use [nvidia/Nemotron-Pretraining-Dataset-sample](https://huggingface.co/datasets/nvidia/Nemotron-Pretraining-Dataset-sample). - -For **very large datasets** (tens of millions of documents), add `--hf_streaming --hf_max_samples_per_split ` to avoid downloading the full dataset — only the rows actually consumed are fetched. - -> **Performance note:** Non-streaming mode downloads all Parquet shards once and caches them as Arrow files on disk. -> Re-runs read from cache and are much faster. -> Streaming re-downloads on every run with no cache, so it is slower for full-dataset processing. - -### Nemotron Post-Training v3 (`reasoning_content`) - -v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from -the final answer). Use `--reasoning_content` to control how it is handled: - -| Value | Behaviour | -| --- | --- | -| `strip` (default) | Field is discarded before `apply_chat_template`. Safe for any tokenizer. | -| `inline` | Wrapped as `…` and prepended to `content`. Preserves reasoning in a tokenizer-agnostic way. | -| `native` | Passed unchanged. Requires the tokenizer's chat template to handle the field (e.g. Qwen3). | - -```bash -python -m modelopt.torch.utils.plugins.megatron_preprocess_data \ - --hf_dataset nvidia/Nemotron-Post-Training-Dataset-v3 \ - --json_keys messages \ - --tokenizer Qwen/Qwen3-0.6B \ - --output_dir tokenized_qwen3 \ - --workers 32 \ - --reasoning_content inline -``` +See **[MEGATRON_DATA_PREP.md](MEGATRON_DATA_PREP.md)** for full documentation: general usage with JSONL and Hugging Face Hub datasets, handling of Nemotron Post-Training v3 `reasoning_content` fields, and ready-to-run tokenization commands for all Nemotron Pre/Post-Training datasets. ## Synthetic Test Dataset diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 571a0c49884..1e384acfb19 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -47,7 +47,7 @@ hf auth login --token ``` > [!WARNING] -> Use `python -m pip` instead of `pip` to avoid conflicts with the system-wide installed packages in the NeMo containers. +> Use `python -m pip` instead of `pip` to avoid conflicts with the system-wide installed packages in the NeMo containers. You may also refer to this [doc](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docker/common/README.md#installing-packages-inside-the-container) on how to correctly install packages in the NeMo containers without breaking existing torch installation. ## Pruning @@ -189,7 +189,7 @@ For more details, see the [Megatron-Bridge conversion README](https://github.com ### Distillation Results -See [results/puzzletron.md](results/puzzletron.md) for MMLU results demonstrating knowledge distillation on Puzzletron-compressed student models. +See [examples/pruning/](../pruning/README.md#tutorials--results) for distillation experiment results covering Minitron and Puzzletron pruning algorithms. ## Post-Training Quantization diff --git a/examples/pruning/README.md b/examples/pruning/README.md index 9e84622269c..294f00031dd 100644 --- a/examples/pruning/README.md +++ b/examples/pruning/README.md @@ -20,6 +20,7 @@ This section focuses on applying Model Optimizer's state-of-the-art complementar | 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)\] | | @@ -186,16 +187,28 @@ If your model parameters are already sorted and you just want to prune the weigh ## Examples -### Minitron Pruning for Megatron-Bridge/ Megatron-LM Framework LLMs (e.g. Qwen 3, Nemotron Nano) +### Minitron Pruning for Megatron-Bridge/ Megatron-LM Framework LLMs (e.g. Qwen3, Nemotron 3 Nano) Checkout the Minitron pruning example for [Megatron-Bridge Framework](../megatron_bridge/README.md#pruning) or [Megatron-LM Framework](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/post_training/modelopt#-pruning) which showcases the usage of the powerful Minitron pruning algorithm developed by NVIDIA Research for pruning LLMs like Llama-3.1-8B, Qwen3-8B, Nemotron-Nano-9B-v2, Nemotron-3-Nano-30B-A3B, etc. Both frameworks support importing from a Hugging Face pretrained checkpoint. -Some of the models pruned using Minitron method followed by distillation and post-training are: +Some of the official models pruned using Minitron method followed by distillation and post-training are: - [Minitron Collection on Hugging Face](https://huggingface.co/collections/nvidia/minitron) - [NVIDIA-Nemotron-Nano-9B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2) +See [minitron/](minitron/README.md) for end-to-end tutorials and results. + +### Puzzletron Pruning for LLMs (e.g. Llama, Qwen, Nemotron) + +Checkout the [Puzzletron README](../puzzletron/README.md) which showcases MIP-based NAS pruning that produces heterogeneous model architectures — varying FFN intermediate sizes per layer and selectively removing attention layers — to meet a target parameter count or memory budget. + +Supported models include Llama-3.1-8B-Instruct, Qwen3-8B, Qwen2.5-7B-Instruct, Nemotron-Nano-12B-v2, Mistral-Small-24B-Instruct-2501, and others via the [configs](../puzzletron/configs/) directory. See the [Puzzletron README](../puzzletron/README.md) for more details. + +After compression, use [Megatron-Bridge distillation](../megatron_bridge/README.md#distillation) to recover accuracy. + +See [puzzletron/](puzzletron/README.md) for distillation results on Puzzletron-compressed models. + ### FastNAS Pruning for PyTorch Computer Vision Models Check out the FastNAS pruning example usage in the [documentation](https://nvidia.github.io/Model-Optimizer/guides/3_pruning.html#pruning-and-subnet-search). @@ -279,16 +292,23 @@ After pruning, distillation is required to recover model accuracy. Below are rec | **Hyperparameter** | **Recommendation** | | :---: | :---: | | **Sequence Length** | 8192 (or 4096 if dataset has smaller sequences) | -| **Global Batch Size (GBS)** | 768 | +| **Global Batch Size (GBS)** | same as the original training or 768 if unsure | | **Micro Batch Size (MBS)** | As large as your GPU memory can accommodate | | **Learning Rate (LR)** | 1e-4 → 1e-5 (linear decay) for 30-50% pruning
• More compression → higher LR
• Less compression → lower LR
• As model gets larger → reduce LR to avoid divergence | | **Warmup Steps** | 100 | -| **Training Max Steps** | Num training tokens / (Seq len × GBS)
• Recommended: 80-100B tokens | +| **Training Max Steps** | Num training tokens / (Seq len × GBS)
• Recommended: 80-100B tokens for best results. | | **Data Composition** | • Standard models: 100% pre-training data
• Reasoning models: 70% reasoning data + 30% pre-training data | > [!TIP] > If you know the maximum learning rate used during the original training, a good rule of thumb for knowledge distillation is to use **1/5th of that maximum LR** when compressing by ~50%. +## Tutorials / Results + +End-to-end distillation results with Megatron-Bridge after Minitron and Puzzletron pruning: + +- **[Minitron — Nemotron-Nano-9B-v2](minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md)**: End-to-end tutorial of structured pruning for Nemotron-Nano-9B-v2 to 7B followed by knowledge distillation up to 80B tokens, quantization, and vLLM deployment. Achieves near-parity with the official 9B model across popular pretraining and reasoning benchmarks. +- **[Puzzletron — Qwen3-8B and Llama-3.1-8B-Instruct](puzzletron/Llama-3.1-8B-Instruct.md)**: MIP-based compression followed by short distillation runs on WikiText-103. Shows MMLU recovery and illustrates the importance of using larger datasets to avoid overfitting. + ## Resources - 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/146) diff --git a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/ABLATIONS.md b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/ABLATIONS.md new file mode 100644 index 00000000000..1786e88fdda --- /dev/null +++ b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/ABLATIONS.md @@ -0,0 +1,76 @@ +# Distillation Blend Ablations + +All experiments prune Nemotron-Nano-9B-v2 → 7B and distill with teacher = Nemotron-Nano-9B-v2 (official). The final chosen blend (**30pre_70post_v1v3**) is in [README.md](README.md). + +--- + +## Baseline: Pre-SFT-v1 Only (no post-training data) + +Pure Nemotron-Pretraining-SFT-v1 data only (no post-training reasoning traces). + +| Tokens | MMLU | MMLU Pro | GPQA Diamond | LCB v6 | AIME 2025 | Math 500 | IFEval | SciCode | +|---|---|---|---|---|---|---|---|---| +| 19B | 72.7 | 70.5 | 53.9 | 58.8 | 63.4 | 94.4 | 57.9 | 19.2 | +| 56B | 73.3 | 71.9 | 54.3 | 62.0 | 63.8 | 95.0 | 58.7 | 17.9 | + +**Notes:** Highest MMLU of any blend, but AIME stagnates and LCB lags. Pretraining data alone insufficient for reasoning benchmarks. + +--- + +## Baseline: Pure Post-Training Data (pt-v1v2) + +100% post-training data (no pretraining data), Nemotron-v1/v2 blend. + +| Tokens | MMLU | MMLU Pro | GPQA Diamond | LCB v6 | AIME 2025 | Math 500 | IFEval | SciCode | +|---|---|---|---|---|---|---|---|---| +| 2.5B | 71.0 | 69.3 | 52.6 | 54.8 | 58.2 | 94.1 | 51.7 | 14.4 | +| 5B | 70.8 | 70.7 | 53.6 | 57.2 | 63.8 | 94.1 | 50.5 | 14.2 | +| 20B | 69.8 | 71.7 | 54.7 | 57.5 | 64.7 | 94.6 | 41.9 | 13.4 | +| 40B | 70.0 | 71.7 | 53.2 | 57.4 | 67.6 | 95.2 | 43.3 | 16.2 | + +**Notes:** IFEval degrades badly at longer training (41.9 at 20B). LCB lags behind other blends. + +--- + +## 30% Pretraining / 70% Post-Training: v1v2 Blend + +30% Nemotron-Pretraining-SFT-v1 + 70% Nemotron-v1/v2 post-training data. + +| Tokens | MMLU | MMLU Pro | GPQA Diamond | LCB v6 | AIME 2025 | Math 500 | IFEval | SciCode | +|---|---|---|---|---|---|---|---|---| +| 2.5B | 71.9 | 68.9 | 49.8 | 56.4 | 55.3 | 93.3 | 58.2 | 14.6 | +| 5B | — | — | — | — | — | — | — | — | +| 20B | 71.6 | 71.2 | 52.7 | 58.0 | 65.1 | 94.0 | 55.7 | 14.2 | +| 40B | 72.7 | 71.1 | 54.0 | 59.7 | 65.5 | 95.2 | 53.8 | 19.2 | +| 60B | 73.0 | 71.9 | 55.9 | 60.0 | 67.8 | 95.4 | 56.4 | 21.7 | +| 80B | 73.4 | 72.7 | 54.7 | 61.8 | 70.7 | 95.3 | 57.8 | 19.9 | +| 100B | 73.5 | 72.8 | 56.4 | 62.4 | 71.9 | 95.8 | 59.1 | 19.4 | + +**Notes:** Best MMLU of the 30/70 blends (~1% above v3 blends). IFEval ~56–59 (lower than v3 blends). GPQA shows instability at longer runs. + +--- + +## 30% Pretraining / 70% Post-Training: v3 Blend + +Refined v3 blend: dropped exercism/text2sql, added Nemotron-Math-v2 part01, boosted Math to 30% total. + +| Tokens | MMLU | MMLU Pro | GPQA Diamond | LCB v6 | AIME 2025 | Math 500 | IFEval | SciCode | +|---|---|---|---|---|---|---|---|---| +| 2.5B | 70.5 | 69.0 | 51.2 | 59.1 | 62.9 | 94.3 | 62.2 | 11.6 | +| 5B | 71.0 | 69.8 | 53.0 | 59.4 | 65.0 | 94.4 | 66.8 | 20.3 | +| 20B | 71.2 | 70.8 | 53.3 | 60.0 | 69.1 | 95.3 | 63.8 | 22.6 | +| 40B | 71.0 | 71.7 | 54.0 | 62.3 | 71.3 | 95.3 | 66.8 | 17.9 | +| 60B | 72.0 | 72.3 | 56.3 | 62.0 | 71.6 | 95.6 | 65.5 | 21.5 | +| 80B | 72.3 | 73.0 | 53.9 | 63.0 | 72.4 | 96.2 | 65.5 | 21.3 | + +**Notes:** Better AIME and LCB than blend 1 at 40B+. GPQA still unstable (53.9 at 80B). MMLU ~1% below v1v2 blend. + +--- + +## Blend Design Notes + +**Why MMLU is ~1% lower with v3 blends:** The heavy reasoning-trace format (chain-of-thought, TIR) in v3 data suppresses general knowledge recall measured by MMLU. This is structural — v1v2 post-training data has a more knowledge-dense format. Upweighting Pretraining-SFT-v1 General (to 20%) partially mitigates this. Given that MMLU Pro is better with v3 blends, lower MMLU is acceptable. + +**Why GPQA is unstable in blend 1:** Science-v1 MCQ (497M tokens) and RQA (278M tokens) are repeated ~14× over 100B training steps, causing overfitting to MCQ format. Fix in v1v3: add Nemotron-Post-Training-Dataset-v1 STEM (~60B tokens, ~0.13 epochs at 80B) as primary science source; reduce Science-v1 to low weights (3+2) for format alignment only. + +**Why 80B is the recommended stopping point:** SciCode degrades or crashes at 100B (blend2: 1.6; AIME also degrades). Best overall profile is at 60–80B tokens. diff --git a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md new file mode 100644 index 00000000000..620c5780a4b --- /dev/null +++ b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md @@ -0,0 +1,326 @@ +# Nemotron-Nano-9B-v2: Prune + Distill + Quantize + vLLM Deployment + +End-to-end optimization of [Nemotron-Nano-9B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2) demonstrating how ModelOpt techniques stack: Minitron structured pruning to 7B → Megatron-Bridge knowledge distillation to recover accuracy → FP8 quantization → vLLM deployment and throughput benchmarking. This document covers: + +1. **[Data Preparation](#1-data-preparation)** — tokenizing the training blend for distillation +2. **[Pruning](#2-pruning)** — Minitron structured pruning from 9B to 7B +3. **[Distillation](#3-distillation)** — recovering accuracy via Megatron-Bridge knowledge distillation (up to 80B tokens) +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/llm_ptq/hf_ptq.py` script +6. **[vLLM Inference Benchmarking](#6-vllm-inference-benchmarking)** — throughput comparison of BF16 vs FP8 on a single H100 + +**Environment:** Container `nvcr.io/nvidia/nemo:26.02`, ModelOpt 0.44.0. See the [Megatron-Bridge README](../../../megatron_bridge/README.md) for environment setup (including ModelOpt mount path) and container usage. + +## Results + +![Benchmark Recovery During Knowledge Distillation](figures/learning_curves.png) + +| Model | MMLU | MMLU Pro | GPQA Diamond | LiveCodeBench v6 | AIME 2025 | Math 500 | IFEval | SciCode (Subtask) | Average | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | +| Pruned 7B (no distillation) | 67.8 | 11.9 | 17.7 | 1.4 | 0.3 | 6.0 | 41.8 | 0.1 | 18.4 | +| Pruned 7B + distill 2.5B tokens (400 iters) | 70.7 | 68.4 | 52.7 | 57.0 | 63.0 | 93.7 | 63.2 | 11.6 | 60.0 | +| Pruned 7B + distill 20B tokens (3200 iters) | 71.3 | 71.7 | 54.8 | 62.0 | 69.1 | 95.2 | 63.8 | 20.9 | 63.6 | +| Pruned 7B + distill 40B tokens (6400 iters) | 71.1 | 71.6 | 53.7 | 60.9 | 70.4 | 95.6 | 68.0 | 21.1 | 64.1 | +| Pruned 7B + distill 60B tokens (9600 iters) | 72.1 | 72.1 | 54.9 | 61.6 | 70.3 | 95.4 | 64.7 | 24.1 | 64.4 | +| Pruned 7B + distill 80B tokens (12800 iters) | 72.2 | 73.0 | 56.9 | 62.6 | 72.0 | 95.8 | 66.2 | 22.2 | 65.1 | +| Nemotron-Nano-9B-v2 (official, pruned from 12B) | 74.7 | 74.9 | 56.1 | 64.4 | 73.2 | 95.9 | 65.8 | 21.9 | 65.9 | +| Nemotron-Nano-12B-v2 (official) | 78.5 | 77.9 | 58.2 | 66.6 | 76.1 | 96.9 | 67.9 | 28.4 | 68.8 | + +**Key observations:** + +- **All benchmarks recover dramatically within the first checkpoint (2.5B tokens).** The pruned-only model is essentially non-functional, but a single distillation run recovers most capabilities. +- **Math 500 and IFEval plateau quickly** — essentially saturated after 2.5B tokens, with minimal gains over the remaining training. +- **MMLU also largely plateaus** after the first checkpoint. +- **AIME, MMLU Pro, GPQA, and SciCode continue improving** throughout the full run and benefit meaningfully from longer training. +- **The 7B model at 80B tokens closes most of the gap to the official 9B**, and actually exceeds it on GPQA, IFEval, and SciCode. The table below compares the 7B→9B gap against the 9B→12B gap — both are ~25% compression — showing that the second pruning round recovers more efficiently: + +| Benchmark | 7B (80B tokens) vs 9B | 9B (official) vs 12B | +| --- | --- | --- | +| MMLU | −2.5 | −3.8 | +| MMLU Pro | −1.9 | −3.0 | +| GPQA Diamond | **+0.8** | −2.1 | +| LiveCodeBench v6 | −1.8 | −2.2 | +| AIME 2025 | −1.2 | −2.9 | +| Math 500 | −0.1 | −1.0 | +| IFEval | **+0.4** | −2.1 | +| SciCode (Subtask) | **+0.3** | −6.5 | +| Average | −0.8 | −2.9 | + +Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post-training v1/v3 (Math 30, Coding 20, Science 15, IF 5)** blend (see [Data Blend](#data-blend) below). Blend ablations are in [ABLATIONS.md](ABLATIONS.md). + +> [!NOTE] +> Exact numbers may vary depending on deployment and evaluation setup. All models above — including the official 9B and 12B — were evaluated with the same [nemo_evaluator.yaml](nemo_evaluator.yaml) for fair comparison. These numbers may differ from those reported on the official [Nemotron-Nano-9B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2) and [Nemotron-Nano-12B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2) HuggingFace model cards. + +> [!NOTE] +> The official Nemotron-Nano-9B-v2 model was itself produced by pruning Nemotron-Nano-12B-v2 using Minitron. See [arxiv:2508.14444](https://arxiv.org/abs/2508.14444) for details on the exact steps used there. + +--- + +## Steps to Reproduce + +### 1. Data Preparation + +See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for tokenization commands for all datasets used in this blend. + +For this experiment: `TOKENIZER=nvidia/NVIDIA-Nemotron-Nano-9B-v2`, `OUTPUT_DIR=tokenized_nemotron_v2`. + +#### Data Blend + +**30% Pretraining (Code 5, General 20, MATH 5) + 70% Post-training v1/v3 (Math 30, Coding 20, Science 15, IF 5)** + +```bash +DATA_BLEND=" \ +5 tokenized_nemotron_v2/nvidia--Nemotron-Pretraining-SFT-v1_Nemotron-SFT-Code_train_text_max10000000 \ +20 tokenized_nemotron_v2/nvidia--Nemotron-Pretraining-SFT-v1_Nemotron-SFT-General_train_text_max10000000 \ +5 tokenized_nemotron_v2/nvidia--Nemotron-Pretraining-SFT-v1_Nemotron-SFT-MATH_train_text_max10000000 \ +15 tokenized_nemotron_v2/nvidia--Nemotron-Math-v2_default_high_part00_messages \ +15 tokenized_nemotron_v2/nvidia--Nemotron-Math-v2_default_high_part01_messages \ +15 tokenized_nemotron_v2/competitive_programming_python_00_messages \ +5 tokenized_nemotron_v2/competitive_programming_cpp_00_messages \ +10 tokenized_nemotron_v2/nvidia--Nemotron-Post-Training-Dataset-v1_default_stem_messages_max5000000 \ +3 tokenized_nemotron_v2/MCQ_messages \ +2 tokenized_nemotron_v2/RQA_messages \ +3 tokenized_nemotron_v2/reasoning_on_messages \ +2 tokenized_nemotron_v2/reasoning_off_messages \ +" +``` + +| Dataset | Tokens | Weight | Notes | +| --- | --- | --- | --- | +| Nemotron-Pretraining-SFT-v1 / Code (10M samples) | 7B | 5 | Pretraining code | +| Nemotron-Pretraining-SFT-v1 / General (10M samples) | 16B | 20 | Upweighted to better close MMLU gap | +| Nemotron-Pretraining-SFT-v1 / MATH (10M samples) | 12B | 5 | Pretraining math | +| Nemotron-Math-v2 / high_part00 | 9B | 15 | Hard math reasoning | +| Nemotron-Math-v2 / high_part01 | 11B | 15 | Hard math reasoning | +| Nemotron-SFT-Competitive-Programming-v2 / python_00 | 7B | 15 | Python reasoning traces | +| Nemotron-SFT-Competitive-Programming-v2 / cpp_00 | 7B | 5 | C++ reasoning traces | +| Nemotron-Post-Training-Dataset-v1 / stem (5M samples) | 20B | 10 | Broad STEM | +| Nemotron-Science-v1 / MCQ | 0.5B | 3 | GPQA MCQ format alignment | +| Nemotron-Science-v1 / RQA | 0.3B | 2 | GPQA format diversity | +| Nemotron-SFT-IF-Chat-v2 / reasoning_on | 2B | 3 | Instruction following (thinking on) | +| Nemotron-SFT-IF-Chat-v2 / reasoning_off | 1B | 2 | Instruction following (thinking off) | + +#### 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 important benchmarks reported in the Nemotron-Nano-9B-v2 model card. 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 (30%)** is the largest post-training category because AIME and MMLU Pro respond strongly to more math reasoning tokens. Two `Nemotron-Math-v2` splits are used to avoid repetition at longer token budgets. +- **Science (15%)** 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 — IFEval reaches 60+% within 2.5B tokens — so a small allocation is sufficient. + +This blend intentionally omits capabilities not targeted in this experiment (e.g. long context and multilingual benchmarks). 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: + +| Capability | Relevant datasets | +| --- | --- | +| Multilingual | `Nemotron-SFT-Multilingual-v1` | +| Agentic / tool use | `Nemotron-SFT-Tool-Call-v1`, `Nemotron-SFT-Tool-Call-v2` | +| Software engineering (SWE) | `Nemotron-SFT-SWE-v1` | +| Safety / alignment | `Nemotron-SFT-Safety-v1` | +| Long context | `Nemotron-SFT-Long-Context-v1` | + +When adding new datasets, reduce weights of lower-priority categories proportionally to keep the total at 100%. + +--- + +### 2. Pruning + +Run on **1 node with 8x H100** (~1 hour) + +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) since dataset has longer sequences. All other arguments use defaults i.e. we optimize for MMLU (10% subset, 0-shot) for the pruned model (without distillation). + +```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-Nano-9B-v2 \ + --trust_remote_code \ + --prune_target_params 7e9 \ + --hparams_to_skip num_attention_heads \ + --seq_length 8192 \ + --output_hf_path /path/to/Nemotron-Nano-9B-v2-Pruned-7B +``` + +Important pruning logs: + +```text +Only considering atmost 40% for width and 20% for depth pruning hparams +Skipping hparams_to_skip=['num_attention_heads'] during search space generation... + Search space for num_layers: [46, 48, 50, 52, 54, 56] + Search space for hidden_size: [2816, 3072, 3328, 3584, 3840, 4096, 4352, 4480] + Search space for mamba_num_heads: [80, 88, 96, 104, 112, 120, 128] + Search space for mamba_head_dim: [56, 64, 72, 80] + Search space for ffn_hidden_size: [9728, 10240, 10752, 11264, 11776, 12288, 12800, 13312, 13824, 14336, 14848, 15360, 15680] + Total search space in consideration: 17472 + +Top 10 candidates with scores: +{'num_layers': 50, 'hidden_size': 4480, 'mamba_num_heads': 128, 'mamba_head_dim': 56, 'ffn_hidden_size': 15680} -> 7.00B params, 0.2019 score +{'num_layers': 56, 'hidden_size': 4096, 'mamba_num_heads': 96, 'mamba_head_dim': 80, 'ffn_hidden_size': 14336} -> 7.00B params, 0.4363 score +{'num_layers': 48, 'hidden_size': 4352, 'mamba_num_heads': 120, 'mamba_head_dim': 80, 'ffn_hidden_size': 13824} -> 7.00B params, 0.6789 score [BEST SUBNET] +{'num_layers': 56, 'hidden_size': 4352, 'mamba_num_heads': 112, 'mamba_head_dim': 80, 'ffn_hidden_size': 10240} -> 7.00B params, 0.5203 score +{'num_layers': 54, 'hidden_size': 4480, 'mamba_num_heads': 104, 'mamba_head_dim': 80, 'ffn_hidden_size': 11264} -> 7.00B params, 0.2615 score +{'num_layers': 46, 'hidden_size': 4480, 'mamba_num_heads': 128, 'mamba_head_dim': 72, 'ffn_hidden_size': 14848} -> 7.00B params, 0.6165 score +{'num_layers': 50, 'hidden_size': 4480, 'mamba_num_heads': 112, 'mamba_head_dim': 64, 'ffn_hidden_size': 15680} -> 7.00B params, 0.4214 score +{'num_layers': 54, 'hidden_size': 4096, 'mamba_num_heads': 112, 'mamba_head_dim': 80, 'ffn_hidden_size': 13312} -> 7.00B params, 0.5894 score +{'num_layers': 56, 'hidden_size': 4352, 'mamba_num_heads': 120, 'mamba_head_dim': 72, 'ffn_hidden_size': 10752} -> 7.00B params, 0.4688 score +{'num_layers': 52, 'hidden_size': 4352, 'mamba_num_heads': 120, 'mamba_head_dim': 72, 'ffn_hidden_size': 12800} -> 7.00B params, 0.5596 score + +Dropping decoder layers [43, 44, 45, 46, 47, 48, 50, 52] from model. +Original hybrid_override_pattern: M-M-M-MM-M-M-M*-M-M-M*-M-M-M-M*-M-M-M-M*-M-MM-M-M-M-M-M- +Pruned hybrid_override_pattern: M-M-M-MM-M-M-M*-M-M-M*-M-M-M-M*-M-M-M-M*-MMMM-M- +``` + +> [!TIP] +> Here we skip the Knowledge Distillation (KD) step for candidates for simplicity. If you want to find a better pruned model, you can take the top K candidates' `export_config` from the logs above and then export all models separately and perform KD for ~2B tokens on each of them before selecting the best subnet based on your desired metrics. + +--- + +### 3. Distillation + +Non-default arguments: `--seq_length 8192` (default: 4096), `--mbs 4` (default: 1), `--train_iters 16000` (train upto ~100B tokens — can stop earlier and take intermediate checkpoints for smaller runs), `--lr_warmup_iters 100` (default: 50), `--eval_interval 400` (default: 100). All other arguments use defaults. + +Run on **96 nodes × 8x H100 (768 GPUs total)**. ~600 H100 GPU-hours per 1k steps (~6.3B tokens), i.e. ~45 min wall-clock per 1k steps. Full 80B token run (~13k steps) takes ~9k H100 GPU-hours (~10 hours wall-clock). + +>[!TIP] +> While we use 96 nodes here for faster training, you can also run with 1 node. If you dont want to do full distillation run, you can stop earlier and take intermediate checkpoints as well. + +```bash +torchrun --nproc_per_node 8 /opt/Model-Optimizer/examples/megatron_bridge/distill_minitron.py \ + --teacher_hf_path nvidia/NVIDIA-Nemotron-Nano-9B-v2 \ + --student_hf_path /path/to/Nemotron-Nano-9B-v2-Pruned-7B \ + --trust_remote_code \ + --tp_size 8 \ + --pp_size 1 \ + --data_paths "${DATA_BLEND}" \ + --data_path_to_cache /path/to/cache \ + --seq_length 8192 \ + --mbs 4 \ + --gbs 768 \ + --train_iters 16000 \ + --lr 1e-4 \ + --min_lr 1e-5 \ + --lr_warmup_iters 100 \ + --eval_interval 400 \ + --eval_iters 32 \ + --log_interval 10 \ + --output_dir + +# Optional: Weights & Biases logging +# --wandb_project \ +# --wandb_entity \ +# --wandb_exp_name +``` + +For multi-node Slurm runs, see the [Megatron-Bridge README](../../../megatron_bridge/README.md#slurm-usage) for details. + +Distillation saves checkpoints in Megatron distributed format under `/checkpoints/iter_XXXXXXX`. You can convert any intermediate checkpoint to HuggingFace format using the Megatron-Bridge conversion script (see [Megatron Bridge README](../../../megatron_bridge/README.md) for full details): + +```bash +python /opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py export \ + --hf-model /path/to/Nemotron-Nano-9B-v2-Pruned-7B \ + --megatron-path /checkpoints/iter_ \ + --hf-path /checkpoints/hf_iter_ +``` + +--- + +### 4. Evaluation + +The eval config xin [nemo_evaluator.yaml](nemo_evaluator.yaml) is for Slurm-based evaluation — it submits a vLLM serving job 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). + +Before running, update the following fields in the yaml: + +- `execution.hostname` — your Slurm login node hostname +- `execution.account` — your Slurm account +- `deployment.checkpoint_path` — Hugging Face checkpoint path (original, pruned or quantized) +- `evaluation.nemo_evaluator_config.config.params.extra.tokenizer` — same path as `checkpoint_path` + +> [!TIP] +> Uncomment `limit_samples` under any task to run a small subset and verify the end-to-end eval pipeline before launching full evals. + +```bash +pip install "nemo-evaluator-launcher[all]==0.1.90" + +# Set required environment variables: +export HF_TOKEN= +export SLURM_JOB_DIR= +export HF_HOME= +export VLLM_CACHE_ROOT= + +# Set additional unused but required environment variables: +export API_KEY=xxxxxx +export INFERENCE_API_KEY=xxxxxx +export OPENAI_CLIENT_ID=xxxxxx +export OPENAI_CLIENT_SECRET=xxxxxx + +nemo-evaluator-launcher run --config nemo_evaluator.yaml +``` + +**Tasks and exact metric names reported in the results table:** + +| Benchmark | Tool | Metric name | +| --- | --- | --- | +| MMLU | [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) (5-shot) | `mmlu` | +| MMLU Pro | NeMo Evaluator | `mmlu-pro_pass_at_1_symbolic_correct` | +| GPQA Diamond | NeMo Evaluator | `gpqa_pass_at_1_symbolic_correct` | +| LiveCodeBench v6 | NeMo Evaluator | `livecodebench_pass_at_1_accuracy` | +| AIME 2025 | NeMo Evaluator | `aime25_pass_at_1_symbolic_correct` | +| Math 500 | NeMo Evaluator | `AA_math_test_500_score_micro_avg_of_5` | +| IFEval | NeMo Evaluator | `ifeval_pass_at_1_average_score` | +| SciCode (Subtask) | NeMo Evaluator | `scicode_pass_at_1_subtask_accuracy` | + +**Key vLLM settings:** Tool calling is not enabled in these evals. + +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/). + +### 5. Quantization + +ModelOpt allows stacking multiple optimization techniques. Here we stack FP8 quantization on top of the pruned and distilled model to get an even more optimized model. See [examples/llm_ptq/README.md](../../../llm_ptq/README.md) for the full PTQ documentation. + +Similar to the official [Nemotron-Nano-9B-v2-FP8](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2-FP8) model, if you want to quantize the pruned 7B model to FP8, the Mamba and MLP layers are quantized to FP8, while all 4 attention layers and the Conv1d components within the Mamba layers are kept in BF16 to avoid accuracy degradation. + +This is done with the `mtq.MAMBA_MOE_FP8_AGGRESSIVE_CFG` config defined in [`modelopt/torch/quantization/config.py`](../../../../modelopt/torch/quantization/config.py). To apply this, you need to modify `QUANT_CFG_CHOICES["fp8"]` in [`examples/llm_ptq/hf_ptq.py`](../../../llm_ptq/hf_ptq.py) to use `mtq.MAMBA_MOE_FP8_AGGRESSIVE_CFG`. You may also consider using `mtq.MAMBA_MOE_FP8_CONSERVATIVE_CFG` for more conservative quantization. + +> [!NOTE] +> You can also quantize to NVFP4 using `mtq.MAMBA_MOE_NVFP4_AGGRESSIVE_CFG` or `mtq.MAMBA_MOE_NVFP4_CONSERVATIVE_CFG`, which may require further distillation (QAD) to recover accuracy and Blackwell GPU for deployment. + +Calibrate and export the HF checkpoint from iteration 12800 to FP8 (takes 1-2 mins on 8x H100): + +```bash +python /opt/Model-Optimizer/examples/llm_ptq/hf_ptq.py \ + --pyt_ckpt_path /checkpoints/hf_iter_12800 \ + --export_path /checkpoints/hf_iter_12800_fp8_aggressive \ + --qformat fp8 \ + --trust_remote_code +``` + +The quantized checkpoint is directly deployable with [vLLM](https://github.com/vllm-project/vllm), [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) and [SGLang](https://github.com/sgl-project/sglang). + +> [!TIP] +> You can run the evaluation using the same `nemo_evaluator.yaml` file for the quantized checkpoint also! + +### 6. vLLM Inference Benchmarking + +Benchmark throughput using [vLLM](https://github.com/vllm-project/vllm) on a single H100 GPU. Run the command once for each HuggingFace checkpoint. vLLM automatically detects FP8 quantization from the embedded `quantization_config` in `config.json` and applies it with no extra flags needed. + +Results on a single H100 (ISL=32768, OSL=1024): + +```bash +vllm bench throughput \ + --model \ + --random-input-len 32768 \ + --random-output-len 1024 \ + --trust-remote-code \ + --mamba_ssm_cache_dtype float32 \ + --kv-cache-dtype fp8 \ + --load-format safetensors +``` + +| Checkpoint | Model loading memory | Output tokens/s | Speedup vs Nemotron-Nano-9B-v2 BF16 | +| --- | --- | --- | --- | +| Nemotron-Nano-12B-v2 (official) | 22.9 GiB | 585 | 0.74× | +| Nemotron-Nano-9B-v2 (official) | 16.6 GiB | 794 | 1.00× | +| Nemotron-Nano-9B-v2-FP8 (official) | 9.6 GiB | 1,012 | 1.27× | +| Nemotron-Nano-9B-v2-Pruned-7B | 13.1 GiB | 963 | 1.21× | +| Nemotron-Nano-9B-v2-Pruned-7B-FP8 | 7.8 GiB | 1,147 | 1.44× | + +In this case, FP8 delivers a ~20-30% throughput gain over BF16 at the same parameter count. The NemotronH hybrid architecture (Mamba + attention) moderates this gain relative to pure-transformer models, since Attention and Conv1d layers are not quantized. diff --git a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/figures/learning_curves.png b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/figures/learning_curves.png new file mode 100644 index 00000000000..40c507bd1b8 Binary files /dev/null and b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/figures/learning_curves.png differ diff --git a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/nemo_evaluator.yaml b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/nemo_evaluator.yaml new file mode 100644 index 00000000000..256a4031be1 --- /dev/null +++ b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/nemo_evaluator.yaml @@ -0,0 +1,194 @@ +# NeMo Evaluator Launcher config for Nemotron-Nano-9B-v2 and Pruned variants +# -------------------------------------------------------------------------- +# Before running, update the following fields in the yaml: +# - `execution.hostname` — your Slurm login node hostname +# - `execution.account` — your Slurm account +# - `deployment.checkpoint_path` — Hugging Face checkpoint path (original, pruned or quantized) +# - `evaluation.nemo_evaluator_config.config.params.extra.tokenizer` — same path as `checkpoint_path` +# +# Usage: +# pip install "nemo-evaluator-launcher[all]==0.1.90" +# +# # Set required environment variables: +# export HF_TOKEN= +# export SLURM_JOB_DIR= +# export HF_HOME= +# export VLLM_CACHE_ROOT= +# +# # Set additional unused but required environment variables: +# export API_KEY=xxxxxx +# export INFERENCE_API_KEY=xxxxxx +# export OPENAI_CLIENT_ID=xxxxxx +# export OPENAI_CLIENT_SECRET=xxxxxx +# +# nemo-evaluator-launcher run --config nemo_evaluator.yaml +# + +defaults: + - execution: slurm/default + - deployment: vllm + - _self_ + +execution: + type: slurm + hostname: + username: ${oc.env:USER} + account: + partition: batch + num_nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + gres: "gpu:8" + walltime: 04:00:00 + sbatch_comment: "{\"OccupiedIdleGPUsJobReaper\":{\"exemptIdleTimeMins\":\"1920\",\"reason\":\"benchmarking\",\"description\":\"Some evals need idle time\ + \ else gets cancelled\"}}" + subproject: nel + output_dir: ${oc.env:SLURM_JOB_DIR} + mode: sequential + + mounts: + mount_home: false + deployment: + n_tasks: 1 + batch_comment: "{\"OccupiedIdleGPUsJobReaper\":{\"exemptIdleTimeMins\":\"1920\",\"reason\":\"benchmarking\",\"description\":\"Required data validation\ + \ and evaluation\"}}" + +# Note: Only tp=1 works for Nano (Mamba-based architecture) +deployment: + # Update this to your Hugging Face checkpoint path (original, pruned or quantized) + checkpoint_path: + served_model_name: Nemotron-Nano-9B-v2 + port: 8000 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + data_parallel_size: 8 + gpu_memory_utilization: 0.8 + extra_args: "--trust-remote-code --no-enable-prefix-caching --mamba_ssm_cache_dtype float32 --model-loader-extra-config '{\"enable_multithread_load\"\ + : true, \"num_threads\": 96}' --kv-cache-dtype fp8 " + env_vars: + VLLM_ATTENTION_BACKEND: FLASH_ATTN + endpoints: + chat: /v1/chat/completions + completions: /v1/completions + health: /health + multiple_instances: true + +evaluation: + nemo_evaluator_config: + target: + api_endpoint: + adapter_config: + use_system_prompt: true + use_reasoning: false + params_to_add: + chat_template_kwargs: + enable_thinking: true + skip_special_tokens: false + use_caching: true + tracking_requests_stats: true + log_failed_requests: true + use_request_logging: true + max_logged_requests: 10 + use_response_logging: true + max_logged_responses: 10 + config: + params: + parallelism: 64 + max_new_tokens: 32768 + temperature: 0.6 + top_p: 0.95 + request_timeout: 3600 + max_retries: 10 + extra: + tokenizer_backend: huggingface + # Update tokenizer path to match checkpoint_path above + tokenizer: + env_vars: + HF_TOKEN: HF_TOKEN + HF_HOME: HF_HOME + VLLM_CACHE_ROOT: VLLM_CACHE_ROOT + API_KEY: API_KEY + INFERENCE_API_KEY: INFERENCE_API_KEY + OPENAI_CLIENT_ID: OPENAI_CLIENT_ID + OPENAI_CLIENT_SECRET: OPENAI_CLIENT_SECRET + + tasks: + # 1. MMLU Pro + - name: ns_mmlu_pro + env_vars: + HF_TOKEN: HF_TOKEN + nemo_evaluator_config: + config: + params: + # limit_samples: 8 + extra: + num_repeats: 1 + args: "++prompt_config=eval/aai/mcq-10choices-boxed" + + # 2. GPQA Diamond + - name: ns_gpqa + env_vars: + HF_TOKEN: HF_TOKEN + nemo_evaluator_config: + config: + params: + # limit_samples: 8 + extra: + num_repeats: 8 + args: "++prompt_config=eval/aai/mcq-4choices" + + # 3. LiveCodeBench + - name: ns_livecodebench + env_vars: + HF_TOKEN: HF_TOKEN + nemo_evaluator_config: + config: + params: + # limit_samples: 8 + extra: + num_repeats: 8 + dataset_split: test_v6_2408_2505 + + # 4. AIME 2025 + - name: ns_aime2025 + env_vars: + HF_TOKEN: HF_TOKEN + nemo_evaluator_config: + config: + params: + # limit_samples: 8 + extra: + num_repeats: 64 + + # 5. MATH500 (Requires JUDGE_API_KEY) + # - name: AA_math_test_500 + # env_vars: + # HF_TOKEN: HF_TOKEN + # JUDGE_API_KEY: JUDGE_API_KEY + # nemo_evaluator_config: + # config: + # params: + # # limit_samples: 8 + # extra: + # n_samples: 5 + + # 6. IFEval + - name: ns_ifeval + env_vars: + HF_TOKEN: HF_TOKEN + # nemo_evaluator_config: + # config: + # params: + # limit_samples: 8 + + # 7. SciCode + - name: ns_scicode + env_vars: + HF_TOKEN: HF_TOKEN + nemo_evaluator_config: + config: + params: + # limit_samples: 8 + max_new_tokens: 8192 + extra: + num_repeats: 8 diff --git a/examples/pruning/minitron/README.md b/examples/pruning/minitron/README.md new file mode 100644 index 00000000000..8749c366a70 --- /dev/null +++ b/examples/pruning/minitron/README.md @@ -0,0 +1,11 @@ +# Minitron Pruning — End-to-End Tutorials + +End-to-end tutorials for [Minitron](https://arxiv.org/abs/2407.14679) structured pruning followed by knowledge distillation, quantization, evaluation,and vLLM deployment. + +Each subdirectory covers a specific source model and target size, including the full data blend, pruning config, distillation hyperparameters, evaluation results, and throughput benchmarks. + +## Related + +- [Minitron pruning instructions](../../megatron_bridge/README.md#pruning) and [Megatron-Bridge distillation instructions](../../megatron_bridge/README.md#distillation) +- [Megatron dataset tokenization](../../dataset/MEGATRON_DATA_PREP.md) +- [Puzzletron pruning algorithm](../../puzzletron/README.md) diff --git a/examples/megatron_bridge/results/puzzletron.md b/examples/pruning/puzzletron/Llama-3.1-8B-Instruct.md similarity index 100% rename from examples/megatron_bridge/results/puzzletron.md rename to examples/pruning/puzzletron/Llama-3.1-8B-Instruct.md diff --git a/examples/pruning/puzzletron/README.md b/examples/pruning/puzzletron/README.md new file mode 100644 index 00000000000..426ced00c4c --- /dev/null +++ b/examples/pruning/puzzletron/README.md @@ -0,0 +1,16 @@ +# Puzzletron Pruning — Distillation Results + +Distillation results for models compressed with [Puzzletron](../../puzzletron/README.md) MIP-based heterogeneous pruning, followed by Megatron-Bridge knowledge distillation. + +## Results + +| Model | File | +| --- | --- | +| Llama-3.1-8B-Instruct and Qwen3-8B | [Llama-3.1-8B-Instruct.md](Llama-3.1-8B-Instruct.md) | + +## Related + +- [Puzzletron pruning example](../../puzzletron/README.md) +- [Megatron-Bridge distillation instructions](../../megatron_bridge/README.md#distillation) +- [Megatron dataset tokenization](../../dataset/MEGATRON_DATA_PREP.md) +- [Minitron pruning instructions](../../pruning/README.md#minitron) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 89183073399..571b40ca499 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -341,6 +341,8 @@ To recover degradation in the quality of the compressed model, we can use knowle See [Megatron-Bridge distillation](../megatron_bridge/README.md#distillation) for instructions on using Megatron-Bridge for knowledge distillation. The distillation script supports both standard HuggingFace and Puzzletron AnyModel checkpoints. +For distillation results on Puzzletron-compressed models, see [examples/pruning/puzzletron/](../pruning/puzzletron/README.md). + ## Advanced Usage Modify `llama-3_1-8B_pruneffn_memory.yaml` file for advanced compression scenarios. diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 0c9a121f696..81dac1580b0 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -78,8 +78,9 @@ --strip_newlines ``` -Note: ``--hf_streaming`` without ``--hf_max_samples_per_split`` falls back to non-streaming, -since streaming the full dataset is slower than the cached non-streaming path. +Note: streaming does not cache to disk, so re-runs re-download. For full-dataset streaming +without a sample cap this is slower than non-streaming mode, but it avoids Arrow schema +compatibility issues with complex nested message types. """ import argparse @@ -191,7 +192,14 @@ def encode(self, json_line: str): if tools: kwargs["tools"] = tools value = self._process_messages(value) - text = _Encoder.tokenizer.apply_chat_template(value, tokenize=False, **kwargs) + try: + text = _Encoder.tokenizer.apply_chat_template(value, tokenize=False, **kwargs) + except Exception as e: + print( + f"apply_chat_template failed: {e}\nData:\n{json.dumps(data, indent=2, default=str)}", + flush=True, + ) + raise # chat template already embeds all special tokens; don't add BOS again add_special_tokens = False else: @@ -452,8 +460,9 @@ def megatron_preprocess_data( hf_split: Hugging Face Hub dataset split. Defaults to None (all splits). hf_max_samples_per_split: Maximum number of rows to consume per split. hf_streaming: Load HuggingFace datasets in streaming mode. Only consumed rows are - downloaded — useful for very large pretraining datasets. Note: streaming does not - cache to disk, so re-runs re-download. Defaults to False. + downloaded — useful for very large pretraining datasets or datasets with complex + nested message schemas that cause Arrow type-cast errors in non-streaming mode. + Note: streaming does not cache to disk, so re-runs re-download. Defaults to False. output_dir: Path to directory to save binary output files. tokenizer_name_or_path: Name or path of the Hugging Face tokenizer to use. json_keys: Key or list of keys to extract from json. Defaults to ["text"]. @@ -485,10 +494,9 @@ def megatron_preprocess_data( warnings.warn( "--hf_streaming is set but --hf_max_samples_per_split is not. " "Streaming without a sample cap re-downloads the full dataset on every run with no " - "disk cache, which is slower than non-streaming mode. Falling back to streaming=False.", + "disk cache, which is slower than the cached non-streaming path.", stacklevel=2, ) - hf_streaming = False Path(output_dir).mkdir(parents=True, exist_ok=True) vocab_size = AutoTokenizer.from_pretrained(tokenizer_name_or_path).vocab_size