Add support for dLLM encoder-decoder models (DiffusionGemma) [tied-weight PTQ export support ] [OmniML-5167]

CHANGELOG.rst

@@ -20,6 +20,13 @@ Changelog
   - **GatedDeltaNet** (linear attention) and **gated attention** (``attention_output_gate``), such as Qwen3.5 (hybrid GatedDeltaNet + gated-attention) language models, including MoE variants — attention / linear-attention heads are not pruned.
   - **Multi-Latent Attention (MLA)**, such as DeepSeek — MLA latent ranks are not pruned.
   - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned.
+- Add dLLM (tied-weight PTQ and HF-checkpoint export) support for diffusion-based encoder-decoder LLMs (e.g. DiffusionGemma) whose encoder/decoder stacks share parameters via HF ``_tied_weights_keys``.
+
+  - **Deduplicate the modules shared at source** in the quantized export step: ``_export_quantized_weight`` and ``_export_fused_experts`` now alias bit-identical packed ``weight`` / ``weight_scale`` / ``weight_scale_2`` buffers across modules sharing a source weight ``data_ptr()`` so the downstream ``postprocess_state_dict`` dedup catches them (~42% storage reduction on ``nvfp4_experts_only`` for tied 26B MoE checkpoints).
+  - New ``sync_tied_input_amax`` helper max-merges per-side ``input_quantizer.amax`` across tied modules before export so single-backbone consumers that load one ``input_scale`` per parameter don't clip either side.
+  - The exported state_dict is also **reordered (decoder keys win instead of encoder)** so canonical-side keys per HF's ``_tied_weights_keys`` declaration win the data_ptr dedup; gated to the DiffusionGemma model class in ``_reorder_canonical_first``, no-op for every other model.
+  - New DiffusionGemma model-specific recipe under ``modelopt_recipes/huggingface/diffusion_gemma/ptq/`` (``nvfp4_experts_only.yaml`` + its ``disabled_quantizers.yaml`` unit) adds the ``*self_conditioning*`` exclude on top of the standard default, leaving the shared ``default_disabled_quantizers`` unit clean for non-diffusion models — pattern matches the existing ``phi4mm`` / ``nemotron_vl`` model-specific recipes.
+  - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change.
 
 0.45 (2026-07-02)
 ^^^^^^^^^^^^^^^^^

examples/llm_ptq/example_utils.py

@@ -786,7 +786,13 @@ def is_model_on_gpu(model) -> bool:
 
 
 def is_enc_dec(model_type) -> bool:
-    """Return if the model is a encoder-decoder model."""
+    """Return whether the model_type uses encoder-decoder-style preview decode.
+
+    Controls whether ``hf_ptq.py`` slices off the prompt prefix from
+    ``.generate()`` output. ``diffusion_gemma`` is structurally encoder-decoder
+    but returns prompt+canvas concatenated, so it stays OFF this list (AR-style
+    decode applies).
+    """
     return model_type in ["t5", "bart", "whisper"]
 
 

examples/llm_ptq/hf_ptq.py

@@ -941,6 +941,11 @@ def input_decode(input_ids):
             raise ValueError("The processor or tokenizer must be set")
 
     def output_decode(generated_ids, input_shape):
+        # Some `.generate()` returns a ModelOutput dataclass (e.g. DiffusionGemma);
+        # unwrap to the token tensor so downstream slicing works uniformly.
+        if hasattr(generated_ids, "sequences"):
+            generated_ids = generated_ids.sequences
+
         if is_enc_dec(model_type):
             if processor is not None and isinstance(processor, WhisperProcessor):
                 return processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

modelopt/torch/export/model_utils.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 """Utility functions for model type detection and classification."""
 
