(self, requests=None, pos=0)
| 174 | pass |
| 175 | |
| 176 | def _detect_batch_size(self, requests=None, pos=0): |
| 177 | if requests: |
| 178 | _, context_enc, continuation_enc = requests[pos] |
| 179 | max_length = len( |
| 180 | (context_enc + continuation_enc)[-(self.max_length + 1) :][:-1] |
| 181 | ) |
| 182 | else: |
| 183 | max_length = self.max_length |
| 184 | |
| 185 | # if OOM, then halves batch_size and tries again |
| 186 | @find_executable_batch_size(starting_batch_size=self.max_batch_size) |
| 187 | def forward_batch(batch_size): |
| 188 | test_batch = torch.ones((batch_size, max_length), device=self.device).long() |
| 189 | for _ in range(5): |
| 190 | _ = F.log_softmax(self._model_call(test_batch), dim=-1).cpu() |
| 191 | return batch_size |
| 192 | |
| 193 | batch_size = forward_batch() |
| 194 | utils.clear_torch_cache() |
| 195 | |
| 196 | return batch_size |
| 197 | |
| 198 | # subclass must implement properties vocab_size, eot_token_id, max_gen_toks, batch_size, device, max_length. |
| 199 | # TODO: enforce this somehow |
no outgoing calls
no test coverage detected