Add files using upload-large-folder tool

lib/python3.12/site-packages/deepspeed/compression/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .compress import init_compression, redundancy_clean
+from .scheduler import compression_scheduler
+from .helper import convert_conv1d_to_linear

lib/python3.12/site-packages/deepspeed/compression/basic_layer.py

@@ -0,0 +1,840 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import math
+from torch import nn
+from torch.nn import init
+import deepspeed.comm as dist
+from .utils import TopKBinarizer, SymQuantizer, AsymQuantizer, TernaryQuantizer, BinaryQuantizer
+from deepspeed.utils import logger
+
+g_mpu = None
+
+
+class QuantAct(nn.Module):
+    """
+    Class to quantize given activations. Note that when using this function, the input activation quantization range will be fixed for all
+    tokens/images for inference. This generally will affect some accuracy but achieve better latency performance.
+    Parameters:
+    ----------
+    act_range_momentum : float, default 0.95
+        Momentum for updating the activation quantization range.
+    quant_mode : str, default 'symmetric'
+    """
+
+    def __init__(self, act_range_momentum=0.95, quant_mode='symmetric'):
+        super(QuantAct, self).__init__()
+
+        self.act_range_momentum = act_range_momentum
+        self.quant_mode = quant_mode
+        if quant_mode == 'symmetric':
+            self.act_function = SymQuantizer.apply
+        else:
+            self.act_function = AsymQuantizer.apply
+
+        self.register_buffer('x_min_max', torch.zeros(2))
+
+    def forward(self, x, num_bits, *args):
+        """
+        x: the activation that we need to quantize
+        num_bits: the number of bits we need to quantize the activation to
+        *args: some extra arguments that are useless but needed for align with the interface of other quantization functions
+        """
+
+        if self.training:
+            x_min = x.data.min()
+            x_max = x.data.max()
+
+            # Initialization
+            if self.x_min_max[0] == self.x_min_max[1]:
+                self.x_min_max[0] = x_min
+                self.x_min_max[1] = x_max
+
+            # if do not need momentum, please set self.act_range_momentum = 0
+            self.x_min_max[0] = self.x_min_max[0] * self.act_range_momentum + x_min * (1 - self.act_range_momentum)
+            self.x_min_max[1] = self.x_min_max[1] * self.act_range_momentum + x_max * (1 - self.act_range_momentum)
+
+        x_q = self.act_function(x, num_bits, self.x_min_max[0], self.x_min_max[1])
+
+        return x_q
+
+
+class Embedding_Compress(nn.Embedding):
+
+    def __init__(self, *kargs):
+        super(Embedding_Compress, self).__init__(*kargs)
+        self.weight.start_bits = None
+        self.weight.target_bits = None
+        self.weight.q_period = None
+        self.weight_quantization_enabled_in_forward = False
+        self.weight_quantization_enabled = False
+
+    def extra_repr(self):
+        return 'num_embeddings={}, embedding_dim={}, weight_quantization={}'.format(
+            self.num_embeddings, self.embedding_dim, self.weight.target_bits)
+
+    def enable_weight_quantization(self, start_bits, target_bits, quantization_period,
+                                   weight_quantization_enabled_in_forward, quantization_type, num_groups):
+        self.weight.start_bits = start_bits
+        self.weight.target_bits = target_bits
+        self.weight.q_period = quantization_period
+        self.weight_quantization_enabled_in_forward = weight_quantization_enabled_in_forward
+        if self.weight_quantization_enabled_in_forward:
+            logger.warning(
+                "************ A lot of MoQ features are not supported in quantize_weight_in_forward mode, please consider to use DS-FP16 optimizer************"
+            )
+            if self.weight.target_bits >= 3:
+                if quantization_type == 'symmetric':
+                    self.weight_quantizer = SymQuantizer.apply
+                else:
+                    self.weight_quantizer = AsymQuantizer.apply
+            elif self.weight.target_bits == 2:
+                assert quantization_type == 'symmetric', 'Only symmetric quantization is supported for ternary weight quantization'
+                self.weight_quantizer = TernaryQuantizer.apply
+            elif self.weight.target_bits == 1:
+                assert quantization_type == 'symmetric', 'Only symmetric quantization is supported for binary weight quantization'
+                self.weight_quantizer = BinaryQuantizer.apply
+            # for embedding, we always use token-wise quantization
+            self.weight_quantize_num_groups = self.weight.size(0)
+
+    def fix_weight_quantization(self):
+        self.weight.data = self.weight_quantizer(self.weight, self.weight.target_bits, None, None,
+                                                 self.weight_quantize_num_groups).data
+        self.weight_quantization_enabled_in_forward = False
+        return None
+
+    def forward(self, input):
+        if self.weight_quantization_enabled_in_forward and self.weight_quantization_enabled:
+            weight = self.weight_quantizer(self.weight, self.weight.target_bits, None, None,
+                                           self.weight_quantize_num_groups)
+        else:
+            weight = self.weight
+
+        out = nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type,
+                                      self.scale_grad_by_freq, self.sparse)
+        return out
+
+
+class LinearLayer_Compress(nn.Linear):
+    """
+    Linear layer with compression.
+    """
+
+    def __init__(self, *kargs, bias=True):
+        super(LinearLayer_Compress, self).__init__(*kargs, bias=bias)
+        self.sparse_pruning_method = None
+        self.row_pruning_method = None
+        self.head_pruning_method = None
+        self.activation_quantization_method = None
+        self.weight.start_bits = None
+        self.weight.target_bits = None
+        self.weight.q_period = None
+        self.weight_quantization_enabled_in_forward = False
+        self.weight_quantization_enabled = False
+        self.sparse_pruning_enabled = False
+        self.row_pruning_enabled = False
+        self.head_pruning_enabled = False
+        self.activation_quantization_enabled = False
+
+    def extra_repr(self):
+        return 'in_features={}, out_features={}, bias={}, sparse pruning={}, row pruning={}, head pruning={}, activation quantization={}, weight_quantization={}'.format(
+            self.in_features, self.out_features, self.bias is not None, self.sparse_pruning_method is not None, \
+            self.row_pruning_method is not None, self.head_pruning_method is not None, self.activation_quantization_method is not None, self.weight.target_bits)
+
+    def enable_sparse_pruning(self, ratio, method):
+        # Here, we support two cases: L1 norm based pruning and topk based pruning
+        self.sparse_pruning_ratio = ratio
+        self.sparse_pruning_method = method
+        if method == 'l1':
+            weight_norm = torch.abs(self.weight.data)
+            mask = TopKBinarizer.apply(weight_norm, self.sparse_pruning_ratio, False)
+            mask = mask.view(self.weight.size())
+            mask = mask.to(self.weight.device)
+        elif method == 'topk':
+            self.sparse_mask_scores = nn.Parameter(torch.Tensor(self.weight.size()))
+            self.sparse_mask_scores.data = self.sparse_mask_scores.data.to(self.weight.device)
+            init.kaiming_uniform_(self.sparse_mask_scores, a=math.sqrt(5))
+            mask = None
+        else:
+            raise NotImplementedError
+
+        self.register_buffer('sparse_pruning_mask', mask)
+
+    def enable_row_pruning(self, ratio, method):
+        # Here, we support two cases: L1 norm based pruning and topk based pruning
+        self.row_pruning_ratio = ratio
+        self.row_pruning_method = method
+
+        if method == 'l1':
+            # compute the l1 norm of each column
+            weight_norm = torch.linalg.norm(self.weight.data, ord=1, dim=1)
+            mask = TopKBinarizer.apply(weight_norm, self.row_pruning_ratio, False)
+            mask = mask.view(-1, 1)
+            mask = mask.to(self.weight.device)
+        elif method == 'topk':
+            self.row_mask_scores = nn.Parameter(torch.Tensor(self.weight.size(0), 1))
+            self.row_mask_scores.data = self.row_mask_scores.data.to(self.weight.device)
+            init.kaiming_uniform_(self.row_mask_scores, a=math.sqrt(5))
+            mask = None
+        else:
+            raise NotImplementedError
+
+        self.register_buffer('row_pruning_mask', mask)
+
+    def enable_head_pruning(self, ratio, method, num_heads):
+        # Here, we support only topk based pruning
+        self.num_heads = num_heads
+        self.head_pruning_ratio = ratio
+        self.head_pruning_method = method
+
+        if method not in ['topk']:
+            raise NotImplementedError
+        else:
+            self.head_pruning_ratio = ratio
+            self.head_pruning_scores = nn.Parameter(torch.Tensor(1,
+                                                                 self.num_heads))  # we apply the pruning to O matrix
+            self.head_pruning_scores.data = self.head_pruning_scores.data.to(self.weight.device)
+            init.kaiming_uniform_(self.head_pruning_scores, a=math.sqrt(5))
+
+    def fix_sparse_pruning_helper(self):
+        mask = self.get_mask(pruning_type='sparse')
+        self.weight.data = self.weight.data * mask
+        del self.sparse_pruning_mask
+        if self.sparse_pruning_method == 'topk':
+            del self.sparse_mask_scores
+        self.sparse_pruning_method = None
+        self.sparse_pruning_enabled = False
+        return None
+
+    def fix_row_col_pruning_helper(self, mask=None, dim_reduction=False):
+        # This function is used for row/col pruning
+        # particularly, if we have two back-to-back layers, F1 and F2; when
+        # we remove rows from F1, we also need to remove columns from F2
+        # However, if we only have one layer, F1, then we only need to mask pruned
+        # rows as 0 in F1
+        if mask is None:
+            mask = self.get_mask(pruning_type='row').bool()
+            if dim_reduction:
+                start_bits = self.weight.start_bits
+                target_bits = self.weight.target_bits
+                q_period = self.weight.q_period
+                self.weight = nn.Parameter(self.weight.data[mask.view(-1), :])
+                self.weight.start_bits = start_bits
+                self.weight.target_bits = target_bits
+                self.weight.q_period = q_period
+                if self.bias is not None:
+                    self.bias = nn.Parameter(self.bias.data[mask.view(-1)])
+                self.out_features = self.weight.size(0)
+            else:
+                self.weight.data = self.weight.data * mask.view(-1, 1)
+                if self.bias is not None:
+                    self.bias.data = self.bias.data * mask.view(-1)
+
+            del self.row_pruning_mask
+            if self.row_pruning_method == 'topk':
+                del self.row_mask_scores
+            self.row_pruning_method = None
+        else:
+            # this is generally for column pruning
+            start_bits = self.weight.start_bits
+            target_bits = self.weight.target_bits
+            q_period = self.weight.q_period
+            self.weight = nn.Parameter(self.weight.data[:, mask.view(-1)])
+            self.weight.start_bits = start_bits
+            self.weight.target_bits = target_bits
+            self.weight.q_period = q_period
+            self.in_features = self.weight.size(1)
+            mask = None
+        self.row_pruning_enabled = False
+        return mask
+
+    def fix_head_pruning_helper(self, mask=None, num_heads=None, dim_reduction=False):
+        # similar as row/col pruning, head pruning also needs to prune QKV which is associated with O matrix
+        num_heads = num_heads if num_heads else self.num_heads
+        if mask is None:
+            if self.head_pruning_method == 'topk':
+                mask = self.get_mask(pruning_type='head').bool()
+                if dim_reduction:
+                    shape = self.weight.size(0)
+                    start_bits = self.weight.start_bits
+                    target_bits = self.weight.target_bits
+                    q_period = self.weight.q_period
+                    self.weight = nn.Parameter(self.weight.data.t().reshape(num_heads,
+                                                                            -1)[mask.view(-1), :].reshape(-1,
+                                                                                                          shape).t())
+                    self.weight.start_bits = start_bits
+                    self.weight.target_bits = target_bits
+                    self.weight.q_period = q_period
+                else:
+
+                    shape = self.weight.size()
+                    self.weight.data = (self.weight.data.t().reshape(self.num_heads, -1) * mask.view(-1, 1)).reshape(
+                        shape[1], shape[0]).t()
+
+                if self.head_pruning_method == 'topk':
+                    del self.head_pruning_scores
+                self.head_pruning_method = None
+            else:
+                raise NotImplementedError
+        else:
+            start_bits = self.weight.start_bits
+            target_bits = self.weight.target_bits
+            q_period = self.weight.q_period
+            shape = self.weight.size(1)
+            self.weight = nn.Parameter(self.weight.data.reshape(num_heads, -1)[mask.view(-1), :].reshape(-1, shape))
+            self.weight.start_bits = start_bits
+            self.weight.target_bits = target_bits
+            self.weight.q_period = q_period
+            if self.bias is not None:
+                self.bias = nn.Parameter(self.bias.data.reshape(num_heads, -1)[mask.view(-1), :].reshape(-1))
+        self.head_pruning_enabled = False
+        return mask
+
+    def get_mask(self, pruning_type='row'):
+        if pruning_type == 'sparse':
+            if self.sparse_pruning_method == 'l1':
+                return self.sparse_pruning_mask.to(self.weight.device)
+            elif self.sparse_pruning_method == 'topk':
+                return TopKBinarizer.apply(self.sparse_mask_scores, self.sparse_pruning_ratio, False)
+            else:
+                raise NotImplementedError
+        if pruning_type == 'row':
+            if self.row_pruning_method == 'l1':
+                return self.row_pruning_mask.to(self.weight.device)
+            elif self.row_pruning_method == 'topk':
+                return TopKBinarizer.apply(self.row_mask_scores, self.row_pruning_ratio, False)
+            else:
+                raise NotImplementedError
+        elif pruning_type == 'head':
+            if self.head_pruning_method == 'topk':
+                return TopKBinarizer.apply(self.head_pruning_scores, self.head_pruning_ratio, False)
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError
+
+    def enable_weight_quantization(self, start_bits, target_bits, quantization_period,
+                                   weight_quantization_enabled_in_forward, quantization_type, num_groups):
+        self.weight.start_bits = start_bits
+        self.weight.target_bits = target_bits
+        self.weight.q_period = quantization_period
+        self.weight_quantization_enabled_in_forward = weight_quantization_enabled_in_forward
+        if self.weight_quantization_enabled_in_forward:
+            logger.warning(
+                "************ A lot of MoQ features are not supported in quantize_weight_in_forward mode, please consider to use DS-FP16 optimizer************"
+            )
+            if self.weight.target_bits >= 3:
+                if quantization_type == 'symmetric':
+                    self.weight_quantizer = SymQuantizer.apply
+                else:
+                    self.weight_quantizer = AsymQuantizer.apply
+            elif self.weight.target_bits == 2:
+                assert quantization_type == 'symmetric', 'Only symmetric quantization is supported for ternary weight quantization'
+                self.weight_quantizer = TernaryQuantizer.apply
+            elif self.weight.target_bits == 1:
+                assert quantization_type == 'symmetric', 'Only symmetric quantization is supported for binary weight quantization'
+                self.weight_quantizer = BinaryQuantizer.apply
+            self.weight_quantize_num_groups = num_groups
+
+    def fix_weight_quantization(self):
+        self.weight.data = self.weight_quantizer(self.weight, self.weight.target_bits, None, None,
+                                                 self.weight_quantize_num_groups).data
+        self.weight_quantization_enabled_in_forward = False
+        return None
+
+    def enable_activation_quantization(self, bits, quantization_type, range_calibration):
+        assert bits in [4, 8], 'Only 4/8 bits activation quantization are supported for now'
+        self.activation_quantization_bits = bits
+        self.activation_quantization_method = f"{quantization_type}_{range_calibration}"
+        if range_calibration == 'static':
+            self.activation_quantizer = QuantAct(quant_mode=quantization_type)
+        else:
+            if quantization_type == 'symmetric':
+                self.activation_quantizer = SymQuantizer.apply
+            else:
+                self.activation_quantizer = AsymQuantizer.apply
+
+    def head_pruning_reshape(self, w, mask):
+        shape = w.shape
+        return (w.t().reshape(self.num_heads, -1) * mask.view(-1, 1)).reshape(shape[1], shape[0]).t()
+
+    def forward(self, input, skip_bias_add=False):
+
+        if self.weight_quantization_enabled_in_forward and self.weight_quantization_enabled:
+            weight = self.weight_quantizer(self.weight, self.weight.target_bits, None, None,
+                                           self.weight_quantize_num_groups)
+            bias = self.bias
+        else:
+            weight = self.weight
+            bias = self.bias
+
+        if self.sparse_pruning_enabled and self.sparse_pruning_method:
+            mask = self.get_mask(pruning_type='sparse')
+            weight = weight * mask.view(self.weight.size())
+
+        if self.row_pruning_enabled and self.row_pruning_method:
+            mask = self.get_mask(pruning_type='row')
+            weight = weight * mask.view(-1, 1)
+            if bias is not None:
+                bias = bias * mask.view(-1)
+
+        if self.head_pruning_enabled and self.head_pruning_method:
+            mask = self.get_mask(pruning_type='head')
+            weight = self.head_pruning_reshape(weight, mask)
+
+        if self.activation_quantization_enabled:
+            if 'dynamic' in self.activation_quantization_method:
+                num_groups = input.numel() // input.size(-1)
+            else:
+                num_groups = 1
+            input = self.activation_quantizer(input, self.activation_quantization_bits, None, None, num_groups)
+
+        if skip_bias_add:
+            # used for mpu linear layers
+            output = nn.functional.linear(input, weight, None)
+            return output, bias
+        else:
+            output = nn.functional.linear(input, weight, bias)
+            return output
+
+
+class Conv2dLayer_Compress(nn.Conv2d):
+    """
+    Conv2D layer with compression.
+    """
+
+    def __init__(self, *kargs):
+        super(Conv2dLayer_Compress, self).__init__(*kargs)
+        self.sparse_pruning_method = None
+        self.channel_pruning_method = None
+        self.activation_quantization_method = None
+        self.weight.start_bits = None
+        self.weight.target_bits = None
+        self.weight.q_period = None
+        self.weight_quantization_enabled_in_forward = False
+        self.sparse_pruning_enabled = False
+        self.channel_pruning_enabled = False
+        self.activation_quantization_enabled = False
+
+    def __repr__(self):
+        s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}'
+             ', stride={stride}')
+        if self.padding != (0, ) * len(self.padding):
+            s += ', padding={padding}'
+        if self.dilation != (1, ) * len(self.dilation):
+            s += ', dilation={dilation}'
+        if self.output_padding != (0, ) * len(self.output_padding):
+            s += ', output_padding={output_padding}'
+        if self.groups != 1:
+            s += ', groups={groups}'
+        if self.bias is None:
+            s += ', bias=False'
+        if self.padding_mode != 'zeros':
+            s += ', padding_mode={padding_mode}'
+        output = s.format(**self.__dict__)
+
+        return output + ' sparse pruning={}, channel pruning={}, activation quantization={}, weight_quantization={}'.format(
+            self.sparse_pruning_method is not None, self.channel_pruning_method is not None,
+            self.activation_quantization_method is not None, self.weight.target_bits)
+
+    def enable_sparse_pruning(self, ratio, method):
+        self.sparse_pruning_ratio = ratio
+        self.sparse_pruning_method = method
+        if method == 'l1':
+            weight_norm = torch.abs(self.weight.data)
+            mask = TopKBinarizer.apply(weight_norm, self.sparse_pruning_ratio, False)
+            mask = mask.view(self.weight.size())
+            mask = mask.to(self.weight.device)
+        elif method == 'topk':
+            self.sparse_mask_scores = nn.Parameter(torch.Tensor(self.weight.size()))
+            self.sparse_mask_scores.data = self.sparse_mask_scores.data.to(self.weight.device)
+            init.kaiming_uniform_(self.sparse_mask_scores, a=math.sqrt(5))
+            mask = None
+        else:
+            raise NotImplementedError
+
+        self.register_buffer('sparse_pruning_mask', mask)
+
+    def enable_channel_pruning(self, ratio, method):
+        # Here, we support two cases: L1 norm based pruning and topk based pruning
+        self.channel_pruning_ratio = ratio
+        self.channel_pruning_method = method
+
+        if method == 'l1':
+            # compute the l1 norm of each conv2d kernel (the last three dimension)
+            weight_norm = torch.linalg.norm(self.weight.data, ord=1, dim=[1, 2, 3])
+            mask = TopKBinarizer.apply(weight_norm, self.channel_pruning_ratio, False)
+            mask = mask.view(-1, 1, 1, 1)
+            mask = mask.to(self.weight.device)
+        elif method == 'topk':
+            self.channel_mask_scores = nn.Parameter(torch.Tensor(self.weight.size(0), 1, 1, 1))
+            self.channel_mask_scores.data = self.channel_mask_scores.data.to(self.weight.device)
+            init.kaiming_uniform_(self.channel_mask_scores, a=math.sqrt(5))
+            mask = None
+        else:
+            raise NotImplementedError
+
+        self.register_buffer('channel_pruning_mask', mask)
+
+    def fix_sparse_pruning_helper(self):
+        mask = self.get_mask(pruning_type='sparse')
+        self.weight.data = self.weight.data * mask
+        del self.sparse_pruning_mask
+        if self.sparse_pruning_method == 'topk':
+            del self.sparse_mask_scores
+        self.sparse_pruning_method = None
+        self.sparse_pruning_enabled = False
+        return None
+
+    def fix_channel_pruning_helper(self, mask=None, dim_reduction=False):
+        if mask is None:
+            if self.channel_pruning_method in ['l1', 'topk']:
+                mask = self.get_mask(pruning_type='channel').bool()
+                if dim_reduction:
+                    start_bits = self.weight.start_bits
+                    target_bits = self.weight.target_bits
+                    q_period = self.weight.q_period
+                    self.weight = nn.Parameter(self.weight.data[mask.view(-1), ...])
+                    self.weight.start_bits = start_bits
+                    self.weight.target_bits = target_bits
+                    self.weight.q_period = q_period
+                    if self.bias is not None:
+                        self.bias = nn.Parameter(self.bias.data[mask.view(-1)])
+                else:
+                    self.weight.data = self.weight.data * mask.view(-1, 1, 1, 1)
+                    if self.bias is not None:
+                        self.bias.data = self.bias.data * mask.view(-1)
+                del self.channel_pruning_mask
+                if self.channel_pruning_method == 'topk':
+                    del self.channel_mask_scores
+                self.channel_pruning_method = None
+            else:
+                raise NotImplementedError
+        else:
+            start_bits = self.weight.start_bits
+            target_bits = self.weight.target_bits
+            q_period = self.weight.q_period
+            self.weight = nn.Parameter(self.weight.data[:, mask.view(-1), ...])
+            self.weight.start_bits = start_bits
+            self.weight.target_bits = target_bits
+            self.weight.q_period = q_period
+            mask = None
+        self.channel_pruning_enabled = False
+        return mask
+
+    def get_mask(self, pruning_type='sparse'):
+        if pruning_type == 'sparse':
+            if self.sparse_pruning_method == 'l1':
+                return self.sparse_pruning_mask.to(self.weight.device)
+            elif self.sparse_pruning_method == 'topk':
+                return TopKBinarizer.apply(self.sparse_mask_scores, self.sparse_pruning_ratio, False)
+            else:
+                raise NotImplementedError
+        elif pruning_type == 'channel':
+            if self.channel_pruning_method == 'l1':
+                return self.channel_pruning_mask.to(self.weight.device)
+            elif self.channel_pruning_method == 'topk':
+                return TopKBinarizer.apply(self.channel_mask_scores, self.channel_pruning_ratio, False)
+            else:
+                raise NotImplementedError
+        else:
+            raise NotImplementedError
+
+    def fix_weight_quantization(self):
+        self.weight.data = self.weight_quantizer(self.weight, self.weight.target_bits, None, None,
+                                                 self.weight_quantize_num_groups).data
+        self.weight_quantization_enabled_in_forward = False
+        return None
+
+    def enable_weight_quantization(self, start_bits, target_bits, quantization_period,
+                                   weight_quantization_enabled_in_forward, quantization_type, num_groups):
+        self.weight.start_bits = start_bits
+        self.weight.target_bits = target_bits
+        self.weight.q_period = quantization_period
+        self.weight_quantization_enabled_in_forward = weight_quantization_enabled_in_forward
+        if self.weight_quantization_enabled_in_forward:
+            assert self.weight.target_bits >= 4, 'Only >=4 bits weight quantization are supported during forward pass for now'
+            logger.warning(
+                "************ A lot of MoQ features are not supported in quantize_weight_in_forward mode, please consider to use DS-FP16 optimizer************"
+            )
+            if quantization_type == 'symmetric':
+                self.weight_quantizer = SymQuantizer.apply
+            else:
+                self.weight_quantizer = AsymQuantizer.apply
+            self.weight_quantize_num_groups = num_groups
+
+    def enable_activation_quantization(self, bits, quantization_type, range_calibration):
+        assert bits in [4, 8], 'Only 4/8 bits activation quantization are supported for now'
+        self.activation_quantization_bits = bits
+        self.activation_quantization_method = f"{quantization_type}_{range_calibration}"
+        if range_calibration == 'static':
+            self.activation_quantizer = QuantAct(quant_mode=quantization_type)
+        else:
+            if quantization_type == 'symmetric':
+                self.activation_quantizer = SymQuantizer.apply
+            else:
+                self.activation_quantizer = AsymQuantizer.apply
+
+    def forward(self, input):
+
+        if self.weight_quantization_enabled_in_forward and self.weight_quantization_enabled:
+            weight = self.weight_quantizer(self.weight, self.weight.target_bits, None, None,
+                                           self.weight_quantize_num_groups)
+            bias = self.bias
+        else:
+            weight = self.weight
+            bias = self.bias
+
+        if self.sparse_pruning_enabled and self.sparse_pruning_method:
+            mask = self.get_mask(pruning_type='sparse')
+            weight = weight * mask.view(self.weight.size())
+
+        if self.channel_pruning_enabled:
+            mask = self.get_mask(pruning_type='channel')
+            weight = weight * mask.view(-1, 1, 1, 1)
+            if bias is not None:
+                bias = bias * mask.view(-1)
+
+        if self.activation_quantization_enabled:
+            if 'dynamic' in self.activation_quantization_method:
+                num_groups = input.numel() // input[0].numel()
+            else:
+                num_groups = 1
+            input = self.activation_quantizer(input, self.activation_quantization_bits, None, None, num_groups)
+
+        return nn.functional.conv2d(input, weight, bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+class BNLayer_Compress(nn.BatchNorm2d):
+
+    def fix_channel_pruning_helper(self, mask, dim_reduction=True):
+        self.weight = nn.Parameter(self.weight.data[mask.view(-1)])
+        self.bias = nn.Parameter(self.bias.data[mask.view(-1)])
+        self.running_mean = self.running_mean[mask.view(-1)]
+        self.running_var = self.running_var[mask.view(-1)]
+
+
+def _reduce(input_):
+    """All-reduce the input tensor across model parallel group."""
+    group = g_mpu.get_model_parallel_group()
+
+    # Bypass the function if we are using only 1 GPU.
+    if dist.get_world_size(group=group) == 1:
+        return input_
+
+    # All-reduce.
+    dist.all_reduce(input_, group=group)
+
+    return input_
+
+
+def split_tensor_along_last_dim(tensor, num_partitions, contiguous_split_chunks=False):
+    """Split a tensor along its last dimension.
+    Arguments:
+        tensor: input tensor.
+        num_partitions: number of partitions to split the tensor
+        contiguous_split_chunks: If True, make each chunk contiguous
+                                 in memory.
+    """
+    # Get the size and dimension.
+    last_dim = tensor.dim() - 1
+    assert tensor.size()[last_dim] % num_partitions == 0
+    last_dim_size = tensor.size()[last_dim] // num_partitions
+    # Split.
+    tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
+    # Note: torch.split does not create contiguous tensors by default.
+    if contiguous_split_chunks:
+        return tuple(chunk.contiguous() for chunk in tensor_list)
+
+    return tensor_list
+
+
+def _split(input_):
+    """Split the tensor along its last dimension and keep the
+    corresponding slice."""
+    group = g_mpu.get_model_parallel_group()
+
+    # Bypass the function if we are using only 1 GPU.
+    if dist.get_world_size(group=group) == 1:
+        return input_
+
+    # Split along last dimension.
+    world_size = dist.get_world_size(group=group)
+    input_list = split_tensor_along_last_dim(input_, world_size)
+
+    # Note: torch.split does not create contiguous tensors by default.
+    rank = dist.get_rank(group=group)
+    output = input_list[rank].contiguous()
+
+    return output
+
+
+def _gather(input_):
+    """Gather tensors and concatenate along the last dimension."""
+    group = g_mpu.get_model_parallel_group()
+
+    # Bypass the function if we are using only 1 GPU.
+    if dist.get_world_size(group=group) == 1:
+        return input_
+
+    # Size and dimension.
+    last_dim = input_.dim() - 1
+    rank = dist.get_rank(group=group)
+    world_size = dist.get_world_size(group=group)
+
+    tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+    tensor_list[rank] = input_
+    dist.all_gather(tensor_list, input_, group=group)
+
+    # Note: torch.cat already creates a contiguous tensor.
+    output = torch.cat(tensor_list, dim=last_dim).contiguous()
+
+    return output
+
+
+class _CopyToModelParallelRegion(torch.autograd.Function):
+    """Pass the input to the model parallel region."""
+
+    @staticmethod
+    def forward(ctx, input_):
+        return input_
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return _reduce(grad_output)
+
+
+class _ReduceFromModelParallelRegion(torch.autograd.Function):
+    """All-reduce the input from the model parallel region."""
+
+    @staticmethod
+    def forward(ctx, input_):
+        return _reduce(input_)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+class _ScatterToModelParallelRegion(torch.autograd.Function):
+    """Split the input and keep only the corresponding chuck to the rank."""
+
+    @staticmethod
+    def forward(ctx, input_):
+        return _split(input_)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return _gather(grad_output)
+
+
+class _GatherFromModelParallelRegion(torch.autograd.Function):
+    """Gather the input from model parallel region and concatenate."""
+
+    @staticmethod
+    def forward(ctx, input_):
+        return _gather(input_)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return _split(grad_output)
+
+
+# -----------------
+# Helper functions.
+# -----------------
+
+
+def copy_to_model_parallel_region(input_):
+    return _CopyToModelParallelRegion.apply(input_)
+
+
+def reduce_from_model_parallel_region(input_):
+    return _ReduceFromModelParallelRegion.apply(input_)
+
+
+def scatter_to_model_parallel_region(input_):
+    return _ScatterToModelParallelRegion.apply(input_)
+
+
+def gather_from_model_parallel_region(input_):
+    return _GatherFromModelParallelRegion.apply(input_)
+
+
+class ColumnParallelLinear_Compress(LinearLayer_Compress):
+
+    def __init__(self, mpu, input_size, output_size, bias=True, gather_output=True, skip_bias_add=False):
+        # Keep input parameters
+        global g_mpu
+        g_mpu = mpu
+        self.input_size = input_size
+        self.output_size = output_size
+        self.gather_output = gather_output
+        self.skip_bias_add = skip_bias_add
+
+        # Divide the weight matrix along the last dimension.
+        world_size = mpu.get_model_parallel_world_size()
+        assert output_size % world_size == 0
+        self.output_size_per_partition = output_size // world_size
+
+        super(ColumnParallelLinear_Compress, self).__init__(self.input_size, self.output_size_per_partition, bias=bias)
+
+    def forward(self, input_):
+        # Set up backprop all-reduce.
+        input_parallel = copy_to_model_parallel_region(input_)
+        # Matrix multiply.
+        if self.skip_bias_add:
+            output_parallel, bias = super().forward(input_parallel, True)
+        else:
+            output_parallel = super().forward(input_parallel)
+            bias = None
+        if self.gather_output:
+            # All-gather across the partitions.
+            output = gather_from_model_parallel_region(output_parallel)
+        else:
+            output = output_parallel
+        return output, bias
+
+
+class RowParallelLinear_Compress(LinearLayer_Compress):
+
+    def __init__(self, mpu, input_size, output_size, bias=True, input_is_parallel=False, skip_bias_add=False):
+        # Keep input parameters
+        global g_mpu
+        g_mpu = mpu
+        self.input_size = input_size
+        self.output_size = output_size
+        self.input_is_parallel = input_is_parallel
+        self.skip_bias_add = skip_bias_add
+
+        # Divide the weight matrix along the last dimension.
+        world_size = mpu.get_model_parallel_world_size()
+        assert input_size % world_size == 0
+        self.input_size_per_partition = input_size // world_size
+
+        super(RowParallelLinear_Compress, self).__init__(self.input_size_per_partition, self.output_size, bias=bias)
+
+    def forward(self, input_):
+        # Set up backprop all-reduce.
+        if self.input_is_parallel:
+            input_parallel = input_
+        else:
+            input_parallel = scatter_to_model_parallel_region(input_)
+        # Matrix multiply.
+        output_parallel, bias = super().forward(input_parallel, True)
+
+        # All-reduce across all the partitions.
+        output_ = reduce_from_model_parallel_region(output_parallel)
+        if not self.skip_bias_add:
+            if bias is not None:
+                output = output_ + bias
+            else:
+                output = output_
+            output_bias = None
+        else:
+            output = output_
+            output_bias = bias
+        return output, output_bias

lib/python3.12/site-packages/deepspeed/compression/compress.py

@@ -0,0 +1,239 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import re
+from .helper import compression_preparation, fix_compression, recursive_getattr, is_module_compressible
+from .config import get_compression_config
+from ..runtime.config_utils import dict_raise_error_on_duplicate_keys
+from .constants import *
+import os
+import json
+
+try:
+    import neural_compressor as nc
+except ImportError as e:
+    nc = None
+
+
+def check_deepspeed_config(config):
+    if isinstance(config, dict):
+        return config
+    elif os.path.exists(config):
+        return json.load(open(config, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys)
+    else:
+        raise ValueError(
+            f"Expected a string path to an existing deepspeed config, or a dictionary. Received: {config}")
+
+
+def get_module_name(group_name, model, key_word, exist_module_name, mpu=None, verbose=True):
+    '''
+    get the associated module name from the model based on the key_word provided by users
+    '''
+    return_module_name = []
+    for name, module in model.named_modules():
+
+        module_check = is_module_compressible(module, mpu)
+
+        if re.search(key_word, name) is not None and module_check:
+            if name in exist_module_name and verbose:
+                # logger.warning
+                raise ValueError(
+                    f"{name} is already added to compression, please check your config file for {group_name}.")
+            if name not in exist_module_name:
+                exist_module_name.add(name)
+                return_module_name.append(name)
+    return return_module_name, exist_module_name
+
+
+def get_compress_methods(model, compress_methods, mpu=None):
+    # extract the compression module for each method in compress_methods
+    layer_added_compress_methods = []
+    for method, method_content in compress_methods.items():
+        if LAYER_REDUCTION in method:
+            continue
+        # for loop different methods, i.e., weight quantization, activation quantization etc
+        exist_module_name = set()
+        shared_parameters = method_content[SHARED_PARAMETERS]  # get all the shared parameters
+        for group_name, method_parameters in method_content[DIFFERENT_GROUPS].items():
+            # for loop different groups, i.e., weight quantization group 1, weight quantization group 2 etc
+            module_name_list = []
+            related_module_name_list = []
+            if method_parameters[DIFFERENT_GROUPS_RELATED_MODULE_SCOPE]:
+                # this is used for head/row/channel pruning, if users provide the related module scope, we can shrink the layer dim for them
+                # otherwise we just mask those as zeros
+                for key_word, related_key_words in zip(method_parameters[DIFFERENT_GROUPS_MODULE_SCOPE],
+                                                       method_parameters[DIFFERENT_GROUPS_RELATED_MODULE_SCOPE]):
+                    module_name, exist_module_name = get_module_name(group_name,
+                                                                     model,
+                                                                     key_word,
+                                                                     exist_module_name,
+                                                                     mpu=mpu)
+                    module_name_list.append(module_name)
+                    tmp_related_module_name_list = []
+                    for rkw in related_key_words:
+                        # related key word can be a list, for instance the QKV for O matrix in Attention
+                        module_name, _ = get_module_name(group_name, model, rkw, set(), mpu=mpu)
+                        tmp_related_module_name_list.append(module_name)
+                    related_module_name_list.append(tmp_related_module_name_list)
+            else:
+                for key_word in method_parameters[DIFFERENT_GROUPS_MODULE_SCOPE]:
+                    module_name, exist_module_name = get_module_name(group_name,
+                                                                     model,
+                                                                     key_word,
+                                                                     exist_module_name,
+                                                                     mpu=mpu)
+                    module_name_list.append(module_name)
+
+            if module_name_list:
+                # combine shared parameters with each group
+                combined_method_parameters = {
+                    **(method_parameters.copy().pop(DIFFERENT_GROUPS_PARAMETERS)),
+                    **shared_parameters
+                }
+                compression_item = [module_name_list, related_module_name_list, {method: combined_method_parameters}]
+                layer_added_compress_methods.append(compression_item)
+    return layer_added_compress_methods
+
+
+def init_compression(model, deepspeed_config, teacher_model=None, mpu=None):
+    """
+    Compress a model: replace linear/conv2d layer with deepspeed compression-aware modules
+    Args:
+        model (`torch.nn.Module`)
+            The model to compress.
+        deepspeed_config (`DeepSpeedConfig`)
+            The path of ds_config
+        mpu
+            The mpu module for Row/Column parallelism
+    """
+    compress_methods = get_compression_config(check_deepspeed_config(deepspeed_config))
+    if hasattr(model, 'module'):
+        c_model = model.module
+    else:
+        c_model = model
+
+    # For layer reduction
+    if compress_methods[LAYER_REDUCTION][LAYER_REDUCTION_ENABLED]:
+        assert teacher_model is not None, "Teacher model is required for layer reduction"
+        student_initialization(c_model, teacher_model, deepspeed_config)
+
+    layer_added_compress_methods = get_compress_methods(c_model, compress_methods, mpu=mpu)
+    compression_preparation(c_model, layer_added_compress_methods, mpu)
+
+    # For sparse pruning snip_momentum method
+    shared_parameters = compress_methods[SPARSE_PRUNING][SHARED_PARAMETERS]
+    if shared_parameters[SPARSE_PRUNING_ENABLED] and \
+        shared_parameters[SPARSE_PRUNING_METHOD] == SPARSE_PRUNING_METHOD_SNIP_MOMENTUM:
+
+        assert nc is not None, "please ensure the neural_compressor python package is installed by pip or conda if user wants to use snip_momentum sparse pruning"
+
+        from .helper import generate_pruners, register_on_step_begin
+        from nc import WeightPruningConfig
+
+        config = WeightPruningConfig(target_sparsity=1 - shared_parameters[SPARSE_PRUNING_DENSE_RATIO],
+                                     pattern=shared_parameters[SPARSE_PRUNING_BLOCK_PATTERN],
+                                     pruning_frequency=shared_parameters[SPARSE_PRUNING_SCHEDULE_OFFSET_STRIDE],
+                                     start_step=shared_parameters[SPARSE_PRUNING_SCHEDULE_OFFSET],
+                                     end_step=shared_parameters[SPARSE_PRUNING_SCHEDULE_OFFSET_END],
+                                     excluded_op_names=shared_parameters[SPARSE_PRUNING_EXCLUDED_MODULES])
+        pruners = generate_pruners(config, c_model)
+        c_model.pruners = pruners
+        register_on_step_begin(c_model)
+
+    return model
+
+
+def redundancy_clean(model, deepspeed_config, mpu=None):
+    """
+    Remove the redundancy of a model
+    Args:
+        model (`torch.nn.Module`)
+            The model to compress.
+        deepspeed_config (`DeepSpeedConfig`)
+            The path of ds_config
+        mpu
+            The mpu module for Row/Column parallelism
+    """
+    compress_methods = get_compression_config(check_deepspeed_config(deepspeed_config))
+    if hasattr(model, 'module'):
+        c_model = model.module
+    else:
+        c_model = model
+
+    layer_added_compress_methods_tmp = get_compress_methods(c_model, compress_methods, mpu=mpu)
+    # sort methods
+    order_list = [
+        WEIGHT_QUANTIZATION, SPARSE_PRUNING, ROW_PRUNING, HEAD_PRUNING, CHANNEL_PRUNING, ACTIVATION_QUANTIZATION
+    ]
+    layer_added_compress_methods = sorted(layer_added_compress_methods_tmp,
+                                          key=lambda x: order_list.index(list(x[2].keys())[0]))
+
+    for module_name_lists, related_module_name_lists, compression_technique in layer_added_compress_methods:
+        stored_mask = []
+        need_mask = True if related_module_name_lists else False
+        for i, mnl in enumerate(module_name_lists):
+            for module_name in mnl:
+                mask = fix_compression(c_model, module_name, compression_technique, dim_reduction=need_mask)
+                if need_mask:
+                    stored_mask.append(mask)
+            if need_mask:
+                for rmnl in related_module_name_lists[i]:
+                    for j, module_name in enumerate(rmnl):
+                        mask = fix_compression(c_model,
+                                               module_name,
+                                               compression_technique,
+                                               mask=stored_mask[j],
+                                               dim_reduction=True)
+    return model
+
+
+def student_initialization(student_model, teacher_model, deepspeed_config):
+    '''
+    Given a student model and a teacher model, select the
+    Args:
+        student_model (`torch.nn.Module`)
+            The model we will update weight
+        teacher_model (`torch.nn.Module`)
+            The model guide the student to learn
+        deepspeed_config (`DeepSpeedConfig`)
+            The path of ds_config
+    '''
+    config = get_compression_config(check_deepspeed_config(deepspeed_config))
+    compress_methods = config[LAYER_REDUCTION]
+
+    module_name_prefix = compress_methods[MODULE_NAME_PREFIX]
+    teacher_layer = compress_methods[TEACHER_LAYER]
+    student_layer = [i for i in range(len(teacher_layer))]
+    other_module_name = compress_methods[OTHER_MODULE_NAME]
+    '''
+        name_prefix (`str`)
+            The prefix name before the layer #.
+            Example 1: bert.encoder.layer, for BERT_base model's prefix name
+            Example 2: transformer.h, for GPT-2 hugging face prefix name
+        teacher_layer (`list of integers`)
+            The layer of teacher will be used for student's reinitialization
+            Example 1: [1,3,5,7,9], means we want to matches the 2nd/4th/6th/8th/10th layer of teacher to the first 5 layers of student
+        student_layer (`list` or None)
+            The layer of student need to be re-initialized
+            Example 1: None, means we want to reinitialize all the layers
+            Example 1: [0,1,2,3,4], means  we want to reinitialize the first 5 layers
+        other_module_name (`list of string`)
+            The modules will be used for student's reinitialization
+            Example 1: ['bert.pooler', 'bert.embeddings', 'classifier'], means we want to apply the weight in teacher's embedding/pooler/classier module to the student
+            Example 2: ['transformer.w', 'transformer.ln_f', 'lm_head'], means we want to apply the weight in teacher's embedding layers module to the student
+    Note that teacher_layer should matches student layer
+    '''
+    assert len(student_layer) == len(teacher_layer)
+    for s_name, t_name in zip(student_layer, teacher_layer):
+        s_module = recursive_getattr(student_model, module_name_prefix + '.' + str(s_name))
+        t_module = recursive_getattr(teacher_model, module_name_prefix + '.' + str(t_name))
+        for s_param, t_param in zip(s_module.parameters(), t_module.parameters()):
+            s_param.data.copy_(t_param.data)
+    for name in other_module_name:
+        s_module = recursive_getattr(student_model, name)
+        t_module = recursive_getattr(teacher_model, name)
+        print(name)
+        for s_param, t_param in zip(s_module.parameters(), t_module.parameters()):
+            s_param.data.copy_(t_param.data)

lib/python3.12/site-packages/deepspeed/compression/config.py

@@ -0,0 +1,452 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .constants import *
+import copy
+from ..runtime.config_utils import get_scalar_param, get_list_param
+
+
+def get_compression_config(param_dict):
+    #
+    output = {}
+
+    if COMPRESSION_TRAINING not in param_dict.keys():
+        param_dict[COMPRESSION_TRAINING] = {}
+    sub_param_dict = param_dict[COMPRESSION_TRAINING]
+    output[WEIGHT_QUANTIZATION] = get_weight_quantization(sub_param_dict)
+    output[ACTIVATION_QUANTIZATION] = get_activation_quantization(sub_param_dict)
+    output[SPARSE_PRUNING] = get_sparse_pruning(sub_param_dict)
+    output[ROW_PRUNING] = get_row_pruning(sub_param_dict)
+    output[HEAD_PRUNING] = get_head_pruning(sub_param_dict)
+    output[CHANNEL_PRUNING] = get_channel_pruning(sub_param_dict)
+
+    output[LAYER_REDUCTION] = get_layer_reduction(sub_param_dict)
+
+    return output
+
+
+def get_layer_reduction(param_dict):
+    output = {}
+    output[LAYER_REDUCTION_ENABLED] = LAYER_REDUCTION_ENABLED_DEFAULT
+    if get_layer_reduction_enabled(param_dict):
+        output[LAYER_REDUCTION_ENABLED] = get_layer_reduction_enabled(param_dict)
+        for key, val in get_layer_reduction_params(param_dict).items():
+            output[key] = val
+    return output
+
+
+def get_layer_reduction_enabled(param_dict):
+    if LAYER_REDUCTION in param_dict.keys():
+        return get_scalar_param(param_dict[LAYER_REDUCTION], LAYER_REDUCTION_ENABLED, LAYER_REDUCTION_ENABLED_DEFAULT)
+    else:
+        return False
+
+
+def get_layer_reduction_params(param_dict):
+    if LAYER_REDUCTION in param_dict.keys():
+        layer_reduction_params = copy.copy(param_dict[LAYER_REDUCTION])
+        layer_reduction_params.pop(LAYER_REDUCTION_ENABLED)
+        return layer_reduction_params
+    else:
+        return False
+
+
+def get_quantize_enabled(param_dict):
+    if COMPRESSION_TRAINING not in param_dict.keys():
+        return False
+
+    sub_param_dict = param_dict[COMPRESSION_TRAINING]
+    output = get_weight_quantization_shared_parameters(sub_param_dict)
+    return output[WEIGHT_QUANTIZE_ENABLED]
+
+
+def get_weight_quantization(param_dict):
+    output = {}
+    if WEIGHT_QUANTIZATION not in param_dict.keys():
+        param_dict[WEIGHT_QUANTIZATION] = {SHARED_PARAMETERS: {}, DIFFERENT_GROUPS: {}}
+    sub_param_dict = param_dict[WEIGHT_QUANTIZATION]
+    # shared parameters
+    output[SHARED_PARAMETERS] = get_weight_quantization_shared_parameters(sub_param_dict)
+    # each sub-groups
+    if output[SHARED_PARAMETERS][WEIGHT_QUANTIZE_ENABLED]:
+        assert DIFFERENT_GROUPS in sub_param_dict.keys(
+        ), f"Weigh Quantization is enabled, {DIFFERENT_GROUPS} must be specified"
+    output[DIFFERENT_GROUPS] = get_weight_quantization_different_groups(sub_param_dict)
+    return output
+
+
+def get_weight_quantization_shared_parameters(param_dict):
+    output = {}
+    if SHARED_PARAMETERS in param_dict.keys():
+        sub_param_dict = param_dict[SHARED_PARAMETERS]
+        output[WEIGHT_QUANTIZE_ENABLED] = get_scalar_param(sub_param_dict, WEIGHT_QUANTIZE_ENABLED,
+                                                           WEIGHT_QUANTIZE_ENABLED_DEFAULT)
+        output[WEIGHT_QUANTIZE_KERNEL] = get_scalar_param(sub_param_dict, WEIGHT_QUANTIZE_KERNEL,
+                                                          WEIGHT_QUANTIZE_KERNEL_DEFAULT)
+        output[WEIGHT_QUANTIZE_SCHEDULE_OFFSET] = get_scalar_param(sub_param_dict, WEIGHT_QUANTIZE_SCHEDULE_OFFSET,
+                                                                   WEIGHT_QUANTIZE_SCHEDULE_OFFSET_DEFAULT)
+        output[WEIGHT_QUANTIZE_GROUPS] = get_scalar_param(sub_param_dict, WEIGHT_QUANTIZE_GROUPS,
+                                                          WEIGHT_QUANTIZE_GROUPS_DEFAULT)
+        output[WEIGHT_QUANTIZE_VERBOSE] = get_scalar_param(sub_param_dict, WEIGHT_QUANTIZE_VERBOSE,
+                                                           WEIGHT_QUANTIZE_VERBOSE_DEFAULT)
+        output[WEIGHT_QUANTIZE_TYPE] = get_scalar_param(sub_param_dict, WEIGHT_QUANTIZE_TYPE,
+                                                        WEIGHT_QUANTIZE_TYPE_DEFAULT)
+        output[WEIGHT_QUANTIZE_IN_FORWARD_ENABLED] = get_scalar_param(sub_param_dict,
+                                                                      WEIGHT_QUANTIZE_IN_FORWARD_ENABLED,
+                                                                      WEIGHT_QUANTIZE_IN_FORWARD_ENABLED_DEFAULT)
+        assert output[WEIGHT_QUANTIZE_TYPE] in [
+            WEIGHT_QUANTIZE_SYMMETRIC, WEIGHT_QUANTIZE_ASYMMETRIC
+        ], f"Invalid weight quantize type. Supported types: [{WEIGHT_QUANTIZE_SYMMETRIC}, {WEIGHT_QUANTIZE_ASYMMETRIC}]"
+        output[WEIGHT_QUANTIZE_ROUNDING] = get_scalar_param(sub_param_dict, WEIGHT_QUANTIZE_ROUNDING,
+                                                            WEIGHT_QUANTIZE_ROUNDING_DEFAULT)
+        assert output[WEIGHT_QUANTIZE_ROUNDING] in [
+            WEIGHT_QUANTIZE_NEAREST_ROUNDING, WEIGHT_QUANTIZE_STOCHASTIC_ROUNDING
+        ], f"Invalid weight quantize rounding. Supported types: [{WEIGHT_QUANTIZE_NEAREST_ROUNDING}, {WEIGHT_QUANTIZE_STOCHASTIC_ROUNDING}]"
+        if WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE in sub_param_dict.keys():
+            output[WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE] = get_scalar_param(
+                sub_param_dict[WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE], WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE_ENABLED,
+                WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE_ENABLED_DEFAULT)
+            output[WEIGHT_QUANTIZE_CHANGE_RATIO] = get_scalar_param(
+                sub_param_dict[WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE], WEIGHT_QUANTIZE_CHANGE_RATIO,
+                WEIGHT_QUANTIZE_CHANGE_RATIO_DEFAULT)
+        else:
+            output[WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE] = WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE_ENABLED_DEFAULT
+            output[WEIGHT_QUANTIZE_CHANGE_RATIO] = WEIGHT_QUANTIZE_CHANGE_RATIO_DEFAULT
+    else:
+        output[WEIGHT_QUANTIZE_ENABLED] = WEIGHT_QUANTIZE_ENABLED_DEFAULT
+        output[WEIGHT_QUANTIZE_KERNEL] = WEIGHT_QUANTIZE_KERNEL_DEFAULT
+        output[WEIGHT_QUANTIZE_SCHEDULE_OFFSET] = WEIGHT_QUANTIZE_SCHEDULE_OFFSET_DEFAULT
+        output[WEIGHT_QUANTIZE_GROUPS] = WEIGHT_QUANTIZE_GROUPS_DEFAULT
+        output[WEIGHT_QUANTIZE_VERBOSE] = WEIGHT_QUANTIZE_VERBOSE_DEFAULT
+        output[WEIGHT_QUANTIZE_TYPE] = WEIGHT_QUANTIZE_TYPE_DEFAULT
+        output[WEIGHT_QUANTIZE_ROUNDING] = WEIGHT_QUANTIZE_ROUNDING_DEFAULT
+        output[WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE] = WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE_ENABLED_DEFAULT
+        output[WEIGHT_QUANTIZE_CHANGE_RATIO] = WEIGHT_QUANTIZE_CHANGE_RATIO_DEFAULT
+    return output
+
+
+def get_weight_quantization_different_groups(param_dict):
+    output = {}
+    sub_param_dict = param_dict[DIFFERENT_GROUPS]
+
+    def get_params(name, group_dict):
+        assert WEIGHT_QUANTIZE_START_BITS in group_dict.keys(
+        ), f"{WEIGHT_QUANTIZE_START_BITS} must be specified for weight quantization group {name}"
+        assert WEIGHT_QUANTIZE_TARGET_BITS in group_dict.keys(
+        ), f"{WEIGHT_QUANTIZE_TARGET_BITS} must be specified for weight quantization group {name}"
+        group_dict[WEIGHT_QUANTIZATION_PERIOD] = get_scalar_param(group_dict, WEIGHT_QUANTIZATION_PERIOD,
+                                                                  WEIGHT_QUANTIZATION_PERIOD_DEFAULT)
+        return group_dict
+
+    for k, v in sub_param_dict.items():
+        output[k] = {}
+        output[k][DIFFERENT_GROUPS_PARAMETERS] = get_params(k, sub_param_dict[k][DIFFERENT_GROUPS_PARAMETERS])
+        output[k][DIFFERENT_GROUPS_MODULE_SCOPE] = get_scalar_param(sub_param_dict[k], DIFFERENT_GROUPS_MODULE_SCOPE,
+                                                                    DIFFERENT_GROUPS_MODULE_SCOPE_DEFAULT)
+        output[k][DIFFERENT_GROUPS_RELATED_MODULE_SCOPE] = get_scalar_param(
+            sub_param_dict[k], DIFFERENT_GROUPS_RELATED_MODULE_SCOPE, DIFFERENT_GROUPS_RELATED_MODULE_SCOPE_DEFAULT)
+
+    return output
+
+
+def get_activation_quantization(param_dict):
+    output = {}
+    if ACTIVATION_QUANTIZATION not in param_dict.keys():
+        param_dict[ACTIVATION_QUANTIZATION] = {SHARED_PARAMETERS: {}, DIFFERENT_GROUPS: {}}
+    sub_param_dict = param_dict[ACTIVATION_QUANTIZATION]
+    # shared parameters
+    output[SHARED_PARAMETERS] = get_activation_quantization_shared_parameters(sub_param_dict)
+    # each sub-groups
+    if output[SHARED_PARAMETERS][ACTIVATION_QUANTIZATION_ENABLED]:
+        assert DIFFERENT_GROUPS in sub_param_dict.keys(
+        ), f"Activation Quantization is enabled, {DIFFERENT_GROUPS} must be specified"
+    output[DIFFERENT_GROUPS] = get_activation_quantization_different_groups(sub_param_dict)
+    return output
+
+
+def get_activation_quantization_shared_parameters(param_dict):
+    output = {}
+    if SHARED_PARAMETERS in param_dict.keys():
+        sub_param_dict = param_dict[SHARED_PARAMETERS]
+        output[ACTIVATION_QUANTIZATION_ENABLED] = get_scalar_param(sub_param_dict, ACTIVATION_QUANTIZATION_ENABLED,
+                                                                   ACTIVATION_QUANTIZATION_ENABLED_DEFAULT)
+        output[ACTIVATION_QUANTIZE_TYPE] = get_scalar_param(sub_param_dict, ACTIVATION_QUANTIZE_TYPE,
+                                                            ACTIVATION_QUANTIZE_TYPE_DEFAULT)
+        assert output[ACTIVATION_QUANTIZE_TYPE] in [
+            ACTIVATION_QUANTIZE_SYMMETRIC, ACTIVATION_QUANTIZE_ASYMMETRIC
+        ], f"Invalid activation quantize type. Supported types: [{ACTIVATION_QUANTIZE_SYMMETRIC}, {ACTIVATION_QUANTIZE_ASYMMETRIC}]"
+        output[ACTIVATION_QUANTIZE_RANGE] = get_scalar_param(sub_param_dict, ACTIVATION_QUANTIZE_RANGE,
+                                                             ACTIVATION_QUANTIZE_RANGE_DEFAULT)
+        assert output[ACTIVATION_QUANTIZE_RANGE] in [
+            ACTIVATION_QUANTIZE_RANGE_DYNAMIC, ACTIVATION_QUANTIZE_RANGE_STATIC
+        ], f"Invalid activation quantize range calibration. Supported types: [{ACTIVATION_QUANTIZE_RANGE_DYNAMIC}, {ACTIVATION_QUANTIZE_RANGE_STATIC}]"
+        output[ACTIVATION_QUANTIZE_SCHEDULE_OFFSET] = get_scalar_param(sub_param_dict,
+                                                                       ACTIVATION_QUANTIZE_SCHEDULE_OFFSET,
+                                                                       ACTIVATION_QUANTIZE_SCHEDULE_OFFSET_DEFAULT)
+    else:
+        output[ACTIVATION_QUANTIZATION_ENABLED] = ACTIVATION_QUANTIZATION_ENABLED_DEFAULT
+        output[ACTIVATION_QUANTIZE_TYPE] = ACTIVATION_QUANTIZE_TYPE_DEFAULT
+        output[ACTIVATION_QUANTIZE_RANGE] = ACTIVATION_QUANTIZE_RANGE_DEFAULT
+        output[ACTIVATION_QUANTIZE_SCHEDULE_OFFSET] = ACTIVATION_QUANTIZE_SCHEDULE_OFFSET_DEFAULT
+    return output
+
+
+def get_activation_quantization_different_groups(param_dict):
+    output = {}
+    sub_param_dict = param_dict[DIFFERENT_GROUPS]
+
+    def get_params(name, group_dict):
+        assert ACTIVATION_QUANTIZE_BITS in group_dict.keys(
+        ), f"{ACTIVATION_QUANTIZE_BITS} must be specified for activation quantization group {name}"
+        return group_dict
+
+    for k, v in sub_param_dict.items():
+        output[k] = {}
+        output[k][DIFFERENT_GROUPS_PARAMETERS] = get_params(k, sub_param_dict[k][DIFFERENT_GROUPS_PARAMETERS])
+        output[k][DIFFERENT_GROUPS_MODULE_SCOPE] = get_scalar_param(sub_param_dict[k], DIFFERENT_GROUPS_MODULE_SCOPE,
+                                                                    DIFFERENT_GROUPS_MODULE_SCOPE_DEFAULT)
+        output[k][DIFFERENT_GROUPS_RELATED_MODULE_SCOPE] = get_scalar_param(
+            sub_param_dict[k], DIFFERENT_GROUPS_RELATED_MODULE_SCOPE, DIFFERENT_GROUPS_RELATED_MODULE_SCOPE_DEFAULT)
+
+    return output
+
+
+def get_sparse_pruning(param_dict):
+    output = {}
+    if SPARSE_PRUNING not in param_dict.keys():
+        param_dict[SPARSE_PRUNING] = {SHARED_PARAMETERS: {}, DIFFERENT_GROUPS: {}}
+    sub_param_dict = param_dict[SPARSE_PRUNING]
+    # shared parameters
+    output[SHARED_PARAMETERS] = get_sparse_pruning_shared_parameters(sub_param_dict)
+    # each sub-groups
+    if output[SHARED_PARAMETERS][SPARSE_PRUNING_ENABLED] and output[SHARED_PARAMETERS][
+            SPARSE_PRUNING_METHOD] != SPARSE_PRUNING_METHOD_SNIP_MOMENTUM:
+        assert DIFFERENT_GROUPS in sub_param_dict.keys(
+        ), f"Sparse Pruning is enabled and not snip_momentum method, {DIFFERENT_GROUPS} must be specified"
+    output[DIFFERENT_GROUPS] = get_sparse_pruning_different_groups(sub_param_dict)
+    return output
+
+
+def get_sparse_pruning_shared_parameters(param_dict):
+    output = {}
+
+    if SHARED_PARAMETERS in param_dict.keys():
+        sub_param_dict = param_dict[SHARED_PARAMETERS]
+        output[SPARSE_PRUNING_ENABLED] = get_scalar_param(sub_param_dict, SPARSE_PRUNING_ENABLED,
+                                                          SPARSE_PRUNING_ENABLED_DEFAULT)
+        output[SPARSE_PRUNING_METHOD] = get_scalar_param(sub_param_dict, SPARSE_PRUNING_METHOD,
+                                                         SPARSE_PRUNING_METHOD_DEFAULT)
+        assert output[SPARSE_PRUNING_METHOD] in [
+            SPARSE_PRUNING_METHOD_L1, SPARSE_PRUNING_METHOD_TOPK, SPARSE_PRUNING_METHOD_SNIP_MOMENTUM
+        ], f"Invalid sparse pruning method. Supported types: [{SPARSE_PRUNING_METHOD_L1}, {SPARSE_PRUNING_METHOD_TOPK}, {SPARSE_PRUNING_METHOD_SNIP_MOMENTUM}]"
+        output[SPARSE_PRUNING_SCHEDULE_OFFSET] = get_scalar_param(sub_param_dict, SPARSE_PRUNING_SCHEDULE_OFFSET,
+                                                                  SPARSE_PRUNING_SCHEDULE_OFFSET_DEFAULT)
+        if output[SPARSE_PRUNING_METHOD] == SPARSE_PRUNING_METHOD_SNIP_MOMENTUM:
+            output[SPARSE_PRUNING_BLOCK_PATTERN] = get_scalar_param(sub_param_dict, SPARSE_PRUNING_BLOCK_PATTERN,
+                                                                    SPARSE_PRUNING_BLOCK_PATTERN_DEFAULT)
+            output[SPARSE_PRUNING_DENSE_RATIO] = get_scalar_param(sub_param_dict, SPARSE_PRUNING_DENSE_RATIO,
+                                                                  SPARSE_PRUNING_DENSE_RATIO_DEFAULT)
+            assert output[SPARSE_PRUNING_DENSE_RATIO] > 0 and output[
+                SPARSE_PRUNING_DENSE_RATIO] < 1, f"Invalid dense_ratio value. Must be less than 1"
+            output[SPARSE_PRUNING_SCHEDULE_OFFSET_STRIDE] = get_scalar_param(
+                sub_param_dict, SPARSE_PRUNING_SCHEDULE_OFFSET_STRIDE, SPARSE_PRUNING_SCHEDULE_OFFSET_STRIDE_DEFAULT)
+            output[SPARSE_PRUNING_EXCLUDED_MODULES] = get_list_param(sub_param_dict, SPARSE_PRUNING_EXCLUDED_MODULES,
+                                                                     SPARSE_PRUNING_EXCLUDED_MODULES_DEFAULT)
+            output[SPARSE_PRUNING_SCHEDULE_OFFSET_END] = get_scalar_param(sub_param_dict,
+                                                                          SPARSE_PRUNING_SCHEDULE_OFFSET_END,
+                                                                          output[SPARSE_PRUNING_SCHEDULE_OFFSET])
+            assert output[SPARSE_PRUNING_SCHEDULE_OFFSET] <= output[
+                SPARSE_PRUNING_SCHEDULE_OFFSET_END], f"Invalid schedule_offset and schedule_offset_end values"
+    else:
+        output[SPARSE_PRUNING_ENABLED] = SPARSE_PRUNING_ENABLED_DEFAULT
+        output[SPARSE_PRUNING_METHOD] = SPARSE_PRUNING_METHOD_DEFAULT
+        output[SPARSE_PRUNING_SCHEDULE_OFFSET] = SPARSE_PRUNING_SCHEDULE_OFFSET_DEFAULT
+    return output
+
+
+def get_sparse_pruning_different_groups(param_dict):
+    output = {}
+    sub_param_dict = param_dict[DIFFERENT_GROUPS]
+
+    def get_params(name, group_dict):
+        assert SPARSE_PRUNING_DENSE_RATIO in group_dict.keys(
+        ), f"{SPARSE_PRUNING_DENSE_RATIO} must be specified for sparse pruning group {name}"
+        return group_dict
+
+    for k, v in sub_param_dict.items():
+        output[k] = {}
+        output[k][DIFFERENT_GROUPS_PARAMETERS] = get_params(k, sub_param_dict[k][DIFFERENT_GROUPS_PARAMETERS])
+        output[k][DIFFERENT_GROUPS_MODULE_SCOPE] = get_scalar_param(sub_param_dict[k], DIFFERENT_GROUPS_MODULE_SCOPE,
+                                                                    DIFFERENT_GROUPS_MODULE_SCOPE_DEFAULT)
+        output[k][DIFFERENT_GROUPS_RELATED_MODULE_SCOPE] = get_scalar_param(
+            sub_param_dict[k], DIFFERENT_GROUPS_RELATED_MODULE_SCOPE, DIFFERENT_GROUPS_RELATED_MODULE_SCOPE_DEFAULT)
+
+    return output
+
+
+def get_row_pruning(param_dict):
+    output = {}
+    if ROW_PRUNING not in param_dict.keys():
+        param_dict[ROW_PRUNING] = {SHARED_PARAMETERS: {}, DIFFERENT_GROUPS: {}}
+    sub_param_dict = param_dict[ROW_PRUNING]
+    # shared parameters
+    output[SHARED_PARAMETERS] = get_row_pruning_shared_parameters(sub_param_dict)
+    # each sub-groups
+    if output[SHARED_PARAMETERS][ROW_PRUNING_ENABLED]:
+        assert DIFFERENT_GROUPS in sub_param_dict.keys(
+        ), f"Row Pruning is enabled, {DIFFERENT_GROUPS} must be specified"
+    output[DIFFERENT_GROUPS] = get_row_pruning_different_groups(sub_param_dict)
+    return output
+
+
+def get_row_pruning_shared_parameters(param_dict):
+    output = {}
+    if SHARED_PARAMETERS in param_dict.keys():
+        sub_param_dict = param_dict[SHARED_PARAMETERS]
+        output[ROW_PRUNING_ENABLED] = get_scalar_param(sub_param_dict, ROW_PRUNING_ENABLED,
+                                                       ROW_PRUNING_ENABLED_DEFAULT)
+        output[ROW_PRUNING_METHOD] = get_scalar_param(sub_param_dict, ROW_PRUNING_METHOD, ROW_PRUNING_METHOD_DEFAULT)
+        assert output[ROW_PRUNING_METHOD] in [
+            ROW_PRUNING_METHOD_L1, ROW_PRUNING_METHOD_TOPK
+        ], f"Invalid row pruning method. Supported types: [{ROW_PRUNING_METHOD_L1}, {ROW_PRUNING_METHOD_TOPK}]"
+        output[ROW_PRUNING_SCHEDULE_OFFSET] = get_scalar_param(sub_param_dict, ROW_PRUNING_SCHEDULE_OFFSET,
+                                                               ROW_PRUNING_SCHEDULE_OFFSET_DEFAULT)
+    else:
+        output[ROW_PRUNING_ENABLED] = ROW_PRUNING_ENABLED_DEFAULT
+        output[ROW_PRUNING_METHOD] = ROW_PRUNING_METHOD_DEFAULT
+        output[ROW_PRUNING_SCHEDULE_OFFSET] = ROW_PRUNING_SCHEDULE_OFFSET_DEFAULT
+    return output
+
+
+def get_row_pruning_different_groups(param_dict):
+    output = {}
+    sub_param_dict = param_dict[DIFFERENT_GROUPS]
+
+    def get_params(name, group_dict):
+        assert ROW_PRUNING_DENSE_RATIO in group_dict.keys(
+        ), f"{ROW_PRUNING_DENSE_RATIO} must be specified for row pruning group {name}"
+        return group_dict
+
+    for k, v in sub_param_dict.items():
+        output[k] = {}
+        output[k][DIFFERENT_GROUPS_PARAMETERS] = get_params(k, sub_param_dict[k][DIFFERENT_GROUPS_PARAMETERS])
+        output[k][DIFFERENT_GROUPS_MODULE_SCOPE] = get_scalar_param(sub_param_dict[k], DIFFERENT_GROUPS_MODULE_SCOPE,
+                                                                    DIFFERENT_GROUPS_MODULE_SCOPE_DEFAULT)
+        output[k][DIFFERENT_GROUPS_RELATED_MODULE_SCOPE] = get_scalar_param(
+            sub_param_dict[k], DIFFERENT_GROUPS_RELATED_MODULE_SCOPE, DIFFERENT_GROUPS_RELATED_MODULE_SCOPE_DEFAULT)
+    return output
+
+
+def get_head_pruning(param_dict):
+    output = {}
+    if HEAD_PRUNING not in param_dict.keys():
+        param_dict[HEAD_PRUNING] = {SHARED_PARAMETERS: {}, DIFFERENT_GROUPS: {}}
+    sub_param_dict = param_dict[HEAD_PRUNING]
+    # shared parameters
+    output[SHARED_PARAMETERS] = get_head_pruning_shared_parameters(sub_param_dict)
+    # each sub-groups
+    if output[SHARED_PARAMETERS][HEAD_PRUNING_ENABLED]:
+        assert DIFFERENT_GROUPS in sub_param_dict.keys(
+        ), f"Head Pruning is enabled, {DIFFERENT_GROUPS} must be specified"
+    output[DIFFERENT_GROUPS] = get_head_pruning_different_groups(sub_param_dict)
+    return output
+
+
+def get_head_pruning_shared_parameters(param_dict):
+    output = {}
+    if SHARED_PARAMETERS in param_dict.keys():
+        sub_param_dict = param_dict[SHARED_PARAMETERS]
+        output[HEAD_PRUNING_ENABLED] = get_scalar_param(sub_param_dict, HEAD_PRUNING_ENABLED,
+                                                        HEAD_PRUNING_ENABLED_DEFAULT)
+        output[HEAD_PRUNING_METHOD] = get_scalar_param(sub_param_dict, HEAD_PRUNING_METHOD,
+                                                       HEAD_PRUNING_METHOD_DEFAULT)
+        assert output[HEAD_PRUNING_METHOD] in [
+            HEAD_PRUNING_METHOD_L1, HEAD_PRUNING_METHOD_TOPK
+        ], f"Invalid head pruning method. Supported types: [{HEAD_PRUNING_METHOD_L1}, {HEAD_PRUNING_METHOD_TOPK}]"
+        output[HEAD_PRUNING_SCHEDULE_OFFSET] = get_scalar_param(sub_param_dict, HEAD_PRUNING_SCHEDULE_OFFSET,
+                                                                HEAD_PRUNING_SCHEDULE_OFFSET_DEFAULT)
+        if output[HEAD_PRUNING_ENABLED]:
+            assert HEAD_PRUNING_NUM_HEADS in sub_param_dict.keys(
+            ), f"{HEAD_PRUNING_NUM_HEADS} must be specified for head pruning"
+            output[HEAD_PRUNING_NUM_HEADS] = sub_param_dict[HEAD_PRUNING_NUM_HEADS]
+    else:
+        output[HEAD_PRUNING_ENABLED] = HEAD_PRUNING_ENABLED_DEFAULT
+        output[HEAD_PRUNING_METHOD] = HEAD_PRUNING_METHOD_DEFAULT
+        output[HEAD_PRUNING_SCHEDULE_OFFSET] = HEAD_PRUNING_SCHEDULE_OFFSET_DEFAULT
+    return output
+
+
+def get_head_pruning_different_groups(param_dict):
+    output = {}
+    sub_param_dict = param_dict[DIFFERENT_GROUPS]
+
+    def get_params(name, group_dict):
+        assert HEAD_PRUNING_DENSE_RATIO in group_dict.keys(
+        ), f"dense_ratio must be specified for head pruning group {name}"
+        return group_dict
+
+    for k, v in sub_param_dict.items():
+        output[k] = {}
+        output[k][DIFFERENT_GROUPS_PARAMETERS] = get_params(k, sub_param_dict[k][DIFFERENT_GROUPS_PARAMETERS])
+        output[k][DIFFERENT_GROUPS_MODULE_SCOPE] = get_scalar_param(sub_param_dict[k], DIFFERENT_GROUPS_MODULE_SCOPE,
+                                                                    DIFFERENT_GROUPS_MODULE_SCOPE_DEFAULT)
+        output[k][DIFFERENT_GROUPS_RELATED_MODULE_SCOPE] = get_scalar_param(
+            sub_param_dict[k], DIFFERENT_GROUPS_RELATED_MODULE_SCOPE, DIFFERENT_GROUPS_RELATED_MODULE_SCOPE_DEFAULT)
+    return output
+
+
+def get_channel_pruning(param_dict):
+    output = {}
+    if CHANNEL_PRUNING not in param_dict.keys():
+        param_dict[CHANNEL_PRUNING] = {SHARED_PARAMETERS: {}, DIFFERENT_GROUPS: {}}
+    sub_param_dict = param_dict[CHANNEL_PRUNING]
+    # shared parameters
+    output[SHARED_PARAMETERS] = get_channel_pruning_shared_parameters(sub_param_dict)
+    # each sub-groups
+    if output[SHARED_PARAMETERS][CHANNEL_PRUNING_ENABLED]:
+        assert DIFFERENT_GROUPS in sub_param_dict.keys(
+        ), f"Sparse Pruning is enabled, {DIFFERENT_GROUPS} must be specified"
+    output[DIFFERENT_GROUPS] = get_channel_pruning_different_groups(sub_param_dict)
+    return output
+
+
+def get_channel_pruning_shared_parameters(param_dict):
+    output = {}
+    if SHARED_PARAMETERS in param_dict.keys():
+        sub_param_dict = param_dict[SHARED_PARAMETERS]
+        output[CHANNEL_PRUNING_ENABLED] = get_scalar_param(sub_param_dict, CHANNEL_PRUNING_ENABLED,
+                                                           CHANNEL_PRUNING_ENABLED_DEFAULT)
+        output[CHANNEL_PRUNING_METHOD] = get_scalar_param(sub_param_dict, CHANNEL_PRUNING_METHOD,
+                                                          CHANNEL_PRUNING_METHOD_DEFAULT)
+        assert output[CHANNEL_PRUNING_METHOD] in [
+            CHANNEL_PRUNING_METHOD_L1, CHANNEL_PRUNING_METHOD_TOPK
+        ], f"Invalid channel pruning method. Supported types: [{CHANNEL_PRUNING_METHOD_L1}, {CHANNEL_PRUNING_METHOD_TOPK}]"
+        output[CHANNEL_PRUNING_SCHEDULE_OFFSET] = get_scalar_param(sub_param_dict, CHANNEL_PRUNING_SCHEDULE_OFFSET,
+                                                                   CHANNEL_PRUNING_SCHEDULE_OFFSET_DEFAULT)
+    else:
+        output[CHANNEL_PRUNING_ENABLED] = CHANNEL_PRUNING_ENABLED_DEFAULT
+        output[CHANNEL_PRUNING_METHOD] = CHANNEL_PRUNING_METHOD_DEFAULT
+        output[CHANNEL_PRUNING_SCHEDULE_OFFSET] = CHANNEL_PRUNING_SCHEDULE_OFFSET_DEFAULT
+    return output
+
+
+def get_channel_pruning_different_groups(param_dict):
+    output = {}
+    sub_param_dict = param_dict[DIFFERENT_GROUPS]
+
+    def get_params(name, group_dict):
+        assert CHANNEL_PRUNING_DENSE_RATIO in group_dict.keys(
+        ), f"{CHANNEL_PRUNING_DENSE_RATIO} must be specified for channel pruning group {name}"
+        return group_dict
+
+    for k, v in sub_param_dict.items():
+        output[k] = {}
+        output[k][DIFFERENT_GROUPS_PARAMETERS] = get_params(k, sub_param_dict[k][DIFFERENT_GROUPS_PARAMETERS])
+        output[k][DIFFERENT_GROUPS_MODULE_SCOPE] = get_scalar_param(sub_param_dict[k], DIFFERENT_GROUPS_MODULE_SCOPE,
+                                                                    DIFFERENT_GROUPS_MODULE_SCOPE_DEFAULT)
+        output[k][DIFFERENT_GROUPS_RELATED_MODULE_SCOPE] = get_scalar_param(
+            sub_param_dict[k], DIFFERENT_GROUPS_RELATED_MODULE_SCOPE, DIFFERENT_GROUPS_RELATED_MODULE_SCOPE_DEFAULT)
+
+    return output

lib/python3.12/site-packages/deepspeed/compression/constants.py

@@ -0,0 +1,188 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+#########################################
+# Compression Methods
+# It has several sub-components
+# #########################################
+COMPRESSION_TRAINING = "compression_training"
+SHARED_PARAMETERS = "shared_parameters"
+DIFFERENT_GROUPS = "different_groups"
+TECHNIQUE_ENABLED = "enabled"
+TECHNIQUE_SCHEDULE_OFFSET = "schedule_offset"
+TECHNIQUE_SCHEDULE_OFFSET_END = "schedule_offset_end"
+DIFFERENT_GROUPS_PARAMETERS = "params"
+DIFFERENT_GROUPS_MODULE_SCOPE = "modules"
+DIFFERENT_GROUPS_MODULE_SCOPE_DEFAULT = "*"
+DIFFERENT_GROUPS_RELATED_MODULE_SCOPE = "related_modules"
+DIFFERENT_GROUPS_RELATED_MODULE_SCOPE_DEFAULT = None
+# COMPRESSION_TRAINING_ENABLED = "enabled"
+# COMPRESSION_TRAINING_ENABLED_DEFAULT = False
+
+####
+# Layer Reduction
+####
+LAYER_REDUCTION = "layer_reduction"
+LAYER_REDUCTION_ENABLED = "enabled"
+LAYER_REDUCTION_ENABLED_DEFAULT = False
+KEEP_NUMBER_LAYER = "keep_number_layer"
+MODULE_NAME_PREFIX = "module_name_prefix"
+TEACHER_LAYER = "teacher_layer"
+OTHER_MODULE_NAME = "other_module_name"
+
+####
+# Weight Quantization
+####
+WEIGHT_QUANTIZATION = "weight_quantization"
+
+WEIGHT_QUANTIZATION_PERIOD = "quantization_period"
+WEIGHT_QUANTIZATION_PERIOD_DEFAULT = 1
+
+WEIGHT_QUANTIZE_IN_FORWARD_ENABLED = "quantize_weight_in_forward"
+WEIGHT_QUANTIZE_IN_FORWARD_ENABLED_DEFAULT = False
+
+WEIGHT_QUANTIZE_ENABLED = TECHNIQUE_ENABLED
+WEIGHT_QUANTIZE_ENABLED_DEFAULT = False
+
+WEIGHT_QUANTIZE_KERNEL = "quantizer_kernel"
+WEIGHT_QUANTIZE_KERNEL_DEFAULT = False
+
+WEIGHT_QUANTIZE_SCHEDULE_OFFSET = TECHNIQUE_SCHEDULE_OFFSET
+WEIGHT_QUANTIZE_SCHEDULE_OFFSET_DEFAULT = 0
+
+WEIGHT_QUANTIZE_GROUPS = "quantize_groups"
+WEIGHT_QUANTIZE_GROUPS_DEFAULT = 1
+
+WEIGHT_QUANTIZE_VERBOSE = "quantize_verbose"
+WEIGHT_QUANTIZE_VERBOSE_DEFAULT = False
+
+WEIGHT_QUANTIZE_TYPE = "quantization_type"
+WEIGHT_QUANTIZE_TYPE_DEFAULT = "symmetric"
+WEIGHT_QUANTIZE_SYMMETRIC = "symmetric"
+WEIGHT_QUANTIZE_ASYMMETRIC = "asymmetric"
+
+WEIGHT_QUANTIZE_ROUNDING = "rounding"
+WEIGHT_QUANTIZE_ROUNDING_DEFAULT = "nearest"
+WEIGHT_QUANTIZE_STOCHASTIC_ROUNDING = "stochastic"
+WEIGHT_QUANTIZE_NEAREST_ROUNDING = "nearest"
+# maybe deleted for a cleaner version
+WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE = "fp16_mixed_quantize"
+
+WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE_ENABLED = "enabled"
+WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE_ENABLED_DEFAULT = False
+
+WEIGHT_QUANTIZE_CHANGE_RATIO = "quantize_change_ratio"
+WEIGHT_QUANTIZE_CHANGE_RATIO_DEFAULT = 0.001
+
+WEIGHT_QUANTIZE_START_BITS = "start_bits"
+WEIGHT_QUANTIZE_TARGET_BITS = "target_bits"
+###
+# Activation Quantization
+###
+ACTIVATION_QUANTIZATION = "activation_quantization"
+
+ACTIVATION_QUANTIZATION_ENABLED = TECHNIQUE_ENABLED
+ACTIVATION_QUANTIZATION_ENABLED_DEFAULT = False
+
+ACTIVATION_QUANTIZE_SCHEDULE_OFFSET = TECHNIQUE_SCHEDULE_OFFSET
+ACTIVATION_QUANTIZE_SCHEDULE_OFFSET_DEFAULT = 1000
+
+ACTIVATION_QUANTIZE_TYPE = "quantization_type"
+ACTIVATION_QUANTIZE_TYPE_DEFAULT = "symmetric"
+ACTIVATION_QUANTIZE_SYMMETRIC = "symmetric"
+ACTIVATION_QUANTIZE_ASYMMETRIC = "asymmetric"
+
+ACTIVATION_QUANTIZE_RANGE = 'range_calibration'
+ACTIVATION_QUANTIZE_RANGE_DEFAULT = 'dynamic'
+ACTIVATION_QUANTIZE_RANGE_STATIC = 'static'
+ACTIVATION_QUANTIZE_RANGE_DYNAMIC = 'dynamic'
+
+ACTIVATION_QUANTIZE_BITS = "bits"
+###
+# Sparse Pruning
+###
+SPARSE_PRUNING = "sparse_pruning"
+
+SPARSE_PRUNING_ENABLED = TECHNIQUE_ENABLED
+SPARSE_PRUNING_ENABLED_DEFAULT = False
+
+SPARSE_PRUNING_METHOD = "method"
+SPARSE_PRUNING_METHOD_DEFAULT = "l1"
+SPARSE_PRUNING_METHOD_L1 = "l1"
+SPARSE_PRUNING_METHOD_TOPK = "topk"
+SPARSE_PRUNING_METHOD_SNIP_MOMENTUM = "snip_momentum"
+
+SPARSE_PRUNING_BLOCK_PATTERN = "block_pattern"
+SPARSE_PRUNING_BLOCK_PATTERN_DEFAULT = "4x1"
+
+SPARSE_PRUNING_SCHEDULE_OFFSET_STRIDE = "schedule_offset_stride"
+SPARSE_PRUNING_SCHEDULE_OFFSET_STRIDE_DEFAULT = 1
+
+SPARSE_PRUNING_SCHEDULE_OFFSET = TECHNIQUE_SCHEDULE_OFFSET
+SPARSE_PRUNING_SCHEDULE_OFFSET_DEFAULT = 1000
+
+SPARSE_PRUNING_SCHEDULE_OFFSET_END = TECHNIQUE_SCHEDULE_OFFSET_END
+SPARSE_PRUNING_SCHEDULE_OFFSET_END_DEFAULT = SPARSE_PRUNING_SCHEDULE_OFFSET_DEFAULT
+
+SPARSE_PRUNING_DENSE_RATIO = "dense_ratio"
+SPARSE_PRUNING_DENSE_RATIO_DEFAULT = 0.1
+
+SPARSE_PRUNING_EXCLUDED_MODULES = "excluded_modules"
+SPARSE_PRUNING_EXCLUDED_MODULES_DEFAULT = []
+###
+# Row Pruning
+###
+ROW_PRUNING = "row_pruning"
+
+ROW_PRUNING_ENABLED = TECHNIQUE_ENABLED
+ROW_PRUNING_ENABLED_DEFAULT = False
+
+ROW_PRUNING_METHOD = "method"
+ROW_PRUNING_METHOD_DEFAULT = "l1"
+ROW_PRUNING_METHOD_L1 = "l1"
+ROW_PRUNING_METHOD_TOPK = "topk"
+
+ROW_PRUNING_SCHEDULE_OFFSET = TECHNIQUE_SCHEDULE_OFFSET
+ROW_PRUNING_SCHEDULE_OFFSET_DEFAULT = 1000
+
+ROW_PRUNING_DENSE_RATIO = "dense_ratio"
+
+###
+# Head Pruning
+###
+HEAD_PRUNING = "head_pruning"
+
+HEAD_PRUNING_ENABLED = TECHNIQUE_ENABLED
+HEAD_PRUNING_ENABLED_DEFAULT = False
+
+HEAD_PRUNING_METHOD = "method"
+HEAD_PRUNING_METHOD_DEFAULT = "topk"
+HEAD_PRUNING_METHOD_L1 = "l1"
+HEAD_PRUNING_METHOD_TOPK = "topk"
+
+HEAD_PRUNING_SCHEDULE_OFFSET = TECHNIQUE_SCHEDULE_OFFSET
+HEAD_PRUNING_SCHEDULE_OFFSET_DEFAULT = 1000
+
+HEAD_PRUNING_NUM_HEADS = "num_heads"
+
+HEAD_PRUNING_DENSE_RATIO = "dense_ratio"
+
+###
+# Channel Pruning
+###
+CHANNEL_PRUNING = "channel_pruning"
+
+CHANNEL_PRUNING_ENABLED = TECHNIQUE_ENABLED
+CHANNEL_PRUNING_ENABLED_DEFAULT = False
+
+CHANNEL_PRUNING_METHOD = "method"
+CHANNEL_PRUNING_METHOD_DEFAULT = "l1"
+CHANNEL_PRUNING_METHOD_L1 = "l1"
+CHANNEL_PRUNING_METHOD_TOPK = "topk"
+
+CHANNEL_PRUNING_SCHEDULE_OFFSET = TECHNIQUE_SCHEDULE_OFFSET
+CHANNEL_PRUNING_SCHEDULE_OFFSET_DEFAULT = 1000
+
+CHANNEL_PRUNING_DENSE_RATIO = "dense_ratio"

lib/python3.12/site-packages/deepspeed/compression/helper.py

@@ -0,0 +1,322 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from .basic_layer import Embedding_Compress, LinearLayer_Compress, Conv2dLayer_Compress, BNLayer_Compress, ColumnParallelLinear_Compress, RowParallelLinear_Compress
+from .constants import *
+from deepspeed.utils import logger
+
+try:
+    from neural_compressor.compression import pruner as nc_pruner
+except ImportError as e:
+    nc_pruner = None
+
+
+def recursive_getattr(model, module_name):
+    """
+    Recursively get the attribute of a module.
+    Args:
+        model (`torch.nn.Module`)
+            The model to get the attribute from.
+        module_name (`str`)
+            The name of the module to get the attribute from.
+    """
+    split_list = module_name.split('.')
+    output = model
+    for name in split_list:
+        output = getattr(output, name)
+    return output
+
+
+def recursive_setattr(model, module_name, module):
+    """
+    Recursively set the attribute of a module.
+    Args:
+        model (`torch.nn.Module`)
+            The model to set the attribute in.
+        module_name (`str`)
+            The name of the module to set the attribute in.
+        module (`torch.nn.Module`)
+            The module to set the attribute to.
+    """
+    split_list = module_name.split('.')
+    output = model
+    for name in split_list[:-1]:
+        output = getattr(output, name)
+    output.__setattr__(split_list[-1], module)
+
+
+def module_replacement(model, module_name, compression_technique=None, mpu=None):
+    """
+    Replace a module with a new module.
+    Args:
+        model (`torch.nn.Module`)
+            The model to replace the module in.
+        module_name (`str`)
+            The name of the module to replace.
+        compression_technique (`str`)
+            The compression technique to use for the new module.
+    """
+
+    # Get the old module
+    old_module = recursive_getattr(model, module_name)
+
+    need_bias = False
+    if hasattr(old_module, 'bias') and old_module.bias is not None:
+        need_bias = True
+
+    # Initialize the new module
+    if isinstance(old_module, LinearLayer_Compress) or isinstance(old_module, torch.nn.Linear):
+        if isinstance(old_module, LinearLayer_Compress):
+            new_module = old_module
+        else:
+            new_module = LinearLayer_Compress(old_module.in_features, old_module.out_features,
+                                              bias=need_bias).to(device=old_module.weight.device,
+                                                                 dtype=old_module.weight.dtype)
+            new_module.weight.data = old_module.weight.data
+            if need_bias:
+                new_module.bias.data = old_module.bias.data
+    elif isinstance(old_module, Conv2dLayer_Compress) or isinstance(old_module, torch.nn.Conv2d):
+        if isinstance(old_module, Conv2dLayer_Compress):
+            new_module = old_module
+        else:
+            new_module = Conv2dLayer_Compress(old_module.in_channels, old_module.out_channels, old_module.kernel_size, old_module.stride, old_module.padding, \
+                                            old_module.dilation, old_module.groups, need_bias, \
+                                            old_module.padding_mode).to(device=old_module.weight.device, dtype=old_module.weight.dtype)
+            new_module.weight.data = old_module.weight.data
+            if need_bias:
+                new_module.bias.data = old_module.bias.data
+    elif isinstance(old_module, torch.nn.BatchNorm2d):
+        new_module = BNLayer_Compress(old_module.num_features, old_module.eps, old_module.momentum, old_module.affine,
+                                      old_module.track_running_stats).to(old_module.weight.device,
+                                                                         old_module.weight.dtype)
+        new_module.weight.data = old_module.weight.data
+        if need_bias:
+            new_module.bias.data = old_module.bias.data
+        new_module.running_mean.data = old_module.running_mean.data
+        new_module.running_var.data = old_module.running_var.data
+    elif isinstance(old_module, Embedding_Compress) or isinstance(old_module, torch.nn.Embedding):
+        if isinstance(old_module, Embedding_Compress):
+            new_module = old_module
+        else:
+            new_module = Embedding_Compress(old_module.num_embeddings, old_module.embedding_dim, old_module.padding_idx, old_module.max_norm, old_module.norm_type, \
+                                        old_module.scale_grad_by_freq, old_module.sparse).to(device=old_module.weight.device, dtype=old_module.weight.dtype)
+            new_module.weight.data = old_module.weight.data
+    elif mpu is not None and (isinstance(old_module, ColumnParallelLinear_Compress)
+                              or isinstance(old_module, mpu.ColumnParallelLinear)):
+        if isinstance(old_module, ColumnParallelLinear_Compress):
+            new_module = old_module
+        else:
+            new_module = ColumnParallelLinear_Compress(mpu,
+                                                       old_module.input_size,
+                                                       old_module.output_size,
+                                                       gather_output=old_module.gather_output,
+                                                       skip_bias_add=old_module.skip_bias_add,
+                                                       bias=need_bias).to(device=old_module.weight.device,
+                                                                          dtype=old_module.weight.dtype)
+            new_module.weight.data = old_module.weight.data
+            if need_bias:
+                new_module.bias.data = old_module.bias.data
+    elif mpu is not None and (isinstance(old_module, RowParallelLinear_Compress)
+                              or isinstance(old_module, mpu.RowParallelLinear)):
+        if isinstance(old_module, RowParallelLinear_Compress):
+            new_module = old_module
+        else:
+            new_module = RowParallelLinear_Compress(mpu,
+                                                    old_module.input_size,
+                                                    old_module.output_size,
+                                                    input_is_parallel=old_module.input_is_parallel,
+                                                    skip_bias_add=old_module.skip_bias_add,
+                                                    bias=need_bias).to(device=old_module.weight.device,
+                                                                       dtype=old_module.weight.dtype)
+            new_module.weight.data = old_module.weight.data
+            if need_bias:
+                new_module.bias.data = old_module.bias.data
+    else:
+        new_module = None
+
+    if compression_technique is not None:
+        for k, v in compression_technique.items():
+            if k == SPARSE_PRUNING:
+                if v[SPARSE_PRUNING_ENABLED]:
+                    new_module.enable_sparse_pruning(v[SPARSE_PRUNING_DENSE_RATIO], v[SPARSE_PRUNING_METHOD])
+            elif k == ROW_PRUNING:
+                if v[ROW_PRUNING_ENABLED]:
+                    new_module.enable_row_pruning(v[ROW_PRUNING_DENSE_RATIO], v[ROW_PRUNING_METHOD])
+            elif k == HEAD_PRUNING:
+                if v[HEAD_PRUNING_ENABLED]:
+                    new_module.enable_head_pruning(v[HEAD_PRUNING_DENSE_RATIO], v[HEAD_PRUNING_METHOD],
+                                                   v[HEAD_PRUNING_NUM_HEADS])
+            elif k == ACTIVATION_QUANTIZATION:
+                if v[ACTIVATION_QUANTIZATION_ENABLED]:
+                    new_module.enable_activation_quantization(v[ACTIVATION_QUANTIZE_BITS], v[ACTIVATION_QUANTIZE_TYPE],
+                                                              v[ACTIVATION_QUANTIZE_RANGE])
+            elif k == WEIGHT_QUANTIZATION:
+                if v[WEIGHT_QUANTIZE_ENABLED]:
+                    new_module.enable_weight_quantization(v[WEIGHT_QUANTIZE_START_BITS],
+                                                          v[WEIGHT_QUANTIZE_TARGET_BITS],
+                                                          v[WEIGHT_QUANTIZATION_PERIOD],
+                                                          v[WEIGHT_QUANTIZE_IN_FORWARD_ENABLED],
+                                                          v[WEIGHT_QUANTIZE_TYPE], v[WEIGHT_QUANTIZE_GROUPS])
+            elif k == CHANNEL_PRUNING:
+                if v[CHANNEL_PRUNING_ENABLED]:
+                    new_module.enable_channel_pruning(v[CHANNEL_PRUNING_DENSE_RATIO], v[CHANNEL_PRUNING_METHOD])
+            else:
+                raise NotImplementedError('Compression technique {} is not implemented'.format(k))
+
+    # Replace the old module with the new one
+    recursive_setattr(model, module_name, new_module)
+
+
+def is_module_compressible(module, mpu=None):
+    ret = isinstance(module, torch.nn.Linear) or \
+          isinstance(module, torch.nn.Conv2d) or \
+          isinstance(module, torch.nn.Embedding) or \
+          isinstance(module, torch.nn.BatchNorm2d)
+
+    if mpu is not None:
+        ret = ret or isinstance(module, mpu.RowParallelLinear) or isinstance(module, mpu.ColumnParallelLinear)
+
+    return ret
+
+
+def compression_preparation(model, compression_technique_list, mpu):
+    """
+    Prepare the compression techniques of a model.
+    Args:
+        model (`torch.nn.Module`)
+            The model to prepare the compression techniques of.
+        compression_technique_list (`list`)
+            The list of compression techniques to prepare the model to.
+            list[]
+    """
+    # Here we first replace all module with our linear wrapper
+    for module_name, module in model.named_modules():
+        if is_module_compressible(module, mpu):
+            module_replacement(model, module_name, mpu=mpu)
+    for module_name_lists, _, compression_technique in compression_technique_list:
+        for mnl in module_name_lists:
+            for module_name in mnl:
+                module_replacement(model, module_name, compression_technique)
+
+    return model
+
+
+def fix_compression(model, module_name, compression_technique, mask=None, dim_reduction=False):
+    """
+    Fix the compression technique of a module.
+    Args:
+        model (`torch.nn.Module`)
+            The model to fix the compression technique of.
+        module_name (`str`)
+            The name of the module to fix the compression technique of.
+        compression_technique (`str`)
+            The compression technique to fix the module to.
+    """
+    # Here we can make things much simpler by just replacing the module
+    module = recursive_getattr(model, module_name)
+    for k, v in compression_technique.items():
+        if k == WEIGHT_QUANTIZATION and v[WEIGHT_QUANTIZE_IN_FORWARD_ENABLED] and v[WEIGHT_QUANTIZE_ENABLED]:
+            return module.fix_weight_quantization()
+        elif k == SPARSE_PRUNING and v[SPARSE_PRUNING_ENABLED]:
+            return module.fix_sparse_pruning_helper()
+        elif k == ROW_PRUNING and (v[ROW_PRUNING_ENABLED] or mask is not None):
+            return module.fix_row_col_pruning_helper(mask, dim_reduction=dim_reduction)
+        elif k == HEAD_PRUNING and (v[HEAD_PRUNING_ENABLED] or mask is not None):
+            return module.fix_head_pruning_helper(mask, v[HEAD_PRUNING_NUM_HEADS], dim_reduction=dim_reduction)
+        elif k == CHANNEL_PRUNING and (v[CHANNEL_PRUNING_ENABLED] or mask is not None):
+            return module.fix_channel_pruning_helper(mask, dim_reduction=dim_reduction)
+
+
+def convert_conv1d_to_linear(model, convert_type):
+    '''
+    This is a help function to convert conv1d to linear (e.g., convert GPT2 from HF)
+    '''
+    if hasattr(model, 'module'):
+        c_model = model.module
+    else:
+        c_model = model
+
+    for name, module in c_model.named_modules():
+        if isinstance(module, convert_type):
+            old_module = recursive_getattr(c_model, name)
+            new_module = torch.nn.Linear(old_module.weight.data.size(0),
+                                         old_module.weight.data.size(1),
+                                         bias=True if old_module.bias is not None else False)
+            new_module.weight.data = old_module.weight.data.t().contiguous()
+            if new_module.bias is not None:
+                new_module.bias.data = old_module.bias.data.view(-1)
+
+            recursive_setattr(c_model, name, new_module)
+
+    return model
+
+
+def generate_pruners(config, model):
+    """Generate pruners.
+    Args:
+        config (`neural_compressor.WeightPruningConfig`)
+            The object to the class WeightPruningConfig.
+        model (`torch.nn.module`)
+            The torch module object to be pruned.
+    """
+    assert nc_pruner is not None, "please ensure the neural_compressor python package is installed by pip or conda if user wants to use snip_momentum sparse pruning"
+    from nc_pruner.utils import process_config, parse_to_prune
+    from nc_pruner.pruners import get_pruner
+    assert isinstance(model, torch.nn.Module)
+    pruners_info = process_config(config)
+    pruners = []
+    for info in pruners_info:
+        modules = parse_to_prune(info, model)
+        if modules == {}:
+            logger.warning("one pruner hooks no layers, please have a check")
+
+        pruners.append(get_pruner(info, modules))
+        info['modules'] = [key for key in modules.keys()]
+        info['len_of_modules'] = len(info['modules'])
+        logger.info(info)
+    return pruners
+
+
+def register_on_step_begin(model):
+    """Mount on_step_begin to the model.
+    Args:
+        model (`torch.nn.module`)
+            The torch module object to be pruned.
+    """
+
+    def hook(module, input):
+        for pruner in module.pruners:
+            pruner.on_step_begin(0)
+
+    hook_handle = model.register_forward_pre_hook(hook)
+    return hook_handle
+
+
+def rewrite_optimizer_step(opt: torch.optim.Optimizer):
+    """Mount on_before/after_optimizer_step to the optimizer.
+    Args:
+        model (`torch.opt.Optimizer`)
+            The torch optimizer object to be hooked.
+    """
+
+    def new_step(self, closure=None):
+        if hasattr(self, "pruners"):
+            for pruner in self.pruners:
+                pruner.on_before_optimizer_step()
+
+        if closure is not None:
+            res = self.orig_step(closure)
+        else:
+            res = self.orig_step()
+        if hasattr(self, "pruners"):
+            for pruner in self.pruners:
+                pruner.on_after_optimizer_step()
+        return res
+
+    opt.orig_step = opt.step
+    import types
+    opt.step = types.MethodType(new_step, opt)
+    return opt

lib/python3.12/site-packages/deepspeed/compression/scheduler.py

@@ -0,0 +1,173 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .compress import get_module_name
+from .constants import *
+from .helper import recursive_getattr
+from deepspeed.utils import logger
+
+
+class compression_scheduler():
+    '''
+    Used to schedule different compression methods
+    '''
+
+    def __init__(self, model, compression_config):
+        self.model = model
+        self.compression_config = compression_config
+        self.make_init()
+        self.training_steps = 0
+        self.weight_quantization_enabled = False
+
+        self.verbose = {
+            WEIGHT_QUANTIZATION: False,
+            ACTIVATION_QUANTIZATION: False,
+            SPARSE_PRUNING: False,
+            HEAD_PRUNING: False,
+            ROW_PRUNING: False,
+            CHANNEL_PRUNING: False
+        }
+
+    def make_init(self):
+        self.different_compression_methods = {}
+        for method, method_content in self.compression_config.items():
+            if LAYER_REDUCTION in method:
+                continue
+            self.different_compression_methods[method] = {
+                TECHNIQUE_ENABLED: False,
+                SHARED_PARAMETERS: None,
+                DIFFERENT_GROUPS: []
+            }
+            exist_module_name = set()
+            shared_parameters = method_content[SHARED_PARAMETERS]
+            self.different_compression_methods[method][TECHNIQUE_ENABLED] = shared_parameters[TECHNIQUE_ENABLED]
+            self.different_compression_methods[method][SHARED_PARAMETERS] = shared_parameters
+
+            for group_name, method_parameters in method_content[DIFFERENT_GROUPS].items():
+                module_name_list = []
+                for key_word in method_parameters[DIFFERENT_GROUPS_MODULE_SCOPE]:
+                    module_name, exist_module_name = get_module_name(group_name,
+                                                                     self.model,
+                                                                     key_word,
+                                                                     exist_module_name,
+                                                                     verbose=False)
+                    module_name_list.extend(module_name)
+                if module_name_list:
+                    self.different_compression_methods[method][DIFFERENT_GROUPS].append(
+                        [group_name, module_name_list,
+                         method_parameters.copy().pop('params')])
+
+    def check_weight_quantization(self):
+        # check weight quantization
+        wq = self.different_compression_methods[WEIGHT_QUANTIZATION]
+        if not wq[TECHNIQUE_ENABLED]:
+            return
+        else:
+            shared_parameters = wq[SHARED_PARAMETERS]
+            if self.training_steps >= shared_parameters[TECHNIQUE_SCHEDULE_OFFSET]:
+                for group_name, module_name_list, method_parameters in wq[DIFFERENT_GROUPS]:
+                    for module_name in module_name_list:
+                        module = recursive_getattr(self.model, module_name)
+                        module.weight_quantization_enabled = True
+
+                if not self.verbose[WEIGHT_QUANTIZATION]:
+                    logger.info(f'Weight quantization is enabled at step {self.training_steps}')
+                    self.weight_quantization_enabled = True
+                    self.verbose[WEIGHT_QUANTIZATION] = True
+
+    def check_activation_quantization(self):
+        # check activation quantization
+        aq = self.different_compression_methods[ACTIVATION_QUANTIZATION]
+        if not aq[TECHNIQUE_ENABLED]:
+            return
+        else:
+            shared_parameters = aq[SHARED_PARAMETERS]
+            if self.training_steps >= shared_parameters[TECHNIQUE_SCHEDULE_OFFSET]:
+                for group_name, module_name_list, method_parameters in aq[DIFFERENT_GROUPS]:
+                    for module_name in module_name_list:
+                        module = recursive_getattr(self.model, module_name)
+                        module.activation_quantization_enabled = True
+                if not self.verbose[ACTIVATION_QUANTIZATION]:
+                    logger.info(f'Activation quantization is enabled at step {self.training_steps}')
+                    self.verbose[ACTIVATION_QUANTIZATION] = True
+
+    def check_sparse_pruning(self):
+        # check sparse pruning
+        sp = self.different_compression_methods[SPARSE_PRUNING]
+        if not sp[TECHNIQUE_ENABLED]:
+            return
+        else:
+            shared_parameters = sp[SHARED_PARAMETERS]
+            if shared_parameters[TECHNIQUE_SCHEDULE_OFFSET] <= self.training_steps <= shared_parameters[
+                    TECHNIQUE_SCHEDULE_OFFSET_END]:
+                for group_name, module_name_list, method_parameters in sp[DIFFERENT_GROUPS]:
+                    for module_name in module_name_list:
+                        module = recursive_getattr(self.model, module_name)
+                        module.sparse_pruning_enabled = True
+                if not self.verbose[SPARSE_PRUNING]:
+                    logger.info(f'Sparse pruning is enabled at step {self.training_steps}')
+                    self.verbose[SPARSE_PRUNING] = True
+
+    def check_head_pruning(self):
+        # check head pruning
+        hp = self.different_compression_methods[HEAD_PRUNING]
+        if not hp[TECHNIQUE_ENABLED]:
+            return
+        else:
+            shared_parameters = hp[SHARED_PARAMETERS]
+            if self.training_steps >= shared_parameters[TECHNIQUE_SCHEDULE_OFFSET]:
+                for group_name, module_name_list, method_parameters in hp[DIFFERENT_GROUPS]:
+                    for module_name in module_name_list:
+                        module = recursive_getattr(self.model, module_name)
+                        module.head_pruning_enabled = True
+                if not self.verbose[HEAD_PRUNING]:
+                    logger.info(f'Head pruning is enabled at step {self.training_steps}')
+                    self.verbose[HEAD_PRUNING] = True
+
+    def check_row_pruning(self):
+        # check row pruning
+        rp = self.different_compression_methods[ROW_PRUNING]
+        if not rp[TECHNIQUE_ENABLED]:
+            return
+        else:
+            shared_parameters = rp[SHARED_PARAMETERS]
+            if self.training_steps >= shared_parameters[TECHNIQUE_SCHEDULE_OFFSET]:
+                for group_name, module_name_list, method_parameters in rp[DIFFERENT_GROUPS]:
+                    for module_name in module_name_list:
+                        module = recursive_getattr(self.model, module_name)
+                        module.row_pruning_enabled = True
+                if not self.verbose[ROW_PRUNING]:
+                    logger.info(f'Row pruning is enabled at step {self.training_steps}')
+                    self.verbose[ROW_PRUNING] = True
+
+    def check_channel_pruning(self):
+        # check channel pruning
+        cp = self.different_compression_methods[CHANNEL_PRUNING]
+        if not cp[TECHNIQUE_ENABLED]:
+            return
+        else:
+            shared_parameters = cp[SHARED_PARAMETERS]
+            if self.training_steps >= shared_parameters[TECHNIQUE_SCHEDULE_OFFSET]:
+                for group_name, module_name_list, method_parameters in cp[DIFFERENT_GROUPS]:
+                    for module_name in module_name_list:
+                        module = recursive_getattr(self.model, module_name)
+                        module.channel_pruning_enabled = True
+                if not self.verbose[CHANNEL_PRUNING]:
+                    logger.info(f'Channel pruning is enabled at step {self.training_steps}')
+                    self.verbose[CHANNEL_PRUNING] = True
+
+    def check_all_modules(self):
+        # check all different compression methods we have
+        self.check_weight_quantization()
+        self.check_activation_quantization()
+        self.check_sparse_pruning()
+        self.check_head_pruning()
+        self.check_row_pruning()
+        self.check_channel_pruning()
+
+    def step(self, step_zero_check=False):
+        if not step_zero_check:
+            self.training_steps += 1
+        self.check_all_modules()

lib/python3.12/site-packages/deepspeed/compression/utils.py

@@ -0,0 +1,222 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from torch import autograd
+import math
+
+
+class TopKBinarizer(autograd.Function):
+    """
+    Top-k Binarizer.
+    Computes a binary mask M from a real value matrix S such that `M_{i,j} = 1` if and only if `S_{i,j}`
+    is among the k% highest values of S.
+    Implementation is inspired from:
+        https://github.com/yaozhewei/MLPruning
+    """
+
+    @staticmethod
+    def forward(ctx, inputs: torch.tensor, threshold: float, sigmoid: bool):
+        """
+        Args:
+            inputs (`torch.FloatTensor`)
+                The input matrix from which the binarizer computes the binary mask.
+            threshold (`float`)
+                The percentage of weights to keep (the rest is pruned).
+                `threshold` is a float between 0 and 1.
+            sigmoid (`bool`)
+                Whether to apply a sigmoid on the threshold
+        Returns:
+            mask (`torch.FloatTensor`)
+                Binary matrix of the same size as `inputs` acting as a mask (1 - the associated weight is
+                retained, 0 - the associated weight is pruned).
+        """
+        # Get the subnetwork by sorting the inputs and using the top threshold
+        if sigmoid:
+            threshold = torch.sigmoid(threshold).item()
+        ctx.sigmoid = sigmoid
+        mask = inputs.clone()
+
+        _, idx = inputs.flatten().sort(descending=True)
+        j = math.ceil(threshold * inputs.numel())
+
+        # flat_out and mask access the same memory.
+        flat_out = mask.flatten()
+        flat_out[idx[j:]] = 0.
+        flat_out[idx[:j]] = 1.
+        ctx.save_for_backward(mask)
+
+        return mask
+
+    @staticmethod
+    def backward(ctx, gradOutput):
+        mask, = ctx.saved_tensors
+        if ctx.sigmoid:
+            return gradOutput.clone(), ((gradOutput * mask).sum()).view(-1), None
+        else:
+            return gradOutput.clone(), None, None
+
+
+class SymQuantizer(torch.autograd.Function):
+    """
+    Symmetric quantization
+    """
+
+    @staticmethod
+    def forward(ctx, input, num_bits, min_value=None, max_value=None, num_groups=1):
+        """
+        Args:
+            inputs (`torch.FloatTensor`)
+                The input which needs to be quantized
+            num_bits (int, >=4)
+                Number of bits to use for quantization
+            min_value/max_value (torch.FloatTensor)
+                Used for static activation quantization
+            num_groups (int)
+                How many groups to partition the quantization into
+        Returns:
+            quantized_input (`torch.FloatTensor`)
+                Quantized input
+        """
+        assert (min_value is None and max_value is None) or (min_value is not None and max_value is not None
+                                                             and num_groups == 1)
+        q_range = 2**num_bits
+        input_shape = input.shape
+        if min_value is None:
+            input = input.reshape(num_groups, -1)
+            max_input = torch.amax(torch.abs(input), dim=-1).view(num_groups, -1)
+        else:
+            max_input = torch.max(min_value.abs(), max_value).view(-1)
+
+        scale = 2 * max_input / q_range
+        output = (input / scale).round().clamp(-q_range // 2, q_range // 2 - 1) * scale
+        output = output.reshape(input_shape).contiguous()
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = grad_output.clone()
+        return grad_input, None, None, None, None
+
+
+class AsymQuantizer(torch.autograd.Function):
+    """
+    Asymmetric quantization
+    """
+
+    @staticmethod
+    def forward(ctx, input, num_bits, min_value=None, max_value=None, num_groups=1):
+        """
+        Args:
+            inputs (`torch.FloatTensor`)
+                The input which needs to be quantized
+            num_bits (int, >=4)
+                Number of bits to use for quantization
+            min_value/max_value (torch.FloatTensor)
+                Used for static activation quantization
+            num_groups (int)
+                How many groups to partition the quantization into
+        Returns:
+            quantized_input (`torch.FloatTensor`)
+                Quantized input
+        """
+
+        assert (min_value is None and max_value is None) or (min_value is not None and max_value is not None
+                                                             and num_groups == 1)
+        q_range = 2**num_bits
+        input_shape = input.shape
+        if min_value is None:
+            input = input.reshape(num_groups, -1)
+            min_value = input.amin(dim=-1, keepdim=True)
+            max_value = input.amax(dim=-1, keepdim=True)
+
+        scale = (max_value - min_value) / q_range
+        zero_point = (min_value / scale).round() * scale
+
+        output = ((input - zero_point) / scale).round().clamp(0, q_range - 1) * scale + zero_point
+        output = output.reshape(input_shape).contiguous()
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = grad_output.clone()
+        return grad_input, None, None, None, None
+
+
+class TernaryQuantizer(torch.autograd.Function):
+    """
+    Ternary quantization
+    """
+
+    @staticmethod
+    def forward(ctx, input, num_bits, min_value=None, max_value=None, num_groups=1):
+        """
+        Args:
+            inputs (`torch.FloatTensor`)
+                The input which needs to be quantized
+            num_bits (int)
+                Dummy variable
+            min_value/max_value (torch.FloatTensor)
+                Used for static activation quantization; for now they are dummy variable
+            num_groups (int)
+                How many groups to partition the quantization into
+        Returns:
+            quantized_input (`torch.FloatTensor`)
+                Quantized input
+        """
+
+        assert (min_value is None and max_value is None)
+        input_flat = input.reshape(num_groups, -1)
+        n = input_flat.shape[1]
+        m = input_flat.norm(p=1, dim=1).div(n)
+        thres = (0.7 * m).view(-1, 1)
+        pos = (input_flat > thres).type(input.type())
+        neg = (input_flat < -thres).type(input.type())
+        mask = (input_flat.abs() > thres).type(input.type())
+        alpha = ((mask * input_flat).abs().sum(dim=1) / mask.sum(dim=1)).view(-1, 1)
+        output = alpha * pos - alpha * neg
+        output = output.reshape(input.shape).contiguous()
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = grad_output.clone()
+        return grad_input, None, None, None, None
+
+
+class BinaryQuantizer(torch.autograd.Function):
+    """
+    Binary quantization
+    """
+
+    @staticmethod
+    def forward(ctx, input, num_bits, min_value=None, max_value=None, num_groups=1):
+        """
+        Args:
+            inputs (`torch.FloatTensor`)
+                The input which needs to be quantized
+            num_bits (int)
+                Dummy variable
+            min_value/max_value (torch.FloatTensor)
+                Used for static activation quantization; for now they are dummy variable
+            num_groups (int)
+                How many groups to partition the quantization into
+        Returns:
+            quantized_input (`torch.FloatTensor`)
+                Quantized input
+        """
+
+        assert (min_value is None and max_value is None)
+        input_flat = input.reshape(num_groups, -1)
+        n = input_flat.shape[1]
+        m = input_flat.norm(p=1, dim=1, keepdim=True).div(n)
+        output = input_flat.sign().mul(m)
+        output = output.reshape(input.shape).contiguous()
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = grad_output.clone()
+        return grad_input, None, None, None, None

lib/python3.12/site-packages/deepspeed/ops/compile/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from ..op_builder import DeepCompileBuilder

lib/python3.12/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp

@@ -0,0 +1,215 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "cpu_adagrad.h"
+#include <torch/extension.h>
+#include <functional>
+#include <iostream>
+#include <map>
+#include <memory>
+#include <type_traits>
+#include <unordered_map>
+
+using namespace std::string_literals;
+static std::unordered_map<int, std::shared_ptr<void>> s_optimizers;
+
+// C++ interface
+
+template <typename ds_params_precision_t, typename ds_state_precision_t>
+void Adagrad_Optimizer::Step_1(ds_params_precision_t* _params,
+                               ds_params_precision_t* grads,
+                               ds_state_precision_t* _exp_avg_sq,
+                               size_t _param_size)
+{
+    size_t rounded_size = 0;
+#if defined(__AVX512__) or defined(__AVX256__)
+    Step_AVX<1>(&rounded_size, _params, grads, _exp_avg_sq, _param_size);
+#endif
+    if (_param_size > rounded_size) {
+        float step_size = -1 * _alpha;
+        for (size_t t = rounded_size; t < _param_size; t += TILE) {
+            size_t copy_size = TILE;
+            if ((t + TILE) > _param_size) copy_size = _param_size - t;
+            size_t offset = copy_size + t;
+#pragma omp parallel for
+            for (size_t k = t; k < offset; k++) {
+                float grad = (float)grads[k];
+                float param = (float)_params[k];
+                float momentum = grads[k];
+                float variance = _exp_avg_sq[k];
+                if (_weight_decay > 0) { grad = param * _weight_decay + grad; }
+
+                variance += grad * grad;
+
+                grad = sqrt(variance);
+                grad += _eps;
+                grad = momentum / grad;
+                param = grad * step_size + param;
+                _params[k] = param;
+                // STORE UPDATE TERM TO GRAD'S MEMORY
+                grads[k] = grad * step_size;
+                _exp_avg_sq[k] = variance;
+            }
+        }
+    }
+}
+
+template <typename ds_params_precision_t, typename ds_state_precision_t>
+void Adagrad_Optimizer::Step_4(ds_params_precision_t* _params,
+                               ds_params_precision_t* grads,
+                               ds_state_precision_t* _exp_avg_sq,
+                               size_t _param_size)
+{
+    size_t rounded_size = 0;
+#if defined(__AVX512__) or defined(__AVX256__)
+    Step_AVX<4>(&rounded_size, _params, grads, _exp_avg_sq, _param_size);
+#endif
+    if (_param_size > rounded_size)
+        Step_1((_params + rounded_size),
+               (grads + rounded_size),
+               (_exp_avg_sq + rounded_size),
+               (_param_size - rounded_size));
+}
+
+int create_adagrad_optimizer(int optimizer_id,
+                             float alpha = 1e-2,
+                             float eps = 1e-8,
+                             float weight_decay = 0,
+                             bool should_log = false)
+{
+    auto opt = std::make_shared<Adagrad_Optimizer>(alpha, eps, weight_decay);
+
+    s_optimizers[optimizer_id] = opt;
+
+    if (should_log) {
+        std::string avx_type = "";
+#if defined(__AVX512__)
+        avx_type = "AVX512";
+#else
+#if defined(__AVX256__)
+        avx_type = "AVX2";
+#else
+        avx_type = "scalar";
+#endif
+#endif
+
+        printf("Adagrad Optimizer #%d is created with %s arithmetic capability.\n",
+               optimizer_id,
+               avx_type.c_str());
+        printf("Config: alpha=%f, weight_decay=%f\n", alpha, weight_decay);
+    }
+
+    return 0;
+}
+
+template <typename ds_params_precision_t, typename ds_state_precision_t>
+void Adagrad_Optimizer::Step_8(ds_params_precision_t* _params,
+                               ds_params_precision_t* grads,
+                               ds_state_precision_t* _exp_avg_sq,
+                               size_t _param_size)
+{
+    size_t rounded_size = 0;
+#if defined(__AVX512__) or defined(__AVX256__)
+    Step_AVX<8>(&rounded_size, _params, grads, _exp_avg_sq, _param_size);
+#endif
+    if (_param_size > rounded_size)
+        Step_4((_params + rounded_size),
+               (grads + rounded_size),
+               (_exp_avg_sq + rounded_size),
+               (_param_size - rounded_size));
+}
+
+template <typename ds_params_precision_t, typename ds_state_precision_t>
+void step_invoker(std::shared_ptr<Adagrad_Optimizer> opt,
+                  void* _params,
+                  void* grads,
+                  void* _exp_avg_sq,
+                  size_t _param_size)
+{
+    opt->Step_8((ds_params_precision_t*)(_params),
+                (ds_params_precision_t*)(grads),
+                (ds_state_precision_t*)(_exp_avg_sq),
+                _param_size);
+}
+
+std::map<std::tuple<c10::ScalarType, c10::ScalarType>,
+         std::function<void(std::shared_ptr<Adagrad_Optimizer>, void*, void*, void*, size_t)>>
+    invokers;
+
+// Fill map with template functions for each type
+template <class ds_params_precision_t, class ds_state_precision_t>
+void create_invoker()
+{
+    invokers[std::tuple(c10::CppTypeToScalarType<ds_params_precision_t>(),
+                        c10::CppTypeToScalarType<ds_state_precision_t>())] =
+        step_invoker<ds_params_precision_t, ds_state_precision_t>;
+}
+struct InvokerInitializer {
+    InvokerInitializer()
+    {
+        create_invoker<c10::Half, float>();
+        create_invoker<c10::Half, c10::Half>();
+        create_invoker<c10::BFloat16, float>();
+        create_invoker<c10::BFloat16, c10::BFloat16>();
+        create_invoker<float, float>();
+    }
+} _invoker_initializer;
+
+void invoke(std::shared_ptr<Adagrad_Optimizer> opt,
+            torch::Tensor& params,
+            torch::Tensor& grads,
+            torch::Tensor& exp_avg_sq,
+            size_t param_size)
+{
+    c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype());
+    c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg_sq.options().dtype());
+
+    auto it = invokers.find(std::tuple(params_type, state_type));
+    if (it == invokers.end()) {
+        throw std::runtime_error("Adagrad optimizer with param type "s +
+                                 c10::toString(params_type) + " and state type "s +
+                                 c10::toString(state_type) +
+                                 " is not supported on current hardware"s);
+    }
+
+    it->second(opt, params.data_ptr(), grads.data_ptr(), exp_avg_sq.data_ptr(), param_size);
+}
+
+int ds_adagrad_step(int optimizer_id,
+                    size_t step,
+                    float lr,
+                    float epsilon,
+                    float weight_decay,
+                    torch::Tensor& params,
+                    torch::Tensor& grads,
+                    torch::Tensor& exp_avg_sq)
+{
+    auto params_c = params.contiguous();
+    auto grads_c = grads.contiguous();
+    auto exp_avg_sq_c = exp_avg_sq.contiguous();
+
+    std::shared_ptr<Adagrad_Optimizer> opt =
+        std::static_pointer_cast<Adagrad_Optimizer>(s_optimizers[optimizer_id]);
+    opt->IncrementStep(step);
+    opt->update_state(lr, epsilon, weight_decay);
+
+    invoke(opt, params_c, grads_c, exp_avg_sq_c, params_c.numel());
+
+    return 0;
+}
+
+int destroy_adagrad_optimizer(int optimizer_id)
+{
+    s_optimizers.erase(optimizer_id);
+
+    return 0;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("adagrad_update", &ds_adagrad_step, "DeepSpeed CPU Adagrad update (C++)");
+    m.def("create_adagrad", &create_adagrad_optimizer, "DeepSpeed CPU Adagrad (C++)");
+    m.def("destroy_adagrad", &destroy_adagrad_optimizer, "DeepSpeed CPU Adagrad destroy (C++)");
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp

@@ -0,0 +1,13 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "cpu_adam.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("adam_update", &ds_adam_step, "DeepSpeed CPU Adam update (C++)");
+    m.def("create_adam", &create_adam_optimizer, "DeepSpeed CPU Adam (C++)");
+    m.def("destroy_adam", &destroy_adam_optimizer, "DeepSpeed CPU Adam destroy (C++)");
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp

@@ -0,0 +1,244 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <torch/extension.h>
+#include <cassert>
+#include <functional>
+#include <iostream>
+#include <map>
+#include <memory>
+#include <type_traits>
+#include <unordered_map>
+#include "cpu_adam.h"
+
+using namespace std::string_literals;
+static std::unordered_map<int, std::shared_ptr<void>> s_optimizers;
+
+// C++ interface
+
+template <typename ds_params_precision_t, typename ds_state_precision_t>
+void Adam_Optimizer::Step_1(ds_params_precision_t* _params,
+                            ds_params_precision_t* grads,
+                            ds_state_precision_t* _exp_avg,
+                            ds_state_precision_t* _exp_avg_sq,
+                            size_t _param_size)
+{
+    size_t rounded_size = 0;
+#if defined(__AVX512__) or defined(__AVX256__)
+    Step_AVX<1>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size);
+#endif
+    if (_param_size > rounded_size) {
+        float betta1_minus1 = 1 - _betta1;
+        float betta2_minus1 = 1 - _betta2;
+
+        float step_size = -1 * _alpha / _bias_correction1;
+        float w_decay = -1 * _alpha * _weight_decay;
+
+        for (size_t t = rounded_size; t < _param_size; t += TILE) {
+            size_t copy_size = TILE;
+            if ((t + TILE) > _param_size) copy_size = _param_size - t;
+            size_t offset = copy_size + t;
+#pragma omp parallel for
+            for (size_t k = t; k < offset; k++) {
+                float grad = (float)grads[k];
+                float param = (float)_params[k];
+                float momentum = _exp_avg[k];
+                float variance = _exp_avg_sq[k];
+                if (_weight_decay > 0 && !_adamw_mode) { grad = param * _weight_decay + grad; }
+                momentum = momentum * _betta1;
+                momentum = grad * betta1_minus1 + momentum;
+
+                variance = variance * _betta2;
+                grad = grad * grad;
+                variance = grad * betta2_minus1 + variance;
+
+                grad = sqrt(variance);
+                grad = grad * _bias_correction2 + _eps;
+                grad = momentum / grad;
+                if (_weight_decay > 0 && _adamw_mode) { param += w_decay * param; }
+                param = grad * step_size + param;
+                _params[k] = param;
+                _exp_avg[k] = momentum;
+                _exp_avg_sq[k] = variance;
+            }
+        }
+    }
+}
+
+template <typename ds_params_precision_t, typename ds_state_precision_t>
+void Adam_Optimizer::Step_4(ds_params_precision_t* _params,
+                            ds_params_precision_t* grads,
+                            ds_state_precision_t* _exp_avg,
+                            ds_state_precision_t* _exp_avg_sq,
+                            size_t _param_size)
+{
+    size_t rounded_size = 0;
+#if defined(__AVX512__) or defined(__AVX256__)
+    Step_AVX<4>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size);
+#endif
+    if (_param_size > rounded_size)
+        Step_1((_params + rounded_size),
+               (grads + rounded_size),
+               (_exp_avg + rounded_size),
+               (_exp_avg_sq + rounded_size),
+               (_param_size - rounded_size));
+}
+
+int create_adam_optimizer(int optimizer_id,
+                          float alpha,
+                          float betta1,
+                          float betta2,
+                          float eps,
+                          float weight_decay,
+                          bool adamw_mode,
+                          bool should_log)
+{
+    auto opt =
+        std::make_shared<Adam_Optimizer>(alpha, betta1, betta2, eps, weight_decay, adamw_mode);
+
+    s_optimizers[optimizer_id] = opt;
+
+    if (should_log) {
+        std::string avx_type = "";
+#if defined(__AVX512__)
+        avx_type = "AVX512";
+#else
+#if defined(__AVX256__)
+        avx_type = "AVX2";
+#else
+        avx_type = "scalar";
+#endif
+#endif
+
+        printf("Adam Optimizer #%d is created with %s arithmetic capability.\n",
+               optimizer_id,
+               avx_type.c_str());
+        printf("Config: alpha=%f, betas=(%f, %f), weight_decay=%f, adam_w=%d\n",
+               alpha,
+               betta1,
+               betta2,
+               weight_decay,
+               (int)adamw_mode);
+    }
+
+    return 0;
+}
+
+template <typename ds_params_precision_t, typename ds_state_precision_t>
+void Adam_Optimizer::Step_8(ds_params_precision_t* _params,
+                            ds_params_precision_t* grads,
+                            ds_state_precision_t* _exp_avg,
+                            ds_state_precision_t* _exp_avg_sq,
+                            size_t _param_size)
+{
+    size_t rounded_size = 0;
+#if defined(__AVX512__) or defined(__AVX256__)
+    Step_AVX<8>(&rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size);
+#endif
+    if (_param_size > rounded_size)
+        Step_4((_params + rounded_size),
+               (grads + rounded_size),
+               (_exp_avg + rounded_size),
+               (_exp_avg_sq + rounded_size),
+               (_param_size - rounded_size));
+}
+
+template <typename ds_params_precision_t, typename ds_state_precision_t>
+void step_invoker(std::shared_ptr<Adam_Optimizer> opt,
+                  void* _params,
+                  void* grads,
+                  void* _exp_avg,
+                  void* _exp_avg_sq,
+                  size_t _param_size)
+{
+    opt->Step_8((ds_params_precision_t*)(_params),
+                (ds_params_precision_t*)(grads),
+                (ds_state_precision_t*)(_exp_avg),
+                (ds_state_precision_t*)(_exp_avg_sq),
+                _param_size);
+}
+
+std::map<std::tuple<c10::ScalarType, c10::ScalarType>,
+         std::function<void(std::shared_ptr<Adam_Optimizer>, void*, void*, void*, void*, size_t)>>
+    invokers;
+
+// Fill map with template functions for each type
+template <class ds_params_precision_t, class ds_state_precision_t>
+void create_invoker()
+{
+    invokers[std::tuple(c10::CppTypeToScalarType<ds_params_precision_t>(),
+                        c10::CppTypeToScalarType<ds_state_precision_t>())] =
+        step_invoker<ds_params_precision_t, ds_state_precision_t>;
+}
+struct InvokerInitializer {
+    InvokerInitializer()
+    {
+        create_invoker<c10::Half, float>();
+        create_invoker<c10::Half, c10::Half>();
+        create_invoker<c10::BFloat16, float>();
+        create_invoker<c10::BFloat16, c10::BFloat16>();
+        create_invoker<float, float>();
+    }
+} _invoker_initializer;
+
+void invoke(std::shared_ptr<Adam_Optimizer> opt,
+            torch::Tensor& params,
+            torch::Tensor& grads,
+            torch::Tensor& exp_avg,
+            torch::Tensor& exp_avg_sq,
+            size_t param_size)
+{
+    c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype());
+    c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg.options().dtype());
+
+    auto it = invokers.find(std::tuple(params_type, state_type));
+    if (it == invokers.end()) {
+        throw std::runtime_error("Adam optimizer with param type "s + c10::toString(params_type) +
+                                 " and state type "s + c10::toString(state_type) +
+                                 " is not supported on current hardware"s);
+    }
+
+    it->second(opt,
+               params.data_ptr(),
+               grads.data_ptr(),
+               exp_avg.data_ptr(),
+               exp_avg_sq.data_ptr(),
+               param_size);
+}
+
+int ds_adam_step(int optimizer_id,
+                 size_t step,
+                 float lr,
+                 float beta1,
+                 float beta2,
+                 float epsilon,
+                 float weight_decay,
+                 bool bias_correction,
+                 torch::Tensor& params,
+                 torch::Tensor& grads,
+                 torch::Tensor& exp_avg,
+                 torch::Tensor& exp_avg_sq)
+{
+    auto params_c = params.contiguous();
+    auto grads_c = grads.contiguous();
+    auto exp_avg_c = exp_avg.contiguous();
+    auto exp_avg_sq_c = exp_avg_sq.contiguous();
+
+    std::shared_ptr<Adam_Optimizer> opt =
+        std::static_pointer_cast<Adam_Optimizer>(s_optimizers[optimizer_id]);
+    opt->IncrementStep(step, beta1, beta2);
+    opt->update_state(lr, epsilon, weight_decay, bias_correction);
+
+    invoke(opt, params_c, grads_c, exp_avg_c, exp_avg_sq_c, params_c.numel());
+
+    return 0;
+}
+
+int destroy_adam_optimizer(int optimizer_id)
+{
+    s_optimizers.erase(optimizer_id);
+
+    return 0;
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <torch/extension.h>
+
+void multi_tensor_adam_cuda(int chunk_size,
+                            at::Tensor noop_flag,
+                            std::vector<std::vector<at::Tensor>> tensor_lists,
+                            const float lr,
+                            const float beta1,
+                            const float beta2,
+                            const float epsilon,
+                            const int step,
+                            const int mode,
+                            const int bias_correction,
+                            const float weight_decay);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("multi_tensor_adam",
+          &multi_tensor_adam_cuda,
+          "Compute and apply gradient update to parameters for Adam optimizer");
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu

@@ -0,0 +1,203 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Copyright NVIDIA/apex
+This file is adapted from fused adam in NVIDIA/apex, commit a109f85
+*/
+
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "multi_tensor_apply.cuh"
+#include "type_shim.h"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+typedef enum : int {
+    ADAM_MODE_0 = 0,  // L2 regularization mode
+    ADAM_MODE_1 = 1   // Decoupled weight decay mode(AdamW)
+} adamMode_t;
+
+using MATH_T = float;
+
+template <typename T, typename index_t>
+struct AdamFunctor {
+    __device__ __forceinline__ void operator()(int chunk_size,
+                                               volatile int* noop_gmem,
+                                               TensorListMetadata<4>& tl,
+                                               const float beta1,
+                                               const float beta2,
+                                               const float beta1_correction,
+                                               const float beta2_correction,
+                                               const float epsilon,
+                                               const float lr,
+                                               adamMode_t mode,
+                                               const float decay)
+    {
+        // I'd like this kernel to propagate infs/nans.
+        // if(*noop_gmem == 1)
+        //   return;
+
+        index_t tensor_loc = tl.block_to_tensor[blockIdx.x];
+
+        // potentially use to pass in list of scalar
+        // int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+
+        index_t chunk_idx = tl.block_to_chunk[blockIdx.x];
+        index_t n = tl.sizes[tensor_loc];
+
+        T* g = (T*)tl.addresses[0][tensor_loc];
+        g += chunk_idx * chunk_size;
+
+        T* p = (T*)tl.addresses[1][tensor_loc];
+        p += chunk_idx * chunk_size;
+
+        T* m = (T*)tl.addresses[2][tensor_loc];
+        m += chunk_idx * chunk_size;
+
+        T* v = (T*)tl.addresses[3][tensor_loc];
+        v += chunk_idx * chunk_size;
+
+        n -= chunk_idx * chunk_size;
+
+        // see note in multi_tensor_scale_kernel.cu
+        for (index_t i_start = 0; i_start < n && i_start < chunk_size;
+             i_start += blockDim.x * ILP) {
+            MATH_T r_g[ILP];
+            MATH_T r_p[ILP];
+            MATH_T r_m[ILP];
+            MATH_T r_v[ILP];
+#pragma unroll
+            for (int ii = 0; ii < ILP; ii++) {
+                int i = i_start + threadIdx.x + ii * blockDim.x;
+                if (i < n && i < chunk_size) {
+                    r_g[ii] = g[i];
+                    r_p[ii] = p[i];
+                    r_m[ii] = m[i];
+                    r_v[ii] = v[i];
+                } else {
+                    r_g[ii] = MATH_T(0);
+                    r_p[ii] = MATH_T(0);
+                    r_m[ii] = MATH_T(0);
+                    r_v[ii] = MATH_T(0);
+                }
+            }
+#pragma unroll
+            for (int ii = 0; ii < ILP; ii++) {
+                if (mode == ADAM_MODE_0) {  // L2
+                    r_g[ii] = r_g[ii] + (decay * r_p[ii]);
+                    r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii];
+                    r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii];
+                    MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+                    MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+                    MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+                    MATH_T update = next_m_unbiased / denom;
+                    r_p[ii] = r_p[ii] - (lr * update);
+                } else {  // weight decay
+                    r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii];
+                    r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii];
+                    MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+                    MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+                    MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+                    MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]);
+                    r_p[ii] = r_p[ii] - (lr * update);
+                }
+            }
+#pragma unroll
+            for (int ii = 0; ii < ILP; ii++) {
+                int i = i_start + threadIdx.x + ii * blockDim.x;
+                if (i < n && i < chunk_size) {
+                    p[i] = r_p[ii];
+                    m[i] = r_m[ii];
+                    v[i] = r_v[ii];
+                }
+            }
+        }
+    }
+};
+
+void multi_tensor_adam_cuda(int chunk_size,
+                            at::Tensor noop_flag,
+                            std::vector<std::vector<at::Tensor>> tensor_lists,
+                            const float lr,
+                            const float beta1,
+                            const float beta2,
+                            const float epsilon,
+                            const int step,
+                            const int mode,
+                            const int bias_correction,
+                            const float weight_decay)
+{
+    using namespace at;
+
+    // Handle bias correction mode
+    float bias_correction1 = 1.0f, bias_correction2 = 1.0f;
+    if (bias_correction == 1) {
+        bias_correction1 = 1 - std::pow(beta1, step);
+        bias_correction2 = 1 - std::pow(beta2, step);
+    }
+
+    size_t max_size = 0;
+    bool requires_64bit_indexing = false;
+    for (auto it = tensor_lists.begin(); it != tensor_lists.end(); it++) {
+        for (auto it2 = it->begin(); it2 != it->end(); it2++) {
+            if (it2->numel() > max_size) {
+                max_size = it2->numel();
+                if (max_size >= INT_MAX) {
+                    requires_64bit_indexing = true;
+                    break;
+                }
+            }
+        }
+        if (requires_64bit_indexing) { break; }
+    }
+
+    // Assume single type across p,g,m1,m2 now
+    if (requires_64bit_indexing) {
+        DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(),
+                                       0,
+                                       "adam",
+                                       multi_tensor_apply<4>((int64_t)BLOCK_SIZE,
+                                                             (int64_t)chunk_size,
+                                                             noop_flag,
+                                                             tensor_lists,
+                                                             AdamFunctor<scalar_t_0, int64_t>(),
+                                                             beta1,
+                                                             beta2,
+                                                             bias_correction1,
+                                                             bias_correction2,
+                                                             epsilon,
+                                                             lr,
+                                                             (adamMode_t)mode,
+                                                             weight_decay);)
+    } else {
+        DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(),
+                                       0,
+                                       "adam",
+                                       multi_tensor_apply<4>(BLOCK_SIZE,
+                                                             chunk_size,
+                                                             noop_flag,
+                                                             tensor_lists,
+                                                             AdamFunctor<scalar_t_0, int32_t>(),
+                                                             beta1,
+                                                             beta2,
+                                                             bias_correction1,
+                                                             bias_correction2,
+                                                             epsilon,
+                                                             lr,
+                                                             (adamMode_t)mode,
+                                                             weight_decay);)
+    }
+
+    AT_CUDA_CHECK(cudaGetLastError());
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh

