@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4afe687a82
Fork 0

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, 02:15 UTC

Enable ModelOpt Llama4 fp8 checkpoint deployment (#20419)

Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>

Zhiyu committed Jul 11, 2025 at 23:07 UTC 4afe687a8291397c9744187fbf4b01a79b89564a
5 files changed +501 -35
vllm/model_executor/layers/fused_moe/layer.py
+31 -6
@@ -81,6 +81,16 @@ class FusedMoEMethodBase(QuantizeMethodBase):
81 params_dtype: torch.dtype, **extra_weight_attrs):
82 raise NotImplementedError
83
84 + def uses_weight_scale_2_pattern(self) -> bool:
85 + """
86 + Returns True if this quantization method uses 'weight_scale_2' pattern
87 + for per-tensor weight scales (e.g., FP4 variants), False otherwise.
88 +
89 + This method should be overridden by subclasses that use the
90 + 'weight_scale_2' pattern instead of the standard 'weight_scale' pattern.
91 + """
92 + return False
93 +
94 @staticmethod
95 def maybe_make_prepare_finalize(
96 moe: FusedMoEConfig) -> Optional[FusedMoEPrepareAndFinalize]:
@@ -1081,12 +1091,23 @@ class FusedMoE(torch.nn.Module):
1091
1092 # TODO @dsikka: ModelOpt should follow the proper MoE loading pattern
1093 if "ModelOpt" in quant_method_name:
1084 - if ('weight_scale_2' in weight_name
1085 - or 'input_scale' in weight_name):
1086 - self._load_per_tensor_weight_scale(shard_id=shard_id,
1087 - param=param,
1088 - loaded_weight=loaded_weight,
1089 - expert_id=expert_id)
1094 + # Determine per-tensor weight scale patterns based on variant
1095 + # Use the dedicated method instead of brittle string matching
1096 + uses_weight_scale_2 = self.quant_method.uses_weight_scale_2_pattern(
1097 + )
1098 +
1099 + # For per-tensor, FP4 uses "weight_scale_2", FP8 uses "weight_scale"
1100 + per_tensor_conditions = (
1101 + "weight_scale_2" in weight_name if uses_weight_scale_2 else
1102 + "weight_scale" in weight_name) or "input_scale" in weight_name
1103 +
1104 + if per_tensor_conditions:
1105 + self._load_per_tensor_weight_scale(
1106 + shard_id=shard_id,
1107 + param=param,
1108 + loaded_weight=loaded_weight,
1109 + expert_id=expert_id,
1110 + )
1111 elif "weight" in weight_name:
1112 self._load_model_weight_or_group_weight_scale(
1113 shard_id=shard_id,
@@ -1558,3 +1579,7 @@ direct_register_custom_op(
1579 dispatch_key=current_platform.dispatch_key,
1580 tags=(torch.Tag.needs_fixed_stride_order, ),
1581 )
1582 +
1583 +# Mark the FusedMoE weight_loader as supporting MoE-specific parameters
1584 +# to avoid expensive runtime reflection in model loading code
1585 +FusedMoE.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
vllm/model_executor/layers/quantization/modelopt.py
+261 -5
@@ -42,9 +42,13 @@ class ModelOptFp8Config(QuantizationConfig):
42 def __init__(
43 self,
44 is_checkpoint_fp8_serialized: bool = False,
45 + kv_cache_quant_method: Optional[str] = None,
46 + exclude_modules: Optional[list[str]] = None,
47 ) -> None:
48 super().__init__()
49 self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
50 + self.kv_cache_quant_method = kv_cache_quant_method
51 + self.exclude_modules = exclude_modules
52 if is_checkpoint_fp8_serialized:
53 logger.warning("Detected ModelOpt fp8 checkpoint. Please note that"
54 " the format is experimental and could change.")
@@ -69,6 +73,11 @@ class ModelOptFp8Config(QuantizationConfig):
73 def from_config(cls, config: dict[str, Any]) -> "ModelOptFp8Config":
74 quant_config = cls.get_from_keys(config, ["quantization"])
75 quant_method = quant_config["quant_algo"]
76 + kv_cache_quant_method = cls.get_from_keys(
77 + config, ["quantization"]).get("kv_cache_quant_algo")
78 + exclude_modules = cls.get_from_keys(
79 + config, ["quantization"]).get("exclude_modules")
80 +
81 if quant_method not in QUANT_ALGOS:
82 raise ValueError(f"ModelOpt currently only supports: {QUANT_ALGOS}"
83 " quantizations in vLLM. Please check the "
@@ -76,27 +85,51 @@ class ModelOptFp8Config(QuantizationConfig):
85 "quant configuration.")
86 is_checkpoint_fp8_serialized = ("FP8" in quant_method)
87
79 - return cls(is_checkpoint_fp8_serialized)
88 + return cls(is_checkpoint_fp8_serialized, kv_cache_quant_method,
89 + exclude_modules)
90 +
91 + def is_layer_excluded(self, prefix: str) -> bool:
92 + """
93 + Check if a layer should be excluded from quantization.
94 +
95 + This method handles both regular models and multimodal models that use
96 + the language_model prefix. For multimodal models, it checks if the
97 + module name (without the language_model prefix) is in the exclude list.
98 + """
99 + if self.exclude_modules is None:
100 + return False
101 +
102 + # Check if any excluded module matches the prefix
103 + for module in self.exclude_modules:
104 + if (module in prefix
105 + or (prefix.startswith("language_model.")
106 + and module in prefix.removeprefix("language_model."))):
107 + return True
108 + return False
109
110 def get_quant_method(self, layer: torch.nn.Module,
111 prefix: str) -> Optional["QuantizeMethodBase"]:
112 from vllm.attention.layer import Attention # Avoid circular import
113 if isinstance(layer, LinearBase):
114 + if self.is_layer_excluded(prefix):
115 + return UnquantizedLinearMethod()
116 return ModelOptFp8LinearMethod(self)
117 elif isinstance(layer, Attention):
118 return ModelOptFp8KVCacheMethod(self)
119 + elif isinstance(layer, FusedMoE):
120 + return ModelOptFp8MoEMethod(self)
121 return None
122
123
124 class ModelOptFp8LinearMethod(LinearMethodBase):
125 """Linear method for Model Optimizer static quantization.
126 Supports loading FP8 checkpoints with static weight scale and
94 - activation scale. Future support might be added for dynamic
127 + activation scale. Future support might be added for dynamic
128 scales.
129
130 Limitations:
131 1. Only support per-tensor quantization due to torch._scaled_mm support.
99 - 2. Only support float8_e4m3fn datatype
132 + 2. Only support float8_e4m3fn datatype
133 Args: quant_config: The ModelOpt quantization config.
134 """
135
@@ -172,6 +205,223 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
205 bias=bias)
206
207
208 +class ModelOptFp8MoEMethod(FusedMoEMethodBase):
209 + """MoE method for ModelOpt FP8.
210 + Supports loading FP8 checkpoints with static weight scale and
211 + activation scale.
212 + Args:
213 + quant_config: The ModelOpt quantization config.
214 + """
215 +
216 + def __init__(self, quant_config: ModelOptFp8Config):
217 + self.quant_config = quant_config
218 + from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
219 + cutlass_fp8_supported)
220 + self.cutlass_fp8_supported = cutlass_fp8_supported()
221 +
222 + def create_weights(
223 + self,
224 + layer: torch.nn.Module,
225 + num_experts: int,
226 + hidden_size: int,
227 + intermediate_size_per_partition: int,
228 + params_dtype: torch.dtype,
229 + **extra_weight_attrs,
230 + ):
231 +
232 + # Use FP8 dtype if checkpoint is serialized
233 + weight_dtype = (torch.float8_e4m3fn
234 + if self.quant_config.is_checkpoint_fp8_serialized else
235 + params_dtype)
236 + weight_loader = extra_weight_attrs.get("weight_loader")
237 +
238 + w13_weight = ModelWeightParameter(
239 + data=torch.empty(num_experts,
240 + 2 * intermediate_size_per_partition,
241 + hidden_size,
242 + dtype=weight_dtype),
243 + input_dim=2,
244 + output_dim=1,
245 + weight_loader=weight_loader,
246 + )
247 + layer.register_parameter("w13_weight", w13_weight)
248 +
249 + w2_weight = ModelWeightParameter(
250 + data=torch.empty(num_experts,
251 + hidden_size,
252 + intermediate_size_per_partition,
253 + dtype=weight_dtype),
254 + input_dim=2,
255 + output_dim=1,
256 + weight_loader=weight_loader,
257 + )
258 + layer.register_parameter("w2_weight", w2_weight)
259 +
260 + if self.quant_config.is_checkpoint_fp8_serialized:
261 + # WEIGHT SCALES - Per-tensor scaling for ModelOpts
262 + # Allocate 2 scales for w1 and w3 respectively.
263 + # They will be combined to a single scale after weight loading.
264 + w13_weight_scale = PerTensorScaleParameter(
265 + data=torch.full(
266 + (num_experts, 2),
267 + 1.0,
268 + dtype=torch.float32,
269 + ),
270 + weight_loader=weight_loader,
271 + )
272 + w2_weight_scale = PerTensorScaleParameter(
273 + data=torch.full((num_experts, ), 1.0, dtype=torch.float32),
274 + weight_loader=weight_loader,
275 + )
276 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
277 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
278 +
279 + # Set weight loader attributes for scales
280 + extra_weight_attrs.update(
281 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
282 +
283 + # INPUT SCALES - Per-tensor scaling for ModelOpt
284 + w13_input_scale = PerTensorScaleParameter(
285 + data=torch.full((num_experts, ), 1.0, dtype=torch.float32),
286 + weight_loader=weight_loader,
287 + )
288 + w2_input_scale = PerTensorScaleParameter(
289 + data=torch.full((num_experts, ), 1.0, dtype=torch.float32),
290 + weight_loader=weight_loader,
291 + )
292 + layer.register_parameter("w13_input_scale", w13_input_scale)
293 + layer.register_parameter("w2_input_scale", w2_input_scale)
294 +
295 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
296 + """Process FP8 MoE weights after loading from serialized checkpoint.
297 + Only supports pre-quantized checkpoints with FP8 weights and scales.
298 + """
299 +
300 + layer.w13_weight = Parameter(layer.w13_weight.data,
301 + requires_grad=False)
302 + layer.w2_weight = Parameter(layer.w2_weight.data, requires_grad=False)
303 +
304 + from vllm._custom_ops import scaled_fp8_quant
305 + from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
306 + per_tensor_dequantize)
307 +
308 + # Handle scale parameters
309 + if hasattr(layer,
310 + "w13_weight_scale") and layer.w13_weight_scale is not None:
311 + # Fp8 moe kernel needs single weight scale for w13 per expert.
312 + # We take the max of the w1 and w3 scales
313 + # then dequant and requant each expert.
314 + if layer.w13_weight_scale.dim() == 2:
315 +
316 + # Get the maximum scale across w1 and w3 for each expert
317 + max_w13_scales = layer.w13_weight_scale.max(dim=1).values
318 +
319 + # Requantize each expert's weights using the combined scale
320 + # w13_weight (num_experts, 2 * intermediate_size, hidden_size)
321 + # where the first intermediate_size rows are w1, the next are w3
322 + intermediate_size = layer.w13_weight.shape[1] // 2
323 + for expert_id in range(layer.w13_weight.shape[0]):
324 + start = 0
325 + for shard_id in range(2): # w1 and w3
326 + # Dequantize using the original scale for this shard
327 + dq_weight = per_tensor_dequantize(
328 + layer.w13_weight[expert_id][start:start +
329 + intermediate_size, :],
330 + layer.w13_weight_scale[expert_id][shard_id],
331 + )
332 + # Requantize using the combined max scale
333 +
334 + (
335 + layer.w13_weight[expert_id][start:start +
336 + intermediate_size, :],
337 + _,
338 + ) = scaled_fp8_quant(dq_weight,
339 + max_w13_scales[expert_id])
340 +
341 + start += intermediate_size
342 +
343 + # Update the scale parameter to be per-expert
344 + layer.w13_weight_scale = Parameter(max_w13_scales,
345 + requires_grad=False)
346 + else:
347 + layer.w13_weight_scale = Parameter(layer.w13_weight_scale.data,
348 + requires_grad=False)
349 +
350 + if hasattr(layer,
351 + "w2_weight_scale") and layer.w2_weight_scale is not None:
352 + layer.w2_weight_scale = Parameter(layer.w2_weight_scale.data,
353 + requires_grad=False)
354 + # Input scales must be equal for each expert in fp8 MoE layers.
355 + if hasattr(layer,
356 + "w13_input_scale") and layer.w13_input_scale is not None:
357 + layer.w13_input_scale = Parameter(layer.w13_input_scale.max(),
358 + requires_grad=False)
359 + if hasattr(layer,
360 + "w2_input_scale") and layer.w2_input_scale is not None:
361 + layer.w2_input_scale = Parameter(layer.w2_input_scale.max(),
362 + requires_grad=False)
363 +
364 + def apply(
365 + self,
366 + layer: torch.nn.Module,
367 + x: torch.Tensor,
368 + router_logits: torch.Tensor,
369 + top_k: int,
370 + renormalize: bool,
371 + use_grouped_topk: bool = False,
372 + topk_group: Optional[int] = None,
373 + num_expert_group: Optional[int] = None,
374 + global_num_experts: int = -1,
375 + expert_map: Optional[torch.Tensor] = None,
376 + custom_routing_function: Optional[Callable] = None,
377 + scoring_func: str = "softmax",
378 + e_score_correction_bias: Optional[torch.Tensor] = None,
379 + apply_router_weight_on_input: bool = False,
380 + activation: str = "silu",
381 + enable_eplb: bool = False,
382 + expert_load_view: Optional[torch.Tensor] = None,
383 + logical_to_physical_map: Optional[torch.Tensor] = None,
384 + logical_replica_count: Optional[torch.Tensor] = None,
385 + ) -> torch.Tensor:
386 + if enable_eplb:
387 + raise NotImplementedError(
388 + "EPLB not supported for `ModelOptFp8MoEMethod` yet.")
389 +
390 + # Expert selection
391 + topk_weights, topk_ids = FusedMoE.select_experts(
392 + hidden_states=x,
393 + router_logits=router_logits,
394 + use_grouped_topk=use_grouped_topk,
395 + top_k=top_k,
396 + renormalize=renormalize,
397 + topk_group=topk_group,
398 + num_expert_group=num_expert_group,
399 + custom_routing_function=custom_routing_function,
400 + scoring_func=scoring_func,
401 + e_score_correction_bias=e_score_correction_bias,
402 + )
403 + from vllm.model_executor.layers.fused_moe.fused_moe import (
404 + fused_experts)
405 + return fused_experts(
406 + x,
407 + layer.w13_weight,
408 + layer.w2_weight,
409 + topk_weights=topk_weights,
410 + topk_ids=topk_ids,
411 + inplace=True,
412 + activation=activation,
413 + use_fp8_w8a8=True,
414 + per_channel_quant=False,
415 + global_num_experts=global_num_experts,
416 + expert_map=expert_map,
417 + w1_scale=layer.w13_weight_scale,
418 + w2_scale=layer.w2_weight_scale,
419 + a1_scale=layer.w13_input_scale,
420 + a2_scale=layer.w2_input_scale,
421 + apply_router_weight_on_input=apply_router_weight_on_input,
422 + )
423 +
424 +
425 class ModelOptNvFp4Config(QuantizationConfig):
426 """Config class for ModelOpt FP4."""
427
@@ -274,7 +524,7 @@ class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):
524 class ModelOptNvFp4LinearMethod(LinearMethodBase):
525 """Linear method for Model Optimizer NVFP4.
526 Supports loading NVFP4 checkpoints with the following structure:
277 -
527 +
528 input_scale: torch.float32, scalar ,
529 weight: NVFP4(represented as byte) Shape: [1, X, y/2]
530 weight_scale: FP8-E4M3, Shape: [X, Y], aka per block scale,
@@ -455,7 +705,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
705 class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
706 """
707 MoE Method for FP4 Quantization.
458 - Args:
708 + Args:
709 quant_config: NVFP4 Quant Config
710 """
711
@@ -472,6 +722,12 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
722 " quantization. Please use Blackwell and"
723 " above.")
724
725 + def uses_weight_scale_2_pattern(self) -> bool:
726 + """
727 + FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales.
728 + """
729 + return True
730 +
731 def create_weights(self, layer: torch.nn.Module, num_experts: int,
732 hidden_size: int, intermediate_size_per_partition: int,
733 params_dtype: torch.dtype, **extra_weight_attrs):
vllm/model_executor/model_loader/weight_utils.py
+10
@@ -762,6 +762,10 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
762 modelopt_scale_names = [
763 ".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"
764 ]
765 + # Also support qkv_proj scale parameters (from stacked parameter processing)
766 + qkv_proj_scale_names = [
767 + ".self_attn.qkv_proj.k_scale", ".self_attn.qkv_proj.v_scale"
768 + ]
769 for scale_name in possible_scale_names:
770 if name.endswith(scale_name):
771 if any(mo_scale_name in name
@@ -769,6 +773,12 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
773 remapped_name = name.replace(
774 f".self_attn.{scale_name[1]}_proj{scale_name}",
775 f".self_attn.attn{scale_name}")
776 + elif any(qkv_scale_name in name
777 + for qkv_scale_name in qkv_proj_scale_names):
778 + # Handle qkv_proj scale parameters
779 + remapped_name = name.replace(
780 + f".self_attn.qkv_proj{scale_name}",
781 + f".self_attn.attn{scale_name}")
782 else:
783 remapped_name = name.replace(scale_name, f".attn{scale_name}")
784 if remapped_name not in params_dict:
vllm/model_executor/models/llama4.py
+55 -4
@@ -35,7 +35,8 @@ from vllm.model_executor.layers.linear import (QKVParallelLinear,
35 RowParallelLinear)
36 from vllm.model_executor.layers.quantization import QuantizationConfig
37 from vllm.model_executor.layers.rotary_embedding import get_rope
38 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38 +from vllm.model_executor.model_loader.weight_utils import (
39 + default_weight_loader, maybe_remap_kv_scale_name)
40
41 from .llama import LlamaForCausalLM, LlamaMLP, LlamaModel
42 from .utils import (AutoWeightsLoader, extract_layer_index, fast_topk,
@@ -432,12 +433,24 @@ class Llama4Model(LlamaModel):
433 for param_name, weight_name, shard_id in stacked_params_mapping:
434 if weight_name not in name or "experts" in name:
435 continue
435 - name = name.replace(weight_name, param_name)
436 + # This check is for ModelOpt ckpts with kv cache quant enabled
437 + if not (name.endswith(
438 + (".k_scale", ".v_scale")) and "self_attn" in name):
439 + name = name.replace(weight_name, param_name)
440 if is_pp_missing_parameter(name, self):
441 continue
442 + if name.endswith("scale") and "expert" not in name:
443 + # Remapping the name of FP8 kv-scale.
444 + name = maybe_remap_kv_scale_name(name, params_dict)
445 + if name is None:
446 + continue
447 param = params_dict[name]
439 - weight_loader = param.weight_loader
440 - weight_loader(param, loaded_weight, shard_id)
448 + weight_loader = getattr(param, "weight_loader",
449 + default_weight_loader)
450 + if weight_loader == default_weight_loader:
451 + weight_loader(param, loaded_weight)
452 + else:
453 + weight_loader(param, loaded_weight, shard_id)
454 loaded_params.add(name)
455 break
456 else:
@@ -452,6 +465,44 @@ class Llama4Model(LlamaModel):
465 if not moe_loaded:
466 if is_pp_missing_parameter(name, self):
467 continue
468 +
469 + # Handle flat expert scale parameters that
470 + # don't match per-expert patterns
471 + if ("experts." in name and ("w13_input_scale" in name
472 + or "w13_weight_scale" in name
473 + or "w2_input_scale" in name
474 + or "w2_weight_scale" in name)):
475 + # These are flat expert scales that apply to all experts
476 + param = params_dict[name]
477 + weight_loader = getattr(param, "weight_loader",
478 + default_weight_loader)
479 +
480 + # Check for MoE-specific loading support via
481 + # attribute instead of expensive runtime reflection
482 + supports_moe = getattr(weight_loader,
483 + 'supports_moe_loading', False)
484 +
485 + if supports_moe:
486 + # This is a MoE weight loader
487 + if "w13_" in name:
488 + shard_id = "w1"
489 + elif "w2_" in name:
490 + shard_id = "w2"
491 + else:
492 + shard_id = "w1"
493 +
494 + weight_loader(param,
495 + loaded_weight,
496 + name,
497 + shard_id=shard_id,
498 + expert_id=0)
499 + else:
500 + # Regular weight loader (handles both
501 + # param.weight_loader and default_weight_loader)
502 + weight_loader(param, loaded_weight)
503 + loaded_params.add(name)
504 + continue
505 +
506 param = params_dict[name]
507 weight_loader = getattr(param, "weight_loader",
508 default_weight_loader)
vllm/model_executor/models/mllama4.py
+144 -20
@@ -717,6 +717,7 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,
717 SupportsPP):
718 packed_modules_mapping = {
719 "qkv_proj": ["q_proj", "k_proj", "v_proj"],
720 + "gate_up_proj": ["gate_proj", "up_proj"],
721 }
722
723 @classmethod
@@ -902,32 +903,109 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,
903 qkv_weight = torch.cat(weight, dim=0)
904 yield key, qkv_weight
905
905 - def load_weights(self, weights: Iterable[tuple[str,
906 - torch.Tensor]]) -> set[str]:
906 + def _rename_weight_for_modelopt_checkpoint(self, name: str) -> str:
907 + """Rename weights from ModelOpt llama4 fp8 checkpoints to vLLM
908 + format."""
909 + if name.startswith("model."):
910 + # Handle expert scale parameters with flat naming
911 + if "feed_forward.experts." in name and ("_input_scale" in name or
912 + "_weight_scale" in name):
913 + renamed = name.replace("model.", "language_model.model.", 1)
914 + # Map checkpoint naming to vLLM's expected naming
915 + if "down_proj_input_scale" in renamed:
916 + return renamed.replace("down_proj_input_scale",
917 + "w2_input_scale")
918 + elif "down_proj_weight_scale" in renamed:
919 + return renamed.replace("down_proj_weight_scale",
920 + "w2_weight_scale")
921 + elif "gate_up_proj_input_scale" in renamed:
922 + return renamed.replace("gate_up_proj_input_scale",
923 + "w13_input_scale")
924 + elif "gate_up_proj_weight_scale" in renamed:
925 + return renamed.replace("gate_up_proj_weight_scale",
926 + "w13_weight_scale")
927 + return renamed
928 +
929 + # Handle attention scale parameters
930 + elif "self_attn." in name and (".k_scale" in name
931 + or ".v_scale" in name):
932 + renamed = name.replace("model.", "language_model.model.", 1)
933 + if ".k_proj.k_scale" in renamed:
934 + return renamed.replace(".k_proj.k_scale", ".attn.k_scale")
935 + elif ".v_proj.v_scale" in renamed:
936 + return renamed.replace(".v_proj.v_scale", ".attn.v_scale")
937 + return renamed
938 +
939 + # Standard model.* to language_model.model.* renaming
940 + return name.replace("model.", "language_model.model.", 1)
941 +
942 + elif name.startswith("lm_head.weight"):
943 + return name.replace("lm_head.weight",
944 + "language_model.lm_head.weight")
945 +
946 + return name
947 +
948 + def _separate_and_rename_weights(
949 + self, weights: Iterable[tuple[str, torch.Tensor]]
950 + ) -> tuple[list[tuple[str, torch.Tensor]], list[tuple[str, torch.Tensor]]]:
951 + """Rename weights and separate them into language_model and other
952 + weights."""
953 + language_model_weights = []
954 + other_weights = []
955
908 - stacked_params_mapping = [
909 - # (param_name, shard_name, shard_id)
910 - (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
911 - (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
912 - (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
913 - ]
914 - params_dict = dict(self.named_parameters())
915 - updated_params: set[str] = set()
956 + for name, weight in weights:
957 + renamed = self._rename_weight_for_modelopt_checkpoint(name)
958
917 - # language_model is an Llama4ForCausalLM instance. We load it's
918 - # using llama4's load_weights routine.
919 - language_model_weights, other_weights = self.separate_weights(
920 - weights, prefix="language_model.")
921 - loader = AutoWeightsLoader(self)
922 - loaded_language_model_params = loader.load_weights(
923 - language_model_weights)
924 - assert loaded_language_model_params is not None
925 - updated_params.update(loaded_language_model_params)
959 + if renamed.startswith("language_model."):
960 + language_model_weights.append((renamed, weight))
961 + else:
962 + other_weights.append((renamed, weight))
963 +
964 + return language_model_weights, other_weights
965 +
966 + def _handle_expert_scale_broadcasting(
967 + self, weights: list[tuple[str, torch.Tensor]], params_dict: dict
968 + ) -> tuple[list[tuple[str, torch.Tensor]], set[str]]:
969 + """Handle expert scale parameters that need broadcasting.
970 +
971 + ModelOpt checkpoints use a single value tensor scalar for BMM style
972 + experts, vLLM expects the scale to be broadcasted across all experts.
973 + """
974 + regular_weights = []
975 + expert_scale_weights = []
976 + updated_params = set()
977 +
978 + for name, weight in weights:
979 + # Check if this is an expert scale parameter that needs broadcasting
980 + if ("feed_forward.experts." in name and "scale" in name
981 + and ".shared_expert" not in name):
982 + if name in params_dict:
983 + param = params_dict[name]
984 + if (hasattr(param, 'data') and param.data.numel() > 1
985 + and weight.numel() == 1):
986 + # Broadcast single value to all experts
987 + param.data.fill_(weight.item())
988 + updated_params.add(name)
989 + continue
990 +
991 + expert_scale_weights.append((name, weight))
992 + else:
993 + regular_weights.append((name, weight))
994 +
995 + return regular_weights, expert_scale_weights, updated_params
996 +
997 + def _load_other_weights(self, other_weights: Iterable[tuple[str,
998 + torch.Tensor]],
999 + params_dict: dict,
1000 + stacked_params_mapping: list) -> set[str]:
1001 + """Load non-language-model weights with stacking support."""
1002 + updated_params = set()
1003
1004 if self.use_data_parallel:
1005 other_weights = self._consolidate_qkv_weights(other_weights)
1006
1007 for name, loaded_weight in other_weights:
1008 + # Try stacked parameter mapping first
1009 for param_name, weight_name, shard_id in stacked_params_mapping:
1010 if weight_name not in name or self.use_data_parallel:
1011 continue
@@ -938,10 +1016,56 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,
1016 weight_loader(param, loaded_weight, shard_id)
1017 break
1018 else:
1019 + # Use regular weight loading
1020 param = params_dict[name]
1021 weight_loader = getattr(param, "weight_loader",
1022 default_weight_loader)
944 -
1023 weight_loader(param, loaded_weight)
1024 updated_params.add(name)
1025 +
1026 + return updated_params
1027 +
1028 + def load_weights(self, weights: Iterable[tuple[str,
1029 + torch.Tensor]]) -> set[str]:
1030 +
1031 + stacked_params_mapping = [
1032 + # (param_name, shard_name, shard_id)
1033 + (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
1034 + (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
1035 + (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
1036 + # Shared expert gate_up_proj stacking
1037 + (".shared_expert.gate_up_proj", ".shared_expert.gate_proj", 0),
1038 + (".shared_expert.gate_up_proj", ".shared_expert.up_proj", 1),
1039 + # Feed forward gate_up_proj stacking (for non-MoE layers if any)
1040 + (".feed_forward.gate_up_proj", ".feed_forward.gate_proj", 0),
1041 + (".feed_forward.gate_up_proj", ".feed_forward.up_proj", 1),
1042 + ]
1043 + params_dict = dict(self.named_parameters())
1044 + updated_params: set[str] = set()
1045 +
1046 + # Separate and rename weights
1047 + language_model_weights, other_weights = (
1048 + self._separate_and_rename_weights(weights))
1049 +
1050 + # Handle expert scale parameters
1051 + regular_weights, expert_scale_weights, updated_params_from_experts = (
1052 + self._handle_expert_scale_broadcasting(language_model_weights,
1053 + params_dict))
1054 + updated_params.update(updated_params_from_experts)
1055 +
1056 + loader = AutoWeightsLoader(self)
1057 + loaded_language_model_params = loader.load_weights(regular_weights)
1058 + assert loaded_language_model_params is not None
1059 + updated_params.update(loaded_language_model_params)
1060 +
1061 + if expert_scale_weights:
1062 + loaded_expert_scale_params = loader.load_weights(
1063 + expert_scale_weights)
1064 + if loaded_expert_scale_params:
1065 + updated_params.update(loaded_expert_scale_params)
1066 +
1067 + updated_params.update(
1068 + self._load_other_weights(other_weights, params_dict,
1069 + stacked_params_mapping))
1070 +
1071 return updated_params