diff --git a/.gitignore b/.gitignore index 44f18c0..24495b0 100644 --- a/.gitignore +++ b/.gitignore @@ -3,7 +3,6 @@ dist/ __pycache__/ *.egg-info/ checkpoints/ -peft_checkpoints/ wandb/ datasets/ results/ diff --git a/README.md b/README.md index d7a2baf..56e8253 100644 --- a/README.md +++ b/README.md @@ -22,12 +22,12 @@ This open version is not an officially supported Google product. install timesfm==1.3.0` to install an older version of this package to load them. -## Update - Apr. 8, 2026 +## Update - Apr. 9, 2026 -Added PEFT (LoRA/DoRA) fine-tuning pipeline for TimesFM 2.5 with multi-GPU -support. See [`peft/`](peft/) for docs and usage. Also added unit tests -(`tests/`), fixed per-input ridge regression in XReg to prevent data leakage, -and incorporated several community fixes. +Added fine-tuning example using HuggingFace Transformers + PEFT (LoRA) — see +[`timesfm-forecasting/examples/finetuning/`](timesfm-forecasting/examples/finetuning/). +Also added unit tests (`tests/`), fixed per-input ridge regression in XReg to +prevent data leakage, and incorporated several community fixes. ## Update - Mar. 19, 2026 @@ -56,7 +56,7 @@ Since the Sept. 2025 launch, the following improvements have been completed: 1. ✅ Flax version of the model for faster inference. 2. ✅ Covariate support via XReg (see Oct. 2025 update). 3. ✅ Documentation, examples, and agent skill (see `timesfm-forecasting/`). -4. ✅ PEFT fine-tuning pipeline with LoRA/DoRA and multi-GPU support (see `peft/`). +4. ✅ Fine-tuning example with LoRA via HuggingFace Transformers + PEFT (see `timesfm-forecasting/examples/finetuning/`). 5. ✅ Unit tests for core layers, configs, and utilities (see `tests/`). ### Install diff --git a/peft/README.md b/peft/README.md deleted file mode 100644 index fdc96f6..0000000 --- a/peft/README.md +++ /dev/null @@ -1,201 +0,0 @@ -# TimesFM 2.5 — PEFT Fine-Tuning Pipeline - -Production-grade **LoRA / DoRA** fine-tuning for -[TimesFM 2.5](https://github.com/google-research/timesfm) (200M PyTorch) -with **multi-GPU** support via PyTorch DDP. - -## Features - -| Strategy | Description | -|---|---| -| **LoRA** | Low-Rank Adaptation — adds trainable A/B matrices to frozen linear layers ([paper](https://arxiv.org/abs/2106.09685)) | -| **DoRA** | Weight-Decomposed LoRA — decomposes adapted weights into magnitude + direction ([paper](https://arxiv.org/abs/2402.09353)) | -| **Linear Probing** | Train only the output heads (`--train_output_head`) with `--lora_rank 0` | - -Additional capabilities: - -- **Multi-GPU** via `torchrun` (DDP) -- **Mixed precision** — fp16 or bf16 -- **Gradient checkpointing** — trade compute for memory on long contexts -- **Cosine-with-warmup** LR schedule -- **Early stopping** on validation loss -- **Adapter-only** checkpoint saving / loading (safetensors) -- **Weight merging** — fold adapters back into base weights for zero-overhead inference -- **Quantile loss** — optional pinball loss on all 9 quantile channels -- **W&B logging** (opt-in) - -## Quick Start - -### 1. Install - -```bash -# From the repo root -pip install -e ".[torch]" -``` - -### 2. Prepare Data - -Your CSV can be in either format: - -- **Long format** — columns: `[id, timestamp, value]` -- **Wide format** — each numeric column is an independent series - -### 3. Single-GPU Training - -```bash -python -m peft.finetune \ - --data_path data.csv \ - --value_col y \ - --context_len 512 \ - --horizon_len 128 \ - --adapter_type lora \ - --lora_rank 8 \ - --num_epochs 10 \ - --batch_size 32 -``` - -### 4. Multi-GPU Training - -```bash -torchrun --nproc_per_node=4 -m peft.finetune \ - --data_path data.csv \ - --value_col y \ - --adapter_type dora \ - --lora_rank 16 \ - --mixed_precision bf16 \ - --gradient_checkpointing -``` - -### 5. Using the Launch Script - -```bash -# Edit environment variables to taste -DATA_PATH=data.csv VALUE_COL=y NUM_GPUS=4 bash peft/finetune.sh -``` - -## Python API - -```python -import numpy as np -from timesfm.timesfm_2p5.timesfm_2p5_torch import TimesFM_2p5_200M_torch -from timesfm.configs import ForecastConfig - -from peft import PEFTConfig, PEFTTrainer, TimeSeriesDataset - -# 1. Load pretrained model (no torch.compile for training) -wrapper = TimesFM_2p5_200M_torch.from_pretrained( - "google/timesfm-2.5-200m-pytorch", - torch_compile=False, -) - -# 2. Configure PEFT -config = PEFTConfig( - adapter_type="lora", # or "dora" - lora_rank=8, - lora_alpha=16, - target_modules="all", # "all" | "attention" | "ffn" - learning_rate=1e-4, - num_epochs=10, - batch_size=32, - context_len=512, - horizon_len=128, - mixed_precision="bf16", # "no" | "fp16" | "bf16" -) - -# 3. Create datasets -train_series = [np.random.randn(2000).astype(np.float32) for _ in range(100)] -val_series = [np.random.randn(800).astype(np.float32) for _ in range(100)] - -train_ds = TimeSeriesDataset(train_series, context_len=512, horizon_len=128, stride=32) -val_ds = TimeSeriesDataset(val_series, context_len=512, horizon_len=128, stride=128) - -# 4. Train -trainer = PEFTTrainer(wrapper.model, config) -history = trainer.fit(train_ds, val_ds) - -# 5. Save adapter-only checkpoint (~2 MB for rank-8 LoRA) -trainer.save_adapter("./my_adapter/adapter.safetensors") - -# 6. Merge adapter into base model for zero-overhead inference -trainer.merge_adapter() -wrapper.compile(ForecastConfig(max_context=512, max_horizon=128)) -point, quantiles = wrapper.forecast(horizon=128, inputs=[my_series]) -``` - -## Loading a Saved Adapter - -```python -from peft import PEFTConfig, inject_adapters, load_adapter_weights - -wrapper = TimesFM_2p5_200M_torch.from_pretrained( - "google/timesfm-2.5-200m-pytorch", torch_compile=False -) - -# Must inject adapters with the *same* config before loading weights. -config = PEFTConfig(adapter_type="lora", lora_rank=8, target_modules="all") -inject_adapters(wrapper.model, config) -load_adapter_weights(wrapper.model, "./my_adapter/adapter.safetensors") - -# Option A: use with adapters active -# Option B: merge for maximum inference throughput -from peft import merge_adapters -merge_adapters(wrapper.model) -``` - -## Architecture - -TimesFM 2.5 (200M) has 20 transformer layers, each containing: - -| Linear Layer | Shape | LoRA params (rank 8) | -|---|---|---| -| `attn.qkv_proj` (fused Q/K/V) | 1280 → 3840 | 40,960 | -| `attn.out` | 1280 → 1280 | 20,480 | -| `ff0` | 1280 → 1280 | 20,480 | -| `ff1` | 1280 → 1280 | 20,480 | - -With `target_modules="all"` and `lora_rank=8`: - -- **2,048,000** trainable adapter parameters (~1% of the 200M total) -- DoRA adds ~102,400 magnitude parameters (negligible overhead) - -## CLI Options - -``` -python -m peft.finetune --help -``` - -| Flag | Default | Description | -|---|---|---| -| `--model_id` | `google/timesfm-2.5-200m-pytorch` | HF repo or local path | -| `--data_path` | *(required)* | Path to CSV | -| `--id_col` | `None` | Series identifier column (long format) | -| `--value_col` | `None` | Value column (long format) | -| `--context_len` | 512 | Context window (rounded to multiple of 32) | -| `--horizon_len` | 128 | Prediction horizon (≤ 128) | -| `--adapter_type` | `lora` | `lora` or `dora` | -| `--lora_rank` | 8 | Low-rank dimension | -| `--lora_alpha` | 16 | Scaling factor | -| `--target_modules` | `all` | `all`, `attention`, or `ffn` | -| `--train_output_head` | off | Also train output projections | -| `--num_epochs` | 10 | Training epochs | -| `--batch_size` | 32 | Per-GPU batch size | -| `--learning_rate` | 1e-4 | Peak learning rate | -| `--mixed_precision` | `no` | `no`, `fp16`, or `bf16` | -| `--gradient_checkpointing` | off | Activation checkpointing | -| `--use_quantile_loss` | off | Add pinball loss | -| `--use_wandb` | off | W&B logging | -| `--early_stopping_patience` | 5 | Patience epochs | - -## File Layout - -``` -peft/ -├── __init__.py # Public API -├── adapters.py # LoRA / DoRA layers + inject / merge / save / load -├── config.py # PEFTConfig dataclass -├── data.py # TimeSeriesDataset -├── trainer.py # PEFTTrainer (DDP, AMP, checkpointing) -├── finetune.py # CLI entry-point -├── finetune.sh # Example launch script -└── README.md # This file -``` diff --git a/peft/__init__.py b/peft/__init__.py deleted file mode 100644 index 2cadf0d..0000000 --- a/peft/__init__.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""PEFT (LoRA/DoRA) fine-tuning pipeline for TimesFM 2.5.""" - -from .adapters import ( - DoRALinear, - LoRALinear, - get_adapter_params, - inject_adapters, - load_adapter_weights, - merge_adapters, - save_adapter_weights, -) -from .config import PEFTConfig -from .data import TimeSeriesDataset -from .trainer import PEFTTrainer - -__all__ = [ - "PEFTConfig", - "PEFTTrainer", - "TimeSeriesDataset", - "LoRALinear", - "DoRALinear", - "inject_adapters", - "merge_adapters", - "save_adapter_weights", - "load_adapter_weights", - "get_adapter_params", -] diff --git a/peft/adapters.py b/peft/adapters.py deleted file mode 100644 index acc6eee..0000000 --- a/peft/adapters.py +++ /dev/null @@ -1,285 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""LoRA and DoRA adapter layers for PyTorch, plus injection / merging helpers. - -References: - LoRA — https://arxiv.org/abs/2106.09685 - DoRA — https://arxiv.org/abs/2402.09353 -""" - -import math -import os -from collections import OrderedDict -from typing import Dict - -import torch -import torch.nn as nn -import torch.nn.functional as F -from safetensors.torch import load_file, save_file - -from .config import PEFTConfig - - -# --------------------------------------------------------------------------- -# Adapter layers -# --------------------------------------------------------------------------- - - -class LoRALinear(nn.Module): - """Drop-in replacement for ``nn.Linear`` that adds a low-rank branch. - - ``output = base_linear(x) + (dropout(x) @ A @ B) * (alpha / rank)`` - - *A* is Kaiming-uniform initialised; *B* is zero-initialised so the - effective delta is zero at init and the pretrained model is preserved. - """ - - def __init__( - self, - base_linear: nn.Linear, - rank: int = 8, - alpha: float = 16.0, - dropout: float = 0.0, - ): - super().__init__() - self.base_linear = base_linear - self.rank = rank - self.scaling = alpha / rank - - in_f = base_linear.in_features - out_f = base_linear.out_features - - self.lora_A = nn.Parameter(torch.empty(in_f, rank)) - self.lora_B = nn.Parameter(torch.zeros(rank, out_f)) - nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) - - self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() - - # Freeze the pretrained weight. - self.base_linear.weight.requires_grad = False - if self.base_linear.bias is not None: - self.base_linear.bias.requires_grad = False - - def forward(self, x: torch.Tensor) -> torch.Tensor: - base_out = self.base_linear(x) - lora_out = self.dropout(x) @ self.lora_A @ self.lora_B * self.scaling - return base_out + lora_out - - def merge_weights(self) -> nn.Linear: - """Fold the LoRA delta into the base ``nn.Linear`` and return it.""" - with torch.no_grad(): - delta = (self.lora_A @ self.lora_B * self.scaling).T # (out, in) - self.base_linear.weight.add_(delta) - return self.base_linear - - -class DoRALinear(nn.Module): - """Weight-Decomposed Low-Rank Adaptation (DoRA). - - Decomposes the adapted weight into *magnitude* and *direction*:: - - W' = m · (W + ΔW) / ‖W + ΔW‖_col - - ``m`` is initialised from the pretrained column norms so the model - starts at the same operating point. - """ - - def __init__( - self, - base_linear: nn.Linear, - rank: int = 8, - alpha: float = 16.0, - dropout: float = 0.0, - ): - super().__init__() - self.base_linear = base_linear - self.rank = rank - self.scaling = alpha / rank - - in_f = base_linear.in_features - out_f = base_linear.out_features - dev = base_linear.weight.device - - self.lora_A = nn.Parameter(torch.empty(in_f, rank, device=dev)) - self.lora_B = nn.Parameter(torch.zeros(rank, out_f, device=dev)) - nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) - - # Magnitude vector — initialised from pretrained column norms. - with torch.no_grad(): - col_norms = base_linear.weight.norm(dim=1) - self.magnitude = nn.Parameter(col_norms.clone()) - - self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() - - self.base_linear.weight.requires_grad = False - if self.base_linear.bias is not None: - self.base_linear.bias.requires_grad = False - - def forward(self, x: torch.Tensor) -> torch.Tensor: - delta_W = (self.lora_A @ self.lora_B * self.scaling).T # (out, in) - adapted_W = self.base_linear.weight + delta_W - col_norm = adapted_W.norm(dim=1, keepdim=True).clamp(min=1e-8) - W_prime = self.magnitude.unsqueeze(1) * (adapted_W / col_norm) - return F.linear(x, W_prime, self.base_linear.bias) - - def merge_weights(self) -> nn.Linear: - """Fold DoRA into the base ``nn.Linear`` and return it.""" - with torch.no_grad(): - delta_W = (self.lora_A @ self.lora_B * self.scaling).T - adapted_W = self.base_linear.weight + delta_W - col_norm = adapted_W.norm(dim=1, keepdim=True).clamp(min=1e-8) - self.base_linear.weight.copy_( - self.magnitude.unsqueeze(1) * (adapted_W / col_norm) - ) - return self.base_linear - - -# --------------------------------------------------------------------------- -# Injection / merge helpers -# --------------------------------------------------------------------------- - -_ADAPTER_CLS = {"lora": LoRALinear, "dora": DoRALinear} - - -def inject_adapters( - model: nn.Module, - config: PEFTConfig, -) -> nn.Module: - """Inject LoRA / DoRA adapters into a ``TimesFM_2p5_200M_torch_module``. - - All base parameters are frozen. Only adapter parameters (and, optionally, - the output-projection heads) remain trainable. - - Args: - model: The ``TimesFM_2p5_200M_torch_module`` instance. - config: PEFT configuration. - - Returns: - The same model, mutated in-place with adapter wrappers. - """ - adapter_cls = _ADAPTER_CLS[config.adapter_type] - kwargs = dict(rank=config.lora_rank, alpha=config.lora_alpha, dropout=config.lora_dropout) - target = config.target_modules - - # 1. Freeze everything. - for p in model.parameters(): - p.requires_grad = False - - # 2. Determine which layers get adapters. - total_layers = model.x # 20 - if config.num_adapter_layers > 0: - first_adapter_layer = total_layers - config.num_adapter_layers - else: - first_adapter_layer = 0 - - # 3. Wrap target nn.Linear modules with adapters. - for layer_idx in range(total_layers): - if layer_idx < first_adapter_layer: - continue - xf = model.stacked_xf[layer_idx] - - if target in ("all", "attention"): - # Fused QKV projection (TimesFM 2.5 always uses fuse_qkv=True). - if hasattr(xf.attn, "qkv_proj") and isinstance(xf.attn.qkv_proj, nn.Linear): - xf.attn.qkv_proj = adapter_cls(xf.attn.qkv_proj, **kwargs) - else: - # Fallback for non-fused Q / K / V. - for attr in ("query", "key", "value"): - orig = getattr(xf.attn, attr, None) - if isinstance(orig, nn.Linear): - setattr(xf.attn, attr, adapter_cls(orig, **kwargs)) - # Output projection. - if isinstance(xf.attn.out, nn.Linear): - xf.attn.out = adapter_cls(xf.attn.out, **kwargs) - - if target in ("all", "ffn"): - if isinstance(xf.ff0, nn.Linear): - xf.ff0 = adapter_cls(xf.ff0, **kwargs) - if isinstance(xf.ff1, nn.Linear): - xf.ff1 = adapter_cls(xf.ff1, **kwargs) - - # 4. Optionally unfreeze output heads. - if config.train_output_head: - for p in model.output_projection_point.parameters(): - p.requires_grad = True - for p in model.output_projection_quantiles.parameters(): - p.requires_grad = True - - return model - - -def merge_adapters(model: nn.Module) -> nn.Module: - """Fold all adapter weights back into base ``nn.Linear`` layers. - - After merging, the model has standard ``nn.Linear`` modules and can be - used for normal inference or saved as a regular checkpoint. - """ - for layer_idx in range(model.x): - xf = model.stacked_xf[layer_idx] - - for attr in ("qkv_proj", "out"): - layer = getattr(xf.attn, attr, None) - if isinstance(layer, (LoRALinear, DoRALinear)): - setattr(xf.attn, attr, layer.merge_weights()) - for attr in ("query", "key", "value"): - layer = getattr(xf.attn, attr, None) - if isinstance(layer, (LoRALinear, DoRALinear)): - setattr(xf.attn, attr, layer.merge_weights()) - for attr in ("ff0", "ff1"): - layer = getattr(xf, attr, None) - if isinstance(layer, (LoRALinear, DoRALinear)): - setattr(xf, attr, layer.merge_weights()) - - # Unfreeze everything so the merged model can be retrained if desired. - for p in model.parameters(): - p.requires_grad = True - - return model - - -# --------------------------------------------------------------------------- -# Save / load adapter-only weights -# --------------------------------------------------------------------------- - - -def get_adapter_params(model: nn.Module) -> Dict[str, torch.Tensor]: - """Return an ``OrderedDict`` of all trainable (adapter) parameters.""" - return OrderedDict( - (n, p.data) for n, p in model.named_parameters() if p.requires_grad - ) - - -def save_adapter_weights(model: nn.Module, path: str) -> None: - """Save adapter weights to a ``safetensors`` file.""" - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - save_file(get_adapter_params(model), path) - - -def load_adapter_weights(model: nn.Module, path: str) -> None: - """Load adapter weights from a ``safetensors`` file. - - The model must already have adapters injected (via ``inject_adapters``) - before calling this function. - """ - device = str(next(model.parameters()).device) - tensors = load_file(path, device=device) - trainable = {n for n, p in model.named_parameters() if p.requires_grad} - missing = trainable - set(tensors.keys()) - if missing: - raise ValueError(f"Adapter checkpoint is missing keys: {missing}") - - state = model.state_dict() - state.update(tensors) - model.load_state_dict(state, strict=True) diff --git a/peft/config.py b/peft/config.py deleted file mode 100644 index 643be89..0000000 --- a/peft/config.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Configuration for the TimesFM 2.5 PEFT fine-tuning pipeline.""" - -from dataclasses import dataclass, field -from typing import List, Literal, Optional - - -@dataclass -class PEFTConfig: - """Full configuration for PEFT fine-tuning of TimesFM 2.5. - - Attributes: - adapter_type: Type of adapter — "lora" or "dora". - lora_rank: Rank of the low-rank decomposition. - lora_alpha: Scaling factor (effective lr multiplier = alpha / rank). - lora_dropout: Dropout applied to the LoRA path. - target_modules: Which layers to adapt — "all", "attention", or "ffn". - num_adapter_layers: How many transformer layers (from the top) to adapt. - 0 means all 20 layers. E.g. 4 means only layers 16-19 get adapters. - The advisor recommends 2–4 for financial data to avoid overfitting. - train_output_head: Whether to also unfreeze and train the output - projection heads (point + quantile). - - learning_rate: Peak learning rate for AdamW. - weight_decay: L2 regularization coefficient. - num_epochs: Number of training epochs. - batch_size: Per-device batch size. - gradient_clip_norm: Max gradient norm for clipping. - warmup_ratio: Fraction of total steps used for linear warmup. - - context_len: Context window length (padded up to a multiple of 32). - horizon_len: Prediction horizon (must be <= 128 for single-step training). - - use_quantile_loss: Whether to add pinball loss on quantile channels. - quantile_loss_weight: Relative weight of the quantile loss term. - - mixed_precision: AMP dtype — "no", "fp16", or "bf16". - gradient_checkpointing: Trade compute for memory in the transformer stack. - - use_wandb: Enable Weights & Biases logging (rank-0 only). - wandb_project: W&B project name. - log_every_n_steps: Console / W&B logging frequency. - - checkpoint_dir: Directory for adapter checkpoints. - save_every_n_epochs: Checkpoint save frequency. - early_stopping_patience: Epochs without val-loss improvement before stop. - - num_workers: DataLoader workers per process. - seed: Random seed for reproducibility. - """ - - # --- Adapter --- - adapter_type: Literal["lora", "dora"] = "lora" - lora_rank: int = 8 - lora_alpha: float = 16.0 - lora_dropout: float = 0.0 - target_modules: Literal["all", "attention", "ffn"] = "all" - num_adapter_layers: int = 0 # 0 = all 20 layers; N > 0 = only last N layers - train_output_head: bool = False - - # --- Optimiser --- - learning_rate: float = 1e-4 - weight_decay: float = 0.01 - num_epochs: int = 10 - batch_size: int = 32 - gradient_clip_norm: float = 1.0 - warmup_ratio: float = 0.05 - - # --- Data --- - context_len: int = 512 - horizon_len: int = 128 - - # --- Loss --- - use_quantile_loss: bool = False - quantile_loss_weight: float = 0.5 - - # --- Performance --- - mixed_precision: Literal["no", "fp16", "bf16"] = "no" - gradient_checkpointing: bool = False - - # --- Logging --- - use_wandb: bool = False - wandb_project: str = "timesfm-2.5-peft" - log_every_n_steps: int = 50 - - # --- Checkpointing --- - checkpoint_dir: str = "./peft_checkpoints" - save_every_n_epochs: int = 1 - early_stopping_patience: int = 5 - - # --- Misc --- - num_workers: int = 4 - seed: int = 42 diff --git a/peft/data.py b/peft/data.py deleted file mode 100644 index 34d6dae..0000000 --- a/peft/data.py +++ /dev/null @@ -1,150 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Time-series dataset for fine-tuning TimesFM 2.5.""" - -import math -from typing import List, Optional, Sequence, Union - -import numpy as np -import pandas as pd -import torch -from torch.utils.data import Dataset - - -class TimeSeriesDataset(Dataset): - """Sliding-window dataset that produces (context, mask, target) tuples. - - Accepts data in several formats: - - * **list of arrays** — each element is a 1-D NumPy array or Python list - representing a single time series. - * **long-format DataFrame** — columns ``[id_col, value_col]`` where each - unique ``id_col`` identifies a series. - * **wide-format DataFrame** — every numeric column is treated as an - independent time series. - - For each series the dataset generates sliding windows of length - ``context_len + horizon_len`` with the given ``stride``. Series shorter - than the window are left-padded with zeros and masked. - - Args: - data: Time-series data (see above). - context_len: Context (input) length. Will be rounded up to a multiple - of ``patch_len`` (32). - horizon_len: Prediction horizon. Must be ≤ 128. - stride: Step size between consecutive windows. - patch_len: Patch size used by the model (default 32). - id_col: Column name for series identifier (long-format DataFrames). - value_col: Column name for values (long-format DataFrames). - """ - - PATCH_LEN = 32 - MAX_HORIZON = 128 - - def __init__( - self, - data: Union[List[np.ndarray], pd.DataFrame], - context_len: int = 512, - horizon_len: int = 128, - stride: int = 1, - patch_len: int = PATCH_LEN, - id_col: Optional[str] = None, - value_col: Optional[str] = None, - ): - if horizon_len > self.MAX_HORIZON: - raise ValueError( - f"horizon_len={horizon_len} exceeds the single-step maximum of " - f"{self.MAX_HORIZON}. Use a shorter horizon for fine-tuning; the " - f"model handles longer horizons via autoregressive decoding at " - f"inference time." - ) - - self.patch_len = patch_len - # Round context_len up to a multiple of patch_len. - self.context_len = math.ceil(context_len / patch_len) * patch_len - self.horizon_len = horizon_len - self.window_len = self.context_len + horizon_len - - self.series: List[np.ndarray] = self._parse_data(data, id_col, value_col) - self.windows = self._build_windows(stride) - - # -- Data parsing -------------------------------------------------------- - - @staticmethod - def _parse_data( - data: Union[List[np.ndarray], pd.DataFrame], - id_col: Optional[str], - value_col: Optional[str], - ) -> List[np.ndarray]: - if isinstance(data, pd.DataFrame): - if id_col is not None and value_col is not None: - # Long format. - return [ - grp[value_col].to_numpy(dtype=np.float32) - for _, grp in data.groupby(id_col, sort=False) - ] - # Wide format — each numeric column is a series. - return [ - data[c].to_numpy(dtype=np.float32) - for c in data.select_dtypes(include="number").columns - ] - # List / sequence of arrays. - return [np.asarray(s, dtype=np.float32) for s in data] - - def _build_windows(self, stride: int) -> List[tuple]: - windows = [] - for sidx, series in enumerate(self.series): - slen = len(series) - if slen < self.window_len: - # Single (padded) window. - windows.append((sidx, 0, slen)) - else: - for start in range(0, slen - self.window_len + 1, stride): - windows.append((sidx, start, start + self.window_len)) - return windows - - # -- torch Dataset interface --------------------------------------------- - - def __len__(self) -> int: - return len(self.windows) - - def __getitem__(self, idx: int): - sidx, start, end = self.windows[idx] - raw = self.series[sidx][start:end] - - if len(raw) < self.window_len: - # Left-pad context; target uses whatever tail is available. - available_ctx = max(0, len(raw) - self.horizon_len) - target = raw[available_ctx:].copy() - if len(target) < self.horizon_len: - target = np.pad(target, (0, self.horizon_len - len(target))) - - ctx_raw = raw[:available_ctx] - pad_len = self.context_len - len(ctx_raw) - context = np.pad(ctx_raw, (pad_len, 0)).astype(np.float32) - mask = np.zeros(self.context_len, dtype=bool) - mask[:pad_len] = True - else: - context = raw[: self.context_len].astype(np.float32) - mask = np.zeros(self.context_len, dtype=bool) - target = raw[self.context_len : self.context_len + self.horizon_len].astype( - np.float32 - ) - - return ( - torch.from_numpy(context), - torch.from_numpy(mask), - torch.from_numpy(target), - ) diff --git a/peft/finetune.py b/peft/finetune.py deleted file mode 100644 index 70d379c..0000000 --- a/peft/finetune.py +++ /dev/null @@ -1,252 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""CLI entry-point for TimesFM 2.5 PEFT fine-tuning. - -Single-GPU:: - - python peft/finetune.py --data_path data.csv --value_col y - -Multi-GPU (4 GPUs):: - - torchrun --nproc_per_node=4 peft/finetune.py --data_path data.csv --value_col y -""" - -import argparse -import logging -import sys - -import numpy as np -import pandas as pd - -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(levelname)s] %(name)s — %(message)s", - datefmt="%H:%M:%S", -) -logger = logging.getLogger("peft.finetune") - - -def parse_args(argv=None): - p = argparse.ArgumentParser( - description="Fine-tune TimesFM 2.5 with LoRA / DoRA (multi-GPU ready)." - ) - - # -- Model --------------------------------------------------------------- - g = p.add_argument_group("Model") - g.add_argument( - "--model_id", - default="google/timesfm-2.5-200m-pytorch", - help="HuggingFace repo-id or local directory for the base model.", - ) - - # -- Data ---------------------------------------------------------------- - g = p.add_argument_group("Data") - g.add_argument("--data_path", required=True, help="Path to a CSV file.") - g.add_argument( - "--id_col", - default=None, - help="Column identifying individual time series (long format).", - ) - g.add_argument( - "--value_col", - default=None, - help="Column with the values to forecast (long format).", - ) - g.add_argument("--context_len", type=int, default=512) - g.add_argument( - "--horizon_len", - type=int, - default=128, - help="Prediction horizon (max 128 for single-step training).", - ) - g.add_argument( - "--stride", - type=int, - default=32, - help="Stride for the sliding-window dataset.", - ) - g.add_argument( - "--val_split", - type=float, - default=0.2, - help="Fraction of each series reserved for validation.", - ) - - # -- Adapter ------------------------------------------------------------- - g = p.add_argument_group("Adapter") - g.add_argument( - "--adapter_type", - choices=["lora", "dora"], - default="lora", - ) - g.add_argument("--lora_rank", type=int, default=8) - g.add_argument("--lora_alpha", type=float, default=16.0) - g.add_argument("--lora_dropout", type=float, default=0.0) - g.add_argument( - "--target_modules", - choices=["all", "attention", "ffn"], - default="all", - ) - g.add_argument( - "--num_adapter_layers", - type=int, - default=0, - help="Only adapt the last N transformer layers (0 = all 20). " - "Advisor recommends 2-4 for financial data.", - ) - g.add_argument( - "--train_output_head", - action="store_true", - help="Also train the output projection heads.", - ) - - # -- Training ------------------------------------------------------------ - g = p.add_argument_group("Training") - g.add_argument("--num_epochs", type=int, default=10) - g.add_argument("--batch_size", type=int, default=32) - g.add_argument("--learning_rate", type=float, default=1e-4) - g.add_argument("--weight_decay", type=float, default=0.01) - g.add_argument("--gradient_clip_norm", type=float, default=1.0) - g.add_argument("--warmup_ratio", type=float, default=0.05) - g.add_argument( - "--mixed_precision", - choices=["no", "fp16", "bf16"], - default="no", - ) - g.add_argument("--gradient_checkpointing", action="store_true") - g.add_argument("--use_quantile_loss", action="store_true") - g.add_argument("--quantile_loss_weight", type=float, default=0.5) - - # -- Logging / checkpointing -------------------------------------------- - g = p.add_argument_group("Logging") - g.add_argument("--use_wandb", action="store_true") - g.add_argument("--wandb_project", default="timesfm-2.5-peft") - g.add_argument("--log_every_n_steps", type=int, default=50) - g.add_argument("--checkpoint_dir", default="./peft_checkpoints") - g.add_argument("--save_every_n_epochs", type=int, default=1) - g.add_argument("--early_stopping_patience", type=int, default=5) - - # -- Misc ---------------------------------------------------------------- - g = p.add_argument_group("Misc") - g.add_argument("--num_workers", type=int, default=4) - g.add_argument("--seed", type=int, default=42) - - return p.parse_args(argv) - - -def main(argv=None): - args = parse_args(argv) - - # Lazy imports so --help is fast. - from timesfm.timesfm_2p5.timesfm_2p5_torch import TimesFM_2p5_200M_torch - - from .config import PEFTConfig - from .data import TimeSeriesDataset - from .trainer import PEFTTrainer - - # -- Load model ---------------------------------------------------------- - logger.info("Loading base model from %s …", args.model_id) - wrapper = TimesFM_2p5_200M_torch.from_pretrained( - args.model_id, torch_compile=False - ) - - # -- Build config -------------------------------------------------------- - config = PEFTConfig( - adapter_type=args.adapter_type, - lora_rank=args.lora_rank, - lora_alpha=args.lora_alpha, - lora_dropout=args.lora_dropout, - target_modules=args.target_modules, - num_adapter_layers=args.num_adapter_layers, - train_output_head=args.train_output_head, - learning_rate=args.learning_rate, - weight_decay=args.weight_decay, - num_epochs=args.num_epochs, - batch_size=args.batch_size, - gradient_clip_norm=args.gradient_clip_norm, - warmup_ratio=args.warmup_ratio, - context_len=args.context_len, - horizon_len=args.horizon_len, - use_quantile_loss=args.use_quantile_loss, - quantile_loss_weight=args.quantile_loss_weight, - mixed_precision=args.mixed_precision, - gradient_checkpointing=args.gradient_checkpointing, - use_wandb=args.use_wandb, - wandb_project=args.wandb_project, - log_every_n_steps=args.log_every_n_steps, - checkpoint_dir=args.checkpoint_dir, - save_every_n_epochs=args.save_every_n_epochs, - early_stopping_patience=args.early_stopping_patience, - num_workers=args.num_workers, - seed=args.seed, - ) - - # -- Load data ----------------------------------------------------------- - logger.info("Reading data from %s …", args.data_path) - df = pd.read_csv(args.data_path) - - # Parse series from DataFrame. - if args.id_col and args.value_col: - all_series = [ - grp[args.value_col].to_numpy(dtype=np.float32) - for _, grp in df.groupby(args.id_col, sort=False) - ] - elif args.value_col: - all_series = [df[args.value_col].to_numpy(dtype=np.float32)] - else: - all_series = [ - df[c].to_numpy(dtype=np.float32) - for c in df.select_dtypes(include="number").columns - ] - - # Train / val split (tail of each series → val). - train_series, val_series = [], [] - for s in all_series: - split_idx = max(1, int(len(s) * (1 - args.val_split))) - train_series.append(s[:split_idx]) - val_series.append(s[split_idx - config.context_len :]) # overlap for context - - train_ds = TimeSeriesDataset( - train_series, - context_len=config.context_len, - horizon_len=config.horizon_len, - stride=args.stride, - ) - val_ds = TimeSeriesDataset( - val_series, - context_len=config.context_len, - horizon_len=config.horizon_len, - stride=config.horizon_len, # non-overlapping for val - ) - - logger.info( - "Dataset: %d train windows, %d val windows", len(train_ds), len(val_ds) - ) - - # -- Train --------------------------------------------------------------- - trainer = PEFTTrainer(wrapper.model, config) - history = trainer.fit(train_ds, val_ds) - - # -- Save final adapter -------------------------------------------------- - final_path = f"{config.checkpoint_dir}/final_adapter.safetensors" - trainer.save_adapter(final_path) - logger.info("Final adapter saved → %s", final_path) - - return history - - -if __name__ == "__main__": - main() diff --git a/peft/finetune.sh b/peft/finetune.sh deleted file mode 100644 index 37a16e7..0000000 --- a/peft/finetune.sh +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env bash -# ============================================================================ -# Example launch script for TimesFM 2.5 PEFT fine-tuning. -# -# Single GPU: -# bash peft/finetune.sh -# -# Multi-GPU (e.g. 4 GPUs): -# NUM_GPUS=4 bash peft/finetune.sh -# ============================================================================ - -set -euo pipefail - -NUM_GPUS="${NUM_GPUS:-1}" - -# --- Data ------------------------------------------------------------------- -DATA_PATH="${DATA_PATH:-data.csv}" # path to your CSV -ID_COL="${ID_COL:-}" # series-id column (long format), leave empty for wide -VALUE_COL="${VALUE_COL:-}" # value column (long format), leave empty for wide -CONTEXT_LEN="${CONTEXT_LEN:-512}" -HORIZON_LEN="${HORIZON_LEN:-128}" -STRIDE="${STRIDE:-32}" -VAL_SPLIT="${VAL_SPLIT:-0.2}" - -# --- Adapter ---------------------------------------------------------------- -ADAPTER_TYPE="${ADAPTER_TYPE:-lora}" # lora | dora -LORA_RANK="${LORA_RANK:-8}" -LORA_ALPHA="${LORA_ALPHA:-16}" -TARGET_MODULES="${TARGET_MODULES:-all}" # all | attention | ffn -NUM_ADAPTER_LAYERS="${NUM_ADAPTER_LAYERS:-4}" # 0=all 20, advisor recommends 2-4 - -# --- Training --------------------------------------------------------------- -NUM_EPOCHS="${NUM_EPOCHS:-10}" -BATCH_SIZE="${BATCH_SIZE:-32}" -LR="${LR:-1e-4}" -MIXED_PRECISION="${MIXED_PRECISION:-no}" # no | fp16 | bf16 - -# --- Logging / checkpoint --------------------------------------------------- -CHECKPOINT_DIR="${CHECKPOINT_DIR:-./peft_checkpoints}" - -# ============================================================================ - -CMD_ARGS=( - peft/finetune.py - --data_path "$DATA_PATH" - --context_len "$CONTEXT_LEN" - --horizon_len "$HORIZON_LEN" - --stride "$STRIDE" - --val_split "$VAL_SPLIT" - --adapter_type "$ADAPTER_TYPE" - --lora_rank "$LORA_RANK" - --lora_alpha "$LORA_ALPHA" - --target_modules "$TARGET_MODULES" - --num_adapter_layers "$NUM_ADAPTER_LAYERS" - --train_output_head - --num_epochs "$NUM_EPOCHS" - --batch_size "$BATCH_SIZE" - --learning_rate "$LR" - --mixed_precision "$MIXED_PRECISION" - --checkpoint_dir "$CHECKPOINT_DIR" -) - -# Optional columns. -[[ -n "$ID_COL" ]] && CMD_ARGS+=(--id_col "$ID_COL") -[[ -n "$VALUE_COL" ]] && CMD_ARGS+=(--value_col "$VALUE_COL") - -if [[ "$NUM_GPUS" -gt 1 ]]; then - echo "Launching multi-GPU training on $NUM_GPUS GPUs …" - torchrun --nproc_per_node="$NUM_GPUS" "${CMD_ARGS[@]}" -else - echo "Launching single-GPU training …" - python "${CMD_ARGS[@]}" -fi diff --git a/peft/trainer.py b/peft/trainer.py deleted file mode 100644 index 4b3ccc1..0000000 --- a/peft/trainer.py +++ /dev/null @@ -1,578 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Multi-GPU PEFT trainer for TimesFM 2.5. - -Supports: -* LoRA / DoRA adapters (via ``adapters.inject_adapters``) -* PyTorch DDP multi-GPU (``torchrun``) -* Mixed-precision training (fp16 / bf16) -* Gradient checkpointing -* Cosine-with-warmup LR schedule -* Early stopping & adapter-only checkpointing -* Optional W&B logging -""" - -import logging -import math -import os -import time -from typing import Dict, Optional - -import torch -import torch.distributed as dist -import torch.nn as nn -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data import DataLoader, Dataset - -from .adapters import ( - inject_adapters, - load_adapter_weights, - merge_adapters, - save_adapter_weights, -) -from .config import PEFTConfig - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# Utility: access the raw model under potential DDP wrapper -# --------------------------------------------------------------------------- - - -def _unwrap(model: nn.Module) -> nn.Module: - return model.module if isinstance(model, DDP) else model - - -# --------------------------------------------------------------------------- -# Training forward — replicates the model's inference preprocessing so that -# gradients flow through the transformer + adapter parameters. -# --------------------------------------------------------------------------- - - -def _training_forward( - model: nn.Module, - context: torch.Tensor, - masks: torch.Tensor, - gradient_checkpointing: bool = False, -): - """Run a differentiable forward pass for fine-tuning. - - This mirrors the pre-processing that ``TimesFM_2p5_200M_torch_module.decode`` - performs (patching → RevIN → transformer → output projections → un-RevIN), - but without ``torch.no_grad()`` and without KV-cache / AR decoding. - - Args: - model: The (possibly DDP-wrapped) model. - context: ``(B, context_len)`` raw time-series values. - masks: ``(B, context_len)`` boolean mask (``True`` = padding). - gradient_checkpointing: Use activation checkpointing on transformer layers. - - Returns: - ``(output_ts, output_qs)`` — *un-normalised* predictions, each of shape - ``(B, N, output_patch_len, num_quantiles)``. - """ - from timesfm.torch.util import revin, update_running_stats - - raw = _unwrap(model) - B = context.shape[0] - p = raw.p # 32 - o = raw.o # 128 - q = raw.q # 10 - os_ = raw.os # 1024 - - # 1. Patch ---------------------------------------------------------------- - patched = context.reshape(B, -1, p) # (B, N, 32) - patched_masks = masks.reshape(B, -1, p) # (B, N, 32) - N = patched.shape[1] - - # 2. Running RevIN stats -------------------------------------------------- - n = torch.zeros(B, device=context.device) - mu = torch.zeros(B, device=context.device) - sigma = torch.zeros(B, device=context.device) - patch_mus, patch_sigmas = [], [] - for i in range(N): - (n, mu, sigma), _ = update_running_stats( - n, mu, sigma, patched[:, i], patched_masks[:, i] - ) - patch_mus.append(mu) - patch_sigmas.append(sigma) - ctx_mu = torch.stack(patch_mus, dim=1) # (B, N) - ctx_sigma = torch.stack(patch_sigmas, dim=1) # (B, N) - - # 3. Normalise + mask ----------------------------------------------------- - normed = revin(patched, ctx_mu, ctx_sigma, reverse=False) - normed = torch.where(patched_masks, 0.0, normed) - - # 4. Tokenise ------------------------------------------------------------- - tok_in = torch.cat([normed, patched_masks.to(normed.dtype)], dim=-1) - embeddings = raw.tokenizer(tok_in) # (B, N, model_dims) - - # 5. Transformer stack ---------------------------------------------------- - patch_mask = patched_masks[..., -1] # (B, N) per-patch mask - x = embeddings - for layer in raw.stacked_xf: - if gradient_checkpointing: - x = torch.utils.checkpoint.checkpoint( - _transformer_layer_fn, layer, x, patch_mask, use_reentrant=False - ) - else: - x, _ = layer(x, patch_mask) - - # 6. Output projections --------------------------------------------------- - normed_ts = raw.output_projection_point(x) # (B, N, o*q) - normed_qs = raw.output_projection_quantiles(x) # (B, N, os*q) - - # 7. Un-normalise --------------------------------------------------------- - output_ts = revin( - normed_ts.reshape(B, N, o, q), ctx_mu, ctx_sigma, reverse=True - ) - output_qs = revin( - normed_qs.reshape(B, N, os_, q), ctx_mu, ctx_sigma, reverse=True - ) - - return output_ts, output_qs - - -def _transformer_layer_fn( - layer: nn.Module, x: torch.Tensor, mask: torch.Tensor -) -> torch.Tensor: - out, _ = layer(x, mask) - return out - - -# --------------------------------------------------------------------------- -# Loss computation -# --------------------------------------------------------------------------- - -_DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] - - -def _quantile_loss( - pred: torch.Tensor, target: torch.Tensor, tau: float -) -> torch.Tensor: - """Pinball (quantile) loss.""" - diff = target - pred - return 2.0 * torch.where(diff >= 0, tau * diff, (tau - 1.0) * diff) - - -def _compute_loss( - output_ts: torch.Tensor, - target: torch.Tensor, - horizon_len: int, - use_quantile_loss: bool = False, - quantile_loss_weight: float = 0.5, -): - """Compute MSE (+ optional quantile) loss on the last-patch prediction. - - Args: - output_ts: ``(B, N, 128, 10)`` denormalised forecast tensor. - target: ``(B, horizon_len)`` ground-truth future values. - horizon_len: Number of steps to compare. - use_quantile_loss: Add pinball loss on quantile channels. - quantile_loss_weight: Relative weight of the quantile term. - - Returns: - Scalar loss tensor. - """ - # Last input-patch → first horizon_len steps, median channel (idx 5). - pred_median = output_ts[:, -1, :horizon_len, 5] # (B, H) - loss = torch.nn.functional.mse_loss(pred_median, target) - - if use_quantile_loss: - q_loss = torch.tensor(0.0, device=loss.device) - for qi, tau in enumerate(_DEFAULT_QUANTILES): - pred_q = output_ts[:, -1, :horizon_len, qi + 1] # channels 1-9 - q_loss = q_loss + _quantile_loss(pred_q, target, tau).mean() - loss = loss + quantile_loss_weight * q_loss - - return loss - - -# --------------------------------------------------------------------------- -# PEFTTrainer -# --------------------------------------------------------------------------- - - -class PEFTTrainer: - """Production-grade PEFT trainer for TimesFM 2.5 (PyTorch). - - Typical usage:: - - from timesfm.timesfm_2p5.timesfm_2p5_torch import TimesFM_2p5_200M_torch - model = TimesFM_2p5_200M_torch.from_pretrained( - "google/timesfm-2.5-200m-pytorch", torch_compile=False - ) - trainer = PEFTTrainer(model.model, PEFTConfig(...)) - history = trainer.fit(train_dataset, val_dataset) - trainer.save_adapter("./adapter/adapter.safetensors") - """ - - def __init__(self, model: nn.Module, config: PEFTConfig): - self.config = config - self._setup_distributed() - self._setup_seed(config.seed) - - # Inject adapters and freeze base weights. - inject_adapters(model, config) - - # Move to device. - self.device = torch.device( - f"cuda:{self.local_rank}" if torch.cuda.is_available() else "cpu" - ) - model.to(self.device) - - self.raw_model = model - if self.is_distributed: - self.model = DDP(model, device_ids=[self.local_rank]) - else: - self.model = model - - # Optimizer — only trainable (adapter) parameters. - trainable = [p for p in model.parameters() if p.requires_grad] - self.optimizer = torch.optim.AdamW( - trainable, - lr=config.learning_rate, - weight_decay=config.weight_decay, - ) - - # AMP setup. - self.autocast_dtype = { - "fp16": torch.float16, - "bf16": torch.bfloat16, - "no": None, - }[config.mixed_precision] - self.scaler = ( - torch.amp.GradScaler("cuda") - if config.mixed_precision == "fp16" - else None - ) - - # Logging. - self._wandb = None - if config.use_wandb and self.is_main: - try: - import wandb - - wandb.init(project=config.wandb_project, config=config.__dict__) - self._wandb = wandb - except ImportError: - logger.warning("wandb not installed — skipping W&B logging.") - - n_trainable = sum(p.numel() for p in trainable) - n_total = sum(p.numel() for p in model.parameters()) - if self.is_main: - logger.info( - "Trainable parameters: %s / %s (%.2f%%)", - f"{n_trainable:,}", - f"{n_total:,}", - 100 * n_trainable / n_total, - ) - - # -- Distributed setup --------------------------------------------------- - - def _setup_distributed(self): - self.local_rank = int(os.environ.get("LOCAL_RANK", 0)) - self.world_size = int(os.environ.get("WORLD_SIZE", 1)) - self.is_distributed = self.world_size > 1 - self.is_main = self.local_rank == 0 - - if self.is_distributed and not dist.is_initialized(): - dist.init_process_group("nccl") - torch.cuda.set_device(self.local_rank) - - @staticmethod - def _setup_seed(seed: int): - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - - # -- Data loaders -------------------------------------------------------- - - def _make_loader(self, dataset: Dataset, is_train: bool) -> DataLoader: - cfg = self.config - sampler = None - shuffle = is_train - if self.is_distributed: - sampler = torch.utils.data.distributed.DistributedSampler( - dataset, - num_replicas=self.world_size, - rank=self.local_rank, - shuffle=is_train, - ) - shuffle = False - - return DataLoader( - dataset, - batch_size=cfg.batch_size, - shuffle=shuffle, - sampler=sampler, - num_workers=cfg.num_workers, - pin_memory=True, - drop_last=is_train, - ) - - # -- LR schedule --------------------------------------------------------- - - def _build_scheduler(self, total_steps: int): - warmup_steps = int(self.config.warmup_ratio * total_steps) - - def lr_lambda(step: int) -> float: - if step < warmup_steps: - return step / max(1, warmup_steps) - progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) - return 0.5 * (1.0 + math.cos(math.pi * progress)) - - return torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda) - - # -- Training / validation ----------------------------------------------- - - def _train_step(self, batch): - context, masks, target = [t.to(self.device, non_blocking=True) for t in batch] - - ctx_manager = ( - torch.amp.autocast("cuda", dtype=self.autocast_dtype) - if self.autocast_dtype is not None - else _nullcontext() - ) - - with ctx_manager: - output_ts, _ = _training_forward( - self.model, - context, - masks, - gradient_checkpointing=self.config.gradient_checkpointing, - ) - loss = _compute_loss( - output_ts, - target, - self.config.horizon_len, - use_quantile_loss=self.config.use_quantile_loss, - quantile_loss_weight=self.config.quantile_loss_weight, - ) - - self.optimizer.zero_grad(set_to_none=True) - if self.scaler is not None: - self.scaler.scale(loss).backward() - self.scaler.unscale_(self.optimizer) - nn.utils.clip_grad_norm_( - (p for p in self.raw_model.parameters() if p.requires_grad), - self.config.gradient_clip_norm, - ) - self.scaler.step(self.optimizer) - self.scaler.update() - else: - loss.backward() - nn.utils.clip_grad_norm_( - (p for p in self.raw_model.parameters() if p.requires_grad), - self.config.gradient_clip_norm, - ) - self.optimizer.step() - - return loss.detach() - - @torch.no_grad() - def _validate(self, val_loader: DataLoader) -> float: - self.model.eval() - total_loss = 0.0 - n = 0 - - for batch in val_loader: - context, masks, target = [t.to(self.device, non_blocking=True) for t in batch] - - ctx_manager = ( - torch.amp.autocast("cuda", dtype=self.autocast_dtype) - if self.autocast_dtype is not None - else _nullcontext() - ) - with ctx_manager: - output_ts, _ = _training_forward( - self.model, - context, - masks, - gradient_checkpointing=False, - ) - loss = _compute_loss( - output_ts, - target, - self.config.horizon_len, - use_quantile_loss=self.config.use_quantile_loss, - quantile_loss_weight=self.config.quantile_loss_weight, - ) - total_loss += loss.item() - n += 1 - - avg = total_loss / max(n, 1) - if self.is_distributed: - t = torch.tensor(avg, device=self.device) - dist.all_reduce(t, op=dist.ReduceOp.SUM) - avg = (t / self.world_size).item() - return avg - - # -- Main loop ----------------------------------------------------------- - - def fit( - self, - train_dataset: Dataset, - val_dataset: Optional[Dataset] = None, - ) -> Dict[str, list]: - """Run the full training loop. - - Args: - train_dataset: Training data (``TimeSeriesDataset`` or any - ``Dataset`` returning ``(context, mask, target)`` tensors). - val_dataset: Optional validation data. - - Returns: - Dictionary with ``train_loss``, ``val_loss``, ``lr`` histories. - """ - cfg = self.config - train_loader = self._make_loader(train_dataset, is_train=True) - val_loader = ( - self._make_loader(val_dataset, is_train=False) if val_dataset else None - ) - - steps_per_epoch = len(train_loader) - total_steps = cfg.num_epochs * steps_per_epoch - scheduler = self._build_scheduler(total_steps) - - history: Dict[str, list] = {"train_loss": [], "val_loss": [], "lr": []} - best_val_loss = float("inf") - patience_counter = 0 - global_step = 0 - - if self.is_main: - logger.info( - "Training: %d epochs, %d steps/epoch, %d total steps", - cfg.num_epochs, - steps_per_epoch, - total_steps, - ) - - for epoch in range(cfg.num_epochs): - self.model.train() - if self.is_distributed: - train_loader.sampler.set_epoch(epoch) - - epoch_loss = 0.0 - t0 = time.time() - - for step, batch in enumerate(train_loader): - loss = self._train_step(batch) - scheduler.step() - global_step += 1 - epoch_loss += loss.item() - - if self.is_main and global_step % cfg.log_every_n_steps == 0: - lr = scheduler.get_last_lr()[0] - logger.info( - "[epoch %d step %d/%d] loss=%.5f lr=%.2e", - epoch + 1, - step + 1, - steps_per_epoch, - loss.item(), - lr, - ) - if self._wandb is not None: - self._wandb.log( - {"train/loss": loss.item(), "train/lr": lr}, - step=global_step, - ) - - avg_train_loss = epoch_loss / max(steps_per_epoch, 1) - history["train_loss"].append(avg_train_loss) - history["lr"].append(scheduler.get_last_lr()[0]) - - # Validation. - val_loss = None - if val_loader is not None: - val_loss = self._validate(val_loader) - history["val_loss"].append(val_loss) - - elapsed = time.time() - t0 - if self.is_main: - msg = ( - f"[Epoch {epoch + 1}/{cfg.num_epochs}] " - f"train_loss={avg_train_loss:.5f}" - ) - if val_loss is not None: - msg += f" val_loss={val_loss:.5f}" - msg += f" ({elapsed:.1f}s)" - logger.info(msg) - if self._wandb is not None: - metrics = {"epoch": epoch + 1, "train/epoch_loss": avg_train_loss} - if val_loss is not None: - metrics["val/loss"] = val_loss - self._wandb.log(metrics, step=global_step) - - # Checkpoint + early stopping. - if val_loss is not None and val_loss < best_val_loss: - best_val_loss = val_loss - patience_counter = 0 - if self.is_main and cfg.save_every_n_epochs > 0: - ckpt_path = os.path.join(cfg.checkpoint_dir, "best_adapter.safetensors") - save_adapter_weights(self.raw_model, ckpt_path) - logger.info(" ↳ Saved best adapter → %s", ckpt_path) - elif val_loss is not None: - patience_counter += 1 - if patience_counter >= cfg.early_stopping_patience: - if self.is_main: - logger.info("Early stopping triggered (patience=%d).", cfg.early_stopping_patience) - break - - if ( - self.is_main - and cfg.save_every_n_epochs > 0 - and (epoch + 1) % cfg.save_every_n_epochs == 0 - ): - ep_path = os.path.join( - cfg.checkpoint_dir, f"adapter_epoch{epoch + 1}.safetensors" - ) - save_adapter_weights(self.raw_model, ep_path) - - # Cleanup. - if self.is_distributed: - dist.destroy_process_group() - if self._wandb is not None: - self._wandb.finish() - - return history - - # -- Convenience wrappers ------------------------------------------------ - - def save_adapter(self, path: str) -> None: - """Save adapter weights to *path* (safetensors format).""" - save_adapter_weights(self.raw_model, path) - - def load_adapter(self, path: str) -> None: - """Load adapter weights from *path*.""" - load_adapter_weights(self.raw_model, path) - - def merge_adapter(self) -> nn.Module: - """Fold adapter weights into base model and return the raw model.""" - return merge_adapters(self.raw_model) - - -# --------------------------------------------------------------------------- -# Tiny helper to replace contextlib.nullcontext (available ≥3.7 but -# with async generics issues) for the AMP autocast conditional. -# --------------------------------------------------------------------------- - -class _nullcontext: - def __enter__(self): - return None - - def __exit__(self, *_): - return False diff --git a/timesfm-forecasting/examples/finetuning/README.md b/timesfm-forecasting/examples/finetuning/README.md new file mode 100644 index 0000000..6097f59 --- /dev/null +++ b/timesfm-forecasting/examples/finetuning/README.md @@ -0,0 +1,102 @@ +# Fine-Tuning TimesFM 2.5 with LoRA + +Parameter-efficient fine-tuning of +[TimesFM 2.5](https://huggingface.co/google/timesfm-2.5-200m-transformers) +using **HuggingFace Transformers** and **PEFT (LoRA)**. + +This approach is based on the fine-tuning workflow by +[@kashif](https://github.com/kashif) at HuggingFace +([notebook](https://github.com/huggingface/notebooks/blob/main/examples/timesfm2_5.ipynb)). + +## How It Works + +TimesFM 2.5 is available as a standard +[Transformers](https://github.com/huggingface/transformers) model +(`TimesFm2_5ModelForPrediction`). This means it supports the full Transformers +ecosystem out of the box, including: + +- **PEFT adapters** — LoRA, QLoRA, etc. via the + [`peft`](https://github.com/huggingface/peft) library +- **All attention backends** — eager, SDPA, Flash Attention 2/3, Flex Attention +- **Standard `from_pretrained` / `save_pretrained` workflow** + +The model's forward pass natively computes a training loss when `future_values` +are provided, so fine-tuning requires nothing more than a standard PyTorch +training loop. + +## Quick Start + +### Install + +```bash +pip install transformers accelerate peft pandas pyarrow scikit-learn +``` + +### Train + +```bash +# Fine-tune with default settings on the retail sales dataset +python finetune_lora.py + +# Custom hyperparameters +python finetune_lora.py \ + --epochs 20 \ + --batch_size 64 \ + --lr 5e-5 \ + --lora_r 8 \ + --lora_alpha 16 \ + --context_len 64 \ + --horizon_len 13 \ + --output_dir my-retail-adapter +``` + +### Evaluate + +```bash +# Evaluate a previously trained adapter (skip training) +python finetune_lora.py --eval_only --output_dir timesfm2_5-retail-lora +``` + +## Key Concepts + +### No External Normalisation + +TimesFM 2.5 applies its own internal instance normalisation (RevIN). **Do not** +normalise your data externally — feed raw values and let the model handle it. + +### Random Window Sampling + +Following [Chronos-2](https://github.com/amazon-science/chronos-forecasting), +each training example is a random `(context, horizon)` window sliced from one of +the input series. This is more data-efficient than always using the same +fixed window per series. + +### LoRA Target Modules + +Using `target_modules="all-linear"` applies LoRA to every linear layer in the +model. With `r=4` this adds only ~0.6% trainable parameters (~1.4M out of +~232M), which is enough to meaningfully adapt the model to a new domain. + +## CLI Options + +| Flag | Default | Description | +|------|---------|-------------| +| `--model_id` | `google/timesfm-2.5-200m-transformers` | HuggingFace model ID | +| `--context_len` | `64` | Context length for training windows | +| `--horizon_len` | `13` | Forecast horizon in time steps | +| `--epochs` | `10` | Training epochs | +| `--batch_size` | `32` | Batch size | +| `--lr` | `1e-4` | Learning rate | +| `--lora_r` | `4` | LoRA rank | +| `--lora_alpha` | `8` | LoRA alpha | +| `--lora_dropout` | `0.05` | LoRA dropout | +| `--num_samples` | `5000` | Random training windows to pre-sample | +| `--output_dir` | `timesfm2_5-retail-lora` | Where to save the adapter | +| `--seed` | `42` | Random seed | +| `--eval_only` | — | Skip training; evaluate existing adapter | + +## Acknowledgements + +The Transformers integration and fine-tuning approach were developed by +[@kashif](https://github.com/kashif) at HuggingFace. See the original notebook: + diff --git a/timesfm-forecasting/examples/finetuning/finetune_lora.py b/timesfm-forecasting/examples/finetuning/finetune_lora.py new file mode 100644 index 0000000..b9dcead --- /dev/null +++ b/timesfm-forecasting/examples/finetuning/finetune_lora.py @@ -0,0 +1,446 @@ +#!/usr/bin/env python3 +"""Fine-tune TimesFM 2.5 with LoRA using HuggingFace Transformers + PEFT. + +This script demonstrates parameter-efficient fine-tuning of TimesFM 2.5 on a +retail demand forecasting dataset (weekly store sales). It uses the HuggingFace +Transformers checkpoint and the standard PEFT library for LoRA adapters. + +The approach is based on the fine-tuning workflow by @kashif at HuggingFace: +https://github.com/huggingface/notebooks/blob/main/examples/timesfm2_5.ipynb + +The dataset is the same one used in the Chronos-2 quickstart notebook. Each +store has ~120 weekly data points. The goal is to forecast the next 13 weeks +(one quarter) of sales per store. + +Requirements: + pip install transformers accelerate peft pandas pyarrow scikit-learn + +Usage: + python finetune_lora.py [OPTIONS] + + Options: + --model_id HuggingFace model ID (default: google/timesfm-2.5-200m-transformers) + --context_len Context length for training windows (default: 64, must be multiple of 32) + --horizon_len Forecast horizon in time steps (default: 13) + --epochs Number of training epochs (default: 10) + --batch_size Training batch size (default: 32) + --lr Learning rate (default: 1e-4) + --lora_r LoRA rank (default: 4) + --lora_alpha LoRA alpha (default: 8) + --lora_dropout LoRA dropout (default: 0.05) + --num_samples Number of random training windows to pre-sample (default: 5000) + --output_dir Directory to save the LoRA adapter (default: timesfm2_5-retail-lora) + --seed Random seed (default: 42) +""" + +import argparse +import logging +import os + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import DataLoader, Dataset + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", +) +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Dataset +# --------------------------------------------------------------------------- + +class TimeSeriesRandomWindowDataset(Dataset): + """Random-window dataset for time series fine-tuning. + + Pre-samples random (series, split-point) windows similar to Chronos-2's + random slicing. Each window has a full *context_len* context (no + zero-padding) to avoid corrupting TimesFM's internal RevIN normalisation + statistics. + + No external normalisation is needed — TimesFM handles instance + normalisation internally. The loss is computed in the original data scale. + """ + + def __init__( + self, + series_list: list[np.ndarray], + context_len: int, + horizon_len: int, + num_samples: int = 5000, + seed: int = 42, + ): + self.series_list = series_list + self.context_len = context_len + self.horizon_len = horizon_len + self.samples: list[tuple[int, int]] = [] + + rng = np.random.default_rng(seed) + min_len = context_len + horizon_len + valid = [i for i, s in enumerate(series_list) if len(s) >= min_len] + if not valid: + raise ValueError( + f"No series long enough for context_len={context_len} + " + f"horizon_len={horizon_len}. Shortest series: " + f"{min(len(s) for s in series_list)}" + ) + + for _ in range(num_samples): + idx = rng.choice(valid) + series = series_list[idx] + max_start = len(series) - min_len + start = rng.integers(0, max_start + 1) + self.samples.append((idx, start)) + + def __len__(self) -> int: + return len(self.samples) + + def __getitem__(self, i: int): + idx, start = self.samples[i] + series = self.series_list[idx] + end = start + self.context_len + self.horizon_len + + context = torch.tensor( + series[start : start + self.context_len], dtype=torch.float32 + ) + target = torch.tensor( + series[start + self.context_len : end], dtype=torch.float32 + ) + return context, target + + +class TimeSeriesLastWindowDataset(Dataset): + """Validation dataset using the last window of each series.""" + + def __init__( + self, + series_list: list[np.ndarray], + context_len: int, + horizon_len: int, + ): + self.items: list[tuple[torch.Tensor, torch.Tensor]] = [] + min_len = context_len + horizon_len + for s in series_list: + if len(s) >= min_len: + ctx = torch.tensor(s[-min_len:-horizon_len], dtype=torch.float32) + tgt = torch.tensor(s[-horizon_len:], dtype=torch.float32) + self.items.append((ctx, tgt)) + + def __len__(self) -> int: + return len(self.items) + + def __getitem__(self, i: int): + return self.items[i] + + +# --------------------------------------------------------------------------- +# Data loading helpers +# --------------------------------------------------------------------------- + +def load_retail_sales( + context_len: int, + horizon_len: int, + num_samples: int, + seed: int, +) -> tuple[TimeSeriesRandomWindowDataset, TimeSeriesLastWindowDataset]: + """Download and prepare the retail sales dataset. + + This is the same dataset used in the Chronos-2 quickstart notebook and + in @kashif's TimesFM 2.5 fine-tuning example. Each store has ~120 weekly + data points; the target column is ``Sales``. + + Returns train dataset and val dataset. + """ + logger.info("Loading retail sales dataset …") + sales_train_df = pd.read_parquet( + "https://autogluon.s3.amazonaws.com/datasets/timeseries/" + "retail_sales/train.parquet" + ) + target = "Sales" + + all_series: list[np.ndarray] = [] + for _, group in sales_train_df.groupby("id"): + values = group[target].values.astype(np.float32) + if len(values) >= context_len + horizon_len: + all_series.append(values) + + logger.info( + "Valid stores: %d (need >= %d data points)", + len(all_series), + context_len + horizon_len, + ) + + train_ds = TimeSeriesRandomWindowDataset( + all_series, context_len, horizon_len, num_samples=num_samples, seed=seed + ) + val_ds = TimeSeriesLastWindowDataset(all_series, context_len, horizon_len) + return train_ds, val_ds + + +# --------------------------------------------------------------------------- +# Training +# --------------------------------------------------------------------------- + +def train(args: argparse.Namespace) -> None: + from peft import LoraConfig, get_peft_model + from transformers import TimesFm2_5ModelForPrediction + + device = "cuda" if torch.cuda.is_available() else "cpu" + logger.info("Using device: %s", device) + + # ------------------------------------------------------------------ + # Load model + # ------------------------------------------------------------------ + logger.info("Loading model: %s", args.model_id) + model = TimesFm2_5ModelForPrediction.from_pretrained( + args.model_id, + torch_dtype=torch.bfloat16, + device_map=device, + ) + horizon_len = args.horizon_len + context_len = min(args.context_len, model.config.context_length) + + # ------------------------------------------------------------------ + # Apply LoRA + # ------------------------------------------------------------------ + lora_config = LoraConfig( + r=args.lora_r, + lora_alpha=args.lora_alpha, + target_modules="all-linear", + lora_dropout=args.lora_dropout, + bias="none", + ) + model = get_peft_model(model, lora_config) + model.print_trainable_parameters() + + # ------------------------------------------------------------------ + # Prepare data + # ------------------------------------------------------------------ + train_ds, val_ds = load_retail_sales( + context_len, horizon_len, num_samples=args.num_samples, seed=args.seed + ) + train_loader = DataLoader( + train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True + ) + val_loader = DataLoader(val_ds, batch_size=args.batch_size) + + logger.info( + "Train samples: %d (%d batches) | Val samples: %d", + len(train_ds), + len(train_loader), + len(val_ds), + ) + + # ------------------------------------------------------------------ + # Optimiser & scheduler + # ------------------------------------------------------------------ + optimizer = torch.optim.AdamW( + model.parameters(), lr=args.lr, weight_decay=0.01 + ) + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, T_max=args.epochs * len(train_loader) + ) + + # ------------------------------------------------------------------ + # Training loop + # ------------------------------------------------------------------ + best_val_loss = float("inf") + + for epoch in range(1, args.epochs + 1): + model.train() + epoch_loss = 0.0 + n_batches = 0 + + for context, target_vals in train_loader: + context = context.to(device) + target_vals = target_vals.to(device) + + outputs = model( + past_values=context, + future_values=target_vals, + forecast_context_len=context_len, + ) + loss = outputs.loss + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + optimizer.step() + optimizer.zero_grad() + scheduler.step() + + epoch_loss += loss.item() + n_batches += 1 + + avg_train_loss = epoch_loss / max(n_batches, 1) + + # Validation + model.eval() + val_loss = 0.0 + val_batches = 0 + with torch.no_grad(): + for context, target_vals in val_loader: + context = context.to(device) + target_vals = target_vals.to(device) + outputs = model( + past_values=context, + future_values=target_vals, + forecast_context_len=context_len, + ) + val_loss += outputs.loss.item() + val_batches += 1 + + avg_val_loss = val_loss / max(val_batches, 1) + + logger.info( + "Epoch %d/%d (%d steps) — train loss: %.4f, val loss: %.4f", + epoch, + args.epochs, + n_batches, + avg_train_loss, + avg_val_loss, + ) + + if avg_val_loss < best_val_loss: + best_val_loss = avg_val_loss + model.save_pretrained(args.output_dir) + logger.info(" ✓ saved best adapter → %s", args.output_dir) + + logger.info("Training complete. Best val loss: %.4f", best_val_loss) + + +# --------------------------------------------------------------------------- +# Evaluation +# --------------------------------------------------------------------------- + +def evaluate(args: argparse.Namespace) -> None: + """Compare zero-shot vs fine-tuned on a subset of stores.""" + from peft import PeftModel + from transformers import TimesFm2_5ModelForPrediction + + device = "cuda" if torch.cuda.is_available() else "cpu" + + logger.info("Loading base model …") + base_model = TimesFm2_5ModelForPrediction.from_pretrained( + args.model_id, + torch_dtype=torch.bfloat16, + device_map=device, + ) + base_model.eval() + horizon_len = args.horizon_len + context_len = min(args.context_len, base_model.config.context_length) + + logger.info("Loading LoRA adapter from %s …", args.output_dir) + ft_model = PeftModel.from_pretrained(base_model, args.output_dir) + ft_model.eval() + + # --- Load data --- + sales_train_df = pd.read_parquet( + "https://autogluon.s3.amazonaws.com/datasets/timeseries/" + "retail_sales/train.parquet" + ) + sales_test_df = pd.read_parquet( + "https://autogluon.s3.amazonaws.com/datasets/timeseries/" + "retail_sales/test.parquet" + ) + target = "Sales" + + store_ids = sales_train_df["id"].unique()[:8] + + base_maes: list[float] = [] + ft_maes: list[float] = [] + + for store_id in store_ids: + store_train = ( + sales_train_df[sales_train_df["id"] == store_id][target] + .values.astype(np.float32) + ) + store_test = ( + sales_test_df[sales_test_df["id"] == store_id][target] + .values.astype(np.float32) + ) + ground_truth = store_test[:horizon_len] + if len(ground_truth) < horizon_len or len(store_train) < context_len: + continue + + test_input = torch.tensor( + store_train[-context_len:], dtype=torch.float32, device=device + ).unsqueeze(0) + + with torch.no_grad(): + base_out = base_model(past_values=test_input) + ft_out = ft_model(past_values=test_input) + + base_forecast = base_out.mean_predictions[0, :horizon_len].float().cpu().numpy() + ft_forecast = ft_out.mean_predictions[0, :horizon_len].float().cpu().numpy() + + base_mae = float(np.abs(base_forecast - ground_truth).mean()) + ft_mae = float(np.abs(ft_forecast - ground_truth).mean()) + base_maes.append(base_mae) + ft_maes.append(ft_mae) + + logger.info( + "Store %s — zero-shot MAE: %.2f, LoRA MAE: %.2f", + store_id, + base_mae, + ft_mae, + ) + + if base_maes: + avg_base = np.mean(base_maes) + avg_ft = np.mean(ft_maes) + improvement = (avg_base - avg_ft) / avg_base * 100 + logger.info("Average zero-shot MAE: %.2f", avg_base) + logger.info("Average LoRA MAE: %.2f", avg_ft) + logger.info("Improvement: %.1f%%", improvement) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description="Fine-tune TimesFM 2.5 with LoRA (Transformers + PEFT)" + ) + p.add_argument( + "--model_id", + default="google/timesfm-2.5-200m-transformers", + help="HuggingFace model ID", + ) + p.add_argument("--context_len", type=int, default=64) + p.add_argument("--horizon_len", type=int, default=13) + p.add_argument("--epochs", type=int, default=10) + p.add_argument("--batch_size", type=int, default=32) + p.add_argument("--lr", type=float, default=1e-4) + p.add_argument("--lora_r", type=int, default=4) + p.add_argument("--lora_alpha", type=int, default=8) + p.add_argument("--lora_dropout", type=float, default=0.05) + p.add_argument("--num_samples", type=int, default=5000) + p.add_argument("--output_dir", default="timesfm2_5-retail-lora") + p.add_argument("--seed", type=int, default=42) + p.add_argument( + "--eval_only", + action="store_true", + help="Skip training and only run evaluation", + ) + return p.parse_args() + + +def main() -> None: + args = parse_args() + + if not args.eval_only: + train(args) + + if os.path.isdir(args.output_dir): + evaluate(args) + else: + logger.warning( + "No adapter found at %s — skipping evaluation.", args.output_dir + ) + + +if __name__ == "__main__": + main()