diff --git a/docs/performance.md b/docs/performance.md index 6d30e86..7b9eb12 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -358,6 +358,59 @@ batched-traversal-inference work in Optimization Priorities #5 lands — that is the change that would put compile on the dominant phase, not just on the trainer's optimization steps. Not enabled on `main`. +### AMP on trainer networks (2026-05-07, regression) + +Wrapped the trainer optimization phases with `torch.autocast(fp16)` and +`torch.amp.GradScaler`: `_train_advantage` and `_train_strategy` now run +their network forward/backward/optimizer step through the AMP path when +`run.use_amp=true` and the trainer device is CUDA. + +Safety mitigations included in the implementation: + +- `GradScaler.unscale_(optimizer)` is called before `clip_grad_norm_`. +- Non-finite loss guard increments `amp/nonfinite_loss_count` and skips the + bad step instead of applying it. +- Advantage squared loss computes `diff.float().square()` so the loss + reduction is fp32 even when the forward path is autocast to fp16. +- Strategy logits are cast back to fp32 before `masked_fill` and + `log_softmax`. +- Metrics now expose `amp/grad_scale` and `amp/nonfinite_loss_count`. + +Measurement used the small `smoke.yaml` config with synthetic replay-memory +samples via: + +```bash +uv run python scripts/bench_amp_trainer.py \ + --config configs/deep_cfr/smoke.yaml \ + --runs 3 \ + --warmup 1 \ + --device cuda +``` + +| | mean ms/call | speedup vs fp32 | +| --- | ---: | ---: | +| fp32 | 3.22 | 1.00× | +| AMP (fp16) | 3.92 | 0.82× | + +Net result: regression. This matches the same dispatch-overhead-vs-kernel +benefit dynamic as the `torch.compile` regression above: the current trainer +model and smoke workload are too small for AMP's lower-precision kernels to +pay back autocast and scaler bookkeeping overhead. + +The full `default.yaml` 100-iteration A/B was intentionally skipped. Given +the small-model regression and the matching `torch.compile` precedent on the +same model family, there is no current evidence that spending GPU time on the +longer A/B would produce a different decision. The infrastructure is kept +merged but default-off: `run.use_amp=false` remains the default, and +re-enabling is a one-field config flip. + +Re-measure AMP only after the model grows to at least `hidden_size >= 1024` +or `num_layers >= 6`. At that point run both the fast +`scripts/bench_amp_trainer.py` micro-bench and the formal 100-iteration +fp32-vs-AMP A/B. If AMP still provides less than 5% speedup at that larger +model size, keep it default-off and raise the next re-measure trigger to an +even larger model. + ### GPU forward profiling for batched traversal (2026-05-07, decision support) To decide whether Optimization Priorities #5 (batched traversal inference) is @@ -556,6 +609,8 @@ Do this in order. Skipping ahead is the failure mode that creates misleading 2. **Next**: experiment with a larger network config. Measure compute vs learning-curve trade-off with the existing toolchain (no compile/TRT yet). This step decides the model size that future optimizations target. + It is also the prerequisite for revisiting AMP, `torch.compile`, and + TensorRT: all three are dispatch-overhead-bound on the current small model. 3. **Then**: re-measure `torch.compile` on the trainer at the chosen model size. The earlier regression was size-bound; expect a different result. 4. **Then**: integrate TensorRT into the inference server (covers traversal diff --git a/scripts/bench_amp_trainer.py b/scripts/bench_amp_trainer.py new file mode 100644 index 0000000..656c08d --- /dev/null +++ b/scripts/bench_amp_trainer.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import argparse +import math +import statistics +import time +from pathlib import Path + +import numpy as np +import torch + +from coolrl_lost_cities.games.classic.deep_cfr.cli import _with_overrides +from coolrl_lost_cities.games.classic.deep_cfr.config import load_config +from coolrl_lost_cities.games.classic.deep_cfr.memory import TrainingSample +from coolrl_lost_cities.games.classic.deep_cfr.tracking import NullRunTracker +from coolrl_lost_cities.games.classic.deep_cfr.trainer import DeepCFRTrainer + +REPO_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_CONFIG = REPO_ROOT / "configs" / "deep_cfr" / "default.yaml" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Benchmark trainer-side AMP train phases.") + parser.add_argument("--config", default=str(DEFAULT_CONFIG)) + parser.add_argument("--runs", type=int, default=20) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--device", default="cuda") + return parser.parse_args() + + +def _seed_memories(trainer: DeepCFRTrainer) -> None: + rng = np.random.default_rng(trainer.config.run.seed + 909) + action_size = trainer.action_size + input_dim = trainer.input_dim + advantage_count = ( + trainer.config.optimization.advantage_batch_size + * trainer.config.optimization.advantage_updates_per_iteration + ) + strategy_count = ( + trainer.config.optimization.strategy_batch_size + * trainer.config.optimization.strategy_updates_per_iteration + ) + for player in range(2): + for index in range(max(advantage_count, trainer.config.optimization.advantage_batch_size)): + legal = rng.random(action_size) > 0.25 + legal[int(rng.integers(0, action_size))] = True + trainer.advantage_memories[player].add( + TrainingSample( + info_state=rng.normal(size=input_dim).astype(np.float32), + target=rng.normal(scale=20.0, size=action_size).astype(np.float32), + legal_mask=legal, + iteration=index + 1, + player=player, + ), + rng, + ) + for index in range(max(strategy_count, trainer.config.optimization.strategy_batch_size)): + legal = rng.random(action_size) > 0.25 + legal[int(rng.integers(0, action_size))] = True + target = np.zeros(action_size, dtype=np.float32) + weights = rng.random(np.count_nonzero(legal)).astype(np.float32) + weights /= weights.sum() + target[legal] = weights + trainer.strategy_memory.add( + TrainingSample( + info_state=rng.normal(size=input_dim).astype(np.float32), + target=target, + legal_mask=legal, + iteration=index + 1, + player=-1, + ), + rng, + ) + + +def _new_trainer(*, use_amp: bool, config_path: str, device: str) -> DeepCFRTrainer: + config = _with_overrides( + load_config(config_path), + { + "run": {"use_amp": use_amp, "device": device, "max_iterations": 1}, + "checkpoint": {"save_every": 0, "save_latest": False}, + "evaluation": {"eval_every": 0}, + }, + ) + trainer = DeepCFRTrainer( + config=config, + game_config=config.rules.to_lost_cities_config(seed=config.run.seed), + device=device, + tracker=NullRunTracker(), + ) + _seed_memories(trainer) + return trainer + + +def _measure(trainer: DeepCFRTrainer, runs: int) -> list[float]: + durations: list[float] = [] + for iteration in range(runs): + trainer.iteration = iteration + 1 + trainer._runtime_metrics = {} + if trainer.device.type == "cuda": + torch.cuda.synchronize() + started = time.perf_counter() + trainer._train_advantage_networks() + trainer._train_strategy_network() + if trainer.device.type == "cuda": + torch.cuda.synchronize() + durations.append(time.perf_counter() - started) + return durations + + +def _summary(values: list[float], warmup: int) -> tuple[float, float, float]: + measured = values[warmup:] + if not measured: + raise ValueError("warmup must be less than runs") + return ( + statistics.mean(measured), + statistics.median(measured), + sorted(measured)[math.ceil(0.95 * (len(measured) - 1))], + ) + + +def main() -> None: + args = parse_args() + if args.runs <= 0: + raise SystemExit("--runs must be positive") + if args.warmup < 0 or args.warmup >= args.runs: + raise SystemExit("--warmup must be non-negative and less than --runs") + if args.device == "cuda" and not torch.cuda.is_available(): + raise SystemExit("CUDA is not available; pass --device cpu for CPU fallback smoke.") + + print(f"Config: {args.config}") + print(f"Device: {args.device}") + print(f"Runs: {args.runs} Warmup: {args.warmup}") + print() + rows: list[tuple[str, float, float, float]] = [] + for use_amp, label in ((False, "fp32"), (True, "amp")): + trainer = _new_trainer(use_amp=use_amp, config_path=args.config, device=args.device) + values = _measure(trainer, args.runs) + mean, p50, p95 = _summary(values, args.warmup) + rows.append((label, mean, p50, p95)) + del trainer + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + print(f"{'mode':<8} {'mean_ms':>10} {'p50_ms':>10} {'p95_ms':>10}") + for label, mean, p50, p95 in rows: + print(f"{label:<8} {mean * 1000.0:>10.2f} {p50 * 1000.0:>10.2f} {p95 * 1000.0:>10.2f}") + fp32_mean = rows[0][1] + amp_mean = rows[1][1] + print() + print(f"speedup: {fp32_mean / amp_mean:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py index 9274422..8f0395d 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py @@ -230,6 +230,9 @@ class DeepCFRTrainer: lr=self.config.optimization.learning_rate, weight_decay=self.config.optimization.weight_decay, ) + self._amp_enabled = bool(self.config.run.use_amp) and self.device.type == "cuda" + self._amp_dtype = torch.float16 + self._scaler = torch.amp.GradScaler("cuda", enabled=self._amp_enabled) self.advantage_memories = [ ReservoirMemory(self.config.memory.advantage_capacity) for _ in range(2) ] @@ -258,6 +261,8 @@ class DeepCFRTrainer: self._runtime_metrics: dict[str, float | int] = {} self._inference_server: InferenceServerController | None = None self._last_inference_weight_sync_iteration: int | None = None + if self.config.run.use_amp and not self._amp_enabled: + self.tracker.log_event("AMP requested but trainer device is not CUDA; running fp32.") def checkpoint_payload(self, metrics: IterationMetrics | None = None) -> dict: return { @@ -339,6 +344,9 @@ class DeepCFRTrainer: self._runtime_metrics["time/strategy_train_seconds"] = ( time.perf_counter() - strategy_started ) + if self.config.run.use_amp: + self._runtime_metrics["amp/grad_scale"] = float(self._scaler.get_scale()) + self._runtime_metrics.setdefault("amp/nonfinite_loss_count", 0) eval_started = time.perf_counter() eval_metrics = self._evaluate(iteration) @@ -953,40 +961,52 @@ class DeepCFRTrainer: - sample_started ) x, y, legal, sample_iterations = self._batch_tensors(batch) - pred = network(x) - diff = (pred - y).masked_fill(~legal, 0.0) - if self.config.training_weighting.mode == "none": - loss = diff.square().sum() / legal.sum().clamp_min(1) - elif self.config.training_weighting.mode == "lcfr": - sample_weights = self._iteration_weights( - sample_iterations, self.config.training_weighting.lcfr_alpha - ) - action_weights = sample_weights[:, None] * legal.float() - loss = (diff.square() * action_weights).sum() / action_weights.sum().clamp_min( - 1.0e-12 - ) - else: - positive_weights = self._iteration_weights( - sample_iterations, self.config.training_weighting.dcfr_alpha - ) - negative_weights = self._iteration_weights( - sample_iterations, self.config.training_weighting.dcfr_beta - ) - target_weights = torch.where( - y >= 0.0, positive_weights[:, None], negative_weights[:, None] - ) - action_weights = target_weights * legal.float() - loss = (diff.square() * action_weights).sum() / action_weights.sum().clamp_min( - 1.0e-12 - ) + with torch.autocast( + device_type="cuda", + dtype=self._amp_dtype, + enabled=self._amp_enabled, + ): + pred = network(x) + diff = (pred - y).masked_fill(~legal, 0.0).float() + squared = diff.square() + if self.config.training_weighting.mode == "none": + loss = squared.sum() / legal.sum().clamp_min(1) + elif self.config.training_weighting.mode == "lcfr": + sample_weights = self._iteration_weights( + sample_iterations, self.config.training_weighting.lcfr_alpha + ) + action_weights = sample_weights[:, None] * legal.float() + loss = (squared * action_weights).sum() / action_weights.sum().clamp_min( + 1.0e-12 + ) + else: + positive_weights = self._iteration_weights( + sample_iterations, self.config.training_weighting.dcfr_alpha + ) + negative_weights = self._iteration_weights( + sample_iterations, self.config.training_weighting.dcfr_beta + ) + target_weights = torch.where( + y >= 0.0, positive_weights[:, None], negative_weights[:, None] + ) + action_weights = target_weights * legal.float() + loss = (squared * action_weights).sum() / action_weights.sum().clamp_min( + 1.0e-12 + ) + if not torch.isfinite(loss): + self._record_nonfinite_loss() + optimizer.zero_grad(set_to_none=True) + continue optimizer.zero_grad(set_to_none=True) - loss.backward() + self._scaler.scale(loss).backward() if self.config.optimization.grad_clip > 0.0: + self._scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_( network.parameters(), self.config.optimization.grad_clip ) - optimizer.step() - losses.append(float(loss.detach().cpu())) + self._scaler.step(optimizer) + self._scaler.update() + losses.append(float(loss.detach().float().cpu())) return float(np.mean(losses)) if losses else 0.0 def _train_strategy( @@ -1007,7 +1027,13 @@ class DeepCFRTrainer: - sample_started ) x, y, legal, sample_iterations = self._batch_tensors(batch) - logits = network(x).masked_fill(~legal, torch.finfo(torch.float32).min) + with torch.autocast( + device_type="cuda", + dtype=self._amp_dtype, + enabled=self._amp_enabled, + ): + logits = network(x) + logits = logits.float().masked_fill(~legal, torch.finfo(torch.float32).min) log_probs = nn.functional.log_softmax(logits, dim=-1).masked_fill(~legal, 0.0) per_sample_loss = -(y * log_probs).sum(dim=-1) if self.config.training_weighting.mode == "none": @@ -1026,12 +1052,23 @@ class DeepCFRTrainer: loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum().clamp_min( 1.0e-12 ) + if not torch.isfinite(loss): + self._record_nonfinite_loss() + optimizer.zero_grad(set_to_none=True) + continue optimizer.zero_grad(set_to_none=True) - loss.backward() + self._scaler.scale(loss).backward() if self.config.optimization.grad_clip > 0.0: + self._scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_( network.parameters(), self.config.optimization.grad_clip ) - optimizer.step() - last_loss = float(loss.detach().cpu()) + self._scaler.step(optimizer) + self._scaler.update() + last_loss = float(loss.detach().float().cpu()) return last_loss + + def _record_nonfinite_loss(self) -> None: + self._runtime_metrics["amp/nonfinite_loss_count"] = ( + int(self._runtime_metrics.get("amp/nonfinite_loss_count", 0)) + 1 + ) diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index c38c44e..ff18c53 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/tests/games/classic/test_deep_cfr_trainer.py @@ -4,6 +4,7 @@ import re from pathlib import Path import numpy as np +import pytest import torch from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state, input_dim from coolrl_lost_cities.games.classic.deep_cfr.traversal import CythonDeepCFRTraverser @@ -399,6 +400,76 @@ def test_deep_cfr_trainer_supports_lcfr_and_dcfr_loss_weighting() -> None: assert metrics[0].strategy_loss >= 0.0 +def test_deep_cfr_trainer_amp_cpu_falls_back_to_fp32() -> None: + trainer = DeepCFRTrainer( + _deep_cfr_config( + { + "run": {"max_iterations": 1, "seed": 25, "device": "cpu", "use_amp": True}, + "network": {"hidden_size": 16}, + "traversal": { + "traversals_per_player": 1, + "max_depth": 2, + "max_nodes_per_traversal": 32, + }, + "optimization": { + "advantage_updates_per_iteration": 1, + "strategy_updates_per_iteration": 1, + "advantage_batch_size": 2, + "strategy_batch_size": 2, + }, + "checkpoint": {"save_every": 0}, + "evaluation": {"eval_every": 0}, + } + ), + LostCitiesConfig(seed=25), + device="cpu", + ) + + metrics = trainer.train() + + assert len(metrics) == 1 + assert metrics[0].runtime_metrics["amp/grad_scale"] == 1.0 + assert metrics[0].runtime_metrics["amp/nonfinite_loss_count"] == 0 + assert metrics[0].advantage_loss >= 0.0 + assert metrics[0].strategy_loss >= 0.0 + + +def test_deep_cfr_trainer_amp_cuda_smoke() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + trainer = DeepCFRTrainer( + _deep_cfr_config( + { + "run": {"max_iterations": 1, "seed": 26, "device": "cuda", "use_amp": True}, + "network": {"hidden_size": 16}, + "traversal": { + "traversals_per_player": 1, + "max_depth": 2, + "max_nodes_per_traversal": 32, + }, + "optimization": { + "advantage_updates_per_iteration": 1, + "strategy_updates_per_iteration": 1, + "advantage_batch_size": 2, + "strategy_batch_size": 2, + }, + "checkpoint": {"save_every": 0}, + "evaluation": {"eval_every": 0}, + } + ), + LostCitiesConfig(seed=26), + device="cuda", + ) + + metrics = trainer.train() + + assert len(metrics) == 1 + assert metrics[0].runtime_metrics["amp/grad_scale"] > 0.0 + assert metrics[0].runtime_metrics["amp/nonfinite_loss_count"] == 0 + assert np.isfinite(metrics[0].advantage_loss) + assert np.isfinite(metrics[0].strategy_loss) + + def test_deep_cfr_cython_traverser_restores_state_and_collects_samples() -> None: trainer = DeepCFRTrainer( _deep_cfr_config(