chore: import upstream snapshot with attribution

tests/.coveragerc

@@ -0,0 +1,5 @@
+# .coveragerc to control coverage.py
+[run]
+parallel = True
+sigterm = True
+source = deepspeed

tests/accelerator/ds_config.json

@@ -0,0 +1,19 @@
+{
+  "train_batch_size": 1,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "optimizer": {
+    "type": "Adam",
+    "params": {
+      "lr": 0.00015,
+      "weight_decay": 1e-2
+    }
+  },
+  "fp16": {
+    "enabled": false,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/accelerator/test_ds_init.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import torch
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+
+
+class OneLayerNet(torch.nn.Module):
+
+    def __init__(self, D_in, D_out):
+        """
+        In the constructor we instantiate two nn.Linear modules and assign them as
+        member variables.
+        """
+        super(OneLayerNet, self).__init__()
+        self.linear1 = torch.nn.Linear(D_in, D_out)
+
+    def forward(self, x):
+        """
+        In the forward function we accept a Variable of input data and we must return
+        a Variable of output data. We can use Modules defined in the constructor as
+        well as arbitrary operators on Variables.
+        """
+        h_relu = self.linear1(x).clamp(min=0)
+        y_pred = self.linear1(h_relu)
+        return y_pred
+
+
+def test_literal_device():
+    model = OneLayerNet(128, 128)
+
+    os.environ['RANK'] = '0'
+    os.environ['WORLD_SIZE'] = '1'
+    os.environ['MASTER_ADDR'] = '127.0.0.1'
+    os.environ['MASTER_PORT'] = '8088'
+    os.environ['LOCAL_RANK'] = '0'
+    deepspeed.init_distributed(get_accelerator().communication_backend_name())
+    deepspeed.initialize(model=model, config='ds_config.json')
+    string = get_accelerator().device_name()  #'xpu' or 'cuda'
+    string0 = get_accelerator().device_name(0)  #'xpu:0' or 'cuda:0'
+    string1 = get_accelerator().device_name(1)  #'xpu:1' or 'cuda:1'
+    assert string == 'xpu' or string == 'cuda'
+    assert string0 == 'xpu:0' or string0 == 'cuda:0'
+    assert string1 == 'xpu:1' or string1 == 'cuda:1'

tests/benchmarks/DS4Sci_EvoformerAttention_bench.py

@@ -0,0 +1,107 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+This script is to test the performance of the DS4Sci_EvoformerAttention op.
+To run the script,
+1. Clone the CUTLASS repo. E.g. git clone https://github.com/NVIDIA/cutlass.git
+2. DeepSpeed will detect a local or installed CUTLASS. If needed, set CUTLASS_PATH explicitly.
+3. Run the script. E.g. python DS4Sci_EvoformerAttention_bench.py
+"""
+
+import contextlib
+import torch
+from typing import List
+from torch.nn import functional as F
+from deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention
+from deepspeed.accelerator import get_accelerator
+
+
+def attention_reference(
+        q_input: torch.Tensor,  # [*, Dim_Q, H, C_hid]
+        k_input: torch.Tensor,  # [*, Dim_Q, H, C_hid]
+        v_input: torch.Tensor,  # [*, Dim_Q, H, C_hid]
+        biases: List[torch.Tensor],
+        sm_scale: float) -> torch.Tensor:
+    # Original shape: [*, Dim_Q, H, C_hid] -> Transpose to: [*, H, Dim_Q, C_hid]
+    q = q_input.transpose(-2, -3)
+    k = k_input.transpose(-2, -3)
+    v = v_input.transpose(-2, -3)
+
+    # Now, q, k, v are in shape: [*, H, Dim_Q, C_hid]
+
+    # Transpose k to shape [*, H, C_hid, Dim_Q]
+    k_t = k.transpose(-1, -2)
+
+    # Now, q and k_t are in shapes: [*, H, Dim_Q, C_hid] and [*, H, C_hid, Dim_Q] respectively
+
+    # [*, H, Dim_Q, Dim_Q]
+    a = torch.matmul(q, k_t) * sm_scale
+
+    for b in biases:
+        a += b
+
+    a = F.softmax(a, dim=-1)
+
+    # Now, a is in shape [*, H, Dim_Q, Dim_Q], v is in shape [*, H, Dim_Q, C_hid]
+
+    # Matmul operation results in [*, H, Dim_Q, C_hid]
+    a_v = torch.matmul(a, v)
+
+    # [*, Dim_Q, H, C_hid]
+    o = a_v.transpose(-2, -3)
+
+    return o
+
+
+dtype = torch.float16
+
+N = 256
+heads = 4
+dim = 32
+seq_len = 256
+
+
+@contextlib.contextmanager
+def cuda_timer(res_list):
+    start = get_accelerator().Event(enable_timing=True)
+    end = get_accelerator().Event(enable_timing=True)
+    start.record()
+    yield
+    end.record()
+    get_accelerator().synchronize()
+    res_list.append(start.elapsed_time(end))
+
+
+def benchmark():
+    ours_fw = []
+    ours_bw = []
+    baseline_fw = []
+    baseline_bw = []
+    for batch in range(1, 17):
+        Q = torch.randn(batch, N, seq_len, heads, dim, dtype=dtype, device="cuda", requires_grad=True)
+        K = torch.randn(batch, N, seq_len, heads, dim, dtype=dtype, device="cuda", requires_grad=True)
+        V = torch.randn(batch, N, seq_len, heads, dim, dtype=dtype, device="cuda", requires_grad=True)
+        bias1 = torch.randn(batch, N, 1, 1, seq_len, dtype=dtype, device="cuda", requires_grad=False)
+        bias2 = torch.randn(batch, 1, heads, seq_len, seq_len, dtype=dtype, device="cuda", requires_grad=True)
+        # warm up
+        DS4Sci_EvoformerAttention(Q, K, V, [bias1, bias2])
+        with cuda_timer(ours_fw):
+            out = DS4Sci_EvoformerAttention(Q, K, V, [bias1, bias2])
+        d_out = torch.rand_like(out)
+        with cuda_timer(ours_bw):
+            out.backward(d_out)
+        # warm up
+        attention_reference(Q, K, V, [bias1, bias2], 1 / (dim**0.5))
+        with cuda_timer(baseline_fw):
+            ref_out = attention_reference(Q, K, V, [bias1, bias2], 1 / (dim**0.5))
+        with cuda_timer(baseline_bw):
+            ref_out.backward(d_out)
+
+    print("batch size\tours (FW)\tbaseline (FW)\tours (BW)\tbaseline (BW)")
+    for i in range(len(ours_fw)):
+        print(f"{i+1}\t{ours_fw[i]}\t{baseline_fw[i]}\t{ours_bw[i]}\t{baseline_bw[i]}")
+
+
+benchmark()

tests/benchmarks/flatten_bench.py

@@ -0,0 +1,137 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+#!/usr/bin/env python
+# run the benchmark under timeit (-t), cProfile (-c), line_profiler (-l)
+#
+# usage:
+# ./flatten_bench.py -t
+# ./flatten_bench.py -c
+# kernprof -l flatten_bench.py -l; python -m line_profiler  flatten_bench.py.lprof
+
+import argparse
+
+import gc
+
+import torch
+from torch._utils import _flatten_dense_tensors
+from deepspeed.accelerator import get_accelerator
+from deepspeed.ops.op_builder import UtilsBuilder
+
+from apex_C import flatten as flatten_apex
+
+util_ops = UtilsBuilder().load()
+flatten = util_ops.flatten
+unflatten = util_ops.unflatten
+
+torch.manual_seed(0)
+# emulate a small typical model weights
+x = [
+    torch.rand((512, 512)).to(get_accelerator().device_name()),
+    torch.rand((512, 1024)).to(get_accelerator().device_name()),
+    torch.rand((512, 30000)).to(get_accelerator().device_name())
+]
+t = x * 30
+
+# warm up and check that the same output is produced
+flat_py = _flatten_dense_tensors(t)
+flat_cpp = flatten(t)
+flat_apex = flatten_apex(t)
+#numel = flat_cpp.numel()
+assert torch.eq(flat_py, flat_cpp).all(), "both produce the same tensor"
+assert torch.eq(flat_py, flat_apex).all(), "both produce the same tensor"
+
+TIMES = 1000
+
+
+# the programs being tested
+def py():
+    for i in range(TIMES):
+        flat = _flatten_dense_tensors(t)
+
+
+def cpp():
+    for i in range(TIMES):
+        flat = flatten(t)
+
+
+def apex():
+    for i in range(TIMES):
+        flat = flatten_apex(t)
+
+
+#### cProfile ####
+
+import cProfile
+
+
+def cprofileme():
+    print("--------------- cProfile -----------------")
+    print("py")
+    cProfile.run("py()", sort=-1)
+    gc.collect()
+    get_accelerator().empty_cache()
+    print("cpp")
+    cProfile.run("cpp()", sort=-1)
+    gc.collect()
+    get_accelerator().empty_cache()
+    print("apex")
+    cProfile.run("apex()", sort=-1)
+    gc.collect()
+    get_accelerator().empty_cache()
+
+
+#### timeit ####
+
+import timeit
+
+
+def timeme():
+    print("--------------- timeit -----------------")
+    print(f'py  ={timeit.Timer("py()", globals=globals()).timeit(number=1)}')
+    gc.collect()
+    get_accelerator().empty_cache()
+    print(f'cpp ={timeit.Timer("cpp()", globals=globals()).timeit(number=1)}')
+    gc.collect()
+    get_accelerator().empty_cache()
+    print(f'apex={timeit.Timer("apex()", globals=globals()).timeit(number=1)}')
+    gc.collect()
+    get_accelerator().empty_cache()
+
+
+#### line_profiler ####
+# this one requires a special way to be called
+# pip install line_profiler
+# kernprof -l flatten_bench.py -l; python -m line_profiler  flatten_bench.py.lprof
+
+
+def line_profileme():
+    print("--------------- line_profiler -----------------")
+    print("py")
+    profile(py)()  # noqa: F821 # type: ignore
+    gc.collect()
+    get_accelerator().empty_cache()
+    print("cpp")
+    profile(cpp)()  # noqa: F821 # type: ignore
+    gc.collect()
+    get_accelerator().empty_cache()
+    print("apex")
+    profile(apex)()  # noqa: F821 # type: ignore
+    gc.collect()
+    get_accelerator().empty_cache()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-l", action='store_true')
+    parser.add_argument("-c", action='store_true')
+    parser.add_argument("-t", action='store_true')
+    args = parser.parse_args()
+    if args.l:
+        line_profileme()
+    elif args.c:
+        cprofileme()
+    elif args.t:
+        timeme()

tests/benchmarks/unflatten_bench.py

@@ -0,0 +1,146 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+#!/usr/bin/env python
+
+# run the benchmark under timeit (-t), cProfile (-c), line_profiler (-l)
+#
+# usage:
+# ./unflatten_bench.py -t
+# ./unflatten_bench.py -c
+# kernprof -l unflatten_bench.py -l; python -m line_profiler  unflatten_bench.py.lprof
+
+import argparse
+import gc
+import torch
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+from deepspeed.accelerator import get_accelerator
+from deepspeed.ops.op_builder import UtilsBuilder
+
+from apex_C import flatten as flatten_apex
+from apex_C import unflatten as unflatten_apex
+
+util_ops = UtilsBuilder().load()
+flatten = util_ops.flatten
+unflatten = util_ops.unflatten
+
+torch.manual_seed(0)
+# emulate a small typical model weights
+x = [
+    torch.rand((512, 512)).to(get_accelerator().device_name()),
+    torch.rand((512, 1024)).to(get_accelerator().device_name()),
+    torch.rand((512, 30000)).to(get_accelerator().device_name())
+]
+unflat_t = x * 30
+
+# warm up and check that the same output is produced
+flat_py = _flatten_dense_tensors(unflat_t)
+flat_cpp = flatten(unflat_t)
+flat_apex = flatten_apex(unflat_t)
+#numel = flat_cpp.numel()
+assert torch.eq(flat_py, flat_cpp).all(), "both produce the same tensor"
+assert torch.eq(flat_py, flat_apex).all(), "both produce the same tensor"
+
+flat_t = flat_py
+unflat_py = _unflatten_dense_tensors(flat_py, unflat_t)
+for i in range(len(unflat_t)):
+    assert torch.eq(unflat_t[i], unflat_py[i]).all()
+unflat_cpp = _unflatten_dense_tensors(flat_cpp, unflat_t)
+for i in range(len(unflat_t)):
+    assert torch.eq(unflat_t[i], unflat_cpp[i]).all()
+unflat_apex = _unflatten_dense_tensors(flat_apex, unflat_t)
+for i in range(len(unflat_t)):
+    assert torch.eq(unflat_t[i], unflat_apex[i]).all()
+
+
+# the programs being tested
+def py():
+    for i in range(1000):
+        unflat = _unflatten_dense_tensors(flat_t, unflat_t)
+
+
+def cpp():
+    for i in range(1000):
+        unflat = unflatten(flat_t, unflat_t)
+
+
+def apex():
+    for i in range(1000):
+        unflat = unflatten_apex(flat_t, unflat_t)
+
+
+#### cProfile ####
+
+import cProfile
+
+
+def cprofileme():
+    print("--------------- cProfile -----------------")
+    print("py")
+    cProfile.run("py()", sort=-1)
+    gc.collect()
+    get_accelerator().empty_cache()
+    print("cpp")
+    cProfile.run("cpp()", sort=-1)
+    gc.collect()
+    get_accelerator().empty_cache()
+    print("apex")
+    cProfile.run("apex()", sort=-1)
+    gc.collect()
+    get_accelerator().empty_cache()
+
+
+#### timeit ####
+
+import timeit
+
+
+def timeme():
+    print("--------------- timeit -----------------")
+    print(f'py  ={timeit.Timer("py()", globals=globals()).timeit(number=1)}')
+    gc.collect()
+    get_accelerator().empty_cache()
+    print(f'cpp ={timeit.Timer("cpp()", globals=globals()).timeit(number=1)}')
+    gc.collect()
+    get_accelerator().empty_cache()
+    print(f'apex={timeit.Timer("apex()", globals=globals()).timeit(number=1)}')
+    gc.collect()
+    get_accelerator().empty_cache()
+
+
+#### line_profiler ####
+# this one requires a special way to be called
+# pip install line_profiler
+# kernprof -l unflatten_bench.py -l; python -m line_profiler unflatten_bench.py.lprof
+
+
+def line_profileme():
+    print("--------------- line_profier -----------------")
+    print("py")
+    profile(py)()  # noqa: F821 # type: ignore
+    gc.collect()
+    get_accelerator().empty_cache()
+    print("cpp")
+    profile(cpp)()  # noqa: F821 # type: ignore
+    gc.collect()
+    get_accelerator().empty_cache()
+    print("apex")
+    profile(apex)()  # noqa: F821 # type: ignore
+    gc.collect()
+    get_accelerator().empty_cache()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-l", action='store_true')
+    parser.add_argument("-c", action='store_true')
+    parser.add_argument("-t", action='store_true')
+    args = parser.parse_args()
+    if args.l:
+        line_profileme()
+    elif args.c:
+        cprofileme()
+    elif args.t:
+        timeme()

tests/conftest.py

@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# tests directory-specific settings - this file is run automatically by pytest before any tests are run
+
+import sys
+import pytest
+import os
+from os.path import abspath, dirname
+import torch
+import warnings
+
+# Set this environment variable for the T5 inference unittest(s) (e.g. google/t5-v1_1-small)
+os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
+
+# allow having multiple repository checkouts and not needing to remember to rerun
+# 'pip install -e .[dev]' when switching between checkouts and running tests.
+git_repo_path = abspath(dirname(dirname(__file__)))
+sys.path.insert(1, git_repo_path)
+
+
+def pytest_configure(config):
+    config.option.color = "yes"
+    config.option.durations = 0
+    config.option.durations_min = 1
+    config.option.verbose = True
+
+
+def pytest_addoption(parser):
+    parser.addoption("--torch_ver", default=None, type=str)
+    parser.addoption("--cuda_ver", default=None, type=str)
+
+
+def validate_version(expected, found):
+    version_depth = expected.count('.') + 1
+    found = '.'.join(found.split('.')[:version_depth])
+    return found == expected
+
+
+@pytest.fixture(scope="session", autouse=True)
+def check_environment(pytestconfig):
+    expected_torch_version = pytestconfig.getoption("torch_ver")
+    expected_cuda_version = pytestconfig.getoption("cuda_ver")
+    if expected_torch_version is None:
+        warnings.warn(
+            "Running test without verifying torch version, please provide an expected torch version with --torch_ver")
+    elif not validate_version(expected_torch_version, torch.__version__):
+        pytest.exit(
+            f"expected torch version {expected_torch_version} did not match found torch version {torch.__version__}",
+            returncode=2)
+    if expected_cuda_version is None:
+        warnings.warn(
+            "Running test without verifying cuda version, please provide an expected cuda version with --cuda_ver")
+    elif not validate_version(expected_cuda_version, torch.version.cuda):
+        pytest.exit(
+            f"expected cuda version {expected_cuda_version} did not match found cuda version {torch.version.cuda}",
+            returncode=2)
+
+
+# Override of pytest "runtest" for DistributedTest class
+# This hook is run before the default pytest_runtest_call
+@pytest.hookimpl(tryfirst=True)
+def pytest_runtest_call(item):
+    # We want to use our own launching function for distributed tests
+    if getattr(item.cls, "is_dist_test", False):
+        dist_test_class = item.cls()
+        dist_test_class(item._request)
+        item.runtest = lambda: True  # Dummy function so test is not run twice
+
+
+# We allow DistributedTest to reuse distributed environments. When the last
+# test for a class is run, we want to make sure those distributed environments
+# are destroyed.
+def pytest_runtest_teardown(item, nextitem):
+    if getattr(item.cls, "reuse_dist_env", False) and not nextitem:
+        dist_test_class = item.cls()
+        for num_procs, pool in dist_test_class._pool_cache.items():
+            dist_test_class._close_pool(pool, num_procs, force=True)
+
+
+@pytest.hookimpl(tryfirst=True)
+def pytest_fixture_setup(fixturedef, request):
+    if getattr(fixturedef.func, "is_dist_fixture", False):
+        dist_fixture_class = fixturedef.func()
+        dist_fixture_class(request)

tests/hybrid_engine/hybrid_engine_config.json

@@ -0,0 +1,19 @@
+{
+  "train_batch_size" : 32,
+  "train_micro_batch_size_per_gpu": 2,
+  "steps_per_print": 10,
+  "zero_optimization": {
+    "stage": 0,
+    "offload_param": {
+        "device": "cpu"
+    },
+    "stage3_param_persistence_threshold": 0
+  },
+  "fp16":{
+    "enabled": true,
+    "loss_scale_window": 100
+  },
+  "gradient_clipping": 1.0,
+  "prescale_gradients": false,
+  "wall_clock_breakdown" : false
+}

tests/hybrid_engine/hybrid_engine_test.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from transformers import AutoModelForCausalLM
+import deepspeed
+import argparse
+from deepspeed.accelerator import get_accelerator
+
+deepspeed.runtime.utils.see_memory_usage('pre test', force=True)
+
+model = AutoModelForCausalLM.from_pretrained('facebook/opt-350M').half().to(get_accelerator().device_name())
+parser = argparse.ArgumentParser()
+parser = deepspeed.add_config_arguments(parser)
+args = parser.parse_args()
+
+deepspeed.runtime.utils.see_memory_usage('post test', force=True)
+
+m, _, _, _ = deepspeed.initialize(model=model, args=args, enable_hybrid_engine=True)
+
+m.eval()
+input = torch.ones(1, 16, device='cuda', dtype=torch.long)
+out = m(input)
+
+m.train()
+out = m(input)
+print(out['logits'], out['logits'].norm())

tests/lightning/test_simple.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.strategies import DeepSpeedStrategy
+from torch.utils.data import DataLoader, Dataset
+
+
+class RandomDataset(Dataset):
+
+    def __init__(self, size, length):
+        self.len = length
+        self.data = torch.randn(length, size)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+    def __len__(self):
+        return self.len
+
+
+class BoringModel(LightningModule):
+
+    def __init__(self):
+        super().__init__()
+        self.layer = torch.nn.Linear(32, 2)
+
+    def forward(self, x):
+        return self.layer(x)
+
+    def training_step(self, batch, batch_idx):
+        loss = self(batch).sum()
+        self.log("train_loss", loss)
+        return {"loss": loss}
+
+    def validation_step(self, batch, batch_idx):
+        loss = self(batch).sum()
+        self.log("valid_loss", loss)
+
+    def test_step(self, batch, batch_idx):
+        loss = self(batch).sum()
+        self.log("test_loss", loss)
+
+    def configure_optimizers(self):
+        return torch.optim.SGD(self.layer.parameters(), lr=0.1)
+
+    def train_dataloader(self):
+        return DataLoader(RandomDataset(32, 64), batch_size=2)
+
+    def val_dataloader(self):
+        return DataLoader(RandomDataset(32, 64), batch_size=2)
+
+
+def test_lightning_model():
+    """Test that DeepSpeed works with a simple LightningModule and LightningDataModule."""
+
+    model = BoringModel()
+    trainer = Trainer(strategy=DeepSpeedStrategy(), max_epochs=1, precision=16, accelerator="gpu", devices=1)
+    trainer.fit(model)

tests/model/BingBertSquad/BingBertSquad_run_func_test.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Note: please copy webtext data to "Megatron-LM" folder, before running this script.
+"""
+
+import unittest
+import os
+import re
+from .BingBertSquad_test_common import BaseTestCase
+
+
+def grep_loss_from_file(file_name):
+    loss = 0.0
+
+    with open(file_name, 'r') as f:
+        lines = f.readlines()
+        line_filter = "bert_squad_progress: step="
+        match_number = re.compile(r'loss=([-+]?[0-9]+\.?[0-9]*(?:[Ee][-+]?[0-9]+)?)')
+
+        for line in lines:
+            if line_filter in line:
+                loss = re.findall(match_number, line)
+                loss = float(loss[0])
+
+    if loss == 0.0:
+        print("no loss found in file ", file_name)
+
+    return loss
+
+
+class BingBertSquadFuncTestCase(BaseTestCase):
+
+    def __init__(self, methodName="DeepSpeed function test on BingBertSquad model"):
+        super(BingBertSquadFuncTestCase, self).__init__(methodName)
+
+    def setUp(self):
+        self.save_dir = os.getcwd()
+        new_dir = os.path.dirname(__file__)
+        if new_dir:
+            os.chdir(new_dir)
+
+    def tearDown(self):
+        os.chdir(self.save_dir)
+
+    def test_gpu4_fp16(self):
+        test_config = {
+            "gpus": 4,
+            "deepspeed": False,
+            "json": "deepspeed_bsz24_fp16_config.json",
+            "max_steps": 8,
+            "max_epoch_steps": 4,
+            "other_args": "--fp16 --print_steps 1"
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_gpu4_fp16_zero2(self):
+        test_config = {
+            "gpus": 4,
+            "deepspeed": False,
+            "json": "deepspeed_bsz24_fp16_zero2_config.json",
+            "max_steps": 8,
+            "max_epoch_steps": 4,
+            "other_args": "--fp16 --print_steps 1"
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_gpu1_fp16(self):
+        test_config = {
+            "gpus": 1,
+            "deepspeed": False,
+            "json": "deepspeed_bsz24_fp16_config.json",
+            "max_steps": 8,
+            "max_epoch_steps": 4,
+            "other_args": "--fp16 --print_steps 1"
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_gpu4_fp32(self):
+        test_config = {
+            "gpus": 4,
+            "deepspeed": False,
+            "json": "deepspeed_bsz24_fp32_config.json",
+            "max_steps": 8,
+            "max_epoch_steps": 4,
+            "other_args": "--print_steps 1"
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_gpu1_fp32(self):
+        test_config = {
+            "gpus": 1,
+            "deepspeed": False,
+            "json": "deepspeed_bsz24_fp32_config.json",
+            "max_steps": 8,
+            "max_epoch_steps": 4,
+            "other_args": "--print_steps 1"
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def run_test(self, test_config, r_tol):
+        print("\n")
+        print("{0}: starting......".format(self.id()))
+
+        prefix = "BingBertSquad_func"
+
+        test_config['other_args'] += f" --max_steps {test_config['max_steps']}"
+        test_config['other_args'] += f" --max_steps_per_epoch {test_config['max_epoch_steps']}"
+
+        # baseline run...
+        test_config["deepspeed"] = False
+        base_file = self.gen_output_name(test_config, prefix)
+
+        # skip baseline run if it exists.
+        if not self.has_loss_data(base_file):
+            print("{0}: baseline run.".format(self.id()))
+            self.run_BingBertSquad_test(test_config, base_file)
+        else:
+            print("{0}: baseline exists.".format(self.id()))
+
+        # DeepSpeed run...
+        test_config["deepspeed"] = True
+        print("{0}: DeepSpeed run.".format(self.id()))
+        test_file = self.gen_output_name(test_config, prefix)
+        self.run_BingBertSquad_test(test_config, test_file)
+
+        return self.check_parity(base_file, test_file, r_tol)
+
+    def has_loss_data(self, file_name):
+        has_loss = False
+        if os.path.exists(file_name):
+            loss = grep_loss_from_file(file_name)
+            if loss != 0.0:
+                has_loss = True
+
+        return has_loss
+
+    def check_parity(self, base_file, test_file, r_tol):
+        base_loss = grep_loss_from_file(base_file)
+        test_loss = grep_loss_from_file(test_file)
+
+        print("baseline loss: {0}, test loss: {1}".format(base_loss, test_loss))
+
+        if base_loss == 0.0 or test_loss == 0.0:
+            return False
+
+        if abs((base_loss - test_loss) / base_loss) > r_tol:
+            return False
+
+        return True
+
+
+def suite():
+    suite = unittest.TestSuite()
+    suite.addTest(BingBertSquadFuncTestCase('test_gpu4_fp16'))
+    suite.addTest(BingBertSquadFuncTestCase('test_gpu4_fp16_zero2'))
+    suite.addTest(BingBertSquadFuncTestCase('test_gpu1_fp16'))
+    suite.addTest(BingBertSquadFuncTestCase('test_gpu4_fp32'))
+    suite.addTest(BingBertSquadFuncTestCase('test_gpu1_fp32'))
+    return suite
+
+
+if __name__ == '__main__':
+    runner = unittest.TextTestRunner(failfast=True)
+    runner.run(suite())

tests/model/BingBertSquad/BingBertSquad_test_common.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import unittest
+import subprocess
+import os
+import time
+import shlex
+
+
+class BaseTestCase(unittest.TestCase):
+
+    def __init__(self, methodName="DeepSpeed performance test"):
+        super(BaseTestCase, self).__init__(methodName)
+        self.test_dir = "./test"
+        self.baseline_dir = "./baseline"
+        self.timestr = time.strftime("%Y%m%d-%H%M%S")
+
+    def gen_output_name(self, test_config, prefix):
+        other_args = test_config["other_args"] if "other_args" in test_config else ""
+        zero_args = "_zero" if "zero" in test_config and test_config["zero"] else ""
+        other_args = other_args.strip(' -\\').replace(" ", "").replace("\"", "")
+
+        if other_args:
+            other_args = "_" + other_args
+
+        if test_config["deepspeed"]:
+            file_name = "_gpu{0}_{1}_ds{2}-{3}.log".format(test_config["gpus"], other_args, zero_args, self.timestr)
+            save_dir = self.test_dir
+        else:
+            file_name = "_gpu{0}_{1}.log".format(test_config["gpus"], other_args)
+            save_dir = self.baseline_dir
+
+        return os.path.join(save_dir, prefix + file_name)
+
+    def ensure_directory_exists(self, filename):
+        dirname = os.path.dirname(filename)
+        if not os.path.exists(dirname):
+            os.makedirs(dirname)
+
+    def clean_test_env(self):
+        cmd = shlex.split("dlts_ssh pkill -9 -f /usr/bin/python")
+        print(cmd)
+        subprocess.run(cmd, check=False, executable='/bin/bash')
+        time.sleep(20)
+
+    def run_BingBertSquad_test(self, test_config, output):
+        ds_flag = " -d --deepspeed_config " + test_config["json"] if test_config["deepspeed"] else " "
+        other_args = " " + test_config["other_args"] if "other_args" in test_config else " "
+
+        cmd = "./run_BingBertSquad_sanity.sh -e 1 -g {0} {1} {2}".format(test_config["gpus"], other_args, ds_flag)
+        cmd = shlex.split(cmd)
+        self.ensure_directory_exists(output)
+        with open(output, "w") as f:
+            print(cmd)
+            subprocess.run(cmd, check=False, executable='/bin/bash', stdout=f, stderr=f)

tests/model/BingBertSquad/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .BingBertSquad_run_func_test import BingBertSquadFuncTestCase
+from .BingBertSquad_run_func_test import suite

tests/model/BingBertSquad/deepspeed_bsz24_fp16_config.json

@@ -0,0 +1,20 @@
+{
+  "train_batch_size": 24,
+  "train_micro_batch_size_per_gpu": 3,
+  "steps_per_print": 1,
+  "optimizer": {
+    "type": "Adam",
+    "params": {
+      "lr": 3e-5,
+      "weight_decay": 0.0,
+      "bias_correction": false
+    }
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true
+  },
+  "zero_optimization": {
+      "stage": 1
+  }
+}

tests/model/BingBertSquad/deepspeed_bsz24_fp16_eigenvalue_quantize_config.json

@@ -0,0 +1,49 @@
+{
+  "train_batch_size": 24,
+  "train_micro_batch_size_per_gpu": 3,
+  "steps_per_print": 1,
+  "optimizer": {
+    "type": "Adam",
+    "params": {
+      "lr": 3e-5,
+      "weight_decay": 0.0,
+      "bias_correction": false
+    }
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true
+  },
+  "tensorboard": {
+    "enabled": true,
+    "output_path": "/tmp/eigenvalue_quantize_output",
+    "job_name": "eigenvalue_quantize"
+  },
+  "eigenvalue": {
+      "enabled": true,
+      "verbose": true,
+      "max_iter": 50,
+      "tol": 1e-2,
+      "stability": 0,
+      "gas_boundary_resolution": 1,
+      "model_name": "bert-large"
+  },
+  "quantize_training": {
+    "quantize_bits": {
+      "start_bits": 12,
+      "target_bits": 4
+    },
+    "quantize_type": "symmetric",
+    "quantize_schedule": {
+      "quantize_period": 400,
+      "schedule_offset": 400
+    },
+    "quantize_groups": 16,
+    "fp16_mixed_quantize": {
+      "enabled": true,
+      "quantize_change_ratio": 0.001
+    },
+    "quantize_verbose": true,
+    "quantize_eigenvalue": true
+  }
+}

tests/model/BingBertSquad/deepspeed_bsz24_fp16_zero2_config.json

@@ -0,0 +1,20 @@
+{
+  "train_batch_size": 24,
+  "train_micro_batch_size_per_gpu": 3,
+  "steps_per_print": 1,
+  "optimizer": {
+    "type": "Adam",
+    "params": {
+      "lr": 3e-5,
+      "weight_decay": 0.0,
+      "bias_correction": false
+    }
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true
+  },
+  "zero_optimization": {
+      "stage": 2
+  }
+}

tests/model/BingBertSquad/deepspeed_bsz24_fp32_config.json

@@ -0,0 +1,17 @@
+{
+  "train_batch_size": 24,
+  "train_micro_batch_size_per_gpu": 3,
+  "steps_per_print": 1,
+  "optimizer": {
+    "type": "Adam",
+    "params": {
+      "lr": 3e-5,
+      "weight_decay": 0.0,
+      "bias_correction": false
+    }
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": false
+  }
+}

tests/model/BingBertSquad/run_BingBertSquad.sh

@@ -0,0 +1,171 @@
+#!/bin/bash
+
+usage() {
+  echo """
+Usage: $0 [defined arguments...] [other arguments...]
+
+[defined]
+	-g, --num_gpus          num gpus per node to use
+	-h, --help              this help text
+	-n, --num_nodes         num nodes to use
+	-e, --epochs		        num of training epochs
+	-b, --batch_size 	      training batch size
+  -p, --port              master port for nccl
+
+[other arguments]
+	all undefined arguments will be passed to the user's application
+  """
+}
+
+validate_folder() {
+	dir=$1
+	dir_name=$2
+
+	if [[ -d ${dir} ]]; then
+		echo "Using ${dir_name}: ${dir}"
+	else
+		echo "${dir} folder not found"
+		exit 1
+	fi
+}
+
+remove_folder() {
+	dir=$1
+	dir_name=$2
+
+	if [[ -d ${dir} ]]; then
+		echo "The variable ${dir_name} is set to ${dir} which already exists, so removing and creating a fresh one"
+    rm -rvf ${dir}
+	fi
+}
+
+num_nodes=1
+num_gpus=8
+epochs=2
+batch_size=24
+enable_deepspeed=false
+master_port=$((20000+RANDOM%5000))
+LR=3e-5
+
+while [[ $# -gt 0 ]]
+do
+key="$1"
+case $key in
+    -g|--num_gpus)
+    num_gpus="$2"
+    shift
+    shift
+    ;;
+    -n|--num_nodes)
+    num_nodes="$2"
+    shift
+    shift
+    ;;
+    -e|--epochs)
+    epochs="$2"
+    shift
+    shift
+    ;;
+    -b|--batch_size)
+    batch_size="$2"
+    shift
+    shift
+    ;;
+    -p|--master_port)
+    master_port="$2"
+    shift
+    shift
+    ;;
+    -d|--deepspeed)
+    enable_deepspeed=true
+    shift
+    ;;
+    -h|--help)
+    usage
+    exit 0
+    ;;
+    *) # other arguments
+    other_args="${other_args} $1"
+    shift
+    ;;
+esac
+done
+
+# Validate path to BingBertSquad script
+if [ -z "${BingBertSquad_DIR+x}" ]; then
+  export BingBertSquad_DIR=../../../../DeepSpeedExamples/training/BingBertSquad
+  echo "BingBertSquad_DIR environment variable not set; trying default: ${BingBertSquad_DIR}"
+fi
+validate_folder ${BingBertSquad_DIR} "BingBertSquad_DIR"
+
+# Validate path to processed Squad data
+if [ -z "${SQUAD_DIR+x}" ]; then
+  export SQUAD_DIR=/data/BingBertSquad
+  echo "SQUAD_DIR environment variable not set; trying default: ${SQUAD_DIR}"
+fi
+validate_folder ${SQUAD_DIR} "SQUAD_DIR"
+
+# Set output path
+if [ -z "${OUTPUT_DIR+x}" ]; then
+  export OUTPUT_DIR=/tmp/BingBertSquad-Output
+  echo "OUTPUT_DIR environment variable not set; trying default: ${OUTPUT_DIR}"
+fi
+remove_folder ${OUTPUT_DIR} "OUTPUT_DIR"
+
+echo "num_nodes: ${num_nodes}"
+echo "num_gpus: ${num_gpus}"
+echo "epochs: ${epochs}"
+echo "batch_size: ${batch_size}"
+echo "master_port: ${master_port}"
+echo "deepspeed: ${enable_deepspeed}"
+echo "other_args: ${other_args}"
+
+EFFECTIVE_BATCH_SIZE=${batch_size}
+MAX_GPU_BATCH_SIZE=3
+PER_GPU_BATCH_SIZE=$((EFFECTIVE_BATCH_SIZE/num_gpus))
+if [[ $PER_GPU_BATCH_SIZE -lt $MAX_GPU_BATCH_SIZE ]]; then
+       GRAD_ACCUM_STEPS=1
+else
+       GRAD_ACCUM_STEPS=$((PER_GPU_BATCH_SIZE/MAX_GPU_BATCH_SIZE))
+fi
+
+if [[ ${enable_deepspeed} == true ]]; then
+  BingBertSquad_script=${BingBertSquad_DIR}/nvidia_run_squad_deepspeed.py
+else
+  BingBertSquad_script=${BingBertSquad_DIR}/nvidia_run_squad_baseline.py
+fi
+
+JOB_NAME="BingBertSquad_ds-${enable_deepspeed}_${num_gpus}-gpu"
+
+squad_args="--bert_model bert-large-uncased \
+            --do_train \
+            --do_lower_case \
+            --train_file ${SQUAD_DIR}/train-v1.1.json \
+            --predict_file ${SQUAD_DIR}/dev-v1.1.json \
+            --train_batch_size ${PER_GPU_BATCH_SIZE} \
+            --learning_rate ${LR} \
+            --num_train_epochs ${epochs} \
+            --max_seq_length 384 \
+            --doc_stride 128 \
+            --do_predict \
+            --output_dir ${OUTPUT_DIR} \
+            --gradient_accumulation_steps ${GRAD_ACCUM_STEPS} \
+		      	--job_name ${JOB_NAME} \
+            --model_file ${SQUAD_DIR}/training_state_checkpoint_162.tar
+            "
+
+run_cmd="deepspeed.pt \
+      --num_nodes ${num_nodes} \
+      --num_gpus ${num_gpus} \
+      --master_port ${master_port}
+      ${BingBertSquad_script} ${other_args} ${squad_args}"
+
+# Sanitize input before running eval()
+safe_cmd=$(printf '%q' "$run_cmd")
+
+echo ${safe_cmd}
+eval ${safe_cmd}
+
+set +x
+
+#python ${BingBertSquad_DIR}/evaluate-v1.1.py ${SQUAD_DIR}/dev-v1.1.json ${OUTPUT_DIR}/predictions.json > ${OUTPUT_DIR}/CorrectnessScores.txt

tests/model/BingBertSquad/run_BingBertSquad_sanity.sh

@@ -0,0 +1,168 @@
+#!/bin/bash
+
+usage() {
+  echo """
+Usage: $0 [defined arguments...] [other arguments...]
+
+[defined]
+	-g, --num_gpus          num gpus per node to use
+	-h, --help              this help text
+	-n, --num_nodes         num nodes to use
+	-e, --epochs		        num of training epochs
+	-b, --batch_size 	      training batch size
+  -p, --port              master port for nccl
+
+[other arguments]
+	all undefined arguments will be passed to the user's application
+  """
+}
+
+validate_folder() {
+	dir=$1
+	dir_name=$2
+
+	if [[ -d ${dir} ]]; then
+		echo "Using ${dir_name}: ${dir}"
+	else
+		echo "${dir} folder not found"
+		exit 1
+	fi
+}
+
+remove_folder() {
+	dir=$1
+	dir_name=$2
+
+	if [[ -d ${dir} ]]; then
+		echo "The variable ${dir_name} is set to ${dir} which already exists, so removing and creating a fresh one"
+    rm -rvf ${dir}
+	fi
+}
+
+num_nodes=1
+num_gpus=8
+epochs=2
+batch_size=24
+enable_deepspeed=false
+master_port=$((20000+RANDOM%5000))
+LR=3e-5
+
+while [[ $# -gt 0 ]]
+do
+key="$1"
+case $key in
+    -g|--num_gpus)
+    num_gpus="$2"
+    shift
+    shift
+    ;;
+    -n|--num_nodes)
+    num_nodes="$2"
+    shift
+    shift
+    ;;
+    -e|--epochs)
+    epochs="$2"
+    shift
+    shift
+    ;;
+    -b|--batch_size)
+    batch_size="$2"
+    shift
+    shift
+    ;;
+    -p|--master_port)
+    master_port="$2"
+    shift
+    shift
+    ;;
+    -d|--deepspeed)
+    enable_deepspeed=true
+    echo "Found deespcale flag"
+    shift
+    ;;
+    -h|--help)
+    usage
+    exit 0
+    ;;
+    *) # other arguments
+    other_args="${other_args} $1"
+    shift
+    ;;
+esac
+done
+
+# Validate path to BingBertSquad script
+if [ -z "${BingBertSquad_DIR+x}" ]; then
+  export BingBertSquad_DIR=../../../DeepSpeedExamples/training/BingBertSquad
+  echo "BingBertSquad_DIR environment variable not set; trying default: ${BingBertSquad_DIR}"
+fi
+validate_folder ${BingBertSquad_DIR} "BingBertSquad_DIR"
+
+# Validate path to processed Squad data
+if [ -z "${SQUAD_DIR+x}" ]; then
+  export SQUAD_DIR=/data/BingBertSquad
+  echo "SQUAD_DIR environment variable not set; trying default: ${SQUAD_DIR}"
+fi
+validate_folder ${SQUAD_DIR} "SQUAD_DIR"
+
+# Set output path
+if [ -z "${OUTPUT_DIR+x}" ]; then
+  export OUTPUT_DIR=/tmp/BingBertSquad-Output
+  echo "OUTPUT_DIR environment variable not set; trying default: ${OUTPUT_DIR}"
+fi
+remove_folder ${OUTPUT_DIR} "OUTPUT_DIR"
+
+echo "num_nodes: ${num_nodes}"
+echo "num_gpus: ${num_gpus}"
+echo "epochs: ${epochs}"
+echo "batch_size: ${batch_size}"
+echo "master_port: ${master_port}"
+echo "deepspeed: ${enable_deepspeed}"
+echo "other_args: ${other_args}"
+
+EFFECTIVE_BATCH_SIZE=${batch_size}
+MAX_GPU_BATCH_SIZE=3
+PER_GPU_BATCH_SIZE=$((EFFECTIVE_BATCH_SIZE/num_gpus))
+if [[ $PER_GPU_BATCH_SIZE -lt $MAX_GPU_BATCH_SIZE ]]; then
+       GRAD_ACCUM_STEPS=1
+else
+       GRAD_ACCUM_STEPS=$((PER_GPU_BATCH_SIZE/MAX_GPU_BATCH_SIZE))
+fi
+
+if [[ ${enable_deepspeed} == true ]]; then
+  BingBertSquad_script=${BingBertSquad_DIR}/nvidia_run_squad_deepspeed.py
+else
+  BingBertSquad_script=${BingBertSquad_DIR}/nvidia_run_squad_baseline.py
+fi
+
+JOB_NAME="BingBertSquad_ds-${enable_deepspeed}_${num_gpus}-gpu"
+
+squad_args="--bert_model bert-large-uncased \
+            --do_train \
+            --do_lower_case \
+            --train_file ${SQUAD_DIR}/train-v1.1.json \
+            --predict_file ${SQUAD_DIR}/dev-v1.1.json \
+            --train_batch_size ${PER_GPU_BATCH_SIZE} \
+            --learning_rate ${LR} \
+            --num_train_epochs ${epochs} \
+            --max_seq_length 384 \
+            --doc_stride 128 \
+            --output_dir ${OUTPUT_DIR} \
+            --gradient_accumulation_steps ${GRAD_ACCUM_STEPS} \
+		      	--job_name ${JOB_NAME} \
+            --model_file ${SQUAD_DIR}/training_state_checkpoint_162.tar
+            "
+
+run_cmd="deepspeed.pt \
+      --num_nodes ${num_nodes} \
+      --num_gpus ${num_gpus} \
+      --master_port ${master_port}
+      ${BingBertSquad_script} ${other_args} ${squad_args}"
+
+echo ${run_cmd}
+eval ${run_cmd}
+
+set +x
+
+#python ${BingBertSquad_DIR}/evaluate-v1.1.py ${SQUAD_DIR}/dev-v1.1.json ${OUTPUT_DIR}/predictions.json > ${OUTPUT_DIR}/CorrectnessScores.txt

tests/model/BingBertSquad/run_tests.sh

@@ -0,0 +1,88 @@
+#!/bin/bash
+
+if [[ ! -d logs ]]
+then
+        mkdir logs
+fi
+
+validate_file() {
+  file=$1
+  file_name=$2
+
+  if [[ -f $file ]]; then
+    echo "Using ${file_name}: ${file}"
+  else
+    echo "${file} not found"
+    exit 1
+  fi
+}
+
+validate_folder() {
+	dir=$1
+	dir_name=$2
+
+	if [[ -d ${dir} ]]; then
+		echo "Using ${dir_name}: ${dir}"
+	else
+		echo "${dir} folder not found"
+		exit 1
+	fi
+}
+
+# Validate path to BingBertSquad script
+if [ -z "${BingBertSquad_DIR+x}" ]; then
+  export BingBertSquad_DIR=../../../DeepSpeedExamples/training/BingBertSquad
+  echo "BingBertSquad_DIR environment variable not set; trying default: ${BingBertSquad_DIR}"
+fi
+validate_folder ${BingBertSquad_DIR} "BingBertSquad_DIR"
+
+fp16_config_json=deepspeed_bsz24_fp16_config.json
+validate_file ${fp16_config_json} "fp16_config_json"
+fp32_config_json=deepspeed_bsz24_fp32_config.json
+validate_file ${fp32_config_json} "fp32_config_json"
+
+start_time=`date +"%D %T"`
+echo "---------------begin @ ${start_time}--------------"
+# Note: you may play around with commented parts below (num_gpus and nohup command) for simultaneous runs; just make sure your hardware allocation can support it
+for num_gpus in 8 1 # 4 2
+do
+        #run_cmd="nohup bash run_BingBertSquad.sh -g ${num_gpus} -d --deepspeed_config ${fp16_config_json} --fp16 > logs/deepspeed_fp16_${num_gpus}_`date +"%Y%m%d%H%M%S"`.out 2> logs/deepspeed_fp16_${num_gpus}_`date +"%Y%m%d%H%M%S"`.err &"
+        run_cmd="bash run_BingBertSquad.sh -g ${num_gpus} -d --deepspeed_config ${fp16_config_json} --fp16"
+        start_time=`date +"%D %T"`
+        echo "---------------begin @ ${start_time}--------------"
+        echo ${run_cmd}
+        eval ${run_cmd}
+        end_time=`date +"%D %T"`
+        echo "---------------finish @ ${end_time} --------------"
+
+        #run_cmd="nohup bash run_BingBertSquad.sh -g ${num_gpus} -d --deepspeed_config ${fp32_config_json} > logs/deepspeed_fp32_${num_gpus}_`date +"%Y%m%d%H%M%S"`.out 2> logs/deepspeed_fp32_${num_gpus}_`date +"%Y%m%d%H%M%S"`.err &"
+        run_cmd="bash run_BingBertSquad.sh -g ${num_gpus} -d --deepspeed_config ${fp32_config_json}"
+        start_time=`date +"%D %T"`
+        echo "---------------begin @ ${start_time}--------------"
+        echo ${run_cmd}
+        eval ${run_cmd}
+        end_time=`date +"%D %T"`
+        echo "---------------finish @ ${end_time} --------------"
+
+        #run_cmd="nohup bash run_BingBertSquad.sh -g ${num_gpus} --fp16 > logs/baseline_fp16_${num_gpus}_`date +"%Y%m%d%H%M%S"`.out 2> logs/baseline_fp16_${num_gpus}_`date +"%Y%m%d%H%M%S"`.err &"
+        run_cmd="bash run_BingBertSquad.sh -g ${num_gpus} --fp16"
+        start_time=`date +"%D %T"`
+        echo "---------------begin @ ${start_time}--------------"
+        echo ${run_cmd}
+        eval ${run_cmd}
+        end_time=`date +"%D %T"`
+        echo "---------------finish @ ${end_time} --------------"
+
+        #run_cmd="nohup bash run_BingBertSquad.sh -g ${num_gpus} > logs/baseline_fp32_${num_gpus}_`date +"%Y%m%d%H%M%S"`.out 2> logs/baseline_fp32_${num_gpus}_`date +"%Y%m%d%H%M%S"`.err &"
+        run_cmd="bash run_BingBertSquad.sh -g ${num_gpus}"
+        start_time=`date +"%D %T"`
+        echo "---------------begin @ ${start_time}--------------"
+        echo ${run_cmd}
+        eval ${run_cmd}
+        end_time=`date +"%D %T"`
+        echo "---------------finish @ ${end_time} --------------"
+done
+end_time=`date +"%D %T"`
+echo "---------------finish @ ${end_time} --------------"
+
+set +x

tests/model/BingBertSquad/test_e2e_squad.py

@@ -0,0 +1,125 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import subprocess as sp
+import os
+from math import isclose
+import sys
+import pytest
+import json
+
+sys.path.append("../../../DeepSpeedExamples/training/BingBertSquad")
+import evaluate as eval
+
+squad_dir = "/data/BingBertSquad"
+base_dir = "../../../DeepSpeedExamples/training/BingBertSquad"
+
+script_file_name = "run_squad_deepspeed.sh"
+model_file_name = "training_state_checkpoint_162.tar"
+eval_file_name = "dev-v1.1.json"
+pred_file_name = "predictions.json"
+
+num_gpus = "4"
+timeout_sec = 5 * 60 * 60  # 5 hours
+
+eval_version = "1.1"
+
+
+def create_config_file(tmpdir, zeroenabled=False):
+    config_dict = {
+        "train_batch_size": 24,
+        "train_micro_batch_size_per_gpu": 6,
+        "steps_per_print": 10,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 3e-5,
+                "weight_decay": 0.0,
+                "bias_correction": False
+            }
+        },
+        "gradient_clipping": 1.0,
+        "fp16": {
+            "enabled": True
+        }
+    }
+    config_dict["zero_optimization"] = zeroenabled
+
+    config_path = os.path.join(tmpdir, 'temp_config.json')
+    with open(config_path, 'w') as fd:
+        json.dump(config_dict, fd)
+    return config_path
+
+
+def test_e2e_squad_deepspeed_base(tmpdir):
+    config_file = create_config_file(tmpdir)
+
+    # base run results => {"exact_match": 83.9829706717124, "f1": 90.71138132004097}
+    expected_exact_match = 83.98
+    expected_f1 = 90.71
+
+    model_file = os.path.join(squad_dir, model_file_name)
+    eval_file = os.path.join(squad_dir, eval_file_name)
+
+    output_dir = os.path.join(tmpdir, "output")
+    pred_file = os.path.join(output_dir, pred_file_name)
+
+    proc = sp.Popen(["bash", script_file_name, num_gpus, model_file, squad_dir, output_dir, config_file], cwd=base_dir)
+
+    try:
+        proc.communicate(timeout=timeout_sec)
+
+        if os.path.exists(pred_file):
+            eval_result = eval.evaluate(eval_version, eval_file, pred_file)
+
+            print("evaluation result: ", json.dumps(eval_result))
+
+            assert isclose(eval_result["exact_match"], expected_exact_match, abs_tol=1e-2)
+            assert isclose(eval_result["f1"], expected_f1, abs_tol=1e-2)
+
+        else:
+            pytest.fail("Error: Run Failed")
+
+    except sp.TimeoutExpired:
+        proc.kill()
+        pytest.fail("Error: Timeout")
+    except sp.CalledProcessError:
+        pytest.fail("Error: Run Failed")
+
+
+def test_e2e_squad_deepspeed_zero(tmpdir):
+    config_file = create_config_file(tmpdir, True)
+
+    # base run results => {"exact_match": 84.1438032166509, "f1": 90.89776136505441}
+    expected_exact_match = 84.14
+    expected_f1 = 90.89
+
+    model_file = os.path.join(squad_dir, model_file_name)
+    eval_file = os.path.join(squad_dir, eval_file_name)
+
+    output_dir = os.path.join(tmpdir, "output")
+    pred_file = os.path.join(output_dir, pred_file_name)
+
+    proc = sp.Popen(["bash", script_file_name, num_gpus, model_file, squad_dir, output_dir, config_file], cwd=base_dir)
+
+    try:
+        proc.communicate(timeout=timeout_sec)
+
+        if os.path.exists(pred_file):
+            eval_result = eval.evaluate(eval_version, eval_file, pred_file)
+
+            print("evaluation result: ", json.dumps(eval_result))
+
+            assert isclose(eval_result["exact_match"], expected_exact_match, abs_tol=1e-2)
+            assert isclose(eval_result["f1"], expected_f1, abs_tol=1e-2)
+
+        else:
+            pytest.fail("Error: Run Failed")
+
+    except sp.TimeoutExpired:
+        proc.kill()
+        pytest.fail("Error: Timeout")
+    except sp.CalledProcessError:
+        pytest.fail("Error: Run Failed")

tests/model/Megatron_GPT2/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Note: please copy webtext data to "Megatron-LM" folder, before running this script.
+"""
+
+from .run_func_test import GPT2FuncTestCase
+from .run_checkpoint_test import GPT2CheckpointTestCase, checkpoint_suite
+from .run_func_test import suite

tests/model/Megatron_GPT2/ds_config_func_bs4_zero1.json

@@ -0,0 +1,16 @@
+{
+  "train_batch_size": 4,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 1
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_bs4_zero2.json

@@ -0,0 +1,19 @@
+{
+  "train_batch_size": 4,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 2,
+    "reduce_bucket_size": 7000000,
+    "allgather_bucket_size": 7000000,
+    "reduce_scatter": true
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_bs4_zero2_offload.json

@@ -0,0 +1,20 @@
+{
+  "train_batch_size": 4,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 2,
+    "reduce_bucket_size": 7000000,
+    "allgather_bucket_size": 7000000,
+    "reduce_scatter": true,
+    "cpu_offload": true
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_bs8_no_zero.json

@@ -0,0 +1,16 @@
+{
+  "train_batch_size": 8,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 0
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_bs8_zero0_gas3.json

@@ -0,0 +1,23 @@
+{
+  "train_micro_batch_size_per_gpu": 8,
+  "gradient_accumulation_steps": 3,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 0,
+    "reduce_bucket_size": 7000000,
+    "allgather_bucket_size": 7000000,
+    "reduce_scatter": true
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "activation_checkpointing": {
+    "partition_activations": true,
+    "contiguous_memory_optimization": true
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_bs8_zero1.json

@@ -0,0 +1,16 @@
+{
+  "train_batch_size": 8,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 1
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_bs8_zero2.json

@@ -0,0 +1,23 @@
+{
+  "train_batch_size": 8,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 2,
+    "reduce_bucket_size": 7000000,
+    "allgather_bucket_size": 7000000,
+    "reduce_scatter": true
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "activation_checkpointing": {
+    "partition_activations": true,
+    "contiguous_memory_optimization": true
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_bs8_zero2_gas3.json

@@ -0,0 +1,23 @@
+{
+  "train_micro_batch_size_per_gpu": 8,
+  "gradient_accumulation_steps": 3,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 2,
+    "reduce_bucket_size": 7000000,
+    "allgather_bucket_size": 7000000,
+    "reduce_scatter": true
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "activation_checkpointing": {
+    "partition_activations": true,
+    "contiguous_memory_optimization": true
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_bs8_zero2_offload.json

@@ -0,0 +1,24 @@
+{
+  "train_batch_size": 8,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 2,
+    "reduce_bucket_size": 7000000,
+    "allgather_bucket_size": 7000000,
+    "reduce_scatter": true,
+    "cpu_offload": true
+  },
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "activation_checkpointing": {
+    "partition_activations": true,
+    "contiguous_memory_optimization": true
+  }
+}

tests/model/Megatron_GPT2/ds_config_func_scheduler.json

@@ -0,0 +1,24 @@
+{
+  "train_batch_size": 4,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 2
+  },
+  "gradient_clipping": 1.0,
+  "scheduler": {
+    "type": "WarmupLR",
+    "params": {
+      "warmup_min_lr": 0,
+      "warmup_max_lr": 0.001,
+      "warmup_num_steps": 10
+    }
+  },
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_config_perf_bs16.json

@@ -0,0 +1,17 @@
+{
+  "train_batch_size": 16,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 1
+  },
+  "disable_allgather": true,
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_config_perf_bs32.json

@@ -0,0 +1,17 @@
+{
+  "train_batch_size": 32,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 1
+  },
+  "disable_allgather": true,
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_config_perf_bs8.json

@@ -0,0 +1,17 @@
+{
+  "train_batch_size": 8,
+  "gradient_accumulation_steps": 1,
+  "steps_per_print": 1,
+  "zero_optimization": {
+    "stage": 1
+  },
+  "disable_allgather": true,
+  "gradient_clipping": 1.0,
+  "fp16": {
+    "enabled": true,
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  }
+}

tests/model/Megatron_GPT2/ds_gpt2_test.sh

@@ -0,0 +1,103 @@
+#! /bin/bash
+
+helpFunction()
+{
+    echo ""
+    echo "Usage: $0 -m model-parallelism -g gpu-per-node -n node# -b batch-size -s steps -l layers -h hidden_size -q seq_length -e heads -c ckpt_num_layers -p [-d]"
+    echo -e "\t-m model parallelism"
+    echo -e "\t-g gpus per node"
+    echo -e "\t-n node count"
+    echo -e "\t-b batch size"
+    echo -e "\t-s training steps"
+    echo -e "\t-l layers"
+    echo -e "\t-h hidden size"
+    echo -e "\t-q sequence length"
+    echo -e "\t-e attention heads"
+    echo -e "\t-c checkpoint num_layers"
+    echo -e "\t-o other args"
+    echo -e "\t-d DeepSpeed config json file"
+    echo -e "\t-z Enable Zero optimization"
+    echo -e "\t-p DeepSpeed master port"
+    exit 1
+}
+
+layers=2
+hidden_size=128
+seq_length=1024
+ckpt_num_layers=1
+other_args=""
+ds_opt=""
+zero_opt=""
+master_port=29600
+
+script_path=$(realpath $0)
+script_dir=$(dirname $script_path)
+
+while getopts "m:g:n:b:s:l:h:q:e:c:o:d:z" opt
+do
+    case "$opt" in
+        m ) mp="$OPTARG" ;;
+        g ) gpus="$OPTARG" ;;
+        n ) nodes="$OPTARG" ;;
+        b ) bs="$OPTARG" ;;
+        s ) steps="$OPTARG" ;;
+        l ) layers="$OPTARG" ;;
+        h ) hidden_size="$OPTARG" ;;
+        q ) seq_length="$OPTARG" ;;
+        e ) heads="$OPTARG" ;;
+        c ) ckpt_num_layers="$OPTARG" ;;
+        p ) master_port="$OPTARG" ;;
+        o ) other_args="$OPTARG" ;;
+        d ) ds_opt="--deepspeed --deepspeed_config $script_dir/$OPTARG" ;;
+        z ) zero_opt="--zero_optimization" ;;
+        ? ) helpFunction ;;
+    esac
+done
+
+# Print helpFunction in case parameters are empty
+if [ -z "$mp" ] || [ -z "$gpus" ] || [ -z "$nodes" ] || [ -z "$bs" ] || [ -z "$steps" ]
+then
+    echo "Some or all of the parameters are empty";
+    helpFunction
+fi
+
+# Change for multinode config
+MASTER_ADDR=localhost
+MASTER_PORT=6000
+
+gpt_options=" \
+       --model-parallel-size ${mp} \
+       --num-layers ${layers} \
+       --hidden-size ${hidden_size} \
+       --num-attention-heads ${heads} \
+       --batch-size ${bs} \
+       --seq-length ${seq_length} \
+       --max-position-embeddings ${seq_length} \
+       --train-iters ${steps} \
+       --train-data webtext \
+       --lazy-loader \
+       --tokenizer-type GPT2BPETokenizer \
+       --split 949,50,1 \
+       --distributed-backend nccl \
+       --lr 0.00015 \
+       --no-load-optim \
+       --lr-decay-style cosine \
+       --weight-decay 1e-2 \
+       --clip-grad 1.0 \
+       --warmup .01 \
+       --checkpoint-activations \
+       --checkpoint-num-layers ${ckpt_num_layers} \
+       --fp16 \
+       --cache-dir /tmp/cache_dir \
+       --log-interval 1 \
+       ${other_args} \
+       ${ds_opt} \
+       ${zero_opt} \
+"
+
+work_dir="../../../DeepSpeedExamples/Megatron-LM/"
+run_cmd="(cd ${work_dir} && deepspeed --master_port ${master_port} --num_nodes $nodes --num_gpus $gpus pretrain_gpt2.py ${gpt_options})"
+echo ${run_cmd}
+eval ${run_cmd}
+
+set +x

tests/model/Megatron_GPT2/run_checkpoint_test.py

@@ -0,0 +1,573 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Note: please copy webtext data to "Megatron-LM" folder, before running this script.
+"""
+
+import unittest
+import subprocess
+import os
+import re
+import shlex
+from .test_common import BaseTestCase
+
+LAYERS = 2
+HIDDEN_SIZE = 128
+ATTN_HEADS = 8
+
+
+def remove_file(test_id, filename):
+    cmd = shlex.split(f"if [ -f {filename} ] ; then rm -v {filename}; fi")
+    print(f"{test_id} cmd: {cmd}")
+    subprocess.run(cmd, check=False, executable='/bin/bash')
+
+
+def grep_loss_from_file(file_name):
+    loss = 0.0
+
+    with open(file_name, 'r') as f:
+        lines = f.readlines()
+        line_filter = "validation loss at the end of training for test data | LM loss:"
+        match_number = re.compile(r'LM loss: ([-+]?[0-9]+\.?[0-9]*(?:[Ee][-+]?[0-9]+)?)')
+
+        for line in lines:
+            if line_filter in line:
+                loss = re.findall(match_number, line)
+                loss = float(loss[0])
+
+    if loss == 0.0:
+        print("no loss found in file ", file_name)
+
+    return loss
+
+
+class GPT2CheckpointTestCase(BaseTestCase):
+
+    def __init__(self, methodName="DeepSpeed function test on GPT2 model"):
+        super(GPT2CheckpointTestCase, self).__init__(methodName)
+
+    def setUp(self):
+        self.save_dir = os.getcwd()
+        new_dir = os.path.dirname(__file__)
+        if new_dir:
+            os.chdir(new_dir)
+
+    def tearDown(self):
+        os.chdir(self.save_dir)
+
+    def test_mp2_gpu4_node1_with_zero1(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero1",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu8_w_zero1",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero1.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_node1_with_zero2(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu8_w_zero2",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_node1_with_zero2_offload(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2_offload",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu8_w_zero2_offload",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_load_gpu1_node1_with_zero1(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "load_gpus": 1,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero1",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp1_gpu2_gpu1_w_zero1",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero1.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_load_gpu4_node1_with_zero1(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "load_gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero1",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp1_gpu2_gpu4_w_zero1",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero1.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_load_gpu1_node1_with_zero2(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "load_gpus": 1,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp1_gpu2_gpu1_w_zero2",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_load_gpu1_node1_with_zero2_offload(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "load_gpus": 1,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2_offload",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp1_gpu2_gpu1_w_zero2_offload",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_load_gpu4_node1_with_zero2(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "load_gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp1_gpu2_gpu4_w_zero2",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_load_gpu4_node1_with_zero2_offload(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "load_gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2_offload",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp1_gpu2_gpu4_w_zero2_offload",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_load_gpu2_node1_with_zero1(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "load_gpus": 2,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero1",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu4_gpu2_w_zero1",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero1.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu2_load_gpu4_node1_with_zero1(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 2,
+            "load_gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero1",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu2_gpu4_w_zero1",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero1.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_load_gpu2_node1_with_zero2(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "load_gpus": 2,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu4_gpu2_w_zero2",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_load_gpu2_node1_with_zero2_offload(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "load_gpus": 2,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2_offload",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu4_gpu2_w_zero2_offload",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu2_load_gpu4_node1_with_zero2(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 2,
+            "load_gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu2_gpu4_w_zero2",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu2_load_gpu4_node1_with_zero2_offload(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 2,
+            "load_gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "tag": "ds_zero2_offload",
+            "zero": True,
+            "other_args": "",
+            "checkpoint_name": "ckpt_mp2_gpu2_gpu4_w_zero2_offload",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_node1_without_zero(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1100,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": 256,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "zero": False,
+            "other_args": "",
+            "tag": "ds_without_zero",
+            "checkpoint_name": "ckpt_mp4_gpu16_wo_zero",
+            "checkpoint_interval": 1000,
+            "json": "ds_config_func_bs8_no_zero.json",
+        }
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def gen_name(self, test_config, prefix):
+        save_dir = "checkpoint_test_logs"
+        tag = test_config["tag"]
+        checkpoint_name = test_config["checkpoint_name"]
+        file_name = f"_{tag}_{checkpoint_name}.log"
+        return os.path.join(save_dir, prefix + file_name)
+
+    def run_test(self, test_config, r_tol):
+        print("\n")
+
+        print("{0}: starting......".format(self.id()))
+
+        # Cache save and load gpu counts
+        save_gpus = test_config["gpus"]
+        if "load_gpus" in test_config:
+            load_gpus = test_config["load_gpus"]
+            del test_config["load_gpus"]
+        else:
+            load_gpus = test_config["gpus"]
+
+        # save to current directory.
+        checkpoint_folder = test_config["checkpoint_name"]
+        checkpoint_interval = test_config["checkpoint_interval"]
+        checkpoint_name = test_config["checkpoint_name"]
+        #---------------remove old checkpoint---------------#
+        try:
+            cmd = shlex.split(f"rm -rf {checkpoint_name}")
+            print(f"{self.id()} cmd: {cmd}")
+            subprocess.run(cmd, check=False, executable='/bin/bash')
+        except Exception:
+            print("No old checkpoint")
+
+        if "cpu_optimizer" in test_config and test_config["cpu_optimizer"]:
+            cpu_optimizer_flag = " --cpu-optimizer"
+        else:
+            cpu_optimizer_flag = ""
+
+        #-----------------Saving Checkpoint-----------------#
+        # building checkpoint arguments
+        test_config[
+            "other_args"] = f"\"--save {checkpoint_folder} --save-interval {checkpoint_interval} {cpu_optimizer_flag}\""
+
+        prefix = "gpt2_saving_checkpoint"
+
+        # create checkpoint run...
+        base_file = self.gen_name(test_config, prefix)
+
+        # remove previous test log
+        try:
+            cmd = shlex.split(f"rm {base_file}")
+            subprocess.run(cmd, check=False, executable='/bin/bash')
+        except Exception:
+            print(f"{self.id()} No old logs")
+
+        print("{0}: Run for saving checkpoint".format(self.id()))
+        self.run_gpt2_test(test_config, base_file)
+
+        #-----------------Loading Checkpoint-----------------#
+
+        # building checkpoint arguments
+        test_config["other_args"] = f"\"--load {checkpoint_folder} {cpu_optimizer_flag} \""
+
+        # set checkpoint load iteration
+        try:
+            cmd = shlex.split(f"echo {checkpoint_interval} > {checkpoint_name}/latest_checkpointed_iteration.txt")
+            print(f"{self.id()} running cmd: {cmd}")
+            subprocess.run(cmd, check=False, executable='/bin/bash')
+        except Exception:
+            print(f"{self.id()} Failed to update the checkpoint iteration file")
+            return False
+
+        prefix = "gpt2_loading_checkpoint"
+
+        # set load gpus
+        test_config["gpus"] = load_gpus
+
+        print("{0}: Second run loading checkpoint and continuing.".format(self.id()))
+        test_file = self.gen_name(test_config, prefix)
+
+        # remove previous test log
+        try:
+            cmd = shlex.split(f"rm {test_file}")
+            subprocess.run(cmd, check=False, executable='/bin/bash')
+        except Exception:
+            print(f"{self.id()} no previous logs for")
+        self.run_gpt2_test(test_config, test_file)
+        return self.check_parity(base_file, test_file, r_tol)
+
+    def has_loss_data(self, file_name):
+        has_loss = False
+        if os.path.exists(file_name):
+            loss = grep_loss_from_file(file_name)
+            if loss != 0.0:
+                has_loss = True
+
+        return has_loss
+
+    def check_parity(self, base_file, test_file, r_tol):
+        base_loss = grep_loss_from_file(base_file)
+        test_loss = grep_loss_from_file(test_file)
+
+        print("baseline loss: {0}, test loss: {1}".format(base_loss, test_loss))
+
+        if base_loss == 0.0 or test_loss == 0.0:
+            return False
+
+        if abs((base_loss - test_loss) / base_loss) > r_tol:
+            return False
+
+        return True
+
+
+def checkpoint_suite():
+    suite = unittest.TestSuite()
+
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu4_node1_with_zero1'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu4_node1_with_zero2'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu4_node1_with_zero2_offload'))
+
+    # Shrink DP
+    suite.addTest(GPT2CheckpointTestCase('test_mp1_gpu2_load_gpu1_node1_with_zero1'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp1_gpu2_load_gpu1_node1_with_zero2'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp1_gpu2_load_gpu1_node1_with_zero2_offload'))
+
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu4_load_gpu2_node1_with_zero1'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu4_load_gpu2_node1_with_zero2'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu4_load_gpu2_node1_with_zero2_offload'))
+
+    # Expand DP
+    suite.addTest(GPT2CheckpointTestCase('test_mp1_gpu2_load_gpu4_node1_with_zero1'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp1_gpu2_load_gpu4_node1_with_zero2'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp1_gpu2_load_gpu4_node1_with_zero2_offload'))
+
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu2_load_gpu4_node1_with_zero1'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu2_load_gpu4_node1_with_zero2'))
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu2_load_gpu4_node1_with_zero2_offload'))
+
+    suite.addTest(GPT2CheckpointTestCase('test_mp2_gpu4_node1_without_zero'))
+
+    return suite
+
+
+if __name__ == '__main__':
+    runner = unittest.TextTestRunner(failfast=True)
+    runner.run(checkpoint_suite())

tests/model/Megatron_GPT2/run_func_test.py

@@ -0,0 +1,603 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Note: please copy webtext data to "Megatron-LM" folder, before running this script.
+"""
+
+import unittest
+import os
+import re
+from .test_common import BaseTestCase
+
+LAYERS = 2
+HIDDEN_SIZE = 128
+ATTN_HEADS = 8
+SEQ_LEN = 64
+MASTER_PORT = 29700
+
+
+def grep_loss_from_file(file_name):
+    loss = 0.0
+    print(f'grepping {file_name}')
+    with open(file_name, 'r') as f:
+        lines = f.readlines()
+        line_filter = "validation loss at the end of training for test data | LM loss:"
+        match_number = re.compile(r'LM loss: ([-+]?[0-9]+\.?[0-9]*(?:[Ee][-+]?[0-9]+)?)')
+
+        for line in lines:
+            if line_filter in line:
+                loss = re.findall(match_number, line)
+                loss = float(loss[0])
+
+    if loss == 0.0:
+        print("no loss found in file ", file_name)
+
+    return loss
+
+
+class GPT2FuncTestCase(BaseTestCase):
+
+    def __init__(self, methodName="DeepSpeed function test on GPT2 model"):
+        super(GPT2FuncTestCase, self).__init__(methodName)
+
+    def setUp(self):
+        self.save_dir = os.getcwd()
+        new_dir = os.path.dirname(__file__)
+        if new_dir:
+            os.chdir(new_dir)
+
+    def tearDown(self):
+        os.chdir(self.save_dir)
+
+    def test_mp1_gpu2_node1_fp16(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_no_zero.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu1_node1_zero1(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 1,
+            "nodes": 1,
+            "bs": 4,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs4_zero1.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_node1_zero1(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero1.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_node1_zero1(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero1.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp4_gpu4_node1_zero1(self):
+        test_config = {
+            "mp": 4,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero1.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu1_node1_zero2(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 1,
+            "nodes": 1,
+            "bs": 4,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs4_zero2.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_node1_zero2(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_node1_zero2(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2.json",
+        }
+
+        basic_run_config = test_config
+        succ = self.run_test(basic_run_config, 0.01)
+        self.assertTrue(succ)
+
+        partition_activation_config = test_config
+        succ = self.run_partition_activations_test(partition_activation_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp4_gpu4_node1_zero2(self):
+        test_config = {
+            "mp": 4,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2.json",
+        }
+
+        basic_run_config = test_config
+        succ = self.run_test(basic_run_config, 0.01)
+        self.assertTrue(succ)
+
+        partition_activation_config = test_config
+        succ = self.run_partition_activations_test(partition_activation_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu1_node1_zero2_ds_offload(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 1,
+            "nodes": 1,
+            "bs": 4,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs4_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+        succ = self.run_test(test_config, 0.02)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_node1_zero2_ds_offload(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+        succ = self.run_test(test_config, 0.02)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_node1_zero2_gas(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": True,
+            "json": "ds_config_func_bs8_zero2_gas3.json",
+            "baseline": "ds_config_func_bs8_zero0_gas3.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+        succ = self.run_partition_activations_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_node1_zero2_ds_offload(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+
+        basic_run_config = test_config
+        succ = self.run_test(basic_run_config, 0.02)
+        self.assertTrue(succ)
+
+        partition_activation_config = test_config
+        succ = self.run_partition_activations_test(partition_activation_config, 0.02)
+        self.assertTrue(succ)
+
+    def test_mp4_gpu4_node1_zero2_ds_offload(self):
+        test_config = {
+            "mp": 4,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+        }
+
+        basic_run_config = test_config
+        succ = self.run_test(basic_run_config, 0.02)
+        self.assertTrue(succ)
+
+        partition_activation_config = test_config
+        succ = self.run_partition_activations_test(partition_activation_config, 0.02)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu1_node1_zero2_torch_offload(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 1,
+            "nodes": 1,
+            "bs": 4,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs4_zero2_offload.json",
+            "cpu_optimizer": True,
+            "test_torch_offload": True,
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp1_gpu2_node1_zero2_torch_offload(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 2,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+            "test_torch_offload": True,
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp2_gpu4_node1_zero2_torch_offload(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+            "test_torch_offload": True,
+        }
+
+        basic_run_config = test_config
+        succ = self.run_test(basic_run_config, 0.01)
+        self.assertTrue(succ)
+
+        partition_activation_config = test_config
+        succ = self.run_partition_activations_test(partition_activation_config, 0.01)
+        self.assertTrue(succ)
+
+    def test_mp4_gpu4_node1_zero2_torch_offload(self):
+        test_config = {
+            "mp": 4,
+            "gpus": 4,
+            "nodes": 1,
+            "bs": 8,
+            "steps": 1000,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_bs8_zero2_offload.json",
+            "cpu_optimizer": True,
+            "test_torch_offload": True,
+        }
+
+        basic_run_config = test_config
+        succ = self.run_test(basic_run_config, 0.01)
+        self.assertTrue(succ)
+
+        partition_activation_config = test_config
+        succ = self.run_partition_activations_test(partition_activation_config, 0.01)
+
+    def test_optimizer_scheduler(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 1,
+            "nodes": 1,
+            "bs": 4,
+            "steps": 20,
+            "layers": LAYERS,
+            "hidden_size": HIDDEN_SIZE,
+            "seq_length": SEQ_LEN,
+            "heads": ATTN_HEADS,
+            "deepspeed": False,
+            "json": "ds_config_func_scheduler.json",
+        }
+
+        succ = self.run_test(test_config, 0.01)
+        # assure no crash.
+        self.assertTrue(True)
+
+    def run_partition_activations_test(self, test_config, r_tol):
+        print("\n")
+        print("{0}: starting......".format(self.id()))
+
+        baseline_prefix = "gpt2_func_"
+        prefix = "gpt2_partition_activation_"
+
+        deepspeed_config = test_config["json"]
+        baseline_deepspeed_config = False
+        cpu_optimizer_flag = self.gen_cpu_optimizer_flag(test_config, True)
+
+        # baseline run...
+        # turnoff deepspeed if baseline deepspeed config
+        # is not provided
+        if not "baseline" in test_config:
+            test_config["deepspeed"] = False
+        else:
+            test_config["json"] = test_config["baseline"]
+            baseline_prefix += test_config["json"][0:-5]
+            baseline_deepspeed_config = True
+
+        test_config["other_args"] = f"\"{cpu_optimizer_flag}\""
+        base_file = self.gen_output_name(test_config, baseline_prefix, baseline_config=baseline_deepspeed_config)
+
+        # skip baseline run if it exists.
+        if not self.has_loss_data(base_file):
+            print("{0}: baseline run.".format(self.id()))
+            self.run_gpt2_test(test_config, base_file)
+        else:
+            print("{0}: baseline exists.".format(self.id()))
+
+        # DeepSpeed run...
+        test_config["deepspeed"] = True
+        cpu_optimizer_flag = self.gen_cpu_optimizer_flag(test_config, False)
+        test_config["other_args"] = f"\"--deepspeed-activation-checkpointing {cpu_optimizer_flag}\""
+        test_config["json"] = deepspeed_config
+
+        print("{0}: DeepSpeed run.".format(self.id()))
+        test_file = self.gen_output_name(test_config, prefix)
+        self.run_gpt2_test(test_config, test_file)
+
+        return self.check_parity(base_file, test_file, r_tol)
+
+    def run_test(self, test_config, r_tol):
+        print("\n")
+        print("{0}: starting......".format(self.id()))
+
+        prefix = "gpt2_func"
+        baseline_prefix = prefix
+
+        deepspeed_config = test_config["json"]
+        baseline_deepspeed_config = False
+        cpu_optimizer_flag = self.gen_cpu_optimizer_flag(test_config, True)
+
+        # baseline run...
+        # turn off deepspeed if a baseline deepspeed config
+        # is not provided
+        if not "baseline" in test_config:
+            test_config["deepspeed"] = False
+        else:
+            test_config["json"] = test_config["baseline"]
+            baseline_prefix = prefix + test_config["json"][0:-5]
+            baseline_deepspeed_config = True
+
+        test_config["other_args"] = f"\"{cpu_optimizer_flag}\""
+
+        # baseline run...
+        base_file = self.gen_output_name(test_config, baseline_prefix, baseline_config=baseline_deepspeed_config)
+
+        # skip baseline run if it exists.
+        if not self.has_loss_data(base_file):
+            print("{0}: baseline run.".format(self.id()))
+            self.run_gpt2_test(test_config, base_file)
+        else:
+            print("{0}: baseline exists.".format(self.id()))
+
+        # DeepSpeed run...
+        test_config["deepspeed"] = True
+        cpu_optimizer_flag = self.gen_cpu_optimizer_flag(test_config, False)
+        test_config["other_args"] = f"\"{cpu_optimizer_flag}\""
+
+        print("{0}: DeepSpeed run.".format(self.id()))
+        test_file = self.gen_output_name(test_config, prefix)
+        self.run_gpt2_test(test_config, test_file)
+
+        return self.check_parity(base_file, test_file, r_tol)
+
+    def has_loss_data(self, file_name):
+        has_loss = False
+        if os.path.exists(file_name):
+            loss = grep_loss_from_file(file_name)
+            if loss != 0.0:
+                has_loss = True
+
+        return has_loss
+
+    def check_parity(self, base_file, test_file, r_tol):
+        base_loss = grep_loss_from_file(base_file)
+        test_loss = grep_loss_from_file(test_file)
+
+        print("baseline loss: {0}, test loss: {1}".format(base_loss, test_loss))
+
+        if base_loss == 0.0 or test_loss == 0.0:
+            return False
+
+        if abs((base_loss - test_loss) / base_loss) > r_tol:
+            return False
+
+        return True
+
+    def gen_cpu_optimizer_flag(self, test_config, is_baseline):
+        if 'cpu_optimizer' in test_config and test_config['cpu_optimizer']:
+            cpu_optimizer_flag = "--cpu-optimizer"
+            if is_baseline:
+                cpu_optimizer_flag += " --cpu_torch_adam"
+                return cpu_optimizer_flag
+            if 'test_torch_offload' in test_config and test_config['test_torch_offload']:
+                cpu_optimizer_flag += " --cpu_torch_adam"
+                return cpu_optimizer_flag
+        else:
+            cpu_optimizer_flag = ""
+
+        return cpu_optimizer_flag
+
+
+def suite():
+    suite = unittest.TestSuite()
+
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu2_node1_fp16'))
+
+    # Baseline = Megatron + Torch.Optim.Adam
+    # Test = Megatron + Torch.Optim.Adam + ZeRO-Offload
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu1_node1_zero2_torch_offload'))
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu2_node1_zero2_torch_offload'))
+    suite.addTest(GPT2FuncTestCase('test_mp2_gpu4_node1_zero2_torch_offload'))
+    suite.addTest(GPT2FuncTestCase('test_mp4_gpu4_node1_zero2_torch_offload'))
+
+    # Baseline = Megatron + Torch.Optim.Adam
+    # Test = Megatron + DeepSpeedAdam + ZeRO-Offload
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu1_node1_zero2_ds_offload'))
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu2_node1_zero2_ds_offload'))
+    suite.addTest(GPT2FuncTestCase('test_mp2_gpu4_node1_zero2_ds_offload'))
+    suite.addTest(GPT2FuncTestCase('test_mp4_gpu4_node1_zero2_ds_offload'))
+
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu1_node1_zero1'))
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu2_node1_zero1'))
+    suite.addTest(GPT2FuncTestCase('test_mp2_gpu4_node1_zero1'))
+    suite.addTest(GPT2FuncTestCase('test_mp4_gpu4_node1_zero1'))
+
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu1_node1_zero2'))
+    suite.addTest(GPT2FuncTestCase('test_mp1_gpu2_node1_zero2'))
+    suite.addTest(GPT2FuncTestCase('test_mp2_gpu4_node1_zero2'))
+    suite.addTest(GPT2FuncTestCase('test_mp4_gpu4_node1_zero2'))
+
+    suite.addTest(GPT2FuncTestCase('test_mp2_gpu4_node1_zero2_gas'))
+
+    suite.addTest(GPT2FuncTestCase('test_optimizer_scheduler'))
+
+    return suite
+
+
+if __name__ == '__main__':
+    runner = unittest.TextTestRunner(failfast=True)
+    runner.run(suite())

tests/model/Megatron_GPT2/run_perf_baseline.py

@@ -0,0 +1,130 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Note: please copy webtext data to "Megatron-LM" folder, before running this script.
+"""
+
+import unittest
+import re
+from test_common import BaseTestCase
+
+
+class GPT2PerfBaselineTestCase(BaseTestCase):
+
+    def __init__(self, methodName="DeepSpeed performance test on GPT2 model"):
+        super(GPT2PerfBaselineTestCase, self).__init__(methodName)
+
+    def test_perf_1_5B(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 16,
+            "nodes": 4,
+            "bs": 16,
+            "steps": 100,
+            "layers": 48,
+            "hidden_size": 1600,
+            "seq_length": 1024,
+            "heads": 16,
+            "deepspeed": False,
+        }
+
+        self.run_test(test_config)
+
+    def test_perf_4B(self):
+        test_config = {
+            "mp": 4,
+            "gpus": 16,
+            "nodes": 4,
+            "bs": 8,
+            "steps": 100,
+            "layers": 64,
+            "hidden_size": 2304,
+            "seq_length": 1024,
+            "heads": 16,
+            "deepspeed": False,
+        }
+
+        self.run_test(test_config)
+
+    def test_perf_8B(self):
+        test_config = {
+            "mp": 4,
+            "gpus": 16,
+            "nodes": 4,
+            "bs": 8,
+            "steps": 100,
+            "layers": 72,
+            "hidden_size": 3072,
+            "seq_length": 1024,
+            "heads": 24,
+            "deepspeed": False,
+        }
+
+        self.run_test(test_config)
+
+    def test_perf_20B(self):
+        test_config = {
+            "mp": 16,
+            "gpus": 16,
+            "nodes": 4,
+            "bs": 4,
+            "steps": 50,
+            "layers": 111,
+            "hidden_size": 3808,
+            "seq_length": 1024,
+            "heads": 32,
+            "ckpt_num_layers": 1,
+            "deepspeed": False,
+        }
+
+        self.run_test(test_config)
+
+    def run_test(self, test_config):
+        print("\n")
+        print("{0}: starting......".format(self.id()))
+        prefix = "gpt2_perf"
+
+        test_file = self.gen_output_name(test_config, prefix)
+        self.run_gpt2_test(test_config, test_file)
+        exec_time = self.grep_latency_from_file(test_file)
+
+        if exec_time == 0.0:
+            print("{0}: no latency found in file {1}".format(self.id(), test_file))
+        else:
+            print("{0}: execution time per iteration is {1}ms.".format(self.id(), exec_time))
+
+    def grep_latency_from_file(self, file_name):
+        latency = 0.0
+        count = 0
+
+        with open(file_name, 'r') as f:
+            lines = f.readlines()
+            line_filter = "elapsed time per iteration"
+            match_number = re.compile(r'elapsed time per iteration \(ms\): ([-+]?[0-9]+\.?[0-9]*(?:[Ee][-+]?[0-9]+)?)')
+
+            for line in lines:
+                if line_filter in line:
+                    ms_per_iter = re.findall(match_number, line)
+                    latency += float(ms_per_iter[0])
+                    count += 1
+
+        if count > 0:
+            latency /= count
+
+        return latency
+
+
+def suite():
+    suite = unittest.TestSuite()
+    suite.addTest(GPT2PerfBaselineTestCase('test_perf_1_5B'))
+    suite.addTest(GPT2PerfBaselineTestCase('test_perf_4B'))
+    suite.addTest(GPT2PerfBaselineTestCase('test_perf_8B'))
+    suite.addTest(GPT2PerfBaselineTestCase('test_perf_20B'))
+    return suite
+
+
+if __name__ == '__main__':
+    runner = unittest.TextTestRunner(failfast=True)
+    runner.run(suite())

tests/model/Megatron_GPT2/run_perf_test.py

@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Note: please copy webtext data to "Megatron-LM" folder, before running this script.
+"""
+
+import unittest
+import re
+from test_common import BaseTestCase
+
+
+class GPT2PerfTestCase(BaseTestCase):
+
+    def __init__(self, methodName="DeepSpeed performance test on GPT2 model"):
+        super(GPT2PerfTestCase, self).__init__(methodName)
+
+    def test_perf_1_5B(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 16,
+            "nodes": 4,
+            "bs": 32,
+            "steps": 100,
+            "layers": 48,
+            "hidden_size": 1600,
+            "seq_length": 1024,
+            "heads": 16,
+            "deepspeed": True,
+            "json": "ds_config_perf_bs32.json",
+        }
+
+        self.run_test(test_config)
+
+    def test_perf_4B(self):
+        test_config = {
+            "mp": 1,
+            "gpus": 16,
+            "nodes": 4,
+            "bs": 8,
+            "steps": 100,
+            "layers": 64,
+            "hidden_size": 2304,
+            "seq_length": 1024,
+            "heads": 16,
+            "deepspeed": True,
+            "json": "ds_config_perf_bs8.json",
+        }
+
+        self.run_test(test_config)
+
+    def test_perf_8B(self):
+        test_config = {
+            "mp": 2,
+            "gpus": 16,
+            "nodes": 4,
+            "bs": 16,
+            "steps": 100,
+            "layers": 72,
+            "hidden_size": 3072,
+            "seq_length": 1024,
+            "heads": 24,
+            "deepspeed": True,
+            "json": "ds_config_perf_bs16.json",
+        }
+
+        self.run_test(test_config)
+
+    def test_perf_20B(self):
+        test_config = {
+            "mp": 4,
+            "gpus": 16,
+            "nodes": 4,
+            "bs": 8,
+            "steps": 50,
+            "layers": 111,
+            "hidden_size": 3808,
+            "seq_length": 1024,
+            "heads": 32,
+            "ckpt_num_layers": 1,
+            "deepspeed": True,
+            "json": "ds_config_perf_bs8.json",
+        }
+
+        self.run_test(test_config)
+
+    def run_test(self, test_config):
+        print("\n")
+        print("{0}: starting......".format(self.id()))
+        prefix = "gpt2_perf"
+
+        test_file = self.gen_output_name(test_config, prefix)
+        self.run_gpt2_test(test_config, test_file)
+        exec_time = self.grep_latency_from_file(test_file)
+
+        if exec_time == 0.0:
+            print("{0}: no latency found in file {1}".format(self.id(), test_file))
+        else:
+            print("{0}: execution time per iteration is {1}ms.".format(self.id(), exec_time))
+
+    def grep_latency_from_file(self, file_name):
+        latency = 0.0
+        count = 0
+
+        with open(file_name, 'r') as f:
+            lines = f.readlines()
+            line_filter = "elapsed time per iteration"
+            match_number = re.compile(r'elapsed time per iteration \(ms\): ([-+]?[0-9]+\.?[0-9]*(?:[Ee][-+]?[0-9]+)?)')
+
+            for line in lines:
+                if line_filter in line:
+                    ms_per_iter = re.findall(match_number, line)
+                    latency += float(ms_per_iter[0])
+                    count += 1
+
+        if count > 0:
+            latency /= count
+
+        return latency
+
+
+def suite():
+    suite = unittest.TestSuite()
+    suite.addTest(GPT2PerfTestCase('test_perf_1_5B'))
+    suite.addTest(GPT2PerfTestCase('test_perf_4B'))
+    suite.addTest(GPT2PerfTestCase('test_perf_8B'))
+    suite.addTest(GPT2PerfTestCase('test_perf_20B'))
+    return suite
+
+
+if __name__ == '__main__':
+    runner = unittest.TextTestRunner(failfast=True)
+    runner.run(suite())

tests/model/Megatron_GPT2/test_common.py

@@ -0,0 +1,68 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import unittest
+import subprocess
+import os
+import time
+import shlex
+
+
+class BaseTestCase(unittest.TestCase):
+
+    def __init__(self, methodName="DeepSpeed performance test"):
+        super(BaseTestCase, self).__init__(methodName)
+        self.test_dir = "./test"
+        self.baseline_dir = "./baseline"
+        self.timestr = time.strftime("%Y%m%d-%H%M%S")
+
+    def gen_output_name(self, test_config, prefix, baseline_config=False):
+        other_args = test_config["other_args"] if "other_args" in test_config else ""
+        zero_args = "_zero" if "zero" in test_config and test_config["zero"] else ""
+        other_args = other_args.strip(' -\\').replace(" ", "").replace("\"", "")
+
+        if other_args:
+            other_args = "_" + other_args
+
+        if test_config["deepspeed"] and not baseline_config:
+            file_name = "_mp{0}_gpu{1}_node{2}_bs{3}_step{4}_layer{5}_hidden{6}_seq{7}_head{8}{9}_ds{10}-{11}.log".format(
+                test_config["mp"], test_config["gpus"], test_config["nodes"], test_config["bs"], test_config["steps"],
+                test_config["layers"], test_config["hidden_size"], test_config["seq_length"], test_config["heads"],
+                other_args, zero_args, self.timestr)
+            save_dir = self.test_dir
+        else:
+            file_name = "_mp{0}_gpu{1}_node{2}_bs{3}_step{4}_layer{5}_hidden{6}_seq{7}_head{8}{9}.log".format(
+                test_config["mp"], test_config["gpus"], test_config["nodes"], test_config["bs"], test_config["steps"],
+                test_config["layers"], test_config["hidden_size"], test_config["seq_length"], test_config["heads"],
+                other_args)
+            save_dir = self.baseline_dir
+
+        return os.path.join(save_dir, prefix + file_name)
+
+    def ensure_directory_exists(self, filename):
+        dirname = os.path.dirname(filename)
+        if not os.path.exists(dirname):
+            os.makedirs(dirname)
+
+    def clean_test_env(self):
+        cmd = shlex.split("dlts_ssh pkill -9 -f /usr/bin/python")
+        print(cmd)
+        subprocess.run(cmd, check=False, executable='/bin/bash')
+        time.sleep(20)
+
+    def run_gpt2_test(self, test_config, output):
+        ds_flag = "-d " + test_config["json"] if test_config["deepspeed"] else ""
+        ckpt_num = test_config["ckpt_num_layers"] if "ckpt_num_layers" in test_config else 1
+        other_args = "-o " + test_config["other_args"] if "other_args" in test_config else ""
+
+        cmd = "./ds_gpt2_test.sh -m {0} -g {1} -n {2} -b {3} -s {4} -l {5} -h {6} -q {7} -e {8} -c {9} {10} {11}".format(
+            test_config["mp"], test_config["gpus"], test_config["nodes"], test_config["bs"], test_config["steps"],
+            test_config["layers"], test_config["hidden_size"], test_config["seq_length"], test_config["heads"],
+            ckpt_num, other_args, ds_flag)
+        cmd = shlex.split(cmd)
+        self.ensure_directory_exists(output)
+        with open(output, "w") as f:
+            print(cmd)
+            subprocess.run(cmd, check=False, executable='/bin/bash', stdout=f, stderr=f)

tests/model/run_sanity_check.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Note: please copy webtext data to "Megatron-LM" folder, before running this script.
+"""
+
+import sys
+import unittest
+
+sys.path.append('../DeepSpeedExamples/Megatron_GPT2')
+sys.path.append('../DeepSpeedExamples/BingBertSquad')
+
+# Import the test cases here.
+import Megatron_GPT2
+import BingBertSquad
+
+
+def pytest_hack(runner_result):
+    '''This is an ugly hack to get the unittest suites to play nicely with
+    pytest. Otherwise failed tests are not reported by pytest for some reason.
+
+    Long-term, these model tests should be adapted to pytest.
+    '''
+    if not runner_result.wasSuccessful():
+        print('SUITE UNSUCCESSFUL:', file=sys.stderr)
+        for fails in runner_result.failures:
+            print(fails, file=sys.stderr)
+        assert runner_result.wasSuccessful()  # fail the test
+
+
+def test_megatron():
+    runner = unittest.TextTestRunner(failfast=True)
+    pytest_hack(runner.run(Megatron_GPT2.suite()))
+
+
+def test_megatron_checkpoint():
+    runner = unittest.TextTestRunner(failfast=True)
+    pytest_hack(runner.run(Megatron_GPT2.checkpoint_suite()))
+
+
+def test_squad():
+    runner = unittest.TextTestRunner(failfast=True)
+    pytest_hack(runner.run(BingBertSquad.suite()))

tests/onebit/README.md

@@ -0,0 +1,31 @@
+# One-Bit tests
+
+In this folder, you can test the functionality and performance of different backend for doing compressed allreduce, which is the main algorithm in one-bit optimizers like [One-Bit Adam](https://www.deepspeed.ai/tutorials/onebit-adam/), [One-Bit Lamb](https://www.deepspeed.ai/tutorials/onebit-lamb/) and [Zero-One Adam](https://www.deepspeed.ai/tutorials/zero-one-adam/).
+
+## How to run
+
+### NCCL and MPI backend
+
+Basically it requires your environment have relative communication backend installed, the NCCL backend of PyTorch distributed or Message Passing Interface (MPI) like MVAPICH2-GDR and OpenMPI. [Detailed Pre-requisites](https://www.deepspeed.ai/tutorials/zero-one-adam/#12-pre-requisites-for-01-adam).
+
+To test accuracy and performance of NCCL backend:
+```bash
+python test_nccl_backend.py
+python test_nccl_perf.py
+```
+Similarly, for MPI backend:
+```bash
+python test_mpi_backend.py
+python test_mpi_perf.py
+```
+
+### Compressed backend
+
+This backend provides an approach to abstract the generic part of one-bit optimizers and implements accelerator dependent part with DeepSpeed custom op builder. To use this `CompressedBackend` and test it, you should make sure that your current accelerator supports `PackbitsBuilder`, so that it could be loaded to do high performance packing and unpacking between float and Byte datatype.
+An example can be found in `Deepspeed/op_builder/xpu/packbits.py`.
+
+The test usage is same as others:
+```bash
+python test_compressed_backend.py
+python test_compressed_perf.py
+```

tests/onebit/test_compressed_backend.py

@@ -0,0 +1,96 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed.comm as dist
+import numpy as np
+import argparse
+import deepspeed
+import os
+
+from deepspeed.runtime.comm.compressed import CompressedBackend
+from deepspeed.accelerator import get_accelerator
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--local_rank', type=int, default=-1)
+args = parser.parse_args()
+
+deepspeed.init_distributed(dist_backend=get_accelerator().communication_backend_name())
+args.local_rank = int(os.environ['LOCAL_RANK'])
+
+get_accelerator().set_device(args.local_rank)
+device = torch.device(get_accelerator().device_name(), args.local_rank)
+
+size = dist.get_world_size()
+rank = dist.get_rank()
+
+backend = CompressedBackend()
+local_rank = args.local_rank
+
+
+# A simulated compression function using deepspeed.comm
+def torch_sim(a):
+    a_sign = a.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)
+    scale = a.norm() / np.sqrt(a.numel())
+    a_compressed = scale * a_sign
+    a_sign = None
+    worker_error = a - a_compressed
+    dist.all_reduce(a_compressed)
+    a_compressed.mul_(1 / dist.get_world_size())
+    a_server_sign = a_compressed.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)
+    a_list = torch.chunk(a_compressed, chunks=dist.get_world_size())
+    server_scale = [chunk_a.norm() / np.sqrt(chunk_a.numel()) for chunk_a in a_list]
+    a_sign_list = torch.chunk(a_server_sign, dist.get_world_size())
+    a_server_compressed = torch.cat([server_scale[i] * a_sign_list[i] for i in range(dist.get_world_size())])
+    rank = dist.get_rank()
+    server_error = a_list[rank] - server_scale[rank] * a_sign_list[rank]
+    get_accelerator().synchronize()
+    dist.barrier()
+    return a_server_compressed, worker_error, server_error
+
+
+tensor_size = 300 * 2**20
+server_size = int(tensor_size / size)
+if tensor_size % (8 * size) != 0:
+    right_tensor_size = tensor_size + (8 * size - (tensor_size % (8 * size)))
+else:
+    right_tensor_size = tensor_size
+right_server_size = right_tensor_size // size
+
+# Adding bias to the initialization of the gradient we are communicating
+# In order to get rid of the case where some elements in the gradient are too small
+a = (torch.rand(tensor_size, device=device) - 0.5) + 0.01 * rank
+
+worker_error = torch.zeros(right_tensor_size, device=device)
+server_error = torch.zeros(right_server_size, device=device)
+
+a_torch, worker_error_torch, server_error_torch = torch_sim(a)
+get_accelerator().empty_cache()
+
+a_after = backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+
+print(a_torch.cpu())
+print(a_after.cpu())
+
+threshold = 1e-6
+magnitude_threshold = 1e-6
+diff_mask = (a_after - a_torch) > threshold
+diff_server_mask = torch.chunk(diff_mask, size)[rank]
+mpi_server = torch.chunk(a_after, size)[rank] + server_error
+torch_server = torch.chunk(a_torch, size)[rank] + server_error_torch
+
+test_correctness = True
+
+# If the number in the compensated_server_m is too small (e.g 1e-8), then calling sign() might be problematic
+# The test would skip those numbers that are too small in compensated_server_m
+if test_correctness:
+    if torch.sum(diff_server_mask) == 0:
+        print('Successfully passed the test for Compressed Backend at Rank {}'.format(rank))
+    else:
+        check_mag_mask = mpi_server[diff_server_mask] > magnitude_threshold
+        if torch.sum(check_mag_mask) == 0:
+            print('Successfully passed the test for Compressed Backend at Rank {}'.format(rank))
+        else:
+            print('Fails at {} of positions'.format(torch.sum(check_mag_mask)))

tests/onebit/test_compressed_perf.py

@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed.comm as dist
+import numpy as np
+import argparse
+import deepspeed
+import os
+
+from deepspeed.runtime.comm.compressed import CompressedBackend
+from deepspeed.utils.timer import SynchronizedWallClockTimer
+from deepspeed.accelerator import get_accelerator
+from statistics import mean
+
+timers = SynchronizedWallClockTimer()
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--local_rank', type=int, default=-1)
+args = parser.parse_args()
+
+deepspeed.init_distributed(dist_backend=get_accelerator().communication_backend_name())
+args.local_rank = int(os.environ['LOCAL_RANK'])
+
+get_accelerator().set_device(args.local_rank)
+device = torch.device(get_accelerator().device_name(), args.local_rank)
+
+size = dist.get_world_size()
+rank = dist.get_rank()
+
+backend = CompressedBackend()
+local_rank = args.local_rank
+
+# Setting tensor_size (BERT-Large)
+tensor_size = 300 * 2**20
+server_size = int(tensor_size / size)
+if tensor_size % (8 * size) != 0:
+    right_tensor_size = tensor_size + (8 * size - (tensor_size % (8 * size)))
+else:
+    right_tensor_size = tensor_size
+right_server_size = right_tensor_size // size
+
+# Adding bias to the initialization of the gradient we are communicating
+# In order to get rid of the case where some elements in the gradient are too small
+a = (torch.rand(tensor_size, device=device) - 0.5) + 0.01 * rank
+
+worker_error = torch.zeros(right_tensor_size, device=device)
+server_error = torch.zeros(right_server_size, device=device)
+
+warmup = 10
+iters = 10
+
+# Warmup
+for i in range(warmup):
+    backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+
+time_list = []
+
+a_sign = a.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)
+scale = a.norm() / np.sqrt(a.numel())
+a_compressed = scale * a_sign
+
+print("Shape of the compressed buffer:", a_compressed.shape) if rank == 0 else None
+
+for i in range(iters):
+    timers('compressed_allreduce').start()
+    backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+    #deepspeed.comm.all_reduce(a_compressed)
+    timers('compressed_allreduce').stop()
+    time_list.append(timers('compressed_allreduce').elapsed())
+
+#timer_names = ['compressed_allreduce']
+#timers.log(names=timer_names, normalizer=1, memory_breakdown=None)
+
+places = 2
+convert = 1e3
+float_size = 4
+
+if rank == 0:
+    for i in range(iters):
+        lat = time_list[i]
+        print("latency = ", lat * convert)
+
+minlat = round(min(time_list) * convert)
+maxlat = round(max(time_list) * convert)
+meanlat = round(mean(time_list) * convert, places)
+print("min, max, and mean = {} ms, {} ms, {} ms".format(minlat, maxlat, meanlat)) if rank == 0 else None
+#print("tensor shape", a.shape)
+duration = meanlat / 1e3
+tput = ((tensor_size * 4) / duration)
+print("algo throughput: %f Bytes/s, %f GB/s" % (tput, tput / 1e9)) if rank == 0 else None
+size = tensor_size * 4
+n = dist.get_world_size()
+busbw = (size / duration) * (2 * (n - 1) / n)
+print("busbw: %f GB/s" % (busbw / 1e9)) if rank == 0 else None

tests/onebit/test_mpi_backend.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from mpi4py import MPI
+import torch
+import deepspeed.comm as dist
+import numpy as np
+import deepspeed
+
+from deepspeed.runtime.comm.mpi import MpiBackend
+from deepspeed.accelerator import get_accelerator
+
+comm = MPI.COMM_WORLD
+size = comm.Get_size()
+rank = comm.Get_rank()
+
+deepspeed.init_distributed(dist_backend=get_accelerator().communication_backend_name())
+
+# Change cuda_aware to True to test out CUDA-Aware MPI communication
+backend = MpiBackend(cuda_aware=False)
+
+local_rank = rank % get_accelerator().device_count()
+device = torch.device(get_accelerator().device_name(), local_rank)
+
+
+# A simulated compression function using deepspeed.comm
+def torch_sim(a):
+    a_sign = a.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)
+    scale = a.norm() / np.sqrt(a.numel())
+    a_compressed = scale * a_sign
+    a_sign = None
+    worker_error = a - a_compressed
+    dist.all_reduce(a_compressed)
+    a_compressed.mul_(1 / dist.get_world_size())
+    a_server_sign = a_compressed.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)
+    a_list = torch.chunk(a_compressed, chunks=dist.get_world_size())
+    server_scale = [chunk_a.norm() / np.sqrt(chunk_a.numel()) for chunk_a in a_list]
+    a_sign_list = torch.chunk(a_server_sign, dist.get_world_size())
+    a_server_compressed = torch.cat([server_scale[i] * a_sign_list[i] for i in range(dist.get_world_size())])
+    rank = dist.get_rank()
+    server_error = a_list[rank] - server_scale[rank] * a_sign_list[rank]
+    get_accelerator().synchronize()
+    dist.barrier()
+    return a_server_compressed, worker_error, server_error
+
+
+tensor_size = 100 * 2**20
+server_size = int(tensor_size / size)
+if tensor_size % (8 * size) != 0:
+    right_tensor_size = tensor_size + (8 * size - (tensor_size % (8 * size)))
+else:
+    right_tensor_size = tensor_size
+right_server_size = right_tensor_size // size
+
+# Adding bias to the initialization of the gradient we are communicating
+# In order to get rid of the case where some elements in the gradient are too small
+a = (torch.rand(tensor_size, device=device) - 0.5) + 0.01 * rank
+
+worker_error = torch.zeros(right_tensor_size, device=device)
+server_error = torch.zeros(right_server_size, device=device)
+
+a_torch, worker_error_torch, server_error_torch = torch_sim(a)
+get_accelerator().empty_cache()
+
+a_after = backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+
+threshold = 1e-6
+magnitude_threshold = 1e-6
+diff_mask = (a_after - a_torch) > threshold
+diff_server_mask = torch.chunk(diff_mask, size)[rank]
+mpi_server = torch.chunk(a_after, size)[rank] + server_error
+torch_server = torch.chunk(a_torch, size)[rank] + server_error_torch
+
+test_correctness = True
+
+# If the number in the compensated_server_m is too small (e.g 1e-8), then calling sign() might be problematic
+# The test would skip those numbers that are too small in compensated_server_m
+if test_correctness:
+    if torch.sum(diff_server_mask) == 0:
+        print('Successfully passed the test for MPI Backend at Rank {}'.format(rank))
+    else:
+        check_mag_mask = mpi_server[diff_server_mask] > magnitude_threshold
+        if torch.sum(check_mag_mask) == 0:
+            print('Successfully passed the test for MPI Backend at Rank {}'.format(rank))
+        else:
+            print('Fails at {} of positions'.format(torch.sum(check_mag_mask)))

tests/onebit/test_mpi_perf.py

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from mpi4py import MPI
+import torch
+import deepspeed
+
+from deepspeed.runtime.comm.mpi import MpiBackend
+
+# Configure wall clock timer
+from deepspeed.utils.timer import SynchronizedWallClockTimer
+from deepspeed.accelerator import get_accelerator
+
+from statistics import mean
+
+timers = SynchronizedWallClockTimer()
+
+comm = MPI.COMM_WORLD
+size = comm.Get_size()
+rank = comm.Get_rank()
+
+deepspeed.init_distributed(dist_backend=get_accelerator().communication_backend_name())
+# Change cuda_aware to True to test out CUDA-Aware MPI communication
+backend = MpiBackend(cuda_aware=False)
+
+local_rank = rank % get_accelerator().device_count()
+device = torch.device(get_accelerator().device_name(), local_rank)
+
+tensor_size = 300 * 2**20
+server_size = int(tensor_size / size)
+if tensor_size % (8 * size) != 0:
+    right_tensor_size = tensor_size + (8 * size - (tensor_size % (8 * size)))
+else:
+    right_tensor_size = tensor_size
+right_server_size = right_tensor_size // size
+
+# Adding bias to the initialization of the gradient we are communicating
+# In order to get rid of the case where some elements in the gradient are too small
+a = (torch.rand(tensor_size, device=device) - 0.5) + 0.01 * rank
+
+worker_error = torch.zeros(right_tensor_size, device=device)
+server_error = torch.zeros(right_server_size, device=device)
+
+warmup = 10
+iters = 10
+
+# Warmup
+for i in range(warmup):
+    backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+
+time_list = []
+
+for i in range(iters):
+    timers('compressed_allreduce').start()
+    backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+    timers('compressed_allreduce').stop()
+    time_list.append(timers('compressed_allreduce').elapsed())
+
+timer_names = ['compressed_allreduce']
+timers.log(names=timer_names, normalizer=1, memory_breakdown=None)
+
+places = 2
+convert = 1e3
+float_size = 4
+
+if rank == 0:
+    for i in range(iters):
+        lat = time_list[i]
+        print("latency = ", lat * convert)
+
+minlat = round(min(time_list) * convert)
+maxlat = round(max(time_list) * convert)
+meanlat = round(mean(time_list) * convert, places)
+print("min, max, and mean = {} ms, {} ms, {} ms".format(minlat, maxlat, meanlat))

tests/onebit/test_nccl_backend.py

@@ -0,0 +1,93 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed.comm as dist
+import numpy as np
+import argparse
+import deepspeed
+import os
+
+from deepspeed.runtime.comm.nccl import NcclBackend
+from deepspeed.accelerator import get_accelerator
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--local_rank', type=int, default=-1)
+args = parser.parse_args()
+
+deepspeed.init_distributed(dist_backend=get_accelerator().communication_backend_name())
+args.local_rank = int(os.environ['LOCAL_RANK'])
+
+get_accelerator().set_device(args.local_rank)
+device = torch.device(get_accelerator().device_name(), args.local_rank)
+
+size = dist.get_world_size()
+rank = dist.get_rank()
+
+backend = NcclBackend()
+local_rank = args.local_rank
+
+
+# A simulated compression function using deepspeed.comm
+def torch_sim(a):
+    a_sign = a.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)
+    scale = a.norm() / np.sqrt(a.numel())
+    a_compressed = scale * a_sign
+    a_sign = None
+    worker_error = a - a_compressed
+    dist.all_reduce(a_compressed)
+    a_compressed.mul_(1 / dist.get_world_size())
+    a_server_sign = a_compressed.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)
+    a_list = torch.chunk(a_compressed, chunks=dist.get_world_size())
+    server_scale = [chunk_a.norm() / np.sqrt(chunk_a.numel()) for chunk_a in a_list]
+    a_sign_list = torch.chunk(a_server_sign, dist.get_world_size())
+    a_server_compressed = torch.cat([server_scale[i] * a_sign_list[i] for i in range(dist.get_world_size())])
+    rank = dist.get_rank()
+    server_error = a_list[rank] - server_scale[rank] * a_sign_list[rank]
+    get_accelerator().synchronize()
+    dist.barrier()
+    return a_server_compressed, worker_error, server_error
+
+
+tensor_size = 300 * 2**20
+server_size = int(tensor_size / size)
+if tensor_size % (8 * size) != 0:
+    right_tensor_size = tensor_size + (8 * size - (tensor_size % (8 * size)))
+else:
+    right_tensor_size = tensor_size
+right_server_size = right_tensor_size // size
+
+# Adding bias to the initialization of the gradient we are communicating
+# In order to get rid of the case where some elements in the gradient are too small
+a = (torch.rand(tensor_size, device=device) - 0.5) + 0.01 * rank
+
+worker_error = torch.zeros(right_tensor_size, device=device)
+server_error = torch.zeros(right_server_size, device=device)
+
+a_torch, worker_error_torch, server_error_torch = torch_sim(a)
+get_accelerator().empty_cache()
+
+a_after = backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+
+threshold = 1e-6
+magnitude_threshold = 1e-6
+diff_mask = (a_after - a_torch) > threshold
+diff_server_mask = torch.chunk(diff_mask, size)[rank]
+mpi_server = torch.chunk(a_after, size)[rank] + server_error
+torch_server = torch.chunk(a_torch, size)[rank] + server_error_torch
+
+test_correctness = True
+
+# If the number in the compensated_server_m is too small (e.g 1e-8), then calling sign() might be problematic
+# The test would skip those numbers that are too small in compensated_server_m
+if test_correctness:
+    if torch.sum(diff_server_mask) == 0:
+        print('Successfully passed the test for NCCL Backend at Rank {}'.format(rank))
+    else:
+        check_mag_mask = mpi_server[diff_server_mask] > magnitude_threshold
+        if torch.sum(check_mag_mask) == 0:
+            print('Successfully passed the test for NCCL Backend at Rank {}'.format(rank))
+        else:
+            print('Fails at {} of positions'.format(torch.sum(check_mag_mask)))

tests/onebit/test_nccl_perf.py

@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed.comm as dist
+import numpy as np
+import argparse
+import deepspeed
+import os
+
+from deepspeed.runtime.comm.nccl import NcclBackend
+from deepspeed.utils.timer import SynchronizedWallClockTimer
+from deepspeed.accelerator import get_accelerator
+from statistics import mean
+
+timers = SynchronizedWallClockTimer()
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--local_rank', type=int, default=-1)
+args = parser.parse_args()
+
+deepspeed.init_distributed(dist_backend=get_accelerator().communication_backend_name())
+args.local_rank = int(os.environ['LOCAL_RANK'])
+
+get_accelerator().set_device(args.local_rank)
+device = torch.device(get_accelerator().device_name(), args.local_rank)
+
+size = dist.get_world_size()
+rank = dist.get_rank()
+
+backend = NcclBackend()
+local_rank = args.local_rank
+
+# Setting tensor_size (BERT-Large)
+tensor_size = 300 * 2**20
+server_size = int(tensor_size / size)
+if tensor_size % (8 * size) != 0:
+    right_tensor_size = tensor_size + (8 * size - (tensor_size % (8 * size)))
+else:
+    right_tensor_size = tensor_size
+right_server_size = right_tensor_size // size
+
+# Adding bias to the initialization of the gradient we are communicating
+# In order to get rid of the case where some elements in the gradient are too small
+a = (torch.rand(tensor_size, device=device) - 0.5) + 0.01 * rank
+
+worker_error = torch.zeros(right_tensor_size, device=device)
+server_error = torch.zeros(right_server_size, device=device)
+
+warmup = 10
+iters = 10
+
+# Warmup
+for i in range(warmup):
+    backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+
+time_list = []
+
+a_sign = a.sign().add_(1).bool().float().add_(-0.5).mul_(2.0)
+scale = a.norm() / np.sqrt(a.numel())
+a_compressed = scale * a_sign
+
+print("Shape of the compressed buffer:", a_compressed.shape) if rank == 0 else None
+
+for i in range(iters):
+    timers('compressed_allreduce').start()
+    backend.compressed_allreduce(a, worker_error, server_error, local_rank)
+    #deepspeed.comm.all_reduce(a_compressed)
+    timers('compressed_allreduce').stop()
+    time_list.append(timers('compressed_allreduce').elapsed())
+
+#timer_names = ['compressed_allreduce']
+#timers.log(names=timer_names, normalizer=1, memory_breakdown=None)
+
+places = 2
+convert = 1e3
+float_size = 4
+
+if rank == 0:
+    for i in range(iters):
+        lat = time_list[i]
+        print("latency = ", lat * convert)
+
+minlat = round(min(time_list) * convert)
+maxlat = round(max(time_list) * convert)
+meanlat = round(mean(time_list) * convert, places)
+print("min, max, and mean = {} ms, {} ms, {} ms".format(minlat, maxlat, meanlat)) if rank == 0 else None
+#print("tensor shape", a.shape)
+duration = meanlat / 1e3
+tput = ((tensor_size * 4) / duration)
+print("algo throughput: %f Bytes/s, %f GB/s" % (tput, tput / 1e9)) if rank == 0 else None
+size = tensor_size * 4
+n = dist.get_world_size()
+busbw = (size / duration) * (2 * (n - 1) / n)
+print("busbw: %f GB/s" % (busbw / 1e9)) if rank == 0 else None

tests/perf/adagrad_test.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad
+import time
+
+NUM_ITERS = 100
+
+
+def _test_perf(param, optimizer_func):
+    optimizer = optimizer_func(param)
+    avg = 0
+    for i in range(NUM_ITERS):
+        for i, p in enumerate(param):
+            p.grad = torch.ones_like(p) * 2
+        start = time.time()
+        optimizer.step()
+        stop = time.time()
+        avg += (stop - start)
+
+    return avg / NUM_ITERS
+
+
+def _main():
+    device = 'cpu'
+    model_size = 1 * 1024**3
+    group_size = [model_size, 274432]
+    param = [torch.nn.Parameter(torch.ones(size, device=device)) for size in group_size]
+    torch_time = _test_perf(param, torch.optim.Adagrad)
+    ds_time = _test_perf(param, DeepSpeedCPUAdagrad)
+    print(f"Step time: {torch_time=} {ds_time=}")
+
+
+_main()

tests/perf/adam_test.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from deepspeed.ops.adam import DeepSpeedCPUAdam
+import time
+
+NUM_ITERS = 100
+
+
+def _test_perf(param, optimizer_func):
+    optimizer = optimizer_func(param)
+    avg = 0
+    for i in range(NUM_ITERS):
+        for i, p in enumerate(param):
+            p.grad = torch.ones_like(p) * 2
+        start = time.time()
+        optimizer.step()
+        stop = time.time()
+        avg += (stop - start)
+
+    return avg / NUM_ITERS
+
+
+def _main():
+    device = 'cpu'
+    model_size = 1 * 1024**3
+    group_size = [model_size, 274432]
+    param = [torch.nn.Parameter(torch.ones(size, device=device)) for size in group_size]
+    torch_time = _test_perf(param, torch.optim.Adam)
+    ds_time = _test_perf(param, DeepSpeedCPUAdam)
+    print(f"Step time: {torch_time=} {ds_time=}")
+
+
+_main()

tests/perf/adam_test1.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from deepspeed.ops.adam import DeepSpeedCPUAdam
+import time
+
+device = 'cpu'
+model_size = 1 * 1024**3
+param = torch.nn.Parameter(torch.ones(model_size, device=device))
+
+optimizer = DeepSpeedCPUAdam([param])
+#torch.set_num_threads(128)
+param.grad = torch.ones(model_size, device=device)
+avg = 0
+for i in range(100):
+    start = time.time()
+    optimizer.step()
+    stop = time.time()
+    avg += (stop - start)
+    param.grad = torch.ones(model_size, device=device) * 2
+print("Elapsed Time is ", avg / 100)

tests/pytest.ini

@@ -0,0 +1,13 @@
+[pytest]
+addopts = -m "not sequential and not nightly and not inference and not seq_inference and not inference_ops and not inference_v2 and not inference_v2_ops and not stable_diffusion and not evaluation"
+markers =
+    sequential:Tests that need to be run sequentially
+    inference:Inference model tests
+    inference_ops:Individual inference operator tests
+    inference_v2:Inference tests for the v2 stack
+    inference_v2_ops:Op tests for the v2 stack
+    seq_inference:Inference model tests to run sequentially
+    nightly:Tests that should be run nightly
+    world_size:Change world size of individual tests in a class
+    stable_diffusion:Tests that run Stable Diffusion
+    evaluation:Tests that evaluate model correctness

tests/small_model_debugging/partial_offload_test.py

@@ -0,0 +1,128 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import json
+import argparse
+import torch
+import deepspeed
+from torch.utils.data.distributed import DistributedSampler
+import deepspeed.comm as dist
+
+
+class SimpleModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim, empty_grad=False):
+        super(SimpleModel, self).__init__()
+        self.linear = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.linear3 = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.linear4 = torch.nn.Linear(hidden_dim, hidden_dim)
+        if empty_grad:
+            self.layers2 = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim)])
+        self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        hidden = x
+        hidden = self.linear(hidden)
+        hidden = self.linear2(hidden)
+        hidden = self.linear3(hidden)
+        hidden = self.linear4(hidden)
+        return self.cross_entropy_loss(hidden, y)
+
+
+def create_config_from_dict(tmpdir, config_dict):
+    config_path = os.path.join(tmpdir, 'temp_config.json')
+    with open(config_path, 'w') as fd:
+        json.dump(config_dict, fd)
+    return config_path
+
+
+def get_data_loader(model, total_samples, hidden_dim, device):
+    batch_size = model.train_micro_batch_size_per_gpu()
+    train_data = torch.randn(total_samples, hidden_dim, device=device, dtype=torch.half)
+    train_label = torch.empty(total_samples, dtype=torch.long, device=device).random_(hidden_dim)
+    train_dataset = torch.utils.data.TensorDataset(train_data, train_label)
+    sampler = DistributedSampler(train_dataset)
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=sampler)
+    return train_loader
+
+
+def get_args(tmpdir, config_dict):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_rank", type=int, default=0)
+    parser.add_argument('--zero', type=int, default=0)
+    args = parser.parse_args()  #args=''
+
+    config_dict["zero_optimization"]["stage"] = args.zero
+    print('config_dict["zero_optimization"]', config_dict["zero_optimization"])
+    config_path = create_config_from_dict(tmpdir, config_dict)
+
+    args.deepspeed_config = config_path
+    return args
+
+
+def print0(msg):
+    if dist.get_rank() == 0:
+        print(msg, flush=True)
+
+
+rank = int(os.environ['RANK'])
+print('seed:', 2222 + rank)
+torch.random.manual_seed(2222 + rank)
+
+config_dict = {
+    "train_batch_size": 256,
+    "steps_per_print": 1,
+    "optimizer": {
+        "type": "Adam",
+        "params": {
+            "lr": 0.00015,
+        }
+    },
+    "fp16": {
+        "enabled": True,
+        "initial_scale_power": 15
+    },
+    "zero_optimization": {
+        "stage": 0,
+        "sub_group_size": 8,
+        "reduce_bucket_size": 20,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": True,
+            "ratio": 0.3
+        }
+    }
+}
+#        "initial_scale_power": 15
+args = get_args('/tmp/', config_dict)
+hidden_dim = 4 * 1024
+
+model = SimpleModel(hidden_dim, empty_grad=False)
+
+model, _, _, _ = deepspeed.initialize(args=args,
+                                      model=model,
+                                      model_parameters=model.parameters(),
+                                      dist_init_required=True)
+
+
+def print_params(tag, model):
+    if dist.get_rank() == 0:
+        for n, p in model.named_parameters():
+            print0("{} {}:{}".format(tag, n, p))
+
+
+data_loader = get_data_loader(model=model, total_samples=1000, hidden_dim=hidden_dim, device=model.device)
+#print_params('pre-train', model)
+#while True:
+for n, batch in enumerate(data_loader):
+    loss = model(batch[0], batch[1])
+    if dist.get_rank() == 0:
+        print("LOSS:", loss.item())
+    model.backward(loss)
+    model.step()
+    #print_params('step={}'.format(n), model)
+    if n == 2: break

tests/small_model_debugging/stage3_test.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+import deepspeed
+
+###################################
+# Setup
+###################################
+
+
+class VerboseLinear(torch.nn.Linear):
+
+    def __init__(self, **kwargs):
+        print('Begin VerboseLinear.__init__')
+        super().__init__(**kwargs)
+        print('End VerboseLinear.__init__')
+
+
+class LinearStack(torch.nn.Module):
+
+    def __init__(self, input_dim=2, hidden_dim=4, output_dim=4, num_layers=2):
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.hidden_dim = hidden_dim
+
+        self.input_layer = VerboseLinear(in_features=self.input_dim, out_features=self.hidden_dim)
+        self.layers = torch.nn.ModuleList([
+            torch.nn.Linear(in_features=self.hidden_dim, out_features=self.hidden_dim, bias=False)
+            for x in range(num_layers)
+        ])
+        self.output_layer = torch.nn.Linear(in_features=self.hidden_dim, out_features=self.output_dim)
+        self.identity = torch.nn.Identity()
+
+    def forward(self, x):
+        x = self.input_layer(x)
+        for layer in self.layers:
+            x = layer(x)
+        x = self.output_layer(x)
+        x = self.identity(x)
+        return x
+
+
+###################################
+# DRIVER
+###################################
+
+
+def test_driver():
+    print()
+    print('BUILDING MODEL')
+    with deepspeed.zero.Init():
+        model = LinearStack()
+    print()
+
+    # parted = [name for (name, p) in model.named_parameters() if p._partitioned]
+    # not_parted = [name for (name, p) in model.named_parameters() if not p._partitioned]
+    # print('partitioned: ', parted)
+    # print('full: ', not_parted)
+    # print()
+
+    model.train()
+
+    test_input = torch.rand(1, model.input_dim)
+    grad_output = torch.rand(1, model.output_dim)
+
+    grad_output.requires_grad = False
+    test_input.requires_grad = False
+
+    print()
+    print('BEGINNING FORWARD')
+    print()
+
+    output = model(test_input)
+    output.backward(grad_output)
+
+    # parted = [name for (name, p) in model.named_parameters() if p._partitioned]
+    # not_parted = [name for (name, p) in model.named_parameters() if not p._partitioned]
+    # print('partitioned: ', parted)
+    # print('full:' , not_parted)
+    # print()
+
+    #samyamspeed.disable()
+
+
+test_driver()

tests/small_model_debugging/test.py

@@ -0,0 +1,52 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from deepspeed.pt.deepspeed_linear import LinearModuleForZeroStage3
+from deepspeed.pt.log_utils import logger
+from deepspeed.accelerator import get_accelerator
+
+
+def see_memory_usage(message):
+
+    # Print message except when distributed but not rank 0
+    logger.info(message)
+    logger.info(
+        "Memory Allocated %s GigaBytes ",
+        get_accelerator().memory_allocated() / (1024 * 1024 * 1024),
+    )
+    logger.info(
+        "Max Memory Allocated %s GigaBytes",
+        get_accelerator().max_memory_allocated() / (1024 * 1024 * 1024),
+    )
+    logger.info(
+        "Cache Allocated %s GigaBytes",
+        get_accelerator().memory_cached() / (1024 * 1024 * 1024),
+    )
+    logger.info(
+        "Max cache Allocated %s GigaBytes",
+        get_accelerator().max_memory_cached() / (1024 * 1024 * 1024),
+    )
+
+
+tens = torch.rand(1024, 16384, dtype=torch.half, device=torch.device(get_accelerator().device_name()))
+tens_back = tens.detach().clone()
+
+#linear_bk = torch.nn.functional.linear
+#torch.nn.functional.linear = deepspeed.pt.deepspeed_linear.LinearFunctionForZeroStage3.apply
+model = LinearModuleForZeroStage3(16384, 16384)
+
+model.to(get_accelerator().device_name()).half()
+
+see_memory_usage("Before forward")
+y = model(tens)
+
+see_memory_usage("After forward")
+
+model.weight.data = torch.zeros(1, dtype=torch.half, device=torch.device(get_accelerator().device_name()))
+
+see_memory_usage("After weight zero")
+
+y.backward(tens_back)

tests/small_model_debugging/test_mics.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+deepspeed test_mics_config.py --mics_shard_size=1
+
+deepspeed test_mics_config.py --mics_shard_size=2
+
+# for debugging the hierarchical params gathering
+export NDEV_PER_NODE=2
+deepspeed test_mics_config.py --mics_shard_size=4 --mics_hierarchical_params_gather

tests/small_model_debugging/test_mics_config.py

@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Testing on a 8 GPUs node
+NDEV_PER_NODE=2 torchrun --nnodes 1 --nproc-per-node 8 test_mics_config.py
+"""
+
+import os
+import json
+import argparse
+import torch
+import deepspeed
+from torch.utils.data.distributed import DistributedSampler
+import deepspeed.comm as dist
+
+
+class SimpleModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim, empty_grad=False):
+        super(SimpleModel, self).__init__()
+        self.linear = torch.nn.Linear(hidden_dim, hidden_dim)
+        if empty_grad:
+            self.layers2 = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim)])
+        self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        hidden = x
+        hidden = self.linear(hidden)
+        return self.cross_entropy_loss(hidden, y)
+
+
+def create_config_from_dict(tmpdir, config_dict):
+    config_path = os.path.join(tmpdir, 'temp_config.json')
+    with open(config_path, 'w') as fd:
+        json.dump(config_dict, fd)
+    return config_path
+
+
+def get_data_loader(model, total_samples, hidden_dim, device):
+    batch_size = model.train_micro_batch_size_per_gpu()
+    train_data = torch.randn(total_samples, hidden_dim, device=device, dtype=torch.float)
+    train_label = torch.empty(total_samples, dtype=torch.long, device=device).random_(hidden_dim)
+    train_dataset = torch.utils.data.TensorDataset(train_data, train_label)
+    sampler = DistributedSampler(train_dataset)
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=sampler)
+    return train_loader
+
+
+def get_args(tmpdir, config_dict):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--zero', type=int, default=3)
+    parser.add_argument('--local_rank', type=int)
+
+    parser.add_argument('--mics_shard_size', default=2, type=int)
+    parser.add_argument('--mics_hierarchical_params_gather', default=False, action='store_true')
+    args = parser.parse_args()  #args=''
+
+    config_dict["zero_optimization"]["stage"] = args.zero
+    config_dict["zero_optimization"]["mics_shard_size"] = args.mics_shard_size
+    config_dict["zero_optimization"]["mics_hierarchical_params_gather"] = args.mics_hierarchical_params_gather
+
+    # print('config_dict["zero_optimization"]', config_dict["zero_optimization"])
+    config_path = create_config_from_dict(tmpdir, config_dict)
+
+    args.deepspeed_config = config_path
+    return args
+
+
+def print0(msg):
+    if dist.get_rank() == 0:
+        print(msg, flush=True)
+
+
+rank = int(os.environ['RANK'])
+print('seed:', 2222 + rank)
+torch.random.manual_seed(2222 + rank)
+
+config_dict = {
+    "train_batch_size": 8,
+    "steps_per_print": 1,
+    "optimizer": {
+        "type": "Adam",
+        "params": {
+            "lr": 0.00015,
+        }
+    },
+    "fp16": {
+        "enabled": False,
+        "initial_scale_power": 15
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "reduce_bucket_size": 20,
+        "mics_shard_size": 4,
+        "mics_hierarchical_params_gather": True,
+        "stage3_model_persistence_threshold": 10
+    }
+}
+#        "initial_scale_power": 15
+args = get_args('/tmp/', config_dict)
+hidden_dim = 32
+
+with deepspeed.zero.MiCS_Init(config_dict_or_path=config_dict):
+    model = SimpleModel(hidden_dim, empty_grad=False)
+# print('------> init model with deepspeed.zero.Init()')
+
+model, _, _, _ = deepspeed.initialize(args=args,
+                                      model=model,
+                                      model_parameters=model.parameters(),
+                                      dist_init_required=True)
+
+
+def print_params(tag, model):
+    if dist.get_rank() == 0:
+        for n, p in model.named_parameters():
+            print0("{} {}:{}".format(tag, n, p))
+
+
+data_loader = get_data_loader(model=model, total_samples=1000, hidden_dim=hidden_dim, device=model.device)
+#print_params('pre-train', model)
+for n, batch in enumerate(data_loader):
+    loss = model(batch[0], batch[1])
+    if dist.get_rank() == 0:
+        print("LOSS:", loss.item())
+    model.backward(loss)
+    model.step()
+    #print_params('step={}'.format(n), model)
+    if n == 5: break

tests/small_model_debugging/test_model.py

@@ -0,0 +1,125 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import json
+import argparse
+import torch
+import deepspeed
+from torch.utils.data.distributed import DistributedSampler
+import deepspeed.comm as dist
+
+
+class SimpleModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim, empty_grad=False):
+        super(SimpleModel, self).__init__()
+        self.linear = torch.nn.Linear(hidden_dim, hidden_dim, bias=True)
+        self.linear = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
+        if empty_grad:
+            self.layers2 = torch.nn.ModuleList([torch.nn.Linear(hidden_dim,
+                                                                hidden_dim)])  #QuantizeLinear(hidden_dim, hidden_dim)
+        self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        hidden = x
+        hidden1 = self.linear(hidden)
+        hidden2 = self.linear(hidden1)
+        return self.cross_entropy_loss(hidden2, y)
+
+
+def create_config_from_dict(tmpdir, config_dict):
+    config_path = os.path.join(tmpdir, 'temp_config.json')
+    with open(config_path, 'w') as fd:
+        json.dump(config_dict, fd)
+    return config_path
+
+
+def get_data_loader(model, total_samples, hidden_dim, device):
+    batch_size = model.train_micro_batch_size_per_gpu()
+    train_data = torch.randn(total_samples, hidden_dim, device=device, dtype=torch.half)
+    train_label = torch.empty(total_samples, dtype=torch.long, device=device).random_(hidden_dim)
+    train_dataset = torch.utils.data.TensorDataset(train_data, train_label)
+    sampler = DistributedSampler(train_dataset)
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=sampler)
+    return train_loader
+
+
+def get_args(tmpdir, config_dict):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_rank", type=int, default=0)
+    parser.add_argument('--zero', type=int, default=0)
+    parser.add_argument('--zero_hpz_partition_size', type=int, default=1)
+    args = parser.parse_args()  #args=''
+
+    config_dict["zero_optimization"]["stage"] = args.zero
+    config_dict["zero_optimization"]["zero_hpz_partition_size"] = args.zero_hpz_partition_size
+    print('config_dict["zero_optimization"]', config_dict["zero_optimization"])
+    config_path = create_config_from_dict(tmpdir, config_dict)
+
+    args.deepspeed_config = config_path
+    return args
+
+
+def print0(msg):
+    if dist.get_rank() == 0:
+        print(msg, flush=True)
+
+
+rank = int(os.environ['RANK'])
+print('seed:', 2222 + rank)
+torch.random.manual_seed(2222 + rank)
+
+config_dict = {
+    "train_batch_size": 256,
+    "steps_per_print": 1,
+    "optimizer": {
+        "type": "Adam",
+        "params": {
+            "lr": 0.00015,
+        }
+    },
+    "fp16": {
+        "enabled": True,
+        "initial_scale_power": 8
+    },
+    "zero_optimization": {
+        "stage": 0,
+        "reduce_bucket_size": 20,
+        "zero_hpz_partition_size": 1,
+        "reduce_scatter": True,
+        "zero_quantized_weights": False,
+        "zero_quantized_gradients": False
+    }
+}
+#        "initial_scale_power": 15
+args = get_args('/tmp/', config_dict)
+hidden_dim = 4 * 1024
+
+model = SimpleModel(hidden_dim, empty_grad=False)
+
+model, _, _, _ = deepspeed.initialize(args=args,
+                                      model=model,
+                                      model_parameters=model.parameters(),
+                                      dist_init_required=True)
+
+
+def print_params(tag, model):
+    if dist.get_rank() == 0:
+        for n, p in model.named_parameters():
+            print0("{} {}:{}".format(tag, n, p))
+
+
+data_loader = get_data_loader(model=model, total_samples=256, hidden_dim=hidden_dim, device=model.device)
+#print_params('pre-train', model)
+
+for n, batch in enumerate(data_loader):
+    loss = model(batch[0], batch[1])
+    if dist.get_rank() == 0:
+        print("LOSS:", loss.item())
+    model.backward(loss)
+    model.step()
+    #print_params('step={}'.format(n), model)
+    #if n == 5: break

tests/torch_compile/ds_config_z2.json

@@ -0,0 +1,35 @@
+{
+  "train_batch_size": 8,
+  "steps_per_print": 2000,
+  "optimizer": {
+    "type": "Adam",
+    "params": {
+      "lr": 0.001,
+      "betas": [
+        0.8,
+        0.999
+      ],
+      "eps": 1e-8,
+      "weight_decay": 3e-7
+    }
+  },
+  "scheduler": {
+    "type": "WarmupLR",
+    "params": {
+      "warmup_min_lr": 0,
+      "warmup_max_lr": 0.001,
+      "warmup_num_steps": 1000
+    }
+  },
+  "gradient_clipping": 1.0,
+  "prescale_gradients": false,
+  "bf16": {
+      "enabled": true
+  },
+  "wall_clock_breakdown": false,
+  "zero_optimization": {
+      "stage": 2,
+      "overlap_comm": false,
+      "contiguous_gradients": false
+  }
+}

tests/torch_compile/ds_config_z3.json

@@ -0,0 +1,36 @@
+{
+  "train_batch_size": 8,
+  "steps_per_print": 2000,
+  "optimizer": {
+    "type": "Adam",
+    "params": {
+      "lr": 0.001,
+      "betas": [
+        0.8,
+        0.999
+      ],
+      "eps": 1e-8,
+      "weight_decay": 3e-7
+    }
+  },
+  "scheduler": {
+    "type": "WarmupLR",
+    "params": {
+      "warmup_min_lr": 0,
+      "warmup_max_lr": 0.001,
+      "warmup_num_steps": 1000
+    }
+  },
+  "gradient_clipping": 1.0,
+  "prescale_gradients": false,
+  "bf16": {
+      "enabled": true
+  },
+  "wall_clock_breakdown": false,
+  "zero_optimization": {
+      "stage": 3,
+      "reduce_scatter": true,
+      "overlap_comm": false,
+      "contiguous_gradients": false
+  }
+}

tests/torch_compile/ds_config_z3_deepcompile_no_persist.json

@@ -0,0 +1,41 @@
+{
+  "train_batch_size": 8,
+  "steps_per_print": 2000,
+  "optimizer": {
+    "type": "Adam",
+    "params": {
+      "lr": 0.001,
+      "betas": [
+        0.8,
+        0.999
+      ],
+      "eps": 1e-8,
+      "weight_decay": 3e-7
+    }
+  },
+  "scheduler": {
+    "type": "WarmupLR",
+    "params": {
+      "warmup_min_lr": 0,
+      "warmup_max_lr": 0.001,
+      "warmup_num_steps": 1000
+    }
+  },
+  "gradient_clipping": 1.0,
+  "prescale_gradients": false,
+  "bf16": {
+    "enabled": true
+  },
+  "wall_clock_breakdown": false,
+  "compile": {
+    "deepcompile": true
+  },
+  "zero_optimization": {
+    "stage": 3,
+    "reduce_scatter": true,
+    "overlap_comm": false,
+    "contiguous_gradients": false,
+    "stage3_param_persistence_threshold": 0,
+    "stage3_model_persistence_threshold": 0
+  }
+}

tests/torch_compile/test_compile.py

@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import argparse
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+from deepspeed import comm
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+torch._dynamo.config.cache_size_limit = 100
+
+
+def get_dynamo_stats():
+    return torch._dynamo.utils.counters["graph_break"]
+
+
+class RandomDataset(Dataset):
+
+    def __init__(self, size, length):
+        self.len = length
+        self.data = torch.randn(length, size).to(torch.bfloat16)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+    def __len__(self):
+        return self.len
+
+
+data_size = 1024
+data_length = 100
+rand_loader = DataLoader(dataset=RandomDataset(data_size, data_length), batch_size=1, shuffle=False)
+
+
+class MyModule(torch.nn.Module):
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.fc0 = torch.nn.Linear(1024, 256, bias=False)
+        self.fc1 = torch.nn.Linear(256, 256, bias=False)
+        self.dropout = torch.nn.Dropout(0.5)
+
+    def forward(self, data, residual):
+        output = residual + self.fc1(self.fc0(self.dropout(data))) * 0.5
+        return output
+
+
+model = MyModule()
+params = model.parameters()
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher')
+parser.add_argument('--deepspeed_config',
+                    type=str,
+                    default='ds_config_z3.json',
+                    help='path to DeepSpeed configuration file')
+cmd_args = parser.parse_args()
+
+# initialize the DeepSpeed engine
+model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args, model=model, model_parameters=params)
+model_engine.compile()
+
+residual = torch.rand(256, 256, dtype=torch.float).to(get_accelerator().current_device_name())
+
+start_stats = get_dynamo_stats()
+
+if comm.get_rank() == 0:
+    #print(dynamo_stats['graph_breaks'])
+    for item in start_stats.items():
+        print(item)
+
+for step, batch in enumerate(rand_loader):
+    if step % 10 == 0 and comm.get_rank() == 0:
+        print(f'step={step}')
+    # forward() method
+    loss = model_engine(batch.to(get_accelerator().current_device_name()), residual).sum()
+    # runs backpropagation
+    model_engine.backward(loss)
+    # weight update
+    model_engine.step()
+
+dynamo_stats = get_dynamo_stats()
+
+if comm.get_rank() == 0:
+    # print break down of graph break stats with markdown, print in table format, start with reason, then count
+    # print a tag 'dynamo_output' before each line to allow post processing
+    print("dynamo_output | Reason | Count |")
+    print("dynamo_output | ------ | ----- |")
+    for item in dynamo_stats.items():
+        # replace '|' in item[0] with a literal '|' to avoid mess with table format
+        item = (item[0].replace('|', r'\|'), item[1])
+        print(f"dynamo_output | {item[0]} | {item[1]} |")
+    print(f"dynamo_output | Total | {sum(dynamo_stats.values())} |")

tests/torch_compile/test_deepcompile_skipped_frame.py

@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Regression test for skipped eager frames under DeepCompile ZeRO-3."""
+
+import argparse
+import logging
+
+import torch
+
+import deepspeed
+from deepspeed import comm
+from deepspeed.accelerator import get_accelerator
+
+torch._dynamo.config.cache_size_limit = 100
+
+
+def configure_dynamo_logging():
+    try:
+        import torch._logging
+
+        torch._logging.set_logs(dynamo=logging.INFO, graph_breaks=True)
+        torch._dynamo.config.verbose = True
+    except Exception:
+        pass
+
+
+def dynamo_counter_text():
+    counters = torch._dynamo.utils.counters
+    return repr({str(key): dict(value) for key, value in counters.items()})
+
+
+class SkippedFrameModel(torch.nn.Module):
+
+    def __init__(self, vocab_size=384, hidden=384, n_layers=3):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.embed_tokens = torch.nn.Embedding(vocab_size, hidden)
+        self.layers = torch.nn.ModuleList([torch.nn.Linear(hidden, hidden, bias=False) for _ in range(n_layers)])
+        self.head = torch.nn.Linear(hidden, vocab_size, bias=False)
+
+    @torch.compiler.disable
+    def _compiler_disabled_forward(self, input_ids):
+        h = self.embed_tokens(input_ids)
+        for layer in self.layers:
+            h = layer(h)
+            h = torch.relu(h)
+        return self.head(h)
+
+    def forward(self, input_ids):
+        return self._compiler_disabled_forward(input_ids)
+
+
+def main():
+    configure_dynamo_logging()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_rank", type=int, default=-1)
+    parser.add_argument("--deepspeed_config", type=str, default="ds_config_z3_deepcompile_no_persist.json")
+    args = parser.parse_args()
+
+    model = SkippedFrameModel()
+    assert all(p.numel() > 100000 for p in model.parameters())
+
+    engine, _, _, _ = deepspeed.initialize(args=args, model=model, model_parameters=model.parameters())
+    torch._dynamo.reset()
+    torch._dynamo.utils.counters.clear()
+    engine.compile()
+
+    device = get_accelerator().current_device_name()
+    input_ids = torch.randint(0, model.vocab_size, (1, 16), device=device)
+
+    for step in range(3):
+        loss = engine(input_ids).sum()
+        engine.backward(loss)
+        engine.step()
+        if comm.get_rank() == 0:
+            print(f"step={step} loss={loss.item():.4f}")
+
+    counters = dynamo_counter_text()
+    fallback = getattr(engine, "_deepcompile_z3_eager_fallback", None)
+    fallback_stats = fallback.stats() if fallback is not None else {}
+    if comm.get_rank() == 0:
+        print(f"dynamo_counters={counters}")
+        print(f"fallback_stats={fallback_stats}")
+    assert "compiler.disable" in counters or "Skip inlining" in counters
+    assert fallback_stats.get("total_gathered_params", 0) > 0
+
+    if comm.get_rank() == 0:
+        print("PASS")
+
+
+if __name__ == "__main__":
+    main()

tests/torch_compile/test_deepcompile_z3_release.py

@@ -0,0 +1,130 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import math
+
+import pytest
+import torch
+
+import deepspeed.comm as dist
+from deepspeed.accelerator import get_accelerator
+from deepspeed.compile.config import CompileConfig
+from deepspeed.compile.util import get_deepcompile_handle, is_deepcompile_supported
+from unit.common import DistributedTest
+
+pytestmark = pytest.mark.skipif(not is_deepcompile_supported(),
+                                reason="DeepCompile requires CUDA and supported PyTorch")
+
+
+class TestDeepCompileZ3ReleaseStorage(DistributedTest):
+    world_size = 2
+    non_daemonic_procs = True
+
+    def _device(self):
+        return torch.device(get_accelerator().current_device_name())
+
+    def _init_dc(self):
+        dc = get_deepcompile_handle()
+        dc.init(dist.get_world_group(), CompileConfig(deepcompile=True), 1024)
+        return dc
+
+    def _register_param(self, dc, graph_id, ds_id, shape, persistent=False):
+        device = self._device()
+        world_size = dist.get_world_size()
+        true_numel = math.prod(shape)
+        shard_numel = math.ceil(true_numel / world_size)
+        rank = dist.get_rank()
+        values = torch.arange(rank * shard_numel, (rank + 1) * shard_numel, device=device, dtype=torch.float32)
+        grad_buffer = torch.zeros_like(values)
+        dc.register_z3_param(ds_id, list(shape), values, grad_buffer, persistent, values.dtype)
+        dc.register_graph_z3(graph_id, [ds_id])
+        return values
+
+    def _gather_view_and_storage(self, shard, graph_id, ds_id):
+        gathered = torch.ops.dc.allgather_param.default(shard, graph_id, ds_id)
+        gathered = torch.ops.dc.wait_allgather.default(gathered, graph_id, ds_id)
+        view = gathered.reshape(-1).narrow(0, 0, gathered.numel() - 1)
+        assert view.untyped_storage().data_ptr() == gathered.untyped_storage().data_ptr()
+        storage = view.untyped_storage()
+        assert storage.nbytes() >= gathered.numel() * gathered.element_size()
+        return view, storage
+
+    def _release(self, view, graph_id, ds_id, n_users, synchronize=True):
+        torch.ops.dc.release_param.default(view, graph_id, ds_id, n_users)
+        if synchronize:
+            get_accelerator().synchronize()
+
+    def _expected_view_sum(self, shape):
+        world_size = dist.get_world_size()
+        shard_numel = math.ceil(math.prod(shape) / world_size)
+        values = torch.arange(0, world_size * shard_numel, dtype=torch.float32, device=self._device())
+        values = values[:math.prod(shape)].reshape(-1)
+        return values.narrow(0, 0, values.numel() - 1).sum()
+
+    def test_storage_resized_to_zero_after_release_single_use(self):
+        graph_id, ds_id = 9010, 9011
+        dc = self._init_dc()
+        try:
+            shard = self._register_param(dc, graph_id, ds_id, [4097])
+            view, storage = self._gather_view_and_storage(shard, graph_id, ds_id)
+            self._release(view, graph_id, ds_id, 1)
+            assert storage.nbytes() == 0
+        finally:
+            dc.cleanup()
+
+    def test_storage_nonzero_until_final_release_when_multi_use(self):
+        graph_id, ds_id = 9020, 9021
+        dc = self._init_dc()
+        try:
+            shard = self._register_param(dc, graph_id, ds_id, [3])
+            view, storage = self._gather_view_and_storage(shard, graph_id, ds_id)
+            before_release_nbytes = storage.nbytes()
+            self._release(view, graph_id, ds_id, 2)
+            assert storage.nbytes() == before_release_nbytes
+            self._release(view, graph_id, ds_id, 2)
+            assert storage.nbytes() == 0
+        finally:
+            dc.cleanup()
+
+    def test_persistent_param_storage_unchanged_across_release(self):
+        graph_id, ds_id = 9030, 9031
+        dc = self._init_dc()
+        try:
+            shard = self._register_param(dc, graph_id, ds_id, [4], persistent=True)
+            view, storage = self._gather_view_and_storage(shard, graph_id, ds_id)
+            before_ptr = storage.data_ptr()
+            before_nbytes = storage.nbytes()
+            self._release(view, graph_id, ds_id, 1)
+            assert storage.data_ptr() == before_ptr
+            assert storage.nbytes() == before_nbytes
+        finally:
+            dc.cleanup()
+
+    def test_consumer_stream_can_finish_before_storage_reuse(self):
+        graph_id, ds_id = 9040, 9041
+        if not hasattr(torch.cuda, "_sleep"):  #ignore-cuda
+            pytest.skip("CUDA sleep helper is unavailable")
+        dc = self._init_dc()
+        try:
+            shard = self._register_param(dc, graph_id, ds_id, [4097])
+            view, storage = self._gather_view_and_storage(shard, graph_id, ds_id)
+            padded_bytes = storage.nbytes()
+            result = torch.empty((), device=self._device(), dtype=view.dtype)
+            consumer_stream = get_accelerator().Stream()
+            with get_accelerator().stream(consumer_stream):
+                torch.cuda._sleep(int(1e8))  #ignore-cuda
+                result.copy_(view.sum())
+                self._release(view, graph_id, ds_id, 1, synchronize=False)
+
+            scratch = torch.empty((padded_bytes // view.element_size()) + 1024,
+                                  device=self._device(),
+                                  dtype=view.dtype)
+            scratch.fill_(17)
+            get_accelerator().synchronize()
+            assert torch.allclose(result, self._expected_view_sum([4097]))
+            assert storage.nbytes() == 0
+            del scratch
+        finally:
+            dc.cleanup()

tests/unit/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+'''Copyright The Microsoft DeepSpeed Team'''

tests/unit/accelerator/test_accelerator.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+import os
+import sys
+import importlib
+import re
+
+import deepspeed
+
+DS_ACCEL_PATH = "deepspeed.accelerator"
+IGNORE_FILES = ["abstract_accelerator.py", "real_accelerator.py"]
+
+
+@pytest.fixture
+def accel_class_name(module_name):
+    class_list = []
+    mocked_modules = []
+
+    # Get the accelerator class name for a given module
+    while True:
+        try:
+            module = importlib.import_module(module_name)
+            break
+        except ModuleNotFoundError as e:
+            # If the environment is missing a module, mock it so we can still
+            # test importing the accelerator class
+            missing_module = re.search(r"\'(.*)\'", e.msg).group().strip("'")
+            sys.modules[missing_module] = lambda x: None
+            mocked_modules.append(missing_module)
+    for name in dir(module):
+        if name.endswith("_Accelerator"):
+            class_list.append(name)
+
+    assert len(class_list) == 1, f"Multiple accelerator classes found in {module_name}"
+
+    yield class_list[0]
+
+    # Clean up mocked modules so as to not impact other tests
+    for module in mocked_modules:
+        del sys.modules[module]
+
+
+@pytest.mark.parametrize(
+    "module_name",
+    [
+        DS_ACCEL_PATH + "." + f.rstrip(".py") for f in os.listdir(deepspeed.accelerator.__path__[0])
+        if f.endswith("_accelerator.py") and f not in IGNORE_FILES
+    ],
+)
+def test_abstract_methods_defined(module_name, accel_class_name):
+    module = importlib.import_module(module_name)
+    accel_class = getattr(module, accel_class_name)
+    accel_class.__init__ = lambda self: None
+    _ = accel_class()

tests/unit/alexnet_model.py

@@ -0,0 +1,150 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import deepspeed
+import deepspeed.comm as dist
+import deepspeed.runtime.utils as ds_utils
+from deepspeed.utils.torch import required_torch_version
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.pipe.module import PipelineModule, LayerSpec
+from .util import no_child_process_in_deepspeed_io
+
+
+class AlexNet(nn.Module):
+
+    def __init__(self, num_classes=10):
+        super(AlexNet, self).__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=5),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(64, 192, kernel_size=5, padding=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(192, 384, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        )
+        self.classifier = nn.Linear(256, num_classes)
+        self.loss_fn = nn.CrossEntropyLoss()
+
+    def forward(self, x, y):
+        x = self.features(x)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return self.loss_fn(x, y)
+
+
+class AlexNetPipe(AlexNet):
+
+    def to_layers(self):
+        layers = [*self.features, lambda x: x.view(x.size(0), -1), self.classifier]
+        return layers
+
+
+class AlexNetPipeSpec(PipelineModule):
+
+    def __init__(self, num_classes=10, **kwargs):
+        self.num_classes = num_classes
+        specs = [
+            LayerSpec(nn.Conv2d, 3, 64, kernel_size=11, stride=4, padding=5),
+            LayerSpec(nn.ReLU, inplace=True),
+            LayerSpec(nn.MaxPool2d, kernel_size=2, stride=2),
+            LayerSpec(nn.Conv2d, 64, 192, kernel_size=5, padding=2),
+            F.relu,
+            LayerSpec(nn.MaxPool2d, kernel_size=2, stride=2),
+            LayerSpec(nn.Conv2d, 192, 384, kernel_size=3, padding=1),
+            F.relu,
+            LayerSpec(nn.Conv2d, 384, 256, kernel_size=3, padding=1),
+            F.relu,
+            LayerSpec(nn.Conv2d, 256, 256, kernel_size=3, padding=1),
+            F.relu,
+            LayerSpec(nn.MaxPool2d, kernel_size=2, stride=2),
+            lambda x: x.view(x.size(0), -1),
+            LayerSpec(nn.Linear, 256, self.num_classes),  # classifier
+        ]
+        super().__init__(layers=specs, loss_fn=nn.CrossEntropyLoss(), **kwargs)
+
+
+# Define this here because we cannot pickle local lambda functions
+def cast_to_half(x):
+    return x.half()
+
+
+def cifar_trainset(fp16=False):
+    torchvision = pytest.importorskip("torchvision", minversion="0.5.0")
+    from torchvision import transforms
+
+    transform_list = [
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ]
+    if fp16:
+        transform_list.append(torchvision.transforms.Lambda(cast_to_half))
+
+    transform = transforms.Compose(transform_list)
+
+    local_rank = get_accelerator().current_device()
+
+    # Only one rank per machine downloads.
+    dist.barrier()
+    if local_rank != 0:
+        dist.barrier()
+    data_root = os.getenv("TEST_DATA_DIR", "/tmp/")
+    if os.getenv("CIFAR10_DATASET_PATH"):
+        data_root = os.getenv("CIFAR10_DATASET_PATH")
+        download = False
+    else:
+        data_root = os.path.join(os.getenv("TEST_DATA_DIR", "/tmp"), "cifar10-data")
+        download = True
+    trainset = torchvision.datasets.CIFAR10(root=data_root, train=True, download=download, transform=transform)
+    if local_rank == 0:
+        dist.barrier()
+    return trainset
+
+
+def train_cifar(model, config, num_steps=400, average_dp_losses=True, fp16=True, seed=123):
+    if required_torch_version(min_version=2.1):
+        fork_kwargs = {"device_type": get_accelerator().device_name()}
+    else:
+        fork_kwargs = {}
+    with get_accelerator().random().fork_rng(devices=[get_accelerator().current_device_name()], **fork_kwargs):
+        ds_utils.set_random_seed(seed)
+
+        # disable dropout
+        model.eval()
+
+        trainset = cifar_trainset(fp16=fp16)
+        config['local_rank'] = dist.get_rank()
+
+        with no_child_process_in_deepspeed_io():
+            engine, _, _, _ = deepspeed.initialize(config=config,
+                                                   model=model,
+                                                   model_parameters=[p for p in model.parameters()],
+                                                   training_data=trainset)
+
+        losses = []
+        for step in range(num_steps):
+            loss = engine.train_batch()
+            losses.append(loss.item())
+            if step % 50 == 0 and dist.get_rank() == 0:
+                print(f'STEP={step} LOSS={loss.item()}')
+
+        if average_dp_losses:
+            loss_tensor = torch.tensor(losses).to(get_accelerator().device_name())
+            dist.all_reduce(loss_tensor)
+            loss_tensor /= dist.get_world_size()
+            losses = loss_tensor.tolist()
+
+    return losses

tests/unit/autotuning/test_autotuning.py

@@ -0,0 +1,78 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import pytest
+from unit.simple_model import create_config_from_dict
+from deepspeed.launcher import runner as dsrun
+from deepspeed.autotuning.autotuner import Autotuner
+from deepspeed.autotuning.scheduler import ResourceManager
+
+RUN_OPTION = 'run'
+TUNE_OPTION = 'tune'
+
+
+def test_command_line():
+    '''Validate handling of command line arguments'''
+    for opt in [RUN_OPTION, TUNE_OPTION]:
+        dsrun.parse_args(args=f"--num_nodes 1 --num_gpus 1 --autotuning {opt} foo.py".split())
+
+    for error_opts in [
+            "--autotuning --num_nodes 1 --num_gpus 1 foo.py".split(),
+            "--autotuning test --num_nodes 1 -- num_gpus 1 foo.py".split(), "--autotuning".split()
+    ]:
+        with pytest.raises(SystemExit):
+            dsrun.parse_args(args=error_opts)
+
+
+@pytest.mark.parametrize("arg_mappings",
+                        [
+                            None,
+                            {
+                            },
+                            {
+                                "train_micro_batch_size_per_gpu": "--per_device_train_batch_size"
+                            },
+                            {
+                                "train_micro_batch_size_per_gpu": "--per_device_train_batch_size",
+                                "gradient_accumulation_steps": "--gradient_accumulation_steps"
+                            },
+                            {
+                                "train_batch_size": "-tbs"
+                            }
+                        ]) # yapf: disable
+def test_resource_manager_arg_mappings(arg_mappings):
+    rm = ResourceManager(args=None,
+                         hosts="worker-0, worker-1",
+                         num_gpus_per_node=4,
+                         results_dir=None,
+                         exps_dir=None,
+                         arg_mappings=arg_mappings)
+
+    if arg_mappings is not None:
+        for k, v in arg_mappings.items():
+            assert k.strip() in rm.arg_mappings.keys()
+            assert arg_mappings[k.strip()].strip() == rm.arg_mappings[k.strip()]
+
+
+@pytest.mark.parametrize("active_resources",
+                        [
+                           {"worker-0": [0, 1, 2, 3]},
+                           {"worker-0": [0, 1, 2, 3], "worker-1": [0, 1, 2, 3]},
+                           {"worker-0": [0], "worker-1": [0, 1, 2], "worker-2": [0, 1, 2]},
+                           {"worker-0": [0, 1], "worker-2": [4, 5]}
+                        ]
+                        ) # yapf: disable
+def test_autotuner_resources(tmpdir, active_resources):
+    config_dict = {"autotuning": {"enabled": True, "exps_dir": os.path.join(tmpdir, 'exps_dir'), "arg_mappings": {}}}
+    config_path = create_config_from_dict(tmpdir, config_dict)
+    args = dsrun.parse_args(args=f'--autotuning {TUNE_OPTION} foo.py --deepspeed_config {config_path}'.split())
+    tuner = Autotuner(args=args, active_resources=active_resources)
+
+    expected_num_nodes = len(list(active_resources.keys()))
+    assert expected_num_nodes == tuner.exp_num_nodes
+
+    expected_num_gpus = min([len(v) for v in active_resources.values()])
+    assert expected_num_gpus == tuner.exp_num_gpus

tests/unit/checkpoint/common.py

@@ -0,0 +1,255 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import torch
+import numbers
+
+import deepspeed
+from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
+from deepspeed.runtime.fp16.fused_optimizer import FP16_Optimizer
+from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer
+from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
+from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+
+from unit.common import preferred_dtype
+from unit.simple_model import *
+from unittest.mock import MagicMock, patch
+
+
+def compare_deepspeed_states(saved_model, loaded_model):
+    # These are compared in more depth in other places
+    assert hasattr(loaded_model, 'module')
+
+    assert saved_model.sparse_tensor_module_names == loaded_model.sparse_tensor_module_names
+    assert saved_model.skipped_steps == loaded_model.skipped_steps
+    assert saved_model.global_steps == loaded_model.global_steps
+
+
+def zero3_params_to_fetch(param_list):
+    return [p for p in param_list if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE]
+
+
+def compare_model_states(saved_model, loaded_model, compare_optimizer=True, load_module_only=False):
+    if not load_module_only:
+        compare_deepspeed_states(saved_model, loaded_model)
+
+    params_to_fetch = zero3_params_to_fetch(
+        list(saved_model.module.named_parameters()) + list(loaded_model.module.named_parameters()))
+    enable_gather = len(params_to_fetch) > 0
+    with deepspeed.zero.GatheredParameters(params_to_fetch, enabled=enable_gather):
+        for p0, p1 in zip(saved_model.module.named_parameters(), loaded_model.module.named_parameters()):
+            np0, p0 = p0
+            np1, p1 = p1
+            if 'deepspeed_moe.gate.wg' in np0:
+                # these params are converted to float at runtime, cast to half for comparison
+                p1 = p1.half()
+                p0 = p0.half()
+            assert id(p0) != id(p1), f'Comparing fp16 model state tensor against itself : {id(p0)} <====> {id(p1)}'
+            try:
+                assert torch.allclose(p0, p1,
+                                      atol=1e-07), f"FP16 model state {p0} is not equal to {p1}, names:{np0}, {np1}"
+            except RuntimeError as err:
+                print(f"FP16 model state {p0} is not equal to {p1}, names:{np0}, {np1}")
+                raise err
+
+    if not compare_optimizer:
+        return
+
+    if DeepSpeedZeroOptimizer_Stage3 is not None and isinstance(saved_model.optimizer, DeepSpeedZeroOptimizer_Stage3):
+        for p0, p1 in zip(saved_model.optimizer.fp32_partitioned_groups_flat,
+                          loaded_model.optimizer.fp32_partitioned_groups_flat):
+            assert torch.allclose(p0, p1, atol=1e-07), f"Fp32 model states {p0} is not equal to {p1}"
+
+    elif isinstance(saved_model.optimizer, DeepSpeedZeroOptimizer):
+        for p0, p1 in zip(saved_model.optimizer.single_partition_of_fp32_groups,
+                          loaded_model.optimizer.single_partition_of_fp32_groups):
+            assert id(p0) != id(p1), f'Comparing fp32 model state tensor against itself: {id(p0)} <====> {id(p1)}'
+            assert torch.allclose(p0, p1, atol=1e-07), f"Fp32 model states {p0} is not equal to {p1}"
+
+    elif isinstance(saved_model.optimizer, FP16_Optimizer):
+        for p0, p1 in zip(saved_model.optimizer.fp32_groups_flat, loaded_model.optimizer.fp32_groups_flat):
+            assert id(p0) != id(p1), f'Comparing fp32 model state tensor against itself: {id(p0)} <====> {id(p1)}'
+            assert torch.allclose(p0, p1, atol=1e-07), f"FP32 model states {p0} is not equal to {p1}"
+
+    elif isinstance(saved_model.optimizer, FP16_UnfusedOptimizer):
+        for params0, params1 in zip(saved_model.optimizer.fp32_groups, loaded_model.optimizer.fp32_groups):
+            for p0, p1 in zip(params0, params1):
+                assert id(p0) != id(p1), f'Comparing fp32 model state tensor against itself: {id(p0)} <====> {id(p1)}'
+                assert torch.allclose(p0, p1, atol=1e-07), f"FP32 model states {p0} is not equal to {p1}"
+    elif isinstance(saved_model.optimizer, torch.optim.Optimizer):
+        pass
+    else:
+        assert False, f'Unexpected Optimizer Type: {saved_model.optimizer}'
+
+
+def compare_state_dicts(state0, state1, expected_mismatch_keys=[]):
+    key_set0 = set(k for k in state0.keys() if k not in expected_mismatch_keys)
+    key_set1 = set(k for k in state1.keys() if k not in expected_mismatch_keys)
+    assert key_set0 == key_set1, f'failure due to key mismatch {key_set0} != {key_set1}'
+
+    for k in key_set0:
+        s0 = state0[k]
+        s1 = state1[k]
+        if k in expected_mismatch_keys:
+            continue
+        if isinstance(s0, torch.Tensor) and isinstance(s1, torch.Tensor):
+            assert id(s0) != id(s1), f'Comparing optimizer state tensor against itself: {id(s0)} <====> {id(s1)}'
+            assert torch.equal(s0.to('cpu'), s1.to('cpu'))
+        else:
+            assert s0 == s1, f'failures with keys = {k}, {k}, values = {s0} and {s1}'
+
+
+def compare_opt_state_dicts(state0, state1, expected_mismatch_keys=[]):
+    for param_group0, saved_param_group1 in zip(state0['param_groups'], state1['param_groups']):
+        compare_state_dicts(param_group0, saved_param_group1, expected_mismatch_keys)
+
+    assert "state" in state0
+    assert "state" in state1
+    assert len([state0["state"].keys()]) == len([state1["state"].keys()])
+
+    for (k0, s0), (k1, s1) in zip(state0["state"].items(), state1["state"].items()):
+        assert k0 == k1, f'failure due to key mismatch {k0} != {k1}'
+        compare_state_dicts(s0, s1, expected_mismatch_keys)
+
+
+def compare_optimizer_states(saved_model, loaded_model, hidden_dim, fp16=True):
+    saved_optimizer = saved_model.optimizer.optimizer if fp16 else saved_model.optimizer
+    loaded_optimizer = loaded_model.optimizer.optimizer if fp16 else loaded_model.optimizer
+
+    for state0, state1 in zip(saved_optimizer.state.values(), loaded_optimizer.state.values()):
+        compare_state_dicts(state0, state1)
+
+
+def compare_lr_scheduler_states(saved_model, loaded_model):
+    assert hasattr(saved_model, 'lr_scheduler')
+    assert hasattr(loaded_model, 'lr_scheduler')
+
+    saved_scheduler = saved_model.lr_scheduler
+    loaded_scheduler = loaded_model.lr_scheduler
+
+    assert hasattr(saved_scheduler, 'state_dict')
+    assert hasattr(loaded_scheduler, 'state_dict')
+
+    saved_sd = saved_scheduler.state_dict()
+    loaded_sd = loaded_scheduler.state_dict()
+
+    print(f"saved_sd = {saved_sd}")
+    print(f"loaded_sd = {loaded_sd}")
+
+    assert saved_sd.keys() == loaded_sd.keys()
+
+    for state0, state1 in zip(saved_sd.values(), loaded_sd.values()):
+        if isinstance(state0, numbers.Number) and isinstance(state1, numbers.Number):
+            assert state0 == state1
+
+
+# following mixture-of-experts.md
+def create_moe_param_groups(model):
+    from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer
+
+    parameters = {'params': [p for p in model.parameters()], 'name': 'parameters'}
+    return split_params_into_different_moe_groups_for_optimizer(parameters)
+
+
+def create_deepspeed_model(config_dict, model, base_optimizer):
+    ds_model, _, _, _ = deepspeed.initialize(config=config_dict,
+                                             model=model,
+                                             model_parameters=create_moe_param_groups(model),
+                                             optimizer=base_optimizer)
+    ds_model.empty_partition_cache()
+    return ds_model
+
+
+def checkpoint_correctness_verification(config_dict,
+                                        models,
+                                        hidden_dim,
+                                        tmpdir,
+                                        load_optimizer_states=False,
+                                        load_lr_scheduler_states=False,
+                                        train_batch=False,
+                                        base_optimizers=[None, None],
+                                        empty_tag=False,
+                                        seq_dataloader=False,
+                                        load_module_only=False,
+                                        dtype=None):
+    if dtype is None:
+        dtype = preferred_dtype()
+
+    ds_model = create_deepspeed_model(config_dict=config_dict, model=models[0], base_optimizer=base_optimizers[0])
+
+    if seq_dataloader:
+        data_loader = sequence_dataloader(model=ds_model,
+                                          total_samples=50,
+                                          hidden_dim=hidden_dim,
+                                          device=ds_model.device,
+                                          dtype=dtype)
+    else:
+        data_loader = random_dataloader(model=ds_model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=ds_model.device,
+                                        dtype=dtype)
+
+    if train_batch:
+        ds_model.set_dataloader(data_loader)
+        for _, batch in enumerate(data_loader):
+            loss = ds_model.train_batch()
+    else:
+        for _, batch in enumerate(data_loader):
+            loss = ds_model(batch[0], batch[1])
+            ds_model.backward(loss)
+            ds_model.step()
+
+    # Flush zero stage 3 cache
+    ds_model.empty_partition_cache()
+
+    trained_model = ds_model
+
+    save_folder = os.path.join(tmpdir, 'saved_checkpoint')
+    save_tag = None if empty_tag else '1'
+
+    trained_model.save_checkpoint(save_folder, tag=save_tag)
+
+    dist.barrier()
+
+    for root, _, files in os.walk(save_folder):
+        for f in files:
+            if "_expert_" in f and "_model_states" in f:
+                expert = torch.load(os.path.join(root, f), weights_only=False)
+                needed, storages = 0, {}
+                for name, tensor in expert.items():
+                    needed += tensor.size().numel()
+                    storage = tensor.storage()
+                    # some storage can be shared within an expert's checkpoint
+                    storages[storage.data_ptr()] = storage.size()
+                stored = sum(v for _, v in storages.items())
+                assert needed == stored, f"MoE expert checkpoint uses more storage than required: {f}"
+
+    loaded_model = create_deepspeed_model(config_dict=config_dict, model=models[1], base_optimizer=base_optimizers[1])
+    assert list(trained_model.parameters())[0].dtype == list(loaded_model.parameters())[0].dtype
+
+    context = patch.object(loaded_model, "_get_optimizer_ckpt_name",
+                           wraps=loaded_model._get_optimizer_ckpt_name) if not load_optimizer_states else MagicMock()
+    with context as optim_load_state_dict_mock:
+        loaded_model.load_checkpoint(save_folder,
+                                     tag=save_tag,
+                                     load_optimizer_states=load_optimizer_states,
+                                     load_lr_scheduler_states=load_lr_scheduler_states,
+                                     load_module_only=load_module_only)
+        if not load_optimizer_states:
+            # should not attempt to get the file name to load it
+            optim_load_state_dict_mock.assert_not_called()
+
+    compare_model_states(trained_model,
+                         loaded_model,
+                         compare_optimizer=load_optimizer_states,
+                         load_module_only=load_module_only)
+
+    if load_optimizer_states:
+        compare_optimizer_states(trained_model, loaded_model, hidden_dim, dtype == torch.float16)
+
+    if load_lr_scheduler_states:
+        compare_lr_scheduler_states(trained_model, loaded_model)

tests/unit/checkpoint/test_autotp_uc_checkpoint.py

@@ -0,0 +1,312 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import types
+from types import SimpleNamespace
+
+import torch
+
+from deepspeed.checkpoint.constants import (CAT_DIM, FP32_WEIGHT_KEY, PARAM, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS,
+                                            PARAMETER_WITH_SUB_PARAMS, SUB_PARAM_SHAPE,
+                                            TP_REPLICATED_PARAMETER_PATTERNS, UNIVERSAL_CHECKPOINT_INFO)
+from deepspeed.checkpoint.universal_checkpoint import SubparamShape as CheckpointSubparamShape
+from deepspeed.checkpoint.ds_to_universal import merge_tp_slices
+from deepspeed.checkpoint.universal_checkpoint import (_get_param_uc_restore_meta, _resolve_autotp_partition,
+                                                       load_hp_checkpoint_state)
+from deepspeed.runtime.bf16_optimizer import BF16_Optimizer
+from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
+
+
+class _DummyAddress:
+
+    def __init__(self, start, numel):
+        self.start = start
+        self.numel = numel
+
+
+class _DummyHPMapping:
+
+    def __init__(self, param):
+        self.lp_fragment_address = _DummyAddress(0, param.numel())
+        self._param = param
+        self.optim_fragment = {}
+
+    def get_hp_fragment(self):
+        return self._param.view(-1)
+
+    def get_optim_state_keys(self):
+        return []
+
+
+def _make_param(shape, meta=None):
+    param = torch.nn.Parameter(torch.zeros(shape, dtype=torch.float32))
+    param._hp_mapping = _DummyHPMapping(param)
+    if meta is not None:
+        setattr(param, 'ds_autotp_universal_checkpoint_meta', meta)
+    return param
+
+
+def test_resolve_autotp_partition_row_parallel_weight():
+    param = _make_param(
+        (4, 4), {
+            'partition_type': 'row',
+            'partition_dim': 1,
+            'logical_shape': (4, 8),
+            'output_shape': (4, ),
+            'sub_param_shape': None,
+            'original_shape': (4, 8),
+            'is_bias': False,
+            'replicated': False,
+        })
+    full_hp_param = torch.arange(32, dtype=torch.float32).view(4, 8)
+
+    slice_flat = _resolve_autotp_partition(param, {PARAM: full_hp_param}, full_hp_param, tp_rank=1, tp_world_size=2)
+
+    expected = full_hp_param.chunk(2, dim=1)[1].flatten()
+    assert torch.equal(slice_flat, expected)
+
+
+def test_resolve_autotp_partition_subparam_column_weight():
+    param = _make_param(
+        (3, 4), {
+            'partition_type': 'column',
+            'partition_dim': 0,
+            'logical_shape': (6, 4),
+            'output_shape': (6, ),
+            'sub_param_shape': ((2, 2, 2), 4),
+            'original_shape': (6, 4),
+            'is_bias': False,
+            'replicated': False,
+        })
+    full_hp_param = torch.arange(24, dtype=torch.float32).view(6, 4)
+
+    slice_flat = _resolve_autotp_partition(param, {PARAM: full_hp_param}, full_hp_param, tp_rank=0, tp_world_size=2)
+
+    chunks = [sub.chunk(2, dim=0)[0] for sub in full_hp_param.view(3, 2, 4)]
+    expected = torch.cat(chunks, dim=0).flatten()
+    assert torch.equal(slice_flat, expected)
+
+
+def test_resolve_autotp_partition_subparam_sizes_uneven_gqa_like():
+    # Simulate a fused QKV weight where Q/K/V have uneven sizes along partition_dim=0.
+    # Example (GQA-like):
+    #   Q: 8
+    #   K: 4
+    #   V: 4
+    # Total: 16
+    #
+    # With tp_world_size=2, correct slicing is:
+    #   Q chunk -> 4 per rank
+    #   K chunk -> 2 per rank
+    #   V chunk -> 2 per rank
+    # Each rank gets 8 rows total, but importantly boundaries must align with Q/K/V.
+    sub_param_sizes = [8, 4, 4]
+    tp_world_size = 2
+    tp_rank = 1
+
+    param = _make_param(
+        (8, 2),
+        {
+            "partition_type": "column",
+            "partition_dim": 0,
+            "logical_shape": (sum(sub_param_sizes), 2),  # (16, 2)
+            "output_shape": (sum(sub_param_sizes), ),  # (16,)
+            "sub_param_shape": (tuple(sub_param_sizes), 2),
+            "sub_param_sizes": sub_param_sizes,
+            "original_shape": (sum(sub_param_sizes), 2),
+            "is_bias": False,
+            "replicated": False,
+        })
+
+    # Full (unsharded) HP parameter: shape (16, 2)
+    full_hp_param = torch.arange(sum(sub_param_sizes) * 2, dtype=torch.float32).view(sum(sub_param_sizes), 2)
+
+    slice_flat = _resolve_autotp_partition(param, {PARAM: full_hp_param},
+                                           full_hp_param,
+                                           tp_rank=tp_rank,
+                                           tp_world_size=tp_world_size)
+
+    # Expected: split into Q/K/V blocks, chunk each block by TP, take tp_rank slice, concat back.
+    q, k, v = torch.split(full_hp_param, sub_param_sizes, dim=0)
+    expected = torch.cat([
+        q.chunk(tp_world_size, dim=0)[tp_rank],
+        k.chunk(tp_world_size, dim=0)[tp_rank],
+        v.chunk(tp_world_size, dim=0)[tp_rank]
+    ],
+                         dim=0).flatten()
+
+    assert torch.equal(slice_flat, expected)
+
+
+def test_resolve_autotp_partition_replicated_bias():
+    full_hp_param = torch.arange(8, dtype=torch.float32)
+    param = _make_param(
+        (8, ), {
+            'partition_type': 'row',
+            'partition_dim': None,
+            'logical_shape': (8, ),
+            'output_shape': (8, ),
+            'sub_param_shape': None,
+            'original_shape': (8, ),
+            'is_bias': True,
+            'replicated': True,
+        })
+
+    slice_flat = _resolve_autotp_partition(param, {PARAM: full_hp_param}, full_hp_param, tp_rank=1, tp_world_size=2)
+
+    assert torch.equal(slice_flat, full_hp_param)
+
+
+def test_load_hp_checkpoint_state_prefers_autotp_metadata(tmp_path, monkeypatch):
+    param = _make_param(
+        (4, 4), {
+            'partition_type': 'row',
+            'partition_dim': 1,
+            'logical_shape': (4, 8),
+            'output_shape': (4, ),
+            'sub_param_shape': None,
+            'original_shape': (4, 8),
+            'is_bias': False,
+            'replicated': False,
+        })
+    param.load_hp_checkpoint_state = types.MethodType(load_hp_checkpoint_state, param)
+
+    import deepspeed.checkpoint.universal_checkpoint as uc
+    monkeypatch.setattr(uc, "current_param", param, raising=False)
+
+    ckpt_dir = tmp_path / "weight"
+    ckpt_dir.mkdir(parents=True)
+    full_hp_param = torch.arange(32, dtype=torch.float32).view(4, 8)
+    torch.save({PARAM: full_hp_param}, ckpt_dir / f"{FP32_WEIGHT_KEY}.pt")
+
+    monkeypatch.setattr(
+        torch,
+        "load",
+        lambda *args, **kwargs: {PARAM: full_hp_param} if str(args[0]).endswith("fp32.pt") else 0,
+    )
+
+    step = param.load_hp_checkpoint_state(str(ckpt_dir), tp_rank=1, tp_world_size=2)
+
+    assert step is None
+    expected = full_hp_param.chunk(2, dim=1)[1].flatten()
+    assert torch.equal(param.data.flatten(), expected)
+
+
+def _write_tp_slice(base_dir, param_name, tp_idx, state_name, tensor):
+    shard_dir = base_dir / param_name / str(tp_idx)
+    shard_dir.mkdir(parents=True, exist_ok=True)
+    torch.save(tensor.reshape(-1), shard_dir / f"{state_name}.00")
+
+
+def _write_tp_states(base_dir, param_name, tp_idx, fp32_tensor):
+    # merge_tp_slices attempts to merge these three states, so the test must write all of them.
+    _write_tp_slice(base_dir, param_name, tp_idx, "fp32", fp32_tensor)
+    _write_tp_slice(base_dir, param_name, tp_idx, "exp_avg", torch.zeros_like(fp32_tensor))
+    _write_tp_slice(base_dir, param_name, tp_idx, "exp_avg_sq", torch.zeros_like(fp32_tensor))
+
+
+def test_merge_tp_slices_emits_subparam_shape_metadata(tmp_path):
+    slice_dir = tmp_path / "slices"
+    output_dir = tmp_path / "out"
+    param_name = "module.qkv.weight"
+
+    tp0 = torch.arange(12, dtype=torch.float32).view(3, 4)
+    tp1 = torch.arange(12, 24, dtype=torch.float32).view(3, 4)
+    _write_tp_states(slice_dir, param_name, 0, tp0)
+    _write_tp_states(slice_dir, param_name, 1, tp1)
+
+    uc_info = {
+        PARAMETER_WITH_ROW_PARALLELISM_PATTERNS: [],
+        TP_REPLICATED_PARAMETER_PATTERNS: [],
+        PARAMETER_WITH_SUB_PARAMS: [{
+            "patterns": [rf"^{param_name}$"],
+            "shape": [(2, 2, 2), 4],
+            "partition_dim": 0,
+        }],
+    }
+
+    ds_checkpoint = SimpleNamespace(
+        get_checkpoint_info=lambda key: uc_info if key == UNIVERSAL_CHECKPOINT_INFO else {})
+
+    unmatched = merge_tp_slices(ds_checkpoint, str(output_dir), str(slice_dir), 2, (param_name, torch.Size([3, 4])))
+
+    ckpt = torch.load(output_dir / param_name / "fp32.pt", weights_only=False)
+    assert not unmatched
+    assert isinstance(ckpt[SUB_PARAM_SHAPE], CheckpointSubparamShape)
+    assert ckpt[SUB_PARAM_SHAPE].partition_dim == 0
+
+
+def test_merge_tp_slices_uses_row_parallel_cat_dim(tmp_path):
+    slice_dir = tmp_path / "slices"
+    output_dir = tmp_path / "out"
+    param_name = "module.proj.weight"
+
+    tp0 = torch.arange(16, dtype=torch.float32).view(4, 4)
+    tp1 = torch.arange(16, 32, dtype=torch.float32).view(4, 4)
+    _write_tp_states(slice_dir, param_name, 0, tp0)
+    _write_tp_states(slice_dir, param_name, 1, tp1)
+
+    uc_info = {
+        PARAMETER_WITH_ROW_PARALLELISM_PATTERNS: [rf"^{param_name}$"],
+        TP_REPLICATED_PARAMETER_PATTERNS: [],
+        PARAMETER_WITH_SUB_PARAMS: [],
+    }
+
+    ds_checkpoint = SimpleNamespace(
+        get_checkpoint_info=lambda key: uc_info if key == UNIVERSAL_CHECKPOINT_INFO else {})
+
+    merge_tp_slices(ds_checkpoint, str(output_dir), str(slice_dir), 2, (param_name, torch.Size([4, 4])))
+
+    ckpt = torch.load(output_dir / param_name / "fp32.pt", weights_only=False)
+    assert ckpt[CAT_DIM] == 1
+    assert torch.equal(ckpt[PARAM], torch.cat([tp0, tp1], dim=1))
+
+
+def test_zero_optimizer_uc_info_comes_from_cached_state():
+    param = _make_param((2, 2))
+    expected_uc_info = {"key": "value"}
+    setattr(param, UNIVERSAL_CHECKPOINT_INFO, expected_uc_info)
+
+    optimizer = object.__new__(DeepSpeedZeroOptimizer)
+    optimizer.bit16_groups = [[param]]
+    optimizer._enable_universal_checkpoint()
+    delattr(param, UNIVERSAL_CHECKPOINT_INFO)
+
+    assert optimizer._get_universal_checkpoint_info() == expected_uc_info
+
+
+def test_bf16_optimizer_uc_info_comes_from_cached_state():
+    param = _make_param((2, 2))
+    expected_uc_info = {"key": "value"}
+    setattr(param, UNIVERSAL_CHECKPOINT_INFO, expected_uc_info)
+
+    optimizer = object.__new__(BF16_Optimizer)
+    optimizer.bf16_groups = [[param]]
+    optimizer._enable_universal_checkpoint()
+    delattr(param, UNIVERSAL_CHECKPOINT_INFO)
+
+    assert optimizer._get_universal_checkpoint_info() == expected_uc_info
+
+
+def test_get_param_uc_restore_meta_returns_top_level_restore_schema():
+    meta = {
+        "partition_dim": 1,
+        "logical_shape": (4, 8),
+        "output_shape": (4, ),
+        "sub_param_shape": None,
+        "sub_param_sizes": None,
+        "target_partition_shape": (4, 4),
+        "is_bias": False,
+        "replicated": False,
+        "conversion": {
+            "partition_dim": 999
+        },
+    }
+    param = _make_param((4, 4), meta)
+
+    restore_meta = _get_param_uc_restore_meta(param)
+
+    assert restore_meta["partition_dim"] == 1
+    assert restore_meta["conversion"]["partition_dim"] == 999

tests/unit/checkpoint/test_convert_checkpoint.py

@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import torch.nn as nn
+
+import deepspeed
+from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict
+from unit.common import DistributedTest
+
+
+class ModelWithSharedWeights(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.layer0 = nn.Linear(100, 100)
+        self.layer1 = nn.Linear(200, 200)
+        self.layer2 = nn.Linear(300, 300)
+        # tie layer 1 and layer 2
+        self.layer1.weight = self.layer2.weight
+
+
+class TestCheckpointConvert(DistributedTest):
+    world_size = 2
+
+    def test_convert_zero_checkpoint_to_fp32_state_dict(self, tmpdir):
+        config = {
+            "train_micro_batch_size_per_gpu": 2,
+            "zero_allow_untested_optimizer": True,
+            "zero_optimization": {
+                "stage": 3
+            },
+        }
+        model = ModelWithSharedWeights()
+        optimizer = torch.optim.Adam(model.parameters())
+
+        deepspeed_engine, _, _, _ = deepspeed.initialize(
+            config=config,
+            model=model,
+            optimizer=optimizer,
+        )
+        ds_save_dir = tmpdir / "checkpoint_ds"
+        deepspeed_engine.save_checkpoint(ds_save_dir, tag="checkpoint")
+
+        model = ModelWithSharedWeights()
+
+        # save checkpoint
+        fp32_save_dir = tmpdir / "checkpoint_fp32"
+        convert_zero_checkpoint_to_fp32_state_dict(ds_save_dir, fp32_save_dir)
+
+        # load state_dict from fp32 checkpoint
+        state_dict = torch.load(fp32_save_dir / 'pytorch_model.bin')
+
+        # check shared tensor
+        assert id(state_dict['layer1.weight']) == id(state_dict['layer2.weight'])
+
+        # load state_dict into model
+        model.load_state_dict(state_dict, strict=True)

tests/unit/checkpoint/test_latest_checkpoint.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed
+
+import pytest
+from unit.common import DistributedTest
+from unit.simple_model import *
+
+from unit.checkpoint.common import checkpoint_correctness_verification
+from deepspeed.ops.op_builder import FusedAdamBuilder
+
+if not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME]:
+    pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+
+class TestLatestCheckpoint(DistributedTest):
+    world_size = 1
+
+    def test_existing_latest(self, tmpdir):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            }
+        }
+        hidden_dim = 10
+        models = [SimpleModel(hidden_dim=hidden_dim) for _ in range(2)]
+        checkpoint_correctness_verification(config_dict=config_dict,
+                                            models=models,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=True,
+                                            load_lr_scheduler_states=False,
+                                            empty_tag=True,
+                                            dtype=torch.float)
+
+    def test_missing_latest(self, tmpdir):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            }
+        }
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        # should be no-op, since latest doesn't exist
+        model.load_checkpoint(tmpdir)

tests/unit/checkpoint/test_lr_scheduler.py

@@ -0,0 +1,120 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed
+from deepspeed.ops.op_builder import CPUAdamBuilder
+from deepspeed.accelerator import get_accelerator
+
+from unit.common import DistributedTest
+from unit.simple_model import *
+
+from unit.checkpoint.common import checkpoint_correctness_verification
+
+import pytest
+
+
+@pytest.mark.parametrize('zero_stage, use_cpu_offload', [(0, False), (1, False), (2, False), (2, True), (3, False),
+                                                         (3, True)])
+class TestLRSchedulerCheckpoint(DistributedTest):
+    world_size = 2
+
+    def test_checkpoint_lr_scheduler(self, tmpdir, zero_stage, use_cpu_offload):
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+        if get_accelerator().device_name() == 'cpu':
+            pytest.skip("CPU accelerator does not support this test.")
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "cpu_offload": use_cpu_offload
+            },
+            "scheduler": {
+                "type": "WarmupLR",
+                "params": {
+                    "warmup_min_lr": 0,
+                    "warmup_max_lr": 0.001,
+                    "warmup_num_steps": 1000
+                }
+            }
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+        hidden_dim = 10
+
+        if zero_stage == 3:
+            global DeepSpeedZeroOptimizer_Stage3
+            from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
+            with deepspeed.zero.Init(config_dict_or_path=config_dict):
+                models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+        else:
+            models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict,
+                                            models,
+                                            hidden_dim,
+                                            tmpdir,
+                                            load_optimizer_states=False,
+                                            load_lr_scheduler_states=True)
+
+    def test_checkpoint_no_lr_scheduler(self, tmpdir, zero_stage, use_cpu_offload):
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+        if get_accelerator().device_name() == 'cpu':
+            pytest.skip("CPU accelerator does not support this test.")
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 1e-5
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "cpu_offload": use_cpu_offload
+            },
+            "scheduler": {
+                "type": "WarmupLR",
+                "params": {
+                    "warmup_min_lr": 0,
+                    "warmup_max_lr": 0.001,
+                    "warmup_num_steps": 1000
+                }
+            },
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+        hidden_dim = 10
+
+        if zero_stage == 3:
+            with deepspeed.zero.Init(config_dict_or_path=config_dict):
+                models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+        else:
+            models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict,
+                                            models,
+                                            hidden_dim,
+                                            tmpdir,
+                                            load_optimizer_states=False,
+                                            load_lr_scheduler_states=False)

tests/unit/checkpoint/test_mics_optimizer.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import deepspeed
+
+from deepspeed.utils.torch import required_torch_version
+from unit.common import DistributedTest
+from unit.simple_model import *
+from unit.checkpoint.common import *
+
+import pytest
+
+if not required_torch_version(max_version=2.0):
+    pytest.skip("Skipping until we resolve problems with torch 2.1", allow_module_level=True)
+
+
+class TestMiCSCheckpoint(DistributedTest):
+    world_size = 4
+
+    def _toy_model_config(self, shard_size):
+
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 8
+            },
+            "wall_clock_breakdown": True,
+            "zero_optimization": {
+                "stage": 3,
+                "mics_shard_size": shard_size
+            }
+        }
+
+        hidden_dim = 10
+        with deepspeed.zero.MiCS_Init(config_dict_or_path=config_dict):
+            models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        return config_dict, hidden_dim, models
+
+    @pytest.mark.parametrize('shard_size', [1, 2, 4])
+    def test_load_optimizer_state(self, tmpdir, shard_size):
+        config_dict, hidden_dim, models = self._toy_model_config(shard_size)
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=True)
+
+    @pytest.mark.parametrize('shard_size', [1, 2, 4])
+    def test_not_load_optimizer_state(self, tmpdir, shard_size):
+        config_dict, hidden_dim, models = self._toy_model_config(shard_size)
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=False)
+
+    @pytest.mark.parametrize('shard_size', [1, 2, 4])
+    def test_load_module_only(self, tmpdir, shard_size):
+        config_dict, hidden_dim, models = self._toy_model_config(shard_size)
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_module_only=True)
+
+    @pytest.mark.parametrize('shard_size', [1, 2, 4])
+    def test_save_checkpoint_on_first_partition_group(self, tmpdir, shard_size):
+        config_dict, _, models = self._toy_model_config(shard_size)
+        ds_engine, _, _, _ = deepspeed.initialize(config=config_dict,
+                                                  model=models[0],
+                                                  model_parameters=models[0].parameters(),
+                                                  optimizer=None)
+
+        ds_engine.save_checkpoint(tmpdir)
+        if ds_engine.global_rank < shard_size:
+            assert ds_engine.save_non_zero_checkpoint == True
+        else:
+            assert ds_engine.save_non_zero_checkpoint == False

tests/unit/checkpoint/test_moe_checkpoint.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer
+from deepspeed.utils.torch import required_torch_version
+
+from unit.common import DistributedTest
+from unit.simple_model import *
+
+from unit.checkpoint.common import checkpoint_correctness_verification
+
+import pytest
+
+
+class TestMoECheckpoint(DistributedTest):
+    world_size = 4
+
+    @pytest.mark.parametrize("ep_size", [4])
+    def test_checkpoint_moe(self, tmpdir, ep_size):
+        if not required_torch_version(min_version=1.8):
+            pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly")
+
+        config_dict = {"train_batch_size": 8, "steps_per_print": 1, "fp16": {"enabled": True}}
+        hidden_dim = 16
+
+        models = [SimpleMoEModel(hidden_dim=hidden_dim, num_experts=ep_size, ep_size=ep_size) for _ in range(2)]
+        optimizers = [torch.optim.AdamW(params=model.parameters()) for model in models]
+        checkpoint_correctness_verification(config_dict,
+                                            models=models,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=True,
+                                            load_lr_scheduler_states=False,
+                                            empty_tag=True,
+                                            base_optimizers=optimizers,
+                                            seq_dataloader=True,
+                                            dtype=torch.float16)
+
+    @pytest.mark.parametrize("ep_size, load_optim_states", [(4, True), (4, False), (2, True), (2, False)])
+    def test_checkpoint_moe_and_zero(self, tmpdir, ep_size, load_optim_states):
+        if not required_torch_version(min_version=1.8):
+            pytest.skip("DeepSpeed MoE tests need torch 1.8 or higher to run correctly")
+
+        config_dict = {
+            "train_batch_size": 8,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 8
+            },
+            "zero_optimization": {
+                "stage": 2,
+            }
+        }
+        hidden_dim = 16
+
+        models = [SimpleMoEModel(hidden_dim=hidden_dim, num_experts=ep_size, ep_size=ep_size) for _ in range(2)]
+        # param group must have a random unique name (for now)
+        # TODO: clean-up this requirement, the unique name should not be required here
+        param_groups = [{'params': [p for p in model.parameters()], 'name': 'random-unique-name'} for model in models]
+        params = [split_params_into_different_moe_groups_for_optimizer(group) for group in param_groups]
+        optimizers = [torch.optim.AdamW(params=param) for param in params]
+        checkpoint_correctness_verification(config_dict,
+                                            models=models,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=load_optim_states,
+                                            load_lr_scheduler_states=False,
+                                            empty_tag=True,
+                                            base_optimizers=optimizers,
+                                            seq_dataloader=True,
+                                            dtype=torch.float16)

tests/unit/checkpoint/test_other_optimizer.py

@@ -0,0 +1,146 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed
+from deepspeed.ops.op_builder import FusedLambBuilder
+
+from unit.common import DistributedTest
+from unit.simple_model import *
+
+from unit.checkpoint.common import checkpoint_correctness_verification
+
+import pytest
+
+
+class TestOtherOptimizerCheckpoint(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME], reason="lamb is not compatible")
+    def test_checkpoint_unfused_optimizer(self, tmpdir):
+        #if not get_accelerator().is_fp16_supported():
+        #    pytest.skip("fp16 is not supported")
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0,
+            "scheduler": {
+                "type": "OneCycle",
+                "params": {
+                    "cycle_first_step_size": 1000,
+                    "cycle_first_stair_count": 500,
+                    "cycle_second_step_size": 1000,
+                    "cycle_second_stair_count": 500,
+                    "decay_step_size": 1000,
+                    "cycle_min_lr": 0.0001,
+                    "cycle_max_lr": 0.0010,
+                    "decay_lr_rate": 0.001,
+                    "cycle_min_mom": 0.85,
+                    "cycle_max_mom": 0.99,
+                    "decay_mom_rate": 0.0
+                }
+            }
+        }
+        dtype = torch.float
+        if get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+            dtype = torch.float16
+
+        # with bf16 fails with: DeepSpeed lamb optimizer requires dynamic loss scaling
+        # if get_accelerator().is_bf16_supported():
+        #     config_dict["bf16"] = {"enabled": True}
+
+        args = args_from_dict(tmpdir, config_dict)
+        hidden_dim = 10
+        models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        # Load & verify optimizer states
+        checkpoint_correctness_verification(config_dict,
+                                            models=models,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=True,
+                                            dtype=dtype)
+
+        # Ignore optimizer states
+        checkpoint_correctness_verification(config_dict,
+                                            models=models,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=False,
+                                            dtype=dtype)
+
+    def test_checkpoint_fused_optimizer(self, tmpdir):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU accelerator does not support this test")
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+        }
+        dtype = torch.float
+        if get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+            dtype = torch.float16
+
+        args = args_from_dict(tmpdir, config_dict)
+        hidden_dim = 10
+        models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        # Load & verify optimizer states
+        checkpoint_correctness_verification(config_dict,
+                                            models=models,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=True,
+                                            dtype=dtype)
+
+        # Ignore optimizer states
+        checkpoint_correctness_verification(config_dict,
+                                            models=models,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=False,
+                                            dtype=dtype)
+
+    def test_checkpoint_fp32_optimizer(self, tmpdir):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+            "fp16": {
+                "enabled": False
+            }
+        }
+
+        args = args_from_dict(tmpdir, config_dict)
+        hidden_dim = 10
+        models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+        checkpoint_correctness_verification(config_dict,
+                                            models=models,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            dtype=torch.float32)

tests/unit/checkpoint/test_pipeline.py

@@ -0,0 +1,114 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
+from unit.common import DistributedTest
+from unit.simple_model import *
+from unit.checkpoint.common import checkpoint_correctness_verification
+from unit.util import skip_on_arch
+
+import pytest
+
+
+class TestPipelineCheckpoint(DistributedTest):
+    world_size = 4
+
+    @pytest.mark.parametrize("zero_stage", [0, 1])
+    def test_checkpoint_pipe_engine(self, zero_stage, tmpdir):
+        skip_on_arch(min_arch=7)
+
+        config_dict = {
+            "train_batch_size": 2,
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-5
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage
+            },
+            "fp16": {
+                "enabled": zero_stage > 0
+            },
+            "scheduler": {
+                "type": "OneCycle",
+                "params": {
+                    "cycle_first_step_size": 1000,
+                    "cycle_first_stair_count": 500,
+                    "cycle_second_step_size": 1000,
+                    "cycle_second_stair_count": 500,
+                    "decay_step_size": 1000,
+                    "cycle_min_lr": 0.0001,
+                    "cycle_max_lr": 0.0010,
+                    "decay_lr_rate": 0.001,
+                    "cycle_min_mom": 0.85,
+                    "cycle_max_mom": 0.99,
+                    "decay_mom_rate": 0.0
+                }
+            }
+        }
+
+        models = [LinearStackPipe(num_stages=2) for _ in range(2)]
+        checkpoint_correctness_verification(config_dict=config_dict,
+                                            models=models,
+                                            hidden_dim=models[0].hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=True,
+                                            load_lr_scheduler_states=True,
+                                            train_batch=True,
+                                            dtype=torch.float16 if zero_stage > 0 else torch.float32)
+
+    @pytest.mark.parametrize(
+        "base_topo,test_topo",
+        [
+            #(PipeTopo(num_pp=1,
+            #          num_dp=4),
+            # PipeTopo(num_pp=4,
+            #          num_dp=1)),
+            #(PipeTopo(num_pp=2,
+            #          num_dp=2),
+            # PipeTopo(num_pp=2,
+            #          num_dp=2)),
+            #(PipeTopo(num_pp=4,
+            #          num_dp=1),
+            # PipeTopo(num_pp=2,
+            #          num_dp=2)),
+        ])
+    def test_checkpoint_pipe_module(self, base_topo, test_topo, tmpdir):
+        checkpoint_engine = TorchCheckpointEngine()
+        base_model = LinearStackPipe(topology=base_topo)
+        base_model.save_state_dict(tmpdir, checkpoint_engine=checkpoint_engine)
+
+        dist.barrier()
+
+        test_model = LinearStackPipe(topology=test_topo)
+        test_model.load_state_dir(tmpdir, checkpoint_engine=checkpoint_engine)
+
+        # Base and test can have different lengths, so make sure we map from the
+        # smaller to larger model
+        if len(base_model.forward_funcs) < len(test_model.forward_funcs):
+            A = base_model
+            B = test_model
+        else:
+            A = test_model
+            B = base_model
+
+        # Compare layers individually since partitions are different
+        for idx, A_layer in enumerate(A.forward_funcs):
+            if not hasattr(A_layer, 'parameters'):
+                # Skip functionals, etc.
+                continue
+
+            # Find the corresponding layer in B
+            global_idx = idx + A._local_start
+            B_local_idx = global_idx - B._local_start
+            B_layer = B.forward_funcs[B_local_idx]
+
+            # Compare layer parameters
+            for p0, p1 in zip(A_layer.parameters(), B_layer.parameters()):
+                assert torch.allclose(p0, p1, atol=1e-07), f"Model state {p0} is not equal to {p1}"

tests/unit/checkpoint/test_reshape_checkpoint.py

@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.checkpoint import model_3d_desc
+
+
+def _do_reshape(src_3d, tgt_3d):
+    assert src_3d.can_reshape(tgt_3d)
+    new_3d_map = src_3d.reshape(tgt_3d)
+
+    assert len(new_3d_map) == tgt_3d.dp_degree
+    for new_2d_map in new_3d_map:
+        assert new_2d_map.pp_degree == tgt_3d.pp_degree
+        assert new_2d_map.tp_degree == tgt_3d.tp_degree
+
+    return new_3d_map
+
+
+# Specify 3d shape as pp/tp/dp
+def test_reshape_222_to_111():
+    src_3d = model_3d_desc(pp_degree=2, tp_degree=2, dp_degree=2)
+    tgt_3d = model_3d_desc(pp_degree=1, tp_degree=1, dp_degree=1)
+
+    new_3d_map = _do_reshape(src_3d, tgt_3d)
+
+    assert new_3d_map[0].get_data(pp_index=0, tp_index=0) == [0, 4, 1, 5, 2, 6, 3, 7]
+
+
+def test_reshape_222_to_121():
+    src_3d = model_3d_desc(pp_degree=2, tp_degree=2, dp_degree=2)
+    tgt_3d = model_3d_desc(pp_degree=1, tp_degree=2, dp_degree=1)
+
+    new_3d_map = _do_reshape(src_3d, tgt_3d)
+
+    assert new_3d_map[0].get_data(pp_index=0, tp_index=0) == [0, 4, 2, 6]
+    assert new_3d_map[0].get_data(pp_index=0, tp_index=1) == [1, 5, 3, 7]
+
+
+def test_reshape_222_to_122():
+    src_3d = model_3d_desc(pp_degree=2, tp_degree=2, dp_degree=2)
+    tgt_3d = model_3d_desc(pp_degree=1, tp_degree=2, dp_degree=2)
+
+    new_3d_map = _do_reshape(src_3d, tgt_3d)
+
+    assert new_3d_map[0].get_data(pp_index=0, tp_index=0) == [0, 4]
+    assert new_3d_map[0].get_data(pp_index=0, tp_index=1) == [1, 5]
+    assert new_3d_map[1].get_data(pp_index=0, tp_index=0) == [2, 6]
+    assert new_3d_map[1].get_data(pp_index=0, tp_index=1) == [3, 7]
+
+
+def test_reshape_222_to_211():
+    src_3d = model_3d_desc(pp_degree=2, tp_degree=2, dp_degree=2)
+    tgt_3d = model_3d_desc(pp_degree=2, tp_degree=1, dp_degree=1)
+
+    new_3d_map = _do_reshape(src_3d, tgt_3d)
+
+    assert new_3d_map[0].get_data(pp_index=0, tp_index=0) == [0, 4, 1, 5]
+    assert new_3d_map[0].get_data(pp_index=1, tp_index=0) == [2, 6, 3, 7]

tests/unit/checkpoint/test_shared_weights.py

@@ -0,0 +1,49 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import torch.nn as nn
+
+import deepspeed
+from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+from unit.common import DistributedTest
+
+
+class ModelWithSharedWeights(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.layer0 = nn.Linear(100, 100)
+        self.layer1 = nn.Linear(200, 200)
+        self.layer2 = nn.Linear(300, 300)
+        # tie layer 1 and layer 2
+        self.layer1.weight = self.layer2.weight
+
+
+class TestCheckpointSharedWeights(DistributedTest):
+    world_size = 2
+
+    def test_checkpoint_shared_weights(self, tmp_path):
+        config = {
+            "train_micro_batch_size_per_gpu": 2,
+            "zero_allow_untested_optimizer": True,
+            "zero_optimization": {
+                "stage": 2
+            },
+        }
+        model = ModelWithSharedWeights()
+        optimizer = torch.optim.Adam(model.parameters())
+
+        deepspeed_engine, _, _, _ = deepspeed.initialize(
+            config=config,
+            model=model,
+            optimizer=optimizer,
+        )
+        filename = tmp_path / "checkpoint.pt"
+        deepspeed_engine.save_checkpoint(filename, tag="checkpoint")
+
+        model = ModelWithSharedWeights()
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(filename, tag="checkpoint")
+        model.load_state_dict(state_dict, strict=True)

tests/unit/checkpoint/test_sparse.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed
+
+from unit.common import DistributedTest
+from unit.simple_model import *
+
+import pytest
+
+
+class TestSparseCheckpoint(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize(["to_save_model_has_embedding", "to_save_model_sparse"], [
+        [False, False],
+        [True, False],
+        [True, True],
+    ])
+    @pytest.mark.parametrize(["destination_has_embedding", "destination_sparse"], [
+        [False, False],
+        [True, False],
+        [True, True],
+    ])
+    def test_non_strict_load_sparse(self, tmpdir, to_save_model_has_embedding, to_save_model_sparse,
+                                    destination_has_embedding, destination_sparse):
+
+        class ModelNoEmbedding(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 1)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        class ModelEmbedding(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.emb = torch.nn.Embedding(10, 3)
+                self.linear = torch.nn.Linear(3, 1)
+
+            def forward(self, x, offsets):
+                return self.linear(self.emb(x, offsets))
+
+        if to_save_model_has_embedding:
+            model_to_save = ModelEmbedding()
+        else:
+            model_to_save = ModelNoEmbedding()
+        if destination_has_embedding:
+            model_destination = ModelEmbedding()
+        else:
+            model_destination = ModelNoEmbedding()
+
+        engine_to_save, _, _, _ = deepspeed.initialize(model=model_to_save,
+                                                       config={
+                                                           "train_batch_size": 2,
+                                                           "sparse_gradients": to_save_model_sparse
+                                                       })
+        engine_destination, _, _, _ = deepspeed.initialize(model=model_destination,
+                                                           config={
+                                                               "train_batch_size": 2,
+                                                               "sparse_gradients": destination_sparse
+                                                           })
+
+        save_folder = os.path.join(tmpdir, 'saved_checkpoint')
+        save_tag = '1'
+
+        engine_to_save.save_checkpoint(save_folder, tag=save_tag)
+
+        is_sparse_destination = isinstance(model_destination, ModelEmbedding) and destination_sparse
+        if isinstance(model_destination, ModelEmbedding) and model_destination.emb.sparse:
+            assert "emb.weight" in engine_destination.sparse_tensor_module_names
+        engine_destination.load_checkpoint(save_folder,
+                                           tag=save_tag,
+                                           load_module_strict=False,
+                                           load_optimizer_states=False,
+                                           load_lr_scheduler_states=False,
+                                           load_module_only=False)
+        if isinstance(model_destination, ModelEmbedding) and isinstance(model_to_save, ModelEmbedding):
+            assert engine_destination.sparse_tensor_module_names == engine_to_save.sparse_tensor_module_names
+        elif isinstance(model_destination, ModelEmbedding):
+            assert not is_sparse_destination or "emb.weight" in engine_destination.sparse_tensor_module_names
+        else:
+            assert len(engine_destination.sparse_tensor_module_names) == 0

tests/unit/checkpoint/test_tag_validation.py

@@ -0,0 +1,85 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed
+
+from unit.common import DistributedTest
+from unit.simple_model import *
+
+import pytest
+
+
+class TestCheckpointValidationTag(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize('valid_mode', ["FAIL", "WARN", "IGNORE"])
+    def test_checkpoint_unique_tag(self, tmpdir, valid_mode):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "checkpoint": {
+                "tag_validation": valid_mode
+            }
+        }
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        if valid_mode == "FAIL":
+            with pytest.raises(AssertionError):
+                model.save_checkpoint(save_dir=tmpdir, tag=f"tag-{dist.get_rank()}")
+        else:
+            model.save_checkpoint(save_dir=tmpdir, tag=f"tag-{dist.get_rank()}")
+
+    def test_checkpoint_unknown_tag_validation(self, tmpdir):
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "checkpoint": {
+                "tag_validation": "foo"
+            }
+        }
+        hidden_dim = 10
+        args = args_from_dict(tmpdir, config_dict)
+        model = SimpleModel(hidden_dim)
+
+        with pytest.raises(deepspeed.DeepSpeedConfigError):
+            model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+
+class TestSaveCheckpointInvalidDir(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize('save_dir', [None, ""])
+    def test_save_checkpoint_empty_dir(self, save_dir):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            }
+        }
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        with pytest.raises(ValueError):
+            model.save_checkpoint(save_dir=save_dir)

tests/unit/checkpoint/test_universal_checkpoint.py

@@ -0,0 +1,272 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import math
+
+import deepspeed
+from types import SimpleNamespace
+from torch.utils._pytree import tree_map
+
+from deepspeed.utils.torch import required_torch_version
+from deepspeed.checkpoint import UNIVERSAL_CHECKPOINT_INFO
+from deepspeed.checkpoint.ds_to_universal import main as convert_to_universal
+
+from unit.common import DistributedTest, DistributedFixture
+from unit.simple_model import *
+from unit.util import bf16_required_version_check
+
+from unit.checkpoint.common import compare_opt_state_dicts, compare_state_dicts
+
+import pytest
+import deepspeed.comm as dist
+
+
+def get_expected_mismatch_keys():
+    # torch 1.2.* stores raw tensor id numbers in checkpoint state which leads to
+    # false positive mismatches in checkpoint state comparisons.
+    # Newer torch versions store tensor ids as 0, 1, 2, ...
+    return [] if required_torch_version(min_version=1.4) else ['params']
+
+
+def maybe_step(t):
+    return not torch.is_tensor(t) or (t.device.type == 'cpu' and t.numel() == 1)
+
+
+def gather_opt_state(optimizer_state):
+
+    def gather_tensor(t):
+
+        if maybe_step(t):
+            return t
+        else:
+            buffer = [torch.zeros_like(t.flatten()) for _ in range(dist.get_world_size())]
+            dist.all_gather(buffer, t.flatten())
+            return torch.cat(buffer)
+
+    return tree_map(gather_tensor, optimizer_state)
+
+
+def remove_pad_in_opt_state(optimizer_state, num_params):
+
+    def remove_pad(t):
+        if maybe_step(t):
+            return t
+        else:
+            return t[:num_params]
+
+    return tree_map(remove_pad, optimizer_state)
+
+
+CP_TAG = "test_tag"
+
+
+def init_ds_engine(model, ds_config, use_torch_adam):
+
+    if use_torch_adam:
+        ds_optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+        del ds_config["optimizer"]
+        model, _, _, _ = deepspeed.initialize(config=ds_config, model=model, optimizer=ds_optimizer)
+    else:
+        model, _, _, _ = deepspeed.initialize(config=ds_config, model=model, model_parameters=model.parameters())
+
+    return model
+
+
+def train_save_convert(ds_config, hidden_dim, load_optim, use_torch_adam, dtype, tmpdir, world_size):
+    if dtype == torch.bfloat16 and not bf16_required_version_check():
+        return
+
+    test_step = 8
+
+    model = SimpleModel(hidden_dim, nlayers=2)
+    model = init_ds_engine(model, ds_config, use_torch_adam)
+    data_loader = random_dataloader(model=model,
+                                    total_samples=test_step,
+                                    hidden_dim=hidden_dim,
+                                    device=model.device,
+                                    dtype=dtype)
+    for batch in data_loader:
+        loss = model(batch[0], batch[1])
+        model.backward(loss)
+        model.step()
+
+    if ds_config["zero_optimization"]["stage"] == 3:
+        model.optimizer._set_fp32_optimizer_param_groups()
+        sd = model.optimizer.optimizer.state_dict() if load_optim else None
+        model.optimizer._clear_fp32_optimizer_param_groups()
+    else:
+        sd = model.optimizer.optimizer.state_dict() if load_optim else None
+
+    client_state = {}
+    client_state[UNIVERSAL_CHECKPOINT_INFO] = {}
+    client_state['iteration'] = test_step
+    model.save_checkpoint(tmpdir, tag=CP_TAG, client_state=client_state)
+
+    cp_dir = os.path.join(tmpdir, CP_TAG)
+    univ_cp_dir = f"{cp_dir}_universal"
+
+    args = SimpleNamespace(input_folder=cp_dir,
+                           output_folder=univ_cp_dir,
+                           num_extract_workers=1,
+                           num_merge_workers=1,
+                           keep_temp_folder=False,
+                           strict=True,
+                           inject_missing_state=False)
+
+    dist.barrier()
+    if dist.get_rank() == 0:
+        convert_to_universal(args)
+
+    model_state = model.state_dict()
+    optimizer_state = None
+    if load_optim:
+        if ds_config["zero_optimization"]["stage"] == 3:
+            model.optimizer._set_fp32_optimizer_param_groups()
+            optimizer_state = gather_opt_state(model.optimizer.optimizer.state_dict())
+            model.optimizer._clear_fp32_optimizer_param_groups()
+            update_gathered_stage3_optimizer(optimizer_state, model._get_zero_param_shapes(), world_size)
+        else:
+            optimizer_state = gather_opt_state(model.optimizer.optimizer.state_dict())
+
+    if dist.get_rank() == 0:
+        torch.save((model_state, optimizer_state), os.path.join(tmpdir, "baseline_state.pt"))
+
+    dist.barrier()
+    model.destroy()
+
+
+@pytest.fixture
+def ds_config(zero_stage, dtype, sub_group_size):
+    ds_config = {
+        "train_batch_size": 8,
+        "optimizer": {
+            "type": 'Adam'
+        },
+        "zero_optimization": {
+            "stage": zero_stage,
+        }
+    }
+    if dtype == torch.float16:
+        ds_config["fp16"] = {"enabled": True, "initial_scale_power": 8}
+    elif dtype == torch.bfloat16:
+        ds_config["bf16"] = {"enabled": True}
+    if sub_group_size > 0:
+        ds_config["zero_optimization"]["sub_group_size"] = sub_group_size
+    return ds_config
+
+
+class _baseline(DistributedFixture):
+    world_size = None
+
+    def run(self, tmpdir, ds_config, zero_stage, dtype, load_optim, use_torch_adam):
+        hidden_dim = 10
+        train_save_convert(ds_config, hidden_dim, load_optim, use_torch_adam, dtype, tmpdir, self.world_size)
+
+
+class baseline_ws2(_baseline):
+    world_size = 2
+
+
+class baseline_ws4(_baseline):
+    world_size = 4
+
+
+# Stage3 use shard parameter, need to reorganize the optimizer parameters.
+def update_gathered_stage3_optimizer(optimizer_state, param_shapes, world_size):
+    for sub_group_id, group in enumerate(optimizer_state["param_groups"]):
+        group["params"] = None
+
+    new_state = {}
+    for sub_group_id, sub_group_param_shape in enumerate(param_shapes):
+        total_numel = optimizer_state['state'][sub_group_id]['exp_avg'].numel()
+        assert total_numel % world_size == 0
+        numel_per_rank = total_numel // world_size
+        param_offset_in_current_rank = 0
+        for param_name, param_shape in sub_group_param_shape.items():
+            param_numel = param_shape.numel()
+            param_partition_numel = math.ceil(param_numel / world_size)
+            param_optimizer_tensor = {
+                "exp_avg": torch.zeros(param_numel),
+                "exp_avg_sq": torch.zeros(param_numel),
+                "step": optimizer_state['state'][sub_group_id]['step'],
+            }
+            for key in ["exp_avg", "exp_avg_sq"]:
+                write_offset = 0
+                for rank in range(world_size):
+                    offset = param_offset_in_current_rank + rank * numel_per_rank
+                    length = min(param_partition_numel, param_numel - rank * param_partition_numel)
+                    tmp = optimizer_state['state'][sub_group_id][key].narrow(0, offset, length)
+                    param_optimizer_tensor[key].narrow(0, write_offset, length).copy_(tmp)
+                    write_offset += length
+            param_offset_in_current_rank += param_partition_numel
+            new_state[param_name] = param_optimizer_tensor
+    optimizer_state["state"] = new_state
+
+
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("zero_stage", [1, 3])
+@pytest.mark.parametrize("use_torch_adam", [False, True])
+@pytest.mark.parametrize("load_optim", [False, True])
+@pytest.mark.parametrize("sub_group_size", [-1, 100])
+class TestZeROUniversalCheckpointDP(DistributedTest):
+
+    def _run_test(self, tmpdir, dtype, ds_config, load_optim, use_torch_adam, world_size):
+        if dtype == torch.bfloat16 and not bf16_required_version_check():
+            pytest.skip(
+                " DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
+            )
+
+        hidden_dim = 10
+        loaded_model_state, loaded_optimizer_state = torch.load(f"{tmpdir}/baseline_state.pt", weights_only=False)
+
+        ds_config["checkpoint"] = {"load_universal": True}
+        univ_model = SimpleModel(hidden_dim, nlayers=2)
+        univ_model = init_ds_engine(univ_model, ds_config, use_torch_adam)
+        univ_model.load_checkpoint(tmpdir, tag=f"{CP_TAG}_universal", load_optimizer_states=load_optim)
+
+        model_state = univ_model.state_dict()
+        compare_state_dicts(model_state, loaded_model_state)
+
+        if load_optim:
+            if ds_config["zero_optimization"]["stage"] == 3:
+                univ_model.optimizer._set_fp32_optimizer_param_groups()
+                optimizer_state = gather_opt_state(univ_model.optimizer.optimizer.state_dict())
+                univ_model.optimizer._clear_fp32_optimizer_param_groups()
+                update_gathered_stage3_optimizer(optimizer_state, univ_model._get_zero_param_shapes(), world_size)
+            else:
+                optimizer_state = gather_opt_state(univ_model.optimizer.optimizer.state_dict())
+            # padding sizes may differ when dp sizes are different
+            param_count = sum(p.numel() for p in univ_model.parameters())
+            optimizer_state = remove_pad_in_opt_state(optimizer_state, param_count)
+            loaded_optimizer_state = remove_pad_in_opt_state(loaded_optimizer_state, param_count)
+
+            compare_opt_state_dicts(optimizer_state, loaded_optimizer_state, get_expected_mismatch_keys())
+
+        # Run training again to verify that the optimizer has necessary states
+        test_step = 8
+        data_loader = random_dataloader(model=univ_model,
+                                        total_samples=test_step,
+                                        hidden_dim=hidden_dim,
+                                        device=univ_model.device,
+                                        dtype=dtype)
+        for batch in data_loader:
+            loss = univ_model(batch[0], batch[1])
+            univ_model.backward(loss)
+            univ_model.step()
+
+        univ_model.destroy()
+
+    @pytest.mark.world_size(2)
+    def test_dp_world_size_2to2(self, baseline_ws2, tmpdir, dtype, ds_config, load_optim, use_torch_adam):
+        self._run_test(tmpdir, dtype, ds_config, load_optim, use_torch_adam, 2)
+
+    @pytest.mark.world_size(2)
+    def test_dp_world_size_4to2(self, baseline_ws4, tmpdir, dtype, ds_config, load_optim, use_torch_adam):
+        self._run_test(tmpdir, dtype, ds_config, load_optim, use_torch_adam, 2)
+
+    @pytest.mark.world_size(4)
+    def test_dp_world_size_2to4(self, baseline_ws2, tmpdir, dtype, ds_config, load_optim, use_torch_adam):
+        self._run_test(tmpdir, dtype, ds_config, load_optim, use_torch_adam, 4)

tests/unit/checkpoint/test_zero_optimizer.py

@@ -0,0 +1,776 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed
+from types import SimpleNamespace
+from deepspeed.ops.op_builder import CPUAdamBuilder
+from deepspeed.checkpoint.utils import clone_tensors_for_torch_save, get_model_ckpt_name_for_rank
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.zero import ZeroParamStatus
+from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
+from deepspeed.utils.torch import required_torch_version
+
+from unit.common import DistributedTest, DistributedFixture
+from unit.simple_model import *
+
+from unit.checkpoint.common import *
+
+import pytest
+
+
+class TestZeROCheckpoint(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize('zero_stage', [3])
+    def test_pipeline_checkpoint_loading(self, tmpdir, zero_stage):
+        config_dict = {
+            "train_batch_size": 2,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "pipeline_loading_checkpoint": True,
+            }
+        }
+        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 = 10
+
+        with deepspeed.zero.Init(config_dict_or_path=config_dict):
+            models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_module_only=True)
+
+    @pytest.mark.parametrize('zero_stage, use_cpu_offload, adam_optimizer', [(0, False, 'Adam'), (1, False, 'Adam'),
+                                                                             (2, False, 'Adam'),
+                                                                             (2, True, 'deepspeed_adam'),
+                                                                             (3, False, 'Adam'),
+                                                                             (3, True, 'deepspeed_adam')])
+    def test_load_optimizer_state(self, tmpdir, zero_stage, use_cpu_offload, adam_optimizer):
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+            "wall_clock_breakdown": True,
+            "zero_optimization": {
+                "stage": zero_stage,
+                "cpu_offload": use_cpu_offload
+            }
+        }
+        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 = 10
+
+        if zero_stage == 3:
+            with deepspeed.zero.Init(config_dict_or_path=config_dict):
+                models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+        else:
+            models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=True)
+
+    @pytest.mark.parametrize('zero_stage, use_cpu_offload, adam_optimizer', [(1, False, "Adam"), (2, False, "Adam"),
+                                                                             (2, True, 'deepspeed_adam'),
+                                                                             (3, False, 'Adam'),
+                                                                             (3, True, 'deepspeed_adam')])
+    def test_not_load_optimizer_state(self, tmpdir, zero_stage, use_cpu_offload, adam_optimizer):
+        if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
+            pytest.skip("cpu-adam is not compatible")
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "cpu_offload": use_cpu_offload
+            }
+        }
+        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 = 10
+
+        if zero_stage == 3:
+            global DeepSpeedZeroOptimizer_Stage3
+            from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
+            with deepspeed.zero.Init(config_dict_or_path=config_dict):
+                models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+        else:
+            models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=False)
+
+    @pytest.mark.parametrize('zero_stage', [1, 2])
+    def test_hybrid_optimizer_state(self, tmpdir, zero_stage):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 2,
+            "gradient_accumulation_steps": 2,
+            "steps_per_print": 1,
+            "zero_optimization": {
+                "stage": zero_stage
+            },
+            "zero_allow_untested_optimizer": True,
+        }
+        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 = 10
+        models = [SimpleModel(hidden_dim=hidden_dim) for _ in range(2)]
+        optimizers = [HybridStateOptimizer(model.parameters()) for model in models]
+
+        checkpoint_correctness_verification(config_dict,
+                                            models=models,
+                                            base_optimizers=optimizers,
+                                            hidden_dim=hidden_dim,
+                                            tmpdir=tmpdir,
+                                            load_optimizer_states=True)
+
+    @pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
+    def test_load_module_only(self, tmpdir, zero_stage):
+        if zero_stage == 0 and get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU Accelerator does not support this test")
+        config_dict = {
+            "train_batch_size": 2,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+        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 = 10
+
+        if zero_stage == 3:
+            with deepspeed.zero.Init(config_dict_or_path=config_dict):
+                models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+        else:
+            models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_module_only=True)
+
+
+class ws4_model_checkpoint(DistributedFixture):
+    world_size = 4
+
+    def run(self, class_tmpdir, elastic_save, load_optim):
+        config_dict = {
+            "train_batch_size": 4,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": 2,
+                "elastic_checkpoint": elastic_save
+            }
+        }
+        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 = 10
+        model = SimpleModel(hidden_dim)
+
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        data_loader = random_dataloader(model=model, total_samples=8, hidden_dim=hidden_dim, device=model.device)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+        if load_optim:
+            torch.save(model.optimizer.optimizer.state_dict(), os.path.join(class_tmpdir, 'opt-state-dict'))
+        model.save_checkpoint(class_tmpdir)
+
+
+class ws4_model_checkpoint_zeropp(DistributedFixture):
+
+    world_size = 4
+
+    def run(self, class_tmpdir):
+        config_dict = {
+            "train_batch_size": 4,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "zero_hpz_partition_size": 2,
+            }
+        }
+
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        for param in model.parameters():
+            param.data = torch.ones_like(param.data, device=param.data.device, requires_grad=False)
+
+        # save model and zero checkpoint
+        torch.save(model.state_dict(), os.path.join(class_tmpdir, "model.pt"))
+        ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
+        ds_model.save_checkpoint(class_tmpdir)
+
+
+@pytest.mark.parametrize("elastic_save", [True, False])
+@pytest.mark.parametrize("elastic_load", [True, False])
+@pytest.mark.parametrize("load_optim", [True, False])
+class TestZeROElasticCheckpoint(DistributedTest):
+    world_size = 2
+
+    def test_elastic_checkpoint_fixed_dp(self, tmpdir, elastic_save, elastic_load, load_optim):
+        config_dict = {
+            "train_batch_size": 2,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": 2,
+                "elastic_checkpoint": elastic_save
+            }
+        }
+        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 = 10
+
+        # torch 1.2.* stores raw tensor id numbers in checkpoint state which leads to
+        # false positive mismatches in checkpoint state comparisons.
+        # Newer torch versions store tensor ids as 0, 1, 2, ...
+        expected_mismatch_keys = [] if required_torch_version(min_version=1.4) else ['params']
+        models = [SimpleModel(hidden_dim) for _ in range(2)]
+        model, _, _, _ = deepspeed.initialize(config=config_dict,
+                                              model=models[0],
+                                              model_parameters=models[0].parameters())
+        run_steps = 8
+        data_loader = random_dataloader(model=model,
+                                        total_samples=run_steps,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+        if load_optim:
+            opt_state_dict_file = f'opt-state-dict_rank{dist.get_rank()}'
+            torch.save(model.optimizer.optimizer.state_dict(), os.path.join(tmpdir, opt_state_dict_file))
+        model.save_checkpoint(tmpdir)
+
+        config_dict["zero_optimization"]["elastic_checkpoint"] = elastic_load
+        model, _, _, _ = deepspeed.initialize(config=config_dict,
+                                              model=models[1],
+                                              model_parameters=models[1].parameters())
+        model.load_checkpoint(tmpdir, load_optimizer_states=load_optim)
+
+        if load_optim:
+            saved_sd = torch.load(os.path.join(tmpdir, opt_state_dict_file), weights_only=False)
+            curr_sd = model.optimizer.optimizer.state_dict()
+            compare_opt_state_dicts(curr_sd, saved_sd, expected_mismatch_keys)
+
+        data_loader = random_dataloader(model=model, total_samples=8, hidden_dim=hidden_dim, device=model.device)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+    def test_elastic_checkpoint_change_dp(self, ws4_model_checkpoint, class_tmpdir, elastic_save, elastic_load,
+                                          load_optim):
+        config_dict = {
+            "train_batch_size": 4,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": 2,
+                "elastic_checkpoint": elastic_load
+            }
+        }
+        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 = 10
+        model = SimpleModel(hidden_dim)
+
+        # Load checkpoint with dp world size = 2
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        if load_optim:
+            with pytest.raises(deepspeed.runtime.zero.utils.ZeRORuntimeException):
+                model.load_checkpoint(class_tmpdir, load_optimizer_states=load_optim)
+        else:
+            model.load_checkpoint(class_tmpdir, load_optimizer_states=load_optim)
+
+
+class TestZeROSaveLoadEdgeCase(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
+    def test_immediate_save_load(self, tmpdir, zero_stage):
+        config_dict = {
+            "train_batch_size": 4,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+        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 = 10
+        model = SimpleModel(hidden_dim)
+
+        ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
+        ds_model.save_checkpoint(tmpdir)
+        ds_model.load_checkpoint(tmpdir,
+                                 load_optimizer_states=False,
+                                 load_lr_scheduler_states=False,
+                                 load_module_only=False)
+
+    @pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
+    def test_load_immediate_save(self, tmpdir, zero_stage):
+        if zero_stage == 0 and get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU Accelerator does not support this test")
+        config_dict = {
+            "train_batch_size": 4,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+        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 = 10
+        model = SimpleModel(hidden_dim)
+
+        # 1. pretrain a model and save it
+        ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
+        data_loader = random_dataloader(model=ds_model, total_samples=1, hidden_dim=hidden_dim, device=ds_model.device)
+        for _, batch in enumerate(data_loader):
+            loss = ds_model(batch[0], batch[1])
+            ds_model.backward(loss)
+            ds_model.step()
+
+        ds_model.empty_partition_cache()
+        ds_model.save_checkpoint(tmpdir)
+
+        # 2. load and immediately save a model with a fresh ds engine
+        ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
+        ds_model.load_checkpoint(tmpdir,
+                                 load_optimizer_states=False,
+                                 load_lr_scheduler_states=False,
+                                 load_module_only=False)
+        ds_model.save_checkpoint(tmpdir)
+
+    @pytest.mark.parametrize('zero_stage', [0, 1, 2, 3])
+    def test_save_before_accum_grad_is_done(self, tmpdir, zero_stage):
+        config_dict = {
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "stage3_gather_fp16_weights_on_model_save": True,
+            },
+            "gradient_accumulation_steps": 2,
+            "train_micro_batch_size_per_gpu": 1,
+            "train_batch_size": 4,
+        }
+        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 = 10
+        model = SimpleModel(hidden_dim)
+
+        # This test reproduces a bug where one tries to retrieve a 16bit model before grad_accum
+        # cycle was completed.
+        # So we config grad_accum=2 and step only once and save_16bit_model
+        ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
+
+        data_loader = random_dataloader(model=ds_model, total_samples=2, hidden_dim=hidden_dim, device=ds_model.device)
+
+        batch = next(iter(data_loader))
+        loss = ds_model(batch[0], batch[1])
+        ds_model.backward(loss)
+        ds_model.step()
+
+        ds_model.empty_partition_cache()
+
+        # we stepped only once, and now save 16bit model before gradient_accumulation_steps=2 is complete
+        ds_model.save_16bit_model(tmpdir, "model.pt")
+
+        # let's test just as well that we can save the checkpoint too
+        ds_model.save_checkpoint(tmpdir)
+
+
+class TestZeROCheckpointFrozenWeights(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    def test_load_optimizer_state(self, tmpdir, zero_stage):
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+            "wall_clock_breakdown": True,
+            "zero_optimization": {
+                "stage": zero_stage
+            }
+        }
+        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 = 10
+
+        with deepspeed.zero.Init(enabled=zero_stage == 3, config_dict_or_path=config_dict):
+            models = [SimpleFrozenModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=True)
+
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    def test_not_load_optimizer_state(self, tmpdir, zero_stage):
+
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": 'Adam',
+                "params": {
+                    "lr": 0.00015,
+                    "betas": [0.8, 0.999],
+                    "eps": 1e-8,
+                    "weight_decay": 3e-7
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage
+            }
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+        hidden_dim = 10
+
+        with deepspeed.zero.Init(enabled=zero_stage == 3, config_dict_or_path=config_dict):
+            models = [SimpleFrozenModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_optimizer_states=False)
+
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    def test_load_module_only(self, tmpdir, zero_stage):
+        config_dict = {
+            "train_batch_size": 2,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+        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 = 10
+
+        with deepspeed.zero.Init(enabled=zero_stage == 3, config_dict_or_path=config_dict):
+            models = [SimpleFrozenModel(hidden_dim, empty_grad=False) for _ in range(2)]
+
+        checkpoint_correctness_verification(config_dict, models, hidden_dim, tmpdir, load_module_only=True)
+
+    @pytest.mark.parametrize('zero_stage', [1, 2])
+    def test_save_exclude_frozen_weights(self, tmpdir, zero_stage):
+        world_size = 1
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+        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 = 10
+
+        model = SimpleFrozenModel(hidden_dim, empty_grad=False)
+
+        ds_engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        # Validate backwards-compatibility of including frozen parameters in checkpoint
+        all_ckpt_folder = os.path.join(tmpdir, 'all_params')
+        ds_engine.save_checkpoint(all_ckpt_folder)
+        all_params_ckpt_file = get_model_ckpt_name_for_rank(os.path.join(all_ckpt_folder, 'global_step0'), '00')
+        loaded_all_param_model = torch.load(all_params_ckpt_file, weights_only=False)['module']
+        all_param_names = set([n for n, p in model.named_parameters()])
+        assert set(loaded_all_param_model.keys()) == all_param_names
+
+        # Validate exclusion of frozen parameters
+        trainable_ckpt_folder = os.path.join(tmpdir, 'no_frozen_params')
+        ds_engine.save_checkpoint(trainable_ckpt_folder, exclude_frozen_parameters=True)
+
+        trainable_ckpt_file = get_model_ckpt_name_for_rank(os.path.join(trainable_ckpt_folder, 'global_step0'), '00')
+
+        # Excluding frozen parameters should reduce checkpoint size
+        assert os.path.getsize(all_params_ckpt_file) > os.path.getsize(trainable_ckpt_file)
+
+        loaded_trainable_param_model = torch.load(trainable_ckpt_file, weights_only=False)['module']
+        frozen_param_names = set([n for n, p in model.named_parameters() if not p.requires_grad])
+        loaded_trainable_param_names = set(loaded_trainable_param_model.keys())
+        overlap_names = set.intersection(loaded_trainable_param_names, frozen_param_names)
+        assert len(overlap_names) == 0
+
+        trainable_param_names = set([n for n, p in model.named_parameters() if p.requires_grad])
+        assert loaded_trainable_param_names == trainable_param_names
+
+    @pytest.mark.parametrize('zero_stage', [1, 2])
+    def test_save_exclude_custom_frozen_weights(self, tmpdir, zero_stage):
+        world_size = 1
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+        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 = 10
+
+        model = SimpleFrozenModel(hidden_dim, empty_grad=False)
+
+        ds_engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        # Validate custom state_dict model
+        state_dict_bk = model.state_dict
+        model.state_dict = model.custom_state_dict
+        custom_state_dict_ckpt_folder = os.path.join(tmpdir, 'custom_state_dict')
+        ds_engine.save_checkpoint(custom_state_dict_ckpt_folder, exclude_frozen_parameters=True)
+
+        custom_state_dict_ckpt_file = get_model_ckpt_name_for_rank(
+            os.path.join(custom_state_dict_ckpt_folder, 'global_step0'), '00')
+        loaded_custom_state_dict_param_model = torch.load(custom_state_dict_ckpt_file, weights_only=False)['module']
+        loaded_custom_state_dict_param_names = set(loaded_custom_state_dict_param_model.keys())
+
+        custom_state_dict_param_names = set([k for k, v in model.state_dict().items()])
+        trainable_param_names = set([n for n, p in model.named_parameters() if p.requires_grad])
+        overlap_names = set.intersection(custom_state_dict_param_names, trainable_param_names)
+
+        assert loaded_custom_state_dict_param_names == overlap_names
+
+        model.state_dict = state_dict_bk
+
+
+class TestSaveTensorClone(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.parametrize('zero_stage', [1, 2])
+    @pytest.mark.parametrize('use_cpu_device', [True, False])
+    def test_save_tensor_clone(self, tmpdir, zero_stage, use_cpu_device):
+
+        config_dict = {
+            "optimizer": {
+                "type": "AdamW",
+            },
+            "zero_optimization": {
+                "stage": zero_stage
+            },
+            "train_batch_size": 1,
+            "train_micro_batch_size_per_gpu": 1
+        }
+        hidden_dim = 1024
+        model = SimpleModel(hidden_dim, nlayers=4).half()
+        ref_model_state_dict = model.state_dict()
+
+        ds_engine, _, _, _ = deepspeed.initialize(model=model, config_params=config_dict)
+        clone_device = torch.device('cpu') if use_cpu_device else get_accelerator().current_device()
+        clone_state_dict = clone_tensors_for_torch_save(ds_engine.module.state_dict())
+        compare_state_dicts(ref_model_state_dict, clone_state_dict)
+
+        ref_ckpt_file = os.path.join(tmpdir, 'ref_ckpt.pt')
+        torch.save(ref_model_state_dict, ref_ckpt_file)
+        clone_ckpt_file = os.path.join(tmpdir, 'clone_ckpt.pt')
+        torch.save(clone_state_dict, clone_ckpt_file)
+
+        compare_state_dicts(torch.load(ref_ckpt_file, weights_only=False),
+                            torch.load(clone_ckpt_file, weights_only=False))
+
+
+def test_elastic_checkpoint_is_deprecated_for_zero3(monkeypatch):
+    warning_messages = []
+
+    def mock_logger_warning(message, *args, **kwargs):
+        warning_messages.append(message)
+
+    monkeypatch.setattr("deepspeed.utils.logger.warning", mock_logger_warning)
+
+    DeepSpeedZeroConfig(stage=3, elastic_checkpoint=True)
+
+    assert any("elastic checkpointing is deprecated" in str(message).lower() for message in warning_messages)
+
+
+class TestZeRONonDistributed(DistributedTest):
+    world_size = 1
+    # This test calls deepspeed.initialize(), so use the harness' file-store
+    # initialization instead of env:// TCP rendezvous ports under xdist.
+    init_distributed = True
+
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    def test_chmod_exception_handling(self, monkeypatch, zero_stage):
+
+        config_dict = {
+            "optimizer": {
+                "type": "AdamW"
+            },
+            "train_batch_size": 1,
+            "zero_optimization": {
+                "stage": zero_stage
+            }
+        }
+        args = SimpleNamespace(local_rank=0)
+        net = SimpleModel(hidden_dim=4)
+        engine, _, _, _ = deepspeed.initialize(args=args,
+                                               config=config_dict,
+                                               model=net,
+                                               model_parameters=net.parameters())
+
+        log_called = False
+
+        def mock_logger_info(message, *args, **kwargs):
+            nonlocal log_called
+            log_called = True
+
+        monkeypatch.setattr("deepspeed.utils.logger.info", mock_logger_info)
+        """
+            This is presented for use-cases like Azure Storage File Share (where permissions are not allowed)
+            We use a fake file for this test (file not existing would present a similar issue as not being able to chmod)
+        """
+        fake_recovery_script_dst = os.path.join("tmp", "zero_to_fp32.py")
+        engine._change_recovery_script_permissions(fake_recovery_script_dst)
+
+        assert log_called, "Expected deepspeed.utils.logger.info to be called."
+
+
+class TestZeROPPLoadCheckpoint(DistributedTest):
+
+    world_size = 4
+
+    def test_load_zeropp_model(self, ws4_model_checkpoint_zeropp, class_tmpdir):
+        config_dict = {
+            "train_batch_size": 4,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "zero_hpz_partition_size": 2,
+                "stage3_param_persistence_threshold": 1
+            }
+        }
+
+        # Init model and load saved model
+        hidden_dim = 10
+        with deepspeed.zero.Init(config_dict_or_path=config_dict):
+            model = SimpleModel(hidden_dim)
+        ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
+
+        with deepspeed.zero.GatheredParameters(ds_model.module.parameters(), modifier_rank=0):
+            if dist.get_rank() == 0:
+                state_dict = torch.load(os.path.join(class_tmpdir, "model.pt"))
+                ds_model.module.load_state_dict(state_dict)
+
+        # Check the parameters after gather
+        params_to_gather = [p for p in ds_model.module.parameters() if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]
+        if len(params_to_gather) > 0:
+            handle = params_to_gather[0].all_gather_coalesced(params_to_gather)
+            handle.wait()
+        for ds_param in params_to_gather:
+            for v in ds_param.data.cpu().flatten().numpy():
+                assert v == 1.0
+
+    def test_load_zeropp_checkpoint(self, ws4_model_checkpoint_zeropp, class_tmpdir):
+        config_dict = {
+            "train_batch_size": 4,
+            "optimizer": {
+                "type": 'Adam'
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "zero_hpz_partition_size": 2,
+                "stage3_param_persistence_threshold": 1
+            }
+        }
+
+        # Init model and load zero checkpoint
+        hidden_dim = 10
+        model = SimpleModel(hidden_dim)
+        ds_model = create_deepspeed_model(config_dict=config_dict, model=model, base_optimizer=None)
+        ds_model.load_checkpoint(class_tmpdir,
+                                 load_optimizer_states=True,
+                                 load_lr_scheduler_states=False,
+                                 load_module_only=False)
+
+        # Check the parameters after gather
+        params_to_gather = [p for p in ds_model.module.parameters() if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]
+        if len(params_to_gather) > 0:
+            handle = params_to_gather[0].all_gather_coalesced(params_to_gather)
+            handle.wait()
+        for ds_param in params_to_gather:
+            for v in ds_param.data.cpu().flatten().numpy():
+                assert v == 1.0

tests/unit/comm/test_dist.py

@@ -0,0 +1,213 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import torch
+import deepspeed.comm as dist
+import deepspeed
+
+from unit.common import DistributedTest, DistributedFixture, get_master_port
+from unit.simple_model import SimpleModel
+from deepspeed.accelerator import get_accelerator
+
+import pytest
+from deepspeed.ops.op_builder import FusedAdamBuilder
+
+if not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME]:
+    pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+
+class TestInit(DistributedTest):
+    world_size = 3
+
+    def test(self):
+        assert dist.is_initialized()
+        assert dist.get_world_size() == 3
+        assert dist.get_rank() < 3
+
+
+# Demonstration of pytest's parameterization and fixtures
+@pytest.fixture(params=["hello"])
+def greeting(request):
+    return request.param
+
+
+@pytest.mark.parametrize("number,color", [(1138, "purple")])
+class TestDistArgs(DistributedTest):
+    world_size = 2
+    """ Classes that use DistributedTest class must define a test* method """
+
+    @pytest.mark.parametrize("shape", ["icosahedron"])
+    def test(self, number, color, shape, greeting):
+        """Ensure that we can parse args to DistributedTest methods. """
+        assert dist.get_world_size() == 2
+        assert number == 1138
+        assert color == "purple"
+        assert shape == "icosahedron"
+        assert greeting == "hello"
+
+
+# Demonstration of distributed tests grouped in single class
+@pytest.mark.parametrize("number", [1138])
+class TestGroupedDistTest(DistributedTest):
+    world_size = 2
+
+    def test_one(self, number):
+        assert dist.get_world_size() == 2
+        assert number == 1138
+
+    def test_two(self, number, color="purple"):
+        assert dist.get_world_size() == 2
+        assert number == 1138
+        assert color == "purple"
+
+
+# Demonstration of world_size override
+class TestWorldSizeOverrideDistTest(DistributedTest):
+    world_size = 2
+
+    def test_world_size_2(self):
+        assert dist.get_world_size() == 2
+
+    @pytest.mark.world_size(1)
+    def test_world_size_1(self):
+        assert dist.get_world_size() == 1
+
+
+# Demonstration of the DistributedFixture class
+@pytest.fixture(params=[2, 4])
+def val1(request):
+    return request.param
+
+
+@pytest.fixture(params=[16, 32])
+def val2(request):
+    return request.param
+
+
+class distributed_fixture(DistributedFixture):
+    world_size = 2
+
+    def run(self, class_tmpdir, val1, val2):
+        assert int(os.environ["WORLD_SIZE"]) == self.world_size
+        local_rank = os.environ["LOCAL_RANK"]
+        file_path = os.path.join(class_tmpdir, f"checkpoint-{local_rank}.pt")
+        with open(file_path, "w") as f:
+            f.write(f"{local_rank},{val1},{val2}")
+
+
+class TestDistributedFixture(DistributedTest):
+    world_size = 1
+
+    def test(self, distributed_fixture, class_tmpdir, val1, val2):
+        for rank in range(2):
+            file_path = os.path.join(class_tmpdir, f"checkpoint-{rank}.pt")
+            with open(file_path, "r") as f:
+                chkpt = f.read()
+            assert chkpt == f"{rank},{val1},{val2}"
+        assert int(os.environ["WORLD_SIZE"]) == 1
+
+
+@pytest.mark.parametrize("num_elements", [128, 3])
+class TestDistAllReduce(DistributedTest):
+    device_count = get_accelerator().device_count()
+    if device_count >= 4:
+        world_size = [1, 2, 4]
+    elif device_count >= 2:
+        world_size = [1, 2]
+    else:
+        world_size = [1]
+
+    def test(self, num_elements):
+        x = torch.ones(1, num_elements).to(get_accelerator().device_name()) * (dist.get_rank() + 1)
+        sum_of_ranks = (dist.get_world_size() * (dist.get_world_size() + 1)) // 2
+        result = torch.ones(1, num_elements).to(get_accelerator().device_name()) * sum_of_ranks
+        dist.all_reduce(x)
+        assert torch.all(x == result)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("num_elements", [128, 3])
+class TestDistInferenceAllReduce(DistributedTest):
+    device_count = get_accelerator().device_count()
+    if device_count >= 4:
+        world_size = [1, 2, 4]
+    elif device_count >= 2:
+        world_size = [1, 2]
+    else:
+        world_size = [1]
+
+    def test(self, dtype, num_elements):
+        x = torch.ones(1, num_elements).to(get_accelerator().device_name()) * (dist.get_rank() + 1)
+        sum_of_ranks = (dist.get_world_size() * (dist.get_world_size() + 1)) // 2
+        result = torch.ones(1, num_elements).to(get_accelerator().device_name()) * sum_of_ranks
+        result = result.to(dtype)
+        x = x.to(dtype)
+        dist.inference_all_reduce(x)
+        assert torch.all(x == result)
+
+
+@pytest.mark.parametrize("dist_init_required", [True, False, None])
+class TestDistInit(DistributedTest):
+    init_distributed = False
+
+    def test_already_init(self, dist_init_required):
+        torch.distributed.init_process_group(get_accelerator().communication_backend_name())
+        deepspeed.init_distributed(get_accelerator().communication_backend_name(),
+                                   dist_init_required=dist_init_required)
+
+    def test_no_init(self, dist_init_required):
+        if dist_init_required or dist_init_required is None:
+            deepspeed.init_distributed(get_accelerator().communication_backend_name(),
+                                       dist_init_required=dist_init_required)
+        else:
+            # torch.dist is not done and for some reason the user says they don't want it done
+            with pytest.raises(Exception):
+                deepspeed.init_distributed(get_accelerator().communication_backend_name(),
+                                           dist_init_required=dist_init_required)
+
+
+class TestDistInitNoEnv(DistributedTest):
+    world_size = 1
+    init_distributed = False
+    set_dist_env = False
+
+    def test(self):
+        torch.distributed.init_process_group(backend=get_accelerator().communication_backend_name(),
+                                             init_method=f"tcp://127.0.0.1:{get_master_port()}",
+                                             world_size=1,
+                                             rank=0)
+        assert torch.distributed.is_initialized()
+        deepspeed.init_distributed(get_accelerator().communication_backend_name(), auto_mpi_discovery=True)
+
+
+@pytest.mark.parametrize("dist_init_required", [True, False])
+class TestDistInitWithModel(DistributedTest):
+    init_distributed = False
+
+    def test_already_init(self, dist_init_required):
+        torch.distributed.init_process_group(get_accelerator().communication_backend_name())
+        model = SimpleModel(4)
+        config_dict = {"train_micro_batch_size_per_gpu": 1, "optimizer": {"type": "Adam", "params": {}}}
+        engine, *_ = deepspeed.initialize(model=model,
+                                          config=config_dict,
+                                          model_parameters=model.parameters(),
+                                          dist_init_required=dist_init_required)
+
+    def test_no_init(self, dist_init_required):
+        model = SimpleModel(4)
+        config_dict = {"train_micro_batch_size_per_gpu": 1, "optimizer": {"type": "Adam", "params": {}}}
+        if dist_init_required:
+            engine, *_ = deepspeed.initialize(model=model,
+                                              config=config_dict,
+                                              model_parameters=model.parameters(),
+                                              dist_init_required=dist_init_required)
+        else:
+            # torch.dist is not done and for some reason the user says they don't want it done
+            with pytest.raises(Exception):
+                engine, *_ = deepspeed.initialize(model=model,
+                                                  config=config_dict,
+                                                  model_parameters=model.parameters(),
+                                                  dist_init_required=dist_init_required)

tests/unit/common.py

@@ -0,0 +1,602 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import re
+import time
+import inspect
+import socket
+import subprocess
+from abc import ABC, abstractmethod
+from pathlib import Path
+import random
+import tempfile
+import numpy as np
+from typing import Callable, Any
+
+import torch
+import torch.multiprocessing as mp
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+import deepspeed.comm as dist
+
+from .util import torch_assert_close
+
+import pytest
+from _pytest.outcomes import Skipped
+from _pytest.fixtures import FixtureLookupError, FixtureFunctionMarker
+
+# Worker timeout for tests that hang
+DEEPSPEED_TEST_TIMEOUT = int(os.environ.get('DS_UNITTEST_TIMEOUT', '600'))
+
+
+def is_rocm_pytorch():
+    return hasattr(torch.version, 'hip') and torch.version.hip is not None
+
+
+def get_xdist_worker_id():
+    xdist_worker = os.environ.get('PYTEST_XDIST_WORKER', None)
+    if xdist_worker is not None:
+        xdist_worker_id = xdist_worker.replace('gw', '')
+        return int(xdist_worker_id)
+    return None
+
+
+def get_master_port(base_port=29500, port_range_size=1000):
+    xdist_worker_id = get_xdist_worker_id()
+    if xdist_worker_id is not None:
+        # Make xdist workers use different port ranges to avoid race conditions
+        base_port += port_range_size * xdist_worker_id
+
+    # Select first open port in range
+    port = base_port
+    max_port = base_port + port_range_size
+    sock = socket.socket()
+    while port < max_port:
+        try:
+            sock.bind(('', port))
+            sock.close()
+            return str(port)
+        except OSError:
+            port += 1
+    raise IOError('no free ports')
+
+
+def _get_cpu_socket_count():
+    import shlex
+    p1 = subprocess.Popen(shlex.split("cat /proc/cpuinfo"), stdout=subprocess.PIPE)
+    p2 = subprocess.Popen(["grep", "physical id"], stdin=p1.stdout, stdout=subprocess.PIPE)
+    p1.stdout.close()
+    p3 = subprocess.Popen(shlex.split("sort -u"), stdin=p2.stdout, stdout=subprocess.PIPE)
+    p2.stdout.close()
+    p4 = subprocess.Popen(shlex.split("wc -l"), stdin=p3.stdout, stdout=subprocess.PIPE)
+    p3.stdout.close()
+    r = int(p4.communicate()[0])
+    p4.stdout.close()
+    return r
+
+
+def set_accelerator_visible():
+    cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+    xdist_worker_id = get_xdist_worker_id()
+    if xdist_worker_id is None:
+        xdist_worker_id = 0
+    if cuda_visible is None:
+        # CUDA_VISIBLE_DEVICES is not set, discover it using accelerator specific command instead
+        if get_accelerator().device_name() == 'cuda':
+            if is_rocm_pytorch():
+                rocm_smi = subprocess.check_output(['rocm-smi', '--showid'])
+                gpu_ids = filter(lambda s: 'GPU' in s, rocm_smi.decode('utf-8').strip().split('\n'))
+                num_accelerators = len(list(gpu_ids))
+            else:
+                nvidia_smi = subprocess.check_output(['nvidia-smi', '--list-gpus'])
+                num_accelerators = len(nvidia_smi.decode('utf-8').strip().split('\n'))
+        elif get_accelerator().device_name() == 'xpu':
+            clinfo = subprocess.check_output(['clinfo'])
+            lines = clinfo.decode('utf-8').strip().split('\n')
+            num_accelerators = 0
+            for line in lines:
+                match = re.search('Device Type.*GPU', line)
+                if match:
+                    num_accelerators += 1
+        elif get_accelerator().device_name() == 'hpu':
+            try:
+                hl_smi = subprocess.check_output(['hl-smi', "-L"])
+                num_accelerators = re.findall(r"Module ID\s+:\s+(\d+)", hl_smi.decode())
+            except FileNotFoundError:
+                sim_list = subprocess.check_output(['ls', '-1', '/dev/accel'])
+                num_accelerators = re.findall(r"accel(\d+)", sim_list.decode())
+            num_accelerators = sorted(num_accelerators, key=int)
+            os.environ["HABANA_VISIBLE_MODULES"] = ",".join(num_accelerators)
+        elif get_accelerator().device_name() == 'npu':
+            npu_smi = subprocess.check_output(['npu-smi', 'info', '-l'])
+            num_accelerators = int(npu_smi.decode('utf-8').strip().split('\n')[0].split(':')[1].strip())
+        elif get_accelerator().device_name() == 'supa':
+            br_smi = subprocess.check_output(['brsmi', 'gpu', 'list'])
+            gpu_ids = filter(lambda s: 'GPU' in s, br_smi.decode('utf-8').strip().split('\n'))
+            num_accelerators = len(list(gpu_ids))
+        else:
+            assert get_accelerator().device_name() == 'cpu'
+            num_accelerators = _get_cpu_socket_count()
+
+        if isinstance(num_accelerators, list):
+            cuda_visible = ",".join(num_accelerators)
+        else:
+            cuda_visible = ",".join(map(str, range(num_accelerators)))
+
+    # rotate list based on xdist worker id, example below
+    # wid=0 -> ['0', '1', '2', '3']
+    # wid=1 -> ['1', '2', '3', '0']
+    # wid=2 -> ['2', '3', '0', '1']
+    # wid=3 -> ['3', '0', '1', '2']
+    dev_id_list = cuda_visible.split(",")
+    dev_id_list = dev_id_list[xdist_worker_id:] + dev_id_list[:xdist_worker_id]
+    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(dev_id_list)
+
+
+class DistributedExec(ABC):
+    """
+    Base class for distributed execution of functions/methods. Contains common
+    methods needed for DistributedTest and DistributedFixture.
+    """
+    world_size = 2
+    backend = get_accelerator().communication_backend_name()
+    init_distributed = True
+    set_dist_env = True
+    requires_cuda_env = True
+    reuse_dist_env = False
+    non_daemonic_procs = False
+    _pool_cache = {}
+    exec_timeout = DEEPSPEED_TEST_TIMEOUT
+
+    @abstractmethod
+    def run(self):
+        ...
+
+    def __call__(self, request):
+        self._fixture_kwargs = self._get_fixture_kwargs(request, self.run)
+        world_size = self.world_size
+        if self.requires_cuda_env and not get_accelerator().is_available():
+            pytest.skip("only supported in accelerator environments.")
+
+        self._launch_with_file_store(request, world_size)
+
+    def _get_fixture_kwargs(self, request, func):
+        if not request:
+            return {}
+        # Grab fixture / parametrize kwargs from pytest request object
+        fixture_kwargs = {}
+        params = inspect.getfullargspec(func).args
+        params.remove("self")
+        for p in params:
+            try:
+                fixture_kwargs[p] = request.getfixturevalue(p)
+            except FixtureLookupError:
+                pass  # test methods can have kwargs that are not fixtures
+        return fixture_kwargs
+
+    def _launch_daemonic_procs(self, num_procs, init_method):
+        # Create process pool or use cached one
+        master_port = None
+
+        if get_accelerator().device_name() == 'hpu':
+            if self.reuse_dist_env:
+                print("Ignoring reuse_dist_env for hpu")
+                self.reuse_dist_env = False
+
+        if self.reuse_dist_env:
+            if num_procs not in self._pool_cache:
+                self._pool_cache[num_procs] = mp.Pool(processes=num_procs)
+                master_port = get_master_port()
+            pool = self._pool_cache[num_procs]
+        else:
+            pool = mp.Pool(processes=num_procs)
+            master_port = get_master_port()
+
+        # Run the test
+        args = [(local_rank, num_procs, master_port, init_method) for local_rank in range(num_procs)]
+        skip_msgs_async = pool.starmap_async(self._dist_run, args)
+
+        try:
+            skip_msgs = skip_msgs_async.get(self.exec_timeout)
+        except mp.TimeoutError:
+            # Shortcut to exit pytest in the case of a hanged test. This
+            # usually means an environment error and the rest of tests will
+            # hang (causing super long unit test runtimes)
+            pytest.exit("Test hanged, exiting", returncode=1)
+        finally:
+            # Regardless of the outcome, ensure proper teardown
+            # Tear down distributed environment and close process pools
+            self._close_pool(pool, num_procs)
+
+        # If we skipped a test, propagate that to this process
+        if any(skip_msgs):
+            assert len(set(skip_msgs)) == 1, "Multiple different skip messages received"
+            pytest.skip(skip_msgs[0])
+
+    def _launch_non_daemonic_procs(self, num_procs, init_method):
+        assert not self.reuse_dist_env, "Cannot reuse distributed environment with non-daemonic processes"
+
+        master_port = get_master_port()
+        skip_msg = mp.Queue()  # Allows forked processes to share pytest.skip reason
+        processes = []
+        prev_start_method = mp.get_start_method()
+        mp.set_start_method('spawn', force=True)
+        for local_rank in range(num_procs):
+            p = mp.Process(target=self._dist_run, args=(local_rank, num_procs, master_port, init_method, skip_msg))
+            p.start()
+            processes.append(p)
+        mp.set_start_method(prev_start_method, force=True)
+
+        # Now loop and wait for a test to complete. The spin-wait here isn't a big
+        # deal because the number of processes will be O(#GPUs) << O(#CPUs).
+        any_done = False
+        start = time.time()
+        while (not any_done) and ((time.time() - start) < self.exec_timeout):
+            for p in processes:
+                if not p.is_alive():
+                    any_done = True
+                    break
+            time.sleep(.1)  # So we don't hog CPU
+
+        # If we hit the timeout, then presume a test is hanged
+        if not any_done:
+            for p in processes:
+                p.terminate()
+            pytest.exit("Test hanged, exiting", returncode=1)
+
+        # Wait for all other processes to complete
+        for p in processes:
+            p.join(self.exec_timeout)
+
+        failed = [(rank, p) for rank, p in enumerate(processes) if p.exitcode != 0]
+        for rank, p in failed:
+            # If it still hasn't terminated, kill it because it hung.
+            if p.exitcode is None:
+                p.terminate()
+                pytest.fail(f'Worker {rank} hung.', pytrace=False)
+            if p.exitcode < 0:
+                pytest.fail(f'Worker {rank} killed by signal {-p.exitcode}', pytrace=False)
+            if p.exitcode > 0:
+                pytest.fail(f'Worker {rank} exited with code {p.exitcode}', pytrace=False)
+
+        if not skip_msg.empty():
+            # This assumed all skip messages are the same, it may be useful to
+            # add a check here to assert all exit messages are equal
+            pytest.skip(skip_msg.get())
+
+    def _launch_procs(self, num_procs, init_method):
+        # Verify we have enough accelerator devices to run this test
+        if get_accelerator().is_available() and get_accelerator().device_count() < num_procs:
+            pytest.skip(
+                f"Skipping test because not enough GPUs are available: {num_procs} required, {get_accelerator().device_count()} available"
+            )
+
+        if get_accelerator().device_name() == 'xpu':
+            self.non_daemonic_procs = True
+            self.reuse_dist_env = False
+
+        # Allow disabling reuse_dist_env via environment variable.
+        # This is useful for CI full test runs where reusing distributed environment
+        # can cause pool worker cleanup to hang after tests complete.
+        if os.environ.get('DS_DISABLE_REUSE_DIST_ENV', '0') == '1':
+            self.reuse_dist_env = False
+
+        # Set start method to `forkserver` (or `fork`)
+        mp.set_start_method('forkserver', force=True)
+
+        if self.non_daemonic_procs:
+            self._launch_non_daemonic_procs(num_procs, init_method)
+        else:
+            self._launch_daemonic_procs(num_procs, init_method)
+
+    def _dist_run(self, local_rank, num_procs, master_port, init_method, skip_msg=""):
+        if dist.is_initialized():
+            if get_accelerator().is_available():
+                # local_rank might not match the rank in the previous run if you are reusing the environment
+                get_accelerator().set_device(dist.get_rank())
+        else:
+            """ Initialize deepspeed.comm and execute the user function. """
+            if self.set_dist_env:
+                os.environ['MASTER_ADDR'] = '127.0.0.1'
+                os.environ['MASTER_PORT'] = str(master_port)
+                os.environ['LOCAL_RANK'] = str(local_rank)
+                # NOTE: unit tests don't support multi-node so local_rank == global rank
+                os.environ['RANK'] = str(local_rank)
+                # In case of multiprocess launching LOCAL_SIZE should be same as WORLD_SIZE
+                # DeepSpeed single node launcher would also set LOCAL_SIZE accordingly
+                os.environ['LOCAL_SIZE'] = str(num_procs)
+                os.environ['WORLD_SIZE'] = str(num_procs)
+
+            # turn off NCCL logging if set
+            os.environ.pop('NCCL_DEBUG', None)
+
+            if get_accelerator().is_available():
+                set_accelerator_visible()
+
+            if get_accelerator().is_available():
+                get_accelerator().set_device(local_rank)
+
+            if self.init_distributed:
+                deepspeed.init_distributed(dist_backend=self.backend,
+                                           init_method=init_method,
+                                           rank=local_rank,
+                                           world_size=num_procs)
+                dist.barrier()
+
+        try:
+            self.run(**self._fixture_kwargs)
+        except BaseException as e:
+            if isinstance(e, Skipped):
+                if self.non_daemonic_procs:
+                    skip_msg.put(e.msg)
+                else:
+                    skip_msg = e.msg
+            else:
+                raise e
+
+        return skip_msg
+
+    def _launch_with_file_store(self, request, world_size):
+        tmpdir = request.getfixturevalue("tmpdir")
+
+        if isinstance(world_size, int):
+            world_size = [world_size]
+        for procs in world_size:
+            with tempfile.NamedTemporaryFile(delete=False, dir=str(tmpdir), suffix='_filestore') as fp:
+                init_method = f"file://{fp.name}"
+            self._launch_procs(procs, init_method)
+            time.sleep(0.5)
+
+    def _dist_destroy(self):
+        if (dist is not None) and dist.is_initialized():
+            dist.barrier()
+            dist.destroy_process_group()
+
+    def _close_pool(self, pool, num_procs, force=False):
+        if force or not self.reuse_dist_env:
+            pool.starmap(self._dist_destroy, [() for _ in range(num_procs)])
+            pool.close()
+            pool.join()
+
+
+class DistributedFixture(DistributedExec):
+    """
+    Implementation that extends @pytest.fixture to allow for distributed execution.
+    This is primarily meant to be used when a test requires executing two pieces of
+    code with different world sizes.
+
+    There are 2 parameters that can be modified:
+        - world_size: int = 2 -- the number of processes to launch
+        - backend: Literal['nccl','mpi','gloo'] = 'nccl' -- which backend to use
+
+    Features:
+        - able to call pytest.skip() inside fixture
+        - can be reused by multiple tests
+        - can accept other fixtures as input
+
+    Limitations:
+        - cannot use @pytest.mark.parametrize
+        - world_size cannot be modified after definition and only one world_size value is accepted
+        - any fixtures used must also be used in the test that uses this fixture (see example below)
+        - return values cannot be returned. Passing values to a DistributedTest
+          object can be achieved using class_tmpdir and writing to file (see example below)
+
+    Usage:
+        - must implement a run(self, ...) method
+        - fixture can be used by making the class name input to a test function
+
+    Example:
+        @pytest.fixture(params=[10,20])
+        def regular_pytest_fixture(request):
+            return request.param
+
+        class distributed_fixture_example(DistributedFixture):
+            world_size = 4
+
+            def run(self, regular_pytest_fixture, class_tmpdir):
+                assert int(os.environ["WORLD_SIZE"]) == self.world_size
+                local_rank = os.environ["LOCAL_RANK"]
+                print(f"Rank {local_rank} with value {regular_pytest_fixture}")
+                with open(os.path.join(class_tmpdir, f"{local_rank}.txt"), "w") as f:
+                    f.write(f"{local_rank},{regular_pytest_fixture}")
+
+        class TestExample(DistributedTest):
+            world_size = 1
+
+            def test(self, distributed_fixture_example, regular_pytest_fixture, class_tmpdir):
+                assert int(os.environ["WORLD_SIZE"]) == self.world_size
+                for rank in range(4):
+                    with open(os.path.join(class_tmpdir, f"{rank}.txt"), "r") as f:
+                        assert f.read() == f"{rank},{regular_pytest_fixture}"
+    """
+    is_dist_fixture = True
+
+    # These values are just placeholders so that pytest recognizes this as a fixture
+    _pytestfixturefunction = FixtureFunctionMarker(scope="function", params=None)
+    __name__ = ""
+
+    def __init__(self):
+        assert isinstance(self.world_size, int), "Only one world size is allowed for distributed fixtures"
+        self.__name__ = type(self).__name__
+        _pytestfixturefunction = FixtureFunctionMarker(scope="function", params=None, name=self.__name__)
+
+
+class DistributedTest(DistributedExec):
+    """
+    Implementation for running pytest with distributed execution.
+
+    There are 2 parameters that can be modified:
+        - world_size: Union[int,List[int]] = 2 -- the number of processes to launch
+        - backend: Literal['nccl','mpi','gloo'] = 'nccl' -- which backend to use
+
+    Features:
+        - able to call pytest.skip() inside tests
+        - works with pytest fixtures, parametrize, mark, etc.
+        - can contain multiple tests (each of which can be parametrized separately)
+        - class methods can be fixtures (usable by tests in this class only)
+        - world_size can be changed for individual tests using @pytest.mark.world_size(world_size)
+        - class_tmpdir is a fixture that can be used to get a tmpdir shared among
+          all tests (including DistributedFixture)
+
+    Usage:
+        - class name must start with "Test"
+        - must implement one or more test*(self, ...) methods
+
+    Example:
+        @pytest.fixture(params=[10,20])
+        def val1(request):
+            return request.param
+
+        @pytest.mark.fast
+        @pytest.mark.parametrize("val2", [30,40])
+        class TestExample(DistributedTest):
+            world_size = 2
+
+            @pytest.fixture(params=[50,60])
+            def val3(self, request):
+                return request.param
+
+            def test_1(self, val1, val2, str1="hello world"):
+                assert int(os.environ["WORLD_SIZE"]) == self.world_size
+                assert all(val1, val2, str1)
+
+            @pytest.mark.world_size(1)
+            @pytest.mark.parametrize("val4", [70,80])
+            def test_2(self, val1, val2, val3, val4):
+                assert int(os.environ["WORLD_SIZE"]) == 1
+                assert all(val1, val2, val3, val4)
+    """
+    is_dist_test = True
+
+    # Temporary directory that is shared among test methods in a class
+    @pytest.fixture(autouse=True, scope="class")
+    def class_tmpdir(self, tmpdir_factory):
+        fn = tmpdir_factory.mktemp(self.__class__.__name__)
+        return fn
+
+    def run(self, **fixture_kwargs):
+        self._current_test(**fixture_kwargs)
+
+    def __call__(self, request):
+        self._current_test = self._get_current_test_func(request)
+        self._fixture_kwargs = self._get_fixture_kwargs(request, self._current_test)
+
+        if self.requires_cuda_env and not get_accelerator().is_available():
+            pytest.skip("only supported in accelerator environments.")
+
+        # Catch world_size override pytest mark
+        for mark in getattr(request.function, "pytestmark", []):
+            if mark.name == "world_size":
+                world_size = mark.args[0]
+                break
+        else:
+            world_size = self._fixture_kwargs.get("world_size", self.world_size)
+
+        self._launch_with_file_store(request, world_size)
+
+    def _get_current_test_func(self, request):
+        # DistributedTest subclasses may have multiple test methods
+        func_name = request.function.__name__
+        return getattr(self, func_name)
+
+
+def get_test_path(filename):
+    curr_path = Path(__file__).parent
+    return str(curr_path.joinpath(filename))
+
+
+# bf16 > fp16 > fp32
+def preferred_dtype():
+    if get_accelerator().is_bf16_supported():
+        return torch.bfloat16
+    elif get_accelerator().is_fp16_supported():
+        return torch.float16
+    else:
+        return torch.float32
+
+
+class EnableDeterminism:
+
+    def __init__(self, seed: int):
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+        self.seed = seed + local_rank
+        self.saved_random_state = None
+        self.saved_np_random_state = None
+        self.saved_cuda_launch_blocking = None
+        self.saved_cublas_workspace_config = None
+        self.saved_deterministic_algorithms = None
+
+    def __enter__(self):
+        self.saved_random_state = random.getstate()
+        self.saved_np_random_state = np.random.get_state()
+        self.saved_acc_rng_state = get_accelerator().get_rng_state()
+        self.saved_cuda_launch_blocking = os.environ.get("CUDA_LAUNCH_BLOCKING", "")
+        self.saved_cublas_workspace_config = os.environ.get("CUBLAS_WORKSPACE_CONFIG", "")
+        self.saved_deterministic_algorithms = torch.are_deterministic_algorithms_enabled()
+
+        random.seed(self.seed)
+        np.random.seed(self.seed)
+        get_accelerator().manual_seed(self.seed)
+        get_accelerator().manual_seed_all(self.seed)
+
+        os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
+        torch.use_deterministic_algorithms(True)
+
+    def __exit__(self, type, value, traceback):
+        random.setstate(self.saved_random_state)
+        np.random.set_state(self.saved_np_random_state)
+        get_accelerator().set_rng_state(self.saved_acc_rng_state)
+        os.environ["CUDA_LAUNCH_BLOCKING"] = self.saved_cuda_launch_blocking
+        os.environ["CUBLAS_WORKSPACE_CONFIG"] = self.saved_cublas_workspace_config
+        torch.use_deterministic_algorithms(self.saved_deterministic_algorithms)
+
+
+def enable_determinism(seed: int):
+
+    def decorator(func: Callable) -> Callable:
+
+        def wrapper(*args: Any, **kwargs: Any):
+            with EnableDeterminism(seed):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def reduce_boolean_flags(flag: bool, op=all) -> bool:
+    if not dist.is_initialized():
+        return flag
+    device = get_accelerator().current_device()
+    tensor_flag = torch.tensor(1 if flag else 0, dtype=torch.int, device=device)
+    world_size = dist.get_world_size()
+    tensor_flag_buf = torch.zeros(world_size, dtype=torch.int, device=device)
+    dist.all_gather_into_tensor(tensor_flag_buf, tensor_flag)
+    list_flags = [bool(f) for f in tensor_flag_buf.tolist()]
+    return op(list_flags)
+
+
+def allclose_on_all_ranks(actual, expected, assert_message=None, **kwargs) -> None:
+    """
+    Compare two tensors across all ranks.
+    We want to make sure that all ranks succeed or fail together.
+    """
+    allclose_local = False
+    allclose_global = False
+    mismatch_msg = ""
+    try:
+        torch_assert_close(actual, expected, **kwargs)
+        allclose_local = True
+        allclose_global = reduce_boolean_flags(allclose_local, all)
+    except AssertionError:
+        allclose_global = reduce_boolean_flags(allclose_local, all)
+        mismatch_msg = f"Tensors are not close: {actual=}, {expected=} {kwargs=}"
+
+    if not allclose_global:
+        message = "Tensors are not close on all ranks." if assert_message is None else assert_message
+        raise AssertionError(f"{message} {mismatch_msg}")

tests/unit/compile/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team

tests/unit/compile/test_inductor_aot_kwargs.py

@@ -0,0 +1,141 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed.compile.inductor as inductor
+
+
+def _compiler(name):
+
+    def compiler(*args, **kwargs):
+        return name, args, kwargs
+
+    compiler.__name__ = name
+    return compiler
+
+
+def _install_patch_spies(monkeypatch):
+    compiler_calls = []
+    partition_calls = []
+
+    def fake_patch_compiler(original_compiler, dc_compiler, z3_partition, graph_id, graph_param_manager, bwd):
+        wrapped = {
+            "original_compiler": original_compiler,
+            "dc_compiler": dc_compiler,
+            "z3_partition": z3_partition,
+            "graph_id": graph_id,
+            "graph_param_manager": graph_param_manager,
+            "bwd": bwd,
+        }
+        compiler_calls.append(wrapped)
+        return wrapped
+
+    def fake_wrap_partition_fn(z3_partition, partition_fn, real_inputs, param_indices, frame_id, frames_partitioned):
+        wrapped = {
+            "z3_partition": z3_partition,
+            "partition_fn": partition_fn,
+            "real_inputs": real_inputs,
+            "param_indices": param_indices,
+            "frame_id": frame_id,
+            "frames_partitioned": frames_partitioned,
+        }
+        partition_calls.append(wrapped)
+        return wrapped
+
+    monkeypatch.setattr(inductor, "patch_compiler", fake_patch_compiler)
+    monkeypatch.setattr(inductor, "wrap_partition_fn", fake_wrap_partition_fn)
+    return compiler_calls, partition_calls
+
+
+def _patch_kwargs(kwargs, monkeypatch):
+    compiler_calls, partition_calls = _install_patch_spies(monkeypatch)
+    make_fw_graph = object()
+    make_bw_graph = object()
+    real_inputs = object()
+    param_indices = object()
+    param_manager = object()
+    frames_partitioned = set()
+
+    applied = inductor._patch_deepcompile_aot_kwargs(kwargs,
+                                                     graph_id=7,
+                                                     z3_partition=True,
+                                                     make_fw_graph=make_fw_graph,
+                                                     make_bw_graph=make_bw_graph,
+                                                     real_inputs=real_inputs,
+                                                     param_indices=param_indices,
+                                                     param_manager=param_manager,
+                                                     frame_id=11,
+                                                     frames_partitioned=frames_partitioned)
+
+    return {
+        "applied": applied,
+        "compiler_calls": compiler_calls,
+        "partition_calls": partition_calls,
+        "make_fw_graph": make_fw_graph,
+        "make_bw_graph": make_bw_graph,
+        "real_inputs": real_inputs,
+        "param_indices": param_indices,
+        "param_manager": param_manager,
+        "frames_partitioned": frames_partitioned,
+    }
+
+
+def test_legacy_inductor_shape_wraps_explicit_bw_compiler(monkeypatch):
+    fw_compiler = _compiler("fw")
+    bw_compiler = _compiler("bw")
+    inference_compiler = _compiler("inference")
+    partition_fn = _compiler("partition")
+    kwargs = {
+        "fw_compiler": fw_compiler,
+        "bw_compiler": bw_compiler,
+        "inference_compiler": inference_compiler,
+        "partition_fn": partition_fn,
+    }
+
+    result = _patch_kwargs(kwargs, monkeypatch)
+
+    assert result["applied"] is True
+    assert len(result["compiler_calls"]) == 2
+    assert result["compiler_calls"][0]["original_compiler"] is fw_compiler
+    assert result["compiler_calls"][0]["dc_compiler"] is result["make_fw_graph"]
+    assert result["compiler_calls"][0]["bwd"] is False
+    assert result["compiler_calls"][1]["original_compiler"] is bw_compiler
+    assert result["compiler_calls"][1]["dc_compiler"] is result["make_bw_graph"]
+    assert result["compiler_calls"][1]["bwd"] is True
+    assert kwargs["fw_compiler"] is result["compiler_calls"][0]
+    assert kwargs["bw_compiler"] is result["compiler_calls"][1]
+    assert kwargs["inference_compiler"] is result["compiler_calls"][0]
+    assert kwargs["partition_fn"] is result["partition_calls"][0]
+    assert result["partition_calls"][0]["partition_fn"] is partition_fn
+
+
+def test_missing_bw_compiler_uses_original_fw_compiler_for_backward(monkeypatch):
+    fw_compiler = _compiler("fw")
+    partition_fn = _compiler("partition")
+    kwargs = {
+        "fw_compiler": fw_compiler,
+        "partition_fn": partition_fn,
+    }
+
+    result = _patch_kwargs(kwargs, monkeypatch)
+
+    assert result["applied"] is True
+    assert result["compiler_calls"][0]["original_compiler"] is fw_compiler
+    assert result["compiler_calls"][0]["bwd"] is False
+    assert result["compiler_calls"][1]["original_compiler"] is fw_compiler
+    assert result["compiler_calls"][1]["dc_compiler"] is result["make_bw_graph"]
+    assert result["compiler_calls"][1]["bwd"] is True
+    assert kwargs["bw_compiler"] is result["compiler_calls"][1]
+
+
+def test_torchxla_openxla_shape_passes_through_unchanged(monkeypatch):
+    kwargs = {"fw_compiler": _compiler("openxla_eval_boxed")}
+    original_kwargs = dict(kwargs)
+
+    result = _patch_kwargs(kwargs, monkeypatch)
+
+    assert result["applied"] is False
+    assert result["compiler_calls"] == []
+    assert result["partition_calls"] == []
+    assert kwargs == original_kwargs

tests/unit/compile/test_list_schedule.py

@@ -0,0 +1,406 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import operator
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch.fx import Graph, GraphModule
+
+import deepspeed.compile.util as compile_util
+from deepspeed.compile import backend as backend_mod
+from deepspeed.compile import inductor as inductor_mod
+from deepspeed.compile import list_schedule as schedule_mod
+from deepspeed.compile.passes import prefetch as prefetch_mod
+from deepspeed.compile.passes import selective_gather as selective_gather_mod
+from deepspeed.compile.profilers import ProfilingResult
+from deepspeed.compile.profilers.graph_profile import _backfill_missing_profile_metadata, is_profile_incomplete
+
+_DC_LIBRARIES = []
+
+
+def _define_dc_ops():
+    try:
+        torch.ops.dc.allgather_param.default
+        torch.ops.dc.wait_allgather.default
+        torch.ops.dc.release_param.default
+        torch.ops.dc.reduce_grad.default
+        return
+    except AttributeError:
+        pass
+
+    lib = torch.library.Library("dc", "DEF")
+    for schema in (
+            "allgather_param(Tensor a, int graph_id, int id, ScalarType? dtype = None) -> Tensor",
+            "wait_allgather(Tensor(a) a, int graph_id, int id) -> Tensor(a)",
+            "release_param(Tensor(a) a, int graph_id, int id, int n_users) -> Tensor(a)",
+            "reduce_grad(Tensor a, int graph_id, int id) -> Tensor",
+            "free_tensors(Tensor[] tensors) -> ()",
+            "end_backward(Tensor[] tensors, int graph_id, bool release_reduce_buckets = True) -> ()",
+    ):
+        try:
+            lib.define(schema)
+        except RuntimeError as exc:
+            if "already been registered" not in str(exc):
+                raise
+    _DC_LIBRARIES.append(lib)
+
+
+@pytest.fixture(autouse=True)
+def stub_deepcompile_ops(monkeypatch):
+    _define_dc_ops()
+    no_copy_ops = {torch.ops.dc.wait_allgather.default}
+    monkeypatch.setattr(compile_util, "get_no_copy_ops", lambda: no_copy_ops)
+
+
+def _with_meta(node, tensor_size=0, device_time=0):
+    node.meta["tensor_size"] = tensor_size
+    if device_time is not None:
+        node.meta["device_time"] = device_time
+    return node
+
+
+def _placeholder(graph, name):
+    return _with_meta(graph.placeholder(name))
+
+
+def test_sync_memory_profile_complete_noops_without_distributed(monkeypatch):
+    monkeypatch.setattr(backend_mod.dist, "is_initialized", lambda: False)
+
+    def fail_all_reduce(*args, **kwargs):
+        raise AssertionError("all_reduce should not run without distributed init")
+
+    monkeypatch.setattr(backend_mod.dist, "all_reduce", fail_all_reduce)
+
+    assert backend_mod._sync_memory_profile_complete(True)
+    assert not backend_mod._sync_memory_profile_complete(False)
+
+
+def test_sync_memory_profile_complete_reduces_asymmetric_failure(monkeypatch):
+    monkeypatch.setattr(backend_mod.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(backend_mod, "get_accelerator", lambda: SimpleNamespace(current_device=lambda: "cpu"))
+
+    def mark_any_rank_failed(tensor, op):
+        assert op == backend_mod.dist.ReduceOp.MIN
+        tensor[0] = 0
+
+    monkeypatch.setattr(backend_mod.dist, "all_reduce", mark_any_rank_failed)
+
+    assert not backend_mod._sync_memory_profile_complete(True)
+
+
+def _allgather(graph, arg, ds_id, name, tensor_size=1, device_time=1):
+    return _with_meta(
+        graph.call_function(torch.ops.dc.allgather_param.default, (arg, 0, ds_id), {"dtype": torch.float16},
+                            name=f"allgather_ds_param_{name}_{ds_id}"),
+        tensor_size=tensor_size,
+        device_time=device_time,
+    )
+
+
+def _wait(graph, arg, ds_id, name):
+    return _with_meta(
+        graph.call_function(torch.ops.dc.wait_allgather.default, (arg, 0, ds_id),
+                            name=f"wait_allgather_ds_param_{name}_{ds_id}"))
+
+
+def _neg(graph, arg, name, device_time=0):
+    return _with_meta(graph.call_function(operator.neg, (arg, ), name=name), device_time=device_time)
+
+
+def _add(graph, lhs, rhs, name, device_time=0):
+    return _with_meta(graph.call_function(operator.add, (lhs, rhs), name=name), device_time=device_time)
+
+
+def _release(graph, arg, ds_id, name):
+    return _with_meta(
+        graph.call_function(torch.ops.dc.release_param.default, (arg, 0, ds_id, 1),
+                            name=f"release_ds_param_{name}_{ds_id}"))
+
+
+def _scheduled_names(graph):
+    return [node.name for node in schedule_mod.fast_free_schedule(graph, 0, 0, debug_log=True).nodes]
+
+
+def test_fast_free_schedule_keeps_zero_free_acc_filter():
+    graph = Graph()
+
+    safe_param = _placeholder(graph, "safe_param")
+    safe_pre_param = _placeholder(graph, "safe_pre_param")
+    unsafe_param = _placeholder(graph, "unsafe_param")
+    unsafe_extra_param = _placeholder(graph, "unsafe_extra_param")
+
+    safe_pre_ag = _allgather(graph, safe_pre_param, 10, "safe_pre")
+    safe_pre_wait = _wait(graph, safe_pre_ag, 10, "safe_pre")
+    safe_pre_use = _neg(graph, safe_pre_wait, "safe_pre_use")
+    safe_ag = _allgather(graph, _add(graph, safe_param, safe_pre_use, "safe_param_dep"), 11, "safe")
+    safe_wait = _wait(graph, safe_ag, 11, "safe")
+    safe_use = _neg(graph, safe_wait, "safe_use", device_time=100)
+    safe_release = _release(graph, safe_use, 11, "safe")
+
+    unsafe_ag = _allgather(graph, unsafe_param, 20, "unsafe")
+    unsafe_wait = _wait(graph, unsafe_ag, 20, "unsafe")
+    unsafe_extra_ag = _allgather(graph, unsafe_extra_param, 21, "unsafe_extra")
+    unsafe_extra_wait = _wait(graph, unsafe_extra_ag, 21, "unsafe_extra")
+    unsafe_use = _add(graph, unsafe_wait, unsafe_extra_wait, "unsafe_use", device_time=1)
+    unsafe_release = _release(graph, unsafe_use, 20, "unsafe")
+
+    graph.output((safe_release, unsafe_release))
+    graph.lint()
+
+    names = _scheduled_names(graph)
+
+    assert names.index(safe_release.name) < names.index(unsafe_ag.name)
+    assert names.index(safe_release.name) < names.index(unsafe_extra_ag.name)
+
+
+def test_fast_free_schedule_prefers_lower_allgather_pressure_in_zero_free_acc_bucket():
+    graph = Graph()
+
+    high_param = _placeholder(graph, "high_param")
+    high_pre_param = _placeholder(graph, "high_pre_param")
+    low_param = _placeholder(graph, "low_param")
+    low_pre_param = _placeholder(graph, "low_pre_param")
+
+    high_pre_ag = _allgather(graph, high_pre_param, 30, "high_pre", tensor_size=100)
+    high_pre_wait = _wait(graph, high_pre_ag, 30, "high_pre")
+    high_ag = _allgather(graph, _add(graph, high_param, high_pre_wait, "high_param_dep"), 31, "high")
+    high_wait = _wait(graph, high_ag, 31, "high")
+    high_use = _neg(graph, high_wait, "high_use", device_time=1)
+    high_release = _release(graph, high_use, 31, "high")
+
+    low_pre_ag = _allgather(graph, low_pre_param, 40, "low_pre", tensor_size=1)
+    low_pre_wait = _wait(graph, low_pre_ag, 40, "low_pre")
+    low_ag = _allgather(graph, _add(graph, low_param, low_pre_wait, "low_param_dep"), 41, "low")
+    low_wait = _wait(graph, low_ag, 41, "low")
+    low_use = _neg(graph, low_wait, "low_use", device_time=100)
+    low_release = _release(graph, low_use, 41, "low")
+
+    graph.output((high_release, low_release))
+    graph.lint()
+
+    names = _scheduled_names(graph)
+
+    assert names.index(low_release.name) < names.index(high_ag.name)
+
+
+def test_fast_free_schedule_uses_pressure_tiebreaker_in_fallback_bucket():
+    graph = Graph()
+
+    high_param = _placeholder(graph, "fallback_high_param")
+    high_extra_param = _placeholder(graph, "fallback_high_extra_param")
+    low_param = _placeholder(graph, "fallback_low_param")
+    low_extra_param = _placeholder(graph, "fallback_low_extra_param")
+
+    high_ag = _allgather(graph, high_param, 50, "fallback_high", tensor_size=100)
+    high_wait = _wait(graph, high_ag, 50, "fallback_high")
+    high_extra_ag = _allgather(graph, high_extra_param, 51, "fallback_high_extra", tensor_size=10)
+    high_extra_wait = _wait(graph, high_extra_ag, 51, "fallback_high_extra")
+    high_use = _add(graph, high_wait, high_extra_wait, "fallback_high_use", device_time=1)
+    high_release = _release(graph, high_use, 50, "fallback_high")
+
+    low_ag = _allgather(graph, low_param, 60, "fallback_low", tensor_size=1)
+    low_wait = _wait(graph, low_ag, 60, "fallback_low")
+    low_extra_ag = _allgather(graph, low_extra_param, 61, "fallback_low_extra", tensor_size=10)
+    low_extra_wait = _wait(graph, low_extra_ag, 61, "fallback_low_extra")
+    low_use = _add(graph, low_wait, low_extra_wait, "fallback_low_use", device_time=100)
+    low_release = _release(graph, low_use, 60, "fallback_low")
+
+    graph.output((high_release, low_release))
+    graph.lint()
+
+    names = _scheduled_names(graph)
+
+    assert names.index(low_ag.name) < names.index(high_ag.name)
+
+
+def test_fast_free_schedule_keeps_single_allgather_release_order():
+    graph = Graph()
+
+    param = _placeholder(graph, "param")
+    ag = _allgather(graph, param, 70, "single")
+    wait = _wait(graph, ag, 70, "single")
+    use = _neg(graph, wait, "single_use")
+    release = _release(graph, use, 70, "single")
+
+    graph.output((release, ))
+    graph.lint()
+
+    names = _scheduled_names(graph)
+
+    assert names.index(ag.name) < names.index(wait.name)
+    assert names.index(wait.name) < names.index(use.name)
+    assert names.index(use.name) < names.index(release.name)
+
+
+def test_profile_backfill_makes_partial_profile_safe_for_profile_dependent_passes(monkeypatch):
+    graph = Graph()
+
+    param = _placeholder(graph, "partial_profile_param")
+    ag = _allgather(graph, param, 90, "partial_profile", device_time=None)
+    wait = _wait(graph, ag, 90, "partial_profile")
+    use = _neg(graph, wait, "partial_profile_use", device_time=None)
+    release = _release(graph, use, 90, "partial_profile")
+
+    ag.meta.pop("tensor_size", None)
+    for node in (ag, use):
+        node.meta.pop("wall_time", None)
+        node.meta.pop("alloc_mem", None)
+        node.meta.pop("max_mem", None)
+
+    graph.output((release, ))
+    graph.lint()
+
+    _backfill_missing_profile_metadata(graph)
+    assert is_profile_incomplete(graph)
+
+    for node in graph.nodes:
+        if node in (ag, use):
+            assert node.meta["device_time"] == 0.0
+        else:
+            assert "device_time" in node.meta
+        assert "wall_time" in node.meta
+        assert "tensor_size" in node.meta
+        assert "alloc_mem" in node.meta
+        assert "max_mem" in node.meta
+    assert ag.meta["tensor_size"] == 0
+
+    names = _scheduled_names(graph)
+    assert names.index(ag.name) < names.index(wait.name)
+    assert names.index(wait.name) < names.index(use.name)
+    assert names.index(use.name) < names.index(release.name)
+
+    class FakeAccelerator:
+
+        def current_device(self):
+            return "cpu"
+
+        def total_memory(self):
+            return 1024
+
+        def available_memory(self):
+            return 1024
+
+    fake_ds_param = SimpleNamespace(numel=7,
+                                    dtype=torch.float16,
+                                    param=SimpleNamespace(ds_persist=False, ds_shape=(1, )))
+    fake_param_manager = {
+        0: SimpleNamespace(params={"partial_profile_param": fake_ds_param}, ds_ids={"partial_profile_param": 90})
+    }
+    profiling_results = {
+        0: ProfilingResult(fwd_graph=graph, bwd_graph=None, fwd_mem=[("profiled_before_abort", 0, 0, 0)])
+    }
+    gm = GraphModule(torch.nn.Module(), graph)
+    logs = []
+    prefetch_logs = []
+    persisted = []
+
+    monkeypatch.setattr(prefetch_mod, "print_rank_0", lambda message: prefetch_logs.append(message))
+    assert prefetch_mod.schedule_prefetch(gm,
+                                          graph_id=0,
+                                          graph_order=[(0, True)],
+                                          profiling_results=profiling_results,
+                                          create_inputs_fn=lambda: (),
+                                          mem_budget=0,
+                                          param_manager=fake_param_manager,
+                                          bwd=False) is gm
+    assert any("incomplete profiling data" in message for message in prefetch_logs)
+
+    monkeypatch.setattr(selective_gather_mod, "print_rank_0", lambda message: logs.append(message))
+    monkeypatch.setattr(selective_gather_mod, "get_accelerator", lambda: FakeAccelerator())
+    monkeypatch.setattr(selective_gather_mod, "get_deepcompile_handle",
+                        lambda: SimpleNamespace(set_persistent=persisted.append))
+    monkeypatch.setattr(selective_gather_mod.dist, "all_reduce", lambda *args, **kwargs: None)
+
+    selective_gather_mod.selective_gather(gm,
+                                          graph_id=0,
+                                          graph_order=[(0, True)],
+                                          profiling_results=profiling_results,
+                                          create_inputs_fn=lambda: (),
+                                          mem_budget=0,
+                                          param_manager=fake_param_manager,
+                                          bwd=True)
+    assert persisted == []
+    assert any("incomplete profiling data" in message for message in logs)
+
+
+def test_schedule_prefetch_skips_when_memory_profile_incomplete(monkeypatch):
+    graph = Graph()
+
+    param = _placeholder(graph, "mem_incomplete_param")
+    ag = _allgather(graph, param, 91, "mem_incomplete")
+    wait = _wait(graph, ag, 91, "mem_incomplete")
+    use = _neg(graph, wait, "mem_incomplete_use")
+    release = _release(graph, use, 91, "mem_incomplete")
+
+    graph.output((release, ))
+    graph.lint()
+
+    profiling_results = {
+        0:
+        ProfilingResult(fwd_graph=graph,
+                        bwd_graph=None,
+                        fwd_mem=[("profiled_before_abort", 0, 0, 0)],
+                        fwd_mem_complete=False)
+    }
+    gm = GraphModule(torch.nn.Module(), graph)
+    logs = []
+
+    monkeypatch.setattr(prefetch_mod, "print_rank_0", lambda message: logs.append(message))
+
+    assert prefetch_mod.schedule_prefetch(gm,
+                                          graph_id=0,
+                                          graph_order=[(0, False)],
+                                          profiling_results=profiling_results,
+                                          create_inputs_fn=lambda: (),
+                                          mem_budget=0,
+                                          param_manager={},
+                                          bwd=False) is gm
+    assert gm.graph is graph
+    assert any("incomplete profiling data" in message for message in logs)
+
+
+def test_graphsafe_rng_state_outputs_are_registered_no_reuse():
+    graphsafe_run_with_rng_state = inductor_mod._get_graphsafe_run_with_rng_state()
+    if graphsafe_run_with_rng_state is None:
+        pytest.skip("graphsafe_run_with_rng_state is unavailable in this torch build")
+
+    calls = []
+
+    def fake_register(op_overload, **kwargs):
+        calls.append((op_overload, kwargs))
+
+    assert inductor_mod._register_graphsafe_rng_state_no_reuse(fake_register)
+    assert calls == [(graphsafe_run_with_rng_state, {"never_reuse_output": True})]
+
+
+def test_register_custom_ops_includes_graphsafe_rng_state_no_reuse(monkeypatch):
+    graphsafe_run_with_rng_state = inductor_mod._get_graphsafe_run_with_rng_state()
+    if graphsafe_run_with_rng_state is None:
+        pytest.skip("graphsafe_run_with_rng_state is unavailable in this torch build")
+
+    _define_dc_ops()
+    registered_ops = []
+
+    def fake_add_needs_realized_inputs(_op_overload):
+        return None
+
+    def fake_register_lowering(op_overload, **_kwargs):
+
+        def record_handler(handler):
+            registered_ops.append(op_overload)
+            return handler
+
+        return record_handler
+
+    monkeypatch.setattr(inductor_mod, "add_needs_realized_inputs", fake_add_needs_realized_inputs)
+    monkeypatch.setattr(inductor_mod, "register_lowering", fake_register_lowering)
+    monkeypatch.setattr(inductor_mod, "fallbacks", set())
+    monkeypatch.setattr(inductor_mod.Scheduler, "is_dc_patched", True, raising=False)
+
+    inductor_mod.register_custom_ops()
+
+    assert graphsafe_run_with_rng_state in registered_ops

tests/unit/compile/test_zero3_grad_dtype.py

@@ -0,0 +1,30 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed.compile.init_z3 import _resolve_expected_grad_dtype
+
+
+def test_missing_grad_dtype_attribute_falls_back_to_param_dtype():
+
+    class FakeParam:
+        dtype = torch.bfloat16
+
+    assert _resolve_expected_grad_dtype(FakeParam()) is torch.bfloat16
+
+
+def test_explicit_none_grad_dtype_allows_raw_grad_dtype():
+    param = torch.empty((2, 3), dtype=torch.bfloat16)
+    param.grad_dtype = None
+
+    assert _resolve_expected_grad_dtype(param) is None
+
+
+def test_explicit_grad_dtype_is_preserved():
+    param = torch.empty((2, 3), dtype=torch.bfloat16)
+    param.grad_dtype = torch.float32
+
+    assert _resolve_expected_grad_dtype(param) is torch.float32

tests/unit/compression/test_compression.py

@@ -0,0 +1,258 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import pytest
+import random
+import numpy as np
+from unit.megatron_model import get_gpt2_model
+from deepspeed.compression.compress import init_compression
+from unit.modeling import BertConfig
+from unit.modelingpreln import BertEncoder as BertEncoderPreln
+from deepspeed.compression.basic_layer import LinearLayer_Compress, ColumnParallelLinear_Compress, RowParallelLinear_Compress
+from deepspeed.compression.helper import convert_conv1d_to_linear
+from deepspeed.accelerator import get_accelerator
+from deepspeed.utils.torch import required_torch_version
+from unit.common import DistributedTest
+
+pytestmark = pytest.mark.skipif(not required_torch_version(min_version=1.5),
+                                reason='Megatron-LM package requires Pytorch version 1.5 or above')
+
+
+def reset_random(seed=1234):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    get_accelerator().manual_seed_all(seed)
+
+
+def create_bert_model():
+    hidden_size = 384
+    num_layers = 2
+    heads = 12
+    dropout_ratio = 0.1
+    bert_config = BertConfig(vocab_size_or_config_json_file=119547,
+                             hidden_size=hidden_size,
+                             num_hidden_layers=num_layers,
+                             num_attention_heads=heads,
+                             intermediate_size=hidden_size * 4,
+                             hidden_act="gelu",
+                             hidden_dropout_prob=dropout_ratio,
+                             attention_probs_dropout_prob=dropout_ratio,
+                             max_position_embeddings=512,
+                             type_vocab_size=2,
+                             initializer_range=0.2)
+
+    weights = []
+    biases = []
+
+    for i in range(4):
+        weights.append(torch.nn.Parameter(torch.Tensor(hidden_size, hidden_size)))
+
+    weights.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
+    weights.append(torch.nn.Parameter(torch.Tensor(hidden_size * 4, hidden_size)))
+    weights.append(torch.nn.Parameter(torch.Tensor(hidden_size, hidden_size * 4)))
+    weights.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
+
+    biases.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
+    for i in range(4):
+        biases.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
+    biases.append(torch.nn.Parameter(torch.Tensor(hidden_size * 4)))
+    biases.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
+    biases.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
+
+    return BertEncoderPreln(bert_config, weights, biases)
+
+
+class Conv1D(torch.nn.Module):
+    """
+    1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).
+    Basically works like a linear layer but the weights are transposed.
+    Args:
+        nf (`int`): The number of output features.
+        nx (`int`): The number of input features.
+    """
+
+    def __init__(self, nf, nx):
+        super().__init__()
+        self.nf = nf
+        w = torch.empty(nx, nf)
+        self.weight = torch.nn.Parameter(w)
+        self.bias = torch.nn.Parameter(torch.zeros(nf))
+
+    def forward(self, x):
+        size_out = x.size()[:-1] + (self.nf, )
+        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
+        x = x.view(size_out)
+        return x
+
+
+def create_conv1d_model():
+    nf = 128
+    nx = 128
+
+    return torch.nn.ModuleList([Conv1D(nf, nx) for i in range(4)])
+
+
+class TestCompression(DistributedTest):
+
+    def setup_method(self, method):
+        reset_random()
+
+    def get_ds_config(self):
+        ds_config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "fp16": {
+                "enabled": True
+            },
+            "compression_training": {
+                "weight_quantization": {
+                    "shared_parameters": {
+                        "enabled": True,
+                        "quantizer_kernel": False,
+                        "schedule_offset": 50,
+                        "quantize_groups": 1,
+                        "quantize_verbose": False,
+                        "quantization_type": "asymmetric",
+                        "rounding": "nearest",
+                        "fp16_mixed_quantize": {
+                            "enabled": False,
+                            "quantize_change_ratio": 0.001
+                        }
+                    },
+                    "different_groups": {
+                        "wq1": {
+                            "params": {
+                                "start_bits": 12,
+                                "target_bits": 8,
+                                "quantization_period": 50
+                            },
+                            "modules": ["attention.self", "intermediate"]
+                        },
+                        "wq2": {
+                            "params": {
+                                "start_bits": 12,
+                                "target_bits": 4,
+                                "quantization_period": 50
+                            },
+                            "modules": ["attention.output"]
+                        }
+                    }
+                },
+                "activation_quantization": {
+                    "shared_parameters": {
+                        "enabled": True,
+                        "quantization_type": "asymmetric",
+                        "range_calibration": "dynamic",
+                        "schedule_offset": 50
+                    },
+                    "different_groups": {
+                        "aq1": {
+                            "params": {
+                                "bits": 8
+                            },
+                            "modules": ["attention.output"]
+                        }
+                    }
+                },
+                "sparse_pruning": {
+                    "shared_parameters": {
+                        "enabled": True,
+                        "schedule_offset": 30,
+                        "method": "l1"
+                    },
+                    "different_groups": {
+                        "sp1": {
+                            "params": {
+                                "dense_ratio": 0.5
+                            },
+                            "modules": ["attention.self"]
+                        }
+                    }
+                },
+                "row_pruning": {
+                    "shared_parameters": {
+                        "enabled": True,
+                        "schedule_offset": 20,
+                        "method": "topk"
+                    },
+                    "different_groups": {
+                        "rp1": {
+                            "params": {
+                                "dense_ratio": 0.5
+                            },
+                            "modules": ["intermediate.dense"],
+                            "related_modules": [["layer.\\w+.output.dense"]]
+                        }
+                    }
+                },
+                "head_pruning": {
+                    "shared_parameters": {
+                        "enabled": True,
+                        "schedule_offset": 10,
+                        "method": "topk",
+                        "num_heads": 12
+                    },
+                    "different_groups": {
+                        "rp1": {
+                            "params": {
+                                "dense_ratio": 0.5
+                            },
+                            "modules": ["attention.output.dense"],
+                            "related_modules": [["self.query", "self.key", "self.value"]]
+                        }
+                    }
+                }
+            }
+        }
+
+        return ds_config_dict
+
+    def test_linear_layer_compress(self, tmpdir):
+        model = create_bert_model()
+        compressed_model = init_compression(model, self.get_ds_config())
+
+        assert isinstance(compressed_model.layer[0].attention.self.query, LinearLayer_Compress)
+        assert isinstance(compressed_model.layer[0].attention.self.key, LinearLayer_Compress)
+        assert isinstance(compressed_model.layer[0].attention.self.value, LinearLayer_Compress)
+
+    @pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
+    def test_mpu_compress(self, tmpdir):
+        if not required_torch_version(max_version=1.13):
+            pytest.skip("megatron not compatible with torch >1.13")
+        from megatron import mpu
+        args_defaults = {
+            'num_layers': 2,
+            'hidden_size': 128,
+            'num_attention_heads': 8,
+            'max_position_embeddings': 128,
+        }
+
+        model = get_gpt2_model(args_defaults)
+        compressed_model = init_compression(model, self.get_ds_config(), mpu=mpu)
+
+        assert isinstance(compressed_model.module.language_model.transformer.layers[0].attention.query_key_value,
+                          ColumnParallelLinear_Compress)
+        assert isinstance(compressed_model.module.language_model.transformer.layers[0].attention.dense,
+                          RowParallelLinear_Compress)
+        assert isinstance(compressed_model.module.language_model.transformer.layers[0].mlp.dense_h_to_4h,
+                          ColumnParallelLinear_Compress)
+        assert isinstance(compressed_model.module.language_model.transformer.layers[0].mlp.dense_4h_to_h,
+                          RowParallelLinear_Compress)
+
+    def test_conv1d_convertion(self, tmpdir):
+        model = create_conv1d_model()
+        compressed_model = convert_conv1d_to_linear(model, Conv1D)
+
+        assert isinstance(compressed_model[0], torch.nn.Linear)
+        assert isinstance(compressed_model[1], torch.nn.Linear)
+        assert isinstance(compressed_model[2], torch.nn.Linear)
+        assert isinstance(compressed_model[3], torch.nn.Linear)

tests/unit/compression/test_dequantization.py

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# Copyright (c) 2023, 2023, Oracle and/or its affiliates.
+
+import os
+import torch
+import pytest
+from unit.common import DistributedTest
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+
+
+class TestDequantization(DistributedTest):
+
+    def init(self):
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        self.device = torch.device(get_accelerator().device_name(local_rank))
+
+        from deepspeed.ops.op_builder import InferenceBuilder
+        if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+            pytest.skip("InferenceBuilder is not implemented")
+        else:
+            self.dequantize_func = InferenceBuilder().load().dequantize_fp16
+
+    def run_dequantize_test(self, M, N, num_groups):
+        weight = torch.randint(-255, 255, (M, N)).to(dtype=torch.int8, device=self.device)
+        scale = torch.rand(num_groups, 1).to(device=self.device)
+
+        weight_deq = (weight.reshape(num_groups, -1) * scale).reshape(M, N).to(torch.float16).contiguous()
+        weight_deq_backend = self.dequantize_func(weight, scale, num_groups)
+
+        assert torch.allclose(weight_deq, weight_deq_backend)
+
+    def test_dequantize(self):
+        self.init()
+
+        self.run_dequantize_test(14336, 7168, 32)
+        self.run_dequantize_test(14336, 1792, 32)
+        self.run_dequantize_test(768, 768, 32)
+        self.run_dequantize_test(768, 768, 48)

tests/unit/ds_batch_config.json

@@ -0,0 +1,15 @@
+{
+    "train_batch_size": 2,
+    "gradient_accumulation_steps": 1,
+    "optimizer": {
+      "type": "Adam",
+      "params": {
+        "lr": 0.00015
+      }
+    },
+    "gradient_clipping": 1.0,
+    "fp16": {
+      "enabled": true,
+      "loss_scale": 0
+    }
+  }

tests/unit/elasticity/test_elastic.py

@@ -0,0 +1,253 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import deepspeed
+from unit.common import DistributedTest
+from deepspeed.git_version_info import version as ds_version
+import os
+from unit.simple_model import SimpleModel
+from deepspeed.ops.op_builder import FusedAdamBuilder, FusedLambBuilder
+
+if not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME]:
+    pytest.skip("This op has not been implemented on this system.", allow_module_level=True)
+
+
+@pytest.fixture
+def ds_config():
+    config_dict = {
+        "elasticity": {
+            "enabled": True,
+            "max_train_batch_size": 10000,
+            "micro_batch_sizes": [8, 12, 16, 17],
+            "min_gpus": 32,
+            "max_gpus": 1500,
+            "min_time": 20,
+            "version": 0.1
+        }
+    }
+    return config_dict
+
+
+def test_basic_10k(ds_config):
+    final_batch_size, valid_gpus = deepspeed.elasticity.compute_elastic_config(ds_config=ds_config,
+                                                                               target_deepspeed_version=ds_version)
+
+    for gpu_num in valid_gpus:
+        assert final_batch_size % gpu_num == 0, f"Batch {final_batch_size} is not divisible by GPU count {gpu_num}"
+        batch_per_gpu = final_batch_size // gpu_num
+        found_valid_mbsize = False
+
+        for mb in ds_config['elasticity']['micro_batch_sizes']:
+            if batch_per_gpu % mb == 0:
+                found_valid_mb = True
+                break
+        assert found_valid_mb, "No valid mb found"
+
+    assert len(valid_gpus) == 23
+    assert final_batch_size == 9792
+
+
+def test_old_version(ds_config):
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        final_batch_size, valid_gpus = deepspeed.elasticity.compute_elastic_config(ds_config=ds_config,
+                                                                                   target_deepspeed_version="0.2")
+
+
+def test_disabled(ds_config):
+    ds_config['elasticity']['enabled'] = False
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        final_batch_size, valid_gpus = deepspeed.elasticity.compute_elastic_config(ds_config=ds_config,
+                                                                                   target_deepspeed_version=ds_version)
+
+
+def test_valid_world_size(ds_config):
+    final_batch_size, valid_gpus, mbsize = deepspeed.elasticity.compute_elastic_config(
+        ds_config=ds_config, target_deepspeed_version=ds_version, world_size=64)
+    assert mbsize == 17
+
+
+def test_invalid_world_size(ds_config):
+    with pytest.raises(deepspeed.elasticity.config.ElasticityIncompatibleWorldSize):
+        final_batch_size, valid_gpus, mbsize = deepspeed.elasticity.compute_elastic_config(
+            ds_config=ds_config, target_deepspeed_version=ds_version, world_size=128)
+
+
+def test_future_elastic_version(ds_config):
+    ds_config['elasticity']['version'] = '0.3'
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)
+
+
+def test_missing_max_batch(ds_config):
+    del ds_config['elasticity']['max_train_batch_size']
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)
+
+
+def test_missing_micro_batch(ds_config):
+    del ds_config['elasticity']['micro_batch_sizes']
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)
+
+
+def test_empty_config():
+    ds_config = {"elasticity": {"enabled": True}}
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)
+
+
+def test_model_parallel_v1_invalid(ds_config):
+    ds_config["elasticity"]["model_parallel_size"] = 4
+    ds_config["elasticity"]["num_gpus_per_node"] = 8
+    ds_config["elasticity"]["version"] = 0.1
+
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)
+
+
+def test_model_parallel_v2_invalid(ds_config):
+    ds_config["elasticity"]["model_parallel_size"] = 16
+    ds_config["elasticity"]["num_gpus_per_node"] = 8
+    ds_config["elasticity"]["version"] = 0.2
+
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        deepspeed.elasticity.compute_elastic_config(ds_config=ds_config,
+                                                    target_deepspeed_version=ds_version,
+                                                    world_size=16)
+
+
+def test_model_parallel_v2_valid(ds_config):
+    ds_config["elasticity"]["model_parallel_size"] = 4
+    ds_config["elasticity"]["num_gpus_per_node"] = 8
+    ds_config["elasticity"]["version"] = 0.2
+
+    os.environ["WORLD_SIZE"] = str(16)
+    deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)
+    os.environ.pop("WORLD_SIZE")
+
+
+@pytest.mark.parametrize('key, value', [('micro_batch_sizes', [1, 4, -1, 2, -10]), ('min_gpus', -1), ('max_gpus', -1),
+                                        ('micro_batch_sizes', 5), ('micro_batch_sizes', ['a', None, 0.5]),
+                                        ('micro_batch_sizes', [2, 0.5, 4])])
+def test_invalid_config_values(key, value, ds_config):
+    ds_config['elasticity'][key] = value
+    with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+        deepspeed.elasticity.compute_elastic_config(ds_config=ds_config, target_deepspeed_version=ds_version)
+
+
+def test_proper_mbsz(ds_config):
+    ds_config["elasticity"]["max_train_batch_size"] = 32
+    ds_config["elasticity"]["micro_batch_sizes"] = [1, 2, 3, 7]
+    ds_config["elasticity"]["min_gpus"] = 1
+    final_batch_size, valid_gpus, mbsize = deepspeed.elasticity.compute_elastic_config(
+        ds_config=ds_config, target_deepspeed_version=ds_version, world_size=7)
+    assert mbsize == 3
+
+
+class TestNonElasticBatchParams(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME],
+                        reason="FusedLambBuilder has not been implemented on this system.")
+    def test(self):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0,
+            "elasticity": {
+                "enabled": True,
+                "max_train_batch_size": 4,
+                "micro_batch_sizes": [1, 2, 3, 4],
+                "min_gpus": 1,
+                "max_gpus": 4,
+                "min_time": 20,
+                "version": 0.1
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+
+        with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+            model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+
+class TestNonElasticBatchParamsWithOverride(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME],
+                        reason="FusedLambBuilder has not been implemented on this system.")
+    def test(self):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0,
+            "elasticity": {
+                "enabled": True,
+                "max_train_batch_size": 4,
+                "micro_batch_sizes": [1, 2, 3, 4],
+                "min_gpus": 1,
+                "max_gpus": 4,
+                "min_time": 20,
+                "version": 0.1,
+                "ignore_non_elastic_batch_info": True
+            }
+        }
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+
+class TestElasticConfigChanged(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.skipif(not deepspeed.ops.__compatible_ops__[FusedLambBuilder.NAME],
+                        reason="FusedLambBuilder has not been implemented on this system.")
+    def test(self):
+        config_dict = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Lamb",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "gradient_clipping": 1.0,
+            "elasticity": {
+                "enabled": True,
+                "max_train_batch_size": 4,
+                "micro_batch_sizes": [1, 2, 3, 4],
+                "min_gpus": 1,
+                "max_gpus": 4,
+                "min_time": 20,
+                "version": 0.1,
+                "ignore_non_elastic_batch_info": True
+            }
+        }
+        import json, os
+        scheduler_elastic_config = config_dict.copy()
+        scheduler_elastic_config["elasticity"]["max_train_batch_size"] = 27
+        os.environ['DEEPSPEED_ELASTICITY_CONFIG'] = json.dumps(scheduler_elastic_config)
+        hidden_dim = 10
+
+        model = SimpleModel(hidden_dim, empty_grad=False)
+
+        with pytest.raises(deepspeed.elasticity.config.ElasticityError):
+            model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())

tests/unit/hybrid_engine/test_he_all.py

@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import torch
+import pytest
+import deepspeed
+from deepspeed.ops.op_builder import OpBuilder
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+
+from transformers import (AutoConfig, AutoTokenizer, AutoModelForCausalLM)
+from deepspeed.ops.op_builder import InferenceBuilder
+
+if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+    pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+rocm_version = OpBuilder.installed_rocm_version()
+if rocm_version != (0, 0):
+    pytest.skip("skip inference tests on rocm for now", allow_module_level=True)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize("batch_size", [1, 2], ids=["bsz=1", "bsz=2"])
+@pytest.mark.parametrize("model_name", ["EleutherAI/gpt-neo-1.3B", "facebook/opt-1.3b"])
+class TestHybridEngineTextGen(DistributedTest):
+    world_size = 1
+
+    def _generate(self, model, tokenizer, prompt):
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        tokens = tokenizer.batch_encode_plus(prompt, return_tensors="pt", padding=True)
+        for t in tokens:
+            if torch.is_tensor(tokens[t]):
+                tokens[t] = tokens[t].to(f'{get_accelerator().device_name()}:{local_rank}')
+        output = model.generate(**tokens, do_sample=False, max_length=100)
+        outputs = tokenizer.batch_decode(output, skip_special_tokens=True)
+        return outputs
+
+    def get_model(self, model_name):
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        model_config = AutoConfig.from_pretrained(model_name)
+        model_config.dropout = 0.0
+        model = AutoModelForCausalLM.from_pretrained(model_name, config=model_config)
+        model = model.half()
+        model = model.to(f'{get_accelerator().device_name()}:{local_rank}')
+        return model
+
+    def get_tokenizer(self, model_name):
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        tokenizer.pad_token = tokenizer.eos_token
+        return tokenizer
+
+    def get_prompt(self, batch_size):
+        if batch_size == 1:
+            prompt = ["Microsoft is in Washington"]
+        elif batch_size == 2:
+            prompt = ["DeepSpeed is", "Microsoft is in Washington"]
+        else:
+            raise NotImplementedError(f"batch_size {batch_size} not implemented")
+        return prompt
+
+    def test_correctness(self, batch_size, model_name):
+        pytest.skip("skip test for now, will fix in follow-up PR")
+        model = self.get_model(model_name)
+        tokenizer = self.get_tokenizer(model_name)
+        prompt = self.get_prompt(batch_size)
+
+        base_out = self._generate(model, tokenizer, prompt)
+
+        ds_config = {"train_batch_size": 1, "fp16": {"enabled": True}, "hybrid_engine": {"enabled": True}}
+        model, *_ = deepspeed.initialize(model=model, config=ds_config)
+
+        model.eval()
+        ds1_out = self._generate(model, tokenizer, prompt)
+        assert base_out == ds1_out, f"base_out: {base_out}, ds1_out: {ds1_out}"
+
+        model.train()
+        model.eval()
+        ds2_out = self._generate(model, tokenizer, prompt)
+        assert base_out == ds2_out
+
+    def test_functionality(self, batch_size, model_name):
+        model = self.get_model(model_name)
+        tokenizer = self.get_tokenizer(model_name)
+        prompt = self.get_prompt(batch_size)
+
+        ds_config = {"train_batch_size": 1, "fp16": {"enabled": True}, "hybrid_engine": {"enabled": True}}
+        model, *_ = deepspeed.initialize(model=model, config=ds_config)
+
+        model.eval()
+        ds1_out = self._generate(model, tokenizer, prompt)
+
+        model.train()
+        model.eval()
+        ds2_out = self._generate(model, tokenizer, prompt)
+
+        assert ds1_out == ds2_out, f"ds1_out: {ds1_out}, ds2_out: {ds2_out}"

tests/unit/hybrid_engine/test_he_llama.py

@@ -0,0 +1,102 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import torch
+import pytest
+import deepspeed
+from deepspeed.ops.op_builder import OpBuilder
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+
+from transformers import (AutoConfig, AutoTokenizer, AutoModelForCausalLM)
+from deepspeed.ops.op_builder import InferenceBuilder
+
+if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+    pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+rocm_version = OpBuilder.installed_rocm_version()
+if rocm_version != (0, 0):
+    pytest.skip("skip inference tests on rocm for now", allow_module_level=True)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize("batch_size", [1, 2], ids=["bsz=1", "bsz=2"])
+@pytest.mark.parametrize("model_name", ["huggyllama/llama-7b"])
+class TestHybridEngineLlama(DistributedTest):
+    world_size = 1
+
+    def _generate(self, model, tokenizer, prompt):
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        tokens = tokenizer.batch_encode_plus(prompt, return_tensors="pt", padding=True)
+        for t in tokens:
+            if torch.is_tensor(tokens[t]):
+                tokens[t] = tokens[t].to(f'{get_accelerator().device_name()}:{local_rank}')
+        #output = model.generate(**tokens, do_sample=False, max_length=100)
+        output = model.generate(tokens.input_ids, do_sample=False, max_length=100)
+        outputs = tokenizer.batch_decode(output, skip_special_tokens=True)
+        return outputs
+
+    def get_model(self, model_name):
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        model_config = AutoConfig.from_pretrained(model_name)
+        model_config.dropout = 0.0
+        model = AutoModelForCausalLM.from_pretrained(model_name, config=model_config)
+        # Make the model smaller so we can run it on a single GPU in CI
+        _ = [model.model.layers.pop(-1) for _ in range(8)]
+        model = model.half()
+        model = model.to(f'{get_accelerator().device_name()}:{local_rank}')
+        return model
+
+    def get_tokenizer(self, model_name):
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        tokenizer.pad_token = tokenizer.eos_token
+        return tokenizer
+
+    def get_prompt(self, batch_size):
+        if batch_size == 1:
+            prompt = ["Microsoft is in Washington"]
+        elif batch_size == 2:
+            prompt = ["DeepSpeed is", "Microsoft is in Washington"]
+        else:
+            raise NotImplementedError(f"batch_size {batch_size} not implemented")
+        return prompt
+
+    def test_correctness(self, batch_size, model_name):
+        pytest.skip("skip test for now, will fix in follow-up PR")
+        model = self.get_model(model_name)
+        tokenizer = self.get_tokenizer(model_name)
+        prompt = self.get_prompt(batch_size)
+
+        base_out = self._generate(model, tokenizer, prompt)
+
+        ds_config = {"train_batch_size": 1, "fp16": {"enabled": True}, "hybrid_engine": {"enabled": True}}
+        model, *_ = deepspeed.initialize(model=model, config=ds_config)
+
+        model.eval()
+        ds1_out = self._generate(model, tokenizer, prompt)
+        assert base_out == ds1_out, f"base_out: {base_out}, ds1_out: {ds1_out}"
+
+        model.train()
+        model.eval()
+        ds2_out = self._generate(model, tokenizer, prompt)
+        assert base_out == ds2_out
+
+    def test_functionality(self, batch_size, model_name):
+        model = self.get_model(model_name)
+        tokenizer = self.get_tokenizer(model_name)
+        prompt = self.get_prompt(batch_size)
+
+        ds_config = {"train_batch_size": 1, "fp16": {"enabled": True}, "hybrid_engine": {"enabled": True}}
+        model, *_ = deepspeed.initialize(model=model, config=ds_config)
+
+        model.eval()
+        ds1_out = self._generate(model, tokenizer, prompt)
+
+        model.train()
+        model.eval()
+        ds2_out = self._generate(model, tokenizer, prompt)
+
+        assert ds1_out == ds2_out, f"ds1_out: {ds1_out}, ds2_out: {ds2_out}"

tests/unit/hybrid_engine/test_he_lora.py

@@ -0,0 +1,234 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import math
+import torch
+import torch.nn.functional as F
+import pytest
+import deepspeed
+from deepspeed.runtime.zero import GatheredParameters
+from deepspeed.ops.op_builder import OpBuilder
+from deepspeed.utils import safe_get_full_grad
+import numpy.testing as npt
+from unit.common import DistributedTest
+from deepspeed.ops.op_builder import InferenceBuilder
+from deepspeed.accelerator import get_accelerator
+
+if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+    pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+from transformers import (AutoConfig, AutoTokenizer, AutoModelForCausalLM)
+
+rocm_version = OpBuilder.installed_rocm_version()
+if rocm_version != (0, 0):
+    pytest.skip("skip inference tests on rocm for now", allow_module_level=True)
+
+
+def to_device(batch, device):
+    output = {}
+    for k, v in batch.items():
+        try:
+            output[k] = v.to(device)
+        except Exception:
+            output[k] = v
+    return output
+
+
+def convert_linear_layer_to_lora(model, part_module_name, lora_dim=0, lora_scaling=1, lora_droppout=0):
+    from deepspeed.compression.helper import recursive_getattr, recursive_setattr
+
+    repalce_name = []
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Linear) and part_module_name in name:
+            repalce_name.append(name)
+    for name in repalce_name:
+        module = recursive_getattr(model, name)
+        tmp = LinearLayer_LoRA(module.weight, lora_dim, lora_scaling, lora_droppout,
+                               module.bias).to(module.weight.device).to(module.weight.dtype)
+        recursive_setattr(model, name, tmp)
+    return model
+
+
+class LinearLayer_LoRA(torch.nn.Module):
+    # an simple implementation of LoRA
+    # for now only support Linear Layer
+    def __init__(self, weight, lora_dim=0, lora_scaling=1, lora_droppout=0, bias=None):
+        super(LinearLayer_LoRA, self).__init__()
+        self.weight = weight
+        self.bias = bias
+
+        if lora_dim <= 0:
+            raise ValueError("You are training to use LoRA, whose reduced dim should be larger than 1")
+
+        try:
+            # for zero stage 3
+            rows, columns = weight.ds_shape
+        except Exception:
+            rows, columns = weight.shape
+        self.lora_right_weight = torch.nn.Parameter(torch.zeros(
+            columns, lora_dim))  # apply transpose so in forward we do not need to transpose again
+        self.lora_left_weight = torch.nn.Parameter(torch.zeros(lora_dim, rows))
+        self.lora_scaling = lora_scaling / lora_dim
+
+        if lora_droppout > 0:
+            self.lora_dropout = torch.nn.Dropout(lora_droppout)
+        else:
+            self.lora_dropout = torch.nn.Identity()
+
+        self.reset_parameters()
+        # disable the original weight gradient
+        self.weight.requires_grad = False
+        # fuse LoRA to the original weight
+        self.fuse_lora = False
+
+    def eval(self):
+        self.lora_dropout.eval()
+
+    def train(self, mode=True):
+        self.lora_dropout.train(mode)
+
+    def reset_parameters(self):
+        torch.nn.init.kaiming_uniform_(self.lora_right_weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_left_weight)
+
+    def forward(self, input):
+        if self.fuse_lora:
+            return F.linear(input, self.weight, self.bias)
+        else:
+            return F.linear(input, self.weight, self.bias) + (
+                self.lora_dropout(input) @ self.lora_right_weight @ self.lora_left_weight) * self.lora_scaling
+
+
+def only_optimize_lora_parameters(model):
+    # turn off the gradient of all the parameters except the LoRA parameters
+    for name, param in model.named_parameters():
+        if "lora_right_weight" in name or "lora_left_weight" in name:
+            param.requires_grad = True
+        else:
+            param.requires_grad = False
+    return model
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize("batch_size", [1], ids=["bsz=1"])
+@pytest.mark.parametrize("zero_stage", [2, 3], ids=["zero_stage=2", "zero_stage=3"])
+@pytest.mark.parametrize("model_name", ["EleutherAI/gpt-neo-125m", "facebook/opt-350m", "bigscience/bloom-560m"])
+@pytest.mark.parametrize("offload_device", ["none", "cpu"])
+class TestHybridEngineLoRA(DistributedTest):
+    world_size = 1
+
+    def get_model(self, model_name):
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        model_config = AutoConfig.from_pretrained(model_name)
+        model_config.dropout = 0.0
+        model = AutoModelForCausalLM.from_pretrained(model_name, config=model_config)
+        model = model.half()
+        device = get_accelerator().device_name()
+        model = model.to(f'{device}:{local_rank}')
+        return model
+
+    def get_tokenizer(self, model_name):
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        tokenizer.pad_token = tokenizer.eos_token
+        return tokenizer
+
+    def get_train_sentences(self, batch_size):
+        sentences = [
+            r"\n\nHuman: I am trying to write a fairy tale. What is the most popular plot?\n\n"
+            r"Assistant: The most popular plot might be a princess goes to a faraway land, falls in love",
+            r"\n\nHuman: What flowers should I grow to attract bees?\n\nAssistant: The reason you want bees "
+            r"in your garden is to attract pollinators and get more fruit or vegetable production."
+        ]
+        if batch_size <= 2:
+            return sentences[:batch_size]
+        else:
+            raise NotImplementedError(f"batch_size {batch_size} not implemented")
+
+    def test_lora(self, batch_size, model_name, zero_stage, offload_device):
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        model = self.get_model(model_name)
+        tokenizer = self.get_tokenizer(model_name)
+        train_sentences = self.get_train_sentences(batch_size)
+
+        # Inject LoRA
+        model = convert_linear_layer_to_lora(model, "", 8)
+        model = only_optimize_lora_parameters(model)
+
+        ds_config = {
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1.0,
+                    "betas": [0.9, 0.95]
+                }
+            },
+            "train_batch_size": batch_size,
+            "fp16": {
+                "enabled": True,
+                "initial_scale_power": 12
+            },
+            "hybrid_engine": {
+                "enabled": True,
+                "pin_parameters": True
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "offload_optimizer": {
+                    "device": offload_device
+                }
+            }
+        }
+
+        model, *_ = deepspeed.initialize(model=model, config=ds_config)
+
+        # Verify gradient norm is larger than 0
+        before_grad_update_layer0_params = [
+            ele.detach().cpu().float().numpy() for ele in model.layer_params[0]
+            if ele is not None and len(ele.shape) > 1
+        ]
+
+        model.train()
+        batch = tokenizer(train_sentences, max_length=16, padding="max_length", truncation=True, return_tensors="pt")
+        device = get_accelerator().device_name()
+        batch = to_device(batch, f'{device}:{local_rank}')
+        batch["labels"] = batch["input_ids"]
+        outputs = model(**batch, use_cache=False)
+        loss = outputs.loss
+        model.backward(loss)
+
+        grad_norm_dict = dict()
+        for name, param in model.named_parameters():
+            if param.requires_grad is True:
+                grad_norm_dict[name] = torch.linalg.norm(safe_get_full_grad(param))
+
+        model.step()
+        grad_norm = sum([ele.detach().cpu().numpy() for ele in grad_norm_dict.values()])
+        assert grad_norm > 1E-5
+
+        # Verify parameter remains the same
+        after_grad_update_layer0_params = [
+            ele.detach().cpu().float().numpy() for ele in model.layer_params[0]
+            if ele is not None and len(ele.shape) > 1
+        ]
+        for lhs, rhs in zip(before_grad_update_layer0_params, after_grad_update_layer0_params):
+            npt.assert_allclose(lhs, rhs, 1E-5, 1E-5)
+
+        # Verify fuse will mutate layer_params
+        model.eval()
+        with GatheredParameters(model.parameters()):
+            model.fuse_lora_weight()
+
+        after_grad_update_layer0_params_lora_fused = [
+            ele.detach().cpu().float().numpy() for ele in model.layer_params[0]
+            if ele is not None and len(ele.shape) > 1
+        ]
+
+        for lhs, rhs in zip(before_grad_update_layer0_params, after_grad_update_layer0_params_lora_fused):
+            with pytest.raises(AssertionError):
+                npt.assert_allclose(lhs, rhs, 1E-5, 1E-5)
+
+        with GatheredParameters(model.parameters()):
+            model.unfuse_lora_weight()

tests/unit/inference/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team