Merge branch 'main' into jingyux/ltx-2-update

Author: jingyu-ml

.pre-commit-config.yaml

@@ -24,7 +24,9 @@ repos:
     hooks:
       - id: ruff-check
         args: [--fix, --exit-non-zero-on-fix]
+        exclude: ^examples/specdec_bench/specdec_bench/datasets/speed\.py$
       - id: ruff-format
+        exclude: ^examples/specdec_bench/specdec_bench/datasets/speed\.py$
 
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.17.1
@@ -93,6 +95,7 @@ repos:
               examples/llm_eval/modeling.py|
               examples/llm_qat/main.py|
               examples/llm_sparsity/weight_sparsity/finetune.py|
+              examples/specdec_bench/specdec_bench/models/specbench_medusa.py|
               examples/speculative_decoding/main.py|
               examples/speculative_decoding/medusa_utils.py|
               examples/speculative_decoding/server_generate.py|

CHANGELOG.rst

@@ -8,6 +8,7 @@ NVIDIA Model Optimizer Changelog (Linux)
 
 - User does not need to manually register MOE modules to cover experts calibration coverage in PTQ workflow.
 - ``hf_ptq.py`` now saves the quantization summary and moe expert token count table to the export directory.
+- Add ``--moe_calib_experts_ratio`` flag in ``hf_ptq.py`` to specify the ratio of experts to calibrate during forward pass to improve expert coverage during calibration. Default to all the experts.
 - Add sparse attention optimization for transformer models (``modelopt.torch.sparsity.attention_sparsity``). This reduces computational cost by skipping attention computation. Supports calibration for threshold selection on HuggingFace models. See `examples/llm_sparsity/attention_sparsity/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_sparsity/attention_sparsity>`_ for usage.
 - Add support for rotating the input before quantization for RHT.
 

examples/llm_ptq/example_utils.py

@@ -201,6 +201,7 @@ def build_quant_cfg(
     model_type,
     quant_cfg_choices,
     kv_quant_cfg_choices,
+    moe_calib_experts_ratio: float | None = None,
 ) -> dict[str, Any]:
     quant_cfg = {}
     assert qformat in quant_cfg_choices, (
@@ -232,6 +233,20 @@ def build_quant_cfg(
             getattr(mtq, kv_quant_cfg_choices[kv_cache_qformat])["quant_cfg"],
         )
 
+    if moe_calib_experts_ratio:
+        assert 0 < moe_calib_experts_ratio <= 1, "moe_calib_experts_ratio must be between 0 and 1"
+        if isinstance(quant_cfg["algorithm"], str):
+            quant_cfg["algorithm"] = {
+                "method": quant_cfg["algorithm"],
+                "moe_calib_experts_ratio": moe_calib_experts_ratio,
+            }
+        elif isinstance(quant_cfg["algorithm"], dict):
+            quant_cfg["algorithm"]["moe_calib_experts_ratio"] = moe_calib_experts_ratio
+        else:
+            warnings.warn(
+                f"Quantization algorithm: {quant_cfg['algorithm']} does not support setting moe_calib_experts_ratio"
+            )
+
     # Gemma 7B has accuracy regression using alpha 1. We set 0.5 instead.
     if model_type == "gemma" and "int8_sq" in qformat:
         quant_cfg["algorithm"] = {"method": "smoothquant", "alpha": 0.5}

examples/llm_ptq/hf_ptq.py

@@ -906,6 +906,7 @@ def quantize_main(
             model_type,
             QUANT_CFG_CHOICES,
             KV_QUANT_CFG_CHOICES,
+            args.moe_calib_experts_ratio,
         )
 
         # Exclude MTP layers from quantization if detected (e.g., GLM-4.7's layer 92)
@@ -1126,8 +1127,21 @@ def parse_args() -> argparse.Namespace:
             "(sensitivity scores, costs, etc.). Only used when auto_quantize_bits is specified."
         ),
     )
+    parser.add_argument(
+        "--moe_calib_experts_ratio",
+        type=float,
+        default=1.0,
+        help=(
+            "Fraction of experts to calibrate during forward pass (ratio in (0.0, 1.0]). "
+            "Only used for MOE models; used to reduce the number of experts calibrated during the forward pass."
+            "Does not impact non-MOE models."
+        ),
+    )
 
-    return parser.parse_args()
+    args = parser.parse_args()
+    if not (0.0 < args.moe_calib_experts_ratio <= 1.0):
+        parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].")
+    return args
 
 
 def main(args: argparse.Namespace):

examples/megatron_bridge/README.md

@@ -128,6 +128,9 @@ python -m modelopt.torch.utils.plugins.megatron_preprocess_data \
     --max_sequence_length 256_000
 ```
 
+The [Nemotron-Pretraining-SFT-v1](https://huggingface.co/datasets/nvidia/Nemotron-Pretraining-SFT-v1) dataset is huge, so it will take a while to download and tokenize. You can also split the large `.jsonl` into multiple files (e.g. 10M samples per file using `split -l 10000000 -d --additional-suffix=.jsonl <file>.jsonl <file>_part`) and tokenize them parallelly.
+To quickly test the script, you can try the [nvidia/Nemotron-Pretraining-Dataset-sample](https://huggingface.co/datasets/nvidia/Nemotron-Pretraining-Dataset-sample) dataset.
+
 If you skip `--hf_name`, it will download and tokenize all subsets for the dataset.
 If you skip `--hf_split`, it will download and tokenize all splits for the subset.
 If you skip `--hf_max_samples_per_split`, it will download and tokenize all samples for the split.

examples/specdec_bench/README.md

@@ -28,17 +28,121 @@ MTBench is available [here](https://huggingface.co/datasets/HuggingFaceH4/mt_ben
 Download `nvidia/gpt-oss-120b-Eagle3` to a local directory `/path/to/eagle`.
 
 ```bash
