| 61 | |
| 62 | @dataclass |
| 63 | class Case: |
| 64 | name: str |
| 65 | n_runs_benchmark: int |
| 66 | n_runs_warmup: int |
| 67 | n_hidden_states: int |
| 68 | n_samples: int |
| 69 | vocab_size: int |
| 70 | hidden_size: int |
| 71 | top_k: int | None = None |
| 72 | top_p: float | None = None |
| 73 | tp: TPInfo = TP1 |
| 74 | |
| 75 | def make_fn_kwargs(self) -> dict: |
| 76 | """This function can be slow because it allocates tensors.""" |
| 77 | weights = torch.randn( |
| 78 | (self.vocab_size, self.hidden_size), dtype=torch.bfloat16, device=device |
| 79 | ) |
| 80 | weights = shard_weights(weights, self.tp) |
| 81 | kwargs = dict( |
| 82 | hidden_states=torch.randn( |
| 83 | (self.n_hidden_states, self.hidden_size), dtype=torch.bfloat16, device=device |
| 84 | ), |
| 85 | weights=weights, |
| 86 | num_samples=self.n_samples, |
| 87 | temperature=torch.tensor(1.0, device=device), |
| 88 | tp=self.tp, |
| 89 | ) |
| 90 | if self.top_k is not None: |
| 91 | kwargs["top_k"] = self.top_k |
| 92 | if self.top_p is not None: |
| 93 | kwargs["top_p"] = self.top_p |
| 94 | return kwargs |
| 95 | |
| 96 | |
| 97 | def _prepare_case(case: Case): |