2 Commits

Author	SHA1	Message	Date
lars	c984d0a19d	chore: bump uv.lock and fix ruff 0.16 default-rule lint findings ... uv.lock was stale (ty 0.0.50 -> 0.0.78, ruff 0.15 -> 0.16, polars, numpy, typer, wandb, pytest, and others), all within existing pyproject.toml bounds. ruff 0.16 widened its default rule selection, taking this repo from 0 to 274 lint errors under the same config; --fix handled most of it (import sorting, Optional[X] -> X | None, ...), and the remainder (unused unpacked variables, dict()-as-literal, subprocess.run without explicit check=, a couple of intentional broad excepts/naive datetimes) were fixed or annotated by hand. Also fixes a real type-narrowing gap ty 0.0.78 caught in test_config_consumed_keys.py's `or`-combined isinstance check. torch stays pinned to 2.3.x (deliberate, see CLAUDE.md); pyarrow's <25 ceiling is left as a separate decision. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TMdZFqXXig7i3XkirSUxef	2026-09-04 14:09:29 +02:00
lars	78978769f6	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>	2026-08-17 15:36:12 +02:00