| 116 | |
| 117 | |
| 118 | def run_hf( |
| 119 | requests: list[SampleRequest], |
| 120 | model: str, |
| 121 | tokenizer: PreTrainedTokenizerBase, |
| 122 | n: int, |
| 123 | max_batch_size: int, |
| 124 | trust_remote_code: bool, |
| 125 | disable_detokenize: bool = False, |
| 126 | ) -> float: |
| 127 | llm = AutoModelForCausalLM.from_pretrained(model, torch_dtype=torch.float16, trust_remote_code=trust_remote_code) |
| 128 | if llm.config.model_type == "llama": |
| 129 | # To enable padding in the HF backend. |
| 130 | tokenizer.pad_token = tokenizer.eos_token |
| 131 | llm = llm.cuda() |
| 132 | |
| 133 | pbar = tqdm(total=len(requests)) |
| 134 | start = time.perf_counter() |
| 135 | batch: list[str] = [] |
| 136 | max_prompt_len = 0 |
| 137 | max_output_len = 0 |
| 138 | for i in range(len(requests)): |
| 139 | prompt = requests[i].prompt |
| 140 | prompt_len = requests[i].prompt_len |
| 141 | output_len = requests[i].expected_output_len |
| 142 | # Add the prompt to the batch. |
| 143 | batch.append(prompt) |
| 144 | max_prompt_len = max(max_prompt_len, prompt_len) |
| 145 | max_output_len = max(max_output_len, output_len) |
| 146 | if len(batch) < max_batch_size and i != len(requests) - 1: |
| 147 | # Check if we can add more requests to the batch. |
| 148 | next_prompt_len = requests[i + 1].prompt_len |
| 149 | next_output_len = requests[i + 1].expected_output_len |
| 150 | if (max(max_prompt_len, next_prompt_len) + max(max_output_len, next_output_len)) <= 2048: |
| 151 | # We can add more requests to the batch. |
| 152 | continue |
| 153 | |
| 154 | # Generate the sequences. |
| 155 | input_ids = tokenizer(batch, return_tensors="pt", padding=True).input_ids |
| 156 | llm_outputs = llm.generate( |
| 157 | input_ids=input_ids.cuda(), |
| 158 | do_sample=True, |
| 159 | num_return_sequences=n, |
| 160 | temperature=1.0, |
| 161 | top_p=1.0, |
| 162 | use_cache=True, |
| 163 | max_new_tokens=max_output_len, |
| 164 | ) |
| 165 | if not disable_detokenize: |
| 166 | # Include the decoding time. |
| 167 | tokenizer.batch_decode(llm_outputs, skip_special_tokens=True) |
| 168 | pbar.update(len(batch)) |
| 169 | |
| 170 | # Clear the batch. |
| 171 | batch = [] |
| 172 | max_prompt_len = 0 |
| 173 | max_output_len = 0 |
| 174 | end = time.perf_counter() |
| 175 | return end - start |