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>
48 lines
2.2 KiB
Python
48 lines
2.2 KiB
Python
"""Mixed-precision training support (`train.precision`, gitea #47).
|
|
|
|
Only `"fp32"` (no autocast) and `"bf16"` are supported — no `"fp16"`/
|
|
`GradScaler`. bf16 needs no gradient scaler and covers every training GPU in
|
|
the fleet (Ampere and newer: A100, L40S, H200, RTX 4070); fp16 would need a
|
|
scaler *and* fixes to two fragile spots that stay correct under bf16 but break
|
|
under fp16's narrower range — `giant.model.routers`' `1e-8` epsilons (below
|
|
fp16's ~6e-8 subnormal floor) and `giant.model.wgan.gradient_penalty`'s
|
|
sum-of-squares gradient norm (overflows fp16 above ~65504). Revisit if a
|
|
pre-Ampere (V100) training target ever shows up.
|
|
"""
|
|
|
|
import torch
|
|
|
|
_SUPPORTED_DEVICE_TYPES = ("cuda", "cpu")
|
|
|
|
|
|
def resolve_autocast(precision: str, device: torch.device) -> tuple[str, torch.dtype, bool]:
|
|
"""Resolves `train.precision` + a target device into the
|
|
`(device_type, dtype, enabled)` triple `torch.autocast` takes as kwargs —
|
|
computed once per `StageTrainer` rather than re-derived every step.
|
|
|
|
Raises `ValueError` rather than silently falling back to fp32: a training
|
|
run that's quietly not using the mixed precision it was configured for is
|
|
a wasted GPU-week, not a warning.
|
|
"""
|
|
if precision == "fp32":
|
|
return device.type, torch.float32, False
|
|
if precision != "bf16":
|
|
raise ValueError(f"unknown precision {precision!r}; must be 'fp32' or 'bf16'")
|
|
|
|
if device.type == "cuda":
|
|
if not torch.cuda.is_bf16_supported():
|
|
cap = torch.cuda.get_device_capability(device)
|
|
raise ValueError(
|
|
f"train.precision = 'bf16' but {torch.cuda.get_device_name(device)} "
|
|
f"(compute capability {cap[0]}.{cap[1]}) has no native bf16 support "
|
|
"(needs Ampere/sm_80 or newer) — use train.precision = 'fp32' instead"
|
|
)
|
|
return "cuda", torch.bfloat16, True
|
|
if device.type == "cpu":
|
|
# torch 2.3's CPU autocast supports bf16 unconditionally — this is
|
|
# also what lets the bf16 training path be tested without a GPU.
|
|
return "cpu", torch.bfloat16, True
|
|
raise ValueError(
|
|
f"train.precision = 'bf16' is not supported on device type {device.type!r} (only {_SUPPORTED_DEVICE_TYPES} are)"
|
|
)
|