train/train.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Fine-tune Toto-2 on cleaned TLE daily series with space-weather channels (v2).

Objective: next-patch quantile (pinball) loss in Toto's asinh-scaled space, the
same as v1, but the loss is applied ONLY to the orbital channels -- the solar
channels (F10.7, Ap) are input-only context (LOSS_CHANNEL_MASK), so the model
uses them to inform drag but is not asked to forecast them.

Run (smoke, 4m, one year):
  python v2/train/train.py \
    --years 2020 --model Datadog/Toto-2.0-4m \
    --sw-csv v2/data/SW-All.csv --window-patches 3 \
    --batch-size 64 --max-steps 800

Then full all-years:
  python v2/train/train.py --years 2005 ... 2024 --model Datadog/Toto-2.0-2.5B ...
"""

from __future__ import annotations

import argparse
import dataclasses
import json
import math
import sys
import time
from pathlib import Path

import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

UTILS = Path(__file__).resolve().parent.parent / "utils"
sys.path.insert(0, str(UTILS))

from tle_dataset import (  # noqa: E402
    TLEDatasetV2, series_collate_fn, N_CHANNELS, LOSS_CHANNEL_MASK, DRIFT_CHANNELS,
)
from toto2 import Toto2Model, Toto2ModelConfig  # noqa: E402


def build_model(model_id: str, init: str) -> Toto2Model:
    """Load Toto-2 for (A) continued pretraining from Toto weights, or
    (B) from-scratch pretraining (same architecture, random init)."""
    if init == "pretrained":
        return Toto2Model.from_pretrained(model_id)
    # scratch: fetch only the architecture config, randomly initialize weights
    from huggingface_hub import hf_hub_download
    raw = json.loads(Path(hf_hub_download(model_id, "config.json")).read_text())
    known = {f.name for f in dataclasses.fields(Toto2ModelConfig)}
    cfg = Toto2ModelConfig(**{k: v for k, v in raw.items() if k in known})
    return Toto2Model(cfg)


def quantile_pinball_loss(quantiles, target_scaled, valid, knots):
    pred = quantiles[..., :-1, :]          # (Q,B,C,S-1,P) position i predicts i+1
    tgt = target_scaled[..., 1:, :].unsqueeze(0)
    m = valid[..., 1:, :].unsqueeze(0)
    err = tgt - pred
    k = knots.view(-1, 1, 1, 1, 1)
    pin = torch.maximum(k * err, (k - 1.0) * err)
    m = m.expand_as(pin)
    return (pin * m).sum() / m.sum().clamp_min(1.0)


def compute_loss(model, batch, device, patch_size, chan_weight):
    target = batch["target"].to(device)            # (B,C,T)
    mask = batch["target_mask"].to(device)         # (B,C,T)
    series_ids = batch["series_ids"].to(device)
    cpm_mask = torch.ones_like(mask)
    B, C, T = target.shape
    S = T // patch_size

    with torch.no_grad():
        scaled, _, _ = model.scaler(target, mask & cpm_mask)
        target_scaled = scaled.asinh().view(B, C, S, patch_size)
    # per-channel loss WEIGHT (float): solar=0, orbital=1, drift channels upweighted.
    # Weighted pinball = sum(w*pin)/sum(w), so the denominator stays a proper mean.
    valid = (mask.float() * chan_weight.to(device).view(1, C, 1)).view(B, C, S, patch_size)

    out = model.forward(target, mask, cpm_mask, series_ids, num_return_steps=None)
    knots = torch.tensor(model.output_head.knots, device=device, dtype=torch.float32)
    return quantile_pinball_loss(out.quantiles.float(), target_scaled.float(), valid, knots)


def build_loader(args, split, shuffle, verbose):
    ds = TLEDatasetV2(
        input_dir=args.input_dir, cache_dir=args.cache_dir, cache_file=args.cache_file,
        sw_csv=args.sw_csv, years=args.years, patch_size=args.patch_size,
        window_patches=args.window_patches, stride_patches=args.stride_patches,
        split=split, clean=not args.no_clean, leo_only=not args.no_leo,
        split_mode=args.split_mode,
        train_until=args.train_until, valid_until=args.valid_until,
        max_satellites=args.max_satellites, verbose=verbose,
    )
    if len(ds) == 0:
        return None
    return DataLoader(ds, batch_size=args.batch_size, shuffle=shuffle,
                      num_workers=args.num_workers, collate_fn=series_collate_fn,
                      drop_last=shuffle)


def lr_at(step, base_lr, warmup, max_steps, min_lr, schedule):
    """Linear warmup, then constant or cosine decay to min_lr."""
    if warmup > 0 and step < warmup:
        return base_lr * (step + 1) / warmup
    if schedule == "cosine":
        prog = min(1.0, (step - warmup) / max(1, max_steps - warmup))
        return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + math.cos(math.pi * prog))
    return base_lr


@torch.no_grad()
def validate(model, loader, device, patch_size, chan_weight, max_batches):
    model.eval()
    tot, n = 0.0, 0
    for i, b in enumerate(loader):
        if i >= max_batches:
            break
        tot += float(compute_loss(model, b, device, patch_size, chan_weight)); n += 1
    model.train()
    return tot / max(1, n)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input-dir", default="/home/irteam/data-vol1/models/OrbitGPT/data/TLEs")
    ap.add_argument("--cache-dir", default="/home/irteam/data-vol1/models/OrbitGPT/v2/cache")
    ap.add_argument("--cache-file", default=None,
                    help="explicit prebuilt cache npz (e.g. the full 2005-2024 superset); "
                         "skips parsing, ignores --years/--no-clean/--sw-csv, filters by --split")
    ap.add_argument("--sw-csv", default="/home/irteam/data-vol1/models/OrbitGPT/v2/data/SW-All.csv")
    ap.add_argument("--years", type=int, nargs="+", default=[2020])
    ap.add_argument("--model", default="Datadog/Toto-2.0-4m")
    ap.add_argument("--init", default="pretrained", choices=["pretrained", "scratch"],
                    help="pretrained = continue-pretrain from Toto weights (recommended); "
                         "scratch = same architecture, random init")
    ap.add_argument("--no-clean", action="store_true")
    ap.add_argument("--no-leo", action="store_true",
                    help="disable cm-tle-pred LEO-only filter (affects cache key)")
    ap.add_argument("--split-mode", default="time", choices=["time", "satellite"],
                    help="time = epoch cutoffs (forecast-honest); satellite = cm-tle-pred 70/15/15")
    # default context = 8 patches (256 days) for pretraining; pass smaller for quick smokes
    ap.add_argument("--window-patches", type=int, default=8)
    ap.add_argument("--stride-patches", type=int, default=4)
    ap.add_argument("--train-until", default="2022-01-01")
    ap.add_argument("--valid-until", default="2023-01-01")
    ap.add_argument("--max-satellites", type=int, default=None)
    ap.add_argument("--batch-size", type=int, default=64)
    ap.add_argument("--num-workers", type=int, default=4)
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--min-lr", type=float, default=None, help="cosine floor (default lr/10)")
    ap.add_argument("--schedule", default="cosine", choices=["cosine", "constant"])
    ap.add_argument("--weight-decay", type=float, default=0.0)
    ap.add_argument("--drift-loss-weight", type=float, default=4.0,
                    help="loss weight multiplier for the position-critical drift channels "
                         "(d_bstar, d_mean_motion); 1.0 = uniform (old behavior)")
    ap.add_argument("--warmup", type=int, default=40)
    ap.add_argument("--max-steps", type=int, default=800)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--val-every", type=int, default=250,
                    help="run validation every N steps (0 = off). Needs a non-empty "
                         "valid split (train_until <= epoch < valid_until).")
    ap.add_argument("--val-batches", type=int, default=20)
    ap.add_argument("--amp", default="bf16", choices=["bf16", "fp32"])
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--out", default="/home/irteam/data-vol1/models/OrbitGPT/v2/ckpt")
    args = ap.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    torch.manual_seed(42)
    if args.min_lr is None:
        args.min_lr = args.lr / 10.0
    print(f"[model] {args.init} init of {args.model}")
    model = build_model(args.model, args.init).to(device)
    patch_size = model.config.patch_size
    args.patch_size = patch_size
    model.train()
    print(f"[model] patch_size={patch_size} params={sum(p.numel() for p in model.parameters())/1e6:.1f}M "
          f"| schedule={args.schedule} lr={args.lr}->{args.min_lr} window={args.window_patches}patch")

    # per-channel loss weights: orbital=1, solar=0, position-critical drift channels
    # (d_bstar, d_mean_motion) upweighted so the objective tracks SGP4 along-track/drag.
    chan_weight = torch.from_numpy(LOSS_CHANNEL_MASK.astype("float32")).clone()
    for c in DRIFT_CHANNELS:
        chan_weight[c] = chan_weight[c] * args.drift_loss_weight
    print(f"[loss] channel weights = {chan_weight.tolist()} (drift x{args.drift_loss_weight})")
    train_loader = build_loader(args, "train", True, True)
    if train_loader is None:
        print("[data] train split empty -> split=all")
        ds = TLEDatasetV2(input_dir=args.input_dir, cache_dir=args.cache_dir, cache_file=args.cache_file,
                          sw_csv=args.sw_csv, years=args.years, patch_size=patch_size,
                          window_patches=args.window_patches, stride_patches=args.stride_patches,
                          split="all", clean=not args.no_clean, leo_only=not args.no_leo,
                          split_mode=args.split_mode,
                          max_satellites=args.max_satellites, verbose=True)
        train_loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True,
                                  num_workers=args.num_workers, collate_fn=series_collate_fn, drop_last=True)
    val_loader = build_loader(args, "valid", False, False) if args.val_every else None
    if args.val_every and val_loader is None:
        print("[valid] no valid-split windows (e.g. single-year smoke) -> validation disabled")

    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95))
    use_amp = args.amp == "bf16" and device.type == "cuda"

    Path(args.out).mkdir(parents=True, exist_ok=True)
    ckpt = Path(args.out) / f"toto_v2_{Path(args.model).name}.pt"
    best_ckpt = Path(args.out) / f"toto_v2_{Path(args.model).name}_best.pt"
    step, t0 = 0, time.time()
    best_val, last_val = float("inf"), None
    pbar = tqdm(total=args.max_steps, desc="train", unit="step")
    while step < args.max_steps:
        for batch in train_loader:
            if step >= args.max_steps:
                break
            for g in opt.param_groups:
                g["lr"] = lr_at(step, args.lr, args.warmup, args.max_steps, args.min_lr, args.schedule)
            opt.zero_grad(set_to_none=True)
            if use_amp:
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    loss = compute_loss(model, batch, device, patch_size, chan_weight)
            else:
                loss = compute_loss(model, batch, device, patch_size, chan_weight)
            loss.backward()
            gn = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            opt.step()
            pbar.update(1)
            post = {"loss": f"{float(loss):.4f}", "gnorm": f"{float(gn):.2f}",
                    "lr": f"{opt.param_groups[0]['lr']:.1e}"}
            if last_val is not None:
                post["val"] = f"{last_val:.4f}"
            pbar.set_postfix(**post)

            if val_loader is not None and step > 0 and step % args.val_every == 0:
                last_val = validate(model, val_loader, device, patch_size, chan_weight, args.val_batches)
                improved = last_val < best_val
                if improved:
                    best_val = last_val
                    torch.save({"model": model.state_dict(), "config": vars(args),
                                "step": step, "val_loss": last_val}, best_ckpt)
                pbar.write(f"[valid] step {step:6d}  val_loss {last_val:.5f}  "
                           f"train_loss {float(loss):.5f}{'  (best, saved)' if improved else ''}")
            step += 1
    pbar.close()

    if val_loader is not None:
        last_val = validate(model, val_loader, device, patch_size, chan_weight, args.val_batches)
        print(f"[valid] final val_loss {last_val:.5f} (best {best_val:.5f} -> {best_ckpt.name})")

    torch.save({"model": model.state_dict(), "config": vars(args)}, ckpt)
    print(f"[done] {step} steps in {time.time()-t0:.1f}s -> {ckpt}")


if __name__ == "__main__":
    main()