Fix llama2 benchmarks (Lightning-AI#2767)

riccardofelluga · web-flow · commit fb989d485cf5 · 2025-12-12T14:28:18.000Z
diff --git a/thunder/benchmarks/__init__.py b/thunder/benchmarks/__init__.py
@@ -2081,7 +2081,7 @@ def fn(self) -> Callable:
         from litgpt.model import CausalSelfAttention
 
         module = (
-            CausalSelfAttention(self.config)
+            CausalSelfAttention(self.config, 0)
             .to(device=self.device, dtype=self.tdtype)
             .requires_grad_(self.requires_grad)
         )

Original file line number	Diff line number	Diff line change
`@@ -2081,7 +2081,7 @@ def fn(self) -> Callable:`
`2081`	`2081`	`from litgpt.model import CausalSelfAttention`
`2082`	`2082`
`2083`	`2083`	`module = (`
`2084`		`- CausalSelfAttention(self.config)`
	`2084`	`+ CausalSelfAttention(self.config, 0)`
`2085`	`2085`	`.to(device=self.device, dtype=self.tdtype)`
`2086`	`2086`	`.requires_grad_(self.requires_grad)`
`2087`	`2087`	`)`