Coverage for transformer_lens/benchmarks/activation_cache.py: 46%

1"""Activation cache benchmarks for TransformerBridge.""" 

2 

3from typing import Optional 

4 

5import torch 

6 

7from transformer_lens.ActivationCache import ActivationCache 

8from transformer_lens.benchmarks.utils import ( 

9 BenchmarkResult, 

10 BenchmarkSeverity, 

11 safe_allclose, 

12) 

13from transformer_lens.model_bridge import TransformerBridge 

14 

15 

16def benchmark_run_with_cache( 

17 bridge: TransformerBridge, 

18 test_text: str, 

19) -> BenchmarkResult: 

20 """Benchmark run_with_cache functionality (structural self-check). 

21 

22 Args: 

23 bridge: TransformerBridge model to test 

24 test_text: Input text for testing 

25 

26 Returns: 

27 BenchmarkResult with cache functionality details 

28 """ 

29 try: 

30 output, cache = bridge.run_with_cache(test_text) 

31 

32 # Verify output and cache 

33 if not isinstance(output, torch.Tensor): 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 return BenchmarkResult( 

35 name="run_with_cache", 

36 severity=BenchmarkSeverity.DANGER, 

37 message="Output is not a tensor", 

38 passed=False, 

39 ) 

40 

41 if not isinstance(cache, ActivationCache): 41 ↛ 42line 41 didn't jump to line 42 because the condition on line 41 was never true

42 return BenchmarkResult( 

43 name="run_with_cache", 

44 severity=BenchmarkSeverity.DANGER, 

45 message="Cache is not an ActivationCache object", 

46 passed=False, 

47 ) 

48 

49 if len(cache) == 0: 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true

50 return BenchmarkResult( 

51 name="run_with_cache", 

52 severity=BenchmarkSeverity.DANGER, 

53 message="Cache is empty", 

54 passed=False, 

55 ) 

56 

57 # Verify cache contains expected keys 

58 cache_keys = list(cache.keys()) 

59 expected_patterns = ["embed", "unembed"] 

60 # Not all architectures have ln_final (e.g., OPT-350m). 

61 has_ln_final = ( 

62 hasattr(bridge, "adapter") 

63 and bridge.adapter.component_mapping 

64 and "ln_final" in bridge.adapter.component_mapping 

65 ) 

66 if has_ln_final: 66 ↛ 69line 66 didn't jump to line 69 because the condition on line 66 was always true

67 expected_patterns.append("ln_final") 

68 

69 missing_patterns = [] 

70 for pattern in expected_patterns: 

71 if not any(pattern in key for key in cache_keys): 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true

72 missing_patterns.append(pattern) 

73 

74 if missing_patterns: 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true

75 return BenchmarkResult( 

76 name="run_with_cache", 

77 severity=BenchmarkSeverity.DANGER, 

78 message=f"Cache missing expected patterns: {missing_patterns}", 

79 details={"missing": missing_patterns, "cache_keys_count": len(cache_keys)}, 

80 passed=False, 

81 ) 

82 

83 # Verify cached tensors are actually tensors 

84 non_tensor_keys = [] 

85 for key, value in cache.items(): 

86 if not isinstance(value, torch.Tensor): 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true

87 non_tensor_keys.append(key) 

88 

89 if non_tensor_keys: 89 ↛ 90line 89 didn't jump to line 90 because the condition on line 89 was never true

90 return BenchmarkResult( 

91 name="run_with_cache", 

92 severity=BenchmarkSeverity.DANGER, 

93 message=f"Cache contains {len(non_tensor_keys)} non-tensor values", 

94 details={"non_tensor_keys": non_tensor_keys[:5]}, 

95 passed=False, 

96 ) 

97 

98 return BenchmarkResult( 

99 name="run_with_cache", 

100 severity=BenchmarkSeverity.INFO, 

101 message=f"run_with_cache successful with {len(cache)} cached activations", 

102 details={"cache_size": len(cache)}, 

103 ) 

104 

105 except Exception as e: 

