No mixed precision (autocast/GradScaler), no torch.compile #47

New Issue

2026-08-13T15:08:05+02:00

lars commented

2026-08-13 15:08:05 +02:00

grep finds no autocast, no GradScaler, no bfloat16, no torch.compile
anywhere in giant/. For a project whose stated target is a ~10×
native-Geant4 eval budget, and whose flow/AR path costs k_max × steps
sequential model calls per physics step, bf16 autocast on the training loop and
torch.compile on the sampler are likely worth more wall-clock than any
architectural change in the modularity batch of issues.

Two caveats to check first: _route_forward
(giant/model/trunks.py:58, 65) allocates its accumulator with
torch.zeros(..., device=x.device) and no dtype, which will need
dtype=x.dtype under autocast; and torch.compile interacts poorly with the
eval-mode grouped out[mask] = expert(...) dispatch (data-dependent shapes →
recompiles), so the routed path likely wants dynamic=True or an explicit
compile opt-out.

Migrated from issues.md (v0.3.0 branch review, 2026-08-13), Issue 22.

`grep` finds no `autocast`, no `GradScaler`, no `bfloat16`, no `torch.compile` anywhere in `giant/`. For a project whose stated target is a ~10× native-Geant4 eval budget, and whose flow/AR path costs `k_max × steps` sequential model calls per physics step, bf16 autocast on the training loop and `torch.compile` on the sampler are likely worth more wall-clock than any architectural change in the modularity batch of issues. Two caveats to check first: `_route_forward` (`giant/model/trunks.py:58, 65`) allocates its accumulator with `torch.zeros(..., device=x.device)` and **no `dtype`**, which will need `dtype=x.dtype` under autocast; and `torch.compile` interacts poorly with the eval-mode grouped `out[mask] = expert(...)` dispatch (data-dependent shapes → recompiles), so the routed path likely wants `dynamic=True` or an explicit compile opt-out. --- Migrated from `issues.md` (v0.3.0 branch review, 2026-08-13), Issue 22.

lars added the performance label 2026-08-13 15:08:05 +02:00

lars commented

2026-08-17 15:36:30 +02:00

Fixed in 7897876 on fix/issue-47.

Implemented the bf16-autocast half of this issue: a new train.precision config key ("fp32" default, "bf16" opt-in) wraps the training step (both FlowDDPMStageTrainer and WGANStageTrainer) in torch.autocast, resolved via a new giant.training.amp.resolve_autocast.

Decisions made during planning/implementation:

fp32 + bf16 only, no fp16/GradScaler — fp16 breaks giant/model/routers.py's three 1e-8 epsilons (below fp16's ~6e-8 subnormal floor) and overflows gradient_penalty's grad norm at ordinary early-WGAN-GP magnitudes. Every training GPU in the fleet (A100/L40S/H200/RTX 4070) has native bf16; fp16 would only matter for pre-Ampere V100s.
resolve_autocast raises loudly (naming the device) if bf16 is requested on hardware that can't do it, instead of silently training in fp32.
Autocast wraps the training step only — val_loss/best-checkpoint selection stays fp32, so it's comparable across every run recorded so far. Inference/the samplers are untouched (also sidesteps .cpu().numpy() calls that would otherwise hit TypeError: Got unsupported ScalarType BFloat16).
Fixed a real bug this surfaced: _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 — same config knob, different model output dtype. Fixed to match the experts' dtype (regression test in tests/test_router.py).
Added explicit fp32 guards (autocast(enabled=False)) around spots that are correct in fp32 but degrade quietly (not a 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's own measurements) benefiting most on both axes.

Deliberately left open, per discussion during planning: torch.compile — the other half of this issue — is a much larger surface (the routed trunk's data-dependent out[mask] = expert(x[mask]) dispatch, the autoregressive sampler's if finished.all(): break and per-slot recompiles, sample_ddim's .item() calls, and rollout.py's arbitrary post-filter batch dimension all need real restructuring, not just a decorator). Happy to file that as its own follow-up issue if wanted.

Fixed in 7897876 on fix/issue-47. Implemented the bf16-autocast half of this issue: a new `train.precision` config key (`"fp32"` default, `"bf16"` opt-in) wraps the training step (both `FlowDDPMStageTrainer` and `WGANStageTrainer`) in `torch.autocast`, resolved via a new `giant.training.amp.resolve_autocast`. Decisions made during planning/implementation: - fp32 + bf16 only, no fp16/GradScaler — fp16 breaks `giant/model/routers.py`'s three `1e-8` epsilons (below fp16's ~6e-8 subnormal floor) and overflows `gradient_penalty`'s grad norm at ordinary early-WGAN-GP magnitudes. Every training GPU in the fleet (A100/L40S/H200/RTX 4070) has native bf16; fp16 would only matter for pre-Ampere V100s. - `resolve_autocast` raises loudly (naming the device) if bf16 is requested on hardware that can't do it, instead of silently training in fp32. - Autocast wraps the training step only — `val_loss`/best-checkpoint selection stays fp32, so it's comparable across every run recorded so far. Inference/the samplers are untouched (also sidesteps `.cpu().numpy()` calls that would otherwise hit `TypeError: Got unsupported ScalarType BFloat16`). - Fixed a real bug this surfaced: `_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 — same config knob, different model output dtype. Fixed to match the experts' dtype (regression test in `tests/test_router.py`). - Added explicit fp32 guards (`autocast(enabled=False)`) around spots that are correct in fp32 but degrade quietly (not a 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's own measurements) benefiting most on both axes. **Deliberately left open, per discussion during planning:** `torch.compile` — the other half of this issue — is a much larger surface (the routed trunk's data-dependent `out[mask] = expert(x[mask])` dispatch, the autoregressive sampler's `if finished.all(): break` and per-slot recompiles, `sample_ddim`'s `.item()` calls, and `rollout.py`'s arbitrary post-filter batch dimension all need real restructuring, not just a decorator). Happy to file that as its own follow-up issue if wanted.

lars closed this issue

2026-08-17 15:36:30 +02:00

lars referenced this issue from a commit

2026-08-17 15:39:03 +02:00

Add bf16 autocast to the training loop (gitea #47)

lars referenced this issue

2026-08-17 15:39:34 +02:00

Add bf16 autocast to the training loop (gitea #47) #68

lars referenced this issue from a commit

2026-08-17 15:45:52 +02:00

Merge pull request 'Add bf16 autocast to the training loop (gitea #47)' (#68) from fix/issue-47 into master

Sign in to join this conversation.