Memoryfix

fairscale/nn/data_parallel/fully_sharded_data_parallel.py

@@ -1172,6 +1172,8 @@ def _reset_lazy_init(self) -> None:
         self._is_root: Optional[bool] = None
         self._streams: Dict[str, torch.cuda.Stream] = {}
         self._reducer: Optional[ReduceScatterBucketer] = None
+        self._forward_ordering: List[FullyShardedDataParallel] = []
+        self._backward_rebuild_ordering: List[FullyShardedDataParallel] = []
         for p in self.params:
             if hasattr(p, "_fp32_shard"):
                 del p._fp32_shard  # reset _init_param_attributes
@@ -1344,6 +1346,8 @@ def _set_is_root(self) -> None:
                 m.no_broadcast_optim_state = m.no_broadcast_optim_state or (
                     (m.world_size == 1) and (m.world_size < self.world_size) and (m.process_group != self.process_group)
                 )
+                m._forward_ordering = self._forward_ordering
+                m._backward_rebuild_ordering = self._backward_rebuild_ordering
 
     def _setup_streams(self) -> None:
         """Create streams to overlap data transfer and computation."""
@@ -1419,6 +1423,13 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
 
         outputs = self.module(*args, **kwargs)
 
+        # In the first forward pass, track the order that modules are computed.
+        # In the following passes, we assume that the order remains the same to
+        # to kick off the all-gather for the next module in the list, in case we
+        # are waiting for the computation to finish.
+        if self not in self._forward_ordering:
+            self._forward_ordering.append(self)
+
         if self.reshard_after_forward:
             self._free_full_params()
             if self.mixed_precision or self.move_params_to_cpu:
@@ -1482,6 +1493,7 @@ def _pre_backward_hook(*unused: Any) -> None:
             # that final backward callback is attached to the outer most
             # backward graph task and called after all the backward
             # calls are completed.
+
             if self._is_root:
                 self._queue_wait_for_post_backward()
 
@@ -1501,6 +1513,13 @@ def _pre_backward_hook(*unused: Any) -> None:
             # overhead.
             if self.reshard_after_forward:
                 self._rebuild_full_params()
+
+                # Similar to _forward_ordering, in the first backward pass we track the order
+                # that weights were gathered for modules in the backward pass. Then, we use
+                # this order in future passes to kick off the all-gather for the next module
+                # in case we are waiting for the computation of the current module to finish.
+                if self not in self._backward_rebuild_ordering:
+                    self._backward_rebuild_ordering.append(self)
             else:
                 self._use_full_params()
 
@@ -1864,7 +1883,9 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                     self._output_pre_backward_hook_registered.clear()
 
     @torch.no_grad()
-    def _rebuild_full_params(self, force_full_precision: bool = False) -> Optional[List[Tuple[torch.Tensor, bool]]]:
+    def _rebuild_full_params(
+        self, force_full_precision: bool = False, wait_for_all_gather: bool = True
+    ) -> Optional[List[Tuple[torch.Tensor, bool]]]:
         """
         Gather all shards of params.
 
@@ -1935,6 +1956,9 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
         # Early exit if we already have full params and don't need full precision.
         if self.has_full_params and not force_full_precision:
+            if not wait_for_all_gather:
+                return None
+            torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
             for p in self.params:
                 update_p_data()
             return output_tensors
@@ -1997,7 +2021,8 @@ def update_p_data(custom_output_tensor: Optional[torch.Tensor] = None) -> None:
 
                     if self.move_params_to_cpu and (self.params[0].dtype == self.compute_dtype):
                         self._free_fp16_param_shard([p])
-
+        if not wait_for_all_gather:
+            return None
         torch.cuda.current_stream().wait_stream(self._streams["all_gather"])
         return output_tensors
 
@@ -2052,6 +2077,9 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         self.has_full_params = False
         current_stream = torch.cuda.current_stream()
         for p in params:
+            # Shared params are not owned by this FSDP instance.
+            if hasattr(p, "_is_shared") and p._is_shared:
+                continue
             if not p._is_sharded:  # e.g., world_size == 1
                 if self.mixed_precision or self.move_params_to_cpu:
                     self._free_fp16_param_shard([p])
@@ -2067,6 +2095,42 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
             # the Storage to 0 to save memory.
             free_storage_(p._full_param_padded)
 
+        # When we are memory bound (which is the case here as we are freeing up
+        # params), we are not able to let the CPU run completely free because
+        # it will end up scheduling GPU operations required to compute all
+        # future modules. This causes significant increases in GPU reserved
+        # memory and potential thrashing.
+        # So instead, we simply schedule the all-gather for the next module
+        # to be executed and wait for the computations of the current module
+        # to finish before moving forward.
+        self._schedule_next_all_gather_and_synchronize()
+
+    @torch.no_grad()
+    def _schedule_next_all_gather_and_synchronize(self) -> None:
+        self.assert_state([TrainingState.FORWARD, TrainingState.BACKWARD_POST])
+        ordering = self._forward_ordering
+        if self.training_state == TrainingState.BACKWARD_POST:
+            ordering = self._backward_rebuild_ordering
+        # Not all modules require rebuilding in backward pass, so this check is required.
+        if self in ordering:
+            next_idx = ordering.index(self) + 1
+            if next_idx < len(ordering):
+                next_module = ordering[next_idx]
+                # _pre_backward_hook_has_run prevents us from kicking off all-gather on a forward pass happening due to activation
+                # checkpointing. In these scenarios, forward only runs up until the module that already ran their backward hook.
+                # In addition, if the module to be scheduled has a shared param, there is a potential race condition where params for
+                # the current module are freed and all gather for the next module is happening. Similarly, modules with ssd_offload
+                # are not supported because ssd_offload happens before every all-gather call. So we just skip such modules.
+                if (
+                    not next_module._pre_backward_hook_has_run
+                    and not next_module._has_shared_params
+                    and not next_module.ssd_offload
+                ):
+                    # Kick-off all gather for the next module without waiting.
+                    next_module._rebuild_full_params(wait_for_all_gather=False)
+        # Wait for computation kernels to finish running.
+        torch.cuda.current_stream().synchronize()
+
     def local_metadata_dict(self) -> Dict[str, Any]:
         """
         Get the information needed to reconstruct the model from shards offline.

