Allow skipping warmup in bench_offline_throughput.py (#2103)

This commit is contained in:
Lianmin Zheng
2024-11-20 01:25:21 -08:00
committed by GitHub
parent 5942dfc00a
commit 3295cd8af2
2 changed files with 18 additions and 11 deletions

View File

@@ -156,9 +156,6 @@ class TpModelWorkerClient:
return logits_output, next_token_ids
def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
# A cuda stream sync here to avoid the cuda illegal memory access error.
torch.cuda.current_stream().synchronize()
# Create a new copy of sampling_info because it will be updated in-place by the scheduler for the next batch.
sampling_info = model_worker_batch.sampling_info
sampling_info.update_penalties()
@@ -169,6 +166,9 @@ class TpModelWorkerClient:
linear_penalties=sampling_info.linear_penalties,
)
# A cuda stream sync here to avoid the cuda illegal memory access error.
torch.cuda.current_stream().synchronize()
# Push a new batch to the queue
self.input_queue.put((model_worker_batch, self.future_token_ids_ct))