From 6da5ff9e47ac3eefc592b80222814f11c5db8e56 Mon Sep 17 00:00:00 2001 From: Chad Voegele Date: Wed, 5 Aug 2026 16:41:51 +0000 Subject: [PATCH] Fix false tied-weight detection during export Signed-off-by: Chad Voegele --- modelopt/torch/export/hf_export_handlers.py | 59 +++++++-- modelopt/torch/export/moe_utils.py | 78 +----------- modelopt/torch/export/quant_utils.py | 9 +- modelopt/torch/export/registry.py | 50 ++++++-- modelopt/torch/export/unified_export_hf.py | 112 +++++++++--------- .../unit/torch/export/test_export_registry.py | 9 +- .../torch/export/test_unified_export_hf.py | 106 +++++++++++------ .../plugins/test_fused_experts.py | 64 +++------- 8 files changed, 233 insertions(+), 254 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 800b51daca9..1a7c8169f83 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/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) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 787e173959e..f760269227a 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/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 - ``(.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. diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index ab2ef0d9029..9a1be9eecb9 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/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]}'. " diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 260cb32eea3..4acce2f859e 100644 --- a/modelopt/torch/export/registry.py +++ b/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] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index ecb69a3f906..28614d4903a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -569,25 +569,8 @@ def _export_quantized_weight( sub_module: nn.Module, dtype: torch.dtype, weight_name: str = "weight", - _tied_cache: dict[int, nn.Module] | None = None, ): - """For the given weight attr of the sub_module, export the quantization info of it. - - The export includes converting weight tensor to correct quantized values and quantized dtype, - and registering scaling factors. - - Tied-weight dedup is opt-in via ``_tied_cache``: the setattr below replaces - ``.weight`` with a fresh ``nn.Parameter`` wrapping packed bytes, breaking - any HF-level tie. When the caller passes a ``_tied_cache`` dict (keyed by - the pre-pack ``weight.data_ptr()``), the alias step at the end re-points - ``weight`` / ``weight_scale`` / ``weight_scale_2`` at a previously-processed - module sharing the same source memory so the downstream data_ptr dedup can - collapse them. The cache is owned by the caller (typically - ``_export_transformers_checkpoint``) and scoped to one export invocation; - when ``_tied_cache`` is ``None`` (the default) the alias step is skipped - entirely. Uses memory identity only — no ``_tied_weights_keys`` lookup, - no-op for non-tied modules. - """ + """Export a module weight and its quantization metadata.""" quantization_format = get_quantization_format(sub_module) if quantization_format == QUANTIZATION_NONE: return @@ -595,13 +578,6 @@ def _export_quantized_weight( block_size = get_weight_block_size(sub_module, weight_name) quantizer_attrs = quantizer_attr_names(weight_name) weight: nn.Parameter = getattr(sub_module, weight_name) - - # Capture source identity BEFORE any tensor-creating operation below. - # For HF-tied weights this matches across all modules sharing the - # underlying Parameter; the cache lookup at the end of this function - # uses it to detect ties whose Python identity is about to be broken - # by the setattr on `weight_name` further down. - _tied_source_data_ptr = weight.data_ptr() weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr( sub_module, quantizer_attrs.weight_quantizer ) @@ -812,32 +788,6 @@ def _export_quantized_weight( if weight_scale is not None: sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale) - # Tied-weight dedup: if a previously-processed module shared the same - # source weight memory, alias the packed weight + scale buffers so the - # downstream data_ptr dedup in postprocess_state_dict can collapse them. - # input_scale is safe to alias because sync_tied_input_amax (earlier in - # this export) already max-merged the per-side amaxes. Gated on the - # caller-owned _tied_cache so the dedup state is scoped to one export. - if _tied_cache is not None: - _prior = _tied_cache.get(_tied_source_data_ptr) - if _prior is not None and _prior is not sub_module: - if hasattr(_prior, weight_name): - setattr(sub_module, weight_name, getattr(_prior, weight_name)) - for _attr in ( - quantizer_attrs.weight_scale, - quantizer_attrs.weight_scale_2, - quantizer_attrs.input_scale, - ): - if not hasattr(_prior, _attr): - continue - if _attr in sub_module._buffers: - del sub_module._buffers[_attr] - elif hasattr(sub_module, _attr): - delattr(sub_module, _attr) - sub_module.register_buffer(_attr, getattr(_prior, _attr)) - else: - _tied_cache[_tied_source_data_ptr] = sub_module - torch.cuda.empty_cache() @@ -845,7 +795,8 @@ def _process_quantized_modules( model: nn.Module, dtype: torch.dtype, is_modelopt_qlora: bool = False, -) -> None: + ctx: ExportContext | None = None, +) -> ExportContext: """Process all quantized modules in model, export weights in-place. This function iterates through all modules in the model and invokes the first matching @@ -856,11 +807,13 @@ def _process_quantized_modules( dtype: The data type for weight conversion. is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. If True, modules with base_layer attribute are skipped. + ctx: Optional context prepared before export-time model mutation. """ - # 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) + if ctx is None: + ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + ctx.weight_formats = { + name: get_quantization_format(module) for name, (module, _) in ctx.weight_locations.items() + } fsdp_module_to_reshard = None for name, sub_module in model.named_modules(): @@ -889,6 +842,45 @@ def _process_quantized_modules( if handler is not None: handler(name, sub_module, ctx) + return ctx + + +def _remove_skipped_duplicate_weights( + state_dict: dict[str, torch.Tensor], ctx: ExportContext +) -> dict[str, torch.Tensor]: + """Remove source weights and quantizer state omitted from packing as tied aliases.""" + if not ctx.skipped_weight_names: + return state_dict + + remove_exact: set[str] = set() + remove_prefixes: set[str] = set() + for full_name in ctx.skipped_weight_names: + _, weight_name = ctx.weight_locations[full_name] + module_name = full_name[: -(len(weight_name) + 1)] if "." in full_name else "" + prefix = f"{module_name}." if module_name else "" + attrs = quantizer_attr_names(weight_name) + remove_exact.update( + { + full_name, + prefix + attrs.weight_scale, + prefix + attrs.weight_scale_2, + prefix + attrs.input_scale, + } + ) + remove_prefixes.update( + { + prefix + attrs.weight_quantizer + ".", + prefix + attrs.input_quantizer + ".", + prefix + attrs.output_quantizer + ".", + } + ) + + return { + name: tensor + for name, tensor in state_dict.items() + if name not in remove_exact and not any(name.startswith(p) for p in remove_prefixes) + } + def _export_transformers_checkpoint( model: nn.Module, @@ -918,7 +910,7 @@ def _export_transformers_checkpoint( # 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) + export_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"): handler = PrepareMoEInputsRegistry.match(sub_module.experts) @@ -928,7 +920,7 @@ def _export_transformers_checkpoint( 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) + handler(name, sub_module, export_ctx) # Resmooth and requantize fused layers # TODO: Handle mixed precision @@ -978,7 +970,7 @@ def _export_transformers_checkpoint( ) # Process all quantized modules and export weights - _process_quantized_modules(model, dtype, is_modelopt_qlora) + _process_quantized_modules(model, dtype, is_modelopt_qlora, export_ctx) # Reconstruct fused MoELinear: per-expert _QuantLinear weights → original 3D format from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear @@ -995,6 +987,8 @@ def _export_transformers_checkpoint( # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict). quantized_state_dict = model.state_dict() + quantized_state_dict = _remove_skipped_duplicate_weights(quantized_state_dict, export_ctx) + # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization. kv_cache_max_bound = 448 kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] diff --git a/tests/unit/torch/export/test_export_registry.py b/tests/unit/torch/export/test_export_registry.py index 67647f9c31f..824892d7bc7 100644 --- a/tests/unit/torch/export/test_export_registry.py +++ b/tests/unit/torch/export/test_export_registry.py @@ -303,10 +303,11 @@ def test_process_quantized_modules_exports_via_registry(): assert weight.dtype == torch.float8_e4m3fn -def test_export_context_caches_are_per_instance(): +def test_export_context_duplicate_state_is_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 == {} + ctx_a.duplicate_weight_map["alias.weight"] = "canonical.weight" + ctx_a.mark_skipped("alias.weight") + assert ctx_b.duplicate_weight_map == {} + assert ctx_b.skipped_weight_names == set() diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 118331ce3d9..3360e40f410 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -28,8 +28,16 @@ _collect_canonical_tied_patterns, _reorder_canonical_first, ) -from modelopt.torch.export.quant_utils import fuse_prequant_layernorm, sync_tied_input_amax -from modelopt.torch.export.unified_export_hf import _export_quantized_weight +from modelopt.torch.export.quant_utils import ( + fuse_prequant_layernorm, + postprocess_state_dict, + sync_tied_input_amax, +) +from modelopt.torch.export.registry import ExportContext +from modelopt.torch.export.unified_export_hf import ( + _process_quantized_modules, + _remove_skipped_duplicate_weights, +) from modelopt.torch.quantization.nn import TensorQuantizer @@ -127,60 +135,86 @@ def forward_loop(m): mtq.quantize(parent, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop) -def test_export_quantized_weight_aliases_packed_weight_for_tied_linears(): - """Tied Linears share data_ptr for packed .weight and scale buffers after export.""" +def test_export_context_maps_tied_alias_to_canonical_weight_name(): enc, dec = make_tied_linear_pair() parent = wrap_in_parent_with_tied_keys(enc, dec) - _calibrate_through_both_children(parent) - # Per-call dedup cache (the production pattern: caller owns the cache, scoped - # to one export invocation). Threaded through both sides of the tied pair so - # the alias step at the end of _export_quantized_weight catches the dedup. - tied_cache: dict = {} - _export_quantized_weight(enc, torch.float16, "weight", _tied_cache=tied_cache) - _export_quantized_weight(dec, torch.float16, "weight", _tied_cache=tied_cache) + ctx = ExportContext(parent, torch.float16) + + assert ctx.duplicate_weight_map == {"encoder.weight": "decoder.weight"} + + +def test_export_context_duplicate_map_survives_weight_replacement(): + enc, dec = make_tied_linear_pair() + parent = wrap_in_parent_with_tied_keys(enc, dec) + ctx = ExportContext(parent, torch.float16) + + enc.weight = torch.nn.Parameter(enc.weight.detach().clone()) + dec.weight = torch.nn.Parameter(dec.weight.detach().clone()) - assert enc.weight.data_ptr() == dec.weight.data_ptr() - for scale_attr in ("weight_scale", "weight_scale_2"): - if hasattr(enc, scale_attr) and hasattr(dec, scale_attr): - assert getattr(enc, scale_attr).data_ptr() == getattr(dec, scale_attr).data_ptr() + assert ctx.duplicate_of("encoder.weight") == "decoder.weight" -def test_export_quantized_weight_no_alias_for_untied_linears(): - """Untied Linears keep independent data_ptrs after export — no false-positive aliasing.""" +def test_export_context_has_no_duplicates_for_untied_linears(): parent = torch.nn.Module() parent.encoder = torch.nn.Linear(16, 32, bias=False) parent.decoder = torch.nn.Linear(16, 32, bias=False) - assert parent.encoder.weight.data_ptr() != parent.decoder.weight.data_ptr() - _calibrate_through_both_children(parent) - # Same fresh cache shape as the positive case — confirms that even with - # dedup enabled, untied modules with distinct source data_ptrs do not get - # falsely aliased. - tied_cache: dict = {} - _export_quantized_weight(parent.encoder, torch.float16, "weight", _tied_cache=tied_cache) - _export_quantized_weight(parent.decoder, torch.float16, "weight", _tied_cache=tied_cache) + ctx = ExportContext(parent, torch.float16) - assert parent.encoder.weight.data_ptr() != parent.decoder.weight.data_ptr() + assert ctx.duplicate_weight_map == {} -def test_export_quantized_weight_skips_alias_when_one_tied_side_is_unquantized(): - """Unquantized side early-returns; its .weight stays at the original shared Parameter.""" +def test_process_quantized_modules_skips_and_filters_tied_alias(): enc, dec = make_tied_linear_pair() parent = wrap_in_parent_with_tied_keys(enc, dec) - original_shared_data_ptr = enc.weight.data_ptr() + _calibrate_through_both_children(parent) + original_weight = enc.weight + + ctx = _process_quantized_modules(parent, torch.float16) + state_dict = _remove_skipped_duplicate_weights(parent.state_dict(), ctx) + assert parent.encoder.weight is original_weight + assert parent.decoder.weight is not original_weight + assert ctx.skipped_weight_names == {"encoder.weight"} + assert "encoder.weight" not in state_dict + assert "decoder.weight" in state_dict + assert "decoder.weight_scale" in state_dict + + +def test_process_quantized_modules_keeps_differently_quantized_tied_weights(): + enc, dec = make_tied_linear_pair() + parent = wrap_in_parent_with_tied_keys(enc, dec) _calibrate_through_both_children(parent) - # is_enabled is a read-only property; .disable() is the canonical bypass. dec.weight_quantizer.disable() - tied_cache: dict = {} - _export_quantized_weight(enc, torch.float16, "weight", _tied_cache=tied_cache) - _export_quantized_weight(dec, torch.float16, "weight", _tied_cache=tied_cache) + ctx = _process_quantized_modules(parent, torch.float16) + state_dict = _remove_skipped_duplicate_weights(parent.state_dict(), ctx) + + assert ctx.skipped_weight_names == set() + assert "encoder.weight" in state_dict + assert "decoder.weight" in state_dict + + +def test_postprocess_state_dict_preserves_tensors_with_different_byte_ranges(): + storage = torch.arange(4) + state_dict = {"short": storage[:2], "long": storage} + assert state_dict["short"].data_ptr() == state_dict["long"].data_ptr() + + processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None) + + assert set(processed) == set(state_dict) + + +def test_postprocess_state_dict_preserves_zero_pointer_tensors(): + state_dict = { + "first": torch.empty(4, device="meta"), + "second": torch.empty(4, device="meta"), + } + + processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None) - assert enc.weight.data_ptr() != original_shared_data_ptr # encoder got fresh packed - assert dec.weight.data_ptr() == original_shared_data_ptr # decoder untouched - assert enc.weight.data_ptr() != dec.weight.data_ptr() + assert set(processed) == set(state_dict) def _linear_with_input_quantizer(): diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index c435b3698be..607255dac8c 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -28,6 +28,7 @@ import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module from modelopt.torch.export.moe_utils import _export_fused_experts from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format +from modelopt.torch.export.unified_export_hf import _process_quantized_modules from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name from modelopt.torch.quantization.model_calib import local_hessian_calibrate @@ -685,73 +686,36 @@ def _cleanup_registry(mod_type): if QuantModuleRegistry.get(mod_type) is not None: QuantModuleRegistry.unregister(mod_type) - def test_per_expert_buffers_share_data_ptr_for_tied_fused_experts(self): - """Two tied FusedExperts modules: every per-expert .weight + scale buffer shares data_ptr.""" + def test_tied_fused_experts_expand_only_canonical_subtree(self): parent = _build_two_moe_blocks(tie=True) expert_type = type(parent.encoder.experts) self._cleanup_registry(expert_type) try: _calibrate_two_moe_blocks(parent) - # Per-call dedup caches threaded through both export calls; int keys - # for per-expert wrapper dedup, tuple keys for module-level dedup. - tied_cache: dict = {} - moe_tied_cache: dict = {} - _export_fused_experts( - parent.encoder.experts, - torch.float16, - _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, - ) - _export_fused_experts( - parent.decoder.experts, - torch.float16, - _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, - ) + ctx = _process_quantized_modules(parent, torch.float16) - for idx in range(NUM_EXPERTS): - enc_expert = getattr(parent.encoder.experts, str(idx)) - dec_expert = getattr(parent.decoder.experts, str(idx)) - for proj_name in ("gate_proj", "up_proj", "down_proj"): - enc_proj = getattr(enc_expert, proj_name) - dec_proj = getattr(dec_expert, proj_name) - assert enc_proj.weight.data_ptr() == dec_proj.weight.data_ptr() - for scale_attr in ("weight_scale", "weight_scale_2"): - if hasattr(enc_proj, scale_attr) and hasattr(dec_proj, scale_attr): - assert ( - getattr(enc_proj, scale_attr).data_ptr() - == getattr(dec_proj, scale_attr).data_ptr() - ) + assert ctx.duplicate_of("decoder.experts.gate_up_proj") == ( + "encoder.experts.gate_up_proj" + ) + assert ctx.duplicate_of("decoder.experts.down_proj") == "encoder.experts.down_proj" + assert hasattr(parent.encoder.experts, "0") + assert not hasattr(parent.decoder.experts, "0") + assert not hasattr(parent.decoder.experts, "gate_up_proj") + assert not hasattr(parent.decoder.experts, "down_proj") finally: self._cleanup_registry(expert_type) - def test_per_expert_buffers_have_independent_data_ptrs_for_untied_fused_experts(self): - """Two untied FusedExperts modules: per-expert buffers stay independent (no false-positive alias).""" + def test_untied_fused_experts_expand_independently(self): parent = _build_two_moe_blocks(tie=False) expert_type = type(parent.encoder.experts) self._cleanup_registry(expert_type) try: _calibrate_two_moe_blocks(parent) - # Same fresh caches as the positive case — confirms that even with - # dedup enabled, untied modules with distinct source data_ptrs do - # not get falsely aliased. - tied_cache: dict = {} - moe_tied_cache: dict = {} - _export_fused_experts( - parent.encoder.experts, - torch.float16, - _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, - ) - _export_fused_experts( - parent.decoder.experts, - torch.float16, - _moe_tied_cache=moe_tied_cache, - _tied_cache=tied_cache, - ) + ctx = _process_quantized_modules(parent, torch.float16) + assert ctx.duplicate_weight_map == {} for idx in range(NUM_EXPERTS): enc_expert = getattr(parent.encoder.experts, str(idx)) dec_expert = getattr(parent.decoder.experts, str(idx))