Coverage for transformer_lens/benchmarks/generation.py: 48%
65 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Generation and KV cache benchmarks for TransformerBridge."""
3from typing import Any, Optional
5from transformer_lens import HookedTransformer
6from transformer_lens.benchmarks.utils import (
7 BenchmarkResult,
8 BenchmarkSeverity,
9 deterministic_rng,
10 is_tiny_test_model,
11)
12from transformer_lens.model_bridge import TransformerBridge
15def resolve_text_generator(bridge: Any):
16 """Return the callable that produces text for this architecture, or None
17 (diffusion LMs delegate to their native sampler)."""
18 if getattr(bridge.adapter, "supports_generation", True):
19 return bridge.generate
20 if getattr(bridge.adapter, "native_sampler", None):
21 return bridge.diffusion_generate
22 return None
25def benchmark_generation(
26 bridge: TransformerBridge,
27 test_text: str,
28 max_new_tokens: int = 10,
29 reference_model: Optional[HookedTransformer] = None,
30) -> BenchmarkResult:
31 """Benchmark basic text generation."""
32 try:
33 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""): 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true
34 return BenchmarkResult(
35 name="generation",
36 severity=BenchmarkSeverity.INFO,
37 message="Skipped for tiny/test model (random weights produce degenerate generation)",
38 )
39 generator = resolve_text_generator(bridge)
40 if generator is None: 40 ↛ 41line 40 didn't jump to line 41 because the condition on line 40 was never true
41 return BenchmarkResult(
42 name="generation",
43 severity=BenchmarkSeverity.INFO,
44 message="Skipped: architecture supports no text generation",
45 )
46 # Greedy (deterministic, no seeding): tests the loop's mechanics, not
47 # sampling quality. stop_at_eos=False because some models argmax EOS on a
48 # bare prompt (EXAONE-4, raw HF) — a choice to stop, not a stall. BOS is
49 # left alone: forcing it would override default_prepend_bos=False adapters
50 # and derail checkpoints whose BOS token is their EOS.
51 gen_kwargs: dict[str, Any] = {"max_new_tokens": max_new_tokens, "temperature": 0.0}
52 if getattr(bridge.adapter, "supports_generation", True): 52 ↛ 55line 52 didn't jump to line 55 because the condition on line 52 was always true
53 # Native diffusion samplers take neither kwarg through **kwargs.
54 gen_kwargs["stop_at_eos"] = False
55 output = generator(test_text, **gen_kwargs)
57 if not isinstance(output, str): 57 ↛ 58line 57 didn't jump to line 58 because the condition on line 57 was never true
58 return BenchmarkResult(
59 name="generation",
60 severity=BenchmarkSeverity.DANGER,
61 message="Generated output is not a string",
62 passed=False,
63 )
65 # Check token count instead of character count to handle whitespace-only generation
66 input_tokens = bridge.to_tokens(test_text)
67 output_tokens = bridge.to_tokens(output)
69 # Strip leading BOS token if present for fair comparison
70 input_len = input_tokens.shape[-1]
71 output_len = output_tokens.shape[-1]
73 # Compare like with like. decode->re-tokenize is lossy when the prompt
74 # holds out-of-vocabulary characters (a DNA model given English drops
75 # them as [UNK]), which reads as "no new tokens" though generation ran.
76 prompt_roundtrip = test_text
77 if bridge.tokenizer is not None: 77 ↛ 79line 77 didn't jump to line 79 because the condition on line 77 was always true
78 prompt_roundtrip = bridge.tokenizer.decode(input_tokens[0], skip_special_tokens=True)
79 if output_len <= input_len and len(output) <= len(prompt_roundtrip):
80 return BenchmarkResult(
81 name="generation",
82 severity=BenchmarkSeverity.DANGER,
83 message="Generated text has no new tokens",
84 details={
85 "input_tokens": input_len,
86 "output_tokens": output_len,
87 "input_chars": len(test_text),
88 "output_chars": len(output),
89 },
90 passed=False,
91 )
93 return BenchmarkResult(
94 name="generation",
95 severity=BenchmarkSeverity.INFO,
96 message=f"Generation successful: {input_len} -> {output_len} tokens ({len(test_text)} -> {len(output)} chars)",
97 details={
98 "input_tokens": input_len,
99 "output_tokens": output_len,
100 "input_chars": len(test_text),
101 "output_chars": len(output),
102 "max_new_tokens": max_new_tokens,
103 },
104 )
106 except Exception as e:
107 return BenchmarkResult(
108 name="generation",
109 severity=BenchmarkSeverity.ERROR,
110 message=f"Generation failed: {str(e)}",
111 passed=False,
112 )
115def benchmark_generation_with_kv_cache(
116 bridge: TransformerBridge,
117 test_text: str,
118 max_new_tokens: int = 10,
119 reference_model: Optional[HookedTransformer] = None,
120) -> BenchmarkResult:
121 """Benchmark text generation with KV caching enabled.
123 This ensures that the KV cache is properly passed through attention layers
124 during generation, and that the cache update logic works correctly.
126 Args:
127 bridge: TransformerBridge model to test
128 test_text: Input text for generation
129 max_new_tokens: Number of tokens to generate
130 reference_model: Optional HookedTransformer reference model (not used)
132 Returns:
133 BenchmarkResult with generation details
134 """
135 try:
136 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
137 return BenchmarkResult(
138 name="generation_with_kv_cache",
139 severity=BenchmarkSeverity.INFO,
140 message="Skipped for tiny/test model (random weights produce degenerate generation)",
141 )
143 # Cache-free architectures (RWKV, HyenaDNA) would "pass" this vacuously —
144 # the recompute path produces output while exercising no cache at all.
145 # Diffusion samplers have no KV cache concept either.
146 if not getattr(bridge.adapter, "supports_kv_cache", True) or not getattr(
147 bridge.adapter, "supports_generation", True
148 ):
149 return BenchmarkResult(
150 name="generation_with_kv_cache",
151 severity=BenchmarkSeverity.INFO,
152 message="Skipped: architecture has no KV cache (generation recomputes each step)",
153 )
155 # Generate with KV cache (should be enabled by default for max_new_tokens > 1)
156 with deterministic_rng():
157 output = bridge.generate(
158 test_text,
159 max_new_tokens=max_new_tokens,
160 temperature=0.7,
161 prepend_bos=True,
162 )
164 if output is None or len(output) == 0:
165 return BenchmarkResult(
166 name="generation_with_kv_cache",
167 severity=BenchmarkSeverity.DANGER,
168 message="Generation with KV cache produced no output",
169 passed=False,
170 )
172 return BenchmarkResult(
173 name="generation_with_kv_cache",
174 severity=BenchmarkSeverity.INFO,
175 message=f"KV cache generation successful ({len(output)} chars)",
176 details={"output_len": len(output), "max_new_tokens": max_new_tokens},
177 )
179 except Exception as e:
180 return BenchmarkResult(
181 name="generation_with_kv_cache",
182 severity=BenchmarkSeverity.ERROR,
183 message=f"KV cache generation failed: {str(e)}",
184 passed=False,
185 )
188def benchmark_multiple_generation_calls(
189 bridge: TransformerBridge,
190 test_prompts: list,
191 max_new_tokens: int = 5,
192 reference_model: Optional[HookedTransformer] = None,
193) -> BenchmarkResult:
194 """Benchmark multiple generation calls to ensure KV cache handling is robust.
196 Args:
197 bridge: TransformerBridge model to test
198 test_prompts: List of input prompts for generation
199 max_new_tokens: Number of tokens to generate per prompt
200 reference_model: Optional HookedTransformer reference model (not used)
202 Returns:
203 BenchmarkResult with multiple generation details
204 """
205 try:
206 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
207 return BenchmarkResult(
208 name="multiple_generation_calls",
209 severity=BenchmarkSeverity.INFO,
210 message="Skipped for tiny/test model (random weights produce degenerate generation)",
211 )
213 generator = resolve_text_generator(bridge)
214 if generator is None:
215 return BenchmarkResult(
216 name="multiple_generation_calls",
217 severity=BenchmarkSeverity.INFO,
218 message="Skipped: architecture supports no text generation",
219 )
221 outputs = []
222 with deterministic_rng():
223 for prompt in test_prompts:
224 output = generator(
225 prompt,
226 max_new_tokens=max_new_tokens,
227 temperature=0.7,
228 prepend_bos=True,
229 )
230 if output is None or len(output) == 0:
231 return BenchmarkResult(
232 name="multiple_generation_calls",
233 severity=BenchmarkSeverity.DANGER,
234 message=f"Generation failed for prompt: {prompt[:50]}...",
235 passed=False,
236 )
237 outputs.append(output)
239 return BenchmarkResult(
240 name="multiple_generation_calls",
241 severity=BenchmarkSeverity.INFO,
242 message=f"All {len(test_prompts)} generation calls successful",
243 details={
244 "prompt_count": len(test_prompts),
245 "max_new_tokens": max_new_tokens,
246 "output_lens": [len(out) for out in outputs],
247 },
248 )
250 except Exception as e:
251 return BenchmarkResult(
252 name="multiple_generation_calls",
253 severity=BenchmarkSeverity.ERROR,
254 message=f"Multiple generation calls failed: {str(e)}",
255 passed=False,
256 )