diff --git a/modelopt/torch/export/__init__.py b/modelopt/torch/export/__init__.py index 004788bbf14..9a4624a25e7 100644 --- a/modelopt/torch/export/__init__.py +++ b/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 * diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py new file mode 100644 index 00000000000..800b51daca9 --- /dev/null +++ b/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) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py new file mode 100644 index 00000000000..8e2cda63df9 --- /dev/null +++ b/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 `, +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() diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index a903adfc846..cee64c22c05 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -15,7 +15,6 @@ """Code that export quantized Hugging Face models for deployment.""" -import collections.abc import json import re import tempfile @@ -65,14 +64,14 @@ except ImportError: export_sparse_attention_config = None +# Importing the built-in handlers installs their entries in the two registries. +from . import hf_export_handlers as _hf_export_handlers # noqa: F401 from .convert_hf_config import convert_hf_quant_config_format from .layer_utils import ( - get_expert_linear_names, get_experts_list, is_layernorm, is_moe, is_quantlinear, - set_expert_quantizer_amax, sync_moe_gate_up_amax, ) from .model_config import ( @@ -89,7 +88,6 @@ QUANTIZATION_W4A16_NVFP4, ) from .model_utils import _reorder_canonical_first, get_language_model_from_vl, is_multimodal_model -from .moe_utils import _export_fused_experts from .plugins import SpeculativeDecodingExporter, has_spec_opt, sanitize_hf_config_for_deployment from .quant_aware_conversion import ( build_reverse_name_mapper, @@ -112,6 +110,7 @@ sync_tied_input_amax, to_quantized_weight, ) +from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry __all__ = ["export_hf_checkpoint", "export_speculative_decoding"] @@ -786,9 +785,8 @@ def _process_quantized_modules( ) -> None: """Process all quantized modules in model, export weights in-place. - This function iterates through all modules in the model and exports quantized weights - for modules that have quantization enabled. It handles both standard linear layers - and specialized expert modules (Llama4TextExperts, GptOssExperts). + This function iterates through all modules in the model and invokes the first matching + handler in :data:`ExportModuleRegistry`. Modules matching no handler are left untouched. Args: model: The model containing quantized modules. @@ -796,14 +794,10 @@ def _process_quantized_modules( is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. If True, modules with base_layer attribute are skipped. """ - # Per-call tied-weight dedup caches. Created fresh on every invocation - # so cache state is scoped to one export and cannot leak into a later - # call (a process-global cache would carry stale entries whose data_ptr - # keys can be recycled by PyTorch's allocator across exports — silent - # false-positive aliasing). int keys hold dense Linear / per-expert - # wrapper dedup; tuple keys hold MoE fused-experts module dedup. - _tied_cache: dict[int, nn.Module] = {} - _moe_tied_cache: dict[tuple[int, int], nn.Module] = {} + # Per-call tied-weight dedup caches inside the context. Created fresh on + # every invocation so cache state is scoped to one export and cannot leak + # into a later call (see ExportContext). + ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) fsdp_module_to_reshard = None for name, sub_module in model.named_modules(): @@ -818,97 +812,19 @@ def _process_quantized_modules( fsdp_module_to_reshard = sub_module # We skip QuantLoraLinear module for modelopt QLoRA - if is_modelopt_qlora and (hasattr(sub_module, "base_layer")): - continue - - # Step-3.5 QuantMoELinear reconstructs packed MoE tensors from child - # expert QuantLinears after export. Fill missing input amax here, before - # named_modules() reaches those children, so every expert emits input_scale. - if type(sub_module).__name__ == "QuantMoELinear" and hasattr(sub_module, "experts"): - set_expert_quantizer_amax(list(sub_module.experts), quantizer_attrs="input_quantizer") + if ctx.is_modelopt_qlora and hasattr(sub_module, "base_layer"): continue # Preprocessing: restore unpacked weight so the export path can read - # the live quantizer state. Falls through to the export branches below. + # the live quantizer state. Falls through to the handler dispatch below. if hasattr(sub_module, "weight_packed") or ( "QuantFP8Linear" in type(sub_module).__name__ and sub_module.weight.element_size() <= 1 ): sub_module.unpack_weight() - first_proj_attr = getattr(sub_module, "_first_proj_attr", "gate_up_proj") - if hasattr(sub_module, f"{first_proj_attr}_weight_quantizers"): - # _QuantFusedExperts uses plural `_weight_quantizers` - # (ModuleList), which get_quantization_format's singular-weight_quantizer - # check misses. Handle it explicitly before the format gate so fused-experts - # get split + quantized. - with fsdp2_aware_weight_update(model, sub_module, reshard=False): - _export_fused_experts( - sub_module, - dtype, - _moe_tied_cache=_moe_tied_cache, - _tied_cache=_tied_cache, - ) - elif get_quantization_format(sub_module) != QUANTIZATION_NONE: - # Skip QuantMoELinear - it's handled separately in _reconstruct_fused_moe_linear - if type(sub_module).__name__ == "QuantMoELinear": - continue - if is_quantlinear(sub_module): - try: - with fsdp2_aware_weight_update(model, sub_module, reshard=False): - _export_quantized_weight(sub_module, dtype, _tied_cache=_tied_cache) - except AssertionError as e: - raise AssertionError( - f"Failed to export module '{name}' (type={type(sub_module).__name__}): {e}" - ) from e - elif isinstance(sub_module, nn.Embedding) and hasattr(sub_module, "weight_quantizer"): - # Quantized nn.Embedding: pack the embedding table the same way as Linear - # weights so downstream loaders see the NVFP4/FP8/INT-packed bytes + scales. - # Skip packing when the embedding's weight is tied to another module - # (e.g. tied_word_embeddings → lm_head): _export_quantized_weight reassigns - # the .weight attribute to a new uint8 Parameter, which severs the Python- - # level tie and leaves the other module pointing at a stale float Parameter. - tied_to = [ - other_name - for other_name, other_module in model.named_modules() - if other_module is not sub_module - and getattr(other_module, "weight", None) is sub_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." - ) - else: - try: - with fsdp2_aware_weight_update(model, sub_module, reshard=False): - _export_quantized_weight(sub_module, dtype, _tied_cache=_tied_cache) - except AssertionError as e: - raise AssertionError( - f"Failed to export embedding '{name}' (type={type(sub_module).__name__}): {e}" - ) from e - elif ( - "Llama4TextExperts" in type(sub_module).__name__ - or "GptOssExperts" in type(sub_module).__name__ - ): - # TODO: consolidate uncalibrated experts handling logic - # Handle weight quantizers amax values using smart fallback logic - set_expert_quantizer_amax( - modules=sub_module, - quantizer_attrs=["gate_up_proj_weight_quantizer", "down_proj_weight_quantizer"], - ) - # Handle input quantizers amax values using smart fallback logic - set_expert_quantizer_amax( - modules=sub_module, - quantizer_attrs=["gate_up_proj_input_quantizer", "down_proj_input_quantizer"], - ) - # Export the quantized weights - with fsdp2_aware_weight_update(model, sub_module, reshard=False): - for weight_name in ["gate_up_proj", "down_proj"]: - _export_quantized_weight( - sub_module, dtype, weight_name, _tied_cache=_tied_cache - ) + handler = ExportModuleRegistry.match(sub_module) + if handler is not None: + handler(name, sub_module, ctx) def _export_transformers_checkpoint( @@ -940,71 +856,19 @@ def _export_transformers_checkpoint( accelerator = kwargs.get("accelerator") - # Handle input quantizers of experts that are not calibrated - for _, sub_module in model.named_modules(): + # Handle input quantizers of experts that are not calibrated. Each MoE block is + # dispatched by its experts container to the matching preparation handler. + prepare_ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + for name, sub_module in model.named_modules(): if is_moe(sub_module) and hasattr(sub_module, "experts"): - expert_linear_names = get_expert_linear_names(sub_module) - first_proj_attr = getattr(sub_module.experts, "_first_proj_attr", "gate_up_proj") - has_fused_experts_quantizers = hasattr( - sub_module.experts, f"{first_proj_attr}_weight_quantizers" - ) - for linear_name in expert_linear_names: - # Handle DBRX experts specifically - if "QuantDbrxExperts" in type(sub_module.experts).__name__: - # For DBRX, experts are in sub_module.experts.mlp and linear layers are ModuleLists - experts_mlp = sub_module.experts.mlp - 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"], - ) - elif has_fused_experts_quantizers: - # _QuantFusedExperts: amax fallback is handled in _export_fused_experts - break - elif ( - "QuantGptOssExperts" in type(sub_module.experts).__name__ - or "QuantLlama4TextExperts" in type(sub_module.experts).__name__ - ): - # Handle GPT-OSS / Llama4 fused experts specifically. - # Both use gate_up_proj and down_proj with singular input quantizers - # (gate_up_proj_input_quantizer/down_proj_input_quantizer); the actual - # amax fallback and weight export is performed in _process_quantized_modules. - gpt_oss_linear_names = ["gate_up_proj", "down_proj"] - for linear_name in gpt_oss_linear_names: - if hasattr(sub_module.experts, linear_name): - linear_module = getattr(sub_module.experts, linear_name) - if hasattr(linear_module, "input_quantizer"): - set_expert_quantizer_amax( - modules=[linear_module], - quantizer_attrs=["input_quantizer"], - ) - elif isinstance(sub_module.experts, collections.abc.Iterable): - # For other MoE models (like Mixtral) with iterable experts - try: - set_expert_quantizer_amax( - modules=[getattr(expert, linear_name) for expert in sub_module.experts], - quantizer_attrs=["input_quantizer"], - ) - except AttributeError as e: - # Provide more helpful debugging information - expert_types = [type(expert).__name__ for expert in sub_module.experts] - raise AttributeError( - f"Failed to access attribute '{linear_name}' on experts. " - f"MoE module type: {type(sub_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 - else: - # Unsupported MoE model structure - raise NotImplementedError( - f"MoE model with experts type '{type(sub_module.experts).__name__}' is not supported in export." - f"Please file an issue or add support for this model architecture." - ) + handler = PrepareMoEInputsRegistry.match(sub_module.experts) + if handler is None: + # Unsupported MoE model structure + raise NotImplementedError( + f"MoE model with experts type '{type(sub_module.experts).__name__}' is not supported in export." + f"Please file an issue or add support for this model architecture." + ) + handler(name, sub_module, prepare_ctx) # Resmooth and requantize fused layers # TODO: Handle mixed precision diff --git a/tests/unit/torch/export/test_export_registry.py b/tests/unit/torch/export/test_export_registry.py new file mode 100644 index 00000000000..67647f9c31f --- /dev/null +++ b/tests/unit/torch/export/test_export_registry.py @@ -0,0 +1,312 @@ +# 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. + +"""Tests for the registries dispatching unified Hugging Face export handlers.""" + +import pytest +import torch +import torch.nn as nn +from _test_utils.torch.export.utils import ToyModel, partial_fp8_config + +import modelopt.torch.quantization as mtq +from modelopt.torch.export import hf_export_handlers, unified_export_hf +from modelopt.torch.export.hf_export_handlers import ( + _export_bmm_experts, + _export_fused_experts_module, + _export_moe_linear, + _export_quant_embedding, + _export_quant_linear, + _prepare_bmm_experts, + _prepare_dbrx_experts, + _prepare_fused_experts, + _prepare_iterable_experts, +) +from modelopt.torch.export.registry import ( + ExportContext, + ExportModuleRegistry, + PrepareMoEInputsRegistry, + _ExportHandlerRegistryCls, +) +from modelopt.torch.export.unified_export_hf import _process_quantized_modules + + +class _Experts(nn.Module): + pass + + +def _make_dynamic_subclass(base: type[nn.Module], prefix: str = "Quant") -> type[nn.Module]: + """Mimic _DMRegistryCls's on-the-fly class generation (e.g. QuantLinear).""" + return type(f"{prefix}{base.__name__}", (base,), {}) + + +def test_unified_export_imports_builtin_handlers(): + assert unified_export_hf._hf_export_handlers is hf_export_handlers + + +def test_class_key_matches_generated_subclass_via_mro(): + registry = _ExportHandlerRegistryCls() + + @registry.register(nn.Linear) + def linear_handler(name, module, ctx): + pass + + quant_linear_cls = _make_dynamic_subclass(nn.Linear) + module = nn.Linear(2, 2) + module.__class__ = quant_linear_cls + + assert registry.match(module) is linear_handler + assert registry.match(nn.Linear(2, 2)) is linear_handler + assert registry.match(nn.Embedding(2, 2)) is None + + +def test_name_key_matches_original_and_generated_class(): + registry = _ExportHandlerRegistryCls() + + @registry.register("_Experts") + def experts_handler(name, module, ctx): + pass + + raw = _Experts() + generated = _Experts() + generated.__class__ = _make_dynamic_subclass(_Experts) + + # The original class name appears in the generated class's MRO. + assert registry.match(raw) is experts_handler + assert registry.match(generated) is experts_handler + assert registry.match(nn.Linear(2, 2)) is None + + +def test_keys_and_predicate_are_both_required(): + registry = _ExportHandlerRegistryCls() + + @registry.register("_Experts", predicate=lambda module: hasattr(module, "experts")) + def experts_handler(name, module, ctx): + pass + + without_experts = _Experts() + with_experts = _Experts() + with_experts.experts = nn.ModuleList([nn.Linear(2, 2)]) + + assert registry.match(without_experts) is None + assert registry.match(with_experts) is experts_handler + + +def test_first_registered_entry_wins(): + registry = _ExportHandlerRegistryCls() + + @registry.register(predicate=lambda module: isinstance(module, nn.Linear)) + def specific_handler(name, module, ctx): + pass + + @registry.register(nn.Module) + def generic_handler(name, module, ctx): + pass + + assert registry.match(nn.Linear(2, 2)) is specific_handler + assert registry.match(nn.Embedding(2, 2)) is generic_handler + + +def test_registration_order_is_independent_between_registries(): + prepare_registry = _ExportHandlerRegistryCls() + export_registry = _ExportHandlerRegistryCls() + + def handler_a(name, module, ctx): + pass + + def handler_b(name, module, ctx): + pass + + def handler_c(name, module, ctx): + pass + + for handler in [handler_a, handler_b, handler_c]: + prepare_registry.register(nn.Module)(handler) + for handler in [handler_b, handler_a, handler_c]: + export_registry.register(nn.Module)(handler) + + assert [handler for _, _, handler in prepare_registry._entries] == [ + handler_a, + handler_b, + handler_c, + ] + assert [handler for _, _, handler in export_registry._entries] == [ + handler_b, + handler_a, + handler_c, + ] + + module = nn.Linear(2, 2) + assert prepare_registry.match(module) is handler_a + assert export_registry.match(module) is handler_b + + +def test_register_requires_key_or_predicate(): + registry = _ExportHandlerRegistryCls() + + def handler(name, module, ctx): + pass + + with pytest.raises(AssertionError): + registry.register()(handler) + + +def test_prepend_registers_before_existing_entries(): + registry = _ExportHandlerRegistryCls() + + @registry.register(predicate=lambda module: True) + def catch_all_handler(name, module, ctx): + pass + + @registry.register(nn.Linear, prepend=True) + def specific_handler(name, module, ctx): + pass + + assert registry.match(nn.Linear(2, 2)) is specific_handler + assert registry.match(nn.Embedding(2, 2)) is catch_all_handler + + +def test_reregistering_same_handler_replaces_entry_in_place(): + registry = _ExportHandlerRegistryCls() + + @registry.register(nn.Linear) + def first_handler(name, module, ctx): + pass + + @registry.register(predicate=lambda module: True) + def catch_all_handler(name, module, ctx): + pass + + # Simulate a module reload changing the same function's registration: + # the entry is replaced in place, keeping its winning position. + registry.register(nn.Embedding)(first_handler) + assert len(registry._entries) == 2 + assert registry.match(nn.Linear(2, 2)) is catch_all_handler + assert registry.match(nn.Embedding(2, 2)) is first_handler + + +def _named_module(name: str, base_name: str | None = None, **attrs) -> nn.Module: + """Create a module instance whose class (and optionally base class) has a given name.""" + base = type(base_name, (nn.Module,), {}) if base_name else nn.Module + module = type(name, (base,), {})() + for attr, value in attrs.items(): + setattr(module, attr, value) + return module + + +def test_builtin_dispatch_covers_all_handler_shapes(): + # Step-3.5 QuantMoELinear wrapper: exact leaf-class name plus experts attr. + moe_linear = _named_module("QuantMoELinear", experts=nn.ModuleList([nn.Linear(2, 2)])) + assert ExportModuleRegistry.match(moe_linear) is _export_moe_linear + assert PrepareMoEInputsRegistry.match(moe_linear) is None + assert ExportModuleRegistry.match(_named_module("QuantMoELinear")) is None + + # DBRX experts container: the usual generated class name... + dbrx = _named_module("QuantDbrxExperts", base_name="DbrxExperts") + assert PrepareMoEInputsRegistry.match(dbrx) is _prepare_dbrx_experts + assert ExportModuleRegistry.match(dbrx) is None + # ...and the _DMRegistryCls collision-fallback name, where only the mixin + # class name ("_QuantDbrxExperts") remains recognizable in the MRO. + fallback = _named_module( + "transformers_modules_modeling_dbrx_QuantDbrxExperts", base_name="_QuantDbrxExperts" + ) + assert PrepareMoEInputsRegistry.match(fallback) is _prepare_dbrx_experts + + # DBRX preparation remains type-specific even if a future DBRX variant gains + # plural quantizers; the independent export registry takes the fused path. + fused_dbrx = _named_module( + "QuantDbrxExperts", + base_name="DbrxExperts", + gate_up_proj_weight_quantizers=nn.ModuleList(), + ) + assert PrepareMoEInputsRegistry.match(fused_dbrx) is _prepare_dbrx_experts + assert ExportModuleRegistry.match(fused_dbrx) is _export_fused_experts_module + + # Structural fused-experts matching wins over BMM name matching independently + # in both registries. + fused = _named_module( + "QuantGptOssExperts", + base_name="GptOssExperts", + gate_up_proj_weight_quantizers=nn.ModuleList(), + ) + assert PrepareMoEInputsRegistry.match(fused) is _prepare_fused_experts + assert ExportModuleRegistry.match(fused) is _export_fused_experts_module + + # Structural fused matching also takes precedence over the iterable fallback. + fused_iterable = nn.ModuleList() + fused_iterable.gate_up_proj_weight_quantizers = nn.ModuleList() + assert PrepareMoEInputsRegistry.match(fused_iterable) is _prepare_fused_experts + assert ExportModuleRegistry.match(fused_iterable) is _export_fused_experts_module + + # BMM-style experts match through raw or quant-generated class names. + for bmm in [ + _named_module("GptOssExperts"), + _named_module("QuantLlama4TextExperts", base_name="Llama4TextExperts"), + ]: + assert PrepareMoEInputsRegistry.match(bmm) is _prepare_bmm_experts + assert ExportModuleRegistry.match(bmm) is _export_bmm_experts + + opaque = _named_module("OpaqueExperts") + assert PrepareMoEInputsRegistry.match(opaque) is None + assert ExportModuleRegistry.match(opaque) is None + + +@pytest.mark.parametrize( + "iterable_type", + [nn.ModuleList, nn.Sequential, nn.ModuleDict, nn.ParameterList], +) +def test_iterable_modules_only_match_preparation_registry(iterable_type): + iterable = iterable_type() + assert PrepareMoEInputsRegistry.match(iterable) is _prepare_iterable_experts + assert ExportModuleRegistry.match(iterable) is None + + +def test_builtin_registry_dispatches_quantized_modules(): + model = ToyModel(dims=[16, 32, 16]) + mtq.quantize(model, partial_fp8_config, lambda module: module(torch.randn(2, 4, 16))) + + quantized = [module for module in model.modules() if type(module).__name__ == "QuantLinear"] + assert quantized, "expected at least one quantized linear in the toy model" + for module in quantized: + assert ExportModuleRegistry.match(module) is _export_quant_linear + + # Plain (unquantized) modules match no export handler. + assert ExportModuleRegistry.match(nn.Linear(2, 2)) is None + + embedding = nn.Embedding(4, 4) + embedding.weight_quantizer = None + assert ExportModuleRegistry.match(embedding) is _export_quant_embedding + + +def test_process_quantized_modules_exports_via_registry(): + model = ToyModel(dims=[16, 32, 16]) + mtq.quantize(model, partial_fp8_config, lambda module: module(torch.randn(2, 4, 16))) + + _process_quantized_modules(model, torch.float16) + + state_dict = model.state_dict() + fp8_weights = [key for key in state_dict if key.endswith("weight_scale")] + assert fp8_weights, "expected weight_scale buffers registered by the linear handler" + for key in fp8_weights: + weight = state_dict[key.replace("weight_scale", "weight")] + assert weight.dtype == torch.float8_e4m3fn + + +def test_export_context_caches_are_per_instance(): + model = nn.Linear(2, 2) + ctx_a = ExportContext(model=model, dtype=torch.float16) + ctx_b = ExportContext(model=model, dtype=torch.float16) + ctx_a.tied_cache[123] = model + assert ctx_b.tied_cache == {} + assert ctx_b.moe_tied_cache == {}