chore: import upstream snapshot with attribution

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