Addressed MR comments

Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>

Author: kinjalpatel27

modelopt/torch/quantization/model_calib.py

@@ -118,7 +118,7 @@ def sync_quantizer_amax_across_tp(
         # Syncing amax across TP for sequential quantizer
         if isinstance(quantizer, SequentialQuantizer):
             for _q in quantizer:
-                "Syncing amax across TP for sequential quantizer"
+                # Syncing amax across TP for sequential quantizer
                 sync_quantizer_amax_across_tp(
                     _q, linear_name, quantizer_type, axes_for_sync, parallel_state
                 )

modelopt/torch/quantization/plugins/transformer_engine.py

@@ -73,6 +73,7 @@ def _setup(self):
         # GroupedMLP stores the weights as weight0, weight1, etc. To run setup in order to
         # initialize the quantizer states, self.weight is used to extract shape, dtype etc. Assigning
         # self.weight0 to self.weight to run the quantizer states initialization.
+        assert not hasattr(self, "weight"), "self.weight should not exist for TEGroupedLinear"
         self.weight = self.weight0
         # Memorize the original weight.dtype for modelopt_post_restore given that
         # the dtype can change later.
@@ -84,6 +85,7 @@ def modelopt_post_restore(self, prefix: str = ""):
         # GroupedMLP stores the weights as weight0, weight1, etc. To run post_restore in order to
         # initialize the quantizer states, self.weight is used to extract shape, dtype etc. Assigning
         # self.weight0 to self.weight to run the quantizer states initialization.
+        assert not hasattr(self, "weight"), "self.weight should not exist for TEGroupedLinear"
         self.weight = self.weight0
         super().modelopt_post_restore(prefix=prefix)
         # Remove self.weight after post_restore.

modelopt/torch/quantization/utils.py

@@ -253,6 +253,7 @@ def is_quantized_linear(module):
         and hasattr(module, "weight_quantizer")
         and (
             (getattr(module, "weight", None) is not None and module.weight.dim() == 2)
+            # module.weight0 check is required to support TEGroupedLinear
             or (getattr(module, "weight0", None) is not None and module.weight0.dim() == 2)
         )
     )