Comments

pytorch · Dec 28, 2023 · f3daad1 · f3daad1
1 parent e55eef4
commit f3daad1
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/examples/llama/2d_llama.py b/examples/llama/2d_llama.py
@@ -6,12 +6,15 @@
 from torch.distributed._tensor import init_device_mesh
 
 
+# Utility
 def modify_view(
     gm: torch.fx.GraphModule,
     tp: int
 ):
     """
-    Adjust dimension size of view ops to make them compatible with tensor parallelism.
+    Adjust dimension size of view ops to make them compatible with tensor
+    parallelism.  For example, when TP is 4, we need to adjust `num_heads` from
+    32 to 8.  This is needed for attention layers.
     """
     for node in gm.graph.nodes:
         if node.op == "call_method" and (