diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 9a23e8181b0..14cf7247e04 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -155,7 +155,7 @@ jobs: with: docker_image: "nvcr.io/nvidia/nemo:26.08" example: megatron_bridge - timeout_minutes: 60 + timeout_minutes: 75 pip_install_extras: "[hf,puzzletron,dev-test]" runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), ',megatron_bridge,') }} diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 95e3f12ef58..16ec41ae920 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,8 @@ Changelog - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root `` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``. - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD. - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager. +- Add HuggingFace unified export of quantized Qwen3-VL and Qwen3.5-VL checkpoints (PTQ or QAD) via ``examples/megatron_bridge/export_quantized_megatron_to_hf.py``, Qwen3.5-VL additionally covering GatedDeltaNet linear-attention layers and MoE shared experts. Only the language model is quantized; the vision tower is copied from the source HuggingFace checkpoint. +- Megatron-Bridge scripts now choose the MoE expert layout automatically from the model config: the faster fused ``TEGroupedMLP`` (grouped GEMM) unless the architecture cannot export it to HuggingFace, in which case ``SequentialMLP`` keeps the checkpoint exportable and ``--no_moe_grouped_gemm`` forces it explicitly. For the affected architectures this changes MoE activation scales from one shared scale to per-expert. *Misc* @@ -46,6 +48,11 @@ Changelog - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error. +- Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint. +- Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. The exporter now raises instead of writing that checkpoint; the scripts also avoid the situation by selecting ``SequentialMLP`` for those architectures. +- Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name). +- Megatron-Core HuggingFace export now verifies its own output: if the exported checkpoint is missing tensors the source checkpoint has (an architecture whose module has no export rule), it raises instead of writing a valid-looking checkpoint. Depth-pruned models and tied embeddings are accounted for. +- Loading a Megatron checkpoint that holds quantizer tensors but no restorable ModelOpt state now raises instead of silently loading the model unquantized. - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7). - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs. diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 9c5992ead13..06b02d8e804 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -120,7 +120,7 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automati - A **text** dataset runs text-only calibration of the language model (vision tower idle). > [!NOTE] -> HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM is not yet supported; the quantized VLM is saved in Megatron checkpoint format only. +> HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM covers **Qwen3-VL** and **Qwen3.5-VL**. Other VLMs such as Gemma3-VL are saved in Megatron checkpoint format only. ## Distillation diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 7a88172f2e2..25bd9242818 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -50,12 +50,18 @@ from megatron.core.datasets.utils import get_blend_from_list from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.utils import unwrap_model -from transformers import AutoConfig, AutoTokenizer +from transformers import AutoTokenizer import modelopt.torch.distill as mtd import modelopt.torch.utils.distributed as dist +from modelopt.torch.opt.conversion import ModeloptStateManager from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 -from modelopt.torch.utils.plugins.mbridge import load_modelopt_megatron_checkpoint +from modelopt.torch.utils.plugins.mbridge import ( + is_vlm_config, + load_modelopt_megatron_checkpoint, + set_moe_expert_layout, + use_moe_grouped_gemm, +) with contextlib.suppress(ModuleNotFoundError): import modelopt.torch.puzzletron.plugins.mbridge # noqa: F401 @@ -92,6 +98,15 @@ def get_args(): help="HuggingFace model name or path for the teacher (e.g. Qwen/Qwen3-8B)", ) parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code") + parser.add_argument( + "--no_moe_grouped_gemm", + action="store_true", + help=( + "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). " + "By default grouped GEMM is used unless the architecture cannot export it to " + "HuggingFace, in which case SequentialMLP is selected automatically." + ), + ) parser.add_argument( "--student_megatron_path", type=str, @@ -326,11 +341,26 @@ def _tokenizer_prepends_bos(args) -> bool: def main(args: argparse.Namespace): + student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state( + args.student_megatron_path + ) + # A quantized student pins the layout: it must match what quantize.py wrote, so reuse the same + # data-driven choice. An unquantized (e.g. pruned) student exports via Megatron-Bridge, which + # reads either layout, so it keeps the faster grouped GEMM. + moe_grouped_gemm = ( + use_moe_grouped_gemm( + args.student_hf_path, + trust_remote_code=args.trust_remote_code, + force_sequential=args.no_moe_grouped_gemm, + ) + if student_has_modelopt_state + else not args.no_moe_grouped_gemm + ) checkpoint_dir = os.path.join(args.output_dir, "checkpoints") tensorboard_dir = os.path.join(args.output_dir, "tb_logs") # Build student and teacher model providers - def _build_model_provider(hf_path, load_weights=True): + def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True): bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code) provider = bridge.to_megatron_provider(load_weights=load_weights) @@ -343,6 +373,7 @@ def _build_model_provider(hf_path, load_weights=True): provider.expert_model_parallel_size = args.ep_size provider.expert_tensor_parallel_size = 1 # Expert tensor parallelism is not supported provider.seq_length = args.seq_length + set_moe_expert_layout(provider, moe_grouped_gemm) if args.sft: # A response-only loss mask needs per-token reduction to combine across CP ranks. # Must stay in sync with ``average_in_collective=not args.sft`` on the DDP config. @@ -358,17 +389,21 @@ def _build_model_provider(hf_path, load_weights=True): # The student structure is always built from --student_hf_path. When --student_megatron_path is # given, the HF weights are skipped (they are overwritten by the Megatron checkpoint, loaded into # the built student inside the patched provide() below). - student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state( - args.student_megatron_path - ) + # Only the student's layout is pinned -- it must match --student_megatron_path (see quantize.py). student_provider = _build_model_provider( - args.student_hf_path, load_weights=args.student_megatron_path is None + args.student_hf_path, + load_weights=args.student_megatron_path is None, + moe_grouped_gemm=moe_grouped_gemm, ) if student_has_modelopt_state: # Gradient accumulation fusion is not supported with ModelOpt quantized models. Disable it # before the model is built so the student's linear layers are constructed accordingly. student_provider.gradient_accumulation_fusion = False - teacher_provider = _build_model_provider(args.teacher_hf_path) + # The teacher only runs forward, is loaded from HF, and is hidden from the checkpoint + # (``expose_minimal_state_dict``), so it keeps the faster grouped GEMM regardless. + teacher_provider = _build_model_provider( + args.teacher_hf_path, moe_grouped_gemm=not args.no_moe_grouped_gemm + ) # The KD losses compare logits elementwise over the vocab dim, so both output layers must have # the same padded width. A shared tokenizer does not imply it: the HF configs can disagree. @@ -390,10 +425,7 @@ def _build_model_provider(hf_path, load_weights=True): # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under # ``language_model`` (used as ``distill_submodule`` below). - is_vlm = hasattr( - AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), - "vision_config", - ) + is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code) if is_vlm: warn_rank_0( @@ -420,9 +452,14 @@ def _restore_student_hook(model_chunks): print_rank_0( f"Loading student weights from Megatron checkpoint {args.student_megatron_path}" ) - load_modelopt_megatron_checkpoint( - [unwrap_model(model_chunks[0])], args.student_megatron_path - ) + student = unwrap_model(model_chunks[0]) + loaded = load_modelopt_megatron_checkpoint([student], args.student_megatron_path) + if is_vlm and student_has_modelopt_state and loaded[0] is student: + # PTQ stores the state on the VLM root (it quantizes and saves the whole VLM), but + # only ``language_model`` is distilled and checkpointed here, so move it there to + # keep the quantizers across the QAD checkpoint's save / restore. Resuming from a + # language-model-only checkpoint already restores it there. + ModeloptStateManager.transfer_state_dict(student, student.language_model) return model_chunks distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True) diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py index d5e95ff9d18..5957f49b389 100644 --- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -59,14 +59,17 @@ import torch from megatron.bridge import AutoBridge +from megatron.bridge.training.post_training.checkpointing import has_modelopt_state from transformers import AutoConfig import modelopt.torch.utils.distributed as dist from modelopt.torch.export import copy_hf_ckpt_remote_code from modelopt.torch.utils import print_args, print_rank_0 from modelopt.torch.utils.plugins.mbridge import ( + is_vlm_config, load_mbridge_model_from_hf, load_modelopt_megatron_checkpoint, + use_moe_grouped_gemm, ) # Megatron-Bridge checkpoint iteration directories use names like ``iter_0000100``. @@ -213,6 +216,16 @@ def get_args() -> argparse.Namespace: "correct for homogeneous students; unused for VLMs.", ) parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code") + parser.add_argument( + "--no_moe_grouped_gemm", + action="store_true", + help=( + "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). " + "By default grouped GEMM is used unless the architecture cannot export it to " + "HuggingFace, in which case SequentialMLP is selected automatically. VLMs only: the " + "LLM path reads the expert layout from the checkpoint and ignores this flag." + ), + ) parser.add_argument("--tp_size", type=int, default=1, help="Tensor parallel size") parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size") parser.add_argument("--ep_size", type=int, default=1, help="Expert parallel size") @@ -226,10 +239,15 @@ def get_args() -> argparse.Namespace: def main(args: argparse.Namespace): checkpoint_export_paths: list[tuple[Path, Path]] = _get_checkpoint_export_paths(args) - is_vlm = hasattr( - AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), - "vision_config", - ) + # This path drops quantization, so a QAD checkpoint would export silently unquantized. + # ``has_modelopt_state`` ignores ``kd_loss``, so plain distillation still passes. + quantized = [str(p) for p, _ in checkpoint_export_paths if has_modelopt_state(str(p))] + if quantized: + raise ValueError( + f"{quantized[0]} is quantized; this script exports full precision only and would drop " + "the quantizers. Use export_quantized_megatron_to_hf.py instead." + ) + is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code) if is_vlm: # Build the full VLM (vision tower / projector + original LM from HF), then overwrite the LM @@ -238,6 +256,11 @@ def main(args: argparse.Namespace): _bridge, _provider, _model, full_model, _tokenizer = load_mbridge_model_from_hf( hf_model_name_or_path=args.student_hf_path, trust_remote_code=args.trust_remote_code, + moe_grouped_gemm=use_moe_grouped_gemm( + args.student_hf_path, + trust_remote_code=args.trust_remote_code, + force_sequential=args.no_moe_grouped_gemm, + ), provider_overrides={ "tensor_model_parallel_size": args.tp_size, "pipeline_model_parallel_size": args.pp_size, diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index e4e3703d8a5..8889e1983ae 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -47,6 +47,7 @@ from modelopt.torch.utils.plugins.mbridge import ( load_mbridge_model_from_hf, load_modelopt_megatron_checkpoint, + use_moe_grouped_gemm, ) @@ -71,6 +72,15 @@ def get_args() -> argparse.Namespace: help="Directory to write the exported HuggingFace (unified) checkpoint to.", ) parser.add_argument("--trust_remote_code", action="store_true") + parser.add_argument( + "--no_moe_grouped_gemm", + action="store_true", + help=( + "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). " + "By default grouped GEMM is used unless the architecture cannot export it to " + "HuggingFace, in which case SequentialMLP is selected automatically." + ), + ) parser.add_argument( "--export_extra_modules", action="store_true", @@ -108,6 +118,11 @@ def main(args: argparse.Namespace): _bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf( hf_model_name_or_path=args.hf_model_name_or_path, trust_remote_code=trust_remote_code, + moe_grouped_gemm=use_moe_grouped_gemm( + args.hf_model_name_or_path, + trust_remote_code=trust_remote_code, + force_sequential=args.no_moe_grouped_gemm, + ), provider_overrides={ "tensor_model_parallel_size": 1, # Tensor parallelism is not supported "pipeline_model_parallel_size": args.pp_size, @@ -144,9 +159,7 @@ def main(args: argparse.Namespace): print_rank_0( f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..." ) - # TODO (OMNIML-5366): quantized-VLM HF export. export_mcore_gpt_to_hf's per-arch mappings don't - # cover Qwen3.5-VL / Gemma3-VL; See if Megatron-Bridge's AutoBridge.export_hf_weights_quant can be - # used instead. + # TODO: Gemma3-VL is not in export_mcore_gpt_to_hf's per-arch mappings yet. export_mcore_gpt_to_hf( unwrapped_model, args.hf_model_name_or_path, diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 42cb78f2a7c..974de2192a8 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -66,7 +66,7 @@ print_rank_0, warn_rank_0, ) -from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf +from modelopt.torch.utils.plugins.mbridge import get_language_model, load_mbridge_model_from_hf from modelopt.torch.utils.plugins.megatron_calibration import ( get_megatron_calibration_forward_loop, get_megatron_vlm_calibration_forward_loop, @@ -430,8 +430,7 @@ def main(args: argparse.Namespace): # For VLMs (e.g. Qwen3-VL), only the language model is pruned; the vision tower is left intact. # hidden_size is shared with the vision->LM projector, so it is skipped - language_model = getattr(unwrapped_model, "language_model", unwrapped_model) - is_vlm = language_model is not unwrapped_model + language_model, is_vlm = get_language_model(unwrapped_model) if is_vlm: warn_rank_0( "VLM detected: pruning model.language_model only; all non-language-model components " diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index fedd7bf62d0..954fa0a4702 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -62,6 +62,7 @@ import gc import torch +from megatron.bridge.models.hf_pretrained.utils import is_safe_repo from transformers import AutoProcessor import modelopt.torch.quantization as mtq @@ -70,7 +71,11 @@ from modelopt.recipe.presets import KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0 from modelopt.torch.utils.dataset_utils import get_supported_datasets -from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_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_calibration import ( get_megatron_calibration_forward_loop, get_megatron_vlm_calibration_forward_loop, @@ -96,6 +101,15 @@ def get_args() -> argparse.Namespace: parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument("--hf_model_name_or_path", type=str, required=True) parser.add_argument("--trust_remote_code", action="store_true") + parser.add_argument( + "--no_moe_grouped_gemm", + action="store_true", + help=( + "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). " + "By default grouped GEMM is used unless the architecture cannot export it to " + "HuggingFace, in which case SequentialMLP is selected automatically." + ), + ) parser.add_argument( "--export_megatron_path", type=str, @@ -276,9 +290,20 @@ def get_quant_config(args: argparse.Namespace) -> dict: def main(args: argparse.Namespace): + trust_remote_code = is_safe_repo( + trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path + ) + + moe_grouped_gemm = use_moe_grouped_gemm( + args.hf_model_name_or_path, + trust_remote_code=trust_remote_code, + force_sequential=args.no_moe_grouped_gemm, + ) + bridge, _provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf( hf_model_name_or_path=args.hf_model_name_or_path, - trust_remote_code=args.trust_remote_code, + trust_remote_code=trust_remote_code, + moe_grouped_gemm=moe_grouped_gemm, provider_overrides={ "tensor_model_parallel_size": args.tp_size, "pipeline_model_parallel_size": args.pp_size, @@ -293,8 +318,7 @@ def main(args: argparse.Namespace): ) # Only the language model is quantized (vision tower + projector stay full precision) - language_model = getattr(unwrapped_model, "language_model", unwrapped_model) - is_vlm = language_model is not unwrapped_model + language_model, is_vlm = get_language_model(unwrapped_model) if is_vlm: warn_rank_0( "VLM detected: quantizing `model.language_model` only (vision tower left in full precision)." @@ -377,7 +401,7 @@ def forward_loop(_model=None): # VLMs: drive the full VLM forward on image-text pairs so the language model's quantizers # see vision-conditioned activations (we still quantize the LM only). processor = AutoProcessor.from_pretrained( - args.hf_model_name_or_path, trust_remote_code=args.trust_remote_code + args.hf_model_name_or_path, trust_remote_code=trust_remote_code ) forward_loop = get_megatron_vlm_calibration_forward_loop( unwrapped_model, # full VLM (vision encoder + projector + language model) @@ -412,12 +436,12 @@ def forward_loop(_model=None): model, args.export_megatron_path, hf_tokenizer_path=args.hf_model_name_or_path, - hf_tokenizer_kwargs={"trust_remote_code": args.trust_remote_code}, + hf_tokenizer_kwargs={"trust_remote_code": trust_remote_code}, ) if is_vlm: print_rank_0( - f"\nSaved quantized VLM to {args.export_megatron_path} in Megatron format " - "(HuggingFace unified export of a quantized VLM is not yet supported)." + f"\nSaved quantized VLM to {args.export_megatron_path} in Megatron format. To deploy this " + "model, convert it to a Unified HF ckpt with export_quantized_megatron_to_hf.py." ) else: print_rank_0( diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py index 54058c742db..b2046fe081a 100644 --- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py @@ -205,7 +205,7 @@ def load_multimodal_components( """Load multimodal components from safetensors file. Args: - pretrained_model_path: Path to the pretrained model. + pretrained_model_path: Directory or HuggingFace repo id of the pretrained model. prefixes: Tensor key prefixes to select. Defaults to the LLaVA-style ``multi_modal_projector`` / ``vision_model`` prefixes. Pass ``("model.visual.",)`` for Qwen3-VL checkpoints. @@ -215,9 +215,21 @@ def load_multimodal_components( """ hf_checkpoint_path = Path(pretrained_model_path) if not hf_checkpoint_path.is_dir(): - raise ValueError( - f"Invalid pretrained model path: {pretrained_model_path}. It should be a directory." - ) + # Also accept a repo id, which is what the example scripts pass to quantize.py. + local_files_only = _is_hf_hub_offline() + try: + hf_checkpoint_path = Path( + snapshot_download( + repo_id=str(pretrained_model_path), + allow_patterns=["*.safetensors", "*.safetensors.index.json"], + local_files_only=local_files_only, + ) + ) + except (LocalEntryNotFoundError, OSError, ValueError) as exc: + raise ValueError( + f"Invalid pretrained model path: {pretrained_model_path}. It should be a " + "directory or an available HuggingFace repo id." + ) from exc safetensors_file = Path(hf_checkpoint_path) / "model.safetensors" safetensors_index_file = Path(hf_checkpoint_path) / "model.safetensors.index.json" diff --git a/modelopt/torch/export/plugins/mcore_common.py b/modelopt/torch/export/plugins/mcore_common.py index 15395b7a1e5..d275aee4c0a 100644 --- a/modelopt/torch/export/plugins/mcore_common.py +++ b/modelopt/torch/export/plugins/mcore_common.py @@ -39,7 +39,12 @@ qwen25_causal_lm_export, qwen25_causal_lm_import, ) -from .mcore_qwen3vl import qwen3vl_causal_lm_export, qwen3vl_causal_lm_import +from .mcore_qwen3vl import ( + QWEN3VL_VISION_PREFIXES, + qwen3vl_causal_lm_export, + qwen3vl_causal_lm_import, +) +from .mcore_qwen35vl import QWEN3_5_VL_VISION_PREFIXES, qwen3_5_vl_causal_lm_export all_mcore_hf_export_mapping: dict[str, Any] = { "DeepseekV2ForCausalLM": deepseek_causal_lm_export, @@ -56,6 +61,17 @@ "Qwen2ForCausalLM": qwen25_causal_lm_export, "GptOssForCausalLM": gptoss_causal_lm_export, "Qwen3VLForConditionalGeneration": qwen3vl_causal_lm_export, + "Qwen3_5ForConditionalGeneration": qwen3_5_vl_causal_lm_export, + "Qwen3_5MoeForConditionalGeneration": qwen3_5_vl_causal_lm_export, +} + +# VLM architectures whose Megatron export covers the language model only: the vision tower is copied +# verbatim from the HF checkpoint under these key prefixes. Architectures absent here fall back to +# ``LLAVA_VISION_PREFIXES`` when the model is an MCore ``LLaVAModel``, and to no copy otherwise. +all_mcore_hf_vision_passthrough_mapping: dict[str, tuple[str, ...]] = { + "Qwen3VLForConditionalGeneration": QWEN3VL_VISION_PREFIXES, + "Qwen3_5ForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES, + "Qwen3_5MoeForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES, } all_mcore_hf_import_mapping: dict[str, Any] = { diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py index f231da48433..87b696e641f 100644 --- a/modelopt/torch/export/plugins/mcore_custom.py +++ b/modelopt/torch/export/plugins/mcore_custom.py @@ -16,6 +16,7 @@ """Custom Megatron mapping and safetensors utility.""" +import copy import json import math import os @@ -127,6 +128,18 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] ) +class GroupedMLPPacking(CustomModuleMapping): + """A custom module mapping that packs grouped MoE experts into one stacked tensor.""" + + def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}): + """Create a custom module mapping that packs grouped experts.""" + super().__init__( + func_name="grouped_mlp_packing", + target_name_or_prefix=target_name_or_prefix, + func_kwargs=func_kwargs, + ) + + class GatedMLPMerging(CustomModuleMapping): """A custom module mapping that merges gate_proj and up_proj.""" @@ -175,6 +188,18 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] ) +class GatedDeltaNetSlicing(CustomModuleMapping): + """A custom module mapping that splits GatedDeltaNet's fused ``in_proj``.""" + + def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}): + """Create a custom module mapping that splits the fused GatedDeltaNet input projection.""" + super().__init__( + func_name="gated_delta_net_slicing", + target_name_or_prefix=target_name_or_prefix, + func_kwargs=func_kwargs, + ) + + class PackNameRemapping(CustomModuleMapping): """A custom module mapping that packs module after name remapping.""" @@ -223,6 +248,29 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] ) +# LLaVA-style checkpoints keep the vision tower under these prefixes; see +# ``all_mcore_hf_vision_passthrough_mapping`` for the per-architecture overrides. +LLAVA_VISION_PREFIXES = ("multi_modal_projector", "vision_model") + + +def with_language_model_prefix( + mapping: dict[str, CustomModuleMapping], +) -> dict[str, CustomModuleMapping]: + """Nest a text-model mapping under ``model.language_model.``; other prefixes are unchanged.""" + result = {} + for key, m in mapping.items(): + if not isinstance(m, CustomModuleMapping): + result[key] = m # plain flags such as ``use_packed_local_experts`` + continue + prefix = m.target_name_or_prefix + if prefix.startswith("model."): + prefix = "model.language_model." + prefix[len("model.") :] + result[key] = type(m)( + target_name_or_prefix=prefix, func_kwargs=copy.deepcopy(m.func_kwargs) + ) + return result + + def save_safetensors(state_dict, save_directory: str | os.PathLike): """Save safetensors with pipeline model parallel support.""" pp_rank = get_pipeline_model_parallel_rank() diff --git a/modelopt/torch/export/plugins/mcore_qwen35vl.py b/modelopt/torch/export/plugins/mcore_qwen35vl.py new file mode 100644 index 00000000000..737ffdebb22 --- /dev/null +++ b/modelopt/torch/export/plugins/mcore_qwen35vl.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. + +"""Custom mapping from Qwen3.5-VL Hugging Face models to Megatron Core models. + +Qwen3.5 interleaves GatedDeltaNet linear-attention layers (fused ``in_proj``, split here into +HF's ``in_proj_qkv`` / ``_z`` / ``_b`` / ``_a``) with gated full-attention layers, and adds MoE +shared experts. Only the language model is exported; the vision tower is copied from HF. +Routed experts are packed as ``[num_experts, out, in]``, exported from either expert layout. +""" + +from .mcore_custom import ( + GatedDeltaNetSlicing, + GatedMLPSlicing, + GroupedMLPPacking, + NameRemapping, + PackNameRemapping, + with_language_model_prefix, +) +from .mcore_qwen import qwen3_causal_lm_export + +# Vision-tower weights copied straight from the HF checkpoint (never quantized). +QWEN3_5_VL_VISION_PREFIXES = ("model.visual.",) + +# Qwen3.5 adds linear-attention layers and shared experts on top of the Qwen3 rules. +_qwen3_5_extra_export: dict = { + # Linear attention (GatedDeltaNet). ``linear_attn`` splits the fused in_proj; the rest + # are plain renames of the surrounding parameters. + "linear_attn": GatedDeltaNetSlicing("model.layers.{}.linear_attn."), + "linear_attn.conv1d": NameRemapping("model.layers.{}.linear_attn.conv1d."), + "linear_attn.A_log": NameRemapping("model.layers.{}.linear_attn.A_log"), + "linear_attn.dt_bias": NameRemapping("model.layers.{}.linear_attn.dt_bias"), + # Megatron's GDN output norm is zero-centered; HF's RMSNorm gamma is centered on 1. + "linear_attn.out_norm": NameRemapping( + "model.layers.{}.linear_attn.norm.", {"zero_centered_gamma": True} + ), + "linear_attn.out_proj": NameRemapping("model.layers.{}.linear_attn.out_proj."), + # Routed experts are stored packed: [num_experts, out, in], keeping Megatron's orientation. + "use_packed_local_experts": True, + "local_experts.linear_fc1": PackNameRemapping( + "model.layers.{}.mlp.experts.gate_up_proj", + {"layer_type": "linear_fc1", "transpose": False}, + ), + "local_experts.linear_fc2": PackNameRemapping( + "model.layers.{}.mlp.experts.down_proj", + {"layer_type": "linear_fc2", "transpose": False}, + ), + # Same packed layout from fused TEGroupedMLP, so grouped GEMM stays usable. + "experts.linear_fc1": GroupedMLPPacking("model.layers.{}.mlp.experts.gate_up_proj"), + "experts.linear_fc2": GroupedMLPPacking("model.layers.{}.mlp.experts.down_proj"), + # MoE shared experts (routed experts + router come from the Qwen3 rules). + "shared_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.shared_expert."), + "shared_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.shared_expert.down_proj."), + "shared_experts.gate_weight": NameRemapping("model.layers.{}.mlp.shared_expert_gate.weight"), +} + +qwen3_5_vl_causal_lm_export = with_language_model_prefix( + {**qwen3_causal_lm_export, **_qwen3_5_extra_export} +) diff --git a/modelopt/torch/export/plugins/mcore_qwen3vl.py b/modelopt/torch/export/plugins/mcore_qwen3vl.py index 1f2d3830d61..6312be1abb5 100644 --- a/modelopt/torch/export/plugins/mcore_qwen3vl.py +++ b/modelopt/torch/export/plugins/mcore_qwen3vl.py @@ -15,48 +15,17 @@ """Custom mapping from Qwen3-VL Hugging Face models to Megatron Core models. -Qwen3-VL differs from Qwen3 in one structural way: language-model weights live -under ``model.language_model.`` instead of ``model.``, while ``lm_head.weight`` -remains at the root level. The mappings below are derived automatically from -the Qwen3 mappings by inserting ``language_model.`` after ``model.`` for every -prefix that starts with ``model.``. - -Note: the visual encoder (``model.visual.*``) is intentionally excluded — this -mapping covers only the language-model decoder used for quantization and export. - -Note: ``Qwen3VLMoeForConditionalGeneration`` is **not** supported here. The MoE -variant stores expert weights as 3-D tensors (``mlp.experts.gate_up_proj``, -``mlp.experts.down_proj``) that require a dedicated fused-expert mapping and -cannot reuse the dense Qwen3 rules. - -Reference: https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct/blob/main/model.safetensors.index.json +Qwen3-VL nests the language model under ``model.language_model.`` while ``lm_head`` stays at the +root, so the mappings are derived from Qwen3's. The visual encoder is copied verbatim from HF via +``QWEN3VL_VISION_PREFIXES`` rather than mapped. ``Qwen3VLMoeForConditionalGeneration`` is not +supported: its 3-D fused expert weights cannot reuse the dense Qwen3 rules. """ -import copy - -from .mcore_custom import CustomModuleMapping +from .mcore_custom import with_language_model_prefix from .mcore_qwen import qwen3_causal_lm_export, qwen3_causal_lm_import +# Vision-tower weights copied straight from the HF checkpoint (never quantized). +QWEN3VL_VISION_PREFIXES = ("model.visual.",) -def _with_language_model_prefix( - mapping: dict[str, CustomModuleMapping], -) -> dict[str, CustomModuleMapping]: - """Derive a VL mapping from a base Qwen3 mapping. - - Rewrites every ``target_name_or_prefix`` that starts with ``model.`` to - ``model.language_model.``. Prefixes that do not start with - ``model.`` (e.g. ``lm_head.``) are left unchanged. - """ - result = {} - for key, m in mapping.items(): - prefix = m.target_name_or_prefix - if prefix.startswith("model."): - prefix = "model.language_model." + prefix[len("model.") :] - result[key] = type(m)( - target_name_or_prefix=prefix, func_kwargs=copy.deepcopy(m.func_kwargs) - ) - return result - - -qwen3vl_causal_lm_import = _with_language_model_prefix(qwen3_causal_lm_import) -qwen3vl_causal_lm_export = _with_language_model_prefix(qwen3_causal_lm_export) +qwen3vl_causal_lm_import = with_language_model_prefix(qwen3_causal_lm_import) +qwen3vl_causal_lm_export = with_language_model_prefix(qwen3_causal_lm_export) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index bbbd21ad244..070ad8d34e4 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -21,6 +21,7 @@ import io import json import os +import re import tempfile from collections import OrderedDict from pathlib import Path @@ -53,8 +54,12 @@ copy_non_safetensor_files_from_ckpt, load_multimodal_components, ) -from .plugins.mcore_common import all_mcore_hf_export_mapping +from .plugins.mcore_common import ( + all_mcore_hf_export_mapping, + all_mcore_hf_vision_passthrough_mapping, +) from .plugins.mcore_custom import ( + LLAVA_VISION_PREFIXES, CustomModuleMapping, get_safetensor, save_safetensors_by_layer_index, @@ -133,7 +138,14 @@ def __init__( moe_router_dtype: str | None = None, ): """Create a GPTModel exporter instance.""" - if not isinstance(model, (GPTModel, HybridModel, LLaVAModel)): + # VLM wrappers keep the decoder under ``.language_model``; only that is exported, the + # vision tower being copied from the HF checkpoint as-is. + language_model = ( + model + if isinstance(model, (GPTModel, HybridModel)) + else getattr(model, "language_model", None) + ) + if not isinstance(language_model, (GPTModel, HybridModel)): raise ValueError("Input to GPTModelExport must be a megatron.core.models.GPTModel!") self._state_dict = OrderedDict() @@ -153,22 +165,25 @@ def __init__( self._hf_text_config = getattr(self._hf_config, "text_config", self._hf_config) # Update hf_config - self._hf_text_config.num_hidden_layers = model.config.num_layers - self._hf_text_config.hidden_size = model.config.hidden_size - self._hf_text_config.head_dim = model.config.kv_channels - self._hf_text_config.num_attention_heads = model.config.num_attention_heads - self._hf_text_config.num_key_value_heads = model.config.num_query_groups + self._hf_text_config.num_hidden_layers = language_model.config.num_layers + self._hf_text_config.hidden_size = language_model.config.hidden_size + self._hf_text_config.head_dim = language_model.config.kv_channels + self._hf_text_config.num_attention_heads = language_model.config.num_attention_heads + self._hf_text_config.num_key_value_heads = language_model.config.num_query_groups self.is_multimodal = isinstance(model, LLaVAModel) if not self.is_multimodal: - self._hf_text_config.intermediate_size = model.config.ffn_hidden_size + self._hf_text_config.intermediate_size = language_model.config.ffn_hidden_size self._hf_quant_config: dict = {} self._hf_extra_config = None self.export_extra_modules = export_extra_modules - self.is_multimodal = isinstance(model, LLaVAModel) - self.model = model.language_model if self.is_multimodal else model + self.model = language_model self.dtype = dtype self.trust_remote_code = trust_remote_code self.arch = self._hf_config.architectures[0] + # ``None`` when there is no vision tower to copy through. + self.vision_passthrough_prefixes = all_mcore_hf_vision_passthrough_mapping.get( + self.arch, LLAVA_VISION_PREFIXES if self.is_multimodal else None + ) # TODO: May modify this later according to what quantization exported ckpt is, currently only support BF16. if self.arch == "GptOssForCausalLM": if hasattr(self._hf_config, "quantization_config"): @@ -393,14 +408,13 @@ def save_pretrained( # Merge the multimodal components into that shard so they land in a file # the index builder picks up (it scans shards 1..num_layers). first_layer_key = next(iter(layer_state_dicts)) - if self.is_multimodal: - multimodal_state_dict = load_multimodal_components(pretrained_model_name_or_path) - layer_state_dicts[first_layer_key].update(multimodal_state_dict) - elif self.arch == "Qwen3VLForConditionalGeneration": - vision_state_dict = load_multimodal_components( - pretrained_model_name_or_path, prefixes=("model.visual.",) + if self.vision_passthrough_prefixes is not None: + layer_state_dicts[first_layer_key].update( + load_multimodal_components( + pretrained_model_name_or_path, + prefixes=self.vision_passthrough_prefixes, + ) ) - layer_state_dicts[first_layer_key].update(vision_state_dict) # Bracket the writer's config.json read-modify-write with barriers so peers # never observe a truncated file (also ensures export_dir exists). @@ -424,6 +438,51 @@ def save_pretrained( name_template="model-{:05d}-of-{:05d}", ) + # Every rank has written its shards; one rank now checks nothing was dropped. + torch.distributed.barrier() + if is_writer_rank: + self._verify_exported_keys(save_directory, pretrained_model_name_or_path) + + def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) -> None: + """Raise if the export dropped tensors the source has: a missing rule emits nothing.""" + if pretrained_model_name_or_path is None or not os.path.isdir( + str(pretrained_model_name_or_path) + ): + return # hub id: not worth a download inside export + index_file = Path(save_directory) / "model.safetensors.index.json" + if not index_file.exists(): + return + with open(index_file) as f: + exported = set(json.load(f)["weight_map"]) + source = _read_checkpoint_keys(pretrained_model_name_or_path) + if not source: + return + + # Narrow on purpose: compare module prefixes, not tensor names, since a quantized source + # carries extras with no export counterpart, and only inside decoder layers, whose naming + # is stable. A dropped decoder module is the case that loads fine and produces garbage. + num_layers = self.model.config.num_layers + exported_modules = {key.rsplit(".", 1)[0] for key in exported} + missing = set() + for key in source - exported: + layer = re.search(r"\.layers\.(\d+)\.", key) + if layer is None: + continue # see the note above: decoder layers only + if int(layer.group(1)) >= num_layers: + continue # depth-pruned model: the source has layers this export does not + if key.rsplit(".", 1)[0] in exported_modules: + continue # module is exported; this name is a source-side quantization artifact + if "rotary_emb" in key: + continue # non-persistent buffer some conversions still ship + missing.add(key) + if missing: + raise RuntimeError( + f"Export dropped {len(missing)} tensor(s) present in " + f"{pretrained_model_name_or_path}, e.g. {sorted(missing)[:8]}. The checkpoint " + f"written to {save_directory} is incomplete -- the architecture has no export " + "rule for one of its decoder modules." + ) + @property def state_dict(self): """Return the real quantized state_dict of the base model.""" @@ -497,8 +556,12 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): if not isinstance(layer.input_layernorm, IdentityOp): self.rules["input_layernorm"](layer.input_layernorm, layer_id, is_mtp=is_mtp) else: + # GatedDeltaNet fuses the input layernorm into ``in_proj`` rather than ``linear_qkv``. + qkv_module = getattr(layer.self_attention, "linear_qkv", None) + if qkv_module is None: + qkv_module = getattr(layer.self_attention, "in_proj", None) fused_key, norm_weight = self._get_fused_norm_weight( - getattr(layer.self_attention, "linear_qkv", None), + qkv_module, primary_key="fused_input_layernorm", ) if norm_weight is not None: @@ -531,6 +594,9 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): layer.self_attention.linear_kv_up_proj, layer_id, is_mtp=is_mtp ) self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id, is_mtp=is_mtp) + elif "linear_attn" in self.rules and hasattr(layer.self_attention, "in_proj"): + # GatedDeltaNet (Qwen3.5 linear attention): no q/k layernorm, no core_attention. + self._get_gated_delta_net_state_dict(layer, layer_id, is_mtp=is_mtp) else: if layer.self_attention.q_layernorm is not None and not isinstance( layer.self_attention.q_layernorm, (IdentityOp, L2Norm) @@ -585,6 +651,13 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): self.rules["shared_experts.linear_fc2"]( layer.mlp.shared_experts.linear_fc2, layer_id, is_mtp=is_mtp ) + if ( + "shared_experts.gate_weight" in self.rules + and getattr(layer.mlp.shared_experts, "gate_weight", None) is not None + ): + self.rules["shared_experts.gate_weight"]( + layer.mlp.shared_experts.gate_weight, layer_id, is_mtp=is_mtp + ) if hasattr(layer.mlp.experts, "local_experts"): if not self.rules.get("use_packed_local_experts", False): for expert_id, expert in enumerate(layer.mlp.experts.local_experts): @@ -613,6 +686,14 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): self.rules["experts.linear_fc2"]( layer.mlp.experts.linear_fc2, layer_id, is_mtp=is_mtp ) + else: + # Otherwise the routed experts are dropped and the checkpoint looks valid. + raise NotImplementedError( + f"No export rule for {type(layer.mlp.experts).__name__} experts of " + f"{self.arch}: fused (grouped GEMM) experts need an 'experts.linear_fc1' " + "rule. Re-run quantization and export with --no_moe_grouped_gemm to build " + "the experts as SequentialMLP instead." + ) else: self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id, is_mtp=is_mtp) self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id, is_mtp=is_mtp) @@ -629,7 +710,9 @@ def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]: self._state_dict = OrderedDict() try: for mtp_layer in mtp.layers: - inner_layers = mtp_layer.mtp_model_layer.layers + # Some architectures (Qwen3.5) put a single TransformerLayer here, not a container. + inner = mtp_layer.mtp_model_layer + inner_layers = getattr(inner, "layers", None) or [inner] first_id = inner_layers[0].layer_number - 1 last_id = inner_layers[-1].layer_number - 1 @@ -728,6 +811,16 @@ def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]: self.exclude_modules.append("mtp*") return mtp_state_dict + def _get_gated_delta_net_state_dict(self, layer, layer_id, is_mtp=False): + """Export a GatedDeltaNet (Qwen3.5 linear-attention) layer's ``self_attention``.""" + gdn = layer.self_attention + self.rules["linear_attn"](gdn, layer_id, is_mtp=is_mtp) + self.rules["linear_attn.conv1d"](gdn.conv1d, layer_id, is_mtp=is_mtp) + self.rules["linear_attn.A_log"](gdn.A_log, layer_id, is_mtp=is_mtp) + self.rules["linear_attn.dt_bias"](gdn.dt_bias, layer_id, is_mtp=is_mtp) + self.rules["linear_attn.out_norm"](gdn.out_norm, layer_id, is_mtp=is_mtp) + self.rules["linear_attn.out_proj"](gdn.out_proj, layer_id, is_mtp=is_mtp) + def _get_mamba_layer_state_dict(self, layer, layer_id, is_mtp=False): if not isinstance(layer.norm, IdentityOp): self.rules["norm"](layer.norm, layer_id, is_mtp=is_mtp) @@ -863,7 +956,9 @@ def _custom_mapping_to_lambda(mapping): "qkv_slicing": self._qkv_slicing, "self_attention_scaling": self._self_attention_scaling, "gated_mlp_slicing": self._gated_mlp_slicing, + "gated_delta_net_slicing": self._gated_delta_net_slicing, "grouped_mlp_slicing": self._grouped_mlp_slicing, + "grouped_mlp_packing": self._grouped_mlp_packing, "pack_name_remapping": self._pack_name_remapping, "pack_name_remapping_gpt_oss": self._pack_name_remapping_gpt_oss, } @@ -1010,8 +1105,11 @@ def _record_excluded_module(self, prefix: str): def _mtp_prefix(prefix: str) -> str: """Rewrite a base-model target prefix (backbone/model root) to its MTP counterpart.""" if "backbone" in prefix: - return prefix.replace("backbone", "mtp") - return prefix.replace("model", "mtp") + return prefix.replace("backbone", "mtp", 1) + # Replace the root only: a VLM's "model.language_model." must not become "mtp.language_mtp.". + if prefix.startswith("model.language_model."): + return "mtp." + prefix[len("model.language_model.") :] + return prefix.replace("model", "mtp", 1) def _name_remapping( self, @@ -1021,6 +1119,7 @@ def _name_remapping( mapping={}, dtype: torch.dtype | None = None, is_mtp: bool = False, + zero_centered_gamma: bool = False, ): if is_mtp: prefix = self._mtp_prefix(prefix) @@ -1028,13 +1127,21 @@ def _name_remapping( dtype = self.dtype if isinstance(module, torch.Tensor): - self._state_dict[prefix] = module + self._state_dict[prefix] = (module + 1.0) if zero_centered_gamma else module return name_to_value, qformat, block_size = self._get_quantized_state(module, dtype, prefix=prefix) self._record_layer_quant_config(prefix, qformat, block_size) weight = name_to_value.pop("weight") + if zero_centered_gamma: + # Megatron centres this gamma on 0, HF on 1. Assert, don't derive: the config flag + # is model-wide while this one is per-norm. + assert getattr(module, "zero_centered_gamma", False), ( + f"{prefix} is mapped as zero-centered gamma but the module reports otherwise; " + "exporting would shift the weights by 1.0" + ) + weight = weight + 1.0 weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) if weight_scale is None: @@ -1125,9 +1232,87 @@ def _gated_mlp_slicing( self._state_dict[gate_proj_key] = val.detach().clone() self._state_dict[up_proj_key] = val.detach().clone() - def _grouped_mlp_slicing(self, module, prefix, parallel_config=None, is_mtp=False): + def _grouped_mlp_packing(self, module, prefix, parallel_config=None, is_mtp=False): + """Pack TEGroupedMLP experts into one ``[num_experts, out, in]`` tensor (Qwen3.5 layout). + + Reuses the per-expert path for the EP gather and per-expert quantizers, then stacks by + global expert id and re-quantizes once with the max scale, as ``_pack_name_remapping`` does. + """ + if is_mtp: + prefix = self._mtp_prefix(prefix) + marker = "\x00pack\x00" + saved_state_dict = self._state_dict + self._state_dict = OrderedDict() + try: + qformat, block_size = self._grouped_mlp_slicing( + module, + marker + "{}", + parallel_config=parallel_config, + is_mtp=False, + quantize=False, + record_quant_config=False, + ) + per_expert = self._state_dict + finally: + self._state_dict = saved_state_dict + + def collect(suffix): + found = {} + for key, value in per_expert.items(): + if not key.startswith(marker) or not key.endswith(suffix): + continue + found[int(key[len(marker) :].split(".", 1)[0])] = value + return [found[i] for i in sorted(found)] + + weights = collect(".weight") + if not weights: + return + # Record against the packed prefix, as _pack_name_remapping does for the other packed path. + if qformat in (None, QUANTIZATION_NONE): + self._record_excluded_module(prefix) + else: + assert block_size is not None + self._record_layer_quant_config(prefix, qformat, block_size) + scales, scales_2 = collect(".weight_scale"), collect(".weight_scale_2") + input_scales = collect(".input_scale") + + # Quantize once over the stack, exactly as _pack_name_remapping does. + merged_weight = torch.stack(weights, dim=0) + if not scales: + self._state_dict[prefix] = merged_weight + else: + if scales_2: + # NVFP4 keeps each expert's block scales; only the global scale is merged. + merged_scale = torch.stack(scales, dim=0) + merged_scale_2 = torch.max(torch.stack(scales_2, dim=0), dim=0)[0] + else: + merged_scale, merged_scale_2 = torch.max(torch.stack(scales, dim=0), dim=0)[0], None + self._state_dict[prefix] = to_quantized_weight( + merged_weight, merged_scale, qformat, merged_scale_2, block_size + ) + # Same suffixes as _pack_name_remapping so both packed paths agree. + self._state_dict[prefix + "_weight_scale"] = merged_scale + if merged_scale_2 is not None: + self._state_dict[prefix + "_weight_scale_2"] = merged_scale_2 + if input_scales: + self._state_dict[prefix + "_input_scale"] = torch.max( + torch.stack(input_scales, dim=0), dim=0 + )[0] + + def _grouped_mlp_slicing( + self, + module, + prefix, + parallel_config=None, + is_mtp=False, + quantize=True, + record_quant_config=True, + ): """Export TEGroupedMLP weight0..weight{N-1} as one HF-style entry per expert. + ``quantize=False`` emits unquantized weights alongside the scales, which + ``_grouped_mlp_packing`` needs so it can quantize once over the stacked tensor. + At EP>1, local ids are mapped to global via ``module.local_expert_indices`` and per-expert state is ``all_gather_object``-ed across the EP group. All EP ranks MUST enter this method for the same layer in lockstep or the gather hangs. @@ -1245,12 +1430,16 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None, is_mtp=Fals if weight_scale_cpu is None: local_expert_state[expert_prefix + "weight"] = weight else: - local_expert_state[expert_prefix + "weight"] = to_quantized_weight( - weight, - weight_scale_cpu, - qformat, - weight_scale_2_cpu, - block_size, + local_expert_state[expert_prefix + "weight"] = ( + weight + if not quantize + else to_quantized_weight( + weight, + weight_scale_cpu, + qformat, + weight_scale_2_cpu, + block_size, + ) ) local_expert_state[expert_prefix + "weight_scale"] = weight_scale_cpu.clone() @@ -1274,7 +1463,7 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None, is_mtp=Fals # Record quant config for ALL global experts on every rank; otherwise the writer's # hf_quant_config.json would miss (EP-1)/EP of the routed experts. All experts in # a TEGroupedMLP layer share qformat/block_size, so local values apply globally. - if seen_qformat is not None: + if seen_qformat is not None and record_quant_config: assert seen_block_size is not None num_total_experts = num_experts * ep_size for global_id in range(num_total_experts): @@ -1302,6 +1491,7 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None, is_mtp=Fals del gathered_bytes else: self._state_dict.update(local_expert_state) + return seen_qformat, seen_block_size def _qkv_slicing( self, @@ -1338,7 +1528,11 @@ def _qkv_slicing( head_num = config.num_attention_heads head_size = config.kv_channels heads_per_group = head_num // num_query_groups - qkv_total_dim = head_num + 2 * num_query_groups + # Gated attention (Qwen3.5) packs a gate beside every query head, so a group holds + # [q, gate, k, v]; HF keeps the gate inside ``q_proj``. + output_gate = getattr(config, "attention_output_gate", False) + group_dim = (2 * heads_per_group if output_gate else heads_per_group) + 2 + qkv_total_dim = num_query_groups * group_dim weight = name_to_value.pop("weight") @@ -1360,12 +1554,24 @@ def _qkv_slicing( q_slice = torch.cat( [ - torch.arange((heads_per_group + 2) * i, (heads_per_group + 2) * i + heads_per_group) + torch.arange(group_dim * i, group_dim * i + heads_per_group) for i in range(num_query_groups_local) ] ) - k_slice = torch.arange(heads_per_group, per_rank_qkv_dim, (heads_per_group + 2)) - v_slice = torch.arange(heads_per_group + 1, per_rank_qkv_dim, (heads_per_group + 2)) + gate_slice = ( + torch.cat( + [ + torch.arange( + group_dim * i + heads_per_group, group_dim * i + 2 * heads_per_group + ) + for i in range(num_query_groups_local) + ] + ) + if output_gate + else None + ) + k_slice = torch.arange(group_dim - 2, per_rank_qkv_dim, group_dim) + v_slice = torch.arange(group_dim - 1, per_rank_qkv_dim, group_dim) ## Example of slices ## 7b: num_query_groups = head_num = 32, ## q_slice = [0, 3, 6, 9 , ... 90, 93] @@ -1378,7 +1584,16 @@ def _qkv_slicing( slices = [q_slice, k_slice, v_slice] prefixes = [q_proj_prefix, k_proj_prefix, v_proj_prefix] - proj_weights = [weight[s].reshape(-1, hidden_size) for s in slices] + def _take(tensor, index, last_dim, with_gate=False): + """Gather ``index`` heads, appending the gate heads for q under gated attention.""" + taken = tensor[index] + if with_gate: + taken = torch.cat([taken, tensor[gate_slice]], dim=1) + return taken.reshape(-1, last_dim) + + gated = [output_gate, False, False] # q carries the gate; k and v do not + + proj_weights = [_take(weight, s, hidden_size, g) for s, g in zip(slices, gated)] proj_keys = [p + "weight" for p in prefixes] if weight_scale is None: @@ -1393,10 +1608,8 @@ def _qkv_slicing( [per_rank_qkv_dim, head_size, weight_scale_hidden_size] ) proj_weight_scales = [ - weight_scale[s] - .reshape(-1, weight_scale_hidden_size) - .to(dtype=weight_scale_dtype) - for s in slices + _take(weight_scale, s, weight_scale_hidden_size, g).to(dtype=weight_scale_dtype) + for s, g in zip(slices, gated) ] else: # per-tensor scaling @@ -1432,7 +1645,7 @@ def _qkv_slicing( # Slice bias similar to weight bias = val.detach().clone() bias = bias.reshape([per_rank_qkv_dim, head_size]) - proj_biases = [bias[s].reshape(-1) for s in slices] + proj_biases = [_take(bias, s, 1, g).reshape(-1) for s, g in zip(slices, gated)] proj_bias_keys = [q_proj_prefix + key, k_proj_prefix + key, v_proj_prefix + key] for bias_tensor, bias_key in zip(proj_biases, proj_bias_keys): self._state_dict[bias_key] = bias_tensor @@ -1441,6 +1654,104 @@ def _qkv_slicing( self._state_dict[k_proj_key] = val.detach().clone() self._state_dict[v_proj_key] = val.detach().clone() + def _gated_delta_net_slicing(self, module, prefix, is_mtp=False): + """Split GatedDeltaNet's fused ``in_proj`` into HF's qkv / z / b / a projections. + + Megatron packs ``[query, key, value, z, beta, alpha]``; sizes come from the module so TP + sharding needs no re-derivation. + """ + if is_mtp: + prefix = self._mtp_prefix(prefix) + in_proj = module.in_proj + name_to_value, qformat, block_size = self._get_quantized_state( + in_proj, self.dtype, prefix=prefix + ) + + assert tuple(module.in_proj_split_names) == ( + "query", + "key", + "value", + "z", + "beta", + "alpha", + ), ( + f"Unexpected GatedDeltaNet in_proj layout {tuple(module.in_proj_split_names)}; the " + "split below assumes [query, key, value, z, beta, alpha]" + ) + sections = dict(zip(module.in_proj_split_names, module.in_proj_split_sections)) + split_sizes = [ + sections["query"] + sections["key"] + sections["value"], + sections["z"], + sections["beta"], + sections["alpha"], + ] + proj_names = ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a") + proj_prefixes = [prefix + name + "." for name in proj_names] + # The recipes keep the alpha / beta gates in BF16, but Megatron fuses all six sections + # behind one quantizer, so they can only be dropped here rather than by a quantizer_name. + keep_bf16 = { + p for p, n in zip(proj_prefixes, proj_names) if n in ("in_proj_a", "in_proj_b") + } + + for proj_prefix in proj_prefixes: + if proj_prefix in keep_bf16: + self._record_excluded_module(proj_prefix) + else: + self._record_layer_quant_config(proj_prefix, qformat, block_size) + if qformat in (None, QUANTIZATION_NONE): + # Split the fused in_proj exclude entry into the per-HF-name projections. + self.exclude_modules = [ + m for m in self.exclude_modules if m != prefix.removesuffix(".") + ] + for proj_prefix in proj_prefixes: + self._record_excluded_module(proj_prefix) + + weight = name_to_value.pop("weight") + proj_weights = list(torch.split(weight, split_sizes, dim=0)) + proj_keys = [p + "weight" for p in proj_prefixes] + weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) + + if weight_scale is None: + for key, proj_weight in zip(proj_keys, proj_weights): + self._state_dict[key] = proj_weight + else: + if len(weight_scale.shape) > 0: + # Per-channel / per-block scales are laid out along the same (output) dim. + proj_scales = list(torch.split(weight_scale, split_sizes, dim=0)) + else: + proj_scales = [weight_scale.detach().clone() for _ in proj_keys] + for proj_prefix, proj_weight, scale, key in zip( + proj_prefixes, proj_weights, proj_scales, proj_keys + ): + if proj_prefix in keep_bf16: + self._state_dict[key] = proj_weight + continue + self._state_dict[key] = to_quantized_weight( + proj_weight, scale, qformat, weight_scale_2, block_size + ) + self._state_dict[key + "_scale"] = scale + + if weight_scale_2 is not None: + if len(weight_scale_2.shape) > 0: + raise ValueError("weight_scale_2 must be a scalar!") + for proj_prefix, key in zip(proj_prefixes, proj_keys): + if proj_prefix not in keep_bf16: + self._state_dict[key + "_scale_2"] = weight_scale_2.detach().clone() + + # weight and weight_scale have been popped; the rest (bias, input_scale, ...) is + # either split like the weight or replicated onto every projection. + for key, val in name_to_value.items(): + if key == "bias": + for proj_bias, proj_prefix in zip( + torch.split(val.detach().clone(), split_sizes, dim=0), proj_prefixes + ): + self._state_dict[proj_prefix + key] = proj_bias + else: + for proj_prefix in proj_prefixes: + if proj_prefix in keep_bf16: + continue + self._state_dict[proj_prefix + key] = val.detach().clone() + def _self_attention_scaling( self, module, prefix, k_scale_name="k_scale", v_scale_name="v_scale", is_mtp=False ): @@ -1460,8 +1771,8 @@ def _self_attention_scaling( # FP8 KV Cache is supported in VLLM; NVFP4 supported in TRTLLM self.kv_cache_dtype = kv_cache_dtype - def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False): - """Pack name remapping into one tensor.""" + def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, transpose=True): + """Pack per-expert weights into one tensor; ``transpose`` for HF [E, in, out] layouts.""" if is_mtp: prefix = self._mtp_prefix(prefix) weight_list = [] @@ -1488,10 +1799,10 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False): merged_weight = torch.stack(weight_list, dim=0) - # Transpose the last two dimensions to match HuggingFace format - # Megatron format: [num_experts, out_features, in_features] - # HF format: [num_experts, in_features, out_features] - merged_weight = merged_weight.transpose(-2, -1).contiguous() + # Megatron is [num_experts, out, in]; most HF layouts want [num_experts, in, out], but + # Qwen3.5 keeps Megatron's orientation. + if transpose: + merged_weight = merged_weight.transpose(-2, -1).contiguous() if weight_scale_2_list[0] is None: merged_weight_scale_2 = None @@ -1503,8 +1814,8 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False): # NVFP4 merged_weight_scale_2 = torch.max(torch.stack(weight_scale_2_list, dim=0), dim=0)[0] merged_weight_scale = torch.stack(weight_scale_list, dim=0) - # Transpose the scaling factors to match the transposed weights - merged_weight_scale = merged_weight_scale.transpose(-2, -1).contiguous() + if transpose: + merged_weight_scale = merged_weight_scale.transpose(-2, -1).contiguous() if input_scale_list[0] is not None: merged_input_scale = torch.max(torch.stack(input_scale_list, dim=0), dim=0)[0] @@ -1684,6 +1995,20 @@ def _gather_kv_cache_dtype(self): return None +def _read_checkpoint_keys(checkpoint_dir) -> set[str]: + """Tensor names in a local HuggingFace checkpoint, from its index or single safetensors file.""" + directory = Path(checkpoint_dir) + index_file = directory / "model.safetensors.index.json" + if index_file.exists(): + with open(index_file) as f: + return set(json.load(f)["weight_map"]) + single_file = directory / "model.safetensors" + if single_file.exists(): + with safe_open(str(single_file), framework="pt", device="cpu") as f: + return set(f.keys()) + return set() + + def export_mcore_gpt_to_hf( model: torch.nn.Module, pretrained_model_name_or_path: str | os.PathLike, diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index c0aa048ebcd..23fdaa8b324 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -14,6 +14,7 @@ # limitations under the License. """Megatron-Bridge plugins for using with Model-Optimizer.""" +from functools import cache from typing import Any from megatron.bridge import AutoBridge @@ -30,12 +31,94 @@ from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.transformer.module import MegatronModule from megatron.core.utils import unwrap_model -from transformers import AutoTokenizer +from torch.distributed.checkpoint import FileSystemReader +from transformers import AutoConfig, AutoTokenizer +from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec -from modelopt.torch.utils import print_rank_0 +from modelopt.torch.utils import print_rank_0, warn_rank_0 -__all__ = ["load_mbridge_model_from_hf", "load_modelopt_megatron_checkpoint"] +__all__ = [ + "get_language_model", + "is_vlm_config", + "load_mbridge_model_from_hf", + "load_modelopt_megatron_checkpoint", + "set_moe_expert_layout", + "use_moe_grouped_gemm", +] + + +def get_language_model(model: MegatronModule) -> tuple[MegatronModule, bool]: + """Return ``(language_model, is_vlm)``; VLM wrappers nest it under ``.language_model``. + + Must agree with :func:`is_vlm_config`, which answers the same question before the Megatron + model exists, and with ``GPTModelExporter``, which repeats this because ``modelopt.torch.export`` + cannot import this module. + """ + if isinstance(model, GPTModel | HybridModel): + return model, False + language_model = getattr(model, "language_model", None) + return (model, False) if language_model is None else (language_model, True) + + +@cache +def _hf_config(hf_model_name_or_path: str, trust_remote_code: bool = False): + """Cached ``AutoConfig`` for the helpers below; a script asks the same questions repeatedly. + + Shared, so treat the result as read-only. + """ + return AutoConfig.from_pretrained(hf_model_name_or_path, trust_remote_code=trust_remote_code) + + +def is_vlm_config(hf_model_name_or_path: str, trust_remote_code: bool = False) -> bool: + """Whether a HuggingFace checkpoint describes a VLM, from its config alone.""" + return hasattr(_hf_config(hf_model_name_or_path, trust_remote_code), "vision_config") + + +def use_moe_grouped_gemm( + hf_model_name_or_path: str, + trust_remote_code: bool = False, + force_sequential: bool = False, +) -> bool: + """Pick the MoE expert layout: grouped GEMM unless that would not be HF-exportable. + + Grouped GEMM calibrates faster, but only architectures with an ``experts.linear_fc1`` export + rule can be converted to HuggingFace from it. Every script that builds the model must agree, + since the layout is baked into the Megatron checkpoint -- hence a pure function of the config. + """ + if force_sequential: + return False + config = _hf_config(hf_model_name_or_path, trust_remote_code) + text_config = getattr(config, "text_config", config) + is_moe = any( + getattr(text_config, name, None) + for name in ("num_experts", "num_local_experts", "n_routed_experts") + ) + if not is_moe: + return True # ignored for dense models + architectures = getattr(config, "architectures", None) or [""] + exportable = "experts.linear_fc1" in all_mcore_hf_export_mapping.get(architectures[0], {}) + if not exportable: + warn_rank_0( + f"{architectures[0]} has no export rule for fused (grouped GEMM) MoE experts; " + "building them as SequentialMLP so the checkpoint stays exportable." + ) + return exportable + + +def set_moe_expert_layout(provider, moe_grouped_gemm: bool) -> None: + """Apply the MoE expert layout to a provider, hybrid stack spec included. + + Set ``moe_grouped_gemm`` on the provider (the bridge's native, possibly custom/hybrid spec + reads it at build time) rather than replacing the whole layer spec -- overwriting it would + drop custom layers (e.g. Qwen3.5's GatedDeltaNet or Gemma3's custom spec). A hybrid provider + additionally needs its stack spec rebuilt, since the native one pins ``TEGroupedMLP``. + """ + if isinstance(provider, HybridModelProvider): + provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm) + provider.moe_grouped_gemm = moe_grouped_gemm + elif (provider.num_moe_experts or 0) > 0: + provider.moe_grouped_gemm = moe_grouped_gemm def load_mbridge_model_from_hf( @@ -61,7 +144,6 @@ def load_mbridge_model_from_hf( provider_overrides: Overrides for the provider. init_model_parallel: Whether to initialize model parallel. moe_grouped_gemm: Whether to use grouped GEMM for MoE. - Pruning does not support grouped GEMM yet. load_weights: Whether to load the HF weights into the model. Set to ``False`` when the weights will be loaded from a Megatron checkpoint instead (e.g. for export), in which case only the model structure (with the correct layer spec) is built. @@ -84,14 +166,7 @@ def load_mbridge_model_from_hf( assert hasattr(provider, key), f"{type(provider)} does not have attribute {key}" setattr(provider, key, value) - # Set moe_grouped_gemm on the provider (the bridge's native, possibly custom/hybrid spec reads - # it at build time) rather than replacing the whole layer spec -- overwriting it would drop - # custom layers (e.g. Qwen3.5's GatedDeltaNet + gated-attention or Gemma3's custom spec). - if isinstance(provider, HybridModelProvider): - provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm) - provider.moe_grouped_gemm = moe_grouped_gemm - elif (provider.num_moe_experts or 0) > 0: - provider.moe_grouped_gemm = moe_grouped_gemm + set_moe_expert_layout(provider, moe_grouped_gemm) provider.finalize() if init_model_parallel: provider.initialize_model_parallel(seed=0) @@ -99,9 +174,9 @@ def load_mbridge_model_from_hf( model = provider.provide_distributed_model(wrap_with_ddp=False) assert len(model) == 1 unwrapped_model = unwrap_model(model[0]) - # VLMs (e.g. Qwen3-VL) wrap the language model as ``.language_model``; the pruning target is the - # inner GPTModel/HybridModel, but we still return the full wrapper so callers can save the VLM. - language_model = getattr(unwrapped_model, "language_model", unwrapped_model) + # The optimization target is the inner GPTModel/HybridModel, but callers get the full + # wrapper back so they can save the whole VLM. + language_model, _ = get_language_model(unwrapped_model) assert isinstance(language_model, GPTModel | HybridModel), ( f"Expected a GPTModel/HybridModel (optionally wrapped as `.language_model`), " f"got {type(unwrapped_model)}" @@ -118,9 +193,14 @@ def load_mbridge_model_from_hf( return bridge, provider, model, unwrapped_model, tokenizer +def _checkpoint_keys(checkpoint_path: str) -> list[str]: + """Tensor names in a Megatron distributed checkpoint.""" + return list(FileSystemReader(checkpoint_path).read_metadata().state_dict_metadata) + + def load_modelopt_megatron_checkpoint( model: list[MegatronModule], megatron_path: str, restore_modelopt_state: bool = True -) -> None: +) -> list[MegatronModule]: """Load Megatron checkpoint weights (with modelopt_state). Args: @@ -129,10 +209,39 @@ def load_modelopt_megatron_checkpoint( restore_modelopt_state: Whether to restore the ModelOpt state (e.g. quantizers) before loading weights. Set ``False`` to load weights only -- e.g. to reload a full-precision distilled student without reconstructing the ``kd_loss`` mode (which would require a teacher model). + + Returns: + The modules loaded into: ``.language_model`` for a language-model-only VLM checkpoint, + otherwise the modules passed in. Callers that then move the ModelOpt state need this to + know where it landed. """ + # _load_model_weights_from_checkpoint does not resolve the latest iter_* directory, so resolve it explicitly + checkpoint_path = _get_modelopt_checkpoint_path(megatron_path) + + # ``distill.py`` distills a VLM's language model only, so its checkpoint holds the language + # model at the root and must be loaded into ``.language_model``, not the full VLM wrapper. + # Keyed off the same ``language_model.`` prefix ``get_language_model`` navigates, so a renamed + # vision tower cannot make a full-VLM checkpoint look language-model-only. + checkpoint_keys = _checkpoint_keys(checkpoint_path) + unwrapped_model = unwrap_model(model) + if any(get_language_model(m)[1] for m in unwrapped_model) and not any( + key.startswith("language_model.") for key in checkpoint_keys + ): + print_rank_0("Language-model-only checkpoint: loading into the VLM's `.language_model`.") + model = [get_language_model(m)[0] for m in unwrapped_model] + # Restore the ModelOpt state before loading weights. # has_modelopt_state / load_modelopt_state resolves the latest iter_* directory - if restore_modelopt_state and has_modelopt_state(megatron_path): - load_modelopt_state(model, megatron_path) - # _load_model_weights_from_checkpoint does not resolve the latest iter_* directory, so resolve it explicitly - _load_model_weights_from_checkpoint(_get_modelopt_checkpoint_path(megatron_path), model) + if restore_modelopt_state: + if has_modelopt_state(megatron_path): + load_modelopt_state(model, megatron_path) + elif any("_quantizer." in key for key in checkpoint_keys): + # The quantizers cannot be rebuilt without the state, and the loader ignores the + # leftover amax tensors, so the model would silently load unquantized. + raise RuntimeError( + f"{megatron_path} holds quantizer tensors but no restorable ModelOpt state. " + "The state was dropped when the checkpoint was written -- re-run the step that " + "produced it." + ) + _load_model_weights_from_checkpoint(checkpoint_path, model) + return model diff --git a/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml index fd27046d608..fa34eda4a44 100644 --- a/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml +++ b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml @@ -23,6 +23,8 @@ - "*linear_attn.in_proj_a*" - "*linear_attn.in_proj_b*" - "*mixer.conv1d*" + # Megatron-Core names the GatedDeltaNet block ``self_attention``, HF ``linear_attn``. + - "*self_attention.conv1d*" - "*mlp.gate.*" - "*mlp.shared_expert_gate.*" - "*output_layer*" diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index 3aadadd289c..e819c82ab4e 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -28,6 +28,9 @@ enable: false - quantizer_name: '*mixer.conv1d*' enable: false + # Megatron-Core names the GatedDeltaNet block ``self_attention``; HF calls it ``linear_attn``. + - quantizer_name: '*self_attention.conv1d*' + enable: false - quantizer_name: '*mlp.gate.*' enable: false - quantizer_name: '*mlp.shared_expert_gate.*' diff --git a/tests/_test_utils/torch/export/unified_checkpoint.py b/tests/_test_utils/torch/export/unified_checkpoint.py new file mode 100644 index 00000000000..e91e400f1c1 --- /dev/null +++ b/tests/_test_utils/torch/export/unified_checkpoint.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. +"""Content checks for an exported unified HuggingFace checkpoint. + +Existence checks pass even when the exporter silently drops a whole module family, so compare +the exported tensors against the HuggingFace checkpoint the model came from. +""" + +import json +from pathlib import Path + +import torch +from safetensors.torch import load_file + +from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor + +__all__ = [ + "assert_exported_checkpoint_matches", + "assert_safetensors_index_consistent", + "load_safetensors_dir", +] + +# Per-tensor / per-channel scales the exporter adds for quantized weights. These have no +# counterpart in the (unquantized) reference checkpoint. +QUANT_SUFFIXES = ( + "weight_scale", + "weight_scale_2", + "input_scale", + "output_scale", + "k_scale", + "v_scale", +) + + +def load_safetensors_dir(path: Path | str) -> dict[str, torch.Tensor]: + """Load every safetensors shard under ``path`` into one dict.""" + state_dict: dict[str, torch.Tensor] = {} + for shard in sorted(Path(path).glob("*.safetensors")): + state_dict.update(load_file(str(shard))) + assert state_dict, f"No safetensors tensors found in {path}" + return state_dict + + +def assert_safetensors_index_consistent(export_dir: Path | str) -> None: + """Assert ``model.safetensors.index.json`` matches the shards actually written.""" + export_dir = Path(export_dir) + index_file = export_dir / "model.safetensors.index.json" + if not index_file.exists(): # single unsharded file: nothing to cross-check + return + weight_map = json.loads(index_file.read_text())["weight_map"] + missing_files = {f for f in set(weight_map.values()) if not (export_dir / f).exists()} + assert not missing_files, f"index.json references missing shards: {sorted(missing_files)}" + exported = set(load_safetensors_dir(export_dir)) + assert set(weight_map) == exported, ( + f"index.json disagrees with the shards: {sorted(set(weight_map) - exported)[:10]} only in " + f"index, {sorted(exported - set(weight_map))[:10]} only in shards" + ) + + +def _is_packed(tensor: torch.Tensor) -> bool: + """Whether a tensor holds sub-byte weights packed two-per-``uint8`` (NVFP4 / INT4).""" + return tensor.dtype == torch.uint8 + + +def _expected_shape(exported: torch.Tensor, reference: torch.Tensor) -> tuple[int, ...]: + """Reference shape, with the last dim halved when the export packs two 4-bit values per byte.""" + shape = tuple(reference.shape) + if _is_packed(exported) and shape: + return (*shape[:-1], shape[-1] // 2) + return shape + + +def _unpack_nvfp4(packed: torch.Tensor) -> torch.Tensor: + """Expand two E2M1 values per ``uint8`` back into a float tensor (low nibble first).""" + codes = torch.empty( + (*packed.shape[:-1], packed.shape[-1] * 2), dtype=torch.long, device=packed.device + ) + codes[..., 0::2] = packed & 0x0F + codes[..., 1::2] = packed >> 4 + return NVFP4QTensor.get_e2m1_values(packed.device).to(torch.float32)[codes] + + +def _scale_key(key: str, suffix: str) -> str: + """Scale key for the dotted (``...proj.weight``) and packed (``...gate_up_proj``) layouts.""" + return key.replace(".weight", f".{suffix}") if key.endswith(".weight") else f"{key}_{suffix}" + + +def _dequantize(key: str, exported: dict[str, torch.Tensor]) -> torch.Tensor: + """Undo the exporter's weight scaling so the value can be compared to the source.""" + scale = exported.get(_scale_key(key, "weight_scale")) + if _is_packed(exported[key]): + # NVFP4: E2M1 codes, an E4M3 scale per block along the last dim, and a global FP32 scale. + weight = _unpack_nvfp4(exported[key]) + block_size = weight.shape[-1] // scale.shape[-1] + weight = weight * scale.to(torch.float32).repeat_interleave(block_size, dim=-1) + scale_2 = exported.get(_scale_key(key, "weight_scale_2")) + return weight if scale_2 is None else weight * scale_2.to(torch.float32) + weight = exported[key].to(torch.float32) + if scale is None: + return weight + scale = scale.to(torch.float32) + return weight * scale.reshape(*scale.shape, *([1] * (weight.ndim - scale.ndim))) + + +def assert_exported_checkpoint_matches( + export_dir: Path | str, + ref_hf_dir: Path | str, + *, + allow_missing: tuple[str, ...] = (), + allow_unexpected: tuple[str, ...] = (), + check_values: bool = True, + bit_exact_prefixes: tuple[str, ...] = (), + rtol: float = 0.15, + packed_rtol: float = 0.3, +) -> None: + """Assert an exported unified HF checkpoint reproduces the reference model it came from. + + Args: + export_dir: Exported unified HF checkpoint. + ref_hf_dir: HuggingFace checkpoint the Megatron model was built from. + allow_missing: Substrings of reference keys the export may omit. + allow_unexpected: Substrings of extra exported keys, beyond quantization scales. + check_values: Compare values too; ``False`` when the Megatron weights are random. + bit_exact_prefixes: Keys that must match exactly even when ``check_values`` is off. + rtol: Max relative error for quantized tensors. + packed_rtol: Same for sub-byte formats, whose grid is coarser. + """ + exported = load_safetensors_dir(export_dir) + reference = load_safetensors_dir(ref_hf_dir) + assert_safetensors_index_consistent(export_dir) + + missing = {k for k in set(reference) - set(exported) if not any(a in k for a in allow_missing)} + assert not missing, ( + f"{len(missing)} reference tensor(s) absent from the export, e.g. {sorted(missing)[:8]}" + ) + + # Anything extra must be a quantization scale; a stray weight means a mis-named rule. + unexpected = { + k + for k in set(exported) - set(reference) + if not k.endswith(QUANT_SUFFIXES) and not any(a in k for a in allow_unexpected) + } + assert not unexpected, f"Export produced unexpected tensors: {sorted(unexpected)[:8]}" + + shared = sorted(set(exported) & set(reference)) + mismatched = [ + (k, tuple(exported[k].shape), tuple(reference[k].shape)) + for k in shared + if tuple(exported[k].shape) != _expected_shape(exported[k], reference[k]) + ] + assert not mismatched, f"Shape mismatches (key, exported, reference): {mismatched[:8]}" + + if bit_exact_prefixes: + drifted = [ + k + for k in shared + if k.startswith(bit_exact_prefixes) and not torch.equal(exported[k], reference[k]) + ] + assert not drifted, ( + f"{len(drifted)} copied-through tensor(s) differ from the reference: {drifted[:8]}" + ) + checked = [k for k in shared if k.startswith(bit_exact_prefixes)] + assert checked, f"bit_exact_prefixes {bit_exact_prefixes} matched no exported tensor" + + if not check_values: + return + + wrong = [] + for key in shared: + got, want = _dequantize(key, exported), reference[key].to(torch.float32) + if ( + exported[key].dtype == reference[key].dtype + and _scale_key(key, "weight_scale") not in exported + ): + # Copied through untouched (norms, router, vision tower): must be bit-exact. + if not torch.equal(exported[key], reference[key]): + wrong.append((key, "not bit-exact")) + else: + denom = want.abs().max().clamp_min(torch.finfo(torch.float32).tiny) + rel = ((got - want).abs().max() / denom).item() + tol = packed_rtol if _is_packed(exported[key]) else rtol + if rel > tol: + wrong.append((key, f"max_rel_err={rel:.4f} > {tol}")) + assert not wrong, f"{len(wrong)} tensor(s) differ from the reference, e.g. {wrong[:8]}" diff --git a/tests/_test_utils/torch/megatron/modelopt_state.py b/tests/_test_utils/torch/megatron/modelopt_state.py new file mode 100644 index 00000000000..2c5f22bd5d5 --- /dev/null +++ b/tests/_test_utils/torch/megatron/modelopt_state.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. +"""ModelOpt-state checks for Megatron distributed checkpoints.""" + +from pathlib import Path + +from megatron.bridge.training.post_training.checkpointing import has_modelopt_state +from torch.distributed.checkpoint import FileSystemReader + +__all__ = ["assert_has_modelopt_state", "assert_no_quantizers_matching"] + + +def assert_has_modelopt_state(megatron_path: Path | str) -> None: + """Assert a Megatron checkpoint carries restorable ModelOpt state. + + ``rglob("modelopt_state")`` passes on an empty state, which exports unquantized. + """ + + state_dirs = list(Path(megatron_path).rglob("modelopt_state")) + assert state_dirs, f"No modelopt_state directory under {megatron_path}" + assert has_modelopt_state(str(megatron_path)), ( + f"modelopt_state under {megatron_path} holds no restorable mode (only 'kd_loss' or " + "empty), so the quantizers would not survive a reload" + ) + + +def assert_no_quantizers_matching(megatron_path: Path | str, *substrings: str) -> None: + """Assert no calibrated quantizer under ``megatron_path`` matches ``substrings``. + + Disabled-quantizer patterns are written against HuggingFace names, so they silently + no-op wherever Megatron names the module differently. + """ + iter_dirs = sorted(Path(megatron_path).glob("iter_*")) + assert iter_dirs, f"No iter_* checkpoint under {megatron_path}" + keys = FileSystemReader(str(iter_dirs[-1])).read_metadata().state_dict_metadata + quantizers = [k for k in keys if "_quantizer." in k] + assert quantizers, f"No quantizers at all under {megatron_path}; was the model quantized?" + for substring in substrings: + hits = sorted(k for k in quantizers if substring in k) + assert not hits, f"Expected no quantizer matching {substring!r}, found: {hits[:4]}" diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index cf75e50e107..712cece2055 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -14,12 +14,15 @@ # limitations under the License. import contextlib +import re +from collections import defaultdict from functools import partial from pathlib import Path import pytest import torch from _test_utils.torch.misc import set_seed +from safetensors.torch import load_file, save_file transformers = pytest.importorskip("transformers") from transformers import ( @@ -390,6 +393,45 @@ def _get_tiny_qwen3_5_vl(moe: bool = False, **config_kwargs) -> PreTrainedModel: return AutoModelForImageTextToText.from_config(cfg) +def _pack_qwen3_5_moe_experts(dir_path: Path | str) -> None: + """Repack a saved Qwen3.5-MoE fixture's experts the way released checkpoints store them. + + transformers unpacks routed experts into per-expert ``experts.{i}.*`` on save, but every real + Qwen3.5 checkpoint stores them packed as ``experts.gate_up_proj`` / ``experts.down_proj``. + Without this the fixture would push the exporter toward a layout no real checkpoint uses. + """ + dir_path = Path(dir_path) + shards = sorted(dir_path.glob("*.safetensors")) + state_dict: dict[str, torch.Tensor] = {} + for shard in shards: + state_dict.update(load_file(str(shard))) + + per_expert = re.compile(r"^(.*\.mlp\.experts)\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$") + grouped: dict[str, dict[int, dict[str, torch.Tensor]]] = defaultdict(lambda: defaultdict(dict)) + for key in list(state_dict): + match = per_expert.match(key) + if match: + base, expert_id, proj = match.groups() + grouped[base][int(expert_id)][proj] = state_dict.pop(key) + if not grouped: + return + + for base, experts in grouped.items(): + ids = sorted(experts) + # gate first, then up, along the output dim -- the order the packed tensor is split on. + state_dict[f"{base}.gate_up_proj"] = torch.stack( + [torch.cat([experts[i]["gate_proj"], experts[i]["up_proj"]], dim=0) for i in ids] + ) + state_dict[f"{base}.down_proj"] = torch.stack([experts[i]["down_proj"] for i in ids]) + + for shard in shards: + shard.unlink() + save_file(state_dict, str(dir_path / "model.safetensors"), metadata={"format": "pt"}) + index = dir_path / "model.safetensors.index.json" + if index.exists(): + index.unlink() + + def _create_tiny_qwen3_5_vl_dir( tmp_path: Path | str, with_processor: bool = False, @@ -398,7 +440,7 @@ def _create_tiny_qwen3_5_vl_dir( moe: bool = False, **config_kwargs, ) -> Path | tuple[Path, PreTrainedModel]: - return _create_tiny_vlm_dir( + result = _create_tiny_vlm_dir( Path(tmp_path) / ("tiny_qwen3_5_moe_vl" if moe else "tiny_qwen3_5_vl"), QWEN3_5_VL_REF, _get_tiny_qwen3_5_vl, @@ -407,6 +449,9 @@ def _create_tiny_qwen3_5_vl_dir( moe=moe, **config_kwargs, ) + if moe: + _pack_qwen3_5_moe_experts(result[0] if isinstance(result, tuple) else result) + return result get_tiny_qwen3_5_vl = partial(_get_tiny_qwen3_5_vl, moe=False) @@ -447,6 +492,31 @@ def create_tiny_nemotron_dir( ##### NEMOTRON-H (Mamba + Attention + MoE/MLP hybrid) ##### +def _match_released_nemotron_h_embedding_name() -> None: + """Save NemotronH's embedding as ``backbone.embeddings`` like released checkpoints do. + + transformers <= 5.15 renames it to the legacy singular on save, so a saved fixture disagrees + with every real checkpoint on one tensor. Dropped upstream in 5.16, where this is a no-op + since there is no such rule left to remove. + """ + # Optional dependency: conversion_mapping is transformers 5.x only, while NemotronH also + # exists on the 4.57 floor of the support matrix. + try: + from transformers.conversion_mapping import ( + get_checkpoint_conversion_mapping, + register_checkpoint_conversion_mapping, + ) + except ImportError: + return + + mapping = get_checkpoint_conversion_mapping("nemotron_h") + if mapping is None: + return + kept = [m for m in mapping if getattr(m, "source_patterns", None) != ["embedding.weight"]] + if len(kept) != len(mapping): + register_checkpoint_conversion_mapping("nemotron_h", kept, overwrite=True) + + def get_tiny_nemotron_h(**config_kwargs) -> PreTrainedModel: set_seed(SEED) @@ -454,6 +524,8 @@ def get_tiny_nemotron_h(**config_kwargs) -> PreTrainedModel: # module is imported broadly (including by the min-transformers CI job). from transformers import NemotronHConfig + _match_released_nemotron_h_embedding_name() + # Tiny NemotronH hybrid. hybrid_override_pattern letters: M=Mamba, E=MoE/FFN, *=Attention. # "ME*E" matches the NemotronH default and exercises Mamba + MoE + attention layers. kwargs = { diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py index 6128b5ad6ae..11907936210 100644 --- a/tests/examples/megatron_bridge/test_qad.py +++ b/tests/examples/megatron_bridge/test_qad.py @@ -18,8 +18,12 @@ import pytest 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.megatron.modelopt_state import ( + assert_has_modelopt_state, + assert_no_quantizers_matching, +) from _test_utils.torch.transformers_models import ( - create_tiny_gemma3vl_dir, create_tiny_qwen3_5_moe_vl_dir, create_tiny_qwen3_dir, ) @@ -27,41 +31,29 @@ @pytest.mark.timeout(720) # Multiple steps in one test hence takes longer than the default timeout @pytest.mark.parametrize( - ("create_student", "is_vlm", "is_moe"), + "create_student", [ - (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), False, False), - # Dense-VLM QAD path; the MoE VLM below covers it in CI, so run this one on demand only. + lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), pytest.param( - lambda tmp_path: create_tiny_gemma3vl_dir( + lambda tmp_path: create_tiny_qwen3_5_moe_vl_dir( tmp_path, with_processor=True, + # Cover both Qwen3.5 decoder kinds at the same layer count num_hidden_layers=2, - intermediate_size=128, - max_position_embeddings=512, - ), - True, - False, - marks=pytest.mark.manual, - ), - pytest.param( - lambda tmp_path: create_tiny_qwen3_5_moe_vl_dir( - tmp_path, with_processor=True, num_hidden_layers=2 + layer_types=["linear_attention", "full_attention"], ), - True, - True, ), ], - ids=["qwen3", "gemma3vl", "qwen3_5_moe_vl"], + ids=["qwen3", "qwen3_5_moe_vl"], ) -def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): +def test_qad(tmp_path: Path, num_gpus, create_student): """Quantize a tiny model, run QAD from the quantized student, and export the result. - For VLMs only the language model is quantized and distilled (vision tower / projector untouched), - and a text calibration dataset infers text-only LM calibration. VLM quantized-HF export is - unsupported, so the VLM case stops at the distilled Megatron checkpoint and verifies the ModelOpt - (quantize) state survived distillation; the LLM case additionally exports a unified HF checkpoint. + Covers what only QAD exercises: that the ModelOpt state survives distillation. Per-architecture + export is covered more cheaply by test_quantize_export.py, so keep this to one LLM and one VLM. """ hf_model_path = create_student(tmp_path) + is_vlm = "vision_config" in (hf_model_path / "config.json").read_text() quantized_megatron_path = tmp_path / "quantized_megatron" distill_output_dir = tmp_path / "qad_output" train_iters = 3 @@ -81,9 +73,9 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): export_megatron_path=quantized_megatron_path, ) run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True) - assert list(quantized_megatron_path.rglob("modelopt_state")), ( - "Expected modelopt_state in the quantized Megatron checkpoint" - ) + assert_has_modelopt_state(quantized_megatron_path) + # Megatron names these differently from HF, so the recipe's patterns must have aliases. + assert_no_quantizers_matching(quantized_megatron_path, "conv1d", "mlp.router", "output_layer") # Step 2: QAD -- load the quantized student from the Megatron checkpoint (restoring the ModelOpt # quantizers) and distill from the (unquantized) HF teacher. The distilled checkpoint must keep the @@ -113,18 +105,17 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): tracker = distilled_megatron_path / "latest_checkpointed_iteration.txt" assert tracker.read_text(encoding="utf-8").strip() == str(early_exit_iter) assert (distilled_megatron_path / "iter_0000001").is_dir() - assert list(distilled_megatron_path.rglob("modelopt_state")), ( - "Expected modelopt_state to be preserved in the distilled (QAD) checkpoint" - ) - - if is_vlm: - return # VLM quantized-HF export is unsupported; stop at the distilled Megatron checkpoint + assert_has_modelopt_state(distilled_megatron_path) # Step 3: export the distilled quantized checkpoint to a unified HF checkpoint. hf_quant_config.json # is only written for a quantized model, so its presence confirms the quantizers survived QAD. hf_export_path = tmp_path / "qad_fp8_hf" export_cmd = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "export_quantized_megatron_to_hf.py"], + [ + "torchrun", + f"--nproc_per_node={num_gpus}", + "export_quantized_megatron_to_hf.py", + ], hf_model_name_or_path=hf_model_path, megatron_path=distilled_megatron_path, export_unified_hf_path=hf_export_path, @@ -133,4 +124,11 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe): run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True) assert (hf_export_path / "config.json").exists() assert (hf_export_path / "hf_quant_config.json").exists() - assert list(hf_export_path.glob("*.safetensors")), "Expected exported safetensors weights" + # QAD trains the student, so language-model weights drift from the reference; the vision + # tower is never trained and must still come through byte for byte. + assert_exported_checkpoint_matches( + hf_export_path, + hf_model_path, + check_values=False, + bit_exact_prefixes=("model.visual.",) if is_vlm else (), + ) diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index 55a97b49d00..61f258c3759 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -16,26 +16,49 @@ from pathlib import Path +import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command -from _test_utils.torch.transformers_models import create_tiny_qwen3_dir +from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches +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_dir, + create_tiny_qwen3_moe_dir, + create_tiny_qwen3vl_dir, +) +# 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. +# Use a vLLM-friendly head_dim (64): the default tiny config (head_dim=2) is unsupported. +_DENSE_KWARGS = { + "hidden_size": 128, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "num_hidden_layers": 2, + "intermediate_size": 256, + "max_position_embeddings": 512, +} -# NOTE: Qwen3.5-VL covered by test_qad.py -def test_quantize_and_export(tmp_path: Path, num_gpus): - """Quantize a tiny Qwen3 via a YAML recipe and export it to a unified HF checkpoint.""" - # Use a vLLM-friendly head_dim (64) since the default tiny config (head_dim=2) is unsupported. - hf_model_path = create_tiny_qwen3_dir( - tmp_path, - with_tokenizer=True, - hidden_size=128, - num_attention_heads=2, - num_key_value_heads=2, - num_hidden_layers=2, - intermediate_size=256, - max_position_embeddings=512, - ) - megatron_path = tmp_path / "qwen3_fp8_megatron" - hf_export_path = tmp_path / "qwen3_fp8_hf" + +@pytest.mark.parametrize( + ("create_model", "model_kwargs"), + [ + (create_tiny_qwen3_dir, _DENSE_KWARGS), + # MoE: routed experts used to be dropped silently from the export. + (create_tiny_qwen3_moe_dir, _DENSE_KWARGS), + # Dense VLM: only the language model is quantized, vision is copied through. + (create_tiny_qwen3vl_dir, {}), + # Mamba hybrid + MoE, and the one architecture that keeps grouped-GEMM experts. + (create_tiny_nemotron_h_dir, {}), + ], + ids=["qwen3", "qwen3_moe", "qwen3vl", "nemotron_h"], +) +def test_quantize_and_export(tmp_path: Path, num_gpus, create_model, model_kwargs): + """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint.""" + hf_model_path = create_model(tmp_path, with_tokenizer=True, **model_kwargs) + megatron_path = tmp_path / "fp8_megatron" + hf_export_path = tmp_path / "fp8_hf" # Step 1: quantize and save a Megatron checkpoint quantize_cmd = extend_cmd_parts( @@ -51,9 +74,7 @@ def test_quantize_and_export(tmp_path: Path, num_gpus): ) run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True) assert (megatron_path / "latest_checkpointed_iteration.txt").exists() - assert list(megatron_path.rglob("modelopt_state")), ( - "Expected modelopt_state in the Megatron checkpoint" - ) + assert_has_modelopt_state(megatron_path) # Step 2: export to HF export_cmd = extend_cmd_parts( @@ -65,7 +86,7 @@ def test_quantize_and_export(tmp_path: Path, num_gpus): run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True) assert (hf_export_path / "config.json").exists() assert (hf_export_path / "hf_quant_config.json").exists() - assert list(hf_export_path.glob("*.safetensors")), "Expected exported safetensors weights" + assert_exported_checkpoint_matches(hf_export_path, hf_model_path) # The exported unified checkpoint should be loadable and runnable by vLLM. The deployment check below # is disabled because it takes too long in CI (likely because of first run) diff --git a/tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py b/tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py new file mode 100644 index 00000000000..197f5923dac --- /dev/null +++ b/tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. +"""The MoE expert layout must be chosen identically by every script that builds the model.""" + +import pytest + +from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping + + +@pytest.mark.parametrize( + ("arch", "grouped_is_exportable"), + [ + # These map fused grouped-GEMM experts, so they keep the faster layout. + ("NemotronHForCausalLM", True), + ("Qwen3_5MoeForConditionalGeneration", True), + ("Qwen3MoeForCausalLM", False), + ("DeepseekV3ForCausalLM", False), + ("GptOssForCausalLM", False), + ("Llama4ForConditionalGeneration", False), + ], +) +def test_grouped_gemm_exportability(arch, grouped_is_exportable): + assert arch in all_mcore_hf_export_mapping, f"{arch} is no longer an exported architecture" + has_rule = "experts.linear_fc1" in all_mcore_hf_export_mapping[arch] + assert has_rule is grouped_is_exportable, ( + f"{arch} grouped-GEMM exportability changed; use_moe_grouped_gemm would now pick the " + "other expert layout, which silently breaks checkpoints written by the previous default" + ) diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 30d592e94ba..6b5620206ff 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -21,11 +21,15 @@ import pytest import torch import transformers -from _test_utils.torch.megatron.models import get_mcore_gpt_model +from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches +from _test_utils.torch.megatron.models import get_mcore_gpt_model, get_mcore_hybrid_model from _test_utils.torch.megatron.utils import get_forward from _test_utils.torch.transformers_models import ( create_tiny_llama_dir, create_tiny_nemotron_dir, + create_tiny_nemotron_h_dir, + create_tiny_qwen3_5_moe_vl_dir, + create_tiny_qwen3_moe_dir, create_tiny_qwen3vl_dir, ) from safetensors import safe_open @@ -87,7 +91,34 @@ def _test_unified_export_megatron( size, model_dir=None, ): - if model_type == "qwen3vl": + if model_type == "nemotron_h": + config = transformers.AutoConfig.from_pretrained(model_dir) + model = get_mcore_hybrid_model( + tensor_model_parallel_size=size, + pipeline_model_parallel_size=1, + initialize_megatron=True, + num_layers=config.num_hidden_layers, + hybrid_layer_pattern=config.hybrid_override_pattern, + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_query_groups=config.num_key_value_heads, + ffn_hidden_size=config.intermediate_size, + max_sequence_length=config.max_position_embeddings, + vocab_size=config.vocab_size, + mamba_state_dim=config.ssm_state_size, + mamba_num_heads=config.mamba_num_heads, + mamba_head_dim=config.mamba_head_dim, + mamba_num_groups=config.n_groups, + num_moe_experts=config.n_routed_experts, + moe_ffn_hidden_size=config.moe_intermediate_size, + moe_shared_expert_intermediate_size=config.moe_shared_expert_intermediate_size, + # NemotronH is the only arch that exports fused grouped-GEMM experts. + moe_grouped_gemm=True, + # NemotronH norms are RMSNorm; the builder defaults to LayerNorm, whose biases have no + # counterpart in the HF checkpoint. + normalization="RMSNorm", + ).cuda() + elif model_type == "qwen3vl": config = transformers.AutoConfig.from_pretrained(model_dir) text_cfg = config.text_config num_layers = text_cfg.num_hidden_layers @@ -98,6 +129,47 @@ def _test_unified_export_megatron( max_sequence_length = text_cfg.max_position_embeddings vocab_size = text_cfg.vocab_size extra_kwargs = {"kv_channels": text_cfg.head_dim, "qk_layernorm": True} + elif model_type in {"qwen3_5_moe_vl_grouped", "qwen3_5_moe_vl_sequential"}: + text_cfg = transformers.AutoConfig.from_pretrained(model_dir).text_config + num_layers = text_cfg.num_hidden_layers + hidden_size = text_cfg.hidden_size + num_attention_heads = text_cfg.num_attention_heads + num_query_groups = text_cfg.num_key_value_heads + ffn_hidden_size = text_cfg.intermediate_size + max_sequence_length = text_cfg.max_position_embeddings + vocab_size = text_cfg.vocab_size + # Hybrid GatedDeltaNet + gated attention, with routed experts stored packed. + extra_kwargs = { + "kv_channels": text_cfg.head_dim, + "qk_layernorm": True, + "experimental_attention_variant": "gated_delta_net", + "num_moe_experts": text_cfg.num_experts, + "moe_ffn_hidden_size": text_cfg.moe_intermediate_size, + "moe_shared_expert_intermediate_size": text_cfg.shared_expert_intermediate_size, + "moe_shared_expert_gate": True, + # Match the HF layer_types pattern (every Nth layer is full attention, rest GDN). + "linear_attention_freq": len(text_cfg.layer_types), + # Both layouts must reach the same packed HF tensors, via GroupedMLPPacking + # (TEGroupedMLP) and PackNameRemapping (SequentialMLP) respectively. + "moe_grouped_gemm": model_type.endswith("grouped"), + } + elif model_type == "qwen3_moe": + config = transformers.AutoConfig.from_pretrained(model_dir) + num_layers = config.num_hidden_layers + hidden_size = config.hidden_size + num_attention_heads = config.num_attention_heads + num_query_groups = config.num_key_value_heads + ffn_hidden_size = config.intermediate_size + max_sequence_length = config.max_position_embeddings + vocab_size = config.vocab_size + # SequentialMLP: the Qwen3 rules only cover per-expert (``local_experts``) MoE. + extra_kwargs = { + "kv_channels": config.hidden_size // config.num_attention_heads, + "qk_layernorm": True, + "num_moe_experts": config.num_experts, + "moe_ffn_hidden_size": config.moe_intermediate_size, + "moe_grouped_gemm": False, + } elif model_type in {"llama", "nemotron"}: config = transformers.AutoConfig.from_pretrained(model_dir) num_layers = config.num_hidden_layers @@ -114,22 +186,26 @@ def _test_unified_export_megatron( activation_func = "squared_relu" if model_type == "nemotron" else "swiglu" normalization = "LayerNorm" if model_type == "nemotron" else "RMSNorm" - model = get_mcore_gpt_model( - tensor_model_parallel_size=size, - pipeline_model_parallel_size=1, - initialize_megatron=True, - num_layers=num_layers, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - num_query_groups=num_query_groups, - ffn_hidden_size=ffn_hidden_size, - max_sequence_length=max_sequence_length, - vocab_size=vocab_size, - activation_func=activation_func, - normalization=normalization, - transformer_impl="modelopt", - **extra_kwargs, - ).cuda() + model = ( + model + if model_type == "nemotron_h" + else get_mcore_gpt_model( + tensor_model_parallel_size=size, + pipeline_model_parallel_size=1, + initialize_megatron=True, + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_query_groups=num_query_groups, + ffn_hidden_size=ffn_hidden_size, + max_sequence_length=max_sequence_length, + vocab_size=vocab_size, + activation_func=activation_func, + normalization=normalization, + transformer_impl="modelopt", + **extra_kwargs, + ).cuda() + ) if quant_config: quant_config_dict = getattr(mtq, quant_config) @@ -166,20 +242,17 @@ def _test_unified_export_megatron( if quant_config: _verify_model_quant_config(tmp_export_dir, quant_config, kv_cache_quant_cfg) - if model_type == "qwen3vl" and rank == 0: - # sanity check that vision weights were merged by export_mcore_gpt_to_hf - keys = [] - for sf in sorted(tmp_export_dir.glob("*.safetensors")): - with safe_open(str(sf), framework="pt", device="cpu") as f: - keys.extend(f.keys()) - # every decoder layer should be present, not just some - for i in range(num_layers): - assert any(k.startswith(f"model.language_model.layers.{i}.") for k in keys), ( - f"language model layer {i} keys missing from export" - ) - assert any(k.startswith("model.visual.") for k in keys), ( - "vision encoder keys missing from export" + if rank == 0 and extra_module is None: + # Names / shapes only: these Megatron weights are random, not loaded from model_dir. + assert_exported_checkpoint_matches( + tmp_export_dir, + model_dir, + check_values=False, + # get_mcore_gpt_model always enables the MoE router bias; tiny HF configs have none. + allow_unexpected=("mlp.gate.expert_bias",), ) + + if model_type == "qwen3vl" and rank == 0: # try to load the model and run a forward pass vl_model = Qwen3VLForConditionalGeneration.from_pretrained( tmp_export_dir, torch_dtype=torch.bfloat16 @@ -194,8 +267,10 @@ def _test_unified_export_megatron( ("model_type", "extra_module", "quant_config", "kv_cache_quant_cfg"), [ ("nemotron", None, None, None), - ("nemotron", None, "NVFP4_DEFAULT_CFG", None), - ("nemotron", None, "NVFP4_DEFAULT_CFG", "FP8_KV_CFG"), + # NemotronH (Mamba + attention + grouped-GEMM MoE) supersedes the older Nemotron for + # quantized coverage; the plain/eagle/medusa cases below still smoke-test the old arch. + ("nemotron_h", None, "NVFP4_DEFAULT_CFG", None), + ("nemotron_h", None, "NVFP4_DEFAULT_CFG", "FP8_KV_CFG"), ("nemotron", "eagle", None, None), ("nemotron", "medusa", None, None), ("llama", None, None, None), @@ -205,6 +280,14 @@ def _test_unified_export_megatron( ("llama", "medusa", None, None), ("qwen3vl", None, None, None), ("qwen3vl", None, "FP8_DEFAULT_CFG", None), + # Regression guard: routed experts used to be dropped silently from the export. + ("qwen3_moe", None, None, None), + ("qwen3_moe", None, "FP8_DEFAULT_CFG", None), + # Packed routed experts (Qwen3.5). NVFP4 keeps per-expert block scales while FP8 merges a + # single scale, so the packing rules only get full coverage across both formats. + ("qwen3_5_moe_vl_grouped", None, "NVFP4_DEFAULT_CFG", None), + ("qwen3_5_moe_vl_grouped", None, "FP8_DEFAULT_CFG", None), + ("qwen3_5_moe_vl_sequential", None, "NVFP4_DEFAULT_CFG", None), ], ) def test_unified_export_megatron( @@ -216,6 +299,12 @@ def test_unified_export_megatron( model_dir = create_tiny_qwen3vl_dir(tmp_path) elif model_type == "nemotron": model_dir = create_tiny_nemotron_dir(tmp_path) + elif model_type == "nemotron_h": + model_dir = create_tiny_nemotron_h_dir(tmp_path) + elif model_type == "qwen3_moe": + model_dir = create_tiny_qwen3_moe_dir(tmp_path) + elif model_type.startswith("qwen3_5_moe_vl"): + model_dir = create_tiny_qwen3_5_moe_vl_dir(tmp_path) else: raise ValueError(f"Unsupported model_type: {model_type}") # TODO: Fix TP>1 failures