Skip to content

Commit 3d94208

Browse files
physicsrobclaude
andcommitted
Docs: compile entry points section; schedule-cache caveat in ONNX triage
- CLAUDE.md gains a 'Compile entry points' section: the new compile_headless(graph, pos_encoding, *, ...) convention with optimize/assume_zero_init parity, the OnnxArtifact return handle (paths + scalars only — the streaming memory bound stays hard), and the load_onnx front door (doom's OnnxTokenRuntime remains the CUDA-graph perf runtime). - The three-cause ONNX self-consistency triage note — in CLAUDE.md AND in _ONNX_CONSISTENCY_CAUSES (the error text people actually see) — now says the discrimination recompile must run with TW_SCHEDULE_CACHE_DIR unset: a cache-replayed schedule bug reproduces on both backends and would masquerade as cause 2/3. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 5780f3f commit 3d94208

2 files changed

Lines changed: 49 additions & 2 deletions

File tree

CLAUDE.md

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,47 @@ against.
259259
and `/tmp/torchwright-modal-run.log` symlinks to the latest run.
260260
Grep that file instead of spending another Modal round-trip.
261261

262+
# Compile entry points
263+
264+
Three entry points in `torchwright/compiler/export.py`, one loader
265+
front door in `torchwright/compiler/onnx_load.py`:
266+
267+
compiled = compile_headless(graph, pos_encoding, *, d=..., d_head=...,
268+
optimize=0, assume_zero_init=False, ...)
269+
270+
In-process `CompiledHeadless` for tests and debugging. `graph` is
271+
either a single output `Node` (outputs gathered at the node's natural
272+
residual columns) or an `io` dict `{"name": (input_node, output_node)}`
273+
(overlay mode: outputs land at input columns via delta transfer, for
274+
autoregressive feedback). All other parameters are keyword-only.
275+
`optimize` and `assume_zero_init` thread straight to `forward_compile`,
276+
so this backend can reproduce a production `optimize=2` schedule
277+
exactly. Passing the `PosEncoding` first (the pre-2026 argument
278+
order) raises a `TypeError` naming the new order.
279+
280+
artifact = compile_to_onnx(output_node, pos_encoding, embedding, path, ...)
281+
artifact = compile_headless_to_onnx(output_node, pos_encoding, path, ...)
282+
283+
Both exporters return an **`OnnxArtifact`**: the written paths
284+
(`path`, `meta_path`, `debug_path`) plus small build metadata (`kind`,
285+
`n_layers`, `per_layer_n_heads`, `d`, `d_head`, `cache_stride`,
286+
`cache_window`; token exports add `d_embed`/`vocab_size`). It is
287+
built strictly from paths and scalars after export completes — it
288+
holds no graph, no weights, no exporter state (the exporters'
289+
streaming memory bound is a hard invariant). `artifact.load()`
290+
returns the matching runtime module; `artifact.debug_session(
291+
output_node, pos_encoding)` opens an `OnnxDebugSession` (see
292+
*Debugging the ONNX artifact* below).
293+
294+
model = load_onnx(path) # torchwright.compiler.onnx_load
295+
296+
Loads any torchwright ONNX export by dispatching on the sidecar's
297+
format key: `OnnxHeadlessModule` for float-I/O exports,
298+
`OnnxTokenModule` for token-I/O exports (vocab tokenizer + argmax
299+
`generate` loop). torchwright_doom's `OnnxTokenRuntime` remains the
300+
CUDA-graph perf runtime; these loaders are the contract-correctness
301+
harness.
302+
262303
# Debugging compiled graphs
263304

264305
When a compiled graph produces wrong output, the cause is almost
@@ -346,6 +387,8 @@ seconds; the compile it replaces is minutes):
346387

347388
output_node, pos_encoding = build_my_graph() # deterministic rebuild
348389
sess = OnnxDebugSession("model.onnx", output_node, pos_encoding)
390+
# or, holding the export's OnnxArtifact:
391+
# sess = artifact.debug_session(output_node, pos_encoding)
349392

350393
out, past = sess.step(inputs, sess.empty_past(), debug=True)
351394
val = sess.debug_value(node)
@@ -374,7 +417,9 @@ Requirements and caveats:
374417
stop and report; (2) ONNX-emission bug in `compiler/export.py`
375418
also D1; (3) a debug-sidecar/canonical-id remap bug. Recompiling
376419
via `compile_headless` and re-running `debug=True` discriminates
377-
(1) from (2)+(3).
420+
(1) from (2)+(3) — run the discrimination compile with
421+
`TW_SCHEDULE_CACHE_DIR` unset: a cache-replayed schedule bug
422+
reproduces on both backends and would masquerade as cause 2/3.
378423
- The debug session is separate from any production session: the
379424
promoted outputs defeat onnxruntime's memory-reuse planning. Never
380425
put it on a hot path.

torchwright/debug/onnx_debug.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,9 @@
8383
"\n (3) debug-sidecar / canonical-id remap bug"
8484
" (torchwright/compiler/graph_identity.py)."
8585
"\n Recompiling via compile_headless and re-running debug=True"
86-
" discriminates (1) from (2)+(3)."
86+
" discriminates (1) from (2)+(3) — with TW_SCHEDULE_CACHE_DIR unset"
87+
" for that run: a cache-replayed schedule bug reproduces on both"
88+
" backends and would masquerade as cause 2/3."
8789
)
8890

8991

0 commit comments

Comments
 (0)