fix EP fsdp gradient divide factor (#1551)

issue pointed out in
#1534 (comment)
pytorch/pytorch#160285

solution given by @rakkit in
#1534 (comment)

Author: tianyu-l

torchtitan/distributed/parallel_dims.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 from dataclasses import dataclass
-from functools import cached_property
 
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 
@@ -219,11 +218,18 @@ def pp_enabled(self):
     def ep_enabled(self):
         return self.ep > 1
 
-    @cached_property
+    @property
+    def fsdp_gradient_divide_factor(self) -> int:
+        # This is needed for FSDP-sharded experts when Expert Parallel is enabled.
+        # Although the FSDP sharding of experts is done on a mesh of a different size than
+        # other parameters, the gradient division factor should be consistent with data.
+        return self.dp_replicate * self.dp_shard * self.cp
+
+    @property
     def non_data_parallel_size(self):
         return self.cp * self.tp * self.pp
 
-    @cached_property
+    @property
     def seq_len_divisor(self):
         # Sequence Parallel requires that seq_len be divisible by TP degree.
         # https://github.com/pytorch/torchtitan/pull/640#discussion_r1849481001

torchtitan/experiments/llama4/infra/parallelize.py

@@ -139,6 +139,7 @@ def parallelize_llama(
                 if dp_mod_ep_mesh_dim_names
                 else None
             ),
+            gradient_divide_factor=parallel_dims.fsdp_gradient_divide_factor,
         )
 
         if parallel_dims.dp_replicate_enabled:
@@ -270,6 +271,7 @@ def apply_fsdp(
     cpu_offload: bool = False,
     reshard_after_forward_policy: str = "default",
     dp_mod_ep_mesh: DeviceMesh | None = None,
+    gradient_divide_factor: int | None = None,
 ):
     """
     Apply data parallelism (via FSDP2) to the model.
@@ -322,6 +324,12 @@ def apply_fsdp(
                 **fsdp_mod_ep_config,
                 reshard_after_forward=reshard_after_forward,
             )
+            # NOTE: # Although the FSDP sharding of experts is done on a mesh of
+            #       a different size than other parameters, the gradient division
+            #       factor should be consistent with data.
+            transformer_block.moe.experts.set_gradient_divide_factor(
+                gradient_divide_factor,
+            )
 
         fully_shard(
             transformer_block,

torchtitan/models/deepseek_v3/infra/parallelize.py

@@ -122,6 +122,7 @@ def parallelize_deepseekv3(
                 if dp_mod_ep_mesh_dim_names
                 else None
             ),
+            gradient_divide_factor=parallel_dims.fsdp_gradient_divide_factor,
         )
 
         if parallel_dims.dp_replicate_enabled: