Document torch.compile experiment result (regression)

Wrapping trainer networks with torch.compile produced a 4.8% regression
in iteration time on default.yaml (17.93s → 18.79s). Two causes: (1)
the dominant phase is CPU traversal which bypasses the compiled
wrapper, (2) DeepCFRMLP is too small for compile dispatch overhead to
pay back. Implementation kept on experiments/torch-compile for future
revisits when the trainer model or inference path changes.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-07 18:17:46 +09:00
co-authored by Claude Opus 4.7
parent 1de8e8e2b3
commit 079d9d916a
+33
View File
@@ -322,3 +322,36 @@ Based on the current metrics, the more plausible performance work is:
6. For eval-heavy runs, optimize the safe-heuristic opponents and policy 6. For eval-heavy runs, optimize the safe-heuristic opponents and policy
post-processing before assuming TensorRT is the main lever. post-processing before assuming TensorRT is the main lever.
The inspected eval row shows those costs dominate the slowest opponents. The inspected eval row shows those costs dominate the slowest opponents.
## Experiments
### `torch.compile` on trainer networks (2026-05-07, regression)
Wrapped both advantage networks and the strategy network with
`torch.compile()` at trainer construction time. Implementation also
required a `_clean_state_dict()` helper to strip the `_orig_mod.` prefix
that compiled modules add to `state_dict()`, plus a `_orig_mod`-routed
path for `load_state_dict()` so multiprocessing traversal workers and
checkpoint restoration could keep using the uncompiled `DeepCFRMLP`
class.
Measurement (8 iterations on `default.yaml`, eval and checkpoint
disabled, iteration 1 dropped as compile warm-up):
| | iter mean | 1000-iter projection |
| --- | ---: | ---: |
| Baseline (no compile) | 17.93s | 4.98h |
| `torch.compile` on trainer nets | 18.79s | 5.22h |
| Effect | +0.86s (+4.8%) | +14 min |
Net result: regression. Two reasons:
- Traversal is ~60% of iteration time and runs in CPU multiprocessing
workers that reconstruct networks from cleaned `state_dict`s, so they
bypass the compiled wrapper entirely.
- `DeepCFRMLP` (512-hidden, 3-layer) is small enough that the compiled
call dispatch overhead exceeds the kernel-fusion benefit.
Implementation preserved on branch `experiments/torch-compile` for
revisiting if the trainer model grows substantially or traversal moves
to GPU-batched inference. Not enabled on `main`.