[Tensor Parallel] update examples to simplify embedding + first transformer block

tianyu-l · tianyu-l · commit 8c2c2f073063 · 2024-05-15T17:57:49.000-07:00
ghstack-source-id: 49c0b6f Pull Request resolved: #1259
diff --git a/distributed/tensor_parallelism/fsdp_tp_example.py b/distributed/tensor_parallelism/fsdp_tp_example.py
@@ -107,17 +107,13 @@
     {
         "tok_embeddings": RowwiseParallel(
             input_layouts=Replicate(),
+            output_layouts=Shard(1),
         ),
         "output": ColwiseParallel(
             input_layouts=Shard(1),
             output_layouts=Replicate()
         ),
         "norm": SequenceParallel(),
-        "layers.0": PrepareModuleInput(
-            input_layouts=(Replicate(), None),
-            desired_input_layouts=(Shard(1), None),
-            use_local_output=True,
-        ),
     }
 )
 

Original file line number	Diff line number	Diff line change
`@@ -107,17 +107,13 @@`
`107`	`107`	`{`
`108`	`108`	`"tok_embeddings": RowwiseParallel(`
`109`	`109`	`input_layouts=Replicate(),`
	`110`	`+ output_layouts=Shard(1),`
`110`	`111`	`),`
`111`	`112`	`"output": ColwiseParallel(`
`112`	`113`	`input_layouts=Shard(1),`
`113`	`114`	`output_layouts=Replicate()`
`114`	`115`	`),`
`115`	`116`	`"norm": SequenceParallel(),`
`116`		`- "layers.0": PrepareModuleInput(`
`117`		`- input_layouts=(Replicate(), None),`
`118`		`- desired_input_layouts=(Shard(1), None),`
`119`		`- use_local_output=True,`
`120`		`- ),`
`121`	`117`	`}`
`122`	`118`	`)`
`123`	`119`