Add quantization support for DiffusionGemma

auto_round/calibration/llm.py

@@ -38,6 +38,7 @@
     hook_ngram_embeddings_on_cpu,
     is_quantized_input_module,
     mv_module_from_gpu,
+    safe_tie_weights,
     to_device,
     to_dtype,
 )
@@ -135,7 +136,7 @@ def collect(self, block_names, nsamples, layer_names=None, last_cache_name=None)
                             no_split_module_classes=no_split_modules,
                         )
                         if hasattr(c.model_context.model, "tie_weights"):
-                            c.model_context.model.tie_weights()
+                            safe_tie_weights(c.model_context.model)
                         device_map = infer_auto_device_map(
                             c.model_context.model,
                             max_memory=new_max_memory,

auto_round/inference/convert_model.py

@@ -47,6 +47,7 @@
     is_transformers_version_greater_or_equal_5,
     set_module,
 )
+from auto_round.utils.model import prune_stale_tied_weights_keys
 
 supported_devices = ("cpu", "hpu", "xpu", "cuda", "mps")
 
@@ -367,8 +368,10 @@ def get_layer_config(model, quantization_config):
         modules_in_block_to_quantize = flatten_list(
             quantization_config.modules_in_block_to_quantize
         )  # Flatten the list
+        # Pre-compile patterns once instead of recompiling them for every layer name.
+        compiled_modules_in_block = [re.compile(n) for n in modules_in_block_to_quantize]
         for layer_name in layer_names:
-            if not any([re.search(re.compile(n), layer_name) is not None for n in modules_in_block_to_quantize]):
+            if not any(pattern.search(layer_name) is not None for pattern in compiled_modules_in_block):
                 extra_config[layer_name] = {"bits": 16}  # Default to 16-bit for unquantized layers
 
     # Expand GPTQ 'dynamic' config (regex-based)
@@ -383,8 +386,9 @@ def get_layer_config(model, quantization_config):
             model=model,
         )
 
