fix(hf_ptq): use no_grad instead of inference_mode in export_quantized (NVBug 6537702)

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)
 ^^^^^^^^^^^^^^^^^

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__}"

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