coorl-lost-cities/docs/plans/torch_compile.md

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_dicts 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_dicts.

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_dicts 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_dicts. 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:

@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:

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():

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.