[Distributed] Add repr methods for ParallelStyles (pytorch#149478)

Fixes pytorch#149470

Pull Request resolved: pytorch#149478
Approved by: https://github.com/wanchaol

Author: shink

torch/distributed/tensor/parallel/style.py

@@ -166,6 +166,14 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
             ),
         )
 
+    def __repr__(self) -> str:
+        tmpstr = self.__class__.__name__ + "("
+        tmpstr += f"input_layouts={self.input_layouts}, "
+        tmpstr += f"output_layouts={self.output_layouts}, "
+        tmpstr += f"use_local_output={self.use_local_output}"
+        tmpstr += ")"
+        return tmpstr
+
 
 class RowwiseParallel(ParallelStyle):
     """
@@ -303,6 +311,14 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
             ),
         )
 
+    def __repr__(self) -> str:
+        tmpstr = self.__class__.__name__ + "("
+        tmpstr += f"input_layouts={self.input_layouts}, "
+        tmpstr += f"output_layouts={self.output_layouts}, "
+        tmpstr += f"use_local_output={self.use_local_output}"
+        tmpstr += ")"
+        return tmpstr
+
 
 class SequenceParallel(ParallelStyle):
     """
@@ -398,6 +414,14 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
             partial(self._prepare_output_fn, self.use_local_output),
         )
 
+    def __repr__(self) -> str:
+        tmpstr = self.__class__.__name__ + "("
+        if len(self.sequence_sharding) == 1:
+            tmpstr += f"sequence_dim={self.sequence_sharding[0].dim}, "
+        tmpstr += f"use_local_output={self.use_local_output}"
+        tmpstr += ")"
+        return tmpstr
+
 
 class PrepareModuleInput(ParallelStyle):
     """
@@ -557,6 +581,16 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
             )  # type: ignore[misc, call-arg]
         return module
 
+    def __repr__(self) -> str:
+        tmpstr = self.__class__.__name__ + "("
+        tmpstr += f"input_layouts={self.input_layouts}, "
+        tmpstr += f"desired_input_layouts={self.desired_input_layouts}, "
+        tmpstr += f"input_kwarg_layouts={self.input_kwarg_layouts}, "
+        tmpstr += f"desired_input_kwarg_layouts={self.desired_input_kwarg_layouts}, "
+        tmpstr += f"use_local_output={self.use_local_output}"
+        tmpstr += ")"
+        return tmpstr
+
 
 class PrepareModuleOutput(ParallelStyle):
     """
@@ -656,3 +690,11 @@ def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
             lambda _, inputs, outputs: self._prepare_out_fn(outputs, device_mesh)
         )  # type: ignore[misc, call-arg]
         return module
+
+    def __repr__(self) -> str:
+        tmpstr = self.__class__.__name__ + "("
+        tmpstr += f"output_layouts={self.output_layouts}, "
+        tmpstr += f"desired_output_layouts={self.desired_output_layouts}, "
+        tmpstr += f"use_local_output={self.use_local_output}"
+        tmpstr += ")"
+        return tmpstr