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 4, 11:15 UTC
[UX][Bugfix] Fix OOM by setting PyTorch `max_split_size_mb` during model loading (#41268)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Apr 30, 2026 at 10:46 UTC
f03d82efdd88fbd85ddf7a5475e237ae3abaf01e
1 file changed
+28
-1
vllm/v1/worker/gpu_worker.py
+28
-1
@@ -5,12 +5,13 @@
5
import gc
6
import os
7
from collections.abc import Callable
8
-from contextlib import AbstractContextManager, nullcontext
8
+from contextlib import AbstractContextManager, contextmanager, nullcontext
9
from datetime import timedelta
10
from types import NoneType
11
from typing import TYPE_CHECKING, Any
12
13
import numpy as np
14
+import regex as re
15
import torch
16
import torch.nn as nn
17
@@ -208,6 +209,30 @@ class Worker(WorkerBase):
209
)
210
return allocator.use_memory_pool(tag=tag)
211
212
+ @contextmanager
213
+ def _scoped_allocator_max_split(self, max_split_size_mb: int):
214
+ """Temporarily set max_split_size_mb to reduce allocator fragmentation at the
215
+ cost of more cudaMalloc calls (negligible in practice). Restores the original
216
+ value on exit."""
217
+ if not current_platform.is_cuda():
218
+ yield
219
+ return
220
+
221
+ conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
222
+ match = re.search(r"max_split_size_mb:(\d+)", conf)
223
+ original_value = match.group(1) if match else None
224
+
225
+ torch._C._accelerator_setAllocatorSettings(
226
+ f"max_split_size_mb:{max_split_size_mb}"
227
+ )
228
+ try:
229
+ yield
230
+ finally:
231
+ # PyTorch defaults to SIZE_MAX (no limit).
232
+ _SIZE_MAX_MB = (2**64 - 1) // (1024 * 1024)
233
+ restore = original_value if original_value else str(_SIZE_MAX_MB)
234
+ torch._C._accelerator_setAllocatorSettings(f"max_split_size_mb:{restore}")
235
+
236
@instrument(span_name="Init device")
237
def init_device(self):
238
if self.device_config.device_type == "cuda":
@@ -312,6 +337,8 @@ class Worker(WorkerBase):
337
with (
338
self._maybe_get_memory_pool_context(tag="weights"),
339
set_current_vllm_config(self.vllm_config),
340
+ # 20 MiB is the minimum PyTorch allows for max_split_size_mb.
341
+ self._scoped_allocator_max_split(max_split_size_mb=20),
342
):
343
self.model_runner.load_model(load_dummy_weights=load_dummy_weights)
344