From b190d284b3d1e2a3b23cd4cdb508f5bc49b431b8 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:47:21 -0700 Subject: [PATCH 01/11] Export vLLM fake-quant checkpoints through Megatron-Bridge export_quantized_megatron_to_hf.py --vllm_fakequant_export now checks the quantizers, folds the weight quantizers into BF16 weights and saves with bridge.save_hf_pretrained, which writes the enabled activation and KV-cache quantizers to quantizer_state.pth and a versioned quant_recipe.yaml under HF names for vllm_serve_fakequant.py. - Any TP / EP / PP (new --tp_size, fake-quant only) and every architecture Megatron-Bridge maps; a VLM's language-model-only QAD checkpoint keeps the HF vision tower. - Weight-only formats (e.g. GGML IQ2) get no sidecars and serve with stock vLLM. - check_mcore_vllm_fakequant_quantizers shares the existing exporter's checks. - Tests: parity with export_mcore_gpt_to_hf_vllm_fq (tiny Qwen3, Nemotron-H), export parallelism invariance (Qwen3.5-VL-MoE EP, Nemotron-H PP), and a fake-quant export step in the QAD example test. They skip without a Megatron-Bridge that writes the sidecars. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 2 +- examples/megatron_bridge/README.md | 21 ++++ .../export_quantized_megatron_to_hf.py | 87 ++++++++++++--- examples/vllm_serve/README.md | 1 + .../export/plugins/vllm_fakequant_megatron.py | 33 +++++- tests/examples/megatron_bridge/test_qad.py | 33 ++++++ .../megatron_bridge/test_quantize_export.py | 90 +++++++++++++++- .../test_vllm_fakequant_mbridge_export.py | 100 ++++++++++++++++++ 8 files changed, 352 insertions(+), 15 deletions(-) create mode 100644 tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1b4c6cbe2be..8145bfa1ca0 100755 --- a/CHANGELOG.rst +++ b/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``. diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 66683526675..78264741691 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -371,6 +371,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. + ### LoRA QAD Add `--lora_rank` to PTQ to train low-rank weight updates with the usual QAD flow: diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index 08a58102ce7..2e996d37024 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/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,32 @@ 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); update Megatron-Bridge." + ) + 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 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 +218,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 +233,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 +242,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,6 +256,13 @@ def main(args: argparse.Namespace): unwrapped_model = unwrap_model(model[0]) mtq.merge_quant_lora(unwrapped_model) + 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) + args.checkpoint_exported = True + print_rank_0(f"Exported vLLM fake-quant checkpoint to {args.export_unified_hf_path}") + return + # 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. @@ -215,10 +281,7 @@ def main(args: argparse.Namespace): 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( + export_mcore_gpt_to_hf( unwrapped_model, args.local_checkpoint_path, export_extra_modules=export_extra_modules, diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index b958646a140..93a6aae13c9 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -258,6 +258,7 @@ python3 examples/hf_ptq/hf_ptq.py \ Note: `--pyt_ckpt_path` can point to either an HF checkpoint or a ModelOpt-saved checkpoint (e.g., a QAT/QAD checkpoint produced by `examples/llm_qat/train.py`). If the input checkpoint is already quantized, the script will **skip re-quantization** and only export artifacts for vLLM fakequant reload. +- For **Megatron-Bridge** models, export with `export_quantized_megatron_to_hf.py --vllm_fakequant_export` as described in the [Megatron-Bridge example README](../megatron_bridge/README.md#evaluating-a-qad-checkpoint-in-vllm-with-fake-quantization). This generates `quant_recipe.yaml` (versioned, one ModelOpt quantizer state per HF quantizer name) and `quantizer_state.pth` for vLLM fakequant reload. - For **MCore** models, export the model with flag `--export-vllm-fq` as described in [Megatron-LM README](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/post_training/modelopt#-nvfp4-quantization-qauntization-aware-training-and-model-export). This generates `quant_recipe.yaml` and `quantizer_state.pth` for vLLM fakequant reload. Step 2: use the exported artifacts when serving: diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 30723624739..e063f849749 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -15,6 +15,7 @@ """Export Megatron Core Model to HuggingFace vLLM fakequant checkpoint.""" import os +import re import tempfile from collections.abc import Callable, Mapping from contextlib import contextmanager @@ -35,7 +36,7 @@ from modelopt.torch.utils import get_unwrapped_name from modelopt.torch.utils.distributed import DistributedProcessGroup, is_master -__all__ = ["export_mcore_gpt_to_hf_vllm_fq"] +__all__ = ["check_mcore_vllm_fakequant_quantizers", "export_mcore_gpt_to_hf_vllm_fq"] def _quantizer_configs(module: torch.nn.Module) -> tuple[dict[str, dict], str]: @@ -97,6 +98,36 @@ def _quantizer_configs(module: torch.nn.Module) -> tuple[dict[str, dict], str]: return configs, "" +def check_mcore_vllm_fakequant_quantizers(model: torch.nn.Module) -> None: + """Raise if vLLM fake-quant reload cannot reproduce a quantizer of a Megatron model. + + Checks the settings the vLLM fake-quant export and reload support, with weight quantizers + folded into the weights. Collective: a failure on any rank raises on every rank. + """ + _, failure = _quantizer_configs(model) + tied = any(getattr(m, "share_embeddings_and_output_weights", False) for m in model.modules()) + for name, quantizer in model.named_modules(): + if failure: + break + if not isinstance(quantizer, TensorQuantizer): + continue + active = quantizer.is_enabled and quantizer._if_quant + transform = quantizer.rotate_is_enabled or quantizer.pre_quant_scale is not None + if re.search(r"(^|\.)mtp\.", name) and (active or transform): + failure = f"MTP quantization is not supported by vLLM fakequant export/reload: {name}" + elif re.search(r"(^|\.)output_layer\.", name) and ( + (active and (tied or "weight_quantizer" not in name)) or (tied and transform) + ): + failure = f"Unsupported lm_head quantization for vLLM fakequant: {name}" + failure = DistributedProcessGroup.get_dist_syncd_obj( + failure, + DistributedProcessGroup(None), + lambda failures: next((message for message in failures if message), ""), + ) + if failure: + raise ValueError(failure) + + def _save_quantizer_state(path: Path, save: Callable[[Path], object]) -> None: """Save quantizer state or recipes and share write completion or failure across ranks.""" failure = "" diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index b3df875f3ce..d25cd8e0c0e 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -21,6 +21,7 @@ import pytest import torch import torch.distributed.checkpoint as dcp +import yaml from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.export.unified_checkpoint import ( assert_exported_checkpoint_matches, @@ -34,8 +35,12 @@ create_tiny_qwen3_5_moe_vl_dir, create_tiny_qwen3_dir, ) +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from torch.distributed.checkpoint import FileSystemReader +# TODO: drop this probe once the minimum Megatron-Bridge writes the ModelOpt quantizer sidecars. +HAS_BRIDGE_QUANTIZER_SIDECARS = hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe") + # The fix ships in the nemo:26.10 container. # TODO(Megatron-Bridge#6243): drop this probe once the minimum Megatron-Bridge carries it. try: @@ -244,3 +249,31 @@ def test_qad(tmp_path: Path, num_gpus, create_student, lora_rank): if is_moe: # allow_unexpected waives the expert names wholesale; this re-tightens it. assert_per_expert_experts_complete(hf_export_path) + + # Step 4: export for vLLM fake-quant evaluation. Megatron-Bridge writes the whole model, so + # a VLM's language-model-only QAD checkpoint must still come out with the original vision tower. + if HAS_BRIDGE_QUANTIZER_SIDECARS: + fakequant_path = tmp_path / "qad_vllm_fakequant" + fakequant_cmd = extend_cmd_parts( + [ + "torchrun", + f"--nproc_per_node={num_gpus}", + "export_quantized_megatron_to_hf.py", + "--vllm_fakequant_export", + ], + hf_model_name_or_path=hf_model_path, + megatron_path=distilled_megatron_path, + export_unified_hf_path=fakequant_path, + tp_size=num_gpus, + ) + run_example_command(fakequant_cmd, example_path="megatron_bridge", setup_free_port=True) + assert_exported_checkpoint_matches( + fakequant_path, + hf_model_path, + check_values=False, + bit_exact_prefixes=("model.visual.",) if is_vlm else (), + ) + recipe = yaml.safe_load((fakequant_path / "quant_recipe.yaml").read_text()) + state = torch.load(fakequant_path / "quantizer_state.pth", weights_only=True) + assert any(name.endswith("k_bmm_quantizer") for name in recipe["quantizers"]) + assert {key.rpartition(".")[0] for key in state} <= recipe["quantizers"].keys() diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index a60fb417acd..770e4d8fdfc 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -14,17 +14,25 @@ # limitations under the License. """Tests for quantize.py and export_quantized_megatron_to_hf.py scripts.""" +from functools import partial from pathlib import Path import pytest +import torch +import yaml from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches +from _test_utils.torch.export.unified_checkpoint import ( + assert_exported_checkpoint_matches, + load_safetensors_dir, +) from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state from _test_utils.torch.transformers_models import ( create_tiny_nemotron_h_dir, + create_tiny_qwen3_5_moe_vl_dir, create_tiny_qwen3_moe_dir, create_tiny_qwen3vl_dir, ) +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge # Per-architecture export *mappings* are covered in-process by # tests/gpu_megatron/torch/export/test_unified_export_megatron.py; these cases cover the script @@ -134,3 +142,83 @@ def test_quantize_and_export( # ) # outputs = llm.generate(["Hello!"], vllm.SamplingParams(max_tokens=4)) # assert outputs and outputs[0].outputs and outputs[0].outputs[0].text + + +@pytest.mark.skipif( + not hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe"), + reason="Megatron-Bridge does not write the ModelOpt quantizer sidecars", +) +@pytest.mark.parametrize( + ("create_model", "export_parallelism"), + [ + ( + partial( + create_tiny_qwen3_5_moe_vl_dir, + with_processor=True, + num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], + ), + "ep_size", + ), + (partial(create_tiny_nemotron_h_dir, with_tokenizer=True), "pp_size"), + ], + ids=["qwen3_5_moe_vl", "nemotron_h"], +) +@pytest.mark.timeout(480) +def test_quantize_and_vllm_fakequant_export( + tmp_path: Path, num_gpus, create_model, export_parallelism +): + """Fake-quant exports through Megatron-Bridge do not depend on the export parallelism.""" + hf_model_path = create_model(tmp_path) + hf_model_path = hf_model_path[0] if isinstance(hf_model_path, tuple) else hf_model_path + megatron_path = tmp_path / "quantized_megatron" + quantize_cmd = extend_cmd_parts( + ["torchrun", "--nproc_per_node=1", "quantize.py", "--skip_generate"], + hf_model_name_or_path=hf_model_path, + recipe="general/ptq/nvfp4_default-kv_fp8", + calib_dataset_name="cnn_dailymail", + calib_num_samples=4, + calib_batch_size=2, + seq_length=16, + export_megatron_path=megatron_path, + ) + run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True) + + exports = [] + for size in sorted({1, num_gpus}): + export_path = tmp_path / f"vllm_fakequant_{size}" + export_cmd = extend_cmd_parts( + [ + "torchrun", + f"--nproc_per_node={size}", + "export_quantized_megatron_to_hf.py", + "--vllm_fakequant_export", + ], + hf_model_name_or_path=hf_model_path, + megatron_path=megatron_path, + export_unified_hf_path=export_path, + **{export_parallelism: size}, + ) + run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True) + exports.append( + ( + load_safetensors_dir(export_path), + torch.load(export_path / "quantizer_state.pth", weights_only=True), + yaml.safe_load((export_path / "quant_recipe.yaml").read_text()), + ) + ) + + weights, state, recipe = exports[0] + assert_exported_checkpoint_matches( + tmp_path / "vllm_fakequant_1", hf_model_path, check_values=False + ) + assert not any("quantizer" in name for name in weights) + quantizers = recipe["quantizers"] + assert {key.rpartition(".")[0] for key in state} <= quantizers.keys() + assert any(name.endswith("k_bmm_quantizer") for name in quantizers) + for other_weights, other_state, other_recipe in exports[1:]: + assert other_recipe == recipe + for reference, other in ((weights, other_weights), (state, other_state)): + assert other.keys() == reference.keys() + for name, tensor in reference.items(): + torch.testing.assert_close(other[name], tensor, rtol=0, atol=0, msg=name) diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py new file mode 100644 index 00000000000..ecff1abd29c --- /dev/null +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""vLLM fake-quant export through Megatron-Bridge matches ModelOpt's Megatron exporter.""" + +import os +from functools import partial +from pathlib import Path + +import pytest +import torch +import yaml +from _test_utils.torch.transformers_models import create_tiny_nemotron_h_dir, create_tiny_qwen3_dir +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.core.utils import unwrap_model +from safetensors.torch import load_file + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.export 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.quantization.nn import TensorQuantizer +from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf +from modelopt.torch.utils.plugins.megatron_generate import megatron_prefill + +if not hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe"): + pytest.skip( + "Megatron-Bridge does not write the ModelOpt quantizer sidecars", allow_module_level=True + ) + + +def _load_export(export_dir: Path): + weights = {} + for shard in sorted(export_dir.glob("*.safetensors")): + weights.update(load_file(shard)) + state = torch.load(export_dir / "quantizer_state.pth", weights_only=True) + recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text()) + return weights, state, recipe + + +def _test_matches_megatron_exporter(source, tmp_path, rank, size): + # Models in a reused worker pick their attention backend through process-wide variables. + for name in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"): + os.environ.pop(name, None) + bridge, _, model, unwrapped, _ = load_mbridge_model_from_hf( + hf_model_name_or_path=str(source), provider_overrides={"pipeline_dtype": torch.bfloat16} + ) + torch.manual_seed(0) + input_ids = torch.randint(0, unwrapped.vocab_size, (2, 16), device="cuda") + mtq.quantize( + unwrapped, + load_recipe("general/ptq/nvfp4_default-kv_fp8").quantize.model_dump(), + lambda m: megatron_prefill(m, input_ids), + ) + reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge" + export_mcore_gpt_to_hf_vllm_fq(unwrapped, str(source), export_dir=str(reference_dir)) + for chunk in model: + check_mcore_vllm_fakequant_quantizers(unwrap_model(chunk)) + mtq.fold_weight(unwrap_model(chunk)) + bridge.save_hf_pretrained(model, str(export_dir)) + if rank != 0: + return + + ref_weights, ref_state, ref_recipe = _load_export(reference_dir) + weights, state, recipe = _load_export(export_dir) + assert weights.keys() == ref_weights.keys() + for name, weight in weights.items(): + torch.testing.assert_close(weight, ref_weights[name], rtol=0, atol=0, msg=name) + + assert recipe["format"] == "modelopt_quantizer_recipe" and recipe["version"] == 1 + quantizers = recipe["quantizers"] + enabled = {name for name, entry in ref_recipe.items() if not entry["_disabled"]} + assert quantizers.keys() == enabled + assert any(name.endswith("k_bmm_quantizer") for name in quantizers) + to_cfg = TensorQuantizer.attribute_config_from_modelopt_state + for name, entry in quantizers.items(): + assert to_cfg(ref_recipe[name]).items() <= to_cfg(entry).items(), name + + assert state.keys() == {key for key in ref_state if key.rpartition(".")[0] in quantizers} + for key, value in state.items(): + torch.testing.assert_close(value, ref_state[key], rtol=0, atol=0, msg=key) + + +@pytest.mark.parametrize("create_source", [create_tiny_qwen3_dir, create_tiny_nemotron_h_dir]) +def test_matches_megatron_exporter(dist_workers_size_1, tmp_path, create_source): + source = create_source(tmp_path, with_tokenizer=True) + dist_workers_size_1.run(partial(_test_matches_megatron_exporter, source, tmp_path)) From c5c3433118822679a15bbfae165afb4245d2b89f Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 19 Oct 2026 13:07:23 -0700 Subject: [PATCH 02/11] Drop every weight-quantizer buffer when folding weights mtq.fold_weight bakes the weight quantizer into the weight, so any state the quantizer keeps to re-derive that weight (e.g. a pinned GGML payload) is stale afterwards. Dropping all of its buffers, not only _amax and _pre_quant_scale, lets Megatron-Bridge save a folded IQ2 model without unmapped quantizer tensors. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../torch/quantization/nn/modules/quant_module.py | 12 +++++++----- tests/unit/torch/quantization/test_ggml_gptq.py | 2 ++ 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 8b1d8a4d349..e87acf7dff4 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -164,7 +164,9 @@ def _fold_weight_quantizer( # Disable pre-quant scaling so it is not applied twice. quantizer._enable_pre_quant_scale = False elif not keep_attrs: - for attr_name in ("_pre_quant_scale", "_amax"): + # Each of these only re-derives the folded weight (e.g. a pinned GGML payload). + buffers = {name for name, _ in quantizer.named_buffers(recurse=False)} + for attr_name in buffers | {"_pre_quant_scale", "_amax"}: if hasattr(quantizer, attr_name): delattr(quantizer, attr_name) @@ -176,10 +178,10 @@ def fold_weight(self, keep_attrs: bool = False): Every fake-quant weight quantizer is folded, including disabled quantizers whose pre-quant scale or rotation remains active. The folded transform is baked into the stored weight and - then disabled, so subsequent forwards use the stored weight directly. Calibration buffers - (``_pre_quant_scale``, ``_amax``) are dropped unless ``keep_attrs``. A retained pre-quant - scale remains stored but is made inactive because it has already been applied to the folded - weight. + then disabled, so subsequent forwards use the stored weight directly. The quantizer's + buffers (e.g. ``_pre_quant_scale``, ``_amax``) are dropped unless ``keep_attrs``. A + retained pre-quant scale remains stored but is made inactive because it has already been + applied to the folded weight. """ # Handle all attributes that end with _weight_quantizer for name in dir(self): diff --git a/tests/unit/torch/quantization/test_ggml_gptq.py b/tests/unit/torch/quantization/test_ggml_gptq.py index 1106d330e3b..fc5a01ca080 100644 --- a/tests/unit/torch/quantization/test_ggml_gptq.py +++ b/tests/unit/torch/quantization/test_ggml_gptq.py @@ -319,6 +319,8 @@ def test_folding_a_bf16_weight_decodes_the_pinned_payload(): assert torch.equal(model.weight_quantizer(model.weight.float()).to(torch.bfloat16), decoded) mtq.fold_weight(model) assert torch.equal(model.weight, decoded) + # The folded weight is the decoding, so the payload is dropped with the other quantizer state. + assert _pin(model) is None def test_pin_survives_save_and_restore(): From 4c5fd634c5e3059b3f2c905fa8a6baaef2484586 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:31:03 -0700 Subject: [PATCH 03/11] Cover per-expert Qwen3-MoE in the fake-quant export parallelism test Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/examples/megatron_bridge/test_quantize_export.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index 770e4d8fdfc..79c5f26cb74 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -160,9 +160,11 @@ def test_quantize_and_export( ), "ep_size", ), + # Per-expert HF names with SequentialMLP experts: each EP rank owns some of the names. + (partial(create_tiny_qwen3_moe_dir, with_tokenizer=True), "ep_size"), (partial(create_tiny_nemotron_h_dir, with_tokenizer=True), "pp_size"), ], - ids=["qwen3_5_moe_vl", "nemotron_h"], + ids=["qwen3_5_moe_vl", "qwen3_moe", "nemotron_h"], ) @pytest.mark.timeout(480) def test_quantize_and_vllm_fakequant_export( From 916395d389555b82f775f1b81d1fb1863560955d Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:13:26 -0700 Subject: [PATCH 04/11] Keep the SVDQuant low-rank term when folding weights Folding now drops every weight-quantizer buffer, which included the SVDQuant LoRA factors that SVDQuantLinear.fold_weight read afterwards, so the folded weight silently lost the low-rank correction. Read the factors before folding. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../quantization/nn/modules/quant_linear.py | 26 +++++-------------- .../quantization/test_tensor_quant_cpu.py | 23 ++++++++++++++++ 2 files changed, 30 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py index d0569ff5b71..a817b3025cd 100644 --- a/modelopt/torch/quantization/nn/modules/quant_linear.py +++ b/modelopt/torch/quantization/nn/modules/quant_linear.py @@ -171,26 +171,14 @@ def fold_weight(self, keep_attrs: bool = False): and isinstance(self.weight_quantizer, TensorQuantizer) and self.weight_quantizer.fake_quant ) + # Read the LoRA factors first: folding drops the weight quantizer's buffers. + lora_a = lora_b = None + if should_fold and not keep_attrs and self._not_sequential_quantizers(): + lora_a = self.weight_quantizer.svdquant_lora_a + lora_b = self.weight_quantizer.svdquant_lora_b super().fold_weight(keep_attrs) - if should_fold: - if ( - not keep_attrs - and self._not_sequential_quantizers() - and self.weight_quantizer.svdquant_lora_a is not None - and self.weight_quantizer.svdquant_lora_b is not None - ): - self.weight.data.copy_( - self.weight - + self.weight_quantizer.svdquant_lora_b @ self.weight_quantizer.svdquant_lora_a - ) - if not keep_attrs: - _attrs = [ - "_svdquant_lora_a", - "_svdquant_lora_b", - ] - for attr in _attrs: - if hasattr(self.weight_quantizer, attr): - delattr(self.weight_quantizer, attr) + if lora_a is not None and lora_b is not None: + self.weight.data.copy_(self.weight + lora_b @ lora_a) class RealQuantLinear(QuantModule): diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index 31352e42bb5..2546c692816 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -26,6 +26,7 @@ import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module from modelopt.torch.quantization import QuantModuleRegistry, tensor_quant from modelopt.torch.quantization.config import QuantizerAttributeConfig, RotateConfig +from modelopt.torch.quantization.conversion import create_and_replace_svdquant_linear_on_the_fly from modelopt.torch.quantization.nn import ( SequentialQuantizer, TensorQuantizer, @@ -403,6 +404,28 @@ def test_fold_weight_keep_attrs_keeps_calibration_attrs_inactive(monkeypatch): unregister_quant_backend(backend_name) +def test_fold_weight_folds_svdquant_lora(): + model = torch.nn.Sequential(torch.nn.Linear(4, 3)) + inputs = torch.randn(2, 4) + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + ], + "algorithm": "max", + } + mtq.quantize(model, config, lambda m: m(inputs)) + create_and_replace_svdquant_linear_on_the_fly(model) + model[0].weight_quantizer.svdquant_lora_a = torch.randn(2, 4) + model[0].weight_quantizer.svdquant_lora_b = torch.randn(3, 2) + output_before = model(inputs) + + mtq.fold_weight(model) + + assert not hasattr(model[0].weight_quantizer, "_svdquant_lora_a") + assert torch.allclose(model(inputs), output_before, atol=1e-6) + + def test_temporarily_fold_weights_restores_after_exception(monkeypatch): calls = [] backend_name = "test_temporary_fold_backend" From 45cca84182db420b57f636b6d744e7e66ca544d6 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:21:38 -0700 Subject: [PATCH 05/11] Mark the export as saved in one place for both export paths test_mlflow_utils requires a single args.checkpoint_exported = True per save; the vLLM fake-quant branch added a second one. Move the unified export into export_unified_hf so both paths share it. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../export_quantized_megatron_to_hf.py | 62 ++++++++++--------- 1 file changed, 32 insertions(+), 30 deletions(-) diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index 2e996d37024..c9db893743b 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -208,6 +208,36 @@ def export_vllm_fakequant(bridge, model: list, export_dir: str) -> None: 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 @@ -259,36 +289,8 @@ def main(args: argparse.Namespace): 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) - args.checkpoint_exported = True - print_rank_0(f"Exported vLLM fake-quant checkpoint to {args.export_unified_hf_path}") - return - - # 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, - ) + 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}") From f882803a37652d0717e5b1c20daf419440439599 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:16:08 -0700 Subject: [PATCH 06/11] Fold sequential weight quantizers (e.g. W4A8) in mtq.fold_weight fold_weight skipped SequentialQuantizer weight quantizers, so the Megatron-Bridge fake-quant export saved W4A8 weights unquantized. Fold the whole chain, member by member in the weight dtype as the forward pass does, then disable and clean each member. The export now matches the Megatron exporter bitwise for W4A8 too. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + .../quantization/nn/modules/quant_module.py | 35 +++++++++++-------- .../plugins/transformer_engine.py | 7 ++-- .../test_vllm_fakequant_mbridge_export.py | 34 ++++++++++++------ .../quantization/test_tensor_quant_cpu.py | 19 ++++++++++ 5 files changed, 70 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8145bfa1ca0..4200b39beee 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -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. diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index e87acf7dff4..903a44d7b2e 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -142,7 +142,7 @@ def iter_weight_quantizers_for_calibration(self): @staticmethod @torch.no_grad() def _fold_weight_quantizer( - quantizer: TensorQuantizer, + quantizer: TensorQuantizer | SequentialQuantizer, weights: Iterable[torch.Tensor], keep_attrs: bool = False, ): @@ -150,25 +150,32 @@ def _fold_weight_quantizer( ``weights`` is an iterable so a single quantizer shared across views (e.g. one per-tensor quantizer over all experts of a fused MoE weight) can be folded view by - view while disabling and dropping its calibration attrs exactly once. + view while disabling and dropping its calibration attrs exactly once. A + :class:`SequentialQuantizer` is folded as its whole chain. """ if not quantizer.fake_quant: return + members = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] for weight in weights: - weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype)) + folded = weight + # Each member's output is in the weight dtype, as in the forward pass. + for member in members: + folded = member(folded.float().contiguous()).to(weight.dtype) + weight.data.copy_(folded) quantizer.disable() quantizer.disable_rotate() - if keep_attrs and hasattr(quantizer, "_pre_quant_scale"): - # The scale is already baked into the folded weight. - # Disable pre-quant scaling so it is not applied twice. - quantizer._enable_pre_quant_scale = False - elif not keep_attrs: - # Each of these only re-derives the folded weight (e.g. a pinned GGML payload). - buffers = {name for name, _ in quantizer.named_buffers(recurse=False)} - for attr_name in buffers | {"_pre_quant_scale", "_amax"}: - if hasattr(quantizer, attr_name): - delattr(quantizer, attr_name) + for member in members: + if keep_attrs and hasattr(member, "_pre_quant_scale"): + # The scale is already baked into the folded weight. + # Disable pre-quant scaling so it is not applied twice. + member._enable_pre_quant_scale = False + elif not keep_attrs: + # Each of these only re-derives the folded weight (e.g. a pinned GGML payload). + buffers = {name for name, _ in member.named_buffers(recurse=False)} + for attr_name in buffers | {"_pre_quant_scale", "_amax"}: + if hasattr(member, attr_name): + delattr(member, attr_name) def _validate_fold_weight(self): """Validate that weight folding is supported before any module is modified.""" @@ -188,7 +195,7 @@ def fold_weight(self, keep_attrs: bool = False): attr = getattr(self, name) if ( name.endswith("weight_quantizer") - and isinstance(attr, TensorQuantizer) + and isinstance(attr, (TensorQuantizer, SequentialQuantizer)) and attr.fake_quant ): # Get the corresponding weight name by removing _weight_quantizer suffix diff --git a/modelopt/torch/quantization/plugins/transformer_engine.py b/modelopt/torch/quantization/plugins/transformer_engine.py index c2ae02a7f93..c7517eb559c 100644 --- a/modelopt/torch/quantization/plugins/transformer_engine.py +++ b/modelopt/torch/quantization/plugins/transformer_engine.py @@ -247,9 +247,12 @@ def iter_weights_for_calibration(self): def fold_weight(self, keep_attrs: bool = False): """Fold each grouped weight with its corresponding fake-quant quantizer.""" - quantizer_weights: dict[TensorQuantizer, list[torch.Tensor]] = {} + quantizer_weights: dict[TensorQuantizer | SequentialQuantizer, list[torch.Tensor]] = {} for weight, quantizer in self.iter_weights_for_calibration(): - if isinstance(quantizer, TensorQuantizer) and quantizer.fake_quant: + if ( + isinstance(quantizer, (TensorQuantizer, SequentialQuantizer)) + and quantizer.fake_quant + ): quantizer_weights.setdefault(quantizer, []).append(weight) for quantizer, weights in quantizer_weights.items(): diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py index ecff1abd29c..a039694da52 100644 --- a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py @@ -51,7 +51,14 @@ def _load_export(export_dir: Path): return weights, state, recipe -def _test_matches_megatron_exporter(source, tmp_path, rank, size): +QUANT_CFGS = { + "nvfp4-kv_fp8": load_recipe("general/ptq/nvfp4_default-kv_fp8").quantize.model_dump(), + # INT4 + FP8 weights: a SequentialQuantizer chain the export folds as a whole. + "w4a8": {**mtq.W4A8_AWQ_BETA_CFG, "algorithm": "max"}, +} + + +def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size): # Models in a reused worker pick their attention backend through process-wide variables. for name in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"): os.environ.pop(name, None) @@ -60,11 +67,7 @@ def _test_matches_megatron_exporter(source, tmp_path, rank, size): ) torch.manual_seed(0) input_ids = torch.randint(0, unwrapped.vocab_size, (2, 16), device="cuda") - mtq.quantize( - unwrapped, - load_recipe("general/ptq/nvfp4_default-kv_fp8").quantize.model_dump(), - lambda m: megatron_prefill(m, input_ids), - ) + mtq.quantize(unwrapped, QUANT_CFGS[quant_cfg_name], lambda m: megatron_prefill(m, input_ids)) reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge" export_mcore_gpt_to_hf_vllm_fq(unwrapped, str(source), export_dir=str(reference_dir)) for chunk in model: @@ -84,7 +87,9 @@ def _test_matches_megatron_exporter(source, tmp_path, rank, size): quantizers = recipe["quantizers"] enabled = {name for name, entry in ref_recipe.items() if not entry["_disabled"]} assert quantizers.keys() == enabled - assert any(name.endswith("k_bmm_quantizer") for name in quantizers) + assert any(name.endswith("k_bmm_quantizer") for name in quantizers) == quant_cfg_name.endswith( + "kv_fp8" + ) to_cfg = TensorQuantizer.attribute_config_from_modelopt_state for name, entry in quantizers.items(): assert to_cfg(ref_recipe[name]).items() <= to_cfg(entry).items(), name @@ -94,7 +99,16 @@ def _test_matches_megatron_exporter(source, tmp_path, rank, size): torch.testing.assert_close(value, ref_state[key], rtol=0, atol=0, msg=key) -@pytest.mark.parametrize("create_source", [create_tiny_qwen3_dir, create_tiny_nemotron_h_dir]) -def test_matches_megatron_exporter(dist_workers_size_1, tmp_path, create_source): +@pytest.mark.parametrize( + ("create_source", "quant_cfg_name"), + [ + (create_tiny_qwen3_dir, "nvfp4-kv_fp8"), + (create_tiny_nemotron_h_dir, "nvfp4-kv_fp8"), + (create_tiny_qwen3_dir, "w4a8"), + ], +) +def test_matches_megatron_exporter(dist_workers_size_1, tmp_path, create_source, quant_cfg_name): source = create_source(tmp_path, with_tokenizer=True) - dist_workers_size_1.run(partial(_test_matches_megatron_exporter, source, tmp_path)) + dist_workers_size_1.run( + partial(_test_matches_megatron_exporter, source, quant_cfg_name, tmp_path) + ) diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index 2546c692816..81191fd6ae4 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -426,6 +426,25 @@ def test_fold_weight_folds_svdquant_lora(): assert torch.allclose(model(inputs), output_before, atol=1e-6) +@torch.no_grad() +def test_fold_weight_folds_sequential_weight_quantizer(): + model = SimpleLinear() + inputs = model.get_input() + mtq.quantize(model, {**WINT4INT8_CFG, "algorithm": "max"}, lambda m: m(inputs)) + expected = { + module: module.weight_quantizer(module.weight) + for module in model.modules() + if isinstance(getattr(module, "weight_quantizer", None), SequentialQuantizer) + } + assert expected + + mtq.fold_weight(model) + + for module, weight in expected.items(): + assert torch.equal(module.weight, weight) + assert not any(q.is_enabled or hasattr(q, "_amax") for q in module.weight_quantizer) + + def test_temporarily_fold_weights_restores_after_exception(monkeypatch): calls = [] backend_name = "test_temporary_fold_backend" From 5c1fc9edaa5c3afff1d566c3fdefa231cea8d98b Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:00:18 -0700 Subject: [PATCH 07/11] Fold the fake-quant export tests into the existing Megatron example tests test_quantize_and_export now also exports each quantized checkpoint with --vllm_fakequant_export (EP, TP and PP across its three cases) instead of a separate test quantizing three more models; test_qad keeps the fake-quant export only for the VLM, the one case only QAD produces. The parity test keeps one NVFP4 + FP8 KV case and the W4A8 case. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/examples/megatron_bridge/test_qad.py | 8 +- .../megatron_bridge/test_quantize_export.py | 131 +++++++----------- .../test_vllm_fakequant_mbridge_export.py | 6 +- 3 files changed, 57 insertions(+), 88 deletions(-) diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index d25cd8e0c0e..c343f2f241f 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -250,9 +250,9 @@ def test_qad(tmp_path: Path, num_gpus, create_student, lora_rank): # allow_unexpected waives the expert names wholesale; this re-tightens it. assert_per_expert_experts_complete(hf_export_path) - # Step 4: export for vLLM fake-quant evaluation. Megatron-Bridge writes the whole model, so - # a VLM's language-model-only QAD checkpoint must still come out with the original vision tower. - if HAS_BRIDGE_QUANTIZER_SIDECARS: + # Step 4: export for vLLM fake-quant evaluation (test_quantize_export.py covers the rest). + # Only QAD saves a VLM's language model alone; the export must still carry the vision tower. + if is_vlm and HAS_BRIDGE_QUANTIZER_SIDECARS: fakequant_path = tmp_path / "qad_vllm_fakequant" fakequant_cmd = extend_cmd_parts( [ @@ -271,7 +271,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, lora_rank): fakequant_path, hf_model_path, check_values=False, - bit_exact_prefixes=("model.visual.",) if is_vlm else (), + bit_exact_prefixes=("model.visual.",), ) recipe = yaml.safe_load((fakequant_path / "quant_recipe.yaml").read_text()) state = torch.load(fakequant_path / "quantizer_state.pth", weights_only=True) diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index 79c5f26cb74..92f140a1518 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -14,7 +14,6 @@ # limitations under the License. """Tests for quantize.py and export_quantized_megatron_to_hf.py scripts.""" -from functools import partial from pathlib import Path import pytest @@ -28,12 +27,14 @@ from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state from _test_utils.torch.transformers_models import ( create_tiny_nemotron_h_dir, - create_tiny_qwen3_5_moe_vl_dir, create_tiny_qwen3_moe_dir, create_tiny_qwen3vl_dir, ) from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +# TODO: drop this probe once the minimum Megatron-Bridge writes the ModelOpt quantizer sidecars. +HAS_BRIDGE_QUANTIZER_SIDECARS = hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe") + # Per-architecture export *mappings* are covered in-process by # tests/gpu_megatron/torch/export/test_unified_export_megatron.py; these cases cover the script # wiring (CLI, recipe, checkpoint hand-off) that only running quantize.py + the exporter exercises. @@ -49,7 +50,14 @@ @pytest.mark.parametrize( - ("create_model", "model_kwargs", "recipe", "quantize_parallelism", "export_parallelism"), + ( + "create_model", + "model_kwargs", + "recipe", + "quantize_parallelism", + "export_parallelism", + "fakequant_export_parallelism", + ), [ # MoE: routed experts used to be dropped silently from the export. Calibrates with context # parallelism (not on the VLM: M-RoPE CP needs nemo:26.10). @@ -59,17 +67,33 @@ "general/ptq/nvfp4_default-kv_fp8", "cp_size", "pp_size", + # Per-expert HF names, each EP rank owning some of them. + "ep_size", ), # Dense VLM: only the language model is quantized, vision is copied through. Keeps the # FP8 script-to-checkpoint coverage. - (create_tiny_qwen3vl_dir, {}, "general/ptq/fp8_default-kv_fp8", "tp_size", "pp_size"), + ( + create_tiny_qwen3vl_dir, + {}, + "general/ptq/fp8_default-kv_fp8", + "tp_size", + "pp_size", + "tp_size", + ), # Mamba hybrid + MoE with grouped-GEMM experts: exported with its experts sharded # across ranks (EP), resharded from the TP-quantized checkpoint. - (create_tiny_nemotron_h_dir, {}, "general/ptq/nvfp4_default-kv_fp8", "tp_size", "ep_size"), + ( + create_tiny_nemotron_h_dir, + {}, + "general/ptq/nvfp4_default-kv_fp8", + "tp_size", + "ep_size", + "pp_size", + ), ], ids=["qwen3_moe", "qwen3vl", "nemotron_h"], ) -@pytest.mark.timeout(360) # quantize + export in one test; 1-gpu CI exceeds the default 300s +@pytest.mark.timeout(420) # quantize + 2 exports in one test; 1-gpu CI exceeds the default 300s def test_quantize_and_export( tmp_path: Path, num_gpus, @@ -78,6 +102,7 @@ def test_quantize_and_export( recipe, quantize_parallelism, export_parallelism, + fakequant_export_parallelism, fake_hub_checkpoint, ): """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint. @@ -143,84 +168,32 @@ def test_quantize_and_export( # outputs = llm.generate(["Hello!"], vllm.SamplingParams(max_tokens=4)) # assert outputs and outputs[0].outputs and outputs[0].outputs[0].text - -@pytest.mark.skipif( - not hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe"), - reason="Megatron-Bridge does not write the ModelOpt quantizer sidecars", -) -@pytest.mark.parametrize( - ("create_model", "export_parallelism"), - [ - ( - partial( - create_tiny_qwen3_5_moe_vl_dir, - with_processor=True, - num_hidden_layers=2, - layer_types=["linear_attention", "full_attention"], - ), - "ep_size", - ), - # Per-expert HF names with SequentialMLP experts: each EP rank owns some of the names. - (partial(create_tiny_qwen3_moe_dir, with_tokenizer=True), "ep_size"), - (partial(create_tiny_nemotron_h_dir, with_tokenizer=True), "pp_size"), - ], - ids=["qwen3_5_moe_vl", "qwen3_moe", "nemotron_h"], -) -@pytest.mark.timeout(480) -def test_quantize_and_vllm_fakequant_export( - tmp_path: Path, num_gpus, create_model, export_parallelism -): - """Fake-quant exports through Megatron-Bridge do not depend on the export parallelism.""" - hf_model_path = create_model(tmp_path) - hf_model_path = hf_model_path[0] if isinstance(hf_model_path, tuple) else hf_model_path - megatron_path = tmp_path / "quantized_megatron" - quantize_cmd = extend_cmd_parts( - ["torchrun", "--nproc_per_node=1", "quantize.py", "--skip_generate"], - hf_model_name_or_path=hf_model_path, - recipe="general/ptq/nvfp4_default-kv_fp8", - calib_dataset_name="cnn_dailymail", - calib_num_samples=4, - calib_batch_size=2, - seq_length=16, - export_megatron_path=megatron_path, - ) - run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True) - - exports = [] - for size in sorted({1, num_gpus}): - export_path = tmp_path / f"vllm_fakequant_{size}" - export_cmd = extend_cmd_parts( + # Step 3: export folded weights + quantizer sidecars through Megatron-Bridge for vLLM fake quant + if HAS_BRIDGE_QUANTIZER_SIDECARS: + fakequant_path = tmp_path / "vllm_fakequant" + fakequant_cmd = extend_cmd_parts( [ "torchrun", - f"--nproc_per_node={size}", + f"--nproc_per_node={num_gpus}", "export_quantized_megatron_to_hf.py", "--vllm_fakequant_export", ], - hf_model_name_or_path=hf_model_path, + hf_model_name_or_path=model_arg, megatron_path=megatron_path, - export_unified_hf_path=export_path, - **{export_parallelism: size}, + export_unified_hf_path=fakequant_path, + **{fakequant_export_parallelism: num_gpus}, ) - run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True) - exports.append( - ( - load_safetensors_dir(export_path), - torch.load(export_path / "quantizer_state.pth", weights_only=True), - yaml.safe_load((export_path / "quant_recipe.yaml").read_text()), - ) + run_example_command(fakequant_cmd, example_path="megatron_bridge", setup_free_port=True) + is_vlm = "vision_config" in (hf_model_path / "config.json").read_text() + assert_exported_checkpoint_matches( + fakequant_path, + hf_model_path, + check_values=False, + bit_exact_prefixes=("model.visual.",) if is_vlm else (), ) - - weights, state, recipe = exports[0] - assert_exported_checkpoint_matches( - tmp_path / "vllm_fakequant_1", hf_model_path, check_values=False - ) - assert not any("quantizer" in name for name in weights) - quantizers = recipe["quantizers"] - assert {key.rpartition(".")[0] for key in state} <= quantizers.keys() - assert any(name.endswith("k_bmm_quantizer") for name in quantizers) - for other_weights, other_state, other_recipe in exports[1:]: - assert other_recipe == recipe - for reference, other in ((weights, other_weights), (state, other_state)): - assert other.keys() == reference.keys() - for name, tensor in reference.items(): - torch.testing.assert_close(other[name], tensor, rtol=0, atol=0, msg=name) + assert not any("quantizer" in name for name in load_safetensors_dir(fakequant_path)) + recipe = yaml.safe_load((fakequant_path / "quant_recipe.yaml").read_text()) + quantizers = recipe["quantizers"] + state = torch.load(fakequant_path / "quantizer_state.pth", weights_only=True) + assert any(name.endswith("k_bmm_quantizer") for name in quantizers) + assert {key.rpartition(".")[0] for key in state} <= quantizers.keys() diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py index a039694da52..408888cc3f1 100644 --- a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py @@ -101,11 +101,7 @@ def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size @pytest.mark.parametrize( ("create_source", "quant_cfg_name"), - [ - (create_tiny_qwen3_dir, "nvfp4-kv_fp8"), - (create_tiny_nemotron_h_dir, "nvfp4-kv_fp8"), - (create_tiny_qwen3_dir, "w4a8"), - ], + [(create_tiny_nemotron_h_dir, "nvfp4-kv_fp8"), (create_tiny_qwen3_dir, "w4a8")], ) def test_matches_megatron_exporter(dist_workers_size_1, tmp_path, create_source, quant_cfg_name): source = create_source(tmp_path, with_tokenizer=True) From 08ce4b2d833170455ebfd0495f9d0af58c1877e4 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:06:33 -0700 Subject: [PATCH 08/11] Compare against the versioned recipe of the Megatron-LM fake-quant export Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../export/test_vllm_fakequant_mbridge_export.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py index 408888cc3f1..4438ed56f80 100644 --- a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py @@ -29,6 +29,10 @@ import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe from modelopt.torch.export import export_mcore_gpt_to_hf_vllm_fq +from modelopt.torch.export.plugins.vllm_fakequant_hf import ( + QUANTIZER_RECIPE_FORMAT, + QUANTIZER_RECIPE_VERSION, +) from modelopt.torch.export.plugins.vllm_fakequant_megatron import ( check_mcore_vllm_fakequant_quantizers, ) @@ -48,7 +52,9 @@ def _load_export(export_dir: Path): weights.update(load_file(shard)) state = torch.load(export_dir / "quantizer_state.pth", weights_only=True) recipe = yaml.safe_load((export_dir / "quant_recipe.yaml").read_text()) - return weights, state, recipe + assert recipe["format"] == QUANTIZER_RECIPE_FORMAT + assert recipe["version"] == QUANTIZER_RECIPE_VERSION + return weights, state, recipe["quantizers"] QUANT_CFGS = { @@ -78,13 +84,11 @@ def _test_matches_megatron_exporter(source, quant_cfg_name, tmp_path, rank, size return ref_weights, ref_state, ref_recipe = _load_export(reference_dir) - weights, state, recipe = _load_export(export_dir) + weights, state, quantizers = _load_export(export_dir) assert weights.keys() == ref_weights.keys() for name, weight in weights.items(): torch.testing.assert_close(weight, ref_weights[name], rtol=0, atol=0, msg=name) - assert recipe["format"] == "modelopt_quantizer_recipe" and recipe["version"] == 1 - quantizers = recipe["quantizers"] enabled = {name for name, entry in ref_recipe.items() if not entry["_disabled"]} assert quantizers.keys() == enabled assert any(name.endswith("k_bmm_quantizer") for name in quantizers) == quant_cfg_name.endswith( From f13e5cce03334c566531474ae1d456997cf761f0 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:42:11 -0700 Subject: [PATCH 09/11] Check the quantizers once before the Megatron-LM fake-quant export export_mcore_gpt_to_hf_vllm_fq now runs check_mcore_vllm_fakequant_quantizers up front, as the Megatron-Bridge path does, instead of collecting the same MTP, lm_head and unsupported-setting failures during the export. The check reports MTP and lm_head quantizers before their individual settings. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../export/plugins/vllm_fakequant_megatron.py | 47 +++++-------------- 1 file changed, 11 insertions(+), 36 deletions(-) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index e063f849749..64c4deb4ed4 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -104,11 +104,10 @@ def check_mcore_vllm_fakequant_quantizers(model: torch.nn.Module) -> None: Checks the settings the vLLM fake-quant export and reload support, with weight quantizers folded into the weights. Collective: a failure on any rank raises on every rank. """ - _, failure = _quantizer_configs(model) + failure = "" tied = any(getattr(m, "share_embeddings_and_output_weights", False) for m in model.modules()) + # MTP and lm_head quantizers are rejected outright, before their individual settings. for name, quantizer in model.named_modules(): - if failure: - break if not isinstance(quantizer, TensorQuantizer): continue active = quantizer.is_enabled and quantizer._if_quant @@ -119,6 +118,10 @@ def check_mcore_vllm_fakequant_quantizers(model: torch.nn.Module) -> None: (active and (tied or "weight_quantizer" not in name)) or (tied and transform) ): failure = f"Unsupported lm_head quantization for vLLM fakequant: {name}" + if failure: + break + if not failure: + _, failure = _quantizer_configs(model) failure = DistributedProcessGroup.get_dist_syncd_obj( failure, DistributedProcessGroup(None), @@ -257,24 +260,8 @@ def __init__(self, *args, **kwargs): self._grouped_fq_modules: set[torch.nn.Module] = set() def _gather_exclude_modules(self) -> list[str]: - """Validate settings before gathering excluded modules.""" - # Tied embeddings bypass the output-layer route, so validate the live head here too. - head = getattr(self.model, "output_layer", None) - if head is not None: - for name, quantizer in head.named_modules(): - if not isinstance(quantizer, TensorQuantizer): - continue - tied = self.model.share_embeddings_and_output_weights - active = quantizer.is_enabled and quantizer._if_quant - if (active and (tied or "weight_quantizer" not in name)) or ( - tied and (quantizer.rotate_is_enabled or quantizer.pre_quant_scale is not None) - ): - self._quantizer_validation_failure = ( - self._quantizer_validation_failure - or f"Unsupported lm_head quantization for vLLM fakequant: {name}" - ) - break - # All ranks reach this after MTP collection and before writing weights. + """Raise export-time failures (e.g. conflicting packed recipes) before gathering.""" + # All ranks reach this after collecting the shards and before writing weights. failure = DistributedProcessGroup.get_dist_syncd_obj( self._quantizer_validation_failure, DistributedProcessGroup(None), @@ -470,23 +457,10 @@ def _get_quantized_state( """ name_to_value = {} qformat: str = self._get_quantization_format(module) + # MTP quantizers are rejected up front (check_mcore_vllm_fakequant_quantizers). if prefix.startswith("mtp."): - for name, quantizer in module.named_modules(): - if isinstance(quantizer, TensorQuantizer) and ( - (quantizer.is_enabled and quantizer._if_quant) - or quantizer.rotate_is_enabled - or quantizer.pre_quant_scale is not None - ): - self._quantizer_validation_failure = self._quantizer_validation_failure or ( - f"MTP quantization is not supported by vLLM fakequant export/reload: {name}" - ) - break - return self._get_weight_bias(module, dtype), qformat, 0 - quantizer_configs, error = _quantizer_configs(module) - if error: - self._quantizer_validation_failure = self._quantizer_validation_failure or error - # Skip unsupported kernels while peers complete the export collectives. return self._get_weight_bias(module, dtype), qformat, 0 + quantizer_configs, _ = _quantizer_configs(module) source_prefix = prefix if not prefix or prefix.endswith(".") else prefix + "." for qname, qstate in quantizer_configs.items(): marker_id = len(self._quantizer_recipe_markers) @@ -571,6 +545,7 @@ def export_mcore_gpt_to_hf_vllm_fq( Raises: ValueError: If a quantizer uses settings that cannot be restored. """ + check_mcore_vllm_fakequant_quantizers(model) exporter = VllmFqGPTModelExporter( model, pretrained_model_name_or_path, From 65ca9f2907f2e5e26bdabc9f4ce414a3c9d165bf Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:45:45 -0700 Subject: [PATCH 10/11] Test the Megatron examples and GPU tests against a newer Megatron-Bridge Overlay a pinned Megatron-Bridge commit and its Megatron-LM on every CI job that runs in a NeMo container, so the fake-quant export tests run before a NeMo container ships it. The tests no longer skip without the quantizer sidecars, so an older Megatron-Bridge fails them instead of skipping them silently. The README shows how to mount the same commit. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../actions/megatron-from-source/action.yml | 27 +++++++++ .github/workflows/_example_tests_runner.yml | 2 + .github/workflows/gpu_tests.yml | 2 + examples/megatron_bridge/README.md | 18 +++++- .../export_quantized_megatron_to_hf.py | 3 +- tests/examples/megatron_bridge/test_qad.py | 6 +- .../megatron_bridge/test_quantize_export.py | 59 +++++++++---------- .../test_vllm_fakequant_mbridge_export.py | 6 -- 8 files changed, 78 insertions(+), 45 deletions(-) create mode 100644 .github/actions/megatron-from-source/action.yml diff --git a/.github/actions/megatron-from-source/action.yml b/.github/actions/megatron-from-source/action.yml new file mode 100644 index 00000000000..50eef77d41d --- /dev/null +++ b/.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" diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml index 18979ce0f48..67f6eb7935e 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.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: | diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index 6340788026f..fd8b8278cbd 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.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 diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 78264741691..4d67602186c 100644 --- a/examples/megatron_bridge/README.md +++ b/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. @@ -390,7 +406,7 @@ Serve the export with [vllm_serve_fakequant.py](../vllm_serve/vllm_serve_fakequa 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. +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 diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index c9db893743b..3b7871e5ebb 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -192,7 +192,8 @@ def export_vllm_fakequant(bridge, model: list, export_dir: str) -> None: 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); update Megatron-Bridge." + "(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)) diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index c343f2f241f..e9bc834520f 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -35,12 +35,8 @@ create_tiny_qwen3_5_moe_vl_dir, create_tiny_qwen3_dir, ) -from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from torch.distributed.checkpoint import FileSystemReader -# TODO: drop this probe once the minimum Megatron-Bridge writes the ModelOpt quantizer sidecars. -HAS_BRIDGE_QUANTIZER_SIDECARS = hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe") - # The fix ships in the nemo:26.10 container. # TODO(Megatron-Bridge#6243): drop this probe once the minimum Megatron-Bridge carries it. try: @@ -252,7 +248,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, lora_rank): # Step 4: export for vLLM fake-quant evaluation (test_quantize_export.py covers the rest). # Only QAD saves a VLM's language model alone; the export must still carry the vision tower. - if is_vlm and HAS_BRIDGE_QUANTIZER_SIDECARS: + if is_vlm: fakequant_path = tmp_path / "qad_vllm_fakequant" fakequant_cmd = extend_cmd_parts( [ diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index 92f140a1518..d031990e76a 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -30,10 +30,6 @@ create_tiny_qwen3_moe_dir, create_tiny_qwen3vl_dir, ) -from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge - -# TODO: drop this probe once the minimum Megatron-Bridge writes the ModelOpt quantizer sidecars. -HAS_BRIDGE_QUANTIZER_SIDECARS = hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe") # Per-architecture export *mappings* are covered in-process by # tests/gpu_megatron/torch/export/test_unified_export_megatron.py; these cases cover the script @@ -169,31 +165,30 @@ def test_quantize_and_export( # assert outputs and outputs[0].outputs and outputs[0].outputs[0].text # Step 3: export folded weights + quantizer sidecars through Megatron-Bridge for vLLM fake quant - if HAS_BRIDGE_QUANTIZER_SIDECARS: - fakequant_path = tmp_path / "vllm_fakequant" - fakequant_cmd = extend_cmd_parts( - [ - "torchrun", - f"--nproc_per_node={num_gpus}", - "export_quantized_megatron_to_hf.py", - "--vllm_fakequant_export", - ], - hf_model_name_or_path=model_arg, - megatron_path=megatron_path, - export_unified_hf_path=fakequant_path, - **{fakequant_export_parallelism: num_gpus}, - ) - run_example_command(fakequant_cmd, example_path="megatron_bridge", setup_free_port=True) - is_vlm = "vision_config" in (hf_model_path / "config.json").read_text() - assert_exported_checkpoint_matches( - fakequant_path, - hf_model_path, - check_values=False, - bit_exact_prefixes=("model.visual.",) if is_vlm else (), - ) - assert not any("quantizer" in name for name in load_safetensors_dir(fakequant_path)) - recipe = yaml.safe_load((fakequant_path / "quant_recipe.yaml").read_text()) - quantizers = recipe["quantizers"] - state = torch.load(fakequant_path / "quantizer_state.pth", weights_only=True) - assert any(name.endswith("k_bmm_quantizer") for name in quantizers) - assert {key.rpartition(".")[0] for key in state} <= quantizers.keys() + fakequant_path = tmp_path / "vllm_fakequant" + fakequant_cmd = extend_cmd_parts( + [ + "torchrun", + f"--nproc_per_node={num_gpus}", + "export_quantized_megatron_to_hf.py", + "--vllm_fakequant_export", + ], + hf_model_name_or_path=model_arg, + megatron_path=megatron_path, + export_unified_hf_path=fakequant_path, + **{fakequant_export_parallelism: num_gpus}, + ) + run_example_command(fakequant_cmd, example_path="megatron_bridge", setup_free_port=True) + is_vlm = "vision_config" in (hf_model_path / "config.json").read_text() + assert_exported_checkpoint_matches( + fakequant_path, + hf_model_path, + check_values=False, + bit_exact_prefixes=("model.visual.",) if is_vlm else (), + ) + assert not any("quantizer" in name for name in load_safetensors_dir(fakequant_path)) + recipe = yaml.safe_load((fakequant_path / "quant_recipe.yaml").read_text()) + quantizers = recipe["quantizers"] + state = torch.load(fakequant_path / "quantizer_state.pth", weights_only=True) + assert any(name.endswith("k_bmm_quantizer") for name in quantizers) + assert {key.rpartition(".")[0] for key in state} <= quantizers.keys() diff --git a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py index 4438ed56f80..0eab41a1f63 100644 --- a/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py +++ b/tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py @@ -22,7 +22,6 @@ import torch import yaml from _test_utils.torch.transformers_models import create_tiny_nemotron_h_dir, create_tiny_qwen3_dir -from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from megatron.core.utils import unwrap_model from safetensors.torch import load_file @@ -40,11 +39,6 @@ from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf from modelopt.torch.utils.plugins.megatron_generate import megatron_prefill -if not hasattr(MegatronModelBridge, "build_modelopt_quantizer_recipe"): - pytest.skip( - "Megatron-Bridge does not write the ModelOpt quantizer sidecars", allow_module_level=True - ) - def _load_export(export_dir: Path): weights = {} From 449cffb8bbfbb45975b6d0aac7d76c3ab9a92638 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 23 Oct 2026 23:01:36 -0700 Subject: [PATCH 11/11] Keep a quantized output_layer's checkpoint state on newer Megatron-Core The CI jobs now run newer Megatron-Core, whose GPTModel.sharded_state_dict also accepts an empty tensor placeholder. Its body no longer matched the one keep_gpt_output_layer_extra_state recognises, so the patch was skipped and a quantized output_layer could not be checkpointed. Recognise both bodies. Skip test_layerwise_calibrate_vlm_text_forward_loop and test_fp8_real_quantize on Blackwell like the other tests that hit the flaky CUDA illegal memory access there; CI hit it in a different one of them on each of three runs. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- modelopt/torch/quantization/plugins/megatron.py | 10 +++++++--- .../torch/quantization/plugins/test_megatron.py | 2 ++ 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index a96e7fd2a33..84fd26cc4a3 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -372,8 +372,12 @@ def _output_layer_untied(config) -> bool: return False -# Statement kinds of the upstream GPTModel.sharded_state_dict body patched below. -_GPT_SSD_STATEMENTS = ["Assign", "Assign", "Assign", "Assert", "Return"] +# Statement kinds of the upstream GPTModel.sharded_state_dict bodies patched below. Newer +# megatron-core also accepts an empty tensor placeholder (one more assignment). +_GPT_SSD_STATEMENTS = ( + ["Assign", "Assign", "Assign", "Assert", "Return"], + ["Assign", "Assign", "Assign", "Assign", "Assert", "Return"], +) def _output_layer_extra_state_has_data(entry: Any) -> bool: @@ -402,7 +406,7 @@ def keep_gpt_output_layer_extra_state() -> bool: body = func.body if isinstance(func, ast.FunctionDef) else [] if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant): body = body[1:] # docstring - if [type(stmt).__name__ for stmt in body] != _GPT_SSD_STATEMENTS: + if [type(stmt).__name__ for stmt in body] not in _GPT_SSD_STATEMENTS: warn_rank_0( "GPTModel.sharded_state_dict is not the version ModelOpt patches; leaving it as is. " "If it does not keep a populated output_layer._extra_state, saving a quantized " diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 9e2a7be8a0b..c53a87276e6 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -691,6 +691,7 @@ def _test_fp8_real_quantize_helper(rank, size): assert real_quant_mem < cur_mem +@skip_flaky_on_blackwell def test_fp8_real_quantize(dist_workers): dist_workers.run(_test_fp8_real_quantize_helper) @@ -1725,6 +1726,7 @@ def forward_loop(_model): assert all(m.input_quantizer.amax is not None for m in calibrated) +@skip_flaky_on_blackwell @pytest.mark.parametrize("drive_wrapper", [False, True]) def test_layerwise_calibrate_vlm_text_forward_loop(dist_workers_size_1, drive_wrapper): dist_workers_size_1.run(partial(_test_layerwise_calibrate_vlm_text_forward_loop, drive_wrapper))