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