chore: import upstream snapshot with attribution

tests/unit/runtime/activation_checkpointing/test_activation_checkpointing.py

@@ -0,0 +1,353 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# TODO: add tests with model parallelism for activation partitioning and other features.
+
+import pytest
+import torch
+import deepspeed
+from deepspeed.pipe import PipelineModule, LayerSpec
+from deepspeed.accelerator import get_accelerator
+from copy import deepcopy
+from unit.common import DistributedTest
+
+ckpt = deepspeed.checkpointing.checkpoint
+
+
+def _compute(module, *inputs, do_checkpoint=False):
+    if do_checkpoint:
+        outputs = ckpt(module, *inputs)
+    else:
+        outputs = module(*inputs)
+
+    if torch.is_tensor(outputs):
+        outputs = (outputs, )
+
+    sum(o.sum() for o in outputs if torch.is_tensor(o) and o.requires_grad).backward()
+
+    grads = [p.grad for p in module.parameters()]
+    input_grads = [inp.grad for inp in inputs if torch.is_tensor(inp)]
+
+    return {
+        'outputs': outputs,
+        'module_grads': grads,
+        'input_grads': input_grads,
+    }
+
+
+def _prep_inputs(*inputs):
+    _inputs = []
+
+    for inp in inputs:
+        inp = deepcopy(inp)
+        if torch.is_tensor(inp):
+            inp = inp.to(get_accelerator().device_name())
+        _inputs.append(inp)
+
+    return tuple(_inputs)
+
+
+def _match_outputs(ref, tgt):
+    assert type(ref) == type(tgt)
+    if type(ref) in [list, tuple]:
+        for x, y in zip(ref, tgt):
+            _match_outputs(x, y)
+    elif not torch.is_tensor(ref):
+        assert ref == tgt
+    elif ref.is_floating_point():
+        assert torch.allclose(ref, tgt)
+    else:
+        assert torch.equal(ref, tgt)
+
+
+def _test_activation_checkpoint(module, *inputs):
+    if get_accelerator().device_name() == "cpu":
+        pytest.skip("CPU accelerator does not support this test yet")
+    # Move to device
+    module.to(get_accelerator().device_name())
+
+    # Get rid of dropouts until we fork the RNG between tests.
+    module.eval()
+
+    module_ = deepcopy(module)
+    inputs_ = _prep_inputs(*inputs)
+    base = _compute(module_, *inputs_, do_checkpoint=False)
+
+    module_ = deepcopy(module)
+    inputs_ = _prep_inputs(*inputs)
+    test = _compute(module_, *inputs_, do_checkpoint=True)
+
+    for group in base.keys():
+        for b, t in zip(base[group], test[group]):
+            _match_outputs(b, t)
+
+
+def _test_activation_checkpoint_ordering(module, expected_ordering, *inputs):
+    if get_accelerator().device_name() == "cpu":
+        pytest.skip("CPU accelerator does not support this test yet")
+    # Move to device
+    module.to(get_accelerator().device_name())
+
+    # Get rid of dropouts until we fork the RNG between tests.
+    module.eval()
+
+    module_ = deepcopy(module)
+    inputs_ = _prep_inputs(*inputs)
+    test = _compute(module_, *inputs_, do_checkpoint=True)
+
+    outputs = test['outputs']
+    test_ordering = []
+    for item in outputs:
+        if type(item) in [list, tuple]:
+            test_ordering += [torch.is_tensor(t) for t in item]
+        else:
+            test_ordering += [torch.is_tensor(item)]
+
+    assert expected_ordering == test_ordering
+
+
+#
+# Helpers
+#
+
+
+class MaskedLinear(torch.nn.Linear):
+
+    def forward(self, x, mask):
+        out = super().forward(x)
+        if mask.is_floating_point():
+            out = out * mask
+        else:
+            # must cast BoolTensor in older torch versions
+            out = out * mask.type_as(out)
+        return out
+
+
+class MaskedLinearSeq(MaskedLinear):
+    """Tests pipeline modules by also returning the mask."""
+
+    def forward(self, x, mask):
+        return super().forward(x, mask), mask
+
+
+class MaskedLinearSeqDup(MaskedLinearSeq):
+    """MaskedLinearSeq, but with more outputs than inputs and in a different order."""
+
+    def forward(self, x, mask):
+        dup = x.clone().detach() * 1.38  # just an arbitrary scaling
+        x, mask = super().forward(x, mask)
+        return dup, x, mask
+
+
+class DropMaskLinear(torch.nn.Linear):
+
+    def forward(self, x, mask):
+        return super().forward(x)
+
+
+class LinearNonTensorInput(torch.nn.Linear):
+
+    def forward(self, x, non_tensor_input):
+        return super().forward(x)
+
+
+class LinearNonTensorOutput(torch.nn.Linear):
+
+    def __init__(self, non_tensor_output):
+        super().__init__(HIDDEN_DIM, HIDDEN_DIM)
+        self.non_tensor_output = non_tensor_output
+
+    def forward(self, x):
+        out = super().forward(x)
+        return out, self.non_tensor_output
+
+
+HIDDEN_DIM = 20
+
+
+def _mixed_mask(size=HIDDEN_DIM):
+    entries = torch.randn(size)
+    mask = torch.where(entries > 0, torch.ones(size), torch.zeros(size))
+    mask = mask.bool()
+    return mask
+
+
+def _bool_to_float(btensor, dtype=torch.float32):
+    """Converts a torch.BoolTensor to an equivalent dtype. """
+    ones = torch.ones(size=btensor.size(), dtype=dtype)
+    zeros = torch.zeros(size=btensor.size(), dtype=dtype)
+    return torch.where(btensor, ones, zeros)
+
+
+#
+# Tests
+#
+
+
+# both bool and float are important, as bool is not differentiable
+@pytest.mark.parametrize('mask', [
+    _mixed_mask(),
+    _bool_to_float(_mixed_mask()),
+])
+class TestActivationCheckpoint(DistributedTest):
+    world_size = 1
+
+    def test_ckpt_inputs1_outputs1(self, mask):
+        module = torch.nn.Linear(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        inputs.requires_grad = True
+        _test_activation_checkpoint(module, inputs)
+
+    def test_ckpt_inputs2_outputs1(self, mask):
+        module = MaskedLinear(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        inputs.requires_grad = True
+        _test_activation_checkpoint(module, inputs, mask)
+
+    def test_ckpt_inputs2_outputs2(self, mask):
+        module = MaskedLinearSeq(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        inputs.requires_grad = True
+        _test_activation_checkpoint(module, inputs, mask)
+
+    def test_ckpt_inputs2_outputs3(self, mask):
+        module = MaskedLinearSeqDup(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        inputs.requires_grad = True
+        _test_activation_checkpoint(module, inputs, mask)
+
+    def test_ckpt_arg_none(self, mask):
+        module = DropMaskLinear(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = (torch.rand(HIDDEN_DIM), None)
+        inputs[0].requires_grad = True
+        _test_activation_checkpoint(module, *inputs)
+
+
+@pytest.mark.parametrize('non_tensor', [None, 2, True, (None, 2.5), (None, True, torch.randn(HIDDEN_DIM))])
+class TestCheckpointNonTensor(DistributedTest):
+    world_size = 1
+
+    def test_ckpt_non_tensor_input(self, non_tensor):
+        module = LinearNonTensorInput(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        inputs.requires_grad = True
+        _test_activation_checkpoint(module, inputs, non_tensor)
+
+    def test_ckpt_non_tensor_output(self, non_tensor):
+        module = LinearNonTensorOutput(non_tensor)
+        inputs = torch.rand(HIDDEN_DIM)
+        inputs.requires_grad = True
+        _test_activation_checkpoint(module, inputs)
+
+
+@pytest.mark.parametrize('non_tensor_output', [
+    None, (torch.randn(HIDDEN_DIM), 2.5), (None, torch.randn(HIDDEN_DIM), True), (None, True, torch.randn(HIDDEN_DIM))
+])
+class TestCheckpointNonTensorOutputOrdering(DistributedTest):
+    world_size = 1
+
+    def test_ckpt_non_tensor_output_ordering(self, non_tensor_output):
+        module = LinearNonTensorOutput(non_tensor_output)
+        inputs = torch.rand(HIDDEN_DIM)
+        inputs.requires_grad = True
+
+        # First return is a tensor
+        ordering = [True]
+        if type(non_tensor_output) in [list, tuple]:
+            ordering += [torch.is_tensor(t) for t in non_tensor_output]
+        else:
+            ordering += [torch.is_tensor(non_tensor_output)]
+        _test_activation_checkpoint_ordering(module, ordering, inputs)
+
+
+class TestCheckpointableLayersConfig(DistributedTest):
+    world_size = 1
+
+    def test_gpt2_checkpointable_layers(self):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU accelerator does not support this test yet")
+
+        # Create a simple topology for testing
+        from deepspeed.runtime.pipe.topology import PipeModelDataParallelTopology
+        topo = PipeModelDataParallelTopology(num_pp=1, num_mp=1, num_dp=1)
+
+        # Create test classes that we want to checkpoint
+        class TestTransformerLayer(torch.nn.Module):
+
+            def forward(self, x):
+                return x
+
+        class ParallelTransformerLayerPipe(TestTransformerLayer):
+            pass
+
+        class GMLPBlock(TestTransformerLayer):
+            pass
+
+        # Create a mock GPT2 model with different layer types
+        class TestGPT2ModelPipe(PipelineModule):
+
+            def __init__(self):
+                self.layers_spec = [
+                    LayerSpec(ParallelTransformerLayerPipe),
+                    LayerSpec(GMLPBlock),
+                    LayerSpec(torch.nn.Linear, 10, 10),  # Should not be checkpointed
+                ]
+
+                super().__init__(layers=self.layers_spec,
+                                 topology=topo,
+                                 checkpointable_layers=["GMLPBlock", "ParallelTransformerLayerPipe"])
+
+        model = TestGPT2ModelPipe()
+        model.to(get_accelerator().device_name())
+
+        # Build layers manually for testing
+        layers = [spec.build() for spec in model.layers_spec]
+
+        # Test that _is_checkpointable returns correct values
+        assert model._is_checkpointable([layers[0]]) == True  # ParallelTransformerLayerPipe
+        assert model._is_checkpointable([layers[1]]) == True  # GMLPBlock
+        assert model._is_checkpointable([layers[2]]) == False  # Linear layer
+
+
+def test_configure_with_contiguous_checkpointing_requires_num_checkpoints():
+    # Regression: ``_configure_defaults`` previously initialized ``num_layers``
+    # to ``False`` while the assert below uses ``is not None``; ``False is not
+    # None`` is True, so the missing-config assert silently passed and a
+    # cryptic ``IndexError`` surfaced later from ``range(num_layers)``. With
+    # the default switched to ``None`` (matching the module-level default),
+    # the helpful assert message fires at the configure() call site.
+    #
+    # ``configure()`` mutates module globals before raising, so snapshot and
+    # restore them around the call to avoid order-dependent failures in other
+    # activation-checkpointing tests sharing the same pytest worker.
+    cp = deepspeed.checkpointing
+    saved = (
+        cp.PARTITION_ACTIVATIONS,
+        cp.CONTIGUOUS_CHECKPOINTING,
+        cp.num_layers,
+        cp.CPU_CHECKPOINT,
+        cp.SYNCHRONIZE,
+        cp.PROFILE_TIME,
+        cp.mpu,
+        cp.deepspeed_checkpointing_enabled,
+    )
+    try:
+        with pytest.raises(AssertionError, match="number of layers"):
+            deepspeed.checkpointing.configure(
+                mpu_=None,
+                partition_activations=True,
+                contiguous_checkpointing=True,
+            )
+    finally:
+        (
+            cp.PARTITION_ACTIVATIONS,
+            cp.CONTIGUOUS_CHECKPOINTING,
+            cp.num_layers,
+            cp.CPU_CHECKPOINT,
+            cp.SYNCHRONIZE,
+            cp.PROFILE_TIME,
+            cp.mpu,
+            cp.deepspeed_checkpointing_enabled,
+        ) = saved

tests/unit/runtime/activation_checkpointing/test_activation_checkpointing_non_reentrant.py

@@ -0,0 +1,122 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# TODO: add tests with model parallelism for activation partitioning and other features.
+
+import sys
+import torch
+import pytest
+from importlib import util
+
+from deepspeed.runtime.activation_checkpointing.checkpointing import non_reentrant_checkpoint
+from unit.common import DistributedTest
+
+# the hack to clone the module `test_activation_checkpointing` and inject
+# `non_reentrant_checkpoint` as the `ckpt` of the origin test module
+ORG_SPEC = util.find_spec('test_activation_checkpointing')
+test_act_ckpt = util.module_from_spec(ORG_SPEC)
+ORG_SPEC.loader.exec_module(test_act_ckpt)
+sys.modules['test_act_ckpt'] = test_act_ckpt
+test_act_ckpt.ckpt = non_reentrant_checkpoint
+
+HIDDEN_DIM = test_act_ckpt.HIDDEN_DIM
+
+MaskedLinear = test_act_ckpt.MaskedLinear
+MaskedLinearSeq = test_act_ckpt.MaskedLinearSeq
+MaskedLinearSeqDup = test_act_ckpt.MaskedLinearSeqDup
+DropMaskLinear = test_act_ckpt.DropMaskLinear
+LinearNonTensorInput = test_act_ckpt.LinearNonTensorInput
+LinearNonTensorOutput = test_act_ckpt.LinearNonTensorOutput
+
+_test_activation_checkpoint = test_act_ckpt._test_activation_checkpoint
+_mixed_mask = test_act_ckpt._mixed_mask
+_bool_to_float = test_act_ckpt._bool_to_float
+_test_activation_checkpoint_ordering = test_act_ckpt._test_activation_checkpoint_ordering
+
+
+class TestActivationCheckpointWithGrad(test_act_ckpt.TestActivationCheckpoint):
+    """test `non_reentrant_checkpoint` can still checkpoint activations for inputs with grad"""
+    pass
+
+
+class TestCheckpointNonTensorWithGrad(test_act_ckpt.TestCheckpointNonTensor):
+    """test `non_reentrant_checkpoint` can still checkpoint activations for inputs with grad"""
+    pass
+
+
+class TestCheckpointNonTensorOutputOrderingWithGrad(test_act_ckpt.TestCheckpointNonTensorOutputOrdering):
+    """test `non_reentrant_checkpoint` can still checkpoint activations for inputs with grad"""
+    pass
+
+
+# below classes are used to test the graph with inputs have no grad and parameters has grad, namely partial graph?
+@pytest.mark.parametrize('mask', [
+    _mixed_mask(),
+    _bool_to_float(_mixed_mask()),
+])
+class TestActivationCheckpointWithoutGrad(DistributedTest):
+    """test all input tensors without grad"""
+    world_size = 1
+
+    def test_ckpt_inputs1_outputs1(self, mask):
+        module = torch.nn.Linear(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        _test_activation_checkpoint(module, inputs)
+
+    def test_ckpt_inputs2_outputs1(self, mask):
+        module = MaskedLinear(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        _test_activation_checkpoint(module, inputs, mask)
+
+    def test_ckpt_inputs2_outputs2(self, mask):
+        module = MaskedLinearSeq(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        _test_activation_checkpoint(module, inputs, mask)
+
+    def test_ckpt_inputs2_outputs3(self, mask):
+        module = MaskedLinearSeqDup(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        _test_activation_checkpoint(module, inputs, mask)
+
+    def test_ckpt_arg_none(self, mask):
+        module = DropMaskLinear(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = (torch.rand(HIDDEN_DIM), None)
+        _test_activation_checkpoint(module, *inputs)
+
+
+@pytest.mark.parametrize('non_tensor', [None, 2, True, (None, 2.5), (None, True, torch.randn(HIDDEN_DIM))])
+class TestCheckpointNonTensorWithoutGrad(DistributedTest):
+    """test all input tensors without grad"""
+    world_size = 1
+
+    def test_ckpt_non_tensor_input(self, non_tensor):
+        module = LinearNonTensorInput(HIDDEN_DIM, HIDDEN_DIM)
+        inputs = torch.rand(HIDDEN_DIM)
+        _test_activation_checkpoint(module, inputs, non_tensor)
+
+    def test_ckpt_non_tensor_output(self, non_tensor):
+        module = LinearNonTensorOutput(non_tensor)
+        inputs = torch.rand(HIDDEN_DIM)
+        _test_activation_checkpoint(module, inputs)
+
+
+@pytest.mark.parametrize('non_tensor_output', [
+    None, (torch.randn(HIDDEN_DIM), 2.5), (None, torch.randn(HIDDEN_DIM), True), (None, True, torch.randn(HIDDEN_DIM))
+])
+class TestCheckpointNonTensorOutputOrderingWithoutGrad(DistributedTest):
+    """test all input tensors without grad"""
+    world_size = 1
+
+    def test_ckpt_non_tensor_output_ordering(self, non_tensor_output):
+        module = LinearNonTensorOutput(non_tensor_output)
+        inputs = torch.rand(HIDDEN_DIM)
+
+        # First return is a tensor
+        ordering = [True]
+        if type(non_tensor_output) in [list, tuple]:
+            ordering += [torch.is_tensor(t) for t in non_tensor_output]
+        else:
+            ordering += [torch.is_tensor(non_tensor_output)]
+        _test_activation_checkpoint_ordering(module, ordering, inputs)

tests/unit/runtime/comm/test_coalesced_collectives.py

@@ -0,0 +1,161 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+unit tests for coalesced collectives
+"""
+
+import torch
+import deepspeed
+import deepspeed.comm as dist
+from deepspeed.runtime.comm.coalesced_collectives import reduce_scatter_coalesced, all_to_all_quant_reduce
+from deepspeed.accelerator import get_accelerator
+import pytest
+
+from unit.common import DistributedTest
+
+
+class TestReduceScatterCoalesced(DistributedTest):
+    world_size = 2
+
+    def test_single_input(self):
+        input = torch.full((6, ), dist.get_rank(), dtype=torch.half, device=get_accelerator().current_device_name())
+
+        (output, ) = reduce_scatter_coalesced([input], dist.get_world_group())
+
+        assert output.shape == (3, )
+        assert torch.allclose(output, torch.full_like(output, 0.5))
+
+    def test_two_inputs(self):
+        tensor_kwargs = {"device": get_accelerator().current_device_name(), "dtype": torch.half}
+        inputs = [
+            dist.get_rank() * torch.arange(0, 6, **tensor_kwargs),
+            dist.get_rank() * torch.arange(6, 9, **tensor_kwargs),
+        ]
+
+        output1, output2 = reduce_scatter_coalesced(inputs, dist.get_world_group())
+
+        if dist.get_rank() == 0:
+            assert output1.shape == (3, )
+            assert torch.allclose(output1, torch.arange(0, 3, **tensor_kwargs) / 2)
+            assert output2.shape == (2, )
+            assert torch.allclose(output2, torch.arange(6, 8, **tensor_kwargs) / 2)
+        elif dist.get_rank() == 1:
+            assert output1.shape == (3, )
+            assert torch.allclose(output1, torch.arange(3, 6, **tensor_kwargs) / 2)
+            assert output2.shape == (1, )
+            assert torch.allclose(output2, torch.arange(8, 9, **tensor_kwargs) / 2)
+
+
+class TestReduceScatterCoalescedTensorSmallerThanWorldSize(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        input = torch.zeros((1, ), dtype=torch.half, device=get_accelerator().current_device_name())
+
+        (output, ) = reduce_scatter_coalesced([input], dist.get_world_group())
+
+        if dist.get_rank() == 0:
+            assert output.shape == (1, )
+            assert torch.allclose(output, torch.zeros_like(output))
+        elif dist.get_rank() == 1:
+            assert output.shape == (0, )
+
+
+# Currently we cannot test all_to_all_quant_reduce in non-fallback cases because we don't support multinodes tests.
+class TestAllToAllQuantReduceFallback(DistributedTest):
+    world_size = 2
+
+    def test_1d_tensor(self):
+        # case 1: 1D tensor
+        input = torch.zeros((10, ), dtype=torch.half, device=get_accelerator().current_device_name())
+        from deepspeed.ops.op_builder import QuantizerBuilder
+        if not deepspeed.ops.__compatible_ops__[QuantizerBuilder.NAME]:
+            pytest.skip("QuantizerBuilder is not implemented")
+        output = all_to_all_quant_reduce([input], {})[0]
+
+        if dist.get_rank() == 0:
+            assert output.shape == (5, )
+            assert torch.allclose(output, torch.zeros_like(output))
+        elif dist.get_rank() == 1:
+            assert output.shape == (5, )
+            assert torch.allclose(output, torch.zeros_like(output))
+
+    def test_non_divisible(self):
+        # case 2: tensor size not divisible by global_world_size
+        input = torch.zeros((7, 7), dtype=torch.half, device=get_accelerator().current_device_name())
+        from deepspeed.ops.op_builder import QuantizerBuilder
+        if not deepspeed.ops.__compatible_ops__[QuantizerBuilder.NAME]:
+            pytest.skip("QuantizerBuilder is not implemented")
+        output = all_to_all_quant_reduce([input], {})[0]
+
+        if dist.get_rank() == 0:
+            assert output.shape == (25, )
+            assert torch.allclose(output, torch.zeros_like(output))
+        elif dist.get_rank() == 1:
+            assert output.shape == (24, )
+            assert torch.allclose(output, torch.zeros_like(output))
+
+
+class TestLocoQuantized(DistributedTest):
+
+    world_size = 1
+
+    @pytest.mark.parametrize("num_bits", [4, 8])
+    @pytest.mark.parametrize("tensor_size", [(16, 16), (64, 64)])
+    @pytest.mark.parametrize("devices_per_node", [4, 8])
+    def test_loco_quantized_reduction(self, num_bits, tensor_size, devices_per_node):
+        from deepspeed.ops.op_builder import QuantizerBuilder
+        if not deepspeed.ops.__compatible_ops__[QuantizerBuilder.NAME]:
+            pytest.skip("QuantizerBuilder is not implemented")
+
+        quantizer_module = QuantizerBuilder().load()
+
+        tensor = torch.randn(tensor_size, device='cuda', dtype=torch.half)
+
+        num_nodes = 2  # Fake world size
+        total_elements = tensor.numel()
+        total_devices = devices_per_node * num_nodes
+        num_groups = max(tensor.shape[0], tensor.shape[1], total_devices)
+
+        # Initialize error_feedback tensor
+        error_feedback = torch.randn(tensor_size, device=tensor.device, dtype=tensor.dtype)
+        error_feedback_ori = error_feedback.clone()
+        # Swizzle the original tensor
+        tensor_reshaped = tensor.reshape(num_nodes, devices_per_node, total_elements // total_devices)
+        swizzled_tensor = tensor_reshaped.permute(1, 0, 2).reshape(tensor.size())
+
+        # Perform loco_swizzle_quant
+        output, scales = quantizer_module.loco_swizzle_quant(tensor, error_feedback, 0.0, num_groups, num_bits,
+                                                             quantizer_module.Symmetric, 1, num_nodes,
+                                                             devices_per_node)
+
+        # Compare swizzled_tensor with the output of loco_swizzle_quant
+        dequantized = quantizer_module.dequantize(output, scales, scales.numel(), num_bits,
+                                                  quantizer_module.Symmetric).view(tensor.size())
+
+        assert torch.allclose(swizzled_tensor + error_feedback_ori, dequantized + error_feedback)
+
+        # Calculate elements per group and groups per partition
+        elements_per_group = total_elements // num_groups
+        groups_per_partition = num_groups // devices_per_node
+
+        # Reshape dequantized data to match the grouping in loco_quantized_reduction
+        dequantized_reshaped = dequantized.view(devices_per_node, groups_per_partition, elements_per_group)
+
+        # Perform reduction across devices_per_node dimension
+        reduced_dequantized = dequantized_reshaped.cumsum(dim=0)[-1]
+        # Initialize error_feedback tensor
+        error_feedback = torch.randn(reduced_dequantized.shape, device=tensor.device, dtype=dequantized.dtype)
+        error_feedback_ori = error_feedback.clone()
+
+        # perform loco_quantized_reduction
+        output, scales = quantizer_module.loco_quantized_reduction(output, scales, error_feedback, 0.0, num_groups,
+                                                                   num_groups // devices_per_node, num_bits,
+                                                                   quantizer_module.Symmetric, devices_per_node)
+
+        dequantized_reduced = quantizer_module.dequantize(output, scales, scales.numel(), num_bits,
+                                                          quantizer_module.Symmetric).view(error_feedback.size())
+
+        assert torch.allclose(reduced_dequantized + error_feedback_ori, dequantized_reduced + error_feedback)

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

tests/unit/runtime/pipe/test_pipe.py

@@ -0,0 +1,252 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import copy
+import torch.nn as nn
+import pytest
+
+import torch
+
+import deepspeed
+import deepspeed.comm as dist
+from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
+from deepspeed.runtime.pipe.module import PipelineModule
+from unit.alexnet_model import AlexNetPipe, train_cifar
+from unit.common import DistributedTest
+from unit.util import skip_on_arch, no_child_process_in_deepspeed_io
+
+PipeTopo = PipeDataParallelTopology
+
+config_dict = {
+    "train_batch_size": 4,
+    "grandient_accumulation_steps": 1,
+    "steps_per_print": 20,
+    "optimizer": {
+        "type": "Adam",
+        "params": {
+            "lr": 0.001,
+            "betas": [0.9, 0.999],
+            "eps": 1e-8,
+            "weight_decay": 3e-7
+        }
+    },
+    "zero_optimization": {
+        "stage": 0
+    },
+    "fp16": {
+        "enabled": False
+    },
+    "pipeline": {
+        "seed_layers": True,
+        "activation_checkpoint_interval": 1
+    }
+}
+
+
+def rel_diff(A, B):
+    return abs(A - B) / abs(A)
+
+
+@pytest.mark.parametrize('topo_config', [
+    {
+        "num_pp": 1,
+        "num_dp": 4
+    },
+    {
+        "num_pp": 2,
+        "num_dp": 2
+    },
+    {
+        "num_pp": 4,
+        "num_dp": 1
+    },
+])
+class TestPipeCifar10(DistributedTest):
+    world_size = 4
+
+    def test_pipe_base(self, topo_config):
+        skip_on_arch(min_arch=7)
+        topo = PipeTopo(**topo_config)
+        steps = 100  # must be >=100
+
+        # Allocate model for consistent initial weights.
+        init_net = AlexNetPipe()
+
+        base_net = copy.deepcopy(init_net)
+        base_model = PipelineModule(layers=base_net.to_layers(), num_stages=1, loss_fn=nn.CrossEntropyLoss())
+
+        # Train with just data parallelism
+        base_losses = train_cifar(base_model, config=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled'])
+
+        test_net = copy.deepcopy(init_net)
+        test_model = PipelineModule(layers=test_net.to_layers(), topology=topo, loss_fn=nn.CrossEntropyLoss())
+
+        test_losses = train_cifar(test_model, config=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled'])
+
+        abs_diffs = [l0 - l1 for l0, l1 in zip(base_losses, test_losses)]
+        rel_diffs = [rel_diff(l0, l1) for l0, l1 in zip(base_losses, test_losses)]
+        if dist.get_rank() == 0:
+            print(f'abs min={min(abs_diffs)} max={max(abs_diffs)} avg={sum(abs_diffs)/len(abs_diffs)}')
+            print(f'rel min={min(rel_diffs)} max={max(rel_diffs)} avg={sum(rel_diffs)/len(rel_diffs)}')
+            print(f'first: base={base_losses[0]} test={test_losses[0]} abs={abs_diffs[0]} rel={rel_diffs[0]}')
+
+            for lastX in [1, 10, 100]:
+                base_avg = sum(base_losses[-lastX:]) / lastX
+                test_avg = sum(test_losses[-lastX:]) / lastX
+                print(
+                    f'last-{lastX}: base={base_avg} test={test_avg} abs={base_avg - test_avg} rel={rel_diff(base_avg, test_avg)}'
+                )
+
+        lastX = 100
+        base = base_losses[-lastX:]
+        base_avg = sum(base) / len(base)
+        test = test_losses[-lastX:]
+        test_avg = sum(test) / len(test)
+        assert rel_diff(base_avg, test_avg) < 0.05  # Originally 0.03, but seeing instability with AMD results
+
+    # def _check_model_params_equal(self, model1, model2):
+    #     for p1, p2 in zip(model1.parameters(), model2.parameters()):
+    #         if p1.data.ne(p2.data).sum() > 0:
+    #             assert False, f"model params not equal"
+
+    def test_pipe_use_reentrant(self, topo_config):
+        skip_on_arch(min_arch=7)
+
+        topo = PipeTopo(**topo_config)
+        steps = 100  # must be >=100
+
+        # Allocate model for consistent initial weights.
+        init_net = AlexNetPipe()
+
+        # Train with not set use_reentrant, default: True
+        base_net = copy.deepcopy(init_net)
+        base_model = PipelineModule(layers=base_net.to_layers(), topology=topo, loss_fn=nn.CrossEntropyLoss())
+        base_losses = train_cifar(base_model, config=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled'])
+
+        # Train with set use_reentrant=False, this will use ``non_reentrant_checkpoint``
+        test_config_dict = copy.deepcopy(config_dict)
+        test_config_dict['pipeline']['use_reentrant'] = False
+        test_net = copy.deepcopy(init_net)
+        test_model = PipelineModule(layers=test_net.to_layers(), topology=topo, loss_fn=nn.CrossEntropyLoss())
+        test_losses = train_cifar(test_model,
+                                  config=test_config_dict,
+                                  num_steps=steps,
+                                  fp16=config_dict['fp16']['enabled'])
+
+        abs_diffs = [l0 - l1 for l0, l1 in zip(base_losses, test_losses)]
+        rel_diffs = [rel_diff(l0, l1) for l0, l1 in zip(base_losses, test_losses)]
+        if dist.get_rank() == 0:
+            print(f'abs min={min(abs_diffs)} max={max(abs_diffs)} avg={sum(abs_diffs)/len(abs_diffs)}')
+            print(f'rel min={min(rel_diffs)} max={max(rel_diffs)} avg={sum(rel_diffs)/len(rel_diffs)}')
+            print(f'first: base={base_losses[0]} test={test_losses[0]} abs={abs_diffs[0]} rel={rel_diffs[0]}')
+
+            for lastX in [1, 10, 100]:
+                base_avg = sum(base_losses[-lastX:]) / lastX
+                test_avg = sum(test_losses[-lastX:]) / lastX
+                print(
+                    f'last-{lastX}: base={base_avg} test={test_avg} abs={base_avg - test_avg} rel={rel_diff(base_avg, test_avg)}'
+                )
+        lastX = 100
+        base = base_losses[-lastX:]
+        base_avg = sum(base) / len(base)
+        test = test_losses[-lastX:]
+        test_avg = sum(test) / len(test)
+        assert rel_diff(base_avg, test_avg) < 0.05
+
+        # the following check could passed on higher version docker: nvcr.io/nvidia/pytorch:23.07-py3(torch2.1.0 cuda12.1)
+        # Check if models have same weights after training
+        # self._check_model_params_equal(base_model, test_model)
+
+
+class DynamicShapeTestLayer(nn.Module):
+
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.fc = nn.Linear(hidden_size, hidden_size)
+        self.shapes = set()
+
+    def forward(self, x):
+        self.shapes.add(x.shape)
+        y = self.fc(x)
+        return y
+
+
+class DynamicShapeTestModel(nn.Module):
+
+    def __init__(self, n_layers, hidden_size):
+        super().__init__()
+        self.layers = nn.ModuleList([DynamicShapeTestLayer(hidden_size) for _ in range(n_layers)])
+
+
+@pytest.mark.parametrize('topo_config', [
+    {
+        "num_pp": 1,
+        "num_dp": 4
+    },
+    {
+        "num_pp": 2,
+        "num_dp": 2
+    },
+    {
+        "num_pp": 4,
+        "num_dp": 1
+    },
+])
+class TestPipeDynamicShape(DistributedTest):
+    world_size = 4
+
+    def test_pipe_base(self, topo_config):
+        """This test checks if the pipeline engine can handle dynamic shapes correctly.
+        We pass inputs of different shapes to the pipeline engine.
+        """
+
+        n_iter = 10
+        n_layers = 4
+        n_samples = 1024
+        batch_size = 4
+        channel_dims = [8, 16, 32, 64]
+        hidden_size = 16
+
+        topo = PipeTopo(**topo_config)
+
+        model = DynamicShapeTestModel(n_layers, hidden_size)
+        model = PipelineModule(layers=model.layers, topology=topo, loss_fn=nn.MSELoss(), dynamic_shape=True)
+
+        # Each batch has different channel dim but we use the same channel dim in the same batch
+        xs = [
+            torch.randn(channel_dims[(i // batch_size) % len(channel_dims)], hidden_size, dtype=torch.float32)
+            for i in range(n_samples)
+        ]
+        ys = [torch.randn_like(x) for x in xs]
+
+        class CustomDataset(torch.utils.data.Dataset):
+
+            def __init__(self, xs, ys):
+                self.xs = xs
+                self.ys = ys
+
+            def __len__(self):
+                return len(self.xs)
+
+            def __getitem__(self, idx):
+                return self.xs[idx], self.ys[idx]
+
+        dataset = CustomDataset(xs, ys)
+
+        config_dict["train_batch_size"] = batch_size
+
+        with no_child_process_in_deepspeed_io():
+            engine, _, _, _ = deepspeed.initialize(config=config_dict,
+                                                   model=model,
+                                                   model_parameters=[p for p in model.parameters()],
+                                                   training_data=dataset)
+
+        for _ in range(n_iter):
+            _ = engine.train_batch()
+
+        # Check if all layers have seen different shapes
+        for layer in model.modules():
+            if isinstance(layer, DynamicShapeTestLayer):
+                assert len(layer.shapes) > 1

tests/unit/runtime/pipe/test_pipe_schedule.py

@@ -0,0 +1,143 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import deepspeed.runtime.pipe.schedule as schedule
+
+
+def _count_type(cmds, classtype):
+    return len(list(filter(lambda c: type(c) == classtype, cmds)))
+
+
+def test_pipe_inference_schedule_singlestage():
+    sched = schedule.InferenceSchedule(micro_batches=4, stages=1, stage_id=0)
+    assert sched.num_micro_batches == 4
+    full = list(iter(sched))
+    for idx, cmds in enumerate(full):
+        assert len(cmds) == 2
+        assert type(cmds[0]) == schedule.LoadMicroBatch
+        assert type(cmds[1]) == schedule.ForwardPass
+        assert cmds[0].buffer_id == cmds[1].buffer_id
+    assert len(full) == sched.num_micro_batches
+
+
+def test_pipe_train_schedule_singlestage():
+    sched = schedule.TrainSchedule(micro_batches=4, stages=1, stage_id=0)
+    assert sched.num_micro_batches == 4
+    full = list(iter(sched))
+    for idx, cmds in enumerate(full):
+        if (idx % 2) != 0:
+            assert (len(cmds) == 1) or (len(cmds) == 4)
+            assert type(cmds[0]) == schedule.BackwardPass
+        else:
+            assert len(cmds) == 2
+            assert type(cmds[0]) == schedule.LoadMicroBatch
+            assert type(cmds[1]) == schedule.ForwardPass
+            assert cmds[0].buffer_id == cmds[1].buffer_id
+    assert len(full) == sched.num_micro_batches * 2
+
+
+@pytest.mark.parametrize('micro_batches', [1, 3, 8, 10])
+def test_pipe_inference_schedule_firststage(micro_batches, stages=3):
+    sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=0)
+    assert sched.num_micro_batches == micro_batches
+    full = list(iter(sched))
+    for idx, cmds in enumerate(full):
+        # Ensure we don't send an activation the first step
+        if idx == 0:
+            assert len(cmds) == 2
+            assert type(cmds[0]) == schedule.LoadMicroBatch
+            assert type(cmds[1]) == schedule.ForwardPass
+            assert cmds[0].buffer_id == cmds[1].buffer_id
+            continue
+
+        # the last active step is only a send
+        if idx == sched.num_micro_batches:
+            assert len(cmds) == 1
+            assert type(cmds[0]) == schedule.SendActivation
+            continue
+
+        # no work later on
+        if idx > sched.num_micro_batches:
+            assert len(cmds) == 0
+            continue
+
+        # Normally we need to load/forward/send
+        assert len(cmds) == 3
+        assert _count_type(cmds, schedule.LoadMicroBatch) == 1
+        assert _count_type(cmds, schedule.ForwardPass) == 1
+        assert _count_type(cmds, schedule.SendActivation) == 1
+    assert len(full) == micro_batches + stages - 1
+
+
+@pytest.mark.parametrize('micro_batches', [1, 3, 8, 10])
+def test_pipe_inference_schedule_midstage(micro_batches, stages=3):
+    sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=1)
+
+    full = list(iter(sched))
+    for idx, cmds in enumerate(full):
+        if idx < sched.stage:
+            assert len(cmds) == 0
+            continue
+        if idx == sched.stage + sched.num_micro_batches:
+            assert len(cmds) == 1
+            assert type(cmds[0]) == schedule.SendActivation
+            continue
+        if idx > sched.stage + sched.num_micro_batches:
+            assert len(cmds) == 0
+            continue
+        assert _count_type(cmds, schedule.LoadMicroBatch) == 0
+        assert _count_type(cmds, schedule.ForwardPass) == 1
+        assert _count_type(cmds, schedule.RecvActivation) == 1
+        if idx > sched.stage:
+            assert _count_type(cmds, schedule.SendActivation) == 1
+    assert len(full) == micro_batches + stages - 1
+
+
+@pytest.mark.parametrize('micro_batches', [1, 3, 8, 10])
+def test_pipe_inference_schedule_laststage(micro_batches, stages=3):
+    sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=2)
+    full = list(iter(sched))
+    for idx, cmds in enumerate(full):
+        if idx < sched.stage or idx > sched.stage + sched.num_micro_batches:
+            assert len(cmds) == 0
+            continue
+        assert _count_type(cmds, schedule.LoadMicroBatch) == 1
+        assert _count_type(cmds, schedule.ForwardPass) == 1
+        assert _count_type(cmds, schedule.RecvActivation) == 1
+        assert _count_type(cmds, schedule.SendActivation) == 0
+    assert len(full) == micro_batches + stages - 1
+
+
+def test_pipe_schedule_firststage():
+    sched = schedule.TrainSchedule(micro_batches=8, stages=3, stage_id=0)
+    for cmds in sched:
+        assert all(instr.__class__ != schedule.SendGrad for instr in cmds)
+        assert all(instr.__class__ != schedule.RecvActivation for instr in cmds)
+        for instr in cmds:
+            if isinstance(instr, schedule.BufferOpInstruction):
+                assert 0 <= instr.buffer_id < sched.num_pipe_buffers()
+
+
+def test_pipe_schedule_laststage():
+    sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=2)
+    assert len(list(iter(sched))) == 2 * (sched.micro_batches + sched.stages - 1)
+    for cmds in sched:
+        assert all(instr.__class__ != schedule.SendActivation for instr in cmds)
+        assert all(instr.__class__ != schedule.RecvGrad for instr in cmds)
+
+
+def test_pipe_stagequery():
+    sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=0)
+    assert sched.is_first_stage
+    assert not sched.is_last_stage
+
+    sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=1)
+    assert not sched.is_first_stage
+    assert not sched.is_last_stage
+
+    sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=2)
+    assert not sched.is_first_stage
+    assert sched.is_last_stage

tests/unit/runtime/pipe/test_topology.py

@@ -0,0 +1,226 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+import torch
+import deepspeed.comm as dist
+
+from deepspeed.runtime.pipe.topology import PipelineParallelGrid as Grid
+from deepspeed.runtime.pipe.topology import ProcessTopology as Topo
+from deepspeed.runtime.pipe.topology import _prime_factors
+
+from deepspeed.accelerator import get_accelerator
+from unit.common import DistributedTest
+
+
+def test_topology_2d():
+    topo = Topo(axes=['row', 'col'], dims=[2, 2])
+
+    assert topo.world_size() == 4
+
+    assert topo.get_rank(row=0, col=0) == 0
+    assert topo.get_rank(row=0, col=1) == 1
+    assert topo.get_rank(row=1, col=0) == 2
+    assert topo.get_rank(row=1, col=1) == 3
+
+    assert topo.get_axis_list(axis='row', idx=0) == [0, 1]
+    assert topo.get_axis_list(axis='row', idx=1) == [2, 3]
+    assert topo.get_axis_list(axis='col', idx=0) == [0, 2]
+    assert topo.get_axis_list(axis='col', idx=1) == [1, 3]
+
+
+def test_topology_dims():
+    topo = Topo(axes=['a', 'b', 'c'], dims=[2, 3, 4])
+    assert topo.world_size() == 24
+    assert topo.get_dim('a') == 2
+    assert topo.get_dim('b') == 3
+    assert topo.get_dim('c') == 4
+
+
+def test_topology_match():
+    topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2])
+    print(topo.filter_match(pipe=0, data=1))
+    assert topo.filter_match(pipe=0, data=1) == [2, 3]
+    print([topo.get_coord(r) for r in topo.filter_match(pipe=0, data=1)])
+
+
+def test_topology_rank_repr():
+    topo = Topo(axes=['a', 'b'], dims=[2, 2])
+    assert topo.get_rank_repr(rank=0) == 'a_00-b_00'
+    assert topo.get_rank_repr(rank=1) == 'a_00-b_01'
+    assert topo.get_rank_repr(rank=2) == 'a_01-b_00'
+    assert topo.get_rank_repr(rank=3) == 'a_01-b_01'
+
+    assert topo.get_rank_repr(rank=3, inner_sep='+') == 'a+01-b+01'
+    assert topo.get_rank_repr(rank=3, inner_sep='🤗', outer_sep='_JEFF_') == 'a🤗01_JEFF_b🤗01'
+
+    topo = Topo(axes=['pipe', 'data'], dims=[2, 2])
+    assert topo.get_rank_repr(rank=0) == ''
+    assert topo.get_rank_repr(rank=1) == ''
+    assert topo.get_rank_repr(rank=2) == ''
+    assert topo.get_rank_repr(rank=3) == ''
+
+    assert topo.get_rank_repr(rank=0, omit_axes=['pipe']) == 'data_00'
+    assert topo.get_rank_repr(rank=1, omit_axes=['pipe']) == 'data_01'
+    assert topo.get_rank_repr(rank=2, omit_axes=['pipe']) == 'data_00'
+    assert topo.get_rank_repr(rank=3, omit_axes=['pipe']) == 'data_01'
+
+    assert topo.get_rank_repr(rank=0, omit_axes=[]) == 'pipe_00-data_00'
+    assert topo.get_rank_repr(rank=1, omit_axes=[]) == 'pipe_00-data_01'
+    assert topo.get_rank_repr(rank=2, omit_axes=[]) == 'pipe_01-data_00'
+    assert topo.get_rank_repr(rank=3, omit_axes=[]) == 'pipe_01-data_01'
+
+    topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2])
+    assert topo.get_rank_repr(rank=0) == 'model_00'
+    assert topo.get_rank_repr(rank=1) == 'model_01'
+    assert topo.get_rank_repr(rank=2) == 'model_00'
+    assert topo.get_rank_repr(rank=3) == 'model_01'
+    assert topo.get_rank_repr(rank=4) == 'model_00'
+    assert topo.get_rank_repr(rank=5) == 'model_01'
+    assert topo.get_rank_repr(rank=6) == 'model_00'
+    assert topo.get_rank_repr(rank=7) == 'model_01'
+
+
+def test_topology_3d():
+    topo = Topo(axes=['a', 'b', 'c'], dims=[2, 2, 2])
+
+    assert topo.get_rank(a=0, b=0, c=0) == 0
+    assert topo.get_rank(a=0, b=0, c=1) == 1
+    assert topo.get_rank(a=0, b=1, c=0) == 2
+    assert topo.get_rank(a=0, b=1, c=1) == 3
+    assert topo.get_rank(a=1, b=0, c=0) == 4
+    assert topo.get_rank(a=1, b=0, c=1) == 5
+    assert topo.get_rank(a=1, b=1, c=0) == 6
+    assert topo.get_rank(a=1, b=1, c=1) == 7
+
+    assert topo.get_axis_list('a', 0) == [0, 1, 2, 3]
+    assert topo.get_axis_list('a', 1) == [4, 5, 6, 7]
+    assert topo.get_axis_list('b', 0) == [0, 1, 4, 5]
+    assert topo.get_axis_list('b', 1) == [2, 3, 6, 7]
+    assert topo.get_axis_list('c', 0) == [0, 2, 4, 6]
+    assert topo.get_axis_list('c', 1) == [1, 3, 5, 7]
+
+    assert topo.get_coord(0) == topo.ProcessCoord(0, 0, 0)
+    assert topo.get_coord(1) == topo.ProcessCoord(0, 0, 1)
+    assert topo.get_coord(2) == topo.ProcessCoord(0, 1, 0)
+    assert topo.get_coord(3) == topo.ProcessCoord(0, 1, 1)
+    assert topo.get_coord(4) == topo.ProcessCoord(1, 0, 0)
+    assert topo.get_coord(5) == topo.ProcessCoord(1, 0, 1)
+    assert topo.get_coord(6) == topo.ProcessCoord(1, 1, 0)
+    assert topo.get_coord(7) == topo.ProcessCoord(1, 1, 1)
+
+    assert topo.filter_match(a=0) == [0, 1, 2, 3]
+    assert topo.filter_match(b=1, c=1) == [3, 7]
+    assert topo.filter_match(a=1, b=1, c=1) == [7]
+
+    # Easy access method
+    assert topo.get_coord(0).a == 0
+
+
+def test_topology_comm_list():
+    topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2])
+
+    assert topo.get_rank(pipe=0, data=0, model=0) == 0
+    assert topo.get_rank(pipe=0, data=0, model=1) == 1
+    assert topo.get_rank(pipe=0, data=1, model=0) == 2
+    assert topo.get_rank(pipe=0, data=1, model=1) == 3
+    assert topo.get_rank(pipe=1, data=0, model=0) == 4
+    assert topo.get_rank(pipe=1, data=0, model=1) == 5
+    assert topo.get_rank(pipe=1, data=1, model=0) == 6
+    assert topo.get_rank(pipe=1, data=1, model=1) == 7
+
+    pipe_list = [
+        [0, 4],  # data=0, model=0
+        [1, 5],  # data=0, model=1
+        [2, 6],  # data=1, model=0
+        [3, 7],  # data=1, model=1
+    ]
+    assert topo.get_axis_comm_lists('pipe') == pipe_list
+
+    data_list = [
+        [0, 2],  # pipe=0, model=0
+        [1, 3],  # pipe=0, model=1
+        [4, 6],  # pipe=1, model=0
+        [5, 7],  # pipe=1, model=1
+    ]
+    assert topo.get_axis_comm_lists('data') == data_list
+
+    model_list = [
+        [0, 1],  # pipe=0, data=0
+        [2, 3],  # pipe=0, data=1
+        [4, 5],  # pipe=1, data=0
+        [6, 7],  # pipe=1, data=1
+    ]
+    assert topo.get_axis_comm_lists('model') == model_list
+
+    # Handle nonsense. We don't want to RuntimeError because it allows us to write more
+    # generalized code for data/model/pipe parallelism
+    assert topo.get_axis_comm_lists('jeff') == []
+
+
+class TestDistributedTopology(DistributedTest):
+    world_size = 4
+
+    def test_grid_pipe_data(self):
+        topo = Topo(axes=['pipe', 'data'], dims=[2, 2])
+        grid = Grid(topology=topo)
+
+        assert grid._is_grid_valid()
+
+        rank = dist.get_rank()
+
+        assert grid.is_first_stage == (grid.get_stage_id() == 0)
+        assert grid.is_last_stage == (grid.get_stage_id() == grid.get_pipe_parallel_world_size() - 1)
+
+        # Test collectives along the pipeline parallel process groups
+        rank_tensor = torch.LongTensor(data=[rank]).to(get_accelerator().device_name())
+        dist.all_reduce(rank_tensor, group=grid.get_pipe_parallel_group())
+        pipe_group = grid.pp_group
+        assert torch.all(rank_tensor == sum(pipe_group))
+
+        # Test collectives along the data parallel process groups
+        rank_tensor = torch.LongTensor(data=[rank]).to(get_accelerator().device_name())
+        dist.all_reduce(rank_tensor, group=grid.get_data_parallel_group())
+        data_group = grid.dp_group
+        assert torch.all(rank_tensor == sum(data_group))
+
+    def test_stage_to_global(self):
+        topo = Topo(axes=['pipe', 'data'], dims=[2, 2])
+        grid = Grid(topology=topo)
+
+        assert grid._is_grid_valid()
+
+        assert grid.stage_to_global(stage_id=0, data=0) == 0
+        assert grid.stage_to_global(stage_id=0, data=1) == 1
+        assert grid.stage_to_global(stage_id=1, data=0) == 2
+        assert grid.stage_to_global(stage_id=1, data=1) == 3
+
+        me = topo.get_coord(rank=dist.get_rank())
+        if me.data == 0:
+            assert grid.stage_to_global(stage_id=0) == 0
+            assert grid.stage_to_global(stage_id=1) == 2
+        else:
+            assert grid.stage_to_global(stage_id=0) == 1
+            assert grid.stage_to_global(stage_id=1) == 3
+
+
+def test_primes():
+    """ Test prime factorizations. """
+
+    def _product(ps):
+        p = 1
+        for num in ps:
+            p *= num
+        return p
+
+    with pytest.raises(ValueError):
+        _prime_factors(0)
+
+    for x in range(1, 30):
+        primes = _prime_factors(x)
+        assert _product(primes) == x
+        for p in primes:
+            assert _prime_factors(p) == [p]

tests/unit/runtime/rollout/test_hybrid_engine_rollout.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""CPU-only unit tests for HybridEngineRollout (no GPU needed).
+
+Tests cover configuration defaults and the pure-tensor sampling helper.
+"""
+
+from unittest.mock import MagicMock
+
+import torch
+
+from deepspeed.runtime.rollout.hybrid_engine_rollout import (
+    HybridEngineRollout,
+    HybridEngineRolloutConfig,
+)
+
+
+def _make_engine():
+    engine = MagicMock()
+    engine.module = MagicMock()
+    engine.module.parameters.return_value = iter([])
+    return engine
+
+
+def _make_tokenizer():
+    tok = MagicMock()
+    tok.pad_token_id = 0
+    tok.eos_token_id = 2
+    return tok
+
+
+# -- config defaults ----------------------------------------------------
+
+
+def test_config_defaults():
+    cfg = HybridEngineRolloutConfig()
+    assert cfg.use_graph_capture is False
+
+
+# -- constructor --------------------------------------------------------
+
+
+def test_constructor_stores_config():
+    engine = _make_engine()
+    tok = _make_tokenizer()
+    cfg = HybridEngineRolloutConfig(use_graph_capture=True)
+    rollout = HybridEngineRollout(engine, tok, cfg=cfg)
+    assert rollout.use_graph_capture is True
+    assert rollout.engine is engine
+    assert rollout.tokenizer is tok
+
+
+def test_constructor_defaults_without_cfg():
+    rollout = HybridEngineRollout(_make_engine(), _make_tokenizer())
+    assert rollout.use_graph_capture is False
+
+
+# -- _sample_top_p ------------------------------------------------------
+
+
+def test_sample_top_p_returns_correct_shape():
+    logits = torch.randn(4, 100)
+    tokens = HybridEngineRollout._sample_top_p(logits, temperature=1.0, top_p=1.0)
+    assert tokens.shape == (4, 1)
+
+
+def test_sample_top_p_deterministic_with_low_temp():
+    logits = torch.tensor([[1.0, 10.0, 2.0]])
+    tok = HybridEngineRollout._sample_top_p(logits, temperature=1e-10, top_p=1.0)
+    assert tok.item() == 1
+
+
+def test_sample_top_p_top_p_filters():
+    logits = torch.tensor([[0.0, 0.0, 100.0]])
+    tok = HybridEngineRollout._sample_top_p(logits, temperature=1.0, top_p=0.5)
+    assert tok.item() == 2
+
+
+def test_sample_top_p_batch():
+    logits = torch.randn(8, 50)
+    tokens = HybridEngineRollout._sample_top_p(logits, temperature=0.8, top_p=0.9)
+    assert tokens.shape == (8, 1)
+    assert (tokens >= 0).all() and (tokens < 50).all()
+
+
+# -- sync_weights is no-op ---------------------------------------------
+
+
+def test_sync_weights_is_noop():
+    rollout = HybridEngineRollout(_make_engine(), _make_tokenizer())
+    assert rollout.sync_weights(step=0) is None
+
+
+# -- generate dispatches correctly -------------------------------------
+
+
+def test_generate_calls_graph_capture_when_enabled():
+    engine = _make_engine()
+    tok = _make_tokenizer()
+    cfg = HybridEngineRolloutConfig(use_graph_capture=True)
+    rollout = HybridEngineRollout(engine, tok, cfg=cfg)
+    rollout._generate_graph = MagicMock(return_value=torch.zeros(1, 5, dtype=torch.long))
+
+    req = MagicMock()
+    req.prompt_ids = torch.tensor([[1, 2]])
+    req.prompt_attention_mask = torch.ones(1, 2, dtype=torch.long)
+    sampling = MagicMock()
+    sampling.temperature = 0
+    sampling.n_samples_per_prompt = 1
+    sampling.max_new_tokens = 3
+
+    rollout.generate(req, sampling)
+    rollout._generate_graph.assert_called_once()

tests/unit/runtime/rollout/test_rollout_interface.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Conformance tests for the RolloutEngine interface.
+
+Validates the dataclass invariants and exercises the interface against a
+``FakeRollout`` so the contract is testable without GPUs or a model. The real
+backends are tested manually with a launched training script (see README).
+"""
+
+import pytest
+import torch
+
+from deepspeed.runtime.rollout import (
+    RolloutBatch,
+    RolloutEngine,
+    RolloutRequest,
+    SamplingConfig,
+    build_rollout,
+)
+
+# --- dataclass invariants ---------------------------------------------------
+
+
+def test_rollout_request_validates_shapes():
+    with pytest.raises(ValueError, match="must be 2-D"):
+        RolloutRequest(prompt_ids=torch.zeros(8), prompt_attention_mask=torch.ones(8))
+    with pytest.raises(ValueError, match="does not match"):
+        RolloutRequest(prompt_ids=torch.zeros(2, 4, dtype=torch.long), prompt_attention_mask=torch.ones(2, 5))
+
+
+def test_rollout_batch_validates_shapes():
+    with pytest.raises(ValueError, match="must be 2-D"):
+        RolloutBatch(input_ids=torch.zeros(8, dtype=torch.long),
+                     attention_mask=torch.ones(8),
+                     response_start_idx=torch.tensor([4]))
+    with pytest.raises(ValueError, match="does not match"):
+        RolloutBatch(input_ids=torch.zeros(2, 4, dtype=torch.long),
+                     attention_mask=torch.ones(2, 5),
+                     response_start_idx=torch.tensor([4, 4]))
+    with pytest.raises(ValueError, match="1-D of length"):
+        RolloutBatch(input_ids=torch.zeros(2, 4, dtype=torch.long),
+                     attention_mask=torch.ones(2, 4),
+                     response_start_idx=torch.tensor([4]))
+
+
+def test_rollout_batch_accessors():
+    batch = RolloutBatch(
+        input_ids=torch.zeros(3, 12, dtype=torch.long),
+        attention_mask=torch.ones(3, 12),
+        response_start_idx=torch.tensor([4, 5, 6]),
+    )
+    assert batch.batch_size == 3
+    assert batch.seq_len == 12
+
+
+def test_sampling_config_defaults():
+    cfg = SamplingConfig(max_new_tokens=32)
+    assert cfg.temperature == 1.0
+    assert cfg.top_p == 1.0
+    assert cfg.top_k == -1
+    assert cfg.n_samples_per_prompt == 1
+
+
+# --- interface conformance via FakeRollout ---------------------------------
+
+
+class FakeRollout(RolloutEngine):
+    """Deterministic stub: appends ``[42] * max_new_tokens`` to each prompt."""
+
+    name = "fake"
+
+    def __init__(self, response_token: int = 42):
+        self.response_token = response_token
+        self.sync_calls: list = []
+
+    def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> RolloutBatch:
+        B, T_p = request.prompt_ids.shape
+        n = sampling.n_samples_per_prompt
+        T_r = sampling.max_new_tokens
+
+        prompts_expanded = request.prompt_ids.repeat_interleave(n, dim=0)
+        attn_p_expanded = request.prompt_attention_mask.repeat_interleave(n, dim=0)
+        response = torch.full((B * n, T_r), self.response_token, dtype=request.prompt_ids.dtype)
+        response_attn = torch.ones((B * n, T_r), dtype=attn_p_expanded.dtype)
+
+        input_ids = torch.cat([prompts_expanded, response], dim=1)
+        attention_mask = torch.cat([attn_p_expanded, response_attn], dim=1)
+        response_start_idx = torch.full((B * n, ), T_p, dtype=torch.long)
+        return RolloutBatch(input_ids=input_ids, attention_mask=attention_mask, response_start_idx=response_start_idx)
+
+    def sync_weights(self, step: int) -> None:
+        self.sync_calls.append(step)
+
+
+def test_fake_rollout_shape_basic():
+    fake = FakeRollout()
+    req = RolloutRequest(prompt_ids=torch.tensor([[1, 2, 3], [4, 5, 6]]),
+                         prompt_attention_mask=torch.ones(2, 3, dtype=torch.long))
+    out = fake.generate(req, SamplingConfig(max_new_tokens=4))
+    assert out.input_ids.shape == (2, 7)
+    assert out.attention_mask.shape == (2, 7)
+    # With left-padded (fully real here) prompts of width 3, response begins
+    # at column 3 for every sample.
+    assert out.response_start_idx.tolist() == [3, 3]
+
+
+def test_fake_rollout_with_n_samples():
+    fake = FakeRollout()
+    req = RolloutRequest(prompt_ids=torch.tensor([[1, 2], [3, 4]]),
+                         prompt_attention_mask=torch.ones(2, 2, dtype=torch.long))
+    out = fake.generate(req, SamplingConfig(max_new_tokens=3, n_samples_per_prompt=4))
+    assert out.input_ids.shape == (8, 5)
+    assert out.response_start_idx.tolist() == [2] * 8
+
+
+def test_fake_rollout_left_padded_prompts():
+    fake = FakeRollout()
+    # left-padded prompts: prompt B has only the last 2 positions real, but
+    # response_start_idx still equals the prompt column width T_p.
+    prompt_ids = torch.tensor([[1, 2, 3, 4], [0, 0, 5, 6]])
+    attn = torch.tensor([[1, 1, 1, 1], [0, 0, 1, 1]], dtype=torch.long)
+    req = RolloutRequest(prompt_ids=prompt_ids, prompt_attention_mask=attn)
+    out = fake.generate(req, SamplingConfig(max_new_tokens=2))
+    assert out.response_start_idx.tolist() == [4, 4]
+
+
+def test_sync_records_steps():
+    fake = FakeRollout()
+    fake.sync_weights(0)
+    fake.sync_weights(5)
+    assert fake.sync_calls == [0, 5]
+
+
+def test_engine_factory_unknown_raises():
+    from deepspeed.runtime.rollout.base import RolloutConfig
+
+    with pytest.raises(ValueError, match="Unknown rollout engine"):
+        build_rollout(RolloutConfig(engine="totally_made_up"))
+
+
+def test_engine_factory_hybrid_requires_student_engine():
+    from deepspeed.runtime.rollout.base import RolloutConfig
+
+    with pytest.raises(ValueError, match="needs both"):
+        build_rollout(RolloutConfig(engine="hybrid_engine"))

tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import pytest
+import deepspeed
+from unit.common import DistributedTest
+from unit.util import skip_on_arch
+from deepspeed.accelerator import get_accelerator
+
+if get_accelerator().device_name() == 'hpu':
+    pytest.skip("sparse_gradients not supported by HPU.", allow_module_level=True)
+
+
+class Model(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.emb = torch.nn.EmbeddingBag(10, 3, mode="sum", sparse=True)
+        self.linear = torch.nn.Linear(3, 1)
+
+    def forward(self, x, offsets):
+        return self.linear(self.emb(x, offsets))
+
+
+class Adam(torch.optim.Optimizer):
+
+    def __init__(self, dense_params, sparse_params):
+        super().__init__(dense_params + sparse_params, defaults={})
+        self.adam = torch.optim.Adam(dense_params)
+        self.adam_sparse = torch.optim.SparseAdam(sparse_params)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss_1 = self.adam.step(closure)
+        loss_2 = self.adam_sparse.step(closure)
+
+        if loss_1 is not None and loss_2 is not None:
+            return loss_1 + loss_2
+        return loss_1 or loss_2
+
+
+def get_model_optimizer():
+    torch.manual_seed(0)
+    model = Model()
+    optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters()))
+    return model, optimizer
+
+
+def get_data(device):
+    x = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long, device=device)
+    offsets = torch.tensor([0, 4], dtype=torch.long, device=device)
+    y = torch.tensor([[1.0], [0.0]], device=device)
+    return x, offsets, y
+
+
+class TestSparseAdam(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        skip_on_arch(min_arch=7)
+
+        config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True}
+        model, optimizer = get_model_optimizer()
+        loss = torch.nn.BCEWithLogitsLoss()
+        engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config_dict)
+
+        x, offsets, y = get_data(engine.device)
+
+        engine.gradient_average = True
+        res = engine(x, offsets)
+        engine.backward(loss(res, y))
+
+        averaged_grads = {}
+        for k, v in engine.named_parameters():
+            grad = v.grad.to_dense() if v.grad.is_sparse else v.grad
+            averaged_grads[k] = grad
+            v.grad = None
+
+        engine.gradient_average = False
+        res = engine(x, offsets)
+        engine.backward(loss(res, y))
+
+        for k, v in engine.named_parameters():
+            grad = v.grad.to_dense() if v.grad.is_sparse else v.grad
+            assert torch.allclose(grad, averaged_grads[k] * engine.world_size)

tests/unit/runtime/sparse_tensor/test_csr.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import random
+from deepspeed.runtime.sparse_tensor import SparseTensor
+
+
+def test_csr_addition_self():
+    row_count = 10
+    random.seed(1234)
+
+    x = torch.ones(1, 5)
+    for i in range(row_count - 1):
+        if random.random() > 0.75:
+            x = torch.cat([x, torch.ones(1, 5)])
+        else:
+            x = torch.cat([x, torch.zeros(1, 5)])
+    dense_x = x.clone()
+    cx = SparseTensor(x)
+
+    assert torch.all(dense_x == cx.to_dense())
+
+    cx.add(cx)
+    assert torch.all(dense_x + dense_x == cx.to_dense())
+
+
+def test_csr_addition_different():
+    row_count = 10
+    random.seed(1234)
+
+    x = torch.ones(1, 5)
+    for i in range(row_count - 1):
+        if random.random() > 0.75:
+            x = torch.cat([x, torch.ones(1, 5)])
+        else:
+            x = torch.cat([x, torch.zeros(1, 5)])
+    dense_x = x.clone()
+    cx = SparseTensor(x)
+
+    y = torch.ones(1, 5)
+    for i in range(row_count - 1):
+        if random.random() > 0.75:
+            y = torch.cat([y, torch.ones(1, 5)])
+        else:
+            y = torch.cat([y, torch.zeros(1, 5)])
+    dense_y = y.clone()
+    cy = SparseTensor(y)
+
+    dense_sum = dense_x + dense_y
+    cx.add(cy)
+
+    assert torch.all(dense_sum == cx.to_dense())

tests/unit/runtime/sparse_tensor/test_sparse_grads.py

@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import pytest
+import deepspeed
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+import deepspeed.utils.groups as groups
+
+if get_accelerator().device_name() == 'hpu':
+    pytest.skip("sparse_gradients not supported by HPU.", allow_module_level=True)
+
+
+class Model(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.emb = torch.nn.EmbeddingBag(10, 3, mode="sum", sparse=True)
+        self.linear = torch.nn.Linear(3, 1)
+
+    def forward(self, x, offsets):
+        return self.linear(self.emb(x, offsets))
+
+
+class Adam(torch.optim.Optimizer):
+
+    def __init__(self, dense_params, sparse_params):
+        super().__init__(dense_params + sparse_params, defaults={})
+        self.adam = torch.optim.Adam(dense_params)
+        self.adam_sparse = torch.optim.SparseAdam(sparse_params)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss_1 = self.adam.step(closure)
+        loss_2 = self.adam_sparse.step(closure)
+
+        if loss_1 is not None and loss_2 is not None:
+            return loss_1 + loss_2
+        return loss_1 or loss_2
+
+
+class TestSparseAdam(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True}
+
+        model = Model()
+        optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters()))
+        engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config_dict)
+        loss = torch.nn.BCEWithLogitsLoss()
+        x = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long, device=engine.device)
+        offsets = torch.tensor([0, 4], dtype=torch.long, device=engine.device)
+        y = torch.tensor([[1.0], [0.0]], device=engine.device)
+        res = engine(x, offsets)
+        engine.backward(loss(res, y))
+        engine.step()
+
+        results = [engine.all_gather_scalar(i, groups._get_data_parallel_group()) for i in model.emb.parameters()]
+        for res in results:
+            assert torch.allclose(res[0], res[1])

tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py

@@ -0,0 +1,160 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed.checkpoint.constants import (PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_SUB_PARAMS,
+                                            TP_REPLICATED_PARAMETER_PATTERNS, DS_AUTOTP_UC_META)
+from deepspeed.module_inject.layers import (_build_param_uc_restore_meta, _get_param_uc_conversion_meta,
+                                            LinearAllreduce, LinearLayer, SubParamLinearLayer,
+                                            collect_autotp_universal_checkpoint_info)
+
+
+def test_collect_autotp_universal_checkpoint_info_row_parallel():
+    layer = LinearAllreduce(torch.nn.Linear(16, 8, bias=True), mp_group=None, name="proj")
+    model = torch.nn.Module()
+    model.proj = layer
+
+    uc_info = collect_autotp_universal_checkpoint_info(model)
+
+    # collect_autotp_universal_checkpoint_info() stores regex patterns like r"^proj\.weight$"
+    assert r"^proj\.weight$" in uc_info[PARAMETER_WITH_ROW_PARALLELISM_PATTERNS]
+    # bias in LinearAllreduce is marked replicated, so it should appear in replicated patterns
+    assert r"^proj\.bias$" in uc_info[TP_REPLICATED_PARAMETER_PATTERNS]
+
+
+def test_collect_autotp_universal_checkpoint_info_subparams():
+    layer = SubParamLinearLayer(torch.nn.Linear(12, 12, bias=True),
+                                mp_group=None,
+                                shape=(3, -1),
+                                partition_dim=0,
+                                name="qkv")
+    model = torch.nn.Module()
+    model.qkv = layer
+
+    uc_info = collect_autotp_universal_checkpoint_info(model)
+
+    assert len(uc_info[PARAMETER_WITH_SUB_PARAMS]) == 1
+    assert uc_info[PARAMETER_WITH_SUB_PARAMS][0]["partition_dim"] == 0
+
+
+def test_collect_autotp_universal_checkpoint_info_column_parallel_bias_not_replicated():
+    layer = LinearLayer(torch.nn.Linear(16, 8, bias=True), mp_group=None, name="dense")
+    model = torch.nn.Module()
+    model.dense = layer
+
+    uc_info = collect_autotp_universal_checkpoint_info(model)
+
+    assert not any("dense.weight" in p for p in uc_info[PARAMETER_WITH_ROW_PARALLELISM_PATTERNS])
+    assert not any("dense.bias" in p for p in uc_info[TP_REPLICATED_PARAMETER_PATTERNS])
+
+
+def test_collect_autotp_universal_checkpoint_info_subparams_preserves_shape_metadata():
+    layer = SubParamLinearLayer(torch.nn.Linear(12, 12, bias=True),
+                                mp_group=None,
+                                shape=((2, 10), 12),
+                                partition_dim=0,
+                                name="fused")
+    model = torch.nn.Module()
+    model.fused = layer
+
+    uc_info = collect_autotp_universal_checkpoint_info(model)
+
+    assert uc_info[PARAMETER_WITH_SUB_PARAMS][0]["shape"] == [(2, 10), 12]
+
+
+def test_subparam_layer_marks_standardized_param_metadata():
+    layer = SubParamLinearLayer(torch.nn.Linear(12, 12, bias=True),
+                                mp_group=None,
+                                shape=(3, -1),
+                                partition_dim=0,
+                                name="packed")
+
+    weight_meta = getattr(layer.weight, DS_AUTOTP_UC_META)
+    bias_meta = getattr(layer.bias, DS_AUTOTP_UC_META)
+
+    assert weight_meta["sub_param_sizes"] == (4, 4, 4)
+    assert tuple(weight_meta["target_partition_shape"]) == tuple(layer.weight.shape)
+    assert tuple(bias_meta["target_partition_shape"]) == tuple(layer.bias.shape)
+
+
+def test_universal_checkpoint_info_excludes_param_level_recovery_fields():
+    layer = SubParamLinearLayer(torch.nn.Linear(12, 12, bias=True),
+                                mp_group=None,
+                                shape=(3, -1),
+                                partition_dim=0,
+                                name="packed")
+    model = torch.nn.Module()
+    model.packed = layer
+
+    uc_info = collect_autotp_universal_checkpoint_info(model)
+    subparam_entry = uc_info[PARAMETER_WITH_SUB_PARAMS][0]
+
+    assert "shape" in subparam_entry
+    assert "partition_dim" in subparam_entry
+    assert "patterns" in subparam_entry
+    assert "sub_param_sizes" not in subparam_entry
+    assert "target_partition_shape" not in subparam_entry
+
+
+def test_collect_uses_conversion_view_not_recovery_fields():
+    layer = SubParamLinearLayer(torch.nn.Linear(12, 12, bias=True),
+                                mp_group=None,
+                                shape=(3, -1),
+                                partition_dim=0,
+                                name="packed")
+    model = torch.nn.Module()
+    model.packed = layer
+
+    meta = getattr(layer.weight, "ds_autotp_universal_checkpoint_meta")
+    meta["partition_dim"] = 99
+    meta["sub_param_shape"] = (999, -1)
+
+    uc_info = collect_autotp_universal_checkpoint_info(model)
+    subparam_entry = uc_info[PARAMETER_WITH_SUB_PARAMS][0]
+
+    assert subparam_entry["partition_dim"] == 0
+    assert subparam_entry["shape"] == [3, -1]
+
+
+def test_param_uc_restore_builder_normalizes_shapes_and_nests_conversion_view():
+    restore_meta = _build_param_uc_restore_meta(partition_type="column",
+                                                partition_dim=0,
+                                                logical_shape=[12, 8],
+                                                output_shape=[12],
+                                                sub_param_shape=[3, -1],
+                                                sub_param_sizes=[4, 4, 4],
+                                                target_partition_shape=torch.Size([4, 8]),
+                                                original_shape=torch.Size([12, 8]),
+                                                is_bias=False,
+                                                replicated=False)
+
+    assert restore_meta["logical_shape"] == (12, 8)
+    assert restore_meta["output_shape"] == (12, )
+    assert restore_meta["sub_param_shape"] == (3, -1)
+    assert restore_meta["sub_param_sizes"] == (4, 4, 4)
+    assert restore_meta["target_partition_shape"] == (4, 8)
+    assert restore_meta["original_shape"] == (12, 8)
+    assert restore_meta["conversion"] == {
+        "partition_type": "column",
+        "partition_dim": 0,
+        "sub_param_shape": (3, -1),
+        "original_shape": (12, 8),
+        "is_bias": False,
+        "replicated": False,
+    }
+
+
+def test_conversion_helper_reads_builder_nested_view():
+    param = torch.nn.Parameter(torch.zeros(4, 8))
+    param.ds_autotp_universal_checkpoint_meta = _build_param_uc_restore_meta(partition_type="row",
+                                                                             partition_dim=1,
+                                                                             logical_shape=[4, 16],
+                                                                             output_shape=[4],
+                                                                             original_shape=[4, 16],
+                                                                             is_bias=False,
+                                                                             replicated=False)
+
+    assert _get_param_uc_conversion_meta(param) == param.ds_autotp_universal_checkpoint_meta["conversion"]

tests/unit/runtime/test_autocast.py

@@ -0,0 +1,89 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import functools
+
+import pytest
+import torch
+import deepspeed.runtime.zero.linear as zero_linear
+from deepspeed.runtime.zero.linear import LinearModuleForZeroStage3
+from deepspeed.accelerator import get_accelerator
+from deepspeed.utils.torch import required_torch_version
+from unit.common import DistributedTest
+
+
+@pytest.mark.parametrize('half_op', [False, True])
+class TestAutoCastDisable(DistributedTest):
+
+    def test_missing_amp_autocast(self, half_op):
+        hidden_dim = 4
+        if half_op:
+            input = torch.randn(hidden_dim).to(get_accelerator().device_name()).half()
+            ds_linear = LinearModuleForZeroStage3(hidden_dim, hidden_dim).to(get_accelerator().device_name()).half()
+        else:
+            input = torch.randn(hidden_dim).to(get_accelerator().device_name())
+            ds_linear = LinearModuleForZeroStage3(hidden_dim, hidden_dim).to(get_accelerator().device_name())
+
+        output = ds_linear(input)
+        assert output.dtype == ds_linear.weight.dtype
+
+    def test_disable_autocast_linear(self, half_op):
+        hidden_dim = 4
+        if half_op:
+            input = torch.randn(hidden_dim).to(get_accelerator().device_name()).half()
+            ds_linear = LinearModuleForZeroStage3(hidden_dim, hidden_dim).to(get_accelerator().device_name()).half()
+        else:
+            input = torch.randn(hidden_dim).to(get_accelerator().device_name())
+            ds_linear = LinearModuleForZeroStage3(hidden_dim, hidden_dim).to(get_accelerator().device_name())
+
+        with torch.amp.autocast(device_type=get_accelerator().device_name(), enabled=False):
+            output = ds_linear(input)
+            assert output.dtype == ds_linear.weight.dtype
+
+
+@pytest.mark.parametrize('half_input, half_weight', [(False, False), (False, True), (True, False), (True, True)])
+class TestAutoCastEnable(DistributedTest):
+
+    def test_autocast_linear(self, tmpdir, half_input, half_weight):
+        hidden_dim = 4
+        input = torch.randn(hidden_dim).to(get_accelerator().device_name())
+        ds_linear = LinearModuleForZeroStage3(hidden_dim, hidden_dim).to(get_accelerator().device_name())
+
+        if half_input:
+            input = input.half()
+
+        if half_weight:
+            ds_linear = ds_linear.half()
+
+        with torch.amp.autocast(device_type=get_accelerator().device_name()):
+            output = ds_linear(input)
+            assert output.dtype == torch.half or output.dtype == torch.bfloat16
+
+
+def test_get_autocast_decorators_use_torch_amp_on_torch_2_4_or_newer():
+    if not required_torch_version(min_version=2.4):
+        pytest.skip('torch.amp.custom_fwd/custom_bwd are only available on torch >= 2.4')
+
+    device_type = get_accelerator().device_name()
+
+    assert isinstance(zero_linear.autocast_custom_fwd, functools.partial)
+    assert isinstance(zero_linear.autocast_custom_bwd, functools.partial)
+    assert zero_linear.autocast_custom_fwd.func is torch.amp.custom_fwd
+    assert zero_linear.autocast_custom_bwd.func is torch.amp.custom_bwd
+    assert zero_linear.autocast_custom_fwd.keywords == {'device_type': device_type}
+    assert zero_linear.autocast_custom_bwd.keywords == {'device_type': device_type}
+
+
+def test_get_autocast_decorators_use_legacy_amp_or_noop_before_torch_2_4():
+    if required_torch_version(min_version=2.4):
+        pytest.skip('legacy AMP fallback only applies on torch < 2.4')
+
+    device_type = get_accelerator().device_name()
+    legacy_amp = getattr(getattr(torch, device_type, None), 'amp', None)
+    expected_custom_fwd = getattr(legacy_amp, 'custom_fwd', zero_linear.noop_decorator)
+    expected_custom_bwd = getattr(legacy_amp, 'custom_bwd', zero_linear.noop_decorator)
+
+    assert zero_linear.autocast_custom_fwd is expected_custom_fwd
+    assert zero_linear.autocast_custom_bwd is expected_custom_bwd

tests/unit/runtime/test_data.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.utils import RepeatingLoader
+import torch
+import pytest
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataset
+
+
+def test_repeating_loader():
+    loader = [1, 2, 3]
+    loader = RepeatingLoader(loader)
+
+    for idx in range(50):
+        assert next(loader) == 1
+        assert next(loader) == 2
+        assert next(loader) == 3
+
+
+@pytest.mark.parametrize('train_batch_size, drop_last', [(1, True), (4, True), (1, False), (4, False)])
+class TestDataLoaderDropLast(DistributedTest):
+    world_size = 1
+
+    def test(self, train_batch_size, drop_last):
+        config_dict = {"train_batch_size": train_batch_size, "dataloader_drop_last": drop_last, "steps_per_print": 1}
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        optimizer = torch.optim.AdamW(params=model.parameters())
+        # TODO: no way to set DeepSpeedEngine.deepspeed_io params, need to use
+        # pin_memory=False for cuda device
+        train_dataset = random_dataset(total_samples=50,
+                                       hidden_dim=hidden_dim,
+                                       device=torch.device('cpu'),
+                                       dtype=torch.float32)
+        model, _, training_dataloader, _ = deepspeed.initialize(config=config_dict,
+                                                                model=model,
+                                                                training_data=train_dataset,
+                                                                optimizer=optimizer)
+        training_dataloader.num_local_io_workers = 0  # We can't do nested mp.pool
+        for n, batch in enumerate(training_dataloader):
+            x = batch[0].to(get_accelerator().current_device_name())
+            y = batch[1].to(get_accelerator().current_device_name())
+            loss = model(x, y)
+            model.backward(loss)
+            model.step()

tests/unit/runtime/test_data_efficiency.py

@@ -0,0 +1,239 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import os
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+import pytest
+from unit.common import DistributedTest
+from unit.simple_model import Curriculum_SimpleModel, SimpleModel, random_dataloader, random_dataset
+
+
+class MPU():
+
+    def __init__(self, tp_world_size):
+        self.rank = deepspeed.comm.get_rank()
+        self.world_size = deepspeed.comm.get_world_size()
+        self.tp_world_size = tp_world_size
+
+        for i in range(0, self.world_size, tp_world_size):
+            ranks = range(i, i + tp_world_size)
+            group = deepspeed.comm.new_group(ranks)
+            if self.rank in ranks:
+                self.tp_group = group
+
+        for i in range(0, tp_world_size):
+            ranks = range(i, self.world_size, tp_world_size)
+            group = deepspeed.comm.new_group(ranks)
+            if self.rank in ranks:
+                self.dp_group = group
+
+    def get_model_parallel_rank(self):
+        return self.rank % self.tp_world_size
+
+    def get_model_parallel_world_size(self):
+        return self.tp_world_size
+
+    def get_data_parallel_rank(self):
+        return self.rank // self.tp_world_size
+
+    def get_data_parallel_world_size(self):
+        return self.world_size // self.tp_world_size
+
+    def get_data_parallel_group(self):
+        return self.dp_group
+
+    def get_model_parallel_group(self):
+        return self.tp_group
+
+
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16])
+class TestDataEfficiency(DistributedTest):
+    world_size = 2
+
+    def test_curriculum_learning(self, dtype):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU accelerator does not support this test yet")
+        if not dtype in get_accelerator().supported_dtypes():
+            pytest.skip(f"This test does not support {dtype=}.")
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015,
+                    "weight_decay": 0.01
+                }
+            },
+            "gradient_clipping": 1.0,
+            "data_efficiency": {
+                "enabled": True,
+                "seed": 1234,
+                "data_sampling": {
+                    "enabled": True,
+                    "num_workers": 0,
+                    "curriculum_learning": {
+                        "enabled": True,
+                        "data_cluster_path": "/tmp",
+                        "curriculum_metrics": {
+                            "dummy_metric": {
+                                "index_to_sample_path": "dummy",
+                                "index_to_metric_path": "dummy",
+                                "difficulty_type": "value",
+                                "clustering_type": "single_cluster",
+                                "min_difficulty": 2,
+                                "max_difficulty": 10,
+                                "schedule_type": "fixed_root",
+                                "schedule_config": {
+                                    "total_curriculum_step": 8,
+                                    "difficulty_step": 2,
+                                    "root_degree": 1
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        if dtype == torch.float16:
+            config_dict["fp16"] = {"enabled": True, "loss_scale": 0, "initial_scale_power": 8}
+        else:
+            config_dict["bf16"] = {"enabled": True}
+
+        def data_post_process(data, data_sampler_state_dict):
+            assert 'dummy_metric' in data_sampler_state_dict['current_difficulties']
+            return data
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        dataset = random_dataset(20, hidden_dim, torch.device('cpu'), dtype=dtype)
+        model, _, data_loader, _ = deepspeed.initialize(config=config_dict,
+                                                        model=model,
+                                                        training_data=dataset,
+                                                        model_parameters=model.parameters(),
+                                                        mpu=MPU(1))
+        if model.mpu.get_data_parallel_rank() == 0 and not os.path.exists('/tmp'):
+            os.makedirs('/tmp')
+        model.set_data_post_process_func(data_post_process)
+        for n, batch in enumerate(data_loader):
+            x = batch[0].to(get_accelerator().current_device_name())
+            y = batch[1].to(get_accelerator().current_device_name())
+            loss = model(x, y)
+            model.backward(loss)
+            model.step()
+            if n >= 10:
+                break
+
+
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16])
+class TestLegacyCurriculumScheduler(DistributedTest):
+    world_size = 2
+
+    def test_fixed_discrete(self, dtype):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU accelerator does not support this test yet")
+        if not dtype in get_accelerator().supported_dtypes():
+            pytest.skip(f"This test does not support {dtype=}.")
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015,
+                    "weight_decay": 0.01
+                }
+            },
+            "gradient_clipping": 1.0,
+            "curriculum_learning": {
+                "enabled": True,
+                "curriculum_type": "seqlen",
+                "min_difficulty": 1,
+                "max_difficulty": 5,
+                "schedule_type": "fixed_discrete",
+                "schedule_config": {
+                    "difficulty": [1, 2, 3, 4, 5],
+                    "max_step": [2, 4, 6, 8]
+                }
+            }
+        }
+        if dtype == torch.float16:
+            config_dict["fp16"] = {"enabled": True, "loss_scale": 0, "initial_scale_power": 8}
+        else:
+            config_dict["bf16"] = {"enabled": True}
+        hidden_dim = 10
+        ground_truths = {1: 1, 2: 1, 3: 2, 4: 2, 5: 3, 6: 3, 7: 4, 8: 4}
+
+        model = Curriculum_SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=20,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=dtype)
+        for n, batch in enumerate(data_loader):
+            loss, seqlen = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+            true_seqlen = 5
+            if n + 1 in ground_truths:
+                true_seqlen = ground_truths[n + 1]
+            assert seqlen == true_seqlen, f"Incorrect curriculum schedule {n=}, {seqlen=}, {true_seqlen=}"
+
+    def test_fixed_linear(self, dtype):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU accelerator does not support this test yet")
+        if not dtype in get_accelerator().supported_dtypes():
+            pytest.skip(f"This test does not support {dtype=}.")
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015,
+                    "weight_decay": 0.01
+                }
+            },
+            "gradient_clipping": 1.0,
+            "curriculum_learning": {
+                "enabled": True,
+                "curriculum_type": "seqlen",
+                "min_difficulty": 2,
+                "max_difficulty": 10,
+                "schedule_type": "fixed_linear",
+                "schedule_config": {
+                    "total_curriculum_step": 8,
+                    "difficulty_step": 2
+                }
+            }
+        }
+        if dtype == torch.float16:
+            config_dict["fp16"] = {"enabled": True, "loss_scale": 0, "initial_scale_power": 8}
+        else:
+            config_dict["bf16"] = {"enabled": True}
+        hidden_dim = 10
+        ground_truths = {1: 2, 2: 4, 3: 4, 4: 6, 5: 6, 6: 8, 7: 8, 8: 10, 9: 10, 10: 10}
+
+        model = Curriculum_SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=20,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=dtype)
+        for n, batch in enumerate(data_loader):
+            loss, seqlen = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+            if n + 1 in ground_truths:
+                true_seqlen = ground_truths[n + 1]
+                assert seqlen == true_seqlen, "Incorrect curriculum schedule"

tests/unit/runtime/test_ds_config_dict.py

@@ -0,0 +1,348 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# A test on its own
+import os
+import pytest
+import json
+import hjson
+import argparse
+import torch
+
+from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
+from deepspeed.accelerator import get_accelerator
+
+from unit.common import DistributedTest, get_test_path
+from unit.simple_model import SimpleModel, create_config_from_dict, random_dataloader
+import deepspeed.comm as dist
+
+# A test on its own
+import deepspeed
+from deepspeed.runtime.config import DeepSpeedConfig
+from deepspeed.runtime.precision_config import get_bfloat16_config
+
+
+class TestBasicConfig(DistributedTest):
+    world_size = 1
+
+    def test_accelerator(self):
+        assert (get_accelerator().is_available())
+
+    def test_check_version(self):
+        assert hasattr(deepspeed, "__git_hash__")
+        assert hasattr(deepspeed, "__git_branch__")
+        assert hasattr(deepspeed, "__version__")
+        assert hasattr(deepspeed, "__version_major__")
+        assert hasattr(deepspeed, "__version_minor__")
+        assert hasattr(deepspeed, "__version_patch__")
+
+
+@pytest.fixture
+def base_config():
+    config_dict = {
+        "train_batch_size": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 0.00015
+            }
+        },
+    }
+    return config_dict
+
+
+def _run_batch_config(ds_config, train_batch=None, micro_batch=None, gas=None):
+    ds_config.train_batch_size = train_batch
+    ds_config.train_micro_batch_size_per_gpu = micro_batch
+    ds_config.gradient_accumulation_steps = gas
+    success = True
+    try:
+        ds_config._configure_train_batch_size()
+    except AssertionError:
+        success = False
+    return success
+
+
+def _batch_assert(status, ds_config, batch, micro_batch, gas, success):
+
+    if not success:
+        assert not status
+        return
+
+    assert ds_config.train_batch_size == batch
+    assert ds_config.train_micro_batch_size_per_gpu == micro_batch
+    assert ds_config.gradient_accumulation_steps == gas
+
+
+#Tests different batch config provided in deepspeed json file
+@pytest.mark.parametrize('num_ranks,batch,micro_batch,gas,success',
+                         [(2,32,16,1,True),
+                         (2,32,8,2,True),
+                         (2,33,17,2,False),
+                         (2,32,18,1,False)]) # yapf: disable
+class TestBatchConfig(DistributedTest):
+    world_size = 2
+
+    def test(self, num_ranks, batch, micro_batch, gas, success):
+        assert dist.get_world_size() == num_ranks, \
+        f'The test assumes a world size of {num_ranks}'
+
+        ds_batch_config = get_test_path('ds_batch_config.json')
+        ds_config = DeepSpeedConfig(ds_batch_config)
+
+        #test cases when all parameters are provided
+        status = _run_batch_config(ds_config, train_batch=batch, micro_batch=micro_batch, gas=gas)
+        _batch_assert(status, ds_config, batch, micro_batch, gas, success)
+
+        #test cases when two out of three parameters are provided
+        status = _run_batch_config(ds_config, train_batch=batch, micro_batch=micro_batch)
+        _batch_assert(status, ds_config, batch, micro_batch, gas, success)
+
+        if success:
+            #when gas is provided with one more parameter
+            status = _run_batch_config(ds_config, train_batch=batch, gas=gas)
+            _batch_assert(status, ds_config, batch, micro_batch, gas, success)
+
+            status = _run_batch_config(ds_config, micro_batch=micro_batch, gas=gas)
+            _batch_assert(status, ds_config, batch, micro_batch, gas, success)
+
+            #test the case when only micro_batch or train_batch is provided
+            if gas == 1:
+                status = _run_batch_config(ds_config, micro_batch=micro_batch)
+                _batch_assert(status, ds_config, batch, micro_batch, gas, success)
+
+                status = _run_batch_config(ds_config, train_batch=batch)
+                _batch_assert(status, ds_config, batch, micro_batch, gas, success)
+        else:
+            #when only gas is provided
+            status = _run_batch_config(ds_config, gas=gas)
+            _batch_assert(status, ds_config, batch, micro_batch, gas, success)
+
+            #when gas is provided with something else and gas does not divide batch
+            if gas != 1:
+                status = _run_batch_config(ds_config, train_batch=batch, gas=gas)
+                _batch_assert(status, ds_config, batch, micro_batch, gas, success)
+
+
+def test_temp_config_json(tmpdir):
+    config_dict = {
+        "train_batch_size": 1,
+    }
+    config_path = create_config_from_dict(tmpdir, config_dict)
+    config_json = json.load(open(config_path, 'r'))
+    assert 'train_batch_size' in config_json
+
+
+@pytest.mark.parametrize("gather_weights_key",
+                         ["stage3_gather_16bit_weights_on_model_save", "stage3_gather_fp16_weights_on_model_save"])
+def test_gather_16bit_params_on_model_save(gather_weights_key):
+    config_dict = {
+        gather_weights_key: True,
+    }
+    config = DeepSpeedZeroConfig(**config_dict)
+
+    assert config.gather_16bit_weights_on_model_save == True
+
+
+@pytest.mark.parametrize("bf16_key", ["bf16", "bfloat16"])
+def test_get_bfloat16_enabled(bf16_key):
+    cfg = {
+        bf16_key: {
+            "enabled": True,
+        },
+    }
+    assert get_bfloat16_config(cfg).enabled == True
+
+
+def test_quantized_eigenvalue_config_parses():
+    ds_config_path = get_test_path('../model/BingBertSquad/deepspeed_bsz24_fp16_eigenvalue_quantize_config.json')
+
+    ds_config = DeepSpeedConfig(ds_config_path)
+
+    assert ds_config._param_dict["quantize_training"]["quantize_eigenvalue"] is True
+
+
+def test_compression_training_without_legacy_quantize_training_uses_defaults():
+    config_dict = {
+        "train_micro_batch_size_per_gpu": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 1e-4,
+            },
+        },
+        "compression_training": {
+            "weight_quantization": {
+                "shared_parameters": {
+                    "enabled": True,
+                },
+                "different_groups": {},
+            }
+        },
+    }
+
+    ds_config = DeepSpeedConfig(config_dict)
+
+    assert ds_config.eigenvalue_enabled is False
+    assert ds_config.eigenvalue_verbose is False
+
+
+class TestConfigLoad(DistributedTest):
+    world_size = 1
+
+    def test_dict(self, base_config):
+        if get_accelerator().is_bf16_supported():
+            base_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            base_config["fp16"] = {"enabled": True}
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=base_config, model=model, model_parameters=model.parameters())
+
+    def test_json(self, base_config, tmpdir):
+        if get_accelerator().is_bf16_supported():
+            base_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            base_config["fp16"] = {"enabled": True}
+        config_path = os.path.join(tmpdir, "config.json")
+        with open(config_path, 'w') as fp:
+            json.dump(base_config, fp)
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_path, model=model, model_parameters=model.parameters())
+
+    def test_hjson(self, base_config, tmpdir):
+        if get_accelerator().is_bf16_supported():
+            base_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            base_config["fp16"] = {"enabled": True}
+        config_path = os.path.join(tmpdir, "config.json")
+        with open(config_path, 'w') as fp:
+            hjson.dump(base_config, fp)
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_path, model=model, model_parameters=model.parameters())
+
+
+class TestDeprecatedDeepScaleConfig(DistributedTest):
+    world_size = 1
+
+    def test(self, base_config, tmpdir):
+        if get_accelerator().is_bf16_supported():
+            base_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            base_config["fp16"] = {"enabled": True}
+        config_path = create_config_from_dict(tmpdir, base_config)
+        parser = argparse.ArgumentParser()
+        args = parser.parse_args(args='')
+        args.deepscale_config = config_path
+        args.local_rank = 0
+
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(args=args, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model, total_samples=5, hidden_dim=hidden_dim, device=model.device)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+class TestDistInit(DistributedTest):
+    world_size = 1
+
+    def test(self, base_config):
+        if get_accelerator().is_bf16_supported():
+            base_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            base_config["fp16"] = {"enabled": True}
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=base_config,
+                                              model=model,
+                                              model_parameters=model.parameters(),
+                                              dist_init_required=True)
+        data_loader = random_dataloader(model=model, total_samples=5, hidden_dim=hidden_dim, device=model.device)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+class TestInitNoOptimizer(DistributedTest):
+    world_size = 1
+
+    def test(self, base_config):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("This test timesout with CPU accelerator")
+
+        # XXX: the bf16 path w/ no optimizer needs to be fixed
+        # if get_accelerator().is_bf16_supported():
+        #     base_config["bf16"] = {"enabled": True}
+        dtype = torch.float
+        if get_accelerator().is_fp16_supported():
+            dtype = torch.float16
+            base_config["fp16"] = {"enabled": True}
+
+        del base_config["optimizer"]
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim=hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=base_config, model=model)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=5,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=dtype)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            with pytest.raises(AssertionError):
+                model.backward(loss)
+            with pytest.raises(AssertionError):
+                model.step()
+
+
+class TestArgs(DistributedTest):
+    world_size = 1
+
+    def test_none_args(self, base_config):
+        if get_accelerator().is_bf16_supported():
+            base_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            base_config["fp16"] = {"enabled": True}
+        model = SimpleModel(hidden_dim=10)
+        model, _, _, _ = deepspeed.initialize(args=None, model=model, config=base_config)
+        data_loader = random_dataloader(model=model, total_samples=5, hidden_dim=10, device=model.device)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+
+    def test_no_args(self, base_config):
+        if get_accelerator().is_bf16_supported():
+            base_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            base_config["fp16"] = {"enabled": True}
+        model = SimpleModel(hidden_dim=10)
+        model, _, _, _ = deepspeed.initialize(model=model, config=base_config)
+        data_loader = random_dataloader(model=model, total_samples=5, hidden_dim=10, device=model.device)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+
+
+class TestNoModel(DistributedTest):
+    world_size = 1
+
+    def test(self, base_config):
+        if get_accelerator().is_bf16_supported():
+            base_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            base_config["fp16"] = {"enabled": True}
+        model = SimpleModel(hidden_dim=10)
+        with pytest.raises(AssertionError):
+            model, _, _, _ = deepspeed.initialize(model=None, config=base_config)
+
+        with pytest.raises(AssertionError):
+            model, _, _, _ = deepspeed.initialize(model, config=base_config)

tests/unit/runtime/test_ds_config_model.py

@@ -0,0 +1,86 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import json
+import os
+from typing import List, Optional
+
+from pydantic import Field, ValidationError
+
+from deepspeed.runtime import config as ds_config
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+
+
+class SimpleConf(DeepSpeedConfigModel):
+    param_1: int = 0
+    param_2_old: Optional[str] = Field(None,
+                                       json_schema_extra={
+                                           "deprecated": True,
+                                           "new_param": "param_2",
+                                           "new_param_fn": (lambda x: [x])
+                                       })
+    param_2: Optional[List[str]] = None
+    param_3: int = Field(0, alias="param_3_alias")
+
+
+def test_only_required_fields(tmpdir):
+    '''Ensure that config containing only the required fields is accepted. '''
+    cfg_json = tmpdir.mkdir('ds_config_unit_test').join('minimal.json')
+
+    with open(cfg_json, 'w') as f:
+        required_fields = {'train_batch_size': 64}
+        json.dump(required_fields, f)
+
+    run_cfg = ds_config.DeepSpeedConfig(cfg_json)
+    assert run_cfg is not None
+    assert run_cfg.train_batch_size == 64
+    assert run_cfg.train_micro_batch_size_per_gpu == 64
+    assert run_cfg.gradient_accumulation_steps == 1
+
+
+def test_config_duplicate_key(tmpdir):
+    config_dict = '''
+    {
+        "train_batch_size": 24,
+        "train_batch_size": 24,
+    }
+    '''
+    config_path = os.path.join(tmpdir, 'temp_config.json')
+
+    with open(config_path, 'w') as jf:
+        jf.write("%s" % config_dict)
+
+    with pytest.raises(ValueError):
+        run_cfg = ds_config.DeepSpeedConfig(config_path)
+
+
+def test_config_base():
+    config = SimpleConf(**{"param_1": 42})
+    assert config.param_1 == 42
+
+
+def test_config_base_deprecatedfield():
+    config = SimpleConf(**{"param_2_old": "DS"})
+    assert config.param_2 == ["DS"]
+
+
+def test_config_base_aliasfield():
+    config = SimpleConf(**{"param_3": 10})
+    assert config.param_3 == 10
+
+    config = SimpleConf(**{"param_3_alias": 10})
+    assert config.param_3 == 10
+
+
+@pytest.mark.parametrize("config_dict", [{"param_1": "DS"}, {"param_2": "DS"}, {"param_1_typo": 0}])
+def test_config_base_literalfail(config_dict):
+    with pytest.raises(ValidationError):
+        config = SimpleConf(**config_dict)
+
+
+def test_config_base_deprecatedfail():
+    with pytest.raises(AssertionError):
+        config = SimpleConf(**{"param_2": ["DS"], "param_2_old": "DS"})

tests/unit/runtime/test_ds_initialize.py

@@ -0,0 +1,493 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+from typing import Callable
+import torch
+from torch.optim import Optimizer, Adam, AdamW
+from torch.optim.lr_scheduler import _LRScheduler, LambdaLR
+
+from unit.simple_model import SimpleModel, random_dataloader
+from unit.common import DistributedTest
+from unit.util import bf16_required_version_check, required_amp_check
+
+import deepspeed
+from deepspeed.ops.adam import FusedAdam
+from deepspeed.runtime.lr_schedules import WARMUP_LR, WarmupLR
+from deepspeed.runtime.config import ADAM_OPTIMIZER
+from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer
+from deepspeed.runtime.utils import see_memory_usage
+from deepspeed.utils.torch import required_torch_version
+from deepspeed.accelerator import get_accelerator
+from deepspeed.ops.op_builder import FusedAdamBuilder
+
+
+# Ensure client multiprocessing is not broken by deepspeed import
+@pytest.mark.parametrize('method', ['spawn', 'fork', 'forkserver'])
+def test_start_method_safety(method):
+    import torch.multiprocessing as mp
+    mp.set_start_method(method, force=True)
+
+
+@pytest.mark.parametrize('zero_stage', [0, 3])
+class TestNoOptim(DistributedTest):
+    world_size = 1
+
+    def test(self, zero_stage):
+        if zero_stage == 3 and not required_torch_version(min_version=1.8):
+            pytest.skip("zero-3 param offload requires at least torch 1.8")
+
+        ds_config = {
+            'train_batch_size': self.world_size,
+            'zero_optimization': {
+                "stage": zero_stage,
+                "offload_param": {
+                    "device": "cpu"
+                }
+            }
+        }
+        if get_accelerator().is_bf16_supported():
+            ds_config["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            ds_config["fp16"] = {"enabled": True}
+        # 20B test
+        #hidden_dim = 16 * 1024
+        hidden_dim = 4
+
+        with deepspeed.zero.Init(enabled=zero_stage == 3, config_dict_or_path=ds_config):
+            model = SimpleModel(hidden_dim, nlayers=78)
+        see_memory_usage('pre-init', force=True)
+        model, _, _, _ = deepspeed.initialize(model=model, config=ds_config)
+        see_memory_usage('post-init', force=True)
+        data_loader = random_dataloader(model=model, total_samples=50, hidden_dim=hidden_dim, device=model.device)
+        for batch in data_loader:
+            model(batch[0], batch[1])
+        see_memory_usage('post-fwds', force=True)
+
+
+@pytest.mark.parametrize('optimizer_type', [None, Optimizer, Callable])
+class TestClientOptimizer(DistributedTest):
+    world_size = 1
+
+    def test(self, optimizer_type):
+
+        def _optimizer_callable(params) -> Optimizer:
+            return AdamW(params=params)
+
+        if (optimizer_type is None) and (not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME]):
+            pytest.skip("FusedAdam is not compatible")
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+
+        config_dict = {'train_batch_size': 1}
+        if optimizer_type is None:
+            client_optimizer = None
+            config_dict['optimizer'] = {'type': ADAM_OPTIMIZER}
+        elif optimizer_type is Optimizer:
+            client_optimizer = Adam(model.parameters())
+        else:
+            client_optimizer = _optimizer_callable
+
+        _, ds_optimizer, _, _ = deepspeed.initialize(config=config_dict,
+                                                     model=model,
+                                                     model_parameters=list(model.parameters()),
+                                                     optimizer=client_optimizer)
+        if client_optimizer is None:
+            assert isinstance(ds_optimizer, FusedAdam)
+        elif isinstance(client_optimizer, Optimizer):
+            assert ds_optimizer == client_optimizer
+        else:
+            assert isinstance(ds_optimizer, AdamW)
+
+
+@pytest.mark.parametrize('client_parameters', [True, False])
+class TestConfigOptimizer(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME],
+                        reason="FusedAdam is not compatible")
+    def test(self, client_parameters):
+        ds_config = {"train_batch_size": 1, "optimizer": {"type": "Adam", "params": {"lr": 0.001}}}
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+
+        if client_parameters:
+            model_parameters = list(model.parameters())
+        else:
+            model_parameters = None
+
+        _, ds_optimizer, _, _ = deepspeed.initialize(config=ds_config, model=model, model_parameters=model_parameters)
+
+        assert isinstance(ds_optimizer, FusedAdam)
+
+
+@pytest.mark.parametrize('optimizer_extension', ['zero1', 'zero2', 'zero3', 'amp', None])
+@pytest.mark.parametrize('model_dtype', ['fp16', 'bf16', 'fp32'])
+@pytest.mark.parametrize('grad_accum_dtype', [None, 'fp16', 'bf16', 'fp32'])
+class TestOptimizerImplementation(DistributedTest):
+    world_size = 1
+    reuse_dist_env = True
+
+    def test(self, optimizer_extension, model_dtype, grad_accum_dtype):
+        if not get_accelerator().is_fp16_supported():
+            if model_dtype == 'fp16' or grad_accum_dtype == 'fp16':
+                pytest.skip("fp16 is not supported")
+        if optimizer_extension == 'zero1':
+            zero_stage = 1
+        elif optimizer_extension == 'zero2':
+            zero_stage = 2
+        elif optimizer_extension == 'zero3':
+            zero_stage = 3
+        else:
+            zero_stage = 0
+        amp = (optimizer_extension == 'amp')
+        fp16 = (model_dtype == 'fp16')
+        bf16 = (model_dtype == 'bf16')
+        # Skip checks
+        if bf16 and not bf16_required_version_check():
+            pytest.skip(
+                "DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
+            )
+        if amp and not required_amp_check():
+            pytest.skip("Amp is not installed can't run amp check")
+        # Config declaration
+        ds_config = {
+            "train_batch_size": 1,
+            'fp16': {
+                'enabled': fp16
+            },
+            'bf16': {
+                'enabled': bf16
+            },
+            'amp': {
+                'enabled': amp
+            },
+            'zero_optimization': {
+                "stage": zero_stage
+            },
+            "data_types": {
+                "grad_accum_dtype": grad_accum_dtype
+            },
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.001
+                }
+            }
+        }
+
+        key = (optimizer_extension, model_dtype, grad_accum_dtype)
+
+        # Enumerate supported configurations
+        is_supported = {}
+        # ZeRO 1 Wrapper
+        is_supported[('zero1', 'fp16', None)] = True
+        is_supported[('zero1', 'fp16', 'fp16')] = True
+        is_supported[('zero1', 'fp16', 'bf16')] = True
+        is_supported[('zero1', 'fp16', 'fp32')] = True
+        is_supported[('zero1', 'bf16', None)] = True
+        is_supported[('zero1', 'bf16', 'fp16')] = True
+        is_supported[('zero1', 'bf16', 'bf16')] = True
+        is_supported[('zero1', 'bf16', 'fp32')] = True
+        is_supported[('zero1', 'fp32', None)] = True
+        is_supported[('zero1', 'fp32', 'fp16')] = True
+        is_supported[('zero1', 'fp32', 'bf16')] = True
+        is_supported[('zero1', 'fp32', 'fp32')] = True
+        # ZeRO 2 Wrapper
+        is_supported[('zero2', 'fp16', None)] = True
+        is_supported[('zero2', 'fp16', 'fp16')] = True
+        is_supported[('zero2', 'fp16', 'bf16')] = True
+        is_supported[('zero2', 'fp16', 'fp32')] = True
+        is_supported[('zero2', 'bf16', None)] = True
+        is_supported[('zero2', 'bf16', 'fp16')] = True
+        is_supported[('zero2', 'bf16', 'bf16')] = True
+        is_supported[('zero2', 'bf16', 'fp32')] = True
+        is_supported[('zero2', 'fp32', None)] = True
+        is_supported[('zero2', 'fp32', 'fp16')] = True
+        is_supported[('zero2', 'fp32', 'bf16')] = True
+        is_supported[('zero2', 'fp32', 'fp32')] = True
+        # ZeRO 3 Wrapper
+        is_supported[('zero3', 'fp16', None)] = True
+        is_supported[('zero3', 'fp16', 'fp16')] = True
+        is_supported[('zero3', 'fp16', 'bf16')] = True
+        is_supported[('zero3', 'fp16', 'fp32')] = True
+        is_supported[('zero3', 'bf16', None)] = True
+        is_supported[('zero3', 'bf16', 'fp16')] = True
+        is_supported[('zero3', 'bf16', 'bf16')] = True
+        is_supported[('zero3', 'bf16', 'fp32')] = True
+        is_supported[('zero3', 'fp32', None)] = True
+        is_supported[('zero3', 'fp32', 'fp16')] = True
+        is_supported[('zero3', 'fp32', 'bf16')] = True
+        is_supported[('zero3', 'fp32', 'fp32')] = True
+        # Amp Wrapper
+        is_supported[('amp', 'fp32', None)] = True
+        is_supported[('amp', 'fp32', 'fp32')] = True
+        # FP16 Wrapper
+        is_supported[(None, 'fp16', None)] = True
+        is_supported[(None, 'fp16', 'fp16')] = True
+        # BF16 Wrapper
+        is_supported[(None, 'bf16', 'bf16')] = True
+        is_supported[(None, 'bf16', None)] = True
+        # No Wrapper
+        is_supported[(None, 'fp32', None)] = True
+        is_supported[(None, 'fp32', 'fp32')] = True
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        model_parameters = list(model.parameters())
+
+        if key in is_supported:
+            _, ds_optimizer, _, _ = deepspeed.initialize(config=ds_config,
+                                                         model=model,
+                                                         model_parameters=model_parameters)
+            assert True
+        else:
+            with pytest.raises(NotImplementedError):
+                _, ds_optimizer, _, _ = deepspeed.initialize(config=ds_config,
+                                                             model=model,
+                                                             model_parameters=model_parameters)
+
+
+class TestBf16ZeRO0UnfusedOptimizer(DistributedTest):
+    world_size = 1
+    reuse_dist_env = True
+
+    def test_static_scale_and_zero_grad_after_step(self):
+        if not bf16_required_version_check():
+            pytest.skip(
+                "DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
+            )
+
+        hidden_dim = 16
+        model = SimpleModel(hidden_dim)
+        client_optimizer = AdamW(model.parameters(), lr=1e-4)
+        ds_config = {
+            "train_batch_size": 1,
+            "train_micro_batch_size_per_gpu": 1,
+            "bf16": {
+                "enabled": True
+            },
+            "zero_optimization": {
+                "stage": 0
+            },
+        }
+
+        engine, _, _, _ = deepspeed.initialize(config=ds_config,
+                                               model=model,
+                                               model_parameters=list(model.parameters()),
+                                               optimizer=client_optimizer)
+
+        assert isinstance(engine.optimizer, FP16_UnfusedOptimizer)
+        assert engine.optimizer.low_precision_dtype == torch.bfloat16
+        assert engine.optimizer.loss_scale_config.dynamic_loss_scale is False
+        assert engine.optimizer.loss_scale_config.cur_scale == 1
+
+        data_loader = random_dataloader(model=engine,
+                                        total_samples=1,
+                                        hidden_dim=hidden_dim,
+                                        device=engine.device,
+                                        dtype=torch.bfloat16)
+        batch = next(iter(data_loader))
+
+        loss = engine(batch[0], batch[1])
+        engine.backward(loss)
+        assert any(param.grad is not None for param in engine.module.parameters() if param.requires_grad)
+
+        engine.step()
+        assert all(param.grad is None for param in engine.module.parameters() if param.requires_grad)
+
+
+@pytest.mark.parametrize("scheduler_type", [None, _LRScheduler, Callable])
+@pytest.mark.parametrize("optimizer_type", [None, Optimizer, Callable])
+class TestClientLrScheduler(DistributedTest):
+    world_size = 1
+
+    def test(self, scheduler_type, optimizer_type):
+
+        def _my_lambda(epoch):
+            return epoch // 10
+
+        def _optimizer_callable(params) -> Optimizer:
+            return torch.optim.AdamW(params=params)
+
+        def _lr_scheduler_callable(optimizer) -> _LRScheduler:
+            return LambdaLR(optimizer, _my_lambda)
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+
+        config_dict = {'train_batch_size': 1}
+
+        client_optimizer = None
+        client_scheduler = None
+
+        if optimizer_type is None:
+            config_dict['optimizer'] = {'type': ADAM_OPTIMIZER}
+        elif optimizer_type is Optimizer:
+            client_optimizer = torch.optim.Adam(model.parameters())
+        else:
+            client_optimizer = _optimizer_callable
+
+        if scheduler_type is None:
+            config_dict['scheduler'] = {'type': WARMUP_LR, 'params': {}}
+        elif scheduler_type == _LRScheduler:
+            if isinstance(client_optimizer, Optimizer):
+                client_scheduler = LambdaLR(client_optimizer, _my_lambda)
+            else:
+                # Verify invalid combination is correctly handled
+                client_scheduler = LambdaLR(torch.optim.Adam(model.parameters()), _my_lambda)
+        else:
+            client_scheduler = _lr_scheduler_callable
+
+        if isinstance(client_scheduler, _LRScheduler) and not isinstance(client_optimizer, Optimizer):
+            with pytest.raises(AssertionError):
+                _, _, _, _ = deepspeed.initialize(config=config_dict,
+                                                  model=model,
+                                                  model_parameters=list(model.parameters()),
+                                                  optimizer=client_optimizer,
+                                                  lr_scheduler=client_scheduler)
+        else:
+            _, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                            model=model,
+                                                            model_parameters=list(model.parameters()),
+                                                            optimizer=client_optimizer,
+                                                            lr_scheduler=client_scheduler)
+            if client_scheduler is None:
+                assert isinstance(ds_lr_scheduler, WarmupLR)
+            elif isinstance(client_scheduler, _LRScheduler):
+                assert ds_lr_scheduler == client_scheduler
+            else:
+                assert isinstance(ds_lr_scheduler, LambdaLR)
+
+
+@pytest.mark.parametrize("scheduler_type", [None, _LRScheduler, Callable])
+class TestClientLrSchedulerInit(DistributedTest):
+    world_size = 1
+
+    def test_same_lrscheler_and_callable(self, scheduler_type):
+        """
+        Expect behavior
+
+        if lr scheduler is defined in code and passed into initialize as arg,
+        it will be used even this is a lr scheduler has been defined in config.
+
+        Initialize lr scheduler from config when no lr scheduler is defined in code.
+        """
+
+        def _my_lambda(epoch):
+            return epoch // 10
+
+        def _lr_scheduler_callable(optimizer) -> _LRScheduler:
+            return LambdaLR(optimizer, _my_lambda)
+
+        config_dict = {'train_batch_size': 1}
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+
+        client_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+        if scheduler_type is None:
+            config_dict['scheduler'] = {'type': WARMUP_LR, 'params': {}}
+            client_scheduler = None
+        elif scheduler_type == _LRScheduler:
+            client_scheduler = LambdaLR(client_optimizer, _my_lambda)
+        else:
+            client_scheduler = _lr_scheduler_callable
+
+        _, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                        model=model,
+                                                        model_parameters=list(model.parameters()),
+                                                        optimizer=client_optimizer,
+                                                        lr_scheduler=client_scheduler)
+        if scheduler_type is None:
+            # in this case, we initialize from config
+            assert not isinstance(ds_lr_scheduler, LambdaLR)
+            assert isinstance(ds_lr_scheduler, WarmupLR)
+        else:
+            # in this case, we initialize from passed-in scheduler
+            assert isinstance(ds_lr_scheduler, LambdaLR)
+            assert not isinstance(ds_lr_scheduler, WarmupLR)
+
+    def test_diff_lrscheler_and_callable(self, scheduler_type):
+        """
+        In this test,
+        the LambdaLR will be used for lrscheduler type
+        and the StepLR will be used for callable type
+        """
+
+        from torch.optim.lr_scheduler import StepLR
+
+        def _my_lambda(epoch):
+            return epoch // 10
+
+        def _lr_scheduler_callable(optimizer) -> _LRScheduler:
+            return StepLR(optimizer, step_size=30)
+
+        config_dict = {'train_batch_size': 1}
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+
+        client_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+        if scheduler_type is None:
+            config_dict['scheduler'] = {'type': WARMUP_LR, 'params': {}}
+            client_scheduler = None
+        elif scheduler_type == _LRScheduler:
+            client_scheduler = LambdaLR(client_optimizer, _my_lambda)
+        else:
+            client_scheduler = _lr_scheduler_callable
+
+        _, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                        model=model,
+                                                        model_parameters=list(model.parameters()),
+                                                        optimizer=client_optimizer,
+                                                        lr_scheduler=client_scheduler)
+        if scheduler_type is None:
+            assert isinstance(ds_lr_scheduler, WarmupLR)
+        elif scheduler_type == _LRScheduler:
+            assert isinstance(ds_lr_scheduler, LambdaLR)
+        else:
+            # callable
+            assert isinstance(ds_lr_scheduler, StepLR)
+
+    def test_diff_lrscheler_and_callable_onecyclelr_steplr(self, scheduler_type):
+
+        from deepspeed.runtime.lr_schedules import OneCycle, ONE_CYCLE, CYCLE_MIN_LR, CYCLE_MAX_LR
+        from torch.optim.lr_scheduler import OneCycleLR, StepLR
+
+        def _lr_scheduler_callable(optimizer) -> _LRScheduler:
+            return OneCycleLR(optimizer, max_lr=0.01, total_steps=200)
+
+        config_dict = {'train_batch_size': 1}
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+
+        client_optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+        if scheduler_type is None:
+            config_dict['scheduler'] = {'type': ONE_CYCLE, 'params': {CYCLE_MIN_LR: 0, CYCLE_MAX_LR: 0.1}}
+            client_scheduler = None
+        elif scheduler_type == _LRScheduler:
+            client_scheduler = StepLR(client_optimizer, step_size=30)
+        else:
+            client_scheduler = _lr_scheduler_callable
+
+        _, _, _, ds_lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                        model=model,
+                                                        model_parameters=list(model.parameters()),
+                                                        optimizer=client_optimizer,
+                                                        lr_scheduler=client_scheduler)
+        if scheduler_type is None:
+            assert isinstance(ds_lr_scheduler, OneCycle)
+        elif scheduler_type == _LRScheduler:
+            assert isinstance(ds_lr_scheduler, StepLR)
+        else:
+            # callable
+            assert isinstance(ds_lr_scheduler, OneCycleLR)

tests/unit/runtime/test_engine_eigenvalue.py

@@ -0,0 +1,18 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.runtime.engine import _eigenvalue_summary_events
+
+
+def test_eigenvalue_summary_events_use_block_values():
+    block_eigenvalue = {
+        "layer0.weight": (1.25, 0),
+        "layer1.weight": (0.5, 1),
+    }
+
+    assert _eigenvalue_summary_events(block_eigenvalue, global_samples=128) == [
+        ("Train/Eigenvalues/ModelBlockParam_0", 1.25, 128),
+        ("Train/Eigenvalues/ModelBlockParam_1", 0.5, 128),
+    ]

tests/unit/runtime/test_lr_schedulers.py

@@ -0,0 +1,560 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import math
+
+import torch
+import deepspeed
+import pytest
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataloader
+from deepspeed.runtime.lr_schedules import LR_RANGE_TEST, LR_RANGE_TEST_MIN_LR, LR_RANGE_TEST_STEP_RATE, LR_RANGE_TEST_STEP_SIZE, LR_RANGE_TEST_STAIRCASE
+from deepspeed.runtime.lr_schedules import WARMUP_LR, WARMUP_MIN_LR, WARMUP_MAX_LR, WARMUP_NUM_STEPS, WARMUP_TYPE, WARMUP_LOG_RATE, WARMUP_LINEAR_RATE
+from deepspeed.runtime.lr_schedules import ONE_CYCLE, CYCLE_MIN_LR, CYCLE_MAX_LR, CYCLE_FIRST_STEP_SIZE, DECAY_LR_RATE, DECAY_STEP_SIZE
+from deepspeed.runtime.lr_schedules import CYCLE_MIN_MOM, CYCLE_MAX_MOM, DECAY_MOM_RATE
+from deepspeed.runtime.lr_schedules import WARMUP_DECAY_LR, TOTAL_NUM_STEPS
+from deepspeed.runtime.lr_schedules import WARMUP_COSINE_LR, WARMUP_MIN_RATIO, COS_MIN_RATIO, WarmupCosineLR
+from deepspeed.runtime.lr_schedules import WarmupLR, WarmupDecayLR
+
+
+def _verify_continuous_decrease(values):
+    for i in range(len(values) - 1):
+        assert values[i] > values[i + 1]
+
+
+def _verify_continuous_increase(values):
+    for i in range(len(values) - 1):
+        assert values[i] < values[i + 1]
+
+
+def _verify_staircase_increase(values, step_size):
+    num_values = len(values)
+    for i in range(0, num_values, step_size):
+        j = min(i + step_size, num_values)
+        assert all([values[i] == v for v in values[i:j]])
+
+
+@pytest.mark.parametrize("scheduler_type,params", [(WARMUP_LR, {}),
+                                                   (WARMUP_DECAY_LR, {
+                                                       WARMUP_NUM_STEPS: 10,
+                                                       TOTAL_NUM_STEPS: 20
+                                                   }), (WARMUP_COSINE_LR, {
+                                                       WARMUP_NUM_STEPS: 10,
+                                                       TOTAL_NUM_STEPS: 20
+                                                   }), (ONE_CYCLE, {
+                                                       CYCLE_MIN_LR: 0,
+                                                       CYCLE_MAX_LR: 0.1
+                                                   }), (LR_RANGE_TEST, {})])
+class TestGetLrBeforeTrain(DistributedTest):
+    world_size = 1
+
+    def test(self, scheduler_type, params):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                },
+            },
+            "scheduler": {
+                "type": scheduler_type,
+                "params": params
+            },
+            "gradient_clipping": 1.0
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, lr_scheduler = 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)
+
+        true_lrs = lr_scheduler.get_lr()
+        for group, true_lr in zip(model.optimizer.param_groups, true_lrs):
+            assert group['lr'] == true_lr, f"True lr {true_lr}, optimizer lr {group['lr']}"
+
+        for n, batch in enumerate(data_loader):
+            # get lr before training starts
+            lr_scheduler.get_lr()
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+@pytest.mark.parametrize("warmup_num_steps", [10, 15, 19, 33])
+@pytest.mark.parametrize("warmup_type", [WARMUP_LOG_RATE, WARMUP_LINEAR_RATE])
+class TestLrSchedule(DistributedTest):
+    world_size = 1
+
+    def test_lr_warmup_schedule(self, warmup_num_steps, warmup_type):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                },
+            },
+            "scheduler": {
+                "type": WARMUP_LR,
+                "params": {
+                    WARMUP_MIN_LR: 0.1,
+                    WARMUP_MAX_LR: 0.2,
+                    WARMUP_NUM_STEPS: warmup_num_steps,
+                    WARMUP_TYPE: warmup_type,
+                }
+            },
+            "gradient_clipping": 1.0
+        }
+        schedule_params = config_dict["scheduler"]["params"]
+        total_num_steps = 2 * warmup_num_steps
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                         model=model,
+                                                         model_parameters=model.parameters())
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=total_num_steps * 2,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float)
+        step_lrs = []
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+            step_lrs.append(lr_scheduler.get_lr())
+
+        # Verify initial lr
+        assert step_lrs[0] == [schedule_params[WARMUP_MIN_LR]]
+
+        # Verify warmup completion
+        warmup_num_steps = schedule_params[WARMUP_NUM_STEPS]
+        warmup_max_lr = [schedule_params[WARMUP_MAX_LR]]
+        assert step_lrs[warmup_num_steps] == warmup_max_lr
+
+        # Verify post-warmup completion
+        assert all([warmup_max_lr == lr for lr in step_lrs[warmup_num_steps:]])
+
+    def test_lr_warmup_decay_schedule(self, warmup_num_steps, warmup_type):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                },
+            },
+            "scheduler": {
+                "type": WARMUP_DECAY_LR,
+                "params": {
+                    WARMUP_MIN_LR: 0.1,
+                    WARMUP_MAX_LR: 0.2,
+                    WARMUP_NUM_STEPS: warmup_num_steps,
+                    TOTAL_NUM_STEPS: warmup_num_steps * 2,
+                    WARMUP_TYPE: warmup_type
+                }
+            },
+            "gradient_clipping": 1.0
+        }
+        schedule_params = config_dict["scheduler"]["params"]
+        total_num_steps = schedule_params[TOTAL_NUM_STEPS]
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                         model=model,
+                                                         model_parameters=model.parameters())
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=total_num_steps * 2,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float)
+        step_lrs = []
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+            step_lrs.append(lr_scheduler.get_lr())
+
+        # Verify initial lr
+        assert step_lrs[0] == [schedule_params[WARMUP_MIN_LR]]
+
+        # Verify lr at warmup completion
+        warmup_num_steps = schedule_params[WARMUP_NUM_STEPS]
+        warmup_max_lr = [schedule_params[WARMUP_MAX_LR]]
+        assert step_lrs[warmup_num_steps] == warmup_max_lr
+
+        # Verify decay phase
+        previous_lr = warmup_max_lr
+        for lr in step_lrs[warmup_num_steps + 1:]:
+            assert lr < previous_lr
+            previous_lr = lr
+
+
+@pytest.mark.parametrize("scheduler_type,params", [(WARMUP_LR, {}),
+                                                   (WARMUP_DECAY_LR, {
+                                                       WARMUP_NUM_STEPS: 5,
+                                                       TOTAL_NUM_STEPS: 10
+                                                   }),
+                                                   (ONE_CYCLE, {
+                                                       CYCLE_MIN_LR: 0,
+                                                       CYCLE_MAX_LR: 0.1,
+                                                       CYCLE_FIRST_STEP_SIZE: 5,
+                                                       DECAY_STEP_SIZE: 5
+                                                   }),
+                                                   (LR_RANGE_TEST, {
+                                                       LR_RANGE_TEST_MIN_LR: 1e-4,
+                                                       LR_RANGE_TEST_STEP_SIZE: 1
+                                                   })])
+class TestSchedulerOptimizerParity(DistributedTest):
+    world_size = 1
+
+    def test(self, scheduler_type, params):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                },
+            },
+            "scheduler": {
+                "type": scheduler_type,
+                "params": params
+            },
+            "gradient_clipping": 1.0
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, lr_scheduler = 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()
+            assert lr_scheduler.get_lr() == model.get_lr()
+
+
+@pytest.mark.parametrize("min_lr, step_rate, step_size, staircase",
+                         [(1e-4, 1e-5, 1, True),
+                          (1e-5, 1e-5, 1, False),
+                          (1e-4, 1e-3, 10, True),
+                          (1e-3, 1e-3, 10, False),
+                          (1e-2, 1e-2, 19, True),
+                          (1e-2, 1e-2, 19, False)
+                           ])# yapf: disable
+class TestLrRange(DistributedTest):
+    world_size = 1
+
+    def test(self, min_lr, step_rate, step_size, staircase):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                },
+            },
+            "scheduler": {
+                "type": LR_RANGE_TEST,
+                "params": {
+                    LR_RANGE_TEST_MIN_LR: min_lr,
+                    LR_RANGE_TEST_STEP_RATE: step_rate,
+                    LR_RANGE_TEST_STEP_SIZE: step_size,
+                    LR_RANGE_TEST_STAIRCASE: staircase
+                }
+            },
+            "gradient_clipping": 1.0
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                         model=model,
+                                                         model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=max(50, step_size * 2),
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float)
+
+        step_lrs = []
+        for _, batch in enumerate(data_loader):
+            step_lrs.extend(lr_scheduler.get_lr())
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        # Verify starting lr
+        assert step_lrs[0] == min_lr
+
+        if staircase:
+            # Verify staircase increasing lr
+            _verify_staircase_increase(step_lrs, step_size)
+        else:
+            # Verify continuous increasing lr
+            _verify_continuous_increase(step_lrs)
+
+
+class TestOneCycle(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.parametrize("min_lr, max_lr, decay_rate, cycle_step_size, decay_step_size",
+                             [
+                                 (1e-5, 1e-2, 1e-3, 10, 10),
+                                 (1e-3, 1e-1, 0, 21, 21),
+                                 (1e-5, 1e-2, 1e-3, 10, 10),
+                                 (1e-3, 1e-1, 1e-1, 21, 21),
+                                 (1e-5, 1e-1, 0, 10, 0),
+                             ])  # yapf: disable
+    def test_lr(self, min_lr, max_lr, decay_rate, cycle_step_size, decay_step_size):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                },
+            },
+            "scheduler": {
+                "type": ONE_CYCLE,
+                "params": {
+                    CYCLE_MIN_LR: min_lr,
+                    CYCLE_MAX_LR: max_lr,
+                    DECAY_LR_RATE: decay_rate,
+                    CYCLE_FIRST_STEP_SIZE: cycle_step_size,
+                    DECAY_STEP_SIZE: decay_step_size
+                }
+            },
+            "gradient_clipping": 1.0
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                         model=model,
+                                                         model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=max(50, cycle_step_size * 3),
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float)
+
+        step_lrs = []
+        for _, batch in enumerate(data_loader):
+            step_lrs.extend(lr_scheduler.get_lr())
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        # Verify starting lr
+        assert step_lrs[0] == min_lr
+
+        # Verify peak lr
+        assert step_lrs[cycle_step_size] == max_lr
+
+        # Verify increasing phase
+        _verify_continuous_increase(step_lrs[:cycle_step_size])
+
+        # Verify decreasing phase
+        _verify_continuous_decrease(step_lrs[cycle_step_size:(cycle_step_size * 2)])
+
+        # Verify decay phase
+        if decay_rate > 0:
+            _verify_continuous_decrease(step_lrs[(cycle_step_size * 2):])
+
+    @pytest.mark.parametrize("min_mom, max_mom, decay_rate, step_size",
+                             [
+                                 (0.08, 0.09, 1e-3, 10),
+                                 (0.08, 0.09, 0, 21),
+                                 (0.08, 0.09, 1e-3, 10),
+                                 (0.08, 0.09, 0, 21),
+                             ]) # yapf: disable
+    def test_mom(self, min_mom, max_mom, decay_rate, step_size):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                },
+            },
+            "scheduler": {
+                "type": ONE_CYCLE,
+                "params": {
+                    CYCLE_MIN_LR: 1e-3,
+                    CYCLE_MAX_LR: 1e-2,
+                    CYCLE_MIN_MOM: min_mom,
+                    CYCLE_MAX_MOM: max_mom,
+                    DECAY_MOM_RATE: decay_rate,
+                    CYCLE_FIRST_STEP_SIZE: step_size,
+                    DECAY_STEP_SIZE: step_size
+                }
+            },
+            "gradient_clipping": 1.0
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                         model=model,
+                                                         model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=max(50, step_size * 3),
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float)
+
+        step_moms = []
+        for _, batch in enumerate(data_loader):
+            step_moms.append(lr_scheduler.get_mom())
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        # Verify starting lr
+        assert step_moms[0][0][0] == max_mom
+
+        # Verify peak lr
+        assert step_moms[step_size][0][0] == min_mom
+
+        # Verify decreasing phase
+        _verify_continuous_decrease(step_moms[:step_size])
+
+        # Verify increasing phase
+        _verify_continuous_increase(step_moms[step_size:(step_size * 2)])
+
+        # Verify decay phase
+        if decay_rate > 0:
+            _verify_continuous_increase(step_moms[(step_size * 2):])
+
+
+class TestWarmupCosineLR(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.parametrize("total_num_steps, warmup_num_steps, cos_min_ratio, warmup_min_ratio",
+                             [
+                                 (100, 10, 0.1, 0.2),
+                                 (200, 20, 0.1, 0.2),
+                                 (500, 30, 0.0, 0.2),
+                                 (600, 300, 0.1, 0.0),
+                                 (600, 550, 0.0, 0.0),
+                             ])  # yapf: disable
+    def test_lr(self, total_num_steps, warmup_num_steps, cos_min_ratio, warmup_min_ratio):
+        opt_lr = 0.0015
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": opt_lr
+                },
+            },
+            "scheduler": {
+                "type": WARMUP_COSINE_LR,
+                "params": {
+                    TOTAL_NUM_STEPS: total_num_steps,
+                    WARMUP_MIN_RATIO: warmup_min_ratio,
+                    WARMUP_NUM_STEPS: warmup_num_steps,
+                    COS_MIN_RATIO: cos_min_ratio,
+                }
+            },
+            "gradient_clipping": 1.0
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, lr_scheduler = deepspeed.initialize(config=config_dict,
+                                                         model=model,
+                                                         model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model,
+                                        total_samples=max(50, total_num_steps * 3),
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float)
+
+        step_lrs = []
+        for _, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+            step_lrs.extend(lr_scheduler.get_lr())
+
+        # Verify starting lr
+        assert abs(step_lrs[0] - opt_lr * warmup_min_ratio) < 1e-7
+
+        # Verify peak lr
+        assert abs(step_lrs[warmup_num_steps - 1] - opt_lr) < 1e-7
+
+        # Verify end lr
+        assert abs(step_lrs[total_num_steps - 1] - opt_lr * cos_min_ratio) < 1e-7
+
+        # Verify increasing phase
+        _verify_continuous_increase(step_lrs[:warmup_num_steps])
+
+        # Verify decreasing phase
+        _verify_continuous_decrease(step_lrs[warmup_num_steps:total_num_steps])
+
+
+def test_warmup_cosine_lr_initializes_all_param_groups():
+    dense = torch.nn.Parameter(torch.zeros(1))
+    expert = torch.nn.Parameter(torch.zeros(1))
+    optimizer = torch.optim.Adam([{"params": [dense], "lr": 0.0015}, {"params": [expert], "lr": 0.003}])
+
+    scheduler = WarmupCosineLR(optimizer=optimizer, total_num_steps=100, warmup_num_steps=10, warmup_min_ratio=0.0)
+
+    assert scheduler.get_lr_ratio() == 0.0
+    assert scheduler.get_lr() == [0.0, 0.0]
+    assert scheduler.get_last_lr() == [0.0, 0.0]
+    assert [group["lr"] for group in optimizer.param_groups] == [0.0, 0.0]
+
+    scheduler.step(1)
+
+    expected_ratio = math.log(2) / math.log(10)
+    expected_lrs = [0.0015 * expected_ratio, 0.003 * expected_ratio]
+
+    assert scheduler.get_lr_ratio() == pytest.approx(expected_ratio)
+    assert scheduler.get_lr() == pytest.approx(expected_lrs)
+    assert scheduler.get_last_lr() == pytest.approx(expected_lrs)
+    assert [group["lr"] for group in optimizer.param_groups] == pytest.approx(expected_lrs)
+
+
+@pytest.mark.parametrize("scheduler_cls", [WarmupLR, WarmupDecayLR, WarmupCosineLR])
+@pytest.mark.parametrize("bad_warmup_num_steps", [None, -5])
+def test_warmup_schedulers_reject_invalid_warmup_num_steps(scheduler_cls, bad_warmup_num_steps):
+    param = torch.nn.Parameter(torch.zeros(1))
+    optimizer = torch.optim.Adam([param], lr=0.001)
+
+    kwargs = {"optimizer": optimizer, "warmup_num_steps": bad_warmup_num_steps}
+    if scheduler_cls in (WarmupDecayLR, WarmupCosineLR):
+        kwargs["total_num_steps"] = 100
+
+    with pytest.raises(ValueError):
+        scheduler_cls(**kwargs)

tests/unit/runtime/test_multi_output_model.py

@@ -0,0 +1,129 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+from pytest import approx
+from unit.util import torch_assert_close
+from unit.common import DistributedTest, preferred_dtype
+from unit.multi_output_model import MultiOutputModel, multi_output_dataloader
+
+
+class TestTwoOutputModel(DistributedTest):
+    world_size = 1
+
+    def test(self, tmpdir):
+        grad_accumulation_steps = 2
+        micro_batch_size = 1
+        world_size = self.world_size
+        config_dict = {
+            "train_micro_batch_size_per_gpu": micro_batch_size,
+            "gradient_accumulation_steps": grad_accumulation_steps,
+            "train_batch_size": micro_batch_size * grad_accumulation_steps * world_size,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+
+        hidden_dim = 10
+        weight_value = 0.1
+
+        model = MultiOutputModel(hidden_dim, weight_value)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        total_samples = 4
+        data_loader = multi_output_dataloader(model=model,
+                                              total_samples=total_samples,
+                                              hidden_dim=hidden_dim,
+                                              device=model.device,
+                                              inputs=[1.0, 2.0],
+                                              targets=[1, 2])
+        for n, batch in enumerate(data_loader):
+            assert len(batch) % 2 == 0, \
+                 "multi_output_dataloader failed to return even number of data samples (input+target)"
+
+            midpoint = len(batch) // 2
+            inputs, targets = batch[:midpoint], batch[midpoint:]
+            loss_tuple = model(inputs, targets)
+
+            expected_loss = torch.tensor(2.302734375, dtype=preferred_dtype(), device=model.device)
+            for loss in loss_tuple:
+                assert loss.shape == torch.Size([])
+                assert loss.item() == approx(expected_loss.item())
+
+            summed_loss = sum(loss_tuple)
+            scaled_loss = model.backward(summed_loss)
+            expected_scaled_loss = summed_loss / grad_accumulation_steps
+            torch_assert_close(scaled_loss, expected_scaled_loss)
+
+            model.step()
+
+
+class TestThreeOutputModel(DistributedTest):
+    world_size = 1
+
+    def test(self, tmpdir):
+        grad_accumulation_steps = 3
+        micro_batch_size = 1
+        world_size = 1
+        config_dict = {
+            "train_micro_batch_size_per_gpu": micro_batch_size,
+            "gradient_accumulation_steps": grad_accumulation_steps,
+            "train_batch_size": micro_batch_size * grad_accumulation_steps * world_size,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+
+        hidden_dim = 10
+        weight_value = 0.1
+
+        model = MultiOutputModel(hidden_dim, weight_value)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+        total_samples = grad_accumulation_steps * micro_batch_size * 2
+        data_loader = multi_output_dataloader(model=model,
+                                              total_samples=total_samples,
+                                              hidden_dim=hidden_dim,
+                                              device=model.device,
+                                              inputs=[1.0, 2.0, 3.0],
+                                              targets=[1, 2, 3])
+        for n, batch in enumerate(data_loader):
+            assert len(batch) % 2 == 0, \
+                 "multi_output_dataloader failed to return even number of data samples (input+target)"
+
+            midpoint = len(batch) // 2
+            inputs, targets = batch[:midpoint], batch[midpoint:]
+            loss_tuple = model(inputs, targets)
+            assert len(loss_tuple) == 3
+
+            expected_loss = torch.tensor(2.302734375, dtype=preferred_dtype(), device=model.device)
+
+            for loss in loss_tuple:
+                assert loss.shape == torch.Size([])
+                assert loss.item() == approx(expected_loss.item())
+
+            summed_loss = sum(loss_tuple)
+            scaled_loss = model.backward(summed_loss)
+            expected_scaled_loss = summed_loss / grad_accumulation_steps
+            torch_assert_close(scaled_loss, expected_scaled_loss)
+
+            model.step()

tests/unit/runtime/test_multiple_models.py

@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import deepspeed
+import deepspeed.comm as dist
+import torch
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataloader
+
+
+def create_model(config_dict):
+    hidden_dim = 64
+    model = SimpleModel(hidden_dim)
+    model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+    return model
+
+
+def train_shared_loss(num_models, config_dict, dtype):
+    hidden_dim = 64
+
+    models = [create_model(config_dict) for _ in range(num_models)]
+    data_loader = random_dataloader(model=models[0],
+                                    total_samples=4,
+                                    hidden_dim=hidden_dim,
+                                    device=models[0].device,
+                                    dtype=dtype)
+    dist.barrier()
+    for _, batch in enumerate(data_loader):
+        losses = [m.module(batch[0], batch[1]) for m in models]
+        loss = sum(l / (i + 1) for i, l in enumerate(losses))
+        loss.backward()
+
+        for m in models:
+            m._backward_epilogue()
+
+        for m in models:
+            m.step()
+
+        for m in models:
+            m.optimizer.zero_grad()
+
+    for m in models:
+        m.destroy()
+
+
+def train_independent_loss(num_models, config_dict, dtype):
+    hidden_dim = 64
+
+    models = [create_model(config_dict) for _ in range(num_models)]
+    data_loader = random_dataloader(model=models[0],
+                                    total_samples=4,
+                                    hidden_dim=hidden_dim,
+                                    device=models[0].device,
+                                    dtype=dtype)
+    dist.barrier()
+    for _, batch in enumerate(data_loader):
+        losses = [m.module(batch[0], batch[1]) for m in models]
+        for m, loss in zip(models, losses):
+            m.backward(loss)
+            m.step()
+
+    for m in models:
+        m.destroy()
+
+
+@pytest.mark.parametrize('num_models', [1, 2, 3])
+class TestMultipleModels(DistributedTest):
+    world_size = 2
+    reuse_dist_env = True
+
+    @pytest.mark.parametrize('shared_loss', [False, True])
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    @pytest.mark.parametrize('fp32_grad_accum', [False, True])
+    @pytest.mark.parametrize('contiguous_gradients', [False, True])
+    @pytest.mark.parametrize('overlap_comm', [False, True])
+    def test_zero_optimizer(self, num_models, shared_loss, zero_stage, fp32_grad_accum, contiguous_gradients,
+                            overlap_comm):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-4
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "contiguous_gradients": contiguous_gradients,
+                "overlap_comm": overlap_comm,
+            },
+            "fp16": {
+                "initial_scale_power": 8,
+                "enabled": True
+            },
+        }
+        if fp32_grad_accum:
+            config_dict["data_types"] = {"grad_accum_dtype": "fp32"}
+
+        if shared_loss:
+            train_shared_loss(num_models=num_models, config_dict=config_dict, dtype=torch.float16)
+        else:
+            train_independent_loss(num_models=num_models, config_dict=config_dict, dtype=torch.float16)
+
+    # TODO: Combination of shared_loss==True and bf16.immediate_grad_update==False is currently broken
+    @pytest.mark.parametrize('shared_loss', [False, True])
+    def test_bf16_optimizer(self, num_models, shared_loss):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-4
+                }
+            },
+            "zero_optimization": {
+                "stage": 1,
+            },
+            "bf16": {
+                "enabled": True,
+                "immediate_grad_update": True,
+            },
+            "data_types": {
+                "grad_accum_dtype": "fp32"
+            }
+        }
+
+        if shared_loss:
+            train_shared_loss(num_models=num_models, config_dict=config_dict, dtype=torch.bfloat16)
+        else:
+            train_independent_loss(num_models=num_models, config_dict=config_dict, dtype=torch.bfloat16)

tests/unit/runtime/test_mup_optimizers.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed
+import torch
+import pytest
+
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataloader
+from mup.shape import set_base_shapes
+from deepspeed.accelerator import get_accelerator
+
+
+@pytest.mark.parametrize("optimizer, expected_opt_class", [("MuAdam", torch.optim.Adam),
+                                                           ("MuAdamW", torch.optim.AdamW), ("MuSGD", torch.optim.SGD)]) # yapf: disable
+@pytest.mark.parametrize("zero_offload", [True, False]) # yapf: disable
+class TestMuPOptimizers(DistributedTest):
+    world_size = 1
+    reuse_dist_env = True
+
+    def test(self, optimizer, expected_opt_class, zero_offload):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "zero_allow_untested_optimizer": True,
+            "optimizer": {
+                "type": optimizer,
+                "params": {
+                    "lr": 0.00015,
+                }
+            },
+            "gradient_clipping": 1.0,
+            "zero_optimization": {
+                "stage": 2,
+                "cpu_offload": zero_offload
+            }
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        set_base_shapes(model, None)
+        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)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        ds_optimizer = model.optimizer.optimizer
+        assert isinstance(ds_optimizer, expected_opt_class)

tests/unit/runtime/test_no_sync_ctxt.py

@@ -0,0 +1,197 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+from contextlib import nullcontext
+import torch
+
+from unit.simple_model import SimpleModel, random_dataloader
+from unit.common import DistributedTest
+
+import deepspeed
+import deepspeed.comm as dist
+from deepspeed.utils import safe_get_full_grad
+
+
+class TestNoSyncCtxt(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+    @pytest.mark.parametrize("zero_stage", [0, 1, 2, 3])
+    def test_zero_stage(self, zero_stage, dtype):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            },
+        }
+
+        invalid_cfg = zero_stage > 1
+        if dtype == torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+        elif dtype == torch.float16:
+            config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
+
+        hidden_dim = 64
+        total_samples = 32
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=total_samples,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=dtype)
+        dist.barrier()
+
+        with pytest.raises(AssertionError) if invalid_cfg else nullcontext() as assertinfo:
+            with model.no_sync():
+                for _, batch in enumerate(data_loader):
+                    loss = model(batch[0], batch[1])
+                    model.backward(loss)
+        if invalid_cfg:
+            assert ("no_sync context manager is incompatible" in str(assertinfo))
+
+    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+    @pytest.mark.parametrize("zero_stage", [0, 1])
+    def test_engine_step(self, zero_stage, dtype):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            },
+        }
+
+        if dtype == torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+        elif dtype == torch.float16:
+            config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
+
+        hidden_dim = 64
+        total_samples = 32
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=total_samples,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=dtype)
+        dist.barrier()
+
+        with model.no_sync():
+            for _, batch in enumerate(data_loader):
+                loss = model(batch[0], batch[1])
+                model.backward(loss)
+                with pytest.raises(AssertionError) as assertinfo:
+                    model.step()
+                assert ("It is illegal to call Engine.step() inside no_sync context manager" in str(assertinfo))
+
+    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+    @pytest.mark.parametrize("zero_stage", [0, 1])
+    def test_multiple_ctxts(self, zero_stage, dtype):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            },
+        }
+
+        if dtype == torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+        elif dtype == torch.float16:
+            config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
+
+        hidden_dim = 64
+        total_samples = 32
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=total_samples,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=dtype)
+        dist.barrier()
+
+        param_list = list(model.parameters())
+        first_losses = []
+        first_grad_norms = []
+        with model.no_sync():
+            for _, batch in enumerate(data_loader):
+                loss = model(batch[0], batch[1])
+                first_losses.append(loss.item())
+                model.backward(loss)
+                grad_norm = sum([safe_get_full_grad(p).norm() for p in param_list])
+                first_grad_norms.append(grad_norm.item())
+
+        second_losses = []
+        second_grad_norms = []
+
+        model.zero_grad()
+        with model.no_sync():
+            for _, batch in enumerate(data_loader):
+                loss = model(batch[0], batch[1])
+                second_losses.append(loss.item())
+                model.backward(loss)
+                grad_norm = sum([safe_get_full_grad(p).norm() for p in param_list])
+                second_grad_norms.append(grad_norm.item())
+
+        assert len(first_losses) == len(second_losses)
+        for x, y in zip(first_losses, second_losses):
+            assert x == y
+
+        assert len(first_grad_norms) == len(second_grad_norms)
+        for x, y in zip(first_grad_norms, second_grad_norms):
+            assert x == y
+
+    def test_reentry(self):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3
+                }
+            },
+            "zero_optimization": {
+                "stage": 1,
+            },
+        }
+
+        hidden_dim = 64
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        dist.barrier()
+
+        with model.no_sync():
+            with pytest.raises(AssertionError) as assertinfo:
+                with model.no_sync():
+                    pass
+            assert ("no_sync context manager reentry is unsupported" in str(assertinfo))

tests/unit/runtime/test_pld.py

@@ -0,0 +1,104 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import numpy as np
+import deepspeed
+import pytest
+from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop
+
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, PLD_SimpleModel, random_dataloader
+from deepspeed.accelerator import get_accelerator
+
+
+@pytest.mark.parametrize('theta', [0, 0.1, 0.9, 1.0])
+def test_pld_schedule(tmpdir, theta):
+    gamma = 0.001
+
+    pld_scheduler = ProgressiveLayerDrop(theta, gamma)
+    for i in range(10):
+        pld_scheduler.update_state(i)
+        expected_theta = (1. - theta) * np.exp(-gamma * i) + theta
+        actual_theta = pld_scheduler.get_theta()
+        assert expected_theta == actual_theta
+
+
+@pytest.mark.parametrize('theta', [0, 0.1, 0.9, 1.0])
+class TestPLDModel(DistributedTest):
+    world_size = 1
+
+    def test_pld_model(self, theta):
+        gamma = 0.001
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.0001
+                }
+            },
+            "progressive_layer_drop": {
+                "enabled": True,
+                "theta": theta,
+                "gamma": gamma
+            }
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+        hidden_dim = 10
+
+        model = PLD_SimpleModel(hidden_dim, empty_grad=False)
+        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)
+
+        for i, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+            expected_theta = (1. - theta) * np.exp(-gamma * i) + theta
+            actual_theta = model.get_pld_theta()
+            assert expected_theta == actual_theta
+
+
+class TestNonPLDModel(DistributedTest):
+    world_size = 1
+
+    def test_non_pld_model(self):
+        gamma = 0.001
+        theta = 0.5
+        config_dict = {
+            "train_batch_size": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.0001
+                }
+            },
+            "progressive_layer_drop": {
+                "enabled": True,
+                "theta": theta,
+                "gamma": gamma
+            }
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+        data_loader = random_dataloader(model=model, total_samples=1, hidden_dim=hidden_dim, device=model.device)
+
+        for i, batch in enumerate(data_loader):
+            with pytest.raises(TypeError):
+                loss = model(batch[0], batch[1])

tests/unit/runtime/test_precision_config_dynamic_scale.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+from pydantic import ValidationError
+
+from deepspeed.runtime.precision_config import DeepSpeedFP16Config
+
+
+@pytest.mark.parametrize("field", ["loss_scale_window", "min_loss_scale"])
+@pytest.mark.parametrize("value", [0, -1])
+def test_fp16_dynamic_scale_rejects_nonpositive_when_dynamic(field, value):
+    # Dynamic loss scaling is active when fp16 is enabled and loss_scale == 0.
+    with pytest.raises(ValidationError):
+        DeepSpeedFP16Config(enabled=True, loss_scale=0, **{field: value})
+
+
+@pytest.mark.parametrize("field", ["loss_scale_window", "min_loss_scale"])
+@pytest.mark.parametrize("value", [1, 1000])
+def test_fp16_dynamic_scale_accepts_positive_when_dynamic(field, value):
+    cfg = DeepSpeedFP16Config(enabled=True, loss_scale=0, **{field: value})
+    assert getattr(cfg, field) > 0
+
+
+@pytest.mark.parametrize("field", ["loss_scale_window", "min_loss_scale"])
+@pytest.mark.parametrize("value", [0, -1])
+def test_fp16_dynamic_scale_ignored_with_static_loss_scale(field, value):
+    # With a static loss scale (loss_scale > 0) these fields are unused, so a
+    # non-positive value must not fail config construction (compatibility).
+    cfg = DeepSpeedFP16Config(enabled=True, loss_scale=128, **{field: value})
+    assert getattr(cfg, field) == value
+
+
+@pytest.mark.parametrize("field", ["loss_scale_window", "min_loss_scale"])
+@pytest.mark.parametrize("value", [0, -1])
+def test_fp16_dynamic_scale_ignored_when_fp16_disabled(field, value):
+    # When fp16 is disabled the dynamic scaling fields are unused.
+    cfg = DeepSpeedFP16Config(enabled=False, loss_scale=0, **{field: value})
+    assert getattr(cfg, field) == value
+
+
+@pytest.mark.parametrize("field", ["loss_scale_window", "min_loss_scale"])
+@pytest.mark.parametrize("value", [True, False])
+def test_fp16_dynamic_scale_rejects_bool(field, value):
+    # Pydantic coerces bool to int (True -> 1), which would otherwise slip past
+    # the positivity check. Bools must be rejected before coercion.
+    with pytest.raises(ValidationError):
+        DeepSpeedFP16Config(enabled=True, loss_scale=0, **{field: value})
+
+
+@pytest.mark.parametrize("field", ["loss_scale_window", "min_loss_scale"])
+@pytest.mark.parametrize("value", [float("inf"), float("nan"), "abc", None])
+def test_fp16_dynamic_scale_rejects_non_integer(field, value):
+    # Non-finite and non-numeric values must be rejected rather than coerced.
+    with pytest.raises(ValidationError):
+        DeepSpeedFP16Config(enabled=True, loss_scale=0, **{field: value})

tests/unit/runtime/test_precision_config_loss_scale.py

@@ -0,0 +1,31 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import math
+
+import pytest
+from pydantic import ValidationError
+
+from deepspeed.runtime.precision_config import DeepSpeedFP16Config
+
+
+@pytest.mark.parametrize("loss_scale", [-1, float("inf"), float("nan"), True])
+def test_fp16_loss_scale_rejects_invalid_values(loss_scale):
+    with pytest.raises(ValidationError):
+        DeepSpeedFP16Config(loss_scale=loss_scale)
+
+
+@pytest.mark.parametrize("loss_scale", [0, 1, 2.0, "3"])
+def test_fp16_loss_scale_accepts_valid_values(loss_scale):
+    cfg = DeepSpeedFP16Config(loss_scale=loss_scale)
+    assert math.isfinite(cfg.loss_scale)
+    assert cfg.loss_scale >= 0
+
+
+@pytest.mark.parametrize("loss_scale", [[], {}])
+def test_fp16_loss_scale_invalid_type_has_clear_error(loss_scale):
+    with pytest.raises(ValidationError) as excinfo:
+        DeepSpeedFP16Config(loss_scale=loss_scale)
+    assert "must be a number" in str(excinfo.value)

tests/unit/runtime/test_runtime_utils.py

@@ -0,0 +1,128 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from torch._utils import _flatten_dense_tensors
+import deepspeed.comm as dist
+import pytest
+from typing import Dict
+
+import deepspeed.runtime.utils as ds_utils
+import deepspeed.utils.groups as groups
+from deepspeed.accelerator import get_accelerator
+
+from unit.common import DistributedTest
+
+
+def test_call_to_str():
+    c2s = ds_utils.call_to_str
+
+    assert c2s('int') == 'int()'
+    assert c2s('int', 3) == 'int(3)'
+    assert c2s('int', 3, 'jeff') == 'int(3, \'jeff\')'
+
+    assert c2s('hello', val=3) == 'hello(val=3)'
+    assert c2s('hello', 1138, val=3) == 'hello(1138, val=3)'
+
+
+class TestClipGradNorm(DistributedTest):
+    world_size = 2
+
+    def test_gather(self):
+        param1 = torch.nn.Parameter(torch.Tensor([0]))
+        param1.grad = torch.Tensor([1])
+        param2 = torch.nn.Parameter(torch.Tensor([0]))
+        param2.grad = torch.Tensor([dist.get_rank() + 1])
+        # param2 is now MoE parameter
+        param2.allreduce = False
+
+        parameters = [param1, param2]
+
+        groups._create_expert_and_data_parallel(2)
+
+        norm = ds_utils.clip_grad_norm_(parameters, max_norm=0.1)
+        norm = torch.Tensor([norm]).to(get_accelerator().device_name(dist.get_rank()))
+        world_size = dist.get_world_size()
+        gathered_norm = [torch.zeros(1).to(get_accelerator().device_name()) for i in range(world_size)]
+
+        dist.all_gather(gathered_norm, norm)
+
+        assert gathered_norm[0] == gathered_norm[1], "norm at rank 0 does not match the norm at rank 1"
+
+    def test_clipped_val(self):
+        max_norm = 0.1
+
+        def test_params():
+            param1 = torch.nn.Parameter(torch.Tensor([0]))
+            param1.grad = torch.Tensor([1])
+            param2 = torch.nn.Parameter(torch.Tensor([0]))
+            param2.grad = torch.Tensor([1])
+            return [param1, param2]
+
+        # This assumes gradients are same on all the ranks and doesn't consider multiple ranks
+        params_expected = test_params()
+        torch.nn.utils.clip_grad_norm_(params_expected, max_norm)
+
+        params_actual = test_params()
+        ds_utils.clip_grad_norm_(params_actual, max_norm=max_norm)
+
+        # This can be allclose
+        assert torch.equal(params_expected[0].grad, params_actual[0].grad)
+        assert torch.equal(params_expected[1].grad, params_actual[1].grad)
+
+
+@pytest.mark.parametrize("check_using_norm", [(False), (True)])
+class TestCheckOverflow(DistributedTest):
+    world_size = 2
+
+    def test(self, check_using_norm):
+        groups._create_expert_and_data_parallel(2)
+
+        param1 = torch.nn.Parameter(torch.Tensor([0]))
+        param1.grad = torch.Tensor([1])
+        param2 = torch.nn.Parameter(torch.Tensor([0]))
+        if dist.get_rank() == 0:
+            param2.grad = torch.Tensor([1])
+        else:
+            param2.grad = torch.Tensor([float("inf")])
+        param2.allreduce = False
+        # param2 is now MoE parameter
+        parameters = [param1, param2]
+        if check_using_norm:
+            grads_group_flat = [_flatten_dense_tensors([p.grad for p in parameters])]
+            norm = ds_utils.get_weight_norm(grads_group_flat)
+            overflow_checker = ds_utils.CheckOverflow([parameters])
+            overflow = overflow_checker.check_using_norm([norm], reduce_overflow=False)
+        else:
+            overflow_checker = ds_utils.CheckOverflow([parameters])
+            overflow = overflow_checker.check()
+        assert overflow
+
+
+@pytest.mark.skipif(not hasattr(torch.autograd.graph, "_get_grad_fn_or_grad_acc"),
+                    reason="requires torch.autograd.graph._get_grad_fn_or_grad_acc")
+def test_count_used_parameters_enables_grad_for_grad_acc_lookup(monkeypatch):
+    """count_used_parameters_in_backward should enable grad for grad-acc lookup."""
+    param = torch.nn.Parameter(torch.tensor([1.0], requires_grad=True))
+    seen: Dict[str, int] = {"lookup_calls": 0}
+    original_getter = torch.autograd.graph._get_grad_fn_or_grad_acc
+
+    def _require_grad_enabled(t):
+        seen["lookup_calls"] += 1
+        if not torch.is_grad_enabled():
+            raise RuntimeError("grad mode must be enabled for grad-acc lookup")
+        return original_getter(t)
+
+    monkeypatch.setattr(torch.autograd.graph, "_get_grad_fn_or_grad_acc", _require_grad_enabled)
+
+    def _hook(grad):
+        seen["count"] = ds_utils.count_used_parameters_in_backward([param])
+        return grad
+
+    param.register_hook(_hook)
+    loss = (param * 2.0).sum()
+    loss.backward()
+    assert seen["lookup_calls"] > 0
+    assert "count" in seen

tests/unit/runtime/test_tp_plan_extraction.py

@@ -0,0 +1,110 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed.runtime.tensor_parallel.config import _get_hf_tp_plan
+
+
+class TestTPPlanExtraction:
+
+    def test_extract_tp_plan_from_mock_model(self):
+
+        class MockHFModel:
+
+            def __init__(self):
+                self._tp_plan = {"layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise"}
+
+        model = MockHFModel()
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is not None
+        assert "layers.*.self_attn.q_proj" in tp_plan
+        assert tp_plan["layers.*.self_attn.q_proj"] == "colwise"
+
+    def test_extract_tp_plan_from_model_with_config(self):
+
+        class MockHFConfig:
+            base_model_tp_plan = {"layers.*.self_attn.q_proj": "colwise"}
+
+        class MockHFModel:
+
+            def __init__(self, config):
+                self.config = config
+
+        config = MockHFConfig()
+        model = MockHFModel(config)
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is not None
+        assert "layers.*.self_attn.q_proj" in tp_plan
+
+    def test_no_tp_plan_model(self):
+        model = torch.nn.Linear(10, 10)
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is None
+
+    def test_empty_tp_plan(self):
+
+        class MockHFModel:
+
+            def __init__(self):
+                self._tp_plan = {}
+
+        model = MockHFModel()
+        tp_plan = _get_hf_tp_plan(model)
+
+        # Empty _tp_plan is falsy, so falls through to config then None
+        assert tp_plan is None
+
+    def test_none_tp_plan_falls_back_to_config(self):
+
+        class MockHFConfig:
+            base_model_tp_plan = {"layers.*.self_attn.q_proj": "colwise"}
+
+        class MockHFModel:
+
+            def __init__(self, config):
+                self.config = config
+                self._tp_plan = None
+
+        config = MockHFConfig()
+        model = MockHFModel(config)
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is not None
+        assert "layers.*.self_attn.q_proj" in tp_plan
+
+    def test_none_tp_plan(self):
+
+        class MockHFModel:
+
+            def __init__(self):
+                pass
+
+        model = MockHFModel()
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is None
+
+    def test_priority_config_over_model(self):
+
+        class MockHFConfig:
+            base_model_tp_plan = {"config_plan": "colwise"}
+
+        class MockHFModel:
+
+            def __init__(self, config):
+                self.config = config
+                self._tp_plan = {"model_plan": "colwise"}
+
+        config = MockHFConfig()
+        model = MockHFModel(config)
+        tp_plan = _get_hf_tp_plan(model)
+
+        assert tp_plan is not None
+        assert "config_plan" in tp_plan
+        assert "model_plan" not in tp_plan

tests/unit/runtime/utils/test_partition.py

@@ -0,0 +1,195 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+import torch
+import deepspeed.comm as dist
+
+from deepspeed.runtime.utils import partition_uniform
+from deepspeed.runtime.utils import partition_balanced
+from deepspeed.runtime.utils import prefix_sum_inc
+from deepspeed.runtime.utils import PartitionedTensor
+from deepspeed.accelerator import get_accelerator
+
+from unit.common import DistributedTest
+
+
+class TestPartitionedTensor(DistributedTest):
+    world_size = 4
+
+    def test(self):
+        world = dist.get_world_size()
+
+        group = dist.new_group(ranks=list(range(world)))
+
+        rows = world * 4
+        cols = 3
+
+        full = torch.rand(rows, cols).to(get_accelerator().device_name())
+        dist.broadcast(full, src=0, group=group)
+        part = PartitionedTensor(full, group=group)
+
+        assert len(part.local_size()) == 1
+        assert part.local_size()[0] * world == full.numel()
+
+        reconstructed = part.full()
+        assert torch.equal(full, reconstructed)
+
+
+class TestPartitionedTensorUnEven(DistributedTest):
+    world_size = 4
+
+    def test(self):
+        world = dist.get_world_size()
+
+        group = dist.new_group(ranks=list(range(world)))
+
+        rows = world * 4 - 1
+        cols = world + 1
+
+        full = torch.rand(rows, cols).to(get_accelerator().device_name())
+        dist.broadcast(full, src=0, group=group)
+        part = PartitionedTensor(full, group=group)
+
+        assert len(part.local_size()) == 1
+
+        reconstructed = part.full()
+        assert torch.equal(full, reconstructed)
+
+
+class TestPartitionedTensorMeta(DistributedTest):
+    world_size = 4
+
+    def test(self):
+        world = dist.get_world_size()
+
+        group = dist.new_group(ranks=list(range(world)))
+
+        rows = world * 7
+        cols = 3
+
+        full = torch.rand(rows, cols).to(get_accelerator().device_name())
+        dist.broadcast(full, src=0, group=group)
+        part = PartitionedTensor(full, group=group)
+
+        my_meta = PartitionedTensor.from_meta(part.to_meta(), part.local_data, group)
+        assert torch.equal(full, my_meta.full())
+
+
+def assert_valid_partition(weights, parts, P):
+    N = len(weights)
+    assert len(parts) == P + 1
+    assert parts[0] == 0
+    assert parts[P] == N
+    for idx in range(P):
+        assert parts[idx] <= parts[idx + 1]
+
+
+def get_partition_weights(weights, parts):
+    """ Return the amount of weight in each partition. """
+    costs = [0] * (len(parts) - 1)
+    P = len(parts) - 1
+    for p in range(P):
+        start = parts[p]
+        stop = parts[p + 1]
+        costs[p] = sum(weights[start:stop])
+    return costs
+
+
+def test_prefix_sum():
+    x = [3, 4, 5]
+    psum = prefix_sum_inc(x)
+    assert psum == [3, 7, 12]
+
+
+def test_valid_partition():
+    N = 10
+    P = 1
+    weights = [1] * N
+    parts = partition_balanced(weights, P)
+    assert_valid_partition(weights, parts, P)
+
+
+def test_short_partition_uniform():
+    N = 2
+    P = 4
+    weights = [1] * N
+    parts = partition_uniform(len(weights), P)
+    assert_valid_partition(weights, parts, P)
+
+
+def test_short_partition():
+    N = 2
+    P = 4
+    weights = [1] * N
+    parts = partition_balanced(weights, P)
+    assert_valid_partition(weights, parts, P)
+
+
+def test_easy_balance_uniform():
+    weights = [1] * 8
+    P = 4
+    parts = partition_uniform(len(weights), P)
+    assert_valid_partition(weights, parts, P)
+    costs = get_partition_weights(weights, parts)
+    assert all(c == 2 for c in costs)
+
+
+def test_easy_balance_balanced():
+    weights = [1] * 8
+    P = 4
+    parts = partition_balanced(weights, P)
+    assert_valid_partition(weights, parts, P)
+    costs = get_partition_weights(weights, parts)
+    assert all(c == 2 for c in costs), costs
+
+
+def test_int_balanced():
+    weights = [0, 1, 2, 3, 3, 3]
+    P = 4
+    parts = partition_balanced(weights, P)
+    assert parts == [0, 3, 4, 5, 6]
+
+    assert_valid_partition(weights, parts, P)
+    costs = get_partition_weights(weights, parts)
+    assert all(c == 3 for c in costs)
+
+
+def test_float_balanced():
+    weights = [0., 1.1, 1.9, 3., 3., 3.]
+    P = 4
+    parts = partition_balanced(weights, P)
+    assert_valid_partition(weights, parts, P)
+    assert parts == [0, 3, 4, 5, 6]
+
+
+@pytest.mark.skip(reason="Variance-minimizing partitioning returns different result.")
+def test_float_lastheavy():
+    weights = [0., 1.1, 1.9, 3., 30.]
+    P = 2
+    parts = partition_balanced(weights, P)
+    assert_valid_partition(weights, parts, P)
+    assert parts == [0, 4, 5]
+
+
+def test_float_midheavy():
+    weights = [0., 1.1, 30, 3.]
+    P = 3
+    parts = partition_balanced(weights, P)
+    assert_valid_partition(weights, parts, P)
+    assert parts == [0, 2, 3, 4]
+
+
+def test_balance_bert():
+    # Parameters per layer for a transformer model with 24 transformers and hidden dim 1024
+    weights = [
+        52559872, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224,
+        12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224, 12596224,
+        12596224, 12596224, 12596224, 0, 52559872
+    ]
+    P = 8
+    parts = partition_balanced(weights, P)
+    assert_valid_partition(weights, parts, P)

tests/unit/runtime/zenflow/test_zf.py

@@ -0,0 +1,136 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import deepspeed.comm as dist
+from deepspeed.accelerator import get_accelerator
+
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataloader
+import deepspeed
+
+
+class BaseZenFlowTest:
+    hidden_dim = 10
+    batch_size = 4
+    grad_acc_steps = 1
+
+    def get_config_dict(self, stage, offload_selective_optimizer, select_strategy, select_interval, update_interval,
+                        full_warm_up_rounds):
+        config = {
+            "train_batch_size": self.batch_size,
+            "gradient_accumulation_steps": self.grad_acc_steps,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-4
+                }
+            },
+            "zero_optimization": {
+                "stage": stage,
+                "offload_optimizer": {
+                    "device": "cpu"
+                },
+                "overlap_comm": True,
+                "zenflow": {
+                    "topk_ratio": 0.2,
+                    "select_strategy": select_strategy,
+                    "select_interval": select_interval,
+                    "update_interval": update_interval,
+                    "overlap_step": False,
+                    "offload": offload_selective_optimizer,
+                    "auto_ratio": 0.99,
+                    "full_warm_up_rounds": full_warm_up_rounds,
+                }
+            },
+            "zero_allow_untested_optimizer": True,
+        }
+
+        if get_accelerator().is_bf16_supported():
+            config["bf16"] = {"enabled": True}
+        return config
+
+    def run_training_distributed(self, config_dict):
+
+        if get_accelerator().device_name() == "cpu":
+            return
+
+        model = SimpleModel(self.hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        train_dataloader = random_dataloader(model=model,
+                                             total_samples=20,
+                                             hidden_dim=self.hidden_dim,
+                                             device=model.device)
+
+        dist.barrier()
+
+        for step, batch in enumerate(train_dataloader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+        model.destroy()
+
+
+@pytest.mark.parametrize("stage", [1, 2, 3])
+@pytest.mark.parametrize("full_warm_up_rounds", [0, 3])
+@pytest.mark.parametrize("offload_selective_optimizer", [True, False])
+@pytest.mark.parametrize("select_strategy,select_interval,update_interval", [
+    ("auto", "auto", "auto"),
+    ("step", 10, 3),
+    ("epoch", 1, 4),
+])
+class TestZenFlowSingleGPU(DistributedTest, BaseZenFlowTest):
+    world_size = 1
+
+    def test_zenflow_single_gpu(self, stage, offload_selective_optimizer, select_strategy, select_interval,
+                                update_interval, full_warm_up_rounds):
+        tester = BaseZenFlowTest()
+        config_dict = tester.get_config_dict(stage, offload_selective_optimizer, select_strategy, select_interval,
+                                             update_interval, full_warm_up_rounds)
+        tester.run_training_distributed(config_dict)
+
+
+@pytest.mark.parametrize("stage", [1, 2, 3])
+@pytest.mark.parametrize("full_warm_up_rounds", [0, 3])
+@pytest.mark.parametrize("offload_selective_optimizer", [True, False])
+@pytest.mark.parametrize("select_strategy,select_interval,update_interval", [
+    ("auto", "auto", "auto"),
+    ("step", 10, 3),
+    ("epoch", 1, 4),
+])
+class TestZenFlowDistributed(DistributedTest, BaseZenFlowTest):
+    world_size = 2
+
+    def test_zenflow_distributed(self, stage, offload_selective_optimizer, select_strategy, select_interval,
+                                 update_interval, full_warm_up_rounds):
+        config_dict = self.get_config_dict(stage, offload_selective_optimizer, select_strategy, select_interval,
+                                           update_interval, full_warm_up_rounds)
+        self.run_training_distributed(config_dict)
+
+
+@pytest.mark.parametrize(
+    "cores,perc,expected_zf,expected_pt",
+    [
+        # Normal split: ceil(0.25 * 8) = 2 cores reserved for training.
+        ([0, 1, 2, 3, 4, 5, 6, 7], 0.25, [2, 3, 4, 5, 6, 7], [0, 1]),
+        # Rounds up: ceil(0.1 * 8) = 1.
+        ([0, 1, 2, 3, 4, 5, 6, 7], 0.1, [1, 2, 3, 4, 5, 6, 7], [0]),
+        # Two cores, half each.
+        ([10, 11], 0.5, [11], [10]),
+        # Reserve rounds to 0 -> both sides share the full set.
+        ([0, 1, 2, 3], 0.0, [0, 1, 2, 3], [0, 1, 2, 3]),
+        # Reserve rounds to every core -> both sides share the full set.
+        ([0, 1, 2, 3], 1.0, [0, 1, 2, 3], [0, 1, 2, 3]),
+    ])
+def test_split_affinity(cores, perc, expected_zf, expected_pt):
+    from deepspeed.runtime.zenflow.zenflow_utils import _split_affinity
+    zf, pt = _split_affinity(cores, perc)
+    assert zf == expected_zf
+    assert pt == expected_pt
+    # When the sides are actually isolated they must partition the cores exactly.
+    if zf != pt:
+        assert sorted(zf + pt) == sorted(cores)
+        assert not (set(zf) & set(pt))

tests/unit/runtime/zenflow/test_zf_config.py

@@ -0,0 +1,86 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+from pydantic import ValidationError
+
+from deepspeed.runtime.zero.config import DeepSpeedZeroConfig, ZeroStageEnum
+from deepspeed.runtime.zenflow.zenflow_config import ZenFlowConfig
+from deepspeed.runtime.zero.offload_config import DeepSpeedZeroOffloadOptimizerConfig
+
+
+def test_stage_enum_accepts_int_and_enum():
+    """`stage` can be passed as either an int or the ZeroStageEnum."""
+    c1 = DeepSpeedZeroConfig(stage=2)
+    assert c1.stage == ZeroStageEnum.gradients
+    c2 = DeepSpeedZeroConfig(stage=ZeroStageEnum.weights)
+    assert c2.stage == ZeroStageEnum.weights
+
+
+def test_offload_optimizer_config_from_dict():
+    """A dict for offload_optimizer should be coerced into DeepSpeedZeroOffloadOptimizerConfig."""
+    cfg = DeepSpeedZeroConfig(offload_optimizer={"device": "cpu", "pin_memory": True})
+    assert isinstance(cfg.offload_optimizer, DeepSpeedZeroOffloadOptimizerConfig)
+    assert cfg.offload_optimizer.device == "cpu"
+    assert cfg.offload_optimizer.pin_memory is True
+
+
+def test_invalid_offload_optimizer_type_raises():
+    """Passing a non-dict to offload_optimizer must error out."""
+    with pytest.raises(ValidationError):
+        DeepSpeedZeroConfig(offload_optimizer="not a dict")
+
+
+def test_zenflow_config_from_dict():
+    """A dict for zenflow should be coerced into ZenFlowConfig."""
+    zenflow_payload = {
+        "topk_ratio": 0.25,
+        "select_strategy": "auto",
+        "select_interval": 4,
+        "update_interval": 8,
+        "full_warm_up_rounds": 1,
+        "overlap_step": True
+    }
+    cfg = DeepSpeedZeroConfig(zenflow=zenflow_payload)
+    assert isinstance(cfg.zenflow, ZenFlowConfig)
+    assert cfg.zenflow.topk_ratio == 0.25
+    assert cfg.zenflow.select_strategy == "auto"
+    assert cfg.zenflow.select_interval == 4
+    assert cfg.zenflow.update_interval == 8
+    assert cfg.zenflow.full_warm_up_rounds == 1
+    assert cfg.zenflow.overlap_step is True
+
+
+def test_invalid_zenflow_type_raises():
+    """Passing a non-dict to zenflow must error out."""
+    with pytest.raises(ValidationError):
+        DeepSpeedZeroConfig(zenflow=123)
+
+
+def test_offload_and_zenflow_combined():
+    """
+    offload_optimizer and zenflow can be used together under stage 2
+    without validation errors.
+    """
+    payload = {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": True
+        },
+        "zenflow": {
+            "topk_ratio": 0.3,
+            "select_strategy": "epoch",
+            "select_interval": 3,
+            "update_interval": 6,
+            "full_warm_up_rounds": 0,
+            "overlap_step": False
+        }
+    }
+    cfg = DeepSpeedZeroConfig(**payload)
+    assert isinstance(cfg.offload_optimizer, DeepSpeedZeroOffloadOptimizerConfig)
+    assert cfg.offload_optimizer.device == "cpu"
+    assert isinstance(cfg.zenflow, ZenFlowConfig)
+    assert cfg.zenflow.select_strategy == "epoch"

tests/unit/runtime/zero/test_ignore_unused_parameters.py

@@ -0,0 +1,63 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+from unit.common import DistributedTest
+from unit.simple_model import UnusedParametersModel, random_dataloader
+from deepspeed.ops.op_builder import CPUAdamBuilder
+
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+
+
+@pytest.mark.parametrize('ignore_unused_parameters', [False, True])
+class TestStage2IgnoreUnusedParameters(DistributedTest):
+    world_size = 1
+
+    def test(self, ignore_unused_parameters):
+        use_cpu_offload = True
+
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 2,
+            "gradient_accumulation_steps": 2,
+            "steps_per_print": 1,
+            "zero_optimization": {
+                "stage": 2,
+                "cpu_offload": use_cpu_offload,
+                "ignore_unused_parameters": ignore_unused_parameters
+            },
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3
+                }
+            },
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
+        hidden_dim = 4
+
+        model = UnusedParametersModel(hidden_dim=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)
+
+        def _loop():
+            for n, batch in enumerate(data_loader):
+                loss = model(batch[0], batch[1])
+                model.backward(loss)
+                model.step()
+
+        if ignore_unused_parameters:
+            _loop()
+        else:
+            with pytest.raises(AssertionError) as e:
+                _loop()
+            assert e.value.args and 'ignore_unused_parameters' in e.value.args[0]

tests/unit/runtime/zero/test_nvme_checkpointing.py

@@ -0,0 +1,147 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import pytest
+import deepspeed.comm as dist
+import torch
+
+from unit.common import DistributedTest
+from unit.simple_model import random_dataloader, SimpleModel
+
+import deepspeed
+from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
+from deepspeed.runtime.zero.partition_parameters import Init
+from deepspeed.ops.aio import AsyncIOBuilder
+from deepspeed.accelerator import get_accelerator
+
+
+@pytest.mark.sequential
+class TestNVMeCheckpointing(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.parametrize('param_offload_device, optim_offload_device',
+                             [(OffloadDeviceEnum.none, OffloadDeviceEnum.nvme),
+                              (OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme),
+                              (OffloadDeviceEnum.nvme, OffloadDeviceEnum.none),
+                              (OffloadDeviceEnum.nvme, OffloadDeviceEnum.cpu),
+                              (OffloadDeviceEnum.nvme, OffloadDeviceEnum.nvme)])
+    def test_nvme_checkpointing(self, tmpdir, param_offload_device, optim_offload_device):
+        zero_dir, ckpt_dir = os.path.join(tmpdir, "zero"), os.path.join(tmpdir, "checkpoint")
+
+        first_stage_steps, second_stage_steps = 2, 2
+
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+
+        if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
+            pytest.skip('Skip tests since async-io is not compatible')
+
+        torch.manual_seed(123)
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015,
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 8
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "offload_param": {
+                    "device": param_offload_device,
+                    "nvme_path": str(zero_dir)
+                },
+                "offload_optimizer": {
+                    "device": optim_offload_device,
+                    "nvme_path": str(zero_dir)
+                },
+                "sub_group_size": 100,
+                "stage3_max_live_parameters": 100,
+                "stage3_param_persistence_threshold": 0,
+            },
+            "aio": {
+                "block_size": 1048576  # Minimum AIO bytes, anything smaller than this will not be offloaded
+            }
+        }
+
+        hidden_dim, nlayers = 2048, 2
+        with deepspeed.zero.Init(config_dict_or_path=config_dict):
+            model = SimpleModel(hidden_dim, nlayers=nlayers, empty_grad=False)
+
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        model.empty_partition_cache()
+
+        assert first_stage_steps > 0
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=first_stage_steps,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        dist.barrier()
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        dist.barrier()
+        model.save_checkpoint(ckpt_dir)
+
+        if second_stage_steps > 0:
+            second_stage_batches = list(
+                random_dataloader(model=model,
+                                  total_samples=second_stage_steps,
+                                  hidden_dim=hidden_dim,
+                                  device=model.device,
+                                  dtype=torch.float16))
+            dist.barrier()
+            for n, batch in enumerate(second_stage_batches):
+                loss = model(batch[0], batch[1])
+                model.backward(loss)
+                model.step()
+            dist.barrier()
+
+        final_batch = next(
+            iter(
+                random_dataloader(model=model,
+                                  total_samples=1,
+                                  hidden_dim=hidden_dim,
+                                  device=model.device,
+                                  dtype=torch.float16)))
+        dist.barrier()
+        loss_before = float(model(final_batch[0], final_batch[1]))
+
+        # Needed in ZeRO 3. Not doing so can give memory leak
+        model.destroy()
+
+        # TODO: This should be on the engine? There needs to be a better way.
+        Init.param_id = 0
+
+        with deepspeed.zero.Init(config_dict_or_path=config_dict):
+            model = SimpleModel(hidden_dim, nlayers=nlayers, empty_grad=False)
+
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        model.load_checkpoint(ckpt_dir)
+
+        if second_stage_steps > 0:
+            dist.barrier()
+            for n, batch in enumerate(second_stage_batches):
+                loss = model(batch[0], batch[1])
+                model.backward(loss)
+                model.step()
+            dist.barrier()
+
+        dist.barrier()
+        loss_after = float(model(final_batch[0], final_batch[1]))
+
+        assert loss_before == loss_after

tests/unit/runtime/zero/test_unwrap_model.py

@@ -0,0 +1,93 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+import deepspeed
+from deepspeed.runtime.zero import unwrap_model_for_generation
+from deepspeed.accelerator import get_accelerator
+
+from unit.common import DistributedTest, preferred_dtype
+from unit.simple_model import SimpleModel
+
+config = {
+    "train_batch_size": 2,
+    "steps_per_print": 1,
+    "optimizer": {
+        "type": "Adam",
+        "params": {
+            "lr": 0.00015
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_param_persistence_threshold": 1,
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": True
+        }
+    }
+}
+
+if get_accelerator().is_bf16_supported():
+    config["bf16"] = {"enabled": True}
+elif get_accelerator().is_fp16_supported():
+    config["fp16"] = {"enabled": True, "loss_scale": 138.}
+
+
+class TestUnwrapModel(DistributedTest):
+    # gather across more than 1 gpu
+    world_size = 2
+
+    def test(self):
+
+        def hooks_exist(engine):
+            if engine.optimizer is not None and hasattr(engine.optimizer, "parameter_offload"):
+                optimizer_offload = engine.optimizer.parameter_offload
+            elif engine.optimizer is not None:
+                optimizer_offload = engine.optimizer
+
+            hooks = 0
+            for hook in optimizer_offload.forward_hooks:
+                hooks += 1
+            if hooks > 0:
+                return True
+            return False
+
+        model = SimpleModel(hidden_dim=100)
+        engine, _, _, _ = deepspeed.initialize(args=None, model=model, config=config)
+
+        with unwrap_model_for_generation(engine):
+            # assert no hooks
+            assert not hooks_exist(engine)
+            # assert parameters gathered
+            assert model.linears[0].weight.numel() != 0, "GatheredParameters should give a non-0-sized tensor"
+
+        # assert hooks
+        assert hooks_exist(engine)
+
+
+class TestUnwrapModelTraceInvalidate(DistributedTest):
+    # unwrap_model_for_generation re-registers the ZeRO-3 hooks; without trace
+    # invalidation the next training step pops an empty fetch deque.
+    world_size = 2
+
+    def test(self):
+        model = SimpleModel(hidden_dim=100)
+        engine, _, _, _ = deepspeed.initialize(args=None, model=model, config=config)
+
+        x = torch.randn(2, 100, device=engine.device, dtype=preferred_dtype())
+        y = torch.empty(2, dtype=torch.long, device=engine.device).random_(100)
+
+        loss = engine(x, y)
+        engine.backward(loss)
+        engine.step()
+
+        with unwrap_model_for_generation(engine):
+            pass
+
+        loss = engine(x, y)
+        engine.backward(loss)
+        engine.step()

tests/unit/runtime/zero/test_zero_autocast_mixed_dtype.py

@@ -0,0 +1,129 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.torch_autocast import get_comm_dtype, has_comm_dtype
+from deepspeed.runtime.zero.partition_parameters import get_allgather_dtype
+from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad
+from unit.common import DistributedTest
+from unit.util import bf16_required_version_check
+
+
+def _safe_module_name():
+    return f"{MixedDtypeAdapterModule.__module__}.{MixedDtypeAdapterModule.__name__}"
+
+
+def _zero3_bf16_autocast_config():
+    return {
+        "train_micro_batch_size_per_gpu": 1,
+        "steps_per_print": 1,
+        "bf16": {
+            "enabled": True,
+            "bf16_master_weights_and_grads": True,
+            "bf16_optimizer_states": True,
+        },
+        "zero_optimization": {
+            "stage": 3,
+            "stage3_param_persistence_threshold": 0,
+            "stage3_module_granularity_threshold": 0,
+            "stage3_use_all_reduce_for_fetch_params": False,
+        },
+        "torch_autocast": {
+            "enabled": True,
+            "dtype": str(torch.bfloat16),
+            "lower_precision_safe_modules": [_safe_module_name()],
+        },
+    }
+
+
+class MixedDtypeAdapterModule(torch.nn.Module):
+
+    def __init__(self, hidden_dim):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.base_weight = torch.nn.Parameter(torch.randn(hidden_dim, hidden_dim) * 0.01, requires_grad=False)
+
+    def attach_fp32_adapter(self, rank):
+        device = get_accelerator().current_device_name()
+        self.adapter_a = torch.nn.Parameter(
+            torch.randn(rank, self.hidden_dim, device=device, dtype=torch.float32) * 0.01)
+        self.adapter_b = torch.nn.Parameter(
+            torch.randn(self.hidden_dim, rank, device=device, dtype=torch.float32) * 0.01)
+
+        assert hasattr(self.base_weight, "convert_to_zero_parameters")
+        self.base_weight.convert_to_zero_parameters([self.adapter_a, self.adapter_b])
+
+    def forward(self, x, target):
+        base = torch.nn.functional.linear(x, self.base_weight)
+        adapter_hidden = torch.nn.functional.linear(x, self.adapter_a)
+        adapter = torch.nn.functional.linear(adapter_hidden, self.adapter_b) / self.adapter_a.shape[0]
+        output = base + adapter
+        return torch.nn.functional.mse_loss(output.float(), target.float())
+
+
+def _assert_mixed_partition_dtypes(model):
+    assert model.base_weight.dtype == torch.bfloat16
+    assert model.base_weight.ds_tensor.dtype == torch.bfloat16
+
+    for adapter_param in [model.adapter_a, model.adapter_b]:
+        assert adapter_param.dtype == torch.float32
+        assert adapter_param.ds_tensor.dtype == torch.float32
+
+
+def _assert_autocast_comm_dtype(model):
+    for param in [model.base_weight, model.adapter_a, model.adapter_b]:
+        assert has_comm_dtype(param)
+        assert get_comm_dtype(param) == torch.bfloat16
+        assert get_allgather_dtype(param, param.ds_tensor) == torch.bfloat16
+
+
+class TestZero3AutocastMixedDtype(DistributedTest):
+    world_size = 2
+
+    def test_fp32_adapter_with_bf16_base_params(self):
+        if not bf16_required_version_check():
+            pytest.skip("BF16 ZeRO-3 autocast test requires BF16 accelerator support.")
+
+        hidden_dim = 8
+        config = _zero3_bf16_autocast_config()
+
+        with deepspeed.zero.Init(config_dict_or_path=config):
+            model = MixedDtypeAdapterModule(hidden_dim)
+        model.attach_fp32_adapter(rank=4)
+        _assert_mixed_partition_dtypes(model)
+
+        trainable_params = [p for p in model.parameters() if p.requires_grad]
+        optimizer = torch.optim.AdamW(trainable_params, lr=0.1)
+        engine, _, _, _ = deepspeed.initialize(config=config,
+                                               model=model,
+                                               model_parameters=trainable_params,
+                                               optimizer=optimizer)
+        try:
+            _assert_mixed_partition_dtypes(engine.module)
+            _assert_autocast_comm_dtype(engine.module)
+
+            adapter_a_before = safe_get_full_fp32_param(engine.module.adapter_a).detach().clone()
+            device = engine.device
+            x = torch.randn(2, hidden_dim, device=device, dtype=torch.float32)
+            target = torch.randn(2, hidden_dim, device=device, dtype=torch.float32)
+
+            loss = engine(x, target)
+            engine.backward(loss)
+
+            adapter_a_grad = safe_get_full_grad(engine.module.adapter_a)
+            assert adapter_a_grad is not None
+            assert torch.count_nonzero(adapter_a_grad).item() > 0
+
+            engine.step()
+
+            adapter_a_after = safe_get_full_fp32_param(engine.module.adapter_a)
+            assert not torch.equal(adapter_a_before, adapter_a_after)
+            _assert_mixed_partition_dtypes(engine.module)
+        finally:
+            engine.destroy()

tests/unit/runtime/zero/test_zero_config.py

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.runtime.zero.config import DeepSpeedZeroConfig, DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig
+
+
+def test_zero_config_deprecatedfields():
+    config = DeepSpeedZeroConfig(**{"cpu_offload_param": True})
+    assert isinstance(config.offload_param, DeepSpeedZeroOffloadParamConfig)
+
+    config = DeepSpeedZeroConfig(**{"cpu_offload": True})
+    assert isinstance(config.offload_optimizer, DeepSpeedZeroOffloadOptimizerConfig)
+
+    config = DeepSpeedZeroConfig(**{"stage3_gather_fp16_weights_on_model_save": True})
+    assert config.gather_16bit_weights_on_model_save == True
+
+
+def test_zero_config_aliasfields():
+    config = DeepSpeedZeroConfig(**{"stage3_prefetch_bucket_size": 12345})
+    assert config.prefetch_bucket_size == 12345
+
+    config = DeepSpeedZeroConfig(**{"stage3_param_persistence_threshold": 12345})
+    assert config.param_persistence_threshold == 12345
+
+    config = DeepSpeedZeroConfig(**{"stage3_max_reuse_distance": 12345})
+    assert config.max_reuse_distance == 12345
+
+    config = DeepSpeedZeroConfig(**{"stage3_gather_16bit_weights_on_model_save": True})
+    assert config.gather_16bit_weights_on_model_save == True
+
+
+def test_zero_config_pipeline_loading_checkpoint():
+    for stage in [0, 1, 2]:
+        config = DeepSpeedZeroConfig(**{"stage": stage})
+        assert config.pipeline_loading_checkpoint == False
+
+
+def test_zero_config_overlapcomm():
+    for stage in [0, 1, 2]:
+        config = DeepSpeedZeroConfig(**{"stage": stage})
+        assert config.overlap_comm == False
+
+    config = DeepSpeedZeroConfig(**{"stage": 3})
+    assert config.overlap_comm == True
+
+
+def test_zero_config_offload_configs():
+    config = DeepSpeedZeroConfig()
+    assert config.offload_param is None
+    assert config.offload_optimizer is None
+
+    config = DeepSpeedZeroConfig(**{"offload_param": None, "offload_optimizer": None})
+    assert config.offload_param is None
+    assert config.offload_optimizer is None
+
+    config = DeepSpeedZeroConfig(**{"offload_param": {}, "offload_optimizer": {}})
+    assert isinstance(config.offload_param, DeepSpeedZeroOffloadParamConfig)
+    assert isinstance(config.offload_optimizer, DeepSpeedZeroOffloadOptimizerConfig)
+
+
+def test_zero_offload_optimizer_config_pipeline():
+    config = DeepSpeedZeroOffloadOptimizerConfig()
+    assert config.pipeline == False
+
+    config = DeepSpeedZeroOffloadOptimizerConfig(**{"pipeline_read": True, "pipeline_write": False})
+    assert config.pipeline == True
+
+    config = DeepSpeedZeroOffloadOptimizerConfig(**{"pipeline_read": False, "pipeline_write": True})
+    assert config.pipeline == True
+
+    config = DeepSpeedZeroOffloadOptimizerConfig(**{"pipeline_read": True, "pipeline_write": True})
+    assert config.pipeline == True

tests/unit/runtime/zero/test_zero_context.py

@@ -0,0 +1,374 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from types import SimpleNamespace
+
+import torch
+import pytest
+import deepspeed
+import deepspeed.comm as dist
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.zero.partition_parameters import (MultipleAllGatherHandles, ZeroParamStatus,
+                                                         partitioned_param_data_shape)
+
+from unit.common import DistributedTest, preferred_dtype, reduce_boolean_flags
+from unit.simple_model import SimpleModel
+from utils import setup_serial_env
+
+
+# Test that no sub-class or super-class is missed
+class ConvX(torch.nn.Conv1d):
+
+    def __init__(self, *args):
+        super().__init__(*args)
+        # This would not be partitioned before bugfix 5ca8167
+        self.param_in = torch.nn.Parameter(torch.FloatTensor(5).uniform_())
+
+    def forward(self, x):
+        return x
+
+
+class ConvNet(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv1 = ConvX(1, 3, 4)
+        self.param = torch.nn.Parameter(torch.FloatTensor(5).uniform_())
+
+    def forward(self, x):
+        return x
+
+
+config = {
+    "train_batch_size": 1,
+    "steps_per_print": 1,
+    "optimizer": {
+        "type": "Adam",
+        "params": {
+            "lr": 0.00015
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_param_persistence_threshold": 1,
+    }
+}
+
+if get_accelerator().is_bf16_supported():
+    config["bf16"] = {"enabled": True}
+elif get_accelerator().is_fp16_supported():
+    config["fp16"] = {"enabled": True, "loss_scale": 138.}
+
+
+def test_multiple_all_gather_handles_wait_passes_dependency_by_keyword():
+
+    class PositionalWaitHandle:
+
+        def __init__(self):
+            self.handle_dependency = None
+
+        def wait(self, handle_dependency=True):
+            self.handle_dependency = handle_dependency
+
+    class KeywordOnlyWaitHandle:
+
+        def __init__(self):
+            self.handle_dependency = None
+
+        def wait(self, *, handle_dependency=True):
+            self.handle_dependency = handle_dependency
+
+    class KwargsWaitHandle:
+
+        def __init__(self):
+            self.kwargs = None
+
+        def wait(self, **kwargs):
+            self.kwargs = kwargs
+
+    handles = [PositionalWaitHandle(), KeywordOnlyWaitHandle(), KwargsWaitHandle()]
+
+    MultipleAllGatherHandles(handles).wait(handle_dependency=False)
+
+    assert handles[0].handle_dependency is False
+    assert handles[1].handle_dependency is False
+    assert handles[2].kwargs == {"handle_dependency": False}
+
+
+class TestZeroGatheredParametersFree(DistributedTest):
+    world_size = 1
+
+    def test(self):
+        config_dict = {"train_batch_size": 1, "zero_optimization": {"stage": 3}}
+        hidden_dim = 10
+
+        class MyModel(torch.nn.Module):
+
+            def __init__(self, hidden_dim):
+                super(MyModel, self).__init__()
+                self.l1 = torch.nn.Linear(hidden_dim, hidden_dim)
+
+        with deepspeed.zero.Init(config_dict_or_path=config_dict):
+            model = MyModel(hidden_dim)
+
+        with deepspeed.zero.GatheredParameters(list(model.parameters())):
+            assert model.l1.weight.numel() != 0, "GatheredParameters should give a non-0-sized tensor"
+
+        # on exit from `GatheredParameters` the gathered params should be freed and not leak memory
+        assert model.l1.weight.numel() == 0, "outside of GatheredParameters the param should go back to be 0-sized"
+
+
+class TestMiCSGatheredParametersFree(DistributedTest):
+    world_size = 1
+
+    def test(self):
+        config_dict = {"train_batch_size": 1, "zero_optimization": {"stage": 3, "mics_shard_size": 1}}
+        hidden_dim = 10
+
+        class MyModel(torch.nn.Module):
+
+            def __init__(self, hidden_dim):
+                super(MyModel, self).__init__()
+                self.l1 = torch.nn.Linear(hidden_dim, hidden_dim)
+
+        with deepspeed.zero.MiCS_Init(config_dict_or_path=config_dict):
+            model = MyModel(hidden_dim)
+
+        with deepspeed.zero.GatheredParameters(list(model.parameters())):
+            assert model.l1.weight.numel() != 0, "GatheredParameters should give a non-0-sized tensor"
+
+        # on exit from `GatheredParameters` the gathered params should be freed and not leak memory
+        assert model.l1.weight.numel() == 0, "outside of GatheredParameters the param should go back to be 0-sized"
+
+
+class TestGatheredParametersAllRanksErrorOnModification(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "zero_optimization": {
+                "stage": 3,
+                "enable_sanity_checks": True
+            }
+        }
+        hidden_dim = 10
+
+        class MyModel(torch.nn.Module):
+
+            def __init__(self, hidden_dim):
+                super(MyModel, self).__init__()
+                self.l1 = torch.nn.Linear(hidden_dim, hidden_dim)
+                self.l2 = torch.nn.Linear(hidden_dim, hidden_dim)
+
+        with deepspeed.zero.Init(config_dict_or_path=config_dict):
+            model = MyModel(hidden_dim)
+
+        error_local = False
+        try:
+            with deepspeed.zero.GatheredParameters([model.l1.weight, model.l2.weight], modifier_rank=None):
+                with torch.no_grad():
+                    model.l1.weight.add_(0.0)
+        except RuntimeError as exc:
+            if "in-place modification" in str(exc):
+                error_local = True
+
+        error_global = reduce_boolean_flags(error_local, all)
+        if not error_global:
+            raise AssertionError("Expected in-place modification error on all ranks.")
+
+
+class TestSerialContext(DistributedTest):
+    world_size = 1
+    init_distributed = False
+    set_dist_env = False
+
+    def test_subclass_param(self):
+        setup_serial_env()
+        with deepspeed.zero.Init(config=config):
+            model = ConvNet()
+
+        assert model.param.ds_status == ZeroParamStatus.NOT_AVAILABLE
+        assert model.conv1.param_in.ds_status == ZeroParamStatus.NOT_AVAILABLE
+
+    def test_scattered_init_dist(self):
+        setup_serial_env()
+        assert not dist.is_initialized()
+        with deepspeed.zero.Init():
+            assert dist.is_initialized()
+
+    def test_scatter_halftype(self):
+        if not get_accelerator().is_fp16_supported():
+            pytest.skip("fp16 is not supported")
+        setup_serial_env()
+
+        with deepspeed.zero.Init():
+            l = torch.nn.Linear(10, 10)
+            assert l.weight.ds_tensor.dtype == torch.float16
+
+            y = torch.LongTensor([3, 3])
+            assert y.dtype == torch.long
+
+    def test_throughput_calculation(self):
+        setup_serial_env()
+
+        train_micro_batch_size_per_gpu = 7
+        gradient_accumulation_steps = 6
+        config_dict = {
+            "train_micro_batch_size_per_gpu": train_micro_batch_size_per_gpu,
+            "gradient_accumulation_steps": gradient_accumulation_steps,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.001,
+                }
+            },
+            "zero_optimization": {
+                "stage": 0
+            },
+        }
+
+        args = SimpleNamespace(local_rank=0)
+        net = SimpleModel(hidden_dim=4)
+        engine, _, _, _ = deepspeed.initialize(args=args,
+                                               config=config_dict,
+                                               model=net,
+                                               model_parameters=net.parameters())
+        assert engine.tput_timer.batch_size == train_micro_batch_size_per_gpu * gradient_accumulation_steps
+
+        assert not engine.tput_timer.initialized
+        assert not engine.tput_timer.started
+        assert engine.tput_timer.start_step == 2
+        assert engine.tput_timer.start_time == 0
+        assert engine.tput_timer.micro_step_count == 0
+        assert engine.tput_timer.global_step_count == 0
+        assert engine.tput_timer.total_elapsed_time == 0
+
+        # calling stop() while uninitialized - has no effect
+        engine.tput_timer.stop()
+        assert not engine.tput_timer.initialized
+        assert not engine.tput_timer.started
+        assert engine.tput_timer.start_time == 0
+        assert engine.tput_timer.micro_step_count == 0
+        assert engine.tput_timer.global_step_count == 0
+        assert engine.tput_timer.total_elapsed_time == 0
+
+        # any call to start() (from dataloader or not) initializes the timer
+        engine.tput_timer.start()
+        assert engine.tput_timer.initialized
+        assert engine.tput_timer.started
+        assert engine.tput_timer.start_time == 0
+        assert engine.tput_timer.micro_step_count == 0
+        assert engine.tput_timer.global_step_count == 0
+        assert engine.tput_timer.total_elapsed_time == 0
+
+        # calling stop() after initialized - increments the local micro step counter
+        engine.tput_timer.stop()
+        assert engine.tput_timer.initialized
+        assert not engine.tput_timer.started
+        assert engine.tput_timer.start_time == 0
+        assert engine.tput_timer.micro_step_count == 1
+        assert engine.tput_timer.global_step_count == 0
+        assert engine.tput_timer.total_elapsed_time == 0
+
+        # calling start()/stop() to increment the step counter until start_step
+        while engine.tput_timer.micro_step_count < (gradient_accumulation_steps * engine.tput_timer.start_step):
+            engine.tput_timer.start()
+            global_step = (engine.tput_timer.micro_step_count + 1) % gradient_accumulation_steps == 0
+            engine.tput_timer.stop(global_step=global_step)
+        assert engine.tput_timer.global_step_count == engine.tput_timer.start_step
+        assert engine.tput_timer.total_elapsed_time == 0
+
+        # calling start()/stop() accumulates duration during gradient accumulation
+        while engine.tput_timer.global_step_count == engine.tput_timer.start_step:
+            engine.tput_timer.start()
+            current_duration = engine.tput_timer.step_elapsed_time
+            total_duration = engine.tput_timer.total_elapsed_time
+
+            global_step = (engine.tput_timer.micro_step_count + 1) % gradient_accumulation_steps == 0
+            engine.tput_timer.stop(global_step=global_step)
+            duration = engine.tput_timer.end_time - engine.tput_timer.start_time
+            # step elapsed time is reset after gradient accumulation steps
+            assert engine.tput_timer.step_elapsed_time == (0 if engine.tput_timer.global_step_count
+                                                           != engine.tput_timer.start_step else current_duration +
+                                                           duration)
+            assert engine.tput_timer.total_elapsed_time == total_duration + duration
+
+    def test_ext_param_getattr(self):
+        setup_serial_env()
+
+        class ExtLinear(torch.nn.Module):
+
+            def __init__(self, dim=16):
+                super().__init__()
+                self.dim = dim
+                self.linear1 = torch.nn.Linear(dim, dim)
+                self.linear2 = torch.nn.Linear(dim, dim)
+
+            def forward(self, input):
+                A = self.linear1(input)
+                B = self.linear2(A)
+
+                # external use of self.linear1.weight
+                C = torch.nn.functional.linear(B, self.linear1.weight)
+                return C.sum()
+
+        net = ExtLinear()
+
+        args = SimpleNamespace(local_rank=0)
+        engine, optim, _, _ = deepspeed.initialize(args=args,
+                                                   model=net,
+                                                   model_parameters=net.parameters(),
+                                                   config=config)
+
+        with deepspeed.zero.GatheredParameters(net.linear1.weight):
+            assert net.linear1.weight.numel() == net.dim**2
+
+        input = torch.rand(net.dim).to(engine.device).to(preferred_dtype())
+        loss = engine(input)
+        engine.backward(loss)
+        engine.step()
+
+
+class TestScatterGather(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        with deepspeed.zero.Init():
+            l = torch.nn.Linear(6, 3)
+        assert l.weight.ds_status == ZeroParamStatus.NOT_AVAILABLE
+        assert l.weight.shape == torch.Size(partitioned_param_data_shape)
+
+        # Ensure there is no impact outside the context
+        l2 = torch.nn.Linear(6, 3)
+        assert not hasattr(l2.weight, 'ds_status')
+        assert l2.weight.numel() == l2.in_features * l2.out_features
+
+        with deepspeed.zero.GatheredParameters(l.weight):
+            assert l.weight.ds_status == ZeroParamStatus.AVAILABLE
+            assert l.weight.numel() == l.in_features * l.out_features
+
+
+class TestGatherUpdate(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        with deepspeed.zero.Init():
+            l = torch.nn.Linear(4, 2)
+        assert l.weight.ds_status == ZeroParamStatus.NOT_AVAILABLE
+
+        # Gather and make a change
+        with deepspeed.zero.GatheredParameters(l.weight, modifier_rank=1):
+            assert l.weight.ds_status == ZeroParamStatus.AVAILABLE
+            if dist.get_rank() == 1:
+                with torch.no_grad():
+                    l.weight.zero_()
+
+        # should now be scattered again
+
+        # Now gather again and ensure the change is global
+        with deepspeed.zero.GatheredParameters(l.weight):
+            # all ranks compare
+            assert torch.equal(l.weight, torch.zeros_like(l.weight))

tests/unit/runtime/zero/test_zero_context_ancestry.py

@@ -0,0 +1,113 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+from deepspeed.accelerator import get_accelerator
+
+from utils import setup_serial_env
+from unit.common import DistributedTest
+
+config = {
+    "train_batch_size": 1,
+    "steps_per_print": 1,
+    "optimizer": {
+        "type": "Adam",
+        "params": {
+            "lr": 0.00015
+        }
+    },
+    "fp16": {
+        "enabled": True,
+        "loss_scale": 138.
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_param_persistence_threshold": 1,
+    }
+}
+
+
+# test that sub-classes get params that aren't prematurely partitioned and thus requiring gathering
+# fixed by https://github.com/deepspeedai/DeepSpeed/pull/1202
+class GrandPa(torch.nn.Module):
+
+    def __init__(self, *args):
+        super().__init__(*args)
+        self.param_grandpa = torch.nn.Parameter(torch.ones(5))
+        self.param_grandpa.data = (self.param_grandpa.data + 1).data  # test param is not yet partitioned
+
+
+class Pa(GrandPa):
+
+    def __init__(self, *args):
+        super().__init__(*args)
+        self.param_pa = torch.nn.Parameter(torch.ones(5))
+        self.param_pa.data = (self.param_pa.data + 1).data  # test param is not yet partitioned
+        self.param_grandpa.data = (self.param_grandpa.data + 1).data  # test param is not yet partitioned
+
+
+class Son(Pa):
+
+    def __init__(self):
+        super().__init__()
+        self.param = torch.nn.Parameter(torch.ones(5))
+        self.param.data = (self.param.data + 1).data  # test param is not yet partitioned
+        self.param_pa.data = (self.param_pa.data + 1).data  # test param is not yet partitioned
+        self.param_grandpa.data = (self.param_grandpa.data + 1).data  # test param is not yet partitioned
+
+
+class TestSerialParamInit(DistributedTest):
+    world_size = 1
+    init_distributed = False
+    set_dist_env = False
+
+    def test_subclass_param_init(self):
+        setup_serial_env()
+        with deepspeed.zero.Init(config=config):
+            model = Son().cpu()
+
+        # test that all params have been partitioned
+        assert model.param_grandpa.ds_status == ZeroParamStatus.NOT_AVAILABLE
+        assert model.param_pa.ds_status == ZeroParamStatus.NOT_AVAILABLE
+        assert model.param.ds_status == ZeroParamStatus.NOT_AVAILABLE
+
+        # test that the weights manipulation during each __init__ worked in all w/o needing gathering
+        ones = torch.ones(5).half().to(get_accelerator().device_name())
+        with deepspeed.zero.GatheredParameters(list(model.parameters(recurse=False))):
+            assert torch.equal(model.param, ones + 1)
+            assert torch.equal(model.param_pa, ones + 2)
+            assert torch.equal(model.param_grandpa, ones + 3)
+
+
+class TestDSInitWZinit(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        ds_config = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            }
+        }
+
+        class Model(torch.nn.Module):
+
+            def __init__(self):
+                super(Model, self).__init__()
+                self.linear = torch.nn.Linear(4, 4)
+
+            def magic(self):
+                return 42
+
+        with deepspeed.zero.Init():
+            model = Model()
+            engine, *_ = deepspeed.initialize(model=model, config=ds_config, model_parameters=model.parameters())
+        assert engine.magic() == 42

tests/unit/runtime/zero/test_zero_context_return.py

@@ -0,0 +1,187 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from types import SimpleNamespace
+import torch
+import pytest
+import deepspeed
+from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+from deepspeed.accelerator import get_accelerator
+
+from utils import setup_serial_env
+from unit.common import DistributedTest, preferred_dtype
+
+
+class DanglingBias(torch.nn.Linear):
+
+    def forward(self, *inputs):
+        out = super().forward(*inputs)
+        # return the bias to trigger a dangling external param
+        return out, self.bias
+
+
+class DataClass:
+    """Just wraps data in an object. """
+
+    def __init__(self, out=None, bias=None):
+        self.out = out
+        self.bias = bias
+
+
+class DanglingBiasClass(DanglingBias):
+
+    def forward(self, *inputs):
+        out, bias = super().forward(*inputs)
+        return DataClass(out=out, bias=bias)
+
+
+class DanglingAttention(torch.nn.Linear):
+
+    def __init__(self, dim=16, return_obj=False):
+        super().__init__(dim, dim)
+        self.dim = dim
+        self.return_obj = return_obj
+        if return_obj:
+            self.d_linear = DanglingBiasClass(dim, dim)
+        else:
+            self.d_linear = DanglingBias(dim, dim)
+
+    def forward(self, input):
+        out = super().forward(input)
+        if self.return_obj:
+            out_obj = self.d_linear(out)
+            assert out_obj.bias.ds_status == ZeroParamStatus.AVAILABLE
+            # forward the external param
+            return out_obj.out, out_obj.bias
+        else:
+            out, bias = self.d_linear(out)
+            assert hasattr(bias, 'ds_status') or hasattr(bias, 'ds_param_alias')
+            z3_bias = bias if hasattr(bias, 'ds_status') else bias.ds_param_alias
+            assert z3_bias.ds_status == ZeroParamStatus.AVAILABLE
+            return out, bias
+
+
+class ModelContainer(torch.nn.Module):
+
+    def __init__(self, dim=16, return_obj=False):
+        super().__init__()
+        self.dim = dim
+        self.linear1 = torch.nn.Linear(dim, dim)
+        self.dangler = DanglingAttention(dim, return_obj=return_obj)
+
+    def forward(self, input):
+        act1 = self.linear1(input)
+        # bias is actually dangler.d_linear1.bias
+        act2, bias = self.dangler(act1)
+        return (act2 + bias).sum()
+
+
+class DanglingExt(torch.nn.Module):
+
+    def __init__(self, dim=16):
+        super().__init__()
+        self.dim = dim
+        self.container = ModelContainer(dim)
+
+    def forward(self, input):
+        out = self.container(input)
+
+        # Make sure it's at the right level of the stack
+        assert len(self._external_params) == 0
+        assert len(self.container._external_params) == 1
+        assert len(self.container.dangler._external_params) == 0
+        return out
+
+
+class ModelContainerVariableOutputType(ModelContainer):
+
+    def __init__(self, dim=16, output_type=dict):
+        super().__init__()
+        self.output_type = output_type
+        self.dim = dim
+        self.linear1 = torch.nn.Linear(dim, dim)
+
+    def forward(self, input):
+        act1 = self.linear1(input)
+        if self.output_type is dict:
+            return {'loss': act1.sum()}
+        if self.output_type is torch.tensor:
+            return act1.sum()
+
+
+config = {
+    "train_batch_size": 1,
+    "steps_per_print": 1,
+    "optimizer": {
+        "type": "Adam",
+        "params": {
+            "lr": 0.00015
+        }
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "stage3_param_persistence_threshold": 1,
+    }
+}
+
+if get_accelerator().is_bf16_supported():
+    config["bf16"] = {"enabled": True}
+elif get_accelerator().is_fp16_supported():
+    config["fp16"] = {"enabled": True, "loss_scale": 138.}
+
+
+class TestReturnParam(DistributedTest):
+    world_size = 1
+
+    def test_ext_param_return(self):
+        setup_serial_env()
+
+        net = DanglingExt()
+
+        args = SimpleNamespace(local_rank=0)
+        engine, _, _, _ = deepspeed.initialize(args=args, model=net, model_parameters=net.parameters(), config=config)
+
+        for _ in range(5):
+            input = torch.rand(net.dim).to(engine.device).to(preferred_dtype())
+            loss = engine(input)
+            engine.backward(loss)
+            engine.step()
+
+    @pytest.mark.skip('WIP')
+    def test_ext_param_returnobj(self):
+        setup_serial_env()
+        print()
+
+        net = ModelContainer(return_obj=True)
+
+        args = SimpleNamespace(local_rank=0)
+        engine, _, _, _ = deepspeed.initialize(args=args, model=net, model_parameters=net.parameters(), config=config)
+
+        for _ in range(5):
+            input = torch.rand(net.dim).to(engine.device).to(preferred_dtype())
+            loss = engine(input)
+            assert len(net._external_params) == 1
+            assert len(net.dangler._external_params) == 0
+            engine.backward(loss)
+            engine.step()
+
+    @pytest.mark.parametrize('output_type', [torch.tensor, dict, None])
+    def test_stage_3_output_type(self, output_type):
+        setup_serial_env()
+        print()
+
+        net = ModelContainerVariableOutputType(output_type=output_type)
+
+        args = SimpleNamespace(local_rank=0)
+        engine, _, _, _ = deepspeed.initialize(args=args, model=net, model_parameters=net.parameters(), config=config)
+
+        for _ in range(1):
+            input = torch.rand(net.dim).to(engine.device).to(preferred_dtype())
+            loss = engine(input)
+            if loss is not None:
+                if isinstance(loss, dict):
+                    loss = loss['loss']
+                engine.backward(loss)
+                engine.step()

tests/unit/runtime/zero/test_zero_dynamic_class.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from unit.common import DistributedTest
+
+import deepspeed
+
+
+class TestNewClassDeclaredNestingInit(DistributedTest):
+    world_size = 1
+
+    def test_new_class_declared_nesting_init(self):
+        ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
+
+        with deepspeed.zero.Init(config_dict_or_path=ds_config):
+
+            class MyModel(torch.nn.Module):
+
+                def __init__(self):
+                    super().__init__()
+                    self.fc = torch.nn.Linear(4, 4)
+
+            with deepspeed.zero.Init(config_dict_or_path=ds_config):
+                model = MyModel()
+
+        # ensure that zero3 processed the parameter
+        assert hasattr(model.fc.weight, "ds_id")
+        deepspeed_engine, *_ = deepspeed.initialize(model=model, config_params=ds_config)
+
+
+class TestNewClassDeclaredInsideNestingInit(DistributedTest):
+    world_size = 1
+
+    def test_new_class_declared_inside_nesting_init(self):
+        ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
+
+        with deepspeed.zero.Init(config_dict_or_path=ds_config):
+
+            class MyModel(torch.nn.Module):
+
+                def __init__(self):
+                    super().__init__()
+                    self.fc = torch.nn.Linear(1, 1)
+
+            model = MyModel()
+
+        # ensure that zero3 processed the parameter
+        assert hasattr(model.fc.weight, "ds_id")
+        deepspeed_engine, *_ = deepspeed.initialize(model=model, config_params=ds_config)

tests/unit/runtime/zero/test_zero_grad_clip.py

@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import pytest
+import deepspeed
+from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
+from deepspeed.utils import safe_get_local_grad, safe_set_local_grad
+from deepspeed.accelerator import get_accelerator
+from unit.simple_model import SimpleModel
+import os
+
+
+def get_config(precision, clip_value, offload_device="cpu"):
+    config = {
+        "train_batch_size": 8,
+        "steps_per_print": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 1e-4
+            }
+        },
+        "zero_optimization": {
+            "stage": 3,
+            "offload_optimizer": {
+                "device": offload_device
+            },
+            "contiguous_gradients": True,
+            "overlap_comm": False,
+        },
+        "gradient_clipping": 1.0,
+    }
+
+    if precision == "fp16":
+        config["fp16"] = {
+            "enabled": True,
+            "loss_scale": 1024,
+            "initial_scale_power": 10,
+        }
+    elif precision == "bf16":
+        config["bf16"] = {
+            "enabled": True,
+        }
+
+    return config
+
+
+@pytest.mark.parametrize("precision,clip_value,offload_device", [
+    ("fp16", 0.5, "cpu"),
+    ("bf16", 0.05, "cpu"),
+    ("fp16", 0.5, "none"),
+    ("bf16", 0.05, "none"),
+])
+class TestZeroGradClip():
+    world_size = 1
+
+    def test_grad_clip_and_norm_update(self, precision, clip_value, offload_device):
+        """Test custom gradient clipping with configurations and to check if the norm_groups are updated correctly"""
+        config_dict = get_config(precision, clip_value, offload_device)
+
+        model = SimpleModel(hidden_dim=10)
+
+        # Set up distributed environment variables
+        os.environ['LOCAL_RANK'] = '0'
+        os.environ['RANK'] = '0'
+        os.environ['WORLD_SIZE'] = '1'
+        os.environ['MASTER_ADDR'] = '127.0.0.1'
+        os.environ['MASTER_PORT'] = '29500'
+
+        try:
+            model_engine, optimizer, _, _ = deepspeed.initialize(args=None,
+                                                                 model=model,
+                                                                 config=config_dict,
+                                                                 model_parameters=model.parameters(),
+                                                                 dist_init_required=True)
+        except Exception as e:
+            pytest.skip("Could not initialize deepspeed")
+
+        assert isinstance(optimizer, DeepSpeedZeroOptimizer_Stage3)
+
+        torch.manual_seed(1670)
+        inputs = torch.randn(8, 10, device=model_engine.device)
+        targets = torch.randn(8, 10, device=model_engine.device)
+
+        if model_engine.fp16_enabled() and get_accelerator().is_fp16_supported():
+            inputs = inputs.half()
+            targets = targets.half()
+        elif model_engine.bfloat16_enabled() and get_accelerator().is_bf16_supported():
+            inputs = inputs.bfloat16()
+            targets = targets.bfloat16()
+        else:
+            pytest.skip("Unsupported precision")
+
+        loss = model_engine(inputs, targets)
+        model_engine.backward(loss)
+
+        pre_clip_norm_groups = optimizer._get_norm_groups()
+        pre_clip_global_norm = torch.linalg.vector_norm(torch.stack(pre_clip_norm_groups))
+
+        modified_count = 0
+
+        for param in model_engine.parameters():
+            if not hasattr(param, 'ds_id'):
+                continue
+
+            grad = safe_get_local_grad(param)
+            if grad is not None:
+                pre_clip_norm = grad.norm().item()
+                clamped_grad = torch.clamp(grad, -clip_value, clip_value)
+                post_clip_norm = clamped_grad.norm().item()
+
+                if pre_clip_norm > clip_value:
+                    # Checks if the post-clip norm is less than the pre-clip norm
+                    assert post_clip_norm < pre_clip_norm, f"Post-clip norm should be < pre-clip norm for param {param.ds_id}"
+
+                safe_set_local_grad(param, clamped_grad)
+                modified_count += 1
+
+        # Get post-clip state
+        post_clip_norm_groups = optimizer._get_norm_groups()
+        post_clip_global_norm = torch.linalg.vector_norm(torch.stack(post_clip_norm_groups))
+
+        assert modified_count > 0, "No parameters were modified during clipping"
+        assert post_clip_global_norm.item() < pre_clip_global_norm.item(
+        ), f"Post-clip norm {post_clip_global_norm.item():.6f} should be < pre-clip norm {pre_clip_global_norm.item():.6f}"
+
+        model_engine.step()
+        final_norm = optimizer._global_grad_norm
+        if pre_clip_global_norm.item() > clip_value:
+            assert post_clip_global_norm.item() < pre_clip_global_norm.item(
+            ), "Global norm should be reduced after clipping when pre-clip norm > clip_value"

tests/unit/runtime/zero/test_zero_late_module_attach.py

@@ -0,0 +1,136 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Regression tests for issue #6961.
+
+ZeRO-3 forward used to crash with ``AttributeError: 'dict' object has no
+attribute '_in_forward'`` when a submodule's ``_parameters`` was a plain
+``dict`` instead of a ``ZeROOrderedDict``. PyTorch 2.5+ defaults
+``nn.Module._parameters`` to ``dict`` (pytorch/pytorch#129164), and any
+module not converted at ``DeepSpeedZeRoOffload`` init time hits the crash.
+The tests force the plain-dict condition explicitly so they exercise the
+fix on every supported torch version, not only torch 2.5+.
+"""
+
+import torch
+
+import deepspeed
+from deepspeed.runtime.zero.parameter_offload import (ZeROOrderedDict, ensure_zero_ordered_dict)
+
+from unit.common import DistributedTest, preferred_dtype
+
+
+class _Tiny(torch.nn.Module):
+
+    def __init__(self, hidden_dim=16):
+        super().__init__()
+        self.fc = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+def _zero3_config(dtype):
+    return {
+        "train_batch_size": 1,
+        "fp16": {
+            "enabled": dtype is torch.float16
+        },
+        "bf16": {
+            "enabled": dtype is torch.bfloat16
+        },
+        "zero_optimization": {
+            "stage": 3
+        },
+    }
+
+
+class TestZero3LateModuleAttach(DistributedTest):
+    world_size = 1
+
+    def test_forward_after_late_submodule_attach(self):
+        """Attaching a fresh ``nn.Linear`` after ``initialize`` must not crash."""
+        hidden = 16
+        dtype = preferred_dtype()
+        model = _Tiny(hidden)
+        engine, *_ = deepspeed.initialize(model=model,
+                                          config=_zero3_config(dtype),
+                                          model_parameters=list(model.parameters()))
+
+        late = torch.nn.Linear(hidden, hidden, bias=False).to(device=engine.device, dtype=dtype)
+        # Force the post-pytorch/pytorch#129164 condition deterministically so
+        # the test exercises the fix regardless of the installed torch version.
+        late._parameters = dict(late._parameters)
+        engine.module.late = late
+
+        x = torch.randn(2, hidden, dtype=dtype, device=engine.device)
+        engine(x)
+
+        # Prologue must have lazily converted the late submodule.
+        assert isinstance(engine.module.late._parameters, ZeROOrderedDict)
+
+    def test_idempotent_on_already_injected_modules(self):
+        """Repeated forwards must not re-wrap an already-converted ``_parameters``."""
+        hidden = 16
+        dtype = preferred_dtype()
+        model = _Tiny(hidden)
+        engine, *_ = deepspeed.initialize(model=model,
+                                          config=_zero3_config(dtype),
+                                          model_parameters=list(model.parameters()))
+
+        first_pdict = engine.module.fc._parameters
+        assert isinstance(first_pdict, ZeROOrderedDict)
+
+        x = torch.randn(2, hidden, dtype=dtype, device=engine.device)
+        engine(x)
+        engine(x)
+
+        assert engine.module.fc._parameters is first_pdict
+
+
+class TestEnsureZeroOrderedDict:
+    """Direct unit tests for the helper. No distributed harness needed."""
+
+    def test_skips_already_converted(self):
+        m = torch.nn.Linear(4, 4, bias=False)
+        m._parameters = ZeROOrderedDict(parent_module=m)
+        before = m._parameters
+        ensure_zero_ordered_dict(m)
+        assert m._parameters is before
+
+    def test_wraps_plain_dict(self):
+        m = torch.nn.Linear(4, 4, bias=False)
+        m._parameters = dict(m._parameters)
+        ensure_zero_ordered_dict(m)
+        assert isinstance(m._parameters, ZeROOrderedDict)
+        assert "weight" in m._parameters
+        assert m._original_parameters is not m._parameters
+
+    def test_preserves_existing_original_parameters(self):
+        """Subsequent wraps must not clobber the first-saved original.
+
+        ``_inject_parameters`` at engine init records the true torch-native
+        container in ``_original_parameters``; the deepcompile path in
+        ``init_z3.py`` reads it back to un-inject. If the helper later runs
+        after some intermediate replacement of ``_parameters``, it must not
+        overwrite that saved reference.
+        """
+        m = torch.nn.Linear(4, 4, bias=False)
+        sentinel = m._parameters
+        m._original_parameters = sentinel
+        m._parameters = dict(sentinel)  # different object, same contents
+        ensure_zero_ordered_dict(m)
+        assert m._original_parameters is sentinel
+
+    def test_noop_when_parameters_missing(self):
+        """Helper must not raise when ``_parameters`` is missing or None."""
+
+        class Bare:
+            pass
+
+        m = Bare()
+        ensure_zero_ordered_dict(m)  # no-op, no exception
+        m._parameters = None
+        ensure_zero_ordered_dict(m)  # no-op, no exception
+        assert m._parameters is None

tests/unit/runtime/zero/test_zero_leaf_module.py

@@ -0,0 +1,544 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import deepspeed.comm as dist
+import torch
+
+from unit.common import DistributedTest, preferred_dtype
+from unit.simple_model import random_dataloader
+
+import deepspeed
+from deepspeed.utils import set_z3_leaf_modules, unset_z3_leaf_modules, get_z3_leaf_modules, z3_leaf_module, \
+    set_z3_leaf_modules_by_name, set_z3_leaf_modules_by_suffix
+from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
+from deepspeed.runtime.zero.leaf_module_config import (DEFAULT_LEAF_MODULE_CLASSES, DEFAULT_LEAF_MODULE_NAMES,
+                                                       DEFAULT_LEAF_MODULE_NAME_SUFFIXES)
+from deepspeed.accelerator import get_accelerator
+from torch import nn
+import time
+
+
+class ChooseModuleByCounter(torch.nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(ChooseModuleByCounter, self).__init__()
+        self.linears = torch.nn.ModuleList(
+            [torch.nn.Linear(hidden_dim, hidden_dim, bias=False),
+             torch.nn.Linear(hidden_dim, hidden_dim, bias=False)])
+        self.act = torch.nn.ReLU()
+        self.cel = torch.nn.CrossEntropyLoss()
+        self.counter = 0
+
+    def forward(self, x, y):
+        # This fails without setting this module as a leaf module.
+        # See the comment in `set_z3_leaf_modules()`.
+        x = self.linears[self.counter % len(self.linears)](x)
+        x = self.act(x)
+        loss = self.cel(x, y)
+        self.counter += 1
+        return x, loss
+
+
+class ChooseModuleByRankModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(ChooseModuleByRankModel, self).__init__()
+        self.linears = torch.nn.ModuleList(
+            [torch.nn.Linear(hidden_dim, hidden_dim, bias=False),
+             torch.nn.Linear(hidden_dim, hidden_dim, bias=False)])
+        self.act = torch.nn.ReLU()
+        self.cel = torch.nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        # Each rank runs only one of the linear layers
+        x = self.linears[dist.get_rank() % len(self.linears)](x)
+        x = self.act(x)
+        loss = self.cel(x, y)
+        return x, loss
+
+
+class MultiOutputMoEBlock(nn.Module):
+    """A simplified MoE block that returns multiple tensors.
+
+    This model mimics Qwen3 MoE which returns (hidden_states, router_logits).
+    When used with ZeRO3 leaf modules and autograd multithreading enabled,
+    this pattern previously caused race conditions in fetch_sub_module
+    because backward hooks could be triggered concurrently from multiple threads.
+
+    See: https://github.com/deepspeedai/DeepSpeed/issues/7824
+    """
+
+    def __init__(self, hidden_dim, num_experts=4):
+        super(MultiOutputMoEBlock, self).__init__()
+        self.num_experts = num_experts
+        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
+        self.experts = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim, bias=False) for _ in range(num_experts)])
+        self.act = nn.ReLU()
+        self.cel = nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        # Compute router logits - this tensor will have gradients flowing through it
+        router_logits = self.gate(x)
+
+        # Process through experts
+        for expert in self.experts:
+            x = expert(x)
+        x = self.act(x)
+        loss = self.cel(x, y)
+
+        # Return multiple tensors - this triggers concurrent backward hooks
+        # when autograd multithreading is enabled
+        return x, loss, router_logits
+
+
+class MLPBlock(nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(MLPBlock, self).__init__()
+        self.gate_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.act_fn = nn.GELU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class FineGrainedBlock(nn.Module):
+
+    def __init__(self, hidden_dim, num_block):
+        super(FineGrainedBlock, self).__init__()
+        self.num_block = num_block
+        self.mlp_layers = torch.nn.ModuleList([MLPBlock(hidden_dim=hidden_dim) for _ in range(self.num_block)])
+
+    def forward(self, x):
+        for i in range(self.num_block):
+            x = self.mlp_layers[i](x)
+        return x
+
+
+class BaseLeafModule(nn.Module):
+
+    def __init__(self):
+        super(BaseLeafModule, self).__init__()
+
+
+class SubLeafModule(BaseLeafModule):
+
+    def __init__(self, hidden_dim):
+        super(SubLeafModule, self).__init__()
+        self.proj = nn.Linear(hidden_dim, hidden_dim)
+
+    def forward(self, x):
+        return self.proj(x)
+
+
+class WrapperLeafModule(nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(WrapperLeafModule, self).__init__()
+        self.child = SubLeafModule(hidden_dim)
+
+    def forward(self, x):
+        return self.child(x)
+
+
+def test_set_leaf_modules_with_fully_qualified_name():
+    hidden_dim = 16
+    model = WrapperLeafModule(hidden_dim)
+    fq_name = f"{SubLeafModule.__module__}.{SubLeafModule.__qualname__}"
+
+    matched = set_z3_leaf_modules(model, [fq_name])
+
+    assert len(matched) == 1
+    assert matched[0] is model.child
+    assert z3_leaf_module(model.child)
+    assert not z3_leaf_module(model)
+
+
+def test_set_leaf_modules_no_raise_when_missing():
+    hidden_dim = 16
+    model = WrapperLeafModule(hidden_dim)
+
+    matched = set_z3_leaf_modules(model, ["NonExistentClass"], raise_if_not_found=False)
+
+    assert matched == []
+    assert not z3_leaf_module(model.child)
+
+
+def test_set_leaf_modules_by_name():
+    hidden_dim = 16
+    model = WrapperLeafModule(hidden_dim)
+
+    matched, missing = set_z3_leaf_modules_by_name(model, ["child"])
+
+    assert matched == [model.child]
+    assert missing == []
+    assert z3_leaf_module(model.child)
+
+
+def test_set_leaf_modules_by_name_missing():
+    hidden_dim = 16
+    model = WrapperLeafModule(hidden_dim)
+
+    matched, missing = set_z3_leaf_modules_by_name(model, ["missing"], raise_if_not_found=False)
+
+    assert matched == []
+    assert missing == ["missing"]
+
+
+def test_set_leaf_modules_by_suffix():
+    hidden_dim = 16
+    model = WrapperLeafModule(hidden_dim)
+
+    matched, missing = set_z3_leaf_modules_by_suffix(model, ["child"])
+
+    assert missing == []
+    assert matched == [model.child]
+    assert z3_leaf_module(model.child)
+
+
+def test_set_leaf_modules_by_suffix_missing():
+    hidden_dim = 16
+    model = WrapperLeafModule(hidden_dim)
+
+    matched, missing = set_z3_leaf_modules_by_suffix(model, ["missing"], raise_if_not_found=False)
+
+    assert matched == []
+    assert missing == ["missing"]
+
+
+def test_zero_leaf_module_default_config():
+    config = DeepSpeedZeroConfig()
+    assert config.leaf_module.classes == DEFAULT_LEAF_MODULE_CLASSES
+    assert config.leaf_module.names == DEFAULT_LEAF_MODULE_NAMES
+    assert config.leaf_module.name_suffixes == DEFAULT_LEAF_MODULE_NAME_SUFFIXES
+
+
+def test_zero_leaf_module_custom_config():
+    payload = {
+        "leaf_module": {
+            "classes": ["custom.module.CustomClass"],
+            "names": ["transformer.layer"],
+            "name_suffixes": ["experts"]
+        }
+    }
+
+    config = DeepSpeedZeroConfig(**payload)
+
+    assert config.leaf_module.classes == ["custom.module.CustomClass"]
+    assert config.leaf_module.names == ["transformer.layer"]
+    assert config.leaf_module.name_suffixes == ["experts"]
+
+
+def test_zero_leaf_module_string_coercion():
+    payload = {"leaf_module": {"classes": "my.Class", "names": "submodule", "name_suffixes": "tail"}}
+
+    config = DeepSpeedZeroConfig(**payload)
+
+    assert config.leaf_module.classes == ["my.Class"]
+    assert config.leaf_module.names == ["submodule"]
+    assert config.leaf_module.name_suffixes == ["tail"]
+
+
+@pytest.mark.skip(reason="Requires Hugging Face transformers; run manually when validating defaults.")
+def test_default_leaf_module_classes_exist():
+    import importlib
+
+    from deepspeed.runtime.zero.leaf_module_config import DEFAULT_LEAF_MODULE_CLASSES
+
+    for cls_path in DEFAULT_LEAF_MODULE_CLASSES:
+        module_name, _, class_name = cls_path.rpartition('.')
+        module = importlib.import_module(module_name)
+        assert hasattr(module, class_name), f"Expected {class_name} in {module_name}"
+
+
+class modelWithFineGrainedBlock(nn.Module):
+
+    def __init__(self, hidden_dim, num_block):
+        super(modelWithFineGrainedBlock, self).__init__()
+        self.coarse_grained_layer1 = nn.Linear(hidden_dim, 8 * hidden_dim)
+        self.coarse_grained_layer2 = nn.Linear(8 * hidden_dim, hidden_dim)
+        self.fine_grained_layer = FineGrainedBlock(hidden_dim, num_block)
+        self.cel = torch.nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        x = self.coarse_grained_layer1(x)
+        x = self.coarse_grained_layer2(x)
+        x = self.fine_grained_layer(x)
+        loss = self.cel(x, y)
+        return x, loss
+
+
+def run_model(model, config_dict, hidden_dim, dtype, requires_grad):
+    model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+    data_loader = random_dataloader(model=model,
+                                    total_samples=10,
+                                    hidden_dim=hidden_dim,
+                                    device=model.device,
+                                    dtype=dtype)
+    dist.barrier()
+    for batch in data_loader:
+        batch[0].requires_grad = requires_grad
+        loss = model(batch[0], batch[1])
+        loss = loss[1]
+        model.backward(loss)
+        model.step()
+
+    # Needed in ZeRO 3. Not doing so can give memory leak
+    model.destroy()
+
+
+class TestSetZ3LeafModule(DistributedTest):
+    # Need multiple gpus to test possible hanging
+    world_size = 2
+    reuse_dist_env = True
+
+    def _create_zero_config(self, hidden_dim, leaf_module=None):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_prefetch_bucket_size": hidden_dim**2,
+                "stage3_param_persistence_threshold": 0,
+                "stage3_max_reuse_distance": 0,
+            }
+        }
+        if leaf_module is not None:
+            config_dict["zero_optimization"]["leaf_module"] = leaf_module
+
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        return config_dict
+
+    def _test_set_z3_leaf_modules(self, cls, requires_grad):
+        hidden_dim = 128
+        config_dict = self._create_zero_config(hidden_dim)
+
+        model = cls(hidden_dim)
+
+        assert not z3_leaf_module(model)
+        set_z3_leaf_modules(model, [cls])
+        assert z3_leaf_module(model)
+
+        run_model(model, config_dict, hidden_dim, preferred_dtype(), requires_grad)
+
+    def test_choose_module_by_counter(self):
+        self._test_set_z3_leaf_modules(ChooseModuleByCounter, True)
+
+    def test_choose_module_by_rank(self):
+        self._test_set_z3_leaf_modules(ChooseModuleByRankModel, True)
+
+    def test_multi_output_leaf_module_thread_safety(self):
+        """Test that leaf modules returning multiple tensors work correctly with autograd multithreading.
+
+        This tests the fix for https://github.com/deepspeedai/DeepSpeed/issues/7824
+        where MoE models (like Qwen3) returning multiple tensors caused race conditions
+        in fetch_sub_module when autograd executed backward hooks from multiple threads.
+        """
+        # Ensure autograd multithreading is enabled (this is the default, but be explicit)
+        torch.autograd.set_multithreading_enabled(True)
+
+        hidden_dim = 128
+        config_dict = self._create_zero_config(hidden_dim)
+
+        model = MultiOutputMoEBlock(hidden_dim, num_experts=4)
+
+        assert not z3_leaf_module(model)
+        set_z3_leaf_modules(model, [MultiOutputMoEBlock])
+        assert z3_leaf_module(model)
+
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=10,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=preferred_dtype())
+        dist.barrier()
+
+        # Run multiple iterations to increase chance of hitting race conditions
+        for batch in data_loader:
+            batch[0].requires_grad = True
+            # Model returns (output, loss, router_logits)
+            output, loss, router_logits = model(batch[0], batch[1])
+            # Include router_logits in the loss to ensure multiple backward paths
+            total_loss = loss + 0.01 * router_logits.mean()
+            model.backward(total_loss)
+            model.step()
+
+        model.destroy()
+
+    def test_multi_output_non_leaf_module_thread_safety(self):
+        """Ensure non-leaf modules returning multiple tensors remain thread-safe.
+
+        This covers the multi-output autograd multithreading case without marking the
+        module as a ZeRO leaf module.
+        """
+        torch.autograd.set_multithreading_enabled(True)
+
+        hidden_dim = 128
+        config_dict = self._create_zero_config(hidden_dim)
+
+        model = MultiOutputMoEBlock(hidden_dim, num_experts=4)
+        assert not z3_leaf_module(model)
+
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=10,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=preferred_dtype())
+        dist.barrier()
+
+        for batch in data_loader:
+            batch[0].requires_grad = True
+            output, loss, router_logits = model(batch[0], batch[1])
+            total_loss = loss + 0.01 * router_logits.mean()
+            model.backward(total_loss)
+            model.step()
+
+        model.destroy()
+
+    def test_no_grad_input_error(self):
+        try:
+            self._test_set_z3_leaf_modules(ChooseModuleByCounter, False)
+            raise AssertionError(
+                "Expected RuntimeError: inputs with requires_grad=False is not supported for a leaf module")
+        except RuntimeError as e:
+            pass
+
+    def test_set_unset_leaf_modules(self):
+        hidden_dim = 128
+        model = ChooseModuleByCounter(hidden_dim)
+        assert len(set_z3_leaf_modules(model, [torch.nn.ModuleList])) == 1, \
+            "Expected only one module to be set as a leaf module"
+        assert len(get_z3_leaf_modules(model)) == 1, "Expected there is only one leaf module"
+
+        assert len(unset_z3_leaf_modules(model, [torch.nn.ModuleList])) == 1, \
+            "Expected only one module to be unset as a leaf module"
+        assert len(get_z3_leaf_modules(model)) == 0, "Expected there is no leaf module"
+
+    def test_set_leaf_modules_with_subclass(self):
+        hidden_dim = 32
+        model = WrapperLeafModule(hidden_dim)
+
+        leaf_modules = set_z3_leaf_modules(model, [BaseLeafModule])
+
+        assert len(leaf_modules) == 1, "Expected the subclass instance to be marked as leaf"
+        assert leaf_modules[0] is model.child, "Expected the subclass instance to be returned"
+        assert z3_leaf_module(model.child), "Expected subclass instance flagged as leaf"
+        assert not z3_leaf_module(model), "Expected wrapper module to remain non-leaf"
+
+    def test_set_no_match_class(self):
+        hidden_dim = 128
+        model = ChooseModuleByCounter(hidden_dim)
+        try:
+            set_z3_leaf_modules(model, [torch.nn.Conv2d])
+            raise AssertionError("Expected error that no module is set as a leaf module")
+        except ValueError as e:
+            pass
+
+    def test_leaf_module_enabled_via_config(self):
+        hidden_dim = 128
+        leaf_class_fqn = f"{ChooseModuleByCounter.__module__}.{ChooseModuleByCounter.__qualname__}"
+        config_dict = self._create_zero_config(hidden_dim,
+                                               leaf_module={
+                                                   "classes": [leaf_class_fqn],
+                                                   "name_suffixes": ["linears"]
+                                               })
+
+        model = ChooseModuleByCounter(hidden_dim)
+        assert not z3_leaf_module(model)
+
+        run_model(model, config_dict, hidden_dim, preferred_dtype(), True)
+
+        assert z3_leaf_module(model)
+        modules_by_name = dict(model.named_modules())
+        assert "linears" in modules_by_name
+        assert z3_leaf_module(modules_by_name["linears"])
+
+
+@pytest.mark.parametrize("module_granularity_threshold", [0, 100, 12100, 10000000])
+class TestZ3LeafOptimization(DistributedTest):
+    world_size = 2
+    reuse_dist_env = True
+
+    def test_finegrained_optimization(self, module_granularity_threshold: int):
+        hidden_dim = 128
+        num_block = 16
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_prefetch_bucket_size": hidden_dim**2,
+                "stage3_param_persistence_threshold": 0,
+                "stage3_max_reuse_distance": 0,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        def bench_loss_and_time(config):
+            warm_up_step = 10
+            model = modelWithFineGrainedBlock(hidden_dim, num_block)
+            model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
+            data_loader = random_dataloader(model=model,
+                                            total_samples=20,
+                                            hidden_dim=hidden_dim,
+                                            device=model.device,
+                                            dtype=preferred_dtype())
+            dist.barrier()
+            loss_list = []
+
+            for i, batch in enumerate(data_loader):
+                if i == warm_up_step:
+                    dist.barrier()
+                    get_accelerator().synchronize()
+                    start_time = time.time()
+                batch[0].requires_grad = True
+                loss = model(batch[0], batch[1])
+                loss = loss[1]
+                loss_list.append(loss)
+                model.backward(loss)
+                model.step()
+            get_accelerator().synchronize()
+            end_time = time.time()
+            duration = end_time - start_time
+            model.destroy()
+            return loss_list, duration
+
+        baseline_loss_list, baseline_exec_time = bench_loss_and_time(config_dict)
+
+        config_dict["zero_optimization"]["stage3_module_granularity_threshold"] = module_granularity_threshold
+        loss, duration = bench_loss_and_time(config_dict)
+
+        if dist.get_rank() == 0:
+            print("baseline exec time:", baseline_exec_time)
+            print(
+                f"finegrained optimziation exec time: {duration},granularity threshold:{module_granularity_threshold} "
+            )
+            assert baseline_loss_list == loss, f"incorrect loss value with threshold:{module_granularity_threshold}"

tests/unit/runtime/zero/test_zero_nesting_init.py

@@ -0,0 +1,93 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from unit.common import DistributedTest
+
+from transformers import GPT2Config, VisionEncoderDecoderConfig, VisionEncoderDecoderModel, ViTConfig
+from transformers.integrations.deepspeed import HfDeepSpeedConfig
+
+import deepspeed
+
+
+def _create_tiny_vision_encoder_decoder_model(model_path):
+    encoder_config = ViTConfig(image_size=8,
+                               patch_size=4,
+                               num_hidden_layers=1,
+                               hidden_size=8,
+                               num_attention_heads=2,
+                               intermediate_size=16)
+    decoder_config = GPT2Config(vocab_size=32,
+                                n_positions=16,
+                                n_embd=8,
+                                n_layer=1,
+                                n_head=2,
+                                bos_token_id=0,
+                                eos_token_id=1,
+                                add_cross_attention=True,
+                                is_decoder=True)
+    config = VisionEncoderDecoderConfig.from_encoder_decoder_configs(encoder_config, decoder_config)
+    model = VisionEncoderDecoderModel(config)
+    model.save_pretrained(model_path, safe_serialization=False)
+
+
+class TestNestingInit(DistributedTest):
+    world_size = 1
+
+    def test_nesting_init(self):
+        ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
+
+        with deepspeed.zero.Init(config_dict_or_path=ds_config):
+            with deepspeed.zero.Init(config_dict_or_path=ds_config):
+                model = torch.nn.Linear(4, 4)
+
+        # ensure that zero3 processed the parameter
+        assert hasattr(model.weight, "ds_id")
+
+        deepspeed_engine, *_ = deepspeed.initialize(model=model, config_params=ds_config)
+
+
+class TestShutdownInNestingInit(DistributedTest):
+    world_size = 1
+
+    def test_shutdown_in_nesting_init(self):
+        ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
+
+        with deepspeed.zero.Init(config_dict_or_path=ds_config):
+            with deepspeed.zero.Init(config_dict_or_path=ds_config):
+                model1 = torch.nn.Linear(4, 4)
+
+            assert hasattr(model1.weight, "ds_id")
+            deepspeed_engine1, *_ = deepspeed.initialize(model=model1, config_params=ds_config)
+            with deepspeed.zero.Init(config_dict_or_path=ds_config):
+                model2 = torch.nn.Linear(4, 4)
+
+        # ensure that zero3 processed the parameter
+        assert hasattr(model2.weight, "ds_id")
+        deepspeed_engine2, *_ = deepspeed.initialize(model=model2, config_params=ds_config)
+
+
+class TestNestedParallelInit(DistributedTest):
+    world_size = 1
+
+    # Testing a model with composed and nested zero.Inits, with 3 zero.Init contexts, 1 parent and 2 children.
+    # The skeleton of the model is like so
+    #
+    # class VisionEncoderDecoderModel(...)::
+    #     def __init__(self):
+    #             encoder = AutoModel.from_config(config.encoder)
+    #             decoder = AutoModelForCausalLM.from_config(config.decoder)
+    #
+    # And the user calls like below:
+    # VisionEncoderDecoderModel.from_pretrained(...)
+    # which calls this constructor inside zero.Init
+
+    def test_nested_parallel_init(self, tmp_path):
+        ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
+        _create_tiny_vision_encoder_decoder_model(tmp_path)
+        dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive
+        model = VisionEncoderDecoderModel.from_pretrained(str(tmp_path), local_files_only=True)
+        assert all([hasattr(p, 'ds_id') for p in model.parameters()])

tests/unit/runtime/zero/test_zero_offloadpp.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+import pytest
+import deepspeed.comm as dist
+from unit.common import DistributedTest
+from unit.simple_model import random_dataloader
+
+import deepspeed
+import torch
+from deepspeed.runtime.zero.offload_config import DeepSpeedZeroOffloadOptimizerConfig
+
+import torch.nn as nn
+
+
+class NNModel(nn.Module):
+
+    def __init__(self, h_dim=1024, n_layers=2):
+        super(NNModel, self).__init__()
+        self.layers = nn.ModuleList([nn.Linear(h_dim, h_dim) for i in range(n_layers)])
+        self.cross_entropy_loss = nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        for layer in self.layers:
+            x = layer(x)
+        return self.cross_entropy_loss(x, y)
+
+
+def test_zero_partial_offload_config():
+    config = DeepSpeedZeroOffloadOptimizerConfig(**{"ratio": 0.3})
+    assert config.ratio == 0.3
+
+
+#Large sweep along hidden dim, num_layers of different sizes
+@pytest.mark.parametrize("h_dim", [1024])
+@pytest.mark.parametrize("n_layers", [4, 8])
+class TestZeroPartialOffloadConfigSweep(DistributedTest):
+    world_size = 4
+
+    def test(self, h_dim: int, n_layers: int) -> None:
+
+        config_dict = {
+            "train_batch_size": 256,
+            "steps_per_print": 1,
+            "gradient_clipping": 1.0,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015,
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 15
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "sub_group_size": 8,
+                "reduce_bucket_size": 20,
+                "offload_optimizer": {
+                    "device": "cpu",
+                    "pin_memory": True,
+                    "ratio": 0.3
+                }
+            }
+        }
+
+        model = NNModel(h_dim, n_layers)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=20,
+                                        hidden_dim=h_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        dist.barrier()
+
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()

tests/unit/runtime/zero/test_zero_tensor_fragment.py

@@ -0,0 +1,460 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import deepspeed.comm as dist
+import torch
+import math
+
+from unit.common import DistributedTest
+from unit.simple_model import random_dataloader, SimpleModel
+from unit.util import bf16_required_version_check
+
+import deepspeed
+from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
+from deepspeed.utils import safe_set_full_fp32_param, safe_set_full_grad, safe_set_full_optimizer_state
+from deepspeed.utils import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state
+from deepspeed.utils import safe_set_local_fp32_param, safe_set_local_grad, safe_set_local_optimizer_state
+from deepspeed.utils import safe_update_full_grad_vectorized
+from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
+from deepspeed.ops.aio import AsyncIOBuilder
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.swap_tensor import MIN_SWAPPABLE_BYTES
+
+WEIGHT_KEY = 'weight'
+FIRST_ORDER_KEY = 'exp_avg'
+SECOND_ORDER_KEY = 'exp_avg_sq'
+GRADIENT_KEY = 'gradient'
+
+
+def validate_tensor(model, api_type, opt_states):
+    assert api_type in ["full", "local"]
+    for _, lp in model.named_parameters():
+        param_list = []
+        if opt_states:
+            param_list.append(
+                safe_get_full_optimizer_state(lp, 'exp_avg') if api_type ==
+                "full" else safe_get_local_optimizer_state(lp, 'exp_avg'))
+            param_list.append(
+                safe_get_full_optimizer_state(lp, 'exp_avg_sq') if api_type ==
+                "full" else safe_get_local_optimizer_state(lp, 'exp_avg_sq'))
+        else:
+            param_list.append(safe_get_full_fp32_param(lp) if api_type == "full" else safe_get_local_fp32_param(lp))
+            param_list.append(safe_get_full_grad(lp) if api_type == "full" else safe_get_local_grad(lp))
+        if lp.requires_grad:
+            assert all([p is not None for p in param_list])
+        else:
+            assert all([p is None for p in param_list])
+
+
+class MyModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim, frozen_weights):
+        super(MyModel, self).__init__()
+        self.act = torch.nn.ReLU()
+        self.cel = torch.nn.CrossEntropyLoss()
+        self.linears = torch.nn.ModuleList(
+            [torch.nn.Linear(hidden_dim, 1),
+             torch.nn.Linear(1, 1),
+             torch.nn.Linear(1, hidden_dim)])
+        if frozen_weights:
+            self.linears[0].weight.requires_grad = False
+            self.linears[0].bias.requires_grad = False
+
+    def forward(self, x, y):
+        for l in self.linears:
+            x = l(x)
+            x = self.act(x)
+        return self.cel(x, y)
+
+
+def run_fragmented_model(model, config_dict, hidden_dim, dtype, validate_after_bwd, validate_after_step):
+    model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+    data_loader = random_dataloader(model=model,
+                                    total_samples=10,
+                                    hidden_dim=hidden_dim,
+                                    device=model.device,
+                                    dtype=dtype)
+    dist.barrier()
+    for n, batch in enumerate(data_loader):
+        loss = model(batch[0], batch[1])
+        model.backward(loss)
+        validate_after_bwd(model)
+        model.step()
+        validate_after_step(model)
+
+    # Needed in ZeRO 3. Not doing so can give memory leak
+    model.destroy()
+
+
+@pytest.mark.parametrize('frozen_weights', [True, False])
+class TestTensorFragmentGet(DistributedTest):
+    # Need multiple gpus to test possible hanging
+    world_size = 2
+    reuse_dist_env = True
+
+    @pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32])
+    @pytest.mark.parametrize('api_type', ['local', 'full'])
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    @pytest.mark.parametrize('offload_device', [OffloadDeviceEnum.none, OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme])
+    def test_zero_fragments(self, tmpdir, dtype, api_type, zero_stage, offload_device, frozen_weights):
+        if not dtype in get_accelerator().supported_dtypes():
+            pytest.skip(f"{get_accelerator()._name} does not support {dtype} data type")
+
+        if offload_device == OffloadDeviceEnum.nvme:
+            if zero_stage != 3:
+                pytest.skip(f"Nvme offload not supported for zero stage {zero_stage}")
+            if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
+                pytest.skip('Skip tests since async-io is not compatible')
+
+        if api_type == "local" and zero_stage != 3:
+            pytest.skip(f"Local APIs only for zero stage 3 but current stage is {zero_stage}")
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+
+        if dtype == torch.half:
+            config_dict["fp16"] = {"enabled": True, "initial_scale_power": 2}
+        elif dtype == torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        if offload_device == OffloadDeviceEnum.cpu:
+            config_dict["zero_optimization"]["offload_optimizer"] = {"device": offload_device}
+        elif offload_device == OffloadDeviceEnum.nvme:
+            config_dict["zero_optimization"]["offload_optimizer"] = {
+                "device": offload_device,
+                "nvme_path": str(tmpdir)
+            }
+
+        hidden_dim = MIN_SWAPPABLE_BYTES
+        if zero_stage == 3:
+            with deepspeed.zero.Init(config_dict_or_path=config_dict):
+                model = MyModel(hidden_dim, frozen_weights)
+        else:
+            model = MyModel(hidden_dim, frozen_weights)
+
+        validate_after_bwd = lambda model: validate_tensor(model, api_type, opt_states=False)
+        validate_after_step = lambda model: validate_tensor(model, api_type, opt_states=True)
+
+        run_fragmented_model(model, config_dict, hidden_dim, dtype, validate_after_bwd, validate_after_step)
+
+    def test_bf16_optimizer_fragments(self, frozen_weights):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU accelerator does not support this test yet.")
+        if frozen_weights:
+            pytest.skip("TODO: Frozen weights not currently supported by BF16 Optimizer")
+
+        if not bf16_required_version_check():
+            pytest.skip(
+                " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
+            )
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "bf16": {
+                "enabled": True
+            },
+            # Use fp32 gradient accumulation to ensure BF16_Optimizer is used
+            # (bf16 model + bf16 grad_accum uses FP16_Optimizer which doesn't support tensor fragment APIs)
+            "data_types": {
+                "grad_accum_dtype": "fp32"
+            },
+            "zero_optimization": {
+                "stage": 1,
+            }
+        }
+
+        hidden_dim = 128
+        model = MyModel(hidden_dim, frozen_weights)
+
+        api_type = "full"
+        validate_after_bwd = lambda model: validate_tensor(model, api_type, opt_states=False)
+        validate_after_step = lambda model: validate_tensor(model, api_type, opt_states=True)
+
+        run_fragmented_model(model, config_dict, hidden_dim, torch.bfloat16, validate_after_bwd, validate_after_step)
+
+
+def create_random_values(model, key_list, group, grad_dtype):
+    param_values = {}
+    for n, lp in model.named_parameters():
+        param_shape = lp.ds_shape if hasattr(lp, 'ds_id') else lp.shape
+        param_values[n] = {}
+        for key in key_list:
+            dtype = grad_dtype if key == GRADIENT_KEY else torch.float32
+            rand_value = torch.rand(param_shape, dtype=dtype, device=model.device)
+            dist.broadcast(rand_value, src=0, group=group)
+            param_values[n][key] = rand_value
+    return param_values
+
+
+def set_param_values_with_dict(model, value_dict):
+    for n, lp in model.named_parameters():
+        for key, value_tensor in value_dict[n].items():
+            if key == GRADIENT_KEY:
+                safe_set_full_grad(lp, value_tensor)
+            elif key == WEIGHT_KEY:
+                safe_set_full_fp32_param(lp, value_tensor)
+            else:
+                safe_set_full_optimizer_state(lp, value_tensor, key)
+
+
+def update_param_values_with_dict(model, value_dict):
+    new_grad_values = {}
+    for n, lp in model.named_parameters():
+        if GRADIENT_KEY in value_dict[n]:
+            new_grad_values[id(lp)] = value_dict[n][GRADIENT_KEY]
+
+    def update_gradient_callback(old_value, param):
+        return new_grad_values[id(param)]
+
+    update_param_list = []
+    for n, lp in model.named_parameters():
+        for key, value_tensor in value_dict[n].items():
+            if key == GRADIENT_KEY:
+                update_param_list.append(lp)
+
+    if len(update_param_list) > 0:
+        safe_update_full_grad_vectorized(update_param_list, update_gradient_callback)
+
+
+def validate_param_values_with_dict(model, value_dict):
+    for n, lp in model.named_parameters():
+        for key, expected_tensor in value_dict[n].items():
+            if key == GRADIENT_KEY:
+                actual_tensor = safe_get_full_grad(lp)
+            elif key == WEIGHT_KEY:
+                actual_tensor = safe_get_full_fp32_param(lp)
+            else:
+                actual_tensor = safe_get_full_optimizer_state(lp, key)
+
+            assert torch.equal(expected_tensor, actual_tensor)
+
+
+def create_random_values_for_local(model, key_list, group, grad_dtype):
+    param_values = {}
+    for n, lp in model.named_parameters():
+        param_shape = lp.ds_tensor.shape
+        param_values[n] = {}
+        for key in key_list:
+            dtype = grad_dtype if key == GRADIENT_KEY else torch.float32
+            rand_value = torch.rand(param_shape, dtype=dtype, device=model.device)
+            param_values[n][key] = rand_value
+    return param_values
+
+
+def set_local_param_values_with_dict(model, value_dict):
+    for n, lp in model.named_parameters():
+
+        for key, value_tensor in value_dict[n].items():
+            if key == GRADIENT_KEY:
+                safe_set_local_grad(lp, value_tensor)
+            elif key == WEIGHT_KEY:
+                safe_set_local_fp32_param(lp, value_tensor)
+            else:
+                safe_set_local_optimizer_state(lp, value_tensor, key)
+
+
+def validate_local_param_values_with_dict(model, value_dict):
+    for n, lp in model.named_parameters():
+        for key, expected_tensor in value_dict[n].items():
+            if key == GRADIENT_KEY:
+                actual_tensor = safe_get_local_grad(lp)
+            elif key == WEIGHT_KEY:
+                actual_tensor = safe_get_local_fp32_param(lp)
+            else:
+                actual_tensor = safe_get_local_optimizer_state(lp, key)
+
+            assert torch.equal(expected_tensor, actual_tensor)
+
+
+helper_funcs_mapping = {
+    "full": {
+        "create_random_values": create_random_values,
+        "set_param_values_with_dict": set_param_values_with_dict,
+        "update_param_values_with_dict": update_param_values_with_dict,
+        "validate_param_values_with_dict": validate_param_values_with_dict,
+    },
+    "local": {
+        "create_random_values": create_random_values_for_local,
+        "set_param_values_with_dict": set_local_param_values_with_dict,
+        "validate_param_values_with_dict": validate_local_param_values_with_dict
+    }
+}
+
+
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32])
+class TestTensorFragmentSet(DistributedTest):
+    # Need multiple gpus to test possible hanging
+    world_size = 2
+    reuse_dist_env = True
+
+    @pytest.mark.parametrize('api_type', ['local', 'full'])
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    @pytest.mark.parametrize('offload_device', [OffloadDeviceEnum.none, OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme])
+    def test_zero_fragments(self, tmpdir, api_type, zero_stage, offload_device, dtype):
+        if not dtype in get_accelerator().supported_dtypes():
+            pytest.skip(f"{get_accelerator()._name} does not support {dtype} data type")
+
+        if dtype == torch.bfloat16 and not bf16_required_version_check(accelerator_check=False):
+            pytest.skip(
+                " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
+            )
+
+        if api_type == "local" and zero_stage != 3:
+            pytest.skip(f"Local APIs only for zero stage 3 but current stage is {zero_stage}")
+
+        if offload_device == OffloadDeviceEnum.nvme:
+            if zero_stage != 3:
+                pytest.skip(f"Nvme offload not supported for zero stage {zero_stage}")
+            if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
+                pytest.skip('Skip tests since async-io is not compatible')
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+
+        if offload_device == OffloadDeviceEnum.cpu:
+            config_dict["zero_optimization"]["offload_optimizer"] = {"device": offload_device}
+        elif offload_device == OffloadDeviceEnum.nvme:
+            config_dict["zero_optimization"]["offload_optimizer"] = {
+                "device": offload_device,
+                "nvme_path": str(tmpdir)
+            }
+
+        if dtype == torch.float16:
+            config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
+        elif dtype == torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        hidden_dim = int(math.sqrt(MIN_SWAPPABLE_BYTES))
+        if zero_stage == 3:
+            config_dict["zero_optimization"]["param_persistence_threshold"] = hidden_dim
+            with deepspeed.zero.Init(config_dict_or_path=config_dict):
+                model = SimpleModel(hidden_dim)
+        else:
+            model = SimpleModel(hidden_dim)
+
+        world = dist.get_world_size()
+        group = dist.new_group(ranks=list(range(world)))
+
+        dist.barrier()
+
+        def after_bwd_validate_func(model):
+            state_keys = [WEIGHT_KEY, GRADIENT_KEY]
+            helper_funcs = helper_funcs_mapping[api_type]
+            optim_state_values = helper_funcs["create_random_values"](model, state_keys, group, grad_dtype=dtype)
+            helper_funcs["set_param_values_with_dict"](model, optim_state_values)
+            helper_funcs["validate_param_values_with_dict"](model, optim_state_values)
+
+        def after_step_validate_func(model):
+            state_keys = [WEIGHT_KEY, FIRST_ORDER_KEY, SECOND_ORDER_KEY]
+            helper_funcs = helper_funcs_mapping[api_type]
+            optim_state_values = helper_funcs["create_random_values"](model, state_keys, group, grad_dtype=dtype)
+            helper_funcs["set_param_values_with_dict"](model, optim_state_values)
+            helper_funcs["validate_param_values_with_dict"](model, optim_state_values)
+
+        run_fragmented_model(model, config_dict, hidden_dim, dtype, after_bwd_validate_func, after_step_validate_func)
+
+
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32])
+class TestTensorFragmentUpdate(DistributedTest):
+    # Need multiple gpus to test possible hanging
+    world_size = 2
+    reuse_dist_env = True
+
+    @pytest.mark.parametrize('torch_adam', [False, True])
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    @pytest.mark.parametrize('offload_device', [OffloadDeviceEnum.none, OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme])
+    def test_zero_fragments(self, tmpdir, torch_adam, zero_stage, offload_device, dtype):
+        if not dtype in get_accelerator().supported_dtypes():
+            pytest.skip(f"{get_accelerator()._name} does not support {dtype} data type")
+
+        if offload_device == OffloadDeviceEnum.nvme:
+            if zero_stage != 3:
+                pytest.skip(f"Nvme offload not supported for zero stage {zero_stage}")
+            if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
+                pytest.skip('Skip tests since async-io is not compatible')
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6,
+                    "torch_adam": torch_adam
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+
+        if offload_device == OffloadDeviceEnum.cpu:
+            config_dict["zero_optimization"]["offload_optimizer"] = {"device": offload_device}
+        elif offload_device == OffloadDeviceEnum.nvme:
+            config_dict["zero_optimization"]["offload_optimizer"] = {
+                "device": offload_device,
+                "nvme_path": str(tmpdir)
+            }
+
+        if dtype == torch.float16:
+            config_dict["fp16"] = {"enabled": True, "initial_scale_power": 8}
+        elif dtype == torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        hidden_dim = int(math.sqrt(MIN_SWAPPABLE_BYTES))
+        if zero_stage == 3:
+            config_dict["zero_optimization"]["param_persistence_threshold"] = hidden_dim
+            with deepspeed.zero.Init(config_dict_or_path=config_dict):
+                model = SimpleModel(hidden_dim)
+        else:
+            model = SimpleModel(hidden_dim)
+
+        world = dist.get_world_size()
+        group = dist.new_group(ranks=list(range(world)))
+
+        dist.barrier()
+
+        api_type = "full"
+
+        def after_bwd_validate_func(model):
+            state_keys = [GRADIENT_KEY]
+            helper_funcs = helper_funcs_mapping[api_type]
+            optim_state_values = helper_funcs["create_random_values"](model, state_keys, group, grad_dtype=dtype)
+            helper_funcs["update_param_values_with_dict"](model, optim_state_values)
+            helper_funcs["validate_param_values_with_dict"](model, optim_state_values)
+
+        def after_step_validate_func(model):
+            pass
+
+        run_fragmented_model(model, config_dict, hidden_dim, dtype, after_bwd_validate_func, after_step_validate_func)

