diff --git a/auto_round/export/utils.py b/auto_round/export/utils.py index f76d308091..ae212d8b07 100644 --- a/auto_round/export/utils.py +++ b/auto_round/export/utils.py @@ -44,8 +44,39 @@ def save_pretrained_artifact(artifact, output_dir: str, artifact_name: str = "ar return True +def _save_source_config(model: nn.Module, save_dir: str) -> bool: + """Preserve the source JSON, replacing only its quantization metadata when available.""" + try: + source_dir = _resolve_model_source_dir(model) + if source_dir is None: + return False + config = getattr(model, "config", None) + if os.path.isdir(source_dir): + config_path = os.path.join(source_dir, "config.json") + else: + from transformers.utils.hub import cached_file + + revision = getattr(config, "_commit_hash", None) + config_path = cached_file(source_dir, "config.json", revision=revision) + with open(config_path, encoding="utf-8") as source_file: + config_dict = json.load(source_file) + if not isinstance(config_dict, dict): + raise TypeError("Source config.json must contain a JSON object") + quantization_config = getattr(config, "quantization_config", None) + if quantization_config is not None: + config_dict["quantization_config"] = quantization_config + config_json = json.dumps(config_dict, indent=2, ensure_ascii=False) + os.makedirs(save_dir, exist_ok=True) + with open(os.path.join(save_dir, "config.json"), "w", encoding="utf-8") as config_file: + config_file.write(config_json) + return True + except Exception as error: + logger.warning("Could not save original model config; using model.config instead: %s", error) + return False + + def save_config_artifact(model: nn.Module, save_dir: str) -> None: - """Write ``model.config`` to ``save_dir``, for transformers and diffusers models alike. + """Prefer the source config JSON, falling back to serializing ``model.config``. A diffusers ``ModelMixin`` keeps its config in a ``FrozenDict``, which has no ``save_pretrained``; ``ModelMixin.save_config`` is the equivalent writer. @@ -53,6 +84,8 @@ def save_config_artifact(model: nn.Module, save_dir: str) -> None: config = getattr(model, "config", None) if config is None: return + if _save_source_config(model, save_dir): + return if not hasattr(config, "save_pretrained") and hasattr(model, "save_config"): model.save_config(save_dir) @@ -432,6 +465,7 @@ def save_model( model.save_pretrained( save_dir, max_shard_size=max_shard_size, safe_serialization=safe_serialization, **save_kwargs ) + _save_source_config(model, save_dir) except (KeyError, TypeError) as e: # Some third-party configs fail during config serialization in save_pretrained. # Fall back to saving weights separately + config without diff. diff --git a/test/unit/common/export/test_export_utils.py b/test/unit/common/export/test_export_utils.py index 786bc87c8a..07c01bb81d 100644 --- a/test/unit/common/export/test_export_utils.py +++ b/test/unit/common/export/test_export_utils.py @@ -124,6 +124,87 @@ def test_fallback_on_save_pretrained_failure(self): # ============================================================================== +class TestSourceConfig: + @pytest.mark.parametrize("immediate_saving", [False, True]) + @pytest.mark.parametrize("dtype", [None, torch.bfloat16]) + def test_preserves_original_config(self, tmp_path, immediate_saving, dtype): + source_dir = tmp_path / "source" + source_dir.mkdir() + original = { + "model_type": "original", + "text_config": {"custom_field": [1, 2]}, + "torch_dtype": "float32", + "dtype": "float32", + "quantization_config": {"quant_method": "old"}, + } + (source_dir / "config.json").write_text(json.dumps(original), encoding="utf-8") + quantization_config = {"quant_method": "auto-round", "bits": 4} + model = nn.Linear(4, 4) + model.dtype = torch.float32 + model.config = SimpleNamespace( + _name_or_path=str(source_dir), + model_type="changed", + quantization_config=quantization_config, + save_pretrained=MagicMock(), + ) + output_dir = tmp_path / "output" + + def save_reconstructed_config(*args, **kwargs): + (output_dir / "config.json").write_text('{"model_type": "changed"}', encoding="utf-8") + + model.save_pretrained = MagicMock(side_effect=save_reconstructed_config) + with patch("auto_round.export.utils.apply_post_save_source_fixes"): + save_model(model, str(output_dir), immediate_saving=immediate_saving, dtype=dtype) + + expected = {**original, "quantization_config": quantization_config} + if dtype is not None: + expected.update(torch_dtype="bfloat16", dtype="bfloat16") + assert json.loads((output_dir / "config.json").read_text()) == expected + assert json.loads((output_dir / "quantization_config.json").read_text()) == quantization_config + assert json.loads((source_dir / "config.json").read_text()) == original + model.config.save_pretrained.assert_not_called() + + def test_remote_failure_falls_back(self, tmp_path): + model = nn.Module() + output_dir = str(tmp_path / "output") + model.config = SimpleNamespace(_name_or_path="org/model", save_pretrained=MagicMock()) + with patch("transformers.utils.hub.cached_file", side_effect=OSError("download failed")): + _save_model_configs(model, output_dir) + model.config.save_pretrained.assert_called_once_with(output_dir) + + @pytest.mark.parametrize("source_content", [None, "invalid json", "[]"]) + def test_invalid_source_falls_back(self, tmp_path, source_content): + source_dir = tmp_path / "source" + source_dir.mkdir() + if source_content is not None: + (source_dir / "config.json").write_text(source_content, encoding="utf-8") + model = nn.Module() + model.config = SimpleNamespace(_name_or_path=str(source_dir), save_pretrained=MagicMock()) + _save_model_configs(model, str(tmp_path / "output")) + model.config.save_pretrained.assert_called_once_with(str(tmp_path / "output")) + + def test_remote_source_config(self, tmp_path): + source_config = tmp_path / "original.json" + source_config.write_text('{"model_type": "original", "custom_field": true}', encoding="utf-8") + model = nn.Module() + model.config = SimpleNamespace( + _name_or_path="org/model", + _commit_hash="source-revision", + quantization_config={"bits": 4}, + save_pretrained=MagicMock(), + ) + output_dir = tmp_path / "output" + with patch("transformers.utils.hub.cached_file", return_value=str(source_config)) as cached_file: + _save_model_configs(model, str(output_dir)) + cached_file.assert_called_once_with("org/model", "config.json", revision="source-revision") + assert json.loads((output_dir / "config.json").read_text()) == { + "model_type": "original", + "custom_field": True, + "quantization_config": {"bits": 4}, + } + model.config.save_pretrained.assert_not_called() + + class TestStateDictHasMetaTensor: """Test _state_dict_has_meta_tensor function."""