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][Core] Fix memory issue with logits & sampling (#13776)
Signed-off-by: Roger Wang <ywang@roblox.com>
Roger Wang committed
Mar 8, 2025 at 06:11 UTC
8d5aa466fbff258159ecb0ac85134f5356b5344c
3 files changed
+69
-29
tests/basic_correctness/test_cumem.py
+10
-1
@@ -142,7 +142,16 @@ def test_end_to_end(model: str, use_v1: bool):
142
used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
143
# now the memory usage is mostly cudagraph memory pool,
144
# and it should be less than the model weights (1B model, 2GiB weights)
145
- assert used_bytes < 2 * GiB_bytes
145
+
146
+ # NOTE: In V1, the memory buffer for logits (max_num_reqs x vocab_size)
147
+ # is captured but cannot be releasesd from PyTorch due to a known bug,
148
+ # therefore high memory usage after `llm.sleep` is called is expected.
149
+ # FIXME(youkaichao & ywang96): Fix memory buffer issue with sleep mode
150
+ # in V1.
151
+ if use_v1:
152
+ assert used_bytes < 7 * GiB_bytes
153
+ else:
154
+ assert used_bytes < 2 * GiB_bytes
155
156
llm.wake_up()
157
output2 = llm.generate(prompt, sampling_params)
vllm/v1/worker/gpu_model_runner.py
+38
-28
@@ -1238,6 +1238,42 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1238
)
1239
return hidden_states
1240
1241
+ @torch.inference_mode()
1242
+ def _dummy_sampler_run(
1243
+ self,
1244
+ hidden_states: torch.Tensor,
1245
+ ) -> torch.Tensor:
1246
+
1247
+ logits = self.model.compute_logits(hidden_states, None)
1248
+ num_reqs = logits.size(0)
1249
+
1250
+ dummy_tensors = lambda v: torch.full(
1251
+ (num_reqs, ), v, device=self.device)
1252
+
1253
+ dummy_metadata = SamplingMetadata(
1254
+ temperature=dummy_tensors(0.5),
1255
+ all_greedy=False,
1256
+ all_random=False,
1257
+ top_p=dummy_tensors(0.9),
1258
+ top_k=dummy_tensors(logits.size(1) - 1),
1259
+ min_p=None,
1260
+ generators={},
1261
+ max_num_logprobs=None,
1262
+ no_penalties=True,
1263
+ prompt_token_ids=None,
1264
+ frequency_penalties=dummy_tensors(0.1),
1265
+ presence_penalties=dummy_tensors(0.1),
1266
+ repetition_penalties=dummy_tensors(0.1),
1267
+ output_token_ids=[[] for _ in range(num_reqs)],
1268
+ min_tokens={},
1269
+ logit_bias=[None for _ in range(num_reqs)],
1270
+ allowed_token_ids_mask=None,
1271
+ )
1272
+ sampler_output = self.model.sample(logits=logits,
1273
+ sampling_metadata=dummy_metadata)
1274
+
1275
+ return sampler_output
1276
+
1277
def profile_run(self) -> None:
1278
# Profile with multimodal encoder & encoder cache.
1279
# TODO: handle encoder-decoder models once we support them.
@@ -1353,37 +1389,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1389
hidden_states = self._dummy_run(self.max_num_tokens)
1390
if get_pp_group().is_last_rank:
1391
hidden_states = hidden_states[logit_indices]
1356
- logits = self.model.compute_logits(hidden_states, None)
1357
- dummy_tensors = lambda v: torch.full(
1358
- (num_reqs, ), v, device=self.device)
1359
- dummy_metadata = SamplingMetadata(
1360
- temperature=dummy_tensors(0.5),
1361
- all_greedy=False,
1362
- all_random=False,
1363
- top_p=dummy_tensors(0.9),
1364
- top_k=dummy_tensors(logits.size(1) - 1),
1365
- min_p=None,
1366
- generators={},
1367
- max_num_logprobs=None,
1368
- no_penalties=True,
1369
- prompt_token_ids=torch.ones_like(logits,
1370
- dtype=torch.int64),
1371
- frequency_penalties=dummy_tensors(0.1),
1372
- presence_penalties=dummy_tensors(0.1),
1373
- repetition_penalties=dummy_tensors(0.1),
1374
- output_token_ids=[[] for _ in range(num_reqs)],
1375
- min_tokens={},
1376
- logit_bias=[None for _ in range(num_reqs)],
1377
- allowed_token_ids_mask=None,
1378
- )
1379
- sampler_output = self.model.sample(
1380
- logits=logits, sampling_metadata=dummy_metadata)
1392
+ sampler_output = self._dummy_sampler_run(hidden_states)
1393
else:
1382
- logits = None
1394
sampler_output = None
1384
- dummy_metadata = None
1395
torch.cuda.synchronize()
1386
- del hidden_states, logits, sampler_output, dummy_metadata
1396
+ del hidden_states, sampler_output
1397
self.encoder_cache.clear()
1398
gc.collect()
1399
vllm/v1/worker/gpu_worker.py
+21
@@ -119,6 +119,8 @@ class Worker(WorkerBase):
119
self.model_runner: GPUModelRunner = GPUModelRunner(
120
self.vllm_config, self.device)
121
122
+ # FIXME(youkaichao & ywang96): Use TorchDispatchMode instead of memory pool
123
+ # to hijack tensor allocation.
124
def load_model(self) -> None:
125
if self.vllm_config.model_config.enable_sleep_mode:
126
allocator = CuMemAllocator.get_instance()
@@ -211,6 +213,25 @@ class Worker(WorkerBase):
213
self.model_runner._dummy_run(size)
214
if not self.model_config.enforce_eager:
215
self.model_runner.capture_model()
216
+
217
+ # Warm up sampler and preallocate memory buffer for logits and other
218
+ # sampling related tensors of max possible shape to avoid memory
219
+ # fragmentation issue.
220
+ # NOTE: This is called after `capture_model` on purpose to prevent
221
+ # memory buffers from being cleared by `torch.cuda.empty_cache`.
222
+ try:
223
+ self.model_runner._dummy_sampler_run(
224
+ hidden_states=self.model_runner._dummy_run(
225
+ num_tokens=self.scheduler_config.max_num_seqs))
226
+ except RuntimeError as e:
227
+ if 'out of memory' in str(e):
228
+ raise RuntimeError(
229
+ "CUDA out of memory occurred when warming up sampler. "
230
+ "Please try lowering `gpu_memory_utilization` when "
231
+ "initializing the engine.") from None
232
+ else:
233
+ raise e
234
+
235
# Reset the seed to ensure that the random state is not affected by
236
# the model initialization and profiling.
237
set_random_seed(self.model_config.seed)