# Owner(s): ["oncall: distributed"]

# 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.

import copy
import os
import sys
import unittest
from contextlib import suppress
from typing import Any, List, cast

import numpy as np

import torch
import torch.distributed as dist

if not dist.is_available():
    print("Distributed not available, skipping tests", file=sys.stderr)
    sys.exit(0)
from torch.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook import (
    hook_with_zero_step,
    hook_with_zero_step_interleaved,
)
from torch.distributed.algorithms.ddp_comm_hooks.default_hooks import (
    allreduce_hook,
)
from torch.distributed.algorithms.join import Join, Joinable, JoinHook
from torch.distributed.optim import ZeroRedundancyOptimizer
from torch.distributed.optim.zero_redundancy_optimizer import _broadcast_object
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import SGD, AdamW
from torch.testing._internal import common_distributed
from torch.testing._internal.common_utils import (
    IS_WINDOWS,
    TEST_WITH_ASAN,
    TEST_WITH_DEV_DBG_ASAN,
    instantiate_parametrized_tests,
    parametrize,
    run_tests,
)

try:
    import torchvision
    HAS_TORCHVISION = True
except ImportError:
    HAS_TORCHVISION = False

# Use GLOO on GPU when running CUDA + Windows
def _get_backend_for_tests():
    return (
        dist.Backend.NCCL if not IS_WINDOWS and torch.cuda.is_available()
        # Windows only has GLOO, but GLOO GPU works. And use GLOO CPU when
        # no GPUs are available.
        else dist.Backend.GLOO
    )

BACKEND = _get_backend_for_tests()

@unittest.skipIf(
    TEST_WITH_ASAN or TEST_WITH_DEV_DBG_ASAN, "CUDA + ASAN does not work."
)
class TestZeroRedundancyOptimizer(common_distributed.MultiProcessTestCase):
    def setUp(self):
        super(TestZeroRedundancyOptimizer, self).setUp()
        os.environ["WORLD_SIZE"] = str(self.world_size)
        self._spawn_processes()

    @property
    def device(self):
        return torch.device("cuda") if torch.cuda.is_available() \
            else torch.device("cpu")

    @property
    def world_size(self):
        return 1

    def tearDown(self):
        try:
            torch.distributed.destroy_process_group()
        except AssertionError:
            pass
        try:
            os.remove(self.file_name)
        except OSError:
            pass

    def dist_init(self, rank, world_size=-1, backend=BACKEND):
        if (world_size < 1):
            world_size = self.world_size
        store = dist.FileStore(self.file_name, world_size)
        return dist.init_process_group(
            backend=backend, store=store, rank=rank, world_size=world_size,
        )


