From 65700643df8940f63d4bb016f1a7fb052671f42c Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Mon, 20 Jan 2025 21:48:53 +0000
Subject: [PATCH 01/12] Add finetuning support

---
 notebooks/finetuning_torch.py | 274 ++++++++++++++++++++++++++++++++++
 1 file changed, 274 insertions(+)
 create mode 100644 notebooks/finetuning_torch.py

diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py
new file mode 100644
index 0000000..50a7f4d
--- /dev/null
+++ b/notebooks/finetuning_torch.py
@@ -0,0 +1,274 @@
+# Filename: tutorial_timesfm.py
+
+import yfinance as yf
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset, DataLoader
+import torch.optim as optim
+import timesfm
+from os import path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+
+
+from timesfm.pytorch_patched_decoder import TimesFMConfig, PatchedTimeSeriesDecoder
+
+import torch
+import matplotlib.pyplot as plt
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# --------------------------------------------------
+# 1. Download stock data via yfinance
+# --------------------------------------------------
+def download_yfinance_data(ticker="AAPL", start="2020-01-01", end="2022-01-01"):
+    """
+    Download daily stock data for a given ticker from Yahoo Finance.
+    Returns a pandas DataFrame with columns like 'Open', 'High', 'Low', 'Close', 'Volume'.
+    """
+    df = yf.download(ticker, start=start, end=end)
+    df = df.dropna()
+    return df["Close"].reset_index(drop=True)
+
+
+# --------------------------------------------------
+# 2. Create a dataset class for TimesFM
+# --------------------------------------------------
+class FinancialDataset(Dataset):
+    def __init__(
+        self,
+        series: pd.Series,
+        config: TimesFMConfig,
+        context_length=128,  # how many past timesteps as input
+        horizon_length=32,  # how many future steps to predict
+    ):
+        super().__init__()
+
+        self.series = series.values.astype(np.float32)
+        self.context_length = context_length
+        self.horizon_length = horizon_length
+        self.config = config
+
+        self.samples = []
+        # We want to ensure we have at least context_length + horizon_length points.
+        for start_idx in range(0, len(self.series) - (context_length + horizon_length)):
+            end_idx = start_idx + context_length
+            # context slice
+            x_context = self.series[start_idx:end_idx]
+            # future/horizon slice
+            x_future = self.series[end_idx : end_idx + horizon_length]
+            self.samples.append((x_context, x_future))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        x_context, x_future = self.samples[index]
+        # Convert to torch
+        x_context = torch.tensor(x_context, dtype=torch.float32)
+        x_future = torch.tensor(x_future, dtype=torch.float32)
+
+        input_padding = torch.zeros_like(x_context)
+
+        freq = torch.zeros(1, dtype=torch.long)
+
+        return x_context, input_padding, freq, x_future
+
+
+def collate_fn(batch):
+    xs_context = [item[0] for item in batch]
+    xs_padding = [item[1] for item in batch]
+    freqs = [item[2] for item in batch]
+    xs_future = [item[3] for item in batch]
+
+    x_context = torch.stack(xs_context, dim=0)
+    input_pad = torch.stack(xs_padding, dim=0)
+    freq = torch.stack(freqs, dim=0)  # shape [B, 1]
+    x_future = torch.stack(xs_future, dim=0)
+
+    return x_context, input_pad, freq, x_future
+
+
+def get_model(*, load_weights: bool = False):
+    # standard model hack
+    repo_id = "google/timesfm-2.0-500m-pytorch"
+    tfm = timesfm.TimesFm(
+        hparams=timesfm.TimesFmHparams(
+            backend="cuda",
+            per_core_batch_size=32,
+            horizon_len=128,
+            num_layers=50,
+            use_positional_embedding=False,
+            context_len=192,
+        ),
+        checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=repo_id),
+    )
+
+    model = PatchedTimeSeriesDecoder(tfm._model_config)
+
+    if load_weights:
+        checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
+        print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
+        loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
+        model.load_state_dict(loaded_checkpoint)
+        print("After loading:")
+        print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
+    model = model.to(device)
+
+    # import sys
+    # sys.exit(-1)
+    # repo_id = "google/timesfm-1.0-200m"
+    return model, tfm._model_config
+
+
+def train_model(
+    ticker="AAPL", start="2015-01-01", end="2022-01-01", train_split=0.8, batch_size=8, num_epochs=20, pretrained=False
+):
+    df_close = download_yfinance_data(ticker, start=start, end=end)
+    model, config = get_model(load_weights=pretrained)
+
+    total_len = len(df_close)
+    train_size = int(total_len * train_split)
+    val_size = total_len - train_size
+
+    train_series = df_close.iloc[:train_size].reset_index(drop=True)
+    val_series = df_close.iloc[train_size:].reset_index(drop=True)
+
+    train_dataset = FinancialDataset(
+        series=train_series, config=config, context_length=128, horizon_length=config.horizon_len
+    )
+    val_dataset = FinancialDataset(
+        series=val_series, config=config, context_length=128, horizon_length=config.horizon_len
+    )
+    print("Train samples:", len(train_dataset))
+    print("Val samples:", len(val_dataset))
+    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
+    val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
+
+    optimizer = optim.Adam(model.parameters(), lr=1e-4)
+
+    for epoch in range(num_epochs):
+        model.train()
+        total_train_loss = 0.0
+
+        for x_context, x_padding, freq, x_future in train_dataloader:
+            x_context, x_padding, freq, x_future = (
+                x_context.to(device),
+                x_padding.to(device),
+                freq.to(device),
+                x_future.to(device),
+            )
+            predictions = model(x_context, x_padding.float(), freq)
+            # predictions shape => [B, N, horizon_len, (1 + #quantiles)]
+            predictions_mean = predictions[..., 0]  # => [B, N, horizon_len]
+            last_patch_pred = predictions_mean[:, -1, :]  # => [B, horizon_len]
+
+            # x_future => [B, horizon_len]
+            loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            total_train_loss += loss.item()
+
+        avg_train_loss = total_train_loss / len(train_dataloader)
+
+        # -------- Compute validation loss --------
+        model.eval()
+        total_val_loss = 0.0
+        with torch.no_grad():
+            for x_context, x_padding, freq, x_future in val_dataloader:
+                x_context, x_padding, freq, x_future = (
+                    x_context.to(device),
+                    x_padding.to(device),
+                    freq.to(device),
+                    x_future.to(device),
+                )
+                predictions = model(x_context, x_padding.float(), freq)
+                predictions_mean = predictions[..., 0]
+                last_patch_pred = predictions_mean[:, -1, :]
+                val_loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
+                total_val_loss += val_loss.item()
+
+        avg_val_loss = total_val_loss / max(len(val_dataloader), 1)
+
+        print(f"[Epoch {epoch+1}] Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
+
+    torch.save(model.state_dict(), "timesfm_finetuned.pth")
+    return model, train_dataloader, val_dataloader
+
+
+def plot_predictions(model, dataloader):
+    model.eval()
+    with torch.no_grad():
+        x_context, x_padding, freq, x_future = next(iter(dataloader))
+        x_context, x_padding, freq, x_future = (
+            x_context.to(device),
+            x_padding.to(device),
+            freq.to(device),
+            x_future.to(device),
+        )
+        # Forward pass
+        predictions = model(x_context, x_padding.float(), freq)
+        # => [B, N, horizon_len, (1 + #quantiles)]
+        predictions_mean = predictions[..., 0]  # => [B, N, horizon_len]
+        last_patch_prediction = predictions_mean[:, -1, :]  # => [B, horizon_len]
+
+        # We'll plot only the first sample in the batch
+        i = 0
+        pred_vals = last_patch_prediction[i].cpu().numpy()  # [horizon_len]
+        context_vals = x_context[i].cpu().numpy()  # [context_len]
+        future_vals = x_future[i].cpu().numpy()  # [horizon_len]
+
+        horizon_len = future_vals.shape[0]
+        context_len = context_vals.shape[0]
+
+        plt.figure(figsize=(10, 5))
+
+        # Plot context
+        plt.plot(range(context_len), context_vals, label="Context (History)", color="blue")
+
+        # Plot predicted future
+        plt.plot(
+            range(context_len, context_len + horizon_len),
+            pred_vals,
+            label="Predicted Future",
+            color="orange",
+        )
+
+        # Plot ground truth future
+        plt.plot(
+            range(context_len, context_len + horizon_len),
+            future_vals,
+            label="Ground Truth Future",
+            color="green",
+            linestyle="--",
+        )
+
+        plt.xlabel("Time")
+        plt.ylabel("Value")
+        plt.title("Model Forecast vs. Ground Truth")
+        plt.legend()
+        plt.show()
+        plt.savefig("pic_predictions.png")
+
+
+if __name__ == "__main__":
+    # Example usage
+    model, train_dl, val_dl = train_model(
+        ticker="AAPL",
+        start="2012-01-01",
+        end="2019-01-01",
+        train_split=0.8,
+        batch_size=256,
+        num_epochs=50,
+        pretrained=True,
+    )
+
+    plot_predictions(model, val_dl)

From 4942cda83b7a59eb649834bc82363c5d4ce2ab85 Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Tue, 21 Jan 2025 01:45:30 +0000
Subject: [PATCH 02/12] Refactor into 2 examples

---
 notebooks/finetuning_example.py | 224 +++++++++++++++++++
 notebooks/finetuning_torch.py   | 381 +++++++++++++-------------------
 2 files changed, 377 insertions(+), 228 deletions(-)
 create mode 100644 notebooks/finetuning_example.py

diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
new file mode 100644
index 0000000..5664456
--- /dev/null
+++ b/notebooks/finetuning_example.py
@@ -0,0 +1,224 @@
+"""
+Example usage of the TimesFM Finetuning Framework.
+"""
+
+import yfinance as yf
+import torch
+from os import path
+import numpy as np
+from torch.utils.data import Dataset
+from timesfm import TimesFm, TimesFmHparams, TimesFmCheckpoint
+from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder
+from finetuning_torch import FinetuningConfig, TimesFMFinetuner
+from huggingface_hub import snapshot_download
+import numpy as np
+import pandas as pd
+from torch.utils.data import Dataset
+import torch
+import yfinance as yf
+from typing import Tuple, Optional
+
+from timesfm import TimesFm, TimesFmHparams
+
+
+class TimeSeriesDataset(Dataset):
+    """Dataset for time series data compatible with TimesFM."""
+
+    def __init__(self, series: np.ndarray, context_length: int, horizon_length: int):
+        """
+        Initialize dataset.
+
+        Args:
+            series: Time series data
+            context_length: Number of past timesteps to use as input
+            horizon_length: Number of future timesteps to predict
+        """
+        self.series = series
+        self.context_length = context_length
+        self.horizon_length = horizon_length
+        self._prepare_samples()
+
+    def _prepare_samples(self) -> None:
+        """Prepare sliding window samples from the time series."""
+        self.samples = []
+        total_length = self.context_length + self.horizon_length
+
+        for start_idx in range(0, len(self.series) - total_length + 1):
+            end_idx = start_idx + self.context_length
+            x_context = self.series[start_idx:end_idx]
+            x_future = self.series[end_idx : end_idx + self.horizon_length]
+            self.samples.append((x_context, x_future))
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        x_context, x_future = self.samples[index]
+
+        x_context = torch.tensor(x_context, dtype=torch.float32)
+        x_future = torch.tensor(x_future, dtype=torch.float32)
+
+        input_padding = torch.zeros_like(x_context)
+        freq = torch.zeros(1, dtype=torch.long)
+
+        return x_context, input_padding, freq, x_future
+
+
+def prepare_datasets(
+    series: np.ndarray, context_length: int, horizon_length: int, train_split: float = 0.8
+) -> Tuple[Dataset, Dataset]:
+    """
+    Prepare training and validation datasets from time series data.
+
+    Args:
+        series: Input time series data
+        context_length: Number of past timesteps to use
+        horizon_length: Number of future timesteps to predict
+        train_split: Fraction of data to use for training
+
+    Returns:
+        Tuple of (train_dataset, val_dataset)
+    """
+    train_size = int(len(series) * train_split)
+    train_data = series[:train_size]
+    val_data = series[train_size:]
+
+    # Create datasets
+    train_dataset = TimeSeriesDataset(train_data, context_length=context_length, horizon_length=horizon_length)
+
+    val_dataset = TimeSeriesDataset(val_data, context_length=context_length, horizon_length=horizon_length)
+
+    return train_dataset, val_dataset
+
+
+def get_model(load_weights: bool = False):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    repo_id = "google/timesfm-2.0-500m-pytorch"
+    hparams = TimesFmHparams(
+        backend=device,
+        per_core_batch_size=32,
+        horizon_len=128,
+        num_layers=50,
+        use_positional_embedding=False,
+        context_len=192,
+    )
+    tfm = TimesFm(hparams=hparams, checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
+
+    model = PatchedTimeSeriesDecoder(tfm._model_config)
+    if load_weights:
+        checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
+        loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
+        model.load_state_dict(loaded_checkpoint)
+    model = model.to(device)
+    return model, hparams, tfm._model_config
+
+
+def plot_predictions(
+    model: TimesFm,
+    val_dataset: Dataset,
+    save_path: Optional[str] = "predictions.png",
+) -> None:
+    """
+    Plot model predictions against ground truth for a batch of validation data.
+
+    Args:
+        model: Trained TimesFM model
+        val_dataset: Validation dataset
+        save_path: Path to save the plot
+    """
+    import matplotlib.pyplot as plt
+
+    model.eval()
+
+    x_context, x_padding, freq, x_future = val_dataset[0]
+    x_context = x_context.unsqueeze(0)  # Add batch dimension
+    x_padding = x_padding.unsqueeze(0)
+    freq = freq.unsqueeze(0)
+    x_future = x_future.unsqueeze(0)
+
+    device = next(model.parameters()).device
+    x_context = x_context.to(device)
+    x_padding = x_padding.to(device)
+    freq = freq.to(device)
+    x_future = x_future.to(device)
+
+    with torch.no_grad():
+        predictions = model(x_context, x_padding.float(), freq)
+        predictions_mean = predictions[..., 0]  # [B, N, horizon_len]
+        last_patch_pred = predictions_mean[:, -1, :]  # [B, horizon_len]
+
+    context_vals = x_context[0].cpu().numpy()
+    future_vals = x_future[0].cpu().numpy()
+    pred_vals = last_patch_pred[0].cpu().numpy()
+
+    context_len = len(context_vals)
+    horizon_len = len(future_vals)
+
+    plt.figure(figsize=(12, 6))
+
+    plt.plot(range(context_len), context_vals, label="Historical Data", color="blue", linewidth=2)
+
+    plt.plot(
+        range(context_len, context_len + horizon_len),
+        future_vals,
+        label="Ground Truth",
+        color="green",
+        linestyle="--",
+        linewidth=2,
+    )
+
+    plt.plot(range(context_len, context_len + horizon_len), pred_vals, label="Prediction", color="red", linewidth=2)
+
+    plt.xlabel("Time Step")
+    plt.ylabel("Value")
+    plt.title("TimesFM Predictions vs Ground Truth")
+    plt.legend()
+    plt.grid(True)
+
+    if save_path:
+        plt.savefig(save_path)
+        print(f"Plot saved to {save_path}")
+
+    plt.close()
+
+
+def get_data(context_len: int, horizon_len: int) -> Tuple[Dataset, Dataset]:
+    df = yf.download("AAPL", start="2010-01-01", end="2019-01-01")
+    time_series = df["Close"].values
+
+    train_dataset, val_dataset = prepare_datasets(
+        series=time_series,
+        context_length=context_len,
+        horizon_length=horizon_len,
+        train_split=0.8,
+    )
+
+    print(f"Created datasets:")
+    print(f"- Training samples: {len(train_dataset)}")
+    print(f"- Validation samples: {len(val_dataset)}")
+    return train_dataset, val_dataset
+
+
+def basic_example():
+    """Basic example of finetuning TimesFM on stock data."""
+    model, hparams, tfm_config = get_model(load_weights=True)
+    config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=False)
+
+    train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
+    finetuner = TimesFMFinetuner(model, config)
+
+    print("\nStarting finetuning...")
+    results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
+
+    print("\nFinetuning completed!")
+    print(f"Training history: {len(results['history']['train_loss'])} epochs")
+
+    plot_predictions(
+        model=model,
+        val_dataset=val_dataset,
+        save_path="timesfm_predictions.png",
+    )
+
+
+if __name__ == "__main__":
+    basic_example()
diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py
index 50a7f4d..bf908e7 100644
--- a/notebooks/finetuning_torch.py
+++ b/notebooks/finetuning_torch.py
@@ -1,274 +1,199 @@
-# Filename: tutorial_timesfm.py
+"""
+TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
+
+Example usage:
+    ```python
+    # Prepare datasets
+    train_dataset = TimeSeriesDataset(train_data, context_length=128, horizon_length=32)
+    val_dataset = TimeSeriesDataset(val_data, context_length=128, horizon_length=32)
+    
+    # Initialize model and configuration
+    model = TimesFm(...)
+    config = FinetuningConfig(
+        batch_size=64,
+        num_epochs=50,
+        learning_rate=1e-4,
+        use_wandb=True
+    )
+    
+    # Create finetuner
+    finetuner = TimesFMFinetuner(model, config)
+    
+    # Finetune model
+    results = finetuner.finetune(train_dataset, val_dataset)
+    ```
+"""
+
+import abc
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Dict, Any
 
-import yfinance as yf
-import numpy as np
-import pandas as pd
 import torch
 from torch.utils.data import Dataset, DataLoader
 import torch.optim as optim
-import timesfm
-from os import path
-from typing import Any, Sequence
+from torch.nn.parallel import DistributedDataParallel as DDP
+import wandb
+import multiprocessing as mp
 
-import numpy as np
-import torch
-from huggingface_hub import snapshot_download
+from timesfm import TimesFm
 
 
-from timesfm.pytorch_patched_decoder import TimesFMConfig, PatchedTimeSeriesDecoder
+@dataclass
+class FinetuningConfig:
+    """Configuration for TimesFM finetuning process."""
 
-import torch
-import matplotlib.pyplot as plt
+    # Training parameters
+    batch_size: int = 32
+    num_epochs: int = 20
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.01
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Hardware parameters
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    distributed: bool = False
+    world_size: int = 1
+
+    # Logging parameters
+    use_wandb: bool = False
+    wandb_project: str = "timesfm-finetuning"
 
 
-# --------------------------------------------------
-# 1. Download stock data via yfinance
-# --------------------------------------------------
-def download_yfinance_data(ticker="AAPL", start="2020-01-01", end="2022-01-01"):
-    """
-    Download daily stock data for a given ticker from Yahoo Finance.
-    Returns a pandas DataFrame with columns like 'Open', 'High', 'Low', 'Close', 'Volume'.
-    """
-    df = yf.download(ticker, start=start, end=end)
-    df = df.dropna()
-    return df["Close"].reset_index(drop=True)
+class TimesFMFinetuner:
+    """Main class for finetuning TimesFM models."""
 
-
-# --------------------------------------------------
-# 2. Create a dataset class for TimesFM
-# --------------------------------------------------
-class FinancialDataset(Dataset):
     def __init__(
         self,
-        series: pd.Series,
-        config: TimesFMConfig,
-        context_length=128,  # how many past timesteps as input
-        horizon_length=32,  # how many future steps to predict
+        model: TimesFm,
+        config: FinetuningConfig,
+        loss_fn: Optional[callable] = None,
+        logger: Optional[logging.Logger] = None,
     ):
-        super().__init__()
+        """
+        Initialize TimesFM finetuner.
 
-        self.series = series.values.astype(np.float32)
-        self.context_length = context_length
-        self.horizon_length = horizon_length
+        Args:
+            model: TimesFM model to finetune
+            config: Finetuning configuration
+            logger: Optional logger instance
+        """
+        self.model = model
         self.config = config
+        self.logger = logger or logging.getLogger(__name__)
 
-        self.samples = []
-        # We want to ensure we have at least context_length + horizon_length points.
-        for start_idx in range(0, len(self.series) - (context_length + horizon_length)):
-            end_idx = start_idx + context_length
-            # context slice
-            x_context = self.series[start_idx:end_idx]
-            # future/horizon slice
-            x_future = self.series[end_idx : end_idx + horizon_length]
-            self.samples.append((x_context, x_future))
+        self.device = torch.device(config.device)
+        self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))  # MSELoss()
 
-    def __len__(self):
-        return len(self.samples)
+        if config.use_wandb:
+            self._setup_wandb()
 
-    def __getitem__(self, index):
-        x_context, x_future = self.samples[index]
-        # Convert to torch
-        x_context = torch.tensor(x_context, dtype=torch.float32)
-        x_future = torch.tensor(x_future, dtype=torch.float32)
+    def _setup_wandb(self) -> None:
+        """Initialize Weights & Biases logging."""
+        wandb.init(project=self.config.wandb_project, entity=self.config.wandb_entity, config=self.config.__dict__)
 
-        input_padding = torch.zeros_like(x_context)
+    def _create_dataloader(self, dataset: Dataset, name: str) -> DataLoader:
+        """Create a dataloader from a dataset."""
+        return DataLoader(
+            dataset,
+            batch_size=self.config.batch_size,
+            shuffle=name == "train",
+            num_workers=mp.cpu_count(),
+            pin_memory=self.device.type == "cuda",
+            persistent_workers=True,
+            prefetch_factor=2,
+        )
 
