Repository: facebookresearch/fairscale

Issue:
FSDP does not overlap all gather with forward pass
## 🐛 Bug

FSDP doesn't seem to overlap all-gather calls with forward pass when wrapping individual layers with `wrap()` or `auto_wrap()`. However, the reduce-scatter calls do overlap with computation in the backward pass. 

## To Reproduce

An example of wrapping individual layers with `wrap()` is, 

```
    layers=[]
    fsdp_params = dict(flatten_parameters=True, reshard_after_forward=False)
    with enable_wrap(wrapper_cls=FSDP, process_group=group, **fsdp_params):
        for i in range(num_layers):
            layer=wrap(BERTLayer(hidden_size, hidden_size, num_heads, i))
            assert isinstance(layer,FSDP)
            layers.append(layer)
    m = nn.Sequential(*layers).to(rank)
    opt = base_optimizer(params=m.parameters(), **base_optimizer_arguments)
```

## Problem

The Nsys profile of the above code with 8 layers (num_layers = 8) is: 

![image](https://user-images.githubusercontent.com/8636796/115657626-10a28500-a2ec-11eb-8892-45f16a033bbb.png)

The highlighted green slice is the first layer of the model in the forward pass. The AllGather calls do not overlap with `Kernels`, but in the backward pass, ReduceScatter calls do overlap with `Kernels` as desired. 

This symptom could be caused by the line: 

https://github.com/facebookresearch/fairscale/blob/b0048b2873de5a1752226977a24a140da22030c7/fairscale/nn/data_parallel/fully_sharded_data_parallel.py#L909

This line adds dependency for AllGather calls to wait on the previous optimizer step (which makes sense), but it also stops the next layer of the model from gathering model parameters from different processors, while the current layer is running computations. This `wait_stream` may be a bit too heavy-handed. It could probably be changed to wait on an event of the previous opt step, rather than on the `current_stream()`, serializing operations across layers. 

## Environment

 - PyTorch Version (e.g., 1.0): 1.9.0a0+df837d0
 - OS (e.g., Linux): Ubuntu 20.04.1 LTS (x86_64)
 - How you installed PyTorch (`conda`, `pip`, source): `conda`
 - Build command you used (if compiling from source): 
 - Python version: 3.8 (64-bit runtime)
 - CUDA/cuDNN version: 11.2.142/8.1.1
 - GPU models and configuration: A100-SXM4-40GB
 - Any other relevant information:



Discussion:
@cyanguwa Thank you so much for this excellent issue! CC: @myleott 

> This line adds dependency for AllGather calls to wait on the previous optimizer step

If the dependency is added only to the optimizer, then between layers in a forward pass, the overlap should still be happening.  However, your profile seems to show it wasn't the case. Do you have a small test case and steps to reproduce this so that we can investigate it using the case? Thanks!
We share a single "all_gather" stream across all layers, without synchronization, so we get overlapping as long as there are no device-to-host transfers in the forward pass (that way subsequent layers queue up their all-gathers while the preceeding layers are executing).

Re: the all_gather stream synchronizing with the current stream, that it's only called once from the "root" FSDP instance: https://github.com/facebookresearch/fairscale/blob/b0048b2873de5a1752226977a24a140da22030c7/fairscale/nn/data_parallel/fully_sharded_data_parallel.py#L758-L770.

Most likely you have device-to-host transfers in your forward pass, which is preventing overlapping. You can test this by putting a `torch.cuda._sleep(3000000000)` at the start of your forward pass and making sure you get all the way to the end of the forward pass without blocking (since any blocking would indicate a device-to-host transfer, which would then block on the sleep command).

More generally, we should fix this by supporting explicit prefetching of parameters in FSDP, which is on my TODO list, but won't be ready for a little while.
I am taking a closer look at this next. @cyanguwa, any new findings on your side? How do you setup nsys for profiling and generate the graph in your first comment? Do you mind sharing? Thanks a lot!
Hi @min-xu-ai, thanks for looking into this!

This is a dialed-down version of BERT, and the Nsight profiling command line is `nsys profile -o fsdp.wrap python fsdp.py`. You can install `nvtx` via `pip install nvtx`. Let me know if there's anything else you need!

```
$ cat fsdp.py
import sys
import torch
import argparse
import fairscale
import time
import os
import nvtx

import torch.nn as nn
import torch.multiprocessing as mp
import torch.distributed as dist
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
from fairscale.nn.wrap import auto_wrap, enable_wrap, wrap

def parse_args():

    parser = argparse.ArgumentParser()
    parser.add_argument('--batch-size', type=int, default=1)
    parser.add_argument('--lr', default=1e-4)
    parser.add_argument('--seq-length', default=1024, type=int)
    parser.add_argument('--num-heads', default=24, type=int)
    parser.add_argument('--hidden-size', default=2304, type=int)
    parser.add_argument('--num-layers', default=8, type=int)
    parser.add_argument('--iterations', default=5, type=int)

    return parser.parse_args()

class BERTLayer(nn.Module):

    def __init__(self, embed_size, hidden_size, num_heads, layer_id):
        super(BERTLayer, self).__init__()
        self.id = layer_id
        self.mha = nn.MultiheadAttention(embed_size, num_heads)
        self.self_output = nn.Linear(hidden_size, hidden_size)
        self.dropout1 = nn.Dropout(0.1)
        self.layer_norm1 = nn.LayerNorm(embed_size)
        self.intermediate = nn.Linear(hidden_size, hidden_size)
        self.act = nn.GELU()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.dropout2 = nn.Dropout(0.1)
        self.layer_norm2 = nn.LayerNorm(embed_size)

    def forward(self, input):
        t1, _ = self.mha(input, input, input)
        t2 = self.self_output(t1)
        t2 = self.dropout1(t2)
        t3 = self.layer_norm1(input + t2)
        t4 = self.intermediate(t3)
        t4 = self.act(t4)
        t5 = self.dense(t4)
        t5 = self.dropout2(t5)
        t6 = self.layer_norm2(t3 + t5)
        return t6

def run(rank, world_size):

    args = parse_args()
    batch_size = args.batch_size
    seq_length = args.seq_length
    hidden_size = args.hidden_size
    num_heads = args.num_heads
    num_layers = args.num_layers

    layers = [BERTLayer(hidden_size, hidden_size, num_heads, i) for i in range(num_layers)]
    m = nn.Sequential(*layers)

    torch.cuda.set_device(rank)

    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
    group = dist.new_group()

    base_optimizer = torch.optim.Adam
    base_optimizer_arguments = {"lr": args.lr}

    layers=[]
    fsdp_params = dict(flatten_parameters=True, reshard_after_forward=False)
    with enable_wrap(wrapper_cls=FSDP, process_group=group, **fsdp_params):
        for i in range(num_layers):
            layer=wrap(BERTLayer(hidden_size, hidden_size, num_heads, i))
            layers.append(layer)
    m = nn.Sequential(*layers).to(rank)

    opt = base_optimizer(params=m.parameters(), **base_optimizer_arguments)

    fake_input = torch.randn(args.seq_length, args.batch_size, hidden_size).to(rank)

    # optimisation loop
    start_time = time.time()
    for i in range(args.iterations):
        # reset gradient buffers
        opt.zero_grad()

        # fprop
        with nvtx.annotate(message="forward"+str(i), color="green"):
            y = m(fake_input)

        # fake a loss via summation
        loss = y.sum().abs()

        # Run bprop
        with nvtx.annotate(message="backward"+str(i), color="red"):
            loss.backward()

        # Update weights
        with nvtx.annotate(message="optimizer"+str(i), color="blue"):
            opt.step()

    torch.cuda.synchronize()
    stop_time = time.time()

    if rank == 0:
        print("Time taken per training iteration: {:.3f} s".format((stop_time - start_time)/args.iterations))

if __name__ == "__main__":
    world_size=8
    mp.spawn(run, args=(world_size,),nprocs=world_size)
```
Thanks @cyanguwa. I think I have recreated this on my machine. I see a bunch of DtoD `memcpy` that seems to be in the forward and in the backward I don't see them and I only see a bunch of `memset`. I will investigate more.

<img width="435" alt="image" src="https://user-images.githubusercontent.com/24926999/116603488-d33c8980-a8e1-11eb-9b27-fba3d40f5096.png">

Just note it down, installing the profiling tools were not trivial for me:

1. I was using py3.9 but I needed to switch to py3.8 so that pip install nvtx can use a prebuilt wheel package
2. I am on arch linux, so I needed cuda-tools package for the nsight system tools
3. with my GPUs and cuda version, I needed to use pytorch 1.9 nightly instead of 1.8 for fsdp to work
CC: @cyanguwa @min-xu-ai 

Jumping on the issue, to add some details I gathered during my investigations.

I noticed that two FSDP block A and B will not overlap with each other because the AllGather of B will wait for A to finish its computation. This explains why @cyanguwa does not see any AllGather overlap between the blocks in the BERT model.

### Guilty line

I traced the lack of overlap to the `wait_stream` below:

```
    @torch.no_grad()
    def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
        ...
        self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(self._streams["all_gather"]):
            ...
```

This line was introduced in https://github.com/facebookresearch/fairscale/pull/501 to correct an issue with a MoE model, but it is the line responsible for the stream "all_gather" to wait for the current_stream which holds the computation of the previous FSDP layer (because they seem to all share the same "all_gather" stream").

Commenting this line will lead to AllGather correctly interleaving with compute (tested on my side).

### Workaround

Leaving a small layer Z (for example a ReLU + a BN) between two FSDP wrapped blocks A and B (so A then Z then B) will lead to the computation of Z to be overlapped with the AllGather of B.

I built a reduced test case with 3 layers one after the other in which you can observe this phenomena and the workaround.

### How to reproduce

Run the script with `-t 1` to use FSDP and with `-t 2` with a different variant of wrapping FSDP in which there is a gap between two FSDP wrapped layers.

```
import argparse
from enum import Enum

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from fairscale.nn import FullyShardedDataParallel
from torch.nn.parallel import DistributedDataParallel


def get_global_group():
    """
    Singleton pytorch distributed group
    Inspired by https://github.com/pytorch/fairseq
    """
    if dist.is_initialized():
        # '''
        if not hasattr(get_global_group, "_global_group"):
            # ideally we could use torch.distributed.group.WORLD, but it seems
            # to cause random NCCL hangs in some cases
            get_global_group._global_group = dist.new_group()
        return get_global_group._global_group
        # '''
        # return dist.new_group()
    else:
        return None


class ModelType(Enum):
    DDP = 0
    FSDP = 1
    FSDP_OPT = 2


class Model(nn.Module):
    def __init__(self, model_type: ModelType):
        super().__init__()
        fc1 = nn.Linear(10, 32)
        fc2 = nn.Linear(32, 32)
        fc3 = nn.Linear(32, 10)
        if model_type != ModelType.DDP:
            self.fc1 = FullyShardedDataParallel(fc1, process_group=get_global_group())
            if model_type == ModelType.FSDP_OPT:
                self.fc2 = fc2
            else:
                self.fc2 = FullyShardedDataParallel(fc2, process_group=get_global_group())
            self.fc3 = FullyShardedDataParallel(fc3, process_group=get_global_group())
        else:
            self.fc1 = fc1
            self.fc2 = fc2
            self.fc3 = fc3
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return self.fc3(x)


def _worker(gpu_id: int, world_size: int, model_type_index: int):
    torch.cuda.set_device(gpu_id)
    dist.init_process_group(
        backend="nccl",
        init_method="tcp://127.0.0.1:9099",
        world_size=world_size,
        rank=gpu_id,
    )

    model_type = ModelType(model_type_index)

    batch_size = 250_000
    torch.manual_seed(0)
    input = torch.randn(size=(batch_size, 10)).cuda()
    target = torch.zeros(size=(batch_size, 10)).cuda()
    model = Model(model_type=model_type).cuda()
    if model_type == ModelType.DDP:
        model = DistributedDataParallel(model, device_ids=[gpu_id])
    else:
        model = FullyShardedDataParallel(model, process_group=get_global_group())
    optimizer = optim.SGD(model.parameters(), lr=1e-2)

    with torch.profiler.profile(schedule=torch.profiler.schedule(wait=2, warmup=3, active=3)) as prof:
        for _ in range(8):
            criterion = nn.MSELoss()
            with torch.autograd.profiler.record_function("forward"):
                out = model(input)
                loss = criterion(out, target)
            print(loss.item())
            with torch.autograd.profiler.record_function("zero_grad"):
                optimizer.zero_grad()
            with torch.autograd.profiler.record_function("backward"):
                loss.backward()
            with torch.autograd.profiler.record_function("optimizer"):
                optimizer.step()
            prof.step()

    if gpu_id == 0:
        prof.export_chrome_trace(f"trace_{model_type.name}_{gpu_id}.json")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-t", "--type", default=0, type=int, help="type of model (0: ddp, 1: fsdp, 2: fsdp_opt)")
    parser.add_argument('-g', '--gpus', default=2, type=int, help='number of gpus per node')
    args = parser.parse_args()
    mp.spawn(_worker, (args.gpus, args.type), nprocs=args.gpus)
```

Running the variant 1 will lead to no overlap:

<img width="1473" alt="Screenshot 2021-05-07 at 14 55 37" src="https://user-images.githubusercontent.com/7412790/117506280-551c5a80-af53-11eb-9177-35334cb5b104.png">

Running with the variant 2 will lead to overlap of the second linear layer compute with the layer3 AllGather:

<img width="1054" alt="Screenshot 2021-05-07 at 14 59 13" src="https://user-images.githubusercontent.com/7412790/117506342-71b89280-af53-11eb-9895-601b45e26413.png">

The new profiling graphs with dependencies are so nice!
@QuentinDuval Thanks for the great debugging!

@myleott, is there some mechanism that prevents the all_gather to go too far into the future ahead of the computation? If not, then when compute is super slow and/or networking is super fast, forward pass can run into OOM situation for very big models, right?

Resolving patch:
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1448,20 +1448,22 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         if params is None:
             params = self.params
         self.has_full_params = False
-        self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
-        with torch.cuda.stream(self._streams["all_gather"]):
-            for p in params:
-                if not p._is_sharded:  # e.g., world_size == 1
-                    if self.mixed_precision:
-                        self._free_fp16_param_shard([p])
-                    continue
-                # There may be external references to the Tensor Storage that we
-                # can't modify, such as references that are created by
-                # ctx.save_for_backward in the forward pass. Thus when we
-                # unshard parameters, we should reuse the original Tensor
-                # Storage object and unshard it in-place. For now, just resize
-                # the Storage to 0 to save memory.
-                free_storage_(p._full_param_padded)
+        current_stream = torch.cuda.current_stream()
+        for p in params:
+            if not p._is_sharded:  # e.g., world_size == 1
+                if self.mixed_precision:
+                    self._free_fp16_param_shard([p])
+                continue
+            # Don't let PyTorch reuse this memory until all work in the current
+            # stream is complete.
+            p._full_param_padded.record_stream(current_stream)
+            # There may be external references to the Tensor Storage that we
+            # can't modify, such as references that are created by
+            # ctx.save_for_backward in the forward pass. Thus when we
+            # unshard parameters, we should reuse the original Tensor
+            # Storage object and unshard it in-place. For now, just resize
+            # the Storage to 0 to save memory.
+            free_storage_(p._full_param_padded)
 
     @torch.no_grad()
     def _use_fp32_param_shard(self, params: Optional[List[Parameter]] = None) -> None:
@@ -1692,16 +1694,23 @@ def _get_default_cuda_device(module: nn.Module) -> torch.device:
 def cast_floats_to_right_precision(to_fp16: bool, no_grad: bool, *args: Any, **kwargs: Any) -> Tuple[Any, Any]:
     """
     Cast floating point Tensors in *args or **kwargs to FP16 or FP32 if they are not.
+    We also retain the requires_grad flag so that casting doesn't affect the autograd graph.
     """
 
     def fn_fp16(x: torch.Tensor) -> torch.Tensor:
         if x.dtype is torch.float32:
-            return x.half()
+            y = x.half()
+            if x.is_leaf:
+                y.requires_grad = x.requires_grad
+            return y
         return x
 
     def fn_fp32(x: torch.Tensor) -> torch.Tensor:
         if x.dtype is torch.float16:
-            return x.float()
+            y = x.float()
+            if x.is_leaf:
+                y.requires_grad = x.requires_grad
+            return y
         return x
 
     fn = fn_fp16 if to_fp16 else fn_fp32


Test patch:
diff --git a/fairscale/utils/testing.py b/fairscale/utils/testing.py
--- a/fairscale/utils/testing.py
+++ b/fairscale/utils/testing.py
@@ -34,6 +34,7 @@
 import multiprocessing
 import os
 import random
+from statistics import mean
 import subprocess
 import sys
 import tempfile
@@ -577,30 +578,35 @@ def forward(self, *input: Tensor, **kwargs: Any) -> Tensor:
 
 @functools.lru_cache()
 def get_cycles_per_ms() -> float:
-    """Approximate number of cycles per millisecond for torch.cuda._sleep
+    """Measure and return approximate number of cycles per millisecond for torch.cuda._sleep
 
     Copied from: github.com/pytorch/pytorch/blob/master/test/test_cuda.py
-
-    ..note::
-        This doesn't seems to return consistent cycles on desktop GPUs likely
-        due to frequency scaling.
-        >>> get_cycles_per_ms()
-        227.6441091140009
-        # new python process
-        >>> get_cycles_per_ms()
-        564.652154766248
-        # new python process
-        >>> get_cycles_per_ms()
-        245.56459442962856
     """
-    start = torch.cuda.Event(enable_timing=True)
-    end = torch.cuda.Event(enable_timing=True)
-    start.record()
-    torch.cuda._sleep(1000000)
-    end.record()
-    end.synchronize()
-    cycles_per_ms = 1000000 / start.elapsed_time(end)
-    return cycles_per_ms
+
+    def measure() -> float:
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        torch.cuda._sleep(1000000)
+        end.record()
+        end.synchronize()
+        cycles_per_ms = 1000000 / start.elapsed_time(end)
+        return cycles_per_ms
+
+    # Get 10 values and remove the 2 max and 2 min and return the avg.
+    # This is to avoid system disturbance that skew the results, e.g.
+    # the very first cuda call likely does a bunch of init, which takes
+    # much longer than subsequent calls.
+    #
+    # Tested on both Tesla V100, Quadro GP100, Titan RTX, RTX 3090 GPUs
+    # and seems to return stable values. Therefore, we enable caching
+    # using lru_cache decorator above.
+    num = 10
+    vals = []
+    for _ in range(num):
+        vals.append(measure())
+    vals = sorted(vals)
+    return mean(vals[2 : num - 2])
 
 
 class DummyProcessGroup:
@@ -681,3 +687,15 @@ def dump_all_tensors(rank: int) -> None:
         except Exception as e:
             pass
     print(torch.cuda.memory_summary())
+
+
+def get_smi_memory() -> float:
+    """Return process's GPU memory in MB."""
+    pid = os.getpid()
+    info_string = torch.cuda.list_gpu_processes()
+    for line in info_string.splitlines():
+        if str(pid) in line:
+            toks = line.split()
+            return float(toks[3])
+    # If the process is not in the list, we are not using the GPU.
+    return 0.0
diff --git a/tests/ci_test_list_2.txt b/tests/ci_test_list_2.txt
--- a/tests/ci_test_list_2.txt
+++ b/tests/ci_test_list_2.txt
@@ -1,5 +1,6 @@
 tests/nn/misc/test_checkpoint_activations.py
 tests/nn/misc/test_checkpoint_activations_norm.py
+tests/nn/data_parallel/test_fsdp_overlap.py
 tests/nn/data_parallel/test_fsdp_multiple_forward.py
 tests/nn/data_parallel/test_fsdp_apply.py
 tests/nn/data_parallel/test_fsdp_state_dict.py
diff --git a/tests/nn/data_parallel/test_fsdp_overlap.py b/tests/nn/data_parallel/test_fsdp_overlap.py
new file mode 100644
--- /dev/null
+++ b/tests/nn/data_parallel/test_fsdp_overlap.py
@@ -0,0 +1,245 @@
+# 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
+
+""" Test FSDP and ensure expected overlapping between all_gather and forward. """
+
+from statistics import mean
+import time
+from unittest.mock import patch
+
+import pytest
+import torch
+from torch.cuda import Event
+import torch.multiprocessing as mp
+import torch.nn as nn
+
+from fairscale.nn import enable_wrap, wrap
+from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
+from fairscale.utils.testing import (
+    dist_init,
+    get_cycles_per_ms,
+    skip_if_single_gpu,
+    teardown,
+    temp_files_ctx,
+    torch_version,
+)
+
+
+class Layer(nn.Module):
+    def __init__(self, compute_cycles, has_params: bool):
+        super().__init__()
+        self.sleep_cycles = compute_cycles
+        self.optional_param = None
+        if has_params:
+            self.optional_param = nn.Parameter(torch.rand(1))
+
+    def forward(self, x):
+        # Get 2 events.
+        self.e1 = Event(enable_timing=True)
+        self.e2 = Event(enable_timing=True)
+
+        # Record the fake forward compute time.
+        self.e1.record()
+        if self.sleep_cycles > 0:
+            torch.cuda._sleep(self.sleep_cycles)
+        if self.optional_param is not None:
+            x = x + self.optional_param  # force the param to be part of the graph
+        self.e2.record()
+        return x
+
+    def get_time(self):
+        # return the recorded duration.
+        return self.e1.elapsed_time(self.e2)
+
+
+def _create_model(fsdp_config, compute_cycles, has_params: bool):
+    with enable_wrap(wrapper_cls=FSDP, **fsdp_config):
+        model = wrap(
+            nn.Sequential(
+                wrap(Layer(compute_cycles, has_params)),
+                wrap(Layer(compute_cycles, has_params)),
+                wrap(Layer(compute_cycles, has_params)),
+                wrap(Layer(compute_cycles, has_params)),
+            )
+        ).cuda()
+    return model
+
+
+class Min10:
+    def __init__(self):
+        self.data = []
+
+    def add(self, new_data):
+        if len(self.data) < 10:
+            self.data.append(new_data)
+        else:
+            self.data = sorted(self.data)
+            if new_data < self.data[-1]:
+                self.data[-1] = new_data
+
+    def avg(self):
+        return mean(self.data)
+
+
+def _distributed_worker(
+    gpu_id, world_size, fsdp_config, tempfile, tempfile_rpc,
+):
+    torch.cuda.set_device(gpu_id)
+
+    rank = gpu_id
+    result = dist_init(rank, world_size, tempfile, tempfile_rpc)
+    assert result, "Dist init failed"
+
+    # Save the original torch.distributed.all_gather function since we will
+    # patch it to include an artificial delay.
+    orig_all_gather = torch.distributed.all_gather
+
+    def run(compute_cycles, all_gather_cycles):
+        has_params = all_gather_cycles > 0
+        model = _create_model(fsdp_config, compute_cycles, has_params)
+
+        # Get the input and sets the input's requires_grad to True because
+        # we have a fake compute in the forward pass.
+        batch = torch.rand(1).cuda()
+        batch.requires_grad = True
+
+        # 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):
+            # Get two events for measuring the overall time.
+            e1 = Event(enable_timing=True)
+            e2 = Event(enable_timing=True)
+
+            cpu_start = time.process_time()
+
+            all_gather_called = False
+
+            def _delayed_all_gather(*args, **kwargs):
+                nonlocal all_gather_called
+                all_gather_called = True
+                torch.cuda._sleep(all_gather_cycles)
+                return orig_all_gather(*args, **kwargs)
+
+            # forward pass
+            #
+            # Even though both e1 & e2 are on the compute stream, since
+            # compute depends on all_gather, e2-e1 includes all_gather time.
+            e1.record()
+            with patch("torch.distributed.all_gather", _delayed_all_gather):
+                out = model(batch)
+                if has_params and world_size > 1:
+                    assert all_gather_called
+                else:
+                    assert not all_gather_called
+            e2.record()
+
+            # backward pass
+            out.backward()
+            if torch_version() >= (1, 7, 0):
+                model.zero_grad(set_to_none=True)
+            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 = []
+            for mod in model.modules():
+                if not isinstance(mod, Layer):
+                    continue
+                times.append(mod.get_time())
+
+            # 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(),
+        }
+
+    sleep_cycles = int(100 * get_cycles_per_ms())
+
+    e1 = run(0, 0)  # no compute, no all-gather
+    e2 = run(0, sleep_cycles)  # no compute, only all-gather
+    e3 = run(sleep_cycles, 0)  # only compute, no all-gather
+    e4 = run(sleep_cycles, sleep_cycles)  # both compute and all-gather
+    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:
+            # 10X longer is a safe margin, since the GPU work timing is around 100X more
+            # of that of the CPU.
+            assert s * 10 < l, f"{s} * 10 < {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.
+    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.
+            assert s * 10 < l, f"{s} * 10 < {l} in " + debug_string
+
+    # Check the GPU overlapping when there is all-gather.
+    if world_size > 1:
+        compute_only = e3["gpu_compute"]
+        all_gather_only = e2["gpu_total"]
+        both = e4["gpu_total"]
+        assert compute_only + all_gather_only > 1.1 * both, (
+            f"{compute_only} + {all_gather_only} > 1.1 * {both} in " + debug_string
+        )
+
+    teardown()
+
+
+@skip_if_single_gpu
+@pytest.mark.parametrize("world_size", [1, 2])
+@pytest.mark.parametrize("flatten", ["flatten", "no_flatten"])
+@pytest.mark.parametrize("mixed", ["mixed", "full"])
+def test_forward_overlap(world_size, flatten, mixed):
+    fsdp_config = {
+        "flatten_parameters": flatten == "flatten",
+        "mixed_precision": mixed == "mixed",
+    }
+    with temp_files_ctx(2) as temp_files:
+        mp.spawn(
+            _distributed_worker, (world_size, fsdp_config, temp_files[0], temp_files[1]), nprocs=world_size,
+        )