Allow skipping warmup in bench_offline_throughput.py (#2103)
This commit is contained in:
@@ -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))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user