From fb737b5b03a14d09a0a0c983419a9856ef87d04a Mon Sep 17 00:00:00 2001 From: Zhiyu Cheng Date: Sun, 2 Aug 2026 23:18:36 -0700 Subject: [PATCH] fix(hf_ptq): use no_grad instead of inference_mode in export_quantized `export_quantized` wrapped its whole body in `torch.inference_mode()`. On the FSDP2 path (`--use_fsdp2`), `get_model_state_dict(full_state_dict=True)` gathers the full params inside that context, so the gathered tensors are inference tensors; the subsequent `state_dict()` -> `param.detach()` then fails with `RuntimeError: Cannot set version_counter for inference tensor`. Switching the export context to `torch.no_grad()` keeps the gathered params as normal tensors (version counter intact) so `detach()` works, while still disabling autograd. Original fix by Shengliang Xu, verified end-to-end on 2 nodes with dense Qwen3-8B and Qwen3-30B-A3B (MoE) FSDP2 PTQ fp8 exports. Fixes NVBug 6537702 (Llama-3.1-8B-Instruct, 2x8 GB200, fp8_default-kv_fp8). Co-authored-by: Shengliang Xu Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Zhiyu Cheng --- CHANGELOG.rst | 1 + examples/hf_ptq/hf_ptq.py | 4 +- .../hf_ptq/test_export_quantized_context.py | 58 +++++++++++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 tests/examples/hf_ptq/test_export_quantized_context.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6e6289aca13..7a188669cb8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -101,6 +101,7 @@ Changelog - Fix vLLM 0.24+ support, where ``FusedMoE`` became a factory function and the expert weights moved onto a ``RoutedExperts`` submodule. Registering the old class broke every ``QuantModuleRegistry`` lookup with ``TypeError: issubclass() arg 2 must be a class``; the plugin now registers whichever fused-MoE module class the installed release provides. The registry key moves ``vllm_FusedMoE`` to ``vllm_RoutedExperts`` and quantizer paths gain ``.routed_experts``, so an older ``modelopt_state`` does not restore onto 0.24+ as-is. - Fix ``examples/vllm_serve`` serving shared experts uncalibrated: their ``gate_proj``/``up_proj`` quantizer keys were not merged into ``gate_up_proj`` on reload, so they matched no module and were dropped. - Fix Qwen3-VL MoE PTQ failing on ``transformers>=5.12`` with ``AttributeError: 'QuantQwen3VLMoeTextExperts' object has no attribute 'hidden_size'`` (NVBug 6518551). transformers 5.12 moved ``Qwen3VLMoeTextExperts`` onto the standard ``@use_experts_implementation`` fused layout (``hidden_size``/``expert_dim`` renamed to ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` transposed to ``(num_experts, 2*intermediate_dim, hidden_dim)``, two ``F.linear`` calls per expert), but the legacy ``_QuantQwen3VLMoeTextExperts`` wrapper stayed statically registered and shadowed on-the-fly detection. The new layout is now left to ``register_fused_experts_on_the_fly``, which claims it with the generic ``_QuantFusedExperts``; the legacy wrapper is still registered on ``transformers<5.12``, whose ``torch.bmm``-based forward the generic wrapper cannot intercept. +- Fix ``examples/hf_ptq`` multi-node FSDP2 export (``--use_fsdp2``) failing with ``RuntimeError: Cannot set version_counter for inference tensor``. ``export_quantized`` wrapped its whole body in ``torch.inference_mode()``, so the full params gathered by ``get_model_state_dict(full_state_dict=True)`` were inference tensors and the subsequent ``state_dict()`` -> ``param.detach()`` could not set their version counter. The export context is now ``torch.no_grad()``, which still disables autograd but keeps the gathered params as normal tensors. 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index d4eebc97ddf..6ddfba6614c 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -820,7 +820,9 @@ def export_quantized( default_padding_side, default_pad_token, ): - with torch.inference_mode(): + # Not inference_mode: the FSDP2 path gathers full params in this context and + # inference tensors break the subsequent state_dict() -> param.detach(). + with torch.no_grad(): if model_type is None: print(f"Unknown model type {type(language_model).__name__}. Continue exporting...") model_type = f"unknown:{type(language_model).__name__}" diff --git a/tests/examples/hf_ptq/test_export_quantized_context.py b/tests/examples/hf_ptq/test_export_quantized_context.py new file mode 100644 index 00000000000..f4364296b85 --- /dev/null +++ b/tests/examples/hf_ptq/test_export_quantized_context.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +import ast +import importlib +import inspect +import textwrap +from pathlib import Path + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" + + +def _import_hf_ptq(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + return importlib.import_module("hf_ptq") + + +def _context_manager_calls(func): + """Return the dotted names of every context manager entered by ``with`` in ``func``.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(func))) + names = [] + for node in ast.walk(tree): + if isinstance(node, ast.With): + for item in node.items: + expr = item.context_expr + target = expr.func if isinstance(expr, ast.Call) else expr + if isinstance(target, (ast.Attribute, ast.Name)): + names.append(ast.unparse(target)) + return names + + +def test_export_quantized_does_not_use_inference_mode(monkeypatch): + """``export_quantized`` must not run under ``torch.inference_mode()``. + + On the FSDP2 path (``--use_fsdp2``, multi-node) the export gathers the full params + via ``get_model_state_dict(full_state_dict=True)`` inside this context. Tensors + allocated under ``inference_mode`` are inference tensors whose version counter + cannot be set, so the subsequent ``state_dict()`` -> ``param.detach()`` fails with + ``RuntimeError: Cannot set version_counter for inference tensor``. ``torch.no_grad()`` + disables autograd just the same but keeps the gathered params as normal tensors. + """ + hf_ptq = _import_hf_ptq(monkeypatch) + contexts = _context_manager_calls(hf_ptq.export_quantized) + + assert "torch.inference_mode" not in contexts + assert "torch.no_grad" in contexts