# TODO: sandcastle_skip_if does not work here.
@unittest.skipIf(
    TEST_WITH_ASAN or TEST_WITH_DEV_DBG_ASAN, "CUDA + ASAN does not work."
)
class TestZeroRedundancyOptimizerSingleRank(TestZeroRedundancyOptimizer):
    def test_state_dict(self):
        """Check that ZeroRedundancyOptimizer exposes the expected state dict
        interface, irrespective of the sharding."""
        self.dist_init(self.rank)
        LR1 = 0.1
        LR2 = 0.01
        MOMENTUM = 0.9
        RECIPIENT_RANK = 0  # rank 0 is the only rank since the world size is 1
        x = torch.tensor([1.0], device=self.device, requires_grad=True)
        o = ZeroRedundancyOptimizer(
            [x], optimizer_class=SGD, lr=LR1, momentum=MOMENTUM,
        )
        x.backward()
        o.step()
        self.assertEqual(x, torch.tensor([0.9], device=self.device))
        self.assertEqual(
            o.optim.state[x]["momentum_buffer"],
            torch.tensor([1.0], device=self.device),
        )

        o.zero_grad()
        o.consolidate_state_dict(to=RECIPIENT_RANK)
        state_dict = o.state_dict()

        # Check that the state dict has keys compliant with PyTorch
        self.assertIn("param_groups", state_dict.keys())
        self.assertIn("state", state_dict.keys())

        # Check that the state has the expected keys
        self.assertEqual(state_dict["param_groups"][0]["lr"], 0.1)
        self.assertEqual(state_dict["param_groups"][0]["momentum"], 0.9)
        self.assertFalse(state_dict["param_groups"][0]["nesterov"])
        self.assertEqual(state_dict["param_groups"][0]["weight_decay"], 0.0)
        self.assertEqual(state_dict["param_groups"][0]["dampening"], 0.0)

        # Check that the state and the `param_groups` attribute are in sync
        for k in state_dict["param_groups"][0]:
            if k != "params":
                self.assertEqual(
                    state_dict["param_groups"][0][k],
                    o.param_groups[0][k],
                )

        # Check that the state is reloaded with the correct values and device
        o = ZeroRedundancyOptimizer([x], optimizer_class=SGD, lr=LR2)
        o.load_state_dict(state_dict)
        self.assertEqual(
            o.optim.state[x]["momentum_buffer"],
            torch.tensor([1.0], device=self.device),
        )

        # We should we using `LR1` and not `LR2` after reloading, both within
        # the optimizer and as exposed by the `param_groups` attribute
        self.assertEqual(o.param_groups[0]["lr"], LR1)
        x.backward()
        o.step()
        self.assertEqual(x, torch.tensor([0.71], device=self.device))
        self.assertEqual(
            o.optim.state[x]["momentum_buffer"],
            torch.tensor([1.9], device=self.device),
        )

        # Check that the exposed `param_groups`` are on the proper device
        self.assertEqual(o.param_groups[0]["params"][0].device, x.device)

    def test_lr_scheduler(self):
        """Check that a normal PyTorch ``lr_scheduler`` is usable with
        ZeroRedundancyOptimizer."""
        self.dist_init(self.rank)
        NUM_ITERS = 5
        LR = 0.01
        x = torch.tensor([1.0], device=self.device, requires_grad=True)
        x2 = torch.tensor([1.0], device=self.device, requires_grad=True)
        o = ZeroRedundancyOptimizer([x], optimizer_class=SGD, lr=LR)
        o2 = torch.optim.SGD([x2], lr=LR)
        s = torch.optim.lr_scheduler.StepLR(o, 1)
        s2 = torch.optim.lr_scheduler.StepLR(o2, 1)
        for _ in range(NUM_ITERS):
            x.backward()
            o.zero_grad()
            o.step()
            s.step()
            x2.backward()
            o2.zero_grad()
            o2.step()
            s2.step()
            self.assertEqual(x, x2)

    def test_step_with_kwargs(self):
        """Check that the ``step(**kwargs)`` interface is properly exposed."""
        self.dist_init(self.rank)
        LR = 0.1

        class SGDWithStepKWArg(torch.optim.SGD):
            def step(self, closure=None, kwarg=None):
                super().step()
                kwarg.append(5)

        kwarg: List[Any] = []
        x = torch.tensor([1.0], device=self.device, requires_grad=True)
        o = ZeroRedundancyOptimizer(
            [x], optimizer_class=SGDWithStepKWArg, lr=LR,
        )
        x.backward()
        o.step(0, kwarg=kwarg)
        self.assertEqual(kwarg, [5])
        self.assertEqual(x, torch.tensor([0.9], device=self.device))

    def test_step_with_extra_inner_key(self):
        """Check that ZeroRedundancyOptimizer wrapping an optimizer that adds
        extra keys to ``param_groups`` exposes those keys through ZeRO's own
        ``param_groups``."""
        self.dist_init(self.rank)
        LR = 0.1

        class SGDWithNewKey(torch.optim.SGD):
            # Dummy optimizer which adds a new key to the param groups
            def step(self, closure=None):
                super().step()
                self.param_groups[0]["new_key"] = 0.1

        x = torch.tensor([1.0], device=self.device, requires_grad=True)
        o = ZeroRedundancyOptimizer([x], optimizer_class=SGDWithNewKey, lr=LR)
        x.backward()
        o.step()
        self.assertEqual(o.param_groups[0]["new_key"], 0.1)
        self.assertEqual(x, torch.tensor([0.9], device=self.device))

    def test_step_without_closure(self):
        """Check that the ``step()`` method (without closure) is handled as
        expected."""
        self.dist_init(self.rank)
        LR = 0.1

        class SGDWithoutClosure(torch.optim.SGD):
            def step(self):
                return super().step()

        x = torch.tensor([1.0], device=self.device, requires_grad=True)
        o = ZeroRedundancyOptimizer(
            [x], optimizer_class=SGDWithoutClosure, lr=LR,
        )
        x.backward()
        o.step()
        self.assertEqual(x, torch.tensor([0.9], device=self.device))

    def test_zero_grad(self):
        """Check that the ``zero_grad`` method is properly handled."""
        self.dist_init(self.rank)
        LR = 0.01
        x = torch.rand(1)
        m = torch.nn.Linear(1, 1)
        o = ZeroRedundancyOptimizer(m.parameters(), optimizer_class=SGD, lr=LR)
        y = m(x)
        y.backward(x)
        self.assertNotEqual(m.weight.grad, torch.zeros_like(m.weight))
        self.assertNotEqual(m.weight.grad, torch.zeros_like(m.weight))
        o.zero_grad()
        self.assertFalse(m.weight.grad)
        self.assertFalse(m.bias.grad)

    def test_constructor(self):
        """Check the robustness of the ZeroRedundancyOptimizer constructor by
        passing different values for the ``params`` argument."""
        self.dist_init(self.rank)
        LR = 0.01
        m = torch.nn.Sequential(
            torch.nn.Linear(5, 10),
            torch.nn.Linear(10, 10),
            torch.nn.Linear(10, 10),
        )
        # Test various constructor inputs in the form: (input, expected error)
        ctor_inputs = [
            ([], ValueError),                          # empty parameter list
            (torch.randn(1), TypeError),               # non-iterable: `torch.Tensor`
            (1.2, TypeError),                          # non-iterable: `float`
            ([
                {"params": [l.weight for l in m]},
                {"params": [l.bias for l in m]},
            ], None),                                  # iterable of dict
            (list(m.parameters()) + [42], TypeError),  # iterable containing invalid type
            (m.parameters(), None),                    # `params` as a generator
            (list(m.parameters()), None)               # `params` as a list
        ]
        for ctor_input, error in ctor_inputs:
            context = self.assertRaises(error) if error else suppress()
            with context:
                ZeroRedundancyOptimizer(
                    ctor_input, optimizer_class=SGD, lr=LR,
                )

        # Test constructing with multiple parameter groups more thoroughly
        WD = 0.01
        BETAS = (0.9, 0.999)
        EPS = 1e-8
        params = [
            {"params": [l.weight for l in m], "weight_decay": 0.},
            {"params": [l.bias for l in m], "weight_decay": WD},
        ]
        o = ZeroRedundancyOptimizer(
            params, optimizer_class=AdamW,
            lr=LR, betas=BETAS, eps=EPS,
        )
        assert len(o.param_groups) == 2, \
            f"Expected 2 ZeRO param groups, but got {len(o.param_groups)}"
        assert len(o.optim.param_groups) == 2, \
            "Expected 2 local optimizer param groups, but got " \
            f"{len(o.optim.param_groups)}"

    def test_same_dense_param_type(self):
        """Check that ZeroRedundancyOptimizer raises an exception if the input
        parameters include sparse tensors or different dense types.

        NOTE: This test should be removed once support for sparse parameters
        and varying parameter types is added.
        """
        self.dist_init(self.rank)
        LR = 0.01
        inputs = [
            [torch.sparse_coo_tensor(size=(2, 3))],
            [torch.FloatTensor(1), torch.DoubleTensor(1)],
            [torch.FloatTensor(1), torch.FloatTensor(1),
                torch.sparse_coo_tensor(size=(2, 3))]
        ]
        for input in inputs:
            with self.assertRaises(ValueError):
                ZeroRedundancyOptimizer(input, optimizer_class=SGD, lr=LR)