-        freq = torch.zeros(1, dtype=torch.long)
+    def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
+        """Train for one epoch."""
+        self.model.train()
+        total_loss = 0.0
+        n_batches = len(train_loader)
 
-        return x_context, input_padding, freq, x_future
+        for batch in train_loader:
+            x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
 
-
-def collate_fn(batch):
-    xs_context = [item[0] for item in batch]
-    xs_padding = [item[1] for item in batch]
-    freqs = [item[2] for item in batch]
-    xs_future = [item[3] for item in batch]
-
-    x_context = torch.stack(xs_context, dim=0)
-    input_pad = torch.stack(xs_padding, dim=0)
-    freq = torch.stack(freqs, dim=0)  # shape [B, 1]
-    x_future = torch.stack(xs_future, dim=0)
-
-    return x_context, input_pad, freq, x_future
-
-
-def get_model(*, load_weights: bool = False):
-    # standard model hack
-    repo_id = "google/timesfm-2.0-500m-pytorch"
-    tfm = timesfm.TimesFm(
-        hparams=timesfm.TimesFmHparams(
-            backend="cuda",
-            per_core_batch_size=32,
-            horizon_len=128,
-            num_layers=50,
-            use_positional_embedding=False,
-            context_len=192,
-        ),
-        checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=repo_id),
-    )
-
-    model = PatchedTimeSeriesDecoder(tfm._model_config)
-
-    if load_weights:
-        checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
-        print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
-        loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
-        model.load_state_dict(loaded_checkpoint)
-        print("After loading:")
-        print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
-    model = model.to(device)
-
-    # import sys
-    # sys.exit(-1)
-    # repo_id = "google/timesfm-1.0-200m"
-    return model, tfm._model_config
-
-
-def train_model(
-    ticker="AAPL", start="2015-01-01", end="2022-01-01", train_split=0.8, batch_size=8, num_epochs=20, pretrained=False
-):
-    df_close = download_yfinance_data(ticker, start=start, end=end)
-    model, config = get_model(load_weights=pretrained)
-
-    total_len = len(df_close)
-    train_size = int(total_len * train_split)
-    val_size = total_len - train_size
-
-    train_series = df_close.iloc[:train_size].reset_index(drop=True)
-    val_series = df_close.iloc[train_size:].reset_index(drop=True)
-
-    train_dataset = FinancialDataset(
-        series=train_series, config=config, context_length=128, horizon_length=config.horizon_len
-    )
-    val_dataset = FinancialDataset(
-        series=val_series, config=config, context_length=128, horizon_length=config.horizon_len
-    )
-    print("Train samples:", len(train_dataset))
-    print("Val samples:", len(val_dataset))
-    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
-    val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
-
-    optimizer = optim.Adam(model.parameters(), lr=1e-4)
-
-    for epoch in range(num_epochs):
-        model.train()
-        total_train_loss = 0.0
-
-        for x_context, x_padding, freq, x_future in train_dataloader:
-            x_context, x_padding, freq, x_future = (
-                x_context.to(device),
-                x_padding.to(device),
-                freq.to(device),
-                x_future.to(device),
-            )
-            predictions = model(x_context, x_padding.float(), freq)
-            # predictions shape => [B, N, horizon_len, (1 + #quantiles)]
-            predictions_mean = predictions[..., 0]  # => [B, N, horizon_len]
-            last_patch_pred = predictions_mean[:, -1, :]  # => [B, horizon_len]
-
-            # x_future => [B, horizon_len]
-            loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
+            predictions = self.model(x_context, x_padding.float(), freq)
+            predictions_mean = predictions[..., 0]
+            last_patch_pred = predictions_mean[:, -1, :]
+            loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
 
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
 
-            total_train_loss += loss.item()
+            total_loss += loss.item()
 
-        avg_train_loss = total_train_loss / len(train_dataloader)
+        return total_loss / n_batches
 
-        # -------- Compute validation loss --------
-        model.eval()
-        total_val_loss = 0.0
-        with torch.no_grad():
-            for x_context, x_padding, freq, x_future in val_dataloader:
-                x_context, x_padding, freq, x_future = (
-                    x_context.to(device),
-                    x_padding.to(device),
-                    freq.to(device),
-                    x_future.to(device),
-                )
-                predictions = model(x_context, x_padding.float(), freq)
-                predictions_mean = predictions[..., 0]
-                last_patch_pred = predictions_mean[:, -1, :]
-                val_loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
-                total_val_loss += val_loss.item()
+    @torch.no_grad()
+    def _validate(self, val_loader: DataLoader) -> float:
+        """Perform validation."""
+        self.model.eval()
+        total_loss = 0.0
 
-        avg_val_loss = total_val_loss / max(len(val_dataloader), 1)
+        for batch in val_loader:
+            x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch]
 
-        print(f"[Epoch {epoch+1}] Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
+            predictions = self.model(x_context, x_padding.float(), freq)
+            predictions_mean = predictions[..., 0]
+            last_patch_pred = predictions_mean[:, -1, :]
 
-    torch.save(model.state_dict(), "timesfm_finetuned.pth")
-    return model, train_dataloader, val_dataloader
+            loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+            total_loss += loss.item()
 
+        return total_loss / len(val_loader)
 
-def plot_predictions(model, dataloader):
-    model.eval()
-    with torch.no_grad():
-        x_context, x_padding, freq, x_future = next(iter(dataloader))
-        x_context, x_padding, freq, x_future = (
-            x_context.to(device),
-            x_padding.to(device),
-            freq.to(device),
-            x_future.to(device),
-        )
-        # Forward pass
-        predictions = model(x_context, x_padding.float(), freq)
-        # => [B, N, horizon_len, (1 + #quantiles)]
-        predictions_mean = predictions[..., 0]  # => [B, N, horizon_len]
-        last_patch_prediction = predictions_mean[:, -1, :]  # => [B, horizon_len]
+    def finetune(self, train_dataset: Dataset, val_dataset: Dataset) -> Dict[str, Any]:
+        """
+        Finetune the TimesFM model on the provided datasets.
 
-        # We'll plot only the first sample in the batch
-        i = 0
-        pred_vals = last_patch_prediction[i].cpu().numpy()  # [horizon_len]
-        context_vals = x_context[i].cpu().numpy()  # [context_len]
-        future_vals = x_future[i].cpu().numpy()  # [horizon_len]
+        Args:
+            train_dataset: Training dataset
+            val_dataset: Validation dataset
 
-        horizon_len = future_vals.shape[0]
-        context_len = context_vals.shape[0]
+        Returns:
+            Dict containing training history and best model path
+        """
+        self.model = self.model.to(self.device)
 
-        plt.figure(figsize=(10, 5))
+        train_loader = self._create_dataloader(train_dataset, "train")
+        val_loader = self._create_dataloader(val_dataset, "val")
 
-        # Plot context
-        plt.plot(range(context_len), context_vals, label="Context (History)", color="blue")
-
-        # Plot predicted future
-        plt.plot(
-            range(context_len, context_len + horizon_len),
-            pred_vals,
-            label="Predicted Future",
-            color="orange",
+        optimizer = optim.Adam(
+            self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
         )
 
-        # Plot ground truth future
-        plt.plot(
-            range(context_len, context_len + horizon_len),
-            future_vals,
-            label="Ground Truth Future",
-            color="green",
-            linestyle="--",
-        )
+        history = {"train_loss": [], "val_loss": [], "learning_rate": []}
 
-        plt.xlabel("Time")
-        plt.ylabel("Value")
-        plt.title("Model Forecast vs. Ground Truth")
-        plt.legend()
-        plt.show()
-        plt.savefig("pic_predictions.png")
+        self.logger.info(f"Starting training for {self.config.num_epochs} epochs...")
+        self.logger.info(f"Training samples: {len(train_dataset)}")
+        self.logger.info(f"Validation samples: {len(val_dataset)}")
 
+        try:
+            for epoch in range(self.config.num_epochs):
+                train_loss = self._train_epoch(train_loader, optimizer)
 
-if __name__ == "__main__":
-    # Example usage
-    model, train_dl, val_dl = train_model(
-        ticker="AAPL",
-        start="2012-01-01",
-        end="2019-01-01",
-        train_split=0.8,
-        batch_size=256,
-        num_epochs=50,
-        pretrained=True,
-    )
+                val_loss = self._validate(val_loader)
 
-    plot_predictions(model, val_dl)
+                current_lr = optimizer.param_groups[0]["lr"]
+
+                history["train_loss"].append(train_loss)
+                history["val_loss"].append(val_loss)
+                history["learning_rate"].append(current_lr)
+
+                metrics = {
+                    "train_loss": train_loss,
+                    "val_loss": val_loss,
+                    "learning_rate": current_lr,
+                    "epoch": epoch + 1,
+                }
+
+                if self.config.use_wandb:
+                    wandb.log(metrics)
+
+                print(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
+        except KeyboardInterrupt:
+            self.logger.info("Training interrupted by user")
+
+        return {"history": history}

From fb6213ed59313d5f73b95916aa9e92dadb87ccef Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Tue, 21 Jan 2025 01:50:18 +0000
Subject: [PATCH 03/12] Fix Wandb errro

---
 notebooks/finetuning_example.py | 2 +-
 notebooks/finetuning_torch.py   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
index 5664456..966c980 100644
--- a/notebooks/finetuning_example.py
+++ b/notebooks/finetuning_example.py
@@ -202,7 +202,7 @@ def get_data(context_len: int, horizon_len: int) -> Tuple[Dataset, Dataset]:
 def basic_example():
     """Basic example of finetuning TimesFM on stock data."""
     model, hparams, tfm_config = get_model(load_weights=True)
-    config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=False)
+    config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=True)
 
     train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
     finetuner = TimesFMFinetuner(model, config)
diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py
index bf908e7..e923136 100644
--- a/notebooks/finetuning_torch.py
+++ b/notebooks/finetuning_torch.py
@@ -90,7 +90,7 @@ class TimesFMFinetuner:
 
     def _setup_wandb(self) -> None:
         """Initialize Weights & Biases logging."""
-        wandb.init(project=self.config.wandb_project, entity=self.config.wandb_entity, config=self.config.__dict__)
+        wandb.init(project=self.config.wandb_project, config=self.config.__dict__)
 
     def _create_dataloader(self, dataset: Dataset, name: str) -> DataLoader:
         """Create a dataloader from a dataset."""

From a559718c664ed67b68e58d77f1cd050643c77ef9 Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Tue, 21 Jan 2025 02:19:35 +0000
Subject: [PATCH 04/12] GPU support added, almost done

---
 notebooks/finetuning_example.py | 95 ++++++++++++++++++++++++++++-----
 notebooks/finetuning_torch.py   | 69 ++++++++++++++++--------
 2 files changed, 130 insertions(+), 34 deletions(-)

diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
index 966c980..c42cf88 100644
--- a/notebooks/finetuning_example.py
+++ b/notebooks/finetuning_example.py
@@ -2,23 +2,21 @@
 Example usage of the TimesFM Finetuning Framework.
 """
 
-import yfinance as yf
-import torch
 from os import path
-import numpy as np
-from torch.utils.data import Dataset
-from timesfm import TimesFm, TimesFmHparams, TimesFmCheckpoint
-from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder
-from finetuning_torch import FinetuningConfig, TimesFMFinetuner
-from huggingface_hub import snapshot_download
+from typing import Optional, Tuple
+
 import numpy as np
 import pandas as pd
-from torch.utils.data import Dataset
 import torch
+import torch.multiprocessing as mp
 import yfinance as yf
-from typing import Tuple, Optional
+from finetuning_torch import FinetuningConfig, TimesFMFinetuner
+from huggingface_hub import snapshot_download
+from torch.utils.data import Dataset
 
-from timesfm import TimesFm, TimesFmHparams
+from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams
+from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder
+import os
 
 
 class TimeSeriesDataset(Dataset):
@@ -199,7 +197,7 @@ def get_data(context_len: int, horizon_len: int) -> Tuple[Dataset, Dataset]:
     return train_dataset, val_dataset
 
 
-def basic_example():
+def single_gpu_example():
     """Basic example of finetuning TimesFM on stock data."""
     model, hparams, tfm_config = get_model(load_weights=True)
     config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=True)
@@ -220,5 +218,76 @@ def basic_example():
     )
 
 
+def setup_process(rank, world_size, model, config, train_dataset, val_dataset, return_dict):
+    """Initialize the distributed process."""
+    # Set up the process group
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"
+
+    # Initialize the process group
+    torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=rank)
+
+    # Set the device for this process
+    torch.cuda.set_device(rank)
+
+    try:
+        finetuner = TimesFMFinetuner(model, config, rank=rank)
+        results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
+
+        if rank == 0:  # Only store results and plot from the main process
+            return_dict["results"] = results
+            plot_predictions(
+                model=model,
+                val_dataset=val_dataset,
+                save_path="timesfm_predictions.png",
+            )
+    finally:
+        # Cleanup - important!
+        torch.distributed.destroy_process_group()
+
+
+def multi_gpu_example():
+    """Example of finetuning TimesFM using multiple GPUs."""
+    # Define which GPUs to use
+    gpu_ids = [0]  # Just using one GPU
+    world_size = len(gpu_ids)
+
+    # Initialize model and config
+    model, hparams, tfm_config = get_model(load_weights=True)
+    config = FinetuningConfig(
+        batch_size=256,
+        num_epochs=5,
+        learning_rate=1e-4,
+        use_wandb=False,
+        distributed=True,
+        gpu_ids=gpu_ids,
+    )
+
+    # Get datasets
+    train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
+
+    # Create a multiprocessing manager to share results between processes
+    manager = mp.Manager()
+    return_dict = manager.dict()
+
+    # Launch processes
+    mp.spawn(
+        setup_process,
+        args=(world_size, model, config, train_dataset, val_dataset, return_dict),
+        nprocs=world_size,
+        join=True,
+    )
+
+    # Get results from the main process
+    results = return_dict.get("results", None)
+    print("\nFinetuning completed!")
+    if results:
+        print(f"Training history: {len(results['history']['train_loss'])} epochs")
+
+    return results
+
+
 if __name__ == "__main__":
-    basic_example()
+    # Use either single GPU or multi-GPU example
+    # basic_example()  # Single GPU
+    multi_gpu_example()  # Multi-GPU
diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py
index e923136..a10d99e 100644
--- a/notebooks/finetuning_torch.py
+++ b/notebooks/finetuning_torch.py
@@ -26,17 +26,19 @@ Example usage:
 
 import abc
 import logging
-from dataclasses import dataclass
+import multiprocessing as mp
+import os
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional, Dict, Any
+from typing import Any, Dict, List, Optional
 
 import torch
-from torch.utils.data import Dataset, DataLoader
+import torch.distributed as dist
 import torch.optim as optim
 from torch.nn.parallel import DistributedDataParallel as DDP
-import wandb
-import multiprocessing as mp
+from torch.utils.data import DataLoader, Dataset
 
+import wandb
 from timesfm import TimesFm
 
 
@@ -59,49 +61,65 @@ class FinetuningConfig:
     use_wandb: bool = False
     wandb_project: str = "timesfm-finetuning"
 
+    gpu_ids: List[int] = field(default_factory=lambda: [0])  # List of GPU IDs to use
+    distributed: bool = False
+    master_port: str = "12355"
+    master_addr: str = "localhost"
+
 
 class TimesFMFinetuner:
-    """Main class for finetuning TimesFM models."""
-
     def __init__(
         self,
         model: TimesFm,
         config: FinetuningConfig,
+        rank: int = 0,
         loss_fn: Optional[callable] = None,
         logger: Optional[logging.Logger] = None,
     ):
-        """
-        Initialize TimesFM finetuner.
-
-        Args:
-            model: TimesFM model to finetune
-            config: Finetuning configuration
-            logger: Optional logger instance
-        """
         self.model = model
         self.config = config
         self.logger = logger or logging.getLogger(__name__)
+        self.rank = rank
 
-        self.device = torch.device(config.device)
-        self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))  # MSELoss()
+        if config.distributed:
+            self._setup_distributed(rank)
 
-        if config.use_wandb:
+        self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
+        self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))
+
+        if config.use_wandb and rank == 0:  # Only initialize wandb on main process
             self._setup_wandb()
 
+    def _setup_distributed(self, rank):
+        """Setup distributed training environment."""
+        os.environ["MASTER_ADDR"] = self.config.master_addr
+        os.environ["MASTER_PORT"] = self.config.master_port
+
+        if not dist.is_initialized():
+            dist.init_process_group(backend="nccl", world_size=len(self.config.gpu_ids), rank=rank)
+
     def _setup_wandb(self) -> None:
         """Initialize Weights & Biases logging."""
         wandb.init(project=self.config.wandb_project, config=self.config.__dict__)
 
     def _create_dataloader(self, dataset: Dataset, name: str) -> DataLoader:
         """Create a dataloader from a dataset."""
+        if self.config.distributed:
+            sampler = torch.utils.data.distributed.DistributedSampler(
+                dataset, num_replicas=len(self.config.gpu_ids), rank=dist.get_rank(), shuffle=name == "train"
+            )
+        else:
+            sampler = None
+
         return DataLoader(
             dataset,
             batch_size=self.config.batch_size,
-            shuffle=name == "train",
-            num_workers=mp.cpu_count(),
+            shuffle=(name == "train" and not self.config.distributed),
+            num_workers=mp.cpu_count() // len(self.config.gpu_ids),
             pin_memory=self.device.type == "cuda",
             persistent_workers=True,
             prefetch_factor=2,
+            sampler=sampler,
         )
 
     def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
@@ -157,13 +175,19 @@ class TimesFMFinetuner:
         """
         self.model = self.model.to(self.device)
 
+        if self.config.distributed:
+            self.model = DDP(
+                self.model,
+                device_ids=[self.config.gpu_ids[dist.get_rank()]],
+                output_device=self.config.gpu_ids[dist.get_rank()],
+            )
+
         train_loader = self._create_dataloader(train_dataset, "train")
         val_loader = self._create_dataloader(val_dataset, "val")
 
         optimizer = optim.Adam(
             self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
         )
-
         history = {"train_loss": [], "val_loss": [], "learning_rate": []}
 
         self.logger.info(f"Starting training for {self.config.num_epochs} epochs...")
@@ -196,4 +220,7 @@ class TimesFMFinetuner:
         except KeyboardInterrupt:
             self.logger.info("Training interrupted by user")
 
+        if self.config.distributed:
+            dist.destroy_process_group()
+
         return {"history": history}

From 955d6cda2297c656b96e280fc5a935ea30d7f61b Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Tue, 21 Jan 2025 15:52:08 +0000
Subject: [PATCH 05/12] Gpu support finished

---
 notebooks/finetuning_example.py |  63 ++++++++++---------
 notebooks/finetuning_torch.py   | 107 +++++++++++++++++++++++---------
 2 files changed, 108 insertions(+), 62 deletions(-)

diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
index c42cf88..f1f8a74 100644
--- a/notebooks/finetuning_example.py
+++ b/notebooks/finetuning_example.py
@@ -107,7 +107,6 @@ def get_model(load_weights: bool = False):
         checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
         loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
         model.load_state_dict(loaded_checkpoint)
-    model = model.to(device)
     return model, hparams, tfm._model_config
 
 
@@ -219,54 +218,55 @@ def single_gpu_example():
 
 
 def setup_process(rank, world_size, model, config, train_dataset, val_dataset, return_dict):
-    """Initialize the distributed process."""
-    # Set up the process group
-    os.environ["MASTER_ADDR"] = "localhost"
-    os.environ["MASTER_PORT"] = "12355"
-
-    # Initialize the process group
-    torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=rank)
-
-    # Set the device for this process
-    torch.cuda.set_device(rank)
-
+    """Setup process function with optimized CUDA handling."""
     try:
+        if torch.cuda.is_available():
+            torch.cuda.set_device(rank)
+
+        os.environ["MASTER_ADDR"] = config.master_addr
+        os.environ["MASTER_PORT"] = config.master_port
+        if not torch.distributed.is_initialized():
+            torch.distributed.init_process_group(backend="nccl", world_size=world_size, rank=rank)
+
         finetuner = TimesFMFinetuner(model, config, rank=rank)
+
         results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
 
-        if rank == 0:  # Only store results and plot from the main process
+        if rank == 0:
             return_dict["results"] = results
             plot_predictions(
                 model=model,
                 val_dataset=val_dataset,
                 save_path="timesfm_predictions.png",
             )
+
+    except Exception as e:
+        print(f"Error in process {rank}: {str(e)}")
+        raise e
     finally:
-        # Cleanup - important!
-        torch.distributed.destroy_process_group()
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
 
 
 def multi_gpu_example():
-    """Example of finetuning TimesFM using multiple GPUs."""
-    # Define which GPUs to use
-    gpu_ids = [0]  # Just using one GPU
+    """Example of finetuning TimesFM using multiple GPUs with optimized spawn."""
+    mp.set_start_method("spawn", force=True)
+
+    gpu_ids = [0, 1]
     world_size = len(gpu_ids)
 
-    # Initialize model and config
     model, hparams, tfm_config = get_model(load_weights=True)
+
+    # Create config
     config = FinetuningConfig(
         batch_size=256,
         num_epochs=5,
         learning_rate=1e-4,
-        use_wandb=False,
+        use_wandb=True,
         distributed=True,
         gpu_ids=gpu_ids,
     )
-
-    # Get datasets
     train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
