Merge pull request #712 from TransformerLensOrg/dev

v2.4.1

tests/unit/test_use_attn_result.py

@@ -0,0 +1,73 @@
+import torch
+
+from transformer_lens import HookedTransformer
+from transformer_lens.HookedTransformerConfig import HookedTransformerConfig
+
+
+def test_atten_result_normal_attn_correct():
+ """Verifies that the attn_result flag does not change the output for models with normal attention."""
+ d_model = 128
+ d_head = 8
+ n_heads = 16
+ n_ctx = 128
+ n_layers = 1
+ d_vocab = 10
+
+ cfg = HookedTransformerConfig(
+ d_model=d_model,
+ d_head=d_head,
+ n_heads=n_heads,
+ n_ctx=n_ctx,
+ n_layers=n_layers,
+ attn_only=True,
+ d_vocab=d_vocab,
+ )
+
+ model = HookedTransformer(cfg)
+ assert model.cfg.use_split_qkv_input is False
+
+ x = torch.arange(1, 9).unsqueeze(0)
+ normal_output = model(x)
+
+ model.set_use_attn_result(True)
+ assert model.cfg.use_attn_result is True
+
+ split_output = model(x)
+
+ assert torch.allclose(normal_output, split_output, atol=1e-6)
+
+
+def test_atten_result_grouped_query_attn_correct():
+ """Verifies that the atten_result flag does not change the output for models with grouped query attention."""
+
+ d_model = 128
+ d_head = 8
+ n_heads = 16
+ n_ctx = 128
+ n_key_value_heads = 2
+ n_layers = 1
+ d_vocab = 10
+
+ cfg = HookedTransformerConfig(
+ d_model=d_model,
+ d_head=d_head,
+ n_heads=n_heads,
+ n_ctx=n_ctx,
+ n_key_value_heads=n_key_value_heads,
+ n_layers=n_layers,
+ attn_only=True,
+ d_vocab=d_vocab,
+ )
+
+ model = HookedTransformer(cfg)
+ assert model.cfg.use_split_qkv_input is False
+
+ x = torch.arange(1, 9).unsqueeze(0)
+ normal_output = model(x)
+
+ model.set_use_attn_result(True)
+ assert model.cfg.use_attn_result is True
+
+ split_output = model(x)
+
+ assert torch.allclose(normal_output, split_output, atol=1e-6)

transformer_lens/HookedTransformer.py

@@ -8,7 +8,6 @@
 alteration of activations in individual components like attention heads and MLP layers, facilitating
 a deeper understanding of the internal workings of transformers like GPT-2.
 """
-
 import logging
 import os
 from typing import Dict, List, NamedTuple, Optional, Tuple, Union, cast, overload
@@ -1570,7 +1569,10 @@ def load_and_process_state_dict(
  # so that quantization settings are not lost
  self.load_state_dict(state_dict, assign=True, strict=False)
  else:
- self.load_state_dict(state_dict, strict=False)
+ state_dict_keys = list(state_dict.keys())
+ for key in state_dict_keys:
+ self.load_state_dict({key: state_dict[key]}, strict=False)
+ del state_dict[key]
 
  def fill_missing_keys(self, state_dict):
  return loading.fill_missing_keys(self, state_dict)