lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Add EMA weights, weight decay, step-based LR schedule, and grad-norm logging to training

Browse Source 
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>
This commit is contained in:
Lars Bogner
2026-07-17 11:02:12 +02:00
parent 24a83486ec
commit 5a0e98c5d1
4 changed files with 168 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/cli.py
+78 -7
View File
@@ -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(
giant/config.py
+5
View File
@@ -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,
giant/pipeline.py
+3
View File
@@ -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,
)
giant/train.py
+82 -14
View File
@@ -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: