deepspeedai--deepspeed/tests/unit/runtime/half_precision/test_zero_optim_overflow.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
import deepspeed
from deepspeed.accelerator import get_accelerator
import pytest
import numpy as np
from unit.common import DistributedTest
from unit.simple_model import SimpleModel, random_dataloader
from deepspeed.utils import safe_set_full_grad


def has_inf_or_nan(x):
    float_x = x.float()
    nan = float_x.isnan()
    inf = float_x.isinf()
    inf_or_nan = nan.logical_or(inf)
    return inf_or_nan.float().max()


def run_model_step(model, x_sample, y_label, grad_value):
    loss = model(x_sample, y_label)
    model.backward(loss)
    for p in model.parameters():
        grad = torch.empty_like(p, dtype=p.dtype)
        grad.fill_(grad_value)
        safe_set_full_grad(p, grad)
    model.step()


@pytest.mark.parametrize("zero_stage", [1, 2])
@pytest.mark.parametrize("offload_optimizer", [False, True])
class TestZeROFloat16(DistributedTest):
    world_size = 2

    def test_no_overflow(self, zero_stage, offload_optimizer):
        if not get_accelerator().is_fp16_supported():
            pytest.skip("fp16 is not supported")

        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 0.00015
                }
            },
            "fp16": {
                "enabled": True,
                "loss_scale": 0,
                "initial_scale_power": 8,
                "loss_scale_window": 2
            },
            "zero_optimization": {
                "stage": zero_stage
            }
        }

        if offload_optimizer:
            config_dict["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}

        hidden_dim = 10
        model = SimpleModel(hidden_dim)
        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())

        expected_loss_scale = 2**8
        expected_scale_window = 2
        # Ensure the dynamic loss scaler is correctly configured.
        loss_scaler = optim.loss_scaler

        assert optim.dynamic_loss_scale == True
        assert loss_scaler.cur_scale == expected_loss_scale
        assert loss_scaler.scale_window == expected_scale_window

        num_iterations = 10
        grad_values = np.random.uniform(-0.1, 0.1, num_iterations)
        data_loader = random_dataloader(model=model,
                                        total_samples=num_iterations,
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=torch.float16)
        for i, (batch, grad_value) in enumerate(zip(data_loader, grad_values)):
            run_model_step(model, batch[0], batch[1], grad_value)
            assert loss_scaler.cur_scale == expected_loss_scale
            assert loss_scaler.cur_iter == (i + 1)

            if loss_scaler.cur_iter % expected_scale_window == 0:
                expected_loss_scale *= 2

    def test_all_overflow(self, zero_stage, offload_optimizer):
        if not get_accelerator().is_fp16_supported():
            pytest.skip("fp16 is not supported")

        overflow_gradients = [float('inf'), float('-inf')] + [float('nan')] * 6
        initial_scale_power = len(overflow_gradients)
        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 0.00015
                }
            },
            "fp16": {
                "enabled": True,
                "loss_scale": 0,
                "initial_scale_power": initial_scale_power,
                "loss_scale_window": 2,
                "hysteresis": 1,
            },
            "zero_optimization": {
                "stage": zero_stage,
            }
        }

        if offload_optimizer:
            config_dict["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}

        hidden_dim = 10
        model = SimpleModel(hidden_dim)
        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())

        expected_loss_scale = 2**initial_scale_power
        expected_scale_window = 2
        # Ensure the dynamic loss scaler is correctly configured.
        loss_scaler = optim.loss_scaler

        assert optim.dynamic_loss_scale == True
        assert loss_scaler.cur_scale == expected_loss_scale
        assert loss_scaler.scale_window == expected_scale_window

        data_loader = random_dataloader(model=model,
                                        total_samples=len(overflow_gradients),
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=torch.float16)
        for i, (batch, grad_value) in enumerate(zip(data_loader, overflow_gradients)):
            run_model_step(model, batch[0], batch[1], grad_value)
            expected_loss_scale = max(expected_loss_scale / 2, 1)
            assert loss_scaler.cur_scale == expected_loss_scale
            assert loss_scaler.cur_iter == (i + 1)

    def test_some_overflow(self, zero_stage, offload_optimizer):
        if not get_accelerator().is_fp16_supported():
            pytest.skip("fp16 is not supported")
        initial_scale_power = 8
        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 0.00015
                }
            },
            "fp16": {
                "enabled": True,
                "loss_scale": 0,
                "initial_scale_power": initial_scale_power,
                "loss_scale_window": 2,
                "hysteresis": 1,
            },
            "zero_optimization": {
                "stage": zero_stage,
            }
        }

        if offload_optimizer:
            config_dict["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}

        hidden_dim = 10
        model = SimpleModel(hidden_dim)
        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())

        expected_loss_scale = 2**initial_scale_power
        expected_scale_window = 2
        # Ensure the dynamic loss scaler is correctly configured.
        loss_scaler = optim.loss_scaler

        assert optim.dynamic_loss_scale == True
        assert loss_scaler.cur_scale == expected_loss_scale
        assert loss_scaler.scale_window == expected_scale_window

        expected_iteration = 0

        # Run model with overflows to decrease scale
        overflow_gradients = [float('inf'), float('nan')]
        expected_iteration += len(overflow_gradients)
        data_loader = random_dataloader(model=model,
                                        total_samples=len(overflow_gradients),
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=torch.float16)
        for batch, grad_value in zip(data_loader, overflow_gradients):
            run_model_step(model, batch[0], batch[1], grad_value)

        expected_loss_scale /= (2**len(overflow_gradients))
        assert loss_scaler.cur_scale == expected_loss_scale
        assert loss_scaler.cur_iter == expected_iteration

        # Run model scale_window + 1 times to increase scale once
        normal_gradients = np.random.uniform(-0.1, 0.1, expected_scale_window + 1)
        expected_iteration += len(normal_gradients)
        data_loader = random_dataloader(model=model,
                                        total_samples=len(normal_gradients),
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=torch.float16)
        for batch, grad_value in zip(data_loader, normal_gradients):
            run_model_step(model, batch[0], batch[1], grad_value)

        expected_loss_scale *= 2
        assert loss_scaler.cur_scale == expected_loss_scale
        assert loss_scaler.cur_iter == expected_iteration

        # Run model with overflows to decrease scale
        overflow_gradients = [float('inf')]
        expected_iteration += len(overflow_gradients)
        data_loader = random_dataloader(model=model,
                                        total_samples=len(overflow_gradients),
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=torch.float16)
        for batch, grad_value in zip(data_loader, overflow_gradients):
            run_model_step(model, batch[0], batch[1], grad_value)

        expected_loss_scale /= (2**len(overflow_gradients))
        assert loss_scaler.cur_scale == expected_loss_scale
        assert loss_scaler.cur_iter == expected_iteration


