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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user