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, 12:15 UTC
Transformers backend tweaks (#17365)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Apr 29, 2025 at 17:08 UTC
900edfa8d4081b99ac6891c6b3cc65076ee3f9af
1 file changed
+9
-13
vllm/model_executor/models/transformers.py
+9
-13
@@ -15,7 +15,6 @@
15
# limitations under the License.
16
"""Wrapper around `transformers` models"""
17
import re
18
-from itertools import chain
18
from typing import Iterable, Literal, Optional, Union
19
20
import torch
@@ -166,12 +165,9 @@ class TransformersModel(nn.Module):
165
# Initialize buffers (e.g. rotary embedding inverse frequency)
166
self.init_buffers(self.model)
167
169
- # Initialize parameters
168
+ # Initialize any parameters that have not had their modules replaced
169
self.init_parameters(self.model)
170
172
- # Move remaining meta tensors to device (should happen last)
173
- self.meta_to_empty(self.model)
174
-
171
self.make_empty_intermediate_tensors = (
172
make_empty_intermediate_tensors_factory(["hidden_states"],
173
config.hidden_size))
@@ -296,6 +292,14 @@ class TransformersModel(nn.Module):
292
"""
293
for name, buffer in module.named_buffers(recurse=False):
294
if buffer.device == torch.device("meta"):
295
+ if module == self.model:
296
+ logger.warning(
297
+ "To initialize buffers correctly, we instantiate the "
298
+ "parent module and and extract the value of the "
299
+ "buffer from it. In this case, the parent module is "
300
+ "the base model. Instantiating the entire model here "
301
+ "risks GPU OOM. Could this buffer be moved to a child "
302
+ "module?")
303
new_buffer = getattr(type(module)(self.config), name)
304
setattr(module, name, new_buffer)
305
for child in module.children():
@@ -320,14 +324,6 @@ class TransformersModel(nn.Module):
324
for child in module.children():
325
self.init_parameters(child)
326
323
- def meta_to_empty(self, module: nn.Module):
324
- tensors = list(chain(module.buffers(), module.parameters()))
325
- if tensors and all(t.device == torch.device("meta") for t in tensors):
326
- module.to_empty(device=self.device_config.device)
327
- return # We can stop recursing because to_empty is recursive
328
- for child in module.children():
329
- self.meta_to_empty(child)
330
-
327
def get_input_embeddings(self) -> nn.Module:
328
return self.model.get_input_embeddings()
329