106 return BenchmarkResult( 

107 name="run_with_cache", 

108 severity=BenchmarkSeverity.ERROR, 

109 message=f"run_with_cache failed: {str(e)}", 

110 passed=False, 

111 ) 

112 

113 

114def benchmark_activation_cache( 

115 bridge: TransformerBridge, 

116 test_text: str, 

117 reference_cache: Optional[dict[str, torch.Tensor]] = None, 

118 tolerance: float = 1e-3, 

119) -> BenchmarkResult: 

120 """Benchmark activation cache values against a reference activation snapshot. 

121 

122 Args: 

123 bridge: TransformerBridge model to test 

124 test_text: Input text for testing (must match the snapshot's prompt) 

125 reference_cache: Optional reference activations keyed by hook name (e.g. a 

126 golden fixture snapshot). Structural self-check only if None. 

127 tolerance: Tolerance for activation comparison 

128 

129 Returns: 

130 BenchmarkResult with cache value comparison details 

131 """ 

132 try: 

133 bridge_output, bridge_cache = bridge.run_with_cache(test_text) 

134 

135 if reference_cache is None: 135 ↛ 145line 135 didn't jump to line 145 because the condition on line 135 was always true

136 # No reference - just verify cache structure 

137 return BenchmarkResult( 

138 name="activation_cache", 

139 severity=BenchmarkSeverity.INFO, 

140 message=f"Activation cache created with {len(bridge_cache)} entries", 

141 details={"cache_size": len(bridge_cache)}, 

142 ) 

143 

144 # Find common keys 

145 bridge_keys = set(bridge_cache.keys()) 

146 reference_keys = set(reference_cache.keys()) 

147 common_keys = bridge_keys & reference_keys 

148 

149 if len(common_keys) == 0: 

150 return BenchmarkResult( 

151 name="activation_cache", 

152 severity=BenchmarkSeverity.DANGER, 

153 message="No common keys between Bridge and Reference caches", 

154 details={ 

155 "bridge_keys": len(bridge_keys), 

156 "reference_keys": len(reference_keys), 

157 }, 

158 passed=False, 

159 ) 

160 

161 # Compare activations for common keys 

162 mismatches = [] 

163 for key in sorted(common_keys): 

164 bridge_tensor = bridge_cache[key] 

165 reference_tensor = reference_cache[key] 

166 

167 # Check shapes 

168 if bridge_tensor.shape != reference_tensor.shape: 

169 mismatches.append( 

170 f"{key}: Shape mismatch - Bridge{bridge_tensor.shape} vs Ref{reference_tensor.shape}" 

171 ) 

172 continue 

173 

174 # Check values 

175 if not safe_allclose(bridge_tensor, reference_tensor, atol=tolerance, rtol=0.0): 

176 b = bridge_tensor.cpu().float() 

177 r = reference_tensor.cpu().float() 

178 max_diff = torch.max(torch.abs(b - r)).item() 

179 mean_diff = torch.mean(torch.abs(b - r)).item() 

180 mismatches.append( 

181 f"{key}: Value mismatch - max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}" 

182 ) 

183 

184 if mismatches: 

185 return BenchmarkResult( 

186 name="activation_cache", 

187 severity=BenchmarkSeverity.WARNING, 

188 message=f"Found {len(mismatches)}/{len(common_keys)} cached activations with differences", 

189 details={ 

190 "total_keys": len(common_keys), 

191 "mismatches": len(mismatches), 

192 "sample_mismatches": mismatches[:5], 

193 }, 

194 ) 

195 

196 return BenchmarkResult( 

197 name="activation_cache", 

198 severity=BenchmarkSeverity.INFO, 

199 message=f"All {len(common_keys)} cached activations match within tolerance", 

200 details={"cache_size": len(common_keys), "tolerance": tolerance}, 

201 ) 

202 

203 except Exception as e: 

204 return BenchmarkResult( 

205 name="activation_cache", 

206 severity=BenchmarkSeverity.ERROR, 

207 message=f"Activation cache check failed: {str(e)}", 

208 passed=False, 

209 )