Skip to content

feat(capture): wire mlp_in/mlp_out hooks for transcoder training - #206

Merged
RhizoNymph merged 1 commit into
feat/integrationfrom
feat/capture-mlp-hooks
Jun 28, 2026
Merged

feat(capture): wire mlp_in/mlp_out hooks for transcoder training#206
RhizoNymph merged 1 commit into
feat/integrationfrom
feat/capture-mlp-hooks

Conversation

@RhizoNymph

Copy link
Copy Markdown
Owner

What

Wires the mlp_in and mlp_out capture hooks into the decoder layers of gemma3, gemma4, qwen3, qwen3_moe, and qwen3_next. These were already first-class hook names in the capture type system (HookName, _HOOK_NAME_TO_ID, _VALID_HOOK_NAMES) but had no tap insertion points in any model, so requesting them produced nothing.

  • mlp_in — the normalized activation fed into the MLP/MoE sublayer (output of the pre-MLP norm).
  • mlp_out — the MLP/MoE branch the sublayer writes back to the residual stream (after any post-MLP norm / layer-scale), so the residual-stream decomposition post_block == post_attn + mlp_out holds where post_attn is captured after the deferred add (qwen family, gemma4).

Why

These are the paired activations required to train transcoders (a transcoder learns mlp_in → mlp_out for a layer; cross-layer transcoders consume them across layers, aligned per token). Until now the capture machinery exposed only residual-stream points (pre_attn/post_attn/post_block), from which MLP input/output could only be indirectly reconstructed.

Notes

  • The taps read replicated tensors (mlp_in is the normed input; mlp_out is read after the down-proj all-reduce / MoE combine), so they capture on TP rank 0 like the existing residual-stream hooks — no sharded-gather path needed.
  • mlp_in/mlp_out are intentionally excluded from the :all fan-out so :all stays model-agnostic and never reserves buffers on models that don't wire them. Request them explicitly, e.g. {"mlp_in": [12], "mlp_out": [12]}.
  • On gemma MoE layers mlp_in covers only the dense path (the parallel MoE branch is normed separately); mlp_out always captures the combined branch.
  • Docs updated: capture feature doc gains a Hook-points table; the parallelism design doc clarifies mlp_in/mlp_out are replicated and captured today.

Scope / follow-up

qwen3_next.py is wired but its unit test skips: qwen3_next.py and qwen3_5.py currently fail to import on feat/integration because the upstream mamba refactor (vllm-project#41126) moved mamba.gdn_linear_attn into the mamba.gdn package and the integration merge (#186) left these two files pointing at the dead path. That import breakage is pre-existing and unrelated to this change; it will be fixed separately (it needs GPU validation on the linear-attn path). Once fixed, Qwen3-Next/Qwen3.5 inherit the wiring here and the skipped test runs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant