[2/2] Export vLLM fake-quant checkpoints through Megatron-Bridge

.github/actions/megatron-from-source/action.yml

@@ -0,0 +1,27 @@
+name: Megatron-Bridge from source
+description: >-
+  Replace the NeMo container's Megatron-Bridge and Megatron-Core sources with a Megatron-Bridge commit and
+  its pinned Megatron-LM, so CI covers features newer than the container. Remove once the container has them.
+
+inputs:
+  commit:
+    description: Megatron-Bridge commit. Keep in sync with examples/megatron_bridge/README.md.
+    default: 1d9caf7324e608aca7181de40be26a6387976210
+
+runs:
+  using: composite
+  steps:
+    - shell: bash
+      run: |
+        set -euxo pipefail
+        src=$(mktemp -d)
+        git -C "$src" init -q
+        git -C "$src" fetch -q --depth 1 https://github.com/NVIDIA-NeMo/Megatron-Bridge.git "${{ inputs.commit }}"
+        git -C "$src" checkout -q FETCH_HEAD
+        git -C "$src" submodule update -q --init --depth 1 3rdparty/Megatron-LM
+        # Both are editable installs in the NeMo container, so replacing the sources is enough.
+        for dir in src/megatron/bridge 3rdparty/Megatron-LM/megatron; do
+          rm -rf "/opt/Megatron-Bridge/$dir"
+          cp -r "$src/$dir" "/opt/Megatron-Bridge/$dir"
+        done
+        rm -rf "$src"

.github/workflows/_example_tests_runner.yml

@@ -64,6 +64,8 @@ jobs:
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV
           echo "PATH=${PATH}:/usr/local/tensorrt/targets/x86_64-linux-gnu/bin" >> $GITHUB_ENV
+      - if: contains(inputs.docker_image, 'nvidia/nemo:')
+        uses: ./.github/actions/megatron-from-source
       - name: Install dependencies
         shell: bash
         run: |

.github/workflows/gpu_tests.yml

@@ -90,6 +90,8 @@ jobs:
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV
+      - if: contains(matrix.container_image, 'nvidia/nemo:')
+        uses: ./.github/actions/megatron-from-source
       - name: Run gpu tests
         env:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml

CHANGELOG.rst

@@ -23,7 +23,7 @@ Changelog
 - Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences.
 - Add ``--document_boundaries`` to ``examples/megatron_bridge/distill.py`` to attend within each document of a pre-tokenized window, and ``tokenizer`` / ``max_length`` / ``pack_max_tokens`` to ``chat_jsonl`` blends to drop conversations with nothing to train and pair the rest for ``--sft_in_batch_packing``. Packed distillation of Gated DeltaNet models needs ``flash-linear-attention`` >= 0.5.2 (earlier versions compute it wrongly); install ``examples/megatron_bridge/requirements.txt`` in ``nemo:26.08``.
 - Add ``--num_layers_in_first_pipeline_stage``, ``--num_layers_in_last_pipeline_stage`` and ``--recompute_num_layers_per_stage`` to ``examples/megatron_bridge/distill.py`` to balance pipeline-parallel distillation: give the last stage, which also runs the output layers and KD loss, fewer layers, and recompute fewer layers on later stages.
-- Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``.
+- Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate PTQ / QAD checkpoints in vLLM for formats it has no kernels for (e.g. NVFP4 on Hopper, GGML IQ): it folds the weight quantizers into BF16 weights and saves through Megatron-Bridge at any TP / EP / PP, LLMs and VLMs, with the activation and KV-cache quantizers in sidecars that ``examples/vllm_serve/vllm_serve_fakequant.py`` loads automatically. Also allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``.
 - GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models.
 - The Megatron unified HF export (``export_mcore_gpt_to_hf``) now copies the source checkpoint's non-model files (tokenizer, processor, remote code, chat templates, ...) into the export unchanged, from local disk only. The ``examples/megatron_bridge`` scripts make a Hugging Face Hub model local first with ``modelopt.torch.export.ensure_local_checkpoint``.
 
@@ -45,6 +45,7 @@ Changelog
 - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root.
 - Fix GPT-OSS draft training with pre-normalization teacher hidden states by enabling the target model's final normalization.
 - Fix vanilla DSpark exports to use the Qwen3 DSpark architecture, head names and configuration expected by vLLM. Re-export existing training checkpoints to use this format.
+- ``mtq.fold_weight`` now folds sequential weight quantizers (e.g. W4A8 INT4 + FP8 weights) into the weights; it used to leave them unfolded and enabled.
 - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``.
 - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source.
 - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU.

