Megatron Bridge

This directory contains examples of using Model Optimizer with the NeMo Megatron-Bridge framework for quantization, pruning, and distillation. These workflows can be used on their own or combined.

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.06) which has all the dependencies installed.

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

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.06
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

-v ${MEGATRON_BRIDGE_SRC_DIR}:/opt/Megatron-Bridge

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, which is gated.

hf auth login --token <your token>

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 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 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_sq, int4_awq, w4a8_awq, ...). 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 and nemotron-post-training-dataset-v2) for calibration (sequence length = 4096):

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

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

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

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).

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 / Post-Training Quantization to recover accuracy of the model by distilling from the original model (teacher).

The distill.py script supports both standard HuggingFace checkpoints and Puzzletron AnyModel 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 for Quantization Aware Distillation, 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 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.

Distillation with Real Data

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

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 see all available arguments:

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):

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).

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 instead — see QAD):

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

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:

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

Quantization Aware Distillation (QAD)

To recover the accuracy lost during 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. We also use a smaller learning rate for QAD:

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 (point --megatron_path at /output/qwen3_8b_nvfp4_qad/checkpoints), exactly like the PTQ checkpoint in step 2 above.

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/ 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.

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

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 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.

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):

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):

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):

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:

torchrun --nproc_per_node 1 prune_minitron.py --help

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).

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 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:

# 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

Resources

• 📅 Roadmap
• 📖 Documentation
• 💡 Release Notes
• 🐛 File a bug
• ✨ File a Feature Request