Revert "[FSDP2] Computed grad divide factors at runtime (#125484)"

pytorchmergebot · pytorchmergebot · commit dba689bbfdc6 · 2024-05-05T06:12:01.000Z
This reverts commit 9aa7699. Reverted #125484 on behalf of https://github.com/huydhn due to Sorry for reverting your change, I am trying to restore ROCm distributed failures in trunk https://hud.pytorch.org/pytorch/pytorch/commit/9aa7699185e4ec39077e3046dfd63244dffa9ddb ([comment](#125484 (comment)))
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_comm.py b/test/distributed/_composable/fsdp/test_fully_shard_comm.py
@@ -18,8 +18,6 @@
     OffloadPolicy,
 )
 from torch.distributed._composable.fsdp._fsdp_collectives import (
-    _div_if_needed,
-    _get_gradient_divide_factors,
     foreach_all_gather,
     foreach_all_gather_copy_out,
     foreach_reduce,
@@ -209,18 +207,6 @@ def test_reduce_scatter_fp32(self):
                 reduce_scatter_dtype=torch.float32,
             )
 
-    @unittest.skipIf(not TEST_CUDA, "no cuda")
-    def test_reduce_scatter_fp16(self):
-        param_sizes = self._get_param_sizes()
-        default_stream = torch.cuda.current_stream()
-        stream = torch.cuda.Stream()
-        for reduce_scatter_stream in (default_stream, stream):
-            self._test_reduce_scatter(
-                param_sizes,
-                reduce_scatter_stream=reduce_scatter_stream,
-                reduce_scatter_dtype=torch.float16,
-            )
-
     def _test_reduce_scatter(
         self,
         param_sizes: List[torch.Size],
@@ -252,24 +238,17 @@ def _test_reduce_scatter(
             orig_dtype=orig_params[0].dtype,
             reduce_dtype=reduce_scatter_dtype,
             device=self.device,
+            divide_factors=fsdp_param_group._grad_divide_factors,
             all_reduce_group=None,
             all_reduce_stream=all_reduce_stream,
         )
         torch.cuda.current_stream().wait_event(view_out_event)
 
         # Check reduce-scatter correctness
-        predivide_factor, postdivide_factor = _get_gradient_divide_factors(
-            group, None, reduce_scatter_dtype
-        )
         reduced_grads = [grad.detach().clone() for grad in unsharded_grads]
         for grad in reduced_grads:
-            _div_if_needed(grad, predivide_factor)
-            dist.all_reduce(
-                grad,
-                group=group,
-                op=dist.ReduceOp.AVG if predivide_factor is None else dist.ReduceOp.SUM,
-            )
-            _div_if_needed(grad, postdivide_factor)
+            dist.all_reduce(grad, group=group)
+            grad /= self.world_size
         for fsdp_param, reduced_grad in zip(fsdp_params, reduced_grads):
             sharded_grad = fsdp_param.sharded_param.grad
             self.assertIsInstance(sharded_grad, DTensor)
diff --git a/torch/distributed/_composable/fsdp/_fsdp_collectives.py b/torch/distributed/_composable/fsdp/_fsdp_collectives.py
@@ -125,6 +125,7 @@ def foreach_reduce(
     orig_dtype: torch.dtype,
     reduce_dtype: Optional[torch.dtype],
     device: torch.device,
+    divide_factors: Union[Tuple[None, None], Tuple[float, float]],
     all_reduce_group: Optional[dist.ProcessGroup],
     all_reduce_stream: torch.cuda.Stream,
 ) -> torch.cuda.Event:
@@ -141,9 +142,7 @@ def foreach_reduce(
         )
     grad_dtype = unsharded_grads[0].dtype
     reduce_dtype = reduce_dtype or grad_dtype
-    predivide_factor, postdivide_factor = _get_gradient_divide_factors(
-        reduce_scatter_group, all_reduce_group, reduce_dtype
-    )
+    predivide_factor, postdivide_factor = divide_factors
     world_size = reduce_scatter_group.size()
     padded_unsharded_sizes = tuple(
         _get_dim0_padded_size(grad.size(), world_size) for grad in unsharded_grads
@@ -167,22 +166,18 @@ def foreach_reduce(
             (reduce_scatter_output_numel,)
         )
         _div_if_needed(reduce_scatter_input, predivide_factor)
-        dist.reduce_scatter_tensor(
-            output=post_reduce_output,
-            input=reduce_scatter_input,
-            group=reduce_scatter_group,
-            op=ReduceOp.AVG if predivide_factor is None else ReduceOp.SUM,
+        _reduce_scatter(
+            post_reduce_output,
+            reduce_scatter_input,
+            reduce_scatter_group,
+            divide_factors,
         )
     view_out_stream = reduce_scatter_stream
     if all_reduce_group is not None:
         view_out_stream = all_reduce_stream
         all_reduce_stream.wait_stream(reduce_scatter_stream)
         with torch.cuda.stream(all_reduce_stream):
-            dist.all_reduce(
-                post_reduce_output,
-                group=all_reduce_group,
-                op=ReduceOp.AVG if predivide_factor is None else ReduceOp.SUM,
-            )
+            _all_reduce(post_reduce_output, all_reduce_group, divide_factors)
     with torch.cuda.stream(view_out_stream):
         _div_if_needed(post_reduce_output, postdivide_factor)
         post_reduce_output = _to_dtype_if_needed(post_reduce_output, orig_dtype)
@@ -262,27 +257,30 @@ def _get_all_gather_input_metadatas(
     )
 
 
-def _get_gradient_divide_factors(
-    reduce_scatter_group: dist.ProcessGroup,
-    all_reduce_group: Optional[dist.ProcessGroup],
-    reduce_dtype: torch.dtype,
-) -> Union[Tuple[None, None], Tuple[float, float]]:
-    # For fp32/bf16, we do not need to worry about overflow/underflow, so we
-    # use NCCL's built-in division to avoid separate div kernels
-    if reduce_dtype in (torch.float32, torch.bfloat16):
-        return None, None
-    data_parallel_size = reduce_scatter_group.size()
-    if all_reduce_group is not None:
-        data_parallel_size *= all_reduce_group.size()
-    # 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.
-    factor: int = 1
-    while data_parallel_size % factor == 0 and data_parallel_size / factor > factor:
-        factor *= 2
-    factor = float(factor)
-    return (factor, data_parallel_size / factor)
+def _reduce_scatter(
+    output: torch.Tensor,
+    input: torch.Tensor,
+    group: dist.ProcessGroup,
+    divide_factors: Union[Tuple[None, None], Tuple[float, float]],
+) -> None:
+    if divide_factors[0]:
+        dist.reduce_scatter_tensor(output, input, group=group)
+    else:
+        # Using NCCL's reduce-scatter to do the division by world size saves
+        # extra memory read/write from a separate division kernel
+        dist.reduce_scatter_tensor(output, input, op=ReduceOp.AVG, group=group)
+
+
+def _all_reduce(
+    tensor: torch.Tensor,
+    group: dist.ProcessGroup,
+    divide_factors: Union[Tuple[None, None], Tuple[float, float]],
+) -> None:
+    if divide_factors[0]:
+        dist.all_reduce(tensor, group=group)
+    else:
+        # saves extra memory read/write from a separate division kernel
+        dist.all_reduce(tensor, op=ReduceOp.AVG, group=group)
 
 
 def _div_if_needed(tensor: torch.Tensor, div_factor: Optional[float]) -> None:
diff --git a/torch/distributed/_composable/fsdp/_fsdp_param_group.py b/torch/distributed/_composable/fsdp/_fsdp_param_group.py
@@ -1,6 +1,6 @@
 import contextlib
 
-from typing import Any, cast, Dict, List, NamedTuple, Optional, Set, Tuple
+from typing import Any, cast, Dict, List, NamedTuple, Optional, Set, Tuple, Union
 
 import torch
 import torch.distributed as dist
@@ -164,6 +164,32 @@ def _init_mp_dtypes(self) -> None:
             )
         self._reduce_dtype = next(iter(reduce_dtypes))
 
+    def _init_grad_divide_factors(self):
+        data_parallel_world_size = 1
+        data_parallel_world_size *= self.mesh_info.shard_mesh_size
+        if self._is_hsdp:
+            data_parallel_world_size *= self.mesh_info.replicate_mesh_size
+        if self._reduce_dtype in (torch.float32, torch.bfloat16):
+            # Use NCCL's AVG op to divide after reduction since it is more
+            # performant and fp32 has sufficient precision
+            self._grad_divide_factors: Union[Tuple[None, None], Tuple[float, float]] = (
+                None,
+                None,
+            )
+            return
+        # 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.
+        factor: int = 1
+        while (
+            data_parallel_world_size % factor == 0
+            and data_parallel_world_size / factor > factor
+        ):
+            factor *= 2
+        factor = float(factor)
+        self._grad_divide_factors = (factor, data_parallel_world_size / factor)
+
     def lazy_init(self):
         # Lazy init should be idempotent
         param_names_on_meta = [
@@ -181,6 +207,7 @@ def lazy_init(self):
         # Initialize mixed precision attributes lazily in case the user changes
         # the parameter dtypes after construction time but before forward
         self._init_mp_dtypes()
+        self._init_grad_divide_factors()
         self._register_state_dict_hooks()
 
     # Runtime #
@@ -319,6 +346,7 @@ def post_backward(self, *unused: Any):
                 self._orig_dtype,
                 self._reduce_dtype,
                 self.device,
+                self._grad_divide_factors,
                 self._all_reduce_process_group
                 if self._is_hsdp and self.all_reduce_grads
                 else None,
