diff --git a/docs/performance.md b/docs/performance.md index da523d3..6d30e86 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -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. diff --git a/docs/plans/amp_trainer.md b/docs/plans/amp_trainer.md new file mode 100644 index 0000000..0141bfc --- /dev/null +++ b/docs/plans/amp_trainer.md @@ -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//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//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. diff --git a/docs/plans/batched_traversal_inference_server.md b/docs/plans/batched_traversal_inference_server.md new file mode 100644 index 0000000..ed35e2b --- /dev/null +++ b/docs/plans/batched_traversal_inference_server.md @@ -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. diff --git a/docs/plans/cython_safe_heuristic_bots.md b/docs/plans/cython_safe_heuristic_bots.md new file mode 100644 index 0000000..c3852f7 --- /dev/null +++ b/docs/plans/cython_safe_heuristic_bots.md @@ -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//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 `_` methods to `cdef inline double _(...)` (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. diff --git a/docs/plans/torch_compile.md b/docs/plans/torch_compile.md new file mode 100644 index 0000000..bb3bdbf --- /dev/null +++ b/docs/plans/torch_compile.md @@ -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. diff --git a/docs/reports/cost_cython_nogil_audit_2026-05-07.md b/docs/reports/cost_cython_nogil_audit_2026-05-07.md new file mode 100644 index 0000000..ca51065 --- /dev/null +++ b/docs/reports/cost_cython_nogil_audit_2026-05-07.md @@ -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 diff --git a/docs/reports/cost_pytorch_cuda_multithread_2026-05-07.md b/docs/reports/cost_pytorch_cuda_multithread_2026-05-07.md new file mode 100644 index 0000000..9c7e607 --- /dev/null +++ b/docs/reports/cost_pytorch_cuda_multithread_2026-05-07.md @@ -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` diff --git a/docs/reports/cost_pytorch_free_threaded_2026-05-07.md b/docs/reports/cost_pytorch_free_threaded_2026-05-07.md new file mode 100644 index 0000000..2523bfd --- /dev/null +++ b/docs/reports/cost_pytorch_free_threaded_2026-05-07.md @@ -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"