diff --git a/.github/actions/megatron-from-source/action.yml b/.github/actions/megatron-from-source/action.yml index 3c5b918e268..d8d94592cb1 100644 --- a/.github/actions/megatron-from-source/action.yml +++ b/.github/actions/megatron-from-source/action.yml @@ -6,7 +6,7 @@ description: >- inputs: commit: description: Megatron-Bridge commit. Keep in sync with examples/megatron_bridge/README.md. - default: 033ca6f65633d1f740dffc845dfa53dc4c4433bd + default: 0652ea797ab629e6451c64e1c92f976501b3b8b3 runs: using: composite diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 17ff5c2b231..c16a5c6e147 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -51,7 +51,7 @@ Some features (e.g. [`--vllm_fakequant_export`](#evaluating-a-qad-checkpoint-in- ```bash export MBRIDGE_DIR=${PWD}/Megatron-Bridge git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git ${MBRIDGE_DIR} -git -C ${MBRIDGE_DIR} checkout 033ca6f65633d1f740dffc845dfa53dc4c4433bd +git -C ${MBRIDGE_DIR} checkout 0652ea797ab629e6451c64e1c92f976501b3b8b3 git -C ${MBRIDGE_DIR} submodule update --init 3rdparty/Megatron-LM ``` diff --git a/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py new file mode 100644 index 00000000000..9cfa1549835 --- /dev/null +++ b/tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py @@ -0,0 +1,151 @@ +# 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. +"""Quantized HF export through Megatron-Bridge matches ModelOpt's Megatron exporter.""" + +import fnmatch +import json +import os +from functools import partial +from pathlib import Path + +import pytest +import torch +from _test_utils.torch.export.unified_checkpoint import load_safetensors_dir +from _test_utils.torch.transformers_models import ( + create_tiny_llama_dir, + create_tiny_nemotron_h_dir, + create_tiny_qwen3_5_moe_vl_dir, + create_tiny_qwen3_dir, +) + +import modelopt.torch.quantization as mtq +from modelopt.torch.export import export_mcore_gpt_to_hf +from modelopt.torch.utils.plugins.mbridge import ( + get_language_model, + load_mbridge_model_from_hf, + use_moe_grouped_gemm, +) +from modelopt.torch.utils.plugins.megatron_generate import megatron_prefill + +QUANT_CFGS = { + "nvfp4": mtq.NVFP4_DEFAULT_CFG, + "fp8": mtq.FP8_DEFAULT_CFG, + "nvfp4_mlp_only": mtq.NVFP4_MLP_ONLY_CFG, +} + + +def _load(export_dir: Path): + return load_safetensors_dir(export_dir), json.loads((export_dir / "config.json").read_text()) + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(torch.uint8) if tensor.element_size() == 1 else tensor + + +def _excluded(patterns, layers): + return {layer for layer in layers if any(fnmatch.fnmatch(layer, p) for p in patterns)} + + +def _quantize(source, quant_cfg_name, 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={"tensor_model_parallel_size": size, "pipeline_dtype": torch.bfloat16}, + moe_grouped_gemm=use_moe_grouped_gemm(str(source)), + ) + language_model, is_vlm = get_language_model(unwrapped) + quant_cfg = QUANT_CFGS[quant_cfg_name] + if is_vlm: + # As in quantize.py, the vision tower stays unquantized. + disable_vision = {"quantizer_name": "vision_model.*", "enable": False} + quant_cfg = {**quant_cfg, "quant_cfg": [*quant_cfg["quant_cfg"], disable_vision]} + torch.manual_seed(0) + input_ids = torch.randint(0, language_model.vocab_size, (2, 16), device="cuda") + mtq.quantize(unwrapped, quant_cfg, lambda m: megatron_prefill(language_model, input_ids)) + return bridge, model, unwrapped + + +def _export_with_megatron_exporter(source, quant_cfg_name, export_dir, rank, size): + _, _, unwrapped = _quantize(source, quant_cfg_name, size) + export_mcore_gpt_to_hf(unwrapped, str(source), dtype=torch.bfloat16, export_dir=str(export_dir)) + + +def _export_with_bridge(source, quant_cfg_name, export_dir, rank, size): + bridge, model, _ = _quantize(source, quant_cfg_name, size) + bridge.save_hf_pretrained_modelopt(model, str(export_dir)) + + +@pytest.fixture +def bridge_workers(request): + # TP=2 on multi-GPU hosts, so the bridge also gathers sharded weights and scales. TP stays <= 2, + # the fewest KV heads of these models. A fixture starts the pool before the test timeout. + return request.getfixturevalue(f"dist_workers_size_{min(torch.cuda.device_count(), 2)}") + + +@pytest.mark.parametrize( + ("create_source", "quant_cfg_name"), + [ + (partial(create_tiny_llama_dir, with_tokenizer=True), "nvfp4"), + (partial(create_tiny_qwen3_dir, with_tokenizer=True), "fp8"), + # The bridge can't export quantized Mamba or GatedDeltaNet projections yet. + (partial(create_tiny_nemotron_h_dir, with_tokenizer=True), "nvfp4_mlp_only"), + (partial(create_tiny_qwen3_5_moe_vl_dir, with_processor=True), "nvfp4_mlp_only"), + ], + ids=["llama-nvfp4", "qwen3-fp8", "nemotron_h-nvfp4_mlp_only", "qwen3_5_moe_vl-nvfp4_mlp_only"], +) +def test_matches_megatron_exporter( + dist_workers_size_1, bridge_workers, tmp_path, create_source, quant_cfg_name +): + source = create_source(tmp_path) + reference_dir, export_dir = tmp_path / "megatron_exporter", tmp_path / "bridge" + # The Megatron exporter needs TP=1. + dist_workers_size_1.run( + partial(_export_with_megatron_exporter, source, quant_cfg_name, reference_dir) + ) + bridge_workers.run(partial(_export_with_bridge, source, quant_cfg_name, export_dir)) + + ref_weights, ref_config = _load(reference_dir) + weights, config = _load(export_dir) + assert weights.keys() == ref_weights.keys() + for name, weight in weights.items(): + reference = ref_weights[name] + # The Megatron exporter keeps per-tensor scales 1-d; both hold the same scalar. + if reference.numel() == weight.numel() == 1: + reference = reference.reshape(weight.shape) + assert weight.shape == reference.shape, name + assert weight.dtype == reference.dtype, name + if name.endswith("input_scale") and bridge_workers.world_size > 1: + # Activations reduced across TP ranks calibrate to slightly different amaxes. + torch.testing.assert_close(weight, reference, rtol=2e-2, atol=0, msg=name) + else: + assert torch.equal(_bits(weight), _bits(reference)), name + + quant, ref_quant = config.pop("quantization_config"), ref_config.pop("quantization_config") + assert config == ref_config + assert {k: v for k, v in quant.items() if k != "ignore"} == { + k: v for k, v in ref_quant.items() if k != "ignore" + } + # The bridge must exclude exactly the unquantized linear layers. The Megatron exporter's + # patterns spell them differently and miss some, e.g. Qwen3.5's mlp.shared_expert_gate. + layers = { + name.rpartition(".")[0] + for name, weight in weights.items() + if name.endswith(".weight") and weight.dim() == 2 + } + unquantized = {layer for layer in layers if f"{layer}.weight_scale" not in weights} + assert _excluded(quant["ignore"], layers) == unquantized + assert _excluded(ref_quant["ignore"], layers) <= unquantized