Refresh optimization plans
Archive implemented AMP, Option A inference-server, and Cython heuristic plans. Add the active Option B interleaved traversal plan and update model-size/torch.compile plans to reflect the current traversal scheduling conclusion. Co-Authored-By: Codex <codex@openai.com>
This commit is contained in:
Codex
@@ -0,0 +1,226 @@
|
||||
# Plan: Option B Per-Worker Interleaved Traversal
|
||||
|
||||
**Status:** Planning only; do not implement until this design is reviewed.
|
||||
**Owner:** Codex for prototype design and implementation; operator for long-run
|
||||
benchmarks on `home`.
|
||||
**Background:** Option A, the central traversal inference server, was implemented
|
||||
and benchmarked on 2026-05-07. It regressed traversal because the existing
|
||||
recursive worker path is sync-blocking and can only feed batches near the worker
|
||||
count, not the GPU-efficient bs=64+ regime. See `docs/performance.md`:
|
||||
"Option A Bench Result and Structural Ceiling" and "Clarifying the traversal
|
||||
bottleneck: sync policy boundary, not SIMD."
|
||||
|
||||
## Goal
|
||||
|
||||
Restructure traversal scheduling so each worker advances many traversal
|
||||
instances concurrently, suspends each instance at policy-needed states, batches
|
||||
the pending policy requests, runs one policy forward, and resumes the matching
|
||||
instances.
|
||||
|
||||
The target is to turn the current policy-forward shape:
|
||||
|
||||
```text
|
||||
one traversal -> policy request -> bs=1 forward -> resume
|
||||
```
|
||||
|
||||
into:
|
||||
|
||||
```text
|
||||
N traversal continuations -> collect policy requests -> bs=32..128 forward -> resume
|
||||
```
|
||||
|
||||
without changing CFR math, game rules, replay sample semantics, or public
|
||||
training CLI behavior.
|
||||
|
||||
## Why This Is The Next Optimization
|
||||
|
||||
Microbench evidence from `experiments/traversal_policy_boundary/`:
|
||||
|
||||
| Device | Component | median us/call |
|
||||
| --- | --- | ---: |
|
||||
| CPU | encode + legal | 3.10 |
|
||||
| CPU | push + pop | 0.15 |
|
||||
| CPU | policy boundary bs=1 | 111.50 |
|
||||
| CUDA | policy boundary bs=1 | 181.30 |
|
||||
| CUDA | torch forward bs=64 | 2.55 |
|
||||
|
||||
The game mechanics and encoding are not the dominant cost. The dominant cost is
|
||||
the one-row Python/PyTorch policy boundary. Option A moved that boundary to a
|
||||
server process, but because every worker blocks waiting for one response, the
|
||||
server observed mean batches around 7-8 and regressed end-to-end. Option B is
|
||||
the first design that directly changes the scheduling shape.
|
||||
|
||||
## Non-Goals
|
||||
|
||||
- Do not re-enable `traversal.inference_backend: server` as the default.
|
||||
Option A remains available but structurally capped until traversal can feed
|
||||
larger batches.
|
||||
- Do not port traversal to Julia, C++, or a new game engine.
|
||||
- Do not change Deep CFR sampling math, regret targets, strategy-memory
|
||||
location, weighting, or replay schema.
|
||||
- Do not change model architecture.
|
||||
- Do not implement TensorRT, `torch.compile`, or AMP in this plan.
|
||||
- Do not remove the existing recursive traversal path until the interleaved path
|
||||
has parity and benchmark evidence.
|
||||
|
||||
## Design Sketch
|
||||
|
||||
The current Cython traversal is recursive and calls policy synchronously. Option
|
||||
B needs an explicit continuation representation so policy calls become yield
|
||||
points.
|
||||
|
||||
One worker owns a fixed set of active traversal contexts:
|
||||
|
||||
```text
|
||||
TraversalContext
|
||||
GameState state
|
||||
explicit stack frames replacing recursion
|
||||
RNG state
|
||||
traverser player
|
||||
iteration
|
||||
partial node/action values
|
||||
pending info_state/legal mask
|
||||
output advantage/strategy samples
|
||||
TraversalStats
|
||||
```
|
||||
|
||||
Worker loop:
|
||||
|
||||
1. Initialize `K` traversal contexts from the worker's assigned seeds.
|
||||
2. Advance each context until it reaches one of:
|
||||
- terminal/cutoff/done,
|
||||
- needs policy forward,
|
||||
- error.
|
||||
3. Collect pending policy requests into a batch, grouped by network target:
|
||||
advantage player 0/1, strategy network, or league snapshot if enabled.
|
||||
4. Run batched forward for each group on the worker's selected inference device.
|
||||
5. Scatter logits/advantages back into the contexts.
|
||||
6. Resume contexts until all assigned traversals finish.
|
||||
7. Return the same `(stats, advantage_samples, strategy_samples)` shape as
|
||||
`run_cython_traversal_batch`.
|
||||
|
||||
The first prototype should keep one worker process and one GPU model copy per
|
||||
worker if `device=cuda`. That may duplicate VRAM across workers, so the initial
|
||||
benchmark can run with fewer workers and larger `interleave_width`. A later
|
||||
hybrid can combine Option B's continuation batching with Option A's central
|
||||
server if VRAM pressure dominates.
|
||||
|
||||
## Config Surface
|
||||
|
||||
Add only after the prototype proves parity:
|
||||
|
||||
```yaml
|
||||
traversal:
|
||||
scheduler: recursive # recursive | interleaved
|
||||
interleave_width: 64 # traversal contexts advanced per worker
|
||||
interleave_max_batch: 128 # cap per forward group
|
||||
```
|
||||
|
||||
Default remains `recursive`.
|
||||
|
||||
## Implementation Phases
|
||||
|
||||
### Phase 0: Design Spike
|
||||
|
||||
- Trace current `_traverse` control flow and enumerate every value that must
|
||||
survive across a policy yield point.
|
||||
- Decide whether to implement the explicit stack in Cython (`.pyx`) or as a
|
||||
Python prototype first.
|
||||
- Identify exact parity surfaces:
|
||||
`TraversalStats`, advantage samples, strategy samples, RNG sequence, and
|
||||
terminal/cutoff behavior.
|
||||
|
||||
Deliverable: short design note appended to this plan before code work starts.
|
||||
|
||||
### Phase 1: Python Prototype, No Production Wiring
|
||||
|
||||
- Add an experiment-only traversal prototype under `experiments/` that mimics
|
||||
the current traversal semantics with explicit stacks.
|
||||
- Use small configs (`max_depth`, low traversal count) and compare samples/stats
|
||||
to the recursive path under fixed seeds.
|
||||
- Measure realized batch size and scheduler overhead.
|
||||
|
||||
Success gate: sample/stat parity on small deterministic fixtures, and realized
|
||||
policy batches materially above worker count.
|
||||
|
||||
### Phase 2: Cython Prototype Behind Non-Default Flag
|
||||
|
||||
- Add an interleaved traversal entry point beside the existing recursive one.
|
||||
- Keep the existing recursive path untouched and default.
|
||||
- Wire through `workers.py` only behind `traversal.scheduler: interleaved`.
|
||||
- Add focused tests for parity on smoke configs.
|
||||
|
||||
Success gate: `uv run pytest -q tests/games/classic/test_deep_cfr_trainer.py`
|
||||
and new interleaved traversal tests pass.
|
||||
|
||||
### Phase 3: Benchmark
|
||||
|
||||
Benchmark against current `default.yaml`, eval/checkpoint disabled:
|
||||
|
||||
```bash
|
||||
uv run lost-cities-deep-cfr train \
|
||||
--config configs/deep_cfr/default.yaml \
|
||||
--set run.max_iterations=10 \
|
||||
--set checkpoint.save_latest=false \
|
||||
--set checkpoint.save_every=0 \
|
||||
--set evaluation.eval_every=0
|
||||
```
|
||||
|
||||
Compare:
|
||||
|
||||
- recursive baseline,
|
||||
- interleaved with `interleave_width` in `{16, 32, 64, 128}`,
|
||||
- worker counts in `{1, 2, 4, 8}` as VRAM allows.
|
||||
|
||||
Metrics:
|
||||
|
||||
- `iteration_seconds`
|
||||
- `traversal_seconds`
|
||||
- realized policy batch size mean/p50/p95/max
|
||||
- scheduler overhead if instrumented
|
||||
- `advantage_train_seconds`, `strategy_train_seconds` to confirm no unrelated
|
||||
drift
|
||||
- sample counts and traversal stats
|
||||
|
||||
Success gate: at least **1.5x traversal speedup** with no sample/stat parity
|
||||
failure. Stretch target: **3x traversal speedup** if realized batches reach the
|
||||
bs=64 regime without high scheduler overhead.
|
||||
|
||||
## Risks
|
||||
|
||||
- **State-machine complexity.** Recursive CFR control flow has many local
|
||||
values. Mitigation: prototype with small depth and exhaustive parity before
|
||||
optimizing.
|
||||
- **RNG drift.** Interleaving changes operation order. Mitigation: store RNG
|
||||
state per traversal context and define parity against the recursive path only
|
||||
where ordering is intentionally preserved. If exact ordering is impossible,
|
||||
require distributional/sample-count parity and document the break.
|
||||
- **VRAM duplication.** Per-worker GPU models may not fit at larger model sizes.
|
||||
Mitigation: start with fewer workers and larger interleave width; revisit a
|
||||
central server only after Option B proves the scheduling benefit.
|
||||
- **Sample memory pressure.** More active contexts mean more pending samples.
|
||||
Mitigation: stream completed samples out of contexts as soon as a traversal
|
||||
finishes.
|
||||
- **Scheduler overhead cancels batching.** Mitigation: benchmark light and heavy
|
||||
modes; record realized batch size and overhead explicitly.
|
||||
|
||||
## Decision Tree
|
||||
|
||||
- **Parity fails in Phase 1/2:** stop. Do not optimize. Document the exact
|
||||
mismatch.
|
||||
- **Parity passes, realized batch remains <16:** Option B did not change the
|
||||
structural ceiling enough. Reconsider Option C or a deeper traversal rewrite.
|
||||
- **Parity passes, realized batch >=64, speedup <1.5x:** batching worked but
|
||||
non-forward work dominates. Keep recursive default and document.
|
||||
- **Parity passes, traversal speedup >=1.5x:** keep interleaved behind config,
|
||||
run longer learning-curve A/B.
|
||||
- **Longer A/B is stable and speedup persists:** consider making
|
||||
`traversal.scheduler: interleaved` the default.
|
||||
|
||||
## Definition Of Done
|
||||
|
||||
- Plan reviewed and Phase 0 design note added.
|
||||
- Prototype proves whether explicit continuation batching can preserve traversal
|
||||
semantics.
|
||||
- Bench results are added to `docs/performance.md`.
|
||||
- Default behavior remains unchanged until parity and benchmark gates pass.
|
||||
Reference in New Issue
Block a user