chore: import upstream snapshot with attribution

tests/unit/runtime/half_precision/test_dynamic_loss_scale.py

@@ -0,0 +1,291 @@
+# 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
+from deepspeed.ops.op_builder import FusedLambBuilder
+
+
+def run_model_step(model, gradient_list):
+    for value in gradient_list:
+        for p in model.parameters():
+            p.grad = torch.empty_like(p, dtype=p.dtype)
+            p.grad.fill_(value)
+        model.step()
+
+
+class TestFused(DistributedTest):
+    world_size = 1
+
+    def test_no_overflow(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+
+        config_dict = {
+            "train_batch_size": 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
+            }
+        }
+        hidden_dim = 1
+        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.
+        assert optim.loss_scale_config.dynamic_loss_scale == True
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.scale_window == expected_scale_window
+
+        for i, value in enumerate(np.random.uniform(-0.1, 0.1, 10)):
+            run_model_step(model, [value])
+            assert optim.loss_scale_config.cur_scale == expected_loss_scale
+            assert optim.loss_scale_config.cur_iter == (i + 1)
+            if optim.loss_scale_config.cur_iter % expected_scale_window == 0:
+                expected_loss_scale *= 2
+
+    def test_all_overflow(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 0,
+                "initial_scale_power": 4,
+                "loss_scale_window": 2
+            }
+        }
+        hidden_dim = 1
+        model = SimpleModel(hidden_dim)
+        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+        expected_loss_scale = 2**4
+        # Ensure the dynamic loss scaler is correctly configured.
+        assert optim.loss_scale_config.dynamic_loss_scale == True
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+
+        overflow_gradients = [float('inf'), float('-inf')] + [float('nan')] * 6
+        for i, value in enumerate(overflow_gradients):
+            run_model_step(model, [value])
+            expected_loss_scale = max(expected_loss_scale / 2, 1)
+            assert optim.loss_scale_config.cur_scale == expected_loss_scale
+            assert optim.loss_scale_config.cur_iter == (i + 1)
+
+    def test_some_overflow(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 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
+            }
+        }
+        hidden_dim = 1
+        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
+        expected_iteration = 0
+        # Ensure the dynamic loss scaler is correctly configured.
+        assert optim.loss_scale_config.dynamic_loss_scale == True
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.scale_window == expected_scale_window
+
+        # Run model with overflows to decrease scale
+        overflow_gradients = [float('inf'), float('nan')]
+        expected_iteration += len(overflow_gradients)
+        run_model_step(model, overflow_gradients)
+        expected_loss_scale /= (2**len(overflow_gradients))
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.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)
+        run_model_step(model, normal_gradients)
+        expected_loss_scale *= 2
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.cur_iter == expected_iteration
+
+        # Run model with overflows to decrease scale
+        overflow_gradients = [float('inf')]
+        expected_iteration += len(overflow_gradients)
+        run_model_step(model, overflow_gradients)
+        expected_loss_scale /= (2**len(overflow_gradients))
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.cur_iter == expected_iteration
+
+
+@pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME],
+                    reason="FusedLambBuilder has not been implemented on this system.")
+class TestUnfused(DistributedTest):
+    world_size = 1
+
+    def test_no_overflow(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 0,
+                "initial_scale_power": 8,
+                "loss_scale_window": 2
+            }
+        }
+        hidden_dim = 1
+        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.
+        assert optim.loss_scale_config.dynamic_loss_scale == True
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.scale_window == expected_scale_window
+
+        for i, value in enumerate(np.random.uniform(-0.1, 0.1, 10)):
+            run_model_step(model, [value])
+            assert optim.loss_scale_config.cur_scale == expected_loss_scale
+            assert optim.loss_scale_config.cur_iter == (i + 1)
+            if optim.loss_scale_config.cur_iter % expected_scale_window == 0:
+                expected_loss_scale *= 2
+
+    def test_all_overflow(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+
+        min_loss_scale_value = 2.0
+
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 0,
+                "initial_scale_power": 4,
+                "loss_scale_window": 2,
+                "min_loss_scale": min_loss_scale_value
+            }
+        }
+        hidden_dim = 1
+        model = SimpleModel(hidden_dim)
+        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+        expected_loss_scale = 2**4
+        expected_min_loss_scale = min_loss_scale_value
+        # Ensure the dynamic loss scaler is correctly configured.
+        assert optim.loss_scale_config.dynamic_loss_scale == True
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.min_loss_scale == expected_min_loss_scale
+
+        overflow_gradients = [float('inf'), float('-inf')] + [float('nan')] * 6
+        for i, value in enumerate(overflow_gradients):
+            run_model_step(model, [value])
+            expected_loss_scale = max(expected_loss_scale / 2, expected_min_loss_scale)
+            assert optim.loss_scale_config.cur_scale == expected_loss_scale
+            assert optim.loss_scale_config.cur_iter == (i + 1)
+
+    def test_some_overflow(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 0,
+                "initial_scale_power": 8,
+                "loss_scale_window": 2
+            }
+        }
+        hidden_dim = 1
+        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
+        expected_iteration = 0
+        # Ensure the dynamic loss scaler is correctly configured.
+        assert optim.loss_scale_config.dynamic_loss_scale == True
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.scale_window == expected_scale_window
+
+        # Run model with overflows to decrease scale
+        overflow_gradients = [float('inf'), float('nan')]
+        expected_iteration += len(overflow_gradients)
+        run_model_step(model, overflow_gradients)
+        expected_loss_scale /= (2**len(overflow_gradients))
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.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)
+        run_model_step(model, normal_gradients)
+        expected_loss_scale *= 2
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.cur_iter == expected_iteration
+
+        # Run model with overflows to decrease scale
+        overflow_gradients = [float('inf')]
+        expected_iteration += len(overflow_gradients)
+        run_model_step(model, overflow_gradients)
+        expected_loss_scale /= (2**len(overflow_gradients))
+        assert optim.loss_scale_config.cur_scale == expected_loss_scale
+        assert optim.loss_scale_config.cur_iter == expected_iteration