-
-    # Create a multiprocessing manager to share results between processes
     manager = mp.Manager()
     return_dict = manager.dict()
 
@@ -278,16 +278,17 @@ def multi_gpu_example():
         join=True,
     )
 
-    # Get results from the main process
     results = return_dict.get("results", None)
     print("\nFinetuning completed!")
-    if results:
-        print(f"Training history: {len(results['history']['train_loss'])} epochs")
-
     return results
 
 
 if __name__ == "__main__":
-    # Use either single GPU or multi-GPU example
-    # basic_example()  # Single GPU
-    multi_gpu_example()  # Multi-GPU
+    try:
+        # single_gpu_example()  # Single GPU
+        multi_gpu_example()  # Multi-GPU
+    except Exception as e:
+        print(f"Training failed: {str(e)}")
+    finally:
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py
index a10d99e..29215ef 100644
--- a/notebooks/finetuning_torch.py
+++ b/notebooks/finetuning_torch.py
@@ -39,7 +39,6 @@ from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader, Dataset
 
 import wandb
-from timesfm import TimesFm
 
 
 @dataclass
@@ -63,14 +62,14 @@ class FinetuningConfig:
 
     gpu_ids: List[int] = field(default_factory=lambda: [0])  # List of GPU IDs to use
     distributed: bool = False
-    master_port: str = "12355"
+    master_port: str = "12358"
     master_addr: str = "localhost"
 
 
 class TimesFMFinetuner:
     def __init__(
         self,
-        model: TimesFm,
+        model,
         config: FinetuningConfig,
         rank: int = 0,
         loss_fn: Optional[callable] = None,
@@ -87,7 +86,7 @@ class TimesFMFinetuner:
         self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
         self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))
 
-        if config.use_wandb and rank == 0:  # Only initialize wandb on main process
+        if config.use_wandb and rank == 0:
             self._setup_wandb()
 
     def _setup_distributed(self, rank):
@@ -100,7 +99,11 @@ class TimesFMFinetuner:
 
     def _setup_wandb(self) -> None:
         """Initialize Weights & Biases logging."""
-        wandb.init(project=self.config.wandb_project, config=self.config.__dict__)
+
+    def _setup_wandb(self) -> None:
+        """Initialize Weights & Biases logging only on the main process."""
+        if self.rank == 0:  # Only initialize on main process
+            wandb.init(project=self.config.wandb_project, config=self.config.__dict__)
 
     def _create_dataloader(self, dataset: Dataset, name: str) -> DataLoader:
         """Create a dataloader from a dataset."""
@@ -115,15 +118,11 @@ class TimesFMFinetuner:
             dataset,
             batch_size=self.config.batch_size,
             shuffle=(name == "train" and not self.config.distributed),
-            num_workers=mp.cpu_count() // len(self.config.gpu_ids),
-            pin_memory=self.device.type == "cuda",
-            persistent_workers=True,
-            prefetch_factor=2,
             sampler=sampler,
         )
 
     def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
-        """Train for one epoch."""
+        """Train for one epoch with loss debugging."""
         self.model.train()
         total_loss = 0.0
         n_batches = len(train_loader)
@@ -136,6 +135,10 @@ class TimesFMFinetuner:
             last_patch_pred = predictions_mean[:, -1, :]
             loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
 
+            if self.config.distributed:
+                losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
+                dist.all_gather(losses, loss)
+
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
@@ -144,21 +147,26 @@ class TimesFMFinetuner:
 
         return total_loss / n_batches
 
-    @torch.no_grad()
     def _validate(self, val_loader: DataLoader) -> float:
-        """Perform validation."""
+        """Perform validation with loss debugging."""
         self.model.eval()
         total_loss = 0.0
 
-        for batch in val_loader:
-            x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch]
+        with torch.no_grad():
+            for batch in val_loader:
+                x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch]
 
-            predictions = self.model(x_context, x_padding.float(), freq)
-            predictions_mean = predictions[..., 0]
-            last_patch_pred = predictions_mean[:, -1, :]
+                predictions = self.model(x_context, x_padding.float(), freq)
+                predictions_mean = predictions[..., 0]
+                last_patch_pred = predictions_mean[:, -1, :]
 
-            loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
-            total_loss += loss.item()
+                loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+
+                if self.config.distributed:
+                    losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
+                    dist.all_gather(losses, loss)
+
+                total_loss += loss.item()
 
         return total_loss / len(val_loader)
 
@@ -202,21 +210,58 @@ class TimesFMFinetuner:
 
                 current_lr = optimizer.param_groups[0]["lr"]
 
-                history["train_loss"].append(train_loss)
-                history["val_loss"].append(val_loss)
-                history["learning_rate"].append(current_lr)
+                if self.config.distributed:
+                    train_tensor = torch.tensor(train_loss, device=self.device)
+                    val_tensor = torch.tensor(val_loss, device=self.device)
 
-                metrics = {
-                    "train_loss": train_loss,
-                    "val_loss": val_loss,
-                    "learning_rate": current_lr,
-                    "epoch": epoch + 1,
-                }
+                    world_size = dist.get_world_size()
+                    train_losses = [torch.zeros_like(train_tensor, device=self.device) for _ in range(world_size)]
+                    val_losses = [torch.zeros_like(val_tensor, device=self.device) for _ in range(world_size)]
 
-                if self.config.use_wandb:
-                    wandb.log(metrics)
+                    dist.all_gather(train_losses, train_tensor)
+                    dist.all_gather(val_losses, val_tensor)
 
-                print(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
+                    if self.rank == 0 and self.config.use_wandb:
+                        train_losses = [t.cpu().item() for t in train_losses]
+                        val_losses = [t.cpu().item() for t in val_losses]
+
+                        for gpu_idx, (t_loss, v_loss) in enumerate(zip(train_losses, val_losses)):
+                            wandb.log(
+                                {
+                                    f"train_loss_gpu_{gpu_idx}": t_loss,
+                                    f"val_loss_gpu_{gpu_idx}": v_loss,
+                                },
+                                commit=False,
+                            )
+
+                        wandb.log(
+                            {
+                                "train_loss": train_loss,
+                                "val_loss": val_loss,
+                                "learning_rate": current_lr,
+                                "epoch": epoch + 1,
+                            }
+                        )
+                        history["train_loss"].append(train_loss)
+                        history["val_loss"].append(val_loss)
+                        history["learning_rate"].append(current_lr)
+
+                else:
+                    if self.config.use_wandb:
+                        wandb.log(
+                            {
+                                "train_loss": train_loss,
+                                "val_loss": val_loss,
+                                "learning_rate": current_lr,
+                                "epoch": epoch + 1,
+                            }
+                        )
+                    history["train_loss"].append(train_loss)
+                    history["val_loss"].append(val_loss)
+                    history["learning_rate"].append(current_lr)
+
+                if self.rank == 0:
+                    print(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
         except KeyboardInterrupt:
             self.logger.info("Training interrupted by user")
 

From 4f4649c212c7c2ba244996842f39cdf5fafeab9a Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Tue, 21 Jan 2025 16:09:28 +0000
Subject: [PATCH 06/12] Small refactoring

---
 notebooks/finetuning_example.py |   2 +-
 notebooks/finetuning_torch.py   | 365 +++++++++++++++++++-------------
 2 files changed, 218 insertions(+), 149 deletions(-)

diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
index f1f8a74..07325b9 100644
--- a/notebooks/finetuning_example.py
+++ b/notebooks/finetuning_example.py
@@ -261,7 +261,7 @@ def multi_gpu_example():
     config = FinetuningConfig(
         batch_size=256,
         num_epochs=5,
-        learning_rate=1e-4,
+        learning_rate=3e-5,
         use_wandb=True,
         distributed=True,
         gpu_ids=gpu_ids,
diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py
index 29215ef..a3c7f3e 100644
--- a/notebooks/finetuning_torch.py
+++ b/notebooks/finetuning_torch.py
@@ -1,115 +1,204 @@
 """
 TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
-
-Example usage:
-    ```python
-    # Prepare datasets
-    train_dataset = TimeSeriesDataset(train_data, context_length=128, horizon_length=32)
-    val_dataset = TimeSeriesDataset(val_data, context_length=128, horizon_length=32)
-    
-    # Initialize model and configuration
-    model = TimesFm(...)
-    config = FinetuningConfig(
-        batch_size=64,
-        num_epochs=50,
-        learning_rate=1e-4,
-        use_wandb=True
-    )
-    
-    # Create finetuner
-    finetuner = TimesFMFinetuner(model, config)
-    
-    # Finetune model
-    results = finetuner.finetune(train_dataset, val_dataset)
-    ```
 """
 
-import abc
 import logging
-import multiprocessing as mp
 import os
+from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import torch
 import torch.distributed as dist
-import torch.optim as optim
+import torch.nn as nn
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader, Dataset
 
 import wandb
 
 
+class MetricsLogger(ABC):
+    """Abstract base class for logging metrics during training.
+
+    This class defines the interface for logging metrics during model training.
+    Concrete implementations can log to different backends (e.g., WandB, TensorBoard).
+    """
+
+    @abstractmethod
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        """Log metrics to the specified backend.
+
+        Args:
+            metrics: Dictionary containing metric names and values.
+            step: Optional step number or epoch for the metrics.
+        """
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        """Clean up any resources used by the logger."""
+        pass
+
+
+class WandBLogger(MetricsLogger):
+    """Weights & Biases implementation of metrics logging.
+
+    Args:
+        project: Name of the W&B project.
+        config: Configuration dictionary to log.
+        rank: Process rank in distributed training.
+    """
+
+    def __init__(self, project: str, config: Dict[str, Any], rank: int = 0):
+        self.rank = rank
+        if rank == 0:
+            wandb.init(project=project, config=config)
+
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        """Log metrics to W&B if on the main process.
+
+        Args:
+            metrics: Dictionary of metrics to log.
+            step: Current training step or epoch.
+        """
+        if self.rank == 0:
+            wandb.log(metrics, step=step)
+
+    def close(self) -> None:
+        """Finish the W&B run if on the main process."""
+        if self.rank == 0:
+            wandb.finish()
+
+
+class DistributedManager:
+    """Manages distributed training setup and cleanup.
+
+    Args:
+        world_size: Total number of processes.
+        rank: Process rank.
+        master_addr: Address of the master process.
+        master_port: Port for distributed communication.
+        backend: PyTorch distributed backend to use.
+    """
+
+    def __init__(
+        self,
+        world_size: int,
+        rank: int,
+        master_addr: str = "localhost",
+        master_port: str = "12358",
+        backend: str = "nccl",
+    ):
+        self.world_size = world_size
+        self.rank = rank
+        self.master_addr = master_addr
+        self.master_port = master_port
+        self.backend = backend
+
+    def setup(self) -> None:
+        """Initialize the distributed environment."""
+        os.environ["MASTER_ADDR"] = self.master_addr
+        os.environ["MASTER_PORT"] = self.master_port
+
+        if not dist.is_initialized():
+            dist.init_process_group(backend=self.backend, world_size=self.world_size, rank=self.rank)
+
+    def cleanup(self) -> None:
+        """Clean up the distributed environment."""
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
 @dataclass
 class FinetuningConfig:
-    """Configuration for TimesFM finetuning process."""
+    """Configuration for model training.
+
+    Args:
+        batch_size: Number of samples per batch.
+        num_epochs: Number of training epochs.
+        learning_rate: Initial learning rate.
+        weight_decay: L2 regularization factor.
+        device: Device to train on ('cuda' or 'cpu').
+        distributed: Whether to use distributed training.
+        gpu_ids: List of GPU IDs to use.
+        master_port: Port for distributed training.
+        master_addr: Address for distributed training.
+        use_wandb: Whether to use Weights & Biases logging.
+        wandb_project: W&B project name.
+    """
 
-    # Training parameters
     batch_size: int = 32
     num_epochs: int = 20
     learning_rate: float = 1e-4
     weight_decay: float = 0.01
-
-    # Hardware parameters
     device: str = "cuda" if torch.cuda.is_available() else "cpu"
     distributed: bool = False
-    world_size: int = 1
-
-    # Logging parameters
+    gpu_ids: List[int] = field(default_factory=lambda: [0])
+    master_port: str = "12358"
+    master_addr: str = "localhost"
     use_wandb: bool = False
     wandb_project: str = "timesfm-finetuning"
 
-    gpu_ids: List[int] = field(default_factory=lambda: [0])  # List of GPU IDs to use
-    distributed: bool = False
-    master_port: str = "12358"
-    master_addr: str = "localhost"
-
 
 class TimesFMFinetuner:
+    """Handles model training and validation.
+
+    Args:
+        model: PyTorch model to train.
+        config: Training configuration.
+        rank: Process rank for distributed training.
+        loss_fn: Loss function (defaults to MSE).
+        logger: Optional logging.Logger instance.
+    """
+
     def __init__(
         self,
-        model,
+        model: nn.Module,
         config: FinetuningConfig,
         rank: int = 0,
-        loss_fn: Optional[callable] = None,
+        loss_fn: Optional[Callable] = None,
         logger: Optional[logging.Logger] = None,
     ):
         self.model = model
         self.config = config
-        self.logger = logger or logging.getLogger(__name__)
         self.rank = rank
-
-        if config.distributed:
-            self._setup_distributed(rank)
-
+        self.logger = logger or logging.getLogger(__name__)
         self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
         self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))
 
-        if config.use_wandb and rank == 0:
-            self._setup_wandb()
+        if config.use_wandb:
+            self.metrics_logger = WandBLogger(config.wandb_project, config.__dict__, rank)
 
-    def _setup_distributed(self, rank):
-        """Setup distributed training environment."""
-        os.environ["MASTER_ADDR"] = self.config.master_addr
-        os.environ["MASTER_PORT"] = self.config.master_port
+        if config.distributed:
+            self.dist_manager = DistributedManager(
+                world_size=len(config.gpu_ids),
+                rank=rank,
+                master_addr=config.master_addr,
+                master_port=config.master_port,
+            )
+            self.dist_manager.setup()
+            self.model = self._setup_distributed_model()
 
-        if not dist.is_initialized():
-            dist.init_process_group(backend="nccl", world_size=len(self.config.gpu_ids), rank=rank)
+    def _setup_distributed_model(self) -> nn.Module:
+        """Configure model for distributed training."""
+        self.model = self.model.to(self.device)
+        return DDP(
+            self.model, device_ids=[self.config.gpu_ids[self.rank]], output_device=self.config.gpu_ids[self.rank]
+        )
 
-    def _setup_wandb(self) -> None:
-        """Initialize Weights & Biases logging."""
+    def _create_dataloader(self, dataset: Dataset, is_train: bool) -> DataLoader:
+        """Create appropriate DataLoader based on training configuration.
 
-    def _setup_wandb(self) -> None:
-        """Initialize Weights & Biases logging only on the main process."""
-        if self.rank == 0:  # Only initialize on main process
-            wandb.init(project=self.config.wandb_project, config=self.config.__dict__)
+        Args:
+            dataset: Dataset to create loader for.
+            is_train: Whether this is for training (affects shuffling).
 
-    def _create_dataloader(self, dataset: Dataset, name: str) -> DataLoader:
-        """Create a dataloader from a dataset."""
+        Returns:
+            DataLoader instance.
+        """
         if self.config.distributed:
             sampler = torch.utils.data.distributed.DistributedSampler(
-                dataset, num_replicas=len(self.config.gpu_ids), rank=dist.get_rank(), shuffle=name == "train"
+                dataset, num_replicas=len(self.config.gpu_ids), rank=dist.get_rank(), shuffle=is_train
             )
         else:
             sampler = None
@@ -117,23 +206,44 @@ class TimesFMFinetuner:
         return DataLoader(
             dataset,
             batch_size=self.config.batch_size,
-            shuffle=(name == "train" and not self.config.distributed),
+            shuffle=(is_train and not self.config.distributed),
             sampler=sampler,
         )
 
+    def _process_batch(self, batch: List[torch.Tensor]) -> tuple:
+        """Process a single batch of data.
+
+        Args:
+            batch: List of input tensors.
+
+        Returns:
+            Tuple of (loss, predictions).
+        """
+        x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
+
+        predictions = self.model(x_context, x_padding.float(), freq)
+        predictions_mean = predictions[..., 0]
+        last_patch_pred = predictions_mean[:, -1, :]
+
+        loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+
+        return loss, predictions
+
     def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
-        """Train for one epoch with loss debugging."""
+        """Train for one epoch.
+
+        Args:
+            train_loader: DataLoader for training data.
+            optimizer: Optimizer instance.
+
+        Returns:
+            Average training loss for the epoch.
+        """
         self.model.train()
         total_loss = 0.0
-        n_batches = len(train_loader)
 
         for batch in train_loader:
-            x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
-
-            predictions = self.model(x_context, x_padding.float(), freq)
-            predictions_mean = predictions[..., 0]
-            last_patch_pred = predictions_mean[:, -1, :]
-            loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+            loss, _ = self._process_batch(batch)
 
             if self.config.distributed:
                 losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
@@ -145,22 +255,23 @@ class TimesFMFinetuner:
 
             total_loss += loss.item()
 
-        return total_loss / n_batches
+        return total_loss / len(train_loader)
 
     def _validate(self, val_loader: DataLoader) -> float:
-        """Perform validation with loss debugging."""
+        """Perform validation.
+
+        Args:
+            val_loader: DataLoader for validation data.
+
+        Returns:
+            Average validation loss.
+        """
         self.model.eval()
         total_loss = 0.0
 
         with torch.no_grad():
             for batch in val_loader:
-                x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch]
-
-                predictions = self.model(x_context, x_padding.float(), freq)
-                predictions_mean = predictions[..., 0]
-                last_patch_pred = predictions_mean[:, -1, :]
-
-                loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+                loss, _ = self._process_batch(batch)
 
                 if self.config.distributed:
                     losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
@@ -171,31 +282,23 @@ class TimesFMFinetuner:
         return total_loss / len(val_loader)
 
     def finetune(self, train_dataset: Dataset, val_dataset: Dataset) -> Dict[str, Any]:
-        """
-        Finetune the TimesFM model on the provided datasets.
+        """Train the model.
 
         Args:
-            train_dataset: Training dataset
-            val_dataset: Validation dataset
+            train_dataset: Training dataset.
+            val_dataset: Validation dataset.
 
         Returns:
