Coverage for transformer_lens/cache/key_value_cache.py: 38%

1"""Key-Value cache for TransformerLens. 

2 

3Defines the TransformerLensKeyValueCache which manages a list of per-layer 

4cache entries and attention masks. 

5""" 

6 

7from dataclasses import dataclass 

8from typing import List, Union 

9 

10import torch 

11from jaxtyping import Int 

12 

13from transformer_lens.config.transformer_lens_config import TransformerLensConfig 

14from transformer_lens.utilities.multi_gpu import get_device_for_block_index 

15 

16from .key_value_cache_entry import TransformerLensKeyValueCacheEntry 

17 

18 

19@dataclass 

20class TransformerLensKeyValueCache: 

21 """ 

22 A cache for storing past keys and values for the Transformer. This is important for generating text - we can cache a lot of past computation and avoid repeating ourselves! 

23 

24 This cache is a list of TransformerLensKeyValueCacheEntry objects, one for each layer in the Transformer. Each object stores a [batch, pos_so_far, n_heads, d_head] tensor for both keys and values, and each entry has an append method to add a single new key and value. 

25 

26 The cache can be frozen so that it is not updated during the forward pass. This is useful when we want to run many inputs with the same prefix. 

27 """ 

28 

29 entries: List[TransformerLensKeyValueCacheEntry] 

30 previous_attention_mask: Int[torch.Tensor, "batch pos_so_far"] 

31 frozen: bool = False 

32 

33 @classmethod 

34 def init_cache( 

35 cls, 

36 cfg: TransformerLensConfig, 

37 device: Union[torch.device, str, None], 

38 batch_size: int = 1, 

39 ): 

40 # Determine device for each layer 

41 if hasattr(cfg, "n_devices"): 

42 # Configs that track n_devices (TransformerBridgeConfig): per-block placement 

43 device_for_layer = lambda i: get_device_for_block_index(i, cfg, device) 

44 else: 

45 # Fallback when no model is provided - use single device 

46 fallback_device = device if device is not None else cfg.device 

47 if fallback_device is None: 

48 fallback_device = torch.device("cpu") 

49 device_for_layer = lambda i: fallback_device 

50 

51 return cls( 

52 entries=[ 

53 TransformerLensKeyValueCacheEntry.init_cache_entry( 

54 cfg, 

55 device_for_layer(i), 

56 batch_size, 

57 ) 

58 for i in range(cfg.n_layers) 

59 ], 

60 previous_attention_mask=torch.empty( 

61 # This may actually be an int64, but type promotion will handle it: 

62 # See: https://pytorch.org/docs/stable/tensor_attributes.html#type-promotion-doc 

63 # See: https://github.com/pytorch/pytorch/issues/35014 

64 (batch_size, 0), 

65 device=device, 

66 dtype=torch.int, 

67 ), 

68 ) 

69 

70 def freeze(self): 

71 self.frozen = True 

72 for entry in self.entries: 

73 entry.frozen = True 

74 

75 def unfreeze(self): 

76 self.frozen = False 

77 for entry in self.entries: 

78 entry.frozen = False 

79 

80 def append_attention_mask(self, attention_mask: Int[torch.Tensor, "batch new_tokens"]): 

81 attention_mask = attention_mask.to(self.previous_attention_mask.device) 

82 updated_attention_mask = torch.cat([self.previous_attention_mask, attention_mask], dim=-1) 

83 if not self.frozen: 

84 self.previous_attention_mask = updated_attention_mask 

85 return updated_attention_mask 

86 

87 def __getitem__(self, idx): 

88 return self.entries[idx]