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/performance.md
+304
View File
@@ -386,3 +386,307 @@ bs=1 to 0.34 μs at bs=256 (>230×). The available supply of ~368 states per
traversal sits comfortably in the bs=64–256 range where μs/call plateaus near
90 μs. End-to-end gain will be bounded by encoding and worker-GPU coordination
overhead, but the GPU forward is not the limiter once batching is in place.
## Batched Traversal Inference: Design Decision (2026-05-07)
Three structural options were considered for Priority #5:
- **A. Central inference server.** Workers stay in multiprocessing and
reconstruct nothing on GPU. A separate server process owns the model, batches
incoming policy requests across workers, runs GPU forward, returns logits.
Worker traversal logic and the Cython recursion are untouched.
- **B. Per-worker batching.** Each worker interleaves multiple traversals
internally to form its own batches. GPU process count = worker count, so model
copies and GPU contention scale with workers. Batching efficiency is bounded
by per-worker in-flight count.
- **C. Single-process vectorized traversal.** Drop multiprocessing entirely.
Main process runs N traversals lockstep with explicit recursion stacks,
forming a natural batch dimension across traversal instances. Existing
recursive traversal can be kept and a new `traversal/batched.{py,pyx}` added
as a parallel backend gated by config; existing code is not modified.
### Decision: A
Reasons:
- **Hardware fit dominates.** A central inference server keeps multiprocessing,
so all available CPU cores stay productive on game logic. C is single-process,
so on a 32-core remote machine with a weak GPU it wastes 31 cores while the
weak GPU caps batching gains; A is strictly better there. On a 6-core / RTX
3090 box A and C are competitive but uncertain — C only wins when GPU forward
is the dominant share of traversal, and game logic in CFR traversal is not
negligible.
- **C is not the "ultimate" answer on multi-core machines.** A truly maximal
design would combine C's batched GPU forward with `nogil` threaded game
logic, which is strictly more complex than C alone. Plain C, by being
single-process, gives up CPU parallelism that the existing multiprocessing
path already exploits.
- **A is mostly additive.** New modules: `inference_server.py`,
`inference_client.py`, shared-memory tensor pool, weight-sync hook. Existing
touches are small: worker policy call site (one line), worker spawn/teardown
(server start/stop), trainer (periodic weight push). Cython traversal
recursion, game engine, replay/training paths are unchanged.
- **The hard part is IPC tuning, not code volume.** Latency budget vs GPU
forward, weight-staleness window, backpressure, and shared-memory tensor
layout. Code is small; the design surface is concentrated in one place.
### IPC: what crosses the process boundary
Only the encoded policy input and its response cross IPC:
- Forward request: encoded state vector, ~365 floats ≈ 1.5KB.
- Forward response: action logits, ~22 floats ≈ 88 bytes.
Game state, traversal recursion stack, event log, CFR regret/strategy
accumulators, and chance-node sampling history all stay inside the worker
process. The policy network consumes a flat encoded state (`input_dim=365`),
so the server needs no game-tree context to answer a request.
The replay-buffer write path (workers shipping collected regret/strategy
samples to the trainer) is separate, already exists today, and is reflected in
`memory_add_seconds` ≈ 1.25s/iter; A does not add to it.
### IPC mechanism: multiprocessing + shared memory
- **Big payload (state, logits)**: shared-memory tensors. Either
`torch.multiprocessing` with `tensor.share_memory_()` and a pre-allocated
buffer pool indexed by slot id, or `multiprocessing.shared_memory.SharedMemory`
with manual slot management. Pickle is bypassed for the data itself.
- **Control messages (slot index, request id)**: small `Queue`. Pickle still
happens here but only for ints/tuples, which is sub-microsecond and
negligible against ~90 μs GPU forward.
- The naive path (`multiprocessing.Queue(tensor)` with default pickle) is the
one that is slow and is what causes the "Python IPC is slow" reputation.
With shared memory, multiprocessing IPC is effectively on par with thread
shared-memory access for tensor traffic.
### Why not Cython `nogil` + threading instead of multiprocessing
Threading would avoid IPC entirely, but it requires the game-engine hot path
to be genuinely `nogil`-clean — no Python objects touched anywhere on the path.
Whether the existing Cython traversal qualifies is unknown and likely no:
auditing and migrating it to be fully `nogil`-clean is a substantial,
high-risk change to existing code, contradicting A's "mostly additive"
property. Additional drawbacks: a single segfault kills all threads;
multi-threaded CUDA usage has subtle context-sharing pitfalls; tooling and
prior art are weaker than for the multiprocessing pattern. Revisit only after
free-threaded Python (PEP 703) stabilizes or if a future profile shows the
shared-memory IPC is itself the limiter.
### Implementation plan
1. Prototype A on the 6-core / 3090 host with a single worker: validate
end-to-end correctness and measure IPC round-trip latency vs GPU forward.
2. Scale to multiple workers; tune `batch_window_us`, `max_batch`, and
`sync_every` (weight push frequency).
3. Deploy to the 32-core / weak-GPU remote and confirm CPU-side scaling holds
and the weak GPU is still the right place to keep the model.
4. Defer C. Re-evaluate only if A's measurements show GPU forward is no longer
on the critical path and game-logic CPU cost dominates — in that case the
right next step is C with `nogil` threading, not plain C.
## Post-A Optimization Calculus (forward-looking, 2026-05-07)
Once Option A lands, the bottleneck shape changes. This section records the
expected sequencing for follow-up work. It is forward-looking and has not been
measured yet — verify against bench numbers after A is benchmarked.
### Why compile / TensorRT are negligible *today* but become meaningful later
Today (small model: 3-layer, 512 hidden):
- `torch.compile` on the trainer's networks already regressed (see the
2026-05-07 experiment above). The model is too small for kernel fusion to
beat compile dispatch overhead.
- `torch.compile` / TensorRT on the inference-server forward (post-A) would
shave ~30–50% off ~90μs/call → ~50–70μs/call. With forward share of an iter
reduced to <1% by A's batching, the iter-level multiplier is ~1.00–1.01×.
Negligible.
Two compounding shifts can flip this:
1. **Larger model.** Going from 512 hidden / 3 layers to ~1024 hidden /
~6 layers pushes the forward call out of dispatch-bound territory into
kernel-bound territory. Compile fusion and TensorRT both deliver real
1.5–2× on the forward call itself once the kernel is large enough to
amortize launch overhead. Forward share of iter time also rebalances upward
because per-call time scales with FLOPs while batching gain is fixed.
2. **Denser, larger evaluation.** Moving toward `eval_every: 5` and
`evaluation.games: 1000` makes evaluation about half of iteration wall-clock
(see the amortized eval table earlier in this doc). Eval is pure inference,
so TensorRT on the inference-server's forward path applies directly.
When both shifts happen together, an illustrative future iter (rough order of
magnitude only):
| Configuration | Iter time (rough) |
| --- | ---: |
| Today (small model, eval_every=25) | 17.85s |
| + A (batched traversal inference) | ~14s |
| + larger model (≈4× FLOPs), no compile/TRT | ~50s |
| + dense eval (eval_every=5, games=1000) | ~70s |
| + compile (trainer) + TensorRT (inference) | ~45s |
That last row is where compile/TensorRT contributes ~1.5× iter — the same
tooling that is iter-neutral today. The numbers above are illustrative; real
ratios depend on model size, kernel autotune outcomes, and the eval-vs-train
balance.
### Tooling split
- **TensorRT**: applies only to inference (no backward). Targets:
- inference-server forward in traversal,
- inference-server forward in evaluation.
Both are served by the same A-era server, so a single TensorRT integration
covers both.
- **`torch.compile`**: applies to trainer's advantage/strategy training
(forward+backward+optimizer). The 2026-05-07 regression on a small model
does **not** generalize — it must be re-measured on whatever larger model
config we settle on. Do not conclude "compile is bad" from the small-model
data point.
### Recommended sequencing
Do this in order. Skipping ahead is the failure mode that creates misleading
"compile/TRT didn't help" data.
1. **Now**: benchmark A (`scripts/bench_inference_backend.py`) and confirm the
`local` vs `server` multipliers on `home` and `remote`. Validate the iter
1.2–1.3× / traversal 1.5–2× working estimate.
2. **Next**: experiment with a larger network config. Measure compute vs
learning-curve trade-off with the existing toolchain (no compile/TRT yet).
This step decides the model size that future optimizations target.
3. **Then**: re-measure `torch.compile` on the trainer at the chosen model
size. The earlier regression was size-bound; expect a different result.
4. **Then**: integrate TensorRT into the inference server (covers traversal
and eval forward simultaneously). Bound the gain by the post-step-2
`policy_network_seconds` share, not the headline TensorRT speedup.
5. **In parallel with 2–4**: if denser eval is operationally useful, raise
`evaluation.games` and lower `evaluation.eval_every`. This step does not
require code changes but sharply increases the value of step 4.
Out of scope until A bench numbers are in: Option C, `nogil` threading, async
inference client, compiled encoding.
## Option A Bench Result and Structural Ceiling (2026-05-07)
Option A (`traversal.inference_backend: server`) was implemented and
benchmarked. **Result: regression. A is deferred. `default.yaml` stays on
`local`. The implementation is preserved behind the flag for future revisit.**
### Bench numbers
`scripts/bench_inference_backend.py --device cuda --iterations 5 --warmup 1`,
RTX 3090, after a per-call IPC fix (replaced `multiprocessing.Manager()`
queues/events with spawn-context primitives, slot reuse per worker batch,
shared memory confirmed in use for state/response payloads).
| Backend | iter | traversal | adv_train | strat_train | mem_add | batch_tensor |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| `local` | 16.75s | 10.81s | 3.91s | 2.02s | 0.95s | 3.56s |
| `server` | 57.61s | 51.61s | 3.96s | 2.02s | 0.50s | 3.57s |
| Speedup | 0.29× | 0.21× | 0.99× | 1.00× | 1.89× | 1.00× |
Raw: `runs/bench/2026-05-07_193335_inference_backend/results.json`.
Training and eval phases are unchanged (as expected — A only touches the
traversal forward path). The regression is contained in `traversal_seconds`,
which is ~5× worse.
### Diagnosis
The server emits per-flush batch stats. **Mean batch size: ~7.2–7.9, max 8.**
This is the structural ceiling, not a tunable misconfiguration:
- Traversal recursion is **sync-blocking** at the policy call site. Each
worker has at most one in-flight policy request at a time.
- In-flight requests at the server ≤ `num_workers` = 8.
- The server further splits each batch by `(network_kind, network_index)`,
so the actual GPU forward group size is roughly half of that — about 4
rows per group.
Per-state cost at this realized batch size, from the GPU profile table:
| Realized batch | μs/state |
| ---: | ---: |
| 1 | 80.07 |
| 4 | 20.30 |
| 8 (extrapolated) | ~12 |
| 64 | 1.46 |
| 256 | 0.34 |
So the GPU is doing ~12μs per state instead of the projected ~1.5μs at
bs=64. The IPC round-trip per call (queue post + server scheduler + event
wakeup, even with shared-memory payload) is on the order of hundreds of μs
per call, which exceeds both the local CPU forward (~80–200μs at bs=1 on
this small MLP) and the marginal GPU gain. Net: per-call cost roughly
doubles or triples, compounded across ~205k calls/iter, gives the observed
5× traversal regression.
`batch_window_us` and `max_batch` tuning cannot escape this ceiling —
there are simply not 64 concurrent in-flight requests to coalesce when only
8 workers are blocking-sync.
### What this means for the headline GPU profile (`scripts/profile_gpu_forward.py`)
The earlier "230× speedup at bs=256" is a **per-state GPU forward**
microbenchmark, not an end-to-end traversal speedup. Realizing that gain
requires *actually feeding the GPU* with bs=64+ batches. Sync-blocking
multi-worker traversal cannot do that. Reaching the bs=64 regime needs
either:
- Per-worker traversal interleaving (worker advances `worker_chunk_size`
traversals concurrently, suspending at each policy call — Option B
shape), which requires turning Cython traversal recursion into a
resumable state machine. Same scope as a partial Option C, localized to
worker scope.
- Option C proper (single-process vectorized traversal).
Both require restructuring traversal. Option A's "additive, no traversal
changes" property turned out to also mean "cannot drive the batch size up."
### Why deferring A (not deleting) is the right call
- The plumbing (server process, shared-memory client, weight sync, config
flag) is complete and tested. Re-enabling is a config flip.
- The fundamental issue at this model size is that **GPU forward time is
too small to amortize IPC overhead** at any realistic batch size we can
drive without restructuring traversal. Bigger model changes that
arithmetic; the same plumbing then becomes useful.
- The `mem_add_seconds` row showed a real 1.89× win, suggesting the
shared-memory replay-write path adopted along the way is worth keeping
even with `local` backend. (Confirm separately; this is a side effect.)
### Re-enable A when one of these holds
1. **Model grows** to ~1024 hidden / ~6 layers (compile/TRT discussion
above). Forward time scales with FLOPs while IPC overhead is fixed; at
some point IPC becomes a small fraction.
2. **Per-worker interleaved traversal** ships (Option B-shape refactor).
Drives realized batch toward 64 and reclaims the profile table's gains.
3. **Eval becomes the dominant phase** (`eval_every: 5`,
`evaluation.games: 1000`). Eval is not sync-blocking traversal; it is
already batch_size=64 in eval code. The same inference server can
serve eval directly without the worker-side ceiling.
### Free-threaded Python (3.13t) note
Free-threaded Python + Cython `nogil` would let many threads (well above
core count) run game logic concurrently in one process, with shared memory
and no IPC. With ~64 threads sync-blocking on policy calls, the server
would naturally see bs=64. **In principle this is the cleanest endpoint.**
In practice as of early 2026:
- Free-threaded Python is an opt-in build (`python3.13t`), still
experimental, with measurable single-thread overhead.
- PyTorch's free-threaded compatibility is partial.
- Cython `nogil`-cleanliness audit on the existing game engine is still
required and was the original reason `nogil` threading was deferred in
the design decision above.
- No project-level adoption pressure on `python3.13t` today.
So free-threaded Python does change the architectural answer, but it does
not unblock A *now*. Track the ecosystem; revisit when (a) `python3.13t`
becomes mainstream or (b) the Cython engine is `nogil`-cleaned for other
reasons.
docs/plans/amp_trainer.md
+259
View File
@@ -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.
docs/plans/batched_traversal_inference_server.md
+209
View File
@@ -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.
docs/plans/cython_safe_heuristic_bots.md
+210
View File
@@ -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.
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.
docs/reports/cost_cython_nogil_audit_2026-05-07.md
+444
View File
@@ -0,0 +1,444 @@
# Cython `nogil`-cleanliness Audit
Date: 2026-05-07
Scope: `game.pyx`, `traversal.pyx`, `encoding.pyx`, `cfr_math.pyx` (+ matching `.pxd`)
Trigger: free-threaded Python (3.13t) + threaded traversal as an alternative
to multiprocessing for Deep CFR. See
`docs/performance.md` "Option A Bench Result and Structural Ceiling" and
"Free-threaded Python (3.13t) note".
---
## 한 줄 결론
**Medium effort.** 게임 엔진(`game.pyx`)과 보조 산술(`cfr_math.pyx`,
`encoding.pyx` core encode 함수)은 이미 거의 nogil-clean이다. 진짜
blocker는 한 곳에 모여 있다: **`traversal.pyx`의 `_traverse` 재귀 본체가
PyTorch forward, NumPy 배열 할당, Python `TraversalStats`/`TrainingSample`
객체 mutation, f-string, Python list/dict bucket 누적**을 모두 직접
한다는 것. 이걸 두 단계로 분리(순수 C 시뮬레이션 + Python에서 후처리)하지
않으면 `with nogil:` 블록을 의미 있게 키울 수 없다. 비용은 traversal 한
파일의 mid-scale 리팩터(추정 1–2주, 회귀 위험 큼) + 작은 게임 엔진 정리
(1–2일)이다.
---
## 파일별 현황
### 1. `game.pyx` / `game.pxd` — **거의 nogil-clean (small effort)**
게임 상태 데이터는 모두 raw `int*` (C 힙)에 살고, 핵심 동작(`_apply_*_c`,
`_undo_*_c`, `_legal_actions_c`, `_can_play_encoded_card_c`,
`_score_from_summary_c`, `_recompute_score_caches`,
`_swap_deck_cards_c`, `_push_action_c`, `_pop_action_c`)이 전부
`cdef ... noexcept`/`except *` 시그니처에 C 산술만 한다.
이미 nogil-callable 후보 (시그니처에 `nogil` 키워드만 추가하면 되는 것):
- `_is_legal_action_c` (game.pyx:869)
- `_legal_actions_c` (game.pyx:896)
- `_unified_legal_actions_c` (game.pyx:924)
- `_can_play_encoded_card_c` (game.pyx:953)
- `_fill_undo_c` (game.pyx:964)
- `_score_from_summary_c` (game.pyx:1267)
- `_has_any_legal_draw` (game.pyx:1281)
- `_hand_index` / `_expedition_len_index` / `_expedition_index` /
`_discard_index` / `_encode_card` / `_card_color` / `_card_rank`
(game.pyx:1293–1312)
- `_recompute_score_caches` (game.pyx:1229)
부분 GIL 필요 (작은 수정으로 nogil 가능):
- `_apply_action_unchecked_c` / `_apply_card_action` / `_apply_draw_action`
(game.pyx:1008, 1099, 1140) — 본문은 순수 C이지만 `except *`라
exception propagation을 위해 GIL이 필요. nogil 컨텍스트에서 호출하려면
`noexcept` 또는 명시적 `nogil` + `with gil` 예외 블록이 필요. 본문에는
실제로 raise할 곳이 없으므로 시그니처를 `noexcept`로 바꾸는 게 가장 싸다.
단, `_apply_card_action`/`_apply_draw_action`은 invariant를 깨는 입력이
들어와도 silently 진행하게 되므로 호출 전 검증을 강화해야 한다.
- `_apply_action_with_undo_c`, `_push_action_c`, `_pop_action_c`,
`_swap_deck_cards_c`, `_ensure_undo_capacity_c` (game.pyx:1004–1055) —
`_ensure_undo_capacity_c`만 `realloc` 실패 시 `MemoryError`를 raise.
`with gil:` 짧은 블록으로 분리하거나, traversal 진입 시
capacity를 미리 키워두면 nogil-clean하게 만들 수 있다.
- `_undo_*_c` (game.pyx:1161, 1169, 1203) — `raise ValueError("undo
... mismatch")` 가드가 들어 있음. Production 경로에서는 fire되지 않으므로
guard를 `assert` 또는 디버그 빌드 한정으로 빼면 nogil 가능.
완전 GIL 함수 (nogil 변환 비대상; 호출자가 GIL 가진 채로 부른다):
- `__init__`, `_configure`, `from_snapshot`, `to_snapshot`, `validate_invariants`
(Python config object/dict touch, dataclass, `Counter`, yaml, etc.)
- 모든 property: `phase`, `deck`, `hands`, `expeditions`, `discards` —
Python list-of-Card 생성. 전부 reporting/serialization용이라 hot path
아님.
- `clone()` (game.pyx:550) — `GameState(self.config)` 생성자 호출이
Python object instantiation. nogil 안에서 부르려면 별도 `cdef
GameState _clone_into(self, GameState dst) nogil` 같은 C-only fast clone을
추가해야 한다 (전부 `memcpy`이므로 trivial하지만 새 entry point 필요).
- `to_unified_action`, `hand_slots`, `sort_hand` 등 Python 인터페이스 —
hot path 아님.
요약: game.pyx는 **시그니처 정리 + 작은 helper 추가**로 hot path를 통째로
`nogil` 안에 넣을 수 있다. 게임 엔진 자체는 큰 비용이 아니다.
### 2. `cfr_math.pyx` / `cfr_math.pxd` — **이미 nogil-clean (zero effort)**
`regret_matching_c`, `normalize_legal_policy_c`, `sample_policy_c` 모두
`noexcept` + 순수 C 산술. `nogil` 키워드만 시그니처에 추가하면 끝.
(file:1–89). 파이썬 wrapper 3개(`regret_matching` 등, file:92–146)는
NumPy 인터페이스라 GIL 필요하지만 hot path 아님 — traversal은 이미 C
함수 직접 호출 (traversal.pyx:11, `from ... cimport regret_matching_c`).
### 3. `encoding.pyx` / `encoding.pxd` — **거의 nogil-clean (small effort)**
C-only encoders (전부 `noexcept`/`except -1`, raw float buffer 출력):
- `_base_input_dim_c`, `input_dim_c`, `_input_dim_with_flags_c`,
`_numeric_value_c`, `_max_numeric_sum_c`, `_max_score_estimate_c`
(encoding.pyx:12–57)
- `_color_playability_summary_c` (encoding.pyx:60–) — 본문은 순수 C
(state의 C 필드 참조 + `abs()`), nogil 가능.
- `_append_derived_playability_features_c`,
`_append_slot_aware_playability_features_c` (encoding.pyx:156, 209) —
본문은 순수 C 산술. nogil 가능.
- `encode_info_state_c`, `_encode_info_state_with_flags_c`
(encoding.pyx:287–413) — `except -1`로 Python `ValueError`를 raise할 수
있는 두 곳(file:319, 321)이 있지만 둘 다 정적 sanity 체크
(`player < 0`, `action_size > 64`). 호출 전에 검증되면 제거해도 안전.
`abs(state.expedition_penalty)` (encoding.pyx:53, 95): Cython이 `int`에
대해 C `abs`로 lower하므로 nogil-safe. `bool(encoding.derived_playability)`
(file:420, 432) 같은 건 Python wrapper에서만 호출되므로 무관.
요약: `_encode_info_state_with_flags_c`를 `noexcept nogil`로 바꾸고
input validation을 호출자로 옮기면 nogil-clean. 변환 매우 쉬움.
### 4. `traversal.pyx` / `traversal.pxd` — ****진짜 blocker가 모두 여기 있다 (medium-large effort)****
이미 nogil-callable인 helper들:
- `_next_u32`, `_next_double` (traversal.pyx:24, 29) — `noexcept`,
raw uint32 LCG. 사실상 nogil이지만 키워드 빠짐.
- `_sample_policy_from_actions_c` (traversal.pyx:33) — `noexcept`,
raw pointer.
- `_sampling_policy` (traversal.pyx:582) — `noexcept`, raw pointer.
- `_from_unified_action_c`, `_to_unified_action_c` (traversal.pyx:992, 997)
- `_opened_color_count` (traversal.pyx:766)
- `_self_play_bucket` (traversal.pyx:774) — `noexcept`이지만 `len(self.
league_advantage_networks)`를 본다 → Python list `__len__` (PyObject_Size).
이건 GIL 필요. 단순한 fix: 별도 `cdef int _league_size`를 캐싱.
- `_depth_bucket_start` (traversal.pyx:58)
- `random_rollout_value_c` (traversal.pyx:1092) — 본문은 순수 C이지만
`_push_action_c`, `_pop_action_c`, `_legal_actions_c`가 nogil이 되면
자동으로 nogil-callable. raise 두 줄(file:1106, 1108)을 호출 전 검증으로
옮기면 끝.
반면 hot path인 `_traverse` (traversal.pyx:259) 본문에는 다음과 같은
Python-object touch가 깔려 있다 (per-node, per-iteration):
1. **PyTorch forward 호출** — `_policy_from_networks` (file:439) /
`_policy_from_strategy_network` (file:519). NumPy `np.empty`,
`torch.as_tensor`, `networks[player](x)`, `.detach().cpu().numpy().
astype(np.float32)`. 이게 모든 `_policy` 호출(노드당 1회)에서 일어남.
2. **`stats` mutation** — 모든 카운터 증가가 Python attr 접근:
`stats.nodes += 1`, `stats.terminals += 1`, `stats.max_depth_reached`,
`stats.regret_fallback_*` 등 (file:289, 290, 298, 302, 386, 695–752).
3. **f-string + dict bucket** — `_record_endpoint` (file:980), `_record_
fallback_depth_bucket` (file:753): `f"{start}_{start + width - 1}"`,
`stats.endpoint_depth_buckets[key] = ... .get(key, 0) + 1`. Python
string format + dict lookup.
4. **NumPy 배열 할당 per leaf** — `_record_strategy` (file:877), `_record_
advantage` (file:914), `_record_external_advantage` (file:948): `np.empty(
self.action_size, dtype=np.float32)`, `.append(TrainingSample(...))`.
Sample마다 두 개의 작은 NumPy array + dataclass 인스턴스화.
5. **Python list `.append`** — `self.advantage_samples.append(...)`,
`self.strategy_samples.append(...)` (file:903, 937, 969). list의
PyObject reference 갱신은 free-threaded Python에서도 atomic refcount
비용을 추가로 부담한다.
6. **`SafeHeuristicBot.act(state)`** — `_fixed_opponent_action`
(file:633, 652), `_rollout_value` (file:841). Python class
메서드 호출. `safe_heuristic` 옵션 사용 시만 핫.
7. **`league_advantage_networks` indexing** — `_self_play_snapshot_
networks` (file:802), `[-recent_count:]`, `[:max(0, ...)]` slicing
= Python list slicing.
8. **`f"invalid ..."` raises** — game state 검증 실패 시.
`_traverse`는 game state mutation(전부 C struct 통한
`_push_action_c`/`_pop_action_c`)과 위 Python object 작업을 한 함수에서
교차해서 한다. 즉 `with nogil:`로 감쌀 수 있는 자연스러운 chunk가
없다 — recursion 한 단계 안에서 GIL을 ~6번 release/re-acquire해야
하는데, 그 비용이 forward latency보다 크다.
---
## 주요 blocker 카탈로그
### B1. PyTorch forward 호출 (가장 큰 단일 blocker)
```python
# traversal.pyx:473-475
with torch.inference_mode():
x = torch.as_tensor(info_state, dtype=torch.float32, device=self.device).unsqueeze(0)
advantages = networks[player](x).squeeze(0).detach().cpu().numpy().astype(np.float32)
```
- 빈도: 노드당 1회 (~205k/iter, performance.md 참조).
- 변환 난이도: **High (구조 변경 필수)**. 핵심 통찰은: **이걸 nogil 만들
필요 없다.** PyTorch CUDA 호출 자체가 internally GIL을 잠깐 잡지만
`inference_mode` + CUDA dispatch는 잘 알려진 GIL-friendly 영역이다.
진짜 문제는 *traversal recursion이 forward 호출에서 sync-block*해서
배치가 안 모이는 것 (performance.md "Option B-shape refactor"). nogil로
단일 thread를 빠르게 만들기보다 **traversal을 resumable state machine
으로 깨고 N개 thread를 띄워 동시에 sync-block시키면**, free-threaded
Python 하에서 batch=N forward로 자연 합쳐진다. 즉 nogil-cleaning은
Option B/C와 같은 작업의 일부이지 독립 작업이 아니다.
- 권고: 이 blocker는 nogil audit 단독으로 고치지 말고, "traversal을
state-machine으로 해체" 작업 안에 묶는다.
### B2. Python `TraversalStats` attribute mutation (전 노드 핫)
```python
# traversal.pyx:289-294
stats.nodes += 1
if depth > stats.max_depth_reached:
stats.max_depth_reached = depth
if self.has_max_nodes and stats.nodes >= self.max_nodes:
stats.node_limit_cutoffs += 1
```
- 빈도: 매 노드. 합쳐서 노드당 5–15회 attr access.
- 변환 난이도: **Low–Medium**. `TraversalStats`를 `cdef class`로 바꾸고
필드를 `cdef public long long`로 선언하면 attr access가 C struct field
store가 된다. 단, `stats.regret_fallback_depth_buckets` 같은 dict
필드는 별도로 처리(아래 B3).
- 위치: traversal.pyx:289, 290, 293, 298, 302, 331, 386, 695–752, 853, 856,
912, 946, 978, 985–990 + `_record_*` 전체.
### B3. dict bucket + f-string key (depth/color buckets)
```python
# traversal.pyx:986-990, 753-764, 702-705, 723-725, 742-744
key = f"{start}_{start + width - 1}"
stats.endpoint_depth_buckets[key] = stats.endpoint_depth_buckets.get(key, 0) + 1
```
- 빈도: 매 leaf/cutoff/regret-fallback 노드.
- 변환 난이도: **Medium**. dict + str key + format은 nogil 불가. 해법:
- bucket 인덱스로 미리 정해진 정수 array를 쓴다 (`endpoint_depth_bucket_max
/ endpoint_depth_bucket_width + 1` slot의 `cdef long[:]` 또는 raw
int64 array). string key는 마지막 reporting 단계에서만 생성.
- color/opened_color bucket도 모두 5–6개 정해진 슬롯이므로 `cdef
long[5]`로 충분.
### B4. NumPy 배열 + dataclass 인스턴스 per training sample
```python
# traversal.pyx:896-911, 926-945, 959-977
target = np.empty(self.action_size, dtype=np.float32)
legal_mask = np.empty(self.action_size, dtype=np.bool_)
...
self.advantage_samples.append(TrainingSample(info_state=..., target=..., ...))
```
- 빈도: leaf마다 1개 advantage sample + 노드별 strategy sample
(interval-gated).
- 변환 난이도: **Medium-High**. 두 가지 옵션:
- **(a) Buffer pre-allocate**: traverser가 큰 `cdef float[:, ::1]
advantage_targets`, `cdef uint8[:, ::1] advantage_legal`,
`cdef long[:] advantage_iteration` 등을 미리 잡아두고 row index만 늘린다.
drain 시점에 `TrainingSample` Python 객체로 wrap. **추천**.
- **(b)** PyObject 그대로 두고 `with gil:` 짧게 — sample 누적이
노드당 ~1회라 IPC overhead 분석 그대로 적용된다 (작은 hold라도 thread
contention 발생).
- 추가 고려: `info_state`(`np.empty(input_dim, dtype=np.float32)`)도 노드당
새 NumPy. buffer-pool 또는 batched encoder로 묶어야 한다.
### B5. PyTorch `state_dict()`-share, league list slicing
```python
# traversal.pyx:802-817
candidates = self.league_advantage_networks[-recent_count:]
...
candidates = self.league_advantage_networks[:max(0, len(self.league_advantage_networks) - recent_count)]
```
- 빈도: traversal 진입 시 한 번 (`traverse`에서 미리 픽), 재귀 안에서는
`active_self_play_networks`만 본다. 따라서 cold path. 변환 불필요.
### B6. `SafeHeuristicBot.act(state)` — Python bot
```python
# traversal.pyx:633, 652, 841
return int(self.safe_heuristic_opponent_bot.act(state))
```
- 빈도: `opponent_policy=safe_heuristic` 또는 `cutoff_rollout_policy=
safe_heuristic`일 때만. 현 default는 self_play_league + score_diff
cutoff (per memory의 opponent_policy_network_divergence note + AGENTS).
- 변환 난이도: **Medium-High** (Python class 전체를 cython화). 현 default
config에서는 핫 아님 — 시도하지 않는 게 합리.
### B7. `len(self.league_advantage_networks)`
```python
# traversal.pyx:782, 783, 803, 806, 811, 814
recent_count = min(len(self.league_advantage_networks), self.self_play_recent_window)
```
- 빈도: `_self_play_bucket`가 traversal 진입에 한 번, `_self_play_snapshot_
networks`가 한 번. 노드당이 아님 → cold path. 무시 가능 (단, 두 함수가
`_traverse` 안에서 직접 불리지 않음을 확인했음, file:244–251).
### B8. `_apply_action_unchecked_c` `except *`
게임 엔진 쪽 game.pyx:1008. 본문에 raise 없음 → `noexcept`로 강등하면
`_traverse`의 `state._push_action_c` (game.pyx:1029, `except *`) 호출도
`noexcept`로 만들 수 있다. 단, `_ensure_undo_capacity_c`의 `MemoryError`만
별도 처리 필요.
### B9. Recursion이 그 자체로 `_traverse` (cdef method `except *`)
`_traverse`는 `cdef float ... except *` (file:259). nogil로 만들려면
재귀 호출도 nogil 컨텍스트여야 하고, 모든 파이썬 touch가 제거되어야 한다.
즉 **B1–B4가 전부 해결되기 전엔 `_traverse` 본체를 `nogil`로 못 만든다.**
---
## 작업 단계 (안전한 순서)
1. **단계 0 — 측정 인프라.** Cython annotate (`cython -a`)를 빌드 스크립트에
추가. `.html`에서 노란/빨간 줄 = Python interaction. 반복적으로 본다.
2. **단계 1 — 무비용 청소 (1–2일):**
- `cfr_math.pyx`의 3개 C 함수에 `nogil` 키워드 추가.
- `encoding.pyx`의 `_encode_info_state_with_flags_c`와 모든 helper의
검증을 호출자로 옮기고 `noexcept nogil`로.
- `game.pyx`의 `_legal_actions_c`, `_unified_legal_actions_c`,
`_can_play_encoded_card_c`, `_score_from_summary_c`,
`_has_any_legal_draw`, 모든 `_*_index`/`_card_*` 함수에 `nogil` 추가.
- 회귀 테스트: `uv run pytest -q`.
3. **단계 2 — 게임 엔진 mutation을 nogil로 (2–3일):**
- `_apply_card_action`, `_apply_draw_action`, `_apply_action_unchecked_c`
를 `noexcept`로 강등 (호출 전 legality check가 이미 `_traverse`에서
수행되므로 안전).
- `_undo_*_c`의 `ValueError` mismatch 가드를 debug 빌드 한정 (`IF
DEBUG:` 컴파일 디렉티브 또는 release 시 제거).
- `_ensure_undo_capacity_c`: traversal 진입 시점에 한 번 큰 capacity로
`realloc`해두고, hot path의 `_push_action_c`는 capacity 체크만
(`assert undo_stack_len < undo_stack_capacity` debug only)하게 분리.
- 결과: `_push_action_c`/`_pop_action_c`/`_swap_deck_cards_c` 모두
`nogil`.
4. **단계 3 — TraversalStats를 cdef class로 (3–5일):**
- 모든 정수 카운터를 `cdef public long long` 필드로.
- depth bucket / color bucket dict들을 fixed-size `cdef long[N]` array로
교체하고 reporting 단계에서만 dict로 변환.
- `_record_endpoint`, `_record_fallback_depth_bucket`,
`_record_regret_matching_decision`을 `noexcept nogil`로 다시 작성.
- 회귀 테스트: `metrics.jsonl`의 모든 키가 동일한 값으로 나오는지 비교.
5. **단계 4 — Sample buffer pre-allocate (3–5일):**
- traverser에 `cdef float[:, ::1] advantage_targets`,
`cdef uint8[:, ::1] advantage_legal_masks`,
`cdef float[:, ::1] advantage_info_states`,
`cdef long[:] advantage_iterations`, `cdef int[:] advantage_players`
등을 chunk-grow array로. row index만 nogil에서 늘림.
- `drain_samples()`에서만 GIL 잡고 `TrainingSample` 리스트로 wrap.
- 회귀 테스트: trainer가 받는 sample 분포 동일해야 함.
6. **단계 5 — Forward 호출 분리 (large, 다른 작업과 묶음):**
- `_traverse`를 "forward 직전까지" + "forward 결과 받은 후" 두 구간의
resumable state machine으로 재구성. forward 호출은 외부 batcher가
수행. 이게 Option B/C 본체이므로 별도 design doc 필요.
- 그제서야 `_traverse` 자체를 `nogil`로 선언할 의미가 생긴다.
7. **단계 6 — 검증:**
- `cython -a`로 hot path가 모두 흰색인지 시각 확인.
- micro-bench: 단일 thread에서 traversal 시간이 회귀 없는지.
- free-threaded Python (`uv run --python python3.13t ...`) 또는
`nogil`-제어 micro-bench로 N=2/4/8 thread scaling 확인.
---
## 위험
- **Silent slowdown (GIL re-acquisition)**: `with nogil:` 블록 안에서
Python 객체를 무심코 건드리면 Cython이 `with gil:` 블록을 자동 삽입
(또는 `noexcept nogil` 위반 시 컴파일 에러). 작은 attr touch 하나가
re-acquisition 비용을 부르고, 멀티스레드에선 contention으로 single-thread
대비 더 느려질 수 있다. 검증: `cython -a`가 진실의 원천. 모든 hot 경로가
흰색이어야 함. 추가로 `python -X dev`나 `PYTHONDEVMODE=1`로 thread state
체크.
- **Correctness regression on undo path**: 단계 2의 `_undo_*` 가드 제거가
invariant를 silently 위반시킬 수 있음. 검증: `tests/games/classic/test_
deep_cfr_trainer.py` + `validate_invariants()`를 `--set debug=true` 같은
모드에서 매 100노드마다 호출.
- **Sample buffer overflow**: 단계 4의 chunk-grow가 race 없는지
(single-traverser-per-thread 구조 유지) 확인. 두 thread가 같은 traverser
객체를 공유하면 안 됨.
- **`TraversalStats` API 변경**: `metrics.jsonl` 형식 변경 가능성. 단계
3에서 reporting 어댑터를 명시적으로 보존. 기존 dict 형식과 byte-wise
동일한 테스트 추가.
- **Cython `nogil` + cdef class 라이프타임**: `cdef class` 인스턴스의
refcount는 free-threaded Python에서 atomic이지만 deallocation이 nogil
컨텍스트 안에서 트리거되면 안 됨. 모든 cdef object는 함수 시작에 GIL
잡힌 채로 acquire, nogil 블록 안에서는 raw pointer/struct만 접근.
- **CUDA forward thread-safety**: PyTorch는 같은 device 위 동시
forward에 대해 internal lock을 사용한다. N=64 thread가 동시에 forward를
치면 합쳐주지 않으면 lock contention만 늘 수 있다. 단계 6의 batcher가
필수.
---
## 권고
**현 시점에는 단계 1–3까지만 기회 봐서 진행하고, 단계 4 이상은 보류.**
이유:
1. 단계 1–3은 **나중 단계와 무관하게 단일-thread traversal도 살짝 빠르게**
만들고, `cython -a`상의 visible Python interaction을 줄여 다음 작업의
기반이 된다. 비용 작음(~1주), 회귀 위험 낮음(테스트 충분).
2. 단계 4부터는 free-threaded Python이나 Option B/C 같은 호출자 측 변경이
같이 와야 의미가 있다. 현재 `default.yaml`은 single-process local
backend로 잘 돌고 있고(performance.md), 모델 크기·eval 비중·python3.13t
생태계 모두 트리거가 안 와 있음.
3. 단계 5/6는 Option B-shape refactor와 사실상 같은 작업이므로 **별도
설계 문서가 먼저** 필요하다. nogil audit이 그걸 정당화하는 근거는
되지만 단독 추진 사유는 안 된다.
**다시 볼 트리거** (둘 중 하나라도 만족):
- (a) **Free-threaded Python (3.13t)이 mainstream** 으로 가서 PyTorch 공식
지원이 stable이 되고, `uv`가 3.13t를 1차 시민으로 다룬다.
- (b) **Model이 커진다** — hidden=1024 / depth=6 등으로 forward가
단일 호출 ~수백 μs 영역에 들어가서, traversal 한 번에 한 forward를 GIL
잡고 부르는 게 명백히 bottleneck이 된다.
- (c) **Eval 비중이 dominant**해진다 (`eval_every=5`, `evaluation.games=
1000+`). Eval은 이미 batch-friendly이라 thread pool + nogil game engine만
으로도 큰 win.
위 세 가지가 모두 멀어 보일 때(현 상황)는 단계 1–3만 chip away 하고,
설계 측면에서는 Option B-shape (per-worker interleaved traversal) 쪽이
ROI가 더 높다 (performance.md "Re-enable A when one of these holds" 참조).
### 빠른 우선순위 1순위 (지금 당장 1일)
`cython -a` 빌드 옵션 추가 + cfr_math와 encoding hot path에 `nogil`
키워드만 다는 것. 이건 아무것도 안 깨고 tooling 인프라가 생긴다.
다음 nogil 작업할 때 진단 출발점이 됨.
---
## 참조
- `docs/performance.md`:571 (Option A Bench Result)
- `docs/performance.md`:672 (Free-threaded Python note)
- `src/coolrl_lost_cities/games/classic/game.pxd`
- `src/coolrl_lost_cities/games/classic/game.pyx`:550, 869, 924, 953, 1004,
1099, 1140, 1161, 1229, 1281, 1293
- `src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx`:259 (`_traverse`),
439 (`_policy_from_networks`), 519 (`_policy_from_strategy_network`),
582 (`_sampling_policy`), 753 (`_record_fallback_depth_bucket`),
877 (`_record_strategy`), 914 (`_record_advantage`),
980 (`_record_endpoint`), 1092 (`random_rollout_value_c`)
- `src/coolrl_lost_cities/games/classic/deep_cfr/encoding.pyx`:287, 291,
416, 425
- `src/coolrl_lost_cities/games/classic/deep_cfr/cfr_math.pyx`:5, 37, 68
docs/reports/cost_pytorch_cuda_multithread_2026-05-07.md
+324
View File
@@ -0,0 +1,324 @@
# Cost / Risk Report: PyTorch CUDA from Multiple Threads in One Process
작성일: 2026-05-07
대상 질문: Deep CFR 파이프라인을 multiprocessing → threading(또는 free-threaded
Python)으로 옮길 경우, 같은 프로세스 안에서 여러 스레드가 동시에 PyTorch CUDA
연산을 호출하는 것이 얼마나 비싸고 위험한가?
---
## 한 줄 결론
**Risky — separate-module + 적절한 stream/lock 규율이 있으면 "safe with care",
공유 모듈에 대한 동시 forward/backward/load_state_dict 조합은
serialization 없이는 silent wrong outputs 또는 segfault를 일으킬 수 있음.**
지금 코드 형태(공유 advantage/strategy network + trainer가 매 N step마다
weight push) 그대로 threading으로 옮기면 weight-sync 경계에서 데이터 레이스가
거의 확실히 발생한다.
---
## CUDA 멀티스레드 모델 (PyTorch 2.x 기준, 2026 초 시점)
### 1) Current stream은 thread-local
PyTorch C++ 코어(`c10/cuda/CUDAStream.h`)가 명시적으로 보장한다:
> "the notion of 'current stream for device' is thread local (every OS thread
> has a separate current stream, as one might expect)"
> — [pytorch/c10/cuda/CUDAStream.h](https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDAStream.h)
즉 thread A가 `with torch.cuda.stream(s):` 안에서 띄운 커널은 thread B의
current stream 설정과 독립이다. 하지만 *기본* current stream은 **device의 default
stream**이며, 모든 스레드가 명시적으로 `set_stream`을 하지 않으면 같은 default
stream을 공유한다. 이 경우 GPU 측에서는 직렬화된다(병렬 launch가 안 됨).
### 2) Per-thread default stream(PTDS)는 PyTorch에서 *enable되지 않음*
CUDA driver-level의 `--default-stream per-thread` 컴파일 옵션은 PyTorch
배포 빌드에서 켜져 있지 않다. 관련 트래킹 이슈
[pytorch#25540](https://github.com/pytorch/pytorch/issues/25540) 은 2019년
오픈 이후 미해결 상태로 남아 있고, 결과적으로 "여러 스레드가 default stream을
사용하면 모두 같은 legacy default stream에 직렬화된다"가 현행 동작이다.
([cuda streams run sequentially #59692](https://github.com/pytorch/pytorch/issues/59692),
[default stream is not synchronous #101300](https://github.com/pytorch/pytorch/issues/101300)).
→ **결론: 진짜 GPU 동시성이 필요하면 각 스레드가 명시적으로
`torch.cuda.Stream()`을 만들고 `with torch.cuda.stream(s):` 컨텍스트로
감싸야 한다.** 그렇지 않으면 멀티스레드 = 멀티프로세스 대비 GIL/IPC만
줄어들 뿐 GPU 자체는 직렬화된다.
### 3) Python-level forward/backward 호출의 thread safety
공식 문서/포럼 발언을 종합하면:
- **Tensor 자체는 read-only thread-safe, write는 직렬화 책임이 사용자.**
Edward Yang: "PyTorch underlying C++ library is expected to be thread safe
(although the Tensor object is not thread-safe for multiple writers; you need
to synchronize that yourself)."
([forum #36540](https://discuss.pytorch.org/t/is-pytorch-supposed-to-be-thread-safe/36540))
- **Inference만 한다면 같은 nn.Module 인스턴스를 여러 스레드에서 동시
`forward()` 호출해도 OK** — 단 module state를 mutate하지 않는다는 전제.
([forum #88583](https://discuss.pytorch.org/t/is-inference-thread-safe/88583)).
주의: BatchNorm `train()` 모드, dropout state, lazy module init,
`register_buffer`로 EMA 갱신 같은 건 mutate에 해당.
- **TorchScript/JIT module은 단일 인스턴스를 동시에 forward하면 안 됨**
([pytorch#15210](https://github.com/pytorch/pytorch/issues/15210),
[pytorch#51452](https://github.com/pytorch/pytorch/issues/51452)).
우리는 JIT을 안 쓰지만, `torch.compile`로 래핑된 모듈은 내부 캐시 일관성에서
비슷한 위험을 가질 수 있다
([dev-discuss: compile + multithreading](https://dev-discuss.pytorch.org/t/impact-of-multithreading-and-local-caching-on-torch-compile/2498)).
- **C++ custom 모듈은 여러 스레드에서 동시 호출 시 segfault 보고**
([pytorch#19029](https://github.com/pytorch/pytorch/issues/19029)) — 우리
코드엔 직접 해당하지 않지만, 의존성 라이브러리에 비슷한 게 끼어들 수 있음.
### 4) Backward / autograd
- Autograd 엔진은 **device당 1 스레드**의 worker pool로 backward를 실행한다
([forum #36824](https://discuss.pytorch.org/t/only-1-thread-for-backward/36824),
[autograd notes](https://docs.pytorch.org/docs/stable/notes/autograd.html)).
즉 두 스레드가 *서로 다른* graph에 대해 동시에 `.backward()`를 호출하면
엔진은 receive-side에서 큐잉/locking으로 처리한다 — 코어는 thread-safe.
- 그러나 **두 스레드가 share된 graph 부분을 동시에 backward하면
파괴(graph free) 레이스로 다른 스레드가 crash**한다(같은 forum 답변).
Deep CFR에서 traversal worker는 backward를 안 부르므로 직접 적용은 적지만,
trainer + 동시 inference + replay sample이 동일 텐서를 retain하는 경우
주의해야 한다.
- `torch.no_grad()` / `torch.inference_mode()`는 **thread-local TLS 플래그**다.
스레드별로 따로 켜야 한다. trainer 스레드가 backward 도중에 inference 스레드가
같은 모듈을 forward하더라도 TLS가 분리되므로 모드 자체의 충돌은 없다.
하지만 module state(파라미터)는 공유되므로 아래 weight-update 위험은 그대로다.
### 5) CUDA caching allocator
`CUDACachingAllocator`는 **per-device mutex**로 보호된다
([zdevito guide](https://zdevito.github.io/2022/08/04/cuda-caching-allocator.html),
[CUDACachingAllocator.cpp](https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDACachingAllocator.cpp)).
`cudaEventCreate` 같은 비싼 호출은 EventPool을 둬서 멀티스레드
allocation rate가 높아도 안전하지만, 락 경합은 존재한다. 작은 모델 + 고빈도
forward 패턴(우리 traversal과 정확히 일치)은 allocator lock contention이
스루풋의 실질적 ceiling이 될 수 있다는 점이 위험 항목에 들어간다.
### 6) Free-threaded Python (3.13t / 3.14t) 상황
- PyTorch는 3.13t 빌드(`cp313t` nightly wheels)를 제공하지만 *partial
support*. 트래킹 이슈 [pytorch#130249](https://github.com/pytorch/pytorch/issues/130249).
- DataLoader가 thread-based로 가서 ImageNet iter +74% 같은 사례가 있지만
([Trent Nelson's notes](https://trent.me/articles/pytorch-and-python-free-threading/)),
**"competing Python threads feeding the same CUDA stream still need explicit
synchronization"** — GIL이 사라져도 stream/모듈 동기화 책임은 동일하다.
- 우리 의존성 중 **Cython 확장(`game.pyx`, `traversal.pyx`, `encoding.pyx`,
`cfr_math.pyx`)**은 `nogil` 정합성 감사가 안 된 상태(performance.md "Why not
Cython nogil + threading" 섹션). free-threaded 빌드에서 굴리는 것은
threading 이전 단계의 별도 risk.
---
## 우리 시나리오 매핑
현재 (multiprocessing 기준) GPU touchpoint 위치:
| 위치 | 파일:라인 | 역할 | 현재 격리 수준 |
| --- | --- | --- | --- |
| Trainer forward+backward+optimizer (advantage) | `trainer.py:981-987` | 매 iteration 학습 | 메인 프로세스, 단일 스레드 |
| Trainer forward+backward+optimizer (strategy) | `trainer.py:1028-1034` | 매 iteration 학습 | 메인 프로세스, 단일 스레드 |
| Trainer device 결정 + eval 모듈 deepcopy | `trainer.py:862-870` | eval용 모듈 복제 | 메인 프로세스 |
| Imitation 학습 step | `imitation.py:87-89` | 옵션 path | 메인 프로세스 |
| Policy gradient 학습 step | `policy_gradient.py:69-71` | 옵션 path | 메인 프로세스 |
| Eval forward (batched, inference_mode) | `evaluate.py:211, 249` | eval phase | 메인 프로세스 |
| Inference server forward (batched, inference_mode) | `inference_server.py:231-249` | option A 워커 inference | **별도 spawn 프로세스** |
| Inference server weight load | `inference_server.py:61-67, 173-192` | trainer→server weight push | **별도 프로세스, 큐 직렬화** |
| Worker network reconstruct (CPU only today) | `workers.py:119, 141` | per-worker state_dict 재생성 | 별도 fork/spawn 프로세스 |
핵심 관찰: **모든 GPU 쓰기 경로(optimizer.step, load_state_dict)는 오늘 단일
프로세스 내에서 단일 스레드에 의해 직렬로 발생**한다. Multiprocessing이
"무료로" 보장해주던 격리다.
### Hypothetical threaded design (worst case에 가까운 단순 변환)
| 스레드 | 호출하는 CUDA 연산 | 모듈 공유? | 위험 |
| --- | --- | --- | --- |
| Trainer thread | advantage/strategy forward+backward+`optimizer.step()` | 본인이 own | weight write |
| N traversal threads | advantage/strategy `forward()` (`inference_mode`) | trainer와 share | read-during-write race |
| Eval thread (옵션) | strategy `forward()`, deepcopy | trainer와 share | deepcopy 중 write |
| Weight-sync (없어짐 — 같은 in-process 객체) | `load_state_dict` (만약 league snapshot용으로 남으면) | share | 같음 |
가장 무서운 조합: **trainer가 `optimizer.step()` 도중**(파라미터 텐서가
in-place로 부분 갱신되는 시점)에 traversal 스레드가 같은 파라미터에 대해
`forward()`를 돌리는 경우. PyTorch는 파라미터 텐서 write에 대한 user-side
동기화를 요구하므로(위 §3) **결과는 silent wrong outputs**다 — segfault도
없고 에러도 없고, 그냥 일부 파라미터가 step 전, 일부는 step 후 값으로 섞여
forward가 진행된다. CFR regret 추정 자체에 노이즈를 더해서 학습 발산/품질
저하로 나타난다.
`load_state_dict`도 동일하게 **부분 갱신 + 동시 read** 위험이다.
[forum #224131](https://discuss.pytorch.org/t/thread-safety-between-model-state-dict-and-optimizer-step/224131)
의 동일한 우려가 그대로 적용된다.
---
## 위험 표
| # | 항목 | 종류 | 현실성 (우리 코드) | 증상 |
| --- | --- | --- | --- | --- |
| R1 | 공유 모듈에 대한 forward vs optimizer.step race | silent wrong outputs | **매우 높음** | 학습 noisy, 발산 가능. crash 없음. |
| R2 | 공유 모듈 forward vs `load_state_dict` race | silent wrong outputs | 높음 (league snapshot 갱신 시) | 부분 weight forward, sample contamination |
| R3 | 모든 스레드 default stream → GPU 직렬화 | perf collapse | 매우 높음 (default 동작) | 멀티스레드인데 GPU utilization 그대로 |
| R4 | Allocator lock contention (작은 텐서, 고빈도) | perf collapse | 중간 | trace상 `cudaMalloc`/free 락 대기 증가 |
| R5 | 두 스레드가 공유 autograd graph 일부에 backward | crash / wrong grad | 낮음 (우리는 backward를 trainer만 호출) | graph free race → segfault |
| R6 | `torch.compile`된 모듈 + 멀티스레드 캐시 일관성 | wrong output / 재컴파일 폭주 | 낮음 (현재 main에 compile 없음) | unexpected recompile, dispatcher TLS 충돌 |
| R7 | Cython 확장 `nogil` 미감사 | crash / heap 손상 | 매우 높음 (free-threaded 한정) | segfault, 재현 불가 버그 |
| R8 | CUDA context 상호작용 (단일 context는 OK, 그래도 set_device 누락 시 wrong device) | wrong device error | 낮음 | runtime error |
| R9 | Pinned-memory / `non_blocking=True` H2D 카피 + 다른 스레드의 source tensor mutate | data race on source | 중간 | non-deterministic input bytes ([forum #182924](https://discuss.pytorch.org/t/is-it-safe-to-use-tensor-cuda-non-blocking-true-in-a-thread/182924)) |
| R10 | TLS 누수: `inference_mode`/`no_grad`가 trainer 스레드에서 켜져 있는데 backward 호출 | 잘못된 grad 누락 | 낮음 (코드 명시적으로 with-block 사용) | grad가 안 나와서 학습 정지 |
가장 위험한 건 R1, R2, R3, R7. 이 넷은 "그냥 옮기기"의 직접 결과다.
---
## Mitigations
### M1. Per-thread module copy (제일 안전, 메모리 비용)
각 traversal 스레드가 모듈 deepcopy를 들고 있고, weight-sync 시점에만
`load_state_dict`로 갱신. trainer는 자기 모듈만 만진다.
- 비용: 모듈 사이즈 × N_threads VRAM. 우리 모델(512×3)은 작아서 무시 가능.
- 동기화 지점: weight-sync 때 traversal 스레드를 **잠시 quiesce**해야 안전.
현재 multiprocessing 패턴(`weight_sync_event`)이 그대로 매핑됨.
- 결과: R1, R2 제거. R5도 자동 제거.
### M2. Reader-writer lock on the shared module
Trainer write(step/load_state_dict)는 writer lock, traversal forward는 reader
lock. Pythonic하게는 `threading.RLock` + writer가 모든 reader 종료 대기.
- 비용: writer가 풀릴 때까지 모든 traversal 스레드 정지 → tail latency 증가.
CFR처럼 "조금 stale한 weight도 OK" 알고리즘에선 OK.
- 결과: R1, R2 제거. M1 대비 메모리 절약 / latency 손해.
### M3. Per-thread CUDA stream
각 스레드가 `s = torch.cuda.Stream(); with torch.cuda.stream(s):`로 forward를
감싼다. trainer도 본인 stream에서 step. weight-sync 시 `torch.cuda.synchronize()`
또는 stream event로 ordering 보장.
- 비용: 거의 없음 (stream 객체는 cheap). 코드 변경은 호출 사이트 추가.
- 결과: R3 제거 — 진짜 GPU 동시성 가능. **R1/R2는 해결하지 못함**(stream은
ordering이지 mutual exclusion이 아니다). 반드시 M1 또는 M2와 같이 써야 한다.
### M4. Serialize at boundary (가장 단순, 거의 multiprocessing 효과)
Trainer forward/backward/step 전체를 큰 lock으로 감싸고, traversal forward도
같은 lock으로 감싼다. = 사실상 GIL 흉내.
- 비용: 멀티스레드 의미 사라짐. GPU 사용률 = single thread.
- 결과: 모든 race 제거되지만 free-threaded Python으로 갈 이유가 없어짐.
*threading 도입 자체의 가치가 사라지는 시그널.*
### M5. 현재 inference-server 패턴을 프로세스 → 스레드로 단순 치환
현 `inference_server.py`는 1 spawn process + 워커가 큐로 RequestMessage 송신.
이를 1 server thread + N requester threads로 바꾸는 건 가장 minimal한 변경.
- 모든 GPU 호출이 server thread 1개로 모이므로 R1/R2/R5 자동 회피.
- IPC 비용 사라짐(shared memory가 그냥 메모리). performance.md "Option A
Bench Result"의 IPC 오버헤드(~수백 μs/call) 제거가 가능.
- 단점: server thread가 여전히 single GPU executor → R3와 같은 GPU 직렬화는
유지되나, 그게 *batching의 목표*이기 때문에 해롭지 않음. realized batch
size가 ceiling 8(현재) → free-threaded로 64+ 스레드면 ceiling 64로 올라감 →
performance.md가 기대했던 bs=64 regime에 진입.
이 시나리오에선 **GPU 호출은 여전히 1 스레드만 한다**. 멀티스레드 → CUDA의
복잡도 대부분이 사라진다. *권고 핵심.*
### M6. (보조) `torch.compile` 모듈을 공유 forward 대상에서 제외
향후 trainer가 compile을 다시 쓴다면 inference 스레드들에 노출하지 말 것
([dev-discuss: compile + multithreading](https://dev-discuss.pytorch.org/t/impact-of-multithreading-and-local-caching-on-torch-compile/2498)
의 캐시 일관성 이슈 회피).
---
## 권고
### 지금 위험 수준
코드 그대로(공유 모듈 + 동시 forward + 동시 step) threading 전환하면:
- R1, R2가 거의 확실히 발생 → 학습 품질에 silent regression.
- R3 때문에 GPU 활용도는 multiprocessing 대비 거의 개선 없음.
- Cython 코드(R7) 때문에 free-threaded 빌드에서 segfault 위험.
따라서 "**그냥 threading으로 옮기기**"는 **추천하지 않음**.
### 안전한 경로 (선호 순)
1. **M5 (server-thread pattern) + 기존 multiprocessing 워커 유지**.
가장 작은 변경으로 GPU 호출 스레드를 1개로 묶고, IPC 비용만 줄인다.
하지만 이건 threading으로의 *전환*이 아니라 "Option A의 in-process variant"
라는 점 명심.
2. **Free-threaded Python으로 가야 한다면**: 그 결정은
- (a) Cython 코드의 `nogil` 감사를 끝내고
- (b) 적어도 traversal worker가 thread가 되어 game logic을 동시에 굴리고
- (c) GPU 호출은 M5 패턴으로 단일 server-thread에 위임
세 조건이 동시에 충족될 때만 가치가 있다. performance.md "Free-threaded
Python 노트"의 결론(현재 미적용)과 일치.
3. **공유 모듈을 어쩔 수 없이 여러 스레드에서 호출해야 한다면** M1
(per-thread copy, 모델이 작으니 비용 미미) + M3 (per-thread stream) 조합이
안전. 단순히 lock(M2/M4)만 걸면 GPU 활용도는 single-thread와 동일해진다.
### 안전 검증 체크리스트 (전환 전)
- [ ] Module 공유 여부를 모든 forward 호출 사이트에 대해 표로 만들기.
"이 forward는 trainer가 weight write하는 모듈인가?"가 yes면 M1/M2 필수.
- [ ] 각 GPU-호출 스레드가 `torch.cuda.set_device` 명시.
- [ ] 각 GPU-호출 스레드가 자신의 `torch.cuda.Stream` 보유 + `with` 감싸기.
- [ ] Weight push 경로에서 reader fence (모든 inflight forward 완료 대기) 보장.
현 `weight_sync_event`와 동일한 의미를 in-process로 구현.
- [ ] `torch.compile`된 모듈은 공유 forward 대상에서 제외.
- [ ] Cython 확장이 free-threaded 빌드에서 동작/safe함을 회귀 테스트로 확인.
- [ ] `CUDA_LAUNCH_BLOCKING=1`로 한 번 돌려 silent wrong-output을 동기 에러로
변환해보고 race 미존재 확인.
- [ ] 결정성 테스트: 동일 seed로 multiprocessing 버전과 threading 버전의
iteration 1 forward outputs bit-exact 비교 (allocator/stream 비결정 제외).
- [ ] 부하 테스트: 동시 forward+step을 10⁵ 회 돌려 weight checksum 변화가
예상 범위 내인지 (R1 검출).
### 한 줄 정리
**현재 모델 사이즈에서는 multiprocessing → threading의 위험·복잡도 vs 이득
비율이 나쁘다.** 정말 단일 프로세스가 필요하다면 *모든* GPU 연산을 한
"server thread"로 모으는 M5 형태가 거의 모든 위험을 회피한다 — 그리고 그건
이미 우리가 가진 inference_server.py 구조의 thread 버전일 뿐이다.
---
## 참고 자료
- [PyTorch CUDA semantics docs](https://docs.pytorch.org/docs/stable/notes/cuda.html)
- [PyTorch Autograd mechanics](https://docs.pytorch.org/docs/stable/notes/autograd.html)
- [c10/cuda/CUDAStream.h — current stream is thread-local](https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDAStream.h)
- [c10/cuda/CUDACachingAllocator.cpp](https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDACachingAllocator.cpp)
- [zdevito: A guide to PyTorch's CUDA caching allocator](https://zdevito.github.io/2022/08/04/cuda-caching-allocator.html)
- [pytorch#25540 — per-thread default stream feature request, unresolved](https://github.com/pytorch/pytorch/issues/25540)
- [pytorch#59692 — streams sequentially serialized](https://github.com/pytorch/pytorch/issues/59692)
- [pytorch#101300 — default stream not synchronous](https://github.com/pytorch/pytorch/issues/101300)
- [pytorch#15210 — torch::jit::script::Module not thread-safe](https://github.com/pytorch/pytorch/issues/15210)
- [pytorch#19029 — C++ custom module not thread safe](https://github.com/pytorch/pytorch/issues/19029)
- [pytorch#51452 — JIT module forward thread safety](https://github.com/pytorch/pytorch/issues/51452)
- [pytorch#130249 — Python 3.13t free-threaded support tracking](https://github.com/pytorch/pytorch/issues/130249)
- [forum: Is inference thread-safe?](https://discuss.pytorch.org/t/is-inference-thread-safe/88583)
- [forum: Is PyTorch supposed to be thread-safe?](https://discuss.pytorch.org/t/is-pytorch-supposed-to-be-thread-safe/36540)
- [forum: Only 1 thread for backward?](https://discuss.pytorch.org/t/only-1-thread-for-backward/36824)
- [forum: state_dict vs optimizer.step thread safety](https://discuss.pytorch.org/t/thread-safety-between-model-state-dict-and-optimizer-step/224131)
- [forum: Tensor.cuda(non_blocking=True) in a thread](https://discuss.pytorch.org/t/is-it-safe-to-use-tensor-cuda-non-blocking-true-in-a-thread/182924)
- [forum: Threaded inference c10::CuDNNError](https://discuss.pytorch.org/t/threaded-inference-c10-cudnnerror/182191)
- [dev-discuss: torch.compile + multithreading caches](https://dev-discuss.pytorch.org/t/impact-of-multithreading-and-local-caching-on-torch-compile/2498)
- [Trent Nelson — PyTorch and free-threading](https://trent.me/articles/pytorch-and-python-free-threading/)
- 본 저장소: `docs/performance.md` "Option A Bench Result and Structural
Ceiling", "Free-threaded Python (3.13t) 노트"
- 본 저장소: `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py`,
`trainer.py`, `workers.py`, `evaluate.py`
docs/reports/cost_pytorch_free_threaded_2026-05-07.md
+249
View File
@@ -0,0 +1,249 @@
# Cost Report: Adopting Free-Threaded Python (3.13t / 3.14t) for PyTorch
**Date**: 2026-05-07
**Scope**: Evaluate the cost of migrating this project's Deep CFR training stack
to free-threaded Python (PEP 703 / `python3.13t` / `python3.14t`) so that the
multiprocessing traversal workers + inference server can be replaced by
threads, escaping the structural batching ceiling documented in
`docs/performance.md` ("Option A Bench Result and Structural Ceiling").
## 한 줄 결론
**WAIT (3-6개월).** PyTorch 측 free-threaded 지원은 이미 production 수준
(2.10 cp314t wheel, 2.11 experimental)이지만, **우리 코드 측 진짜 비용은
PyTorch가 아니라 Cython 게임 엔진의 `nogil` 정합성 작업**이다. 그 작업을
하지 않으면 free-threaded Python으로 전환해도 traversal 핫패스에서 GIL이
없어진 효과를 못 본다 — 단지 단일-스레드 perf 회귀(현재 ~20-40%, Python
3.14에서 ~5-10%로 개선 예정)만 떠안게 된다. 더 직접적이고 더 작은 우회로
(Option B-shape per-worker interleaving)가 같은 batch-size 목표를 더 적은
ecosystem-risk로 달성한다.
---
## PyTorch 측 상태 (2026-05 기준)
### Wheel 가용성
- **`torch` 2.10.0** (2026-01): cp313t / cp314t 정식 wheel을 PyPI에 게시.
Linux x86_64 / aarch64, Windows, macOS 모두 커버.
([PyPI torch](https://pypi.org/project/torch/),
[PyTorch Issue #156856](https://github.com/pytorch/pytorch/issues/156856))
- **`torch` 2.11.0** (2026-05 직전): cp314t를 "experimentally supported"로
명시. cp313t는 stable.
([PyTorch 2.11 Release Blog](https://pytorch.org/blog/pytorch-2-11-release-blog/),
[PyTorch 2.11 Release Notes](https://github.com/pytorch/pytorch/releases))
- `pip install torch` on `python3.13t` 환경에서 실제로 동작한다는 트래킹
페이지 보고가 있음.
([Free-threading Compatibility Tracking](https://py-free-threading.github.io/tracking/))
### CUDA wheel 주의사항
- **Python 3.14 (정규/free-threaded)** 의 CUDA wheel 배포 늦음 보고가 있다.
([PyTorch Issue #169929](https://github.com/pytorch/pytorch/issues/169929))
3.14t에서 CUDA를 쓰려면 wheel 인덱스/플랫폼을 확인해야 한다. 우리 RTX 3090
로컬은 cp313t + CUDA가 안전한 조합.
### 안전한 영역
- 기본 텐서 연산 (`torch.as_tensor`, `nn.Linear`, ReLU 등 우리 MLP에서
쓰는 모든 op)은 GIL을 이미 release한다. Quansight/Meta 보고에서
멀티스레드에서 잘 도는 것으로 확인.
([Quansight: free-threaded rollout](https://labs.quansight.org/blog/free-threaded-python-rollout))
- `torch.inference_mode()` 컨텍스트, `state_dict` load, MLP forward —
우리 inference server가 쓰는 모든 경로가 잘 알려진 멀티스레드-가능 영역.
### 위험 영역 / gated
- **`torch.compile`**: free-threaded 빌드에서 "기본 동작은 가능, 진짜
멀티스레드 사용은 미지원"으로 명시.
([PyTorch Issue #156856](https://github.com/pytorch/pytorch/issues/156856))
우리 프로젝트는 `torch.compile` 회귀가 이미 확인되어 비활성 상태이므로
이 제약은 영향 없음 (`docs/performance.md` "torch.compile 실험" 절).
- **`DataLoader`**: 멀티스레드 DataLoader는 SPDL/Meta 측에서 prototype 단계
시연이 있을 뿐, stable contract 아님. 우리 프로젝트는 DataLoader를 쓰지
않으므로 영향 없음.
- **CUDA 멀티스레드 컨텍스트 공유**: long-standing 주의 사항. 한 프로세스
안에서 여러 스레드가 같은 CUDA stream에 일을 던지면 race/순서 문제 가능.
현재 우리 inference server는 단일-스레드 디스패치 루프 (`run_inference_server`,
`inference_server.py:218-246`)이므로 free-threaded로 전환해도 forward
자체는 단일 스레드가 처리하면 안전.
([PyTorch Forums: thread safety + multiprocessing CUDA](https://discuss.pytorch.org/t/thread-safety-in-multiprocessing-cuda-tensors-dont-update-asynchronously/160151),
[CUDA semantics 2.11](https://docs.pytorch.org/docs/2.11/notes/cuda.html))
- **Autograd**: 멀티스레드 backward는 PyTorch에서 historically 락이 많음.
우리 trainer는 단일 스레드에서 backward를 돌릴 것이므로 영향 없음 —
단, 멀티스레드 inference + 동시 backward를 같은 모델에 시도하지 않아야
한다는 일반 원칙은 그대로.
### 단일-스레드 perf 회귀
- **3.13t**: specializing adaptive interpreter 비활성화로 단일-스레드
코드가 ~20-40% 느림.
([CodSpeed: State of 3.13](https://codspeed.io/blog/state-of-python-3-13-performance-free-threading))
- **3.14t**: specializing interpreter 재활성화. 회귀가 ~5-10%로 축소
(예상). 3.14는 PEP 779 통과 후 free-threaded가 "non-experimental,
officially supported"로 격상됨.
- **PyTorch 자체**: native code는 GIL을 이미 release하므로 free-threaded
빌드에서도 텐서 연산 perf는 거의 동일 (Trent Nelson 보고).
([Trent Nelson: PyTorch + free-threading](https://trent.me/articles/pytorch-and-python-free-threading/))
- 우리 traversal hot path는 Cython이라 Python 인터프리터 회귀의 직접
영향이 작지만, Python 콜백 (정책 호출 → numpy → torch) 빈도가 매우
높으므로 (~205k calls/iter) 회귀가 곱 effect로 누적될 수 있다.
### 알려진 이슈 / 트래커
- [pytorch#130249 — Python 3.13 support](https://github.com/pytorch/pytorch/issues/130249) —
지속 업데이트되는 메타 이슈.
- [pytorch#156856 — Python 3.14 support](https://github.com/pytorch/pytorch/issues/156856) —
3.14 / 3.14t 진행 상황.
- [pytorch#169929 — Python 3.14 CUDA wheel 누락 보고](https://github.com/pytorch/pytorch/issues/169929)
### 생태계 신호 (실제로 쓰는 사람들)
- **Optuna**: 3.13t를 정식 지원, 멀티스레드 trial 실행 검증.
([Optuna 3.13t support](https://medium.com/optuna/overview-of-python-free-threading-v3-13t-support-in-optuna-ad9ab62a11ba))
- **Meta SPDL** (DataLoader 대체): ImageNet 이터레이터에서 process →
thread 전환으로 +74% throughput / -50GB 메모리 보고 (8x A100).
단, 이건 고정-비용 비교가 아닌 cherry-picked benchmark.
- **SGLang**: 3.14t 지원 요청 issue가 열려 있음 (open).
([sglang#22889](https://github.com/sgl-project/sglang/issues/22889))
→ 즉, **메이저 LLM serving 프레임워크조차 아직 production 도입을 안
했다**는 신호.
- **Lightning / RLlib**: free-threaded 도입 공식 발표 없음 (검색 시점).
- 일반 보고: PyO3 dependent Rust extensions (pydantic, tiktoken 등) 일부가
free-threaded wheel 없음 → setup이 "fiddly". 우리 프로젝트는 이런
의존성이 사실상 없음 (Cython만 있음 — 이건 free-threaded 지원 wheel
배포 진행 중).
---
## 우리 코드 측 작업량
다음은 multiprocessing → threading 전환 시 만져야 할 곳 / 동시성 가정을
재검토해야 할 곳을 파일·함수 단위로 정리한 것.
### 파일·함수 인벤토리
| 파일 | 역할 | 현재 동시성 가정 | 스레드化 비용 |
| --- | --- | --- | --- |
| `src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py` (`_run_traversal_iteration`, `_evaluate_iteration` 부근, line 461-525, 818-855) | `ProcessPoolExecutor(mp_context=spawn)`로 worker batch dispatch | 워커는 별 프로세스, fork-safe 가정 없음 (spawn) | `ThreadPoolExecutor`로 교체. 각 worker가 trainer의 model state에 read-only 접근 → state_dict copy 시점만 lock으로 보호. 중간. |
| `src/coolrl_lost_cities/games/classic/deep_cfr/workers.py` (`run_traversal_worker_batch`, `_configure_worker_torch_threads` line 28-46) | per-process `torch.set_num_threads(1)`, networks를 매 batch마다 `state_dict`로 load 후 eval | 프로세스마다 격리된 torch state, model copy 1쌍 | thread-shared model로 단순화 가능. `torch.set_num_threads(1)` 호출은 process-global이라 thread 환경에서는 1번만 호출하면 됨 — 약간 작업. **MODEL을 공유하는 순간 weight-update 동시성 문제 신규 발생** (현재는 매 batch 시작 시 state_dict 복사라 자연스레 안전). 중간-높음. |
| `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py` (`InferenceServerController`, `run_inference_server`, line 218-326) | 별 프로세스 + spawn context, shared-memory tensor pool | 프로세스 격리 → thread 환경에서는 server 자체가 불필요. 같은 주소 공간에서 직접 호출. | 서버 자체 삭제 또는 in-process thread-pool 디스패처로 변환. **다만 이 변환이 free-threaded migration의 진짜 목적이므로 비용이라기보다 보상.** 중간. |
| `src/coolrl_lost_cities/games/classic/deep_cfr/inference_client.py` (`InferenceClient.forward`, `NetworkProxy.__call__`) | shared-memory slot 잡고 queue post → event wait | 슬롯 = 워커 1개 가정. 슬롯 free pool은 `mp.Queue` 기반. | 스레드化하면 그냥 직접 model forward 호출 가능. 작은 인터페이스 어댑터만 필요. 작음. |
| `src/coolrl_lost_cities/games/classic/deep_cfr/inference_buffers.py` (`InferenceBuffers`, line 45) | `mp.get_context("spawn")` 기반 shared memory + queues | mp 전용 | 스레드 환경에서는 통째로 불필요. 작음 (삭제). |
| `src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx` (1179줄) | Cython 재귀 traversal, 정책 호출, regret 누적 | **GIL 보유 가정**. Python object 접근 다수. 자유 스레드化 시 race 위험 미평가. | **이게 진짜 비용.** `nogil` 클린업 audit 필요. 1179줄 + `cfr_math.pyx` + `encoding.pyx` 전부. **수일~수주 작업, high risk** (`docs/performance.md`도 같은 결론). |
| `src/coolrl_lost_cities/games/classic/deep_cfr/memory.py` (73줄, `TrainingSample` 추가 경로) | 현재 워커가 결과를 list로 반환, trainer가 단일 스레드에서 `memory.add(...)` | 스레드 추가 시 add 호출이 동시 발생 → 락 필요 | replay buffer add 경로에 `threading.Lock` 추가. 작음. |
| `src/coolrl_lost_cities/games/classic/deep_cfr/inference_server.py` (`run_inference_server` 디스패치 루프, `torch.inference_mode()` + `torch.as_tensor` line 231-246) | 단일 프로세스 단일 스레드 디스패치 | 스레드化 후에도 단일 디스패처 thread 1개로 유지 가능 (CUDA stream 안전) | 단일 thread 보장. 작음. |
| Cython `.pyx` (encoding, cfr_math, traversal) | numpy + Python object 빈번 | `nogil` cleanup 비용 매우 큼 | **자유 스레드化의 critical path.** |
### 새로 필요한 동시성 primitive 위치 추정
다음 곳에 명시적 락 또는 per-thread 격리가 필요해진다:
1. `trainer.py`: 모델 weight 업데이트 ↔ inference 디스패처 read 사이 —
현재는 `weight_queue`가 알아서 sync. 스레드化 후엔 RWLock 또는 epoch
기반 swap 필요. **1곳, 패턴 명확.**
2. `memory.py`: replay buffer `add` 경로. **1곳, 평범한 lock으로 해결.**
3. `traversal_stats.py`: stats 누적 — thread-local 후 머지가 깔끔. **1곳.**
4. `traversal.pyx`: Python-level state mutation (regret/strategy
accumulator dict 등) — `nogil` 영역에서 건드리면 안 됨. **이게 가장
많고 가장 어려운 곳.** 정확한 location 수는 audit 전엔 미상이지만
재귀 호출마다 등장.
### 요약 추정
- **PyTorch 사용 위치만 카운트하면**: 명시적 locking이 새로 필요한 자리는
3-5곳 정도로 매우 작다 (weight swap, replay add, stats merge).
- **진짜 작업량**: Cython 엔진 `nogil` cleanup. 이건 PyTorch 측 free-
threaded 지원과 무관하게 필요한 상수 비용이고, `docs/performance.md`
Decision A의 "free-threaded는 옵션 있지만 cleanup 비용 때문에 deferred"
와 일치한다.
---
## 위험 표
| # | 위험 | 발생 확률 | 영향도 | 검증 방법 |
| -: | --- | --- | --- | --- |
| 1 | Cython traversal 핫패스가 `nogil`-clean이 아니어서 스레드 추가가 곧 GIL 직렬화로 환원 | **높음 (사실상 확정)** | High — 마이그레이션 목적 자체가 사라짐 | `traversal.pyx` 함수에 `nogil` 어노테이션을 시험 적용 → 컴파일 에러로 Python-object 접근 위치 enumerate. 1-2일. |
| 2 | 단일-스레드 perf 회귀 (3.13t: 20-40%) → traversal 절대 시간이 미세하게 더 느려짐 | 높음 (3.13t에서) / 중간 (3.14t에서 ~5-10%) | Medium | 같은 머신에서 `python3.13` vs `python3.13t`로 현 traversal 벤치 (`scripts/bench_inference_backend.py`) 비교. 반나절. |
| 3 | PyTorch CUDA wheel이 cp314t에서 누락/늦음 → CUDA 트레이너에서 import 실패 | 중간 | High (이라면 즉시 블로커) | `pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126` on `python3.14t` 시도. |
| 4 | Cython 의존성 (numpy 등)이 free-threaded wheel 부재 또는 thread-unsafe op | 낮음-중간 | Medium | `pip install` smoke + import 테스트. numpy 2.3+는 free-threaded compat. |
| 5 | 같은 모델에 multi-thread inference + concurrent backward에서 autograd 락 경합 | 중간 (구조에 따라) | Medium | 우리 구조는 단일-스레드 backward라 회피 가능. 디자인 단계에서 "trainer step 동안 inference 중지" 약속만 지키면 됨. |
| 6 | 멀티스레드 CUDA stream 사용 시 race | 낮음 (단일 디스패처 thread 유지하면 0) | High if hit | inference server를 thread 1개로 제한. 검증: `torch.cuda.synchronize()` + functional test. |
| 7 | 생태계 미성숙 — wandb / pytest / ruff / 기타 dev 툴 free-threaded 호환성 | 중간 | Low (개발 환경만 영향, 학습은 OK) | 새 venv 만들어 `uv sync` 시도. |
| 8 | PEP 703 정책 변경 / Python 측 backout (낮지만 0 아님) | 매우 낮음 | High | 트래커 모니터. 3.14에서 phase 2 (officially supported)로 격상되어 위험 감소. |
| 9 | 우리가 적용한 시점이 아직 너무 일러 회귀 발생 시 upstream에 patch 못 받음 | 중간 | Medium | SGLang / Lightning 등 production 도입 신호 대기. 현재 미도입. |
---
## 권고
### 지금 시도하지 말 것
이유 3줄:
1. **PyTorch는 충분히 준비됐지만, 우리 병목은 PyTorch가 아니다.**
`docs/performance.md` Option A 분석은 batching ceiling이 *Cython
traversal의 sync-blocking 구조* 때문이라고 명시. free-threaded Python은
"워커를 스레드로 바꿀 수 있게" 해주지만, Cython이 `nogil`이 아니면
스레드끼리 GIL을 직렬-획득하므로 ceiling이 안 올라간다.
2. **단일-스레드 perf 회귀 (3.13t 20-40%)** 가 traversal 절대 시간을
악화시킬 수 있다. 3.14t에서 5-10%로 개선될 때까지 기다리는 편이 비용
대비 안전.
3. **Production 도입 신호가 아직 약하다**. SGLang은 issue가 open이고,
Lightning/RLlib는 발표 없음. 우리가 early adopter가 되어 디버깅 비용을
짊어질 가치는 단일 프로젝트 입장에서 낮다.
### 더 작은 우회로 (이미 plan에 있음)
`docs/performance.md` "Re-enable A when one of these holds" 항목 #2
**per-worker interleaved traversal (Option B-shape)** 가 같은 batch=64
목표를 free-threaded migration 없이 달성한다. Cython 재귀를 resumable
state machine으로 바꾸는 작업은 `nogil` audit보다 **로컬 영역이고 risk가
낮다** (단일 worker scope, 검증 가능, ecosystem 의존성 0). 자유-스레드
이주의 대안으로 Option B를 먼저 시도하는 것을 권장.
### 다시 보는 조건 (트리거)
다음 중 하나라도 충족되면 재평가:
- **(a) Cython 엔진을 다른 이유로 `nogil`-clean 화한다** — 그 경우
free-threaded Python은 거의 무료 부산물이 된다. 비용 대부분이 그쪽
작업에 흡수됨.
- **(b) Python 3.14.x patch release에서 free-threaded가 stable로 격상되고
단일-스레드 회귀가 ≤5%로 측정됨** + **메이저 ML 프레임워크 (Lightning,
RLlib, vLLM, SGLang 중 2개 이상)** 가 production 도입 발표.
- **(c) 모델이 1024-hidden / 6-layer 이상으로 커져서** GPU forward가
IPC 오버헤드를 흡수할 수 있는 영역으로 들어가면 — Option A 자체가
다시 살아나므로 free-threaded migration이 필요 없어진다.
- **(d) eval이 dominant phase가 됨** (eval_every=5, games=1000). eval은
이미 batch=64이므로 free-threaded와 무관하게 Option A 재활성화로 충분.
### 액션
1. **지금 (0 비용)**: `docs/performance.md`의 free-threaded note를 이
리포트로 cross-link. 트래킹 페이지
([py-free-threading.github.io/tracking](https://py-free-threading.github.io/tracking/))
를 분기별로 확인.
2. **다음 작업 후보**: Option B-shape per-worker interleaving 설계 검토.
3. **장기**: 모델 사이즈 결정이 끝나면 (`docs/performance.md` 권장
sequencing #2) Option A의 batch ceiling이 자연 해소되는지 재측정.
---
## 참고
- [PyTorch Issue #130249 — Python 3.13 support](https://github.com/pytorch/pytorch/issues/130249)
- [PyTorch Issue #156856 — Python 3.14 support](https://github.com/pytorch/pytorch/issues/156856)
- [PyTorch Issue #169929 — Python 3.14 CUDA wheel](https://github.com/pytorch/pytorch/issues/169929)
- [PyTorch 2.11 Release Blog](https://pytorch.org/blog/pytorch-2-11-release-blog/)
- [PyTorch 2.11 Release Notes](https://github.com/pytorch/pytorch/releases)
- [py-free-threading Compatibility Tracking](https://py-free-threading.github.io/tracking/)
- [Quansight Labs — Free-threaded rollout](https://labs.quansight.org/blog/free-threaded-python-rollout)
- [CodSpeed — State of Python 3.13 free-threading perf](https://codspeed.io/blog/state-of-python-3-13-performance-free-threading)
- [Trent Nelson — PyTorch + Free-Threading 실전](https://trent.me/articles/pytorch-and-python-free-threading/)
- [Optuna — 3.13t 지원](https://medium.com/optuna/overview-of-python-free-threading-v3-13t-support-in-optuna-ad9ab62a11ba)
- [SGLang Issue #22889 — 3.14t 지원 요청 (open)](https://github.com/sgl-project/sglang/issues/22889)
- [PyTorch CUDA semantics 2.11](https://docs.pytorch.org/docs/2.11/notes/cuda.html)
- 사내: `docs/performance.md` "Option A Bench Result and Structural Ceiling",
"Free-threaded Python (3.13t) note"