fix: save the config of diffusers models on export

auto_round/export/export_to_llmcompressor/export.py

@@ -18,7 +18,12 @@
 
 import torch
 
-from auto_round.export.utils import is_immediate_saving_mode, save_model, save_pretrained_artifact
+from auto_round.export.utils import (
+    is_immediate_saving_mode,
+    save_config_artifact,
+    save_model,
+    save_pretrained_artifact,
+)
 from auto_round.logger import logger
 from auto_round.utils import (
     SUPPORTED_LAYER_TYPES,
@@ -227,7 +232,7 @@ def save_quantized_as_llmcompressor(
         return model
 
     # save model.config, model.state_dict()
-    model.config.save_pretrained(output_dir)
+    save_config_artifact(model, output_dir)
 
     save_model(model, output_dir, safe_serialization=safe_serialization, immediate_saving=immediate_saving)
 

auto_round/export/utils.py

@@ -42,18 +42,42 @@ def save_pretrained_artifact(artifact, output_dir: str, artifact_name: str = "ar
     return True
 
 
-def _save_model_configs(model: nn.Module, save_dir: str) -> None:
-    if hasattr(model, "config") and model.config is not None:
-        try:
-            model.config.save_pretrained(save_dir)
-        except (KeyError, TypeError):
-            # Some third-party configs (e.g. qwen-tts) fail with use_diff=True
-            # due to missing keys in recursive_diff_dict. Fall back to full config.
-            import json
+def save_config_artifact(model: nn.Module, save_dir: str) -> None:
+    """Write ``model.config`` to ``save_dir``, for transformers and diffusers models alike.
 
+    A diffusers ``ModelMixin`` keeps its config in a ``FrozenDict``, which has no
+    ``save_pretrained``; ``ModelMixin.save_config`` is the equivalent writer.
+    """
+    config = getattr(model, "config", None)
+    if config is None:
+        return
+
+    if not hasattr(config, "save_pretrained") and hasattr(model, "save_config"):
+        model.save_config(save_dir)
+        # save_config serializes the config's own dict, so the quantization_config the
+        # exporter set on the config object afterwards has to be merged back in.
+        quantization_config = getattr(config, "quantization_config", None)
+        if quantization_config is not None:
             config_path = os.path.join(save_dir, "config.json")
+            with open(config_path, encoding="utf-8") as f:
+                config_dict = json.load(f)
+            config_dict["quantization_config"] = quantization_config
             with open(config_path, "w", encoding="utf-8") as f:
-                f.write(model.config.to_json_string(use_diff=False))
+                json.dump(config_dict, f, indent=2, sort_keys=True)
+        return
+
+    try:
+        config.save_pretrained(save_dir)
+    except (KeyError, TypeError):
+        # Some third-party configs (e.g. qwen-tts) fail with use_diff=True
+        # due to missing keys in recursive_diff_dict. Fall back to full config.
+        config_path = os.path.join(save_dir, "config.json")
+        with open(config_path, "w", encoding="utf-8") as f:
+            f.write(config.to_json_string(use_diff=False))
+
+
+def _save_model_configs(model: nn.Module, save_dir: str) -> None:
+    save_config_artifact(model, save_dir)
 
     if hasattr(model, "generation_config") and model.generation_config is not None:
         model.generation_config.save_pretrained(save_dir)

auto_round/formats.py

@@ -358,8 +358,10 @@ def save_quantized(
         if not unsupported_meta_device(model):
             model = model.to("cpu")
             model.save_pretrained(output_dir)
-        elif hasattr(model, "config") and model.config is not None:
-            model.config.save_pretrained(output_dir)
+        else:
+            from auto_round.export.utils import save_config_artifact
+
+            save_config_artifact(model, output_dir)
 
         if tokenizer is not None and hasattr(tokenizer, "save_pretrained"):
             tokenizer.save_pretrained(output_dir)

test/test_cpu/export/test_export.py

@@ -659,3 +659,33 @@ def test_immediate_saving_mode(tiny_opt_model_path, tmp_path, low_cpu_mem_usage,
     with safe_open(os.path.join(quantized_model_path, safetensor_files[0]), framework="pt") as f:
         keys = f.keys()
         assert len(keys) > 0, "Safetensors file has no tensors"
+
+
+def test_save_model_writes_diffusers_config(tmp_path):
+    """A diffusers config is a FrozenDict with no save_pretrained; the export must still write it."""
+    diffusers = pytest.importorskip("diffusers")
+
+    from auto_round.export.utils import save_model
+
+    model = diffusers.SD3Transformer2DModel(
+        sample_size=8,
+        patch_size=2,
+        in_channels=4,
+        num_layers=1,
+        attention_head_dim=32,
+        num_attention_heads=2,
+        joint_attention_dim=64,
+        caption_projection_dim=64,
+        pooled_projection_dim=64,
+        out_channels=4,
+    )
+    assert not hasattr(model.config, "save_pretrained")
+    model.config.quantization_config = {"quant_method": "auto-round", "bits": 4}
+
+    # immediate_saving: weights are already on disk, only the configs are written
+    save_model(model, str(tmp_path), immediate_saving=True)
+
+    with open(os.path.join(tmp_path, "config.json")) as f:
+        config = json.load(f)
+    assert config["_class_name"] == "SD3Transformer2DModel"
+    assert config["quantization_config"] == {"quant_method": "auto-round", "bits": 4}