Enable TP+PP support

wconstab · wconstab · commit 05f08025239d · 2024-05-01T22:21:04.000-07:00
ghstack-source-id: 8177304 Pull Request resolved: #285
diff --git a/torchtitan/parallelisms/parallelize_llama.py b/torchtitan/parallelisms/parallelize_llama.py
@@ -205,7 +205,8 @@ def parallelize_llama(model, world_mesh, parallel_dims, job_config: JobConfig):
         )
 
         # Apply tensor + sequence parallelism to every transformer block
-        for layer_id, transformer_block in enumerate(model.layers):
+        for layer_name, transformer_block in model.layers.named_children():
+            # for layer_id, transformer_block in enumerate(model.layers):
             layer_plan = {
                 "attention": PrepareModuleInput(
                     input_layouts=(Shard(1), None),

Original file line number	Diff line number	Diff line change
`@@ -205,7 +205,8 @@ def parallelize_llama(model, world_mesh, parallel_dims, job_config: JobConfig):`
`205`	`205`	`)`
`206`	`206`
`207`	`207`	`# Apply tensor + sequence parallelism to every transformer block`
`208`		`- for layer_id, transformer_block in enumerate(model.layers):`
	`208`	`+ for layer_name, transformer_block in model.layers.named_children():`
	`209`	`+ # for layer_id, transformer_block in enumerate(model.layers):`
`209`	`210`	`layer_plan = {`
`210`	`211`	`"attention": PrepareModuleInput(`
`211`	`212`	`input_layouts=(Shard(1), None),`