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

@@ -7,9 +7,10 @@
 from collections.abc import Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import cast, List, Optional
 
 import torch
+import torch.distributed._functional_collectives as funcol
 import torch.nn as nn
 
 from torch.distributed._tensor import (
@@ -20,6 +21,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 +51,101 @@ class MixedPrecisionPolicy:
     reduce_dtype: Optional[torch.dtype] = None
 
 
+@dataclass(frozen=True)
+class SimpleFSDPPartial(Partial):
+    gradient_divide_factor: Optional[float] = None
+    reduce_dtype: Optional[torch.dtype] = None
+    data_parallel_size: Optional[int] = None
+    force_sum_reduction_for_comms: bool = False
+
+    def _get_gradient_divide_factors(
+        self,
+    ) -> tuple[Optional[float], Optional[float], str, str,]:
+        """
+        the logic follows
+        https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py#L688
+        """
+        if self.gradient_divide_factor is None:
+            return None, None, None, None
+
+        overflow_risk = self.reduce_dtype not in (torch.float32, torch.bfloat16)
+
+        if not overflow_risk and not self.force_sum_reduction_for_comms:
+            if self.gradient_divide_factor == self.data_parallel_size:
+                if self.data_parallel_size == 1:
+                    return None, None, "SUM", "SUM"
+                return None, None, "AVG", "AVG"
+            else:
+                reduce_scatter_op = torch.distributed._make_nccl_premul_sum(
+                    1 / self.gradient_divide_factor
+                )
+                return None, None, reduce_scatter_op, "SUM"
+
+        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, "SUM", "SUM"
+
+    def _reduce_value(
+        self, tensor: torch.Tensor, mesh: DeviceMesh, mesh_dim: int
+    ) -> torch.Tensor:
+        # for all_reduce in DDP
+        (
+            pre_factor,
+            post_factor,
+            _,
+            all_reduce_op,
+        ) = self._get_gradient_divide_factors()
+        if pre_factor is not None:
+            _div_if_needed(tensor, pre_factor)
+        tensor = funcol.all_reduce(
+            tensor, reduceOp=all_reduce_op, group=(mesh, mesh_dim)
+        )
+        if post_factor is not None:
+            _div_if_needed(tensor, post_factor)
+        return tensor
+
+    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,
+            reduce_scatter_op,
+            _,
+        ) = self._get_gradient_divide_factors()
+
+        if pre_factor is not None:
+            _div_if_needed(tensor, pre_factor)
+        shard_spec = cast(Shard, shard_spec)
+        tensor = shard_spec._reduce_shard_tensor(
+            tensor, mesh, reduce_scatter_op, mesh_dim
+        )
+
+        if post_factor is not None:
+            _div_if_needed(tensor, post_factor)
+        return tensor
+
+
 def _distribute_dtensor(
     tensor: DTensor,
     device_mesh: DeviceMesh,
@@ -192,18 +289,27 @@ 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="avg",
+                gradient_divide_factor=gradient_divide_factor,
+                reduce_dtype=mp_policy.reduce_dtype,
+                data_parallel_size=self.device_mesh.size(),
+            )
+            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 +392,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 +455,7 @@ def data_parallel(
                 mode,
                 regional_ac,
                 mp_policy=mp_policy,
+                gradient_divide_factor=gradient_divide_factor,
             ),
         )
     return model

torchtitan/models/deepseek_v3/__init__.py

@@ -51,6 +51,8 @@
         qk_rope_head_dim=64,
         v_head_dim=128,
         mscale=0.70,
+        use_flex_attn=True,
+        attn_mask_type="block_causal",
     ),
     "debugmodel_flex_attn": DeepSeekV3ModelArgs(
         vocab_size=2048,

torchtitan/models/llama3/__init__.py

@@ -29,7 +29,11 @@
 
 llama3_configs = {
     "debugmodel": TransformerModelArgs(
-        dim=256, n_layers=6, n_heads=16, vocab_size=2048, rope_theta=500000
+        dim=256,
+        n_layers=6,
+        n_heads=16,
+        vocab_size=2048,
+        rope_theta=500000,
     ),
     "debugmodel_flex_attn": TransformerModelArgs(
         dim=256,