-            Dict containing training history and best model path
+            Dictionary containing training history.
         """
         self.model = self.model.to(self.device)
+        train_loader = self._create_dataloader(train_dataset, is_train=True)
+        val_loader = self._create_dataloader(val_dataset, is_train=False)
 
-        if self.config.distributed:
-            self.model = DDP(
-                self.model,
-                device_ids=[self.config.gpu_ids[dist.get_rank()]],
-                output_device=self.config.gpu_ids[dist.get_rank()],
-            )
-
-        train_loader = self._create_dataloader(train_dataset, "train")
-        val_loader = self._create_dataloader(val_dataset, "val")
-
-        optimizer = optim.Adam(
+        optimizer = torch.optim.Adam(
             self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
         )
+
         history = {"train_loss": [], "val_loss": [], "learning_rate": []}
 
         self.logger.info(f"Starting training for {self.config.num_epochs} epochs...")
@@ -205,67 +308,33 @@ class TimesFMFinetuner:
         try:
             for epoch in range(self.config.num_epochs):
                 train_loss = self._train_epoch(train_loader, optimizer)
-
                 val_loss = self._validate(val_loader)
-
                 current_lr = optimizer.param_groups[0]["lr"]
 
-                if self.config.distributed:
-                    train_tensor = torch.tensor(train_loss, device=self.device)
-                    val_tensor = torch.tensor(val_loss, device=self.device)
+                metrics = {
+                    "train_loss": train_loss,
+                    "val_loss": val_loss,
+                    "learning_rate": current_lr,
+                    "epoch": epoch + 1,
+                }
 
-                    world_size = dist.get_world_size()
-                    train_losses = [torch.zeros_like(train_tensor, device=self.device) for _ in range(world_size)]
-                    val_losses = [torch.zeros_like(val_tensor, device=self.device) for _ in range(world_size)]
+                if self.config.use_wandb:
+                    self.metrics_logger.log_metrics(metrics)
 
-                    dist.all_gather(train_losses, train_tensor)
-                    dist.all_gather(val_losses, val_tensor)
-
-                    if self.rank == 0 and self.config.use_wandb:
-                        train_losses = [t.cpu().item() for t in train_losses]
-                        val_losses = [t.cpu().item() for t in val_losses]
-
-                        for gpu_idx, (t_loss, v_loss) in enumerate(zip(train_losses, val_losses)):
-                            wandb.log(
-                                {
-                                    f"train_loss_gpu_{gpu_idx}": t_loss,
-                                    f"val_loss_gpu_{gpu_idx}": v_loss,
-                                },
-                                commit=False,
-                            )
-
-                        wandb.log(
-                            {
-                                "train_loss": train_loss,
-                                "val_loss": val_loss,
-                                "learning_rate": current_lr,
-                                "epoch": epoch + 1,
-                            }
-                        )
-                        history["train_loss"].append(train_loss)
-                        history["val_loss"].append(val_loss)
-                        history["learning_rate"].append(current_lr)
-
-                else:
-                    if self.config.use_wandb:
-                        wandb.log(
-                            {
-                                "train_loss": train_loss,
-                                "val_loss": val_loss,
-                                "learning_rate": current_lr,
-                                "epoch": epoch + 1,
-                            }
-                        )
-                    history["train_loss"].append(train_loss)
-                    history["val_loss"].append(val_loss)
-                    history["learning_rate"].append(current_lr)
+                history["train_loss"].append(train_loss)
+                history["val_loss"].append(val_loss)
+                history["learning_rate"].append(current_lr)
 
                 if self.rank == 0:
-                    print(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
+                    self.logger.info(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
+
         except KeyboardInterrupt:
             self.logger.info("Training interrupted by user")
 
         if self.config.distributed:
-            dist.destroy_process_group()
+            self.dist_manager.cleanup()
+
+        if self.config.use_wandb:
+            self.metrics_logger.close()
 
         return {"history": history}

From f84366e3d10a0a5c046549a98fd8728aace547de Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Tue, 21 Jan 2025 19:29:42 +0000
Subject: [PATCH 07/12] Style fix

---
 notebooks/finetuning_example.py | 402 ++++++++++++------------
 notebooks/finetuning_torch.py   | 532 ++++++++++++++++----------------
 2 files changed, 467 insertions(+), 467 deletions(-)

diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
index 07325b9..7ffb53d 100644
--- a/notebooks/finetuning_example.py
+++ b/notebooks/finetuning_example.py
@@ -20,275 +20,275 @@ import os
 
 
 class TimeSeriesDataset(Dataset):
-    """Dataset for time series data compatible with TimesFM."""
+  """Dataset for time series data compatible with TimesFM."""
 
-    def __init__(self, series: np.ndarray, context_length: int, horizon_length: int):
-        """
-        Initialize dataset.
+  def __init__(self, series: np.ndarray, context_length: int, horizon_length: int):
+    """
+    Initialize dataset.
 
-        Args:
-            series: Time series data
-            context_length: Number of past timesteps to use as input
-            horizon_length: Number of future timesteps to predict
-        """
-        self.series = series
-        self.context_length = context_length
-        self.horizon_length = horizon_length
-        self._prepare_samples()
+    Args:
+      series: Time series data
+      context_length: Number of past timesteps to use as input
+      horizon_length: Number of future timesteps to predict
+    """
+    self.series = series
+    self.context_length = context_length
+    self.horizon_length = horizon_length
+    self._prepare_samples()
 
-    def _prepare_samples(self) -> None:
-        """Prepare sliding window samples from the time series."""
-        self.samples = []
-        total_length = self.context_length + self.horizon_length
+  def _prepare_samples(self) -> None:
+    """Prepare sliding window samples from the time series."""
+    self.samples = []
+    total_length = self.context_length + self.horizon_length
 
-        for start_idx in range(0, len(self.series) - total_length + 1):
-            end_idx = start_idx + self.context_length
-            x_context = self.series[start_idx:end_idx]
-            x_future = self.series[end_idx : end_idx + self.horizon_length]
-            self.samples.append((x_context, x_future))
+    for start_idx in range(0, len(self.series) - total_length + 1):
+      end_idx = start_idx + self.context_length
+      x_context = self.series[start_idx:end_idx]
+      x_future = self.series[end_idx : end_idx + self.horizon_length]
+      self.samples.append((x_context, x_future))
 
-    def __len__(self) -> int:
-        return len(self.samples)
+  def __len__(self) -> int:
+    return len(self.samples)
 
-    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        x_context, x_future = self.samples[index]
+  def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    x_context, x_future = self.samples[index]
 
-        x_context = torch.tensor(x_context, dtype=torch.float32)
-        x_future = torch.tensor(x_future, dtype=torch.float32)
+    x_context = torch.tensor(x_context, dtype=torch.float32)
+    x_future = torch.tensor(x_future, dtype=torch.float32)
 
-        input_padding = torch.zeros_like(x_context)
-        freq = torch.zeros(1, dtype=torch.long)
+    input_padding = torch.zeros_like(x_context)
+    freq = torch.zeros(1, dtype=torch.long)
 
-        return x_context, input_padding, freq, x_future
+    return x_context, input_padding, freq, x_future
 
 
 def prepare_datasets(
-    series: np.ndarray, context_length: int, horizon_length: int, train_split: float = 0.8
+  series: np.ndarray, context_length: int, horizon_length: int, train_split: float = 0.8
 ) -> Tuple[Dataset, Dataset]:
-    """
-    Prepare training and validation datasets from time series data.
+  """
+  Prepare training and validation datasets from time series data.
 
-    Args:
-        series: Input time series data
-        context_length: Number of past timesteps to use
-        horizon_length: Number of future timesteps to predict
-        train_split: Fraction of data to use for training
+  Args:
+    series: Input time series data
+    context_length: Number of past timesteps to use
+    horizon_length: Number of future timesteps to predict
+    train_split: Fraction of data to use for training
 
-    Returns:
-        Tuple of (train_dataset, val_dataset)
-    """
-    train_size = int(len(series) * train_split)
-    train_data = series[:train_size]
-    val_data = series[train_size:]
+  Returns:
+    Tuple of (train_dataset, val_dataset)
+  """
+  train_size = int(len(series) * train_split)
+  train_data = series[:train_size]
+  val_data = series[train_size:]
 
-    # Create datasets
-    train_dataset = TimeSeriesDataset(train_data, context_length=context_length, horizon_length=horizon_length)
+  # Create datasets
+  train_dataset = TimeSeriesDataset(train_data, context_length=context_length, horizon_length=horizon_length)
 
-    val_dataset = TimeSeriesDataset(val_data, context_length=context_length, horizon_length=horizon_length)
+  val_dataset = TimeSeriesDataset(val_data, context_length=context_length, horizon_length=horizon_length)
 
-    return train_dataset, val_dataset
+  return train_dataset, val_dataset
 
 
 def get_model(load_weights: bool = False):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    repo_id = "google/timesfm-2.0-500m-pytorch"
-    hparams = TimesFmHparams(
-        backend=device,
-        per_core_batch_size=32,
-        horizon_len=128,
-        num_layers=50,
-        use_positional_embedding=False,
-        context_len=192,
-    )
-    tfm = TimesFm(hparams=hparams, checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
+  device = "cuda" if torch.cuda.is_available() else "cpu"
+  repo_id = "google/timesfm-2.0-500m-pytorch"
+  hparams = TimesFmHparams(
+    backend=device,
+    per_core_batch_size=32,
+    horizon_len=128,
+    num_layers=50,
+    use_positional_embedding=False,
+    context_len=192,
+  )
+  tfm = TimesFm(hparams=hparams, checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
 
-    model = PatchedTimeSeriesDecoder(tfm._model_config)
-    if load_weights:
-        checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
-        loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
-        model.load_state_dict(loaded_checkpoint)
-    return model, hparams, tfm._model_config
+  model = PatchedTimeSeriesDecoder(tfm._model_config)
+  if load_weights:
+    checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
+    loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
+    model.load_state_dict(loaded_checkpoint)
+  return model, hparams, tfm._model_config
 
 
 def plot_predictions(
-    model: TimesFm,
-    val_dataset: Dataset,
-    save_path: Optional[str] = "predictions.png",
+  model: TimesFm,
+  val_dataset: Dataset,
+  save_path: Optional[str] = "predictions.png",
 ) -> None:
-    """
-    Plot model predictions against ground truth for a batch of validation data.
+  """
+  Plot model predictions against ground truth for a batch of validation data.
 
-    Args:
-        model: Trained TimesFM model
-        val_dataset: Validation dataset
-        save_path: Path to save the plot
-    """
-    import matplotlib.pyplot as plt
+  Args:
+    model: Trained TimesFM model
+    val_dataset: Validation dataset
+    save_path: Path to save the plot
+  """
+  import matplotlib.pyplot as plt
 
-    model.eval()
+  model.eval()
 
-    x_context, x_padding, freq, x_future = val_dataset[0]
-    x_context = x_context.unsqueeze(0)  # Add batch dimension
-    x_padding = x_padding.unsqueeze(0)
-    freq = freq.unsqueeze(0)
-    x_future = x_future.unsqueeze(0)
+  x_context, x_padding, freq, x_future = val_dataset[0]
+  x_context = x_context.unsqueeze(0)  # Add batch dimension
+  x_padding = x_padding.unsqueeze(0)
+  freq = freq.unsqueeze(0)
+  x_future = x_future.unsqueeze(0)
 
-    device = next(model.parameters()).device
-    x_context = x_context.to(device)
-    x_padding = x_padding.to(device)
-    freq = freq.to(device)
-    x_future = x_future.to(device)
+  device = next(model.parameters()).device
+  x_context = x_context.to(device)
+  x_padding = x_padding.to(device)
+  freq = freq.to(device)
+  x_future = x_future.to(device)
 
-    with torch.no_grad():
-        predictions = model(x_context, x_padding.float(), freq)
-        predictions_mean = predictions[..., 0]  # [B, N, horizon_len]
-        last_patch_pred = predictions_mean[:, -1, :]  # [B, horizon_len]
+  with torch.no_grad():
+    predictions = model(x_context, x_padding.float(), freq)
+    predictions_mean = predictions[..., 0]  # [B, N, horizon_len]
+    last_patch_pred = predictions_mean[:, -1, :]  # [B, horizon_len]
 
-    context_vals = x_context[0].cpu().numpy()
-    future_vals = x_future[0].cpu().numpy()
-    pred_vals = last_patch_pred[0].cpu().numpy()
+  context_vals = x_context[0].cpu().numpy()
+  future_vals = x_future[0].cpu().numpy()
+  pred_vals = last_patch_pred[0].cpu().numpy()
 
-    context_len = len(context_vals)
-    horizon_len = len(future_vals)
+  context_len = len(context_vals)
+  horizon_len = len(future_vals)
 
-    plt.figure(figsize=(12, 6))
+  plt.figure(figsize=(12, 6))
 
-    plt.plot(range(context_len), context_vals, label="Historical Data", color="blue", linewidth=2)
+  plt.plot(range(context_len), context_vals, label="Historical Data", color="blue", linewidth=2)
 
-    plt.plot(
-        range(context_len, context_len + horizon_len),
-        future_vals,
-        label="Ground Truth",
-        color="green",
-        linestyle="--",
-        linewidth=2,
-    )
+  plt.plot(
+    range(context_len, context_len + horizon_len),
+    future_vals,
+    label="Ground Truth",
+    color="green",
+    linestyle="--",
+    linewidth=2,
+  )
 
-    plt.plot(range(context_len, context_len + horizon_len), pred_vals, label="Prediction", color="red", linewidth=2)
+  plt.plot(range(context_len, context_len + horizon_len), pred_vals, label="Prediction", color="red", linewidth=2)
 
-    plt.xlabel("Time Step")
-    plt.ylabel("Value")
-    plt.title("TimesFM Predictions vs Ground Truth")
-    plt.legend()
-    plt.grid(True)
+  plt.xlabel("Time Step")
+  plt.ylabel("Value")
+  plt.title("TimesFM Predictions vs Ground Truth")
+  plt.legend()
+  plt.grid(True)
 
-    if save_path:
-        plt.savefig(save_path)
-        print(f"Plot saved to {save_path}")
+  if save_path:
+    plt.savefig(save_path)
+    print(f"Plot saved to {save_path}")
 
-    plt.close()
+  plt.close()
 
 
 def get_data(context_len: int, horizon_len: int) -> Tuple[Dataset, Dataset]:
-    df = yf.download("AAPL", start="2010-01-01", end="2019-01-01")
-    time_series = df["Close"].values
+  df = yf.download("AAPL", start="2010-01-01", end="2019-01-01")
+  time_series = df["Close"].values
 
-    train_dataset, val_dataset = prepare_datasets(
-        series=time_series,
-        context_length=context_len,
-        horizon_length=horizon_len,
-        train_split=0.8,
-    )
+  train_dataset, val_dataset = prepare_datasets(
+    series=time_series,
+    context_length=context_len,
+    horizon_length=horizon_len,
+    train_split=0.8,
+  )
 
-    print(f"Created datasets:")
-    print(f"- Training samples: {len(train_dataset)}")
-    print(f"- Validation samples: {len(val_dataset)}")
-    return train_dataset, val_dataset
+  print(f"Created datasets:")
+  print(f"- Training samples: {len(train_dataset)}")
+  print(f"- Validation samples: {len(val_dataset)}")
+  return train_dataset, val_dataset
 
 
 def single_gpu_example():
-    """Basic example of finetuning TimesFM on stock data."""
-    model, hparams, tfm_config = get_model(load_weights=True)
-    config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=True)
+  """Basic example of finetuning TimesFM on stock data."""
+  model, hparams, tfm_config = get_model(load_weights=True)
+  config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=True)
 
-    train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
-    finetuner = TimesFMFinetuner(model, config)
+  train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
+  finetuner = TimesFMFinetuner(model, config)
 
-    print("\nStarting finetuning...")
-    results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
+  print("\nStarting finetuning...")
+  results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
 
-    print("\nFinetuning completed!")
-    print(f"Training history: {len(results['history']['train_loss'])} epochs")
+  print("\nFinetuning completed!")
+  print(f"Training history: {len(results['history']['train_loss'])} epochs")
 
-    plot_predictions(
-        model=model,
-        val_dataset=val_dataset,
-        save_path="timesfm_predictions.png",
-    )
+  plot_predictions(
+    model=model,
+    val_dataset=val_dataset,
+    save_path="timesfm_predictions.png",
+  )
 
 
 def setup_process(rank, world_size, model, config, train_dataset, val_dataset, return_dict):
-    """Setup process function with optimized CUDA handling."""
-    try:
-        if torch.cuda.is_available():
-            torch.cuda.set_device(rank)
+  """Setup process function with optimized CUDA handling."""
+  try:
+    if torch.cuda.is_available():
+      torch.cuda.set_device(rank)
 
-        os.environ["MASTER_ADDR"] = config.master_addr
-        os.environ["MASTER_PORT"] = config.master_port
-        if not torch.distributed.is_initialized():
-            torch.distributed.init_process_group(backend="nccl", world_size=world_size, rank=rank)
+    os.environ["MASTER_ADDR"] = config.master_addr
+    os.environ["MASTER_PORT"] = config.master_port
+    if not torch.distributed.is_initialized():
+      torch.distributed.init_process_group(backend="nccl", world_size=world_size, rank=rank)
 
-        finetuner = TimesFMFinetuner(model, config, rank=rank)
+    finetuner = TimesFMFinetuner(model, config, rank=rank)
 
-        results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
+    results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
 
-        if rank == 0:
-            return_dict["results"] = results
-            plot_predictions(
-                model=model,
-                val_dataset=val_dataset,
-                save_path="timesfm_predictions.png",
-            )
+    if rank == 0:
+      return_dict["results"] = results
+      plot_predictions(
+        model=model,
+        val_dataset=val_dataset,
+        save_path="timesfm_predictions.png",
+      )
 
-    except Exception as e:
-        print(f"Error in process {rank}: {str(e)}")
-        raise e
-    finally:
-        if torch.distributed.is_initialized():
-            torch.distributed.destroy_process_group()
+  except Exception as e:
+    print(f"Error in process {rank}: {str(e)}")
+    raise e
+  finally:
+    if torch.distributed.is_initialized():
+      torch.distributed.destroy_process_group()
 
 
 def multi_gpu_example():
-    """Example of finetuning TimesFM using multiple GPUs with optimized spawn."""
-    mp.set_start_method("spawn", force=True)
+  """Example of finetuning TimesFM using multiple GPUs with optimized spawn."""
+  mp.set_start_method("spawn", force=True)
 
-    gpu_ids = [0, 1]
-    world_size = len(gpu_ids)
+  gpu_ids = [0, 1]
+  world_size = len(gpu_ids)
 
-    model, hparams, tfm_config = get_model(load_weights=True)
+  model, hparams, tfm_config = get_model(load_weights=True)
 
-    # Create config
-    config = FinetuningConfig(
-        batch_size=256,
-        num_epochs=5,
-        learning_rate=3e-5,
-        use_wandb=True,
-        distributed=True,
-        gpu_ids=gpu_ids,
-    )
-    train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
-    manager = mp.Manager()
-    return_dict = manager.dict()
+  # Create config
+  config = FinetuningConfig(
+    batch_size=256,
+    num_epochs=5,
+    learning_rate=3e-5,
+    use_wandb=True,
+    distributed=True,
+    gpu_ids=gpu_ids,
+  )
+  train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
+  manager = mp.Manager()
+  return_dict = manager.dict()
 
-    # Launch processes
-    mp.spawn(
-        setup_process,
-        args=(world_size, model, config, train_dataset, val_dataset, return_dict),
-        nprocs=world_size,
-        join=True,
-    )
+  # Launch processes
+  mp.spawn(
+    setup_process,
+    args=(world_size, model, config, train_dataset, val_dataset, return_dict),
+    nprocs=world_size,
+    join=True,
+  )
 
-    results = return_dict.get("results", None)
-    print("\nFinetuning completed!")
-    return results
+  results = return_dict.get("results", None)
+  print("\nFinetuning completed!")
+  return results
 
 
 if __name__ == "__main__":
-    try:
-        # single_gpu_example()  # Single GPU
-        multi_gpu_example()  # Multi-GPU
-    except Exception as e:
-        print(f"Training failed: {str(e)}")
-    finally:
-        if torch.distributed.is_initialized():
-            torch.distributed.destroy_process_group()
+  try:
+    # single_gpu_example()  # Single GPU
+    multi_gpu_example()  # Multi-GPU
+  except Exception as e:
+    print(f"Training failed: {str(e)}")
+  finally:
+    if torch.distributed.is_initialized():
+      torch.distributed.destroy_process_group()
\ No newline at end of file
diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py
index a3c7f3e..af2eadb 100644
--- a/notebooks/finetuning_torch.py
+++ b/notebooks/finetuning_torch.py
@@ -18,323 +18,323 @@ import wandb
 
 
 class MetricsLogger(ABC):
-    """Abstract base class for logging metrics during training.
+  """Abstract base class for logging metrics during training.
 
-    This class defines the interface for logging metrics during model training.
-    Concrete implementations can log to different backends (e.g., WandB, TensorBoard).
+  This class defines the interface for logging metrics during model training.
+  Concrete implementations can log to different backends (e.g., WandB, TensorBoard).
+  """
+
+  @abstractmethod
+  def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+    """Log metrics to the specified backend.
+
+    Args:
+      metrics: Dictionary containing metric names and values.
+      step: Optional step number or epoch for the metrics.
     """
+    pass
 
-    @abstractmethod
-    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
-        """Log metrics to the specified backend.
-
-        Args:
-            metrics: Dictionary containing metric names and values.
-            step: Optional step number or epoch for the metrics.
-        """
-        pass
-
-    @abstractmethod
-    def close(self) -> None:
-        """Clean up any resources used by the logger."""
-        pass
+  @abstractmethod
+  def close(self) -> None:
+    """Clean up any resources used by the logger."""
+    pass
 
 
 class WandBLogger(MetricsLogger):
-    """Weights & Biases implementation of metrics logging.
+  """Weights & Biases implementation of metrics logging.
+
+  Args:
+    project: Name of the W&B project.
+    config: Configuration dictionary to log.
+    rank: Process rank in distributed training.
+  """
+
+  def __init__(self, project: str, config: Dict[str, Any], rank: int = 0):
+    self.rank = rank
+    if rank == 0:
+      wandb.init(project=project, config=config)
+
+  def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+    """Log metrics to W&B if on the main process.
 
     Args:
