Conversation
jhnwu3
left a comment
There was a problem hiding this comment.
nice cleanup, will check in again, when you're ready + draft is done.
MateehUllah
left a comment
There was a problem hiding this comment.
I reviewed the attention-capture lifecycle and the new regression coverage. One edge case seems worth testing more directly: test_disable_preserves_results_until_next_forward verifies that disabling capture preserves an outstanding gradient hook through backward, which is good, but the documented lifecycle also says a new captured forward replaces the map and clears the previous gradient. I don't see a direct assertion for that transition before backward.
Could we add a regression case that performs a gradient-captured forward/backward, confirms both map and gradient are populated, then performs another captured forward and verifies the new map is present while the old gradient has been cleared? That would directly lock down the lifecycle contract described in the updated API docs and protect against stale gradients being paired with a newer attention map.
Fixes #1166.
AttentionRollout previously enabled backward hooks despite using only attention maps, causing it to fail inside torch.no_grad(). This change separates map capture from gradient capture in Transformer and StageAttentionNet.
Introduces AttentionInterpretable and its subclass GradientInterpretable, retaining CheferInterpretable as an alias. Rollout explicitly requests forward-only capture and runs without autograd; Chefer explicitly requests gradients. Attribution calculations remain unchanged.
Existing positional register_hook arguments and concrete models’ one-argument hook behavior are preserved. Disabling capture stops both modes while retaining the latest results; the next attention forward clears stale gradients and replaces or clears the map. Chefer now raises clear errors for empty layers, missing tensors, and mismatched map/gradient shapes.
Compatibility: custom models must explicitly inherit the appropriate interface and accept the capture_gradients keyword. The alias does not adapt obsolete method signatures. Overlapping captured forwards/backwards on one model remain unsupported and are documented.
Validation: 36 focused tests passed in Pixi, including 11 new tests covering no-grad execution, unchanged parameter gradients, interpreter switching, cache lifecycle, positional compatibility, migration errors, and failure cleanup. All changed executable statements were covered. The synthetic example, new docstring examples, and #1166-specific contribution checks passed.