#!/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()