From 78978769f6fb85dc0856a71b3553b3421723f2e2 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 17 Aug 2026 15:36:12 +0200 Subject: [PATCH] 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 --- giant/cli.py | 10 +++ giant/config.py | 16 ++++ giant/model/routers.py | 50 ++++++++---- giant/model/trunks.py | 41 ++++++++-- giant/model/wgan.py | 38 +++++---- giant/training/amp.py | 47 +++++++++++ giant/training/stage2_inputs.py | 14 +++- giant/training/trainers.py | 116 +++++++++++++++++---------- tests/test_amp.py | 136 ++++++++++++++++++++++++++++++++ tests/test_config.py | 24 ++++++ tests/test_router.py | 34 ++++++++ 11 files changed, 446 insertions(+), 80 deletions(-) create mode 100644 giant/training/amp.py create mode 100644 tests/test_amp.py diff --git a/giant/cli.py b/giant/cli.py index 99ed4b7..9f199b2 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -612,6 +612,14 @@ def train( "steps (default: 50); per-epoch metrics always log in full", ), ] = None, + precision: Annotated[ + Optional[str], + typer.Option( + "--precision", + help="Training-step autocast precision: 'fp32' (default) or " + "'bf16'. No 'fp16' — see giant.training.amp.resolve_autocast", + ), + ] = None, ) -> None: """Train the GIANT surrogate model.""" batch_size_auto = False @@ -647,6 +655,7 @@ def train( "wandb_project": wandb_project, "wandb_run_name": wandb_run_name, "wandb_log_every": wandb_log_every, + "precision": precision, "hidden_dim": hidden_dim, "n_blocks": n_blocks, "dropout": dropout, @@ -725,6 +734,7 @@ def train( typer.echo(f"device: {_device}") typer.echo(f"out_dir: {out_dir}") + typer.echo(f"precision: {t['precision']}") run_train_job( data=data, diff --git a/giant/config.py b/giant/config.py index 7259975..8e8e7d0 100644 --- a/giant/config.py +++ b/giant/config.py @@ -801,6 +801,12 @@ class TrainConfig: # thousands of steps. Per-epoch metrics (the metrics.csv row) always log # in full. wandb_log_every: int = 50 + # Training-step autocast dtype: "fp32" (default, no autocast) or "bf16". + # No "fp16" — GradScaler and the double-backward in + # giant.model.wgan.gradient_penalty don't mix well, and bf16 alone covers + # every training GPU in the fleet (Ampere and newer). See + # giant.training.amp.resolve_autocast (gitea #47). + precision: str = "fp32" @classmethod def from_dict(cls, d: dict | None) -> "TrainConfig": @@ -822,6 +828,7 @@ class TrainConfig: wandb_project=d.get("wandb_project", "giant"), wandb_run_name=d.get("wandb_run_name", ""), wandb_log_every=d.get("wandb_log_every", 50), + precision=d.get("precision", "fp32"), ) def to_dict(self) -> dict: @@ -842,6 +849,7 @@ class TrainConfig: "wandb_project": self.wandb_project, "wandb_run_name": self.wandb_run_name, "wandb_log_every": self.wandb_log_every, + "precision": self.precision, } @@ -1114,6 +1122,7 @@ FLAG_SPECS: tuple[FlagSpec, ...] = ( FlagSpec("wandb_project", ("train.wandb_project",)), FlagSpec("wandb_run_name", ("train.wandb_run_name",)), FlagSpec("wandb_log_every", ("train.wandb_log_every",)), + FlagSpec("precision", ("train.precision",)), # --hidden-dim/--n-blocks/--dropout are stage-1-only backward-compat # shorthands (they predate stage2_model having its own flags); # --stage1-* wins when both are given. @@ -1473,6 +1482,13 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None: if stop_sampling not in ("greedy", "sample"): raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'") + precision = _get_path(cfg, "train.precision") + if precision not in ("fp32", "bf16"): + raise ValueError( + f"train.precision = {precision!r} — must be 'fp32' or 'bf16' " + "('fp16' is not supported: see giant.training.amp.resolve_autocast)" + ) + stage1_context = _get_path(cfg, "stage2_model.stage1_context") if stage1_context not in ("truth", "sampled"): raise ValueError(f"stage2_model.stage1_context = {stage1_context!r} — must be 'truth' or 'sampled'") diff --git a/giant/model/routers.py b/giant/model/routers.py index 12afa36..34df658 100644 --- a/giant/model/routers.py +++ b/giant/model/routers.py @@ -45,21 +45,36 @@ class Router(nn.Module): straight-through Gumbel-softmax (`gumbel=True`, train mode only): hardens the forward pass to a one-hot sample (matching eval-time top-1 dispatch) while keeping the soft sample's gradient on backward. + + Forced fp32 (`torch.autocast(..., enabled=False)`) regardless of the + caller's ambient `train.precision` autocast region: `clamp_min(1e-8)` + below sits under bf16's precision but *above* fp16's ~6e-8 subnormal + floor, so `log_probs` degrading here is exactly the kind of quiet + drift that cost a whole rollout benchmark before (see the MoE section + of CLAUDE.md's Roadmap) — cheap to rule out (gitea #47). """ - probs = self.gate(cond_cont, cond_cat) - if not (self.gumbel and self.training): - return probs - log_probs = torch.log(probs.clamp_min(1e-8)) - return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1) + with torch.autocast(cond_cont.device.type, enabled=False): + probs = self.gate(cond_cont, cond_cat) + if not (self.gumbel and self.training): + return probs + log_probs = torch.log(probs.clamp_min(1e-8)) + return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1) def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: """(B,) hard expert index, used for eval-time grouped dispatch.""" - return self.gate(cond_cont, cond_cat).argmax(dim=-1) + with torch.autocast(cond_cont.device.type, enabled=False): + return self.gate(cond_cont, cond_cat).argmax(dim=-1) def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017).""" - importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,) - return (importance.std() / (importance.mean() + 1e-8)) ** 2 + """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017). + + Forced fp32 — `importance` sums `gate()` over the whole batch (a + large-magnitude accumulation in reduced precision), then takes a + `std/mean` ratio: a classic catastrophic-cancellation shape (gitea + #47).""" + with torch.autocast(cond_cont.device.type, enabled=False): + importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,) + return (importance.std() / (importance.mean() + 1e-8)) ** 2 def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: """Optional supervised auxiliary loss shaping the router's own belief. @@ -76,12 +91,17 @@ class Router(nn.Module): def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for - the full explanation, unchanged in v0.3.0.""" - gate = self.gate(cond_cont, cond_cat) # (B, n_experts) - row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1) # (B,) - norm_entropy = row_entropy.mean() / math.log(self.n_experts) - importance = gate.sum(dim=0) # (n_experts,) - return norm_entropy, importance + the full explanation, unchanged in v0.3.0. + + Forced fp32, same rationale as `balance_loss`/`combine_weights`: the + `+ 1e-8` epsilon here is `entropy_loss`'s training-loss path too, not + just a diagnostic (gitea #47).""" + with torch.autocast(cond_cont.device.type, enabled=False): + gate = self.gate(cond_cont, cond_cat) # (B, n_experts) + row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1) # (B,) + norm_entropy = row_entropy.mean() / math.log(self.n_experts) + importance = gate.sum(dim=0) # (n_experts,) + return norm_entropy, importance ROUTER_REGISTRY: dict[str, type[Router]] = {} diff --git a/giant/model/trunks.py b/giant/model/trunks.py index 398523e..c8f546c 100644 --- a/giant/model/trunks.py +++ b/giant/model/trunks.py @@ -109,21 +109,50 @@ def _route_forward( N-expert dense compute, fully differentiable (`weight` is `router.combine_weights`). Eval mode: grouped top-1 dispatch — each row runs exactly one expert, the actual source of the per-call speedup. + + The accumulator's dtype is deferred to the first expert call rather than + fixed at fp32: under autocast (`train.precision = "bf16"`, gitea #47) an + expert's `ResBlock` stack returns bf16, and an fp32-fixed accumulator + would silently upcast every mixture term (train mode) or downcast every + dispatched row via `index_put_` (eval mode) — making a `RoutedTrunk` + return a different dtype than the unrouted `ExpertTrunk` it's a drop-in + replacement for, purely because `router.enabled` was set. + + `router.combine_weights` is deliberately fp32 internally (it forces its + own autocast-disabled region — see `Router.combine_weights`'s docstring), + so `weights` itself is always fp32 regardless of the ambient precision. + Left as-is, `weights[:, i:i+1] * expert(x, cond)` would type-promote the + whole mixture back to fp32 by ordinary PyTorch promotion rules — the same + dtype-mismatch bug this function exists to avoid, just moved one line + over. `weights` is cast down to each expert's own output dtype right + before combining: the softmax stays numerically stable at fp32, but its + *result* (values in [0, 1], not precision-sensitive to represent) loses + nothing meaningful by then being used at bf16. """ if training: - weights = router.combine_weights(cond_cont, cond_cat) # (B, n_experts) - out = torch.zeros(x.shape[0], experts[0].out_dim, device=x.device) + weights = router.combine_weights(cond_cont, cond_cat) # (B, n_experts), fp32 + out = None for i, expert in enumerate(experts): - out = out + weights[:, i : i + 1] * expert(x, cond) + expert_out = expert(x, cond) + term = weights[:, i : i + 1].to(expert_out.dtype) * expert_out + out = term if out is None else out + term + assert out is not None, "RoutedTrunk built with zero experts" return out idx = router.top1(cond_cont, cond_cat) # (B,) - out_dim = experts[0].out_dim - out = torch.zeros(x.shape[0], out_dim, device=x.device) + out = None for i, expert in enumerate(experts): mask = idx == i if mask.any(): - out[mask] = expert(x[mask], cond[mask]) + expert_out = expert(x[mask], cond[mask]) + if out is None: + out = torch.zeros(x.shape[0], expert_out.shape[-1], device=x.device, dtype=expert_out.dtype) + out[mask] = expert_out + if out is None: + # No row was ever dispatched (only reachable with an empty batch, + # x.shape[0] == 0) — nothing to infer a dtype from, so fall back to + # x's own, matching this function's pre-autocast behavior. + out = torch.zeros(x.shape[0], experts[0].out_dim, device=x.device, dtype=x.dtype) return out diff --git a/giant/model/wgan.py b/giant/model/wgan.py index ac6a852..a456266 100644 --- a/giant/model/wgan.py +++ b/giant/model/wgan.py @@ -22,21 +22,31 @@ def gradient_penalty( norm to 1 — `x_hat`/`grad` are forced to all-zero for such a row, which would otherwise contribute a constant `(||0|| - 1)^2 == 1` bias to the mean regardless of critic behavior — so they're excluded from the mean. + + Deliberately kept fp32 (`torch.autocast(..., enabled=False)`) regardless + of the caller's ambient `train.precision` autocast region: this is a + `create_graph=True` double-backward, and `grad.norm(2, dim=1)` sums + squares over the critic's full input width (hundreds of dims for stage + 2), which overflows bf16's range at gradient magnitudes well within + normal early-WGAN-GP territory. Disclosed cost: the critic forward + inside this function always runs fp32, even when the rest of the WGAN + stage's step is bf16 (gitea #47). """ - eps = torch.rand(real.size(0), 1, device=real.device) - x_hat = eps * real + (1 - eps) * fake - if mask is not None: - x_hat = x_hat * mask - x_hat = x_hat.requires_grad_(True) - scores = critic_fn(x_hat) - grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0] - if mask is not None: - grad = grad * mask - penalty = (grad.norm(2, dim=1) - 1) ** 2 - if mask is not None: - valid = (mask.sum(dim=1) > 0).float() - return (penalty * valid).sum() / valid.sum().clamp_min(1.0) - return penalty.mean() + with torch.autocast(real.device.type, enabled=False): + eps = torch.rand(real.size(0), 1, device=real.device) + x_hat = eps * real.float() + (1 - eps) * fake.float() + if mask is not None: + x_hat = x_hat * mask + x_hat = x_hat.requires_grad_(True) + scores = critic_fn(x_hat) + grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0] + if mask is not None: + grad = grad * mask + penalty = (grad.norm(2, dim=1) - 1) ** 2 + if mask is not None: + valid = (mask.sum(dim=1) > 0).float() + return (penalty * valid).sum() / valid.sum().clamp_min(1.0) + return penalty.mean() def critic_loss( diff --git a/giant/training/amp.py b/giant/training/amp.py new file mode 100644 index 0000000..05edee8 --- /dev/null +++ b/giant/training/amp.py @@ -0,0 +1,47 @@ +"""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)" + ) diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py index ca0b6fd..342db88 100644 --- a/giant/training/stage2_inputs.py +++ b/giant/training/stage2_inputs.py @@ -120,9 +120,17 @@ def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor: slot i: `1.0` at `i=0`, `prod_{j=1` ("no re-derivation needed": the existing stick-breaking encoding is already scale-free, so this is derivable from - the batch's ground-truth stick logits alone, no `e_sec` required).""" - cumprod = torch.cumprod(1.0 - fraction, dim=1) - return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1) + the batch's ground-truth stick logits alone, no `e_sec` required). + + Forced fp32 regardless of the caller's ambient `train.precision` autocast + region: a `cumprod` over `K_MAX` slots in bf16 underflows to zero within a + handful of slots, killing `remaining_frac` as a conditioning signal — the + numpy encoder (`giant.data.transforms.encode_secondaries`'s stick-breaking + twin) already promotes to float64 for exactly this reason (gitea #47).""" + with torch.autocast(fraction.device.type, enabled=False): + fraction = fraction.float() + cumprod = torch.cumprod(1.0 - fraction, dim=1) + return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1) def _shift_prev(x: torch.Tensor) -> torch.Tensor: diff --git a/giant/training/trainers.py b/giant/training/trainers.py index fbe2799..cfbd540 100644 --- a/giant/training/trainers.py +++ b/giant/training/trainers.py @@ -28,6 +28,7 @@ from giant.data.dataset import StepBatch from giant.model.network import Router, build_objective, resolve_type_n_classes, stage2_type_dim from giant.model.wgan import generator_loss, gradient_penalty from giant.sample import sample_stage1 +from giant.training.amp import resolve_autocast from giant.training.metrics import MetricSpec, stage_metric, train_metric, val_metric from giant.training.stage2_inputs import ( _assemble_stage2_ar_inputs_scheduled, @@ -109,6 +110,7 @@ class StageSpec: warmup_epochs: int = 0 epochs: int = 1 steps_per_epoch: int = 1 + precision: str = "fp32" # routing auxiliaries lambda_balance: float = 0.0 @@ -174,6 +176,7 @@ class StageSpec: warmup_epochs=t.warmup_epochs, epochs=t.epochs, steps_per_epoch=max(steps_per_epoch, 1), + precision=t.precision, lambda_balance=stage_spec.router.lambda_balance, lambda_proc=stage_spec.router.lambda_proc, lambda_entropy=stage_spec.router.lambda_entropy, @@ -256,6 +259,12 @@ class StageTrainer: self.router = _stage_router(self.model) self._modules = (self.model, *extra_modules) + # Resolved once (not re-derived every step) — see + # giant.training.amp.resolve_autocast (gitea #47). + self._autocast_device_type, self._autocast_dtype, self._autocast_enabled = resolve_autocast( + spec.precision, device + ) + self.particle_type_cfg = spec.particle_type self.particle_type_n_classes = spec.particle_type_n_classes self.ema_decay = spec.ema_decay @@ -518,6 +527,22 @@ class StageTrainer: stop_acc = (((logits >= 0).float() == target).float() * mask_f).sum() / denom return l_stop, stop_acc + def _autocast(self) -> torch.autocast: + """The training-step autocast region (`train.precision`, gitea #47). + + Only wraps forward/loss computation — `backward()`/`optimizer.step()` + stay outside, and `val_loss` never calls this at all, so validation + (and the best-checkpoint selection it drives) stays precision- + independent and comparable against every fp32-only run recorded so + far. `enabled=False` under `precision = "fp32"` (the default) makes + this a true no-op, so callers never need to branch on precision + themselves.""" + return torch.autocast( + self._autocast_device_type, + dtype=self._autocast_dtype, + enabled=self._autocast_enabled, + ) + def _step_optimizer(self, optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float: """`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning the pre-clip grad norm. The one place the grad-clip constant lives. @@ -780,7 +805,8 @@ class FlowDDPMStageTrainer(StageTrainer): self.spec.gumbel_tau_end, ) epoch = global_step // self.spec.steps_per_epoch - out = self._compute(batch, device, epoch=epoch) + with self._autocast(): + out = self._compute(batch, device, epoch=epoch) grad_norm = self._step_optimizer(self.optimizer, out["loss"], self.params) if not self.frozen: self.lr_sched.step() @@ -943,49 +969,53 @@ class WGANStageTrainer(StageTrainer): grad_probe: dict[str, float] = {} ar_inputs = None - if not self.is_stage2: - real = x1_s1 + with self._autocast(): + if not self.is_stage2: + real = x1_s1 - def critic_fn(x): - return self.critic(x, cond_cont, cond_cat) + def critic_fn(x): + return self.critic(x, cond_cont, cond_cat) - z = torch.randn(B, self.model.noise_dim, device=device) - fake = self.model(z, cond_cont, cond_cat) - mask = None - else: - real, fake_raw, mask, critic_fn, ar_inputs = self._stage2_real_and_fake( - _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx), - stage1_ctx, - global_step, - device, - ) - if self.particle_type_cfg.target == "onehot": - # Straight-through Gumbel-softmax relaxation of the type - # slice only — the critic must see a hard one-hot forward - # (matching what "real" data looks like) while gradient - # still flows smoothly to the generator. grad_probe captures - # the gradient-magnitude instrumentation — see - # _relax_onehot_type_slice's docstring. - tau = _gumbel_tau( + z = torch.randn(B, self.model.noise_dim, device=device) + fake = self.model(z, cond_cont, cond_cat) + mask = None + else: + real, fake_raw, mask, critic_fn, ar_inputs = self._stage2_real_and_fake( + _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx), + stage1_ctx, global_step, - self.total_steps, - self.spec.type_gumbel_tau_start, - self.spec.type_gumbel_tau_end, + device, ) - fake_raw = _relax_onehot_type_slice( - fake_raw, - sec_cont.size(1), - CONT_SLOT_DIM, - stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes), - tau, - grad_probe=grad_probe, - ) - fake = fake_raw * mask + if self.particle_type_cfg.target == "onehot": + # Straight-through Gumbel-softmax relaxation of the type + # slice only — the critic must see a hard one-hot forward + # (matching what "real" data looks like) while gradient + # still flows smoothly to the generator. grad_probe captures + # the gradient-magnitude instrumentation — see + # _relax_onehot_type_slice's docstring. + tau = _gumbel_tau( + global_step, + self.total_steps, + self.spec.type_gumbel_tau_start, + self.spec.type_gumbel_tau_end, + ) + fake_raw = _relax_onehot_type_slice( + fake_raw, + sec_cont.size(1), + CONT_SLOT_DIM, + stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes), + tau, + grad_probe=grad_probe, + ) + fake = fake_raw * mask - # --- critic step (every batch) --- - fake_detached = fake.detach() - real_score = critic_fn(real) - fake_score = critic_fn(fake_detached) + # --- critic step (every batch) --- + fake_detached = fake.detach() + real_score = critic_fn(real) + fake_score = critic_fn(fake_detached) + + # gradient_penalty forces its own fp32 region internally (see its + # docstring) regardless of the ambient autocast above. gp = gradient_penalty(critic_fn, real, fake_detached, mask=mask) d_loss = fake_score.mean() - real_score.mean() + self.gp_weight * gp wasserstein = (real_score.mean() - fake_score.mean()).detach() @@ -994,8 +1024,9 @@ class WGANStageTrainer(StageTrainer): # --- generator (+ n_sec) step --- did_g_step = global_step % self.n_critic == 0 - l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device) - l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs) + with self._autocast(): + l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device) + l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs) # On a non-generator-step batch with no n_sec_head/stop_head on this # stage (n_sec now defaults to stage 2), there's nothing for @@ -1003,7 +1034,8 @@ class WGANStageTrainer(StageTrainer): # be a graph-less zero tensor, which .backward() rejects outright. skip_g_step = not did_g_step and self.model.n_sec_head is None and self.model.stop_head is None if did_g_step: - g_loss_adv = generator_loss(critic_fn, fake) + with self._autocast(): + g_loss_adv = generator_loss(critic_fn, fake) g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * (l_nsec + l_stop) else: g_loss_adv = torch.zeros((), device=device) diff --git a/tests/test_amp.py b/tests/test_amp.py new file mode 100644 index 0000000..e53ce12 --- /dev/null +++ b/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 diff --git a/tests/test_config.py b/tests/test_config.py index e00dec9..71d8713 100644 --- a/tests/test_config.py +++ b/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"), [ diff --git a/tests/test_router.py b/tests/test_router.py index 64c7325..25284bd 100644 --- a/tests/test_router.py +++ b/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}" + )