chore: import upstream snapshot with attribution

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

tests/unit/inference/quantization/test_intX_quantization.py

@@ -0,0 +1,419 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import numpy as np
+import torch
+import torch.nn as nn
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.quantization.quantization import _init_group_wise_weight_quantization
+from deepspeed.inference.quantization.utils import Quantizer, DeQuantizer
+from deepspeed.inference.quantization.layers import QuantizedLinear
+from deepspeed.utils.torch import required_torch_version
+from transformers.models.opt.modeling_opt import OPTDecoderLayer
+from transformers import AutoConfig, OPTConfig, AutoModel
+import pytest
+from collections import OrderedDict
+from typing import Dict
+from deepspeed.ops.aio import AsyncIOBuilder
+
+device = get_accelerator().device_name() if get_accelerator().is_available() else 'cpu'
+
+if not required_torch_version(min_version=1.11):
+    pytest.skip("torch.Tensor.bitwise_left_shift in INT4 quantizer needs torch 1.11 or above.",
+                allow_module_level=True)
+
+
+def reset_random(seed=1234):
+    np.random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    get_accelerator().manual_seed_all(seed)
+
+
+def quantization_test_helper(pre_quant_type: torch.dtype, num_bits: int):
+    reset_random()
+    num_group = 1024 * 32
+    group_size = 64
+    quantization_config = {'num_bits': num_bits, 'group_size': group_size, 'group_dim': 1, 'symmetric': False}
+
+    quantizer = Quantizer(config=quantization_config)
+    dequantizer = DeQuantizer(config=quantization_config, dtype=pre_quant_type)
+
+    data = torch.randn(num_group, group_size, dtype=pre_quant_type, device=device)
+
+    quantized_data, scale_buf, min_vals = quantizer.quantize(data)
+    dequantized_data = dequantizer.dequantize(quantized_data, scale_buf, min_vals)
+
+    max_diff = torch.max(torch.abs(data - dequantized_data))
+    mean_diff = torch.mean(torch.abs(data - dequantized_data))
+
+    # This threshold value is emperically selected.
+    assert mean_diff < 0.15 and max_diff < 0.5, f'Numeric error exceed threshold, mean diff {mean_diff} (threshold 0.15), max diff {max_diff} (threshold 0.5)'
+
+
+def zero3_post_init_quantization_test_helper(cpu_offload: bool, nvme_offload: bool, bits: int, nvme_path=None):
+    import deepspeed
+    from transformers.integrations.deepspeed import HfDeepSpeedConfig
+
+    if nvme_offload and not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
+        pytest.skip('Skip tests since async-io is not compatible')
+
+    def get_zero3_ds_config(hf_config: OPTConfig, cpu_offload: bool, nvme_offload: bool, bits: int) -> Dict:
+        GB = 1 << 30
+
+        ds_config = {
+            "fp16": {
+                "enabled": True,
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_prefetch_bucket_size": 2 * hf_config.hidden_size * hf_config.hidden_size,
+                "stage3_param_persistence_threshold": hf_config.hidden_size,
+                "stage3_max_live_parameters": 2 * hf_config.hidden_size * hf_config.hidden_size
+            },
+            "steps_per_print": 2000,
+            "train_micro_batch_size_per_gpu": 1,
+            "wall_clock_breakdown": False,
+            'weight_quantization': {
+                'post_init_quant': {
+                    'fc': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'self_attn.q_proj': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'self_attn.k_proj': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'self_attn.v_proj': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'self_attn.out_proj': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'lm_head': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'embed_tokens': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                }
+            }
+        }
+
+        if cpu_offload:
+            ds_config["zero_optimization"]["offload_param"] = dict(device="cpu", pin_memory=1)
+        if nvme_offload:
+            ds_config["zero_optimization"]["offload_param"] = dict(
+                device="nvme",
+                pin_memory=True,
+                nvme_path=nvme_path or '~/tmp_offload_dir',
+                buffer_count=5,
+                buffer_size=1 * GB,
+            )
+            ds_config["aio"] = {
+                "block_size": 1048576,
+                "queue_depth": 8,
+                "thread_count": 1,
+                "single_submit": False,
+                "overlap_events": True,
+            }
+
+        return ds_config
+
+    hf_config = AutoConfig.from_pretrained('facebook/opt-125m')
+    ds_config = get_zero3_ds_config(hf_config=hf_config, cpu_offload=cpu_offload, nvme_offload=nvme_offload, bits=bits)
+
+    input_ids = torch.ones(1, 16, dtype=torch.int32, device=device)
+    attention_mask = torch.ones(1, 16, dtype=torch.float32, device=device)
+
+    with torch.no_grad():
+        ref_model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16).to(device)
+        ref_model.eval()
+        ref_output = ref_model(input_ids=input_ids, attention_mask=attention_mask)
+
+    with torch.no_grad():
+        dschf = HfDeepSpeedConfig(ds_config)
+        model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16)
+        model = model.eval()
+
+        model = _init_group_wise_weight_quantization(model, ds_config)
+        ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
+        ds_engine.module.eval()
+        model = ds_engine.module
+
+    output = model(input_ids=input_ids, attention_mask=attention_mask)
+
+    mean_diff = torch.mean(torch.abs(output.last_hidden_state - ref_output.last_hidden_state))
+
+    # This threshold value is emperically selected.
+    assert mean_diff < 0.4, f'Numeric error exceed threshold, relative error {mean_diff} (threshold 0.4)'
+
+
+def zero3_quantized_initialization_test_helper(cpu_offload: bool, nvme_offload: bool, bits: int, nvme_path=None):
+    import deepspeed
+    from transformers.integrations.deepspeed import HfDeepSpeedConfig
+
+    if nvme_offload and not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
+        pytest.skip('Skip tests since async-io is not compatible')
+
+    def get_zero3_ds_config(hf_config: OPTConfig, cpu_offload: bool, nvme_offload: bool, bits: int) -> Dict:
+        GB = 1 << 30
+
+        ds_config = {
+            "fp16": {
+                "enabled": True,
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_prefetch_bucket_size": 2 * hf_config.hidden_size * hf_config.hidden_size,
+                "stage3_param_persistence_threshold": hf_config.hidden_size,
+                "stage3_max_live_parameters": 2 * hf_config.hidden_size * hf_config.hidden_size
+            },
+            "steps_per_print": 2000,
+            "train_micro_batch_size_per_gpu": 1,
+            "wall_clock_breakdown": False,
+            'weight_quantization': {
+                'quantized_initialization': {
+                    'num_bits': bits,
+                    'group_size': 32,
+                    'group_dim': 1,
+                    'symmetric': False
+                },
+            }
+        }
+
+        if cpu_offload:
+            ds_config["zero_optimization"]["offload_param"] = dict(device="cpu", pin_memory=1)
+        if nvme_offload:
+            ds_config["zero_optimization"]["offload_param"] = dict(
+                device="nvme",
+                pin_memory=True,
+                nvme_path=nvme_path or '~/tmp_offload_dir',
+                buffer_count=5,
+                buffer_size=1 * GB,
+            )
+            ds_config["aio"] = {
+                "block_size": 1048576,
+                "queue_depth": 8,
+                "thread_count": 1,
+                "single_submit": False,
+                "overlap_events": True,
+            }
+
+        return ds_config
+
+    hf_config = AutoConfig.from_pretrained('facebook/opt-125m')
+    ds_config = get_zero3_ds_config(hf_config=hf_config, cpu_offload=cpu_offload, nvme_offload=nvme_offload, bits=bits)
+
+    input_ids = torch.ones(1, 16, dtype=torch.int32, device=device)
+    attention_mask = torch.ones(1, 16, dtype=torch.float32, device=device)
+
+    with torch.no_grad():
+        ref_model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16).to(device)
+        ref_model.eval()
+        ref_output = ref_model(input_ids=input_ids, attention_mask=attention_mask)
+
+    with torch.no_grad():
+        dschf = HfDeepSpeedConfig(ds_config)
+        model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16)
+        model = model.eval()
+        ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
+        ds_engine.module.eval()
+        model = ds_engine.module
+
+    output = model(input_ids=input_ids, attention_mask=attention_mask)
+
+    mean_diff = torch.mean(torch.abs(output.last_hidden_state - ref_output.last_hidden_state))
+
+    # This threshold value is emperically selected.
+    assert mean_diff < 0.4, f'Numeric error exceed threshold, relative error {mean_diff} (threshold 0.4)'
+
+
+@pytest.fixture(params=[4, 8], ids=["4bits", "8bits"])
+def quantization_bits(request):
+    return request.param
+
+
+@pytest.fixture(params=[0, 1], ids=["0", "1"])
+def group_dim(request):
+    return request.param
+
+
+class TestQuantizedInt(DistributedTest):
+
+    def test_model_quantization(self, quantization_bits):
+        reset_random()
+
+        config = AutoConfig.from_pretrained('facebook/opt-125m')
+
+        with torch.no_grad():
+            model = OPTDecoderLayer(config).half().to(device)
+            bits = quantization_bits
+
+            ds_config = {
+                'weight_quantization': {
+                    'post_init_quant': {
+                        'fc': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        },
+                        'self_attn.q_proj': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        },
+                        'self_attn.k_proj': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        },
+                        'self_attn.v_proj': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        },
+                        'self_attn.out_proj': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        }
+                    }
+                }
+            }
+
+            model = _init_group_wise_weight_quantization(model, ds_config)
+
+            assert type(model.fc1) is QuantizedLinear
+            assert type(model.fc2) is QuantizedLinear
+            assert type(model.self_attn.q_proj) is QuantizedLinear
+            assert type(model.self_attn.k_proj) is QuantizedLinear
+            assert type(model.self_attn.v_proj) is QuantizedLinear
+            assert type(model.self_attn.out_proj) is QuantizedLinear
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_quantized_linear(self, quantization_bits, group_dim):
+        reset_random()
+
+        layers = []
+
+        for idx in range(5):
+            layers.append(
+                (f'layer_{idx}', nn.Linear(in_features=128, out_features=128, dtype=torch.float16, device=device)))
+
+        input_tensor = torch.randn(32, 128, dtype=torch.float16, device=device)
+        with torch.no_grad():
+            model = nn.Sequential(OrderedDict(layers))
+
+            ref_output = model(input_tensor)
+
+            ds_config = {
+                'weight_quantization': {
+                    'post_init_quant': {
+                        'layer': {
+                            'num_bits': quantization_bits,
+                            'group_size': 64,
+                            'group_dim': group_dim,
+                            'symmetric': False
+                        }
+                    }
+                }
+            }
+
+            model = _init_group_wise_weight_quantization(model, ds_config)
+
+            assert type(model.layer_0) is QuantizedLinear
+            assert type(model.layer_1) is QuantizedLinear
+            assert type(model.layer_2) is QuantizedLinear
+            assert type(model.layer_3) is QuantizedLinear
+            assert type(model.layer_4) is QuantizedLinear
+
+            output = model(input_tensor)
+
+            mean_diff = torch.mean(torch.abs(ref_output - output))
+
+            # This threshold value is emperically selected.
+            assert mean_diff < 0.15, f'Numeric error exceed threshold, mean diff {mean_diff}'
+
+    def test_float_int4_quantization(self):
+        reset_random()
+        quantization_test_helper(torch.float32, 4)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_half_int4_quantization(self):
+        reset_random()
+        quantization_test_helper(torch.float16, 4)
+
+    def test_float_int8_quantization(self):
+        reset_random()
+        quantization_test_helper(torch.float32, 8)
+
+    def test_half_int8_quantization(self):
+        reset_random()
+        quantization_test_helper(torch.float16, 8)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_post_init_quant(self, quantization_bits):
+        reset_random()
+        zero3_post_init_quantization_test_helper(cpu_offload=False, nvme_offload=False, bits=quantization_bits)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_post_init_quant_cpu_offload(self, quantization_bits):
+        reset_random()
+        zero3_post_init_quantization_test_helper(cpu_offload=True, nvme_offload=False, bits=quantization_bits)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_post_init_quant_nvme_offload(self, tmpdir):
+        reset_random()
+        zero3_post_init_quantization_test_helper(cpu_offload=False,
+                                                 nvme_offload=True,
+                                                 bits=4,
+                                                 nvme_path=str(tmpdir.join("nvme_offload")))
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_quantized_initialization(self, quantization_bits):
+        reset_random()
+        zero3_quantized_initialization_test_helper(cpu_offload=False, nvme_offload=False, bits=quantization_bits)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_quantized_initialization_cpu_offload(self, quantization_bits):
+        reset_random()
+        zero3_quantized_initialization_test_helper(cpu_offload=True, nvme_offload=False, bits=quantization_bits)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_quantized_initialization_nvme_offload(self, tmpdir):
+        reset_random()
+        zero3_quantized_initialization_test_helper(cpu_offload=False,
+                                                   nvme_offload=True,
+                                                   bits=4,
+                                                   nvme_path=str(tmpdir.join("nvme_offload")))

tests/unit/inference/test_checkpoint_sharding.py

@@ -0,0 +1,150 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import pytest
+import torch
+import deepspeed
+from deepspeed.model_implementations import DeepSpeedTransformerInference
+from unit.common import DistributedTest, DistributedFixture
+from transformers import AutoConfig, AutoModelForCausalLM
+import deepspeed.comm as dist
+from huggingface_hub import snapshot_download
+from deepspeed.ops.op_builder import InferenceBuilder
+
+# Handle different versions of transformers
+try:
+    from transformers.utils import is_offline_mode
+except ImportError:
+    # For transformers >= 5.0, is_offline_mode was removed
+    # transformers >= 5.0 requires huggingface_hub >= 1.2.1 which has is_offline_mode
+    from huggingface_hub import is_offline_mode
+
+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)
+
+
+def check_dtype(model, expected_dtype):
+
+    def find_dtype(module):
+        for child in module.children():
+            if isinstance(child, DeepSpeedTransformerInference):
+                return child.attention.attn_qkvw.dtype
+            else:
+                found_dtype = find_dtype(child)
+                if found_dtype:
+                    return found_dtype
+
+    found_dtype = find_dtype(model)
+    assert found_dtype, "Did not find DeepSpeedTransformerInference in model"
+    assert (found_dtype == expected_dtype), f"Expected transformer dtype {expected_dtype}, but found {found_dtype}"
+
+
+@pytest.fixture(params=[
+    "bigscience/bloom-560m", "EleutherAI/gpt-j-6B", "EleutherAI/gpt-neo-125M", "facebook/opt-350m", "facebook/opt-125m"
+])
+def model_name(request):
+    return request.param
+
+
+@pytest.fixture(params=[torch.float16, torch.int8], ids=["fp16", "int8"])
+def dtype(request):
+    if request.param not in get_accelerator().supported_dtypes():
+        pytest.skip(f"{request.param} not supported by {get_accelerator().device_name()}.")
+    return request.param
+
+
+class save_shard(DistributedFixture):
+    world_size = 2
+
+    def run(self, model_name, class_tmpdir):
+        # Only write a checkpoint if one does not exist
+        if not os.path.isdir(os.path.join(class_tmpdir, model_name)):
+            world_size = int(os.getenv("WORLD_SIZE", "1"))
+            inf_config = {
+                "replace_with_kernel_inject": True,
+                "dtype": torch.float16,
+                "enable_cuda_graph": False,
+                "tensor_parallel": {
+                    "tp_size": world_size
+                },
+                "save_mp_checkpoint_path": os.path.join(str(class_tmpdir), model_name),
+            }
+
+            # Load model and save sharded checkpoint
+            model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+            model = deepspeed.init_inference(model, config=inf_config)
+
+
+@pytest.mark.seq_inference
+class TestCheckpointShard(DistributedTest):
+    world_size = 2
+
+    def test(self, model_name, dtype, class_tmpdir, save_shard):
+        world_size = int(os.getenv("WORLD_SIZE", "1"))
+        inf_config = {
+            "replace_with_kernel_inject": True,
+            "dtype": dtype,
+            "enable_cuda_graph": False,
+            "tensor_parallel": {
+                "tp_size": world_size
+            },
+            "checkpoint": os.path.join(class_tmpdir, model_name, "ds_inference_config.json"),
+        }
+
+        # Load model on meta tensors
+        model_config = AutoConfig.from_pretrained(model_name)
+        # Note that we use half precision to load initially, even for int8
+        with deepspeed.OnDevice(dtype=torch.float16, device="meta"):
+            model = AutoModelForCausalLM.from_config(model_config, torch_dtype=torch.bfloat16)
+        model = model.eval()
+        model = deepspeed.init_inference(model, config=inf_config)
+        check_dtype(model, dtype)
+
+
+@pytest.mark.seq_inference
+class TestCheckpointShardinAutoTP(DistributedTest):
+    world_size = 2
+
+    def test(self, model_name, class_tmpdir):
+
+        def write_checkpoints_json(model_name, class_tmpdir):
+            import json
+            from pathlib import Path
+            local_rank = int(os.getenv("LOCAL_RANK", "0"))
+            if local_rank == 0:
+                # download only on first process
+                cached_repo_dir = snapshot_download(
+                    model_name,
+                    local_files_only=is_offline_mode(),
+                    cache_dir=os.getenv("HF_HOME", None),
+                    ignore_patterns=["*.safetensors", "*.msgpack", "*.h5"],
+                )
+                file_list = [str(entry) for entry in Path(cached_repo_dir).rglob("*.[bp][it][n]") if entry.is_file()]
+                data = {"type": "ds_model", "checkpoints": file_list, "version": 1.0}
+                os.makedirs(os.path.join(class_tmpdir, model_name), exist_ok=True)
+                json.dump(data, open(os.path.join(class_tmpdir, model_name, "ds_inference_config.json"), "w"))
+            dist.barrier()
+
+        world_size = int(os.getenv("WORLD_SIZE", "1"))
+        inf_config = {
+            "replace_with_kernel_inject": False,
+            "tensor_parallel": {
+                "tp_size": world_size
+            },
+            "checkpoint": os.path.join(class_tmpdir, model_name, "ds_inference_config.json"),
+        }
+
+        write_checkpoints_json(model_name, class_tmpdir)
+
+        # Load model on meta tensors
+        model_config = AutoConfig.from_pretrained(model_name)
+        # Note that we use half precision to load initially, even for int8
+        with deepspeed.OnDevice(dtype=torch.bfloat16, device="meta"):
+            model = AutoModelForCausalLM.from_config(model_config, torch_dtype=torch.bfloat16)
+        model = model.eval()
+        model = deepspeed.init_inference(model, config=inf_config)

tests/unit/inference/test_inference.py

@@ -0,0 +1,743 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+import os
+import time
+
+import deepspeed
+import torch
+
+from packaging import version as pkg_version
+from torch import nn
+from transformers import pipeline
+from transformers.models.t5.modeling_t5 import T5Block
+from transformers.models.roberta.modeling_roberta import RobertaLayer
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.git_version_info import torch_info
+from deepspeed.model_implementations import DeepSpeedTransformerInference
+from deepspeed.ops.op_builder import InferenceBuilder
+from deepspeed.ops.op_builder import OpBuilder
+
+from unit.common import DistributedTest
+
+rocm_version = OpBuilder.installed_rocm_version()
+if rocm_version != (0, 0):
+    pytest.skip("skip inference tests on rocm for now", allow_module_level=True)
+
+_bert_models = [
+    "google-bert/bert-base-cased",
+    "google-bert/bert-base-uncased",
+    "google-bert/bert-large-cased",
+    "google-bert/bert-large-uncased",
+    "google-bert/bert-base-multilingual-cased",
+    "google-bert/bert-base-multilingual-uncased",
+    "deepset/minilm-uncased-squad2",
+    "cross-encoder/ms-marco-MiniLM-L-12-v2",
+    "dslim/bert-base-NER",
+    "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad",
+    "distilbert/distilbert-base-cased-distilled-squad",
+]
+_roberta_models = [
+    "FacebookAI/roberta-large",
+    "FacebookAI/roberta-base",
+    "deepset/roberta-base-squad2",
+    "j-hartmann/emotion-english-distilroberta-base",
+    "Jean-Baptiste/roberta-large-ner-english",
+]
+_gpt_models = [
+    "openai-community/gpt2",
+    "distilbert/distilgpt2",
+    "Norod78/hebrew-bad_wiki-gpt_neo-tiny",
+    "EleutherAI/gpt-j-6b",
+    "EleutherAI/pythia-70m-deduped",
+    "bigscience/bloom-560m",
+]
+_opt_models = [
+    "facebook/opt-125m",  # 125m, 1.7B, ..., 175B variants have the same model architecture.
+    "facebook/opt-350m",  # 350m applies layer norm after attention layer which is different than other variants.
+]
+_test_models = set(_bert_models + _roberta_models + _gpt_models + _opt_models)
+_test_tasks = [
+    "fill-mask", "question-answering", "text-classification", "token-classification", "text-generation",
+    "text2text-generation", "summarization", "translation"
+]
+_test_model_list = _bert_models + _roberta_models + _gpt_models + _opt_models
+
+_model_pipeline_tags = {
+    "google-bert/bert-base-cased": "fill-mask",
+    "google-bert/bert-base-uncased": "fill-mask",
+    "google-bert/bert-large-cased": "fill-mask",
+    "google-bert/bert-large-uncased": "fill-mask",
+    "google-bert/bert-base-multilingual-cased": "fill-mask",
+    "google-bert/bert-base-multilingual-uncased": "fill-mask",
+    "deepset/minilm-uncased-squad2": "question-answering",
+    "cross-encoder/ms-marco-MiniLM-L-12-v2": "text-ranking",
+    "dslim/bert-base-NER": "token-classification",
+    "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad": "question-answering",
+    "distilbert/distilbert-base-cased-distilled-squad": "question-answering",
+    "FacebookAI/roberta-large": "fill-mask",
+    "FacebookAI/roberta-base": "fill-mask",
+    "deepset/roberta-base-squad2": "question-answering",
+    "j-hartmann/emotion-english-distilroberta-base": "text-classification",
+    "Jean-Baptiste/roberta-large-ner-english": "token-classification",
+    "openai-community/gpt2": "text-generation",
+    "distilbert/distilgpt2": "text-generation",
+    "Norod78/hebrew-bad_wiki-gpt_neo-tiny": "text-generation",
+    "EleutherAI/gpt-j-6b": "text-generation",
+    "EleutherAI/pythia-70m-deduped": "text-generation",
+    "bigscience/bloom-560m": "text-generation",
+    "facebook/opt-125m": "text-generation",
+    "facebook/opt-350m": "text-generation",
+}
+
+# Build the deterministic model/task matrix without querying Hugging Face during collection.
+_model_w_tasks = [(model, task) for model in _test_model_list for task in [_model_pipeline_tags[model]]
+                  if task in _test_tasks]
+
+# Assign to pytest variables for testing
+pytest.model_w_tasks = _model_w_tasks
+pytest.mt_names = [f"{m}-{t}" for m, t in pytest.model_w_tasks]
+
+
+@pytest.fixture(scope="module", autouse=True)
+def verify_models():
+    models_without_tags = _test_models.difference(_model_pipeline_tags)
+    if models_without_tags:
+        pytest.fail(f"Model(s) missing a test pipeline tag: {models_without_tags}")
+
+    unknown_models = set(_model_pipeline_tags).difference(_test_models)
+    if unknown_models:
+        pytest.fail(f"Model pipeline tag mapping contains unknown model(s): {unknown_models}")
+
+    missing_tags = {model for model, tag in _model_pipeline_tags.items() if not tag}
+    if missing_tags:
+        pytest.fail(f"Model pipeline tag mapping contains empty tag(s): {missing_tags}")
+
+
+""" Fixtures for inference config """
+
+
+@pytest.fixture(params=pytest.model_w_tasks, ids=pytest.mt_names)
+def model_w_task(request):
+    return request.param
+
+
+@pytest.fixture(params=[torch.float, torch.half], ids=["fp32", "fp16"])
+def dtype(request):
+    return request.param
+
+
+@pytest.fixture(params=[True, False], ids=["CG", "noCG"])
+def enable_cuda_graph(request):
+    return request.param
+
+
+@pytest.fixture(params=[True, False], ids=["Triton", "noTriton"])
+def enable_triton(request):
+    return request.param
+
+
+@pytest.fixture(params=[1, 2], ids=["ws1", "ws2"])
+def world_size(request):
+    return request.param
+
+
+""" Fixtures for running query """
+
+
+@pytest.fixture
+def query(model_w_task):
+    model, task = model_w_task
+    angle_bracket_mask_models = ["roberta", "camembert", "esm", "ibert", "luke", "mpnet", "yoso", "mpnet"]
+
+    if task == "fill-mask":
+        if any(map(lambda x: x in model, angle_bracket_mask_models)):
+            return "Hello I'm a <mask> model."
+        else:
+            return "Hell I'm a [MASK] model."
+    elif task == "question-answering":
+        return {
+            "question": "What's my name?",
+            "context": "My name is Clara and I live in Berkeley",
+        }
+    elif task == "text-classification":
+        return "DeepSpeed is the greatest"
+    elif task == "token-classification":
+        return "My name is jean-baptiste and I live in montreal."
+    elif task == "text-generation":
+        return "DeepSpeed is the greatest"
+    elif task == "text2text-generation":
+        return "Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy"
+    elif task == "translation" or task == "summarization":
+        return "Hello, my dog is cute"
+    else:
+        NotImplementedError(f'query for task "{task}" is not implemented')
+
+
+@pytest.fixture
+def inf_kwargs(model_w_task):
+    model, task = model_w_task
+    if task == "text-generation":
+        if model == "EleutherAI/gpt-j-6b":
+            # This model on V100 is hitting memory problems that limit the number of output tokens
+            return {"do_sample": False, "temperature": 1.0, "max_length": 12}
+        return {"do_sample": False, "temperature": 1.0, "max_length": 20}
+    else:
+        return {}
+
+
+""" Assertion fixture for verifying model outputs """
+
+
+def fill_mask_assert(x, y):
+    return set(res["token_str"] for res in x) == set(res["token_str"] for res in y)
+
+
+def question_answering_assert(x, y):
+    return x["answer"] == y["answer"]
+
+
+def text_classification_assert(x, y):
+    return set(res["label"] for res in x) == set(res["label"] for res in y)
+
+
+def token_classification_assert(x, y):
+    return set(ent["word"] for ent in x) == set(ent["word"] for ent in y)
+
+
+def text_generation_assert(x, y):
+    return set(res["generated_text"] for res in x) == set(res["generated_text"] for res in y)
+
+
+def text2text_generation_assert(x, y):
+    return set(res["generated_text"] for res in x) == set(res["generated_text"] for res in y)
+
+
+def translation_assert(x, y):
+    return set(res["translation_text"] for res in x) == set(res["translation_text"] for res in y)
+
+
+def summarization_assert(x, y):
+    return set(res["summary_text"] for res in x) == set(res["summary_text"] for res in y)
+
+
+@pytest.fixture
+def assert_fn(model_w_task):
+    model, task = model_w_task
+    assert_fn_dict = {
+        "fill-mask": fill_mask_assert,
+        "question-answering": question_answering_assert,
+        "text-classification": text_classification_assert,
+        "token-classification": token_classification_assert,
+        "text-generation": text_generation_assert,
+        "text2text-generation": text2text_generation_assert,
+        "translation": translation_assert,
+        "summarization": summarization_assert
+    }
+    assert_fn = assert_fn_dict.get(task, None)
+    if assert_fn is None:
+        NotImplementedError(f'assert_fn for task "{task}" is not implemented')
+    return assert_fn
+
+
+# Used to verify DeepSpeed kernel injection worked with a model
+def check_injection(model):
+
+    def verify_injection(module):
+        for child in module.children():
+            if isinstance(child, nn.ModuleList):
+                assert isinstance(child[0], DeepSpeedTransformerInference),\
+                    "DeepSpeed-Inference Transformer kernels has not been injected in the model"
+                break
+            else:
+                verify_injection(child)
+
+    verify_injection(model)
+
+
+# Used to Get Device name
+def getDeviceId(local_rank):
+    device = torch.device(f"{get_accelerator().device_name(local_rank)}")
+    return device
+
+
+# Verify that test is valid
+def validate_test(model_w_task, dtype, enable_cuda_graph, enable_triton):
+    model, task = model_w_task
+    msg = ""
+    if enable_cuda_graph and (torch_info["cuda_version"] == "0.0"):
+        msg = "CUDA not detected, cannot use CUDA Graph"
+    elif enable_cuda_graph and pkg_version.parse(torch.__version__) < pkg_version.parse("1.10"):
+        msg = "CUDA Graph is only available in torch versions >= 1.10"
+    elif "gpt-j-6b" in model:
+        if dtype != torch.half:
+            msg = f"Not enough GPU memory to run {model} with dtype {dtype}"
+        elif enable_cuda_graph:
+            msg = f"Not enough GPU memory to run {model} with CUDA Graph enabled"
+    elif "gpt-neox-20b" in model:  # TODO: remove this when neox issues resolved
+        msg = "Skipping gpt-neox-20b for now"
+    elif ("gpt-neox-20b" in model) and (dtype != torch.half):
+        msg = f"Not enough GPU memory to run {model} with dtype {dtype}"
+    elif ("bloom" in model) and (dtype != torch.half):
+        msg = f"Bloom models only support half precision, cannot use dtype {dtype}"
+    elif (model not in _bert_models + _roberta_models) and enable_cuda_graph:
+        msg = "Non bert/roberta models do no support CUDA Graph"
+    elif enable_triton and not (dtype in [torch.half]):
+        msg = "Triton is for fp16"
+    elif enable_triton and not deepspeed.HAS_TRITON:
+        msg = "triton needs to be installed for the test"
+    elif (model not in _bert_models + _roberta_models) and enable_triton:
+        msg = "Triton kernels do not support Non bert/roberta models yet"
+
+    # These should be removed once we fix several inference tests failing
+    if model in [
+            "EleutherAI/pythia-70m-deduped", "distilbert/distilbert-base-cased-distilled-squad", "EleutherAI/gpt-j-6b"
+    ]:
+        msg = "Test is currently broken"
+    return msg
+
+
+@pytest.mark.inference
+class TestModelTask(DistributedTest):
+    world_size = 1
+
+    def test(
+        self,
+        model_w_task,
+        dtype,
+        enable_cuda_graph,
+        enable_triton,
+        query,
+        inf_kwargs,
+        assert_fn,
+        perf_meas=True,
+    ):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph, enable_triton)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        if dtype not in get_accelerator().supported_dtypes():
+            pytest.skip(f"Acceleraor {get_accelerator().device_name()} does not support {dtype}.")
+
+        if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+            pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+        model, task = model_w_task
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+        # Load the model on CPU first to avoid OOM for large models @fp32
+        pipe = pipeline(task, model=model, device=torch.device("cpu"), framework="pt")
+        if dtype == torch.half:
+            pipe.model.half()
+
+        # Switch device to GPU after converting to half
+        device = torch.device(get_accelerator().device_name(local_rank))
+        pipe.device = device
+        pipe.model.to(device)
+
+        # Warm-up queries for perf measurement
+        #for i in range(10):
+        #    _ = pipe(query, **inf_kwargs)
+        get_accelerator().synchronize()
+        start = time.time()
+        bs_output = pipe(query, **inf_kwargs)
+        get_accelerator().synchronize()
+        bs_time = time.time() - start
+
+        args = {
+            'mp_size': 1,
+            'dtype': dtype,
+            'replace_with_kernel_inject': True,
+            'enable_cuda_graph': enable_cuda_graph,
+            'use_triton': enable_triton,
+            'triton_autotune': False,
+        }
+        if pipe.tokenizer.model_max_length < deepspeed.ops.transformer.inference.config.DeepSpeedInferenceConfig(
+        ).max_out_tokens:
+            args.update({'max_out_tokens': pipe.tokenizer.model_max_length})
+        pipe.model = deepspeed.init_inference(pipe.model, **args)
+        check_injection(pipe.model)
+        # Warm-up queries for perf measurement
+        #for i in range(10):
+        #    _ = pipe(query, **inf_kwargs)
+        get_accelerator().synchronize()
+        start = time.time()
+        ds_output = pipe(query, **inf_kwargs)
+        get_accelerator().synchronize()
+        ds_time = time.time() - start
+
+        if perf_meas:
+            print(
+                f"model={model}, task={task}, dtype={dtype}, cuda_graph={enable_cuda_graph}, triton={enable_triton}, bs_time={bs_time}, ds_time={ds_time}"
+            )
+
+        # facebook/opt* and some bigscient/bloom* models are not matching
+        # baseline exactly, adding an exception to them for now
+        if ("opt" in model) or ("bloom" in model):
+            bs_output = pipe(query, **inf_kwargs)
+
+        # These performance tests are only measuring the time for a single
+        # inference request, we just want to check that performance isn't terrible
+        #assert ds_time <= (bs_time * 1.1)
+
+        assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize("model_w_task", [("EleutherAI/gpt-neo-1.3B", "text-generation"),
+                                          ("EleutherAI/gpt-neox-20b", "text-generation"),
+                                          ("bigscience/bloom-3b", "text-generation"),
+                                          ("EleutherAI/gpt-j-6b", "text-generation")],
+                         ids=["gpt-neo", "gpt-neox", "bloom", "gpt-j"])
+class TestMPSize(DistributedTest):
+    world_size = 2
+
+    def test(
+        self,
+        model_w_task,
+        dtype,
+        query,
+        inf_kwargs,
+        assert_fn,
+    ):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+            pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+        model, task = model_w_task
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+        # We have to load these large models on CPU with pipeline because not
+        # enough GPU memory
+        pipe = pipeline(task, model=model, device=torch.device("cpu"), framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=self.world_size,
+                                              dtype=dtype,
+                                              replace_with_kernel_inject=True)
+        check_injection(pipe.model)
+        # Switch device to GPU so that input tensors are not on CPU
+        pipe.device = torch.device(get_accelerator().device_name(local_rank))
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.inference
+@pytest.mark.parametrize("model_w_task", [("openai-community/gpt2", "text-generation")], ids=["gpt2"])
+class TestLowCpuMemUsage(DistributedTest):
+    world_size = 1
+
+    def test(
+        self,
+        model_w_task,
+        query,
+        inf_kwargs,
+        assert_fn,
+    ):
+        model, task = model_w_task
+        dtype = torch.float16
+        if dtype not in get_accelerator().supported_dtypes():
+            pytest.skip(f"Acceleraor {get_accelerator().device_name()} does not support {dtype}.")
+
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        device = getDeviceId(local_rank)
+        pipe = pipeline(task, model=model, model_kwargs={"low_cpu_mem_usage": True}, device=device, framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=self.world_size,
+                                              dtype=dtype,
+                                              replace_method="auto",
+                                              replace_with_kernel_inject=True)
+
+        ds_output = pipe(query, **inf_kwargs)
+
+        assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize(
+    "model_w_task, injection_policy",
+    [
+        (("google/t5-v1_1-small", "text2text-generation"), {
+            T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')
+        }),
+        (("FacebookAI/roberta-large", "fill-mask"), {
+            RobertaLayer: ('output.dense')
+        }),
+    ],
+    ids=["t5", "roberta"],
+)
+@pytest.mark.parametrize("dtype", [torch.float], ids=["fp32"])
+class TestInjectionPolicy(DistributedTest):
+
+    def test(self, model_w_task, injection_policy, query, inf_kwargs, assert_fn, dtype, world_size):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        model, task = model_w_task
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+        pipe = pipeline(task,
+                        model=model,
+                        device=torch.device(get_accelerator().device_name(local_rank)),
+                        framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=world_size,
+                                              dtype=dtype,
+                                              injection_policy=injection_policy)
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize("model_w_task", [("Felladrin/Llama-160M-Chat-v1", "text-generation")], ids=["llama"])
+@pytest.mark.parametrize("dtype", [torch.half], ids=["fp16"])
+class TestLlamaInjection(DistributedTest):
+    world_size = 1
+
+    def test(self, model_w_task, dtype, query, inf_kwargs, assert_fn):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        if dtype not in get_accelerator().supported_dtypes():
+            pytest.skip(f"Accelerator {get_accelerator().device_name()} does not support {dtype}.")
+
+        if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+            pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+        model, task = model_w_task
+
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        device = torch.device(get_accelerator().device_name(local_rank))
+
+        pipe = pipeline(task,
+                        model=model,
+                        device=torch.device("cpu"),
+                        model_kwargs={"low_cpu_mem_usage": True},
+                        framework="pt")
+
+        if dtype == torch.half:
+            pipe.model.half()
+
+        pipe.device = device
+        pipe.model.to(device)
+        bs_output = pipe(query, **inf_kwargs)
+
+        try:
+            pipe.model = deepspeed.init_inference(pipe.model,
+                                                  mp_size=self.world_size,
+                                                  dtype=dtype,
+                                                  replace_with_kernel_inject=True)
+            check_injection(pipe.model)
+        except AttributeError as e:
+            if "'LlamaAttention' object has no attribute 'num_heads'" in str(e):
+                pytest.skip("Skipping due to transformers version compatibility issue with self-attention")
+            raise e
+
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        # Llama models are not matching baseline exactly
+        # We skip the result check for now, since this is irrelevant to this test
+        # assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize('keep_module_on_host', [True, False])
+@pytest.mark.parametrize(
+    "model_w_task",
+    [("Helsinki-NLP/opus-mt-en-de", "translation"), ("Salesforce/codegen-350M-mono", "text-generation")],
+    ids=["marian", "codegen"],  #codegen has fusedqkv weight.
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
+class TestAutoTensorParallelism(DistributedTest):
+    world_size = [2]
+
+    def test(
+        self,
+        model_w_task,
+        query,
+        inf_kwargs,
+        assert_fn,
+        dtype,
+        keep_module_on_host,
+    ):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        model, task = model_w_task
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        world_size = int(os.getenv("WORLD_SIZE", "2"))
+
+        if dtype not in get_accelerator().supported_dtypes():
+            pytest.skip(f"Acceleraor {get_accelerator().device_name()} does not support {dtype}.")
+
+        if model == "Salesforce/codegen-350M-mono":
+            pytest.skip("Disable Codegen model due to slight result difference")
+            #TODO: re-enable this test once we have a fix for the slight result difference
+
+        pipe = pipeline(task,
+                        model=model,
+                        device=torch.device(get_accelerator().device_name(local_rank)),
+                        framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=world_size,
+                                              dtype=dtype,
+                                              keep_module_on_host=keep_module_on_host)
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        assert assert_fn(bs_output, ds_output)
+
+        if keep_module_on_host:
+            for name, param in model.named_parameters():
+                assert param.device == torch.device('cpu'), f"keep_module_on_host is on but param {name} is not on cpu"
+
+    @pytest.mark.world_size(3)
+    def test_odd_world_size(
+        self,
+        model_w_task,
+        query,
+        inf_kwargs,
+        assert_fn,
+        dtype,
+        keep_module_on_host,
+    ):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        model, task = model_w_task
+        if model == "Salesforce/codegen-350M-mono":
+            pytest.skip("codegen does not supported by odd world_size")
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        world_size = int(os.getenv("WORLD_SIZE", "3"))
+
+        pipe = pipeline(task,
+                        model=model,
+                        device=torch.device(get_accelerator().device_name(local_rank)),
+                        framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=world_size,
+                                              dtype=dtype,
+                                              keep_module_on_host=keep_module_on_host)
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        assert assert_fn(bs_output, ds_output)
+
+        if keep_module_on_host:
+            for name, param in model.named_parameters():
+                assert param.device == torch.device('cpu'), f"keep_module_on_host is on but param {name} is not on cpu"
+
+
+@pytest.mark.nightly
+@pytest.mark.parametrize(
+    "model_family, model_name",
+    (
+        ["gpt2", "EleutherAI/gpt-neo-2.7B"],
+        #["gpt2", "EleutherAI/gpt-j-6b"], # Causing OOM for this test
+        ["gpt2", "openai-community/gpt2-xl"],
+    ),
+)
+@pytest.mark.parametrize("task", ["lambada_standard"])
+class TestLMCorrectness(DistributedTest):
+    world_size = 1
+    exec_timeout = 1200  # Give these tests longer to complete
+
+    def test(self, model_family, model_name, task):
+        # imports here to avoid import errors when pytest collects tests
+        import lm_eval
+        import lm_eval.models
+        import lm_eval.tasks
+        import lm_eval.evaluator
+
+        # The bootstrap_stderr function in lm_eval.metrics uses a
+        # multiprocessing Pool to increase performance. Since we use a Pool for
+        # our distributed tests and cannot nest Pools, we must redefine and
+        # patch this function with a version that does not use Pool.
+        def no_pool_bootstrap_stderr(f, xs, iters):
+            from lm_eval.metrics import _bootstrap_internal
+            from lm_eval.metrics import sample_stddev
+            res = []
+            chunk_size = min(1000, iters)
+            for i in range(iters // chunk_size):
+                res.extend(_bootstrap_internal(f, chunk_size)((i, xs)))
+            return sample_stddev(res)
+
+        lm_eval.metrics.bootstrap_stderr = no_pool_bootstrap_stderr
+
+        local_rank = os.getenv("LOCAL_RANK", "0")
+        device = torch.device(get_accelerator().device_name(local_rank))
+        dtype = torch.float
+        task_dict = lm_eval.tasks.get_task_dict([task])
+
+        if 'gpt-j-6b' in model_name:
+            dtype = torch.half
+            lm = lm_eval.models.get_model(model_family).create_from_arg_string(f"pretrained={model_name}",
+                                                                               {"device": "cpu"})
+            setattr(lm, model_family, getattr(lm, model_family).half().to(device))
+            lm._device = device
+        else:
+            if get_accelerator().device_name() == 'hpu':
+                #lm_eval not supporting HPU device, so get model with CPU and move it to HPU.
+                lm = lm_eval.models.get_model(model_family).create_from_arg_string(f"pretrained={model_name}",
+                                                                                   {"device": "cpu"})
+                setattr(lm, model_family, getattr(lm, model_family).to(device))
+                lm._device = device
+            else:
+                lm = lm_eval.models.get_model(model_family).create_from_arg_string(
+                    f"pretrained={model_name}", {"device": get_accelerator().device_name()})
+
+        get_accelerator().synchronize()
+        start = time.time()
+        bs_output = lm_eval.evaluator.evaluate(lm=lm, task_dict=task_dict)
+        get_accelerator().synchronize()
+        bs_time = time.time() - start
+
+        getattr(lm, model_family).to("cpu")
+        ds_model = deepspeed.init_inference(
+            getattr(lm, model_family),
+            mp_size=1,
+            dtype=dtype,
+            replace_with_kernel_inject=True,
+            enable_cuda_graph=False,
+        )
+        check_injection(ds_model)
+        setattr(lm, model_family, ds_model)
+        get_accelerator().synchronize()
+        start = time.time()
+        ds_output = lm_eval.evaluator.evaluate(lm=lm, task_dict=task_dict)
+        get_accelerator().synchronize()
+        ds_time = time.time() - start
+
+        ppl_diff = abs(bs_output["results"][task]["ppl"] - ds_output["results"][task]["ppl"])
+        #assert ds_time <= bs_time
+        assert ppl_diff < 0.01

tests/unit/inference/test_inference_config.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+import deepspeed
+from unit.common import DistributedTest
+from unit.simple_model import create_config_from_dict
+
+
+@pytest.mark.inference
+class TestInferenceConfig(DistributedTest):
+    world_size = 1
+
+    def test_overlap_kwargs(self):
+        config = {"replace_with_kernel_inject": True, "dtype": torch.float32}
+        kwargs = {"replace_with_kernel_inject": True}
+
+        engine = deepspeed.init_inference(torch.nn.Module(), config=config, **kwargs)
+        assert engine._config.replace_with_kernel_inject
+
+    def test_overlap_kwargs_conflict(self):
+        config = {"replace_with_kernel_inject": True}
+        kwargs = {"replace_with_kernel_inject": False}
+
+        with pytest.raises(ValueError):
+            engine = deepspeed.init_inference(torch.nn.Module(), config=config, **kwargs)
+
+    def test_kwargs_and_config(self):
+        config = {"replace_with_kernel_inject": True}
+        kwargs = {"dtype": torch.float32}
+
+        engine = deepspeed.init_inference(torch.nn.Module(), config=config, **kwargs)
+        assert engine._config.replace_with_kernel_inject
+        assert engine._config.dtype == kwargs["dtype"]
+
+    def test_json_config(self, tmpdir):
+        config = {"replace_with_kernel_inject": True, "dtype": "torch.float32"}
+        config_json = create_config_from_dict(tmpdir, config)
+
+        engine = deepspeed.init_inference(torch.nn.Module(), config=config_json)
+        assert engine._config.replace_with_kernel_inject

tests/unit/inference/test_model_profiling.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import time
+import pytest
+import torch
+import deepspeed
+from transformers import pipeline
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+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)
+
+if torch.half not in get_accelerator().supported_dtypes():
+    pytest.skip(f"fp16 not supported, valid dtype: {get_accelerator().supported_dtypes()}", allow_module_level=True)
+
+
+@pytest.mark.inference
+@pytest.mark.parametrize("use_cuda_events", [True, False])
+@pytest.mark.parametrize("enable_cuda_graph", [True, False])
+class TestModelProfiling(DistributedTest):
+    world_size = 1
+
+    def test(self, enable_cuda_graph, use_cuda_events):
+        task = "fill-mask"
+        model = "bert-base-cased"
+        dtype = torch.float16
+        query = "I am a [MASK] model"
+
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        world_size = int(os.getenv("WORLD_SIZE", "1"))
+
+        pipe = pipeline(task, model, framework="pt", device=get_accelerator().device_name(local_rank))
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              dtype=dtype,
+                                              mp_size=world_size,
+                                              replace_with_kernel_inject=True,
+                                              enable_cuda_graph=enable_cuda_graph)
+        pipe.model.profile_model_time(use_cuda_events=use_cuda_events)
+
+        e2e_times = []
+        model_times = []
+        for _ in range(10):
+            get_accelerator().synchronize()
+            start = time.perf_counter_ns()
+
+            r = pipe(query)
+
+            get_accelerator().synchronize()
+            end = time.perf_counter_ns()
+
+            e2e_times.append((end - start) / 1e6)  # convert ns to ms
+            model_times.extend(pipe.model.model_times())
+
+        for e2e_t, model_t in zip(e2e_times, model_times):
+            assert e2e_t >= model_t

