Add bf16 autocast to the training loop (gitea #47) #68
Assignees
lars (Lars Bogner)
Clear assignees
lars (Lars Bogner)
No Assignees
Notifications
Due Date
No due date set.
Dependencies
No dependencies set.
Reference: lars/giant#68
Reference in New Issue
Block a user
Delete Branch "fix/issue-47"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
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:
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.
can't do it, rather than silently falling back to fp32.
best-checkpoint selection it drives) stays fp32 so it's comparable
across every run recorded so far.
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.
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
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>