tests/ci_test_list_1.txt

@@ -9,4 +9,5 @@ tests/nn/data_parallel/test_fsdp_input.py
 tests/nn/data_parallel/test_fsdp_optimizer_utils.py
 tests/nn/data_parallel/test_fsdp.py
 tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py
+tests/nn/data_parallel/test_fsdp_prefetch_order.py
 tests/optim/test_layerwise_gradient_scaler.py

tests/nn/data_parallel/test_fsdp_overlap.py

@@ -112,8 +112,6 @@ def run(compute_cycles, all_gather_cycles):
 
         # We run 20 iterations but only collect timing data from the minimal 10
         # data points because nondeterministic system events can disturb the timing.
-        cpu_iter = Min10()
-        cpu_wait = Min10()
         gpu_compute = Min10()
         gpu_total = Min10()
         for _ in range(20):
@@ -165,12 +163,7 @@ def _delayed_all_gather_base(*args, **kwargs):
             else:
                 for p in model.parameters():
                     p.grad = None
-
-            cpu_iter_time = time.process_time() - cpu_start
-
-            # wait for gpu
             out.item()
-            cpu_wait_for_gpu_time = time.process_time() - cpu_start - cpu_iter_time
 
             # get sum of the compute time
             times = []
@@ -182,15 +175,11 @@ def _delayed_all_gather_base(*args, **kwargs):
             # get gpu compute + all_gather time
             overall_gpu_time = e1.elapsed_time(e2)
 
-            cpu_iter.add(cpu_iter_time)
-            cpu_wait.add(cpu_wait_for_gpu_time)
             gpu_compute.add(sum(times))
             gpu_total.add(overall_gpu_time)
 
         del model
         return {
-            "cpu_iter": cpu_iter.avg(),
-            "cpu_wait": cpu_wait.avg(),
             "gpu_compute": gpu_compute.avg(),
             "gpu_total": gpu_total.avg(),
         }
@@ -204,33 +193,17 @@ def _delayed_all_gather_base(*args, **kwargs):
     debug_string = f"\nrank{rank}:\n  e1: {e1}\n  e2: {e2}\n  e3: {e3}\n  e4: {e4}"
     print(debug_string)
 
-    # Check the cpu/gpu timing. CPU should run ahead of GPU. Therefore, cpu-gpu
-    # wait should be long, except when there is no real work on GPU.
-    #
-    # If the assertions fail below, we likely have a cpu-gpu wait in the forward/backward pass.
-    short = [e1["cpu_iter"], e2["cpu_iter"], e3["cpu_iter"], e4["cpu_iter"], e1["cpu_wait"]]
-    long = [e3["cpu_wait"], e4["cpu_wait"]]
-    if world_size == 1:
-        short.append(e2["cpu_wait"])  # all gather should not be happening.
-    else:
-        long.append(e2["cpu_wait"])  # all gather should happen and prolong the cpu-gpu wait.
-    for s in short:
-        for l in long:
-            # 5X longer is a safe margin, since the GPU work timing is around 100X more
-            # of that of the CPU.
-            assert s * 5 < l, f"{s} * 5 < {l} in " + debug_string
-
     # Check the GPU timing.
     short = [e1["gpu_compute"], e1["gpu_total"], e2["gpu_compute"]]
     long = [e3["gpu_compute"], e3["gpu_total"], e4["gpu_compute"], e4["gpu_total"]]
     if world_size == 1:
         short.append(e2["gpu_total"])  # all gather should not be happening.
     else:
