diff --git a/auto_round/export/export_to_llmcompressor/export.py b/auto_round/export/export_to_llmcompressor/export.py index cd0cc2186..b4c73423c 100644 --- a/auto_round/export/export_to_llmcompressor/export.py +++ b/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) diff --git a/auto_round/export/utils.py b/auto_round/export/utils.py index d62192999..3751c47f7 100644 --- a/auto_round/export/utils.py +++ b/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) diff --git a/auto_round/formats.py b/auto_round/formats.py index 314ceb9cf..76c1ed002 100644 --- a/auto_round/formats.py +++ b/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) diff --git a/test/test_cpu/export/test_export.py b/test/test_cpu/export/test_export.py index e23d2afbd..9b27d76cc 100644 --- a/test/test_cpu/export/test_export.py +++ b/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}