[V1] Move OOM check into sampler run

[ywang96]

Move the OOM check from model warmup into _dummy_sampler_run since OOM can also happen during profile_run.

FIXES #12529

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ywang96.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ywang96]

FWIW - this is what it looks like when the OOM check triggers, which isn't too bad!

INFO 03-14 02:45:05 [gpu_model_runner.py:1434] Graph capturing finished in 22 secs, took 0.64 GiB
ERROR 03-14 02:45:06 [core.py:337] EngineCore hit an exception: Traceback (most recent call last):
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/worker/gpu_model_runner.py", line 1292, in _dummy_sampler_run
ERROR 03-14 02:45:06 [core.py:337]     sampler_output = self.model.sample(
ERROR 03-14 02:45:06 [core.py:337]                      ^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/model_executor/models/llama.py", line 544, in sample
ERROR 03-14 02:45:06 [core.py:337]     next_tokens = self.sampler(logits, sampling_metadata)
ERROR 03-14 02:45:06 [core.py:337]                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/tmp-nvme/myenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
ERROR 03-14 02:45:06 [core.py:337]     return self._call_impl(*args, **kwargs)
ERROR 03-14 02:45:06 [core.py:337]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/tmp-nvme/myenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
ERROR 03-14 02:45:06 [core.py:337]     return forward_call(*args, **kwargs)
ERROR 03-14 02:45:06 [core.py:337]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/sample/sampler.py", line 49, in forward
ERROR 03-14 02:45:06 [core.py:337]     sampled = self.sample(logits, sampling_metadata)
ERROR 03-14 02:45:06 [core.py:337]               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/sample/sampler.py", line 104, in sample
ERROR 03-14 02:45:06 [core.py:337]     random_sampled = self.topk_topp_sampler(
ERROR 03-14 02:45:06 [core.py:337]                      ^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/tmp-nvme/myenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
ERROR 03-14 02:45:06 [core.py:337]     return self._call_impl(*args, **kwargs)
ERROR 03-14 02:45:06 [core.py:337]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/tmp-nvme/myenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
ERROR 03-14 02:45:06 [core.py:337]     return forward_call(*args, **kwargs)
ERROR 03-14 02:45:06 [core.py:337]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/sample/ops/topk_topp_sampler.py", line 62, in forward_native
ERROR 03-14 02:45:06 [core.py:337]     logits = apply_top_k_top_p(logits, k, p)
ERROR 03-14 02:45:06 [core.py:337]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/sample/ops/topk_topp_sampler.py", line 94, in apply_top_k_top_p
ERROR 03-14 02:45:06 [core.py:337]     logits_sort, logits_idx = logits.sort(dim=-1, descending=False)
ERROR 03-14 02:45:06 [core.py:337]                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337] torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.47 GiB. GPU 0 has a total capacity of 79.10 GiB of which 1.21 GiB is free. Including non-PyTorch memory, this process has 0 bytes memory in use. Of the allocated memory 76.70 GiB is allocated by PyTorch, with 84.00 MiB allocated in private pools (e.g., CUDA Graphs), and 24.77 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
ERROR 03-14 02:45:06 [core.py:337] 
ERROR 03-14 02:45:06 [core.py:337] The above exception was the direct cause of the following exception:
ERROR 03-14 02:45:06 [core.py:337] 
ERROR 03-14 02:45:06 [core.py:337] Traceback (most recent call last):
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/engine/core.py", line 329, in run_engine_core
ERROR 03-14 02:45:06 [core.py:337]     engine_core = EngineCoreProc(*args, **kwargs)
ERROR 03-14 02:45:06 [core.py:337]                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/engine/core.py", line 284, in __init__
ERROR 03-14 02:45:06 [core.py:337]     super().__init__(vllm_config, executor_class, log_stats)
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/engine/core.py", line 62, in __init__
ERROR 03-14 02:45:06 [core.py:337]     num_gpu_blocks, num_cpu_blocks = self._initialize_kv_caches(
ERROR 03-14 02:45:06 [core.py:337]                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/engine/core.py", line 135, in _initialize_kv_caches
ERROR 03-14 02:45:06 [core.py:337]     self.model_executor.initialize_from_config(kv_cache_configs)
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/executor/abstract.py", line 63, in initialize_from_config
ERROR 03-14 02:45:06 [core.py:337]     self.collective_rpc("compile_or_warm_up_model")
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/executor/uniproc_executor.py", line 56, in collective_rpc
ERROR 03-14 02:45:06 [core.py:337]     answer = run_method(self.driver_worker, method, args, kwargs)
ERROR 03-14 02:45:06 [core.py:337]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/utils.py", line 2238, in run_method
ERROR 03-14 02:45:06 [core.py:337]     return func(*args, **kwargs)
ERROR 03-14 02:45:06 [core.py:337]            ^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/worker/gpu_worker.py", line 226, in compile_or_warm_up_model
ERROR 03-14 02:45:06 [core.py:337]     self.model_runner._dummy_sampler_run(
ERROR 03-14 02:45:06 [core.py:337]   File "/tmp-nvme/myenv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
ERROR 03-14 02:45:06 [core.py:337]     return func(*args, **kwargs)
ERROR 03-14 02:45:06 [core.py:337]            ^^^^^^^^^^^^^^^^^^^^^
ERROR 03-14 02:45:06 [core.py:337]   File "/home/jovyan/vllm/vllm/v1/worker/gpu_model_runner.py", line 1296, in _dummy_sampler_run
ERROR 03-14 02:45:06 [core.py:337]     raise RuntimeError(
ERROR 03-14 02:45:06 [core.py:337] RuntimeError: CUDA out of memory occurred when warming up sampler with 1024 dummy requests. Please try lowering `max_num_seqs` or `gpu_memory_utilization` when initializing the engine.