tests/unit/runtime/half_precision/test_fp16.py

@@ -0,0 +1,717 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed.comm as dist
+import deepspeed
+import pytest
+from deepspeed.ops.adam import FusedAdam
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, SimpleOptimizer, random_dataloader, SimpleMoEModel, sequence_dataloader
+from deepspeed.utils.torch import required_torch_version
+from deepspeed.accelerator import get_accelerator
+from deepspeed.ops.op_builder import CPUAdamBuilder, FusedLambBuilder
+from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer
+
+if torch.half not in get_accelerator().supported_dtypes():
+    pytest.skip(f"fp16 not supported, valid dtype: {get_accelerator().supported_dtypes()}", allow_module_level=True)
+
+
+class TestLambFP32GradClip(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME],
+                        reason="FusedLambBuilder has not been implemented on this system.")
+    def test(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+class TestLambFP16(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME],
+                        reason="FusedLambBuilder has not been implemented on this system.")
+    def test__basic(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0,
+            "fp16": {
+                "enabled": True
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME],
+                        reason="FusedLambBuilder has not been implemented on this system.")
+    def test_empty_grad(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0,
+            "fp16": {
+                "enabled": True
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=True)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+class TestAdamFP32EmptyGrad(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0,
+            "fp16": {
+                "enabled": False
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=True)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+class TestAdamwFP16Basic(DistributedTest):
+    world_size = 1
+
+    def test(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {"train_batch_size": 1, "steps_per_print": 1, "fp16": {"enabled": True}}
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        optimizer = torch.optim.AdamW(params=model.parameters())
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+class TestFP16OptimizerForMoE(DistributedTest):
+    world_size = 2
+
+    def test_unfused_gradnorm(self, monkeypatch):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        if not required_torch_version(min_version=1.8):
+            pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly")
+
+        config_dict = {"train_batch_size": 2, "steps_per_print": 1, "fp16": {"enabled": True}}
+        hidden_dim = 10
+
+        def mock_unscale_and_clip_grads(total_norm, apply_scale=True):
+            torch_norm_tensor = get_accelerator().FloatTensor([total_norm])
+            all_gather_results = [torch.zeros_like(torch_norm_tensor) for _ in range(dist.get_world_size())]
+            dist.all_gather(all_gather_results, torch_norm_tensor)
+            assert len(set([x.item() for x in all_gather_results])) == 1
+            return 1.0
+
+        # initialize MoE
+        model = SimpleMoEModel(hidden_dim, ep_size=2)
+        optimizer = torch.optim.AdamW(params=model.parameters())
+        engine, optimizer, _, _ = deepspeed.initialize(config=config_dict,
+                                                       model=model,
+                                                       optimizer=optimizer,
+                                                       dist_init_required=False)
+        monkeypatch.setattr(optimizer, 'unscale_and_clip_grads', mock_unscale_and_clip_grads)
+        data_loader = sequence_dataloader(model=engine,
+                                          total_samples=50,
+                                          hidden_dim=hidden_dim,
+                                          device=engine.device,
+                                          dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = engine(batch[0], batch[1])
+            engine.backward(loss)
+            engine.step()
+
+    def test_fused_gradnorm(self, monkeypatch):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        if not required_torch_version(min_version=1.8):
+            pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly")
+
+        config_dict = {"train_batch_size": 2, "steps_per_print": 1, "fp16": {"enabled": True}}
+        hidden_dim = 10
+
+        def mock_unscale_and_clip_grads(grads_groups_flat, total_norm, apply_scale=True):
+            torch_norm_tensor = get_accelerator().FloatTensor([total_norm])
+            all_gather_results = [torch.zeros_like(torch_norm_tensor) for _ in range(dist.get_world_size())]
+            dist.all_gather(all_gather_results, torch_norm_tensor)
+            assert len(set([x.item() for x in all_gather_results])) == 1
+            return 1.0
+
+        # initialize MoE
+        model = SimpleMoEModel(hidden_dim, ep_size=2)
+        param_group = {'params': [p for p in model.parameters()], 'name': 'random-unique-name'}
+        params = split_params_into_different_moe_groups_for_optimizer(param_group)
+        # optimizer = torch.optim.AdamW(params=model.parameters())
+        optimizer = FusedAdam(params=params)
+        engine, optimizer, _, _ = deepspeed.initialize(config=config_dict,
+                                                       model=model,
+                                                       optimizer=optimizer,
+                                                       dist_init_required=False)
+        monkeypatch.setattr(optimizer, 'unscale_and_clip_grads', mock_unscale_and_clip_grads)
+        data_loader = sequence_dataloader(model=engine,
+                                          total_samples=50,
+                                          hidden_dim=hidden_dim,
+                                          device=engine.device,
+                                          dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = engine(batch[0], batch[1])
+            engine.backward(loss)
+            engine.step()
+
+    @pytest.mark.parametrize("fused_lamb_legacy", [(False), (True)])
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME],
+                        reason="FusedLambBuilder has not been implemented on this system.")
+    def test_lamb_gradnorm(self, monkeypatch, fused_lamb_legacy: bool):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        if not required_torch_version(min_version=1.8):
+            pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly")
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "fp16": {
+                "enabled": True
+            },
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            }
+        }
+        hidden_dim = 10
+
+        def mock_unscale_and_clip_grads(total_norm, apply_scale=True):
+            torch_norm_tensor = get_accelerator().FloatTensor([total_norm])
+            all_gather_results = [torch.zeros_like(torch_norm_tensor) for _ in range(dist.get_world_size())]
+            dist.all_gather(all_gather_results, torch_norm_tensor)
+            assert len(set([x.item() for x in all_gather_results])) == 1
+            return 1.0
+
+        # initialize MoE
+        model = SimpleMoEModel(hidden_dim, ep_size=2)
+        engine, optimizer, _, _ = deepspeed.initialize(config=config_dict,
+                                                       model=model,
+                                                       model_parameters=model.parameters(),
+                                                       dist_init_required=False)
+        monkeypatch.setattr(optimizer, 'unscale_and_clip_grads', mock_unscale_and_clip_grads)
+        optimizer.fused_lamb_legacy = fused_lamb_legacy
+        data_loader = sequence_dataloader(model=engine,
+                                          total_samples=50,
+                                          hidden_dim=hidden_dim,
+                                          device=engine.device,
+                                          dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = engine(batch[0], batch[1])
+            engine.backward(loss)
+            engine.step()
+
+
+class TestAdamwFP16EmptyGrad(DistributedTest):
+    world_size = 1
+
+    def test(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {"train_batch_size": 1, "steps_per_print": 1, "fp16": {"enabled": True}}
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        optimizer = torch.optim.AdamW(params=model.parameters())
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+@pytest.mark.parametrize("zero_stage", [1, 2, 3])
+@pytest.mark.parametrize("use_cpu_offload", [True, False])
+class TestAdamFP16ZeroOneCycleCompatibility(DistributedTest):
+    world_size = 1
+
+    def test(self, zero_stage, use_cpu_offload):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "scheduler": {
+                "type": "OneCycle",
+                "params": {
+                    "cycle_first_step_size": 16000,
+                    "cycle_first_stair_count": 8000,
+                    "decay_step_size": 16000,
+                    "cycle_min_lr": 1e-06,
+                    "cycle_max_lr": 3e-05,
+                    "decay_lr_rate": 1e-07,
+                    "cycle_min_mom": 0.85,
+                    "cycle_max_mom": 0.99,
+                    "decay_mom_rate": 0.0
+                }
+            },
+            "fp16": {
+                "enabled": True
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "cpu_offload": use_cpu_offload
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=10,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        model.destroy()
+
+
+@pytest.mark.parametrize("zero_stage", [1, 2, 3])
+@pytest.mark.parametrize("use_cpu_offload", [True, False])
+class TestZeroStaticScale(DistributedTest):
+    world_size = 1
+
+    def test(self, zero_stage, use_cpu_offload, hidden_dim=4):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+
+        config_dict = {
+            "train_batch_size": 4,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 138.
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "cpu_offload": use_cpu_offload
+            }
+        }
+
+        model = SimpleModel(hidden_dim)
+        model, optim, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+        # Ensure the static scaler is configured.
+        assert optim.dynamic_loss_scale == False
+        assert optim.loss_scaler.loss_scale == 138.
+
+        # Now make sure things work..
+        data_loader = random_dataloader(model=model,
+                                        total_samples=10,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        model.destroy()
+
+
+@pytest.mark.parametrize("zero_stage", [1, 2, 3])
+@pytest.mark.parametrize("use_cpu_offload", [True, False])
+class TestZeroAllowUntestedOptimizer(DistributedTest):
+    world_size = 1
+
+    def test(self, zero_stage, use_cpu_offload):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+
+        config_dict = {
+            "train_batch_size": 4,
+            "steps_per_print": 1,
+            "fp16": {
+                "enabled": True,
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "cpu_offload": use_cpu_offload
+            },
+            "zero_allow_untested_optimizer": False,
+            "zero_force_ds_cpu_optimizer": False
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        optimizer = SimpleOptimizer(model.parameters())
+        with pytest.raises(AssertionError):
+            model, optim, _, _ = deepspeed.initialize(config=config_dict,
+                                                      model=model,
+                                                      optimizer=optimizer,
+                                                      model_parameters=model.parameters())
+            model.destroy()
+
+
+@pytest.mark.parametrize("zero_stage", [1, 2, 3])
+@pytest.mark.parametrize("use_cpu_offload", [True, False])
+class TestZeroEmptyPartition(DistributedTest):
+    world_size = 3
+
+    def test(self, zero_stage, use_cpu_offload):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+
+        if zero_stage == 3:
+            pytest.skip("skip for now")
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": 1,
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 8
+            },
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "cpu_offload": use_cpu_offload,
+                "reduce_bucket_size": 100,
+                "allgather_bucket_size": 100
+            }
+        }
+        hidden_dim = 1
+        model = SimpleModel(hidden_dim)
+
+        # Ensure model has 2 parameters, to cause empty partition with DP=3
+        assert len(list(model.parameters())) == 2
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+        # Now make sure things work..
+        data_loader = random_dataloader(model=model,
+                                        total_samples=1,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        model.destroy()
+
+
+@pytest.mark.parametrize("zero_stage", [1, 2, 3])
+@pytest.mark.parametrize("optimizer_constructor", [FusedAdam, torch.optim.Adam])
+class TestZeroSupportedClientOptimizer(DistributedTest):
+    world_size = 1
+
+    def test(self, zero_stage, optimizer_constructor):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "fp16": {
+                "enabled": True
+            },
+            "zero_optimization": {
+                "stage": zero_stage
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        client_optimizer = optimizer_constructor(params=model.parameters())
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=client_optimizer)
+        model.destroy()
+
+
+class TestZero2ReduceScatterOff(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0,
+            "zero_optimization": {
+                "stage": 2,
+                "contiguous_gradients": True,
+                "allgather_bucket_size": 2000000000,
+                "reduce_bucket_size": 200000000,
+                "overlap_comm": False,
+                "reduce_scatter": False
+            },
+            "fp16": {
+                "enabled": True
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+@pytest.mark.parametrize("adam_type", ["Adam", "AdamW"])
+@pytest.mark.parametrize("torch_impl", [True, False])
+class TestFP16AdamTypes(DistributedTest):
+    world_size = 1
+
+    def test(self, adam_type, torch_impl):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 10
+            },
+            "optimizer": {
+                "type": adam_type,
+                "torch_adam": torch_impl,
+                "params": {
+                    "lr": 0.00015
+                }
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=10,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+
+        for _, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+class TestZero3LazyScatter(DistributedTest):
+    world_size = 1
+
+    def test(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 10
+            },
+            "optimizer": {
+                "type": "AdamW",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "zero_optimization": {
+                "stage": 3
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(
+            config=config_dict,
+            model=model,
+            model_parameters=model.parameters(),
+        )
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=10,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+
+        for _, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        model.destroy()
+
+
+@pytest.mark.parametrize('stage', [1, 2, 3])
+class TestZeroEmptyGrad(DistributedTest):
+    world_size = 1
+
+    def test(self, stage):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "fp16": {
+                "enabled": True
+            },
+            "zero_optimization": {
+                "stage": stage
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        optimizer = torch.optim.Adam(model.parameters())
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, optimizer=optimizer)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        model.destroy()

tests/unit/runtime/half_precision/test_fp8.py

@@ -0,0 +1,98 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+import pytest
+from unit.common import DistributedTest, is_rocm_pytorch
+from unit.util import skip_on_arch
+
+try:
+    import transformer_engine.pytorch as transformer_engine
+    from transformer_engine.common import recipe
+except ImportError:
+    pytest.skip("Transformer Engine package is missing, skipping tests", allow_module_level=True)
+
+
+@pytest.mark.parametrize("base_datatype", ["fp16", "bf16", "fp32"])
+class TestFp8ComposabilityAcrossZero(DistributedTest):
+    world_size = 1
+
+    def test(self, base_datatype):
+        skip_on_arch(min_arch=9)
+
+        def run_zero(stage, model_dtype):
+            num_batches = 128
+            batch_size = 16
+            hidden_dim = 768
+            # Have to set seed before model
+            torch.random.manual_seed(42)
+            enable_fp16 = model_dtype == torch.float16
+            enable_bf16 = model_dtype == torch.bfloat16
+            # TransformerEngine Model
+            model = transformer_engine.Linear(hidden_dim, hidden_dim, bias=True, params_dtype=model_dtype)
+
+            # Create FP8 recipe. Note: All input args are optional.
+            fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID,
+                                               amax_history_len=16,
+                                               amax_compute_algo="max")
+            config = {
+                "train_batch_size": batch_size,
+                "gradient_accumulation_steps": 1,
+                "optimizer": {
+                    "type": "Adam",
+                    "params": {
+                        "lr": 0.00001
+                    }
+                },
+                "zero_optimization": {
+                    "stage": stage
+                },
+                "fp16": {
+                    "enabled": enable_fp16,
+                    "loss_scale": 0.1
+                },
+                "bf16": {
+                    "enabled": enable_bf16
+                }
+            }
+            # Init DeepSpeed
+            model, optimizer, _, _ = deepspeed.initialize(args=None,
+                                                          model=model,
+                                                          model_parameters=model.parameters(),
+                                                          config=config)
+
+            batches = torch.randn(num_batches, batch_size, hidden_dim, device=model.device, dtype=model_dtype)
+            for batch in batches:
+                # Enables autocasting for the forward pass
+                with transformer_engine.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+                    out = model(batch)
+                loss = out.mean()
+                model.backward(loss)
+                model.step()
+            return loss
+
+        if base_datatype == "fp16":
+            model_dtype = torch.float16
+        elif base_datatype == "bf16":
+            model_dtype = torch.bfloat16
+        else:
+            model_dtype = torch.float32
+
+        # Set default tolerances
+        rtol, atol = 1e-07, 1e-05
+
+        # Relax tolerance only for ROCm + FP16
+        if is_rocm_pytorch() and base_datatype in ["fp16", "bf16"]:
+            rtol, atol = 1e-07, 1e-04
+
+        # config
+        zero_stage = [0, 1, 2, 3]
+        losses = []
+        for stage in zero_stage:
+            loss = run_zero(stage, model_dtype)
+            losses.append(loss)
+        all_equal = all(torch.allclose(loss, losses[0], rtol, atol) for loss in losses)
+        assert (all_equal)

tests/unit/runtime/half_precision/test_zero_optim_overflow.py

@@ -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()])