Registry-based module dispatch for unified_export_hf.py

modelopt/torch/export/__init__.py

@@ -21,6 +21,7 @@
 from .model_utils import *
 from .moe_utils import *
 from .plugins import *
+from .registry import *
 from .transformer_engine import *
 from .unified_export_hf import *
 from .unified_export_megatron import *

modelopt/torch/export/hf_export_handlers.py

@@ -0,0 +1,202 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Built-in module handlers for unified Hugging Face export."""
+
+import collections.abc
+import warnings
+
+import torch.nn as nn
+
+from modelopt.torch.quantization.utils import fsdp2_aware_weight_update
+
+from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax
+from .model_config import QUANTIZATION_NONE
+from .moe_utils import _export_fused_experts
+from .quant_utils import get_quantization_format
+from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry
+
+__all__: list[str] = []
+
+
+def _has_fused_experts_quantizers(module: nn.Module) -> bool:
+    first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
+    return hasattr(module, f"{first_proj_attr}_weight_quantizers")
+
+
+def _export_weight(
+    module: nn.Module,
+    ctx: ExportContext,
+    weight_name: str = "weight",
+) -> None:
+    # Imported lazily to avoid a cycle: unified_export_hf imports this module to
+    # install the built-in handlers while retaining this legacy helper's import path.
+    from .unified_export_hf import _export_quantized_weight
+
+    _export_quantized_weight(module, ctx.dtype, weight_name, _tied_cache=ctx.tied_cache)
+
+
+# Preparation handlers are registered in the same precedence as the legacy MoE prepass.
+
+
+# Keyed on the mixin class name too: the generated class is normally named
+# "QuantDbrxExperts", but _DMRegistryCls falls back to a module-prefixed name on
+# collision, while "_QuantDbrxExperts" remains in the generated class's MRO.
+@PrepareMoEInputsRegistry.register("QuantDbrxExperts", "_QuantDbrxExperts")
+def _prepare_dbrx_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None:
+    """Fill missing input amax values for DBRX per-expert ModuleLists."""
+    experts_mlp = moe_module.experts.mlp
+    for linear_name in get_expert_linear_names(moe_module):
+        if hasattr(experts_mlp, linear_name):
+            linear_modulelist = getattr(experts_mlp, linear_name)
+            if hasattr(linear_modulelist, "__iter__"):
+                set_expert_quantizer_amax(
+                    modules=list(linear_modulelist),
+                    quantizer_attrs=["input_quantizer"],
+                )
+
+
+@PrepareMoEInputsRegistry.register(predicate=_has_fused_experts_quantizers)
+def _prepare_fused_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None:
+    """Mark fused experts handled; their missing amax fallback occurs during export."""
+
+
+@PrepareMoEInputsRegistry.register("Llama4TextExperts", "GptOssExperts")
+def _prepare_bmm_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None:
+    """Fill missing input amax values for fused BMM-style experts."""
+    # Both use gate_up_proj and down_proj with singular input quantizers
+    # (gate_up_proj_input_quantizer/down_proj_input_quantizer); the weight-side
+    # amax fallback and weight export happen in _export_bmm_experts.
+    for linear_name in ["gate_up_proj", "down_proj"]:
+        if hasattr(moe_module.experts, linear_name):
+            linear_module = getattr(moe_module.experts, linear_name)
+            if hasattr(linear_module, "input_quantizer"):
+                set_expert_quantizer_amax(
+                    modules=[linear_module],
+                    quantizer_attrs=["input_quantizer"],
+                )
+
+
+@PrepareMoEInputsRegistry.register(
+    predicate=lambda module: isinstance(module, collections.abc.Iterable)
+)
+def _prepare_iterable_experts(name: str, moe_module: nn.Module, ctx: ExportContext) -> None:
+    """Fill missing input amax values for iterable per-expert submodules."""
+    expert_linear_names = get_expert_linear_names(moe_module)
+    linear_name = None
+    try:
+        for linear_name in expert_linear_names:
+            set_expert_quantizer_amax(
+                modules=[getattr(expert, linear_name) for expert in moe_module.experts],
+                quantizer_attrs=["input_quantizer"],
+            )
+    except AttributeError as e:
+        expert_types = [type(expert).__name__ for expert in moe_module.experts]
+        raise AttributeError(
+            f"Failed to access attribute '{linear_name}' on experts. "
+            f"MoE module type: {type(moe_module).__name__}, "
+            f"Expert types: {expert_types}, "
+            f"Expected linear names: {expert_linear_names}. "
+            f"This suggests the get_expert_linear_names function may need "
+            f"to be updated for this model architecture. "
+            f"Original error: {e}"
+        ) from e
+
+
+# Export handlers are registered in the same precedence as the legacy model walk.
+
+
+@ExportModuleRegistry.register(
+    "QuantMoELinear", predicate=lambda module: hasattr(module, "experts")
+)
+def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None:
+    """Fill missing input amax before child expert QuantLinears are exported."""
+    set_expert_quantizer_amax(list(module.experts), quantizer_attrs="input_quantizer")
+
+
+@ExportModuleRegistry.register(predicate=_has_fused_experts_quantizers)
+def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContext) -> None:
+    """Split and quantize a fused-experts module with plural weight quantizers."""
+    with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
+        _export_fused_experts(
+            module,
+            ctx.dtype,
+            _moe_tied_cache=ctx.moe_tied_cache,
+            _tied_cache=ctx.tied_cache,
+        )
+
+
+@ExportModuleRegistry.register(predicate=is_quantlinear)
+def _export_quant_linear(name: str, module: nn.Module, ctx: ExportContext) -> None:
+    """Export a standard quantized linear layer."""
+    if get_quantization_format(module) == QUANTIZATION_NONE:
+        return
+    try:
+        with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
+            _export_weight(module, ctx)
+    except AssertionError as e:
+        raise AssertionError(
+            f"Failed to export module '{name}' (type={type(module).__name__}): {e}"
+        ) from e
+
+
+@ExportModuleRegistry.register(
+    nn.Embedding, predicate=lambda module: hasattr(module, "weight_quantizer")
+)
+def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) -> None:
+    """Export a quantized embedding table unless its weight is tied."""
+    if get_quantization_format(module) == QUANTIZATION_NONE:
+        return
+    # Packing replaces .weight, which would sever any Python-level weight tie and
+    # leave the other module pointing at a stale float Parameter.
+    tied_to = [
+        other_name
+        for other_name, other_module in ctx.model.named_modules()
+        if other_module is not module and getattr(other_module, "weight", None) is module.weight
+    ]
+    if tied_to:
+        warnings.warn(
+            f"Skipping quantized weight packing for embedding '{name}': its "
+            f"weight Parameter is shared with {tied_to} (weight tying). Packing "
+            "would break the tie and produce stale weights in the tied module(s). "
+            "The embedding will be exported as its fake-quantized float weight."
+        )
+        return
+    try:
+        with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
+            _export_weight(module, ctx)
+    except AssertionError as e:
+        raise AssertionError(
+            f"Failed to export embedding '{name}' (type={type(module).__name__}): {e}"
+        ) from e
+
+
+@ExportModuleRegistry.register("Llama4TextExperts", "GptOssExperts")
+def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> None:
+    """Export fused BMM-style expert weights and quantization metadata."""
+    if get_quantization_format(module) == QUANTIZATION_NONE:
+        return
+    # TODO: consolidate uncalibrated experts handling logic
+    set_expert_quantizer_amax(
+        modules=module,
+        quantizer_attrs=["gate_up_proj_weight_quantizer", "down_proj_weight_quantizer"],
+    )
+    set_expert_quantizer_amax(
+        modules=module,
+        quantizer_attrs=["gate_up_proj_input_quantizer", "down_proj_input_quantizer"],
+    )
+    with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
+        for weight_name in ["gate_up_proj", "down_proj"]:
+            _export_weight(module, ctx, weight_name)

modelopt/torch/export/registry.py

@@ -0,0 +1,146 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Registries dispatching per-module logic for the unified HF export path.
+
+This mirrors the registration-and-dispatch idiom of
+:class:`QuantModuleRegistry <modelopt.torch.quantization.nn.modules.quant_module.QuantModuleRegistry>`,
+but not its mechanism: quantization registers replacement classes and converts modules
+in place, whereas export registers functions that emit compressed weights and scale
+buffers for a module without changing its class.
+
+Preparation and export use separate registries because they have independent matching
+precedence. Registering a handler for a new module type replaces what previously required
+editing if/elif chains inside ``unified_export_hf.py``.
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+
+__all__ = [
+    "ExportContext",
+    "ExportHandler",
+    "ExportModuleRegistry",
+    "PrepareMoEInputsRegistry",
+]
+
+
+@dataclass
+class ExportContext:
+    """Shared state for a single export invocation, passed to every handler call.
+
+    The tied-weight dedup caches must be scoped to one export invocation: a
+    process-global cache would carry stale entries whose ``data_ptr`` keys can be
+    recycled by PyTorch's allocator across exports, causing silent false-positive
+    aliasing. ``tied_cache`` (int keys) holds dense Linear / per-expert wrapper
+    dedup; ``moe_tied_cache`` (tuple keys) holds MoE fused-experts module dedup.
+    """
+
+    model: nn.Module
+    dtype: torch.dtype
+    is_modelopt_qlora: bool = False
+    tied_cache: dict[int, nn.Module] = field(default_factory=dict)
+    moe_tied_cache: dict[tuple[int, int], nn.Module] = field(default_factory=dict)
+
+
+ExportHandler = Callable[[str, nn.Module, ExportContext], None]
+
+
+class _ExportHandlerRegistryCls:
+    """Ordered, first-match-wins registry mapping modules to handler functions.
+
+    An entry can match a module by any combination of:
+
+    - a class key: the registered class appears in ``type(module).__mro__``, so
+      dynamically generated quantized classes (e.g. ``QuantLinear``) match through
+      their original base class;
+    - a class-name string key: the string equals the ``__name__`` of a class in the
+      MRO — for classes that cannot be imported statically (trust_remote_code models
+      or on-the-fly generated quantized classes);
+    - a predicate on the module instance, for structural detection.
+
+    When keys and a predicate are both given, both must match. Entries are tried in
+    registration order and the first match wins, so more specific handlers must be
+    registered before generic ones. External handlers can use ``prepend=True`` to take
+    precedence over built-in entries.
+    """
+
+    def __init__(self) -> None:
+        self._entries: list[
+            tuple[
+                tuple[type | str, ...],
+                Callable[[nn.Module], bool] | None,
+                ExportHandler,
+            ]
+        ] = []
+
+    def register(
+        self,
+        *keys: type | str,
+        predicate: Callable[[nn.Module], bool] | None = None,
+        prepend: bool = False,
+    ) -> Callable[[ExportHandler], ExportHandler]:
+        """Return a decorator registering a handler function.
+
+        Re-registering the same handler (e.g. on module reload) replaces its existing
+        entry in place instead of appending a duplicate.
+
+        Usage::
+
+            @ExportModuleRegistry.register(
+                "Llama4TextExperts",
+                "GptOssExperts",
+            )
+            def _export_bmm_experts(name, module, ctx): ...
+        """
+        assert keys or predicate is not None, "register() requires at least one key or a predicate"
+
+        def decorator(handler: ExportHandler) -> ExportHandler:
+            entry = (keys, predicate, handler)
+            identity = (handler.__module__, handler.__qualname__)
+            for i, (_, _, existing) in enumerate(self._entries):
+                if (existing.__module__, existing.__qualname__) == identity:
+                    self._entries[i] = entry
+                    return handler
+            if prepend:
+                self._entries.insert(0, entry)
+            else:
+                self._entries.append(entry)
+            return handler
+
+        return decorator
+
+    def match(self, module: nn.Module) -> ExportHandler | None:
+        """Return the first registered handler matching ``module``, or ``None``."""
+        mro = type(module).__mro__
+        mro_names = {cls.__name__ for cls in mro}
+        for keys, predicate, handler in self._entries:
+            if keys and not any(
+                key in mro_names if isinstance(key, str) else key in mro for key in keys
+            ):
+                continue
+            if predicate is not None and not predicate(module):
+                continue
+            return handler
+        return None
+
+
+# Matches an MoE block's ``.experts`` container; the handler receives the enclosing block.
+PrepareMoEInputsRegistry = _ExportHandlerRegistryCls()
+# Matches and passes the same module to the handler during the whole-model export walk.
+ExportModuleRegistry = _ExportHandlerRegistryCls()