From c15ec0d6c43ee561560358958461a82c094462af Mon Sep 17 00:00:00 2001 From: Chenghao Liu Date: Thu, 10 Sep 2026 00:53:31 +0800 Subject: [PATCH 1/3] Fix folding transposed GPT-OSS and Llama 4 expert weights Signed-off-by: Chenghao Liu --- CHANGELOG.rst | 1 + .../quantization/nn/modules/quant_module.py | 4 +- .../torch/quantization/plugins/huggingface.py | 8 +++ .../quantization/plugins/test_huggingface.py | 59 +++++++++++++++++++ 4 files changed, 71 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 335678eac5a..e96d9c94cf5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,7 @@ Changelog **Bug Fixes** +- Fix ``fold_weight`` failing on quantized GPT-OSS and Llama 4 models with transposed expert weights. - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet. diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 522f2d0db09..0771702bc44 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -145,6 +145,8 @@ def _fold_weight_quantizer( quantizer: TensorQuantizer, weights: Iterable[torch.Tensor], keep_attrs: bool = False, + *, + quantize_dtype: torch.dtype = torch.float32, ): """Fold ``quantizer`` into each weight view in place, then disable and clean it once. @@ -156,7 +158,7 @@ def _fold_weight_quantizer( return for weight in weights: - weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype)) + weight.data.copy_(quantizer(weight.to(quantize_dtype).contiguous()).to(weight.dtype)) quantizer.disable() quantizer.disable_rotate() if keep_attrs and hasattr(quantizer, "_pre_quant_scale"): diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 2515910ec6e..f08ceb0bdd6 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -620,6 +620,14 @@ def iter_weights_for_calibration(self): weight = getattr(self, weight_name) yield weight.transpose(-1, -2), getattr(self, f"{weight_name}_weight_quantizer") + def fold_weight(self, keep_attrs: bool = False): + """Fold expert weights in the transposed orientation used by their forward.""" + for weight, quantizer in self.iter_weights_for_calibration(): + if isinstance(quantizer, TensorQuantizer): + QuantModule._fold_weight_quantizer( + quantizer, (weight,), keep_attrs, quantize_dtype=weight.dtype + ) + class _QuantSparseSequentialMoe(QuantModule): """Quantization wrapper for HuggingFace sparse MoE blocks. diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 0da6cfbb1fc..37671f23a74 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -382,6 +382,65 @@ def __init__(self): assert down_q is experts.down_proj_weight_quantizer +@pytest.mark.parametrize("model_type", ["gpt_oss", "llama4"]) +@pytest.mark.parametrize("keep_attrs", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_fold_transposed_expert_weights_preserves_forward(model_type, keep_attrs, dtype): + if model_type == "gpt_oss": + model = get_tiny_gpt_oss(num_hidden_layers=1, hidden_size=32, intermediate_size=48) + else: + model = AutoModelForCausalLM.from_config( + transformers.Llama4TextConfig( + num_hidden_layers=1, + num_local_experts=4, + hidden_size=32, + intermediate_size=48, + intermediate_size_mlp=48, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=16, + vocab_size=32, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + no_rope_layers=[0], + ) + ) + model = model.to(dtype=dtype).eval() + input_ids = torch.randint(0, 32, (2, 7)) + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*.experts.*weight_quantizer", + "cfg": {"num_bits": 8, "axis": (0, 1)}, + }, + ], + "algorithm": "max", + } + mtq.quantize(model, config, lambda m: m(input_ids)) + experts = next( + module for module in model.modules() if isinstance(module, _TransposedExpertsCalibMixin) + ) + pairs = list(experts.iter_weights_for_calibration()) + for weight, quantizer in pairs: + quantizer.pre_quant_scale = torch.linspace(0.75, 1.25, weight.shape[-1], dtype=dtype) + pointers = [weight.data_ptr() for weight, _ in pairs] + + with torch.no_grad(): + expected = model(input_ids).logits + mtq.fold_weight(model, keep_attrs=keep_attrs) + torch.testing.assert_close(model(input_ids).logits, expected) + for (weight, quantizer), pointer in zip(pairs, pointers): + assert weight.data_ptr() == pointer + assert not quantizer.is_enabled + assert quantizer.pre_quant_scale is None + assert hasattr(quantizer, "_amax") == keep_attrs + assert hasattr(quantizer, "_pre_quant_scale") == keep_attrs + mtq.fold_weight(model, keep_attrs=keep_attrs) + torch.testing.assert_close(model(input_ids).logits, expected) + + def test_hf_decoder_discoverer_registration_path(): model = get_tiny_llama() assert any( From 4b96b515bc1ef1dd4973e9679c74583f0911c60a Mon Sep 17 00:00:00 2001 From: Chenghao Liu Date: Thu, 10 Sep 2026 22:55:30 +0800 Subject: [PATCH 2/3] fix: materialize offloaded weights while folding quantizers Signed-off-by: Chenghao Liu --- modelopt/torch/quantization/model_quant.py | 6 +++-- .../torch/quantization/plugins/accelerate.py | 3 ++- .../quantization/plugins/test_huggingface.py | 22 +++++++++++++++++-- 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 3f6040fd4ef..e32604d609d 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -43,7 +43,7 @@ from .config import QuantizeAlgoCfgType from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg from .nn import QuantModule, SequentialQuantizer, TensorQuantizer -from .utils import is_quantized +from .utils import enable_weight_access_and_writeback, is_quantized, module_name_maps __all__ = [ "auto_quantize", @@ -731,9 +731,11 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False): Any weight-quantizer rotation is folded into the weights and disabled so subsequent forwards do not re-rotate the already-folded weights. """ + names = module_name_maps(model) for name, module in model.named_modules(): if isinstance(module, QuantModule): - module.fold_weight(keep_attrs) + with enable_weight_access_and_writeback(module, model, names): + module.fold_weight(keep_attrs) @contextmanager diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py index 156669bfc8f..377afe6d067 100644 --- a/modelopt/torch/quantization/plugins/accelerate.py +++ b/modelopt/torch/quantization/plugins/accelerate.py @@ -50,7 +50,8 @@ def _writeback_params_to_weights_map(module, align_hook): continue if isinstance(align_hook.weights_map, PrefixedDataset): key = align_hook.weights_map.prefix + name - w_map = align_hook.weights_map.dataset.state_dict + dataset = align_hook.weights_map.dataset + w_map = getattr(dataset, "state_dict", dataset) else: w_map = align_hook.weights_map key = name diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 37671f23a74..ef6c20456eb 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -385,7 +385,11 @@ def __init__(self): @pytest.mark.parametrize("model_type", ["gpt_oss", "llama4"]) @pytest.mark.parametrize("keep_attrs", [False, True]) @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -def test_fold_transposed_expert_weights_preserves_forward(model_type, keep_attrs, dtype): +@pytest.mark.parametrize("offload", [None, "cpu", "disk"]) +def test_fold_transposed_expert_weights_preserves_forward( + model_type, keep_attrs, dtype, offload, tmp_path +): + set_seed() if model_type == "gpt_oss": model = get_tiny_gpt_oss(num_hidden_layers=1, hidden_size=32, intermediate_size=48) else: @@ -429,16 +433,30 @@ def test_fold_transposed_expert_weights_preserves_forward(model_type, keep_attrs with torch.no_grad(): expected = model(input_ids).logits + if offload is not None: + accelerate = pytest.importorskip("accelerate") + if offload == "cpu": + accelerate.cpu_offload(experts, execution_device=torch.device("cpu")) + else: + accelerate.disk_offload( + experts, tmp_path / "offload", execution_device=torch.device("cpu") + ) + assert experts.gate_up_proj.is_meta + torch.testing.assert_close(model(input_ids).logits, expected) mtq.fold_weight(model, keep_attrs=keep_attrs) torch.testing.assert_close(model(input_ids).logits, expected) for (weight, quantizer), pointer in zip(pairs, pointers): - assert weight.data_ptr() == pointer + if offload is None: + assert weight.data_ptr() == pointer assert not quantizer.is_enabled assert quantizer.pre_quant_scale is None assert hasattr(quantizer, "_amax") == keep_attrs assert hasattr(quantizer, "_pre_quant_scale") == keep_attrs mtq.fold_weight(model, keep_attrs=keep_attrs) torch.testing.assert_close(model(input_ids).logits, expected) + if offload is not None: + assert experts.gate_up_proj.is_meta + assert experts.down_proj.is_meta def test_hf_decoder_discoverer_registration_path(): From 424c3e2b332c419dc309edad9b6023cd69381010 Mon Sep 17 00:00:00 2001 From: Chenghao Liu Date: Fri, 11 Sep 2026 16:16:53 +0800 Subject: [PATCH 3/3] test: explain optional Accelerate import Signed-off-by: Chenghao Liu --- tests/unit/torch/quantization/plugins/test_huggingface.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index ef6c20456eb..21427de414f 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -434,6 +434,7 @@ def test_fold_transposed_expert_weights_preserves_forward( with torch.no_grad(): expected = model(input_ids).logits if offload is not None: + # Accelerate is optional; only offload cases require it. accelerate = pytest.importorskip("accelerate") if offload == "cpu": accelerate.cpu_offload(experts, execution_device=torch.device("cpu"))