cleanup

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

examples/vllm_serve/README.md

@@ -35,23 +35,50 @@ for details about installing partial dependency sets.
 
 ## Calibrate and serve fake quant model in vLLM
 
-Step 1: Configure quantization settings.  
-You can either edit the `quant_config` dictionary in `vllm_serve_fakequant.py`, or set the following environment variables to control quantization behavior:
+Step 1: Configure quantization with the ModelOpt CLI flags below. Each flag falls back to its corresponding environment variable when omitted:
 
-| Variable        | Description                                      | Default             |
-|-----------------|--------------------------------------------------|---------------------|
-| QUANT_DATASET   | Dataset name for calibration                     | cnn_dailymail       |
-| QUANT_CALIB_SIZE| Number of samples used for calibration           | 512                 |
-| QUANT_CFG       | Quantization config                              | None                |
-| KV_QUANT_CFG    | KV-cache quantization config                     | None                |
-| QUANT_FILE_PATH | Optional path to exported quantizer state dict `quantizer_state.pth` | None |
-| MODELOPT_STATE_PATH | Optional path to exported `vllm_fq_modelopt_state.pth` (restores quantizer state and parameters) | None |
-| CALIB_BATCH_SIZE | Calibration batch size                           | 1                  |
-| RECIPE_PATH      | Optional path to a ModelOpt PTQ recipe YAML  | None |
+| CLI flag | Environment fallback | Description | Default / auto-detection |
+| --- | --- | --- | --- |
+| `--modelopt-quant-cfg` | `QUANT_CFG` | ModelOpt weight/activation quantization config name | Unset |
+| `--modelopt-kv-quant-cfg` | `KV_QUANT_CFG` | KV-cache quantization config name | Unset |
+| `--modelopt-quant-file-path` | `QUANT_FILE_PATH` | Megatron-exported quantizer tensor state | `<model_dir>/quantizer_state.pth` when present |
+| `--modelopt-state-path` | `MODELOPT_STATE_PATH` | Full HF ModelOpt fakequant state | `<model_dir>/vllm_fq_modelopt_state.pth` when present |
+| `--modelopt-recipe-path` | `RECIPE_PATH` | PTQ recipe YAML or Megatron per-quantizer resolved config | `<model_dir>/vllm_fq_quantizer_state.yaml` when present |
+| `--modelopt-quant-dataset` | `QUANT_DATASET` | Calibration dataset | `cnn_dailymail` |
+| `--modelopt-quant-calib-size` | `QUANT_CALIB_SIZE` | Number of calibration samples | `512` |
+| `--modelopt-calib-batch-size` | `CALIB_BATCH_SIZE` | Calibration batch size | `1` |
 
-Set these variables in your shell or Docker environment as needed to customize calibration.
+CLI values take precedence over their environment fallbacks. `--modelopt-quant-cfg` /
+`--modelopt-kv-quant-cfg` and `--modelopt-recipe-path` are mutually exclusive because a recipe
+already carries its quantization configuration. For a local model directory, HF full state
+auto-detection takes precedence over Megatron quantizer-state/recipe sidecars.
 
