misc

Author: fegin

torchtitan/distributed/parallel_dims.py

@@ -128,7 +128,7 @@ def unflatten_mesh(
             ("pp", "batch", "cp", "tp"),
             (self.pp, batch, self.cp, self.tp),
         )
-        loss_mesh = dataloading_mesh["batch", "cp"].flatten("loss_mesh")
+        loss_mesh = dataloading_mesh["batch", "cp"]._flatten("loss_mesh")
         dense_mesh = unflatten_mesh(
             self._world_mesh,
             ("pp", "dp_replicate", "fsdp", "tp"),
@@ -143,7 +143,7 @@ def unflatten_mesh(
         self._meshes = {
             "pp": dataloading_mesh["pp"],
             "batch": dataloading_mesh["batch"],
-            "loss": loss_mesh["loss"],
+            "loss": loss_mesh,
             "dp_replicate": dense_mesh["dp_replicate"],
             "fsdp": dense_mesh["fsdp"],
             "cp": dataloading_mesh["cp"],

torchtitan/distributed/utils.py

@@ -191,6 +191,8 @@ def set_determinism(
     # As long as we are not in the 1-D (PP-only) case, we will have a seed to use for all ranks of the SPMD mesh.
     # IF PP is also used, this seed is unique per PP rank.
     if duplicate_seed_meshes:
+        # We only need to pass one submesh as DTensor manual_seed uses that to verify
+        # if the current rank is actually in the submesh.
         torch.distributed.tensor._random.manual_seed(seed, duplicate_seed_meshes[0])