-        project: Name of the W&B project.
-        config: Configuration dictionary to log.
-        rank: Process rank in distributed training.
+      metrics: Dictionary of metrics to log.
+      step: Current training step or epoch.
     """
+    if self.rank == 0:
+      wandb.log(metrics, step=step)
 
-    def __init__(self, project: str, config: Dict[str, Any], rank: int = 0):
-        self.rank = rank
-        if rank == 0:
-            wandb.init(project=project, config=config)
-
-    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
-        """Log metrics to W&B if on the main process.
-
-        Args:
-            metrics: Dictionary of metrics to log.
-            step: Current training step or epoch.
-        """
-        if self.rank == 0:
-            wandb.log(metrics, step=step)
-
-    def close(self) -> None:
-        """Finish the W&B run if on the main process."""
-        if self.rank == 0:
-            wandb.finish()
+  def close(self) -> None:
+    """Finish the W&B run if on the main process."""
+    if self.rank == 0:
+      wandb.finish()
 
 
 class DistributedManager:
-    """Manages distributed training setup and cleanup.
+  """Manages distributed training setup and cleanup.
 
-    Args:
-        world_size: Total number of processes.
-        rank: Process rank.
-        master_addr: Address of the master process.
-        master_port: Port for distributed communication.
-        backend: PyTorch distributed backend to use.
-    """
+  Args:
+    world_size: Total number of processes.
+    rank: Process rank.
+    master_addr: Address of the master process.
+    master_port: Port for distributed communication.
+    backend: PyTorch distributed backend to use.
+  """
 
-    def __init__(
-        self,
-        world_size: int,
-        rank: int,
-        master_addr: str = "localhost",
-        master_port: str = "12358",
-        backend: str = "nccl",
-    ):
-        self.world_size = world_size
-        self.rank = rank
-        self.master_addr = master_addr
-        self.master_port = master_port
-        self.backend = backend
+  def __init__(
+    self,
+    world_size: int,
+    rank: int,
+    master_addr: str = "localhost",
+    master_port: str = "12358",
+    backend: str = "nccl",
+  ):
+    self.world_size = world_size
+    self.rank = rank
+    self.master_addr = master_addr
+    self.master_port = master_port
+    self.backend = backend
 
-    def setup(self) -> None:
-        """Initialize the distributed environment."""
-        os.environ["MASTER_ADDR"] = self.master_addr
-        os.environ["MASTER_PORT"] = self.master_port
+  def setup(self) -> None:
+    """Initialize the distributed environment."""
+    os.environ["MASTER_ADDR"] = self.master_addr
+    os.environ["MASTER_PORT"] = self.master_port
 
-        if not dist.is_initialized():
-            dist.init_process_group(backend=self.backend, world_size=self.world_size, rank=self.rank)
+    if not dist.is_initialized():
+      dist.init_process_group(backend=self.backend, world_size=self.world_size, rank=self.rank)
 
-    def cleanup(self) -> None:
-        """Clean up the distributed environment."""
-        if dist.is_initialized():
-            dist.destroy_process_group()
+  def cleanup(self) -> None:
+    """Clean up the distributed environment."""
+    if dist.is_initialized():
+      dist.destroy_process_group()
 
 
 @dataclass
 class FinetuningConfig:
-    """Configuration for model training.
+  """Configuration for model training.
 
-    Args:
-        batch_size: Number of samples per batch.
-        num_epochs: Number of training epochs.
-        learning_rate: Initial learning rate.
-        weight_decay: L2 regularization factor.
-        device: Device to train on ('cuda' or 'cpu').
-        distributed: Whether to use distributed training.
-        gpu_ids: List of GPU IDs to use.
-        master_port: Port for distributed training.
-        master_addr: Address for distributed training.
-        use_wandb: Whether to use Weights & Biases logging.
-        wandb_project: W&B project name.
-    """
+  Args:
+    batch_size: Number of samples per batch.
+    num_epochs: Number of training epochs.
+    learning_rate: Initial learning rate.
+    weight_decay: L2 regularization factor.
+    device: Device to train on ('cuda' or 'cpu').
+    distributed: Whether to use distributed training.
+    gpu_ids: List of GPU IDs to use.
+    master_port: Port for distributed training.
+    master_addr: Address for distributed training.
+    use_wandb: Whether to use Weights & Biases logging.
+    wandb_project: W&B project name.
+  """
 
-    batch_size: int = 32
-    num_epochs: int = 20
-    learning_rate: float = 1e-4
-    weight_decay: float = 0.01
-    device: str = "cuda" if torch.cuda.is_available() else "cpu"
-    distributed: bool = False
-    gpu_ids: List[int] = field(default_factory=lambda: [0])
-    master_port: str = "12358"
-    master_addr: str = "localhost"
-    use_wandb: bool = False
-    wandb_project: str = "timesfm-finetuning"
+  batch_size: int = 32
+  num_epochs: int = 20
+  learning_rate: float = 1e-4
+  weight_decay: float = 0.01
+  device: str = "cuda" if torch.cuda.is_available() else "cpu"
+  distributed: bool = False
+  gpu_ids: List[int] = field(default_factory=lambda: [0])
+  master_port: str = "12358"
+  master_addr: str = "localhost"
+  use_wandb: bool = False
+  wandb_project: str = "timesfm-finetuning"
 
 
 class TimesFMFinetuner:
-    """Handles model training and validation.
+  """Handles model training and validation.
+
+  Args:
+    model: PyTorch model to train.
+    config: Training configuration.
+    rank: Process rank for distributed training.
+    loss_fn: Loss function (defaults to MSE).
+    logger: Optional logging.Logger instance.
+  """
+
+  def __init__(
+    self,
+    model: nn.Module,
+    config: FinetuningConfig,
+    rank: int = 0,
+    loss_fn: Optional[Callable] = None,
+    logger: Optional[logging.Logger] = None,
+  ):
+    self.model = model
+    self.config = config
+    self.rank = rank
+    self.logger = logger or logging.getLogger(__name__)
+    self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
+    self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))
+
+    if config.use_wandb:
+      self.metrics_logger = WandBLogger(config.wandb_project, config.__dict__, rank)
+
+    if config.distributed:
+      self.dist_manager = DistributedManager(
+        world_size=len(config.gpu_ids),
+        rank=rank,
+        master_addr=config.master_addr,
+        master_port=config.master_port,
+      )
+      self.dist_manager.setup()
+      self.model = self._setup_distributed_model()
+
+  def _setup_distributed_model(self) -> nn.Module:
+    """Configure model for distributed training."""
+    self.model = self.model.to(self.device)
+    return DDP(
+      self.model, device_ids=[self.config.gpu_ids[self.rank]], output_device=self.config.gpu_ids[self.rank]
+    )
+
+  def _create_dataloader(self, dataset: Dataset, is_train: bool) -> DataLoader:
+    """Create appropriate DataLoader based on training configuration.
 
     Args:
-        model: PyTorch model to train.
-        config: Training configuration.
-        rank: Process rank for distributed training.
-        loss_fn: Loss function (defaults to MSE).
-        logger: Optional logging.Logger instance.
+      dataset: Dataset to create loader for.
+      is_train: Whether this is for training (affects shuffling).
+
+    Returns:
+      DataLoader instance.
     """
+    if self.config.distributed:
+      sampler = torch.utils.data.distributed.DistributedSampler(
+        dataset, num_replicas=len(self.config.gpu_ids), rank=dist.get_rank(), shuffle=is_train
+      )
+    else:
+      sampler = None
 
-    def __init__(
-        self,
-        model: nn.Module,
-        config: FinetuningConfig,
-        rank: int = 0,
-        loss_fn: Optional[Callable] = None,
-        logger: Optional[logging.Logger] = None,
-    ):
-        self.model = model
-        self.config = config
-        self.rank = rank
-        self.logger = logger or logging.getLogger(__name__)
-        self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
-        self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))
+    return DataLoader(
+      dataset,
+      batch_size=self.config.batch_size,
+      shuffle=(is_train and not self.config.distributed),
+      sampler=sampler,
+    )
 
-        if config.use_wandb:
-            self.metrics_logger = WandBLogger(config.wandb_project, config.__dict__, rank)
+  def _process_batch(self, batch: List[torch.Tensor]) -> tuple:
+    """Process a single batch of data.
 
-        if config.distributed:
-            self.dist_manager = DistributedManager(
-                world_size=len(config.gpu_ids),
-                rank=rank,
-                master_addr=config.master_addr,
-                master_port=config.master_port,
-            )
-            self.dist_manager.setup()
-            self.model = self._setup_distributed_model()
+    Args:
+      batch: List of input tensors.
 
-    def _setup_distributed_model(self) -> nn.Module:
-        """Configure model for distributed training."""
-        self.model = self.model.to(self.device)
-        return DDP(
-            self.model, device_ids=[self.config.gpu_ids[self.rank]], output_device=self.config.gpu_ids[self.rank]
-        )
+    Returns:
+      Tuple of (loss, predictions).
+    """
+    x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
 
-    def _create_dataloader(self, dataset: Dataset, is_train: bool) -> DataLoader:
-        """Create appropriate DataLoader based on training configuration.
+    predictions = self.model(x_context, x_padding.float(), freq)
+    predictions_mean = predictions[..., 0]
+    last_patch_pred = predictions_mean[:, -1, :]
 
-        Args:
-            dataset: Dataset to create loader for.
-            is_train: Whether this is for training (affects shuffling).
+    loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
 
-        Returns:
-            DataLoader instance.
-        """
-        if self.config.distributed:
-            sampler = torch.utils.data.distributed.DistributedSampler(
-                dataset, num_replicas=len(self.config.gpu_ids), rank=dist.get_rank(), shuffle=is_train
-            )
-        else:
-            sampler = None
+    return loss, predictions
 
-        return DataLoader(
-            dataset,
-            batch_size=self.config.batch_size,
-            shuffle=(is_train and not self.config.distributed),
-            sampler=sampler,
-        )
+  def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
+    """Train for one epoch.
 
-    def _process_batch(self, batch: List[torch.Tensor]) -> tuple:
-        """Process a single batch of data.
+    Args:
+      train_loader: DataLoader for training data.
+      optimizer: Optimizer instance.
 
-        Args:
-            batch: List of input tensors.
+    Returns:
+      Average training loss for the epoch.
+    """
+    self.model.train()
+    total_loss = 0.0
 
-        Returns:
-            Tuple of (loss, predictions).
-        """
-        x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
+    for batch in train_loader:
+      loss, _ = self._process_batch(batch)
 
-        predictions = self.model(x_context, x_padding.float(), freq)
-        predictions_mean = predictions[..., 0]
-        last_patch_pred = predictions_mean[:, -1, :]
+      if self.config.distributed:
+        losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
+        dist.all_gather(losses, loss)
 
-        loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+      optimizer.zero_grad()
+      loss.backward()
+      optimizer.step()
 
-        return loss, predictions
+      total_loss += loss.item()
 
-    def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
-        """Train for one epoch.
+    return total_loss / len(train_loader)
 
-        Args:
-            train_loader: DataLoader for training data.
-            optimizer: Optimizer instance.
+  def _validate(self, val_loader: DataLoader) -> float:
+    """Perform validation.
 
-        Returns:
-            Average training loss for the epoch.
-        """
-        self.model.train()
-        total_loss = 0.0
+    Args:
+      val_loader: DataLoader for validation data.
 
-        for batch in train_loader:
-            loss, _ = self._process_batch(batch)
+    Returns:
+      Average validation loss.
+    """
+    self.model.eval()
+    total_loss = 0.0
 
-            if self.config.distributed:
-                losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
-                dist.all_gather(losses, loss)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-
-            total_loss += loss.item()
-
-        return total_loss / len(train_loader)
-
-    def _validate(self, val_loader: DataLoader) -> float:
-        """Perform validation.
-
-        Args:
-            val_loader: DataLoader for validation data.
-
-        Returns:
-            Average validation loss.
-        """
-        self.model.eval()
-        total_loss = 0.0
-
-        with torch.no_grad():
-            for batch in val_loader:
-                loss, _ = self._process_batch(batch)
-
-                if self.config.distributed:
-                    losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
-                    dist.all_gather(losses, loss)
-
-                total_loss += loss.item()
-
-        return total_loss / len(val_loader)
-
-    def finetune(self, train_dataset: Dataset, val_dataset: Dataset) -> Dict[str, Any]:
-        """Train the model.
-
-        Args:
-            train_dataset: Training dataset.
-            val_dataset: Validation dataset.
-
-        Returns:
-            Dictionary containing training history.
-        """
-        self.model = self.model.to(self.device)
-        train_loader = self._create_dataloader(train_dataset, is_train=True)
-        val_loader = self._create_dataloader(val_dataset, is_train=False)
-
-        optimizer = torch.optim.Adam(
-            self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
-        )
-
-        history = {"train_loss": [], "val_loss": [], "learning_rate": []}
-
-        self.logger.info(f"Starting training for {self.config.num_epochs} epochs...")
-        self.logger.info(f"Training samples: {len(train_dataset)}")
-        self.logger.info(f"Validation samples: {len(val_dataset)}")
-
-        try:
-            for epoch in range(self.config.num_epochs):
-                train_loss = self._train_epoch(train_loader, optimizer)
-                val_loss = self._validate(val_loader)
-                current_lr = optimizer.param_groups[0]["lr"]
-
-                metrics = {
-                    "train_loss": train_loss,
-                    "val_loss": val_loss,
-                    "learning_rate": current_lr,
-                    "epoch": epoch + 1,
-                }
-
-                if self.config.use_wandb:
-                    self.metrics_logger.log_metrics(metrics)
-
-                history["train_loss"].append(train_loss)
-                history["val_loss"].append(val_loss)
-                history["learning_rate"].append(current_lr)
-
-                if self.rank == 0:
-                    self.logger.info(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
-
-        except KeyboardInterrupt:
-            self.logger.info("Training interrupted by user")
+    with torch.no_grad():
+      for batch in val_loader:
+        loss, _ = self._process_batch(batch)
 
         if self.config.distributed:
-            self.dist_manager.cleanup()
+          losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
+          dist.all_gather(losses, loss)
+
+        total_loss += loss.item()
+
+    return total_loss / len(val_loader)
+
+  def finetune(self, train_dataset: Dataset, val_dataset: Dataset) -> Dict[str, Any]:
+    """Train the model.
+
+    Args:
+      train_dataset: Training dataset.
+      val_dataset: Validation dataset.
+
+    Returns:
+      Dictionary containing training history.
+    """
+    self.model = self.model.to(self.device)
+    train_loader = self._create_dataloader(train_dataset, is_train=True)
+    val_loader = self._create_dataloader(val_dataset, is_train=False)
+
+    optimizer = torch.optim.Adam(
+      self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
+    )
+
+    history = {"train_loss": [], "val_loss": [], "learning_rate": []}
+
+    self.logger.info(f"Starting training for {self.config.num_epochs} epochs...")
+    self.logger.info(f"Training samples: {len(train_dataset)}")
+    self.logger.info(f"Validation samples: {len(val_dataset)}")
+
+    try:
+      for epoch in range(self.config.num_epochs):
+        train_loss = self._train_epoch(train_loader, optimizer)
+        val_loss = self._validate(val_loader)
+        current_lr = optimizer.param_groups[0]["lr"]
+
+        metrics = {
+          "train_loss": train_loss,
+          "val_loss": val_loss,
+          "learning_rate": current_lr,
+          "epoch": epoch + 1,
+        }
 
         if self.config.use_wandb:
-            self.metrics_logger.close()
+          self.metrics_logger.log_metrics(metrics)
 
-        return {"history": history}
+        history["train_loss"].append(train_loss)
+        history["val_loss"].append(val_loss)
+        history["learning_rate"].append(current_lr)
+
+        if self.rank == 0:
+          self.logger.info(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
+
+    except KeyboardInterrupt:
+      self.logger.info("Training interrupted by user")
+
+    if self.config.distributed:
+      self.dist_manager.cleanup()
+
+    if self.config.use_wandb:
+      self.metrics_logger.close()
+
+    return {"history": history}
\ No newline at end of file

From 86551f761af790d70e9869beb73939a313666c03 Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Tue, 21 Jan 2025 19:57:50 +0000
Subject: [PATCH 08/12] Add jupyter notebook

---
 notebooks/finetuning_example.py               |   2 +-
 notebooks/finetuning_torch.ipynb              | 292 ++++++++++++++++++
 .../timesfm}/finetuning_torch.py              |   0
 3 files changed, 293 insertions(+), 1 deletion(-)
 create mode 100644 notebooks/finetuning_torch.ipynb
 rename {notebooks => src/timesfm}/finetuning_torch.py (100%)

diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
index 7ffb53d..5f166b3 100644
--- a/notebooks/finetuning_example.py
+++ b/notebooks/finetuning_example.py
@@ -10,7 +10,7 @@ import pandas as pd
 import torch
 import torch.multiprocessing as mp
 import yfinance as yf
-from finetuning_torch import FinetuningConfig, TimesFMFinetuner
+from timesfm.finetuning_torch import FinetuningConfig, TimesFMFinetuner
 from huggingface_hub import snapshot_download
 from torch.utils.data import Dataset
 
diff --git a/notebooks/finetuning_torch.ipynb b/notebooks/finetuning_torch.ipynb
new file mode 100644
index 0000000..9515e7e
--- /dev/null
+++ b/notebooks/finetuning_torch.ipynb
@@ -0,0 +1,292 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Introduction\n",
+    "This notebook shows how to use TimesFM with finetuning. \n",
+    "\n",
+    "In order to perform finetuning, you need to create the Pytorch Dataset in a proper format. The example of the Dataset is provided below.\n",
+    "The finetuning code can be found in timesfm.finetuning_torch.py. This notebook just imports the methods from finetuning"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Dataset Creation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from os import path\n",
+    "from typing import Optional, Tuple\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import torch\n",
+    "import torch.multiprocessing as mp\n",
+    "import yfinance as yf\n",
+    "from timesfm.finetuning_torch import FinetuningConfig, TimesFMFinetuner\n",
+    "from huggingface_hub import snapshot_download\n",
+    "from torch.utils.data import Dataset\n",
+    "\n",
+    "from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams\n",
+    "from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "class TimeSeriesDataset(Dataset):\n",
+    "  \"\"\"Dataset for time series data compatible with TimesFM.\"\"\"\n",
+    "\n",
+    "  def __init__(self, series: np.ndarray, context_length: int, horizon_length: int):\n",
+    "    \"\"\"\n",
+    "    Initialize dataset.\n",
+    "\n",
+    "    Args:\n",
+    "      series: Time series data\n",
+    "      context_length: Number of past timesteps to use as input\n",
+    "      horizon_length: Number of future timesteps to predict\n",
+    "    \"\"\"\n",
+    "    self.series = series\n",
+    "    self.context_length = context_length\n",
+    "    self.horizon_length = horizon_length\n",
+    "    self._prepare_samples()\n",
+    "\n",
+    "  def _prepare_samples(self) -> None:\n",
+    "    \"\"\"Prepare sliding window samples from the time series.\"\"\"\n",
+    "    self.samples = []\n",
+    "    total_length = self.context_length + self.horizon_length\n",
+    "\n",
+    "    for start_idx in range(0, len(self.series) - total_length + 1):\n",
+    "      end_idx = start_idx + self.context_length\n",
+    "      x_context = self.series[start_idx:end_idx]\n",
+    "      x_future = self.series[end_idx : end_idx + self.horizon_length]\n",
+    "      self.samples.append((x_context, x_future))\n",
+    "\n",
+    "  def __len__(self) -> int:\n",
+    "    return len(self.samples)\n",
+    "\n",
+    "  def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n",
+    "    x_context, x_future = self.samples[index]\n",
+    "\n",
+    "    x_context = torch.tensor(x_context, dtype=torch.float32)\n",
+    "    x_future = torch.tensor(x_future, dtype=torch.float32)\n",
+    "\n",
+    "    input_padding = torch.zeros_like(x_context)\n",
+    "    freq = torch.zeros(1, dtype=torch.long)\n",
+    "\n",
+    "    return x_context, input_padding, freq, x_future\n",
+    "\n",
+    "\n",
+    "def prepare_datasets(\n",
+    "  series: np.ndarray, context_length: int, horizon_length: int, train_split: float = 0.8\n",
+    ") -> Tuple[Dataset, Dataset]:\n",
+    "  \"\"\"\n",
+    "  Prepare training and validation datasets from time series data.\n",
+    "\n",
+    "  Args:\n",
+    "    series: Input time series data\n",
+    "    context_length: Number of past timesteps to use\n",
+    "    horizon_length: Number of future timesteps to predict\n",
+    "    train_split: Fraction of data to use for training\n",
+    "\n",
+    "  Returns:\n",
+    "    Tuple of (train_dataset, val_dataset)\n",
+    "  \"\"\"\n",
+    "  train_size = int(len(series) * train_split)\n",
+    "  train_data = series[:train_size]\n",
+    "  val_data = series[train_size:]\n",
+    "\n",
+    "  # Create datasets\n",
+    "  train_dataset = TimeSeriesDataset(train_data, context_length=context_length, horizon_length=horizon_length)\n",
+    "\n",
+    "  val_dataset = TimeSeriesDataset(val_data, context_length=context_length, horizon_length=horizon_length)\n",
+    "\n",
+    "  return train_dataset, val_dataset\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Model Creation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_model(load_weights: bool = False):\n",
+    "  device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "  repo_id = \"google/timesfm-2.0-500m-pytorch\"\n",
+    "  hparams = TimesFmHparams(\n",
+    "    backend=device,\n",
+    "    per_core_batch_size=32,\n",
+    "    horizon_len=128,\n",
+    "    num_layers=50,\n",
+    "    use_positional_embedding=False,\n",
+    "    context_len=192,\n",
+    "  )\n",
+    "  tfm = TimesFm(hparams=hparams, checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))\n",
+    "\n",
+    "  model = PatchedTimeSeriesDecoder(tfm._model_config)\n",
+    "  if load_weights:\n",
+    "    checkpoint_path = path.join(snapshot_download(repo_id), \"torch_model.ckpt\")\n",
+    "    loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)\n",
+    "    model.load_state_dict(loaded_checkpoint)\n",
+    "  return model, hparams, tfm._model_config\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def plot_predictions(\n",
+    "  model: TimesFm,\n",
+    "  val_dataset: Dataset,\n",
+    "  save_path: Optional[str] = \"predictions.png\",\n",
+    ") -> None:\n",
+    "  \"\"\"\n",
+    "  Plot model predictions against ground truth for a batch of validation data.\n",
+    "\n",
+    "  Args:\n",
+    "    model: Trained TimesFM model\n",
+    "    val_dataset: Validation dataset\n",
+    "    save_path: Path to save the plot\n",
+    "  \"\"\"\n",
+    "  import matplotlib.pyplot as plt\n",
+    "\n",
+    "  model.eval()\n",
+    "\n",
+    "  x_context, x_padding, freq, x_future = val_dataset[0]\n",
+    "  x_context = x_context.unsqueeze(0)  # Add batch dimension\n",
+    "  x_padding = x_padding.unsqueeze(0)\n",
+    "  freq = freq.unsqueeze(0)\n",
+    "  x_future = x_future.unsqueeze(0)\n",
+    "\n",
+    "  device = next(model.parameters()).device\n",
+    "  x_context = x_context.to(device)\n",
+    "  x_padding = x_padding.to(device)\n",
+    "  freq = freq.to(device)\n",
+    "  x_future = x_future.to(device)\n",
+    "\n",
+    "  with torch.no_grad():\n",
+    "    predictions = model(x_context, x_padding.float(), freq)\n",
+    "    predictions_mean = predictions[..., 0]  # [B, N, horizon_len]\n",
+    "    last_patch_pred = predictions_mean[:, -1, :]  # [B, horizon_len]\n",
+    "\n",
+    "  context_vals = x_context[0].cpu().numpy()\n",
+    "  future_vals = x_future[0].cpu().numpy()\n",
+    "  pred_vals = last_patch_pred[0].cpu().numpy()\n",
+    "\n",
+    "  context_len = len(context_vals)\n",
+    "  horizon_len = len(future_vals)\n",
+    "\n",
+    "  plt.figure(figsize=(12, 6))\n",
+    "\n",
+    "  plt.plot(range(context_len), context_vals, label=\"Historical Data\", color=\"blue\", linewidth=2)\n",
+    "\n",
+    "  plt.plot(\n",
+    "    range(context_len, context_len + horizon_len),\n",
+    "    future_vals,\n",
+    "    label=\"Ground Truth\",\n",
+    "    color=\"green\",\n",
+    "    linestyle=\"--\",\n",
+    "    linewidth=2,\n",
+    "  )\n",
+    "\n",
+    "  plt.plot(range(context_len, context_len + horizon_len), pred_vals, label=\"Prediction\", color=\"red\", linewidth=2)\n",
+    "\n",
+    "  plt.xlabel(\"Time Step\")\n",
+    "  plt.ylabel(\"Value\")\n",
+    "  plt.title(\"TimesFM Predictions vs Ground Truth\")\n",
+    "  plt.legend()\n",
+    "  plt.grid(True)\n",
+    "\n",
+    "  if save_path:\n",
+    "    plt.savefig(save_path)\n",
+    "    print(f\"Plot saved to {save_path}\")\n",
+    "\n",
+    "  plt.close()\n",
+    "\n",
+    "\n",
+    "def get_data(context_len: int, horizon_len: int) -> Tuple[Dataset, Dataset]:\n",
+    "  df = yf.download(\"AAPL\", start=\"2010-01-01\", end=\"2019-01-01\")\n",
+    "  time_series = df[\"Close\"].values\n",
+    "\n",
+    "  train_dataset, val_dataset = prepare_datasets(\n",
+    "    series=time_series,\n",
+    "    context_length=context_len,\n",
+    "    horizon_length=horizon_len,\n",
+    "    train_split=0.8,\n",
+    "  )\n",
+    "\n",
+    "  print(f\"Created datasets:\")\n",
+    "  print(f\"- Training samples: {len(train_dataset)}\")\n",
+    "  print(f\"- Validation samples: {len(val_dataset)}\")\n",
+    "  return train_dataset, val_dataset\n",
+    "\n",
+    "\n",
+    "def single_gpu_example():\n",
+    "  \"\"\"Basic example of finetuning TimesFM on stock data.\"\"\"\n",
+    "  model, hparams, tfm_config = get_model(load_weights=True)\n",
+    "  config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=True)\n",
+    "\n",
+    "  train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)\n",
+    "  finetuner = TimesFMFinetuner(model, config)\n",
+    "\n",
+    "  print(\"\\nStarting finetuning...\")\n",
+    "  results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)\n",
+    "\n",
+    "  print(\"\\nFinetuning completed!\")\n",
+    "  print(f\"Training history: {len(results['history']['train_loss'])} epochs\")\n",
+    "\n",
+    "  plot_predictions(\n",
+    "    model=model,\n",
+    "    val_dataset=val_dataset,\n",
+    "    save_path=\"timesfm_predictions.png\",\n",
+    "  )\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "single_gpu_example()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "timesfm-311",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/finetuning_torch.py b/src/timesfm/finetuning_torch.py
similarity index 100%
rename from notebooks/finetuning_torch.py
rename to src/timesfm/finetuning_torch.py

From bca190dae68acda8fd30e6e9293d70e00fd344e3 Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkin@google.com>
Date: Sat, 1 Feb 2025 02:26:39 +0000
Subject: [PATCH 09/12] PR Feedback

---
 notebooks/finetuning_example.py      | 294 --------------------
 pyproject.toml                       |   1 +
 src/finetuning/__init__.py           |   0
 src/finetuning/finetuning_example.py | 388 ++++++++++++++++++++++++++
 src/finetuning/finetuning_torch.py   | 398 +++++++++++++++++++++++++++
 src/timesfm/finetuning_torch.py      | 340 -----------------------
 6 files changed, 787 insertions(+), 634 deletions(-)
 delete mode 100644 notebooks/finetuning_example.py
 create mode 100644 src/finetuning/__init__.py
 create mode 100644 src/finetuning/finetuning_example.py
 create mode 100644 src/finetuning/finetuning_torch.py
 delete mode 100644 src/timesfm/finetuning_torch.py

diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
deleted file mode 100644
index 5f166b3..0000000
--- a/notebooks/finetuning_example.py
+++ /dev/null
@@ -1,294 +0,0 @@
-"""
-Example usage of the TimesFM Finetuning Framework.
-"""
-
-from os import path
-from typing import Optional, Tuple
-
-import numpy as np
-import pandas as pd
-import torch
-import torch.multiprocessing as mp
-import yfinance as yf
-from timesfm.finetuning_torch import FinetuningConfig, TimesFMFinetuner
-from huggingface_hub import snapshot_download
-from torch.utils.data import Dataset
-
-from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams
-from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder
-import os
-
-
-class TimeSeriesDataset(Dataset):
-  """Dataset for time series data compatible with TimesFM."""
-
-  def __init__(self, series: np.ndarray, context_length: int, horizon_length: int):
-    """
-    Initialize dataset.
-
-    Args:
-      series: Time series data
-      context_length: Number of past timesteps to use as input
-      horizon_length: Number of future timesteps to predict
-    """
-    self.series = series
-    self.context_length = context_length
-    self.horizon_length = horizon_length
-    self._prepare_samples()
-
-  def _prepare_samples(self) -> None:
-    """Prepare sliding window samples from the time series."""
-    self.samples = []
-    total_length = self.context_length + self.horizon_length
-
-    for start_idx in range(0, len(self.series) - total_length + 1):
-      end_idx = start_idx + self.context_length
-      x_context = self.series[start_idx:end_idx]
-      x_future = self.series[end_idx : end_idx + self.horizon_length]
-      self.samples.append((x_context, x_future))
-
-  def __len__(self) -> int:
-    return len(self.samples)
-
-  def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    x_context, x_future = self.samples[index]
-
-    x_context = torch.tensor(x_context, dtype=torch.float32)
-    x_future = torch.tensor(x_future, dtype=torch.float32)
-
-    input_padding = torch.zeros_like(x_context)
-    freq = torch.zeros(1, dtype=torch.long)
-
-    return x_context, input_padding, freq, x_future
-
-
-def prepare_datasets(
-  series: np.ndarray, context_length: int, horizon_length: int, train_split: float = 0.8
-) -> Tuple[Dataset, Dataset]:
-  """
-  Prepare training and validation datasets from time series data.
-
-  Args:
-    series: Input time series data
-    context_length: Number of past timesteps to use
-    horizon_length: Number of future timesteps to predict
-    train_split: Fraction of data to use for training
-
-  Returns:
-    Tuple of (train_dataset, val_dataset)
-  """
-  train_size = int(len(series) * train_split)
-  train_data = series[:train_size]
-  val_data = series[train_size:]
-
-  # Create datasets
-  train_dataset = TimeSeriesDataset(train_data, context_length=context_length, horizon_length=horizon_length)
-
-  val_dataset = TimeSeriesDataset(val_data, context_length=context_length, horizon_length=horizon_length)
-
-  return train_dataset, val_dataset
-
-
-def get_model(load_weights: bool = False):
-  device = "cuda" if torch.cuda.is_available() else "cpu"
-  repo_id = "google/timesfm-2.0-500m-pytorch"
-  hparams = TimesFmHparams(
-    backend=device,
-    per_core_batch_size=32,
-    horizon_len=128,
-    num_layers=50,
-    use_positional_embedding=False,
-    context_len=192,
-  )
-  tfm = TimesFm(hparams=hparams, checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
-
-  model = PatchedTimeSeriesDecoder(tfm._model_config)
-  if load_weights:
-    checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
-    loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
-    model.load_state_dict(loaded_checkpoint)
-  return model, hparams, tfm._model_config
-
-
-def plot_predictions(
-  model: TimesFm,
-  val_dataset: Dataset,
-  save_path: Optional[str] = "predictions.png",
-) -> None:
-  """
-  Plot model predictions against ground truth for a batch of validation data.
-
-  Args:
-    model: Trained TimesFM model
-    val_dataset: Validation dataset
-    save_path: Path to save the plot
-  """
-  import matplotlib.pyplot as plt
-
-  model.eval()
-
-  x_context, x_padding, freq, x_future = val_dataset[0]
-  x_context = x_context.unsqueeze(0)  # Add batch dimension
-  x_padding = x_padding.unsqueeze(0)
-  freq = freq.unsqueeze(0)
-  x_future = x_future.unsqueeze(0)
-
-  device = next(model.parameters()).device
-  x_context = x_context.to(device)
-  x_padding = x_padding.to(device)
-  freq = freq.to(device)
-  x_future = x_future.to(device)
-
-  with torch.no_grad():
-    predictions = model(x_context, x_padding.float(), freq)
-    predictions_mean = predictions[..., 0]  # [B, N, horizon_len]
-    last_patch_pred = predictions_mean[:, -1, :]  # [B, horizon_len]
-
-  context_vals = x_context[0].cpu().numpy()
-  future_vals = x_future[0].cpu().numpy()
-  pred_vals = last_patch_pred[0].cpu().numpy()
-
-  context_len = len(context_vals)
-  horizon_len = len(future_vals)
-
-  plt.figure(figsize=(12, 6))
-
-  plt.plot(range(context_len), context_vals, label="Historical Data", color="blue", linewidth=2)
-
-  plt.plot(
-    range(context_len, context_len + horizon_len),
-    future_vals,
-    label="Ground Truth",
-    color="green",
-    linestyle="--",
-    linewidth=2,
-  )
-
-  plt.plot(range(context_len, context_len + horizon_len), pred_vals, label="Prediction", color="red", linewidth=2)
-
-  plt.xlabel("Time Step")
-  plt.ylabel("Value")
-  plt.title("TimesFM Predictions vs Ground Truth")
-  plt.legend()
-  plt.grid(True)
-
-  if save_path:
-    plt.savefig(save_path)
-    print(f"Plot saved to {save_path}")
-
-  plt.close()
-
-
-def get_data(context_len: int, horizon_len: int) -> Tuple[Dataset, Dataset]:
-  df = yf.download("AAPL", start="2010-01-01", end="2019-01-01")
-  time_series = df["Close"].values
-
-  train_dataset, val_dataset = prepare_datasets(
-    series=time_series,
-    context_length=context_len,
-    horizon_length=horizon_len,
-    train_split=0.8,
-  )
-
-  print(f"Created datasets:")
-  print(f"- Training samples: {len(train_dataset)}")
-  print(f"- Validation samples: {len(val_dataset)}")
-  return train_dataset, val_dataset
-
-
-def single_gpu_example():
-  """Basic example of finetuning TimesFM on stock data."""
-  model, hparams, tfm_config = get_model(load_weights=True)
-  config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=True)
-
-  train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
-  finetuner = TimesFMFinetuner(model, config)
-
-  print("\nStarting finetuning...")
-  results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
-
-  print("\nFinetuning completed!")
-  print(f"Training history: {len(results['history']['train_loss'])} epochs")
-
-  plot_predictions(
-    model=model,
-    val_dataset=val_dataset,
-    save_path="timesfm_predictions.png",
-  )
-
-
-def setup_process(rank, world_size, model, config, train_dataset, val_dataset, return_dict):
-  """Setup process function with optimized CUDA handling."""
-  try:
-    if torch.cuda.is_available():
-      torch.cuda.set_device(rank)
-
-    os.environ["MASTER_ADDR"] = config.master_addr
-    os.environ["MASTER_PORT"] = config.master_port
-    if not torch.distributed.is_initialized():
-      torch.distributed.init_process_group(backend="nccl", world_size=world_size, rank=rank)
-
-    finetuner = TimesFMFinetuner(model, config, rank=rank)
-
-    results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
-
-    if rank == 0:
-      return_dict["results"] = results
-      plot_predictions(
-        model=model,
-        val_dataset=val_dataset,
-        save_path="timesfm_predictions.png",
-      )
-
-  except Exception as e:
-    print(f"Error in process {rank}: {str(e)}")
-    raise e
-  finally:
-    if torch.distributed.is_initialized():
-      torch.distributed.destroy_process_group()
-
-
-def multi_gpu_example():
-  """Example of finetuning TimesFM using multiple GPUs with optimized spawn."""
-  mp.set_start_method("spawn", force=True)
-
-  gpu_ids = [0, 1]
-  world_size = len(gpu_ids)
-
-  model, hparams, tfm_config = get_model(load_weights=True)
-
-  # Create config
-  config = FinetuningConfig(
-    batch_size=256,
-    num_epochs=5,
-    learning_rate=3e-5,
-    use_wandb=True,
-    distributed=True,
-    gpu_ids=gpu_ids,
-  )
-  train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
-  manager = mp.Manager()
-  return_dict = manager.dict()
-
-  # Launch processes
-  mp.spawn(
-    setup_process,
-    args=(world_size, model, config, train_dataset, val_dataset, return_dict),
-    nprocs=world_size,
-    join=True,
-  )
-
-  results = return_dict.get("results", None)
-  print("\nFinetuning completed!")
-  return results
-
-
-if __name__ == "__main__":
-  try:
-    # single_gpu_example()  # Single GPU
-    multi_gpu_example()  # Multi-GPU
-  except Exception as e:
-    print(f"Training failed: {str(e)}")
-  finally:
-    if torch.distributed.is_initialized():
-      torch.distributed.destroy_process_group()
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index b9dd150..81aaa54 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,6 +9,7 @@ authors = [
     "Abhimanyu Das <abhidas@google.com>",
     "Petros Mol <pmol@google.com>",
     "Justin Güse <guese.justin@gmail.com>",
+    "Michael Chertushkin <chertushkinmichael@gmail.com>"
 ]
 readme = "README.md"
 keywords = ["time series", "timesfm", "forecast", "time series model"]
diff --git a/src/finetuning/__init__.py b/src/finetuning/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/finetuning/finetuning_example.py b/src/finetuning/finetuning_example.py
new file mode 100644
index 0000000..2c396d6
--- /dev/null
+++ b/src/finetuning/finetuning_example.py
@@ -0,0 +1,388 @@
+"""
+Example usage of the TimesFM Finetuning Framework.
+
+For single GPU:
+python script.py --training_mode=single
+
+For multiple GPUs:
+python script.py --training_mode=multi --gpu_ids=0,1,2
+"""
+
+import os
+from os import path
+from typing import Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.multiprocessing as mp
+import yfinance as yf
+from absl import app, flags
+from huggingface_hub import snapshot_download
+from torch.utils.data import Dataset
+
+from finetuning.finetuning_torch import FinetuningConfig, TimesFMFinetuner
+from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams
+from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_enum(
+    "training_mode",
+    "single",
+    ["single", "multi"],
+    'Training mode: "single" for single-GPU or "multi" for multi-GPU training.',
+)
+
+flags.DEFINE_list(
+    "gpu_ids", ["0"],
+    "Comma-separated list of GPU IDs to use for multi-GPU training. Example: 0,1,2"
+)
+
+
+class TimeSeriesDataset(Dataset):
+  """Dataset for time series data compatible with TimesFM."""
+
+  def __init__(self,
+               series: np.ndarray,
+               context_length: int,
+               horizon_length: int,
+               freq_type: int = 0):
+    """
+        Initialize dataset.
+
+        Args:
+            series: Time series data
+            context_length: Number of past timesteps to use as input
+            horizon_length: Number of future timesteps to predict
+            freq_type: Frequency type (0, 1, or 2)
+        """
+    if freq_type not in [0, 1, 2]:
+      raise ValueError("freq_type must be 0, 1, or 2")
+
+    self.series = series
+    self.context_length = context_length
+    self.horizon_length = horizon_length
+    self.freq_type = freq_type
+    self._prepare_samples()
+
+  def _prepare_samples(self) -> None:
+    """Prepare sliding window samples from the time series."""
+    self.samples = []
+    total_length = self.context_length + self.horizon_length
+
+    for start_idx in range(0, len(self.series) - total_length + 1):
+      end_idx = start_idx + self.context_length
+      x_context = self.series[start_idx:end_idx]
+      x_future = self.series[end_idx:end_idx + self.horizon_length]
+      self.samples.append((x_context, x_future))
+
+  def __len__(self) -> int:
+    return len(self.samples)
+
+  def __getitem__(
+      self, index: int
+  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    x_context, x_future = self.samples[index]
+
+    x_context = torch.tensor(x_context, dtype=torch.float32)
+    x_future = torch.tensor(x_future, dtype=torch.float32)
+
+    input_padding = torch.zeros_like(x_context)
+    freq = torch.tensor([self.freq_type], dtype=torch.long)
+
+    return x_context, input_padding, freq, x_future
+
+
+def prepare_datasets(series: np.ndarray,
+                     context_length: int,
+                     horizon_length: int,
+                     freq_type: int = 0,
+                     train_split: float = 0.8) -> Tuple[Dataset, Dataset]:
+  """
+    Prepare training and validation datasets from time series data.
+
+    Args:
+        series: Input time series data
+        context_length: Number of past timesteps to use
+        horizon_length: Number of future timesteps to predict
+        freq_type: Frequency type (0, 1, or 2)
+        train_split: Fraction of data to use for training
+
+    Returns:
+        Tuple of (train_dataset, val_dataset)
+    """
+  train_size = int(len(series) * train_split)
+  train_data = series[:train_size]
+  val_data = series[train_size:]
+
+  # Create datasets with specified frequency type
+  train_dataset = TimeSeriesDataset(train_data,
+                                    context_length=context_length,
+                                    horizon_length=horizon_length,
+                                    freq_type=freq_type)
+
+  val_dataset = TimeSeriesDataset(val_data,
+                                  context_length=context_length,
+                                  horizon_length=horizon_length,
+                                  freq_type=freq_type)
+
+  return train_dataset, val_dataset
+
+
+def get_model(load_weights: bool = False):
+  device = "cuda" if torch.cuda.is_available() else "cpu"
+  repo_id = "google/timesfm-2.0-500m-pytorch"
+  hparams = TimesFmHparams(
+      backend=device,
+      per_core_batch_size=32,
+      horizon_len=128,
+      num_layers=50,
+      use_positional_embedding=False,
+      context_len=
+      192,  # Context length can be anything up to 2048 in multiples of 32
+  )
+  tfm = TimesFm(hparams=hparams,
+                checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
+
+  model = PatchedTimeSeriesDecoder(tfm._model_config)
+  if load_weights:
+    checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
+    loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
+    model.load_state_dict(loaded_checkpoint)
+  return model, hparams, tfm._model_config
+
+
+def plot_predictions(
+    model: TimesFm,
+    val_dataset: Dataset,
+    save_path: Optional[str] = "predictions.png",
+) -> None:
+  """
+    Plot model predictions against ground truth for a batch of validation data.
+
+    Args:
+      model: Trained TimesFM model
+      val_dataset: Validation dataset
+      save_path: Path to save the plot
+    """
+  import matplotlib.pyplot as plt
+
+  model.eval()
+
+  x_context, x_padding, freq, x_future = val_dataset[0]
+  x_context = x_context.unsqueeze(0)  # Add batch dimension
+  x_padding = x_padding.unsqueeze(0)
+  freq = freq.unsqueeze(0)
+  x_future = x_future.unsqueeze(0)
+
+  device = next(model.parameters()).device
+  x_context = x_context.to(device)
+  x_padding = x_padding.to(device)
+  freq = freq.to(device)
+  x_future = x_future.to(device)
+
+  with torch.no_grad():
+    predictions = model(x_context, x_padding.float(), freq)
+    predictions_mean = predictions[..., 0]  # [B, N, horizon_len]
+    last_patch_pred = predictions_mean[:, -1, :]  # [B, horizon_len]
+
+  context_vals = x_context[0].cpu().numpy()
+  future_vals = x_future[0].cpu().numpy()
+  pred_vals = last_patch_pred[0].cpu().numpy()
+
+  context_len = len(context_vals)
+  horizon_len = len(future_vals)
+
+  plt.figure(figsize=(12, 6))
+
+  plt.plot(range(context_len),
+           context_vals,
+           label="Historical Data",
+           color="blue",
+           linewidth=2)
+
+  plt.plot(
+      range(context_len, context_len + horizon_len),
+      future_vals,
+      label="Ground Truth",
+      color="green",
+      linestyle="--",
+      linewidth=2,
+  )
+
+  plt.plot(range(context_len, context_len + horizon_len),
+           pred_vals,
+           label="Prediction",
+           color="red",
+           linewidth=2)
+
+  plt.xlabel("Time Step")
+  plt.ylabel("Value")
+  plt.title("TimesFM Predictions vs Ground Truth")
+  plt.legend()
+  plt.grid(True)
+
+  if save_path:
+    plt.savefig(save_path)
+    print(f"Plot saved to {save_path}")
+
+  plt.close()
+
+
+def get_data(context_len: int,
+             horizon_len: int,
+             freq_type: int = 0) -> Tuple[Dataset, Dataset]:
+  df = yf.download("AAPL", start="2010-01-01", end="2019-01-01")
+  time_series = df["Close"].values
+
+  train_dataset, val_dataset = prepare_datasets(
+      series=time_series,
+      context_length=context_len,
+      horizon_length=horizon_len,
+      freq_type=freq_type,
+      train_split=0.8,
+  )
+
+  print(f"Created datasets:")
+  print(f"- Training samples: {len(train_dataset)}")
+  print(f"- Validation samples: {len(val_dataset)}")
+  print(f"- Using frequency type: {freq_type}")
+  return train_dataset, val_dataset
+
+
+def single_gpu_example():
+  """Basic example of finetuning TimesFM on stock data."""
+  model, hparams, tfm_config = get_model(load_weights=True)
+  config = FinetuningConfig(batch_size=256,
+                            num_epochs=5,
+                            learning_rate=1e-4,
+                            use_wandb=True,
+                            freq_type=1,
+                            log_every_n_steps=10,
+                            val_check_interval=0.5,
+                            use_quantile_loss=True)
+
+  train_dataset, val_dataset = get_data(128,
+                                        tfm_config.horizon_len,
+                                        freq_type=config.freq_type)
+  finetuner = TimesFMFinetuner(model, config)
+
+  print("\nStarting finetuning...")
+  results = finetuner.finetune(train_dataset=train_dataset,
+                               val_dataset=val_dataset)
+
+  print("\nFinetuning completed!")
+  print(f"Training history: {len(results['history']['train_loss'])} epochs")
+
+  plot_predictions(
+      model=model,
+      val_dataset=val_dataset,
+      save_path="timesfm_predictions.png",
+  )
+
+
+def setup_process(rank, world_size, model, config, train_dataset, val_dataset,
+                  return_dict):
+  """Setup process function with optimized CUDA handling."""
+  try:
+    if torch.cuda.is_available():
+      torch.cuda.set_device(rank)
+
+    os.environ["MASTER_ADDR"] = config.master_addr
+    os.environ["MASTER_PORT"] = config.master_port
+    if not torch.distributed.is_initialized():
+      torch.distributed.init_process_group(backend="nccl",
+                                           world_size=world_size,
+                                           rank=rank)
+
+    finetuner = TimesFMFinetuner(model, config, rank=rank)
+
+    results = finetuner.finetune(train_dataset=train_dataset,
+                                 val_dataset=val_dataset)
+
+    if rank == 0:
+      return_dict["results"] = results
+      plot_predictions(
+          model=model,
+          val_dataset=val_dataset,
+          save_path="timesfm_predictions.png",
+      )
+
+  except Exception as e:
+    print(f"Error in process {rank}: {str(e)}")
+    raise e
+  finally:
+    if torch.distributed.is_initialized():
+      torch.distributed.destroy_process_group()
+
+
+def multi_gpu_example():
+  """Example of finetuning TimesFM using multiple GPUs with optimized spawn."""
+  mp.set_start_method("spawn", force=True)
+
+  gpu_ids = [0, 1]
+  world_size = len(gpu_ids)
+
+  model, hparams, tfm_config = get_model(load_weights=True)
+
+  # Create config
+  config = FinetuningConfig(
+      batch_size=256,
+      num_epochs=5,
+      learning_rate=3e-5,
+      use_wandb=True,
+      distributed=True,
+      gpu_ids=gpu_ids,
+      log_every_n_steps=50,
+      val_check_interval=0.5,
+  )
+  train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
+  manager = mp.Manager()
+  return_dict = manager.dict()
+
+  # Launch processes
+  mp.spawn(
+      setup_process,
+      args=(world_size, model, config, train_dataset, val_dataset, return_dict),
+      nprocs=world_size,
+      join=True,
+  )
+
+  results = return_dict.get("results", None)
+  print("\nFinetuning completed!")
+  return results
+
+
+def main(argv):
+  """Main function that selects and runs the appropriate training mode."""
+
+  try:
+    if FLAGS.training_mode == "single":
+      print("\nStarting single-GPU training...")
+      single_gpu_example()
+    else:
+      gpu_ids = [int(id) for id in FLAGS.gpu_ids]
+      print(f"\nStarting multi-GPU training using GPUs: {gpu_ids}...")
+
+      config = FinetuningConfig(
+          batch_size=256,
+          num_epochs=5,
+          learning_rate=3e-5,
+          use_wandb=True,
+          distributed=True,
+          gpu_ids=gpu_ids,
+      )
+
+      results = multi_gpu_example(config)
+      print("\nMulti-GPU training completed!")
+
+  except Exception as e:
+    print(f"Training failed: {str(e)}")
+  finally:
+    if torch.distributed.is_initialized():
+      torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/src/finetuning/finetuning_torch.py b/src/finetuning/finetuning_torch.py
new file mode 100644
index 0000000..c065234
--- /dev/null
+++ b/src/finetuning/finetuning_torch.py
@@ -0,0 +1,398 @@
+"""
+TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
+"""
+
+import logging
+import os
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, 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 timesfm.patched_decoder import DEFAULT_QUANTILES
+
+import wandb
+
+
+class MetricsLogger(ABC):
+  """Abstract base class for logging metrics during training.
+
+    This class defines the interface for logging metrics during model training.
+    Concrete implementations can log to different backends (e.g., WandB, TensorBoard).
+    """
+
+  @abstractmethod
+  def log_metrics(self,
+                  metrics: Dict[str, Any],
+                  step: Optional[int] = None) -> None:
+    """Log metrics to the specified backend.
+
+        Args:
+          metrics: Dictionary containing metric names and values.
+          step: Optional step number or epoch for the metrics.
+        """
+    pass
+
+  @abstractmethod
+  def close(self) -> None:
+    """Clean up any resources used by the logger."""
+    pass
+
+
+class WandBLogger(MetricsLogger):
+  """Weights & Biases implementation of metrics logging.
+
+    Args:
+      project: Name of the W&B project.
+      config: Configuration dictionary to log.
+      rank: Process rank in distributed training.
+    """
+
+  def __init__(self, project: str, config: Dict[str, Any], rank: int = 0):
+    self.rank = rank
+    if rank == 0:
+      wandb.init(project=project, config=config)
+
+  def log_metrics(self,
+                  metrics: Dict[str, Any],
+                  step: Optional[int] = None) -> None:
+    """Log metrics to W&B if on the main process.
+
+        Args:
+          metrics: Dictionary of metrics to log.
+          step: Current training step or epoch.
+        """
+    if self.rank == 0:
+      wandb.log(metrics, step=step)
+
+  def close(self) -> None:
+    """Finish the W&B run if on the main process."""
+    if self.rank == 0:
+      wandb.finish()
+
+
+class DistributedManager:
+  """Manages distributed training setup and cleanup.
+
+    Args:
+      world_size: Total number of processes.
+      rank: Process rank.
+      master_addr: Address of the master process.
+      master_port: Port for distributed communication.
+      backend: PyTorch distributed backend to use.
+    """
+
+  def __init__(
+      self,
+      world_size: int,
+      rank: int,
+      master_addr: str = "localhost",
+      master_port: str = "12358",
+      backend: str = "nccl",
+  ):
+    self.world_size = world_size
+    self.rank = rank
+    self.master_addr = master_addr
+    self.master_port = master_port
+    self.backend = backend
+
+  def setup(self) -> None:
+    """Initialize the distributed environment."""
+    os.environ["MASTER_ADDR"] = self.master_addr
+    os.environ["MASTER_PORT"] = self.master_port
+
+    if not dist.is_initialized():
+      dist.init_process_group(backend=self.backend,
+                              world_size=self.world_size,
+                              rank=self.rank)
+
+  def cleanup(self) -> None:
+    """Clean up the distributed environment."""
+    if dist.is_initialized():
+      dist.destroy_process_group()
+
+
+@dataclass
+class FinetuningConfig:
+  """Configuration for model training.
+
+    Args:
+      batch_size: Number of samples per batch.
+      num_epochs: Number of training epochs.
+      learning_rate: Initial learning rate.
+      weight_decay: L2 regularization factor.
+      freq_type: Frequency, can be [0, 1, 2].
+      use_quantile_loss: bool = False  # Flag to enable/disable quantile loss
+      quantiles: List[float] = field(default_factory=lambda: [0.1, 0.5, 0.9])
+      device: Device to train on ('cuda' or 'cpu').
+      distributed: Whether to use distributed training.
+      gpu_ids: List of GPU IDs to use.
+      master_port: Port for distributed training.
+      master_addr: Address for distributed training.
+      use_wandb: Whether to use Weights & Biases logging.
+      wandb_project: W&B project name.
+      log_every_n_steps: Log metrics every N steps (batches), this is inspired from Pytorch Lightning
+      val_check_interval: How often within one training epoch to check val metrics. (also from Pytorch Lightning)
+        Can be: float (0.0-1.0): fraction of epoch (e.g., 0.5 = validate twice per epoch)
+                int: validate every N batches
+    """
+
+  batch_size: int = 32
+  num_epochs: int = 20
+  learning_rate: float = 1e-4
+  weight_decay: float = 0.01
+  freq_type: int = 0
+  use_quantile_loss: bool = False
+  device: str = "cuda" if torch.cuda.is_available() else "cpu"
+  distributed: bool = False
+  gpu_ids: List[int] = field(default_factory=lambda: [0])
+  master_port: str = "12358"
+  master_addr: str = "localhost"
+  use_wandb: bool = False
+  wandb_project: str = "timesfm-finetuning"
+  log_every_n_steps: int = 50
+  val_check_interval: float = 0.5
+
+
+class TimesFMFinetuner:
+  """Handles model training and validation.
+
+    Args:
+      model: PyTorch model to train.
+      config: Training configuration.
+      rank: Process rank for distributed training.
+      loss_fn: Loss function (defaults to MSE).
+      logger: Optional logging.Logger instance.
+    """
+
+  def __init__(
+      self,
+      model: nn.Module,
+      config: FinetuningConfig,
+      rank: int = 0,
+      loss_fn: Optional[Callable] = None,
+      logger: Optional[logging.Logger] = None,
+  ):
+    self.model = model
+    self.config = config
+    self.rank = rank
+    self.logger = logger or logging.getLogger(__name__)
+    self.device = torch.device(
+        f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
+    self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1))**2))
+
+    if config.use_wandb:
+      self.metrics_logger = WandBLogger(config.wandb_project, config.__dict__,
+                                        rank)
+
+    if config.distributed:
+      self.dist_manager = DistributedManager(
+          world_size=len(config.gpu_ids),
+          rank=rank,
+          master_addr=config.master_addr,
+          master_port=config.master_port,
+      )
+      self.dist_manager.setup()
+      self.model = self._setup_distributed_model()
+
+  def _setup_distributed_model(self) -> nn.Module:
+    """Configure model for distributed training."""
+    self.model = self.model.to(self.device)
+    return DDP(self.model,
+               device_ids=[self.config.gpu_ids[self.rank]],
+               output_device=self.config.gpu_ids[self.rank])
+
+  def _create_dataloader(self, dataset: Dataset, is_train: bool) -> DataLoader:
+    """Create appropriate DataLoader based on training configuration.
+
+        Args:
+          dataset: Dataset to create loader for.
+          is_train: Whether this is for training (affects shuffling).
+
+        Returns:
+          DataLoader instance.
+        """
+    if self.config.distributed:
+      sampler = torch.utils.data.distributed.DistributedSampler(
+          dataset,
+          num_replicas=len(self.config.gpu_ids),
+          rank=dist.get_rank(),
+          shuffle=is_train)
+    else:
+      sampler = None
+
+    return DataLoader(
+        dataset,
+        batch_size=self.config.batch_size,
+        shuffle=(is_train and not self.config.distributed),
+        sampler=sampler,
+    )
+
+  def _quantile_loss(self, pred: torch.Tensor, actual: torch.Tensor,
+                     quantile: float) -> torch.Tensor:
+    """Calculates quantile loss.
+        Args:
+            pred: Predicted values
+            actual: Actual values
+            quantile: Quantile at which loss is computed
+        Returns:
+            Quantile loss
+        """
+    dev = actual - pred
+    loss_first = dev * quantile
+    loss_second = -dev * (1.0 - quantile)
+    return 2 * torch.where(loss_first >= 0, loss_first, loss_second)
+
+  def _process_batch(self, batch: List[torch.Tensor]) -> tuple:
+    """Process a single batch of data.
+
+        Args:
+          batch: List of input tensors.
+
+        Returns:
+          Tuple of (loss, predictions).
+        """
+    x_context, x_padding, freq, x_future = [
+        t.to(self.device, non_blocking=True) for t in batch
+    ]
+
+    predictions = self.model(x_context, x_padding.float(), freq)
+    predictions_mean = predictions[..., 0]
+    last_patch_pred = predictions_mean[:, -1, :]
+
+    loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+    if self.config.use_quantile_loss:
+      quantiles = self.config.quantiles or DEFAULT_QUANTILES
+      for i, quantile in enumerate(quantiles):
+        last_patch_quantile = predictions[:, -1, :, i + 1]
+        loss += torch.mean(
+            self._quantile_loss(last_patch_quantile, x_future.squeeze(-1),
+                                quantile))
+
+    return loss, predictions
+
+  def _train_epoch(self, train_loader: DataLoader,
+                   optimizer: torch.optim.Optimizer) -> float:
+    """Train for one epoch in a distributed setting.
+
+        Args:
+            train_loader: DataLoader for training data.
+            optimizer: Optimizer instance.
+
+        Returns:
+            Average training loss for the epoch.
+        """
+    self.model.train()
+    total_loss = 0.0
+    num_batches = len(train_loader)
+
+    for batch in train_loader:
+      loss, _ = self._process_batch(batch)
+
+      optimizer.zero_grad()
+      loss.backward()
+      optimizer.step()
+
+      total_loss += loss.item()
+
+    avg_loss = total_loss / num_batches
+
+    if self.config.distributed:
+      avg_loss_tensor = torch.tensor(avg_loss, device=self.device)
+      dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM)
+      avg_loss = (avg_loss_tensor / dist.get_world_size()).item()
+
+    return avg_loss
+
+  def _validate(self, val_loader: DataLoader) -> float:
+    """Perform validation.
+
+        Args:
+            val_loader: DataLoader for validation data.
+
+        Returns:
+            Average validation loss.
+        """
+    self.model.eval()
+    total_loss = 0.0
+    num_batches = len(val_loader)
+
+    with torch.no_grad():
+      for batch in val_loader:
+        loss, _ = self._process_batch(batch)
+        total_loss += loss.item()
+
+    avg_loss = total_loss / num_batches
+
+    if self.config.distributed:
+      avg_loss_tensor = torch.tensor(avg_loss, device=self.device)
+      dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM)
+      avg_loss = (avg_loss_tensor / dist.get_world_size()).item()
+
+    return avg_loss
+
+  def finetune(self, train_dataset: Dataset,
+               val_dataset: Dataset) -> Dict[str, Any]:
+    """Train the model.
+
+        Args:
+          train_dataset: Training dataset.
+          val_dataset: Validation dataset.
+
+        Returns:
+          Dictionary containing training history.
+        """
+    self.model = self.model.to(self.device)
+    train_loader = self._create_dataloader(train_dataset, is_train=True)
+    val_loader = self._create_dataloader(val_dataset, is_train=False)
+
+    optimizer = torch.optim.Adam(self.model.parameters(),
+                                 lr=self.config.learning_rate,
+                                 weight_decay=self.config.weight_decay)
+
+    history = {"train_loss": [], "val_loss": [], "learning_rate": []}
+
+    self.logger.info(
+        f"Starting training for {self.config.num_epochs} epochs...")
+    self.logger.info(f"Training samples: {len(train_dataset)}")
+    self.logger.info(f"Validation samples: {len(val_dataset)}")
+
+    try:
+      for epoch in range(self.config.num_epochs):
+        train_loss = self._train_epoch(train_loader, optimizer)
+        val_loss = self._validate(val_loader)
+        current_lr = optimizer.param_groups[0]["lr"]
+
+        metrics = {
+            "train_loss": train_loss,
+            "val_loss": val_loss,
+            "learning_rate": current_lr,
+            "epoch": epoch + 1,
+        }
+
+        if self.config.use_wandb:
+          self.metrics_logger.log_metrics(metrics)
+
+        history["train_loss"].append(train_loss)
+        history["val_loss"].append(val_loss)
+        history["learning_rate"].append(current_lr)
+
+        if self.rank == 0:
+          self.logger.info(
+              f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}"
+          )
+
+    except KeyboardInterrupt:
+      self.logger.info("Training interrupted by user")
+
+    if self.config.distributed:
+      self.dist_manager.cleanup()
+
+    if self.config.use_wandb:
+      self.metrics_logger.close()
+
+    return {"history": history}
diff --git a/src/timesfm/finetuning_torch.py b/src/timesfm/finetuning_torch.py
deleted file mode 100644
index af2eadb..0000000
--- a/src/timesfm/finetuning_torch.py
+++ /dev/null
@@ -1,340 +0,0 @@
-"""
-TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
-"""
-
-import logging
-import os
-from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, 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
-
-import wandb
-
-
-class MetricsLogger(ABC):
-  """Abstract base class for logging metrics during training.
-
-  This class defines the interface for logging metrics during model training.
-  Concrete implementations can log to different backends (e.g., WandB, TensorBoard).
-  """
-
-  @abstractmethod
-  def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
-    """Log metrics to the specified backend.
-
-    Args:
-      metrics: Dictionary containing metric names and values.
-      step: Optional step number or epoch for the metrics.
-    """
-    pass
-
-  @abstractmethod
-  def close(self) -> None:
-    """Clean up any resources used by the logger."""
-    pass
-
-
-class WandBLogger(MetricsLogger):
-  """Weights & Biases implementation of metrics logging.
-
-  Args:
-    project: Name of the W&B project.
-    config: Configuration dictionary to log.
-    rank: Process rank in distributed training.
-  """
-
-  def __init__(self, project: str, config: Dict[str, Any], rank: int = 0):
-    self.rank = rank
-    if rank == 0:
-      wandb.init(project=project, config=config)
-
-  def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
-    """Log metrics to W&B if on the main process.
-
-    Args:
-      metrics: Dictionary of metrics to log.
-      step: Current training step or epoch.
-    """
-    if self.rank == 0:
-      wandb.log(metrics, step=step)
-
-  def close(self) -> None:
-    """Finish the W&B run if on the main process."""
-    if self.rank == 0:
-      wandb.finish()
-
-
-class DistributedManager:
-  """Manages distributed training setup and cleanup.
-
-  Args:
-    world_size: Total number of processes.
-    rank: Process rank.
-    master_addr: Address of the master process.
-    master_port: Port for distributed communication.
-    backend: PyTorch distributed backend to use.
-  """
-
-  def __init__(
-    self,
-    world_size: int,
-    rank: int,
-    master_addr: str = "localhost",
-    master_port: str = "12358",
-    backend: str = "nccl",
-  ):
-    self.world_size = world_size
-    self.rank = rank
-    self.master_addr = master_addr
-    self.master_port = master_port
-    self.backend = backend
-
-  def setup(self) -> None:
-    """Initialize the distributed environment."""
-    os.environ["MASTER_ADDR"] = self.master_addr
-    os.environ["MASTER_PORT"] = self.master_port
-
-    if not dist.is_initialized():
-      dist.init_process_group(backend=self.backend, world_size=self.world_size, rank=self.rank)
-
-  def cleanup(self) -> None:
-    """Clean up the distributed environment."""
-    if dist.is_initialized():
-      dist.destroy_process_group()
-
-
-@dataclass
-class FinetuningConfig:
-  """Configuration for model training.
-
-  Args:
-    batch_size: Number of samples per batch.
-    num_epochs: Number of training epochs.
-    learning_rate: Initial learning rate.
-    weight_decay: L2 regularization factor.
-    device: Device to train on ('cuda' or 'cpu').
-    distributed: Whether to use distributed training.
-    gpu_ids: List of GPU IDs to use.
-    master_port: Port for distributed training.
-    master_addr: Address for distributed training.
-    use_wandb: Whether to use Weights & Biases logging.
-    wandb_project: W&B project name.
-  """
-
-  batch_size: int = 32
-  num_epochs: int = 20
-  learning_rate: float = 1e-4
-  weight_decay: float = 0.01
-  device: str = "cuda" if torch.cuda.is_available() else "cpu"
-  distributed: bool = False
-  gpu_ids: List[int] = field(default_factory=lambda: [0])
-  master_port: str = "12358"
-  master_addr: str = "localhost"
-  use_wandb: bool = False
-  wandb_project: str = "timesfm-finetuning"
-
-
-class TimesFMFinetuner:
-  """Handles model training and validation.
-
-  Args:
-    model: PyTorch model to train.
-    config: Training configuration.
-    rank: Process rank for distributed training.
-    loss_fn: Loss function (defaults to MSE).
-    logger: Optional logging.Logger instance.
-  """
-
-  def __init__(
-    self,
-    model: nn.Module,
-    config: FinetuningConfig,
-    rank: int = 0,
-    loss_fn: Optional[Callable] = None,
-    logger: Optional[logging.Logger] = None,
-  ):
-    self.model = model
-    self.config = config
-    self.rank = rank
-    self.logger = logger or logging.getLogger(__name__)
-    self.device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
-    self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))
-
-    if config.use_wandb:
-      self.metrics_logger = WandBLogger(config.wandb_project, config.__dict__, rank)
-
-    if config.distributed:
-      self.dist_manager = DistributedManager(
-        world_size=len(config.gpu_ids),
-        rank=rank,
-        master_addr=config.master_addr,
-        master_port=config.master_port,
-      )
-      self.dist_manager.setup()
-      self.model = self._setup_distributed_model()
-
-  def _setup_distributed_model(self) -> nn.Module:
-    """Configure model for distributed training."""
-    self.model = self.model.to(self.device)
-    return DDP(
-      self.model, device_ids=[self.config.gpu_ids[self.rank]], output_device=self.config.gpu_ids[self.rank]
-    )
-
-  def _create_dataloader(self, dataset: Dataset, is_train: bool) -> DataLoader:
-    """Create appropriate DataLoader based on training configuration.
-
-    Args:
-      dataset: Dataset to create loader for.
-      is_train: Whether this is for training (affects shuffling).
-
-    Returns:
-      DataLoader instance.
-    """
-    if self.config.distributed:
-      sampler = torch.utils.data.distributed.DistributedSampler(
-        dataset, num_replicas=len(self.config.gpu_ids), rank=dist.get_rank(), shuffle=is_train
-      )
-    else:
-      sampler = None
-
-    return DataLoader(
-      dataset,
-      batch_size=self.config.batch_size,
-      shuffle=(is_train and not self.config.distributed),
-      sampler=sampler,
-    )
-
-  def _process_batch(self, batch: List[torch.Tensor]) -> tuple:
-    """Process a single batch of data.
-
-    Args:
-      batch: List of input tensors.
-
-    Returns:
-      Tuple of (loss, predictions).
-    """
-    x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
-
-    predictions = self.model(x_context, x_padding.float(), freq)
-    predictions_mean = predictions[..., 0]
-    last_patch_pred = predictions_mean[:, -1, :]
-
-    loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
-
-    return loss, predictions
-
-  def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
-    """Train for one epoch.
-
-    Args:
-      train_loader: DataLoader for training data.
-      optimizer: Optimizer instance.
-
-    Returns:
-      Average training loss for the epoch.
-    """
-    self.model.train()
-    total_loss = 0.0
-
-    for batch in train_loader:
-      loss, _ = self._process_batch(batch)
-
-      if self.config.distributed:
-        losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
-        dist.all_gather(losses, loss)
-
-      optimizer.zero_grad()
-      loss.backward()
-      optimizer.step()
-
-      total_loss += loss.item()
-
-    return total_loss / len(train_loader)
-
-  def _validate(self, val_loader: DataLoader) -> float:
-    """Perform validation.
-
-    Args:
-      val_loader: DataLoader for validation data.
-
-    Returns:
-      Average validation loss.
-    """
-    self.model.eval()
-    total_loss = 0.0
-
-    with torch.no_grad():
-      for batch in val_loader:
-        loss, _ = self._process_batch(batch)
-
-        if self.config.distributed:
-          losses = [torch.zeros_like(loss) for _ in range(dist.get_world_size())]
-          dist.all_gather(losses, loss)
-
-        total_loss += loss.item()
-
-    return total_loss / len(val_loader)
-
-  def finetune(self, train_dataset: Dataset, val_dataset: Dataset) -> Dict[str, Any]:
-    """Train the model.
-
-    Args:
-      train_dataset: Training dataset.
-      val_dataset: Validation dataset.
-
-    Returns:
-      Dictionary containing training history.
-    """
-    self.model = self.model.to(self.device)
-    train_loader = self._create_dataloader(train_dataset, is_train=True)
-    val_loader = self._create_dataloader(val_dataset, is_train=False)
-
-    optimizer = torch.optim.Adam(
-      self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
-    )
-
-    history = {"train_loss": [], "val_loss": [], "learning_rate": []}
-
-    self.logger.info(f"Starting training for {self.config.num_epochs} epochs...")
-    self.logger.info(f"Training samples: {len(train_dataset)}")
-    self.logger.info(f"Validation samples: {len(val_dataset)}")
-
-    try:
-      for epoch in range(self.config.num_epochs):
-        train_loss = self._train_epoch(train_loader, optimizer)
-        val_loss = self._validate(val_loader)
-        current_lr = optimizer.param_groups[0]["lr"]
-
-        metrics = {
-          "train_loss": train_loss,
-          "val_loss": val_loss,
-          "learning_rate": current_lr,
-          "epoch": epoch + 1,
-        }
-
-        if self.config.use_wandb:
-          self.metrics_logger.log_metrics(metrics)
-
-        history["train_loss"].append(train_loss)
-        history["val_loss"].append(val_loss)
-        history["learning_rate"].append(current_lr)
-
-        if self.rank == 0:
-          self.logger.info(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
-
-    except KeyboardInterrupt:
-      self.logger.info("Training interrupted by user")
-
-    if self.config.distributed:
-      self.dist_manager.cleanup()
-
-    if self.config.use_wandb:
-      self.metrics_logger.close()
-
-    return {"history": history}
\ No newline at end of file

From ca87a438a0341e25acdba7a0cfe658a441fd7850 Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkin@google.com>
Date: Sat, 1 Feb 2025 02:30:17 +0000
Subject: [PATCH 10/12] Update examples with new feedback

---
 notebooks/finetuning_torch.ipynb | 178 ++++++++++++++++++-------------
 1 file changed, 104 insertions(+), 74 deletions(-)

diff --git a/notebooks/finetuning_torch.ipynb b/notebooks/finetuning_torch.ipynb
index 9515e7e..c50dca1 100644
--- a/notebooks/finetuning_torch.ipynb
+++ b/notebooks/finetuning_torch.ipynb
@@ -44,18 +44,27 @@
     "class TimeSeriesDataset(Dataset):\n",
     "  \"\"\"Dataset for time series data compatible with TimesFM.\"\"\"\n",
     "\n",
-    "  def __init__(self, series: np.ndarray, context_length: int, horizon_length: int):\n",
+    "  def __init__(self,\n",
+    "               series: np.ndarray,\n",
+    "               context_length: int,\n",
+    "               horizon_length: int,\n",
+    "               freq_type: int = 0):\n",
     "    \"\"\"\n",
-    "    Initialize dataset.\n",
+    "        Initialize dataset.\n",
+    "\n",
+    "        Args:\n",
+    "            series: Time series data\n",
+    "            context_length: Number of past timesteps to use as input\n",
+    "            horizon_length: Number of future timesteps to predict\n",
+    "            freq_type: Frequency type (0, 1, or 2)\n",
+    "        \"\"\"\n",
+    "    if freq_type not in [0, 1, 2]:\n",
+    "      raise ValueError(\"freq_type must be 0, 1, or 2\")\n",
     "\n",
-    "    Args:\n",
-    "      series: Time series data\n",
-    "      context_length: Number of past timesteps to use as input\n",
-    "      horizon_length: Number of future timesteps to predict\n",
-    "    \"\"\"\n",
     "    self.series = series\n",
     "    self.context_length = context_length\n",
     "    self.horizon_length = horizon_length\n",
+    "    self.freq_type = freq_type\n",
     "    self._prepare_samples()\n",
     "\n",
     "  def _prepare_samples(self) -> None:\n",
@@ -66,47 +75,57 @@
     "    for start_idx in range(0, len(self.series) - total_length + 1):\n",
     "      end_idx = start_idx + self.context_length\n",
     "      x_context = self.series[start_idx:end_idx]\n",
-    "      x_future = self.series[end_idx : end_idx + self.horizon_length]\n",
+    "      x_future = self.series[end_idx:end_idx + self.horizon_length]\n",
     "      self.samples.append((x_context, x_future))\n",
     "\n",
     "  def __len__(self) -> int:\n",
     "    return len(self.samples)\n",
     "\n",
-    "  def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n",
+    "  def __getitem__(\n",
+    "      self, index: int\n",
+    "  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n",
     "    x_context, x_future = self.samples[index]\n",
     "\n",
     "    x_context = torch.tensor(x_context, dtype=torch.float32)\n",
     "    x_future = torch.tensor(x_future, dtype=torch.float32)\n",
     "\n",
     "    input_padding = torch.zeros_like(x_context)\n",
-    "    freq = torch.zeros(1, dtype=torch.long)\n",
+    "    freq = torch.tensor([self.freq_type], dtype=torch.long)\n",
     "\n",
     "    return x_context, input_padding, freq, x_future\n",
     "\n",
-    "\n",
-    "def prepare_datasets(\n",
-    "  series: np.ndarray, context_length: int, horizon_length: int, train_split: float = 0.8\n",
-    ") -> Tuple[Dataset, Dataset]:\n",
+    "def prepare_datasets(series: np.ndarray,\n",
+    "                     context_length: int,\n",
+    "                     horizon_length: int,\n",
+    "                     freq_type: int = 0,\n",
+    "                     train_split: float = 0.8) -> Tuple[Dataset, Dataset]:\n",
     "  \"\"\"\n",
-    "  Prepare training and validation datasets from time series data.\n",
+    "    Prepare training and validation datasets from time series data.\n",
     "\n",
-    "  Args:\n",
-    "    series: Input time series data\n",
-    "    context_length: Number of past timesteps to use\n",
-    "    horizon_length: Number of future timesteps to predict\n",
-    "    train_split: Fraction of data to use for training\n",
+    "    Args:\n",
+    "        series: Input time series data\n",
+    "        context_length: Number of past timesteps to use\n",
+    "        horizon_length: Number of future timesteps to predict\n",
+    "        freq_type: Frequency type (0, 1, or 2)\n",
+    "        train_split: Fraction of data to use for training\n",
     "\n",
-    "  Returns:\n",
-    "    Tuple of (train_dataset, val_dataset)\n",
-    "  \"\"\"\n",
+    "    Returns:\n",
+    "        Tuple of (train_dataset, val_dataset)\n",
+    "    \"\"\"\n",
     "  train_size = int(len(series) * train_split)\n",
     "  train_data = series[:train_size]\n",
     "  val_data = series[train_size:]\n",
     "\n",
-    "  # Create datasets\n",
-    "  train_dataset = TimeSeriesDataset(train_data, context_length=context_length, horizon_length=horizon_length)\n",
+    "  # Create datasets with specified frequency type\n",
+    "  train_dataset = TimeSeriesDataset(train_data,\n",
+    "                                    context_length=context_length,\n",
+    "                                    horizon_length=horizon_length,\n",
+    "                                    freq_type=freq_type)\n",
     "\n",
-    "  val_dataset = TimeSeriesDataset(val_data, context_length=context_length, horizon_length=horizon_length)\n",
+    "  val_dataset = TimeSeriesDataset(val_data,\n",
+    "                                  context_length=context_length,\n",
+    "                                  horizon_length=horizon_length,\n",
+    "                                  freq_type=freq_type)\n",
     "\n",
     "  return train_dataset, val_dataset\n"
    ]
