From aa2b17f0b3ad3c83f688ac6cad6a4516c8b197c1 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 13:51:37 -0400 Subject: [PATCH 01/17] chore: update .gitignore for egg-info, uv.lock, peft_checkpoints - Replace timesfm_jax.egg-info/ with generic *.egg-info/ glob - Add uv.lock (lockfile is environment-specific) - Add peft_checkpoints/ (training output directory) --- .gitignore | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 589bbfc..44f18c0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,9 +1,11 @@ .venv/ dist/ __pycache__/ +*.egg-info/ checkpoints/ +peft_checkpoints/ wandb/ datasets/ results/ -timesfm_jax.egg-info/ +uv.lock development_setup.md From 7357458e45f259fe74f26a91b072bf06e2dcc110 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 13:51:55 -0400 Subject: [PATCH 02/17] feat: add LoRA/DoRA adapter layers for TimesFM 2.5 (PyTorch) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement production-grade PEFT adapters targeting the 2.5 architecture: - LoRALinear: low-rank A/B decomposition with scaling (alpha/rank) - DoRALinear: weight-decomposed LoRA (magnitude + direction) - inject_adapters(): freezes base weights, wraps target nn.Linear modules - Supports fused QKV (qkv_proj), attention output, and FFN layers - num_adapter_layers controls how many top layers get adapters (0=all) - target_modules selects 'all', 'attention', or 'ffn' - merge_adapters(): folds adapter deltas back into base nn.Linear - save/load_adapter_weights(): safetensors adapter-only checkpoints - PEFTConfig dataclass with all hyperparameters References: LoRA — https://arxiv.org/abs/2106.09685 DoRA — https://arxiv.org/abs/2402.09353 --- peft/__init__.py | 41 +++++++ peft/adapters.py | 283 +++++++++++++++++++++++++++++++++++++++++++++++ peft/config.py | 106 ++++++++++++++++++ 3 files changed, 430 insertions(+) create mode 100644 peft/__init__.py create mode 100644 peft/adapters.py create mode 100644 peft/config.py diff --git a/peft/__init__.py b/peft/__init__.py new file mode 100644 index 0000000..2cadf0d --- /dev/null +++ b/peft/__init__.py @@ -0,0 +1,41 @@ +# 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 new file mode 100644 index 0000000..951cfe1 --- /dev/null +++ b/peft/adapters.py @@ -0,0 +1,283 @@ +# 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 + + 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)) + + # 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. + """ + tensors = load_file(path, device="cpu") + 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 new file mode 100644 index 0000000..643be89 --- /dev/null +++ b/peft/config.py @@ -0,0 +1,106 @@ +# 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 From 9875d926fe87b11306e9cba7ae8cb7ffa43e382f Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 13:52:19 -0400 Subject: [PATCH 03/17] feat: add TimeSeriesDataset for PEFT fine-tuning Sliding-window dataset that produces (context, mask, target) tuples: - Accepts list of arrays, long-format, or wide-format DataFrames - Context length auto-rounded to multiple of patch_len (32) - Left-pads short series with proper masking - Configurable stride for window overlap --- peft/data.py | 150 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 150 insertions(+) create mode 100644 peft/data.py diff --git a/peft/data.py b/peft/data.py new file mode 100644 index 0000000..34d6dae --- /dev/null +++ b/peft/data.py @@ -0,0 +1,150 @@ +# 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), + ) From eca7ca342810fa1f72b43f331437c9c4eb7ef6c9 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 13:52:59 -0400 Subject: [PATCH 04/17] feat: add multi-GPU PEFT trainer for TimesFM 2.5 PEFTTrainer with production-grade training loop: - PyTorch DDP multi-GPU via torchrun - Mixed-precision training (fp16/bf16) with GradScaler - Gradient checkpointing for long contexts - Cosine-with-warmup LR schedule - MSE loss + optional pinball quantile loss (9 channels) - Early stopping on validation loss - Adapter-only checkpointing (safetensors) - W&B logging (rank-0 only) - Differentiable training forward that replicates the 2.5 patch -> RevIN -> transformer -> output-head -> un-RevIN path --- peft/trainer.py | 578 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 578 insertions(+) create mode 100644 peft/trainer.py diff --git a/peft/trainer.py b/peft/trainer.py new file mode 100644 index 0000000..4b3ccc1 --- /dev/null +++ b/peft/trainer.py @@ -0,0 +1,578 @@ +# 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 From a67eeb2e7d84e31e2ba95af1ecd05e1cf5ce436d Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 13:53:35 -0400 Subject: [PATCH 05/17] feat: add CLI entry-point and launch script for PEFT fine-tuning - finetune.py: argparse CLI with all config flags, CSV data loading, chronological train/val split, and full training pipeline Usage: python -m peft.finetune --data_path data.csv --value_col y Multi-GPU: torchrun --nproc_per_node=4 -m peft.finetune ... - finetune.sh: env-var driven launch script for single/multi-GPU Usage: DATA_PATH=data.csv NUM_GPUS=4 bash peft/finetune.sh --- peft/finetune.py | 252 +++++++++++++++++++++++++++++++++++++++++++++++ peft/finetune.sh | 73 ++++++++++++++ 2 files changed, 325 insertions(+) create mode 100644 peft/finetune.py create mode 100644 peft/finetune.sh diff --git a/peft/finetune.py b/peft/finetune.py new file mode 100644 index 0000000..70d379c --- /dev/null +++ b/peft/finetune.py @@ -0,0 +1,252 @@ +#!/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 new file mode 100644 index 0000000..37a16e7 --- /dev/null +++ b/peft/finetune.sh @@ -0,0 +1,73 @@ +#!/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 From b6ac2b3559c521ec4482ff8fb8db3c11d4b2ccef Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 13:53:43 -0400 Subject: [PATCH 06/17] docs: add README for the PEFT fine-tuning pipeline Covers quick start, Python API, CLI reference, adapter loading/merging, architecture overview with parameter counts, and file layout. --- peft/README.md | 201 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 peft/README.md diff --git a/peft/README.md b/peft/README.md new file mode 100644 index 0000000..fdc96f6 --- /dev/null +++ b/peft/README.md @@ -0,0 +1,201 @@ +# 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 +``` From bc03b77e9e9fba779e35deab78471a9b87342fe3 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 14:09:14 -0400 Subject: [PATCH 07/17] fix: correct 'complied' typo and replace print with logging Apply changes from PR #396 by @shahrukhx01: - Fix typo 'complied' -> 'compiled' in ForecastConfig docstrings - Replace bare print() with logging.info() in load_checkpoint() --- src/timesfm/configs.py | 4 ++-- src/timesfm/timesfm_2p5/timesfm_2p5_torch.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/timesfm/configs.py b/src/timesfm/configs.py index 88a63af..9a263de 100644 --- a/src/timesfm/configs.py +++ b/src/timesfm/configs.py @@ -23,11 +23,11 @@ class ForecastConfig: """Options for forecasting. Attributes: - max_context: The maximum context length. This is used by the complied decode + max_context: The maximum context length. This is used by the compiled decode function at inference time during batched inference. Any input time series with length less than max_context will be padded with zeros, and with length greater than max_context will be truncated. - max_horizon: The maximum horizon length. This is used by the complied decode + max_horizon: The maximum horizon length. This is used by the compiled decode function at inference time during batched inference. The compiled cached decoding function will by default forecast till max_horizon. normalize_inputs: Whether to normalize the inputs. This is useful when the diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index 9bbd771..3e7c9f1 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py @@ -85,7 +85,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module): if "torch_compile" in kwargs: torch_compile = kwargs["torch_compile"] if torch_compile: - print("Compiling model...") + logging.info("Compiling model...") self = torch.compile(self) self.eval() From c10494a4c570321b50b69655216c3d54d91f0753 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 14:13:37 -0400 Subject: [PATCH 08/17] test: add unit tests for configs, torch layers, utils, and base utils Apply changes from PR #394 by @cj-wong: - tests/__init__.py: package marker - tests/test_base_utils.py: strip_leading_nans + linear_interpolation tests - tests/test_configs.py: frozen dataclass, defaults, replace, equality tests - tests/test_torch_layers.py: ResidualBlock, RMSNorm, RandomFourierFeatures - tests/test_torch_utils.py: update_running_stats, revin, DecodeCache tests --- tests/__init__.py | 1 + tests/test_base_utils.py | 168 +++++++++++++++++++ tests/test_configs.py | 196 ++++++++++++++++++++++ tests/test_torch_layers.py | 262 +++++++++++++++++++++++++++++ tests/test_torch_utils.py | 336 +++++++++++++++++++++++++++++++++++++ 5 files changed, 963 insertions(+) create mode 100644 tests/__init__.py create mode 100644 tests/test_base_utils.py create mode 100644 tests/test_configs.py create mode 100644 tests/test_torch_layers.py create mode 100644 tests/test_torch_utils.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..f4c7e19 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# Tests for TimesFM. diff --git a/tests/test_base_utils.py b/tests/test_base_utils.py new file mode 100644 index 0000000..ace6c5f --- /dev/null +++ b/tests/test_base_utils.py @@ -0,0 +1,168 @@ +# 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. + +"""Tests for NaN-handling and interpolation utilities in the base module. + +``strip_leading_nans`` and ``linear_interpolation`` sit on the critical +inference path: every user input passes through them before being +patched and fed to the transformer. Incorrect behavior here — silently +keeping NaN values or interpolating the wrong indices — causes NaN +propagation through the entire model and produces garbage forecasts. +""" + +import numpy as np + +from timesfm.timesfm_2p5.timesfm_2p5_base import ( + linear_interpolation, + strip_leading_nans, +) + + +# --------------------------------------------------------------------------- +# strip_leading_nans +# --------------------------------------------------------------------------- + + +class TestStripLeadingNans: + """Tests for strip_leading_nans — removes leading NaN prefix.""" + + def test_no_nans_returns_unchanged(self): + """An array without NaN values must pass through unmodified.""" + arr = np.array([1.0, 2.0, 3.0]) + result = strip_leading_nans(arr) + np.testing.assert_array_equal(result, arr) + + def test_strips_leading_nans_only(self): + """Leading NaNs are removed; NaNs embedded in the middle are kept.""" + arr = np.array([np.nan, np.nan, 1.0, np.nan, 3.0]) + result = strip_leading_nans(arr) + expected = np.array([1.0, np.nan, 3.0]) + np.testing.assert_array_equal(result, expected) + + def test_single_leading_nan(self): + """Edge case: exactly one leading NaN.""" + arr = np.array([np.nan, 5.0, 6.0]) + result = strip_leading_nans(arr) + np.testing.assert_array_equal(result, np.array([5.0, 6.0])) + + def test_no_leading_nan_with_internal_nans(self): + """If the first element is valid, nothing is stripped regardless of + internal NaNs.""" + arr = np.array([1.0, np.nan, np.nan, 4.0]) + result = strip_leading_nans(arr) + np.testing.assert_array_equal(result, arr) + + def test_single_valid_element(self): + """A single non-NaN element must be returned as-is.""" + arr = np.array([42.0]) + result = strip_leading_nans(arr) + np.testing.assert_array_equal(result, np.array([42.0])) + + def test_all_nans_returns_full_array(self): + """When every element is NaN, ``np.argmax`` on an all-False mask + returns 0 — so the implementation returns the original array, not an + empty one. + + This documents the *actual* behavior (which differs from the + docstring claim of returning an empty array). Downstream code + (``linear_interpolation``) is designed to handle this case. + """ + arr = np.array([np.nan, np.nan, np.nan]) + result = strip_leading_nans(arr) + # Actual behavior: argmax(~isnan) = 0 when all NaN → returns full array. + assert len(result) == 3 + assert np.all(np.isnan(result)) + + def test_preserves_dtype(self): + """Output dtype must match input dtype (float32 stays float32).""" + arr = np.array([np.nan, 1.0, 2.0], dtype=np.float32) + result = strip_leading_nans(arr) + assert result.dtype == np.float32 + + +# --------------------------------------------------------------------------- +# linear_interpolation +# --------------------------------------------------------------------------- + + +class TestLinearInterpolation: + """Tests for linear_interpolation — fills NaN gaps via ``np.interp``.""" + + def test_no_nans_returns_identical(self): + """Without NaN values the array is returned as-is (fast path).""" + arr = np.array([1.0, 2.0, 3.0]) + result = linear_interpolation(arr.copy()) + np.testing.assert_array_equal(result, arr) + + def test_interpolates_single_interior_nan(self): + """A single interior NaN is linearly interpolated from neighbors.""" + arr = np.array([0.0, np.nan, 2.0]) + result = linear_interpolation(arr) + np.testing.assert_allclose(result, [0.0, 1.0, 2.0]) + + def test_interpolates_multiple_interior_nans(self): + """Multiple consecutive interior NaN values are interpolated.""" + arr = np.array([0.0, np.nan, np.nan, 3.0]) + result = linear_interpolation(arr) + np.testing.assert_allclose(result, [0.0, 1.0, 2.0, 3.0]) + + def test_extrapolates_trailing_nans(self): + """Trailing NaN values are filled via ``np.interp`` which holds the + last known value (nearest-neighbor extrapolation).""" + arr = np.array([1.0, 2.0, np.nan, np.nan]) + result = linear_interpolation(arr) + # np.interp extrapolates by clamping to boundary values. + np.testing.assert_allclose(result, [1.0, 2.0, 2.0, 2.0]) + + def test_extrapolates_leading_nans(self): + """Leading NaN values are filled with the first valid value. + + In practice ``strip_leading_nans`` runs first, but the function must + still be robust on its own. + """ + arr = np.array([np.nan, np.nan, 3.0, 4.0]) + result = linear_interpolation(arr) + np.testing.assert_allclose(result, [3.0, 3.0, 3.0, 4.0]) + + def test_output_has_no_nans(self): + """After interpolation, no NaN values should remain.""" + arr = np.array([np.nan, 1.0, np.nan, np.nan, 4.0, np.nan]) + result = linear_interpolation(arr) + assert not np.any(np.isnan(result)) + + def test_preserves_non_nan_values(self): + """Non-NaN values in the original array must never be modified.""" + arr = np.array([10.0, np.nan, 30.0, np.nan, 50.0]) + original_valid = arr[~np.isnan(arr)].copy() + result = linear_interpolation(arr) + np.testing.assert_array_equal( + result[~np.isnan(np.array([10.0, np.nan, 30.0, np.nan, 50.0]))], + original_valid, + ) + + def test_interpolation_is_monotone_for_monotone_input(self): + """If the known values are strictly increasing, the interpolated + result must also be non-decreasing — a basic sanity check on the + interpolation direction.""" + arr = np.array([1.0, np.nan, np.nan, 4.0, np.nan, 6.0]) + result = linear_interpolation(arr) + diffs = np.diff(result) + assert np.all(diffs >= 0) + + def test_single_non_nan_fills_all_gaps(self): + """With only one valid value, every NaN is replaced by that value + (np.interp clamps to the single known point).""" + arr = np.array([np.nan, 5.0, np.nan]) + result = linear_interpolation(arr) + np.testing.assert_allclose(result, [5.0, 5.0, 5.0]) diff --git a/tests/test_configs.py b/tests/test_configs.py new file mode 100644 index 0000000..bd07c6f --- /dev/null +++ b/tests/test_configs.py @@ -0,0 +1,196 @@ +# 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. + +"""Tests for TimesFM configuration dataclasses. + +These tests verify that config dataclasses enforce immutability, compose +correctly, and carry the exact default values the model implementation +relies on. Catching a silent default-value drift here prevents subtle +inference regressions that would otherwise only surface as degraded +forecast quality. +""" + +import dataclasses + +import pytest + +from timesfm.configs import ( + ForecastConfig, + RandomFourierFeaturesConfig, + ResidualBlockConfig, + StackedTransformersConfig, + TransformerConfig, +) + + +# --------------------------------------------------------------------------- +# ForecastConfig +# --------------------------------------------------------------------------- + + +class TestForecastConfig: + """Tests for ForecastConfig — the primary user-facing configuration.""" + + def test_defaults_match_safe_inference_settings(self): + """Default config must be conservative: no normalization, no fancy heads. + + These defaults are what users get when they call ``ForecastConfig()`` + without arguments. Changing them silently would break all existing + code that relies on the defaults. + """ + cfg = ForecastConfig() + assert cfg.max_context == 0 + assert cfg.max_horizon == 0 + assert cfg.normalize_inputs is False + assert cfg.per_core_batch_size == 1 + assert cfg.use_continuous_quantile_head is False + assert cfg.force_flip_invariance is True + assert cfg.infer_is_positive is True + assert cfg.fix_quantile_crossing is False + assert cfg.return_backcast is False + + def test_frozen_prevents_mutation(self): + """Configs are frozen dataclasses — mutating them must raise. + + This is critical because ``compile()`` captures the config object and + the compiled decode closure relies on its values never changing. + """ + cfg = ForecastConfig(max_context=512) + with pytest.raises(dataclasses.FrozenInstanceError): + cfg.max_context = 1024 + + def test_replace_creates_independent_copy(self): + """``dataclasses.replace`` must yield a new object with updated fields. + + The compile path uses ``replace`` to adjust context/horizon to valid + multiples; the original config must remain untouched. + """ + original = ForecastConfig(max_context=512, max_horizon=128) + replaced = dataclasses.replace(original, max_context=1024) + + assert replaced.max_context == 1024 + assert replaced.max_horizon == 128 # untouched + assert original.max_context == 512 # original unchanged + + def test_equality_is_structural(self): + """Two configs with identical fields must be equal (value semantics).""" + a = ForecastConfig(max_context=256, normalize_inputs=True) + b = ForecastConfig(max_context=256, normalize_inputs=True) + assert a == b + + def test_inequality_on_any_field_difference(self): + """A single differing field must break equality.""" + a = ForecastConfig(max_context=256) + b = ForecastConfig(max_context=512) + assert a != b + + +# --------------------------------------------------------------------------- +# ResidualBlockConfig +# --------------------------------------------------------------------------- + + +class TestResidualBlockConfig: + """Tests for ResidualBlockConfig used by tokenizer and output projections.""" + + def test_frozen_prevents_mutation(self): + cfg = ResidualBlockConfig( + input_dims=64, + hidden_dims=128, + output_dims=128, + use_bias=True, + activation="swish", + ) + with pytest.raises(dataclasses.FrozenInstanceError): + cfg.input_dims = 32 + + def test_activation_accepts_all_valid_literals(self): + """All three activation modes must be constructable without error.""" + for act in ("relu", "swish", "none"): + cfg = ResidualBlockConfig( + input_dims=8, + hidden_dims=16, + output_dims=8, + use_bias=False, + activation=act, + ) + assert cfg.activation == act + + +# --------------------------------------------------------------------------- +# TransformerConfig & StackedTransformersConfig +# --------------------------------------------------------------------------- + + +class TestTransformerConfig: + """Tests for TransformerConfig — architecture-level hyperparameters.""" + + def test_model_dims_must_be_divisible_by_num_heads(self): + """The model instantiation will fail if this invariant is broken. + + We verify the config at least *carries* the right values that the + TimesFM 2.5 definition uses (1280 dims, 16 heads → 80 head_dim). + """ + cfg = TransformerConfig( + model_dims=1280, + hidden_dims=1280, + num_heads=16, + attention_norm="rms", + feedforward_norm="rms", + qk_norm="rms", + use_bias=False, + use_rotary_position_embeddings=True, + ff_activation="swish", + fuse_qkv=True, + ) + assert cfg.model_dims % cfg.num_heads == 0 + assert cfg.model_dims // cfg.num_heads == 80 # head_dim + + def test_stacked_config_composes_correctly(self): + """StackedTransformersConfig must wrap a TransformerConfig cleanly.""" + xf = TransformerConfig( + model_dims=64, + hidden_dims=64, + num_heads=4, + attention_norm="rms", + feedforward_norm="rms", + qk_norm="none", + use_bias=True, + use_rotary_position_embeddings=False, + ff_activation="relu", + fuse_qkv=False, + ) + stacked = StackedTransformersConfig(num_layers=6, transformer=xf) + assert stacked.num_layers == 6 + assert stacked.transformer is xf + assert stacked.transformer.model_dims == 64 + + +# --------------------------------------------------------------------------- +# RandomFourierFeaturesConfig +# --------------------------------------------------------------------------- + + +class TestRandomFourierFeaturesConfig: + """Tests for RandomFourierFeaturesConfig.""" + + def test_frozen_prevents_mutation(self): + cfg = RandomFourierFeaturesConfig( + input_dims=32, + output_dims=64, + projection_stddev=1.0, + use_bias=True, + ) + with pytest.raises(dataclasses.FrozenInstanceError): + cfg.output_dims = 128 diff --git a/tests/test_torch_layers.py b/tests/test_torch_layers.py new file mode 100644 index 0000000..e16fdb4 --- /dev/null +++ b/tests/test_torch_layers.py @@ -0,0 +1,262 @@ +# 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. + +"""Tests for PyTorch layer building blocks: ResidualBlock, RMSNorm, +RandomFourierFeatures. + +These layers are the atoms of the TimesFM architecture. Verifying their +output shapes, numerical properties, and failure modes protects against +regressions during refactors. All tests use small dimensions and run +on CPU — no model checkpoint or GPU required. +""" + +import torch +import pytest + +from timesfm.configs import RandomFourierFeaturesConfig, ResidualBlockConfig +from timesfm.torch.dense import RandomFourierFeatures, ResidualBlock +from timesfm.torch.normalization import RMSNorm + + +# --------------------------------------------------------------------------- +# ResidualBlock +# --------------------------------------------------------------------------- + + +class TestResidualBlock: + """Tests for the residual block: hidden → activation → output + skip.""" + + @pytest.fixture + def swish_block(self): + """A small residual block with SiLU/Swish activation (matches TimesFM).""" + cfg = ResidualBlockConfig( + input_dims=16, + hidden_dims=32, + output_dims=8, + use_bias=True, + activation="swish", + ) + return ResidualBlock(cfg) + + def test_output_shape(self, swish_block): + """Output must have the config's ``output_dims`` as the last dimension, + regardless of input batch shape.""" + x = torch.randn(4, 16) + out = swish_block(x) + assert out.shape == (4, 8) + + def test_output_shape_3d(self, swish_block): + """The block must handle (batch, seq, features) inputs — the layout + used when processing patched time series.""" + x = torch.randn(2, 10, 16) + out = swish_block(x) + assert out.shape == (2, 10, 8) + + def test_residual_connection_nonzero(self): + """The residual connection must contribute to the output. + + We verify this by comparing the output when the hidden path is + zeroed out vs. the full output. + """ + cfg = ResidualBlockConfig( + input_dims=8, + hidden_dims=16, + output_dims=8, + use_bias=False, + activation="none", + ) + block = ResidualBlock(cfg) + x = torch.randn(2, 8) + + with torch.no_grad(): + # Residual path only: zero out hidden and output layers. + block.hidden_layer.weight.zero_() + block.output_layer.weight.zero_() + residual_only = block(x) + + # Must equal the residual layer output. + expected = block.residual_layer(x) + torch.testing.assert_close(residual_only, expected) + + @pytest.mark.parametrize("activation", ["relu", "swish", "none"]) + def test_all_activations_produce_valid_output(self, activation): + """All supported activations must produce finite, non-NaN output.""" + cfg = ResidualBlockConfig( + input_dims=8, + hidden_dims=16, + output_dims=8, + use_bias=True, + activation=activation, + ) + block = ResidualBlock(cfg) + x = torch.randn(4, 8) + out = block(x) + assert not torch.any(torch.isnan(out)) + assert not torch.any(torch.isinf(out)) + + def test_invalid_activation_raises(self): + """Unsupported activation must raise ``ValueError`` immediately — + fail fast rather than producing garbage at inference time.""" + cfg = ResidualBlockConfig( + input_dims=8, + hidden_dims=16, + output_dims=8, + use_bias=True, + activation="gelu", + ) + with pytest.raises(ValueError, match="not supported"): + ResidualBlock(cfg) + + def test_gradient_flows_through_both_paths(self): + """Gradients must reach both the main path and the residual path. + + Dead gradients on either path would prevent the layer from learning. + """ + cfg = ResidualBlockConfig( + input_dims=8, + hidden_dims=16, + output_dims=8, + use_bias=True, + activation="swish", + ) + block = ResidualBlock(cfg) + x = torch.randn(2, 8, requires_grad=True) + out = block(x) + loss = out.sum() + loss.backward() + + assert block.hidden_layer.weight.grad is not None + assert block.residual_layer.weight.grad is not None + assert torch.any(block.hidden_layer.weight.grad != 0) + assert torch.any(block.residual_layer.weight.grad != 0) + + +# --------------------------------------------------------------------------- +# RMSNorm +# --------------------------------------------------------------------------- + + +class TestRMSNorm: + """Tests for RMS normalization used in transformer attention/FF blocks.""" + + def test_output_shape_preserved(self): + """RMSNorm must not change the tensor shape.""" + norm = RMSNorm(num_features=64) + x = torch.randn(2, 10, 64) + out = norm(x) + assert out.shape == x.shape + + def test_zero_scale_produces_zeros(self): + """With default scale (initialized to zeros), output must be all zeros. + + This is a critical initialization property: at init, each transformer + layer's post-norm effectively passes through zeros, relying on the + residual connection to carry signal. + """ + norm = RMSNorm(num_features=8) + # scale is initialized to zeros by default. + x = torch.randn(4, 8) + out = norm(x) + torch.testing.assert_close(out, torch.zeros_like(out)) + + def test_unit_scale_preserves_rms_magnitude(self): + """With scale = 1, output should have approximately unit RMS along + the feature dimension — that's the point of RMS normalization.""" + norm = RMSNorm(num_features=64) + with torch.no_grad(): + norm.scale.fill_(1.0) + + x = torch.randn(8, 64) * 100 # large magnitude + out = norm(x) + + rms = torch.sqrt(torch.mean(out**2, dim=-1)) + # After normalization, RMS should be close to 1.0. + torch.testing.assert_close( + rms, + torch.ones(8), + atol=0.1, + rtol=0.1, + ) + + def test_no_nan_on_zero_input(self): + """A zero-valued input must not cause NaN (epsilon prevents div-by-0).""" + norm = RMSNorm(num_features=8, epsilon=1e-6) + with torch.no_grad(): + norm.scale.fill_(1.0) + x = torch.zeros(2, 8) + out = norm(x) + assert not torch.any(torch.isnan(out)) + + +# --------------------------------------------------------------------------- +# RandomFourierFeatures +# --------------------------------------------------------------------------- + + +class TestRandomFourierFeatures: + """Tests for the random Fourier feature layer.""" + + def test_output_shape(self): + """Output dims must be exactly ``config.output_dims``.""" + cfg = RandomFourierFeaturesConfig( + input_dims=8, + output_dims=32, + projection_stddev=1.0, + use_bias=True, + ) + layer = RandomFourierFeatures(cfg) + x = torch.randn(4, 8) + out = layer(x) + assert out.shape == (4, 32) + + def test_output_dims_not_multiple_of_4_raises(self): + """The four Fourier components (cos, sin, sq_wave_1, sq_wave_2) + require ``output_dims`` to be divisible by 4.""" + cfg = RandomFourierFeaturesConfig( + input_dims=8, + output_dims=30, # not divisible by 4 + projection_stddev=1.0, + use_bias=True, + ) + with pytest.raises(ValueError, match="multiple of 4"): + RandomFourierFeatures(cfg) + + def test_fourier_components_bounded(self): + """cos and sin outputs are bounded in [-1, 1]; sign outputs are + bounded in {-1, 0, 1}. The total Fourier part (before residual) + is thus bounded. We verify the output stays finite.""" + cfg = RandomFourierFeaturesConfig( + input_dims=8, + output_dims=32, + projection_stddev=1.0, + use_bias=False, + ) + layer = RandomFourierFeatures(cfg) + x = torch.randn(16, 8) * 10 # moderately large input + out = layer(x) + assert not torch.any(torch.isnan(out)) + assert not torch.any(torch.isinf(out)) + + def test_3d_input_supported(self): + """The layer must handle (batch, seq, features) tensors.""" + cfg = RandomFourierFeaturesConfig( + input_dims=8, + output_dims=16, + projection_stddev=1.0, + use_bias=True, + ) + layer = RandomFourierFeatures(cfg) + x = torch.randn(2, 5, 8) + out = layer(x) + assert out.shape == (2, 5, 16) diff --git a/tests/test_torch_utils.py b/tests/test_torch_utils.py new file mode 100644 index 0000000..76f2296 --- /dev/null +++ b/tests/test_torch_utils.py @@ -0,0 +1,336 @@ +# 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. + +"""Tests for PyTorch utility functions: running statistics and RevIN. + +These utilities are invoked at every patch boundary during autoregressive +decoding. Bugs here cause silent numerical drift that compounds over +long horizons, making them especially hard to diagnose from forecast +output alone. +""" + +import torch +import numpy as np +import pytest + +from timesfm.torch.util import ( + DecodeCache, + _TOLERANCE, + revin, + update_running_stats, +) + + +# --------------------------------------------------------------------------- +# update_running_stats +# --------------------------------------------------------------------------- + + +class TestUpdateRunningStats: + """Tests for Welford-style online mean / variance accumulation.""" + + def test_single_batch_matches_numpy(self): + """A single update with no mask must match numpy's mean and std. + + This is the most basic correctness check: feed all values at once + and compare against the ground-truth statistics. + """ + x = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]]) + mask = torch.zeros_like(x, dtype=torch.bool) + n0 = torch.zeros(1) + mu0 = torch.zeros(1) + sigma0 = torch.zeros(1) + + (new_n, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask) + + np_values = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + expected_mu = np.mean(np_values) + # Population std (ddof=0), same as PyTorch default. + expected_sigma = np.std(np_values, ddof=0) + + assert new_n.item() == pytest.approx(5.0) + assert new_mu.item() == pytest.approx(expected_mu, abs=1e-5) + assert new_sigma.item() == pytest.approx(expected_sigma, abs=1e-5) + + def test_incremental_accumulation_matches_full_computation(self): + """Accumulating two batches incrementally must yield the same result + as computing statistics over all values at once. + + This is the defining property of online/streaming statistics. + """ + all_values = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]) + batch1 = torch.tensor([[1.0, 2.0, 3.0]]) + batch2 = torch.tensor([[4.0, 5.0, 6.0]]) + + def no_mask(t): + return torch.zeros_like(t, dtype=torch.bool) + + # Full computation. + n0 = torch.zeros(1) + mu0 = torch.zeros(1) + sigma0 = torch.zeros(1) + (full_n, full_mu, full_sigma), _ = update_running_stats( + n0, mu0, sigma0, all_values, no_mask(all_values) + ) + + # Incremental computation. + (n1, mu1, sigma1), _ = update_running_stats( + n0, mu0, sigma0, batch1, no_mask(batch1) + ) + (inc_n, inc_mu, inc_sigma), _ = update_running_stats( + n1, mu1, sigma1, batch2, no_mask(batch2) + ) + + assert inc_n.item() == pytest.approx(full_n.item()) + assert inc_mu.item() == pytest.approx(full_mu.item(), abs=1e-5) + assert inc_sigma.item() == pytest.approx(full_sigma.item(), abs=1e-5) + + def test_masked_elements_excluded_from_statistics(self): + """Masked positions must be completely ignored — as if they don't exist. + + In TimesFM, leading padding is masked. If mask handling is broken, + the zero-padding values pollute the running mean and variance. + """ + # Two values: 10 and 20 are valid; 0 is masked. + x = torch.tensor([[0.0, 10.0, 20.0]]) + mask = torch.tensor([[True, False, False]]) + n0 = torch.zeros(1) + mu0 = torch.zeros(1) + sigma0 = torch.zeros(1) + + (new_n, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask) + + assert new_n.item() == pytest.approx(2.0) + assert new_mu.item() == pytest.approx(15.0, abs=1e-5) + expected_sigma = np.std([10.0, 20.0], ddof=0) + assert new_sigma.item() == pytest.approx(expected_sigma, abs=1e-5) + + def test_all_masked_yields_zero_stats(self): + """When every element is masked, the function must return zeros + rather than NaN or raise an error. + + This happens when an input series is entirely padding. + """ + x = torch.tensor([[99.0, 99.0, 99.0]]) + mask = torch.ones_like(x, dtype=torch.bool) + n0 = torch.zeros(1) + mu0 = torch.zeros(1) + sigma0 = torch.zeros(1) + + (new_n, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask) + + assert new_n.item() == 0.0 + assert new_mu.item() == 0.0 + assert new_sigma.item() == 0.0 + + def test_batched_computation_independent(self): + """Each sample in the batch must be computed independently. + + Cross-sample leakage would corrupt multi-series forecasting. + """ + x = torch.tensor( + [ + [1.0, 2.0, 3.0], + [100.0, 200.0, 300.0], + ] + ) + mask = torch.zeros_like(x, dtype=torch.bool) + n0 = torch.zeros(2) + mu0 = torch.zeros(2) + sigma0 = torch.zeros(2) + + (new_n, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask) + + assert new_mu[0].item() == pytest.approx(2.0, abs=1e-5) + assert new_mu[1].item() == pytest.approx(200.0, abs=1e-5) + + expected_sigma_0 = np.std([1.0, 2.0, 3.0], ddof=0) + expected_sigma_1 = np.std([100.0, 200.0, 300.0], ddof=0) + assert new_sigma[0].item() == pytest.approx(expected_sigma_0, abs=1e-5) + assert new_sigma[1].item() == pytest.approx(expected_sigma_1, abs=1e-5) + + def test_constant_input_yields_zero_sigma(self): + """A constant series has zero variance — sigma must be exactly 0. + + This is important because ``revin`` guards against division-by-zero + using ``_TOLERANCE`` when sigma is near zero. + """ + x = torch.tensor([[7.0, 7.0, 7.0, 7.0]]) + mask = torch.zeros_like(x, dtype=torch.bool) + n0 = torch.zeros(1) + mu0 = torch.zeros(1) + sigma0 = torch.zeros(1) + + (_, new_mu, new_sigma), _ = update_running_stats(n0, mu0, sigma0, x, mask) + + assert new_mu.item() == pytest.approx(7.0) + assert new_sigma.item() == pytest.approx(0.0) + + +# --------------------------------------------------------------------------- +# revin (Reversible Instance Normalization) +# --------------------------------------------------------------------------- + + +class TestRevIN: + """Tests for the RevIN normalization used in patched decoding.""" + + def test_forward_then_reverse_is_identity(self): + """normalize → denormalize must reconstruct the original tensor. + + This is the fundamental invariant of reversible normalization: the + model operates in normalized space, and the output is mapped back + to the original scale. Any deviation here directly corrupts the + final forecast values. + """ + x = torch.tensor([[10.0, 20.0, 30.0]]) + mu = torch.tensor([20.0]) + sigma = torch.tensor([10.0]) + + normed = revin(x, mu, sigma, reverse=False) + recovered = revin(normed, mu, sigma, reverse=True) + + torch.testing.assert_close(recovered, x, atol=1e-5, rtol=1e-5) + + def test_forward_produces_correct_normalization(self): + """After forward normalization: (x - mu) / sigma.""" + x = torch.tensor([[10.0, 20.0, 30.0]]) + mu = torch.tensor([20.0]) + sigma = torch.tensor([10.0]) + + normed = revin(x, mu, sigma, reverse=False) + + expected = torch.tensor([[-1.0, 0.0, 1.0]]) + torch.testing.assert_close(normed, expected, atol=1e-5, rtol=1e-5) + + def test_reverse_produces_correct_denormalization(self): + """After reverse: x * sigma + mu.""" + normed = torch.tensor([[-1.0, 0.0, 1.0]]) + mu = torch.tensor([20.0]) + sigma = torch.tensor([10.0]) + + recovered = revin(normed, mu, sigma, reverse=True) + + expected = torch.tensor([[10.0, 20.0, 30.0]]) + torch.testing.assert_close(recovered, expected, atol=1e-5, rtol=1e-5) + + def test_zero_sigma_does_not_produce_nan(self): + """When sigma < tolerance, the function substitutes 1.0 to avoid + division by zero. This occurs for constant-valued input series. + + NaN propagation from here would poison the entire transformer + forward pass. + """ + x = torch.tensor([[5.0, 5.0, 5.0]]) + mu = torch.tensor([5.0]) + sigma = torch.tensor([0.0]) # zero variance + + normed = revin(x, mu, sigma, reverse=False) + + assert not torch.any(torch.isnan(normed)) + assert not torch.any(torch.isinf(normed)) + + def test_near_zero_sigma_guarded_by_tolerance(self): + """Sigma values just below ``_TOLERANCE`` must trigger the guard.""" + x = torch.tensor([[1.0, 2.0, 3.0]]) + mu = torch.tensor([2.0]) + sigma = torch.tensor([_TOLERANCE / 2]) # below threshold + + normed = revin(x, mu, sigma, reverse=False) + + assert not torch.any(torch.isnan(normed)) + # With sigma replaced by 1.0: result = x - mu + expected = torch.tensor([[-1.0, 0.0, 1.0]]) + torch.testing.assert_close(normed, expected, atol=1e-5, rtol=1e-5) + + def test_roundtrip_with_batched_3d_input(self): + """RevIN must broadcast correctly for (batch, patches, patch_len) + tensors — the actual shape used during patched decoding.""" + batch, patches, patch_len = 2, 4, 32 + x = torch.randn(batch, patches, patch_len) + mu = torch.tensor([1.0, 2.0]) # (batch,) + sigma = torch.tensor([3.0, 4.0]) # (batch,) + + normed = revin(x, mu, sigma, reverse=False) + recovered = revin(normed, mu, sigma, reverse=True) + + torch.testing.assert_close(recovered, x, atol=1e-5, rtol=1e-5) + + def test_roundtrip_with_batched_4d_input(self): + """RevIN must broadcast correctly for (batch, patches, patch_len, q) + tensors — the shape used for quantile outputs. + + In the actual decode path, mu/sigma have shape (batch, patches) for + 4D tensors, so the ``len(mu.shape) == len(x.shape) - 2`` branch + fires and adds two trailing singleton dimensions. + """ + batch, patches, patch_len, q = 2, 4, 32, 10 + x = torch.randn(batch, patches, patch_len, q) + # Match the actual call-site shape: (batch, patches) + mu = torch.randn(batch, patches) + sigma = torch.abs(torch.randn(batch, patches)) + 1.0 # ensure positive + + normed = revin(x, mu, sigma, reverse=False) + recovered = revin(normed, mu, sigma, reverse=True) + + torch.testing.assert_close(recovered, x, atol=1e-4, rtol=1e-4) + + def test_negative_values_handled_correctly(self): + """RevIN must work for series with negative values (e.g. temperature, + financial returns). ``infer_is_positive`` is a separate downstream + flag and does not affect RevIN itself. + """ + x = torch.tensor([[-10.0, -5.0, 0.0, 5.0, 10.0]]) + mu = torch.tensor([0.0]) + sigma = torch.tensor([7.07]) + + normed = revin(x, mu, sigma, reverse=False) + recovered = revin(normed, mu, sigma, reverse=True) + + torch.testing.assert_close(recovered, x, atol=1e-3, rtol=1e-3) + + +# --------------------------------------------------------------------------- +# DecodeCache +# --------------------------------------------------------------------------- + + +class TestDecodeCache: + """Tests for the DecodeCache dataclass used in KV-cache decoding.""" + + def test_is_mutable(self): + """DecodeCache is *not* frozen — the attention loop mutates + ``next_index`` and ``num_masked`` in-place during autoregressive + decoding.""" + cache = DecodeCache( + next_index=torch.tensor([0]), + num_masked=torch.tensor([0]), + key=torch.zeros(1, 10, 4, 8), + value=torch.zeros(1, 10, 4, 8), + ) + cache.next_index = torch.tensor([5]) + assert cache.next_index.item() == 5 + + def test_key_value_shape_consistency(self): + """Key and value tensors must have identical shapes — they are + indexed in parallel during attention computation.""" + batch, seq, heads, head_dim = 2, 64, 16, 80 + cache = DecodeCache( + next_index=torch.zeros(batch, dtype=torch.int32), + num_masked=torch.zeros(batch, dtype=torch.int32), + key=torch.zeros(batch, seq, heads, head_dim), + value=torch.zeros(batch, seq, heads, head_dim), + ) + assert cache.key.shape == cache.value.shape + assert cache.key.shape == (batch, seq, heads, head_dim) From a63360a57c974a5de44e697e9b6dd8ecc66589a7 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 14:15:44 -0400 Subject: [PATCH 09/17] fix: per-input ridge regression to prevent data leakage in xreg Apply changes from PR #393 by @MarcoGorworworelli: - Normalize covariates per-input instead of batch-wide to prevent each input's result from depending on batch composition - Fit separate ridge regressions per time series instead of a single batched regression, preventing cross-series data leakage - Applied to both src/timesfm/utils/xreg_lib.py and v1/src/timesfm/xreg_lib.py --- src/timesfm/utils/xreg_lib.py | 110 ++++++++++++++++++---------------- v1/src/timesfm/xreg_lib.py | 107 +++++++++++++++++---------------- 2 files changed, 116 insertions(+), 101 deletions(-) diff --git a/src/timesfm/utils/xreg_lib.py b/src/timesfm/utils/xreg_lib.py index 7a1b19b..4355261 100644 --- a/src/timesfm/utils/xreg_lib.py +++ b/src/timesfm/utils/xreg_lib.py @@ -370,11 +370,20 @@ class BatchedInContextXRegBase: x_train = np.concatenate(x_train, axis=1) x_test = np.concatenate(x_test, axis=1) - # Normalize for robustness. - x_mean = np.mean(x_train, axis=0, keepdims=True) - x_std = np.where((w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w, 1.0) - x_train = [(x_train - x_mean) / x_std] - x_test = [(x_test - x_mean) / x_std] + # Normalize per-input for robustness (batch-wide normalization + # would make each input's result depend on batch composition). + train_splits = np.cumsum(self.train_lens)[:-1] + test_splits = np.cumsum(self.test_lens)[:-1] + train_parts = np.split(x_train, train_splits, axis=0) + test_parts = np.split(x_test, test_splits, axis=0) + norm_train, norm_test = [], [] + for tr, te in zip(train_parts, test_parts): + m = np.mean(tr, axis=0, keepdims=True) + s = np.where((w := np.std(tr, axis=0, keepdims=True)) > _TOL, w, 1.0) + norm_train.append((tr - m) / s) + norm_test.append((te - m) / s) + x_train = [np.concatenate(norm_train, axis=0)] + x_test = [np.concatenate(norm_test, axis=0)] # Categorical features. Encode one by one. one_hot_encoder = preprocessing.OneHotEncoder( @@ -463,58 +472,57 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase): assert_covariate_shapes=assert_covariate_shapes, ) - x_train = x_train_raw.copy() - if max_rows_per_col: - nrows, ncols = x_train.shape - if nrows > (w := ncols * max_rows_per_col): - subsample = jax.random.choice( - jax.random.PRNGKey(max_rows_per_col_sample_seed), - nrows, - (w,), - replace=False, - ) - x_train = x_train[subsample] - flat_targets = flat_targets[subsample] - device = jax.devices("cpu")[0] if force_on_cpu else None - # Runs jitted version of the solvers which are quicker at the cost of - # running jitting during the first time calling. Re-jitting happens whenever - # new (padded) shapes are encountered. - # Ocassionally it helps with the speed and the accuracy if we force single - # thread execution on cpu for accelerator machines: - # 1. Avoid moving data to accelarator memory. - # 2. Avoid precision loss if any. - with jax.default_device(device): - x_train_raw = _to_padded_jax_array(x_train_raw) - x_train = _to_padded_jax_array(x_train) - flat_targets = _to_padded_jax_array(flat_targets) - x_test = _to_padded_jax_array(x_test) - beta_hat = ( - jnp.linalg.pinv( - x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]), - hermitian=True, - ) - @ x_train.T - @ flat_targets - ) - y_hat = x_test @ beta_hat - y_hat_context = x_train_raw @ beta_hat if debug_info else None - outputs = [] outputs_context = [] + train_idx, test_idx = 0, 0 - # Reconstruct the ragged 2-dim batched forecasts from flattened linear fits. - train_index, test_index = 0, 0 - for train_index_delta, test_index_delta in zip(self.train_lens, self.test_lens): - outputs.append(np.array(y_hat[test_index : (test_index + test_index_delta)])) - if debug_info: - outputs_context.append( - np.array(y_hat_context[train_index : (train_index + train_index_delta)]) + with jax.default_device(device): + for trl, tel in zip(self.train_lens, self.test_lens): + x_tr = x_train_raw[train_idx : train_idx + trl] + x_te = x_test[test_idx : test_idx + tel] + y_tr = flat_targets[train_idx : train_idx + trl] + + x_tr_fit = x_tr.copy() + if max_rows_per_col: + nrows, ncols = x_tr_fit.shape + if nrows > (w := ncols * max_rows_per_col): + subsample = jax.random.choice( + jax.random.PRNGKey(max_rows_per_col_sample_seed), + nrows, + (w,), + replace=False, + ) + x_tr_fit = x_tr_fit[subsample] + y_tr = y_tr[subsample] + + x_tr_raw_j = _to_padded_jax_array(x_tr) + x_tr_j = _to_padded_jax_array(x_tr_fit) + y_tr_j = _to_padded_jax_array(y_tr) + x_te_j = _to_padded_jax_array(x_te) + + beta_hat = ( + jnp.linalg.pinv( + x_tr_j.T @ x_tr_j + ridge * jnp.eye(x_tr_j.shape[1]), + hermitian=True, + ) + @ x_tr_j.T + @ y_tr_j ) - train_index += train_index_delta - test_index += test_index_delta + outputs.append(np.array(x_te_j @ beta_hat)) + if debug_info: + outputs_context.append(np.array(x_tr_raw_j @ beta_hat)) + + train_idx += trl + test_idx += tel if debug_info: - return outputs, outputs_context, flat_targets, x_train, x_test + return ( + outputs, + outputs_context, + _to_padded_jax_array(flat_targets), + _to_padded_jax_array(x_train_raw), + _to_padded_jax_array(x_test), + ) else: return outputs diff --git a/v1/src/timesfm/xreg_lib.py b/v1/src/timesfm/xreg_lib.py index 0062a22..460b6e4 100644 --- a/v1/src/timesfm/xreg_lib.py +++ b/v1/src/timesfm/xreg_lib.py @@ -339,12 +339,20 @@ class BatchedInContextXRegBase: x_train = np.concatenate(x_train, axis=1) x_test = np.concatenate(x_test, axis=1) - # Normalize for robustness. - x_mean = np.mean(x_train, axis=0, keepdims=True) - x_std = np.where((w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w, - 1.0) - x_train = [(x_train - x_mean) / x_std] - x_test = [(x_test - x_mean) / x_std] + # Normalize per-input for robustness (batch-wide normalization + # would make each input's result depend on batch composition). + train_splits = np.cumsum(self.train_lens)[:-1] + test_splits = np.cumsum(self.test_lens)[:-1] + train_parts = np.split(x_train, train_splits, axis=0) + test_parts = np.split(x_test, test_splits, axis=0) + norm_train, norm_test = [], [] + for tr, te in zip(train_parts, test_parts): + m = np.mean(tr, axis=0, keepdims=True) + s = np.where((w := np.std(tr, axis=0, keepdims=True)) > _TOL, w, 1.0) + norm_train.append((tr - m) / s) + norm_test.append((te - m) / s) + x_train = [np.concatenate(norm_train, axis=0)] + x_test = [np.concatenate(norm_test, axis=0)] # Categorical features. Encode one by one. one_hot_encoder = preprocessing.OneHotEncoder( @@ -431,56 +439,55 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase): assert_covariate_shapes=assert_covariate_shapes, ) - x_train = x_train_raw.copy() - if max_rows_per_col: - nrows, ncols = x_train.shape - if nrows > (w := ncols * max_rows_per_col): - subsample = jax.random.choice( - jax.random.PRNGKey(max_rows_per_col_sample_seed), - nrows, - (w,), - replace=False, - ) - x_train = x_train[subsample] - flat_targets = flat_targets[subsample] - device = jax.devices("cpu")[0] if force_on_cpu else None - # Runs jitted version of the solvers which are quicker at the cost of - # running jitting during the first time calling. Re-jitting happens whenever - # new (padded) shapes are encountered. - # Ocassionally it helps with the speed and the accuracy if we force single - # thread execution on cpu for accelerator machines: - # 1. Avoid moving data to accelarator memory. - # 2. Avoid precision loss if any. - with jax.default_device(device): - x_train_raw = _to_padded_jax_array(x_train_raw) - x_train = _to_padded_jax_array(x_train) - flat_targets = _to_padded_jax_array(flat_targets) - x_test = _to_padded_jax_array(x_test) - beta_hat = (jnp.linalg.pinv( - x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]), - hermitian=True, - ) @ x_train.T @ flat_targets) - y_hat = x_test @ beta_hat - y_hat_context = x_train_raw @ beta_hat if debug_info else None + # Fit per-input regressions to prevent data leakage across batch items. outputs = [] outputs_context = [] + train_idx, test_idx = 0, 0 - # Reconstruct the ragged 2-dim batched forecasts from flattened linear fits. - train_index, test_index = 0, 0 - for train_index_delta, test_index_delta in zip(self.train_lens, - self.test_lens): - outputs.append(np.array(y_hat[test_index:(test_index + - test_index_delta)])) - if debug_info: - outputs_context.append( - np.array(y_hat_context[train_index:(train_index + - train_index_delta)])) - train_index += train_index_delta - test_index += test_index_delta + with jax.default_device(device): + for trl, tel in zip(self.train_lens, self.test_lens): + x_tr = x_train_raw[train_idx : train_idx + trl] + x_te = x_test[test_idx : test_idx + tel] + y_tr = flat_targets[train_idx : train_idx + trl] + + x_tr_fit = x_tr.copy() + if max_rows_per_col: + nrows, ncols = x_tr_fit.shape + if nrows > (w := ncols * max_rows_per_col): + subsample = jax.random.choice( + jax.random.PRNGKey(max_rows_per_col_sample_seed), + nrows, + (w,), + replace=False, + ) + x_tr_fit = x_tr_fit[subsample] + y_tr = y_tr[subsample] + + x_tr_raw_j = _to_padded_jax_array(x_tr) + x_tr_j = _to_padded_jax_array(x_tr_fit) + y_tr_j = _to_padded_jax_array(y_tr) + x_te_j = _to_padded_jax_array(x_te) + + beta_hat = (jnp.linalg.pinv( + x_tr_j.T @ x_tr_j + ridge * jnp.eye(x_tr_j.shape[1]), + hermitian=True, + ) @ x_tr_j.T @ y_tr_j) + outputs.append(np.array(x_te_j @ beta_hat)) + if debug_info: + outputs_context.append(np.array(x_tr_raw_j @ beta_hat)) + + train_idx += trl + test_idx += tel if debug_info: - return outputs, outputs_context, flat_targets, x_train, x_test + return ( + outputs, + outputs_context, + _to_padded_jax_array(flat_targets), + _to_padded_jax_array(x_train_raw), + _to_padded_jax_array(x_test), + ) else: return outputs From 1bb44d5eefa7821cf13c9c5e06f8f68ef3a55847 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 14:18:02 -0400 Subject: [PATCH 10/17] fix: respect batch_size in v1 data_loader when permute=False Apply changes from PR #391 by @MarcoGorworworelli: - Fix train_gen() to iterate in proper batch_size chunks instead of yielding all time series at once when permute=False - Add test_data_loader.py to verify batch boundaries --- v1/src/timesfm/data_loader.py | 5 ++-- v1/tests/test_data_loader.py | 50 +++++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) create mode 100644 v1/tests/test_data_loader.py diff --git a/v1/src/timesfm/data_loader.py b/v1/src/timesfm/data_loader.py index d81b130..e2fc259 100644 --- a/v1/src/timesfm/data_loader.py +++ b/v1/src/timesfm/data_loader.py @@ -149,11 +149,12 @@ class TimeSeriesdata(object): else: epoch_len = self.epoch_len for idx in perm[0:epoch_len]: - for _ in range(num_ts // self.batch_size + 1): + batch_indices = range(0, num_ts, self.batch_size) + for batch_idx in batch_indices: if self.permute: tsidx = np.random.choice(num_ts, size=self.batch_size, replace=False) else: - tsidx = np.arange(num_ts) + tsidx = np.arange(batch_idx, min(batch_idx + self.batch_size, num_ts)) dtimes = np.arange(idx - hist_len, idx + self.pred_len) ( bts_train, diff --git a/v1/tests/test_data_loader.py b/v1/tests/test_data_loader.py new file mode 100644 index 0000000..ffee145 --- /dev/null +++ b/v1/tests/test_data_loader.py @@ -0,0 +1,50 @@ +from pathlib import Path + +import numpy as np +import pandas as pd + +from timesfm.data_loader import TimeSeriesdata + + +def test_train_gen_respects_batch_size_when_permute_is_false(tmp_path: Path) -> None: + rows = 12 + df = pd.DataFrame( + { + "ds": pd.date_range("2024-01-01", periods=rows, freq="D"), + "ts_1": np.arange(rows), + "ts_2": np.arange(rows) + 10, + "ts_3": np.arange(rows) + 20, + "ts_4": np.arange(rows) + 30, + "ts_5": np.arange(rows) + 40, + } + ) + data_path = tmp_path / "sample.csv" + df.to_csv(data_path, index=False) + + loader = TimeSeriesdata( + data_path=str(data_path), + datetime_col="ds", + num_cov_cols=None, + cat_cov_cols=None, + ts_cols=np.array(["ts_1", "ts_2", "ts_3", "ts_4", "ts_5"]), + train_range=[0, 8], + val_range=[8, 10], + test_range=[10, 12], + hist_len=3, + pred_len=2, + batch_size=2, + freq="D", + normalize=False, + epoch_len=1, + holiday=False, + permute=False, + ) + + batches = list(loader.train_gen()) + ts_indices = [batch[-1].tolist() for batch in batches] + + assert ts_indices == [[0, 1], [2, 3], [4]] + for batch in batches: + assert len(batch[-1]) <= 2 + assert batch[0].shape[0] == len(batch[-1]) + assert batch[3].shape[0] == len(batch[-1]) From 30f28a1b1d0c31070f93462742575d64e53069e2 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 14:18:15 -0400 Subject: [PATCH 11/17] fix: correct SKILL.md link in README Apply changes from PR #390 by @amansinghbais: - Fix link to point to the actual SKILL.md file instead of the directory --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 40afca3..2d12958 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ This open version is not an officially supported Google product. ## Update - Mar. 19, 2026 -Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/tree/master/timesfm-forecasting) is out. +Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/blob/master/timesfm-forecasting/SKILL.md) is out. ## Update - Oct. 29, 2025 From 13a8eb2a25000777260d466f19f8c23191599a56 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 14:19:08 -0400 Subject: [PATCH 12/17] ci: upgrade GitHub Actions to v6 Apply changes from PR #367 by @Copilot: - actions/checkout v2 -> v6 - actions/setup-python v2 -> v6 --- .github/workflows/main.yml | 4 ++-- .github/workflows/manual_publish.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 6bee2a1..30c3423 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -10,9 +10,9 @@ jobs: build: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v2 + uses: actions/setup-python@v6 with: python-version: '3.11' - name: Install uv diff --git a/.github/workflows/manual_publish.yml b/.github/workflows/manual_publish.yml index 3cd30e9..82de6fc 100644 --- a/.github/workflows/manual_publish.yml +++ b/.github/workflows/manual_publish.yml @@ -7,9 +7,9 @@ jobs: build-and-publish: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v2 + uses: actions/setup-python@v6 with: python-version: '3.11' - name: Install uv From 54f5405b7da738bc35751c4e12b48e0beb4930f1 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 14:20:40 -0400 Subject: [PATCH 13/17] docs: fix swapped xreg_mode descriptions and typo in error message Apply changes from PR #366 by @cj-wong: - Correct xreg_mode docstring: descriptions for 'xreg + timesfm' and 'timesfm + xreg' were swapped - Fix 'covaraites' -> 'covariates' typo in error message --- src/timesfm/timesfm_2p5/timesfm_2p5_base.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py index 22f0df7..151abcb 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py @@ -221,9 +221,10 @@ class TimesFM_2p5: dynamic_categorical_covariates: A dict of dynamic categorical covariates. static_numerical_covariates: A dict of static numerical covariates. static_categorical_covariates: A dict of static categorical covariates. - xreg_mode: one of "xreg + timesfm" or "timesfm + xreg". "timesfm + xreg" - fits a model on the residuals of the TimesFM forecast. "xreg + timesfm" - fits a model on the targets then forecasts on the residuals via TimesFM. + xreg_mode: one of "xreg + timesfm" or "timesfm + xreg". "xreg + timesfm" + first fits an XReg model on the targets, then uses TimesFM to forecast + the residuals. "timesfm + xreg" first runs TimesFM to get a forecast, + then fits an XReg model on the residuals of that forecast. normalize_xreg_target_per_input: whether to normalize the xreg target per input in the given batch. ridge: ridge penalty for the linear model. @@ -285,7 +286,7 @@ class TimesFM_2p5: if test_lens[-1] > self.forecast_config.max_horizon: raise ValueError( - "Forecast horizon length inferred from the dynamic covaraites is longer than the" + "Forecast horizon length inferred from the dynamic covariates is longer than the" f"max_horizon defined in the forecast config: {test_lens[-1]} > {self.forecast_config.max_horizon=}." ) From ad192b7954dd6edf502563c43c8613c88bdb7c41 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 15:25:53 -0400 Subject: [PATCH 14/17] =?UTF-8?q?docs:=20update=20README=20=E2=80=94=20rep?= =?UTF-8?q?lace=20'under=20construction'=20with=20completed=20status?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add Apr. 2026 update entry for PEFT pipeline, unit tests, and community fixes - Replace 'under construction' numbered list with checklist of completed items: Flax model, covariate support, docs/examples, PEFT pipeline, unit tests --- README.md | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 2d12958..d7a2baf 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,13 @@ 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 + +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. + ## Update - Mar. 19, 2026 Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/blob/master/timesfm-forecasting/SKILL.md) is out. @@ -44,12 +51,13 @@ Comparing to TimesFM 2.0, this new 2.5 model: - gets rid of the `frequency` indicator. - has a couple of new forecasting flags. -Along with the model upgrade we have also upgraded the inference API. This repo -will be under construction over the next few weeks to +Since the Sept. 2025 launch, the following improvements have been completed: -1. add support for an upcoming Flax version of the model (faster inference). -2. add back covariate support. -3. populate more docstrings, docs and notebook. +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/`). +5. ✅ Unit tests for core layers, configs, and utilities (see `tests/`). ### Install From 18d5eb2d441e4340d9700f4007628a221f0ac145 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Wed, 8 Apr 2026 21:40:06 -0400 Subject: [PATCH 15/17] fix: improve PEFT device consistency and XReg output slicing - Initialize LoRA parameters on the same device as the base linear layer - Load adapter weights directly to the model device instead of hardcoded CPU - Slice XReg linear regression outputs to match the specified sequence lengths - Replace batch-wide covariate normalization with per-input normalization in create_covariate_matrix to prevent cross-input scale leakage. - Refactor BatchedInContextXRegLinear.fit to solve ridge regression per instance rather than as a single global matrix solve, avoiding cross-contamination between batched inputs. - Truncate JAX regression outputs to the actual train/test lengths after the padded matrix multiply, fixing shape mismatches for non-power-of-2 horizons (e.g. horizon=24 was returning 32 elements). --- peft/adapters.py | 8 +++++--- src/timesfm/utils/xreg_lib.py | 4 ++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/peft/adapters.py b/peft/adapters.py index 951cfe1..acc6eee 100644 --- a/peft/adapters.py +++ b/peft/adapters.py @@ -110,9 +110,10 @@ class DoRALinear(nn.Module): 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)) - self.lora_B = nn.Parameter(torch.zeros(rank, out_f)) + 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. @@ -272,7 +273,8 @@ def load_adapter_weights(model: nn.Module, path: str) -> None: The model must already have adapters injected (via ``inject_adapters``) before calling this function. """ - tensors = load_file(path, device="cpu") + 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: diff --git a/src/timesfm/utils/xreg_lib.py b/src/timesfm/utils/xreg_lib.py index 4355261..2759b67 100644 --- a/src/timesfm/utils/xreg_lib.py +++ b/src/timesfm/utils/xreg_lib.py @@ -509,9 +509,9 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase): @ x_tr_j.T @ y_tr_j ) - outputs.append(np.array(x_te_j @ beta_hat)) + outputs.append(np.array(x_te_j @ beta_hat)[:tel]) if debug_info: - outputs_context.append(np.array(x_tr_raw_j @ beta_hat)) + outputs_context.append(np.array(x_tr_raw_j @ beta_hat)[:trl]) train_idx += trl test_idx += tel From caddef1db8ea24434212b7bada409c0a34c87b95 Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Thu, 9 Apr 2026 11:15:58 -0400 Subject: [PATCH 16/17] refactor: replace custom PEFT pipeline with Transformers+PEFT example 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/) --- .gitignore | 1 - README.md | 12 +- peft/README.md | 201 ------ peft/__init__.py | 41 -- peft/adapters.py | 285 --------- peft/config.py | 106 ---- peft/data.py | 150 ----- peft/finetune.py | 252 -------- peft/finetune.sh | 73 --- peft/trainer.py | 578 ------------------ .../examples/finetuning/README.md | 102 ++++ .../examples/finetuning/finetune_lora.py | 446 ++++++++++++++ 12 files changed, 554 insertions(+), 1693 deletions(-) delete mode 100644 peft/README.md delete mode 100644 peft/__init__.py delete mode 100644 peft/adapters.py delete mode 100644 peft/config.py delete mode 100644 peft/data.py delete mode 100644 peft/finetune.py delete mode 100644 peft/finetune.sh delete mode 100644 peft/trainer.py create mode 100644 timesfm-forecasting/examples/finetuning/README.md create mode 100644 timesfm-forecasting/examples/finetuning/finetune_lora.py 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() From 6ae67d41d813fcdab0a1bc785b79053c3769a63e Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Thu, 9 Apr 2026 23:00:36 -0400 Subject: [PATCH 17/17] revert: drop PR #393 (xreg batch behavior) and PR #390 (SKILL.md link) per maintainer feedback --- .gitignore | 1 + README.md | 5 +- src/timesfm/utils/xreg_lib.py | 110 ++++++++++++++++------------------ 3 files changed, 54 insertions(+), 62 deletions(-) diff --git a/.gitignore b/.gitignore index 24495b0..00dc636 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ datasets/ results/ uv.lock development_setup.md +debug.log diff --git a/README.md b/README.md index 56e8253..c7b2d50 100644 --- a/README.md +++ b/README.md @@ -26,12 +26,11 @@ This open version is not an officially supported Google product. 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. +Also added unit tests (`tests/`) and incorporated several community fixes. ## Update - Mar. 19, 2026 -Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/blob/master/timesfm-forecasting/SKILL.md) is out. +Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/tree/master/timesfm-forecasting) is out. ## Update - Oct. 29, 2025 diff --git a/src/timesfm/utils/xreg_lib.py b/src/timesfm/utils/xreg_lib.py index 2759b67..7a1b19b 100644 --- a/src/timesfm/utils/xreg_lib.py +++ b/src/timesfm/utils/xreg_lib.py @@ -370,20 +370,11 @@ class BatchedInContextXRegBase: x_train = np.concatenate(x_train, axis=1) x_test = np.concatenate(x_test, axis=1) - # Normalize per-input for robustness (batch-wide normalization - # would make each input's result depend on batch composition). - train_splits = np.cumsum(self.train_lens)[:-1] - test_splits = np.cumsum(self.test_lens)[:-1] - train_parts = np.split(x_train, train_splits, axis=0) - test_parts = np.split(x_test, test_splits, axis=0) - norm_train, norm_test = [], [] - for tr, te in zip(train_parts, test_parts): - m = np.mean(tr, axis=0, keepdims=True) - s = np.where((w := np.std(tr, axis=0, keepdims=True)) > _TOL, w, 1.0) - norm_train.append((tr - m) / s) - norm_test.append((te - m) / s) - x_train = [np.concatenate(norm_train, axis=0)] - x_test = [np.concatenate(norm_test, axis=0)] + # Normalize for robustness. + x_mean = np.mean(x_train, axis=0, keepdims=True) + x_std = np.where((w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w, 1.0) + x_train = [(x_train - x_mean) / x_std] + x_test = [(x_test - x_mean) / x_std] # Categorical features. Encode one by one. one_hot_encoder = preprocessing.OneHotEncoder( @@ -472,57 +463,58 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase): assert_covariate_shapes=assert_covariate_shapes, ) + x_train = x_train_raw.copy() + if max_rows_per_col: + nrows, ncols = x_train.shape + if nrows > (w := ncols * max_rows_per_col): + subsample = jax.random.choice( + jax.random.PRNGKey(max_rows_per_col_sample_seed), + nrows, + (w,), + replace=False, + ) + x_train = x_train[subsample] + flat_targets = flat_targets[subsample] + device = jax.devices("cpu")[0] if force_on_cpu else None + # Runs jitted version of the solvers which are quicker at the cost of + # running jitting during the first time calling. Re-jitting happens whenever + # new (padded) shapes are encountered. + # Ocassionally it helps with the speed and the accuracy if we force single + # thread execution on cpu for accelerator machines: + # 1. Avoid moving data to accelarator memory. + # 2. Avoid precision loss if any. + with jax.default_device(device): + x_train_raw = _to_padded_jax_array(x_train_raw) + x_train = _to_padded_jax_array(x_train) + flat_targets = _to_padded_jax_array(flat_targets) + x_test = _to_padded_jax_array(x_test) + beta_hat = ( + jnp.linalg.pinv( + x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]), + hermitian=True, + ) + @ x_train.T + @ flat_targets + ) + y_hat = x_test @ beta_hat + y_hat_context = x_train_raw @ beta_hat if debug_info else None + outputs = [] outputs_context = [] - train_idx, test_idx = 0, 0 - with jax.default_device(device): - for trl, tel in zip(self.train_lens, self.test_lens): - x_tr = x_train_raw[train_idx : train_idx + trl] - x_te = x_test[test_idx : test_idx + tel] - y_tr = flat_targets[train_idx : train_idx + trl] - - x_tr_fit = x_tr.copy() - if max_rows_per_col: - nrows, ncols = x_tr_fit.shape - if nrows > (w := ncols * max_rows_per_col): - subsample = jax.random.choice( - jax.random.PRNGKey(max_rows_per_col_sample_seed), - nrows, - (w,), - replace=False, - ) - x_tr_fit = x_tr_fit[subsample] - y_tr = y_tr[subsample] - - x_tr_raw_j = _to_padded_jax_array(x_tr) - x_tr_j = _to_padded_jax_array(x_tr_fit) - y_tr_j = _to_padded_jax_array(y_tr) - x_te_j = _to_padded_jax_array(x_te) - - beta_hat = ( - jnp.linalg.pinv( - x_tr_j.T @ x_tr_j + ridge * jnp.eye(x_tr_j.shape[1]), - hermitian=True, - ) - @ x_tr_j.T - @ y_tr_j + # Reconstruct the ragged 2-dim batched forecasts from flattened linear fits. + train_index, test_index = 0, 0 + for train_index_delta, test_index_delta in zip(self.train_lens, self.test_lens): + outputs.append(np.array(y_hat[test_index : (test_index + test_index_delta)])) + if debug_info: + outputs_context.append( + np.array(y_hat_context[train_index : (train_index + train_index_delta)]) ) - outputs.append(np.array(x_te_j @ beta_hat)[:tel]) - if debug_info: - outputs_context.append(np.array(x_tr_raw_j @ beta_hat)[:trl]) - - train_idx += trl - test_idx += tel + train_index += train_index_delta + test_index += test_index_delta if debug_info: - return ( - outputs, - outputs_context, - _to_padded_jax_array(flat_targets), - _to_padded_jax_array(x_train_raw), - _to_padded_jax_array(x_test), - ) + return outputs, outputs_context, flat_targets, x_train, x_test else: return outputs