class TestZeroRedundancyOptimizerDistributed(TestZeroRedundancyOptimizer):
    @property
    def device(self):
        return torch.device(self.rank) if torch.cuda.is_available() \
            else torch.device("cpu")

    @property
    def world_size(self):
        return min(4, max(2, torch.cuda.device_count()))

    @property
    def context(self):
        return suppress() if not torch.cuda.is_available() \
            else torch.cuda.device(self.rank)

    def _check_same_model_params(
        self,
        model_a: torch.nn.Module,
        model_b: torch.nn.Module,
        message: str = "",
    ) -> None:
        # Check that model parameters match
        for p_a, p_b in zip(model_a.parameters(), model_b.parameters()):
            torch.testing.assert_close(
                p_a, p_b, atol=1e-3, rtol=1e-5,
                msg=f"Model parameters differ:\n{p_a} {p_b}\n" + message,
            )
        # Check that model buffers match
        for b_a, b_b in zip(model_a.buffers(), model_b.buffers()):
            torch.testing.assert_close(
                b_a, b_b,
                msg=f"Model buffers differ:\n{b_a} {b_b}\n" + message,
            )

    @common_distributed.skip_if_no_gpu
    @common_distributed.skip_if_rocm
    def test_step(self):
        """Check that ZeroRedundancyOptimizer properly exposes the ``step()``
        interface."""
        self.dist_init(self.rank, world_size=self.world_size)
        LR = 0.01

        with self.context:
            x = torch.tensor([float(self.rank + 1)], device=self.device)
            m = torch.nn.Linear(1, 1)
            m.weight.data = torch.tensor([[1.0]])
            m.bias.data = torch.tensor([2.0])
            m = m.to(self.device)
            m_zero = copy.deepcopy(m).to(self.device)

            o = SGD(m.parameters(), lr=LR)
            o_zero = ZeroRedundancyOptimizer(
                m_zero.parameters(), optimizer_class=SGD, lr=LR,
            )

            y = m(x)
            y.backward(x)
            y_zero = m_zero(x)
            y_zero.backward(x)

            for p in m.parameters():
                dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
                p.grad.data /= self.world_size
            o.step()
            for p in m_zero.parameters():
                dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
                p.grad.data /= self.world_size
            o_zero.step()

            self.assertEqual(m.weight, m_zero.weight)
            self.assertEqual(m.bias, m_zero.bias)

    @common_distributed.skip_if_no_gpu
    @common_distributed.skip_if_rocm
    def test_step_with_closure(self):
        """Check that ZeroRedundancyOptimizer properly exposes the
        ``step(closure)`` interface."""
        self.dist_init(self.rank, world_size=self.world_size)

        with self.context:
            for bucket_view in [False, True]:
                x_val = self.rank + 1
                weight = 1.0
                bias = 2.0
                error = 1.0
                target = torch.tensor(
                    [x_val * weight + bias + error],
                    device=self.device,
                )
                loss_fn = torch.nn.L1Loss()

                x = torch.tensor([float(x_val)], device=self.device)
                m = torch.nn.Linear(1, 1)
                m.weight.data = torch.tensor([[weight]])
                m.bias.data = torch.tensor([bias])
                m.to(self.device)

                o = ZeroRedundancyOptimizer(
                    m.parameters(),
                    optimizer_class=SGD,
                    parameters_as_bucket_view=bucket_view,
                    lr=0.1,
                )

                y = m(x)
                y.backward(x)
                for p in m.parameters():
                    dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
                    p.grad.data /= self.world_size

                def closure():
                    o.zero_grad()
                    output = m(x)
                    loss = loss_fn(output, target)
                    loss.backward()
                    return loss

                loss = o.step(closure=closure)

                self.assertEqual(loss, torch.tensor(error))
                self.assertEqual(m.weight, torch.tensor([[1.1]]))
                self.assertEqual(m.bias, torch.tensor([2.1]))

    @common_distributed.skip_if_no_gpu
    def test_lr_scheduler(self):
        """Check that a normal PyTorch ``lr_scheduler`` is usable with
        ZeroRedundancyOptimizer."""
        self.dist_init(self.rank)
        x = torch.tensor([1.0], device=self.device, requires_grad=True)
        x2 = torch.tensor([1.0], device=self.device, requires_grad=True)
        o = ZeroRedundancyOptimizer([x], optimizer_class=SGD, lr=0.01)
        o2 = torch.optim.SGD([x2], lr=0.01)
        s = torch.optim.lr_scheduler.StepLR(o, 1)
        s2 = torch.optim.lr_scheduler.StepLR(o2, 1)
        for _ in range(5):
            x.backward()
            o.zero_grad()
            o.step()
            s.step()
            x2.backward()
            o2.zero_grad()
            o2.step()
            s2.step()
            self.assertEqual(x, x2)

    def test_sharding(self):
        """
        Check ZeroRedundancyOptimizer's parameter sharding at construction
        time.

        NOTE: The correctness of this test depends on the ZeRO implementation
        using the sorted-greedy partitioning algorithm. For details, see
        ``ZeroRedundancyOptimizer._partition_parameters()`` in
        zero_redundancy_optimizer.py.
        """
        self.dist_init(self.rank)
        LR = 0.01
        sizes = [9, 7, 5, 3]
        params = []
        for size in sizes * self.world_size:
            params.append(torch.rand(size, 1))
        o = ZeroRedundancyOptimizer(params, optimizer_class=SGD, lr=LR)
        self.assertEqual(
            sum([x.numel() for x in o.optim.param_groups[0]["params"]]),
            sum(sizes),
        )

    def test_add_param_group(self):
        """Check that ZeroRedundancyOptimizer properly handles adding a new
        parameter group a posteriori and that all ranks get a shard of the
        contained parameters.

        NOTE: The correctness of this test depends on the ZeRO implementation
        using the sorted-greedy partitioning algorithm. For details, see
        ``ZeroRedundancyOptimizer._partition_parameters()`` in
        zero_redundancy_optimizer.py.
        """
        self.dist_init(self.rank)
        LR = 0.01

        # Test with all parameters trainable to begin with
        def all_trainable():
            params = []
            sizes = [9, 7, 5, 3]
            sizes_world = sizes * self.world_size
            for size in sizes_world[:-1]:
                params.append(torch.rand(size, 1))

            # Make sure that the params are trainable so that they are factored
            # into the size-based parameter partitioning
            for p in params:
                p.requires_grad = True

            o = ZeroRedundancyOptimizer(params, optimizer_class=SGD, lr=LR)
            self.assertEqual(len(o.param_groups), 1)
            o.add_param_group({"params": [torch.rand(3, 1)]})
            # Verify that new group is added to the correct partition, making
            # all partitions have the same elements
            self.assertEqual(len(o.param_groups), 2)
            self.assertEqual(
                sum([
                    x.numel()
                    for g in o.optim.param_groups
                    for x in g["params"]
                ]),
                sum(sizes),
            )
            self.assertEqual(len(o.optim.param_groups), 2)

        # Test a pathological config with a first big non-trainable param
        def some_trainable():
            params = []
            for size in [100, 3, 5, 2, 6, 4]:
                params.append(torch.rand(size, 1))

            # Make sure that all but the first param are trainable so that they
            # are factored into the size-based parameter partitioning
            for p in params[1:]:
                p.requires_grad = True

            o = ZeroRedundancyOptimizer(params, optimizer_class=SGD, lr=LR)
            self.assertEqual(len(o.param_groups), 1)
            o.add_param_group({"params": [torch.rand(3, 1)]})
            self.assertEqual(len(o.param_groups), 2)
            self.assertEqual(len(o.optim.param_groups), 2)

        all_trainable()
        some_trainable()

    @common_distributed.skip_if_no_gpu
    def test_multiple_param_groups(self):
        """
        Check parity between constructing ZeRO with multiple parameter groups
        upfront versus adding parameter groups to ZeRO after construction
        versus a non-sharded optimizer.
        """
        self.dist_init(self.rank)
        BATCH_SIZE, NUM_ITERS = 8, 3
        INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM = 5, 10, 5
        WD, LR = 0.01, 0.01
        model1 = torch.nn.Sequential(
            torch.nn.Linear(INPUT_DIM, HIDDEN_DIM),
            torch.nn.Linear(HIDDEN_DIM, HIDDEN_DIM),
            torch.nn.Linear(HIDDEN_DIM, OUTPUT_DIM),
        )
        model2 = copy.deepcopy(model1)
        model3 = copy.deepcopy(model1)
        model1 = model1.to(self.device)
        model2 = model2.to(self.device)
        model3 = model3.to(self.device)
        inputs = [
            torch.randn(BATCH_SIZE, INPUT_DIM).to(self.device)
            for _ in range(NUM_ITERS)
        ]
        # Construct `optim1` with both parameter groups upfront
        optim1 = ZeroRedundancyOptimizer(
            [
                {"params": [l.weight for l in model1], "weight_decay": 0.},
                {"params": [l.bias for l in model1], "weight_decay": WD},
            ],
            optimizer_class=AdamW, lr=LR,
        )
        # Construct `optim2` by adding the second parameter after
        optim2 = ZeroRedundancyOptimizer(
            [l.weight for l in model2],
            optimizer_class=AdamW, lr=LR, weight_decay=0.,
        )
        optim2.add_param_group(
            {"params": [l.bias for l in model2], "weight_decay": WD}
        )
        # Construct `optim3` as a non-sharded optimizer
        optim3 = AdamW(
            [
                {"params": [l.weight for l in model3], "weight_decay": 0.},
                {"params": [l.bias for l in model3], "weight_decay": WD},
            ], lr=LR,
        )
        # Check parity over a few iterations
        for input in inputs:
            for model, optim in (
                (model1, optim1), (model2, optim2), (model3, optim3),
            ):
                optim.zero_grad()
                out = model(input)
                loss = out.sum()
                loss.backward()
                optim.step()
            for layer1, layer2, layer3 in zip(model1, model2, model3):
                torch.testing.assert_close(layer1.weight, layer2.weight)
                torch.testing.assert_close(layer1.weight, layer3.weight)
                torch.testing.assert_close(layer1.bias, layer2.bias)
                torch.testing.assert_close(layer1.bias, layer3.bias)

    @common_distributed.skip_if_no_gpu
    @common_distributed.skip_if_rocm
    def test_collect_shards(self):
        """Check the state consolidation mechanism and the state dict exposed
        by ZeroRedundancyOptimizer."""
        self.dist_init(self.rank)
        LR = 1e-3
        MOMENTUM = 0.99
        BATCH_SIZE, INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM = 3, 20, 10, 5
        REFERENCE_RANK = 0
        target = torch.rand((BATCH_SIZE, OUTPUT_DIM), device=self.device)
        inputs = torch.rand((BATCH_SIZE, INPUT_DIM), device=self.device)
        model = torch.nn.Sequential(
            torch.nn.Linear(INPUT_DIM, HIDDEN_DIM),
            torch.nn.Linear(HIDDEN_DIM, OUTPUT_DIM),
        ).to(self.device)
        loss_fn = torch.nn.L1Loss()
        loss_fn.to(self.device)
        optimizer = ZeroRedundancyOptimizer(
            model.parameters(),
            optimizer_class=SGD,
            lr=LR,
            momentum=MOMENTUM,  # ensure there exists state to shard
        )

        def closure():
            optimizer.zero_grad()
            output = model(inputs)
            loss = loss_fn(output, target)
            loss.backward()
            return loss

        # Run a dummy step so that the optimizer state dict exists
        _ = optimizer.step(closure=closure)

        # Get the optimizer state on the reference rank
        optimizer.consolidate_state_dict(to=REFERENCE_RANK)
        if self.rank == REFERENCE_RANK:
            # Check that the state has the correct size
            optimizer_state_dict = optimizer.state_dict()
            self.assertEqual(
                len(optimizer_state_dict["state"]),
                len(list(model.parameters())),
            )
        else:
            optimizer_state_dict = {}

        # Load the optimizer state on all ranks without any exceptions
        optimizer_state_dict = _broadcast_object(
            optimizer_state_dict,
            src_rank=REFERENCE_RANK,
            group=dist.group.WORLD,
            device=self.device,
        )
        optimizer.load_state_dict(optimizer_state_dict)

    def test_nondefault_process_group(self):
        """Check that ZeroRedundancyOptimizer works with a non-default process
        group consisting only of even ranks."""
        # Skip the test if below the minimum world size since then the test is
        # trivial
        MIN_WORLD_SIZE = 4
        if self.world_size < MIN_WORLD_SIZE:
            common_distributed.logger.info(
                "Skipping `test_nondefault_process_group()` since world size "
                f"of {self.world_size} is less than {MIN_WORLD_SIZE}"
            )
            return
        BACKEND = dist.Backend.GLOO
        self.dist_init(self.rank, self.world_size, BACKEND)
        # Use GPU if enough are available, or fall back to CPU otherwise, which
        # is fine since Gloo backend supports both
        if torch.cuda.is_available() and \
                torch.cuda.device_count() >= self.world_size:
            device = torch.device(self.rank)
        else:
            device = torch.device("cpu")
        # Create a new process group consisting of the even ranks to exercise
        # the case where the global and local ranks do not necessarily match
        subgroup_ranks = [r for r in range(self.world_size) if r % 2 == 0]
        process_group = dist.new_group(
            ranks=subgroup_ranks, backend=BACKEND,
        )
        # Ranks not participating in the new process group are no longer needed
        if self.rank not in subgroup_ranks:
            return

        # Set different seeds across ranks so that each rank gets different
        # training data and hence the model sync check is meaningful
        torch.manual_seed(self.rank)
        np.random.seed(self.rank)

        EPOCHS, BATCH_SIZE, INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM = 5, 3, 20, 10, 5
        LR = 1e-3
        MOMENTUM = 0.99
        REFERENCE_RANK = 0
        assert REFERENCE_RANK in subgroup_ranks, \
            "Reference rank must be in the new process group"
        loss_fn = torch.nn.L1Loss().to(device)

        def check(optimizer):
            for _ in range(EPOCHS):
                target = torch.rand((BATCH_SIZE, OUTPUT_DIM), device=device)
                inputs = torch.rand((BATCH_SIZE, INPUT_DIM), device=device)

                def closure():
                    optimizer.zero_grad()
                    output = model(inputs)
                    loss = loss_fn(output, target)
                    loss /= self.world_size
                    loss.backward()
                    dist.all_reduce(loss, group=process_group)
                    return loss

                _ = optimizer.step(closure=closure)

                # Check that the parameters match across ranks after a step
                for pg in optimizer.param_groups:
                    for p in pg["params"]:
                        receptacle = [
                            p.clone() for _ in subgroup_ranks
                        ] if self.rank == REFERENCE_RANK else []
                        dist.gather(
                            p, receptacle, dst=REFERENCE_RANK,
                            group=process_group,
                        )
                        if self.rank == REFERENCE_RANK:
                            reference_param = receptacle[0]
                            for param in receptacle[1:]:
                                torch.testing.assert_close(
                                    reference_param,
                                    param,
                                    msg="Models differ between ranks",
                                )

        model = torch.nn.Sequential(
            torch.nn.Linear(INPUT_DIM, HIDDEN_DIM),
            torch.nn.Linear(HIDDEN_DIM, OUTPUT_DIM),
        ).to(device)
        optimizer = ZeroRedundancyOptimizer(
            model.parameters(),
            optimizer_class=SGD,
            lr=LR,
            momentum=MOMENTUM,  # ensure there exists state to shard
            process_group=process_group,
        )
        check(optimizer)

    @common_distributed.skip_if_no_gpu
    @parametrize(
        "optimizer_class_str",
        ["Adam", "AdamW", "SGD"],
        # Use string to appease the internal test name parser
    )
    @parametrize(
        "maximize",
        [False, True],
    )
    def test_local_optimizer_parity(
        self,
        optimizer_class_str: str,
        maximize: bool,
    ):
        """When combined with DDP, check that a local optimizer gives the same
        results as wrapping that optimizer with ZeroRedundancyOptimizer."""
        self.dist_init(self.rank)
        BATCHES = 20
        BATCH_SIZE = 64
        LR = 1e-3
        INPUT_DIM = 2
        HIDDEN_DIM = 3
        OUTPUT_DIM = 3
        torch.manual_seed(self.rank)
        np.random.seed(self.rank)
        if optimizer_class_str == "Adam":
            optimizer_class = torch.optim.Adam
        elif optimizer_class_str == "AdamW":
            optimizer_class = torch.optim.AdamW
        elif optimizer_class_str == "SGD":
            optimizer_class = torch.optim.SGD
        else:
            assert 0, f"Unsupported optimizer class: {optimizer_class_str}"

        with self.context:
            # Define a base model with a different buffer for each rank
            model = torch.nn.Sequential(
                torch.nn.Linear(INPUT_DIM, HIDDEN_DIM),
                torch.nn.Linear(HIDDEN_DIM, HIDDEN_DIM),
                torch.nn.Linear(HIDDEN_DIM, OUTPUT_DIM),
            ).to(self.device)
            model.register_buffer(
                "test_buffer", torch.ones((1), device=self.device) * self.rank,
            )
            # Define models/optimizers for DDP with ZeRO and DDP with local
            # optimizer
            defaults = {"maximize": True} if maximize else {}
            sharded_optimizer = ZeroRedundancyOptimizer(
                params=model.parameters(), optimizer_class=optimizer_class,
                lr=LR, **defaults,
            )
            sharded_ddp_model = DDP(
                module=model, device_ids=[self.rank],
                broadcast_buffers=True, find_unused_parameters=True,
            )
            local_model = copy.deepcopy(model).to(self.device)
            ddp_optimizer = optimizer_class(
                local_model.parameters(), lr=LR, **defaults,
            )
            ddp_model = DDP(
                local_model, device_ids=[self.rank],
                broadcast_buffers=True, find_unused_parameters=True,
            )
            # Check that the model is properly synchronized between ranks
            # at construction time
            self._check_same_model_params(
                sharded_ddp_model, ddp_model,
                "Models differ from the start",
            )

            def check_step():
                input_tensor = torch.rand((BATCH_SIZE, INPUT_DIM))

                def closure_ddp(input_tensor=input_tensor):
                    ddp_optimizer.zero_grad()
                    ddp_loss = ddp_model(input_tensor).abs().sum()
                    ddp_loss.backward()
                    return ddp_loss

                def closure_sharded(input_tensor=input_tensor):
                    sharded_optimizer.zero_grad()
                    sharded_loss = sharded_ddp_model(input_tensor).abs().sum()
                    sharded_loss.backward()
                    return sharded_loss

                loss_ddp = cast(
                    torch.Tensor, ddp_optimizer.step(closure=closure_ddp),
                )
                loss_sharded_optim = cast(
                    torch.Tensor,
                    sharded_optimizer.step(closure=closure_sharded),
                )
                torch.testing.assert_close(
                    loss_ddp, loss_sharded_optim,
                    msg="Losses differ between local optimizer and ZeRO",
                )
                self._check_same_model_params(
                    sharded_ddp_model, ddp_model,
                    "Models differ after a step",
                )

            # Check that parity is maintained
            for i in range(BATCHES):
                check_step()
                # For the second half of batches, change the parameter
                # trainability to further test parity
                if i > BATCHES // 2:
                    next(ddp_model.parameters()).requires_grad = bool(i % 2)
                    next(sharded_ddp_model.parameters()).requires_grad = bool(i % 2)

            # Check that the `state_dict` checkpoints are compatible between
            # the local optimizer and ZeRO
            REFERENCE_RANK = 0
            # - Get states
            ddp_state_dict = ddp_optimizer.state_dict()
            sharded_optimizer.consolidate_state_dict(to=REFERENCE_RANK)
            sharded_optim_state_dict = [
                sharded_optimizer.state_dict()
                if self.rank == REFERENCE_RANK else {}
            ]
            dist.broadcast_object_list(
                sharded_optim_state_dict, src=REFERENCE_RANK,
                group=dist.group.WORLD,
            )
            sharded_optim_state_dict = sharded_optim_state_dict[0]

            # - Cross-load the states
            # Run one step and check that the models are still the same
            ddp_state_dict_ref = copy.deepcopy(ddp_state_dict)
            ddp_optimizer.load_state_dict(sharded_optim_state_dict)
            sharded_optimizer.load_state_dict(ddp_state_dict)
            check_step()

            # - Reload their respective states
            # Run one step and check that the models are still the same
            ddp_optimizer.load_state_dict(ddp_state_dict_ref)
            sharded_optimizer.load_state_dict(sharded_optim_state_dict)
            check_step()

    def _test_zero_join(self, device):
        """Check that the ZeRO join hook allows training with uneven inputs
        when using the given device."""
        NUM_INPUTS = 3
        NUM_EPOCHS = 2
        LR = 0.01
        torch.manual_seed(0)
        torch.cuda.manual_seed(0)

        rank = self.rank
        world_size = self.world_size
        is_gpu = device.type == "cuda"
        backend = _get_backend_for_tests() if is_gpu else dist.Backend.GLOO
        self.dist_init(rank, world_size, backend)

        model = torch.nn.Sequential(
            torch.nn.Linear(2, 3),
            torch.nn.Linear(3, 3),
            torch.nn.Linear(3, 3),
        )
        model.to(device)

        # DDP ensures correct gradients in data parallel training, so DDP with
        # local optimizers on uneven inputs should be equivalent to ZeRO on
        # uneven inputs with gradients being manually set
        ddp_model = DDP(model, device_ids=[rank]) if is_gpu else DDP(model)
        local_optim = torch.optim.Adam(ddp_model.parameters(), lr=LR)
        zero_model = copy.deepcopy(model)
        zero_model.to(device)
        zero_optim = ZeroRedundancyOptimizer(
            zero_model.parameters(), torch.optim.Adam, lr=LR,
        )
        loss_fn = torch.nn.MSELoss()

        # Use uneven inputs: rank i has i extra inputs
        inputs = [
            torch.randn(20, 2).to(device) for _ in range(NUM_INPUTS + rank)
        ]
        labels = torch.randn(20, 3).to(device)

        # Save the gradients and parameters from DDP as the ground truth; do
        # so on the last-joining rank (in this case, the largest rank)
        grads_at_each_iter = []
        params_at_each_iter = []
        with ddp_model.join():
            for _ in range(NUM_EPOCHS):
                for input in inputs:
                    output = ddp_model(input)
                    loss_fn(output, labels).backward()
                    if rank == world_size - 1:
                        grads = []
                        for p in ddp_model.parameters():
                            grads.append(p.grad.detach().clone().to(device))
                    local_optim.step()
                    if rank == world_size - 1:
                        params = []
                        for p in ddp_model.parameters():
                            params.append(p.detach().clone().to(device))
                        grads_at_each_iter.append(grads)
                        params_at_each_iter.append(params)

        # Broadcast the saved gradients and parameters to all of the other
        # ranks (which joined early)
        grads_and_params = [grads_at_each_iter, params_at_each_iter]
        grads_and_params = _broadcast_object(
            grads_and_params, src_rank=world_size - 1, group=dist.group.WORLD,
            device=device,
        )
        grads_at_each_iter = grads_and_params[0]
        params_at_each_iter = grads_and_params[1]
        # TODO: Replace this `_broadcast_object` with `broadcast_object_list`
        # once the latter supports loading to the destination device instead
        # of the source device

        # A process must still set the remaining gradients after joining, so we
        # define a join hook to do this before the ZeRO join hook
        class _JoinGradInfo():
            def __init__(self, grads):
                self.grads = grads  # remaining gradients to set (in order)
                self.index = 0

        class _SetGradsJoinHook(JoinHook):
            def __init__(self, zero_optim, grads):
                zero_optim._join_grad_info = _JoinGradInfo(grads)
                self.zero = zero_optim
                super().__init__()

            def main_hook(self):
                join_grad_info = self.zero._join_grad_info
                grads = self.zero._join_grad_info.grads[join_grad_info.index]
                join_grad_info.index += 1
                for p, grad in zip(self.zero._all_params, grads):
                    p.grad = grad.detach().clone().to(device)

        class _GradientSetter(Joinable):
            def __init__(self):
                super().__init__()

            def join_hook(self, **kwargs):
                assert "zero_optim" in kwargs
                assert "grads" in kwargs
                zero_optim = kwargs["zero_optim"]
                grads = kwargs["grads"]
                return _SetGradsJoinHook(zero_optim, grads)

            @property
            def join_device(self):
                return device

            @property
            def join_process_group(self):
                return dist.group.WORLD

        num_grads_after_joining = NUM_EPOCHS * (world_size - rank - 1)
        grads = grads_at_each_iter[-num_grads_after_joining:]
        gradient_setter = _GradientSetter()
        iter = 0
        with Join(
            [gradient_setter, zero_optim], zero_optim=zero_optim, grads=grads,
        ):
            for _ in range(NUM_EPOCHS):
                for input in inputs:
                    # Notify join context that this process has not joined
                    Join.notify_join_context(gradient_setter)
                    # Set gradients manually
                    for p, grad in zip(
                        zero_model.parameters(), grads_at_each_iter[iter],
                    ):
                        p.grad = grad.detach().clone().to(device)
                    # Perform optimizer step and check parity
                    zero_optim.step()
                    for p, ddp_p in zip(
                        zero_model.parameters(), params_at_each_iter[iter],
                    ):
                        torch.testing.assert_close(
                            p, ddp_p,
                            msg="Parameters differ between using ZeRO and "
                            "local optimizer",
                        )
                    iter += 1

    @common_distributed.requires_nccl()
    @common_distributed.skip_if_no_gpu
    def test_zero_join_gpu(self):
        """Check that the ZeRO join hook allows training with uneven inputs
        on GPU."""
        self._test_zero_join(self.device)

    @common_distributed.requires_gloo()
    def test_zero_join_cpu(self):
        """Check that the ZeRO join hook allows training with uneven inputs
        on CPU."""
        self._test_zero_join(torch.device("cpu"))

    def _test_zero_model_parallel(self, parameters_as_bucket_view: bool):
        # Use two processes each with two GPUs
        assert self.rank < 2
        NUM_EPOCHS = 2
        NUM_INPUTS = 4
        LR = 0.01
        torch.manual_seed(0)
        torch.cuda.manual_seed(0)

        class ModelParallelModel(torch.nn.Module):
            def __init__(self, dev0, dev1):
                super().__init__()
                self.dev0 = dev0
                self.dev1 = dev1
                self.net0 = torch.nn.Linear(10, 10).to(dev0)
                self.relu = torch.nn.ReLU()
                self.net1 = torch.nn.Linear(10, 5).to(dev1)

            def forward(self, x):
                x = x.to(self.dev0)
                x = self.relu(self.net0(x))
                x = x.to(self.dev1)
                return self.net1(x)

        class LocalModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.net0 = torch.nn.Linear(10, 10)
                self.relu = torch.nn.ReLU()
                self.net1 = torch.nn.Linear(10, 5)

            def forward(self, x):
                return self.net1(self.relu(self.net0(x)))

        dev0 = torch.device(2 * self.rank)
        dev1 = torch.device(2 * self.rank + 1)
        mp_model = ModelParallelModel(dev0, dev1)
        ddp_model = DDP(mp_model)
        local_model = LocalModel().to(dev0)

        # Ensure the parameters are the same across the two models
        def copy_param(p):
            return torch.nn.Parameter(p.detach().clone().to(dev0))

        local_model.net0.weight = copy_param(mp_model.net0.weight)
        local_model.net0.bias = copy_param(mp_model.net0.bias)
        local_model.net1.weight = copy_param(mp_model.net1.weight)
        local_model.net1.bias = copy_param(mp_model.net1.bias)

        # Compare parity between DDP with model parallelism using ZeRO and
        # a local model using a local optimizer
        zero_optim = ZeroRedundancyOptimizer(
            ddp_model.parameters(),
            optimizer_class=torch.optim.Adam,
            parameters_as_bucket_view=parameters_as_bucket_view,
            lr=LR,
        )
        local_optim = torch.optim.Adam(local_model.parameters(), lr=LR)
        inputs = [torch.randn(20, 10).to(dev0) for _ in range(NUM_INPUTS)]

        for _ in range(NUM_EPOCHS):
            for input in inputs:
                def closure_local():
                    local_optim.zero_grad()
                    local_loss = local_model(input).abs().sum()
                    local_loss.backward()
                    return local_loss

                def closure_ddp():
                    zero_optim.zero_grad()
                    ddp_loss = ddp_model(input).abs().sum()
                    ddp_loss.backward()
                    return ddp_loss

                local_loss = cast(
                    torch.Tensor, local_optim.step(closure=closure_local)
                )
                ddp_loss = cast(
                    torch.Tensor, zero_optim.step(closure=closure_ddp)
                )

                # Increased tolerances are needed to pass when using TF32
                # See: https://github.com/pytorch/pytorch/issues/67764
                torch.testing.assert_close(
                    local_loss.cpu(), ddp_loss.cpu(), rtol=1e-03, atol=1e-08,
                ), "Losses differ between local optimizer and ZeRO"

                for local_p, ddp_p in zip(
                    local_model.parameters(),
                    ddp_model.parameters()
                ):
                    torch.testing.assert_close(
                        local_p.cpu(), ddp_p.cpu(), rtol=1e-03, atol=1e-04,
                    ), "Models differ after a step"

    @common_distributed.skip_if_lt_x_gpu(4)
    @parametrize(
        "parameters_as_bucket_view",
        [False, True],
    )
    def test_zero_model_parallel(
        self,
        parameters_as_bucket_view: bool,
    ):
        """Check that ZeRO works with model parallelism where the model's
        layers are assigned to different devices."""
        if self.rank >= 2:
            return
        # Disable DDP + ReplicatedTensor when `parameter_as_bucket_view=True`
        # since then ZeroRedundancyOptimizer modifies the model parameters in
        # place.
        from torch.nn.parallel._replicated_tensor_ddp_utils import _ddp_replicated_tensor
        context = _ddp_replicated_tensor(False) if parameters_as_bucket_view \
            else suppress()
        with context:
            self.dist_init(self.rank, world_size=2)
            self._test_zero_model_parallel(parameters_as_bucket_view)

    def _test_ddp_zero_overlap(
        self,
        device,
        hook_constructor,
        gradient_as_bucket_view,
        static_graph,
        **kwargs,
    ):
        SGD_LR = 0.01
        SGD_MOMENTUM = 0.9
        SGD_WEIGHT_DECAY = 0.001
        NUM_INPUTS = 5
        torch.manual_seed(0)
        torch.cuda.manual_seed(0)

        rank = self.rank
        is_gpu = device.type == "cuda"
        if is_gpu:
            torch.cuda.set_device(device)
        models_to_test = [(
            torch.nn.Sequential(
                torch.nn.Linear(1000, 2000),
                torch.nn.Linear(2000, 500),
            ),
            [torch.randn(1, 1000).to(device) for _ in range(NUM_INPUTS)],
        )]
        if HAS_TORCHVISION:
            models_to_test.append((
                torchvision.models.resnet50(),
                [
                    torch.randn(1, 3, 3, 1000).to(device)
                    for _ in range(NUM_INPUTS)
                ]
            ))
        for (model, inputs) in models_to_test:
            # Enable determinism in cudnn operators
            with torch.backends.cudnn.flags(
                enabled=True, deterministic=True, benchmark=False
            ):
                device_ids = [rank] if is_gpu else None
                # Set up the DDP model overlapping with ZeRO
                ddp_model_overlap = DDP(
                    copy.deepcopy(model).to(device),
                    device_ids=device_ids,
                    gradient_as_bucket_view=gradient_as_bucket_view
                )
                if static_graph:
                    ddp_model_overlap._set_static_graph()
                zero_optim = ZeroRedundancyOptimizer(
                    ddp_model_overlap.parameters(),
                    optimizer_class=torch.optim.SGD,
                    overlap_with_ddp=True,
                    lr=SGD_LR,
                    momentum=SGD_MOMENTUM,
                    weight_decay=SGD_WEIGHT_DECAY,
                )
                ddp_model_overlap.register_comm_hook(
                    None,
                    hook_constructor(
                        allreduce_hook, ddp_model_overlap, zero_optim,
                        **kwargs,
                    )
                )

                # Set up the DDP model with local optimizer
                ddp_model_local = DDP(
                    copy.deepcopy(model).to(device),
                    device_ids=device_ids,
                    gradient_as_bucket_view=gradient_as_bucket_view
                )
                if static_graph:
                    ddp_model_local._set_static_graph()
                local_optim = torch.optim.SGD(
                    ddp_model_local.parameters(),
                    lr=SGD_LR,
                    momentum=SGD_MOMENTUM,
                    weight_decay=SGD_WEIGHT_DECAY
                )

                # Check that the parameters match initially
                for p1, p2 in zip(
                    ddp_model_overlap.parameters(),
                    ddp_model_local.parameters()
                ):
                    self.assertEqual(p1, p2)

                # Save the parameters to ensure they were updated
                init_params_overlap = copy.deepcopy(
                    list(ddp_model_overlap.parameters())
                )

                # Ensure that this test runs independently
                dist.barrier()

                # Run the DDP model overlapping with ZeRO
                # NOTE: Overlapping currently requires 2 or 3 warmup iterations
                # to ensure DDP buckets have been rebuilt (depending on the
                # value of `static_graph`)
                num_warmup_inputs = 2 if not static_graph else 3
                for input in inputs[:num_warmup_inputs]:
                    output = ddp_model_overlap(input)
                    loss = output.sum()
                    loss.backward()
                for input in inputs:
                    zero_optim.zero_grad()
                    output = ddp_model_overlap(input)
                    loss = output.sum()
                    loss.backward()

                # Run the DDP model with local optimizer
                for input in inputs:
                    local_optim.zero_grad()
                    output = ddp_model_local(input)
                    loss = output.sum()
                    loss.backward()
                    local_optim.step()
                dist.barrier()

                # Check that the parameters are equal
                for p1, p2 in zip(
                    ddp_model_overlap.parameters(),
                    ddp_model_local.parameters()
                ):
                    self.assertEqual(p1, p2)

                # Check that the parameters were updated
                self.assertNotEqual(
                    init_params_overlap, list(ddp_model_overlap.parameters()),
                )

                # Ensure that this test runs independently
                dist.barrier()

    # NOTE: The test is skipped if using Windows since functional optimizers
    # are not currently supported.
    @common_distributed.skip_if_win32()
    @common_distributed.requires_nccl()
    @common_distributed.skip_if_no_gpu
    @common_distributed.skip_if_rocm
    @parametrize(
        "use_gpu",
        [True],
        # Add `False` once the Gloo sync issue causing hangs is fixed
        # See: https://github.com/pytorch/pytorch/issues/62300
    )
    @parametrize(
        "use_interleaved_hook",
        [False, True],
    )
    @parametrize(
        "gradient_as_bucket_view",
        [False, True],
    )
    @parametrize(
        "static_graph",
        [False, True],
    )
    @parametrize(
        "shard_buckets",
        [False, True],
    )
    def test_ddp_zero_overlap(
        self,
        use_gpu: bool,
        use_interleaved_hook: bool,
        gradient_as_bucket_view: bool,
        static_graph: bool,
        shard_buckets: bool,
    ):
        """
        Check that overlapping DDP with ZeRO using the given method determined
        by ``hook_constructor`` and ``shard_buckets`` and using the given ZeRO
        and DDP arguments achieves parity with DDP using a local optimizer.
        """
        device = torch.device(self.rank) if use_gpu else torch.device("cpu")
        backend = _get_backend_for_tests()
        self.dist_init(self.rank, self.world_size, backend)
        hook_constructor = hook_with_zero_step if not use_interleaved_hook \
            else hook_with_zero_step_interleaved

        # Disable DDP + ReplicatedTensor since ZeroRedundancyOptimizer
        # modifies the model parameters in place.
        from torch.nn.parallel._replicated_tensor_ddp_utils import _ddp_replicated_tensor
        with _ddp_replicated_tensor(False):
            self._test_ddp_zero_overlap(
                device, hook_constructor, gradient_as_bucket_view, static_graph,
                shard_buckets=shard_buckets,
            )


instantiate_parametrized_tests(TestZeroRedundancyOptimizerSingleRank)
instantiate_parametrized_tests(TestZeroRedundancyOptimizerDistributed)

if __name__ == "__main__":
    # ! unittest should not be used here, else the tests are not properly registered
    run_tests()