Document Option A bench result, post-A calculus, plans, and cost reports
performance.md additions: - Batched Traversal Inference design decision (A vs B vs C with rationale). - Option A bench result and structural ceiling (realized batch ~7.2, IPC overhead exceeds GPU gain at small model size). - Post-A optimization calculus: why compile/TensorRT remain iter-neutral today and become meaningful only after model growth and/or denser eval. Sequencing matters; do not retest these on the current small model. - Free-threaded Python (3.13t/3.14t) note: cleanest endpoint in principle, but PyTorch maturity + Cython nogil audit cost block near-term adoption. docs/plans/ (4 plan documents for Codex execution): - batched_traversal_inference_server.md (executed; deferred). - amp_trainer.md. - torch_compile.md. - cython_safe_heuristic_bots.md (executed; first-pass landed). docs/reports/ (3 cost reports): - cost_pytorch_free_threaded_2026-05-07.md: WAIT 3-6 months; PyTorch wheels exist but our Cython is the gating cost. - cost_cython_nogil_audit_2026-05-07.md: medium effort, traversal.pyx carries 90% of blockers; Steps 1-3 (cfr_math/encoding nogil keywords, TraversalStats cdef class) are safe and cheap, Steps 4-6 wait for triggers. - cost_pytorch_cuda_multithread_2026-05-07.md: risky; optimizer.step / load_state_dict race silently with concurrent forward; per-thread default streams unset means naive threading serializes on default stream anyway. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.7
@@ -0,0 +1,259 @@
|
||||
# Plan: Trainer-side AMP (Automatic Mixed Precision)
|
||||
|
||||
**Status:** Ready for implementation
|
||||
**Owner:** Codex
|
||||
**Background:** See `docs/performance.md` → "AMP Status" (currently a no-op flag) and "Post-A Optimization Calculus" (AMP becomes more meaningful only at larger model sizes; today's gains are bounded by the small `DeepCFRMLP` and the dominant traversal phase). Also note the `torch.compile` regression experiment in the same doc — small models punish low-level kernel optimizations because dispatch overhead outweighs fused-kernel gains. AMP can hit the same wall.
|
||||
|
||||
## Goal
|
||||
|
||||
Wire `run.use_amp` into the Deep CFR trainer's optimization phases (`_train_advantage` and `_train_strategy`) using `torch.autocast` + `torch.amp.GradScaler`. Cut `advantage_train_seconds + strategy_train_seconds` (~7.1s/iter combined, ~40% of a non-eval iteration on the inspected default run) without regressing the learning curve.
|
||||
|
||||
## Non-goals
|
||||
|
||||
- Do not change traversal worker behavior. Traversal workers run on CPU today; AMP does not apply.
|
||||
- Do not change the inference server's forward path. The server has its own `inference_server.use_amp` flag (eval-only AMP, no GradScaler). Trainer AMP and server AMP are independent.
|
||||
- Do not change replay buffer dtype. Samples remain float32 in shared memory and on the host side; only the trainer's forward/backward switches to mixed precision.
|
||||
- Do not change network dtype, parameter dtype, or optimizer dtype. AMP only autocasts forward; parameters and master weights remain fp32.
|
||||
- Do not implement bf16 (no GradScaler needed) as the default. fp16 is the primary target; bf16 is a follow-up.
|
||||
|
||||
## Success criteria
|
||||
|
||||
1. With `--set run.use_amp=true` on `configs/deep_cfr/default.yaml` running on CUDA, `advantage_train_seconds + strategy_train_seconds` drops by at least 15% averaged over 20 non-eval iterations on `home` (RTX 3090) compared to `--set run.use_amp=false` baseline. **Realistic upper bound:** ~30% on `home`. If measured speedup is below 5%, document the result in `docs/performance.md` and leave the flag default-off (analogous to the `torch.compile` regression).
|
||||
2. With `--set run.use_amp=true`, `loss/advantage` and `loss/strategy` trajectories track the fp32 baseline within seed noise over at least 50 iterations on `default.yaml`. No NaN/Inf appears in `loss/*` rows of `metrics.jsonl`.
|
||||
3. Eval-side win-rate trajectories (`eval/<opponent>/win_rate0`) under AMP are indistinguishable from the fp32 baseline at iteration 50, 100, and 200 within the noise band of a single-seed comparison.
|
||||
4. `--set run.use_amp=false` (the default) produces byte-identical training to `main` for the same seed: no AMP-related code path runs.
|
||||
5. `run.use_amp=true` on CPU is a documented no-op (CUDA not available → skip autocast/scaler) and does not crash. Same for non-CUDA `run.device`.
|
||||
6. All existing tests pass. New unit test covering the AMP code path passes.
|
||||
|
||||
## Key files (current)
|
||||
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py`:
|
||||
- `DeepCFRTrainer.__init__` (lines ~187-258) — instantiates networks, optimizers; this is where the `GradScaler` should be created.
|
||||
- `DeepCFRTrainer._train_advantage` (lines ~892-948) — advantage forward+backward+step. AMP wrapping target.
|
||||
- `DeepCFRTrainer._train_strategy` (lines ~950-995) — strategy forward+backward+step. AMP wrapping target.
|
||||
- `DeepCFRTrainer._batch_tensors` (lines ~853-883) — tensors stay float32 (do not cast inputs).
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/config.py`:
|
||||
- `RunConfig.use_amp: bool = False` (line 25) — flag already exists.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/networks.py` — no changes; `DeepCFRMLP` and `ColorSharedNetwork` work under autocast as-is.
|
||||
- `tests/games/classic/test_deep_cfr_trainer.py` — extend with AMP smoke test.
|
||||
|
||||
## New surface
|
||||
|
||||
No new files. All changes live in `trainer.py` (and one optional config addition for AMP dtype selection). The plan extends `RunConfig` minimally if we want bf16 selectability; otherwise `run.use_amp` alone is sufficient and dtype is fp16 by default.
|
||||
|
||||
Optional `RunConfig` extension (deferred — only add if Step 5 measurement motivates it):
|
||||
|
||||
```python
|
||||
class RunConfig(StrictModel):
|
||||
# ... existing fields ...
|
||||
use_amp: bool = False
|
||||
amp_dtype: str = "float16" # or "bfloat16"
|
||||
```
|
||||
|
||||
## Where AMP code goes
|
||||
|
||||
### Construction (in `__init__`)
|
||||
|
||||
After optimizers are created:
|
||||
|
||||
```python
|
||||
self._amp_enabled = bool(self.config.run.use_amp) and self.device.type == "cuda"
|
||||
self._amp_dtype = torch.float16 # bf16 deferred
|
||||
self._scaler = torch.amp.GradScaler("cuda", enabled=self._amp_enabled)
|
||||
```
|
||||
|
||||
Single shared `GradScaler` across both advantage and strategy networks is fine — it tracks one global loss-scale that adjusts based on observed Inf/NaN gradients and applies to whichever optimizer it is told to step. Sharing the scaler matches PyTorch's recommended pattern for multi-network training and avoids two competing scale schedules.
|
||||
|
||||
### Advantage train loop (`_train_advantage`)
|
||||
|
||||
Replace the forward + backward + step block with:
|
||||
|
||||
```python
|
||||
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)
|
||||
if self.config.training_weighting.mode == "none":
|
||||
loss = diff.square().sum() / legal.sum().clamp_min(1)
|
||||
elif self.config.training_weighting.mode == "lcfr":
|
||||
...
|
||||
else:
|
||||
...
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
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)
|
||||
self._scaler.step(optimizer)
|
||||
self._scaler.update()
|
||||
losses.append(float(loss.detach().float().cpu()))
|
||||
```
|
||||
|
||||
Notes:
|
||||
- `legal.sum().clamp_min(1)` and the weighted-mean denominators are int/float reductions that stay fp32 — `autocast` only casts ops that are on the autocast-allow-list. The division is safe.
|
||||
- `loss.detach().float().cpu()` — explicit `.float()` to avoid logging an fp16 NaN that prints surprisingly.
|
||||
- The `unscale_` before `clip_grad_norm_` is mandatory; otherwise the clip threshold is applied to scaled gradients. This is the standard PyTorch AMP idiom.
|
||||
|
||||
### Strategy train loop (`_train_strategy`)
|
||||
|
||||
Same pattern. The `log_softmax` over masked logits remains numerically stable under fp16 because `masked_fill(~legal, finfo.min)` uses fp32-min before autocast can downcast — verify this; if autocast downcasts the mask fill value, replace with `torch.finfo(self._amp_dtype if self._amp_enabled else torch.float32).min` or apply the mask after autocast. Safer: apply `masked_fill` *outside* the autocast block (operating on the post-cast logits is fine — `masked_fill` is allow-listed).
|
||||
|
||||
### Disabled path
|
||||
|
||||
When `self._amp_enabled is False`, `torch.autocast(..., enabled=False)` is a true no-op (does not change dispatcher state), and `GradScaler(enabled=False)` makes `scale`, `unscale_`, `step`, and `update` all forward to plain optimizer behavior. So the same code runs in both modes — no branching needed in the hot loop.
|
||||
|
||||
## Flag semantics
|
||||
|
||||
| `run.use_amp` | `run.device` resolves to | Behavior |
|
||||
| --- | --- | --- |
|
||||
| `false` (default) | any | No autocast, no scaler. Behavior identical to current `main`. |
|
||||
| `true` | `cuda` | Autocast(fp16) + GradScaler active in advantage/strategy train loops. |
|
||||
| `true` | `cpu` | Logs a one-time warning "AMP requested but device is CPU; running fp32." Trainer behaves as `false`. |
|
||||
|
||||
The warning is logged via `self.tracker.log_event(...)` once at run start, after `_amp_enabled` is computed.
|
||||
|
||||
Trainer-side AMP is fully independent of `inference_server.use_amp`. Both can be on, both off, or either alone:
|
||||
|
||||
- `inference_server.use_amp` → server-process `torch.autocast` around the inference forward inside `inference_mode()`. No GradScaler (no backward pass). Affects traversal forward latency only.
|
||||
- `run.use_amp` → trainer-process autocast + GradScaler around `_train_advantage` and `_train_strategy` forward+backward+step. Affects optimization phases only.
|
||||
|
||||
## Numerical stability considerations
|
||||
|
||||
CFR regret targets can have wide dynamic range — both very small (regrets near zero for converged actions) and large in magnitude (early-iteration noisy estimates). fp16's 1e-5 to 6.5e4 representable range is narrow compared to fp32, so two failure modes are realistic:
|
||||
|
||||
1. **Gradient overflow → NaN.** Mitigated by `GradScaler`. The scaler observes Inf/NaN in unscaled gradients, skips the step, and halves the scale. Standard.
|
||||
2. **Forward overflow in the squared-error loss.** `diff.square()` on fp16 inputs can overflow if `(pred - y)` magnitude exceeds ~256. This is upstream of GradScaler — the loss itself becomes Inf in the autocast region and the entire batch is wasted. If this happens, the scaler will skip the step but the loss-scale scheduler will not recover because the issue is in the forward, not the gradient.
|
||||
|
||||
Mitigations for (2):
|
||||
|
||||
- Compute the squared error in fp32 explicitly: cast `diff` to fp32 via `diff.float()` before `.square()` if measurement shows fp16 overflow on real CFR samples. This is cheap and confined to the trainer.
|
||||
- Clamp `pred - y` to a wide-but-safe range before squaring: e.g. `diff.clamp(-128, 128)`. Only do this if measurement shows it's needed — clamping silently changes loss semantics.
|
||||
- Default plan: do not pre-mitigate. Add a NaN/Inf guard (see below); if it triggers in the wild, switch to the explicit `.float()` cast in the loss computation.
|
||||
|
||||
### NaN/Inf guard
|
||||
|
||||
After computing `loss` in each train step:
|
||||
|
||||
```python
|
||||
if not torch.isfinite(loss):
|
||||
self._runtime_metrics["amp/nonfinite_loss_count"] = (
|
||||
int(self._runtime_metrics.get("amp/nonfinite_loss_count", 0)) + 1
|
||||
)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
continue
|
||||
```
|
||||
|
||||
This makes overflow visible in `metrics.jsonl` per iteration without aborting training. Codex should also expose `_scaler.get_scale()` as `amp/grad_scale` per iteration so we can see the scaler's scale schedule in W&B / metrics plots.
|
||||
|
||||
## Validation: A/B learning-curve test
|
||||
|
||||
After implementation, run two short trainings with the same seed:
|
||||
|
||||
```bash
|
||||
# Baseline
|
||||
uv run lost-cities-deep-cfr train \
|
||||
--config configs/deep_cfr/default.yaml \
|
||||
--keep \
|
||||
--set run.max_iterations=100 \
|
||||
--set run.use_amp=false \
|
||||
--set run.experiment_name=amp-baseline
|
||||
|
||||
# AMP
|
||||
uv run lost-cities-deep-cfr train \
|
||||
--config configs/deep_cfr/default.yaml \
|
||||
--keep \
|
||||
--set run.max_iterations=100 \
|
||||
--set run.use_amp=true \
|
||||
--set run.experiment_name=amp-on
|
||||
```
|
||||
|
||||
Compare with `lost-cities-deep-cfr analyze` on each, plus a side-by-side table of:
|
||||
|
||||
- `loss/advantage`, `loss/strategy` at iters {25, 50, 75, 100}
|
||||
- `eval/<opponent>/win_rate0` at iters {50, 100} (eval_every=25 is default)
|
||||
- `time/advantage_train_seconds`, `time/strategy_train_seconds`, `time/iteration_seconds` means over iters 5-100 (iter 1-4 dropped as warm-up)
|
||||
- `amp/nonfinite_loss_count`, `amp/grad_scale` (AMP run only)
|
||||
|
||||
Acceptance: AMP win-rates within fp32 win-rates ± noise, training-phase wall clock down by ≥15%, no nonfinite losses, scaler scale stable (not collapsing toward 1).
|
||||
|
||||
## Implementation steps (ordered, each independently mergeable)
|
||||
|
||||
### Step 1: scaler + flag plumbing in `__init__`
|
||||
|
||||
- In `DeepCFRTrainer.__init__`, after the optimizer block, compute `self._amp_enabled`, `self._amp_dtype`, `self._scaler`.
|
||||
- Log a one-time warning if `config.run.use_amp` is requested but `self.device.type != "cuda"`.
|
||||
- Add an `amp/grad_scale` runtime metric emitted once per iteration in `run_iteration` (read `self._scaler.get_scale()` after train phases).
|
||||
- No behavior change yet — scaler is created but unused.
|
||||
|
||||
### Step 2: wrap `_train_advantage`
|
||||
|
||||
- Add `with torch.autocast(...)` around the forward + loss computation.
|
||||
- Replace `loss.backward()` → `self._scaler.scale(loss).backward()`.
|
||||
- If `grad_clip > 0`: `self._scaler.unscale_(optimizer)` before `clip_grad_norm_`.
|
||||
- Replace `optimizer.step()` → `self._scaler.step(optimizer)` + `self._scaler.update()`.
|
||||
- Add the NaN/Inf guard described above; emit `amp/nonfinite_loss_count` to runtime metrics.
|
||||
|
||||
### Step 3: wrap `_train_strategy`
|
||||
|
||||
- Same pattern as Step 2.
|
||||
- Verify `masked_fill(~legal, torch.finfo(torch.float32).min)` semantics under autocast. Either keep the mask fill outside autocast or use a dtype-aware `finfo`.
|
||||
|
||||
### Step 4: unit test
|
||||
|
||||
- Extend `tests/games/classic/test_deep_cfr_trainer.py`:
|
||||
- `test_train_advantage_amp_smoke`: build a tiny trainer with `run.use_amp=True`, `device="cuda"` (skip if CUDA unavailable via `pytest.skip`). Seed memories with synthetic samples. Run one iteration. Assert no exception, finite `loss/advantage`, `amp/grad_scale` present in runtime metrics.
|
||||
- `test_amp_cpu_falls_back`: `run.use_amp=True`, `device="cpu"`. Assert iteration runs without error and behaves as fp32 (e.g. compare loss to a fp32 reference run on the same seed).
|
||||
|
||||
### Step 5: bench + learning-curve A/B
|
||||
|
||||
- Add a short script `scripts/bench_amp_trainer.py` (mirrors `scripts/profile_gpu_forward.py` style):
|
||||
- Construct trainer with `default.yaml`.
|
||||
- Pre-fill advantage and strategy memories with `optimization.advantage_batch_size * advantage_updates_per_iteration` synthetic samples (so the train phases run end-to-end without needing a real traversal).
|
||||
- Run `_train_advantage` and `_train_strategy` 20 times each under `use_amp=False` and `use_amp=True`. Drop the first 2 as warm-up. Print mean ms per call.
|
||||
- Run the full A/B (Validation section above) and append a date-stamped subsection to `docs/performance.md` → "Experiments" with: hardware, iter-time delta, loss/win-rate parity table, NaN counts, scale schedule, decision (default-on / default-off / disabled).
|
||||
|
||||
### Step 6: documentation
|
||||
|
||||
- Update `docs/performance.md` § "AMP Status" to point to the new experiment subsection and remove the "treat as a no-op" line if AMP becomes default-on, or leave it and explain why if AMP measurement was a regression.
|
||||
- No CLAUDE.md / AGENTS.md updates needed — `--set run.use_amp=true` already works syntactically.
|
||||
|
||||
## Definition of done
|
||||
|
||||
- `run.use_amp=true` on CUDA produces measurable speedup on advantage+strategy train phases on `home`, with the A/B learning curve showing parity within noise over 100 iterations.
|
||||
- `run.use_amp=false` is byte-identical to `main`.
|
||||
- AMP code path covered by tests; CPU fallback documented and tested.
|
||||
- `metrics.jsonl` exposes `amp/grad_scale` and `amp/nonfinite_loss_count` so future runs are self-diagnosing.
|
||||
- An experiment subsection is appended to `docs/performance.md` recording the speedup and the parity check.
|
||||
- `uv run ruff check .` passes; `uv run pytest -q tests/games/classic/test_deep_cfr_trainer.py` passes.
|
||||
|
||||
## Risks and mitigations
|
||||
|
||||
- **Small-model regression (analogous to `torch.compile` 2026-05-07).** `DeepCFRMLP` at hidden=512, 3 layers is small. AMP overhead per call (cast/uncast, GradScaler bookkeeping) may exceed the kernel speedup at this size — the same dynamic that bit `torch.compile`. Mitigation: measure first; if speedup is below 5%, keep default off and document. Preserve the implementation on a branch (`experiments/amp-trainer`) for revisiting if `NetworkConfig.hidden_size` or `num_layers` increases.
|
||||
- **fp16 overflow in `diff.square()`.** Mitigation: NaN/Inf guard logs the count; if observed, cast `diff` to fp32 before squaring inside the autocast region (`diff.float().square()`). PyTorch will not redowncast it.
|
||||
- **GradScaler scale collapse.** If many consecutive steps overflow, the scaler can drop to scale=1 and stay there, defeating the point. Mitigation: log `amp/grad_scale` every iteration; if it stays ≤16 for >10 iterations, switch the loss-side cast (above) on.
|
||||
- **Interaction with `clip_grad_norm_`.** Forgetting `unscale_` before clipping silently changes the effective clip threshold. Mitigation: explicit step in the implementation; covered by code review.
|
||||
- **Determinism breakage.** Autocast can change op kernels and therefore reduction order; bit-exact reproducibility vs fp32 is not preserved. This is expected. Mitigation: A/B is on learning trajectories within seed noise, not on bit-identity.
|
||||
- **Inference server contention.** If `inference_server.use_amp` and `run.use_amp` are both on, two autocast regions exist in two processes — they do not conflict. Trainer GradScaler does not affect server inference.
|
||||
|
||||
## Bench plan
|
||||
|
||||
Two artifacts:
|
||||
|
||||
1. `scripts/bench_amp_trainer.py` (new, small) — micro-bench the `_train_advantage` and `_train_strategy` methods in isolation under `use_amp=False` vs `True`. Reports per-call ms with mean + p50 + p95 over 20 runs. This is the fast feedback loop during implementation.
|
||||
2. End-to-end run pair (above) — the real signal. 100-iter `default.yaml` baseline vs AMP, same seed, single GPU. Compare iter-time and learning curves.
|
||||
|
||||
Compare pattern matches the `torch.compile` experiment write-up in `docs/performance.md`. Use a similar table:
|
||||
|
||||
| | iter mean | adv+strat mean | adv+strat share | 1000-iter projection |
|
||||
| --- | ---: | ---: | ---: | ---: |
|
||||
| Baseline (use_amp=false) | TBD | TBD | TBD% | TBD h |
|
||||
| AMP (use_amp=true) | TBD | TBD | TBD% | TBD h |
|
||||
| Effect | TBD | TBD | TBD pp | TBD min |
|
||||
|
||||
## Out-of-scope follow-ups (do not start in this plan)
|
||||
|
||||
- bf16 dtype option (`amp_dtype: bfloat16`). bf16 sidesteps GradScaler entirely and avoids the fp16 overflow class. Add only if (a) hardware supports it efficiently (Ampere+) and (b) fp16 measurement shows scale collapse or frequent nonfinite losses.
|
||||
- AMP on the inference server's training-side weight push path. Server already has `inference_server.use_amp` for forward; backward is not its job.
|
||||
- AMP on evaluation forward. Evaluation runs `eval()` + no-grad; if eval forward becomes a bottleneck (see "Evaluation Optimization Options" #7 in `docs/performance.md`), wrap there separately.
|
||||
- `torch.compile` retry. The 2026-05-07 regression was size-bound. Re-evaluate only after `NetworkConfig` grows substantially (per "Post-A Optimization Calculus" in `docs/performance.md`); coordinate with that work, not this plan.
|
||||
- Larger model config experiment. AMP becomes meaningfully more useful at hidden≈1024 / layers≈6 per the post-A calculus. That is a separate model-architecture work item, not an AMP work item.
|
||||
@@ -0,0 +1,209 @@
|
||||
# Plan: Batched Traversal Inference Server (Priority #5, Option A)
|
||||
|
||||
**Status:** Ready for implementation
|
||||
**Owner:** Codex
|
||||
**Background:** See `docs/performance.md` → "Batched Traversal Inference: Design Decision (2026-05-07)" for the A/B/C analysis and rationale. This plan implements Option A.
|
||||
|
||||
## Goal
|
||||
|
||||
Replace per-worker single-state CPU policy forwards in Deep CFR traversal with a central GPU inference server that batches policy requests across all workers. Targets the dominant phase (`traversal_seconds` ≈ 60% of iteration time).
|
||||
|
||||
## Non-goals
|
||||
|
||||
- Do not modify the Cython traversal recursion structure.
|
||||
- Do not modify the game engine, replay buffer, or training loop math.
|
||||
- Do not implement Option C (single-process vectorized traversal). Keep it as future work.
|
||||
- Do not require Cython `nogil`-cleanliness.
|
||||
- Do not change the public CLI surface.
|
||||
|
||||
## Success criteria
|
||||
|
||||
1. With `traversal.inference_backend: server` enabled on `configs/deep_cfr/default.yaml`, end-to-end training produces eval-winrate trajectories indistinguishable (within seed noise) from the current `local` backend over at least 50 iterations on `home` hardware.
|
||||
2. On `home` (6-core + RTX 3090), `traversal_seconds` decreases by at least 30% compared to the current run profile in `docs/performance.md`.
|
||||
3. On `remote` (32-core + weak GPU), `traversal_seconds` decreases or stays within 10% of current; if it regresses more, fall back to `local` is the operator's choice — the plan still ships.
|
||||
4. With `traversal.inference_backend: local`, behavior is byte-identical to current `main`.
|
||||
5. All existing tests pass. New unit tests for the inference client/server round-trip pass.
|
||||
|
||||
## Key files (current)
|
||||
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/workers.py` — `run_traversal_worker_batch`. Spawns CPU networks per worker (`device = torch.device("cpu")` at line 62). This is the worker entry point that must learn about the inference server.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx` — recursive traversal. Two policy call sites:
|
||||
- Lines 473–475: `networks[player](x).squeeze(0).detach().cpu().numpy()` — advantage-net forward.
|
||||
- Lines 550–552: `self.strategy_network(x).squeeze(0).detach().cpu().numpy()` — strategy-net forward.
|
||||
- Both are the seams that route through the inference client when the backend is `server`.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/networks.py` — `DeepCFRMLP`.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py` — main loop; spawns the worker pool and owns the trainer-side networks. Must also start/stop the inference server and push weights periodically.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/config.py` — config schema. Add `traversal.inference_backend` and an `inference_server` block.
|
||||
|
||||
## New files
|
||||
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py` — server process: owns models on GPU, drains the request queue, runs batched forward, writes responses.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_client.py` — client used inside workers: claims a request slot, posts the encoded state, waits for the response. Provides a `forward(network_id, player, state)` API.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_buffers.py` — shared-memory tensor pool: pre-allocated `[num_slots, input_dim]` request buffer and `[num_slots, action_size]` response buffer, plus per-slot ready events and a free-slot stack.
|
||||
- `tests/games/classic/deep_cfr/test_inference_server.py` — unit/integration tests for the request path.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Process layout
|
||||
|
||||
```
|
||||
main process (trainer)
|
||||
├─ inference server process (GPU)
|
||||
│ - owns advantage networks (per-player), strategy network, league snapshots
|
||||
│ - request buffer (shared mem): [num_slots, input_dim] float32
|
||||
│ - response buffer (shared mem): [num_slots, action_size] float32
|
||||
│ - control queue (mp.Queue): control messages only
|
||||
│ - per-slot ready event (mp.Event[num_slots])
|
||||
├─ traversal worker 1 ─┐
|
||||
├─ ... ├─ inference clients post requests, wait on per-slot event
|
||||
└─ traversal worker N ─┘
|
||||
```
|
||||
|
||||
### Request protocol
|
||||
|
||||
A worker policy call becomes:
|
||||
1. Client pops a free slot id from a shared `mp.Queue`-backed free-slot stack (small int).
|
||||
2. Client writes the encoded state into `request_buffer[slot_id]` (shared memory; no pickle).
|
||||
3. Client puts a `RequestMessage(slot_id, network_kind, player, network_index)` onto the request control queue. `network_kind ∈ {ADVANTAGE, STRATEGY, LEAGUE}`. Pickle cost is negligible (small struct of ints).
|
||||
4. Server drains the request queue with a short batch window (`batch_window_us`, default 200μs) up to `max_batch` (default 256). Empty drain blocks on the queue with a small timeout.
|
||||
5. Server stacks the requested rows from `request_buffer`, runs forward, writes outputs back to `response_buffer[slot_id]` for each request.
|
||||
6. Server fires the per-slot ready event for each completed request.
|
||||
7. Client wakes on its slot's event, reads `response_buffer[slot_id]`, copies to a local numpy array, returns the slot to the free-slot stack.
|
||||
|
||||
### Weight sync
|
||||
|
||||
- Trainer holds master weights. On a fixed cadence (`inference_server.weight_sync_every` iterations, default 1 — push every iter), trainer sends a `WeightUpdateMessage` containing `state_dict`s via `torch.multiprocessing` (auto-shares tensor storage; cheap once and copied into server's GPU model).
|
||||
- Server applies `load_state_dict` and signals `weight_sync_complete`. Trainer waits before kicking off the next traversal iteration.
|
||||
- League snapshots: passed via the same channel when the league updates. League list is small relative to per-iter cost.
|
||||
|
||||
### Backpressure & lifecycle
|
||||
|
||||
- Free-slot stack size = `num_slots` (default `max(64, 4 * num_workers * worker_chunk_size)`). Workers block on slot allocation if all slots in flight; this is the natural backpressure.
|
||||
- Server shutdown: trainer puts a `Shutdown` sentinel on the control queue at training end. Server drains, exits.
|
||||
- Crash isolation: if the server dies, workers will hang on their slot events. Trainer monitors the server process; on death, it raises and tears down the pool. No silent corruption.
|
||||
|
||||
### IPC choice
|
||||
|
||||
- `torch.multiprocessing` for weight passing (auto-shares tensors).
|
||||
- `multiprocessing.shared_memory.SharedMemory` (numpy view) for request/response buffers — manual slot management. Pre-allocated once at startup; no per-call allocation.
|
||||
- `multiprocessing.Queue` for control messages only (slot ids and small structs).
|
||||
- `multiprocessing.Event` array for per-slot wakeups.
|
||||
|
||||
Rationale: see `docs/performance.md` § "IPC mechanism: multiprocessing + shared memory".
|
||||
|
||||
## Config schema
|
||||
|
||||
Extend `config.py`:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class InferenceServerConfig:
|
||||
enabled: bool = False # if False, behave like current code
|
||||
device: str = "cuda" # server-side device
|
||||
num_slots: int | None = None # None → auto: max(64, 4 * num_workers * worker_chunk_size)
|
||||
max_batch: int = 256
|
||||
batch_window_us: int = 200
|
||||
weight_sync_every: int = 1 # iterations
|
||||
use_amp: bool = False # eval-only AMP for forward (no grad)
|
||||
|
||||
@dataclass
|
||||
class TraversalConfig:
|
||||
# ... existing fields ...
|
||||
inference_backend: Literal["local", "server"] = "local"
|
||||
```
|
||||
|
||||
In `default.yaml`, leave `inference_backend: local` for now. Add an explicit `configs/deep_cfr/default_server.yaml` variant that flips it on for benchmarking.
|
||||
|
||||
## Worker integration
|
||||
|
||||
`run_traversal_worker_batch` (workers.py) currently:
|
||||
- Builds CPU `DeepCFRMLP` instances and loads `state_dict`s.
|
||||
- Passes them as positional `networks` into the `Traversal` Cython object.
|
||||
|
||||
When `inference_backend == "server"`:
|
||||
- Skip building local networks. Instead, build an `InferenceClient` bound to the shared buffers and queues that the trainer wires in via `TraversalWorkerBatch`.
|
||||
- Pass a small `NetworkProxy` object with the same call signature as the current network: `proxy(x: torch.Tensor) -> torch.Tensor`. Internally it converts to numpy, calls `client.forward(...)`, returns a torch tensor.
|
||||
- Critically: the Cython traversal sites at `traversal.pyx:473-475` and `:550-552` should not need source changes if `NetworkProxy` is a callable returning a 2-D tensor. The existing `.squeeze(0).detach().cpu().numpy()` chain still works on the proxy's returned tensor (which can just be a CPU tensor wrapping the numpy result). **Verify this; if Cython has typed assumptions that reject a Python proxy, fall back to a thin Python helper invoked from the `.pyx` instead.**
|
||||
|
||||
Add to `TraversalWorkerBatch`:
|
||||
- `inference_handles: InferenceClientHandles | None` — shared-memory names, queue handles, event arrays. None when `inference_backend == "local"`.
|
||||
|
||||
## Trainer integration
|
||||
|
||||
`trainer.py`:
|
||||
1. On run start, if `inference_backend == "server"`: instantiate `InferenceServer` (spawns process), build `InferenceClientHandles`, push initial weights, wait for `weight_sync_complete`.
|
||||
2. Per iteration: before `pool.starmap(run_traversal_worker_batch, ...)`, push fresh weights if `iter % weight_sync_every == 0`. Pass `inference_handles` into each `TraversalWorkerBatch`.
|
||||
3. After traversal: same as today.
|
||||
4. On run end: send shutdown sentinel; join the server process.
|
||||
|
||||
## Implementation steps (ordered, each independently mergeable)
|
||||
|
||||
### Step 1: shared-memory buffer module
|
||||
|
||||
- Create `inference_buffers.py` with `InferenceBuffers` class: pre-allocates request/response numpy arrays via `SharedMemory`, exposes `attach(name)` for child processes, has `release()` cleanup.
|
||||
- Free-slot management: `mp.Queue` of slot ids, populated at startup with `range(num_slots)`.
|
||||
- Per-slot ready events: `[mp.Event() for _ in range(num_slots)]`.
|
||||
- Unit test: parent creates buffers, child attaches, writes a row, parent reads. Verify zero-copy semantics.
|
||||
|
||||
### Step 2: inference client
|
||||
|
||||
- Create `inference_client.py`. `InferenceClient.forward(network_kind, player, network_index, state_np: np.ndarray) -> np.ndarray`:
|
||||
- Pop free slot, write state, post request, wait on event, copy response, return slot.
|
||||
- Add a `NetworkProxy` callable that wraps `client.forward` to look like a `torch.nn.Module` for traversal call sites.
|
||||
- Unit test: stub a server thread that echoes state*2; assert client gets the expected output.
|
||||
|
||||
### Step 3: inference server
|
||||
|
||||
- Create `inference_server.py`. Spawn-friendly entry function `run_inference_server(handles, model_config, control_queue, weight_queue)`.
|
||||
- Owns models on `device`. Sets `eval()` and `inference_mode()`.
|
||||
- Main loop: drain request queue with `batch_window_us` deadline, group by `(network_kind, network_index)`, run batched forward per group, scatter outputs to response buffer slots, fire events.
|
||||
- Handles `WeightUpdateMessage` and `Shutdown`.
|
||||
- Optional `use_amp`: wrap forward in `torch.autocast` when configured.
|
||||
- Unit test: send N requests across multiple kinds; verify outputs match `model(stacked_input)`.
|
||||
|
||||
### Step 4: config + trainer wiring
|
||||
|
||||
- Extend `config.py` with `InferenceServerConfig` and `TraversalConfig.inference_backend`.
|
||||
- Update `trainer.py` to start/stop server, push weights, attach handles to `TraversalWorkerBatch`.
|
||||
- Update `workers.py`: when `inference_backend == "server"`, build proxies instead of CPU networks.
|
||||
- Add `configs/deep_cfr/default_server.yaml` flipping `traversal.inference_backend: server`.
|
||||
|
||||
### Step 5: traversal call-site verification
|
||||
|
||||
- Run with `inference_backend: server` and a tiny config (1 worker, 1 traversal). Confirm the proxy is callable from `traversal.pyx:473-475` and `:550-552` without Cython type errors.
|
||||
- If Cython rejects the proxy: refactor those two call sites to invoke a Python helper that takes `(networks, player, info_state)` and returns a numpy array. The helper picks `local` or `server` path by inspecting the object. This is a 2-line change per site.
|
||||
|
||||
### Step 6: integration tests
|
||||
|
||||
- `tests/games/classic/deep_cfr/test_inference_server.py`:
|
||||
- End-to-end smoke: 1 iteration of training with `server` backend on CPU device; assert no crash, replay buffer populated.
|
||||
- Equivalence: same seed, same initial weights, both backends → assert traversal samples match within numerical tolerance for at least 1 iteration. (May require `weight_sync_every = 1` and deterministic GPU forward; if exact match is fragile, accept distributional equivalence over 10 iterations.)
|
||||
|
||||
### Step 7: benchmarking
|
||||
|
||||
- Add `scripts/bench_inference_backend.py` (mirrors `scripts/profile_gpu_forward.py` style):
|
||||
- Run 10 iterations on `default.yaml` with `inference_backend=local`.
|
||||
- Run 10 iterations on `default_server.yaml`.
|
||||
- Print iteration-mean and per-phase mean times.
|
||||
- Document results in a new `docs/performance.md` experiment subsection (date-stamped).
|
||||
|
||||
## Risks and mitigations
|
||||
|
||||
- **Cython proxy incompatibility (Step 5).** Mitigation: pre-prototype with a 5-line Python script that imports `Traversal` and passes a stub callable in place of a `DeepCFRMLP`. If it works, the rest of the plan stands.
|
||||
- **Weight sync staleness invalidates CFR.** Mitigation: default `weight_sync_every: 1` (every iteration). Only loosen after measuring.
|
||||
- **GPU contention with eval.** Eval already runs on the trainer's device; the inference server adds another tenant on the same GPU. Mitigation: serialize eval and traversal phases (they already are sequential in the iteration loop). Document this in `inference_server.py`.
|
||||
- **Deadlock on server crash.** Mitigation: trainer monitors server `process.is_alive()` between iterations and raises if dead. Workers waiting on events will be torn down with the pool.
|
||||
- **Slot exhaustion under bursty load.** Mitigation: default `num_slots = 4 * num_workers * worker_chunk_size` to absorb bursts. Operators can tune.
|
||||
- **Remote (weak GPU) regression.** If `traversal_seconds` regresses on remote, the operator can flip `inference_backend: local` and ship without it. The plan succeeds either way.
|
||||
|
||||
## What this plan explicitly does not do
|
||||
|
||||
- Does not implement async client (workers stay sync-blocking on slot events). Async client is a future optimization if IPC round-trip dominates after measurement.
|
||||
- Does not implement `nogil` threading or Option C. See `docs/performance.md` for the deferred path.
|
||||
- Does not optimize encoding (`policy_encoding_seconds`). If post-A measurements show encoding dominates, that is a separate work item.
|
||||
|
||||
## Out-of-scope follow-ups (do not start)
|
||||
|
||||
- Option C (vectorized traversal) — only re-evaluate if A's measurements show GPU forward is no longer on the critical path.
|
||||
- TensorRT / `torch.compile` on the inference server's models — defer until A baseline numbers are collected.
|
||||
- Batched encoding inside the server (compute encoding from raw game state on GPU) — only if `policy_encoding_seconds` becomes the new bottleneck.
|
||||
@@ -0,0 +1,210 @@
|
||||
# Plan: Cython Port of the Safe-Heuristic Bot Family
|
||||
|
||||
**Status:** Ready for implementation
|
||||
**Owner:** Codex
|
||||
**Background:** See `docs/performance.md` → "Evaluation Breakdown" and "Post-A Optimization Calculus". The safe-heuristic family dominates eval wall-clock through `opponent_act_seconds` (7.57s / 4.63s / 9.22s for `safe_heuristic` / `safe_heuristic_loose` / `safe_heuristic_strict` in the inspected eval row), not GPU forward. As denser eval becomes operationally useful (`eval_every: 5`, `evaluation.games: 1000`), this cost becomes a first-order wall-clock concern.
|
||||
|
||||
## Goal
|
||||
|
||||
Port `SafeHeuristicBot` (and its loose / strict parameterizations) from pure-Python `bots/heuristic.py` to Cython, consuming game state directly through the existing typed Cython `GameState` interface, so that the per-iteration evaluation cost dominated by `opponent_act_seconds` shrinks substantially without any change to bot decision behavior.
|
||||
|
||||
## Non-goals
|
||||
|
||||
- Do not change bot decision logic. The Cython port must be byte-identical to the Python implementation under fixed seeds for the full action sequence of every benchmarked seeded game.
|
||||
- Do not port `RandomBot`, `PassiveDiscardBot`, or `NoisyPolicy`. Their per-iter cost is small (see `docs/performance.md` table: `random` 0.06s, `passive_discard` 0.00s, `noisy_safe` 0.57s).
|
||||
- Do not change the public bot registry names or CLI surface.
|
||||
- Do not touch the Cython traversal pipeline, encoding, training, or replay paths. This is a bot-only change.
|
||||
- Do not redesign `LostCitiesPolicy` or the `PolicyInput` / `Snapshot` interfaces.
|
||||
|
||||
## Success criteria
|
||||
|
||||
1. **Equivalence (hard).** For a fixed corpus of seeded games (see "Equivalence strategy" below), the Cython implementation produces a byte-identical action id sequence to the Python implementation for each of `safe_heuristic`, `safe_heuristic_loose`, `safe_heuristic_strict`. Deviation in any single action fails CI.
|
||||
2. **Per-opponent speedup.** On `configs/deep_cfr/default.yaml` evaluation, per-opponent `opponent_act_seconds` for each safe-heuristic variant decreases by **≥3×** versus the pre-port baseline measured on the same hardware. Concretely, target post-port values:
|
||||
- `safe_heuristic`: ≤ 2.5s (from 7.57s)
|
||||
- `safe_heuristic_loose`: ≤ 1.5s (from 4.63s)
|
||||
- `safe_heuristic_strict`: ≤ 3.1s (from 9.22s)
|
||||
3. **Eval wall-clock.** With the same `evaluation.games` and `evaluation.eval_every`, `eval/<variant>/elapsed_seconds` for each safe-heuristic variant decreases proportionally, and total `evaluation_seconds` for an iteration that runs the full opponent suite decreases meaningfully (the safe-heuristic opponents are the dominant tail per the performance doc).
|
||||
4. **All existing tests pass**, including every `test_safe_heuristic_*` test in `tests/games/classic/test_bots.py`. New equivalence tests (see below) also pass.
|
||||
5. **No behavioral regression in training.** A short training run on `default.yaml` with the Cython bots in eval reproduces the same eval-winrate trajectory (within seed noise) as the Python implementation over at least one full eval cadence.
|
||||
|
||||
## Independence claim
|
||||
|
||||
This work is **fully independent** of:
|
||||
|
||||
- the batched traversal inference server plan (`docs/plans/batched_traversal_inference_server.md`),
|
||||
- the AMP wiring (`run.use_amp`),
|
||||
- the `torch.compile` experiment branch (`experiments/torch-compile`).
|
||||
|
||||
It touches no traversal, training, or networks code. It can ship in parallel with any of those efforts. The eval pipeline already calls `bot.act(state)` per opponent turn; replacing that bot's implementation language is local to the bots package.
|
||||
|
||||
## Key files (current)
|
||||
|
||||
- `src/coolrl_lost_cities/games/classic/bots/heuristic.py` — pure-Python implementation. ~1160 lines. Fully read and analyzed. Contains:
|
||||
- Module-level helpers `play_action`, `discard_action`, `draw_from_discard_action` (used by tests).
|
||||
- `SafeHeuristicParams` (frozen dataclass, behavioral knobs).
|
||||
- `DerivedHeuristicConfig` (frozen dataclass, derived constants).
|
||||
- `derive_heuristic_config(config, params)` — `lru_cache(maxsize=64)` cached.
|
||||
- `SafeHeuristicBot(LostCitiesPolicy)` — the policy itself, with `_act_card` / `_act_draw` and ~20 helper methods.
|
||||
- `src/coolrl_lost_cities/games/classic/bots/base.py` — `legal_from_obs`, `first_legal`. The fallback path (non-`GameState` input) routes through these. The Cython port must preserve this fallback behavior unchanged.
|
||||
- `src/coolrl_lost_cities/games/classic/bots/registry.py` — `BOT_REGISTRY`, `LOOSE_SAFE_HEURISTIC_PARAMS`, `STRICT_SAFE_HEURISTIC_PARAMS`. Currently constructs `SafeHeuristicBot` from the Python class.
|
||||
- `src/coolrl_lost_cities/games/classic/game.pxd` — typed `GameState` declaration. The bot will consume this through Python attribute access (sufficient — see "Cython tactics" below) and via direct `cpdef` calls (`legal_card_mask`, `legal_draw_mask`, `score_diff`, `can_play_card`, etc.).
|
||||
- `src/coolrl_lost_cities/games/classic/policy.py` — `LostCitiesPolicy` ABC and `PolicyInput` typedef.
|
||||
- `setup.py` — Cython extension list.
|
||||
- `tests/games/classic/test_bots.py` — existing safe-heuristic tests (`test_safe_heuristic_mirror_match_finishes`, `test_safe_heuristic_opponent_value_ignores_hidden_hand`, `test_safe_heuristic_started_expedition_value_ignores_invalid_lower_followup`, `test_safe_heuristic_draws_playable_discard_instead_of_deck`, `test_safe_heuristic_can_draw_discard_to_deny_opponent_when_losing`, `test_safe_heuristic_classic_self_play_opens_expeditions`, `test_safe_heuristic_avoids_opening_weak_fifth_color`, `test_safe_heuristic_prefers_followup_on_started_expedition`, `test_safe_heuristic_avoids_unopened_discard_draw_after_four_opens`).
|
||||
|
||||
## New files
|
||||
|
||||
- `src/coolrl_lost_cities/games/classic/bots/heuristic.pyx` — single Cython file containing the ported `SafeHeuristicBot` class plus `cdef` helpers. One file (not split per variant): the variants differ only in `SafeHeuristicParams` values, not in code.
|
||||
- `src/coolrl_lost_cities/games/classic/bots/heuristic.pxd` — minimal typed declarations for the bot class and its hot helper signatures, so future extensions (or other Cython modules) can import typed entry points. Optional in Step 1; required in Step 3 if Cython call-site overhead from Python attribute lookup dominates measurement.
|
||||
- `tests/games/classic/test_safe_heuristic_equivalence.py` — equivalence tests (see "Equivalence strategy").
|
||||
|
||||
## Files to touch
|
||||
|
||||
- `src/coolrl_lost_cities/games/classic/bots/heuristic.py` — **keep as a thin wrapper** that re-exports `SafeHeuristicBot`, `SafeHeuristicParams`, `DerivedHeuristicConfig`, `derive_heuristic_config`, `play_action`, `discard_action`, `draw_from_discard_action`, `PLAY_OR_DISCARD_ACTIONS_PER_SLOT`, `DRAW_FROM_DECK_ACTION` from `heuristic.pyx`. This preserves all existing imports (`tests/games/classic/test_bots.py` imports `from coolrl_lost_cities.games.classic.bots.heuristic import draw_from_discard_action`). Do not delete the file. Do not duplicate logic.
|
||||
- `src/coolrl_lost_cities/games/classic/bots/registry.py` — no functional change. Imports continue to resolve through the `heuristic.py` shim.
|
||||
- `setup.py` — add the new extension:
|
||||
```python
|
||||
Extension(
|
||||
"coolrl_lost_cities.games.classic.bots.heuristic",
|
||||
["src/coolrl_lost_cities/games/classic/bots/heuristic.pyx"],
|
||||
),
|
||||
```
|
||||
|
||||
## Equivalence strategy
|
||||
|
||||
Behavior drift is the dominant risk. Equivalence tests gate the merge of every step.
|
||||
|
||||
### Test design
|
||||
|
||||
`tests/games/classic/test_safe_heuristic_equivalence.py` builds a fixed corpus of seeded `LostCitiesConfig` × seed pairs, plays full games to terminal, and asserts the Cython and Python bots emit byte-identical action ids at every turn.
|
||||
|
||||
Corpus:
|
||||
|
||||
- **Configs**: at minimum, the default `LostCitiesConfig()` from `tests/games/classic/helpers.py`. Add the small-tier config used by `test_safe_heuristic_started_expedition_value_ignores_invalid_lower_followup` if it exists, plus a tier where `n_handshakes == 0` (exercises the `state.config.n_handshakes <= 0` early return in `_best_handshake_play`).
|
||||
- **Seeds**: `range(0, 50)` per config. Fifty full games per config gives broad coverage of decision branches without making CI slow.
|
||||
- **Variants**: all three (`SafeHeuristicBot()`, `SafeHeuristicBot(LOOSE_SAFE_HEURISTIC_PARAMS)`, `SafeHeuristicBot(STRICT_SAFE_HEURISTIC_PARAMS)`). Loose and strict shift only the parameter dataclass, but each must be tested independently because their thresholds drive different branches in `_should_open_expedition`, `_best_handshake_play`, and the visible-draw support logic.
|
||||
|
||||
### Driver
|
||||
|
||||
Use a self-play harness (mirror match) that constructs two clone `GameState` instances per seed: one driven by the Python bot, one by the Cython bot. At each turn, both bots receive the same `GameState`. The test asserts:
|
||||
|
||||
```python
|
||||
assert cython_action == python_action, (turn, seed, variant, state_summary)
|
||||
```
|
||||
|
||||
If either bot diverges mid-game, the states diverge and subsequent comparisons are meaningless — so abort the comparison on first disagreement and report `(seed, variant, turn, action_py, action_cy)`.
|
||||
|
||||
### Branch coverage targets
|
||||
|
||||
The corpus must hit all decision paths. Each of these branches must be exercised by at least one (variant, seed, turn) tuple:
|
||||
|
||||
- `_act_card` → handshake play taken.
|
||||
- `_act_card` → number play on a started expedition.
|
||||
- `_act_card` → speculative open (`speculative_open=True`).
|
||||
- `_act_card` → strong open (`strong_open=True`).
|
||||
- `_act_card` → exceptional open (`exceptional_open=True`).
|
||||
- `_act_card` → single-late open (`single_late_open=True`).
|
||||
- `_act_card` → forced open path (no expeditions started, no normal play).
|
||||
- `_act_card` → discard path with `unusable_discard_bonus`.
|
||||
- `_act_card` → discard path with `discard_safety_bonus`.
|
||||
- `_act_draw` → draw from deck.
|
||||
- `_act_draw` → draw from discard (handshake top card).
|
||||
- `_act_draw` → draw from discard (numeric top card).
|
||||
- `_act_draw` → unopened-color discard penalty triggers (`opened_colors >= 4`).
|
||||
- `_visible_draw_value` → exceptional-support short-circuit.
|
||||
- Loose vs strict thresholds: at least one (seed, turn) where loose opens an expedition that strict declines, and one where strict's `late_open_block_threshold` blocks an opening that loose would take.
|
||||
|
||||
A coverage assertion at the end of the test sweep records which of the above branches fired (via a side-channel counter inside a debug build of the bot, or by analyzing the action stream); CI fails if any branch was not hit.
|
||||
|
||||
### Floating-point determinism
|
||||
|
||||
The bot uses Python `float` arithmetic with `max(...)` over `(value, action)` tuples. In Python, ties break by `action` ordering (since the value is the first element of the tuple). The Cython port **must use the same tie-break order**: stable `(value, action)` comparison, where ties prefer the lower-numbered action that came first in the iteration. Implementation: collect `(value, action)` pairs in the same iteration order as Python, then linear-scan for the maximum, replacing only on strict `>`. This matches Python's `max` semantics on a list-of-tuples.
|
||||
|
||||
All arithmetic must be `double` (matches Python `float`). Avoid `cdivision` floating-point divergence — there is essentially no division in the hot path, but `0.25 * number_sum` and similar must remain `double`.
|
||||
|
||||
## Cython-ization tactics
|
||||
|
||||
The hot loops are inside `_best_number_play`, `_best_discard`, `_visible_draw_value`, `_color_commitment`, `_public_color_commitment_for_opponent`, `_bonus_potential`, and `_opening_plan_value`. They iterate over hand cards (≤ ~10) and colors (≤ 6) per call, and the bot is called once per opponent turn (~30 turns/game × N games × N opponents).
|
||||
|
||||
### Tactics, in priority order
|
||||
|
||||
1. **Type the bot class as `cdef class SafeHeuristicBot`** with `cdef` helpers. Convert all `_<name>` methods to `cdef inline double _<name>(...)` (or `cdef int` for action ids) where possible. Keep `act` as `cpdef` so the Python registry can construct and call it.
|
||||
2. **Type all hot locals** as `int` / `double` / `Py_ssize_t`. The dominant cost in the Python version is per-card list comprehensions building intermediate `list[Card]` objects; replace with explicit indexed loops over `state.hands[player]` and short fixed-size `cdef double` accumulators. No intermediate Python lists in the hot path.
|
||||
3. **Cache `state.config` accessors once per `act` call.** `state.config.expedition_penalty`, `bonus_threshold`, `bonus_amount`, `n_colors`, `min_rank`, `max_rank`, `n_handshakes`, `n_ranks`, `hand_size`, `deck_size` — pull all of these into typed locals at the top of `act`.
|
||||
4. **Cache `derive_heuristic_config` per `act`.** Already cached via `lru_cache`, but each lookup re-hashes the config. Pull the resulting `DerivedHeuristicConfig` into a typed local; access its fields once.
|
||||
5. **Avoid Python-object operations in inner loops.** `Card` objects expose `color`, `rank`, `is_handshake`, `numeric_value(min_rank)`. In the Cython port, when the Python `Card` object is unavoidable (it is part of the public `GameState.hands` shape), bind its three relevant attributes to typed locals at the top of each loop iteration. Do not call `card.numeric_value(min_rank)` inside conditions repeatedly — compute it once per card per iteration.
|
||||
6. **Direct typed access where possible.** `GameState.score_diff(player)` is `cpdef int` and `can_play_card` is a Python wrapper around `can_play_encoded_card` (`cpdef bint`). Where the bot calls `state.can_play_card(player, card)` for a `Card` object, the Cython port can encode the card once (`color * cards_per_color + (rank - min_rank)` per game.pyx encoding) and call `state.can_play_encoded_card(player, encoded)` directly. **Verify the encoding formula against `game.pyx` `_encode_card` before relying on it; otherwise keep the Python `can_play_card` call and accept the small overhead.**
|
||||
7. **`cpdef list legal_card_mask(self)` returns a Python list of `bool`.** Type the binding as `list legal` in the Cython bot, and index it with typed `Py_ssize_t`. If profiling shows mask access dominates, convert callers to use `unified_legal_mask_np()` and read it as a typed numpy view — but only after Step 4 measurements show this is needed.
|
||||
8. **Inline tiny helpers.** `_num`, `_late_penalty`, `_new_color_open_penalty` are one-liners. Inline them with `cdef inline`.
|
||||
9. **Compiler directives.** Use the same directives as `setup.py` already applies project-wide: `boundscheck=False`, `wraparound=False`, `cdivision=True`, `initializedcheck=False`. The bot file inherits these from `setup.py`'s `compiler_directives` block — do not need to override per-file.
|
||||
|
||||
### What NOT to optimize
|
||||
|
||||
- Do not rewrite `derive_heuristic_config` to drop its `lru_cache`. It is called once per `act` and the cache hit is fast.
|
||||
- Do not replace `GameState.hands[player]` with raw int-array access. The `Card` Python object boundary is the public API, and crossing it is what `_card_obj` already costs in Cython. Keep the boundary; just don't re-cross it inside tight loops.
|
||||
- Do not memoize across calls. The bot is stateless; each `act` call gets a fresh state.
|
||||
|
||||
## Risks and mitigations
|
||||
|
||||
- **Behavior drift introduced silently.** Mitigation: equivalence tests (above) run in CI on every PR. They are the gate. Benchmarks are not run until equivalence is green.
|
||||
- **Floating-point ordering difference between Python and Cython.** Mitigation: explicit `(value, action)` linear-scan max with strict `>` comparison; identical iteration order. Equivalence tests would surface any divergence.
|
||||
- **`Card` object identity vs equality.** The Python code uses `is not` comparisons (`if other is not card`, `if followup is not card`, `if card is not exclude_card`). Mitigation: in Cython, preserve `is`-comparison semantics by using object-pointer identity (`PyObject*` compare) or by passing the slot index instead of the Card. Slot-index passing is preferred — it sidesteps identity altogether.
|
||||
- **`max(candidates)[1]` Python tuple semantics.** Mitigation: documented above. Strict `>` linear scan.
|
||||
- **Cached `derive_heuristic_config` shared between Python and Cython invocations.** Mitigation: the Cython port must import `derive_heuristic_config` from the same module path so the LRU cache is shared (or rebuild a Cython-side cache keyed identically). Easiest: keep `derive_heuristic_config` and `DerivedHeuristicConfig` in the same `.pyx` and re-export from the `.py` shim.
|
||||
- **Build-system regression.** Adding an extension that fails to compile in CI on a system without the right toolchain. Mitigation: the project already builds three `.pyx` files (`game`, `cfr_math`, `encoding`, `traversal`); the toolchain is established. Add the new extension and run `uv run python -c "import coolrl_lost_cities.games.classic.bots.heuristic"` locally before pushing.
|
||||
- **Performance regression on small tiers.** The bot may not actually be the bottleneck in micro-config evaluation. Mitigation: benchmarks are run on `default.yaml` (the contract surface), not on smoke configs.
|
||||
|
||||
## Implementation steps (each independently mergeable)
|
||||
|
||||
### Step 1: scaffolding (no behavior change)
|
||||
|
||||
- Add `bots/heuristic.pyx` containing the entire current `bots/heuristic.py` body verbatim, with no `cdef` or typing changes — just renamed.
|
||||
- Convert `bots/heuristic.py` to a re-export shim (`from .heuristic_impl import *` style) — but to avoid a circular import via the `.pyx` module name, keep both as `heuristic.{py,pyx}` is not workable. Instead, name the new file `heuristic_cy.pyx` (Cython compiles to `heuristic_cy`), and have `heuristic.py` import the public symbols from `heuristic_cy`. Update `setup.py` accordingly.
|
||||
- **Decision point:** if Cython supports a `.pyx` shadowing a `.py` in the same package, prefer that for a cleaner import path. Otherwise use the `heuristic_cy` naming. Resolve at Step 1 implementation time and update this plan inline.
|
||||
- Run all existing tests to confirm zero behavior change.
|
||||
- Merge.
|
||||
|
||||
### Step 2: type the base `SafeHeuristicBot`
|
||||
|
||||
- Convert `SafeHeuristicBot` to `cdef class`. Add typed locals to the hot helpers listed under "Cython tactics". Keep the data classes (`SafeHeuristicParams`, `DerivedHeuristicConfig`) as Python frozen dataclasses — they are not hot.
|
||||
- Run the equivalence test suite (built in Step 4 — but for this step, run the existing `test_safe_heuristic_*` tests as a proxy gate; the full equivalence suite lands in Step 4).
|
||||
- Merge.
|
||||
|
||||
### Step 3: type `cdef inline` helpers and inner loops
|
||||
|
||||
- Inline `_num`, `_late_penalty`, `_new_color_open_penalty`. Convert `_color_commitment`, `_public_color_commitment_for_opponent`, `_bonus_potential`, `_opening_plan_value`, `_visible_draw_value`, `_visible_open_support_value`, `_visible_number_can_help_open` to typed `cdef double` / `cdef bint` helpers.
|
||||
- Replace list comprehensions with indexed loops over `state.hands[player]`, accumulating into typed scalars.
|
||||
- Run equivalence suite. Merge only if green.
|
||||
|
||||
### Step 4: equivalence test suite
|
||||
|
||||
- Add `tests/games/classic/test_safe_heuristic_equivalence.py` with the corpus and branch-coverage assertions described under "Equivalence strategy".
|
||||
- Wire into `pytest -q tests/games/classic/test_safe_heuristic_equivalence.py`. Confirm the test passes against the Step 3 build. If it surfaces drift, fix Step 3 before continuing.
|
||||
- Merge.
|
||||
- **Note on ordering vs Step 2/3:** ideally Step 4 lands before Step 2 so equivalence is a green gate the entire time. Recommended order: 1 → 4 (against the verbatim port, must be a no-op pass) → 2 → 3.
|
||||
|
||||
### Step 5: variant validation
|
||||
|
||||
- Loose and strict differ only by params, but exercise their threshold differences explicitly. Confirm `LOOSE_SAFE_HEURISTIC_PARAMS` and `STRICT_SAFE_HEURISTIC_PARAMS` in `registry.py` flow through the Cython implementation unchanged.
|
||||
- Confirm `NoisyPolicy(SafeHeuristicBot(), RandomBot(seed))` still works (`NoisyPolicy` lives in `registry.py` and wraps the bot; the wrapper does not need porting).
|
||||
- Add a small targeted test that constructs all three variants and runs one full game each.
|
||||
- Merge.
|
||||
|
||||
### Step 6: benchmark and documentation
|
||||
|
||||
- Benchmark protocol: run `configs/deep_cfr/default.yaml` evaluation (full opponent suite, default `evaluation.games`) once on the pre-port commit and once on the post-port commit, on the same `home` machine, with no other GPU load. Record:
|
||||
- `eval/safe_heuristic/elapsed_seconds`, `eval/safe_heuristic/opponent_act_seconds`,
|
||||
- `eval/safe_heuristic_loose/elapsed_seconds`, `eval/safe_heuristic_loose/opponent_act_seconds`,
|
||||
- `eval/safe_heuristic_strict/elapsed_seconds`, `eval/safe_heuristic_strict/opponent_act_seconds`,
|
||||
- `evaluation_seconds` for the iteration.
|
||||
- Confirm success criteria 2 and 3.
|
||||
- Add a date-stamped subsection to `docs/performance.md` recording results, methodology, and any branch-coverage gaps surfaced during equivalence testing.
|
||||
- Merge.
|
||||
|
||||
## Out-of-scope follow-ups (do not start)
|
||||
|
||||
- Porting `RandomBot`, `PassiveDiscardBot`, or `NoisyPolicy` to Cython. Their measured cost is small. Re-evaluate only if a future profile shows them on the critical path.
|
||||
- Replacing `Card` Python objects with raw int encoding inside the bot interface. This would require touching `LostCitiesPolicy` and `PolicyInput` more broadly, expanding scope.
|
||||
- Caching across `act` calls (e.g., remembering opened-color counts). The bot is stateless by design; introducing state risks correctness for marginal speedup.
|
||||
- TensorRT / `torch.compile` on policy networks during eval. Tracked separately under the inference-server plan and `docs/performance.md` § "Post-A Optimization Calculus".
|
||||
- Vectorizing eval to compute many game states' bot actions simultaneously. Substantial rework of the eval driver; out of scope here.
|
||||
@@ -0,0 +1,268 @@
|
||||
# Plan: Revisit `torch.compile` on Deep CFR Trainer (and Inference Server Forward)
|
||||
|
||||
**Status:** Conditional — gated on a model-size bump. Do not implement against the current `default.yaml` (512 hidden / 3 layers); a regression has already been measured at that size.
|
||||
**Owner:** Codex
|
||||
**Background:** See `docs/performance.md`:
|
||||
- "Experiments → `torch.compile` on trainer networks (2026-05-07, regression)" — the prior attempt regressed iter time by +4.8% on the small default model. Implementation is preserved on branch `experiments/torch-compile`.
|
||||
- "Post-A Optimization Calculus (forward-looking, 2026-05-07)" — argues compile/TensorRT become meaningful only after model size grows out of the dispatch-bound regime and/or eval density rises. This plan honors that sequencing.
|
||||
|
||||
This plan is the follow-up referenced in step 3 of "Recommended sequencing".
|
||||
|
||||
## Goal
|
||||
|
||||
Re-enable `torch.compile` on the Deep CFR trainer's networks at a model size where kernel work amortizes compile dispatch overhead, and (optionally, secondary) on the inference server's eval-mode forward path introduced by `docs/plans/batched_traversal_inference_server.md`. The goal is iter-time speedup with **no learning-curve drift**.
|
||||
|
||||
## Non-goals
|
||||
|
||||
- Do not enable compile on the current 512-hidden / 3-layer `default.yaml`. The regression is already measured.
|
||||
- Do not compile the Cython traversal call-site networks (workers' CPU networks under `inference_backend: local`). Their per-call shape is `batch_size == 1` and dispatch-dominated; compile cannot help and recompilation triggers are higher-risk.
|
||||
- Do not introduce TensorRT here. TensorRT is a separate work item also gated on the inference server (see `docs/performance.md` § "Tooling split").
|
||||
- Do not change network architectures. This plan picks up whatever larger model the project settles on in step 2 of the post-A sequencing.
|
||||
- Do not change the public CLI surface. Compile toggles via config only.
|
||||
- Do not enable `mode="max-autotune"` by default. It is opt-in for benchmarking.
|
||||
|
||||
## Success criteria
|
||||
|
||||
1. **Model-size precondition met.** The active `default.yaml` (or the targeted variant) has `network.hidden_size ≥ 1024` *or* `network.num_layers ≥ 6`, *or* an architecture (e.g. `color_shared` with non-trivial `color_attention_layers`) whose per-call forward time exceeds ~150 μs at the trainer's training batch size on the target GPU. If neither condition holds, this plan is **not merged**; the branch is parked.
|
||||
2. **Iter-time improvement.** With `compile.trainer.enabled: true` on the chosen larger model and the same seed, the 1000-iter projection improves by at least **5%** vs the no-compile baseline on the same machine (measured on `home`). Eval and checkpointing should be disabled for the bench window, matching the protocol used in the 2026-05-07 experiment.
|
||||
3. **No learning-curve drift.** Over at least 100 iterations with `compile.trainer.enabled: true` vs `false` (same seed, same config), the eval win-rate trajectories against `random` and `safe_heuristic` are within seed noise. If trajectories visibly diverge, the plan does not ship even if iter time improves.
|
||||
4. **No checkpoint-format break.** Checkpoints saved with compile enabled must load cleanly when compile is disabled, and vice versa. (Handled via `_clean_state_dict()`; see Risks.)
|
||||
5. **No multiprocessing-worker break.** Whether `inference_backend` is `local` or `server`, traversal workers must continue to receive uncompiled `state_dict`s without `_orig_mod.` prefixes.
|
||||
6. **(Secondary) Inference-server forward.** If step B below is taken, the server's `policy_network_seconds` decreases by at least 20% at the chosen model size, with no traversal-path correctness regression. If step B does not produce a measurable win, it is left disabled and the plan still ships with step A only.
|
||||
7. All existing tests pass. Lint clean (`uv run ruff check .`). The relevant Deep CFR test subset passes.
|
||||
|
||||
## Why this is gated on a larger model (explicit dependency)
|
||||
|
||||
The 2026-05-07 experiment already isolated the failure mode: at 512 hidden / 3 layers / ReLU MLP, the per-call forward is ~80 μs at `batch_size=1` and ~90 μs in the bs≤256 plateau (see "GPU forward profiling for batched traversal", `docs/performance.md`). Compile dispatch overhead, plus the cudagraph/decomposition path, costs more than the kernel-fusion benefit at that size. The trainer's training batches *are* larger than 1 (`optimization.advantage_batch_size` / `optimization.strategy_batch_size`), so the dispatch-dominated regime ends sooner there than for traversal — but the prior measurement shows it still does not pay at 512×3.
|
||||
|
||||
Compile becomes meaningful when one of these is true:
|
||||
|
||||
- **Wider/deeper MLP.** `hidden_size ≥ 1024` *or* `num_layers ≥ 6`. At that point the per-layer matmul is large enough that fused epilogue gains (linear+activation) and reduced Python-level dispatch overhead exceed compile's per-call cost.
|
||||
- **Non-trivial attention.** `color_shared` with `color_attention_layers ≥ 2` introduces `nn.TransformerEncoderLayer`, which is one of the architectures where `torch.compile` reliably wins (LayerNorm + softmax + GEMM fusion).
|
||||
- **Backward+optimizer fusion.** Compile can fuse parts of the optimizer step on larger models. The trainer phase compiles forward+backward+optimizer together; this is where the biggest absolute wins live, but only above the dispatch-bound floor.
|
||||
|
||||
**Threshold rule (hard gate):** do not merge this plan unless at least one of the following is true on the config it targets:
|
||||
|
||||
- `network.hidden_size ≥ 1024`, or
|
||||
- `network.num_layers ≥ 6`, or
|
||||
- `network.kind == "color_shared"` with `color_attention_layers ≥ 2`.
|
||||
|
||||
If none hold, park the branch and re-evaluate at the next model-size bump. This explicit gate is required by the "Post-A Optimization Calculus" reasoning in `docs/performance.md`.
|
||||
|
||||
## Two compile targets (treated separately)
|
||||
|
||||
These have different dispatch profiles and different failure modes. Each is independently mergeable; step A is the primary objective, step B is secondary.
|
||||
|
||||
### A. Trainer networks (forward + backward + optimizer)
|
||||
|
||||
Wraps `advantage_networks[player]` and `strategy_network` at trainer construction time with `torch.compile(...)`. The compiled wrapper sits in front of the optimization loop in `trainer.py`. Workers (under `inference_backend: local`) and the inference server (under `inference_backend: server`) keep using uncompiled networks built from cleaned `state_dict`s.
|
||||
|
||||
Why this target: the trainer's optimization step has fixed batch shapes (no dynamic shapes), runs many steps per iteration, and exercises forward+backward+optimizer — the regime where compile pays best when the kernel is large enough.
|
||||
|
||||
### B. Inference-server forward (eval-mode, no grad)
|
||||
|
||||
Optional follow-up. After Option A from the batched-traversal plan lands and is benchmarked, optionally compile the server's `model.forward` path with `mode="reduce-overhead"` or `mode="default"`. The server already uses `eval()` + `inference_mode()`. Batches are bounded by `max_batch` (default 256) but variable in size up to that — this introduces a dynamic-shape concern (see Risks).
|
||||
|
||||
Why this target is secondary: per `docs/performance.md` § "Why compile / TensorRT are negligible *today* but become meaningful later", the inference-server forward share of an iter is <1% post-A on the small model. Even a 50% forward speedup is iter-level negligible until the model grows. At that point compile and TensorRT compete for the same role; this plan covers compile, and the TensorRT plan (separate, future) covers TensorRT.
|
||||
|
||||
## Compile mode selection
|
||||
|
||||
Start with `mode="default"` for both targets. Evaluate `mode="reduce-overhead"` and `mode="max-autotune"` only after a baseline number is in.
|
||||
|
||||
- **`default`** — safe. Lowest compile time. First measurement.
|
||||
- **`reduce-overhead`** — uses CUDA graphs. Helps small-batch regimes by amortizing launch overhead. **Incompatible with multiprocessing in non-trivial ways**: the captured graph holds CUDA stream/state from the capturing process. The trainer process is the only place this mode would be used (workers do not compile); confirm the trainer's compiled call is not entangled with worker process spawn (it should not be — workers are already spawned with cleaned `state_dict`s before any compiled call). Use `reduce-overhead` only on the trainer phase, never on a target that crosses a `mp.spawn` boundary.
|
||||
- **`max-autotune`** — autotunes kernel selection. Long compile time (minutes). Only worth it on a stable, frozen model config that will be trained for many hours. Run as an A/B against `default` after the headline result is established.
|
||||
|
||||
For the inference server (target B), `default` is the only safe mode initially. `reduce-overhead` requires fixed-shape inputs; the server's batch dimension varies up to `max_batch`, so cudagraphs would either recompile per-shape or require pre-padding to `max_batch`. Treat that as a separate experiment.
|
||||
|
||||
## Key files
|
||||
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py` — owns trainer-side networks, optimizer, and the train loop. This is where target A's `torch.compile(...)` calls go, plus the `_clean_state_dict()` helper and the `_orig_mod`-routed `load_state_dict` shim. Reference implementation lives on branch `experiments/torch-compile` (commit `05cc02a`).
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/networks.py` — network classes. Not modified by this plan; compile wraps the constructed module from outside.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/workers.py` — multiprocessing traversal worker entry. Reconstructs CPU networks from `state_dict`s. Must keep receiving cleaned (no-`_orig_mod.`) state dicts.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py` — created by the batched-traversal plan. Target B wraps `model.forward` here. If the inference-server plan has not landed, target B is deferred.
|
||||
- `src/coolrl_lost_cities/games/classic/deep_cfr/config.py` — extend with a `CompileConfig` block.
|
||||
|
||||
## Reference implementation (branch `experiments/torch-compile`)
|
||||
|
||||
The prior implementation (commit `05cc02a` on `experiments/torch-compile`) provides the working pattern. Specifically:
|
||||
|
||||
- `_clean_state_dict()` helper strips the `_orig_mod.` prefix that `torch.compile`'s `OptimizedModule` adds to `state_dict()` keys. All checkpoint writes and all weight pushes to multiprocessing workers / the inference server go through this helper.
|
||||
- `load_state_dict` on a compiled wrapper is routed via `module._orig_mod.load_state_dict(...)` so it accepts clean state dicts (saved checkpoints have no prefix).
|
||||
|
||||
This plan **reuses these helpers verbatim**. Cherry-pick the trainer changes from that branch as the starting point and rebase onto current `main`.
|
||||
|
||||
## Config schema
|
||||
|
||||
Extend `config.py`:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class CompileTrainerConfig:
|
||||
enabled: bool = False
|
||||
mode: Literal["default", "reduce-overhead", "max-autotune"] = "default"
|
||||
fullgraph: bool = False # safer to start False; may flip True after stabilization
|
||||
dynamic: bool = False # trainer batches are fixed shape
|
||||
|
||||
@dataclass
|
||||
class CompileInferenceServerConfig:
|
||||
enabled: bool = False
|
||||
mode: Literal["default", "reduce-overhead"] = "default"
|
||||
fullgraph: bool = False
|
||||
dynamic: bool = True # batch dim varies up to max_batch
|
||||
|
||||
@dataclass
|
||||
class CompileConfig:
|
||||
trainer: CompileTrainerConfig = field(default_factory=CompileTrainerConfig)
|
||||
inference_server: CompileInferenceServerConfig = field(default_factory=CompileInferenceServerConfig)
|
||||
```
|
||||
|
||||
In `default.yaml`, leave both `enabled: false`. Add an explicit benchmarking variant (e.g. `configs/deep_cfr/default_compile.yaml`) that flips `compile.trainer.enabled: true` and bumps the model size to meet the threshold rule.
|
||||
|
||||
## Checkpoint and weight-sync handling (the `_orig_mod.` prefix)
|
||||
|
||||
This is the single most error-prone part of the work. The contract is:
|
||||
|
||||
1. **Saved checkpoints are always uncompiled-shaped.** Before saving, run `_clean_state_dict(module.state_dict())` to strip `_orig_mod.`. Files saved with compile enabled must load fine when compile is disabled.
|
||||
2. **Loaded checkpoints are always uncompiled-shaped.** When loading into a compiled wrapper, route through `module._orig_mod.load_state_dict(clean_dict)`.
|
||||
3. **Weights pushed to multiprocessing traversal workers** (under `inference_backend: local`) are always uncompiled-shaped. Workers reconstruct an uncompiled `DeepCFRMLP` and call `load_state_dict` on it.
|
||||
4. **Weights pushed to the inference server** (under `inference_backend: server`) are always uncompiled-shaped. The server may *separately* wrap its loaded model with `torch.compile` (target B). The wire format is uncompiled.
|
||||
5. **Resume from a checkpoint trained with a different `compile.*` setting** must work in either direction. Test this explicitly.
|
||||
|
||||
These rules mean compile is purely a runtime optimization; it never appears in any persisted artifact and never crosses a process boundary.
|
||||
|
||||
## Trainer integration (target A)
|
||||
|
||||
In `trainer.py`, after constructing `advantage_networks[player]` and `strategy_network` and binding their optimizers:
|
||||
|
||||
```python
|
||||
if config.compile.trainer.enabled:
|
||||
compile_kwargs = dict(
|
||||
mode=config.compile.trainer.mode,
|
||||
fullgraph=config.compile.trainer.fullgraph,
|
||||
dynamic=config.compile.trainer.dynamic,
|
||||
)
|
||||
self._advantage_networks = [
|
||||
torch.compile(net, **compile_kwargs) for net in self._advantage_networks
|
||||
]
|
||||
self._strategy_network = torch.compile(self._strategy_network, **compile_kwargs)
|
||||
```
|
||||
|
||||
Caveats:
|
||||
|
||||
- The training step calls `loss.backward()` and `optimizer.step()`. Compile fuses across the forward, but backward and optimizer paths are separately traced via Dynamo's autograd hooks. Confirm by inspecting `torch._dynamo.config.cache_size_limit` is not being hit (recompilation noise).
|
||||
- All `state_dict()` writes (checkpoint, weight push, eval snapshot) go through `_clean_state_dict()`.
|
||||
- All `load_state_dict()` reads route through the `_orig_mod` shim when targeting a compiled wrapper.
|
||||
|
||||
## Inference-server integration (target B, optional)
|
||||
|
||||
In `inference_server.py`, after the server process loads the model on `device` and sets `eval()` + `inference_mode()`:
|
||||
|
||||
```python
|
||||
if compile_cfg.inference_server.enabled:
|
||||
self._model = torch.compile(
|
||||
self._model,
|
||||
mode=compile_cfg.inference_server.mode,
|
||||
fullgraph=compile_cfg.inference_server.fullgraph,
|
||||
dynamic=compile_cfg.inference_server.dynamic,
|
||||
)
|
||||
```
|
||||
|
||||
Caveats:
|
||||
|
||||
- Server already runs in a child process spawned via `torch.multiprocessing`. Compile is invoked **inside the child**, never in the parent. This avoids the cudagraph-across-spawn issue.
|
||||
- The server batches request rows into a `[k, input_dim]` tensor where `k ∈ [1, max_batch]`. Set `dynamic=True` so Dynamo does not recompile per `k`. If recompilation is observed, pad batches up to a small set of bucket sizes (e.g. powers of two) before forward.
|
||||
- `mode="reduce-overhead"` is **not safe** here without bucketing; cudagraphs require fixed shapes.
|
||||
- Weight sync: when the trainer pushes a fresh `state_dict`, the server must apply it via `self._model._orig_mod.load_state_dict(clean_dict)` if compile is enabled, else `self._model.load_state_dict(clean_dict)`.
|
||||
|
||||
## Implementation steps (ordered, each independently mergeable)
|
||||
|
||||
### Step 0: precondition check
|
||||
|
||||
- Before doing anything, confirm the active model config (or the variant being targeted) meets the threshold rule above. If it does not, **stop**. This plan does not ship against the small default.
|
||||
|
||||
### Step 1: cherry-pick reference implementation (target A)
|
||||
|
||||
- Cherry-pick `experiments/torch-compile` (commit `05cc02a`) onto a fresh branch. Resolve any conflicts against current `main`.
|
||||
- Move the unconditional `torch.compile(...)` calls behind `config.compile.trainer.enabled`. Default is `false`.
|
||||
- Add the `CompileConfig` and `CompileTrainerConfig` dataclasses to `config.py`. Wire `--set compile.trainer.enabled=true` through.
|
||||
- Confirm `_clean_state_dict()` and the `_orig_mod`-routed `load_state_dict` paths are correctly invoked on every checkpoint save, every checkpoint load, every weight push to workers, and (post-A) every weight push to the inference server.
|
||||
|
||||
### Step 2: tests for state-dict round-tripping
|
||||
|
||||
- Unit test: build a `DeepCFRMLP`, wrap with `torch.compile`, call `_clean_state_dict(model.state_dict())`, build a fresh uncompiled `DeepCFRMLP`, `load_state_dict` from the clean dict, assert parameter equality.
|
||||
- Unit test: build a compiled model, `load_state_dict` from a clean dict, assert no error and parameters match.
|
||||
- Integration test: short training run with `compile.trainer.enabled=true`, save checkpoint, resume training with `compile.trainer.enabled=false`, assert no parameter mismatch on load.
|
||||
|
||||
### Step 3: precondition guardrail
|
||||
|
||||
- In `trainer.py`, when `compile.trainer.enabled=true`, log a warning at startup if the model fails the threshold rule (hidden < 1024 and layers < 6 and not color-attention). Do not error — operators may want to bench against the threshold — but make the regression risk explicit.
|
||||
|
||||
### Step 4: bench (target A)
|
||||
|
||||
See "Bench plan" below. This is the gating step. If results do not clear success criterion 2 and 3, do not merge target A.
|
||||
|
||||
### Step 5: trainer compile shipped behind config flag
|
||||
|
||||
- Once Step 4 is green, enable `compile.trainer.enabled=true` in the larger-model config that became the new `default.yaml` (or the dedicated `default_compile.yaml`). Do not flip it on for any small-model config.
|
||||
|
||||
### Step 6: (optional) inference-server forward compile (target B)
|
||||
|
||||
- Only if the batched-traversal-inference server has landed (per `docs/plans/batched_traversal_inference_server.md`).
|
||||
- Add `CompileInferenceServerConfig` wiring. Add `torch.compile(...)` invocation inside the server process. Route weight sync through `_orig_mod` when enabled.
|
||||
- Add a small inference-server bench script (or extend `scripts/bench_inference_backend.py` if it has landed) to measure server-side `policy_network_seconds` with and without compile at the chosen larger model size.
|
||||
- Ship only if criterion 6 is met. Otherwise leave disabled.
|
||||
|
||||
### Step 7: documentation
|
||||
|
||||
- Add a date-stamped experiment subsection to `docs/performance.md` recording the bench result at the new model size, mirroring the existing 2026-05-07 entry's structure.
|
||||
- If the result is a regression at the chosen model size, document it and park the branch again with a note about which threshold to revisit.
|
||||
|
||||
## Bench plan
|
||||
|
||||
Mirror the 2026-05-07 protocol exactly so results are comparable to the prior data point.
|
||||
|
||||
- **Hardware**: `home` (6-core + RTX 3090). Confirm baseline numbers on `remote` separately if/when applicable.
|
||||
- **Config**: the larger-model variant that meets the threshold rule. Disable eval and checkpointing for the bench window (`--set checkpoint.save_latest=false --set checkpoint.save_every=0` and an eval-disabling override).
|
||||
- **Length**: at least 8 iterations measured, with iteration 1 dropped as compile warm-up. Replicate the 2026-05-07 table format:
|
||||
|
||||
| | iter mean | 1000-iter projection |
|
||||
| --- | ---: | ---: |
|
||||
| Baseline (no compile) | … | … |
|
||||
| `torch.compile` trainer (mode=default) | … | … |
|
||||
| Effect | … | … |
|
||||
|
||||
- **A/B**: same seed both runs. Same machine, same GPU, no concurrent jobs (per AGENTS.md).
|
||||
- **Drift check (criterion 3)**: run a separate 100-iter A/B with `eval.eval_every=25` enabled, same seed, and overlay the eval win-rate trajectories from `metrics.jsonl`. If the trajectories diverge beyond seed noise, target A does not ship even if the iter-time A/B looked good.
|
||||
- **Mode sweep**: only after `mode="default"` clears the bar, also bench `mode="reduce-overhead"` (trainer only) and `mode="max-autotune"` (trainer only). Record both.
|
||||
|
||||
For target B, bench `policy_network_seconds` from the inference-server side with and without compile at the chosen larger model size. The report does not need to wait for a learning-curve A/B — the server is eval-mode only and pushes uncompiled weights every iter, so it cannot drift training.
|
||||
|
||||
## Risks
|
||||
|
||||
- **Recompilation triggers.** If `torch._dynamo.config.cache_size_limit` is hit, the compiled wrapper falls back to eager and the run silently regresses. Mitigation: log Dynamo recompile events at startup; fail loudly if recompile count exceeds a small threshold during the bench window. Trainer batches are fixed-shape (`optimization.*_batch_size`), so this should not trigger for target A. For target B, set `dynamic=True` and watch for recompilation across batch sizes.
|
||||
- **`_orig_mod.` prefix leaking into checkpoints or worker state_dicts.** Most likely bug. Mitigation: tests in Step 2 explicitly guard this. The reference implementation already handles it.
|
||||
- **`reduce-overhead` × multiprocessing.** Cudagraphs in `reduce-overhead` mode capture CUDA stream/context state. They are safe inside the trainer process (no spawn after compile) and inside the inference-server child process (compile happens after spawn). They are **not** safe if compile is invoked before `mp.spawn` and the child inherits compiled state. The plan only invokes compile after spawn boundaries.
|
||||
- **Dynamic shapes on inference server (target B).** Server batches vary in size up to `max_batch`. Mitigation: `dynamic=True`. Fallback: pad to fixed bucket sizes.
|
||||
- **Backward path not benefiting.** Compile's biggest theoretical wins on the trainer phase come from fusing forward+backward+optimizer. In practice on simple MLPs the optimizer is already fused (e.g. `torch.optim.Adam(..., fused=True)` if available); compile may add little on top. Mitigation: bench is the answer — if iter time does not move 5%, the plan does not merge. This is a real possibility.
|
||||
- **CFR variable-length traversal does not feed compile.** Confirmed: traversal call sites use `batch_size == 1` and run on CPU workers (or via the inference server, which is target B not target A). Target A only sees the trainer optimization batches, which are fixed-shape. So dynamic-shape concerns do not apply to target A.
|
||||
- **GPU contention with eval.** Eval already runs on the trainer's device. Compile increases peak memory during compile/autotune phases (especially `max-autotune`). Mitigation: compile happens once per process at startup; eval runs after warm-up. Bench with eval disabled to isolate iter time, then re-enable for the drift check.
|
||||
- **AMP interaction.** `run.use_amp` is currently a no-op (per `docs/performance.md`). If AMP is implemented before this plan ships, re-bench compile under AMP — the two interact and prior numbers do not transfer.
|
||||
|
||||
## Decision tree
|
||||
|
||||
- **Threshold rule fails on the active model.** Park the branch. Do not merge. Re-evaluate at the next model-size bump.
|
||||
- **Threshold rule passes; bench shows ≥5% iter improvement and no learning-curve drift.** Ship target A behind the config flag, enable on the larger-model config.
|
||||
- **Threshold rule passes; bench shows iter improvement but learning-curve drift.** Do not ship. Investigate determinism (compile mode, autograd path, optimizer fusion). Likely cause: optimizer-fusion change altering update order. If cause cannot be isolated within reasonable effort, park.
|
||||
- **Threshold rule passes; bench shows <5% iter improvement.** Do not ship. The cost (state-dict gymnastics, recompile risk, checkpoint compatibility surface) is not justified by sub-5% wins. Park.
|
||||
- **Threshold rule passes; bench shows regression.** Park the branch with a documented `docs/performance.md` entry. Note the model size at which the regression was observed and the next threshold to try.
|
||||
- **Target A shipped; target B (inference-server compile) bench shows <20% server-forward improvement.** Leave target B disabled. Revisit alongside the TensorRT plan, since they target the same surface.
|
||||
|
||||
## Out-of-scope
|
||||
|
||||
- **TensorRT.** Separate plan. TensorRT and `torch.compile` on the inference-server forward are alternatives covering the same surface (target B). This plan covers compile only; the TensorRT plan covers TensorRT.
|
||||
- **Compiling the Cython traversal call-site networks.** Per the 2026-05-07 experiment, this is dispatch-bound and unhelped by compile. Workers under `inference_backend: local` will keep using uncompiled CPU networks indefinitely.
|
||||
- **Compiling the encoding path.** Encoding is numpy/Cython, not a `torch.nn.Module`.
|
||||
- **`torch.export` / AOTInductor.** Out of scope until the model architecture is frozen and an offline-compiled artifact is operationally needed.
|
||||
- **Multi-GPU.** Not relevant to the current single-GPU setup.
|
||||
- **`fullgraph=True`.** Default `fullgraph=False` until stabilization. Promoting to `fullgraph=True` is a follow-up after one full training run completes cleanly with compile enabled.
|
||||
Reference in New Issue
Block a user