test_fsdp_overlap.py

# 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,
        )