+import re
+
 import torch.nn as nn
 
 MODEL_NAME_TO_TYPE = {
@@ -33,6 +35,9 @@
     "Qwen3Next": "qwen3next",
     "QWen": "qwen",
     "RecurrentGemma": "recurrentgemma",
+    # DiffusionGemma must come before "Gemma" — get_model_type substring-matches
+    # in order, and "gemma" is a substring of "diffusiongemma".
+    "DiffusionGemma": "diffusion_gemma",
     "Gemma3": "gemma3",
     "Gemma2": "gemma2",
     "Gemma": "gemma",
@@ -157,3 +162,84 @@ def get_language_model_from_vl(model) -> list[nn.Module] | None:
 
     # Pattern 4: No language_model found
     return None
+
+
+def _collect_canonical_tied_patterns(
+    model: nn.Module,
+) -> tuple[list[re.Pattern], list[str]]:
+    """Walk the model and collect canonical-side tied-weight matchers.
+
+    Patterns are submodule-prefixed regexes from each module's
+    ``_tied_weights_keys`` dict-style declaration (the prefix matters
+    for nested models where the dict lives on an inner submodule).
+    Side substrings are dot-separated tokens that appear only on the
+    canonical side of those declarations — needed because modelopt's
+    per-expert unpacking creates post-export keys (e.g.
+    ``…experts.Y.gate_proj.input_scale``) that HF's regexes never knew
+    about. List-style (legacy) declarations are skipped.
+    """
+    patterns: list[re.Pattern] = []
+    alias_token_set: set[str] = set()
+    canonical_token_set: set[str] = set()
+
+    def _tokens(s: str) -> set[str]:
+        """Identifiers in a regex string, with regex specials as separators."""
+        return {tok for tok in re.split(r"[^A-Za-z0-9_]+", s) if tok}
+
+    for name, submodule in model.named_modules():
+        tied = getattr(submodule, "_tied_weights_keys", None)
+        if not isinstance(tied, dict) or not tied:
+            continue
+        prefix = f"{name}." if name else ""
+        for alias_pat, canonical_pat in tied.items():
+            patterns.append(re.compile(prefix + canonical_pat))
+            alias_token_set.update(_tokens(prefix + alias_pat))
+            canonical_token_set.update(_tokens(prefix + canonical_pat))
+
+    # Tokens unique to the canonical side become substring matchers.
+    side_substrings = sorted(canonical_token_set - alias_token_set)
+    return patterns, side_substrings
+
+
+def _reorder_canonical_first(state_dict: dict, model: nn.Module) -> dict:
+    r"""Reorder ``state_dict`` so canonical-side tied keys iterate first.
+
+    Lets the downstream first-wins data_ptr dedup keep canonical names.
+    Uses both regex patterns and substring matchers from
+    :func:`_collect_canonical_tied_patterns`. Gated on the model class
+    name to scope the reorder to DiffusionGemma; other tied
+    encoder-decoder models that ship dict-style ``_tied_weights_keys``
+    can be added to the allowlist here. Mirrors the ``model_type``
+    dispatch used for the Whisper and Nemotron-VL branches elsewhere
+    in ``unified_export_hf.py``.
+    """
+    model_type = type(model).__name__.lower()
+    if "diffusiongemma" not in model_type and "diffusion_gemma" not in model_type:
+        return state_dict
+
+    canonical_patterns, side_substrings = _collect_canonical_tied_patterns(model)
+    if not canonical_patterns and not side_substrings:
+        return state_dict
+
+    def _has_side_substring(key: str) -> bool:
+        # Require the token to appear as a proper dot-separated path
+        # component, not just as a substring of an unrelated identifier.
+        for tok in side_substrings:
+            if (
+                f".{tok}." in key
+                or key.startswith(f"{tok}.")
+                or key.endswith(f".{tok}")
+                or key == tok
+            ):
+                return True
+        return False
+
+    head: dict = {}
+    tail: dict = {}
+    for k, v in state_dict.items():
+        if any(p.search(k) for p in canonical_patterns) or _has_side_substring(k):
+            head[k] = v
+        else:
+            tail[k] = v
+    head.update(tail)
+    return head

modelopt/torch/export/moe_utils.py

@@ -23,7 +23,60 @@
 import torch.nn as nn
 
 
-def _export_fused_experts(module: nn.Module, dtype: torch.dtype) -> None:
+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.
+
+    Called once the per-expert subtree exists (either via the fast-path
+    aliases or via the full unpack/pack path) so the redundant fused form
+    doesn't appear in the exported state_dict alongside the per-expert form.
+    """
+    for attr in (
+        "gate_up_proj",
+        "down_proj",
+        "gate_up_proj_weight_quantizers",
+        "gate_up_proj_input_quantizer",
+        "down_proj_weight_quantizers",
+        "down_proj_input_quantizer",
+    ):
+        if hasattr(module, attr):
+            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:
     """Split fused MoE expert weights and export per-expert quantization scales.
 
     Works with any module wrapped by ``_QuantFusedExperts`` — i.e. any HF
@@ -42,13 +95,43 @@ def _export_fused_experts(module: nn.Module, dtype: torch.dtype) -> None:
            {E}.gate_proj.weight, {E}.gate_proj.weight_scale, ...
            {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
+    ``(gate_up_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
 
     n = module.num_experts
     expert_dim = _get_fused_expert_intermediate_dim(module)
 
+    # Capture source tensor identities BEFORE unpacking (the source
+    # attrs are deleted at the end of this function).
+    _source_key = (module.gate_up_proj.data_ptr(), module.down_proj.data_ptr())
+
+    # Tied-experts fast path: if this exact (gate_up, 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.
     gate_up_input_q = module.gate_up_proj_input_quantizer
     down_input_q = module.down_proj_input_quantizer
@@ -154,7 +237,7 @@ def _export_fused_experts(module: nn.Module, dtype: torch.dtype) -> None:
             wrapper.weight_quantizer = w_quantizer
             wrapper.input_quantizer = i_quantizer
 
-            _export_quantized_weight(wrapper, dtype)
+            _export_quantized_weight(wrapper, dtype, _tied_cache=_tied_cache)
 
             proj = nn.Module()
             proj.weight = wrapper.weight
@@ -167,16 +250,14 @@ def _export_fused_experts(module: nn.Module, dtype: torch.dtype) -> None:
         module.add_module(str(idx), expert)
 
     # 4. Remove fused params and quantizer lists — replaced by per-expert submodules
-    for attr in (
-        "gate_up_proj",
-        "down_proj",
-        "gate_up_proj_weight_quantizers",
-        "gate_up_proj_input_quantizer",
-        "down_proj_weight_quantizers",
-        "down_proj_input_quantizer",
-    ):
-        if hasattr(module, attr):
-            delattr(module, attr)
+    _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):

modelopt/torch/export/quant_utils.py

@@ -1479,3 +1479,79 @@ def has_quantized_modules(model: nn.Module) -> bool:
         get_quantization_format(sub_module) != QUANTIZATION_NONE
         for _, sub_module in model.named_modules()
     )
+
+
+def sync_tied_input_amax(model: nn.Module) -> int:
+    """Max-merge input_quantizer amaxes across modules sharing a weight ``data_ptr``.
+
+    Mutates ``model`` in place: overwrites the ``.amax`` buffer on every
+    affected ``input_quantizer`` with the per-group maximum. Intended to
+    run as part of an export pipeline that already replaces weights with
+    packed bytes downstream — i.e. the model is not expected to be reused
+    after this helper runs.
+
+    Closes the loop on ``input_scale`` for HF-tied modules whose forward
+    paths see different activation distributions (encoder vs decoder in
+    YOCO-style models). Must run BEFORE per-module export so the merged
+    amax flows into ``input_scale`` derivation. Handles both dense
+    Linears (keyed by ``weight.data_ptr()``) and fused MoE (keyed by
+    ``(gate_up_proj, down_proj)`` data_ptr tuple). Returns the number of
+    tied groups merged.
+    """
+    from collections import defaultdict
+
+    by_dp: dict = defaultdict(list)
+    for _, m in model.named_modules():
+        # Fused MoE: 3-D source tensors with shared input quantizers
+        if (
+            hasattr(m, "gate_up_proj_input_quantizer")
+            and hasattr(m, "gate_up_proj")
+            and hasattr(m, "down_proj")
+            and m.gate_up_proj.dim() == 3
+        ):
+            key = ("moe", m.gate_up_proj.data_ptr(), m.down_proj.data_ptr())
+            by_dp[key].append(m)
+        # Dense quantized Linear with an input_quantizer
+        elif (
+            hasattr(m, "input_quantizer")
+            and hasattr(m, "weight")
+            and isinstance(m.weight, torch.nn.Parameter)
+        ):
+            by_dp[("dense", m.weight.data_ptr())].append(m)
+
+    def _merge(quantizers: list) -> bool:
+        """Max-merge amaxes across the quantizer list. Returns True on merge."""
+        valid = [
+            q
+            for q in quantizers
+            if q is not None
+            and getattr(q, "is_enabled", False)
+            and getattr(q, "_amax", None) is not None
+            and not q._amax.is_meta
+        ]
+        if len(valid) < 2:
+            return False
+        # Require scalar (per-tensor) amax — matches preprocess_linear_fusion.
+        if any(q._amax.numel() != 1 for q in valid):
+            warn(
+                "sync_tied_input_amax: non-scalar input_quantizer amax encountered "
+                "in a tied group; skipping. Only per-tensor input quantizers are "
+                "supported for tied-modules merging."
+            )
+            return False
+        merged = torch.max(torch.stack([q.amax for q in valid]))
+        for q in valid:
+            q.amax = merged.clone()
+        return True
+
+    synced = 0
+    for key, modules in by_dp.items():
+        if len(modules) < 2:
+            continue
+        if key[0] == "moe":
+            for q_name in ("gate_up_proj_input_quantizer", "down_proj_input_quantizer"):
+                if _merge([getattr(m, q_name, None) for m in modules]):
+                    synced += 1
+        elif _merge([m.input_quantizer for m in modules]):
+            synced += 1
+    return synced