examples/megatron_bridge/README.md

@@ -46,6 +46,22 @@ docker run \
   ${DOCKER_IMAGE} bash
 ```
 
+Some features (e.g. [`--vllm_fakequant_export`](#evaluating-a-qad-checkpoint-in-vllm-with-fake-quantization)) need a newer Megatron-Bridge than the one in `nvcr.io/nvidia/nemo:26.08`. Until a NeMo container includes it, check out Megatron-Bridge at the commit ModelOpt's CI tests against, together with its Megatron-LM submodule:
+
+```bash
+export MBRIDGE_DIR=${PWD}/Megatron-Bridge
+git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git ${MBRIDGE_DIR}
+git -C ${MBRIDGE_DIR} checkout 1d9caf7324e608aca7181de40be26a6387976210
+git -C ${MBRIDGE_DIR} submodule update --init 3rdparty/Megatron-LM
+```
+
+and add these mounts to the `docker run` command above:
+
+```bash
+  -v ${MBRIDGE_DIR}/src/megatron/bridge:/opt/Megatron-Bridge/src/megatron/bridge \
+  -v ${MBRIDGE_DIR}/3rdparty/Megatron-LM/megatron:/opt/Megatron-Bridge/3rdparty/Megatron-LM/megatron \
+```
+
 > [!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.
 
@@ -371,6 +387,27 @@ sequence length, and GPU memory: measure per-stage time and memory over a few it
 
 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).
 
+#### Evaluating a QAD checkpoint in vLLM with fake quantization
+
+To evaluate a format vLLM has no kernels for (e.g. NVFP4 on Hopper, GGML IQ2, or a custom research format), export with `--vllm_fakequant_export`. The weight quantizers are folded into BF16 weights and Megatron-Bridge saves the model at any TP / EP / PP, VLMs included. The enabled activation and KV-cache quantizers go to `quantizer_state.pth` (ranges) and `quant_recipe.yaml` (settings), both under HF names:
+
+```bash
+torchrun --nproc_per_node 8 export_quantized_megatron_to_hf.py \
+    --hf_model_name_or_path Qwen/Qwen3-8B \
+    --megatron_path /output/qwen3_8b_nvfp4_qad/checkpoints \
+    --tp_size 8 \
+    --vllm_fakequant_export \
+    --export_unified_hf_path /tmp/Qwen3-8B-NVFP4-QAD-vllm-fq
+```
+
+Serve the export with [vllm_serve_fakequant.py](../vllm_serve/vllm_serve_fakequant.py), which loads both sidecars automatically, then run your evaluations against the server:
+
+```bash
+python ../vllm_serve/vllm_serve_fakequant.py /tmp/Qwen3-8B-NVFP4-QAD-vllm-fq -tp 8 --port 8000
+```
+
+Weight-only formats leave no quantizer enabled after folding, so no sidecars are written and stock `vllm serve` loads the export as a plain BF16 checkpoint. This mode needs a Megatron-Bridge that writes the quantizer sidecars; see [Pre-Requisites](#pre-requisites).
+
 ### LoRA QAD
 
 Add `--lora_rank` to PTQ to train low-rank weight updates with the usual QAD flow:

examples/megatron_bridge/export_quantized_megatron_to_hf.py

@@ -33,12 +33,28 @@
         --pp_size 2 \
         --export_unified_hf_path /tmp/Qwen3-8B-FP8-hf
 
+With --vllm_fakequant_export, the weight quantizers are instead folded into BF16 weights and the
+model is saved through Megatron-Bridge (any TP / EP / PP, LLMs and VLMs). The enabled activation and
+KV-cache quantizers go to quantizer_state.pth and quant_recipe.yaml sidecars, which
+examples/vllm_serve/vllm_serve_fakequant.py loads to evaluate the checkpoint with vLLM. Formats
+without them (e.g. weight-only GGML IQ2) need no sidecars and serve with stock vLLM:
+
+    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-QAD/checkpoints \
+        --tp_size 2 \
+        --vllm_fakequant_export \
+        --export_unified_hf_path /tmp/Qwen3-8B-NVFP4-QAD-vllm-fq
+
 See `README.md` in this directory for more details.
 """
 
 import argparse
+import json
+from pathlib import Path
 
 import torch
+from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
 from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
 from megatron.core.utils import unwrap_model
 from mlflow_utils import (
@@ -53,10 +69,13 @@
 import modelopt.torch.quantization as mtq
 import modelopt.torch.utils.distributed as dist
 from modelopt.torch.export import ensure_local_checkpoint, export_mcore_gpt_to_hf
-from modelopt.torch.export.plugins.vllm_fakequant_megatron import export_mcore_gpt_to_hf_vllm_fq
+from modelopt.torch.export.plugins.vllm_fakequant_megatron import (
+    check_mcore_vllm_fakequant_quantizers,
+)
 from modelopt.torch.utils import print_args, print_rank_0
 from modelopt.torch.utils.mlflow import Tool, masked_args
 from modelopt.torch.utils.plugins.mbridge import (
+    is_vlm_config,
     load_mbridge_model_from_hf,
     load_modelopt_megatron_checkpoint,
     use_moe_grouped_gemm,
@@ -111,9 +130,10 @@ def get_args() -> argparse.Namespace:
         "--vllm_fakequant_export",
         action="store_true",
         help=(
-            "Fold the weight quantizers into BF16 weights, for accuracy evaluation in a stock vLLM "
-            "of formats it has no kernels for (e.g. GGML IQ2). Activation quantizers go to a "
-            "quantizer_state.pth sidecar."
+            "Fold the weight quantizers into BF16 weights and save through Megatron-Bridge, for "
+            "accuracy evaluation in vLLM of formats it has no kernels for (e.g. NVFP4 QAD on "
+            "Hopper, GGML IQ2). Enabled activation and KV-cache quantizers go to "
+            "quantizer_state.pth and quant_recipe.yaml sidecars for vllm_serve_fakequant.py."
         ),
     )
     parser.add_argument(
@@ -122,7 +142,12 @@ def get_args() -> argparse.Namespace:
         help="Export extra modules such as Medusa heads, EAGLE, or MTP.",
     )
 
-    # Only pipeline and expert parallelism are supported for export
+    parser.add_argument(
+        "--tp_size",
+        type=int,
+        default=1,
+        help="Tensor parallel size (--vllm_fakequant_export only; the unified exporter needs 1)",
+    )
     parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size")
     parser.add_argument(
         "--ep_size",
@@ -157,6 +182,63 @@ def get_args() -> argparse.Namespace:
     return args
 
 
+def export_vllm_fakequant(bridge, model: list, export_dir: str) -> None:
+    """Fold the weight quantizers into BF16 weights and save them with Megatron-Bridge.
+
+    Megatron-Bridge writes the enabled quantizers' ranges and settings under HF names to
+    ``quantizer_state.pth`` and ``quant_recipe.yaml``, the sidecars ``vllm_serve_fakequant.py``
+    loads; it writes none when no quantizer is left enabled (e.g. weight-only formats).
+    """
+    if not hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe"):
+        raise RuntimeError(
+            "--vllm_fakequant_export needs a Megatron-Bridge that writes the quantizer sidecars "
+            "(quantizer_state.pth and quant_recipe.yaml); see the Pre-Requisites in "
+            "examples/megatron_bridge/README.md."
+        )
+    for chunk in model:
+        check_mcore_vllm_fakequant_quantizers(unwrap_model(chunk))
+        mtq.fold_weight(unwrap_model(chunk))
+    bridge.save_hf_pretrained(model, export_dir)
+    if dist.is_master():
+        # The weights are BF16, so a pre-quantized source's quantization config would misdescribe them.
+        config_path = Path(export_dir) / "config.json"
+        config = json.loads(config_path.read_text())
+        if config.pop("quantization_config", None) is not None:
+            config_path.write_text(json.dumps(config, indent=2))
+        (Path(export_dir) / "hf_quant_config.json").unlink(missing_ok=True)
+    dist.barrier()
+
+
+def export_unified_hf(unwrapped_model, args: argparse.Namespace, trust_remote_code: bool) -> None:
+    """Export the quantized model to a unified HF checkpoint with ModelOpt's Megatron exporter."""
+    # Extra modules (Medusa / EAGLE / MTP) only exist on the last pipeline stage. Use an all-reduce
+    # MAX over all ranks (rather than a broadcast from a hard-coded source rank) so the decision is
+    # correct regardless of pipeline placement / global rank ordering.
+    has_extra_modules = hasattr(unwrapped_model, "eagle_module") or hasattr(
+        unwrapped_model, "medusa_heads"
+    )
+    if torch.distributed.is_initialized():
+        flag = torch.tensor(
+            [int(has_extra_modules)], dtype=torch.int, device=torch.cuda.current_device()
+        )
+        torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MAX)
+        has_extra_modules = bool(flag.item())
+    export_extra_modules = has_extra_modules and args.export_extra_modules
+
+    print_rank_0(
+        f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..."
+    )
+    # TODO: Gemma3-VL is not in export_mcore_gpt_to_hf's per-arch mappings yet.
+    export_mcore_gpt_to_hf(
+        unwrapped_model,
+        args.local_checkpoint_path,
+        export_extra_modules=export_extra_modules,
+        dtype=torch.bfloat16,
+        export_dir=args.export_unified_hf_path,
+        trust_remote_code=trust_remote_code,
+    )
+
+
 def main(args: argparse.Namespace):
     # --hf_model_name_or_path stays as given. Later steps read the local copy of the whole
     # checkpoint in local_checkpoint_path, since exporters read the source's files from local disk
@@ -167,9 +249,13 @@ def main(args: argparse.Namespace):
     trust_remote_code = is_safe_repo(
         trust_remote_code=args.trust_remote_code, hf_path=args.local_checkpoint_path
     )
+    if not args.vllm_fakequant_export and args.tp_size != 1:
+        raise ValueError("--tp_size > 1 needs --vllm_fakequant_export.")
+    if args.vllm_fakequant_export and args.export_extra_modules:
+        raise ValueError("--export_extra_modules is not supported with --vllm_fakequant_export.")
 
     # Build the model structure from HF
-    _bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf(
+    bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf(
         hf_model_name_or_path=args.local_checkpoint_path,
         trust_remote_code=trust_remote_code,
         moe_grouped_gemm=use_moe_grouped_gemm(
@@ -178,7 +264,7 @@ def main(args: argparse.Namespace):
             force_sequential=args.no_moe_grouped_gemm,
         ),
         provider_overrides={
-            "tensor_model_parallel_size": 1,  # Tensor parallelism is not supported
+            "tensor_model_parallel_size": args.tp_size,
             "pipeline_model_parallel_size": args.pp_size,
             "expert_model_parallel_size": args.ep_size,
             "expert_tensor_parallel_size": 1,  # Expert tensor parallelism is not supported
@@ -187,7 +273,11 @@ def main(args: argparse.Namespace):
             "pipeline_dtype": torch.bfloat16,
         },
         init_model_parallel=True,
-        load_weights=False,  # The weights come from the Megatron checkpoint, so HF weights are not loaded
+        # The weights come from the Megatron checkpoint. Megatron-Bridge saves the whole in-memory
+        # model, and a VLM's QAD checkpoint holds only the language model, so load the HF vision
+        # tower first.
+        load_weights=args.vllm_fakequant_export
+        and is_vlm_config(args.local_checkpoint_path, trust_remote_code=trust_remote_code),
     )
 
     # Load the quantized checkpoint (with the correct layer spec) rather than reconstructing it from the checkpoint
@@ -197,35 +287,11 @@ def main(args: argparse.Namespace):
     unwrapped_model = unwrap_model(model[0])
     mtq.merge_quant_lora(unwrapped_model)
 
-    # Extra modules (Medusa / EAGLE / MTP) only exist on the last pipeline stage. Use an all-reduce
-    # MAX over all ranks (rather than a broadcast from a hard-coded source rank) so the decision is
-    # correct regardless of pipeline placement / global rank ordering.
-    has_extra_modules = hasattr(unwrapped_model, "eagle_module") or hasattr(
-        unwrapped_model, "medusa_heads"
-    )
-    if torch.distributed.is_initialized():
-        flag = torch.tensor(
-            [int(has_extra_modules)], dtype=torch.int, device=torch.cuda.current_device()
-        )
-        torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MAX)
-        has_extra_modules = bool(flag.item())
-    export_extra_modules = has_extra_modules and args.export_extra_modules
-
-    print_rank_0(
-        f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..."
-    )
-    # TODO: Gemma3-VL is not in export_mcore_gpt_to_hf's per-arch mappings yet.
-    export_fn = (
-        export_mcore_gpt_to_hf_vllm_fq if args.vllm_fakequant_export else export_mcore_gpt_to_hf
-    )
-    export_fn(
-        unwrapped_model,
-        args.local_checkpoint_path,
-        export_extra_modules=export_extra_modules,
-        dtype=torch.bfloat16,
-        export_dir=args.export_unified_hf_path,
-        trust_remote_code=trust_remote_code,
-    )
+    if args.vllm_fakequant_export:
+        print_rank_0(f"Exporting vLLM fake-quant checkpoint to {args.export_unified_hf_path}...")
+        export_vllm_fakequant(bridge, model, args.export_unified_hf_path)
+    else:
+        export_unified_hf(unwrapped_model, args, trust_remote_code)
     args.checkpoint_exported = True
     print_rank_0(f"Exported HuggingFace checkpoint to {args.export_unified_hf_path}")