Chronos-2: Add LoRA fine-tuning support (#393)

*Issue #, if available:*

*Description of changes:* Adds support for LoRA fine-tuning.

- [x] Move peft/pandas dependency to an extra
- [x] Add tests for LoRA
- [x] Update notebook with LoRA info
- [x] Enable automatic recognition and loading of LoRA adapters


By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Author: abdulfatir

notebooks/chronos-2-quickstart.ipynb

@@ -37,7 +37,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%pip install 'chronos-forecasting>=2.0' 'pandas[pyarrow]' 'matplotlib'"
+    "%pip install 'chronos-forecasting>=2.1' 'pandas[pyarrow]' 'matplotlib'"
    ]
   },
   {

pyproject.toml

@@ -20,7 +20,6 @@ dependencies = [
   "numpy>=1.21,<3",
   "einops>=0.7.0,<1",
   "scikit-learn>=1.6.0,<2",
-  "boto3",
 ]
 classifiers = [
   "Programming Language :: Python :: 3",
@@ -40,7 +39,19 @@ packages = ["src/chronos"]
 path = "src/chronos/__about__.py"
 
 [project.optional-dependencies]
-test = ["pytest~=8.0", "numpy>=1.21,<3", "fev>=0.6.1", "pandas>=2.0,<2.4"]
+extras = [
+  "boto3>=1.10,<2",
+  "peft>=0.13.0,<1",
+  "fev>=0.6.1",
+  "pandas[pyarrow]>=2.0,<2.4",
+]
+test = [
+  "pytest~=8.0",
+  "boto3>=1.10,<2",
+  "peft>=0.13.0,<1",
+  "fev>=0.6.1",
+  "pandas[pyarrow]>=2.0,<2.4",
+]
 typecheck = ["mypy~=1.9"]
 dev = [
   "gluonts[pro]~=0.16",

src/chronos/chronos2/pipeline.py

@@ -9,13 +9,15 @@
 import warnings
 from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence
 
 import numpy as np
 import torch
 from einops import rearrange, repeat
 from torch.utils.data import DataLoader
 from transformers import AutoConfig
+from transformers.utils.import_utils import is_peft_available
+from transformers.utils.peft_utils import find_adapter_config_file
 
 import chronos.chronos2
 from chronos.base import BaseChronosPipeline, ForecastType
@@ -28,6 +30,7 @@
     import datasets
     import fev
     import pandas as pd
+    from peft import LoraConfig
 
 logger = logging.getLogger(__name__)
 
@@ -99,6 +102,8 @@ def fit(
         | Sequence[TensorOrArray]
         | Sequence[Mapping[str, TensorOrArray | Mapping[str, TensorOrArray | None]]]
         | None = None,
+        finetune_mode: Literal["full", "lora"] = "full",
+        lora_config: "LoraConfig | dict | None" = None,
         context_length: int | None = None,
         learning_rate: float = 1e-6,
         num_steps: int = 1000,
@@ -123,10 +128,16 @@ def fit(
         validation_inputs
             The time series used for validation and model selection. The format of `validation_inputs` is exactly the same as `inputs`, by default None which
             means that no validation is performed. Note that enabling validation may slow down fine-tuning for large datasets.
+        finetune_mode
+            One of "full" (performs full fine-tuning) or "lora" (performs Low Rank Adaptation (LoRA) fine-tuning), by default "full"
+        lora_config
+            The configuration to use for LoRA fine-tuning when finetune_mode="lora". Can be a `LoraConfig` object or a dict which is used to initialize `LoraConfig`.
+            When unspecified and finetune_mode="lora", a default configuration is used
         context_length
             The maximum context length used during fine-tuning, by default set to the model's default context length
         learning_rate
             The learning rate for the optimizer, by default 1e-6
+            When finetune_mode="lora", we recommend using a higher value of the learning rate, such as 1e-5
         num_steps
             The number of steps to fine-tune for, by default 1000
         batch_size
@@ -151,13 +162,55 @@ def fit(
         import torch.cuda
         from transformers.training_args import TrainingArguments
 
+        if finetune_mode == "lora":
+            if is_peft_available():
+                from peft import LoraConfig, get_peft_model
+            else:
+                warnings.warn(
+                    "`peft` is required for `finetune_mode='lora'`. Please install it with `pip install peft`. Falling back to `finetune_mode='full'`."
+                )
+                finetune_mode = "full"
+
         from chronos.chronos2.trainer import Chronos2Trainer, EvaluateAndSaveFinalStepCallback
 
+        assert finetune_mode in ["full", "lora"], f"finetune_mode must be one of ['full', 'lora'], got {finetune_mode}"
+
+        if finetune_mode == "full" and lora_config is not None:
+            raise ValueError(
+                "lora_config should not be specified when `finetune_mode='full'`. To enable LoRA, set `finetune_mode='lora'`."
+            )
+
         # Create a copy of the model to avoid modifying the original
         config = deepcopy(self.model.config)
         model = Chronos2Model(config).to(self.model.device)  # type: ignore
         model.load_state_dict(self.model.state_dict())
 
+        if finetune_mode == "lora":
+            if lora_config is None:
+                lora_config = LoraConfig(
+                    r=8,
+                    lora_alpha=16,
+                    target_modules=[
+                        "self_attention.q",
+                        "self_attention.v",
+                        "self_attention.k",
+                        "self_attention.o",
+                        "output_patch_embedding.output_layer",
+                    ],
+                )
+            elif isinstance(lora_config, dict):
+                lora_config = LoraConfig(**lora_config)
+            else:
+                assert isinstance(lora_config, LoraConfig), (
+                    f"lora_config must be an instance of LoraConfig or a dict, got {type(lora_config)}"
+                )
+
+            model = get_peft_model(model, lora_config)
+            n_trainable_params, n_params = model.get_nb_trainable_parameters()
+            logger.info(
+                f"Using LoRA. Number of trainable parameters: {n_trainable_params}, total parameters: {n_params}."
+            )
+
         if context_length is None:
             context_length = self.model_context_length
 
@@ -1064,9 +1117,25 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
         Supports the same arguments as ``AutoConfig`` and ``AutoModel`` from ``transformers``.
         """
 
+        # Check if the model is on S3 and cache it locally first
+        # NOTE: Only base models (not LoRA adapters) are supported via S3
         if str(pretrained_model_name_or_path).startswith("s3://"):
             return BaseChronosPipeline.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
 
+        # Check if the hub model_id or local path is a LoRA adapter
+        if find_adapter_config_file(pretrained_model_name_or_path) is not None:
+            if not is_peft_available():
+                raise ImportError(
+                    f"The model at {pretrained_model_name_or_path} is a `peft` adaptor, but `peft` is not available. "
+                    f"Please install `peft` with `pip install peft` to use this model. "
+                )
+            from peft import AutoPeftModel
+
+            model = AutoPeftModel.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+            model = model.merge_and_unload()
+            return cls(model=model)
+
+        # Handle the case for the base model
         config = AutoConfig.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
         assert hasattr(config, "chronos_config"), "Not a Chronos config file"
 

test/dummy-chronos2-lora/adapter_config.json

@@ -0,0 +1,35 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": {
+    "base_model_class": "Chronos2Model",
+    "parent_library": "chronos.chronos2.model"
+  },
+  "base_model_name_or_path": "test/dummy-chronos2-model",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 8,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "self_attention.q",
+    "self_attention.k",
+    "self_attention.o",
+    "output_patch_embedding.output_layer",
+    "self_attention.v"
+  ],
+  "task_type": null,
+  "use_dora": false,
+  "use_rslora": false
+}

test/dummy-chronos2-lora/adapter_model.safetensors

@@ -38,6 +38,10 @@ def test_base_chronos2_pipeline_loads_from_hf():
     BaseChronosPipeline.from_pretrained("amazon/chronos-2", device_map="cpu")
 
 
+def test_chronos2_lora_pipeline_loads_from_disk():
+    Chronos2Pipeline.from_pretrained(Path(__file__).parent / "dummy-chronos2-lora", device_map="cpu")
+
+
 @pytest.mark.parametrize(
     "inputs, prediction_length, expected_output_shapes",
     [
@@ -671,12 +675,20 @@ def test_predict_df_with_future_df_with_different_freq_raises_error(pipeline):
         ),
     ],
 )
+@pytest.mark.parametrize("finetune_mode", ["full", "lora"])
 def test_when_input_is_valid_then_pipeline_can_be_finetuned(
-    pipeline, inputs, prediction_length, expected_output_shapes
+    pipeline, inputs, prediction_length, expected_output_shapes, finetune_mode
 ):
     # Get outputs before fine-tuning
     orig_outputs_before = pipeline.predict(inputs, prediction_length=prediction_length)
-    ft_pipeline = pipeline.fit(inputs, prediction_length=prediction_length, num_steps=5, min_past=1, batch_size=32)
+    ft_pipeline = pipeline.fit(
+        inputs,
+        prediction_length=prediction_length,
+        num_steps=5,
+        min_past=1,
+        batch_size=32,
+        finetune_mode=finetune_mode,
+    )
     # Get outputs from fine-tuned pipeline
     ft_outputs = ft_pipeline.predict(inputs, prediction_length=prediction_length)
     # Get outputs from original pipeline after fine-tuning