Fix/tied weight export identity

modelopt/torch/export/hf_export_handlers.py

@@ -24,7 +24,7 @@
 
 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 .moe_utils import _delete_fused_moe_source_attrs, _export_fused_experts
 from .quant_utils import get_quantization_format
 from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry
 
@@ -36,16 +36,39 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool:
     return hasattr(module, f"{first_proj_attr}_weight_quantizers")
 
 
+def _full_weight_name(module_name: str, weight_name: str) -> str:
+    return f"{module_name}.{weight_name}" if module_name else weight_name
+
+
+def _is_duplicate_with_same_format(
+    module_name: str,
+    ctx: ExportContext,
+    weight_name: str,
+) -> bool:
+    """Whether this source weight can be omitted in favor of its canonical tied name."""
+    full_name = _full_weight_name(module_name, weight_name)
+    canonical_name = ctx.duplicate_of(full_name)
+    if canonical_name is None:
+        return False
+    return ctx.weight_formats[canonical_name] == ctx.weight_formats[full_name]
+
+
 def _export_weight(
+    module_name: str,
     module: nn.Module,
     ctx: ExportContext,
     weight_name: str = "weight",
 ) -> None:
+    full_name = _full_weight_name(module_name, weight_name)
+    if _is_duplicate_with_same_format(module_name, ctx, weight_name):
+        ctx.mark_skipped(full_name)
+        return
+
     # 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)
+    _export_quantized_weight(module, ctx.dtype, weight_name)
 
 
 # Preparation handlers are registered in the same precedence as the legacy MoE prepass.
@@ -129,13 +152,27 @@ def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None
 @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."""
+    first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
+    first_name = _full_weight_name(name, first_proj_attr)
+    down_name = _full_weight_name(name, "down_proj")
+    first_canonical = ctx.duplicate_of(first_name)
+    down_canonical = ctx.duplicate_of(down_name)
+
+    # Omit the entire tied subtree when both source projections point to one canonical
+    # fused-experts module and use the same quantization representation.
+    if (
+        first_canonical is not None
+        and down_canonical is not None
+        and first_canonical.rsplit(".", 1)[0] == down_canonical.rsplit(".", 1)[0]
+        and _is_duplicate_with_same_format(name, ctx, first_proj_attr)
+        and _is_duplicate_with_same_format(name, ctx, "down_proj")
+    ):
+        ctx.mark_skipped(first_name, down_name)
+        _delete_fused_moe_source_attrs(module)
+        return
+
     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,
-        )
+        _export_fused_experts(module, ctx.dtype)
 
 
 @ExportModuleRegistry.register(predicate=is_quantlinear)
@@ -145,7 +182,7 @@ def _export_quant_linear(name: str, module: nn.Module, ctx: ExportContext) -> No
         return
     try:
         with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
-            _export_weight(module, ctx)
+            _export_weight(name, module, ctx)
     except AssertionError as e:
         raise AssertionError(
             f"Failed to export module '{name}' (type={type(module).__name__}): {e}"
@@ -176,7 +213,7 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) ->
         return
     try:
         with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
-            _export_weight(module, ctx)
+            _export_weight(name, module, ctx)
     except AssertionError as e:
         raise AssertionError(
             f"Failed to export embedding '{name}' (type={type(module).__name__}): {e}"
@@ -199,4 +236,4 @@ def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> Non
     )
     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)
+            _export_weight(name, module, ctx, weight_name)

modelopt/torch/export/moe_utils.py

@@ -23,35 +23,6 @@
 import torch.nn as nn
 
 
-def _alias_per_expert_subtree_from_prior(module: nn.Module, prior: nn.Module, n: int) -> None:
-    """Build per-expert subtree on ``module`` by aliasing ``prior``'s packed buffers.
-
-    For each expert ``idx`` in ``0..n-1``, creates ``module.{idx}.{gate,up,down}_proj``
-    sub-modules whose ``weight`` / ``weight_scale`` / ``weight_scale_2`` /
-    ``input_scale`` are aliased to the prior side's already-packed tensors.
-    data_ptr equality is preserved so the downstream
-    ``postprocess_state_dict`` dedup collapses the duplicates at write time.
-    Called by ``_export_fused_experts`` on the tied-experts cache-hit fast path.
-    """
-    for _idx in range(n):
-        _prior_expert = getattr(prior, str(_idx), None)
-        if _prior_expert is None:
-            continue
-        _cur_expert = nn.Module()
-        for _proj_name in ("gate_proj", "up_proj", "down_proj"):
-            _prior_proj = getattr(_prior_expert, _proj_name, None)
-            if _prior_proj is None:
-                continue
-            _cur_proj = nn.Module()
-            if hasattr(_prior_proj, "weight"):
-                _cur_proj.weight = _prior_proj.weight
-            for _attr in ("weight_scale", "weight_scale_2", "input_scale"):
-                if hasattr(_prior_proj, _attr):
-                    _cur_proj.register_buffer(_attr, getattr(_prior_proj, _attr))
-            _cur_expert.add_module(_proj_name, _cur_proj)
-        module.add_module(str(_idx), _cur_expert)
-
-
 def _delete_fused_moe_source_attrs(module: nn.Module) -> None:
     """Remove the 3-D fused source params and per-expert quantizer ModuleLists.
 
