diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py index c4659457e..c97cdeeae 100644 --- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py +++ b/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. diff --git a/tests/ci_test_list_1.txt b/tests/ci_test_list_1.txt index aa2fe4460..45678bd57 100644 --- a/tests/ci_test_list_1.txt +++ b/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 diff --git a/tests/nn/data_parallel/test_fsdp_overlap.py b/tests/nn/data_parallel/test_fsdp_overlap.py index e1e584609..87f86cb4f 100644 --- a/tests/nn/data_parallel/test_fsdp_overlap.py +++ b/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. diff --git a/tests/nn/data_parallel/test_fsdp_prefetch_order.py b/tests/nn/data_parallel/test_fsdp_prefetch_order.py new file mode 100644 index 000000000..631aa8107 --- /dev/null +++ b/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, + )