tests/unit/inference/test_stable_diffusion.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import torch
+import pytest
+import deepspeed
+import numpy
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+
+
+# Setup for these models is different from other pipelines, so we add a separate test
+@pytest.mark.stable_diffusion
+class TestStableDiffusion(DistributedTest):
+    world_size = 1
+
+    def test(self):
+        from diffusers import DiffusionPipeline
+        from image_similarity_measures.quality_metrics import rmse
+        dev = get_accelerator().device_name()
+        generator = torch.Generator(device=dev)
+        seed = 0xABEDABE7
+        generator.manual_seed(seed)
+        prompt = "a dog on a rocket"
+        model = "prompthero/midjourney-v4-diffusion"
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        device = torch.device(f"{dev}:{local_rank}")
+        pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.half)
+        pipe = pipe.to(device)
+        baseline_image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0]
+
+        pipe = deepspeed.init_inference(
+            pipe,
+            mp_size=1,
+            dtype=torch.half,
+            replace_with_kernel_inject=True,
+            enable_cuda_graph=True,
+        )
+        generator.manual_seed(seed)
+        deepspeed_image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0]
+
+        rmse_value = rmse(org_img=numpy.asarray(baseline_image), pred_img=numpy.asarray(deepspeed_image))
+
+        # RMSE threshold value is arbitrary, may need to adjust as needed
+        assert rmse_value <= 0.01

tests/unit/inference/v2/__init__.py

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

tests/unit/inference/v2/inference_test_utils.py

@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import torch
+from deepspeed.accelerator import get_accelerator
+
+TOLERANCES = None
+
+
+def get_tolerances():
+    global TOLERANCES
+    if TOLERANCES is None:
+        TOLERANCES = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}
+        if get_accelerator().is_bf16_supported():
+            # Note: BF16 tolerance is higher than FP16 because of the lower precision (7 (+1) bits vs
+            # 10 (+1) bits)
+            TOLERANCES[torch.bfloat16] = (4.8e-1, 3.2e-2)
+    return TOLERANCES
+
+
+DTYPES = None
+
+
+def get_dtypes(include_float=True):
+    global DTYPES
+    if DTYPES is None:
+        DTYPES = [torch.float16, torch.float32] if include_float else [torch.float16]
+        try:
+            if get_accelerator().is_bf16_supported():
+                DTYPES.append(torch.bfloat16)
+        except (AssertionError, AttributeError):
+            pass
+    return DTYPES
+
+
+def allclose(x, y, tolerances: Tuple[int, int] = None):
+    assert x.dtype == y.dtype
+    if tolerances is None:
+        rtol, atol = get_tolerances()[x.dtype]
+    else:
+        rtol, atol = tolerances
+    return torch.allclose(x, y, rtol=rtol, atol=atol)

tests/unit/inference/v2/kernels/__init__.py

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

tests/unit/inference/v2/kernels/core_ops/__init__.py

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

tests/unit/inference/v2/kernels/core_ops/test_bias_activation.py

