[FSDP] Add initial summon_full_params(with_grads=True)

Andrew Gu · Andrew Gu · commit dabe7889dd9b · 2022-09-27T21:03:15.000Z
ghstack-source-id: 9c80c33 Pull Request resolved: #85738
diff --git a/test/distributed/fsdp/test_fsdp_summon_full_params.py b/test/distributed/fsdp/test_fsdp_summon_full_params.py
@@ -1,25 +1,28 @@
 # Owner(s): ["oncall: distributed"]
+import contextlib
 import itertools
 import math
 import sys
 from copy import deepcopy
-from typing import Optional
+from typing import List, Optional
 
 import torch
 import torch.nn as nn
 from torch import distributed as dist
 from torch.distributed.fsdp import CPUOffload
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp import MixedPrecision
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
 from torch.distributed.fsdp.flat_param import FlatParamHandle
 from torch.distributed.fsdp.wrap import enable_wrap, wrap
+from torch.nn.parallel.distributed import DistributedDataParallel as DDP
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import (
     CUDAInitMode,
     DeterministicModel,
     FSDPInitMode,
     FSDPTest,
     NestedWrappedModule,
+    TransformerWithSharedParams,
 )
 from torch.testing._internal.common_utils import (
     TEST_WITH_DEV_DBG_ASAN,
@@ -578,6 +581,124 @@ def test_named_parameters_buffers(self, prefix: str, recurse: bool):
                     self.assertEqual(n1, n2)
                     self.assertEqual(p1, p2)
 
+    @skip_if_lt_x_gpu(2)
+    def test_with_grads(self):
+        self.run_subtests(
+            {
+                "writeback": [False, True],
+                "offload_to_cpu": [False, True],
+                "sharding_strategy": [
+                    ShardingStrategy.FULL_SHARD,
+                    ShardingStrategy.SHARD_GRAD_OP,
+                    ShardingStrategy.NO_SHARD,
+                ],
+                "use_orig_params": [True],
+            },
+            self._test_with_grads,
+        )
+
+    def _test_with_grads(
+        self,
+        writeback: bool,
+        offload_to_cpu: bool,
+        sharding_strategy: ShardingStrategy,
+        use_orig_params: bool,
+    ):
+        def _check_grads(
+            ddp_model: DDP,
+            fsdp_model: FSDP,
+            old_fsdp_grads: Optional[List[torch.Tensor]],
+        ):
+            WRITEBACK_FACTOR = 2
+            with FSDP.summon_full_params(
+                fsdp_model,
+                writeback=writeback,
+                offload_to_cpu=offload_to_cpu,
+                with_grads=True,
+            ):
+                for (n1, p1), (n2, p2) in zip(
+                    ddp_model.module.named_parameters(),
+                    fsdp_model.named_parameters(),
+                ):
+                    # Parameter names are only expected to match because
+                    # `fsdp_model` has top-level FSDP, so its
+                    # `named_parameters()` cleans *all* of the names
+                    self.assertEqual(n1, n2)
+                    assert p1.grad is not None
+                    torch.testing.assert_close(p1.grad, p2.grad)
+                    # Ensure that the tensor is not all zeros, which would
+                    # mean that the multiplication is vacuous
+                    assert torch.count_nonzero(p2.grad) > 0
+                    p2.grad *= WRITEBACK_FACTOR
+            new_fsdp_grads = [
+                param.grad for param in fsdp_model.parameters()
+                if param.grad is not None
+            ]
+            writeback_persists = writeback or sharding_strategy == ShardingStrategy.NO_SHARD
+            for old_grad, new_grad in zip(old_fsdp_grads, new_fsdp_grads):
+                if writeback_persists:
+                    torch.testing.assert_close(old_grad * WRITEBACK_FACTOR, new_grad)
+                else:
+                    torch.testing.assert_close(old_grad, new_grad)
+            if writeback_persists:
+                # Modify the DDP gradients for parity
+                for param in ddp_model.parameters():
+                    param.grad *= WRITEBACK_FACTOR
+
+        def _get_error_context(is_supported: bool):
+            return (
+                contextlib.suppress() if is_supported
+                else self.assertRaises(NotImplementedError)
+            )  # some configs not implemented yet
+
+        def _get_fsdp_grads(fsdp_model: FSDP, is_supported: bool):
+            if is_supported:
+                return [
+                    param.grad.clone() for param in fsdp_model.parameters()
+                    if param.grad is not None
+                ]
+            return None  # unused
+
+        is_supported = use_orig_params and not offload_to_cpu
+        model = TransformerWithSharedParams.init(
+            self.process_group,
+            FSDPInitMode.NO_FSDP,
+            CUDAInitMode.CUDA_BEFORE,
+            deterministic=True,
+        )
+        ddp_model = DDP(model, device_ids=[self.rank])
+        fsdp_model = TransformerWithSharedParams.init(
+            self.process_group,
+            FSDPInitMode.RECURSIVE,
+            CUDAInitMode.CUDA_BEFORE,
+            deterministic=True,
+            fsdp_kwargs={
+                "use_orig_params": use_orig_params,
+                "sharding_strategy": sharding_strategy,
+            },
+        )
+        with FSDP.summon_full_params(fsdp_model):
+            for p1, p2 in zip(ddp_model.module.parameters(), fsdp_model.parameters()):
+                assert torch.all(torch.isclose(p1, p2))
+
+        # Check `summon_full_params()` after backward
+        inp = fsdp_model.get_input(torch.device("cuda"))
+        ddp_out = ddp_model(*inp)
+        fsdp_out = fsdp_model(*inp)
+        ddp_out.sum().backward()
+        fsdp_out.sum().backward()
+        old_fsdp_grads = _get_fsdp_grads(fsdp_model, is_supported)
+        with _get_error_context(is_supported):
+            _check_grads(ddp_model, fsdp_model, old_fsdp_grads)
+
+        # Check `summon_full_params()` between forward and backward
+        inp = fsdp_model.get_input(torch.device("cuda"))
+        ddp_out = ddp_model(*inp)
+        fsdp_out = fsdp_model(*inp)
+        old_fsdp_grads = _get_fsdp_grads(fsdp_model, is_supported)
+        with _get_error_context(is_supported):
+            _check_grads(ddp_model, fsdp_model, old_fsdp_grads)
+
 
 instantiate_parametrized_tests(TestSummonFullParams)
 instantiate_parametrized_tests(TestSummonFullParamsNoShard)
diff --git a/torch/distributed/fsdp/flat_param.py b/torch/distributed/fsdp/flat_param.py
@@ -846,6 +846,37 @@ def _free_low_precision_sharded_param(self):
         self._check_low_precision_shard()
         _free_storage(self.flat_param._mp_shard)  # type: ignore[attr-defined]
 
+    @torch.no_grad()
+    def unshard_grad(self):
+        if not self.uses_sharded_strategy:
+            self._use_unsharded_grad_views()
+            return
+        flat_param = self.flat_param
+        self._check_unsharded(flat_param)
+        padded_unsharded_grad = torch.empty(
+            flat_param._padded_unsharded_size,  # type: ignore[attr-defined]
+            device=self.device,
+        )
+        if flat_param.grad is None:
+            flat_param._saved_grad_shard = None  # type: ignore[attr-defined]
+            sharded_grad = torch.zeros_like(flat_param)  # type: ignore[attr-defined]
+        else:
+            self._check_sharded(flat_param.grad)
+            flat_param._saved_grad_shard = flat_param.grad  # type: ignore[attr-defined]
+            sharded_grad = flat_param._saved_grad_shard  # type: ignore[attr-defined]
+        dist._all_gather_base(padded_unsharded_grad, sharded_grad, self.process_group)
+        unsharded_size = self.flat_param._unpadded_unsharded_size
+        flat_param.grad = padded_unsharded_grad[:unsharded_size.numel()].view(unsharded_size)
+        self._use_unsharded_grad_views()
+
+    def reshard_grad(self):
+        if self._use_orig_params:
+            self._use_sharded_grad_views()
+        if not self.uses_sharded_strategy:
+            return
+        self.flat_param.grad = self.flat_param._saved_grad_shard  # type: ignore[attr-defined]
+        delattr(self.flat_param, "_saved_grad_shard")
+
     def prepare_gradient_for_backward(self):
         """
         Prepares the gradient for the backward computation by saving and
@@ -1093,7 +1124,7 @@ def _use_unsharded_views(self, as_params: bool) -> None:
                 be used during forward/backward computation and when hiding the
                 original parameters from :meth:`nn.Module.named_parameters`.
         """
-        self._check_unsharded()
+        self._check_unsharded(self.flat_param)
         views = self._get_unflat_views(self.flat_param)
         for i, (view, (param_name, module, _)) in enumerate(
             zip(views, self.flat_param._param_infos)
@@ -1139,6 +1170,41 @@ def _use_unsharded_views(self, as_params: bool) -> None:
             else:
                 setattr(module, param_name, prim_param)
 
+    def _use_unsharded_grad_views(self) -> None:
+        """
+        Unflattens the unsharded flattened parameter's gradient by setting the
+        original module parameter variables' gradients to be views into it.
+        """
+        # Expects the gradient to be in `flat_param.grad`
+        if self.flat_param.grad is None:
+            return
+        self._check_unsharded(self.flat_param.grad)
+        views = self._get_unflat_views(self.flat_param, self.flat_param.grad)
+        for i, (view, (param_name, module, _)) in enumerate(
+            zip(views, self.flat_param._param_infos)
+        ):
+            p_assert(
+                hasattr(module, param_name),
+                f"{self.flat_param._prefixed_param_names[i]} is missing",
+            )
+            param = getattr(module, param_name)
+            param.grad = view
+        for i, (
+            param_name,
+            module,
+            module_name,
+            prim_param_name,
+            prim_module,
+            _,
+        ) in enumerate(self.flat_param._shared_param_infos):
+            p_assert(
+                hasattr(module, param_name),
+                f"{module_name + '.' + param_name if module_name else param_name} is missing",
+            )  # did not save prefixed name
+            param = getattr(module, param_name)
+            prim_param = getattr(prim_module, prim_param_name)
+            param.grad = prim_param.grad
+
     @contextlib.contextmanager
     def unflatten_as_params(self) -> Generator:
         """
@@ -1223,16 +1289,7 @@ def _use_sharded_grad_views(self) -> None:
         """
         flat_param = self.flat_param
         self._check_sharded(flat_param)
-        # Priority: `_cpu_grad` > `_saved_grad_shard` > `grad`
-        # - CPU offloading: `_cpu_grad`
-        # - No CPU offloading + sharded strategies: `_saved_grad_shard`
-        # - No CPU offloading + `NO_SHARD`: `grad`
-        if hasattr(flat_param, "_cpu_grad"):
-            grad = flat_param._cpu_grad  # type: ignore[attr-defined]
-        elif hasattr(flat_param, "_saved_grad_shard"):
-            grad = flat_param._saved_grad_shard  # type: ignore[attr-defined]
-        else:
-            grad = flat_param.grad
+        grad = self.sharded_grad
         if grad is None:
             return  # no-op
         self._check_sharded(grad)
@@ -1474,6 +1531,26 @@ def parameter_module_names(self) -> Iterator[Tuple[str, str]]:
         ):
             yield (param_name, module_name)
 
+    @property
+    def sharded_grad(self) -> Optional[Tensor]:
+        """Returns the handle's sharded gradient."""
+        flat_param = self.flat_param
+        # Priority for non-`None`: `_cpu_grad` > `_saved_grad_shard` > `grad`
+        # - CPU offloading: `_cpu_grad`
+        # - No CPU offloading + sharded strategies: `_saved_grad_shard`
+        # - No CPU offloading + `NO_SHARD`: `grad`
+        if hasattr(flat_param, "_cpu_grad"):
+            grad = flat_param._cpu_grad  # type: ignore[attr-defined]
+        elif hasattr(flat_param, "_saved_grad_shard"):
+            grad = flat_param._saved_grad_shard  # type: ignore[attr-defined]
+        else:
+            p_assert(
+                flat_param.grad is None or not self.uses_sharded_strategy,
+                "Sharded strategies should use `_cpu_grad` or `_saved_grad_shard`",
+            )
+            grad = flat_param.grad
+        return grad
+
     #######################
     # CHECKS & INVARIANTS #
     #######################
@@ -1520,13 +1597,13 @@ def _check_low_precision_shard(self):
             f"Expects the low precision shard to be on {self.device} but got {device}",
         )
 
-    def _check_unsharded(self):
-        msg_prefix = "Expects the flattened parameter to be unsharded "
-        p_assert(self.flat_param is not None, msg_prefix + "but got `None`")
+    def _check_unsharded(self, tensor: Tensor):
+        msg_prefix = "Expects tensor to be unsharded "
+        p_assert(tensor is not None, msg_prefix + "but got `None`")
         unsharded_size = self.flat_param._unpadded_unsharded_size
         p_assert(
-            self.flat_param.size() == unsharded_size,
-            msg_prefix + f"with size {unsharded_size} but got {self.flat_param.size()}",
+            tensor.size() == unsharded_size,
+            msg_prefix + f"with size {unsharded_size} but got {tensor.size()}",
         )
 
     def _check_sharded(self, tensor: Tensor):
diff --git a/torch/distributed/fsdp/fully_sharded_data_parallel.py b/torch/distributed/fsdp/fully_sharded_data_parallel.py
