lars
5a0e98c5d1
Gives flow-matching sampling a cleaner EMA shadow copy to draw from (--ema-decay, --weights raw|ema in predict/rollout), fixes the LR warmup/cosine schedule stepping once per epoch even when an epoch is tens of thousands of steps, and caps the per-epoch val-loss pass (--max-val-batches) so large val sets don't dominate epoch time. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
520 lines
19 KiB
Python
520 lines
19 KiB
Python
import copy
|
|
import csv
|
|
import math
|
|
import os
|
|
import signal
|
|
import time
|
|
from pathlib import Path
|
|
from types import FrameType
|
|
from typing import Callable
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader
|
|
from tqdm import tqdm
|
|
|
|
from giant.model.schedule import (
|
|
CosineSchedule,
|
|
flow_matching_loss,
|
|
flow_matching_loss_secondary,
|
|
)
|
|
from giant.validate import validate_marginals
|
|
|
|
_METRICS_FIELDS = [
|
|
"epoch",
|
|
"train_loss",
|
|
"train_loss_s1",
|
|
"train_loss_nsec",
|
|
"train_loss_s2",
|
|
"train_loss_balance",
|
|
"train_loss_proc",
|
|
"val_loss",
|
|
"val_loss_s1",
|
|
"val_loss_nsec",
|
|
"val_loss_s2",
|
|
"val_loss_balance",
|
|
"val_loss_proc",
|
|
"lr",
|
|
"grad_norm",
|
|
"epoch_time_s",
|
|
]
|
|
|
|
_CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)
|
|
|
|
|
|
class _GracefulShutdown:
|
|
"""Turns SIGINT/SIGTERM into a flag check instead of an immediate crash.
|
|
|
|
A second signal while already shutting down restores the default
|
|
handler and re-sends the signal, so an unresponsive run can still be
|
|
force-killed.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.requested = False
|
|
self._previous: dict[
|
|
int,
|
|
Callable[[int, FrameType | None], object] | signal.Handlers | int | None,
|
|
] = {}
|
|
|
|
def __enter__(self) -> "_GracefulShutdown":
|
|
for sig in _CATCHABLE_SIGNALS:
|
|
self._previous[sig] = signal.getsignal(sig)
|
|
signal.signal(sig, self._handle)
|
|
return self
|
|
|
|
def __exit__(self, *exc_info) -> None:
|
|
for sig, handler in self._previous.items():
|
|
signal.signal(sig, handler)
|
|
|
|
def _handle(self, signum: int, frame) -> None:
|
|
if self.requested:
|
|
signal.signal(signum, self._previous[signum])
|
|
os.kill(os.getpid(), signum)
|
|
return
|
|
self.requested = True
|
|
print(
|
|
f"\nreceived {signal.Signals(signum).name} — finishing the current "
|
|
"batch, then saving a checkpoint and exiting (send again to force-quit)"
|
|
)
|
|
|
|
|
|
def _build_sec_x1(
|
|
sec_cont: torch.Tensor,
|
|
sec_pdg_idx: torch.Tensor,
|
|
pdg_emb_weight: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Assemble the Stage-2 flow target by appending type embeddings.
|
|
|
|
sec_cont: (B, K_MAX, 4) — [stick_logit, dir_local]
|
|
sec_pdg_idx: (B, K_MAX) — integer PDG model-indices
|
|
pdg_emb_weight: (pdg_vocab, emb_dim) — live embedding table weights
|
|
|
|
Returns (B, SEC_DIM) = (B, K_MAX * (4 + emb_dim)).
|
|
|
|
Detaches the looked-up rows: this tensor becomes x1 in the flow-matching
|
|
loss (u_t = x1 - x0), so without detaching, the Stage-2 loss could pull
|
|
the embedding table itself toward whatever the decoder already predicts
|
|
(a moving, self-referential regression target) instead of only pulling
|
|
the decoder toward the table. The table is still trained normally via
|
|
its Stage-1 conditioning role and `predict_n_sec`.
|
|
"""
|
|
type_emb = pdg_emb_weight[sec_pdg_idx].detach() # (B, K_MAX, emb_dim)
|
|
x1_s2 = torch.cat([sec_cont, type_emb], dim=-1) # (B, K_MAX, 4+emb_dim)
|
|
return x1_s2.flatten(1) # (B, SEC_DIM)
|
|
|
|
|
|
@torch.no_grad()
|
|
def _update_ema(
|
|
ema_model: torch.nn.Module, model: torch.nn.Module, decay: float
|
|
) -> None:
|
|
for ema_p, p in zip(ema_model.parameters(), model.parameters()):
|
|
ema_p.mul_(decay).add_(p, alpha=1 - decay)
|
|
|
|
|
|
def _compute_losses(
|
|
stage1_model: torch.nn.Module,
|
|
sec_decoder: torch.nn.Module,
|
|
batch: tuple,
|
|
mode: str,
|
|
ddpm_schedule,
|
|
device: torch.device,
|
|
lambda_nsec: float,
|
|
lambda_s2: float,
|
|
lambda_balance: float = 0.0,
|
|
lambda_proc: float = 0.0,
|
|
) -> tuple[
|
|
torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
|
|
]:
|
|
"""Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc) for one batch."""
|
|
cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx = batch
|
|
cond_cont = cond_cont.to(device)
|
|
cond_cat = cond_cat.to(device)
|
|
x1_s1 = x1_s1.to(device)
|
|
n_sec = n_sec.to(device)
|
|
sec_cont = sec_cont.to(device)
|
|
sec_pdg_idx = sec_pdg_idx.to(device)
|
|
proc_idx = proc_idx.to(device)
|
|
|
|
# Stage-1 flow loss
|
|
if mode == "flow":
|
|
l_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat)
|
|
else:
|
|
assert ddpm_schedule is not None
|
|
l_s1 = ddpm_schedule.loss(stage1_model, x1_s1, cond_cont, cond_cat)
|
|
|
|
# n_sec classification loss
|
|
n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat)
|
|
l_nsec = F.cross_entropy(n_sec_logits, n_sec)
|
|
|
|
# Stage-2 secondary flow loss
|
|
# Use a noiseless Stage-1 target as context (detach to avoid back-prop
|
|
# coupling between the two flow paths through the same embedding table).
|
|
# The type-embedding lookup itself is also detached inside _build_sec_x1,
|
|
# so the shared PDG table is shaped only by its Stage-1 conditioning role
|
|
# and predict_n_sec, not by chasing the Stage-2 decoder's predictions.
|
|
from giant.constants import K_MAX
|
|
|
|
pdg_emb_weight = stage1_model.pdg_embedding_weight()
|
|
x1_s2 = _build_sec_x1(sec_cont, sec_pdg_idx, pdg_emb_weight)
|
|
|
|
sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
|
|
l_s2 = flow_matching_loss_secondary(
|
|
sec_decoder,
|
|
x1_s2,
|
|
cond_cont,
|
|
cond_cat,
|
|
x1_s1.detach(),
|
|
sec_mask,
|
|
)
|
|
|
|
# Optional MoE load-balance auxiliary loss: only present when both stages
|
|
# are routed (RoutedDenoisingMLP/RoutedSecondaryDecoder carry `.router`,
|
|
# the monolith models don't), computed on cond_cont alone (cheap — no
|
|
# trunk compute) so it's reported even when lambda_balance == 0.
|
|
if hasattr(stage1_model, "router") and hasattr(sec_decoder, "router"):
|
|
l_balance = stage1_model.router.balance_loss(
|
|
cond_cont, cond_cat
|
|
) + sec_decoder.router.balance_loss(cond_cont, cond_cat)
|
|
# Supervised router auxiliary loss (e.g. ProcessRouter's process
|
|
# classifier); a scalar 0 for routers with no such loss (EnergyRouter).
|
|
l_proc = stage1_model.router.classify_loss(
|
|
cond_cont, cond_cat, proc_idx
|
|
) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx)
|
|
else:
|
|
l_balance = torch.zeros((), device=device)
|
|
l_proc = torch.zeros((), device=device)
|
|
|
|
total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
|
|
if lambda_balance > 0:
|
|
total = total + lambda_balance * l_balance
|
|
if lambda_proc > 0:
|
|
total = total + lambda_proc * l_proc
|
|
return total, l_s1, l_nsec, l_s2, l_balance, l_proc
|
|
|
|
|
|
def train(
|
|
stage1_model: torch.nn.Module,
|
|
sec_decoder: torch.nn.Module,
|
|
train_loader: DataLoader,
|
|
val_loader: DataLoader,
|
|
mode: str,
|
|
epochs: int,
|
|
lr: float,
|
|
warmup_epochs: int,
|
|
device: torch.device,
|
|
out_dir: str | Path,
|
|
weight_decay: float = 0.01,
|
|
ema_decay: float = 0.9999,
|
|
lambda_nsec: float = 0.1,
|
|
lambda_s2: float = 1.0,
|
|
lambda_balance: float = 0.0,
|
|
lambda_proc: float = 0.0,
|
|
normalizer_dict: dict | None = None,
|
|
pdg_map: dict | None = None,
|
|
mat_map: dict | None = None,
|
|
proc_map: dict | None = None,
|
|
model_config: dict | None = None,
|
|
resume_path: str | Path | None = None,
|
|
validate_every: int = 0,
|
|
validate_steps: int = 10,
|
|
max_val_batches: int = 0,
|
|
total_train_batches: int = 0,
|
|
) -> None:
|
|
out_dir = Path(out_dir)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
stage1_model = stage1_model.to(device)
|
|
sec_decoder = sec_decoder.to(device)
|
|
|
|
# Flow-matching/diffusion models sample noticeably better from an EMA of
|
|
# the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed
|
|
# sinusoidal-embedding freqs, or non-learned router centers) never change
|
|
# after this initial copy, so only parameters need the running average.
|
|
ema_stage1_model: torch.nn.Module | None = None
|
|
ema_sec_decoder: torch.nn.Module | None = None
|
|
if ema_decay > 0:
|
|
ema_stage1_model = copy.deepcopy(stage1_model).eval()
|
|
ema_sec_decoder = copy.deepcopy(sec_decoder).eval()
|
|
for p in ema_stage1_model.parameters():
|
|
p.requires_grad_(False)
|
|
for p in ema_sec_decoder.parameters():
|
|
p.requires_grad_(False)
|
|
|
|
all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters())
|
|
optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay)
|
|
|
|
# Warmup/decay in units of optimizer steps rather than epochs: at large
|
|
# dataset sizes a single epoch can be tens of thousands of steps, and an
|
|
# epoch-granularity schedule would leave warmup/cosine decay unable to
|
|
# move within it. Requires an accurate `total_train_batches` (steps per
|
|
# epoch); the only caller, run_train_job, always supplies one.
|
|
steps_per_epoch = max(total_train_batches, 1)
|
|
warmup_steps = warmup_epochs * steps_per_epoch
|
|
total_steps = max(epochs * steps_per_epoch, 1)
|
|
|
|
def _lr_lambda(step: int) -> float:
|
|
if warmup_steps > 0 and step < warmup_steps:
|
|
return (step + 1) / warmup_steps
|
|
t = step - warmup_steps
|
|
T = max(total_steps - warmup_steps, 1)
|
|
return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T))
|
|
|
|
lr_sched = optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda)
|
|
|
|
ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None
|
|
|
|
start_epoch = 1
|
|
best_val_loss = float("inf")
|
|
if resume_path is not None:
|
|
ckpt = torch.load(resume_path, map_location=device, weights_only=False)
|
|
stage1_model.load_state_dict(ckpt["model"])
|
|
sec_decoder.load_state_dict(ckpt["sec_decoder"])
|
|
if ema_decay > 0:
|
|
assert ema_stage1_model is not None and ema_sec_decoder is not None
|
|
ema_stage1_model.load_state_dict(ckpt.get("model_ema", ckpt["model"]))
|
|
ema_sec_decoder.load_state_dict(
|
|
ckpt.get("sec_decoder_ema", ckpt["sec_decoder"])
|
|
)
|
|
optimizer.load_state_dict(ckpt["optimizer"])
|
|
lr_sched.load_state_dict(ckpt["lr_sched"])
|
|
start_epoch = ckpt.get("epoch", 0) + 1
|
|
best_val_loss = ckpt.get("best_val_loss", float("inf"))
|
|
|
|
# optimizer/lr_sched.load_state_dict() above restore the checkpoint's
|
|
# own base LR, which would otherwise silently override an explicit
|
|
# `lr` argument. Make `lr` authoritative again, applied at whatever
|
|
# point the cosine/warmup schedule has already reached.
|
|
lr_sched.base_lrs = [lr for _ in lr_sched.base_lrs]
|
|
resumed_lr = lr * _lr_lambda(lr_sched.last_epoch)
|
|
for group in optimizer.param_groups:
|
|
group["lr"] = resumed_lr
|
|
|
|
if start_epoch > epochs:
|
|
print(
|
|
f"checkpoint already completed epoch {start_epoch - 1} "
|
|
f"(>= --epochs {epochs}) — nothing to train"
|
|
)
|
|
return
|
|
|
|
metrics_path = out_dir / "metrics.csv"
|
|
resuming_existing_metrics = resume_path is not None and metrics_path.exists()
|
|
write_header = not resuming_existing_metrics
|
|
metrics_file = open(
|
|
metrics_path, "a" if resuming_existing_metrics else "w", newline=""
|
|
)
|
|
metrics_writer = csv.DictWriter(metrics_file, fieldnames=_METRICS_FIELDS)
|
|
if write_header:
|
|
metrics_writer.writeheader()
|
|
|
|
epoch_w = len(str(epochs))
|
|
last_completed_epoch = start_epoch - 1
|
|
with _GracefulShutdown() as shutdown:
|
|
for epoch in range(start_epoch, epochs + 1):
|
|
epoch_start = time.monotonic()
|
|
stage1_model.train()
|
|
sec_decoder.train()
|
|
train_loss_sum = 0.0
|
|
train_s1_sum = 0.0
|
|
train_nsec_sum = 0.0
|
|
train_s2_sum = 0.0
|
|
train_balance_sum = 0.0
|
|
train_proc_sum = 0.0
|
|
train_n = 0
|
|
train_batches = 0
|
|
grad_norm_sum = 0.0
|
|
ema_loss = 0.0
|
|
ema_grad_norm = 0.0
|
|
bar = tqdm(
|
|
train_loader,
|
|
desc=f" epoch {epoch:{epoch_w}d}/{epochs}",
|
|
total=total_train_batches or None,
|
|
leave=False,
|
|
unit="batch",
|
|
dynamic_ncols=True,
|
|
)
|
|
for batch in bar:
|
|
loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
|
|
stage1_model,
|
|
sec_decoder,
|
|
batch,
|
|
mode,
|
|
ddpm_schedule,
|
|
device,
|
|
lambda_nsec,
|
|
lambda_s2,
|
|
lambda_balance,
|
|
lambda_proc,
|
|
)
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0)
|
|
optimizer.step()
|
|
lr_sched.step()
|
|
if ema_decay > 0:
|
|
assert ema_stage1_model is not None and ema_sec_decoder is not None
|
|
_update_ema(ema_stage1_model, stage1_model, ema_decay)
|
|
_update_ema(ema_sec_decoder, sec_decoder, ema_decay)
|
|
|
|
B = batch[0].size(0)
|
|
batch_loss = loss.item()
|
|
batch_grad_norm = grad_norm.item()
|
|
train_loss_sum += batch_loss * B
|
|
train_s1_sum += l_s1.item() * B
|
|
train_nsec_sum += l_nsec.item() * B
|
|
train_s2_sum += l_s2.item() * B
|
|
train_balance_sum += l_balance.item() * B
|
|
train_proc_sum += l_proc.item() * B
|
|
train_n += B
|
|
train_batches += 1
|
|
grad_norm_sum += batch_grad_norm
|
|
ema_loss = (
|
|
batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss
|
|
)
|
|
ema_grad_norm = (
|
|
batch_grad_norm
|
|
if train_batches == 1
|
|
else 0.95 * ema_grad_norm + 0.05 * batch_grad_norm
|
|
)
|
|
bar.set_postfix_str(
|
|
f"loss={ema_loss:.4f} gnorm={ema_grad_norm:.3f}", refresh=False
|
|
)
|
|
|
|
if shutdown.requested:
|
|
break
|
|
bar.close()
|
|
|
|
if shutdown.requested:
|
|
break
|
|
|
|
train_loss = train_loss_sum / max(train_n, 1)
|
|
train_grad_norm = grad_norm_sum / max(train_batches, 1)
|
|
current_lr = optimizer.param_groups[0]["lr"]
|
|
|
|
stage1_model.eval()
|
|
sec_decoder.eval()
|
|
val_loss_sum = 0.0
|
|
val_s1_sum = 0.0
|
|
val_nsec_sum = 0.0
|
|
val_s2_sum = 0.0
|
|
val_balance_sum = 0.0
|
|
val_proc_sum = 0.0
|
|
val_n = 0
|
|
with torch.no_grad():
|
|
for val_batch_idx, batch in enumerate(val_loader):
|
|
if max_val_batches > 0 and val_batch_idx >= max_val_batches:
|
|
break
|
|
loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
|
|
stage1_model,
|
|
sec_decoder,
|
|
batch,
|
|
mode,
|
|
ddpm_schedule,
|
|
device,
|
|
lambda_nsec,
|
|
lambda_s2,
|
|
lambda_balance,
|
|
lambda_proc,
|
|
)
|
|
B = batch[0].size(0)
|
|
val_loss_sum += loss.item() * B
|
|
val_s1_sum += l_s1.item() * B
|
|
val_nsec_sum += l_nsec.item() * B
|
|
val_s2_sum += l_s2.item() * B
|
|
val_balance_sum += l_balance.item() * B
|
|
val_proc_sum += l_proc.item() * B
|
|
val_n += B
|
|
val_loss = val_loss_sum / max(val_n, 1)
|
|
epoch_time = time.monotonic() - epoch_start
|
|
|
|
is_best = val_loss < best_val_loss
|
|
marker = " [best]" if is_best else ""
|
|
print(
|
|
f"epoch {epoch:{epoch_w}d}/{epochs}"
|
|
f" train {train_loss:.4f}"
|
|
f" (s1={train_s1_sum / max(train_n, 1):.3f}"
|
|
f" nsec={train_nsec_sum / max(train_n, 1):.3f}"
|
|
f" s2={train_s2_sum / max(train_n, 1):.3f}"
|
|
f" bal={train_balance_sum / max(train_n, 1):.3f}"
|
|
f" proc={train_proc_sum / max(train_n, 1):.3f})"
|
|
f" val {val_loss:.4f}"
|
|
f" lr {current_lr:.2e} gnorm {train_grad_norm:.3f}"
|
|
f" {epoch_time:.1f}s{marker}"
|
|
)
|
|
metrics_writer.writerow(
|
|
{
|
|
"epoch": epoch,
|
|
"train_loss": train_loss,
|
|
"train_loss_s1": train_s1_sum / max(train_n, 1),
|
|
"train_loss_nsec": train_nsec_sum / max(train_n, 1),
|
|
"train_loss_s2": train_s2_sum / max(train_n, 1),
|
|
"train_loss_balance": train_balance_sum / max(train_n, 1),
|
|
"train_loss_proc": train_proc_sum / max(train_n, 1),
|
|
"val_loss": val_loss,
|
|
"val_loss_s1": val_s1_sum / max(val_n, 1),
|
|
"val_loss_nsec": val_nsec_sum / max(val_n, 1),
|
|
"val_loss_s2": val_s2_sum / max(val_n, 1),
|
|
"val_loss_balance": val_balance_sum / max(val_n, 1),
|
|
"val_loss_proc": val_proc_sum / max(val_n, 1),
|
|
"lr": current_lr,
|
|
"grad_norm": train_grad_norm,
|
|
"epoch_time_s": epoch_time,
|
|
}
|
|
)
|
|
metrics_file.flush()
|
|
|
|
if validate_every > 0 and epoch % validate_every == 0:
|
|
print(f"[epoch {epoch}] marginal validation:")
|
|
validate_marginals(
|
|
stage1_model,
|
|
val_loader,
|
|
mode=mode,
|
|
schedule=ddpm_schedule,
|
|
device=device,
|
|
steps=validate_steps,
|
|
sec_decoder=sec_decoder,
|
|
)
|
|
|
|
ckpt: dict = {
|
|
"model": stage1_model.state_dict(),
|
|
"sec_decoder": sec_decoder.state_dict(),
|
|
"optimizer": optimizer.state_dict(),
|
|
"lr_sched": lr_sched.state_dict(),
|
|
"epoch": epoch,
|
|
"best_val_loss": best_val_loss,
|
|
}
|
|
if ema_decay > 0:
|
|
assert ema_stage1_model is not None and ema_sec_decoder is not None
|
|
ckpt["model_ema"] = ema_stage1_model.state_dict()
|
|
ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict()
|
|
if normalizer_dict is not None:
|
|
ckpt["normalizer"] = normalizer_dict
|
|
if pdg_map is not None:
|
|
ckpt["pdg_map"] = pdg_map
|
|
if mat_map is not None:
|
|
ckpt["mat_map"] = mat_map
|
|
if proc_map is not None:
|
|
ckpt["proc_map"] = proc_map
|
|
if model_config is not None:
|
|
ckpt["model_config"] = model_config
|
|
|
|
if val_loss < best_val_loss:
|
|
best_val_loss = val_loss
|
|
ckpt["best_val_loss"] = best_val_loss
|
|
torch.save(ckpt, out_dir / "best.pt")
|
|
|
|
torch.save(ckpt, out_dir / "last.pt")
|
|
last_completed_epoch = epoch
|
|
|
|
if shutdown.requested:
|
|
break
|
|
|
|
metrics_file.close()
|
|
|
|
if shutdown.requested:
|
|
print(
|
|
f"stopped after epoch {last_completed_epoch} due to shutdown signal — "
|
|
f"resume with --resume {out_dir / 'last.pt'}"
|
|
)
|