A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Stream Llama4 weight loading to avoid host-OOM with copy-returning loaders (#44645)
Signed-off-by: Noa Neria <nneria@nvidia.com>
Noa Neria committed
Jun 15, 2026 at 05:23 UTC
1801fad0ba6238381430794d83d4c5540c2d73aa
2 files changed
+64
-72
vllm/model_executor/models/llama4.py
+6
-2
@@ -798,10 +798,14 @@ class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
798
self,
799
skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
800
)
801
- weights = [
801
+ # Use a generator (not a list comprehension) so the weights iterator is
802
+ # consumed lazily by AutoWeightsLoader. Materializing it here would hold
803
+ # the entire language-model checkpoint in host memory at once, which can
804
+ # OOM loaders that return private copies rather than mmap views.
805
+ weights = (
806
self.permute_qk_weight_for_rotary(name, loaded_weight)
807
for name, loaded_weight in weights
804
- ]
808
+ )
809
return loader.load_weights(weights)
810
811
def permute_qk_weight_for_rotary(
vllm/model_executor/models/mllama4.py
+58
-70
@@ -1131,66 +1131,6 @@ class Llama4ForConditionalGeneration(
1131
1132
return name
1133
1134
- def _separate_and_rename_weights(
1135
- self, weights: Iterable[tuple[str, torch.Tensor]]
1136
- ) -> tuple[list[tuple[str, torch.Tensor]], list[tuple[str, torch.Tensor]]]:
1137
- """Rename weights and separate them into language_model and other
1138
- weights."""
1139
- language_model_weights = []
1140
- other_weights = []
1141
-
1142
- for name, weight in weights:
1143
- renamed = self._rename_weight_for_modelopt_checkpoint(name)
1144
-
1145
- attr = renamed.split(".", 1)[0]
1146
- if isinstance(getattr(self, attr), StageMissingLayer):
1147
- continue
1148
-
1149
- if renamed.startswith("language_model."):
1150
- language_model_weights.append((renamed, weight))
1151
- else:
1152
- other_weights.append((renamed, weight))
1153
-
1154
- return language_model_weights, other_weights
1155
-
1156
- def _handle_expert_scale_broadcasting(
1157
- self, weights: list[tuple[str, torch.Tensor]], params_dict: dict
1158
- ) -> tuple[list[tuple[str, torch.Tensor]], set[str]]:
1159
- """Handle expert scale parameters that need broadcasting.
1160
-
1161
- ModelOpt checkpoints use a single value tensor scalar for BMM style
1162
- experts, vLLM expects the scale to be broadcasted across all experts.
1163
- """
1164
- regular_weights = []
1165
- expert_scale_weights = []
1166
- updated_params = set()
1167
-
1168
- for name, weight in weights:
1169
- # Check if this is an expert scale parameter that needs broadcasting
1170
- if (
1171
- "feed_forward.experts." in name
1172
- and "scale" in name
1173
- and ".shared_expert" not in name
1174
- ):
1175
- name = maybe_remap_moe_expert_param_name(name, params_dict)
1176
- if name in params_dict:
1177
- param = params_dict[name]
1178
- if (
1179
- hasattr(param, "data")
1180
- and param.data.numel() > 1
1181
- and weight.numel() == 1
1182
- ):
1183
- # Broadcast single value to all experts
1184
- param.data.fill_(weight.item())
1185
- updated_params.add(name)
1186
- continue
1187
-
1188
- expert_scale_weights.append((name, weight))
1189
- else:
1190
- regular_weights.append((name, weight))
1191
-
1192
- return regular_weights, expert_scale_weights, updated_params
1193
-
1134
def _load_other_weights(
1135
self,
1136
other_weights: Iterable[tuple[str, torch.Tensor]],
@@ -1251,19 +1191,67 @@ class Llama4ForConditionalGeneration(
1191
params_dict = dict(self.named_parameters())
1192
updated_params: set[str] = set()
1193
1254
- # Separate and rename weights
1255
- language_model_weights, other_weights = self._separate_and_rename_weights(
1256
- weights
1257
- )
1194
+ # Stream thelanguage-model weights straight into
1195
+ # AutoWeightsLoader so each tensor is loaded and released as we iterate,
1196
+ # instead of materializing the whole checkpoint in host memory first.
1197
+ # Only the small vision/projector and scalar expert-scale groups are
1198
+ # buffered.
1199
+ other_weights: list[tuple[str, torch.Tensor]] = []
1200
+ expert_scale_weights: list[tuple[str, torch.Tensor]] = []
1201
+
1202
+ def regular_language_model_weights() -> Iterable[tuple[str, torch.Tensor]]:
1203
+ """Rename weights and separate them into language_model and other
1204
+ weights.
1205
+
1206
+ Yields the (large) language_model weights for streaming; the small
1207
+ groups (vision/projector and scalar expert scales) are buffered into
1208
+ the lists above.
1209
+ """
1210
+ for name, weight in weights:
1211
+ renamed = self._rename_weight_for_modelopt_checkpoint(name)
1212
+
1213
+ attr = renamed.split(".", 1)[0]
1214
+ if isinstance(getattr(self, attr), StageMissingLayer):
1215
+ continue
1216
1259
- # Handle expert scale parameters
1260
- regular_weights, expert_scale_weights, updated_params_from_experts = (
1261
- self._handle_expert_scale_broadcasting(language_model_weights, params_dict)
1262
- )
1263
- updated_params.update(updated_params_from_experts)
1217
+ if not renamed.startswith("language_model."):
1218
+ other_weights.append((renamed, weight))
1219
+ continue
1220
+
1221
+ # Handle expert scale parameters that need broadcasting.
1222
+ # ModelOpt checkpoints use a single value tensor scalar for BMM
1223
+ # style experts, vLLM expects the scale to be broadcasted across
1224
+ # all experts.
1225
+ if (
1226
+ "feed_forward.experts." in renamed
1227
+ and "scale" in renamed
1228
+ and ".shared_expert" not in renamed
1229
+ ):
1230
+ renamed = maybe_remap_moe_expert_param_name(renamed, params_dict)
1231
+ if renamed in params_dict:
1232
+ param = params_dict[renamed]
1233
+ if (
1234
+ hasattr(param, "data")
1235
+ and param.data.numel() > 1
1236
+ and weight.numel() == 1
1237
+ ):
1238
+ # Broadcast single value to all experts
1239
+ param.data.fill_(weight.item())
1240
+ updated_params.add(renamed)
1241
+ continue
1242
+
1243
+ expert_scale_weights.append((renamed, weight))
1244
+ continue
1245
+
1246
+ yield renamed, weight
1247
1248
loader = AutoWeightsLoader(self)
1266
- loaded_language_model_params = loader.load_weights(regular_weights)
1249
+ # AutoWeightsLoader consumes its input lazily and runs to exhaustion,
1250
+ # so other_weights / expert_scale_weights are fully populated as a side
1251
+ # effect by the time this returns.
1252
+ loaded_language_model_params = loader.load_weights(
1253
+ regular_language_model_weights()
1254
+ )
1255
assert loaded_language_model_params is not None
1256
updated_params.update(loaded_language_model_params)
1257