Update on "Enable CP"

This PR adds experimental flags and functions to enable context parallelism. We currently support on    ly FSDP + CP and CP only. CP + TP is being tested.



[ghstack-poisoned]

Author: fegin

torchtitan/models/llama/model.py

@@ -410,7 +410,7 @@ def _precompute_freqs_cis(self) -> torch.Tensor:
             self.model_args.dim // self.model_args.n_heads,
             # Need to compute until at least the max token limit for generation
             # (use 2x max sequence length to be safe)
-            self.model_args.max_seq_len * 2,
+            self.model_args.max_seq_len,
             self.model_args.rope_theta,
         )
 

torchtitan/parallelisms/__init__.py

@@ -61,8 +61,6 @@ def build_mesh(self, device_type):
         logger.info(f"Building {len(dims)}-D device mesh with {names}, {dims}")
         names = tuple(names)
         world_mesh = init_device_mesh(device_type, dims, mesh_dim_names=names)
-        if self.cp > 1 and self.dp > 1:
-            world_mesh.create_view_dim(dims=("dp", "cp"), name="dp_cp")
         return world_mesh
 
     @property

torchtitan/parallelisms/parallelize_llama.py

@@ -506,7 +506,7 @@ def apply_fsdp(
     if parallel_dims.cp_enabled:
         # Temporary solution to enable FSDP + CP
         if parallel_dims.dp_enabled:
-            dp_mesh = world_mesh["dp_cp"]
+            dp_mesh = world_mesh["dp", "cp"]._flatten()
         else:
             dp_mesh = world_mesh["cp"]
     else: