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

import copy
import functools
import itertools
import sys
import unittest
from typing import List, Optional

import torch
from torch import distributed as dist
from torch.cuda.amp.common import amp_definitely_not_available
from torch.distributed.fsdp import CPUOffload, MixedPrecision
from torch.distributed.fsdp.fully_sharded_data_parallel import (
    FullyShardedDataParallel as FSDP,
    ShardingStrategy,
)
from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
from torch.distributed.fsdp.wrap import ModuleWrapPolicy
from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer
from torch.nn.parallel.distributed import DistributedDataParallel as DDP
from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
from torch.testing._internal.common_fsdp import (
    DEVICEInitMode,
    DummyProcessGroup,
    FSDPInitMode,
    FSDPTest,
    NestedWrappedModule,
    NonUniformReqGradNWM,
    subtest_name,
    TransformerWithSharedParams,
)
from torch.testing._internal.common_utils import (
    instantiate_parametrized_tests,
    parametrize,
    run_tests,
    TEST_WITH_DEV_DBG_ASAN,
    TestCase,
)


if not dist.is_available():
    print("Distributed not available, skipping tests", file=sys.stderr)
    sys.exit(0)


if TEST_WITH_DEV_DBG_ASAN:
    print(
        "Skip dev-asan as torch + multiprocessing spawn have known issues",
        file=sys.stderr,
    )
    sys.exit(0)


params = "cpu_offload,sharding_strategy,mixed_precision,use_orig_params"
cpu_offload_config = [CPUOffload(offload_params=True), CPUOffload(offload_params=False)]
sharding_strategy_config = [ShardingStrategy.SHARD_GRAD_OP, None]
mixed_precision = ["enable_mixed_precision", None]
use_orig_params = ["enable_use_orig_params", None]

configs = list(
    itertools.product(
        cpu_offload_config, sharding_strategy_config, mixed_precision, use_orig_params
    )
)
test_name_mapping = {
    str(CPUOffload(offload_params=True)): "offload_true",
    str(CPUOffload(offload_params=False)): "offload_false",
    str(ShardingStrategy.SHARD_GRAD_OP): "shard_grad_op",
    "enable_mixed_precision": "mixed_precision",
    "enable_use_orig_params": "use_orig_params",
}

subtest_name = functools.partial(subtest_name, test_name_mapping)


class TestShardGradScaler(TestCase):
    @unittest.skipIf(
        amp_definitely_not_available(), "no supported device (cuda, xla) found"
    )
    def test_grad_scaling(self):
        pg = DummyProcessGroup(0, 1)
        scaler = ShardedGradScaler(init_scale=2.0, process_group=pg, enabled=True)
        t0 = torch.full((1,), 4.0, dtype=torch.float32, device="cpu")
        t1 = torch.full((1,), 8.0, dtype=torch.float32, device="cpu")
        outputs = [t1.clone(), (t0.clone(), t1.clone()), [t0.clone(), t1.clone()]]
        outputs = scaler.scale(outputs)
        self.assertTrue(
            outputs[0] == 16.0 and outputs[1][0] == 8.0 and outputs[1][1] == 16.0
        )
        self.assertTrue(outputs[2][0] == 8.0 and outputs[2][1] == 16.0)
        self.assertTrue(scaler._scale.device == t1.device)

    @unittest.skipIf(
        amp_definitely_not_available(), "no supported device (cuda, xla) found"
    )
    def test_scaling_unscaling_sparse(self):
        pg = DummyProcessGroup(0, 1)
        scaler = ShardedGradScaler(init_scale=2.0, process_group=pg, enabled=True)
        inv_scale = torch.full((1,), 0.5, dtype=torch.float, device="cpu")
        found_inf = torch.full((1,), 0, dtype=torch.float, device="cpu")

        i = torch.tensor([[0, 1, 1], [2, 0, 2]], device="cpu", dtype=torch.int64)
        v = torch.tensor([16.0, 32.0, 64.0], dtype=torch.float, device="cpu")
        s = torch.sparse_coo_tensor(
            i, v, torch.Size([2, 3]), device="cpu", dtype=torch.float
        )

        # unscale sparse tensors
        s1 = s.clone()
        s1.grad = s.clone()
        opt = torch.optim.SGD([s1], lr=1.0)
        found_inf.zero_()
        found_inf = scaler._unscale_grads_(opt, inv_scale, found_inf)[s1.device]
        self.assertEqual(found_inf, 0.0)
        self.assertEqual(s1.grad.to_dense(), (s / 2).to_dense())

        # unscale sparse tensor: inf
        v = torch.tensor([16.0, 32.0, float("inf")], dtype=torch.float, device="cpu")
        s1.grad = torch.sparse_coo_tensor(
            i, v, torch.Size([2, 3]), device="cpu", dtype=torch.float
        )
        found_inf.zero_()
        found_inf = scaler._unscale_grads_(opt, inv_scale, found_inf)[s1.device]
        self.assertEqual(found_inf, 1.0)

        # unscale sparse tensor: overflow (marked as inf)
        i = torch.tensor([[1, 1, 1], [0, 0, 2]], device="cpu", dtype=torch.int64)
        # coalescing sparse tensor here will cause the value to be Inf
        v = torch.tensor([2**15, 2**15, 1.0], dtype=torch.float16, device="cpu")
        s1 = torch.sparse_coo_tensor(
            i, v, torch.Size([2, 3]), device="cpu", dtype=torch.float16
        )
        s1.grad = s1.clone()
        found_inf.zero_()
        found_inf = scaler._unscale_grads_(opt, inv_scale, found_inf)[s1.device]
        self.assertEqual(found_inf, 1.0)

    @unittest.skipIf(
        amp_definitely_not_available(), "no supported device (cuda, xla) found"
    )
    def test_inf_gradients_skip_optim_step(self):
        pg = DummyProcessGroup(0, 1)
        scaler = ShardedGradScaler(init_scale=2.0, process_group=pg, enabled=True)
        loss = torch.full((1,), 4.0, dtype=torch.float32, device="cpu")
        t0 = torch.tensor([float("inf")], dtype=torch.float32, device="cpu")
        t0.grad = t0.clone()
        opt = torch.optim.SGD([t0], lr=1.0)
        scaler.scale(loss)
        ret_val = scaler.step(opt)
        self.assertTrue(ret_val is None)


class TestShardedGradScalerParityWithDDP(FSDPTest):
    def _get_init_modes_for_test(self, cpu_offload):
        modes = [DEVICEInitMode.DEVICE_AFTER, DEVICEInitMode.DEVICE_BEFORE]
        # Note that DEVICEInitMode.DEVICE_NEVER works currently only with CPU
        # offload as we explicitly bring the param back to CUDA device. In
        # general, it will not work since we try to all_gather p.data which is
        # on CPU but NCCL only supports GPU.
        if cpu_offload.offload_params:
            modes.append(DEVICEInitMode.DEVICE_NEVER)

        return modes

    @skip_if_lt_x_gpu(2)
    @parametrize(params, configs, subtest_name)
    def test_fsdp_ddp_parity_with_grad_scaler(
        self,
        cpu_offload: CPUOffload,
        sharding_strategy: Optional[ShardingStrategy],
        mixed_precision: Optional[str],
        use_orig_params: Optional[str],
    ):
        init_modes = self._get_init_modes_for_test(cpu_offload)
        mp = (
            MixedPrecision(
                param_dtype=torch.float16,
                reduce_dtype=torch.float16,
                buffer_dtype=torch.float16,
            )
            if mixed_precision is not None
            else None
        )
        # the ``NonUniformReqGradNWM`` model requires we set `init_scale`
        # more conservatively than default to avoid infs with the initial steps
        if use_orig_params == "enable_use_orig_params":
            use_orig = True
            model_cls = NonUniformReqGradNWM
            sharded_grad_scaler_kwargs = {"init_scale": 2.0**11}
        else:
            use_orig = False
            model_cls = NestedWrappedModule  # type: ignore[assignment]
            sharded_grad_scaler_kwargs = None
        for device_init_mode in init_modes:
            self._test_fsdp_parity(
                model_cls,
                FSDPInitMode.RECURSIVE,
                device_init_mode=device_init_mode,
                cpu_offload=cpu_offload,
                sharding_strategy=sharding_strategy,
                mixed_precision=mp,
                enable_sharded_grad_scaler=True,
                use_orig_params=use_orig,
                sharded_grad_scaler_kwargs=sharded_grad_scaler_kwargs,
            )

    def _build_model_and_optim(
        self,
        cpu_offload: CPUOffload = CPUOffload(offload_params=False),
        use_orig_params: bool = False,
    ):
        model = TransformerWithSharedParams.init(
            self.process_group,
            FSDPInitMode.NO_FSDP,
            DEVICEInitMode.DEVICE_BEFORE,
            deterministic=True,
        )
        ref_model = DDP(
            copy.deepcopy(model),
            device_ids=[self.rank],
        )
        ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
        fsdp_kwargs = {
            "use_orig_params": use_orig_params,
            "cpu_offload": cpu_offload,
            "auto_wrap_policy": ModuleWrapPolicy(
                {
                    TransformerEncoderLayer,
                    TransformerDecoderLayer,
                }
            ),
        }
        model = FSDP(model, **fsdp_kwargs)
        optim = torch.optim.Adam(model.parameters(), lr=1e-2)
        return model, optim, ref_model, ref_optim

    @skip_if_lt_x_gpu(2)
    def test_sharded_grad_scaler_found_inf(self):
        self.run_subtests(
            {
                "use_orig_params": [False, True],
                "cpu_offload": [
                    CPUOffload(offload_params=True),
                    CPUOffload(offload_params=False),
                ],
            },
            self._test_sharded_grad_scaler_found_inf,
        )

    def _test_sharded_grad_scaler_found_inf(
        self,
        use_orig_params: bool,
        cpu_offload: CPUOffload,
    ):
        model, optim, ref_model, ref_optim = self._build_model_and_optim(
            cpu_offload=cpu_offload,
            use_orig_params=use_orig_params,
        )
        grad_scaler = ShardedGradScaler(init_scale=2.0)
        ref_grad_scaler = torch.amp.GradScaler(device="cuda", init_scale=2.0)
        scaled_losses: List[torch.Tensor] = []
        device = torch.device("cuda")
        torch.manual_seed(42 + self.rank + 1)

        for iter in range(10):
            for _model, _optim, _grad_scaler in (
                (ref_model, ref_optim, ref_grad_scaler),
                (model, optim, grad_scaler),
            ):
                module = _model.module
                inp = module.get_input(device)
                _optim.zero_grad()
                output = _model(*inp)
                loss = module.get_loss(inp, output)
                scaled_loss = _grad_scaler.scale(loss)
                scaled_losses.append(scaled_loss)
                scaled_loss.backward()
                orig_params = [
                    param.detach().clone()
                    for param in _model.parameters()
                    if param.grad is not None
                ]
                should_find_inf = iter % 2 == 0
                if should_find_inf and (
                    _model is ref_model or (_model is model and self.rank == 0)
                ):
                    # other ranks should find infs from rank 0
                    # after collectives
                    for param in _model.parameters():
                        if param.grad is None:
                            continue
                        param.grad.fill_(float("inf"))
                        break
                _grad_scaler.step(_optim)
                orig_scale = _grad_scaler.get_scale()
                _grad_scaler.update()
                if should_find_inf:
                    self.assertEqual(
                        _grad_scaler.get_scale(),
                        orig_scale * _grad_scaler.get_backoff_factor(),
                        (
                            f"rank: {self.rank} iter: {iter} expect origin scale {orig_scale} "
                            f"to be backed off by {_grad_scaler.get_backoff_factor()} "
                            f"but got {_grad_scaler.get_scale()}"
                        ),
                    )
                else:
                    self.assertEqual(
                        _grad_scaler.get_scale(),
                        orig_scale,
                        (
                            f"rank: {self.rank} iter: {iter} expect same scale {orig_scale} "
                            f"but got {_grad_scaler.get_scale()}"
                        ),
                    )
                for param, orig_param in zip(
                    [param for param in _model.parameters() if param.grad is not None],
                    orig_params,
                ):
                    if should_find_inf:
                        self.assertEqual(
                            param,
                            orig_param,
                            (
                                f"rank: {self.rank} iter: {iter} expect the same params before "
                                f"and after optim.step but got {param} vs {orig_param}"
                            ),
                        )
                    else:
                        self.assertNotEqual(
                            param,
                            orig_param,
                            (
                                f"rank: {self.rank} iter: {iter} expect the updated params after "
                                f"optim.step but got {param} vs {orig_param}"
                            ),
                        )
            self.assertEqual(
                scaled_losses[0],
                scaled_losses[1],
                f"iter: {iter} {scaled_losses[0]} vs {scaled_losses[1]}",
            )


instantiate_parametrized_tests(TestShardGradScaler)
instantiate_parametrized_tests(TestShardedGradScalerParityWithDDP)

if __name__ == "__main__":
    run_tests()