fix simplefsdp gradient_divide_factor

Author: ruisizhang123

torchtitan/experiments/simple_fsdp/deepseek_v3/parallelize.py

@@ -125,18 +125,19 @@ def parallelize_deepseekv3(
                 ):
                     experts_shard_dim = 1
 
+                # when EP is enable, the shared experts' gradient reduce is done over
+                # dp_mod_ep_mesh instead of whole dp_mesh.
+                # we add a `fsdp_gradient_divide_factor` to scale gradient over dp_mesh
+                # to be consistent with data.
                 transformer_block.moe.experts = data_parallel(
                     transformer_block.moe.experts,
                     dp_mod_ep_mesh,
                     dp_mode,
                     ac_mode=job_config.activation_checkpoint.mode,
                     mp_policy=mp_policy,
                     shard_dim=experts_shard_dim,
+                    reduction_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,69 @@ class MixedPrecisionPolicy:
     reduce_dtype: Optional[torch.dtype] = None
 
 
+class _ScaledPartial(Partial):
+    # A subclass of Partial placement that allows user to perform gradient reduction with a custom
+    # factor (reduction_divide_factor) other than the default world size.
+    def __init__(
+        self,
+        reduction_divide_factor: float,
+        reduce_dtype: Optional[torch.dtype] = None,
+    ):
+        self.reduction_divide_factor = reduction_divide_factor
+        self.reduce_dtype = reduce_dtype
+        super().__init__(reduce_op="sum")
+
+    def _get_reduction_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
+        """
+        assert self.reduce_dtype in (
+            torch.float32,
+            torch.bfloat16,
+        ), "only support reduce_dtype to be fp32/bf16"
+        pre_factor, post_factor = None, self.reduction_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_reduction_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_reduction_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 +256,25 @@ def __init__(
         mode,
         regional_ac,
         mp_policy,
+        reduction_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 = [
+            _ScaledPartial(
+                reduction_divide_factor=reduction_divide_factor,
+                reduce_dtype=mp_policy.reduce_dtype,
+            )
+            if reduction_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 +357,7 @@ def data_parallel(
     ac_mode: str = "none",
     mp_policy: Optional[MixedPrecisionPolicy] = None,
     shard_dim: int = 0,
+    reduction_divide_factor: Optional[float] = None,
 ):
     if mode == "replicate":
         param_sharding = (Replicate(),)
@@ -348,6 +420,7 @@ def data_parallel(
                 mode,
                 regional_ac,
                 mp_policy=mp_policy,
+                reduction_divide_factor=reduction_divide_factor,
             ),
         )
     return model