Skip to content

Fix Granite residual contribution hooks - #1658

Merged
jlarson4 merged 6 commits into
devfrom
bug/granite-1648-residual-multiplier
Aug 12, 2026
Merged

Fix Granite residual contribution hooks#1658
jlarson4 merged 6 commits into
devfrom
bug/granite-1648-residual-multiplier

Conversation

@jlarson4

Copy link
Copy Markdown
Collaborator

Description

  • Granite-family HF blocks compute residual + sublayer_out * residual_multiplier inline, so blocks.N.hook_attn_out / hook_mlp_out exposed raw module outputs. The value was off by the multiplier and wrong as intervention points.
  • New ScaledResidualBlockBridge makes hook_attn_out / hook_mlp_out real HookPoints firing on the scaled contribution. Reads see raw * scale with a bit-exact forward, a write of v (including in-place mutation) reconstructs v as the contribution, and backward hooks always route through the HookPoint.
  • Adopted by the granite, granite_moe, and granite_moe_hybrid adapters, with residual_multiplier added to the config passthrough. Zero scale is rejected at construction.
  • On hybrid mamba layers hook_attn_out is pruned at bind time. With experts, hook_mlp_out stays absent, the block_sparse_moe + shared_mlp sum has no single capture point.
  • Renamed the block capture-hook lifecycle (_maybe_wire_capture_hooks etc.) since it now wires more than pre-LN capture.
  • Tests use local tiny fixtures at residual_multiplier=0.22: residual identities, contribution = raw × scale, zero-ablation, write/in-place-mutation/bit-exactness regressions, backward-hook gradient flow, and HF forward parity across dense + MoE + hybrid.
  • Tests verified to fail on the unfixed adapter, allowing us to remove the GraniteMoeHybridForCausalLM strict-xfail in the alias-resolution audit.

Fixes #1648 – Granite portion

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • This change requires a documentation update

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4
jlarson4 merged commit 8c7d690 into dev Aug 12, 2026
25 checks passed
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.

[Bug Report] hook_attn_out / hook_mlp_out are pre-transform module outputs on OLMo 2/3 and Granite

1 participant