minor fix

Signed-off-by: larkz <larkz@nvidia.com>

Author: jQizhang

modelopt/torch/quantization/model_calib.py

@@ -72,13 +72,6 @@ def weight_only_quantize(model: nn.Module):
             with enable_weight_access_and_writeback(module, model, name_to_module):
                 for weight, weight_quantizer in module.iter_weights_for_calibration():
                     weight_quantizer(weight)
-        else:
-            for weight_name in weight_attr_names(module):
-                with enable_weight_access_and_writeback(module, model, name_to_module):
-                    weight_quantizer = getattr(
-                        module, quantizer_attr_names(weight_name).weight_quantizer
-                    )
-                    weight_quantizer(getattr(module, weight_name))
         seen_modules.add(module)
 
 

modelopt/torch/quantization/nn/modules/quant_module.py

@@ -120,14 +120,7 @@ def modelopt_post_restore(self, prefix: str = ""):
                 module.to(non_tq_param_or_buffer.device)
 
     def iter_weights_for_calibration(self):
-        """Yield ``(weight, weight_quantizer)`` pairs for weight-only calibration.
-
-        The default implementation iterates over all weights returned by
-        :func:`~modelopt.torch.quantization.utils.weight_attr_names`.  Subclasses that
-        store weights under non-standard attribute names (e.g.
-        ``_QuantTEGroupedLinear`` uses ``weight0``, ``weight1``, …) should
-        override this method.
-        """
+        """Yield ``(weight, weight_quantizer)`` pairs for weight-only calibration."""
         from modelopt.torch.quantization.utils import quantizer_attr_names, weight_attr_names
 
         for weight_name in weight_attr_names(self):

modelopt/torch/quantization/plugins/transformer_engine.py

@@ -152,12 +152,7 @@ def modelopt_post_restore(self, prefix: str = ""):
         delattr(self, "weight")
 
     def iter_weights_for_calibration(self):
-        """Yield ``(weight_i, weight_quantizer)`` for each of the ``num_gemms`` grouped weights.
-
-        Override is needed because ``self.weight`` is removed in ``_setup``, so the
-        base-class implementation (which relies on ``weight_attr_names``) would find
-        no weights.  Here we iterate over ``weight0``, ``weight1``, … directly.
-        """
+        """Yield ``(weight_i, weight_quantizer)`` for each of the ``num_gemms`` grouped weights."""
         for i in range(self.num_gemms):
             weight_i = getattr(self, f"weight{i}", None)
             if weight_i is not None: