fix(compile): skip in-place weight trims when streaming (on_layer_compiled)

The ONNX streaming path (compile_to_onnx) installs an on_layer_compiled callback
that extracts each layer's weights and NULLs the tensors per layer, bounding peak
memory to one dense layer's worth. The post-loop trim_unused_heads /
trim_unused_slots then sliced those now-None tensors -> 'NoneType' object is not
subscriptable (heads) / 'bool' object has no attribute 'any' (slots), but only for
graphs with UNUSED heads/slots (n < n_heads); small graphs where every head is used
took the no-slice branch, so the onnxruntime-skipped compile_to_onnx tests never
caught it. The DOOM forward (sparse attention: ~1-6 of 40 heads/layer used) trips it.

Skip both in-place trims when on_layer_compiled is set: the streamed sparse export
already drops the zero heads/slots, so the in-memory trim is both impossible (weights
freed) and redundant. Non-streaming compile_headless is unchanged.

Adds tests/compile/forward/test_streaming_trim.py (no onnxruntime needed) — fails
without the fix (TypeError at attn.py:161), passes with it. Validated: full
forward-compile dir green (235 passed); compile_to_onnx of the DOOM token-I/O forward
now completes at 2.18GB peak / 55MB ONNX (was an OOM via dense compile_headless).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: physicsrob

tests/compile/forward/test_streaming_trim.py

@@ -0,0 +1,70 @@
+"""Regression: a streaming compile must SKIP the post-loop in-place weight trims.
+
+``compile_to_onnx`` passes an ``on_layer_compiled`` callback that extracts each
+layer's weights and NULLs the tensors per layer (to bound peak memory to one
+dense layer's worth regardless of depth — see export._make_stream_layer_weights_cb).
+The post-loop ``trim_unused_heads`` / ``trim_unused_slots`` then slice those
+tensors; if they run after the streaming null they hit ``NoneType`` and crash
+(``'NoneType' object is not subscriptable`` / ``'bool' object has no attribute
+'any'``).
+
+The onnxruntime-based ``compile_to_onnx`` tests are skipped when onnxruntime is
+absent, so this no-onnxruntime test pins the trim/null ordering directly: with a
+streaming callback present, ``forward_compile`` must complete (trims skipped).
+The streamed sparse export already drops the zero heads/slots, so the in-memory
+trim is both impossible and redundant in that mode.
+"""
+
+import torch
+
+from torchwright.compiler.forward.compile import forward_compile
+from torchwright.graph import PosEncoding
+from torchwright.ops.inout_nodes import create_input
+from torchwright.ops.linear_relu_linear import linear_relu_linear
+
+
+def _null_layer_weights(layer) -> None:
+    """Mimic the ONNX streaming callback: free each layer's dense weights."""
+    layer.attn.attn.query_matrix = None
+    layer.attn.attn.key_matrix = None
+    layer.attn.attn.value_matrix = None
+    layer.attn.attn.output_matrix = None
+    layer.mlp.linear1.output_matrix = None
+    layer.mlp.linear2.output_matrix = None
+
+
+def test_streaming_callback_skips_inplace_trim():
+    inp = create_input("x", 16)
+    pos = PosEncoding(9)  # trig_width 8 == d_head below
+    torch.manual_seed(0)
+    out = linear_relu_linear(
+        inp,
+        torch.randn(24, 16),
+        torch.randn(24),
+        torch.randn(24, 16),
+        torch.randn(16),
+        name="chain",
+    )
+
+    fired: list[int] = []
+
+    def on_layer_compiled(i, layer):
+        _null_layer_weights(layer)
+        fired.append(i)
+
+    # trim_heads=True + a streaming callback that nulls weights: the post-loop
+    # trims must be skipped. Before the fix this raised inside trim_unused_heads
+    # / trim_unused_slots; it must now complete.
+    net = forward_compile(
+        d=64,
+        d_head=8,
+        output_node=out,
+        pos_encoding=pos,
+        on_layer_compiled=on_layer_compiled,
+        trim_heads=True,
+        verbose=False,
+        device="cpu",
+    )
+
+    assert fired, "streaming callback never fired"
+    assert net is not None

torchwright/compiler/forward/compile.py

@@ -945,7 +945,12 @@ def forward_compile(
     net.residual_assignment = ra
     net.assert_aliases = graph.get_assert_aliases()
 
-    if trim_heads:
+    # Skip in-place weight trimming when streaming: ``on_layer_compiled`` (the
+    # ONNX streaming path) extracts each layer's weights and NULLs the tensors
+    # per layer, so the post-loop trims here would slice ``None``. The streamed
+    # sparse export already drops the unused (zero) heads/slots, making the
+    # in-memory trim both impossible and redundant in that mode.
+    if trim_heads and on_layer_compiled is None:
         max_heads = d // d_head
         for layer in net.layers:
             layer.attn.attn.trim_unused_heads()
@@ -987,17 +992,20 @@ def forward_compile(
                     f"{kv_depth:>10}  {cp:>10}  {sa:>10}"
                 )
 
-    # Trim trailing unused MLP slots (same idea as head trimming).
-    for layer in net.layers:
-        layer.mlp.trim_unused_slots()
-    if verbose:
-        mlp_before = d_hidden * len(net.layers)
-        mlp_after = sum(layer.mlp.d_hidden for layer in net.layers)
-        mlp_saved = (mlp_before - mlp_after) * (2 * d + 2)
-        print(
-            f"\n  MLP trimming: {mlp_before - mlp_after}/{mlp_before} "
-            f"slots trimmed ({mlp_saved:,} params saved)"
-        )
+    # Trim trailing unused MLP slots (same idea as head trimming). Skipped when
+    # streaming for the same reason as head trimming above (weights already
+    # extracted + nulled per layer by ``on_layer_compiled``).
+    if on_layer_compiled is None:
+        for layer in net.layers:
+            layer.mlp.trim_unused_slots()
+        if verbose:
+            mlp_before = d_hidden * len(net.layers)
+            mlp_after = sum(layer.mlp.d_hidden for layer in net.layers)
+            mlp_saved = (mlp_before - mlp_after) * (2 * d + 2)
+            print(
+                f"\n  MLP trimming: {mlp_before - mlp_after}/{mlp_before} "
+                f"slots trimmed ({mlp_saved:,} params saved)"
+            )
 
     if device == "auto":
         net.to(get_device(verbose=verbose))