This function can be slow because it allocates tensors.
(self)
| 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): |
no test coverage detected