freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: replace custom PEFT pipeline with Transformers+PEFT example

Browse Source 
Remove the custom peft/ directory (LoRA/DoRA adapters, trainer, data
pipeline) in favor of a lightweight fine-tuning example that uses the
standard HuggingFace Transformers + PEFT ecosystem.

The new example at timesfm-forecasting/examples/finetuning/ demonstrates
LoRA fine-tuning via TimesFm2_5ModelForPrediction and the peft library,
based on the approach by @kashif at HuggingFace.

- Remove peft/ (8 files)
- Add timesfm-forecasting/examples/finetuning/finetune_lora.py
- Add timesfm-forecasting/examples/finetuning/README.md
- Update README.md to reference new example
- Clean up .gitignore (remove peft_checkpoints/)
This commit is contained in:
darkpowerxo
2026-04-09 11:15:58 -04:00
parent 18d5eb2d44
commit caddef1db8
12 changed files with 554 additions and 1693 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
-1
View File
@@ -3,7 +3,6 @@ dist/
__pycache__/
*.egg-info/
checkpoints/
peft_checkpoints/
wandb/
datasets/
results/
README.md
+6 -6
View File
@@ -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
peft/README.md
-201
View File
@@ -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
```
peft/__init__.py
-41
View File
@@ -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",
]
peft/adapters.py
-285
View File
@@ -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)
peft/config.py
-106
View File
@@ -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
peft/data.py
-150
View File
@@ -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),
)
peft/finetune.py
-252
View File
@@ -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()
peft/finetune.sh
-73
View File
@@ -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
peft/trainer.py
-578
View File
@@ -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
timesfm-forecasting/examples/finetuning/README.md
+102
View File
@@ -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:
<https://github.com/huggingface/notebooks/blob/main/examples/timesfm2_5.ipynb>
timesfm-forecasting/examples/finetuning/finetune_lora.py
+446
View File
@@ -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()