feat: add LoRA/DoRA adapter layers for TimesFM 2.5 (PyTorch)

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

@@ -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",
+]

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)

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