@pytest.mark.parametrize("zero_stage", [1, 2])
@pytest.mark.parametrize("offload_optimizer", [False, True])
class TestZeROBFloat16(DistributedTest):
    world_size = 2

    def test_no_overflow(self, zero_stage, offload_optimizer):
        if not get_accelerator().is_bf16_supported():
            pytest.skip("bf16 is not supported")

        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 0.00015
                }
            },
            "bf16": {
                "enabled": True,
            },
            "zero_optimization": {
                "stage": zero_stage
            }
        }

        if offload_optimizer:
            config_dict["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}

        hidden_dim = 10
        model = SimpleModel(hidden_dim)
        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())

        num_iterations = 10
        grad_values = np.random.uniform(-0.1, 0.1, num_iterations)
        data_loader = random_dataloader(model=model,
                                        total_samples=num_iterations,
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=torch.bfloat16)
        for i, (batch, grad_value) in enumerate(zip(data_loader, grad_values)):
            run_model_step(model, batch[0], batch[1], grad_value)

        assert model.skipped_steps == 0
        assert all([not has_inf_or_nan(p) for p in model.parameters()])

    def test_detect_grad_overflow(self, zero_stage, offload_optimizer):
        if not get_accelerator().is_bf16_supported():
            pytest.skip("bf16 is not supported")

        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 0.00015
                }
            },
            "bf16": {
                "enabled": True,
                "check_grad_overflow": True
            },
            "zero_optimization": {
                "stage": zero_stage,
            }
        }

        if offload_optimizer:
            config_dict["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}

        hidden_dim = 10
        model = SimpleModel(hidden_dim)
        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())

        overflow_gradients = [float('inf'), float('-inf')] + [float('nan')] * 6
        data_loader = random_dataloader(model=model,
                                        total_samples=len(overflow_gradients),
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=torch.bfloat16)

        for i, (batch, grad_value) in enumerate(zip(data_loader, overflow_gradients)):
            run_model_step(model, batch[0], batch[1], grad_value)
            assert model.skipped_steps == (i + 1)

        assert all([not has_inf_or_nan(p) for p in model.parameters()])

    def test_ignore_grad_overflow(self, zero_stage, offload_optimizer):
        if not get_accelerator().is_bf16_supported():
            pytest.skip("bf16 is not supported")

        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 0.00015
                }
            },
            "bf16": {
                "enabled": True,
                "check_grad_overflow": False
            },
            "zero_optimization": {
                "stage": zero_stage,
            }
        }

        if offload_optimizer:
            config_dict["zero_optimization"]["offload_optimizer"] = {"device": "cpu"}

        hidden_dim = 10
        model = SimpleModel(hidden_dim)
        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())

        overflow_gradients = [float('inf'), float('-inf')] + [float('nan')] * 6
        data_loader = random_dataloader(model=model,
                                        total_samples=len(overflow_gradients),
                                        hidden_dim=hidden_dim,
                                        device=model.device,
                                        dtype=torch.bfloat16)

        for i, (batch, grad_value) in enumerate(zip(data_loader, overflow_gradients)):
            run_model_step(model, batch[0], batch[1], grad_value)

        assert model.skipped_steps == 0
        assert all([has_inf_or_nan(p) for p in model.parameters()])