[OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them

CHANGELOG.rst

@@ -98,6 +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.
+- 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)
 ^^^^^^^^^^^^^^^^^^^

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. 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
 
 #### Slurm (recommended)

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,12 @@ 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 (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)")
 

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,16 @@ def compress_convert(
     Returns:
         The compressed model.
     """
+    # 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 any(
+            isinstance(q, SequentialQuantizer) for q in module.weight_quantizer.modules()
+        ):
+            raise NotImplementedError(
+                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:
             class_to_register = RealQuantLinear

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,6 +72,31 @@
 ]
 
 
+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
+    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:
+        # 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"{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."
+        )
+
+
 # TODO: Descriptors for the supported algorithms
 def calibrate(
     model: nn.Module,
@@ -99,6 +126,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 +467,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:
@@ -996,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:

modelopt/torch/quantization/utils/core_utils.py

@@ -1005,17 +1005,26 @@ 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
-    yield
-    for m in module.modules():
-        if hasattr(m, "weight_quantizer"):
-            m.weight_quantizer._fake_quant = original_fake_quant.pop(0)
+    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 isinstance(q, TensorQuantizer)
+    ]
+    original_fake_quant = [q._fake_quant for q in quantizers]
+    for q in quantizers:
+        q._fake_quant = True
+    try:
+        yield
+    finally:
+        for q, fake_quant in zip(quantizers, original_fake_quant):
+            q._fake_quant = fake_quant
 
 
 @contextmanager

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

tests/gpu/torch/export/test_fsdp2_export.py

@@ -177,13 +177,16 @@ def _fuse_layers(rank, size, quant_config, bias):
         def calib_fn(x):
             return x(calib_data)
 
-        # Shard model
-        fully_shard(model)
-        torch.distributed.barrier()
+        # 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()
 
-        # Quantize model
         mtq.quantize(model, quant_config, calib_fn)
         mtq.quantize(non_fsdp_model, quant_config, calib_fn)
+        if awq_on_ranks:
+            fully_shard(model)
 
         torch.distributed.barrier()
 
@@ -217,13 +220,16 @@ 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()
+        # 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()
 
-        # Quantize model
         mtq.quantize(model, quant_config, calib_fn)
         mtq.quantize(non_fsdp_model, quant_config, calib_fn)
+        if awq_on_ranks:
+            fully_shard(model)
 
         torch.distributed.barrier()
 
@@ -305,7 +311,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 +333,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

tests/gpu/torch/quantization/test_fsdp2.py

@@ -102,31 +102,36 @@ 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()
 
     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):

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 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. 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.
+    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)
+
+    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