capture: true block-output at the post_block hook#176
Draft
maxsloef-goodfire wants to merge 2 commits into
Draft
capture: true block-output at the post_block hook#176maxsloef-goodfire wants to merge 2 commits into
maxsloef-goodfire wants to merge 2 commits into
Conversation
…ange) Token-level rename across code, tests, docs, and examples. The hook fires after the mlp() call in program order, but in deferred-residual architectures the captured/steered `residual` does not yet include the MLP contribution (the add happens in the NEXT layer's fused add+norm) -- so 'post_mlp' misdescribes the dataflow. 'post_block' names the position in the layer, not a dataflow claim. No functional change: same tensors captured/steered as before. The semantic correction (capturing the true block output, residual + hidden) is stacked on top of this PR. Clients sending 'post_mlp' in capture specs must switch to 'post_block' (pre-release branch; no alias kept). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Introduce apply_block_steering(): captures residual + hidden_states (the true block output, what HF exposes as hidden_states[L+1]) instead of the bare pre-MLP-add residual, which was byte-identical to post_attn's capture in deferred-residual architectures. The sum is computed only when a capture manager is installed on the rank (static for the process lifetime, so torch.compile traces it as a constant branch); steering still applies to `residual` -- identical propagation to the old behavior. The 43 deferred-residual model files are a mechanical one-line call-site conversion + import each (verified: no other changes). The 30 materialized-stream models (opt/gpt_neox/gemma4/...) already hooked the materialized stream and need no change. Validated on qwen3-8b vs HF transformers: post_block[L] == hs[L+1] at mean cos 0.99995 (layers 0-34; HF's final tuple entry is post-final-norm so the last layer is not directly comparable). Stacked on the mechanical post_mlp->post_block rename. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
maxsloef-goodfire
force-pushed
the
post-block-semantics
branch
from
d6f90af to
995f6c0
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR 2 of 3 (stacked on #175 — view the top commit only: 45 files, +153/−44).
Introduces
apply_block_steering()and converts the 43 deferred-residual model files to it. In those architectures (llama/qwen/deepseek-style), each branch-add is folded into the next layer's fused add+norm, so the old hook captured aresidualthat did not yet include the MLP output — byte-identical to whatpost_attnalready captures (measured cos ~0.4–0.97 vs true hidden states). The new helper capturesresidual + hidden_states= the true block output (hs[L+1]), validated vs HF transformers at cos 0.99995 (layers 0–34) on qwen3-8b.torch.compileconstant-folds the branch and non-capture servers pay nothing.residualinto the next fused add).The capture-path fixes and features (
capture_wait, v2-runner guard, global filesystem capture, fsync docs) are stacked on top in #174.🤖 Generated with Claude Code