@@ -128,14 +147,16 @@
     "  device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
     "  repo_id = \"google/timesfm-2.0-500m-pytorch\"\n",
     "  hparams = TimesFmHparams(\n",
-    "    backend=device,\n",
-    "    per_core_batch_size=32,\n",
-    "    horizon_len=128,\n",
-    "    num_layers=50,\n",
-    "    use_positional_embedding=False,\n",
-    "    context_len=192,\n",
+    "      backend=device,\n",
+    "      per_core_batch_size=32,\n",
+    "      horizon_len=128,\n",
+    "      num_layers=50,\n",
+    "      use_positional_embedding=False,\n",
+    "      context_len=\n",
+    "      192,  # Context length can be anything up to 2048 in multiples of 32\n",
     "  )\n",
-    "  tfm = TimesFm(hparams=hparams, checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))\n",
+    "  tfm = TimesFm(hparams=hparams,\n",
+    "                checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))\n",
     "\n",
     "  model = PatchedTimeSeriesDecoder(tfm._model_config)\n",
     "  if load_weights:\n",
@@ -152,18 +173,18 @@
    "outputs": [],
    "source": [
     "def plot_predictions(\n",
-    "  model: TimesFm,\n",
-    "  val_dataset: Dataset,\n",
-    "  save_path: Optional[str] = \"predictions.png\",\n",
+    "    model: TimesFm,\n",
+    "    val_dataset: Dataset,\n",
+    "    save_path: Optional[str] = \"predictions.png\",\n",
     ") -> None:\n",
     "  \"\"\"\n",
-    "  Plot model predictions against ground truth for a batch of validation data.\n",
+    "    Plot model predictions against ground truth for a batch of validation data.\n",
     "\n",
-    "  Args:\n",
-    "    model: Trained TimesFM model\n",
-    "    val_dataset: Validation dataset\n",
-    "    save_path: Path to save the plot\n",
-    "  \"\"\"\n",
+    "    Args:\n",
+    "      model: Trained TimesFM model\n",
+    "      val_dataset: Validation dataset\n",
+    "      save_path: Path to save the plot\n",
+    "    \"\"\"\n",
     "  import matplotlib.pyplot as plt\n",
     "\n",
     "  model.eval()\n",
@@ -194,18 +215,26 @@
     "\n",
     "  plt.figure(figsize=(12, 6))\n",
     "\n",
-    "  plt.plot(range(context_len), context_vals, label=\"Historical Data\", color=\"blue\", linewidth=2)\n",
+    "  plt.plot(range(context_len),\n",
+    "           context_vals,\n",
+    "           label=\"Historical Data\",\n",
+    "           color=\"blue\",\n",
+    "           linewidth=2)\n",
     "\n",
     "  plt.plot(\n",
-    "    range(context_len, context_len + horizon_len),\n",
-    "    future_vals,\n",
-    "    label=\"Ground Truth\",\n",
-    "    color=\"green\",\n",
-    "    linestyle=\"--\",\n",
-    "    linewidth=2,\n",
+    "      range(context_len, context_len + horizon_len),\n",
+    "      future_vals,\n",
+    "      label=\"Ground Truth\",\n",
+    "      color=\"green\",\n",
+    "      linestyle=\"--\",\n",
+    "      linewidth=2,\n",
     "  )\n",
     "\n",
-    "  plt.plot(range(context_len, context_len + horizon_len), pred_vals, label=\"Prediction\", color=\"red\", linewidth=2)\n",
+    "  plt.plot(range(context_len, context_len + horizon_len),\n",
+    "           pred_vals,\n",
+    "           label=\"Prediction\",\n",
+    "           color=\"red\",\n",
+    "           linewidth=2)\n",
     "\n",
     "  plt.xlabel(\"Time Step\")\n",
     "  plt.ylabel(\"Value\")\n",
@@ -218,43 +247,44 @@
     "    print(f\"Plot saved to {save_path}\")\n",
     "\n",
     "  plt.close()\n",
-    "\n",
-    "\n",
-    "def get_data(context_len: int, horizon_len: int) -> Tuple[Dataset, Dataset]:\n",
-    "  df = yf.download(\"AAPL\", start=\"2010-01-01\", end=\"2019-01-01\")\n",
-    "  time_series = df[\"Close\"].values\n",
-    "\n",
-    "  train_dataset, val_dataset = prepare_datasets(\n",
-    "    series=time_series,\n",
-    "    context_length=context_len,\n",
-    "    horizon_length=horizon_len,\n",
-    "    train_split=0.8,\n",
-    "  )\n",
-    "\n",
-    "  print(f\"Created datasets:\")\n",
-    "  print(f\"- Training samples: {len(train_dataset)}\")\n",
-    "  print(f\"- Validation samples: {len(val_dataset)}\")\n",
-    "  return train_dataset, val_dataset\n",
-    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "\n",
     "def single_gpu_example():\n",
     "  \"\"\"Basic example of finetuning TimesFM on stock data.\"\"\"\n",
     "  model, hparams, tfm_config = get_model(load_weights=True)\n",
-    "  config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=True)\n",
+    "  config = FinetuningConfig(batch_size=256,\n",
+    "                            num_epochs=5,\n",
+    "                            learning_rate=1e-4,\n",
+    "                            use_wandb=True,\n",
+    "                            freq_type=1,\n",
+    "                            log_every_n_steps=10,\n",
+    "                            val_check_interval=0.5,\n",
+    "                            use_quantile_loss=True)\n",
     "\n",
-    "  train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)\n",
+    "  train_dataset, val_dataset = get_data(128,\n",
+    "                                        tfm_config.horizon_len,\n",
+    "                                        freq_type=config.freq_type)\n",
     "  finetuner = TimesFMFinetuner(model, config)\n",
     "\n",
     "  print(\"\\nStarting finetuning...\")\n",
-    "  results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)\n",
+    "  results = finetuner.finetune(train_dataset=train_dataset,\n",
+    "                               val_dataset=val_dataset)\n",
     "\n",
     "  print(\"\\nFinetuning completed!\")\n",
     "  print(f\"Training history: {len(results['history']['train_loss'])} epochs\")\n",
     "\n",
     "  plot_predictions(\n",
-    "    model=model,\n",
-    "    val_dataset=val_dataset,\n",
-    "    save_path=\"timesfm_predictions.png\",\n",
+    "      model=model,\n",
+    "      val_dataset=val_dataset,\n",
+    "      save_path=\"timesfm_predictions.png\",\n",
     "  )\n"
    ]
   },