@@ -0,0 +1,132 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Copyright NVIDIA/apex
+This file is adapted from fused adam in NVIDIA/apex, commit a109f85
+*/
+
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "compat.h"
+
+#include <assert.h>
+
+// #include <iostream>
+
+// This header is the one-stop shop for all your multi-tensor apply needs.
+
+// TODO:  Kernel arg size limit may be <4KB for some other cards (ie Jetson)
+constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30};
+constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320};
+
+template <int n>
+struct TensorListMetadata {
+    void* addresses[n][depth_to_max_tensors[n - 1]];
+    int sizes[depth_to_max_tensors[n - 1]];
+    unsigned char block_to_tensor[depth_to_max_blocks[n - 1]];
+    int block_to_chunk[depth_to_max_blocks[n - 1]];  // I fear this needs to be a full int.
+    int start_tensor_this_launch;
+};
+
+template <typename T, typename U, typename... ArgTypes>
+__global__ void multi_tensor_apply_kernel(int64_t chunk_size,
+                                          volatile int* noop_flag,
+                                          T tl,
+                                          U callable,
+                                          ArgTypes... args)
+{
+    // Hand the chunk information to the user-supplied functor to process however it likes.
+    callable(chunk_size, noop_flag, tl, args...);
+}
+
+template <int depth, typename T, typename... ArgTypes>
+void multi_tensor_apply(int64_t block_size,
+                        int64_t chunk_size,
+                        const at::Tensor& noop_flag,
+                        const std::vector<std::vector<at::Tensor>>& tensor_lists,
+                        T callable,
+                        ArgTypes... args)
+{
+    TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth");
+    int len0 = tensor_lists[0].size();
+    TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0");
+    auto ref_device = tensor_lists[0][0].device();
+    TORCH_CHECK(ref_device.type() == at::kCUDA, "expected input to be on cuda");
+    for (int l = 0; l < tensor_lists.size(); l++)  // No range-based for because I need indices
+    {
+        TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists");
+        for (int t = 0; t < tensor_lists[l].size(); t++) {
+            // TODO:  Print which tensor fails.
+            bool contiguous_memory = tensor_lists[l][t].is_contiguous();
+#ifdef VERSION_GE_1_5
+            contiguous_memory = (contiguous_memory ||
+                                 tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast));
+#endif
+            TORCH_CHECK(contiguous_memory, "A tensor was not contiguous.");
+            TORCH_CHECK(tensor_lists[l][t].device() == ref_device,
+                        "A tensor was not on the same device as the first tensor");
+            TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch");
+        }
+    }
+
+    int ntensors = tensor_lists[0].size();
+
+    TensorListMetadata<depth> tl;
+
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0]));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    tl.start_tensor_this_launch = 0;
+    int loc_block_info = 0;
+    int loc_tensor_info = 0;
+    for (int t = 0; t < ntensors; t++) {
+        tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel();
+        for (int d = 0; d < depth; d++)
+            tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr();
+        loc_tensor_info++;
+
+        auto chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size;
+
+        for (auto chunk = 0; chunk < chunks_this_tensor; chunk++) {
+            // std::cout << chunks_this_tensor << std::endl;
+            tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1;
+            tl.block_to_chunk[loc_block_info] = chunk;
+            loc_block_info++;
+
+            bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth - 1] &&
+                                 chunk == chunks_this_tensor - 1);
+            bool blocks_full = (loc_block_info == depth_to_max_blocks[depth - 1]);
+            bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1);
+            if (tensors_full || blocks_full || last_chunk) {
+                // using accscalar_t = acc_type<scalar_t, true>;
+                multi_tensor_apply_kernel<<<loc_block_info, block_size, 0, stream>>>(
+                    chunk_size, noop_flag.DATA_PTR<int>(), tl, callable, args...);
+
+                AT_CUDA_CHECK(cudaGetLastError());
+
+                // Reset.  The control flow possibilities here make my brain hurt.
+                loc_block_info = 0;
+                if (chunk == chunks_this_tensor - 1) {
+                    // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 <<
+                    // std::endl;
+                    loc_tensor_info = 0;
+                    tl.start_tensor_this_launch = t + 1;
+                } else {
+                    // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 <<
+                    // std::endl;
+                    tl.sizes[0] = tl.sizes[loc_tensor_info - 1];
+                    for (int d = 0; d < depth; d++)
+                        tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1];
+                    loc_tensor_info = 1;
+                    tl.start_tensor_this_launch = t;
+                }
+            }
+        }
+    }
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp

@@ -0,0 +1,341 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <assert.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <fcntl.h>
+#include <libaio.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <time.h>
+#include <unistd.h>
+
+#include <algorithm>
+#include <chrono>
+#include <cmath>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <limits>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "deepspeed_aio_common.h"
+
+using namespace std;
+using namespace std::chrono;
+
+#define DEBUG_DS_AIO_PERF 0
+#define DEBUG_DS_AIO_SUBMIT_PERF 0
+
+static const std::string c_library_name = "deepspeed_aio";
+
+static void _report_aio_statistics(const char* tag,
+                                   const std::vector<std::chrono::duration<double>>& latencies)
+    __attribute__((unused));
+
+static void _report_aio_statistics(const char* tag,
+                                   const std::vector<std::chrono::duration<double>>& latencies)
+{
+    std::vector<double> lat_usec;
+    for (auto& lat : latencies) { lat_usec.push_back(lat.count() * 1e6); }
+    const auto min_lat = *(std::min_element(lat_usec.begin(), lat_usec.end()));
+    const auto max_lat = *(std::max_element(lat_usec.begin(), lat_usec.end()));
+    const auto avg_lat = std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size();
+
+    std::cout << c_library_name << ": latency statistics(usec) " << tag
+              << " min/max/avg = " << min_lat << " " << max_lat << " " << avg_lat << std::endl;
+}
+
+static void _get_aio_latencies(std::vector<std::chrono::duration<double>>& raw_latencies,
+                               struct deepspeed_aio_latency_t& summary_latencies)
+{
+    std::vector<double> lat_usec;
+    for (auto& lat : raw_latencies) { lat_usec.push_back(lat.count() * 1e6); }
+    summary_latencies._min_usec = *(std::min_element(lat_usec.begin(), lat_usec.end()));
+    summary_latencies._max_usec = *(std::max_element(lat_usec.begin(), lat_usec.end()));
+    summary_latencies._avg_usec =
+        std::accumulate(lat_usec.begin(), lat_usec.end(), 0) / lat_usec.size();
+}
+
+static void _do_io_submit_singles(const int64_t n_iocbs,
+                                  const int64_t iocb_index,
+                                  std::unique_ptr<aio_context>& aio_ctxt,
+                                  std::vector<std::chrono::duration<double>>& submit_times)
+{
+    for (auto i = 0; i < n_iocbs; ++i) {
+        const auto st = std::chrono::high_resolution_clock::now();
+        const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, 1, aio_ctxt->_iocbs.data() + i);
+        submit_times.push_back(std::chrono::high_resolution_clock::now() - st);
+#if DEBUG_DS_AIO_SUBMIT_PERF
+        printf("submit(usec) %f io_index=%lld buf=%p len=%lu off=%llu \n",
+               submit_times.back().count() * 1e6,
+               iocb_index,
+               aio_ctxt->_iocbs[i]->u.c.buf,
+               aio_ctxt->_iocbs[i]->u.c.nbytes,
+               aio_ctxt->_iocbs[i]->u.c.offset);
+#endif
+        assert(submit_ret > 0);
+    }
+}
+
+static void _do_io_submit_block(const int64_t n_iocbs,
+                                const int64_t iocb_index,
+                                std::unique_ptr<aio_context>& aio_ctxt,
+                                std::vector<std::chrono::duration<double>>& submit_times)
+{
+    const auto st = std::chrono::high_resolution_clock::now();
+    const auto submit_ret = io_submit(aio_ctxt->_io_ctxt, n_iocbs, aio_ctxt->_iocbs.data());
+    submit_times.push_back(std::chrono::high_resolution_clock::now() - st);
+#if DEBUG_DS_AIO_SUBMIT_PERF
+    printf("submit(usec) %f io_index=%lld nr=%lld buf=%p len=%lu off=%llu \n",
+           submit_times.back().count() * 1e6,
+           iocb_index,
+           n_iocbs,
+           aio_ctxt->_iocbs[0]->u.c.buf,
+           aio_ctxt->_iocbs[0]->u.c.nbytes,
+           aio_ctxt->_iocbs[0]->u.c.offset);
+#endif
+    assert(submit_ret > 0);
+}
+
+static int _do_io_complete(const int64_t min_completes,
+                           const int64_t max_completes,
+                           std::unique_ptr<aio_context>& aio_ctxt,
+                           std::vector<std::chrono::duration<double>>& reap_times)
+{
+    const auto start_time = std::chrono::high_resolution_clock::now();
+    int64_t n_completes = io_pgetevents(aio_ctxt->_io_ctxt,
+                                        min_completes,
+                                        max_completes,
+                                        aio_ctxt->_io_events.data(),
+                                        nullptr,
+                                        nullptr);
+    reap_times.push_back(std::chrono::high_resolution_clock::now() - start_time);
+    assert(n_completes >= min_completes);
+    return n_completes;
+}
+
+void do_aio_operation_sequential(const bool read_op,
+                                 std::unique_ptr<aio_context>& aio_ctxt,
+                                 std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
+                                 deepspeed_aio_config_t* config,
+                                 deepspeed_aio_perf_t* perf)
+{
+    struct io_prep_context prep_ctxt(read_op, xfer_ctxt, aio_ctxt->_block_size, &aio_ctxt->_iocbs);
+
+    const auto num_io_blocks = static_cast<int64_t>(
+        ceil(static_cast<double>(xfer_ctxt->_num_bytes) / aio_ctxt->_block_size));
+#if DEBUG_DS_AIO_PERF
+    const auto io_op_name = std::string(read_op ? "read" : "write");
+    std::cout << c_library_name << ": start " << io_op_name << " " << xfer_ctxt->_num_bytes
+              << " bytes with " << num_io_blocks << " io blocks" << std::endl;
+#endif
+
+    std::vector<std::chrono::duration<double>> submit_times;
+    std::vector<std::chrono::duration<double>> reap_times;
+    const auto max_queue_bytes =
+        static_cast<int64_t>(aio_ctxt->_queue_depth * aio_ctxt->_block_size);
+
+    auto start = std::chrono::high_resolution_clock::now();
+    for (int64_t iocb_index = 0; iocb_index < num_io_blocks; iocb_index += aio_ctxt->_queue_depth) {
+        const auto start_offset = iocb_index * aio_ctxt->_block_size;
+        const auto start_buffer = (char*)xfer_ctxt->_mem_buffer + start_offset;
+        const auto n_iocbs =
+            min(static_cast<int64_t>(aio_ctxt->_queue_depth), (num_io_blocks - iocb_index));
+        const auto num_bytes = min(max_queue_bytes, (xfer_ctxt->_num_bytes - start_offset));
+        prep_ctxt.prep_iocbs(n_iocbs, num_bytes, start_buffer, start_offset);
+
+        if (config->_single_submit) {
+            _do_io_submit_singles(n_iocbs, iocb_index, aio_ctxt, submit_times);
+        } else {
+            _do_io_submit_block(n_iocbs, iocb_index, aio_ctxt, submit_times);
+        }
+
+        _do_io_complete(n_iocbs, n_iocbs, aio_ctxt, reap_times);
+    }
+    const std::chrono::duration<double> elapsed = std::chrono::high_resolution_clock::now() - start;
+
+    if (perf) {
+        _get_aio_latencies(submit_times, perf->_submit);
+        _get_aio_latencies(reap_times, perf->_complete);
+        perf->_e2e_usec = elapsed.count() * 1e6;
+        perf->_e2e_rate_GB = (xfer_ctxt->_num_bytes / elapsed.count() / 1e9);
+    }
+
+#if DEBUG_DS_AIO_PERF
+    _report_aio_statistics("submit", submit_times);
+    _report_aio_statistics("complete", reap_times);
+#endif
+
+#if DEBUG_DS_AIO_PERF
+    std::cout << c_library_name << ": runtime(usec) " << elapsed.count() * 1e6
+              << " rate(GB/sec) = " << (xfer_ctxt->_num_bytes / elapsed.count() / 1e9) << std::endl;
+#endif
+
+#if DEBUG_DS_AIO_PERF
+    std::cout << c_library_name << ": finish " << io_op_name << " " << xfer_ctxt->_num_bytes
+              << " bytes " << std::endl;
+#endif
+}
+
+void do_aio_operation_overlap(const bool read_op,
+                              std::unique_ptr<aio_context>& aio_ctxt,
+                              std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
+                              deepspeed_aio_config_t* config,
+                              deepspeed_aio_perf_t* perf)
+{
+    struct io_prep_generator io_gen(read_op, xfer_ctxt, aio_ctxt->_block_size);
+
+#if DEBUG_DS_AIO_PERF
+    const auto io_op_name = std::string(read_op ? "read" : "write");
+    std::cout << c_library_name << ": start " << io_op_name << " " << xfer_ctxt->_num_bytes
+              << " bytes with " << io_gen._num_io_blocks << " io blocks" << std::endl;
+#endif
+
+    std::vector<std::chrono::duration<double>> submit_times;
+    std::vector<std::chrono::duration<double>> reap_times;
+
+    auto request_iocbs = aio_ctxt->_queue_depth;
+    auto n_pending_iocbs = 0;
+    const auto min_completes = 1;
+    auto start = std::chrono::high_resolution_clock::now();
+    while (true) {
+        const auto n_iocbs = io_gen.prep_iocbs(request_iocbs - n_pending_iocbs, &aio_ctxt->_iocbs);
+        if (n_iocbs > 0) {
+            if (config->_single_submit) {
+                _do_io_submit_singles(
+                    n_iocbs, (io_gen._next_iocb_index - n_iocbs), aio_ctxt, submit_times);
+            } else {
+                _do_io_submit_block(
+                    n_iocbs, (io_gen._next_iocb_index - n_iocbs), aio_ctxt, submit_times);
+            }
+        }
+
+        n_pending_iocbs += n_iocbs;
+        assert(n_pending_iocbs <= aio_ctxt->_queue_depth);
+
+        if (n_pending_iocbs == 0) { break; }
+
+        const auto n_complete =
+            _do_io_complete(min_completes, n_pending_iocbs, aio_ctxt, reap_times);
+        n_pending_iocbs -= n_complete;
+    }
+
+    const std::chrono::duration<double> elapsed = std::chrono::high_resolution_clock::now() - start;
+
+    if (perf) {
+        _get_aio_latencies(submit_times, perf->_submit);
+        _get_aio_latencies(reap_times, perf->_complete);
+        perf->_e2e_usec = elapsed.count() * 1e6;
+        perf->_e2e_rate_GB = (xfer_ctxt->_num_bytes / elapsed.count() / 1e9);
+    }
+
+#if DEBUG_DS_AIO_PERF
+    _report_aio_statistics("submit", submit_times);
+    _report_aio_statistics("complete", reap_times);
+#endif
+
+#if DEBUG_DS_AIO_PERF
+    std::cout << c_library_name << ": runtime(usec) " << elapsed.count() * 1e6
+              << " rate(GB/sec) = " << (xfer_ctxt->_num_bytes / elapsed.count() / 1e9) << std::endl;
+#endif
+
+#if DEBUG_DS_AIO_PERF
+    std::cout << c_library_name << ": finish " << io_op_name << " " << xfer_ctxt->_num_bytes
+              << " bytes " << std::endl;
+#endif
+}
+
+void report_file_error(const char* filename, const std::string file_op, const int error_code)
+{
+    std::string err_msg = file_op + std::string(" failed on ") + std::string(filename) +
+                          " error = " + std::to_string(error_code);
+    std::cerr << c_library_name << ":  " << err_msg << std::endl;
+}
+
+int open_file(const char* filename, const bool read_op)
+{
+    const int flags = read_op ? (O_RDONLY | O_DIRECT) : (O_WRONLY | O_CREAT | O_DIRECT);
+#if defined(__ENABLE_CANN__)
+    int* flags_ptr = (int*)&flags;
+    *flags_ptr = read_op ? (O_RDONLY) : (O_WRONLY | O_CREAT);
+#endif
+    const int mode = 0600;
+    const auto fd = open(filename, flags, mode);
+    if (fd == -1) {
+        const auto error_code = errno;
+        const auto error_msg = read_op ? " open for read " : " open for write ";
+        report_file_error(filename, error_msg, error_code);
+        return -1;
+    }
+    return fd;
+}
+
+int regular_read(const char* filename, std::vector<char>& buffer)
+{
+    const auto fd = open(filename, O_RDONLY, 0600);
+    assert(fd != -1);
+    struct stat fs;
+    const auto result = fstat(fd, &fs);
+    assert(result != -1);
+    int64_t num_bytes = fs.st_size;
+    buffer.resize(num_bytes);
+    int64_t read_bytes = 0;
+    auto r = 0;
+    do {
+        const auto buffer_ptr = buffer.data() + read_bytes;
+        const auto bytes_to_read = num_bytes - read_bytes;
+        r = read(fd, buffer_ptr, bytes_to_read);
+        read_bytes += r;
+    } while (r > 0);
+
+    if (read_bytes != num_bytes) {
+        std::cerr << "read error " << " read_bytes (read) = " << read_bytes
+                  << " num_bytes (fstat) = " << num_bytes << std::endl;
+    }
+    assert(read_bytes == num_bytes);
+    close(fd);
+    return 0;
+}
+
+static bool _validate_buffer(const char* filename, void* aio_buffer, const int64_t num_bytes)
+{
+    std::vector<char> regular_buffer;
+    const auto reg_ret = regular_read(filename, regular_buffer);
+    assert(0 == reg_ret);
+    std::cout << "regular read of " << filename << " returned " << regular_buffer.size() << " bytes"
+              << std::endl;
+
+    if (static_cast<int64_t>(regular_buffer.size()) != num_bytes) { return false; }
+
+    return (0 == memcmp(aio_buffer, regular_buffer.data(), regular_buffer.size()));
+}
+
+bool validate_aio_operation(const bool read_op,
+                            const char* filename,
+                            void* aio_buffer,
+                            const int64_t num_bytes)
+{
+    const auto msg_suffix = std::string("deepspeed_aio_") +
+                            std::string(read_op ? "read()" : "write()") +
+                            std::string("using read()");
+
+    if (false == _validate_buffer(filename, aio_buffer, num_bytes)) {
+        std::cout << "Fail: correctness of " << msg_suffix << std::endl;
+        return false;
+    }
+
+    std::cout << "Pass: correctness of  " << msg_suffix << std::endl;
+    return true;
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <deepspeed_aio_utils.h>
+#include <stdlib.h>
+#include <memory>
+#include <string>
+
+using namespace std;
+
+void do_aio_operation_sequential(const bool read_op,
+                                 std::unique_ptr<aio_context>& aio_ctxt,
+                                 std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
+                                 deepspeed_aio_config_t* config,
+                                 deepspeed_aio_perf_t* perf);
+
+void do_aio_operation_overlap(const bool read_op,
+                              std::unique_ptr<aio_context>& aio_ctxt,
+                              std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
+                              deepspeed_aio_config_t* config,
+                              deepspeed_aio_perf_t* perf);
+
+int open_file(const char* filename, const bool read_op);
+
+void report_file_error(const char* filename, const std::string file_op, const int error_code);
+
+int regular_read(const char* filename, std::vector<char>& buffer);
+
+bool validate_aio_operation(const bool read_op,
+                            const char* filename,
+                            void* aio_buffer,
+                            const int64_t num_bytes);

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp

@@ -0,0 +1,76 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <cmath>
+
+#include "deepspeed_aio_utils.h"
+
+using namespace std;
+
+const int c_block_size = 128 * 1024;
+const int c_io_queue_depth = 8;
+
+deepspeed_aio_config_t::deepspeed_aio_config_t()
+    : _block_size(c_block_size),
+      _queue_depth(c_io_queue_depth),
+      _single_submit(false),
+      _overlap_events(false),
+      _lock_memory(false)
+{
+}
+
+deepspeed_aio_config_t::deepspeed_aio_config_t(const int block_size,
+                                               const int queue_depth,
+                                               const bool single_submit,
+                                               const bool overlap_events,
+                                               const bool lock_memory)
+    : _block_size(block_size),
+      _queue_depth(queue_depth),
+      _single_submit(single_submit),
+      _overlap_events(overlap_events),
+      _lock_memory(lock_memory)
+{
+}
+
+void deepspeed_aio_latency_t::dump(const std::string tag)
+{
+    std::cout << tag << _min_usec << " " << _max_usec << " " << _avg_usec << " " << std::endl;
+}
+
+void deepspeed_aio_latency_t::accumulate(const struct deepspeed_aio_latency_t& other)
+{
+    _min_usec += other._min_usec;
+    _max_usec += other._max_usec;
+    _avg_usec += other._avg_usec;
+}
+
+void deepspeed_aio_latency_t::scale(const float scaler)
+{
+    _min_usec *= scaler;
+    _max_usec *= scaler;
+    _avg_usec *= scaler;
+}
+
+aio_context::aio_context(const int block_size, const int queue_depth)
+{
+    _block_size = block_size;
+    _queue_depth = queue_depth;
+    for (auto i = 0; i < queue_depth; ++i) {
+        _iocbs.push_back((struct iocb*)calloc(1, sizeof(struct iocb)));
+    }
+    _io_events.resize(queue_depth);
+    io_queue_init(queue_depth, &_io_ctxt);
+}
+
+aio_context::~aio_context()
+{
+    for (auto& iocb : _iocbs) { free(iocb); }
+    _io_events.resize(0);
+    io_queue_release(_io_ctxt);
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h

@@ -0,0 +1,59 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <libaio.h>
+#include <stdlib.h>
+
+#include <string>
+#include <vector>
+
+using namespace std;
+
+struct deepspeed_aio_latency_t {
+    double _min_usec;
+    double _max_usec;
+    double _avg_usec;
+
+    void dump(const std::string tag);
+    void accumulate(const deepspeed_aio_latency_t&);
+    void scale(const float value);
+};
+
+struct deepspeed_aio_perf_t {
+    deepspeed_aio_latency_t _submit;
+    deepspeed_aio_latency_t _complete;
+    double _e2e_usec;
+    double _e2e_rate_GB;
+};
+
+struct deepspeed_aio_config_t {
+    const int _block_size;
+    const int _queue_depth;
+    const bool _single_submit;
+    const bool _overlap_events;
+    const bool _lock_memory;
+
+    deepspeed_aio_config_t();
+    deepspeed_aio_config_t(const int block_size,
+                           const int queue_depth,
+                           const bool single_submit,
+                           const bool overlap_events,
+                           const bool lock_memory);
+};
+
+struct aio_context {
+    io_context_t _io_ctxt;
+    std::vector<struct io_event> _io_events;
+    std::vector<struct iocb*> _iocbs;
+    int _block_size;
+    int _queue_depth;
+
+    aio_context(const int block_size, const int queue_depth);
+    ~aio_context();
+};

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp

@@ -0,0 +1,132 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <cmath>
+#include <iostream>
+
+#include "deepspeed_aio_utils.h"
+
+using namespace std;
+
+const int c_block_size = 128 * 1024;
+const int c_io_queue_depth = 8;
+
+io_xfer_ctxt::io_xfer_ctxt(const int fd,
+                           const int64_t file_offset,
+                           const int64_t buffer_offset,
+                           const int64_t num_bytes,
+                           const void* buffer)
+    : _fd(fd),
+      _file_base_offset(file_offset),
+      _buffer_base_offset(buffer_offset),
+      _mem_buffer(buffer),
+      _num_bytes(num_bytes)
+{
+}
+
+io_prep_context::io_prep_context(const bool read_op,
+                                 const std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
+                                 const size_t block_size,
+                                 const std::vector<struct iocb*>* iocbs)
+    : _read_op(read_op), _xfer_ctxt(xfer_ctxt), _block_size(block_size), _iocbs(iocbs)
+{
+}
+
+void io_prep_context::prep_iocbs(const int n_iocbs,
+                                 const size_t num_bytes,
+                                 const void* start_buffer,
+                                 const int64_t start_offset)
+{
+    assert(static_cast<size_t>(n_iocbs) <= _iocbs->size());
+    for (auto i = 0; i < n_iocbs; ++i) {
+        const auto shift = i * _block_size;
+        const auto xfer_buffer = (char*)start_buffer + _xfer_ctxt->_buffer_base_offset + shift;
+        const auto xfer_offset = _xfer_ctxt->_file_base_offset + start_offset + shift;
+        auto byte_count = _block_size;
+
+        if ((shift + _block_size) > num_bytes) { byte_count = num_bytes - shift; }
+
+        if (_read_op) {
+            io_prep_pread(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset);
+        } else {
+            io_prep_pwrite(_iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, byte_count, xfer_offset);
+        }
+    }
+}
+
+io_prep_generator::io_prep_generator(const bool read_op,
+                                     const std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
+                                     const size_t block_size)
+    : _read_op(read_op),
+      _xfer_ctxt(xfer_ctxt),
+      _block_size(block_size),
+      _remaining_bytes(xfer_ctxt->_num_bytes),
+      _next_iocb_index(0)
+{
+    _num_io_blocks =
+        static_cast<int64_t>(ceil(static_cast<double>(xfer_ctxt->_num_bytes) / block_size));
+    _remaining_io_blocks = _num_io_blocks;
+}
+
+int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector<struct iocb*>* iocbs)
+{
+    if ((_remaining_bytes) == 0 || (_remaining_io_blocks == 0)) {
+        assert(static_cast<int64_t>(_remaining_bytes) == _remaining_io_blocks);
+        return 0;
+    }
+
+    assert(static_cast<size_t>(n_iocbs) <= iocbs->size());
+
+    auto actual_n_iocbs = min(static_cast<int64_t>(n_iocbs), _remaining_io_blocks);
+    for (auto i = 0; i < actual_n_iocbs; ++i, ++_next_iocb_index) {
+        const auto xfer_buffer = (char*)_xfer_ctxt->_mem_buffer + _xfer_ctxt->_buffer_base_offset +
+                                 (_next_iocb_index * _block_size);
+        const auto xfer_offset = _xfer_ctxt->_file_base_offset + (_next_iocb_index * _block_size);
+        const auto num_bytes = min(static_cast<int64_t>(_block_size), _remaining_bytes);
+        if (_read_op) {
+            io_prep_pread(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset);
+        } else {
+            io_prep_pwrite(iocbs->at(i), _xfer_ctxt->_fd, xfer_buffer, num_bytes, xfer_offset);
+        }
+        _remaining_bytes -= num_bytes;
+    }
+    _remaining_io_blocks -= actual_n_iocbs;
+
+    return actual_n_iocbs;
+}
+
+int get_file_size(const char* filename, int64_t& size)
+{
+    struct stat st;
+    if (stat(filename, &st) == -1) { return -1; }
+    size = st.st_size;
+    return 0;
+}
+
+void* ds_page_aligned_alloc(const int64_t size, const bool lock)
+{
+    void* ptr;
+    int retval;
+
+    retval = posix_memalign(&ptr, (size_t)sysconf(_SC_PAGESIZE), size);
+    if (retval) { return nullptr; }
+
+    if (lock == false) { return ptr; }
+
+    auto mlock_ret = mlock(ptr, size);
+    if (mlock_ret != 0) {
+        auto mlock_error = errno;
+        std::cerr << "mlock failed to allocate " << size << " bytes with error no " << mlock_error
+                  << " msg " << strerror(mlock_error) << std::endl;
+        free(ptr);
+        return nullptr;
+    }
+
+    return ptr;
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h

@@ -0,0 +1,81 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#pragma once
+
+#include <assert.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <fcntl.h>
+#include <libaio.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <deepspeed_aio_types.h>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <vector>
+
+struct io_xfer_ctxt {
+    const int _fd;
+    const int64_t _file_base_offset;
+    const int64_t _buffer_base_offset;
+    const void* _mem_buffer;
+    const int64_t _num_bytes;
+
+    io_xfer_ctxt(const int fd,
+                 const int64_t file_offset,
+                 const int64_t buffer_offset,
+                 const int64_t num_bytes,
+                 const void* buffer);
+};
+
+struct io_prep_context {
+    const bool _read_op;
+    const std::unique_ptr<io_xfer_ctxt>& _xfer_ctxt;
+    const size_t _block_size;
+    const std::vector<struct iocb*>* _iocbs;
+
+    io_prep_context(const bool read_op,
+                    const std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
+                    const size_t block_size,
+                    const std::vector<struct iocb*>* iocbs);
+
+    void prep_iocbs(const int n_iocbs,
+                    const size_t num_bytes,
+                    const void* start_buffer,
+                    const int64_t start_offset);
+};
+
+struct io_prep_generator {
+    const bool _read_op;
+    const std::unique_ptr<io_xfer_ctxt>& _xfer_ctxt;
+    const size_t _block_size;
+
+    int64_t _remaining_bytes;
+    int64_t _num_io_blocks;
+    int64_t _remaining_io_blocks;
+    int64_t _next_iocb_index;
+
+    io_prep_generator(const bool read_op,
+                      const std::unique_ptr<io_xfer_ctxt>& xfer_ctxt,
+                      const size_t block_size);
+
+    int prep_iocbs(const int n_iocbs, std::vector<struct iocb*>* iocbs);
+};
+
+void* ds_page_aligned_alloc(const int64_t size, const bool lock = false);
+
+int get_file_size(const char* filename, int64_t& size);

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "deepspeed_aio_op_desc.h"
+
+using namespace std;
+
+io_op_desc_t::io_op_desc_t(const bool read_op,
+                           const torch::Tensor& buffer,
+                           const int fd,
+                           const char* filename,
+                           const int64_t file_num_bytes,
+                           const int intra_op_parallelism,
+                           const bool validate,
+                           const int64_t file_offset)
+    : _read_op(read_op),
+      _buffer(buffer),
+      _fd(fd),
+      _filename(filename),
+      _file_num_bytes(file_num_bytes),
+      _file_offset(file_offset),
+      _intra_op_parallelism(intra_op_parallelism),
+      _num_bytes_per_thread(static_cast<int64_t>(buffer.nbytes()) / intra_op_parallelism),
+      _validate(validate)
+{
+}
+
+char* io_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); }
+
+void io_op_desc_t::finish() {}
+
+void io_op_desc_t::validate() {}
+
+void io_op_desc_t::run(const int tid,
+                       std::unique_ptr<aio_context>& aio_ctxt,
+                       deepspeed_aio_config_t* aio_config)
+{
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#ifndef _IO_OP_DESC_T_
+#define _IO_OP_DESC_T_
+#include <memory>
+#include <queue>
+#include "deepspeed_py_aio.h"
+
+struct io_op_desc_t {
+    const bool _read_op;
+    torch::Tensor _buffer;
+    int _fd;
+    const std::string _filename;
+    const int64_t _file_num_bytes;
+    const int _intra_op_parallelism;
+    const int64_t _num_bytes_per_thread;
+    torch::Tensor _contiguous_buffer;
+    const bool _validate;
+    const int64_t _file_offset;
+
+    io_op_desc_t(const bool read_op,
+                 const torch::Tensor& buffer,
+                 const int fd,
+                 const char* filename,
+                 const int64_t file_num_bytes,
+                 const int intra_op_parallelism,
+                 const bool validate,
+                 const int64_t file_offset);
+
+    virtual void run(const int tid,
+                     std::unique_ptr<aio_context>& aio_ctxt,
+                     deepspeed_aio_config_t* aio_config);
+
+    virtual char* data_ptr() const;
+
+    virtual void validate();
+
+    virtual void finish();
+};
+#endif  // _IO_OP_DESC_T_

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include "deepspeed_aio_thread.h"
+
+using namespace std;
+
+deepspeed_aio_thread_t::deepspeed_aio_thread_t(const int tid, deepspeed_aio_config_t& aio_config)
+    : _tid(tid),
+      _aio_config(aio_config),
+      _aio_ctxt(new aio_context(aio_config._block_size, aio_config._queue_depth)),
+      _time_to_exit(false)
+{
+}
+
+deepspeed_aio_thread_t::~deepspeed_aio_thread_t() {}
+
+void deepspeed_aio_thread_t::run()
+{
+    while (true) {
+        std::shared_ptr<struct io_op_desc_t> next_io_op = nullptr;
+
+        {
+            std::unique_lock<std::mutex> lock(_work_sync._mutex);
+            _work_sync._cond_var.wait(lock,
+                                      [this] { return (!_work_queue.empty() || _time_to_exit); });
+            if (!_work_queue.empty()) {
+                next_io_op = _work_queue.front();
+                _work_queue.pop();
+            }
+        }
+
+        if (next_io_op) {
+            next_io_op->run(_tid, _aio_ctxt, &_aio_config);
+
+            {
+                std::lock_guard<std::mutex> lock(_complete_sync._mutex);
+                _complete_queue.push(next_io_op);
+            }
+            _complete_sync._cond_var.notify_one();
+        }
+
+        if (_time_to_exit) { break; }
+    }
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <condition_variable>
+#include <memory>
+#include <queue>
+#include "deepspeed_cpu_op.h"
+
+struct thread_sync_t {
+    std::mutex _mutex;
+    std::condition_variable _cond_var;
+};
+
+struct deepspeed_aio_thread_t {
+    const int _tid;
+    deepspeed_aio_config_t& _aio_config;
+
+    std::unique_ptr<struct aio_context> _aio_ctxt;
+    std::queue<std::shared_ptr<struct io_op_desc_t>> _work_queue;
+    std::queue<std::shared_ptr<struct io_op_desc_t>> _complete_queue;
+
+    bool _time_to_exit;
+
+    struct thread_sync_t _work_sync;
+    struct thread_sync_t _complete_sync;
+
+    deepspeed_aio_thread_t(const int tid, deepspeed_aio_config_t& aio_config);
+
+    ~deepspeed_aio_thread_t();
+
+    void run();
+};

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp

@@ -0,0 +1,109 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "deepspeed_cpu_op.h"
+#include "deepspeed_pin_tensor.h"
+
+using namespace std;
+
+cpu_op_desc_t::cpu_op_desc_t(
+    const bool read_op,
+    const torch::Tensor& buffer,
+    const std::unique_ptr<struct deepspeed_pin_tensor_t>& pinned_tensor_mgr,
+    const int fd,
+    const char* filename,
+    const int64_t file_num_bytes,
+    const int intra_op_parallelism,
+    const bool validate,
+    const int64_t file_offset)
+    : io_op_desc_t(read_op,
+                   buffer,
+                   fd,
+                   filename,
+                   file_num_bytes,
+                   intra_op_parallelism,
+                   validate,
+                   file_offset),
+      _cpu_buffer(buffer),
+      _pinned_tensor_mgr(pinned_tensor_mgr),
+      _is_managed_bounce_buffer(false)
+{
+    // Need to use CPU bounce buffer if buffer is not a page-locked DRAM memory.
+    _use_bounce_buffer =
+        !(_buffer.is_cpu() && (_buffer.is_pinned() || _pinned_tensor_mgr->is_managed(_buffer)));
+    if (_use_bounce_buffer) {
+        _alloc_bounce_buffer();
+        if (!_read_op) { _cpu_buffer.copy_(_buffer); }
+    }
+    _contiguous_buffer = _cpu_buffer.contiguous();
+}
+
+char* cpu_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); }
+
+void cpu_op_desc_t::finish()
+{
+    if (_use_bounce_buffer) {
+        if (_read_op) {
+            if (_buffer.is_cuda()) {
+                _buffer.copy_(_cpu_buffer.to(torch::Device(torch::kCUDA, _buffer.get_device()),
+                                             /*non_blocking=*/true));
+            }
+            if (_buffer.is_xpu()) { _buffer.copy_(_cpu_buffer.to(torch::kXPU)); }
+            if (_buffer.is_cpu()) { _buffer.copy_(_cpu_buffer); }
+#if defined(__ENABLE_CANN__)
+            if (torch_npu::utils::is_npu(_buffer)) {
+                auto device = at::Device("npu:0");
+                _buffer.copy_(_cpu_buffer.to(device));
+            }
+#endif
+        }
+
+        _free_bounce_buffer();
+    }
+}
+
+void cpu_op_desc_t::validate()
+{
+    validate_aio_operation(_read_op, _filename.c_str(), data_ptr(), _file_num_bytes);
+}
+
+void cpu_op_desc_t::run(const int tid,
+                        std::unique_ptr<aio_context>& aio_ctxt,
+                        deepspeed_aio_config_t* aio_config)
+{
+    assert(tid < _intra_op_parallelism);
+    const auto buffer_base_offset = _num_bytes_per_thread * tid;
+    const auto file_base_offset = _file_offset + (_num_bytes_per_thread * tid);
+
+    std::unique_ptr<io_xfer_ctxt> xfer_ctxt(new io_xfer_ctxt(
+        _fd, file_base_offset, buffer_base_offset, _num_bytes_per_thread, data_ptr()));
+
+    if (aio_config->_overlap_events) {
+        do_aio_operation_overlap(_read_op, aio_ctxt, xfer_ctxt, aio_config, nullptr);
+    } else {
+        do_aio_operation_sequential(_read_op, aio_ctxt, xfer_ctxt, aio_config, nullptr);
+    }
+}
+
+void cpu_op_desc_t::_alloc_bounce_buffer()
+{
+    auto options = torch::TensorOptions()
+                       .dtype(_buffer.dtype())
+                       .layout(_buffer.layout())
+                       .device(torch::kCPU)
+                       .requires_grad(false);
+
+#if defined(__CUDA_ARCH__)
+    _cpu_buffer = torch::empty(_buffer.numel(), options).pin_memory();
+#else
+    _is_managed_bounce_buffer = true;
+    _cpu_buffer = _pinned_tensor_mgr->alloc(_buffer.numel(), options);
+#endif
+}
+
+void cpu_op_desc_t::_free_bounce_buffer()
+{
+    if (_is_managed_bounce_buffer) { _pinned_tensor_mgr->free(_cpu_buffer); }
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <memory>
+#include <queue>
+#include "deepspeed_aio_op_desc.h"
+
+struct cpu_op_desc_t : io_op_desc_t {
+    torch::Tensor _cpu_buffer;
+    bool _use_bounce_buffer;
+    bool _is_managed_bounce_buffer;
+    const std::unique_ptr<struct deepspeed_pin_tensor_t>& _pinned_tensor_mgr;
+
+    cpu_op_desc_t(const bool read_op,
+                  const torch::Tensor& buffer,
+                  const std::unique_ptr<struct deepspeed_pin_tensor_t>& pinned_tensor_mgr,
+                  const int fd,
+                  const char* filename,
+                  const int64_t file_num_bytes,
+                  const int intra_op_parallelism,
+                  const bool validate,
+                  const int64_t file_offset);
+
+    void run(const int tid,
+             std::unique_ptr<aio_context>& aio_ctxt,
+             deepspeed_aio_config_t* aio_config);
+
+    char* data_ptr() const;
+
+    void validate();
+
+    void finish();
+
+    void _alloc_bounce_buffer();
+    void _free_bounce_buffer();
+};

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp

@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for managing CPU tensors occupying page-locked memory.
+*/
+
+#include "deepspeed_pin_tensor.h"
+
+using namespace std;
+
+deepspeed_pin_tensor_t::~deepspeed_pin_tensor_t()
+{
+    for (auto iter = _locked_tensors.begin(); iter != _locked_tensors.end(); ++iter) {
+        munlock(iter->first, iter->second);
+        std::free((void*)iter->first);
+    }
+    _locked_tensors.clear();
+}
+
+torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem,
+                                            const torch::TensorOptions& options)
+{
+    const auto scalar_dtype = torch::typeMetaToScalarType(options.dtype());
+    const auto num_bytes = num_elem * torch::elementSize(scalar_dtype);
+    auto pinned_buffer = ds_page_aligned_alloc(num_bytes, true);
+    assert(nullptr != pinned_buffer);
+
+    _locked_tensors[pinned_buffer] = num_bytes;
+
+    return at::from_blob(pinned_buffer, static_cast<int64_t>(num_elem), options);
+}
+
+torch::Tensor deepspeed_pin_tensor_t::alloc(const int64_t num_elem, const at::ScalarType& elem_type)
+{
+    auto options = torch::TensorOptions().dtype(elem_type).device(torch::kCPU).requires_grad(false);
+    return alloc(num_elem, options);
+}
+
+bool deepspeed_pin_tensor_t::free(torch::Tensor& locked_tensor)
+{
+    auto addr = locked_tensor.data_ptr();
+    if (_locked_tensors.find(addr) != _locked_tensors.end()) {
+        munlock(addr, _locked_tensors[addr]);
+        std::free(addr);
+        _locked_tensors.erase(addr);
+        return true;
+    }
+
+    return false;
+}
+
+bool deepspeed_pin_tensor_t::is_managed(const torch::Tensor& buffer)
+{
+    if (!buffer.is_cpu()) { return false; }
+    auto addr = buffer.data_ptr();
+    if (_locked_tensors.find(addr) != _locked_tensors.end()) { return true; }
+    return false;
+};

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h