-Step 2: Run the following command, with all supported flag as `vllm serve`:
+Install the shim to expose these options through the ordinary `vllm` command:
+
+```bash
+pip install -e examples/vllm_serve
+```
+
+Step 2: Serve with the shim. It accepts every stock vLLM flag plus the ModelOpt flags above:
+
+```bash
+vllm serve <model_path> -tp 8 --host 0.0.0.0 --port 8000 \
+  --modelopt-quant-cfg NVFP4_DEFAULT_CFG \
+  --modelopt-quant-dataset cnn_dailymail \
+  --modelopt-quant-calib-size 512
+```
+
+For an exported HF or Megatron fakequant directory containing the standard sidecars, no
+ModelOpt path flags are required:
+
+```bash
+vllm serve <export_dir> -tp 8 --host 0.0.0.0 --port 8000
+```
+
+When fakequant is requested explicitly or auto-detected, the shim selects
+`fakequant_worker.FakeQuantWorker` unless `--worker-cls` is supplied. With no ModelOpt flags,
+environment settings, or recognized sidecars, it delegates to the stock vLLM CLI unchanged.
+The legacy direct invocation remains available:
 
 ```bash
 python vllm_serve_fakequant.py <model_path> -tp 8 --host 0.0.0.0 --port 8000

examples/vllm_serve/vllm_reload_utils.py

@@ -122,6 +122,24 @@ def _convert_key_for_vllm(key: str, value: Any) -> tuple[str, str | None, Any]:
         )
         return ("group", group_key, value)
 
+    # Fused HF expert containers keep one activation quantizer beside their packed 3-D
+    # expert weights instead of one quantizer below each ``experts.<index>`` module.
+    # Nemotron-H's non-gated experts use ``up_proj_*``; gated variants use
+    # ``gate_up_proj_*``. Both feed vLLM's routed-expert w13 projection.
+    fused_expert_first_match = re.search(
+        r"(.*\.experts)\.(?:gate_up_proj|up_proj)_([^.]+_quantizer)(\..+)?$", key
+    )
+    if fused_expert_first_match:
+        suffix = fused_expert_first_match.group(3) or ""
+        group_key = (
+            fused_expert_first_match.group(1)
+            + _EXPERTS_INFIX
+            + ".w13_"
+            + fused_expert_first_match.group(2)
+            + suffix
+        )
+        return ("group", group_key, value)
+
     # Check if this is a non-expert gate/up projection that needs merging. Only *routed* experts
     # (``experts.<i>.``) merge into w13/w2 above; shared experts are a plain MLP whose gate/up
     # still merge into ``gate_up_proj``, so they must not be excluded by the "experts" substring.
@@ -144,6 +162,20 @@ def _convert_key_for_vllm(key: str, value: Any) -> tuple[str, str | None, Any]:
         )
         return ("group", group_key, value)
 
+    fused_expert_down_match = re.search(
+        r"(.*\.experts)\.down_proj_([^.]+_quantizer)(\..+)?$", key
+    )
+    if fused_expert_down_match:
+        suffix = fused_expert_down_match.group(3) or ""
+        group_key = (
+            fused_expert_down_match.group(1)
+            + _EXPERTS_INFIX
+            + ".w2_"
+            + fused_expert_down_match.group(2)
+            + suffix
+        )
+        return ("group", group_key, value)
+
     # Transform bmm_quantizer keys: self_attn.q/k/v_bmm_quantizer -> self_attn.attn.q/k/v_bmm_quantizer
     bmm_match = re.search(r"(.*\.self_attn)\.([qkv]_bmm_quantizer.*)$", key) or re.search(
         r"(.*\.mixer)\.([qkv]_bmm_quantizer.*)$", key

modelopt/torch/export/plugins/vllm_fakequant_hf.py

@@ -18,6 +18,7 @@ import copy
 import json
 import logging
 import re
+import shutil
 import warnings
 from collections.abc import Callable
 from contextlib import ExitStack, contextmanager
@@ -569,6 +570,53 @@ def _carry_over_unplaced_weights(export_dir: Path, model: nn.Module) -> None:
     (export_dir / "model.safetensors.index.json").write_text(json.dumps(index, indent=2))
 
 
+def _copy_source_config(model: nn.Module, export_dir: Path) -> None:
+    """Preserve the source checkpoint's config.json without reserializing it.
+
+    A model loaded by a newer Transformers release can use a native config class whose
+    serialized schema differs from the checkpoint's remote-code config class. Since an HF
+    fakequant export retains that remote code, its config must come from the same checkpoint
+    revision. save_pretrained writes its generated config first; this helper replaces it
+    with the original bytes.
+    """
+    source = getattr(getattr(model, "config", None), "_name_or_path", None)
+    if not source:
+        warnings.warn(
+            "Could not identify the source checkpoint config; keeping the config generated "
+            "by save_pretrained."
+        )
+        return
+
+    source_config = Path(source) / "config.json"
+    if not source_config.is_file():
+        try:
+            from transformers.utils import cached_file
+
+            resolved = cached_file(
+                source,
+                "config.json",
+                _raise_exceptions_for_gated_repo=False,
+                _raise_exceptions_for_missing_entries=False,
+                _raise_exceptions_for_connection_errors=False,
+            )
+            source_config = Path(resolved) if resolved else source_config
+        except Exception as e:
+            logging.getLogger(__name__).warning(
+                "Could not resolve source config for %r: %s", source, e
+            )
+
+    if not source_config.is_file():
+        warnings.warn(
+            f"Source checkpoint config not found for {source!r}; keeping the config generated "
+            "by save_pretrained."
+        )
+        return
+
+    export_config = export_dir / "config.json"
+    if source_config.resolve() != export_config.resolve():
+        shutil.copyfile(source_config, export_config)
+
+
 def export_hf_vllm_fq_checkpoint(
     model: nn.Module,
     export_dir: Path | str,
@@ -755,6 +803,7 @@ def export_hf_vllm_fq_checkpoint(
         # inplace_mem_efficient branch (it deliberately omits state_dict= there -- see the
         # comment above -- so there is no state_dict to merge extras into).
         _carry_over_unplaced_weights(export_dir, model)
+        _copy_source_config(model, export_dir)
 
     finally:
         if not inplace_mem_efficient:

tests/gpu/torch/export/test_vllm_fakequant_hf_export.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from copy import deepcopy
+from pathlib import Path
 
 import pytest
 import torch
@@ -39,6 +40,7 @@ def _test_hf_vllm_export(tmp_path, quant_cfg, model_dir):
     """
 
     # Load the model
+    source_config = (Path(model_dir) / "config.json").read_bytes()
     model = AutoModelForCausalLM.from_pretrained(model_dir)
     model = model.cuda()
     model.eval()
@@ -84,6 +86,9 @@ def _test_hf_vllm_export(tmp_path, quant_cfg, model_dir):
 
     export_hf_vllm_fq_checkpoint(model, export_dir=export_dir)
 
+    # The exporter must preserve the source checkpoint's config bytes for remote-code parity.
+    assert (export_dir / "config.json").read_bytes() == source_config
+
     # Verify the input model is not mutated: all state dict values unchanged
     state_dict_after_export = model.state_dict()
     for key, param_before in state_dict_before_export.items():

tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py

@@ -906,6 +906,8 @@ def test_tiny_qwen3_moe_quantize(tiny_qwen3_moe_llm):
     for hf_key, expected_quantizer in (
         ("model.layers.0.mlp.experts.0.gate_proj.input_quantizer._amax", "w13_input_quantizer"),
         ("model.layers.0.mlp.experts.0.down_proj.weight_quantizer._amax", "w2_weight_quantizer"),
+        ("model.layers.0.mlp.experts.up_proj_input_quantizer._amax", "w13_input_quantizer"),
+        ("model.layers.0.mlp.experts.down_proj_input_quantizer._amax", "w2_input_quantizer"),
     ):
         action, vllm_key, _ = reload_utils._convert_key_for_vllm(hf_key, 1.0)
         assert action == "group", (hf_key, action)