ADLR/megatron-lm!2358 - Add repr for parallel linear layers

megatron/core/extensions/transformer_engine.py

@@ -377,6 +377,12 @@ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
             state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
         )
 
+    def __repr__(self):
+        return (
+            f"{type(self).__name__}(in_features={self.in_features}, "
+            f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
+        )
+
 
 class TEColumnParallelLinear(TELinear):
     """
@@ -451,6 +457,12 @@ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
             state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
         )
 
+    def __repr__(self):
+        return (
+            f"{type(self).__name__}(in_features={self.in_features}, "
+            f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
+        )
+
 
 class TERowParallelLinear(TELinear):
     """
@@ -525,6 +537,12 @@ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
             state_dict, prefix, {'weight': 1}, sharded_offsets
         )
 
+    def __repr__(self):
+        return (
+            f"{type(self).__name__}(in_features={self.in_features}, "
+            f"out_features={self.out_features}, bias={self.use_bias}, TP={self.tp_size})"
+        )
+
 
 class TEDotProductAttention(te.pytorch.DotProductAttention):
     """

megatron/core/tensor_parallel/layers.py

@@ -989,6 +989,14 @@ def get_extra_state(self) -> None:
         """Keep compatibility with TE state dict."""
         return None
 
+    def __repr__(self):
+        tp = self.output_size // self.output_size_per_partition
+        use_bias = self.bias is not None and self.bias is True
+        return (
+            f"{type(self).__name__}(in_features={self.input_size}, "
+            f"out_features={self.output_size}, bias={use_bias}, TP={tp})"
+        )
+
 
 class RowParallelLinear(torch.nn.Module):
     """Linear layer with row parallelism.
@@ -1213,3 +1221,11 @@ def set_extra_state(self, state: Any):
     def get_extra_state(self) -> None:
         """Keep compatibility with TE state dict."""
         return None
+
+    def __repr__(self):
+        tp = self.input_size // self.input_size_per_partition
+        use_bias = self.bias is not None and self.bias is True
+        return (
+            f"{type(self).__name__}(in_features={self.input_size}, "
+            f"out_features={self.output_size}, bias={use_bias}, TP={tp})"
+        )