@@ -74,12 +45,7 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None:
             delattr(module, attr)
 
 
-def _export_fused_experts(
-    module: nn.Module,
-    dtype: torch.dtype,
-    _moe_tied_cache: dict[tuple[int, int], nn.Module] | None = None,
-    _tied_cache: dict[int, nn.Module] | None = None,
-) -> None:
+def _export_fused_experts(module: nn.Module, dtype: torch.dtype) -> None:
     """Split fused MoE expert weights and export per-expert quantization scales.
 
     Works with any module wrapped by ``_QuantFusedExperts`` (gated, with a fused
@@ -99,20 +65,6 @@ def _export_fused_experts(
            {E}.gate_proj.weight, {E}.gate_proj.weight_scale, ...  # gated only
            {E}.up_proj.weight, {E}.up_proj.weight_scale, ...
            {E}.down_proj.weight, {E}.down_proj.weight_scale, ...
-
-    Tied-experts dedup is opt-in via ``_moe_tied_cache``: when multiple
-    fused-expert modules share their 3-D source params via HF
-    ``_tied_weights_keys``, the unpacking creates fresh per-expert tensors
-    that break the tie. With ``_moe_tied_cache`` provided (tuple-keyed by
-    ``(<first_proj>.data_ptr(), down_proj.data_ptr())``), the alias step
-    at the end re-points the per-expert ``weight`` / ``weight_scale`` /
-    ``weight_scale_2`` / ``input_scale`` buffers at a previously-processed
-    module sharing the same source memory. ``_tied_cache`` (int-keyed) is
-    threaded through to the per-projection ``_export_quantized_weight``
-    calls so wrapper-level dedup uses the same scope as standalone Linears.
-    Both caches are owned by the caller (typically
-    ``_export_transformers_checkpoint``) and scoped to one export
-    invocation; when ``None`` the corresponding alias step is skipped.
     """
     from modelopt.torch.export.unified_export_hf import _export_quantized_weight
     from modelopt.torch.quantization.plugins.huggingface import _get_fused_expert_intermediate_dim
