Add a Class AttentionScores that carries out the computation from q,k to attn_scores. This allows us to hook and cache this computation.
Motivation and Context
This was required for interp.
How Has This Been Tested?
Added a test that gets attn scores for Pythia and (as a simple test that we really got attention scores) checks the causal mask test_pythia_attn_causal_mask
Does this PR introduce a breaking change?
No
Confusions
I am a bit confused why my cached activations were a list (of length 1) and I had to do
AttentionScores module
Description
Add a Class
AttentionScores
that carries out the computation from q,k toattn_scores
. This allows us to hook and cache this computation.Motivation and Context
This was required for interp.
How Has This Been Tested?
test_pythia_attn_causal_mask
Does this PR introduce a breaking change?
No
Confusions
I am a bit confused why my cached activations were a list (of length 1) and I had to do
but I class this as "that's just how our hooks work"