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