@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for managing CPU tensors occupying page-locked memory.
+TODO: Implement a full-featured manager that
+1. Avoid page-locked memory leaks
+2. Minimize page-locked memory usage by reducing internal fragmentation
+Functionality for managing CPU tensors occupying page-locked memory.
+*/
+
+#include <map>
+#include "deepspeed_py_aio.h"
+
+struct deepspeed_pin_tensor_t {
+    std::map<void*, int64_t> _locked_tensors;
+
+    deepspeed_pin_tensor_t() = default;
+
+    ~deepspeed_pin_tensor_t();
+
+    torch::Tensor alloc(const int64_t num_elem, const at::ScalarType& elem_type);
+    torch::Tensor alloc(const int64_t num_elem, const torch::TensorOptions& options);
+
+    bool free(torch::Tensor& locked_tensor);
+
+    bool is_managed(const torch::Tensor& buffer);
+};

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp

@@ -0,0 +1,123 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <assert.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <cassert>
+#include <chrono>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "deepspeed_py_aio.h"
+
+using namespace std;
+using namespace std::chrono;
+
+#define DEBUG_DS_AIO_READ 0
+#define DEBUG_DS_AIO_WRITE 0
+
+static const std::string c_library_name = "deepspeed_aio";
+
+int deepspeed_py_aio_write(const torch::Tensor& buffer,
+                           const char* filename,
+                           const int block_size,
+                           const int queue_depth,
+                           const bool single_submit,
+                           const bool overlap_events,
+                           const bool validate)
+{
+    const auto start_time = std::chrono::high_resolution_clock::now();
+    deepspeed_aio_config_t config(block_size, queue_depth, single_submit, overlap_events, false);
+
+    const auto fd = open_file(filename, false);
+    if (fd == -1) { return -1; }
+
+    auto write_buffer = (char*)buffer.data_ptr();
+    const auto num_write_bytes = static_cast<int64_t>(buffer.nbytes());
+
+    std::unique_ptr<io_xfer_ctxt> xfer_ctxt(
+        new io_xfer_ctxt(fd, 0, 0, num_write_bytes, write_buffer));
+    std::unique_ptr<aio_context> aio_ctxt(new aio_context(config._block_size, config._queue_depth));
+
+    if (config._overlap_events) {
+        do_aio_operation_overlap(false, aio_ctxt, xfer_ctxt, &config, nullptr);
+    } else {
+        do_aio_operation_sequential(false, aio_ctxt, xfer_ctxt, &config, nullptr);
+    }
+    const std::chrono::duration<double> aio_time =
+        std::chrono::high_resolution_clock::now() - start_time;
+
+    close(fd);
+
+    if (validate) { validate_aio_operation(false, filename, write_buffer, num_write_bytes); }
+
+    const std::chrono::duration<double> fn_time =
+        std::chrono::high_resolution_clock::now() - start_time;
+    std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6
+              << " call = " << fn_time.count() * 1e6 << std::endl;
+    return 0;
+}
+
+int deepspeed_py_aio_read(torch::Tensor& buffer,
+                          const char* filename,
+                          const int block_size,
+                          const int queue_depth,
+                          const bool single_submit,
+                          const bool overlap_events,
+                          const bool validate)
+{
+    const auto start_time = std::chrono::high_resolution_clock::now();
+    int64_t num_file_bytes;
+    if (-1 == get_file_size(filename, num_file_bytes)) {
+        const auto error_code = errno;
+        report_file_error(filename, " fstat for read", error_code);
+        return -1;
+    }
+
+    deepspeed_aio_config_t config(block_size, queue_depth, single_submit, overlap_events, false);
+    const auto fd = open_file(filename, true);
+    if (fd == -1) { return -1; }
+
+    auto read_buffer = (char*)buffer.data_ptr();
+    assert(static_cast<int64_t>(buffer.nbytes()) == num_file_bytes);
+
+    std::unique_ptr<io_xfer_ctxt> xfer_ctxt(
+        new io_xfer_ctxt(fd, 0, 0, num_file_bytes, read_buffer));
+    std::unique_ptr<aio_context> aio_ctxt(new aio_context(config._block_size, config._queue_depth));
+
+    if (config._overlap_events) {
+        do_aio_operation_overlap(true, aio_ctxt, xfer_ctxt, &config, nullptr);
+    } else {
+        do_aio_operation_sequential(true, aio_ctxt, xfer_ctxt, &config, nullptr);
+    }
+    const std::chrono::duration<double> aio_time =
+        std::chrono::high_resolution_clock::now() - start_time;
+
+    close(fd);
+
+    if (validate) { validate_aio_operation(true, filename, read_buffer, num_file_bytes); }
+
+    const std::chrono::duration<double> fn_time =
+        std::chrono::high_resolution_clock::now() - start_time;
+    std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6
+              << " call = " << fn_time.count() * 1e6 << std::endl;
+    return 0;
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping tensors to/from (NVMe) storage devices.
+*/
+
+#include <deepspeed_aio_common.h>
+#include <stdlib.h>
+#include <torch/extension.h>
+
+int deepspeed_py_aio_write(const torch::Tensor& buffer,
+                           const char* filename,
+                           const int block_size,
+                           const int queue_depth,
+                           const bool single_submit,
+                           const bool overlap_events,
+                           const bool validate);
+
+int deepspeed_py_aio_read(torch::Tensor& buffer,
+                          const char* filename,
+                          const int block_size,
+                          const int queue_depth,
+                          const bool single_submit,
+                          const bool overlap_events,
+                          const bool validate);

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include "deepspeed_py_aio_handle.h"
+#include <cstdlib>
+
+using namespace std;
+
+deepspeed_aio_handle_t::deepspeed_aio_handle_t(const int block_size,
+                                               const int queue_depth,
+                                               const bool single_submit,
+                                               const bool overlap_events,
+                                               const int intra_op_parallelism)
+    : deepspeed_io_handle_t(block_size,
+                            queue_depth,
+                            single_submit,
+                            overlap_events,
+                            intra_op_parallelism)
+{
+}
+
+deepspeed_aio_handle_t::~deepspeed_aio_handle_t() {}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h

@@ -0,0 +1,22 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <condition_variable>
+#include <memory>
+#include "deepspeed_py_io_handle.h"
+
+struct deepspeed_aio_handle_t : deepspeed_io_handle_t {
+    deepspeed_aio_handle_t(const int block_size,
+                           const int queue_depth,
+                           const bool single_submit,
+                           const bool overlap_events,
+                           const int intra_op_parallelism);
+
+    ~deepspeed_aio_handle_t();
+};

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp

@@ -0,0 +1,135 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping tensors to/from (NVMe) storage devices.
+*/
+
+#include "deepspeed_py_copy.h"
+#include <omp.h>
+
+#define ROUND_DOWN(size, step) ((size) & ~((step) - 1))
+
+#if defined(__AVX512__) or defined(__AVX256__)
+union AVX_Data {
+#if defined(__AVX512__)
+    __m512 data;
+#else
+    __m256 data;
+#endif
+};
+#endif
+
+static void helper_memcpy_1(float* dest, float* src, size_t param_size)
+{
+    size_t rounded_size = 0;
+
+#if defined(__AVX512__) or defined(__AVX256__)
+
+    rounded_size = ROUND_DOWN(param_size, SIMD_WIDTH);
+
+    for (size_t t = 0; t < rounded_size; t += TILE) {
+        size_t copy_size = TILE;
+        if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
+        size_t offset = copy_size + t;
+#pragma omp parallel for
+        for (size_t i = t; i < offset; i += SIMD_WIDTH) {
+            AVX_Data src_4;
+            src_4.data = SIMD_LOAD(src + i);
+
+            SIMD_STORE(dest + i, src_4.data);
+        }
+    }
+
+#endif
+
+    if (param_size > rounded_size) {
+#pragma omp parallel for
+        for (size_t k = rounded_size; k < param_size; k++) { dest[k] = src[k]; }
+    }
+}
+
+static void helper_memcpy_4(float* dest, float* src, size_t param_size)
+{
+    size_t rounded_size = 0;
+
+#if defined(__AVX512__) or defined(__AVX256__)
+
+    rounded_size = ROUND_DOWN(param_size, (SIMD_WIDTH << 2));
+
+    for (size_t t = 0; t < rounded_size; t += TILE) {
+        size_t copy_size = TILE;
+        if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
+        size_t offset = copy_size + t;
+#pragma omp parallel for
+        for (size_t i = t; i < offset; i += (SIMD_WIDTH << 2)) {
+            AVX_Data src_4[4];
+            src_4[0].data = SIMD_LOAD(src + i);
+            src_4[1].data = SIMD_LOAD(src + i + SIMD_WIDTH);
+            src_4[2].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 1));
+            src_4[3].data = SIMD_LOAD(src + i + SIMD_WIDTH * 3);
+
+            SIMD_STORE(dest + i, src_4[0].data);
+            SIMD_STORE(dest + i + SIMD_WIDTH, src_4[1].data);
+            SIMD_STORE(dest + i + (SIMD_WIDTH << 1), src_4[2].data);
+            SIMD_STORE(dest + i + SIMD_WIDTH * 3, src_4[3].data);
+        }
+    }
+#endif
+    if (param_size > rounded_size)
+        helper_memcpy_1((dest + rounded_size), (src + rounded_size), (param_size - rounded_size));
+}
+
+static void helper_mempcy_8(float* dest, float* src, size_t param_size)
+{
+    size_t rounded_size = 0;
+
+#if defined(__AVX512__) or defined(__AVX256__)
+
+    rounded_size = ROUND_DOWN(param_size, (SIMD_WIDTH << 2));
+
+    for (size_t t = 0; t < rounded_size; t += TILE) {
+        size_t copy_size = TILE;
+        if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
+        size_t offset = copy_size + t;
+#pragma omp parallel for
+        for (size_t i = t; i < offset; i += (SIMD_WIDTH << 3)) {
+            AVX_Data src_4[8];
+            src_4[0].data = SIMD_LOAD(src + i);
+            src_4[1].data = SIMD_LOAD(src + i + SIMD_WIDTH);
+            src_4[2].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 1));
+            src_4[3].data = SIMD_LOAD(src + i + SIMD_WIDTH * 3);
+            src_4[4].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 2));
+            src_4[5].data = SIMD_LOAD(src + i + SIMD_WIDTH * 5);
+            src_4[6].data = SIMD_LOAD(src + i + SIMD_WIDTH * 6);
+            src_4[7].data = SIMD_LOAD(src + i + SIMD_WIDTH * 7);
+
+            SIMD_STORE(dest + i, src_4[0].data);
+            SIMD_STORE(dest + i + SIMD_WIDTH, src_4[1].data);
+            SIMD_STORE(dest + i + (SIMD_WIDTH << 1), src_4[2].data);
+            SIMD_STORE(dest + i + SIMD_WIDTH * 3, src_4[3].data);
+            SIMD_STORE(dest + i + (SIMD_WIDTH << 2), src_4[4].data);
+            SIMD_STORE(dest + i + SIMD_WIDTH * 5, src_4[5].data);
+            SIMD_STORE(dest + i + SIMD_WIDTH * 6, src_4[6].data);
+            SIMD_STORE(dest + i + SIMD_WIDTH * 7, src_4[7].data);
+        }
+    }
+#endif
+    if (param_size > rounded_size)
+        helper_memcpy_4((dest + rounded_size), (src + rounded_size), (param_size - rounded_size));
+}
+
+int deepspeed_py_memcpy(torch::Tensor& dest, const torch::Tensor& src)
+{
+    auto dest_c = dest.contiguous();
+    auto src_c = src.contiguous();
+
+    float* dest_ptr = (float*)dest_c.data_ptr();
+    float* src_ptr = (float*)src_c.data_ptr();
+
+    helper_mempcy_8(dest_ptr, src_ptr, dest_c.size(0));
+
+    return 0;
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#if (__x86_64__ || __i386__)
+#include <cpuid.h>
+#include <x86intrin.h>
+#endif
+
+#include <deepspeed_aio_common.h>
+#include <stdlib.h>
+#include <torch/extension.h>
+
+#define TILE (1024 * 1024 * 1024)
+
+#if defined(__AVX512__)
+#define SIMD_STORE(a, d) _mm512_storeu_ps(a, d)
+#define SIMD_LOAD(x) _mm512_loadu_ps(x)
+#define SIMD_SET(x) _mm512_set1_ps(x)
+#define SIMD_MUL(x, y) _mm512_mul_ps(x, y)
+#define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c)
+#define SIMD_SQRT(x) _mm512_sqrt_ps(x)
+#define SIMD_DIV(x, y) _mm512_div_ps(x, y)
+#define SIMD_WIDTH 16
+#else
+#if defined(__AVX256__)
+#define SIMD_STORE(a, d) _mm256_storeu_ps(a, d)
+#define SIMD_LOAD(x) _mm256_loadu_ps(x)
+#define SIMD_SET(x) _mm256_set1_ps(x)
+#define SIMD_MUL(x, y) _mm256_mul_ps(x, y)
+#define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c)
+#define SIMD_SQRT(x) _mm256_sqrt_ps(x)
+#define SIMD_DIV(x, y) _mm256_div_ps(x, y)
+#define SIMD_WIDTH 8
+#endif
+#endif
+
+int deepspeed_py_memcpy(torch::Tensor& dest, const torch::Tensor& src);

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp

@@ -0,0 +1,322 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include "deepspeed_py_io_handle.h"
+#include <cstdlib>
+
+using namespace std;
+
+static void _start_aio_thread(std::shared_ptr<struct deepspeed_aio_thread_t> ctxt) { ctxt->run(); }
+
+deepspeed_io_handle_t::deepspeed_io_handle_t(const int block_size,
+                                             const int queue_depth,
+                                             const bool single_submit,
+                                             const bool overlap_events,
+                                             const int intra_op_parallelism)
+    : _aio_ctxt(new aio_context(block_size, queue_depth)),
+      _single_submit(single_submit),
+      _overlap_events(overlap_events),
+      _intra_op_parallelism(intra_op_parallelism),
+      _aio_config(block_size, queue_depth, single_submit, overlap_events, false),
+      _num_pending_ops(0),
+      _pinned_tensor_mgr(new deepspeed_pin_tensor_t())
+{
+    for (auto i = 0; i < intra_op_parallelism; ++i) {
+        _thread_contexts.push_back(std::make_shared<deepspeed_aio_thread_t>(i, _aio_config));
+    }
+
+    for (auto& ctxt : _thread_contexts) {
+        _threads.push_back(std::thread(_start_aio_thread, ctxt));
+    }
+}
+
+deepspeed_io_handle_t::~deepspeed_io_handle_t()
+{
+    _stop_threads();
+    for (auto& thr : _threads) { thr.join(); }
+}
+
+const int deepspeed_io_handle_t::get_block_size() const
+{
+    return _aio_ctxt ? _aio_ctxt->_block_size : -1;
+}
+
+const int deepspeed_io_handle_t::get_queue_depth() const
+{
+    return _aio_ctxt ? _aio_ctxt->_queue_depth : -1;
+}
+
+const bool deepspeed_io_handle_t::get_single_submit() const { return _single_submit; }
+
+const bool deepspeed_io_handle_t::get_overlap_events() const { return _overlap_events; }
+
+const int deepspeed_io_handle_t::get_intra_op_parallelism() const { return _intra_op_parallelism; }
+
+int deepspeed_io_handle_t::read(torch::Tensor& buffer,
+                                const char* filename,
+                                const bool validate,
+                                const int64_t file_offset)
+{
+    const auto start_time = std::chrono::high_resolution_clock::now();
+
+    assert(_aio_ctxt);
+
+    int64_t num_file_bytes;
+    if (-1 == get_file_size(filename, num_file_bytes)) {
+        const auto error_code = errno;
+        report_file_error(filename, " fstat for read", error_code);
+        return -1;
+    }
+    assert(static_cast<int64_t>(buffer.nbytes()) == num_file_bytes);
+
+    const auto fd = open_file(filename, true);
+    if (fd == -1) { return -1; }
+
+    auto read_buffer = (char*)buffer.data_ptr();
+    std::unique_ptr<io_xfer_ctxt> xfer_ctxt(
+        new io_xfer_ctxt(fd, file_offset, 0, num_file_bytes, read_buffer));
+
+    if (_aio_config._overlap_events) {
+        do_aio_operation_overlap(true, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr);
+    } else {
+        do_aio_operation_sequential(true, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr);
+    }
+
+    close(fd);
+    const std::chrono::duration<double> aio_time =
+        std::chrono::high_resolution_clock::now() - start_time;
+
+    if (validate) { validate_aio_operation(true, filename, read_buffer, num_file_bytes); }
+    const std::chrono::duration<double> fn_time =
+        std::chrono::high_resolution_clock::now() - start_time;
+    std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6
+              << " call = " << fn_time.count() * 1e6 << std::endl;
+    return 0;
+}
+
+int deepspeed_io_handle_t::write(const torch::Tensor& buffer,
+                                 const char* filename,
+                                 const bool validate,
+                                 const int64_t file_offset)
+{
+    assert(_aio_ctxt);
+
+    const auto start_time = std::chrono::high_resolution_clock::now();
+
+    const auto fd = open_file(filename, false);
+    if (fd == -1) { return -1; }
+
+    auto write_buffer = (char*)buffer.data_ptr();
+    const auto num_write_bytes = static_cast<int64_t>(buffer.nbytes());
+    std::unique_ptr<io_xfer_ctxt> xfer_ctxt(
+        new io_xfer_ctxt(fd, file_offset, 0, num_write_bytes, write_buffer));
+
+    if (_aio_config._overlap_events) {
+        do_aio_operation_overlap(false, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr);
+    } else {
+        do_aio_operation_sequential(false, _aio_ctxt, xfer_ctxt, &_aio_config, nullptr);
+    }
+    const std::chrono::duration<double> aio_time =
+        std::chrono::high_resolution_clock::now() - start_time;
+
+    close(fd);
+
+    if (validate) { validate_aio_operation(false, filename, write_buffer, num_write_bytes); }
+
+    const std::chrono::duration<double> fn_time =
+        std::chrono::high_resolution_clock::now() - start_time;
+    std::cout << "Elapsed time(usec): " << "aio = " << aio_time.count() * 1e6
+              << " call = " << fn_time.count() * 1e6 << std::endl;
+    return 0;
+}
+
+void deepspeed_io_handle_t::_schedule_aio_work(std::shared_ptr<struct io_op_desc_t> scheduled_op)
+{
+    for (auto& ctxt : _thread_contexts) {
+        {
+            std::lock_guard<std::mutex> lock(ctxt->_work_sync._mutex);
+            ctxt->_work_queue.push(scheduled_op);
+        }
+        ctxt->_work_sync._cond_var.notify_one();
+    }
+    _num_pending_ops++;
+}
+
+std::shared_ptr<struct io_op_desc_t> deepspeed_io_handle_t::_wait_for_aio_work()
+{
+    std::shared_ptr<struct io_op_desc_t> completed_op = nullptr;
+    for (auto& ctxt : _thread_contexts) {
+        std::unique_lock<std::mutex> lock(ctxt->_complete_sync._mutex);
+        ctxt->_complete_sync._cond_var.wait(lock,
+                                            [ctxt] { return !ctxt->_complete_queue.empty(); });
+        completed_op = ctxt->_complete_queue.front();
+        ctxt->_complete_queue.pop();
+    }
+    return completed_op;
+}
+
+void deepspeed_io_handle_t::_stop_threads()
+{
+    assert(0 == _num_pending_ops);
+    for (auto& ctxt : _thread_contexts) {
+        {
+            std::lock_guard<std::mutex> lock(ctxt->_work_sync._mutex);
+            ctxt->_time_to_exit = true;
+        }
+        ctxt->_work_sync._cond_var.notify_one();
+    }
+}
+
+int deepspeed_io_handle_t::wait()
+{
+    assert(_num_pending_ops > 0);
+    auto num_completed_ops = 0;
+
+    while (_num_pending_ops > 0) {
+        auto completed_op = _wait_for_aio_work();
+
+        if (completed_op->_validate) { completed_op->validate(); }
+
+        completed_op->finish();
+
+        close(completed_op->_fd);
+
+        --_num_pending_ops;
+        ++num_completed_ops;
+    }
+
+    return num_completed_ops;
+}
+
+bool deepspeed_io_handle_t::_is_valid_parallel_aio_op(const bool read_op, const int64_t num_bytes)
+{
+    const auto op_string = read_op ? "Read" : "Write";
+    if (num_bytes % get_intra_op_parallelism()) {
+        std::cout << "deepspeed_aio failure: parallel " << op_string << " num_bytes = " << num_bytes
+                  << " not divisible by thread count = " << get_intra_op_parallelism() << std::endl;
+        return false;
+    }
+
+    return true;
+}
+
+std::shared_ptr<struct io_op_desc_t> deepspeed_io_handle_t::_create_io_op_desc(
+    const bool read_op,
+    const torch::Tensor& buffer,
+    const int fd,
+    const char* filename,
+    const int64_t file_num_bytes,
+    const bool validate,
+    const int64_t file_offset)
+{
+    return std::make_shared<cpu_op_desc_t>(read_op,
+                                           buffer,
+                                           _pinned_tensor_mgr,
+                                           fd,
+                                           filename,
+                                           file_num_bytes,
+                                           _intra_op_parallelism,
+                                           validate,
+                                           file_offset);
+}
+
+int deepspeed_io_handle_t::pread(const torch::Tensor& buffer,
+                                 const char* filename,
+                                 const bool validate,
+                                 const bool async,
+                                 const int64_t file_offset)
+{
+    int64_t num_file_bytes;
+    if (-1 == get_file_size(filename, num_file_bytes)) {
+        const auto error_code = errno;
+        report_file_error(filename, " fstat for read", error_code);
+        return -1;
+    }
+
+    // buffer can exceed file size to enable 4k alignment
+    const auto buffer_bytes = static_cast<int64_t>(buffer.nbytes());
+    assert((num_file_bytes % _intra_op_parallelism) == 0);
+
+    if (!_is_valid_parallel_aio_op(true, buffer_bytes)) { return -1; }
+
+    const auto fd = open_file(filename, true);
+    if (fd == -1) { return -1; }
+
+    auto scheduled_op =
+        _create_io_op_desc(true, buffer, fd, filename, num_file_bytes, validate, file_offset);
+
+    _schedule_aio_work(scheduled_op);
+
+    if (async) { return 0; }
+
+    return wait();
+}
+
+int deepspeed_io_handle_t::pwrite(const torch::Tensor& buffer,
+                                  const char* filename,
+                                  const bool validate,
+                                  const bool async,
+                                  const int64_t file_offset)
+{
+    const auto num_write_bytes = static_cast<int64_t>(buffer.nbytes());
+    assert((num_write_bytes % _intra_op_parallelism) == 0);
+
+    if (!_is_valid_parallel_aio_op(false, num_write_bytes)) { return -1; }
+
+    const auto fd = open_file(filename, false);
+    if (fd == -1) { return -1; }
+
+    auto scheduled_op =
+        _create_io_op_desc(false, buffer, fd, filename, num_write_bytes, validate, file_offset);
+
+    _schedule_aio_work(scheduled_op);
+
+    if (async) { return 0; }
+
+    return wait();
+}
+
+int deepspeed_io_handle_t::sync_pread(torch::Tensor& buffer,
+                                      const char* filename,
+                                      const int64_t file_offset)
+{
+    return pread(buffer, filename, false, false, file_offset);
+}
+
+int deepspeed_io_handle_t::sync_pwrite(const torch::Tensor& buffer,
+                                       const char* filename,
+                                       const int64_t file_offset)
+{
+    return pwrite(buffer, filename, false, false, file_offset);
+}
+
+int deepspeed_io_handle_t::async_pread(torch::Tensor& buffer,
+                                       const char* filename,
+                                       const int64_t file_offset)
+{
+    return pread(buffer, filename, false, true, file_offset);
+}
+
+int deepspeed_io_handle_t::async_pwrite(const torch::Tensor& buffer,
+                                        const char* filename,
+                                        const int64_t file_offset)
+{
+    return pwrite(buffer, filename, false, true, file_offset);
+}
+
+at::Tensor deepspeed_io_handle_t::new_cpu_locked_tensor(const int64_t num_elem,
+                                                        const torch::Tensor& example_tensor)
+{
+    return _pinned_tensor_mgr->alloc(num_elem, example_tensor.scalar_type());
+}
+
+bool deepspeed_io_handle_t::free_cpu_locked_tensor(torch::Tensor& locked_tensor)
+{
+    return _pinned_tensor_mgr->free(locked_tensor);
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h

@@ -0,0 +1,94 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <condition_variable>
+#include <memory>
+#include "deepspeed_aio_thread.h"
+#include "deepspeed_pin_tensor.h"
+
+struct deepspeed_io_handle_t {
+    std::unique_ptr<struct aio_context> _aio_ctxt;
+    const bool _single_submit;
+    const bool _overlap_events;
+    const int _intra_op_parallelism;
+    deepspeed_aio_config_t _aio_config;
+
+    std::vector<std::shared_ptr<struct deepspeed_aio_thread_t>> _thread_contexts;
+    std::vector<std::thread> _threads;
+    int _num_pending_ops;
+    std::unique_ptr<struct deepspeed_pin_tensor_t> _pinned_tensor_mgr;
+
+    deepspeed_io_handle_t(const int block_size,
+                          const int queue_depth,
+                          const bool single_submit,
+                          const bool overlap_events,
+                          const int intra_op_parallelism);
+
+    virtual ~deepspeed_io_handle_t() = 0;
+
+    const int get_block_size() const;
+    const int get_queue_depth() const;
+    const bool get_single_submit() const;
+    const bool get_overlap_events() const;
+    const int get_intra_op_parallelism() const;
+
+    int read(torch::Tensor& buffer,
+             const char* filename,
+             const bool validate,
+             const int64_t file_offset);
+
+    int write(const torch::Tensor& buffer,
+              const char* filename,
+              const bool validate,
+              const int64_t file_offset);
+
+    int pread(const torch::Tensor& buffer,
+              const char* filename,
+              const bool validate,
+              const bool async,
+              const int64_t file_offset);
+
+    int pwrite(const torch::Tensor& buffer,
+               const char* filename,
+               const bool validate,
+               const bool async,
+               const int64_t file_offset);
+
+    int sync_pread(torch::Tensor& buffer, const char* filename, const int64_t file_offset);
+
+    int sync_pwrite(const torch::Tensor& buffer, const char* filename, const int64_t file_offset);
+
+    int async_pread(torch::Tensor& buffer, const char* filename, const int64_t file_offset);
+
+    int async_pwrite(const torch::Tensor& buffer, const char* filename, const int64_t file_offset);
+
+    // TODO: Make API's args to be shape and dtype.
+    torch::Tensor new_cpu_locked_tensor(const int64_t num_elem,
+                                        const torch::Tensor& example_tensor);
+
+    bool free_cpu_locked_tensor(torch::Tensor&);
+
+    int wait();
+
+    void _stop_threads();
+
+    void _schedule_aio_work(std::shared_ptr<struct io_op_desc_t> scheduled_op);
+
+    std::shared_ptr<struct io_op_desc_t> _wait_for_aio_work();
+
+    bool _is_valid_parallel_aio_op(const bool read_op, const int64_t num_bytes);
+
+    virtual std::shared_ptr<struct io_op_desc_t> _create_io_op_desc(const bool read_op,
+                                                                    const torch::Tensor& buffer,
+                                                                    const int fd,
+                                                                    const char* filename,
+                                                                    const int64_t file_num_bytes,
+                                                                    const bool validate,
+                                                                    const int64_t file_offset);
+};

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp

@@ -0,0 +1,116 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
+*/
+
+#include <torch/extension.h>
+#include "deepspeed_py_aio_handle.h"
+#include "deepspeed_py_copy.h"
+using namespace pybind11::literals;
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("aio_read", &deepspeed_py_aio_read, "DeepSpeed Asynchronous I/O Read");
+
+    m.def("aio_write", &deepspeed_py_aio_write, "DeepSpeed Asynchronous I/O Write");
+
+    m.def("deepspeed_memcpy", &deepspeed_py_memcpy, "DeepSpeed Memory Copy");
+
+    py::class_<deepspeed_aio_handle_t>(m, "aio_handle")
+        .def(py::init<const int, const int, const bool, const bool, const int>(),
+             "AIO handle constructor",
+             "block_size"_a = 1024 * 1024,
+             "queue_depth"_a = 128,
+             "single_submit"_a = false,
+             "overlap_events"_a = false,
+             "intra_op_parallelism"_a = 1)
+
+        .def("get_block_size", &deepspeed_aio_handle_t::get_block_size)
+        .def("get_queue_depth", &deepspeed_aio_handle_t::get_queue_depth)
+        .def("get_single_submit", &deepspeed_aio_handle_t::get_single_submit)
+        .def("get_overlap_events", &deepspeed_aio_handle_t::get_overlap_events)
+        .def("get_intra_op_parallelism", &deepspeed_aio_handle_t::get_intra_op_parallelism)
+
+        .def("read",
+             &deepspeed_aio_handle_t::read,
+             "Synchronous and non-parallel file read. Returns count of completed read ops",
+             "buffer"_a,
+             "filename"_a,
+             "validate"_a,
+             "file_offset"_a = 0)
+
+        .def("write",
+             &deepspeed_aio_handle_t::write,
+             "Synchronous and non-parallel file write. Returns count of completed write ops",
+             "buffer"_a,
+             "filename"_a,
+             "validate"_a,
+             "file_offset"_a = 0)
+
+        .def("pread",
+             &deepspeed_aio_handle_t::pread,
+             "Parallel file read with option of parallelism. Returns count of completed read ops",
+             "buffer"_a,
+             "filename"_a,
+             "validate"_a,
+             "async"_a,
+             "file_offset"_a = 0)
+
+        .def("pwrite",
+             &deepspeed_aio_handle_t::pwrite,
+             "Parallel file write with option of parallelism. Returns count of completed write ops",
+             "buffer"_a,
+             "filename"_a,
+             "validate"_a,
+             "async"_a,
+             "file_offset"_a = 0)
+
+        .def("sync_pread",
+             &deepspeed_aio_handle_t::sync_pread,
+             "Synchrononous parallel file read. Returns count of completed read ops",
+             "buffer"_a,
+             "filename"_a,
+             "file_offset"_a = 0)
+
+        .def("sync_pwrite",
+             &deepspeed_aio_handle_t::sync_pwrite,
+             "Synchronous parallel file write. Returns count of completed write ops",
+             "buffer"_a,
+             "filename"_a,
+             "file_offset"_a = 0)
+
+        .def("async_pread",
+             &deepspeed_aio_handle_t::async_pread,
+             "Asynchronous parallel file read. Returns 0 on success. Returns 0 on success, and "
+             "following wait() returns count of completed ops.",
+             "buffer"_a,
+             "filename"_a,
+             "file_offset"_a = 0)
+
+        .def("async_pwrite",
+             &deepspeed_aio_handle_t::async_pwrite,
+             "Asynchronous parallel file write. Returns 0 on success, and following wait() returns "
+             "count of completed ops.",
+             "buffer"_a,
+             "filename"_a,
+             "file_offset"_a = 0)
+
+        .def("new_cpu_locked_tensor",
+             &deepspeed_aio_handle_t::new_cpu_locked_tensor,
+             "Allocate pinned CPU tensor.",
+             "num_elem"_a,
+             "example_tenosr"_a)
+
+        .def("free_cpu_locked_tensor",
+             &deepspeed_aio_handle_t::free_cpu_locked_tensor,
+             "Free pinned CPU tensor.",
+             "tensor"_a)
+
+        .def("wait",
+             &deepspeed_aio_handle_t::wait,
+             "Wait for (ongoing) asynchronous operations to complete");
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json

@@ -0,0 +1,29 @@
+{
+    "block_size": [
+        "128K",
+        "256K",
+        "1M"
+    ],
+    "queue_depth": [
+        4,
+        16,
+        32
+    ],
+    "io_parallel": [
+        1,
+        2,
+        4,
+        8
+    ],
+    "single_submit": [
+        true,
+        false
+    ],
+    "overlap_events": [
+        true,
+        false
+    ],
+    "threads": [
+        1
+    ]
+}

lib/python3.12/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp

@@ -0,0 +1,188 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "deepcompile.h"
+
+#define USE_C10D_NCCL
+
+namespace dc {
+
+std::shared_ptr<DSParamRegistry> param_registry;
+std::unordered_map<long, std::shared_ptr<CustomOpExecutor>> executors;
+std::shared_ptr<DoubleBufferedReduceBucket> reduce_buckets = nullptr;
+
+c10::intrusive_ptr<c10d::ProcessGroup> process_group = nullptr;
+c10::intrusive_ptr<c10d::symmetric_memory::SymmetricMemory> symm_mem = nullptr;
+ncclComm_t nccl_comm;
+bool use_symm_mem;
+bool clone_custom_op_output;
+bool profile = false;
+bool pre_div_reduce = true;
+
+bool sync_before_reduce;     // for debugging
+bool sync_after_reduce;      // for debugging
+bool sync_before_allgather;  // for debugging
+bool sync_after_allgather;   // for debugging
+
+std::vector<int64_t> sizes_to_int_vector(at::IntArrayRef sizes)
+{
+    std::vector<int64_t> result;
+    for (int i = 0; i < sizes.size(); i++) { result.push_back(sizes[i]); }
+    return result;
+}
+
+void enable_profiling(bool enable) { profile = enable; }
+
+bool is_profiling() { return profile; }
+
+c10::intrusive_ptr<c10d::symmetric_memory::SymmetricMemory> getSymmMemWorkspace(int64_t size)
+{
+    c10::Device device = c10::Device(c10::kCUDA, c10::cuda::current_device());
+    std::vector<int64_t> sizes = {size};
+    std::vector<int64_t> strides = {1};
+    at::Tensor sym_mem_ws = c10d::symmetric_memory::empty_strided_p2p(
+        {size}, {1}, c10::ScalarType::Byte, device, process_group->getGroupName(), std::nullopt);
+    return c10d::symmetric_memory::rendezvous(sym_mem_ws);
+}
+
+void lazy_init_symm_memory()
+{
+    if (use_symm_mem && !symm_mem) {
+        int64_t max_param_size = 0;
+        for (const auto& it : param_registry->getParams()) {
+            int64_t size = it.second.getDSTensor().numel() * it.second.getDSTensor().element_size();
+            if (size > max_param_size) { max_param_size = size; }
+        }
+        symm_mem = getSymmMemWorkspace(max_param_size);
+    }
+}
+
+ncclDataType_t get_nccl_data_type(at::ScalarType scalar_type)
+{
+    switch (scalar_type) {
+        case at::kFloat: return ncclFloat;
+        case at::kHalf: return ncclHalf;
+        case at::kDouble: return ncclDouble;
+        case at::kBFloat16: return ncclBfloat16;
+        case at::kLong: return ncclInt64;
+        case at::kInt: return ncclInt;
+        case at::kChar: return ncclInt8;
+        default: throw std::runtime_error("Unsupported scalar type");
+    }
+}
+
+void reset()
+{
+    executors.clear();
+    // We keep the buckets for memory estimation
+    // reduce_buckets->clear();
+}
+
+void cleanup()
+{
+    reset();
+
+    ncclCommDestroy(nccl_comm);
+    process_group = nullptr;
+    symm_mem = nullptr;
+}
+
+at::Tensor reduce_grad(at::Tensor grad_tensor, long graph_id, long ds_id)
+{
+    if (sync_before_reduce) { c10::cuda::device_synchronize(); }
+
+    assert(hasKey(executors, graph_id));
+    if (!profile) { executors[graph_id]->reduceGrad(grad_tensor, ds_id); }
+
+    if (sync_after_reduce) { c10::cuda::device_synchronize(); }
+
+    return at::Tensor();
+}
+
+at::Tensor reduce_grad_meta(at::Tensor grad_tensor, long graph_id, long ds_id)
+{
+    return at::Tensor();
+}
+
+void free_tensors(std::vector<at::Tensor> tensors)
+{
+    int64_t THRESHOLD = 10 * 1024 * 1024;
+
+    if (!profile) {
+        for (auto& tensor : tensors) {
+            if (tensor.is_cuda() && tensor.numel() > THRESHOLD) {
+                tensor.record_stream(at::cuda::getCurrentCUDAStream());
+                tensor.set_data(torch::empty({0}, tensor.options()));
+            }
+        }
+    }
+}
+
+void free_tensors_meta(std::vector<at::Tensor> tensors) {}
+
+void init(c10::intrusive_ptr<c10d::ProcessGroup> pg,
+          int64_t initial_reduce_bucket_size,
+          bool enable_double_buffer,
+          bool _use_symm_mem,
+          bool _clone_custom_op_output,
+          bool _sync_before_reduce,
+          bool _sync_after_reduce,
+          bool _sync_before_allgather,
+          bool _sync_after_allgather)
+{
+    process_group = pg;
+
+    ncclUniqueId ncclID;
+    ncclGetUniqueId(&ncclID);
+
+    // ProcessGroup doesn't have an API to get the CUDA stream for comm calls.
+    // So we create a NCCL communicator and call NCCL APIs directly.
+    auto vec = std::vector<uint8_t>(reinterpret_cast<uint8_t*>(&ncclID),
+                                    reinterpret_cast<uint8_t*>(&ncclID) + NCCL_UNIQUE_ID_BYTES);
+    auto device = torch::Device(torch::kCUDA);
+    at::Tensor tensor = torch::from_blob(vec.data(), {static_cast<long>(vec.size())}, torch::kUInt8)
+                            .to(torch::Device(torch::kCUDA));
+    std::vector<at::Tensor> bcast_input = {tensor};
+
+    process_group->broadcast(bcast_input, c10d::BroadcastOptions())->wait();
+
+    // create a new nccl communicator
+    std::memcpy(&ncclID, tensor.to(torch::Device(torch::kCPU)).data_ptr(), NCCL_UNIQUE_ID_BYTES);
+    ncclCommInitRank(&nccl_comm, process_group->getSize(), ncclID, process_group->getRank());
+
+    param_registry = std::make_shared<DSParamRegistry>();
+    reduce_buckets = std::make_shared<DoubleBufferedReduceBucket>(initial_reduce_bucket_size,
+                                                                  enable_double_buffer);
+    use_symm_mem = _use_symm_mem;
+    clone_custom_op_output = _clone_custom_op_output;
+
+    sync_before_reduce = _sync_before_reduce;
+    sync_after_reduce = _sync_after_reduce;
+    sync_before_allgather = _sync_before_allgather;
+    sync_after_allgather = _sync_after_allgather;
+}
+
+void start_forward()
+{
+    lazy_init_symm_memory();
+    for (auto& it : executors) { it.second->startForward(); }
+}
+
+void end_forward()
+{
+    for (auto& it : executors) { it.second->endForward(); }
+}
+
+void start_backward(bool update)
+{
+    for (auto& it : executors) { it.second->startBackward(update); }
+}
+
+// We don't call this
+// void end_backward(bool update)
+// {
+// }
+
+}  // namespace dc