fix simplefsdp gradient_divide_factor

Author: ruisizhang123

torchtitan/experiments/simple_fsdp/deepseek_v3/parallelize.py

@@ -132,11 +132,8 @@ def parallelize_deepseekv3(
                     ac_mode=job_config.activation_checkpoint.mode,
                     mp_policy=mp_policy,
                     shard_dim=experts_shard_dim,
+                    gradient_divide_factor=parallel_dims.fsdp_gradient_divide_factor,
                 )
-                # TODO(ruisizhang123): support set_gradient_divide_factor in simplefsdp
-                # transformer_block.moe.experts.set_gradient_divide_factor(
-                #     parallel_dims.fsdp_gradient_divide_factor,
-                # )
 
         model = data_parallel(
             model,

torchtitan/experiments/simple_fsdp/simple_fsdp.py

@@ -20,6 +20,7 @@
     Shard,
 )
 from torch.distributed.device_mesh import _mesh_resources, DeviceMesh
+from torch.distributed.fsdp._fully_shard._fsdp_collectives import _div_if_needed
 from torch.distributed.tensor._dtensor_spec import DTensorSpec
 from torch.distributed.tensor._redistribute import redistribute_local_tensor
 from torch.distributed.tensor.placement_types import _StridedShard, Placement
@@ -49,6 +50,84 @@ class MixedPrecisionPolicy:
     reduce_dtype: Optional[torch.dtype] = None
 
 
+class SimpleFSDPPartial(Partial):
+    # A subclass of Partial placement that allows user to perform gradient reduction with a custom
+    # factor (gradient_divide_factor) other than the default world size.
+    def __init__(
+        self,
+        reduce_op: str,
+        gradient_divide_factor: float,
+        reduce_dtype: Optional[torch.dtype] = None,
+    ):
+        self.gradient_divide_factor = gradient_divide_factor
+        self.reduce_dtype = reduce_dtype
+        super().__init__(reduce_op)
+
+    def _get_gradient_divide_factors(
+        self,
+    ) -> tuple[Optional[float], Optional[float]]:
+        """
+        the logic follows
+        https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py#L688
+        """
+        overflow_risk = self.reduce_dtype not in (torch.float32, torch.bfloat16)
+
+        pre_factor: Optional[float]
+        if overflow_risk:
+            # Since fp16 has smaller dynamic range than fp32/bf16, we want to avoid
+            # overflow/underflow. For N data parallel workers, each worker computes
+            # g_i, and they collectively reduce (g_1 + ... + g_N) / N. To avoid
+            # overflow/underflow, we divide by ~sqrt(N) before/after the reduction.
+            pre_factor = 1
+            while (
+                self.gradient_divide_factor % pre_factor == 0
+                and self.gradient_divide_factor / pre_factor > pre_factor
+            ):
+                pre_factor *= 2
+            post_factor = self.gradient_divide_factor / pre_factor
+        else:
+            # Prefer post-multiplying as it operates on less data and is thus faster
+            pre_factor, post_factor = None, self.gradient_divide_factor
+
+        return pre_factor, post_factor
+
+    def _reduce_value(
+        self, tensor: torch.Tensor, mesh: DeviceMesh, mesh_dim: int
+    ) -> torch.Tensor:
+        # for all_reduce in DDP
+        (
+            pre_factor,
+            post_factor,
+        ) = self._get_gradient_divide_factors()
+        if pre_factor is not None:
+            _div_if_needed(tensor, pre_factor)
+        reduced = super()._reduce_value(tensor, mesh, mesh_dim)
+        if post_factor is not None:
+            _div_if_needed(reduced, post_factor)
+        return reduced
+
+    def _reduce_shard_value(
+        self,
+        tensor: torch.Tensor,
+        mesh: DeviceMesh,
+        mesh_dim: int,
+        shard_spec: Placement,
+    ) -> torch.Tensor:
+        # for reduce_scatter in FSDP
+        (
+            pre_factor,
+            post_factor,
+        ) = self._get_gradient_divide_factors()
+
+        if pre_factor is not None:
+            _div_if_needed(tensor, pre_factor)
+        reduced = super()._reduce_shard_value(tensor, mesh, mesh_dim, shard_spec)
+
+        if post_factor is not None:
+            _div_if_needed(reduced, post_factor)
+        return reduced
+
+
 def _distribute_dtensor(
     tensor: DTensor,
     device_mesh: DeviceMesh,
@@ -192,18 +271,26 @@ def __init__(
         mode,
         regional_ac,
         mp_policy,
+        gradient_divide_factor,
     ):
         super().__init__()
         self.device_mesh = device_mesh
         self.param_sharding = param_sharding
         self.mode = mode
         self.compute_placements = [Replicate()] * self.device_mesh.ndim
-        self.grad_placements = [Partial(reduce_op="avg")] * self.device_mesh.ndim
+        self.grad_placements = [
+            SimpleFSDPPartial(
+                reduce_op="sum",
+                gradient_divide_factor=gradient_divide_factor,
+                reduce_dtype=mp_policy.reduce_dtype,
+            )
+            if gradient_divide_factor is not None
+            else Partial(reduce_op="avg")
+        ] * self.device_mesh.ndim
         self.regional_ac = regional_ac
         mp_policy = mp_policy or MixedPrecisionPolicy()
         self.param_dtype = mp_policy.param_dtype
         self.reduce_dtype = mp_policy.reduce_dtype
-        self.ep_mesh_name, self.tp_mesh_name = "ep", "tp"
 
     def replicate_compute(self, x):
         # data parallel runtime replicate parameters and do local compute
@@ -286,6 +373,7 @@ def data_parallel(
     ac_mode: str = "none",
     mp_policy: Optional[MixedPrecisionPolicy] = None,
     shard_dim: int = 0,
+    gradient_divide_factor: Optional[float] = None,
 ):
     if mode == "replicate":
         param_sharding = (Replicate(),)
@@ -348,6 +436,7 @@ def data_parallel(
                 mode,
                 regional_ac,
                 mp_policy=mp_policy,
+                gradient_divide_factor=gradient_divide_factor,
             ),
         )
     return model