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 14, 05:30 UTC
[NVIDIA] Auto detect modelopt quant and fix DSR1-FP4 weight loading (#22073)
Po-Han Huang (NVIDIA) committed
Aug 5, 2025 at 09:02 UTC
bdcb42e45db5cbbc02b0f69ac304c87d7a8cb6b6
3 files changed
+67
-15
vllm/config.py
+15
@@ -1108,6 +1108,21 @@ class ModelConfig:
1108
if quant_cfg is None:
1109
# compressed-tensors uses a "compression_config" key
1110
quant_cfg = getattr(self.hf_config, "compression_config", None)
1111
+
1112
+ else:
1113
+ # Set quant_method for ModelOpt models.
1114
+ producer_name = quant_cfg.get("producer", {}).get("name")
1115
+ if producer_name == "modelopt":
1116
+ quant_algo = quant_cfg.get("quantization",
1117
+ {}).get("quant_algo")
1118
+ if quant_algo == "FP8":
1119
+ quant_cfg["quant_method"] = "modelopt"
1120
+ elif quant_algo == "NVFP4":
1121
+ quant_cfg["quant_method"] = "modelopt_fp4"
1122
+ elif quant_algo is not None:
1123
+ raise ValueError(
1124
+ f"Unknown ModelOpt quant algo: {quant_algo}")
1125
+
1126
return quant_cfg
1127
1128
def _verify_quantization(self) -> None:
vllm/model_executor/layers/fused_moe/layer.py
+38
-15
@@ -919,9 +919,13 @@ class FusedMoE(torch.nn.Module):
919
elif shard_id == "w2":
920
param_data[expert_id] = loaded_weight
921
922
- def _load_w13_weight_scale(self, shard_dim: int,
923
- loaded_weight: torch.Tensor,
924
- param: torch.Tensor, tp_rank: int):
922
+ def _load_combined_w13_weight_scale(self, shard_dim: int,
923
+ loaded_weight: torch.Tensor,
924
+ param: torch.Tensor, tp_rank: int):
925
+ """
926
+ Load w13 weight scales assuming that w1 weight scales and w3 weight
927
+ scales are stored in the same loaded_weight tensor.
928
+ """
929
shard_size = param.shape[shard_dim]
930
loaded_weight = loaded_weight.narrow(shard_dim, shard_size * tp_rank,
931
shard_size)
@@ -1168,24 +1172,43 @@ class FusedMoE(torch.nn.Module):
1172
uses_weight_scale_2 = self.quant_method.uses_weight_scale_2_pattern(
1173
)
1174
1171
- # For per-tensor, FP4 uses "weight_scale_2", FP8 uses "weight_scale"
1172
- per_tensor_conditions = (
1173
- "weight_scale_2" in weight_name if uses_weight_scale_2 else
1174
- "weight_scale" in weight_name) or "input_scale" in weight_name
1175
-
1176
- if "w13_weight_scale" in weight_name:
1177
- self._load_w13_weight_scale(shard_dim=shard_dim,
1178
- loaded_weight=loaded_weight,
1179
- param=param,
1180
- tp_rank=self.tp_rank)
1181
- elif per_tensor_conditions:
1175
+ # Call _load_per_tensor_weight_scale() to load per-tensor (scalar)
1176
+ # weights scales.
1177
+ # Input scales are always per-tensor.
1178
+ # Weight scales: FP4 uses "weight_scale_2" and FP8 uses
1179
+ # "weight_scale" for per-tensor scales.
1180
+ is_per_tensor = ("weight_scale_2" in weight_name
1181
+ if uses_weight_scale_2 else "weight_scale"
1182
+ in weight_name) or "input_scale" in weight_name
1183
+ if is_per_tensor:
1184
self._load_per_tensor_weight_scale(
1185
shard_id=shard_id,
1186
param=param,
1187
loaded_weight=loaded_weight,
1188
expert_id=expert_id,
1189
)
1188
- elif "weight" in weight_name:
1190
+ return True if return_success else None
1191
+
1192
+ # If the weight is w13_weight_scale and w13_weight_scales are
1193
+ # combined into single loaded_weight, call
1194
+ # _load_combined_w13_weight_scale() to load it.
1195
+ # This is checked by comparing the hidden_out dims of the
1196
+ # loaded_weight and the param.
1197
+ if "w13_weight_scale" in weight_name:
1198
+ loaded_weight_hidden_out = loaded_weight.shape[-2]
1199
+ param_hidden_out = param.data.shape[-2] * self.tp_size
1200
+ if loaded_weight_hidden_out == param_hidden_out:
1201
+ self._load_combined_w13_weight_scale(
1202
+ shard_dim=shard_dim,
1203
+ loaded_weight=loaded_weight,
1204
+ param=param,
1205
+ tp_rank=self.tp_rank,
1206
+ )
1207
+ return True if return_success else None
1208
+
1209
+ # For other weights, call _load_model_weight_or_group_weight_scale()
1210
+ # to load it.
1211
+ if "weight" in weight_name:
1212
self._load_model_weight_or_group_weight_scale(
1213
shard_id=shard_id,
1214
shard_dim=shard_dim,
vllm/transformers_utils/config.py
+14
@@ -449,6 +449,20 @@ def get_config(
449
model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type]
450
config.update({"architectures": [model_type]})
451
452
+ # ModelOpt 0.31.0 and after saves the quantization config in the model
453
+ # config file.
454
+ quantization_config = config_dict.get("quantization_config", None)
455
+
456
+ # ModelOpt 0.29.0 and before saves the quantization config in a separate
457
+ # "hf_quant_config.json" in the same directory as the model config file.
458
+ if quantization_config is None \
459
+ and file_or_path_exists(model, "hf_quant_config.json", revision):
460
+ quantization_config = get_hf_file_to_dict("hf_quant_config.json",
461
+ model, revision)
462
+
463
+ if quantization_config is not None:
464
+ config.quantization_config = quantization_config
465
+
466
if hf_overrides_kw:
467
logger.debug("Overriding HF config with %s", hf_overrides_kw)
468
config.update(hf_overrides_kw)