[dtensor][debug] tutorial showing users how to use commdebugmode and giving access to visual browser

Author: sinhaanshul

distributed_comm_debug_mode.rst

@@ -11,7 +11,7 @@ Prerequisites:
 
 
 What is CommDebugMode and why is it useful
-------------------
+------------------------------------------
 As the size of models continues to increase, users are seeking to leverage various combinations of parallel strategies to scale up distributed training. However, the lack of interoperability between existing solutions poses a significant challenge, primarily due to the absence of a unified abstraction that can bridge these different parallelism strategies. To address this issue, PyTorch has proposed DistributedTensor (DTensor)which abstracts away the complexities of tensor communication in distributed training, providing a seamless user experience. However, this abstraction creates a lack of transparency that can make it challenging for users to identify and resolve issues. To address this challenge, my internship project aims to develop and enhance CommDebugMode, a Python context manager that will serve as one of the primary debugging tools for DTensors. CommDebugMode is a python context manager that enables users to view when and why collective operations are happening when using DTensors, addressing this problem.
 
 
@@ -26,7 +26,7 @@ Using CommDebugMode and getting its output is very simple.
             output = model(inp)
 
     # print the operation level collective tracing information
-    print(comm_mode.generate_comm_debug_tracing_table(noise_level=2))
+    print(comm_mode.generate_comm_debug_tracing_table(noise_level=0))
 
     # log the operation level collective tracing information to a file
     comm_mode.log_comm_debug_tracing_table_to_file(
@@ -36,16 +36,35 @@ Using CommDebugMode and getting its output is very simple.
     # dump the operation level collective tracing information to json file,
     # used in the visual browser below
     comm_mode.generate_json_dump(noise_level=2)
-.. code-block:: python
 
-All users have to do is wrap the code running the model in CommDebugMode and call the API that they want to use to display the data.
-Documentation Title
-===================
+    """
+    This is what the output looks like for a MLPModule at noise level 0
+    Expected Output:
+        Global
+          FORWARD PASS
+            *c10d_functional.all_reduce: 1
+            MLPModule
+              FORWARD PASS
+                *c10d_functional.all_reduce: 1
+                MLPModule.net1
+                MLPModule.relu
+                MLPModule.net2
+                  FORWARD PASS
+                    *c10d_functional.all_reduce: 1
+    """
+
+All users have to do is wrap the code running the model in CommDebugMode and call the API that they want to use to display the data. One important thing to note
+is that the users can use a noise_level arguement to control how much information is displayed to the user. You can see what each noise_level will display to the user.
+
+| 0. prints module-level collective counts
+| 1. prints dTensor operations not included in trivial operations, module information
+| 2. prints operations not included in trivial operations
+| 3. prints all operations
 
-Introduction to the Module
---------------------------
+In the example above, users can see in the first picture that the collective operation, all_reduce, occurs once in the forward pass of the MLPModule. The second picture provides a greater level of detail, allowing users to pinpoint that the all-reduce operation happens in the second linear layer of the MLPModule.
 
-Below is the interactive module tree visualization:
+
+Below is the interactive module tree visualization that users can upload their JSON dump to:
 
 .. raw:: html
 
@@ -155,5 +174,3 @@ Below is the interactive module tree visualization:
         <script src="https://cdn.jsdelivr.net/gh/pytorch/pytorch@main/torch/distributed/_tensor/debug/comm_mode_broswer_visual.js"></script>
     </body>
     </html>
-
-.. raw:: html