-        long.append(e2["gpu_total"])  # all gather should happen and prolong the cpu-gpu wait.
+        long.append(e2["gpu_total"])  # all gather should happen and prolong the gpu wait.
     for s in short:
         for l in long:
             # 10X longer is a safe margin, since the time is around 100X longer
-            # when there is work on GPU vs. no work.
+            # when there is compute work on GPU vs. no work.
             assert s * 10 < l, f"{s} * 10 < {l} in " + debug_string
 
     # Check the GPU overlapping when there is all-gather.

tests/nn/data_parallel/test_fsdp_prefetch_order.py

@@ -0,0 +1,111 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pylint: disable=missing-module-docstring
+# pylint: disable=missing-class-docstring
+# pylint: disable=missing-function-docstring
+
+""" Check that the ordering used for prefetching model weights
+    matches the expected execution order for the model.
+"""
+
+import tempfile
+
+import pytest
+import torch
+import torch.multiprocessing as mp
+import torch.nn as nn
+from torch.optim import SGD
+
+from fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown
+from fairscale.internal import torch_version
+from fairscale.nn import checkpoint_wrapper
+from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
+from fairscale.nn.data_parallel import TrainingState, auto_wrap_bn
+from fairscale.nn.wrap import enable_wrap, wrap
+
+
+def _get_module_type(fsdp):
+    m = fsdp._fsdp_wrapped_module._fpw_module
+    if isinstance(m, nn.Sequential):
+        return type(m[0])
+    return type(m)
+
+
+def _test_func(rank, world_size, fsdp_config, tempfile_name, unused):
+    result = dist_init(rank, world_size, tempfile_name, unused)
+    assert result, "Dist init failed"
+
+    assert isinstance(fsdp_config, dict), str(fsdp_config)
+
+    torch.cuda.set_device(rank)
+
+    class Model(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.block1 = nn.Sequential(nn.Conv2d(3, 4, kernel_size=3), nn.BatchNorm2d(4), nn.ReLU(inplace=True))
+            self.block2 = nn.Sequential(nn.Conv2d(4, 4, kernel_size=3), nn.BatchNorm2d(4), nn.ReLU(inplace=False))
+            self.block3 = nn.Linear(12, 8)
+            self.head = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), nn.Linear(4, 10))
+
+        def forward(self, x):
+            return self.head(self.block3(self.block2(self.block1(x))))
+
+    model = Model()
+    # Wrapping BatchNorm as separate modules for the forward pass.
+    model.block1 = auto_wrap_bn(model.block1, fsdp_config={"reshard_after_forward": True})
+    model.block2 = auto_wrap_bn(model.block2, fsdp_config={"reshard_after_forward": True})
+
+    # Checkpoints shouldn't affect the ordering.
+    model.block2 = checkpoint_wrapper(model.block2)
+
+    with enable_wrap(
+        wrapper_cls=FSDP,
+    ):
+        model.block1 = wrap(model.block1)
+        model.block2 = wrap(model.block2)
+        model.block3 = wrap(model.block3)
+        model = wrap(model)
+
+    optim = SGD(model.parameters(), lr=0.1)
+    model = model.cuda()
+
+    # Orderings are stored in the first pass.
+    in_data = torch.randn(size=(2, 3, 16, 16)).cuda()
+    in_data.requires_grad = True
+    out = model(in_data)
+    out.sum().backward()
+    optim.step()
+
+    expected_forward_ordering = [nn.BatchNorm2d, nn.Conv2d, nn.BatchNorm2d, nn.Conv2d, nn.Linear, Model]
+    actual_forward_ordering = [_get_module_type(m) for m in model._forward_ordering]
+    assert expected_forward_ordering == actual_forward_ordering
+
+    expected_backward_ordering = [nn.Linear, nn.Conv2d, nn.BatchNorm2d, nn.BatchNorm2d, nn.Conv2d]
+    actual_backward_ordering = [_get_module_type(m) for m in model._backward_rebuild_ordering]
+    assert expected_backward_ordering == actual_backward_ordering
+
+    model.assert_state(TrainingState.IDLE)
+    teardown()
+
+
+@skip_if_single_gpu
+def test():
+    if torch_version() < (1, 6, 0):
+        pytest.skip("older pytorch doesn't support reduce_scatter")
+
+    temp_file_name = tempfile.mkstemp()[1]
+    unused = tempfile.mkstemp()[1]
+
+    fsdp_config = {}
+
+    # Using world_size > 1 to trigger all-gathers.
+    world_size = 2
+    mp.spawn(
+        _test_func,
+        args=(world_size, fsdp_config, temp_file_name, unused),
+        nprocs=world_size,
+        join=True,
+    )