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