@@ -0,0 +1,101 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum
+from deepspeed.inference.v2.kernels.core_ops import CUDABiasActivation
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_bias_act_implementation(input: torch.Tensor, bias: Optional[torch.Tensor],
+                                      act_type: ActivationType) -> torch.Tensor:
+    bias_func_map = {
+        ActivationType.RELU: torch.nn.functional.relu,
+        ActivationType.GELU: torch.nn.functional.gelu,
+        ActivationType.SILU: torch.nn.functional.silu,
+        ActivationType.IDENTITY: lambda x: x,
+    }
+
+    dtype = input.dtype
+    input_f = input.to(torch.float32)
+    if bias is not None:
+        bias_f = bias.to(torch.float32)
+        output_f = input_f + bias_f
+    else:
+        output_f = input_f
+    output_f = bias_func_map[act_type](output_f)
+
+    return output_f.to(dtype)
+
+
+def _bias_activation_test_helper(tokens: int,
+                                 channels: int,
+                                 act_fn: ActivationType,
+                                 dtype: DtypeEnum,
+                                 use_bias: bool = True) -> None:
+    """
+    Fully parameterized testing entry point.
+    """
+    # Input vals
+    input_tensor = torch.randn((tokens, channels), dtype=dtype.value, device=get_accelerator().current_device_name())
+    if use_bias:
+        bias = torch.randn((channels), dtype=dtype.value, device=get_accelerator().current_device_name())
+    else:
+        bias = None
+
+    # Reference output
+    ref_output = reference_bias_act_implementation(input_tensor, bias, act_fn)
+
+    bias_act = CUDABiasActivation(channels, dtype, act_fn)
+
+    # New output
+    ds_tensor = input_tensor.clone()
+    bias_act(ds_tensor, bias)
+
+    # Check
+    assert allclose(ds_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 4096), (37, 2048), (112, 14432), (1024, 6144)])
+@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
+def test_token_channels_permutations(tokens: int, channels: int, dtype: torch.dtype) -> None:
+    """
+    Validate bias activation kernel with different token and channel permutations when using the RELU
+    activation function.
+    """
+    act_fn = ActivationType.RELU
+    dtype = DtypeEnum(dtype)
+    _bias_activation_test_helper(tokens, channels, act_fn, dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("act_fn",
+                         [ActivationType.RELU, ActivationType.GELU, ActivationType.SILU, ActivationType.IDENTITY])
+def test_act_fns(act_fn: ActivationType) -> None:
+    """
+    Validate bias activation kernel with different activation functions.
+    """
+    tokens = 223
+    channels = 4096
+    dtype = DtypeEnum.fp16
+    _bias_activation_test_helper(tokens, channels, act_fn, dtype)
+
+
+@pytest.mark.inference_v2_ops
+def test_no_bias() -> None:
+    """
+    Validate bias activation kernel with no bias.
+    """
+    tokens = 223
+    channels = 4096
+    dtype = DtypeEnum.fp16
+    act_fn = ActivationType.IDENTITY
+    _bias_activation_test_helper(tokens, channels, act_fn, dtype, use_bias=False)

tests/unit/inference/v2/kernels/core_ops/test_blas_linear.py

@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.core_ops import BlasLibLinear
+from ....v2.inference_test_utils import allclose
+
+# Note: only testing with FP16 and BF16 because we use TF32 on Ampere and we don't have a good
+# set of tolerances. Since this is just on top of BLAS though, the test is more about
+# making sure the stride/contiguity is correct and that's data type agnostic.
+
+
+def reference_implementation(hidden_states, weights):
+    return hidden_states @ weights.t()
+
+
+problem_shapes = [
+    (1, 1, 1024, 1024),
+    (1, 1024, 1024, 1024),
+    (2, 1024, 1024, 1024),
+    (1, 128, 768, 3072),
+    (1, 128, 3072, 768),
+    (1, 1024, 8192, 8192),
+    (1, 733, 8192, 32768),
+    (1, 13, 32768, 8192),
+]
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("fp_dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("problem_shape", problem_shapes)
+def test_blas_linear(fp_dtype: torch.dtype, problem_shape: Tuple[int, int, int, int]):
+    batch, seq_len, in_features, out_features = problem_shape
+    hidden_states = torch.randn(batch, seq_len, in_features, dtype=fp_dtype,
+                                device=get_accelerator().current_device()) * 0.1
+    weights = torch.randn(out_features, in_features, dtype=fp_dtype, device=get_accelerator().current_device()) * 0.01
+    ds_output = torch.empty(batch, seq_len, out_features, dtype=fp_dtype, device=get_accelerator().current_device())
+
+    ds_kernel = BlasLibLinear(fp_dtype)
+
+    ds_output = ds_kernel(ds_output, hidden_states, weights)
+    ref_output = reference_implementation(hidden_states, weights)
+
+    assert allclose(ds_output, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("fp_dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("problem_shape", problem_shapes)
+def test_blas_linear_t(fp_dtype: torch.dtype, problem_shape: Tuple[int, int, int, int]):
+    batch, seq_len, in_features, out_features = problem_shape
+    hidden_states = torch.randn(batch, seq_len, in_features, dtype=fp_dtype,
+                                device=get_accelerator().current_device()) * 0.1
+    weights = torch.randn(out_features, in_features, dtype=fp_dtype, device=get_accelerator().current_device()) * 0.01
+    ds_output = torch.empty(batch, seq_len, out_features, dtype=fp_dtype, device=get_accelerator().current_device())
+
+    ds_kernel = BlasLibLinear(fp_dtype)
+
+    # Transpose the weights then revert to the format we expect.
+    weights = weights.t().contiguous()
+    weights = weights.t()
+    ds_output = ds_kernel(ds_output, hidden_states, weights)
+
+    ref_output = reference_implementation(hidden_states, weights)
+
+    assert allclose(ds_output, ref_output)

tests/unit/inference/v2/kernels/core_ops/test_gated_activation.py

@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Iterable, Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.core_ops import CUDAGatedActivation
+from deepspeed.inference.v2.inference_utils import ActivationType
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_geglu_implementation(input: torch.Tensor,
+                                   bias: Optional[torch.Tensor] = None,
+                                   act_fn: Optional[ActivationType] = ActivationType.GEGLU) -> torch.Tensor:
+    act_func_map = {
+        ActivationType.ReGLU: torch.nn.functional.relu,
+        ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
+        ActivationType.SiGLU: torch.nn.functional.silu,
+    }
+
+    dtype = input.dtype
+    input = input.to(torch.float32)
+
+    if bias is not None:
+        bias = bias.to(torch.float32)
+        input = input + bias
+
+    act_act = input[..., ::2]
+    act_linear = input[..., 1::2]
+
+    act_act = act_func_map[act_fn](act_act)
+
+    return (act_act * act_linear).to(dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("shape", [(1372, 16384), (2, 743, 22016)])
+@pytest.mark.parametrize("dtype", get_dtypes())
+def test_dtypes(shape: Iterable[int], dtype: torch.dtype) -> None:
+    input_tensor = torch.randn(shape, dtype=dtype, device=get_accelerator().current_device_name())
+
+    # Reference output
+    ref_output = reference_geglu_implementation(input_tensor, act_fn=ActivationType.GEGLU)
+
+    # Build kernel
+    geglu = CUDAGatedActivation(input_tensor.size(-1), input_tensor.dtype, ActivationType.GEGLU)
+
+    # New output
+    output_shape = list(input_tensor.shape)
+    output_shape[-1] //= 2
+    output_tensor = torch.empty(output_shape, dtype=input_tensor.dtype, device=get_accelerator().current_device_name())
+    geglu(output_tensor, input_tensor)
+
+    # Check
+    assert allclose(output_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("act_fn", [ActivationType.GEGLU, ActivationType.ReGLU, ActivationType.SiGLU])
+def test_act_fn(act_fn: ActivationType) -> None:
+    input_tensor = torch.randn(832, 4096, dtype=torch.float16, device=get_accelerator().current_device())
+
+    # Reference output
+    ref_output = reference_geglu_implementation(input_tensor, act_fn=act_fn)
+
+    cuda_act = CUDAGatedActivation(4096, torch.float16, act_fn)
+
+    # New output
+    output_tensor = torch.empty(832, 2048, dtype=torch.float16, device=get_accelerator().current_device())
+    cuda_act(output_tensor, input_tensor)
+
+    assert allclose(output_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+def test_act_with_bias():
+    input_tensor = torch.randn(832, 4096, dtype=torch.float16, device=get_accelerator().current_device())
+    bias = torch.randn(4096, dtype=torch.float16, device=get_accelerator().current_device())
+
+    # Reference output
+    ref_output = reference_geglu_implementation(input_tensor, bias=bias, act_fn=ActivationType.GEGLU)
+
+    cuda_act = CUDAGatedActivation(4096, torch.float16, ActivationType.GEGLU)
+
+    # New output
+    output_tensor = torch.empty(832, 2048, dtype=torch.float16, device=get_accelerator().current_device())
+
+    cuda_act(output_tensor, input_tensor, bias)
+
+    assert allclose(output_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+def test_max_channels():
+    input_tensor = torch.randn(832, 48152, dtype=torch.float16, device=get_accelerator().current_device())
+
+    ref_output = reference_geglu_implementation(input_tensor, act_fn=ActivationType.GEGLU)
+
+    cuda_act = CUDAGatedActivation(48152, torch.float16, ActivationType.GEGLU)
+
+    output_tensor = torch.empty(832, 24076, dtype=torch.float16, device=get_accelerator().current_device())
+    cuda_act(output_tensor, input_tensor)
+
+    assert allclose(output_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+def test_bad_dtype() -> None:
+    with pytest.raises(ValueError):
+        CUDAGatedActivation(128, torch.int8, ActivationType.GEGLU)
+
+
+@pytest.mark.inference_v2_ops
+def test_bad_act_fn() -> None:
+    with pytest.raises(ValueError):
+        CUDAGatedActivation(128, torch.float16, ActivationType.RELU)
+
+
+@pytest.mark.inference_v2_ops
+def test_bad_alignment() -> None:
+    with pytest.raises(ValueError):
+        CUDAGatedActivation(127, torch.float16, ActivationType.GEGLU)
+
+
+@pytest.mark.inference_v2_ops
+def test_too_many_channels() -> None:
+    with pytest.raises(ValueError):
+        CUDAGatedActivation(49160, torch.float16, ActivationType.GEGLU)

tests/unit/inference/v2/kernels/core_ops/test_post_ln.py

@@ -0,0 +1,47 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.core_ops import CUDAFPPostLN
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> torch.Tensor:
+    residual_f = residual.to(torch.float32)
+    hidden_states_f = hidden_states.to(torch.float32)
+    gamma_f = gamma.to(torch.float32)
+    beta_f = beta.to(torch.float32)
+    return torch.nn.functional.layer_norm(residual_f + hidden_states_f, (hidden_states_f.size(-1), ),
+                                          weight=gamma_f,
+                                          bias=beta_f,
+                                          eps=epsilon).to(hidden_states.dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 4096), (37, 2048), (112, 14432), (1024, 6144)])
+@pytest.mark.parametrize("dtype", get_dtypes())
+def test_cuda_post_ln(tokens: int, channels: int, dtype: torch.dtype) -> None:
+
+    # Input vals
+    hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((channels), dtype=dtype, device=get_accelerator().current_device_name())
+    beta = torch.rand((channels), dtype=dtype, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    post_ln_kernel = CUDAFPPostLN(hidden_states.size(-1), residual.dtype)
+    ds_output = torch.empty_like(residual)
+    post_ln_kernel(ds_output, residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_output, ref_output)

tests/unit/inference/v2/kernels/core_ops/test_pre_ln.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.core_ops import CUDAFPPreLN
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> torch.Tensor:
+    residual_f = residual.to(torch.float32)
+    hidden_states_f = hidden_states.to(torch.float32)
+    gamma_f = gamma.to(torch.float32)
+    beta_f = beta.to(torch.float32)
+    residual_out = residual_f + hidden_states_f
+    hidden_out = torch.nn.functional.layer_norm(residual_out, (hidden_states_f.size(-1), ),
+                                                weight=gamma_f,
+                                                bias=beta_f,
+                                                eps=epsilon)
+    return residual_out.to(hidden_states.dtype), hidden_out.to(hidden_states.dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 4096), (37, 2048), (112, 14432), (1024, 6144)])
+@pytest.mark.parametrize("dtype", get_dtypes())
+def test_cuda_pre_ln(tokens: int, channels: int, dtype: torch.dtype) -> None:
+
+    # Input vals
+    hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((channels), dtype=dtype, device=get_accelerator().current_device_name())
+    beta = torch.rand((channels), dtype=dtype, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_output_res, ref_output_hid = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    pre_ln_kernel = CUDAFPPreLN(hidden_states.size(-1), residual.dtype)
+    ds_output_res = torch.empty_like(residual)
+    ds_output_hid = torch.empty_like(hidden_states)
+    pre_ln_kernel(ds_output_res, ds_output_hid, residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_output_res, ref_output_res)
+    assert allclose(ds_output_hid, ref_output_hid)

tests/unit/inference/v2/kernels/core_ops/test_rms_norm.py

@@ -0,0 +1,77 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.core_ops import CUDARMSNorm, CUDARMSPreNorm
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_rms_norm(vals: torch.Tensor, gamma: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    variance = vals.to(torch.float32).pow(2).mean(-1, keepdim=True)
+    vals = vals * torch.rsqrt(variance + epsilon)
+
+    if gamma.dtype in [torch.float16, torch.bfloat16]:
+        vals = vals.to(gamma.dtype)
+
+    return gamma * vals
+
+
+def reference_rms_pre_norm(vals: torch.Tensor,
+                           residual: torch.Tensor,
+                           gamma: torch.Tensor,
+                           epsilon: float = 1e-5) -> torch.Tensor:
+    residual = residual + vals
+    return residual, reference_rms_norm(residual, gamma, epsilon)
+
+
+def _rms_norm_testing_helper(rows: int, channels: int, do_residual: bool, dtype: DtypeEnum) -> None:
+    device = get_accelerator().current_device_name()
+    t_dtype = dtype.value
+
+    vals = torch.randn((rows, channels), dtype=t_dtype, device=device)
+    gamma = torch.randn((channels), dtype=t_dtype, device=device)
+    epsilon = 1e-5
+
+    if do_residual:
+        residual_in = torch.randn((rows, channels), dtype=t_dtype, device=device)
+        ds_residual = residual_in.clone()
+
+        ref_residual, ref_output = reference_rms_pre_norm(vals, residual_in, gamma, epsilon)
+
+        kernel = CUDARMSPreNorm(channels, t_dtype, epsilon=epsilon)
+        ds_out = torch.empty_like(ds_residual)
+
+        kernel(ds_residual, ds_out, residual_in, vals, gamma)
+
+        assert allclose(ds_out, ref_output)
+        assert allclose(ds_residual, ref_residual)
+    else:
+
+        ref_output = reference_rms_norm(vals, gamma, epsilon)
+
+        kernel = CUDARMSNorm(channels, t_dtype, epsilon=epsilon)
+        ds_out = torch.empty_like(vals)
+
+        kernel(ds_out, vals, gamma)
+
+        assert allclose(ds_out, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", get_dtypes())
+@pytest.mark.parametrize("do_residual", [True, False])
+def test_rms_dtypes(dtype: DtypeEnum, do_residual: bool) -> None:
+    _rms_norm_testing_helper(883, 1024, do_residual, DtypeEnum(dtype))
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("rows, cols", [(1, 4096), (37, 2048), (112, 14432), (1024, 6144)])
+@pytest.mark.parametrize("do_residual", [True, False])
+def test_rms_shapes(rows: int, cols: int, do_residual: bool) -> None:
+    _rms_norm_testing_helper(rows, cols, do_residual, DtypeEnum.fp16)

tests/unit/inference/v2/kernels/cutlass_ops/__init__.py

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

tests/unit/inference/v2/kernels/cutlass_ops/test_moe_gemm.py

@@ -0,0 +1,113 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum
+from deepspeed.inference.v2.kernels.cutlass_ops import MoEGEMM
+from ....v2.inference_test_utils import allclose
+
+SINGLE_EXPERT_CASES = [(13, 2048, 2048), (256, 1024, 4096), (278, 5120, 2048), (893, 5120, 2560)]
+
+PYTORCH_ACT_FN_MAP = {
+    ActivationType.GELU: torch.nn.functional.gelu,
+    ActivationType.SILU: torch.nn.functional.silu,
+    ActivationType.RELU: torch.nn.functional.relu
+}
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, in_neurons, out_neurons", SINGLE_EXPERT_CASES)
+def test_single_expert(n_tokens: int, in_neurons: int, out_neurons: int) -> None:
+    """
+    Validate that the GEMM kernel produces identical results for a single GEMM instance.
+    """
+    device = get_accelerator().current_device()
+
+    activations = torch.rand((n_tokens, in_neurons), device=device, dtype=torch.float16) - 0.5
+    weights = torch.rand((1, in_neurons, out_neurons), device=device, dtype=torch.float16) - 0.5
+    biases = torch.randn((1, out_neurons), device=device, dtype=torch.float16)
+
+    weights_ref = weights.reshape(in_neurons, out_neurons)
+    biases_ref = biases.reshape(out_neurons)
+    ref_output = torch.matmul(activations, weights_ref) + biases_ref
+
+    moe_gemm = MoEGEMM(DtypeEnum.fp16, ActivationType.IDENTITY)
+    output = torch.empty((n_tokens, out_neurons), device=device, dtype=torch.float16)
+    cumsum_rows = torch.tensor([n_tokens], dtype=torch.int64, device=device)
+
+    moe_gemm(output, activations, weights, cumsum_rows, biases)
+    assert allclose(output, ref_output, tolerances=(1e-2, 1e-2))
+    get_accelerator().synchronize()
+
+
+def moe_test_helper(in_neurons: int, out_neurons: int, n_experts: int, max_tokens_per_expert: int,
+                    act_fn: ActivationType, dtype: DtypeEnum) -> None:
+    """
+    Helper function for validating the GEMM kernel for a single expert.
+    """
+    device = get_accelerator().current_device()
+
+    expert_allocations = torch.randint(0, max_tokens_per_expert, (n_experts, ), device=device, dtype=torch.int32)
+    cumsum_rows = expert_allocations.cumsum(dim=0)
+    print(cumsum_rows.dtype)
+
+    activations = torch.rand((cumsum_rows[-1], in_neurons), device=device, dtype=dtype.value) - 0.5
+    weights = torch.rand((n_experts, in_neurons, out_neurons), device=device, dtype=dtype.value) - 0.5
+    biases = torch.randn((n_experts, out_neurons), device=device, dtype=dtype.value)
+
+    out_ref = torch.empty((cumsum_rows[-1], out_neurons), device=device, dtype=dtype.value)
+
+    for expert_idx in range(n_experts):
+        start = cumsum_rows[expert_idx - 1] if expert_idx > 0 else 0
+        end = cumsum_rows[expert_idx]
+        activations_slice = activations[start:end]
+        weights_slice = weights[expert_idx]
+        biases_slice = biases[expert_idx]
+        out_ref[start:end] = torch.matmul(activations_slice, weights_slice) + biases_slice
+
+    if act_fn != ActivationType.IDENTITY:
+        act_fn_fn = PYTORCH_ACT_FN_MAP[act_fn]
+        out_ref = act_fn_fn(out_ref)
+
+    moe_gemm = MoEGEMM(DtypeEnum.fp16, act_fn)
+    output = torch.empty((cumsum_rows[-1], out_neurons), device=device, dtype=dtype.value)
+
+    moe_gemm(output, activations, weights, cumsum_rows, biases)
+
+    if dtype == DtypeEnum.bf16:
+        assert allclose(output, out_ref, tolerances=(1e-1, 1e-1))
+    else:
+        assert allclose(output, out_ref, tolerances=(1e-2, 1e-2))
+    get_accelerator().synchronize()
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("max_tokens_per_expert", [1, 4, 16, 64, 128])
+def test_multi_expert(max_tokens_per_expert: int) -> None:
+    """
+    Validate for multi-expert GEMM instances that the output is identical to the reference.
+    """
+    moe_test_helper(5120, 2048, 64, max_tokens_per_expert, ActivationType.IDENTITY, DtypeEnum.fp16)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("act_fn", [ActivationType.GELU, ActivationType.SILU, ActivationType.RELU])
+def test_act_fns(act_fn: ActivationType) -> None:
+    """
+    Validate activation function behavior.
+    """
+    moe_test_helper(5120, 2048, 64, 32, act_fn, DtypeEnum.fp16)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", [DtypeEnum.fp16, DtypeEnum.bf16])
+def test_dtypes(dtype: DtypeEnum) -> None:
+    """
+    Validate data type behavior.
+    """
+    moe_test_helper(5120, 2048, 64, 32, ActivationType.IDENTITY, dtype)

tests/unit/inference/v2/kernels/ragged_ops/__init__.py

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

tests/unit/inference/v2/kernels/ragged_ops/ragged_testing_utils.py

@@ -0,0 +1,300 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import random
+from typing import List, Optional, Tuple
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.ragged import (
+    AllocationMode,
+    DSSequenceDescriptor,
+    DSStateManager,
+    DSStateManagerConfig,
+    KVCacheConfig,
+    MemoryConfig,
+    PlaceholderSequenceDescriptor,
+    RaggedBatchWrapper,
+)
+from ....v2.inference_test_utils import allclose
+
+
+def build_simple_batch(seq_lens: List[int],
+                       vocab_range: Optional[int] = 100,
+                       padding: Optional[bool] = False) -> RaggedBatchWrapper:
+    """
+    Construct a simple batch with the given sequence lengths. This method should not
+    be used for for testing scenarios that require information about KV or sequence
+    history.
+    """
+    total_tokens = max(sum(seq_lens), 1024)
+    n_seqs = max(len(seq_lens), 128)
+
+    config = DSStateManagerConfig(max_tracked_sequences=n_seqs,
+                                  max_ragged_sequence_count=n_seqs,
+                                  max_ragged_batch_size=total_tokens)
+    batch = RaggedBatchWrapper(config)
+
+    batch.clear()
+
+    for seq_len in seq_lens:
+        seq_desc = PlaceholderSequenceDescriptor()
+        tokens = torch.randint(0, vocab_range, (seq_len, ))
+        batch.insert_sequence(seq_desc, tokens)
+
+    batch.finalize(padding=padding)
+
+    return batch
+
+
+def build_complex_batch(seq_params: List[Tuple[int, int, int]],
+                        kv_block_size: int,
+                        vocab_range: Optional[int] = 100,
+                        padding: Optional[bool] = False) -> Tuple[RaggedBatchWrapper, int]:
+    """
+    Construct a fully paramtrized batch with the given sequence lengths. This method
+    can be used to construct more realistic inputs for testing scenarios that will interact
+    with all the members of the RaggedBatchWrapper.
+    """
+    seq_lens = [seq_param[0] for seq_param in seq_params]
+    total_tokens = max(sum(seq_lens), 1024)
+    n_seqs = max(len(seq_lens), 128)
+
+    config = DSStateManagerConfig(max_tracked_sequences=n_seqs,
+                                  max_ragged_sequence_count=n_seqs,
+                                  max_ragged_batch_size=total_tokens)
+    batch = RaggedBatchWrapper(config)
+
+    batch.clear()
+
+    total_kv_blocks = 0
+
+    for seq_len, n_seen_tokens, kv_ptr in seq_params:
+        n_kv_blocks = (seq_len + n_seen_tokens + kv_block_size - 1) // kv_block_size
+        seq_desc = PlaceholderSequenceDescriptor(seen_tokens=n_seen_tokens,
+                                                 cur_allocated_blocks=n_kv_blocks,
+                                                 kv_blocks_ptr=kv_ptr)
+        tokens = torch.randint(0, vocab_range, (seq_len, ))
+        batch.insert_sequence(seq_desc, tokens)
+        total_kv_blocks += n_kv_blocks
+
+    batch.finalize(padding=padding)
+
+    return batch, total_kv_blocks
+
+
+def build_batch_and_manager(
+    seq_params: List[Tuple[int, int]],
+    head_size: int,
+    n_heads_kv: int,
+    kv_block_size: int,
+    vocab_range: Optional[int] = 100,
+    padding: Optional[bool] = False,
+    kv_fill: Optional[List[torch.Tensor]] = None
+) -> Tuple[RaggedBatchWrapper, DSStateManager, List[DSSequenceDescriptor]]:
+    """
+    Will construct and populate a batch and KVCache with the given sequence parameters.
+
+    Arguments:
+        seq_params (List[Tuple[int, int]]): A list of tuples containing the sequence length and
+            the number of tokens that have already been seen for that sequence.
+        head_size (int): The size of each attention head.
+        n_heads_kv (int): The number of attention heads for the KV-cache.
+        kv_block_size (int): The size of each block in the KV-cache.
+        vocab_range (Optional[int]): The range of the vocabulary. Defaults to 100.
+        padding (Optional[bool]): Whether to pad the batch. Defaults to False.
+        kv_fill (Optional[List[torch.Tensor]]): A list of tensors to use to populate the KV-cache.
+            If this is not provided, the KV-cache will be treated as empty and the contents should
+            not be relied upon. NOTE(cmikeh2): This functionality relies on the functionality
+            of LinearBlockedKVCopy. If tests relying on this feature are failing, make sure that
+            LinearBlockedKVCopy is working correctly.
+    """
+    seq_lens = [seq_param[0] for seq_param in seq_params]
+    fill_lens = [seq_param[1] for seq_param in seq_params]
+    max_created_batch_len = max(sum(seq_lens), sum(fill_lens))
+    total_tokens = max(max_created_batch_len, 1024)
+    n_seqs = max(len(seq_lens), 128)
+
+    req_kv_blocks = [None] * n_seqs
+    total_kv_blocks = 0
+    for i, (seq_len, n_seen_tokens) in enumerate(seq_params):
+        req_kv_blocks[i] = (seq_len + n_seen_tokens + kv_block_size - 1) // kv_block_size
+        total_kv_blocks += req_kv_blocks[i]
+
+    kv_config = KVCacheConfig(block_size=kv_block_size,
+                              num_allocation_groups=1,
+                              cache_shape=(1, n_heads_kv, head_size))
+    memory_config = MemoryConfig(mode=AllocationMode.ALLOCATE, size=total_kv_blocks)
+
+    config = DSStateManagerConfig(max_tracked_sequences=n_seqs,
+                                  max_ragged_sequence_count=n_seqs,
+                                  max_ragged_batch_size=total_tokens,
+                                  memory_config=memory_config)
+
+    batch = RaggedBatchWrapper(config)
+    state_manager = DSStateManager(config, (kv_config, ))
+
+    # At the beginning of operation, the design of the allocator is such that it will return
+    # linear blocks of memory. The following will "warm up" the allocator so that we can be
+    # more certain that code is not dependent on this behavior.
+    all_allocs = []
+    for _ in range(20):
+        decision = random.randint(0, 1)
+
+        if decision == 0:
+            blocks_to_allocate = random.randint(0, total_kv_blocks)
+            if blocks_to_allocate <= state_manager.free_blocks[0] and blocks_to_allocate > 0:
+                all_allocs.append(state_manager.allocate_blocks(blocks_to_allocate))
+        else:
+            if len(all_allocs) > 0:
+                idx = random.randint(0, len(all_allocs) - 1)
+                state_manager._kv_cache.free(all_allocs[idx])
+
+                del all_allocs[idx]
+
+    for alloc in all_allocs:
+        state_manager._kv_cache.free(alloc)
+
+    assert state_manager.free_blocks[0] == total_kv_blocks
+
+    batch.clear()
+    seq_descs = []
+
+    if kv_fill is None or sum(fill_lens) == 0:
+        for i, (seq_len, n_seen_tokens) in enumerate(seq_params):
+            # Create empty descriptor
+            seq_desc = state_manager.get_or_create_sequence(i)
+
+            # Update `seen_tokens` in the descriptor
+            seq_desc.pre_forward(n_seen_tokens)
+            seq_desc.post_forward()
+
+            # Ensure there's enough KV-cache for the sequence
+            kv_block_ids = state_manager.allocate_blocks(req_kv_blocks[i])
+            print(f"Allocated {req_kv_blocks[i]} blocks for sequence {i}: {kv_block_ids}")
+            seq_desc.extend_kv_cache(kv_block_ids)
+
+            # Insert sequence into batch
+            tokens = torch.randint(0, vocab_range, (seq_len, ))
+            batch.insert_sequence(seq_desc, tokens)
+            seq_desc.pre_forward(seq_len)
+            seq_descs.append(seq_desc)
+    else:
+        qkv = torch.empty((total_tokens, (n_heads_kv * 3) * head_size),
+                          dtype=torch.float16,
+                          device=get_accelerator().current_device())
+        fills_as_tensor = torch.tensor(fill_lens, dtype=torch.int32)
+        fill_cumsum = torch.cat((torch.tensor([0], dtype=torch.int32), torch.cumsum(fills_as_tensor, dim=0)))
+
+        for i, (_, n_seen_tokens) in enumerate(seq_params):
+            # Create empty descriptor
+            seq_desc = state_manager.get_or_create_sequence(i)
+
+            # Update `seen_tokens` in the descriptor
+            if n_seen_tokens > 0:
+                dummy_fill_toks = torch.randint(0, vocab_range, (n_seen_tokens, ))
+                batch.insert_sequence(seq_desc, dummy_fill_toks)
+                seq_desc.pre_forward(n_seen_tokens)
+
+            # Ensure there's enough KV-cache for the sequence
+            kv_block_ids = state_manager.allocate_blocks(req_kv_blocks[i])
+            print(f"Allocated {req_kv_blocks[i]} blocks for sequence {i}: {kv_block_ids}")
+            seq_desc.extend_kv_cache(kv_block_ids)
+            seq_descs.append(seq_desc)
+
+            if n_seen_tokens == 0:
+                continue
+
+            assert kv_fill[i].shape[0] == n_seen_tokens
+            assert kv_fill[i].shape[1] == n_heads_kv * head_size * 2
+
+            local_q = torch.randn((n_seen_tokens, n_heads_kv * head_size), dtype=torch.float16, device=qkv.device)
+            local_qkv = torch.cat((local_q, kv_fill[i]), dim=1)
+            qkv[fill_cumsum[i]:fill_cumsum[i + 1]] = local_qkv
+
+        batch.finalize(padding=padding)
+
+        from deepspeed.inference.v2.kernels.ragged_ops import LinearBlockedKVCopy
+        kv_copy = LinearBlockedKVCopy(head_size, n_heads_kv, n_heads_kv, torch.float16)
+        kv_cache = state_manager.get_cache(0)
+        kv_copy(kv_cache, qkv, batch)
+
+        for seq_desc in seq_descs:
+            if seq_desc.in_flight_tokens > 0:
+                seq_desc.post_forward()
+
+        batch.clear()
+
+        for i, (seq_len, _) in enumerate(seq_params):
+            seq_desc = state_manager.get_or_create_sequence(i)
+            tokens = torch.randint(0, vocab_range, (seq_len, ))
+            batch.insert_sequence(seq_desc, tokens)
+            seq_desc.pre_forward(seq_len)
+
+            # We will skip KV cache allocation here because we did a lump allocation above
+            # for both the fill and the sequence itself.
+
+    batch.finalize(padding=padding)
+
+    return batch, state_manager, seq_descs
+
+
+def validate_kv_cache(kv_cache: torch.Tensor,
+                      k: torch.Tensor,
+                      v: torch.Tensor,
+                      seq_descs: List[DSSequenceDescriptor],
+                      batch: RaggedBatchWrapper,
+                      exact: bool = True) -> None:
+    """
+    Given a QKV tensor and a KV cache, validate that the cache contains the correct values.
+    """
+    block_size = kv_cache.shape[1]
+    n_kv_heads = kv_cache.shape[3]
+    head_size = kv_cache.shape[4]
+
+    inflight_descs = batch.inflight_seq_descriptors(on_device=False)[:batch.current_sequences]
+
+    if inflight_descs.shape[0] != len(seq_descs):
+        raise ValueError("The number of sequence descriptors does not match the number of sequences in the batch.")
+
+    for seq_desc, inflight_seq in zip(seq_descs, inflight_descs):
+        start_idx = inflight_seq[0]
+        assigned_kv_blocks = seq_desc.kv_cache_ids(on_device=False)
+
+        real_k_values = k[start_idx:start_idx + seq_desc.in_flight_tokens]
+        real_v_values = v[start_idx:start_idx + seq_desc.in_flight_tokens]
+
+        start_block_idx = seq_desc.seen_tokens // block_size
+        local_start_idx = 0
+        cur_start_idx = seq_desc.seen_tokens
+
+        for block_idx in range(start_block_idx, seq_desc.cur_allocated_blocks):
+            block = kv_cache[assigned_kv_blocks[0, block_idx].item()]
+            block_start_idx = cur_start_idx % block_size
+            n_tokens_to_check = min(block_size - block_start_idx, seq_desc.in_flight_tokens - local_start_idx)
+            block_end_idx = block_start_idx + n_tokens_to_check
+
+            if exact:
+                assert torch.equal(
+                    block[block_start_idx:block_end_idx, 0, :, :],
+                    real_k_values[local_start_idx:local_start_idx + n_tokens_to_check].reshape(
+                        n_tokens_to_check, n_kv_heads, head_size))
+                assert torch.equal(
+                    block[block_start_idx:block_end_idx, 1, :, :],
+                    real_v_values[local_start_idx:local_start_idx + n_tokens_to_check].reshape(
+                        n_tokens_to_check, n_kv_heads, head_size))
+            else:
+                assert allclose(
+                    block[block_start_idx:block_end_idx, 0, :, :],
+                    real_k_values[local_start_idx:local_start_idx + n_tokens_to_check].reshape(
+                        n_tokens_to_check, n_kv_heads, head_size))
+                assert allclose(
+                    block[block_start_idx:block_end_idx, 1, :, :],
+                    real_v_values[local_start_idx:local_start_idx + n_tokens_to_check].reshape(
+                        n_tokens_to_check, n_kv_heads, head_size))
+
+            local_start_idx += n_tokens_to_check
+            cur_start_idx += n_tokens_to_check

tests/unit/inference/v2/kernels/ragged_ops/test_atom_builder.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.kernels.ragged_ops import AtomBuilder
+from .ragged_testing_utils import build_complex_batch
+
+Q_BLOCK_SIZE = 128
+KV_BLOCK_SIZE = 128
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('seq_params', [(1, 0, 0), (1, 228, 0), (383, 0, 0), (1, 494, 0)])
+def test_single_sequence(seq_params) -> None:
+    seq_len, n_seen_tokens, _ = seq_params
+
+    batch, _ = build_complex_batch([seq_params], kv_block_size=KV_BLOCK_SIZE, padding=False)
+    atom_builder = AtomBuilder()
+
+    atoms = torch.empty((8, 8), dtype=torch.int32, device=torch.device("cpu"))
+    atoms, n_atoms = atom_builder(atoms, batch, Q_BLOCK_SIZE, KV_BLOCK_SIZE)
+
+    calc_n_atoms = (seq_len + 127) // 128
+
+    assert n_atoms == calc_n_atoms
+
+    for i, atom in enumerate(atoms[:n_atoms]):
+        # Since the ptr was 0, first 2 elements should be 0
+        assert atom[0] == 0
+        assert atom[1] == 0
+
+        # Since we have a single sequence, the q_start_idx should always be
+        # whichever atom we're on multiplied by the block size
+        assert atom[2] == i * Q_BLOCK_SIZE
+        assert atom[3] == min(Q_BLOCK_SIZE, seq_len - i * Q_BLOCK_SIZE)
+        total_toks = i * Q_BLOCK_SIZE + min(Q_BLOCK_SIZE, seq_len - i * Q_BLOCK_SIZE)
+
+        assert atom[4] == (total_toks + n_seen_tokens + KV_BLOCK_SIZE - 1) // KV_BLOCK_SIZE
+        assert atom[5] == (total_toks + n_seen_tokens)
+
+        assert atom[6] == n_seen_tokens + i * Q_BLOCK_SIZE

tests/unit/inference/v2/kernels/ragged_ops/test_blocked_flash.py

@@ -0,0 +1,197 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import itertools
+
+from typing import List, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.ragged_ops import (
+    AtomBuilder,
+    BlockedFlashAttn,
+    get_q_block_size,
+    get_kv_block_size,
+    LinearBlockedKVCopy,
+)
+from deepspeed.inference.v2.ragged import split_kv
+from deepspeed.ops.op_builder import RaggedUtilsBuilder
+
+from .ragged_testing_utils import build_batch_and_manager
+from ....v2.inference_test_utils import allclose
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+    validate_accuracy = True
+except ImportError:
+    validate_accuracy = False
+"""
+NOTE(cmikeh2): These tests depend on atom construction and KV-cache copying to behave correctly.
+If one or the other of those is not working, then these tests will fail. Before debugging here,
+make sure that the atom construction and KV-cache copying tests are passing.
+"""
+
+
+def _blocked_flash_testing_helper(head_size: int, n_heads_q: int, n_heads_kv: int,
+                                  seq_params: List[Tuple[int, int]]) -> None:
+    """
+    Helper function for testing blocked flash attention. Used to enable parametrize to only set up
+    a subset of parameters before being passed to the unified test function.
+    """
+    q_block_size = get_q_block_size(head_size)
+    kv_block_size = get_kv_block_size(head_size)
+
+    kvs = []
+    for _, history_len in seq_params:
+        if history_len > 0:
+            kvs.append(
+                torch.randn((history_len, 2 * n_heads_kv * head_size),
+                            device=get_accelerator().current_device(),
+                            dtype=torch.float16))
+        else:
+            kvs.append(None)
+
+    batch, state_manager, _ = build_batch_and_manager(seq_params, head_size, n_heads_kv, kv_block_size, kv_fill=kvs)
+
+    atom_builder = AtomBuilder()
+    kv_copy = LinearBlockedKVCopy(head_size, n_heads_q, n_heads_kv, DtypeEnum.fp16)
+    atom_flash = BlockedFlashAttn(head_size, DtypeEnum.fp16)
+
+    total_atoms = sum((seq[0] + q_block_size - 1) // q_block_size for seq in seq_params)
+    atoms = torch.empty((total_atoms, 8), dtype=torch.int32, device=get_accelerator().current_device())
+    alloc_func = RaggedUtilsBuilder().load().allocate_fast_host_buffer
+    atoms_host = alloc_func(atoms)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=get_accelerator().current_device(),
+                      dtype=torch.float16)
+
+    atoms_host, n_atoms = atom_builder(atoms_host, batch, q_block_size, kv_block_size)
+    atoms.copy_(atoms_host[:n_atoms])
+
+    kv_cache = state_manager.get_cache(0)
+    kv_copy(kv_cache, qkv, batch)
+
+    out = torch.empty((batch.current_tokens, head_size * n_heads_q),
+                      device=get_accelerator().current_device(),
+                      dtype=torch.float16)
+    k_cache, v_cache = split_kv(kv_cache)
+    q = qkv[:, :head_size * n_heads_q]
+
+    atom_flash(out, q, k_cache, v_cache, atoms, 1.0)
+
+    if validate_accuracy:
+        cu_seqlens_q = torch.tensor([0] + list(itertools.accumulate([seq[0] for seq in seq_params])),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+        cu_seqlens_kv = torch.tensor([0] + list(itertools.accumulate([seq[1] + seq[0] for seq in seq_params])),
+                                     dtype=torch.int32,
+                                     device=get_accelerator().current_device())
+
+        inflight_kv = qkv[:, head_size * n_heads_q:]
+        full_kvs = []
+        for i, kv in enumerate(kvs):
+            if kv is not None:
+                full_kvs.append(torch.cat([kv, inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]]], dim=0))
+            else:
+                full_kvs.append(inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]])
+        run_kvs = torch.cat(full_kvs, dim=0)
+        k = run_kvs[:, :head_size * n_heads_kv]
+        v = run_kvs[:, head_size * n_heads_kv:]
+
+        q_ref = q.reshape((batch.current_tokens, n_heads_q, head_size))
+        k_ref = k.reshape((k.shape[0], n_heads_kv, head_size))
+        v_ref = v.reshape((v.shape[0], n_heads_kv, head_size))
+
+        max_seqlen_q = max([seq[0] for seq in seq_params])
+        max_seqlen_kv = max([seq[1] + seq[0] for seq in seq_params])
+
+        ref_o = flash_attn_varlen_func(q_ref,
+                                       k_ref,
+                                       v_ref,
+                                       cu_seqlens_q,
+                                       cu_seqlens_kv,
+                                       max_seqlen_q,
+                                       max_seqlen_kv,
+                                       softmax_scale=1.0,
+                                       causal=True)
+
+        ref_o = ref_o.reshape(batch.current_tokens, head_size * n_heads_q)
+
+        assert allclose(out, ref_o)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens", [2, 33, 65, 128, 256, 2037])
+def test_single_prompt(n_tokens: int) -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(n_tokens, 0)]
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("prompt_lengths", [(128, 128), (192, 38), (514, 713), (83, 312, 610)])
+def test_multiple_prompts(prompt_lengths: Tuple[int, int]) -> None:
+    """
+    Test multiple prompts in a single batch.
+    """
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(prompt_lengths[i], 0) for i in range(len(prompt_lengths))]
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("seq_params", [(1, 34), (43, 40), (1, 144), (64, 128), (332, 628)])
+def test_continuation(seq_params: Tuple[int, int]) -> None:
+    """
+    Test continued generation/prompt processing.
+    """
+    head_size = 64
+    n_heads_q = 32
+    n_heads_kv = 32
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, [seq_params])
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_size", [64, 128])
+def test_head_size(head_size: int) -> None:
+    n_heads_q = 16
+    n_heads_kv = 16
+    seq_params = [(128, 128), (192, 38), (1, 814)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_config", [(32, 8), (64, 16), (40, 8)])
+def test_gqa(head_config: Tuple[int, int]) -> None:
+    head_size = 128
+    n_heads_q = head_config[0]
+    n_heads_kv = head_config[1]
+
+    seq_params = [(128, 128), (192, 38), (1, 814)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+def test_fully_composed() -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(332, 628), (1, 718), (1, 323), (180, 5), (224, 0)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)

tests/unit/inference/v2/kernels/ragged_ops/test_blocked_kv_copy.py

@@ -0,0 +1,112 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.ragged_ops import LinearBlockedKVCopy
+from .ragged_testing_utils import build_batch_and_manager, validate_kv_cache
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, history_size", [(1, 0), (17, 0), (33, 8), (63, 1)])
+@pytest.mark.parametrize("head_size", [64, 80, 96, 128])
+def test_single_sequence_single_block(n_tokens: int, history_size: int, head_size: int):
+    """
+    Validate that the copy works correctly
+    """
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch, state_manager, seq_descs = build_batch_and_manager([(n_tokens, history_size)], head_size, n_heads_kv,
+                                                              kv_block_size)
+
+    assert batch.current_sequences == 1
+    assert batch.current_tokens == n_tokens
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+
+    kv_cache = state_manager.get_cache(0)
+
+    copy_impl = LinearBlockedKVCopy(head_size, n_heads_q, n_heads_kv, torch.float16)
+    copy_impl(kv_cache, qkv, batch)
+
+    k = qkv[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv[:, head_size * (n_heads_q + n_heads_kv):]
+
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, history_size", [(128, 0), (177, 0), (169, 8), (117, 88)])
+@pytest.mark.parametrize("head_size", [64, 80, 96, 128])
+def test_single_sequence_multiple_blocks(n_tokens: int, history_size: int, head_size: int):
+    """
+    Validate that the copy works correctly
+    """
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch, state_manager, seq_descs = build_batch_and_manager([(n_tokens, history_size)], head_size, n_heads_kv,
+                                                              kv_block_size)
+
+    assert batch.current_sequences == 1
+    assert batch.current_tokens == n_tokens
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+
+    kv_cache = state_manager.get_cache(0)
+
+    copy_impl = LinearBlockedKVCopy(head_size, n_heads_q, n_heads_kv, torch.float16)
+    copy_impl(kv_cache, qkv, batch)
+
+    k = qkv[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv[:, head_size * (n_heads_q + n_heads_kv):]
+
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_size", [64, 80, 96, 128])
+def test_multi_sequence(head_size: int) -> None:
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch_config = [
+        (128, 0),
+        (177, 0),
+        (169, 8),
+        (117, 88),
+        (1, 293),
+        (1, 733),
+        (1, 33),
+    ]
+
+    batch, state_manager, seq_descs = build_batch_and_manager(batch_config, head_size, n_heads_kv, kv_block_size)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+
+    kv_cache = state_manager.get_cache(0)
+
+    copy_impl = LinearBlockedKVCopy(head_size, n_heads_q, n_heads_kv, torch.float16)
+    copy_impl(kv_cache, qkv, batch)
+
+    k = qkv[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv[:, head_size * (n_heads_q + n_heads_kv):]
+
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch)

tests/unit/inference/v2/kernels/ragged_ops/test_blocked_rotary_emb.py

@@ -0,0 +1,258 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.ragged_ops import BlockedRotaryEmbeddings, BlockedTrainedRotaryEmbeddings
+from deepspeed.inference.v2.ragged import RaggedBatchWrapper, DSSequenceDescriptor
+from .ragged_testing_utils import build_batch_and_manager, validate_kv_cache
+from ....v2.inference_test_utils import allclose
+"""
+NOTE(cmikeh2): It is very possible to see unit test failures (even on FP16) depending on when
+certain values are casted up to or down from float32. If we are seeing accuracy issues, we should
+make sure we are aligning on the training implementation's cast pattern here, given these tolerances
+tend to be sufficient elsewhere.
+"""
+
+
+def rotary_pos_embs(q: torch.Tensor,
+                    k: torch.Tensor,
+                    seq_descs: List[DSSequenceDescriptor],
+                    batch: RaggedBatchWrapper,
+                    head_size: int,
+                    rotary_dim: int = -1) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    rotary_dim = rotary_dim if rotary_dim >= 0 else head_size
+
+    def make_cos_sin_emb(seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        t = torch.arange(seq_len, dtype=torch.float32, device=get_accelerator().current_device())
+        inv_freq = (1.0 / (10000.0**(torch.arange(
+            0, rotary_dim, 2, dtype=torch.float32, device=get_accelerator().current_device()) / rotary_dim))).half()
+
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+
+        return torch.cos(emb)[:, None, :], torch.sin(emb)[:, None, :], inv_freq
+
+    def rotate_half(x: torch.Tensor) -> torch.Tensor:
+        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., :x.shape[-1] // 2]), dim=-1)
+
+    cos, sin, freqs = make_cos_sin_emb(1024)
+
+    q_out = torch.empty_like(q)
+    k_out = torch.empty_like(k)
+    n_heads_q = q.shape[1] // head_size
+    n_heads_kv = k.shape[1] // head_size
+
+    inflight_descs = batch.inflight_seq_descriptors(on_device=False)[:batch.current_sequences]
+
+    if inflight_descs.shape[0] != len(seq_descs):
+        raise ValueError("The number of sequence descriptors does not match the number of sequences in the batch.")
+
+    for seq_desc, inflight_seq in zip(seq_descs, inflight_descs):
+        start_idx = inflight_seq[0]
+        n_tokens = seq_desc.in_flight_tokens
+
+        q_src = q[start_idx:start_idx + n_tokens].reshape(n_tokens, n_heads_q, head_size).float()
+        k_src = k[start_idx:start_idx + n_tokens].reshape(n_tokens, n_heads_kv, head_size).float()
+        freq_start_offset = seq_desc.seen_tokens
+
+        q_src_rot = q_src[:, :, :rotary_dim]
+        k_src_rot = k_src[:, :, :rotary_dim]
+
+        cos_chunk = cos[range(freq_start_offset, freq_start_offset + n_tokens)]
+        sin_chunk = sin[range(freq_start_offset, freq_start_offset + n_tokens)]
+
+        q_rot = q_src_rot * cos_chunk + rotate_half(q_src_rot) * sin_chunk
+        k_rot = k_src_rot * cos_chunk + rotate_half(k_src_rot) * sin_chunk
+
+        q_emb = torch.cat((q_rot, q_src[:, :, rotary_dim:]), dim=-1)
+        k_emb = torch.cat((k_rot, k_src[:, :, rotary_dim:]), dim=-1)
+
+        q_out[start_idx:start_idx + n_tokens] = q_emb.reshape(n_tokens, n_heads_q * head_size).to(q_out.dtype)
+        k_out[start_idx:start_idx + n_tokens] = k_emb.reshape(n_tokens, n_heads_kv * head_size).to(k_out.dtype)
+
+    return q_out, k_out, freqs
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, history_size", [(1, 0), (17, 0), (33, 15), (1, 63)])
+@pytest.mark.parametrize("trained_emb", [False, True])
+@pytest.mark.parametrize("head_size", [64, 80, 96])
+def test_single_sequence_single_block(n_tokens: int, history_size: int, trained_emb: bool, head_size: int):
+    """
+    Validate that the copy works correctly
+    """
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch, state_manager, seq_descs = build_batch_and_manager([(n_tokens, history_size)], head_size, n_heads_kv,
+                                                              kv_block_size)
+
+    assert batch.current_sequences == 1
+    assert batch.current_tokens == n_tokens
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+    qkv_ref = qkv.clone()
+
+    q = qkv_ref[:, :head_size * n_heads_q]
+    k = qkv_ref[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv):]
+
+    q_ref, k, freqs = rotary_pos_embs(q, k, seq_descs, batch, head_size)
+    freqs = freqs.half()
+
+    kv_cache = state_manager.get_cache(0)
+
+    if trained_emb:
+        copy_impl = BlockedTrainedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16)
+        copy_impl(kv_cache, qkv, batch, freqs)
+    else:
+        copy_impl = BlockedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16, head_size, 10000.0)
+        copy_impl(kv_cache, qkv, batch)
+
+    assert allclose(qkv[:, :head_size * n_heads_q], q_ref)
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, history_size", [(128, 0), (177, 0), (169, 8), (117, 88)])
+@pytest.mark.parametrize("trained_emb", [False, True])
+@pytest.mark.parametrize("head_size", [64, 80, 96])
+def test_single_sequence_multiple_blocks(n_tokens: int, history_size: int, trained_emb: bool, head_size: int):
+    """
+    Validate that the copy works correctly
+    """
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch, state_manager, seq_descs = build_batch_and_manager([(n_tokens, history_size)], head_size, n_heads_kv,
+                                                              kv_block_size)
+
+    assert batch.current_sequences == 1
+    assert batch.current_tokens == n_tokens
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+    qkv_ref = qkv.clone()
+
+    q = qkv_ref[:, :head_size * n_heads_q]
+    k = qkv_ref[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv):]
+
+    q_ref, k, freqs = rotary_pos_embs(q, k, seq_descs, batch, head_size)
+    freqs = freqs.half()
+
+    kv_cache = state_manager.get_cache(0)
+
+    if trained_emb:
+        copy_impl = BlockedTrainedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16)
+        copy_impl(kv_cache, qkv, batch, freqs)
+    else:
+        copy_impl = BlockedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16, head_size, 10000.0)
+        copy_impl(kv_cache, qkv, batch)
+
+    assert allclose(qkv[:, :head_size * n_heads_q], q_ref)
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("trained_emb", [False, True])
+@pytest.mark.parametrize("head_size", [64, 80, 96])
+def test_multi_sequences(trained_emb: bool, head_size: int) -> None:
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch_config = [
+        (128, 0),
+        (177, 0),
+        (169, 8),
+        (117, 88),
+        (1, 293),
+        (1, 733),
+        (1, 33),
+    ]
+
+    batch, state_manager, seq_descs = build_batch_and_manager(batch_config, head_size, n_heads_kv, kv_block_size)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+    qkv_ref = qkv.clone()
+
+    q = qkv_ref[:, :head_size * n_heads_q]
+    k = qkv_ref[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv):]
+
+    q_ref, k, freqs = rotary_pos_embs(q, k, seq_descs, batch, head_size)
+    freqs = freqs.half()
+
+    kv_cache = state_manager.get_cache(0)
+
+    if trained_emb:
+        copy_impl = BlockedTrainedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16)
+        copy_impl(kv_cache, qkv, batch, freqs)
+    else:
+        copy_impl = BlockedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16, head_size, 10000.0)
+        copy_impl(kv_cache, qkv, batch)
+
+    assert allclose(qkv[:, :head_size * n_heads_q], q_ref)
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_size", [80, 96])
+def test_rotary_dim(head_size: int) -> None:
+    trained_emb = False
+    rotary_dim = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch_config = [
+        (128, 0),
+        (177, 0),
+        (169, 8),
+        (117, 88),
+        (1, 293),
+        (1, 733),
+        (1, 33),
+    ]
+
+    batch, state_manager, seq_descs = build_batch_and_manager(batch_config, head_size, n_heads_kv, kv_block_size)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+    qkv_ref = qkv.clone()
+
+    q = qkv_ref[:, :head_size * n_heads_q]
+    k = qkv_ref[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv):]
+
+    q_ref, k, freqs = rotary_pos_embs(q, k, seq_descs, batch, head_size, rotary_dim=rotary_dim)
+    freqs = freqs.half()
+
+    kv_cache = state_manager.get_cache(0)
+
+    copy_impl = BlockedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16, rotary_dim, 10000.0)
+    copy_impl(kv_cache, qkv, batch)
+
+    assert allclose(qkv[:, :head_size * n_heads_q], q_ref)
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)

tests/unit/inference/v2/kernels/ragged_ops/test_logits_gather.py

@@ -0,0 +1,96 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.ragged_ops import RaggedLogitsGather
+from ....v2.inference_test_utils import allclose, get_dtypes
+from .ragged_testing_utils import build_simple_batch
+
+
+def baseline_implementation(hidden_states: torch.Tensor, seq_lens: List[int]) -> torch.Tensor:
+    output = torch.empty((len(seq_lens), hidden_states.shape[1]),
+                         dtype=hidden_states.dtype,
+                         device=hidden_states.device)
+
+    offset = 0
+    for i, seq_len in enumerate(seq_lens):
+        output[i] = hidden_states[offset + seq_len - 1]
+        offset += seq_len
+
+    return output
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('dtype', get_dtypes())
+def test_supported_dtypes(dtype: torch.dtype) -> None:
+    """
+    Validate support on nominally supported data types.
+    """
+    model_dim = 4096
+
+    batch = build_simple_batch([256], padding=False)
+    hidden_states = torch.randn((batch.current_tokens, model_dim),
+                                dtype=dtype,
+                                device=get_accelerator().current_device())
+
+    reference_result = baseline_implementation(hidden_states, [256])
+
+    kernel = RaggedLogitsGather(model_dim, dtype)
+    output = torch.empty_like(reference_result)
+    kernel(output, hidden_states, batch)
+
+    assert allclose(output, reference_result)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('seq_lens', [[128, 64, 192, 32], [57, 112, 63, 89, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1],
+                                      [63, 27, 74, 83, 32, 17, 1, 1, 1, 1, 1]])
+def test_multiple_sequences(seq_lens: List[int]) -> None:
+    """
+    Validate support on more multi-sequence inputs.
+    """
+    model_dim = 4096
+    dtype = torch.float16
+
+    batch = build_simple_batch(seq_lens, padding=False)
+    hidden_states = torch.randn((batch.current_tokens, model_dim),
+                                dtype=dtype,
+                                device=get_accelerator().current_device())
+
+    reference_result = baseline_implementation(hidden_states, seq_lens)
+
+    kernel = RaggedLogitsGather(model_dim, dtype)
+    output = torch.empty_like(reference_result)
+    kernel(output, hidden_states, batch)
+
+    assert allclose(output, reference_result)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("model_dim", [1024, 6144, 6784])
+def test_problem_size_permutations(model_dim: int) -> None:
+    """
+    Validate for different embedding sizes.
+    """
+    dtype = torch.float16
+    seq_lens = [128, 64, 192, 32]
+
+    batch = build_simple_batch(seq_lens, padding=False)
+    hidden_states = torch.randn((batch.current_tokens, model_dim),
+                                dtype=dtype,
+                                device=get_accelerator().current_device())
+
+    reference_result = baseline_implementation(hidden_states, seq_lens)
+
+    kernel = RaggedLogitsGather(model_dim, dtype)
+    output = torch.empty_like(reference_result)
+    kernel(output, hidden_states, batch)
+
+    assert allclose(output, reference_result)

tests/unit/inference/v2/kernels/ragged_ops/test_moe_gather.py