tests/unit/runtime/zero/test_zero_tiled.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import copy
+
+import torch
+from deepspeed.runtime.zero.tiling import TiledLinear, TiledLinearReturnBias
+
+import pytest
+
+
+@pytest.mark.parametrize('in_splits,out_splits', [(1, 1), (2, 2), (5, 5), (32, 32)])
+def test_tiled_init(in_splits, out_splits):
+    in_f = 32
+    out_f = 40
+    base = torch.nn.Linear(in_f, out_f, bias=True)
+    l = TiledLinear(in_f,
+                    out_f,
+                    bias=True,
+                    init_linear=copy.deepcopy(base),
+                    out_splits=out_splits,
+                    in_splits=in_splits)
+
+    for out_id in range(out_splits):
+        for in_id in range(in_splits):
+            local_l = l.linears[out_id][in_id]
+            assert isinstance(local_l, torch.nn.Linear)
+
+            rstart = l.out_parts[out_id]
+            rstop = l.out_parts[out_id + 1]
+            cstart = l.in_parts[in_id]
+            cstop = l.in_parts[in_id + 1]
+
+            local_out = rstop - rstart
+            local_in = cstop - cstart
+            assert local_l.weight.size()[1] == local_in, f'local[{out_id}][{in_id}].size {local_l.weight.size()}'
+            assert local_l.weight.size()[0] == local_out
+
+            test = base.weight[rstart:rstop, cstart:cstop]
+
+            assert local_l.weight.size() == test.size()
+            assert torch.equal(local_l.weight.data, test.data)
+
+            if in_id == in_splits - 1:
+                assert local_l.bias is not None
+                assert local_l.bias.size()[0] == local_out
+            else:
+                assert local_l.bias is None
+
+
+@pytest.mark.parametrize('in_splits,out_splits', [(0, 0), (33, 33)])
+def test_tiled_baddim(in_splits, out_splits):
+    dim = 32
+    with pytest.raises(RuntimeError):
+        l = TiledLinear(dim, dim, out_splits=out_splits, in_splits=in_splits)
+
+
+@pytest.mark.skip(reason="seeing nondeterministic failures, skipping for now")
+@pytest.mark.parametrize('bias', [False, True])
+@pytest.mark.parametrize('in_splits,out_splits', [(1, 1), (2, 2)])
+@pytest.mark.parametrize('in_f,out_f', [(32, 32), (23, 29), (29, 23)])
+def test_tiled_forward(in_splits, out_splits, bias, in_f, out_f):
+    base = torch.nn.Linear(in_f, out_f, bias=bias)
+    test = TiledLinear(in_f,
+                       out_f,
+                       bias=bias,
+                       init_linear=copy.deepcopy(base),
+                       out_splits=out_splits,
+                       in_splits=in_splits)
+
+    inp = torch.rand(in_f)
+
+    base_out = base(copy.deepcopy(inp))
+    test_out = test(copy.deepcopy(inp))
+
+    assert torch.allclose(base_out, test_out, rtol=1e-4)
+
+
+@pytest.mark.skip(reason="seeing nondeterministic failures, skipping for now")
+@pytest.mark.parametrize('bias', [False, True])
+@pytest.mark.parametrize('in_splits,out_splits', [(1, 1), (2, 2)])
+@pytest.mark.parametrize('in_f,out_f', [(32, 32), (23, 29), (29, 23)])
+def test_tiled_backward(in_splits, out_splits, bias, in_f, out_f):
+    base = torch.nn.Linear(in_f, out_f, bias=bias)
+    test = TiledLinear(in_f,
+                       out_f,
+                       bias=bias,
+                       init_linear=copy.deepcopy(base),
+                       out_splits=out_splits,
+                       in_splits=in_splits)
+
+    inp = torch.rand(in_f)
+
+    base_out = base(copy.deepcopy(inp))
+    test_out = test(copy.deepcopy(inp))
+    assert torch.allclose(base_out, test_out, rtol=1e-4)
+
+    base_out.sum().backward()
+    test_out.sum().backward()
+
+    # compare grads
+    for row in range(out_splits):
+        rstart = test.out_parts[row]
+        rstop = test.out_parts[row + 1]
+
+        for col in range(in_splits):
+            cstart = test.in_parts[col]
+            cstop = test.in_parts[col + 1]
+
+            local = test.linears[row][col]
+            base_grad = base.weight.grad[rstart:rstop, cstart:cstop]
+            assert torch.allclose(base_grad, local.weight.grad, rtol=1e-4)
+
+            if local.bias is not None:
+                base_grad = base.bias.grad[rstart:rstop]
+                assert torch.allclose(base_grad, local.bias.grad, rtol=1e-4)
+
+
+class LinearWrapper(torch.nn.Linear):
+    """Returns its own bias to simulate Megatron-LM's behavior.
+
+    Megatron-LM optionally delays the bias addition to fuse with a proceeding kernel.
+    """
+
+    def forward(self, input):
+        out = super().forward(input)
+        return out, self.bias
+
+
+@pytest.mark.skip(reason="seeing nondeterministic failures, skipping for now")
+@pytest.mark.parametrize('bias', [False, True])
+@pytest.mark.parametrize('in_splits,out_splits', [(1, 1), (2, 2)])
+@pytest.mark.parametrize('in_f,out_f', [(32, 32), (23, 29), (29, 23)])
+def test_tiled_returnbias_backward(in_splits, out_splits, bias, in_f, out_f):
+    base = LinearWrapper(in_f, out_f, bias=bias)
+    test = TiledLinearReturnBias(in_f,
+                                 out_f,
+                                 bias=bias,
+                                 linear_cls=LinearWrapper,
+                                 init_linear=copy.deepcopy(base),
+                                 out_splits=out_splits,
+                                 in_splits=in_splits)
+
+    inp = torch.rand(in_f)
+
+    base_out_t, base_out_b = base(copy.deepcopy(inp))
+    test_out_t, test_out_b = test(copy.deepcopy(inp))
+    assert torch.allclose(base_out_t, test_out_t, rtol=1e-4)
+    if base_out_b is None:
+        assert test_out_b is None
+        base_out_b = torch.zeros_like(base_out_t)
+        test_out_b = torch.zeros_like(test_out_t)
+    else:
+        assert test_out_b is not None
+        assert torch.allclose(base_out_b, test_out_b, rtol=1e-4)
+
+    (base_out_t + base_out_b).sum().backward()
+    (test_out_t + test_out_b).sum().backward()
+
+    # compare grads
+    for row in range(out_splits):
+        rstart = test.out_parts[row]
+        rstop = test.out_parts[row + 1]
+
+        for col in range(in_splits):
+            cstart = test.in_parts[col]
+            cstop = test.in_parts[col + 1]
+
+            local = test.linears[row][col]
+            base_grad = base.weight.grad[rstart:rstop, cstart:cstop]
+            assert torch.allclose(base_grad, local.weight.grad, rtol=1e-4)
+
+            if local.bias is not None:
+                base_grad = base.bias.grad[rstart:rstop]
+                assert torch.allclose(base_grad, local.bias.grad, rtol=1e-4)

