From 5a0e98c5d1940da0b0025ff00522836098471780 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 17 Jul 2026 11:02:12 +0200 Subject: [PATCH] Add EMA weights, weight decay, step-based LR schedule, and grad-norm logging to training 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 --- giant/cli.py | 85 +++++++++++++++++++++++++++++++++++++---- giant/config.py | 5 +++ giant/pipeline.py | 3 ++ giant/train.py | 96 ++++++++++++++++++++++++++++++++++++++++------- 4 files changed, 168 insertions(+), 21 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index 82fe028..935bc98 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -183,6 +183,39 @@ class Coord(str, Enum): local = "local" +class Weights(str, Enum): + raw = "raw" + ema = "ema" + + +def _load_model_weights( + model: torch.nn.Module, + sec_decoder: torch.nn.Module, + ckpt: dict, + weights: "Weights", + checkpoint_path: Path, +) -> None: + """Load either the raw or EMA state dicts from a training checkpoint. + + EMA weights (giant.train's shadow copy, see --ema-decay) only exist in + checkpoints written after that feature landed, so `ema` fails loudly + rather than silently falling back to raw weights a caller didn't ask for. + """ + if weights == Weights.raw: + model_key, sec_key = "model", "sec_decoder" + else: + model_key, sec_key = "model_ema", "sec_decoder_ema" + if model_key not in ckpt or sec_key not in ckpt: + typer.echo( + f"error: {checkpoint_path} has no EMA weights (trained before " + "--ema-decay, or with --ema-decay 0) — use --weights raw", + err=True, + ) + raise typer.Exit(1) + model.load_state_dict(ckpt[model_key]) + sec_decoder.load_state_dict(ckpt[sec_key]) + + @app.command() def train( data: Annotated[ @@ -209,6 +242,18 @@ def train( ), ] = None, lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None, + weight_decay: Annotated[ + Optional[float], + typer.Option("--weight-decay", "-W", help="AdamW weight decay (default: 0.01)"), + ] = None, + ema_decay: Annotated[ + Optional[float], + typer.Option( + "--ema-decay", + help="EMA decay for a shadow copy of the model weights, saved " + "alongside the raw weights in checkpoints (0 disables; default: 0.9999)", + ), + ] = None, warmup_epochs: Annotated[ Optional[int], typer.Option("--warmup-epochs", "-w") ] = None, @@ -272,6 +317,14 @@ def train( "(ignored in ddpm mode, which always runs the full schedule)", ), ] = None, + max_val_batches: Annotated[ + Optional[int], + typer.Option( + "--max-val-batches", + help="Cap the per-epoch val-loss pass to N batches (0 = full " + "val set every epoch; default: 200)", + ), + ] = None, shuffle_buffer: Annotated[ int, typer.Option( @@ -318,12 +371,15 @@ def train( "epochs": epochs, "batch_size": batch_size_value, "lr": lr, + "weight_decay": weight_decay, + "ema_decay": ema_decay, "warmup_epochs": warmup_epochs, "val_fraction": val_fraction, "num_workers": num_workers, "seed": seed, "validate_every": validate_every, "validate_steps": validate_steps, + "max_val_batches": max_val_batches, }.items() if v is not None } @@ -439,6 +495,15 @@ def predict( steps: Annotated[ int, typer.Option("--steps", "-s", help="Flow matching ODE steps") ] = 10, + weights: Annotated[ + Weights, + typer.Option( + "--weights", + help="raw: the live training weights. ema: the EMA shadow copy " + "(see --ema-decay in `giant train`) — usually cleaner samples, " + "requires a checkpoint trained with EMA enabled.", + ), + ] = Weights.raw, device: Annotated[ Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"), @@ -515,13 +580,11 @@ def predict( tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) model, sec_decoder = build_models(model_cfg) - model.load_state_dict(ckpt["model"]) + _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) model.to(_device).eval() - - sec_decoder.load_state_dict(ckpt["sec_decoder"]) sec_decoder.to(_device).eval() - typer.echo(f"loaded checkpoint: {checkpoint}") + typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})") gconfig.warn_if_checkpoint_config_mismatch(checkpoint) # --- Output path --- @@ -810,6 +873,15 @@ def rollout( int, typer.Option("--steps", "-s", help="Flow matching ODE steps per model call"), ] = 10, + weights: Annotated[ + Weights, + typer.Option( + "--weights", + help="raw: the live training weights. ema: the EMA shadow copy " + "(see --ema-decay in `giant train`) — usually cleaner samples, " + "requires a checkpoint trained with EMA enabled.", + ), + ] = Weights.raw, batch_size: Annotated[ int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward") ] = 4096, @@ -865,11 +937,10 @@ def rollout( tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) model, sec_decoder = build_models(model_cfg) - model.load_state_dict(ckpt["model"]) + _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) model.to(_device).eval() - sec_decoder.load_state_dict(ckpt["sec_decoder"]) sec_decoder.to(_device).eval() - typer.echo(f"loaded checkpoint: {checkpoint}") + typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})") oracle = GeometryOracle.load(geometry) typer.echo( diff --git a/giant/config.py b/giant/config.py index 8ce67f5..4834f3c 100644 --- a/giant/config.py +++ b/giant/config.py @@ -14,6 +14,11 @@ DEFAULT_CONFIG: dict = { "epochs": 100, "batch_size": 4096, "lr": 3e-4, + "weight_decay": 0.01, # AdamW default — exposed so it can be tuned + "ema_decay": 0.9999, # EMA of model weights for sampling; 0 disables + # per-epoch val loss (not the marginal/KL validate_every pass) is + # capped to this many batches; 0 = full val set every epoch + "max_val_batches": 200, "val_fraction": 0.1, "num_workers": 4, "seed": 0, diff --git a/giant/pipeline.py b/giant/pipeline.py index f160231..d6522c7 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -174,6 +174,8 @@ def run_train_job( mode=t["mode"], epochs=t["epochs"], lr=t["lr"], + weight_decay=t["weight_decay"], + ema_decay=t["ema_decay"], warmup_epochs=t["warmup_epochs"], device=device, out_dir=out_dir, @@ -189,5 +191,6 @@ def run_train_job( resume_path=resume, validate_every=t["validate_every"], validate_steps=t["validate_steps"], + max_val_batches=t["max_val_batches"], total_train_batches=total_train_batches, ) diff --git a/giant/train.py b/giant/train.py index 35f3d8c..d796956 100644 --- a/giant/train.py +++ b/giant/train.py @@ -1,3 +1,4 @@ +import copy import csv import math import os @@ -35,6 +36,7 @@ _METRICS_FIELDS = [ "val_loss_balance", "val_loss_proc", "lr", + "grad_norm", "epoch_time_s", ] @@ -103,6 +105,14 @@ def _build_sec_x1( 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, @@ -195,6 +205,8 @@ def train( 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, @@ -207,6 +219,7 @@ def train( 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) @@ -215,15 +228,38 @@ def train( stage1_model = stage1_model.to(device) sec_decoder = sec_decoder.to(device) - all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters()) - optimizer = optim.AdamW(all_params, lr=lr) + # 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) - def _lr_lambda(epoch: int) -> float: - if warmup_epochs > 0 and epoch < warmup_epochs: - return (epoch + 1) / warmup_epochs - t = epoch - warmup_epochs - T = max(epochs - warmup_epochs, 1) - return 0.5 * (1.0 + math.cos(math.pi * t / T)) + 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) @@ -235,6 +271,12 @@ def train( 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 @@ -271,7 +313,6 @@ def train( with _GracefulShutdown() as shutdown: for epoch in range(start_epoch, epochs + 1): epoch_start = time.monotonic() - current_lr = optimizer.param_groups[0]["lr"] stage1_model.train() sec_decoder.train() train_loss_sum = 0.0 @@ -281,7 +322,10 @@ def train( 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}", @@ -305,11 +349,17 @@ def train( ) optimizer.zero_grad() loss.backward() - torch.nn.utils.clip_grad_norm_(all_params, 1.0) + 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 @@ -317,10 +367,19 @@ def train( 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 ) - bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False) + 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 @@ -330,7 +389,8 @@ def train( break train_loss = train_loss_sum / max(train_n, 1) - lr_sched.step() + train_grad_norm = grad_norm_sum / max(train_batches, 1) + current_lr = optimizer.param_groups[0]["lr"] stage1_model.eval() sec_decoder.eval() @@ -342,7 +402,9 @@ def train( val_proc_sum = 0.0 val_n = 0 with torch.no_grad(): - for batch in val_loader: + 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, @@ -377,7 +439,8 @@ def train( 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} {epoch_time:.1f}s{marker}" + f" lr {current_lr:.2e} gnorm {train_grad_norm:.3f}" + f" {epoch_time:.1f}s{marker}" ) metrics_writer.writerow( { @@ -395,6 +458,7 @@ def train( "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, } ) @@ -420,6 +484,10 @@ def train( "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: