chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,365 @@
+# 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()])