diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 329e5d21f05..ae2a415217c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -85,6 +85,8 @@ Changelog - Fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) no longer requires an ``act_fn`` attribute. Some fused-expert modules (e.g. ``MiniMaxM3VLExperts``) apply a custom gated activation between the two ``F.linear`` calls instead of exposing ``act_fn``; they were silently skipped, leaving routed experts unquantized (an experts-only recipe matched nothing) and failing HF export with ``NotImplementedError``. ``_QuantFusedExperts`` is activation-agnostic (it only intercepts the two ``F.linear`` calls), so the requirement was unnecessary. This enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. - Fix unified HF export emitting transformers' *in-memory* (post-``conversion_mapping``) tensor names instead of the original model-hub names, breaking the unified-checkpoint contract (observed on MiniMax-M3: exported ``model.language_model.*`` / ``mlp.experts.*.gate_proj`` instead of hub ``language_model.model.*`` / ``block_sparse_moe.experts.*.w{1,2,3}``). transformers' own save-side ``revert_weight_conversion`` is disabled by ModelOpt because it raises ``RuntimeError`` on 0-d scalar scale tensors, so a new quant-aware reverse conversion (``modelopt/torch/export/quant_aware_conversion.py``) derives rename/split rules from the model's conversion mapping via transformers' ``reverse_transform()`` and carries each weight's companion scale tensors (``weight_scale``, ``weight_scale_2``, ``input_scale``, ``weight_scale_inv``, ``bias``) through the renames and un-fusions, so quantized exports round-trip to the hub names. Any mapping op that cannot be reversed quant-aware yet (e.g. still-stacked fused experts) falls back to the previous in-memory names instead of aborting the export. + Nested submodel reverse mappings are now scoped against registered model namespaces, preventing text-only mappings from capturing an already nested VLM's ``model.visual.*`` namespace or double-prefixing ``model.language_model.*`` (observed on Qwen3.5). + 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index 516c0e90851..ece6e335ca8 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -287,6 +287,37 @@ def _assert_experts_pre_expanded( ) +def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[RenameRule]: + """Drop child-model reverse renames when the child namespace already exists. + + Transformers collects conversions recursively, so a text-only ``model.*`` -> + ``model.language_model.*`` reverse can also reach its parent VLM. In that model, + ``model.language_model`` is already registered and applying the rule globally + would capture both that namespace and siblings such as ``model.visual``. + """ + named_modules = getattr(model, "named_modules", None) + if not callable(named_modules): + return rules + + module_names = {name for name, _ in named_modules() if name} + probe_suffix = ".\x00modelopt_namespace_probe" + kept: list[RenameRule] = [] + for rule in rules: + pattern = re.compile(rule.pattern) + shadowed = False + for module_name in module_names: + mapped = pattern.sub(rule.repl, module_name + probe_suffix) + if not mapped.endswith(probe_suffix): + continue + mapped_parent = mapped[: -len(probe_suffix)] + if mapped_parent in module_names and mapped_parent.startswith(module_name + "."): + shadowed = True + break + if not shadowed: + kept.append(rule) + return kept + + def _build_reverse_rules(model) -> tuple[list[SplitRule], list[RenameRule], list[str]]: """Derive reverse rules from the model's transformers conversion mapping. @@ -373,6 +404,7 @@ def _build_reverse_rules(model) -> tuple[list[SplitRule], list[RenameRule], list # reorder rename runs last and does not destroy the anchor the MoE container/gate # renames rely on. Expert leaf renames act on disjoint ``.experts..`` # substrings and are applied first. + weight_renamings = _drop_shadowed_prefix_renames(model, weight_renamings) rename_rules = leaf_renamings + list(reversed(weight_renamings)) return split_rules, rename_rules, expert_fused_leaves diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index ac6df25e6dc..98bf6c250c2 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -264,6 +264,58 @@ def test_build_reverse_rules_orders_prefix_reorder_after_container(): assert not any(".mlp.experts." in k for k in out) +def test_nested_text_prefix_reverse_does_not_capture_vlm_siblings(): + """A nested text-model conversion must not rewrite the full VLM namespace.""" + pytest.importorskip("transformers.core_model_loading") + from transformers.core_model_loading import WeightRenaming + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.visual = torch.nn.Module() + model.model.visual.patch_embed = torch.nn.Linear(2, 2, bias=False) + model.model.language_model = torch.nn.Module() + model.model.language_model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)]) + model._weight_conversions = [ + WeightRenaming( + source_patterns=r"^model.language_model.", + target_patterns=r"^model.(?!language_model.)", + ) + ] + + state_dict = { + "model.visual.patch_embed.weight": torch.randn(2, 2), + "model.language_model.layers.0.weight": torch.randn(2, 2), + } + reverted = revert_weight_conversion_quant_aware(model, state_dict) + + assert set(reverted) == set(state_dict) + assert build_reverse_name_mapper(model) is None + + +def test_nested_text_prefix_reverse_still_applies_to_text_model(): + """The same conversion remains valid when the nested VLM namespace is absent.""" + pytest.importorskip("transformers.core_model_loading") + from transformers.core_model_loading import WeightRenaming + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)]) + model._weight_conversions = [ + WeightRenaming( + source_patterns=r"^model.language_model.", + target_patterns=r"^model.(?!language_model.)", + ) + ] + + state_dict = {"model.layers.0.weight": torch.randn(2, 2)} + reverted = revert_weight_conversion_quant_aware(model, state_dict) + + assert set(reverted) == {"model.language_model.layers.0.weight"} + mapper = build_reverse_name_mapper(model) + assert mapper is not None + assert mapper("model.layers.0") == "model.language_model.layers.0" + + def test_split_collision_raises(): """A split whose target key already exists must fail instead of overwriting.""" sd = _nvfp4_linear("m.gate_up_proj", 8, 16)