enhance: add source config preservation in save_model and related tests

auto_round/export/utils.py

@@ -44,15 +44,48 @@ 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.
     """
     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.

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."""