[FT] Change the integration from using ManagedDeviceMesh to set_all_r… (meta-pytorch#1109)

Fixes pytorch/torchtitan#1105

While using `ManagedDeviceMesh` makes the integration code cleaner,
`ManagedDeviceMesh` currently suffers from the composability issue with
TP due to the limitations of `DeviceMesh`.

This PR changes the integration to using `FSDP.set_all_reduce_hook()`.
We will revisit the `ManagedDeviceMesh` once `DeviceMesh` becomes more
friendly to the inheritance use cases.

Author: fegin

torchtitan/components/ft.py

@@ -10,8 +10,11 @@
 from typing import Optional
 
 import torch
+import torch.distributed as dist
 import torch.distributed._functional_collectives as funcol
+from torch.distributed._composable.fsdp.fully_shard import FSDPModule
 from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.distributed_c10d import ReduceOp
 from torch.distributed.tensor import DTensor
 from torchtitan.config_manager import JobConfig
 from torchtitan.distributed import ParallelDims
@@ -34,6 +37,9 @@ def __init__(
         self._manager = manager
         self.group_size = group_size
         self.replica_id = replica_id
+        if has_torchft and manager is not None:
+            self.replicate_pg = ft.process_group.ManagedProcessGroup(self._manager)
+            self.replicate_pg.register("dp_replicate")
 
     @property
     def enabled(self) -> bool:
@@ -47,6 +53,17 @@ def manager(self) -> "ft.Manager":
     def get_dp_info(self, dp_degree: int, dp_rank: int) -> tuple[int, int]:
         return dp_degree * self.group_size, dp_degree * self.replica_id + dp_rank
 
+    def set_all_reduce_hook(self, model_parts: list[torch.nn.Module]) -> None:
+        def all_reduce_hook(output):
+            dist.all_reduce(output, group=self.replicate_pg, op=ReduceOp.AVG)
+
+        def apply_set_all_reduce_hook(m):
+            if isinstance(m, FSDPModule):
+                m.set_all_reduce_hook(all_reduce_hook)
+
+        for part in model_parts:
+            part.apply(apply_set_all_reduce_hook)
+
 
 def init_ft_manager(job: JobConfig) -> FTManager:
     """Initialize the FT manager if TorchFT is enabled.
@@ -55,7 +72,7 @@ def init_ft_manager(job: JobConfig) -> FTManager:
         job (JobConfig): The job configuration.
 
     Returns:
-        Optional[ft.Manager]: The FT manager if TorchFT is enabled, otherwise None.
+        FTManager: A wrapper around TorchFT.Manager
     """
     if not job.fault_tolerance.enable:
         return FTManager(None)
@@ -66,7 +83,7 @@ def init_ft_manager(job: JobConfig) -> FTManager:
     if job.fault_tolerance.min_replica_size < 1:
         raise ValueError("At least one FT replica is required.")
 
-    pg = ft.ProcessGroupBabyNCCL()
+    pg = ft.ProcessGroupNCCL()
 
     return FTManager(
         ft.Manager(

torchtitan/distributed/utils.py

@@ -17,28 +17,51 @@
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import DTensor
 
-from torchtitan.components.ft import ft_clip_grad_norm_util, ft_dist_reduce
 from torchtitan.tools.logging import logger
 from torchtitan.tools.utils import device_module, device_type
 
 
-def _dist_reduce(x: torch.Tensor, reduceOp: str, mesh: DeviceMesh) -> float:
-    # Remove FT replicate dimension if it exists.
-    x, reduceOp, mesh = ft_dist_reduce(x, reduceOp, mesh)
+def _dist_reduce(
+    x: torch.Tensor,
+    reduceOp: str,
+    mesh: DeviceMesh,
+    extra_pg: dist.ProcessGroup | None = None,
+) -> float:
+    """Perform distributed reduction on a tensor.
 
+    Args:
+        x (torch.Tensor): Input tensor.
+        reduceOp (str): Reduce operation to perform.
+        mesh (DeviceMesh): Device mesh to use for reduction.
+        extra_pg (dist.ProcessGroup, optional): Extra process group to use for reduction.
+            Defaults to None. If provided, this all_reduce will be called for the extra
+            process group, and then the result will be all_reduced for the mesh.
+    """
     if isinstance(x, DTensor):
         # functional collectives do not support DTensor inputs
         x = x.full_tensor()
+
+    if extra_pg is not None:
+        x = funcol.all_reduce(x, reduceOp=reduceOp, group=extra_pg)
+
     assert x.numel() == 1  # required by `.item()`
     return funcol.all_reduce(x, reduceOp=reduceOp, group=mesh).item()
 
 
-def dist_max(x: torch.Tensor, mesh: DeviceMesh) -> float:
-    return _dist_reduce(x, reduceOp=c10d.ReduceOp.MAX.name, mesh=mesh)
+def dist_max(
+    x: torch.Tensor, mesh: DeviceMesh, extra_pg: dist.ProcessGroup | None
+) -> float:
+    return _dist_reduce(
+        x, reduceOp=c10d.ReduceOp.MAX.name, mesh=mesh, extra_pg=extra_pg
+    )
 
 
-def dist_mean(x: torch.Tensor, mesh: DeviceMesh) -> float:
-    return _dist_reduce(x, reduceOp=c10d.ReduceOp.AVG.name, mesh=mesh)
+def dist_mean(
+    x: torch.Tensor, mesh: DeviceMesh, extra_pg: dist.ProcessGroup | None
+) -> float:
+    return _dist_reduce(
+        x, reduceOp=c10d.ReduceOp.AVG.name, mesh=mesh, extra_pg=extra_pg
+    )
 
 
 def set_determinism(
@@ -301,8 +324,6 @@ def clip_grad_norm_(
         # Will reach here if any non-PP parallelism is used.
         # If only using PP, total_norm will be a local tensor.
 
-        # Remove FT replicate dimension if it exists.
-        total_norm = ft_clip_grad_norm_util(total_norm)
         total_norm = total_norm.full_tensor()
 
     if pp_mesh is not None:

torchtitan/train.py

@@ -78,32 +78,19 @@ def __init__(self, job_config: JobConfig):
         self.device = torch.device(f"{device_type}:{int(os.environ['LOCAL_RANK'])}")
         # Device has to be set before creating TorchFT manager.
         device_module.set_device(self.device)
-        ft_manager = ft.init_ft_manager(job_config)
 
         # init distributed
         world_size = int(os.environ["WORLD_SIZE"])
         parallelism_config = job_config.parallelism
-        if not ft_manager.enabled:
-            self.parallel_dims = parallel_dims = ParallelDims(
-                dp_shard=parallelism_config.data_parallel_shard_degree,
-                dp_replicate=parallelism_config.data_parallel_replicate_degree,
-                cp=parallelism_config.context_parallel_degree,
-                tp=parallelism_config.tensor_parallel_degree,
-                pp=parallelism_config.pipeline_parallel_degree,
-                world_size=world_size,
-                enable_loss_parallel=not parallelism_config.disable_loss_parallel,
-            )
-        else:
-            self.parallel_dims = parallel_dims = ft.FTParallelDims(
-                dp_shard=parallelism_config.data_parallel_shard_degree,
-                dp_replicate=parallelism_config.data_parallel_replicate_degree,
-                cp=parallelism_config.context_parallel_degree,
-                tp=parallelism_config.tensor_parallel_degree,
-                pp=parallelism_config.pipeline_parallel_degree,
-                world_size=world_size,
-                enable_loss_parallel=not parallelism_config.disable_loss_parallel,
-                ft_manager=ft_manager,
-            )
+        self.parallel_dims = parallel_dims = ParallelDims(
+            dp_shard=parallelism_config.data_parallel_shard_degree,
+            dp_replicate=parallelism_config.data_parallel_replicate_degree,
+            cp=parallelism_config.context_parallel_degree,
+            tp=parallelism_config.tensor_parallel_degree,
+            pp=parallelism_config.pipeline_parallel_degree,
+            world_size=world_size,
+            enable_loss_parallel=not parallelism_config.disable_loss_parallel,
+        )
         dist_utils.init_distributed(job_config)
 
         # build meshes
@@ -114,6 +101,12 @@ def __init__(self, job_config: JobConfig):
         else:
             dp_degree, dp_rank = 1, 0
 
+        self.ft_manager = ft.init_ft_manager(job_config)
+        # If TorchFT is enabled, the dp_rank and dp_degree, which are used for
+        # dataloader must be changed.
+        if self.ft_manager.enabled:
+            dp_degree, dp_rank = self.ft_manager.get_dp_info(dp_degree, dp_rank)
+
         # Set random seed, and maybe enable deterministic mode
         # (mainly for debugging, expect perf loss).
         dist_utils.set_determinism(
@@ -131,11 +124,6 @@ def __init__(self, job_config: JobConfig):
             else None
         )
 
-        # If TorchFT is enabled, the dp_rank and dp_degree, which are used for
-        # dataloader must be changed.
-        if ft_manager.enabled:
-            dp_degree, dp_rank = ft_manager.get_dp_info(dp_degree, dp_rank)
-
         self.dataloader = self.train_spec.build_dataloader_fn(
             dp_world_size=dp_degree,
             dp_rank=dp_rank,
@@ -241,6 +229,9 @@ def __init__(self, job_config: JobConfig):
 
             self.model_parts = [model]
 
+        if self.ft_manager.enabled:
+            self.ft_manager.set_all_reduce_hook(self.model_parts)
+
         # initialize device memory monitor and get peak flops for MFU calculation
         device_memory_monitor = self.metrics_processor.device_memory_monitor
         gpu_peak_flops = utils.get_peak_flops(device_memory_monitor.device_name)
@@ -254,7 +245,7 @@ def __init__(self, job_config: JobConfig):
 
         # build optimizer after applying parallelisms to the model
         self.optimizers = self.train_spec.build_optimizers_fn(
-            self.model_parts, job_config, ft_manager
+            self.model_parts, job_config, self.ft_manager
         )
         self.lr_schedulers = self.train_spec.build_lr_schedulers_fn(
             self.optimizers, job_config
@@ -280,7 +271,7 @@ def __init__(self, job_config: JobConfig):
             lr_schedulers=self.lr_schedulers,
             states={"train_state": self},
             job_config=job_config,
-            ft_manager=ft_manager,
+            ft_manager=self.ft_manager,
         )
 
         self.train_context = dist_utils.get_train_context(
@@ -384,11 +375,13 @@ def train_step(self, input_dict: dict[str, torch.Tensor], labels: torch.Tensor):
             parallel_dims.dp_replicate_enabled
             or parallel_dims.dp_shard_enabled
             or parallel_dims.cp_enabled
+            or self.ft_manager.enabled
         ):
             loss = loss.detach()
+            ft_pg = self.ft_manager.replicate_pg if self.ft_manager.enabled else None
             global_avg_loss, global_max_loss = (
-                dist_utils.dist_mean(loss, world_mesh["dp_cp"]),
-                dist_utils.dist_max(loss, world_mesh["dp_cp"]),
+                dist_utils.dist_mean(loss, world_mesh["dp_cp"], ft_pg),
+                dist_utils.dist_max(loss, world_mesh["dp_cp"], ft_pg),
             )
         else:
             global_avg_loss = global_max_loss = loss.detach().item()