@@ -124,25 +76,6 @@ def _export_fused_experts(
     # Only the gated split needs the per-expert intermediate dim (gate|up boundary).
     expert_dim = _get_fused_expert_intermediate_dim(module) if is_gated else None
 
-    # Capture source tensor identities BEFORE unpacking (the source
-    # attrs are deleted at the end of this function).
-    _source_key = (
-        getattr(module, first_proj_attr).data_ptr(),
-        module.down_proj.data_ptr(),
-    )
-
-    # Tied-experts fast path: if this exact (first_proj, down) source-tensor pair
-    # has been processed before, alias all per-expert buffers directly from the
-    # prior module — no unpacking, no per-expert packing, no transient buffers
-    # thrown away. Cache miss falls through to the full unpack/pack below and
-    # registers this module as the prior for any later tied module.
-    if _moe_tied_cache is not None:
-        _prior = _moe_tied_cache.get(_source_key)
-        if _prior is not None and _prior is not module:
-            _alias_per_expert_subtree_from_prior(module, _prior, n)
-            _delete_fused_moe_source_attrs(module)
-            return
-
     # 1. Shared input quantizers — one per projection type, shared across all experts.
     first_proj_input_q = getattr(module, f"{first_proj_attr}_input_quantizer")
     first_proj_weight_quantizers = getattr(module, f"{first_proj_attr}_weight_quantizers")
@@ -271,7 +204,7 @@ def _export_fused_experts(
             wrapper.weight_quantizer = w_quantizer
             wrapper.input_quantizer = i_quantizer
 
-            _export_quantized_weight(wrapper, dtype, _tied_cache=_tied_cache)
+            _export_quantized_weight(wrapper, dtype)
 
             proj = nn.Module()
             proj.weight = wrapper.weight
@@ -286,13 +219,6 @@ def _export_fused_experts(
     # 4. Remove fused params and quantizer lists — replaced by per-expert submodules
     _delete_fused_moe_source_attrs(module)
 
-    # 5. Register this module in the dedup cache so any later tied module
-    # (same source data_ptr pair) takes the fast path at the top of this
-    # function. Reached only on cache miss; cache-hit modules early-exited
-    # above before any unpack work.
-    if _moe_tied_cache is not None:
-        _moe_tied_cache[_source_key] = module
-
 
 def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | None = None):
     """Collect expert_token_count from all quantized MoE layers and save as an HTML table.

modelopt/torch/export/quant_utils.py

@@ -1063,13 +1063,12 @@ def postprocess_state_dict(
     # Check for tied weights and remove duplicates
     seen_tensors = {}
 
-    # Remove any tied weights if found.
+    # Remove any tied weights if found. Device and size distinguish independent tensors whose
+    # allocator addresses happen to match. Zero-pointer tensors are left for serialization to reject.
     for key, value in post_state_dict.items():
-        if isinstance(value, torch.Tensor):
-            # Use tensor data pointer to identify tied weights
-            tensor_id = value.data_ptr()
+        if isinstance(value, torch.Tensor) and value.data_ptr() != 0:
+            tensor_id = (value.device, value.data_ptr(), value.numel() * value.element_size())
             if tensor_id in seen_tensors:
-                # This is a tied weight, mark for deletion and warn
                 keys_to_delete.append(key)
                 logger.warning(
                     f"Found tied weight: '{key}' is tied to '{seen_tensors[tensor_id]}'. "

modelopt/torch/export/registry.py

@@ -34,6 +34,8 @@
 
 from modelopt.torch.utils.distributed import is_fsdp2_model
 
+from .model_utils import _reorder_canonical_first
+
 __all__ = [
     "ExportContext",
     "ExportHandler",
@@ -46,27 +48,49 @@
 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.
+    ``duplicate_weight_map`` maps each duplicate source-parameter name to the canonical
+    name that should survive export. It is snapshotted before export mutates Parameters.
     """
 
     model: nn.Module
     dtype: torch.dtype
     is_modelopt_qlora: bool = False
-    tied_cache: dict[int, nn.Module] | None = field(default_factory=dict)
-    moe_tied_cache: dict[tuple[int, int], nn.Module] | None = field(default_factory=dict)
+    duplicate_weight_map: dict[str, str] = field(init=False, default_factory=dict)
+    weight_locations: dict[str, tuple[nn.Module, str]] = field(init=False, default_factory=dict)
+    weight_formats: dict[str, str | None] = field(init=False, default_factory=dict)
+    skipped_weight_names: set[str] = field(init=False, default_factory=set)
 
     def __post_init__(self) -> None:
-        # FSDP2 may recycle data_ptr() values as modules are resharded, so pointer-keyed dedup can
-        # falsely alias distinct weights. Disable it for FSDP2; consequently, legitimately tied
-        # packed weights and scale buffers are not re-aliased and may be stored as duplicates.
-        # TODO: replace this with stable, name-based tied-group deduplication.
+        # FSDP2 replaces Parameters while resharding. Keep its existing behavior of writing
+        # tied packed weights independently instead of relying on pre-reshard identities.
         if is_fsdp2_model(self.model):
-            self.tied_cache = None
-            self.moe_tied_cache = None
+            return
+
+        names_by_source_id: dict[int, list[str]] = {}
+        for module_name, module in self.model.named_modules(remove_duplicate=False):
+            for weight_name, parameter in module._parameters.items():
+                if parameter is None:
+                    continue
+                full_name = f"{module_name}.{weight_name}" if module_name else weight_name
+                self.weight_locations[full_name] = (module, weight_name)
+                names_by_source_id.setdefault(id(parameter), []).append(full_name)
+
+        for names in names_by_source_id.values():
+            if len(names) < 2:
+                continue
+            # Reuse the existing HF canonical-name policy. For models without an explicit
+            # canonical side this preserves traversal order and keeps the first name.
+            ordered = list(_reorder_canonical_first(dict.fromkeys(names), self.model))
+            canonical = ordered[0]
+            self.duplicate_weight_map.update(dict.fromkeys(ordered[1:], canonical))
+
+    def duplicate_of(self, weight_name: str) -> str | None:
+        """Return the canonical source weight name, or ``None`` if it is not duplicated."""
+        return self.duplicate_weight_map.get(weight_name)
+
+    def mark_skipped(self, *weight_names: str) -> None:
+        """Record duplicate source weights omitted from packing and the state dict."""
+        self.skipped_weight_names.update(weight_names)
 
 
 ExportHandler = Callable[[str, nn.Module, ExportContext], None]