-python3 run.py --model_dir openai/gpt-oss-120b --tokenizer openai/gpt-oss-120b --draft_model_dir /path/to/eagle --mtbench question.jsonl --tp_size 1 --ep_size 1 --draft_length 3 --output_length 4096 --num_requests 80 --engine TRTLLM --concurrency 1 --postprocess gptoss
-
+python3 run.py \
+    --model_dir openai/gpt-oss-120b \
+    --tokenizer openai/gpt-oss-120b \
+    --draft_model_dir /path/to/eagle \
+    --mtbench question.jsonl \
+    --tp_size 1 \
+    --ep_size 1 \
+    --draft_length 3 \
+    --output_length 4096 \
+    --num_requests 80 \
+    --engine TRTLLM \
+    --concurrency 1 \
+    --postprocess gptoss
 ```
 
 ### Running Random ids on GPT OSS + Eagle3
 
 Download `nvidia/gpt-oss-120b-Eagle3` to a local directory `/path/to/eagle`.
 
 ```bash
-python3 run.py --model_dir openai/gpt-oss-120b --tokenizer openai/gpt-oss-120b --draft_model_dir /path/to/eagle --random_isl 1024 --tp_size 1 --ep_size 1 --draft_length 3 --output_length 4096 --num_requests 40 --engine TRTLLM --concurrency 1
+python3 run.py \
+    --model_dir openai/gpt-oss-120b \
+    --tokenizer openai/gpt-oss-120b \
+    --draft_model_dir /path/to/eagle \
+    --random_isl 1024 \
+    --tp_size 1 \
+    --ep_size 1 \
+    --draft_length 3 \
+    --output_length 4096 \
+    --num_requests 40 \
+    --engine TRTLLM \
+    --concurrency 1
+```
+
+### Running [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) on Llama 3.3 70B + Eagle 3
+
+1. Install the requirements file using `pip install -r requirements_speed.txt`
+
+2. Prepare the data using the provided script:
+
+```bash
+python3 prepare_data.py --dataset speed --config all
+```
+
+The data will be saved to `data/` directory, each config type (qualitative, throughput_1k, ...) to each own directory.
+
+#### License
+
+GOVERNING TERMS: This dataset is governed by the NVIDIA Evaluation Dataset License Agreement.
+
+ADDITIONAL INFORMATION: MIT for bigcode/humanevalpack, RUCAIBox/MMATH, RUCAIBox/BAMBOO and EQ-Bench. Apache 2.0 for Writing Bench and Spec-Bench. CC BY 4.0 for FBK-MT/MCIF. MIT and Apache 2.0 for tianyang/repobench_python_v1.1, JetBrains-Research/lca-project-level-code-completion and tianyang/repobench_java_v1.1.
+
+NOTICE: For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare_data.py` script automatically fetches data from all the source datasets.
+
+Additional details are in [HuggingFace dataset repository](https://huggingface.co/datasets/nvidia/SPEED-Bench).
+
+#### Qualitative split
+
+```bash
+python3 run.py \
+    --model_dir meta-llama/Llama-3.3-70B-Instruct \
+    --tokenizer meta-llama/Llama-3.3-70B-Instruct \
+    --draft_model_dir yuhuili/EAGLE3-LLaMA3.3-Instruct-70B \
+    --dataset speed \
+    --dataset_path data/speed/qualitative \
+    --tp_size 8 \
+    --ep_size 1 \
+    --draft_length 3 \
+    --output_length 4096 \
+    --engine TRTLLM \
+    --concurrency 32 \
+    --show_progress
+```
+
+#### Throughput split
 
+```bash
+python3 run.py \
+    --model_dir meta-llama/Llama-3.3-70B-Instruct \
+    --tokenizer meta-llama/Llama-3.3-70B-Instruct \
+    --draft_model_dir yuhuili/EAGLE3-LLaMA3.3-Instruct-70B \
+    --dataset speed \
+    --dataset_path data/speed/throughput_1k \
+    --tp_size 8 \
+    --ep_size 1 \
+    --draft_length 3 \
+    --output_length 4096 \
+    --engine TRTLLM \
+    --concurrency 32 \
+    --show_progress
+```
+
+For longer context (>8192 tokens), please use the following configuration when using TRTLLM:
+
+```yaml
+engine_args:
+  max_seq_len: 131072   # Model max context length (for Llama 3.3 70B)
+  enable_chunked_prefill: true
+```
+
+```bash
+python3 run.py \
+    --model_dir meta-llama/Llama-3.3-70B-Instruct \
+    --tokenizer meta-llama/Llama-3.3-70B-Instruct \
+    --draft_model_dir yuhuili/EAGLE3-LLaMA3.3-Instruct-70B \
+    --dataset speed \
+    --dataset_path data/speed/throughput_16k \
+    --tp_size 8 \
+    --ep_size 1 \
+    --draft_length 3 \
+    --output_length 4096 \
+    --engine TRTLLM \
+    --concurrency 32 \
+    --show_progress \
+    --runtime_params runtime_args_long_context.yaml
 ```
 
 ## Notes