diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index aa8d3a5388b..6b658551f15 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/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) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index c8e45be3650..85952f8dc58 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/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(): diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 89b718623da..62053e549c8 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/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 diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 6721a8f798f..95ca3240b73 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/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