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
Support loading transformers models with named parameters (#16868)
Signed-off-by: Alex <alexwu@character.ai>
Alex Wu committed
Apr 28, 2025 at 15:15 UTC
6e74fd4945f1b0ccf72fda35daeedc8d6e632d0f
1 file changed
+23
vllm/model_executor/models/transformers.py
+23
@@ -166,6 +166,9 @@ class TransformersModel(nn.Module):
166
# Initialize buffers (e.g. rotary embedding inverse frequency)
167
self.init_buffers(self.model)
168
169
+ # Initialize parameters
170
+ self.init_parameters(self.model)
171
+
172
# Move remaining meta tensors to device (should happen last)
173
self.meta_to_empty(self.model)
174
@@ -298,6 +301,25 @@ class TransformersModel(nn.Module):
301
for child in module.children():
302
self.init_buffers(child)
303
304
+ def init_parameters(self, module: nn.Module):
305
+ """
306
+ If a `parameter` is on the `meta` device, then its parent
307
+ `module` is the original module created by:
308
+
309
+ ```python
310
+ with torch.device("meta"):
311
+ self.model: PreTrainedModel = AutoModel.from_config(...)
312
+ ```
313
+ """
314
+ for name, param in module.named_parameters(recurse=False):
315
+ if param.device == torch.device("meta"):
316
+ new_param = nn.Parameter(
317
+ torch.empty_like(param.data,
318
+ device=self.device_config.device))
319
+ setattr(module, name, new_param)
320
+ for child in module.children():
321
+ self.init_parameters(child)
322
+
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):
@@ -342,6 +364,7 @@ class TransformersModel(nn.Module):
364
def load_weights(self, weights: Iterable[tuple[str,
365
torch.Tensor]]) -> set[str]:
366
params_dict = dict(self.named_parameters())
367
+
368
loaded_params = set[str]()
369
for name, loaded_weight in weights:
370
# Use "model" instead of base_model_prefix because