modify perplexity test to use sdpa

dan-garvey · dan-garvey · commit 5782f8b3994e · 2024-10-31T20:58:55.000-05:00
diff --git a/sharktank/sharktank/evaluate/perplexity_torch.py b/sharktank/sharktank/evaluate/perplexity_torch.py
@@ -111,7 +111,7 @@ def load_model(self, dataset, tokenizer, tensor_parallelism_size, attention_kern
             attention_dtype=self.attention_dtype,
             tensor_parallelism_size=tensor_parallelism_size,
         )
-
+        config.attention_kernel="torch"
         if config.tensor_parallelism_size > 1:
             dataset.root_theta = shard_theta(dataset.root_theta, config)
 

Original file line number	Diff line number	Diff line change
`@@ -111,7 +111,7 @@ def load_model(self, dataset, tokenizer, tensor_parallelism_size, attention_kern`
`111`	`111`	`attention_dtype=self.attention_dtype,`
`112`	`112`	`tensor_parallelism_size=tensor_parallelism_size,`
`113`	`113`	`)`
`114`		`-`
	`114`	`+ config.attention_kernel="torch"`
`115`	`115`	`if config.tensor_parallelism_size > 1:`
`116`	`116`	`dataset.root_theta = shard_theta(dataset.root_theta, config)`
`117`	`117`