tests/unit/runtime/zero/test_zeropp.py

@@ -0,0 +1,288 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+import pytest
+import deepspeed.comm as dist
+from torch.nn import Module
+
+from unit.common import DistributedTest
+from unit.simple_model import random_dataloader
+
+import deepspeed
+
+from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
+
+import torch.nn as nn
+import torch
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from torch.utils.data import DataLoader
+
+import numpy as np
+
+
+class NNModel(nn.Module):
+
+    def __init__(self, h_dim=1024, n_layers=2):
+        super(NNModel, self).__init__()
+        self.layers = nn.ModuleList([nn.Linear(h_dim, h_dim) for i in range(n_layers)])
+        self.cross_entropy_loss = nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        for layer in self.layers:
+            x = layer(x)
+        return self.cross_entropy_loss(x, y)
+
+
+def test_zero_hpz_partition_size_config():
+    config = DeepSpeedZeroConfig(**{"zero_hpz_partition_size": 4})
+    assert config.zero_hpz_partition_size == 4
+
+
+def _assert_no_secondary_tensor_group(model: Module) -> None:
+    for _, param in model.named_parameters():
+        assert param.ds_secondary_tensor is None
+        assert param.ds_zero_param_process_group is None
+
+
+def _check_secondary_tensor_existence(model: Module) -> None:
+    for _, param in model.named_parameters():
+        if param.ds_secondary_tensor is not None:
+            return True
+    return False
+
+
+def _assert_secondary_tensor_size(model: Module) -> None:
+    for name, param in model.named_parameters():
+        assert param.ds_secondary_tensor is not None, f"param {param.ds_id}:{name} does not have secondary tensor"
+        assert param.ds_secondary_tensor.size()[0] % param.ds_tensor.size()[0] == 0
+
+
+#Large sweep along hidden dim, num_layers, and zpg of different sizes
+#Assert when zpg=1 that secondary group and tensors are invalid
+@pytest.mark.sequential
+@pytest.mark.parametrize("h_dim", [1024])
+@pytest.mark.parametrize("n_layers", [9])
+@pytest.mark.parametrize("zpg", [1, 2, 4])
+class TestZeroPPConfigSweep(DistributedTest):
+    world_size = 4
+
+    def test(self, h_dim: int, n_layers: int, zpg: int) -> None:
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_max_reuse_distance": 0,
+                "zero_hpz_partition_size": zpg,
+                "zero_quantized_weights": True,
+                "zero_quantized_gradients": True,
+                "contiguous_gradients": True,
+                "overlap_comm": True,
+            },
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1.
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 1.,
+            }
+        }
+
+        model = NNModel(h_dim, n_layers)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=20,
+                                        hidden_dim=h_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        dist.barrier()
+        if zpg == 1:
+            _assert_no_secondary_tensor_group(model)
+
+        for n, batch in enumerate(data_loader):
+            if n == 0 and zpg != 1:
+                _assert_secondary_tensor_size(model)
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+    def test_eval(self, h_dim: int, n_layers: int, zpg: int) -> None:
+        # in this test case, we are testing that hpz should be enabled when eval mode is on
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_max_reuse_distance": 0,
+                "zero_hpz_partition_size": zpg,
+                "contiguous_gradients": True,
+                "overlap_comm": True,
+            },
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1.
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 1.,
+            }
+        }
+
+        model = NNModel(h_dim, n_layers)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=20,
+                                        hidden_dim=h_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        dist.barrier()
+        if zpg == 1:
+            _assert_no_secondary_tensor_group(model)
+
+        for n, batch in enumerate(data_loader):
+            if zpg != 1:
+                # here we check that the hpz is enabled when the previous iteration does not update the model
+                _assert_secondary_tensor_size(model)
+            with torch.no_grad():
+                loss = model(batch[0], batch[1])
+
+    def test_gradient_accumulation(self, h_dim: int, n_layers: int, zpg: int) -> None:
+        # in this test case, we are testing that hpz should be enabled for the intermediate gradient accumulation steps
+        # In this test, we should disable loss_scale
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": 3,
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_max_reuse_distance": 0,
+                "zero_hpz_partition_size": zpg,
+                "contiguous_gradients": True,
+                "overlap_comm": True,
+            },
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1.
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "loss_scale": 0.,
+            }
+        }
+
+        model = NNModel(h_dim, n_layers)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=20,
+                                        hidden_dim=h_dim,
+                                        device=model.device,
+                                        dtype=torch.float16)
+        dist.barrier()
+        if zpg == 1:
+            _assert_no_secondary_tensor_group(model)
+
+        for n, batch in enumerate(data_loader):
+            if n == 0 and zpg != 1:
+                _assert_secondary_tensor_size(model)
+            # here we cannot assert that secondary tensor does not exist because the gradient is likely overflowed as we use random data
+            if n > 0 and n % 3 != 0 and zpg != 1:
+                # if the previous iteration does not update the model, then the hpz should be enabled
+                assert _check_secondary_tensor_existence(model), f"n={n}"
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+@pytest.mark.nightly
+@pytest.mark.parametrize("model_name", ["gpt2"])
+class TestZeroPPConvergence(DistributedTest):
+    world_size = 4
+
+    def load_and_prepare_data(self, model_name):
+        """Load model, tokenizer and dataset, and prepare data loader."""
+        from datasets import load_dataset
+
+        # Load model and tokenizer
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        tokenizer.pad_token = tokenizer.eos_token
+
+        # Load and tokenize dataset
+        dataset = load_dataset("wikitext", 'wikitext-103-raw-v1', split='train[:1%]').filter(lambda x: x["text"])
+
+        def tokenize_function(examples):
+            # Tokenize and ensure 'labels' are the same as 'input_ids'
+            tokenized_output = tokenizer(examples["text"], padding="max_length", truncation=True, return_tensors='pt')
+            tokenized_output["labels"] = tokenized_output["input_ids"].clone()
+            return tokenized_output
+
+        tokenized_dataset = dataset.map(tokenize_function, batched=True)
+        tokenized_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])
+
+        # Create data loader
+        data_loader = DataLoader(tokenized_dataset, batch_size=1, shuffle=False)
+        return model, data_loader
+
+    def get_loss(self, model, data_loader, config_dict, step=500):
+        """Train the model and calculate average loss."""
+        # Initialize DeepSpeed
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        dist.barrier()
+        model.train()
+
+        # Training loop
+        losses = []
+        for n, batch in enumerate(data_loader):
+            if n >= step:
+                break
+            batch = {k: v.to(model.device) for k, v in batch.items()}
+            outputs = model(**batch)
+            loss = outputs.loss
+            model.backward(loss)
+            model.step()
+            losses.append(loss.item())
+
+        return np.nanmean(losses[-100:])
+
+    def get_config_dict(self, use_quantized_weights=False, use_hpz=False):
+        """Generate the configuration dictionary for DeepSpeed."""
+        config = {
+            "train_micro_batch_size_per_gpu": 1,
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_max_reuse_distance": 0,
+                "contiguous_gradients": True,
+                "overlap_comm": True,
+            },
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-5
+                }
+            },
+            "fp16": {
+                "enabled": True
+            }
+        }
+        if use_quantized_weights:
+            config["zero_optimization"]["zero_quantized_weights"] = True
+        if use_hpz:
+            config["zero_optimization"]["zero_hpz_partition_size"] = self.world_size // 2
+        return config
+
+    def test(self, model_name):
+        torch.manual_seed(0)
+        model, data_loader = self.load_and_prepare_data(model_name)
+        zeropp_loss = self.get_loss(model, data_loader, self.get_config_dict(use_quantized_weights=True, use_hpz=True))
+        model, data_loader = self.load_and_prepare_data(model_name)
+        baseline_loss = self.get_loss(model, data_loader, self.get_config_dict())
+
+        # Output and assert
+        print(f"zeropp_loss={zeropp_loss}, baseline_loss={baseline_loss}")
+        assert zeropp_loss < baseline_loss * 1.1, f"zeropp_loss={zeropp_loss}, baseline_loss={baseline_loss}"

tests/unit/runtime/zero/utils.py

@@ -0,0 +1,16 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+from unit.common import get_master_port
+
+
+def setup_serial_env():
+    # Setup for a serial run
+    os.environ['MASTER_ADDR'] = '127.0.0.1'
+    os.environ['MASTER_PORT'] = get_master_port()
+    os.environ['LOCAL_RANK'] = '0'
+    os.environ['RANK'] = '0'
+    os.environ['WORLD_SIZE'] = '1'