@@ -0,0 +1,122 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.ragged_ops import (
+    MoEGather,
+    MoEScatter,
+    RaggedTopKGating,
+)
+from .ragged_testing_utils import build_simple_batch
+"""
+For simplicity's sake, these tests do rely on ``RaggedTopKGating``  and
+``MoEScatter`` to produce correct inputs. If either of these kernels is broken
+these tests will fail, so double check the unit test results there before
+debugging here.
+"""
+
+TEST_CASES = [
+    # (n_tokens, n_experts, n_top_k)
+    (13, 64, 1),
+    (278, 64, 1),
+    (1977, 64, 1),
+    (13, 8, 2),
+    (278, 8, 2),
+    (1977, 8, 2),
+]
+
+
+def build_inputs(n_tokens: int, n_experts: int, n_top_k: int, do_padding: bool):
+
+    assert n_tokens <= 2048, "This test will break if n_tokens > 2048"
+
+    # Sequence composition shouldn't matter here
+    batch = build_simple_batch([n_tokens], padding=do_padding)
+
+    logits = torch.randn((batch.tensor_toks, n_experts),
+                         dtype=torch.float16,
+                         device=get_accelerator().current_device())
+
+    # This will make each token's value equal to its index. NOTE: This will break for
+    # tokens with index > 2048.
+    hidden_states = torch.arange(batch.tensor_toks, dtype=torch.float16,
+                                 device=get_accelerator().current_device()).repeat_interleave(4096, dim=0).reshape(
+                                     batch.tensor_toks, 4096).contiguous()
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    # Gating outputs
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((batch.tensor_toks, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((batch.tensor_toks, n_top_k),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+    expert_offset = torch.empty((batch.tensor_toks, n_top_k),
+                                dtype=torch.int32,
+                                device=get_accelerator().current_device())
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    # Scatter outputs
+    moe_input = torch.empty((batch.tensor_toks * n_top_k, 4096),
+                            dtype=torch.float16,
+                            device=get_accelerator().current_device())
+    expert_cumsum = torch.empty((n_experts, ), dtype=torch.int64, device=get_accelerator().current_device())
+    mapped_slots = torch.empty((batch.tensor_toks, n_top_k),
+                               dtype=torch.int32,
+                               device=get_accelerator().current_device())
+
+    scatter = MoEScatter(DtypeEnum.fp16, 4096)
+    scatter(moe_input, expert_cumsum, mapped_slots, hidden_states, expert_counts, expert_assignment, expert_offset)
+
+    return batch, moe_input, scores, mapped_slots, expert_counts
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, n_experts, n_top_k", TEST_CASES)
+@pytest.mark.parametrize("do_padding", [False])
+def test_moe_gather(n_tokens: int, n_experts: int, n_top_k: int, do_padding: bool):
+    get_accelerator().manual_seed(0xC0FFEE)
+
+    batch, moe_input, scores, mapped_slots, expert_counts = build_inputs(n_tokens, n_experts, n_top_k, do_padding)
+
+    output = torch.randn((batch.tensor_toks, 4096), dtype=torch.float16, device=get_accelerator().current_device())
+
+    gather = MoEGather(DtypeEnum.fp16, 4096)
+    gather(output, moe_input, scores, mapped_slots, expert_counts)
+
+    for token_idx in range(n_tokens):
+        effective_score = scores[token_idx].sum().item()
+        assert torch.equal(
+            output[token_idx],
+            torch.full((4096, ),
+                       token_idx * effective_score,
+                       dtype=torch.float16,
+                       device=get_accelerator().current_device()))
+
+
+@pytest.mark.inference_v2_ops
+def test_moe_gather_normalize_scales():
+    get_accelerator().manual_seed(0xC0FFEE)
+
+    n_tokens = 72
+    n_experts = 8
+    n_top_k = 2
+    do_padding = False
+
+    batch, moe_input, scores, mapped_slots, expert_counts = build_inputs(n_tokens, n_experts, n_top_k, do_padding)
+    output = torch.randn((batch.tensor_toks, 4096), dtype=torch.float16, device=get_accelerator().current_device())
+
+    gather = MoEGather(DtypeEnum.fp16, 4096, normalize_scores=True)
+    gather(output, moe_input, scores, mapped_slots, expert_counts)
+
+    for token_idx in range(n_tokens):
+        assert torch.equal(
+            output[token_idx],
+            torch.full((4096, ), token_idx, dtype=torch.float16, device=get_accelerator().current_device()))

tests/unit/inference/v2/kernels/ragged_ops/test_moe_scatter.py

@@ -0,0 +1,95 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.ragged_ops import MoEScatter, RaggedTopKGating
+from .ragged_testing_utils import build_simple_batch
+"""
+For simplicity's sake, these tests do rely on ``RaggedTopKGating`` to produce correct
+inputs. If ``RaggedTopKGating`` is broken, these tests will fail, so double check
+the unit test results there before debugging here.
+"""
+
+TEST_CONFIGS = [
+    (13, 64, 1),
+    (278, 64, 1),
+    (1977, 64, 1),
+    (13, 8, 2),
+    (278, 8, 2),
+    (1977, 8, 2),
+]
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, n_experts, n_top_k", TEST_CONFIGS)
+@pytest.mark.parametrize("do_padding", [False, True])
+def test_moe_scatter(n_tokens, n_experts, n_top_k, do_padding):
+
+    # Sequence composition shouldn't matter here
+    batch = build_simple_batch([n_tokens], padding=do_padding)
+
+    logits = torch.randn((batch.tensor_toks, n_experts),
+                         dtype=torch.float16,
+                         device=get_accelerator().current_device())
+
+    # This will make each token's value equal to its index. NOTE: This will break for
+    # tokens with index > 2048.
+    hidden_states = torch.arange(batch.tensor_toks, dtype=torch.float16,
+                                 device=get_accelerator().current_device()).repeat_interleave(4096, dim=0).reshape(
+                                     batch.tensor_toks, 4096).contiguous()
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    # Gating outputs
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((batch.tensor_toks, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((batch.tensor_toks, n_top_k),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+    expert_offset = torch.empty((batch.tensor_toks, n_top_k),
+                                dtype=torch.int32,
+                                device=get_accelerator().current_device())
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    # Scatter outputs
+    moe_input = torch.empty((batch.tensor_toks * n_top_k, 4096),
+                            dtype=torch.float16,
+                            device=get_accelerator().current_device())
+    expert_cumsum = torch.empty((n_experts, ), dtype=torch.int64, device=get_accelerator().current_device())
+    mapped_slots = torch.empty((batch.tensor_toks, n_top_k),
+                               dtype=torch.int32,
+                               device=get_accelerator().current_device())
+
+    scatter = MoEScatter(DtypeEnum.fp16, 4096)
+    scatter(moe_input, expert_cumsum, mapped_slots, hidden_states, expert_counts, expert_assignment, expert_offset)
+    get_accelerator().synchronize()
+    assert torch.equal(expert_cumsum, torch.cumsum(expert_counts, dim=0).to(torch.int64))
+
+    if not do_padding:
+        assert torch.unique(mapped_slots).size(0) == n_top_k * n_tokens
+
+    for token_idx in range(batch.tensor_toks):
+        if token_idx < n_tokens:
+            for k in range(n_top_k):
+                expert_idx = expert_assignment[token_idx][k].item()
+                if expert_idx == 0:
+                    expert_cumsum_val = 0
+                else:
+                    expert_cumsum_val = expert_cumsum[expert_idx - 1]
+                offset = expert_offset[token_idx][k]
+                total_offset = offset + expert_cumsum_val
+
+                assert total_offset == mapped_slots[token_idx][k].item()
+                assert torch.equal(moe_input[total_offset], hidden_states[token_idx])
+        else:
+            for k in range(n_top_k):
+                assert mapped_slots[token_idx][k].item() == -1
+
+    assert expert_cumsum[-1] == n_tokens * n_top_k

tests/unit/inference/v2/kernels/ragged_ops/test_ragged_embed.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List, Optional, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.ragged_ops import RaggedEmbeddingKernel
+from ....v2.inference_test_utils import allclose, get_dtypes
+from .ragged_testing_utils import build_batch_and_manager
+
+
+def baseline_implementation(token_ids: torch.Tensor,
+                            embedding_table: torch.Tensor,
+                            unpadded_size: int,
+                            positional_embedding_table: Optional[torch.Tensor] = None,
+                            positional_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """
+    Baseline implementation for our ragged embedding kernel.
+    """
+    if unpadded_size == token_ids.shape[0]:
+        token_embed = torch.nn.functional.embedding(token_ids, embedding_table)
+
+        if positional_embedding_table is not None:
+            pos_embed = torch.nn.functional.embedding(positional_ids, positional_embedding_table)
+            token_embed += pos_embed
+        return token_embed
+    else:
+        real_token_ids = token_ids[:unpadded_size]
+        output = torch.empty((token_ids.shape[0], embedding_table.shape[1]),
+                             dtype=embedding_table.dtype,
+                             device=get_accelerator().current_device())
+        unpadded_output = torch.nn.functional.embedding(real_token_ids, embedding_table)
+
+        # Positional embeddings aren't padded because it's simulated
+        if positional_embedding_table is not None:
+            pos_embed = torch.nn.functional.embedding(positional_ids, positional_embedding_table)
+            unpadded_output += pos_embed
+
+        output[:unpadded_size] = unpadded_output
+        return output
+
+
+def _ragged_embed_test_helper(sequence_config: List[Tuple[int, int]],
+                              embed_dtype: torch.dtype,
+                              token_dtype: torch.dtype,
+                              embed_dim: int,
+                              vocab_size: int,
+                              do_padding: bool = False,
+                              pos_embed_size: int = -1,
+                              pos_embed_offset: int = 0) -> None:
+    """
+    Helper for embedding test to limit the number of tests to run.
+
+    Params:
+        embed_dim (int): Model dimension
+        vocab_size (int): Leading dimension on embedding weight
+        pos_embed_size (int): Size of positional embedding. If negative, no positional embedding
+            is used.
+        pos_embed_offset (int): Offset for positional embedding. Effectively, the raw offsets
+            of a token into a sequence are offset by this amount into the embedding matrix. (
+            i.e. the shape of the positional embeddings is (pos_embed_size + pos_embed_offset
+            embed_dim)
+    """
+    device = get_accelerator().current_device()
+
+    # Heads/Block size are irrelevant here but need something.
+    batch, _, _, = build_batch_and_manager(sequence_config, 64, 16, 64, vocab_range=vocab_size, padding=do_padding)
+
+    embedding_table = torch.randn((vocab_size, embed_dim), dtype=embed_dtype, device=device)
+
+    if pos_embed_size > 0:
+        pos_embedding_table = torch.randn((pos_embed_size + pos_embed_offset, embed_dim),
+                                          dtype=embed_dtype,
+                                          device=device)
+        positional_ids = torch.cat([
+            torch.arange(start_idx, start_idx + seq_len, dtype=token_dtype, device=device)
+            for seq_len, start_idx in sequence_config
+        ]) + pos_embed_offset
+    else:
+        pos_embedding_table = None
+        positional_ids = None
+
+    baseline_output = baseline_implementation(batch.input_ids().to(token_dtype), embedding_table, batch.current_tokens,
+                                              pos_embedding_table, positional_ids)
+
+    kernel = RaggedEmbeddingKernel(embed_dtype, token_dtype, embed_dim)
+    output = torch.empty_like(baseline_output)
+
+    kernel(output,
+           batch,
+           embedding_table,
+           position_embed_weight=pos_embedding_table,
+           position_embed_offset=pos_embed_offset)
+
+    if do_padding:
+        assert output.shape[0] != batch.current_tokens
+
+    assert allclose(output[:batch.current_tokens], baseline_output[:batch.current_tokens])
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('token_dtype', [torch.int32, torch.int64])
+@pytest.mark.parametrize('embed_dtype', get_dtypes())
+def test_dtype_permutations(token_dtype: torch.dtype, embed_dtype: torch.dtype) -> None:
+    """
+    Validate (on a single problem size) that the kernel support for different data types
+    is correct.
+    """
+    embed_dim = 4096
+    vocab_size = 50304
+
+    _ragged_embed_test_helper([(256, 0)], embed_dtype, token_dtype, embed_dim, vocab_size)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('vocab_size, embed_dim', [(1024, 1024), (32000, 5120), (50304, 6144)])
+def test_problem_size_permutations(vocab_size: int, embed_dim: int) -> None:
+    """
+    Validate on wider range of problem sizes.
+    """
+
+    _ragged_embed_test_helper([(256, 0)], torch.float16, torch.int32, embed_dim, vocab_size)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('seq_lens', [[128, 64, 192, 32], [57, 112, 63, 89, 1, 1, 1, 1]])
+@pytest.mark.parametrize('do_padding', [True, False])
+def test_complex_sequences(seq_lens: List[int], do_padding: bool) -> None:
+    """
+    Validate on different ragged batch construction scenarios.
+    """
+    embed_dim = 4096
+    vocab_size = 50304
+
+    _ragged_embed_test_helper([(seq_len, 0) for seq_len in seq_lens],
+                              torch.float16,
+                              torch.int32,
+                              embed_dim,
+                              vocab_size,
+                              do_padding=do_padding)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("seq_lens", [[(256, 0)], [(256, 0),
+                                                   (128, 0)], [(256, 0), (128, 0),
+                                                               (64, 0)], [(1, 877), (619, 0), (213, 372), (1, 45)]])
+def test_positional_embedding(seq_lens: List[Tuple[int, int]]) -> None:
+    """
+    Validate that positional embedding works correctly.
+    """
+    embed_dim = 4096
+    vocab_size = 50304
+
+    _ragged_embed_test_helper(seq_lens, torch.float16, torch.int32, embed_dim, vocab_size, pos_embed_size=2048)
+
+
+@pytest.mark.inference_v2_ops
+def test_positional_embedding_offset() -> None:
+    """
+    Validate that positional embedding works correctly with an offset.
+    """
+    embed_dim = 4096
+    vocab_size = 50304
+    seq_config = [(1, 877), (619, 0), (213, 372), (1, 45)]
+
+    _ragged_embed_test_helper(seq_config,
+                              torch.float16,
+                              torch.int32,
+                              embed_dim,
+                              vocab_size,
+                              pos_embed_size=2048,
+                              pos_embed_offset=2)

tests/unit/inference/v2/kernels/ragged_ops/test_top_k_gating.py

@@ -0,0 +1,171 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.ragged_ops import RaggedTopKGating
+from .ragged_testing_utils import build_simple_batch
+from ...inference_test_utils import allclose
+
+
+def _top_k_gating_testing_helper(n_tokens: int, n_experts: int, n_top_k: int, seed: int = 0xC0FFEE) -> None:
+
+    torch.manual_seed(seed)
+    logits = torch.randn((n_tokens, n_experts), dtype=torch.float16, device=get_accelerator().current_device())
+    batch = build_simple_batch([n_tokens], padding=False)
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((n_tokens, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+    expert_offset = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    ref_weights = F.softmax(logits, dim=-1, dtype=torch.float32)
+    ref_scores, ref_indices = torch.topk(ref_weights, n_top_k, dim=-1)
+
+    assert allclose(scores, ref_scores), f"expected {ref_scores}, got {scores}"
+    assert torch.equal(expert_assignment,
+                       ref_indices.to(torch.int32)), f"expected {ref_indices}, got {expert_assignment}"
+    assert expert_counts.sum(
+    ) == n_tokens * n_top_k, f"expected {n_tokens * n_top_k} tokens, got {expert_counts.sum()}"
+
+    # Ensure that the expert offsets are unique
+    for i in range(n_experts):
+        expert_idxs = torch.where(expert_assignment == i, expert_offset, 0)
+        if expert_counts[i] > 0:
+            assert expert_idxs.unique().shape[0] == expert_counts[
+                i], f"expected {expert_counts[i]} unique offsets, got {expert_idxs.unique().shape[0]}"
+            assert expert_idxs.max(
+            ) == expert_counts[i] - 1, f"expected max offset {expert_counts[i] - 1}, got {expert_idxs.max()}"
+        else:
+            # Should have all 0's so one unique value
+            assert expert_idxs.unique().shape[0] == 1
+            assert expert_idxs.max() == 0
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('n_tokens', [1, 17, 32, 89, 433])
+def test_top_2_e_8_gating(n_tokens: int) -> None:
+    _top_k_gating_testing_helper(n_tokens=n_tokens, n_experts=8, n_top_k=2)
+
+
+def _test_single_mapping_helper(n_tokens: int,
+                                n_experts: int,
+                                assigned_expert: int,
+                                logit_fill: float = 0.0,
+                                match_fill: float = 1.0) -> None:
+
+    n_top_k = 1
+    logits = torch.full((n_tokens, n_experts),
+                        logit_fill,
+                        dtype=torch.float16,
+                        device=get_accelerator().current_device())
+
+    logits[:, assigned_expert] = match_fill
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((n_tokens, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+    expert_offset = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+    batch = build_simple_batch([n_tokens], padding=False)
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    assert expert_counts[assigned_expert] == n_tokens
+    assert torch.all(expert_assignment == assigned_expert)
+    assert torch.unique(expert_offset).shape[0] == n_tokens
+    assert allclose(scores, F.softmax(logits.float(), dim=1)[:, assigned_expert].reshape(-1, n_top_k))
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('n_tokens, n_experts', [(1, 16), (17, 16), (32, 128), (89, 128), (433, 128)])
+def test_single_mapping_gating(n_tokens: int, n_experts: int) -> None:
+    """
+    Evaluate our expert stacking behavior in complete isolation. This ensures all tokens
+    mapped to the same expert are getting unique offsets and identical scores.
+    """
+    assigned_expert = 13
+    _test_single_mapping_helper(n_tokens, n_experts, assigned_expert)
+
+
+@pytest.mark.inference_v2_ops
+def test_negative_logits():
+    """
+    Ensure that scores/values are propagated correctly when all the logits are negative. An
+    earlier implementation of the scoring would return NaN for this case.
+    """
+    _test_single_mapping_helper(128, 32, 13, logit_fill=-2.0, match_fill=-1.0)
+
+
+@pytest.mark.inference_v2_ops
+def test_determinism():
+    """
+    Ensure that ties between two logits are broken deterministically. This is essential when
+    the gating is distributed across multiple devices that need to map the same token to
+    the same expert.
+    """
+
+    n_tokens = 512
+    n_experts = 64
+    n_top_k = 1
+
+    logits = torch.zeros((n_tokens, n_experts), dtype=torch.float16, device=get_accelerator().current_device())
+    batch = build_simple_batch([n_tokens], padding=False)
+
+    logits[:, 19] = 1.0
+    logits[:, 26] = 1.0
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    for _ in range(1024):
+        expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+        scores = torch.empty((n_tokens, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+        expert_assignment = torch.empty((n_tokens, n_top_k),
+                                        dtype=torch.int32,
+                                        device=get_accelerator().current_device())
+        expert_offset = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+        batch = build_simple_batch([n_tokens], padding=False)
+
+        gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+        assert expert_counts[19] == n_tokens
+        assert expert_counts[26] == 0
+        assert torch.all(expert_assignment == 19)
+        assert torch.unique(expert_offset).shape[0] == n_tokens
+        assert allclose(scores, F.softmax(logits.float(), dim=1)[:, 19].reshape(-1, 1))
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('n_tokens, n_experts', [(1, 16), (17, 16), (32, 128), (89, 128), (433, 2)])
+def test_score_accuracy(n_tokens: int, n_experts: int) -> None:
+    """
+    Validate expert scores are correct.
+    """
+    logits = torch.randn((n_tokens, n_experts), dtype=torch.float16, device=get_accelerator().current_device())
+    batch = build_simple_batch([n_tokens], padding=False)
+    n_top_k = 1
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((n_tokens, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+    expert_offset = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+
+    ref_scores = F.softmax(logits.float(), dim=1).max(dim=1).values
+    ref_scores = ref_scores.reshape(-1, 1)
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    assert allclose(scores, ref_scores)
+    assert expert_counts.sum() == n_tokens

tests/unit/inference/v2/model_implementations/__init__.py

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

tests/unit/inference/v2/model_implementations/parameters/__init__.py

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

tests/unit/inference/v2/model_implementations/parameters/test_contiguify.py

@@ -0,0 +1,120 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.model_implementations.flat_model_helpers import (
+    flatten_inference_model,
+    restore_inference_model,
+)
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+from .utils import SimpleParam, DummyInferenceModel
+
+
+class TransformerLayerContainer(LayerContainer):
+    """
+    Stub layer container
+    """
+    PARAM_MAPPING = {
+        "param_1": "param_1.param",
+        "param_2": "param_2.param",
+    }
+
+    param_1: SimpleParam
+
+    param_2: SimpleParam
+
+
+class NonTransformerContainer(LayerContainer):
+    """
+    Stub layer container
+    """
+    PARAM_MAPPING = {
+        "param_1": "param_1.param",
+        "param_2": "param_2.param",
+        "param_3": "param_3.param",
+    }
+
+    param_1: SimpleParam
+
+    param_2: SimpleParam
+
+    param_3: SimpleParam
+
+
+@pytest.mark.inference_v2
+def test_contiguify_roundtrip():
+    """
+    Validate that contiguify round trips and reconstructions are correct.
+    """
+    model = DummyInferenceModel()
+
+    n_layers = 2
+    transformer_params = []
+    transformer_containers = []
+
+    # Create parameters and populate them into the containers
+    for i in range(n_layers):
+        transformer_containers.append(TransformerLayerContainer(model))
+        layer_params = []
+        for j in range(2):
+            layer_params.append(torch.rand(16, 16))
+            transformer_containers[i].set_dependency(f"param_{j+1}", layer_params[j])
+
+        layer_params = [p.to(get_accelerator().current_device()) for p in layer_params]
+
+        transformer_params.append(layer_params)
+        assert transformer_containers[i].is_populated == True
+
+    non_transformer_params = []
+    non_transformer_container = NonTransformerContainer(model)
+
+    for i in range(3):
+        non_transformer_params.append(torch.rand(16, 16).permute(1, 0))
+        non_transformer_container.set_dependency(f"param_{i+1}", non_transformer_params[i])
+
+    non_transformer_params = [p.to(get_accelerator().current_device()) for p in non_transformer_params]
+
+    def validate_containers(t_containers: List[LayerContainer], n_t_containers: LayerContainer,
+                            t_params: List[List[torch.Tensor]], n_t_params: List[torch.Tensor]):
+        """
+        Validate params match what is on the containers.
+        """
+        for i in range(n_layers):
+            l_c = t_containers[i]
+
+            assert l_c.is_initialized == True
+
+            assert torch.equal(l_c.param_1, t_params[i][0])
+            assert torch.equal(l_c.param_2, t_params[i][1])
+
+        assert n_t_containers.is_initialized == True
+        assert torch.equal(n_t_containers.param_1, n_t_params[0])
+        assert torch.equal(n_t_containers.param_2, n_t_params[1])
+        assert torch.equal(n_t_containers.param_3, n_t_params[2])
+        assert not n_t_containers.param_1.is_contiguous()
+        assert not n_t_containers.param_2.is_contiguous()
+        assert not n_t_containers.param_3.is_contiguous()
+
+    buffer, metadata = flatten_inference_model(transformer_containers, non_transformer_container, "NoOpPolicy")
+
+    # Validate containers before contiguify
+    validate_containers(transformer_containers, non_transformer_container, transformer_params, non_transformer_params)
+
+    # Validate restore pass
+    transformer_containers_r = []
+    for i in range(n_layers):
+        transformer_containers_r.append(TransformerLayerContainer(model))
+
+    non_transformer_container_r = NonTransformerContainer(model)
+
+    restore_inference_model(buffer, metadata, transformer_containers_r, non_transformer_container_r)
+
+    validate_containers(transformer_containers_r, non_transformer_container_r, transformer_params,
+                        non_transformer_params)

tests/unit/inference/v2/model_implementations/parameters/test_layer_inheritance.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+
+from .utils import SimpleParam, DummyInferenceModel
+
+
+class ParentLayer(LayerContainer):
+    """
+    A layer that has a dependency on a simple parameter.
+    """
+
+    param_1: SimpleParam
+
+
+class ChildLayer(ParentLayer):
+    """
+    A layer that inherits from another layer.
+    """
+
+    param_2: SimpleParam
+
+
+@pytest.mark.inference_v2
+def test_layer_inheritance():
+    inference_model = DummyInferenceModel()
+
+    multi_param_layer = ChildLayer(inference_model)
+
+    assert multi_param_layer.n_params == 2
+    assert multi_param_layer.is_initialized is False
+
+    multi_param_layer.param_1.param = torch.ones(16, 16)
+
+    assert multi_param_layer.is_initialized is False
+
+    multi_param_layer.param_2.param = torch.full((16, 16), 2.0)
+
+    assert multi_param_layer.is_populated is True
+    assert isinstance(multi_param_layer.param_1, InferenceParameter)
+    assert isinstance(multi_param_layer.param_2, InferenceParameter)

tests/unit/inference/v2/model_implementations/parameters/test_mapping.py

@@ -0,0 +1,168 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.allocator import on_device
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase, ParamList
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+
+
+class MultiDependencyContainer(ParameterBase):
+
+    dependency_1: torch.Tensor
+
+    dependency_2: torch.Tensor
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        param = torch.cat([self.dependency_1, self.dependency_2])
+        return InferenceParameter.initialize(param)
+
+
+class ListDependencyContainer(ParameterBase):
+
+    dependencies: ParamList("list_items")  # noqa: F821
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        param = torch.cat(tuple(self.dependencies))
+        return InferenceParameter.initialize(param)
+
+
+class MappingLayer(LayerContainer):
+    PARAM_MAPPING = {
+        "model.val.item.d_1": "multi_depend.dependency_1",
+        "model.val.item.d_2": "multi_depend.dependency_2",
+        "model.list_vals.*.d": "list_depend.dependencies"
+    }
+
+    multi_depend: MultiDependencyContainer
+
+    list_depend: ListDependencyContainer
+
+
+class SubMappingLayer(MappingLayer):
+    PARAM_MAPPING = {
+        "model.val.item2.d_1": "multi_depend2.dependency_1",
+        "model.val.item2.d_2": "multi_depend2.dependency_2",
+    }
+
+    multi_depend2: MultiDependencyContainer
+
+
+class DoubleMappingLayer(LayerContainer):
+    PARAM_MAPPING = {
+        "model.val.item.d_1": ["multi_depend.dependency_1", "multi_depend.dependency_2"],
+    }
+
+    multi_depend: MultiDependencyContainer
+
+
+class InferenceModel:
+
+    @property
+    def list_items(self) -> int:
+        return 16
+
+
+@pytest.mark.inference_v2
+def test_mapping_syntax():
+    model = InferenceModel()
+
+    mapping_layer = MappingLayer(model)
+
+    mapping_layer.set_dependency("model.val.item.d_1", torch.ones(1))
+    mapping_layer.set_dependency("model.val.item.d_2", torch.ones(1) * 2)
+
+    assert isinstance(mapping_layer.multi_depend, torch.Tensor)
+
+    for i in range(16):
+        mapping_layer.set_dependency(f"model.list_vals.{i}.d", torch.ones(1) * i)
+        if i != 16 - 1:
+            assert mapping_layer.is_populated == False
+
+    assert isinstance(mapping_layer.list_depend, InferenceParameter)
+    assert mapping_layer.is_populated == True
+
+
+@pytest.mark.inference_v2
+def test_sub_mapping_syntax():
+    model = InferenceModel()
+
+    mapping_layer = SubMappingLayer(model)
+
+    mapping_layer.set_dependency("model.val.item.d_1", torch.ones(1))
+    mapping_layer.set_dependency("model.val.item.d_2", torch.ones(1) * 2)
+
+    assert isinstance(mapping_layer.multi_depend, InferenceParameter)
+
+    mapping_layer.set_dependency("model.val.item2.d_1", torch.ones(1))
+    mapping_layer.set_dependency("model.val.item2.d_2", torch.ones(1) * 2)
+
+    assert isinstance(mapping_layer.multi_depend2, InferenceParameter)
+
+    # We want to check into double digits to make sure that this isn't specific
+    # to single difit indexing.
+    for i in range(16):
+        mapping_layer.set_dependency(f"model.list_vals.{i}.d", torch.ones(1) * i)
+        if i != 16 - 1:
+            assert mapping_layer.is_populated == False
+
+    assert isinstance(mapping_layer.list_depend, InferenceParameter)
+    assert mapping_layer.is_populated == True
+
+
+@pytest.mark.inference_v2
+def test_double_mapping_syntax():
+    model = InferenceModel()
+
+    mapping_layer = DoubleMappingLayer(model)
+    mapping_layer.set_dependency("model.val.item.d_1", torch.ones(1))
+
+    # The single parameter setting should immediately make the parameter finalized
+    # and the whole layer initialized.
+    assert isinstance(mapping_layer.multi_depend, InferenceParameter)
+    assert mapping_layer.is_populated == True
+
+
+@pytest.mark.inference_v2
+def test_insufficient_mapping_syntax():
+    """
+    In the above example, we don't have a mapping for `multi_depend2.dependency_2`.
+    """
+
+    with pytest.raises(ValueError):
+
+        class InsuffienctMappingLayer(LayerContainer):
+            PARAM_MAPPING = {
+                "model.val.item.d_1": "multi_depend1.dependency_1",
+                "model.val.item.d_2": "multi_depend1.dependency_2",
+                "model.val.item2.d_1": "multi_depend2.dependency_1",
+            }
+
+            multi_depend1: MultiDependencyContainer
+
+            multi_depend2: MultiDependencyContainer
+
+
+@pytest.mark.inference_v2
+def test_unknown_target_mapping_syntax():
+    """
+    In the above example, `multi_depend_unknown` does not exist
+    """
+
+    with pytest.raises(ValueError):
+
+        class UnknownTargetMappingLayer(LayerContainer):
+            PARAM_MAPPING = {
+                "model.val.item.d_1": "multi_depend1.dependency_1",
+                "model.val.item.d_2": "multi_depend1.dependency_2",
+                "model.val.item2.d_1": "multi_depend_unknown.dependency_1",
+            }
+
+            multi_depend: MultiDependencyContainer

tests/unit/inference/v2/model_implementations/parameters/test_multi_parameter_layer.py

@@ -0,0 +1,79 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+
+from .utils import validate_device, SimpleParam, ListParam, DummyInferenceModel
+
+
+class MultiParameterLayer(LayerContainer):
+    """
+    Two dependencies, both of which are simple parameters.
+    """
+
+    param_1: SimpleParam
+
+    param_2: SimpleParam
+
+
+class MixedMultiParameterLayer(LayerContainer):
+    """
+    Two dependencies, one of which is a simple parameter, the other is a list parameter.
+    """
+
+    param_1: SimpleParam
+
+    param_2: ListParam
+
+
+@pytest.mark.inference_v2
+def test_multi_parameter_layer():
+    inference_model = DummyInferenceModel()
+
+    multi_param_layer = MultiParameterLayer(inference_model)
+
+    assert multi_param_layer.n_params == 2
+    assert multi_param_layer.is_populated is False
+
+    multi_param_layer.param_1.param = torch.ones(16, 16)
+
+    assert multi_param_layer.is_populated is False
+
+    multi_param_layer.param_2.param = torch.full((16, 16), 2.0)
+
+    assert multi_param_layer.is_populated is True
+    assert isinstance(multi_param_layer.param_1, InferenceParameter)
+    assert isinstance(multi_param_layer.param_2, InferenceParameter)
+
+
+@pytest.mark.inference_v2
+def test_mixed_multi_parameter_layer():
+    inference_model = DummyInferenceModel()
+
+    mixed_multi_param_layer = MixedMultiParameterLayer(inference_model)
+
+    assert mixed_multi_param_layer.n_params == 2
+    assert mixed_multi_param_layer.is_populated is False
+
+    mixed_multi_param_layer.param_2.params[1] = torch.full((16, 16), 2.0)
+    assert mixed_multi_param_layer.is_populated is False
+    assert not isinstance(mixed_multi_param_layer.param_2, InferenceParameter)
+
+    mixed_multi_param_layer.param_1.param = torch.ones(16, 16)
+    assert mixed_multi_param_layer.is_populated is False
+    assert isinstance(mixed_multi_param_layer.param_1, InferenceParameter)
+
+    validate_device(mixed_multi_param_layer.param_1)
+
+    mixed_multi_param_layer.param_2.params[0] = torch.full((16, 16), 2.0)
+
+    assert mixed_multi_param_layer.is_populated is True
+    assert isinstance(mixed_multi_param_layer.param_2, InferenceParameter)
+
+    validate_device(mixed_multi_param_layer.param_2)

tests/unit/inference/v2/model_implementations/parameters/test_parameter_list.py

@@ -0,0 +1,105 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.allocator import on_device
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase, ParamList
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+from deepspeed.inference.v2.model_implementations.common_parameters import *
+
+from .utils import validate_device
+
+
+class SimpleMoELayer(LayerContainer):
+
+    moe_mlp_1: UnfusedMoEMLP1Parameter
+
+
+class DummyInferenceModel:
+
+    def __init__(self, experts_per_rank: int) -> None:
+        self._num_experts = experts_per_rank
+
+    @property
+    def n_experts(self) -> int:
+        return self._num_experts
+
+    @on_device
+    def transform_moe_mlp_1_param(self, param: torch.Tensor) -> torch.Tensor:
+        return InferenceParameter.initialize(param)
+
+
+@pytest.mark.inference_v2
+def test_simple_moe_layer():
+
+    inference_model = DummyInferenceModel(experts_per_rank=2)
+
+    simple_moe_layer = SimpleMoELayer(inference_model)
+
+    assert simple_moe_layer.moe_mlp_1.experts[0] is None
+    assert simple_moe_layer.moe_mlp_1.experts[1] is None
+
+    # Set the first expert
+    simple_moe_layer.moe_mlp_1.experts[0] = torch.zeros(16, 16)
+
+    assert simple_moe_layer.moe_mlp_1.experts[0] is not None
+    assert simple_moe_layer.moe_mlp_1.experts[1] is None
+
+    assert not simple_moe_layer.is_initialized
+
+    # Set the second expert
+    simple_moe_layer.moe_mlp_1.experts[1] = torch.ones(16, 16)
+
+    # We have all the experts, so the layer should be initialized
+    assert simple_moe_layer.is_initialized
+    assert isinstance(simple_moe_layer.moe_mlp_1, torch.Tensor)
+
+    validate_device(simple_moe_layer.moe_mlp_1)
+
+
+"""
+Check that we can mix the number of elements in lists in the same context and have that
+be tracked correctly.
+"""
+
+
+class CustomListParam1(ParameterBase):
+
+    deps: ParamList("attr_1")
+
+
+class CustomListParam2(ParameterBase):
+
+    deps: ParamList("attr_2")
+
+
+class MixedLayer(LayerContainer):
+
+    list_1: CustomListParam1
+    list_2: CustomListParam2
+
+
+class MixedInferenceModel:
+
+    @property
+    def attr_1(self) -> int:
+        return 1
+
+    @property
+    def attr_2(self) -> int:
+        return 2
+
+
+@pytest.mark.inference_v2
+def test_mixed_param_lists():
+    model = MixedInferenceModel()
+
+    layer = MixedLayer(model)
+
+    assert layer.list_1.deps.n_params == 1
+    assert layer.list_2.deps.n_params == 2

tests/unit/inference/v2/model_implementations/parameters/utils.py

@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.allocator import on_device
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase, ParametrizedList
+
+
+class SimpleParam(ParameterBase):
+    """
+    Parameter with single dependency.
+    """
+
+    param: torch.Tensor
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform(self.param)
+
+
+class SimpleParametrizedList(ParametrizedList):
+    """
+    Parameter list based on `num_dependencies` attribute.
+    """
+
+    count_attr: str = "num_dependencies"
+
+
+class ListParam(ParameterBase):
+    """
+    Parameter with list dependency.
+
+    NOTE: This uses the tuple workaround for the `ParametrizedList` class
+    as described in the docstring of `ParametrizedList`.
+    """
+
+    params: SimpleParametrizedList
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform(torch.cat(tuple(self.params)))
+
+
+class DummyInferenceModel:
+
+    @property
+    def num_dependencies(self) -> int:
+        return 2
+
+    def transform(self, param: torch.Tensor) -> torch.Tensor:
+        return InferenceParameter.initialize(param)
+
+
+def validate_device(tensor: torch.Tensor):
+    assert tensor.device == torch.device(get_accelerator().current_device())

tests/unit/inference/v2/model_implementations/sharding/__init__.py

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

tests/unit/inference/v2/model_implementations/sharding/test_attn_out_sharding.py

@@ -0,0 +1,129 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.model_implementations.sharding import *
+
+# None of the logic should be dependent on head size.
+HEAD_SIZE = 64
+
+
+def fill_with_head_ids(head_size: int, n_heads: int) -> torch.Tensor:
+    """
+    Fills a tensor with the associated head ids. All columns should have the same value.
+    """
+    head_ids = torch.arange(n_heads, dtype=torch.half, device=get_accelerator().current_device())
+
+    head_ids = head_ids.repeat_interleave(head_size).repeat(head_size * n_heads).reshape(n_heads * head_size, -1)
+    return head_ids
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("n_heads, n_shards", [(1, 1), (8, 4), (32, 8)])
+def test_mha_even_sharding(n_heads: int, n_shards: int):
+    """
+    Even head sharding for MHA.
+
+    Args:
+        n_heads (int): The number QKV heads.
+        n_shards (int): The number of shards to test for.
+    """
+    param = fill_with_head_ids(HEAD_SIZE, n_heads)
+
+    n_local_heads = n_heads // n_shards
+    sharded_shape = (HEAD_SIZE * n_heads, HEAD_SIZE * n_local_heads)
+
+    for shard_rank in range(n_shards):
+        sharded_param = shard_attn_out_param(param, shard_rank, n_shards, HEAD_SIZE)
+        n_heads_local_q, _ = get_local_heads(shard_rank, n_shards, n_heads)
+
+        assert sharded_param.shape[-1] == HEAD_SIZE * n_heads_local_q
+        assert sharded_param.shape == sharded_shape
+
+        heads = torch.chunk(sharded_param, n_local_heads, dim=1)
+
+        for i, head in enumerate(heads):
+            assert torch.all(head == i + shard_rank * n_local_heads)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("n_heads, n_shards", [(3, 2), (20, 8)])
+def test_mha_unbalanced_sharding(n_heads: int, n_shards: int):
+    """
+    Unbalanced head sharding for MHA.
+
+    Args:
+        n_heads (int): The number QKV heads.
+        n_shards (int): The number of shards to test for.
+    """
+    param = fill_with_head_ids(HEAD_SIZE, n_heads)
+
+    max_heads = 0
+    min_heads = n_heads
+    seen_heads = set()
+    total_heads = 0
+
+    for shard_rank in range(n_shards):
+        sharded_param = shard_attn_out_param(param, shard_rank, n_shards, HEAD_SIZE)
+        n_heads_local_q, _ = get_local_heads(shard_rank, n_shards, n_heads)
+
+        assert sharded_param.shape[-1] == HEAD_SIZE * n_heads_local_q
+
+        n_local_heads = sharded_param.shape[1] // HEAD_SIZE
+        total_heads += n_local_heads
+        max_heads = max(max_heads, n_local_heads)
+        min_heads = min(min_heads, n_local_heads)
+
+        for i in range(n_local_heads):
+            head_ids = torch.unique_consecutive(sharded_param[:, i * HEAD_SIZE:(i + 1) * HEAD_SIZE])
+            assert len(head_ids) == 1
+            seen_heads.add(head_ids.item())
+
+    assert max_heads == min_heads + 1
+    assert total_heads == n_heads
+    assert len(seen_heads) == n_heads
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("n_heads_q, n_heads_kv, n_shards", [(20, 4, 8)])
+def test_gqa_uneven_sharding(n_heads_q: int, n_heads_kv: int, n_shards: int):
+    """
+    We only test the uneven GQA test case because even GQA shards the attention output
+    in the exact same manner as MHA.
+
+    Args:
+        n_heads_q (int): The number of query heads.
+        n_heads_kv (int): The number of key/value heads.
+        n_shards (int): The number of shards to test for.
+    """
+    param = fill_with_head_ids(HEAD_SIZE, n_heads_q)
+
+    min_heads = n_heads_q
+    max_heads = 0
+    seen_heads = set()
+    total_heads = 0
+
+    for shard_rank in range(n_shards):
+        sharded_param = shard_attn_out_param(param, shard_rank, n_shards, HEAD_SIZE, n_heads_q, n_heads_kv)
+        n_heads_local_q, _ = get_local_heads(shard_rank, n_shards, n_heads_q, n_heads_kv)
+
+        assert sharded_param.shape[-1] == HEAD_SIZE * n_heads_local_q
+
+        n_local_heads = sharded_param.shape[1] // HEAD_SIZE
+        total_heads += n_local_heads
+        max_heads = max(max_heads, n_local_heads)
+        min_heads = min(min_heads, n_local_heads)
+
+        for i in range(n_local_heads):
+            head_id = torch.unique_consecutive(sharded_param[:, i * HEAD_SIZE:(i + 1) * HEAD_SIZE])
+            assert len(head_id) == 1
+            seen_heads.add(head_id.item())
+
+    assert max_heads == min_heads + 1
+    assert total_heads == n_heads_q
+    assert len(seen_heads) == n_heads_q

tests/unit/inference/v2/model_implementations/sharding/test_mlp_sharding.py

@@ -0,0 +1,116 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.model_implementations.sharding import *
+
+
+def round_up_to_256(x: int) -> int:
+    """
+    Round up to the nearest multiple of 256.
+    """
+    return x + (256 - x % 256)
+
+
+def make_params(model_dim: int, ffn_multiplier: int, n_experts: int, gated: bool = False) -> torch.Tensor:
+    """
+
+    """
+    if gated:
+        mlp_1_intermediate = round_up_to_256(int(model_dim * ffn_multiplier * 4 / 3))
+        mlp_2_intermediate = mlp_1_intermediate // 2
+    else:
+        mlp_1_intermediate = ffn_multiplier * model_dim
+        mlp_2_intermediate = ffn_multiplier * model_dim
+
+    mlp_1_shared_dim = torch.arange(mlp_1_intermediate, dtype=torch.float32, device=get_accelerator().current_device())
+
+    mlp_1_w = mlp_1_shared_dim.repeat_interleave(model_dim).reshape(mlp_1_intermediate, model_dim)
+    mlp_1_b = mlp_1_shared_dim
+
+    mlp_2_shared_dim = torch.arange(mlp_2_intermediate, dtype=torch.float32, device=get_accelerator().current_device())
+    mlp_2_w = mlp_2_shared_dim.repeat(model_dim).reshape(model_dim, mlp_2_intermediate)
+    mlp_2_b = torch.ones(model_dim, dtype=torch.float32, device=get_accelerator().current_device())
+
+    if n_experts > 1:
+        mlp_1_w = mlp_1_w.expand(n_experts, -1, -1)
+        mlp_1_b = mlp_1_b.expand(n_experts, -1)
+        mlp_2_w = mlp_2_w.expand(n_experts, -1, -1)
+        mlp_2_b = mlp_2_b.expand(n_experts, -1)
+
+    return (mlp_1_w, mlp_1_b, mlp_2_w, mlp_2_b)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("model_dim, ffn_multiplier, n_shards", [(1024, 4, 1), (1024, 4, 8), (1024, 4, 6)])
+@pytest.mark.parametrize("n_experts", [1, 16])
+def test_even_ffn_sharding(model_dim: int, ffn_multiplier: int, n_shards: int, n_experts: int):
+    """
+    FFN sharding tends to be much simpler than attention sharding since it works on larger granularities.
+    While the test case of (1024, 4, 6) is not a use case we're likely to see, this does ensure that
+    the sharding logic will round correctly for the alignments we care about.
+    """
+    mlp_1_w, mlp_1_b, mlp_2_w, mlp_2_b = make_params(model_dim, ffn_multiplier, n_experts)
+
+    total_ffn_dim = model_dim * ffn_multiplier
+    mapped_neurons = 0
+
+    is_moe = n_experts > 1
+
+    for shard_rank in range(n_shards):
+        shard_1_w = shard_mlp_1_param(mlp_1_w, shard_rank, n_shards, is_moe=is_moe)
+        shard_1_b = shard_mlp_1_param(mlp_1_b, shard_rank, n_shards, is_moe=is_moe)
+        shard_2_w = shard_mlp_2_param(mlp_2_w, shard_rank, n_shards, is_moe=is_moe)
+        shard_2_b = shard_mlp_2_param(mlp_2_b, shard_rank, n_shards, is_moe=is_moe)
+
+        assert shard_1_w.shape[-2] == shard_2_w.shape[-1]
+        assert shard_1_w.shape[-2] % DEFAULT_SHARD_GRANULARITY == 0
+        assert shard_1_w.shape[-2] == shard_1_b.shape[-1]
+
+        mapped_neurons += shard_1_w.shape[-2]
+
+        if shard_rank != 0:
+            assert shard_2_b is None
+        else:
+            assert shard_2_b.shape[-1] == model_dim
+
+    assert mapped_neurons == total_ffn_dim
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("model_dim, ffn_multiplier, n_shards", [(1024, 4, 1), (1024, 4, 8), (1024, 4, 6)])
+@pytest.mark.parametrize("n_experts", [1, 16])
+def test_gated_ffn_sharding(model_dim: int, ffn_multiplier: int, n_shards: int, n_experts: int):
+    """
+    Test the same cases assuming a gated regime.
+    """
+    mlp_1_w, mlp_1_b, mlp_2_w, mlp_2_b = make_params(model_dim, ffn_multiplier, n_experts, gated=True)
+
+    total_ffn_dim = round_up_to_256(int(model_dim * ffn_multiplier * 4 / 3))
+    mapped_neurons = 0
+
+    is_moe = n_experts > 1
+
+    for shard_rank in range(n_shards):
+        shard_1_w = shard_mlp_1_param(mlp_1_w, shard_rank, n_shards, gated=True, is_moe=is_moe)
+        shard_1_b = shard_mlp_1_param(mlp_1_b, shard_rank, n_shards, gated=True, is_moe=is_moe)
+        shard_2_w = shard_mlp_2_param(mlp_2_w, shard_rank, n_shards, is_moe=is_moe)
+        shard_2_b = shard_mlp_2_param(mlp_2_b, shard_rank, n_shards, is_moe=is_moe)
+
+        assert shard_1_w.shape[-2] == shard_2_w.shape[-1] * 2
+        assert shard_1_w.shape[-2] % DEFAULT_SHARD_GRANULARITY == 0
+        assert shard_1_w.shape[-2] == shard_1_b.shape[-1]
+
+        mapped_neurons += shard_1_w.shape[-2]
+
+        if shard_rank != 0:
+            assert shard_2_b is None
+        else:
+            assert shard_2_b.shape[-1] == model_dim
+
+    assert mapped_neurons == total_ffn_dim

tests/unit/inference/v2/model_implementations/sharding/test_qkv_sharding.py

@@ -0,0 +1,251 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.model_implementations.sharding import *
+
+
+def fill_with_head_ids(head_size: int, n_heads_q: int, n_heads_kv: Optional[int] = None) -> torch.Tensor:
+    """
+
+    """
+    head_ids_q = torch.arange(n_heads_q, dtype=torch.half, device=get_accelerator().current_device())
+    head_vals_q = head_ids_q.repeat_interleave(head_size * head_size * n_heads_q).reshape(n_heads_q * head_size, -1)
+
+    if n_heads_kv is None:
+        return torch.cat([head_vals_q, head_vals_q, head_vals_q], dim=0)
+
+    head_ids_k = torch.arange(n_heads_kv, dtype=torch.half, device=get_accelerator().current_device())
+    head_vals_k = head_ids_k.repeat_interleave(head_size * head_size * n_heads_q).reshape(n_heads_kv * head_size, -1)
+
+    return torch.cat([head_vals_q, head_vals_k, head_vals_k], dim=0)
+
+
+def validate_inferred_shape(shard: torch.Tensor, head_size: int, n_local_q_heads: int, n_local_kv_heads: int):
+    """
+    Validate that the leading dim of the shard is of the expected size and aligns with the sharding
+    logic for the attention computation itself.
+    """
+    inferred_leading_dim = head_size * (n_local_q_heads + 2 * n_local_kv_heads)
+    assert shard.shape[0] == inferred_leading_dim
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads,n_shards", [(1, 1), (32, 1), (32, 8)])
+def test_even_mha_sharding(head_size: int, n_heads: int, n_shards: int):
+    """
+    Test for MHA sharding. In these scenarios, we expect that each of the shards
+    should be the same size.
+    """
+    param = fill_with_head_ids(head_size, n_heads)
+
+    heads_per_shard = n_heads // n_shards
+
+    for shard_rank in range(n_shards):
+
+        shard = shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads, n_heads)
+        n_local_q_heads, n_local_kv_heads = get_local_heads(shard_rank, n_shards, n_heads, n_heads)
+        validate_inferred_shape(shard, head_size, n_local_q_heads, n_local_kv_heads)
+
+        assert shard.shape == (3 * head_size * heads_per_shard, head_size * n_heads)
+
+        heads = shard.chunk(heads_per_shard * 3, dim=0)
+        for i in range(heads_per_shard):
+            assert torch.all(heads[i] == i + shard_rank * heads_per_shard)
+            assert torch.all(heads[i + heads_per_shard] == i + shard_rank * heads_per_shard)
+            assert torch.all(heads[i + heads_per_shard * 2] == i + shard_rank * heads_per_shard)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads, n_shards", [(3, 2), (20, 8)])
+def test_unbalanced_mha_sharding(head_size: int, n_heads: int, n_shards: int):
+    """
+    Test MHA sharding when the distribution of heads will not be equal across all ranks.
+    """
+    param = fill_with_head_ids(head_size, n_heads)
+
+    max_heads = 0
+    min_heads = n_heads
+    total_heads = 0
+    seen_heads = set()
+
+    for shard_rank in range(n_shards):
+        shard = shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads, n_heads)
+        n_local_q_heads, n_local_kv_heads = get_local_heads(shard_rank, n_shards, n_heads, n_heads)
+        validate_inferred_shape(shard, head_size, n_local_q_heads, n_local_kv_heads)
+
+        n_heads_in_shard = shard.shape[0] // head_size // 3
+
+        max_heads = max(max_heads, n_heads_in_shard)
+        min_heads = min(min_heads, n_heads_in_shard)
+        total_heads += n_heads_in_shard
+
+        heads = shard.chunk(n_heads_in_shard * 3, dim=0)
+
+        for local_head_id in range(n_heads_in_shard):
+            head_qkv = torch.cat([
+                heads[local_head_id], heads[local_head_id + n_heads_in_shard],
+                heads[local_head_id + 2 * n_heads_in_shard]
+            ],
+                                 dim=0)
+            assert head_qkv.shape == (3 * head_size, head_size * n_heads)
+
+            global_head_id = torch.unique_consecutive(head_qkv)
+            assert len(global_head_id) == 1
+
+            seen_heads.add(global_head_id.item())
+
+    assert max_heads - min_heads <= 1
+    assert total_heads == n_heads
+    assert len(seen_heads) == n_heads
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads_q, n_heads_kv, n_shards", [(4, 2, 1), (8, 2, 1), (64, 16, 8)])
+def test_gqa_even_sharding(head_size: int, n_heads_q: int, n_heads_kv: int, n_shards: int):
+    """
+    Test GQA sharding when the KV heads are evenly divisible by the number of shards.
+    """
+    param = fill_with_head_ids(head_size, n_heads_q, n_heads_kv)
+
+    n_kv_heads_in_shard = n_heads_kv // n_shards
+    n_q_heads_in_shard = n_heads_q // n_shards
+
+    for shard_rank in range(n_shards):
+        shard = shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads_q, n_heads_kv)
+        n_local_q_heads, n_local_kv_heads = get_local_heads(shard_rank, n_shards, n_heads_q, n_heads_kv)
+        validate_inferred_shape(shard, head_size, n_local_q_heads, n_local_kv_heads)
+
+        assert shard.shape[0] == (n_q_heads_in_shard + n_kv_heads_in_shard * 2) * head_size
+
+        q = shard[:n_q_heads_in_shard * head_size]
+        k = shard[n_q_heads_in_shard * head_size:(n_q_heads_in_shard + n_kv_heads_in_shard) * head_size]
+        v = shard[(n_q_heads_in_shard + n_kv_heads_in_shard) * head_size:]
+
+        for local_head_id in range(n_q_heads_in_shard):
+            assert torch.all(q[local_head_id * head_size:(local_head_id + 1) * head_size] == local_head_id +
+                             shard_rank * n_q_heads_in_shard)
+
+        for local_head_id in range(n_kv_heads_in_shard):
+            assert torch.all(k[local_head_id * head_size:(local_head_id + 1) * head_size] == local_head_id +
+                             shard_rank * n_kv_heads_in_shard)
+            assert torch.all(v[local_head_id * head_size:(local_head_id + 1) * head_size] == local_head_id +
+                             shard_rank * n_kv_heads_in_shard)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads_q, n_heads_kv, n_shards", [(4, 2, 4), (20, 4, 8)])
+def test_gqa_uneven_sharding(head_size: int, n_heads_q: int, n_heads_kv: int, n_shards: int):
+    """
+    Test GQA sharding when there are more shards than KV heads.
+    """
+    param = fill_with_head_ids(head_size, n_heads_q, n_heads_kv)
+
+    n_kv_heads_in_shard = 1
+    n_shards_per_kv_head = n_shards // n_heads_kv
+
+    max_heads = 0
+    min_heads = n_heads_q
+    total_heads = 0
+    seen_heads = set()
+
+    for shard_rank in range(n_shards):
+        shard = shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads_q, n_heads_kv)
+        n_local_q_heads, n_local_kv_heads = get_local_heads(shard_rank, n_shards, n_heads_q, n_heads_kv)
+        validate_inferred_shape(shard, head_size, n_local_q_heads, n_local_kv_heads)
+
+        local_n_heads_q = (shard.shape[0] - 2 * n_kv_heads_in_shard * head_size) // head_size
+
+        max_heads = max(max_heads, local_n_heads_q)
+        min_heads = min(min_heads, local_n_heads_q)
+        total_heads += local_n_heads_q
+
+        q = shard[:local_n_heads_q * head_size]
+        kv = shard[local_n_heads_q * head_size:]
+
+        for local_head_id in range(local_n_heads_q):
+            q_head_id = torch.unique_consecutive(q[local_head_id * head_size:(local_head_id + 1) * head_size])
+            assert len(q_head_id) == 1
+
+            seen_heads.add(q_head_id.item())
+
+        kv_id_calc = shard_rank // n_shards_per_kv_head
+        kv_id = torch.unique_consecutive(kv)
+        assert len(kv_id) == 1
+        assert kv_id.item() == kv_id_calc
+
+    assert max_heads - min_heads <= 1
+    assert total_heads == n_heads_q
+    assert len(seen_heads) == n_heads_q
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads, n_shards", [(6, 8)])
+def test_unsupported_mha_configs(head_size: int, n_heads: int, n_shards: int):
+    """
+    Sharding should fail if there are fewer heads than shards.
+
+    TODO(cmikeh2): Look to support this configuration.
+    """
+    param = fill_with_head_ids(head_size, n_heads)
+
+    for shard_rank in range(n_shards):
+        with pytest.raises(ValueError):
+            shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads, n_heads)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads_q, n_heads_kv, n_shards", [(5, 2, 1), (40, 10, 8), (30, 5, 8)])
+def test_unsupported_gqa_configs(head_size: int, n_heads_q: int, n_heads_kv: int, n_shards: int):
+    """
+    GQA has stricter requirements. We must be able to evenly shard or distribute the KV heads.
+
+    Test cases are to test the following preconditions specifically:
+        1. n_heads_q % n_heads_kv == 0
+        2. We must be able to evenly distribute KV heads
+        3. We must be able to evely split KV heads
+    """
+    param = fill_with_head_ids(head_size, n_heads_q, n_heads_kv)
+
+    for shard_rank in range(n_shards):
+        with pytest.raises(ValueError):
+            shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads_q, n_heads_kv)
+
+
+@pytest.mark.inference_v2
+def test_mha_input_shape_error():
+
+    param = torch.empty(256, 128)
+
+    n_heads = 2
+    head_size = 64
+
+    with pytest.raises(ValueError):
+        shard_qkv_param(param, 0, 1, 64)
+
+
+@pytest.mark.inference_v2
+def test_gqa_input_shape_error():
+
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 4
+
+    # Correct shape is 1536 (=16 * 64 + 2 * 4 * 64), 1024
+    param = torch.empty(2048, 1024)
+
+    with pytest.raises(ValueError):
+        shard_qkv_param(param, 0, 1, head_size, n_heads_q, n_heads_kv)

tests/unit/inference/v2/modules/__init__.py

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

tests/unit/inference/v2/modules/test_blas_linear_module.py

@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum, is_gated
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSLinearConfig
+from deepspeed.inference.v2.modules.interfaces import DSLinearRegistry
+from ...v2.inference_test_utils import allclose
+
+
+def reference_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
+                             act_type: ActivationType) -> torch.Tensor:
+    dtype = hidden_states.dtype
+    out_states = torch.nn.functional.linear(hidden_states, weight, bias)
+    out_states.float()
+
+    if is_gated(act_type):
+        act_func_map = {
+            ActivationType.ReGLU: torch.nn.functional.relu,
+            ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
+            ActivationType.SiGLU: torch.nn.functional.silu,
+        }
+
+        act_act = out_states[..., ::2]
+        act_linear = out_states[..., 1::2]
+
+        act_act = act_func_map[act_type](act_act)
+        out_states = act_act * act_linear
+    else:
+        act_func_map = {
+            ActivationType.RELU: torch.nn.functional.relu,
+            ActivationType.GELU: torch.nn.functional.gelu,
+            ActivationType.SILU: torch.nn.functional.silu,
+            ActivationType.IDENTITY: lambda x: x,
+        }
+
+        out_states = act_func_map[act_type](out_states)
+    return out_states.to(dtype)
+
+
+def _blas_linear_helper(tokens: int,
+                        in_channels: int,
+                        out_channels: int,
+                        dtype: DtypeEnum,
+                        act_fn: ActivationType,
+                        use_bias: bool = True) -> None:
+    linear_config = DSLinearConfig(max_tokens=2048,
+                                   in_channels=in_channels,
+                                   out_channels=out_channels,
+                                   activation=act_fn,
+                                   input_dtype=dtype,
+                                   output_dtype=dtype)
+
+    bundle = ConfigBundle(name='blas_fp_linear', config=linear_config)
+
+    module = DSLinearRegistry.instantiate_config(bundle)
+
+    # Input vals
+    hidden_states = torch.randn(
+        (tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+
+    weight_out_channels = 2 * out_channels if is_gated(act_fn) else out_channels
+    weight = torch.randn(
+        (weight_out_channels, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+    if use_bias:
+        bias = torch.randn(
+            (weight_out_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+    else:
+        bias = None
+
+    # Reference output
+    ref_output = reference_implementation(hidden_states, weight, bias, act_fn)
+
+    # New output
+    ds_output = module(hidden_states, weight, bias)
+
+    # Check
+    assert allclose(ds_output, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, in_channels, out_channels", [(1, 4608, 1728), (37, 8192, 4096), (1280, 3072, 6144)])
+def test_blas_linear_shapes(tokens: int, in_channels: int, out_channels: int) -> None:
+
+    _blas_linear_helper(tokens, in_channels, out_channels, DtypeEnum.fp16, ActivationType.IDENTITY)
+
+
+all_acts = [
+    ActivationType.RELU,
+    ActivationType.GELU,
+    ActivationType.SILU,
+    ActivationType.GEGLU,
+    ActivationType.ReGLU,
+    ActivationType.SiGLU,
+]
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("act_fn", all_acts)
+@pytest.mark.parametrize("use_bias", [True, False])
+def test_blas_linear_act_fn(act_fn: ActivationType, use_bias: bool) -> None:
+
+    _blas_linear_helper(283, 512, 4096, DtypeEnum.fp16, act_fn, use_bias=use_bias)

tests/unit/inference/v2/modules/test_blocked_attn.py

@@ -0,0 +1,207 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import itertools
+
+from typing import List, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSSelfAttentionConfig, PositionalEmbeddingType, RotateHalfConfig
+from deepspeed.inference.v2.modules.interfaces import DSSelfAttentionRegistry, DSSelfAttentionBase
+
+from ..kernels.ragged_ops.ragged_testing_utils import build_batch_and_manager
+from ...v2.inference_test_utils import allclose
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+    validate_accuracy = True
+except ImportError:
+    validate_accuracy = False
+
+
+def _blocked_flash_testing_helper(head_size: int,
+                                  n_heads_q: int,
+                                  n_heads_kv: int,
+                                  seq_params: List[Tuple[int, int]],
+                                  trained_freqs: bool = None) -> None:
+    """
+    Helper function for testing blocked flash attention. This implementation is based on
+    the implemnentation in ``unit.inference.kernels.ragged_ops.test_blocked_flash`` but
+    integrates functionality to validate the composability.
+    """
+    if trained_freqs is None:
+        embed_type = PositionalEmbeddingType.none
+        embed_args = None
+    else:
+        embed_type = PositionalEmbeddingType.rotate_half
+        embed_args = RotateHalfConfig(use_trained_freqs=trained_freqs)
+
+    attn_config = DSSelfAttentionConfig(max_tokens=2048,
+                                        n_heads_q=n_heads_q,
+                                        n_heads_kv=n_heads_kv,
+                                        head_size=head_size,
+                                        max_sequences=32,
+                                        positional_embedding_type=embed_type,
+                                        positional_embedding_config=embed_args)
+
+    config = ConfigBundle(name='dense_blocked_attention', config=attn_config)
+    attn_module: DSSelfAttentionBase = DSSelfAttentionRegistry.instantiate_config(config)
+
+    kv_block_size = attn_module.kv_block_size
+
+    kvs = []
+    for _, history_len in seq_params:
+        if history_len > 0:
+            kvs.append(
+                torch.randn((history_len, 2 * n_heads_kv * head_size),
+                            device=get_accelerator().current_device(),
+                            dtype=torch.float16))
+        else:
+            kvs.append(None)
+
+    batch, state_manager, _ = build_batch_and_manager(seq_params, head_size, n_heads_kv, kv_block_size, kv_fill=kvs)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=get_accelerator().current_device(),
+                      dtype=torch.float16)
+
+    kv_cache = state_manager.get_cache(0)
+
+    attn_module.build_atoms(batch)
+    if not trained_freqs:
+        out = attn_module(qkv, kv_cache, batch)
+    else:
+        inv_freqs = torch.randn((head_size // 2, ), device=get_accelerator().current_device(), dtype=torch.float16)
+        out = attn_module(qkv, kv_cache, batch, inv_freqs)
+
+    if validate_accuracy and trained_freqs is None:
+        cu_seqlens_q = torch.tensor([0] + list(itertools.accumulate([seq[0] for seq in seq_params])),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+        cu_seqlens_kv = torch.tensor([0] + list(itertools.accumulate([seq[1] + seq[0] for seq in seq_params])),
+                                     dtype=torch.int32,
+                                     device=get_accelerator().current_device())
+
+        inflight_kv = qkv[:, head_size * n_heads_q:]
+        full_kvs = []
+        for i, kv in enumerate(kvs):
+            if kv is not None:
+                full_kvs.append(torch.cat([kv, inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]]], dim=0))
+            else:
+                full_kvs.append(inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]])
+        run_kvs = torch.cat(full_kvs, dim=0)
+        k = run_kvs[:, :head_size * n_heads_kv]
+        v = run_kvs[:, head_size * n_heads_kv:]
+
+        q = qkv[:, :head_size * n_heads_q]
+        q_ref = q.reshape((batch.current_tokens, n_heads_q, head_size))
+        k_ref = k.reshape((k.shape[0], n_heads_kv, head_size))
+        v_ref = v.reshape((v.shape[0], n_heads_kv, head_size))
+
+        max_seqlen_q = max([seq[0] for seq in seq_params])
+        max_seqlen_kv = max([seq[1] + seq[0] for seq in seq_params])
+
+        ref_o = flash_attn_varlen_func(q_ref,
+                                       k_ref,
+                                       v_ref,
+                                       cu_seqlens_q,
+                                       cu_seqlens_kv,
+                                       max_seqlen_q,
+                                       max_seqlen_kv,
+                                       softmax_scale=1.0,
+                                       causal=True)
+
+        ref_o = ref_o.reshape(batch.current_tokens, head_size * n_heads_q)
+
+        assert allclose(out, ref_o)
+
+    get_accelerator().synchronize()
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens", [2, 33, 65, 128, 256, 2037])
+def test_single_prompt(n_tokens: int) -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(n_tokens, 0)]
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("prompt_lengths", [(128, 128), (192, 38), (514, 713), (83, 312, 610)])
+def test_multiple_prompts(prompt_lengths: Tuple[int, int]) -> None:
+    """
+    Test multiple prompts in a single batch.
+    """
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(prompt_lengths[i], 0) for i in range(len(prompt_lengths))]
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("seq_params", [(1, 34), (43, 40), (1, 144), (64, 128), (332, 628)])
+def test_continuation(seq_params: Tuple[int, int]) -> None:
+    """
+    Test continued generation/prompt processing.
+    """
+    head_size = 64
+    n_heads_q = 32
+    n_heads_kv = 32
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, [seq_params])
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_size", [64, 128])
+def test_head_size(head_size: int) -> None:
+    n_heads_q = 16
+    n_heads_kv = 16
+    seq_params = [(128, 128), (192, 38), (1, 814)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_config", [(32, 8), (64, 16), (40, 8)])
+def test_gqa(head_config: Tuple[int, int]) -> None:
+    head_size = 128
+    n_heads_q = head_config[0]
+    n_heads_kv = head_config[1]
+
+    seq_params = [(128, 128), (192, 38), (1, 814)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+def test_fully_composed() -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(332, 628), (1, 718), (1, 323), (180, 5), (224, 0)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("trained_freqs", [True, False])
+def test_rotary_emb(trained_freqs: bool) -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(332, 628), (1, 718), (1, 323), (180, 5), (224, 0)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params, trained_freqs=trained_freqs)

tests/unit/inference/v2/modules/test_cuda_pre_ln_module.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSNormConfig
+from deepspeed.inference.v2.modules.interfaces import DSPreNormRegistry
+from ...v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: Optional[torch.Tensor], gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:
+    dtype = residual.dtype
+
+    residual = residual.to(torch.float32)
+    gamma = gamma.to(torch.float32)
+    beta = beta.to(torch.float32)
+
+    if hidden_states is not None:
+        hidden_states = hidden_states.to(torch.float32)
+        residual = residual + hidden_states
+    hidden_states = torch.nn.functional.layer_norm(residual, (residual.size(-1), ),
+                                                   weight=gamma,
+                                                   bias=beta,
+                                                   eps=epsilon)
+    return residual.to(dtype), hidden_states.to(dtype)
+
+
+def _pre_ln_test_helper(n_tokens: int, n_channels: int, dtype: torch.dtype, res_add: bool = False):
+    config = DSNormConfig(max_tokens=2048,
+                          type="layer_norm",
+                          channels=n_channels,
+                          residual_dtype=dtype,
+                          input_dtype=dtype,
+                          output_dtype=dtype,
+                          eps=1e-5)
+    bundle = ConfigBundle(name='cuda_pre_ln', config=config)
+
+    # Input vals
+    if res_add:
+        hidden_states = torch.randn((n_tokens, n_channels),
+                                    dtype=dtype,
+                                    device=get_accelerator().current_device_name())
+    else:
+        hidden_states = None
+
+    residual = torch.randn((n_tokens, n_channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((n_channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    beta = torch.rand((n_channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_residual, ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    pre_ln_module = DSPreNormRegistry.instantiate_config(bundle)
+    gamma = pre_ln_module.transform_param(gamma)
+    beta = pre_ln_module.transform_param(beta)
+
+    ds_residual, ds_output = pre_ln_module(residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_residual, ref_residual)
+    assert allclose(ds_output, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 2048), (37, 8192), (1280, 768), (2048, 5120)])
+def test_token_channels(tokens: int, channels: int) -> None:
+    _pre_ln_test_helper(tokens, channels, torch.float16)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
+def test_dtype(dtype: torch.dtype) -> None:
+    _pre_ln_test_helper(733, 2560, dtype)
+
+
+@pytest.mark.inference_v2_ops
+def test_no_res_add():
+    _pre_ln_test_helper(733, 2560, torch.float16, res_add=False)

tests/unit/inference/v2/modules/test_custom_module.py

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.interfaces import DSPostNormRegistry
+from deepspeed.inference.v2.modules.configs import DSNormConfig
+from deepspeed.inference.v2.modules.implementations import cuda_post_ln
+from ...v2.inference_test_utils import allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> torch.Tensor:
+    residual_f = residual.to(torch.float32)
+    hidden_states_f = hidden_states.to(torch.float32)
+    gamma_f = gamma.to(torch.float32)
+    beta_f = beta.to(torch.float32)
+    return torch.nn.functional.layer_norm(residual_f + hidden_states_f, (hidden_states_f.size(-1), ),
+                                          weight=gamma_f,
+                                          bias=beta_f,
+                                          eps=epsilon).to(hidden_states.dtype)
+
+
+@DSPostNormRegistry.register_module
+class CustomPostLNModule(cuda_post_ln.DSPostLNCUDAModule):
+
+    @staticmethod
+    def name():
+        return 'custom_post_ln'
+
+
+"""
+Here, we explicitly register an LN implementation outside the core deepspeed repo. This should
+validate that the registry is working as expected and we can implement modules outside the core
+repo.
+"""
+
+
+@pytest.mark.inference_v2_ops
+def test_custom_registration():
+    channels = 4096
+    dtype = torch.float16
+    tokens = 1024
+
+    config = DSNormConfig(max_tokens=2048,
+                          type="layer_norm",
+                          channels=channels,
+                          residual_dtype=dtype,
+                          input_dtype=dtype,
+                          output_dtype=dtype,
+                          eps=1e-5)
+    bundle = ConfigBundle(name='custom_post_ln', config=config)
+
+    # Input vals
+    hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    beta = torch.rand((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    post_ln_module = DSPostNormRegistry.instantiate_config(bundle)
+    gamma = post_ln_module.transform_param(gamma)
+    beta = post_ln_module.transform_param(beta)
+    ds_output, _ = post_ln_module(residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_output, ref_output)

tests/unit/inference/v2/modules/test_cutlass_moe.py

@@ -0,0 +1,328 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSMoEConfig
+from deepspeed.inference.v2.modules.interfaces import DSMoERegistry
+
+from ..kernels.ragged_ops.ragged_testing_utils import build_simple_batch
+from ...v2.inference_test_utils import allclose, get_dtypes
+
+
+def _gating_reference(logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Reference gating code.
+    """
+    logits = logits.float()
+    probs = torch.nn.functional.softmax(logits, dim=1)
+
+    indices1_s = torch.argmax(probs, dim=-1)
+    mask1 = torch.nn.functional.one_hot(indices1_s, num_classes=logits.shape[-1])
+    indices_mask = mask1.sum(dim=1) * logits.shape[-1] - 1
+    indices1_s = torch.min(indices1_s, indices_mask)
+
+    gates1_s = (probs * mask1).sum(dim=1)
+
+    sorted_indices = indices1_s.sort()[1]
+    original_indices = sorted_indices.sort()[1]
+
+    exp_count = torch.bincount(indices1_s, minlength=logits.shape[-1]).long()
+    exp_count_cumsum = exp_count.cumsum(dim=0)
+
+    return sorted_indices, original_indices, exp_count_cumsum, gates1_s
+
+
+def _reference_impl(hidden_states: torch.Tensor, gate_weight: torch.Tensor, mlp_1_w: torch.Tensor,
+                    mlp_2_w: torch.Tensor, mlp_1_b: torch.Tensor, mlp_2_b: torch.Tensor,
+                    act_fn: ActivationType) -> torch.Tensor:
+    """
+    Reference implementation of the MoE module.
+    """
+
+    act_fn_dict = {
+        ActivationType.GELU: torch.nn.functional.gelu,
+        ActivationType.RELU: torch.nn.functional.relu,
+        ActivationType.SILU: torch.nn.functional.silu,
+        ActivationType.IDENTITY: lambda x: x,
+    }
+
+    logits = torch.matmul(hidden_states, gate_weight.t())
+    sorted_indices, original_indices, exp_count_cumsum, gate_scales = _gating_reference(logits)
+
+    moe_input = hidden_states[sorted_indices]
+
+    output_unordered = torch.empty_like(hidden_states)
+
+    for expert_idx in range(mlp_1_w.shape[0]):
+        min_bound = 0 if expert_idx == 0 else exp_count_cumsum[expert_idx - 1]
+        max_bound = exp_count_cumsum[expert_idx]
+
+        input_slice = moe_input[min_bound:max_bound]
+        intermediate = torch.nn.functional.linear(input_slice, mlp_1_w[expert_idx], mlp_1_b[expert_idx])
+
+        intermediate = act_fn_dict[act_fn](intermediate)
+        output_slice = torch.nn.functional.linear(intermediate, mlp_2_w[expert_idx], mlp_2_b[expert_idx])
+
+        output_unordered[min_bound:max_bound] = output_slice
+
+    output = output_unordered[original_indices]
+
+    output.mul_(gate_scales.unsqueeze(-1)).reshape(hidden_states.shape)
+    return output
+
+
+def _cutlass_moe_testing_helper(tokens: int,
+                                in_channels: int,
+                                intermediate_dim: int,
+                                experts: int,
+                                dtype: int,
+                                activation_type: ActivationType = ActivationType.GELU,
+                                use_bias: bool = True,
+                                iters: int = 1) -> None:
+
+    config = DSMoEConfig(max_tokens=4096,
+                         model_dim=in_channels,
+                         intermediate_features=intermediate_dim,
+                         n_experts=experts,
+                         activation=activation_type,
+                         input_dtype=dtype,
+                         output_dtype=dtype)
+
+    implementation_config = {"weight_dtype": DtypeEnum(dtype)}
+
+    bundle = ConfigBundle(name='cutlass_multi_gemm_moe', config=config, implementation_config=implementation_config)
+    moe_module = DSMoERegistry.instantiate_config(bundle)
+
+    batch = build_simple_batch([tokens])
+
+    # Parameters
+    gate_weight = torch.randn(
+        (experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    mlp_1_w = torch.randn(
+        (experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+    mlp_2_w = torch.randn(
+        (experts, in_channels, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    if use_bias:
+        mlp_1_b = torch.randn(
+            (experts, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+        mlp_2_b = torch.randn(
+            (experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+    else:
+        mlp_1_b = None
+        mlp_2_b = None
+
+    gate_ds = moe_module.transform_gate_param(gate_weight)
+    mlp_1_w_ds = moe_module.transform_moe_mlp_1_param(mlp_1_w)
+    mlp_1_b_ds = moe_module.transform_moe_mlp_1_param(mlp_1_b)
+    mlp_2_w_ds = moe_module.transform_moe_mlp_2_param(mlp_2_w)
+    mlp_2_b_ds = moe_module.transform_moe_mlp_2_param(mlp_2_b)
+
+    for _ in range(iters):
+        # Input vals
+        hidden_states = torch.randn(
+            (tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+        # Reference implementation
+        ref_output = _reference_impl(hidden_states, gate_weight, mlp_1_w, mlp_2_w, mlp_1_b, mlp_2_b, activation_type)
+
+        output = moe_module(hidden_states,
+                            batch,
+                            gate_ds,
+                            mlp_1_w_ds,
+                            mlp_2_w_ds,
+                            mlp_1_b=mlp_1_b_ds,
+                            mlp_2_b=mlp_2_b_ds)
+
+        # Increase the tolerance for larger meta ops since the error is additive
+        assert allclose(output, ref_output, tolerances=(1e-2, 1e-2))
+
+    get_accelerator().synchronize()
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("experts", [2, 32, 64])
+def test_expert_variance(experts: int) -> None:
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=4096,
+                                intermediate_dim=2048,
+                                experts=experts,
+                                dtype=DtypeEnum.fp16,
+                                activation_type=ActivationType.IDENTITY,
+                                use_bias=True)
+
+
+@pytest.mark.inference_v2_ops
+def test_successive_inputs():
+    """
+    The CUTLASS MoE uses persistent state (expert counts) that is assumed to be cleared
+    on each forward pass. This ensures that the module is clearing that metadata.
+    """
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=4096,
+                                intermediate_dim=2048,
+                                experts=64,
+                                dtype=DtypeEnum.fp16,
+                                activation_type=ActivationType.IDENTITY,
+                                use_bias=True,
+                                iters=10)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
+def test_dtypes(dtype: torch.dtype) -> None:
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=4096,
+                                intermediate_dim=2048,
+                                experts=64,
+                                dtype=DtypeEnum(dtype),
+                                activation_type=ActivationType.IDENTITY,
+                                use_bias=True)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("activation_type", [ActivationType.GELU, ActivationType.RELU, ActivationType.SILU])
+def test_activation_types(activation_type: ActivationType) -> None:
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=4096,
+                                intermediate_dim=2048,
+                                experts=64,
+                                dtype=DtypeEnum.fp16,
+                                activation_type=activation_type,
+                                use_bias=True)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("in_channels, out_channels", [(4096, 2048), (2048, 8192), (6144, 3072)])
+def test_in_out_channels(in_channels: int, out_channels: int) -> None:
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=in_channels,
+                                intermediate_dim=out_channels,
+                                experts=64,
+                                dtype=DtypeEnum.fp16,
+                                activation_type=ActivationType.IDENTITY,
+                                use_bias=True)
+
+
+def _mixtral_moe_baseline(hidden_states: torch.Tensor,
+                          gate_weight: torch.Tensor,
+                          mlp_w1: torch.Tensor,
+                          mlp_w2: torch.Tensor,
+                          mlp_w3: torch.Tensor,
+                          force_float: bool = False) -> torch.Tensor:
+    """
+    Baseline implementation for mixtral MoE module.
+
+    Based on transformers implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py
+    """
+    output_dtype = hidden_states.dtype
+    if force_float:
+        hidden_states = hidden_states.float()
+        gate_weight = gate_weight.float()
+        mlp_w1 = mlp_w1.float()
+        mlp_w2 = mlp_w2.float()
+        mlp_w3 = mlp_w3.float()
+
+    router_logits = torch.nn.functional.linear(hidden_states, gate_weight)
+    routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
+    routing_weights, selected_experts = routing_weights.topk(k=2, dim=-1)
+    routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+
+    # NOTE(cmikeh2): This is a difference implementation, ours will preserve the original scale
+    # as float32 and perform in-kernel fused FP16->FP32->FP16 conversion.
+    routing_weights = routing_weights.to(hidden_states.dtype)
+
+    final_hidden_states = torch.zeros_like(hidden_states)
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=gate_weight.shape[0]).permute(2, 1, 0)
+    get_accelerator().synchronize()
+
+    for expert_idx in range(gate_weight.shape[0]):
+        exp_mlp_w1 = mlp_w1[expert_idx]
+        exp_mlp_w2 = mlp_w2[expert_idx]
+        exp_mlp_w3 = mlp_w3[expert_idx]
+
+        idx, top_x = torch.where(expert_mask[expert_idx])
+
+        if top_x.shape[0] == 0:
+            continue
+
+        top_x_list = top_x.tolist()
+        idx_list = idx.tolist()
+
+        current_state = hidden_states[top_x_list]
+
+        linear = torch.nn.functional.linear
+        intermediate = torch.nn.functional.silu(linear(current_state, exp_mlp_w1)) * linear(current_state, exp_mlp_w3)
+        output = linear(intermediate, exp_mlp_w2) * routing_weights[top_x_list, idx_list].unsqueeze(-1)
+        final_hidden_states.index_add_(0, top_x, output.to(final_hidden_states.dtype))
+
+    return final_hidden_states.to(output_dtype)
+
+
+@pytest.mark.inference_v2_ops
+def test_mixtral_moe_config():
+
+    experts = 8
+    n_top_k = 2
+    in_channels = 4096
+    intermediate_dim = 2048
+    dtype = DtypeEnum.bf16
+
+    # Parameters
+    gate_weight = torch.randn(
+        (experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    mlp_w1 = torch.randn(
+        (experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+    mlp_w3 = torch.randn(
+        (experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+    mlp_w2 = torch.randn(
+        (experts, in_channels, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    n_tokens = 256
+    hidden_states = torch.randn(
+        (n_tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    baseline = _mixtral_moe_baseline(hidden_states, gate_weight, mlp_w1, mlp_w2, mlp_w3)
+
+    mlp_w13_fused = torch.cat([mlp_w1, mlp_w3], dim=-1).reshape(experts, 2 * intermediate_dim, in_channels)
+
+    config = DSMoEConfig(max_tokens=4096,
+                         model_dim=in_channels,
+                         intermediate_features=intermediate_dim,
+                         n_experts=experts,
+                         activation=ActivationType.SiGLU,
+                         input_dtype=dtype,
+                         output_dtype=dtype,
+                         top_k=n_top_k,
+                         normalize_scores=True)
+
+    implementation_config = {"weight_dtype": DtypeEnum(dtype)}
+
+    bundle = ConfigBundle(name='cutlass_multi_gemm_moe', config=config, implementation_config=implementation_config)
+    moe_module = DSMoERegistry.instantiate_config(bundle)
+
+    batch = build_simple_batch([n_tokens])
+
+    gate_ds = moe_module.transform_gate_param(gate_weight)
+    mlp_w1_ds = moe_module.transform_moe_mlp_1_param(mlp_w13_fused)
+    mlp_w2_ds = moe_module.transform_moe_mlp_2_param(mlp_w2)
+
+    output = moe_module(hidden_states, batch, gate_ds, mlp_w1_ds, mlp_w2_ds)
+
+    # NOTE(cmikeh2): These are higher than the other tests for reasons that aren't quite
+    # clear to me. My best guess is that the SiGLU activation is causing larger numerical
+    # divergence. The thresholds chosen here is based on the observed error between the
+    # float and bfloat16 reference implementations.
+    assert allclose(output, baseline.to(dtype.value), tolerances=(5e-2, 5e-2))

tests/unit/inference/v2/modules/test_post_ln_module.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSNormConfig
+from deepspeed.inference.v2.modules.interfaces import DSPostNormRegistry
+from ...v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> torch.Tensor:
+    residual_f = residual.to(torch.float32)
+    hidden_states_f = hidden_states.to(torch.float32)
+    gamma_f = gamma.to(torch.float32)
+    beta_f = beta.to(torch.float32)
+    return torch.nn.functional.layer_norm(residual_f + hidden_states_f, (hidden_states_f.size(-1), ),
+                                          weight=gamma_f,
+                                          bias=beta_f,
+                                          eps=epsilon).to(hidden_states.dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 2048), (37, 8192), (1280, 768), (2048, 5120)])
+@pytest.mark.parametrize("dtype", get_dtypes())
+def test_cuda_post_ln_module(tokens: int, channels: int, dtype: torch.dtype) -> None:
+    config = DSNormConfig(max_tokens=2048,
+                          type="layer_norm",
+                          channels=channels,
+                          residual_dtype=dtype,
+                          input_dtype=dtype,
+                          output_dtype=dtype,
+                          eps=1e-5)
+    bundle = ConfigBundle(name='cuda_post_ln', config=config)
+
+    # Input vals
+    hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    beta = torch.rand((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    post_ln_module = DSPostNormRegistry.instantiate_config(bundle)
+    gamma = post_ln_module.transform_param(gamma)
+    beta = post_ln_module.transform_param(beta)
+    ds_output, _ = post_ln_module(residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_output, ref_output)

tests/unit/inference/v2/modules/test_pre_rms_module.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSNormConfig
+from deepspeed.inference.v2.modules.interfaces import DSPreNormRegistry
+from ...v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: Optional[torch.Tensor], gamma: torch.Tensor,
+                             epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:
+    dtype = residual.dtype
+
+    if hidden_states is not None:
+        hidden_states = hidden_states
+        residual = residual + hidden_states
+
+    rms_vals = residual.to(torch.float32)
+    variance = rms_vals.pow(2).mean(-1, keepdim=True)
+    rms_vals = rms_vals * torch.rsqrt(variance + epsilon)
+
+    if gamma.dtype in [torch.float16, torch.bfloat16]:
+        rms_vals = rms_vals.to(gamma.dtype)
+
+    hidden_states = gamma * rms_vals
+
+    return residual.to(dtype), hidden_states.to(dtype)
+
+
+def _pre_rms_test_helper(n_tokens: int, n_channels: int, dtype: torch.dtype, res_add: bool = False):
+    config = DSNormConfig(max_tokens=2048,
+                          type="rms_norm",
+                          channels=n_channels,
+                          residual_dtype=dtype,
+                          input_dtype=dtype,
+                          output_dtype=dtype,
+                          eps=1e-5)
+    bundle = ConfigBundle(name='cuda_pre_rms', config=config)
+
+    # Input vals
+    if res_add:
+        hidden_states = torch.randn((n_tokens, n_channels),
+                                    dtype=dtype,
+                                    device=get_accelerator().current_device_name())
+    else:
+        hidden_states = None
+
+    residual = torch.randn((n_tokens, n_channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((n_channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_residual, ref_output = reference_implementation(residual, hidden_states, gamma, epsilon)
+
+    # New output
+    pre_ln_module = DSPreNormRegistry.instantiate_config(bundle)
+    gamma = pre_ln_module.transform_param(gamma)
+
+    ds_residual, ds_output = pre_ln_module(residual, hidden_states, gamma)
+
+    # Check
+    assert allclose(ds_residual, ref_residual)
+    assert allclose(ds_output, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 2048), (37, 8192), (1280, 768), (2048, 5120)])
+def test_token_channels(tokens: int, channels: int) -> None:
+    _pre_rms_test_helper(tokens, channels, torch.float16)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
+def test_dtype(dtype: torch.dtype) -> None:
+    _pre_rms_test_helper(733, 2560, dtype)
+
+
+@pytest.mark.inference_v2_ops
+def test_no_res_add():
+    _pre_rms_test_helper(733, 2560, torch.float16, res_add=False)

tests/unit/inference/v2/modules/test_quantized_linear_module.py

@@ -0,0 +1,183 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum, is_gated
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSLinearConfig
+from deepspeed.inference.v2.modules.interfaces import DSLinearRegistry
+from ...v2.inference_test_utils import allclose
+
+
+def reference_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
+                             act_type: ActivationType) -> torch.Tensor:
+    dtype = hidden_states.dtype
+    out_states = torch.nn.functional.linear(hidden_states, weight, bias)
+    out_states.float()
+
+    if is_gated(act_type):
+        act_func_map = {
+            ActivationType.ReGLU: torch.nn.functional.relu,
+            ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
+            ActivationType.SiGLU: torch.nn.functional.silu,
+        }
+
+        act_act = out_states[..., ::2]
+        act_linear = out_states[..., 1::2]
+
+        act_act = act_func_map[act_type](act_act)
+        out_states = act_act * act_linear
+    else:
+        act_func_map = {
+            ActivationType.RELU: torch.nn.functional.relu,
+            ActivationType.GELU: torch.nn.functional.gelu,
+            ActivationType.SILU: torch.nn.functional.silu,
+            ActivationType.IDENTITY: lambda x: x,
+        }
+
+        out_states = act_func_map[act_type](out_states)
+    return out_states.to(dtype)
+
+
+def _fp6_quant_dequant_weights(weight: torch.Tensor) -> torch.Tensor:
+    from deepspeed.inference.v2.modules.implementations.linear.quantized_linear import fp_quantize
+    weight_quantized_fake_fp6, scales = fp_quantize(weight, num_bits=6, exp_bits=3)
+    return weight_quantized_fake_fp6 * scales
+
+
+def quant_dequant_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
+                                 act_type: ActivationType) -> torch.Tensor:
+    dtype = hidden_states.dtype
+    weight_dequantized = _fp6_quant_dequant_weights(weight)
+    out_states = torch.nn.functional.linear(hidden_states, weight_dequantized, bias)
+    out_states.float()
+
+    if is_gated(act_type):
+        act_func_map = {
+            ActivationType.ReGLU: torch.nn.functional.relu,
+            ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
+            ActivationType.SiGLU: torch.nn.functional.silu,
+        }
+
+        act_act = out_states[..., ::2]
+        act_linear = out_states[..., 1::2]
+
+        act_act = act_func_map[act_type](act_act)
+        out_states = act_act * act_linear
+    else:
+        act_func_map = {
+            ActivationType.RELU: torch.nn.functional.relu,
+            ActivationType.GELU: torch.nn.functional.gelu,
+            ActivationType.SILU: torch.nn.functional.silu,
+            ActivationType.IDENTITY: lambda x: x,
+        }
+
+        out_states = act_func_map[act_type](out_states)
+    return out_states.to(dtype)
+
+
+def _fp6_quantized_linear_helper(tokens: int,
+                                 in_channels: int,
+                                 out_channels: int,
+                                 dtype: DtypeEnum,
+                                 act_fn: ActivationType,
+                                 use_bias: bool = True,
+                                 expect_failure: bool = False) -> None:
+    # The current FP6 kernel only supports NVIDIA Ampere GPUs.
+    if not 'cuda' in get_accelerator().current_device_name():
+        return
+    major, _ = torch.cuda.get_device_capability()  #ignore-cuda
+    if major != 8:
+        return
+
+    # Input vals
+    hidden_states = torch.randn(
+        (tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+
+    weight_out_channels = 2 * \
+        out_channels if is_gated(act_fn) else out_channels
+    weight = torch.randn(
+        (weight_out_channels, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+    if use_bias:
+        bias = torch.randn(
+            (weight_out_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+    else:
+        bias = None
+
+    # quantize and dequantize output
+    ref_quant_dequant_output = quant_dequant_implementation(hidden_states, weight, bias, act_fn)
+
+    linear_config = DSLinearConfig(max_tokens=2048,
+                                   in_channels=in_channels,
+                                   out_channels=out_channels,
+                                   activation=act_fn,
+                                   input_dtype=dtype,
+                                   output_dtype=dtype)
+    bundle = ConfigBundle(name='quantized_wf6af16_linear', config=linear_config)
+    fp6_linear_module = DSLinearRegistry.instantiate_config(bundle)
+    weight_fp6 = fp6_linear_module.transform_param(weight.clone().cpu()).to(get_accelerator().current_device_name())
+
+    if expect_failure:
+        with pytest.raises(ValueError) as excinfo:
+            ds_output = fp6_linear_module(hidden_states, weight_fp6, bias)
+        assert "The out and in channel should be multiple of 256 and 64 respectively." in str(excinfo.value)
+    else:
+        ds_output = fp6_linear_module(hidden_states, weight_fp6, bias)
+        # The current FP6 kernel uses FP16 Tensor Core.
+        tolerances = (3e-2, 2e-3)  # tolerances for fp16
+
+        # Check DeepSpeed implementation
+        assert allclose(ds_output, ref_quant_dequant_output, tolerances=tolerances)
+
+
+all_acts = [
+    ActivationType.RELU,
+    ActivationType.GELU,
+    ActivationType.SILU,
+    ActivationType.GEGLU,
+    ActivationType.ReGLU,
+    ActivationType.SiGLU,
+]
+all_tokens = [37]
+all_in_out_channels = [
+    (4096, 4096),
+]
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens", all_tokens)
+@pytest.mark.parametrize("in_channels, out_channels", all_in_out_channels)
+@pytest.mark.parametrize("act_fn", all_acts)
+@pytest.mark.parametrize("use_bias", [True, False])
+def test_fp6_quantized_linear_act_fn(tokens: int, in_channels: int, out_channels: int, act_fn: ActivationType,
+                                     use_bias: bool) -> None:
+    _fp6_quantized_linear_helper(tokens=tokens,
+                                 in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 dtype=DtypeEnum.fp16,
+                                 act_fn=act_fn,
+                                 use_bias=use_bias)
+
+
+# Other shapes, not supported by FP6 kernels. Will raise ValueError.
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens", all_tokens)
+@pytest.mark.parametrize("in_channels, out_channels", [(4608, 1728)])
+@pytest.mark.parametrize("act_fn", all_acts)
+@pytest.mark.parametrize("use_bias", [True, False])
+def test_fp6_quantized_linear_act_fn_fail(tokens: int, in_channels: int, out_channels: int, act_fn: ActivationType,
+                                          use_bias: bool) -> None:
+    _fp6_quantized_linear_helper(tokens=tokens,
+                                 in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 dtype=DtypeEnum.fp16,
+                                 act_fn=act_fn,
+                                 use_bias=use_bias,
+                                 expect_failure=True)

tests/unit/inference/v2/ragged/__init__.py

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

tests/unit/inference/v2/ragged/test_blocked_allocator.py

@@ -0,0 +1,166 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import random
+from typing import List
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.ragged.blocked_allocator import BlockedAllocator
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('bad_size', [0, -1])
+def test_bad_initialization(bad_size: int) -> None:
+    with pytest.raises(ValueError):
+        BlockedAllocator(bad_size)
+
+
+@pytest.mark.inference_v2
+def test_allocation() -> None:
+
+    allocator = BlockedAllocator(16)
+
+    a1 = allocator.allocate(4)
+    assert a1.numel() == 4
+    assert allocator.free_blocks == 12
+
+    a2_allocs = []
+    for i in range(3):
+        a2_allocs.append(allocator.allocate(2))
+        assert allocator.free_blocks == 12 - (i + 1) * 2
+
+    a3 = allocator.allocate(6)
+    assert a3.numel() == 6
+
+    assert allocator.free_blocks == 0
+
+    # Test that we can't allocate more blocks than we have.
+    with pytest.raises(ValueError):
+        allocator.allocate(1)
+
+    all_vals = torch.cat([a1, *a2_allocs, a3], dim=0)
+    unique_vals = torch.unique(all_vals, sorted=False)
+    assert unique_vals.numel() == all_vals.numel()
+
+
+@pytest.mark.inference_v2
+def test_too_large_allocation():
+    allocator = BlockedAllocator(16)
+
+    with pytest.raises(ValueError):
+        allocator.allocate(17)
+
+
+@pytest.mark.inference_v2
+def test_deallocation() -> None:
+    allocator = BlockedAllocator(16)
+
+    # Allocate
+    all_blocks = allocator.allocate(16)
+    assert allocator.free_blocks == 0
+
+    # Deallocate all blocks
+    allocator.free(all_blocks)
+    assert allocator.free_blocks == 16
+
+    # Get all the blocks again
+    all_blocks = allocator.allocate(16)
+
+    # Deallocate in chunks
+    c1 = all_blocks[:4]
+    c2 = all_blocks[4:8]
+
+    allocator.free(c1)
+    assert allocator.free_blocks == 4
+
+    allocator.free(c2)
+    assert allocator.free_blocks == 8
+
+    with pytest.raises(ValueError):
+        allocator.free(c1)
+
+    with pytest.raises(ValueError):
+        allocator.free(c2)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('index', [-1, 2])
+def test_invalid_dealloc_indices(index: int):
+    allocator = BlockedAllocator(1)
+
+    with pytest.raises(ValueError):
+        allocator.free(torch.tensor([index]))
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('index', [-1, 2])
+def test_invalid_alloc_indices(index: int):
+    allocator = BlockedAllocator(1)
+    allocator.allocate(1)
+
+    to_free = [0, index]
+
+    with pytest.raises(ValueError):
+        allocator.free(torch.tensor(to_free))
+
+    # Block 0 should not be freed if passed with an invalid index.
+    assert allocator.free_blocks == 0
+
+    allocator.free(torch.tensor([0]))
+    assert allocator.free_blocks == 1
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('test_iters', [8192])
+def test_long_running_allocation(test_iters: int) -> None:
+    """
+    Evaluate the stability of the allocator over a longer sequence of allocations/deallocations.
+    """
+    TOTAL_BLOCKS = 128
+
+    allocator = BlockedAllocator(TOTAL_BLOCKS)
+
+    def validate_uniqueness(all_blocks: List[torch.Tensor]) -> None:
+        all_vals = torch.cat(all_blocks, dim=0)
+        assert all_vals.numel() <= TOTAL_BLOCKS
+
+        unique_vals = torch.unique(all_vals, sorted=False)
+        assert unique_vals.numel() == all_vals.numel()
+
+    all_allocs: List[torch.Tensor] = []
+    num_allocs = 0
+    num_frees = 0
+    num_blocks_allocated = 0
+    num_blocks_freed = 0
+
+    for _ in range(test_iters):
+        decision = random.randint(0, 1)
+
+        if decision == 0:
+            blocks_to_allocate = random.randint(1, 24)
+            if blocks_to_allocate > allocator.free_blocks:
+                with pytest.raises(ValueError):
+                    allocator.allocate(blocks_to_allocate)
+            else:
+                all_allocs.append(allocator.allocate(blocks_to_allocate))
+                num_allocs += 1
+                num_blocks_allocated += blocks_to_allocate
+        else:
+            if len(all_allocs) > 0:
+                idx = random.randint(0, len(all_allocs) - 1)
+                allocator.free(all_allocs[idx])
+
+                num_frees += 1
+                num_blocks_freed += all_allocs[idx].numel()
+
+                del all_allocs[idx]
+
+        if len(all_allocs) > 0:
+            validate_uniqueness(all_allocs)
+
+    assert num_allocs == num_frees + len(all_allocs)
+    assert num_blocks_allocated == num_blocks_freed + (TOTAL_BLOCKS - allocator.free_blocks)

tests/unit/inference/v2/ragged/test_manager_configs.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+from pydantic import ValidationError
+
+from deepspeed.inference.v2.ragged import DSStateManagerConfig
+
+
+@pytest.mark.inference_v2
+def test_negative_max_tracked_sequences() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_tracked_sequences=-1)
+
+
+@pytest.mark.inference_v2
+def test_zero_max_tracked_sequences() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_tracked_sequences=0)
+
+
+@pytest.mark.inference_v2
+def test_negative_max_ragged_batch_size() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_batch_size=-1)
+
+
+@pytest.mark.inference_v2
+def test_zero_max_ragged_batch_size() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_batch_size=0)
+
+
+@pytest.mark.inference_v2
+def test_negative_max_ragged_sequence_count() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_sequence_count=-1)
+
+
+@pytest.mark.inference_v2
+def test_zero_max_ragged_sequence_count() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_sequence_count=0)
+
+
+@pytest.mark.inference_v2
+def test_too_small_max_ragged_batch_size() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_batch_size=512, max_ragged_sequence_count=1024)
+
+
+@pytest.mark.inference_v2
+def test_too_small_max_tracked_sequences() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_tracked_sequences=512, max_ragged_sequence_count=1024)

tests/unit/inference/v2/ragged/test_ragged_wrapper.py

@@ -0,0 +1,112 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.ragged import (
+    PlaceholderSequenceDescriptor,
+    RaggedBatchWrapper,
+    DSStateManagerConfig,
+)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('max_ragged_sequence_count, max_ragged_batch_size', [(128, 512), (128, 1024)])
+def test_wrapper_initialization(max_ragged_sequence_count: int, max_ragged_batch_size: int) -> None:
+    config = DSStateManagerConfig(max_tracked_sequences=max_ragged_sequence_count,
+                                  max_ragged_batch_size=max_ragged_batch_size,
+                                  max_ragged_sequence_count=max_ragged_sequence_count)
+
+    batch = RaggedBatchWrapper(config)
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('seq_len', [1, 37, 128, 512])
+def test_single_sequence_batch(seq_len: int) -> None:
+    """
+    Test we successfully construct single sequence batches and the on device metadata is accurate.
+    """
+
+    config = DSStateManagerConfig()
+    batch = RaggedBatchWrapper(config)
+
+    batch.clear()
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0
+
+    seq_desc = PlaceholderSequenceDescriptor()
+    tokens = torch.randint(0, 100, (seq_len, ))
+    batch.insert_sequence(seq_desc, tokens)
+
+    batch.finalize()
+
+    assert batch.current_tokens == seq_len
+    assert batch.current_sequences == 1
+    assert torch.equal(batch.input_ids(), tokens.to(get_accelerator().current_device()))
+    assert torch.equal(batch.tokens_to_seq(), torch.zeros_like(tokens, device=get_accelerator().current_device()))
+    assert torch.equal(batch.batch_metadata_buffer(),
+                       torch.tensor([seq_len, 1], device=get_accelerator().current_device()))
+
+    batch.clear()
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('seq_lens', [[128, 128], [1, 32, 243], [64, 1, 1, 1, 1, 393, 27, 2]])
+def test_multi_sequence_batch(seq_lens: List[int]) -> None:
+    """
+    Test sequentially adding new tokens to a batch and validate device data structures hold
+    the appropriate data.
+    """
+    config = DSStateManagerConfig()
+    batch = RaggedBatchWrapper(config)
+
+    batch.clear()
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0
+
+    all_toks = [torch.randint(0, 100, (seq_len, )) for seq_len in seq_lens]
+
+    for i, toks in enumerate(all_toks):
+        seq_desc = PlaceholderSequenceDescriptor()
+        batch.insert_sequence(seq_desc, toks)
+
+        assert batch.current_tokens == sum(seq_lens[:i + 1])
+        assert batch.current_sequences == i + 1
+
+    batch.finalize()
+
+    assert batch.current_tokens == sum(seq_lens)
+    assert batch.current_sequences == len(seq_lens)
+
+    assert torch.equal(batch.input_ids(), torch.cat(all_toks, dim=0).to(get_accelerator().current_device()))
+    assert torch.equal(
+        batch.tokens_to_seq(),
+        torch.cat([torch.full((seq_len, ), i, dtype=torch.int32) for i, seq_len in enumerate(seq_lens)],
+                  dim=0).to(get_accelerator().current_device()))
+
+    for i, seq_len in enumerate(seq_lens):
+        assert batch.inflight_seq_descriptors()[i][0] == sum(seq_lens[:i])
+        assert batch.inflight_seq_descriptors()[i][1] == seq_len
+        assert batch.inflight_seq_descriptors()[i][2] == 0
+
+    assert torch.equal(batch.batch_metadata_buffer(),
+                       torch.tensor([sum(seq_lens), len(seq_lens)], device=get_accelerator().current_device()))
+
+    batch.clear()
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0

tests/unit/launcher/test_ds_arguments.py

@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import argparse
+import pytest
+import deepspeed
+from deepspeed.utils.numa import parse_range_list
+
+
+def basic_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num_epochs', type=int)
+    return parser
+
+
+def test_no_ds_arguments_no_ds_parser():
+    parser = basic_parser()
+    args = parser.parse_args(['--num_epochs', '2'])
+    assert args
+
+    assert hasattr(args, 'num_epochs')
+    assert args.num_epochs == 2
+
+    assert not hasattr(args, 'deepspeed')
+    assert not hasattr(args, 'deepspeed_config')
+
+
+def test_no_ds_arguments():
+    parser = basic_parser()
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args(['--num_epochs', '2'])
+    assert args
+
+    assert hasattr(args, 'num_epochs')
+    assert args.num_epochs == 2
+
+    assert hasattr(args, 'deepspeed')
+    assert args.deepspeed == False
+
+    assert hasattr(args, 'deepspeed_config')
+    assert args.deepspeed_config is None
+
+
+def test_no_ds_enable_argument():
+    parser = basic_parser()
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args(['--num_epochs', '2', '--deepspeed_config', 'foo.json'])
+    assert args
+
+    assert hasattr(args, 'num_epochs')
+    assert args.num_epochs == 2
+
+    assert hasattr(args, 'deepspeed')
+    assert args.deepspeed == False
+
+    assert hasattr(args, 'deepspeed_config')
+    assert type(args.deepspeed_config) == str
+    assert args.deepspeed_config == 'foo.json'
+
+
+def test_no_ds_config_argument():
+    parser = basic_parser()
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args(['--num_epochs', '2', '--deepspeed'])
+    assert args
+
+    assert hasattr(args, 'num_epochs')
+    assert args.num_epochs == 2
+
+    assert hasattr(args, 'deepspeed')
+    assert type(args.deepspeed) == bool
+    assert args.deepspeed == True
+
+    assert hasattr(args, 'deepspeed_config')
+    assert args.deepspeed_config is None
+
+
+def test_no_ds_parser():
+    parser = basic_parser()
+    with pytest.raises(SystemExit):
+        args = parser.parse_args(['--num_epochs', '2', '--deepspeed'])
+
+
+def test_core_deepscale_arguments():
+    parser = basic_parser()
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args(['--num_epochs', '2', '--deepspeed', '--deepspeed_config', 'foo.json'])
+    assert args
+
+    assert hasattr(args, 'num_epochs')
+    assert args.num_epochs == 2
+
+    assert hasattr(args, 'deepspeed')
+    assert type(args.deepspeed) == bool
+    assert args.deepspeed == True
+
+    assert hasattr(args, 'deepspeed_config')
+    assert type(args.deepspeed_config) == str
+    assert args.deepspeed_config == 'foo.json'
+
+
+def test_core_binding_arguments():
+    core_list = parse_range_list("0,2-4,6,8-9")
+    assert core_list == [0, 2, 3, 4, 6, 8, 9]
+
+    try:
+        # negative case for range overlapping
+        core_list = parse_range_list("0,2-6,5-9")
+    except ValueError as e:
+        pass
+    else:
+        # invalid core list must fail
+        assert False
+
+    try:
+        # negative case for reverse order -- case 1
+        core_list = parse_range_list("8,2-6")
+    except ValueError as e:
+        pass
+    else:
+        # invalid core list must fail
+        assert False
+
+    try:
+        # negative case for reverse order -- case 2
+        core_list = parse_range_list("1,6-2")
+    except ValueError as e:
+        pass
+    else:
+        # invalid core list must fail
+        assert False

tests/unit/launcher/test_multinode_runner.py

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from copy import deepcopy
+from deepspeed.launcher import multinode_runner as mnrunner
+from deepspeed.launcher.runner import encode_world_info, parse_args
+import os
+import pytest
+
+
+@pytest.fixture
+def runner_info():
+    hosts = {'worker-0': 4, 'worker-1': 4}
+    world_info = encode_world_info(hosts)
+    env = deepcopy(os.environ)
+    args = parse_args(['test_launcher.py'])
+    return env, hosts, world_info, args
+
+
+def test_pdsh_runner(runner_info):
+    env, resource_pool, world_info, args = runner_info
+    runner = mnrunner.PDSHRunner(args, world_info)
+    cmd, kill_cmd, env = runner.get_cmd(env, resource_pool)
+    assert cmd[0] == 'pdsh'
+    assert env['PDSH_RCMD_TYPE'] == 'ssh'
+
+
+def test_openmpi_runner(runner_info):
+    env, resource_pool, world_info, args = runner_info
+    runner = mnrunner.OpenMPIRunner(args, world_info, resource_pool)
+    cmd = runner.get_cmd(env, resource_pool)
+    assert cmd[0] == 'mpirun'
+    assert 'eth0' in cmd
+
+
+def test_btl_nic_openmpi_runner(runner_info):
+    env, resource_pool, world_info, _ = runner_info
+    args = parse_args(['--launcher_arg', '-mca btl_tcp_if_include eth1', 'test_launcher.py'])
+    runner = mnrunner.OpenMPIRunner(args, world_info, resource_pool)
+    cmd = runner.get_cmd(env, resource_pool)
+    assert 'eth0' not in cmd
+    assert 'eth1' in cmd
+
+
+def test_btl_nic_two_dashes_openmpi_runner(runner_info):
+    env, resource_pool, world_info, _ = runner_info
+    args = parse_args(['--launcher_arg', '--mca btl_tcp_if_include eth1', 'test_launcher.py'])
+    runner = mnrunner.OpenMPIRunner(args, world_info, resource_pool)
+    cmd = runner.get_cmd(env, resource_pool)
+    assert 'eth0' not in cmd
+    assert 'eth1' in cmd
+
+
+def test_mpich_runner(runner_info):
+    env, resource_pool, world_info, args = runner_info
+    runner = mnrunner.MPICHRunner(args, world_info, resource_pool)
+    cmd = runner.get_cmd(env, resource_pool)
+    assert cmd[0] == 'mpirun'
+
+
+def test_slurm_runner(runner_info):
+    env, resource_pool, world_info, args = runner_info
+    runner = mnrunner.SlurmRunner(args, world_info, resource_pool)
+    cmd = runner.get_cmd(env, resource_pool)
+    assert cmd[0] == 'srun'
+
+
+def test_mvapich_runner(runner_info):
+    env, resource_pool, world_info, args = runner_info
+    runner = mnrunner.MVAPICHRunner(args, world_info, resource_pool)
+    cmd = runner.get_cmd(env, resource_pool)
+    assert cmd[0] == 'mpirun'

tests/unit/launcher/test_run.py

@@ -0,0 +1,180 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+from deepspeed.launcher import runner as dsrun
+
+
+def test_parser_mutual_exclusive():
+    '''Ensure dsrun.parse_resource_filter() raises a ValueError when include_str and
+    exclude_str are both provided.
+    '''
+    with pytest.raises(ValueError):
+        dsrun.parse_resource_filter({}, include_str='A', exclude_str='B')
+
+
+def test_parser_local():
+    ''' Test cases with only one node. '''
+    # First try no include/exclude
+    hosts = {'worker-0': [0, 1, 2, 3]}
+    ret = dsrun.parse_resource_filter(hosts)
+    assert (ret == hosts)
+
+    # exclude slots
+    ret = dsrun.parse_resource_filter(hosts, exclude_str='worker-0:1')
+    assert (ret == {'worker-0': [0, 2, 3]})
+
+    ret = dsrun.parse_resource_filter(hosts, exclude_str='worker-0:1,2')
+    assert (ret == {'worker-0': [0, 3]})
+
+    # only use one slot
+    ret = dsrun.parse_resource_filter(hosts, include_str='worker-0:1')
+    assert (ret == {'worker-0': [1]})
+
+    # including slots multiple times shouldn't break things
+    ret = dsrun.parse_resource_filter(hosts, include_str='worker-0:1,1')
+    assert (ret == {'worker-0': [1]})
+    ret = dsrun.parse_resource_filter(hosts, include_str='worker-0:1@worker-0:0,1')
+    assert (ret == {'worker-0': [0, 1]})
+
+    # including just 'worker-0' without : should still use all GPUs
+    ret = dsrun.parse_resource_filter(hosts, include_str='worker-0')
+    assert (ret == hosts)
+
+    # excluding just 'worker-0' without : should eliminate everything
+    ret = dsrun.parse_resource_filter(hosts, exclude_str='worker-0')
+    assert (ret == {})
+
+    # exclude all slots manually
+    ret = dsrun.parse_resource_filter(hosts, exclude_str='worker-0:0,1,2,3')
+    assert (ret == {})
+
+
+def test_parser_multinode():
+    # First try no include/exclude
+    hosts = {'worker-0': [0, 1, 2, 3], 'worker-1': [0, 1, 2, 3]}
+    ret = dsrun.parse_resource_filter(hosts)
+    assert (ret == hosts)
+
+    # include a node
+    ret = dsrun.parse_resource_filter(hosts, include_str='worker-1:0,3')
+    assert (ret == {'worker-1': [0, 3]})
+
+    # exclude a node
+    ret = dsrun.parse_resource_filter(hosts, exclude_str='worker-1')
+    assert (ret == {'worker-0': [0, 1, 2, 3]})
+
+    # exclude part of each node
+    ret = dsrun.parse_resource_filter(hosts, exclude_str='worker-0:0,1@worker-1:3')
+    assert (ret == {'worker-0': [2, 3], 'worker-1': [0, 1, 2]})
+
+
+def test_parser_errors():
+    '''Ensure we catch errors. '''
+    hosts = {'worker-0': [0, 1, 2, 3], 'worker-1': [0, 1, 2, 3]}
+
+    # host does not exist
+    with pytest.raises(ValueError):
+        dsrun.parse_resource_filter(hosts, include_str='jeff')
+    with pytest.raises(ValueError):
+        dsrun.parse_resource_filter(hosts, exclude_str='jeff')
+
+    # slot does not exist
+    with pytest.raises(ValueError):
+        dsrun.parse_resource_filter(hosts, include_str='worker-1:4')
+    with pytest.raises(ValueError):
+        dsrun.parse_resource_filter(hosts, exclude_str='worker-1:4')
+
+    # formatting
+    with pytest.raises(ValueError):
+        dsrun.parse_resource_filter(hosts, exclude_str='worker-1@worker-0:1@5')
+
+
+def test_num_plus_parser():
+    ''' Ensure we catch errors relating to num_nodes/num_gpus + -i/-e being mutually exclusive'''
+
+    # inclusion
+    with pytest.raises(ValueError):
+        dsrun.main(args="--num_nodes 1 -i localhost foo.py".split())
+    with pytest.raises(ValueError):
+        dsrun.main(args="--num_nodes 1 --num_gpus 1 -i localhost foo.py".split())
+    with pytest.raises(ValueError):
+        dsrun.main(args="--num_gpus 1 -i localhost foo.py".split())
+
+    # exclusion
+    with pytest.raises(ValueError):
+        dsrun.main(args="--num_nodes 1 -e localhost foo.py".split())
+    with pytest.raises(ValueError):
+        dsrun.main(args="--num_nodes 1 --num_gpus 1 -e localhost foo.py".split())
+    with pytest.raises(ValueError):
+        dsrun.main(args="--num_gpus 1 -e localhost foo.py".split())
+
+
+def test_hostfile_good():
+    # good hostfile w. empty lines and comment
+    hostfile = """
+    worker-1 slots=2
+    worker-2 slots=2
+
+    localhost slots=1
+    123.23.12.10 slots=2
+
+    #worker-1 slots=3
+    # this is a comment
+
+    """
+    r = dsrun._parse_hostfile(hostfile.splitlines())
+    assert "worker-1" in r
+    assert "worker-2" in r
+    assert "localhost" in r
+    assert "123.23.12.10" in r
+    assert r["worker-1"] == 2
+    assert r["worker-2"] == 2
+    assert r["localhost"] == 1
+    assert r["123.23.12.10"] == 2
+    assert len(r) == 4
+
+
+def test_hostfiles_bad():
+    # duplicate host
+    hostfile = """
+    worker-1 slots=2
+    worker-2 slots=1
+    worker-1 slots=1
+    """
+    with pytest.raises(ValueError):
+        dsrun._parse_hostfile(hostfile.splitlines())
+
+    # incorrect whitespace
+    hostfile = """
+    this is bad slots=1
+    """
+    with pytest.raises(ValueError):
+        dsrun._parse_hostfile(hostfile.splitlines())
+
+    # no whitespace
+    hostfile = """
+    missingslots
+    """
+    with pytest.raises(ValueError):
+        dsrun._parse_hostfile(hostfile.splitlines())
+
+    # empty
+    hostfile = """
+    """
+    with pytest.raises(ValueError):
+        dsrun._parse_hostfile(hostfile.splitlines())
+
+    # mix of good/bad
+    hostfile = """
+    worker-1 slots=2
+    this is bad slots=1
+    worker-2 slots=4
+    missingslots
+
+    """
+    with pytest.raises(ValueError):
+        dsrun._parse_hostfile(hostfile.splitlines())

tests/unit/launcher/test_user_args.py

@@ -0,0 +1,102 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import subprocess
+
+from types import SimpleNamespace
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.launcher.multinode_runner import MultiNodeRunner
+
+
+class DummyRunner(MultiNodeRunner):
+
+    def backend_exists(self):
+        return True
+
+    def get_cmd(self, environment, active_resources):
+        return []
+
+
+if not get_accelerator().is_available():
+    pytest.skip("only supported in accelerator environments.", allow_module_level=True)
+
+user_arg_test_script = """import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--prompt", type=str)
+parser.add_argument("--local_rank", type=int, default=0)
+parser.add_argument("--world_size", type=int, default=1)
+args = parser.parse_args()
+print("ARG PARSE SUCCESS")
+"""
+
+
+@pytest.fixture(scope="function")
+def user_script_fp(tmpdir):
+    script_fp = tmpdir.join("user_arg_test.py")
+    with open(script_fp, "w") as f:
+        f.write(user_arg_test_script)
+    return script_fp
+
+
+@pytest.fixture(scope="function")
+def cmd(user_script_fp, prompt, multi_node):
+    if multi_node:
+        cmd = ("deepspeed", "--force_multi", "--num_nodes", "1", "--num_gpus", "1", user_script_fp, "--prompt", prompt)
+    else:
+        cmd = ("deepspeed", "--num_nodes", "1", "--num_gpus", "1", user_script_fp, "--prompt", prompt)
+    return cmd
+
+
+@pytest.fixture
+def dummy_runner():
+    args = SimpleNamespace(user_args=[], user_script="dummy_script.py")
+    return DummyRunner(args, "dummy_world_info")
+
+
+@pytest.mark.parametrize("prompt", [
+    '''"I am 6' tall"''', """'I am 72" tall'""", """'"translate English to Romanian: "'""",
+    '''I'm going to tell them "DeepSpeed is the best"'''
+])
+@pytest.mark.parametrize("multi_node", [True, False])
+def test_user_args(cmd, multi_node):
+    if multi_node and get_accelerator().device_name() == "cpu":
+        pytest.skip("CPU accelerator does not support this test yet")
+    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    out, err = p.communicate()
+    assert "ARG PARSE SUCCESS" in out.decode("utf-8"), f"User args not parsed correctly: {err.decode('utf-8')}"
+
+
+def test_bash_string_args(tmpdir, user_script_fp):
+    bash_script = f"""
+    ARGS="--prompt 'DeepSpeed is the best'"
+    echo ${{ARGS}}|xargs deepspeed --num_nodes 1 --num_gpus 1 {user_script_fp}
+    """
+
+    bash_fp = tmpdir.join("bash_script.sh")
+    with open(bash_fp, "w") as f:
+        f.write(bash_script)
+
+    p = subprocess.Popen(["bash", bash_fp], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    out, err = p.communicate()
+    assert "ARG PARSE SUCCESS" in out.decode("utf-8"), f"User args not parsed correctly: {err.decode('utf-8')}"
+
+
+def test_add_export_with_special_characters(dummy_runner):
+    """
+    Values with special characters (e.g., 64(x2)) must be quoted to avoid bash syntax errors.
+    """
+    dummy_runner.add_export("SLURM_JOB_CPUS_PER_NODE", "64(x2)")
+    assert dummy_runner.exports["SLURM_JOB_CPUS_PER_NODE"] == "\"64(x2)\""
+
+
+def test_add_export_no_special_characters(dummy_runner):
+    """
+    Values without special characters should remain unquoted (e.g., PYTHONPATH).
+    This avoids issues where unnecessary quotes break module imports.
+    """
+    dummy_runner.add_export("PYTHONPATH", "/usr/local/lib/python3.9/site-packages")
+    assert dummy_runner.exports["PYTHONPATH"] == "/usr/local/lib/python3.9/site-packages"

tests/unit/linear/test_ctx.py

@@ -0,0 +1,109 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+import pytest
+from unit.common import DistributedTest
+
+import deepspeed.comm as dist
+from deepspeed.linear import LoRAConfig, init_lora
+from deepspeed.linear.optimized_linear import LoRAOptimizedLinear
+from unit.simple_model import random_dataloader, SimpleModel
+
+try:
+    import transformers
+except ImportError:
+    transformers = None
+
+if transformers is None:
+    pytest.skip("transformers is required for this test", allow_module_level=True)
+
+
+def injection_assert(model):
+    # pick out random linear that should have been replaced and initialized
+    q_proj = model.model.layers[1].self_attn.q_proj
+
+    assert isinstance(q_proj, LoRAOptimizedLinear), "injection did not happen"
+    assert q_proj._initialized, "lora was not initialized properly"
+    assert isinstance(q_proj.lora_weight_1, torch.nn.Linear)
+    assert isinstance(q_proj.lora_weight_2, torch.nn.Linear)
+
+
+class TestEngine(DistributedTest):
+    world_size = 2
+
+    def test_model(self):
+        lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=2)
+        quant_config = None
+        hidden_dim = 64
+        nlayers = 4
+
+        with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config):
+            model = SimpleModel(hidden_dim=hidden_dim, nlayers=nlayers)
+
+        init_lora(model)
+
+        model_norms = [model.linears[i].weight.norm().item() for i in range(nlayers)]
+
+        ds_config = {
+            "train_batch_size": 2,
+            "steps_per_print": 1,
+            "bf16": {
+                "enabled": True
+            },
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+            "zero_optimization": {
+                "stage": 1
+            }
+        }
+        model, *_ = deepspeed.initialize(config=ds_config, model=model, model_parameters=model.parameters())
+
+        engine_norms = [model.module.linears[i].weight.norm().item() for i in range(nlayers)]
+
+        # Ensure that sharded weights are not broadcast during engine init
+        assert engine_norms == model_norms, f"{dist.get_rank()=} base weight norms are not the same after engine init, {engine_norms=} != {model_norms=}"
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=50,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=torch.bfloat16)
+        for n, batch in enumerate(data_loader):
+            loss = model(batch[0], batch[1])
+            model.backward(loss)
+            model.step()
+
+
+@pytest.mark.skip(
+    "Skipping test for now - the context manager has an issue with ._initialized and .disabled - worked with older transformers probably because it was setting some flags with the same name"
+)
+class TestInitTransformers(DistributedTest):
+    world_size = 2
+
+    def test_pretrained_init(self):
+        lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=2)
+        quant_config = None
+
+        with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config):
+            model = transformers.AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-Llama-3")
+
+        injection_assert(model)
+
+    def test_config_init(self):
+        lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=2)
+        quant_config = None
+
+        config = transformers.AutoConfig.from_pretrained("llamafactory/tiny-random-Llama-3")
+
+        with deepspeed.linear.Init(lora_config=lora_config, quant_config=quant_config):
+            model = transformers.AutoModelForCausalLM.from_config(config)
+
+        injection_assert(model)

tests/unit/linear/test_linear.py

@@ -0,0 +1,123 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+import deepspeed
+import deepspeed.comm as dist
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.linear import OptimizedLinear, LoRAConfig, QuantizationConfig
+from unit.common import DistributedTest
+
+from deepspeed.ops.op_builder import FPQuantizerBuilder
+
+if not deepspeed.ops.__compatible_ops__[FPQuantizerBuilder.NAME]:
+    pytest.skip("FPQuantizer op is not available on this system", allow_module_level=True)
+
+
+class TestBasicLinear(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        lora_config = None
+        quantization_config = None
+
+        input_features = 64  # Number of input features
+        output_features = 64  # Number of output features
+        batch_size = 1  # Number of samples in a batch
+
+        linear_layer = OptimizedLinear(input_dim=input_features,
+                                       output_dim=output_features,
+                                       lora_config=lora_config,
+                                       quantization_config=quantization_config,
+                                       dtype=torch.bfloat16)
+
+        dummy_input = torch.rand(batch_size, input_features, dtype=torch.bfloat16)
+        output = linear_layer(dummy_input)
+        assert output.shape == (batch_size, output_features)
+
+
+@pytest.mark.parametrize("base_weight_sharding", [1, 2])
+class TestLoRALinear(DistributedTest):
+    world_size = 2
+
+    def test(self, base_weight_sharding):
+        rank = dist.get_rank()
+        quantization_config = None
+
+        input_features = 64  # Number of input features
+        output_features = 64  # Number of output features
+        batch_size = 5  # Number of samples in a batch
+
+        lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=base_weight_sharding)
+
+        linear_layer = OptimizedLinear(input_dim=input_features,
+                                       output_dim=output_features,
+                                       lora_config=lora_config,
+                                       quantization_config=quantization_config,
+                                       dtype=torch.bfloat16)
+        device = get_accelerator().current_device_name()
+        linear_layer = linear_layer.to(device)
+        if rank == 0:
+            for n, p in linear_layer.named_parameters():
+                print(f"{n}, {p.shape}")
+
+        dummy_input = torch.rand(batch_size, input_features, device=device, dtype=torch.bfloat16)
+
+        output = linear_layer(dummy_input)
+        assert output.shape == (batch_size, output_features)
+
+
+@pytest.mark.parametrize("q_bits", [8, 6])
+class TestQuantLinear(DistributedTest):
+    world_size = 2
+
+    def test(self, q_bits):
+        input_features = 64  # Number of input features
+        output_features = 64  # Number of output features
+        batch_size = 5  # Number of samples in a batch
+
+        lora_config = None
+        quantization_config = QuantizationConfig(q_bits=q_bits)
+        quantization_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype()
+
+        linear_layer = OptimizedLinear(input_dim=input_features,
+                                       output_dim=output_features,
+                                       lora_config=lora_config,
+                                       quantization_config=quantization_config,
+                                       dtype=torch.bfloat16)
+        device = get_accelerator().current_device_name()
+        linear_layer = linear_layer.to(device)
+        dummy_input = torch.rand([batch_size, input_features], device=device, dtype=torch.bfloat16)
+
+        output = linear_layer(dummy_input)
+        assert output.shape == (batch_size, output_features)
+
+
+@pytest.mark.parametrize("base_weight_sharding", [1, 2], ids=['bws1', 'bws2'])
+@pytest.mark.parametrize("q_bits", [8, 6], ids=['qbit8', 'qbit6'])
+class TestOptimizedLinear(DistributedTest):
+    world_size = 2
+
+    def test(self, base_weight_sharding, q_bits):
+        input_features = 64  # Number of input features
+        output_features = 64  # Number of output features
+        batch_size = 5  # Number of samples in a batch
+
+        lora_config = LoRAConfig(lora_r=16, lora_alpha=16, base_weight_sharding=base_weight_sharding)
+        quantization_config = QuantizationConfig(q_bits=q_bits)
+        quantization_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype()
+
+        linear_layer = OptimizedLinear(input_dim=input_features,
+                                       output_dim=output_features,
+                                       lora_config=lora_config,
+                                       quantization_config=quantization_config,
+                                       dtype=torch.bfloat16)
+        device = get_accelerator().current_device_name()
+        linear_layer = linear_layer.to(device)
+        dummy_input = torch.rand([batch_size, input_features], device=device, dtype=torch.bfloat16)
+        output = linear_layer(dummy_input)
+        assert output.shape == (batch_size, output_features)

tests/unit/linear/test_quant_param.py

@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+import deepspeed
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.linear.quantization import QuantizedParameter
+from deepspeed.linear.config import QuantizationConfig
+
+from deepspeed.ops.op_builder import FPQuantizerBuilder
+
+from unit.common import DistributedTest
+
+if not deepspeed.ops.__compatible_ops__[FPQuantizerBuilder.NAME]:
+    pytest.skip("FPQuantizer op is not available on this system", allow_module_level=True)
+
+
+class TestQuantParam(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.parametrize('dtype', [torch.half, torch.float])
+    def test_unsupported_dtypes(self, dtype):
+        device = get_accelerator().current_device_name()
+        data = torch.rand(5, 5, device='cpu', dtype=dtype)
+        qp = QuantizedParameter(data)
+        with pytest.raises(AssertionError):
+            qp.to(device)
+
+    def test_requires_grad(self):
+        data = torch.rand(5, 5, dtype=torch.bfloat16)
+        with pytest.raises(ValueError):
+            QuantizedParameter(data, requires_grad=True)
+
+    def test_move_to_accelerator(self):
+        device = get_accelerator().current_device()
+        data = torch.rand(5, 5, device='cpu', dtype=torch.bfloat16)
+        quantization_config = QuantizationConfig()
+        quantization_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype()
+        qp = QuantizedParameter(data, quantization_config=quantization_config)
+        assert qp.device == torch.device('cpu')
+        qp = qp.to(get_accelerator().current_device_name())
+        assert qp.device == torch.device(device)
+        assert qp.dtype == quantization_config.q_dtype
+
+    def test_hf_clone(self):
+        device = get_accelerator().current_device_name()
+        data = torch.rand(5, 5, device=device, dtype=torch.bfloat16)
+
+        quantization_config = QuantizationConfig(q_bits=6)
+        qp = QuantizedParameter(data, quantization_config=quantization_config)
+
+        # should be able to clone parameter via dict, HF expects this to work
+        qp_copy = QuantizedParameter(qp.data, **qp.__dict__)
+
+        assert all(qp.data == qp_copy.data)
+        assert qp.quantization_config == qp_copy.quantization_config

tests/unit/megatron_model.py

@@ -0,0 +1,103 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import os
+import sys
+import math
+
+from .common import get_test_path
+from deepspeed.pipe import PipelineModule, LayerSpec
+from deepspeed.accelerator import get_accelerator
+
+
+def get_megatron_version():
+    p = os.popen("pip list --format=columns | grep megatron-lm")
+    pip_list = p.read()
+    assert 'megatron-lm' in pip_list, 'Please install Megatron-LM before getting its version'
+    ver_str = pip_list.split()[1]
+    return float(ver_str[0])
+
+
+def get_gpt2_model(args_others, mp_size=1):
+    from megatron.model import GPT2Model
+    from megatron.initialize import initialize_megatron
+
+    args_defaults = {
+        'vocab_file': get_test_path('gpt2-vocab.json'),
+        'merge_file': get_test_path('gpt2-merges.txt'),
+        'tokenizer_type': 'GPT2BPETokenizer',
+    }
+
+    args_defaults.update(args_others)
+
+    # setting "make-vocab-size-divisible-by" to avoid word-embedding size change in resizing testing.
+    sys.argv.extend(['--model-parallel-size', str(mp_size), '--make-vocab-size-divisible-by', str(1)])
+
+    initialize_megatron(args_defaults=args_defaults, ignore_unknown_args=True)
+    model = GPT2Model(num_tokentypes=0, parallel_output=False)
+    model.to(get_accelerator().device_name())
+    from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
+    from megatron import mpu
+    i = get_accelerator().current_device_name()
+    model = torchDDP(model, device_ids=[i], output_device=i, process_group=mpu.get_data_parallel_group())
+
+    return model
+
+
+class MockGPT2ModelPipe(PipelineModule):
+
+    def __init__(self, num_layers, mp_size, args_others, topo, **kwargs):
+        from megatron.initialize import initialize_megatron
+
+        args_defaults = {
+            'vocab_file': get_test_path('gpt2-vocab.json'),
+            'merge_file': get_test_path('gpt2-merges.txt'),
+            'tokenizer_type': 'GPT2BPETokenizer',
+        }
+
+        args_defaults.update(args_others)
+
+        # setting "make-vocab-size-divisible-by" to avoid word-embedding size change in resizing testing.
+        sys.argv.extend(['--model-parallel-size', str(mp_size), '--make-vocab-size-divisible-by', str(1)])
+
+        initialize_megatron(args_defaults=args_defaults, ignore_unknown_args=True)
+
+        from megatron.model.transformer import ParallelTransformerLayer
+
+        class ParallelTransformerLayerPipe(ParallelTransformerLayer):
+
+            def forward(self, args):
+                # hardcode attn mask for testing, PP requires the attn_mask to be stashed
+                attention_mask = torch.tensor([[True]], device=get_accelerator().current_device_name())
+                return super().forward(args, attention_mask)
+
+        layers = []
+        for x in range(num_layers):
+            layers.append(
+                LayerSpec(ParallelTransformerLayerPipe, self.gpt2_attention_mask_func, self.init_method_normal(0.02),
+                          self.scaled_init_method_normal(0.02, num_layers), x))
+        super().__init__(layers=layers, loss_fn=torch.nn.CrossEntropyLoss(), topology=topo, **kwargs)
+
+    def gpt2_attention_mask_func(self, attention_scores, ltor_mask):
+        attention_scores.masked_fill_(ltor_mask, -10000.0)
+        return attention_scores
+
+    def init_method_normal(self, sigma):
+        """Init method based on N(0, sigma)."""
+
+        def init_(tensor):
+            return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)
+
+        return init_
+
+    def scaled_init_method_normal(self, sigma, num_layers):
+        """Init method based on N(0, sigma/sqrt(2*num_layers)."""
+        std = sigma / math.sqrt(2.0 * num_layers)
+
+        def init_(tensor):
+            return torch.nn.init.normal_(tensor, mean=0.0, std=std)
+
+        return init_

tests/unit/model_parallelism/test_autotp_custom_patterns.py

@@ -0,0 +1,569 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+import deepspeed.comm as dist
+import deepspeed
+from copy import deepcopy
+from torch import nn
+
+from unit.common import DistributedTest, preferred_dtype
+from deepspeed.accelerator import get_accelerator
+from deepspeed.utils import groups
+from deepspeed.module_inject.layers import (LinearAllreduce, LinearLayer, SubParamLinearLayer, fused_LinearLayer)
+from deepspeed.module_inject.autotp_config import AutoTPConfig
+from deepspeed.module_inject.auto_tp import AutoTP
+
+
+def skip_on_device():
+    if get_accelerator().device_name() == 'xpu':
+        pytest.skip("XPU requires a higher version for test")
+
+
+class SequentialLinearModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim, nlayers=1):
+        super(SequentialLinearModel, self).__init__()
+        self.linears = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim) for _ in range(nlayers)])
+
+    def forward(self, x):
+        for layer in self.linears:
+            x = layer(x)
+        return x
+
+
+class CustomLinearModule(torch.nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(CustomLinearModule, self).__init__()
+        self.weight = torch.nn.Parameter(torch.empty(hidden_dim, hidden_dim))
+        self.bias = torch.nn.Parameter(torch.empty(hidden_dim))
+        torch.nn.init.uniform_(self.weight, -0.02, 0.02)
+        torch.nn.init.uniform_(self.bias, -0.02, 0.02)
+
+    def forward(self, x):
+        return torch.matmul(x, self.weight.transpose(-1, -2)) + self.bias
+
+
+class CustomLinearModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(CustomLinearModel, self).__init__()
+        self.custom = CustomLinearModule(hidden_dim)
+
+    def forward(self, x):
+        return self.custom(x)
+
+
+class QKVLinearModule(torch.nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(QKVLinearModule, self).__init__()
+        self.qkv_proj = torch.nn.Linear(hidden_dim, hidden_dim * 3)
+
+    def forward(self, x):
+        return self.qkv_proj(x)
+
+
+class QKVLinearModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(QKVLinearModel, self).__init__()
+        self.self_attn = QKVLinearModule(hidden_dim)
+
+    def forward(self, x):
+        return self.self_attn(x)
+
+
+class DeepAttention(torch.nn.Module):
+    """Mimics HF attention module with separate projection layers."""
+
+    def __init__(self, hidden_dim):
+        super().__init__()
+        self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.o_proj = torch.nn.Linear(hidden_dim, hidden_dim)
+
+    def forward(self, x):
+        return self.o_proj(self.q_proj(x))
+
+
+class DeepBlock(torch.nn.Module):
+    """Mimics a single HF transformer block."""
+
+    def __init__(self, hidden_dim):
+        super().__init__()
+        self.self_attn = DeepAttention(hidden_dim)
+
+    def forward(self, x):
+        return self.self_attn(x)
+
+
+class DeepModel(torch.nn.Module):
+    """Mimics HF transformer structure: model.layers.[N].self_attn.{q,o}_proj.
+
+    This creates a 4-level-deep module hierarchy to test that _replace_module
+    correctly propagates the full module path during recursion.
+    """
+
+    def __init__(self, hidden_dim, nlayers=2):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([DeepBlock(hidden_dim) for _ in range(nlayers)])
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+def init_tp_engine(tp_size, partition_config=None):
+    config_dict = {
+        "train_micro_batch_size_per_gpu": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 1e-6
+            }
+        },
+        "tensor_parallel": {
+            "autotp_size": tp_size,
+        },
+        "zero_optimization": {
+            "stage": 0,
+        }
+    }
+    if partition_config is not None:
+        config_dict["tensor_parallel"]["partition_config"] = partition_config
+    else:
+        config_dict["tensor_parallel"]["partition_config"] = {
+            "use_default_specs": False,
+            "layer_specs": [{
+                "patterns": [".*\\.weight$"],
+                "partition_type": "skip",
+            }],
+        }
+    if preferred_dtype() is torch.float16:
+        config_dict["fp16"] = {"enabled": True}
+    elif preferred_dtype() is torch.bfloat16:
+        config_dict["bf16"] = {"enabled": True}
+
+    model = SequentialLinearModel(hidden_dim=8)
+    deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+
+def apply_autotp_with_partition_config(model, tp_size, partition_config):
+    groups._init_tp_mesh_device(tensor_model_parallel_size=tp_size)
+    autotp_config = AutoTPConfig.from_dict(partition_config)
+    autotp = AutoTP(module=model,
+                    all_reduce_linears=[],
+                    prefix="",
+                    state_dict=None,
+                    linear_layer_setting=None,
+                    orig_layer_impl=None,
+                    keep_module_on_host=False,
+                    partition_config=autotp_config)
+    autotp.set_tensor_parallel_config(tp_size, groups.get_tensor_model_parallel_group())
+    autotp.update_linear_policies()
+    autotp._replace_module(model)
+    return model
+
+
+def gather_subparam_output(output, subparam_sizes, mp_group):
+    tp_world_size = dist.get_world_size(group=mp_group)
+    local_sizes = [size // tp_world_size for size in subparam_sizes]
+    output_chunks = torch.split(output, local_sizes, dim=-1)
+    gathered_chunks = []
+    for chunk in output_chunks:
+        chunk = chunk.contiguous()
+        gathered = [torch.empty_like(chunk) for _ in range(tp_world_size)]
+        dist.all_gather(gathered, chunk, group=mp_group)
+        gathered_chunks.append(torch.cat(gathered, dim=-1))
+    return torch.cat(gathered_chunks, dim=-1)
+
+
+def assert_close_for_preferred_dtype(actual, expected):
+    atol = 1e-3
+    rtol = 2e-2
+    if preferred_dtype() is torch.float32:
+        atol = 1e-5
+        rtol = 1e-5
+    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+class TestAutoTPCustomPatterns(DistributedTest):
+    world_size = 2
+    reuse_dist_env = False
+
+    def test_custom_pattern_replacement(self):
+        skip_on_device()
+        partition_config = {
+            "use_default_specs":
+            False,
+            "layer_specs": [
+                {
+                    "patterns": [".*linears\\.0\\.weight$"],
+                    "partition_type": "row",
+                },
+                {
+                    "patterns": [".*linears\\.1\\.weight$"],
+                    "partition_type": "column",
+                },
+                {
+                    "patterns": [".*linears\\.2\\.weight$"],
+                    "partition_type": "skip",
+                },
+            ],
+        }
+        model = SequentialLinearModel(hidden_dim=16, nlayers=3)
+        model = apply_autotp_with_partition_config(model, tp_size=2, partition_config=partition_config)
+
+        assert isinstance(model.linears[0], LinearAllreduce)
+        assert isinstance(model.linears[1], LinearLayer)
+        assert isinstance(model.linears[2], nn.Linear)
+
+    def test_custom_patterns_applied_via_config(self):
+        skip_on_device()
+        partition_config = {
+            "use_default_specs":
+            False,
+            "layer_specs": [
+                {
+                    "patterns": [".*linears\\.0\\.weight$"],
+                    "partition_type": "row",
+                },
+                {
+                    "patterns": [".*linears\\.1\\.weight$"],
+                    "partition_type": "column",
+                },
+                {
+                    "patterns": [".*linears\\.2\\.weight$"],
+                    "partition_type": "skip",
+                },
+            ],
+        }
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": 2,
+                "partition_config": partition_config,
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        model = SequentialLinearModel(hidden_dim=16, nlayers=3)
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        assert isinstance(engine.module.linears[0], LinearAllreduce)
+        assert isinstance(engine.module.linears[1], LinearLayer)
+        assert isinstance(engine.module.linears[2], nn.Linear)
+
+    def test_use_default_specs_false_skips_unmatched_layers(self):
+        skip_on_device()
+        # Verify unmatched layers remain unsharded when defaults are disabled.
+        partition_config = {
+            "use_default_specs":
+            False,
+            "layer_specs": [
+                {
+                    "patterns": [".*linears\\.0\\.weight$"],
+                    "partition_type": "row",
+                },
+                {
+                    "patterns": [".*linears\\.1\\.weight$"],
+                    "partition_type": "column",
+                },
+            ],
+        }
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": 2,
+                "partition_config": partition_config,
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        model = SequentialLinearModel(hidden_dim=16, nlayers=3)
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        assert isinstance(engine.module.linears[0], LinearAllreduce)
+        assert isinstance(engine.module.linears[1], LinearLayer)
+        assert isinstance(engine.module.linears[2], nn.Linear)
+
+    def test_custom_module_replacement_with_patterns(self):
+        skip_on_device()
+        # Verify custom linear-like modules are partitioned via patterns.
+        partition_config = {
+            "use_default_specs": False,
+            "layer_specs": [
+                {
+                    "patterns": [".*custom\\.weight$"],
+                    "partition_type": "column",
+                },
+            ],
+        }
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": 2,
+                "partition_config": partition_config,
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        model = CustomLinearModel(hidden_dim=16)
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        assert isinstance(engine.module.custom, LinearLayer)
+
+    def test_custom_pattern_disables_fused_qkv_heuristic(self):
+        skip_on_device()
+        # Use a qkv_proj name that would trigger the fused-QKV heuristic, then
+        # verify custom patterns override that path and preserve correctness.
+        torch.manual_seed(1234)
+        hidden_dim = 16
+        qkv_sizes = (hidden_dim, hidden_dim, hidden_dim)
+        partition_config = {
+            "use_default_specs":
+            False,
+            "layer_specs": [
+                {
+                    "patterns": [".*self_attn\\.qkv_proj\\.weight$"],
+                    "partition_type": "column",
+                    "shape": [list(qkv_sizes), -1],
+                    "partition_dim": 0,
+                },
+            ],
+        }
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": 2,
+                "partition_config": partition_config,
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        model = QKVLinearModel(hidden_dim=hidden_dim)
+        baseline = deepcopy(model).to(get_accelerator().current_device(), dtype=preferred_dtype())
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        qkv_layer = engine.module.self_attn.qkv_proj
+        # Custom pattern should force SubParamLinearLayer (shape-based path),
+        # and avoid the legacy fused-QKV heuristic despite the qkv_proj name.
+        assert isinstance(qkv_layer, SubParamLinearLayer)
+        assert not isinstance(qkv_layer, fused_LinearLayer)
+
+        assert qkv_layer.partition_dim == 0
+        assert qkv_layer._subparam_sizes == qkv_sizes
+        assert qkv_layer._orig_weight_shape == (hidden_dim * 3, hidden_dim)
+
+        qkv_layer.gather_params([qkv_layer.weight, qkv_layer.bias])
+        torch.testing.assert_close(qkv_layer.weight, baseline.self_attn.qkv_proj.weight)
+        if qkv_layer.bias is not None:
+            torch.testing.assert_close(qkv_layer.bias, baseline.self_attn.qkv_proj.bias)
+
+        torch.manual_seed(4321)
+        inputs = torch.randn(2, hidden_dim, dtype=preferred_dtype(), device=get_accelerator().current_device())
+        full_output = baseline(inputs)
+        tp_output = engine.module(inputs)
+        assert_close_for_preferred_dtype(tp_output, full_output)
+
+    def test_first_match_precedence(self):
+        skip_on_device()
+        partition_config = {
+            "use_default_specs":
+            False,
+            "layer_specs": [
+                {
+                    "patterns": [".*linears\\.0\\.weight$"],
+                    "partition_type": "skip",
+                },
+                {
+                    "patterns": [".*linears\\.0\\.weight$"],
+                    "partition_type": "column",
+                },
+            ],
+        }
+        model = SequentialLinearModel(hidden_dim=16, nlayers=1)
+        model = apply_autotp_with_partition_config(model, tp_size=2, partition_config=partition_config)
+
+        assert isinstance(model.linears[0], nn.Linear)
+
+    def test_deep_model_full_path_propagation(self):
+        """Verify _replace_module propagates accumulated paths through deep hierarchies.
+
+        Uses a 4-level-deep model (layers.N.self_attn.{q,o}_proj) with patterns
+        that require intermediate path components (layers.N). Without correct
+        full_name propagation, the recursive path is truncated and patterns
+        that include intermediate levels will silently fail to match.
+        """
+        skip_on_device()
+        partition_config = {
+            "use_default_specs":
+            False,
+            "layer_specs": [
+                {
+                    "patterns": [r".*layers\.\d+\.self_attn\.q_proj\.weight$"],
+                    "partition_type": "column",
+                },
+                {
+                    "patterns": [r".*layers\.\d+\.self_attn\.o_proj\.weight$"],
+                    "partition_type": "row",
+                },
+            ],
+        }
+        model = DeepModel(hidden_dim=16, nlayers=2)
+        model = apply_autotp_with_partition_config(model, tp_size=2, partition_config=partition_config)
+
+        # All 4 projections (2 layers x {q_proj, o_proj}) must be replaced.
+        # Before the full_name fix, 0 modules were replaced because the mangled
+        # path "self_attn.q_proj.weight" could not match "layers.N.self_attn...".
+        for i in range(2):
+            assert isinstance(model.layers[i].self_attn.q_proj, LinearLayer), \
+                f"layers.{i}.self_attn.q_proj was not replaced (path propagation bug?)"
+            assert isinstance(model.layers[i].self_attn.o_proj, LinearAllreduce), \
+                f"layers.{i}.self_attn.o_proj was not replaced (path propagation bug?)"
+
+
+def test_invalid_custom_shape_rejected():
+    bad_config = {
+        "layer_specs": [{
+            "patterns": [".*"],
+            "partition_type": "column",
+            "shape": [2, [1, 1]],
+        }]
+    }
+    with pytest.raises(ValueError, match="nested tuple only allowed at partition_dim"):
+        AutoTPConfig.from_dict(bad_config)
+
+
+class TestAutoTPFusedWeights(DistributedTest):
+    world_size = 2
+    reuse_dist_env = False
+
+    def test_gate_up_fused_weight_partition(self):
+        skip_on_device()
+        init_tp_engine(tp_size=2)
+
+        hidden_dim = 8
+        torch.manual_seed(42)
+        linear = nn.Linear(hidden_dim,
+                           hidden_dim * 2,
+                           bias=True,
+                           dtype=preferred_dtype(),
+                           device=get_accelerator().current_device())
+        full_weight = deepcopy(linear.weight.data)
+        full_bias = deepcopy(linear.bias.data)
+
+        layer = SubParamLinearLayer(deepcopy(linear),
+                                    groups.get_tensor_model_parallel_group(),
+                                    shape=(2, -1),
+                                    partition_dim=0,
+                                    name="mlp.gate_up_proj")
+        assert layer._subparam_sizes == (hidden_dim, hidden_dim)
+        assert layer.weight.shape == (hidden_dim, hidden_dim)
+
+        layer.gather_params([layer.weight, layer.bias])
+        torch.testing.assert_close(layer.weight.data, full_weight)
+        torch.testing.assert_close(layer.bias.data, full_bias)
+
+    def test_gqa_uneven_qkv_fused_weight_partition(self):
+        skip_on_device()
+        init_tp_engine(tp_size=2)
+
+        hidden_dim = 8
+        q_size, k_size, v_size = 8, 4, 4
+        torch.manual_seed(123)
+        linear = nn.Linear(hidden_dim,
+                           q_size + k_size + v_size,
+                           bias=True,
+                           dtype=preferred_dtype(),
+                           device=get_accelerator().current_device())
+        full_weight = deepcopy(linear.weight.data)
+        full_bias = deepcopy(linear.bias.data)
+
+        layer = SubParamLinearLayer(deepcopy(linear),
+                                    groups.get_tensor_model_parallel_group(),
+                                    shape=((q_size, k_size, v_size), -1),
+                                    partition_dim=0,
+                                    name="self_attn.qkv_proj")
+        assert layer._subparam_sizes == (q_size, k_size, v_size)
+        assert layer.weight.shape == ((q_size + k_size + v_size) // 2, hidden_dim)
+
+        layer.gather_params([layer.weight, layer.bias])
+        torch.testing.assert_close(layer.weight.data, full_weight)
+        torch.testing.assert_close(layer.bias.data, full_bias)
+
+    def test_gqa_uneven_qkv_fused_forward(self):
+        skip_on_device()
+        groups._init_tp_mesh_device(tensor_model_parallel_size=2)
+
+        hidden_dim = 8
+        q_size, k_size, v_size = 8, 4, 4
+        torch.manual_seed(321)
+        linear = nn.Linear(hidden_dim,
+                           q_size + k_size + v_size,
+                           bias=True,
+                           dtype=preferred_dtype(),
+                           device=get_accelerator().current_device())
+        layer = SubParamLinearLayer(deepcopy(linear),
+                                    groups.get_tensor_model_parallel_group(),
+                                    shape=((q_size, k_size, v_size), -1),
+                                    partition_dim=0,
+                                    name="self_attn.qkv_proj")
+
+        torch.manual_seed(42)
+        inputs = torch.randn(2, hidden_dim, dtype=preferred_dtype(), device=get_accelerator().current_device())
+        full_output = linear(inputs)
+        tp_output = layer(inputs)
+
+        gathered_output = gather_subparam_output(tp_output, (q_size, k_size, v_size),
+                                                 groups.get_tensor_model_parallel_group())
+        assert_close_for_preferred_dtype(gathered_output, full_output)

tests/unit/model_parallelism/test_autotp_training.py

@@ -0,0 +1,822 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import deepspeed.comm as dist
+import torch
+import math
+from copy import deepcopy
+
+from unit.common import DistributedTest, preferred_dtype
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+from unit.simple_model import SimpleModel, random_dataloader
+from deepspeed.utils import groups
+from contextlib import contextmanager
+from torch import nn
+from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, set_autotp_mode, is_autotp_training_mode
+from unit.checkpoint.common import compare_lr_scheduler_states, compare_optimizer_states
+import os
+from deepspeed.runtime.utils import is_model_parallel_parameter
+
+
+def skip_on_device():
+    if get_accelerator().device_name() == 'xpu':
+        pytest.skip("XPU requires a higher version for test")
+
+
+def reset_tp_model_init_state():
+    deepspeed._TP_MODEL_INIT_ARGS = None
+    set_autotp_mode(training=False)
+
+
+def _reset_tp_groups(monkeypatch):
+    monkeypatch.setattr(groups, "_DATA_PARALLEL_GROUP", None)
+    monkeypatch.setattr(groups, "_MODEL_PARALLEL_GROUP", None)
+    monkeypatch.setattr(groups, "_TENSOR_MODEL_PARALLEL_GROUP", None)
+
+
+def _patch_tp_group_creation(monkeypatch, *, rank=2, initialize_mesh_device=None):
+    new_group_calls = []
+
+    def fake_new_group(ranks):
+        ranks = tuple(ranks)
+        new_group_calls.append(ranks)
+        return ranks
+
+    _reset_tp_groups(monkeypatch)
+    monkeypatch.setattr(groups.dist, "get_world_size", lambda group=None: 4)
+    monkeypatch.setattr(groups.dist, "get_rank", lambda group=None: rank)
+    monkeypatch.setattr(groups.dist, "new_group", fake_new_group)
+    monkeypatch.setattr(groups, "log_dist", lambda *args, **kwargs: None)
+    if initialize_mesh_device is not None:
+        monkeypatch.setattr(groups.dist, "initialize_mesh_device", initialize_mesh_device)
+
+    return new_group_calls
+
+
+def test_init_tp_mesh_device_debug_detail_uses_explicit_groups(monkeypatch):
+
+    def fail_initialize_mesh_device(*args, **kwargs):
+        raise AssertionError("DeviceMesh should be skipped when TORCH_DISTRIBUTED_DEBUG=DETAIL")
+
+    new_group_calls = _patch_tp_group_creation(monkeypatch, initialize_mesh_device=fail_initialize_mesh_device)
+    monkeypatch.setenv("TORCH_DISTRIBUTED_DEBUG", "DETAIL")
+
+    data_parallel_group, tensor_parallel_group = groups._init_tp_mesh_device(tensor_model_parallel_size=2)
+
+    assert new_group_calls == [(0, 2), (1, 3), (0, 1), (2, 3)]
+    assert data_parallel_group == (0, 2)
+    assert tensor_parallel_group == (2, 3)
+    assert groups.get_data_parallel_group() == (0, 2)
+    assert groups.get_tensor_model_parallel_group() == (2, 3)
+
+
+def test_init_tp_mesh_device_split_error_falls_back_to_explicit_groups(monkeypatch):
+
+    def raise_split_error(*args, **kwargs):
+        raise RuntimeError(groups._DEVICE_MESH_SPLIT_UNSUPPORTED)
+
+    new_group_calls = _patch_tp_group_creation(monkeypatch, initialize_mesh_device=raise_split_error)
+    monkeypatch.delenv("TORCH_DISTRIBUTED_DEBUG", raising=False)
+
+    data_parallel_group, tensor_parallel_group = groups._init_tp_mesh_device(tensor_model_parallel_size=2)
+
+    assert new_group_calls == [(0, 2), (1, 3), (0, 1), (2, 3)]
+    assert data_parallel_group == (0, 2)
+    assert tensor_parallel_group == (2, 3)
+    assert groups.get_data_parallel_group() == (0, 2)
+    assert groups.get_tensor_model_parallel_group() == (2, 3)
+
+
+class DummyMPU:
+
+    def __init__(self, tp_world_size=1):
+        self.rank = dist.get_rank()
+        self.world_size = dist.get_world_size()
+        self.tp_world_size = tp_world_size
+        self.dp_group = dist.get_world_group()
+        self.tp_group = dist.get_world_group()
+
+    def get_model_parallel_rank(self):
+        return self.rank % self.tp_world_size
+
+    def get_model_parallel_world_size(self):
+        return self.tp_world_size
+
+    def get_data_parallel_rank(self):
+        return self.rank // self.tp_world_size
+
+    def get_data_parallel_world_size(self):
+        return self.world_size // self.tp_world_size
+
+    def get_data_parallel_group(self):
+        return self.dp_group
+
+    def get_model_parallel_group(self):
+        return self.tp_group
+
+
+class SequentialLinearModel(torch.nn.Module):
+
+    def __init__(self, hidden_dim, empty_grad=False, nlayers=1):
+        super(SequentialLinearModel, self).__init__()
+        self.linears = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim) for _ in range(nlayers)])
+        if empty_grad:
+            self.linear2 = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
+        self.empty_grad = empty_grad
+
+    def forward(self, x, y):
+        if len(self.linears) == 1:
+            x = self.linears[0](x)
+        else:
+            for i, l in enumerate(self.linears):
+                x = self.linears[i](x)
+        return self.cross_entropy_loss(x, y)
+
+
+@contextmanager
+def should_assert_with_msg(expected_message):
+    try:
+        yield
+    except AssertionError as e:
+        if dist.get_rank() == 0:
+            print(expected_message)
+            print(str(e))
+        if str(e) == expected_message:
+            pass
+        else:
+            raise e
+    else:
+        raise AssertionError(f"Expected AssertionError with message '{expected_message}' "
+                             "but no exception was raised.")
+
+
+@pytest.mark.parametrize("tp_size", [2, 4])
+class TestTpParallelStates(DistributedTest):
+    world_size = 4
+
+    def test(self, tp_size: int):
+        skip_on_device()
+        dp_size = 4 / tp_size
+        hidden_dim = 128
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "tensor_parallel": {
+                "autotp_size": tp_size,
+                "partition_config": {
+                    "use_default_specs": False,
+                    "layer_specs": [{
+                        "patterns": [".*\\.weight$"],
+                        "partition_type": "skip",
+                    }],
+                }
+            },
+            "zero_optimization": {
+                "stage": 0
+            }
+        }
+        model = SimpleModel(hidden_dim=hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        assert groups.get_tensor_model_parallel_world_size() == tp_size
+        assert groups.get_data_parallel_world_size() == dp_size
+
+
+class TestTpModelInitCompatibility(DistributedTest):
+    world_size = 4
+    reuse_dist_env = False
+
+    def test_tp_model_init_merges_config(self):
+        skip_on_device()
+        reset_tp_model_init_state()
+        model = SimpleModel(hidden_dim=8)
+        deepspeed.tp_model_init(model, tp_size=1, dtype=preferred_dtype())
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        engine, _, _, _ = deepspeed.initialize(model=model,
+                                               model_parameters=model.parameters(),
+                                               config=config_dict,
+                                               mpu=DummyMPU())
+        assert engine.autotp_size() == 1
+        assert is_autotp_training_mode()
+
+    def test_tp_model_init_config_autotp_size_mismatch(self):
+        skip_on_device()
+        reset_tp_model_init_state()
+        model = SimpleModel(hidden_dim=8)
+        deepspeed.tp_model_init(model, tp_size=1, dtype=preferred_dtype())
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "tensor_parallel": {
+                "autotp_size": 2,
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        with pytest.raises(ValueError, match="tensor_parallel.autotp_size"):
+            deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict, mpu=DummyMPU())
+
+    def test_tp_model_init_autocreates_tp_group(self):
+        skip_on_device()
+        reset_tp_model_init_state()
+        # Verify tp_model_init creates TP groups when no mpu is provided.
+        model = SimpleModel(hidden_dim=8)
+        tp_size = 2
+        deepspeed.tp_model_init(model, tp_size=tp_size, dtype=preferred_dtype())
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "tensor_parallel": {
+                "partition_config": {
+                    "use_default_specs": False,
+                    "layer_specs": [{
+                        "patterns": [".*\\.weight$"],
+                        "partition_type": "skip",
+                    }],
+                }
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        assert engine.autotp_size() == tp_size
+        assert groups.get_tensor_model_parallel_world_size() == tp_size
+        assert groups.get_data_parallel_world_size() == dist.get_world_size() // tp_size
+
+    def test_tp_model_init_tp_group_rejects_mpu(self):
+        skip_on_device()
+        reset_tp_model_init_state()
+        model = SimpleModel(hidden_dim=8)
+        tp_group = dist.new_group(ranks=[0])
+        deepspeed.tp_model_init(model, tp_size=1, dtype=preferred_dtype(), tp_group=tp_group)
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        with pytest.raises(ValueError, match="tp_model_init provided tp_group"):
+            deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict, mpu=DummyMPU())
+
+    def test_tp_model_init_dtype_mismatch(self):
+        skip_on_device()
+        reset_tp_model_init_state()
+        model = SimpleModel(hidden_dim=8)
+        deepspeed.tp_model_init(model, tp_size=1, dtype=torch.float16)
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "bf16": {
+                "enabled": True,
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        with pytest.raises(ValueError, match="Conflicting dtype"):
+            deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict, mpu=DummyMPU())
+
+    @pytest.mark.sequential
+    @pytest.mark.parametrize("tp_size", [2, 4])
+    @pytest.mark.parametrize("tp_overlap_comm", [True, False])
+    def test_tp_model_init_row_parallel(self, tp_size: int, tp_overlap_comm: bool):
+        run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=False, use_tp_model_init=True)
+
+    @pytest.mark.sequential
+    @pytest.mark.parametrize("tp_size", [2, 4])
+    @pytest.mark.parametrize("tp_overlap_comm", [True, False])
+    def test_tp_model_init_column_parallel(self, tp_size: int, tp_overlap_comm: bool):
+        run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=True, use_tp_model_init=True)
+
+
+@pytest.mark.parametrize("tp_size", [2, 4])
+class TestTpDataloaderCorrectness(DistributedTest):
+    world_size = 4
+    reuse_dist_env = False
+
+    def test(self, tp_size: int):
+        skip_on_device()
+        hidden_dim = 128
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": tp_size,
+                "partition_config": {
+                    "use_default_specs": False,
+                    "layer_specs": [{
+                        "patterns": [".*\\.weight$"],
+                        "partition_type": "skip",
+                    }],
+                }
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        model = SimpleModel(hidden_dim=hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        torch.manual_seed(42)
+
+        data_loader = random_dataloader(model=model,
+                                        total_samples=3,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=preferred_dtype())
+        dist.barrier()
+        with should_assert_with_msg(
+                "Data inconsistency within the TP group. Please check the Dataloader implementation to ensure consistency."
+        ):
+            for batch in data_loader:
+                # batch[0].requires_grad = requires_grad
+                batch[0] += dist.get_rank()
+                model(batch[0], batch[1])
+
+        model = SimpleModel(hidden_dim=hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+        data_loader = random_dataloader(model=model,
+                                        total_samples=3,
+                                        hidden_dim=hidden_dim,
+                                        device=model.device,
+                                        dtype=preferred_dtype())
+        for batch in data_loader:
+            dist.broadcast(batch[0],
+                           src=groups.get_tensor_model_parallel_src_rank(),
+                           group=groups.get_tensor_model_parallel_group())
+            dist.broadcast(batch[1],
+                           src=groups.get_tensor_model_parallel_src_rank(),
+                           group=groups.get_tensor_model_parallel_group())
+            model(batch[0], batch[1])
+
+
+def process_linear_layer(hidden_dim, input):
+    torch.manual_seed(42)
+    torch_linear = nn.Linear(hidden_dim,
+                             hidden_dim,
+                             dtype=preferred_dtype(),
+                             device=get_accelerator().current_device())
+    torch_out = torch_linear(input)
+    torch_loss = torch_out.sum()
+    torch_loss.backward()
+    return torch_linear, torch_out
+
+
+def run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel, use_tp_model_init=False):
+    skip_on_device()
+    hidden_dim = 128
+    batch_size_per_device = 1
+    config_dict = {
+        "train_micro_batch_size_per_gpu": 1,
+        "steps_per_print": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 1e-6
+            }
+        },
+        "tensor_parallel": {
+            "autotp_size": tp_size,
+            "tp_overlap_comm": tp_overlap_comm
+        },
+        "zero_optimization": {
+            "stage": 0,
+        }
+    }
+    partition_type = "column" if column_parallel else "row"
+    config_dict["tensor_parallel"]["partition_config"] = {
+        "use_default_specs": False,
+        "layer_specs": [{
+            "patterns": [".*\\.weight$"],
+            "partition_type": partition_type,
+        }],
+    }
+    if preferred_dtype() is torch.float16:
+        config_dict["fp16"] = {"enabled": True}
+    elif preferred_dtype() is torch.bfloat16:
+        config_dict["bf16"] = {"enabled": True}
+
+    model = SequentialLinearModel(hidden_dim=hidden_dim)
+    if use_tp_model_init:
+        reset_tp_model_init_state()
+        deepspeed.tp_model_init(model, tp_size=tp_size, dtype=preferred_dtype())
+        mpu = DummyMPU(tp_world_size=tp_size)
+        model, _, _, _ = deepspeed.initialize(model=model,
+                                              model_parameters=model.parameters(),
+                                              config=config_dict,
+                                              mpu=mpu)
+    else:
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+    input = torch.randn(batch_size_per_device,
+                        hidden_dim,
+                        dtype=preferred_dtype(),
+                        requires_grad=True,
+                        device=get_accelerator().current_device())
+    dist.broadcast(input, groups.get_tensor_model_parallel_src_rank(), group=groups.get_tensor_model_parallel_group())
+
+    # Note: correctness checks below use standalone TP wrappers and do not
+    # rely on the model's AutoTP-partitioned parameters.
+    torch_linear, torch_out = process_linear_layer(hidden_dim, input)
+    if column_parallel:
+        linear = LinearLayer(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
+        out = linear(input.to(get_accelerator().current_device()))
+        loss = out.sum()
+        loss.backward()
+
+        cur_device_out = torch.chunk(torch_out, tp_size, dim=-1)[groups.get_tensor_model_parallel_rank()]
+        torch_grad = torch.chunk(torch_linear.weight.grad, tp_size, dim=0)[groups.get_tensor_model_parallel_rank()]
+        torch_bias_grad = torch.chunk(torch_linear.bias.grad, tp_size, dim=0)[groups.get_tensor_model_parallel_rank()]
+
+        torch.testing.assert_close(linear.bias.grad,
+                                   torch_bias_grad.to(get_accelerator().current_device()),
+                                   atol=1e-3,
+                                   rtol=1e-3)
+        torch.testing.assert_close(linear.weight.grad,
+                                   torch_grad.to(get_accelerator().current_device()),
+                                   atol=1e-3,
+                                   rtol=1e-3)
+        torch.testing.assert_close(cur_device_out.to(get_accelerator().current_device()).contiguous(),
+                                   out.contiguous(),
+                                   atol=1e-2,
+                                   rtol=1e-2)
+    else:
+        linear = LinearAllreduce(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
+        input_ = torch.chunk(input, tp_size, dim=-1)[groups.get_tensor_model_parallel_rank()]
+        out = linear(input_.to(get_accelerator().current_device()))
+        loss = out.sum()
+        loss.backward()
+
+        torch_grad = torch.chunk(torch_linear.weight.grad, tp_size, dim=1)[groups.get_tensor_model_parallel_rank()]
+        torch_bias_grad = torch_linear.bias.grad
+        torch.testing.assert_close(linear.bias.grad,
+                                   torch_bias_grad.to(get_accelerator().current_device()),
+                                   atol=1e-3,
+                                   rtol=1e-3)
+        torch.testing.assert_close(linear.weight.grad,
+                                   torch_grad.to(get_accelerator().current_device()),
+                                   atol=1e-3,
+                                   rtol=1e-3)
+        torch.testing.assert_close(out, torch_out.to(get_accelerator().current_device()), atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.sequential
+@pytest.mark.parametrize("tp_size", [2, 4])
+@pytest.mark.parametrize("tp_overlap_comm", [True, False])
+class TestTpLayerFwdBwd(DistributedTest):
+    world_size = 4
+    reuse_dist_env = False
+
+    def testRowParallel(self, tp_size: int, tp_overlap_comm: bool):
+        run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=False)
+
+    def testColumnParallel(self, tp_size: int, tp_overlap_comm: bool):
+        run_tp_layer_fwd_bwd(tp_size, tp_overlap_comm, column_parallel=True)
+
+
+# @pytest.mark.sequential
+class TestParamsGather(DistributedTest):
+    world_size = 4
+    reuse_dist_env = False
+
+    @pytest.mark.parametrize("layer_type", ["linear", "linearallreduce"])
+    def test(self, layer_type):
+        skip_on_device()
+        tp_size = 4
+        hidden_dim = 128
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": tp_size,
+                "partition_config": {
+                    "use_default_specs": False,
+                    "layer_specs": [{
+                        "patterns": [".*\\.weight$"],
+                        "partition_type": "skip",
+                    }],
+                }
+            },
+            "zero_optimization": {
+                "stage": 0,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        torch.manual_seed(42)
+        model = SequentialLinearModel(hidden_dim=hidden_dim)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        torch_linear = nn.Linear(hidden_dim, hidden_dim, dtype=preferred_dtype(), device="cpu")
+        total_params = sum(p.numel() for p in torch_linear.parameters())
+        tp_layer = None
+        if layer_type == "linear":
+            tp_layer = LinearLayer(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
+        elif layer_type == "linearallreduce":
+            tp_layer = LinearAllreduce(deepcopy(torch_linear), groups.get_tensor_model_parallel_group())
+        else:
+            raise ValueError(f"Invalid linear type: {config_dict['linear_type']}")
+
+        tp_params = sum(p.numel() for p in tp_layer.parameters())
+
+        expected_tp_params = 0
+        # compute expected TP params:
+        # - column-parallel (LinearLayer): weight & bias both split => total // tp_size
+        # - row-parallel    (LinearAllreduce): weight split, bias (1d tensors) replicated
+        if layer_type == "linearallreduce":
+            weight_params = torch_linear.weight.numel()
+            bias_params = torch_linear.bias.numel()
+            expected_tp_params = weight_params // tp_size + bias_params
+        else:
+            expected_tp_params = total_params // tp_size
+        assert expected_tp_params == tp_params, (
+            f"{layer_type}: expected {expected_tp_params} tp params, got {tp_params}")
+
+        for name, param in tp_layer.named_parameters(recurse=False):
+            if is_model_parallel_parameter(param):
+                param.gather_params([param])
+
+        torch_linear = torch_linear.to(get_accelerator().current_device())
+        is_same_weights = all(
+            torch.equal(param1, param2) for param1, param2 in zip(tp_layer.parameters(), torch_linear.parameters()))
+
+        assert is_same_weights
+
+        params1 = sum(p.numel() for p in tp_layer.parameters())
+        assert total_params == params1
+
+        for name, param in tp_layer.named_parameters(recurse=False):
+            if is_model_parallel_parameter(param):
+                param._tp_partition([param])
+
+        tp_params2 = sum(p.numel() for p in tp_layer.parameters())
+
+        assert expected_tp_params == tp_params2
+
+
+def dummy_init_engine(config):
+    # This is a dummy initialization function for the DeepSpeed engine.
+    # We only need to use the config to initialize the distributed settings for the test.
+    # Add default partition_config for simple test models if not provided
+    if "tensor_parallel" in config and "partition_config" not in config["tensor_parallel"]:
+        config["tensor_parallel"]["partition_config"] = {
+            "use_default_specs": False,
+            "layer_specs": [{
+                "patterns": [".*\\.weight$"],
+                "partition_type": "skip",
+            }],
+        }
+    model = SequentialLinearModel(hidden_dim=8)
+    model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
+
+
+def prepare_tp_model(hidden_dim, nlayers, linear_indices, allreduce_indices, group, return_global_copy=False):
+    model = SequentialLinearModel(hidden_dim=hidden_dim, nlayers=nlayers).to(preferred_dtype())
+    base_model = None
+    if return_global_copy:
+        base_model = deepcopy(model)
+    for i in linear_indices:
+        layer = LinearLayer(model.linears[i], group)
+        model.linears[i] = layer
+
+    for i in allreduce_indices:
+        layer = LinearAllreduce(model.linears[i], group)
+        model.linears[i] = layer
+
+    return model, base_model
+
+
+@pytest.mark.parametrize("zero_stage", [0, 1, 2])
+@pytest.mark.parametrize("tp_size", [2, 4])
+class TestSave(DistributedTest):
+
+    world_size = 4
+    reuse_dist_env = False
+
+    def test_save_original_weight(self, tp_size: int, zero_stage: int):
+        skip_on_device()
+        hidden_dim = 64
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": tp_size
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+        dummy_init_engine(config_dict)
+        torch.manual_seed(42)
+
+        model, base_model = prepare_tp_model(hidden_dim,
+                                             8, [2, 5], [3, 6],
+                                             groups.get_tensor_model_parallel_group(),
+                                             return_global_copy=True)
+        model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+
+        cur_params_numel = sum(p.numel() for p in model.parameters())
+        base_params_numel = sum(p.numel() for p in base_model.parameters())
+        assert cur_params_numel < base_params_numel
+
+        tp_state_dict = model._consolidated_16bit_state_dict()
+
+        def compare_state_dicts(state_dict1, state_dict2):
+            if state_dict1.keys() != state_dict2.keys():
+                print("The state_dicts have different keys!")
+                return False
+
+            for key in state_dict1:
+                if not torch.allclose(state_dict1[key], state_dict2[key], atol=1e-3):
+                    assert state_dict1[key].device == "cpu"
+                    print(f"Parameters for {key} are different!")
+                    return False
+
+            return True
+
+        base_state_dict = base_model.state_dict()
+        if dist.get_rank() == 0:
+            # we should consider the case when zero3 is used in the future.
+            assert compare_state_dicts(base_state_dict, tp_state_dict), "State_dict is not the same!"
+        else:
+            assert tp_state_dict is None, "noly rank0 should have the state_dict"
+
+    def test_ckpt_save(self, tmpdir, tp_size: int, zero_stage: int):
+        skip_on_device()
+        hidden_dim = 64
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            },
+            "tensor_parallel": {
+                "autotp_size": tp_size
+            },
+            "scheduler": {
+                "type": "WarmupLR",
+                "params": {
+                    "warmup_min_lr": 0,
+                    "warmup_max_lr": 0.001,
+                    "warmup_num_steps": 1000
+                }
+            }
+        }
+
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            config_dict["bf16"] = {"enabled": True}
+
+        dummy_init_engine(config_dict)
+
+        trained_model, _ = prepare_tp_model(hidden_dim, 8, [2, 5], [3, 6], groups.get_tensor_model_parallel_group())
+        loaded_model, _ = prepare_tp_model(hidden_dim, 8, [2, 5], [3, 6], groups.get_tensor_model_parallel_group())
+
+        trained_model, _, _, _ = deepspeed.initialize(model=trained_model,
+                                                      model_parameters=trained_model.parameters(),
+                                                      config=config_dict)
+        torch.manual_seed(42)
+
+        data_loader = random_dataloader(model=trained_model,
+                                        total_samples=3,
+                                        hidden_dim=hidden_dim,
+                                        device=trained_model.device,
+                                        dtype=preferred_dtype())
+        ckpt_path = os.path.join(tmpdir, 'tp_saved_checkpoint')
+        for i, batch in enumerate(data_loader):
+            batch[0].requires_grad = True
+            loss = trained_model(batch[0], batch[1])
+            loss = loss
+            trained_model.backward(loss)
+            trained_model.step()
+        trained_model.save_checkpoint(ckpt_path)
+
+        loaded_model, _, _, _ = deepspeed.initialize(model=loaded_model,
+                                                     model_parameters=loaded_model.parameters(),
+                                                     config=config_dict)
+        loaded_model.load_checkpoint(ckpt_path, load_optimizer_states=True, load_lr_scheduler_states=True)
+        compare_optimizer_states(trained_model, loaded_model, hidden_dim, fp16=(preferred_dtype() == torch.float16))
+        compare_lr_scheduler_states(trained_model, loaded_model)
+
+
+@pytest.mark.parametrize("zero_stage", [0, 1, 2])
+@pytest.mark.parametrize("tp_size", [2, 4])
+class TestTpGradNorm(DistributedTest):
+
+    world_size = 4
+    reuse_dist_env = False
+
+    def test(self, tp_size: int, zero_stage: int):
+        skip_on_device()
+        hidden_dim = 64
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": tp_size
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+            }
+        }
+        if preferred_dtype() is torch.float16:
+            config_dict["fp16"] = {"enabled": True}
+        elif preferred_dtype() is torch.bfloat16:
+            if zero_stage == 0:
+                pytest.skip(
+                    "This test has an overflow data and needs to implement an overflow skip mechanism in BF16_optimizer"
+                )
+            config_dict["bf16"] = {"enabled": True}
+
+        torch.manual_seed(42)
+
+        dummy_init_engine(config=config_dict)
+        tp_model, base_model = prepare_tp_model(hidden_dim,
+                                                8, [2, 5], [3, 6],
+                                                groups.get_tensor_model_parallel_group(),
+                                                return_global_copy=True)
+
+        base_model, base_optimizer, _, _ = deepspeed.initialize(model=base_model,
+                                                                model_parameters=base_model.parameters(),
+                                                                config=config_dict)
+        data_loader = random_dataloader(model=base_model,
+                                        total_samples=20,
+                                        hidden_dim=hidden_dim,
+                                        device=base_model.device,
+                                        dtype=preferred_dtype())
+
+        for i, batch in enumerate(data_loader):
+            batch[0].requires_grad = True
+            loss = base_model(batch[0], batch[1])
+            loss = loss
+            base_model.backward(loss)
+            base_model.step()
+
+        base_norm = base_optimizer._global_grad_norm
+
+        base_model.destroy()
+
+        tp_model, tp_optimizer, _, _ = deepspeed.initialize(model=tp_model,
+                                                            model_parameters=tp_model.parameters(),
+                                                            config=config_dict)
+        for i, batch in enumerate(data_loader):
+            batch[0].requires_grad = True
+            loss = tp_model(batch[0], batch[1])
+            loss = loss
+            tp_model.backward(loss)
+            tp_model.step()
+
+        tp_norm = tp_optimizer._global_grad_norm
+
+        assert math.isclose(base_norm, tp_norm, abs_tol=1e-3), f"base_norm: {base_norm}, tp_norm: {tp_norm}"
+        tp_params_numel = sum(p.numel() for p in tp_model.parameters())
+        base_params_numel = sum(p.numel() for p in base_model.parameters())
+        assert tp_params_numel < base_params_numel, f"tp_params_numel: {tp_params_numel}, base_params_numel: {base_params_numel}"

tests/unit/model_parallelism/test_configurable_parallel_mp.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import torch
+import deepspeed
+import pytest
+import random
+import numpy as np
+import deepspeed.comm as dist
+from deepspeed.accelerator import get_accelerator
+from unit.common import DistributedTest, DistributedFixture
+from unit.megatron_model import get_gpt2_model, get_megatron_version
+from deepspeed.utils.torch import required_torch_version
+
+pytestmark = pytest.mark.skipif(not required_torch_version(min_version=1.5, max_version=1.13),
+                                reason='Megatron-LM package requires Pytorch version >=1.5 and <=1.13')
+
+
+# TODO: integrated testing of TP and ZeRO 1/2/3
+def get_deepspeed_model(model):
+    ds_config_dict = {
+        "train_micro_batch_size_per_gpu": 1,
+        "optimizer": {
+            "type": "Lamb",
+            "params": {
+                "lr": 0.00015
+            }
+        },
+    }
+
+    from megatron import mpu
+    model, _, _, _ = deepspeed.initialize(model=model,
+                                          mpu=mpu,
+                                          model_parameters=model.parameters(),
+                                          config=ds_config_dict)
+    return model
+
+
+class ConfigurableMP(DistributedTest):
+
+    @pytest.fixture(autouse=True)
+    def reset_random(self, seed=1234):
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        get_accelerator().manual_seed_all(seed)
+
+    @pytest.fixture
+    def inputs(self, bs=1, seq_len=20):
+        input_ids = torch.randint(low=0, high=1000, size=(bs, seq_len))
+        position_ids = torch.randint(low=0, high=2, size=(bs, seq_len))
+        attention_mask = torch.randint(low=0, high=2, size=(bs, seq_len), dtype=torch.bool)
+        return [input_ids, position_ids, attention_mask]
+
+
+class TestConfigurableMP(ConfigurableMP):
+
+    @pytest.mark.world_size(1)
+    @pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
+    def test_gpt2_basic(self, tmpdir, inputs):
+        args_defaults = {
+            'num_layers': 2,
+            'hidden_size': 128,
+            'num_attention_heads': 8,
+            'max_position_embeddings': 128,
+        }
+
+        model = get_gpt2_model(args_defaults)
+        model = get_deepspeed_model(model)
+
+        model.eval()
+        device_name = get_accelerator().device_name()
+        baseline = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
+
+        tag = 'mp_1'
+        state_dict = {}
+        state_dict['checkpoint_version'] = get_megatron_version()
+        model.save_checkpoint(tmpdir, tag=tag, client_state=state_dict)
+        dist.barrier()
+        model.load_checkpoint(tmpdir, tag=tag, load_optimizer_states=False, load_lr_scheduler_states=False)
+
+        test = model(inputs[0], inputs[1], inputs[2])
+        assert torch.allclose(baseline, test,
+                              atol=1e-07), f"Baseline output {baseline} is not equal to save-then-load output {test}"
+
+    @pytest.mark.world_size(2)
+    @pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
+    def test_gpt2_mp2_no_resize(self, tmpdir, inputs):
+        args_defaults = {
+            'num_layers': 2,
+            'hidden_size': 128,
+            'num_attention_heads': 8,
+            'max_position_embeddings': 128,
+        }
+
+        model = get_gpt2_model(args_defaults, mp_size=2)
+        model = get_deepspeed_model(model)
+
+        model.eval()
+
+        device_name = get_accelerator().device_name()
+        baseline = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
+
+        tag = 'mp_2'
+        state_dict = {}
+        state_dict['checkpoint_version'] = get_megatron_version()
+        model.save_checkpoint(tmpdir, tag=tag, client_state=state_dict)
+        dist.barrier()
+        model.load_checkpoint(tmpdir, tag=tag, load_optimizer_states=False, load_lr_scheduler_states=False)
+
+        device_name = get_accelerator().device_name()
+        test = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
+        assert torch.allclose(baseline, test, rtol=1.0,
+                              atol=1e-07), f"Baseline output {baseline} is not equal to save-then-load output {test}"
+
+
+# This fixture provides the baseline model with mp=2 to TestConfigurableMPResize
+class baseline_mp2(DistributedFixture):
+    world_size = 2
+
+    def run(self, inputs, class_tmpdir):
+        args_defaults = {
+            'num_layers': 2,
+            'hidden_size': 128,
+            'num_attention_heads': 8,
+            'max_position_embeddings': 128,
+        }
+
+        model = get_gpt2_model(args_defaults, mp_size=self.world_size)
+        model = get_deepspeed_model(model)
+
+        model.eval()
+
+        with torch.no_grad():
+            device_name = get_accelerator().device_name()
+            baseline = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
+            if dist.get_rank() == 0:
+                save_path = os.path.join(class_tmpdir, "output.pt")
+                torch.save(baseline.cpu(), save_path)
+
+            state_dict = {}
+            state_dict['checkpoint_version'] = get_megatron_version()
+            model.save_checkpoint(class_tmpdir, client_state=state_dict)
+
+
+class TestConfigurableResizeMP(ConfigurableMP):
+    world_size = [1, 4]
+
+    @pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
+    def test(self, baseline_mp2, inputs, class_tmpdir):
+        args_defaults = {
+            'num_layers': 2,
+            'hidden_size': 128,
+            'num_attention_heads': 8,
+            'max_position_embeddings': 128,
+        }
+
+        world_size = os.environ["WORLD_SIZE"]
+        model = get_gpt2_model(args_defaults, mp_size=world_size)
+        model = get_deepspeed_model(model)
+
+        model.eval()
+
+        with torch.no_grad():
+            model.load_checkpoint(class_tmpdir, load_optimizer_states=False, load_lr_scheduler_states=False)
+            device_name = get_accelerator().device_name()
+            test = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
+            if dist.get_rank() == 0:
+                load_path = os.path.join(class_tmpdir, "output.pt")
+                baseline = torch.load(load_path, weights_only=False)
+                test = test.cpu()
+                assert torch.allclose(
+                    baseline, test,
+                    atol=1e-03), f"Baseline output {baseline} is not equal to save-then-load output {test}"