From 53bc613244c2b07f244da08be1596476a8b9118a Mon Sep 17 00:00:00 2001 From: Ryuichi Ichinose Date: Thu, 11 Dec 2025 23:51:07 +0900 Subject: [PATCH 1/3] Chronos-2: Ensure base model revision is saved in LoRA adapter config Description: This commit fixes a critical issue where the base model's revision was failing to be recorded in `adapter_config.json` during LoRA fine-tuning. Previously, `fit()` did not propagate the loaded model's revision to `LoraConfig`, resulting in the `revision` field being missing or None in the saved adapter configuration. This omission meant that loading the saved adapter would silently default to the `main` branch of the base model, rather than the specific revision used during training, breaking reproducibility. Key changes: * Persist Revision: Store the `revision` passed to `from_pretrained` in `model.config._source_revision`. * Propagate to Config: Update `fit()` to inject this stored revision into `LoraConfig` so it is correctly saved in `adapter_config.json`. * Validate Integrity: Raise `ValueError` if an explicit `revision` argument in `fit()` conflicts with the loaded model's source revision. Impact: EEnsures that adapter_config.json always contains the correct base model revision. This guarantees that AutoGluon can faithfully reproduce the model state by simply loading the adapter, without risking a silent version mismatch, even in scenarios where the base model receives minimal or infrequent updates. --- src/chronos/chronos2/pipeline.py | 12 ++++++++++ test/test_chronos2.py | 40 ++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/src/chronos/chronos2/pipeline.py b/src/chronos/chronos2/pipeline.py index 3eddcd7f..1f882c31 100644 --- a/src/chronos/chronos2/pipeline.py +++ b/src/chronos/chronos2/pipeline.py @@ -195,6 +195,7 @@ def fit( model.load_state_dict(self.model.state_dict()) if finetune_mode == "lora": + lora_revision = getattr(self.model.config, "_source_revision", None) if lora_config is None: lora_config = LoraConfig( r=8, @@ -206,8 +207,10 @@ def fit( "self_attention.o", "output_patch_embedding.output_layer", ], + revision=lora_revision, ) elif isinstance(lora_config, dict): + lora_config.setdefault("revision", lora_revision) lora_config = LoraConfig(**lora_config) else: assert isinstance(lora_config, LoraConfig), ( @@ -1161,6 +1164,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): Load the model, either from a local path, S3 prefix or from the HuggingFace Hub. Supports the same arguments as ``AutoConfig`` and ``AutoModel`` from ``transformers``. """ + revision = kwargs.get("revision") # Check if the model is on S3 and cache it locally first # NOTE: Only base models (not LoRA adapters) are supported via S3 @@ -1178,6 +1182,10 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): model = AutoPeftModel.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) model = model.merge_and_unload() + + if revision: + model.config._source_revision = revision + return cls(model=model) # Handle the case for the base model @@ -1192,6 +1200,10 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): class_ = Chronos2Model model = class_.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + + if revision: + model.config._source_revision = revision + return cls(model=model) def save_pretrained(self, save_directory: str | Path, *args, **kwargs): diff --git a/test/test_chronos2.py b/test/test_chronos2.py index 3fac7261..a570ad3c 100644 --- a/test/test_chronos2.py +++ b/test/test_chronos2.py @@ -1132,3 +1132,43 @@ def test_eager_and_sdpa_produce_identical_outputs(pipeline): for out_eager, out_sdpa in zip(outputs_eager_grouped, outputs_sdpa_grouped): # Should match exactly or very close (numerical precision) assert torch.allclose(out_eager, out_sdpa, atol=1e-5, rtol=1e-4) + + +@pytest.mark.parametrize("source_revision", ["my-test-branch", None]) +def test_lora_config_uses_source_revision_from_instantiation( + pipeline: Chronos2Pipeline, tmpdir, source_revision +): + """ + Test that fit in 'lora' mode correctly uses the '_source_revision' + that was (notionally) stored in the model's config during instantiation. + """ + output_dir = Path(tmpdir) + dummy_inputs = [torch.rand(100)] + + # Manually set the source revision on the config to simulate + # it being set during from_pretrained + if source_revision: + pipeline.model.config._source_revision = source_revision + elif hasattr(pipeline.model.config, "_source_revision"): + delattr(pipeline.model.config, "_source_revision") + + pipeline.fit( + inputs=dummy_inputs, + prediction_length=10, + finetune_mode="lora", + output_dir=output_dir, + num_steps=1, # Keep it fast + batch_size=32, + ) + + adapter_config_path = output_dir / "finetuned-ckpt" / "adapter_config.json" + assert adapter_config_path.exists(), "adapter_config.json was not created" + + with open(adapter_config_path, "r") as f: + adapter_config = json.load(f) + + if source_revision is not None: + assert "revision" in adapter_config + assert adapter_config["revision"] == source_revision + else: + assert "revision" not in adapter_config or adapter_config["revision"] is None From 7f2f5a1d18aa307e521f8c11fb4d2ef083d5c6d5 Mon Sep 17 00:00:00 2001 From: Ryuichi Ichinose Date: Fri, 12 Dec 2025 23:13:44 +0900 Subject: [PATCH 2/3] Refactor revision handling: remove private attribute hack --- src/chronos/chronos2/pipeline.py | 19 +++++-------------- test/test_chronos2.py | 11 +++-------- 2 files changed, 8 insertions(+), 22 deletions(-) diff --git a/src/chronos/chronos2/pipeline.py b/src/chronos/chronos2/pipeline.py index 1f882c31..7ffafeef 100644 --- a/src/chronos/chronos2/pipeline.py +++ b/src/chronos/chronos2/pipeline.py @@ -40,10 +40,10 @@ class Chronos2Pipeline(BaseChronosPipeline): forecast_type: ForecastType = ForecastType.QUANTILES default_context_length: int = 2048 - def __init__(self, model: Chronos2Model): + def __init__(self, model: Chronos2Model, revision: str | None = None): super().__init__(inner_model=model) self.model = model - + self.revision = revision @staticmethod def _get_prob_mass_per_quantile_level(quantile_levels: torch.Tensor) -> torch.Tensor: """ @@ -195,7 +195,7 @@ def fit( model.load_state_dict(self.model.state_dict()) if finetune_mode == "lora": - lora_revision = getattr(self.model.config, "_source_revision", None) + lora_revision = self.revision if lora_config is None: lora_config = LoraConfig( r=8, @@ -210,7 +210,6 @@ def fit( revision=lora_revision, ) elif isinstance(lora_config, dict): - lora_config.setdefault("revision", lora_revision) lora_config = LoraConfig(**lora_config) else: assert isinstance(lora_config, LoraConfig), ( @@ -1182,11 +1181,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): model = AutoPeftModel.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) model = model.merge_and_unload() - - if revision: - model.config._source_revision = revision - - return cls(model=model) + return cls(model=model, revision=revision) # Handle the case for the base model config = AutoConfig.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) @@ -1200,11 +1195,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): class_ = Chronos2Model model = class_.from_pretrained(pretrained_model_name_or_path, *args, **kwargs) - - if revision: - model.config._source_revision = revision - - return cls(model=model) + return cls(model=model, revision=revision) def save_pretrained(self, save_directory: str | Path, *args, **kwargs): """ diff --git a/test/test_chronos2.py b/test/test_chronos2.py index a570ad3c..14f81796 100644 --- a/test/test_chronos2.py +++ b/test/test_chronos2.py @@ -1139,18 +1139,13 @@ def test_lora_config_uses_source_revision_from_instantiation( pipeline: Chronos2Pipeline, tmpdir, source_revision ): """ - Test that fit in 'lora' mode correctly uses the '_source_revision' - that was (notionally) stored in the model's config during instantiation. + Test that fit in 'lora' mode correctly uses the 'revision' + stored in the pipeline instance. """ output_dir = Path(tmpdir) dummy_inputs = [torch.rand(100)] - # Manually set the source revision on the config to simulate - # it being set during from_pretrained - if source_revision: - pipeline.model.config._source_revision = source_revision - elif hasattr(pipeline.model.config, "_source_revision"): - delattr(pipeline.model.config, "_source_revision") + pipeline.revision = source_revision pipeline.fit( inputs=dummy_inputs, From 6248da86ec295714aee64129d772e8658c5ac13e Mon Sep 17 00:00:00 2001 From: Ryuichi Ichinose Date: Sat, 13 Dec 2025 16:35:24 +0900 Subject: [PATCH 3/3] Add: If dict is given --- src/chronos/chronos2/pipeline.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/chronos/chronos2/pipeline.py b/src/chronos/chronos2/pipeline.py index 7ffafeef..d0a89d8d 100644 --- a/src/chronos/chronos2/pipeline.py +++ b/src/chronos/chronos2/pipeline.py @@ -210,6 +210,7 @@ def fit( revision=lora_revision, ) elif isinstance(lora_config, dict): + lora_config.setdefault("revision", lora_revision) lora_config = LoraConfig(**lora_config) else: assert isinstance(lora_config, LoraConfig), (