From 065cfcaafba54655ce7b2d1e09b23c1f2e7777ac Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Sat, 18 Apr 2026 01:35:13 +0000 Subject: [PATCH 1/6] fixed bug Signed-off-by: Kinjal Patel --- .../export/plugins/vllm_fakequant_megatron.py | 9 ++-- .../torch/export/unified_export_megatron.py | 49 ++++++++++++++----- 2 files changed, 39 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index c8e45be3650..a7c5702bade 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -135,9 +135,9 @@ 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: - weight = module.weight.to(dtype) + name_to_value = self._get_weight_bias(module, dtype, name_to_value) + if "weight" in name_to_value: + weight = name_to_value["weight"] # Fold the weight_quantizer into the weight by applying fake-quantization # (quantize then dequantize). The weight_quantizer amax is not exported; # the vLLM fakequant reload path disables the weight quantizer when absent. @@ -171,9 +171,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..af31d47045c 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -743,6 +743,38 @@ 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] = {}, + ) -> 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. + + Returns: + The dictionary containing the weight and bias. + """ + 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 +799,12 @@ 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 - - if hasattr(module, "bias") and module.bias is not None and module.bias.numel() > 0: - name_to_value["bias"] = module.bias.to(dtype).cpu() + name_to_value = self._get_weight_bias(module, dtype, name_to_value) - if ( - hasattr(module, "expert_bias") - and module.expert_bias is not None - and module.expert_bias.numel() > 0 + if not ( + hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0 ): - name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu() + return name_to_value, qformat, block_size if qformat == QUANTIZATION_NONE: return name_to_value, qformat, block_size From dd948e57d7cd691cbd6b2a213b40c48b709eaf9c Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 21 Apr 2026 01:01:58 +0000 Subject: [PATCH 2/6] bug fix Signed-off-by: Kinjal Patel --- modelopt/torch/quantization/plugins/vllm.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 6721a8f798f..f5cb8d8a1a4 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -385,13 +385,8 @@ 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 + original_weight = self.w13_weight + B = self.w13_weight_quantizer(original_weight) # noqa: N806 try: original_kernel(A, B, C, *args, **kwargs) finally: @@ -403,13 +398,8 @@ 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 + B = self.w2_weight_quantizer(original_weight) # noqa: N806 try: original_kernel(A, B, C, *args, **kwargs) finally: From 482f49215b8e8fe098a49a95b6e12d4922557de5 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 21 Apr 2026 02:24:35 +0000 Subject: [PATCH 3/6] minor Signed-off-by: Kinjal Patel --- modelopt/torch/quantization/plugins/vllm.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index f5cb8d8a1a4..2c731f66497 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -387,10 +387,7 @@ def _invoke_fused_moe_quantized_function( if self.w13_weight_quantizer.is_enabled: # pragma: no cover original_weight = self.w13_weight B = self.w13_weight_quantizer(original_weight) # noqa: N806 - try: - original_kernel(A, B, C, *args, **kwargs) - finally: - self.w13_weight = original_weight + original_kernel(A, B, C, *args, **kwargs) else: original_kernel(A, B, C, *args, **kwargs) if self.w13_output_quantizer.is_enabled: @@ -400,10 +397,7 @@ def _invoke_fused_moe_quantized_function( if self.w2_weight_quantizer.is_enabled: # pragma: no cover original_weight = self.w2_weight B = self.w2_weight_quantizer(original_weight) # noqa: N806 - try: - original_kernel(A, B, C, *args, **kwargs) - finally: - self.w2_weight = original_weight + original_kernel(A, B, C, *args, **kwargs) else: original_kernel(A, B, C, *args, **kwargs) if self.w2_output_quantizer.is_enabled: From 2c50141a6f36b7d028487818680bcb2b0b61c7a1 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 21 Apr 2026 04:36:08 +0000 Subject: [PATCH 4/6] minor Signed-off-by: Kinjal Patel --- .../export/plugins/vllm_fakequant_megatron.py | 4 +++- modelopt/torch/export/unified_export_megatron.py | 14 +++++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index a7c5702bade..85952f8dc58 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -137,7 +137,9 @@ def _get_quantized_state( block_size = 0 name_to_value = self._get_weight_bias(module, dtype, name_to_value) if "weight" in name_to_value: - weight = name_to_value["weight"] + # 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; # the vLLM fakequant reload path disables the weight quantizer when absent. diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index af31d47045c..62053e549c8 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -747,18 +747,24 @@ def _get_weight_bias( self, module: torch.nn.Module, dtype: torch.dtype = torch.float16, - name_to_value: dict[str, torch.Tensor] = {}, + 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. + 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 @@ -801,9 +807,7 @@ def _get_quantized_state( name_to_value = self._get_weight_bias(module, dtype, name_to_value) - if not ( - hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0 - ): + if "weight" not in name_to_value: return name_to_value, qformat, block_size if qformat == QUANTIZATION_NONE: From d01b5f30e6611523dc038d0b2ca98cc175b90549 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 21 Apr 2026 20:24:42 +0000 Subject: [PATCH 5/6] minor Signed-off-by: Kinjal Patel --- modelopt/torch/quantization/plugins/vllm.py | 32 ++++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/quantization/plugins/vllm.py b/modelopt/torch/quantization/plugins/vllm.py index 2c731f66497..95ca3240b73 100644 --- a/modelopt/torch/quantization/plugins/vllm.py +++ b/modelopt/torch/quantization/plugins/vllm.py @@ -385,9 +385,22 @@ 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 + # 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 - B = self.w13_weight_quantizer(original_weight) # noqa: N806 - original_kernel(A, B, C, *args, **kwargs) + 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 else: original_kernel(A, B, C, *args, **kwargs) if self.w13_output_quantizer.is_enabled: @@ -396,8 +409,19 @@ def _invoke_fused_moe_quantized_function( A = self.w2_input_quantizer(A) # noqa: N806 if self.w2_weight_quantizer.is_enabled: # pragma: no cover original_weight = self.w2_weight - B = self.w2_weight_quantizer(original_weight) # noqa: N806 - original_kernel(A, B, C, *args, **kwargs) + 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 else: original_kernel(A, B, C, *args, **kwargs) if self.w2_output_quantizer.is_enabled: From f987f541673ba43a46dca61d2dce4dd2d8c3404d Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Wed, 22 Apr 2026 01:22:54 +0000 Subject: [PATCH 6/6] fixed bug Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_reload_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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)