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

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>

tests/test_amp.py

@@ -0,0 +1,136 @@
+"""Tests for giant/training/amp.py (gitea #47)."""
+
+import tempfile
+from pathlib import Path
+
+import pytest
+import torch
+
+from giant.model.routers import EnergyRouter
+from giant.model.wgan import gradient_penalty
+from giant.training.amp import resolve_autocast
+from giant.training.stage2_inputs import _remaining_energy_fraction
+from test_train import _base_cfg, _run_train
+
+# ---------------------------------------------------------------------------
+# resolve_autocast
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_autocast_fp32_is_disabled():
+    device_type, dtype, enabled = resolve_autocast("fp32", torch.device("cpu"))
+    assert device_type == "cpu"
+    assert dtype is torch.float32
+    assert enabled is False
+
+
+def test_resolve_autocast_bf16_on_cpu_is_enabled():
+    """CPU bf16 autocast is what lets the mixed-precision path be tested
+    without a GPU (torch 2.3 supports it)."""
+    device_type, dtype, enabled = resolve_autocast("bf16", torch.device("cpu"))
+    assert device_type == "cpu"
+    assert dtype is torch.bfloat16
+    assert enabled is True
+
+
+def test_resolve_autocast_bf16_on_unsupported_cuda_raises(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (7, 0))
+    monkeypatch.setattr(torch.cuda, "get_device_name", lambda device=None: "Tesla V100")
+    with pytest.raises(ValueError, match="bf16"):
+        resolve_autocast("bf16", torch.device("cuda"))
+
+
+def test_resolve_autocast_bf16_on_mps_raises():
+    with pytest.raises(ValueError, match="bf16"):
+        resolve_autocast("bf16", torch.device("mps"))
+
+
+def test_resolve_autocast_unknown_precision_raises():
+    with pytest.raises(ValueError, match="fp32.*bf16"):
+        resolve_autocast("fp16", torch.device("cpu"))
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: train() under bf16 on CPU
+# ---------------------------------------------------------------------------
+
+
+def test_train_end_to_end_bf16_cpu_completes_and_stores_fp32_params():
+    """Reuses tests/test_train.py's synthetic-batch harness — train() itself
+    is device-agnostic, and CPU bf16 autocast is real (not mocked) in torch
+    2.3, so this is a genuine exercise of the autocast region added to
+    FlowDDPMStageTrainer.step/WGANStageTrainer.step, not just a config
+    passthrough check.
+
+    Also asserts the checkpoint's stored parameters are fp32: autocast only
+    changes the dtype of intermediate activations, never the model's own
+    stored weights — a regression here would mean something accidentally
+    cast the model itself (e.g. `model.to(dtype=torch.bfloat16)`) rather than
+    using autocast."""
+    cfg = _base_cfg()
+    cfg["train"]["precision"] = "bf16"
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        assert (out_dir / "last.pt").exists()
+        assert (out_dir / "metrics.csv").exists()
+        ckpt = torch.load(out_dir / "last.pt", weights_only=False)
+        for stage_key in ("model", "sec_decoder"):
+            if stage_key not in ckpt:
+                continue
+            for name, tensor in ckpt[stage_key].items():
+                if tensor.is_floating_point():
+                    assert tensor.dtype == torch.float32, f"{stage_key}.{name} is {tensor.dtype}, expected fp32"
+
+
+@pytest.mark.parametrize("generator", ["wgan", "flow"])
+def test_train_end_to_end_bf16_cpu_stage2_generators(generator):
+    """bf16 covers both trainer subclasses (FlowDDPMStageTrainer and
+    WGANStageTrainer) — the wgan default in _base_cfg exercises the
+    generator-forward/critic-scoring autocast region added to
+    WGANStageTrainer.step, and flow exercises the plain _compute wrap."""
+    cfg = _base_cfg()
+    cfg["train"]["precision"] = "bf16"
+    cfg["stage2_model"]["generator"] = generator
+    with tempfile.TemporaryDirectory() as tmp:
+        _run_train(cfg, Path(tmp) / "run")
+
+
+# ---------------------------------------------------------------------------
+# fp32 guards: correct in fp32, quietly degrade in bf16 — stay fp32 even
+# under an active bf16 autocast region.
+# ---------------------------------------------------------------------------
+
+
+def test_remaining_energy_fraction_stays_fp32_under_bf16_autocast():
+    fraction = torch.rand(4, 5).to(torch.bfloat16)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        out = _remaining_energy_fraction(fraction)
+    assert out.dtype == torch.float32
+
+
+def test_gradient_penalty_stays_fp32_under_bf16_autocast():
+    critic = torch.nn.Linear(6, 1)
+
+    def critic_fn(x):
+        return critic(x)
+
+    real = torch.randn(4, 6)
+    fake = torch.randn(4, 6)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        gp = gradient_penalty(critic_fn, real, fake)
+    assert gp.dtype == torch.float32
+
+
+def test_router_balance_and_entropy_loss_stay_fp32_under_bf16_autocast():
+    router = EnergyRouter(n_experts=3)
+    cond_cont = torch.randn(8, 15)
+    cond_cat = torch.zeros(8, 2, dtype=torch.long)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        balance = router.balance_loss(cond_cont, cond_cat)
+        entropy = router.entropy_loss(cond_cont, cond_cat)
+        weights = router.combine_weights(cond_cont, cond_cat)
+    assert balance.dtype == torch.float32
+    assert entropy.dtype == torch.float32
+    assert weights.dtype == torch.float32

