Moved vllm fq export code to separate files (#612)

## What does this PR do?

**Type of change:** : Bug fix

**Overview:** 
moved vLLM fakequant checkpoint export code to separate files:
1. for HF export -> modelopt.torch.export.plugins.vllm_fq_hf
2. for megatron export -> modelopt.torch.export.plugins.vllm_fq_megatron

## Usage
Refer to

[README.md](https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/096ee13ea62bbb0ce0a4e4128c439651374d6235/examples/vllm_serve#load-qatptq-model-and-serve-in-vllm-wip)

## Testing
- Tested HF approach by exporting bf16 model using QAT script and
running vllm server, verified amax values match
- Tested MCore approach by quantizing and exporting bf16 model using
quantize.sh and export.sh script and running vllm server, verified amax
values match
- Tested using unit tests in
`tests/gpu/torch/export/test_vllm_fq_hf_export.py` and
`tests/gpu/torch/export/test_vllm_fq_megatron_export.py`

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: NA
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
NA

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added dedicated export functions for vLLM fakequant checkpoint format,
supporting both HuggingFace and Megatron Core models.

* **Refactor**
* Simplified export API by removing conditional export flags for
cleaner, more predictable behavior.
* Reorganized export functionality into focused plugin modules for
improved maintainability.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>

examples/vllm_serve/README.md

@@ -57,10 +57,10 @@ lm_eval --model local-completions --tasks gsm8k --model_args model=<model_name>,
 
 Overwrite the calibrated amax value with prepared values from either QAT/PTQ.
 
-Step 1: export the model with bf16 weights and amax values.
+Step 1: export the model with bf16 weights and amax values. To export the model:
 
-- For HF model set `export_bf16_weights_amax` to export the model with function `modelopt.torch.export.unified_export_hf.export_hf_checkpoint`.
-- For MCore model use `export_bf16_weights_amax` to export the model with function `modelopt.torch.export.unified_export_megatron.export_mcore_gpt_to_hf`.
+- For HF model use `modelopt.torch.export.export_hf_vllm_fq_checkpoint` function.
+- For MCore model use `modelopt.torch.export.export_mcore_gpt_to_hf_vllm_fq` function.
 
 Step 2: configure <quant_amax.pth> from exported model using AMAX_FILE_PATH environment variable in step 1. For example:
 

modelopt/torch/export/__init__.py

@@ -19,6 +19,7 @@ from .convert_hf_config import *
 from .model_config import *
 from .model_config_export import *
 from .model_utils import *
+from .plugins import *
 from .transformer_engine import *
 from .unified_export_hf import *
 from .unified_export_megatron import *

modelopt/torch/export/plugins/__init__.py

@@ -21,3 +21,7 @@ with import_plugin("megatron_importer"):
     from .megatron_importer import *
 
 from .hf_spec_export import *
+from .vllm_fakequant_hf import *
+
+with import_plugin("vllm_fakequant_megatron"):
+    from .vllm_fakequant_megatron import *

modelopt/torch/export/plugins/vllm_fakequant.py

@@ -1,125 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Export functions for vLLM fakequant."""
-
-import os
-from pathlib import Path
-
-import torch
-import torch.nn as nn
-
-from modelopt.torch.export.layer_utils import is_quantlinear
-from modelopt.torch.export.model_config import QUANTIZATION_NONE
-from modelopt.torch.quantization.utils import get_quantizer_state_dict
-
-
-def export_hf_vllm_fq_checkpoint(
-    model: nn.Module,
-    export_dir: Path | str,
-) -> dict[str, torch.Tensor]:
-    """Exports the torch model weights and amax values separately.
-
-    This function:
-    1. Extracts amax values for calibration
-    2. Deletes all quantizer parameters from state dict to store only weights in original dtype
-
-    Args:
-        model: The quantized model to export
-        export_dir: Directory to save the amax values
-
-    Returns:
-        post_state_dict: Dict containing quantized weights
-    """
-    amax_dict = {
-        name + "._amax": param["_amax"].detach().clone().cpu()
-        for name, param in get_quantizer_state_dict(model).items()
-        if "_amax" in param
-    }
-
-    # remove quantizer from model
-    for _, module in model.named_modules():
-        if is_quantlinear(module):
-            delattr(module, "weight_quantizer")
-            delattr(module, "input_quantizer")
-            delattr(module, "output_quantizer")
-            module.export()
-    torch.save(amax_dict, f"{export_dir}/quant_amax.pth")
-    return model.state_dict()
-
-
-def get_mcore_vllm_fq_quantized_state(
-    module: torch.nn.Module, name_to_value: dict, dtype: torch.dtype = torch.bfloat16
-):
-    """Return a state_dict, quantization format, and block_size of the quantized module.
-
-    Args:
-        module: The target module to perform real quantization.
-        name_to_value: The dictionary to store the quantized state.
-        dtype: The default data type.
-
-    Returns:
-        Tuple: state dict, quantization format, and block_size of the quantized module.
-
-    """
-    qformat: str = QUANTIZATION_NONE
-    block_size = 0
-
-    for name, param in get_quantizer_state_dict(module).items():
-        if "_amax" in param:
-            name_to_value[name + "._amax"] = param["_amax"].to(dtype).cpu()
-    return name_to_value, qformat, block_size
-
-
-def gather_mcore_vllm_fq_quantized_state_dict(
-    state_dict: dict[str, torch.Tensor], save_directory: str | os.PathLike
-):
-    """Gather all quantized state dict from all ranks and save them to a file.
-
-    Args:
-        state_dict: The state dictionary of the module.
-        save_directory: The directory to save the quantized state dict.
-
-    Returns:
-        The state dictionary of the module without quantized state.
-    """
-    amax_state_dict = {
-        k: v.detach().clone().cpu() for k, v in state_dict.items() if k.endswith("_amax")
-    }
-
-    # Gather all amax dicts to rank 0
-    world_size = torch.distributed.get_world_size()
-    rank = torch.distributed.get_rank()
-
-    if rank == 0:
-        # Rank 0 will collect all amax values
-        all_amax_dicts = [None] * world_size
-        torch.distributed.gather_object(amax_state_dict, all_amax_dicts, dst=0)
-
-        # Merge all amax dicts into one
-        merged_amax_dict = {}
-        for amax_dict in all_amax_dicts:
-            if amax_dict is not None:
-                merged_amax_dict.update(amax_dict)
-
-        print(f"Total amax entries from all ranks: {len(merged_amax_dict.keys())}")
-        torch.save(merged_amax_dict, save_directory + "/quant_amax.pth")
-    else:
-        # Other ranks just send their amax values
-        torch.distributed.gather_object(amax_state_dict, None, dst=0)
-
-    torch.distributed.barrier()
-
-    # remove amax values from state_dict
-    return {k: v for k, v in state_dict.items() if not k.endswith("_amax")}

modelopt/torch/export/plugins/vllm_fakequant_hf.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Export HuggingFace model to vLLM fakequant checkpoint."""
+
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+from modelopt.torch.export.layer_utils import is_quantlinear
+from modelopt.torch.quantization.utils import get_quantizer_state_dict
+
+__all__ = ["export_hf_vllm_fq_checkpoint"]
+
+
+def export_hf_vllm_fq_checkpoint(
+    model: nn.Module,
+    export_dir: Path | str,
+):
+    """Exports the torch model weights and amax values separately.
+
+    This function:
+    1. Extracts amax values for calibration
+    2. Deletes all quantizer parameters from state dict to store only weights in original dtype
+    3. Saves the model weights
+
+    Args:
+        model: The quantized model to export
+        export_dir: Directory to save the amax values
+
+    """
+    export_dir = Path(export_dir)
+    export_dir.mkdir(parents=True, exist_ok=True)
+
+    amax_dict = {
+        name + "._amax": param["_amax"].detach().clone().cpu()
+        for name, param in get_quantizer_state_dict(model).items()
+        if "_amax" in param
+    }
+
+    # remove quantizer from model
+    for _, module in model.named_modules():
+        if is_quantlinear(module):
+            for attr in ["weight_quantizer", "input_quantizer", "output_quantizer"]:
+                if hasattr(module, attr):
+                    delattr(module, attr)
+            module.export()
+    torch.save(amax_dict, f"{export_dir}/quant_amax.pth")
+    # Save model
+    model.save_pretrained(export_dir, state_dict=model.state_dict(), save_modelopt_state=False)

modelopt/torch/export/plugins/vllm_fakequant_megatron.py

@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Export Megatron Core Model to HuggingFace vLLM fakequant checkpoint."""
+
+import os
+import tempfile
+from pathlib import Path
+
+import torch
+
+from modelopt.torch.export.model_config import QUANTIZATION_NONE
+from modelopt.torch.export.unified_export_megatron import GPTModelExporter
+
+__all__ = ["export_mcore_gpt_to_hf_vllm_fq"]
+
+
+def gather_mcore_vllm_fq_quantized_state_dict(
+    model, state_dict: dict[str, torch.Tensor], save_directory: str | os.PathLike
+):
+    """Gather all quantized state dict from all ranks and save them to a file.
+
+    Args:
+        state_dict: The state dictionary of the module.
+        save_directory: The directory to save the quantized state dict.
+
+    Returns:
+        The state dictionary of the module without quantized state.
+    """
+    amax_state_dict = {
+        k: v.detach().clone().cpu() for k, v in state_dict.items() if k.endswith("_amax")
+    }
+
+    # Gather all amax dicts to rank 0
+    world_size = torch.distributed.get_world_size()
+    rank = torch.distributed.get_rank()
+
+    if rank == 0:
+        # Rank 0 will collect all amax values
+        all_amax_dicts = [None] * world_size
+        torch.distributed.gather_object(amax_state_dict, all_amax_dicts, dst=0)
+
+        # Merge all amax dicts into one
+        merged_amax_dict = {}
+        for amax_dict in all_amax_dicts:
+            if amax_dict is not None:
+                merged_amax_dict.update(amax_dict)
+
+        print(f"Total amax entries from all ranks: {len(merged_amax_dict.keys())}")
+        torch.save(merged_amax_dict, save_directory + "/quant_amax.pth")
+    else:
+        # Other ranks just send their amax values
+        torch.distributed.gather_object(amax_state_dict, None, dst=0)
+
+    torch.distributed.barrier()
+
+
+class VllmFqGPTModelExporter(GPTModelExporter):
+    """VLLM fakequant GPTModel exporter."""
+
+    def save_pretrained(
+        self,
+        save_directory: str | os.PathLike,
+        pretrained_model_name_or_path: str | os.PathLike | None = None,
+    ):
+        os.makedirs(save_directory, exist_ok=True)
+        gather_mcore_vllm_fq_quantized_state_dict(self.model, self.state_dict, save_directory)
+        assert not (self.is_multimodal and pretrained_model_name_or_path is not None), (
+            "Exporting weights in bf16 and amax values is not supported for multimodal models "
+            "when pretrained_model_name_or_path is not None"
+        )
+        assert not self.export_extra_modules, (
+            "Exporting extra modules is not supported for vLLM fakequant"
+        )
+        super().save_pretrained(save_directory, pretrained_model_name_or_path)
+
+    def _get_quantization_format(self, module: torch.nn.Module):
+        return QUANTIZATION_NONE
+
+
+def export_mcore_gpt_to_hf_vllm_fq(
+    model: torch.nn.Module,
+    pretrained_model_name_or_path: str | os.PathLike | None = None,
+    export_extra_modules: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
+    export_dir: Path | str = tempfile.gettempdir(),
+    moe_router_dtype: torch.dtype | None = None,
+):
+    """Export Megatron Core GPTModel to unified checkpoint and save to export_dir.
+
+    Args:
+        model: The Megatron Core GPTModel instance.
+        pretrained_model_name_or_path: Can be either: the *model id* of a
+            pretrained model hosted inside a model repo on huggingface.co; or
+            a *directory* containing model weights saved using
+            [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+        export_extra_modules: If True, export extra modules like medusa_heads or
+            eagle_module. Otherwise, only export the base model.
+        dtype: The weights data type to export the unquantized layers.
+        export_dir: The target export path.
+    """
+    exporter = VllmFqGPTModelExporter(
+        model,
+        pretrained_model_name_or_path,
+        export_extra_modules=export_extra_modules,
+        dtype=dtype,
+        moe_router_dtype=moe_router_dtype,
+    )
+    exporter.save_pretrained(export_dir, pretrained_model_name_or_path)

modelopt/torch/export/unified_export_hf.py

@@ -59,7 +59,6 @@ from .model_config import (
 )
 from .model_utils import get_language_model_from_vl, is_multimodal_model
 from .plugins import export_spec_ckpt_config, export_spec_ckpt_state_dict, spec_opt_only
-from .plugins.vllm_fakequant import export_hf_vllm_fq_checkpoint
 from .quant_utils import (
     fuse_prequant_layernorm,
     fuse_prequant_to_linear,
@@ -559,7 +558,6 @@ def export_hf_checkpoint(
     dtype: torch.dtype | None = None,
     export_dir: Path | str = tempfile.gettempdir(),
     save_modelopt_state: bool = False,
-    export_vllm_fq_weights_qstate: bool = False,
 ):
     """Exports the torch model to unified checkpoint and saves to export_dir.
 
@@ -568,8 +566,6 @@ def export_hf_checkpoint(
         dtype: the weights data type to export the unquantized layers or the default model data type if None.
         export_dir: the target export path.
         save_modelopt_state: whether to save the modelopt state_dict.
-        export_vllm_fq_weights_qstate: whether to export the weights and quantization state separately for vLLM
-        fakequant serving.
     """
     export_dir = Path(export_dir)
     export_dir.mkdir(parents=True, exist_ok=True)
@@ -583,11 +579,7 @@ def export_hf_checkpoint(
         return
 
     try:
-        if export_vllm_fq_weights_qstate:
-            post_state_dict = export_hf_vllm_fq_checkpoint(model, export_dir)
-            hf_quant_config = None
-        else:
-            post_state_dict, hf_quant_config = _export_hf_checkpoint(model, dtype)
+        post_state_dict, hf_quant_config = _export_hf_checkpoint(model, dtype)
 
         if hf_quant_config is not None:
             # Save hf_quant_config.json for\ backward compatibility

modelopt/torch/export/unified_export_megatron.py

@@ -48,10 +48,6 @@ from .model_config import (
 from .plugins.mcore_common import all_mcore_hf_export_mapping
 from .plugins.mcore_custom import CustomModuleMapping, save_safetensors
 from .plugins.megatron_importer import GPTModelImporter
-from .plugins.vllm_fakequant import (
-    gather_mcore_vllm_fq_quantized_state_dict,
-    get_mcore_vllm_fq_quantized_state,
-)
 from .quant_utils import (
     get_activation_scaling_factor,
     get_kv_cache_dtype,
@@ -116,65 +112,6 @@ def get_kv_cache_scaling_factor(kv_module: nn.Module) -> torch.Tensor:
     return scaling_factor
 
 
-def get_quantized_state(
-    module: torch.nn.Module,
-    dtype: torch.dtype = torch.float16,
-    export_vllm_fq_weights_qstate: bool = False,
-) -> tuple[dict[str, torch.Tensor], str, int]:
-    """Return a state_dict, quantization format, and block_size of the module.
-
-    Args:
-        module: The target module to perform real quantization.
-        dtype: The default data type.
-        export_vllm_fq_weights_qstate: Whether to export the weights in bf16 and amax values.
-
-    Returns:
-        Tuple: state_dict, quantization format, and block_size of the module.
-    """
-    name_to_value = {}
-    qformat: str = get_quantization_format(module)
-    block_size = get_weight_block_size(module)
-
-    if hasattr(module, "weight") and module.weight is not None:
-        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:
-        name_to_value["bias"] = module.bias.to(dtype).cpu()
-
-    if hasattr(module, "expert_bias") and module.expert_bias is not None:
-        name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu()
-
-    if export_vllm_fq_weights_qstate:
-        return get_mcore_vllm_fq_quantized_state(module, name_to_value, dtype)
-
-    # Getting the weight scales
-    weight_scale = get_weight_scaling_factor(module)
-    weight_scale_2 = get_weight_scaling_factor_2(module)
-    if weight_scale is not None:
-        name_to_value["weight_scale"] = weight_scale
-
-    if weight_scale_2 is not None:
-        name_to_value["weight_scale_2"] = weight_scale_2
-
-    # Getting the input scale
-    input_scale = get_activation_scaling_factor(module)
-    if input_scale is not None:
-        name_to_value["input_scale"] = input_scale
-        # TODO (chenhany): support AWQ with pre_quant_scale
-        if hasattr(module.input_quantizer, "_pre_quant_scale"):
-            raise ValueError("Detect pre_quant_scale! SmoothQuant/AWQ are not yet supported!")
-
-    if hasattr(module, "output_quantizer"):
-        output_scale = get_kv_cache_scaling_factor(module)
-        if output_scale is not None:
-            name_to_value["output_scale"] = output_scale
-
-    return name_to_value, qformat, block_size
-
-
 class GPTModelExporter:
     """Megatron Core GPTModel Exporter.
 
@@ -201,7 +138,6 @@ class GPTModelExporter:
         dtype=torch.bfloat16,
         trust_remote_code: bool = True,
         moe_router_dtype: torch.dtype | None = None,
-        export_vllm_fq_weights_qstate: bool = False,
     ):
         """Create a GPTModel exporter instance."""
         if not isinstance(model, (GPTModel, MambaModel, LLaVAModel)):
@@ -237,7 +173,6 @@ class GPTModelExporter:
         self.model = model.language_model if self.is_multimodal else model
         self.dtype = dtype
         self.trust_remote_code = trust_remote_code
-        self.export_vllm_fq_weights_qstate = export_vllm_fq_weights_qstate
         self.arch = self._hf_config.architectures[0]
         # TODO: May modify this later according to what quantization exported ckpt is, currently only support BF16.
         if self.arch == "GptOssForCausalLM":
@@ -347,11 +282,8 @@ class GPTModelExporter:
 
         # Main export process
         state_dict = self.extra_state_dict if self.export_extra_modules else self.state_dict
-        quantization_format = (
-            get_quantization_format(self.model)
-            if not self.export_vllm_fq_weights_qstate
-            else QUANTIZATION_NONE
-        )
+        quantization_format = self._get_quantization_format(self.model)
+
         quantization = None
         kv_cache_quantization = None
 
@@ -398,7 +330,7 @@ class GPTModelExporter:
                 except (OSError, ValueError, ImportError):
                     pass
 
-        if is_last_stage_main_rank and not self.export_vllm_fq_weights_qstate:
+        if is_last_stage_main_rank and quantization is not None:
             hf_quant_config = {
                 "producer": {
                     "name": "modelopt",
@@ -418,9 +350,6 @@ class GPTModelExporter:
             and self.is_multimodal
             and pretrained_model_name_or_path is not None
         ):
-            assert not self.export_vllm_fq_weights_qstate, (
-                "Exporting weights in bf16 and amax values is not supported for multimodal models"
-            )
             hf_checkpoint_path = Path(pretrained_model_name_or_path)
             if not hf_checkpoint_path.is_dir():
                 hf_checkpoint_path = tempfile.gettempdir() + "/" + pretrained_model_name_or_path
@@ -489,9 +418,6 @@ class GPTModelExporter:
         torch.distributed.barrier()
 
         if self.export_extra_modules:
-            assert not self.export_vllm_fq_weights_qstate, (
-                "Exporting weights in bf16 and amax values is not supported for extra modules"
-            )
             if is_last_stage_main_rank:
                 save_file(
                     state_dict, save_directory + "/model.safetensors", metadata={"format": "pt"}
@@ -499,9 +425,6 @@ class GPTModelExporter:
             torch.distributed.barrier()
             return
 
-        if self.export_vllm_fq_weights_qstate:
-            state_dict = gather_mcore_vllm_fq_quantized_state_dict(state_dict, save_directory)
-
         if (
             is_last_stage_main_rank
             and self._hf_config is not None
@@ -581,6 +504,65 @@ class GPTModelExporter:
 
         return all_rules
 
+    def _get_quantized_state(
+        self,
+        module: torch.nn.Module,
+        dtype: torch.dtype = torch.float16,
+    ) -> tuple[dict[str, torch.Tensor], str, int]:
+        """Return a state_dict, quantization format, and block_size of the module.
+
+        Args:
+            module: The target module to perform real quantization.
+            dtype: The default data type.
+
+        Returns:
+            Tuple: state_dict, quantization format, and block_size of the module.
+        """
+        name_to_value = {}
+        qformat: str = self._get_quantization_format(module)
+        block_size = get_weight_block_size(module)
+
+        if hasattr(module, "weight") and module.weight is not None:
+            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:
+            name_to_value["bias"] = module.bias.to(dtype).cpu()
+
+        if hasattr(module, "expert_bias") and module.expert_bias is not None:
+            name_to_value["expert_bias"] = module.expert_bias.to(dtype).cpu()
+
+        if qformat == QUANTIZATION_NONE:
+            return name_to_value, qformat, block_size
+        # Getting the weight scales
+        weight_scale = get_weight_scaling_factor(module)
+        weight_scale_2 = get_weight_scaling_factor_2(module)
+        if weight_scale is not None:
+            name_to_value["weight_scale"] = weight_scale
+
+        if weight_scale_2 is not None:
+            name_to_value["weight_scale_2"] = weight_scale_2
+
+        # Getting the input scale
+        input_scale = get_activation_scaling_factor(module)
+        if input_scale is not None:
+            name_to_value["input_scale"] = input_scale
+            # TODO (chenhany): support AWQ with pre_quant_scale
+            if hasattr(module.input_quantizer, "_pre_quant_scale"):
+                raise ValueError("Detect pre_quant_scale! SmoothQuant/AWQ are not yet supported!")
+
+        if hasattr(module, "output_quantizer"):
+            output_scale = get_kv_cache_scaling_factor(module)
+            if output_scale is not None:
+                name_to_value["output_scale"] = output_scale
+
+        return name_to_value, qformat, block_size
+
+    def _get_quantization_format(self, module: torch.nn.Module):
+        return get_quantization_format(module)
+
     def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str):
         weight_scale = quantized_state.pop("weight_scale", None)
         weight_scale_2 = quantized_state.pop("weight_scale_2", None)
@@ -609,9 +591,7 @@ class GPTModelExporter:
             self._state_dict[prefix] = module
             return
 
-        name_to_value, qformat, block_size = get_quantized_state(
-            module, dtype, self.export_vllm_fq_weights_qstate
-        )
+        name_to_value, qformat, block_size = self._get_quantized_state(module, dtype)
 
         weight = name_to_value.pop("weight")
         weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
@@ -643,9 +623,7 @@ class GPTModelExporter:
     def _gated_mlp_slicing(
         self, module, prefix, gate_proj_name="gate_proj", up_proj_name="up_proj"
     ):
-        name_to_value, qformat, block_size = get_quantized_state(
-            module, self.dtype, self.export_vllm_fq_weights_qstate
-        )
+        name_to_value, qformat, block_size = self._get_quantized_state(module, self.dtype)
 
         weight = name_to_value.pop("weight")
         weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
@@ -710,9 +688,7 @@ class GPTModelExporter:
         k_scale_name="k_scale",
         v_scale_name="v_scale",
     ):
-        name_to_value, qformat, block_size = get_quantized_state(
-            module, self.dtype, self.export_vllm_fq_weights_qstate
-        )
+        name_to_value, qformat, block_size = self._get_quantized_state(module, self.dtype)
 
         q_proj_prefix = prefix + q_proj_name + "."
         k_proj_prefix = prefix + k_proj_name + "."
@@ -834,8 +810,8 @@ class GPTModelExporter:
 
         for expert in module:
             assert layer_type is not None, "layer_type is required for pack_name_remapping"
-            name_to_value, qformat, block_size = get_quantized_state(
-                getattr(expert, layer_type), self.dtype, self.export_vllm_fq_weights_qstate
+            name_to_value, qformat, block_size = self._get_quantized_state(
+                getattr(expert, layer_type), self.dtype
             )
             weight = name_to_value.pop("weight")
             weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
@@ -900,8 +876,8 @@ class GPTModelExporter:
 
         for expert in module:
             assert layer_type is not None, "layer_type is required for pack_name_remapping"
-            name_to_value, qformat, block_size = get_quantized_state(
-                getattr(expert, layer_type), self.dtype, self.export_vllm_fq_weights_qstate
+            name_to_value, qformat, block_size = self._get_quantized_state(
+                getattr(expert, layer_type), self.dtype
             )
             weight = name_to_value.pop("weight")
             bias = name_to_value.pop("bias", None)
@@ -1241,7 +1217,6 @@ def export_mcore_gpt_to_hf(
     dtype: torch.dtype = torch.bfloat16,
     export_dir: Path | str = tempfile.gettempdir(),
     moe_router_dtype: torch.dtype | None = None,
-    export_vllm_fq_weights_qstate: bool = False,
 ):
     """Export Megatron Core GPTModel to unified checkpoint and save to export_dir.
 
@@ -1255,7 +1230,6 @@ def export_mcore_gpt_to_hf(
             eagle_module. Otherwise, only export the base model.
         dtype: The weights data type to export the unquantized layers.
         export_dir: The target export path.
-        export_vllm_fq_weights_qstate: If True, export the weights in bf16 and amax values.
     """
     exporter = GPTModelExporter(
         model,
@@ -1263,7 +1237,6 @@ def export_mcore_gpt_to_hf(
         export_extra_modules=export_extra_modules,
         dtype=dtype,
         moe_router_dtype=moe_router_dtype,
-        export_vllm_fq_weights_qstate=export_vllm_fq_weights_qstate,
     )
     exporter.save_pretrained(export_dir, pretrained_model_name_or_path)
 

tests/gpu/torch/export/test_vllm_fakequant_hf_export.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from copy import deepcopy
+
+import pytest
+import torch
+from _test_utils.torch.transformers_models import create_tiny_llama_dir
+from transformers import AutoModelForCausalLM
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.export import export_hf_vllm_fq_checkpoint
+
+
+@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG])
+def test_hf_vllm_export(tmp_path, quant_cfg):
+    """Test HuggingFace model export for vLLM with fake quantization.
+
+    This test verifies:
+    1. Model weights match before and after export
+    2. quant_amax.pth file is created, huggingface config file does not exist
+    3. Amax values are correctly extracted and saved in quant_amax.pth file
+    """
+
+    # Create a tiny LLaMA model for testing
+    tiny_model_dir = create_tiny_llama_dir(tmp_path, with_tokenizer=True, num_hidden_layers=2)
+
+    # Load the model
+    model = AutoModelForCausalLM.from_pretrained(tiny_model_dir)
+    model = model.cuda()
+    model.eval()
+
+    # Quantize the model
+    def forward_loop(model):
+        input_ids = torch.randint(0, model.config.vocab_size, (1, 128)).cuda()
+        with torch.no_grad():
+            model(input_ids)
+
+    model = mtq.quantize(model, quant_cfg, forward_loop)
+
+    model_state_dict = deepcopy(model.state_dict())
+
+    # Export directory
+    export_dir = tmp_path / "vllm_export"
+    export_dir.mkdir(exist_ok=True)
+
+    # Export for vLLM
+    export_hf_vllm_fq_checkpoint(model, export_dir=export_dir)
+
+    # check if quant_amax.pth file exists
+    quant_amax_file = export_dir / "quant_amax.pth"
+    assert quant_amax_file.exists(), f"quant_amax.pth file should be created in {export_dir}"
+
+    # make sure hf_quant_config.json file does not exist
+    hf_quant_config_file = export_dir / "hf_quant_config.json"
+    assert not hf_quant_config_file.exists(), (
+        f"hf_quant_config.json file should not be created in {export_dir}"
+    )
+
+    # check weights match before and after export
+    model_after = AutoModelForCausalLM.from_pretrained(export_dir)
+    model_after = model_after.cuda()
+    model_after.eval()
+    model_after_state_dict = model_after.state_dict()
+    amax_state_dict = {}
+    for key, param in model_state_dict.items():
+        if key.endswith("_amax"):
+            amax_state_dict[key] = param
+            continue
+
+        assert torch.allclose(param, model_after_state_dict[key], atol=1e-6), (
+            f"Weight mismatch for {key}: "
+            f"before shape={param.shape}, after shape={model_after_state_dict[key].shape}, "
+            f"max diff={torch.abs(param - model_after_state_dict[key]).max()}"
+        )
+
+    # Verify amax values are correct
+    amax_dict = torch.load(quant_amax_file)
+    assert len(amax_dict) > 0, "amax_dict should not be empty"
+    assert amax_dict.keys() == amax_state_dict.keys(), (
+        "amax keys mismatch between before and after export"
+    )

tests/gpu/torch/export/test_vllm_fakequant_megatron_export.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import json
-from copy import deepcopy
 from functools import partial
 
 import pytest
@@ -22,86 +21,13 @@ import torch
 from _test_utils.import_helper import skip_if_no_megatron
 from _test_utils.torch.distributed.utils import spawn_multiprocess_job
 from _test_utils.torch.megatron.models import get_mcore_gpt_model
-from _test_utils.torch.transformers_models import create_tiny_llama_dir
-from transformers import AutoModelForCausalLM
 
 import modelopt.torch.quantization as mtq
-from modelopt.torch.export.unified_export_hf import export_hf_checkpoint
-from modelopt.torch.export.unified_export_megatron import export_mcore_gpt_to_hf
+from modelopt.torch.export import export_mcore_gpt_to_hf_vllm_fq
 
 skip_if_no_megatron(apex_or_te_required=True)
 
 
-@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG])
-def test_hf_vllm_export(tmp_path, quant_cfg):
-    """Test HuggingFace model export for vLLM with fake quantization.
-
-    This test verifies:
-    1. Model weights match before and after export
-    2. quant_amax.pth file is created, huggingface config file does not exist
-    3. Amax values are correctly extracted and saved in quant_amax.pth file
-    """
-
-    # Create a tiny LLaMA model for testing
-    tiny_model_dir = create_tiny_llama_dir(tmp_path, with_tokenizer=True, num_hidden_layers=2)
-
-    # Load the model
-    model = AutoModelForCausalLM.from_pretrained(tiny_model_dir)
-    model = model.cuda()
-    model.eval()
-
-    # Quantize the model
-    def forward_loop(model):
-        input_ids = torch.randint(0, model.config.vocab_size, (1, 128)).cuda()
-        with torch.no_grad():
-            model(input_ids)
-
-    model = mtq.quantize(model, quant_cfg, forward_loop)
-
-    model_state_dict = deepcopy(model.state_dict())
-
-    # Export directory
-    export_dir = tmp_path / "vllm_export"
-    export_dir.mkdir(exist_ok=True)
-
-    # Export for vLLM
-    export_hf_checkpoint(model, export_dir=export_dir, export_vllm_fq_weights_qstate=True)
-
-    # check if quant_amax.pth file exists
-    quant_amax_file = export_dir / "quant_amax.pth"
-    assert quant_amax_file.exists(), f"quant_amax.pth file should be created in {export_dir}"
-
-    # make sure hf_quant_config.json file does not exist
-    hf_quant_config_file = export_dir / "hf_quant_config.json"
-    assert not hf_quant_config_file.exists(), (
-        f"hf_quant_config.json file should not be created in {export_dir}"
-    )
-
-    # check weights match before and after export
-    model_after = AutoModelForCausalLM.from_pretrained(export_dir)
-    model_after = model_after.cuda()
-    model_after.eval()
-    model_after_state_dict = model_after.state_dict()
-    amax_state_dict = {}
-    for key, param in model_state_dict.items():
-        if key.endswith("_amax"):
-            amax_state_dict[key] = param
-            continue
-
-        assert torch.allclose(param, model_after_state_dict[key], atol=1e-6), (
-            f"Weight mismatch for {key}: "
-            f"before shape={param.shape}, after shape={model_after_state_dict[key].shape}, "
-            f"max diff={torch.abs(param - model_after_state_dict[key]).max()}"
-        )
-
-    # Verify amax values are correct
-    amax_dict = torch.load(quant_amax_file)
-    assert len(amax_dict) > 0, "amax_dict should not be empty"
-    assert amax_dict.keys() == amax_state_dict.keys(), (
-        "amax keys mismatch between before and after export"
-    )
-
-
 def _test_mcore_vllm_export(tmp_path, quant_cfg, rank, size):
     """Test megatron-core model export for vLLM with fake quantization."""
     # Create a tiny mcore GPT model
@@ -165,12 +91,11 @@ def _test_mcore_vllm_export(tmp_path, quant_cfg, rank, size):
     export_dir.mkdir(exist_ok=True)
 
     # Export for vLLM
-    export_mcore_gpt_to_hf(
+    export_mcore_gpt_to_hf_vllm_fq(
         model,
         pretrained_model_name_or_path=tmp_path,
         dtype=torch.bfloat16,
         export_dir=str(export_dir),
-        export_vllm_fq_weights_qstate=True,
     )
 
     # check if quant_amax.pth file exists