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>