From 8606f41736243921b0d5e1bf3b8688ca930e66f5 Mon Sep 17 00:00:00 2001 From: babyplutokurt Date: Mon, 10 Aug 2026 11:27:25 -0400 Subject: [PATCH] Skip weight quantizers with no stored weight in fold_weight On Megatron-Core models with tied word embeddings, the output_layer is built with skip_weight_param_allocation: it stores weight = None and borrows the embedding weight at forward time, while ModelOpt still attaches a weight_quantizer to it. fold_weight matched the pair on the *_weight_quantizer attribute name and fake_quant alone, then dereferenced weight.data, crashing with: AttributeError: 'NoneType' object has no attribute 'data' Skip pairs whose weight attribute is not a tensor and leave their quantizer untouched: there is nothing stored to fold, and the shared weight keeps being quantized at forward time through the still-enabled quantizer. The embedding module folds its own stored weight as before. HF models are unaffected; they express tying as a shared tensor rather than None, which is why the crash only surfaced on Megatron. Fixes #2131 Signed-off-by: babyplutokurt --- CHANGELOG.rst | 1 + .../quantization/nn/modules/quant_module.py | 12 ++++++++++++ .../quantization/test_tensor_quant_cpu.py | 19 +++++++++++++++++++ 3 files changed, 32 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 78f8426e5f9..b92b6133ea1 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,6 +25,7 @@ Changelog **Bug Fixes** +- Fix ``mtq.fold_weight`` crashing with ``AttributeError: 'NoneType' object has no attribute 'data'`` on Megatron-Core models with tied word embeddings: the tied ``output_layer`` (built with ``skip_weight_param_allocation``) stores ``weight = None`` and borrows the embedding weight at forward time, yet still carries a ``weight_quantizer``. Weight-quantizer pairs whose weight is not a stored tensor are now skipped and left untouched. - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. 0.46 (2026-08-17) diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 9c9aee478a8..20e5fd12a31 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -157,6 +157,14 @@ def fold_weight(self, keep_attrs: bool = False): transform is baked into the stored weight and then disabled, so subsequent forwards use the stored weight directly. Calibration buffers (``_pre_quant_scale``, ``_amax``) are dropped unless ``keep_attrs``. + + Quantizers whose weight attribute is not a tensor are skipped and left untouched: + modules such as Megatron tied-embedding output layers (built with + ``skip_weight_param_allocation``) store ``weight = None`` and receive the shared + weight as a forward-time argument, so there is nothing stored to fold. Such quantizers + stay enabled after folding, so callers that assert all weight quantizers are disabled + once folding completes (e.g. ``_check_all_weight_quantizers_disabled`` in the vLLM + fakequant export plugin) must account for them. """ # Handle all attributes that end with _weight_quantizer for name in dir(self): @@ -173,6 +181,10 @@ def fold_weight(self, keep_attrs: bool = False): f"{name} doesn't have a corresponding {weight_name} in {self.__class__.__name__}" ) weight = getattr(self, weight_name) + if not isinstance(weight, torch.Tensor): + # e.g. Megatron tied-embedding output_layer: weight is None and + # borrowed at forward time, so there is nothing stored to fold. + continue self._fold_weight_quantizer(attr, (weight,), keep_attrs) diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index ba352ec2162..0346ba4c336 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -332,6 +332,25 @@ def test_fold_weight_keep_attrs_keeps_amax(monkeypatch): unregister_quant_backend(backend_name) +def test_fold_weight_skips_none_weight(): + """A weight quantizer with no stored weight is skipped instead of crashing. + + Megatron-Core tied-embedding output layers are built with + ``skip_weight_param_allocation``: the module stores ``weight = None`` and borrows the + embedding weight at forward time, while still carrying a ``weight_quantizer``. + ``fold_weight`` must skip the pair, leaving the quantizer intact for forward-time use. + """ + qlinear = QuantModuleRegistry.convert(torch.nn.Linear(4, 3)) + qlinear.weight_quantizer.amax = torch.tensor(1.0) + expected_amax = qlinear.weight_quantizer.amax.detach().clone() + qlinear.register_parameter("weight", None) + + qlinear.fold_weight() # must not raise + + assert qlinear.weight_quantizer.is_enabled + assert torch.equal(qlinear.weight_quantizer.amax, expected_amax) + + WINT4INT8_CFG = { "quant_cfg": [ {"quantizer_name": "*", "enable": False},