tests/test_config.py

@@ -816,6 +816,25 @@ def test_validate_config_bad_stop_sampling_rejected():
         assert "stop_sampling" in str(e)
 
 
+def test_validate_config_default_precision_is_fp32():
+    assert gconfig.DEFAULT_CONFIG["train"]["precision"] == "fp32"
+
+
+def test_validate_config_bf16_precision_accepted():
+    cfg = _cfg_with(**{"train.precision": "bf16"})
+    gconfig.validate_config(cfg)  # no raise
+
+
+@pytest.mark.parametrize("bad", ["fp16", "bogus", ""])
+def test_validate_config_bad_precision_rejected(bad):
+    cfg = _cfg_with(**{"train.precision": bad})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "precision" in str(e)
+
+
 def test_validate_config_stage1_context_sampled_accepted_with_both_stages_active():
     """gitea #41: 'sampled' is now implemented, so DEFAULT_CONFIG's
     stage1_model/stage2_model.active = true (both) must let it through."""
@@ -1177,6 +1196,11 @@ def test_overrides_from_flags_train_block_passthrough():
     assert overrides == {"train": {"epochs": 5, "lr": 1e-3}}
 
 
+def test_overrides_from_flags_precision_passthrough():
+    overrides = gconfig.overrides_from_flags({"precision": "bf16"})
+    assert overrides == {"train": {"precision": "bf16"}}
+
+
 @pytest.mark.parametrize(
     ("shorthand", "explicit", "path_key"),
     [

tests/test_router.py

@@ -1195,3 +1195,37 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow():
     )
     assert sec_cont.shape == (B, K_MAX, 4)
     assert sec_valid.shape == (B, K_MAX)
+
+
+def test_routed_and_unrouted_trunk_agree_on_dtype_under_bf16_autocast():
+    """gitea #47 regression: `_route_forward`'s accumulator (giant/model/
+    trunks.py) used to be a hard-fp32 `torch.zeros`, so under autocast a
+    `RoutedTrunk` returned fp32 while an unrouted `ExpertTrunk` returned
+    bf16 — `router.enabled` alone silently changed the model's output dtype.
+    Checked in both train mode (the differentiable mixture sum) and eval
+    mode (the masked `out[mask] = expert(...)` dispatch) — the two branches
+    of `_route_forward` had independent copies of the bug."""
+    torch.manual_seed(0)
+    cond_cont, cond_cat = _cond(B=6)
+    x = torch.randn(6, X_DIM)
+    t = torch.rand(6)
+
+    unrouted = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=2,
+    )
+    routed = _routed_stage1(n_experts=3)
+
+    for train_mode in (True, False):
+        unrouted.train(train_mode)
+        routed.train(train_mode)
+        with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+            out_unrouted = unrouted(x, cond_cont, cond_cat, t=t)
+            out_routed = routed(x, cond_cont, cond_cat, t=t)
+        assert out_unrouted.dtype == out_routed.dtype, (
+            f"train={train_mode}: unrouted returned {out_unrouted.dtype}, routed returned {out_routed.dtype}"
+        )