lars
78978769f6
CI / Lint (ruff check) (push) Successful in 28s
CI / Format (ruff format) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 44s
CI / Lint (ruff check) (pull_request) Successful in 41s
CI / Format (ruff format) (pull_request) Successful in 31s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 32s
CI / Tests (pull_request) Successful in 4m36s
CI / Tests (push) Successful in 4m49s
giant/ had no autocast/GradScaler/torch.compile anywhere despite the
project's ~10x-native-Geant4 eval-budget target. This adds bf16 mixed
precision to the training step (both FlowDDPMStageTrainer and
WGANStageTrainer) via a new train.precision config key ("fp32" default,
"bf16" opt-in) and giant.training.amp.resolve_autocast.
torch.compile is a separate, much larger surface (data-dependent routed
dispatch, the autoregressive sampler's per-token control flow, arbitrary
rollout batch sizes) and is left for a follow-up issue, per discussion.
Scope decisions made during planning:
- fp32 + bf16 only, no fp16/GradScaler. fp16 breaks two things in this
codebase: routers.py's three 1e-8 epsilons sit below fp16's ~6e-8
subnormal floor, and gradient_penalty's grad norm overflows fp16's
range at ordinary early-WGAN-GP gradient magnitudes. Every training
GPU in the fleet (A100/L40S/H200/RTX 4070) has native bf16; only
pre-Ampere V100s would need fp16.
- resolve_autocast raises loudly if bf16 is requested on hardware that
can't do it, rather than silently falling back to fp32.
- Autocast wraps the training step only; val_loss (and the
best-checkpoint selection it drives) stays fp32 so it's comparable
across every run recorded so far.
- _route_forward's mixture accumulator (giant/model/trunks.py) was a
hard-fp32 torch.zeros with no dtype, so under autocast a RoutedTrunk
silently returned a different output dtype than an unrouted
ExpertTrunk purely because router.enabled was set. Fixed to match the
experts' own dtype; the gate weights (forced fp32 for their own
numerical stability) are cast down before combining, so the
mixture's numerics stay solid without reintroducing the dtype split.
- Added explicit fp32 guards (autocast(enabled=False)) around spots
that are correct in fp32 but degrade quietly rather than crash in
bf16: the router's balance/entropy losses and gate softmax, the
stage-2 stick-breaking cumprod, and gradient_penalty's
double-backward + grad norm.
Benchmarked on the local RTX 4070 against configs/baseline.toml's
hyperparams (hidden_dim 512/6 blocks, bs 4096) on a synthetic dataset:
bf16 gave 1.05-1.35x training throughput and 18-33% lower peak GPU
memory across one-shot/routed/autoregressive stage-2 configs, with the
autoregressive path (the dominant cost per baseline.toml) benefiting
most on both axes.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
66 lines
2.7 KiB
Python
66 lines
2.7 KiB
Python
from typing import Callable
|
|
|
|
import torch
|
|
|
|
CriticFn = Callable[[torch.Tensor], torch.Tensor]
|
|
|
|
|
|
def gradient_penalty(
|
|
critic_fn: CriticFn,
|
|
real: torch.Tensor,
|
|
fake: torch.Tensor,
|
|
mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""WGAN-GP penalty (Gulrajani et al. 2017): (||grad||_2 - 1)^2 at a random interpolate.
|
|
|
|
`mask` (same shape as `real`/`fake`, 1 for real content / 0 for padding)
|
|
is for Stage 2's variable-length slot vector: both the interpolate and the
|
|
critic's gradient are zeroed on padded dims first, so the norm target of 1
|
|
is only ever asked of genuine content, not the padding convention shared
|
|
by both `real` and `fake`. Rows fully masked out (e.g. `n_sec == 0`, so
|
|
every slot is padding) have no real content to constrain the gradient
|
|
norm to 1 — `x_hat`/`grad` are forced to all-zero for such a row, which
|
|
would otherwise contribute a constant `(||0|| - 1)^2 == 1` bias to the
|
|
mean regardless of critic behavior — so they're excluded from the mean.
|
|
|
|
Deliberately kept fp32 (`torch.autocast(..., enabled=False)`) regardless
|
|
of the caller's ambient `train.precision` autocast region: this is a
|
|
`create_graph=True` double-backward, and `grad.norm(2, dim=1)` sums
|
|
squares over the critic's full input width (hundreds of dims for stage
|
|
2), which overflows bf16's range at gradient magnitudes well within
|
|
normal early-WGAN-GP territory. Disclosed cost: the critic forward
|
|
inside this function always runs fp32, even when the rest of the WGAN
|
|
stage's step is bf16 (gitea #47).
|
|
"""
|
|
with torch.autocast(real.device.type, enabled=False):
|
|
eps = torch.rand(real.size(0), 1, device=real.device)
|
|
x_hat = eps * real.float() + (1 - eps) * fake.float()
|
|
if mask is not None:
|
|
x_hat = x_hat * mask
|
|
x_hat = x_hat.requires_grad_(True)
|
|
scores = critic_fn(x_hat)
|
|
grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0]
|
|
if mask is not None:
|
|
grad = grad * mask
|
|
penalty = (grad.norm(2, dim=1) - 1) ** 2
|
|
if mask is not None:
|
|
valid = (mask.sum(dim=1) > 0).float()
|
|
return (penalty * valid).sum() / valid.sum().clamp_min(1.0)
|
|
return penalty.mean()
|
|
|
|
|
|
def critic_loss(
|
|
critic_fn: CriticFn,
|
|
real: torch.Tensor,
|
|
fake: torch.Tensor,
|
|
gp_weight: float,
|
|
mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""WGAN-GP critic loss. `fake` must already be `.detach()`'d by the caller."""
|
|
gp = gradient_penalty(critic_fn, real, fake, mask=mask)
|
|
return critic_fn(fake).mean() - critic_fn(real).mean() + gp_weight * gp
|
|
|
|
|
|
def generator_loss(critic_fn: CriticFn, fake: torch.Tensor) -> torch.Tensor:
|
|
return -critic_fn(fake).mean()
|