Fix vLLM fakequant MoE megatron export bug

examples/vllm_serve/vllm_reload_utils.py

@@ -572,7 +572,7 @@ def load_state_dict_from_path(
         saved_quant_dict = {
             key.replace("quantizer_", "quantizer._"): value
             for key, value in saved_quant_dict.items()
-            if "quantizer_" in key
+            if "quantizer" in key
         }
     saved_quant_dict = convert_dict_to_vllm(saved_quant_dict)
 

modelopt/torch/export/plugins/vllm_fakequant_megatron.py

@@ -135,8 +135,10 @@ def _get_quantized_state(
             # string then it usually ends with "." which needs to be removed.
             self.exclude_modules.append(prefix.removesuffix("."))
         block_size = 0
-
-        if hasattr(module, "weight") and module.weight is not None:
+        name_to_value = self._get_weight_bias(module, dtype, name_to_value)
+        if "weight" in name_to_value:
+            # Use the original device (avoid the CPU round-trip introduced by _get_weight_bias;
+            # fake-quantization runs on CUDA and the result is moved to CPU below).
             weight = module.weight.to(dtype)
             # Fold the weight_quantizer into the weight by applying fake-quantization
             # (quantize then dequantize). The weight_quantizer amax is not exported;
@@ -171,9 +173,6 @@ def _get_quantized_state(
         else:
             return name_to_value, qformat, block_size
 
-        if hasattr(module, "bias") and module.bias is not None:
-            name_to_value["bias"] = module.bias.to(dtype).cpu()
-
         # Only save input/output quantizer state; weight_quantizer amax is not exported
         # since it has been folded into the weight above.
         for name, param in get_quantizer_state_dict(module).items():

modelopt/torch/export/unified_export_megatron.py

@@ -743,6 +743,44 @@ def _custom_mapping_to_lambda(mapping):
 
         return all_rules
 
+    def _get_weight_bias(
+        self,
+        module: torch.nn.Module,
+        dtype: torch.dtype = torch.float16,
+        name_to_value: dict[str, torch.Tensor] | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Get the weight and bias of the module.
+
+        Args:
+            module: The target module to get the weight and bias.
+            dtype: The data type of the weight and bias.
+            name_to_value: The dictionary to store the weight and bias. A new dict is created
+                if not provided.
+
+        Returns:
+            The dictionary containing the weight and bias.
+        """
+        if name_to_value is None:
+            name_to_value = {}
+        # numel() > 0 intentionally excludes zero-element weight tensors (e.g. MoE routing
+        # layers whose weight is a placeholder) so callers can use "weight" in name_to_value
+        # as a reliable guard without re-inspecting module.weight.
+        if hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0:
+            weight = module.weight.to(dtype).cpu()
+            name_to_value["weight"] = weight
+
+        if hasattr(module, "bias") and module.bias is not None and module.bias.numel() > 0:
+            name_to_value["bias"] = module.bias.to(dtype).cpu()
+
+        if (
+            hasattr(module, "expert_bias")
+            and module.expert_bias is not None
+            and module.expert_bias.numel() > 0
+        ):
+            name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu()
+
+        return name_to_value
+
     def _get_quantized_state(
         self,
         module: torch.nn.Module,
@@ -767,21 +805,10 @@ def _get_quantized_state(
             self.exclude_modules.append(prefix.removesuffix("."))
         block_size = get_weight_block_size(module)
 
-        if hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0:
-            weight = module.weight.to(dtype).cpu()
-            name_to_value["weight"] = weight
-        else:
-            return name_to_value, qformat, block_size
+        name_to_value = self._get_weight_bias(module, dtype, name_to_value)
 
-        if hasattr(module, "bias") and module.bias is not None and module.bias.numel() > 0:
-            name_to_value["bias"] = module.bias.to(dtype).cpu()
-
-        if (
-            hasattr(module, "expert_bias")
-            and module.expert_bias is not None
-            and module.expert_bias.numel() > 0
-        ):
-            name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu()
+        if "weight" not in name_to_value:
+            return name_to_value, qformat, block_size
 
         if qformat == QUANTIZATION_NONE:
             return name_to_value, qformat, block_size

modelopt/torch/quantization/plugins/vllm.py

@@ -385,14 +385,19 @@ def _invoke_fused_moe_quantized_function(
             # First layer of expert
             A = self.w13_input_quantizer(A)  # noqa: N806
             if self.w13_weight_quantizer.is_enabled:  # pragma: no cover
-                original_weight, self.w13_weight = (
-                    self.w13_weight,
-                    self.w13_weight_quantizer(self.w13_weight),
-                )
-                # In case the weight quantizer isn't folded yet in vllm_serve_fakequant, pass the
-                # quantized weight to the kernel.
-                B = self.w13_weight  # noqa: N806
+                # Same pattern as FakeQuantMethod.apply: wrap as nn.Parameter if needed, swap
+                # w13_weight, call kernel, restore (tensor cannot stay assigned to nn.Parameter slot).
+                original_weight = self.w13_weight
+                quantized_tensor = self.w13_weight_quantizer(original_weight)
                 try:
+                    if isinstance(original_weight, torch.nn.Parameter) and not isinstance(
+                        quantized_tensor, torch.nn.Parameter
+                    ):
+                        quantized_tensor = torch.nn.Parameter(
+                            quantized_tensor, requires_grad=original_weight.requires_grad
+                        )
+                    self.w13_weight = quantized_tensor
+                    B = quantized_tensor  # noqa: N806
                     original_kernel(A, B, C, *args, **kwargs)
                 finally:
                     self.w13_weight = original_weight
@@ -403,14 +408,17 @@ def _invoke_fused_moe_quantized_function(
         elif B is self.w2_weight:
             A = self.w2_input_quantizer(A)  # noqa: N806
             if self.w2_weight_quantizer.is_enabled:  # pragma: no cover
-                original_weight, self.w2_weight = (
-                    self.w2_weight,
-                    self.w2_weight_quantizer(self.w2_weight),
-                )
-                # In case the weight quantizer isn't folded yet in vllm_serve_fakequant, pass the
-                # quantized weight to the kernel.
-                B = self.w2_weight  # noqa: N806
+                original_weight = self.w2_weight
+                quantized_tensor = self.w2_weight_quantizer(original_weight)
                 try:
+                    if isinstance(original_weight, torch.nn.Parameter) and not isinstance(
+                        quantized_tensor, torch.nn.Parameter
+                    ):
+                        quantized_tensor = torch.nn.Parameter(
+                            quantized_tensor, requires_grad=original_weight.requires_grad
+                        )
+                    self.w2_weight = quantized_tensor
+                    B = quantized_tensor  # noqa: N806
                     original_kernel(A, B, C, *args, **kwargs)
                 finally:
                     self.w2_weight = original_weight