Wire AMP into Deep CFR trainer behind run.use_amp flag (default off)

Adds torch.autocast(fp16) + GradScaler around _train_advantage and
_train_strategy when run.use_amp=true and device=cuda. CPU/non-CUDA
falls back to fp32 no-op. Mitigations:
- scaler.unscale_(optimizer) before grad_clip.
- nonfinite-loss guard skips overflowing batches and counts them.
- diff.float().square() in advantage loss to avoid fp16 overflow.
- strategy mask/log_softmax kept in fp32.

New metrics: amp/grad_scale, amp/nonfinite_loss_count.

Tests: AMP CUDA smoke + CPU fallback in test_deep_cfr_trainer.py.

Bench: scripts/bench_amp_trainer.py micro-benches train phases under
synthetic replay memory. smoke.yaml result is fp32 3.22ms / AMP 3.92ms
(0.82×, regression). 100-iter A/B on default.yaml deliberately
skipped: smoke regression mirrors the 2026-05-07 torch.compile
regression dynamic (dispatch overhead > kernel benefit at this model
size) and re-confirming on the same size adds no information.

Default stays run.use_amp: false. Re-enable trigger documented in
docs/performance.md: hidden_size >= 1024 or num_layers >= 6, then run
the bench script + 100-iter A/B before flipping default.
This commit is contained in:
2026-05-07 20:21:17 +09:00
parent ad0be89857
commit 6c976f468a
4 changed files with 351 additions and 33 deletions
@@ -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(