From 402ebf52a5d20cac122c64584b057c4825bc0bbf Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkinmichael@gmail.com>
Date: Sat, 1 Feb 2025 02:43:56 +0000
Subject: [PATCH 11/12] Small fix with default value

---
 src/finetuning/finetuning_torch.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/finetuning/finetuning_torch.py b/src/finetuning/finetuning_torch.py
index c065234..48f7dce 100644
--- a/src/finetuning/finetuning_torch.py
+++ b/src/finetuning/finetuning_torch.py
@@ -13,7 +13,7 @@ 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 timesfm.patched_decoder import DEFAULT_QUANTILES
+from timesfm.pytorch_patched_decoder import _create_quantiles
 
 import wandb
 
@@ -127,7 +127,7 @@ class FinetuningConfig:
       weight_decay: L2 regularization factor.
       freq_type: Frequency, can be [0, 1, 2].
       use_quantile_loss: bool = False  # Flag to enable/disable quantile loss
-      quantiles: List[float] = field(default_factory=lambda: [0.1, 0.5, 0.9])
+      quantiles: Optional[List[float]] = None
       device: Device to train on ('cuda' or 'cpu').
       distributed: Whether to use distributed training.
       gpu_ids: List of GPU IDs to use.
@@ -147,6 +147,7 @@ class FinetuningConfig:
   weight_decay: float = 0.01
   freq_type: int = 0
   use_quantile_loss: bool = False
+  quantiles: Optional[List[float]] = None
   device: str = "cuda" if torch.cuda.is_available() else "cpu"
   distributed: bool = False
   gpu_ids: List[int] = field(default_factory=lambda: [0])
@@ -266,7 +267,7 @@ class TimesFMFinetuner:
 
     loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
     if self.config.use_quantile_loss:
-      quantiles = self.config.quantiles or DEFAULT_QUANTILES
+      quantiles = self.config.quantiles or _create_quantiles()
       for i, quantile in enumerate(quantiles):
         last_patch_quantile = predictions[:, -1, :, i + 1]
         loss += torch.mean(

From 67e0eaaa3d2c366f6c1cd2cb95a3c983ed41e225 Mon Sep 17 00:00:00 2001
From: misha-chertushkin <chertushkin@google.com>
Date: Wed, 5 Feb 2025 13:41:13 +0000
Subject: [PATCH 12/12] Quantiles PR nit fix

---
 src/finetuning/finetuning_torch.py     | 4 ++--
 src/timesfm/pytorch_patched_decoder.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/finetuning/finetuning_torch.py b/src/finetuning/finetuning_torch.py
index 48f7dce..5c2d8b3 100644
--- a/src/finetuning/finetuning_torch.py
+++ b/src/finetuning/finetuning_torch.py
@@ -13,7 +13,7 @@ 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 timesfm.pytorch_patched_decoder import _create_quantiles
+from timesfm.pytorch_patched_decoder import create_quantiles
 
 import wandb
 
@@ -267,7 +267,7 @@ class TimesFMFinetuner:
 
     loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
     if self.config.use_quantile_loss:
-      quantiles = self.config.quantiles or _create_quantiles()
+      quantiles = self.config.quantiles or create_quantiles()
       for i, quantile in enumerate(quantiles):
         last_patch_quantile = predictions[:, -1, :, i + 1]
         loss += torch.mean(
diff --git a/src/timesfm/pytorch_patched_decoder.py b/src/timesfm/pytorch_patched_decoder.py
index 67f6be4..15bf428 100644
--- a/src/timesfm/pytorch_patched_decoder.py
+++ b/src/timesfm/pytorch_patched_decoder.py
@@ -21,7 +21,7 @@ from torch import nn
 import torch.nn.functional as F
 
 
-def _create_quantiles() -> list[float]:
+def create_quantiles() -> list[float]:
   return [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
 
 
@@ -48,7 +48,7 @@ class TimesFMConfig:
   # Horizon length
   horizon_len: int = 128
   # quantiles
-  quantiles: List[float] = dataclasses.field(default_factory=_create_quantiles)
+  quantiles: List[float] = dataclasses.field(default_factory=create_quantiles)
   # Padding value
   pad_val: float = 1123581321.0
   # Tolerance