coolguy/coorl-lost-cities
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Document Option A bench result, post-A calculus, plans, and cost reports

Browse Source 
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:
coolguy
2026-05-07 20:05:38 +09:00
co-authored by Claude Opus 4.7
parent a7ab94e096
commit befe29fc57
8 changed files with 2267 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/plans/torch_compile.md
+268
View File
@@ -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.