From 2561773a470b538f818f3e2c857d05d408671641 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:35:52 +0000 Subject: [PATCH 1/9] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models export_hf_checkpoint on an FSDP2 model quantized with W4A8_AWQ_BETA_CFG failed on every rank with "AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'". The export's AWQ resmooth opens an fsdp2_aware_weight_update window on each sharded layer, and the window's enable_fake_quant read and wrote weight_quantizer._fake_quant, which a SequentialQuantizer (W4A8's INT4 + FP8 weight quantizer) does not have. enable_fake_quant now sets and restores the flag on each stage of a weight quantizer. Re-enable W4A8 in the FSDP2 GPU export tests, which left it out for this reason, and add a one-rank gloo CPU test of the streaming export. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + .../torch/quantization/utils/core_utils.py | 23 +++++++++------- tests/gpu/torch/export/test_fsdp2_export.py | 4 +-- .../export/test_fsdp2_parallel_export.py | 27 ++++++++++++++++--- 4 files changed, 41 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4ea103a0628..8fcc8ca5bae 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -98,6 +98,7 @@ Changelog - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``. - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization. - Fix weight calibration, GPTQ and weight folding of Hugging Face fused MoE experts sharded by the expert parallelism of Transformers 5.16 and later, which used another rank's expert weights on every rank but the first. +- Fix ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``, e.g. ``examples/hf_ptq`` with ``--use_fsdp2 --qformat w4a8_awq_beta``), which failed with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 8bb3ac69288..7c436d93616 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -1005,17 +1005,22 @@ def patched_new(cls, data=None, requires_grad=True): def enable_fake_quant(module): """Temporarily set the fake_quant attribute of a module to True. - This is used to prevent weight compression from being triggered during an unshard() call. + This is used to prevent weight compression from being triggered during an unshard() call. A + sequential weight quantizer (W4A8) holds the flag on each of its stages. """ - original_fake_quant = [] - for m in module.modules(): - if hasattr(m, "weight_quantizer"): - original_fake_quant.append(m.weight_quantizer._fake_quant) - m.weight_quantizer._fake_quant = True + quantizers = [ + q + for m in module.modules() + if hasattr(m, "weight_quantizer") + for q in m.weight_quantizer.modules() + if hasattr(q, "_fake_quant") + ] + original_fake_quant = [q._fake_quant for q in quantizers] + for q in quantizers: + q._fake_quant = True yield - for m in module.modules(): - if hasattr(m, "weight_quantizer"): - m.weight_quantizer._fake_quant = original_fake_quant.pop(0) + for q, fake_quant in zip(quantizers, original_fake_quant): + q._fake_quant = fake_quant @contextmanager diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index ce71b1c4270..3290eb3c849 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -305,7 +305,7 @@ def test_fsdp2_weight_update_context_error_propagates(dist_workers): mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG, mtq.NVFP4_AWQ_LITE_CFG, - # mtq.W4A8_AWQ_BETA_CFG, #TODO: Fix unit test for this case + mtq.W4A8_AWQ_BETA_CFG, # mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, #TODO: Fix unit test for this case mtq.W4A8_MXFP4_FP8_CFG, # NVFP4_MLP_ONLY_CFG is omitted: SmallQKVModel has no MLP, so that config matched @@ -327,7 +327,7 @@ def test_fsdp2_weight_update_context_for_fuse_layers(dist_workers, quant_config, mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG, mtq.NVFP4_AWQ_LITE_CFG, - # mtq.W4A8_AWQ_BETA_CFG, #TODO: Fix unit test for this case + mtq.W4A8_AWQ_BETA_CFG, # mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, #TODO: Fix unit test for this case mtq.W4A8_MXFP4_FP8_CFG, # NVFP4_MLP_ONLY_CFG is omitted: SmallQKVModel has no MLP, so that config matched diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index c80f704bb6d..9885ad0ea15 100644 --- a/tests/unit/torch/export/test_fsdp2_parallel_export.py +++ b/tests/unit/torch/export/test_fsdp2_parallel_export.py @@ -42,13 +42,13 @@ from transformers import AutoModelForCausalLM, LlamaConfig -def _tiny_quantized_llama(quant_cfg=None, tie=True, mesh=None): +def _tiny_quantized_llama(quant_cfg=None, tie=True, mesh=None, hidden_size=64): mto.enable_huggingface_checkpointing() torch.manual_seed(0) cfg = LlamaConfig( vocab_size=128, - hidden_size=64, - intermediate_size=128, + hidden_size=hidden_size, + intermediate_size=2 * hidden_size, num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=4, @@ -94,6 +94,27 @@ def test_streaming_export_matches_resident_path(tmp_path): assert torch.equal(streamed[key].float(), ref[key].cpu().float()), key +def test_streaming_export_of_a_sharded_w4a8_model(tmp_path, gloo_world_of_one): + """An FSDP2 model quantized with W4A8 exports what the unsharded model does. + + W4A8's weight quantizer is a ``SequentialQuantizer``, which the weight-update window the AWQ + resmooth opens on each sharded layer must handle. + """ + # 128 wide, so that W4A8's 128-wide INT4 blocks tile every weight. + model = _tiny_quantized_llama(mtq.W4A8_AWQ_BETA_CFG, hidden_size=128) + ref, _ = _export_transformers_checkpoint(copy.deepcopy(model), torch.bfloat16) + for layer in model.model.layers: + fully_shard(layer, mesh=gloo_world_of_one) + fully_shard(model, mesh=gloo_world_of_one) + + _export_fsdp2_checkpoint_streaming(model, torch.bfloat16, export_dir=tmp_path) + + streamed = _load_all(tmp_path) + assert set(streamed) == set(ref) + for key in ref: + assert torch.equal(streamed[key], ref[key].cpu()), key + + def test_streaming_export_drops_tied_alias_and_writes_config(tmp_path): model = _tiny_quantized_llama() d = tmp_path From 20d76ff495d353a42dfb061ee12c73c379aae697 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:02:25 +0000 Subject: [PATCH 2/9] Restore enable_fake_quant's flags when the window raises An exception inside the window, such as an FSDP2 unshard running out of memory, left every weight quantizer it reached at _fake_quant=True, which later makes compression skip those weights silently. Restore the saved flags in a finally block, as the file's other context managers do, and pick the stages with isinstance(TensorQuantizer), the file's idiom for quantizer modules. Signed-off-by: Shengliang Xu --- .../torch/quantization/utils/core_utils.py | 12 ++++++++---- tests/unit/torch/quantization/test_utils.py | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 7c436d93616..ddf2f4d91d9 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -1008,19 +1008,23 @@ def enable_fake_quant(module): This is used to prevent weight compression from being triggered during an unshard() call. A sequential weight quantizer (W4A8) holds the flag on each of its stages. """ + from ..nn import TensorQuantizer # local: ..nn imports this utils package at module scope + quantizers = [ q for m in module.modules() if hasattr(m, "weight_quantizer") for q in m.weight_quantizer.modules() - if hasattr(q, "_fake_quant") + if isinstance(q, TensorQuantizer) ] original_fake_quant = [q._fake_quant for q in quantizers] for q in quantizers: q._fake_quant = True - yield - for q, fake_quant in zip(quantizers, original_fake_quant): - q._fake_quant = fake_quant + try: + yield + finally: + for q, fake_quant in zip(quantizers, original_fake_quant): + q._fake_quant = fake_quant @contextmanager diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index a22703fc91e..7af9bf7e6bf 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -22,6 +22,8 @@ from torch.distributed.fsdp import fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import ( convert_quantization_axis_to_reduce_axis, core_utils, @@ -33,6 +35,7 @@ _fsdp2_owned_params, _fsdp2_owns, _fsdp2_param_groups, + enable_fake_quant, fsdp2_weight_access_and_writeback_context, persistent_materialization, ) @@ -350,6 +353,21 @@ def _forward_loop(m): assert torch.allclose(inp2[0][0][0], torch.tensor([[10.0 + 1.0 + 2.0]])) +def test_enable_fake_quant_restores_each_stage_when_the_window_raises(): + """A W4A8 weight quantizer holds the flag on each of its stages. A window that raises, such as + an FSDP2 unshard running out of memory, must give every stage its own flag back.""" + module = torch.nn.Module() + module.weight_quantizer = SequentialQuantizer( + TensorQuantizer(QuantizerAttributeConfig(fake_quant=False)), TensorQuantizer() + ) + + with pytest.raises(RuntimeError), enable_fake_quant(module): + assert [q._fake_quant for q in module.weight_quantizer] == [True, True] + raise RuntimeError + + assert [q._fake_quant for q in module.weight_quantizer] == [False, True] + + def _forward_through_other_layers_while_one_is_materialized(rank, size, writeback): mesh = init_device_mesh("cpu", (size,)) torch.manual_seed(0) From db5112133191dca60715f86974e386b452b2ca7a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:14:00 +0000 Subject: [PATCH 3/9] Refuse compressing W4A8 (SequentialQuantizer) weights A W4A8 model compressed with mtq.compress can neither run nor be restored. Its forward dequantizes the INT4 stage and then real-quantizes the FP8 stage, so F.linear receives an FP8QTensor, and compress_restore calls set_from_modelopt_state, which SequentialQuantizer does not have. Every weight assignment also read the container's missing _fake_quant. compress_convert now raises NotImplementedError naming the module before it changes anything, as the Megatron tests already skip W4A8 compress. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/compress.py | 7 ++++ .../unit/torch/quantization/test_compress.py | 34 +++++++++++++++++++ 3 files changed, 42 insertions(+), 1 deletion(-) create mode 100644 tests/unit/torch/quantization/test_compress.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8fcc8ca5bae..9c72a190580 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -98,7 +98,7 @@ Changelog - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``. - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization. - Fix weight calibration, GPTQ and weight folding of Hugging Face fused MoE experts sharded by the expert parallelism of Transformers 5.16 and later, which used another rank's expert weights on every rank but the first. -- Fix ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``, e.g. ``examples/hf_ptq`` with ``--use_fsdp2 --qformat w4a8_awq_beta``), which failed with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``. +- Fix ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``, e.g. ``examples/hf_ptq`` with ``--use_fsdp2 --qformat w4a8_awq_beta``), which failed with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``. ``mtq.compress`` now refuses W4A8 weights with ``NotImplementedError``, since a compressed W4A8 model could neither run nor be reloaded. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/compress.py b/modelopt/torch/quantization/compress.py index ee7e642a0e8..93028c45fb7 100644 --- a/modelopt/torch/quantization/compress.py +++ b/modelopt/torch/quantization/compress.py @@ -32,6 +32,7 @@ from .backends.gemm_registry import disable_real_quant_gemm, enable_real_quant_gemm from .config import CompressCfgType, CompressConfig from .conversion import _replace_quant_module, set_quantizer_attributes_partial +from .nn import SequentialQuantizer from .nn.modules.quant_linear import RealQuantLinear from .qtensor import QTensorWrapper, pack_real_quantize_weight from .utils import is_quantized_linear @@ -69,6 +70,12 @@ def compress_convert( Returns: The compressed model. """ + # A compressed SequentialQuantizer weight can neither run forward nor be restored. + for name, module in model.named_modules(): + if is_quantized_linear(module) and isinstance(module.weight_quantizer, SequentialQuantizer): + raise NotImplementedError( + f"{name}: compressing W4A8 (SequentialQuantizer) weights is not supported." + ) for _, module in model.named_modules(): if is_quantized_linear(module) and type(module) not in RealQuantModuleRegistry: class_to_register = RealQuantLinear diff --git a/tests/unit/torch/quantization/test_compress.py b/tests/unit/torch/quantization/test_compress.py new file mode 100644 index 00000000000..18506be866d --- /dev/null +++ b/tests/unit/torch/quantization/test_compress.py @@ -0,0 +1,34 @@ +# 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. + +"""CPU tests of mtq.compress.""" + +import pytest +import torch +from torch import nn + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.compress import is_real_quantized + + +def test_compress_refuses_w4a8_weights_before_changing_the_model(): + """A compressed W4A8 (SequentialQuantizer) weight can neither run forward nor be restored.""" + model = nn.Sequential(nn.Linear(128, 128, bias=False)) # 128 wide for W4A8's INT4 blocks + mtq.quantize(model, mtq.W4A8_AWQ_BETA_CFG, lambda m: m(torch.randn(2, 128))) + + with pytest.raises(NotImplementedError, match="W4A8"): + mtq.compress(model) + + assert not is_real_quantized(model) From b0ad5bd574f581240b45f9b0e661eda8b0943bc8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:28:15 +0000 Subject: [PATCH 4/9] Refuse a sequential weight quantizer inside a container too Transformer Engine's grouped linear holds a GroupedQuantizer whose per-expert members become SequentialQuantizers under W4A8, which the refusal missed: compress then failed later on the container's missing _fake_quant. Walk the weight quantizer's modules instead. The message now names the mechanism (as in W4A8) and says that a layer the compress config excludes is refused as well, since its RealQuantLinear can neither be restored nor take a new weight either. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/compress.py | 10 +++++++--- tests/unit/torch/quantization/test_compress.py | 6 +++++- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/quantization/compress.py b/modelopt/torch/quantization/compress.py index 93028c45fb7..e1a30e4c079 100644 --- a/modelopt/torch/quantization/compress.py +++ b/modelopt/torch/quantization/compress.py @@ -70,11 +70,15 @@ def compress_convert( Returns: The compressed model. """ - # A compressed SequentialQuantizer weight can neither run forward nor be restored. + # Compressed, a SequentialQuantizer weight cannot run forward; compressed or excluded, its + # RealQuantLinear can neither be restored nor take a new weight. for name, module in model.named_modules(): - if is_quantized_linear(module) and isinstance(module.weight_quantizer, SequentialQuantizer): + if is_quantized_linear(module) and any( + isinstance(q, SequentialQuantizer) for q in module.weight_quantizer.modules() + ): raise NotImplementedError( - f"{name}: compressing W4A8 (SequentialQuantizer) weights is not supported." + f"{name}: compressing a sequential weight quantizer (as in W4A8) is not supported, " + "even when the compress config excludes the layer." ) for _, module in model.named_modules(): if is_quantized_linear(module) and type(module) not in RealQuantModuleRegistry: diff --git a/tests/unit/torch/quantization/test_compress.py b/tests/unit/torch/quantization/test_compress.py index 18506be866d..affd6536f2e 100644 --- a/tests/unit/torch/quantization/test_compress.py +++ b/tests/unit/torch/quantization/test_compress.py @@ -21,12 +21,16 @@ import modelopt.torch.quantization as mtq from modelopt.torch.quantization.compress import is_real_quantized +from modelopt.torch.quantization.nn import GroupedQuantizer -def test_compress_refuses_w4a8_weights_before_changing_the_model(): +@pytest.mark.parametrize("grouped", [False, True], ids=["linear", "grouped"]) +def test_compress_refuses_w4a8_weights_before_changing_the_model(grouped): """A compressed W4A8 (SequentialQuantizer) weight can neither run forward nor be restored.""" model = nn.Sequential(nn.Linear(128, 128, bias=False)) # 128 wide for W4A8's INT4 blocks mtq.quantize(model, mtq.W4A8_AWQ_BETA_CFG, lambda m: m(torch.randn(2, 128))) + if grouped: # one per expert, as Transformer Engine's grouped linear holds them + model[0].weight_quantizer = GroupedQuantizer(model[0].weight_quantizer) with pytest.raises(NotImplementedError, match="W4A8"): mtq.compress(model) From 580cdf73016c896d4ae79a67fd21d2f8901a24ea Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:46:59 +0000 Subject: [PATCH 5/9] Say in the changelog that multi-rank AWQ under FSDP2 is not supported yet The FSDP2 export of a W4A8 AWQ model no longer crashes, but with more than one rank each rank still calibrates AWQ on its own data and folds its own pre_quant_scale into its weight shard. Say that the export no longer fails and that multi-rank AWQ calibration under FSDP2 is not supported yet, instead of presenting the hf_ptq --use_fsdp2 W4A8 flow as fixed. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9c72a190580..0cf345477da 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -98,7 +98,7 @@ Changelog - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``. - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization. - Fix weight calibration, GPTQ and weight folding of Hugging Face fused MoE experts sharded by the expert parallelism of Transformers 5.16 and later, which used another rank's expert weights on every rank but the first. -- Fix ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``, e.g. ``examples/hf_ptq`` with ``--use_fsdp2 --qformat w4a8_awq_beta``), which failed with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``. ``mtq.compress`` now refuses W4A8 weights with ``NotImplementedError``, since a compressed W4A8 model could neither run nor be reloaded. +- ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``) no longer fails with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``; AWQ calibration of an FSDP2 model across more than one rank is not supported yet. ``mtq.compress`` now refuses W4A8 weights with ``NotImplementedError``, since a compressed W4A8 model could neither run nor be reloaded. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ From 78aa164588950114a99dc4c63ae4a9be403ce0f2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:07:13 +0000 Subject: [PATCH 6/9] Refuse AWQ calibration of an FSDP2 model sharded over several ranks Each rank calibrates on its own data, and awq_lite / awq_full fold the smoothing scale a rank picks into that rank's weight shard, so the shards of one weight disagree. The FSDP2 export of such a W4A8 model used to crash; with the fix it would write those mismatched shards. mtq.quantize and mtq.calibrate now raise NotImplementedError for it before anything changes. The decision reads only the run's size, the model's FSDP2 sharding and the algorithm, so every rank raises alike, before any collective. One rank is still allowed. The W4A8 export test now shards before calibrating, as hf_ptq --use_fsdp2 does on one GPU, and a two-rank test checks the refusal. The FSDP2 export GPU tests calibrate before sharding, and the nested FSDP2 backward GPU test skips INT4 AWQ on several GPUs. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/model_quant.py | 21 ++++++++++++++++ tests/gpu/torch/export/test_fsdp2_export.py | 14 ++++------- tests/gpu/torch/quantization/test_fsdp2.py | 2 ++ .../export/test_fsdp2_parallel_export.py | 14 +++++------ .../torch/quantization/test_quantize_cpu.py | 24 +++++++++++++++++++ 6 files changed, 59 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0cf345477da..acbea0cf19d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -98,7 +98,7 @@ Changelog - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``. - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization. - Fix weight calibration, GPTQ and weight folding of Hugging Face fused MoE experts sharded by the expert parallelism of Transformers 5.16 and later, which used another rank's expert weights on every rank but the first. -- ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``) no longer fails with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``; AWQ calibration of an FSDP2 model across more than one rank is not supported yet. ``mtq.compress`` now refuses W4A8 weights with ``NotImplementedError``, since a compressed W4A8 model could neither run nor be reloaded. +- ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``) no longer fails with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``; AWQ calibration of an FSDP2 model sharded over more than one rank, whose ranks would fold different smoothing scales into their weight shards, is rejected with ``NotImplementedError`` until it is supported. ``mtq.compress`` now refuses W4A8 weights with ``NotImplementedError``, since a compressed W4A8 model could neither run nor be reloaded. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 115dcf437ee..ecfaf7d347f 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -70,6 +70,25 @@ ] +def _refuse_awq_on_a_sharded_fsdp2_model(model: nn.Module, algorithm: QuantizeAlgoCfgType) -> None: + """Refuse AWQ calibration of an FSDP2 model sharded over several ranks, before it starts. + + Each rank calibrates on its own data and folds the AWQ scale it picks into its own weight + shard, so the shards would not match. Every rank decides alike, from the run and the model. + """ + if dist.size() == 1 or not dist.is_fsdp2_model(model): + return + methods = [mode.removesuffix("_calibrate") for mode, _ in get_modelike_from_algo_cfg(algorithm)] + awq = [method for method in methods if method in ("awq_lite", "awq_full")] + if awq: + raise NotImplementedError( + f"Calibrating with {', '.join(awq)} is not supported yet on an FSDP2 model sharded over " + f"{dist.size()} ranks: each rank calibrates on its own data and folds the AWQ scale it " + "picks into its own weight shard, so the shards would not match. Calibrate in a single " + "process, or with an algorithm other than AWQ." + ) + + # TODO: Descriptors for the supported algorithms def calibrate( model: nn.Module, @@ -99,6 +118,7 @@ def calibrate( Returns: The calibrated pytorch model. """ + _refuse_awq_on_a_sharded_fsdp2_model(model, algorithm) if forward_loop is not None: # get the number of arguments of forward_loop num_args = len(inspect.signature(forward_loop).parameters) @@ -439,6 +459,7 @@ def forward_loop(model) -> None: Returns: A pytorch model which has been quantized and calibrated. """ quantize_config = QuantizeConfig(**dict(config)) + _refuse_awq_on_a_sharded_fsdp2_model(model, config.get("algorithm")) if not is_quantized(model): model = apply_mode(model, mode=[("quantize", dict(config))], registry=QuantizeModeRegistry) else: diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 3290eb3c849..392c8359246 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -177,13 +177,10 @@ def _fuse_layers(rank, size, quant_config, bias): def calib_fn(x): return x(calib_data) - # Shard model - fully_shard(model) - torch.distributed.barrier() - - # Quantize model + # Quantize, then shard: AWQ calibration of a model sharded over several ranks is refused. mtq.quantize(model, quant_config, calib_fn) mtq.quantize(non_fsdp_model, quant_config, calib_fn) + fully_shard(model) torch.distributed.barrier() @@ -217,13 +214,10 @@ def _export_quantized_weight_test(rank, size, quant_config, bias): def calib_fn(x): return x(calib_data) - # Shard model - fully_shard(model) - torch.distributed.barrier() - - # Quantize model + # Quantize, then shard: AWQ calibration of a model sharded over several ranks is refused. mtq.quantize(model, quant_config, calib_fn) mtq.quantize(non_fsdp_model, quant_config, calib_fn) + fully_shard(model) torch.distributed.barrier() diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index c9e56d0e4fc..4f1dda81a14 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -137,6 +137,8 @@ def test_fsdp_simple_linear(dist_workers): "quant_cfg", [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_SMOOTHQUANT_CFG, mtq.INT4_AWQ_CFG] ) def test_nested_fsdp2_backward(quant_cfg, dist_workers): + if quant_cfg is mtq.INT4_AWQ_CFG and torch.cuda.device_count() > 1: + pytest.skip("AWQ calibration of an FSDP2 model sharded over several ranks is refused") dist_workers.run(partial(_test_nested_fsdp2_backward, quant_cfg=quant_cfg)) diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index 9885ad0ea15..d5bb39d016d 100644 --- a/tests/unit/torch/export/test_fsdp2_parallel_export.py +++ b/tests/unit/torch/export/test_fsdp2_parallel_export.py @@ -95,17 +95,17 @@ def test_streaming_export_matches_resident_path(tmp_path): def test_streaming_export_of_a_sharded_w4a8_model(tmp_path, gloo_world_of_one): - """An FSDP2 model quantized with W4A8 exports what the unsharded model does. + """An FSDP2 model calibrated with W4A8 AWQ on one rank exports what the unsharded model does. W4A8's weight quantizer is a ``SequentialQuantizer``, which the weight-update window the AWQ - resmooth opens on each sharded layer must handle. + resmooth opens on each sharded layer must handle. AWQ calibration of an FSDP2 model is refused + only over several ranks. """ # 128 wide, so that W4A8's 128-wide INT4 blocks tile every weight. - model = _tiny_quantized_llama(mtq.W4A8_AWQ_BETA_CFG, hidden_size=128) - ref, _ = _export_transformers_checkpoint(copy.deepcopy(model), torch.bfloat16) - for layer in model.model.layers: - fully_shard(layer, mesh=gloo_world_of_one) - fully_shard(model, mesh=gloo_world_of_one) + ref_model = _tiny_quantized_llama(mtq.W4A8_AWQ_BETA_CFG, hidden_size=128) + ref, _ = _export_transformers_checkpoint(ref_model, torch.bfloat16) + # Sharded before it is calibrated, as hf_ptq --use_fsdp2 does on one GPU. + model = _tiny_quantized_llama(mtq.W4A8_AWQ_BETA_CFG, hidden_size=128, mesh=gloo_world_of_one) _export_fsdp2_checkpoint_streaming(model, torch.bfloat16, export_dir=tmp_path) diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index f51a514e0e5..101337d9547 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -21,6 +21,7 @@ import pytest import torch +from _test_utils.torch.distributed.utils import bounded_collectives from _test_utils.torch.quantization.models import SimpleConv, SimpleConvLinear, SimpleLinear from _test_utils.torch.quantization.quantize_common import ( INT4_AWQ_CLIP_CFG, @@ -30,6 +31,8 @@ save_restore_test, ) from pydantic import ValidationError +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -42,6 +45,7 @@ register_quant_backend, unregister_quant_backend, ) +from modelopt.torch.quantization.utils import is_quantized from modelopt.torch.utils.distributed import ParallelState # A test config with double-quant (using `SequentialQuantizers`) @@ -729,6 +733,26 @@ def test_indexer_k_check_covers_all_ranks(cpu_dist_workers): cpu_dist_workers.run(_test_indexer_check_covers_all_ranks) +def _awq_on_a_sharded_fsdp2_model_job(rank, size): + model = SimpleLinear() + fully_shard(model, mesh=init_device_mesh("cpu", (size,))) + forward_loop = lambda m: m(m.get_input()) # noqa: E731 + with bounded_collectives(): + for config in (mtq.W4A8_AWQ_BETA_CFG, INT4_AWQ_FULL_CFG): # awq_lite, awq_full + with pytest.raises(NotImplementedError, match=f"FSDP2 model sharded over {size} ranks"): + mtq.quantize(model, config, forward_loop) + assert not is_quantized(model) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop) + with pytest.raises(NotImplementedError, match="awq_lite"): + mtq.calibrate(model, "awq_lite", forward_loop) + + +def test_awq_on_a_sharded_fsdp2_model_is_refused(cpu_dist_workers): + """Each rank would fold its own AWQ scale into its shard: every rank refuses alike, before the + model changes, through ``mtq.quantize`` and ``mtq.calibrate``.""" + cpu_dist_workers.run(_awq_on_a_sharded_fsdp2_model_job) + + @pytest.mark.parametrize( ("quantizer_name", "supported"), [ From d538a02519a0f4b7d093d26e96428a0002e89adf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:29:39 +0000 Subject: [PATCH 7/9] Refuse AWQ AutoQuantize candidates on a sharded FSDP2 model at entry AutoQuantize calibrates its candidates one by one through calibrate(), so on an FSDP2 model sharded over several ranks the AWQ guard fired partway through the search, after earlier candidates had changed the model. auto_quantize now checks every candidate before the search starts and names the AWQ one in the error; calibrate() keeps the check as a backstop. The KV-cache search already rejects any algorithm but max. The message reports the FSDP2 shard mesh's size, which can be narrower than the run's. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- modelopt/torch/quantization/model_quant.py | 16 +++++++++++++--- .../unit/torch/quantization/test_quantize_cpu.py | 12 +++++++++++- 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index acbea0cf19d..ff82cd101dd 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -98,7 +98,7 @@ Changelog - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``. - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization. - Fix weight calibration, GPTQ and weight folding of Hugging Face fused MoE experts sharded by the expert parallelism of Transformers 5.16 and later, which used another rank's expert weights on every rank but the first. -- ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``) no longer fails with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``; AWQ calibration of an FSDP2 model sharded over more than one rank, whose ranks would fold different smoothing scales into their weight shards, is rejected with ``NotImplementedError`` until it is supported. ``mtq.compress`` now refuses W4A8 weights with ``NotImplementedError``, since a compressed W4A8 model could neither run nor be reloaded. +- ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``) no longer fails with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``; AWQ calibration of an FSDP2 model sharded over more than one rank, whose ranks would fold different smoothing scales into their weight shards, is rejected with ``NotImplementedError`` until it is supported, as is ``mtq.auto_quantize`` with an AWQ candidate, such as its default ``NVFP4_AWQ_LITE_CFG``, before the search starts. ``mtq.compress`` now refuses W4A8 weights with ``NotImplementedError``, since a compressed W4A8 model could neither run nor be reloaded. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index ecfaf7d347f..0a69ae9ede9 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -25,6 +25,7 @@ import torch import torch.nn as nn +from torch.distributed.fsdp import FSDPModule import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist @@ -54,6 +55,7 @@ from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg from .nn import AnyQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer from .utils import is_quantized +from .utils.core_utils import _get_fsdp2_mesh __all__ = [ "auto_quantize", @@ -70,7 +72,9 @@ ] -def _refuse_awq_on_a_sharded_fsdp2_model(model: nn.Module, algorithm: QuantizeAlgoCfgType) -> None: +def _refuse_awq_on_a_sharded_fsdp2_model( + model: nn.Module, algorithm: QuantizeAlgoCfgType, candidate: str | None = None +) -> None: """Refuse AWQ calibration of an FSDP2 model sharded over several ranks, before it starts. Each rank calibrates on its own data and folds the AWQ scale it picks into its own weight @@ -81,9 +85,13 @@ def _refuse_awq_on_a_sharded_fsdp2_model(model: nn.Module, algorithm: QuantizeAl methods = [mode.removesuffix("_calibrate") for mode, _ in get_modelike_from_algo_cfg(algorithm)] awq = [method for method in methods if method in ("awq_lite", "awq_full")] if awq: + # For the message only: the FSDP2 shard mesh can be narrower than the run (HSDP, 2D). + mesh = _get_fsdp2_mesh(next(m for m in model.modules() if isinstance(m, FSDPModule))) + ranks = mesh.size(-1) if mesh is not None else dist.size() + what = f"AutoQuantize candidate {candidate!r} calibrates" if candidate else "Calibrating" raise NotImplementedError( - f"Calibrating with {', '.join(awq)} is not supported yet on an FSDP2 model sharded over " - f"{dist.size()} ranks: each rank calibrates on its own data and folds the AWQ scale it " + f"{what} with {', '.join(awq)}, which is not supported yet on an FSDP2 model sharded " + f"over {ranks} ranks: each rank calibrates on its own data and folds the AWQ scale it " "picks into its own weight shard, so the shards would not match. Calibrate in a single " "process, or with an algorithm other than AWQ." ) @@ -1017,6 +1025,8 @@ def forward_backward_step(model, batch) -> None: "Please run auto_quantize with 'max' or 'mse' calibration and use " "get_auto_quantize_config() to obtain a config for mtq.quantize()." ) + # Every candidate, before the search calibrates any: calibrate() would refuse mid-search. + _refuse_awq_on_a_sharded_fsdp2_model(model, algo, candidate=name) # Select the appropriate searcher based on method if method_name not in AUTO_QUANTIZE_SEARCHERS: diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index 101337d9547..d4271942037 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -741,6 +741,16 @@ def _awq_on_a_sharded_fsdp2_model_job(rank, size): for config in (mtq.W4A8_AWQ_BETA_CFG, INT4_AWQ_FULL_CFG): # awq_lite, awq_full with pytest.raises(NotImplementedError, match=f"FSDP2 model sharded over {size} ranks"): mtq.quantize(model, config, forward_loop) + # AutoQuantize checks every candidate before it calibrates the first one. + with pytest.raises(NotImplementedError, match="AutoQuantize candidate 'W4A8_AWQ_BETA_CFG'"): + mtq.auto_quantize( + model, + constraints={"effective_bits": 6.0}, + quantization_formats=[mtq.FP8_DEFAULT_CFG, mtq.W4A8_AWQ_BETA_CFG], + data_loader=[SimpleLinear.get_input()], + forward_step=lambda m, batch: m(batch), + loss_func=lambda output, batch: output.sum(), + ) assert not is_quantized(model) mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop) with pytest.raises(NotImplementedError, match="awq_lite"): @@ -749,7 +759,7 @@ def _awq_on_a_sharded_fsdp2_model_job(rank, size): def test_awq_on_a_sharded_fsdp2_model_is_refused(cpu_dist_workers): """Each rank would fold its own AWQ scale into its shard: every rank refuses alike, before the - model changes, through ``mtq.quantize`` and ``mtq.calibrate``.""" + model changes, through ``mtq.quantize``, ``mtq.auto_quantize`` and ``mtq.calibrate``.""" cpu_dist_workers.run(_awq_on_a_sharded_fsdp2_model_job) From 060836ae59124027d09391b032207a5f271ce82f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 18:50:16 +0000 Subject: [PATCH 8/9] Refuse multi-rank FSDP2 AWQ in hf_ptq up front, and narrow the GPU tests hf_ptq --use_fsdp2 on more than one rank with an AWQ format or recipe, AutoQuantize candidates included, now fails at argument validation, before the distributed load, instead of at calibration minutes later. The check mirrors the library's: awq_lite or awq_full. The README's --use_fsdp2 notes say so. The two GPU tests adapted for the refusal now change only what it refuses, AWQ on more than one rank: test_fsdp2_export.py shards before calibrating again for every other case, and test_nested_fsdp2_backward runs its quantize-then-shard comparison on several GPUs, gating only the shard-then-quantize model. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 ++ examples/hf_ptq/example_utils.py | 29 +++++++++++++++++++ tests/examples/hf_ptq/test_example_utils.py | 23 +++++++++++++++ tests/gpu/torch/export/test_fsdp2_export.py | 20 ++++++++++--- tests/gpu/torch/quantization/test_fsdp2.py | 31 +++++++++++---------- 5 files changed, 87 insertions(+), 18 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 4f282610713..20ebc34b109 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -614,6 +614,8 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *Nemotron-H checkpoints containing MTP weights are not supported with `--use_fsdp2` and are rejected before distributed model loading. Remove `--use_fsdp2` to load and calibrate the MTP modules; the distributed loader does not construct them.* +> *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Run it on one GPU, or use a non-AWQ format.* + ### Usage #### Slurm (recommended) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index c537e553c78..e9fa1fddbe7 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -44,6 +44,8 @@ ProcessorMixin, ) +from modelopt.recipe import load_recipe +from modelopt.recipe.presets import QUANT_CFG_CHOICES from modelopt.torch.export import has_spec_opt from modelopt.torch.export.model_utils import is_multimodal_model from modelopt.torch.models import hf_model_type @@ -52,6 +54,7 @@ prefers_bundled_code, prepare_model_for_loading, ) +from modelopt.torch.quantization.mode import get_modelike_from_algo_cfg from modelopt.torch.utils import distributed as dist_utils from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run from modelopt.torch.utils.plugins.model_load_utils import is_fsdp2_text_only_vlm, load_pretrained @@ -94,6 +97,27 @@ def cleanup_distributed(args): dist_utils.cleanup() +def _awq_calibrations(args) -> set[str]: + """The AWQ calibrations (``awq_lite``, ``awq_full``) the run's --recipe or --qformat uses.""" + if getattr(args, "recipe", None) is not None: + recipe = load_recipe(args.recipe) + configs = [getattr(recipe, "quantize", None)] + search = getattr(recipe, "auto_quantize", None) # the KV-cache search allows only max + if search is not None: + configs += search.candidate_formats + configs += [c for space in search.module_search_spaces for c in space.candidate_formats] + algorithms = [config.algorithm for config in configs if config is not None] + else: + preset = QUANT_CFG_CHOICES.get(getattr(args, "qformat", None) or "") + algorithms = [preset.get("algorithm")] if preset is not None else [] + methods = { + mode.removesuffix("_calibrate") + for algorithm in algorithms + for mode, _ in get_modelike_from_algo_cfg(algorithm) + } + return methods & {"awq_lite", "awq_full"} + + def validate_fsdp2_supported(args, config): """Raise ``NotImplementedError`` for model/CLI combos the FSDP2 path doesn't support yet.""" issues = [] @@ -113,6 +137,11 @@ def validate_fsdp2_supported(args, config): issues.append("speculative decoding (--specdec_offline_dataset)") if getattr(args, "low_memory_mode", False): issues.append("--low_memory_mode (redundant with FSDP2)") + if dist_utils.size() > 1 and (awq := _awq_calibrations(args)): + issues.append( + f"AWQ calibration ({', '.join(sorted(awq))}) on more than one rank: each rank would fold " + "its own smoothing scale into its weight shard (use one GPU, or a non-AWQ format)" + ) if not issues and checkpoint_has_mtp(config.model_type, args.local_checkpoint_path): issues.append("Nemotron-H MTP (auxiliary modules are not constructed by the FSDP2 loader)") diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index caaa625805b..c2435e96973 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -58,6 +58,29 @@ def test_fsdp2_qwen_vlm_requires_text_only_calibration( example_utils.validate_fsdp2_supported(args, config) +@pytest.mark.parametrize( + ("ranks", "qformat", "recipe", "refused"), + [ + (2, "w4a8_awq_beta", None, True), + (2, "fp8", "general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits", True), + (2, "fp8", None, False), + (1, "w4a8_awq_beta", None, False), + ], +) +def test_fsdp2_refuses_awq_on_several_ranks_before_loading( + monkeypatch, ranks, qformat, recipe, refused +): + """Each rank would fold its own AWQ scale into its weight shard, which the library refuses + only once the model is loaded: the example refuses at argument validation.""" + monkeypatch.setattr(example_utils.dist_utils, "size", lambda: ranks) + args = SimpleNamespace(local_checkpoint_path="local-model", qformat=qformat, recipe=recipe) + expected = ( + pytest.raises(NotImplementedError, match="AWQ calibration") if refused else nullcontext() + ) + with expected: + example_utils.validate_fsdp2_supported(args, SimpleNamespace(model_type="llama")) + + def _write_safetensors(path, tensors): save_file(tensors, str(path), metadata={"format": "pt"}) diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 392c8359246..d438ed83ad1 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -177,10 +177,16 @@ def _fuse_layers(rank, size, quant_config, bias): def calib_fn(x): return x(calib_data) - # Quantize, then shard: AWQ calibration of a model sharded over several ranks is refused. + # AWQ calibration of a model sharded over several ranks is refused: there, shard after it. + awq_on_ranks = size > 1 and "awq" in str(quant_config.get("algorithm")) + if not awq_on_ranks: + fully_shard(model) + torch.distributed.barrier() + mtq.quantize(model, quant_config, calib_fn) mtq.quantize(non_fsdp_model, quant_config, calib_fn) - fully_shard(model) + if awq_on_ranks: + fully_shard(model) torch.distributed.barrier() @@ -214,10 +220,16 @@ def _export_quantized_weight_test(rank, size, quant_config, bias): def calib_fn(x): return x(calib_data) - # Quantize, then shard: AWQ calibration of a model sharded over several ranks is refused. + # AWQ calibration of a model sharded over several ranks is refused: there, shard after it. + awq_on_ranks = size > 1 and "awq" in str(quant_config.get("algorithm")) + if not awq_on_ranks: + fully_shard(model) + torch.distributed.barrier() + mtq.quantize(model, quant_config, calib_fn) mtq.quantize(non_fsdp_model, quant_config, calib_fn) - fully_shard(model) + if awq_on_ranks: + fully_shard(model) torch.distributed.barrier() diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index 4f1dda81a14..aec8e37f813 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -102,16 +102,20 @@ def forward_loop(model): out_test = fsdp_model(inputs) out_test.sum().backward() - fully_shard(fsdp_model_quant_after[0]) - fully_shard(fsdp_model_quant_after[1]) - fsdp_model_quant_after = fully_shard(fsdp_model_quant_after) - fsdp_model_quant_after = mtq.quantize(fsdp_model_quant_after, quant_cfg, forward_loop) - optimizer_quant_after = torch.optim.SGD(fsdp_model_quant_after.parameters(), lr=0.1) - out_quant_after = fsdp_model_quant_after(inputs) - out_quant_after.sum().backward() + # AWQ calibration of a model sharded over several ranks is refused: there, quantize first only. + quant_after = size == 1 or "awq" not in str(quant_cfg.get("algorithm")) + if quant_after: + fully_shard(fsdp_model_quant_after[0]) + fully_shard(fsdp_model_quant_after[1]) + fsdp_model_quant_after = fully_shard(fsdp_model_quant_after) + fsdp_model_quant_after = mtq.quantize(fsdp_model_quant_after, quant_cfg, forward_loop) + optimizer_quant_after = torch.optim.SGD(fsdp_model_quant_after.parameters(), lr=0.1) + out_quant_after = fsdp_model_quant_after(inputs) + out_quant_after.sum().backward() assert torch.allclose(out_ref, out_test) - assert torch.allclose(out_ref, out_quant_after) + if quant_after: + assert torch.allclose(out_ref, out_quant_after) optimizer_ref.step() optimizer_ref.zero_grad() @@ -119,14 +123,15 @@ def forward_loop(model): optimizer_test.step() optimizer_test.zero_grad() - optimizer_quant_after.step() - optimizer_quant_after.zero_grad() + if quant_after: + optimizer_quant_after.step() + optimizer_quant_after.zero_grad() out_ref_1 = model(inputss) out_test_1 = fsdp_model(inputss) - out_quant_after_1 = fsdp_model_quant_after(inputss) assert torch.allclose(out_ref_1, out_test_1, rtol=1e-4) - assert torch.allclose(out_ref_1, out_quant_after_1, rtol=1e-4) + if quant_after: + assert torch.allclose(out_ref_1, fsdp_model_quant_after(inputss), rtol=1e-4) def test_fsdp_simple_linear(dist_workers): @@ -137,8 +142,6 @@ def test_fsdp_simple_linear(dist_workers): "quant_cfg", [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_SMOOTHQUANT_CFG, mtq.INT4_AWQ_CFG] ) def test_nested_fsdp2_backward(quant_cfg, dist_workers): - if quant_cfg is mtq.INT4_AWQ_CFG and torch.cuda.device_count() > 1: - pytest.skip("AWQ calibration of an FSDP2 model sharded over several ranks is refused") dist_workers.run(partial(_test_nested_fsdp2_backward, quant_cfg=quant_cfg)) From 75db98ee54a82970db10eded1bb2471df24fde37 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:14:30 +0000 Subject: [PATCH 9/9] Split the changelog entry and name the remedy that keeps the GPUs The one Bug Fixes bullet carried three independent changes; it is now three entries: the FSDP2 W4A8 export fix, the multi-rank FSDP2 AWQ refusal, and the compress refusal. The hf_ptq README note and error message now offer dropping --use_fsdp2 first, as the single-process device_map still spreads the model over all the GPUs, then one rank or a non-AWQ format. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 4 +++- examples/hf_ptq/README.md | 2 +- examples/hf_ptq/example_utils.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ff82cd101dd..638ecd6ed3a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -98,7 +98,9 @@ Changelog - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``. - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization. - Fix weight calibration, GPTQ and weight folding of Hugging Face fused MoE experts sharded by the expert parallelism of Transformers 5.16 and later, which used another rank's expert weights on every rank but the first. -- ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``) no longer fails with ``AttributeError: 'SequentialQuantizer' object has no attribute '_fake_quant'``; AWQ calibration of an FSDP2 model sharded over more than one rank, whose ranks would fold different smoothing scales into their weight shards, is rejected with ``NotImplementedError`` until it is supported, as is ``mtq.auto_quantize`` with an AWQ candidate, such as its default ``NVFP4_AWQ_LITE_CFG``, before the search starts. ``mtq.compress`` now refuses W4A8 weights with ``NotImplementedError``, since a compressed W4A8 model could neither run nor be reloaded. +- Fix ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``), which failed with an ``AttributeError`` on every rank. +- AWQ calibration (``awq_lite``, ``awq_full``) of an FSDP2 model sharded over more than one rank is now rejected with ``NotImplementedError`` by ``mtq.quantize``, ``mtq.calibrate`` and ``mtq.auto_quantize`` (before its search starts, its default formats included), and by ``examples/hf_ptq --use_fsdp2`` before loading the model: each rank would fold a different smoothing scale into its weight shard. Calibrate in a single process, or use a non-AWQ algorithm. +- ``mtq.compress`` now rejects W4A8 weights with ``NotImplementedError``: a compressed W4A8 model could neither run nor be reloaded. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 20ebc34b109..d108dcbf786 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -614,7 +614,7 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *Nemotron-H checkpoints containing MTP weights are not supported with `--use_fsdp2` and are rejected before distributed model loading. Remove `--use_fsdp2` to load and calibrate the MTP modules; the distributed loader does not construct them.* -> *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Run it on one GPU, or use a non-AWQ format.* +> *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Drop `--use_fsdp2` (the single-process `device_map` still spreads the model over all the GPUs), run on one rank, or use a non-AWQ format.* ### Usage diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index e9fa1fddbe7..cca908f4553 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -140,7 +140,8 @@ def validate_fsdp2_supported(args, config): if dist_utils.size() > 1 and (awq := _awq_calibrations(args)): issues.append( f"AWQ calibration ({', '.join(sorted(awq))}) on more than one rank: each rank would fold " - "its own smoothing scale into its weight shard (use one GPU, or a non-AWQ format)" + "its own smoothing scale into its weight shard (drop --use_fsdp2, run on one rank, or " + "use a non-AWQ format)" ) if not issues and checkpoint_has_mtp(config.model_type, args.local_checkpoint_path): issues.append("Nemotron-H MTP (auxiliary modules are not constructed by the FSDP2 loader)")