Update on "[DO NOT MERGE][example] fold batch and sequence dimensions to accelerate Sequence Parallel"

Note: This PR is for showcasing purpose only and is almost a reverse of #190.

At the cost of model code change, we can obtain better Sequence Parallel performance. Without folding and unfolding,  all-gather and reduce-scatter are performed on dim 1 (sequence dim) instead of dim 0 (folded dim), which incurs an extra `aten.cat` after each collective.

Stats from awgu:
> for 8k seq len, batch size 1 on H100, these two cats take about 0.18 ms out of 3 ms of FFN compute (6%)

Experiment on 8-layer `debug_model`
before:
<img width="1023" alt="image" src="https://github.com/pytorch/torchtitan/assets/150487191/04e5ea4b-fa9e-48e5-92be-582841cb2796">
after:
<img width="1023" alt="image" src="https://github.com/pytorch/torchtitan/assets/150487191/38c39506-462d-485a-a16c-48770a28edb0">

[ghstack-poisoned]

Author: tianyu-l

torchtitan/models/llama/model.py

@@ -427,15 +427,10 @@ def forward(self, tokens: torch.Tensor):
             torch.Tensor: Output logits after applying the Transformer model.
 
         """
-        bs = tokens.shape[0]
         # passthrough for nonexistent layers, allows easy configuration of pipeline parallel stage
+        h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens
         # fold batch dimension and sequence dimension for more efficient allgather/reduce_scatter
-        if self.tok_embeddings:
-            tokens = tokens.view(-1)
-            h = self.tok_embeddings(tokens)
-        else:
-            h = tokens
-            h = h.view(-1, self.model_args.dim)
+        h = h.view(-1, self.model_args.dim)
 
         seqlen = self.model_args.max_seq_len
         freqs_cis = self.freqs_cis[0:seqlen]
@@ -444,6 +439,7 @@ def forward(self, tokens: torch.Tensor):
 
         h = self.norm(h) if self.norm else h
         # unfold batch and sequence dimension
+        bs = tokens.shape[0]
         h = h.view(bs, -1, self.model_args.dim)
         output = self.output(h).float() if self.output else h
         return output

torchtitan/parallelisms/parallelize_llama.py

@@ -350,14 +350,18 @@ def apply_tp(model, world_mesh, parallel_dims, job_config: JobConfig):
         {
             "tok_embeddings": RowwiseParallel(
                 input_layouts=Replicate(),
-                output_layouts=Shard(0),
             ),
             "output": col_parallel_strategy(
                 input_layouts=Shard(0),
                 output_layouts=Shard(-1) if loss_parallel else Replicate(),
                 use_local_output=not loss_parallel,
             ),
             "norm": SequenceParallel(sequence_dim=0),
+            "layers.0": PrepareModuleInput(
+                input_layouts=(Replicate(), None),
+                desired_input_layouts=(Shard(0), None),
+                use_local_output=True,
+            ),
         },
     )