Model-Optimizer/examples/megatron_bridge/README.md

# Megatron Bridge

This directory contains examples of using Model Optimizer with the [NeMo Megatron-Bridge](https://github.com/NVIDIA-Nemo/Megatron-Bridge) framework for quantization, pruning, and distillation. These workflows can be used on their own or combined.

<div align="center">

| **Section** | **Description** | **Link** |
| :------------: | :------------: | :------------: |
| Pre-Requisites | Development environment setup | \[[Link](#pre-requisites)\] |
| Post-Training Quantization | Quantizing a model | \[[Link](#post-training-quantization)\] |
| Distillation | Distilling a pruned or quantized model | \[[Link](#distillation)\] |
| Pruning | Pruning a model using Minitron algorithm | \[[Link](#pruning)\] |
| Sanity-Check Generation | Quick generation check with vLLM | \[[Link](#sanity-check-generation)\] |
| Resources | Extra links to relevant resources | \[[Link](#resources)\] |

</div>

> [!TIP]
> Checkout the [Nemotron-3-Nano-30B-A3B pruning + distillation (with data blend prep) + quantization tutorial](tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md) for a complete end-to-end workflow using Megatron-Bridge!
>
> Or the [Qwen3.6-35B-A3B W4A4 NVFP4 + QAD tutorial](tutorials/Qwen3.6-35B-A3B/README.md) for an end-to-end quantization-aware distillation workflow.

## Pre-Requisites

Running these examples requires many additional dependencies to be installed (e.g., Megatron-Bridge, Megatron-core, etc.), hence we strongly recommend directly using the NeMo container (e.g., `nvcr.io/nvidia/nemo:26.08`) which has all the dependencies installed.

To get the ModelOpt examples scripts, mount your Model-Optimizer repo to the container as follows:

```bash
export MODELOPT_DIR=${PWD}/Model-Optimizer # or set to your local Model-Optimizer repository path if you have cloned it
if [ ! -d "${MODELOPT_DIR}" ]; then
  git clone https://github.com/NVIDIA/Model-Optimizer.git ${MODELOPT_DIR}
fi

export DOCKER_IMAGE=nvcr.io/nvidia/nemo:26.08
docker run \
  --gpus all \
  --shm-size=16GB \
  --net=host \
  --ulimit memlock=-1 \
  --rm -it \
  -v ${MODELOPT_DIR}:/opt/Model-Optimizer \
  -v ${MODELOPT_DIR}/modelopt:/opt/venv/lib/python3.12/site-packages/modelopt \
  -v ${MODELOPT_DIR}/modelopt_recipes:/opt/venv/lib/python3.12/site-packages/modelopt_recipes \
  -w /opt/Model-Optimizer/examples/megatron_bridge \
  ${DOCKER_IMAGE} bash
```

> [!WARNING]
> 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.

You also need to login with your HuggingFace token to download gated datasets / models.
Note that the default dataset for pruning and quantization is [`nemotron-post-training-dataset-v2`](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2), which is gated.

```bash
hf auth login --token <your token>
```

### Importing a HuggingFace Checkpoint (optional)

The scripts below take a HuggingFace checkpoint directly, but if you want a Megatron distributed
checkpoint (e.g. to reuse across runs), convert it with Megatron-Bridge's conversion script — use
`--tp` / `--pp` / `--ep` to shard a model that does not fit on one GPU (`--ep` for MoE models), or
`--device cpu` to convert in a single process without GPUs:

```bash
bash /opt/Megatron-Bridge/scripts/conversion/convert.sh import \
    --executor local \
    --device gpu \
    --gpus-per-node 8 \
    --hf-model Qwen/Qwen3-8B \
    --megatron-path /tmp/Qwen3-8B-megatron
```

## Post-Training Quantization

This section shows how to quantize a HuggingFace model using ModelOpt in the Megatron-Bridge framework. Quantization is a two-step flow:

1. [quantize.py](quantize.py) applies post-training quantization (PTQ) with calibration and saves a **Megatron checkpoint** (with ModelOpt state). Tensor / pipeline / expert parallelism are all supported, and the checkpoint can be reloaded for further training (Quantization Aware Training / Quantization Aware Distillation).
2. [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) converts that Megatron checkpoint to a **HuggingFace (unified) checkpoint** that deploys directly with TensorRT-LLM, vLLM, or SGLang.

`quantize.py` supports the following formats via `--quant_cfg` (e.g. `fp8`, `nvfp4`, `int8_smoothquant`, `int4_awq`, `w4a8_awq_beta`, ...). You can also pass any full config name exposed by ModelOpt (e.g. `NVFP4_DEFAULT_CFG`) or a YAML `--recipe` (e.g. `general/ptq/nvfp4_default-kv_fp8`, authoritative for quant_cfg + algorithm + KV-cache). KV-cache quantization can be enabled on top via `--kv_cache_quant` (e.g. `fp8`, `nvfp4`).

**Step 1 — quantize** Qwen3-8B to NVFP4 on 2 GPUs (Tensor Parallelism = 2) using 1024 samples from default dataset (Mix of [`cnn_dailymail`](https://huggingface.co/datasets/abisee/cnn_dailymail) and [`nemotron-post-training-dataset-v2`](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2)) for calibration (sequence length = 4096):

```bash
torchrun --nproc_per_node 2 quantize.py \
    --hf_model_name_or_path Qwen/Qwen3-8B \
    --quant_cfg nvfp4 \
    --tp_size 2 \
    --calib_batch_size 1 \
    --seq_length 4096 \
    --export_megatron_path /tmp/Qwen3-8B-NVFP4-megatron
```

> [!NOTE]
> Data parallelism is implicit: `DP = world_size / (tp_size * pp_size * cp_size)`. Launching with more GPUs than `tp_size * pp_size * cp_size` shards calibration across the extra data-parallel ranks (e.g. `torchrun --nproc_per_node 8 quantize.py --tp_size 2` runs with DP=4).

**Step 2 — export** the Megatron checkpoint to a deployable HuggingFace checkpoint:

```bash
torchrun --nproc_per_node 2 export_quantized_megatron_to_hf.py \
    --hf_model_name_or_path Qwen/Qwen3-8B \
    --megatron_path /tmp/Qwen3-8B-NVFP4-megatron \
    --pp_size 2 \
    --export_unified_hf_path /tmp/Qwen3-8B-NVFP4-hf
```

> [!NOTE]
> The HuggingFace unified exporter can't split weights across GPUs with tensor parallelism. For large models, use `--pp_size` on `export_quantized_megatron_to_hf.py` to shard the export across GPUs with pipeline parallelism instead.

> [!TIP]
> To recover the accuracy lost during quantization, fine-tune the quantized Megatron checkpoint (from step 1) with [Quantization Aware Distillation (QAD)](#quantization-aware-distillation-qad) before running the step 2 export.

To see the full usage for advanced configurations, run `torchrun --nproc_per_node 1 quantize.py --help` (or `export_quantized_megatron_to_hf.py --help`).

### Vision-Language Models (VLMs)

For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automatically quantizes only the **language model** and leaves the vision tower and vision-language projector in full precision, then saves the full VLM back as a Megatron checkpoint. The calibration modality is inferred from `--calib_dataset_name`:

- An **image-text** dataset (the default for VLMs, `nemotron_vlm_dataset_v2`) drives the full VLM forward, so the language model is calibrated on vision-conditioned activations.
- A **text** dataset runs text-only calibration of the language model (vision tower idle).

> [!NOTE]
> HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM covers **Qwen3-VL** and **Qwen3.5-VL**. Other VLMs such as Gemma3-VL are saved in Megatron checkpoint format only.

### Tracking runs with MLflow

Set MLflow's own `MLFLOW_TRACKING_URI`, or pass `--mlflow <tracking-uri>`, to record a run on an MLflow server. Every script here that writes a checkpoint takes the flag — [`prune_minitron.py`](#pruning), `quantize.py`, [`distill.py`](#distillation), `export_quantized_megatron_to_hf.py` and `export_distilled_megatron_to_hf.py` — and they share one experiment-name convention, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be found together. (One gap remains: `distill.py --hf_export_path` writes its HuggingFace checkpoint from rank 0, which is not the rank that owns that run, so it carries no pointer.)

```bash
torchrun --nproc_per_node 2 quantize.py \
    --hf_model_name_or_path Qwen/Qwen3-8B \
    --recipe general/ptq/nvfp4_default-kv_fp8 \
    --tp_size 2 \
    --export_megatron_path /tmp/Qwen3-8B-NVFP4-megatron \
    --mlflow https://<your-mlflow-server>/
```

The run opens *before* the model loads, so a bad URI fails in seconds rather than after a full calibration. Only the master rank uploads: the invocation, every argument as a searchable param, the resolved recipe, that rank's log and the quantizer summary — plus `.experiment.json` written into `--export_megatron_path` once the checkpoint is saved, so a checkpoint on disk names the run that produced it. A failed run is still recorded, with its traceback.

`--mlflow_experiment` defaults to `$USER/<script>/<model basename>-<variant>` — the variant being the pruning target for `prune_minitron.py`, the recipe name for `quantize.py`, the student checkpoint for `distill.py`, and the Megatron checkpoint for the export — and `--mlflow_run_name` to the UTC start time. Authentication uses MLflow's own environment variables. See the [`hf_ptq` README](../hf_ptq/README.md#tracking-runs-with-mlflow) for the full artifact list and the `$MLFLOW_TRACKING_URI` semantics.

`distill.py` — whether it is distilling a pruned BF16 student or running [QAD](#quantization-aware-distillation-qad) from a quantized one — works differently under the hood. It is a training loop, and Megatron-Bridge logs to MLflow from inside it, so the run is opened here and Megatron-Bridge joins it: you get the invocation, every argument and a run log, plus the per-iteration training metrics and the full resolved config that only Megatron-Bridge can see, alongside the existing `--wandb_project` and TensorBoard logging. Two differences worth knowing:

- Uploading checkpoints as MLflow artifacts is **off** by default here, where Megatron-Bridge turns it on: a distilled checkpoint is tens to hundreds of GB and would be pushed over HTTP on every save. Pass `--mlflow_log_checkpoints` to opt in.
- The run is opened on the last rank, because that is the rank Megatron-Bridge looks at for one to join. The uploaded `logs/distill.log` is therefore that rank's output; `print_rank_0` lines, which is most of what the script itself prints, stay on rank 0 and do not reach it.

## Distillation

This section shows how to distill a student model from a teacher model in the Megatron-Bridge framework.

This can be used stand-alone or after [Pruning](#pruning) / [Post-Training Quantization](#post-training-quantization) to recover accuracy of the model by distilling from the original model (teacher).

The [distill.py](distill.py) script supports both standard HuggingFace checkpoints and [Puzzletron AnyModel](../puzzletron/README.md) checkpoints as student/teacher inputs. Just pass the checkpoint path via `--student_hf_path` / `--teacher_hf_path`. The distilled model is saved to `<output_dir>/checkpoints` in Megatron distributed checkpoint format.

To distill a student whose weights live in a **Megatron checkpoint** (e.g. a quantized checkpoint from [quantize.py](quantize.py) for [Quantization Aware Distillation](#quantization-aware-distillation-qad), or a pruned checkpoint), additionally pass `--student_megatron_path` — `--student_hf_path` is still required to build the student architecture.

### Data Preparation

The distillation script expects pre-tokenized data in Megatron's binary format (`.bin` / `.idx` files).

See the **[Dataset Preparation README](../dataset/README.md#tokenizing-for-megatron-frameworks)**
for full instructions on tokenizing JSONL files and Hugging Face datasets and get the list of output prefixes that you can use for `--data_paths` argument.

Alternatively, pass `--sft --sft_dataset_root <dir>` to distill on **raw prompt-completion JSONL**
with the loss masked to the completion. The directory must hold `training.jsonl` (and
`validation.jsonl` when `--eval_iters > 0`) of `{"input": <prompt>, "output": <response>}` records, which are tokenized with
the model's own HuggingFace tokenizer. Both fields are tokenized **as written**, except that
leading and trailing spaces on each field are stripped — no chat template is applied. So if your
model expects role/turn markers, include them in the `"input"` field yourself, and express any
significant separator as a newline rather than a trailing space. A BOS token is prepended
automatically when the tokenizer prepends one at inference, so do not add it yourself; an EOS
token is appended after the response. A record longer than `--seq_length` is truncated from the
**start** of `"input"`, which drops any system prompt or opening role marker baked in there, so
pre-filter or pre-truncate the corpus if that matters.

Teacher and student must share a tokenizer — distillation scores the teacher on the student's
token ids, and the KD losses compare the two models' logits elementwise over the vocab dimension.

### Distillation with Real Data

Example usage to distill a 4B student (HF) from an 8B teacher (HF) on 8 GPUs (TP=8, PP=1):

```bash
torchrun --nnodes 1 --nproc_per_node 8 distill.py \
    --tp_size 8 \
    --teacher_hf_path Qwen/Qwen3-8B \
    --student_hf_path Qwen/Qwen3-4B \
    --data_paths 1.0 tokenized_qwen3/data1_text_document 1.0 tokenized_qwen3/data2_text_document \
    --data_path_to_cache /path/to/cache/dataset_indices_qwen3 \
    --seq_length 8192 \
    --mbs 1 \
    --gbs 768 \
    --train_iters 15000 \
    --lr 1e-4 \
    --min_lr 1e-5 \
    --lr_warmup_iters 50 \
    --eval_interval 100 \
    --eval_iters 32 \
    --log_interval 10 \
    --output_dir /output/qwen3_8b_to_4b_distill
```

Tensorboard logging is enabled by default and logs are saved to `<output_dir>/tensorboard` directory.
To use Weights & Biases for logging, set the `WANDB_API_KEY` environment variable and pass the `--wandb_project` argument.
Optionally, you can also pass `--wandb_entity` and `--wandb_exp_name` arguments to group runs under a project and experiment name.

To measure the initial student's CE and distillation losses, add `--validate_only` to the command.
This skips training and evaluates the student at iteration 0.

To see all available arguments:

```bash
torchrun --nproc_per_node 1 distill.py --help
```

### Quick Test with Mock Data

Example usage with mock data for quick testing (no pre-tokenized data needed):

```bash
torchrun --nproc_per_node 8 distill.py \
    --tp_size 8 \
    --teacher_hf_path Qwen/Qwen3-0.6B \
    --student_hf_path Qwen/Qwen3-0.6B \
    --use_mock_data \
    --seq_length 512 \
    --mbs 1 \
    --gbs 8 \
    --train_iters 100 \
    --eval_interval 10 \
    --eval_iters 4 \
    --output_dir /tmp/test_distill
```

### Vision-Language Models (VLMs)

For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `distill.py` distills only the **language model** (on text data) and leaves the vision tower and projector untouched — matching the pruning and quantization behavior. It composes with pruning and QAD (`--student_megatron_path`) exactly as for LLMs, and the HF export reuses `--student_hf_path` (no `--student_hf_model` needed).

```bash
torchrun --nproc_per_node 8 distill.py \
    --tp_size 8 \
    --teacher_hf_path Qwen/Qwen3-VL-2B-Thinking \
    --student_hf_path Qwen/Qwen3-VL-2B-Thinking \
    ...
```

### Converting to Hugging Face format (optional)

A **non-quantized** distilled checkpoint (LLM or VLM) is saved in Megatron distributed format. If you need a HuggingFace checkpoint, there are two ways to convert it (for a **QAD** checkpoint, which retains quantization state, use [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) instead — see [QAD](#quantization-aware-distillation-qad)):

**Inline** -- add `--hf_export_path` to the `distill.py` command to automatically convert the **final** checkpoint after distillation:

```bash
torchrun --nnodes 1 --nproc_per_node 8 distill.py \
    ... \
    --hf_export_path /path/to/save/distilled_hf_ckpt
```

`--student_hf_path` builds the student and provides the exported config / tokenizer. `--student_hf_model` is a reference HF model with a **homogeneous** architecture, used as the export template only for **heterogeneous** (Puzzletron/NAS) students; for homogeneous models and VLMs, omit it -- it defaults to `--student_hf_path`.

**Separate conversion** -- convert **any** saved iteration (intermediate or final) with [export_distilled_megatron_to_hf.py](export_distilled_megatron_to_hf.py):

```bash
torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \
    --student_hf_path <student_hf_model_or_path> \
    --megatron_path <distill_output_dir>/checkpoints/iter_<iter_number> \
    --hf_export_path /path/to/save/distilled_hf_ckpt
```

Use `--export_iterations` to export multiple saved checkpoints, for example to evaluate how model
quality changes during distillation. To export multiple iterations, keep those Megatron checkpoints
during distillation. The default is to keep the last 5 checkpoints; set `--checkpoint_keep_last -1`
to keep all saved checkpoints.

Then export all retained checkpoints, with one Hugging Face checkpoint written per
`iter_<iteration>` subdirectory:

```bash
torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \
    --student_hf_path <student_hf_model_or_path> \
    --megatron_path <distill_output_dir>/checkpoints \
    --hf_export_path /path/to/save/hf_validation_checkpoints \
    --export_iterations all
```

The export path contains one loadable Hugging Face checkpoint per exported iteration:

```text
hf_validation/
├── iter_0000100/
├── iter_0000200/
└── iter_0000300/
```

To export selected iterations instead, use `--export_iterations 200 400 600`.

### Quantization Aware Distillation (QAD)

To recover the accuracy lost during [Post-Training Quantization](#post-training-quantization), distill the quantized model (student) from the original, unquantized model (teacher). Pass the quantized **Megatron checkpoint** produced by `quantize.py` via `--student_megatron_path` (the ModelOpt quantizers are restored automatically, so distillation trains the fake-quantized student), while `--student_hf_path` provides the student architecture and `--teacher_hf_path` points to the original unquantized model.

If you do not already have a suitable QAD dataset, start with
[data/nemotron-cascade-2-blend.yaml](data/nemotron-cascade-2-blend.yaml). It defines a general-purpose
mixture of SFT data for QAD. Copy it, set the tokenizer for the target model, and adjust the output directory,
sources, and weights as needed before preparing data. Its default 17.3-billion-token budget covers 1000
iterations at global batch size 512 and sequence length 32768, including a 1% validation holdout and margin.
Recalculate the budget when changing those settings, and keep the prepared data unchanged when resuming.

We also use a smaller learning rate for QAD:

```bash
torchrun --nproc_per_node 8 distill.py \
    --tp_size 8 \
    --teacher_hf_path Qwen/Qwen3-8B \
    --student_hf_path Qwen/Qwen3-8B \
    --student_megatron_path /tmp/Qwen3-8B-NVFP4-megatron \
    --data_paths 1.0 tokenized_qwen3/data1_text_document 1.0 tokenized_qwen3/data2_text_document \
    --data_path_to_cache /path/to/cache/dataset_indices_qwen3 \
    --seq_length 8192 \
    --gbs 768 \
    --train_iters 1000 \
    --lr 1e-5 \
    --min_lr 5e-6 \
    --output_dir /output/qwen3_8b_nvfp4_qad
```

The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) (point `--megatron_path` at `/output/qwen3_8b_nvfp4_qad/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization).

### Slurm Usage

To run the distillation script on a Slurm cluster for multi-node training, you just need use `python` instead of `torchrun` and set the number of nodes using `#SBATCH --nodes=<num_nodes>` clause in your Slurm script.

### Distillation Results

See [examples/pruning/](../pruning/README.md#tutorials--results) for distillation experiment results covering Minitron and Puzzletron pruning algorithms.

## Pruning

This section shows how to prune a HuggingFace model using Minitron algorithm in Megatron-Bridge framework. Checkout other available pruning algorithms, supported frameworks and models, and general pruning getting-started in the [pruning README](../pruning/README.md).

The script supports three NAS-based pruning targets and one manual export mode:

| Mode | Flag | Description |
| :---: | :---: | :--- |
| NAS | `--prune_target_params` | Prune to a target total parameter count |
| NAS | `--prune_target_active_params` | Prune to a target active parameter count (useful for MoE models). For non-MoE models, this is equivalent to `--prune_target_params`. |
| NAS | `--prune_target_memory_mb` | Prune to a target memory footprint in MB (weights + KV-cache) for a given batch size and sequence length assuming BF16 precision |
| Manual | `--prune_export_config` | Prune directly to a specified architecture config (no NAS). Useful if you want to take top K candidates and do a short knowledge distillation before selecting the best model. |

Multiple NAS targets can be combined — e.g. `--prune_target_params 6e9 --prune_target_memory_mb 12288` finds the best model with under 6B params and under 12GB memory footprint at (default) batch size 1 and sequence length 4096 assuming BF16 precision.

**Prune by total parameter count** — prune Qwen3-8B to 6B on 2-GPUs (Pipeline Parallelism = 2) while skipping pruning of `num_attention_heads` using following defaults:
    1024 samples from [`nemotron-post-training-dataset-v2`](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2) for calibration,
    at-most 20% depth (`num_layers`) and 40% width is pruned per prunable hparam (`hidden_size`, `ffn_hidden_size`, ...),
    top-10 candidates are evaluated for MMLU score (5% sampled data) to select the best model.

```bash
torchrun --nproc_per_node 2 prune_minitron.py \
    --pp_size 2 \
    --hf_model_name_or_path Qwen/Qwen3-8B \
    --prune_target_params 6e9 \
    --hparams_to_skip num_attention_heads \
    --output_hf_path /tmp/Qwen3-8B-Pruned-6B
```

**Prune by active parameter count** — useful for MoE models where most experts are inactive per token (e.g. prune Nemotron-3-Nano-30B-A3B-BF16 (3.6B active params) to 3B active params):

```bash
torchrun --nproc_per_node 2 prune_minitron.py \
    --pp_size 2 \
    --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
    --prune_target_active_params 3e9 \
    --output_hf_path /tmp/Nemotron-3-Nano-30B-A3B-BF16-Pruned-3B-Active
```

**Prune by memory footprint** — prune to fit a target GPU memory budget (weights + KV-cache at the given sequence length and batch size, assuming BF16):

```bash
torchrun --nproc_per_node 2 prune_minitron.py \
    --pp_size 2 \
    --hf_model_name_or_path Qwen/Qwen3-8B \
    --prune_target_memory_mb 12288 \
    --seq_length 4096 \
    --calib_batch_size 1 \
    --output_hf_path /tmp/Qwen3-8B-Pruned-12GB
```

**Manual pruning** — prune directly to a specified architecture (no NAS, no score evaluation):

```bash
torchrun --nproc_per_node 2 prune_minitron.py \
    --pp_size 2 \
    --hf_model_name_or_path Qwen/Qwen3-8B \
    --prune_export_config '{"hidden_size": 3584, "ffn_hidden_size": 9216}' \
    --output_hf_path /tmp/Qwen3-8B-Pruned-6B-manual
```

To see the full usage for advanced configurations, run:

```bash
torchrun --nproc_per_node 1 prune_minitron.py --help
```

> [!TIP]
> If number of layers in the model is not divisible by number of GPUs i.e. pipeline parallel (PP) size, you can configure
> 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]
> Multi-token-prediction (MTP) heads (e.g. Qwen3.5) are not pruned yet — they are dropped for the prune run and the saved checkpoint has no MTP. Autoregressive inference is unaffected; for speculative decoding, run a short MTP SFT on the pruned model.

### Vision-Language Models (VLMs)

For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `prune_minitron.py` automatically prunes only the **language model** and leaves the vision tower intact, then saves the full VLM back. All the pruning modes above (parameter count, active parameter count, memory footprint, and manual `export_config`) work unchanged, with two VLM-specific caveats:

- The `--prune_target_params` / `--prune_target_active_params` / `--prune_target_memory_mb` targets (and `export_config` dimensions) apply to the **language model only** — the (unpruned) vision tower's parameters are *not* counted, so the full saved VLM will be larger than the target.
- `hidden_size` is never pruned for VLMs (it is shared with the vision projector).

```bash
torchrun --nproc_per_node 2 prune_minitron.py \
    --pp_size 2 \
    --hf_model_name_or_path Qwen/Qwen3.5-4B \
    --prune_target_params 3e9 \
    --output_hf_path /tmp/Qwen3.5-4B-Pruned-3B
```

## Sanity-Check Generation

[generate_vllm.py](generate_vllm.py) runs a quick generation check on an exported HuggingFace checkpoint using vLLM — a useful smoke test for a **quantized**, **pruned**, or **distilled** model to confirm it still produces coherent text. For quantized checkpoints, vLLM auto-detects the ModelOpt quantization from the exported `hf_quant_config.json`, so no extra flags are needed:

```bash
# Quantized model
python generate_vllm.py --model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 --trust_remote_code

# Pruned model
python generate_vllm.py --model /tmp/Qwen3-8B-Pruned-6B
```

> [!NOTE]
> `--trust_remote_code` is only needed for models that ship custom modeling code (e.g. Nemotron); Qwen models don't require it.

## Resources

- 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/1699)
- 📖 [Documentation](https://nvidia.github.io/Model-Optimizer)
- 💡 [Release Notes](https://nvidia.github.io/Model-Optimizer/reference/0_changelog.html)
- 🐛 [File a bug](https://github.com/NVIDIA/Model-Optimizer/issues/new?template=1_bug_report.md)
- ✨ [File a Feature Request](https://github.com/NVIDIA/Model-Optimizer/issues/new?template=2_feature_request.md)