-    # AWQ format: exclude specified modules
-    extra_config = skip_not_convert_modules(model, quantization_config, layer_names, extra_config)
+    # AWQ format: exclude specified modules.
+    if "awq" in (getattr(quantization_config, "quant_method", "") or "").lower():
+        extra_config = skip_not_convert_modules(model, quantization_config, layer_names, extra_config)
 
     # Expand auto_round regex configs (regex-based)
     extra_config = _expand_regex_config(
@@ -873,6 +877,8 @@ def convert_hf_model(model: nn.Module, target_device: str = "cpu") -> tuple[nn.M
     layer_configs = get_layer_config(model, quantization_config)
     used_backends = _replace_by_quant_layers(model, layer_configs, backend, target_device, packing_format)
 
+    prune_stale_tied_weights_keys(model)
+
     # Apply rotation hooks (hadamard, spinquant, quarot, etc.) via unified dispatch.
     _has_rotation = getattr(quantization_config, "rotation_config", None) or getattr(
         quantization_config, "spinquant_config", None

auto_round/special_model_handler.py

@@ -20,6 +20,7 @@
 from auto_round.formats import OutputFormat
 from auto_round.modeling.fused_moe.replace_modules import apply_replacements, release_original_module_
 from auto_round.utils import is_moe_model_via_config, logger
+from auto_round.utils.model import prune_stale_tied_weights_keys
 
 mllms_with_limited_bs = (
     "llava",
@@ -401,6 +402,7 @@ def update_module(
     if cleanup_original:
         release_original_module_(model)
 
+    prune_stale_tied_weights_keys(model)
     return model
 
 
@@ -1138,6 +1140,16 @@ def get_bagel_ignore_layers(model) -> list[str]:
     ],
 )
 
+# diffusion_gemma
+register_ignore_layers(
+    matchers=[
+        ArchitectureMatcher(r"DiffusionGemma", mode="in"),
+    ],
+    ignore_layers=[
+        "router.proj",
+    ],
+)
+
 
 def get_predefined_ignore_layers(model: torch.nn.Module) -> list[str]:
     layers = []
@@ -1270,7 +1282,10 @@ def _nextstep_pipeline_fn(pipe, prompts, guidance_scale=7.5, num_inference_steps
     return pipe, model
 
 
-_PRE_DEFINED_FIXED_ATTR = {"gemma4_unified": {"has_variable_block_shape": True}}
+_PRE_DEFINED_FIXED_ATTR = {
+    "gemma4_unified": {"has_variable_block_shape": True},
+    "diffusion_gemma": {"has_variable_block_shape": True},
+}
 
 
 def get_predefined_fixed_attr(model: torch.nn.Module) -> dict | None:

auto_round/utils/common.py

@@ -28,6 +28,12 @@
 from auto_round.export.export_to_gguf.config import GGUF_CONFIG
 from auto_round.logger import logger
 
+try:
+    from transformers.models.diffusion_gemma.modeling_diffusion_gemma import DiffusionGemmaModel
+except ImportError:
+    # diffusion_gemma is only available in transformers >= 5.11.0.
+    DiffusionGemmaModel = None  # type: ignore[assignment]
+
 
 def download_audiocaps_csv():
     """Download AudioCaps train.csv and return the local cache path.
@@ -238,6 +244,37 @@ def _tensor_get_dtype(self):
     torch.Tensor.get_dtype = _tensor_get_dtype
 
 
+def _patch_diffusion_gemma_tied_weights():
+    """Install a ``__init__`` hook on ``DiffusionGemmaModel`` to prune stale tied-weights keys.
+
+    AutoRound unfuses ``DiffusionGemmaTextExperts.gate_up_proj`` /
+    ``down_proj`` (fused 3D ``nn.Parameter``) into per-expert
+    ``gate_proj / up_proj / down_proj`` ``nn.Linear`` modules at quantize
+    time. The encoder/decoder weight-tying map declared in
+    ``DiffusionGemmaModel._tied_weights_keys`` still references the original
+    fused parameter names (``gate_up_proj``, ``down_proj``).
+    The fix is to prune the stale patterns at the moment the model is constructed.
+    """
+    if DiffusionGemmaModel is None:
+        # diffusion_gemma is unavailable (transformers < 5.11.0); nothing to patch.
+        return
+    if getattr(DiffusionGemmaModel, "_ar_tied_prune_patched", False):
+        return
+    original_init = DiffusionGemmaModel.__init__
+
+    from auto_round.utils.model import prune_stale_tied_weights_keys
+
+    def _patched_init(self, *args, **kwargs):
+        original_init(self, *args, **kwargs)
+        try:
+            prune_stale_tied_weights_keys(self)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning(f"[DiffusionGemma] prune_stale_tied_weights_keys during __init__ failed: {exc}")
+
+    DiffusionGemmaModel.__init__ = _patched_init
+    DiffusionGemmaModel._ar_tied_prune_patched = True
+
+
 def _patch_default_rope_init():
     """Restore legacy ``rope_type='default'`` support for older remote-code models.
 
@@ -351,6 +388,8 @@ def monkey_patch_transformers():
         # transformers 5.3.0 calls tensor.get_dtype() on plain torch.Tensor objects
         # while loading pre-quantized checkpoints.
         _patch_tensor_get_dtype_for_prequantized_loading()
+    if parsed_version >= version.parse("5.11.0"):
+        _patch_diffusion_gemma_tied_weights()
     _patch_default_rope_init()
     _patch_rotary_embedding_init_for_legacy_remote_code()
     if parsed_version >= version.parse("4.56.0"):

auto_round/utils/model.py

@@ -24,6 +24,7 @@
 import torch
 import transformers
 from packaging import version
+from transformers import PreTrainedModel
 
 from auto_round import envs
 from auto_round.export.export_to_gguf.config import ModelType
@@ -76,6 +77,100 @@ def resolve_model_type(model):
     from auto_round.schemes import QuantizationScheme
 
 
+def prune_stale_tied_weights_keys(model: torch.nn.Module) -> int:
+    """Drop ``_tied_weights_keys`` regex patterns that no longer match any parameter.
+
+    AutoRound unfuses MoE experts before calibration, splitting fused projections such as
+    ``gate_up_proj`` ``[num_experts, 2*inter, hidden]`` into per-expert ``gate_proj`` /
+    ``up_proj`` linear weights. The model's declared ``_tied_weights_keys`` still reference
+    the original fused name (e.g. DiffusionGemma ties ``encoder...gate_up_proj`` to
+    ``decoder...gate_up_proj``), which now matches zero parameters.
+
+    Removing only the now-empty patterns lets the remaining ties be established normally.
+    The split ``gate_proj`` / ``up_proj`` weights stay tied through the generic ``*.weight`` pattern
+    that already covers every ``.weight`` parameter under the layers. Already-expanded
+    plain parameter names are kept untouched.
+
+    Args:
+        model (torch.nn.Module): model whose stale tie patterns should be removed.
+
+    Returns:
+        int: number of stale tie entries removed across all submodels.
+    """
+    common_case = re.compile(r"^[A-Za-z0-9_\.]+(weight)|(bias)$")
+    removed = 0
+
+    for _, submodule in model.named_modules(remove_duplicate=False):
+        if not isinstance(submodule, PreTrainedModel):
+            continue
+        tied = getattr(submodule, "_tied_weights_keys", None)
+        if not isinstance(tied, dict) or not tied:
+            continue
+        if not getattr(submodule.config, "tie_word_embeddings", False):
+            continue
+        if all(common_case.match(target) and common_case.match(source) for target, source in tied.items()):
+            continue
+        names = {k for k, _ in submodule.named_parameters(remove_duplicate=False)} | {
+            k for k, _ in submodule.named_buffers(remove_duplicate=False)
+        }
+        pruned = {}
+        for target, source in tied.items():
+            if common_case.match(target) and common_case.match(source):
+                pruned[target] = source
+                continue
+            source_params = [n for n in names if re.search("^" + source, n)]
+            target_params = [n for n in names if re.search("^" + target, n)]
+            if len(source_params) > 0 and len(target_params) > 0 and len(target_params) % len(source_params) == 0:
+                pruned[target] = source
+            else:
+                removed += 1
+                logger.trace(
+                    f"Removing stale tie pattern '{target}' -> '{source}' from "
+                    f"{type(submodule).__name__}: it no longer matches any parameter."
+                )
+        if len(pruned) != len(tied):
+            submodule._tied_weights_keys = pruned
+
+    for _, submodule in model.named_modules(remove_duplicate=False):
+        if not isinstance(submodule, PreTrainedModel):
+            continue
+        cached = getattr(submodule, "all_tied_weights_keys", None)
+        if not isinstance(cached, dict) or not cached:
+            continue
+        existing = {n for n, _ in submodule.named_parameters(remove_duplicate=False)} | {
+            n for n, _ in submodule.named_buffers(remove_duplicate=False)
+        }
+        cached_pruned = {
+            target: source for target, source in cached.items() if target in existing and source in existing
+        }
+        if len(cached_pruned) != len(cached):
+            removed += len(cached) - len(cached_pruned)
+            logger.trace(
+                f"Pruned {len(cached) - len(cached_pruned)} stale entries from "
+                f"{type(submodule).__name__}.all_tied_weights_keys cache "
+                f"(targets/sources no longer resolve after unfuse/quantization)."
+            )
+            submodule.all_tied_weights_keys = cached_pruned
+
+    return removed
+
+
+def safe_tie_weights(model: torch.nn.Module) -> None:
+    """Call ``model.tie_weights()`` defensively.
+
+    Args:
+        model (torch.nn.Module): model whose weights should be tied if supported.
+    """
+    tie_fn = getattr(model, "tie_weights", None)
+    if not callable(tie_fn):
+        return
+    prune_stale_tied_weights_keys(model)
+    try:
+        tie_fn()
+    except ValueError as e:
+        logger.warning(f"model.tie_weights() raised ValueError, skipping weight tying: {e}")
+
+
 def clean_module_parameter(submodule: torch.nn.Module, param_name: str) -> None:
     """This function is recommended to be used instead of module.weight = None.
     For models like `tie_word_embeddings`, setting the embedding weight to None