diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4ea103a0628..638ecd6ed3a 100755 --- a/CHANGELOG.rst +++ b/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) ^^^^^^^^^^^^^^^^^^^ diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 4f282610713..d108dcbf786 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. 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) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index c537e553c78..cca908f4553 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,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)") diff --git a/modelopt/torch/quantization/compress.py b/modelopt/torch/quantization/compress.py index ee7e642a0e8..e1a30e4c079 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,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 diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 115dcf437ee..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,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: diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 8bb3ac69288..ddf2f4d91d9 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/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 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 ce71b1c4270..d438ed83ad1 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/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 diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index c9e56d0e4fc..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): diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index c80f704bb6d..d5bb39d016d 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 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 diff --git a/tests/unit/torch/quantization/test_compress.py b/tests/unit/torch/quantization/test_compress.py new file mode 100644 index 00000000000..affd6536f2e --- /dev/null +++ b/tests/unit/torch/quantization/test_compress.py @@ -0,0 +1,38 @@ +# 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 +from modelopt.torch.quantization.nn import GroupedQuantizer + + +@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) + + assert not is_real_quantized(model) diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index f51a514e0e5..d4271942037 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,36 @@ 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) + # 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"): + 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``, ``mtq.auto_quantize`` and ``mtq.calibrate``.""" + cpu_dist_workers.run(_awq_on_a_sharded_fsdp2_model_job) + + @pytest.mark.parametrize( ("quantizer_name", "supported"), [ 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)