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, 09:45 UTC
[Bugfix][Model] Make Olmo2Model weight loading return loaded weights (#18504)
Signed-off-by: Shane A <shanea@allenai.org>
Shane A committed
May 21, 2025 at 21:17 UTC
51797775c3ffc1277be750eb046c8030f8eca280
1 file changed
+5
-1
vllm/model_executor/models/olmo2.py
+5
-1
@@ -314,7 +314,8 @@ class Olmo2Model(nn.Module):
314
hidden_states = self.norm(hidden_states)
315
return hidden_states
316
317
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
317
+ def load_weights(self, weights: Iterable[tuple[str,
318
+ torch.Tensor]]) -> set[str]:
319
stacked_params_mapping = [
320
# (param_name, shard_name, shard_id)
321
("qkv_proj", "q_proj", "q"),
@@ -325,6 +326,7 @@ class Olmo2Model(nn.Module):
326
]
327
328
params_dict = dict(self.named_parameters(remove_duplicate=False))
329
+ loaded_params: set[str] = set()
330
for name, loaded_weight in weights:
331
if is_pp_missing_parameter(name, self):
332
continue
@@ -347,6 +349,8 @@ class Olmo2Model(nn.Module):
349
weight_loader = getattr(param, "weight_loader",
350
default_weight_loader)
351
weight_loader(param, loaded_weight)
352
+ loaded_params.add(name)
353
+ return loaded_params
354
355
356
class Olmo2ForCausalLM(nn.Module, SupportsPP):