A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[V1] Move OOM check into sampler run (#14728)
Signed-off-by: Roger Wang <ywang@roblox.com> Co-authored-by: Simon Mo <simon.mo@hey.com>
Roger Wang committed
Mar 13, 2025 at 20:40 UTC
ad19c8a00313e9d3f6016ac16333ce0d817a9c2a
2 files changed
+17
-18
vllm/v1/worker/gpu_model_runner.py
+12
-3
@@ -1288,9 +1288,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1288
allowed_token_ids_mask=None,
1289
bad_words_token_ids={},
1290
)
1291
- sampler_output = self.model.sample(logits=logits,
1292
- sampling_metadata=dummy_metadata)
1293
-
1291
+ try:
1292
+ sampler_output = self.model.sample(
1293
+ logits=logits, sampling_metadata=dummy_metadata)
1294
+ except RuntimeError as e:
1295
+ if 'out of memory' in str(e):
1296
+ raise RuntimeError(
1297
+ "CUDA out of memory occurred when warming up sampler with "
1298
+ f"{num_reqs} dummy requests. Please try lowering "
1299
+ "`max_num_seqs` or `gpu_memory_utilization` when "
1300
+ "initializing the engine.") from e
1301
+ else:
1302
+ raise e
1303
return sampler_output
1304
1305
def profile_run(self) -> None:
vllm/v1/worker/gpu_worker.py
+5
-15
@@ -221,21 +221,11 @@ class Worker(WorkerBase):
221
# NOTE: This is called after `capture_model` on purpose to prevent
222
# memory buffers from being cleared by `torch.cuda.empty_cache`.
223
if get_pp_group().is_last_rank:
224
- try:
225
- max_num_reqs = min(
226
- self.scheduler_config.max_num_seqs,
227
- self.scheduler_config.max_num_batched_tokens)
228
- self.model_runner._dummy_sampler_run(
229
- hidden_states=self.model_runner._dummy_run(
230
- num_tokens=max_num_reqs))
231
- except RuntimeError as e:
232
- if 'out of memory' in str(e):
233
- raise RuntimeError(
234
- "CUDA out of memory occurred when warming up sampler. "
235
- "Please try lowering `gpu_memory_utilization` when "
236
- "initializing the engine.") from None
237
- else:
238
- raise e
224
+ max_num_reqs = min(self.scheduler_config.max_num_seqs,
225
+ self.scheduler_config.max_num_batched_tokens)
226
+ self.model_runner._dummy_sampler_run(
227
+ hidden_states=self.model_runner._dummy_run(
228
+ num_tokens=max_num_reqs))
229
230
# Reset the seed to ensure that the random state is not affected by
231
# the model initialization and profiling.