diff --git a/lib/python3.12/site-packages/deepspeed/compression/__init__.py b/lib/python3.12/site-packages/deepspeed/compression/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8c7e1c89387326f7c043f95ff0b2153c4b9f21fe --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/deepspeed/compression/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a6ea5b3fe6ce9ee489fc1d5e7b82d66416d121d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compression/__pycache__/basic_layer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/basic_layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a67599f0e6a2676a6ab38ece287def23875f362 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/basic_layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compression/__pycache__/compress.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/compress.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc51691ee5f8ffc718c400fdc0623b2bbd476731 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/compress.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compression/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d8632cb2659f5fe1a48212ade41921d82168739 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compression/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fea551f7db12ced8407636b7540dc4b98df0153 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compression/__pycache__/helper.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/helper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..301d08d9ecabc9248392f8f04f12c976477681f4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/helper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compression/__pycache__/scheduler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..90e276c87c4dc76c40a638fe61393f24fac974ae Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compression/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..227c3f546236b6ac6a5f2861578be023dd256f61 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/compression/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/compression/basic_layer.py b/lib/python3.12/site-packages/deepspeed/compression/basic_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..bc2b54951bbe9833fb533384cffbef86b513862f --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/deepspeed/compression/compress.py b/lib/python3.12/site-packages/deepspeed/compression/compress.py new file mode 100644 index 0000000000000000000000000000000000000000..2f0e88beee21c040708f6fc7f279826f6c7d5c60 --- /dev/null +++ b/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) diff --git a/lib/python3.12/site-packages/deepspeed/compression/config.py b/lib/python3.12/site-packages/deepspeed/compression/config.py new file mode 100644 index 0000000000000000000000000000000000000000..e1fa5ef4bdb5f01fd609ae30c63406f001aaf781 --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/deepspeed/compression/constants.py b/lib/python3.12/site-packages/deepspeed/compression/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..67375d510a4b0a82c8860040042c5e0719507803 --- /dev/null +++ b/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" diff --git a/lib/python3.12/site-packages/deepspeed/compression/helper.py b/lib/python3.12/site-packages/deepspeed/compression/helper.py new file mode 100644 index 0000000000000000000000000000000000000000..e2d66100a4689255eff210465ff9f8cc156b6b0b --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/deepspeed/compression/scheduler.py b/lib/python3.12/site-packages/deepspeed/compression/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..85fdb67f642f60c13ed1327e42081fc30dce6a04 --- /dev/null +++ b/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() diff --git a/lib/python3.12/site-packages/deepspeed/compression/utils.py b/lib/python3.12/site-packages/deepspeed/compression/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..481e833bdf8ccef9f136155bb83790c52aca8bac --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/deepspeed/ops/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9d2e586ce719b8ac8ed6295c8c1a51e3b5239fb Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/compile/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/compile/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e38d56359fea64b067373b0e56cfed1dca5dab9a --- /dev/null +++ b/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 diff --git a/lib/python3.12/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4bb4ec9796cecd1d5fa140f48a658048a9ea6de Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/compile/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/adagrad/cpu_adagrad.cpp new file mode 100644 index 0000000000000000000000000000000000000000..e276ad0856dd620ecbe47231c5ebc4a8a753cb1f --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include + +using namespace std::string_literals; +static std::unordered_map> s_optimizers; + +// C++ interface + +template +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 +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(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 +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 +void step_invoker(std::shared_ptr 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::function, void*, void*, void*, size_t)>> + invokers; + +// Fill map with template functions for each type +template +void create_invoker() +{ + invokers[std::tuple(c10::CppTypeToScalarType(), + c10::CppTypeToScalarType())] = + step_invoker; +} +struct InvokerInitializer { + InvokerInitializer() + { + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + } +} _invoker_initializer; + +void invoke(std::shared_ptr 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 opt = + std::static_pointer_cast(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++)"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/cpu_adam.cpp new file mode 100644 index 0000000000000000000000000000000000000000..263c443cb4d4e5f8135f208a6fde9ced513ac8ae --- /dev/null +++ b/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++)"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/cpu_adam_impl.cpp new file mode 100644 index 0000000000000000000000000000000000000000..465aae7b9a343d5a3315c1531e11e8648c585b58 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include +#include "cpu_adam.h" + +using namespace std::string_literals; +static std::unordered_map> s_optimizers; + +// C++ interface + +template +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 +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(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 +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 +void step_invoker(std::shared_ptr 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::function, void*, void*, void*, void*, size_t)>> + invokers; + +// Fill map with template functions for each type +template +void create_invoker() +{ + invokers[std::tuple(c10::CppTypeToScalarType(), + c10::CppTypeToScalarType())] = + step_invoker; +} +struct InvokerInitializer { + InvokerInitializer() + { + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + create_invoker(); + } +} _invoker_initializer; + +void invoke(std::shared_ptr 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 opt = + std::static_pointer_cast(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; +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp new file mode 100644 index 0000000000000000000000000000000000000000..13b390248608b046dab443f85346b5446a47d722 --- /dev/null +++ b/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 + +void multi_tensor_adam_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> 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"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu new file mode 100644 index 0000000000000000000000000000000000000000..a1fc7d15aec9ad16ddb242fd62c99cbfdca55d0c --- /dev/null +++ b/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 +#include +#include +#include +// Another possibility: +// #include + +#include + +#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 +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> 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(), + 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(), + beta1, + beta2, + bias_correction1, + bias_correction2, + epsilon, + lr, + (adamMode_t)mode, + weight_decay);) + } + + AT_CUDA_CHECK(cudaGetLastError()); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh b/lib/python3.12/site-packages/deepspeed/ops/csrc/adam/multi_tensor_apply.cuh new file mode 100644 index 0000000000000000000000000000000000000000..342376c141be7578553910df43b003d29e0e4fc4 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include "compat.h" + +#include + +// #include + +// 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 +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 +__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 +void multi_tensor_apply(int64_t block_size, + int64_t chunk_size, + const at::Tensor& noop_flag, + const std::vector>& 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 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; + multi_tensor_apply_kernel<<>>( + chunk_size, noop_flag.DATA_PTR(), 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; + } + } + } + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.cpp new file mode 100644 index 0000000000000000000000000000000000000000..9d7ff5093017f60c06b96554f69630ffe9825918 --- /dev/null +++ b/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 +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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>& latencies) + __attribute__((unused)); + +static void _report_aio_statistics(const char* tag, + const std::vector>& latencies) +{ + std::vector 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>& raw_latencies, + struct deepspeed_aio_latency_t& summary_latencies) +{ + std::vector 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_ctxt, + std::vector>& 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_ctxt, + std::vector>& 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_ctxt, + std::vector>& 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_ctxt, + std::unique_ptr& 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( + ceil(static_cast(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> submit_times; + std::vector> reap_times; + const auto max_queue_bytes = + static_cast(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(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 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_ctxt, + std::unique_ptr& 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> submit_times; + std::vector> 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 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& 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 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(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; +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_common.h new file mode 100644 index 0000000000000000000000000000000000000000..aa4e49f4f4edee792e81bc0581e16ef3ab968811 --- /dev/null +++ b/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 +#include +#include +#include + +using namespace std; + +void do_aio_operation_sequential(const bool read_op, + std::unique_ptr& aio_ctxt, + std::unique_ptr& 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_ctxt, + std::unique_ptr& 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& buffer); + +bool validate_aio_operation(const bool read_op, + const char* filename, + void* aio_buffer, + const int64_t num_bytes); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.cpp new file mode 100644 index 0000000000000000000000000000000000000000..5e34a61065d4b0f16ba904d0e397eeccc149e621 --- /dev/null +++ b/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 + +#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); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_types.h new file mode 100644 index 0000000000000000000000000000000000000000..ce6a4e5cdfa75120c24dad599a8bc717765c220d --- /dev/null +++ b/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 +#include + +#include +#include + +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 _io_events; + std::vector _iocbs; + int _block_size; + int _queue_depth; + + aio_context(const int block_size, const int queue_depth); + ~aio_context(); +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.cpp new file mode 100644 index 0000000000000000000000000000000000000000..fb269b58315fb4653ff296175a71a61ede2e9467 --- /dev/null +++ b/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 +#include + +#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& xfer_ctxt, + const size_t block_size, + const std::vector* 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(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& 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(ceil(static_cast(xfer_ctxt->_num_bytes) / block_size)); + _remaining_io_blocks = _num_io_blocks; +} + +int io_prep_generator::prep_iocbs(const int n_iocbs, std::vector* iocbs) +{ + if ((_remaining_bytes) == 0 || (_remaining_io_blocks == 0)) { + assert(static_cast(_remaining_bytes) == _remaining_io_blocks); + return 0; + } + + assert(static_cast(n_iocbs) <= iocbs->size()); + + auto actual_n_iocbs = min(static_cast(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(_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; +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/common/deepspeed_aio_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..6b7599acecb44189b34257a7e68e2a160ffcdaef --- /dev/null +++ b/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 +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +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& _xfer_ctxt; + const size_t _block_size; + const std::vector* _iocbs; + + io_prep_context(const bool read_op, + const std::unique_ptr& xfer_ctxt, + const size_t block_size, + const std::vector* 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& _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& xfer_ctxt, + const size_t block_size); + + int prep_iocbs(const int n_iocbs, std::vector* iocbs); +}; + +void* ds_page_aligned_alloc(const int64_t size, const bool lock = false); + +int get_file_size(const char* filename, int64_t& size); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.cpp new file mode 100644 index 0000000000000000000000000000000000000000..94525139722571576a991b8414b599174ba4181c --- /dev/null +++ b/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(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_ctxt, + deepspeed_aio_config_t* aio_config) +{ +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_op_desc.h new file mode 100644 index 0000000000000000000000000000000000000000..ac1cdf90f78bee7d0cac45966a5dc8469201f424 --- /dev/null +++ b/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 +#include +#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_ctxt, + deepspeed_aio_config_t* aio_config); + + virtual char* data_ptr() const; + + virtual void validate(); + + virtual void finish(); +}; +#endif // _IO_OP_DESC_T_ diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.cpp new file mode 100644 index 0000000000000000000000000000000000000000..30c3b49143979cd4e71f7e0898b0d454a8099a51 --- /dev/null +++ b/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 next_io_op = nullptr; + + { + std::unique_lock 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 lock(_complete_sync._mutex); + _complete_queue.push(next_io_op); + } + _complete_sync._cond_var.notify_one(); + } + + if (_time_to_exit) { break; } + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_aio_thread.h new file mode 100644 index 0000000000000000000000000000000000000000..a192804db13d8a9a1a160134ce5f90e0ca091cd7 --- /dev/null +++ b/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 +#include +#include +#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 _aio_ctxt; + std::queue> _work_queue; + std::queue> _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(); +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.cpp new file mode 100644 index 0000000000000000000000000000000000000000..56fb33fb188696421c325acc0bd11148761937f7 --- /dev/null +++ b/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& 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_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 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); } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_cpu_op.h new file mode 100644 index 0000000000000000000000000000000000000000..debaf4a90731fff3bcbf1f45e92fb1b0f4b3a8eb --- /dev/null +++ b/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 +#include +#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& _pinned_tensor_mgr; + + cpu_op_desc_t(const bool read_op, + const torch::Tensor& buffer, + const std::unique_ptr& 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_ctxt, + deepspeed_aio_config_t* aio_config); + + char* data_ptr() const; + + void validate(); + + void finish(); + + void _alloc_bounce_buffer(); + void _free_bounce_buffer(); +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.cpp new file mode 100644 index 0000000000000000000000000000000000000000..a97a4ac18ba811dfaad713fdee15248caecf035b --- /dev/null +++ b/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(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; +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_pin_tensor.h new file mode 100644 index 0000000000000000000000000000000000000000..4b8ad7e76085877818aa26e511968efb0dc31fd2 --- /dev/null +++ b/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 +#include "deepspeed_py_aio.h" + +struct deepspeed_pin_tensor_t { + std::map _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); +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.cpp new file mode 100644 index 0000000000000000000000000000000000000000..1ff0397043fae2d904fb22d5d0b9eea5af6b6cd4 --- /dev/null +++ b/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 +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#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(buffer.nbytes()); + + std::unique_ptr xfer_ctxt( + new io_xfer_ctxt(fd, 0, 0, num_write_bytes, write_buffer)); + std::unique_ptr 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 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 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(buffer.nbytes()) == num_file_bytes); + + std::unique_ptr xfer_ctxt( + new io_xfer_ctxt(fd, 0, 0, num_file_bytes, read_buffer)); + std::unique_ptr 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 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 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; +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio.h new file mode 100644 index 0000000000000000000000000000000000000000..ba794db5440d540967053fc2ae4b184cfb8825d8 --- /dev/null +++ b/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 +#include +#include + +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); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.cpp new file mode 100644 index 0000000000000000000000000000000000000000..2b1093e9928659bdaac1c14323928c644061155e --- /dev/null +++ b/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 + +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() {} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_aio_handle.h new file mode 100644 index 0000000000000000000000000000000000000000..1398df9a56c968a6e174fd30586f5afcbbad3421 --- /dev/null +++ b/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 +#include +#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(); +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.cpp new file mode 100644 index 0000000000000000000000000000000000000000..f5480e9d9d836db587cc9ded27c8f437bc294452 --- /dev/null +++ b/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 + +#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; +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_copy.h new file mode 100644 index 0000000000000000000000000000000000000000..f443571a3e7b35f3bc70ec835152c34e5d979137 --- /dev/null +++ b/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 +#include +#endif + +#include +#include +#include + +#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); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.cpp new file mode 100644 index 0000000000000000000000000000000000000000..64d7c2e0541e12a2361b90c63575622e798b6576 --- /dev/null +++ b/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 + +using namespace std; + +static void _start_aio_thread(std::shared_ptr 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(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(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 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 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 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(buffer.nbytes()); + std::unique_ptr 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 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 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 scheduled_op) +{ + for (auto& ctxt : _thread_contexts) { + { + std::lock_guard lock(ctxt->_work_sync._mutex); + ctxt->_work_queue.push(scheduled_op); + } + ctxt->_work_sync._cond_var.notify_one(); + } + _num_pending_ops++; +} + +std::shared_ptr deepspeed_io_handle_t::_wait_for_aio_work() +{ + std::shared_ptr completed_op = nullptr; + for (auto& ctxt : _thread_contexts) { + std::unique_lock 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 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 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(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(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(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); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/deepspeed_py_io_handle.h new file mode 100644 index 0000000000000000000000000000000000000000..dfcb4125ab9ada0d34f18e2a9b8169caba222561 --- /dev/null +++ b/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 +#include +#include "deepspeed_aio_thread.h" +#include "deepspeed_pin_tensor.h" + +struct deepspeed_io_handle_t { + std::unique_ptr _aio_ctxt; + const bool _single_submit; + const bool _overlap_events; + const int _intra_op_parallelism; + deepspeed_aio_config_t _aio_config; + + std::vector> _thread_contexts; + std::vector _threads; + int _num_pending_ops; + std::unique_ptr _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 scheduled_op); + + std::shared_ptr _wait_for_aio_work(); + + bool _is_valid_parallel_aio_op(const bool read_op, const int64_t num_bytes); + + virtual std::shared_ptr _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); +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_lib/py_ds_aio.cpp new file mode 100644 index 0000000000000000000000000000000000000000..bf298b691b814fc51be461ee02938f1c04603102 --- /dev/null +++ b/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 +#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_(m, "aio_handle") + .def(py::init(), + "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"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json b/lib/python3.12/site-packages/deepspeed/ops/csrc/aio/py_test/single_process_config.json new file mode 100644 index 0000000000000000000000000000000000000000..275c54135cd83d3d8508ea1f769b823af9529821 --- /dev/null +++ b/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 + ] +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/deepcompile.cpp new file mode 100644 index 0000000000000000000000000000000000000000..2eca0a33262e69cf8901c8293d6c7859b2619149 --- /dev/null +++ b/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 param_registry; +std::unordered_map> executors; +std::shared_ptr reduce_buckets = nullptr; + +c10::intrusive_ptr process_group = nullptr; +c10::intrusive_ptr 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 sizes_to_int_vector(at::IntArrayRef sizes) +{ + std::vector 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 getSymmMemWorkspace(int64_t size) +{ + c10::Device device = c10::Device(c10::kCUDA, c10::cuda::current_device()); + std::vector sizes = {size}; + std::vector 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 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 tensors) {} + +void init(c10::intrusive_ptr 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(reinterpret_cast(&ncclID), + reinterpret_cast(&ncclID) + NCCL_UNIQUE_ID_BYTES); + auto device = torch::Device(torch::kCUDA); + at::Tensor tensor = torch::from_blob(vec.data(), {static_cast(vec.size())}, torch::kUInt8) + .to(torch::Device(torch::kCUDA)); + std::vector 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(); + reduce_buckets = std::make_shared(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 diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/init.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/init.cpp new file mode 100644 index 0000000000000000000000000000000000000000..ca2538b5f2b187639c10843189ceb439cda19ecd --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/init.cpp @@ -0,0 +1,99 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" +#include "z1.h" +#include "z3.h" + +TORCH_LIBRARY(dc, m) +{ + m.def("allgather_param(Tensor a, int graph_id, int id) -> Tensor"); + m.def("prefetch_params_fused(int graph_id, Tensor[] params, int[] ids) -> ()"); + m.def("wait_allgather(Tensor a, int graph_id, int id) -> Tensor"); + m.def("release_param(Tensor a, int graph_id, int id, int n_users) -> Tensor"); + m.def("reduce_grad(Tensor a, int graph_id, int id) -> Tensor"); + m.def("free_tensors(Tensor[] a) -> ()"); + m.def("offload_tensor(Tensor a, int id, int id) -> Tensor"); + m.def("reload_tensor(Tensor a, int id, int id) -> Tensor"); + m.def("wait_offload(Tensor a, int id, int id) -> Tensor"); + m.def("wait_reload(Tensor a, int id, int id) -> Tensor"); + m.def("offload_parameter(Tensor a, int id, int id) -> ()"); + m.def("reload_parameter(Tensor a, int id, int id) -> ()"); + + m.def("test_call(Tensor a) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(dc, CPU, m) +{ + m.impl("allgather_param", &dc::allgather_param); + m.impl("prefetch_params_fused", &dc::prefetch_params_fused); + m.impl("wait_allgather", &dc::wait_allgather); + m.impl("release_param", &dc::release_param); + m.impl("reduce_grad", &dc::reduce_grad); + m.impl("free_tensors", &dc::free_tensors); + m.impl("offload_tensor", &dc::offload_tensor); + m.impl("reload_tensor", &dc::reload_tensor); + m.impl("wait_offload", &dc::wait_offload); + m.impl("wait_reload", &dc::wait_reload); + m.impl("offload_parameter", &dc::offload_parameter); + m.impl("reload_parameter", &dc::reload_parameter); + + m.impl("test_call", &dc::test_call); +} + +TORCH_LIBRARY_IMPL(dc, CUDA, m) +{ + m.impl("allgather_param", &dc::allgather_param); + m.impl("prefetch_params_fused", &dc::prefetch_params_fused); + m.impl("wait_allgather", &dc::wait_allgather); + m.impl("release_param", &dc::release_param); + m.impl("reduce_grad", &dc::reduce_grad); + m.impl("free_tensors", &dc::free_tensors); + m.impl("offload_tensor", &dc::offload_tensor); + m.impl("reload_tensor", &dc::reload_tensor); + m.impl("wait_offload", &dc::wait_offload); + m.impl("wait_reload", &dc::wait_reload); + m.impl("offload_parameter", &dc::offload_parameter); + m.impl("reload_parameter", &dc::reload_parameter); + + m.impl("test_call", &dc::test_call); +} + +TORCH_LIBRARY_IMPL(dc, Meta, m) +{ + m.impl("allgather_param", &dc::allgather_param_meta); + m.impl("prefetch_params_fused", &dc::prefetch_params_fused_meta); + m.impl("release_param", &dc::release_param_meta); + m.impl("wait_allgather", &dc::wait_allgather_meta); + m.impl("reduce_grad", &dc::reduce_grad_meta); + m.impl("free_tensors", &dc::free_tensors_meta); + m.impl("reload_parameter", &dc::reload_parameter_meta); + m.impl("offload_parameter", &dc::offload_parameter_meta); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("set_persistent", &dc::set_persistent, "Set persistent flag for a parameter"); + m.def("enable_profiling", &dc::enable_profiling, "Enable profiling"); + m.def("is_profiling", &dc::is_profiling, "Check if profiling is enabled"); + m.def("init", &dc::init, "Set the process group"); + m.def("cleanup", &dc::cleanup, "Cleanup the process group"); + m.def("register_z1_param", &dc::register_z1_param, "Register a parameter"); + m.def("register_graph_z1", + &dc::register_graph_z1, + "Register graph with a list of ds parameter ids"); + m.def("register_z3_param", &dc::register_z3_param, "Register a parameter"); + m.def("register_graph_z3", + &dc::register_graph_z3, + "Register graph with a list of ds parameter ids"); + m.def("start_forward", &dc::start_forward, "Start forward pass"); + m.def("end_forward", &dc::end_forward, "End forward pass"); + m.def("start_backward", &dc::start_backward, "Start backward pass"); + // m.def("end_backward", &dc::end_backward, "End backward pass"); + m.def("cleanup", &dc::cleanup, "Clean up DeepCompile"); + m.def("reset", &dc::reset, "Reset the state"); + m.def("invalidate_gathered_param", &dc::invalidate_gathered_param, "Invalidate gathered param"); + m.def("clear_all_gathered_params", &dc::clear_all_gathered_params, "Clear all gathered params"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/util.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/util.cpp new file mode 100644 index 0000000000000000000000000000000000000000..948338028059bc1eff42e00d1d7390c09b58fbb0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/util.cpp @@ -0,0 +1,89 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" + +#include + +namespace dc { + +std::string tensorToString(const at::Tensor& t, size_t max_elem, size_t max_str_len) +{ + auto t_cpu = t.flatten() + .slice(0, 0, std::min((int64_t)max_elem, t.numel())) + .to(c10::Device(c10::kCPU), false, true); + + size_t size = std::min(max_elem, productDim(t.sizes())); + + if (t.scalar_type() == c10::ScalarType::Half || t.scalar_type() == c10::ScalarType::BFloat16) { + auto float_ten = t_cpu.to(c10::ScalarType::Float, false, true).contiguous(); + return tensorPtrToString((float*)float_ten.data_ptr(), size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Float) { + return tensorPtrToString((float*)t_cpu.data_ptr(), size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Double) { + return tensorPtrToString((double*)t_cpu.data_ptr(), size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Int) { + int* ptr = static_cast(t_cpu.data_ptr()); + return tensorPtrToString(ptr, size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Long) { + long* ptr = static_cast(t_cpu.data_ptr()); + return tensorPtrToString(ptr, size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Byte) { + unsigned char* ptr = static_cast(t_cpu.data_ptr()); + std::vector vec; + vec.reserve(size); + for (size_t i = 0; i < size; i++) { + vec.push_back(*ptr); + ptr++; + } + return tensorPtrToString(&vec[0], size, max_str_len); + } else if (t.scalar_type() == c10::ScalarType::Bool) { + bool* ptr = static_cast(t_cpu.data_ptr()); + std::vector vec; + vec.reserve(size); + for (size_t i = 0; i < size; i++) { + vec.push_back(*ptr); + ptr++; + } + return tensorPtrToString(&vec[0], size, max_str_len); + } + std::stringstream ss; + ss << "Failed to convert tensor to string. Invalid type of tensor: " + << toString(t.scalar_type()); + throw std::invalid_argument(ss.str()); +} + +std::string tensorPtrToString(void* ptr, + size_t size, + c10::ScalarType datatype, + size_t max_elem, + size_t max_str_len) +{ + int64_t elem_size = std::min((size_t)max_elem, size); + + if (datatype == c10::ScalarType::Long) { + return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); + } else if (datatype == c10::ScalarType::Int) { + return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); + } else if (datatype == c10::ScalarType::Double) { + return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); + } else if (datatype == c10::ScalarType::Float) { + return tensorPtrToString(static_cast(ptr), elem_size, max_str_len); + } else if (datatype == c10::ScalarType::Half || datatype == c10::ScalarType::BFloat16) { + const auto ten = torch::from_blob(ptr, {(int64_t)elem_size}, datatype); + auto float_ten = ten.to(c10::ScalarType::Float, false, true).contiguous(); + return tensorPtrToString((float*)float_ten.data_ptr(), elem_size, max_str_len); + } + std::stringstream ss; + ss << "Failed to convert tensor ptr to string. Invalid type of tensor: " << toString(datatype); + throw std::invalid_argument(ss.str()); +} + +std::string tensorDimToString(const at::Tensor& t) +{ + const auto dim = t.sizes(); + return join_as_str(dim); +} +} // namespace dc diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z1.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z1.cpp new file mode 100644 index 0000000000000000000000000000000000000000..1fc90839862d32b8d0269ab345210af99f6ec5cb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z1.cpp @@ -0,0 +1,141 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "z1.h" +#include "deepcompile.h" + +#define USE_C10D_NCCL + +#include +#include +#include +#include +#include +#include + +#include + +namespace dc { + +class Z1CustomOpExecutor : public CustomOpExecutor { +public: + Z1CustomOpExecutor(c10::intrusive_ptr process_group, + std::shared_ptr param_registry, + std::shared_ptr reduce_buckets, + std::vector ds_ids, + ncclComm_t nccl_comm, + at::cuda::CUDAStream rs_stream, + at::cuda::CUDAStream copy_stream, + bool pre_div_reduce) + : CustomOpExecutor(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + rs_stream, + copy_stream, + pre_div_reduce) + { + } + ~Z1CustomOpExecutor() {} + + void endBackward() override + { + if (param_updated_) { + for (auto& it : has_acc_grad_) { it.second = false; } + } + } + + void flushReduceBucket(at::ScalarType scalar_type) override + { + int rank = process_group_->getRank(); + + if (!hasKey(reduce_tasks_, scalar_type)) { return; } + + int64_t tmp_recv_numel = 0; + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto copy_done_event = rs_copy_done_events_.at(t.getDSId()); + copy_done_event->block(rs_stream_); + } + + ncclGroupStart(); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + ncclRedOp_t op = pre_div_reduce_ ? ncclSum : ncclAvg; + if (pre_div_reduce_) { + at::cuda::CUDAStreamGuard guard(rs_stream_); + t.getSendBuf().div_(process_group_->getSize()); + } + + // inplace + ncclResult_t result = ncclAllReduce(t.getSendBuf().data_ptr(), + t.getSendBuf().data_ptr(), + t.getSendBuf().numel(), + get_nccl_data_type(scalar_type), + op, + nccl_comm_, + rs_stream_); + if (result != ncclSuccess) { throw std::runtime_error("NCCL AllReduce failed"); } + } + ncclGroupEnd(); + + { + at::cuda::CUDAStreamGuard guard(rs_stream_); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + bool acc_grad = has_acc_grad_.at(t.getDSId()); + auto param = param_registry_->getParam(t.getDSId()); + auto grad_buf = param.getGradBuffer().flatten(); + + if (grad_buf.numel() == 0) { continue; } + + int64_t offset = param.getOffset(); + auto recv_buf = t.getSendBuf().flatten().index( + {torch::indexing::Slice(offset, offset + grad_buf.numel())}); + if (acc_grad) { + grad_buf.add_(recv_buf); + } else { + grad_buf.copy_(recv_buf); + } + has_acc_grad_[t.getDSId()] = true; + } + } + + reduce_buckets_->swap(scalar_type, rs_stream_, copy_stream_); + + // Not very sure if this is necessary + // Want to prevent grad tensor from being released before the copy is done + auto comp_stream = at::cuda::getCurrentCUDAStream(); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto copy_done_event = rs_copy_done_events_.at(t.getDSId()); + copy_done_event->block(comp_stream); + } + reduce_tasks_[scalar_type].clear(); + } +}; + +static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true); + +void register_graph_z1(long graph_id, const std::vector& ds_ids) +{ + executors[graph_id] = std::make_shared(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + rs_stream, + copy_stream, + pre_div_reduce); +} + +void register_z1_param(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + int64_t offset) +{ + param_registry->registerParam(ds_id, ds_shape, ds_tensor, grad_buffer, false, offset, false); +} + +} // namespace dc diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z1.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z1.h new file mode 100644 index 0000000000000000000000000000000000000000..a2f100565eba8791b5c3ed5a1e99205ebe8d45c1 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z1.h @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" + +#pragma once + +namespace dc { + +void register_graph_z1(long graph_id, const std::vector& ds_ids); +void register_z1_param(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + int64_t offset); +} // namespace dc diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z3.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z3.cpp new file mode 100644 index 0000000000000000000000000000000000000000..523bcf2c04b4d6b52ddae939dd7af3f1dc87c928 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z3.cpp @@ -0,0 +1,544 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "z3.h" +#include "deepcompile.h" + +#define USE_C10D_NCCL + +#include +#include +#include +#include +#include +#include + +#include + +namespace dc { + +const size_t TIMEOUT_SYMMETRIC_MEMORY_BARRIER = 60000; + +class Z3CustomOpExecutor : public CustomOpExecutor { +public: + Z3CustomOpExecutor(c10::intrusive_ptr process_group, + std::shared_ptr param_registry, + std::shared_ptr reduce_buckets, + std::vector ds_ids, + ncclComm_t nccl_comm, + at::cuda::CUDAStream ag_stream, + at::cuda::CUDAStream rs_stream, + at::cuda::CUDAStream copy_stream, + at::cuda::CUDAStream offload_stream, + at::cuda::CUDAStream reload_stream, + bool pre_div_reduce) + : CustomOpExecutor(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + rs_stream, + copy_stream, + pre_div_reduce), + ag_stream_(ag_stream), + offload_stream_(offload_stream), + reload_stream_(reload_stream) + { + for (long ds_id : ds_ids_) { + ag_comm_done_events_[ds_id] = + std::make_shared(cudaEventDisableTiming); + ag_comp_done_events_[ds_id] = + std::make_shared(cudaEventDisableTiming); + + param_use_count_[ds_id] = 0; + } + } + ~Z3CustomOpExecutor() {} + + void endBackward() override + { + if (param_updated_) { + for (auto& it : has_acc_grad_) { + it.second = false; + param_registry_->setValid(it.first, false); + } + } + + for (auto& it : reload_buffers_) { + it.second.record_stream(at::cuda::getCurrentCUDAStream()); + } + reload_buffers_.clear(); + } + + void launchAllGather(at::Tensor output_buf, + long ds_id, + c10::intrusive_ptr symm_mem) + { + const DSParam& param = param_registry_->getParam(ds_id); + const at::Tensor& ds_tensor = param.getDSTensor(); + + if (symm_mem == nullptr) { + ncclResult_t result = ncclAllGather(ds_tensor.contiguous().data_ptr(), + output_buf.data_ptr(), + ds_tensor.numel(), + get_nccl_data_type(ds_tensor.scalar_type()), + nccl_comm_, + ag_stream_); + + if (result != ncclSuccess) { throw std::runtime_error("NCCL AllGather failed"); } + } else { + at::cuda::CUDAStreamGuard guard(ag_stream_); + int world_size = process_group_->getSize(); + int rank = process_group_->getRank(); + + at::Tensor local_buf = + symm_mem->get_buffer(rank, ds_tensor.sizes(), ds_tensor.scalar_type(), 0); + local_buf.copy_(ds_tensor, true); + + symm_mem->barrier(0, TIMEOUT_SYMMETRIC_MEMORY_BARRIER); + auto chunks = output_buf.flatten().chunk(world_size); + for (int step = 0; step < world_size; step++) { + int remote_rank = (rank - step + world_size) % world_size; + auto src_buf = symm_mem->get_buffer( + remote_rank, ds_tensor.sizes(), ds_tensor.scalar_type(), 0); + chunks[remote_rank].copy_(src_buf.flatten(), true); + } + symm_mem->barrier(0, TIMEOUT_SYMMETRIC_MEMORY_BARRIER); + } + + param_registry_->registerGatheredParam(ds_id, output_buf); + param_registry_->setValid(ds_id, true); + } + + at::Tensor allgatherParam(long ds_id, + c10::intrusive_ptr symm_mem) + { + if (param_registry_->isValid(ds_id)) { return param_registry_->getGatheredParam(ds_id); } + + const DSParam& param = param_registry_->getParam(ds_id); + const at::Tensor& ds_tensor = param.getDSTensor(); + at::Tensor output_buf = param_registry_->hasGatheredParam(ds_id) + ? param_registry_->getGatheredParam(ds_id) + : torch::empty(param.getShape(), ds_tensor.options()); + + assert(hasKey(ag_comp_done_events_, ds_id)); + ag_comp_done_events_[ds_id]->record(); + ag_comp_done_events_[ds_id]->block(ag_stream_); + + launchAllGather(output_buf, ds_id, symm_mem); + + ag_comm_done_events_[ds_id]->record(ag_stream_); + return output_buf; + } + + void prefetchParamsFused(std::vector ds_ids, + c10::intrusive_ptr symm_mem) + { + std::vector invalid_ds_ids; + for (const auto& ds_id : ds_ids) { + if (!param_registry_->isValid(ds_id)) { invalid_ds_ids.push_back(ds_id); } + } + + std::unordered_map output_bufs; + for (long ds_id : invalid_ds_ids) { + const DSParam& param = param_registry_->getParam(ds_id); + if (param_registry_->hasGatheredParam(ds_id)) { + output_bufs[ds_id] = param_registry_->getGatheredParam(ds_id); + } else { + output_bufs[ds_id] = torch::empty(param.getShape(), param.getDSTensor().options()); + } + } + + for (long ds_id : invalid_ds_ids) { + ag_comp_done_events_[ds_id]->record(); + ag_comp_done_events_[ds_id]->block(ag_stream_); + } + + ncclGroupStart(); + for (long ds_id : invalid_ds_ids) { + assert(hasKey(output_bufs, ds_id)); + launchAllGather(output_bufs.at(ds_id), ds_id, symm_mem); + } + ncclGroupEnd(); + + for (long ds_id : invalid_ds_ids) { ag_comm_done_events_[ds_id]->record(ag_stream_); } + } + + void releaseParam(long ds_id, long n_users) + { + const DSParam& param = param_registry_->getParam(ds_id); + + assert(hasKey(param_use_count_, ds_id)); + if (param_use_count_[ds_id] == 0) { param_use_count_[ds_id] = n_users; } + param_use_count_[ds_id]--; + + if (param_use_count_[ds_id] == 0 && !param.isPersistent()) { + at::Tensor gathered_param = param_registry_->getGatheredParam(ds_id); + + if (gathered_param.defined()) { // gathered param is undefined while profiling + const auto options = gathered_param.options(); + at::Tensor empty_buffer = torch::empty({0}, options); + gathered_param.set_data(empty_buffer); + } + + param_registry_->unregisterGatheredParam(ds_id); + } + } + + at::Tensor waitAllgather(at::Tensor v, long ds_id) + { + assert(hasKey(ag_comm_done_events_, ds_id)); + ag_comm_done_events_[ds_id]->block(at::cuda::getCurrentCUDAStream()); + return v; + } + + void flushReduceBucket(at::ScalarType scalar_type) override + { + if (!hasKey(reduce_tasks_, scalar_type)) { return; } + + int64_t tmp_recv_numel = 0; + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto copy_done_event = rs_copy_done_events_.at(t.getDSId()); + copy_done_event->block(rs_stream_); + + if (has_acc_grad_.at(t.getDSId())) { + tmp_recv_numel += param_registry_->getParam(t.getDSId()).getGradBuffer().numel(); + } + } + + at::Tensor tmp_recv_buf = at::Tensor(); + if (tmp_recv_numel > 0) { + at::cuda::CUDAStreamGuard guard(rs_stream_); + tmp_recv_buf = torch::empty({tmp_recv_numel}, + at::TensorOptions().dtype(scalar_type).device(at::kCUDA)); + } + + ncclGroupStart(); + int64_t offset = 0; + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto recv_buf = param_registry_->getParam(t.getDSId()).getGradBuffer(); + + bool acc_grad = has_acc_grad_.at(t.getDSId()); + + if (acc_grad) { + recv_buf = + tmp_recv_buf.index({torch::indexing::Slice(offset, offset + recv_buf.numel())}); + } + + ncclRedOp_t op = pre_div_reduce_ ? ncclSum : ncclAvg; + if (pre_div_reduce_) { + at::cuda::CUDAStreamGuard guard(rs_stream_); + t.getSendBuf().div_(process_group_->getSize()); + } + ncclResult_t result = ncclReduceScatter(t.getSendBuf().data_ptr(), + recv_buf.data_ptr(), + recv_buf.numel(), + get_nccl_data_type(scalar_type), + op, + nccl_comm_, + rs_stream_); + if (result != ncclSuccess) { throw std::runtime_error("NCCL ReduceScatter failed"); } + + if (acc_grad) { offset += recv_buf.numel(); } + } + ncclGroupEnd(); + + { + at::cuda::CUDAStreamGuard guard(rs_stream_); + int64_t offset = 0; + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + bool acc_grad = has_acc_grad_.at(t.getDSId()); + + if (acc_grad) { + auto recv_buf = param_registry_->getParam(t.getDSId()).getGradBuffer(); + recv_buf.add_(tmp_recv_buf.index( + {torch::indexing::Slice(offset, offset + recv_buf.numel())})); + offset += recv_buf.numel(); + } + has_acc_grad_[t.getDSId()] = true; + } + } + + reduce_buckets_->swap(scalar_type, rs_stream_, copy_stream_); + + // Not very sure if this is necessary + // Want to prevent grad tensor from being released before the copy is done + auto comp_stream = at::cuda::getCurrentCUDAStream(); + for (const ReduceTask& t : reduce_tasks_.at(scalar_type)) { + auto copy_done_event = rs_copy_done_events_.at(t.getDSId()); + copy_done_event->block(comp_stream); + } + reduce_tasks_[scalar_type].clear(); + + if (tmp_recv_numel > 0) { tmp_recv_buf.record_stream(rs_stream_); } + } + + at::Tensor offloadTensor(at::Tensor tensor, long id) + { + if (!hasKey(offload_events_, id)) { + offload_events_[id] = std::make_shared(cudaEventDisableTiming); + offload_comp_done_events_[id] = + std::make_shared(cudaEventDisableTiming); + + const auto options = at::TensorOptions().pinned_memory(true).device(torch::kCPU); + offload_buffers_[id] = at::empty_like(tensor, options); + } + + offload_comp_done_events_[id]->record(); + offload_comp_done_events_[id]->block(offload_stream_); + { + at::cuda::CUDAStreamGuard guard(offload_stream_); + offload_buffers_.at(id).copy_(tensor, true); + } + + tensor.record_stream(offload_stream_); + + offload_events_[id]->record(offload_stream_); + assert(hasKey(offload_buffers_, id)); + return offload_buffers_.at(id); + } + + at::Tensor reloadTensor(at::Tensor tensor, long id) + { + if (!hasKey(reload_events_, id)) { + reload_events_[id] = std::make_shared(cudaEventDisableTiming); + } + + assert(hasKey(offload_buffers_, id)); + offload_events_[id]->block(reload_stream_); + + at::Tensor ten; + { + at::cuda::CUDAStreamGuard guard(reload_stream_); + + assert(hasKey(offload_buffers_, id)); + at::Tensor buf = offload_buffers_.at(id); + const auto options = at::TensorOptions().device(torch::kCUDA); + ten = at::empty_like(buf, options); + ten.copy_(buf, true); + + reload_buffers_[id] = ten; + } + + reload_events_[id]->record(reload_stream_); + return ten; + } + + at::Tensor waitOffload(at::Tensor tensor, long id) + { + assert(hasKey(offload_events_, id)); + offload_events_[id]->block(at::cuda::getCurrentCUDAStream()); + + assert(hasKey(offload_buffers_, id)); + return offload_buffers_.at(id); + } + + at::Tensor waitReload(at::Tensor tensor, long id) + { + assert(hasKey(reload_events_, id)); + reload_events_[id]->block(at::cuda::getCurrentCUDAStream()); + + assert(hasKey(reload_buffers_, id)); + auto ten = reload_buffers_.at(id); + + // We can't release here because the tensor is still being used + // We will need "freeReloadedTensor" after the last user of the tensor to call + // ".record_stream". As it is a bit complicated, we clear the buffer and do at the end of + // the backward pass for now. reload_buffers_.erase(id); + return ten; + } + + void offloadParameter(at::Tensor tensor, long ds_id) { param_registry_->offload(ds_id); } + void reloadParameter(at::Tensor tensor, long ds_id) { param_registry_->reload(ds_id); } + + bool hasReloadBuffer(long id) { return hasKey(reload_buffers_, id); } + + bool hasParam(long ds_id) const { return hasKey(has_acc_grad_, ds_id); } + +private: + at::cuda::CUDAStream ag_stream_; + at::cuda::CUDAStream offload_stream_; + at::cuda::CUDAStream reload_stream_; + + std::unordered_map> ag_comp_done_events_; + std::unordered_map> ag_comm_done_events_; + + std::unordered_map> offload_events_; + std::unordered_map> offload_comp_done_events_; + std::unordered_map> reload_events_; + std::unordered_map offload_buffers_; + std::unordered_map reload_buffers_; + + std::unordered_map param_use_count_; +}; + +static at::cuda::CUDAStream ag_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream rs_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream copy_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream offload_stream = at::cuda::getStreamFromPool(true); +static at::cuda::CUDAStream reload_stream = at::cuda::getStreamFromPool(true); + +void register_graph_z3(long graph_id, const std::vector& ds_ids) +{ + executors[graph_id] = std::make_shared(process_group, + param_registry, + reduce_buckets, + ds_ids, + nccl_comm, + ag_stream, + rs_stream, + copy_stream, + offload_stream, + reload_stream, + pre_div_reduce); +} + +void register_z3_param(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + bool persistent) +{ + param_registry->registerParam(ds_id, ds_shape, ds_tensor, grad_buffer, true, 0, persistent); + if (persistent) { param_registry->registerGatheredParam(ds_id, ds_tensor); } +} + +at::Tensor allgather_param(at::Tensor param_tensor, long graph_id, long ds_id) +{ + auto executor = getExecutor(graph_id, executors); + + if (sync_before_allgather) { c10::cuda::device_synchronize(); } + auto ret = executor->allgatherParam(ds_id, symm_mem); + if (sync_after_allgather) { c10::cuda::device_synchronize(); } + return ret; +} + +void set_persistent(long ds_id) +{ + param_registry->setPersistent(ds_id, true); + + // Allocate buffer here + // Memory fragmentation will be more severe if we allocate in forward/backward + for (auto& it : executors) { + if (it.second->hasParam(ds_id)) { + auto executor = getExecutor(it.first, executors); + executor->allgatherParam(ds_id, symm_mem); + } + } +} + +void prefetch_params_fused(long graph_id, + const std::vector params, + const std::vector& ds_ids) +{ + auto executor = getExecutor(graph_id, executors); + executor->prefetchParamsFused(ds_ids, symm_mem); +} + +void prefetch_params_fused_meta(long graph_id, + const std::vector params, + const std::vector& ds_ids) +{ +} + +// for profiling +void invalidate_gathered_param(long ds_id) +{ + const DSParam& param = param_registry->getParam(ds_id); + if (param.isPersistent()) { return; } + + param_registry->unregisterGatheredParam(ds_id); + param_registry->registerGatheredParam(ds_id, at::Tensor()); +} + +void clear_all_gathered_params() +{ + for (const auto& it : param_registry->getParams()) { + long ds_id = it.first; + const DSParam& param = param_registry->getParam(ds_id); + if (param.isPersistent()) { continue; } + if (param_registry->hasGatheredParam(ds_id)) { + param_registry->unregisterGatheredParam(ds_id); + } + } +} + +at::Tensor allgather_param_meta(at::Tensor param_tensor, long graph_id, long ds_id) +{ + const DSParam& param = param_registry->getParam(ds_id); + auto options = param.getDSTensor().options().device(c10::kMeta); + at::Tensor output_buf = torch::empty(param.getShape(), options); + return output_buf; +} + +at::Tensor release_param(at::Tensor dummy, long graph_id, long ds_id, long n_users) +{ + auto executor = getExecutor(graph_id, executors); + executor->releaseParam(ds_id, n_users); + + if (clone_custom_op_output) { return dummy.clone(); } + return dummy; +} + +at::Tensor release_param_meta(at::Tensor dummy, long graph_id, long ds_id, long n_users) +{ + return dummy; +} + +at::Tensor wait_allgather(at::Tensor v, long graph_id, long ds_id) +{ + auto executor = getExecutor(graph_id, executors); + executor->waitAllgather(v, ds_id); + return v; +} + +at::Tensor wait_allgather_meta(at::Tensor v, long graph_id, long ds_id) { return v; } + +at::Tensor offload_tensor(at::Tensor tensor, long graph_id, long id) +{ + auto executor = getExecutor(graph_id, executors); + return executor->offloadTensor(tensor, id); +} + +at::Tensor reload_tensor(at::Tensor tensor, long graph_id, long id) +{ + auto executor = getExecutor(graph_id, executors); + return executor->reloadTensor(tensor, id); +} + +at::Tensor wait_offload(at::Tensor tensor, long graph_id, long id) +{ + auto executor = getExecutor(graph_id, executors); + return executor->waitOffload(tensor, id); +} + +at::Tensor wait_reload(at::Tensor tensor, long graph_id, long id) +{ + auto executor = getExecutor(graph_id, executors); + if (profile && !executor->hasReloadBuffer(id)) { return tensor; } + return executor->waitReload(tensor, id); +} + +at::Tensor test_call(at::Tensor a) +{ + std::cout << "test_call" << std::endl; + return a; +} + +void reload_parameter(at::Tensor tensor, long graph_id, long ds_id) +{ + auto executor = getExecutor(graph_id, executors); + executor->reloadParameter(tensor, ds_id); +} + +void offload_parameter(at::Tensor tensor, long graph_id, long ds_id) +{ + auto executor = getExecutor(graph_id, executors); + executor->offloadParameter(tensor, ds_id); +} +void reload_parameter_meta(at::Tensor param_tensor, long graph_id, long ds_id) {} +void offload_parameter_meta(at::Tensor tensor, long graph_id, long ds_id) {} + +} // namespace dc diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z3.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z3.h new file mode 100644 index 0000000000000000000000000000000000000000..1031f0c84f7c1447d662c1df4629f50c7c47fce1 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/compile/z3.h @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "deepcompile.h" + +#pragma once + +namespace dc { + +void register_graph_z3(long graph_id, const std::vector& ds_ids); +void register_graph_ops_z3(long graph_id, + const std::vector& op_names, + const std::vector& n_args); +void register_bwd_graph_ops_z3(long graph_id, + const std::vector& op_names, + const std::vector& n_args); +void register_z3_param(long ds_id, + const std::vector& ds_shape, + at::Tensor ds_tensor, + at::Tensor grad_buffer, + bool persistent); +at::Tensor allgather_param(at::Tensor param_tensor, long graph_id, long ds_id); +void set_persistent(long ds_id); +void prefetch_params_fused(long graph_id, + const std::vector params, + const std::vector& ds_ids); +void prefetch_params_fused_meta(long graph_id, + const std::vector params, + const std::vector& ds_ids); +// for profiling +void invalidate_gathered_param(long ds_id); +void clear_all_gathered_params(); +at::Tensor allgather_param_meta(at::Tensor param_tensor, long graph_id, long ds_id); +at::Tensor release_param(at::Tensor dummy, long graph_id, long ds_id, long n_users); +at::Tensor release_param_meta(at::Tensor dummy, long graph_id, long ds_id, long n_users); +at::Tensor wait_allgather(at::Tensor v, long graph_id, const long ds_id); +at::Tensor wait_allgather_meta(at::Tensor v, long graph_id, long ds_id); +at::Tensor offload_tensor(at::Tensor tensor, long graph_id, long id); +at::Tensor reload_tensor(at::Tensor tensor, long graph_id, long id); +at::Tensor wait_offload(at::Tensor tensor, long graph_id, long id); +at::Tensor wait_reload(at::Tensor tensor, long graph_id, long id); +void reload_parameter(at::Tensor tensor, long graph_id, long id); +void offload_parameter(at::Tensor tensor, long graph_id, long id); +void reload_parameter_meta(at::Tensor tensor, long graph_id, long id); +void offload_parameter_meta(at::Tensor tensor, long graph_id, long id); +} // namespace dc diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d25578f410da278dace61fb5b488b14ed1257c9f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/adam/fused_adam.cpp @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_adam.h" + +// C++ interface + +void multi_tensor_adam(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, /*gpmv*/ + 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) +{ + static bool initialized = false; + if (!initialized) { + create_adam_optimizer(0); + initialized = true; + } + for (int i = 0; i < tensor_lists[0].size(); i++) { + ds_adam_step(0, + step, + lr, + beta1, + beta2, + epsilon, + weight_decay, + bias_correction, + tensor_lists[1][i], + tensor_lists[0][i], + tensor_lists[2][i], + tensor_lists[3][i]); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("multi_tensor_adam", + &multi_tensor_adam, + "Compute and apply gradient update to parameters for Adam optimizer"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d28509e592668aca100555e18ae47e3aadacb43a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/ccl.cpp @@ -0,0 +1,344 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +#include +#include "shm.h" + +// #define DO_PROFILE +#ifdef DO_PROFILE +#include +#include +#endif + +// Communication settings +static int world_rank = -1; +static int world_size = -1; + +static std::set _comm_ids; +static std::set _colors; +static std::vector _ccl_comms; +static ccl::shared_ptr_class sub_kvs; +static std::map, int> group_to_comm_id; + +ccl::communicator& _get_comm_from_group() { return _ccl_comms[0]; } +ccl::communicator& _get_comm_from_group(py::object group) { return _ccl_comms[0]; } +ccl::communicator& _get_comm_from_group(std::vector ranks) +{ + if (group_to_comm_id.find(ranks) != group_to_comm_id.end()) { + auto id = group_to_comm_id.find(ranks); + return _ccl_comms[id->second]; + } + return _ccl_comms[0]; +} + +#define CCLCHECK(cmd) \ + do { \ + cmd; \ + } while (0) + +#define KVS_CREATE_SUCCESS 0 +#define KVS_CREATE_FAILURE -1 + +static bool is_initialized = 0; + +static ccl::shared_ptr_class kvs; + +static bool all_ranks_local_p = false; + +void initialize(int size, int rank, torch::Tensor& kvs_data) +{ + if (is_initialized) return; + + // Check whether all ranks is on the same physical machine. + // If true, we will use an SHM based low latency allreduce + + auto ls_string = std::getenv("LOCAL_SIZE"); + int ls = 0; + if (ls_string != NULL) { ls = std::stoi(std::getenv("LOCAL_SIZE")); } + + if (size >= 1 && size == ls) { all_ranks_local_p = true; } + + world_size = size; + world_rank = rank; + is_initialized = 1; + + ccl::kvs::address_type main_addr; + + if (rank != 0) { + memcpy(main_addr.data(), kvs_data.data_ptr(), main_addr.size()); + kvs = ccl::create_kvs(main_addr); + } + + _ccl_comms.emplace_back(ccl::create_communicator(size, rank, kvs)); + + auto addr_string = std::getenv("MASTER_ADDR"); + if (addr_string == NULL) { addr_string = ""; } + auto port_string = std::getenv("MASTER_PORT"); + if (port_string == NULL) { port_string = ""; } + + if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } +} + +/* + rank == 0: create main kvs and return its address + rank == else: return an empty address +*/ +std::vector get_kvs_addr(int rank) +{ + if (rank == 0) { + kvs = ccl::create_main_kvs(); + ccl::kvs::address_type main_addr = kvs->get_address(); + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } else { + ccl::kvs::address_type main_addr; + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } +} + +int get_rank(int group = 0) { return world_rank; } + +int get_world_size(int group = 0) { return world_size; } + +// Find the next ordered, unique value to a set. E.g. <0,1,2,7> --> 3 +int next_unique_val(std::set s) +{ + std::set::iterator itr; + // Base case. Add 0 to start of set. + if (s.empty() || *s.begin() != 0) { + return 0; + // second base case where s = {0} (the case of s = {n != 0} is caught above) + } else if (s.size() == 1) { + return 1; + } else { + int prev_val = *s.begin(); + for (itr = std::next(s.begin()); itr != s.end(); itr++) { + if (*itr != prev_val + 1) { return prev_val + 1; } + prev_val = *itr; + } + return *(s.end()) + 1; + } +} + +std::vector get_sub_kvs_addr(bool first) +{ + if (first) { + sub_kvs = ccl::create_main_kvs(); + ccl::kvs::address_type main_addr = sub_kvs->get_address(); + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } else { + ccl::kvs::address_type main_addr; + auto ccl_kvs_addr = std::vector(main_addr.begin(), main_addr.end()); + return ccl_kvs_addr; + } +} + +void initialize_sub_comm(int size, int rank, torch::Tensor& kvs_data, std::vector ranks) +{ + ccl::kvs::address_type main_addr; + if (rank != 0) { + memcpy(main_addr.data(), kvs_data.data_ptr(), main_addr.size()); + sub_kvs = ccl::create_kvs(main_addr); + } + _ccl_comms.push_back(ccl::create_communicator(size, rank, sub_kvs)); + group_to_comm_id[ranks] = _ccl_comms.size() - 1; +} + +ccl::datatype get_ccl_datatype(c10::ScalarType type) +{ + ccl::datatype ccl_type; + switch (type) { + case c10::ScalarType::Int: ccl_type = ccl::datatype::int32; break; + case c10::ScalarType::Long: ccl_type = ccl::datatype::int64; break; + case c10::ScalarType::Float: ccl_type = ccl::datatype::float32; break; + case c10::ScalarType::Double: ccl_type = ccl::datatype::float64; break; + case c10::ScalarType::BFloat16: ccl_type = ccl::datatype::bfloat16; break; + case c10::ScalarType::Half: ccl_type = ccl::datatype::float16; break; + default: ccl_type = ccl::datatype::int8; + } + return ccl_type; +} + +ccl::reduction get_ccl_reduce_op(py::object op, at::Tensor& input) +{ + py::object ReduceOp = py::module_::import("deepspeed.comm").attr("ReduceOp"); + if (!py::isinstance(op, ReduceOp)) { + throw std::runtime_error("Error: Op must be of type ReduceOp"); + } + + int op_val = py::int_(op.attr("value")); + ccl::reduction ccl_op; + + if (input.scalar_type() == at::kBool) { + if (op_val == (int)py::int_(ReduceOp.attr("SUM").attr("value"))) { + // For bool tensors, map sum to max, which both represent a bitwise or. + // This is to prevent overflow issues with sum, since we use uint8 to + // represent a bool (see cclDataType mapping). + ccl_op = ccl::reduction::max; + } else if (op_val == (int)py::int_(ReduceOp.attr("AVG").attr("value"))) { + throw std::runtime_error("Error: For bool tensors, op must be of type ReduceOp"); + } + } + + if (op_val == (int)py::int_(ReduceOp.attr("SUM").attr("value"))) { + ccl_op = ccl::reduction::sum; + } else if (op_val == (int)py::int_(ReduceOp.attr("MIN").attr("value"))) { + ccl_op = ccl::reduction::min; + } else if (op_val == (int)py::int_(ReduceOp.attr("MAX").attr("value"))) { + ccl_op = ccl::reduction::max; + } else if (op_val == (int)py::int_(ReduceOp.attr("PRODUCT").attr("value"))) { + ccl_op = ccl::reduction::prod; + } else { + throw std::runtime_error("Error: Unrecognized ReduceOp type"); + } + return ccl_op; +} + +void broadcast(torch::Tensor& data, int src, std::vector group, bool async_op) +{ + CCLCHECK(ccl::broadcast(data.data_ptr(), + data.numel(), + get_ccl_datatype(data.scalar_type()), + src, + _get_comm_from_group(group)) + .wait()); +} + +// TODO: implement torch's async_op behavior, document it. +void all_reduce(torch::Tensor& data, py::object op, std::vector group, bool async_op) +{ + CCLCHECK(ccl::allreduce(data.data_ptr(), + data.data_ptr(), + data.numel(), + get_ccl_datatype(data.scalar_type()), + get_ccl_reduce_op(op, data), + _get_comm_from_group(group)) + .wait()); +} + +void all_reduce_caching(torch::Tensor& data, + py::object op, + std::string match_id, + std::vector group, + bool async_op) +{ + ccl::allreduce_attr attr = ccl::default_allreduce_attr; + auto match_str = ccl::v1::string(match_id); + attr.template set(true); + attr.template set(match_str); + // To control this, use operation attribute and set true value for to_cache field and unique + // string (for example, tensor name) for match_id field. Note that: + // match_id should be the same for a specific communication operation across all ranks. + // If the same tensor is a part of different communication operations, match_id should have + // different values for each of these operations. + CCLCHECK(ccl::allreduce(data.data_ptr(), + data.data_ptr(), + data.numel(), + get_ccl_datatype(data.scalar_type()), + get_ccl_reduce_op(op, data), + _get_comm_from_group(group), + attr) + .wait()); +} + +void inference_all_reduce(torch::Tensor& data, py::object op) +{ +#ifdef DO_PROFILE + static double total_time = 0.0; + static double total_time_sq = 0.0; + static int count = -16; // warmup + static double max_time = 0.0; + static double min_time = DBL_MAX; + // make sure all rank reach this point before measuring time + // turn on this if you suspect each rank didn't reach here at the same time (stragger) + // if (all_ranks_local_p) { + // barrier_wait(0, world_size); + //} + auto start = std::chrono::system_clock::now(); +#endif + + static py::object ReduceOp = py::module_::import("deepspeed.comm").attr("ReduceOp"); + static auto ReduceOpSum = (int)py::int_(ReduceOp.attr("SUM").attr("value")); + + assert(py::int_(op.attr("value")) == ReduceOpSum); + + auto numel = data.numel(); + + int data_size = 0; + bool data_type_fallback = false; + + switch (data.scalar_type()) { + case c10::ScalarType::BFloat16: data_size = numel * 2; break; + case c10::ScalarType::Float: data_size = numel * 4; break; + default: data_type_fallback = true; + } + + if (data_type_fallback || !all_ranks_local_p) { + // fallback to oneccl allreduce + CCLCHECK(ccl::allreduce(data.data_ptr(), + data.data_ptr(), + data.numel(), + get_ccl_datatype(data.scalar_type()), + get_ccl_reduce_op(op, data), + _get_comm_from_group()) + .wait()); + } else { + all_reduce_outer_loop(data, numel, data_size); + } + +#ifdef DO_PROFILE + auto end = std::chrono::system_clock::now(); + count++; + if (count > 0) { + double elapsed = std::chrono::duration_cast(end - start).count(); + if (elapsed > max_time) { max_time = elapsed; } + if (elapsed < min_time) { min_time = elapsed; } + total_time += elapsed; + total_time_sq += elapsed * elapsed; + if (world_rank == 0 && count == 1000) { + auto avg = total_time / count; + auto sd = + sqrt(total_time_sq / count - total_time * total_time / (count * count)) / avg * 100; + printf(" C++ kernel\t\t %.2f\t %.2f\t%.2f\t %.2f\n", + min_time, + max_time, + total_time / count, + sd); + } + } +#endif +} + +void barrier(std::vector group, bool async_op) +{ + CCLCHECK(ccl::barrier(_get_comm_from_group(group)).wait()); +} + +std::vector get_available_coll() +{ + std::vector colls{ + "broadcast", "all_reduce", "inference_all_reduce", "all_reduce_caching", "barrier"}; + return colls; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("get_kvs_addr", &get_kvs_addr, "create and get main kvs addr"); + m.def("initialize", &initialize, "ccl initialize"); + m.def("get_rank", &get_rank, "get rank"); + m.def("get_world_size", &get_world_size, "get world size"); + m.def("broadcast", &broadcast, "ccl broadcast"); + m.def("all_reduce", &all_reduce, "ccl all_reduce"); + m.def("inference_all_reduce", &inference_all_reduce, "low latency all_reduce implementation"); + m.def("all_reduce_caching", &all_reduce_caching, "ccl all_reduce with caching"); + m.def("barrier", &barrier, "barrier"); + m.def("initialize_sub_comm", &initialize_sub_comm, "initialize_sub_comm"); + m.def("get_sub_kvs_addr", &get_sub_kvs_addr, "get_sub_kvs_addr"); + m.def("get_available_coll", &get_available_coll, "get_available_coll"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp new file mode 100644 index 0000000000000000000000000000000000000000..be44681ca0626a9dcef9ccf5e7ea09b38bf27d6c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp @@ -0,0 +1,692 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +#include +#include +#include +#include +#include +#include "shm.h" + +// #define DO_PROFILE +#ifdef DO_PROFILE +#include +#include +#endif + +// states for collectives +enum coll_state { + coll_begin = 0, + coll_allreduce_naive__copy_in_done, + coll_allreduce_naive__reduce_done, + // alternative state when allreduce is working on alternative buffer + // of the double buffer. + coll_alt1_allreduce_naive__copy_in_done, + coll_alt2_allreduce_naive__copy_in_done, + coll_alt1_allreduce_naive__reduce_done, +}; + +// SHM building blocks +struct SharedData { + const char* name; + int descriptor; + void* bytes; + size_t nbytes; +}; + +void shared_open(SharedData* data, const char* name, size_t nbytes) +{ + int d = shm_open(name, O_RDWR, S_IRUSR | S_IWUSR); + if (d != -1) { + void* bytes = mmap(NULL, nbytes, PROT_READ | PROT_WRITE, MAP_SHARED, d, 0); + data->name = name; + data->descriptor = d; + data->bytes = bytes; + data->nbytes = nbytes; + } else { + if (errno != ENOENT) { + // don't print if shm can not be found because we want to loop over from + // caller again until the other ranks created the shm + printf("shared_open %s failed, errno=%d\n", name, errno); + } + data->descriptor = -1; + } +} + +void shared_create(SharedData* data, const char* name, void* bytes, size_t nbytes) +{ + int d = shm_open(name, O_CREAT | O_RDWR, S_IRUSR | S_IWUSR); + if (d != -1) { + if (nbytes = write(d, bytes, nbytes)) { shared_open(data, name, nbytes); } + } else { + printf("shared_create %s failed\n", name); + } +} + +void shared_close(SharedData* data) +{ + if (data->descriptor != -1) { + munmap(data->bytes, data->nbytes); + shm_unlink(data->name); + } +} + +static int world_size; + +// SHM based allreduce helper functions +// buffer that holds shm name +#define NAME_BUF_SIZE 1000 +#define MAX_BUF_SIZE 1048576 * 32 +#define NAIVE_ALLREDUCE_THRESHOLD 1048576 +#define SHM_BUFFER_NAME "deepspeed_allreduce_buffer" +struct allreduce_workspace { + enum coll_state states[2]; // idx=0 -- state for symmetric_naive_all_reduce + // idx=1 -- state for distributed_naive_all_reduce + // double buffer to avoid syncing between rounds + // offset=0 -- 2*NAIVE_ALLREDUCE_THRESHOLD : buffer for symmetric_naive_all_reduce + // after that : buffer for distributed_naive_all_reduce + char buffer[2 * NAIVE_ALLREDUCE_THRESHOLD + 2 * MAX_BUF_SIZE]; +}; + +#define BUFFER0_OFFSET(current_buffer) current_buffer* NAIVE_ALLREDUCE_THRESHOLD +#define BUFFER1_OFFSET(current_buffer) 2 * NAIVE_ALLREDUCE_THRESHOLD + current_buffer* MAX_BUF_SIZE + +struct allreduce_workspace** workspace; + +// buffer for small messages, double buffer +char** symmetric_buffer[2]; +// buffer for large messages, double buffer +char** distributed_buffer[2]; + +void wait_buffer_state_until_2(int index, + enum coll_state state0, + enum coll_state state1, + int state_group) +{ + volatile enum coll_state* state_ptr = &(workspace[index]->states[state_group]); + + while (1) { + volatile enum coll_state cur_state = *state_ptr; + if (cur_state == state0 || cur_state == state1) break; + } +} + +__m512 cvt_bf16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); +inline __m512 cvt_bf16_to_fp32(const __m256i src) +{ + auto y = _mm512_cvtepu16_epi32(src); + return _mm512_castsi512_ps(_mm512_bslli_epi128(y, 2)); +} + +inline __m256i cvt_fp32_to_bf16(const __m512 src) __attribute__((target("avx512bw"))); +inline __m256i cvt_fp32_to_bf16(const __m512 src) +{ + __m512i value = _mm512_castps_si512(src); + __m512i nan = _mm512_set1_epi32(0xffff); + auto mask_value = _mm512_cmp_ps_mask(src, src, _CMP_ORD_Q); + __m512i ones = _mm512_set1_epi32(0x1); + __m512i vec_bias = _mm512_set1_epi32(0x7fff); + // uint32_t lsb = (input >> 16) & 1; + auto t_value = _mm512_and_si512(_mm512_srli_epi32(value, 16), ones); + // uint32_t rounding_bias = 0x7fff + lsb; + t_value = _mm512_add_epi32(t_value, vec_bias); + // input += rounding_bias; + t_value = _mm512_add_epi32(t_value, value); + // input = input >> 16; + t_value = _mm512_srli_epi32(t_value, 16); + // Check NaN before converting back to bf16 + t_value = _mm512_mask_blend_epi32(mask_value, nan, t_value); + return _mm512_cvtusepi32_epi16(t_value); +} + +__m512 cvt_fp16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); +inline __m512 cvt_fp16_to_fp32(const __m256i src) { return _mm512_cvtph_ps(src); } + +inline __m256i cvt_fp32_to_fp16(const __m512 src) __attribute__((target("avx512bw"))); +inline __m256i cvt_fp32_to_fp16(const __m512 src) +{ + return _mm512_cvtps_ph(src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC)); +} + +void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) + __attribute__((target("avx512bw"))); + +void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) + __attribute__((target("avx512bw"))); + +void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) + __attribute__((target("avx512bw"))); + +void reduce_all_buffers(int start_elements, + int num_elements, + c10::ScalarType scalar_type, + int to_buffer_idx, + char* to_buffer, + char** buffers) +{ + switch (scalar_type) { + case c10::ScalarType::BFloat16: + reduce_bf16_buffers(start_elements, num_elements, to_buffer, buffers); + break; + case c10::ScalarType::Half: + reduce_fp16_buffers(start_elements, num_elements, to_buffer, buffers); + break; + case c10::ScalarType::Float: + reduce_fp32_buffers(start_elements, num_elements, to_buffer, buffers); + break; + default: assert(!"Should not get here"); + } +} + +#define CVT_ADD_BF16(x) \ + do { \ + auto in##x##_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ + inout_val = _mm512_add_ps(inout_val, in##x##_val); \ + } while (0) + +// Reduce functions down below use vectorized algorithm, the number of bytes processed each +// iteration depends on vector length. 256bit vector ==> 32 bytes, 512bit vector ==> 64 bytes +// If you change implementation of reduce_bf16_buffers, etc. , check whether this number needs +// to be changed +#define VECTOR_LENGTH_IN_BYTES 32 + +void reduce_bf16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) +{ + const int element_size = 2; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); + switch (world_size) { + case 16: CVT_ADD_BF16(15); + case 15: CVT_ADD_BF16(14); + case 14: CVT_ADD_BF16(13); + case 13: CVT_ADD_BF16(12); + case 12: CVT_ADD_BF16(11); + case 11: CVT_ADD_BF16(10); + case 10: CVT_ADD_BF16(9); + case 9: CVT_ADD_BF16(8); + case 8: CVT_ADD_BF16(7); + case 7: CVT_ADD_BF16(6); + case 6: CVT_ADD_BF16(5); + case 5: CVT_ADD_BF16(4); + case 4: CVT_ADD_BF16(3); + case 3: CVT_ADD_BF16(2); + case 2: CVT_ADD_BF16(1); + case 1: break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = cvt_bf16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); + inout_val = _mm512_add_ps(inout_val, in_val); + } + } + _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_bf16(inout_val)); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { val += *(at::BFloat16*)(buffers[j] + i); } + *(at::BFloat16*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +#define CVT_ADD_FP16(x) \ + do { \ + auto in##x##_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[x] + i))); \ + inout_val = _mm512_add_ps(inout_val, in##x##_val); \ + } while (0) + +void reduce_fp16_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) +{ + const int element_size = 2; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[0] + i))); + switch (world_size) { + case 16: CVT_ADD_FP16(15); + case 15: CVT_ADD_FP16(14); + case 14: CVT_ADD_FP16(13); + case 13: CVT_ADD_FP16(12); + case 12: CVT_ADD_FP16(11); + case 11: CVT_ADD_FP16(10); + case 10: CVT_ADD_FP16(9); + case 9: CVT_ADD_FP16(8); + case 8: CVT_ADD_FP16(7); + case 7: CVT_ADD_FP16(6); + case 6: CVT_ADD_FP16(5); + case 5: CVT_ADD_FP16(4); + case 4: CVT_ADD_FP16(3); + case 3: CVT_ADD_FP16(2); + case 2: CVT_ADD_FP16(1); + case 1: break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = cvt_fp16_to_fp32(_mm256_loadu_si256((__m256i*)(buffers[j] + i))); + inout_val = _mm512_add_ps(inout_val, in_val); + } + } + _mm256_storeu_si256((__m256i*)(to_buffer + i), cvt_fp32_to_fp16(inout_val)); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { val += *(at::Half*)(buffers[j] + i); } + *(at::Half*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +#define CVT_ADD_F32(x) \ + do { \ + auto in##x##_val = _mm256_loadu_ps((float*)(buffers[x] + i)); \ + inout_val = _mm256_add_ps(inout_val, in##x##_val); \ + } while (0) + +void reduce_fp32_buffers(int start_elements, int num_elements, char* to_buffer, char** buffers) +{ + const int element_size = 4; + const int vector_length = VECTOR_LENGTH_IN_BYTES / element_size; + int main_elements = num_elements - (num_elements % vector_length); + int remain_elements = num_elements % vector_length; + + // process aligned part +#pragma omp parallel for + for (int i = start_elements * element_size; i < (start_elements + main_elements) * element_size; + i += VECTOR_LENGTH_IN_BYTES) { + auto inout_val = _mm256_loadu_ps((float*)(buffers[0] + i)); + switch (world_size) { + case 16: CVT_ADD_F32(15); + case 15: CVT_ADD_F32(14); + case 14: CVT_ADD_F32(13); + case 13: CVT_ADD_F32(12); + case 12: CVT_ADD_F32(11); + case 11: CVT_ADD_F32(10); + case 10: CVT_ADD_F32(9); + case 9: CVT_ADD_F32(8); + case 8: CVT_ADD_F32(7); + case 7: CVT_ADD_F32(6); + case 6: CVT_ADD_F32(5); + case 5: CVT_ADD_F32(4); + case 4: CVT_ADD_F32(3); + case 3: CVT_ADD_F32(2); + case 2: CVT_ADD_F32(1); + case 1: break; + default: + for (int j = 1; j < world_size; j++) { + auto in_val = _mm256_loadu_ps((float*)(buffers[j] + i)); + inout_val = _mm256_add_ps(inout_val, in_val); + } + } + _mm256_storeu_ps((float*)(to_buffer + i), inout_val); + } + + // process remaining part + int i = (start_elements + main_elements) * element_size; + while (remain_elements > 0) { + float val = 0.0f; + for (int j = 0; j < world_size; j++) { val += *(float*)(buffers[j] + i); } + *(float*)(to_buffer + i) = val; + remain_elements--; + i += element_size; + } +} + +static bool is_initialized = 0; +static int world_rank; + +void shm_initialize(int size, int rank, char* addr_string, char* port_string) +{ + if (is_initialized) return; + is_initialized = 1; + + world_size = size; + world_rank = rank; + + char shm_name_prefix[NAME_BUF_SIZE]; + char shm_name[NAME_BUF_SIZE]; + snprintf(shm_name_prefix, + NAME_BUF_SIZE, + "%s_%d_%s_%s", + SHM_BUFFER_NAME, + getuid(), + addr_string, + port_string); + // create shared workspace for SHM based allreduce + SharedData allreduce_buffer; + // allocate workspace_buf for current rank + struct allreduce_workspace* workspace_buf; + struct allreduce_workspace* workspace_buf_other; + workspace_buf = (struct allreduce_workspace*)malloc(sizeof(struct allreduce_workspace)); + snprintf(shm_name, NAME_BUF_SIZE, "%s_%d", shm_name_prefix, rank); + shared_create(&allreduce_buffer, shm_name, workspace_buf, sizeof(struct allreduce_workspace)); + workspace_buf = (struct allreduce_workspace*)allreduce_buffer.bytes; + workspace_buf->states[0] = coll_alt2_allreduce_naive__copy_in_done; + workspace_buf->states[1] = coll_begin; + + // create the workspace pointer list + workspace = (struct allreduce_workspace**)malloc(size * sizeof(struct allreduce_workspace*)); + symmetric_buffer[0] = (char**)malloc(size * sizeof(char**)); + symmetric_buffer[1] = (char**)malloc(size * sizeof(char**)); + distributed_buffer[0] = (char**)malloc(size * sizeof(char**)); + distributed_buffer[1] = (char**)malloc(size * sizeof(char**)); + + // map shm of all ranks + for (int i = 0; i < size; i++) { + if (i != rank) { + snprintf(shm_name, NAME_BUF_SIZE, "%s_%d", shm_name_prefix, i); + // printf("open %s, %d\n", shm_name, rank); + do { + shared_open(&allreduce_buffer, shm_name, sizeof(struct allreduce_workspace)); + } while (allreduce_buffer.descriptor == -1 && errno == ENOENT); + workspace_buf_other = (struct allreduce_workspace*)allreduce_buffer.bytes; + workspace[i] = workspace_buf_other; + } else { + workspace[i] = workspace_buf; + } + symmetric_buffer[0][i] = workspace[i]->buffer + BUFFER0_OFFSET(0); + symmetric_buffer[1][i] = workspace[i]->buffer + BUFFER0_OFFSET(1); + distributed_buffer[0][i] = workspace[i]->buffer + BUFFER1_OFFSET(0); + distributed_buffer[1][i] = workspace[i]->buffer + BUFFER1_OFFSET(1); + } +} + +static void parallel_memcpy(void* to, void* from, size_t n_bytes) + __attribute__((target("avx512bw"))); +static void parallel_memcpy(void* to, void* from, size_t n_bytes) +{ + auto aligned_bytes = n_bytes - (n_bytes % VECTOR_LENGTH_IN_BYTES); + // process aligned part +#pragma omp parallel for + for (int i = 0; i < aligned_bytes; i += VECTOR_LENGTH_IN_BYTES) { + auto val = _mm256_loadu_si256((__m256i*)((char*)from + i)); + _mm256_storeu_si256((__m256i*)((char*)to + i), val); + } + + // process remaining part + for (int i = aligned_bytes; i < n_bytes; i++) { *((char*)to + i) = *((char*)from + i); } +} + +#define positive_mod(num, mod) ((((num) % (mod)) + (mod)) % (mod)) +#define rank_mod(rank) positive_mod(rank, world_size) +size_t slice_size(size_t chunk_el, int slice_idx) +{ + size_t slice_size = chunk_el / world_size; + return slice_idx == world_size - 1 ? slice_size + (chunk_el % world_size) : slice_size; +} + +char* slice_data(char* data_ptr, size_t chunk_el, int el_size, int slice_idx) +{ + size_t slice_size = chunk_el / world_size; + size_t el_offset = slice_size * slice_idx; + return data_ptr + el_offset * el_size; +} + +size_t slice_el_start(size_t chunk_el, int slice_idx) +{ + size_t slice_size = chunk_el / world_size; + return slice_size * slice_idx; +} + +/* + Symmetrical naive all_reduce + step 0: before enter the function ith times, state is copy(i-1) + step 1: each rank copy data from input (data_ptr) to SHM buffer[i] + step 2: set own state to copy(i) + step 3: wait each other rank's state equal or later than copy(i) + step 4: reduce across SHM buffer(ith) directly into output (data_ptr) +*/ +void symmetric_naive_all_reduce(char* data_ptr, + c10::ScalarType scalar_type, + size_t chunk_size, + size_t chunk_el) +{ +#ifdef DO_PROFILE + static double total_t1_t0 = 0.0; + static double total_t2_t1 = 0.0; + static double total_t3_t2 = 0.0; + static int count = -16; // warmup + auto t0 = std::chrono::system_clock::now(); +#endif + + /* + We can't have infinite number of buffers and states. 2 sets of buffer + and 3 sets of states is just enough. Consider current rank is in step 3, + with it's own state set to copy(i), the other rank will them have the + following situations: + ------------------------------------------------ + my state | can I proceed? | the other rank state + ================================================ + | N | copy(i-1) + |----------------|--------------------- + copy(i) | Y | copy(i) + |----------------|--------------------- + | Y | copy(i+1) + ------------------------------------------------ + * When I have state as copy(i), the other rank cannot have state + copy(i-2) or before. In that case I'll be in state copy(i-1) and cannot + proceed to copy(i). + * The other rank cannot have state copy(i+2) or beyond because my + state is still copy(i), copy(i+1) is as far as the other rank could go. + * From a rank's POV, all the other ranks can be divided into three sets: + - Lagging ranks: ranks that are still working on previous iteration + - Syncing ranks: ranks that are working on current iteration + - Leading ranks: ranks that are working on next iteration + * We can have 3 sets of states, one set for syncing ranks; one set for + lagging ranks; one set of leading ranks. With 3 sets of states, we can + distinguish between lagging and leading ranks. + * Note from any rank's POV, leading ranks and lagging ranks does not + appear at the same time. Either all other ranks are syncing or + lagging, or all other ranks are syncing or leading. Otherwise leading + and lagging ranks will be 2 iterations apart and this should not happen. + * So we have 2 sets of buffers, one buffer is used by current iter; + one buffer used by either lagging ranks or leading ranks. + */ + const int state_group = 0; + static int current_buffer = 0; + static int state_idx = 0; + + enum coll_state copy_current, copy_next; + + switch (state_idx) { + case 0: + copy_current = coll_allreduce_naive__copy_in_done; + copy_next = coll_alt1_allreduce_naive__copy_in_done; + break; + case 1: + copy_current = coll_alt1_allreduce_naive__copy_in_done; + copy_next = coll_alt2_allreduce_naive__copy_in_done; + break; + case 2: + copy_current = coll_alt2_allreduce_naive__copy_in_done; + copy_next = coll_allreduce_naive__copy_in_done; + break; + default: assert(!"Should not get here."); + } + state_idx = (state_idx + 1) % 3; + + parallel_memcpy(symmetric_buffer[current_buffer][world_rank], data_ptr, chunk_size); + std::atomic_thread_fence(std::memory_order_release); + workspace[world_rank]->states[state_group] = copy_current; + +#ifdef DO_PROFILE + auto t1 = std::chrono::system_clock::now(); +#endif + + for (int i = 0; i < world_size; i++) { + // wait until the other rank copy the buffer + if (i != world_rank) { wait_buffer_state_until_2(i, copy_current, copy_next, state_group); } + } +#ifdef DO_PROFILE + auto t2 = std::chrono::system_clock::now(); +#endif + + // each rank reduce the buffer independently so therre is no need for synchronization afterward + reduce_all_buffers( + 0, chunk_el, scalar_type, world_rank, data_ptr, symmetric_buffer[current_buffer]); + + // switch buffer + current_buffer = 1 - current_buffer; + +#ifdef DO_PROFILE + auto t3 = std::chrono::system_clock::now(); + + count++; + if (count > 0) { + total_t1_t0 += std::chrono::duration_cast(t1 - t0).count(); + total_t2_t1 += std::chrono::duration_cast(t2 - t1).count(); + total_t3_t2 += std::chrono::duration_cast(t3 - t2).count(); + if (world_rank == 0 && count == 1000) { + printf("symmetric_naive_all_reduce time breakdown:\n"); + printf("\tcopy input buffer: %.2f\n", total_t1_t0 / count); + printf("\twait for copy: %.2f\n", total_t2_t1 / count); + printf("\treduce: %.2f\n", total_t3_t2 / count); + } + } +#endif +} + +// naive allreduce distributed, each rank do naive reduce on its slice +void distributed_naive_reduce(char* data_ptr, + c10::ScalarType scalar_type, + size_t chunk_size, + size_t chunk_el) +{ +#ifdef DO_PROFILE + static double total_t1_t0 = 0.0; + static double total_t2_t1 = 0.0; + static double total_t3_t2 = 0.0; + static double total_t4_t3 = 0.0; + static double total_t5_t4 = 0.0; + static int count = -16; // warmup + auto t0 = std::chrono::system_clock::now(); +#endif + + const int state_group = 1; + static int current_buffer = 0; + static int state_idx = 0; + + enum coll_state copy_current, copy_next, reduce_current; + + // similar to symmetric_naive_allreduce, but here we only need two sets of + // states, because distributed naive reduce has two barriers in the algorithm + switch (state_idx) { + case 0: + copy_current = coll_allreduce_naive__copy_in_done; + reduce_current = coll_allreduce_naive__reduce_done; + copy_next = coll_alt1_allreduce_naive__copy_in_done; + break; + case 1: + copy_current = coll_alt1_allreduce_naive__copy_in_done; + reduce_current = coll_alt1_allreduce_naive__reduce_done; + copy_next = coll_allreduce_naive__copy_in_done; + break; + default: assert(!"Should not get here."); + } + state_idx = (state_idx + 1) % 2; + + int data_size = chunk_size / chunk_el; + parallel_memcpy(distributed_buffer[current_buffer][world_rank], data_ptr, chunk_size); + std::atomic_thread_fence(std::memory_order_release); + workspace[world_rank]->states[state_group] = copy_current; + +#ifdef DO_PROFILE + auto t1 = std::chrono::system_clock::now(); +#endif + + for (int i = 0; i < world_size; i++) { + // wait until all the other ranks copy the buffer + if (i != world_rank) + wait_buffer_state_until_2(i, copy_current, reduce_current, state_group); + } + +#ifdef DO_PROFILE + auto t2 = std::chrono::system_clock::now(); +#endif + + // reduce scatter + reduce_all_buffers(slice_el_start(chunk_el, world_rank), + slice_size(chunk_el, world_rank), + scalar_type, + world_rank, + distributed_buffer[current_buffer][world_rank], + distributed_buffer[current_buffer]); + std::atomic_thread_fence(std::memory_order_release); + workspace[world_rank]->states[state_group] = reduce_current; + +#ifdef DO_PROFILE + auto t3 = std::chrono::system_clock::now(); +#endif + + for (int i = 0; i < world_size; i++) { + // wait until all the other ranks reduce the buffer + if (i != world_rank) wait_buffer_state_until_2(i, reduce_current, copy_next, state_group); + } + + auto t4 = std::chrono::system_clock::now(); + + for (int i = 0; i < world_size; i++) { + int rank = (i + world_rank) % world_size; + parallel_memcpy( + slice_data(data_ptr, chunk_el, data_size, rank), + slice_data( + distributed_buffer[current_buffer][rank], chunk_el, chunk_size / chunk_el, rank), + slice_size(chunk_el, rank) * data_size); + } + + current_buffer = 1 - current_buffer; + +#ifdef DO_PROFILE + auto t5 = std::chrono::system_clock::now(); + count++; + if (count > 0) { + total_t1_t0 += std::chrono::duration_cast(t1 - t0).count(); + total_t2_t1 += std::chrono::duration_cast(t2 - t1).count(); + total_t3_t2 += std::chrono::duration_cast(t3 - t2).count(); + total_t4_t3 += std::chrono::duration_cast(t4 - t3).count(); + total_t5_t4 += std::chrono::duration_cast(t5 - t4).count(); + if (world_rank == 0 && count == 1000) { + printf("distributed_naive_reduce time breakdown:\n"); + printf("\tcopy input buffer: %.2f\n", total_t1_t0 / count); + printf("\twait for copy: %.2f\n", total_t2_t1 / count); + printf("\treduce: %.2f\n", total_t3_t2 / count); + printf("\twait for reduce finish: %.2f\n", total_t4_t3 / count); + printf("\tcopy out: %.2f\n", total_t5_t4 / count); + } + } +#endif +} + +void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size) +{ + for (int offset = 0; offset < data_size; offset += MAX_BUF_SIZE) { + auto data_ptr = ((char*)(data.data_ptr()) + offset); + size_t chunk_size = data_size - offset > MAX_BUF_SIZE ? MAX_BUF_SIZE : data_size - offset; + size_t chunk_el = chunk_size / (data_size / numel); + if (chunk_size < NAIVE_ALLREDUCE_THRESHOLD) + symmetric_naive_all_reduce(data_ptr, data.scalar_type(), chunk_size, chunk_el); + else + distributed_naive_reduce(data_ptr, data.scalar_type(), chunk_size, chunk_el); + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h new file mode 100644 index 0000000000000000000000000000000000000000..7f73197a8caa5c4487c79c0c30ef926797bae97b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm.h @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#ifndef __SHM_COLLECTIVES__ +#define __SHM_COLLECTIVES__ +#define VECTOR_LENGTH_IN_BYTES 32 +void shm_initialize(int size, int rank, char* addr_string, char* port_string); +void all_reduce_outer_loop(torch::Tensor& data, size_t numel, int data_size); +void barrier_wait(int root_idx, int num_ranks); +#endif diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp new file mode 100644 index 0000000000000000000000000000000000000000..5be5cb799a7bb47f624f5b4952f2947738864d62 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp @@ -0,0 +1,186 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +#include "shm.h" + +// #define DO_PROFILE +#ifdef DO_PROFILE +#include +#include +#endif + +// Communication settings +static int world_rank = -1; +static int world_size = -1; + +static bool is_initialized = 0; + +static bool all_ranks_local_p = false; + +void initialize(int size, int rank) +{ + if (is_initialized) return; + + // Check whether all ranks is on the same physical machine. + // If true, we will use an SHM based low latency allreduce + + auto ls_string = std::getenv("LOCAL_SIZE"); + int ls = 0; + if (ls_string != NULL) { ls = std::stoi(std::getenv("LOCAL_SIZE")); } + + if (size >= 1 && size == ls) { all_ranks_local_p = true; } + + world_size = size; + world_rank = rank; + is_initialized = 1; + + auto addr_string = std::getenv("MASTER_ADDR"); + if (addr_string == NULL) { addr_string = ""; } + auto port_string = std::getenv("MASTER_PORT"); + if (port_string == NULL) { port_string = ""; } + + if (all_ranks_local_p) { shm_initialize(size, rank, addr_string, port_string); } +} + +void inference_all_reduce_(torch::Tensor& data, int op); + +// Success - return 0 +// Fail (cannot hornor the request and need to fall back) - return -1 +void inference_all_reduce_(torch::Tensor& data, int op) +{ + assert(op == 0); +#ifdef DO_PROFILE + static double total_time = 0.0; + static double total_time_sq = 0.0; + static int count = -16; // warmup + static double max_time = 0.0; + static double min_time = DBL_MAX; + // make sure all rank reach this point before measuring time + // turn on this if you suspect each rank didn't reach here at the same time (stragger) + // if (all_ranks_local_p) { barrier_wait(0, world_size); } + auto start = std::chrono::system_clock::now(); +#endif + + auto numel = data.numel(); + + int data_size = 0; + bool data_type_fallback = false; + + switch (data.scalar_type()) { + case c10::ScalarType::BFloat16: data_size = numel * 2; break; + case c10::ScalarType::Half: data_size = numel * 2; break; + case c10::ScalarType::Float: data_size = numel * 4; break; + default: data_type_fallback = true; + } + + if (data_type_fallback) return; + + all_reduce_outer_loop(data, numel, data_size); + +#ifdef DO_PROFILE + auto end = std::chrono::system_clock::now(); + count++; + if (count > 0) { + double elapsed = std::chrono::duration_cast(end - start).count(); + if (elapsed > max_time) { max_time = elapsed; } + if (elapsed < min_time) { min_time = elapsed; } + total_time += elapsed; + total_time_sq += elapsed * elapsed; + if (world_rank == 0 && count == 1000) { + auto avg = total_time / count; + auto sd = + sqrt(total_time_sq / count - total_time * total_time / (count * count)) / avg * 100; + printf(" C++ kernel\t\t %.2f\t %.2f\t%.2f\t %.2f\n", + min_time, + max_time, + total_time / count, + sd); + } + } +#endif + return; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("initialize", &initialize, "shm initialize"); } + +TORCH_LIBRARY(deepspeed, m) +{ + m.def("inference_all_reduce(Tensor self) -> Tensor"); + m.def("inference_all_reduce_(Tensor(a!) self) -> Tensor(a!)"); +} + +torch::Tensor inference_all_reduce_meta(const torch::Tensor& self_) +{ + torch::Tensor result_ = torch::empty_like(self_); + return result_; +} + +torch::Tensor& inference_all_reduce__meta(torch::Tensor& self_) { return self_; } + +torch::Tensor& inference_all_reduce__cpu(torch::Tensor& self_) +{ + TORCH_INTERNAL_ASSERT(self_.device().type() == torch::DeviceType::CPU); + torch::Tensor self_tensor = self_.contiguous(); + inference_all_reduce_(self_tensor, 0); + return self_; +} + +torch::Tensor inference_all_reduce_cpu(const torch::Tensor& self_) +{ + torch::Tensor result = self_.clone(); + inference_all_reduce__cpu(result); + return result; +} + +#include +// The boilerplate functionalization logic, that teaches functionalization +// how to map x_() calls into x() calls. +// Long term, we'd like to not require users to write this logic. +// HOWEVER, if you have a custom op that is mutable, +// You will still need to write an out-of-place version of that op! +at::Tensor& inference_all_reduce__functionalization_glue(at::Tensor& x) +{ + // We expect all tensor inputs to our op to be "functional tensors" + TORCH_INTERNAL_ASSERT(at::functionalization::impl::isFunctionalTensor(x)); + // First, sync and unwrap and functional tensors + at::functionalization::impl::sync(x); + auto x_ = at::functionalization::impl::from_functional_tensor(x); + // Grab the dispatcher entry corresponding to the out-of-place op, "x" + static auto op_handle = c10::Dispatcher::singleton() + // specify namespace::op_name, op_overload_name + .findSchemaOrThrow("deepspeed::inference_all_reduce", "") + // Specify the C++ schema of the out-of-place op. + .typed(); + // Next, redispatch to the out-of-place op, x() (user called x_, we call x) + at::Tensor tmp_output; + { + at::AutoDispatchSkipFunctionalize guard; + tmp_output = op_handle.call(x_); + } + // Finally, tell functionalization about this mutation. + at::functionalization::impl::replace_(x, tmp_output); + at::functionalization::impl::commit_update(x); + at::functionalization::impl::sync(x); + return x; +} + +TORCH_LIBRARY_IMPL(deepspeed, CPU, m) +{ + m.impl("inference_all_reduce", inference_all_reduce_cpu); + m.impl("inference_all_reduce_", inference_all_reduce__cpu); +} + +TORCH_LIBRARY_IMPL(deepspeed, Meta, m) +{ + m.impl("inference_all_reduce", inference_all_reduce_meta); + m.impl("inference_all_reduce_", inference_all_reduce__meta); +} + +TORCH_LIBRARY_IMPL(deepspeed, Functionalize, m) +{ + m.impl("inference_all_reduce_", inference_all_reduce__functionalization_glue); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp new file mode 100644 index 0000000000000000000000000000000000000000..708df7f0146aa996e1652ed938c038331738e149 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/cpu/lion/fused_lion.cpp @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_lion.h" + +// C++ interface + +void multi_tensor_lion(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, /*gpmv*/ + const float lr, + const float beta1, + const float beta2, + const int step, + const int mode, + const float weight_decay) +{ + static bool initialized = false; + if (!initialized) { + create_lion_optimizer(0); + initialized = true; + } + for (int i = 0; i < tensor_lists[0].size(); i++) { + ds_lion_step(0, + step, + lr, + beta1, + beta2, + weight_decay, + tensor_lists[1][i], + tensor_lists[0][i], + tensor_lists[2][i]); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("multi_tensor_lion", + &multi_tensor_lion, + "Compute and apply gradient update to parameters for Lion optimizer"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h new file mode 100644 index 0000000000000000000000000000000000000000..8d4173f1a6a2e5a073d7da4c129f471d4f394632 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_atomic.h @@ -0,0 +1,886 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once +#include +#include +#include +namespace cutlass { +namespace epilogue { +namespace threadblock { + +template +struct atomic_store {}; + +template +struct atomic_store::value>::type> { + using Element = typename AccessType::Element; + static const int kCount = AccessType::kElements; + + CUTLASS_DEVICE + atomic_store(AccessType const& D, void* ptr, bool pred_guard) + { + static_assert(!(kCount % 2), "kCount must be even"); + half2* p = reinterpret_cast(ptr); + uint const* data = reinterpret_cast(&D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + : + : "r"((int)pred_guard)); + for (int i = 0; i < kCount / 2; i++) { + asm volatile(" @p red.relaxed.global.add.noftz.f16x2 [%0], %1;\n" + : + : "l"(p + i), "r"(data[i])); + } + asm volatile("}\n" ::); + } +}; + +template +struct atomic_store::value>::type> { + using Element = typename AccessType::Element; + static const int kCount = AccessType::kElements; + + CUTLASS_DEVICE + atomic_store(AccessType const& D, void* ptr, bool pred_guard) + { + Element* p = reinterpret_cast(ptr); + uint const* data = reinterpret_cast(&D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + : + : "r"((int)pred_guard)); + for (int i = 0; i < kCount; i++) { + asm volatile(" @p red.relaxed.global.add.f32 [%0], %1;\n" + : + : "l"(p + i), "r"(data[i])); + } + asm volatile("}\n" ::); + } +}; + +template +class PredicatedTileIteratorAffineRankNAtomic { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::AffineRankN; + using TensorRef = TensorRef; + using TensorView = TensorView; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = typename Layout::TensorCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert(ThreadMap::Iterations::kRow > 0, "ThreadMap::Iterations::kRow must be > 0"); + static_assert(ThreadMap::Iterations::kGroup > 0, "ThreadMap::Iterations::kGroup must be > 0"); + static_assert(ThreadMap::Iterations::kCluster > 0, + "ThreadMap::Iterations::kCluster must be > 0"); + static_assert(ThreadMap::Iterations::kColumn > 0, "ThreadMap::Iterations::kColumn must be > 0"); + static_assert(!(Layout::kRank % 2), + "Layout rank must be even. This assumes the first half of the " + "modes correspond to the 'row' " + "and the second half of the modes correspond to the 'column'"); + + static bool const kBigEndian = false; + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Parameters structure + struct Params { + // + // Data members + // + + Layout layout; + + /// Stride in units of bytes along M modes + Coord stride_m; + + /// Stride in units of bytes along N modes + Coord stride_n; + + /// Fast divmod objects divided by tensor extents + FastDivmod divmod_m[(Layout::kRank == 2) ? 1 : (Layout::kRank / 2 - 1)]; + + /// Fast divmod objects divided by tensor extents + FastDivmod divmod_n[(Layout::kRank == 2) ? 1 : (Layout::kRank / 2 - 1)]; + + int64_t rank2_inc_col; + int64_t rank2_inc_row; + + // + // Methods + // + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(TensorCoord const& extent, Layout const& layout_) : layout(layout_) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2; ++i) { + stride_m[i] = OffsetBytes(layout_.stride()[i]); + stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2]); + } + + if (kBigEndian) { + // "Big Endian" scheme + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { + divmod_m[i] = FastDivmod(extent[i + 1]); + divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2 + 1]); + } + } else { + // "Little Endian" scheme + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { + divmod_m[i] = FastDivmod(extent[i]); + divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2]); + } + } + } + + CUTLASS_HOST_DEVICE + Params(Layout const& layout_) : layout(layout_) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2; ++i) { + stride_m[i] = OffsetBytes(layout_.stride()[i]); + stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2]); + } + + rank2_inc_col = ThreadMap::Delta::kColumn * stride_n[0]; + rank2_inc_row = ThreadMap::Delta::kRow * stride_m[0]; + } + }; + + /// Mask object + struct Mask { + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + CUTLASS_HOST_DEVICE + Mask() { enable(); } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_HOST_DEVICE void clear() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = false; } + } + + ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask + CUTLASS_DEVICE void enable() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = true; } + } + }; + +private: + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + Params params_; + + /// Byte-level pointer + uint8_t* byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in columns + Index extent_col_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column position (assuming steady-state predicates have + /// been computed) + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + /// Offsets in columns, cached for performance + int64_t offset_modes_n_[ThreadMap::Iterations::kColumn]; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + +private: + // + // Methods + // + +public: + // + // Methods + // + + /// Constructor + CUTLASS_DEVICE + PredicatedTileIteratorAffineRankNAtomic( + Params const& params, + Element* pointer, + MatrixCoord extent, + int thread_idx, + MatrixCoord threadblock_offset = MatrixCoord(), + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : params_(params) + { + MatrixCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_col_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + if (Layout::kRank > 2) { + // Initialize predicates + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + // + // Compute coordinate and decompose into N modes + // + + int coord_n = thread_start_column_ + c * ThreadMap::Delta::kColumn; + + mask_.predicates[c] = coord_n < extent.column(); + + Coord modes_n; + + int64_t offset_modes_n = 0; + + if (kBigEndian) { + modes_n = CoordinateDecomposition(coord_n, params_.divmod_n); + + offset_modes_n = dot(modes_n, params_.stride_n); + } else { + modes_n = CoordinateDecompositionLittleEndian( + coord_n, params_.divmod_n); + + offset_modes_n = dot(modes_n, params_.stride_n); + } + + offset_modes_n_[c] = offset_modes_n; + } + + if (!pointer) { mask_.clear(); } + } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer); + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) + { + uint8_t* byte_pointer = byte_pointer_; + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + int row_begin = thread_start_row_ + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + int64_t offset_modes_m = row_begin * params_.stride_m[0]; + + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + // + // Compute coordinate and decompose into M modes + // + + int coord_m = row * ThreadMap::Delta::kRow + row_begin; + + Coord modes_m; + + if (Layout::kRank > 2) { + if (kBigEndian) { + modes_m = CoordinateDecomposition(coord_m, + params_.divmod_m); + } else { + modes_m = CoordinateDecompositionLittleEndian( + coord_m, params_.divmod_m); + } + + offset_modes_m = dot(modes_m, params_.stride_m); + } + + // + // Compute the offset due to modes M + // + + bool row_guard = (coord_m < extent_row_); + int64_t offset_modes_n = thread_start_column_ * params_.stride_n[0]; + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + // + // Compute coordinate and decompose into N modes + // + + if (Layout::kRank > 2) { offset_modes_n = offset_modes_n_[column]; } + + // + // Compute the pointer and access + // + bool guard; + if (Layout::kRank > 2) { + guard = row_guard && mask_.predicates[column]; + } else { + guard = (coord_m < extent_row_) && + ((thread_start_column_ + ThreadMap::Delta::kColumn * column) < + extent_col_); + } + + atomic_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)(byte_pointer + offset_modes_m + offset_modes_n + byte_offset), + guard); + + if (Layout::kRank == 2) { offset_modes_n += params_.rank2_inc_col; } + } + + if (Layout::kRank == 2) { offset_modes_m += params_.rank2_inc_row; } + } + } + } + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } + + CUTLASS_DEVICE + void load(Fragment& frag) {} + + /// Advances to the next position to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorAffineRankNAtomic& operator++() + { + ++state_[0]; + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + state_[0] = 0; + ++state_[1]; + + thread_start_row_ += + (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + state_[1] = 0; + ++state_[2]; + + thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { state_[2] = 0; } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_DEVICE void clear_mask() { mask_.clear(); } + + ///< Efficiently enables all accesses guarded by mask + CUTLASS_DEVICE void enable_mask() { mask_.enable(); } + + ///< Sets the mask + CUTLASS_DEVICE void get_mask(Mask& mask) { mask = mask_; } + + ///< Sets the mask + CUTLASS_DEVICE void set_mask(Mask const& mask) { mask_ = mask; } +}; + +template +class PredicatedTileIteratorAtomic { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static bool constexpr PermuteD = !layout::is_trivial_permute; + + static_assert(ThreadMap::Iterations::kRow > 0, "ThreadMap::Iterations::kRow must be > 0"); + static_assert(ThreadMap::Iterations::kGroup > 0, "ThreadMap::Iterations::kGroup must be > 0"); + static_assert(ThreadMap::Iterations::kCluster > 0, + "ThreadMap::Iterations::kCluster must be > 0"); + static_assert(ThreadMap::Iterations::kColumn > 0, "ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + using Base = PredicatedTileIteratorParams; + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc()) + { + } + + CUTLASS_HOST_DEVICE + Params(Base const& base) : Base(base) {} + }; + + /// Mask object + struct Mask { + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + CUTLASS_HOST_DEVICE + Mask() { enable(); } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_HOST_DEVICE void clear() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = false; } + } + + ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask + CUTLASS_DEVICE void enable() + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { predicates[i] = true; } + } + }; + +private: + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorParams params_; + + /// Byte-level pointer. This pointer is usually for both load() and store(), + /// unless PermuteD is performed. When having PermuteD, byte_pointer_ is only + /// for load(). + uint8_t* byte_pointer_; + + /// Byte-level pointer for store(). Due to PermuteD Op, store_byte_pointer_ + /// may be with different address computation compared to byte_pointer_. + uint8_t* store_byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + /// Scatter indices + int const* indices_; + + /// PermuteDLayout + PermuteDLayout permute_layout_; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); + +private: + // + // Methods + // + +public: + // + // Methods + // + + /// Constructor + CUTLASS_DEVICE + PredicatedTileIteratorAtomic(PredicatedTileIteratorParams const& params, + Element* pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord(), + int const* indices = nullptr) + : params_(params), + indices_(indices), + permute_layout_(PitchLinearCoord(extent.column(), extent.row()), + params_.stride * kElementsPerAccess / sizeof(AccessType)) + { + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_column_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + mask_.predicates[c] = + ((thread_offset.column() + ThreadMap::Delta::kColumn * c) < extent.column()); + } + + // Null pointer performs no accesses + if (!pointer) { mask_.clear(); } + + if (ScatterD && !indices) { mask_.clear(); } + + // Initialize byte_pointer_ + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.row()) * LongIndex(params_.stride) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + + if (ScatterD) { + byte_pointer_ = + reinterpret_cast(pointer) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + } + + // store_byte_pointer_ is set to be the same with byte_pointer_ unless + // PermuteD is used. + store_byte_pointer_ = PermuteD ? reinterpret_cast(pointer) : byte_pointer_; + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + store_byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, int64_t byte_offset) const + { + uint8_t* byte_pointer = store_byte_pointer_; + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + CUTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + CUTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * + (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType* memory_pointer = + reinterpret_cast(byte_pointer + byte_offset); + + if (ScatterD && row_guard) { + assert(indices_); + + memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + LongIndex(indices_[row_offset + thread_start_row_]) * + LongIndex(params_.stride)); + } + + CUTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + bool guard = row_guard && mask_.predicates[column]; + + if (PermuteD) { + int col_offset = column * ThreadMap::Delta::kColumn; + + int col = col_offset + thread_start_column_; + int row = row_offset + thread_start_row_; + + // Locate memory_pointer + memory_pointer = reinterpret_cast( + byte_pointer + byte_offset + + permute_layout_(PitchLinearCoord(col, row)) * sizeof(AccessType) / + kElementsPerAccess); + } + atomic_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void*)&memory_pointer[0], + guard); + + if (!PermuteD) { + memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess); + } + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD && !PermuteD) { byte_pointer += params_.increment_row; } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) const { store_with_byte_offset(frag, 0); } + + CUTLASS_DEVICE + void load(Fragment& frag) {} + + CUTLASS_DEVICE + MatrixCoord thread_start() const + { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_row() const { return thread_start_row_; } + + /// Need to get the thread start row from the tile iterator + CUTLASS_DEVICE + int32_t thread_start_column() const { return thread_start_column_; } + + /// Extent of the matrix in rows + CUTLASS_DEVICE + Index extent_row() const { return extent_row_; } + + /// Extent of the matrix in columns + CUTLASS_DEVICE + Index extent_column() const { return extent_column_; } + + /// Advances to the next position to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorAtomic& operator++() + { + ++state_[0]; + + if (!ScatterD && !PermuteD) { store_byte_pointer_ += params_.advance_row; } + + if (!ScatterD) { byte_pointer_ += params_.advance_row; } + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + state_[0] = 0; + ++state_[1]; + byte_pointer_ += params_.advance_group; + store_byte_pointer_ += params_.advance_group; + + thread_start_row_ += + (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + state_[1] = 0; + ++state_[2]; + byte_pointer_ += params_.advance_cluster; + store_byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + byte_pointer_ += params_.advance_tile; + store_byte_pointer_ += params_.advance_tile; + + thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow * + ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile; + } + } + } + + return *this; + } + + /// Advances a number of positions to load or store + CUTLASS_HOST_DEVICE + PredicatedTileIteratorAtomic& operator+=(int increment) + { + // Row + state_[0] += increment; + int increment_row = state_[0] / ThreadMap::Count::kRow; + state_[0] = state_[0] % ThreadMap::Count::kRow; + + byte_pointer_ += (params_.advance_row * increment); + store_byte_pointer_ += (params_.advance_row * increment); + thread_start_row_ += (ThreadMap::Shape::kRow * increment); + + // Group + state_[1] += increment_row; + int increment_group = state_[1] / ThreadMap::Count::kGroup; + state_[1] = state_[1] % ThreadMap::Count::kGroup; + + byte_pointer_ += (params_.advance_group * increment_row); + store_byte_pointer_ += (params_.advance_group * increment_row); + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * ThreadMap::Shape::kRow * + ThreadMap::Count::kRow * increment_row; + + // Cluster + state_[2] += increment_group; + int increment_cluster = state_[2] / ThreadMap::Count::kCluster; + state_[2] = state_[2] % ThreadMap::Count::kCluster; + + byte_pointer_ += (params_.advance_cluster * increment_group); + store_byte_pointer_ += (params_.advance_cluster * increment_group); + thread_start_row_ += ThreadMap::Count::kGroup * ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * ThreadMap::Shape::kRow * increment_group; + + // Tile + byte_pointer_ += (params_.advance_tile * increment_cluster); + store_byte_pointer_ += (params_.advance_tile * increment_cluster); + thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow * + ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile * + increment_cluster; + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + CUTLASS_DEVICE void clear_mask() { mask_.clear(); } + + ///< Efficiently enables all accesses guarded by mask + CUTLASS_DEVICE void enable_mask() { mask_.enable(); } + + ///< Sets the mask + CUTLASS_DEVICE void get_mask(Mask& mask) const { mask = mask_; } + + ///< Sets the mask + CUTLASS_DEVICE void set_mask(Mask const& mask) { mask_ = mask; } +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace cutlass diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h new file mode 100644 index 0000000000000000000000000000000000000000..629047dbb057384ac8f0d7fa5d557e0c070cd830 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/predicated_tile_iterator_residual_last.h @@ -0,0 +1,1938 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Templates implementing loading of tiles from pitch-linear rank=2 + tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile + this iterator visits maybe partial, then the remaining tiles are complete. + So, we only need to compute the predicates twice, once before the first tile + and once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "cutlass/arch/memory.h" +#include "cutlass/transform/threadblock/predicated_tile_access_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileIteratorResidualLast +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +/// Regular tile iterator using a precomputed control structure to minimize +/// register liveness and integer arithmetic. +/// +/// Layout is assumed to be invariant at the time the precomputed "Params" +/// object is constructed. +/// +/// Base pointer and tensor extents may be specified at the time the iterator is +/// constructed. Subsequently, they are assumed to be immutable. +/// +/// Adding a logical coordinate offset may be performed at the time the iterator +/// is constructed. Subsequent additions to logical coordinate offset may be +/// performed but are relatively expensive. +/// +/// Visitation order is intended to first visit a "residual" tile that may be +/// partially full in both the advance dimension and the steady-state dimension. +/// This is assumed to be the last tile in the iteration sequence. Advancing an +/// iterator that has just been constructed moves to the first tile that is full +/// in the advance dimension and recomputes predicates. Subsequent accesses may +/// be performed without updating internal predicates and are efficient in terms +/// of live register state and pointer arithmetic instructions. +/// +/// To be efficient, this assumes the iterator will be dereferenced and advanced +/// at least once outside any looping structure to minimize integer arithmetic. +/// +/// Accesses out of bounds are safe so long as `clear_mask()` is called prior to +/// dereferencing the iterator. +/// +/// +/// Example: +/// +/// An efficient pipeline structure may be constructed as follows: +/// +// template +// __global__ void kernel( +// typename Iterator::Params params, +// typename Iterator::Element *ptr, +// TensorCoord extent) { +// +// typename Iterator::Fragment fragment; +// +// TensorCoord threadblock_offset(0, 0); +// +// Iterator iter(params, ptr, extent, threadIdx.x, threadblock_offsets); +// +// +// fragment = *iter; // load "residue" tile first +// ++iter; // advance to first "steady state" tile and update +// internal masks +// +// +// #pragma unroll +// for (int i = Remaining - 1; i >= 0; --i) { +// +// f(fragment); +// +// if (!i) { +// iter.clear_mask(); // light-weight operation to clear masks - +// subsequent loads become NO-OPs. +// } +// +// fragment = *iter; // load tile during "steady state" phase +// ++iter; // advance to next tile - lightweight due to +// steady-state masks +// } +// } +// +// void host(TensorView view) { +// +// using Iterator = +// transform::threadblock::PredicatedTileIteratorResidualLast; +// +// typename Iterator::Params params(view.layout()); +// +// kernel(params, view.data()); +// } +/// +/// +template +class PredicatedTileIteratorResidualLast; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + /// Type used for internal memory accesses + using AccessType = + AlignedArray::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = PredicatedTileAccessIteratorResidualLast; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + using Base = typename TileAccessIterator::Params::Base; + + friend PredicatedTileIteratorResidualLast; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout) {} + + CUTLASS_HOST_DEVICE + Params() {} + + CUTLASS_HOST_DEVICE + Params(Base const& base) : params_(base) {} + }; + +private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + +private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + /// Gather indices + int const* indices = nullptr) + : address_iterator_(params.params_, pointer, extent, thread_id, threadblock_offset, indices) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { address_iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { address_iterator_.get_mask(mask); } + + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const* byte_ptr = + reinterpret_cast(address_iterator_.get()) + byte_offset; + + AccessType const* access_ptr = reinterpret_cast(byte_ptr); + + cutlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_byte_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + address_iterator_.set_iteration_index(0); + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char* byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType* access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { *access_ptr = frag_ptr[idx]; } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = + PredicatedTileIteratorResidualLast, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize, + Gather>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()), + indices) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = + PredicatedTileIteratorResidualLast, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize, + Gather>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< Gather indices + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()), + indices) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank-2 data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRankN<2>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + /// Type used for internal memory accesses + using AccessType = + AlignedArray::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = PredicatedTileAccessIteratorResidualLast; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend PredicatedTileIteratorResidualLast; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout) {} + + CUTLASS_HOST_DEVICE + Params() {} + }; + +private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char*; + +private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : address_iterator_(params.params_, pointer, extent, thread_id, threadblock_offset) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + if (kAdvanceRank) + address_iterator_.add_tile_offset(make_Coord(0, 1)); + else + address_iterator_.add_tile_offset(make_Coord(1, 0)); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { address_iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { address_iterator_.get_mask(mask); } + + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + AccessType* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const* byte_ptr = + reinterpret_cast(address_iterator_.get()) + byte_offset; + + AccessType const* access_ptr = reinterpret_cast(byte_ptr); + + cutlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_byte_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + address_iterator_.set_iteration_index(0); + AccessType const* frag_ptr = reinterpret_cast(&frag); + + CUTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + CUTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char* byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType* access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { *access_ptr = frag_ptr[idx]; } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_byte_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank 2 +/// column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = + PredicatedTileIteratorResidualLast, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))) + { + } + }; + +private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column())) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for affine rank 2 +/// row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = + PredicatedTileIteratorResidualLast, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given an AffineRankN<2> tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) + : params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))) + { + } + }; + +private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const& threadblock_offset, ///< Initial offset of threadblock + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row())) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_byte_offset(Fragment& frag, LongIndex byte_offset) + { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_byte_offset(Fragment const& frag, LongIndex byte_offset) + { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for interleaved data. +/// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template +class PredicatedTileIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorResidualLast for interleaved-32 +/// data. It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorResidualLast, + AdvanceRank, + ThreadMap_, + AccessSize, + false> { +public: + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element*; + using NonConstPointer = typename platform::remove_const::type*; + + using UnderlyingIterator = PredicatedTileIteratorResidualLast< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize>; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = + cutlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIteratorResidualLast; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + CUTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + CUTLASS_HOST_DEVICE + Params(Layout const& layout) : params_(layout::PitchLinear(layout.stride(0))) {} + + CUTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const& base) : params_(base) {} + }; + +private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast( + /// Precomputed parameters object + Params const& params, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const& threadblock_offset, + int const* indices = nullptr ///< gather/scatter indices, note no support for + ///< gather/scatter at this specialization + ) + : iterator_(params.params_, + pointer, + layout::PitchLinearCoord(extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) + { + } + + /// Construct a PredicatedTileIteratorResidualLast with zero threadblock + /// offset + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast(Params const& params, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorResidualLast(params, pointer, extent, thread_id, make_Coord(0, 0)) + { + } + + /// Adds a pointer offset in units of Element + CUTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) + { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast& operator++() + { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + CUTLASS_HOST_DEVICE + PredicatedTileIteratorResidualLast operator++(int) + { + PredicatedTileIteratorResidualLast self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + CUTLASS_HOST_DEVICE + void set_residual_tile(bool enable) { iterator_.set_residual_tile(enable); } + + /// Clears the predicate set efficiently + CUTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + CUTLASS_HOST_DEVICE + void set_mask(Mask const& mask) { iterator_.set_mask(mask); } + + /// Gets the mask + CUTLASS_HOST_DEVICE + void get_mask(Mask& mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load_with_pointer_offset(Fragment& frag, Index pointer_offset) + { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + CUTLASS_DEVICE + void load(Fragment& frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store_with_pointer_offset(Fragment const& frag, Index pointer_offset) + { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + CUTLASS_DEVICE + void store(Fragment const& frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h new file mode 100644 index 0000000000000000000000000000000000000000..2435c07f8989dcd6ad9c99a2fd530fe6092b1e08 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/transpose_warp_iterator.h @@ -0,0 +1,57 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holdvr nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include "warp_iterator_from_smem.h" + +template +struct TransposeWarpIterator { + using Iterator = char; + static bool constexpr kSupportsTranspose = false; +}; + +template < + /// Operand identity + cutlass::gemm::Operand Operand, + /// Data type of A elements + typename Element, + bool kTranspose> +struct TransposeWarpIterator< + cutlass::gemm::warp::WarpIteratorFromSmem> { + using Iterator = cutlass::gemm::warp::WarpIteratorFromSmem; + static bool constexpr kSupportsTranspose = true; +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h new file mode 100644 index 0000000000000000000000000000000000000000..7dd59832b4b0f278408552de38849b676d6b668a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/deepspeed4science/evoformer_attn/iterators/warp_iterator_from_smem.h @@ -0,0 +1,269 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights + *reserved. SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/*! \file + \brief Inspired from + "cutlass/gemm/warp/mma_tensor_op_tile_access_iterator.h" Loads tiles of GEMM + operands from a RowMajor shared-memory layout into registers to use by A100 + TensorCores. + + The difference with "mma_tensor_op_tile_access_iterator.h" is that: + (1) We use "ldmatrix" to load tiles, rather than manual loads (slightly + faster) (2) We support to transpose the operand (eg read `A.transpose()` when + the shared memory holds `A`) + + This is only implemented for the specific shapes. +*/ +#pragma once + +#include + +//////////////////////////////////////////////////////////////////////////////// +namespace cutlass { +namespace gemm { +namespace warp { + +template < + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Element_, + bool kTranspose = false> +class WarpIteratorFromSmem { +public: + /// Shape of tile to load (concept: MatrixShape) + using Shape = cutlass::MatrixShape<32, 32>; + + /// Operand tag + static Operand const kOperand = Operand_; + + /// Basic check + static_assert( + kOperand == Operand::kA || kOperand == Operand::kB, + "WarpIteratorFromSmem may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + static_assert(sizeof_bits::value == 16, "Only supported for half"); + + /// Layout of source tile + using Layout = cutlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = cutlass::MatrixShape<16, 8>; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = 1; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = + (sizeof_bits::value >= 32 ? 1 : 32 / sizeof_bits::value); + + using InstructionCount = MatrixShape; + + static int const kIterations = (kOperand == Operand::kA) ? InstructionCount::kColumn + : InstructionCount::kRow; + +public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + + /// Memory access type + // using AccessType = AlignedArray; + using AccessType = Array; + + static int constexpr kWarpShapeDivisibleInner = + (kOperand == Operand::kA ? InstructionShape::kColumn : InstructionShape::kRow); + static int constexpr kAccessesInner = (kWarpShapeDivisibleInner / kElementsPerAccess) / 4; + static int const kTilesPerInstruction = InstructionShape::kRow / 8; + +private: + /// Underlying tensor reference + TensorRef ref_; + + /// Origin + MatrixCoord origin_; + + /// Iterations in a tile + int iterations_; + +public: + /// Constructor from TensorRef + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem(TensorRef const& ref, int lane_id) + : WarpIteratorFromSmem(ref, {Shape::kRow, Shape::kColumn}, lane_id) + { + } + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem(TensorRef const& ref, TensorCoord extent, int lane_id) + : ref_(ref), iterations_(0) + { + int ldsm_vec_num = (lane_id >> 3); + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id % 8, 0); + static_assert(InstructionCount::kRow * kAccessesInner * kTilesPerInstruction == 4, ""); + CUTLASS_PRAGMA_UNROLL + for (int inst_m_idx = 0; inst_m_idx < InstructionCount::kRow; ++inst_m_idx) { + CUTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + CUTLASS_PRAGMA_UNROLL + for (int access_m_idx = 0; access_m_idx < kTilesPerInstruction; + ++access_m_idx) { + int access_idx = + access_m_idx + + kTilesPerInstruction * (inner_idx + kAccessesInner * inst_m_idx); + + MatrixCoord offset(access_m_idx * 8 + inst_m_idx * InstructionShape::kRow, + inner_idx * 4 * kElementsPerAccess); + + if (access_idx == ldsm_vec_num) { + if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } + origin_ += offset; + } + } + } + } + } else { + origin_ = MatrixCoord(0, lane_id % 8); + static_assert(InstructionCount::kColumn * kAccessesInner == 4, ""); + CUTLASS_PRAGMA_UNROLL + for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; ++inst_n_idx) { + CUTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + int access_idx = inner_idx + kAccessesInner * inst_n_idx; + + MatrixCoord offset(inner_idx * 4 * kElementsPerAccess, inst_n_idx * 8); + + if (access_idx == ldsm_vec_num) { + if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } + origin_ += offset; + } + } + } + } + + ref_.add_coord_offset(origin_); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem& add_tile_offset(TensorCoord const& tile_offset) + { + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, + tile_offset.column() * Shape::kColumn); + if (kTranspose) { coord_offset = TensorCoord{coord_offset.column(), coord_offset.row()}; } + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + CUTLASS_DEVICE + void advance() + { + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } else { + add_tile_offset({1, 0}); + } + + iterations_ = 0; + } + + /// increase iterations in a tile + CUTLASS_HOST_DEVICE + WarpIteratorFromSmem& operator++() + { + iterations_++; + + if (iterations_ >= kIterations) advance(); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + CUTLASS_DEVICE + void load(Fragment& frag) const + { + AccessType* access_ptr = reinterpret_cast(&frag); + using LoadLayout = + typename platform::conditional::type; + + MatrixCoord offset; + if (kOperand == Operand::kA) { + offset = MatrixCoord(0, iterations_ * InstructionShape::kColumn); + } else { + offset = MatrixCoord(iterations_ * InstructionShape::kRow, 0); + } + if (kTranspose) { offset = MatrixCoord(offset.column(), offset.row()); } + cutlass::arch::ldsm(access_ptr[0], ref_.data() + ref_.offset(offset)); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace cutlass +//////////////////////////////////////////////////////////////////////////////// diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp new file mode 100644 index 0000000000000000000000000000000000000000..1a887b50e1a3d06c7e6f554b6df4399de071c712 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cpp @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "fp_quantize.h" + +#include +#include +#include + +#define DISPATCH_QUANTIZE(T_TYPE, C_TYPE, mantisa, exponent) \ + if (val.options().dtype() == torch::T_TYPE) { \ + launch_quantization((C_TYPE*)val.data_ptr(), \ + (uint8_t*)out.data_ptr(), \ + num_groups, \ + group_size, \ + at::cuda::getCurrentCUDAStream(), \ + q_range, \ + q_bits, \ + q_mantisa_bits, \ + stochastic_rounding); \ + } + +at::Tensor quantize(torch::Tensor& out, + torch::Tensor& val, + int group_size, + int stochastic_rounding, + int q_bits, + int q_mantisa_bits) +{ + int total_elems = at::numel(val); + float q_range = q_bits == 8 ? (q_mantisa_bits == 3 ? 480.0 : 114688.0) : // fp8 ranges + (q_bits == 12 ? 510.0 : // fp12 range + (q_bits == 6 ? 28.0 : // fp6 range + 6.0)); // fp4 range (using power 2); TODO (Reza): add the power-4 + // in case accuracy is not matching! + int num_groups = total_elems / group_size; + + DISPATCH_QUANTIZE(kHalf, __half, 23, 8); +#ifdef BF16_AVAILABLE + DISPATCH_QUANTIZE(kBFloat16, __nv_bfloat16, 23, 8); +#endif + + return out; +} + +#define DISPATCH_DEQUANTIZE(T_TYPE, C_TYPE, mantisa) \ + if (val.options().dtype() == torch::T_TYPE) { \ + launch_dequantization((uint8_t*)val_q.data_ptr(), \ + (C_TYPE*)val.data_ptr(), \ + num_groups, \ + group_size, \ + q_mantisa_bits, \ + q_exponent_bits, \ + at::cuda::getCurrentCUDAStream()); \ + return; \ + } + +void dequantize(torch::Tensor& val, + torch::Tensor& val_q, + int group_size, + int q_mantisa_bits, + int q_exponent_bits) +{ + int total_elems = at::numel(val); + + int num_groups = total_elems / group_size; + + DISPATCH_DEQUANTIZE(kHalf, __half, 10); +#ifdef BF16_AVAILABLE + DISPATCH_DEQUANTIZE(kBFloat16, __nv_bfloat16, 7); +#endif +} + +#define DISPATCH_DEQUANTIZE_INDEX(T_TYPE, C_TYPE, mantisa) \ + if (val.options().dtype() == torch::T_TYPE) { \ + launch_selective_dequantization((uint8_t*)val_q.data_ptr(), \ + (C_TYPE*)val.data_ptr(), \ + (int32_t*)indexes.data_ptr(), \ + num_groups, \ + group_size, \ + num_indexes, \ + q_mantisa_bits, \ + q_exponent_bits, \ + at::cuda::getCurrentCUDAStream()); \ + return; \ + } +void selective_dequantize(torch::Tensor& val, + torch::Tensor& val_q, + torch::Tensor& indexes, + int group_size, + int q_mantisa_bits, + int q_exponent_bits) +{ + int total_elems = at::numel(val); + int num_indexes = indexes.size(0); + int num_groups = total_elems / group_size; + + DISPATCH_DEQUANTIZE_INDEX(kHalf, __half, 10); +#ifdef BF16_AVAILABLE + DISPATCH_DEQUANTIZE_INDEX(kBFloat16, __nv_bfloat16, 7); +#endif +} + +at::Tensor get_scales(torch::Tensor& out, int num_groups) +{ + auto options = at::TensorOptions() + .dtype(torch::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + auto scales = + torch::from_blob(out.data_ptr(), {num_groups, 1}, {out.stride(0) / 4, 1}, options); + return scales; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("quantize", &quantize, "quantize function"); + m.def("dequantize", &dequantize, "dequantize function"); + m.def("get_scales", &get_scales, "get scales function"); + m.def("selective_dequantize", &selective_dequantize, "selective dequantize function"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..66ea7392e0115ab9ebbeb66baad01f9d6fa80750 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/fp_quantize.cu @@ -0,0 +1,532 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "context.h" +#include "fp_quantize.h" +#include "memory_access_utils.h" +#include "reduction_utils.h" + +#include +#include + +#include +#include + +#ifdef BF16_AVAILABLE +#include +#endif +#include + +using ROp = reduce::ROpType; + +namespace quantization { + +constexpr int access_granularity = 16; +constexpr int quanitzed_access_granularity = 4; +constexpr int quanitzed_access_granularity_6bits = 2; +constexpr int threads = 256; +constexpr int warps = threads / 32; + +} // namespace quantization + +template +__device__ void round(uint32_t& mantisa, uint32_t& dst_exponent, curandStatePhilox4_32_10_t* state) +{ + constexpr uint32_t mantisa_mask = (1 << (_mantisa_bits - q_mantisa_bits)) - 1; + uint32_t offset = stochastic_rounding ? (curand_poisson(state, 10) & mantisa_mask) + : 1 << (_mantisa_bits - q_mantisa_bits - 1); + mantisa += offset; + dst_exponent += (((mantisa & ~mantisa_mask) == (1 << _mantisa_bits)) ? 1 : 0); +} + +template +__device__ void clip(uint32_t& exponent, uint32_t& mantisa) +{ + constexpr uint32_t max_exponent = (1 << (q_exponent_bits - 1)) + (1 << (_exponent_bits - 1)); + constexpr uint32_t min_exponent = + (1 << (_exponent_bits - 1)) - ((1 << (q_exponent_bits - 1)) - 1); + if (exponent > max_exponent) { + exponent = max_exponent; + mantisa = (((uint32_t)-1) >> (32 - q_mantisa_bits)) << 1; //.11 .. 10 + } + if (exponent < min_exponent) { + exponent = min_exponent; + mantisa = 0; + } +} + +template +__global__ void apply_quantization(T* val, + uint8_t* q_val, + int group_size, + std::pair seed, + float q_range) +{ + int tidx = threadIdx.x; + int wid = tidx >> 5; + int lane = tidx & 0x1f; + int gid = blockIdx.x * quantization::warps + wid; + + constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; + constexpr uint32_t _mantisa_mask = (1 << _mantisa_bits) - 1; + constexpr uint32_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; + constexpr uint32_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); + // CG helpers + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); + constexpr uint32_t load_stride = vector_size * hw_warp_size; + constexpr uint32_t store_stride = (total_q_bits * vector_size / 8) * hw_warp_size; + const uint32_t thread_offset = lane * vector_size; + const uint32_t store_thread_offset = lane * (total_q_bits * vector_size / 8); + const uint32_t base_load_offset = gid * group_size + thread_offset; + const uint32_t base_store_offset = + gid * ((group_size * total_q_bits / 8) + 4) + + store_thread_offset; // 4-byte for saving the scale per group + const T* load_base_ptr = val + base_load_offset; + T tmp_buf[unroll * vector_size]; + T cur_max; + reduce::init(&cur_max); + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + +#pragma unroll + for (int i = 0; i < unroll; i++) { + if (i * load_stride + thread_offset < group_size) { + mem_access::load_global( + &tmp_buf[vector_size * i], load_base_ptr + i * load_stride); + for (int j = 0; j < vector_size; j++) + cur_max = reduce::element(cur_max, __habs(tmp_buf[i * vector_size + j])); + } + } + reduce::_block(tb, warp, &cur_max); + + int mantisa_mask = ((1 << q_mantisa_bits) - 1); + mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); + + uint8_t* store_base_ptr = q_val + base_store_offset; + float scale = (float)q_range / conversion::to(cur_max); +#pragma unroll + for (int i = 0; i < unroll; i++) { + if (i * load_stride + thread_offset < group_size) { + uint64_t q_buf = 0; + uint64_t q_buf1 = 0; +#pragma unroll + for (int j = 0; j < vector_size; j++) { + float val_f = conversion::to(tmp_buf[i * vector_size + j]) * scale; + uint32_t* data = reinterpret_cast(&val_f); + uint32_t sign = (data[0] & _sign_mask) >> (_mantisa_bits + _exponent_bits); + uint32_t cur_exponent = (data[0] & _exponent_mask) >> _mantisa_bits; + uint32_t dst_mantisa = (data[0] & _mantisa_mask); + + uint32_t dst_exponent = cur_exponent; + + round<_mantisa_bits, q_mantisa_bits, stochastic_rounding>( + dst_mantisa, dst_exponent, &state); + if (cur_exponent != 0) + clip<_mantisa_bits, _exponent_bits, q_mantisa_bits, q_exponent_bits>( + dst_exponent, dst_mantisa); + + dst_mantisa = (dst_mantisa & mantisa_mask) >> (_mantisa_bits - q_mantisa_bits); + + if (dst_exponent != (1 << q_exponent_bits) - 1) + dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + + (1 << (q_exponent_bits - 1)) - 1; + if (total_q_bits == 8 || total_q_bits == 4 || total_q_bits == 6) + q_buf = q_buf | + ((uint64_t)((uint8_t)(sign << (q_exponent_bits + q_mantisa_bits) | + (dst_exponent << q_mantisa_bits) | dst_mantisa)) + << j * total_q_bits); + else if (total_q_bits == 12) { + if (j < 5) + q_buf = + q_buf | + ((uint64_t)((uint16_t)(sign << (q_exponent_bits + q_mantisa_bits) | + (dst_exponent << q_mantisa_bits) | dst_mantisa)) + << j * total_q_bits); + else + q_buf1 = + q_buf1 | + ((uint64_t)((uint16_t)(sign << (q_exponent_bits + q_mantisa_bits) | + (dst_exponent << q_mantisa_bits) | dst_mantisa)) + << (j - 5) * total_q_bits); + } + } + if (total_q_bits == 12) { + uint64_t last_nibble_mask = 0xf; + last_nibble_mask = q_buf1 & last_nibble_mask; + q_buf = (last_nibble_mask << 60) | q_buf; + q_buf1 >>= 4; + } + uint8_t* int8_data = reinterpret_cast(&q_buf); + uint8_t* int8_data1 = reinterpret_cast(&q_buf1); + if (total_q_bits == 6) { + mem_access::store_global( + store_base_ptr + i * store_stride, int8_data); + mem_access::store_global( + store_base_ptr + i * store_stride + + quantization::quanitzed_access_granularity_6bits, + int8_data + quantization::quanitzed_access_granularity_6bits); + mem_access::store_global( + store_base_ptr + i * store_stride + + quantization::quanitzed_access_granularity_6bits * 2, + int8_data + 2 * quantization::quanitzed_access_granularity_6bits); + } else { + mem_access::store_global( + store_base_ptr + i * store_stride, int8_data); + + if (total_q_bits > 4) { + mem_access::store_global( + store_base_ptr + i * store_stride + + quantization::quanitzed_access_granularity, + int8_data + quantization::quanitzed_access_granularity); + if (total_q_bits == 12) { + mem_access::store_global( + store_base_ptr + i * store_stride + + quantization::quanitzed_access_granularity * 2, + int8_data1); + } + } + } + } + } + if (lane == 0) { + float q_scale = conversion::to(cur_max) / (float)q_range; + uint8_t* scale_as_int8 = reinterpret_cast(&q_scale); + uint32_t scale_offset = + gid * ((group_size * total_q_bits / 8) + 4) + (group_size * total_q_bits / 8); + if (total_q_bits != 6) + mem_access::store_global( + q_val + scale_offset, scale_as_int8); + else { + mem_access::store_global( + q_val + scale_offset, scale_as_int8); + mem_access::store_global( + q_val + scale_offset + quantization::quanitzed_access_granularity_6bits, + scale_as_int8 + quantization::quanitzed_access_granularity_6bits); + } + } +} + +template +__global__ void apply_dequantization(uint8_t* val, T* q_val, int group_size, int total_num_elements) +{ + constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); + int tidx = (blockIdx.x * blockDim.x + threadIdx.x) * vector_size; + + constexpr int quantized_bits = _mantisa_bits + _exponent_bits + 1; + constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; + constexpr uint16_t _mantisa_mask = (1 << _mantisa_bits) - 1; + constexpr uint16_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; + constexpr uint16_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); + const uint32_t g_index = (tidx / group_size); + const uint32_t group_size_bytes = (group_size * quantized_bits / 8); + const uint8_t* load_base_ptr = + val + g_index * (group_size_bytes + 4) + (tidx % group_size) * quantized_bits / 8; + + int mantisa_mask = ((1 << q_mantisa_bits) - 1); + mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); + + T* store_base_ptr = q_val + tidx; + float scale; + + uint8_t* scale_as_int8 = reinterpret_cast(&scale); + if (quantized_bits == 6) { + mem_access::load_global( + scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); + mem_access::load_global( + scale_as_int8 + quantization::quanitzed_access_granularity_6bits, + val + g_index * (group_size_bytes + 4) + group_size_bytes + + quantization::quanitzed_access_granularity_6bits); + } else + mem_access::load_global( + scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); + + if (tidx < total_num_elements) { + uint64_t q_buf_in; + uint64_t q_buf_in1; + uint8_t* int8_data = reinterpret_cast(&q_buf_in); + uint8_t* int8_data1 = reinterpret_cast(&q_buf_in1); + if (quantized_bits == 6) { + mem_access::load_global( + int8_data, load_base_ptr); + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity_6bits, + load_base_ptr + quantization::quanitzed_access_granularity_6bits); + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity_6bits * 2, + load_base_ptr + quantization::quanitzed_access_granularity_6bits * 2); + + } else { + mem_access::load_global(int8_data, + load_base_ptr); + if (quantized_bits > 4) { + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity, + load_base_ptr + quantization::quanitzed_access_granularity); + if (quantized_bits == 12) { + mem_access::load_global( + int8_data1, load_base_ptr + quantization::quanitzed_access_granularity * 2); + } + } + } + T store_buf[vector_size]; + uint16_t* q_buf = reinterpret_cast(store_buf); +#pragma unroll + for (int j = 0; j < vector_size; j++) { + uint16_t new_data; + if (j < 5 || quantized_bits != 12) { + new_data = (uint16_t)(q_buf_in >> (j * quantized_bits)); + } else { + if (j == 5) { + new_data = (uint16_t)(q_buf_in1); + new_data = (uint16_t)((new_data << 4) | (q_buf_in >> 60)); + } else + new_data = (uint16_t)(q_buf_in1 >> ((j - 6) * quantized_bits + 8)); + } + + uint16_t sign = (new_data & _sign_mask) >> (_mantisa_bits + _exponent_bits); + uint16_t dst_exponent = (new_data & _exponent_mask) >> _mantisa_bits; + uint16_t dst_mantisa = (new_data & _mantisa_mask); + + if (dst_exponent != (1 << q_exponent_bits) - 1) + dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + + (1 << (q_exponent_bits - 1)) - 1; + + q_buf[j] = + ((sign << (q_exponent_bits + q_mantisa_bits)) | (dst_exponent << q_mantisa_bits) | + (dst_mantisa << (q_mantisa_bits - _mantisa_bits))); + float up_cast = conversion::to(store_buf[j]); + store_buf[j] = conversion::to(up_cast * scale); + } + mem_access::store_global(store_base_ptr, store_buf); + } +} + +#define LAUNCH_FOR_QUANTIZATION_UNROLL(COUNT) \ + case COUNT: \ + apply_quantization \ + <<>>(val, q_val, group_size, seed, q_range); \ + break; + +template +void launch_quantization(T* val, + uint8_t* q_val, + int num_groups, + int group_size, + cudaStream_t stream, + float q_range, + int q_bits, + int q_mantisa_bits, + int stochastic_rounding) +{ + const dim3 grid((num_groups + quantization::warps - 1) / quantization::warps); + const dim3 block(quantization::threads); + + std::pair seed = FPContext::Instance().IncrementOffset(16); + + constexpr int vals_per_unroll = hw_warp_size * quantization::access_granularity / sizeof(T); + + const int copy_unroll = (group_size + vals_per_unroll - 1) / vals_per_unroll; + QUANT_SWITCH((q_bits - q_mantisa_bits - 1) * q_mantisa_bits + stochastic_rounding, [&] { + switch (copy_unroll) { + LAUNCH_FOR_QUANTIZATION_UNROLL(1) + LAUNCH_FOR_QUANTIZATION_UNROLL(2) + LAUNCH_FOR_QUANTIZATION_UNROLL(3) + LAUNCH_FOR_QUANTIZATION_UNROLL(4) + LAUNCH_FOR_QUANTIZATION_UNROLL(5) + LAUNCH_FOR_QUANTIZATION_UNROLL(6) + } + }); +} +#define INSTANTIATE_LAUNCH_QUANTIZATION(T, mantisa, exponent) \ + template void launch_quantization( \ + T*, uint8_t*, int, int, cudaStream_t, float q_range, int, int, int); +// fp8(E4M3), nearest-rounding +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_QUANTIZATION(__nv_bfloat16, 23, 8); +#endif +INSTANTIATE_LAUNCH_QUANTIZATION(__half, 23, 8); + +template +void launch_dequantization(uint8_t* val, + T* q_val, + int num_groups, + int group_size, + int q_mantisa_bits, + int q_exponent_bits, + cudaStream_t stream) +{ + int blocks = ((num_groups * group_size) - 1) / + (quantization::threads * (quantization::access_granularity / sizeof(T))) + + 1; + const dim3 grid(blocks); + const dim3 block(quantization::threads); + DEQUANT_SWITCH(q_mantisa_bits * q_exponent_bits, [&] { + apply_dequantization + <<>>(val, q_val, group_size, (num_groups * group_size)); + }); +} +#define INSTANTIATE_LAUNCH_DEQUANTIZATION(T, mantisa) \ + template void launch_dequantization(uint8_t*, T*, int, int, int, int, cudaStream_t); +// fp8(E4M3) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_DEQUANTIZATION(__nv_bfloat16, 7); +#endif +INSTANTIATE_LAUNCH_DEQUANTIZATION(__half, 10); + +template +__global__ void apply_selective_dequantization(uint8_t* val, + T* q_val, + int32_t* indexes, + int group_size, + int total_num_elements) +{ + int index = indexes[blockIdx.x]; + constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T); + int tidx = (blockIdx.y * blockDim.x + threadIdx.x) * vector_size; + int input_index = index * total_num_elements + tidx; + constexpr int quantized_bits = _mantisa_bits + _exponent_bits + 1; + constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1; + constexpr uint16_t _mantisa_mask = (1 << _mantisa_bits) - 1; + constexpr uint16_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits; + constexpr uint16_t _sign_mask = 1 << (_mantisa_bits + _exponent_bits); + const uint32_t g_index = (input_index / group_size); + const uint32_t group_size_bytes = (group_size * quantized_bits / 8); + const uint8_t* load_base_ptr = + val + g_index * (group_size_bytes + 4) + (input_index % group_size) * quantized_bits / 8; + + int mantisa_mask = ((1 << q_mantisa_bits) - 1); + mantisa_mask <<= (_mantisa_bits - q_mantisa_bits); + + T* store_base_ptr = q_val + tidx + blockIdx.x * total_num_elements; + float scale; + + uint8_t* scale_as_int8 = reinterpret_cast(&scale); + if (quantized_bits == 6) { + mem_access::load_global( + scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); + mem_access::load_global( + scale_as_int8 + quantization::quanitzed_access_granularity_6bits, + val + g_index * (group_size_bytes + 4) + group_size_bytes + + quantization::quanitzed_access_granularity_6bits); + } else + mem_access::load_global( + scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes); + + if (tidx < total_num_elements) { + uint64_t q_buf_in; + uint64_t q_buf_in1; + uint8_t* int8_data = reinterpret_cast(&q_buf_in); + uint8_t* int8_data1 = reinterpret_cast(&q_buf_in1); + if (quantized_bits == 6) { + mem_access::load_global( + int8_data, load_base_ptr); + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity_6bits, + load_base_ptr + quantization::quanitzed_access_granularity_6bits); + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity_6bits * 2, + load_base_ptr + quantization::quanitzed_access_granularity_6bits * 2); + } else { + mem_access::load_global(int8_data, + load_base_ptr); + if (quantized_bits > 4) { + mem_access::load_global( + int8_data + quantization::quanitzed_access_granularity, + load_base_ptr + quantization::quanitzed_access_granularity); + if (quantized_bits == 12) { + mem_access::load_global( + int8_data1, load_base_ptr + quantization::quanitzed_access_granularity * 2); + } + } + } + T store_buf[vector_size]; + uint16_t* q_buf = reinterpret_cast(store_buf); +#pragma unroll + for (int j = 0; j < vector_size; j++) { + uint16_t new_data; + if (j < 5 || quantized_bits != 12) { + new_data = (uint16_t)(q_buf_in >> (j * quantized_bits)); + } else { + if (j == 5) { + new_data = (uint16_t)(q_buf_in1); + new_data = (uint16_t)((new_data << 4) | (q_buf_in >> 60)); + } else + new_data = (uint16_t)(q_buf_in1 >> ((j - 6) * quantized_bits + 8)); + } + + uint16_t sign = (new_data & _sign_mask) >> (_mantisa_bits + _exponent_bits); + uint16_t dst_exponent = (new_data & _exponent_mask) >> _mantisa_bits; + uint16_t dst_mantisa = (new_data & _mantisa_mask); + + if (dst_exponent != (1 << q_exponent_bits) - 1) + dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) + + (1 << (q_exponent_bits - 1)) - 1; + + q_buf[j] = + ((sign << (q_exponent_bits + q_mantisa_bits)) | (dst_exponent << q_mantisa_bits) | + (dst_mantisa << (q_mantisa_bits - _mantisa_bits))); + float up_cast = conversion::to(store_buf[j]); + store_buf[j] = conversion::to(up_cast * scale); + } + mem_access::store_global(store_base_ptr, store_buf); + } +} + +template +void launch_selective_dequantization(uint8_t* val, + T* q_val, + int32_t* indexes, + int num_groups, + int group_size, + int num_indexes, + int q_mantisa_bits, + int q_exponent_bits, + cudaStream_t stream) +{ + int total_elements_per_index = (num_groups / num_indexes) * group_size; + int blocks = (total_elements_per_index - 1) / + (quantization::threads * (quantization::access_granularity / sizeof(T))) + + 1; + const dim3 grid(num_indexes, blocks); + const dim3 block(quantization::threads); + DEQUANT_SWITCH(q_mantisa_bits * q_exponent_bits, [&] { + apply_selective_dequantization + <<>>(val, q_val, indexes, group_size, total_elements_per_index); + }); +} +#define INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(T, mantisa) \ + template void launch_selective_dequantization( \ + uint8_t*, T*, int32_t*, int, int, int, int, int, cudaStream_t); +// fp8(E4M3) +#ifdef BF16_AVAILABLE +INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(__nv_bfloat16, 7); +#endif +INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(__half, 10); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h new file mode 100644 index 0000000000000000000000000000000000000000..5bd9badbcb4fa06091a8c01b0b0d386153a24589 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/context.h @@ -0,0 +1,66 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include +#include +#include +#include +#include "cublas_v2.h" +#include "cuda.h" +#include "curand.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#define WARP_SIZE 32 + +class FPContext { +public: + FPContext() : _seed(42) + { + curandCreateGenerator(&_gen, CURAND_RNG_PSEUDO_DEFAULT); + curandSetPseudoRandomGeneratorSeed(_gen, 123); + } + + virtual ~FPContext() {} + + static FPContext& Instance() + { + static FPContext _ctx; + return _ctx; + } + + curandGenerator_t& GetRandGenerator() { return _gen; } + + cudaStream_t GetCurrentStream() + { + // get current pytorch stream. + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + return stream; + } + + std::pair IncrementOffset(uint64_t offset_inc) + { + uint64_t offset = _curr_offset; + _curr_offset += offset_inc; + return std::pair(_seed, offset); + } + + void SetSeed(uint64_t new_seed) { _seed = new_seed; } + +private: + curandGenerator_t _gen; + cublasHandle_t _cublasHandle; + uint64_t _seed; + uint64_t _curr_offset; +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h new file mode 100644 index 0000000000000000000000000000000000000000..60c75541f603f8b58558ecdf1d5658407f4d2857 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/fp_quantizer/includes/fp_quantize.h @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#include +#include + +#include + +#ifdef BF16_AVAILABLE +#include +#endif +#include +#include + +#define QUANT_SWITCH(Q_BITS, ...) \ + [&] { \ + if (12 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 8; \ + constexpr int CONST_Q_MANTISA_BITS = 3; \ + __VA_ARGS__(); \ + } else if (13 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 8; \ + constexpr int CONST_Q_MANTISA_BITS = 3; \ + __VA_ARGS__(); \ + } else if (10 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 8; \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + __VA_ARGS__(); \ + } else if (11 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 8; \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + __VA_ARGS__(); \ + } else if (28 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 12; \ + constexpr int CONST_Q_MANTISA_BITS = 7; \ + __VA_ARGS__(); \ + } else if (29 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 12; \ + constexpr int CONST_Q_MANTISA_BITS = 7; \ + __VA_ARGS__(); \ + } else if (6 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 6; \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + __VA_ARGS__(); \ + } else if (7 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 6; \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + __VA_ARGS__(); \ + } else if (2 == Q_BITS) { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 0; \ + constexpr int CONST_Q_BITS = 4; \ + constexpr int CONST_Q_MANTISA_BITS = 1; \ + __VA_ARGS__(); \ + } else { \ + constexpr int CONST_STOCHASTIC_ROUNDING = 1; \ + constexpr int CONST_Q_BITS = 4; \ + constexpr int CONST_Q_MANTISA_BITS = 1; \ + __VA_ARGS__(); \ + } \ + }() + +#define DEQUANT_SWITCH(Q_MANTISA_EXPONENT_BITS, ...) \ + [&] { \ + if (12 == Q_MANTISA_EXPONENT_BITS) { \ + constexpr int CONST_Q_MANTISA_BITS = 3; \ + constexpr int CONST_Q_EXPONENT_BITS = 4; \ + __VA_ARGS__(); \ + } else if (10 == Q_MANTISA_EXPONENT_BITS) { \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + constexpr int CONST_Q_EXPONENT_BITS = 5; \ + __VA_ARGS__(); \ + } else if (28 == Q_MANTISA_EXPONENT_BITS) { \ + constexpr int CONST_Q_MANTISA_BITS = 7; \ + constexpr int CONST_Q_EXPONENT_BITS = 4; \ + __VA_ARGS__(); \ + } else if (6 == Q_MANTISA_EXPONENT_BITS) { \ + constexpr int CONST_Q_MANTISA_BITS = 2; \ + constexpr int CONST_Q_EXPONENT_BITS = 3; \ + __VA_ARGS__(); \ + } else { \ + constexpr int CONST_Q_MANTISA_BITS = 1; \ + constexpr int CONST_Q_EXPONENT_BITS = 2; \ + __VA_ARGS__(); \ + } \ + }() + +template +void launch_quantization(T* val, + uint8_t* q_val, + int num_groups, + int group_size, + cudaStream_t stream, + float q_range, + int q_bits, + int q_mantisa_bits, + int stochastic_rounding); + +template +void launch_dequantization(uint8_t* val, + T* q_val, + int num_groups, + int group_size, + int q_mantisa_bits, + int q_exponent_bits, + cudaStream_t stream); + +template +void launch_selective_dequantization(uint8_t* val, + T* q_val, + int32_t* indexes, + int num_groups, + int group_size, + int num_indexes, + int q_mantisa_bits, + int q_exponent_bits, + cudaStream_t stream); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp new file mode 100644 index 0000000000000000000000000000000000000000..b7055c8cc72b36df99a8118d8b8085b3bfba11ab --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.cpp @@ -0,0 +1,161 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include "deepspeed_gds_op.h" + +using namespace std; + +// For when there is more than 1 device +static std::map> base_ptr_registry; + +static void _safe_handle_register(const int fd, CUfileDescr_t& cf_descr, CUfileHandle_t& cf_handle) +{ + memset((void*)&cf_descr, 0, sizeof(CUfileDescr_t)); + cf_descr.handle.fd = fd; + cf_descr.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD; + CUfileError_t status = cuFileHandleRegister(&cf_handle, &cf_descr); + if (status.err != CU_FILE_SUCCESS) { + std::cerr << "file register error:" << cuFileGetErrorString(status) << std::endl; + close(fd); + exit(EXIT_FAILURE); + } +} + +static void* _find_base_ptr(const int64_t device, char* buf_ptr) +{ + void* base_ptr = nullptr; + int64_t last = -1; + int64_t ptr_diff; + for (const auto& value : base_ptr_registry[device]) { + ptr_diff = buf_ptr - (char*)value; + if (last == -1 && ptr_diff >= 0) { + last = ptr_diff; + base_ptr = value; + } else if (ptr_diff < last && ptr_diff >= 0) { + last = ptr_diff; + base_ptr = value; + } + } + if (!base_ptr || buf_ptr < base_ptr) { + std::cerr << "BASE PTR ERROR :" << base_ptr << " BUF PTR " << (void*)buf_ptr << std::endl; + for (const auto& value : base_ptr_registry[device]) { + std::cerr << "BASE PTR AVAIL :" << value << std::endl; + } + exit(EXIT_FAILURE); + } + + return base_ptr; +} + +void gds_op_desc_t::add_buffer_to_registry(const torch::Tensor& buffer) +{ + const int64_t device = buffer.get_device(); + void* reg_ptr = buffer.data_ptr(); + + // TODO: add checking to make sure pointer isn't already in set + const auto it = base_ptr_registry.find(device); + if (it == base_ptr_registry.end()) { + std::set new_ptr_set; + new_ptr_set.insert(reg_ptr); + base_ptr_registry.insert(std::pair>(device, new_ptr_set)); + } else { + base_ptr_registry[device].insert(reg_ptr); + } + + check_cudaruntimecall(cudaSetDevice(device)); + CUfileError_t status = cuFileBufRegister(reg_ptr, buffer.nbytes(), 0); + if (status.err != CU_FILE_SUCCESS) { + std::cerr << "buffer register failed:" << cuFileGetErrorString(status) << std::endl; + exit(EXIT_FAILURE); + } +} + +void gds_op_desc_t::remove_buffer_from_registry(const torch::Tensor& buffer) +{ + const int64_t device = buffer.get_device(); + void* reg_ptr = buffer.data_ptr(); + + // std::cout << "DEREG PTR " << reg_ptr << std::endl; + check_cudaruntimecall(cudaSetDevice(device)); + cuFileBufDeregister(reg_ptr); + + // Remove from tracked registry + base_ptr_registry[device].erase(reg_ptr); +} + +gds_op_desc_t::gds_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) + : io_op_desc_t(read_op, + buffer, + fd, + filename, + file_num_bytes, + intra_op_parallelism, + validate, + file_offset) +{ + _contiguous_buffer = _buffer.contiguous(); + const int64_t device = _buffer.get_device(); + check_cudaruntimecall(cudaSetDevice(device)); + _base_ptr = _find_base_ptr(device, (char*)_contiguous_buffer.data_ptr()); + + _safe_handle_register(fd, _cf_descr, _cf_handle); +} + +char* gds_op_desc_t::data_ptr() const { return (char*)_contiguous_buffer.data_ptr(); } + +void gds_op_desc_t::finish() { cuFileHandleDeregister(_cf_handle); } + +void gds_op_desc_t::validate() +{ + check_cudaruntimecall(cudaSetDevice(_buffer.get_device())); + const auto cpu_buffer = _buffer.to(torch::kCPU); + validate_aio_operation( + _read_op, _filename.c_str(), (char*)(cpu_buffer.data_ptr()), _file_num_bytes); +} + +void gds_op_desc_t::run(const int tid, + std::unique_ptr& aio_ctxt, + deepspeed_aio_config_t* aio_config) +{ + assert(tid < _intra_op_parallelism); + check_cudaruntimecall(cudaSetDevice(_buffer.get_device())); + const auto buf_offset = data_ptr() + (_num_bytes_per_thread * tid) - (char*)_base_ptr; + const auto tid_file_offset = _file_offset + (_num_bytes_per_thread * tid); + + if (_read_op) { + auto ret = + cuFileRead(_cf_handle, _base_ptr, _num_bytes_per_thread, tid_file_offset, buf_offset); + if (ret < 0) { _report_error(ret, errno, tid_file_offset); } + } else { + auto ret = + cuFileWrite(_cf_handle, _base_ptr, _num_bytes_per_thread, tid_file_offset, buf_offset); + if (ret < 0) { _report_error(ret, errno, tid_file_offset); } + } +} + +void gds_op_desc_t::_report_error(const ssize_t return_code, + const int error_num, + const off_t offset) +{ + const auto op_string = _read_op ? "read failed with " : "write failed with "; + const auto error_string = IS_CUFILE_ERR(return_code) ? "cuFile error: " : "posix error: "; + const auto error_code = IS_CUFILE_ERR(return_code) ? cuFileGetErrorString(return_code) + : cuFileGetErrorString(error_num); + std::cerr << op_string << error_string << error_code << " return code = " << return_code + << " filename = " << _filename.c_str() << " num bytes = " << _num_bytes_per_thread + << " offset = " << offset << std::endl; + exit(EXIT_FAILURE); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h new file mode 100644 index 0000000000000000000000000000000000000000..d955527b1ba338b82f89b5cb614ee6fc48a1a290 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_op.h @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include +#include + +#include "deepspeed_aio_op_desc.h" +#include "deepspeed_gds_utils.h" + +struct gds_op_desc_t : io_op_desc_t { + CUfileDescr_t _cf_descr; + CUfileHandle_t _cf_handle; + void* _base_ptr; + + gds_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); + + void run(const int tid, + std::unique_ptr& aio_ctxt, + deepspeed_aio_config_t* aio_config); + + char* data_ptr() const; + + void validate(); + + void finish(); + + void _report_error(const ssize_t return_code, const int error_num, const off_t offset); + + static void add_buffer_to_registry(const torch::Tensor& buffer); + + static void remove_buffer_from_registry(const torch::Tensor& buffer); +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..12b014d909880da286ea76b264a72ee15d05d2a7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_gds_utils.h @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +// CUDA/cuFile includes +#include +#include +#include "cufile.h" + +// Macro for checking cuda errors following a cuda launch or api call +#define cudaCheckError() \ + { \ + cudaError_t e = cudaGetLastError(); \ + if (e != cudaSuccess) { \ + printf("Cuda failure %s:%d: '%s'\n", __FILE__, __LINE__, cudaGetErrorString(e)); \ + exit(EXIT_FAILURE); \ + } \ + } + +#define check_cudadrivercall(fn) \ + do { \ + CUresult res = fn; \ + if (res != CUDA_SUCCESS) { \ + const char* str = nullptr; \ + cuGetErrorName(res, &str); \ + std::cerr << "cuda driver api call failed " << #fn << " res : " << res << ", " \ + << __LINE__ << ":" << str << std::endl; \ + std::cerr << "EXITING program!!!" << std::endl; \ + exit(1); \ + } \ + } while (0) + +#define check_cudaruntimecall(fn) \ + do { \ + cudaError_t res = fn; \ + if (res != cudaSuccess) { \ + const char* str = cudaGetErrorName(res); \ + std::cerr << "cuda runtime api call failed " << #fn << __LINE__ << ":" << str \ + << std::endl; \ + std::cerr << "EXITING program!!!" << std::endl; \ + exit(1); \ + } \ + } while (0) + +#define check_cuFileCall(fn, api_msg) \ + do { \ + CUfileError_t status = fn; \ + if (status.err != CU_FILE_SUCCESS) { \ + std::cout << api_msg << " failed with error " << CUFILE_ERRSTR(status.err) \ + << std::endl; \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +// +// cuda driver error description +// +static inline const char* GetCuErrorString(CUresult curesult) +{ + const char* descp; + if (cuGetErrorName(curesult, &descp) != CUDA_SUCCESS) descp = "unknown cuda error"; + return descp; +} + +// +// cuFile APIs return both cuFile specific error codes as well as POSIX error codes +// for ease, the below template can be used for getting the error description depending +// on its type. + +// POSIX +template ::value, std::nullptr_t>::type = nullptr> +std::string cuFileGetErrorString(T status) +{ + status = std::abs(status); + return IS_CUFILE_ERR(status) ? std::string(CUFILE_ERRSTR(status)) + : std::string(std::strerror(status)); +} + +// CUfileError_t +template ::value, std::nullptr_t>::type = nullptr> +std::string cuFileGetErrorString(T status) +{ + std::string errStr = cuFileGetErrorString(static_cast(status.err)); + if (IS_CUDA_ERR(status)) errStr.append(".").append(GetCuErrorString(status.cu_err)); + return errStr; +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp new file mode 100644 index 0000000000000000000000000000000000000000..f11245c75a5e6a4dfdd5c2aa741b40620e8c4a34 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.cpp @@ -0,0 +1,125 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* + GPUDirect Storage functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include "deepspeed_py_gds_handle.h" +#include +#include "deepspeed_gds_op.h" + +using namespace std; + +int deepspeed_gds_handle_t::s_cuFile_init = 0; + +deepspeed_gds_handle_t::deepspeed_gds_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, 1), + _intra_gds_op_parallelism(intra_op_parallelism) +{ + _init_cuFile(block_size, queue_depth); +} + +deepspeed_gds_handle_t::~deepspeed_gds_handle_t() { _close_cuFile(); } + +const int deepspeed_gds_handle_t::get_intra_op_parallelism() const +{ + return _intra_gds_op_parallelism; +} + +void deepspeed_gds_handle_t::_init_cuFile(const int block_size, const int queue_depth) +{ + if (deepspeed_gds_handle_t::s_cuFile_init == 0) { + std::string depthStr = std::to_string(queue_depth); + std::string threadsStr = std::to_string(_intra_gds_op_parallelism); + std::string json1 = R"({"execution": {"max_io_queue_depth": )" + depthStr + ", "; + std::string json2 = R"("max_request_parallelism": )" + threadsStr + ", "; + std::string json3 = R"("max_io_threads": )" + threadsStr + ", "; + std::string json4 = R"("parallel_io": true, "min_io_threshold_size_kb": 8192}})"; + std::ofstream outFile("local_cufile.json"); + if (outFile.is_open()) { + outFile << json1 + json2 + json3 + json4; + outFile.close(); + } else { + std::cerr << "Can't open local cufile" << std::endl; + exit(EXIT_FAILURE); + } + // TODO: Address the following issues with this code + // (1) Fix C++14 warning + // (2) Create file in a different location than PWD + // (3) Handle multi-GPU/multi-rank scenarios: should cufile be shared, is per-rank cufile + // safe? + putenv("CUFILE_ENV_PATH_JSON=$PWD/local_cufile.json"); + cuFileDriverOpen(); + cudaCheckError(); + size_t direct_io_size = (size_t)block_size / 1024; + CUfileError_t status = cuFileDriverSetMaxDirectIOSize(direct_io_size); + if (status.err != CU_FILE_SUCCESS) { + std::cerr << "file register error:" << cuFileGetErrorString(status) << std::endl; + exit(EXIT_FAILURE); + } + } + deepspeed_gds_handle_t::s_cuFile_init++; +} + +void deepspeed_gds_handle_t::_close_cuFile() +{ + deepspeed_gds_handle_t::s_cuFile_init--; + if (deepspeed_gds_handle_t::s_cuFile_init == 0) { cuFileDriverClose(); } +} + +torch::Tensor deepspeed_gds_handle_t::new_pinned_device_tensor(const size_t num_elem, + const torch::Tensor& example_tensor) +{ + auto options = torch::TensorOptions().dtype(example_tensor.scalar_type()).device(torch::kCUDA); + auto dev_tensor = torch::empty(num_elem, options); + pin_device_tensor(dev_tensor); + return dev_tensor; +} + +bool deepspeed_gds_handle_t::free_pinned_device_tensor(torch::Tensor& buffer) +{ + unpin_device_tensor(buffer); + return true; +} + +bool deepspeed_gds_handle_t::pin_device_tensor(const torch::Tensor& buffer) +{ + gds_op_desc_t::add_buffer_to_registry(buffer); + return true; +} + +bool deepspeed_gds_handle_t::unpin_device_tensor(const torch::Tensor& buffer) +{ + gds_op_desc_t::remove_buffer_from_registry(buffer); + return true; +} + +std::shared_ptr deepspeed_gds_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) +{ + if (buffer.is_cuda()) { + return std::make_shared(read_op, + buffer, + fd, + filename, + file_num_bytes, + _intra_op_parallelism, + validate, + file_offset); + } + return deepspeed_io_handle_t::_create_io_op_desc( + read_op, buffer, fd, filename, file_num_bytes, validate, file_offset); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h new file mode 100644 index 0000000000000000000000000000000000000000..25f68e177b2cb2cb731a47e6eb80eb9baddefc13 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/deepspeed_py_gds_handle.h @@ -0,0 +1,49 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include +#include "deepspeed_py_io_handle.h" + +struct deepspeed_gds_handle_t : deepspeed_io_handle_t { + const int _intra_gds_op_parallelism; + + deepspeed_gds_handle_t(const int block_size, + const int queue_depth, + const bool single_submit, + const bool overlap_events, + const int intra_op_parallelism); + + ~deepspeed_gds_handle_t(); + + torch::Tensor new_pinned_device_tensor(const size_t num_elem, + const torch::Tensor& example_tensor); + + bool free_pinned_device_tensor(torch::Tensor&); + + bool pin_device_tensor(const torch::Tensor& buffer); + + bool unpin_device_tensor(const torch::Tensor& buffer); + + void _init_cuFile(const int block_size, const int queue_depth); + + void _close_cuFile(); + + const int get_intra_op_parallelism() const; + + std::shared_ptr _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); + + static int s_cuFile_init; +}; diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp new file mode 100644 index 0000000000000000000000000000000000000000..2f165ee2c32a412c6b2241d97d71ac2422a11c61 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/gds/py_lib/py_ds_gds.cpp @@ -0,0 +1,130 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +Functionality for swapping optimizer tensors to/from (NVMe) storage devices. +*/ + +#include +#include "deepspeed_py_gds_handle.h" +using namespace pybind11::literals; + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + py::class_(m, "gds_handle") + .def(py::init(), + "GDS 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_gds_handle_t::get_block_size) + .def("get_queue_depth", &deepspeed_gds_handle_t::get_queue_depth) + .def("get_single_submit", &deepspeed_gds_handle_t::get_single_submit) + .def("get_overlap_events", &deepspeed_gds_handle_t::get_overlap_events) + .def("get_intra_op_parallelism", &deepspeed_gds_handle_t::get_intra_op_parallelism) + + .def("read", + &deepspeed_gds_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_gds_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_gds_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_gds_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_gds_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_gds_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_gds_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_gds_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_gds_handle_t::new_cpu_locked_tensor, + "Allocate pinned CPU tensor.", + "num_elem"_a, + "example_tenosr"_a) + + .def("free_cpu_locked_tensor", + &deepspeed_gds_handle_t::free_cpu_locked_tensor, + "Free pinned CPU tensor.", + "tensor"_a) + + .def("new_pinned_device_tensor", + &deepspeed_gds_handle_t::new_pinned_device_tensor, + "Allocate pinned device tensor.", + "num_elem"_a, + "example_tenosr"_a) + + .def("free_pinned_device_tensor", + &deepspeed_gds_handle_t::free_pinned_device_tensor, + "Free pinned device tensor.", + "tensor"_a) + + .def("pin_device_tensor", + &deepspeed_gds_handle_t::pin_device_tensor, + "Pin device tensor.", + "tensor"_a) + + .def("unpin_device_tensor", + &deepspeed_gds_handle_t::unpin_device_tensor, + "Unpin device tensor.", + "tensor"_a) + + .def("wait", + &deepspeed_gds_handle_t::wait, + "Wait for (ongoing) asynchronous operations to complete"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c629b93517d278f586613dd1d32cd4ac4dc2867a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda.cpp @@ -0,0 +1,113 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +// CUDA forward declaration +void fused_lamb_cuda(at::Tensor& p, + at::Tensor& p_copy, + at::Tensor& m, + at::Tensor& v, + at::Tensor& g, + float lr, + float beta1, + float beta2, + float max_coeff, + float min_coeff, + float eps, + float grad_scale, + int step, + int mode, + int bias_correction, + float decay, + at::Tensor& w_l2_i, + at::Tensor& u_l2_i, + at::Tensor& lamb_coeff_val); + +#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x) + +// C++ interface +at::Tensor lamb(at::Tensor& p, + at::Tensor& p_copy, + at::Tensor& m, + at::Tensor& v, + at::Tensor& g, + float lr, + float beta1, + float beta2, + float max_coeff, + float min_coeff, + float eps, + float grad_scale, + int step, + int mode, + int bias_correction, + float decay) +{ + CHECK_INPUT(p); + if (p_copy.numel() > 0) CHECK_INPUT(p_copy); + CHECK_INPUT(m); + CHECK_INPUT(v); + CHECK_INPUT(g); + int64_t num_elem = p.numel(); + AT_ASSERTM(m.numel() == num_elem, "number of elements in m and p tensors should be equal"); + AT_ASSERTM(v.numel() == num_elem, "number of elements in v and p tensors should be equal"); + AT_ASSERTM(g.numel() == num_elem, "number of elements in g and p tensors should be equal"); + AT_ASSERTM( + p_copy.numel() == num_elem || p_copy.numel() == 0, + "number of elements in p_copy and p tensors should be equal, or p_copy should be empty"); + + // intermediate for weight L2 reduction + // make sure that the threads per block is at least 512 during the kernel launch otherwise the + // behaviour is unexpected + at::Tensor w_l2_i = at::empty( + {512}, + p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float + : p.type().scalarType())); + + // intermediate for update L2 reduction + // make sure that the threads per block is at least 512 during the kernel launch otherwise the + // behaviour is unexpected + at::Tensor u_l2_i = at::empty( + {512}, + p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float + : p.type().scalarType())); + + at::Tensor lamb_coeff_val = at::empty( + {1}, + p.options().dtype(p.type().scalarType() == at::ScalarType::Half ? at::ScalarType::Float + : p.type().scalarType())); + + fused_lamb_cuda(p, + p_copy, + m, + v, + g, + lr, + beta1, + beta2, + max_coeff, + min_coeff, + eps, + grad_scale, + step, + mode, + bias_correction, + decay, + w_l2_i, + u_l2_i, + lamb_coeff_val); + + return lamb_coeff_val; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("lamb", &lamb, "Adam optimized CUDA implementation with LAMB."); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..d9bacae73457fe1e757f1f6410d64d18ebcf31bb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/lamb/fused_lamb_cuda_kernel.cu @@ -0,0 +1,478 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include "ATen/ATen.h" +#include "ATen/TensorUtils.h" +#include "ATen/cuda/CUDAContext.h" +#include "ATen/cuda/detail/IndexUtils.cuh" +// #include "ATen/Type.h" +#include "ATen/AccumulateType.h" + +#include + +// #include +#if defined(__HIP_PLATFORM_AMD__) && HIP_VERSION > 305 +#include +#else +#include +#endif +#include +#include + +namespace cg = cooperative_groups; + +// Utility class used to avoid linker errors with extern +// unsized shared memory arrays with templated type +namespace { +// This is the un-specialized struct. Note that we prevent instantiation of this +// struct by putting an undefined symbol in the function body so it won't compile. +template +struct SharedMemory { + // Ensure that we won't compile any un-specialized types + __device__ inline operator T*() + { +#ifndef _WIN32 + extern __device__ void error(void); + error(); +#endif + return NULL; + } +}; + +template <> +struct SharedMemory { + __device__ inline operator float*() + { + extern __shared__ float s_float[]; + return s_float; + } +}; + +template <> +struct SharedMemory { + __device__ inline operator double*() + { + extern __shared__ double s_double[]; + return s_double; + } +}; +} // namespace + +#include "type_shim.h" + +typedef enum { + ADAM_MODE_0 = 0, // eps under square root + ADAM_MODE_1 = 1 // eps outside square root +} adamMode_t; + +// s_a and s_b are in shared memory +// g_a and g_b are in shared memory +template +__device__ void reduce_block_in_shared_memory(T* s_a, T* s_b, T* g_a, T* g_b) +{ + // Handle to thread block group + cg::thread_block cta = cg::this_thread_block(); + + // perform block reduction in shared memory, + unsigned int tid = cta.thread_rank(); + + T a_sum = s_a[tid]; + T b_sum = s_b[tid]; + + cg::sync(cta); + + // do reduction in shared mem + if ((blockSize >= 512) && (tid < 256)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 256]; + s_b[tid] = b_sum = b_sum + s_b[tid + 256]; + } + + cg::sync(cta); + + if ((blockSize >= 256) && (tid < 128)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 128]; + s_b[tid] = b_sum = b_sum + s_b[tid + 128]; + } + + cg::sync(cta); + + if ((blockSize >= 128) && (tid < 64)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 64]; + s_b[tid] = b_sum = b_sum + s_b[tid + 64]; + } + + cg::sync(cta); + +#if (__CUDA_ARCH__ >= 300) || (defined(__HIP_PLATFORM_AMD__) && HIP_VERSION >= 502) + if (tid < 32) { + cg::coalesced_group active = cg::coalesced_threads(); + + // Fetch final intermediate sum from 2nd warp + if (blockSize >= 64) { + a_sum = a_sum + s_a[tid + 32]; + b_sum = b_sum + s_b[tid + 32]; + } + + // Reduce final warp using shuffle + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + a_sum += active.shfl_down(a_sum, offset); + b_sum += active.shfl_down(b_sum, offset); + } + } +#else + if ((blockSize >= 64) && (tid < 32)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 32]; + s_b[tid] = b_sum = b_sum + s_b[tid + 32]; + } + + cg::sync(cta); + + if ((blockSize >= 32) && (tid < 16)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 16]; + s_b[tid] = b_sum = b_sum + s_b[tid + 16]; + } + + cg::sync(cta); + + if ((blockSize >= 16) && (tid < 8)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 8]; + s_b[tid] = b_sum = b_sum + s_b[tid + 8]; + } + + cg::sync(cta); + + if ((blockSize >= 8) && (tid < 4)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 4]; + s_b[tid] = b_sum = b_sum + s_b[tid + 4]; + } + + cg::sync(cta); + + if ((blockSize >= 4) && (tid < 2)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 2]; + s_b[tid] = b_sum = b_sum + s_b[tid + 2]; + } + + cg::sync(cta); + + if ((blockSize >= 2) && (tid < 1)) { + s_a[tid] = a_sum = a_sum + s_a[tid + 1]; + s_b[tid] = b_sum = b_sum + s_b[tid + 1]; + } + + cg::sync(cta); + +#endif + + // write result for this block to global mem + if (tid == 0) { + g_a[blockIdx.x] = (T)a_sum; + g_b[blockIdx.x] = (T)b_sum; + } +} + +template +__device__ void reduce_two_vectors_in_register(T a, T b, T* g_a, T* g_b) +{ + const int threadIdInBlock = cg::this_thread_block().thread_rank(); + + T* s_a = SharedMemory(); + T* s_b = SharedMemory() + cg::this_thread_block().size(); + + s_a[threadIdInBlock] = a; + s_b[threadIdInBlock] = b; + + reduce_block_in_shared_memory(s_a, s_b, g_a, g_b); +} + +template +__global__ void lamb_cuda_kernel_part1( + T* __restrict__ p, + GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed + T* __restrict__ m, + T* __restrict__ v, + const GRAD_T* __restrict__ g, + const float b1, + const float b2, + const float eps, + const float grad_scale, + const float step_size, + const size_t tsize, + adamMode_t mode, + const float decay, + T* __restrict__ w_l2_i, + T* __restrict__ u_l2_i) +{ + // Assuming 2D grids and 2D blocks + const int blockId = gridDim.x * blockIdx.y + blockIdx.x; + const int threadsPerBlock = blockDim.x * blockDim.y; + const int threadIdInBlock = cg::this_thread_block().thread_rank(); + const int i = (blockId * threadsPerBlock + threadIdInBlock); + const int totThreads = gridDim.x * gridDim.y * threadsPerBlock; + + T reg_w = 0; + T reg_u = 0; + + for (int j = i; j < tsize; j += totThreads) { + T scaled_grad = g[j] / grad_scale; + T pj = p[j]; + m[j] = b1 * m[j] + (1 - b1) * scaled_grad; + v[j] = b2 * v[j] + (1 - b2) * scaled_grad * scaled_grad; + float denom; + if (mode == ADAM_MODE_0) + denom = sqrtf(v[j] + eps); + else // Mode 1 + denom = sqrtf(v[j]) + eps; + T update = (m[j] / denom) + (decay * p[j]); + + reg_u += update * update; + reg_w += pj * pj; + } + + reduce_two_vectors_in_register(reg_w, reg_u, w_l2_i, u_l2_i); +} + +template +__global__ void lamb_cuda_kernel_part2(const size_t tsize, T* __restrict__ g_a, T* __restrict__ g_b) +{ + T* s_a = SharedMemory(); + T* s_b = SharedMemory() + cg::this_thread_block().size(); + + const int threadIdInBlock = cg::this_thread_block().thread_rank(); + + s_a[threadIdInBlock] = g_a[threadIdInBlock]; + s_b[threadIdInBlock] = g_b[threadIdInBlock]; + + if (threadIdInBlock >= tsize) { + s_a[threadIdInBlock] = 0.0; + s_b[threadIdInBlock] = 0.0; + } + + reduce_block_in_shared_memory(s_a, s_b, g_a, g_b); +} + +template +__global__ void lamb_cuda_kernel_part3( + T* __restrict__ p, + GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed + T* __restrict__ m, + T* __restrict__ v, + const GRAD_T* __restrict__ g, + const float b1, + const float b2, + const float max_coeff, + const float min_coeff, + const float eps, + const float grad_scale, + const float step_size, + const size_t tsize, + adamMode_t mode, + const float decay, + T* __restrict__ w_l2_i, + T* __restrict__ u_l2_i, + T* __restrict__ lamb_coeff_val) +{ + // Assuming 2D grids and 2D blocks + const int blockId = gridDim.x * blockIdx.y + blockIdx.x; + const int threadsPerBlock = blockDim.x * blockDim.y; + const int threadIdInBlock = cg::this_thread_block().thread_rank(); + const int i = (blockId * threadsPerBlock + threadIdInBlock); + const int totThreads = gridDim.x * gridDim.y * threadsPerBlock; + + T reg_w = sqrtf(w_l2_i[0]); + T reg_u = sqrtf(u_l2_i[0]); + + float lamb_coeff = 1.0; + + if (reg_w != 0 && reg_u != 0) { + lamb_coeff = reg_w / reg_u; + if (lamb_coeff > max_coeff) { lamb_coeff = max_coeff; } + if (lamb_coeff < min_coeff) { lamb_coeff = min_coeff; } + } + + if (blockId == 0 && threadIdInBlock == 0) { + lamb_coeff_val[0] = lamb_coeff; + // printf("Cuda Lamb Coeff is %.6f \n",lamb_coeff); + } + + for (int j = i; j < tsize; j += totThreads) { + T pj = (float)p[j]; + T mj = m[j]; + T vj = v[j]; + float denom; + if (mode == ADAM_MODE_0) + denom = sqrtf(vj + eps); + else // Mode 1 + denom = sqrtf(vj) + eps; + T update = (mj / denom) + (decay * pj); + + pj = pj - (step_size * lamb_coeff * update); + p[j] = pj; + if (p_copy != NULL) p_copy[j] = (GRAD_T)pj; + } +} + +void fused_lamb_cuda(at::Tensor& p, + at::Tensor& p_copy, + at::Tensor& m, + at::Tensor& v, + at::Tensor& g, + float lr, + float beta1, + float beta2, + float max_coeff, + float min_coeff, + float eps, + float grad_scale, + int step, + int mode, + int bias_correction, + float decay, + at::Tensor& w_l2_i, + at::Tensor& u_l2_i, + at::Tensor& lamb_coeff) +{ + // using namespace at; + + // Get tensor size + int tsize = p.numel(); + // Determine #threads and #blocks + const int threadsPerBlock = 512; + int num_blocks = (tsize + threadsPerBlock - 1) / threadsPerBlock; + if (num_blocks > 512) num_blocks = 512; + + int smemsize = 0; + + if (p.type().scalarType() == at::ScalarType::Double) + smemsize = 2 * threadsPerBlock * sizeof(double); + else + smemsize = 2 * threadsPerBlock * sizeof(float); + + const dim3 blocks(num_blocks); + const dim3 threads(threadsPerBlock); + + AT_ASSERTM(at::cuda::detail::canUse32BitIndexMath(p), + "parameter tensor is too large to be indexed with int32"); + // Constants + float step_size = 0; + if (bias_correction == 1) { + const float bias_correction1 = 1 - std::pow(beta1, step); + const float bias_correction2 = 1 - std::pow(beta2, step); + step_size = lr * std::sqrt(bias_correction2) / bias_correction1; + } else { + step_size = lr; + } + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + if (g.type().scalarType() == at::ScalarType::Half) { + // all other values should be fp32 for half gradients + AT_ASSERTM(p.type().scalarType() == at::ScalarType::Float, + "expected parameter to be of float type"); + // dispatch is done on the gradient type + using namespace at; // prevents "toString is undefined" errors + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + g.scalar_type(), "lamb_cuda_kernel", ([&] { + using accscalar_t = at::acc_type; + + lamb_cuda_kernel_part1 + <<>>( + p.data(), + p_copy.numel() ? p_copy.data() : NULL, + m.data(), + v.data(), + g.data(), + beta1, + beta2, + eps, + grad_scale, + step_size, + tsize, + (adamMode_t)mode, + decay, + w_l2_i.data(), + u_l2_i.data()); + + lamb_cuda_kernel_part2 + <<<1, threadsPerBlock, smemsize, stream>>>( + num_blocks, w_l2_i.data(), u_l2_i.data()); + + lamb_cuda_kernel_part3 + <<>>( + p.data(), + p_copy.numel() ? p_copy.data() : NULL, + m.data(), + v.data(), + g.data(), + beta1, + beta2, + max_coeff, + min_coeff, + eps, + grad_scale, + step_size, + tsize, + (adamMode_t)mode, + decay, + w_l2_i.data(), + u_l2_i.data(), + lamb_coeff.data()); + })); + } else { + using namespace at; + AT_DISPATCH_FLOATING_TYPES( + g.scalar_type(), "lamb_cuda_kernel", ([&] { + lamb_cuda_kernel_part1 + <<>>( + p.data(), + NULL, // don't output p_copy for fp32, it's wasted write + m.data(), + v.data(), + g.data(), + beta1, + beta2, + eps, + grad_scale, + step_size, + tsize, + (adamMode_t)mode, + decay, + w_l2_i.data(), + u_l2_i.data()); + + lamb_cuda_kernel_part2 + <<<1, threadsPerBlock, smemsize, stream>>>( + num_blocks, w_l2_i.data(), u_l2_i.data()); + + lamb_cuda_kernel_part3 + <<>>( + p.data(), + NULL, // don't output p_copy for fp32, it's wasted write + m.data(), + v.data(), + g.data(), + beta1, + beta2, + max_coeff, + min_coeff, + eps, + grad_scale, + step_size, + tsize, + (adamMode_t)mode, + decay, + w_l2_i.data(), + u_l2_i.data(), + lamb_coeff.data()); + })); + } + C10_CUDA_CHECK(cudaGetLastError()); +} + +// template __device__ void reduce_two_vectors_in_register(float a, float b, float* g_a, +// float* g_b, cg::grid_group &cgg); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..783ffb7a41e396d73e9c0254826012a1ddc6655c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/dequantize.cu @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "dequantization_utils.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; + +template +__global__ void dequantize_kernel(T* __restrict__ dequant_data, + const int8_t* __restrict__ q_data, + const float* __restrict__ q_params, + int elems_per_group, + int total_elems) +{ + dequantize::to_global( + dequant_data, q_data, q_params, elems_per_group, total_elems); +} + +#define LAUNCH_DEQUANT_KERNEL(num_bits, q_type) \ + dequantize_kernel<<>>( \ + dequant_data, q_data, q_params, elems_per_group, total_elems); + +template +void launch_dequantize_kernel(T* dequant_data, + const int8_t* q_data, + const float* q_params, + quantize::Type q_type, + int num_bits, + int elems_per_group, + int total_elems, + cudaStream_t stream) +{ + constexpr int unroll = 8; + constexpr int threads = 512; + constexpr int elems_per_block = unroll * threads * dequantize::granularity / (sizeof(T)); + + const dim3 block(threads); + const dim3 grid((total_elems + elems_per_block - 1) / elems_per_block); + + // TODO(cmikeh2): It may make sense to tune unroll, there is perf benefit for large + // problem sizes with this large unroll value. + if (num_bits == 8 && q_type == quantize::Type::Symmetric) { + LAUNCH_DEQUANT_KERNEL(8, quantize::Type::Symmetric); + } else if (num_bits == 8 && q_type == quantize::Type::Asymmetric) { + LAUNCH_DEQUANT_KERNEL(8, quantize::Type::Asymmetric); + } else if (num_bits == 4 && q_type == quantize::Type::Symmetric) { + LAUNCH_DEQUANT_KERNEL(4, quantize::Type::Symmetric); + } else if (num_bits == 4 && q_type == quantize::Type::Asymmetric) { + LAUNCH_DEQUANT_KERNEL(4, quantize::Type::Asymmetric); + } +} + +template void launch_dequantize_kernel(__half* dequant_data, + const int8_t* q_data, + const float* q_params, + quantize::Type q_type, + int num_bits, + int elems_per_group, + int total_elems, + cudaStream_t stream); + +template void launch_dequantize_kernel(float* dequant_data, + const int8_t* q_data, + const float* q_params, + quantize::Type q_type, + int num_bits, + int elems_per_group, + int total_elems, + cudaStream_t stream); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu new file mode 100644 index 0000000000000000000000000000000000000000..4c08cd4cc3d28d9095813e9190a5034a59c5660e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/fake_quantizer.cu @@ -0,0 +1,1028 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "custom_cuda_layers.h" +#include "memory_access_utils.h" + +namespace cg = cooperative_groups; + +__global__ void fake_quantize_kernel(__half* vals, int group_size, int num_bits) +{ +#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) + + cg::thread_block b = cg::this_thread_block(); // tb + cg::thread_block_tile<32> g = + cg::tiled_partition<32>(b); // warp, 32 not optimal for AMD which should be 64. + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + constexpr int granularity = 16; + constexpr int vals_per_access = granularity / sizeof(__half); + + __half data[vals_per_access]; + + int group_id = blockIdx.x; + + int thread_index = id * vals_per_access; + int reg_count = 0; + int offset = group_id * group_size; + float max = -10000.0; + for (int thread_index = id * vals_per_access; thread_index < group_size; + thread_index += blockDim.x * vals_per_access) { + mem_access::load_global(data, vals + offset + thread_index); + +#pragma unroll + for (int i = 0; i < vals_per_access; i++) { + if (abs((float)data[i]) > max) max = abs((float)data[i]); + } + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + __shared__ float partialMax[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } + + max = g.shfl(max, 0); + + float q_scale = (float)(1 << num_bits) / (2 * max + 1e-5); + float q_scale_inv = 1 / q_scale; + int q_range_max = (1 << (num_bits - 1)) - 1; + int q_range_min = -(1 << (num_bits - 1)); + + for (int thread_index = id * vals_per_access; thread_index < group_size; + thread_index += blockDim.x * vals_per_access) { + mem_access::load_global(data, vals + offset + thread_index); +#pragma unroll + for (int j = 0; j < vals_per_access; j++) { + float q_data; + q_data = __half2float(data[j]); + q_data = __float2int_rn(q_data * q_scale); + q_data = q_data > (q_range_max) ? (q_range_max) + : (q_data < (q_range_min) ? (q_range_min) : q_data); + data[j] = __float2half_rn(q_data * q_scale_inv); + } + mem_access::store_global(vals + offset + thread_index, data); + } + +#endif +} + +__global__ void fake_quantize_kernel(float* vals, int group_size, int num_bits) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + constexpr int granularity = 16; + constexpr int vals_per_access = granularity / sizeof(float); + + float data[vals_per_access]; + + int bid = blockIdx.x; + + int thread_index = id * vals_per_access; + + int reg_count = 0; + + int offset = bid * group_size; + + float max = -10000.0; + + for (int thread_index = id * vals_per_access; thread_index < group_size; + thread_index += blockDim.x * vals_per_access) { + mem_access::load_global(data, vals + offset + thread_index); + +#pragma unroll + for (int i = 0; i < vals_per_access; i++) { + if (abs(data[i]) > max) max = abs(data[i]); + } + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + __shared__ float partialMax[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + + b.sync(); + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } + + max = g.shfl(max, 0); + + float q_scale = (1 << num_bits) / (2 * max + 1e-5); + float q_scale_inv = 1 / q_scale; + + int q_range_max = (1 << (num_bits - 1)) - 1; + int q_range_min = -(1 << (num_bits - 1)); + + for (int thread_index = id * vals_per_access; thread_index < group_size; + thread_index += blockDim.x * vals_per_access) { + mem_access::load_global(data, vals + offset + thread_index); +#pragma unroll + for (int j = 0; j < vals_per_access; j++) { + float q_data; + q_data = __float2int_rn(data[j] * q_scale); + q_data = q_data > (q_range_max) ? (q_range_max) + : (q_data < (q_range_min) ? (q_range_min) : q_data); + data[j] = roundf(q_data * q_scale_inv); + } + mem_access::store_global(vals + offset + thread_index, data); + } +} + +template +void launch_fake_quantize_kernel(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream) +{ + dim3 grid_dim(group_num); + dim3 block_dim(1024); + + fake_quantize_kernel<<>>( + vals, total_count / group_num, num_bits); +} + +template void launch_fake_quantize_kernel(float* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template void launch_fake_quantize_kernel(__half* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); + +__global__ void sr_fake_quantize_kernel(__half* vals, + int token_size, + int token_num, + int num_bits, + std::pair seed) +{ +#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + + float2* vals_cast = reinterpret_cast(vals); + + __half2 data_low[128]; + __half2 data_high[128]; + + int bid = blockIdx.x; + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + unsigned int tid = threadIdx.x; + int reg_count = 0; + int offset = bid * token_size; + int group_index = bid * token_size + tid; + + int total_count = token_size * token_num; + if (group_index < total_count) { + // float min = 10000.0; + float max = -10000.0; + while (tid < token_size) { + float2 data = vals_cast[offset + tid]; + __half2* data_h = reinterpret_cast<__half2*>(&data); + data_low[reg_count] = data_h[0]; + data_high[reg_count] = data_h[1]; + + float2 data_f[2]; + data_f[0] = __half22float2(data_h[0]); + data_f[1] = __half22float2(data_h[1]); + + if (abs((float)data_f[0].x) > max) max = abs((float)data_f[0].x); + if (abs((float)data_f[0].y) > max) max = abs((float)data_f[0].y); + if (abs((float)data_f[1].x) > max) max = abs((float)data_f[1].x); + if (abs((float)data_f[1].y) > max) max = abs((float)data_f[1].y); + + tid += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } + + max = g.shfl(max, 0); + + float q_scale_val = (float)(1 << num_bits) / (max * 2 + 1e-5); + float high_q = (float)((1 << (num_bits - 1)) - 1); + float low_q = (float)(-((1 << (num_bits - 1)))); + + for (int i = 0; i < reg_count; i++) { + int token_index = i * blockDim.x + threadIdx.x; + if (token_index < token_size) { + float2 data_f[2]; + data_f[0] = __half22float2(data_low[i]); + data_f[1] = __half22float2(data_high[i]); + + float2 q_data_int[2]; + q_data_int[0].x = (float)((int)(data_f[0].x * q_scale_val)); + q_data_int[0].y = (float)((int)(data_f[0].y * q_scale_val)); + q_data_int[1].x = (float)((int)(data_f[1].x * q_scale_val)); + q_data_int[1].y = (float)((int)(data_f[1].y * q_scale_val)); + + // Stochastic rounding + float4 rand = curand_uniform4(&state); + + float q_error[4]; + q_error[0] = abs(data_f[0].x - (q_data_int[0].x / q_scale_val)) * q_scale_val; + q_error[1] = abs(data_f[0].y - (q_data_int[0].y / q_scale_val)) * q_scale_val; + q_error[2] = abs(data_f[1].x - (q_data_int[1].x / q_scale_val)) * q_scale_val; + q_error[3] = abs(data_f[1].y - (q_data_int[1].y / q_scale_val)) * q_scale_val; + + q_data_int[0].x = + (rand.x < q_error[0] && q_data_int[0].x > low_q && q_data_int[0].x < high_q) + ? (q_data_int[0].x + (data_f[0].x > 0 ? 1 : -1)) + : q_data_int[0].x; + q_data_int[0].y = + (rand.y < q_error[1] && q_data_int[0].y > low_q && q_data_int[0].y < high_q) + ? (q_data_int[0].y + (data_f[0].y > 0 ? 1 : -1)) + : q_data_int[0].y; + q_data_int[1].x = + (rand.w < q_error[2] && q_data_int[1].x > low_q && q_data_int[1].x < high_q) + ? (q_data_int[1].x + (data_f[1].x > 0 ? 1 : -1)) + : q_data_int[1].x; + q_data_int[1].y = + (rand.z < q_error[3] && q_data_int[1].y > low_q && q_data_int[1].y < high_q) + ? (q_data_int[1].y + (data_f[1].y > 0 ? 1 : -1)) + : q_data_int[1].y; + + data_f[0].x = q_data_int[0].x / q_scale_val; + data_f[0].y = q_data_int[0].y / q_scale_val; + data_f[1].x = q_data_int[1].x / q_scale_val; + data_f[1].y = q_data_int[1].y / q_scale_val; + + float2 result; + __half2* result_h = reinterpret_cast<__half2*>(&result); + result_h[0] = __float22half2_rn(data_f[0]); + result_h[1] = __float22half2_rn(data_f[1]); + + vals_cast[offset + token_index] = result; + } + } + } +#endif +} + +__global__ void sr_fake_quantize_kernel(float* vals, + int token_size, + int token_num, + int num_bits, + std::pair seed) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + int idx = blockIdx.x * blockDim.x + id; + + float4* vals_cast = reinterpret_cast(vals); + + float4 data[128]; + + int bid = blockIdx.x; + int tid = threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + int group_index = bid * token_size + threadIdx.x; + int reg_count = 0; + int total_count = token_size * token_num; + if (group_index < total_count) { + // float min = 10000.0; + float max = -10000.0; + + while (tid < token_size) { + data[reg_count] = vals_cast[group_index]; + + if (abs(data[reg_count].x) > max) max = abs(data[reg_count].x); + if (abs(data[reg_count].y) > max) max = abs(data[reg_count].y); + if (abs(data[reg_count].z) > max) max = abs(data[reg_count].z); + if (abs(data[reg_count].w) > max) max = abs(data[reg_count].w); + + group_index += blockDim.x; + tid += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + __shared__ float partialMax[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } + + max = g.shfl(max, 0); + + float q_scale_val = (float)(1 << num_bits) / (max * 2 + 1e-5); + float high_q = (float)((1 << (num_bits - 1)) - 1); + float low_q = (float)(-((1 << (num_bits - 1)))); + + int offset = (bid)*token_size; + for (int i = 0; i < reg_count; i++) { + group_index = i * blockDim.x + threadIdx.x; + if (group_index < token_size) { + float4 q_data = data[i]; + + float4 q_data_int; + q_data_int.x = (float)((int)(q_data.x * q_scale_val)); + q_data_int.y = (float)((int)(q_data.y * q_scale_val)); + q_data_int.w = (float)((int)(q_data.w * q_scale_val)); + q_data_int.z = (float)((int)(q_data.z * q_scale_val)); + + // Stochastic rounding + float4 rand = curand_uniform4(&state); + + float q_error[4]; + q_error[0] = abs(q_data.x - (q_data_int.x / q_scale_val)) * q_scale_val; + q_error[1] = abs(q_data.y - (q_data_int.y / q_scale_val)) * q_scale_val; + q_error[2] = abs(q_data.w - (q_data_int.w / q_scale_val)) * q_scale_val; + q_error[3] = abs(q_data.z - (q_data_int.z / q_scale_val)) * q_scale_val; + + q_data_int.x = + (rand.x < q_error[0] && q_data_int.x > low_q && q_data_int.x < high_q) + ? (q_data_int.x + (q_data.x > 0 ? 1 : -1)) + : q_data_int.x; + q_data_int.y = + (rand.y < q_error[1] && q_data_int.y > low_q && q_data_int.y < high_q) + ? (q_data_int.y + (q_data.y > 0 ? 1 : -1)) + : q_data_int.y; + q_data_int.w = + (rand.w < q_error[2] && q_data_int.w > low_q && q_data_int.w < high_q) + ? (q_data_int.w + (q_data.w > 0 ? 1 : -1)) + : q_data_int.w; + q_data_int.z = + (rand.z < q_error[3] && q_data_int.z > low_q && q_data_int.z < high_q) + ? (q_data_int.z + (q_data.z > 0 ? 1 : -1)) + : q_data_int.z; + + q_data_int.x /= q_scale_val; + q_data_int.y /= q_scale_val; + q_data_int.w /= q_scale_val; + q_data_int.z /= q_scale_val; + + vals_cast[group_index + offset] = q_data_int; + } + } + } +} + +template +void launch_sr_fake_quantize_kernel(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream) +{ + dim3 block_dim(1024); + dim3 grid_dim(group_num); + + uint64_t inc = total_count / grid_dim.x / block_dim.x; + std::pair seed = TrainingContext::Instance().IncrementOffset(inc); + + sr_fake_quantize_kernel<<>>( + vals, (total_count / group_num) / 4, group_num, num_bits, seed); +} +template void launch_sr_fake_quantize_kernel(float* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template void launch_sr_fake_quantize_kernel(__half* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); + +__global__ void fake_quantize_kernel_asym(__half* vals, int group_size, int num_bits) +{ +#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + float2* vals_cast = reinterpret_cast(vals); + + float2 data[MAX_REG]; + + int group_id = blockIdx.x; + + { + int group_index = id; + int reg_count = 0; + int offset = group_id * group_size; + float max = -10000.0; + float min = 10000.0; + + while (group_index < group_size && reg_count < MAX_REG) { + data[reg_count] = vals_cast[offset + group_index]; + __half* data_h = reinterpret_cast<__half*>(&data[reg_count]); + + if (((float)data_h[0]) > max) max = (float)data_h[0]; + if (((float)data_h[1]) > max) max = (float)data_h[1]; + if (((float)data_h[2]) > max) max = (float)data_h[2]; + if (((float)data_h[3]) > max) max = (float)data_h[3]; + + if (((float)data_h[0]) < min) min = (float)data_h[0]; + if (((float)data_h[1]) < min) min = (float)data_h[1]; + if (((float)data_h[2]) < min) min = (float)data_h[2]; + if (((float)data_h[3]) < min) min = (float)data_h[3]; + + group_index += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(min, i); + if (min > temp) min = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + __shared__ float partialMin[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + if (lane == 0) partialMin[gid] = min; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + if (lane < warp_num) min = partialMin[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(min, i); + if (min > temp) min = temp; + } + + max = g.shfl(max, 0); + min = g.shfl(min, 0); + + float q_scale = ((max - min) + 1e-5) / (float)(1 << num_bits); + float q_scale_inv = 1 / q_scale; + + for (int i = 0; i < reg_count; i++) { + group_index = i * blockDim.x + id; + if (group_index < group_size) { + __half2* data_h = reinterpret_cast<__half2*>(&data[i]); + float2 q_data[2]; + q_data[0] = __half22float2(data_h[0]); + q_data[1] = __half22float2(data_h[1]); + + float2 q_data_int[2]; + + q_data_int[0].x = roundf((q_data[0].x - min) * q_scale_inv); + q_data_int[0].y = roundf((q_data[0].y - min) * q_scale_inv); + q_data_int[1].x = roundf((q_data[1].x - min) * q_scale_inv); + q_data_int[1].y = roundf((q_data[1].y - min) * q_scale_inv); + + q_data_int[0].x = q_data_int[0].x * q_scale + min; + q_data_int[0].y = q_data_int[0].y * q_scale + min; + q_data_int[1].x = q_data_int[1].x * q_scale + min; + q_data_int[1].y = q_data_int[1].y * q_scale + min; + + data_h[0] = __float22half2_rn(q_data_int[0]); + data_h[1] = __float22half2_rn(q_data_int[1]); + + vals_cast[offset + group_index] = data[i]; + } + } + } +#endif +} + +__global__ void fake_quantize_kernel_asym(float* vals, int group_size, int num_bits) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + float4* vals_cast = reinterpret_cast(vals); + + float4 data[MAX_REG]; + + int bid = blockIdx.x; + + int group_index = bid * group_size + id; + int reg_count = 0; + + float max = -10000.0; + float min = 10000.0; + + while (id < group_size && reg_count < MAX_REG) { + float4 data_reg = vals_cast[group_index]; + data[reg_count] = data_reg; + + if (data_reg.x > max) max = data_reg.x; + if (data_reg.y > max) max = data_reg.y; + if (data_reg.w > max) max = data_reg.w; + if (data_reg.z > max) max = data_reg.z; + + if (data_reg.x < min) min = data_reg.x; + if (data_reg.y < min) min = data_reg.y; + if (data_reg.w < min) min = data_reg.w; + if (data_reg.z < min) min = data_reg.z; + + group_index += blockDim.x; + id += blockDim.x; + reg_count++; + } + id = threadIdx.x; + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(min, i); + if (min > temp) min = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + __shared__ float partialMin[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + if (lane == 0) partialMin[gid] = min; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + if (lane < warp_num) min = partialMin[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(min, i); + if (min > temp) min = temp; + } + + max = g.shfl(max, 0); + min = g.shfl(min, 0); + + float q_scale = ((max - min) + 1e-5) / (float)(1 << num_bits); + float q_scale_inv = 1 / q_scale; + for (int i = 0; i < reg_count; i++) { + group_index = i * blockDim.x + id; + if (group_index < group_size) { + float4 q_data; + q_data = data[i]; + + float4 q_data_int; + q_data_int.x = roundf((q_data.x - min) * q_scale_inv); + q_data_int.y = roundf((q_data.y - min) * q_scale_inv); + q_data_int.w = roundf((q_data.w - min) * q_scale_inv); + q_data_int.z = roundf((q_data.z - min) * q_scale_inv); + + q_data.x = q_data_int.x * q_scale + min; + q_data.y = q_data_int.y * q_scale + min; + q_data.w = q_data_int.w * q_scale + min; + q_data.z = q_data_int.z * q_scale + min; + + vals_cast[group_index + bid * group_size] = q_data; + } + } +} + +template +void launch_fake_quantize_kernel_asym(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream) +{ + dim3 grid_dim(group_num); + dim3 block_dim(1024); + + fake_quantize_kernel_asym<<>>( + vals, (total_count / group_num) / 4, num_bits); +} + +template void launch_fake_quantize_kernel_asym(float* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template void launch_fake_quantize_kernel_asym(__half* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); + +__global__ void sr_fake_quantize_kernel_asym(__half* vals, + int token_size, + int token_num, + int num_bits, + std::pair seed) +{ +#if __CUDA_ARCH__ >= 700 || defined(__HIP_PLATFORM_AMD__) + + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + + int idx = blockIdx.x * blockDim.x + threadIdx.x; + + float2* vals_cast = reinterpret_cast(vals); + + __half2 data_low[128]; + __half2 data_high[128]; + + int bid = blockIdx.x; + + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + unsigned int tid = threadIdx.x; + int reg_count = 0; + int offset = bid * token_size; + int group_index = bid * token_size + tid; + + int total_count = token_size * token_num; + if (group_index < total_count) { + float min = 10000.0; + float max = -10000.0; + while (tid < token_size) { + float2 data = vals_cast[offset + tid]; + __half2* data_h = reinterpret_cast<__half2*>(&data); + data_low[reg_count] = data_h[0]; + data_high[reg_count] = data_h[1]; + + float2 data_f[2]; + data_f[0] = __half22float2(data_h[0]); + data_f[1] = __half22float2(data_h[1]); + + if (((float)data_f[0].x) > max) max = (float)data_f[0].x; + if (((float)data_f[0].y) > max) max = (float)data_f[0].y; + if (((float)data_f[1].x) > max) max = (float)data_f[1].x; + if (((float)data_f[1].y) > max) max = (float)data_f[1].y; + + if (((float)data_f[0].x) < min) min = (float)data_f[0].x; + if (((float)data_f[0].y) < min) min = (float)data_f[0].y; + if (((float)data_f[1].x) < min) min = (float)data_f[1].x; + if (((float)data_f[1].y) < min) min = (float)data_f[1].y; + + tid += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(min, i); + if (min > temp) min = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + __shared__ float partialMin[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + if (lane == 0) partialMin[gid] = min; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + if (lane < warp_num) min = partialMin[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(min, i); + if (min > temp) min = temp; + } + + max = g.shfl(max, 0); + min = g.shfl(min, 0); + + float q_scale_val = ((max - min) + 1e-5) / (float)(1 << num_bits); + float q_scale_val_inv = 1 / q_scale_val; + float high_q = (float)((1 << num_bits) - 1); + + for (int i = 0; i < reg_count; i++) { + int token_index = i * blockDim.x + threadIdx.x; + if (token_index < token_size) { + float2 data_f[2]; + data_f[0] = __half22float2(data_low[i]); + data_f[1] = __half22float2(data_high[i]); + + float2 q_data_int[2]; + q_data_int[0].x = (float)((unsigned int)((data_f[0].x - min) * q_scale_val_inv)); + q_data_int[0].y = (float)((unsigned int)((data_f[0].y - min) * q_scale_val_inv)); + q_data_int[1].x = (float)((unsigned int)((data_f[1].x - min) * q_scale_val_inv)); + q_data_int[1].y = (float)((unsigned int)((data_f[1].y - min) * q_scale_val_inv)); + + // Stochastic rounding + float4 rand = curand_uniform4(&state); + + float q_error[4]; + q_error[0] = + abs(data_f[0].x - ((q_data_int[0].x * q_scale_val) + min)) * q_scale_val_inv; + q_error[1] = + abs(data_f[0].y - ((q_data_int[0].y * q_scale_val) + min)) * q_scale_val_inv; + q_error[2] = + abs(data_f[1].x - ((q_data_int[1].x * q_scale_val) + min)) * q_scale_val_inv; + q_error[3] = + abs(data_f[1].y - ((q_data_int[1].y * q_scale_val) + min)) * q_scale_val_inv; + + q_data_int[0].x = (rand.x < q_error[0] && q_data_int[0].x < high_q) + ? (q_data_int[0].x + 1) + : q_data_int[0].x; + q_data_int[0].y = (rand.y < q_error[1] && q_data_int[0].y < high_q) + ? (q_data_int[0].y + 1) + : q_data_int[0].y; + q_data_int[1].x = (rand.w < q_error[2] && q_data_int[1].x < high_q) + ? (q_data_int[1].x + 1) + : q_data_int[1].x; + q_data_int[1].y = (rand.z < q_error[3] && q_data_int[1].y < high_q) + ? (q_data_int[1].y + 1) + : q_data_int[1].y; + + data_f[0].x = q_data_int[0].x * q_scale_val + min; + data_f[0].y = q_data_int[0].y * q_scale_val + min; + data_f[1].x = q_data_int[1].x * q_scale_val + min; + data_f[1].y = q_data_int[1].y * q_scale_val + min; + + float2 result; + __half2* result_h = reinterpret_cast<__half2*>(&result); + result_h[0] = __float22half2_rn(data_f[0]); + result_h[1] = __float22half2_rn(data_f[1]); + + vals_cast[offset + token_index] = result; + } + } + } +#endif +} + +__global__ void sr_fake_quantize_kernel_asym(float* vals, + int token_size, + int token_num, + int num_bits, + std::pair seed) +{ + cg::thread_block b = cg::this_thread_block(); + cg::thread_block_tile<32> g = cg::tiled_partition<32>(b); + + int gid = threadIdx.x >> 5; + int lane = threadIdx.x & 0x1f; + int warp_num = blockDim.x >> 5; + int id = threadIdx.x; + + int idx = blockIdx.x * blockDim.x + id; + + float4* vals_cast = reinterpret_cast(vals); + + float4 data[128]; + + int bid = blockIdx.x; + int tid = threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(seed.first, idx, seed.second, &state); + + int group_index = bid * token_size + threadIdx.x; + int reg_count = 0; + int total_count = token_size * token_num; + if (group_index < total_count) { + float min = 10000.0; + float max = -10000.0; + + while (tid < token_size) { + float4 data_reg = vals_cast[group_index]; + data[reg_count] = data_reg; + if (data_reg.x > max) max = data_reg.x; + if (data_reg.y > max) max = data_reg.y; + if (data_reg.w > max) max = data_reg.w; + if (data_reg.z > max) max = data_reg.z; + + if (data_reg.x < min) min = data_reg.x; + if (data_reg.y < min) min = data_reg.y; + if (data_reg.w < min) min = data_reg.w; + if (data_reg.z < min) min = data_reg.z; + + group_index += blockDim.x; + tid += blockDim.x; + reg_count++; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(max, i); + if (max < temp) max = temp; + } + +#pragma unroll + for (int i = 1; i < WARP_SIZE; i <<= 1) { + auto temp = g.shfl_xor(min, i); + if (min > temp) min = temp; + } + + __shared__ float partialMax[WARP_SIZE]; + __shared__ float partialMin[WARP_SIZE]; + + if (lane == 0) partialMax[gid] = max; + if (lane == 0) partialMin[gid] = min; + + b.sync(); + + if (lane < warp_num) max = partialMax[lane]; + if (lane < warp_num) min = partialMin[lane]; + +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(max, i); + if (max < temp) max = temp; + } +#pragma unroll + for (int i = 1; i < warp_num; i <<= 1) { + auto temp = g.shfl_down(min, i); + if (min > temp) min = temp; + } + + max = g.shfl(max, 0); + min = g.shfl(min, 0); + + float q_scale_val = ((max - min) + 1e-5) / (float)(1 << num_bits); + float high_q = (float)((1 << num_bits) - 1); + + int offset = (bid)*token_size; + for (int i = 0; i < reg_count; i++) { + group_index = i * blockDim.x + threadIdx.x; + if (group_index < token_size) { + float4 q_data = data[i]; + + float4 q_data_int; + q_data_int.x = (float)((int)((q_data.x - min) / q_scale_val)); + q_data_int.y = (float)((int)((q_data.y - min) / q_scale_val)); + q_data_int.w = (float)((int)((q_data.w - min) / q_scale_val)); + q_data_int.z = (float)((int)((q_data.z - min) / q_scale_val)); + + // Stochastic rounding + float4 rand = curand_uniform4(&state); + + float q_error[4]; + q_error[0] = abs(q_data.x - ((q_data_int.x * q_scale_val) + min)) / q_scale_val; + q_error[1] = abs(q_data.y - ((q_data_int.y * q_scale_val) + min)) / q_scale_val; + q_error[2] = abs(q_data.w - ((q_data_int.w * q_scale_val) + min)) / q_scale_val; + q_error[3] = abs(q_data.z - ((q_data_int.z * q_scale_val) + min)) / q_scale_val; + + q_data_int.x = (rand.x < q_error[0] && q_data_int.x < high_q) ? (q_data_int.x + 1) + : q_data_int.x; + q_data_int.y = (rand.y < q_error[1] && q_data_int.y < high_q) ? (q_data_int.y + 1) + : q_data_int.y; + q_data_int.w = (rand.w < q_error[2] && q_data_int.w < high_q) ? (q_data_int.w + 1) + : q_data_int.w; + q_data_int.z = (rand.z < q_error[3] && q_data_int.z < high_q) ? (q_data_int.z + 1) + : q_data_int.z; + + q_data_int.x = q_data_int.x * q_scale_val + min; + q_data_int.y = q_data_int.y * q_scale_val + min; + q_data_int.w = q_data_int.w * q_scale_val + min; + q_data_int.z = q_data_int.z * q_scale_val + min; + + vals_cast[group_index + offset] = q_data_int; + } + } + } +} +template +void launch_sr_fake_quantize_kernel_asym(T* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream) +{ + dim3 block_dim(1024); + dim3 grid_dim(group_num); + + uint64_t inc = total_count / grid_dim.x / block_dim.x; + std::pair seed = TrainingContext::Instance().IncrementOffset(inc); + + sr_fake_quantize_kernel<<>>( + vals, (total_count / group_num) / 4, group_num, num_bits, seed); +} +template void launch_sr_fake_quantize_kernel_asym(float* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); +template void launch_sr_fake_quantize_kernel_asym(__half* vals, + int total_count, + int group_num, + int num_bits, + cudaStream_t stream); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp new file mode 100644 index 0000000000000000000000000000000000000000..b48eaacd0881c68ce06052e894f9f32797233cc8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/pt_binding.cpp @@ -0,0 +1,404 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include "quantization.h" + +template +at::Tensor ds_quantize(at::Tensor& vals, int groups, int bits) +{ + auto t_size = vals.sizes(); + int size = 1; + for (auto dim : t_size) size *= dim; + + if ((((size / groups) - 1) / 4096 + 1) <= 256) { + launch_fake_quantize_kernel( + (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); + } + return vals; +} + +template +at::Tensor ds_sr_quantize(at::Tensor& vals, int groups, int bits) +{ + auto t_size = vals.sizes(); + int size = 1; + for (auto dim : t_size) size *= dim; + + if (((size / groups) / 4 / 1024) <= 256) { + launch_sr_fake_quantize_kernel( + (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); + } + return vals; +} + +template +at::Tensor ds_quantize_asym(at::Tensor& vals, int groups, int bits) +{ + auto t_size = vals.sizes(); + int size = 1; + for (auto dim : t_size) size *= dim; + + if ((((size / groups) - 1) / 4096 + 1) <= 256) { + launch_fake_quantize_kernel_asym( + (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); + } + return vals; +} + +template +at::Tensor ds_sr_quantize_asym(at::Tensor& vals, int groups, int bits) +{ + auto t_size = vals.sizes(); + int size = 1; + for (auto dim : t_size) size *= dim; + + if (((size / groups) / 4 / 1024) <= 256) { + launch_sr_fake_quantize_kernel_asym( + (T*)vals.data_ptr(), size, groups, bits, at::cuda::getCurrentCUDAStream()); + } + return vals; +} + +std::vector quantize_kernel(at::Tensor& input_vals, + int groups, + int numBits, + quantize::Type quantType) +{ + auto dtype = at::kFloat; + auto params_options = at::TensorOptions() + .dtype(dtype) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + const int param_elems = (quantize::requires_offset(quantType)) ? 2 : 1; + auto params = torch::empty({groups, param_elems}, params_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + auto output_sizes = input_vals.sizes().vec(); + output_sizes[output_sizes.size() - 1] /= numBits == 8 ? 1 : 2; + auto output = torch::empty(output_sizes, output_options); + + const int elems_per_group = at::numel(input_vals) / groups; + + launch_quant((int8_t*)output.data_ptr(), + (float*)params.data_ptr(), + (__half*)input_vals.data_ptr(), + groups, + elems_per_group, + numBits, + quantType, + at::cuda::getCurrentCUDAStream()); + + return {output, params}; +} + +template +at::Tensor dequantize(at::Tensor& quantized_data, + at::Tensor& params, + int groups, + int num_bits, + quantize::Type quant_type) +{ + auto dtype = (std::is_same::value) ? torch::kFloat32 : torch::kFloat16; + auto output_options = at::TensorOptions() + .dtype(dtype) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + auto output_sizes = quantized_data.sizes().vec(); + output_sizes[output_sizes.size() - 1] *= num_bits == 8 ? 1 : 2; + auto output = torch::empty(output_sizes, output_options); + + const int total_elems = at::numel(output); + const int elems_per_group = total_elems / groups; + + launch_dequantize_kernel((T*)output.data_ptr(), + (const int8_t*)quantized_data.data_ptr(), + (const float*)params.data_ptr(), + quant_type, + num_bits, + elems_per_group, + total_elems, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +at::Tensor dequantize_int4_to_half_experimental(at::Tensor& data_in, + at::Tensor& scale_buffer, + at::Tensor& min_val_buffer, + int num_group, + int group_size) +{ + auto output_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA); + auto output = torch::empty({num_group, group_size}, output_options); + + launch_dequantize_int4_to_half_experimental((uint8_t*)data_in.data_ptr(), + (half*)output.data_ptr(), + (half*)scale_buffer.data_ptr(), + (half*)min_val_buffer.data_ptr(), + num_group, + group_size, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +at::Tensor dequantize_int8_to_half_experimental(at::Tensor& data_in, + at::Tensor& scale_buffer, + at::Tensor& min_val_buffer, + int num_group, + int group_size) +{ + auto output_options = at::TensorOptions().dtype(at::kHalf).device(at::kCUDA); + auto output = torch::empty({num_group, group_size}, output_options); + + launch_dequantize_int8_to_half_experimental((uint8_t*)data_in.data_ptr(), + (half*)output.data_ptr(), + (half*)scale_buffer.data_ptr(), + (half*)min_val_buffer.data_ptr(), + num_group, + group_size, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +std::vector ds_loco_swizzle_quant(at::Tensor& input_vals, + at::Tensor& error_feedback, + float err_beta, + int groups, + int num_bits, + quantize::Type quant_type, + int pipeline_size, + int nodes, + int devices_per_node) +{ + auto scales_options = at::TensorOptions() + .dtype(at::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; + auto scales = torch::empty({groups, scales_elems}, scales_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + const int quantization_scalar = 8 / num_bits; + const int compressed_vals = at::numel(input_vals) / quantization_scalar; + + auto output = torch::empty({compressed_vals}, output_options); + const int elems_per_group = at::numel(input_vals) / groups; + + launch_loco_swizzled_quant(reinterpret_cast(output.data_ptr()), + reinterpret_cast(scales.data_ptr()), + reinterpret_cast(input_vals.data_ptr()), + reinterpret_cast<__half*>(error_feedback.data_ptr()), + err_beta, + num_bits, + quant_type, + groups, + elems_per_group, + pipeline_size, + nodes, + devices_per_node, + at::cuda::getCurrentCUDAStream()); + + return {output, scales}; +} + +std::vector ds_swizzle_quant(at::Tensor& input_vals, + int groups, + int num_bits, + quantize::Type quant_type, + int pipeline_size, + int nodes, + int devices_per_node) +{ + auto scales_options = at::TensorOptions() + .dtype(at::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; + auto scales = torch::empty({groups, scales_elems}, scales_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + const int quantization_scalar = 8 / num_bits; + const int compressed_vals = at::numel(input_vals) / quantization_scalar; + + auto output = torch::empty({compressed_vals}, output_options); + const int elems_per_group = at::numel(input_vals) / groups; + + launch_swizzled_quant((int8_t*)output.data_ptr(), + (float*)scales.data_ptr(), + (__half*)input_vals.data_ptr(), + num_bits, + quant_type, + groups, + elems_per_group, + pipeline_size, + nodes, + devices_per_node, + at::cuda::getCurrentCUDAStream()); + + return {output, scales}; +} + +std::vector quantized_reduction(at::Tensor& input_vals, + at::Tensor& input_scales, + int in_groups, + int out_groups, + int num_bits, + quantize::Type quant_type, + int devices_per_node) +{ + auto scales_options = at::TensorOptions() + .dtype(at::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; + auto scales = torch::empty({out_groups, scales_elems}, scales_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + std::vector sz(input_vals.sizes().begin(), input_vals.sizes().end()); + sz[sz.size() - 1] = sz.back() / devices_per_node; // num of GPU per nodes + const int elems_per_in_tensor = at::numel(input_vals) / devices_per_node; + auto output = torch::empty(sz, output_options); + + const int elems_per_in_group = elems_per_in_tensor / (in_groups / devices_per_node); + const int elems_per_out_group = elems_per_in_tensor / out_groups; + + launch_dequant_reduce((int8_t*)output.data_ptr(), + (float*)scales.data_ptr(), + (const int8_t*)input_vals.data_ptr(), + (const float*)input_scales.data_ptr(), + devices_per_node, + num_bits, + quant_type, + out_groups, + elems_per_out_group, + elems_per_in_tensor, + in_groups / devices_per_node, + elems_per_in_group, + at::cuda::getCurrentCUDAStream()); + return {output, scales}; +} + +std::vector loco_quantized_reduction(at::Tensor& input_vals, + at::Tensor& input_scales, + at::Tensor& error_feedback, + float err_beta, + int in_groups, + int out_groups, + int num_bits, + quantize::Type quant_type, + int devices_per_node) +{ + auto scales_options = at::TensorOptions() + .dtype(at::kFloat) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + const int scales_elems = (quantize::requires_offset(quant_type)) ? 2 : 1; + + auto scales = torch::empty({out_groups, scales_elems}, scales_options); + + auto output_options = at::TensorOptions() + .dtype(at::kChar) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + + std::vector sz(input_vals.sizes().begin(), input_vals.sizes().end()); + sz[sz.size() - 1] = sz.back() / devices_per_node; + + const int elems_per_in_tensor = at::numel(input_vals) / devices_per_node; + + auto output = torch::empty(sz, output_options); + + const int elems_per_in_group = elems_per_in_tensor / (in_groups / devices_per_node); + const int elems_per_out_group = elems_per_in_tensor / out_groups; + + launch_loco_dequant_reduce((int8_t*)output.data_ptr(), + (float*)scales.data_ptr(), + (const int8_t*)input_vals.data_ptr(), + (const float*)input_scales.data_ptr(), + devices_per_node, + num_bits, + quant_type, + out_groups, + elems_per_out_group, + elems_per_in_tensor, + in_groups / devices_per_node, + elems_per_in_group, + (__half2*)error_feedback.data_ptr(), + err_beta, + at::cuda::getCurrentCUDAStream()); + + return {output, scales}; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("ds_quantize_fp32", &ds_quantize, "DeepSpeed Quantize with fp32 (CUDA)"); + m.def("ds_quantize_fp16", &ds_quantize<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); + m.def("ds_sr_quantize_fp32", &ds_sr_quantize, "DeepSpeed Quantize with fp32 (CUDA)"); + m.def("ds_sr_quantize_fp16", &ds_sr_quantize<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); + m.def("ds_quantize_asym_fp32", &ds_quantize_asym, "DeepSpeed Quantize with fp32 (CUDA)"); + m.def( + "ds_quantize_asym_fp16", &ds_quantize_asym<__half>, "DeepSpeed Quantize with fp16 (CUDA)"); + m.def("ds_sr_quantize_asym_fp32", + &ds_sr_quantize_asym, + "DeepSpeed Quantize with fp32 (CUDA)"); + m.def("ds_sr_quantize_asym_fp16", + &ds_sr_quantize_asym<__half>, + "DeepSpeed Quantize with fp16 (CUDA)"); + pybind11::enum_(m, "QuantizationType") + .value("Symmetric", quantize::Type::Symmetric) + .value("Asymmetric", quantize::Type::Asymmetric) + .export_values(); + m.def("quantize", &quantize_kernel); + m.def("dequantize", &dequantize<__half>); + m.def("dequantize_fp32", &dequantize); + m.def("dequantize_int4_to_half_experimental", + &dequantize_int4_to_half_experimental, + "Dequantize int4 to half (experimental)"); + m.def("dequantize_int8_to_half_experimental", + &dequantize_int8_to_half_experimental, + "Dequantize int8 to half (experimental)"); + m.def("swizzle_quant", &ds_swizzle_quant); + m.def("quantized_reduction", &quantized_reduction); + m.def("loco_swizzle_quant", &ds_loco_swizzle_quant, "LoCo Swizzled Quantization Kernel"); + m.def("loco_quantized_reduction", + &loco_quantized_reduction, + "LoCo Quantization and Reduction Kernel"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu new file mode 100644 index 0000000000000000000000000000000000000000..4100c5174b809e0db4dcf138dfcfde4f42e75843 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quant_reduce.cu @@ -0,0 +1,557 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "dequantization_utils.h" +#include "ds_kernel_utils.h" +#include "memory_access_utils.h" +#include "quantization_utils.h" +#include "reduction_utils.h" + +using rop = reduce::ROpType; + +/* +TODO(cmikeh2): Add implementation that better handles larger nodes. It would like make sense +to leverage some parallel reductions here to improve performance. +*/ + +template +__global__ void __launch_bounds__(1024) dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + int num_tensors) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // NOTE(cmikeh2): This probably could be hardcoded to a larger number, + // but that means even stronger restrictions on the number of elements per group + // A performance analysis here might be beneficial + constexpr int mem_granularity = (numBits == 8) ? 8 : 4; + constexpr int elems_per_load = mem_granularity / sizeof(int8_t); // div by 1 + constexpr int storage_values = 16 / sizeof(__half2); + + const int block_offset = tb.group_index().x * elems_per_out_group; + const int elem_offset = tb.thread_index().x * elems_per_load; + const int base_offset = block_offset + elem_offset; + const int stride = tb.group_dim().x * elems_per_load; + + __half2 local_buffer[totalChunks * storage_values]; + + quantize::GroupStats stats; + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + __half2* iteration_buffer = local_buffer + i * storage_values; + +#pragma unroll + for (int j = 0; j < storage_values; j++) { + iteration_buffer[j] = reduce::init(); + } + + const int iter_offset = i * stride + base_offset; + const int iter_scale_idx = iter_offset / elems_per_in_group; + bool do_loads = i * stride + elem_offset < elems_per_out_group; + + if (numTensors > 0) { +#pragma unroll + for (int j = 0; j < numTensors; j++) { + if (do_loads) { + int8_t load_buffer[elems_per_load]; + + mem_access::load_global( + load_buffer, input_data + j * elems_per_in_tensor + iter_offset); + + quantize::Params params( + input_scales + j * groups_per_in_tensor, iter_scale_idx); + + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, load_buffer, params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = + reduce::element(iteration_buffer[k], dequant_buffer[k]); + } + } + } + } else { +#pragma unroll 4 + for (int j = 0; j < num_tensors; j++) { + if (do_loads) { + int8_t load_buffer[elems_per_load]; + + mem_access::load_global( + load_buffer, input_data + j * elems_per_in_tensor + iter_offset); + + quantize::Params params( + input_scales + j * groups_per_in_tensor, iter_scale_idx); + + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, load_buffer, params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = + reduce::element(iteration_buffer[k], dequant_buffer[k]); + } + } + } + } + +#pragma unroll + for (int j = 0; j < storage_values; j++) { stats.update(iteration_buffer[j]); } + } + + auto params = stats.template get_params(tb, warp); + + if (tb.thread_index().x == 0) { params.store(reduced_scales, tb.group_index().x); } + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + const int iter_offset = i * stride + base_offset; + if (i * stride + elem_offset < elems_per_out_group) { + int8_t local_output[elems_per_load]; + quantize::_chunk( + local_output, local_buffer + i * storage_values, params); + mem_access::store_global(reduced_data + iter_offset, local_output); + } + } +} + +template +int32_t pow2_round(int32_t raw_value) +{ + return (((raw_value - 1) >> Power) + 1) << Power; +} + +#define LAUNCH_DEQUANT_REDUCE(num_chunks) \ + dequant_reduce \ + <<>>(reduced_data, \ + reduced_scales, \ + input_data, \ + input_scales, \ + elems_per_out_group, \ + elems_per_in_tensor, \ + groups_per_in_tensor, \ + elems_per_in_group, \ + num_tensors); + +template +void launch_dequant_reduce_impl(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + int num_tensors, + cudaStream_t stream) +{ + // This is a coincidence. This is derived by 8 halves per 16 bytes with 2-way packing for int4 + constexpr int elems_per_thread = numBits; + const int one_step_threads = + next_pow2((elems_per_out_group + elems_per_thread - 1) / (elems_per_thread)); + // TODO(cmikeh2): Tune this + const int threads = (one_step_threads < 1024) ? one_step_threads : 1024; + + dim3 block(threads); + dim3 grid(out_groups); + + const int elems_per_step = threads * elems_per_thread; + const int unroll_raw = (elems_per_out_group + elems_per_step - 1) / elems_per_step; + + const int unroll = (unroll_raw >= 4) ? pow2_round<1>(unroll_raw) : unroll_raw; + + if (unroll == 1) { + // 0-4096 elems + LAUNCH_DEQUANT_REDUCE(1); + } else if (unroll == 2) { + // 4097-8192 etc... + LAUNCH_DEQUANT_REDUCE(2); + } else if (unroll == 3) { + LAUNCH_DEQUANT_REDUCE(3); + } else if (unroll == 4) { + LAUNCH_DEQUANT_REDUCE(4); + } else if (unroll == 6) { + LAUNCH_DEQUANT_REDUCE(6); + } else if (unroll == 8) { + LAUNCH_DEQUANT_REDUCE(8); + } else if (unroll == 10) { + LAUNCH_DEQUANT_REDUCE(10); + } else if (unroll == 12) { + // 48k limit + LAUNCH_DEQUANT_REDUCE(12); + } else { + assert(false); + } +} + +#define LAUNCH_DEQUANT_REDUCE_IMPL(NUM_BITS, NUM_GPUS, QUANT_TYPE) \ + launch_dequant_reduce_impl(reduced_data, \ + reduced_scales, \ + input_data, \ + input_scales, \ + out_groups, \ + elems_per_out_group, \ + elems_per_in_tensor, \ + groups_per_in_tensor, \ + elems_per_in_group, \ + num_gpus, \ + stream); + +void launch_dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int num_gpus, + int num_bits, + quantize::Type quant_type, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + cudaStream_t stream) +{ + if (quant_type == quantize::Type::Symmetric) { + if (num_bits == 4) { + if (num_gpus == 8) { + LAUNCH_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Symmetric); + } else if (num_gpus == 16) { + LAUNCH_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Symmetric); + } else { + LAUNCH_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Symmetric); + } + } else if (num_bits == 8) { + if (num_gpus == 8) { + LAUNCH_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Symmetric); + } else if (num_gpus == 16) { + LAUNCH_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Symmetric); + } else { + LAUNCH_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Symmetric); + } + } + } else if (quant_type == quantize::Type::Asymmetric) { + if (num_bits == 4) { + if (num_gpus == 8) { + LAUNCH_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Asymmetric); + } else if (num_gpus == 16) { + LAUNCH_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Asymmetric); + } else { + LAUNCH_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Asymmetric); + } + } else if (num_bits == 8) { + if (num_gpus == 8) { + LAUNCH_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Asymmetric); + } else if (num_gpus == 16) { + LAUNCH_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Asymmetric); + } else { + LAUNCH_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Asymmetric); + } + } + } +} + +/* +Modified loco_dequant_reduce function that performs dequantization and reduction, +and incorporates error-feedback by updating the error_feedback tensor in-place. +*/ + +template +__global__ void __launch_bounds__(1024) loco_dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + int num_tensors, + __half2* error_feedback, + const float err_beta) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + constexpr int mem_granularity = (numBits == 8) ? 8 : 4; + constexpr int elems_per_load = mem_granularity / sizeof(int8_t); + constexpr int storage_values = 16 / sizeof(__half2); + + const int block_offset = tb.group_index().x * elems_per_out_group; + const int elem_offset = tb.thread_index().x * elems_per_load; + const int base_offset = block_offset + elem_offset; + const int stride = tb.group_dim().x * elems_per_load; + + constexpr int scaling_factor = elems_per_load / storage_values; + const int block_offset_err = block_offset / scaling_factor; + const int elem_offset_err = tb.thread_index().x * storage_values; + const int base_offset_err = block_offset_err + elem_offset_err; + const int stride_err = tb.group_dim().x * storage_values; + + __half2 local_buffer[totalChunks * storage_values]; + __half2 err_buffer[totalChunks * storage_values]; + + quantize::GroupStats stats; + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + __half2* iteration_buffer = local_buffer + i * storage_values; + __half2* iter_err_buffer = err_buffer + i * storage_values; + +#pragma unroll + for (int j = 0; j < storage_values; j++) { + iteration_buffer[j] = reduce::init(); + } + + const int iter_offset = i * stride + base_offset; + const int iter_offset_err = i * stride_err + base_offset_err; + const int iter_scale_idx = iter_offset / elems_per_in_group; + bool do_loads = i * stride + elem_offset < elems_per_out_group; + + if (numTensors > 0) { +#pragma unroll + for (int j = 0; j < numTensors; j++) { + if (do_loads) { + int8_t load_buffer[elems_per_load]; + + mem_access::load_global( + load_buffer, input_data + j * elems_per_in_tensor + iter_offset); + + quantize::Params params( + input_scales + j * groups_per_in_tensor, iter_scale_idx); + + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, load_buffer, params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = + reduce::element(iteration_buffer[k], dequant_buffer[k]); + } + } + } + } else { +#pragma unroll 4 + for (int j = 0; j < num_tensors; j++) { + if (do_loads) { + int8_t load_buffer[elems_per_load]; + + mem_access::load_global( + load_buffer, input_data + j * elems_per_in_tensor + iter_offset); + + quantize::Params params( + input_scales + j * groups_per_in_tensor, iter_scale_idx); + + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, load_buffer, params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = + reduce::element(iteration_buffer[k], dequant_buffer[k]); + } + } + } + } + mem_access::load_global( + iter_err_buffer, error_feedback + iter_offset_err, do_loads); +#pragma unroll + for (int k = 0; k < storage_values; k++) { + iteration_buffer[k] = __hadd2(iteration_buffer[k], iter_err_buffer[k]); + stats.update(iteration_buffer[k]); + } + } + + auto params = stats.template get_params(tb, warp); + + // Initialize dequantization parameters based on params + auto de_params = params; + de_params.scale = 1.0f / params.scale; + if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; } + + if (tb.thread_index().x == 0) { params.store(reduced_scales, tb.group_index().x); } + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + const int iter_offset = i * stride + base_offset; + const int iter_offset_err = i * stride_err + base_offset_err; + __half2* iteration_buffer = local_buffer + i * storage_values; + __half2* iter_err_buffer = err_buffer + i * storage_values; + + if (i * stride + elem_offset < elems_per_out_group) { + // ----------- Begin Error-Feedback Modification ----------- + int8_t local_output[elems_per_load]; + quantize::_chunk(local_output, iteration_buffer, params); + mem_access::store_global(reduced_data + iter_offset, local_output); + + // Dequantize the quantized output to compute the dequantized value + __half2 dequant_buffer[storage_values]; + dequantize::chunk(dequant_buffer, local_output, de_params); + +#pragma unroll + for (int k = 0; k < storage_values; k++) { + // __half2 to float2 + float2 iter_buf_f = __half22float2(iteration_buffer[k]); + float2 dequant_buf_f = __half22float2(dequant_buffer[k]); + + // Update within float precision + float2 new_error_f; + new_error_f.x = iter_buf_f.x - dequant_buf_f.x; + new_error_f.y = iter_buf_f.y - dequant_buf_f.y; + + float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]); + + iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x; + iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y; + + // float2 back to __half2 + iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f); + } + mem_access::store_global(error_feedback + iter_offset_err, + iter_err_buffer); + } + } +} + +#define LAUNCH_LOCO_DEQUANT_REDUCE(num_chunks) \ + loco_dequant_reduce \ + <<>>(reduced_data, \ + reduced_scales, \ + input_data, \ + input_scales, \ + elems_per_out_group, \ + elems_per_in_tensor, \ + groups_per_in_tensor, \ + elems_per_in_group, \ + num_tensors, \ + error_feedback, \ + err_beta); + +template +void launch_loco_dequant_reduce_impl(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + int num_tensors, + __half2* error_feedback, + const float err_beta, + cudaStream_t stream) +{ + constexpr int elems_per_thread = numBits; + const int one_step_threads = + next_pow2((elems_per_out_group + elems_per_thread - 1) / (elems_per_thread)); + const int threads = (one_step_threads < 1024) ? one_step_threads : 1024; + + dim3 block(threads); + dim3 grid(out_groups); + + const int elems_per_step = threads * elems_per_thread; + const int unroll_raw = (elems_per_out_group + elems_per_step - 1) / elems_per_step; + + const int unroll = (unroll_raw >= 4) ? pow2_round<1>(unroll_raw) : unroll_raw; + + if (unroll == 1) { + LAUNCH_LOCO_DEQUANT_REDUCE(1); + } else if (unroll == 2) { + LAUNCH_LOCO_DEQUANT_REDUCE(2); + } else if (unroll == 3) { + LAUNCH_LOCO_DEQUANT_REDUCE(3); + } else if (unroll == 4) { + LAUNCH_LOCO_DEQUANT_REDUCE(4); + } else if (unroll == 6) { + LAUNCH_LOCO_DEQUANT_REDUCE(6); + } else if (unroll == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE(8); + } else if (unroll == 10) { + LAUNCH_LOCO_DEQUANT_REDUCE(10); + } else if (unroll == 12) { + LAUNCH_LOCO_DEQUANT_REDUCE(12); + } else { + assert(false); + } +} + +#define LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(NUM_BITS, NUM_GPUS, QUANT_TYPE) \ + launch_loco_dequant_reduce_impl(reduced_data, \ + reduced_scales, \ + input_data, \ + input_scales, \ + out_groups, \ + elems_per_out_group, \ + elems_per_in_tensor, \ + groups_per_in_tensor, \ + elems_per_in_group, \ + num_gpus, \ + error_feedback, \ + err_beta, \ + stream); + +void launch_loco_dequant_reduce(int8_t* reduced_data, + float* reduced_scales, + const int8_t* input_data, + const float* input_scales, + int num_gpus, + int num_bits, + quantize::Type quant_type, + int out_groups, + int elems_per_out_group, + int elems_per_in_tensor, + int groups_per_in_tensor, + int elems_per_in_group, + __half2* error_feedback, + const float err_beta, + cudaStream_t stream) +{ + if (quant_type == quantize::Type::Symmetric) { + if (num_bits == 4) { + if (num_gpus == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Symmetric); + } else if (num_gpus == 16) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Symmetric); + } else { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Symmetric); + } + } else if (num_bits == 8) { + if (num_gpus == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Symmetric); + } else if (num_gpus == 16) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Symmetric); + } else { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Symmetric); + } + } + } else if (quant_type == quantize::Type::Asymmetric) { + if (num_bits == 4) { + if (num_gpus == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 8, quantize::Type::Asymmetric); + } else if (num_gpus == 16) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, 16, quantize::Type::Asymmetric); + } else { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(4, -1, quantize::Type::Asymmetric); + } + } else if (num_bits == 8) { + if (num_gpus == 8) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 8, quantize::Type::Asymmetric); + } else if (num_gpus == 16) { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, 16, quantize::Type::Asymmetric); + } else { + LAUNCH_LOCO_DEQUANT_REDUCE_IMPL(8, -1, quantize::Type::Asymmetric); + } + } + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quantize.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..b04d0e968ba58a203f2b5e24790053cbd43bc74a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quantize.cu @@ -0,0 +1,151 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "ds_kernel_utils.h" +#include "memory_access_utils.h" +#include "quantization.h" +#include "quantization_utils.h" +#include "reduction_utils.h" + +namespace cg = cooperative_groups; + +/* +Pure quantization kernel with no fusion. +*/ +template +__global__ void cached_quantization(int8_t* __restrict__ output_data, + float* __restrict__ params, + const __half* __restrict__ input_data, + int groups, + int elems_per_group) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // Indexing offsets + const int block_offset = + (tb.group_index().x * (max_threads / threads_per_group) * elems_per_group) + + (tb.thread_index().y * elems_per_group); + const int elem_offset = tb.thread_index().x * quantize::h_per_load; + const int base_offset = block_offset + elem_offset; + const int stride = tb.size() * quantize::h_per_load; + + const __half* input_base = input_data + base_offset; //.. + + __half2 local_buffer[UNROLL * internal_unroll * quantize::h2_per_load]; + +#pragma unroll + for (int i = 0; i < UNROLL; i++) { + // Convenience helper, should resolve to register indices and not realize. + __half2* iteration_buffer = local_buffer + i * internal_unroll * quantize::h2_per_load; +#pragma unroll + for (int j = 0; j < internal_unroll; j++) { + const int iteration = i * internal_unroll + j; + mem_access::load_global( + iteration_buffer + j * quantize::h2_per_load, + input_base + iteration * stride, + elem_offset + iteration * stride < elems_per_group); + } + } + + quantize:: + local_array( + local_buffer, params, output_data, elems_per_group, groups); +} + +/********* Launcher methods ***********/ +#define LAUNCH_CACHED_QUANT_CALL(q_bits, quant_type) \ + cached_quantization \ + <<>>(output_data, params, input_data, groups, elems_per_group); + +#define LAUNCH_CACHED_QUANT( \ + q_bits, quant_type, unroll_factor_in, internal_unroll_in, threads_per_group_in) \ + const int unroll_factor = unroll_factor_in; \ + const int internal_unroll_l = internal_unroll_in; \ + const int threads_per_group = threads_per_group_in; \ + if (q_bits == 4) { \ + if (quant_type == quantize::Type::Asymmetric) { \ + LAUNCH_CACHED_QUANT_CALL(4, quantize::Type::Asymmetric) \ + } else { \ + LAUNCH_CACHED_QUANT_CALL(4, quantize::Type::Symmetric) \ + } \ + } else { \ + if (quant_type == quantize::Type::Asymmetric) { \ + LAUNCH_CACHED_QUANT_CALL(8, quantize::Type::Asymmetric) \ + } else { \ + LAUNCH_CACHED_QUANT_CALL(8, quantize::Type::Symmetric) \ + } \ + } + +void launch_quant(int8_t* output_data, + float* params, + const __half* input_data, + const int groups, + const int elems_per_group, + const int num_bits, + const quantize::Type quant_type, + cudaStream_t stream) +{ + constexpr int max_threads = 256; + + constexpr int internal_unroll = 2; + + const bool is_subblock_schedule = (elems_per_group <= 128) ? true : false; + const int h_per_step = is_subblock_schedule ? quantize::h_per_load + : quantize::h_per_load * internal_unroll; + + // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of + // warp-sized blocks rather than stepping up to 64/96 threads + const int one_step_threads = next_pow2((elems_per_group + h_per_step - 1) / h_per_step); + const int threads_per_group = (one_step_threads < max_threads) ? one_step_threads : max_threads; + + const int groups_per_block = + is_subblock_schedule ? (max_threads + threads_per_group - 1) / threads_per_group : 1; + const int groups_launch = (groups_per_block + groups - 1) / groups_per_block; + + dim3 block(threads_per_group, groups_per_block); + dim3 grid(groups_launch); + + const int elems_per_step = threads_per_group * h_per_step; + const int external_unroll = (elems_per_group + elems_per_step - 1) / elems_per_step; + + if (is_subblock_schedule) { + // <=128 + if (threads_per_group == 1) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 1); + } else if (threads_per_group == 2) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 2); + } else if (threads_per_group == 4) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 4); + } else if (threads_per_group == 8) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 8); + } else if (threads_per_group == 16) { + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, 1, 16); + } + } else if (external_unroll == 1) { + // 129 - 4096 elems + // (this can launch with 1-7 warps as well) + LAUNCH_CACHED_QUANT(num_bits, quant_type, 1, internal_unroll, max_threads); + } else if (external_unroll == 2) { + // 4097 - 8192 elems + LAUNCH_CACHED_QUANT(num_bits, quant_type, 2, internal_unroll, max_threads); + } else if (external_unroll == 3) { + // 8193 - 12288 elems + LAUNCH_CACHED_QUANT(num_bits, quant_type, 3, internal_unroll, max_threads); + } else if (external_unroll == 4) { + // 12289 - 16384 elems + LAUNCH_CACHED_QUANT(num_bits, quant_type, 4, internal_unroll, max_threads); + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu new file mode 100644 index 0000000000000000000000000000000000000000..b26151ab5c8c327d7b5b54239b73ee3ebd6bea8e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/quantize_intX.cu @@ -0,0 +1,281 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include "memory_access_utils.h" + +template +struct alignas(sizeof(T) * N) AlignedArray { + using Element = T; + static const int kElements = N; + + __device__ __host__ AlignedArray() {} + + __device__ __host__ AlignedArray(const T& rhs) + { +#pragma unroll + for (int idx = 0; idx < kElements; ++idx) { this->at(idx) = rhs; } + } + + __device__ __host__ T& operator[](int offset) + { + return reinterpret_cast(this->buffer[offset]); + } + + __device__ __host__ const T& operator[](int offset) const + { + return reinterpret_cast(this->buffer[offset]); + } + + __device__ __host__ T& at(int offset) { return reinterpret_cast(this->buffer[offset]); } + + __device__ __host__ const T& at(int offset) const + { + return reinterpret_cast(this->buffer[offset]); + } + + __device__ __host__ AlignedArray operator+(const AlignedArray& rhs) const + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < kElements; ++idx) { ret[idx] = this->at(idx) + rhs.at(idx); } + + return ret; + } + + __device__ __forceinline__ void clear() + { +#pragma unroll + for (int idx = 0; idx < kElements; ++idx) { this->at(idx) = Element(0); } + } + + Element buffer[N]; +}; + +template +struct reduce_max { + __device__ __forceinline__ T operator()(const T& lhs, const T& rhs) + { + return lhs > rhs ? lhs : rhs; + } +}; + +template +struct reduce_min { + __device__ __forceinline__ T operator()(const T& lhs, const T& rhs) + { + return lhs < rhs ? lhs : rhs; + } +}; + +template +struct subtract { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& rhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] - rhs; } + + return ret; + } +}; + +template +struct plus { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& rhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] + rhs; } + + return ret; + } +}; + +template +struct multiply { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& rhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] * rhs; } + + return ret; + } +}; + +template +struct clamp { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& min_val, + const T& max_val) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { + ret[idx] = reduce_max()(reduce_min()(lhs[idx], max_val), min_val); + } + + return ret; + } +}; + +template +struct round_int; + +template +struct round_int { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = hrint(lhs[idx]); } + + return ret; + } +}; + +template +struct divide { + __device__ __forceinline__ AlignedArray operator()(const AlignedArray& lhs, + const T& rhs) + { + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; ++idx) { ret[idx] = lhs[idx] / rhs; } + + return ret; + } +}; + +template +__device__ __forceinline__ T to_scalar(const AlignedArray& data) +{ + Reducer re; + T res = data[0]; + +#pragma unroll + for (int idx = 1; idx < N; ++idx) { res = re(res, data[idx]); } + + return res; +} + +template +__device__ __forceinline__ AlignedArray int4_to_half( + const AlignedArray& data) +{ + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N * 2; idx += 2) { + ret[idx] = half(int(data[idx / 2] >> 4)); + ret[idx + 1] = half(int(data[idx / 2] & 0xf)); + } + + return ret; +} + +__global__ void dequantize_int4_to_half(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size) +{ + using AccessType = AlignedArray; + using AccessTypeOut = AlignedArray; + + for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_group * group_size / 8; + idx += blockDim.x * gridDim.x) { + int id_group = idx / (group_size / 8); + AccessType value = reinterpret_cast(data_in)[idx]; + half scale = scale_buffer[id_group]; + half min_value = min_val_buffer[id_group]; + + AccessTypeOut output = int4_to_half(value); + output = divide()(output, scale); + output = plus()(output, min_value); + + reinterpret_cast(data_out)[idx] = output; + } +} + +void launch_dequantize_int4_to_half_experimental(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size, + cudaStream_t stream) +{ + int num_warp = num_group / 4; + int num_block = num_warp / 8; // 256 trd / block + + dequantize_int4_to_half<<>>( + data_in, data_out, scale_buffer, min_val_buffer, num_group, group_size); +} + +template +__device__ __forceinline__ AlignedArray int8_to_half(const AlignedArray& data) +{ + AlignedArray ret; + +#pragma unroll + for (int idx = 0; idx < N; idx += 1) { ret[idx] = half(int(data[idx])); } + + return ret; +} + +__global__ void dequantize_int8_to_half(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size) +{ + using AccessType = AlignedArray; + using AccessTypeOut = AlignedArray; + + for (int idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_group * group_size / 8; + idx += blockDim.x * gridDim.x) { + int id_group = idx / (group_size / 8); + AccessType value = reinterpret_cast(data_in)[idx]; + half scale = scale_buffer[id_group]; + half min_value = min_val_buffer[id_group]; + + AccessTypeOut output = int8_to_half(value); + output = divide()(output, scale); + output = plus()(output, min_value); + + reinterpret_cast(data_out)[idx] = output; + } +} + +void launch_dequantize_int8_to_half_experimental(uint8_t* data_in, + half* data_out, + half* scale_buffer, + half* min_val_buffer, + int num_group, + int group_size, + cudaStream_t stream) +{ + int num_warp = num_group / 4; + int num_block = num_warp / 8; // 256 trd / block + + dequantize_int8_to_half<<>>( + data_in, data_out, scale_buffer, min_val_buffer, num_group, group_size); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu new file mode 100644 index 0000000000000000000000000000000000000000..a4b6096c81af102cab9f71a946024f7117902e2e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/quantization/swizzled_quantize.cu @@ -0,0 +1,427 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "dequantization_utils.h" +#include "memory_access_utils.h" +#include "quantization_utils.h" +#include "reduction_utils.h" + +using rop = reduce::ROpType; + +namespace swiz_quant { +constexpr int max_threads = 512; +constexpr int min_threads = 32; + +constexpr int step_granularity = 2; +constexpr int h_per_step = step_granularity * quantize::h_per_load; +} // namespace swiz_quant + +template +__global__ void swizzled_quant_kernel(int8_t* quantized_data, + float* quantized_scales, + const __half* uncompressed_data, + int elems_per_group, + int nodes, + int devices_per_node) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // Indexing offsets, same as normal quantization for in-case + const int block_rank = blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y; + const int block_offset = block_rank * elems_per_group; + const int elem_offset = tb.thread_index().x * quantize::h_per_load; + const int base_offset = block_offset + elem_offset; + const int stride = tb.size() * quantize::h_per_load; + const __half* input_base = uncompressed_data + base_offset; + + // Local buffer + __half2 local_buffer[totalChunks * quantize::h2_per_load]; + + quantize::GroupStats stats; +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; + + mem_access::load_global( + iteration_buffer, input_base + i * stride, elem_offset + i * stride < elems_per_group); + +#pragma unroll + for (int j = 0; j < quantize::h2_per_load; j++) { stats.update(iteration_buffer[j]); } + } + + auto params = stats.template get_params(tb, warp); + + const int partition_id = blockIdx.z; + const int partition_offset = partition_id / devices_per_node; + const int partition_base = (partition_id % devices_per_node) * nodes; + const int pipelining_offset = blockIdx.y * (devices_per_node * nodes); + const int output_partition = (pipelining_offset + partition_base + partition_offset); + + constexpr int out_scalar_effect = 8 / numBits; + const int out_block_rank = output_partition * gridDim.x + blockIdx.x; + const int out_block_offset = out_block_rank * elems_per_group / out_scalar_effect; + const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect; + int8_t* out_base = quantized_data + out_base_offset; + + const int out_stride = stride / out_scalar_effect; + constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect; + + if (tb.thread_index().x == 0) { params.store(quantized_scales, out_block_rank); } + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + if (i * stride + elem_offset < elems_per_group) { + int8_t local_output[quantize::h_per_load / out_scalar_effect]; + quantize::_chunk( + local_output, local_buffer + i * quantize::h2_per_load, params); + mem_access::store_global(out_base + i * out_stride, local_output); + } + } +} + +#define LAUNCH_SWIZZLE_QUANT(total_chunks, threads) \ + swizzled_quant_kernel<<>>( \ + q_data, q_scales, input_data, elems_per_group, nodes, devices_per_node); + +/* +Swizzled quantization reorganizes the quantized groups in order to better facilitate +communication. As an example of the partitioning scheme we have the following example +of 2 node, 4 device swizzling: + + --- --- --- --- --- --- --- --- +| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | + --- --- --- --- --- --- --- --- +becomes + --- --- --- --- --- --- --- --- +| 0 | 4 | 1 | 5 | 2 | 6 | 3 | 7 | + --- --- --- --- --- --- --- --- + +Multiple quantization groups may be mapped into a single partition. In order to better support +later pipelining, we may also perform an additional slicing. In two-way slicing, for instance, +the first halves of each partition are concatenated. +*/ + +template +void launch_swizzled_quant_impl(int8_t* q_data, + float* q_scales, + const __half* input_data, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream) +{ + const int one_step_threads = + next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step); + const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads + : swiz_quant::max_threads; + const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads + : max_threads; + + dim3 block(threads); + const int groups_per_partition = groups / (nodes * devices_per_node); + assert(groups_per_partition % pipelining == 0); + const int contiguous_groups = groups_per_partition / pipelining; + const int partitions = nodes * devices_per_node; + dim3 grid(contiguous_groups, pipelining, partitions); + + const int elems_per_step = threads * swiz_quant::h_per_step; + const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step); + const int total_unroll = external_unroll * swiz_quant::step_granularity; + + assert(total_unroll % 2 == 0); + + if (threads == 32) { + LAUNCH_SWIZZLE_QUANT(2, 32); + } else if (threads == 64) { + LAUNCH_SWIZZLE_QUANT(2, 64); + } else if (threads == 128) { + LAUNCH_SWIZZLE_QUANT(2, 128); + } else if (threads == 256) { + LAUNCH_SWIZZLE_QUANT(2, 256); + } else if (threads == 512) { + if (total_unroll == 2) { + LAUNCH_SWIZZLE_QUANT(2, 512); + } else if (total_unroll == 4) { + LAUNCH_SWIZZLE_QUANT(4, 512); + } else if (total_unroll == 6) { + LAUNCH_SWIZZLE_QUANT(6, 512); + } else if (total_unroll == 8) { + LAUNCH_SWIZZLE_QUANT(8, 512); + } else if (total_unroll == 10) { + LAUNCH_SWIZZLE_QUANT(10, 512); + } + } +} + +#define DISPATCH_SWIZZLE_QUANT(num_bits, qtype) \ + launch_swizzled_quant_impl(q_data, \ + q_scales, \ + input_data, \ + groups, \ + elems_per_group, \ + pipelining, \ + nodes, \ + devices_per_node, \ + stream); + +void launch_swizzled_quant(int8_t* q_data, + float* q_scales, + const __half* input_data, + int num_bits, + quantize::Type q_type, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream) +{ + if (num_bits == 4) { + if (q_type == quantize::Type::Asymmetric) { + DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Asymmetric); + } else if (q_type == quantize::Type::Symmetric) { + DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Symmetric); + } + } else if (num_bits == 8) { + if (q_type == quantize::Type::Asymmetric) { + DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Asymmetric); + } else if (q_type == quantize::Type::Symmetric) { + DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Symmetric); + } + } +} + +template +__global__ void loco_swizzled_quant_kernel(int8_t* quantized_data, + float* quantized_scales, + const __half* uncompressed_data, + __half* error_feedback, + const float err_beta, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node) +{ + cg::thread_block tb = cg::this_thread_block(); + cg::thread_block_tile warp = cg::tiled_partition(tb); + + // Indexing offsets, same as normal quantization for in-case + const int block_rank_data = + blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y; + const int block_offset_data = block_rank_data * elems_per_group; + const int elem_offset = tb.thread_index().x * quantize::h_per_load; + const int base_offset_data = block_offset_data + elem_offset; + const int stride = tb.size() * quantize::h_per_load; + const __half* uncompressed_data_base = uncompressed_data + base_offset_data; + + const int partition_id = blockIdx.z; + const int partition_offset = partition_id / devices_per_node; + const int partition_base = (partition_id % devices_per_node) * nodes; + const int pipelining_offset = blockIdx.y * (devices_per_node * nodes); + const int output_partition = (pipelining_offset + partition_base + partition_offset); + const int block_rank_err = output_partition * gridDim.x + blockIdx.x; + + const int block_offset_err = block_rank_err * elems_per_group; + const int base_offset_err = block_offset_err + elem_offset; + __half* error_feedback_base = error_feedback + base_offset_err; + + __half2 local_buffer[totalChunks * quantize::h2_per_load]; + __half2 err_buffer[totalChunks * quantize::h2_per_load]; + + quantize::GroupStats stats; + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; + __half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load; + const int i_stride = i * stride; + bool do_loads = (elem_offset + i_stride) < elems_per_group; + + mem_access::load_global( + iteration_buffer, uncompressed_data_base + i_stride, do_loads); + + mem_access::load_global( + iter_err_buffer, error_feedback_base + i_stride, do_loads); + +#pragma unroll + for (int j = 0; j < quantize::h2_per_load; j++) { + iteration_buffer[j] = __hadd2(iteration_buffer[j], iter_err_buffer[j]); + stats.update(iteration_buffer[j]); + } + } + + auto params = stats.template get_params(tb, warp); + + // Initialize dequantization parameters based on params + auto de_params = params; + de_params.scale = 1.0f / params.scale; + if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; } + + if (threadIdx.x == 0) { params.store(quantized_scales, block_rank_err); } + + constexpr int out_scalar_effect = 8 / numBits; + const int out_block_offset = block_rank_err * elems_per_group / out_scalar_effect; + const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect; + int8_t* out_base = quantized_data + out_base_offset; + + const int out_stride = stride / out_scalar_effect; + constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect; + +#pragma unroll + for (int i = 0; i < totalChunks; i++) { + const int i_stride = i * stride; + __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load; + __half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load; + + if (i_stride + elem_offset < elems_per_group) { + int8_t local_output[quantize::h_per_load / out_scalar_effect]; + quantize::_chunk(local_output, iteration_buffer, params); + mem_access::store_global(out_base + i * out_stride, local_output); + + // Dequantize the quantized output to compute the dequantized value + __half2 dequant_buffer[quantize::h2_per_load]; + dequantize::chunk(dequant_buffer, local_output, de_params); + +// Compute new error: sum - dequant_buffer +#pragma unroll + for (int k = 0; k < quantize::h2_per_load; k++) { + // __half2 to float2 + float2 iter_buf_f = __half22float2(iteration_buffer[k]); + float2 dequant_buf_f = __half22float2(dequant_buffer[k]); + + // Update within float precision + float2 new_error_f; + new_error_f.x = iter_buf_f.x - dequant_buf_f.x; + new_error_f.y = iter_buf_f.y - dequant_buf_f.y; + + float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]); + + iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x; + iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y; + + // float2 back to __half2 + iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f); + } + __half2* error_feedback_base_h2 = reinterpret_cast<__half2*>(error_feedback_base); + mem_access::store_global(error_feedback_base_h2 + i_stride / 2, + iter_err_buffer); + } + } +} + +#define LAUNCH_LOCO_SWIZZLE_QUANT(total_chunks, threads) \ + loco_swizzled_quant_kernel \ + <<>>(output_data, \ + params, \ + input_data, \ + error_feedback, \ + err_beta, \ + groups, \ + elems_per_group, \ + pipelining, \ + nodes, \ + devices_per_node); + +template +void launch_loco_swizzled_quant_impl(int8_t* output_data, + float* params, + const __half* input_data, + __half* error_feedback, + const float err_beta, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream) +{ + const int one_step_threads = + next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step); + const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads + : swiz_quant::max_threads; + const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads + : max_threads; + + dim3 block(threads); + const int groups_per_partition = groups / (nodes * devices_per_node); + assert(groups_per_partition % pipelining == 0); + const int contiguous_groups = groups_per_partition / pipelining; + const int partitions = nodes * devices_per_node; + dim3 grid(contiguous_groups, pipelining, partitions); + + const int elems_per_step = threads * swiz_quant::h_per_step; + const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step); + const int total_unroll = external_unroll * swiz_quant::step_granularity; + + assert(total_unroll % 2 == 0); + + if (threads == 32) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 32); + } else if (threads == 64) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 64); + } else if (threads == 128) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 128); + } else if (threads == 256) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 256); + } else if (threads == 512) { + if (total_unroll == 2) { + LAUNCH_LOCO_SWIZZLE_QUANT(2, 512); + } else if (total_unroll == 4) { + LAUNCH_LOCO_SWIZZLE_QUANT(4, 512); + } else if (total_unroll == 6) { + LAUNCH_LOCO_SWIZZLE_QUANT(6, 512); + } else if (total_unroll == 8) { + LAUNCH_LOCO_SWIZZLE_QUANT(8, 512); + } else if (total_unroll == 10) { + LAUNCH_LOCO_SWIZZLE_QUANT(10, 512); + } + } +} + +#define DISPATCH_LOCO_SWIZZLE_QUANT(num_bits, qtype) \ + launch_loco_swizzled_quant_impl(output_data, \ + params, \ + input_data, \ + error_feedback, \ + err_beta, \ + groups, \ + elems_per_group, \ + pipelining, \ + nodes, \ + devices_per_node, \ + stream); + +void launch_loco_swizzled_quant(int8_t* output_data, + float* params, + const __half* input_data, + __half* error_feedback, + const float err_beta, + int num_bits, + quantize::Type q_type, + int groups, + int elems_per_group, + int pipelining, + int nodes, + int devices_per_node, + cudaStream_t stream) +{ + if (num_bits == 4) { + if (q_type == quantize::Type::Asymmetric) { + DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Asymmetric); + } else if (q_type == quantize::Type::Symmetric) { + DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Symmetric); + } + } else if (num_bits == 8) { + if (q_type == quantize::Type::Asymmetric) { + DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Asymmetric); + } else if (q_type == quantize::Type::Symmetric) { + DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Symmetric); + } + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp new file mode 100644 index 0000000000000000000000000000000000000000..352306ba26128b96236f43764c26d2f4191eb391 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/sparse_attention/utils.cpp @@ -0,0 +1,127 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a + https:github.com/ptillet/torch-blocksparse/blob/master/csrc/utils.cpp +*/ + +#include +#include +#include +#include +#ifdef _OPENMP +#include +#endif + +typedef std::vector> ret_t; + +void segment_blocks(torch::Tensor layout, + torch::Tensor idx, + torch::Tensor scratch, + int max_width, + ret_t& ret) +{ + size_t H = layout.size(0); + size_t M = layout.size(1); + size_t N = layout.size(2); + torch::Tensor tmp = torch::zeros_like(layout); + + auto _tmp = tmp.accessor(); + auto _layout = layout.accessor(); + auto _idx = idx.accessor(); + auto _scratch = scratch.accessor(); + std::vector current(H, 0); + +#ifdef _OPENMP +#pragma omp parallel for +#endif + for (size_t h = 0; h < H; h++) { + // surrounding indices + std::vector ii_left(max_width, -1); + std::vector> ii_top(max_width, std::vector(N, -1)); + + for (size_t m = 0; m < M; m++) { + for (size_t n = 0; n < N; n++) { + int v = _layout[h][m][n]; + if (v == 0) continue; + int n_left = ii_left[max_width - 1]; + int m_top = ii_top[max_width - 1][n]; + int top = (m_top >= 0) ? _tmp[h][m_top][n] : 0; + int left = (n_left >= 0) ? _tmp[h][m][n_left] : 0; + int topleft = (m_top >= 0 && n_left >= 0) ? _tmp[h][m_top][n_left] : 0; + int width = std::min(left, std::min(top, topleft)) + 1; + + // reset width if blocks cannot be + // packed together (i.e., there's a 1 "in the middle") + for (int nn = n_left + 1; nn < n; nn++) + if (ii_top[max_width - 1][nn] > ii_top[max_width - 1][n]) width = 1; + _tmp[h][m][n] = width; + + // update n_left ring buffer + for (int k = 0; k < max_width - 1; k++) ii_left[k] = ii_left[k + 1]; + ii_left[max_width - 1] = n; + + // update ii_top ring buffer + for (int k = 0; k < max_width - 1; k++) ii_top[k][n] = ii_top[k + 1][n]; + ii_top[max_width - 1][n] = m; + + // block is too small -- skip + if (width != max_width) continue; + + // retained blocks are set to zeros + for (size_t km = 0; km < max_width; km++) + for (size_t kn = 0; kn < max_width; kn++) { + int mm = ii_top[km][n]; + int nn = ii_left[kn]; + if (mm < 0 || nn < 0) continue; + _layout[h][mm][nn] = 0; + _tmp[h][mm][nn] = 0; + _scratch[h][current[h]][0] = (int)h; + _scratch[h][current[h]][1] = (int)mm; + _scratch[h][current[h]][2] = (int)nn; + _scratch[h][current[h]][3] = _idx[h][mm][nn]; + current[h]++; + } + } + } + } + std::vector to_cat; + for (size_t h = 0; h < H; h++) + if (current[h] > 0) to_cat.push_back(scratch[h].slice(0, 0, current[h])); + if (!to_cat.empty()) ret.push_back({max_width, torch::cat(to_cat)}); +} + +ret_t sdd_segment(torch::Tensor layout, int start_width) +{ + ret_t ret; + + // block index + torch::Tensor idx = torch::zeros_like(layout); + int current = 0; + int64_t H = layout.size(0); + int64_t M = layout.size(1); + int64_t N = layout.size(2); + auto _layout = layout.accessor(); + auto _idx = idx.accessor(); + for (int64_t h = 0; h < H; h++) + for (int64_t m = 0; m < M; m++) + for (int64_t n = 0; n < N; n++) { + if (_layout[h][m][n] == 0) continue; + _idx[h][m][n] = current++; + } + + // scratch memory + torch::Tensor scratch = torch::empty({H, layout.sum().item(), 4}, layout.dtype()); + + for (int max_width = start_width; max_width > 0; max_width /= 2) + segment_blocks(layout, idx, scratch, max_width, ret); + return ret; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("sdd_segment", &sdd_segment, "SDD segmentation handler"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu b/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu new file mode 100644 index 0000000000000000000000000000000000000000..d831b372b65f398b43c5cd343b2bf2db67562f56 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu @@ -0,0 +1,149 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include "memory_access_utils.h" +#include "spatial_cuda_layers.h" + +/* +Fused bias add variants +*/ + +namespace badd_opt { +constexpr int threads = 256; +constexpr int steps = 2; +constexpr int granularity = 16; +constexpr int vals_per_h = granularity / sizeof(__half); +constexpr int vals_per_h2 = granularity / sizeof(__half2); +constexpr int vals_per_block = threads * steps * vals_per_h; +constexpr int stride = vals_per_h * threads; +} // namespace badd_opt + +__global__ void opt_bias_add(__half* result, + const __half* activation, + const __half* bias, + int seq_len, + int channels) +{ + const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; + const int stride = badd_opt::vals_per_h * badd_opt::threads; + + for (int i = 0; i < badd_opt::steps; i++) { + if (id + i * badd_opt::stride < seq_len * channels) { + __half2 act_buffer[badd_opt::vals_per_h2]; + __half2 bias_buffer[badd_opt::vals_per_h2]; + + mem_access::load_global(act_buffer, + activation + id + i * stride); + mem_access::load_global(bias_buffer, + bias + ((id + i * stride) % channels)); + + for (int j = 0; j < badd_opt::vals_per_h2; j++) { act_buffer[j] += bias_buffer[j]; } + + mem_access::store_global(result + id + i * stride, act_buffer); + } + } +} + +__global__ void opt_bias_add_add(__half* result, + const __half* activation, + const __half* bias, + const __half* other, + int seq_len, + int channels) +{ + const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; + const int stride = badd_opt::vals_per_h * badd_opt::threads; + + for (int i = 0; i < badd_opt::steps; i++) { + if (id + i * badd_opt::stride < seq_len * channels) { + __half2 act_buffer[badd_opt::vals_per_h2]; + __half2 bias_buffer[badd_opt::vals_per_h2]; + __half2 other_buffer[badd_opt::vals_per_h2]; + + mem_access::load_global(act_buffer, + activation + id + i * stride); + mem_access::load_global(bias_buffer, + bias + ((id + i * stride) % channels)); + mem_access::load_global(other_buffer, other + id + i * stride); + + for (int j = 0; j < badd_opt::vals_per_h2; j++) { + act_buffer[j] += bias_buffer[j] + other_buffer[j]; + } + + mem_access::store_global(result + id + i * stride, act_buffer); + } + } +} + +__global__ void opt_bias_add_bias_add(__half* result, + const __half* activation, + const __half* bias, + const __half* other, + const __half* other_bias, + int seq_len, + int channels) +{ + const int id = blockIdx.x * badd_opt::vals_per_block + threadIdx.x * badd_opt::vals_per_h; + const int stride = badd_opt::vals_per_h * badd_opt::threads; + + for (int i = 0; i < badd_opt::steps; i++) { + if (id + i * badd_opt::stride < seq_len * channels) { + __half2 act_buffer[badd_opt::vals_per_h2]; + __half2 bias_buffer[badd_opt::vals_per_h2]; + __half2 other_buffer[badd_opt::vals_per_h2]; + __half2 other_bias_buffer[badd_opt::vals_per_h2]; + + mem_access::load_global(act_buffer, + activation + id + i * stride); + mem_access::load_global(bias_buffer, + bias + ((id + i * stride) % channels)); + mem_access::load_global(other_buffer, other + id + i * stride); + mem_access::load_global( + other_bias_buffer, other_bias + ((id + i * stride) % channels)); + + for (int j = 0; j < badd_opt::vals_per_h2; j++) { + act_buffer[j] = + (act_buffer[j] + bias_buffer[j]) + (other_buffer[j] + other_bias_buffer[j]); + } + + mem_access::store_global(result + id + i * stride, act_buffer); + } + } +} + +void launch_opt_bias_add(__half* result, + const __half* activation, + const __half* bias, + const __half* other, + const __half* other_bias, + int batch_size, + int seq_len, + int channels, + cudaStream_t stream) +{ + // Should evaluate `true` for reasonable hidden sizes + assert(channels % badd_opt::vals_per_h == 0); + + const int effective_seq_len = batch_size * seq_len; + const int vals = effective_seq_len * channels; + + dim3 block(badd_opt::threads); + dim3 grid((vals + badd_opt::vals_per_block - 1) / badd_opt::vals_per_block); + + if (!other) { + // We shouldn't have a bias if there's no activation + assert(!other_bias); + + opt_bias_add<<>>( + result, activation, bias, effective_seq_len, channels); + } else if (!other_bias) { + opt_bias_add_add<<>>( + result, activation, bias, other, effective_seq_len, channels); + } else { + opt_bias_add_bias_add<<>>( + result, activation, bias, other, other_bias, effective_seq_len, channels); + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp new file mode 100644 index 0000000000000000000000000000000000000000..cbf6636a6ee4ac5769a88b7cab7d2a222f00d5cc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp @@ -0,0 +1,112 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include +#include "spatial_cuda_layers.h" + +ChannelsLastProblem dimension_problem(at::Tensor& input) +{ + ChannelsLastProblem dims; + + if (input.dim() == 4) { + // In some sense this is unsafe (and a reflection of the assumptions made inside + // the C10 options checker). Basically, there's no great way to be sure that + // a tensor is in channels last because a 1x1 image will appear to be in channels + // last even when it isn't. + assert(input.is_contiguous(at::MemoryFormat::ChannelsLast)); + dims.batch_size = input.size(0); + dims.seq_len = input.size(2) * input.size(3); + dims.channels = input.size(1); + } else { + assert(input.is_contiguous()); + dims.batch_size = input.size(0); + dims.seq_len = input.size(1); + dims.channels = input.size(2); + } + + return dims; +} + +at::Tensor seq_unroll_bias_add(at::Tensor& input, at::Tensor& bias) +{ + assert(input.dtype() == at::kHalf); + + // TODO(cmikeh2): Should probably refactor this into a more portable + // description, since it does generalize for channels-last + ChannelsLastProblem problem = dimension_problem(input); + + auto output = at::empty_like(input); + + launch_opt_bias_add((__half*)output.data_ptr(), + (const __half*)input.data_ptr(), + (const __half*)bias.data_ptr(), + nullptr, + nullptr, + problem.batch_size, + problem.seq_len, + problem.channels, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +at::Tensor seq_bias_add_add(at::Tensor& input, at::Tensor& bias, at::Tensor& other) +{ + assert(input.dtype() == at::kHalf); + + // TODO(cmikeh2): Should probably refactor this into a more portable + // description, since it does generalize for channels-last + ChannelsLastProblem problem = dimension_problem(input); + + auto output = at::empty_like(input); + + launch_opt_bias_add((__half*)output.data_ptr(), + (const __half*)input.data_ptr(), + (const __half*)bias.data_ptr(), + (const __half*)other.data_ptr(), + nullptr, + problem.batch_size, + problem.seq_len, + problem.channels, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +at::Tensor seq_bias_add_bias_add(at::Tensor& input, + at::Tensor& bias, + at::Tensor& other, + at::Tensor& other_bias) +{ + assert(input.dtype() == at::kHalf); + + // TODO(cmikeh2): Should probably refactor this into a more portable + // description, since it does generalize for channels-last + ChannelsLastProblem problem = dimension_problem(input); + + auto output = at::empty_like(input); + + launch_opt_bias_add((__half*)output.data_ptr(), + (const __half*)input.data_ptr(), + (const __half*)bias.data_ptr(), + (const __half*)other.data_ptr(), + (const __half*)other_bias.data_ptr(), + problem.batch_size, + problem.seq_len, + problem.channels, + at::cuda::getCurrentCUDAStream()); + + return output; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("nhwc_bias_add", &seq_unroll_bias_add); + m.def("nhwc_bias_add_add", &seq_bias_add_add); + m.def("nhwc_bias_add_bias_add", &seq_bias_add_bias_add); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h new file mode 100644 index 0000000000000000000000000000000000000000..4f56f89f16cd1a200f057ebd651575224be23b7b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/spatial/includes/spatial_cuda_layers.h @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#if __CUDA_ARCH__ >= 530 +#define HALF_PRECISION_AVAILABLE = 1 +#endif + +#ifdef __HIP_PLATFORM_AMD__ +#include +#else +#include +#endif + +#include +#include + +/*********** Group Norm Kernels, Structs, and Helpers ************/ + +struct { + int64_t batch_size; + int64_t seq_len; + int64_t channels; +} typedef ChannelsLastProblem; + +void launch_opt_bias_add(__half* result, + const __half* activation, + const __half* bias, + const __half* other, + const __half* other_bias, + int batch_size, + int seq_len, + int channels, + cudaStream_t stream); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/utils/flatten_unflatten.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/utils/flatten_unflatten.cpp new file mode 100644 index 0000000000000000000000000000000000000000..ab95ee1914647c85e3e0ad5ebf625877d57a5ad6 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/utils/flatten_unflatten.cpp @@ -0,0 +1,29 @@ +// 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 +#include +// https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/tensor_flatten.h + +at::Tensor flatten(std::vector tensors) +{ + return torch::utils::flatten_dense_tensors(tensors); +} + +std::vector unflatten(at::Tensor flat, std::vector tensors) +{ + return torch::utils::unflatten_dense_tensors(flat, tensors); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("flatten", &flatten, "Flatten dense tensors"); + m.def("unflatten", &unflatten, "Unflatten dense tensors"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp new file mode 100644 index 0000000000000000000000000000000000000000..dc727f8fa2168ef7897f685de6ac5d80aaad8309 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adagrad/cpu_adagrad.cpp @@ -0,0 +1,196 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include "cpu_adagrad.h" +#include +#include +#include +#include +#include +#include + +static std::unordered_map> s_optimizers; + +// C++ interface + +void Adagrad_Optimizer::Step_1(float* _params, + float* grads, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<1>( + &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); +#endif + if (_param_size > rounded_size) { + float step_size = -1 * _alpha; + ds_half_precision_t* grads_cast_h; + ds_half_precision_t* params_cast_h; + if (half_precision) { + grads_cast_h = reinterpret_cast(grads); + params_cast_h = reinterpret_cast(_params); + } + 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 = half_precision ? (float)grads_cast_h[k] : grads[k]; + float param = half_precision ? (float)params_cast_h[k] : _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; + if (half_precision) + params_cast_h[k] = (ds_half_precision_t)param; + else + _params[k] = param; + // STORE UPDATE TERM TO GRAD'S MEMORY + grads[k] = grad * step_size; + _exp_avg_sq[k] = variance; + } + } + } +} + +void Adagrad_Optimizer::Step_4(float* _params, + float* grads, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<4>( + &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); +#endif + if (_param_size > rounded_size) + Step_1((_params + rounded_size), + (grads + rounded_size), + (_exp_avg_sq + rounded_size), + (_param_size - rounded_size), + (dev_params != nullptr ? (dev_params + rounded_size) : dev_params), + half_precision); +} + +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(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; +} + +void Adagrad_Optimizer::Step_8(float* _params, + float* grads, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) + Step_AVX<8>( + &rounded_size, _params, grads, _exp_avg_sq, _param_size, dev_params, half_precision); +#endif + if (_param_size > rounded_size) + Step_4((_params + rounded_size), + (grads + rounded_size), + (_exp_avg_sq + rounded_size), + (_param_size - rounded_size), + (dev_params != nullptr ? (dev_params + rounded_size) : dev_params), + half_precision); +} + +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(); + + float* params_ptr = (float*)params_c.data_ptr(); + float* grads_ptr = (float*)grads_c.data_ptr(); + float* exp_avg_sq_ptr = (float*)exp_avg_sq_c.data_ptr(); + + std::shared_ptr opt = + std::static_pointer_cast(s_optimizers[optimizer_id]); + opt->IncrementStep(step); + opt->update_state(lr, epsilon, weight_decay); + opt->Step_8(params_ptr, grads_ptr, exp_avg_sq_ptr, params_c.numel()); + + return 0; +} + +int ds_adagrad_step_plus_copy(int optimizer_id, + size_t step, + float lr, + float epsilon, + float weight_decay, + torch::Tensor& params, + torch::Tensor& grads, + torch::Tensor& exp_avg_sq, + torch::Tensor& gpu_params) +{ + assert(false); + 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("adagrad_update_copy", + &ds_adagrad_step_plus_copy, + "DeepSpeed CPU Adagrad update and param copy (C++)"); + m.def("create_adagrad", &create_adagrad_optimizer, "DeepSpeed CPU Adagrad (C++)"); + m.def("destroy_adagrad", &destroy_adagrad_optimizer, "DeepSpeed CPU Adagrad destroy (C++)"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp new file mode 100644 index 0000000000000000000000000000000000000000..13b390248608b046dab443f85346b5446a47d722 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/fused_adam_frontend.cpp @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +void multi_tensor_adam_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> 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"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp new file mode 100644 index 0000000000000000000000000000000000000000..0720a020247a3ebf80fc231e11333ea56cb09924 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_adam.dp.cpp @@ -0,0 +1,159 @@ +// 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 +#include +#include + +#include + +#include +#include "multi_tensor_apply.dp.hpp" +#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 +struct AdamFunctor { + __inline__ __attribute__((always_inline)) 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) + { + auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); + int tensor_loc = tl.block_to_tensor[item_ct1.get_group(2)]; + + int chunk_idx = tl.block_to_chunk[item_ct1.get_group(2)]; + int 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 (int i_start = 0; i_start < n && i_start < chunk_size; + i_start += item_ct1.get_local_range(2) * 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 + item_ct1.get_local_id(2) + ii * item_ct1.get_local_range(2); + 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 = sycl::sqrt(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 = sycl::sqrt(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 + item_ct1.get_local_id(2) + ii * item_ct1.get_local_range(2); + 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> 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); + } + + // Assume single type across p,g,m1,m2 now + 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(), + beta1, + beta2, + bias_correction1, + bias_correction2, + epsilon, + lr, + (adamMode_t)mode, + weight_decay);) +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp new file mode 100644 index 0000000000000000000000000000000000000000..14a130e2a23e906ba9b305ee3968aa4e27c17cf8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/adam/multi_tensor_apply.dp.hpp @@ -0,0 +1,226 @@ +// 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 +#include +#include +#include +#include +#include "compat.h" + +#include +#include +#include + +namespace at { +namespace cuda { +sycl::queue* getCurrentCUDAStream() +{ + c10::xpu::XPUStream stream = c10::xpu::getCurrentXPUStream(); + auto& queue = stream.queue(); + return &queue; +} + +sycl::queue* getStreamFromPool(bool) +{ + // not implemented + return nullptr; +} +} // namespace cuda +} // namespace at +// #include + +// 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 +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 +class multi_tensor_apply_kernel { +public: + multi_tensor_apply_kernel(int chunk_size, + volatile int* noop_flag, + T tl, + U callable, + ArgTypes... args) + : chunk_size(chunk_size), noop_flag(noop_flag), tl(tl), callable(callable), args(args...) + { + } + + // This should be identical to original __global__ function + static void inline __global__function(int chunk_size, + volatile int* noop_flag, + T tl, + U callable, + ArgTypes... args) + { + callable(chunk_size, noop_flag, tl, args...); + } + + // If global function template contains parameter pack, + // we only deal with parameter pack at the end of template parameter list + template + static void inline __tuple_expand_driver(int chunk_size, + volatile int* noop_flag, + T tl, + U callable, + Tuple args, + std::index_sequence) + { + __global__function(chunk_size, noop_flag, tl, callable, std::get(args)...); + } + + // + // Because __global__ function can't really use any reference types, we can sure that args + // are all good behaviors + // + void operator()(sycl::nd_item<3>) const + { + __tuple_expand_driver(chunk_size, + noop_flag, + tl, + callable, + args, + std::make_index_sequence()); + } + +private: + int chunk_size; + volatile int* noop_flag; + T tl; + U callable; + std::tuple args; +}; + +// to make sure multi_tensor_apply_kernel can be used in sycl::buffer +namespace sycl { +template +struct is_device_copyable> : std::true_type {}; +} // namespace sycl + +template +void multi_tensor_apply(int block_size, + int chunk_size, + const at::Tensor& noop_flag, + const std::vector>& 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::kXPU, "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 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++; + + int chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; + + for (int 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; + /* multi_tensor_apply_kernel, T, ArgTypes...> + * fn(chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); */ + if constexpr (sizeof(multi_tensor_apply_kernel( + chunk_size, noop_flag.DATA_PTR(), tl, callable, args...)) < + 2048) { + ((sycl::queue*)(stream)) + ->parallel_for( + sycl::nd_range<3>(sycl::range<3>(1, 1, loc_block_info) * + sycl::range<3>(1, 1, block_size), + sycl::range<3>(1, 1, block_size)), + multi_tensor_apply_kernel( + chunk_size, noop_flag.DATA_PTR(), tl, callable, args...)); + } else { + auto capture = multi_tensor_apply_kernel( + chunk_size, noop_flag.DATA_PTR(), tl, callable, args...); + sycl::buffer params(const_cast(&capture), + sycl::range<1>(1)); + stream->submit([&](sycl::handler& cgh) { + auto device_params = + params.template get_access(cgh); + cgh.parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, loc_block_info) * + sycl::range<3>(1, 1, block_size), + sycl::range<3>(1, 1, block_size)), + [=](sycl::nd_item<3> item) { device_params[0](item); }); + }); + } + 0; + + // 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; + } + } + } + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp new file mode 100644 index 0000000000000000000000000000000000000000..cfd004ef13577c02f958fa18d116131398443562 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/common/custom_cuda_kernel.dp.cpp @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include + +inline void has_capability_or_fail(const sycl::device& dev, + const std::initializer_list& props) +{ + for (const auto& it : props) { + if (dev.has(it)) continue; + switch (it) { + case sycl::aspect::fp64: + throw std::runtime_error("'double' is not supported in '" + + dev.get_info() + "' device"); + break; + case sycl::aspect::fp16: + throw std::runtime_error("'half' is not supported in '" + + dev.get_info() + "' device"); + break; + default: +#define __SYCL_ASPECT(ASPECT, ID) \ + case sycl::aspect::ASPECT: return #ASPECT; +#define __SYCL_ASPECT_DEPRECATED(ASPECT, ID, MESSAGE) __SYCL_ASPECT(ASPECT, ID) +#define __SYCL_ASPECT_DEPRECATED_ALIAS(ASPECT, ID, MESSAGE) + auto getAspectNameStr = [](sycl::aspect AspectNum) -> std::string { + switch (AspectNum) { +#include +#include + default: return "unknown aspect"; + } + }; +#undef __SYCL_ASPECT_DEPRECATED_ALIAS +#undef __SYCL_ASPECT_DEPRECATED +#undef __SYCL_ASPECT + throw std::runtime_error("'" + getAspectNameStr(it) + "' is not supported in '" + + dev.get_info() + "' device"); + } + break; + } +} + +void param_update_kernel(const float* input, sycl::half* output, int size) +{ + auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); + int id = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2); + + if (id < size) { output[id] = (sycl::half)input[id]; } +} + +void launch_param_update(const float* input, sycl::half* output, int size, sycl::queue* stream) +{ + int threads = 1024; + + sycl::range<3> grid_dim(1, 1, (size - 1) / threads + 1); + sycl::range<3> block_dim(1, 1, threads); + + { + has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16}); + stream->parallel_for( + sycl::nd_range<3>(grid_dim * block_dim, block_dim), + [=](sycl::nd_item<3> item_ct1) { param_update_kernel(input, output, size); }); + } +} + +void param_update_kernel_half(const float* input, sycl::half* output, int size) +{ + auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); + int id = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2); + sycl::half2* output_cast = reinterpret_cast(output); + if (id < size) { + float input_f = input[id]; + sycl::half2* input_h = reinterpret_cast(&input_f); + output_cast[id] = *input_h; + } +} + +void launch_param_update_half(const float* input, sycl::half* output, int size, sycl::queue* stream) +{ + int threads = 1024; + size /= 2; + sycl::range<3> grid_dim(1, 1, (size - 1) / threads + 1); + sycl::range<3> block_dim(1, 1, threads); + + { + has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16}); + stream->parallel_for( + sycl::nd_range<3>(grid_dim * block_dim, block_dim), + [=](sycl::nd_item<3> item_ct1) { param_update_kernel_half(input, output, size); }); + } +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h new file mode 100644 index 0000000000000000000000000000000000000000..6d54446d472e9dfa6141436f55e0e623be48acb4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/compat.h @@ -0,0 +1,19 @@ +// 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 +*/ + +#ifndef TORCH_CHECK +#define TORCH_CHECK AT_CHECK +#endif + +#ifdef VERSION_GE_1_3 +#define DATA_PTR data_ptr +#else +#define DATA_PTR data +#endif diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h new file mode 100644 index 0000000000000000000000000000000000000000..660f860917f6c8b345825e6110a7508ba19a584f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adagrad.h @@ -0,0 +1,120 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#define NOMINMAX // Windows idiosyncrasy + // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c + +#include +#include +#include "simd.h" + +typedef unsigned short ds_half_precision_t; + +#define STEP(SPAN) \ + void Step_##SPAN(float* _params, \ + float* grads, \ + float* _exp_avg_sq, \ + size_t _param_size, \ + ds_half_precision_t* dev_param = nullptr, \ + bool half_precision = false); + +class Adagrad_Optimizer { +public: + Adagrad_Optimizer(float alpha = 1e-2, float eps = 1e-8, float weight_decay = 0) + : _alpha(alpha), _eps(eps), _weight_decay(weight_decay) + { + } + ~Adagrad_Optimizer() {} +#if defined(__AVX512__) or defined(__AVX256__) + template + void Step_AVX(size_t* rounded_size, + float* _params, + float* grads, + float* _exp_avg_sq, + size_t param_size, + ds_half_precision_t* dev_param = nullptr, + bool half_precision = false); +#endif + STEP(1) + STEP(4) + STEP(8) + inline void IncrementStep(size_t step) + { + _step++; + if (_step != step) { _step = step; } + } + inline void update_state(float lr, float epsilon, float weight_decay) + { + _alpha = lr; + _eps = epsilon; + _weight_decay = weight_decay; + } + +private: + float _alpha; + float _eps; + float _weight_decay; + + float _betta1_t; + float _betta2_t; + size_t _step; +}; + +#if defined(__AVX512__) or defined(__AVX256__) +template +void Adagrad_Optimizer::Step_AVX(size_t* rounded_size, + float* _params, + float* grads, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t new_rounded_size = 0; + AVX_Data eps_4; + eps_4.data = SIMD_SET(_eps); + + float step_size = -1 * _alpha; + AVX_Data step_size_4; + step_size_4.data = SIMD_SET(step_size); + + AVX_Data weight_decay4; + if (_weight_decay > 0) weight_decay4.data = SIMD_SET(_weight_decay); + new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); + for (size_t t = 0; t < new_rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + i, half_precision); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, grads + i, false); + + AVX_Data variance_4[span]; + simd_load(variance_4, _exp_avg_sq + i, false); + + AVX_Data param_4[span]; + simd_load(param_4, _params + i, half_precision); + + if (_weight_decay > 0) { simd_fma(grad_4, param_4, weight_decay4, grad_4); } + + simd_fma(variance_4, grad_4, grad_4, variance_4); + simd_sqrt(grad_4, variance_4); + simd_add(grad_4, grad_4, eps_4); + simd_div(grad_4, momentum_4, grad_4); + simd_fma(param_4, grad_4, step_size_4, param_4); + + simd_store(_params + i, param_4, half_precision); + simd_store(_exp_avg_sq + i, variance_4, false); + } + } + *rounded_size = new_rounded_size; +} +#endif diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h new file mode 100644 index 0000000000000000000000000000000000000000..7bc0364c569d00d01b9f0b526d21d3c57fb478dc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/cpu_adam.h @@ -0,0 +1,237 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#define NOMINMAX // Windows idiosyncrasy + // https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c + +#include +#include +#include +#include "simd.h" + +#include +typedef unsigned short ds_half_precision_t; + +#define STEP(SPAN) \ + void Step_##SPAN(float* _params, \ + float* grads, \ + float* _exp_avg, \ + float* _exp_avg_sq, \ + size_t _param_size, \ + ds_half_precision_t* dev_param = nullptr, \ + bool half_precision = false); + +class Adam_Optimizer { +public: + Adam_Optimizer(float alpha = 1e-3, + float betta1 = 0.9, + float betta2 = 0.999, + float eps = 1e-8, + float weight_decay = 0, + bool adamw_mode = true) + : _alpha(alpha), + _betta1(betta1), + _betta2(betta2), + _eps(eps), + _weight_decay(weight_decay), + _betta1_t(1.0), + _betta2_t(1.0), + _step(0), + _adamw_mode(adamw_mode) + { + } + ~Adam_Optimizer() {} + +#if defined(__AVX512__) or defined(__AVX256__) + template + void Step_AVX(size_t* rounded_size, + float* _params, + float* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t param_size, + ds_half_precision_t* dev_param = nullptr, + bool half_precision = false); +#endif + STEP(1) + STEP(4) + STEP(8) + inline void IncrementStep(size_t step, float beta1, float beta2) + { + if (beta1 != _betta1 || beta2 != _betta2) { + _step = step; + _betta1 = beta1; + _betta2 = beta2; + _betta1_t = std::pow(_betta1, step); + _betta2_t = std::pow(_betta2, step); + } else { + _step++; + if (_step != step) { + _betta1_t = std::pow(_betta1, step); + _betta2_t = std::pow(_betta2, step); + _step = step; + } else { + _betta1_t *= _betta1; + _betta2_t *= _betta2; + } + } + } + inline void update_state(float lr, float epsilon, float weight_decay, bool bias_correction) + { + _alpha = lr; + _eps = epsilon; + _weight_decay = weight_decay; + + _bias_correction1 = 1.0f; + _bias_correction2 = 1.0f; + if (bias_correction == 1) { + _bias_correction1 = 1 - _betta1_t; + _bias_correction2 = 1 / sqrt(1 - _betta2_t); + } + } + +private: + float _alpha; + float _betta1; + float _betta2; + float _eps; + float _weight_decay; + + float _betta1_t; + float _betta2_t; + size_t _step; + + float _bias_correction1; + float _bias_correction2; + + bool _adamw_mode; +}; + +#if defined(__AVX512__) or defined(__AVX256__) +template +void Adam_Optimizer::Step_AVX(size_t* rounded_size, + float* _params, + float* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t _param_size, + ds_half_precision_t* dev_params, + bool half_precision) +{ + size_t new_rounded_size = 0; + int rshft = half_precision ? 1 : 0; + + AVX_Data betta1_4; + betta1_4.data = SIMD_SET(_betta1); + AVX_Data betta2_4; + betta2_4.data = SIMD_SET(_betta2); + + float betta1_minus1 = 1 - _betta1; + float betta2_minus1 = 1 - _betta2; + AVX_Data betta1_minus1_4; + betta1_minus1_4.data = SIMD_SET(betta1_minus1); + AVX_Data betta2_minus1_4; + betta2_minus1_4.data = SIMD_SET(betta2_minus1); + + AVX_Data bias2_sqrt; + bias2_sqrt.data = SIMD_SET(_bias_correction2); + + AVX_Data eps_4; + eps_4.data = SIMD_SET(_eps); + + float step_size = -1 * _alpha / _bias_correction1; + AVX_Data step_size_4; + step_size_4.data = SIMD_SET(step_size); + + float w_decay = -1 * _alpha * _weight_decay; + AVX_Data weight_decay4; + if (_weight_decay > 0) + weight_decay4.data = (_adamw_mode ? SIMD_SET(w_decay) : SIMD_SET(_weight_decay)); + new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span); + for (size_t t = 0; t < new_rounded_size; t += TILE) { + size_t copy_size = TILE; + if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t; + size_t offset = copy_size + t; +#pragma omp parallel for + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + (i >> rshft), half_precision); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, _exp_avg + i, false); + + AVX_Data variance_4[span]; + simd_load(variance_4, _exp_avg_sq + i, false); + + AVX_Data param_4[span]; + simd_load(param_4, _params + (i >> rshft), half_precision); + + if (_weight_decay > 0 && !_adamw_mode) { + simd_fma(grad_4, param_4, weight_decay4, grad_4); + } + + simd_mul(momentum_4, momentum_4, betta1_4); + simd_fma(momentum_4, grad_4, betta1_minus1_4, momentum_4); + simd_mul(variance_4, variance_4, betta2_4); + simd_mul(grad_4, grad_4, grad_4); + simd_fma(variance_4, grad_4, betta2_minus1_4, variance_4); + simd_sqrt(grad_4, variance_4); + simd_fma(grad_4, grad_4, bias2_sqrt, eps_4); + simd_div(grad_4, momentum_4, grad_4); + + if (_weight_decay > 0 && _adamw_mode) { + simd_fma(param_4, param_4, weight_decay4, param_4); + } + + simd_fma(param_4, grad_4, step_size_4, param_4); + + simd_store(_params + (i >> rshft), param_4, half_precision); + simd_store(_exp_avg + i, momentum_4, false); + simd_store(_exp_avg_sq + i, variance_4, false); + } + } + *rounded_size = new_rounded_size; +} +#endif + +int create_adam_optimizer(int optimizer_id, + float alpha = 1e-3, + float betta1 = 0.9, + float betta2 = 0.999, + float eps = 1e-8, + float weight_decay = 0, + bool adamw_mode = true, + bool should_log = false); + +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); + +int ds_adam_step_plus_copy(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, + torch::Tensor& gpu_params); + +int destroy_adam_optimizer(int optimizer_id); diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h new file mode 100644 index 0000000000000000000000000000000000000000..097e2d8585ccbdb83f1c249c1b6d89272a134999 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/simd.h @@ -0,0 +1,198 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#pragma once + +#if (__x86_64__ || __i386__) +#include +#include +#endif + +#define TILE (128 * 1024 * 1024) +#if defined(__AVX512__) or defined(__AVX256__) + +#define ROUND_DOWN(size, step) ((size) & ~((step) - 1)) + +#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_ADD(x, y) _mm512_add_ps(x, y) +#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_AND(x, y) _mm512_and_ps(x, y) +#define SIMD_ANDNOT(x, y) _mm512_andnot_ps(x, y) +#define SIMD_OR(x, y) _mm512_or_ps(x, y) +#define SIMD_XOR(x, y) _mm512_xor_ps(x, y) +#define SIMD_WIDTH 16 + +#define SIMD_LOAD2(x, h) \ + ((h) ? _mm512_cvtph_ps(_mm256_castps_si256(_mm256_loadu_ps(x))) : _mm512_loadu_ps(x)) +#define SIMD_STORE2(x, d, h) \ + ((h) ? _mm256_store_ps(x, _mm256_castsi256_ps(_mm512_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) \ + : _mm512_storeu_ps(x, d)) + +#define INTV __m256i +#elif 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_ADD(x, y) _mm256_add_ps(x, y) +#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_AND(x, y) _mm256_and_ps(x, y) +#define SIMD_ANDNOT(x, y) _mm256_andnot_ps(x, y) +#define SIMD_OR(x, y) _mm256_or_ps(x, y) +#define SIMD_XOR(x, y) _mm256_xor_ps(x, y) +#define SIMD_WIDTH 8 + +#define SIMD_LOAD2(x, h) \ + ((h) ? _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*)x)) : _mm256_loadu_ps(x)) +#define SIMD_STORE2(x, d, h) \ + ((h) ? _mm_store_ps(x, _mm_castsi128_ps(_mm256_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT))) \ + : _mm256_storeu_ps(x, d)) + +#define INTV __m128i +#endif + +union AVX_Data { +#if defined(__AVX512__) + __m512 data; +#elif defined(__AVX256__) + __m256 data; +#endif + // float data_f[16]; +}; + +template +inline void simd_store(float* dst, AVX_Data* src, bool half_precision) +{ + size_t width = (half_precision ? SIMD_WIDTH / 2 : SIMD_WIDTH); +#pragma unroll + for (size_t i = 0; i < span; ++i) { SIMD_STORE2(dst + width * i, src[i].data, half_precision); } +} +template +inline void simd_load(AVX_Data* dst, float* src, bool half_precision) +{ + size_t width = (half_precision ? 1 : SIMD_WIDTH); +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_LOAD2(src + width * i, half_precision); } +} +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data* src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a[i].data); + } +} +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data src_m_r, AVX_Data src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r.data, src_a.data); + } +} +template +inline void simd_fma(AVX_Data* dst, AVX_Data* src_m_l, AVX_Data* src_m_r, AVX_Data* src_a) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_FMA(src_m_l[i].data, src_m_r[i].data, src_a[i].data); + } +} +template +inline void simd_sqrt(AVX_Data* dst, AVX_Data* src) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_SQRT(src[i].data); } +} +template +inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_add(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ADD(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_mul(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MUL(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_div(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_DIV(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_AND(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_andnot(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { + dst[i].data = SIMD_ANDNOT(src_a_l[i].data, src_a_r[i].data); + } +} +template +inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_or(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_OR(src_a_l[i].data, src_a_r[i].data); } +} +template +inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r.data); } +} +template +inline void simd_xor(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_XOR(src_a_l[i].data, src_a_r[i].data); } +} + +#endif diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h new file mode 100644 index 0000000000000000000000000000000000000000..1897afd1fea248f497acf3fb6ba049baa1ca56c3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/includes/type_shim.h @@ -0,0 +1,155 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* Taken from NVIDIA/apex commit 855808f3fc268e9715d613f3c2e56469d8c986d8 */ +#include +/* #include */ +#include + +// Forward/backward compatibility hack around +// https://github.com/pytorch/pytorch/commit/3aeb78079bcd68282fe9117088e138b77318e288 +// pending more future-proof guidance from upstream. +// struct TypeShim +// { +// const at::Type& payload; +// TypeShim(const at::Type& type) : payload(type) {} +// // Enable trivial conversion to a const at::Type& for pre-3aeb78 +// operator const at::Type&(){ return payload; }; +// // Enable dispatch switch statements to take *this directly for post-3aeb78 +// //operator at::ScalarType(){ return payload.; }; +// }; + +#define DISPATCH_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Half: { \ + using scalar_t_##LEVEL = at::Half; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::BFloat16: { \ + using scalar_t_##LEVEL = at::BFloat16; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +#define DISPATCH_DOUBLE_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Double: { \ + using scalar_t_##LEVEL = double; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Half: { \ + using scalar_t_##LEVEL = at::Half; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::BFloat16: { \ + using scalar_t_##LEVEL = at::BFloat16; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +#define DISPATCH_DOUBLE_AND_FLOAT(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Double: { \ + using scalar_t_##LEVEL = double; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Float: { \ + using scalar_t_##LEVEL = float; \ + __VA_ARGS__; \ + break; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +template +__inline__ __attribute__((always_inline)) T +reduce_block_into_lanes(T* x, + T val, + int lanes = 1, + bool share_result = false) // lanes is intended to be <= 32. +{ + auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>(); + int tid = item_ct1.get_local_id(2) + item_ct1.get_local_id(1) * item_ct1.get_local_range(2); + int blockSize = item_ct1.get_local_range(2) * + item_ct1.get_local_range(1); // blockSize is intended to be a multiple of 32. + + if (blockSize >= 64) { + x[tid] = val; + /* + DPCT1118:1: SYCL group functions and algorithms must be encountered in converged control + flow. You may need to adjust the code. + */ + /* + DPCT1065:6: Consider replacing sycl::nd_item::barrier() with + sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if + there is no access to global memory. + */ + item_ct1.barrier(); + } + +#pragma unroll + for (int i = (blockSize >> 1); i >= 64; i >>= 1) { + if (tid < i) x[tid] = x[tid] + x[tid + i]; + /* + DPCT1118:2: SYCL group functions and algorithms must be encountered in converged control + flow. You may need to adjust the code. + */ + /* + DPCT1065:7: Consider replacing sycl::nd_item::barrier() with + sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if + there is no access to global memory. + */ + item_ct1.barrier(); + } + + T final; + + if (tid < 32) { + if (blockSize >= 64) + final = x[tid] + x[tid + 32]; + else + final = val; + // __SYNCWARP(); + +#pragma unroll + for (int i = 16; i >= lanes; i >>= 1) + final = final + __shfl_down_sync(0xffffffff, final, i); + } + + if (share_result) { + if (tid < lanes) x[tid] = final; // EpilogueOp + // Make sure the smem result is visible to all warps. + /* + DPCT1118:3: SYCL group functions and algorithms must be encountered in converged control + flow. You may need to adjust the code. + */ + /* + DPCT1065:8: Consider replacing sycl::nd_item::barrier() with + sycl::nd_item::barrier(sycl::access::fence_space::local_space) for better performance if + there is no access to global memory. + */ + item_ct1.barrier(); + } + + return final; +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp new file mode 100644 index 0000000000000000000000000000000000000000..fb09bb16efdda5b578c239c35be7ee493dc87c75 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/csrc/xpu/packbits/packing.cpp @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +#include +#include +#include +#include + +using namespace sycl; +using namespace xpu; + +void packbitskernel(const float* input, uint8_t* output, const int input_size, id<1> item_ct1) +{ + // get the sign bit of each float and pack them into byte + int i = item_ct1; + for (int j = 0; j < 8; ++j) { + int k = i * 8 + j; + int bit = k < input_size && (!sycl::signbit(input[k])); + output[i] |= bit << (7 - j); + } +} + +void unpackbitskernel(const uint8_t* input, float* output, id<1> item_ct1) +{ + // use the bit value to set float, bit 0 -> float -1, bit 1 -> float 1 + int i = item_ct1; + output[i] = (float((input[i / 8] >> (7 - i % 8)) & 1) - 0.5) * 2; +} + +sycl::queue get_current_queue(at::Device device) +{ + c10::impl::VirtualGuardImpl impl(device.type()); + c10::Stream _stream = impl.getStreamFromGlobalPool(device, /*isHighPriority=*/false); + sycl::queue queue = xpu::get_queue_from_stream(_stream); + return queue; +} + +/* +pack float tensor into uint8 tensor. Every eight float elements get packed into one uint8 +if float x >= 0, will be packed as a '1' bit, or will be packed as '0' +Arguments: + tensor: A bool tensor that get packed. + input_size: numel of input tensor + rank: device id in order to get corresponding stream +*/ +at::Tensor packbits(at::Tensor tensor, int input_size, int rank) +{ + at::Device device = "xpu:" + std::to_string(rank); + sycl::queue q = get_current_queue(device); + + int packed_size = (input_size + 7) / 8; + auto unit8_options = at::TensorOptions().dtype(at::kByte).device(at::kXPU); + at::Tensor packed = torch::zeros({packed_size}, unit8_options); + + float* input = (float*)tensor.data_ptr(); + uint8_t* output = (uint8_t*)packed.data_ptr(); + + auto event = q.submit([&](sycl::handler& cgh) { + cgh.parallel_for<>(range(packed_size), [=](id<1> item_ct1) { + packbitskernel(input, output, input_size, item_ct1); + }); + }); + + return packed; +} + +/* +unpack uint8 tensor into float tensor. Every uint8 element get unpacked into eight float +a '1' bit will be converted to a float(1), a '0' bit will be converted to a float(-1). +Arguments: + tensor: A uint8 tensor that get unpacked. + input_size: numel of input tensor + rank: device id in order to get corresponding stream +*/ +at::Tensor unpackbits(at::Tensor tensor, int input_size, int rank) +{ + at::Device device = "xpu:" + std::to_string(rank); + sycl::queue q = get_current_queue(device); + + auto float_options = at::TensorOptions().dtype(at::kFloat).device(at::kXPU); + at::Tensor unpacked = torch::empty({input_size * 8}, float_options); + + uint8_t* input = (uint8_t*)tensor.data_ptr(); + float* output = (float*)unpacked.data_ptr(); + + auto event = q.submit([&](sycl::handler& cgh) { + cgh.parallel_for<>(range(input_size * 8), + [=](id<1> item_ct1) { unpackbitskernel(input, output, item_ct1); }); + }); + + return unpacked; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("packbits", &packbits, "DeepSpeed XPU packbits (C++)"); + m.def("unpackbits", &unpackbits, "DeepSpeed XPU unpackbits (C++)"); +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9cf23373c265a1c0521a9eaa3b713cdd7aca57c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .quantize import FP_Quantize, Quantizer +from .fp8_gemm import matmul_fp8 diff --git a/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8bde6d4ebcb142d8b771bc6ba37b7fb3c61ac7b8 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0c7ee9e5ac9db15c7c73a0ae43518a47e45973d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac9565ff174d786a3420baa63cee585cd93cfc22 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/fp8_gemm_triton.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dfdb9a48fc0705168890c2a2de790be7be5f9b23 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/__pycache__/quantize.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py new file mode 100644 index 0000000000000000000000000000000000000000..db4fa5ae2c9271b480b6d209da28f0f4fdbbe85d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######## Fused MoE kernel ######### +# These kernels are implemented for +# fusing GeMM with dequantization of +# fp8 weight data when using bit-16 +# activation. +################################### + +import torch + + +def matmul_fp8(inp, weight, scale, quantization_group_size, quantizer): + from deepspeed import get_accelerator + + if not get_accelerator().is_triton_supported(): + return matmul_fp8_fallback(inp, weight, scale, quantization_group_size, quantizer) + else: + # Import dynamically to prevent failures on systems without triton. + from .fp8_gemm_triton import matmul_fp8_triton + return matmul_fp8_triton(inp, weight, scale, quantization_group_size) + + +def matmul_fp8_fallback(inp, weight, scale, quantization_group_size, quantizer): + return torch.matmul(inp, quantizer.dequantize(weight, scale=scale)) diff --git a/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py new file mode 100644 index 0000000000000000000000000000000000000000..086525cc64425558399886cc534c57a039ee0a41 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/fp8_gemm_triton.py @@ -0,0 +1,165 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +######## Fused MoE kernel ######### +# These kernels are implemented for +# fusing GeMM with dequantization of +# fp8 weight data when using bit-16 +# activation. +################################### + +import torch +import triton +import triton.language as tl + + +@triton.jit +def matmul_kernel_fp8_bf16(inp_ptr, weight_ptr, out_ptr, scale_ptr, M, N, K, stride_am, stride_ak, stride_bk, + stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, + quantization_group_size: tl.constexpr): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + + inp_data = inp_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + weight_data = weight_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + weight_ptrs_offset = offs_k[:, None] * (stride_bk // quantization_group_size) + ( + (pid_n * BLOCK_SIZE_N) // quantization_group_size) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + inp = tl.load(inp_data, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + weight = tl.load(weight_data, mask=offs_k[:, None] < K, other=0.0) + scale = tl.load(scale_ptr + weight_ptrs_offset + ((k * BLOCK_SIZE_K * stride_bk) // quantization_group_size)) + # Dequantize weight (fp8 -> bf16) + w = (weight & 0x80).to(tl.uint16) << 8 + w = w | ((weight & 0x7f).to(tl.uint16) << 4) + w = (w + 0x3C00).to(tl.uint16) + w = (w.to(tl.bfloat16, bitcast=True).to(tl.float32) * scale).to(tl.bfloat16) + + inp_data += BLOCK_SIZE_K * stride_ak + weight_data += BLOCK_SIZE_K * stride_bk + + accumulator += tl.dot(inp, w) + + out = accumulator.to(tl.bfloat16) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + out_data = out_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + tl.store(out_data, out, mask=(offs_cm[:, None] < M) & (offs_cn[None, :] < N)) + + +@triton.jit +def matmul_kernel_fp8_fp16(inp_ptr, weight_ptr, out_ptr, scale_ptr, M, N, K, stride_am, stride_ak, stride_bk, + stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, + quantization_group_size: tl.constexpr): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + + inp_data = inp_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + weight_data = weight_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + weight_ptrs_offset = offs_k[:, None] * (stride_bk // quantization_group_size) + ( + (pid_n * BLOCK_SIZE_N) // quantization_group_size) + + weight = tl.load(weight_data, mask=offs_k[:, None] < K, other=0.0) + scale = tl.load(scale_ptr + weight_ptrs_offset) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + inp = tl.load(inp_data, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + # Dequantize weight (fp8 -> fp16) + w = (((weight & 0x80) << 8) | ((weight & 0x7f) << 7)).to(tl.uint16) + w = (w + 0x2000).to(tl.uint16) + w = (w.to(tl.float16, bitcast=True) * scale).to(tl.float16) + + inp_data += BLOCK_SIZE_K * stride_ak + weight_data += BLOCK_SIZE_K * stride_bk + + weight = tl.load(weight_data, mask=offs_k[:, None] < K - (k + 1) * BLOCK_SIZE_K, other=0.0) + scale = tl.load(scale_ptr + (weight_ptrs_offset + + (((k + 1) * BLOCK_SIZE_K * stride_bk) // quantization_group_size))) + + accumulator += tl.dot(inp, w) + + out = accumulator.to(tl.float16) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + out_data = out_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + tl.store(out_data, out, mask=(offs_cm[:, None] < M) & (offs_cn[None, :] < N)) + + +def matmul_fp8_triton(inp, weight, scale, quantization_group_size): + + assert inp.shape[1] == weight.shape[0], \ + f"Incompatible dimensions (input: {inp.shape}, weight: {weight.shape})" + + M, K = inp.shape + K, N = weight.shape + + out = torch.empty((M, N), device=inp.device, dtype=inp.dtype) + + # GEMM tuning parameters! + # TODO: Add a more configurable tuning for selecting the best GeMM + BLOCK_SIZE_M = 16 if M <= 16 else 32 if M <= 32 else 64 if M <= 64 else 128 + BLOCK_SIZE_N = 64 + BLOCK_SIZE_K = max(64, quantization_group_size) + GROUP_SIZE_M = 8 + num_stages = 4 + num_warps = 4 + if M >= 256: + BLOCK_SIZE_M = 256 + BLOCK_SIZE_N = 128 + BLOCK_SIZE_K = max(128, quantization_group_size) + num_stages = 3 + num_warps = 8 + + grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), ) + kernel = matmul_kernel_fp8_bf16 if inp.dtype == torch.bfloat16 else matmul_kernel_fp8_fp16 + kernel[grid](inp, + weight, + out, + scale, + M, + N, + K, + inp.stride(0), + inp.stride(1), + weight.stride(0), + weight.stride(1), + out.stride(0), + out.stride(1), + quantization_group_size=quantization_group_size, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + GROUP_SIZE_M=GROUP_SIZE_M, + num_stages=num_stages, + num_warps=num_warps) + return out diff --git a/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/quantize.py b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..47b3b08c7e03d94a98cc6c79af071beb82e40860 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/fp_quantizer/quantize.py @@ -0,0 +1,167 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import abc +from abc import ABC + +import gc +from deepspeed.ops.op_builder import FPQuantizerBuilder +from deepspeed.accelerator import get_accelerator + +fp_quant_module = None + + +class Quantizer(ABC): + """ + Abstract Quantizer class that implements quantize/dequantize methods. + + Arguments: + group_size (int, optional): number of values or elements that are grouped + together for the quantization process. + """ + + def __init__(self, group_size=512) -> None: + self.group_size = group_size + + @abc.abstractmethod + def quantize(self, + input, + q_bits=8, + q_mantisa_bits=3, + stochastic_mode=False, + return_meta_tensor=False) -> torch.Tensor: + ... + + @abc.abstractmethod + def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor: + ... + + +class FP_Quantize(Quantizer): + + def __init__(self, quantization_config) -> None: + global fp_quant_module + super().__init__(group_size=quantization_config.group_size) + if fp_quant_module is None: + fp_quant_module = FPQuantizerBuilder().load() + self.cuda_impl = getattr(fp_quant_module, "CUDA_IMPL", True) + self.q_config = quantization_config + + self.orig_dtype = None + self.num_groups = None + self.input_q = None + self.scale = None + + def quantize(self, + input, + q_bits=8, + q_mantisa_bits=3, + stochastic_mode=False, + return_meta_tensor=False) -> torch.Tensor: + assert input.dtype == torch.bfloat16, "only support bf16 for now" + if return_meta_tensor: + assert q_bits == 8, "meta tensor is only supported with q_bit=8" + + self.orig_dtype = input.dtype + self.orig_shape = input.shape + + if q_bits == 8: + pass + elif q_bits == 12: + q_mantisa_bits = 4 + elif q_bits == 6: + q_mantisa_bits = 2 + elif q_bits == 4: + q_mantisa_bits = 1 + else: + assert (0), \ + f"Missing {q_bits}-quantization, please add the template arguments for the kernel to support this precision!" + self.num_groups = input.numel() // self.group_size + self.input_q = torch.ones(self.num_groups, + int(self.group_size * q_bits) // 8 + 4, + dtype=torch.uint8, + device=input.device) + out = fp_quant_module.quantize(self.input_q, input, self.group_size, stochastic_mode, q_bits, q_mantisa_bits) + if return_meta_tensor: + data, self.scale = out.split(self.group_size, dim=-1) + data = data.contiguous().reshape(input.shape) + self.scale = self.scale.contiguous() + del self.input_q + del out + gc.collect() + get_accelerator().empty_cache() + return data, self.scale + + return out + + def to(self, *args, **kwargs): + # Intermediate tensors may need to be moved to different devices + if hasattr(self, 'input_q'): + self.input_q = self.input_q.to(*args, **kwargs) + if hasattr(self, 'scale'): + self.scale = self.scale.to(*args, **kwargs) + + def get_scales(self): + return fp_quant_module.get_scales(self.scale, self.num_groups) + + def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=None) -> torch.Tensor: + assert (self.orig_dtype is not None), \ + "[De-quantization Error]: you need to call quantize before dequantizing!" + fp_out = torch.empty(self.orig_shape, dtype=self.orig_dtype, + device=input_q.device) if fp_out is None else fp_out + if q_bits == 8: + pass + elif q_bits == 12: + q_mantisa_bits = 4 + elif q_bits == 6: + q_mantisa_bits = 2 + elif q_bits == 4: + q_mantisa_bits = 1 + else: + assert (0), \ + f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" + + if scale is not None: + assert input_q.numel() == fp_out.numel(), \ + f'[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' + input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() + fp_quant_module.dequantize(fp_out, input_q, self.group_size, q_mantisa_bits, q_bits - q_mantisa_bits - 1) + return fp_out + + def selective_dequantize(self, + input_q, + indexes, + fp_out=None, + q_bits=8, + q_mantisa_bits=3, + scale=None) -> torch.Tensor: + assert (not hasattr(self, 'orig_shape') or len(self.orig_shape) == 3), \ + "Selective-Dequantization works on 3d tensor only! Please reshape the tensor before calling dequantize function." + assert (self.orig_dtype is not None), \ + "[De-quantization Error]: you need to call quantize before dequantizing!" + fp_out = torch.empty( + (indexes.shape[0], + *self.orig_shape[1:]), dtype=self.orig_dtype, device=input_q.device) if fp_out is None else fp_out + if q_bits == 8: + pass + elif q_bits == 12: + q_mantisa_bits = 4 + elif q_bits == 6: + q_mantisa_bits = 2 + elif q_bits == 4: + q_mantisa_bits = 1 + else: + assert (0), \ + f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" + + if scale is not None: + assert input_q.numel() == fp_out.numel(), \ + f'[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' + input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() + + fp_quant_module.selective_dequantize(fp_out, input_q, indexes, self.group_size, q_mantisa_bits, + q_bits - q_mantisa_bits - 1) + return fp_out diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..afe48159933c75045e9445a577d545b7ae260073 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__init__.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import os +import pkgutil +import importlib + +from .builder import get_default_compute_capabilities, OpBuilder + +# Do not remove, required for abstract accelerator to detect if we have a deepspeed or 3p op_builder +__deepspeed__ = True + +# List of all available op builders from deepspeed op_builder +try: + import deepspeed.ops.op_builder # noqa: F401 # type: ignore + op_builder_dir = "deepspeed.ops.op_builder" +except ImportError: + op_builder_dir = "op_builder" + +__op_builders__ = [] + +this_module = sys.modules[__name__] + + +def builder_closure(member_name): + if op_builder_dir == "op_builder": + # during installation time cannot get builder due to torch not installed, + # return closure instead + def _builder(): + from deepspeed.accelerator import get_accelerator + builder = get_accelerator().create_op_builder(member_name) + return builder + + return _builder + else: + # during runtime, return op builder class directly + from deepspeed.accelerator import get_accelerator + builder = get_accelerator().get_op_builder(member_name) + return builder + + +# reflect builder names and add builder closure, such as 'TransformerBuilder()' creates op builder wrt current accelerator +for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(this_module.__file__)]): + if module_name != 'all_ops' and module_name != 'builder': + module = importlib.import_module(f".{module_name}", package=op_builder_dir) + for member_name in module.__dir__(): + if member_name.endswith('Builder') and member_name != "OpBuilder" and member_name != "CUDAOpBuilder": + # assign builder name to variable with same name + # the following is equivalent to i.e. TransformerBuilder = "TransformerBuilder" + this_module.__dict__[member_name] = builder_closure(member_name) diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e343dbd5b8ba7a7b6463762bf07638cacb4f54a0 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f02939f11ffb957fd1c2a0ccbbe55592e874034c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/all_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..701d446829b64a0197a47a027d59d2c967fb54c5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/async_io.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7ee5f640dddbd2db5143f245f0e0026d8fcf5fc Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c4cb057f74326b6b4e134a66d05279d5ce9e669 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adagrad.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a5632c829f0c09f8ef86e497263c726abded1b75 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85d4e29ccc3d13a1a363baf03305fd1849cdb987 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/cpu_lion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e36e9839867e4cd8e1d177dee9e47db54345feee Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/dc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ef058fc840883356e90f4d0036b358b7a98e4e6 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/evoformer_attn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8506e90431e3e1ba27a31a3124a56f4863d8dd37 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fp_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ceeb87f9be8e9e2a776deb5342a27d098385e4f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d612920a253bd9742293b535e6ad8907a6b98a34 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lamb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf9ac6aa6dd4af48ad99c1dba33c131094c2653e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/fused_lion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1667a2b3d236f35ce24b8261e9b534d4442fee91 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/gds.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ed53e9bf7e67568366bf602de02fa4c6c1baa75 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/inference_core_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae773b90571e03b3a8e30b18b1ca11eec63aaa56 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/inference_cutlass_builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5dcde763c5cb8d25448d3e8ca5fbb1d6bc5c66a4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b17d2e3b231cd26c07778407554abceefad2dd7d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_ops.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c835fcea1e4a3a5a0c23ae00f971b587782db84e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/ragged_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8435c09acce814975cbd743705dc136b8871dc15 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/random_ltd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1714a6bdcc0a05b250894eb94e7a7ca9de8439d8 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/sparse_attn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c65917c8dc69c80b96ebebd08b77b14346c6b22e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/spatial_inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95dbc3c3ce7e98799ee711de126f4a4b805b8dd9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/stochastic_transformer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d62ae0522bed15a1f255fba9405f9a3dc75a49e5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/transformer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d24209211b56e0f0d7aaf015c5f31a84e45b6c1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/__pycache__/transformer_inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/all_ops.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/all_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..ff11ca180072f7dc72918b417b5e6e44eb53e3c4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/all_ops.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import pkgutil +import importlib +try: + # during installation time accelerator is visible, otherwise return deepspeed.accelerator + from accelerator import get_accelerator +except ImportError: + from deepspeed.accelerator import get_accelerator + +# List of all available ops + +# reflect all builder names into __op_builders__ +op_builder_dir = get_accelerator().op_builder_dir() +op_builder_module = importlib.import_module(op_builder_dir) +__op_builders__ = [] + +for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(op_builder_module.__file__)]): + # avoid self references + if module_name != 'all_ops' and module_name != 'builder': + module = importlib.import_module("{}.{}".format(op_builder_dir, module_name)) + for member_name in module.__dir__(): + if member_name.endswith('Builder'): + # append builder to __op_builders__ list + builder = get_accelerator().create_op_builder(member_name) + __op_builders__.append(builder) + +ALL_OPS = {op.name: op for op in __op_builders__ if op is not None} +accelerator_name = get_accelerator()._name diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/async_io.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..f59cc6810c6f04a9c9b323093b9bfa202294f9ee --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/async_io.py @@ -0,0 +1,113 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import shutil +import subprocess + +from .builder import TorchCPUOpBuilder + + +class AsyncIOBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_AIO" + NAME = "async_io" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.aio.{self.NAME}_op' + + def lib_sources(self): + src_list = [ + 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', 'csrc/aio/py_lib/deepspeed_py_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', + 'csrc/aio/common/deepspeed_aio_utils.cpp', 'csrc/aio/common/deepspeed_aio_common.cpp', + 'csrc/aio/common/deepspeed_aio_types.cpp', 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', + 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', 'csrc/aio/py_lib/deepspeed_py_copy.cpp', + 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' + ] + return src_list + + def sources(self): + return self.lib_sources() + ['csrc/aio/py_lib/py_ds_aio.cpp'] + + def include_paths(self): + import torch + if self.build_for_cpu: + CUDA_INCLUDE = [] + elif not self.is_rocm_pytorch(): + CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] + else: + CUDA_INCLUDE = [ + os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"), + os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "rocrand"), + os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "hiprand"), + ] + return ['csrc/aio/py_lib', 'csrc/aio/common'] + CUDA_INCLUDE + + def cxx_args(self): + # -O0 for improved debugging, since performance is bound by I/O + args = super().cxx_args() + import torch + TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) + if not (TORCH_MAJOR >= 2 and TORCH_MINOR >= 1): + args.remove('-std=c++17') + args.append('-std=c++14') + args += ['-Wall', '-O0', '-shared', '-fPIC', '-Wno-reorder'] + return args + + def extra_ldflags(self): + if self.build_for_cpu: + return ['-fopenmp'] + + import torch.utils.cpp_extension + CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME + if CUDA_HOME is None: + ldflags = ['-laio'] # the ROCM case + else: + CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") + ldflags = [f'-L{CUDA_HOME}', f'-L{CUDA_LIB64}', '-laio', '-lcuda', '-lcudart'] + return ldflags + + def check_for_libaio_pkg(self): + libs = dict( + dpkg=["-l", "libaio-dev", "apt"], + pacman=["-Q", "libaio", "pacman"], + rpm=["-q", "libaio-devel", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + self.warning(f"{self.NAME}: please install the {lib} package with {tool}") + break + return found + + def is_compatible(self, verbose=False): + # Check for the existence of libaio by using distutils + # to compile and link a test program that calls io_submit, + # which is a function provided by libaio that is used in the async_io op. + # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS + # respectively to specify the directories for libaio.h and libaio.so. + aio_compatible = self.has_function('io_submit', ('aio', )) + if verbose and not aio_compatible: + self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") + + # Check for the libaio package via known package managers + # to print suggestions on which package to install. + self.check_for_libaio_pkg() + + self.warning( + "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." + ) + return super().is_compatible(verbose) and aio_compatible diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/builder.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..f31870a1e4cef12507f37af78ce558b1e5e85d84 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/builder.py @@ -0,0 +1,860 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import re +import sys +import time +import importlib +from pathlib import Path +import subprocess +import shlex +import shutil +import tempfile +import distutils.ccompiler +import distutils.log +import distutils.sysconfig +from distutils.errors import CompileError, LinkError +from abc import ABC, abstractmethod +from typing import List + +YELLOW = '\033[93m' +END = '\033[0m' +WARNING = f"{YELLOW} [WARNING] {END}" + +DEFAULT_TORCH_EXTENSION_PATH = "/tmp/torch_extensions" +DEFAULT_COMPUTE_CAPABILITIES = "6.0;6.1;7.0" + +try: + import torch +except ImportError: + print(f"{WARNING} unable to import torch, please install it if you want to pre-compile any deepspeed ops.") +else: + TORCH_MAJOR = int(torch.__version__.split('.')[0]) + TORCH_MINOR = int(torch.__version__.split('.')[1]) + + +class MissingCUDAException(Exception): + pass + + +class CUDAMismatchException(Exception): + pass + + +def installed_cuda_version(name=""): + import torch.utils.cpp_extension + cuda_home = torch.utils.cpp_extension.CUDA_HOME + if cuda_home is None: + raise MissingCUDAException("CUDA_HOME does not exist, unable to compile CUDA op(s)") + # Ensure there is not a cuda version mismatch between torch and nvcc compiler + output = subprocess.check_output([cuda_home + "/bin/nvcc", "-V"], universal_newlines=True) + output_split = output.split() + release_idx = output_split.index("release") + release = output_split[release_idx + 1].replace(',', '').split(".") + # Ignore patch versions, only look at major + minor + cuda_major, cuda_minor = release[:2] + return int(cuda_major), int(cuda_minor) + + +def get_default_compute_capabilities(): + compute_caps = DEFAULT_COMPUTE_CAPABILITIES + # Update compute capability according to: https://en.wikipedia.org/wiki/CUDA#GPUs_supported + import torch.utils.cpp_extension + if torch.utils.cpp_extension.CUDA_HOME is not None: + if installed_cuda_version()[0] == 11: + if installed_cuda_version()[1] >= 0: + compute_caps += ";8.0" + if installed_cuda_version()[1] >= 1: + compute_caps += ";8.6" + if installed_cuda_version()[1] >= 8: + compute_caps += ";9.0" + elif installed_cuda_version()[0] == 12: + compute_caps += ";8.0;8.6;9.0" + if installed_cuda_version()[1] >= 8: + compute_caps += ";10.0;12.0" + return compute_caps + + +# list compatible minor CUDA versions - so that for example pytorch built with cuda-11.0 can be used +# to build deepspeed and system-wide installed cuda 11.2 +cuda_minor_mismatch_ok = { + 10: ["10.0", "10.1", "10.2"], + 11: ["11.0", "11.1", "11.2", "11.3", "11.4", "11.5", "11.6", "11.7", "11.8"], + 12: ["12.0", "12.1", "12.2", "12.3", "12.4", "12.5", "12.6", + "12.8"], # There does not appear to be a CUDA Toolkit 12.7 +} + + +def assert_no_cuda_mismatch(name=""): + cuda_major, cuda_minor = installed_cuda_version(name) + sys_cuda_version = f'{cuda_major}.{cuda_minor}' + torch_cuda_version = ".".join(torch.version.cuda.split('.')[:2]) + # This is a show-stopping error, should probably not proceed past this + if sys_cuda_version != torch_cuda_version: + if (cuda_major in cuda_minor_mismatch_ok and sys_cuda_version in cuda_minor_mismatch_ok[cuda_major] + and torch_cuda_version in cuda_minor_mismatch_ok[cuda_major]): + print(f"Installed CUDA version {sys_cuda_version} does not match the " + f"version torch was compiled with {torch.version.cuda} " + "but since the APIs are compatible, accepting this combination") + return True + elif os.getenv("DS_SKIP_CUDA_CHECK", "0") == "1": + print( + f"{WARNING} DeepSpeed Op Builder: Installed CUDA version {sys_cuda_version} does not match the " + f"version torch was compiled with {torch.version.cuda}." + "Detected `DS_SKIP_CUDA_CHECK=1`: Allowing this combination of CUDA, but it may result in unexpected behavior." + ) + return True + raise CUDAMismatchException( + f">- DeepSpeed Op Builder: Installed CUDA version {sys_cuda_version} does not match the " + f"version torch was compiled with {torch.version.cuda}, unable to compile " + "cuda/cpp extensions without a matching cuda version.") + return True + + +class OpBuilder(ABC): + _rocm_version = None + _rocm_gpu_arch = None + _rocm_wavefront_size = None + _is_rocm_pytorch = None + _is_sycl_enabled = None + _loaded_ops = {} + + def __init__(self, name): + self.name = name + self.jit_mode = False + self.build_for_cpu = False + self.enable_bf16 = False + self.error_log = None + + @abstractmethod + def absolute_name(self): + ''' + Returns absolute build path for cases where the op is pre-installed, e.g., deepspeed.ops.adam.cpu_adam + will be installed as something like: deepspeed/ops/adam/cpu_adam.so + ''' + pass + + @abstractmethod + def sources(self): + ''' + Returns list of source files for your op, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed) + ''' + pass + + def hipify_extension(self): + pass + + def sycl_extension(self): + pass + + @staticmethod + def validate_torch_version(torch_info): + install_torch_version = torch_info['version'] + current_torch_version = ".".join(torch.__version__.split('.')[:2]) + if install_torch_version != current_torch_version: + raise RuntimeError("PyTorch version mismatch! DeepSpeed ops were compiled and installed " + "with a different version than what is being used at runtime. " + f"Please re-install DeepSpeed or switch torch versions. " + f"Install torch version={install_torch_version}, " + f"Runtime torch version={current_torch_version}") + + @staticmethod + def validate_torch_op_version(torch_info): + if not OpBuilder.is_rocm_pytorch(): + current_cuda_version = ".".join(torch.version.cuda.split('.')[:2]) + install_cuda_version = torch_info['cuda_version'] + if install_cuda_version != current_cuda_version: + raise RuntimeError("CUDA version mismatch! DeepSpeed ops were compiled and installed " + "with a different version than what is being used at runtime. " + f"Please re-install DeepSpeed or switch torch versions. " + f"Install CUDA version={install_cuda_version}, " + f"Runtime CUDA version={current_cuda_version}") + else: + current_hip_version = ".".join(torch.version.hip.split('.')[:2]) + install_hip_version = torch_info['hip_version'] + if install_hip_version != current_hip_version: + raise RuntimeError("HIP version mismatch! DeepSpeed ops were compiled and installed " + "with a different version than what is being used at runtime. " + f"Please re-install DeepSpeed or switch torch versions. " + f"Install HIP version={install_hip_version}, " + f"Runtime HIP version={current_hip_version}") + + @staticmethod + def is_rocm_pytorch(): + if OpBuilder._is_rocm_pytorch is not None: + return OpBuilder._is_rocm_pytorch + + _is_rocm_pytorch = False + try: + import torch + except ImportError: + pass + else: + if TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 5): + _is_rocm_pytorch = hasattr(torch.version, 'hip') and torch.version.hip is not None + if _is_rocm_pytorch: + from torch.utils.cpp_extension import ROCM_HOME + _is_rocm_pytorch = ROCM_HOME is not None + OpBuilder._is_rocm_pytorch = _is_rocm_pytorch + return OpBuilder._is_rocm_pytorch + + @staticmethod + def is_sycl_enabled(): + if OpBuilder._is_sycl_enabled is not None: + return OpBuilder._is_sycl_enabled + + _is_sycl_enabled = False + try: + result = subprocess.run(["c2s", "--version"], capture_output=True) + except: + pass + else: + _is_sycl_enabled = True + + OpBuilder._is_sycl_enabled = _is_sycl_enabled + return OpBuilder._is_sycl_enabled + + @staticmethod + def installed_rocm_version(): + if OpBuilder._rocm_version: + return OpBuilder._rocm_version + + ROCM_MAJOR = '0' + ROCM_MINOR = '0' + ROCM_VERSION_DEV_RAW = "" + if OpBuilder.is_rocm_pytorch(): + from torch.utils.cpp_extension import ROCM_HOME + rocm_ver_file = Path(ROCM_HOME).joinpath(".info/version") + if rocm_ver_file.is_file(): + with open(rocm_ver_file, 'r') as file: + ROCM_VERSION_DEV_RAW = file.read() + elif "rocm" in torch.__version__: + ROCM_VERSION_DEV_RAW = torch.__version__.split("rocm")[1] + if ROCM_VERSION_DEV_RAW != "": + ROCM_MAJOR = ROCM_VERSION_DEV_RAW.split('.')[0] + ROCM_MINOR = ROCM_VERSION_DEV_RAW.split('.')[1] + else: + # Look in /usr/include/rocm-version.h + rocm_ver_file = Path("/usr/include/rocm_version.h") + if rocm_ver_file.is_file(): + with open(rocm_ver_file, 'r') as file: + for ln in file.readlines(): + if "#define ROCM_VERSION_MAJOR" in ln: + ROCM_MAJOR = re.findall(r'\S+', ln)[2] + elif "#define ROCM_VERSION_MINOR" in ln: + ROCM_MINOR = re.findall(r'\S+', ln)[2] + if ROCM_MAJOR == '0': + assert False, "Could not detect ROCm version" + + OpBuilder._rocm_version = (int(ROCM_MAJOR), int(ROCM_MINOR)) + return OpBuilder._rocm_version + + @staticmethod + def get_rocm_gpu_arch(): + if OpBuilder._rocm_gpu_arch: + return OpBuilder._rocm_gpu_arch + rocm_info = Path("/opt/rocm/bin/rocminfo") + if (not rocm_info.is_file()): + rocm_info = Path("rocminfo") + rocm_gpu_arch_cmd = str(rocm_info) + " | grep -o -m 1 'gfx.*'" + try: + result = subprocess.check_output(rocm_gpu_arch_cmd, shell=True) + rocm_gpu_arch = result.decode('utf-8').strip() + except subprocess.CalledProcessError: + rocm_gpu_arch = "" + OpBuilder._rocm_gpu_arch = rocm_gpu_arch + return OpBuilder._rocm_gpu_arch + + @staticmethod + def get_rocm_wavefront_size(): + if OpBuilder._rocm_wavefront_size: + return OpBuilder._rocm_wavefront_size + + rocm_info = Path("/opt/rocm/bin/rocminfo") + if (not rocm_info.is_file()): + rocm_info = Path("rocminfo") + rocm_wavefront_size_cmd = str( + rocm_info) + " | grep -Eo -m1 'Wavefront Size:[[:space:]]+[0-9]+' | grep -Eo '[0-9]+'" + try: + result = subprocess.check_output(rocm_wavefront_size_cmd, shell=True) + rocm_wavefront_size = result.decode('utf-8').strip() + except subprocess.CalledProcessError: + rocm_wavefront_size = "32" + OpBuilder._rocm_wavefront_size = rocm_wavefront_size + return OpBuilder._rocm_wavefront_size + + def include_paths(self): + ''' + Returns list of include paths, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed) + ''' + return [] + + def nvcc_args(self): + ''' + Returns optional list of compiler flags to forward to nvcc when building CUDA sources + ''' + return [] + + def cxx_args(self): + ''' + Returns optional list of compiler flags to forward to the build + ''' + return [] + + def is_compatible(self, verbose=False): + ''' + Check if all non-python dependencies are satisfied to build this op + ''' + return True + + def extra_ldflags(self): + return [] + + def has_function(self, funcname, libraries, library_dirs=None, verbose=False): + ''' + Test for existence of a function within a tuple of libraries. + + This is used as a smoke test to check whether a certain library is available. + As a test, this creates a simple C program that calls the specified function, + and then distutils is used to compile that program and link it with the specified libraries. + Returns True if both the compile and link are successful, False otherwise. + ''' + tempdir = None # we create a temporary directory to hold various files + filestderr = None # handle to open file to which we redirect stderr + oldstderr = None # file descriptor for stderr + try: + # Echo compile and link commands that are used. + if verbose: + distutils.log.set_verbosity(1) + + # Create a compiler object. + compiler = distutils.ccompiler.new_compiler(verbose=verbose) + + # Configure compiler and linker to build according to Python install. + distutils.sysconfig.customize_compiler(compiler) + + # Create a temporary directory to hold test files. + tempdir = tempfile.mkdtemp() + + # Define a simple C program that calls the function in question + prog = "void %s(void); int main(int argc, char** argv) { %s(); return 0; }" % (funcname, funcname) + + # Write the test program to a file. + filename = os.path.join(tempdir, 'test.c') + with open(filename, 'w') as f: + f.write(prog) + + # Redirect stderr file descriptor to a file to silence compile/link warnings. + if not verbose: + filestderr = open(os.path.join(tempdir, 'stderr.txt'), 'w') + oldstderr = os.dup(sys.stderr.fileno()) + os.dup2(filestderr.fileno(), sys.stderr.fileno()) + + # Workaround for behavior in distutils.ccompiler.CCompiler.object_filenames() + # Otherwise, a local directory will be used instead of tempdir + drive, driveless_filename = os.path.splitdrive(filename) + root_dir = driveless_filename[0] if os.path.isabs(driveless_filename) else '' + output_dir = os.path.join(drive, root_dir) + + # Attempt to compile the C program into an object file. + cflags = shlex.split(os.environ.get('CFLAGS', "")) + objs = compiler.compile([filename], output_dir=output_dir, extra_preargs=self.strip_empty_entries(cflags)) + + # Attempt to link the object file into an executable. + # Be sure to tack on any libraries that have been specified. + ldflags = shlex.split(os.environ.get('LDFLAGS', "")) + compiler.link_executable(objs, + os.path.join(tempdir, 'a.out'), + extra_preargs=self.strip_empty_entries(ldflags), + libraries=libraries, + library_dirs=library_dirs) + + # Compile and link succeeded + return True + + except CompileError: + return False + + except LinkError: + return False + + except: + return False + + finally: + # Restore stderr file descriptor and close the stderr redirect file. + if oldstderr is not None: + os.dup2(oldstderr, sys.stderr.fileno()) + if filestderr is not None: + filestderr.close() + + # Delete the temporary directory holding the test program and stderr files. + if tempdir is not None: + shutil.rmtree(tempdir) + + def strip_empty_entries(self, args): + ''' + Drop any empty strings from the list of compile and link flags + ''' + return [x for x in args if len(x) > 0] + + def cpu_arch(self): + try: + from cpuinfo import get_cpu_info + except ImportError as e: + cpu_info = self._backup_cpuinfo() + if cpu_info is None: + return "-march=native" + + try: + cpu_info = get_cpu_info() + except Exception as e: + self.warning(f"{self.name} attempted to use py-cpuinfo but failed (exception type: {type(e)}, {e}), " + "falling back to lscpu to get this information.") + cpu_info = self._backup_cpuinfo() + if cpu_info is None: + return "-march=native" + + if cpu_info['arch'].startswith('PPC_'): + # gcc does not provide -march on PowerPC, use -mcpu instead + return '-mcpu=native' + return '-march=native' + + def get_cuda_compile_flag(self): + try: + if not self.is_rocm_pytorch(): + assert_no_cuda_mismatch(self.name) + return "-D__ENABLE_CUDA__" + except MissingCUDAException: + print(f"{WARNING} {self.name} cuda is missing or is incompatible with installed torch, " + "only cpu ops can be compiled!") + return '-D__DISABLE_CUDA__' + return '-D__DISABLE_CUDA__' + + def _backup_cpuinfo(self): + # Construct cpu_info dict from lscpu that is similar to what py-cpuinfo provides + if not self.command_exists('lscpu'): + self.warning(f"{self.name} attempted to query 'lscpu' after failing to use py-cpuinfo " + "to detect the CPU architecture. 'lscpu' does not appear to exist on " + "your system, will fall back to use -march=native and non-vectorized execution.") + return None + result = subprocess.check_output(['lscpu']) + result = result.decode('utf-8').strip().lower() + + cpu_info = {} + cpu_info['arch'] = None + cpu_info['flags'] = "" + if 'genuineintel' in result or 'authenticamd' in result: + cpu_info['arch'] = 'X86_64' + if 'avx512' in result: + cpu_info['flags'] += 'avx512,' + elif 'avx512f' in result: + cpu_info['flags'] += 'avx512f,' + if 'avx2' in result: + cpu_info['flags'] += 'avx2' + elif 'ppc64le' in result: + cpu_info['arch'] = "PPC_" + + return cpu_info + + def simd_width(self): + try: + from cpuinfo import get_cpu_info + except ImportError as e: + cpu_info = self._backup_cpuinfo() + if cpu_info is None: + return '-D__SCALAR__' + + try: + cpu_info = get_cpu_info() + except Exception as e: + self.warning(f"{self.name} attempted to use py-cpuinfo but failed (exception type: {type(e)}, {e}), " + "falling back to lscpu to get this information.") + cpu_info = self._backup_cpuinfo() + if cpu_info is None: + return '-D__SCALAR__' + + if cpu_info['arch'] == 'X86_64': + if 'avx512' in cpu_info['flags'] or 'avx512f' in cpu_info['flags']: + return '-D__AVX512__' + elif 'avx2' in cpu_info['flags']: + return '-D__AVX256__' + return '-D__SCALAR__' + + def command_exists(self, cmd): + if '|' in cmd: + cmds = cmd.split("|") + else: + cmds = [cmd] + valid = False + for cmd in cmds: + safe_cmd = ["bash", "-c", f"type {cmd}"] + result = subprocess.Popen(safe_cmd, stdout=subprocess.PIPE) + valid = valid or result.wait() == 0 + + if not valid and len(cmds) > 1: + print(f"{WARNING} {self.name} requires one of the following commands '{cmds}', but it does not exist!") + elif not valid and len(cmds) == 1: + print(f"{WARNING} {self.name} requires the '{cmd}' command, but it does not exist!") + return valid + + def warning(self, msg): + self.error_log = f"{msg}" + print(f"{WARNING} {msg}") + + def deepspeed_src_path(self, code_path): + if os.path.isabs(code_path): + return code_path + else: + return os.path.join(Path(__file__).parent.parent.absolute(), code_path) + + def builder(self): + from torch.utils.cpp_extension import CppExtension + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + return CppExtension(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + extra_compile_args={'cxx': self.strip_empty_entries(self.cxx_args())}, + extra_link_args=self.strip_empty_entries(self.extra_ldflags())) + + def load(self, verbose=True): + if self.name in __class__._loaded_ops: + return __class__._loaded_ops[self.name] + + from deepspeed.git_version_info import installed_ops, torch_info, accelerator_name + from deepspeed.accelerator import get_accelerator + if installed_ops.get(self.name, False) and accelerator_name == get_accelerator()._name: + # Ensure the op we're about to load was compiled with the same + # torch/cuda versions we are currently using at runtime. + self.validate_torch_version(torch_info) + if torch.cuda.is_available() and isinstance(self, CUDAOpBuilder): + self.validate_torch_op_version(torch_info) + + op_module = importlib.import_module(self.absolute_name()) + __class__._loaded_ops[self.name] = op_module + return op_module + else: + return self.jit_load(verbose) + + def jit_load(self, verbose=True): + if not self.is_compatible(verbose): + raise RuntimeError( + f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}" + ) + try: + import ninja # noqa: F401 # type: ignore + except ImportError: + raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.") + + if isinstance(self, CUDAOpBuilder) and not self.is_rocm_pytorch(): + self.build_for_cpu = not torch.cuda.is_available() + + self.jit_mode = True + from torch.utils.cpp_extension import load + + start_build = time.time() + sources = [os.path.abspath(self.deepspeed_src_path(path)) for path in self.sources()] + extra_include_paths = [os.path.abspath(self.deepspeed_src_path(path)) for path in self.include_paths()] + + # Torch will try and apply whatever CCs are in the arch list at compile time, + # we have already set the intended targets ourselves we know that will be + # needed at runtime. This prevents CC collisions such as multiple __half + # implementations. Stash arch list to reset after build. + torch_arch_list = None + if "TORCH_CUDA_ARCH_LIST" in os.environ: + torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") + os.environ["TORCH_CUDA_ARCH_LIST"] = "" + + nvcc_args = self.strip_empty_entries(self.nvcc_args()) + cxx_args = self.strip_empty_entries(self.cxx_args()) + + if isinstance(self, CUDAOpBuilder): + if not self.build_for_cpu and self.enable_bf16: + cxx_args.append("-DBF16_AVAILABLE") + nvcc_args.append("-DBF16_AVAILABLE") + nvcc_args.append("-U__CUDA_NO_BFLOAT16_OPERATORS__") + nvcc_args.append("-U__CUDA_NO_BFLOAT162_OPERATORS__") + nvcc_args.append("-U__CUDA_NO_BFLOAT16_CONVERSIONS__") + + if self.is_rocm_pytorch(): + cxx_args.append("-D__HIP_PLATFORM_AMD__=1") + os.environ["PYTORCH_ROCM_ARCH"] = self.get_rocm_gpu_arch() + cxx_args.append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) + + op_module = load(name=self.name, + sources=self.strip_empty_entries(sources), + extra_include_paths=self.strip_empty_entries(extra_include_paths), + extra_cflags=cxx_args, + extra_cuda_cflags=nvcc_args, + extra_ldflags=self.strip_empty_entries(self.extra_ldflags()), + with_cuda=True if (isinstance(self, CUDAOpBuilder) and not self.build_for_cpu) else None, + verbose=verbose) + + build_duration = time.time() - start_build + if verbose: + print(f"Time to load {self.name} op: {build_duration} seconds") + + # Reset arch list so we are not silently removing it for other possible use cases + if torch_arch_list: + os.environ["TORCH_CUDA_ARCH_LIST"] = torch_arch_list + + __class__._loaded_ops[self.name] = op_module + + return op_module + + +class CUDAOpBuilder(OpBuilder): + + def compute_capability_args(self, cross_compile_archs=None): + """ + Returns nvcc compute capability compile flags. + + 1. `TORCH_CUDA_ARCH_LIST` takes priority over `cross_compile_archs`. + 2. If neither is set default compute capabilities will be used + 3. Under `jit_mode` compute capabilities of all visible cards will be used plus PTX + + Format: + + - `TORCH_CUDA_ARCH_LIST` may use ; or whitespace separators. Examples: + + TORCH_CUDA_ARCH_LIST="6.1;7.5;8.6;9.0;10.0" pip install ... + TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 9.0 10.0+PTX" pip install ... + + - `cross_compile_archs` uses ; separator. + + """ + ccs = [] + if self.jit_mode: + # Compile for underlying architectures since we know those at runtime + for i in range(torch.cuda.device_count()): + CC_MAJOR, CC_MINOR = torch.cuda.get_device_capability(i) + cc = f"{CC_MAJOR}.{CC_MINOR}" + if cc not in ccs: + ccs.append(cc) + ccs = sorted(ccs) + ccs[-1] += '+PTX' + else: + # Cross-compile mode, compile for various architectures + # env override takes priority + cross_compile_archs_env = os.environ.get('TORCH_CUDA_ARCH_LIST', None) + if cross_compile_archs_env is not None: + if cross_compile_archs is not None: + print( + f"{WARNING} env var TORCH_CUDA_ARCH_LIST={cross_compile_archs_env} overrides cross_compile_archs={cross_compile_archs}" + ) + cross_compile_archs = cross_compile_archs_env.replace(' ', ';') + else: + if cross_compile_archs is None: + cross_compile_archs = get_default_compute_capabilities() + ccs = cross_compile_archs.split(';') + + ccs = self.filter_ccs(ccs) + if len(ccs) == 0: + raise RuntimeError( + f"Unable to load {self.name} op due to no compute capabilities remaining after filtering") + + args = [] + self.enable_bf16 = True + for cc in ccs: + num = cc[0] + cc[1].split('+')[0] + args.append(f'-gencode=arch=compute_{num},code=sm_{num}') + if cc[1].endswith('+PTX'): + args.append(f'-gencode=arch=compute_{num},code=compute_{num}') + + if int(cc[0]) <= 7: + self.enable_bf16 = False + + return args + + def filter_ccs(self, ccs: List[str]): + """ + Prune any compute capabilities that are not compatible with the builder. Should log + which CCs have been pruned. + """ + return [cc.split('.') for cc in ccs] + + def version_dependent_macros(self): + # Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456 + version_ge_1_1 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0): + version_ge_1_1 = ['-DVERSION_GE_1_1'] + version_ge_1_3 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2): + version_ge_1_3 = ['-DVERSION_GE_1_3'] + version_ge_1_5 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4): + version_ge_1_5 = ['-DVERSION_GE_1_5'] + return version_ge_1_1 + version_ge_1_3 + version_ge_1_5 + + def is_compatible(self, verbose=False): + return super().is_compatible(verbose) + + def builder(self): + try: + if not self.is_rocm_pytorch(): + assert_no_cuda_mismatch(self.name) + self.build_for_cpu = False + except MissingCUDAException: + self.build_for_cpu = True + + if self.build_for_cpu: + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + else: + from torch.utils.cpp_extension import CUDAExtension as ExtensionBuilder + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} if self.build_for_cpu else \ + {'cxx': self.strip_empty_entries(self.cxx_args()), \ + 'nvcc': self.strip_empty_entries(self.nvcc_args())} + + if not self.build_for_cpu and self.enable_bf16: + compile_args['cxx'].append("-DBF16_AVAILABLE") + compile_args['nvcc'].append("-DBF16_AVAILABLE") + + if self.is_rocm_pytorch(): + compile_args['cxx'].append("-D__HIP_PLATFORM_AMD__=1") + #cxx compiler args are required to compile cpp files + compile_args['cxx'].append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) + #nvcc compiler args are required to compile hip files + compile_args['nvcc'].append('-DROCM_WAVEFRONT_SIZE=%s' % self.get_rocm_wavefront_size()) + if self.get_rocm_gpu_arch(): + os.environ["PYTORCH_ROCM_ARCH"] = self.get_rocm_gpu_arch() + + cuda_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args, + extra_link_args=self.strip_empty_entries(self.extra_ldflags())) + + if self.is_rocm_pytorch(): + # hip converts paths to absolute, this converts back to relative + sources = cuda_ext.sources + curr_file = Path(__file__).parent.parent # ds root + for i in range(len(sources)): + src = Path(sources[i]) + if src.is_absolute(): + sources[i] = str(src.relative_to(curr_file)) + else: + sources[i] = str(src) + cuda_ext.sources = sources + return cuda_ext + + def hipify_extension(self): + if self.is_rocm_pytorch(): + from torch.utils.hipify import hipify_python + hipify_python.hipify( + project_directory=os.getcwd(), + output_directory=os.getcwd(), + header_include_dirs=self.include_paths(), + includes=[os.path.join(os.getcwd(), '*')], + extra_files=[os.path.abspath(s) for s in self.sources()], + show_detailed=True, + is_pytorch_extension=True, + hipify_extra_files_only=True, + ) + + def cxx_args(self): + if sys.platform == "win32": + return ['-O2'] + else: + return ['-O3', '-std=c++17', '-g', '-Wno-reorder'] + + def nvcc_args(self): + if self.build_for_cpu: + return [] + args = ['-O3'] + if self.is_rocm_pytorch(): + ROCM_MAJOR, ROCM_MINOR = self.installed_rocm_version() + args += [ + '-std=c++17', '-U__HIP_NO_HALF_OPERATORS__', '-U__HIP_NO_HALF_CONVERSIONS__', + '-U__HIP_NO_HALF2_OPERATORS__', + '-DROCM_VERSION_MAJOR=%s' % ROCM_MAJOR, + '-DROCM_VERSION_MINOR=%s' % ROCM_MINOR + ] + else: + try: + nvcc_threads = int(os.getenv("DS_NVCC_THREADS", "")) + if nvcc_threads <= 0: + raise ValueError("") + except ValueError: + nvcc_threads = min(os.cpu_count(), 8) + + cuda_major, cuda_minor = installed_cuda_version() + if cuda_major > 10: + if cuda_major == 12 and cuda_minor >= 5: + std_lib = '-std=c++20' + else: + std_lib = '-std=c++17' + else: + std_lib = '-std=c++14' + args += [ + '-allow-unsupported-compiler' if sys.platform == "win32" else '', '--use_fast_math', std_lib, + '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__', + f'--threads={nvcc_threads}' + ] + if os.environ.get('DS_DEBUG_CUDA_BUILD', '0') == '1': + args.append('--ptxas-options=-v') + args += self.compute_capability_args() + return args + + def libraries_args(self): + if self.build_for_cpu: + return [] + + if sys.platform == "win32": + return ['cublas', 'curand'] + else: + return [] + + +class TorchCPUOpBuilder(CUDAOpBuilder): + + def get_cuda_lib64_path(self): + import torch + if not self.is_rocm_pytorch(): + CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib64") + if not os.path.exists(CUDA_LIB64): + CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib") + else: + CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.ROCM_HOME, "lib") + return CUDA_LIB64 + + def extra_ldflags(self): + if self.build_for_cpu: + return ['-fopenmp'] + + if not self.is_rocm_pytorch(): + ld_flags = ['-lcurand'] + if not self.build_for_cpu: + ld_flags.append(f'-L{self.get_cuda_lib64_path()}') + return ld_flags + + return [] + + def cxx_args(self): + args = [] + if not self.build_for_cpu: + CUDA_LIB64 = self.get_cuda_lib64_path() + + args += super().cxx_args() + args += [ + f'-L{CUDA_LIB64}', + '-lcudart', + '-lcublas', + '-g', + ] + + CPU_ARCH = self.cpu_arch() + SIMD_WIDTH = self.simd_width() + CUDA_ENABLE = self.get_cuda_compile_flag() + args += [ + CPU_ARCH, + '-fopenmp', + SIMD_WIDTH, + CUDA_ENABLE, + ] + + return args diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7084db8469f1dae201796c44f911f539459e4757 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .comm import CCLCommBuilder, ShareMemCommBuilder +from .fused_adam import FusedAdamBuilder +from .cpu_adam import CPUAdamBuilder +from .no_impl import NotImplementedBuilder +from .async_io import AsyncIOBuilder diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..abac2c8aae6d717f8380b0ead5bdb3e4148d6ab8 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dbc6e745007686fefc63834862b3cf182e6b451b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/async_io.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5d53e1881e4d23c49a249e96f8691f0724f1388 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a622b4a910314bb66813d368cbba92e481e36788 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/comm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4befec2f6bb2811a304cfa0dff8d47ec1dd4817d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/cpu_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad2f98234a66404ec3ce28c316624cf54221a565 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/fused_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93a374c7efa07419167a34860de35a0810372b6a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/__pycache__/no_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/async_io.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..dcb9feabcfc3c7c3081d812fafb7d8816a031eb3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/async_io.py @@ -0,0 +1,91 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import shutil +import subprocess + +from .builder import CPUOpBuilder + + +class AsyncIOBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_AIO" + NAME = "async_io" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.aio.{self.NAME}_op' + + def lib_sources(self): + src_list = [ + 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', 'csrc/aio/py_lib/deepspeed_py_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', + 'csrc/aio/common/deepspeed_aio_utils.cpp', 'csrc/aio/common/deepspeed_aio_common.cpp', + 'csrc/aio/common/deepspeed_aio_types.cpp', 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', + 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', 'csrc/aio/py_lib/deepspeed_py_copy.cpp', + 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' + ] + return src_list + + def sources(self): + return self.lib_sources() + ['csrc/aio/py_lib/py_ds_aio.cpp'] + + def include_paths(self): + return ['csrc/aio/py_lib', 'csrc/aio/common'] + + def cxx_args(self): + # -O0 for improved debugging, since performance is bound by I/O + args = super().cxx_args() + import torch + TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) + if not (TORCH_MAJOR >= 2 and TORCH_MINOR >= 1): + args.remove('-std=c++17') + args.append('-std=c++14') + args += ['-Wall', '-O0', '-shared', '-fPIC', '-Wno-reorder'] + return args + + def extra_ldflags(self): + return ['-laio', '-fopenmp'] + + def check_for_libaio_pkg(self): + libs = dict( + dpkg=["-l", "libaio-dev", "apt"], + pacman=["-Q", "libaio", "pacman"], + rpm=["-q", "libaio-devel", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + self.warning(f"{self.NAME}: please install the {lib} package with {tool}") + break + return found + + def is_compatible(self, verbose=False): + # Check for the existence of libaio by using distutils + # to compile and link a test program that calls io_submit, + # which is a function provided by libaio that is used in the async_io op. + # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS + # respectively to specify the directories for libaio.h and libaio.so. + aio_compatible = self.has_function('io_submit', ('aio', )) + if verbose and not aio_compatible: + self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") + + # Check for the libaio package via known package managers + # to print suggestions on which package to install. + self.check_for_libaio_pkg() + + self.warning( + "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." + ) + return super().is_compatible(verbose) and aio_compatible diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/builder.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..d881842ad0b18654bc26cb3b54d345a4cd160da4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/builder.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class CPUOpBuilder(OpBuilder): + + def builder(self): + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + + cpp_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args) + + return cpp_ext + + def cxx_args(self): + args = ['-O3', '-g', '-Wno-reorder'] + CPU_ARCH = self.cpu_arch() + SIMD_WIDTH = self.simd_width() + args += [CPU_ARCH, '-fopenmp', SIMD_WIDTH] + return args + + def libraries_args(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/comm.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/comm.py new file mode 100644 index 0000000000000000000000000000000000000000..fec960b63b2ec16cb7ab6a92b509a5d97b901cef --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/comm.py @@ -0,0 +1,71 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from .builder import CPUOpBuilder + + +class CCLCommBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_CCL_COMM" + NAME = "deepspeed_ccl_comm" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.comm.{self.NAME}_op' + + def sources(self): + return ['csrc/cpu/comm/ccl.cpp', 'csrc/cpu/comm/shm.cpp'] + + def include_paths(self): + includes = ['csrc/cpu/includes'] + return includes + + def cxx_args(self): + return ['-O2', '-fopenmp'] + + def is_compatible(self, verbose=False): + # TODO: add soft compatibility check for private binary release. + # a soft check, as in we know it can be trivially changed. + return super().is_compatible(verbose) + + def extra_ldflags(self): + ccl_root_path = os.environ.get("CCL_ROOT") + if ccl_root_path is None: + raise ValueError( + "Didn't find CCL_ROOT, install oneCCL from https://github.com/oneapi-src/oneCCL and source its environment variable" + ) + return [] + else: + return ['-lccl', f'-L{ccl_root_path}/lib'] + + +class ShareMemCommBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_SHM_COMM" + NAME = "deepspeed_shm_comm" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.comm.{self.NAME}_op' + + def sources(self): + return ['csrc/cpu/comm/shm_interface.cpp', 'csrc/cpu/comm/shm.cpp'] + + def include_paths(self): + includes = ['csrc/cpu/includes'] + return includes + + def cxx_args(self): + return ['-O2', '-fopenmp'] + + def is_compatible(self, verbose=False): + # TODO: add soft compatibility check for private binary release. + # a soft check, as in we know it can be trivially changed. + return super().is_compatible(verbose) diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..0c8438aea40d4ffd588a496844d692c9a17faaa6 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/cpu_adam.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class CPUAdamBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..34b43825b09024136afdc44e916349ea3f5ce5ad --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/fused_adam.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class FusedAdamBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/cpu/adam/fused_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..69d114a9f1c0b0defc482c1fb143c261fc466125 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu/no_impl.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class NotImplementedBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.comm.{self.NAME}_op' + + def load(self, verbose=True): + raise ValueError("This op had not been implemented on CPU backend.") + + def sources(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..c05f7148895000c10cdfebef7461858a1f987fbb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_adagrad.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import TorchCPUOpBuilder + + +class CPUAdagradBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/adagrad/cpu_adagrad.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..7f4c0847a8c4f32b5f3cacd40f88eed8eea904f9 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_adam.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import TorchCPUOpBuilder + + +class CPUAdamBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_lion.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_lion.py new file mode 100644 index 0000000000000000000000000000000000000000..9a60d99773b31252c9d002a0cf3282185d91a07c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/cpu_lion.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import TorchCPUOpBuilder + + +class CPULionBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_LION" + NAME = "cpu_lion" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lion.{self.NAME}_op' + + def sources(self): + return ['csrc/lion/cpu_lion.cpp', 'csrc/lion/cpu_lion_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/dc.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/dc.py new file mode 100644 index 0000000000000000000000000000000000000000..d05210b8a2b49a28377a83b7be6547e6abe27f9b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/dc.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import TorchCPUOpBuilder + + +class DeepCompileBuilder(TorchCPUOpBuilder): + BUILD_VAR = "DS_BUILD_DEEP_COMPILE" + NAME = "dc" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def sources(self): + return [ + 'csrc/compile/deepcompile.cpp', 'csrc/compile/init.cpp', 'csrc/compile/z1.cpp', 'csrc/compile/z3.cpp', + 'csrc/compile/util.cpp' + ] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + import os + import torch + if self.build_for_cpu: + CUDA_INCLUDE = [] + elif not self.is_rocm_pytorch(): + CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] + else: + CUDA_INCLUDE = [ + os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"), + ] + return ['csrc/includes'] + CUDA_INCLUDE diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/evoformer_attn.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/evoformer_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..7f68ccf8729085028d445fb3d7f76d6f51556bfe --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/evoformer_attn.py @@ -0,0 +1,93 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder, installed_cuda_version +import os + + +class EvoformerAttnBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_EVOFORMER_ATTN" + NAME = "evoformer_attn" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + self.cutlass_path = os.environ.get('CUTLASS_PATH') + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def sources(self): + src_dir = 'csrc/deepspeed4science/evoformer_attn' + return [f'{src_dir}/attention.cpp', f'{src_dir}/attention_back.cu', f'{src_dir}/attention_cu.cu'] + + def nvcc_args(self): + args = super().nvcc_args() + try: + import torch + except ImportError: + self.warning("Please install torch if trying to pre-compile kernels") + return args + major = torch.cuda.get_device_properties(0).major #ignore-cuda + minor = torch.cuda.get_device_properties(0).minor #ignore-cuda + args.append(f"-DGPU_ARCH={major}{minor}") + return args + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile kernels") + return False + if self.cutlass_path is None: + if verbose: + self.warning("Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH") + return False + if os.path.exists(f'{self.cutlass_path}/CHANGELOG.md'): + with open(f'{self.cutlass_path}/CHANGELOG.md', 'r') as f: + if '3.1.0' not in f.read(): + if verbose: + self.warning("Please use CUTLASS version >= 3.1.0") + return False + else: + # pip install nvidia-cutlass package + try: + import cutlass + except ImportError: + if verbose: + self.warning("Please pip install nvidia-cutlass if trying to pre-compile kernels") + return False + cutlass_major, cutlass_minor = cutlass.__version__.split('.')[:2] + cutlass_compatible = (int(cutlass_major) >= 3 and int(cutlass_minor) >= 1) + if not cutlass_compatible: + if verbose: + self.warning("Please use CUTLASS version >= 3.1.0") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 7: + if verbose: + self.warning("Please use a GPU with compute capability >= 7.0") + cuda_okay = False + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("Please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def include_paths(self): + includes = [f'{self.cutlass_path}/include', f'{self.cutlass_path}/tools/util/include'] + return includes diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/fp_quantizer.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/fp_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..2b962ac2c1fea18d1024359623f3faf82228f253 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/fp_quantizer.py @@ -0,0 +1,120 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +try: + from packaging import version as pkg_version +except ImportError: + pkg_version = None + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class FPQuantizerBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_FP_QUANTIZER" + NAME = "fp_quantizer" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.fp_quantizer.{self.NAME}_op' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 8: + if verbose: + self.warning("NVIDIA Inference is only supported on Ampere and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + + try: + import triton + except ImportError: + if verbose: + self.warning( + f"please install triton==2.3.0, 2.3.1 or 3.0.0 if you want to use the FP Quantizer Kernels") + return False + + # triton 2.3.{0,1} and 3.0.0 are ok. + allowed_versions = ("2.3", "3.0", "3.1", "3.2") + if pkg_version: + allowed = (pkg_version.parse(v) for v in allowed_versions) + installed_triton = pkg_version.parse(triton.__version__) + triton_mismatch = all(installed_triton.major != a.major or installed_triton.minor != a.minor + for a in allowed) + else: + installed_triton = triton.__version__ + major, minor, _ = installed_triton.split(".") + allowed = (v.split(".") for v in allowed_versions) + triton_mismatch = all(major != v[0] or minor != v[1] for v in allowed) + + if triton_mismatch: + if verbose: + self.warning( + f"FP Quantizer is using an untested triton version ({installed_triton}), only 2.3.{0,1} and 3.0.0 are known to be compatible with these kernels" + ) + return False + + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 8: + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def sources(self): + return [ + "csrc/fp_quantizer/fp_quantize.cu", + "csrc/fp_quantizer/fp_quantize.cpp", + ] + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def include_paths(self): + return ['csrc/fp_quantizer/includes', 'csrc/includes'] + + @staticmethod + def get_default_quant_dtype(): + import torch + return torch.uint8 + + @staticmethod + def get_quant_range(q_bits=None): + if q_bits == 8: + return 480 + elif q_bits == 6: + return 28. + elif q_bits == 12: + return 510. + else: + assert (0), \ + "Please specify the right quantization range for the selected precision!" diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..ac6e4eeaaea5d9b2c2ee70de3d4261c6348abe94 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_adam.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + +import sys + + +class FusedAdamBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/fused_adam_frontend.cpp', 'csrc/adam/multi_tensor_adam.cu'] + + def include_paths(self): + return ['csrc/includes', 'csrc/adam'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() + + def nvcc_args(self): + nvcc_flags = ['-O3'] + self.version_dependent_macros() + if not self.is_rocm_pytorch(): + nvcc_flags.extend( + ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + + self.compute_capability_args()) + return nvcc_flags diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_lamb.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_lamb.py new file mode 100644 index 0000000000000000000000000000000000000000..f0cb557706b3f7e6d4ab73d7a999a70aebecbeb3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_lamb.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + +import sys + + +class FusedLambBuilder(CUDAOpBuilder): + BUILD_VAR = 'DS_BUILD_FUSED_LAMB' + NAME = "fused_lamb" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lamb.{self.NAME}_op' + + def sources(self): + return ['csrc/lamb/fused_lamb_cuda.cpp', 'csrc/lamb/fused_lamb_cuda_kernel.cu'] + + def include_paths(self): + return ['csrc/includes'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() + + def nvcc_args(self): + nvcc_flags = ['-O3'] + self.version_dependent_macros() + if self.is_rocm_pytorch(): + ROCM_MAJOR, ROCM_MINOR = self.installed_rocm_version() + nvcc_flags += ['-DROCM_VERSION_MAJOR=%s' % ROCM_MAJOR, '-DROCM_VERSION_MINOR=%s' % ROCM_MINOR] + else: + nvcc_flags.extend( + ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + + self.compute_capability_args()) + return nvcc_flags diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_lion.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_lion.py new file mode 100644 index 0000000000000000000000000000000000000000..b900a8f2369df14d65a73235d360db21583ee4ab --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/fused_lion.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + +import sys + + +class FusedLionBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_LION" + NAME = "fused_lion" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lion.{self.NAME}_op' + + def sources(self): + return ['csrc/lion/fused_lion_frontend.cpp', 'csrc/lion/multi_tensor_lion.cu'] + + def include_paths(self): + return ['csrc/includes', 'csrc/lion'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() + + def nvcc_args(self): + nvcc_flags = ['-O3'] + self.version_dependent_macros() + if not self.is_rocm_pytorch(): + nvcc_flags.extend( + ['-allow-unsupported-compiler' if sys.platform == "win32" else '', '-lineinfo', '--use_fast_math'] + + self.compute_capability_args()) + return nvcc_flags diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/gds.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/gds.py new file mode 100644 index 0000000000000000000000000000000000000000..d815ccea387a78a0a6256b6f981b81346a3c230b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/gds.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from .async_io import AsyncIOBuilder + + +class GDSBuilder(AsyncIOBuilder): + BUILD_VAR = "DS_BUILD_GDS" + NAME = "gds" + + def __init__(self): + super().__init__() + + def absolute_name(self): + return f'deepspeed.ops.gds.{self.NAME}_op' + + def lib_sources(self): + src_list = ['csrc/gds/py_lib/deepspeed_py_gds_handle.cpp', 'csrc/gds/py_lib/deepspeed_gds_op.cpp'] + return super().lib_sources() + src_list + + def sources(self): + return self.lib_sources() + ['csrc/gds/py_lib/py_ds_gds.cpp'] + + def cxx_args(self): + return super().cxx_args() + ['-lcufile'] + + def include_paths(self): + import torch + CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")] + return ['csrc/aio/py_lib', 'csrc/aio/common'] + CUDA_INCLUDE + + def extra_ldflags(self): + return super().extra_ldflags() + ['-lcufile'] + + def is_compatible(self, verbose=False): + if self.is_rocm_pytorch(): + if verbose: + self.warning(f'{self.NAME} is not compatible with ROCM') + return False + + try: + import torch.utils.cpp_extension + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile GDS") + return False + + CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME + if CUDA_HOME is None: + if verbose: + self.warning("Please install torch CUDA if trying to pre-compile GDS with CUDA") + return False + + CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") + gds_compatible = self.has_function(funcname="cuFileDriverOpen", + libraries=("cufile", ), + library_dirs=( + CUDA_HOME, + CUDA_LIB64, + ), + verbose=verbose) + + return gds_compatible and super().is_compatible(verbose) diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5ad1b9a7f891eae0ceec8c26109ba2a235ef6349 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .cpu_adam import CPUAdamBuilder +from .fused_adam import FusedAdamBuilder +from .transformer_inference import InferenceBuilder +from .no_impl import NotImplementedBuilder diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa2e5970cd885b8eac57a12dfc589eb74ff7ecbc Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..188c695387dc5512ada80153f5c17fd027e1f1e8 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bb78ffe0002392a393d3a5f6300f836f9906baa Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/cpu_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f753978b4ec1e43e849869e5a9d12bdcc07a0a95 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fp_quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4c18b8be409e00e9ee3d2a93b57b102eb4d11a3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/fused_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f2b5c5029269976c76d2800db74765cf1875e36 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/no_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..677aef5ed5928ff2dd47abc51f90fae2504c464d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/__pycache__/transformer_inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/builder.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..11e710a8ee4833bc66e82c4ed3241f7fe8a2e977 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/builder.py @@ -0,0 +1,38 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class CPUOpBuilder(OpBuilder): + + def builder(self): + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + + cpp_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args) + + return cpp_ext + + def cxx_args(self): + args = ['-O3', '-g', '-Wno-reorder'] + return args + + def libraries_args(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..58eea2698ebb9eb6a8269a527a902792f1c7a6bc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/cpu_adam.py @@ -0,0 +1,28 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class CPUAdamBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c74affb55045a167f8dd3bb51b54f1487a815587 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/fp_quantizer.py @@ -0,0 +1,86 @@ +# Copyright (c) 2024 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class FPQuantizerBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_FP_QUANTIZER" + NAME = "fp_quantizer" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.fp_quantizer.{self.NAME}_op' + + def sources(self): + return [] + + def load(self, verbose=True): + return FPQuantizer + + @staticmethod + def get_default_quant_dtype(): + return torch.float8_e4m3fn + + @staticmethod + def get_quant_range(q_bits=None): + import habana_frameworks.torch.utils.experimental as htexp + if htexp._get_device_type() == htexp.synDeviceType.synDeviceGaudi2: + dtype = torch.float8_e4m3fnuz + else: + dtype = torch.float8_e4m3fn + return torch.finfo(dtype).max + + +class FPQuantizer: + CUDA_IMPL = False + + @classmethod + def selective_dequantize(cls, val_q, scales, indexes, group_size, q_mantisa_bits, q_exponent_bits): + assert False, "Selective dequantize isn't implemented for HPU!" + + @classmethod + def dequantize(cls, fp_out, input_q, scale, group_size, q_mantisa_bits, q_exponent_bits): + orig_shape = fp_out.shape + orig_dtype = fp_out.dtype + dequant_out = torch.ops.hpu.cast_from_fp8(input_q, (1.0 / scale), orig_dtype).view(orig_shape) + fp_out.copy_(dequant_out) + return fp_out + + @classmethod + def quantize(cls, out, val, scale, group_size, stochastic_rounding, q_bits, q_mantisa_bits): + assert q_bits == 8, "Quantize on HPU only supports quantization to FP8" + assert q_mantisa_bits == 3, "Quantize on HPU only supports q_mantissa_bits = 3" + assert out.dtype.is_floating_point, "Quantization on HPU is only to float dtypes" + + num_groups, group_size = out.shape + + # Reshape the tensor + val_reshaped = val.view(num_groups, group_size).float() + # Calculate the scale + max_vals = val_reshaped.abs().max(dim=1, keepdim=True)[0] + q_range = torch.finfo(out.dtype).max + tmp_scale = q_range / max_vals + scale.copy_(tmp_scale) + # Copy quantized + quant, _ = torch.ops.hpu.cast_to_fp8_v2(val_reshaped, scale, stochastic_rounding, dtype=out.dtype) + out.copy_(quant) + + return out + + @classmethod + def get_scales(cls, out, num_groups): + return out diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..5acb121668e37d1f7b8b660b831d35ae863d2a1a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/fused_adam.py @@ -0,0 +1,105 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + +try: + import torch + import math +except ImportError as e: + pass + + +class HPUFusedAdam: + htcore = None + is_lazy_mode = None + + @staticmethod + def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, + bias_correction, weight_decay, *args): + + if HPUFusedAdam.htcore is None: + from habana_frameworks.torch import core as htcore + from habana_frameworks.torch.utils.internal import is_lazy + HPUFusedAdam.htcore = htcore + HPUFusedAdam.is_lazy_mode = is_lazy() + + htcore = HPUFusedAdam.htcore + + htcore.step_closure._mark_step_if_lazy() + step_size = lr + if bias_correction: + bias_correction1 = 1.0 - pow(beta1, step) + bias_correction2 = 1.0 - pow(beta2, step) + step_size = step_size * math.sqrt(bias_correction2) / bias_correction1 + + neg_step = -step_size + neg_step_t = (torch.tensor([neg_step], dtype=torch.float, + requires_grad=False).to(tensor_lists[1][0].dtype).to(tensor_lists[1][0].device, + non_blocking=True)) + + weight_decay = weight_decay if adam_w_mode else 0 + + # since lr is fed into the kernel as tensor, perform the scalar multiplication of wd here + # NOTE: TODO if lr is updated every step, then we need to convert it as tensor and + # perform weight decay unconditonally. + modified_wd = 1.0 - weight_decay * lr + + if HPUFusedAdam.is_lazy_mode: + torch.ops.hpu.optimizer_adamw( + tensor_lists[0], + tensor_lists[1], + tensor_lists[2], + tensor_lists[3], + neg_step_t, + beta1, + beta2, + epsilon, + modified_wd, + ) + else: + modified_wd_t = (torch.tensor([modified_wd], dtype=torch.float, requires_grad=False).to( + tensor_lists[1][0].dtype).to(tensor_lists[1][0].device, non_blocking=True)) + torch.ops.hpu.optimizer_adamw( + tensor_lists[0], + tensor_lists[1], + tensor_lists[2], + tensor_lists[3], + neg_step_t, + beta1, + beta2, + epsilon, + modified_wd_t, + modified_wd != 1.0, + ) + + htcore.step_closure._mark_step_if_lazy() + + +class FusedAdamBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return HPUFusedAdam diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..140d65b48defbea439483b63fd545e52de92209e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/no_impl.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CPUOpBuilder + + +class NotImplementedBuilder(CPUOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.comm.{self.NAME}_op' + + def load(self, verbose=True): + raise ValueError("This op had not been implemented on HPU backend.") + + def sources(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..e397c99200ec9155c933bc0691050b1f1e85cf7b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/hpu/transformer_inference.py @@ -0,0 +1,39 @@ +# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 +import importlib + +# DeepSpeed Team + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class InferenceBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" + NAME = "transformer_inference" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=self.NAME) + + def absolute_name(self): + return f"deepspeed.ops.transformer.inference.{self.NAME}_op" + + def sources(self): + return [] + + def load(self, verbose=True): + if self.name in __class__._loaded_ops: + return __class__._loaded_ops[self.name] + + from deepspeed.git_version_info import installed_ops # noqa: F401 + if installed_ops.get(self.name, False): + op_module = importlib.import_module(self.absolute_name()) + __class__._loaded_ops[self.name] = op_module + return op_module diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/inference_core_ops.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/inference_core_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..b6665ebb76180c533f9700db222fca45314ecd19 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/inference_core_ops.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class InferenceCoreBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_INFERENCE_CORE_OPS" + NAME = "inference_core_ops" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.inference.v2.kernels{self.NAME}' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 6: + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def get_prefix(self): + ds_path = self.deepspeed_src_path("deepspeed") + return "deepspeed" if os.path.isdir(ds_path) else ".." + + def sources(self): + sources = [ + "inference/v2/kernels/core_ops/core_ops.cpp", + "inference/v2/kernels/core_ops/bias_activations/bias_activation.cpp", + "inference/v2/kernels/core_ops/bias_activations/bias_activation_cuda.cu", + "inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.cpp", + "inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cu", + "inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.cpp", + "inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_cuda.cu", + "inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.cpp", + "inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels_cuda.cu", + "inference/v2/kernels/core_ops/cuda_linear/linear_kernels.cpp", + "inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.cu", + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources + + def extra_ldflags(self): + return [] + + def include_paths(self): + sources = [ + 'inference/v2/kernels/core_ops/bias_activations', + 'inference/v2/kernels/core_ops/blas_kernels', + 'inference/v2/kernels/core_ops/cuda_layer_norm', + 'inference/v2/kernels/core_ops/cuda_rms_norm', + 'inference/v2/kernels/core_ops/gated_activations', + 'inference/v2/kernels/core_ops/cuda_linear', + 'inference/v2/kernels/includes', + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + + return sources diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..a4a607288ca857f6812fdd925413493822bfdcb9 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/inference_cutlass_builder.py @@ -0,0 +1,95 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +import os + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class InferenceCutlassBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_CUTLASS_OPS" + NAME = "cutlass_ops" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.inference.v2.kernels.cutlass_ops.{self.NAME}' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 8: + # Only support Ampere and newer + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def get_prefix(self): + ds_path = self.deepspeed_src_path("deepspeed") + return "deepspeed" if os.path.isdir(ds_path) else ".." + + def sources(self): + sources = [ + "inference/v2/kernels/cutlass_ops/cutlass_ops.cpp", + "inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.cu", + "inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.cu", + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources + + def extra_ldflags(self): + import dskernels + lib_path = dskernels.library_path() + prefix = self.get_prefix() + lib_path = os.path.join(prefix, lib_path) + lib_path = self.deepspeed_src_path(lib_path) + + args = [f'-L{lib_path}', '-ldeepspeedft'] + if self.jit_load: + args.append(f'-Wl,-rpath,{lib_path}') + return args + + def include_paths(self): + sources = [ + 'inference/v2/kernels/includes', + 'inference/v2/kernels/cutlass_ops/mixed_gemm', + 'inference/v2/kernels/cutlass_ops/moe_gemm', + 'inference/v2/kernels/cutlass_ops/shared_resources/', + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db12afbbf20ec7a5329b95110dbceaf5e8607d10 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +# MLU related operators will be added in the future. +from .no_impl import NotImplementedBuilder +from .cpu_adagrad import CPUAdagradBuilder +from .cpu_adam import CPUAdamBuilder +from .fused_adam import FusedAdamBuilder diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e97c23b0599c5096632e93bbb70b1e4eb17c3fce Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..48c4b14a82b32324708423a385ea735704757afc Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e428e1f4cf3650401be7029f21a92e9b0f1878aa Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adagrad.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1eb3a46ccb5847c97aae47fd2d2d18bdf18a9286 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/cpu_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87ce89ad6b539151679c74b8e6dd0a2d2e88fca0 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/fused_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6747fd9632cae1bd1c1ab64d7a1eba8cb24ec31e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/__pycache__/no_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/builder.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..17b9723ffcc1e3c81ea2e594d20b352b0a2a5bdf --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/builder.py @@ -0,0 +1,35 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class MLUOpBuilder(OpBuilder): + + def builder(self): + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + + cpp_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=self.strip_empty_entries(self.include_paths()), + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args) + + return cpp_ext + + def cxx_args(self): + return ['-O3', '-g', '-Wno-reorder'] + + def libraries_args(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..68b7bbe514eea6d7714a6f8b6135ebd622a1b7f9 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/cpu_adagrad.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import MLUOpBuilder + + +class CPUAdagradBuilder(MLUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/adagrad/cpu_adagrad.cpp'] + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..b3c8e476bf3960c3ff467f362a72947cf73babf8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/cpu_adam.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import MLUOpBuilder + + +class CPUAdamBuilder(MLUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..2343e55568c489e1ff5d2354e88b676889c7211c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/fused_adam.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import MLUOpBuilder + +try: + import torch +except ImportError as e: + pass + + +class MLUFusedAdam: + + @staticmethod + def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, + bias_correction, weight_decay, *args): + + torch.ops.torch_mlu.fused_adam(noop_flag_buffer, tensor_lists[0], tensor_lists[1], tensor_lists[2], + tensor_lists[3], lr, beta1, beta2, epsilon, step, adam_w_mode, bias_correction, + weight_decay) + + +class FusedAdamBuilder(MLUOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return MLUFusedAdam diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..375c148b4a5e06e53fb87cddbdf988708d425995 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/mlu/no_impl.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# Copyright (c) 2024 Cambricon Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import MLUOpBuilder + + +class NotImplementedBuilder(MLUOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.comm.{self.NAME}_op' + + def load(self, verbose=True): + raise ValueError("This op had not been implemented on MLU backend.") + + def sources(self): + return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e512e25a2a29e1b463bdc373722a64c755b89098 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .fused_adam import FusedAdamBuilder +from .async_io import AsyncIOBuilder +from .no_impl import NotImplementedBuilder +from .cpu_adam import CPUAdamBuilder +from .cpu_adagrad import CPUAdagradBuilder +from .cpu_lion import CPULionBuilder +from .inference import InferenceBuilder diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea58524f7ba5a9b6e4d170c2248e08d238087345 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d395c004f64e3cecc3ae3b0983290fc6fde9c2b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/async_io.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7232ec6d521a7c4947496aa1b3b810710044b0a9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d2f75f11d7080c4b2b57c4f628f5806763e43e5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adagrad.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f84971974b6b2de2ed4a62dfc39b25c7e99e4a3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51af81f693a00557c02c813edfc6a3a9f7d8b5b1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/cpu_lion.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b36bb99f571785475d5cfb6644bb8622dabf97b2 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/fused_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f1cc04489570971a990e1d9bbeee8ac862a637a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0689ddd1a296de68af970b7837b4274664006180 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/__pycache__/no_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/async_io.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..9a574ffe2dfbd2b52f42296460e348c6f619207c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/async_io.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import shutil +import subprocess + +from .builder import NPUOpBuilder + + +class AsyncIOBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_AIO" + NAME = "async_io" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.aio.{self.NAME}_op' + + def sources(self): + return [ + 'csrc/aio/py_lib/deepspeed_py_copy.cpp', 'csrc/aio/py_lib/py_ds_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio.cpp', 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', + 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', 'csrc/aio/common/deepspeed_aio_utils.cpp', + 'csrc/aio/common/deepspeed_aio_common.cpp', 'csrc/aio/common/deepspeed_aio_types.cpp', + 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp' + ] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/aio/py_lib', 'csrc/aio/common'] + return args + + def cxx_args(self): + args = super().cxx_args() + # -O0 for improved debugging, since performance is bound by I/O + CPU_ARCH = self.cpu_arch() + SIMD_WIDTH = self.simd_width() + import torch # Keep this import here to avoid errors when building DeepSpeed wheel without torch installed + TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) + if TORCH_MAJOR >= 2 and TORCH_MINOR >= 1: + CPP_STD = '-std=c++17' + else: + CPP_STD = '-std=c++14' + return args + [ + '-g', + '-Wall', + '-O0', + CPP_STD, + '-shared', + '-fPIC', + '-Wno-reorder', + CPU_ARCH, + '-fopenmp', + SIMD_WIDTH, + '-laio', + ] + + def extra_ldflags(self): + args = super().extra_ldflags() + return args + ['-laio'] + + def check_for_libaio_pkg(self): + libs = dict( + dpkg=["-l", "libaio-dev", "apt"], + pacman=["-Q", "libaio", "pacman"], + rpm=["-q", "libaio-devel", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + self.warning(f"{self.NAME}: please install the {lib} package with {tool}") + break + return found + + def is_compatible(self, verbose=False): + # Check for the existence of libaio by using distutils + # to compile and link a test program that calls io_submit, + # which is a function provided by libaio that is used in the async_io op. + # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS + # respectively to specify the directories for libaio.h and libaio.so. + aio_compatible = self.has_function('io_pgetevents', ('aio', )) + if verbose and not aio_compatible: + self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") + + # Check for the libaio package via known package managers + # to print suggestions on which package to install. + self.check_for_libaio_pkg() + + self.warning( + "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." + ) + return super().is_compatible(verbose) and aio_compatible diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/builder.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..0dea2e78915e2a939d5c1839c2e2d88e948e10f1 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/builder.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import re +import os +try: + import torch_npu +except ImportError as e: + pass + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class NPUOpBuilder(OpBuilder): + _ascend_path = None + _torch_npu_path = None + _cann_version = None + + def __init__(self, name): + super().__init__(name) + self._ascend_path = self.installed_cann_path() + self._torch_npu_path = os.path.join(os.path.dirname(os.path.abspath(torch_npu.__file__))) + try: + self._cann_version = self.installed_cann_version(self.name) + except BaseException: + print(f"{self.name} ascend_cann is missing, npu ops cannot be compiled!") + + def cann_defs(self): + if self._cann_version: + return '-D__ENABLE_CANN__' + return '-D__DISABLE_CANN__' + + def installed_cann_path(self): + if "ASCEND_HOME_PATH" in os.environ or os.path.exists(os.environ["ASCEND_HOME_PATH"]): + return os.environ["ASCEND_HOME_PATH"] + return None + + def installed_cann_version(self, name=""): + ascend_path = self.installed_cann_path() + assert ascend_path is not None, "CANN_HOME does not exist, unable to compile NPU op(s)" + cann_version = "" + for dirpath, _, filenames in os.walk(os.path.realpath(ascend_path)): + if cann_version: + break + install_files = [file for file in filenames if re.match(r"ascend_.*_install\.info", file)] + if install_files: + filepath = os.path.join(dirpath, install_files[0]) + with open(filepath, "r") as f: + for line in f: + if line.find("version") != -1: + cann_version = line.strip().split("=")[-1] + break + return cann_version + + def include_paths(self): + paths = super().include_paths() + paths += [os.path.join(self._ascend_path, 'include'), os.path.join(self._torch_npu_path, 'include')] + return paths + + def cxx_args(self): + args = super().cxx_args() + args += ['-O3', '-std=c++17', '-g', '-Wno-reorder', '-fopenmp'] + args += ['-fstack-protector-all', '-Wl,-z,relro,-z,now,-z,noexecstack', '-Wl,--disable-new-dtags,--rpath'] + args += [ + self.cann_defs(), + self.cpu_arch(), + self.simd_width(), '-L' + os.path.join(self._ascend_path, 'lib64'), + '-L' + os.path.join(self._torch_npu_path, 'lib') + ] + return args + + def extra_ldflags(self): + flags = super().extra_ldflags() + flags += [ + '-L' + os.path.join(self._ascend_path, 'lib64'), '-lascendcl', + '-L' + os.path.join(self._torch_npu_path, 'lib'), '-ltorch_npu' + ] + return flags diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..161bc82efe1ca01660fdeedd30079a8f10f1d269 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_adagrad.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPUAdagradBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/adagrad/cpu_adagrad.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..a4e9569c0f336122cd003a2df5e196527d84666c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_adam.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPUAdamBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py new file mode 100644 index 0000000000000000000000000000000000000000..6917e0fd03d08dec42e71479110224d577b55b5b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/cpu_lion.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPULionBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_LION" + NAME = "cpu_lion" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lion.{self.NAME}_op' + + def sources(self): + return ['csrc/lion/cpu_lion.cpp', 'csrc/lion/cpu_lion_impl.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..d32103db7055669f8ecce578bb4ef7703e3c07ef --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/fused_adam.py @@ -0,0 +1,74 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + +try: + import torch_npu +except ImportError as e: + pass + + +class NPUFusedAdam: + + @staticmethod + def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, + bias_correction, weight_decay, *args): + bias_correction1 = beta1**(step - 1) + bias_correction2 = beta2**(step - 1) + + # iteration group['params'] + for i in range(len(tensor_lists[0])): + grad_flat = tensor_lists[0][i] + param_flat = tensor_lists[1][i] + m_flat = tensor_lists[2][i] + v_flat = tensor_lists[3][i] + + if adam_w_mode: + param_flat.data, m_flat, v_flat = torch_npu.npu_apply_adam_w( + bias_correction1, + bias_correction2, + lr, + weight_decay, + beta1, + beta2, + epsilon, + grad_flat, + None, # max_grad_norm + False, # amsgrad + False, # maximize + out=(param_flat.data, m_flat, v_flat)) + else: + param_flat.data, m_flat, v_flat = torch_npu.npu_apply_adam( + bias_correction1, + bias_correction2, + lr, + beta1, + beta2, + epsilon, + grad_flat, + False, # use_locking + False, # use_nesterov + out=(param_flat.data, m_flat, v_flat)) + + +class FusedAdamBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return NPUFusedAdam diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/inference.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..46f28c0d401161f70431776a5a53387235ebb5ce --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/inference.py @@ -0,0 +1,307 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import IntEnum +from .builder import NPUOpBuilder + +try: + import torch + import torch_npu +except ImportError as e: + pass + + +class ActivationFuncType(IntEnum): + UNKNOWN = 0 + GELU = 1 + ReLU = 2 + GATED_GELU = 3 + GATED_SILU = 4 + + +class InferenceContext: + _workspace = None + + _seed = 42 + _curr_offset = 0 + _stream = 0 + _free_memory_size = 0 + _num_tokens = 1 + _attention_unfused_workspace_offset = 0 + _workSpaceSize = 0 + + workSpaceSize = 0 + kv_caches = None + + @staticmethod + def reset_tokens(initial_tokens=1): + InferenceContext._num_tokens = initial_tokens + + @staticmethod + def current_tokens(): + return InferenceContext._num_tokens + + @staticmethod + def GetWorkSpace(): + return InferenceContext._workspace + + +class NPUInference: + + @staticmethod + def layer_norm(inputs, gamma, beta, epsilon): + return torch.nn.functional.layer_norm(inputs, [inputs.shape[-1]], gamma, beta, eps=epsilon) + + @staticmethod + def _qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + inp_norm = torch.nn.functional.layer_norm(inputs, (inputs.shape[2], ), gamma, beta, eps) + weight = weight.t() if transpose else weight + tmp = torch.matmul(inp_norm, weight) + if add_bias: + tmp += bias + output = [tmp, inp_norm] + return output + + @staticmethod + def qkv_gemm_fp16(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) + + @staticmethod + def qkv_gemm_bf16(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) + + @staticmethod + def qkv_gemm_fp32(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose): + return NPUInference._qkv_gemm(inputs, weight, q_scale, bias, gamma, beta, eps, add_bias, q_int8, transpose) + + @staticmethod + def _bias_add_transform_0213(vals, bias, hidden_dim, seq_length, seq_offset, heads, num_kv, rotary_dim, + rotate_half, rotate_every_two, rope_theta): + bsz, _, _ = vals.shape + q = vals[..., :hidden_dim].reshape(bsz, seq_length, heads, -1) + k = vals[..., hidden_dim:hidden_dim + num_kv * (hidden_dim // heads)].reshape(bsz, seq_length, num_kv, -1) + v = vals[..., hidden_dim + num_kv * (hidden_dim // heads):] + + if rotary_dim > 0 and rotate_every_two: + # sin, cos may use cache + seq_id = torch.arange(0, seq_length).to("npu") + inv_freq = torch.arange(0, rotary_dim, 2) / rotary_dim + inv_freq = inv_freq.to("npu") + inv_freq = 1.0 / torch.pow(rope_theta, inv_freq) + inv_freq = torch.outer(seq_id, inv_freq) + sin = inv_freq.sin() + cos = inv_freq.cos() + # shape: [bsz=1, seq_len, heads=1, rotary_dim] + sin = sin.view(-1, seq_length, 1, rotary_dim // 2).repeat_interleave(2, dim=-1) + cos = cos.view(-1, seq_length, 1, rotary_dim // 2).repeat_interleave(2, dim=-1) + + q_pos, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_pos, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + + q_pos = torch_npu.npu_rotary_mul(q_pos, cos, sin) + q = torch.cat([q_pos, q_pass], dim=-1) + k_pos = torch_npu.npu_rotary_mul(k_pos, cos, sin) + k = torch.cat([k_pos, k_pass], dim=-1) + + output = q.reshape(bsz, seq_length, -1).contiguous() # [b, s, H] + k_cache = k.reshape(bsz, seq_length, heads, -1).transpose(1, 2).contiguous() # [b, n, s, d] + v_cache = v.reshape(bsz, seq_length, heads, -1).transpose(1, 2).contiguous() # [b, n, s, d] + return output, k_cache, v_cache + + @staticmethod + def _softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, + norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, + num_layers, alibi, rope_theta): + bsz, seq_len, k = query_key_value.size() + k = k // (heads + 2 * (num_kv if num_kv > 0 else heads)) + hidden_dim = heads * k + + is_promt = seq_len > 1 + if not InferenceContext.kv_caches: + InferenceContext.kv_caches = [[None, None] for _ in range(num_layers)] + if is_promt: + InferenceContext.reset_tokens(seq_len) + InferenceContext.kv_caches[layer_id] = [None, None] + + soft_len = InferenceContext.current_tokens() + workspace = InferenceContext.GetWorkSpace() + seq_offset = 0 if is_promt else soft_len - 1 + + q, k, v = NPUInference._bias_add_transform_0213(vals=query_key_value, + bias=None, + hidden_dim=hidden_dim, + seq_length=seq_len, + seq_offset=seq_offset, + heads=heads, + num_kv=num_kv if num_kv > 0 else heads, + rotary_dim=rotary_dim, + rotate_half=rotate_half, + rotate_every_two=rotate_every_two, + rope_theta=rope_theta) + + if not is_promt: + k_cache, v_cache = InferenceContext.kv_caches[layer_id] + if k_cache is not None: + k = torch.cat([k_cache, k], dim=2) + v = torch.cat([v_cache, v], dim=2) + InferenceContext.kv_caches[layer_id] = [k, v] + seq_len = k.shape[2] + + layer_scale = max(1, layer_id) if len(alibi.size()) > 1 else 1.0 + alpha = norm_factor * norm_factor / layer_scale + + output = torch_npu.npu_fusion_attention(q, + k.transpose(1, 2).reshape(bsz, seq_len, -1).contiguous(), + v.transpose(1, 2).reshape(bsz, seq_len, -1).contiguous(), + heads, + "BSH", + pse=None, + padding_mask=None, + atten_mask=attn_mask.bool(), + scale=alpha, + pre_tockens=65536, + next_tockens=65536, + keep_prob=1, + inner_precise=0)[0] + + return output, k, v + + @staticmethod + def softmax_context_fp16(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, + norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, + num_layers, alibi, rope_theta): + return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, + heads, num_kv, norm_factor, triangular_masking, local_attention, + window_size, no_masking, layer_id, num_layers, alibi, rope_theta) + + @staticmethod + def softmax_context_bf16(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, + norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, + num_layers, alibi, rope_theta): + return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, + heads, num_kv, norm_factor, triangular_masking, local_attention, + window_size, no_masking, layer_id, num_layers, alibi, rope_theta) + + @staticmethod + def softmax_context_fp32(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads, num_kv, + norm_factor, triangular_masking, local_attention, window_size, no_masking, layer_id, + num_layers, alibi, rope_theta): + return NPUInference._softmax_context(query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, + heads, num_kv, norm_factor, triangular_masking, local_attention, + window_size, no_masking, layer_id, num_layers, alibi, rope_theta) + + @staticmethod + def _vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode): + if transposed_mode: + return torch.matmul(input, weight.t()) + return torch.matmul(input, weight) + + @staticmethod + def vector_matmul_fp16(input, weight, async_op, q_scale, q_int8, transposed_mode): + return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) + + @staticmethod + def vector_matmul_bf16(input, weight, async_op, q_scale, q_int8, transposed_mode): + return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) + + @staticmethod + def vector_matmul_fp32(input, weight, async_op, q_scale, q_int8, transposed_mode): + return NPUInference._vector_matmul(input, weight, async_op, q_scale, q_int8, transposed_mode) + + @staticmethod + def _mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, + mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): + if mlp_after_attn: + residual_add = torch.nn.functional.layer_norm(input + residual + input_bias, (input.shape[-1], ), gamma, + beta, eps) + else: + residual_add = torch.nn.functional.layer_norm(input, (input.shape[-1], ), gamma, beta, eps) + + weight_interm = weight_interm.t() if transpose else weight_interm + tmp = torch.matmul(residual_add, weight_interm) + if mlp_act_func_type == ActivationFuncType.GELU: + tmp = torch.nn.functional.gelu(tmp + bias) + elif mlp_act_func_type == ActivationFuncType.ReLU: + tmp = torch.nn.functional.relu(tmp + bias) + else: + raise Exception('Unsupported ActivationFuncType {}'.format(mlp_act_func_type)) + output = torch.matmul(tmp, weight_out.t()) + return output, residual_add + + @staticmethod + def mlp_gemm_fp16(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, + mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): + return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, + pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, + mlp_act_func_type, transpose) + + @staticmethod + def mlp_gemm_bf16(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, + mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): + return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, + pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, + mlp_act_func_type, transpose) + + @staticmethod + def mlp_gemm_fp32(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, pre_layer_norm, + mlp_after_attn, interm_scale, out_scale, dtype, mlp_act_func_type, transpose): + return NPUInference._mlp_gemm(input, residual, input_bias, weight_interm, weight_out, bias, gamma, beta, eps, + pre_layer_norm, mlp_after_attn, interm_scale, out_scale, dtype, + mlp_act_func_type, transpose) + + @staticmethod + def _residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + if mlp_after_attn: + if pre_layer_norm: + tmp = (residual.float() + attention_output.float() + attention_bias.float() + + final_bias.float()) / mp_size + hidden_state.float() + else: + tmp = residual.float() + hidden_state.float() + final_bias.float() + else: + if add_bias: + residual += attention_bias.float() + tmp = hidden_state.float() + attention_output.float() + (residual.float() + final_bias.float()) / mp_size + + input_dtype = hidden_state.dtype + residual.set_(tmp.to(input_dtype)) + + @staticmethod + def residual_add_bias_fp16(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, + mp_size, mlp_after_attn, add_bias, pre_layer_norm) + + @staticmethod + def residual_add_bias_bf16(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, + mp_size, mlp_after_attn, add_bias, pre_layer_norm) + + @staticmethod + def residual_add_bias_fp32(hidden_state, residual, attention_output, attention_bias, final_bias, mp_size, + mlp_after_attn, add_bias, pre_layer_norm): + return NPUInference._residual_add_bias(hidden_state, residual, attention_output, attention_bias, final_bias, + mp_size, mlp_after_attn, add_bias, pre_layer_norm) + + +class InferenceBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" + NAME = "transformer_inference" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.transformer.inference.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return NPUInference diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/no_impl.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..5b1771fabc22f6ad13161231900a737bba733e68 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/npu/no_impl.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class NotImplementedBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.comm.{self.NAME}_op' + + def load(self, verbose=True): + raise ValueError("This op had not been implemented on NPU backend.") + + def sources(self): + return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/quantizer.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..0b5348e5af96e28e256524b38efaff5d50863d17 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/quantizer.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + + +class QuantizerBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_QUANTIZER" + NAME = "quantizer" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.quantizer.{self.NAME}_op' + + def sources(self): + return [ + 'csrc/quantization/pt_binding.cpp', + 'csrc/quantization/fake_quantizer.cu', + 'csrc/quantization/quantize.cu', + 'csrc/quantization/quantize_intX.cu', + 'csrc/quantization/dequantize.cu', + 'csrc/quantization/swizzled_quantize.cu', + 'csrc/quantization/quant_reduce.cu', + ] + + def include_paths(self): + return ['csrc/includes'] + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/ragged_ops.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/ragged_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..0df28cc2282a0beee59e4e5ae4d84c3a94ef9aa3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/ragged_ops.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class RaggedOpsBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_RAGGED_DEVICE_OPS" + NAME = "ragged_device_ops" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.inference.v2.kernels.ragged_ops.{self.NAME}' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 8: + # Blocked flash has a dependency on Ampere + newer + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def get_prefix(self): + ds_path = self.deepspeed_src_path("deepspeed") + return "deepspeed" if os.path.isdir(ds_path) else ".." + + def sources(self): + sources = [ + "inference/v2/kernels/ragged_ops/ragged_ops.cpp", + "inference/v2/kernels/ragged_ops/atom_builder/atom_builder.cpp", + "inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.cpp", + "inference/v2/kernels/ragged_ops/embed/embed.cpp", + "inference/v2/kernels/ragged_ops/embed/embed_cuda.cu", + "inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.cpp", + "inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary_cuda.cu", + "inference/v2/kernels/ragged_ops/logits_gather/logits_gather.cpp", + "inference/v2/kernels/ragged_ops/logits_gather/logits_gather_cuda.cu", + "inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter.cpp", + "inference/v2/kernels/ragged_ops/moe_scatter/moe_scatter_cuda.cu", + "inference/v2/kernels/ragged_ops/moe_gather/moe_gather.cpp", + "inference/v2/kernels/ragged_ops/moe_gather/moe_gather_cuda.cu", + "inference/v2/kernels/ragged_ops/ragged_helpers/ragged_kernel_helpers.cpp", + "inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating.cpp", + "inference/v2/kernels/ragged_ops/top_k_gating/top_k_gating_cuda.cu", + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources + + def extra_ldflags(self): + import dskernels + lib_path = dskernels.library_path() + + prefix = self.get_prefix() + lib_path = os.path.join(prefix, lib_path) + lib_path = self.deepspeed_src_path(lib_path) + + args = [f'-L{lib_path}', '-lblockedflash'] + if self.jit_load: + args.append(f'-Wl,-rpath,{lib_path}') + return args + + def include_paths(self): + sources = [ + 'inference/v2/kernels/includes', + 'inference/v2/kernels/ragged_ops', + 'inference/v2/kernels/ragged_ops/atom_builder', + 'inference/v2/kernels/ragged_ops/blocked_flash', + 'inference/v2/kernels/ragged_ops/embed', + 'inference/v2/kernels/ragged_ops/includes', + 'inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary', + 'inference/v2/kernels/ragged_ops/logits_gather', + 'inference/v2/kernels/ragged_ops/moe_gather', + 'inference/v2/kernels/ragged_ops/moe_scatter', + 'inference/v2/kernels/ragged_ops/ragged_helpers', + 'inference/v2/kernels/ragged_ops/top_k_gating', + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/ragged_utils.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/ragged_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..208c9f833ebe97c39593e66c46a53b71190fedfd --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/ragged_utils.py @@ -0,0 +1,80 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class RaggedUtilsBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_RAGGED_OPS" + NAME = "ragged_ops" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.inference.v2.{self.NAME}' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 6: + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def get_prefix(self): + ds_path = self.deepspeed_src_path("deepspeed") + return "deepspeed" if os.path.isdir(ds_path) else ".." + + def sources(self): + sources = [ + "inference/v2/ragged/csrc/fast_host_buffer.cu", + "inference/v2/ragged/csrc/ragged_ops.cpp", + ] + + prefix = self.get_prefix() + sources = [os.path.join(prefix, src) for src in sources] + return sources + + def extra_ldflags(self): + return [] + + def include_paths(self): + include_dirs = ['inference/v2/ragged/includes', 'inference/v2/kernels/includes'] + prefix = self.get_prefix() + includes = [os.path.join(prefix, include_dir) for include_dir in include_dirs] + + return includes diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/random_ltd.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/random_ltd.py new file mode 100644 index 0000000000000000000000000000000000000000..54af7150fb36f9eb8bd6a295648d8a11da8d8373 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/random_ltd.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + + +class RandomLTDBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_RANDOM_LTD" + NAME = "random_ltd" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def sources(self): + return [ + 'csrc/random_ltd/pt_binding.cpp', 'csrc/random_ltd/gather_scatter.cu', + 'csrc/random_ltd/slice_attn_masks.cu', 'csrc/random_ltd/token_sort.cu' + ] + + def include_paths(self): + includes = ['csrc/includes'] + return includes diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2a6eb8bbfa2c52c3e61ddbca734011c2ec5d630b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) Microsoft Corporation. + +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' + +from .no_impl import NotImplementedBuilder +from .cpu_adam import CPUAdamBuilder +from .fused_adam import FusedAdamBuilder diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4bf0a33c692b169c4007c85b62cb4e15872274e9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c293f84c319e7c98d639600da4c0c84e8b2f9cb3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e9ea1ffba1bc1dc7db80b13787e0cd951e1dbd1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/cpu_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d995539857978edc74af2965fbc506b017c93651 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/fused_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..75341bbaac19135f521402b9ac1ee9ffb69abd2a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/__pycache__/no_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/builder.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..81f0e98c076837650bd5a3cd9a23493361e6b02e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/builder.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class SDAAOpBuilder(OpBuilder): + + def builder(self): + from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + + compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + + cpp_ext = ExtensionBuilder(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=self.strip_empty_entries(self.include_paths()), + libraries=self.strip_empty_entries(self.libraries_args()), + extra_compile_args=compile_args) + + return cpp_ext + + def cxx_args(self): + return ['-O3', '-g', '-Wno-reorder'] + + def libraries_args(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..b38a71a6275d3b962ed2bd7d8f625ba3298418dd --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/cpu_adam.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team + +from .builder import SDAAOpBuilder + + +class CPUAdamBuilder(SDAAOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..73a2dff4145961a909258278e64a6fd094ca0206 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/fused_adam.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team + +from .builder import SDAAOpBuilder + +try: + import torch +except ImportError as e: + pass + + +class SDAAFusedAdam: + + @staticmethod + def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, + bias_correction, weight_decay, *args): + g_tensor_lis, p_tensor_lis, m_tensor_lis, v_tensor_lis = tensor_lists + torch.ops.sdaa.fused_adam(g_tensor_lis, p_tensor_lis, m_tensor_lis, v_tensor_lis, [], beta1, beta2, epsilon, + lr, weight_decay, adam_w_mode, step, bias_correction) + + +class FusedAdamBuilder(SDAAOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return [] + + def include_paths(self): + return [] + + def load(self, verbose=True): + return SDAAFusedAdam diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..10a8b8f48652a8abbd9978fb444332036571acb3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sdaa/no_impl.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# BSD 3- Clause License Copyright (c) 2023, Tecorigin Co., Ltd. All rights +# reserved. +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) +# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +# STRICT LIABILITY,OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +# OF SUCH DAMAGE. + +# DeepSpeed Team + +from .builder import SDAAOpBuilder + + +class NotImplementedBuilder(SDAAOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.comm.{self.NAME}_op' + + def load(self, verbose=True): + raise ValueError("This op had not been implemented on SDAA backend.") + + def sources(self): + return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/sparse_attn.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sparse_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..b8e5e8a6fd9aaa62372f72c0b78d250994b3e98e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/sparse_attn.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import OpBuilder + +try: + from packaging import version as pkg_version +except ImportError: + pkg_version = None + + +class SparseAttnBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_SPARSE_ATTN" + NAME = "sparse_attn" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.sparse_attention.{self.NAME}_op' + + def sources(self): + return ['csrc/sparse_attention/utils.cpp'] + + def cxx_args(self): + return ['-O2', '-fopenmp'] + + def is_compatible(self, verbose=False): + # Check to see if llvm and cmake are installed since they are dependencies + #required_commands = ['llvm-config|llvm-config-9', 'cmake'] + #command_status = list(map(self.command_exists, required_commands)) + #deps_compatible = all(command_status) + + if self.is_rocm_pytorch(): + if verbose: + self.warning(f'{self.NAME} is not compatible with ROCM') + return False + + try: + import torch + except ImportError: + if verbose: + self.warning(f"unable to import torch, please install it first") + return False + + # torch-cpu will not have a cuda version + if torch.version.cuda is None: + cuda_compatible = False + if verbose: + self.warning(f"{self.NAME} cuda is not available from torch") + else: + major, minor = torch.version.cuda.split('.')[:2] + cuda_compatible = (int(major) == 10 and int(minor) >= 1) or (int(major) >= 11) + if not cuda_compatible: + if verbose: + self.warning(f"{self.NAME} requires CUDA version 10.1+") + + TORCH_MAJOR = int(torch.__version__.split('.')[0]) + TORCH_MINOR = int(torch.__version__.split('.')[1]) + torch_compatible = (TORCH_MAJOR == 1 and TORCH_MINOR >= 5) + if not torch_compatible: + if verbose: + self.warning( + f'{self.NAME} requires a torch version >= 1.5 and < 2.0 but detected {TORCH_MAJOR}.{TORCH_MINOR}') + try: + import triton + except ImportError: + # auto-install of triton is broken on some systems, reverting to manual install for now + # see this issue: https://github.com/deepspeedai/DeepSpeed/issues/1710 + if verbose: + self.warning(f"please install triton==1.0.0 if you want to use sparse attention") + return False + + if pkg_version: + installed_triton = pkg_version.parse(triton.__version__) + triton_mismatch = installed_triton != pkg_version.parse("1.0.0") + else: + installed_triton = triton.__version__ + triton_mismatch = installed_triton != "1.0.0" + + if triton_mismatch: + if verbose: + self.warning( + f"using untested triton version ({installed_triton}), only 1.0.0 is known to be compatible") + return False + + return super().is_compatible(verbose) and torch_compatible and cuda_compatible diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/spatial_inference.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/spatial_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..d6c5fa6611563bc0a45fe712e27f0f06f56bf2d0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/spatial_inference.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class SpatialInferenceBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_SPATIAL_INFERENCE" + NAME = "spatial_inference" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.spatial.{self.NAME}_op' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def sources(self): + return [ + 'csrc/spatial/csrc/opt_bias_add.cu', + 'csrc/spatial/csrc/pt_binding.cpp', + ] + + def include_paths(self): + return ['csrc/spatial/includes', 'csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..52b02a3c629e9d537b8e22139f48f5335396e4ff --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/stochastic_transformer.py @@ -0,0 +1,22 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .transformer import TransformerBuilder + + +class StochasticTransformerBuilder(TransformerBuilder): + BUILD_VAR = "DS_BUILD_STOCHASTIC_TRANSFORMER" + NAME = "stochastic_transformer" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.transformer.{self.NAME}_op' + + def nvcc_args(self): + args = super().nvcc_args() + args.append('-D__STOCHASTIC_MODE__') + return args diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/transformer.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..8db30fdc67919474679715014cbf8d3f470bedf4 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/transformer.py @@ -0,0 +1,36 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder + + +class TransformerBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER" + NAME = "transformer" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.transformer.{self.NAME}_op' + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def sources(self): + return [ + 'csrc/transformer/ds_transformer_cuda.cpp', 'csrc/transformer/cublas_wrappers.cu', + 'csrc/transformer/transform_kernels.cu', 'csrc/transformer/gelu_kernels.cu', + 'csrc/transformer/dropout_kernels.cu', 'csrc/transformer/normalize_kernels.cu', + 'csrc/transformer/softmax_kernels.cu', 'csrc/transformer/general_kernels.cu' + ] + + def include_paths(self): + includes = ['csrc/includes'] + return includes diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/transformer_inference.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/transformer_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..642aed56a192e4b5170a62d8cc457bb451ae8646 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/transformer_inference.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import CUDAOpBuilder, installed_cuda_version + + +class InferenceBuilder(CUDAOpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" + NAME = "transformer_inference" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.transformer.inference.{self.NAME}_op' + + def is_compatible(self, verbose=False): + try: + import torch + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile inference kernels") + return False + + cuda_okay = True + if not self.is_rocm_pytorch() and torch.cuda.is_available(): + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major + if cuda_capability < 6: + if verbose: + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") + cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False + return super().is_compatible(verbose) and cuda_okay + + def filter_ccs(self, ccs): + ccs_retained = [] + ccs_pruned = [] + for cc in [cc.split('.') for cc in ccs]: + if int(cc[0]) >= 6: + ccs_retained.append(cc) + else: + ccs_pruned.append(cc) + if len(ccs_pruned) > 0: + self.warning(f"Filtered compute capabilities {ccs_pruned}") + return ccs_retained + + def sources(self): + return [ + 'csrc/transformer/inference/csrc/pt_binding.cpp', + 'csrc/transformer/inference/csrc/gelu.cu', + 'csrc/transformer/inference/csrc/relu.cu', + 'csrc/transformer/inference/csrc/layer_norm.cu', + 'csrc/transformer/inference/csrc/rms_norm.cu', + 'csrc/transformer/inference/csrc/softmax.cu', + 'csrc/transformer/inference/csrc/dequantize.cu', + 'csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu', + 'csrc/transformer/inference/csrc/transform.cu', + 'csrc/transformer/inference/csrc/pointwise_ops.cu', + ] + + def extra_ldflags(self): + if not self.is_rocm_pytorch(): + return ['-lcurand'] + else: + return [] + + def include_paths(self): + return ['csrc/transformer/inference/includes', 'csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..093f90b30234359d37913bb36919f3259d4cab1c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .cpu_adam import CPUAdamBuilder +from .cpu_adagrad import CPUAdagradBuilder +from .fused_adam import FusedAdamBuilder +from .async_io import AsyncIOBuilder +from .inference import InferenceBuilder +from .flash_attn import FlashAttentionBuilder +from .no_impl import NotImplementedBuilder +from .packbits import PackbitsBuilder diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..26aeccecdb5a5661e77c969effd9d23c09d739a6 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b960d7db74bf4c76c0d117bb929e100ef665161d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/async_io.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ce3d77d1da46735feb5b67ac9981db67a51e88e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3809618648057042423242c04ad706d0e75fd8db Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adagrad.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40023950feb2adb97e76347ae9d89d928ea3ba4c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/cpu_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14cfb224b28d6bac327bfb4e7af4ccc401e7ae23 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/flash_attn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47cad1e43bee5817c22b3cac1619b4611ded3185 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/fused_adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70eedf21f5aa5ed6b76ce02552517f218d9ead2c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/inference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22cb3809e5988ef4d8cf0fe127eff8fcba17802a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/no_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7ee981341ef85528ccac85d58beb08095ac1740 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/__pycache__/packbits.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/async_io.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/async_io.py new file mode 100644 index 0000000000000000000000000000000000000000..8ec030880368dff58337b9e78bfd80bad9f31f97 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/async_io.py @@ -0,0 +1,106 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import shutil +import subprocess + +from .builder import OpBuilder + + +class AsyncIOBuilder(OpBuilder): + BUILD_VAR = "DS_BUILD_AIO" + NAME = "async_io" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.aio.{self.NAME}_op' + + def sources(self): + return [ + 'csrc/aio/py_lib/deepspeed_py_copy.cpp', + 'csrc/aio/py_lib/py_ds_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio.cpp', + 'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', + 'csrc/aio/py_lib/deepspeed_aio_thread.cpp', + 'csrc/aio/common/deepspeed_aio_utils.cpp', + 'csrc/aio/common/deepspeed_aio_common.cpp', + 'csrc/aio/common/deepspeed_aio_types.cpp', + 'csrc/aio/py_lib/deepspeed_pin_tensor.cpp', + 'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', + 'csrc/aio/py_lib/deepspeed_cpu_op.cpp', + 'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', + ] + + def include_paths(self): + return ['csrc/aio/py_lib', 'csrc/aio/common'] + + def cxx_args(self): + import torch + # -O0 for improved debugging, since performance is bound by I/O + CPU_ARCH = self.cpu_arch() + SIMD_WIDTH = self.simd_width() + TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2]) + if TORCH_MAJOR >= 2 and TORCH_MINOR >= 1: + CPP_STD = '-std=c++17' + else: + CPP_STD = '-std=c++14' + return [ + '-g', + '-Wall', + '-O0', + CPP_STD, + '-shared', + '-fPIC', + '-Wno-reorder', + CPU_ARCH, + '-fopenmp', + SIMD_WIDTH, + '-laio', + ] + + def extra_ldflags(self): + return ['-laio'] + + def check_for_libaio_pkg(self): + libs = dict( + dpkg=["-l", "libaio-dev", "apt"], + pacman=["-Q", "libaio", "pacman"], + rpm=["-q", "libaio-devel", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + self.warning(f"{self.NAME}: please install the {lib} package with {tool}") + break + return found + + def is_compatible(self, verbose=False): + # Check for the existence of libaio by using distutils + # to compile and link a test program that calls io_submit, + # which is a function provided by libaio that is used in the async_io op. + # If needed, one can define -I and -L entries in CFLAGS and LDFLAGS + # respectively to specify the directories for libaio.h and libaio.so. + aio_compatible = self.has_function('io_pgetevents', ('aio', )) + if verbose and not aio_compatible: + self.warning(f"{self.NAME} requires the dev libaio .so object and headers but these were not found.") + + # Check for the libaio package via known package managers + # to print suggestions on which package to install. + self.check_for_libaio_pkg() + + self.warning( + "If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found." + ) + return super().is_compatible(verbose) and aio_compatible diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/builder.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..81b15f197f4375b6138f50cf865b2dd64194c629 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/builder.py @@ -0,0 +1,131 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import time +import importlib + +try: + # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed + # if successful this also means we're doing a local install and not JIT compile path + from op_builder import __deepspeed__ # noqa: F401 # type: ignore + from op_builder.builder import OpBuilder +except ImportError: + from deepspeed.ops.op_builder.builder import OpBuilder + + +class SYCLOpBuilder(OpBuilder): + + def builder(self): + try: + from intel_extension_for_pytorch.xpu.cpp_extension import DPCPPExtension + except ImportError: + from intel_extension_for_pytorch.xpu.utils import DPCPPExtension + include_dirs = [os.path.abspath(x) for x in self.strip_empty_entries(self.include_paths())] + print("dpcpp sources = {}".format(self.sources())) + dpcpp_ext = DPCPPExtension(name=self.absolute_name(), + sources=self.strip_empty_entries(self.sources()), + include_dirs=include_dirs, + extra_compile_args={ + 'cxx': self.strip_empty_entries(self.cxx_args()), + }, + extra_link_args=self.strip_empty_entries(self.fixed_aotflags())) + return dpcpp_ext + + def version_dependent_macros(self): + try: + from op_builder.builder import TORCH_MAJOR, TORCH_MINOR + except ImportError: + from deepspeed.ops.op_builder.builder import TORCH_MAJOR, TORCH_MINOR + # Fix from apex that might be relevant for us as well, related to https://github.com/NVIDIA/apex/issues/456 + version_ge_1_1 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0): + version_ge_1_1 = ['-DVERSION_GE_1_1'] + version_ge_1_3 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2): + version_ge_1_3 = ['-DVERSION_GE_1_3'] + version_ge_1_5 = [] + if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4): + version_ge_1_5 = ['-DVERSION_GE_1_5'] + return version_ge_1_1 + version_ge_1_3 + version_ge_1_5 + + def cxx_args(self): + cxx_flags = [ + '-fsycl', '-fsycl-targets=spir64_gen', '-g', '-gdwarf-4', '-O3', '-std=c++17', '-fPIC', '-DMKL_ILP64', + '-fno-strict-aliasing' + ] + if os.environ.get('USE_MKL_GEMM'): + cxx_flags.append('-DUSE_MKL_GEMM') + return cxx_flags + + def extra_ldflags(self): + return [ + '-fPIC', '-fsycl', '-fsycl-targets=spir64_gen', '-fsycl-max-parallel-link-jobs=8', + '-Xs "-options -cl-poison-unsupported-fp64-kernels,cl-intel-enable-auto-large-GRF-mode"', + '-Xs "-device pvc"', '-Wl,-export-dynamic' + ] + + def fixed_aotflags(self): + return [ + '-fsycl', '-fsycl-targets=spir64_gen', '-fsycl-max-parallel-link-jobs=8', '-Xs', + "-options -cl-poison-unsupported-fp64-kernels,cl-intel-enable-auto-large-GRF-mode", '-Xs', "-device pvc" + ] + + def load(self, verbose=True): + from deepspeed.git_version_info import installed_ops, torch_info, accelerator_name # noqa: F401 + from deepspeed.accelerator import get_accelerator + if installed_ops.get(self.name, False) and accelerator_name == get_accelerator()._name: + return importlib.import_module(self.absolute_name()) + else: + return self.jit_load(verbose) + + def jit_load(self, verbose=True): + if not self.is_compatible(verbose): + raise RuntimeError( + f"Unable to JIT load the {self.name} op due to it not being compatible due to hardware/software issue. {self.error_log}" + ) + try: + import ninja # noqa: F401 + except ImportError: + raise RuntimeError(f"Unable to JIT load the {self.name} op due to ninja not being installed.") + + self.jit_mode = True + from intel_extension_for_pytorch.xpu.cpp_extension import load + + start_build = time.time() + # Recognize relative paths as absolute paths for jit load + + sources = [self.deepspeed_src_path(path) for path in self.sources()] + extra_include_paths = [self.deepspeed_src_path(path) for path in self.include_paths()] + + # Torch will try and apply whatever CCs are in the arch list at compile time, + # we have already set the intended targets ourselves we know that will be + # needed at runtime. This prevents CC collisions such as multiple __half + # implementations. Stash arch list to reset after build. + ''' + torch_arch_list = None + if "TORCH_CUDA_ARCH_LIST" in os.environ: + torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST") + os.environ["TORCH_CUDA_ARCH_LIST"] = "" + ''' + + op_module = load( + name=self.name, + sources=self.strip_empty_entries(sources), + extra_include_paths=self.strip_empty_entries(extra_include_paths), + extra_cflags=self.strip_empty_entries(self.cxx_args()), + # extra_cuda_cflags=self.strip_empty_entries(self.nvcc_args()), + extra_ldflags=self.strip_empty_entries(self.extra_ldflags()), + verbose=verbose) + + build_duration = time.time() - start_build + if verbose: + print(f"Time to load {self.name} op: {build_duration} seconds") + ''' + # Reset arch list so we are not silently removing it for other possible use cases + if torch_arch_list: + os.environ["TORCH_CUDA_ARCH_LIST"] = torch_arch_list + ''' + return op_module diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..18f80848e1b800de7f8826405bb9ea9f500db263 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/cpu_adagrad.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import SYCLOpBuilder + + +class CPUAdagradBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/xpu/adagrad/cpu_adagrad.cpp', 'csrc/xpu/common/custom_cuda_kernel.dp.cpp'] + + def include_paths(self): + return ['csrc/xpu/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..9d5fdcd3e8ab51b84bb2fec137b95ee96fee8570 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/cpu_adam.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import SYCLOpBuilder + + +class CPUAdamBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def libraries_args(self): + args = super().libraries_args() + return args + + def include_paths(self): + return ['csrc/includes'] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..c8c2674d5d27c87377d3a64bfbba401a44ae9d52 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/flash_attn.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .builder import SYCLOpBuilder + + +class FlashAttentionBuilderObject(): + + def __init__(self): + pass + + # general functions + def flash_attn_func_v2(self, q, k, v, dropout_p, softmax_scale, is_causal): + try: + import torch + import intel_extension_for_pytorch # noqa + return torch.nn.functional.scaled_dot_product_attention(q, + k, + v, + dropout_p=dropout_p, + is_causal=is_causal, + scale=softmax_scale) + except ImportError: + raise ImportError( + "Please install pytorch and intel_extension_for_pytorch to include scaled dot product attention.") + + +class FlashAttentionBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_FlashAttention" + NAME = "flash_attn" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def sources(self): + return + + def include_paths(self): + return [] + + def extra_ldflags(self): + return [] + + def cxx_args(self): + return [] + + def load(self): + return FlashAttentionBuilderObject() diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..0e0f1a66f8e648305d2ba6fd0ad85a384e6f8bd3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/fused_adam.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .builder import SYCLOpBuilder + + +class FusedAdamBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_FUSED_ADAM" + NAME = "fused_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/xpu/adam/fused_adam_frontend.cpp', 'csrc/xpu/adam/multi_tensor_adam.dp.cpp'] + + def include_paths(self): + return ['csrc/xpu/includes', 'csrc/xpu/adam'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/inference.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..a9ac4f84c2ca9ba885e586934efd406b6469bd9e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/inference.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .builder import SYCLOpBuilder + + +class InferenceBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_TRANSFORMER_INFERENCE" + NAME = "transformer_inference" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.transformer.inference.{self.NAME}_op' + + def sources(self): + return + + def include_paths(self): + return [] + + def extra_ldflags(self): + return [] + + def cxx_args(self): + return [] + + def load(self): + try: + import intel_extension_for_pytorch + if hasattr(intel_extension_for_pytorch, "deepspeed"): + return intel_extension_for_pytorch.deepspeed.transformer_inference.transformer_inference + else: + return intel_extension_for_pytorch.xpu.deepspeed + except ImportError: + raise ImportError("Please install intel-extension-for-pytorch >= 2.1.30 to include DeepSpeed kernels.") diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..8b294f70c2791db9686e8d463ad736ff2d7c90c5 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/no_impl.py @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import SYCLOpBuilder + + +class NotImplementedBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_NOT_IMPLEMENTED" + NAME = "deepspeed_not_implemented" + + def __init__(self, name=None): + name = self.NAME if name is None else name + super().__init__(name=name) + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def load(self, verbose=True): + raise ValueError("This op had not been implemented on XPU backend.") + + def sources(self): + return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return [] diff --git a/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/packbits.py b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/packbits.py new file mode 100644 index 0000000000000000000000000000000000000000..cf5b5ebc59e407ae20f478151846e282680fe4ae --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/op_builder/xpu/packbits.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .builder import SYCLOpBuilder + + +class PackbitsBuilder(SYCLOpBuilder): + BUILD_VAR = "DS_BUILD_PACK_BITS" + NAME = "pack_bits" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.{self.NAME}_op' + + def sources(self): + return ['csrc/xpu/packbits/packing.cpp'] + + def include_paths(self): + return ['csrc/xpu/includes'] + + def cxx_args(self): + args = super().cxx_args() + return args + self.version_dependent_macros() diff --git a/lib/python3.12/site-packages/deepspeed/ops/quantizer/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/quantizer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0a5d1da5e3ae0fa097e7313ddb1328c4f910801d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/quantizer/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .quantizer import ds_quantizer diff --git a/lib/python3.12/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d00cd2c1c35289faa66c4b89d2ef680a34db121 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/quantizer/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41c2226c3839fa0965d544e80c654f553eabda51 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/quantizer/__pycache__/quantizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/quantizer/quantizer.py b/lib/python3.12/site-packages/deepspeed/ops/quantizer/quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..eb4bfd35700075f3b32db329c5f7026b39bef520 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/quantizer/quantizer.py @@ -0,0 +1,29 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + +from deepspeed.ops.op_builder import QuantizerBuilder + +# Cuda modules will be imported if needed +quantizer_cuda_module = None + + +def ds_quantizer(input, groups=1, bit_num=8, sr=False, asym=False): + # Load cuda modules if needed + global quantizer_cuda_module + if quantizer_cuda_module is None: + quantizer_cuda_module = QuantizerBuilder().load() + if sr: + if asym: + quantize_func = quantizer_cuda_module.ds_sr_quantize_asym_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_sr_quantize_asym_fp32 + else: + quantize_func = quantizer_cuda_module.ds_sr_quantize_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_sr_quantize_fp32 + else: + if asym: + quantize_func = quantizer_cuda_module.ds_quantize_asym_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_quantize_asym_fp32 + else: + quantize_func = quantizer_cuda_module.ds_quantize_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_quantize_fp32 + return quantize_func(input, groups, bit_num) diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e357257869f794a06d575bfa378769f8e6d3d43c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .sparsity_config import SparsityConfig, DenseSparsityConfig, FixedSparsityConfig, VariableSparsityConfig, BigBirdSparsityConfig, BSLongformerSparsityConfig, LocalSlidingWindowSparsityConfig +from .sparse_self_attention import SparseSelfAttention +from .bert_sparse_self_attention import BertSparseSelfAttention +from .sparse_attention_utils import SparseAttentionUtils diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddd50a6c1f4a52f97ea48b590410dca35125d585 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9e5ae2435054d35e50b20dafaf643b1b9fa6095 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/bert_sparse_self_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bbba263bf28aea1443e5ab03e27d9e008f6ac4e2 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/matmul.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..879416f848876934033d55f1e183c60f13e1c560 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/softmax.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8d71df576fa25e2f72d362ef4c3eac36d55cc8a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_attention_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97f81fa265291722e1b1caac338888320c449fef Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparse_self_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..377d45e1741788e716bd7e03f015a72e60667704 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/__pycache__/sparsity_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..37f065e4863111275aae82bcd0bc1ae8513a0896 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/bert_sparse_self_attention.py @@ -0,0 +1,77 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from torch import nn +from deepspeed.ops.sparse_attention import SparseSelfAttention, FixedSparsityConfig + + +class BertSparseSelfAttention(nn.Module): + """Implements Sparse Self Attention layer of Bert model based on https://github.com/deepspeedai/DeepSpeedExamples/blob/master/bing_bert/nvidia/modelingpreln.py#L373 + + For more information please see, TODO DeepSpeed Sparse Transformer. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial. + """ + + def __init__( + self, + config, + # SparsityConfig parameters needs to be set accordingly + sparsity_config=FixedSparsityConfig(num_heads=4)): + """Initialize the bert sparse self attention layer. + + Note) you can use any of the provided sparsity configs or simply add yours! + + Arguments: + config: required: Bert model config + sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on FixedSparsityConfig class. + """ + + super(BertSparseSelfAttention, self).__init__() + if config.hidden_size % config.num_attention_heads != 0: + raise ValueError("The hidden size (%d) is not a multiple of the number of attention " + "heads (%d)" % (config.hidden_size, config.num_attention_heads)) + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.sparse_self_attention = SparseSelfAttention(sparsity_config) + + def transpose_for_scores(self, x): + new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) + x = x.view(*new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward(self, hidden_states, attention_mask): + """Applies forward phase of bert sparse self attention + + Arguments: + hidden_states: required: hidden_states tensor of the bert model + attn_mask: required: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported + + Return: + context_layer: a dense tensor containing attention context + """ + mixed_query_layer = self.query(hidden_states) + mixed_key_layer = self.key(hidden_states) + mixed_value_layer = self.value(hidden_states) + + query_layer = self.transpose_for_scores(mixed_query_layer) + key_layer = self.transpose_for_scores(mixed_key_layer) + value_layer = self.transpose_for_scores(mixed_value_layer) + + context_layer = self.sparse_self_attention(query_layer, + key_layer, + value_layer, + key_padding_mask=attention_mask) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size, ) + context_layer = context_layer.view(*new_context_layer_shape) + return context_layer diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/matmul.py b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/matmul.py new file mode 100644 index 0000000000000000000000000000000000000000..b30028fffbaafaf6869156406e0f8d70c8e97538 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/matmul.py @@ -0,0 +1,819 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a +# https://github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py +import importlib +import torch + +import triton +import triton.language as tl +import triton._C.libtriton as libtriton +from deepspeed.accelerator import get_accelerator + + +@triton.jit +def _kernel(A, B, C, stride_za, stride_ha, stride_ma, stride_ka, stride_zb, stride_hb, stride_kb, stride_nb, stride_zc, + stride_hc, stride_mc, stride_nc, DS0, DS1, SDD_K, SDD_off_width, lut, locks, nlocks, **meta): + TM = meta['TM'] + TN = meta['TN'] + TK = meta['TK'] + TZ = meta['TZ'] + BLOCK = meta['BLOCK'] + #------------# + #- Prologue -# + #------------# + pid0 = tl.program_id(0) + pid1 = tl.program_id(1) + pidz = tl.program_id(2) + if meta['SDD']: + pid1 = pid1 + SDD_off_width + blockidm = tl.arange(0, TM) // BLOCK + blockidn = tl.arange(0, TN) // BLOCK + offlutm = blockidm * (TN // BLOCK) * 4 + offlutn = blockidn * 4 + header = lut + pid1 * (TM // BLOCK) * (TN // BLOCK) * 4 + z = tl.load(header + 0) + i = tl.load(header + 1 + offlutm) + j = tl.load(header + 2 + offlutn) + AS1 = SDD_K // TZ + lockid = tl.where(TZ > 1, 1, 0) + offka = pid0 * AS1 + offkb = pid0 * AS1 + offmc = 0 + offnc = 0 + offpa = 0 + offpb = 0 + maxid = TZ + offhc = 0 + offha = z + offhb = z + ram = i * BLOCK + (tl.arange(0, TM) % BLOCK) + rbn = j * BLOCK + (tl.arange(0, TN) % BLOCK) + else: + header = lut + pid0 * 6 + offset = tl.load(header + 0) + AS1 = tl.load(header + 1) + column = tl.load(header + 2) + depth = tl.load(header + 3) + lockid = tl.load(header + 4) + maxid = tl.load(header + 5) + pinc = lut + offset + offhc = depth + if meta['DSD']: + # output offset + offnc = pid1 * TN + offmc = column * TM + offpc = 0 + # dense input offset + offnb = pid1 * TN + offkb = tl.load(pinc) + offkb = tl.multiple_of(offkb, 8) # compiler hint + offpb = 0 + # sparse input offset + offma = 0 + offka = 0 + offpa = tl.load(pinc + 1) + offpa = tl.multiple_of(offpa, 8) # compiler hint + offpa = offpa * BLOCK * BLOCK + offha = 0 + offhb = depth + else: + # output offset + offmc = pid1 * TM + offnc = column * TN + offpc = 0 + # dense input offset + offma = pid1 * TM + offka = tl.load(pinc) + offka = tl.multiple_of(offka, 8) # compiler hint + offpa = 0 + # sparse input offset + offnb = 0 + offkb = 0 + offpb = tl.load(pinc + 1) + offpb = tl.multiple_of(offpb, 8) # compiler hint + offpb = offpb * BLOCK * BLOCK + offha = depth + offhb = 0 + ram = offma + tl.arange(0, TM) + rbn = offnb + tl.arange(0, TN) + + # initialize a, b pointers + rka = offka + tl.arange(0, TK) + rkb = offkb + tl.arange(0, TK) + pa = A + pidz * stride_za + offha * stride_ha + offpa + ram[:, None] * stride_ma + rka[None, :] * stride_ka + pb = B + pidz * stride_zb + offhb * stride_hb + offpb + rbn[None, :] * stride_nb + rkb[:, None] * stride_kb + if meta['DDS']: + checkam = ram[:, None] < DS0 + else: + checkam = AS1 > 0 + if meta['DSD']: + checkbn = rbn[None, :] < DS0 + else: + checkbn = AS1 > 0 + a = tl.load(pa, mask=checkam, other=0.) + b = tl.load(pb, mask=checkbn, other=0.) + + ## ---------------- ## + ## Inner Loop ## + ## ---------------- ## + acc = tl.zeros((TM, TN), dtype=tl.float32) + for k in range(AS1, 0, -TK): + acc += tl.dot(a, b) + if meta['SDD']: + inc_a = TK * stride_ka + inc_b = TK * stride_kb + else: + pinc += 2 + if meta['DSD']: + inc_b = tl.load(pinc) + inc_a = tl.load(pinc + 1) + inc_b = tl.multiple_of(inc_b, 8) + inc_a = tl.multiple_of(inc_a, 8) + inc_b = inc_b * stride_kb + if meta['DDS']: + inc_a = tl.load(pinc) + inc_b = tl.load(pinc + 1) + inc_a = tl.multiple_of(inc_a, 8) + inc_b = tl.multiple_of(inc_b, 8) + inc_a = inc_a * stride_ka + pa += inc_a + pb += inc_b + # pre-fetch + checkak = k > TK + checkbk = k > TK + checka = checkam & checkak + checkb = checkbn & checkbk + a = tl.load(pa, mask=checka) + b = tl.load(pb, mask=checkb) + c = acc.to(C.dtype.element_ty) + + if meta['SDD']: + checkc = True + rr_blockidm = tl.arange(0, TM) // BLOCK + rr_blockidn = tl.arange(0, TN) // BLOCK + rr_offlutm = rr_blockidm * (TN // BLOCK) * 4 + rr_offlutn = rr_blockidn * 4 + off_bkid = 3 + rr_offlutm[:, None] + rr_offlutn[None, :] + bkid = tl.load(header + off_bkid) + offpc = bkid * BLOCK * BLOCK + rcm = tl.arange(0, TM) % BLOCK + rcn = tl.arange(0, TN) % BLOCK + else: + rcm = offmc + tl.arange(0, TM) + rcn = offnc + tl.arange(0, TN) + if meta['DSD']: + checkc = rcn[None, :] < DS0 + if meta['DDS']: + checkc = rcm[:, None] < DS0 + + pc = C + offpc + offhc * stride_hc + pidz * stride_zc + rcm[:, None] * stride_mc + rcn[None, :] * stride_nc + # write-back directly + if lockid == 0: + tl.store(pc, c, mask=checkc) + # accumulate partial results using spin-locks + else: + plock = locks + tl.program_id(2) * nlocks * tl.num_programs(1) + tl.program_id(1) * nlocks + lockid - 1 + pcount = plock + tl.num_programs(2) * tl.num_programs(1) * nlocks + while tl.atomic_cas(plock, 0, 1) == 1: + pass + count = tl.load(pcount) + if count == 0: + tl.store(pc, c, mask=checkc) + else: + d = tl.load(pc, mask=checkc) + tl.store(pc, d + c, mask=checkc) + tl.atomic_xchg(pcount, (count + 1) % maxid) + tl.atomic_xchg(plock, 0) + + +############## +# MAIN API # +############## +class _sparse_matmul(torch.autograd.Function): + + sdd_cache = dict() + dsd_cache = dict() + dds_cache = dict() + locks = dict() + + # Given an array sizes representing reduction size for each + # column of a block-mode matrix multiplication, + # performs load-balancing to achieve more smaller reductions + # between `seg_size` elements + @staticmethod + def load_balance(sizes, block): + #global triton + #if triton is None: + # triton = importlib.import_module('triton') + # segment size + # heuristics taken from OpenAI blocksparse code + # https://github.com/openai/blocksparse/blob/master/blocksparse/matmul.py#L95 + max_size = sizes.max() + min_size = sizes[sizes != 0].min() + #if max_size > min_size * 2.0: + # seg_max = max(triton.cdiv(max_size, 4), min_size*2) + #else: + # seg_max = max_size + seg_max = max_size + seg_min = max(triton.cdiv(seg_max, 4), 4) + # split reduction into segments + div = sizes // seg_max + rem = sizes % seg_max + packs = div + (sizes < seg_min).long() + (rem >= seg_min).long() + width = packs.sum() + segments = torch.empty(width, dtype=sizes.dtype) + column = torch.empty_like(segments) + lockid = torch.zeros_like(segments) + maxid = torch.zeros_like(segments) + nlocks = 0 + current = 0 + col_idx = 0 + for i in range(len(sizes)): + d, r = div[i], rem[i] + isempty = sizes[i] < seg_min + last = current + d + (r >= seg_min) + isempty + # column id + column[current:last] = col_idx + # lock id + if d > 1 or (d == 1 and r >= seg_min): + nlocks += 1 + lockid[current:last] = nlocks + maxid[current:last] = last - current + # segment size + segments[current:current + d] = seg_max + if r < seg_min and not isempty: + segments[current + d - 1] += r + if r >= seg_min or isempty: + segments[current + d] = r + current = last + col_idx += 1 + offsets = torch.zeros_like(segments) + offsets[1:] = torch.cumsum(segments[:-1], dim=0) + return segments, column, lockid, maxid, offsets + + @staticmethod + def get_locks(size, dev): + if dev not in _sparse_matmul.locks or \ + size > _sparse_matmul.locks[dev].size(0): + _sparse_matmul.locks[dev] = torch.zeros(size, dtype=torch.int32, device=dev) + return _sparse_matmul.locks[dev] + + ########################## + # SPARSE = DENSE x DENSE # + ########################## + + @staticmethod + def make_sdd_lut(layout, block, dtype, device): + #_sparse_matmul._load_utils() + #start_width = 64 // block + #segmented = _sparse_matmul.sdd_segment(layout.type(torch.int32), start_width) + start_width = (128 if block > 16 else 32) // block + layout = layout.type(torch.int32) + segmented = libtriton.superblock(layout.data_ptr(), layout.shape[0], layout.shape[1], layout.shape[2], + start_width) + luts, widths, packs = [], [], [] + for size, nnz in segmented: + """ width = nnz.shape[0] // (size * size) + h = nnz[:, 0] + i = nnz[:, 1] + j = nnz[:, 2] + b = nnz[:, 3] + lut = torch.stack((h, i, j, b), dim=1).view(-1).contiguous() + luts.append(lut.type(torch.int32).to(device)) + widths.append(width) + packs.append(size) """ + nnz = nnz.reshape(-1, 4) + width = nnz.shape[0] // (size * size) + luts.append(torch.from_numpy(nnz).type(torch.int32).to(device)) + widths.append(width) + packs.append(size) + # create locks + return luts, None, widths, packs + + @staticmethod + def _sdd_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, luts, num_locks, widths, packs, bench, time): + if trans_c: + a, b = b, a + trans_a, trans_b = not trans_b, not trans_a + AS0 = a.size(0) + # Shape check + a_dim = -2 if trans_a else -1 + b_dim = -1 if trans_b else -2 + a_inner, b_inner = a.shape[a_dim], b.shape[b_dim] + if a_inner != b_inner: + raise ValueError(f"Size of tensor A along the {a_dim} dim ({a_inner}) must match size " + f"of tensor B along the {b_dim} dim ({b_inner})") + if a_inner % 16 != 0: + raise ValueError('Reduction size for SDD must be a multiple of 16') + + batch_size = a.size(0) + a_outer = a.size(3 if trans_a else 2) + dtype = a.dtype + is_16_multiple = a_inner % 16 == 0 + is_32_multiple = a_inner % 32 == 0 + is_64_multiple = a_inner % 64 == 0 + if not is_16_multiple: + raise ValueError('Reduction size for SDD must be a multiple of 16') + device = a.device + # create kernel + total_width = sum([width * pack * pack for width, pack in zip(widths, packs)]) + c = torch.empty((batch_size, total_width, block, block), dtype=dtype, device=a.device) + for lut, width, pack in zip(luts, widths, packs): + F32TK = [8, 16] + F16TK = [16] + F16TK += [32] if is_32_multiple else [] + F16TK += [64] if is_64_multiple else [] + TK = {torch.float32: F32TK, torch.float16: F16TK}[dtype] + num_lock = 1 + meta = { + 'TM': block * pack, + 'TN': block * pack, + 'BLOCK': block, + 'TK': TK[0], + 'TZ': 1, + 'SDD': True, + 'DSD': False, + 'DDS': False + } + # create output + locks = _sparse_matmul.get_locks(2 * width * AS0 * num_lock, a.device) + # maximum grid size is 65535 + # so operation might be decomposed into multiple + # kernel calls + max_width = 49152 + total = 0 if bench else None + for off_width in range(0, width, max_width): + grid = lambda meta: [meta['TZ'], min(max_width, width - off_width), batch_size] + _kernel[grid](a, + b, + c, + a.stride(0), + a.stride(1), + a.stride(3 if trans_a else 2), + a.stride(2 if trans_a else 3), + b.stride(0), + b.stride(1), + b.stride(3 if trans_b else 2), + b.stride(2 if trans_b else 3), + c.stride(0), + c.stride(0), + c.stride(2), + c.stride(3), + a_outer, + a_outer, + a_inner, + off_width, + lut, + locks, + num_lock, + num_warps=4, + **meta) + # save for backward pass + return c + + ########################## + # DENSE = DENSE x SPARSE # + ########################## + + # Given a binary layout of 0s and 1s, + # Construct look-up table for efficient execution on GPUs + @staticmethod + def make_dxx_lut(layout, block, step, trans, device, transform=lambda idx: idx): + # load-balancing + _empty = torch.tensor([], dtype=torch.int64, device=layout.device) + segments = _empty.clone() + column = _empty.clone() + depth = _empty.clone() + lockid = _empty.clone() + maxid = _empty.clone() + offsets = _empty.clone() + current_offset = 0 + current_maxid = 0 + for z in range(layout.size(0)): + if trans: + sizes = torch.sum(layout[z, :, :], 1) + else: + sizes = torch.sum(layout[z, :, :], 0) + z_segments, z_column, z_lockid, z_maxid, z_offsets = _sparse_matmul.load_balance(sizes, block) + z_depth = z * torch.ones_like(z_segments) + z_lockid[z_lockid > 0] += current_maxid + current_maxid = z_lockid.max() + # concatenate depth + segments = torch.cat((segments, z_segments)) + column = torch.cat((column, z_column)) + depth = torch.cat((depth, z_depth)) + maxid = torch.cat((maxid, z_maxid)) + offsets = torch.cat((offsets, current_offset + z_offsets)) + lockid = torch.cat((lockid, z_lockid)) + current_offset += layout[z, :, :].sum() + segments *= step + # pointer increments + if trans: + nnz = layout.nonzero() + else: + nnz = layout.transpose(1, 2).nonzero() + num_blocks = nnz.size(0) + offsets = torch.min(offsets, (num_blocks - 1) * torch.ones_like(offsets)) + idx = transform(nnz[:, 2] * block) + xincs = idx.clone() + xincs[1:] -= idx[:-1] + # divide block into multiple steps + div = block // step + xincs = xincs.view(-1, 1).repeat(1, div) + xincs[:, 1:] = step + xincs[:, 0] -= (div - 1) * step + # first increment for each reduction is actually the offset + xincs[offsets[segments > 0], 0] = idx[offsets[segments > 0]] + xincs = xincs.view(-1) + # block-mode input increments + if trans: + widx = torch.arange(num_blocks) + else: + widx = _empty.clone() + current_offset = 0 + for z in range(layout.size(0)): + layoutw = layout[z, :, :].clone() + msum = layoutw.sum() + layoutw[layoutw > 0] = 1 + torch.arange(msum) + widx = torch.cat((widx, current_offset + layoutw.T[layoutw.T > 0] - 1)) + current_offset += msum + widx = widx + wincs = widx * block * block + wincs[1:] -= widx[:-1] * block * block + wincs = wincs.view(-1, 1).repeat(1, div) + if trans: + wincs[:, 1:] = step + wincs[:, 0] -= (div - 1) * step + else: + wincs[:, 1:] = step * block + wincs[:, 0] -= (div - 1) * step * block + wincs[offsets[segments > 0], 0] = widx[offsets[segments > 0]] + wincs = wincs.view(-1) + # adjust offset and segment size + offsets *= 2 * div + segments *= div + # create header + width = column.size(0) + offsets += 6 * width + header = torch.stack((offsets, segments, column, depth, lockid, maxid), dim=1).view(-1).contiguous() + incs = torch.stack((xincs, wincs), dim=1).view(-1).contiguous() + incs = torch.cat((incs, torch.zeros(2, device=incs.device, dtype=incs.dtype))) + # create lut + lut = torch.cat((header, incs)) + lut = lut.type(torch.int32).to(device) + # create locks + num_locks = max(1, lockid.max()) + return lut, num_locks, width, None + + @staticmethod + def _dds_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, lut, num_locks, width, packs, bench, time): + global triton + if triton is None: + triton = importlib.import_module('triton') + + # shapes / dtypes + AS0 = a.size(0) + AS1 = a.size(1) + AS2 = a.size(3 if trans_a else 2) + AS3 = a.size(2 if trans_a else 3) + BS0 = spdims[0] + BS1 = block * spdims[2 if trans_b else 1] + BS2 = block * spdims[1 if trans_b else 2] + dtype = a.dtype + # kernel + meta = {'TN': block, 'TM': 128, 'TK': 16, 'BLOCK': block, 'TZ': 1, 'SDD': False, 'DSD': False, 'DDS': True} + # output + CS0 = AS0 + CS1 = AS1 + CS2 = BS2 if trans_c else AS2 + CS3 = AS2 if trans_c else BS2 + locks = _sparse_matmul.get_locks(2 * AS0 * AS2 // 32 * num_locks, a.device) + c = torch.empty((CS0, CS1, CS2, CS3), dtype=dtype, device=a.device) + grid = lambda meta: [width, triton.cdiv(AS2, meta['TM']), AS0] + _kernel[grid](a, + b, + c, + a.stride(0), + a.stride(1), + a.stride(3 if trans_a else 2), + a.stride(2 if trans_a else 3), + b.stride(0), + b.stride(1), + b.stride(3 if trans_b else 2), + b.stride(2 if trans_b else 3), + c.stride(0), + c.stride(1), + c.stride(3 if trans_c else 2), + c.stride(2 if trans_c else 3), + AS2, + BS2, + 0, + 0, + lut, + locks, + num_locks, + num_warps=4, + **meta) + return c + + @staticmethod + def _dsd_matmul(a, b, trans_a, trans_b, trans_c, spdims, block, lut, num_locks, width, packs, bench, time): + global triton + if triton is None: + triton = importlib.import_module('triton') + + # shapes / dtypes + AS0 = spdims[0] + AS1 = block * spdims[2 if trans_a else 1] + AS2 = block * spdims[1 if trans_a else 2] + BS0 = b.size(0) + BS1 = b.size(1) + BS2 = b.size(3 if trans_b else 2) + BS3 = b.size(2 if trans_b else 3) + dtype = a.dtype + # kernel + + meta = {'TM': block, 'TN': 128, 'TK': 16, 'BLOCK': block, 'TZ': 1, 'SDD': False, 'DSD': True, 'DDS': False} + # output + CS0 = BS0 + CS1 = BS1 + CS2 = BS3 if trans_c else AS1 + CS3 = AS1 if trans_c else BS3 + locks = _sparse_matmul.get_locks(2 * BS0 * BS3 // 32 * num_locks, a.device) + c = torch.empty((CS0, CS1, CS2, CS3), dtype=dtype, device=a.device) + grid = lambda meta: [width, triton.cdiv(BS3, meta['TN']), BS0] + _kernel[grid](a, + b, + c, + a.stride(0), + a.stride(1), + a.stride(3 if trans_a else 2), + a.stride(2 if trans_a else 3), + b.stride(0), + b.stride(1), + b.stride(3 if trans_b else 2), + b.stride(2 if trans_b else 3), + c.stride(0), + c.stride(1), + c.stride(2), + c.stride(3), + BS3, + AS1, + 0, + 0, + lut, + locks, + num_locks, + num_warps=4, + **meta) + return c + + fn = {'sdd': _sdd_matmul.__get__(object), 'dsd': _dsd_matmul.__get__(object), 'dds': _dds_matmul.__get__(object)} + + @staticmethod + def forward(ctx, a, b, trans_a, trans_b, trans_c, mode, spdims, block, c_lut, c_num_locks, c_width, c_packs, + c_bench, c_time, da_lut, da_num_locks, da_width, da_packs, da_bench, da_time, db_lut, db_num_locks, + db_width, db_packs, db_bench, db_time): + c = _sparse_matmul.fn[mode](a, b, trans_a, trans_b, trans_c, spdims, block, c_lut, c_num_locks, c_width, + c_packs, c_bench, c_time) + # save for backward + ctx.save_for_backward(a, b) + ctx.da_num_locks = da_num_locks + ctx.da_lut = da_lut + ctx.da_width = da_width + ctx.da_packs = da_packs + ctx.da_bench = da_bench + ctx.da_time = da_time + ctx.db_lut = db_lut + ctx.db_num_locks = db_num_locks + ctx.db_width = db_width + ctx.db_bench = db_bench + ctx.db_packs = db_packs + ctx.db_time = db_time + ctx.mode = mode + ctx.spdims = spdims + ctx.block = block + ctx.trans_a = trans_a + ctx.trans_b = trans_b + return c + + @staticmethod + def backward(ctx, dc): + # saved for backward + a, b = ctx.saved_tensors + mode = ctx.mode + # gradients w.r.t. a + if ctx.needs_input_grad[0]: + mode_da = mode[1] + mode[0] + mode[2] + da = _sparse_matmul.fn[mode_da](dc, b, False, not ctx.trans_b, ctx.trans_a, ctx.spdims, ctx.block, + ctx.da_lut, ctx.da_num_locks, ctx.da_width, ctx.da_packs, ctx.da_bench, + ctx.da_time) + # gradients w.r.t. b + if ctx.needs_input_grad[1]: + mode_db = mode[2] + mode[1] + mode[0] + db = _sparse_matmul.fn[mode_db](a, dc, not ctx.trans_a, False, ctx.trans_b, ctx.spdims, ctx.block, + ctx.db_lut, ctx.db_num_locks, ctx.db_width, ctx.db_packs, ctx.db_bench, + ctx.db_time) + return da, db, None, None, None,\ + None, None, None, None,\ + None, None, None, None, None, None,\ + None, None, None, None, None, None,\ + None, None, None, None, None, None + + +class MatMul: + """Block-Sparse MatMul class; this class handles three types of matrix-multiplication: + - sparse = dense X dense + - dense = sparse X dense + - dense = dense X sparse + + For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + """ + + def make_lut(self, dtype, device): + """Generates the sparsity layout/s used in block-sparse matmul + """ + key = (dtype, device) + if key in self.lut_cache: + return self.lut_cache[key] + # C look-up table + layout, block = self.layout, self.block + step = 16 + if self.mode == 'sdd': + c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) + elif self.mode == 'dsd': + c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_dxx_lut(layout, block, step, not self.trans_a, + device) + elif self.mode == 'dds': + c_lut, c_num_locks, c_width, c_packs = _sparse_matmul.make_dxx_lut(layout, block, step, self.trans_b, + device) + # DA look-up table + if self.mode == 'sdd': + da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_dxx_lut(layout, block, step, True, device) + elif self.mode == 'dsd': + da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) + elif self.mode == 'dds': + da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_dxx_lut(layout, block, step, + not self.trans_b, device) + # DB look-up table + if self.mode == 'sdd': + db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_dxx_lut(layout, block, step, False, device) + elif self.mode == 'dsd': + db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_dxx_lut(layout, block, step, self.trans_a, + device) + elif self.mode == 'dds': + db_lut, db_num_locks, db_width, db_packs = _sparse_matmul.make_sdd_lut(layout, block, dtype, device) + self.lut_cache[key] = (c_lut, c_num_locks, c_width, c_packs,\ + da_lut, da_num_locks, da_width, da_packs,\ + db_lut, db_num_locks, db_width, db_packs) + return self.lut_cache[key] + + def __init__(self, layout, block, mode, trans_a=False, trans_b=False, bench=False): + """Initialize the Block-Sparse MatMul class. + + Arguments: + layout: required: sparsity layout tensor + block: required: an integer determining the block size. + mode: required: a string determining type of matmul; ('sdd') sparse = dense X dense, ('dsd') dense = sparse X dense, ('dds') dense = dense X sparse + trans_a: optional: a boolean determining if multiplication needs to be applied on transpose of input a; default is false + trans_b: optional: a boolean determining if multiplication needs to be applied on transpose of input b; default is false + bench: optional: set if you want to do benchmarking + """ + + if mode not in ['sdd', 'dsd', 'dds']: + raise NotImplementedError('Supported modes are: sdd, dsd, dds') + # look-up table cache + self.lut_cache = dict() + # attributes + self.trans_a = trans_a + self.trans_b = trans_b + self.mode = mode + self.block = block + self.layout = layout + layout_dim = layout.ndim + assert layout_dim in (2, 3), "Layout should be a 2 or 3 dimensional tensor of 0s and 1s" + if not mode == 'sdd': + # Dims to be reduced on the 'inside' of the matmul, either -1 or -2 + trans_dense, trans_sparse, sparse_inner = (trans_b, trans_a, -1) if mode == 'dsd' else (trans_a, trans_b, + -2) + self.dense_inner_dim = -((sparse_inner % 2) + 1) if not trans_dense else sparse_inner + sparse_inner = sparse_inner if not trans_sparse else -((sparse_inner % 2) + 1) + + # Inner dim of the dense input should be equal to the inner dim of the sparse input + self.dense_inner_size = layout.shape[sparse_inner] * block + # Expected shape for sparse inputs + self.sparse_shape = (layout.sum().item(), block, block) + + # Support using the same layout across attention heads etc. + if layout_dim == 2: + layout = layout.unsqueeze(0) + + layout = layout.long() # Above code assumes the layout tensor is an integral type + + self.spdims = layout.shape + # timings + self.bench = bench + self.time_c = None + self.time_da = None + self.time_db = None + + # pad shapes of a tensor to make it + # compatible with kernel calls + @staticmethod + def _pad_shape(x, is_sparse): + max_dim = 3 if is_sparse else 4 + for i in range(max_dim - x.dim()): + x = x.unsqueeze(0) + return x + + def __call__(self, a, b): + """Applies Block-Sparse MatMul. + + For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + + Arguments: + a: required: a dense/block-sparse tensor; first input of mat-mul + b: required: a dense/block-sparse tensor; second input of mat-mul + + Return: + c: a dense/block-sparse tensor result of a X b + """ + + + c_lut, c_num_locks, c_width, c_packs,\ + da_lut, da_num_locks, da_width, da_packs,\ + db_lut, db_num_locks, db_width, db_packs = self.make_lut(a.dtype, a.device) + # timings + time_c = [None] + time_da = [None] + time_db = [None] + + original_dims = max(a.ndim, b.ndim) + a, b = self._validate_inputs(a, b) + + # pad shapes with ones + a = MatMul._pad_shape(a, self.mode == 'dsd') + b = MatMul._pad_shape(b, self.mode == 'dds') + # execute + + c = _sparse_matmul.apply(a, b, self.trans_a, self.trans_b, False, self.mode, self.spdims, self.block, c_lut, + c_num_locks, c_width, c_packs, self.bench, time_c, da_lut, da_num_locks, da_width, + da_packs, self.bench, time_da, db_lut, db_num_locks, db_width, db_packs, self.bench, + time_db) + + # This removes any leading singleton dimensions we may have added to the tensor that weren't in the input + dims_to_trim = c.ndim - original_dims + for _ in range(dims_to_trim): + c = c.squeeze(0) + + self.time_c = time_c[0] + self.time_da = time_da[0] + self.time_db = time_db[0] + return c + + def _validate_inputs(self, a, b): + if a.device != b.device: + raise ValueError(f"Inputs must be on the same device; got {a.device} for tensor A " + f"and {b.device} for tensor B") + if not get_accelerator().on_accelerator(a): + raise ValueError("Only GPU devices are supported for now") + + # When autocast is enabled, torch.matmul autocasts to float16, so we do the same here + if torch.is_autocast_enabled(): + a, b = a.half(), b.half() + elif a.dtype != b.dtype: + raise ValueError(f"Inputs must be the same dtype; got {a.dtype} for A and {b.dtype} for B") + + mode, trans_a, trans_b = self.mode, self.trans_a, self.trans_b + if mode != 'sdd': + # One input is sparse + dense, dense_name, sparse, sparse_name = (a, 'A', b, 'B') if mode == 'dds' else (b, 'B', a, 'A') + dense_inner = dense.shape[self.dense_inner_dim] + if dense_inner != self.dense_inner_size: + raise ValueError(f"Expected tensor {dense_name} to have size {self.dense_inner_size} at dim " + f"{self.dense_inner_dim % dense.ndim}, got {dense_inner}.") + + if sparse.shape[-len(self.sparse_shape):] != self.sparse_shape: + raise ValueError(f"Expected tensor with trailing dimensions of shape {self.sparse_shape} for argument " + f"{sparse_name}, got {sparse.shape}") + + def add_extra_dims(x): + # Add extra leading singleton dimensions if needed + dims_needed = 4 - x.ndim + if dims_needed > 0: + singletons = [1] * dims_needed + x = x.view(*singletons, *x.shape) + elif dims_needed < 0: + raise ValueError("Tensors with more than 4 dimensions are not currently supported") + + return x + + # Pad shapes with leading singleton dimensions + a = add_extra_dims(a) + b = add_extra_dims(b) + + return a, b diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/softmax.py b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/softmax.py new file mode 100644 index 0000000000000000000000000000000000000000..debee5688fe3f06699a710dbaa35c09601189a4f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/softmax.py @@ -0,0 +1,296 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a +# https://github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py + +import torch + +import triton +import triton.language as tl + + +def next_power_of_2(n): + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n += 1 + return n + + +def num_warps(n): + if n < 512: + return 4 + if n < 2048: + return 8 + return 16 + + +@triton.heuristics({'num_warps': lambda *args, **meta: num_warps(args[6] * meta['BLOCK'])}) +@triton.heuristics({'TN': lambda *args, **meta: next_power_of_2(args[6] * meta['BLOCK'])}) +@triton.jit +def _forward(X, scale, LUT, RPE, KP_M, ATTN_M, sizemax, stride_zx, stride_zrpe, stride_hrpe, stride_srpe, stride_zkpm, + stride_zattnm, **meta): + TN = meta['TN'] + BLOCK = meta['BLOCK'] + pidhm = tl.program_id(0) + pidz = tl.program_id(1) + # create index ranges + rxm = pidhm % BLOCK + rbm = pidhm // BLOCK + rxn = tl.arange(0, TN) % BLOCK + rbn = tl.arange(0, TN) // BLOCK + # extract information from LUT + header = LUT + rbm * 2 + size = tl.load(header + 0) + offset = tl.load(header + 1) + check = rbn < size + rbmn = tl.where(check, rbn, size - 1) + # block id and column id + blockid = tl.load(LUT + offset + rbmn * 4 + 0) + columnid = tl.load(LUT + offset + rbmn * 4 + 1) + rowid = tl.load(LUT + offset + rbmn * 4 + 2) + headid = tl.load(LUT + offset + rbmn * 4 + 3) + # pointers to X + px = X + pidz * stride_zx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn + x = tl.load(px, mask=check, other=-float('inf')) + x = x.to(tl.float32) + # apply scale + if meta['APPLY_SCALE']: + x = x * scale + # apply RPE + if meta['APPLY_RPE']: + prpe = RPE + pidz * stride_zrpe + headid * stride_hrpe + columnid * BLOCK + rowid * BLOCK * stride_srpe + rxm * stride_srpe + rxn + rpe = tl.load(prpe, mask=check, other=0) + x = x + rpe + # apply key-padding mask + if meta['APPLY_KP_MASK']: + pkp_m = KP_M + pidz * stride_zkpm + columnid * BLOCK + rxn + kp_m = tl.load(pkp_m, mask=check, other=-float('inf')) + if meta['KP_MASK_MUL']: + kp_m = tl.where(kp_m == 0, -float('inf'), 0.) + x = x + kp_m + # apply attention mask + if meta['APPLY_ATTN_MASK']: + pattn_m = ATTN_M + columnid * BLOCK + rowid * BLOCK * stride_zattnm + rxm * stride_zattnm + rxn + attn_m = tl.load(pattn_m, mask=check, other=-float('inf')) + if meta['ATTN_MASK_MUL']: + attn_m = tl.where(attn_m == 0, -float('inf'), 0.) + x = x + attn_m + # computation + x = tl.softmax(x) + tl.store(px, x, mask=check) + + +@triton.heuristics({'num_warps': lambda *args, **meta: num_warps(args[4] * meta['BLOCK'])}) +@triton.heuristics({'TN': lambda *args, **meta: next_power_of_2(args[4]) * meta['BLOCK']}) +@triton.jit +def _backward(X, scale, DX, LUT, sizemax, stride_zx, stride_zdx, **meta): + pidhm = tl.program_id(0) + pidz = tl.program_id(1) + TN = meta['TN'] + BLOCK = meta['BLOCK'] + # create index ranges + rxm = pidhm % BLOCK + rbm = pidhm // BLOCK + rxn = tl.arange(0, TN) % BLOCK + rbn = tl.arange(0, TN) // BLOCK + # extract information from look-up table + header = LUT + rbm * 2 + size = tl.load(header + 0) + offset = tl.load(header + 1) + # bounds checking on lut + check = rbn < size + rbmn = tl.where(check, rbn, size - 1) + # initialize pointers to block-sparse input + blockid = tl.load(LUT + offset + rbmn * 4) + X = X + pidz * stride_zx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn + DX = DX + pidz * stride_zdx + blockid * BLOCK * BLOCK + rxm * BLOCK + rxn + # compute fused softmax backward + x = tl.load(X, mask=check, other=0) + dx = tl.load(DX, mask=check, other=0) + x = x.to(tl.float32) + dx = dx.to(tl.float32) + y = x * (dx - tl.sum(x * dx, 0)) * scale + tl.store(DX, y, mask=check) + + +class _sparse_softmax(torch.autograd.Function): + + bwd_kernels = dict() + + @staticmethod + def make_lut(layout, block, device): + _empty = torch.tensor([], dtype=torch.int64, device=layout.device) + sizes = _empty.clone() + # sizes along rows + for h in range(layout.shape[0]): + sizes = torch.cat((sizes, layout[h, :, :].sum(-1))) + # offsets in block format + offsets = torch.zeros_like(sizes) + offsets[1:] = torch.cumsum(sizes[:-1], dim=0) + # block indices + idx = torch.arange(layout.sum()) + head = layout.nonzero()[:, 0] + rows = layout.nonzero()[:, 1] + columns = layout.nonzero()[:, 2] + core = torch.stack((idx, columns, rows, head), dim=1).view(-1) + # construct look-up table + offsets = offsets * 4 + 2 * sizes.numel() + header = torch.stack((sizes, offsets), dim=1).view(-1) + lut = torch.cat((header, core)).type(torch.int32).to(device) + return lut, int(sizes.max()) + + @staticmethod + def forward(ctx, x, scale, rpe, key_padding_mask, attn_mask, kp_mask_mode, attn_mask_mode, spdims, block, lut, + num_blocks, maxlut, bench, time): + + apply_scale = False if scale == 1.0 else True + + # handle None rpe + if rpe is None: + apply_rpe = False + stride_zrpe, stride_hrpe, stride_srpe = 0, 0, 0 + rpe = torch.empty(0, dtype=x.dtype, device=x.device) + else: + apply_rpe = True + stride_zrpe, stride_hrpe, stride_srpe = rpe.stride(0), rpe.stride(1), rpe.stride(2) + + # handle None key_padding_mask + if key_padding_mask is None: + apply_kp_mask = False + stride_zkpm = 0 + key_padding_mask = torch.empty(0, dtype=x.dtype, device=x.device) + else: + apply_kp_mask = True + stride_zkpm = key_padding_mask.stride(0) + + # handle None attention_mask + if attn_mask is None: + apply_attn_mask = False + stride_zattnm = 0 + attn_mask = torch.empty(0, dtype=x.dtype, device=x.device) + else: + apply_attn_mask = True + stride_zattnm = attn_mask.stride(0) + + # run kernel + M = x.shape[0] + meta = { + 'BLOCK': block, + 'APPLY_SCALE': apply_scale, + 'APPLY_RPE': apply_rpe, + 'APPLY_KP_MASK': apply_kp_mask, + 'APPLY_ATTN_MASK': apply_attn_mask, + 'KP_MASK_MUL': kp_mask_mode == 'mul', + 'ATTN_MASK_MUL': attn_mask_mode == 'mul', + } + grid = lambda opt: [spdims[0] * spdims[1] * block, M] + _forward[grid](x, scale, lut, rpe, key_padding_mask, attn_mask, maxlut, x.stride(0),\ + stride_zrpe, stride_hrpe, stride_srpe, stride_zkpm, stride_zattnm, **meta) + + # save to context + ctx.mark_dirty(x) + ctx.save_for_backward(x, lut) + ctx.spdims = spdims + ctx.block = block + ctx.maxlut = maxlut + ctx.scale = scale + ctx.apply_scale = apply_scale + ctx.apply_rpe = apply_rpe + ctx.apply_kp_mask = apply_kp_mask + ctx.apply_attn_mask = apply_attn_mask + ctx.kp_mask_mode = kp_mask_mode + ctx.attn_mask_mode = attn_mask_mode + return x + + @staticmethod + def backward(ctx, dx): + + # retrieve from context + x, lut = ctx.saved_tensors + # run kernel + M = x.shape[0] + grid = lambda opt: [ctx.spdims[0] * ctx.spdims[1] * ctx.block, M] + _backward[grid](x, ctx.scale, dx, lut, ctx.maxlut, x.stride(0), dx.stride(0), BLOCK=ctx.block) + return dx, None, None, None, None, None, None, None, None, None, None, None, None, None, None + + +class Softmax: + """Block-Sparse Softmax class; this class computes softmax on a block sparse matrix. It is also able to apply either/all of the following masks: + - relative position embedding + - key padding mask + - attention mask + + For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + """ + + def sparse_softmax(*args, **kwargs): + return _sparse_softmax.apply(*args, **kwargs) + + def make_lut(self, device): + """Generates the sparsity layout used in block-sparse softmax + """ + key = (device, ) + if key not in self.lut_cache: + self.lut_cache[key] = _sparse_softmax.make_lut(self.layout, self.block, device) + return self.lut_cache[key] + + def __init__(self, layout, block, bench=False): + """Initialize the Block-Sparse Softmax class. + + Arguments: + layout: required: sparsity layout tensor + block: required: an integer determining the block size. + bench: optional: set if you want to do benchmarking + """ + + self.num_blocks = layout.sum().item() + self.spdims = layout.shape + self.layout = layout + self.block = block + self.bench = bench + self.lut_cache = dict() + + def __call__(self, + x, + scale=1., + rpe=None, + key_padding_mask=None, + attn_mask=None, + key_padding_mask_mode='add', + attn_mask_mode='add'): + """Applies softmax on a Block-Sparse input tensor. + + For more details about sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + + Arguments: + x: required: a block-sparse tensor that softmax is applied on it; computation will be in place and result will be returned in the same tensor + scale: optional: a float value; x values will be multiplied by this value before normalization. Default value is 1.0. + rpe: optional: a tensor same dimension as x that is used as relative position embedding + key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength) + attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported + key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied + attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied + + Return: + x: a block-sparse tensor contains normalized input x using softmax; and masks applied if given + """ + + time_y = [None] + if rpe is not None and rpe.dtype != x.dtype: + raise ValueError('relative position embedding must be %s' % x.dtype) + if attn_mask is not None and attn_mask.dtype != x.dtype: + raise ValueError('Attention mask must be %s' % x.dtype) + if key_padding_mask is not None and key_padding_mask.dtype != x.dtype: + raise ValueError('Key padding mask must be %s' % x.dtype) + lut, maxlut = self.make_lut(x.device) + x = Softmax.sparse_softmax(x, scale, rpe, key_padding_mask, attn_mask, key_padding_mask_mode, attn_mask_mode, + self.spdims, self.block, lut, self.num_blocks, maxlut, self.bench, time_y) + self.time_y = time_y[0] + return x diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ccb0f940dff65839beac579f81c4dfb7e499e6bb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparse_attention_utils.py @@ -0,0 +1,208 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from torch.nn import functional as F +from deepspeed.ops.sparse_attention import BertSparseSelfAttention, SparsityConfig +''' +This file contains few utility functions to handle adapting pretrained model with sparse self-attention module. +''' + + +class SparseAttentionUtils: + """This class provides some utility functions that are use integrating sparse attention into transformer models. + Such utilities include extending position embeddings, replacing current self-attention layer with sparse attention, padding sequences to multiple of block size, etc. + + """ + + @staticmethod + def extend_position_embedding(model, max_position): + """This function extends the position embedding weights of a model loaded from a checkpoint. + It assumes the new max position is bigger than the original max length. + + Arguments: + model: required: a transformer model + max_position: required: an integer determining new position embedding size + Return: + model: updated model; in which position embedding weights have been extended based on new size + """ + + if hasattr(model, 'bert'): + original_max_position = model.bert.embeddings.position_embeddings.weight.size(0) + assert max_position > original_max_position + extend_multiples = max(1, max_position // original_max_position) + model.bert.embeddings.position_embeddings.weight.data = model.bert.embeddings.position_embeddings.weight.repeat( + extend_multiples, 1) + elif hasattr(model, 'roberta'): + # RoBERTa has positions 0 & 1 reserved, so embedding size is max position + 2 + original_max_position, embed_size = model.roberta.embeddings.position_embeddings.weight.shape + original_max_position -= 2 + extend_multiples = max(1, max_position // original_max_position) + assert max_position > original_max_position + max_position += 2 + extended_position_embedding = model.roberta.embeddings.position_embeddings.weight.new_empty( + max_position, embed_size) + k = 2 + for i in range(extend_multiples): + extended_position_embedding[k:( + k + original_max_position)] = model.roberta.embeddings.position_embeddings.weight[2:] + k += original_max_position + model.roberta.embeddings.position_embeddings.weight.data = extended_position_embedding + else: + raise ValueError( + 'Please extend \"extend_position_embedding\" function to support your model type. It currently only supports \"bert\" & \"roberta\"!' + ) + + model.config.max_position_embeddings = max_position + print(f'Extended position embeddings to {original_max_position * extend_multiples}') + + return model + + @staticmethod + def update_tokenizer_model_max_length(tokenizer, max_position): + """This function updates the position embedding length of a tokenizer to a new max position. + + Arguments: + tokenizer: required: a transformer tokenizer + max_position: required: an integer determining new position embedding size + Return: + tokenizer: updated tokenizer; in which model maximum length has been extended based on new size + """ + + tokenizer.model_max_length = max_position + tokenizer.init_kwargs['model_max_length'] = max_position + print(f'updated tokenizer model max imum length to {max_position}') + + return tokenizer + + @staticmethod + def replace_model_self_attention_with_sparse_self_attention( + model, + max_position, + # SparsityConfig parameters needs to be set accordingly + sparsity_config=SparsityConfig(num_heads=4)): + """This function replaces the self attention layers in model encoder with sparse self attention. + It currently supports bert and roberta model and can be easily extended to any other models following similar steps here. + For sparsityConfig, refer to the config class. + + Arguments: + model: required: a transformer model + max_position: required: an integer determining new position embedding size + sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class + + Return: + model: updated model; in which self attention layer has been replaced with DeepSpeed Sparse Self Attention layer. + """ + + if hasattr(model, 'bert'): + model.config.max_position_embeddings = max_position + model.replace_self_attention_layer_with_sparse_self_attention_layer(model.config, model.bert.encoder.layer, + sparsity_config) + elif hasattr(model, 'roberta'): + model.config.max_position_embeddings = max_position + 2 + model.replace_self_attention_layer_with_sparse_self_attention_layer(model.config, + model.roberta.encoder.layer, + sparsity_config) + else: + raise ValueError( + 'Please extend \"update_model_self_attention_to_sparse_self_attention\" function to support \ + your model type. It currently only supports \"bert\" & \"roberta\"!') + return model + + @staticmethod + def replace_self_attention_layer_with_sparse_self_attention_layer( + config, + layers, + # SparsityConfig parameters needs to be set accordingly + sparsity_config=SparsityConfig(num_heads=4)): + """This function replaces the self attention layers in attention layer with sparse self attention. + For sparsityConfig, refer to the config class. + + Arguments: + config: required: transformer model config + layers: required: transformer model attention layers + sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class + + Return: + layers: updated attention layers; in which self attention layers have been replaced with DeepSpeed Sparse Self Attention layer. + """ + + for layer in layers: + deepspeed_sparse_self_attn = BertSparseSelfAttention(config, sparsity_config) + deepspeed_sparse_self_attn.query = layer.attention.self.query + deepspeed_sparse_self_attn.key = layer.attention.self.key + deepspeed_sparse_self_attn.value = layer.attention.self.value + + layer.attention.self = deepspeed_sparse_self_attn + + return layers + + @staticmethod + def pad_to_block_size(block_size, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds, + pad_token_id, model_embeddings): + """This function pads input tokens and attention mask on sequence length dimension to be multiple of block size. + This is a requirement for Sparse Transformer in which the self attention layer works on sequences of length multiple of block size. + It needs to be called in your model, such as BertModel, right before you calculate the embedding outputs. + Note) + 1- instead of passing your embedding layer to this function, you can simply add this function to your model. It can be more simplified if given attention_mask and/or token_type_ids are none. + 2- you need to call unpad function before returning your model output to unpad the encoder sequence output. + + Arguments: + block_size: required: an integer determining the block size of sparsity config. + pad_token_id: required: an integer determining the pad token from the model config; such as bert.config.pad_token_id. + input_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary + attention_mask: a torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. + token_type_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). + position_ids: a torch.LongTensor of shape [batch_size, sequence_length] with the indices of positions of each input sequence tokens in the position embeddings. + inputs_embeds: an optional torch.FloatTensor of shape [batch_size, sequence_length, hidden_size] that contains embedded representation and can be passed instead of input_ids directly. + model_embeddings: an optional object. If inputs_embeds are not none, this will be your model embeddings such as BertEmbeddings from your model such as BertModel. You can move this function inside your model and use self.embeddings instead of passing this parameter. + + Return: + pad_len: an integer determining how much inputs have been padded to transfer sequence length dimension to multiple of block size. + input_ids: if input_ids are not none padded input_ids otherwise none. + attention_mask: if attention_mask is not none padded attention_mask otherwise none. + token_type_ids: if token_type_ids are not none padded token_type_ids otherwise none. + position_ids: if position_ids are not none padded position_ids otherwise none. + inputs_embeds: if inputs_embeds are not none padded inputs_embeds otherwise none. + """ + + batch_size, seq_len = input_ids.shape if input_ids is not None else inputs_embeds.shape[:-1] + + pad_len = (block_size - seq_len % block_size) % block_size + if pad_len > 0: + if inputs_embeds is not None: + pad_input_ids = inputs_embeds.new_full((batch_size, pad_len), pad_token_id, dtype=torch.long) + pad_inputs_embeds = model_embeddings(pad_input_ids) + inputs_embeds = torch.cat([inputs_embeds, pad_inputs_embeds], dim=-2) + # may not be needed as input_ids are not used if inputs_embeds are given + if input_ids is not None: + input_ids = F.pad(input_ids, (0, pad_len), value=pad_token_id) + if position_ids is not None: + # pad position_id with pad_token_id + position_ids = F.pad(position_ids, (0, pad_len), value=pad_token_id) + # pad attention mask without attention on the padding tokens + attention_mask = F.pad(attention_mask, (0, pad_len), value=False) + # pad token_type_ids with token_type_id = 0 + token_type_ids = F.pad(token_type_ids, (0, pad_len), value=0) + + return pad_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds + + @staticmethod + def unpad_sequence_output(pad_len, sequence_output): + """This function unpads sequence output if inputs of the model were padded. + This is a requirement for Sparse Transformer in which the self attention layer works on sequences of length multiple of block size. + It needs to be called in your model, such as BertModel, right before you return the model outputs. + + Arguments: + pad_len: required: an integer determining how much model inputs have been padded to transfer sequence length dimension to multiple of block size. + sequence_output: required: sequence output of the encoder layer. + + Return: + sequence_output: unpaded sequence output of the encoder layer. + """ + + if (pad_len > 0): + sequence_output = sequence_output[:, :-pad_len] + return sequence_output diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..b673c4561902e943981ca3008fae53ec73c0cd73 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py @@ -0,0 +1,149 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch.nn as nn +import torch +from torch import distributed as dist +from deepspeed.ops.sparse_attention import SparsityConfig + + +class SparseSelfAttention(nn.Module): + """Implements an efficient Sparse Self Attention of Transformer layer based on `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509 + + For more information please see, TODO DeepSpeed Sparse Transformer. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial. + """ + + def __init__( + self, + # SparsityConfig parameters needs to be set accordingly + sparsity_config=SparsityConfig(num_heads=4), + key_padding_mask_mode='add', + attn_mask_mode='mul', + max_seq_length=2048): + """Initialize the sparse self attention layer. + Arguments: + sparsity_config: optional: this parameter determines sparsity pattern configuration; it is based on SparsityConfig class. + key_padding_mask_mode: optional: a string determining if key padding mask needs to be added, `add`, or be multiplied, `mul`. + attn_mask_mode: optional: a string determining if attention mask needs to be added, `add`, or be multiplied, `mul`. + max_seq_length: optional: the maximum sequence length this sparse attention module will be applied to; it controls the size of the master_layout. + """ + super().__init__() + + # sparsity information + self.sparsity_config = sparsity_config + + # initialize sparse layout and register as buffer + master_layout = self.sparsity_config.make_layout(max_seq_length) + self.register_buffer("master_layout", master_layout) + self._need_layout_synchronization = True + + # mask modes + self.key_padding_mask_mode = key_padding_mask_mode + self.attn_mask_mode = attn_mask_mode + + ops = dict() + + def get_layout(self, L): + # if layout is never synchronized across GPUs, broadcast the layout from global rank 0 + if self._need_layout_synchronization and dist.is_initialized(): + dist.broadcast(self.master_layout, src=0) + self._need_layout_synchronization = False + + if (L % self.sparsity_config.block != 0): + raise ValueError( + f'Sequence Length, {L}, needs to be dividable by Block size {self.sparsity_config.block}!') + + num_blocks = L // self.sparsity_config.block + return self.master_layout[..., :num_blocks, :num_blocks].cpu() # layout needs to be a CPU tensor + + # add to cache + def get_ops(self, H, L): + from deepspeed.ops.sparse_attention.matmul import MatMul + from deepspeed.ops.sparse_attention.softmax import Softmax + if L not in SparseSelfAttention.ops: + sparsity_layout = self.get_layout(L) + sparse_dot_sdd_nt = MatMul(sparsity_layout, self.sparsity_config.block, 'sdd', trans_a=False, trans_b=True) + + sparse_dot_dsd_nn = MatMul(sparsity_layout, + self.sparsity_config.block, + 'dsd', + trans_a=False, + trans_b=False) + + sparse_softmax = Softmax(sparsity_layout, self.sparsity_config.block) + + SparseSelfAttention.ops[L] = (sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax) + return SparseSelfAttention.ops[L] + + def transpose_key_for_scores(self, x, L): + bsz, num_heads, seq_len, head_dim = x.size() + if seq_len != L: + return x.permute(0, 1, 3, 2) + return x + + def transpose_mask_for_sparse(self, qtype, x, is_key_padding_mask=False): + x = x.type(qtype) + if is_key_padding_mask: + xdim = x.dim() + for d in range(xdim - 1, 0, -1): + x = x.squeeze(dim=d) + return x + return x.squeeze() + + # forward pass + def forward(self, query, key, value, rpe=None, key_padding_mask=None, attn_mask=None): + """Applies forward phase of sparse self attention + + Arguments: + query: required: query tensor + key: required: key tensor + value: required: value tensor + rpe: optional: a tensor same dimension as x that is used as relative position embedding + key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength) + attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported + key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied + attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied + + Return: + attn_output: a dense tensor containing attention context + """ + assert query.dtype == torch.half, "sparse attention only supports training in fp16 currently, please file a github issue if you need fp32 support" + bsz, num_heads, tgt_len, head_dim = query.size() + + # transpose back key if it is already transposed + key = self.transpose_key_for_scores(key, tgt_len) + + # check that operation is supported + if query.shape != key.shape or key.shape != value.shape: + raise NotImplementedError('only self-attention is supported for now') + + # squeeze key_padding_mask if it is given + if key_padding_mask is not None: + key_padding_mask = self.transpose_mask_for_sparse(query.dtype, key_padding_mask, is_key_padding_mask=True) + + # squeeze attn_mask if it is given + if attn_mask is not None: + attn_mask = self.transpose_mask_for_sparse(query.dtype, attn_mask) + + # cache look-up table computations etc + sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax = self.get_ops(num_heads, tgt_len) + + scaling = float(head_dim)**-0.5 + + # attention scores + attn_output_weights = sparse_dot_sdd_nt(query, key) + attn_output_weights = sparse_softmax(attn_output_weights, + scale=scaling, + rpe=rpe, + key_padding_mask=key_padding_mask, + attn_mask=attn_mask, + key_padding_mask_mode=self.key_padding_mask_mode, + attn_mask_mode=self.attn_mask_mode) + + # outputs + attn_output = sparse_dot_dsd_nn(attn_output_weights, value) + return attn_output diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py new file mode 100644 index 0000000000000000000000000000000000000000..1f59c4b469282ebfdb26fdb7c40c61d738231030 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/sparsity_config.py @@ -0,0 +1,727 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import random + + +class SparsityConfig: + """Abstract Configuration class to store `sparsity configuration of a self attention layer`. + It contains shared property of different block-sparse sparsity patterns. However, each class needs to extend it based on required property and functionality. + """ + + def __init__(self, num_heads, block=16, different_layout_per_head=False): + """Initialize the Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. + """ + + self.num_heads = num_heads + self.block = block + self.different_layout_per_head = different_layout_per_head + self.num_layout_heads = num_heads if different_layout_per_head else 1 + + def setup_layout(self, seq_len): + """Create layout tensor for the given sequence length + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) for sparsity layout of all head; initialized with zero + """ + + if (seq_len % self.block != 0): + raise ValueError(f'Sequence Length, {seq_len}, needs to be dividable by Block size {self.block}!') + num_blocks = seq_len // self.block + # TODO Currently we allocate layout per head; needs to be updated if heads share a single layout. + layout = torch.zeros((self.num_heads, num_blocks, num_blocks), dtype=torch.int64) + return layout + + def check_and_propagate_first_head_layout(self, layout): + """If all heads require same sparsity layout, it propagate first head layout to all heads + + Arguments: + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head + """ + + if not self.different_layout_per_head: + layout[1:self.num_heads, :, :] = layout[0, :, :] + return layout + + +class DenseSparsityConfig(SparsityConfig): + """Configuration class to store `Dense` configuration. + In reality, this is not sparse and all blocks are used. We keep it for the sake of comparison and comprehension. + """ + + def __init__(self, num_heads, block=16, different_layout_per_head=False): + """Initialize the Dense Sparsity Pattern Config. + In reality, this is not sparse and all blocks are used. We keep it for the sake of comparison and comprehension. + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + seq_len: required: an integer determining number of attention heads of the layer. + different_layout_per_head: optional: this is just for the sake of consistency with other sparsity formats; can ignore it for DenseSparsityConfig + """ + + super().__init__(num_heads, block, different_layout_per_head) + + def make_layout(self, seq_len): + """Set 1 to all blocks of the layout meaning the pattern is dense; not sparse. + + Arguments: + seq_len: required: an integer determining the underling sequence length; must be <= max sequence length + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; for dense everything is 1 + """ + + layout = self.setup_layout(seq_len) + layout[:, :, :] = 1 + return layout + + +class FixedSparsityConfig(SparsityConfig): + """Configuration class to store `Fixed` sparsity configuration. + For more details about this sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509; this has been customized. + This class extends parent class of `SparsityConfig` and customizes it for `Fixed` sparsity. + """ + + def __init__(self, + num_heads, + block=16, + different_layout_per_head=False, + num_local_blocks=4, + num_global_blocks=1, + attention='bidirectional', + horizontal_global_attention=False, + num_different_global_patterns=1): + """Initialize `Fixed` Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. + num_local_blocks: optional: an integer determining the number of blocks in local attention window. + num_global_blocks: optional: an integer determining how many consecutive blocks in a local window is used as the representative of the window for global attention. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + horizontal_global_attention: optional: a boolean determining if blocks that are global representative of a local window, also attend to all other blocks. This is valid only if attention type is `bidirectional`. Looking at the attention matrix, that means global attention not only includes the vertical blocks, but also horizontal blocks. + num_different_global_patterns: optional: an integer determining number of different global attentions layouts. While global attention can be fixed by which block/s are representative of any local window, since there are multi-heads, each head can use a different global representative. For example, with 4 blocks local window and global attention size of 1 block, we can have 4 different versions in which the first, Second, third, or forth block of each local window can be global representative of that window. This parameter determines how many of such patterns we want. Of course, there is a limitation based on num_local_blocks and num_global_blocks. + """ + + super().__init__(num_heads, block, different_layout_per_head) + + self.num_local_blocks = num_local_blocks + + if (num_local_blocks % num_global_blocks != 0): + raise ValueError( + f'Number of blocks in a local window, {num_local_blocks}, must be dividable by number of global blocks, {num_global_blocks}!' + ) + self.num_global_blocks = num_global_blocks + + if (attention != 'unidirectional' and attention != 'bidirectional'): + raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') + self.attention = attention + + if (attention != 'bidirectional' and horizontal_global_attention): + raise ValueError('only \"bi-directional\" attentions can support horizontal global attention!') + self.horizontal_global_attention = horizontal_global_attention + + if (num_different_global_patterns > 1 and not different_layout_per_head): + raise ValueError( + f'Number of different layouts cannot be more than one when you have set a single layout for all heads! Set different_layout_per_head to True.' + ) + if (num_different_global_patterns > (num_local_blocks // num_global_blocks)): + raise ValueError( + f'Number of layout versions (num_different_global_patterns), {num_different_global_patterns}, cannot be larger than number of local window blocks divided by number of global blocks, {num_local_blocks} / {num_global_blocks} = {num_local_blocks//num_global_blocks}!' + ) + self.num_different_global_patterns = num_different_global_patterns + + def set_local_layout(self, h, layout): + """Sets local attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local layout is set + """ + + num_blocks = layout.shape[1] + for i in range(0, num_blocks, self.num_local_blocks): + end = min(i + self.num_local_blocks, num_blocks) + for row in range(i, end): + for col in range(i, (row + 1 if self.attention == 'unidirectional' else end)): + layout[h, row, col] = 1 + return layout + + def set_global_layout(self, h, layout): + """Sets global attention layout used by the given head in the sparse attention. + + Currently we set global blocks starting from the last block of a local window to the first one. That means if a local window consists of 4 blocks and global attention size is one block, we use block #4 in each local window as global. If we have different layout per head, then other heads will get #3, #2, and #1. And if we have more heads (and different layout has set) than num of global attentions, multiple head may have same global attentions. + Note) if horizontal_global_attention is set, global blocks will be set both horizontally and vertically. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set + """ + + num_blocks = layout.shape[1] + first_global_block_idx = self.num_local_blocks - ( + 1 + h % self.num_different_global_patterns) * self.num_global_blocks + + # set all global blocks except the last one if (in last local window) + end = num_blocks - (num_blocks % self.num_local_blocks) + for i in range(first_global_block_idx, end, self.num_local_blocks): + + # vertical global attention + first_row = 0 if self.attention == 'bidirectional' else i + #(((i // self.num_local_blocks) + 1) * self.num_local_blocks) + #if (first_row < num_blocks): + layout[h, first_row:, i:i + self.num_global_blocks] = 1 + + # horizontal global attention; only in bidirectional attention + if (self.horizontal_global_attention): + layout[h, i:i + self.num_global_blocks, :] = 1 + + # set last global blocks; handle possible short last local window + if (end < num_blocks): + start = min(end + first_global_block_idx, num_blocks - self.num_global_blocks) + end = start + self.num_global_blocks + + # vertical global attention + first_row = 0 if self.attention == 'bidirectional' else start + #(((start // self.num_local_blocks) + 1) * self.num_local_blocks) + #if (first_row < num_blocks): + layout[h, first_row:, start:end] = 1 + + # horizontal global attention + if (self.horizontal_global_attention): + layout[h, start:end, :] = 1 + return layout + + def make_layout(self, seq_len): + """Generates `Fixed` sparsity layout used by each head in the sparse attention. + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `Fixed` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_local_layout(h, layout) + layout = self.set_global_layout(h, layout) + + layout = self.check_and_propagate_first_head_layout(layout) + return layout + + +class VariableSparsityConfig(SparsityConfig): + """Configuration class to store `Variable` sparsity configuration. + This layout is an extension of FixedSparsityConfig in which: + - user can set random layout; default value is zero means no random block + - user can provide a list of local block sizes + - user can provide a list of global block indices. + + For more details about `Fixed` sparsity config, please see `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509; this has been customized. + This class extends parent class of `SparsityConfig` and customizes it for `Fixed` sparsity. + """ + + def __init__(self, + num_heads, + block=16, + different_layout_per_head=False, + num_random_blocks=0, + local_window_blocks=[4], + global_block_indices=[0], + global_block_end_indices=None, + attention='bidirectional', + horizontal_global_attention=False): + """Initialize `Variable` Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. Currently this sparsity config can only assign single layout to all heads; needs to be extended for different layout per head. + num_random_blocks: optional: an integer determining the number of random blocks in each block row. + local_window_blocks: optional: a list of integers determining the number of blocks in each local attention window. It assumes first number determines # of blocks in the first local window, second the second window, ..., and the last number determines the number of blocks in the remaining local windows. + global_block_indices: optional: a list of integers determining which blocks are considered as global attention. Given indices, determine the blocks that all other token blocks attend to and they attend to all other token blocks. Default value is only index 0. Notice that if global_block_end_indices parameter is set, this parameter is used as starting index of each global window. + global_block_end_indices: optional: a list of integers determining end indices of global window blocks. By default this is not used. But if it is set, it must have the same size of global_block_indices parameter, and combining this two parameters, for each index i, blocks from global_block_indices[i] to global_block_end_indices[i] (exclusive) are considered as global attention. + num_global_blocks: optional: an integer determining how many consecutive blocks in a local window is used as the representative of the window for global attention. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + horizontal_global_attention: optional: a boolean determining if blocks that are global representative of a local window, also attend to all other blocks. This is valid only if attention type is `bidirectional`. Looking at the attention matrix, that means global attention not only includes the vertical blocks, but also horizontal blocks. + """ + + super().__init__(num_heads, block, different_layout_per_head) + + self.num_random_blocks = num_random_blocks + self.local_window_blocks = local_window_blocks + self.global_block_indices = global_block_indices + + if (global_block_end_indices is not None): + if (len(global_block_indices) != len(global_block_end_indices)): + raise ValueError( + f'Global block start indices length, {len(global_block_indices)}, must be same as global block end indices length, {len(global_block_end_indices)}!' + ) + for _, (start_idx, end_idx) in enumerate(zip(global_block_indices, global_block_end_indices)): + if start_idx >= end_idx: + raise ValueError( + f'Global block start index, {start_idx}, must be smaller than global block end index, {end_idx}!' + ) + self.global_block_end_indices = global_block_end_indices + + if (attention != 'unidirectional' and attention != 'bidirectional'): + raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') + self.attention = attention + + if (attention != 'bidirectional' and horizontal_global_attention): + raise ValueError('only \"bi-directional\" attentions can support horizontal global attention!') + self.horizontal_global_attention = horizontal_global_attention + + def set_random_layout(self, h, layout): + """Sets random attention layout used by the given head in the sparse attention. + Note) By default, it assumes there will be a unique random block layout for all heads; unless `different_layout_per_head` parameter is set in which each head can have a different random layout. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which random layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_random_blocks): + raise ValueError( + f'Number of random blocks, {self.num_random_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + for row in range(0, num_blocks): + rnd_cols = random.sample(range(0, num_blocks), self.num_random_blocks) + layout[h, row, rnd_cols] = 1 + return layout + + def set_local_layout(self, h, layout): + """Sets local attention layout used by the given head in the sparse attention. + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local layout is set + """ + + num_blocks = layout.shape[1] + start_block_idx = 0 + end_block_idx = 0 + for block_size in self.local_window_blocks: + end_block_idx += block_size + end_block_idx = min(end_block_idx, num_blocks) + for row in range(start_block_idx, end_block_idx): + for col in range(start_block_idx, (row + 1 if self.attention == 'unidirectional' else end_block_idx)): + layout[h, row, col] = 1 + start_block_idx += block_size + + # if there is any remaining not attended part, use the lats local window block size as local window for the remaining applicable local windows + for i in range(start_block_idx, num_blocks, block_size): + end_block_idx = min(i + block_size, num_blocks) + for row in range(i, end_block_idx): + for col in range(i, (row + 1 if self.attention == 'unidirectional' else end_block_idx)): + layout[h, row, col] = 1 + return layout + + def set_global_layout(self, h, layout): + """Sets global attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set + """ + + num_blocks = layout.shape[1] + if (self.global_block_end_indices is None): + for idx in self.global_block_indices: + # if global block idx is in the range of the sequence blocks + if (idx < num_blocks): + #global rows + if (self.horizontal_global_attention): + layout[h, idx, :] = 1 + + #global columns + first_row = 0 if self.attention == 'bidirectional' else idx + layout[h, first_row:, idx] = 1 + else: + for _, (start_idx, end_idx) in enumerate(zip(self.global_block_indices, self.global_block_end_indices)): + # if global block idx is in the range of the sequence blocks + if (start_idx < num_blocks): + end_idx = min(end_idx, num_blocks) + #global rows + if (self.horizontal_global_attention): + layout[h, start_idx:end_idx, :] = 1 + + #global columns + first_row = 0 if self.attention == 'bidirectional' else start_idx + layout[h, first_row:, start_idx:end_idx] = 1 + return layout + + def make_layout(self, seq_len): + """Generates `Variable` sparsity layout used by each head in the sparse attention. + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `Variable` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_random_layout(h, layout) + layout = self.set_local_layout(h, layout) + layout = self.set_global_layout(h, layout) + + layout = self.check_and_propagate_first_head_layout(layout) + return layout + + +class BigBirdSparsityConfig(SparsityConfig): + """Configuration class to store `BigBird` sparsity configuration. + For more details about this sparsity config, please see `Big Bird: Transformers for Longer Sequences`: https://arxiv.org/pdf/2007.14062.pdf + This class extends parent class of `SparsityConfig` and customizes it for `BigBird` sparsity. + """ + + def __init__(self, + num_heads, + block=16, + different_layout_per_head=False, + num_random_blocks=1, + num_sliding_window_blocks=3, + num_global_blocks=1, + attention='bidirectional'): + """Initialize the BigBird Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. + num_random_blocks: optional: an integer determining the number of random blocks in each block row. + num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. + num_global_blocks: optional: an integer determining how many consecutive blocks, starting from index 0, are considered as global attention. Global block tokens will be attended by all other block tokens and will attend to all other block tokens as well. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + """ + + super().__init__(num_heads, block, different_layout_per_head) + + self.num_random_blocks = num_random_blocks + self.num_sliding_window_blocks = num_sliding_window_blocks + self.num_global_blocks = num_global_blocks + + if (attention != 'unidirectional' and attention != 'bidirectional'): + raise NotImplementedError('only \"uni/bi-directional\" attentions are supported for now!') + self.attention = attention + + def set_random_layout(self, h, layout): + """Sets random attention layout used by the given head in the sparse attention. + Note) By default, it assumes there will be a unique random block layout for all heads; unless `different_layout_per_head` parameter is set in which each head can have a different random layout. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which random layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_random_blocks): + raise ValueError( + f'Number of random blocks, {self.num_random_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + for row in range(0, num_blocks): + sample_range = range(0, num_blocks) if self.attention == 'bidirectional' else range(0, row + 1) + rnd_cols = random.sample(sample_range, self.num_random_blocks) + layout[h, row, rnd_cols] = 1 + return layout + + def set_sliding_window_layout(self, h, layout): + """Sets sliding local attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_sliding_window_blocks): + raise ValueError( + f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + w = self.num_sliding_window_blocks // 2 + for row in range(0, num_blocks): + start = max(0, row - w) + end = min(row + w + 1, num_blocks) + layout[h, row, start:end] = 1 + return layout + + def set_global_layout_itc(self, h, layout): + """Sets global attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_global_blocks): + raise ValueError( + f'Number of global blocks, {self.num_global_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + #global rows + layout[h, 0:self.num_global_blocks, :] = 1 + + #global columns + layout[h, :, 0:self.num_global_blocks] = 1 + + if self.attention == 'unidirectional': + # zero out anything attending to the future + layout = torch.tril(layout) + + return layout + + def make_layout(self, seq_len): + """Generates `BigBird` sparsity layout used by each head in the sparse attention. + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BigBird` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_random_layout(h, layout) + layout = self.set_sliding_window_layout(h, layout) + layout = self.set_global_layout_itc(h, layout) + + layout = self.check_and_propagate_first_head_layout(layout) + return layout + + +class BSLongformerSparsityConfig(SparsityConfig): + """Configuration class to store edited `Longformer` sparsity configuration. + + Note) this is a block-sparse version of the Longformer which is slightly different than original Longformer; which is element-wise sparsity. + + For more details about this sparsity config, please see `Longformer: The Long-Document Transformer`: https://arxiv.org/pdf/2004.05150.pdf + This class extends parent class of `SparsityConfig` and customizes it for `Longformer` sparsity. + """ + + def __init__(self, + num_heads, + block=16, + different_layout_per_head=False, + num_sliding_window_blocks=3, + global_block_indices=[0], + global_block_end_indices=None, + attention='bidirectional'): + """Initialize the edited `Longformer` Sparsity Pattern Config. + + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + different_layout_per_head: optional: a boolean determining if each head should be assigned a different sparsity layout; default is false and this will be satisfied based on availability. + + num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. + global_block_indices: optional: a list of integers determining which blocks are considered as global attention. Given indices, determine the blocks that all other token blocks attend to and they attend to all other token blocks. Default value is only index 0. Notice that if global_block_end_indices parameter is set, this parameter is used as starting index of each global window. + global_block_end_indices: optional: a list of integers determining end indices of global window blocks. By default this is not used. But if it is set, it must have the same size of global_block_indices parameter, and combining this two parameters, for each index i, blocks from global_block_indices[i] to global_block_end_indices[i] (exclusive) are considered as global attention. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + """ + + super().__init__(num_heads, block, different_layout_per_head) + + self.num_sliding_window_blocks = num_sliding_window_blocks + self.global_block_indices = global_block_indices + self.attention = attention + + if (global_block_end_indices is not None): + if (len(global_block_indices) != len(global_block_end_indices)): + raise ValueError( + f'Global block start indices length, {len(global_block_indices)}, must be same as global block end indices length, {len(global_block_end_indices)}!' + ) + for _, (start_idx, end_idx) in enumerate(zip(global_block_indices, global_block_end_indices)): + if start_idx >= end_idx: + raise ValueError( + f'Global block start index, {start_idx}, must be smaller than global block end index, {end_idx}!' + ) + self.global_block_end_indices = global_block_end_indices + + def set_sliding_window_layout(self, h, layout): + """Sets sliding local attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_sliding_window_blocks): + raise ValueError( + f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + w = self.num_sliding_window_blocks // 2 + for row in range(0, num_blocks): + start = max(0, row - w) + end = min(row + w + 1, num_blocks) + layout[h, row, start:end] = 1 + return layout + + def set_global_layout(self, h, layout): + """Sets global attention layout used by the given head in the sparse attention. + + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which global layout is set + """ + + num_blocks = layout.shape[1] + if (self.global_block_end_indices is None): + for idx in self.global_block_indices: + # if global block idx is in the range of the sequence blocks + if (idx < num_blocks): + #global rows + layout[h, idx, :] = 1 + + #global columns + layout[h, :, idx] = 1 + else: + for _, (start_idx, end_idx) in enumerate(zip(self.global_block_indices, self.global_block_end_indices)): + # if global block idx is in the range of the sequence blocks + if (start_idx < num_blocks): + end_idx = min(end_idx, num_blocks) + #global rows + layout[h, start_idx:end_idx, :] = 1 + + #global columns + layout[h, :, start_idx:end_idx] = 1 + if self.attention == 'unidirectional': + layout = torch.tril(layout) + return layout + + def make_layout(self, seq_len): + """Generates edited `Longformer` sparsity layout used by each head in the sparse attention. + + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BSLongformer` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_sliding_window_layout(h, layout) + layout = self.set_global_layout(h, layout) + + layout = self.check_and_propagate_first_head_layout(layout) + return layout + + +class LocalSlidingWindowSparsityConfig(SparsityConfig): + """Configuration class to store `Local Sliding Window` sparsity configuration - a purely-local sliding window attention. + This class extends parent class of `SparsityConfig` and customizes it for `Local` sparsity. + """ + + def __init__(self, num_heads, block=16, num_sliding_window_blocks=3, attention='unidirectional'): + """Initialize the Local Sliding Window Sparsity Pattern Config. + For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial + Arguments: + num_heads: required: an integer determining number of attention heads of the layer. + block: optional: an integer determining the block size. Current implementation of sparse self-attention is based on blocked sparse matrices. In which this parameter defines size of such blocks, `Block X Block`. + num_sliding_window_blocks: optional: an integer determining the number of blocks in sliding local attention window. + attention: optional: a string determining attention type. Attention can be `unidirectional`, such as autoregressive models, in which tokens attend only to tokens appear before them in the context. Considering that, the upper triangular of attention matrix is empty as above figure. Or it can be `bidirectional`, such as BERT, in which tokens can attend to any other tokens before or after them. Then, the upper triangular part of the attention matrix is mirror of the lower triangular in the above figure. + """ + + super().__init__(num_heads, block) + self.num_sliding_window_blocks = num_sliding_window_blocks + self.attention = attention + + def set_sliding_window_layout(self, h, layout): + """Sets sliding local attention layout used by the given head in the sparse attention. + Arguments: + h: required: an integer determining head index + layout: required: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head; may not be completely set at this step + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing sparsity layout of all head in which local sliding window layout is set + """ + + num_blocks = layout.shape[1] + if (num_blocks < self.num_sliding_window_blocks): + raise ValueError( + f'Number of sliding window blocks, {self.num_sliding_window_blocks}, must be smaller than overall number of blocks in a row, {num_blocks}!' + ) + + w = self.num_sliding_window_blocks // 2 + for row in range(0, num_blocks): + start = max(0, row - w) + end = min(row + w + 1, num_blocks) if self.attention == "bidirectional" else row + 1 + layout[h, row, start:end] = 1 + return layout + + def make_layout(self, seq_len): + """Generates `Local Sliding Window` sparsity layout used by each head in the sparse attention. + Arguments: + seq_len: required: an integer determining number of attention heads of the layer. + Return: + layout: a tensor of dimension (num_heads, num_blocks, num_blocks) containing `BigBird` sparsity layout of all head + """ + + layout = self.setup_layout(seq_len) + for h in range(0, self.num_layout_heads): + layout = self.set_sliding_window_layout(h, layout) + layout = self.check_and_propagate_first_head_layout(layout) + return layout diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..18d58b5af4bf387ffc8b4e90645e3bd578ee10fe --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import sys +import os + + +def _build_file_index(directory, suffix='.tr'): + """Build an index of source files and their basenames in a given directory. + + Args: + directory (string): the directory to index + suffix (string): index files with this suffix + + Returns: + list: A list of tuples of the form [(basename, absolute path), ...] + """ + + index = [] + + for fname in os.listdir(directory): + if fname.endswith(suffix): + basename = fname[:fname.rfind(suffix)] # strip the suffix + path = os.path.join(directory, fname) + index.append((basename, path)) + + return index + + +# Go over all local source files and parse them as strings +_module = sys.modules[_build_file_index.__module__] +_directory = os.path.dirname(os.path.realpath(__file__)) +for name, fname in _build_file_index(_directory): + with open(fname, 'r') as fin: + setattr(_module, name, fin.read()) diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57a63eb55619c7f08de5c0669ba1bbfc3b6abd38 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr new file mode 100644 index 0000000000000000000000000000000000000000..9be6aa2aed7ccfb0d438d2e4d766d30b21d8f744 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/matmul.tr @@ -0,0 +1,208 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a + https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/matmul.py +*/ + +__global__ void NAME (TYPE* A __readonly __noalias __aligned(16), + TYPE* B __readonly __noalias __aligned(16), + TYPE* C __noalias __aligned(16), + int lda __multipleof(8), + int ldb __multipleof(8), + int ldc __multipleof(8), + long stride_za __multipleof(8), + long stride_zb __multipleof(8), + long stride_zc __multipleof(8), + long stride_ha __multipleof(8), + long stride_hb __multipleof(8), + long stride_hc __multipleof(8), + int DS0, int DS1, + int SDD_K __multipleof(16), + int SDD_off_width, + int* lut, int* locks, int nlocks) { + /* ---------------- */ + /* Prologue */ + /* ---------------- */ + // program ids + int pid0 = get_program_id(0); + int pid1 = get_program_id(1); + int pidz = get_program_id(2); +#ifdef SDD + // load LUT header + pid1 = pid1 + SDD_off_width; + int blockidm[TM] = (0 ... TM) / BLOCK; + int blockidn[TN] = (0 ... TN) / BLOCK; + int offlutm[TM] = blockidm*(TN/BLOCK)*4; + int offlutn[TN] = blockidn*4; + int *header = lut + pid1 * (TM/BLOCK) * (TN/BLOCK) * 4; + int z = *(header + 0); + int i[TM] = *(header + 1 + offlutm); + int j[TN] = *(header + 2 + offlutn); + int AS1 = SDD_K / TZ; + int lockid = select(TZ > 1, 1, 0); + int offka = pid0 * AS1; + int offkb = pid0 * AS1; + int offmc = 0; + int offnc = 0; + int offpa = 0; + int offpb = 0; + int maxid = TZ; + int offhc = 0; + int offha = z; + int offhb = z; + int ram[TM] = i*BLOCK + ((0 ... TM) % BLOCK); + int rbn[TN] = j*BLOCK + ((0 ... TN) % BLOCK); +#else + // load LUT header + int *header = lut + pid0 * 6; + int offset = *(header + 0); + int AS1 = *(header + 1); + int column = *(header + 2); + int depth = *(header + 3); + int lockid = *(header + 4); + int maxid = *(header + 5); + int *pinc = lut + offset; + int offhc = depth; +#ifdef DSD + // output offset + int offnc = pid1 * TN; + int offmc = column * TM; + int offpc = 0; + // dense input offset + int offnb = pid1 * TN; + int offkb __multipleof(8) = *pinc; + int offpb = 0; + // sparse input offset + int offma = 0; + int offka = 0; + long offpa __multipleof(8) = *(pinc + 1); + offpa = offpa * BLOCK * BLOCK; + int offha = 0; + int offhb = depth; +#endif +#ifdef DDS + // output offset + int offmc = pid1 * TM; + int offnc = column * TN; + int offpc = 0; + // dense input offset + int offma = pid1 * TM; + int offka __multipleof(8) = *pinc; + int offpa = 0; + // sparse input offset + int offnb = 0; + int offkb = 0; + long offpb __multipleof(8) = *(pinc + 1); + offpb = offpb * BLOCK * BLOCK; + int offha = depth; + int offhb = 0; +#endif + int ram[TM] = offma + 0 ... TM; + int rbn[TN] = offnb + 0 ... TN; +#endif + // initialize a, b pointers + int rka[TK] = offka + 0 ... TK; + int rkb[TK] = offkb + 0 ... TK; + TYPE* pa[TM, TK] = A + pidz * stride_za + offha * stride_ha + offpa + ram[:, newaxis] * STRIDE_AM + rka[newaxis, :] * STRIDE_AK; + TYPE* pb[TK, TN] = B + pidz * stride_zb + offhb * stride_hb + offpb + rbn[newaxis, :] * STRIDE_BN + rkb[:, newaxis] * STRIDE_BK; + // pre-fetch +#ifdef DDS + bool checkam[TM, TK] = ram[:, newaxis] < DS0; +#else + bool checkam[TM, TK] = AS1 > 0; +#endif +#ifdef DSD + bool checkbn[TK, TN] = rbn[newaxis, :] < DS0; +#else + bool checkbn[TK, TN] = AS1 > 0; +#endif + TYPE a[TM, TK] = checkam ? *pa : 0; + TYPE b[TK, TN] = checkbn ? *pb : 0; + + /* ---------------- */ + /* Inner Loop */ + /* ---------------- */ + // create result tile + float acc[TM, TN] = 0; + int step = TK; + for(int k = AS1; k > 0; k -= step) { + acc += a @ b; + // update pointers +#ifdef SDD + int inc_a = TK * STRIDE_AK; + int inc_b = TK * STRIDE_BK; +#else + pinc += 2; +#ifdef DSD + int inc_b __multipleof(8) = *pinc; + int inc_a __multipleof(8) = *(pinc + 1); + inc_b = inc_b * STRIDE_BK; +#endif +#ifdef DDS + int inc_a __multipleof(8) = *pinc; + int inc_b __multipleof(8) = *(pinc + 1); + inc_a = inc_a * STRIDE_AK; +#endif +#endif + pa += inc_a; + pb += inc_b; + // pre-fetch + bool checkak[TM, TK] = k > TK; + bool checkbk[TK, TN] = k > TK; + bool checka[TM, TK] = checkam && checkak; + bool checkb[TK, TN] = checkbk && checkbn; + a = *?(checka)pa; + b = *?(checkb)pb; + } + TYPE c[TM, TN] = acc; + + /* ---------------- */ + /* Epilogue */ + /* ---------------- */ + // initialize c pointers +#ifdef SDD + bool checkc[TM, TN] = 1; + // rematerialize + int rr_blockidm[TM] = (0 ... TM) / BLOCK; + int rr_blockidn[TN] = (0 ... TN) / BLOCK; + int rr_offlutm[TM] = rr_blockidm*(TN/BLOCK)*4; + int rr_offlutn[TN] = rr_blockidn*4; + int off_bkid[TM, TN] = 3 + rr_offlutm[:, newaxis] + rr_offlutn[newaxis, :]; + int bkid[TM, TN] = *(header + off_bkid); + long offpc[TM, TN] = bkid * BLOCK * BLOCK; + // range within blocks + int rcm[TM] = (0 ... TM) % BLOCK; + int rcn[TN] = (0 ... TN) % BLOCK; +#else + int rcm[TM] = offmc + 0 ... TM; + int rcn[TN] = offnc + 0 ... TN; +#ifdef DSD + bool checkc[TM, TN] = rcn[newaxis, :] < DS0; +#endif +#ifdef DDS + bool checkc[TM, TN] = rcm[:, newaxis] < DS0; +#endif +#endif + TYPE* pc[TM, TN] = C + offpc + offhc*stride_hc + pidz*stride_zc + rcm[:, newaxis]*STRIDE_CM + rcn[newaxis, :]*STRIDE_CN; + // write-back directly + if(lockid == 0) { + *?(checkc) pc = c; + } + // accumulate partial result using spin-locks + else { + int *plock = locks + get_program_id(2)*nlocks*get_num_programs(1) + get_program_id(1)*nlocks + lockid - 1; + int *pcount = plock + get_num_programs(2)*get_num_programs(1)*nlocks; + for(int repeat = 1; repeat == 1; repeat = atomic_cas(plock, 0, 1)); + int count = *pcount; + if(count == 0) + *?(checkc) pc = c; + else + *?(checkc) pc = c + *?(checkc)pc; + atomic_xchg(pcount, (count + 1) % maxid); + atomic_xchg(plock, 0); + } + } diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr new file mode 100644 index 0000000000000000000000000000000000000000..1a90f41d94945e1d6d6f52e6beaea94fa52cdda8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_bwd.tr @@ -0,0 +1,61 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a + https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/softmax.py +*/ + +__global__ void softmax_bwd(TYPE * X __readonly __noalias __aligned(16), + float scale, + TYPE* DX __readonly __noalias __aligned(16), + int* LUT, + int sizemax, + long stride_zx __multipleof(BLOCK), + long stride_zdx __multipleof(BLOCK)) { + int pidhm = get_program_id(0); + int pidz = get_program_id(1); + + // create index ranges + int rxm = pidhm % BLOCK; + int rbm = pidhm / BLOCK; + int rxn[TN] = (0 ... TN) % BLOCK; + int rbn[TN] = (0 ... TN) / BLOCK; + + // extract information from look-up table + int* header = LUT + rbm * 2; + int size = *(header + 0); + int offset = *(header + 1); + + // bounds checking on lut + bool check[TN] = rbn < size; + int rbmn[TN] = check ? rbn : size - 1; + + // initialize pointers to block-sparse input + long blockid[TN] = *(LUT + offset + rbmn*4); + + TYPE* px[TN] = X + pidz * stride_zx + + blockid * BLOCK * BLOCK + + rxm * BLOCK + + rxn; + + TYPE* pdx[TN] = DX + pidz * stride_zdx + + blockid * BLOCK * BLOCK + + rxm * BLOCK + + rxn; + + // compute fused softmax backward + TYPE x[TN] = check ? *px : 0; + TYPE dx[TN] = check ? *pdx : 0; + float Fdx[TN] = dx; + float Fx[TN] = x; + float Fxdx[TN] = Fdx*Fx; + float Fxdxsum = Fxdx[+]; + float Fy[TN] = Fx * (Fdx - Fxdxsum) * scale; + TYPE y[TN] = Fy; + + // write-back + *? (check)pdx = y; +} diff --git a/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr new file mode 100644 index 0000000000000000000000000000000000000000..ebd317d9469b47f7e2ee3032d3aabf57b5620a73 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/ops/sparse_attention/trsrc/softmax_fwd.tr @@ -0,0 +1,143 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +// DeepSpeed Team + +/* +DeepSpeed note, code taken & adapted from commit 9aa94789f13ada713af36cfd8cca2fc9a7f6b79a + https:github.com/ptillet/torch-blocksparse/blob/master/torch_blocksparse/softmax.py +*/ + +__global__ void softmax_fwd(TYPE *X __readonly __noalias __aligned(16), + float scale, + int *LUT __readonly __noalias __aligned(16), + TYPE *RPE __readonly __noalias __aligned(16), + TYPE *KP_M __readonly __noalias __aligned(16), + TYPE *ATTN_M __readonly __noalias __aligned(16), + int num_blocks, + int sizemax, + long stride_zx __multipleof(BLOCK), + long stride_zrpe __multipleof(BLOCK), + int stride_hrpe __multipleof(BLOCK), + int stride_srpe __multipleof(BLOCK), + int stride_zkpm __multipleof(BLOCK), + int stride_zattnm __multipleof(BLOCK)){ + int pidhm = get_program_id(0); + int pidz = get_program_id(1); + + // create index ranges + int rxm = pidhm % BLOCK; + int rbm = pidhm / BLOCK; + int rxn[TN] = (0 ... TN) % BLOCK; + int rbn[TN] = (0 ... TN) / BLOCK; + + // extract information from look-up table + int* header = LUT + rbm * 2; + int size = *(header + 0); + int offset = *(header + 1); + + bool check[TN] = rbn < size; + int rbmn[TN] = check ? rbn : size - 1; + + // block id and column id + long blockid [TN] = *(LUT + offset + rbmn*4 + 0); + long columnid[TN] = *(LUT + offset + rbmn*4 + 1); + long rowid [TN] = *(LUT + offset + rbmn*4 + 2); + long headid [TN] = *(LUT + offset + rbmn*4 + 3); + + // pointers to X + TYPE* px[TN] = X + pidz * stride_zx + + blockid * BLOCK * BLOCK + + rxm * BLOCK + + rxn; +#ifdef APPLY_RPE + // pointers to relative position embedding + TYPE* prpe[TN] = RPE + pidz * stride_zrpe + + headid * stride_hrpe + + columnid * BLOCK + + rowid * BLOCK * stride_srpe + + rxm * stride_srpe + + rxn; +#endif + +#ifdef APPLY_KP_MASK + // pointers to key padding mask + TYPE* pkp_m[TN] = KP_M + pidz * stride_zkpm + + columnid * BLOCK + + rxn; +#endif + +#ifdef APPLY_ATTN_MASK + // pointers to attention mask + TYPE* pattn_m[TN] = ATTN_M + columnid * BLOCK + + rowid * BLOCK * stride_zattnm + + rxm * stride_zattnm + + rxn; +#endif + + // load input + TYPE x[TN] = check ? *px : -INFINITY; + +#ifdef APPLY_RPE + // load relative position embedding + TYPE rpe[TN] = check ? *prpe : 0; +#endif + +#ifdef APPLY_KP_MASK + // load key-padding mask + TYPE kp_m[TN] = check ? *pkp_m : -INFINITY; +#endif + +#ifdef APPLY_ATTN_MASK + // load attention mask + TYPE attn_m[TN] = check ? *pattn_m : -INFINITY; +#endif + + // compute softmax in float +#ifdef APPLY_RPE + float Frpe[TN] = rpe; +#endif + +#ifdef APPLY_KP_MASK + float Fkp_m[TN] = kp_m; +#endif + +#ifdef APPLY_ATTN_MASK + float Fattn_m[TN] = attn_m; +#endif + +#ifdef KP_MASK_MUL + Fkp_m = (Fkp_m == 0) ? (float[TN])-INFINITY : 0; +#endif + +#ifdef ATTN_MASK_MUL + Fattn_m = (Fattn_m == 0) ? (float[TN])-INFINITY : 0; +#endif + + float Fx[TN] = x; + +#ifdef APPLY_SCALE + Fx = Fx * scale; // apply scale +#endif + +#ifdef APPLY_RPE + Fx = Fx + Frpe; // apply relative position embedding +#endif + +#ifdef APPLY_KP_MASK + Fx = Fx + Fkp_m; // apply key padding mask +#endif + +#ifdef APPLY_ATTN_MASK + Fx = Fx + Fattn_m; // apply attention mask +#endif + + float Fxmax = Fx[max]; + float Fy[TN] = exp(Fx - Fxmax); + float Fysum = (check ? Fy : 0)[+]; + + // write-back in half/float + TYPE y[TN] = Fy; + TYPE ysum = Fysum; + *?(check)px = y / ysum; +} diff --git a/lib/python3.12/site-packages/deepspeed/runtime/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..208299fb8c50f73468d293b6fa5dca71649d62e7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32887d5a7a59311dbefd5474ecd6ee7c624ba2bd Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68039e51c1b3f61cdb69285caea09bbb918ac53b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/checkpointing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..794198398c1939f84f264b51624ab62821b0edcc Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py new file mode 100644 index 0000000000000000000000000000000000000000..08c4b81937f918425d6da976b9c8391b30eee741 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py @@ -0,0 +1,1142 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Use to partition the activations stored for backward propagation +Therefore reduces the memory consumption +Also implements CPU checkpointing and contiguous memory checkpointing +Reduces memory consumption and memory fragmentation + +Code for rng checkpointing taken from NVIDIA Megatron-LM mpu/random.py +b886b7bb972afe72bac0f5de4f42a4a7bae8ebef +""" + +# Parts of the code here are adapted from PyTorch +# repo: https://github.com/pytorch/pytorch +import copy +import torch +import contextlib +from deepspeed import comm as dist +import weakref + +import mmap +from torch import _C + +from deepspeed.runtime.config import DeepSpeedConfig +from deepspeed.utils import logger +from deepspeed.runtime.utils import copy_to_device, move_to_device, see_memory_usage +from deepspeed.utils.timer import SynchronizedWallClockTimer as Timers, FORWARD_GLOBAL_TIMER +from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime import compiler + +# DeepSpeed Checkpointing Enabled or Disabled +deepspeed_checkpointing_enabled = False + +# MP parameters +mpu = None + +#set default values +mp_rank = 0 +mp_size = 1 +mp_group = None + +# Model Parameters +num_layers = None + +# Checkpointing buffers +contiguous_data_buffers = [] +data_offsets = [] + +contiguous_size_buffers = [] +size_offsets = [] + +timers = None + +# optimization flags +PARTITION_ACTIVATIONS = False +CPU_CHECKPOINT = False +CONTIGUOUS_CHECKPOINTING = False +SYNCHRONIZE = False +PROFILE_TIME = False + +# Default name for the model parallel rng tracker. +_MODEL_PARALLEL_RNG_TRACKER_NAME = 'model-parallel-rng' + + +def detach_variable(inputs, device=None): + if isinstance(inputs, tuple): + out = [] + for inp in inputs: + if not isinstance(inp, torch.Tensor): + out.append(inp) + continue + + requires_grad = inp.requires_grad + + if device is not None: + x = inp.to(device=device) + else: + x = inp + + x = x.detach() + x.requires_grad = requires_grad + out.append(x) + return tuple(out) + else: + raise RuntimeError("Only tuple of tensors is supported. Got Unsupported input type: ", type(inputs).__name__) + + +def _set_cuda_rng_state(new_state, device=-1): + """Sets the random number generator state of the current GPU. + + Arguments: + new_state (torch.ByteTensor): The desired state + This function is adapted from PyTorch repo (torch.cuda.set_rng_state) #ignore-cuda + with a single change: the input state is not cloned. Cloning caused + major performance issues for +4 GPU cases. + """ + if hasattr(_C, '_cuda_setRNGState') and callable(_C._cuda_setRNGState): + # older PyTorch + def cb(): + with get_accelerator().device(device): + _C._cuda_setRNGState(new_state) + else: + # newer PyTorch + if device == -1: + device = torch.device(get_accelerator().device_name()) + elif isinstance(device, str): + device = torch.device(device) + elif isinstance(device, int): + device = torch.device(get_accelerator().device_name(), device) + + def cb(): + idx = device.index + if idx is None: + idx = get_accelerator().current_device() + default_generator = get_accelerator().default_generator(idx) + default_generator.set_state(new_state) + + get_accelerator().lazy_call(cb) + + +class CudaRNGStatesTracker: + """Tracker for the cuda RNG states. + + Using the `add` method, a cuda rng state is initialized based on + the input `seed` and is assigned to `name`. Later, by forking the + rng state, we can perform operations and return to our starting + cuda state. + """ + + def __init__(self): + # Map from a string name to the cuda rng state. + self.states_ = {} + # Seeds are just for book keeping and ensure no seed is set twice. + self.seeds_ = set() + + def reset(self): + """Set to the initial state (no tracker).""" + self.states_ = {} + self.seeds_ = set() + + def get_states(self): + """Get rng states. Copy the dictionary so we have direct + pointers to the states, not just a pointer to the dictionary.""" + return copy.copy(self.states_) + + def set_states(self, states): + """Set the rng states. For efficiency purposes, we do not check + the size of seed for compatibility.""" + self.states_ = states + + def add(self, name, seed): + """Track the rng state.""" + # Check seed is not already used. + if seed in self.seeds_: + raise Exception('seed {} already exists'.format(seed)) + self.seeds_.add(seed) + # Check that state is not already defined. + if name in self.states_: + raise Exception('cuda rng state {} already exists'.format(name)) + # Get the current rng state. + orig_rng_state = get_accelerator().get_rng_state() + # Set the new state and store it. + get_accelerator().manual_seed(seed) + self.states_[name] = get_accelerator().get_rng_state() + # Reset rng state to what it was. + _set_cuda_rng_state(orig_rng_state) + + @contextlib.contextmanager + def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME): + """Fork the cuda rng state, perform operations, and exit with + the original state.""" + # Check if we have added the state + if name not in self.states_: + raise Exception('cuda rng state {} is not added'.format(name)) + # Store current rng state. + orig_cuda_rng_state = get_accelerator().get_rng_state() + # Set rng state to the desired one + _set_cuda_rng_state(self.states_[name]) + # Do the stuff we wanted to do. + try: + yield + finally: + # Update the current rng state for later use. + self.states_[name] = get_accelerator().get_rng_state() + # And set the state to the original state we started with. + _set_cuda_rng_state(orig_cuda_rng_state) + + +# RNG tracker object. +_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker() + + +def get_cuda_rng_tracker(): + """Get cuda rng tracker.""" + return _CUDA_RNG_STATE_TRACKER + + +def model_parallel_cuda_manual_seed(seed): + """Initialize model parallel cuda seed. + + This function should be called after the model parallel is + initialized. Also, no get_accelerator().manual_seed should be called + after this function. Basically, this is replacement for that + function. + Two set of RNG states are tracked: + default state: This is for data parallelism and is the same among a + set of model parallel GPUs but different across + different model parallel groups. This is used for + example for dropout in the non-model-parallel regions. + model-parallel state: This state is different among a set of model + parallel GPUs, but the same across data parallel + groups. This is used for example for dropout in + model parallel regions. + """ + global mpu + + tp_rank = bwc_tensor_model_parallel_rank(mpu) + + # 2718 is just for fun and any POSITIVE value will work. + offset = seed + 2718 + model_parallel_seed = offset + tp_rank + # Data parallel gets the original seed. + data_parallel_seed = seed + + if dist.get_rank() == 0: + logger.info( + '> initializing model parallel cuda seeds on global rank {}, ' + 'model parallel rank {}, and data parallel rank {} with ' + 'model parallel seed: {} and data parallel seed: {}'.format(dist.get_rank(), tp_rank, + mpu.get_data_parallel_rank(), + model_parallel_seed, data_parallel_seed), ) + _CUDA_RNG_STATE_TRACKER.reset() + # Set the default state. + get_accelerator().manual_seed(data_parallel_seed) + # and model parallel state. + _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, model_parallel_seed) + + +def model_parallel_reconfigure_tp_seed(seed): + global mpu + tp_rank = bwc_tensor_model_parallel_rank(mpu) + model_parallel_seed = seed + 2718 + tp_rank + with _CUDA_RNG_STATE_TRACKER.fork(): + get_accelerator().manual_seed(model_parallel_seed) + + +def get_partition_start(item): + global mp_rank, mp_size, mp_group + size = item.numel() + partition_size = size / mp_size + start = partition_size * mp_rank + return int(start) + + +def get_partition_size(item): + global mp_rank, mp_size, mp_group + size = item.numel() + assert size % mp_size == 0, "Doesn't handle if partition activation if item is not divisible by mp size" + partition_size = size / mp_size + return int(partition_size) + + +def gather_partitioned_activations(tensors, device=None): + global mp_rank, mp_size, mp_group + assert len(tensors) % 2 == 0, f'Expected even count of tensors, instead got {len(tensors)}' + inputs = [] + num_args = int(len(tensors) / 2) + for i in range(num_args): + + item = tensors[2 * i] + size = tensors[2 * i + 1] + + if not is_activation_to_checkpoint(item): + inputs.append(item) + continue + + # don't need to do all_gather if model parallel is not enabled + if mp_group is None or mp_size == 1: + item = item.view(list(size.numpy())) + if device is not None: + item = item.to(device) + inputs.append(item) + continue + + partition_size = item.numel() + tensor_size = partition_size * mp_size + if device is not None: + flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=device) + else: + flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=item.device) + part = flat_tensor.narrow(0, partition_size * mp_rank, partition_size) + part.copy_(item) + dist.all_gather_into_tensor(flat_tensor, part, group=mp_group) + input_tensor = flat_tensor.view(list(size.numpy())) + item.data = input_tensor.data + + inputs.append(item) + + return tuple(inputs) + + +def extract_tensors(all_objects): + """ + Separate objects in list/tuple into tensors and non-tensors and create a mapping to enable re-aggregation. + The order of tensors and non-tensors is preserved in their respective output groups. + + Parameters: + all_objects (list/tuple): Objects containing tensors and non-tensors to be split. + + Returns: + tuple: Containing tensors, non-tensors, and bools of whether each position in original list/tuple was a tensor. + + """ + tensor_objects = [v for v in all_objects if torch.is_tensor(v)] + non_tensor_objects = [v for v in all_objects if not torch.is_tensor(v)] + tensor_flags = [torch.is_tensor(v) for v in all_objects] + if type(all_objects) is tuple: + return tuple(tensor_objects), tuple(non_tensor_objects), tuple(tensor_flags) + return tensor_objects, non_tensor_objects, tensor_flags + + +def merge_tensors(tensor_objects, non_tensor_objects, tensor_flags): + """ + Merge two lists (or tuples) of tensors and non-tensors using a mapping of positions in merged list (or tuple). + + Parameters: + tensor_objects (list/tuple): Tensors to merge. + non_tensor_objects (list/tuple): Non-tensors to merge. + tensor_flags (list/tuple): Indicates whether each position in output is a tensor. + + Returns: + tuple: Merge of tensors and non-tensors + """ + merged_objects = [] + tensor_idx = 0 + non_tensor_idx = 0 + + real_tensor_flags = None + + # remove the flags that are assigned to the size of the flattened tensors + if PARTITION_ACTIVATIONS: + real_tensor_flags = [] + previous_flag = False + for flag in tensor_flags: + if previous_flag: + previous_flag = False + continue + previous_flag = flag + real_tensor_flags.append(flag) + else: + real_tensor_flags = tensor_flags + + for is_tensor in real_tensor_flags: + if is_tensor: + merged_objects.append(tensor_objects[tensor_idx]) + tensor_idx += 1 + else: + merged_objects.append(non_tensor_objects[non_tensor_idx]) + non_tensor_idx += 1 + + return tuple(merged_objects) + + +def is_activation_to_checkpoint(item): + """ + Is an activation to be checkpointed + """ + global mp_size + extra_flag = (not hasattr(item, 'no_checkpointing')) or (hasattr(item, 'no_checkpointing') + and item.no_checkpointing == False) + return torch.is_tensor(item) and item.is_floating_point() and item.numel() >= mp_size and extra_flag + + +def partition_activations(args, cpu_checkpoint, contiguous_checkpoint): + global contiguous_data_buffers, data_offsets + + inputs = [] + num_non_fp_tensors = 0 + + for arg_index, item in enumerate(args): + if not is_activation_to_checkpoint(item): + inputs.append(item) + num_non_fp_tensors += 1 + continue + + i = arg_index - num_non_fp_tensors + partition_size = get_partition_size(item) + partition = item.detach().contiguous().view(-1).narrow(0, get_partition_start(item), partition_size).clone() + + buffer_device = torch.device('cpu') if cpu_checkpoint else partition.device + + if contiguous_checkpoint: + if i >= len(contiguous_data_buffers): + tensor_list = [ + torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) + for _ in range(num_layers) + ] + contiguous_data_buffers.append(tensor_list) + data_offsets.append(0) + elif contiguous_data_buffers[i] is None: + tensor_list = [ + torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device) + for _ in range(num_layers) + ] + contiguous_data_buffers[i] = tensor_list + data_offsets[i] = 0 + + # Because the 'new_empty' returns uninitialized pages, + # the pages need to be populated during the cudaMemcpy time + # which increases the data copy time. To avoid this, we + # pre-populate these pages by simply writing 0 ahead of + # the actual cudaMemcpy operation time. Due to the + # previously launched GPU kernels, there is a small + # window of time here for CPUs to populate pages asynchronously. + contiguous_data_buffers[i][data_offsets[i]].data[range( + 0, contiguous_data_buffers[i][data_offsets[i]].data.shape[0], + int(mmap.PAGESIZE / contiguous_data_buffers[i][data_offsets[i]].data.element_size()))] = 0 + + contiguous_partition = contiguous_data_buffers[i][data_offsets[i]].data.copy_(partition.data) + data_offsets[i] = data_offsets[i] + 1 + inputs.append(contiguous_partition) + else: + partition = partition.cpu() if CPU_CHECKPOINT else partition + inputs.append(partition) + + return inputs + + +def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint): + global contiguous_size_buffers, size_offsets + + new_args = [] + num_non_fp_tensors = 0 + + for arg_index, (arg, inp) in enumerate(zip(args, inputs)): + size = torch.tensor(arg.size()) if torch.is_tensor(arg) else None + if not is_activation_to_checkpoint(arg): + new_args.append(arg) + new_args.append(size) + num_non_fp_tensors += 1 + continue + + arg.data = torch.empty([], device=arg.device).data + arg.saved_data = inp.data + + new_args.append(arg) + i = arg_index - num_non_fp_tensors + + if contiguous_checkpoint: + numel = size.numel() + if i >= len(contiguous_size_buffers): + tmp = torch.tensor(()) + contiguous_size_buffers.append( + tmp.new_empty([numel * num_layers], dtype=size.dtype, device=size.device)) + size_offsets.append(0) + elif contiguous_size_buffers[i] is None: + tmp = torch.tensor(()) + contiguous_size_buffers[i] = tmp.new_empty([numel * num_layers], dtype=size.dtype, device=size.device) + size_offsets[i] = 0 + + contiguous_size = contiguous_size_buffers[i].narrow(0, size_offsets[i], numel).data.copy_(size.data) + contiguous_size = contiguous_size.view_as(size) + size_offsets[i] = size_offsets[i] + numel + new_args.append(contiguous_size) + else: + new_args.append(size) + + return new_args + + +def get_cpu_activations_for_backward(args, inputs): + new_args = [] + for i, (arg, inp) in enumerate(zip(args, inputs)): + if not is_activation_to_checkpoint(arg): + new_args.append(arg) + continue + + arg.data = torch.empty([], device=arg.device).data + arg.saved_data = inp.data + new_args.append(arg) + + return new_args + + +class CheckpointFunction(torch.autograd.Function): + """This function is adapted from torch.utils.checkpoint with + two main changes: + 1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state` #ignore-cuda + 2) the states in the model parallel tracker are also properly + tracked/set/reset. + 3) Performance activation partitioning, contiguous memory optimization + 4) CPU Checkpointing + 5) Profile forward and backward functions + """ + + @staticmethod + def forward(ctx, run_function, all_outputs, *args): + global mpu, timers, SYNCHRONIZE, PROFILE_TIME + + def save_args_for_backward(*all_args): + tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args) + ctx.deepspeed_saved_tensors = tensor_args + ctx.non_tensor_args = non_tensor_args + ctx.tensor_flags = tensor_flags + + if SYNCHRONIZE: + get_accelerator().synchronize() + + if timers is None and PROFILE_TIME: + timers = Timers() + + if PROFILE_TIME: + timers(FORWARD_GLOBAL_TIMER).start() + + ctx.run_function = run_function + global num_layers + global mp_rank, mp_size, mp_group + global contiguous_data_buffers, contiguous_size_buffers + global data_offsets, size_offsets + global PARTITION_ACTIVATIONS, buffer_0, buffer_1, buffer_0_offset, buffer_1_offset + + cuda_device = get_accelerator().current_device_name() + transport_stream = get_accelerator().Stream(device=cuda_device) + + if PARTITION_ACTIVATIONS: + inputs = partition_activations(args, CPU_CHECKPOINT, CONTIGUOUS_CHECKPOINTING) + elif CPU_CHECKPOINT: + inputs = copy_to_device(args, device=torch.device('cpu'), criterion_func=is_activation_to_checkpoint) + + # just in case something funky is happening such as reuse of inputs + inputs_cuda = copy_to_device(args, device=cuda_device, criterion_func=is_activation_to_checkpoint) + + # Copy the rng states. + ctx.fwd_cpu_rng_state = torch.get_rng_state() + ctx.fwd_cuda_rng_state = get_accelerator().get_rng_state() + ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() + + see_memory_usage("Before running forward on the layer", force=False) + # ctx.save_for_backward(*args) + with torch.no_grad(): + outputs = run_function(*inputs_cuda) + + see_memory_usage("After running forward on the layer", force=False) + del inputs_cuda + + if PARTITION_ACTIVATIONS: + new_args = get_partitioned_activations_for_backward(args, inputs, CONTIGUOUS_CHECKPOINTING) + assert len(new_args) % 2 == 0, f'save_for_backward called with odd number of args, {len(new_args)}' + save_args_for_backward(*new_args) + elif CPU_CHECKPOINT: + new_args = get_cpu_activations_for_backward(args, inputs) + save_args_for_backward(*new_args) + else: + save_args_for_backward(*args) + + if PROFILE_TIME: + timers(FORWARD_GLOBAL_TIMER).stop() + timers.log([FORWARD_GLOBAL_TIMER]) + if SYNCHRONIZE: + get_accelerator().synchronize() + + # Tensors returned from forward() may not be differentiable. + if torch.is_tensor(outputs): + non_grad_outputs = [outputs] if not outputs.is_floating_point() else [] + else: + non_grad_outputs = [o for o in outputs if torch.is_tensor(o) and not o.is_floating_point()] + ctx.mark_non_differentiable(*non_grad_outputs) + + if torch.is_tensor(outputs): + all_outputs += [outputs] + return outputs + else: + all_outputs += outputs + outputs, _, _ = extract_tensors(all_objects=outputs) + return tuple(outputs) + + @staticmethod + def backward(ctx, *grads): + global timers + see_memory_usage("In backward", force=False) + # removing pointers to the contiguous buffer memory + # so that they can be garbage collected once the checkpoints + # have been used + if SYNCHRONIZE: + get_accelerator().synchronize() + if PROFILE_TIME: + timers('backward').start() + + if CONTIGUOUS_CHECKPOINTING: + global data_offsets, size_offsets + global contiguous_data_buffers, contiguous_size_buffers + + for buffers in contiguous_data_buffers: + buffers = [] + + # frees up all the pointers to the checkpoints except for the ones + # stored by save for backward + contiguous_data_buffers = [] + contiguous_size_buffers = [] + data_offsets = [] + size_offsets = [] + + see_memory_usage("In backward checkpointing code", force=False) + if not torch.autograd._is_checkpoint_valid(): + raise RuntimeError("Checkpointing is not compatible with .grad(), " + "please use .backward() if possible") + + global PARTITION_ACTIVATIONS + cuda_device = get_accelerator().current_device_name() + transport_stream = get_accelerator().Stream(device=cuda_device) + # Rebuild deepspeed_saved_tensors + for t in ctx.deepspeed_saved_tensors: + if t is not None and hasattr(t, 'saved_data') and t.saved_data is not None: + t.data = t.saved_data.to(t.device) + t.saved_data = None + + if PARTITION_ACTIVATIONS: + # with get_accelerator().stream(transport_stream): + inputs = gather_partitioned_activations(ctx.deepspeed_saved_tensors, + device=cuda_device if CPU_CHECKPOINT else None) + detached_inputs = detach_variable(inputs) + elif CPU_CHECKPOINT: + inputs = move_to_device(ctx.deepspeed_saved_tensors, cuda_device, is_activation_to_checkpoint) + detached_inputs = detach_variable(inputs) + else: + inputs = ctx.deepspeed_saved_tensors + detached_inputs = detach_variable(inputs) + + # Add non tensor input args + detached_inputs = merge_tensors(tensor_objects=detached_inputs, + non_tensor_objects=ctx.non_tensor_args, + tensor_flags=ctx.tensor_flags) + + # Store the current states. + bwd_cpu_rng_state = torch.get_rng_state() + bwd_cuda_rng_state = get_accelerator().get_rng_state() + bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() + + # Set the states to what it used to be before the forward pass. + torch.set_rng_state(ctx.fwd_cpu_rng_state) + _set_cuda_rng_state(ctx.fwd_cuda_rng_state) + get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker) + + # if PARTITION_ACTIVATIONS: + # current_stream=get_accelerator().current_stream() + # current_stream.wait_stream(transport_stream) + + see_memory_usage("In backward checkpointing code before forward", force=False) + + with torch.enable_grad(): + outputs = ctx.run_function(*detached_inputs) + + see_memory_usage("In backward checkpointing code after forward", force=False) + # Set the states back to what it was at the start of this function. + torch.set_rng_state(bwd_cpu_rng_state) + _set_cuda_rng_state(bwd_cuda_rng_state) + get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker) + + if isinstance(outputs, torch.Tensor): + outputs = (outputs, ) + + # Filter out non tensor outputs + outputs, _, _ = extract_tensors(all_objects=outputs) + + # Construct arguments to autograd.backward(). + # This is usually just outputs and grads, but forward() can return tensors that + # are not differentiable. + output_tensors = [] + grad_tensors = [] + for out, grad in zip(outputs, grads): + if out.requires_grad: + output_tensors.append(out) + grad_tensors.append(grad) + + see_memory_usage("In backward checkpointing code before backward", force=False) + + torch.autograd.backward(output_tensors, grad_tensors) + + # Force clear our stashed tensors to prevent a memory leak in certain scenarios + ctx.deepspeed_saved_tensors = None + ctx.non_tensor_args = None + ctx.tensor_flags = None + + see_memory_usage("After backward checkpointing code after backward", force=False) + + if PROFILE_TIME: + timers('backward').stop() + timers.log(['backward']) + if SYNCHRONIZE: + get_accelerator().synchronize() + ret_list = [None, None] # first None for ctx + for inp in detached_inputs: + if torch.is_tensor(inp): + ret_list.append(inp.grad) + else: + ret_list.append(None) + + return tuple(ret_list) + + +def non_reentrant_checkpoint(function, *args): + """This function is union of `torch.utils.checkpoint._checkpoint_without_reentrant` and `CheckpointFunction` in this module + + This function is aim to solve the back probagation error raised from all input requires no grad. + * has already been implemented in pytorch for a while, the solution is stable at most time except for jit module mode. + * can help to solve the issue which is hacked by `deepspeed.runtime.pipe.module.PipelineModule._is_checkpointable` + + Main modifications compared to the implementation of torch: + 1. adapt to the signature of `checkpoint` function in this module + 2. solve the non-deterministic by random state management consistent with deepspeed `CheckpointFunction` + 3. when there is partition or cpu checkpointing, gather them in the unpack_hook during back probagation + 4. make all after backward blocks in the hook which will executed after all leaf nodes backward execution. + 5. above 4. is inspired by `torch.autograd.graph.register_multi_grad_hook`, which is only implemented after 2.0.0 + """ + global mpu, timers, SYNCHRONIZE, PROFILE_TIME + + deepspeed_saved_tensors = None + non_tensor_args = None + tensor_flags = None + + def save_args_for_backward(*all_args): + """keep this function to reduce the modification from original implementation""" + nonlocal deepspeed_saved_tensors, non_tensor_args, tensor_flags + tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args) + deepspeed_saved_tensors = tensor_args + non_tensor_args = non_tensor_args + tensor_flags = tensor_flags + + if SYNCHRONIZE: + get_accelerator().synchronize() + + if timers is None and PROFILE_TIME: + timers = Timers() + + if PROFILE_TIME: + timers(FORWARD_GLOBAL_TIMER).start() + + global num_layers + global mp_rank, mp_size, mp_group + global contiguous_data_buffers, contiguous_size_buffers + global data_offsets, size_offsets + global PARTITION_ACTIVATIONS, buffer_0, buffer_1, buffer_0_offset, buffer_1_offset + + cuda_device = get_accelerator().current_device_name() + transport_stream = get_accelerator().Stream(device=cuda_device) + + if PARTITION_ACTIVATIONS: + inputs = partition_activations(args, CPU_CHECKPOINT, CONTIGUOUS_CHECKPOINTING) + elif CPU_CHECKPOINT: + inputs = copy_to_device(args, device=torch.device('cpu'), criterion_func=is_activation_to_checkpoint) + + # just in case something funky is happening such as reuse of inputs + inputs_cuda = copy_to_device(args, device=cuda_device, criterion_func=is_activation_to_checkpoint) + + # Copy the rng states. + fwd_cpu_rng_state = torch.get_rng_state() + fwd_cuda_rng_state = get_accelerator().get_rng_state() + fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() + + if PARTITION_ACTIVATIONS: + new_args = get_partitioned_activations_for_backward(args, inputs, CONTIGUOUS_CHECKPOINTING) + assert len(new_args) % 2 == 0, f'save_for_backward called with odd number of args, {len(new_args)}' + save_args_for_backward(*new_args) + elif CPU_CHECKPOINT: + new_args = get_cpu_activations_for_backward(args, inputs) + save_args_for_backward(*new_args) + else: + save_args_for_backward(*args) + + class Holder(): + """the place holder object used as activations to save memory""" + pass + + # weakref seems utilized to discover the tensor deletion before a whole + # forward backward pair loop finished + storage: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + weak_holder_list = [] + leaf_tensors = [] + backward_visited_leaf_nodes = 0 + + def checkpoint_pack(tensor_from_forward): + """used to record the activation order in the `weak_holder_list` + + the activation order in holder list is consistent between the first forward and recomputing forward. + * the jit compiled forward will break the order consistency * + """ + res = Holder() + weak_holder_list.append(weakref.ref(res)) + + # if this is a leaf tensor, save it for backward progression trace + # leaf tensor used to be input or parameters, which is not activations and + # has no memory overhead + if tensor_from_forward.requires_grad and tensor_from_forward.is_leaf: + leaf_tensors.append(tensor_from_forward) + return res + + def checkpoint_unpack(holder_from_backward): + """retrieve the activations from recompute""" + nonlocal deepspeed_saved_tensors, non_tensor_args, tensor_flags + + # if this is the first step of backward probagation, recompute the graph and save + # all the activations with the same order as `checkpoint_pack` does + if len(storage) == 0: + unpack_counter = 0 + + def replay_pack(tensor_from_replay): + """save recompute activations""" + nonlocal unpack_counter + unpack_counter += 1 + + if weak_holder_list[unpack_counter - 1]() is None: + return + + detached_activations = tensor_from_replay.detach() + storage[weak_holder_list[unpack_counter - 1]()] = detached_activations + + return + + def replay_unpack(none_value): + """recompute graph need not to backward""" + raise RuntimeError("You are calling backwards on a tensor that is never exposed.") + + global timers + see_memory_usage("In backward", force=False) + # removing pointers to the contiguous buffer memory + # so that they can be garbage collected once the checkpoints + # have been used + if SYNCHRONIZE: + get_accelerator().synchronize() + if PROFILE_TIME: + timers('backward').start() + + if CONTIGUOUS_CHECKPOINTING: + global data_offsets, size_offsets + global contiguous_data_buffers, contiguous_size_buffers + + for buffers in contiguous_data_buffers: + buffers = [] + + # frees up all the pointers to the checkpoints except for the ones + # stored by save for backward + contiguous_data_buffers = [] + contiguous_size_buffers = [] + data_offsets = [] + size_offsets = [] + + see_memory_usage("In backward checkpointing code", force=False) + if not torch.autograd._is_checkpoint_valid(): + raise RuntimeError("Checkpointing is not compatible with .grad(), " + "please use .backward() if possible") + + global PARTITION_ACTIVATIONS + cuda_device = get_accelerator().current_device_name() + transport_stream = get_accelerator().Stream(device=cuda_device) + + # gather inputs which is partitioned or checkpointed before first forward + if PARTITION_ACTIVATIONS: + # with get_accelerator().stream(transport_stream): + inputs = gather_partitioned_activations(deepspeed_saved_tensors, + device=cuda_device if CPU_CHECKPOINT else None) + detached_inputs = detach_variable(inputs) + elif CPU_CHECKPOINT: + inputs = move_to_device(deepspeed_saved_tensors, cuda_device, is_activation_to_checkpoint) + detached_inputs = detach_variable(inputs) + else: + inputs = deepspeed_saved_tensors + detached_inputs = detach_variable(inputs) + + # Add non tensor input args + detached_inputs = merge_tensors(tensor_objects=detached_inputs, + non_tensor_objects=non_tensor_args, + tensor_flags=tensor_flags) + + # Store the current states. + bwd_cpu_rng_state = torch.get_rng_state() + bwd_cuda_rng_state = get_accelerator().get_rng_state() + bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states() + + # Set the states to what it used to be before the forward pass. + torch.set_rng_state(fwd_cpu_rng_state) + _set_cuda_rng_state(fwd_cuda_rng_state) + get_cuda_rng_tracker().set_states(fwd_cuda_rng_state_tracker) + + see_memory_usage("In backward checkpointing code before forward", force=False) + with torch.enable_grad(), torch.autograd.graph.saved_tensors_hooks(replay_pack, replay_unpack): + _unused = function(*detached_inputs) + + see_memory_usage("In backward checkpointing code after forward", force=False) + # Set the states back to what it was at the start of this function. + torch.set_rng_state(bwd_cpu_rng_state) + _set_cuda_rng_state(bwd_cuda_rng_state) + get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker) + + deepspeed_saved_tensors = None + non_tensor_args = None + tensor_flags = None + + if holder_from_backward not in storage: + raise RuntimeError("Attempt to retrieve a tensor saved by autograd multiple times without checkpoint" + " recomputation being triggered in between, this is not currently supported.") + + return storage[holder_from_backward] + + def after_backward_hook(_nonuse_grads): + """the hook registered to all leaf tensors""" + nonlocal leaf_tensors, backward_visited_leaf_nodes + backward_visited_leaf_nodes += 1 + + if backward_visited_leaf_nodes == len(leaf_tensors): + see_memory_usage("After backward checkpointing code after backward", force=False) + + if PROFILE_TIME: + timers('backward').stop() + timers.log(['backward']) + if SYNCHRONIZE: + get_accelerator().synchronize() + + with torch.autograd.graph.saved_tensors_hooks(checkpoint_pack, checkpoint_unpack): + outputs = function(*inputs_cuda) + if PROFILE_TIME or SYNCHRONIZE: + for leaf_tensor in leaf_tensors: + leaf_tensor.register_hook(after_backward_hook) + + see_memory_usage("After running forward on the layer", force=False) + + if PROFILE_TIME: + timers(FORWARD_GLOBAL_TIMER).stop() + timers.log([FORWARD_GLOBAL_TIMER]) + if SYNCHRONIZE: + get_accelerator().synchronize() + + all_outputs = [] + if torch.is_tensor(outputs): + all_outputs += [outputs] + else: + all_outputs += outputs + + if len(all_outputs) == 1: + return all_outputs[0] + else: + return tuple(all_outputs) + + +@compiler.disable # WA from Pytorch repo for compile + zero 3 accuracy issue +def checkpoint(function, *args): + """Checkpoint a model or part of the model. + This has been directly copied from torch.utils.checkpoint. """ + + all_outputs = [] + CheckpointFunction.apply(function, all_outputs, *args) + if len(all_outputs) == 1: + return all_outputs[0] + else: + return tuple(all_outputs) + + +def partition_activations_in_checkpoint(partition_activation): + global PARTITION_ACTIVATIONS + PARTITION_ACTIVATIONS = partition_activation + if dist.get_rank() == 0: + logger.info(f"**************Partition Activations {PARTITION_ACTIVATIONS}************") + + +def set_num_layers(nlayers): + global num_layers + num_layers = nlayers + + +def reset(): + """Resets memory buffers related to contiguous memory optimizations. + Should be called during eval when multiple forward propagations are + computed without any backward propagation that usually clears these + buffers. + Arguments: + None + + Return: + None + """ + if CONTIGUOUS_CHECKPOINTING: + global data_offsets, size_offsets + global contiguous_data_buffers, contiguous_size_buffers + + for buffers in contiguous_data_buffers: + buffers = [] + + # frees up all the pointers to the checkpoints except for the ones + # stored by save for backward + contiguous_data_buffers = [] + contiguous_size_buffers = [] + data_offsets = [] + size_offsets = [] + + +def _configure_using_config_file(config, mpu=None): + global num_layers, PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ + CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME + + config = DeepSpeedConfig(config, mpu=mpu).activation_checkpointing_config + if dist.get_rank() == 0: + logger.info(config.repr()) + PARTITION_ACTIVATIONS = config.partition_activations + CONTIGUOUS_CHECKPOINTING = config.contiguous_memory_optimization + num_layers = config.number_checkpoints + CPU_CHECKPOINT = config.cpu_checkpointing + SYNCHRONIZE = config.synchronize_checkpoint_boundary + PROFILE_TIME = config.profile + + +def _configure_defaults(): + + global mpu, num_layers, deepspeed_checkpointing_enabled + + global PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ + CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME + + PARTITION_ACTIVATIONS = False + CONTIGUOUS_CHECKPOINTING = False + num_layers = False + CPU_CHECKPOINT = False + SYNCHRONIZE = False + PROFILE_TIME = False + deepspeed_checkpointing_enabled = True + + +def configure( + mpu_, + deepspeed_config=None, + partition_activations=None, + contiguous_checkpointing=None, + num_checkpoints=None, + checkpoint_in_cpu=None, + synchronize=None, + profile=None, +): + """Configure DeepSpeed Activation Checkpointing. + + Arguments: + mpu_: Optional: An object that implements the following methods + get_model_parallel_rank/group/world_size, and get_data_parallel_rank/group/world_size + + deepspeed_config: Optional: DeepSpeed Config json file when provided will be used to + configure DeepSpeed Activation Checkpointing + + partition_activations: Optional: Partitions activation checkpoint across model parallel + GPUs when enabled. By default False. Will overwrite deepspeed_config if provided + + contiguous_checkpointing: Optional: Copies activation checkpoints to a contiguous memory + buffer. Works only with homogeneous checkpoints when partition_activations is enabled. + Must provide num_checkpoints. By default False. Will overwrite deepspeed_config if + provided + + num_checkpoints: Optional: Number of activation checkpoints stored during the forward + propagation of the model. Used to calculate the buffer size for contiguous_checkpointing + Will overwrite deepspeed_config if provided + + checkpoint_in_cpu: Optional: Moves the activation checkpoint to CPU. Only works with + partition_activation. Default is false. Will overwrite deepspeed_config if provided + + synchronize: Optional: Performs get_accelerator().synchronize() at the beginning and end of + each call to deepspeed.checkpointing.checkpoint for both forward and backward pass. + By default false. Will overwrite deepspeed_config if provided + + profile: Optional: Logs the forward and backward time for each + deepspeed.checkpointing.checkpoint invocation. Will overwrite deepspeed_config + if provided + + Returns: + None + """ + global mpu, num_layers, deepspeed_checkpointing_enabled + + global PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \ + CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME + + _configure_defaults() + + if mpu_ is not None: + mpu = mpu_ + + if deepspeed_config is not None: + _configure_using_config_file(deepspeed_config, mpu=mpu) + + if partition_activations is not None: + PARTITION_ACTIVATIONS = partition_activations + + if contiguous_checkpointing is not None: + CONTIGUOUS_CHECKPOINTING = contiguous_checkpointing + + if num_checkpoints is not None: + num_layers = num_checkpoints + + if checkpoint_in_cpu is not None: + CPU_CHECKPOINT = checkpoint_in_cpu + + if synchronize is not None: + SYNCHRONIZE = synchronize + + if profile is not None: + PROFILE_TIME = profile + + if CONTIGUOUS_CHECKPOINTING: + assert PARTITION_ACTIVATIONS, "Contiguous Checkpointing is only available with partitioned activations. Set partitioned activations to true in deepspeed config" + if CONTIGUOUS_CHECKPOINTING: + assert num_layers is not None, "Must specify the number of layers with contiguous memory checkpointing" + + global mp_rank, mp_size, mp_group + + if mpu is not None: + if hasattr(mpu, 'get_tensor_model_parallel_rank'): + mp_rank = mpu.get_tensor_model_parallel_rank() + mp_size = mpu.get_tensor_model_parallel_world_size() + mp_group = mpu.get_tensor_model_parallel_group() + else: + mp_rank = mpu.get_model_parallel_rank() + mp_size = mpu.get_model_parallel_world_size() + mp_group = mpu.get_model_parallel_group() + + #print configuration only once + see_memory_usage("After configuration", force=False) + if dist.get_rank() == 0: + logger.info(f"Activation Checkpointing Information") + logger.info(f"----Partition Activations {PARTITION_ACTIVATIONS}, CPU CHECKPOINTING {CPU_CHECKPOINT}") + logger.info(f"----contiguous Memory Checkpointing {CONTIGUOUS_CHECKPOINTING} with {num_layers} total layers") + logger.info(f"----Synchronization {SYNCHRONIZE}") + logger.info(f"----Profiling time in checkpointing {PROFILE_TIME}") + + +def is_configured(): + """True if deepspeed activation checkpointing has been configured + by calling deepspeed.checkpointing.configure, else returns false + + Arguments: + None + + Return: + True of configured, else False + """ + return deepspeed_checkpointing_enabled diff --git a/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/config.py b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/config.py new file mode 100644 index 0000000000000000000000000000000000000000..dc07388a95da039b50bad87b4aa57b12f4e41f6f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/activation_checkpointing/config.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject + +######################################### +# DeepSpeed Activation Checkpointing +######################################### +# Activation Checkpointing Allows to save memory by only keeping a select few +#activations for the backpropagation. +ACTIVATION_CHKPT_FORMAT = ''' +Activation Checkpointing should be configured as: +"session_params": { + "activation_checkpointing": { + "partitioned_activations": [true|false], + "number_checkpoints": 100, + "contiguous_memory_optimization": [true|false], + "cpu_checkpointing": [true|false], + "profile": [true|false], + "synchronize_checkpoint_boundary": [true|false], + } +} +''' + +ACT_CHKPT_PARTITION_ACTIVATIONS = 'partition_activations' +ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT = False + +ACT_CHKPT_NUMBER_CHECKPOINTS = 'number_checkpoints' +ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT = None + +ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION = 'contiguous_memory_optimization' +ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT = False + +ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY = 'synchronize_checkpoint_boundary' +ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT = False + +ACT_CHKPT_PROFILE = 'profile' +ACT_CHKPT_PROFILE_DEFAULT = False + +ACT_CHKPT_CPU_CHECKPOINTING = 'cpu_checkpointing' +ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT = False + +ACT_CHKPT = 'activation_checkpointing' + +ACT_CHKPT_DEFAULT = { + ACT_CHKPT_PARTITION_ACTIVATIONS: ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT, + ACT_CHKPT_NUMBER_CHECKPOINTS: ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT, + ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION: ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT, + ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY: ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT, + ACT_CHKPT_PROFILE: ACT_CHKPT_PROFILE_DEFAULT, + ACT_CHKPT_CPU_CHECKPOINTING: ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT +} + + +class DeepSpeedActivationCheckpointingConfig(DeepSpeedConfigObject): + + def __init__(self, param_dict): + super(DeepSpeedActivationCheckpointingConfig, self).__init__() + + self.partition_activations = None + self.contiguous_memory_optimization = None + self.cpu_checkpointing = None + self.number_checkpoints = None + self.synchronize_checkpoint_boundary = None + self.profile = None + + if ACT_CHKPT in param_dict.keys(): + act_chkpt_config_dict = param_dict[ACT_CHKPT] + else: + act_chkpt_config_dict = ACT_CHKPT_DEFAULT + + self._initialize(act_chkpt_config_dict) + + def _initialize(self, act_chkpt_config_dict): + self.partition_activations = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_PARTITION_ACTIVATIONS, + ACT_CHKPT_PARTITION_ACTIVATIONS_DEFAULT) + + self.contiguous_memory_optimization = get_scalar_param(act_chkpt_config_dict, + ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION, + ACT_CHKPT_CONTIGUOUS_MEMORY_OPTIMIZATION_DEFAULT) + + self.cpu_checkpointing = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_CPU_CHECKPOINTING, + ACT_CHKPT_CPU_CHECKPOINTING_DEFAULT) + + self.number_checkpoints = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_NUMBER_CHECKPOINTS, + ACT_CHKPT_NUMBER_CHECKPOINTS_DEFAULT) + + self.profile = get_scalar_param(act_chkpt_config_dict, ACT_CHKPT_PROFILE, ACT_CHKPT_PROFILE_DEFAULT) + + self.synchronize_checkpoint_boundary = get_scalar_param(act_chkpt_config_dict, + ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY, + ACT_CHKPT_SYNCHRONIZE_CHECKPOINT_BOUNDARY_DEFAULT) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/base_optimizer.py b/lib/python3.12/site-packages/deepspeed/runtime/base_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..d2c54155da89539d878b663ff4eeedc9aa311f53 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/base_optimizer.py @@ -0,0 +1,63 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import torch + +from deepspeed.utils import logger +from deepspeed.utils.tensor_fragment import map_to_flat_opt_states +from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank + + +class DeepSpeedOptimizer(object): + pass + + +class ZeROOptimizer(DeepSpeedOptimizer): + + def load_hp_checkpoint_state_from_checkpoint_dir(self, lp_groups_name: str, checkpoint_dir: str) -> None: + checkpoint_dir = os.path.join(checkpoint_dir, "zero") + optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt") + assert os.path.isfile( + optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.' + optim_sd = torch.load(optim_state_path, weights_only=False) + + self._load_global_state(optim_sd) + + tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) + if self.mpu is None: + logger.warning("MPU is not provided, setting tp size to 1 in checkpoint loading.") + tp_world_size = 1 + else: + tp_world_size = self.mpu.get_slice_parallel_world_size() if hasattr(self.mpu, "get_slice_parallel_world_size") \ + else self.mpu.get_tensor_model_parallel_world_size() + + for i, (param_group, + loaded_param_group) in enumerate(zip(self.optimizer.param_groups, optim_sd['param_groups'])): + # We have an assumption that all params in the same param_group have the same keys + opt_keys = set() + steps = [] + + lp_groups = getattr(self, lp_groups_name) + for lp in lp_groups[i]: + if lp._hp_mapping is not None: + #print(f"Loading {self.param_names[lp]} {tp_rank=} {tp_world_size=}") + step = lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank, + tp_world_size) + for key in lp._hp_mapping.get_optim_state_keys(): + opt_keys.add(key) + steps.append(step) + + hp_param = param_group['params'][0] + assert all(step == steps[0] for step in steps), f"Steps {steps} are not equal" + if steps[0] is not None: + self.optimizer.state[hp_param]['step'] = steps[0] + + map_to_flat_opt_states(hp_param, lp_groups[i], self.optimizer.state, opt_keys) + + for key, value in loaded_param_group.items(): + if key == 'params': + continue + param_group[key] = value diff --git a/lib/python3.12/site-packages/deepspeed/runtime/bf16_optimizer.py b/lib/python3.12/site-packages/deepspeed/runtime/bf16_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c9337a795aeab6bda03f75f00f83583d6ac7091f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/bf16_optimizer.py @@ -0,0 +1,564 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from collections import OrderedDict +import torch +import sys +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from deepspeed import comm as dist +from deepspeed.runtime.constants import PIPE_REPLICATED +from deepspeed.runtime.base_optimizer import ZeROOptimizer +from packaging import version as pkg_version +from deepspeed.git_version_info import version +from deepspeed.runtime.utils import (get_global_norm_of_tensors, clip_tensors_by_global_norm, DummyOptim, + align_dense_tensors, all_gather_dp_groups, is_model_parallel_parameter, + see_memory_usage, graph_process, get_norm_with_moe_layers) +from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state, fragment_address, groups +from deepspeed.moe.utils import is_moe_param, is_moe_param_group +from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank +from deepspeed.utils.torch import register_grad_hook +from deepspeed.checkpoint import enable_universal_checkpoint +from deepspeed.checkpoint.constants import (DS_VERSION, PARTITION_COUNT, BASE_OPTIMIZER_STATE, + SINGLE_PARTITION_OF_FP32_GROUPS, CLIP_GRAD, GROUP_PADDINGS, + PARAM_SLICE_MAPPINGS) + +setattr(sys.modules[__name__], 'fragment_address', fragment_address) + + +def print_rank_0(message, debug=False, force=False): + if dist.get_rank() == 0 and (debug or force): + print(message) + + +class BF16_Optimizer(ZeROOptimizer): + + def __init__(self, + init_optimizer, + param_names, + mpu=None, + clip_grad=0.0, + norm_type=2, + allgather_bucket_size=5000000000, + dp_process_group=None, + timers=None, + grad_acc_dtype=None, + graph_harvesting=False, + immediate_grad_update=True, + has_moe_layers=False): + super().__init__() + see_memory_usage('begin bf16_optimizer', force=True) + self.timers = timers + self.optimizer = init_optimizer + self.param_names = param_names + self.using_real_optimizer = not isinstance(self.optimizer, DummyOptim) + + assert grad_acc_dtype in [torch.float32, torch.bfloat16 + ], f"BF16Optimizer: Unsupported gradient accumulation data type: {grad_acc_dtype}" + self.grad_acc_dtype = grad_acc_dtype + self.immediate_grad_update = immediate_grad_update + + self.clip_grad = clip_grad + self.norm_type = norm_type + self.mpu = mpu + self.allgather_bucket_size = int(allgather_bucket_size) + self.dp_process_group = dp_process_group + self.dp_rank = dist.get_rank(group=self.dp_process_group) + self.has_moe_layers = has_moe_layers + self.non_expert_gradients = [] + self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))] + if self.has_moe_layers: + self._configure_moe_settings() + + # Use torch (un)flatten ops + self.flatten = _flatten_dense_tensors + self.unflatten = _unflatten_dense_tensors + + #align nccl all-gather send buffers to 4-bye boundary + self.nccl_start_alignment_factor = 2 # 4-byte alignment/sizeof(fp16) = 2 + + # Build BF16/FP32 groups + self.bf16_groups = [] + self.bf16_groups_flat = [] + self.bf16_partitioned_groups = [] + + self.fp32_groups_flat_partition = [] + + # Maintain different fp32 gradients views for convenience + self.fp32_groups_gradients = [] + self.fp32_groups_gradient_dict = {} + self.fp32_groups_gradients_flat = [] + self.fp32_groups_actual_gradients_flat = [] + self.fp32_groups_gradient_flat_partition = [] + self.fp32_groups_has_gradients = [] + + self.group_paddings = [] + self.graph_harvesting = graph_harvesting + if self.using_real_optimizer: + self._setup_for_real_optimizer() + + see_memory_usage('end bf16_ optimizer', force=True) + + def destroy(self): + for i, _ in enumerate(self.optimizer.param_groups): + for p in self.bf16_groups[i]: + if getattr(p, '_hp_mapping', None): + p._hp_mapping = None + for hook in self._grad_acc_hooks: + hook.remove() + print_rank_0("Removed grad acc hooks") + + def _configure_moe_settings(self): + assert any( + [is_moe_param_group(group) for group in self.optimizer.param_groups] + ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer" + + for i, group in enumerate(self.optimizer.param_groups): + if is_moe_param_group(group): + assert all([is_moe_param(param) + for param in group['params']]), "All params in MoE group must be MoE params" + self.real_dp_process_group[i] = groups._get_expert_data_parallel_group(group['name']) + self.expert_gradients = {} + if self.has_moe_layers: + for key in groups._get_expert_data_parallel_group_dict().keys(): + self.expert_gradients[key] = [] + + def _setup_for_real_optimizer(self): + self.partition_count = [dist.get_world_size(group=pg) for pg in self.real_dp_process_group] + + for i, param_group in enumerate(self.optimizer.param_groups): + real_dp_world_size = dist.get_world_size(group=self.real_dp_process_group[i]) + see_memory_usage(f'before initializing group {i}', force=True) + + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + + # grab the original list + trainable_parameters = [param for param in param_group['params'] if param.requires_grad] + self.bf16_groups.append(trainable_parameters) + + # create flat bf16 params + self.bf16_groups_flat.append( + self._flatten_dense_tensors_aligned(self.bf16_groups[i], + self.nccl_start_alignment_factor * real_dp_world_size)) + # Make bf16 params point to flat tensor storage + self._update_storage_to_flattened_tensor(tensor_list=self.bf16_groups[i], + flat_tensor=self.bf16_groups_flat[i]) + + # divide flat weights into equal sized partitions + partition_size = self.bf16_groups_flat[i].numel() // real_dp_world_size + bf16_dp_partitions = [ + self.bf16_groups_flat[i].narrow(0, dp_index * partition_size, partition_size) + for dp_index in range(real_dp_world_size) + ] + self.bf16_partitioned_groups.append(bf16_dp_partitions) + + # create fp32 params partition + self.fp32_groups_flat_partition.append(bf16_dp_partitions[partition_id].clone().float().detach()) + self.fp32_groups_flat_partition[i].requires_grad = True + + num_elem_list = [t.numel() for t in self.bf16_groups[i]] + + # create fp32 gradients + fp32_flat_buffer = torch.zeros_like(self.bf16_groups_flat[i], dtype=self.grad_acc_dtype) + self.fp32_groups_gradients_flat.append(fp32_flat_buffer) + if self.has_moe_layers and is_moe_param_group(param_group): + self.expert_gradients[param_group['name']].append(fp32_flat_buffer) + else: + self.non_expert_gradients.append(fp32_flat_buffer) + + # track individual fp32 gradients for entire model + fp32_gradients = self._split_flat_tensor(flat_tensor=self.fp32_groups_gradients_flat[i], + num_elem_list=num_elem_list) + self.fp32_groups_gradients.append(fp32_gradients) + self.fp32_groups_gradient_dict[i] = fp32_gradients + + # flat tensor corresponding to actual fp32 gradients (i.e., minus alignment padding) + length_without_padding = sum(num_elem_list) + self.fp32_groups_actual_gradients_flat.append( + torch.narrow(self.fp32_groups_gradients_flat[i], 0, 0, length_without_padding)) + + # flat tensor corresponding to gradient partition + self.fp32_groups_gradient_flat_partition.append( + torch.narrow(self.fp32_groups_gradients_flat[i], 0, partition_id * partition_size, partition_size)) + + # track fp32 gradient updates + self.fp32_groups_has_gradients.append([False] * len(self.bf16_groups[i])) + + # Record padding required for alignment + if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1: + padding = self.bf16_groups_flat[i].numel() - length_without_padding + else: + padding = 0 + + self.group_paddings.append(padding) + + # update optimizer param groups to reference fp32 params partition + param_group['params'] = [self.fp32_groups_flat_partition[i]] + + see_memory_usage(f'after initializing group {i}', force=True) + + self._grad_acc_hooks = [] + if self.immediate_grad_update: + self.create_grad_acc_hooks() + + # Need optimizer states initialized before linking lp to optimizer state + self._link_all_hp_params() + self._hp_optimizer_states_linked = False + self._enable_universal_checkpoint() + self._param_slice_mappings = self._create_param_mapping() + + def _enable_universal_checkpoint(self): + for lp_param_group in self.bf16_groups: + enable_universal_checkpoint(param_list=lp_param_group) + + def _create_param_mapping(self): + param_mapping = [] + for i, _ in enumerate(self.optimizer.param_groups): + param_mapping_per_group = OrderedDict() + for lp in self.bf16_groups[i]: + if lp._hp_mapping is not None: + lp_name = self.param_names[lp] + param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address() + param_mapping.append(param_mapping_per_group) + + return param_mapping + + def _link_all_hp_params(self): + for i, _ in enumerate(self.optimizer.param_groups): + real_dp_world_size = dist.get_world_size(group=self.real_dp_process_group[i]) + + # Link bf16 and fp32 params in partition + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + partition_size = self.bf16_groups_flat[i].numel() // real_dp_world_size + flat_hp_partition = self.fp32_groups_flat_partition[i] + link_hp_params(lp_param_list=self.bf16_groups[i], + flat_hp_partition=flat_hp_partition, + gradient_dict=self.fp32_groups_gradient_dict, + offload_gradient_dict=None, + use_offload=False, + param_group_index=i, + partition_start=partition_id * partition_size, + partition_size=partition_size, + dp_group=self.real_dp_process_group[i]) + + def _lazy_init_hp_params_optimizer_state(self): + if not self._hp_optimizer_states_linked: + for i, _ in enumerate(self.optimizer.param_groups): + lazy_init_hp_params_optimizer_state(self.bf16_groups[i], self.fp32_groups_flat_partition[i], + self.optimizer.state) + self._hp_optimizer_states_linked = True + + def _split_flat_tensor(self, flat_tensor, num_elem_list): + assert sum(num_elem_list) <= flat_tensor.numel() + tensor_list = [] + offset = 0 + for num_elem in num_elem_list: + dense_tensor = torch.narrow(flat_tensor, 0, offset, num_elem) + tensor_list.append(dense_tensor) + offset += num_elem + + return tensor_list + + def _update_storage_to_flattened_tensor(self, tensor_list, flat_tensor): + updated_params = self.unflatten(flat_tensor, tensor_list) + for p, q in zip(tensor_list, updated_params): + p.data = q.data + + def _flatten_dense_tensors_aligned(self, tensor_list, alignment): + return self.flatten(align_dense_tensors(tensor_list, alignment)) + + @torch.no_grad() + def step(self, closure=None): + if closure is not None: + raise NotImplementedError(f'{self.__class__} does not support closure.') + + non_expert_grads_for_norm, expert_grads_for_norm = self.get_grads_for_norm() + non_expert_groups_norm = get_global_norm_of_tensors(input_tensors=non_expert_grads_for_norm, + mpu=self.mpu, + norm_type=self.norm_type, + use_graph=self.graph_harvesting) + all_groups_norm = non_expert_groups_norm + if self.has_moe_layers: + all_groups_norm = get_norm_with_moe_layers(non_expert_groups_norm, + mpu=self.mpu, + expert_tensors=expert_grads_for_norm, + norm_type=self.norm_type) + + self._global_grad_norm = all_groups_norm + + assert all_groups_norm > 0. + if self.clip_grad > 0.: + clip_tensors_by_global_norm(input_tensors=self.get_grads_for_norm(for_clipping=True), + max_norm=self.clip_grad, + global_norm=all_groups_norm, + mpu=self.mpu, + use_graph=self.graph_harvesting) + + for param_partition, grad_partition in zip(self.fp32_groups_flat_partition, + self.fp32_groups_gradient_flat_partition): + # In case of grad acc dtype different than FP32, need to cast to high precision. + param_partition.grad = grad_partition.to( + param_partition.dtype) if grad_partition.dtype != param_partition.dtype else grad_partition + + self.optimizer.step() + + if self.grad_acc_dtype is not torch.float32: + for param_partition in self.fp32_groups_flat_partition: + param_partition.grad = None + + # We need to link optimizer state after the first step() call + self._lazy_init_hp_params_optimizer_state() + + self.update_lp_params() + + self.clear_hp_grads() + + def backward(self, loss, retain_graph=False, update_hp_grads=True, clear_lp_grads=False, **bwd_kwargs): + """Perform a backward pass and copy the low-precision gradients to the + high-precision copy. + + We copy/accumulate to the high-precision grads now to prevent accumulating in the + bf16 grads after successive backward() calls (i.e., grad accumulation steps > 1) + + The low-precision grads are deallocated during this procedure. + """ + self.clear_lp_grads() + loss.backward(retain_graph=retain_graph, **bwd_kwargs) + + if update_hp_grads: + self.update_hp_grads(clear_lp_grads=clear_lp_grads) + + @torch.no_grad() + def _update_hp_grad(self, lp, group_idx, param_idx, clear_lp_grads): + if lp.grad is None: + return + + hp_grad = self.fp32_groups_gradients[group_idx][param_idx] + assert hp_grad is not None, \ + f'high precision param has no gradient, lp param_id = {id(lp)} group_info = [{group_idx}][{param_idx}]' + + hp_grad.data.add_(lp.grad.data.to(hp_grad.dtype).view(hp_grad.shape)) + lp._hp_grad = hp_grad + self.fp32_groups_has_gradients[group_idx][param_idx] = True + + # clear gradients + if clear_lp_grads: + lp.grad.zero_() + + @torch.no_grad() + def _update_hp_grads_func(self, clear_lp_grads=False): + for i, group in enumerate(self.bf16_groups): + for j, lp in enumerate(group): + self._update_hp_grad(lp, i, j, clear_lp_grads) + + @torch.no_grad() + def update_hp_grads(self, clear_lp_grads=False): + if self.immediate_grad_update: + return + + if self.graph_harvesting: + graph_process(False, self._update_hp_grads_func, clear_lp_grads) + else: + self._update_hp_grads_func(clear_lp_grads) + #cpu op + for i, group in enumerate(self.bf16_groups): + for j, lp in enumerate(group): + if lp.grad is None: + continue + self.fp32_groups_has_gradients[i][j] = True + + @torch.no_grad() + def get_grads_for_reduction(self): + if self.has_moe_layers: + return self.non_expert_gradients, self.expert_gradients + return self.non_expert_gradients, {} + + @torch.no_grad() + def get_grads_for_norm(self, for_clipping=False): + """ + Returns: + tuple[list[Tensor], dict[ep_name, List[Tensor]] | list: + If for_clipping, return all gradients. + Otherwise, separate and return dict of expert_grad and list of non_expert_grad + """ + # (grads, expert_group_name) + expert_grads_for_norm = {} + + # grads + non_expert_grads_for_norm = [] + all_grads_for_clip = [] + + tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu) + assert len(self.bf16_groups) == len(self.optimizer.param_groups) + for i, group in enumerate(self.bf16_groups): + for j, lp in enumerate(group): + if not for_clipping: + if hasattr(lp, PIPE_REPLICATED) and lp.ds_pipe_replicated: + continue + + # skip duplicated parameters. perform norm only on cards with tp_rank=0. + # non-duplicated parameters include: + # - Parameters with tp: Use allreducesum of mp_group. + # - Moe Parameters with ep: Use allreducesum of ep_group. + if not (tensor_mp_rank == 0 or is_model_parallel_parameter(lp) or is_moe_param(lp)): + continue + + if not self.fp32_groups_has_gradients[i][j]: + continue + if not for_clipping: + param_group = self.optimizer.param_groups[i] + if self.has_moe_layers and is_moe_param_group(param_group): + if param_group['name'] not in expert_grads_for_norm: + expert_grads_for_norm[param_group['name']] = [] + expert_grads_for_norm[param_group['name']].append(self.fp32_groups_gradients[i][j]) + else: + non_expert_grads_for_norm.append(self.fp32_groups_gradients[i][j]) + else: + all_grads_for_clip.append(self.fp32_groups_gradients[i][j]) + if not for_clipping: + return non_expert_grads_for_norm, expert_grads_for_norm + return all_grads_for_clip + + @torch.no_grad() + def update_lp_params(self): + for i, (bf16_partitions, + fp32_partition) in enumerate(zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition)): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + bf16_partitions[partition_id].data.copy_(fp32_partition.data) + + all_gather_dp_groups(groups_flat=self.bf16_groups_flat, + partitioned_param_groups=self.bf16_partitioned_groups, + dp_process_group=self.real_dp_process_group, + start_alignment_factor=self.nccl_start_alignment_factor, + allgather_bucket_size=self.allgather_bucket_size) + + def clear_hp_grads(self): + for flat_gradients in self.fp32_groups_gradients_flat: + flat_gradients.zero_() + + for i, group in enumerate(self.fp32_groups_gradients): + self.fp32_groups_has_gradients[i] = [False] * len(group) + + def clear_lp_grads(self, set_to_none=False): + + # using zero_() fixed memory address for graph replay + if self.graph_harvesting: + assert not set_to_none, "graph harvesting is incompatible with setting lp grads to None" + + zero_grads_list = [] + for group in self.bf16_groups: + for param in group: + if set_to_none: + param.grad = None + elif param.grad is not None: + if param.grad.grad_fn is not None: + param.grad.detach_() + zero_grads_list.append(param.grad) + if not set_to_none and len(zero_grads_list) > 0: + torch._foreach_zero_(zero_grads_list) + + def zero_grad(self, set_to_none=True): + self.clear_lp_grads(set_to_none) + self.clear_hp_grads() + + def state_dict(self): + state_dict = {} + state_dict[CLIP_GRAD] = self.clip_grad + state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict() + state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = self.fp32_groups_flat_partition + state_dict[GROUP_PADDINGS] = self.group_paddings + state_dict[PARTITION_COUNT] = self.partition_count + state_dict[DS_VERSION] = version + state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings + + return state_dict + + # Restore base optimizer fp32 weights bfloat16 weights + def _restore_from_bit16_weights(self): + for i, (bf16_partitions, + fp32_partition) in enumerate(zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition)): + partition_id = dist.get_rank(group=self.real_dp_process_group[i]) + fp32_partition.data.copy_(bf16_partitions[partition_id].data) + + def refresh_fp32_params(self): + self._restore_from_bit16_weights() + + def load_state_dict(self, + state_dict_list, + checkpoint_folder=None, + load_optimizer_states=True, + load_from_fp32_weights=False, + load_serial=None, + param_shapes=None): + if checkpoint_folder: + self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights) + else: + self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights) + + def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False): + + dp_rank = dist.get_rank(group=self.dp_process_group) + current_rank_sd = state_dict_list[dp_rank] + + ckpt_version = current_rank_sd.get(DS_VERSION, False) + assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed" + ckpt_version = pkg_version.parse(ckpt_version) + + self.clip_grad = current_rank_sd.get(CLIP_GRAD, self.clip_grad) + + if load_optimizer_states: + print(f"_load_legacy_checkpoint current_rank_sd[BASE_OPTIMIZER_STATE]") + self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE]) + + if load_from_fp32_weights: + for current, saved in zip(self.fp32_groups_flat_partition, + current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]): + src_tensor = _get_padded_tensor(saved, current.numel()) + current.data.copy_(src_tensor.data) + + if load_optimizer_states: + self._link_all_hp_params() + + def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights): + self.load_hp_checkpoint_state_from_checkpoint_dir("bf16_groups", checkpoint_folder) + + def _load_global_state(self, sd): + pass + + @property + def param_groups(self): + """Forward the wrapped optimizer's parameters.""" + return self.optimizer.param_groups + + @property + def state(self): + """Forward the wrapped optimizer's states.""" + return self.optimizer.state + + def accumulate_hp_grads_and_remove_lp(self, lp_param, group_idx, param_idx): + assert self.immediate_grad_update + self._update_hp_grad(lp_param, group_idx, param_idx, clear_lp_grads=False) + + def create_grad_acc_hooks(self): + for i, param_group in enumerate(self.bf16_groups): + for j, param in enumerate(param_group): + if param.requires_grad: + + def wrapper(param, i, j): + + def accumulate_hp_grads_and_remove_lp(*notneeded): + self.accumulate_hp_grads_and_remove_lp(param, i, j) + + self._grad_acc_hooks.append(register_grad_hook(param, accumulate_hp_grads_and_remove_lp)) + + wrapper(param, i, j) + + +def _get_padded_tensor(src_tensor, size): + if src_tensor.numel() >= size: + return src_tensor + padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device) + slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel()) + slice_tensor.data.copy_(src_tensor.data) + return padded_tensor diff --git a/lib/python3.12/site-packages/deepspeed/runtime/compiler.py b/lib/python3.12/site-packages/deepspeed/runtime/compiler.py new file mode 100644 index 0000000000000000000000000000000000000000..be778b83f8bb8f3b9416fdf1d5717532367a8e4b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/compiler.py @@ -0,0 +1,30 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.utils.torch import required_torch_version + +try: + from torch.compiler import is_compiling as torch_is_compiling +except ImportError: + try: + from torch._dynamo.external_utils import is_compiling as torch_is_compiling + except ImportError: + # Torch does not have compiler support + torch_is_compiling = lambda: False + + +def is_compile_supported(): + return required_torch_version(min_version=2.1) + + +def disable(func): + if is_compile_supported(): + return torch.compiler.disable(func) + return func + + +def is_compiling(): + return torch_is_compiling() diff --git a/lib/python3.12/site-packages/deepspeed/runtime/compression/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/compression/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/compression/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8f60246cde7d28c37287efffb0b82d12f432abf Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/compression/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35c9ecd423cc1006200222063b1118ee5301e2d3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/compression/__pycache__/cupy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/compression/cupy.py b/lib/python3.12/site-packages/deepspeed/runtime/compression/cupy.py new file mode 100644 index 0000000000000000000000000000000000000000..7133ac04ed2b65bd496656870aebbf95ac04d3ad --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/compression/cupy.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import cupy +from torch.utils.dlpack import to_dlpack +from torch.utils.dlpack import from_dlpack + + +class CupyBackend(object): + + def __init__(self): + pass + + def torch2cupy(self, tensor): + return cupy.from_dlpack(to_dlpack(tensor)) + + def cupy2torch(self, cupy_tensor): + return from_dlpack(cupy_tensor) + + def compress_by_chunk(self, cupy_bool_tensor, num_chunks): + packed_sign = cupy.packbits(cupy_bool_tensor) + sign_list_packed = cupy.split(packed_sign, num_chunks) + cupy.cuda.get_current_stream().synchronize() + return sign_list_packed diff --git a/lib/python3.12/site-packages/deepspeed/runtime/config.py b/lib/python3.12/site-packages/deepspeed/runtime/config.py new file mode 100644 index 0000000000000000000000000000000000000000..eb4c042447178465f9edf87b07bb41bca2ad61a7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/config.py @@ -0,0 +1,1057 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +from typing import Union +from enum import Enum + +import torch +import json +import hjson +import copy +import base64 + +from .constants import * +from .fp16.loss_scaler import ( + INITIAL_LOSS_SCALE, + SCALE_WINDOW, + DELAYED_SHIFT, + CONSECUTIVE_HYSTERESIS, + MIN_LOSS_SCALE, +) +from .config_utils import ( + get_scalar_param, + dict_raise_error_on_duplicate_keys, + ScientificNotationEncoder, +) +from .zero.config import get_zero_config, ZeroStageEnum +from .activation_checkpointing.config import DeepSpeedActivationCheckpointingConfig +from ..comm.config import DeepSpeedCommsConfig +from ..monitor.config import get_monitor_config +from ..inference.config import WeightQuantConfig +from ..compile.config import CompileConfig + +from deepspeed import comm as dist +from deepspeed.runtime.config_utils import DeepSpeedConfigModel + +from ..git_version_info import version as __version__ +from ..utils import logger + +from ..elasticity import ( + elasticity_enabled, + compute_elastic_config, + ensure_immutable_elastic_config, +) +from ..elasticity.config import ElasticityConfigError +from ..elasticity.constants import ( + ELASTICITY, + IGNORE_NON_ELASTIC_BATCH_INFO, + IGNORE_NON_ELASTIC_BATCH_INFO_DEFAULT, + MODEL_PARALLEL_SIZE, + MODEL_PARALLEL_SIZE_DEFAULT, + NUM_GPUS_PER_NODE, + NUM_GPUS_PER_NODE_DEFAULT, +) + +from ..profiling.config import DeepSpeedFlopsProfilerConfig +from ..autotuning.config import DeepSpeedAutotuningConfig +from ..nebula.config import DeepSpeedNebulaConfig + +from ..compression.config import get_compression_config, get_quantize_enabled +from ..compression.constants import * +from .swap_tensor.aio_config import get_aio_config + +from .tensor_parallel import get_tensor_parallel_config +from .data_pipeline.config import get_data_efficiency_enabled, get_data_efficiency_config, get_curriculum_enabled_legacy, get_curriculum_params_legacy +from .data_pipeline.constants import * + +from ..utils.config import get_timers_config + +TENSOR_CORE_ALIGN_SIZE = 8 + +ADAGRAD_OPTIMIZER = 'adagrad' +ADAM_OPTIMIZER = 'adam' +ADAMW_OPTIMIZER = 'adamw' +LAMB_OPTIMIZER = 'lamb' +ONEBIT_ADAM_OPTIMIZER = 'onebitadam' +ZERO_ONE_ADAM_OPTIMIZER = 'zerooneadam' +ONEBIT_LAMB_OPTIMIZER = 'onebitlamb' +MUADAM_OPTIMIZER = 'muadam' +MUADAMW_OPTIMIZER = 'muadamw' +MUSGD_OPTIMIZER = 'musgd' +LION_OPTIMIZER = 'lion' +DEEPSPEED_OPTIMIZERS = [ + ADAGRAD_OPTIMIZER, ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, + ZERO_ONE_ADAM_OPTIMIZER, MUADAM_OPTIMIZER, MUADAMW_OPTIMIZER, MUSGD_OPTIMIZER, LION_OPTIMIZER +] + +# extra optimizer parameters for adam/adamw +TORCH_ADAM_PARAM = "torch_adam" + +# default to adamw logic for adam/adamw optimizers unless user explicitly opts out +ADAM_W_MODE = "adam_w_mode" +ADAM_W_MODE_DEFAULT = True + + +class DeepSpeedConfigError(Exception): + pass + + +class DtypeEnum(Enum): + # The torch dtype must always be the first value (so we return torch.dtype) + fp16 = torch.float16, "torch.float16", "fp16", "float16", "half" + fp32 = torch.float32, "torch.float32", "fp32", "float32", "float" + int8 = torch.int8, "torch.int8", "int8" + bf16 = torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16" + + # Copied from https://stackoverflow.com/a/43210118 + # Allows us to use multiple values for each Enum index and returns first + # listed value when Enum is called + def __new__(cls, *values): + obj = object.__new__(cls) + # first value is canonical value + obj._value_ = values[0] + for other_value in values[1:]: + cls._value2member_map_[other_value] = obj + obj._all_values = values + return obj + + def __repr__(self): + return "<%s.%s: %s>" % ( + self.__class__.__name__, + self._name_, + ", ".join([repr(v) for v in self._all_values]), + ) + + +def get_pld_enabled(param_dict): + if PROGRESSIVE_LAYER_DROP in param_dict.keys(): + return get_scalar_param(param_dict[PROGRESSIVE_LAYER_DROP], PLD_ENABLED, PLD_ENABLED_DEFAULT) + else: + return False + + +def get_pld_params(param_dict): + if PROGRESSIVE_LAYER_DROP in param_dict.keys(): + pld_params = copy.copy(param_dict[PROGRESSIVE_LAYER_DROP]) + pld_params.pop(PLD_ENABLED) + return pld_params + else: + return False + + +def get_amp_enabled(param_dict): + if AMP in param_dict.keys(): + return get_scalar_param(param_dict[AMP], AMP_ENABLED, AMP_ENABLED_DEFAULT) + else: + return False + + +def get_amp_params(param_dict): + if AMP in param_dict.keys(): + amp_params = copy.copy(param_dict[AMP]) + amp_params.pop(AMP_ENABLED) + return amp_params + else: + return False + + +def get_fp16_enabled(param_dict): + if FP16 in param_dict.keys(): + return get_scalar_param(param_dict[FP16], FP16_ENABLED, FP16_ENABLED_DEFAULT) + else: + return False + + +def get_bfloat16_enabled(param_dict): + for key in [BFLOAT16, BFLOAT16_OLD]: + if key in param_dict.keys(): + return get_scalar_param(param_dict[key], BFLOAT16_ENABLED, BFLOAT16_ENABLED_DEFAULT) + return False + + +def get_bfloat16_immediate_grad_update(param_dict): + for key in [BFLOAT16, BFLOAT16_OLD]: + if key in param_dict.keys(): + return get_scalar_param(param_dict[key], BFLOAT16_IMMEDIATE_GRAD_UPDATE, + BFLOAT16_IMMEDIATE_GRAD_UPDATE_DEFAULT) + return False + + +def get_fp16_master_weights_and_grads_enabled(param_dict): + if get_fp16_enabled(param_dict): + return get_scalar_param(param_dict[FP16], FP16_MASTER_WEIGHTS_AND_GRADS, FP16_MASTER_WEIGHTS_AND_GRADS_DEFAULT) + else: + return False + + +def get_fp16_auto_cast(param_dict): + if get_fp16_enabled(param_dict): + return get_scalar_param(param_dict[FP16], FP16_AUTO_CAST, FP16_AUTO_CAST_DEFAULT) + + +def get_loss_scale(param_dict): + if get_fp16_enabled(param_dict): + return get_scalar_param(param_dict[FP16], FP16_LOSS_SCALE, FP16_LOSS_SCALE_DEFAULT) + elif get_bfloat16_enabled(param_dict): + return 1.0 + else: + return FP16_LOSS_SCALE_DEFAULT + + +def get_initial_dynamic_scale(param_dict): + if get_fp16_enabled(param_dict): + initial_scale_power = get_scalar_param(param_dict[FP16], FP16_INITIAL_SCALE_POWER, + FP16_INITIAL_SCALE_POWER_DEFAULT) + elif get_bfloat16_enabled(param_dict): + initial_scale_power = 0 + else: + initial_scale_power = FP16_INITIAL_SCALE_POWER_DEFAULT + + return 2**initial_scale_power + + +def get_dynamic_loss_scale_args(param_dict): + loss_scale_args = None + if get_fp16_enabled(param_dict): + fp16_dict = param_dict[FP16] + dynamic_loss_args = [ + FP16_INITIAL_SCALE_POWER, + FP16_LOSS_SCALE_WINDOW, + FP16_MIN_LOSS_SCALE, + FP16_HYSTERESIS, + FP16_CONSECUTIVE_HYSTERESIS, + ] + if any(arg in list(fp16_dict.keys()) for arg in dynamic_loss_args): + init_scale = get_scalar_param(fp16_dict, FP16_INITIAL_SCALE_POWER, FP16_INITIAL_SCALE_POWER_DEFAULT) + scale_window = get_scalar_param(fp16_dict, FP16_LOSS_SCALE_WINDOW, FP16_LOSS_SCALE_WINDOW_DEFAULT) + delayed_shift = get_scalar_param(fp16_dict, FP16_HYSTERESIS, FP16_HYSTERESIS_DEFAULT) + consecutive_hysteresis = get_scalar_param(fp16_dict, FP16_CONSECUTIVE_HYSTERESIS, + FP16_CONSECUTIVE_HYSTERESIS_DEFAULT) + min_loss_scale = get_scalar_param(fp16_dict, FP16_MIN_LOSS_SCALE, FP16_MIN_LOSS_SCALE_DEFAULT) + loss_scale_args = { + INITIAL_LOSS_SCALE: 2**init_scale, + SCALE_WINDOW: scale_window, + DELAYED_SHIFT: delayed_shift, + CONSECUTIVE_HYSTERESIS: consecutive_hysteresis, + MIN_LOSS_SCALE: min_loss_scale, + } + + return loss_scale_args + + +def get_gradient_accumulation_steps(param_dict): + return get_scalar_param(param_dict, GRADIENT_ACCUMULATION_STEPS, GRADIENT_ACCUMULATION_STEPS_DEFAULT) + + +def get_sparse_gradients_enabled(param_dict): + return get_scalar_param(param_dict, SPARSE_GRADIENTS, SPARSE_GRADIENTS_DEFAULT) + + +def get_communication_data_type(param_dict, + comm_type=COMMUNICATION_DATA_TYPE, + comm_data_type_default=COMMUNICATION_DATA_TYPE_DEFAULT): + val = get_scalar_param(param_dict, comm_type, comm_data_type_default) + val = val.lower() if val is not None else val + if val is None: + return val # we must determine it by other parameters + elif val == "fp32": + return torch.float32 + elif val == "fp16": + return torch.float16 + elif val == "bf16": + return torch.bfloat16 + + raise ValueError(f"Invalid communication_data_type. Supported data types: ['fp16', 'bf16', 'fp32']. Got: {val}") + + +def get_prescale_gradients(param_dict): + return get_scalar_param(param_dict, PRESCALE_GRADIENTS, PRESCALE_GRADIENTS_DEFAULT) + + +def get_gradient_predivide_factor(param_dict): + return get_scalar_param(param_dict, GRADIENT_PREDIVIDE_FACTOR, GRADIENT_PREDIVIDE_FACTOR_DEFAULT) + + +def get_steps_per_print(param_dict): + return get_scalar_param(param_dict, STEPS_PER_PRINT, STEPS_PER_PRINT_DEFAULT) + + +def get_disable_allgather(param_dict): + return get_scalar_param(param_dict, DISABLE_ALLGATHER, DISABLE_ALLGATHER_DEFAULT) + + +def get_dump_state(param_dict): + return get_scalar_param(param_dict, DUMP_STATE, DUMP_STATE_DEFAULT) + + +def get_gradient_clipping(param_dict): + return get_scalar_param(param_dict, GRADIENT_CLIPPING, GRADIENT_CLIPPING_DEFAULT) + + +def get_graph_harvesting(param_dict): + return get_scalar_param(param_dict, GRAPH_HARVESTING, GRAPH_HARVESTING_DEFAULT) + + +def get_sparse_attention(param_dict): + if SPARSE_ATTENTION in param_dict.keys(): + sparsity = param_dict[SPARSE_ATTENTION] + mode = get_sparse_attention_mode(sparsity) + + if mode == SPARSE_DENSE_MODE: + return get_sparse_dense_config(sparsity) + elif mode == SPARSE_FIXED_MODE: + return get_sparse_fixed_config(sparsity) + elif mode == SPARSE_VARIABLE_MODE: + return get_sparse_variable_config(sparsity) + elif mode == SPARSE_BIGBIRD_MODE: + return get_sparse_bigbird_config(sparsity) + elif mode == SPARSE_BSLONGFORMER_MODE: + return get_sparse_bslongformer_config(sparsity) + else: + raise NotImplementedError(f"Given sparsity mode, {mode}, has not been implemented yet!") + + else: + return None + + +def get_sparse_dense_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + return {SPARSE_MODE: SPARSE_DENSE_MODE, SPARSE_BLOCK: block} + + +def get_sparse_fixed_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + different_layout_per_head = get_scalar_param( + sparsity, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, + ) + num_local_blocks = get_scalar_param(sparsity, SPARSE_NUM_LOCAL_BLOCKS, SPARSE_NUM_LOCAL_BLOCKS_DEFAULT) + num_global_blocks = get_scalar_param(sparsity, SPARSE_NUM_GLOBAL_BLOCKS, SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT) + attention = get_scalar_param(sparsity, SPARSE_ATTENTION_TYPE, SPARSE_ATTENTION_TYPE_DEFAULT) + horizontal_global_attention = get_scalar_param( + sparsity, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT, + ) + num_different_global_patterns = get_scalar_param( + sparsity, + SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS, + SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS_DEFAULT, + ) + + return { + SPARSE_MODE: SPARSE_FIXED_MODE, + SPARSE_BLOCK: block, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, + SPARSE_NUM_LOCAL_BLOCKS: num_local_blocks, + SPARSE_NUM_GLOBAL_BLOCKS: num_global_blocks, + SPARSE_ATTENTION_TYPE: attention, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION: horizontal_global_attention, + SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS: num_different_global_patterns, + } + + +def get_sparse_variable_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + different_layout_per_head = get_scalar_param( + sparsity, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, + ) + num_random_blocks = get_scalar_param(sparsity, SPARSE_NUM_RANDOM_BLOCKS, SPARSE_NUM_RANDOM_BLOCKS_DEFAULT) + local_window_blocks = get_scalar_param(sparsity, SPARSE_LOCAL_WINDOW_BLOCKS, SPARSE_LOCAL_WINDOW_BLOCKS_DEFAULT) + global_block_indices = get_scalar_param(sparsity, SPARSE_GLOBAL_BLOCK_INDICES, SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT) + global_block_end_indices = get_scalar_param( + sparsity, + SPARSE_GLOBAL_BLOCK_END_INDICES, + SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT, + ) + attention = get_scalar_param(sparsity, SPARSE_ATTENTION_TYPE, SPARSE_ATTENTION_TYPE_DEFAULT) + horizontal_global_attention = get_scalar_param( + sparsity, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT, + ) + + return { + SPARSE_MODE: SPARSE_VARIABLE_MODE, + SPARSE_BLOCK: block, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, + SPARSE_NUM_RANDOM_BLOCKS: num_random_blocks, + SPARSE_LOCAL_WINDOW_BLOCKS: local_window_blocks, + SPARSE_GLOBAL_BLOCK_INDICES: global_block_indices, + SPARSE_GLOBAL_BLOCK_END_INDICES: global_block_end_indices, + SPARSE_ATTENTION_TYPE: attention, + SPARSE_HORIZONTAL_GLOBAL_ATTENTION: horizontal_global_attention, + } + + +def get_sparse_bigbird_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + different_layout_per_head = get_scalar_param( + sparsity, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, + ) + num_random_blocks = get_scalar_param(sparsity, SPARSE_NUM_RANDOM_BLOCKS, SPARSE_NUM_RANDOM_BLOCKS_DEFAULT) + num_sliding_window_blocks = get_scalar_param( + sparsity, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT, + ) + num_global_blocks = get_scalar_param(sparsity, SPARSE_NUM_GLOBAL_BLOCKS, SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT) + + return { + SPARSE_MODE: SPARSE_BIGBIRD_MODE, + SPARSE_BLOCK: block, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, + SPARSE_NUM_RANDOM_BLOCKS: num_random_blocks, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS: num_sliding_window_blocks, + SPARSE_NUM_GLOBAL_BLOCKS: num_global_blocks, + } + + +def get_sparse_bslongformer_config(sparsity): + block = get_scalar_param(sparsity, SPARSE_BLOCK, SPARSE_BLOCK_DEFAULT) + different_layout_per_head = get_scalar_param( + sparsity, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT, + ) + num_sliding_window_blocks = get_scalar_param( + sparsity, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT, + ) + global_block_indices = get_scalar_param(sparsity, SPARSE_GLOBAL_BLOCK_INDICES, SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT) + global_block_end_indices = get_scalar_param( + sparsity, + SPARSE_GLOBAL_BLOCK_END_INDICES, + SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT, + ) + + return { + SPARSE_MODE: SPARSE_BSLONGFORMER_MODE, + SPARSE_BLOCK: block, + SPARSE_DIFFERENT_LAYOUT_PER_HEAD: different_layout_per_head, + SPARSE_NUM_SLIDING_WINDOW_BLOCKS: num_sliding_window_blocks, + SPARSE_GLOBAL_BLOCK_INDICES: global_block_indices, + SPARSE_GLOBAL_BLOCK_END_INDICES: global_block_end_indices, + } + + +def get_sparse_attention_mode(param_dict): + if SPARSE_MODE in param_dict.keys(): + return param_dict[SPARSE_MODE] + else: + return SPARSE_MODE_DEFAULT + + +def get_sparse_attention_type(param_dict): + if SPARSE_ATTENTION_TYPE in param_dict.keys(): + return param_dict[SPARSE_ATTENTION_TYPE] + else: + return SPARSE_ATTENTION_TYPE_DEFAULT + + +def get_pipeline_config(param_dict): + """Parses pipeline engine configuration. """ + default_pipeline = { + "stages": "auto", + "partition": "best", + "seed_layers": False, + "activation_checkpoint_interval": 0, + "pipe_partitioned": True, + "grad_partitioned": True, + } + config = default_pipeline + for key, val in param_dict.get("pipeline", {}).items(): + config[key] = val + return config + + +def get_optimizer_name(param_dict): + if OPTIMIZER in param_dict.keys() and TYPE in param_dict[OPTIMIZER].keys(): + return param_dict[OPTIMIZER][TYPE] + else: + return OPTIMIZER_TYPE_DEFAULT + + +def get_optimizer_params(param_dict): + if (get_optimizer_name(param_dict) is not None and OPTIMIZER_PARAMS in param_dict[OPTIMIZER].keys()): + return param_dict[OPTIMIZER][OPTIMIZER_PARAMS] + else: + return None + + +def get_optimizer_gradient_clipping(param_dict): + optimizer_params = get_optimizer_params(param_dict) + if optimizer_params is not None and MAX_GRAD_NORM in optimizer_params.keys(): + return optimizer_params[MAX_GRAD_NORM] + else: + return None + + +def get_optimizer_legacy_fusion(param_dict): + if OPTIMIZER in param_dict.keys() and LEGACY_FUSION in param_dict[OPTIMIZER].keys(): + return param_dict[OPTIMIZER][LEGACY_FUSION] + else: + return LEGACY_FUSION_DEFAULT + + +def get_zero_allow_untested_optimizer(param_dict): + return get_scalar_param(param_dict, ZERO_ALLOW_UNTESTED_OPTIMIZER, ZERO_ALLOW_UNTESTED_OPTIMIZER_DEFAULT) + + +def get_zero_force_ds_cpu_optimizer(param_dict): + return get_scalar_param(param_dict, ZERO_FORCE_DS_CPU_OPTIMIZER, ZERO_FORCE_DS_CPU_OPTIMIZER_DEFAULT) + + +def get_scheduler_name(param_dict): + if SCHEDULER in param_dict.keys() and TYPE in param_dict[SCHEDULER].keys(): + return param_dict[SCHEDULER][TYPE] + else: + return SCHEDULER_TYPE_DEFAULT + + +def get_scheduler_params(param_dict): + if (get_scheduler_name(param_dict) is not None and SCHEDULER_PARAMS in param_dict[SCHEDULER].keys()): + return param_dict[SCHEDULER][SCHEDULER_PARAMS] + else: + return None + + +def get_train_batch_size(param_dict): + return get_scalar_param(param_dict, TRAIN_BATCH_SIZE, TRAIN_BATCH_SIZE_DEFAULT) + + +def get_train_micro_batch_size_per_gpu(param_dict): + return get_scalar_param( + param_dict, + TRAIN_MICRO_BATCH_SIZE_PER_GPU, + TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT, + ) + + +def get_wall_clock_breakdown(param_dict): + return get_scalar_param(param_dict, WALL_CLOCK_BREAKDOWN, WALL_CLOCK_BREAKDOWN_DEFAULT) + + +def get_memory_breakdown(param_dict): + return get_scalar_param(param_dict, MEMORY_BREAKDOWN, MEMORY_BREAKDOWN_DEFAULT) + + +class HybridEngineConfig(DeepSpeedConfigModel): + enabled: bool = False + max_out_tokens: int = 512 + inference_tp_size: int = 1 + release_inference_cache: bool = False + pin_parameters: bool = True + tp_gather_partition_size: int = 8 + + +def get_hybrid_engine_config(param_dict): + hybrid_engine_config_dict = param_dict.get("hybrid_engine", {}) + hybrid_engine_config = HybridEngineConfig(**hybrid_engine_config_dict) + return hybrid_engine_config + + +def get_expert_data_topo_config(param_dict): + return get_scalar_param(param_dict, USE_DATA_BEFORE_EXPERT_PARALLEL, USE_DATA_BEFORE_EXPERT_PARALLEL_DEFAULT) + + +def get_eigenvalue_config(param_dict): + if get_quantize_enabled(param_dict): + param_dict = param_dict[QUANTIZE_TRAINING] + assert not get_eigenvalue_enabled(param_dict), "Eigenvalue based MoQ is temporarily disabled" + return ( + get_eigenvalue_enabled(param_dict), + get_eigenvalue_verbose(param_dict), + get_eigenvalue_max_iter(param_dict), + get_eigenvalue_tol(param_dict), + get_eigenvalue_stability(param_dict), + get_eigenvalue_gas_boundary_resolution(param_dict), + get_eigenvalue_layer_name(param_dict), + get_eigenvalue_layer_num(param_dict), + ) + else: + return ( + EIGENVALUE_ENABLED_DEFAULT, + EIGENVALUE_VERBOSE_DEFAULT, + EIGENVALUE_MAX_ITER_DEFAULT, + EIGENVALUE_TOL_DEFAULT, + EIGENVALUE_STABILITY_DEFAULT, + EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT, + EIGENVALUE_LAYER_NAME_DEFAULT, + EIGENVALUE_LAYER_NUM_DEFAULT, + ) + + +def get_eigenvalue_enabled(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_ENABLED, EIGENVALUE_ENABLED_DEFAULT) + else: + return EIGENVALUE_ENABLED_DEFAULT + + +def get_eigenvalue_verbose(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_VERBOSE, EIGENVALUE_VERBOSE_DEFAULT) + else: + return EIGENVALUE_VERBOSE_DEFAULT + + +def get_eigenvalue_max_iter(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_MAX_ITER, EIGENVALUE_MAX_ITER_DEFAULT) + else: + return EIGENVALUE_MAX_ITER_DEFAULT + + +def get_eigenvalue_tol(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_TOL, EIGENVALUE_TOL_DEFAULT) + else: + return EIGENVALUE_TOL_DEFAULT + + +def get_eigenvalue_stability(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_STABILITY, EIGENVALUE_STABILITY_DEFAULT) + else: + return EIGENVALUE_STABILITY_DEFAULT + + +def get_eigenvalue_gas_boundary_resolution(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param( + param_dict[EIGENVALUE], + EIGENVALUE_GAS_BOUNDARY_RESOLUTION, + EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT, + ) + else: + return EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT + + +def get_eigenvalue_layer_name(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_LAYER_NAME, EIGENVALUE_LAYER_NAME_DEFAULT) + else: + return EIGENVALUE_LAYER_NAME_DEFAULT + + +def get_eigenvalue_layer_num(param_dict): + if EIGENVALUE in param_dict.keys(): + return get_scalar_param(param_dict[EIGENVALUE], EIGENVALUE_LAYER_NUM, EIGENVALUE_LAYER_NUM_DEFAULT) + else: + return EIGENVALUE_LAYER_NUM_DEFAULT + + +def get_checkpoint_params(param_dict): + return param_dict.get(CHECKPOINT, {}) + + +def get_data_types_params(param_dict): + return param_dict.get(DATA_TYPES, {}) + + +def get_checkpoint_tag_validation_mode(checkpoint_params): + tag_validation_mode = checkpoint_params.get(CHECKPOINT_TAG_VALIDATION, CHECKPOINT_TAG_VALIDATION_DEFAULT) + tag_validation_mode = tag_validation_mode.upper() + if tag_validation_mode in CHECKPOINT_TAG_VALIDATION_MODES: + return tag_validation_mode + else: + raise DeepSpeedConfigError( + "Checkpoint config contains invalid tag_validation " + f"value of {tag_validation_mode}, expecting one of {CHECKPOINT_TAG_VALIDATION_MODES}") + + +def get_checkpoint_parallel_write_pipeline(checkpoint_params): + par_write_params = checkpoint_params.get(CHECKPOINT_PARALLEL_WRITE, {}) + par_write_pipeline = par_write_params.get(CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE, + CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE_DEFAULT) + if par_write_pipeline in [True, False]: + return par_write_pipeline + else: + raise DeepSpeedConfigError("checkpoint::parallel_write::pipeline_stage " + f"value of '{par_write_pipeline}' is invalid, expecting: true or false") + + +def get_dataloader_drop_last(param_dict): + return get_scalar_param(param_dict, DATALOADER_DROP_LAST, DATALOADER_DROP_LAST_DEFAULT) + + +'''Write deepspeed config files by modifying basic templates. +Can be used for quickly changing parameters via command line parameters.''' + + +class DeepSpeedConfigWriter: + + def __init__(self, data=None): + self.data = data if data is not None else {} + + def add_config(self, key, value): + self.data[key] = value + + def load_config(self, filename): + self.data = json.load(open(filename, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys) + + def write_config(self, filename): + with open(filename, "w") as outfile: + json.dump(self.data, outfile) + + +class DeepSpeedConfig(object): + + def __init__(self, config: Union[str, dict], mpu=None, mesh_device=None): + super(DeepSpeedConfig, self).__init__() + if isinstance(config, dict): + self._param_dict = config + elif os.path.exists(config): + self._param_dict = hjson.load(open(config, "r"), object_pairs_hook=dict_raise_error_on_duplicate_keys) + else: + try: + config_decoded = base64.urlsafe_b64decode(config).decode('utf-8') + self._param_dict = hjson.loads(config_decoded) + except (UnicodeDecodeError, AttributeError): + raise ValueError( + f"Expected a string path to an existing deepspeed config, or a dictionary or a valid base64. Received: {config}" + ) + + try: + self.global_rank = dist.get_rank() + if mpu is not None: + # Ulysses SP + if not hasattr(mpu, "get_data_parallel_world_size"): + self.world_size = dist.get_world_size() / mpu.get_sequence_parallel_world_size() + else: + self.world_size = mpu.get_data_parallel_world_size() + elif mesh_device is not None: + self.world_size = dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) + else: + # HF zero.init case where there is no mpu + if "sequence_parallel_size" in config: + self.world_size = dist.get_world_size() / config["sequence_parallel_size"] + else: + self.world_size = dist.get_world_size() + except: + self.global_rank = 0 + self.world_size = 1 + logger.info(f"Config mesh_device {mesh_device} world_size = {self.world_size}") + # If elastic-mode enabled, update compute + update _param_dict + self.elasticity_enabled = elasticity_enabled(self._param_dict) + if self.elasticity_enabled: + logger.info("DeepSpeed elasticity support enabled") + final_batch_size, valid_gpus, micro_batch_size = compute_elastic_config( + ds_config=self._param_dict, + target_deepspeed_version=__version__, + world_size=self.world_size, + ) + + elastic_dict = self._param_dict[ELASTICITY] + + # Ensure the resource scheduler saw the same elastic config we are using at runtime + ensure_immutable_elastic_config(runtime_elastic_config_dict=elastic_dict) + + self.elastic_model_parallel_size = elastic_dict.get(MODEL_PARALLEL_SIZE, MODEL_PARALLEL_SIZE_DEFAULT) + if self.elastic_model_parallel_size < 1: + raise ElasticityConfigError("Model-Parallel size cannot be less than 1, " + f"given model-parallel size: {self.elastic_model_parallel_size}") + + self.num_gpus_per_node = elastic_dict.get(NUM_GPUS_PER_NODE, NUM_GPUS_PER_NODE_DEFAULT) + if self.num_gpus_per_node < 1: + raise ElasticityConfigError("NUmber of GPUs per node cannot be less than 1, " + f"given number of GPUs per node: {self.num_gpus_per_node}") + + ignore_non_elastic_batch_info = elastic_dict.get(IGNORE_NON_ELASTIC_BATCH_INFO, + IGNORE_NON_ELASTIC_BATCH_INFO_DEFAULT) + + if not ignore_non_elastic_batch_info: + batch_params = [ + TRAIN_BATCH_SIZE, + TRAIN_MICRO_BATCH_SIZE_PER_GPU, + GRADIENT_ACCUMULATION_STEPS, + ] + if any(map(lambda t: t in self._param_dict, batch_params)): + raise ElasticityConfigError("One or more batch related parameters were found in your " \ + f"ds_config ({TRAIN_BATCH_SIZE}, {TRAIN_MICRO_BATCH_SIZE_PER_GPU}, and/or " \ + f"{GRADIENT_ACCUMULATION_STEPS}). These parameters *will not be used* since " \ + "elastic training is enabled, which takes control of these parameters. " \ + "If you want to suppress this error (the parameters will be silently ignored) " \ + f"please set {IGNORE_NON_ELASTIC_BATCH_INFO}':true in your elasticity config.") + + # micro_bsz * world_size * gas = total_batch_size + # gas = total_batch_size // (micro_bsz * world_size) + gradient_accu_steps = final_batch_size // (micro_batch_size * self.world_size) + + if TRAIN_BATCH_SIZE in self._param_dict: + logger.warning("[Elasticity] overriding training_batch_size: " + f"{self._param_dict[TRAIN_BATCH_SIZE]} -> {final_batch_size}") + if TRAIN_MICRO_BATCH_SIZE_PER_GPU in self._param_dict: + logger.warning("[Elasticity] overriding train_micro_batch_size_per_gpu: " + f"{self._param_dict[TRAIN_MICRO_BATCH_SIZE_PER_GPU]} -> {micro_batch_size}") + if GRADIENT_ACCUMULATION_STEPS in self._param_dict: + logger.warning("[Elasticity] overriding gradient_accumulation_steps: " + f"{self._param_dict[GRADIENT_ACCUMULATION_STEPS]} -> {gradient_accu_steps}") + + logger.info(f"[Elasticity] valid GPU counts: {valid_gpus}") + + self._param_dict[TRAIN_BATCH_SIZE] = final_batch_size + self._param_dict[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = micro_batch_size + self._param_dict[GRADIENT_ACCUMULATION_STEPS] = gradient_accu_steps + + # Pass a copy so that user json is unmodified, e.g. for logging + self._initialize_params(copy.copy(self._param_dict)) + self._configure_train_batch_size() + self._do_sanity_check() + + def _initialize_params(self, param_dict): + self.train_batch_size = get_train_batch_size(param_dict) + self.train_micro_batch_size_per_gpu = get_train_micro_batch_size_per_gpu(param_dict) + self.gradient_accumulation_steps = get_gradient_accumulation_steps(param_dict) + self.steps_per_print = get_steps_per_print(param_dict) + self.dump_state = get_dump_state(param_dict) + + self.disable_allgather = get_disable_allgather(param_dict) + self.communication_data_type = get_communication_data_type(param_dict) + self.seq_parallel_communication_data_type = get_communication_data_type( + param_dict, SEQ_PARALLEL_COMMUNICATION_DATA_TYPE, SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT) + self.prescale_gradients = get_prescale_gradients(param_dict) + self.gradient_predivide_factor = get_gradient_predivide_factor(param_dict) + self.sparse_gradients_enabled = get_sparse_gradients_enabled(param_dict) + + self.zero_config = get_zero_config(param_dict) + self.mics_shard_size = self.zero_config.mics_shard_size + self.mics_hierarchial_params_gather = self.zero_config.mics_hierarchical_params_gather + self.zero_optimization_stage = self.zero_config.stage + self.zero_enabled = self.zero_optimization_stage > 0 + + self.activation_checkpointing_config = DeepSpeedActivationCheckpointingConfig(param_dict) + + self.comms_config = DeepSpeedCommsConfig(param_dict) + self.monitor_config = get_monitor_config(param_dict) + + self.gradient_clipping = get_gradient_clipping(param_dict) + self.fp16_enabled = get_fp16_enabled(param_dict) + self.fp16_auto_cast = get_fp16_auto_cast(param_dict) + self.bfloat16_enabled = get_bfloat16_enabled(param_dict) + self.bfloat16_immediate_grad_update = get_bfloat16_immediate_grad_update(param_dict) + assert not (self.fp16_enabled + and self.bfloat16_enabled), 'bfloat16 and fp16 modes cannot be simultaneously enabled' + self.fp16_master_weights_and_gradients = get_fp16_master_weights_and_grads_enabled(param_dict) + self.amp_enabled = get_amp_enabled(param_dict) + self.amp_params = get_amp_params(param_dict) + self.loss_scale = get_loss_scale(param_dict) + self.initial_dynamic_scale = get_initial_dynamic_scale(param_dict) + self.dynamic_loss_scale_args = get_dynamic_loss_scale_args(param_dict) + + self.compression_config = get_compression_config(param_dict) + self.graph_harvesting = get_graph_harvesting(param_dict) + + self.optimizer_name = get_optimizer_name(param_dict) + if (self.optimizer_name is not None and self.optimizer_name.lower() in DEEPSPEED_OPTIMIZERS): + self.optimizer_name = self.optimizer_name.lower() + + self.optimizer_params = get_optimizer_params(param_dict) + self.optimizer_legacy_fusion = get_optimizer_legacy_fusion(param_dict) + + self.zero_allow_untested_optimizer = get_zero_allow_untested_optimizer(param_dict) + + self.zero_force_ds_cpu_optimizer = get_zero_force_ds_cpu_optimizer(param_dict) + + self.scheduler_name = get_scheduler_name(param_dict) + self.scheduler_params = get_scheduler_params(param_dict) + + self.flops_profiler_config = DeepSpeedFlopsProfilerConfig(param_dict) + self.wall_clock_breakdown = (get_wall_clock_breakdown(param_dict) | self.flops_profiler_config.enabled) + self.memory_breakdown = get_memory_breakdown(param_dict) + self.autotuning_config = DeepSpeedAutotuningConfig(param_dict) + + ( + self.eigenvalue_enabled, + self.eigenvalue_verbose, + self.eigenvalue_max_iter, + self.eigenvalue_tol, + self.eigenvalue_stability, + self.eigenvalue_gas_boundary_resolution, + self.eigenvalue_layer_name, + self.eigenvalue_layer_num, + ) = get_eigenvalue_config(param_dict) + + self.use_data_before_expert_parallel_ = get_expert_data_topo_config(param_dict) + self.hybrid_engine = get_hybrid_engine_config(param_dict) + + self.sparse_attention = get_sparse_attention(param_dict) + self.pipeline = get_pipeline_config(param_dict) + + self.pld_enabled = get_pld_enabled(param_dict) + self.pld_params = get_pld_params(param_dict) + + self.curriculum_enabled_legacy = get_curriculum_enabled_legacy(param_dict) + self.curriculum_params_legacy = get_curriculum_params_legacy(param_dict) + + self.data_efficiency_enabled = get_data_efficiency_enabled(param_dict) + self.data_efficiency_config = get_data_efficiency_config(param_dict) + + checkpoint_params = get_checkpoint_params(param_dict) + validation_mode = get_checkpoint_tag_validation_mode(checkpoint_params) + self.checkpoint_tag_validation_enabled = (validation_mode != ValidationMode.IGNORE) + self.checkpoint_tag_validation_fail = validation_mode == ValidationMode.FAIL + self.load_universal_checkpoint = checkpoint_params.get(LOAD_UNIVERSAL_CHECKPOINT, + LOAD_UNIVERSAL_CHECKPOINT_DEFAULT) + + self.use_node_local_storage = checkpoint_params.get(USE_NODE_LOCAL_STORAGE_CHECKPOINT, + USE_NODE_LOCAL_STORAGE_CHECKPOINT_DEFAULT) + + data_types_params = get_data_types_params(param_dict) + self.grad_accum_dtype = data_types_params.get(GRAD_ACCUM_DTYPE, GRAD_ACCUM_DTYPE_DEFAULT) + + par_write_pipe = get_checkpoint_parallel_write_pipeline(checkpoint_params) + self.checkpoint_parallel_write_pipeline = par_write_pipe + + self.aio_config = get_aio_config(param_dict) + + self.dataloader_drop_last = get_dataloader_drop_last(param_dict) + + self.nebula_config = DeepSpeedNebulaConfig(param_dict) + + self.weight_quantization_config = WeightQuantConfig( + **param_dict['weight_quantization']) if 'weight_quantization' in param_dict else None + + self.compile_config = CompileConfig(**param_dict.get('compile', {})) + + self.timers_config = get_timers_config(param_dict) + self.tensor_parallel_config = get_tensor_parallel_config(param_dict) + + def _batch_assertion(self): + + train_batch = self.train_batch_size + micro_batch = self.train_micro_batch_size_per_gpu + grad_acc = self.gradient_accumulation_steps + + assert (train_batch > 0), f"Train batch size: {train_batch} has to be greater than 0" + + assert (micro_batch > 0), f"Micro batch size per gpu: {micro_batch} has to be greater than 0" + + assert (grad_acc > 0), f"Gradient accumulation steps: {grad_acc} has to be greater than 0" + + assert train_batch == micro_batch * grad_acc * self.world_size, ( + f"Check batch related parameters. train_batch_size is not equal " + "to micro_batch_per_gpu * gradient_acc_step * world_size " + f"{train_batch} != {micro_batch} * {grad_acc} * {self.world_size}") + + def _set_batch_related_parameters(self): + + train_batch = self.train_batch_size + micro_batch = self.train_micro_batch_size_per_gpu + grad_acc = self.gradient_accumulation_steps + + #print(f"in: train_batch = {train_batch}, micro_batch={micro_batch}") + + # all values are provided nothing needs to be set + if train_batch is not None and micro_batch is not None and grad_acc is not None: + return + + # global_accumulation_steps needs to be set + elif train_batch is not None and micro_batch is not None: + grad_acc = train_batch // micro_batch + grad_acc //= self.world_size + self.gradient_accumulation_steps = grad_acc + + # micro_batch_per_gpu needs to be set + elif train_batch is not None and grad_acc is not None: + micro_batch = train_batch // self.world_size + micro_batch //= grad_acc + self.train_micro_batch_size_per_gpu = micro_batch + + # train_batch_size needs to be set + elif micro_batch is not None and grad_acc is not None: + train_batch_size = micro_batch * grad_acc + train_batch_size *= self.world_size + self.train_batch_size = train_batch_size + + # gradient_accumulation_steps and micro_batch_per_gpus is set + elif train_batch is not None: + self.gradient_accumulation_steps = 1 + self.train_micro_batch_size_per_gpu = train_batch // self.world_size + + # train_batch_size and gradient_accumulation_step is set + elif micro_batch is not None: + self.train_batch_size = micro_batch * self.world_size + self.gradient_accumulation_steps = 1 + + # either none of the three parameters are provided or just gradient_accumulation_step is provided + else: + assert False, \ + 'Either train_batch_size or train_micro_batch_size_per_gpu needs to be provided' + + #print(f"final: {self.train_batch_size=} {self.train_micro_batch_size_per_gpu=} {self.gradient_accumulation_steps=}") + + def _configure_train_batch_size(self): + self._set_batch_related_parameters() + self._batch_assertion() + + def _do_sanity_check(self): + self._do_error_check() + + self._do_warning_check() + + def print_user_config(self): + logger.info(" json = {}".format( + json.dumps( + self._param_dict, + sort_keys=True, + indent=4, + cls=ScientificNotationEncoder, + separators=(",", ":"), + ))) + + def print(self, name): + logger.info("{}:".format(name)) + for arg in sorted(vars(self)): + if arg != "_param_dict": + dots = "." * (29 - len(arg)) + logger.info(" {} {} {}".format(arg, dots, getattr(self, arg))) + + self.print_user_config() + + def _do_error_check(self): + assert (self.train_micro_batch_size_per_gpu + ), "DeepSpeedConfig: {} is not defined".format(TRAIN_MICRO_BATCH_SIZE_PER_GPU) + + assert ( + self.gradient_accumulation_steps), "DeepSpeedConfig: {} is not defined".format(GRADIENT_ACCUMULATION_STEPS) + + if self.zero_enabled: + assert (self.zero_optimization_stage + <= ZeroStageEnum.max_stage), "DeepSpeedConfig: Maximum supported ZeRO stage is {}".format( + ZeroStageEnum.max_stage) + + if self.fp16_master_weights_and_gradients: + assert self.zero_enabled and self.zero_optimization_stage == ZeroStageEnum.gradients, "Fp16_master_weights_and_grads is only supported with ZeRO Stage 2 for now." + + def _do_warning_check(self): + fp16_enabled = self.fp16_enabled + + vocabulary_size = self._param_dict.get(VOCABULARY_SIZE, VOCABULARY_SIZE_DEFAULT) + if vocabulary_size and vocabulary_size % TENSOR_CORE_ALIGN_SIZE != 0: + logger.warning( + "DeepSpeedConfig: vocabulary size {} is not aligned to {}, may import tensor core utilization.".format( + vocabulary_size, TENSOR_CORE_ALIGN_SIZE)) + + if (self.optimizer_params is not None and MAX_GRAD_NORM in self.optimizer_params.keys() + and self.optimizer_params[MAX_GRAD_NORM] > 0): + if fp16_enabled: + if self.global_rank == 0: + logger.warning("DeepSpeedConfig: In FP16 mode, DeepSpeed will pass {}:{} to FP16 wrapper".format( + MAX_GRAD_NORM, self.optimizer_params[MAX_GRAD_NORM])) + else: + if self.global_rank == 0: + logger.warning( + "DeepSpeedConfig: In FP32 mode, DeepSpeed does not permit MAX_GRAD_NORM ({}) > 0, setting to zero" + .format(self.optimizer_params[MAX_GRAD_NORM])) + self.optimizer_params[MAX_GRAD_NORM] = 0.0 diff --git a/lib/python3.12/site-packages/deepspeed/runtime/config_utils.py b/lib/python3.12/site-packages/deepspeed/runtime/config_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..54cf813fd7fd60d38e57e3fab8be30ffa2ed63e8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/config_utils.py @@ -0,0 +1,212 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Collection of DeepSpeed configuration utilities +""" +import collections +import json +import torch +from functools import reduce +from pydantic import BaseModel, ConfigDict, field_serializer + +from deepspeed.utils import logger + + +class DeepSpeedConfigModel(BaseModel): + """ + This class should be used as a base for all DeepSpeed configs. It extends + pydantic.BaseModel to allow for deprecated fields. To enable this feature, + add deprecated=True to pydantic.Field: + + my_dep_field: int = Field(0, deprecated=True) + + Deprecated Field kwargs: + - deprecated: [True|False], default False + Enables / Disables deprecated fields + - deprecated_msg: str, default "" + Message to include with deprecation warning + - new_param: str, default "" + Name of the field replacing the deprecated field + - set_new_param: [True|False], default True + If new_param is provided, enables setting the value of that param with + deprecated field value + - new_param_fn: callable, default (lambda x: x) + If new_param is provided and set_new_param is True, this function will + modify the value of the deprecated field before placing that value in + the new_param field + + Example: + my_new_field is replacing a deprecated my_old_field. The expected type + for my_new_field is int while the expected type for my_old_field is + str. We want to maintain backward compatibility with our configs, so we + define the fields with: + + class MyExampleConfig(DeepSpeedConfigModel): + my_new_field: int = 0 + my_old_field: str = Field('0', + deprecated=True, + new_param='my_new_field', + new_param_fn=(lambda x: int(x))) + """ + + def __init__(self, strict=False, **data): + if (not strict): # This is temporary until we refactor all DS configs, allows HF to load models + data = {k: v for k, v in data.items() if (v != "auto" or k == "replace_method")} + super().__init__(**data) + self._deprecated_fields_check() + + def _process_deprecated_field(self, dep_field): + # Get information about the deprecated field + pydantic_config = self + fields_set = pydantic_config.model_fields_set + kwargs = type(pydantic_config).model_fields[dep_field].json_schema_extra + new_param_fn = kwargs.get("new_param_fn", lambda x: x) + param_value = new_param_fn(getattr(pydantic_config, dep_field)) + new_field = kwargs.get("new_param", "") + dep_msg = kwargs.get("deprecated_msg", "") + if dep_field in fields_set: + logger.warning(f"Config parameter {dep_field} is deprecated" + + (f" use {new_field} instead" if new_field else "") + (f". {dep_msg}" if dep_msg else "")) + # Check if there is a new param and if it should be set with a value + if new_field and kwargs.get("set_new_param", True): + # Remove the deprecate field if there is a replacing field + try: + delattr(pydantic_config, dep_field) + except Exception as e: + logger.error(f"Tried removing deprecated '{dep_field}' from config") + raise e + + # Set new param value + new_param_nested = new_field.split(".") + if len(new_param_nested) > 1: + # If the new param exists in a subconfig, we need to get + # the fields set for that subconfig + pydantic_config = reduce(getattr, new_param_nested[:-1], pydantic_config) + fields_set = pydantic_config.model_fields_set + new_param_name = new_param_nested[-1] + assert ( + new_param_name not in fields_set + ), f"Cannot provide deprecated parameter '{dep_field}' and replacing parameter '{new_field}' together" + # A custom function for converting the old param value to new param value can be provided + try: + setattr(pydantic_config, new_param_name, param_value) + except Exception as e: + logger.error(f"Tried setting value for '{new_field}' with value from deprecated '{dep_field}'") + raise e + + def _deprecated_fields_check(self): + fields = self.model_fields + for field_name, field_info in fields.items(): + if field_info.json_schema_extra and field_info.json_schema_extra.get("deprecated", False): + self._process_deprecated_field(field_name) + + model_config = ConfigDict( + validate_default=True, + validate_assignment=True, + use_enum_values=True, + populate_by_name=True, + extra="forbid", + arbitrary_types_allowed=True, + protected_namespaces=(), + ) + + @field_serializer("dtype", check_fields=False) + def serialize_torch_dtype(dtype: torch.dtype) -> str: + return str(dtype) + + +def get_config_default(config, field_name): + assert field_name in config.model_fields, f"'{field_name}' is not a field in {config}" + assert not config.model_fields.get( + field_name).is_required(), f"'{field_name}' is a required field and does not have a default value" + return config.model_fields.get(field_name).get_default() + + +class pp_int(int): + """ + A wrapper for integers that will return a custom string or comma-formatted + string of the integer. For example, print(pp_int(1e5)) will return + "10,000". This is useful mainly for auto-generated documentation purposes. + """ + + def __new__(cls, val, custom_print_str=None): + inst = super().__new__(cls, val) + inst.custom_print_str = custom_print_str + return inst + + def __repr__(self): + if self.custom_print_str: + return self.custom_print_str + return f"{self.real:,}" + + +# adapted from https://stackoverflow.com/a/50701137/9201239 +class ScientificNotationEncoder(json.JSONEncoder): + """ + This class overrides ``json.dumps`` default formatter. + + This version keeps everything as normal except formats numbers bigger than 1e3 using scientific notation. + + Just pass ``cls=ScientificNotationEncoder`` to ``json.dumps`` to activate it + + """ + + def iterencode(self, o, _one_shot=False, level=0): + indent = self.indent if self.indent is not None else 4 + prefix_close = " " * level * indent + level += 1 + prefix = " " * level * indent + if isinstance(o, bool): + return "true" if o else "false" + elif isinstance(o, float) or isinstance(o, int): + if o > 1e3: + return f"{o:e}" + else: + return f"{o}" + elif isinstance(o, collections.abc.Mapping): + x = [f'\n{prefix}"{k}": {self.iterencode(v, level=level)}' for k, v in o.items()] + return "{" + ", ".join(x) + f"\n{prefix_close}" + "}" + elif isinstance(o, collections.abc.Sequence) and not isinstance(o, str): + return f"[{ f', '.join(map(self.iterencode, o)) }]" + return "\n, ".join(super().iterencode(o, _one_shot)) + + +class DeepSpeedConfigObject(object): + """ + For json serialization + """ + + def repr(self): + return self.__dict__ + + def __repr__(self): + return json.dumps( + self.__dict__, + sort_keys=True, + indent=4, + cls=ScientificNotationEncoder, + ) + + +def get_scalar_param(param_dict, param_name, param_default_value): + return param_dict.get(param_name, param_default_value) + + +def get_list_param(param_dict, param_name, param_default_value): + return param_dict.get(param_name, param_default_value) + + +def get_dict_param(param_dict, param_name, param_default_value): + return param_dict.get(param_name, param_default_value) + + +def dict_raise_error_on_duplicate_keys(ordered_pairs): + """Reject duplicate keys.""" + d = dict((k, v) for k, v in ordered_pairs) + if len(d) != len(ordered_pairs): + counter = collections.Counter([pair[0] for pair in ordered_pairs]) + keys = [key for key, value in counter.items() if value > 1] + raise ValueError("Duplicate keys in DeepSpeed config: {}".format(keys)) + return d diff --git a/lib/python3.12/site-packages/deepspeed/runtime/constants.py b/lib/python3.12/site-packages/deepspeed/runtime/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..fa7e9cad73b8e848c490e572d3121e30347c9f0a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/constants.py @@ -0,0 +1,457 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +############################################# +# Routes +############################################# +ROUTE_TRAIN = "train" +ROUTE_EVAL = "eval" +ROUTE_PREDICT = "predict" +ROUTE_ENCODE = "encode" + +############################################# +# Batch size +############################################# +TRAIN_BATCH_SIZE = "train_batch_size" +TRAIN_BATCH_SIZE_DEFAULT = None + +############################################# +# Sparse attention +############################################# +SPARSE_ATTENTION = "sparse_attention" +SPARSE_DENSE_MODE = "dense" +SPARSE_FIXED_MODE = "fixed" +SPARSE_VARIABLE_MODE = "variable" +SPARSE_BIGBIRD_MODE = "bigbird" +SPARSE_BSLONGFORMER_MODE = "bslongformer" +SPARSE_MODE = "mode" +SPARSE_MODE_DEFAULT = SPARSE_FIXED_MODE +SPARSE_BLOCK = "block" +SPARSE_BLOCK_DEFAULT = 16 +SPARSE_DIFFERENT_LAYOUT_PER_HEAD = "different_layout_per_head" +SPARSE_DIFFERENT_LAYOUT_PER_HEAD_DEFAULT = False +SPARSE_NUM_LOCAL_BLOCKS = "num_local_blocks" +SPARSE_NUM_LOCAL_BLOCKS_DEFAULT = 4 +SPARSE_NUM_GLOBAL_BLOCKS = "num_global_blocks" +SPARSE_NUM_GLOBAL_BLOCKS_DEFAULT = 1 +SPARSE_ATTENTION_TYPE = "attention" +SPARSE_ATTENTION_TYPE_DEFAULT = "bidirectional" +SPARSE_HORIZONTAL_GLOBAL_ATTENTION = "horizontal_global_attention" +SPARSE_HORIZONTAL_GLOBAL_ATTENTION_DEFAULT = False +SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS = "num_different_global_patterns" +SPARSE_NUM_DIFFERENT_GLOBAL_PATTERNS_DEFAULT = 1 +SPARSE_NUM_RANDOM_BLOCKS = "num_random_blocks" +SPARSE_NUM_RANDOM_BLOCKS_DEFAULT = 0 +SPARSE_LOCAL_WINDOW_BLOCKS = "local_window_blocks" +SPARSE_LOCAL_WINDOW_BLOCKS_DEFAULT = [4] +SPARSE_GLOBAL_BLOCK_INDICES = "global_block_indices" +SPARSE_GLOBAL_BLOCK_INDICES_DEFAULT = [0] +SPARSE_GLOBAL_BLOCK_END_INDICES = "global_block_end_indices" +SPARSE_GLOBAL_BLOCK_END_INDICES_DEFAULT = None +SPARSE_NUM_SLIDING_WINDOW_BLOCKS = "num_sliding_window_blocks" +SPARSE_NUM_SLIDING_WINDOW_BLOCKS_DEFAULT = 3 + +############################################# +# Optimizer and lr scheduler +############################################# +OPTIMIZER = "optimizer" +OPTIMIZER_TYPE_DEFAULT = None +OPTIMIZER_PARAMS = "params" +TYPE = "type" +LEGACY_FUSION = "legacy_fusion" +LEGACY_FUSION_DEFAULT = False +SCHEDULER = "scheduler" +SCHEDULER_TYPE_DEFAULT = None +SCHEDULER_PARAMS = "params" +MAX_GRAD_NORM = 'max_grad_norm' + +############################################# +# Optimizer and lr scheduler +############################################# +ZERO_ALLOW_UNTESTED_OPTIMIZER = "zero_allow_untested_optimizer" +ZERO_ALLOW_UNTESTED_OPTIMIZER_DEFAULT = False +ZERO_FORCE_DS_CPU_OPTIMIZER = "zero_force_ds_cpu_optimizer" +ZERO_FORCE_DS_CPU_OPTIMIZER_DEFAULT = True + +# Steps +STEPS_PER_PRINT = "steps_per_print" +STEPS_PER_PRINT_DEFAULT = None + +######################################### +# Training micro batch size per GPU +######################################### +# Batch size for one training step. This is used when the +# TRAIN_BATCH_SIZE cannot fit in GPU memory to determine +# the number of gradient accumulation steps. By default, this +# is set to None. Users can configure in ds_config.json as below example: +TRAIN_MICRO_BATCH_SIZE_PER_GPU = ''' +TRAIN_MICRO_BATCH_SIZE_PER_GPU is defined in this format: +"train_micro_batch_size_per_gpu": 1 +''' +TRAIN_MICRO_BATCH_SIZE_PER_GPU = "train_micro_batch_size_per_gpu" +TRAIN_MICRO_BATCH_SIZE_PER_GPU_DEFAULT = None + +######################################### +# Gradient Accumulation +######################################### +# Gradient accumulation feature. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +GRADIENT_ACCUMULATION_FORMAT = ''' +Gradient Accumulation should be of the format: +"gradient_accumulation_steps": 1 +''' +GRADIENT_ACCUMULATION_STEPS = "gradient_accumulation_steps" +GRADIENT_ACCUMULATION_STEPS_DEFAULT = None + +# DeepSpeed CSR gradient sparsity +SPARSE_GRADIENTS = "sparse_gradients" +SPARSE_GRADIENTS_DEFAULT = False + +######################################### +# BFLOAT16 support +######################################### +# BFLOAT16 feature. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +BFLOAT16_FORMAT = ''' +BFLOAT16 parameters should be of the format: +"bf16": { + "enabled": true +} +''' +BFLOAT16 = "bf16" +BFLOAT16_OLD = "bfloat16" # keeping for backwards compatibility + +BFLOAT16_ENABLED = "enabled" +BFLOAT16_ENABLED_DEFAULT = False + +# BFLOAT16 optimizer immediate gradient update +BFLOAT16_IMMEDIATE_GRAD_UPDATE = "immediate_grad_update" +BFLOAT16_IMMEDIATE_GRAD_UPDATE_DEFAULT = True + +######################################### +# FP16 support +######################################### +# FP16 feature. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +FP16_FORMAT = ''' +FP16 parameters should be of the format: +"fp16": { + "enabled": true, + "auto_cast": false, + "loss_scale": 0, + "initial_scale_power": 16, + "loss_scale_window": 1000, + "hysteresis": 2, + "consecutive_hysteresis": false, + "min_loss_scale": 1 +} +''' +FP16 = "fp16" + +FP16_ENABLED = "enabled" +FP16_ENABLED_DEFAULT = False + +# FP16 loss scale, zero means using dynamic scaling +FP16_LOSS_SCALE = "loss_scale" +FP16_LOSS_SCALE_DEFAULT = 0 + +FP16_AUTO_CAST = "auto_cast" +FP16_AUTO_CAST_DEFAULT = False + +# FP16 initial dynamic scale loss power +FP16_INITIAL_SCALE_POWER = "initial_scale_power" +FP16_INITIAL_SCALE_POWER_DEFAULT = 16 + +# FP16 loss scale window +FP16_LOSS_SCALE_WINDOW = "loss_scale_window" +FP16_LOSS_SCALE_WINDOW_DEFAULT = 1000 + +# FP16 hysteresis +FP16_HYSTERESIS = "hysteresis" +FP16_HYSTERESIS_DEFAULT = 2 + +# FP16 consecutive hysteresis +FP16_CONSECUTIVE_HYSTERESIS = "consecutive_hysteresis" +FP16_CONSECUTIVE_HYSTERESIS_DEFAULT = False + +# FP16 min loss scale +FP16_MIN_LOSS_SCALE = "min_loss_scale" +FP16_MIN_LOSS_SCALE_DEFAULT = 1 + +# FP16 master and grads +FP16_MASTER_WEIGHTS_AND_GRADS = "fp16_master_weights_and_grads" +FP16_MASTER_WEIGHTS_AND_GRADS_DEFAULT = False + +######################################### +# Apex AMP support +######################################### +# Use Apex AMP for mixed precision support, all parameters (other than 'enabled') will be passed to +# amp.initialize(model, optimizer, **amp_params) +# See apex documentation for supported parameters/features: https://nvidia.github.io/apex/amp.html#apex.amp.initialize +AMP_FORMAT = ''' +"amp" { + "enabled: true, + "opt_level": "O1", + ... +} +''' +AMP = "amp" + +AMP_ENABLED = "enabled" +AMP_ENABLED_DEFAULT = False + +######################################### +# Gradient clipping +######################################### +# Gradient clipping. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +GRADIENT_CLIPPING_FORMAT = ''' +Gradient clipping should be enabled as: +"gradient_clipping": 1.0 +''' +GRADIENT_CLIPPING = 'gradient_clipping' +GRADIENT_CLIPPING_DEFAULT = 0. + +######################################### +# Capture graph for short kernels sequences +######################################### +# Graph harvesting. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +GRAPH_HARVESTING_FORMAT = ''' +Graph harvesting should be enabled as: +"graph_harvesting": true +''' +GRAPH_HARVESTING = 'graph_harvesting' +GRAPH_HARVESTING_DEFAULT = False + +######################################### +# Communication data type +######################################### +# Supported types: ['none', 'fp16', 'fp32'] +# By default, this feature is not enabled ('none' value) +# Users can configure in ds_config.json as below example: +COMMUNICATION_DATA_TYPE_FORMAT = ''' +Communication data type should be set as: +"communication_data_type": "fp32" +''' +COMMUNICATION_DATA_TYPE = "communication_data_type" +COMMUNICATION_DATA_TYPE_DEFAULT = None + +########################################################### +# Gradient communication data type for sequence parallelism +########################################################### +# Supported types: ['fp16', 'bf16','fp32'] +# Default value is fp32 +# Users can configure in ds_config.json as below example: +SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_FORMAT = ''' +Optional comm data type for seq paralleism should be set as: +"seq_parallel_communication_data_type": "fp32" +''' +SEQ_PARALLEL_COMMUNICATION_DATA_TYPE = "seq_parallel_communication_data_type" +SEQ_PARALLEL_COMMUNICATION_DATA_TYPE_DEFAULT = "fp32" + +######################################### +# Scale/predivide gradients before allreduce +######################################### +# Prescale gradients. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +PRESCALE_GRADIENTS_FORMAT = ''' +Gradient prescaling should be enabled as: +"prescale_gradients": true +''' +PRESCALE_GRADIENTS = "prescale_gradients" +PRESCALE_GRADIENTS_DEFAULT = False + +GRADIENT_PREDIVIDE_FACTOR_FORMAT = ''' +Gradient predivide factor should be enabled as: +"gradient_predivide_factor": 1.0 +''' +GRADIENT_PREDIVIDE_FACTOR = "gradient_predivide_factor" +GRADIENT_PREDIVIDE_FACTOR_DEFAULT = 1.0 + +######################################### +# Disable AllGather +######################################### +# Disable AllGather. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +DISABLE_ALLGATHER_FORMAT = ''' +Disable AllGather should be enabled as: +"disable_allgather": true +''' +DISABLE_ALLGATHER = "disable_allgather" +DISABLE_ALLGATHER_DEFAULT = False + +######################################### +# Dump DeepSpeed state +######################################### +# Dump State. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +DUMP_STATE_FORMAT = ''' +Dump state should be enabled as: +"dump_state": true +''' +DUMP_STATE = 'dump_state' +DUMP_STATE_DEFAULT = False + +######################################### +# Vocabulary size +######################################### +# Vocabulary size. +# Users can configure in ds_config.json as below example: +VOCABULARY_SIZE_FORMAT = ''' +Vocabulary size can be specified as: +"vocabulary_size": 1024 +''' +VOCABULARY_SIZE = 'vocabulary_size' +VOCABULARY_SIZE_DEFAULT = None + +######################################### +# Wall block breakdown +######################################### +# Wall clock breakdown. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +WALL_CLOCK_BREAKDOWN_FORMAT = ''' +Wall block breakdown should be enabled as: +"wall_clock_breakdown": true +''' +WALL_CLOCK_BREAKDOWN = 'wall_clock_breakdown' +WALL_CLOCK_BREAKDOWN_DEFAULT = False + +MEMORY_BREAKDOWN = 'memory_breakdown' +MEMORY_BREAKDOWN_DEFAULT = False + +######################################### +# Eigenvalue +######################################### +# Eigenvalue computation. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +EIGENVALUE_FORMAT = ''' +Tensorboard can be specified as: +"eigenvalue": { + "enabled": true, + "verbose": true, + "max_iter": 100, + "tol": 1e-2, + "stability": 1e-6 +} +''' +EIGENVALUE = "eigenvalue" + +# Tensorboard enable signal +EIGENVALUE_ENABLED = "enabled" +EIGENVALUE_ENABLED_DEFAULT = False + +EIGENVALUE_VERBOSE = "verbose" +EIGENVALUE_VERBOSE_DEFAULT = False + +EIGENVALUE_MAX_ITER = "max_iter" +EIGENVALUE_MAX_ITER_DEFAULT = 100 + +EIGENVALUE_TOL = "tol" +EIGENVALUE_TOL_DEFAULT = 1e-2 + +EIGENVALUE_STABILITY = "stability" +EIGENVALUE_STABILITY_DEFAULT = 1e-6 + +EIGENVALUE_GAS_BOUNDARY_RESOLUTION = "gas_boundary_resolution" +EIGENVALUE_GAS_BOUNDARY_RESOLUTION_DEFAULT = 1 + +EIGENVALUE_LAYER_NAME = "layer_name" +EIGENVALUE_LAYER_NAME_DEFAULT = "bert.encoder.layer" + +EIGENVALUE_LAYER_NUM = "layer_num" +EIGENVALUE_LAYER_NUM_DEFAULT = 0 + +######################################### +# Progressive Layer Drop (PLD) +######################################### +PROGRESSIVE_LAYER_DROP = "progressive_layer_drop" + +# PLD enable signal +PLD_ENABLED = "enabled" +PLD_ENABLED_DEFAULT = False + +PLD_THETA = "theta" +PLD_THETA_DEFAULT = 1.0 + +PLD_GAMMA = "gamma" +PLD_GAMMA_DEFAULT = 0.001 + + +######################################### +# Validation modes +######################################### +class ValidationMode: + WARN = "WARN" + IGNORE = "IGNORE" + FAIL = "FAIL" + + +######################################### +# Checkpoint config params +######################################### +# "checkpoint": { +# tag_validation=["Ignore"|"Warn"|"Fail"] +# load_universal=false +# use_node_local_storage=false +# parallel_write: { +# pipeline_stage: [True|False] +# } +# } +CHECKPOINT = "checkpoint" +CHECKPOINT_TAG_VALIDATION = "tag_validation" +CHECKPOINT_TAG_VALIDATION_DEFAULT = ValidationMode.WARN +CHECKPOINT_TAG_VALIDATION_MODES = [ValidationMode.WARN, ValidationMode.IGNORE, ValidationMode.FAIL] + +LOAD_UNIVERSAL_CHECKPOINT = "load_universal" +LOAD_UNIVERSAL_CHECKPOINT_DEFAULT = False + +USE_NODE_LOCAL_STORAGE_CHECKPOINT = "use_node_local_storage" +USE_NODE_LOCAL_STORAGE_CHECKPOINT_DEFAULT = False + +CHECKPOINT_PARALLEL_WRITE = "parallel_write" +CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE = "pipeline_stage" +CHECKPOINT_PARALLEL_WRITE_PIPELINE_STAGE_DEFAULT = False + +######################################### +# Data types config params +######################################### +# "data_types": { +# grad_accum_dtype=["bf16"|"fp16"|"fp32"] +# } +# } + +DATA_TYPES = "data_types" +GRAD_ACCUM_DTYPE = "grad_accum_dtype" +GRAD_ACCUM_DTYPE_DEFAULT = None + +######################################### +# Drop the last incomplete Batch +# ######################################### +# dataloader_drop_last. By default, this feature is not enabled. +# Users can configure in ds_config.json as below example: +DATALOADER_DROP_LAST_FORMAT = ''' +The last incomplete batch can be dropped by setting: +"dataloader_drop_last": True +''' +DATALOADER_DROP_LAST = "dataloader_drop_last" +DATALOADER_DROP_LAST_DEFAULT = False + +######################################### +# PIPELINE PARALLELISM +######################################### +PIPE_REPLICATED = 'ds_pipe_replicated' + +######################################### +# DATA PARALLELISM +######################################### +DATA_PARALLEL_GROUP = "data_parallel_group" +GLOBAL_RANK = "global_rank" + +######################################### +# EXPERT-DATA PARALLELISM TOPO Config +######################################### +USE_DATA_BEFORE_EXPERT_PARALLEL = "use_data_before_expert_parallelism" +USE_DATA_BEFORE_EXPERT_PARALLEL_DEFAULT = False diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc60faa292a853bdceec0e3dad54ddbd5b2d4dad Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86fce3ea7d66af7530894c26654bd8cef882f4f5 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e06745daae1a233dfb1f58eba444c4b4bca8c7a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87aa04852c9d3e4d84e579c034ec6977c237044d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/__pycache__/curriculum_scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/config.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/config.py new file mode 100644 index 0000000000000000000000000000000000000000..690ce97034e452c2dbb451c3f71ebfaa579ad02d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/config.py @@ -0,0 +1,193 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .constants import * +import copy +from ..config_utils import get_scalar_param + + +# TODO: Reducing config verbosity by returning None or {} when disabled. +# One challenge is that we still need to somehow include the default values, +# for example the *_ENABLED has default of false. +def get_data_efficiency_config(param_dict): + output = {} + output[DATA_EFFICIENCY_ENABLED] = get_data_efficiency_enabled(param_dict) + output[DATA_EFFICIENCY_SEED] = get_data_efficiency_seed(param_dict) + if DATA_EFFICIENCY not in param_dict.keys(): + param_dict[DATA_EFFICIENCY] = {} + sub_param_dict = param_dict[DATA_EFFICIENCY] + output[DATA_SAMPLING] = get_data_sampling(sub_param_dict) + output[DATA_ROUTING] = get_data_routing(sub_param_dict) + return output + + +def get_data_efficiency_enabled(param_dict): + if DATA_EFFICIENCY in param_dict.keys(): + return get_scalar_param(param_dict[DATA_EFFICIENCY], DATA_EFFICIENCY_ENABLED, DATA_EFFICIENCY_ENABLED_DEFAULT) + else: + return False + + +def get_data_efficiency_seed(param_dict): + if DATA_EFFICIENCY in param_dict.keys(): + return get_scalar_param(param_dict[DATA_EFFICIENCY], DATA_EFFICIENCY_SEED, DATA_EFFICIENCY_SEED_DEFAULT) + else: + return DATA_EFFICIENCY_SEED_DEFAULT + + +def get_data_sampling(param_dict): + sub_param_dict = param_dict.get(DATA_SAMPLING, {}) + output = copy.copy(sub_param_dict) + output[DATA_SAMPLING_ENABLED] = get_data_sampling_enabled(param_dict) + output[DATA_SAMPLING_NUM_EPOCHS] = get_data_sampling_num_epochs(param_dict) + output[DATA_SAMPLING_NUM_WORKERS] = get_data_sampling_num_workers(param_dict) + output[DATA_SAMPLING_PIN_MEMORY] = get_data_sampling_pin_memory(param_dict) + output[CURRICULUM_LEARNING] = get_curriculum_learning(sub_param_dict) + output[DYNAMIC_BATCHING] = get_dynamic_batching(sub_param_dict) + return output + + +def get_data_sampling_enabled(param_dict): + if DATA_SAMPLING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_ENABLED, DATA_SAMPLING_ENABLED_DEFAULT) + else: + return False + + +def get_data_sampling_num_epochs(param_dict): + if DATA_SAMPLING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_NUM_EPOCHS, DATA_SAMPLING_NUM_EPOCHS_DEFAULT) + else: + return DATA_SAMPLING_NUM_EPOCHS_DEFAULT + + +def get_data_sampling_num_workers(param_dict): + if DATA_SAMPLING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_NUM_WORKERS, + DATA_SAMPLING_NUM_WORKERS_DEFAULT) + else: + return DATA_SAMPLING_NUM_WORKERS_DEFAULT + + +def get_data_sampling_pin_memory(param_dict): + if DATA_SAMPLING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_SAMPLING], DATA_SAMPLING_PIN_MEMORY, DATA_SAMPLING_PIN_MEMORY_DEFAULT) + else: + return DATA_SAMPLING_PIN_MEMORY_DEFAULT + + +def get_curriculum_learning(param_dict): + output = {} + output[CURRICULUM_LEARNING_ENABLED] = get_curriculum_learning_enabled(param_dict) + if CURRICULUM_LEARNING not in param_dict.keys(): + param_dict[CURRICULUM_LEARNING] = {} + sub_param_dict = param_dict[CURRICULUM_LEARNING] + if output[CURRICULUM_LEARNING_ENABLED]: + assert CURRICULUM_LEARNING_METRICS in sub_param_dict.keys( + ), f"Curriculum learning is enabled, {CURRICULUM_LEARNING_METRICS} must be specified" + for key, val in get_curriculum_learning_params(param_dict).items(): + output[key] = val + return output + + +def get_dynamic_batching(param_dict): + output = copy.copy(param_dict.get(DYNAMIC_BATCHING, {})) + output[DYNAMIC_BATCHING_ENABLED] = bool(output.get(DYNAMIC_BATCHING_ENABLED, DYNAMIC_BATCHING_ENABLED_DEFAULT)) + output[DYNAMIC_BATCHING_LR_SCALING_METHOD] = str( + output.get(DYNAMIC_BATCHING_LR_SCALING_METHOD, DYNAMIC_BATCHING_LR_SCALING_METHOD_DEFAULT)) + output[DYNAMIC_BATCHING_MIN_BATCH_SIZE] = int( + output.get(DYNAMIC_BATCHING_MIN_BATCH_SIZE, DYNAMIC_BATCHING_MIN_BATCH_SIZE_DEFAULT)) + output[DYNAMIC_BATCHING_MAX_BATCH_SIZE] = int(output[DYNAMIC_BATCHING_MAX_BATCH_SIZE]) \ + if DYNAMIC_BATCHING_MAX_BATCH_SIZE in output.keys() \ + else DYNAMIC_BATCHING_MAX_BATCH_SIZE_DEFAULT + output[DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER] = str( + output.get(DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER, DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER_DEFAULT)) + if output[DYNAMIC_BATCHING_ENABLED]: + assert DYNAMIC_BATCHING_MAX_TOKENS in output.keys( + ), f"Dynamic batching is enabled, so {DYNAMIC_BATCHING_MAX_TOKENS} must be specified" + output[DYNAMIC_BATCHING_MAX_TOKENS] = int(output[DYNAMIC_BATCHING_MAX_TOKENS]) + output[DYNAMIC_BATCHING_VERBOSE] = bool(output.get(DYNAMIC_BATCHING_VERBOSE, False)) + return output + + +def get_curriculum_learning_enabled(param_dict): + if CURRICULUM_LEARNING in param_dict.keys(): + return get_scalar_param(param_dict[CURRICULUM_LEARNING], CURRICULUM_LEARNING_ENABLED, + CURRICULUM_LEARNING_ENABLED_DEFAULT) + else: + return False + + +def get_curriculum_learning_params(param_dict): + if CURRICULUM_LEARNING in param_dict.keys(): + curriculum_learning_params = copy.copy(param_dict[CURRICULUM_LEARNING]) + curriculum_learning_params.pop(CURRICULUM_LEARNING_ENABLED) + return curriculum_learning_params + else: + return {} + + +def get_curriculum_enabled_legacy(param_dict): + if CURRICULUM_LEARNING_LEGACY in param_dict.keys(): + return get_scalar_param(param_dict[CURRICULUM_LEARNING_LEGACY], CURRICULUM_ENABLED_LEGACY, + CURRICULUM_ENABLED_DEFAULT_LEGACY) + else: + return False + + +def get_curriculum_params_legacy(param_dict): + if CURRICULUM_LEARNING_LEGACY in param_dict.keys(): + curriculum_params = copy.copy(param_dict[CURRICULUM_LEARNING_LEGACY]) + curriculum_params.pop(CURRICULUM_ENABLED_LEGACY) + return curriculum_params + else: + return False + + +def get_data_routing(param_dict): + output = {} + output[DATA_ROUTING_ENABLED] = get_data_routing_enabled(param_dict) + if DATA_ROUTING not in param_dict.keys(): + param_dict[DATA_ROUTING] = {} + sub_param_dict = param_dict[DATA_ROUTING] + output[RANDOM_LTD] = get_random_ltd(sub_param_dict) + + return output + + +def get_data_routing_enabled(param_dict): + if DATA_ROUTING in param_dict.keys(): + return get_scalar_param(param_dict[DATA_ROUTING], DATA_ROUTING_ENABLED, DATA_ROUTING_ENABLED_DEFAULT) + else: + return False + + +def get_random_ltd(param_dict): + output = {} + output[RANDOM_LTD_ENABLED] = RANDOM_LTD_ENABLED_DEFAULT + output[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE] = {} + output[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][ + RANDOM_LTD_LAYER_TOKEN_LR_ENABLED] = RANDOM_LTD_LAYER_TOKEN_LR_ENABLED_DEFAULT + if get_random_ltd_enabled(param_dict): + output[RANDOM_LTD_ENABLED] = get_random_ltd_enabled(param_dict) + for key, val in get_random_ltd_params(param_dict).items(): + output[key] = val + return output + + +def get_random_ltd_enabled(param_dict): + if RANDOM_LTD in param_dict.keys(): + return get_scalar_param(param_dict[RANDOM_LTD], RANDOM_LTD_ENABLED, RANDOM_LTD_ENABLED_DEFAULT) + else: + return False + + +def get_random_ltd_params(param_dict): + if RANDOM_LTD in param_dict.keys(): + random_ltd_params = copy.copy(param_dict[RANDOM_LTD]) + random_ltd_params.pop(RANDOM_LTD_ENABLED) + return random_ltd_params + else: + return {} diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/constants.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..73cc69c1f606a3c77e6969cc32c57ef7cd8e5111 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/constants.py @@ -0,0 +1,136 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Data efficiency library + See sample config at https://www.deepspeed.ai/docs/config-json/data-efficiency +""" +DATA_EFFICIENCY = "data_efficiency" +DATA_EFFICIENCY_ENABLED = "enabled" +DATA_EFFICIENCY_ENABLED_DEFAULT = False +DATA_EFFICIENCY_SEED = "seed" +DATA_EFFICIENCY_SEED_DEFAULT = 1234 + +######################################### +# Data efficiency - Data Sampling +######################################### +DATA_SAMPLING = "data_sampling" +DATA_SAMPLING_ENABLED = "enabled" +DATA_SAMPLING_ENABLED_DEFAULT = False +DATA_SAMPLING_NUM_EPOCHS = "num_epochs" +DATA_SAMPLING_NUM_EPOCHS_DEFAULT = 1000 +DATA_SAMPLING_NUM_WORKERS = "num_workers" +DATA_SAMPLING_NUM_WORKERS_DEFAULT = 0 +DATA_SAMPLING_PIN_MEMORY = "pin_memory" +DATA_SAMPLING_PIN_MEMORY_DEFAULT = False + +######################################### +# Data efficiency - Data Sampling - Curriculum Learning +######################################### +CURRICULUM_LEARNING = "curriculum_learning" +CURRICULUM_LEARNING_ENABLED = "enabled" +CURRICULUM_LEARNING_ENABLED_DEFAULT = False +CURRICULUM_LEARNING_CLUSTER_PATH = "data_cluster_path" +CURRICULUM_LEARNING_METRICS = "curriculum_metrics" +CURRICULUM_LEARNING_SAMPLE_PATH = "index_to_sample_path" +CURRICULUM_LEARNING_METRIC_PATH = "index_to_metric_path" +CURRICULUM_LEARNING_CLUSTERING_TYPE = "clustering_type" +CURRICULUM_LEARNING_SINGLE_CLUSTER = "single_cluster" +CURRICULUM_LEARNING_CLUSTER_PREFIX = "cluster" +CURRICULUM_LEARNING_DIFFICULTY_TYPE = "difficulty_type" +CURRICULUM_LEARNING_VALUE_BASED = "value" +CURRICULUM_LEARNING_PERCENTILE_BASED = "percentile" +CURRICULUM_LEARNING_MIN_DIFFICULTY = "min_difficulty" +CURRICULUM_LEARNING_MAX_DIFFICULTY = "max_difficulty" +CURRICULUM_LEARNING_SCHEDULE_TYPE = "schedule_type" +CURRICULUM_LEARNING_SCHEDULE_CONFIG = "schedule_config" +CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY = "difficulty" +CURRICULUM_LEARNING_SCHEDULE_MAX_STEP = "max_step" +CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP = "total_curriculum_step" +CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP = "difficulty_step" +CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE = "root_degree" +CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE = "fixed_discrete" +CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT = "fixed_root" +CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR = "fixed_linear" +CURRICULUM_LEARNING_SCHEDULE_CUSTOM = "custom" +CURRICULUM_LEARNING_CURRENT_DIFFICULTY = "current_difficulty" + +CURRICULUM_LEARNING_BATCH = "batch" +CURRICULUM_LEARNING_CONSUMED_SAMPLES = "consumed_samples" +CURRICULUM_LEARNING_STEP = "curriculum_step" +CURRICULUM_LEARNING_CURRENT_DIFFICULTIES = "current_difficulties" +CURRICULUM_LEARNING_DATA_CLUSTER_PATHS = "data_cluster_paths" +CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION = "data_cluster_current_position" +CURRICULUM_LEARNING_NP_RNG_STATE = "np_rng_state" + +######################################### +# Data efficiency - Dynamic batching and LR scaling +######################################### +DYNAMIC_BATCHING = "dynamic_batching" +DYNAMIC_BATCHING_ENABLED = "enabled" +DYNAMIC_BATCHING_ENABLED_DEFAULT = False +DYNAMIC_BATCHING_METRICS_PATH = "metrics_path" +DYNAMIC_BATCHING_LR_SCALING_METHOD = "lr_scaling_method" # "linear" / "sqrt" / "none" +DYNAMIC_BATCHING_LR_SCALING_METHOD_DEFAULT = "linear" +DYNAMIC_BATCHING_MIN_BATCH_SIZE = "min_batch_size" +DYNAMIC_BATCHING_MIN_BATCH_SIZE_DEFAULT = 1 +DYNAMIC_BATCHING_MAX_BATCH_SIZE = "max_batch_size" +DYNAMIC_BATCHING_MAX_BATCH_SIZE_DEFAULT = None +DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER = "sequence_picking_order" # "random" / "seqlen" / "dataloader" +DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER_DEFAULT = "dataloader" # "random" / "seqlen" / "dataloader" +DYNAMIC_BATCHING_MAX_TOKENS = "max_tokens" +DYNAMIC_BATCHING_VERBOSE = "verbose" + +######################################### +# Curriculum Learning legacy implementation +######################################### +CURRICULUM_LEARNING_LEGACY = "curriculum_learning" + +CURRICULUM_ENABLED_LEGACY = "enabled" +CURRICULUM_ENABLED_DEFAULT_LEGACY = False + +######################################### +# Data efficiency - Data Routing +######################################### +DATA_ROUTING = "data_routing" +DATA_ROUTING_ENABLED = "enabled" +DATA_ROUTING_ENABLED_DEFAULT = False + +######################################### +# Data efficiency - Data Routing - Random LTD +######################################### +RANDOM_LTD = "random_ltd" +RANDOM_LTD_ENABLED = "enabled" +RANDOM_LTD_ENABLED_DEFAULT = False + +RANDOM_LTD_MODEL_MASK_NAME = "model_mask_name" +RANDOM_LTD_MODEL_TYPE = "model_type" +RANDOM_LTD_MICRO_BATCH_SIZE = "micro_batch_size" +RANDOM_LTD_GLOBAL_BATCH_SIZE = "global_batch_size" +RANDOM_LTD_SAMPLE_INDEX = "sample_idx" +RANDOM_LTD_ATTENTION_MASK = "attention_mask" +RANDOM_LTD_HIDDEN_STATE_ORDER = "hidden_state_order" +RANDOM_LTD_LAYER_NUM = "random_ltd_layer_num" +RANDOM_LTD_LAYER_ID = "random_ltd_layer_id" +RANDOM_LTD_TOTAL_LAYER_NUM = "total_layer_num" +RANDOM_LTD_CONSUMED_LAYER_TOKENS = "consumed_layer_tokens" + +# scheduler +RANDOM_LTD_SCHEDULER = "random_ltd_schedule" +RANDOM_LTD_MAX_VALUE = "max_value" +RANDOM_LTD_MIN_VALUE = "min_value" +RANDOM_LTD_CURRENT_VALUE = "current_value" +RANDOM_LTD_SCHEDULE_CONFIG = "schedule_config" +RANDOM_LTD_INCREASE_STEP = "seq_per_step" +RANDOM_LTD_REQUIRE_STEP = "require_steps" +RANDOM_LTD_SCHEDULER_TYPE = "schedule_type" +RANDOM_LTD_CURR_STEP = "current_steps" + +# learning rate schedulers +RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE = "layer_token_lr_schedule" +RANDOM_LTD_LAYER_TOKEN_LR_ENABLED = "enabled" +RANDOM_LTD_LAYER_TOKEN_LR_ENABLED_DEFAULT = False +RANDOM_LTD_TOTAL_LAYER_TOKENS = "total_layer_tokens" +RANDOM_LTD_WARMUP_TYPE = "warmup_type" +RANDOM_LTD_WARMUP_LAYER_TOKENS = "warmup_layer_tokens" diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..23d747957dc4647e06fad0a94e5e4b071b6f6e23 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/curriculum_scheduler.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +from deepspeed.utils import logger +from .constants import * + + +class CurriculumScheduler(object): + + def __init__(self, config): + super().__init__() + self.state = {} + assert CURRICULUM_LEARNING_MIN_DIFFICULTY in config, \ + f"Curriculum learning requires the config '{CURRICULUM_LEARNING_MIN_DIFFICULTY}'" + assert CURRICULUM_LEARNING_MAX_DIFFICULTY in config, \ + f"Curriculum learning requires the config '{CURRICULUM_LEARNING_MAX_DIFFICULTY}'" + assert CURRICULUM_LEARNING_SCHEDULE_TYPE in config, \ + f"Curriculum learning requires the config '{CURRICULUM_LEARNING_SCHEDULE_TYPE}'" + self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY] = config[CURRICULUM_LEARNING_MIN_DIFFICULTY] + self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY] = config[CURRICULUM_LEARNING_MAX_DIFFICULTY] + self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = config[CURRICULUM_LEARNING_MIN_DIFFICULTY] + self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] = config[CURRICULUM_LEARNING_SCHEDULE_TYPE] + self.first_step = True + if config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE: + """ + The schedule_config is a list of difficulty and a list of max + step belonging to each difficulty. Example json config: + "schedule_config": { + "difficulty": [1,2,3], + "max_step": [5,10] + } + The "max_step" has one less element than "difficulty", because + the last difficulty will be used for all following steps. + The self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] is a dictionary of + difficulty : [max step for this difficulty, next difficulty]. + """ + assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_discrete schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY}'" + assert CURRICULUM_LEARNING_SCHEDULE_MAX_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_discrete schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_MAX_STEP}'" + assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_MAX_STEP]) > 0 + assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY]) > 0 + assert len(config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY]) == len( + config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_MAX_STEP]) + 1 + self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT: + """ + The schedule_config includes: + total_curriculum_step: how many steps the curriculum learning takes to go + from min difficulty to max difficulty. + difficulty_step: the difficulty level determined every time must + be a multiple of this difficulty_step. This is used to determine + the step of difficulty increase, and to ensure the use of NVIDIA + Tensor Core acceleration (requires multiple of 8 (FP16) or + 16 (INT8)). + root_degree: the degree of the root function. Degree of 2 means + square root and degree of 3 means cube root. Degree of 1 is + equivalent to linear. + "schedule_config": { + "total_curriculum_step": 30000, + "difficulty_step": 8, + "root_degree": 2 + } + """ + assert CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP}'" + assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP}'" + assert CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_root schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE}'" + if config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP] % 8 != 0: + logger.warning( + f'When using seqlen metric, the difficulty_step for curriculum learning has to be multiple of 8 (for FP16 data) or 16 (for INT8 data) to enable NVIDIA Tensor Core acceleration. Disregard this warning if this is unrelated to your metric/hardware.' + ) + self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR: + """ + The schedule_config is the same as CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT but without the + root_degree. + "schedule_config": { + "total_curriculum_step": 30000, + "difficulty_step": 8 + } + """ + assert CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_linear schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP}'" + assert CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP in config[CURRICULUM_LEARNING_SCHEDULE_CONFIG], \ + f"Curriculum learning with fixed_linear schedule requires the schedule_config '{CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP}'" + if config[CURRICULUM_LEARNING_SCHEDULE_CONFIG][CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP] % 8 != 0: + logger.warning( + f'When using seqlen metric, the difficulty_step for curriculum learning has to be multiple of 8 (for FP16 data) or 16 (for INT8 data) to enable NVIDIA Tensor Core acceleration. Disregard this warning if this is unrelated to your metric/hardware.' + ) + self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] = config[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + elif config[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_CUSTOM: + """ + Fully customized schedule. User need to provide a custom schedule + function by using the set_custom_curriculum_learning_schedule API + in deepspeed/runtime/engine.py + """ + self.custom_get_difficulty = None + else: + raise RuntimeError('Unsupported curriculum schedule type') + + def get_current_difficulty(self): + return self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] + + def set_current_difficulty(self, difficulty): + self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = difficulty + + def set_custom_get_difficulty(self, schedule_function): + self.custom_get_difficulty = schedule_function + + def get_state(self): + return self.state + + def set_state(self, state): + self.state = state + + def __fixed_discrete_get_difficulty(self, global_steps): + s_state = self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + if global_steps > s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP][-1]: + return s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY][-1] + for i in range(len(s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP])): + if global_steps <= s_state[CURRICULUM_LEARNING_SCHEDULE_MAX_STEP][i]: + return s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY][i] + + def __fixed_root_get_difficulty(self, global_steps, root_degree=None): + s_state = self.state[CURRICULUM_LEARNING_SCHEDULE_CONFIG] + if root_degree is None: + root_degree = s_state[CURRICULUM_LEARNING_SCHEDULE_ROOT_DEGREE] + next_difficulty = (float(global_steps) / s_state[CURRICULUM_LEARNING_SCHEDULE_TOTAL_STEP])**(1.0 / root_degree) + next_difficulty = math.floor( + next_difficulty * + (self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY] - self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY]) + + self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY]) + next_difficulty -= (next_difficulty % s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP]) + next_difficulty = min(next_difficulty, self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY]) + return next_difficulty + + def get_difficulty(self, global_steps): + if self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_DISCRETE: + return self.__fixed_discrete_get_difficulty(global_steps) + elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_LINEAR: + return self.__fixed_root_get_difficulty(global_steps, 1) + elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_FIXED_ROOT: + return self.__fixed_root_get_difficulty(global_steps) + elif self.state[CURRICULUM_LEARNING_SCHEDULE_TYPE] == CURRICULUM_LEARNING_SCHEDULE_CUSTOM: + return self.custom_get_difficulty(global_steps) + else: + raise RuntimeError('Unsupported curriculum schedule type') + + def update_difficulty(self, global_steps): + if self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] < self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY]: + self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] = self.get_difficulty(global_steps) + return self.state[CURRICULUM_LEARNING_CURRENT_DIFFICULTY] diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4e1652dce44c60634723b204f489081a093f38d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de7746d1267f526866bc1c0fccfe2a1b86cd08c4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/basic_layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..682563b3d65a151732b9c7bd7222744a35689d09 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/helper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..595921a23c7db84531af6b4bb583c0c1395e2695 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b1a6eeb281906b2a859f959376427cfbe86f645 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..780a88c3d5a2fb016f2c59bbe3a88fbaff6499b0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/basic_layer.py @@ -0,0 +1,113 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.utils import logger +from torch import Tensor +from torch.nn import Module +from ..constants import * +from deepspeed.ops.random_ltd.dropping_utils import gpt_sample_tokens, bert_sample_tokens, GatherTokens, ScatterTokens + + +#####based on the paper random-ltd: https://arxiv.org/abs/2211.11586 +class RandomLayerTokenDrop(Module): + """ + A layer wrapper for random LTD + """ + + def __init__(self, layer: Module): + super(RandomLayerTokenDrop, self).__init__() + self.random_ltd_layer = layer + self.reserved_length = None #config['max_value'] + self.random_ltd_scheduler = None + self.max_length = None + self.reserved_length = -1 + self.curr_seq = -1 + self.batch_first = False + + def init_config(self, config, scheduler, random_ltd_layer_id): + self.random_ltd_scheduler = scheduler + self.random_ltd_layer_id = random_ltd_layer_id + self.max_length = self.random_ltd_scheduler.state[RANDOM_LTD_MAX_VALUE] + + self.mask_name = config[RANDOM_LTD_MODEL_MASK_NAME] + self.micro_bs = config[RANDOM_LTD_MICRO_BATCH_SIZE] + self.random_ltd_num_layer = self.random_ltd_scheduler.random_ltd_layer_num + hs_order = config[RANDOM_LTD_HIDDEN_STATE_ORDER] + self.model_type = config[RANDOM_LTD_MODEL_TYPE] + + if hs_order == 'batch_seq_dim': + self.get_hidden_tensor_shape = self.get_bsh + self.batch_first = True + elif hs_order == 'seq_batch_dim': + self.get_hidden_tensor_shape = self.get_sbh + self.batch_first = False + else: + logger.warning( + "************For now, we only support batch_seq_dim or seq_batch_dim inputs. You can easily \ + your own input dimension orders************") + raise NotImplementedError + + if self.model_type == 'encoder': + self.index_generator = bert_sample_tokens + elif self.model_type == 'decoder': + self.index_generator = gpt_sample_tokens + else: + logger.warning("************For now, we only support encoder-only or decoder-only models************") + raise NotImplementedError + + def get_bsh(self, hidden_stats): + self.curr_seq, self.curr_micro_batch = hidden_stats.size()[1], hidden_stats.size()[0] + + def get_sbh(self, hidden_stats): + self.curr_seq, self.curr_micro_batch = hidden_stats.size()[0], hidden_stats.size()[1] + + def forward(self, hidden_states, **kwargs) -> Tensor: + if self.random_ltd_scheduler is not None: + self.reserved_length = self.random_ltd_scheduler.get_current_seq() + self.get_hidden_tensor_shape(hidden_states) + if self.training and self.random_ltd_scheduler is not None and self.reserved_length < self.curr_seq: + if self.mask_name is not None: + mask = kwargs[self.mask_name] + else: + mask = None + if self.random_ltd_layer_id == 0: + sampled_indices, part_attention_mask = self.index_generator(self.reserved_length,\ + self.curr_seq, \ + self.curr_micro_batch, \ + self.random_ltd_num_layer, \ + hidden_states.device, mask) + self.random_ltd_scheduler.state[RANDOM_LTD_SAMPLE_INDEX] = sampled_indices + self.random_ltd_scheduler.state[RANDOM_LTD_ATTENTION_MASK] = part_attention_mask + else: + sampled_indices = self.random_ltd_scheduler.state[RANDOM_LTD_SAMPLE_INDEX] + part_attention_mask = self.random_ltd_scheduler.state[RANDOM_LTD_ATTENTION_MASK] + + hidden_states, part_hidden_states = GatherTokens.apply(hidden_states, + sampled_indices[self.random_ltd_layer_id, :, :], + self.batch_first) + if self.mask_name is not None: + if self.model_type == 'encoder': + kwargs[self.mask_name] = part_attention_mask[self.random_ltd_layer_id] + else: + kwargs[self.mask_name] = part_attention_mask + + outputs = self.random_ltd_layer(part_hidden_states, **kwargs) + + if isinstance(outputs, tuple): + hidden_states = ScatterTokens.apply(hidden_states, outputs[0], + sampled_indices[self.random_ltd_layer_id, :, :], self.batch_first) + my_list = list(outputs) + my_list[0] = hidden_states + return tuple(my_list) + elif isinstance(outputs, Tensor): + hidden_states = ScatterTokens.apply(hidden_states, outputs, + sampled_indices[self.random_ltd_layer_id, :, :], self.batch_first) + return hidden_states + else: + logger.warning("************For now, we only support tuple and tensor output. \ + You need to adjust the output according to the layer in your model************") + raise NotImplementedError + else: + return self.random_ltd_layer(hidden_states, **kwargs) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py new file mode 100644 index 0000000000000000000000000000000000000000..150182d77bcfda20b1aa1aabd4b8785542ca9d1b --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/helper.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .basic_layer import RandomLayerTokenDrop +from collections import OrderedDict +from deepspeed.compression.helper import recursive_getattr, recursive_setattr + + +def convert_to_random_ltd(model, convert_type): + 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 = RandomLayerTokenDrop(old_module) + recursive_setattr(c_model, name, new_module) + + model.random_ltd_initialize() + return model + + +def save_without_random_ltd(model): + if hasattr(model, 'module'): + c_model = model.module + else: + c_model = model + + model_dic = c_model.state_dict() + return remove_random_ltd_state_dict(model_dic) + + +def remove_random_ltd_state_dict(state_dict): + new_state_dict = OrderedDict() + for key, value in state_dict.items(): + if '.random_ltd_layer' in key: + new_key = ''.join(key.split('.random_ltd_layer')) + else: + new_key = key + new_state_dict[new_key] = value + return new_state_dict diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..8a0b031d4f633976f438d5151973e58afa77712e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/scheduler.py @@ -0,0 +1,107 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math + +from deepspeed.utils import logger +# from deepspeed.runtime.lr_schedules import WarmupLR +from ..constants import * + +#####based on the paper random-ltd: https://arxiv.org/abs/2211.11586 + + +class BaseScheduler(object): + + def __init__(self): + self.state = {} + + def __fixed_root_get_value(self, global_steps, root_degree=None): + s_state = self.state[RANDOM_LTD_SCHEDULE_CONFIG] + if root_degree is None: + root_degree = s_state['root_degree'] + next_seq = (float(global_steps) / s_state[RANDOM_LTD_REQUIRE_STEP])**(1.0 / root_degree) + next_seq = math.floor(next_seq * (self.state[RANDOM_LTD_MAX_VALUE] - self.state[RANDOM_LTD_MIN_VALUE]) + + self.state[RANDOM_LTD_MIN_VALUE]) + next_seq -= (next_seq % s_state[RANDOM_LTD_INCREASE_STEP]) + next_seq = min(next_seq, self.state[RANDOM_LTD_MAX_VALUE]) + return next_seq + + def get_value(self, global_steps): + if self.state[RANDOM_LTD_SCHEDULER_TYPE] == 'fixed_linear': + return self.__fixed_root_get_value(global_steps, 1) + else: + raise RuntimeError('Unsupported random LTD schedule type') + + +class RandomLTDScheduler(BaseScheduler): + + def __init__(self, config): + super().__init__() + self.model_layer_num = config[RANDOM_LTD_TOTAL_LAYER_NUM] + self.random_ltd_layer_num = config[RANDOM_LTD_LAYER_NUM] + self.config_schedule = config[RANDOM_LTD_SCHEDULER] + self.global_batch_size = config[RANDOM_LTD_GLOBAL_BATCH_SIZE] + self.reset_to_init() + + if config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]: + logger.warning("**********Work In Progress************") + raise NotImplementedError + + self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = 0 + + # self.first_step = True + def get_total_layer_tokens(self, train_iters): + for step in range(train_iters): + self.update_seq(step) + return self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] + + def reset_to_init(self): + if self.config_schedule is not None: + self.state[RANDOM_LTD_MIN_VALUE] = self.config_schedule[RANDOM_LTD_MIN_VALUE] + self.state[RANDOM_LTD_MAX_VALUE] = self.config_schedule[RANDOM_LTD_MAX_VALUE] + self.state[RANDOM_LTD_CURRENT_VALUE] = self.config_schedule[RANDOM_LTD_MIN_VALUE] + self.state[RANDOM_LTD_SCHEDULE_CONFIG] = self.config_schedule[RANDOM_LTD_SCHEDULE_CONFIG] + self.state[RANDOM_LTD_SCHEDULER_TYPE] = self.config_schedule[RANDOM_LTD_SCHEDULER_TYPE] + self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = 0 + self.state[RANDOM_LTD_CURR_STEP] = -1 + + def get_current_seq(self): + return self.state[RANDOM_LTD_CURRENT_VALUE] + + def set_current_seq(self, seq_length): + self.state[RANDOM_LTD_CURRENT_VALUE] = seq_length + + def get_random_ltd_layer_num(self): + return self.random_ltd_layer_num + + def get_state(self): + return self.state + + def set_state(self, state): + self.state = state + + def update_seq(self, global_steps): + if self.state[RANDOM_LTD_CURRENT_VALUE] < self.state[RANDOM_LTD_MAX_VALUE]: + self.state[RANDOM_LTD_CURRENT_VALUE] = self.get_value(global_steps) + if global_steps != self.state[RANDOM_LTD_CURR_STEP]: + self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] += self.global_batch_size*(self.state[RANDOM_LTD_CURRENT_VALUE] * self.random_ltd_layer_num \ + + self.state[RANDOM_LTD_MAX_VALUE] * (self.model_layer_num - self.random_ltd_layer_num)) + self.state[RANDOM_LTD_CURR_STEP] = global_steps + + def state_dict(self): + return { + RANDOM_LTD_CONSUMED_LAYER_TOKENS: self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS], + RANDOM_LTD_CURR_STEP: self.state[RANDOM_LTD_CURR_STEP], + RANDOM_LTD_CURRENT_VALUE: self.state[RANDOM_LTD_CURRENT_VALUE], + RANDOM_LTD_MIN_VALUE: self.state[RANDOM_LTD_MIN_VALUE], + RANDOM_LTD_MAX_VALUE: self.state[RANDOM_LTD_MAX_VALUE], + } + + def load_state_dict(self, state_dict): + self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = state_dict[RANDOM_LTD_CONSUMED_LAYER_TOKENS] + self.state[RANDOM_LTD_CURR_STEP] = state_dict[RANDOM_LTD_CURR_STEP] + self.state[RANDOM_LTD_CURRENT_VALUE] = state_dict[RANDOM_LTD_CURRENT_VALUE] + self.state[RANDOM_LTD_MIN_VALUE] = state_dict[RANDOM_LTD_MIN_VALUE] + self.state[RANDOM_LTD_MAX_VALUE] = state_dict[RANDOM_LTD_MAX_VALUE] diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..afcfef2ef4dc04279b7fa3ee49916519f33ffbec --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_routing/utils.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch + + +def bsh_decoder_gather(reserved_length, hidden_states, mask): + # random-layer-token-drop + rand_list = [] + part_hidden_states = [] # batch, seq, hidden ## different from megatron + for k in range(hidden_states.size(0)): + B_tmp = torch.randperm(hidden_states.size(1), device=hidden_states.device)[:reserved_length] + B = B_tmp.sort()[0] + rand_list.append(B) + part_hidden_states.append(hidden_states[k:k + 1, B, :]) + + part_hidden_states = torch.cat(part_hidden_states, dim=0) + part_mask = mask[:, :, :reserved_length, :reserved_length] + return part_hidden_states, rand_list, part_mask + + +def bsh_decoder_scatter(hidden_states, part_hidden_states, rand_list): + for k in range(hidden_states.size(0)): + hidden_states[k, rand_list[k], :] = part_hidden_states[k, :, :] + return hidden_states diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84aa2056dfbef98341283f8032a1cbac48152164 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71b8bc745271bbce551150e6ae347239d89a335a Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_analyzer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc47bf240feb87e5c4b198d7761c882b1b8da8b4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/data_sampler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3226f7c458d71732b8810844d66a16297771e96f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/indexed_dataset.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fe9a1773fd9a1ae360c26365d2c8811fcb3738c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b1be8069a290bedc0d8fea02b5f5c25e79f91a0c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/__pycache__/variable_batch_size_and_lr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py new file mode 100644 index 0000000000000000000000000000000000000000..f82c684ec6e229597736831d7ea56ca8c3bdbaa0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py @@ -0,0 +1,885 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import sys +from collections import defaultdict +import csv +import time +from multiprocessing import Process, Manager +import numpy as np +import torch +from torch.utils.data import BatchSampler, SequentialSampler, DataLoader, Subset + +import deepspeed.comm as dist +from deepspeed.utils import logger +from deepspeed.runtime.data_pipeline.data_sampling.indexed_dataset import MMapIndexedDataset, valid_dtypes +from deepspeed.runtime.data_pipeline.data_sampling.utils import split_dataset, split_index, create_mmap_dataset_builder, close_mmap_dataset_builder, find_fit_int_dtype + + +class DataAnalyzer(object): + + def __init__(self, + dataset, + num_workers=1, + worker_id=0, + num_threads=1, + num_threads_reduce=1, + specific_threads=[], + batch_size=1, + metric_names=[], + metric_functions=[], + metric_types=[], + metric_dtypes=[], + save_path="./", + collate_fn=None, + custom_map_init=None, + custom_map_update=None, + custom_map_finalize=None, + custom_reduce=None, + sample_indices=None): + super().__init__() + self.dataset = dataset + self.num_workers = num_workers + self.worker_id = worker_id + self.num_threads = num_threads + self.num_threads_reduce = num_threads_reduce + self.specific_threads = specific_threads + self.batch_size = batch_size + self.metric_names = metric_names + self.metric_functions = metric_functions + self.metric_types = metric_types + self.metric_dtypes = metric_dtypes + self.save_path = save_path + self.collate_fn = collate_fn + self.custom_map_init = custom_map_init + self.custom_map_update = custom_map_update + self.custom_map_finalize = custom_map_finalize + self.custom_reduce = custom_reduce + self.sample_indices = sample_indices + + def init_metric_results(self, thread_id, metric_names, metric_types, metric_dtypes, save_path, worker_id): + metric_results = [] + for m_idx in range(len(metric_names)): + metric_name, metric_type, metric_dtype = metric_names[m_idx], \ + metric_types[m_idx], metric_dtypes[m_idx] + assert metric_dtype in valid_dtypes, f"metric_dtype {metric_dtype} not supported. Supported dtypes {valid_dtypes}" + metric_save_path = f"{save_path}/{metric_name}/worker{worker_id}_thread{thread_id}/" + os.makedirs(metric_save_path, exist_ok=True) + if metric_type == 'single_value_per_sample': + sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" + sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_dtype) + metric_to_sample_fname = f"{metric_save_path}/{metric_name}_metric_to_sample" + os.system(f"rm -rf {metric_to_sample_fname}*") + metric_to_sample_dict = defaultdict(list) + metric_results.append({ + "sample_to_metric_fname": sample_to_metric_fname, + "sample_to_metric_builder": sample_to_metric_builder, + "metric_to_sample_fname": metric_to_sample_fname, + "metric_to_sample_dict": metric_to_sample_dict + }) + elif metric_type == 'accumulate_value_over_samples': + metric_value = None + metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" + metric_results.append({"metric_value": metric_value, "metric_value_fname": metric_value_fname}) + return metric_results + + def update_metric_results(self, + data, + metric_types, + metric_dtypes, + metric_functions, + metric_results, + batch_start_idx=0): + for m_idx in range(len(metric_types)): + metric_type, metric_dtype, metric_function, metric_result = metric_types[m_idx], \ + metric_dtypes[m_idx], metric_functions[m_idx], metric_results[m_idx] + metric_values = metric_function(data) + + assert torch.is_tensor(metric_values) or isinstance(metric_values, np.ndarray), \ + "metric_function must return a tensor or array" + assert metric_values.dtype == metric_dtype, \ + f"metric_function result dtype {metric_values.dtype} does not match metric_dtype {metric_dtype}" + if isinstance(metric_values, np.ndarray): + metric_values = torch.from_numpy(metric_values) + + if metric_type == 'single_value_per_sample': + for row in range(metric_values.size()[0]): + sample_idx = batch_start_idx + row # sample idx following dataset iteration order + if isinstance(data, dict) and 'index' in data: # Megatron use case, idx provided in 'index' field + sample_idx = data['index'][row][0].item() + elif self.sample_indices is not None: # user defined shuffling of indices + sample_idx = self.sample_indices[sample_idx] + metric_result["sample_to_metric_builder"].add_item(metric_values[row].reshape(-1)) + metric_result["metric_to_sample_dict"][metric_values[row].item()].append(sample_idx) + for m_value in metric_result["metric_to_sample_dict"]: + if len(metric_result["metric_to_sample_dict"][m_value]) > 100: + metric_fname = metric_result["metric_to_sample_fname"] + with open(f"{metric_fname}_{m_value}.csv", 'a') as f: + writer = csv.writer(f) + writer.writerows([metric_result["metric_to_sample_dict"][m_value]]) + metric_result["metric_to_sample_dict"][m_value] = [] + elif metric_type == 'accumulate_value_over_samples': + if metric_result["metric_value"] is None: + metric_result["metric_value"] = metric_values + else: + metric_result["metric_value"].add_(metric_values) + + def finalize_metric_results(self, metric_types, metric_dtypes, metric_results): + for m_idx in range(len(metric_types)): + metric_type, metric_dtype, metric_result = metric_types[m_idx], \ + metric_dtypes[m_idx], metric_results[m_idx] + if metric_type == 'single_value_per_sample': + metric_fname = metric_result["sample_to_metric_fname"] + close_mmap_dataset_builder(metric_result["sample_to_metric_builder"], metric_fname) + for m_value in metric_result["metric_to_sample_dict"]: + if len(metric_result["metric_to_sample_dict"][m_value]) > 0: + metric_fname = metric_result["metric_to_sample_fname"] + with open(f"{metric_fname}_{m_value}.csv", 'a') as f: + writer = csv.writer(f) + writer.writerows([metric_result["metric_to_sample_dict"][m_value]]) + metric_result["metric_to_sample_dict"][m_value] = [] + elif metric_type == 'accumulate_value_over_samples': + if metric_result["metric_value"] is not None: + metric_value_builder = create_mmap_dataset_builder(metric_result["metric_value_fname"], + metric_dtype) + metric_value_builder.add_item(metric_result["metric_value"].reshape(-1)) + close_mmap_dataset_builder(metric_value_builder, metric_result["metric_value_fname"]) + + def run_map_helper(self, thread_id): + start_idx, end_idx = self.thread_splits[thread_id][0], \ + self.thread_splits[thread_id][1] + logger.info(f"worker {self.worker_id} thread {thread_id}: start working " \ + f"on data subset {start_idx} to {end_idx}") + thread_dataset = Subset(self.dataset, list(range(start_idx, end_idx))) + sampler = BatchSampler(SequentialSampler(thread_dataset), batch_size=self.batch_size, drop_last=False) + iterator = iter( + DataLoader(thread_dataset, + batch_sampler=sampler, + num_workers=0, + collate_fn=self.collate_fn, + pin_memory=False)) + if self.custom_map_init is None: + metric_results = self.init_metric_results(thread_id, self.metric_names, self.metric_types, + self.metric_dtypes, self.save_path, self.worker_id) + else: + metric_results = self.custom_map_init(thread_id, self.metric_names, self.metric_types, self.metric_dtypes, + self.save_path, self.worker_id) + total_sample = len(thread_dataset) + processed_sample = 0 + start = time.time() + while True: + try: + data = next(iterator) + batch_start_idx = start_idx + processed_sample + if self.custom_map_update is None: + self.update_metric_results(data, self.metric_types, self.metric_dtypes, self.metric_functions, + metric_results, batch_start_idx) + else: + self.custom_map_update(data, self.metric_types, self.metric_dtypes, self.metric_functions, + metric_results, batch_start_idx) + processed_sample += len(data) + duration = (time.time() - start) / 3600.0 + remain_duration = duration * total_sample / processed_sample - duration + logger.info( + f"worker {self.worker_id} thread {thread_id}: {processed_sample} " \ + f"out of {total_sample} processed in {duration:.2f} hr, " \ + f"estimated to finish in {remain_duration:.2f} hr") + except StopIteration: + logger.info(f"worker {self.worker_id} thread {thread_id}: reach end of file") + break + if self.custom_map_finalize is None: + self.finalize_metric_results(self.metric_types, self.metric_dtypes, metric_results) + else: + self.custom_map_finalize(self.metric_types, self.metric_dtypes, metric_results) + logger.info(f"worker {self.worker_id} thread {thread_id}: finished") + + def run_map(self): + self.worker_splits, self.thread_splits = split_dataset(self.dataset, self.num_workers, self.worker_id, + self.num_threads) + if len(self.specific_threads) > 0: + threads_to_run = self.specific_threads + else: + threads_to_run = list(range(self.num_threads)) + if self.num_threads > 1: + p = [] + for thread in threads_to_run: + p.append(Process(target=self.run_map_helper, args=(thread, ))) + p[thread].start() + + for thread in threads_to_run: + p[thread].join() + else: + assert self.num_threads == 1 + self.run_map_helper(0) + + def get_metric_value_percentiles(self, metric_name, num_sample_per_value, total_num_samples): + logger.info(f"Checking the value percentiles of metric {metric_name}...") + processed_samples = 0 + current_percentile = 5 + for key in sorted(num_sample_per_value.keys()): + processed_samples += num_sample_per_value[key] + if processed_samples >= total_num_samples * current_percentile / 100.0: + logger.info(f"Metric {metric_name} {current_percentile}th percentile: {key}") + current_percentile += 5 + + def merge_gather_map_stats(self, num_workers, num_threads, num_threads_reduce, t_idx_reduce, metric_save_path, + metric_name, return_dict): + results = [] + for w_idx in range(num_workers): + for t_idx in range(num_threads): + if (w_idx * num_threads + t_idx) % num_threads_reduce == t_idx_reduce: + w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" + w_sample_to_metric_fname = f"{w_metric_save_path}/{metric_name}_sample_to_metric" + w_sample_to_metric = MMapIndexedDataset(w_sample_to_metric_fname, skip_warmup=True) + unique_v = list(np.unique(w_sample_to_metric)) + sample_to_metric_count = len(w_sample_to_metric) + logger.info(f"Finished gathering map stats from worker {w_idx} thread {t_idx}.") + results.append([unique_v, sample_to_metric_count]) + return_dict[t_idx_reduce] = results + + def merge_sample_to_metric(self, t_idx_reduce, metric_save_path, metric_name, metric_value_dtype, + map_worker_thread): + sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric_thread{t_idx_reduce}" + sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_value_dtype) + for w_t in map_worker_thread: + w_metric_save_path = f"{metric_save_path}/worker{w_t[0]}_thread{w_t[1]}/" + w_sample_to_metric_fname = f"{w_metric_save_path}/{metric_name}_sample_to_metric" + w_data = MMapIndexedDataset(w_sample_to_metric_fname, skip_warmup=True) + for row in range(len(w_data)): + sample_to_metric_builder.add_item(torch.tensor(w_data[row].astype(np.int64), dtype=torch.long)) + logger.info(f"Finished merge_sample_to_metric from worker {w_t[0]} thread {w_t[1]}.") + close_mmap_dataset_builder(sample_to_metric_builder, sample_to_metric_fname) + + def merge_metric_to_sample(self, t_idx_reduce, metric_save_path, metric_name, sample_idx_dtype, metric_value_dtype, + unique_metric_values, num_workers, num_threads): + index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample_thread{t_idx_reduce}" + index_to_sample_builder = create_mmap_dataset_builder(index_to_sample_fname, sample_idx_dtype) + index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric_thread{t_idx_reduce}" + index_to_metric_builder = create_mmap_dataset_builder(index_to_metric_fname, metric_value_dtype) + for unique_v in unique_metric_values: + samples = [] + for w_idx in range(num_workers): + for t_idx in range(num_threads): + w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" + w_metric_to_sample_fname = f"{w_metric_save_path}/{metric_name}_metric_to_sample_{unique_v}.csv" + if os.path.isfile(w_metric_to_sample_fname): + with open(w_metric_to_sample_fname, 'r') as f: + datareader = csv.reader(f) + for row in datareader: + samples += [int(x) for x in row] + index_to_sample_builder.add_item(torch.tensor(samples, dtype=torch.long)) + index_to_metric_builder.add_item(torch.tensor([unique_v], dtype=torch.long)) + logger.info(f"Finished reducing metric {metric_name} value {unique_v}.") + close_mmap_dataset_builder(index_to_sample_builder, index_to_sample_fname) + close_mmap_dataset_builder(index_to_metric_builder, index_to_metric_fname) + + def merge_map_results(self, dataset, metric_names, metric_types, save_path, num_workers, num_threads, + num_threads_reduce): + total_num_samples = len(dataset) + sample_idx_dtype = find_fit_int_dtype(0, total_num_samples - 1) + logger.info( + f"Total number of data samples: {total_num_samples}. Will use {sample_idx_dtype} to store the sample indexes." + ) + for m_idx in range(len(metric_names)): + metric_name, metric_type = metric_names[m_idx], metric_types[m_idx] + if metric_type == 'single_value_per_sample': + metric_save_path = f"{save_path}/{metric_name}/" + sample_to_metric_count = 0 + unique_metric_values = set([]) + manager = Manager() + return_dict = manager.dict() + p = [] + for t_idx_reduce in range(num_threads_reduce): + p.append( + Process(target=self.merge_gather_map_stats, + args=( + num_workers, + num_threads, + num_threads_reduce, + t_idx_reduce, + metric_save_path, + metric_name, + return_dict, + ))) + p[t_idx_reduce].start() + for t_idx_reduce in range(num_threads_reduce): + p[t_idx_reduce].join() + for t_idx_reduce in range(num_threads_reduce): + results = return_dict[t_idx_reduce] + for res in results: + unique_metric_values = unique_metric_values.union(set(res[0])) + sample_to_metric_count += res[1] + value_max = max(unique_metric_values) + value_min = min(unique_metric_values) + assert sample_to_metric_count == total_num_samples, "The number of samples in map result files are not correct. It's possible that some map worker didn't finish successfully." + metric_value_dtype = find_fit_int_dtype(value_min, value_max) + logger.info( + f"Metric {metric_name} has values between {value_min} and {value_max}. Will use {metric_value_dtype} to store the metric values." + ) + + # sample_to_metric + map_worker_thread = [] + for w_idx in range(num_workers): + for t_idx in range(num_threads): + map_worker_thread.append([w_idx, t_idx]) + thread_splits = split_index(0, len(map_worker_thread), num_threads_reduce) + p = [] + for t_idx_reduce in range(num_threads_reduce): + start_idx, end_idx = thread_splits[t_idx_reduce][0], thread_splits[t_idx_reduce][1] + p.append( + Process(target=self.merge_sample_to_metric, + args=( + t_idx_reduce, + metric_save_path, + metric_name, + metric_value_dtype, + map_worker_thread[start_idx:end_idx], + ))) + p[t_idx_reduce].start() + for t_idx_reduce in range(num_threads_reduce): + p[t_idx_reduce].join() + + sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" + sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_value_dtype) + for t_idx_reduce in range(num_threads_reduce): + chunk_fname = f"{metric_save_path}/{metric_name}_sample_to_metric_thread{t_idx_reduce}" + logger.info(f"Merging file {chunk_fname}") + sample_to_metric_builder.merge_file_(chunk_fname) + close_mmap_dataset_builder(sample_to_metric_builder, sample_to_metric_fname) + sample_to_metric = MMapIndexedDataset(sample_to_metric_fname, skip_warmup=True) + assert len(sample_to_metric) == total_num_samples + + # metric_to_sample + unique_metric_values = list(sorted(unique_metric_values)) + thread_splits = split_index(0, len(unique_metric_values), num_threads_reduce) + p = [] + for t_idx_reduce in range(num_threads_reduce): + start_idx, end_idx = thread_splits[t_idx_reduce][0], thread_splits[t_idx_reduce][1] + p.append( + Process(target=self.merge_metric_to_sample, + args=( + t_idx_reduce, + metric_save_path, + metric_name, + sample_idx_dtype, + metric_value_dtype, + unique_metric_values[start_idx:end_idx], + num_workers, + num_threads, + ))) + p[t_idx_reduce].start() + for t_idx_reduce in range(num_threads_reduce): + p[t_idx_reduce].join() + index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample" + index_to_sample_builder = create_mmap_dataset_builder(index_to_sample_fname, sample_idx_dtype) + index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric" + index_to_metric_builder = create_mmap_dataset_builder(index_to_metric_fname, metric_value_dtype) + for t_idx_reduce in range(num_threads_reduce): + chunk_is_fname = f"{metric_save_path}/{metric_name}_index_to_sample_thread{t_idx_reduce}" + logger.info(f"Merging file {chunk_is_fname}") + index_to_sample_builder.merge_file_(chunk_is_fname) + chunk_im_fname = f"{metric_save_path}/{metric_name}_index_to_metric_thread{t_idx_reduce}" + logger.info(f"Merging file {chunk_im_fname}") + index_to_metric_builder.merge_file_(chunk_im_fname) + close_mmap_dataset_builder(index_to_sample_builder, index_to_sample_fname) + close_mmap_dataset_builder(index_to_metric_builder, index_to_metric_fname) + + num_sample_per_value = DataAnalyzer.output_index_to_sample_percentile( + index_to_sample_fname, index_to_metric_fname, metric_name, metric_save_path, total_num_samples, + sample_idx_dtype) + self.get_metric_value_percentiles(metric_name, num_sample_per_value, total_num_samples) + elif metric_type == 'accumulate_value_over_samples': + metric_save_path = f"{save_path}/{metric_name}/" + metric_value = None + for w_idx in range(num_workers): + for t_idx in range(num_threads): + w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/" + w_metric_value_fname = f"{w_metric_save_path}/{metric_name}_metric_value" + w_metric_value = MMapIndexedDataset(w_metric_value_fname, skip_warmup=True) + if metric_value is None: + metric_value = np.copy(w_metric_value[0]) + else: + metric_value += np.copy(w_metric_value[0]) + value_max = int(max(metric_value)) + value_min = int(min(metric_value)) + metric_value_dtype = find_fit_int_dtype(value_min, value_max) + metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" + metric_value_builder = create_mmap_dataset_builder(metric_value_fname, metric_value_dtype) + metric_value_builder.add_item(torch.tensor(metric_value.astype(np.int64), dtype=torch.long)) + close_mmap_dataset_builder(metric_value_builder, metric_value_fname) + + @staticmethod + def output_index_to_sample_percentile(index_to_sample_fname, index_to_metric_fname, metric_name, metric_save_path, + total_num_samples, sample_idx_dtype): + """ read index_to_metric and index_to_sample files and write distribution to index_to_sample_percentage_merged """ + num_sample_per_value = {} + index_to_sample = MMapIndexedDataset(index_to_sample_fname, skip_warmup=True) + index_to_metric = MMapIndexedDataset(index_to_metric_fname, skip_warmup=True) + index_to_sample_merged_fname = f"{metric_save_path}/{metric_name}_index_to_sample_percentile_merged" + index_to_sample_merged_builder = create_mmap_dataset_builder(index_to_sample_merged_fname, sample_idx_dtype) + for v_idx in range(len(index_to_sample)): + if v_idx > 0: + assert index_to_metric[v_idx] > index_to_metric[v_idx - 1] + num_sample_per_value[index_to_metric[v_idx][0]] = len(index_to_sample[v_idx]) + assert sum(list(num_sample_per_value.values())) == total_num_samples + merge_step = max(1, len(index_to_sample) // 100) + for v_idx in range(0, len(index_to_sample), merge_step): + merged_samples = np.copy( + np.concatenate(index_to_sample[v_idx:min(len(index_to_sample), (v_idx + merge_step))], axis=None)) + index_to_sample_merged_builder.add_item(torch.tensor(merged_samples.astype(np.int64), dtype=torch.long)) + logger.info(f"Finished merging index_to_sample {v_idx} to {v_idx+merge_step}.") + close_mmap_dataset_builder(index_to_sample_merged_builder, index_to_sample_merged_fname) + return num_sample_per_value + + def run_reduce(self): + if self.custom_reduce is None: + self.merge_map_results(self.dataset, self.metric_names, self.metric_types, self.save_path, + self.num_workers, self.num_threads, self.num_threads_reduce) + else: + self.custom_reduce(self.dataset, self.metric_names, self.metric_types, self.save_path, self.num_workers, + self.num_threads, self.num_threads_reduce) + + def run_map_reduce(self, comm_group=None): + self.run_map() + # wait for the mapping operation, where all nodes outputs their own (partial) result files + dist.barrier(group=comm_group) + if self.worker_id == 0: + self.run_reduce() + # wait for the reduce, where rank 0 merges all (partial) files. Dataset can then be used by all nodes. + dist.barrier(group=comm_group) + + +class DistributedDataAnalyzer(object): + + def __init__( + self, + dataset, + num_workers=1, + num_threads=1, + worker_id=0, + batch_size=1, + metric_names=[], + metric_functions=[], + metric_types=[], + save_path="./", + collate_fn=None, + device='cuda', + comm_group=None, + sample_indices=None, + ) -> None: + self.dataset = dataset + self.batch_size = batch_size + self.metric_names = metric_names + self.metric_functions = metric_functions + self.metric_types = metric_types + self.save_path = save_path + self.collate_fn = collate_fn + self.device = device + self.sample_indices = sample_indices + self.num_threads = num_threads + self.worker_id = worker_id + + if not dist.is_initialized(): + dist.init_distributed() + + # comm_group and worker_id+num_workers are mutually exclusive + self.comm_group = comm_group + if self.comm_group is None: + # self.comm_group = deepspeed.utils.groups._clone_world_group() + self.num_workers = num_workers + self.worker_id = worker_id + else: + self.num_workers = self.comm_group.size() + self.worker_id = self.comm_group.rank() + + if self.worker_id == 0: + logger.info(f"Distributed data analyzer initialized with {self.num_workers} workers.") + + def run_map_helper(self, thread_id=0, metric_queues=None): + thread_start_idx, thread_end_idx = self.thread_splits[thread_id][0], self.thread_splits[thread_id][1] + worker_dataset = Subset(self.dataset, list(range(thread_start_idx, thread_end_idx))) + sampler = BatchSampler(SequentialSampler(worker_dataset), batch_size=self.batch_size, drop_last=False) + dataloader = DataLoader(dataset=worker_dataset, + batch_sampler=sampler, + num_workers=0, + collate_fn=self.collate_fn, + pin_memory=False) + + # set initial results list + metric_results = [] + for metric_type in self.metric_types: + assert metric_type in ['single_value_per_sample', 'accumulate_value_over_samples'], \ + f"metric_type {metric_type} not implemented." + metric_results.append([] if metric_type == 'single_value_per_sample' else None) + + # iterate dataloader and store metric results + batch_start_idx = thread_start_idx + for data in dataloader: + for m_idx in range(len(self.metric_names)): + metric_type, metric_function = self.metric_types[m_idx], self.metric_functions[m_idx] + metric_values = metric_function(data) + assert torch.is_tensor(metric_values) or isinstance(metric_values, np.ndarray), \ + "metric_function must return a tensor or array" + if isinstance(metric_values, np.ndarray): + metric_values = torch.from_numpy(metric_values) + assert metric_values.dtype in valid_dtypes, \ + f"metric_function result dtype {metric_values.dtype} not supported. Supported dtypes {valid_dtypes}" + + if metric_type == 'single_value_per_sample': + for row in range(metric_values.size()[0]): + value = metric_values[row].item() + sample_idx = batch_start_idx + row # sample idx following dataset iteration order + if isinstance(data, dict) and 'index' in data: # Megatron use case + sample_idx = data['index'][row][0].item() + elif self.sample_indices is not None: # user defined shuffling of indices + sample_idx = self.sample_indices[sample_idx] + metric_results[m_idx].append((value, sample_idx)) + elif metric_type == 'accumulate_value_over_samples': + if metric_results[m_idx] is None: + metric_results[m_idx] = metric_values + else: + metric_results[m_idx].add_(metric_values) + batch_start_idx += len(data) + + if self.num_threads == 1: + return metric_results + + # copy metric_results to the shared queue + assert metric_queues + for m_idx in range(len(self.metric_names)): + results = metric_results[m_idx] + if torch.is_tensor(results): + results = results.item() if results.dim() == 0 else results.tolist() + try: + metric_queues[m_idx].put((thread_id, results)) + except Exception as e: + logger.error(f"Error putting metric results to queue: {e}") + sys.exit(1) + + def run_map_reduce(self): + + # setup individual dataloaders + self.worker_splits, self.thread_splits = split_dataset(self.dataset, + self.num_workers, + self.worker_id, + num_threads=self.num_threads) + node_start_idx, node_end_idx = self.worker_splits[self.worker_id] + logger.info(f"worker {self.worker_id} working on data subset {node_start_idx} to {node_end_idx}.") + + if self.num_threads in [0, 1, None]: + metric_results = self.run_map_helper() + metric_results = [torch.tensor(m).to(self.device) for m in metric_results] + else: + + # create a shared queue of results per metric to be populated by individual threads + with Manager() as manager: + metric_queues = [manager.Queue() for _ in self.metric_names] + threads = [ + Process(target=self.run_map_helper, args=(t, metric_queues)) for t in range(self.num_threads) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + # gather results from shared queues into metric_results + metric_results = [None for _ in self.metric_names] + for m_idx, (queue, metric_type) in enumerate(zip(metric_queues, self.metric_types)): + while not queue.empty(): + t_idx, t_results = queue.get() + t_start_idx, t_end_idx = self.thread_splits[t_idx] + if t_start_idx >= t_end_idx: # no results from this thread + continue #corner case for small datasets and high thread count + t_results = torch.tensor(t_results) + if metric_type == 'single_value_per_sample': + # add thread results to the metric_results list, ordered by thread idx + if metric_results[m_idx] is None: # initialize if needed + metric_results[m_idx] = torch.zeros(node_end_idx - node_start_idx, + t_results.size(1)).to(self.device) + metric_results[m_idx][t_start_idx - node_start_idx:t_end_idx - node_start_idx] = t_results + else: + if metric_results[m_idx] is None: # initialize if needed + metric_results[m_idx] = torch.zeros(t_results.size()).to(self.device) + metric_results[m_idx].add_(t_results) + + # compute dtype for sample ids + total_num_samples = len(self.dataset) + sample_idx_dtype = find_fit_int_dtype(0, total_num_samples - 1) + logger.info(f"Total number of data samples: {total_num_samples}.") + logger.info(f"Will use {sample_idx_dtype} to store the sample indexes.") + + for m_idx in range(len(self.metric_names)): + metric_values, metric_name, metric_type = \ + metric_results[m_idx], self.metric_names[m_idx], self.metric_types[m_idx] + metric_save_path = f"{self.save_path}/{metric_name}/" + os.makedirs(metric_save_path, exist_ok=True) + + if metric_type == 'single_value_per_sample': + + # Compute sample and metric value dtypes based on range + values, samples = metric_values[:, 0], metric_values[:, 1] + value_min, value_max = Dist.min_max(values, self.comm_group) + sample_min, sample_max = Dist.min_max(samples, self.comm_group) + metric_value_dtype = find_fit_int_dtype(value_min, value_max) + sample_value_dtype = find_fit_int_dtype(sample_min, sample_max) + + # sample_to_metric maps sample ids to metric values, as a list of metric values + sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric" + values = [torch.tensor([x]) for x in metric_values[:, 0]] + self.file_write_ordered(values, sample_to_metric_fname, metric_value_dtype) + + # distributed sorting by values, gives an ordered disjoint subset of keys on nodes + metric_values = Dist.sample_sort(metric_values, self.comm_group, self.num_workers) + metric_to_samples_dict = {} + if len(metric_values) > 0: + for value, sample in metric_values: + if value.item() not in metric_to_samples_dict: + metric_to_samples_dict[value.item()] = [] + metric_to_samples_dict[value.item()].append(sample.item()) + + # index_to_metric and index_to_sample serialize a dicitonary from metric to samples + # index_to_metric stores a key per row, index_to_sample stores the values per row + values = [torch.tensor([x]) for x in metric_to_samples_dict.keys()] + samples = [torch.tensor(metric_to_samples_dict[x]) for x in metric_to_samples_dict.keys()] + index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric" #dict keys + index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample" #dict values + self.file_write_ordered(values, index_to_metric_fname, metric_value_dtype) + self.file_write_ordered(samples, index_to_sample_fname, sample_value_dtype) + + if self.worker_id == 0: + DataAnalyzer.output_index_to_sample_percentile(index_to_sample_fname, index_to_metric_fname, + metric_name, metric_save_path, total_num_samples, + sample_idx_dtype) + dist.barrier(self.comm_group) + + elif metric_type == 'accumulate_value_over_samples': + metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value" + dist.reduce(metric_values, dst=0, op=dist.ReduceOp.SUM, group=self.comm_group) + metric_value_dtype = find_fit_int_dtype(metric_values.min(), metric_values.max()) + + if self.worker_id == 0: + builder = create_mmap_dataset_builder(metric_value_fname, metric_value_dtype) + builder.add_item(metric_values.cpu()) + close_mmap_dataset_builder(builder, metric_value_fname) + dist.barrier(self.comm_group) + + def file_write_ordered(self, tensor_list, fname, numpy_dtype): + """ MPI_file_write_ordered extended to write a list of tensors, by one rank, iteratively """ + + # each node has a list of rows (tensors) to be written to the file. + # we will serialize it in order to communicate it in one comm step. + + tkwargs = dict(dtype=torch.int64, device=self.device) + + # 1. gather on rank 0 the number of rows to be sent/recv + row_count = torch.tensor([len(tensor_list)], **tkwargs) + row_counts = torch.zeros(self.num_workers, **tkwargs) + dist.all_gather_into_tensor(row_counts, row_count, group=self.comm_group) + assert row_counts[self.worker_id] == row_count == len(tensor_list), "all_gather failed" + + # 2. gather on rank 0 the sizes of the rows to be sent/recv + row_len = torch.tensor([len(l) for l in tensor_list], **tkwargs) + row_lens = Dist.gather_v(row_len, 0, self.comm_group, self.num_workers, self.worker_id) + + # 4. gather on rank 0 of the total size (sum of all row lengths) to be received + size = torch.tensor([sum(row_len).item()], **tkwargs) + sizes = torch.zeros(self.num_workers, **tkwargs) + dist.all_gather_into_tensor(sizes, size, group=self.comm_group) + assert sizes[self.worker_id] == size.item(), "all_gather did not return the same sizes" #sanity check + + # method to deserializes a buffer into rows of different lengths and write them to file + def write_buffer_to_file(buff, src, builder): + assert self.worker_id == 0, "only rank 0 can write to file" + + # collect all buffers and write them at once + buff = buff.cpu().detach().numpy() + row_offsets = np.cumsum([0] + row_lens[src].tolist()) + arr_list = [] + for i in range(len(row_lens[src])): + arr_list.append(buff[row_offsets[i]:row_offsets[i + 1]]) + builder.add_items(arr_list) + + # 5. rank 0 prepares output folder and file + if self.worker_id == 0: + os.makedirs(os.path.dirname(fname), exist_ok=True) + builder = create_mmap_dataset_builder(fname, numpy_dtype) + + # iterate through ranks that have data to be sent/recv/written + for src in [rank for rank, count in enumerate(row_counts) if count > 0]: + + dist.barrier(group=self.comm_group) + if self.worker_id == 0 and src == 0: # rank 0's write its own data + buffer = torch.cat(tensor_list, dim=0).to(self.device) + write_buffer_to_file(buffer, 0, builder) + elif self.worker_id == 0 and src > 0: # rank 0 receives other rank's data and writes it + buffer = torch.empty(sizes[src].item(), dtype=buffer.dtype, device=buffer.device) + err = dist.recv(buffer, src=src, group=self.comm_group, tag=src) + assert err == src and len(buffer) > 0, "recv failed" + write_buffer_to_file(buffer, src, builder) + elif self.worker_id == src: # current rank sends data to rank 0 + buffer = torch.cat(tensor_list, dim=0).to(self.device) + dist.send(buffer, 0, group=self.comm_group, tag=src) + + # rank 0 closes the file + if self.worker_id == 0: + close_mmap_dataset_builder(builder, fname) # close file + dist.barrier(self.comm_group) + + +class Dist: + """ auxiliary class to perform distributed operations on tensors""" + + @staticmethod + def min_max(tensor, comm_group): + """ given a distributed tensor, return the min/max values across all ranks""" + + value_min, value_max = tensor.min(), tensor.max() + dist.reduce(value_min, 0, op=dist.ReduceOp.MIN, group=comm_group) + dist.reduce(value_max, 0, op=dist.ReduceOp.MAX, group=comm_group) + return value_min.item(), value_max.item() + + @staticmethod + def gather_v(tensor, dst, comm_group, num_workers, worker_id): + """ MPI_Gatherv. gather tensors of variable sizes in a single rank """ + + # gather the number of rows to be sent/recv + size = torch.tensor([len(tensor)], dtype=torch.int64, device=tensor.device) + sizes = torch.zeros(num_workers, dtype=torch.int64, device=tensor.device) + dist.all_gather_into_tensor(sizes, size, group=comm_group) + assert sizes[worker_id] == size, "all_gather failed" + + # all_gather requires all tensors to be of same size so we need to pad them + max_size = max(sizes).item() + buffer = torch.empty(max_size, dtype=tensor.dtype, device=tensor.device) + buffer[0:size] = tensor.data + buffer_list = None + if worker_id == 0: # create padded recv buffers + buffer_list = [torch.empty(max_size, dtype=tensor.dtype, device=tensor.device) for _ in range(num_workers)] + dist.gather(buffer, buffer_list, dst=dst, group=comm_group) + + # revert padding and return value + if worker_id == 0: + buffer_list = [r[:s.item()] for r, s in zip(buffer_list, sizes)] + return buffer_list + + @staticmethod + def sample_sort(tensor, comm_group, num_workers, n_samples=100): + """ perform a distributed random sort of a tensor, and returns the sorted partial tensor""" + device, dims = tensor.device, tensor.size()[1] + + # 1 - sort rows by first column, then second column, then third, etc... + tensor = torch.tensor(sorted(tensor.tolist()), dtype=tensor.dtype, device=tensor.device) + + # 2 - collect few samples per rank + idx = torch.round(torch.linspace(0, len(tensor) - 1, n_samples)).to(int) + samples = tensor[idx][:, 0].contiguous().to(device) #only first column, all but last row + + # 2 - Allgather samples + all_samples = [torch.zeros(n_samples, dtype=samples.dtype, device=device) for _ in range(num_workers)] + dist.all_gather(all_samples, samples, group=comm_group) + all_samples = torch.cat(all_samples, dim=0).to(device) + + # 3 - Sort all samples and collect the ranges of each rank as equidistant + all_samples = all_samples.sort()[0] + idx = torch.round(torch.linspace(0, len(all_samples) - 1, num_workers + 1)).to(int) + ranges = all_samples[idx] # range of each rank r as ranges[r] <= x < ranges[r+1] + ranges[-1] += 1 # increase upper limit of last rank so that x < ranges[r+1]. + + # 4 - collect elements to send to each rank, based on the rank ranges + send = [] + for rank in range(num_workers): + mask = (tensor[:, 0] >= ranges[rank]) & (tensor[:, 0] < ranges[rank + 1]) + send.append(tensor[mask]) + + # 5. all to all to communicate the sizes to be sent/recv + send_count = [torch.tensor([len(s) * dims], dtype=torch.int64, device=device) for s in send] + recv_count = list(torch.empty([num_workers], dtype=torch.int64, device=device).chunk(num_workers)) + dist.all_to_all(recv_count, send_count, group=comm_group) + + # 6. all-to-all-v to communicate the elements to be sent/recv as a single tensor + send = torch.cat(send, dim=0).flatten().to(device) + recv = torch.zeros(sum(recv_count), dtype=send.dtype).to(device) + send_count = [s.item() for s in send_count] # convert to list of ints + recv_count = [r.item() for r in recv_count] + dist.all_to_all_single(recv, send, recv_count, send_count, group=comm_group) + del send + + # 7. the received tensor is the 1D disjoint subset of the distributed tensor. + # We will recover the original dimensionality and sort it by columns again. + recv = recv.view(-1, dims) + recv = torch.tensor(sorted(recv.tolist()), dtype=recv.dtype, device=recv.device) + return recv + + +def test_compare_both_data_analyzers(dataset): + """ given a dataset, compare file and memory based data analyser""" + + id = lambda t: t.to(torch.int64) # identity + batch_sum = lambda t: id(t).sum() #sum batch + num_threads = 4 + kwargs = dict( + dataset=dataset, + batch_size=2**10, + worker_id=int(os.environ['RANK']), + num_workers=int(os.environ['WORLD_SIZE']), + metric_names=["mod", "batch_sum"], + metric_functions=[id, batch_sum], + metric_types=['single_value_per_sample', 'accumulate_value_over_samples'], + num_threads=num_threads, + ) + + dda = DistributedDataAnalyzer( + save_path="./output_dist", + device=f"cuda:{int(os.environ['LOCAL_RANK'])}", + **kwargs, + ) + start_time = time.time() + dda.run_map_reduce() + if dda.worker_id == 0: + print("DistributedDataAnalyzer runtime: %s seconds " % (time.time() - start_time)) + + da = DataAnalyzer(num_threads_reduce=num_threads, + save_path="./output_disk", + metric_dtypes=[torch.int64, torch.int64], + **kwargs) + start_time = time.time() + da.run_map_reduce() + if da.worker_id == 0: + print("DataAnalyzer runtime: %s seconds " % (time.time() - start_time)) + + output_paths = [ + "batch_sum/batch_sum_metric_value.bin", "batch_sum/batch_sum_metric_value.idx", \ + "mod/mod_index_to_metric.bin", "mod/mod_index_to_metric.idx", \ + "mod/mod_index_to_sample.bin", "mod/mod_index_to_sample.idx", \ + "mod/mod_index_to_sample_percentile_merged.bin", "mod/mod_index_to_sample_percentile_merged.idx", \ + "mod/mod_sample_to_metric.bin", "mod/mod_sample_to_metric.idx" + ] + + if dda.worker_id == 0: + for path in output_paths: + with open(os.path.join(da.save_path, path), 'rb') as f1, \ + open(os.path.join(dda.save_path, path), 'rb') as f2: + # if files have suffix .bin, they should be identical + if path.endswith(".bin"): + assert f1.read() == f2.read(), f"files {path} are not identical." + elif f1.read() != f2.read(): + print(f"files {path} are not identical.") + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + + class TestDataset(torch.utils.data.Dataset): + + def __init__(self, size=10_000_000): + self.values = [(x + 7) % 10_000 for x in range(size)] + self.size = size + + __len__ = lambda self: self.size + __getitem__ = lambda self, idx: self.values[idx] + + test_compare_both_data_analyzers(TestDataset()) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..100bef3f7946c8d2b7280e234abc5e625b6c68f9 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/data_sampler.py @@ -0,0 +1,349 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +coding=utf-8 + Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +Part of this code was adopted from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/data_samplers.py +""" + +import torch +import os +import numpy as np + +import deepspeed.comm as dist +from deepspeed.utils import logger +from deepspeed.accelerator import get_accelerator +from ..constants import * +from ..curriculum_scheduler import CurriculumScheduler +from .indexed_dataset import MMapIndexedDataset +from .utils import create_mmap_dataset_builder, close_mmap_dataset_builder, find_fit_int_dtype + + +class DeepSpeedDataSampler(object): + + def __init__(self, + data_efficiency_config, + one_epoch_total_samples, + micro_batch_size, + data_parallel_rank, + data_parallel_size, + data_parallel_group, + gradient_accumulation_steps, + global_rank, + drop_last=True): + # Keep a copy of input params for later use. + self.data_efficiency_config = data_efficiency_config + self.one_epoch_total_samples = one_epoch_total_samples + self.index_dtype = find_fit_int_dtype(0, one_epoch_total_samples) + self.total_samples = one_epoch_total_samples * self.data_efficiency_config[DATA_SAMPLING][ + DATA_SAMPLING_NUM_EPOCHS] + self.micro_batch_size = micro_batch_size + self.data_parallel_rank = data_parallel_rank + self.data_parallel_group = data_parallel_group + self.micro_batch_times_data_parallel_size = \ + self.micro_batch_size * data_parallel_size + self.gradient_accumulation_steps = gradient_accumulation_steps + self.global_batch_size = self.micro_batch_times_data_parallel_size * \ + self.gradient_accumulation_steps + self.global_rank = global_rank + self.drop_last = drop_last + self.np_rng = np.random.default_rng(self.data_efficiency_config[DATA_EFFICIENCY_SEED]) + self.state = {} + self.batch = [] + self.consumed_samples = 0 + if self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]: + self.curriculum_step = 0 + self.current_difficulties = {} + self.data_cluster_paths = [] + self.data_cluster_current_position = [] + self.curriculum_schedulers = {} + self.curriculum_index_to_sample = {} + self.curriculum_index_to_metric = {} + self.difficulty_type = {} + self.clustering_type = {} + self.data_1epoch_size = None + if self.global_rank == 0: + self.data_clusters = [] + self.data_cluster_sizes = [] + cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_CLUSTER_PATH] + if not os.path.exists(cluster_path): + os.makedirs(cluster_path) + for metric in self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS]: + self.curriculum_schedulers[metric] = CurriculumScheduler( + data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS][metric]) + self.difficulty_type[metric] = data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_METRICS][metric][CURRICULUM_LEARNING_DIFFICULTY_TYPE] + self.clustering_type[metric] = data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_METRICS][metric][CURRICULUM_LEARNING_CLUSTERING_TYPE] + if self.global_rank == 0: + if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: + self.curriculum_index_to_sample[metric] = MMapIndexedDataset( + data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS] + [metric][CURRICULUM_LEARNING_SAMPLE_PATH], + skip_warmup=True) + if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: + self.curriculum_index_to_metric[metric] = MMapIndexedDataset( + data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS] + [metric][CURRICULUM_LEARNING_METRIC_PATH], + skip_warmup=True) + + # Sanity checks. + assert self.total_samples > 0, \ + 'no sample to consume: {}'.format(self.total_samples) + assert self.micro_batch_size > 0 + assert data_parallel_size > 0 + assert self.data_parallel_rank < data_parallel_size, \ + 'data_parallel_rank should be smaller than data size: {}, ' \ + '{}'.format(self.data_parallel_rank, data_parallel_size) + + def __len__(self): + return self.total_samples + + def set_custom_curriculum_learning_schedule(self, schedule_func_dict): + for metric in self.curriculum_schedulers: + if metric in schedule_func_dict: + self.curriculum_schedulers[metric].set_custom_get_difficulty(schedule_func_dict[metric]) + + def get_start_end_idx(self, batch_len=None): + """ + given the length of a minibatch (defaults to micro-batch size * data_parallel_size), + return the start and end indices of the current data parallel rank + """ + batch_len = batch_len or self.micro_batch_times_data_parallel_size + start_idx_fn = lambda r: round(r * batch_len / self.data_parallel_group.size()) + start_idx = start_idx_fn(self.data_parallel_rank) + end_idx = start_idx_fn(self.data_parallel_rank + 1) + return start_idx, end_idx + + def get_sample_based_on_metric_value(self, metric, value_start, value_end): + new_samples = None + for row in range(len(self.curriculum_index_to_sample[metric])): + if self.curriculum_index_to_metric[metric][row] <= value_end and self.curriculum_index_to_metric[metric][ + row] > value_start: + row_samples = np.copy(self.curriculum_index_to_sample[metric][row]) + new_samples = row_samples if new_samples is None else np.concatenate( + (new_samples, row_samples), axis=None) + return new_samples + + def get_sample_based_on_metric_percentile(self, metric, percentile_start, percentile_end): + new_samples = None + if self.data_1epoch_size is None: + self.data_1epoch_size = sum(len(x) for x in self.curriculum_index_to_sample[metric]) + max_percentile = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS][ + metric][CURRICULUM_LEARNING_MAX_DIFFICULTY] + sample_per_percentile = self.data_1epoch_size // max_percentile + start_count = sample_per_percentile * percentile_start + end_count = sample_per_percentile * percentile_end + if percentile_end == max_percentile: + end_count = self.data_1epoch_size + current_count = 0 + for row in range(len(self.curriculum_index_to_sample[metric])): + row_size = len(self.curriculum_index_to_sample[metric][row]) + if current_count + row_size > start_count: + row_start = max(0, start_count - current_count) + if current_count + row_size <= end_count: + row_end = row_size + else: + row_end = end_count - current_count + row_samples = np.copy(self.curriculum_index_to_sample[metric][row][row_start:row_end]) + new_samples = row_samples if new_samples is None else np.concatenate( + (new_samples, row_samples), axis=None) + current_count += row_size + if current_count >= end_count: + break + return new_samples + + def get_new_cluster(self, previous_difficulties): + cluster_fname = CURRICULUM_LEARNING_CLUSTER_PREFIX + for metric in self.curriculum_schedulers: + cluster_fname = f"{cluster_fname}_{metric}{self.current_difficulties[metric]}" + cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_CLUSTER_PATH] + cluster_path = f"{cluster_path}/{cluster_fname}" + if self.global_rank == 0: + new_cluster = None + need_clustering = 0 + for metric in self.clustering_type: + if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: + need_clustering += 1 + if need_clustering > 1: + for metric in self.curriculum_schedulers: + if self.clustering_type[metric] == CURRICULUM_LEARNING_SINGLE_CLUSTER: + metric_cluster = np.arange(start=0, + stop=self.one_epoch_total_samples, + step=1, + dtype=self.index_dtype) + else: + if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: + metric_cluster = self.get_sample_based_on_metric_value(metric, float('-inf'), + self.current_difficulties[metric]) + elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: + metric_cluster = self.get_sample_based_on_metric_percentile( + metric, 0, self.current_difficulties[metric]) + new_cluster = metric_cluster if new_cluster is None else \ + np.intersect1d(new_cluster, metric_cluster, assume_unique=True) + for cluster in self.data_clusters: + new_cluster = np.setdiff1d(new_cluster, cluster[0], assume_unique=True) + else: + if len(self.data_clusters) == 0: + new_cluster = np.arange(start=0, stop=self.one_epoch_total_samples, step=1, dtype=self.index_dtype) + for metric in self.curriculum_schedulers: + if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER: + if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: + new_cluster = self.get_sample_based_on_metric_value(metric, previous_difficulties[metric], + self.current_difficulties[metric]) + elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: + new_cluster = self.get_sample_based_on_metric_percentile( + metric, previous_difficulties[metric], self.current_difficulties[metric]) + if new_cluster is not None and len(new_cluster) > 0: + logger.info( + f"new data cluster (previous_difficulties {previous_difficulties}, current_difficulties {self.current_difficulties}) with size {len(new_cluster)} generated." + ) + self.np_rng.shuffle(new_cluster) + cluster_builder = create_mmap_dataset_builder(cluster_path, self.index_dtype) + cluster_builder.add_item_numpy(new_cluster) + close_mmap_dataset_builder(cluster_builder, cluster_path) + self.data_clusters.append(MMapIndexedDataset(cluster_path, skip_warmup=True)) + self.data_cluster_sizes.append(len(self.data_clusters[-1][0])) + else: + logger.info( + f"new data cluster (previous_difficulties {previous_difficulties}, current_difficulties {self.current_difficulties}) has no matched data thus skipped." + ) + dist.barrier(group=self.data_parallel_group) + if os.path.isfile(f"{cluster_path}.bin"): + self.data_cluster_paths.append(cluster_fname) + self.data_cluster_current_position.append(0) + + def sample_from_clusters(self): + num_clusters = len(self.data_clusters) + weight_sum = sum(self.data_cluster_sizes) + weights = [x / weight_sum for x in self.data_cluster_sizes] + samples = self.np_rng.choice(num_clusters, self.global_batch_size, replace=True, p=weights) + samples = np.bincount(samples, minlength=num_clusters) + return samples + + def reshuffle_clusters(self, cidx): + cluster_fname = self.data_cluster_paths[cidx] + cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_CLUSTER_PATH] + cluster_path = f"{cluster_path}/{cluster_fname}" + cluster = np.copy(self.data_clusters[cidx][0]) + self.np_rng.shuffle(cluster) + cluster_builder = create_mmap_dataset_builder(cluster_path, self.index_dtype) + cluster_builder.add_item_numpy(cluster) + close_mmap_dataset_builder(cluster_builder, cluster_path) + self.data_clusters[cidx] = MMapIndexedDataset(cluster_path, skip_warmup=True) + + def get_sample_from_cluster(self, cidx, num_samples): + start_idx = self.data_cluster_current_position[cidx] + samples = list(np.copy(self.data_clusters[cidx][0][start_idx:(start_idx + num_samples)])) + self.data_cluster_current_position[cidx] += num_samples + if len(samples) < num_samples: + num_samples_remained = num_samples - len(samples) + logger.info(f"reshuffling cluster {cidx}.") + self.reshuffle_clusters(cidx) + samples += list(np.copy(self.data_clusters[cidx][0][:num_samples_remained])) + self.data_cluster_current_position[cidx] = num_samples_remained + return samples + + def get_next_global_batch(self): + if self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]: + self.curriculum_step += 1 + new_cluster = False + previous_difficulties = {} + for metric in self.curriculum_schedulers: + next_difficulty = self.curriculum_schedulers[metric].update_difficulty(self.curriculum_step) + if metric not in self.current_difficulties or \ + next_difficulty != self.current_difficulties[metric]: + new_cluster = True + if metric in self.current_difficulties: + previous_difficulties[metric] = self.current_difficulties[metric] + else: + if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED: + previous_difficulties[metric] = float('-inf') + elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED: + previous_difficulties[metric] = 0 + self.current_difficulties[metric] = next_difficulty + if new_cluster: + self.get_new_cluster(previous_difficulties) + if self.global_rank == 0: + samples_per_cluster = self.sample_from_clusters() + batch = [] + for cidx in range(len(samples_per_cluster)): + batch += self.get_sample_from_cluster(cidx, samples_per_cluster[cidx]) + self.np_rng.shuffle(batch) + + # broadcast tensor must have same shape across participants. So we fill batch with -1s when not full + assert len(batch) <= self.global_batch_size + batch += [-1] * (self.global_batch_size - len(batch)) + batch = torch.tensor(batch, device=get_accelerator().current_device_name(), dtype=torch.long).view(-1) + else: + batch = torch.empty(self.global_batch_size, + device=get_accelerator().current_device_name(), + dtype=torch.long) + dist.broadcast(batch, 0, group=self.data_parallel_group) + batch = batch[batch != -1] # remove trailing -1s used to fill incomplete batch tensor + self.batch = batch.tolist() + + def __iter__(self): + while self.consumed_samples <= self.total_samples: + if len(self.batch) == 0: + self.get_next_global_batch() + current_batch = self.batch[:self.micro_batch_times_data_parallel_size] + self.batch = self.batch[self.micro_batch_times_data_parallel_size:] + if len(current_batch) == self.micro_batch_times_data_parallel_size or \ + (len(current_batch) > 0 and not self.drop_last): + start_idx, end_idx = self.get_start_end_idx(len(current_batch)) + yield current_batch[start_idx:end_idx] + self.consumed_samples += len(current_batch) + current_batch = [] + + def state_dict(self): + return { + CURRICULUM_LEARNING_BATCH: self.batch, + CURRICULUM_LEARNING_CONSUMED_SAMPLES: self.consumed_samples, + CURRICULUM_LEARNING_STEP: self.curriculum_step, + CURRICULUM_LEARNING_CURRENT_DIFFICULTIES: self.current_difficulties, + CURRICULUM_LEARNING_DATA_CLUSTER_PATHS: self.data_cluster_paths, + CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION: self.data_cluster_current_position, + CURRICULUM_LEARNING_NP_RNG_STATE: np.random.get_state() + } + + def load_state_dict(self, state_dict): + self.batch = state_dict[CURRICULUM_LEARNING_BATCH] + self.consumed_samples = state_dict[CURRICULUM_LEARNING_CONSUMED_SAMPLES] + self.curriculum_step = state_dict[CURRICULUM_LEARNING_STEP] + self.current_difficulties = state_dict[CURRICULUM_LEARNING_CURRENT_DIFFICULTIES] + self.data_cluster_paths = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_PATHS] + self.data_cluster_current_position = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION] + np.random.set_state(state_dict[CURRICULUM_LEARNING_NP_RNG_STATE]) + cluster_root_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][ + CURRICULUM_LEARNING_CLUSTER_PATH] + # Backward compatibility: previously data_cluster_paths were stored as + # absolute paths. Now we changed it to just the file name so that even + # if user moved the cluster files, the checkpoint loading still works + # as long as user set the correct new CURRICULUM_LEARNING_CLUSTER_PATH + # in deepspeed json config. + for idx in range(len(self.data_cluster_paths)): + if '/' in self.data_cluster_paths[idx]: + self.data_cluster_paths[idx] = self.data_cluster_paths[idx].split('/')[-1] + if self.global_rank == 0: + for cluster_fname in self.data_cluster_paths: + cluster_path = f"{cluster_root_path}/{cluster_fname}" + self.data_clusters.append(MMapIndexedDataset(cluster_path, skip_warmup=True)) + self.data_cluster_sizes.append(len(self.data_clusters[-1][0])) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..872d05de01453aea76fc850de77da3acbbcea76c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/indexed_dataset.py @@ -0,0 +1,627 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Part of this code was adopted from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/indexed_dataset.py +""" + +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# copied from fairseq/fairseq/data/indexed_dataset.py +# Removed IndexedRawTextDataset since it relied on Fairseq dictionary +# other slight modifications to remove fairseq dependencies +# Added document index to index file and made it accessible. +# An empty sentence no longer separates documents. + +# Some of the fixes/improvements are adopted from +# https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/megatron/data/indexed_dataset.py + +from functools import lru_cache +import os +import shutil +import struct +from itertools import accumulate + +import numpy as np +import torch + + +def __best_fitting_dtype(vocab_size=None): + if vocab_size is not None and vocab_size < 65500: + return np.uint16 + else: + return np.int32 + + +def get_available_dataset_impl(): + return ['lazy', 'cached', 'mmap'] + + +def infer_dataset_impl(path): + if IndexedDataset.exists(path): + with open(index_file_path(path), 'rb') as f: + magic = f.read(8) + if magic == IndexedDataset._HDR_MAGIC: + return 'cached' + elif magic == MMapIndexedDataset.Index._HDR_MAGIC[:8]: + return 'mmap' + else: + return None + else: + print(f"Dataset does not exist: {path}") + print("Path should be a basename that both .idx and .bin can be appended to get full filenames.") + return None + + +def make_builder(out_file, impl, vocab_size=None): + if impl == 'mmap': + return MMapIndexedDatasetBuilder(out_file, dtype=__best_fitting_dtype(vocab_size)) + else: + return IndexedDatasetBuilder(out_file) + + +def make_dataset(path, impl, skip_warmup=False): + if not IndexedDataset.exists(path): + print(f"Dataset does not exist: {path}") + print("Path should be a basename that both .idx and .bin can be appended to get full filenames.") + return None + if impl == 'infer': + impl = infer_dataset_impl(path) + if impl == 'lazy' and IndexedDataset.exists(path): + return IndexedDataset(path) + elif impl == 'cached' and IndexedDataset.exists(path): + return IndexedCachedDataset(path) + elif impl == 'mmap' and MMapIndexedDataset.exists(path): + return MMapIndexedDataset(path, skip_warmup) + print(f"Unknown dataset implementation: {impl}") + return None + + +def dataset_exists(path, impl): + if impl == 'mmap': + return MMapIndexedDataset.exists(path) + else: + return IndexedDataset.exists(path) + + +def read_longs(f, n): + a = np.empty(n, dtype=np.int64) + f.readinto(a) + return a + + +def write_longs(f, a): + f.write(np.array(a, dtype=np.int64)) + + +# valid metric_dtypes as numpy and torch types +dtypes = { + 1: (np.uint8, torch.uint8), + 2: (np.int8, torch.int8), + 3: (np.int16, torch.int16), + 4: (np.int32, torch.int32), + 5: (np.int64, torch.int64), + 6: (np.uint16, None), + 7: (np.uint32, None), + 8: (np.uint64, None), +} + +valid_dtypes = set([dt[0] for dt in dtypes.values()] + [dt[1] for dt in dtypes.values() if dt[1] is not None]) + + +def code(dtype): + for c, (np_dt, torch_dt) in dtypes.items(): + if dtype in [np_dt, torch_dt]: + return c + raise ValueError(f"{dtype} not supported. Supported types: {valid_dtypes}") + + +def index_file_path(prefix_path): + return prefix_path + '.idx' + + +def data_file_path(prefix_path): + return prefix_path + '.bin' + + +def create_doc_idx(sizes): + doc_idx = [0] + for i, s in enumerate(sizes): + if s == 0: + doc_idx.append(i + 1) + return doc_idx + + +class IndexedDataset(torch.utils.data.Dataset): + """Loader for IndexedDataset""" + _HDR_MAGIC = b'TNTIDX\x00\x00' + + def __init__(self, path): + super().__init__() + self.path = path + self.data_file = None + self.read_index(path) + + def read_index(self, path): + with open(index_file_path(path), 'rb') as f: + magic = f.read(8) + assert magic == self._HDR_MAGIC, ('Index file doesn\'t match expected format. ' + 'Make sure that --dataset-impl is configured properly.') + version = f.read(8) + assert struct.unpack('= self._len: + raise IndexError('index out of range') + + def __del__(self): + if self.data_file: + self.data_file.close() + + # @lru_cache(maxsize=8) + def __getitem__(self, idx): + if not self.data_file: + self.read_data(self.path) + if isinstance(idx, int): + i = idx + self.check_index(i) + tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]] + a = np.empty(tensor_size, dtype=self.dtype) + self.data_file.seek(self.data_offsets[i] * self.element_size) + self.data_file.readinto(a) + return a + elif isinstance(idx, slice): + start, stop, step = idx.indices(len(self)) + if step != 1: + raise ValueError("Slices into indexed_dataset must be contiguous") + sizes = self.sizes[self.dim_offsets[start]:self.dim_offsets[stop]] + size = sum(sizes) + a = np.empty(size, dtype=self.dtype) + self.data_file.seek(self.data_offsets[start] * self.element_size) + self.data_file.readinto(a) + offsets = list(accumulate(sizes)) + sents = np.split(a, offsets[:-1]) + return sents + + def __len__(self): + return self._len + + def num_tokens(self, index): + return self.sizes[index] + + def size(self, index): + return self.sizes[index] + + @staticmethod + def exists(path): + return (os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path))) + + @property + def supports_prefetch(self): + return False # avoid prefetching to save memory + + +class IndexedCachedDataset(IndexedDataset): + + def __init__(self, path): + super().__init__(path) + self.cache = None + self.cache_index = {} + + @property + def supports_prefetch(self): + return True + + def prefetch(self, indices): + if all(i in self.cache_index for i in indices): + return + if not self.data_file: + self.read_data(self.path) + indices = sorted(set(indices)) + total_size = 0 + for i in indices: + total_size += self.data_offsets[i + 1] - self.data_offsets[i] + self.cache = np.empty(total_size, dtype=self.dtype) + ptx = 0 + self.cache_index.clear() + for i in indices: + self.cache_index[i] = ptx + size = self.data_offsets[i + 1] - self.data_offsets[i] + a = self.cache[ptx:ptx + size] + self.data_file.seek(self.data_offsets[i] * self.element_size) + self.data_file.readinto(a) + ptx += size + if self.data_file: + # close and delete data file after prefetch so we can pickle + self.data_file.close() + self.data_file = None + + # @lru_cache(maxsize=8) + def __getitem__(self, idx): + if isinstance(idx, int): + i = idx + self.check_index(i) + tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]] + a = np.empty(tensor_size, dtype=self.dtype) + ptx = self.cache_index[i] + np.copyto(a, self.cache[ptx:ptx + a.size]) + return a + elif isinstance(idx, slice): + # Hack just to make this work, can optimizer later if necessary + sents = [] + for i in range(*idx.indices(len(self))): + sents.append(self[i]) + return sents + + +class IndexedDatasetBuilder(object): + + def __init__(self, out_file, dtype=np.int32): + self.out_file = open(out_file, 'wb') + self.dtype = dtype + self.data_offsets = [0] + self.dim_offsets = [0] + self.sizes = [] + self.element_size = self.dtype().itemsize + self.doc_idx = [0] + + def add_item(self, tensor): + bytes = self.out_file.write(np.array(tensor.numpy(), dtype=self.dtype)) + self.data_offsets.append(self.data_offsets[-1] + bytes / self.element_size) + for s in tensor.size(): + self.sizes.append(s) + self.dim_offsets.append(self.dim_offsets[-1] + len(tensor.size())) + + def end_document(self): + self.doc_idx.append(len(self.sizes)) + + def merge_file_(self, another_file): + index = IndexedDataset(another_file) + assert index.dtype == self.dtype + + doc_offset = len(self.sizes) + + begin = self.data_offsets[-1] + for data_offset in index.data_offsets[1:]: + self.data_offsets.append(begin + data_offset) + self.sizes.extend(index.sizes) + begin = self.dim_offsets[-1] + for dim_offset in index.dim_offsets[1:]: + self.dim_offsets.append(begin + dim_offset) + self.doc_idx.extend((doc_offset + index.doc_idx)[1:]) + + with open(data_file_path(another_file), 'rb') as f: + while True: + data = f.read(1024) + if data: + self.out_file.write(data) + else: + break + + def finalize(self, index_file): + self.out_file.close() + index = open(index_file, 'wb') + index.write(b'TNTIDX\x00\x00') + index.write(struct.pack(' [0, 10, 30, 35] + if arr.size > 1: + arr[1:] = arr[:-1] + if arr.size > 0: + arr[0] = 0 + + +def get_pointers_with_total(sizes, elemsize, dtype): + """Return a numpy array of type np.dtype giving the byte offsets. + + Multiplies values in the sizes array by elemsize (bytes), + and then computes an exclusive scan to get byte offsets. + Returns the total number of bytes as second item in a tuple. + """ + + # scale values in sizes array by elemsize to get sizes in bytes + pointers = np.array(sizes, dtype=dtype) + pointers *= elemsize + np.cumsum(pointers, axis=0, out=pointers) + + # get total number of bytes from all sizes (last element) + bytes_last = pointers[-1] if len(sizes) > 0 else 0 + + # convert to byte offsets + exscan_from_cumsum_(pointers) + + return pointers, bytes_last + + +class MMapIndexedDataset(torch.utils.data.Dataset): + + class Index(object): + _HDR_MAGIC = b'MMIDIDX\x00\x00' + + @classmethod + def writer(cls, path, dtype): + + class _Writer(object): + + def __enter__(self): + self._file = open(path, 'wb') + + self._file.write(cls._HDR_MAGIC) + self._file.write(struct.pack('= 0: + if max_value <= 255: + return np.uint8 + elif max_value <= 65535: + return np.uint16 + elif max_value <= 4294967295: + return np.uint32 + else: + return np.uint64 + else: + if max_value <= 127 and min_value >= -128: + return np.int8 + elif max_value <= 32767 and min_value >= -32768: + return np.int16 + elif max_value <= 2147483647 and min_value >= -2147483648: + return np.int32 + else: + return np.int64 + + +def split_index(start_idx, end_idx, num_partitions): + partition_boundaries = np.linspace(start_idx, end_idx, dtype=int, num=num_partitions + 1) + return [(partition_boundaries[i], partition_boundaries[i + 1]) for i in range(num_partitions)] + + +def split_dataset(dataset, num_workers, worker_id, num_threads): + worker_splits = split_index(0, len(dataset), num_workers) + thread_splits = split_index(worker_splits[worker_id][0], worker_splits[worker_id][1], num_threads) + return worker_splits, thread_splits + + +def create_mmap_dataset_builder(fname, dtype): + logger.info(f"Creating mmap dataset builder at {fname}.") + return MMapIndexedDatasetBuilder(f"{fname}.bin", dtype=dtype) + + +def close_mmap_dataset_builder(builder, fname): + builder.end_document() + builder.finalize(f"{fname}.idx") + logger.info(f"Finalized mmap dataset builder at {fname}.") diff --git a/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py new file mode 100644 index 0000000000000000000000000000000000000000..c9a39bbc53b509479aa2c18f3ddfc2e0abdb7a3c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr.py @@ -0,0 +1,492 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# support/questions/maintenance: github user @brunomaga or @deepspeedai/deepspeed + +import random +import torch +import os +import numpy as np +from torch.optim.lr_scheduler import LRScheduler +from torch.optim.optimizer import Optimizer +from torch.utils.data import DataLoader, DistributedSampler +from deepspeed.utils import logger +from deepspeed.runtime.pipe.engine import PipelineEngine +from deepspeed.runtime.data_pipeline.constants import * +from deepspeed.runtime.data_pipeline.data_sampling.indexed_dataset import MMapIndexedDataset +from deepspeed.runtime.data_pipeline.data_sampling.data_analyzer import DistributedDataAnalyzer +import pathlib + + +def batch_by_seqlens( + seqlens, + max_tokens, + sequence_ids_per_mb=None, + min_batch_size=1, + max_batch_size=None, + sequence_picking_order="dataloader", + effective_batch_size=1, + required_microbatches_of_same_size=False, + verbose=False, + seed=None, +): + """ + Yield mini-batches of indices bucketed by size. Batches may contain sequences of different lengths. + Similar to "Attention is all you need", Section 5.1: + "sequence pairs were batched together by approximate sequence length. Each training batch + contained a set of sequence pairs containing approximately X source tokens and X target tokens" + + Arguments: + - `seqlens`: a list of difficulties (metric values) for every sample in the dataset; + - `max_tokens`: maximum cap in total difficulty in a batch; + - `min_batch_size`: smallest allowed size of a batch; + - `min_batch_size`: largest allowed size of a batch; + - `sequence_picking_order`: order in which to process samples: "dataloader" (default), "random" or "seqlen" (ascending) + - `effective_batch_size`: effective batch size; + - `required_microbatches_of_same_size`: enable if each mini-batch (in a total of `batch_size_multiple` + micro-batches per batch), should have all micro-batches with the same batch size ie the same + number of sequences. + - `verbose`: print debug information; + - `seed`: random seed for reproducibility; + + Returns: + - `microbatch_ids`: list of tuple of batch id and samples ids per microbatch + - `batch_sizes`: the effective batch size of each batch, used for to compute the scaled LR + - `batch_max_seqlens`: the max seqlen across all microbatches in a batch + """ + + assert sequence_picking_order in ["random", "seqlen", "dataloader"] + if sequence_ids_per_mb is None: + metrics = list(zip(seqlens, range(len(seqlens)))) # use all samples + else: + metrics = list(zip(np.array(seqlens)[sequence_ids_per_mb], sequence_ids_per_mb)) + + if sequence_picking_order == 'random': + metric_random = random.Random(seed) + metric_random.shuffle(metrics) + if sequence_picking_order == 'seqlen': + metrics = sorted(metrics) + + # go through metrics, warn user, and filter samples that alone exceed the max batch threshold + long_ids = [idx for val, idx in metrics if val > max_tokens] + if len(long_ids) > 0: + logger.warning(f"Data indices {long_ids} ignored as metrics exceed {max_tokens}.") + logger.info(f"Original dataset length: {len(metrics)}. New dataset length: {len(long_ids)}") + metrics = [m for m in metrics if m[1] not in long_ids] + + def is_microbatch_valid(metrics): + if min_batch_size and len(metrics) < min_batch_size: return False # insufficient sample count + if max_batch_size and len(metrics) > max_batch_size: return False # too many samples + if sum([m[0] for m in metrics]) > max_tokens: return False # exceeds max + return True + + # go through all samples and pack then in microbatches of metric sums below the threshold + # `required_microbatches_of_same_size` means all minibatches in a batch must be of equal size + equal_size_multiple = effective_batch_size if required_microbatches_of_same_size else 1 + microbatches = [] + batch_init = 0 + while batch_init < len(metrics): + + # we iterate over possible effective batch sizes (groups of microbatches of same size) + valid_batch_end = batch_init + for batch_end in range(batch_init + equal_size_multiple, len(metrics), equal_size_multiple): + + # attempt effective batch + batch = metrics[batch_init:batch_end] + + # pick interleaved samples for each microbatch to help with load balancing + # (in the ordered use case), and to replicate what the distributed sampler does. + mbs = [batch[b::equal_size_multiple] for b in range(equal_size_multiple)] + + # if they are all valid micro-batches, keep them until you find longer mbatches, if any + is_batch_valid = all([is_microbatch_valid(mb) for mb in mbs]) + if is_batch_valid: + valid_batch_end = batch_end + + if batch_init == valid_batch_end: break # last batch is not valid (size zero), so we are done + batch = metrics[batch_init:valid_batch_end] + mbs = [batch[b::equal_size_multiple] for b in range(equal_size_multiple)] + batch_init += sum([len(l) for l in mbs]) + microbatches += mbs + + # make sure we give the same number of (micro-)batches to each dataloader by trimming the dataset + assert len(microbatches) >= effective_batch_size, "not enough datapoints to create a single sample per dataloader" + microbatches = microbatches[:len(microbatches) - len(microbatches) % effective_batch_size] + + #compute the effective batch size for each microbatch. + batch_sizes, batch_max_seqlens, microbatch_ids = [], [], [] + for rank in range(0, len(microbatches), effective_batch_size): + batch_id = rank // effective_batch_size + mbs = microbatches[rank:rank + effective_batch_size] + # compute the number of samples (not tokens) in this batch (not microbatch) + n_sequences = sum([len(mb) for mb in mbs]) + # compute the longest sequence (as number of tokens) in this batch (not microbatch) + sequence_ids_per_mb = [[m[1] for m in metrics] for metrics in mbs] + sequence_lens_per_mb = [[m[0] for m in metrics] for metrics in mbs] + batch_max_seqlen = max([max(seqlens) for seqlens in sequence_lens_per_mb]) + batch_and_mb_ids = zip([batch_id] * effective_batch_size, sequence_ids_per_mb) + batch_sizes.append(n_sequences) + batch_max_seqlens.append(batch_max_seqlen) + microbatch_ids += batch_and_mb_ids + if verbose: + n_tokens_per_mb = [sum([m[0] for m in mb]) for mb in mbs] + n_sequences_per_mb = [len(mb) for mb in mbs] + assert all([n <= max_tokens for n in n_tokens_per_mb]), "size of microbatch exceeds max tokens" + logger.info( + f"Batch id {batch_id} contains in total {len(mbs)} microbatches or {n_sequences} sequences. "\ + f"n_sequences per microbatch {n_sequences_per_mb}. "\ + f"n_tokens per microbatch {n_tokens_per_mb}. "\ + f"sequence ids per microbatch: {sequence_ids_per_mb}. "\ + f"sequence lengths per microbatch: {sequence_lens_per_mb}.") + + # return the sample ids of each microbatch, and the batch sizes + assert len(batch_sizes) == len(microbatch_ids) // effective_batch_size + return microbatch_ids, batch_sizes, batch_max_seqlens + + +def scale_lr(base_batch_size, batch_size, base_lr=1, method="linear"): + """ given a reference lr and batch_size, compute the new LR for a given batch size """ + if method == "linear": + # Linear Scaling Rule: "When the minibatch size is multiplied by k, multiply the learning + # rate by k" (Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, Goyal et al) + return base_lr * batch_size / base_batch_size + if method == "sqrt": + # Square Root scaling: "when multiplying the batch size by k, multiply the learning rate + # by √k, to keep the variance in the gradient expectation constant" + # (A. Krizhevsky. One weird trick for parallelizing convolutional neural networks) + return base_lr * torch.sqrt(batch_size / base_batch_size) + elif method == None or method.upper() == "NONE": + return base_lr + raise ValueError("Unknown scaling method: {}".format(method)) + + +def dataloader_for_variable_batch_size( + dataset, + microbatch_ids, + batch_max_seqlens, + dataloader_rank=0, + dataloader_batch_size=1, + dataloader_num_replicas=1, + dataloader_collate_fn=None, + dataloader_num_workers=2, + dataloader_pin_memory=False, + required_microbatches_of_same_seqlen=False, + sample_padding_fn=None, +): + + # equidistantly distribute the microbatches across the replicas in an interleaved fashion. + sampler = DistributedSampler( + dataset=microbatch_ids, + num_replicas=dataloader_num_replicas, + rank=dataloader_rank, + shuffle=False, + drop_last=False, + ) + + # collate function wraps user-defined collate function to the variable batch data + def collate_fn_wrapper(list_microbatch_ids): + # each batch is a list of sample ids that fill up to the max tokens per batch + # we return the collated batch of all dataset samples of all input batches. + batch = [] + for batch_id, microbatch_ids in list_microbatch_ids: + batch_data = [dataset[idx] for idx in microbatch_ids] + if required_microbatches_of_same_seqlen: + assert sample_padding_fn is not None, \ + "padding dataloader_padding_fn must be provided if required_microbatches_of_same_seqlen is True" + max_seqlen = batch_max_seqlens[batch_id] + assert all([len(sample) <= max_seqlen for sample in batch_data]), \ + "some samples are longer than the computed max seqlen for the batch those samples belong to" + batch_data = [sample_padding_fn(sample, max_seqlen) for sample in batch_data] + batch += batch_data + return dataloader_collate_fn(batch) if dataloader_collate_fn else batch + + dataloader = DataLoader( + dataset=microbatch_ids, + batch_size=dataloader_batch_size, + sampler=sampler, + num_workers=dataloader_num_workers, + collate_fn=collate_fn_wrapper, + pin_memory=dataloader_pin_memory, + ) + + deepspeed_io_kwargs = dict( + dataset=microbatch_ids, + batch_size=dataloader_batch_size, + pin_memory=dataloader_pin_memory, + data_sampler=sampler, + collate_fn=collate_fn_wrapper, + num_local_io_workers=dataloader_num_workers, + ) + + return dataloader, deepspeed_io_kwargs + + +class VariableBatchSizeLR(LRScheduler): + """ an LR scheduler that scales the LR of a given scheduler's LR """ + + @property + def optimizer(self): + return self.base_lr_scheduler.optimizer + + def __init__(self, + lr_scheduler, + base_batch_size, + batch_sizes, + dataloader, + lr_scaling_method="linear", + last_epoch=-1, + verbose=False): + self.batch_sizes = batch_sizes + self.base_batch_size = base_batch_size + self.lr_scaling_method = lr_scaling_method + self.dataloader = dataloader + self.base_lr_scheduler = lr_scheduler + # the following exist in LRScheduler but not in DeepSpeed's LRScheduler so we redefine them here + self.base_lrs = self.base_lr_scheduler.get_lr() + self.last_epoch = last_epoch + self.verbose = verbose + self.step(0) # scale LR for first sample in the dataloader + + def state_dict(self): + return { + 'base_lr_scheduler': self.base_lr_scheduler.state_dict() + } | { + 'base_batch_size': self.base_batch_size, + 'lr_scaling_method': self.lr_scaling_method, + 'batch_sizes': self.batch_sizes, + 'base_lrs': self.base_lrs, + 'last_epoch': self.last_epoch, + 'verbose': self.verbose, + } + + def load_state_dict(self, state_dict): + self.base_lr_scheduler.load_state_dict(state_dict['base_lr_scheduler']) + self.base_batch_size = state_dict['base_batch_size'] + self.lr_scaling_method = state_dict['lr_scaling_method'] + self.batch_sizes = state_dict['batch_sizes'] + self.base_lrs = state_dict['base_lrs'] + self.last_epoch = state_dict['last_epoch'] + self.verbose = state_dict['verbose'] + + def get_last_lr(self): + return self.base_lr_scheduler._last_lr + + def get_lr(self): + return [group['lr'] for group in self.base_lr_scheduler.optimizer.param_groups] + + def step(self, epoch=None): + # call the base scheduler's step method to get LR for next epoch + # Note: optimizer.step precedes lr_scheduler.step(), so the stepping workflow is: + # init: lr_scheduler.step(0) --> set LR for epoch 0 + # epoch 0: optimizer.step(); lr_scheduler.step(1) --> set LR for epoch 1 + # epoch 1: optimizer.step(); lr_scheduler.step(2) --> set LR for epoch 2 + + # reset unscaled LRs (to the original scheduler's one) to be able to step the base LR scheduler + # Note: epoch==0: reset LR scheduler; epoch==None: scale LR for next epoch; + unscaled_lrs = self.base_lrs if epoch == 0 else self.get_last_lr() + for group, lr in zip(self.base_lr_scheduler.optimizer.param_groups, unscaled_lrs): + group['lr'] = lr + + self.base_lr_scheduler.step(epoch) # set unscaled lr, _step_count, last_epoch, _last_lr for new epoch + + # scale the learning rate for the the next iteration for each parameter group. + self.last_epoch = self.last_epoch + 1 if epoch is None else epoch + # batch sizes are precomputed and stored in batch_sizes se we loop around to get the next one + batch_size = self.batch_sizes[self.last_epoch % len(self.batch_sizes)] + for group in self.base_lr_scheduler.optimizer.param_groups: + group['lr'] = scale_lr(self.base_batch_size, batch_size, group['lr'], self.lr_scaling_method) + + if self.verbose: + logger.info( + f"Next batch id {self.last_epoch}. "\ + f"Reference batch_size {self.base_batch_size} and lr {unscaled_lrs}. "\ + f"Scaled batch_size {batch_size} and lr {self.get_lr()}.") + + +def lr_scheduler_for_variable_batch_size(base_batch_size, + batch_sizes, + dataloader, + lr_scheduler_or_optimizer, + lr_scaling_method='linear', + verbose=False): + """ + returns a class that provides an LR scheduler that scales the learning rate at every + iteration taking into account the batch size of that iteration. + If learning rate is constant, ie no LR scheduler, then the base LR will be taken from the + constant LR values in the optimizer param groups. Otherwise from the scheduler's LR. + + Arguments: + - `base_batch_size`: the batch size that the base LR in the optimizer or scheduler refers to; + - `lr_scaling_method`: method to use to scale LR - see `scale_lr()`; + - `lr_scheduler_or_optimizer`: one instance of `LRScheduler` or `Optimizer` to be used as base; + - `batch_sizes`: the effective batch size of each batch in the dataloader; + + Returns the new LRScheduler + """ + + class StubLRScheduler(LRScheduler): + """ a stub LR scheduler that does not change the LR, keeps it constant """ + + def get_lr(self) -> float: + return self.base_lrs + + if isinstance(lr_scheduler_or_optimizer, Optimizer): + lr_scheduler = StubLRScheduler(lr_scheduler_or_optimizer) + elif hasattr(lr_scheduler_or_optimizer, 'optimizer'): #LRScheduler or DeepSpeed 'object' schedulers + assert isinstance(lr_scheduler_or_optimizer.optimizer, Optimizer) + lr_scheduler = lr_scheduler_or_optimizer + else: + raise ValueError("Unknown type for lr_scheduler_or_optimizer: {}".format(type(lr_scheduler_or_optimizer))) + + return VariableBatchSizeLR(lr_scheduler=lr_scheduler, + base_batch_size=base_batch_size, + batch_sizes=batch_sizes, + dataloader=dataloader, + lr_scaling_method=lr_scaling_method, + verbose=verbose) + + +def get_dataloader_and_lr_scheduler_for_variable_batch_size_deepspeed(dataset, + engine, + dataset_seqlens=None, + dataset_filter_ids=None, + dataloader_collate_fn=None, + sample_padding_fn=None, + batch_seqlens_fn=None): + """ + a simplified call to get_dataloader_and_lr_scheduler_for_variable_batch_size for the deepspeed runtime. + Needs the seqlens of every sample. It will try three alternatives: + - if `dataset_seqlens` is provided by user, use that. + - otherwise, looks for the seqlen metric path (in the connfig) that contains the output of the Data Analyzer + - otherwise, use the user-provided function `batch_seqlens_fn` and call Data Analyzer to output seqlen metric + See `batch_by_seqlens()` for arguments and more documentation. + """ + data_efficiency_config = engine._config.data_efficiency_config + data_sampling_config = data_efficiency_config[DATA_SAMPLING] + batching_config = data_sampling_config[DYNAMIC_BATCHING] + assert batching_config[DYNAMIC_BATCHING_ENABLED], "Dynamic batching is not enabled in the config" + + if dataset_seqlens is None: + # In seqlen provided by user, look for the seqlen metric that was output by the Data Analyzer + # (see the main in deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py for an example) + metrics_path = batching_config[DYNAMIC_BATCHING_METRICS_PATH] + sample_to_seqlen_path = os.path.join(metrics_path, "seqlen/seqlen_sample_to_metric") + if not (os.path.exists(f"{sample_to_seqlen_path}.bin") and os.path.exists(f"{sample_to_seqlen_path}.idx")): + # if the metric files are not found, we run the DataAnalyzer to write the metric files + msg = f"Cannot find metric files for sequence length in {sample_to_seqlen_path}.idx or *.bin." + msg += " We will run data analyzer to generated them..." + logger.warning(msg) + + if batch_seqlens_fn is None: + raise ValueError("sample_seqlen_fn must be provided if dataset_seqlens is not provided") + + DistributedDataAnalyzer( + dataset=dataset, + metric_functions=[batch_seqlens_fn], + collate_fn=dataloader_collate_fn, + batch_size=2**10, # batch size for map-reduce, not training + num_workers=engine.world_size, + worker_id=engine.global_rank, + save_path=pathlib.Path(metrics_path), + metric_types=['single_value_per_sample'], + metric_names=["seqlen"], + device=engine.device, + ).run_map_reduce() + + dataset_seqlens = MMapIndexedDataset(sample_to_seqlen_path, skip_warmup=True) + assert len(dataset_seqlens) == len(dataset), \ + "Seqlens size does not match the input dataset size. If you changed the dataset, delete the metrics_path folder." + + # TODO we are copying all seqlens into memory, we should adapt the code to use an iterative streamer + # and use the other files output by DataAnalyzer that returns an ordered dictionary of seqlen to sample ids + dataset_seqlens = np.array(list(dataset_seqlens), dtype=np.int64).flatten() # from Nx1 to N + + dataloader, lr_scheduler, deepspeed_io_kwargs = get_dataloader_and_lr_scheduler_for_variable_batch_size( + dataset=dataset, + dataset_filter_ids=dataset_filter_ids, + dataset_seqlens=dataset_seqlens, + effective_batch_size=engine.train_batch_size(), + max_tokens=batching_config[DYNAMIC_BATCHING_MAX_TOKENS], + lr_scaling_method=batching_config[DYNAMIC_BATCHING_LR_SCALING_METHOD], + sequence_picking_order=batching_config[DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER], + min_batch_size=batching_config[DYNAMIC_BATCHING_MIN_BATCH_SIZE], + max_batch_size=batching_config[DYNAMIC_BATCHING_MAX_BATCH_SIZE], + dataloader_batch_size=engine.train_micro_batch_size_per_gpu(), + dataloader_rank=engine.data_parallel_group.rank(), + dataloader_num_replicas=engine.data_parallel_group.size(), + dataloader_num_workers=data_sampling_config[DATA_SAMPLING_NUM_WORKERS], + dataloader_collate_fn=dataloader_collate_fn, + dataloader_pin_memory=data_sampling_config[DATA_SAMPLING_PIN_MEMORY], + sample_padding_fn=sample_padding_fn, + lr_scheduler_or_optimizer=engine.lr_scheduler or engine.optimizer, + required_microbatches_of_same_size=isinstance(engine, PipelineEngine), + required_microbatches_of_same_seqlen=isinstance(engine, PipelineEngine), + verbose=batching_config[DYNAMIC_BATCHING_VERBOSE], + seed=data_efficiency_config[DATA_EFFICIENCY_SEED], + ) + return dataloader, lr_scheduler, deepspeed_io_kwargs + + +def get_dataloader_and_lr_scheduler_for_variable_batch_size( + dataset, + dataset_seqlens, + max_tokens, + effective_batch_size, + dataset_filter_ids=None, + lr_scaling_method="linear", + min_batch_size=1, + max_batch_size=None, + sequence_picking_order="dataloader", + dataloader_batch_size=1, + dataloader_rank=0, + dataloader_num_replicas=1, + dataloader_num_workers=0, + dataloader_collate_fn=None, + dataloader_pin_memory=False, + lr_scheduler_or_optimizer=None, + required_microbatches_of_same_size=False, + required_microbatches_of_same_seqlen=False, + sample_padding_fn=None, + verbose=False, + seed=None, +): + """ returns a dataloader and LR scheduler for the variable batch size. see `batch_by_seqlens()` for details. """ + + # effective_batch_size = train_micro_batch_size_per_gpu * gradient_accumulation_steps * number of dataloaders + microbatch_ids, batch_sizes, batch_max_seqlens = batch_by_seqlens( + seqlens=dataset_seqlens, + max_tokens=max_tokens, + sequence_ids_per_mb=dataset_filter_ids, + min_batch_size=min_batch_size, + max_batch_size=max_batch_size, + sequence_picking_order=sequence_picking_order, + effective_batch_size=effective_batch_size, + required_microbatches_of_same_size=required_microbatches_of_same_size, + verbose=verbose, + seed=seed, + ) + + dataloader, deepspeed_io_kwargs = dataloader_for_variable_batch_size( + dataset=dataset, + microbatch_ids=microbatch_ids, + batch_max_seqlens=batch_max_seqlens, + dataloader_rank=dataloader_rank, + dataloader_num_replicas=dataloader_num_replicas, + dataloader_batch_size=dataloader_batch_size, + dataloader_collate_fn=dataloader_collate_fn, + dataloader_num_workers=dataloader_num_workers, + dataloader_pin_memory=dataloader_pin_memory, + required_microbatches_of_same_seqlen=required_microbatches_of_same_seqlen, + sample_padding_fn=sample_padding_fn, + ) + + lr_scheduler = lr_scheduler_for_variable_batch_size(base_batch_size=effective_batch_size, + batch_sizes=batch_sizes, + lr_scaling_method=lr_scaling_method, + lr_scheduler_or_optimizer=lr_scheduler_or_optimizer, + dataloader=dataloader, + verbose=verbose) + + return dataloader, lr_scheduler, deepspeed_io_kwargs diff --git a/lib/python3.12/site-packages/deepspeed/runtime/dataloader.py b/lib/python3.12/site-packages/deepspeed/runtime/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..499473b4ced81ba5ae5a447e32e03a2971b63c8f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/dataloader.py @@ -0,0 +1,162 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from torch.utils.data import DataLoader, RandomSampler +from torch.utils.data.distributed import DistributedSampler +from deepspeed.accelerator import get_accelerator + +from deepspeed.runtime.data_pipeline.data_sampling.data_sampler import DeepSpeedDataSampler +from deepspeed.runtime.data_pipeline.constants import CURRICULUM_LEARNING, \ + DATA_EFFICIENCY, DATA_SAMPLING_NUM_WORKERS +from deepspeed.runtime.constants import GRADIENT_ACCUMULATION_STEPS, \ + DATA_PARALLEL_GROUP, GLOBAL_RANK + + +class RepeatingLoader: + + def __init__(self, loader): + """Wraps an iterator to allow for infinite iteration. This is especially useful + for DataLoader types that we wish to automatically restart upon completion. + + Args: + loader (iterator): The data loader to repeat. + """ + self.loader = loader + self.data_iter = iter(self.loader) + + def __iter__(self): + return self + + def __next__(self): + try: + batch = next(self.data_iter) + except StopIteration: + self.data_iter = iter(self.loader) + batch = next(self.data_iter) + return batch + + +class DeepSpeedDataLoader(object): + + def __init__(self, + dataset, + batch_size, + pin_memory, + local_rank, + tput_timer, + collate_fn=None, + num_local_io_workers=None, + data_sampler=None, + data_parallel_world_size=None, + data_parallel_rank=None, + dataloader_drop_last=False, + deepspeed_dataloader_config={}): + self.deepspeed_dataloader_config = deepspeed_dataloader_config + self.tput_timer = tput_timer + self.batch_size = batch_size + self.curriculum_learning_enabled = False + if CURRICULUM_LEARNING in deepspeed_dataloader_config: + self.curriculum_learning_enabled = deepspeed_dataloader_config[CURRICULUM_LEARNING] + + if self.curriculum_learning_enabled: + data_sampler = DeepSpeedDataSampler(self.deepspeed_dataloader_config[DATA_EFFICIENCY], + len(dataset), + self.batch_size, + data_parallel_rank, + data_parallel_world_size, + self.deepspeed_dataloader_config[DATA_PARALLEL_GROUP], + self.deepspeed_dataloader_config[GRADIENT_ACCUMULATION_STEPS], + self.deepspeed_dataloader_config[GLOBAL_RANK], + drop_last=dataloader_drop_last) + device_count = get_accelerator().device_count() + num_local_io_workers = self.deepspeed_dataloader_config[DATA_SAMPLING_NUM_WORKERS] + else: + if local_rank >= 0: + if data_sampler is None: + data_sampler = DistributedSampler(dataset=dataset, + num_replicas=data_parallel_world_size, + rank=data_parallel_rank) + device_count = 1 + else: + if data_sampler is None: + data_sampler = RandomSampler(dataset) + device_count = get_accelerator().device_count() + batch_size *= device_count + + if num_local_io_workers is None: + num_local_io_workers = 2 * device_count + + self.num_local_io_workers = num_local_io_workers + self.data_sampler = data_sampler + self.dataset = dataset + self.collate_fn = collate_fn + self.device_count = device_count + self.batch_size = batch_size + self.pin_memory = pin_memory + self.data = None + self.dataloader_drop_last = dataloader_drop_last + self.post_process_func = None + + if self.dataloader_drop_last: + self.len = len(self.data_sampler) // self.batch_size + else: + from math import ceil + self.len = ceil(len(self.data_sampler) / self.batch_size) + + def __iter__(self): + self._create_dataloader() + return self + + def __len__(self): + return self.len + + def __next__(self): + if self.tput_timer: + self.tput_timer.start() + if self.curriculum_learning_enabled: + data = next(self.data_iterator) + if self.post_process_func is not None: + data = self.post_process_func(data, self.data_sampler.state_dict()) + return data + else: + return next(self.data) + + def _create_dataloader(self): + if self.curriculum_learning_enabled: + if self.collate_fn is None: + self.dataloader = DataLoader(self.dataset, + pin_memory=self.pin_memory, + batch_sampler=self.data_sampler, + num_workers=self.num_local_io_workers) + else: + self.dataloader = DataLoader(self.dataset, + pin_memory=self.pin_memory, + batch_sampler=self.data_sampler, + collate_fn=self.collate_fn, + num_workers=self.num_local_io_workers) + self.data_iterator = iter(self.dataloader) + return self.dataloader + else: + if self.collate_fn is None: + self.dataloader = DataLoader(self.dataset, + batch_size=self.batch_size, + pin_memory=self.pin_memory, + sampler=self.data_sampler, + num_workers=self.num_local_io_workers, + drop_last=self.dataloader_drop_last) + else: + self.dataloader = DataLoader(self.dataset, + batch_size=self.batch_size, + pin_memory=self.pin_memory, + sampler=self.data_sampler, + collate_fn=self.collate_fn, + num_workers=self.num_local_io_workers, + drop_last=self.dataloader_drop_last) + self.data = (x for x in self.dataloader) + + return self.dataloader + + +# DataLoader([(torch.randn(3, 3), torch.tensor(i % 2)) for i in range(10)], batch_size=2)) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/eigenvalue.py b/lib/python3.12/site-packages/deepspeed/runtime/eigenvalue.py new file mode 100644 index 0000000000000000000000000000000000000000..a82d8b1d5c7a04d793940fe8c85d23e15be91c69 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/eigenvalue.py @@ -0,0 +1,153 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.utils import log_dist +import numpy as np +import logging +from deepspeed.utils.torch import required_torch_version + + +class Eigenvalue(object): + + def __init__(self, + verbose=False, + max_iter=100, + tol=1e-2, + stability=0, + gas_boundary_resolution=1, + layer_name='', + layer_num=0): + super().__init__() + + self.verbose = verbose + self.max_iter = max_iter + self.tol = tol + self.stability = stability + self.gas_boundary_resolution = gas_boundary_resolution + self.layer_name = layer_name + self.layer_num = layer_num + + assert len(self.layer_name) > 0 and layer_num > 0 + + log_dist( + f'enabled eigenvalue with verbose={verbose}, max_iter={max_iter}, tol={tol}, stability={stability}, gas_boundary_resolution={gas_boundary_resolution}, layer_name={layer_name}, layer_num={layer_num}', + ranks=[0]) + + # Replace all nan/pos-inf/neg-inf to zero + def nan_to_num(self, x): + if required_torch_version(min_version=1.8): + return torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) + else: + # Fallback to numpy based implementation for backwards-compatibility with PyTorch 1.7 or older versions. + device = x.device + x = x.cpu().numpy() + x = np.nan_to_num(x=x, copy=False, nan=0.0, posinf=0.0, neginf=0.0) + return torch.from_numpy(x).to(device) + + def normalize(self, v): + norm_squared = self.inner_product(v, v) + norm = norm_squared**0.5 + self.stability + normalized_vectors = [vector / norm for vector in v] + normalized_vectors = [self.nan_to_num(vector) for vector in normalized_vectors] + return normalized_vectors + + def inner_product(self, xs, ys): + return sum([torch.sum(x * y) for (x, y) in zip(xs, ys)]) + + def get_layers(self, module): + scope_names = self.layer_name.split('.') + assert len(scope_names) > 0 + + m = module + for name in scope_names: + assert hasattr(m, name), "layer_name configuration is invalid." + m = getattr(m, name) + + return m + + def compute_eigenvalue(self, module, device=None, scale=1.0): + block_eigenvalue = [] + param_keys = [] + layers = self.get_layers(module) + + for block in range(self.layer_num): + model_block = layers[block] + + # We found this randn() has obvious accuracy impact in some cases, save/recover random state here. + rng_state = torch.random.get_rng_state() + if device is None: + v = [ + torch.randn(p.size()) for p in model_block.parameters() + if p.grad is not None and p.grad.grad_fn is not None + ] + else: + v = [ + torch.randn(p.size(), device=device) for p in model_block.parameters() + if p.grad is not None and p.grad.grad_fn is not None + ] + torch.random.set_rng_state(rng_state) + + grads = [ + param.grad for param in model_block.parameters() + if param.grad is not None and param.grad.grad_fn is not None + ] + params = [ + param for param in model_block.parameters() + if param.grad is not None and param.grad.grad_fn is not None + ] + + layer_keys = [id(p) for p in model_block.parameters()] + param_keys.append(layer_keys) + + v = self.normalize(v) + + # Disable eigenvalue if the model doesn't support second order gradients computation, + # e.g. when enabling DS transformer kernel. + if len(grads) == 0 or len(params) == 0: + log_dist(f'The model does NOT support eigenvalue computation.', ranks=[0], level=logging.WARNING) + return [] + + i = 0 + eigenvalue_current, eigenvalue_previous = 1., 0. + + while (i < self.max_iter) and abs(eigenvalue_current) > 0 and (abs( + (eigenvalue_current - eigenvalue_previous) / eigenvalue_current) + >= self.tol): # test convergence criteria + eigenvalue_previous = eigenvalue_current + + Hv = torch.autograd.grad(grads, params, grad_outputs=v, only_inputs=True, retain_graph=True) + #Hv = [hv.float() for hv in Hv] + Hv = [self.nan_to_num(hv).float() for hv in Hv] + + eigenvalue_current = self.inner_product(Hv, v).item() + + v = self.normalize(Hv) + v = [x / scale for x in v] + i += 1 + + eigenvalue_current *= scale + block_eigenvalue.append(eigenvalue_current) + + if self.verbose: + log_dist(f'block: {block}, power iteration: {i}, eigenvalue: {eigenvalue_current}', ranks=[0]) + + block_eigenvalue = self.post_process(block_eigenvalue) + + if self.verbose: + log_dist(f'post processed block_eigenvalue: {block_eigenvalue}', ranks=[0]) + + # {param_id: (eigenvalue, layer_id)} + ev_dict = {} + for i, (layer_keys, value) in enumerate(zip(param_keys, block_eigenvalue)): + ev_dict.update(dict.fromkeys(layer_keys, (value, i))) + + return ev_dict + + # 1. Map all eigenvalues to [0, 1.0]. + # 2. Some layers can't generate valid eigenvalues on fp16 precision, use 1.0 instead. + def post_process(self, value_list): + max_value = abs(max(value_list, key=abs)) + return [abs(v) / max_value if v != 0.0 else 1.0 for v in value_list] diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5067f71c8faf166bc78e88f9b62e8627dda7c7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +'''Copyright The Microsoft DeepSpeed Team''' diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f81294fcbd926fcb6bf91d14169c6a2f1a27ff83 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9aa7220d4ab0111a045c25b8a7a143eb257a9887 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/fused_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8367e22effa4c372a372d59a1e0067578dd8b231 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/loss_scaler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a52b628b6e5bcc9d734c53e1ac01c78cc05fed73 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/fp16/__pycache__/unfused_optimizer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/fused_optimizer.py b/lib/python3.12/site-packages/deepspeed/runtime/fp16/fused_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..49093bb73c8ff64ffb008b20df837766fc54bdcb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/fp16/fused_optimizer.py @@ -0,0 +1,514 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/apex +This file is adapted from FP16_Optimizer in NVIDIA/apex +""" + +import torch +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer +from deepspeed.runtime.utils import get_global_norm, get_flattened_grad_norm, CheckOverflow, get_weight_norm, get_norm_with_moe_layers, is_model_parallel_parameter +from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE +from deepspeed.utils import logger, log_dist +from deepspeed.utils.torch import required_torch_version +from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, CLIP_GRAD +from deepspeed.accelerator import get_accelerator +from deepspeed.moe.utils import is_moe_param_group +from deepspeed.runtime.constants import PIPE_REPLICATED +from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank + +OVERFLOW_CHECK_TIMER = 'overflow_check' +COMPUTE_NORM_TIMER = 'compute_norm' +UNSCALE_AND_CLIP_TIMER = 'unscale_and_clip' +BASIC_STEP_TIMER = 'basic_step' +UPDATE_FP16_TIMER = 'update_fp16' + +OVERFLOW_TIMERS = [COMPUTE_NORM_TIMER, OVERFLOW_CHECK_TIMER] +STEP_TIMERS = OVERFLOW_TIMERS + [UNSCALE_AND_CLIP_TIMER, BASIC_STEP_TIMER, UPDATE_FP16_TIMER] + + +class FP16_Optimizer(DeepSpeedOptimizer): + """ + FP16 Optimizer for training fp16 models. Handles loss scaling. + + For usage example please see, TODO: DeepSpeed V2 Tutorial + """ + + def __init__(self, + init_optimizer, + deepspeed=None, + static_loss_scale=1.0, + dynamic_loss_scale=False, + initial_dynamic_scale=2**32, + dynamic_loss_args=None, + verbose=True, + mpu=None, + clip_grad=0.0, + fused_adam_legacy=False, + has_moe_layers=False, + timers=None): + + self.fused_adam_legacy = fused_adam_legacy + self.timers = timers + self.deepspeed = deepspeed + self.has_moe_layers = has_moe_layers + self.using_pipeline = self.deepspeed.pipeline_parallelism + if not get_accelerator().is_available(): + raise SystemError("Cannot use fp16 without accelerator.") + self.optimizer = init_optimizer + + # param flattened by groups + self.fp16_groups = [] + self.fp16_groups_flat = [] + self.fp32_groups_flat = [] + + self.flatten_grad_norm_mask_list = [] + self.has_executed_step = False + self._global_grad_norm = 0. + + # loop to deal with groups + for i, param_group in enumerate(self.optimizer.param_groups): + # push this group to list before modify + self.fp16_groups.append(param_group['params']) + # init fp16 weight buffer, flattened + self.fp16_groups_flat.append(_flatten_dense_tensors([p.clone().detach() for p in self.fp16_groups[i]])) + # set model fp16 weight to slices of flattened buffer + updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i], self.fp16_groups[i]) + for p, q in zip(self.fp16_groups[i], updated_params): + p.data = q.data + # init master weight, flattened + self.fp32_groups_flat.append(self.fp16_groups_flat[i].clone().float().detach()) + # modify optimizer of have flat master weight + self.fp32_groups_flat[i].requires_grad = True # keep this in case internal optimizer uses it + param_group['params'] = [self.fp32_groups_flat[i]] + + # we may have a way of fusing dynamic scale. Do not support for now + if dynamic_loss_scale: + self.dynamic_loss_scale = True + self.cur_iter = 0 + self.last_overflow_iter = -1 + self.scale_factor = 2 + + if dynamic_loss_args is None: + self.cur_scale = initial_dynamic_scale + self.scale_window = 1000 + self.min_loss_scale = 1 + else: + self.cur_scale = dynamic_loss_args[INITIAL_LOSS_SCALE] + self.scale_window = dynamic_loss_args[SCALE_WINDOW] + self.min_loss_scale = dynamic_loss_args[MIN_LOSS_SCALE] + else: + self.dynamic_loss_scale = False + self.cur_iter = 0 + self.cur_scale = static_loss_scale + self.verbose = verbose + + self.custom_loss_scaler = False + self.external_loss_scale = None + + self.clip_grad = clip_grad + self.norm_type = 2 + + if required_torch_version(max_version=0.4): + self.clip_grad_norm = torch.nn.utils.clip_grad_norm + else: + self.clip_grad_norm = torch.nn.utils.clip_grad_norm_ + + #model parallel object + self.mpu = mpu + + self.overflow = False + self.overflow_checker = CheckOverflow(self.fp16_groups, mpu=self.mpu, deepspeed=deepspeed) + self.initialize_optimizer_states() + + def initialize_optimizer_states(self): + for i, group in enumerate(self.fp16_groups): + self.fp32_groups_flat[i].grad = torch.zeros(self.fp32_groups_flat[i].size(), + device=self.fp32_groups_flat[i].device) + + self.optimizer.step() + + for i, group in enumerate(self.fp16_groups): + self.fp32_groups_flat[i].grad = None + + return + + def zero_grad(self, set_to_none=True): + """ + Zero FP16 parameter grads. + """ + # For speed, set model fp16 grad to None by default + for group in self.fp16_groups: + for p in group: + if set_to_none: + p.grad = None + else: + if p.grad is not None: + p.grad.detach_() + p.grad.zero_() + + def step_fused_adam(self, closure=None): + """ + Not supporting closure. + """ + + # First compute norm for all group so we know if there is overflow + grads_groups_flat = [] + norm_groups = [] + for i, group in enumerate(self.fp16_groups): + grads_groups_flat.append( + _flatten_dense_tensors([ + torch.zeros(p.size(), dtype=p.dtype, device=p.device) if p.grad is None else p.grad for p in group + ])) + norm_groups.append(get_weight_norm(grads_groups_flat[i], mpu=self.mpu)) + + self.overflow = self.overflow_checker.check_using_norm(norm_groups) + prev_scale = self.cur_scale + self._update_scale(self.overflow) + + if self.overflow: + if self.verbose: + logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " + "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) + return self.overflow + + scaled_grad_norm = get_global_norm(norm_list=norm_groups) + + combined_scale = self.unscale_and_clip_grads(grads_groups_flat, scaled_grad_norm, apply_scale=False) + + # Stash unscaled gradient norm + self._global_grad_norm = scaled_grad_norm / self.cur_scale + + # norm is in fact norm*cur_scale + self.optimizer.step(grads=[[g] for g in grads_groups_flat], + output_params=[[p] for p in self.fp16_groups_flat], + scale=combined_scale, + grad_norms=norm_groups) + # TODO: we probably don't need this? just to be safe + for i in range(len(norm_groups)): + updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i], self.fp16_groups[i]) + for p, q in zip(self.fp16_groups[i], updated_params): + p.data = q.data + return self.overflow + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + """Return the current learning rate.""" + return self.optimizer.param_groups[0]["lr"] + + def override_loss_scale(self, loss_scale): + if loss_scale != self.external_loss_scale: + logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') + self.custom_loss_scaler = True + self.external_loss_scale = loss_scale + + def _require_avoid_recompute_norm(self, p, tensor_model_parallel_rank): + # for filtering replicated tensors from tensor + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + return True + if (tensor_model_parallel_rank > 0) and not is_model_parallel_parameter(p): + return True + + def _get_norm_mask_idx(self, group): + """The function preserves the parallel information for norm + from unflattened gradients. + + Args: + group (Iterable[Tensor] ): params group + + Returns: + torch.Tensor: A 2D tensor containing index ranges for each group, + where each row represents a [start index, end index]. + """ + group_mask_idx_list = [] + grad_flat_st_idx = 0 + grad_flat_en_idx = 0 + + for p in group: + grad_flat_en_idx = grad_flat_st_idx + p.numel() + if p.grad is not None and self._require_avoid_recompute_norm(p, bwc_tensor_model_parallel_rank(self.mpu)): + # merge range + if len(group_mask_idx_list) > 0 and grad_flat_st_idx == group_mask_idx_list[-1][-1]: + group_mask_idx_list[-1][-1] = grad_flat_en_idx + else: + group_mask_idx_list.append([grad_flat_st_idx, grad_flat_en_idx]) + grad_flat_st_idx = grad_flat_en_idx + + return torch.tensor(group_mask_idx_list, device=get_accelerator().current_device_name()) + + def step(self, closure=None): + """ + Not supporting closure. + """ + + if self.fused_adam_legacy: + return self.step_fused_adam() + + # First determine if there is overflow. + self.timers(OVERFLOW_CHECK_TIMER).start() + fp16_params = [] + for i, group in enumerate(self.fp16_groups): + fp16_params.extend([p for p in group if p.grad is not None]) + self.overflow = self.overflow_checker.has_overflow(fp16_params) + self.timers(OVERFLOW_CHECK_TIMER).stop() + prev_scale = self.cur_scale + self._update_scale(self.overflow) + if self.overflow: + if self.verbose: + log_dist( + "Overflow detected. Skipping step. Attempted loss " + f"scale: {prev_scale}, reducing to {self.cur_scale}", + ranks=[0]) + # Clear gradients + for i, group in enumerate(self.fp16_groups): + for p in group: + p.grad = None + + self.timers.log(OVERFLOW_TIMERS) + return self.overflow + + grads_groups_flat = [] + non_experts_grads_for_norm = [] + expert_grads_for_norm = {} + assert len(self.fp16_groups) == len(self.optimizer.param_groups) + + for i, group in enumerate(self.fp16_groups): + data_type = self.fp32_groups_flat[i].dtype + + grads_groups_flat.append( + _flatten_dense_tensors([ + torch.zeros(p.size(), dtype=data_type, device=p.device) if p.grad is None else p.grad.to(data_type) + for p in group + ])) + + self.fp32_groups_flat[i].grad = grads_groups_flat[i] + param_group = self.optimizer.param_groups[i] + + # split expert and non_expert grads for norm + if self.has_moe_layers and is_moe_param_group(param_group): + if param_group['name'] not in expert_grads_for_norm: + expert_grads_for_norm[param_group['name']] = [] + + expert_grads_for_norm[param_group['name']].append(self.fp32_groups_flat[i]) + else: + # retrieves the required mask for calculating the norm of flat_grad + # perform this collect operation only once + if not self.has_executed_step: + cur_flat_grad_norm_mask = self._get_norm_mask_idx(group) + self.flatten_grad_norm_mask_list.append(cur_flat_grad_norm_mask) + + non_experts_grads_for_norm.append(self.fp32_groups_flat[i]) + + for p in group: + p.grad = None + + self.timers(COMPUTE_NORM_TIMER).start() + + all_groups_norm = get_flattened_grad_norm(non_experts_grads_for_norm, + mpu=self.mpu, + grad_norm_mask=self.flatten_grad_norm_mask_list) + + if self.has_moe_layers: + all_groups_norm = get_norm_with_moe_layers(all_groups_norm, + mpu=self.mpu, + expert_tensors=expert_grads_for_norm, + norm_type=self.norm_type) + + scaled_global_grad_norm = get_global_norm(norm_list=[all_groups_norm]) + self.timers(COMPUTE_NORM_TIMER).stop() + + # Stash unscaled gradient norm + self._global_grad_norm = scaled_global_grad_norm / self.cur_scale + + self.timers(UNSCALE_AND_CLIP_TIMER).start() + self.unscale_and_clip_grads(grads_groups_flat, scaled_global_grad_norm) + self.timers(UNSCALE_AND_CLIP_TIMER).stop() + + self.timers(BASIC_STEP_TIMER).start() + self.optimizer.step() + self.timers(BASIC_STEP_TIMER).stop() + + #get rid of the fp32 gradients. Not needed anymore + for group in self.fp32_groups_flat: + group.grad = None + + self.timers(UPDATE_FP16_TIMER).start() + + for i in range(len(self.fp16_groups)): + updated_params = _unflatten_dense_tensors(self.fp32_groups_flat[i], self.fp16_groups[i]) + for p, q in zip(self.fp16_groups[i], updated_params): + p.data.copy_(q.data) + self.has_executed_step = True + self.timers(UPDATE_FP16_TIMER).stop() + + self.timers.log(STEP_TIMERS) + + return self.overflow + + def unscale_and_clip_grads(self, grad_groups_flat, total_norm, apply_scale=True): + # compute combined scale factor for this group + combined_scale = self.cur_scale + if self.clip_grad > 0.: + # norm is in fact norm*scale + clip = ((total_norm / self.cur_scale) + 1e-6) / self.clip_grad + if clip > 1: + combined_scale = clip * self.cur_scale + + if apply_scale: + for grad in grad_groups_flat: + grad.data.mul_(1. / combined_scale) + + return combined_scale + + def backward(self, loss, create_graph=False, retain_graph=False): + """ + :attr:`backward` performs the following steps: + + 1. fp32_loss = loss.float() + 2. scaled_loss = fp32_loss*loss_scale + 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves + """ + if self.custom_loss_scaler: + scaled_loss = self.external_loss_scale * loss + scaled_loss.backward() + else: + scaled_loss = (loss.float()) * self.cur_scale + scaled_loss.backward(create_graph=create_graph, retain_graph=retain_graph) + + def _update_scale(self, skip): + if self.dynamic_loss_scale: + prev_scale = self.cur_scale + if skip: + self.cur_scale = max(self.cur_scale / self.scale_factor, self.min_loss_scale) + self.last_overflow_iter = self.cur_iter + if self.verbose: + logger.info(f"\nGrad overflow on iteration {self.cur_iter}") + logger.info(f"Reducing dynamic loss scale from {prev_scale} to {self.cur_scale}") + else: + # Ensure self.scale_window updates since last overflow + stable_interval = (self.cur_iter - self.last_overflow_iter) - 1 + if (stable_interval > 0) and (stable_interval % self.scale_window == 0): + self.cur_scale *= self.scale_factor + if self.verbose: + logger.info(f"No Grad overflow for {self.scale_window} iterations") + logger.info(f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}") + else: + if skip: + logger.info("Grad overflow on iteration: %s", self.cur_iter) + logger.info("Using static loss scale of: %s", self.cur_scale) + self.cur_iter += 1 + return + + # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" + def _get_state(self): + return self.optimizer.state + + def _set_state(self, value): + self.optimizer.state = value + + state = property(_get_state, _set_state) + + # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" + # (for example, to adjust the learning rate) + def _get_param_groups(self): + return self.optimizer.param_groups + + def _set_param_groups(self, value): + self.optimizer.param_groups = value + + param_groups = property(_get_param_groups, _set_param_groups) + + def state_dict(self): + """ + Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. + This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict + of the contained Pytorch optimizer. + Example:: + checkpoint = {} + checkpoint['model'] = model.state_dict() + checkpoint['optimizer'] = optimizer.state_dict() + torch.save(checkpoint, "saved.pth") + """ + state_dict = {} + state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale + state_dict['cur_scale'] = self.cur_scale + state_dict['cur_iter'] = self.cur_iter + if state_dict['dynamic_loss_scale']: + state_dict['last_overflow_iter'] = self.last_overflow_iter + state_dict['scale_factor'] = self.scale_factor + state_dict['scale_window'] = self.scale_window + state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + state_dict['fp32_groups_flat'] = self.fp32_groups_flat + state_dict[CLIP_GRAD] = self.clip_grad + return state_dict + + # Refresh fp32 master params from fp16 copies + def refresh_fp32_params(self): + for current, saved in zip(self.fp32_groups_flat, self.fp16_groups_flat): + current.data.copy_(saved.data) + + def load_state_dict(self, state_dict, load_optimizer_states=True): + """ + Loads a state_dict created by an earlier call to state_dict(). + If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, + whose parameters in turn came from ``model``, it is expected that the user + will call ``model.load_state_dict()`` before + ``fp16_optimizer_instance.load_state_dict()`` is called. + Example:: + model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) + ... + checkpoint = torch.load("saved.pth") + model.load_state_dict(checkpoint['model']) + optimizer.load_state_dict(checkpoint['optimizer']) + """ + # I think it should actually be ok to reload the optimizer before the model. + self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] + self.cur_scale = state_dict['cur_scale'] + self.cur_iter = state_dict['cur_iter'] + if state_dict['dynamic_loss_scale']: + self.last_overflow_iter = state_dict['last_overflow_iter'] + self.scale_factor = state_dict['scale_factor'] + self.scale_window = state_dict['scale_window'] + if load_optimizer_states: + self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) + self.clip_grad = state_dict[CLIP_GRAD] + # At this point, the optimizer's references to the model's fp32 parameters are up to date. + # The optimizer's hyperparameters and internal buffers are also up to date. + # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still + # out of date. There are two options. + # 1: Refresh the master params from the model's fp16 params. + # This requires less storage but incurs precision loss. + # 2: Save and restore the fp32 master copies separately. + # We choose option 2. + # + # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device + # of their associated parameters, because it's possible those buffers might not exist yet in + # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been + # constructed in the same way as the one whose state_dict we are loading, the same master params + # are guaranteed to exist, so we can just copy_() from the saved master params. + for current, saved in zip(self.fp32_groups_flat, state_dict['fp32_groups_flat']): + current.data.copy_(saved.data) + + def __repr__(self): + return repr(self.optimizer) + + # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" + def _get_loss_scale(self): + if self.custom_loss_scaler: + return self.external_loss_scale + else: + return self.cur_scale + + def _set_loss_scale(self, value): + self.loss_scaler.cur_scale = value + + loss_scale = property(_get_loss_scale, _set_loss_scale) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/loss_scaler.py b/lib/python3.12/site-packages/deepspeed/runtime/fp16/loss_scaler.py new file mode 100644 index 0000000000000000000000000000000000000000..451451c51a3218b26667555ba6079536d09b6bda --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/fp16/loss_scaler.py @@ -0,0 +1,270 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +Taken and modified for DeepSpeed from: + https://github.com/NVIDIA/Megatron-LM/blob/master/fp16/loss_scaler.py +Commit: 93ab4bea59dc5cbf97c079d313741866af4deac9 +""" + +import torch +from deepspeed import comm as dist +from deepspeed.utils import logger + +INITIAL_LOSS_SCALE = 'init_scale' +SCALE_WINDOW = 'scale_window' +DELAYED_SHIFT = 'delayed_shift' +CONSECUTIVE_HYSTERESIS = 'consecutive_hysteresis' +MIN_LOSS_SCALE = 'min_scale' + + +# item() is a recent addition, so this helps with backward compatibility. +def to_python_float(t): + if hasattr(t, 'item'): + return t.item() + return t[0] + + +class LossScalerBase: + """LossScalarBase + Base class for a loss scaler + """ + + def __init__(self, cur_scale): + self.cur_scale = cur_scale + self.dynamic = False + + @property + def loss_scale(self): + return self.cur_scale + + def scale_gradient(self, module, grad_in, grad_out): + return tuple(self.loss_scale * g for g in grad_in) + + def update_scale(self, overflow): + pass + + def backward(self, loss, retain_graph=False): + scaled_loss = loss * self.loss_scale + scaled_loss.backward(retain_graph=retain_graph) + # print(f'LossScalerBackward: {scaled_loss=}') + + +class LossScaler(LossScalerBase): + """ + Class that manages a static loss scale. This class is intended to interact with + :class:`FP16_Optimizer`, and should not be directly manipulated by the user. + + Use of :class:`LossScaler` is enabled via the ``static_loss_scale`` argument to + :class:`FP16_Optimizer`'s constructor. + + Args: + scale (float, optional, default=1.0): The loss scale. + """ + + def __init__(self, scale=1): + super(LossScaler, self).__init__(scale) + + # `params` is a list / generator of torch.Variable + def has_overflow(self, params): + return False + + # `x` is a torch.Tensor + def _has_inf_or_nan(x): + return False + + +class DynamicLossScaler(LossScalerBase): + """ + Class that manages dynamic loss scaling. It is recommended to use :class:`DynamicLossScaler` + indirectly, by supplying ``dynamic_loss_scale=True`` to the constructor of + :class:`FP16_Optimizer`. However, it's important to understand how :class:`DynamicLossScaler` + operates, because the default options can be changed using the + the ``dynamic_loss_args`` argument to :class:`FP16_Optimizer`'s constructor. + + Loss scaling is designed to combat the problem of underflowing gradients encountered at long + times when training fp16 networks. Dynamic loss scaling begins by attempting a very high loss + scale. Ironically, this may result in OVERflowing gradients. If overflowing gradients are + encountered, :class:`DynamicLossScaler` informs :class:`FP16_Optimizer` that an overflow has + occurred. + :class:`FP16_Optimizer` then skips the update step for this particular iteration/minibatch, + and :class:`DynamicLossScaler` adjusts the loss scale to a lower value. + If a certain number of iterations occur without overflowing gradients detected, + :class:`DynamicLossScaler` increases the loss scale once more. + In this way :class:`DynamicLossScaler` attempts to "ride the edge" of + always using the highest loss scale possible without incurring overflow. + + Args: + init_scale (float, optional, default=2**32): Initial loss scale attempted by :class:`DynamicLossScaler.` + scale_factor (float, optional, default=2.0): Factor used when adjusting the loss scale. If an overflow is encountered, the loss scale is readjusted to loss scale/``scale_factor``. If ``scale_window`` consecutive iterations take place without an overflow, the loss scale is readjusted to loss_scale*``scale_factor``. + scale_window (int, optional, default=1000): Number of consecutive iterations without an overflow to wait before increasing the loss scale. + consecutive_hysteresis (bool, optional, default=False): Whether to refill hysteresis if we reach an iteration that doesn't overflow + """ + + def __init__(self, + init_scale=2**32, + scale_factor=2., + scale_window=1000, + min_scale=1, + delayed_shift=1, + consecutive_hysteresis=False, + raise_error_at_min_scale=True, + dtype=torch.half): + super(DynamicLossScaler, self).__init__(init_scale) + self.cur_iter = 0 + self.last_overflow_iter = -1 + self.scale_factor = scale_factor + self.scale_window = scale_window + self.min_scale = min_scale + self.delayed_shift = delayed_shift + self.cur_hysteresis = delayed_shift + self.consecutive_hysteresis = consecutive_hysteresis + self.raise_error_at_min_scale = raise_error_at_min_scale + self.dynamic = True + self.dtype = dtype + + # `params` is a list / generator of torch.Variable + def has_overflow_serial(self, params): + for p in params: + if p.grad is not None and self._has_inf_or_nan(p.grad.data): + return True + + return False + + # `x` is a torch.Tensor + def _has_inf_or_nan(x): + try: + # if x is half, the .float() incurs an additional deep copy, but it's necessary if + # Pytorch's .sum() creates a one-element tensor of the same type as x + # (which is true for some recent version of pytorch). + cpu_sum = float(x.float().sum()) + # More efficient version that can be used if .sum() returns a Python scalar + # cpu_sum = float(x.sum()) + except RuntimeError as instance: + # We want to check if inst is actually an overflow exception. + # RuntimeError could come from a different error. + # If so, we still want the exception to propagate. + if "value cannot be converted" not in instance.args[0]: + raise + return True + else: + if cpu_sum in [float('inf'), -float('inf')] or cpu_sum != cpu_sum: + return True + return False + + # `overflow` is boolean indicating whether the gradient overflowed + def update_scale(self, overflow): + if overflow: + # self.cur_scale /= self.scale_factor + if self.delayed_shift == 1 or self.cur_hysteresis == 1: + if (self.cur_scale == self.min_scale) and self.raise_error_at_min_scale: + raise Exception( + "Current loss scale already at minimum - cannot decrease scale anymore. Exiting run.") + else: + next_scale = max(self.cur_scale / self.scale_factor, self.min_scale) + if dist.get_rank() == 0: + overflow_msg = f"[deepspeed] OVERFLOW! Rank {dist.get_rank()} Skipping step." + if self.dtype == torch.half: + overflow_msg += f" Attempted loss scale: {int(self.cur_scale)}, reducing to {int(next_scale)}" + logger.info(overflow_msg) + self.cur_scale = next_scale + else: + if dist.get_rank() == 0: + overflow_msg = f"[deepspeed] OVERFLOW! Rank {dist.get_rank()} Skipping step." + if self.dtype == torch.half: + overflow_msg += f" Attempted loss scale: {int(self.cur_scale)}, but hysteresis is {self.cur_hysteresis}. Reducing hysteresis to {self.cur_hysteresis-1}" + logger.info(overflow_msg) + self.cur_hysteresis -= 1 + self.last_overflow_iter = self.cur_iter + else: + if self.consecutive_hysteresis: + if dist.get_rank() == 0: + hysteresis_msg = f"Consecutive hysteresis is enabled. Restoring hysteresis to {self.delayed_shift}" + logger.info(hysteresis_msg) + self.cur_hysteresis = self.delayed_shift + if (self.cur_iter - self.last_overflow_iter) % self.scale_window == 0: + if not self.consecutive_hysteresis: + self.cur_hysteresis = self.delayed_shift + self.cur_scale *= self.scale_factor + self.cur_iter += 1 + + +# Although loss scaling is only defined for fp16, yet for backwards compatibility +# we still create a scaler for other dtypes (fp32, bf16) which does not perform any scaling. +def CreateLossScaler(dtype, static_loss_scale, dynamic_scaling, dynamic_loss_args): + if dtype == torch.half and dynamic_scaling: + if dynamic_loss_args is None: + return DynamicLossScaler(dtype=dtype) + return DynamicLossScaler(dtype=dtype, **dynamic_loss_args) + + loss_scale_value = static_loss_scale if dtype == torch.half else 1.0 + return LossScaler(scale=loss_scale_value) + + +############################################################## +# Example usage below here -- assuming it's in a separate file +############################################################## +""" +TO-DO separate out into an example. +if __name__ == "__main__": + import torch + from torch.autograd import Variable + from dynamic_loss_scaler import DynamicLossScaler + + # N is batch size; D_in is input dimension; + # H is hidden dimension; D_out is output dimension. + N, D_in, H, D_out = 64, 1000, 100, 10 + + # Create random Tensors to hold inputs and outputs, and wrap them in Variables. + x = Variable(torch.randn(N, D_in), requires_grad=False) + y = Variable(torch.randn(N, D_out), requires_grad=False) + + w1 = Variable(torch.randn(D_in, H), requires_grad=True) + w2 = Variable(torch.randn(H, D_out), requires_grad=True) + parameters = [w1, w2] + + learning_rate = 1e-6 + optimizer = torch.optim.SGD(parameters, lr=learning_rate) + loss_scaler = DynamicLossScaler() + + for t in range(500): + y_pred = x.mm(w1).clamp(min=0).mm(w2) + loss = (y_pred - y).pow(2).sum() * loss_scaler.loss_scale + print('Iter {} loss scale: {}'.format(t, loss_scaler.loss_scale)) + print('Iter {} scaled loss: {}'.format(t, loss.data[0])) + print('Iter {} unscaled loss: {}'.format(t, loss.data[0] / loss_scaler.loss_scale)) + + # Run backprop + optimizer.zero_grad() + loss.backward() + + # Check for overflow + has_overflow = DynamicLossScaler.has_overflow(parameters) + + # If no overflow, unscale grad and update as usual + if not has_overflow: + for param in parameters: + param.grad.data.mul_(1. / loss_scaler.loss_scale) + optimizer.step() + # Otherwise, don't do anything -- ie, skip iteration + else: + print('fp16 dynamic loss scale overflow!') + + # Update loss scale for next iteration + loss_scaler.update_scale(has_overflow) + +""" diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ac6f7e9784ce60f6c1b4a9134b73b4e415337641 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .adam import OnebitAdam +from .lamb import OnebitLamb +from .zoadam import ZeroOneAdam diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e064591e25a37d6b0595414f538118acaad02cc Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..20dd3359857b15eb0e2803cc103aea5ba37a53e2 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/adam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca66c54f1e72933861563282af4b13c5718a2555 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/lamb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e74bfb2319171ffa9458c16fb7cdd4c482b19e2 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/__pycache__/zoadam.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/adam.py b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/adam.py new file mode 100644 index 0000000000000000000000000000000000000000..fa817573f734801070de5391d161695389351ec3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/adam.py @@ -0,0 +1,310 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import types +import torch +import numpy as np +from deepspeed.accelerator import get_accelerator +from deepspeed.utils.torch import required_torch_version +from deepspeed import comm as dist + + +class OnebitAdam(torch.optim.Optimizer): + """Implements the 1-bit Adam algorithm. Currently GPU-only. + For usage example please see https://www.deepspeed.ai/tutorials/onebit-adam/ + For technical details please read https://arxiv.org/abs/2102.02888 + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + freeze_step (int, optional): Number of steps for warmup (uncompressed) + stage before we start using compressed communication. (default 100000) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) NOT SUPPORTED in 1-bit Adam! + eps_inside_sqrt (boolean, optional): in the 'update parameters' step, + adds eps to the bias-corrected second moment estimate before + evaluating square root instead of adding it to the square root of + second moment estimate as in the original paper. (default: False) + cuda_aware (boolean, required): Set True if the underlying MPI implementation + supports CUDA-Aware communication. (default: False) + comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') + .. _Adam\\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, + params, + deepspeed=None, + lr=1e-3, + freeze_step=100000, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + eps_inside_sqrt=False, + weight_decay=0., + max_grad_norm=0., + amsgrad=False, + cuda_aware=False, + comm_backend_name='nccl'): + + if amsgrad: + raise RuntimeError('1-bit Adam does not support the AMSGrad variant.') + + defaults = dict(lr=lr, + bias_correction=bias_correction, + betas=betas, + eps=eps, + weight_decay=weight_decay, + max_grad_norm=max_grad_norm) + + super(OnebitAdam, self).__init__(params, defaults) + self.eps_mode = 0 if eps_inside_sqrt else 1 + self.comm_time = 0.0 + self.step_time = 0.0 + self.ave_step = 1 + self.bk_time = 0.0 + + self.deepspeed = deepspeed + self.adam_freeze_key = False + self.initialize = False + self.freeze_step = freeze_step + self.cuda_aware = cuda_aware + self.using_pipeline = False + + self.comm_backend_name = comm_backend_name + + assert dist.is_initialized(), "Please initialize the torch distributed backend." + # Empty initializer. Set handle based on the comm backend as follows. + self.comm_backend_handle = None + if self.comm_backend_name == 'nccl': + assert ( + required_torch_version(min_version=1.8) + ), "Please use torch 1.8 or greater to enable NCCL backend in 1-bit Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" + from deepspeed.runtime.comm.nccl import NcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'mpi': + from deepspeed.runtime.comm.mpi import MpiBackend + self.comm_backend_handle = MpiBackend(cuda_aware) + elif self.comm_backend_name == 'hccl': + from deepspeed.runtime.comm.hccl import HcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'compressed': + from deepspeed.runtime.comm.compressed import CompressedBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) + self.size = self.comm_backend_handle.size + + self.divider = int(self.size * 8 / np.gcd(self.size, 8)) + + def step(self, closure=None, grads=None): + """Performs a single optimization step. + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + grads (list of tensors, optional): weight gradient to use for the + optimizer update. If gradients have type torch.half, parameters + are expected to be in type torch.float. (default: None) + output params (list of tensors, optional): A reduced precision copy + of the updated weights written out in addition to the regular + updated weights. Have to be of same type as gradients. (default: None) + scale (float, optional): factor to divide gradient tensor values + by before applying to weights. (default: 1) + """ + loss = None + if closure is not None: + loss = closure() + + gather_time = 0 + allgather_time = 0 + all_time = 0 + + if self.adam_freeze_key is False: + v_diff_buffer = 0.0 + + if grads is None: + grads_group = [None] * len(self.param_groups) + # backward compatibility + # assuming a list/generator of parameter means single group + elif isinstance(grads, types.GeneratorType): + grads_group = [grads] + elif type(grads[0]) != list: + grads_group = [grads] + else: + grads_group = grads + + for group, grads_this_group in zip(self.param_groups, grads_group): + if grads_this_group is None: + grads_this_group = [None] * len(group['params']) + + bias_correction = 1 if group['bias_correction'] else 0 + + for p, grad in zip(group['params'], grads_this_group): + if p.grad is None and grad is None: + continue + if grad is None: + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('1-bit Adam does not support sparse gradients') + + state = self.state[p] + + # State initialization + if len(state) == 0: + state['step'] = 0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + + if not self.initialize or (self.adam_freeze_key and 'worker_error' not in state.keys()): + state['tensor_size'] = torch.numel(p.data) + state['corrected_tensor_size'] = state['tensor_size'] + + if state['tensor_size'] % (self.size * self.divider) != 0: + state['corrected_tensor_size'] += ((self.size * self.divider) - (state['tensor_size'] % + (self.size * self.divider))) + state['server_chunk_size'] = state['corrected_tensor_size'] // self.size + get_accelerator().empty_cache() + state['worker_error'] = torch.zeros(state['corrected_tensor_size'], device=p.device) + state['server_error'] = torch.zeros(state['server_chunk_size'], device=p.device) + get_accelerator().empty_cache() + self.adam_freeze_key = True + if not self.initialize and dist.get_rank() == 0: + print("Cupy Buffers Initialized Successfully.") + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + state['step'] += 1 + + if self.adam_freeze_key is False: + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + grad = None + if self.initialize: + update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) + + else: + if 'non_freeze' in group.keys() and group['non_freeze'] is True: + dist.all_reduce(grad) + grad.mul_(1 / dist.get_world_size()) + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + grad = None + else: + if self.initialize is True: + exp_avg.mul_(beta1).add_(1 - beta1, grad) + grad = None + + if self.size > 1: + exp_avg.set_( + self.comm_backend_handle.compressed_allreduce(exp_avg, state['worker_error'], + state['server_error'], + self.deepspeed.local_rank)) + # Because 1-bit compression cannot represent exact zero, it is required to + # provide a momentum mask for those params that have constant exact zeros in their + # momentums, otherwise the compression error would keep accumulating. + # For example, for BERT pre-training seq 128, bert.embeddings.position_embeddings.weight + # always have exact zeros in its momentum for row 129 to 512, because it only + # learns up to seq length 128 while the model supports up to 512 seq length. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) + if 'exp_avg_mask' in group: + if exp_avg.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=exp_avg.device) + exp_avg.mul_(group['exp_avg_mask']) + + if self.initialize: + update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) + + if self.initialize: + if group['weight_decay'] > 0.0: + update += group['weight_decay'] * p.data + with torch.no_grad(): + p.add_(-group['lr'] * update) + + if not self.initialize: + print('Pop out errors', flush=True) + state.pop('worker_error') + state.pop('server_error') + + if not self.initialize: + self.adam_freeze_key = False + self.initialize = True + print(f"Finished the initialization step at rank {dist.get_rank()}") + return loss + + if self.adam_freeze_key is False: + if state['step'] >= self.freeze_step: + print('OnebitAdam - starting compressed communication') + self.adam_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + + return loss + + def load_state_dict(self, state_dict): + """ + Overrides load_state_dict() to add special handling when loading checkpoints + """ + # Because at different stage exp_avg_mask may change (e.g., + # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask + # in checkpoints but always use the one user provided in training script. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) + # Thus here we keep the exp_avg_mask unchanged when loading checkpoint + for i, group in enumerate(self.param_groups): + if 'exp_avg_mask' in group: + state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] + elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: + state_dict['param_groups'][i].pop('exp_avg_mask') + super().load_state_dict(state_dict) + if self.state[self.param_groups[0]['params'][0]]['step'] < self.freeze_step: + if dist.get_rank() == 0: + print("Checkpoint loaded and OnebitAdam warmup stage starts/continues.") + if self.adam_freeze_key is True: + self.adam_freeze_key = False + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = True + else: + self.deepspeed.enable_backward_allreduce = True + else: + if dist.get_rank() == 0: + print("Checkpoint loaded and OnebitAdam compression stage starts/continues.") + if self.adam_freeze_key is False: + self.adam_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + # We reset the compression errors when loading checkpoints for 3 reasons: + # 1) The worker and server error at each GPU are distinct, so in current implementation + # only rank 0's errors are saved in the checkpoint. Thus we have to reset the errors. + # If we want to save them correctly we need O(num_gpu*model_size) memory in order to + # gather all the error, which is a very large memory requirement. It's possible to save + # them in a distributed way, but it will make the checkpoint saving/loading much more complicated. + # 2) Even if we are able to save the compression errors correctly, you need to have the + # exact same number of GPUs in order to load them correctly. + # 3) We verified on BERT pre-training that occasionally resetting the compression error + # at checkpoint loading does not affect the convergence. + # However, please avoid frequent checkpoint loading which could break the error + # compensation mechanism thus affect the convergence. + for group in self.param_groups: + for p in group['params']: + if 'worker_error' in self.state[p]: + self.state[p].pop('worker_error') + if 'server_error' in self.state[p]: + self.state[p].pop('server_error') diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/lamb.py b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/lamb.py new file mode 100644 index 0000000000000000000000000000000000000000..9e7bae816ecdee5aa7e5aa2349a62f1b43ea5e10 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/lamb.py @@ -0,0 +1,447 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import types +import torch +import numpy as np +from deepspeed import comm as dist +from deepspeed.utils.torch import required_torch_version +from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from deepspeed.accelerator import get_accelerator + + +class OnebitLamb(torch.optim.Optimizer): + """Implements the 1-bit Lamb algorithm. Currently GPU-only. + For usage example please see https://www.deepspeed.ai/tutorials/onebit-lamb/ + For technical details please see our paper https://arxiv.org/abs/2104.06069. + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + freeze_step (int, optional): Number of steps for warmup (uncompressed) + stage before we start using compressed communication. (default 100000) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + max_coeff(float, optional): maximum value of the lamb coefficient (default: 10.0) + min_coeff(float, optional): minimum value of the lamb coefficient (default: 0.01) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) NOT SUPPORTED in 1-bit Lamb! + eps_inside_sqrt (boolean, optional): in the 'update parameters' step, + adds eps to the bias-corrected second moment estimate before + evaluating square root instead of adding it to the square root of + second moment estimate as in the original paper. (default: False) + cuda_aware (boolean, required): Set True if the underlying MPI implementation + supports CUDA-Aware communication. (default: False) + comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') + coeff_beta (float, optional): coefficient used for computing + running averages of lamb coefficient (default: 0.9) note that you may want to + increase or decrease this beta depending on the freeze_step you choose, as + 1/(1 - coeff_beta) should be smaller than or equal to freeze_step + factor_max (float, optional): maximum value of scaling factor to the frozen lamb + coefficient during compression stage (default: 4.0) + factor_min (float, optional): minimum value of scaling factor to the frozen lamb + coefficient during compression stage (default: 0.5) + factor_threshold (float, optional): threshold of how much the scaling factor can + fluctuate between steps (default: 0.1) + .. _Large Batch Optimization for Deep Learning\\: Training BERT in 76 minutes: + https://arxiv.org/abs/1904.00962 + .. _Adam\\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, + params, + deepspeed=None, + lr=1e-3, + freeze_step=100000, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + eps_inside_sqrt=False, + weight_decay=0., + max_grad_norm=0., + max_coeff=10.0, + min_coeff=0.01, + amsgrad=False, + cuda_aware=False, + comm_backend_name='nccl', + coeff_beta=0.9, + factor_max=4.0, + factor_min=0.5, + factor_threshold=0.1): + + if amsgrad: + raise RuntimeError('1-bit Lamb does not support the AMSGrad variant.') + + defaults = dict(lr=lr, + bias_correction=bias_correction, + betas=betas, + eps=eps, + weight_decay=weight_decay, + max_grad_norm=max_grad_norm, + max_coeff=max_coeff, + min_coeff=min_coeff) + + super(OnebitLamb, self).__init__(params, defaults) + self.eps_mode = 0 if eps_inside_sqrt else 1 + self.deepspeed = deepspeed + self.lamb_freeze_key = False + self.initialize = False + self.freeze_step = freeze_step + self.cuda_aware = cuda_aware + self.coeff_beta = coeff_beta + self.factor_max = factor_max + self.factor_min = factor_min + self.factor_threshold = factor_threshold + self.using_pipeline = False + + self.comm_backend_name = comm_backend_name + + assert dist.is_initialized(), "Please initialize the torch distributed backend." + # Empty initializer. Set handle based on the comm backend as follows. + self.comm_backend_handle = None + if self.comm_backend_name == 'nccl': + assert ( + required_torch_version(min_version=1.8) + ), "Please use torch 1.8 or greater to enable NCCL backend in 1-bit Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" + from deepspeed.runtime.comm.nccl import NcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'mpi': + from deepspeed.runtime.comm.mpi import MpiBackend + self.comm_backend_handle = MpiBackend(cuda_aware) + elif self.comm_backend_name == 'hccl': + from deepspeed.runtime.comm.hccl import HcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'compressed': + from deepspeed.runtime.comm.compressed import CompressedBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) + + self.size = self.comm_backend_handle.size + + self.divider = int(self.size * 8 / np.gcd(self.size, 8)) + + self.exp_avg_flat = [] + self.dummy_exp_avg = {} + self.corrected_tensor_sizes = [] + self.server_chunk_sizes = [] + self.worker_errors = [] + self.server_errors = [] + + self.lamb_coeffs = [] + + def step(self, closure=None, grads=None): + """Performs a single optimization step. + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + grads (list of tensors, optional): weight gradient to use for the + optimizer update. If gradients have type torch.half, parameters + are expected to be in type torch.float. (default: None) + """ + loss = None + if closure is not None: + loss = closure() + + if grads is None: + grads_group = [None] * len(self.param_groups) + # backward compatibility + # assuming a list/generator of parameter means single group + elif isinstance(grads, types.GeneratorType): + grads_group = [grads] + elif type(grads[0]) != list: + grads_group = [grads] + else: + grads_group = grads + + # remove the previous stats + del self.lamb_coeffs[:] + + if self.lamb_freeze_key: + exp_avg_last_step = [] + for group in self.param_groups: + exp_avg_last_step.append([self.state[p]['exp_avg'].detach().clone() for p in group['params']]) + if 'scaling_coeff' not in self.state[self.param_groups[0]['params'][0]]: + # Compute the scaling_coeff for each momentum at the end of warmup stage. + # This is used to reduce compression error during compression stage. + momentum_scales = [] + for group in self.param_groups: + momentum_scales.append([(torch.linalg.vector_norm(self.state[p]['exp_avg']) / + np.sqrt(torch.numel(self.state[p]['exp_avg']))).item() + for p in group['params']]) + united_scale = sum([sum(x) for x in momentum_scales]) / sum([len(x) for x in momentum_scales]) + for i, group in enumerate(self.param_groups): + for j, p in enumerate(group['params']): + self.state[p]['scaling_coeff'] = united_scale / momentum_scales[i][j] + + for group, grads_this_group in zip(self.param_groups, grads_group): + if grads_this_group is None: + grads_this_group = [None] * len(group['params']) + + bias_correction = 1 if group['bias_correction'] else 0 + + for p, grad in zip(group['params'], grads_this_group): + if p.grad is None and grad is None: + continue + if grad is None: + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('1-bit Lamb does not support sparse gradients') + + state = self.state[p] + + # State initialization + if len(state) == 0 or (len(state) == 1 and 'scaling_coeff' in state.keys()): + state['step'] = 0 + state['lamb_coeff_freeze'] = 0.0 + state['last_factor'] = 1.0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + state['exp_avg_sq_fresh'] = torch.zeros_like(p.data) + + if not self.initialize: + self.lamb_freeze_key = True + + exp_avg, exp_avg_sq, exp_avg_sq_fresh = state['exp_avg'], state['exp_avg_sq'], state[ + 'exp_avg_sq_fresh'] + beta1, beta2 = group['betas'] + max_coeff = group['max_coeff'] + min_coeff = group['min_coeff'] + + state['step'] += 1 + + if self.lamb_freeze_key is False: + # warmup stage, baseline Lamb optimization + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + if state['step'] == self.freeze_step: + exp_avg_sq_fresh.data = exp_avg_sq.detach().clone() + grad = None + if self.initialize: + weight_norm = p.data.pow(2).sum().sqrt() + update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) + if group['weight_decay'] > 0.0: + update += group['weight_decay'] * p.data + update_norm = update.pow(2).sum().sqrt() + lamb_coeff = 1.0 + if weight_norm != 0 and update_norm != 0: + lamb_coeff = (weight_norm / update_norm).item() + if lamb_coeff > max_coeff: + lamb_coeff = max_coeff + if lamb_coeff < min_coeff: + lamb_coeff = min_coeff + if lamb_coeff != 1.0: + state['lamb_coeff_freeze'] = self.coeff_beta * state['lamb_coeff_freeze'] + ( + 1 - self.coeff_beta) * lamb_coeff + self.lamb_coeffs.append(lamb_coeff) + with torch.no_grad(): + p.add_(-group['lr'] * lamb_coeff * update) + else: + # compression stage, update each momentum locally, then + # communicate based on the compressed_allreduce below + if self.initialize: + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg.mul_(self.state[p]['scaling_coeff']) + grad = None + + # init fused momentum + if len(self.exp_avg_flat) == 0: + momentum_groups = [] + tensor_size = 0 + for group in self.param_groups: + for p in group['params']: + momentum_groups.append(self.state[p]['exp_avg']) + tensor_size += torch.numel(p.data) + corrected_tensor_size = tensor_size + if tensor_size % (self.size * self.divider) != 0: + difference = ((self.size * self.divider) - (tensor_size % (self.size * self.divider))) + corrected_tensor_size += difference + self.dummy_exp_avg[0] = torch.zeros(difference, device=momentum_groups[0].data.device) + momentum_groups.append(self.dummy_exp_avg[0]) + self.corrected_tensor_sizes.append(corrected_tensor_size) + self.server_chunk_sizes.append(corrected_tensor_size // self.size) + + self.exp_avg_flat.append(_flatten_dense_tensors([p.detach().clone() for p in momentum_groups])) + updated_params = _unflatten_dense_tensors(self.exp_avg_flat[0], momentum_groups) + for p, q in zip(momentum_groups, updated_params): + p.data = q.data + + if self.initialize and len(self.worker_errors) == 0: + get_accelerator().empty_cache() + for i in range(len(self.exp_avg_flat)): + self.worker_errors.append( + torch.zeros(self.corrected_tensor_sizes[i], device=self.exp_avg_flat[i].device)) + self.server_errors.append(torch.zeros(self.server_chunk_sizes[i], device=self.exp_avg_flat[i].device)) + get_accelerator().empty_cache() + + if self.lamb_freeze_key: + if self.size > 1: + for i in range(len(self.exp_avg_flat)): + if not self.initialize: + get_accelerator().empty_cache() + self.worker_errors.append( + torch.zeros(self.corrected_tensor_sizes[i], device=self.exp_avg_flat[i].device)) + self.server_errors.append( + torch.zeros(self.server_chunk_sizes[i], device=self.exp_avg_flat[i].device)) + get_accelerator().empty_cache() + if dist.get_rank() == 0: + print("Cupy Buffers Initialized Successfully.") + + self.comm_backend_handle.compressed_allreduce(self.exp_avg_flat[i], self.worker_errors[0], + self.server_errors[0], self.deepspeed.local_rank) + + if dist.get_rank() == 0: + print('Pop out errors', flush=True) + del self.worker_errors[:] + del self.server_errors[:] + else: + self.comm_backend_handle.compressed_allreduce(self.exp_avg_flat[i], self.worker_errors[i], + self.server_errors[i], self.deepspeed.local_rank) + + if self.lamb_freeze_key and self.initialize: + for i, group in enumerate(self.param_groups): + bias_correction = 1 if group['bias_correction'] else 0 + + for j, p in enumerate(group['params']): + state = self.state[p] + exp_avg, exp_avg_sq, exp_avg_sq_fresh = state['exp_avg'], state['exp_avg_sq'], state[ + 'exp_avg_sq_fresh'] + beta1, beta2 = group['betas'] + exp_avg.div_(self.state[p]['scaling_coeff']) + # Because 1-bit compression cannot represent exact zero, it is required to + # provide a momentum mask for those params that have constant exact zeros in their + # momentums, otherwise the compression error would keep accumulating. + # For example, for BERT pre-training seq 128, bert.embeddings.position_embeddings.weight + # always have exact zeros in its momentum for row 129 to 512, because it only + # learns up to seq length 128 while the model supports up to 512 seq length. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py about how + # to add this exp_avg_mask for BERT pre-training.) + if 'exp_avg_mask' in group: + if exp_avg.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=exp_avg.device) + exp_avg.mul_(group['exp_avg_mask']) + + grad_reconstruct = ((exp_avg - exp_avg_last_step[i][j] * beta1) / (1 - beta1)) + exp_avg_sq_fresh.mul_(beta2).addcmul_(1 - beta2, grad_reconstruct, grad_reconstruct) + denom = exp_avg_sq.sqrt() + group['eps'] + update_prelim = exp_avg / denom + + if group['weight_decay'] > 0.0: + update = update_prelim + group['weight_decay'] * p.data + else: + update = update_prelim + + lamb_coeff = 1.0 + update_norm = update.pow(2).sum().sqrt() + denom_real = exp_avg_sq_fresh.sqrt() + group['eps'] + factor = (denom / denom_real).max().item() + if group['weight_decay'] > 0.0: + update_ratio = min(1.0, (update_prelim.pow(2).sum().sqrt() / update_norm).item()) + factor = factor * update_ratio + (1.0 - update_ratio) + if factor > self.factor_max: + factor = self.factor_max + if factor < self.factor_min: + factor = self.factor_min + if factor > state['last_factor'] * (1.0 + self.factor_threshold): + factor = state['last_factor'] * (1.0 + self.factor_threshold) + if factor < state['last_factor'] * (1.0 - self.factor_threshold): + factor = state['last_factor'] * (1.0 - self.factor_threshold) + state['last_factor'] = factor + lamb_coeff = state['lamb_coeff_freeze'] * factor + self.lamb_coeffs.append(lamb_coeff) + with torch.no_grad(): + p.add_(-group['lr'] * lamb_coeff * update) + del exp_avg_last_step[:] + exp_avg_last_step = None + + if not self.initialize: + self.lamb_freeze_key = False + self.initialize = True + print(f"Finished the initialization step at rank {dist.get_rank()}") + return loss + + if self.lamb_freeze_key is False: + if state['step'] >= self.freeze_step: + print('OnebitLamb - starting compressed communication') + self.lamb_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + + return loss + + def load_state_dict(self, state_dict): + """ + Overrides load_state_dict() to add special handling when loading checkpoints + """ + # Because at different stage exp_avg_mask may change (e.g., + # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask + # in checkpoints but always use the one user provided in training script. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) + # Thus here we keep the exp_avg_mask unchanged when loading checkpoint + for i, group in enumerate(self.param_groups): + if 'exp_avg_mask' in group: + state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] + elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: + state_dict['param_groups'][i].pop('exp_avg_mask') + super().load_state_dict(state_dict) + # need to reset the fused momentum since loading states will break the linking + del self.exp_avg_flat[:] + self.dummy_exp_avg.clear() + del self.corrected_tensor_sizes[:] + del self.server_chunk_sizes[:] + if self.state[self.param_groups[0]['params'][0]]['step'] < self.freeze_step: + if dist.get_rank() == 0: + print("Checkpoint loaded and OnebitLamb warmup stage starts/continues.") + if self.lamb_freeze_key is True: + self.lamb_freeze_key = False + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = True + else: + self.deepspeed.enable_backward_allreduce = True + for group in self.param_groups: + for p in group['params']: + self.state[p]['lamb_coeff_freeze'] = 0.0 + self.state[p]['last_factor'] = 1.0 + if 'scaling_coeff' in self.state[p]: + self.state[p].pop('scaling_coeff') + else: + if dist.get_rank() == 0: + print("Checkpoint loaded and OnebitLamb compression stage starts/continues.") + if self.lamb_freeze_key is False: + self.lamb_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + # We reset the compression errors when loading checkpoints for 3 reasons: + # 1) The worker and server error at each GPU are distinct, so in current implementation + # only rank 0's errors are saved in the checkpoint. Thus we have to reset the errors. + # If we want to save them correctly we need O(num_gpu*model_size) memory in order to + # gather all the error, which is a very large memory requirement. It's possible to save + # them in a distributed way, but it will make the checkpoint saving/loading much more complicated. + # 2) Even if we are able to save the compression errors correctly, you need to have the + # exact same number of GPUs in order to load them correctly. + # 3) We verified on BERT pre-training that occasionally resetting the compression error + # at checkpoint loading does not affect the convergence. + # However, please avoid frequent checkpoint loading which could break the error + # compensation mechanism thus affect the convergence. + del self.worker_errors[:] + del self.server_errors[:] + + def get_lamb_coeffs(self): + return self.lamb_coeffs diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py new file mode 100644 index 0000000000000000000000000000000000000000..70282ec41714b202a7105925d4b051cd7571ecea --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/fp16/onebit/zoadam.py @@ -0,0 +1,365 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import types +import torch +import numpy as np +from deepspeed.accelerator import get_accelerator +from deepspeed.utils.torch import required_torch_version +from deepspeed import comm as dist + + +class ZeroOneAdam(torch.optim.Optimizer): + """ + Implements the 0/1 Adam algorithm. Currently GPU-only. + For usage example please see https://www.deepspeed.ai/tutorials/zero-one-adam/ + For technical details please read https://arxiv.org/abs/2202.06009 + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups. + lr (float, optional): learning rate. (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability. (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + var_freeze_step (int, optional): The latest step to update the variance, + using the notation from https://arxiv.org/abs/2202.06009, it denotes the + max{i|i in T_v}. Note that this is different from the freeze step from the + 1-bit Adam. The var_freeze_step is usually the end of the learning rate warmup + and thus does not require tuning. (default: 100000) + var_update_scaler (int, optional): The interval to update the variance. Note that + the update policy for variance follows an exponential rule, where var_update_scaler + denotes the kappa in the 0/1 Adam paper. (default: 16) + local_step_scaler (int, optional): The interval to scale the local steps interval + according to the learning rate policy. (default: 32678) + local_step_clipper (int, optional): The largest interval for local steps with + learning rate policy. This corresponds to the variable H in the 0/1 Adam paper. + (default: 16) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) NOT SUPPORTED in 0/1 Adam! + eps_inside_sqrt (boolean, optional): in the 'update parameters' step, + adds eps to the bias-corrected second moment estimate before + evaluating square root instead of adding it to the square root of + second moment estimate as in the original paper. (default: False) + cuda_aware (boolean, required): Set True if the underlying MPI implementation + supports CUDA-Aware communication. (default: False) + comm_backend_name (string, optional): Set to 'mpi' if needed. (default: 'nccl') + .. _Adam\\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, + params, + deepspeed=None, + lr=1e-3, + bias_correction=True, + betas=(0.9, 0.999), + eps=1e-8, + eps_inside_sqrt=False, + weight_decay=0., + max_grad_norm=0., + var_freeze_step=100000, + var_update_scaler=16, + local_step_scaler=32678, + local_step_clipper=16, + amsgrad=False, + cuda_aware=False, + comm_backend_name='nccl'): + + if amsgrad: + raise RuntimeError('0/1 Adam does not support the AMSGrad variant.') + + defaults = dict(lr=lr, + bias_correction=bias_correction, + betas=betas, + eps=eps, + weight_decay=weight_decay, + max_grad_norm=max_grad_norm) + + super(ZeroOneAdam, self).__init__(params, defaults) + self.eps_mode = 0 if eps_inside_sqrt else 1 + self.deepspeed = deepspeed + self.initialize = False + self.cuda_aware = cuda_aware + self.using_pipeline = False + + self.var_freeze_step = var_freeze_step + self.var_update_scaler = var_update_scaler + self.local_step_scaler = local_step_scaler + self.local_step_clipper = local_step_clipper + self.freeze_key = False + self.reinitial_error_buffer = False + + self.comm_backend_name = comm_backend_name + + assert dist.is_initialized(), "Please initialize the torch distributed backend." + # Empty initializer. Set handle based on the comm backend as follows. + self.comm_backend_handle = None + if self.comm_backend_name == 'nccl': + assert ( + required_torch_version(min_version=1.8) + ), "Please use torch 1.8 or greater to enable NCCL backend in 0/1 Adam. Alternatively, please specify 'mpi' as the 'comm_backend_name' in config file to proceed with the MPI backend" + from deepspeed.runtime.comm.nccl import NcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = NcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'mpi': + from deepspeed.runtime.comm.mpi import MpiBackend + self.comm_backend_handle = MpiBackend(cuda_aware) + elif self.comm_backend_name == 'hccl': + from deepspeed.runtime.comm.hccl import HcclBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = HcclBackend(self.deepspeed.mpu) + elif self.comm_backend_name == 'compressed': + from deepspeed.runtime.comm.compressed import CompressedBackend + self.using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + self.comm_backend_handle = CompressedBackend(self.deepspeed.mpu) + self.size = self.comm_backend_handle.size + + self.divider = int(self.size * 8 / np.gcd(self.size, 8)) + + def step(self, closure=None, grads=None): + """Performs a single optimization step. + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + grads (list of tensors, optional): weight gradient to use for the + optimizer update. If gradients have type torch.half, parameters + are expected to be in type torch.float. (default: None) + output params (list of tensors, optional): A reduced precision copy + of the updated weights written out in addition to the regular + updated weights. Have to be of same type as gradients. (default: None) + scale (float, optional): factor to divide gradient tensor values + by before applying to weights. (default: 1) + """ + loss = None + if closure is not None: + loss = closure() + + if grads is None: + grads_group = [None] * len(self.param_groups) + # backward compatibility + # assuming a list/generator of parameter means single group + elif isinstance(grads, types.GeneratorType): + grads_group = [grads] + elif type(grads[0]) != list: + grads_group = [grads] + else: + grads_group = grads + + for group, grads_this_group in zip(self.param_groups, grads_group): + if grads_this_group is None: + grads_this_group = [None] * len(group['params']) + + bias_correction = 1 if group['bias_correction'] else 0 + + for p, grad in zip(group['params'], grads_this_group): + if p.grad is None and grad is None: + continue + if grad is None: + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('0/1 Adam does not support sparse gradients') + + state = self.state[p] + + # State initialization + if len(state) == 0: + state['step'] = 0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + + if not self.initialize or 'worker_error' not in state.keys(): + # Some scalars to help scale the variance update/local step policies + state['var_interval'] = 1 + state['var_counter'] = 0 + state['local_step_interval'] = 1 + state['local_step_counter'] = 0 + state['lrs'] = 0 + state['tensor_size'] = torch.numel(p.data) + state['corrected_tensor_size'] = state['tensor_size'] + + if state['tensor_size'] % (self.size * self.divider) != 0: + state['corrected_tensor_size'] += ((self.size * self.divider) - (state['tensor_size'] % + (self.size * self.divider))) + state['server_chunk_size'] = state['corrected_tensor_size'] // self.size + get_accelerator().empty_cache() + state['worker_error'] = torch.zeros(state['corrected_tensor_size'], device=p.device) + state['server_error'] = torch.zeros(state['server_chunk_size'], device=p.device) + # Accumulation of momentum, i.e., the u variable in the 0/1 Adam paper + state['momentum_accumulator'] = torch.zeros_like(p.data) + get_accelerator().empty_cache() + # self.freeze_key = True + if not self.initialize and dist.get_rank() == 0: + print("Cupy Buffers Initialized Successfully.") + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + comm_buffer = state['momentum_accumulator'] + beta1, beta2 = group['betas'] + + state['step'] += 1 + + if self.initialize: + if self.freeze_key is False: + if state['step'] % state['var_interval'] == 0: + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + else: + if self.size > 1: + with torch.no_grad(): + grad_onebit = self.comm_backend_handle.compressed_allreduce( + grad, state['worker_error'], state['server_error'], self.deepspeed.local_rank) + if 'exp_avg_mask' in group: + if grad_onebit.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=grad_onebit.device) + grad_onebit.mul_(group['exp_avg_mask']) + exp_avg.mul_(beta1).add_(1 - beta1, grad_onebit) + else: + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + state['lrs'] += group['lr'] + grad = None + + if not self.initialize: + if self.size > 1: + comm_buffer.set_( + self.comm_backend_handle.compressed_allreduce(comm_buffer, state['worker_error'], + state['server_error'], + self.deepspeed.local_rank)) + if 'exp_avg_mask' in group: + if comm_buffer.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=comm_buffer.device) + comm_buffer.mul_(group['exp_avg_mask']) + + if self.initialize: + update = exp_avg / (exp_avg_sq.sqrt() + group['eps']) + if group['weight_decay'] > 0.0: + update += group['weight_decay'] * p.data + with torch.no_grad(): + p.data.add_(-group['lr'] * update) + if self.freeze_key is True: + comm_buffer.add_(-group['lr'] * update) + if state['step'] % state['local_step_interval'] == 0 and self.freeze_key: + with torch.no_grad(): + p.data.add_(-1 * comm_buffer) + comm_buffer.mul_(exp_avg_sq.sqrt() + group['eps']) + if self.size > 1: + comm_buffer.copy_( + self.comm_backend_handle.compressed_allreduce(comm_buffer, state['worker_error'], + state['server_error'], + self.deepspeed.local_rank)) + if 'exp_avg_mask' in group: + if comm_buffer.device != group['exp_avg_mask'].device: + group['exp_avg_mask'] = group['exp_avg_mask'].to(device=comm_buffer.device) + comm_buffer.mul_(group['exp_avg_mask']) + exp_avg.zero_().add_(comm_buffer / state['lrs'], alpha=-1) + p.data.add_(comm_buffer / (exp_avg_sq.sqrt() + group['eps'])) + comm_buffer.zero_() + + state['lrs'] = 0 + + # According to 0/1 Adam theory, a fixed variance would allow more accurate estimation of momentum + # However, in practice, we can also disable the manual freezing of variance, since the interval of + # updating variance will increase exponentially, so that it has negligible effect on the estimation. + if self.freeze_key is False: + if state['step'] % state['var_interval'] == 0: + state['var_counter'] += 1 + if state['var_counter'] == self.var_update_scaler: + state['var_counter'] = 0 + state['var_interval'] *= 2 + if (state['step'] + 1) % state['var_interval'] == 0: + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = True + else: + self.deepspeed.enable_backward_allreduce = True + else: + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + else: + state['local_step_counter'] += 1 + if state['local_step_counter'] == self.local_step_scaler: + state['local_step_counter'] = 0 + state['local_step_interval'] = min(self.local_step_clipper, + state['local_step_interval'] * 2) + + if not self.initialize: + print('Pop out errors', flush=True) + self.freeze_key = False + state.pop('worker_error') + state.pop('server_error') + + if not self.initialize: + self.initialize = True + print(f"Finished the initialization step at rank {dist.get_rank()}") + return loss + + if self.state[self.param_groups[0]['params'][0]]['step'] > self.var_freeze_step: + self.freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + + if self.freeze_key is True and self.reinitial_error_buffer is False: + # We need to reinitialize the error buffers when local step > 1 since + # the errors will be logged for different metrics (gradient vs. accumulated momentum). + for group in self.param_groups: + for p in group['params']: + self.state[p]['worker_error'].zero_() + self.state[p]['server_error'].zero_() + self.reinitial_error_buffer = True + + return loss + + def load_state_dict(self, state_dict): + """ + Overrides load_state_dict() to add special handling when loading checkpoints + """ + # Because at different stage exp_avg_mask may change (e.g., + # BERT pre-training seqlen 128 and 512 ), we don't use the exp_avg_mask + # in checkpoints but always use the one user provided in training script. + # (See example in DeepSpeedExamples/bing_bert/deepspeed_train.py.) + # Thus here we keep the exp_avg_mask unchanged when loading checkpoint + for i, group in enumerate(self.param_groups): + if 'exp_avg_mask' in group: + state_dict['param_groups'][i]['exp_avg_mask'] = group['exp_avg_mask'] + elif 'exp_avg_mask' not in group and 'exp_avg_mask' in state_dict['param_groups'][i]: + state_dict['param_groups'][i].pop('exp_avg_mask') + super().load_state_dict(state_dict) + if self.state[self.param_groups[0]['params'][0]]['step'] < self.var_freeze_step: + self.var_freeze_key = False + if (self.state[self.param_groups[0]['params'][0]]['step'] + + 1) % self.state[self.param_groups[0]['params'][0]]['var_interval'] == 0: + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = True + else: + self.deepspeed.enable_backward_allreduce = True + else: + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + else: + self.var_freeze_key = True + if self.using_pipeline: + self.deepspeed.pipeline_enable_backward_allreduce = False + else: + self.deepspeed.enable_backward_allreduce = False + self.reinitial_error_buffer = False + for group in self.param_groups: + for p in group['params']: + if 'worker_error' in self.state[p]: + self.state[p].pop('worker_error') + if 'server_error' in self.state[p]: + self.state[p].pop('server_error') + if 'momentum_accumulator' in self.state[p]: + self.state[p].pop('momentum_accumulator') diff --git a/lib/python3.12/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py b/lib/python3.12/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..530355f846e2c04dfeaaf1886d79f67aa764908f --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/fp16/unfused_optimizer.py @@ -0,0 +1,428 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/apex +This file is adapted from FP16_Optimizer in NVIDIA/apex +""" + +from deepspeed.moe.utils import split_params_grads_into_shared_and_expert_params +import torch +from torch._utils import _flatten_dense_tensors + +from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer +from deepspeed.runtime.utils import get_global_norm, CheckOverflow, get_weight_norm +from deepspeed.runtime.fp16.loss_scaler import INITIAL_LOSS_SCALE, SCALE_WINDOW, MIN_LOSS_SCALE +from deepspeed.utils import logger +from deepspeed.utils.torch import required_torch_version +from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT +from deepspeed.accelerator import get_accelerator +from deepspeed import comm as dist + + +class FP16_UnfusedOptimizer(DeepSpeedOptimizer): + """ + FP16 Optimizer without weight fusion to support LAMB optimizer + + For usage example please see, TODO: DeepSpeed V2 Tutorial + """ + + def __init__(self, + init_optimizer, + deepspeed=None, + static_loss_scale=1.0, + dynamic_loss_scale=False, + dynamic_loss_args=None, + verbose=True, + mpu=None, + clip_grad=0.0, + fused_lamb_legacy=False): + + self.fused_lamb_legacy = fused_lamb_legacy + self._global_grad_norm = 0. + + if dist.get_rank() == 0: + logger.info(f'Fused Lamb Legacy : {self.fused_lamb_legacy} ') + + if not get_accelerator().is_available(): + raise SystemError("Cannot use fp16 without accelerator.") + self.optimizer = init_optimizer + + # param groups + self.fp16_groups = [] + self.fp32_groups = [] + + # loop to deal with groups + for i, param_group in enumerate(self.optimizer.param_groups): + #fp16 weights that represents the actual model weights + self.fp16_groups.append(param_group['params']) + + #creating a fp32 copy of the weights that will be updated first then + #copied to fp16 weights + fp32_group = [p.clone().float().detach() for p in param_group['params']] + + #in case the internal optimizer needs it + for p in fp32_group: + p.requires_grad = True + + #setting the param groups in the optimizer to point to fp32 + #note these are not the weights used by the model + #the model uses the fp16 version that we added to fp16_group + self.fp32_groups.append(fp32_group) + param_group['params'] = self.fp32_groups[i] + + # we may have a way of fusing dynamic scale. Do not support for now + if dynamic_loss_scale: + self.dynamic_loss_scale = True + self.cur_iter = 0 + self.last_overflow_iter = -1 + self.scale_factor = 2.0 + if dynamic_loss_args is None: + self.cur_scale = 1.0 * 2**16 + self.scale_window = 1000 + self.min_loss_scale = 0.25 + else: + self.cur_scale = dynamic_loss_args[INITIAL_LOSS_SCALE] + self.scale_window = dynamic_loss_args[SCALE_WINDOW] + self.min_loss_scale = dynamic_loss_args[MIN_LOSS_SCALE] + else: + self.dynamic_loss_scale = False + self.cur_iter = 0 + self.cur_scale = static_loss_scale + + self.custom_loss_scaler = False + self.external_loss_scale = None + + self.verbose = verbose + + self.clip_grad = clip_grad + self.norm_type = 2 + + if required_torch_version(max_version=0.4): + self.clip_grad_norm = torch.nn.utils.clip_grad_norm + else: + self.clip_grad_norm = torch.nn.utils.clip_grad_norm_ + + self.mpu = mpu + + self.overflow = False + self.overflow_checker = CheckOverflow(self.fp16_groups, mpu=self.mpu, deepspeed=deepspeed) + + self.initialize_optimizer_states() + + def zero_grad(self, set_to_none=True): + """ + Zero FP16 parameter grads. + """ + # FP32 grad should never exist outside of the step function + # For speed, set model fp16 grad to None by default + for group in self.fp16_groups: + for p in group: + if set_to_none: + p.grad = None + else: + if p.grad is not None: + p.grad.detach_() + p.grad.zero_() + + def step_fused_lamb(self, closure=None): + """ + Not supporting closure. + """ + # First compute norm for all group so we know if there is overflow + grads_groups_flat = [] + grads_groups = [] + norm_groups = [] + expert_norm_groups = [] + for i, group in enumerate(self.fp16_groups): + grads = [ + torch.zeros(p.size(), dtype=p.dtype, device=p.device) if p.grad is None else p.grad for p in group + ] + grads_groups.append(grads) + grads_groups_flat.append(_flatten_dense_tensors(grads)) + grads_for_norm, expert_grads_for_norm = split_params_grads_into_shared_and_expert_params(group) + norm_group_value = 0.0 + if len(grads_for_norm) > 0: + norm_group_value = get_weight_norm(_flatten_dense_tensors(grads_for_norm), mpu=self.mpu) + norm_groups.append(norm_group_value) + expert_norm_group_value = 0.0 + if len(expert_grads_for_norm) > 0: + expert_norm_group_value = get_weight_norm(_flatten_dense_tensors(expert_grads_for_norm), mpu=self.mpu) + expert_norm_groups.append(expert_norm_group_value) + + self.overflow = self.overflow_checker.check_using_norm(norm_groups + expert_norm_groups) + prev_scale = self.cur_scale + + self._update_scale(self.overflow) + if self.overflow: + if self.verbose: + logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " + "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) + return self.overflow + + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + combined_scale = self.unscale_and_clip_grads(self._global_grad_norm, apply_scale=False) + self.optimizer.step(grads=grads_groups, output_params=self.fp16_groups, scale=combined_scale) + + for fp32_group, fp16_group in zip(self.fp32_groups, self.fp16_groups): + for idx, (fp32_param, fp16_param) in enumerate(zip(fp32_group, fp16_group)): + + #remove the fp32 grad + fp32_param.grad = None + + #copy data from fp32 to fp16 + fp16_param.data.copy_(fp32_param.data) + + return self.overflow + + def set_lr(self, lr): + """Set the learning rate.""" + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def get_lr(self): + """Return the current learning rate.""" + return self.optimizer.param_groups[0]["lr"] + + def override_loss_scale(self, loss_scale): + if loss_scale != self.external_loss_scale: + logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}') + self.custom_loss_scaler = True + self.external_loss_scale = loss_scale + + def step(self, closure=None): + """ + Not supporting closure. + """ + + if self.fused_lamb_legacy: + return self.step_fused_lamb() + + self.overflow = self.overflow_checker.check() + prev_scale = self.cur_scale + + self._update_scale(self.overflow) + if self.overflow: + if self.verbose: + logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss " + "scale: {}, reducing to {}".format(prev_scale, self.cur_scale)) + return self.overflow + + norm_groups = [] + for i, group in enumerate(self.fp16_groups): + grads_for_norm, _ = split_params_grads_into_shared_and_expert_params(group) + norm_group_value = 0.0 + if len(grads_for_norm) > 0: + norm_group_value = get_weight_norm(grads_for_norm, mpu=self.mpu) + norm_groups.append(norm_group_value) + + # copying gradients to fp32 to work with fp32 parameters + for fp32_param, fp16_param in zip(self.fp32_groups[i], self.fp16_groups[i]): + if fp16_param.grad is None: + fp32_param.grad = torch.zeros(fp16_param.size(), dtype=fp32_param.dtype, device=fp32_param.device) + else: + fp32_param.grad = fp16_param.grad.to(fp32_param.dtype) + + self._global_grad_norm = get_global_norm(norm_list=norm_groups) + self.unscale_and_clip_grads(self._global_grad_norm) + + self.optimizer.step() + + for fp32_group, fp16_group in zip(self.fp32_groups, self.fp16_groups): + for idx, (fp32_param, fp16_param) in enumerate(zip(fp32_group, fp16_group)): + + #remove the fp32 grad + fp32_param.grad = None + + #copy data from fp32 to fp16 + fp16_param.data.copy_(fp32_param.data) + + return self.overflow + + def unscale_and_clip_grads(self, total_norm, apply_scale=True): + # compute combined scale factor for this group + combined_scale = self.cur_scale + if self.clip_grad > 0.: + # norm is in fact norm*scale + clip = ((total_norm / self.cur_scale) + 1e-6) / self.clip_grad + if clip > 1: + combined_scale = clip * self.cur_scale + + if apply_scale: + for group in self.fp32_groups: + for param in group: + if param.grad is not None: + param.grad.data.mul_(1. / combined_scale) + + return combined_scale + + def backward(self, loss, create_graph=False, retain_graph=False): + """ + :attr:`backward` performs the following steps: + + 1. fp32_loss = loss.float() + 2. scaled_loss = fp32_loss*loss_scale + 3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves + """ + if self.custom_loss_scaler: + scaled_loss = self.external_loss_scale * loss + scaled_loss.backward() + else: + scaled_loss = (loss.float()) * self.cur_scale + scaled_loss.backward(create_graph=create_graph, retain_graph=retain_graph) + + def _update_scale(self, skip): + if self.dynamic_loss_scale: + prev_scale = self.cur_scale + if skip: + self.cur_scale = max(self.cur_scale / self.scale_factor, self.min_loss_scale) + self.last_overflow_iter = self.cur_iter + if self.verbose: + logger.info("Grad overflow on iteration: %s", self.cur_iter) + logger.info(f"Reducing dynamic loss scale from {prev_scale} to {self.cur_scale}") + else: + # Ensure self.scale_window updates since last overflow + stable_interval = (self.cur_iter - self.last_overflow_iter) - 1 + if (stable_interval > 0) and (stable_interval % self.scale_window == 0): + self.cur_scale *= self.scale_factor + if self.verbose: + logger.info(f"No Grad overflow for {self.scale_window} iterations") + logger.info(f"Increasing dynamic loss scale from {prev_scale} to {self.cur_scale}") + else: + if skip: + logger.info("Grad overflow on iteration %s", self.cur_iter) + logger.info("Using static loss scale of %s", self.cur_scale) + self.cur_iter += 1 + return + + # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state" + def _get_state(self): + return self.optimizer.state + + def _set_state(self, value): + self.optimizer.state = value + + state = property(_get_state, _set_state) + + # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups" + # (for example, to adjust the learning rate) + def _get_param_groups(self): + return self.optimizer.param_groups + + def _set_param_groups(self, value): + self.optimizer.param_groups = value + + param_groups = property(_get_param_groups, _set_param_groups) + + # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale" + def _get_loss_scale(self): + if self.custom_loss_scaler: + return self.external_loss_scale + else: + return self.cur_scale + + def _set_loss_scale(self, value): + self.loss_scaler.cur_scale = value + + loss_scale = property(_get_loss_scale, _set_loss_scale) + + def state_dict(self): + """ + Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. + This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict + of the contained Pytorch optimizer. + Example:: + checkpoint = {} + checkpoint['model'] = model.state_dict() + checkpoint['optimizer'] = optimizer.state_dict() + torch.save(checkpoint, "saved.pth") + """ + state_dict = {} + state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale + state_dict['cur_scale'] = self.cur_scale + state_dict['cur_iter'] = self.cur_iter + if state_dict['dynamic_loss_scale']: + state_dict['last_overflow_iter'] = self.last_overflow_iter + state_dict['scale_factor'] = self.scale_factor + state_dict['scale_window'] = self.scale_window + state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + state_dict['fp32_groups'] = self.fp32_groups + return state_dict + + # Refresh fp32 master params from fp16 copies + def refresh_fp32_params(self): + for current_group, saved_group in zip(self.fp32_groups, self.fp16_groups): + for current, saved in zip(current_group, saved_group): + current.data.copy_(saved.data) + + def load_state_dict(self, state_dict, load_optimizer_states=True): + """ + Loads a state_dict created by an earlier call to state_dict(). + If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, + whose parameters in turn came from ``model``, it is expected that the user + will call ``model.load_state_dict()`` before + ``fp16_optimizer_instance.load_state_dict()`` is called. + Example:: + model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) + optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0) + ... + checkpoint = torch.load("saved.pth") + model.load_state_dict(checkpoint['model']) + optimizer.load_state_dict(checkpoint['optimizer']) + """ + # I think it should actually be ok to reload the optimizer before the model. + self.dynamic_loss_scale = state_dict['dynamic_loss_scale'] + self.cur_scale = state_dict['cur_scale'] + self.cur_iter = state_dict['cur_iter'] + if state_dict['dynamic_loss_scale']: + self.last_overflow_iter = state_dict['last_overflow_iter'] + self.scale_factor = state_dict['scale_factor'] + self.scale_window = state_dict['scale_window'] + + if load_optimizer_states: + self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) + # At this point, the optimizer's references to the model's fp32 parameters are up to date. + # The optimizer's hyperparameters and internal buffers are also up to date. + # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still + # out of date. There are two options. + # 1: Refresh the master params from the model's fp16 params. + # This requires less storage but incurs precision loss. + # 2: Save and restore the fp32 master copies separately. + # We choose option 2. + # + # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device + # of their associated parameters, because it's possible those buffers might not exist yet in + # the current optimizer instance. In our case, as long as the current FP16_Optimizer has been + # constructed in the same way as the one whose state_dict we are loading, the same master params + # are guaranteed to exist, so we can just copy_() from the saved master params. + for current_group, saved_group in zip(self.fp32_groups, state_dict['fp32_groups']): + for current, saved in zip(current_group, saved_group): + current.data.copy_(saved.data) + + def __repr__(self): + return repr(self.optimizer) + + def initialize_optimizer_states(self): + for i, group in enumerate(self.fp16_groups): + for param in group: + param.grad = torch.zeros(param.size(), + dtype=param.dtype, + device=get_accelerator().current_device_name()) + + for i, group in enumerate(self.fp32_groups): + for param in group: + param.grad = torch.zeros(param.size(), + dtype=param.dtype, + device=get_accelerator().current_device_name()) + + self.optimizer.step() + + for i, group in enumerate(self.fp16_groups): + for param in group: + param.grad = None + + for i, group in enumerate(self.fp32_groups): + for param in group: + param.grad = None diff --git a/lib/python3.12/site-packages/deepspeed/runtime/hybrid_engine.py b/lib/python3.12/site-packages/deepspeed/runtime/hybrid_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..b6e417fd47645333d1eb9b77715ffdda8a851bbd --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/hybrid_engine.py @@ -0,0 +1,445 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from deepspeed.inference.config import DeepSpeedInferenceConfig +from deepspeed.module_inject.replace_policy import replace_policies +from deepspeed.module_inject.utils import policy_to_ds_container +from .engine import DeepSpeedEngine +from .utils import TLinear, get_inactive_params +from deepspeed.runtime.zero import GatheredParameters +import time +import gc +import math +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from torch import nn +from deepspeed.utils import logger +from deepspeed.module_inject.layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding +from ..ops.transformer.inference.op_binding.workspace import WorkspaceOp + +try: + import transformers + OPTLearnedPositionalEmbedding = transformers.models.opt.modeling_opt.OPTLearnedPositionalEmbedding +except: + OPTLearnedPositionalEmbedding = None + + +class DeepSpeedHybridEngine(DeepSpeedEngine): + r"""DeepSpeed engine for training and inference.""" + inference_mp_group = None + + def __init__(self, args, model, **kwargs): + + super().__init__(args, model, **kwargs) + + # synch seed between all GPUs + _rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name()) + dist.broadcast(_rng_state, 0) + get_accelerator().set_rng_state(_rng_state.cpu()) + + self.Z3_enabled = (self._config.zero_config.stage == 3) + self.gather_all_layers = self._config.hybrid_engine.pin_parameters + + # inference containers / fwds + self._inference_containers = [] + self._orig_modules = [] + self._orig_fwds = [] + self.create_inference_module() + + # Performance stats + self._t_start = None + self._total_latency = 0 + self._iters = 0 + self._training_start_time = None + self._generate_latency = 0 + self._training_latency = 0 + self._total_batch_size = None + self._gather_latency = 0 + + self.is_lora_fused = False + self.workspace = WorkspaceOp() + + def convert_to_linear_transposed(self, model): + + def _replace_linear_layer(r_module, parent_type=None, prev_type=None): + for name, child in r_module.named_children(): + if child.__class__ in [torch.nn.Linear] and \ + (parent_type is torch.nn.ModuleList or prev_type is torch.nn.ModuleList): + setattr(r_module, name, TLinear(child, name)) + else: + _replace_linear_layer(child, type(r_module), prev_type=parent_type) + return r_module + + _replace_linear_layer(model) + + def new_inference_container(self, orig_layer, policy_cls, layer_id): + policy = policy_cls(orig_layer, inference=True) + + if self._config.fp16_enabled: + inference_dtype = torch.float16 + elif self._config.bfloat16_enabled: + inference_dtype = torch.bfloat16 + else: + inference_dtype = torch.float32 + + _container = policy_to_ds_container( + policy=policy, + config=DeepSpeedInferenceConfig( + set_empty_params=True, + dtype=inference_dtype, + max_out_tokens=self._config.hybrid_engine.max_out_tokens, + min_out_tokens=self._config.hybrid_engine.max_out_tokens, + transposed_mode=True, + ), + model_config=self.module.config if hasattr(self.module, 'config') else None, + layer_id=layer_id, + child=orig_layer) + + if self.mpu is not None: + if hasattr(self.mpu, 'get_model_parallel_world_size'): + _container.set_tensor_parallel_config(self.mpu.get_model_parallel_world_size(), + self.mpu.get_model_parallel_group()) + else: + _container.set_tensor_parallel_config(self.mpu.get_tensor_model_parallel_world_size(), + self.mpu.get_tensor_model_parallel_group()) + else: + _container.set_tensor_parallel_config(self._config.hybrid_engine.inference_tp_size, self.mp_group) + _container.initialize_tensors(enable_training=True) + _container.create_ds_model_config() + _container.create_module() + _container.set_params_wo_copy(Z3_enabled=self.Z3_enabled) + return _container + + def populate_all_inference_policies(self): + self.inference_policies = {} + for plcy in replace_policies: + _ = plcy(None) + if isinstance(plcy._orig_layer_class, list): + for orig_layer_class in plcy._orig_layer_class: + self.inference_policies.update({orig_layer_class: (self.new_inference_container, plcy)}) + elif plcy._orig_layer_class is not None: + self.inference_policies.update({plcy._orig_layer_class: (self.new_inference_container, plcy)}) + self.inference_policies.update({ + nn.Linear: (LinearLayer, ), + nn.Embedding: (EmbeddingLayer, ), + nn.LayerNorm: (Normalize, ), + OPTLearnedPositionalEmbedding: (OPTEmbedding, ) + }) + + def _fuse_lora_layer(self, layer_id): + self._inference_containers[layer_id].fuse_lora() + + def fuse_lora_weight(self): + for layer_id in range(len(self.layer_params)): + self._fuse_lora_layer(layer_id) + + def _unfuse_lora_layer(self, layer_id): + self._inference_containers[layer_id].unfuse_lora() + + def unfuse_lora_weight(self): + for layer_id in range(len(self.layer_params)): + self._unfuse_lora_layer(layer_id) + + def unfuse_lora_weight_non_pinned(self): + for layer_id in range(len(self.layer_params)): + non_active_params = get_inactive_params(self.layer_params[layer_id]) + non_active_lora_params = get_inactive_params(self.layer_lora_params[layer_id]) + non_active_params.extend(non_active_lora_params) + + with GatheredParameters(non_active_params): + self._unfuse_lora_layer(layer_id) + + def retake_inference_cache(self): + if self._config.hybrid_engine.release_inference_cache: + retake_success = self.workspace.retake_workspace() + + if not retake_success: + logger.warning("Unable to acquire workspace on first attempt, emptying cache and retrying.") + gc.collect() + get_accelerator().empty_cache() + retake_success = self.workspace.retake_workspace() + + if not retake_success: + raise RuntimeError("Unable to retake inference workspace.") + + def generate(self, *inputs, **kwargs): + if self._total_batch_size is None: + bsz = inputs[0].shape[0] if len(inputs) > 0 else \ + kwargs['input_ids'].shape[0] + self._total_batch_size = bsz * dist.get_world_size() + + self._t0 = time.time() + + if self.Z3_enabled and self.gather_all_layers: + if self._config.hybrid_engine.inference_tp_size > 1: + non_tp_params = [] + for other_layer in self._other_layers: + non_tp_params.extend(list(other_layer.parameters())) + + partition_size = self._config.hybrid_engine.tp_gather_partition_size + + layer_groups = math.ceil(len(self.layer_params) / partition_size) + for lg in range(layer_groups): + non_active_params = [] + non_active_lora_params = [] + for layer_id in range(lg * partition_size, min(len(self.layer_params), (lg + 1) * partition_size), + 1): + non_tp_params.extend(self.layer_params[layer_id][:4]) + non_active_params.extend(get_inactive_params(self.layer_params[layer_id])) + non_active_params.extend(get_inactive_params(self.layer_lora_params[layer_id])) + with GatheredParameters(non_active_params): + for layer_id in range(lg * partition_size, + min(len(self.layer_params), (lg + 1) * partition_size), 1): + if len(self.all_lora_params) > 0: + self._fuse_lora_layer(layer_id) + + if self.mpu is not None: + self._inference_containers[layer_id].apply_tensor_parallelism(self.mp_replace, + reversed_dim=True) + + # TODO(cmikeh2) Evaluate if this can be deferred when release_inference_cache + # is enabled. + gc.collect() + get_accelerator().empty_cache() + + self._gather_latency = time.time() - self._t0 + + input_shape = inputs[0].shape if len(inputs) > 0 else \ + kwargs['input_ids'].shape + output = torch.zeros( + (input_shape[0] * self._config.hybrid_engine.inference_tp_size, ) + input_shape[1:], + dtype=inputs[0].dtype if len(inputs) > 0 else kwargs['input_ids'].dtype, + device=inputs[0].device if len(inputs) > 0 else kwargs['input_ids'].device) + input_cont = inputs[0].contiguous() if len(inputs) > 0 else kwargs['input_ids'].contiguous() + dist.all_gather_into_tensor(output, input_cont, group=self.mp_group) + + if len(inputs) > 0: + inputs = (output, *inputs[1:]) + else: + kwargs['input_ids'] = output + + self.retake_inference_cache() + + non_active_params = get_inactive_params(non_tp_params) + with GatheredParameters(non_active_params): + generate_ret_vals = self._generate(*inputs, **kwargs) + + for layer_id in range(len(self.layer_params)): + self._inference_containers[layer_id].release_memory() + + rank = dist.get_rank(group=self.mp_group) + generate_ret_vals = generate_ret_vals[input_shape[0] * rank:input_shape[0] * (rank + 1)] + + else: + non_active_layers = get_inactive_params(self.all_layers_params) + non_active_lora_params = get_inactive_params(self.all_lora_params) + non_active_layers.extend(non_active_lora_params) + with GatheredParameters(non_active_layers): + self._gather_latency = time.time() - self._t0 + + if len(self.all_lora_params) > 0: + self.fuse_lora_weight() + + self.retake_inference_cache() + generate_ret_vals = self._generate(*inputs, **kwargs) + + if len(self.all_lora_params) > 0: + self.unfuse_lora_weight() + else: + if len(self.all_lora_params) > 0 and (not self.Z3_enabled): + self.fuse_lora_weight() + + self.retake_inference_cache() + generate_ret_vals = self._generate(*inputs, **kwargs) + + if len(self.all_lora_params) > 0: + if (not self.Z3_enabled): + self.unfuse_lora_weight() + else: + self.unfuse_lora_weight_non_pinned() + self.is_lora_fused = False + + if self._config.hybrid_engine.release_inference_cache: + self.workspace.release_workspace() + gc.collect() + get_accelerator().empty_cache() + + self._generate_latency = time.time() - self._t0 - self._gather_latency + + return generate_ret_vals + + def create_inference_containers(self, module, layer_id=0): + for name, child in module.named_children(): + if child.__class__ in self.inference_policies: + if self.inference_policies[child.__class__][0] == self.new_inference_container: + self._inference_containers.append(self.inference_policies[child.__class__][0]( + child, self.inference_policies[child.__class__][-1], layer_id)) + self._orig_modules.append(child) + self._orig_fwds.append(child.forward) + + self.layer_params.append(self._inference_containers[layer_id].get_all_params()) + + self.lora_params.append(self._inference_containers[layer_id].get_lora_params()) + self.layer_lora_params.append([]) + for lora_param in self.lora_params[layer_id]: + self.layer_lora_params[layer_id].extend(lora_param[:-1]) + self.all_lora_params.extend(lora_param[:-1]) + + layer_id += 1 + else: + if self.inference_policies[child.__class__][0] == LinearLayer: + self._other_layers.append(self.inference_policies[child.__class__][0](module=child, + mp_group=None, + skip_partition=True)) + else: + self._other_layers.append(self.inference_policies[child.__class__][0]( + weight=child.weight, bias=child.bias if hasattr(child, 'bias') else None)) + self._orig_modules_others.append(child) + self._orig_fwds_others.append(child.forward) + else: + self.create_inference_containers(child, layer_id=layer_id) + + def create_inference_module(self): + self.layer_params = [] + self.layer_lora_params = [] + self.lora_params = [] + self.all_lora_params = [] + + self._other_layers = [] + self._orig_modules_others = [] + self._orig_fwds_others = [] + + if self._config.hybrid_engine.inference_tp_size > 1: + if self.mpu is None: + global_rank = dist.get_rank() + world_size = dist.get_world_size() + mp_group_id = global_rank // self._config.hybrid_engine.inference_tp_size + num_mp_groups = world_size // self._config.hybrid_engine.inference_tp_size + for mp_group_id in range(num_mp_groups): + ranks = list( + range(mp_group_id * self._config.hybrid_engine.inference_tp_size, \ + (mp_group_id + 1) * self._config.hybrid_engine.inference_tp_size, \ + 1) + ) + mp_group = dist.new_group(ranks) + if global_rank in ranks: + # mp_group is used for broader collective + self.mp_group = mp_group + + # mp_replace is used for container tensor slicing + from deepspeed.module_inject import ReplaceWithTensorSlicing + self.mp_replace = ReplaceWithTensorSlicing( + mp_group=self.mp_group, + mp_size=self._config.hybrid_engine.inference_tp_size, + out_dim=0, + in_dim=1) + + else: + self.mp_group = self.mpu.get_model_parallel_group() if hasattr(self.mpu, 'get_model_parallel_group') else \ + self.mpu.get_tensor_model_parallel_group() + + from deepspeed.module_inject import ReplaceWithTensorSlicing + self.mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group, + mp_size=self._config.hybrid_engine.inference_tp_size, + out_dim=0, + in_dim=1) + else: + self.mp_group = None + self.mp_replace = None + self.populate_all_inference_policies() + self.all_layers_params = list(self.module.parameters()) + self.create_inference_containers(self.module) + + if len(self._inference_containers) > 0: + self._generate = self.module.generate + self.module.generate = self.generate + + self._t0 = time.time() + + def _zero3_forward(self, layer_id): + + def run_forward(*inputs, **kwargs): + non_active_params = get_inactive_params(self.layer_params[layer_id]) + non_active_lora_params = get_inactive_params(self.layer_lora_params[layer_id]) + non_active_params.extend(non_active_lora_params) + + with GatheredParameters(non_active_params): + if len(self.all_lora_params) > 0: + # Use the is_lora_fused flag to prevent multiple fusion in Z3 with non-pinned memory + if not self.is_lora_fused: + self._fuse_lora_layer(layer_id) + # Set the is_lora_fused to true when reaching the last layer + if layer_id == len(self.layer_params) - 1: + self.is_lora_fused = True + return self._inference_containers[layer_id].module.forward(*inputs, **kwargs) + + return run_forward + + def eval(self): + if self._t_start is not None: + latency = time.time() - self._t_start + self._total_latency = self._total_latency + latency + self._iters = self._iters + 1 + if not dist.is_initialized() or dist.get_rank() == 0: + if self._total_batch_size is not None: + cur_samples_p_sec = f'|CurSamplesPerSec={(1 / latency * self._total_batch_size):.2f} ' + avg_samples_p_sec = f'|AvgSamplesPerSec={(1 / (self._total_latency / self._iters) * self._total_batch_size):.2f}' + else: + cur_samples_p_sec = '' + avg_samples_p_sec = '' + others = latency - (self._generate_latency + self._training_latency) + print(f'|E2E latency={(latency):.2f}s ' + \ + f'|Gather latency={self._gather_latency:.2f}s ({(self._gather_latency / latency * 100):.2f}%) ' + f'|Generate time={(self._generate_latency):.2f}s ({(self._generate_latency / latency * 100):.2f}%) ' + \ + f'|Training time={(self._training_latency):.2f}s ({(self._training_latency / latency * 100):.2f}%) ' + \ + f'|Others={others:.2f} ({(others / latency * 100):.2f}%)' + \ + cur_samples_p_sec + \ + avg_samples_p_sec) + self._t_start = time.time() + self._training_latency = 0 + super().eval() + if len(self._inference_containers) > 0: + for i, (orig_module, inference_container) in enumerate(zip(self._orig_modules, + self._inference_containers)): + if self.Z3_enabled and not self.gather_all_layers: + orig_module.forward = self._zero3_forward(i) + else: + orig_module.forward = inference_container.module.forward + + inference_container.transform_for_inference() + + if not self.Z3_enabled or self.gather_all_layers: + for orig_module, inference_layer in zip(self._orig_modules_others, self._other_layers): + orig_module.forward = inference_layer.forward + if self.Z3_enabled: + gc.collect() + get_accelerator().empty_cache() + if self._t_start is None: + self._t_start = time.time() + + def train(self, mode=True): + if mode and len(self._orig_modules) > 0: + for inference_container, orig_module, orig_fwd in zip(self._inference_containers, self._orig_modules, + self._orig_fwds): + inference_container.transform_for_training() + orig_module.forward = orig_fwd + for orig_module, orig_fwd in zip(self._orig_modules_others, self._orig_fwds_others): + orig_module.forward = orig_fwd + super().train(mode) + if mode: + self._training_start_time = time.time() + + def step(self, lr_kwargs=None): + super().step(lr_kwargs=lr_kwargs) + + if len(self._inference_containers) > 0: + if not self.Z3_enabled: + for inference_container in self._inference_containers: + inference_container.reset_params() + + if self._training_start_time is not None: + self._training_latency += (time.time() - self._training_start_time) + self._training_start_time = time.time() diff --git a/lib/python3.12/site-packages/deepspeed/runtime/lr_schedules.py b/lib/python3.12/site-packages/deepspeed/runtime/lr_schedules.py new file mode 100644 index 0000000000000000000000000000000000000000..2ffd0bf9f036f5e5ad49521a641545b50f44d4cb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/lr_schedules.py @@ -0,0 +1,878 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Implementation of learning rate schedules. + +Taken and modified from PyTorch v1.0.1 source +https://github.com/pytorch/pytorch/blob/v1.1.0/torch/optim/lr_scheduler.py +""" + +import argparse +from torch.optim import Optimizer +import math +from deepspeed.utils import logger + +LR_SCHEDULE = 'lr_schedule' +LR_RANGE_TEST = 'LRRangeTest' +ONE_CYCLE = 'OneCycle' +WARMUP_LR = 'WarmupLR' +WARMUP_DECAY_LR = 'WarmupDecayLR' +WARMUP_COSINE_LR = 'WarmupCosineLR' +VALID_LR_SCHEDULES = [LR_RANGE_TEST, ONE_CYCLE, WARMUP_LR, WARMUP_DECAY_LR, WARMUP_COSINE_LR] + +LR_RANGE_TEST_MIN_LR = 'lr_range_test_min_lr' +LR_RANGE_TEST_STEP_RATE = 'lr_range_test_step_rate' +LR_RANGE_TEST_STEP_SIZE = 'lr_range_test_step_size' +LR_RANGE_TEST_STAIRCASE = 'lr_range_test_staircase' + +EDGE_VALUE = 'edge_value' +MID_VALUE = 'mid_value' + +CYCLE_FIRST_STEP_SIZE = 'cycle_first_step_size' +CYCLE_FIRST_STAIR_COUNT = 'cycle_first_stair_count' +CYCLE_SECOND_STEP_SIZE = 'cycle_second_step_size' +CYCLE_SECOND_STAIR_COUNT = 'cycle_second_stair_count' +DECAY_STEP_SIZE = 'decay_step_size' + +CYCLE_MIN_LR = 'cycle_min_lr' +CYCLE_MAX_LR = 'cycle_max_lr' +DECAY_LR_RATE = 'decay_lr_rate' + +CYCLE_MIN_MOM = 'cycle_min_mom' +CYCLE_MAX_MOM = 'cycle_max_mom' +DECAY_MOM_RATE = 'decay_mom_rate' + +WARMUP_MIN_LR = 'warmup_min_lr' +WARMUP_MAX_LR = 'warmup_max_lr' +WARMUP_NUM_STEPS = 'warmup_num_steps' +WARMUP_TYPE = 'warmup_type' +WARMUP_LOG_RATE = 'log' +WARMUP_LINEAR_RATE = 'linear' + +WARMUP_MIN_RATIO = 'warmup_min_ratio' +COS_MIN_RATIO = 'cos_min_ratio' + +TOTAL_NUM_STEPS = 'total_num_steps' + + +def add_tuning_arguments(parser): + group = parser.add_argument_group('Convergence Tuning', 'Convergence tuning configurations') + + # LR scheduler + group.add_argument('--lr_schedule', type=str, default=None, help='LR schedule for training.') + + # Learning rate range test + group.add_argument("--lr_range_test_min_lr", type=float, default=0.001, help='Starting lr value.') + group.add_argument("--lr_range_test_step_rate", type=float, default=1.0, help='scaling rate for LR range test.') + group.add_argument("--lr_range_test_step_size", type=int, default=1000, help='training steps per LR change.') + group.add_argument("--lr_range_test_staircase", + type=bool, + default=False, + help='use staircase scaling for LR range test.') + + # OneCycle schedule + group.add_argument("--cycle_first_step_size", + type=int, + default=1000, + help='size of first step of 1Cycle schedule (training steps).') + group.add_argument("--cycle_first_stair_count", + type=int, + default=-1, + help='first stair count for 1Cycle schedule.') + group.add_argument("--cycle_second_step_size", + type=int, + default=-1, + help='size of second step of 1Cycle schedule (default first_step_size).') + group.add_argument("--cycle_second_stair_count", + type=int, + default=-1, + help='second stair count for 1Cycle schedule.') + group.add_argument("--decay_step_size", + type=int, + default=1000, + help='size of intervals for applying post cycle decay (training steps).') + + # 1Cycle LR + group.add_argument("--cycle_min_lr", type=float, default=0.01, help='1Cycle LR lower bound.') + group.add_argument("--cycle_max_lr", type=float, default=0.1, help='1Cycle LR upper bound.') + group.add_argument("--decay_lr_rate", type=float, default=0.0, help='post cycle LR decay rate.') + + # 1Cycle Momentum + group.add_argument('--cycle_momentum', default=False, action='store_true', help='Enable 1Cycle momentum schedule.') + group.add_argument("--cycle_min_mom", type=float, default=0.8, help='1Cycle momentum lower bound.') + group.add_argument("--cycle_max_mom", type=float, default=0.9, help='1Cycle momentum upper bound.') + group.add_argument("--decay_mom_rate", type=float, default=0.0, help='post cycle momentum decay rate.') + + # Warmup LR + group.add_argument('--warmup_min_lr', type=float, default=0, help='WarmupLR minimum/initial LR value') + group.add_argument('--warmup_max_lr', type=float, default=0.001, help='WarmupLR maximum LR value.') + group.add_argument('--warmup_num_steps', type=int, default=1000, help='WarmupLR step count for LR warmup.') + group.add_argument('--warmup_type', + type=str, + default=WARMUP_LOG_RATE, + help='WarmupLR increasing function during warmup') + + # WarmUP cos LR + group.add_argument("--warmup_min_ratio", type=float, default=0.01, help='Cosine LR lower bound.') + group.add_argument("--cos_min_ratio", type=float, default=0.01, help='Cosine LR lower bound.') + + return parser + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser = add_tuning_arguments(parser) + + lr_sched_args, unknown_args = parser.parse_known_args() + return lr_sched_args, unknown_args + + +def override_lr_range_test_params(args, params): + if hasattr(args, LR_RANGE_TEST_MIN_LR) and args.lr_range_test_min_lr is not None: + params[LR_RANGE_TEST_MIN_LR] = args.lr_range_test_min_lr + + if hasattr(args, LR_RANGE_TEST_STEP_RATE) and args.lr_range_test_step_rate is not None: + params[LR_RANGE_TEST_STEP_RATE] = args.lr_range_test_step_rate + + if hasattr(args, LR_RANGE_TEST_STEP_SIZE) and args.lr_range_test_step_size is not None: + params[LR_RANGE_TEST_STEP_SIZE] = args.lr_range_test_step_size + + if hasattr(args, LR_RANGE_TEST_STAIRCASE) and args.lr_range_test_staircase is not None: + params[LR_RANGE_TEST_STAIRCASE] = args.lr_range_test_staircase + + +def override_1cycle_params(args, params): + if hasattr(args, CYCLE_FIRST_STEP_SIZE) and args.cycle_first_step_size is not None: + params[CYCLE_FIRST_STEP_SIZE] = args.cycle_first_step_size + + if hasattr(args, CYCLE_FIRST_STAIR_COUNT) and args.cycle_first_stair_count is not None: + params[CYCLE_FIRST_STAIR_COUNT] = args.cycle_first_stair_count + + if hasattr(args, CYCLE_SECOND_STEP_SIZE) and args.cycle_second_step_size is not None: + params[CYCLE_SECOND_STEP_SIZE] = args.cycle_second_step_size + + if hasattr(args, CYCLE_SECOND_STAIR_COUNT) and args.cycle_second_stair_count is not None: + params[CYCLE_SECOND_STAIR_COUNT] = args.cycle_second_stair_count + + if hasattr(args, DECAY_STEP_SIZE) and args.decay_step_size is not None: + params[DECAY_STEP_SIZE] = args.decay_step_size + + # 1Cycle LR params + if hasattr(args, CYCLE_MIN_LR) and args.cycle_min_lr is not None: + params[CYCLE_MIN_LR] = args.cycle_min_lr + + if hasattr(args, CYCLE_MAX_LR) and args.cycle_max_lr is not None: + params[CYCLE_MAX_LR] = args.cycle_max_lr + + if hasattr(args, DECAY_LR_RATE) and args.decay_lr_rate is not None: + params[DECAY_LR_RATE] = args.decay_lr_rate + + # 1Cycle MOM params + if hasattr(args, CYCLE_MIN_MOM) and args.cycle_min_mom is not None: + params[CYCLE_MIN_MOM] = args.cycle_min_mom + + if hasattr(args, CYCLE_MAX_MOM) and args.cycle_max_mom is not None: + params[CYCLE_MAX_MOM] = args.cycle_max_mom + + if hasattr(args, DECAY_MOM_RATE) and args.decay_mom_rate is not None: + params[DECAY_MOM_RATE] = args.decay_mom_rate + + +def override_warmupLR_params(args, params): + if hasattr(args, WARMUP_MIN_LR) and args.warmup_min_lr is not None: + params[WARMUP_MIN_LR] = args.warmup_min_lr + + if hasattr(args, WARMUP_MAX_LR) and args.warmup_max_lr is not None: + params[WARMUP_MAX_LR] = args.warmup_max_lr + + if hasattr(args, WARMUP_NUM_STEPS) and args.warmup_num_steps is not None: + params[WARMUP_NUM_STEPS] = args.warmup_num_steps + + if hasattr(args, WARMUP_TYPE) and args.warmup_type is not None: + params[WARMUP_TYPE] = args.warmup_type + + +def override_params(args, params): + # LR range test params + override_lr_range_test_params(args, params) + + # 1Cycle params + override_1cycle_params(args, params) + + # WarmupLR params + override_warmupLR_params(args, params) + + +def get_config_from_args(args): + if not hasattr(args, LR_SCHEDULE) or args.lr_schedule is None: + return None, '--{} not specified on command line'.format(LR_SCHEDULE) + + if not args.lr_schedule in VALID_LR_SCHEDULES: + return None, '{} is not supported LR schedule'.format(args.lr_schedule) + + config = {} + config['type'] = args.lr_schedule + config['params'] = {} + + if args.lr_schedule == LR_RANGE_TEST: + override_lr_range_test_params(args, config['params']) + elif args.lr_schedule == ONE_CYCLE: + override_1cycle_params(args, config['params']) + else: + override_warmupLR_params(args, config['params']) + + return config, None + + +def get_lr_from_config(config): + if not 'type' in config: + return None, 'LR schedule type not defined in config' + + if not 'params' in config: + return None, 'LR schedule params not defined in config' + + lr_schedule = config['type'] + lr_params = config['params'] + + if not lr_schedule in VALID_LR_SCHEDULES: + return None, '{} is not a valid LR schedule'.format(lr_schedule) + + if lr_schedule == LR_RANGE_TEST: + return lr_params[LR_RANGE_TEST_MIN_LR], '' + if lr_schedule == ONE_CYCLE: + return lr_params[CYCLE_MAX_LR], '' + # Warmup LR + return lr_params[WARMUP_MAX_LR], '' + + +def update_lr(param_groups, lrs): + for param_group, lr in zip(param_groups, lrs): + param_group['lr'] = lr + return [group['lr'] for group in param_groups] + + +""" +Only optimizers that are subclass of torch.optim.Optimizer are supported. So check the passed optimizer and wrapped +optimizer to see if requirement is satisfied. +TODO: Looking under the hood to examine the wrapped optimizer is a hack that requires a better long-term fix. +""" + + +def get_torch_optimizer(optimizer): + if isinstance(optimizer, Optimizer): + return optimizer + + if hasattr(optimizer, 'optimizer') and isinstance(optimizer.optimizer, Optimizer): + return optimizer.optimizer + + raise TypeError('{} is not a subclass of torch.optim.Optimizer'.format(type(optimizer).__name__)) + + +class LRRangeTest(object): + """Sets the learning rate of each parameter group according to + learning rate range test (LRRT) policy. The policy increases learning + rate starting from a base value with a constant frequency, as detailed in + the paper `A disciplined approach to neural network hyper-parameters: Part 1 `_ + + LRRT policy is used for finding maximum LR that trains a model without divergence, and can be used to + configure the LR boundaries for Cyclic LR schedules. + + LRRT changes the learning rate after every batch. + `step` should be called after a batch has been used for training. + + Args: + optimizer (Optimizer): Wrapped optimizer. + lr_range_test_min_lr (float or list): Initial learning rate which is the + lower boundary in the range test for each parameter group. + lr_range_test_step_size (int): Interval of training steps to increase learning rate. Default: 2000 + lr_range_test_step_rate (float): Scaling rate for range test. Default: 1.0 + lr_range_test_staircase (bool): Scale in staircase fashion, rather than continuous. Default: False. + last_batch_iteration (int): The index of the last batch. This parameter is used when + resuming a training job. Since `step()` should be invoked after each + batch instead of after each epoch, this number represents the total + number of *batches* computed, not the total number of epochs computed. + When last_batch_iteration=-1, the schedule is started from the beginning. + Default: -1 + + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = LRRangeTest(optimizer) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: + https://arxiv.org/abs/1803.09820 +""" + + def __init__(self, + optimizer: Optimizer, + lr_range_test_min_lr: float = 1e-3, + lr_range_test_step_size: int = 2000, + lr_range_test_step_rate: float = 1.0, + lr_range_test_staircase: bool = False, + last_batch_iteration: int = -1): + + self.optimizer = get_torch_optimizer(optimizer) + + if isinstance(lr_range_test_min_lr, list) or isinstance(lr_range_test_min_lr, tuple): + if len(lr_range_test_min_lr) != len(self.optimizer.param_groups): + raise ValueError("expected {} lr_range_test_min_lr, got {}".format(len(self.optimizer.param_groups), + len(lr_range_test_min_lr))) + self.min_lr = list(lr_range_test_min_lr) + else: + self.min_lr = [lr_range_test_min_lr] * len(self.optimizer.param_groups) + + self.step_size = lr_range_test_step_size + self.step_rate = lr_range_test_step_rate + self.last_batch_iteration = last_batch_iteration + self.staircase = lr_range_test_staircase + self.interval_fn = self._staircase_interval if lr_range_test_staircase else self._continuous_interval + + if last_batch_iteration == -1: + self._last_lr = update_lr(self.optimizer.param_groups, self.min_lr) + + def _staircase_interval(self): + return math.floor(float(self.last_batch_iteration + 1) / self.step_size) + + def _continuous_interval(self): + return float(self.last_batch_iteration + 1) / self.step_size + + def _get_increase(self): + return (1 + self.step_rate * self.interval_fn()) + + def get_lr(self): + lr_increase = self._get_increase() + return [lr_range_test_min_lr * lr_increase for lr_range_test_min_lr in self.min_lr] + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + assert getattr(self, '_last_lr', None) is not None, "need to call step() first" + return self._last_lr + + def step(self, batch_iteration=None): + if batch_iteration is None: + batch_iteration = self.last_batch_iteration + 1 + self.last_batch_iteration = batch_iteration + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def state_dict(self): + return {'last_batch_iteration': self.last_batch_iteration} + + def load_state_dict(self, sd): + self.last_batch_iteration = sd['last_batch_iteration'] + + +class OneCycle(object): + """Sets the learning rate of each parameter group according to + 1Cycle learning rate policy (1CLR). 1CLR is a variation of the + Cyclical Learning Rate (CLR) policy that involves one cycle followed by + decay. The policy simultaneously cycles the learning rate (and momentum) + between two boundaries with a constant frequency, as detailed in + the paper `A disciplined approach to neural network hyper-parameters`_. + + 1CLR policy changes the learning rate after every batch. + `step` should be called after a batch has been used for training. + + This implementation was adapted from the github repo: `PyTorch `_. + + Args: + optimizer (Optimizer): Wrapped optimizer. + cycle_min_lr (float or list): Initial learning rate which is the + lower boundary in the cycle for each parameter group. + cycle_max_lr (float or list): Upper learning rate boundaries in the cycle + for each parameter group. Functionally, + it defines the cycle amplitude (cycle_max_lr - cycle_min_lr). + The lr at any cycle is the sum of cycle_min_lr + and some scaling of the amplitude; therefore + cycle_max_lr may not actually be reached depending on + scaling function. + decay_lr_rate(float): Decay rate for learning rate. Default: 0. + cycle_first_step_size (int): Number of training iterations in the + increasing half of a cycle. Default: 2000 + cycle_second_step_size (int): Number of training iterations in the + decreasing half of a cycle. If cycle_second_step_size is None, + it is set to cycle_first_step_size. Default: None + cycle_first_stair_count(int): Number of stairs in first half of cycle phase. This means + lr/mom are changed in staircase fashion. Default 0, means staircase disabled. + cycle_second_stair_count(int): Number of stairs in second half of cycle phase. This means + lr/mom are changed in staircase fashion. Default 0, means staircase disabled. + decay_step_size (int): Intervals for applying decay in decay phase. Default: 0, means no decay. + cycle_momentum (bool): If ``True``, momentum is cycled inversely + to learning rate between 'cycle_min_mom' and 'cycle_max_mom'. + Default: True + cycle_min_mom (float or list): Initial momentum which is the + lower boundary in the cycle for each parameter group. + Default: 0.8 + cycle_max_mom (float or list): Upper momentum boundaries in the cycle + for each parameter group. Functionally, + it defines the cycle amplitude (cycle_max_mom - cycle_min_mom). + The momentum at any cycle is the difference of cycle_max_mom + and some scaling of the amplitude; therefore + cycle_min_mom may not actually be reached depending on + scaling function. Default: 0.9 + decay_mom_rate (float): Decay rate for momentum. Default: 0. + last_batch_iteration (int): The index of the last batch. This parameter is used when + resuming a training job. Since `step()` should be invoked after each + batch instead of after each epoch, this number represents the total + number of *batches* computed, not the total number of epochs computed. + When last_batch_iteration=-1, the schedule is started from the beginning. + Default: -1 + + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = OneCycle(optimizer, 0.0001, 0.0010) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + + .. _A disciplined approach to neural network hyper-parameters: Part 1 -- learning rate, batch size, momentum, and weight decay: https://arxiv.org/abs/1803.09820 + """ + + def __init__(self, + optimizer, + cycle_min_lr, + cycle_max_lr, + decay_lr_rate=0., + cycle_first_step_size=2000, + cycle_second_step_size=None, + cycle_first_stair_count=0, + cycle_second_stair_count=None, + decay_step_size=0, + cycle_momentum=True, + cycle_min_mom=0.8, + cycle_max_mom=0.9, + decay_mom_rate=0., + last_batch_iteration=-1): + + self.optimizer = get_torch_optimizer(optimizer) + + # Initialize cycle shape + self._initialize_cycle(cycle_first_step_size, cycle_second_step_size, cycle_first_stair_count, + cycle_second_stair_count, decay_step_size) + + # Initialize cycle lr + self._initialize_lr(self.optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate, last_batch_iteration) + + # Initialize cyclic momentum + self.cycle_momentum = cycle_momentum + if cycle_momentum: + self._initialize_momentum(self.optimizer, cycle_min_mom, cycle_max_mom, decay_mom_rate, + last_batch_iteration) + # Initialize batch iteration tracker + self.last_batch_iteration = last_batch_iteration + + # Configure cycle shape + + def _initialize_cycle(self, cycle_first_step_size, cycle_second_step_size, cycle_first_stair_count, + cycle_second_stair_count, decay_step_size): + cycle_first_step_size = float(cycle_first_step_size) + cycle_second_step_size = float( + cycle_second_step_size) if cycle_second_step_size is not None else cycle_first_step_size + + self.total_size = cycle_first_step_size + cycle_second_step_size + self.step_ratio = cycle_first_step_size / self.total_size + self.first_stair_count = cycle_first_stair_count + self.second_stair_count = cycle_first_stair_count if cycle_second_stair_count is None else cycle_second_stair_count + self.decay_step_size = decay_step_size + + if math.isclose(self.decay_step_size, 0): + self.skip_lr_decay = True + self.skip_mom_decay = True + else: + self.skip_lr_decay = False + self.skip_mom_decay = False + + # Configure lr schedule + def _initialize_lr(self, optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate, last_batch_iteration): + self.min_lrs = [cycle_min_lr] * len(optimizer.param_groups) + if last_batch_iteration == -1: + for lr, group in zip(self.min_lrs, optimizer.param_groups): + group['lr'] = lr + + self.max_lrs = [cycle_max_lr] * len(optimizer.param_groups) + self.decay_lr_rate = decay_lr_rate + + if math.isclose(self.decay_lr_rate, 0): + self.skip_lr_decay = True + + # Configure momentum schedule + def _initialize_momentum(self, optimizer, cycle_min_mom, cycle_max_mom, decay_mom_rate, last_batch_iteration): + if 'betas' not in optimizer.defaults: + optimizer_name = type(optimizer).__name__ + logger.warning( + f"cycle_momentum is disabled because optimizer {optimizer_name} does not support momentum, no betas attribute in defaults" + ) + self.cycle_momentum = False + return + + self.decay_mom_rate = decay_mom_rate + self.min_moms = [(cycle_min_mom, 0.99)] * len(optimizer.param_groups) + self.max_moms = [(cycle_max_mom, 0.99)] * len(optimizer.param_groups) + + if last_batch_iteration == -1: + for momentum, group in zip(self.min_moms, optimizer.param_groups): + group['betas'] = momentum + + if math.isclose(self.decay_mom_rate, 0): + self.skip_mom_decay = True + + def _get_scale_factor(self): + batch_iteration = (self.last_batch_iteration + 1) + cycle = math.floor(1 + batch_iteration / self.total_size) + x = 1. + batch_iteration / self.total_size - cycle + if x <= self.step_ratio: + scale_factor = x / self.step_ratio + else: + scale_factor = (x - 1) / (self.step_ratio - 1) + + return scale_factor + + def _get_cycle_mom(self): + scale_factor = self._get_scale_factor() + momentums = [] + for base_betas, max_betas in zip(self.min_moms, self.max_moms): + cycle_min_mom = base_betas[0] + cycle_max_mom = max_betas[0] + base_height = (cycle_max_mom - cycle_min_mom) * scale_factor + momentum = cycle_max_mom - base_height + momentums.append((momentum, base_betas[1])) + return momentums + + def _get_cycle_lr(self): + scale_factor = self._get_scale_factor() + lrs = [] + for cycle_min_lr, cycle_max_lr in zip(self.min_lrs, self.max_lrs): + base_height = (cycle_max_lr - cycle_min_lr) * scale_factor + lr = cycle_min_lr + base_height + lrs.append(lr) + + return lrs + + def _get_decay_mom(self, decay_batch_iteration): + if self.skip_mom_decay: + return self.max_moms + + decay_interval = decay_batch_iteration / self.decay_step_size + mom_decay_factor = (1 + self.decay_mom_rate * decay_interval) + momentums = [(beta0 * mom_decay_factor, beta1) for beta0, beta1 in self.max_moms] + + return momentums + + def _get_decay_lr(self, decay_batch_iteration): + """Calculates the learning rate at batch index. This function is used + after the cycle completes and post cycle decaying of lr/mom is enabled. + This function treats `self.last_batch_iteration` as the last batch index. + """ + if self.skip_lr_decay: + return self.min_lrs + + decay_interval = decay_batch_iteration / self.decay_step_size + lr_decay_factor = (1 + self.decay_lr_rate * decay_interval) + lrs = [cycle_min_lr / lr_decay_factor for cycle_min_lr in self.min_lrs] + + return lrs + + def get_lr(self): + """Calculates the learning rate at batch index. This function treats + `self.last_batch_iteration` as the last batch index. + """ + if self.last_batch_iteration < self.total_size: + return self._get_cycle_lr() + return self._get_decay_lr(self.last_batch_iteration - self.total_size + 1) + + def get_mom(self): + """Calculates the momentum at batch index. This function treats + `self.last_batch_iteration` as the last batch index. + """ + if not self.cycle_momentum: + return None + + if self.last_batch_iteration < self.total_size: + return self._get_cycle_mom() + return self._get_decay_mom(self.last_batch_iteration - self.total_size + 1) + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + assert getattr(self, '_last_lr', None) is not None, "need to call step() first" + return self._last_lr + + def step(self, batch_iteration=None): + """ Updates the optimizer with the learning rate for the last batch index. + `self.last_batch_iteration` is treated as the last batch index. + + If self.cycle_momentum is true, also updates optimizer momentum. + """ + if batch_iteration is None: + batch_iteration = self.last_batch_iteration + 1 + + self.last_batch_iteration = batch_iteration + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + if self.cycle_momentum: + momentums = self.get_mom() + for param_group, momentum in zip(self.optimizer.param_groups, momentums): + param_group['betas'] = momentum + + def state_dict(self): + return {'last_batch_iteration': self.last_batch_iteration} + + def load_state_dict(self, sd): + self.last_batch_iteration = sd['last_batch_iteration'] + + +class WarmupLR(object): + """Increase the learning rate of each parameter group from min lr to max lr + over warmup_num_steps steps, and then fix at max lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + warmup_min_lr (float or list): minimum learning rate. Default: 0 + warmup_max_lr (float or list): maximum learning rate. Default: 0.001 + warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000 + warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log + last_batch_iteration (int): The index of the last batch. Default: -1. + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = WarmupLR(optimizer) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + """ + + def __init__(self, + optimizer: Optimizer, + warmup_min_lr: float = 0.0, + warmup_max_lr: float = 0.001, + warmup_num_steps: int = 1000, + warmup_type: str = WARMUP_LOG_RATE, + last_batch_iteration: int = -1): + + self.optimizer = get_torch_optimizer(optimizer) + + self.min_lrs = self._format_param(self.optimizer, warmup_min_lr, "min_lr") + self.max_lrs = self._format_param(self.optimizer, warmup_max_lr, "max_lr") + self.delta_lrs = [big - small for big, small in zip(self.max_lrs, self.min_lrs)] + self.warmup_num_steps = max(2, warmup_num_steps) + # Currently only support linear and log function + if warmup_type not in {WARMUP_LOG_RATE, WARMUP_LINEAR_RATE}: + logger.warning(f"Using unknown warmup_type: {warmup_type}. The increasing function " + f"is set to default (log)") + warmup_type = WARMUP_LOG_RATE + self.warmup_type = warmup_type + self.inverse_log_warm_up = 1.0 / math.log(self.warmup_num_steps) + self.last_batch_iteration = last_batch_iteration + # Initialize lr in optimizer + if last_batch_iteration == -1: + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def get_lr(self): + if self.last_batch_iteration < 0: + logger.warning("Attempting to get learning rate from scheduler before it has started") + return self.min_lrs + gamma = self._get_gamma() + return [min_lr + (delta_lr * gamma) for min_lr, delta_lr in zip(self.min_lrs, self.delta_lrs)] + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + assert getattr(self, '_last_lr', None) is not None, "need to call step() first" + return self._last_lr + + def step(self, last_batch_iteration=None): + if last_batch_iteration is None: + last_batch_iteration = self.last_batch_iteration + 1 + self.last_batch_iteration = last_batch_iteration + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def state_dict(self): + return {'last_batch_iteration': self.last_batch_iteration} + + def load_state_dict(self, sd): + self.last_batch_iteration = sd['last_batch_iteration'] + + def _get_gamma(self): + if self.last_batch_iteration < self.warmup_num_steps: + if self.warmup_type == WARMUP_LOG_RATE: + return self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) + elif self.warmup_type == WARMUP_LINEAR_RATE: + return self.last_batch_iteration / self.warmup_num_steps + return 1.0 + + def _format_param(self, optimizer, param_value, param_name): + if isinstance(param_value, list) or isinstance(param_value, tuple): + if len(param_value) != len(optimizer.param_groups): + raise ValueError("expected {} value for {}, got {}".format(len(optimizer.param_groups), param_name, + FileNotFoundError(param_value))) + return list(param_value) + return [param_value] * len(optimizer.param_groups) + + +class WarmupDecayLR(WarmupLR): + """Increase the learning rate of each parameter group from min lr to max lr + over warmup_num_steps steps, and then decay at linear rate over the remaining training steps. + + Args: + optimizer (Optimizer): Wrapped optimizer. + total_num_steps (int): total number of training steps + warmup_min_lr (float or list): minimum learning rate. Default: 0 + warmup_max_lr (float or list): maximum learning rate. Default: 0.001 + warmup_num_steps (int): number of steps to warm up from min_lr to max_lr. Default: 1000 + warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log + last_batch_iteration (int): The index of the last batch. Default: -1. + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = WarmupDecayLR(optimizer, 1000000) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + """ + + def __init__(self, + optimizer: Optimizer, + total_num_steps: int, + warmup_min_lr: float = 0.0, + warmup_max_lr: float = 0.001, + warmup_num_steps: int = 1000, + warmup_type: str = WARMUP_LOG_RATE, + last_batch_iteration: int = -1): + + self.total_num_steps = total_num_steps + super(WarmupDecayLR, self).__init__(optimizer, warmup_min_lr, warmup_max_lr, warmup_num_steps, warmup_type, + last_batch_iteration) + if self.total_num_steps < self.warmup_num_steps: + logger.warning('total_num_steps {} is less than warmup_num_steps {}'.format( + total_num_steps, warmup_num_steps)) + + def _get_gamma(self): + if self.last_batch_iteration < self.warmup_num_steps: + if self.warmup_type == WARMUP_LOG_RATE: + return self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) + elif self.warmup_type == WARMUP_LINEAR_RATE: + return self.last_batch_iteration / self.warmup_num_steps + return max( + 0.0, + float(self.total_num_steps - self.last_batch_iteration) / + float(max(1.0, self.total_num_steps - self.warmup_num_steps))) + + +class WarmupCosineLR(object): + """Increase the learning rate of each parameter group from min lr ratio to max lr ratio + over warmup_num_steps steps, and then decay at cosine rate over the remaining training steps to min cosine ratio. + + Args: + optimizer (Optimizer): Wrapped optimizer. + total_num_steps (int): total number of training steps + warmup_min_ratio (float or list): warmup start learning rate ratio. Default: 0 + warmup_num_steps (int): number of steps to warm up from warmup_min_ratio to 1.0. Default: 1000 + warmup_type {‘log’, ‘linear’}: increasing function from min_lr to max_lr during warmup. Default: log + cos_min_ratio (float): cosine end learning rate ratio. Default: 0.0001 + last_batch_iteration (int): The index of the last batch. Default: -1. + Example: + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = WarmupCosineLR(optimizer, 1000000) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + """ + + def __init__(self, + optimizer: Optimizer, + total_num_steps: int, + warmup_min_ratio: float = 0.0, + warmup_num_steps: int = 1000, + cos_min_ratio: float = 0.0001, + warmup_type: str = WARMUP_LOG_RATE, + last_batch_iteration: int = -1): + + self.optimizer = get_torch_optimizer(optimizer) + + self.total_num_steps = total_num_steps + self.last_batch_iteration = last_batch_iteration + self.cos_min_ratio = cos_min_ratio + + self.warmup_type = warmup_type + self.warmup_min_ratio = warmup_min_ratio + self.warmup_num_steps = max(2, warmup_num_steps) + self.inverse_log_warm_up = 1.0 / math.log(self.warmup_num_steps) + + if self.total_num_steps < self.warmup_num_steps: + logger.warning('total_num_steps {} is less than warmup_num_steps {}'.format( + total_num_steps, warmup_num_steps)) + self.org_lrs = [group['lr'] for group in self.optimizer.param_groups] + + # Initialize lrs in optimizer groups + if last_batch_iteration == -1: + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def get_lr_ratio(self): + if self.last_batch_iteration < 0: + logger.warning("Attempting to get learning rate from scheduler before it has started") + return [0.0] + + if self.last_batch_iteration < self.warmup_num_steps: + if self.warmup_type == WARMUP_LOG_RATE: + ratio = self.inverse_log_warm_up * math.log(self.last_batch_iteration + 1) + elif self.warmup_type == WARMUP_LINEAR_RATE: + ratio = self.last_batch_iteration / self.warmup_num_steps + ratio_delta = 1. - self.warmup_min_ratio + ratio = self.warmup_min_ratio + ratio * ratio_delta + return ratio + + real_last_step = self.last_batch_iteration - self.warmup_num_steps + 1 + real_total_steps = self.total_num_steps - self.warmup_num_steps + ratio_delta = 1. - self.cos_min_ratio + ratio = (1 + math.cos(math.pi * real_last_step / real_total_steps)) / 2 + ratio = max(0.0, self.cos_min_ratio + ratio_delta * ratio) + return ratio + + def step(self, last_batch_iteration=None): + if last_batch_iteration is None: + last_batch_iteration = self.last_batch_iteration + 1 + self.last_batch_iteration = last_batch_iteration + self._last_lr = update_lr(self.optimizer.param_groups, self.get_lr()) + + def get_lr(self): + if self.last_batch_iteration < 0: + logger.warning("Attempting to get learning rate from scheduler before it has started") + return [0.0] + lr_ratio = self.get_lr_ratio() + return [org_lr * lr_ratio for org_lr in self.org_lrs] + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + assert getattr(self, '_last_lr', None) is not None, "need to call step() first" + return self._last_lr + + def state_dict(self): + return {'last_batch_iteration': self.last_batch_iteration} + + def load_state_dict(self, sd): + self.last_batch_iteration = sd['last_batch_iteration'] + + def _format_param(self, optimizer, param_value, param_name): + if isinstance(param_value, list) or isinstance(param_value, tuple): + if len(param_value) != len(optimizer.param_groups): + raise ValueError("expected {} value for {}, got {}".format(len(optimizer.param_groups), param_name, + FileNotFoundError(param_value))) + return list(param_value) + return [param_value] * len(optimizer.param_groups) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py b/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py new file mode 100644 index 0000000000000000000000000000000000000000..3068247796ef30190295a890aa2d74b671cf651a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py @@ -0,0 +1,1426 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from types import MethodType +from collections import OrderedDict +from functools import reduce +from operator import mul + +import torch +from deepspeed import comm as dist + +from deepspeed.utils import logger +from deepspeed.utils.timer import ThroughputTimer +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.bf16_optimizer import BF16_Optimizer + +from ..engine import DeepSpeedEngine, MEMORY_OPT_ALLREDUCE_SIZE +from deepspeed.utils.timer import FORWARD_MICRO_TIMER, FORWARD_GLOBAL_TIMER, BACKWARD_MICRO_TIMER, \ + BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_INNER_GLOBAL_TIMER, \ + BACKWARD_REDUCE_MICRO_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, \ + STEP_MICRO_TIMER, STEP_GLOBAL_TIMER + +from ..utils import PartitionedTensor +from ..dataloader import RepeatingLoader +from ..zero.config import ZeroStageEnum +from ..activation_checkpointing import checkpointing as ds_checkpointing + +from .module import PipelineModule, PipelineError +from . import p2p +from . import schedule + +TARGET_ID = -2 +LOG_STAGE = -2 +DATA_PARALLEL_ID = -2 + +BATCH_INPUT_TIMER = 'batch_input' +TRAIN_BATCH_TIMER = 'train_batch' +PIPE_SEND_OUTPUT_TIMER = 'pipe_send_output' +PIPE_SEND_GRAD_TIMER = 'pipe_send_grad' +PIPE_RECV_INPUT_TIMER = 'pipe_recv_input' +PIPE_RECV_GRAD_TIMER = 'pipe_recv_grad' + +# The buffer size to store the meta data for each tensor. +TENSOR_META_SIZE = 256 + + +def is_even(number): + return number % 2 == 0 + + +mem_alloced = 0 +mem_cached = 0 + + +def _tensor_bytes(tensor): + return tensor.numel() * tensor.element_size() + + +class PipelineEngine(DeepSpeedEngine): + """ A training engine hybrid pipeline, data, and model parallel training. + + This engine is created by ``deepspeed.initialize()`` when a :class:`PipelineModule` + is provided. + """ + ID_TO_DTYPE = [ + torch.float32, torch.float64, torch.complex64, torch.complex128, torch.float16, torch.bfloat16, torch.uint8, + torch.int8, torch.int16, torch.int32, torch.int64, torch.bool + ] + DTYPE_TO_ID = {dtype: id_ for id_, dtype in enumerate(ID_TO_DTYPE)} + + def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs): + super().__init__(*super_args, **super_kwargs) + assert isinstance(self.module, PipelineModule), "model must base PipelineModule" + + assert self.zero_optimization_stage( + ) < ZeroStageEnum.gradients, "ZeRO-2 and ZeRO-3 are incompatible with pipeline parallelism" + + # We schedule the all-reduces, so disable it in super().backward() + self.enable_backward_allreduce = False + self.has_bool_tensors = has_bool_tensors + self.eval_return_logits = False + self.outputs = None + # BF16 Optimizer is hardcoded for fp32 gradient accumulation + self.using_bf16_optimizer = type(self.optimizer) == BF16_Optimizer + + # used to disable the pipeline all-reduce when used with 1-bit Adam/1-bit LAMB + self.pipeline_enable_backward_allreduce = True + + if self.elasticity_enabled(): + if not self.is_elastic_model_parallel_supported(): + assert not self.elasticity_enabled(), "Elasticity is not currently supported" \ + " with pipeline parallelism." + + # pipeline step for logging + self.log_batch_step_id = -1 + + self.micro_batch_size = self.train_micro_batch_size_per_gpu() + self.micro_batches = self.gradient_accumulation_steps() + + # Set Grid and Communication Groups + self.grid = self.module._grid + if self.grid.get_global_rank() == 0: + logger.info(f'CONFIG: micro_batches={self.micro_batches} ' + f'micro_batch_size={self.micro_batch_size}') + + self.global_rank = self.grid.get_global_rank() + + assert self.dp_world_size == self.grid.data_parallel_size + assert self.train_batch_size() == \ + self.micro_batch_size * self.micro_batches * self.grid.data_parallel_size + + # Set Stage Inf + self.num_stages = self.grid.pipe_parallel_size + self.stage_id = self.grid.get_stage_id() + self.prev_stage = self.stage_id - 1 + self.next_stage = self.stage_id + 1 + + self.data_iterator = None + self.batch_fn = None + + self._force_grad_boundary = False + + self.batch_timer = ThroughputTimer(self._config.timers_config, + batch_size=self.train_batch_size(), + logging_fn=self.tput_log, + monitor_memory=False, + steps_per_output=self.steps_per_print()) + + # PipelineEngine needs to handle data loading specially due to only the first + # and last stages loading inputs/labels. We construct a sampler that uses + if self.training_data: + self._build_data_iter(self.training_data) + + self.is_pipe_parallel = self.grid.pipe_parallel_size > 1 + self.is_data_parallel = self.grid.data_parallel_size > 1 + self.is_model_parallel = self.grid.model_parallel_size > 1 + + # Partition input/output buffers + # XXX temporarily disable while I revert some partition hacks. + assert isinstance(self._config.pipeline['pipe_partitioned'], bool) + assert isinstance(self._config.pipeline['grad_partitioned'], bool) + self.is_pipe_partitioned = self.is_model_parallel and self._config.pipeline['pipe_partitioned'] + self.is_grad_partitioned = self.is_model_parallel and self._config.pipeline['grad_partitioned'] + logger.info(f'is_pipe_partitioned= {self.is_pipe_partitioned} ' + f'is_grad_partitioned= {self.is_grad_partitioned}') + + model_parameters = filter(lambda p: p.requires_grad, self.module.parameters()) + num_params = sum([p.numel() for p in model_parameters]) + unique_params = num_params + # Subtract tied parameters if we don't own them + if self.module.tied_comms: + tied_params = 0 + for key, d in self.module.tied_comms.items(): + if self.global_rank != min(d['ranks']): + tied_params += sum(p.numel() for p in d['module'].parameters()) + unique_params -= tied_params + params_tensor = torch.LongTensor(data=[num_params, unique_params]).to(self.device) + dist.all_reduce(params_tensor, group=self.grid.get_model_parallel_group()) + params_tensor = params_tensor.tolist() + total_params = params_tensor[0] + unique_params = params_tensor[1] + if self.grid.data_parallel_id == 0: + logger.info(f'RANK={self.global_rank} ' + f'STAGE={self.stage_id} ' + f'LAYERS={self.module._local_stop - self.module._local_start} ' + f'[{self.module._local_start}, {self.module._local_stop}) ' + f'STAGE_PARAMS={num_params} ({num_params/1e6:0.3f}M) ' + f'TOTAL_PARAMS={total_params} ({total_params/1e6:0.3f}M) ' + f'UNIQUE_PARAMS={unique_params} ({unique_params/1e6:0.3f}M)') + + #initialize peer-2-peer communication and allreduce groups + if self.is_pipe_parallel: + p2p.init_process_groups(self.grid) + + # Pipeline buffers + self.num_pipe_buffers = 0 + self.pipe_buffers = { + 'inputs': [], # batch input and received activations + 'labels': [], # labels from batch input + 'outputs': [], # activations + 'output_tensors': [], # tensor object to preserve backward graph + } + self.pipe_recv_buf = None + self.grad_layer = None + self._grad_layer_buf = [] + + self.meta_buffer = None + + self.first_output_send = True + self.first_gradient_send = True + self.pipe_partition_input_meta_cache = None + self.pipe_partition_output_meta_cache = None + self.pipe_partition_grad_meta_cache = None + self.grad_partition_grad_layer_meta_cache = None + + #stores the loss for the current micro batch being processed + self.loss = torch.tensor(0.0).to(self.device) + + #stores the loss for the entire batch + self.total_loss = None + self.total_additional_losses = None + self.agg_loss = torch.tensor(0.0, requires_grad=False).to(self.device) + self.dp_group_loss = torch.tensor(0.0, requires_grad=False).to(self.device) + + # stores aggregated-DP train final loss and aggregated-DP additional losses, if any + # additional losses are stored as dict: {loss-name: agg-loss} + self.agg_train_loss = None + self.agg_additional_losses = None + + if self._config.pipeline['activation_checkpoint_interval'] > 0: + self.module.activation_checkpoint_interval = self._config.pipeline['activation_checkpoint_interval'] + # set use_reentrant default to True. + if self._config.pipeline.get('use_reentrant') is None: + self._config.pipeline['use_reentrant'] = True + if self._config.pipeline['use_reentrant'] is False: + # set activation_checkpoint_func to non_reentrant_checkpoint func. + self.module.activation_checkpoint_func = ds_checkpointing.non_reentrant_checkpoint + if self.grid.get_global_rank() == 0: + logger.info(f'CONFIG: activation_checkpoint_func=non_reentrant_checkpoint') + if self.module.activation_checkpoint_interval > 0: + self.module._precompute_checkpointable_values() + + self.module.checkpoint_parallel_write_pipeline = self._config.checkpoint_parallel_write_pipeline + + if self.is_last_stage(): + self.loss_model = self.module.loss_fn + + self.has_attention_mask = self.module.__class__.__name__ == 'GPT2ModelPipe' + # Initialize pipeline communicators. Just send a 0. + if is_even(self.stage_id): + if not self.is_last_stage(): + p2p.send(self.loss, self.next_stage) + if not self.is_first_stage(): + p2p.recv(self.loss, self.prev_stage) + else: + if not self.is_first_stage(): + p2p.recv(self.loss, self.prev_stage) + if not self.is_last_stage(): + p2p.send(self.loss, self.next_stage) + + # XXX look into timer reporting timing + # Initialize some timers because of early weirdness. + if self.wall_clock_breakdown(): + self.timers(FORWARD_MICRO_TIMER).start() + self.timers(FORWARD_MICRO_TIMER).stop() + self.timers(BACKWARD_MICRO_TIMER).start() + self.timers(BACKWARD_MICRO_TIMER).stop() + self.timers(BACKWARD_INNER_MICRO_TIMER).start() + self.timers(BACKWARD_INNER_MICRO_TIMER).stop() + self.timers(BACKWARD_REDUCE_MICRO_TIMER).start() + self.timers(BACKWARD_REDUCE_MICRO_TIMER).stop() + self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).start() + self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).stop() + self.timers(STEP_MICRO_TIMER).start() + self.timers(STEP_MICRO_TIMER).stop() + + self.dynamic_shape = self.module.dynamic_shape + + def set_has_attention_mask(self, value): + assert isinstance(value, bool) + self.has_attention_mask = value + + def _build_data_iter(self, dataset): + sampler = torch.utils.data.distributed.DistributedSampler(dataset, + num_replicas=self.dp_world_size, + rank=self.mpu.get_data_parallel_rank(), + shuffle=False) + # Build a loader and make it repeating. + pipe_dataloader = self.deepspeed_io(dataset, data_sampler=sampler) + pipe_dataloader = RepeatingLoader(pipe_dataloader) + self.set_dataloader(pipe_dataloader) + + def _exec_reduce_tied_grads(self): + # We need to run this first to write to self.averaged_gradients; + # since this class turns `enable_backward_allreduce` off, + # `self.overlapping_partition_gradients_reduce_epilogue()` defined in the DeepSpeedEngine + # never actually runs. I suspect this is because of efficiency problems; get_flat_partition in + # stage2.py might do something expensive; someone will have to look into that later. But + # in the meantime, this fixes ZeRO2 + Pipelining enough to run a demo. Further profiling + # needed to decide if it actually breaks everything. + # (see https://github.com/EleutherAI/gpt-neox/issues/62#issuecomment-761471944) + if self.zero_optimization_partition_gradients(): + self.optimizer.overlapping_partition_gradients_reduce_epilogue() + + weight_group_list = self.module.get_tied_weights_and_groups() + for weight, group in weight_group_list: + grad = weight._hp_grad if self.using_bf16_optimizer else weight.grad + if grad is not None: + dist.all_reduce(grad, group=group) + + def _exec_reduce_grads(self): + self._force_grad_boundary = True + if self.pipeline_enable_backward_allreduce: + if self.using_bf16_optimizer: + # PP+BF16 work for ZeRO Stage 1 + self._bf16_reduce_grads() + else: + self.allreduce_gradients(bucket_size=MEMORY_OPT_ALLREDUCE_SIZE) + self._force_grad_boundary = False + + def _bf16_reduce_grads(self): + self.buffered_allreduce_fallback(grads=None, elements_per_buffer=MEMORY_OPT_ALLREDUCE_SIZE) + + def _reserve_pipe_buffers(self, num_buffers): + """Ensure that each pipeline buffer has at least ``num_buffers`` slots. + + This method only reserves slots and does not allocate tensors. + + Args: + num_buffers (int): The number of buffers to reserve. + """ + if self.num_pipe_buffers >= num_buffers: + return + + num_added = num_buffers - self.num_pipe_buffers + for key in self.pipe_buffers: + self.pipe_buffers[key].extend([None] * num_added) + self.num_pipe_buffers = num_buffers + + def reset_activation_shape(self): + """Reset the buffers when the shape of activation and gradient change. + For example, for curriculum learning that changes the seqlen of each + sample, we need to call this whenever the seqlen is going to change. + """ + self.first_output_send = True + self.pipe_recv_buf = None + self.grad_layer = None + self._grad_layer_buf = [] + self.meta_buffer = None + + self.pipe_partition_input_meta_cache = None + self.pipe_partition_output_meta_cache = None + self.pipe_partition_grad_meta_cache = None + self.grad_partition_grad_layer_meta_cache = None + + def train_batch(self, data_iter=None): + """Progress the pipeline to train the next batch of data. The engine will ingest + ``self.train_batch_size()`` total samples collectively across all workers. + + + An iterator that over training data should be provided as an argument + unless ``deepspeed.initialize()`` was provided a training set. In that event, + the training data will automatically be read. + + + .. warning:: + A total of ``self.gradient_accumulation_steps()`` entries will be pulled + from ``data_iter`` by each pipeline. There must be sufficient + data left in ``data_iter`` or else a ``StopIteration`` will halt training. + + DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader` + that wraps data loaders to automatically restart upon a ``StopIteration``. + + Args: + data_iter (Iterator, optional): Iterator of training data. + + Returns: + The arithmetic mean of the losses computed this batch. + """ + if not torch._C.is_grad_enabled(): + raise RuntimeError(f'train_batch() requires gradients enabled. Use eval_batch() instead.') + + # Curriculum learning could change activation shape + if self.curriculum_enabled_legacy(): + new_difficulty = self.curriculum_scheduler_legacy.update_difficulty( \ + self.global_steps + 1) + if self.global_steps == 0 or self.curriculum_scheduler_legacy.first_step: + self.reset_activation_shape() + self.curriculum_scheduler_legacy.first_step = False + elif new_difficulty != self.curriculum_scheduler_legacy.get_difficulty( \ + self.global_steps): + self.reset_activation_shape() + + if data_iter is not None: + self.set_dataiterator(data_iter) + + self.module.train() + self.total_loss = None + self.total_additional_losses = None + self._compute_loss = True + + # Do the work + self.timers(TRAIN_BATCH_TIMER).start() + sched = schedule.TrainSchedule(micro_batches=self.micro_batches, + stages=self.num_stages, + stage_id=self.stage_id) + self._exec_schedule(sched) + + with torch.no_grad(): + self.agg_train_loss = self._aggregate_total_loss() + + self.timers(TRAIN_BATCH_TIMER).stop() + + if self.steps_per_print() is not None and self.global_steps % self.steps_per_print() == 0: + if self.global_rank == 0: + elapsed = self.timers(TRAIN_BATCH_TIMER).elapsed(reset=True) / 1000.0 + iter_time = elapsed / self.steps_per_print() + tput = self.train_batch_size() / iter_time + log_str = f'steps: {self.global_steps} loss: {self.agg_train_loss:0.4f} ' + if self.agg_additional_losses is not None: + for loss_name, loss_value in self.agg_additional_losses.items(): + log_str += f'{loss_name}: {loss_value.item():0.4f} ' + log_str += f'iter time (s): {iter_time:0.3f} samples/sec: {tput:0.3f}' + print(log_str) + else: + self.timers(TRAIN_BATCH_TIMER).elapsed(reset=True) + + # Monitoring + if self.global_rank == 0 and self.monitor.enabled: + self.summary_events = [(f'Train/Samples/train_loss', self.agg_train_loss.mean().item(), + self.global_samples)] + self.monitor.write_events(self.summary_events) + + if self.steps_per_print() is not None and self.wall_clock_breakdown( + ) and self.global_steps % self.steps_per_print() == 0: + self.timers.log([ + PIPE_SEND_OUTPUT_TIMER, + PIPE_SEND_GRAD_TIMER, + PIPE_RECV_INPUT_TIMER, + PIPE_RECV_GRAD_TIMER, + ]) + + # TODO: should return precisely what loss returned and allow others to be queried? + return self.agg_train_loss + + def eval_batch(self, + data_iter, + return_logits=False, + compute_loss=True, + reduce_output='avg', + bcast_loss=True, + num_micro_batches=None): + """Evaluate the pipeline on a batch of data from ``data_iter``. The + engine will evaluate ``self.train_batch_size()`` total samples + collectively across all workers. + + This method is equivalent to: + + .. code-block:: python + + module.eval() + with torch.no_grad(): + output = module(batch) + + .. warning:: + A total of ``self.gradient_accumulation_steps()`` entries will be pulled + from ``data_iter`` by each pipeline. There must be sufficient + data left in ``data_iter`` or else a ``StopIteration`` will halt training. + + DeepSpeed provides a convenience class :class:`deepspeed.utils.RepeatingLoader` + that wraps data loaders to automatically restart upon a ``StopIteration``. + + Args: + data_iter (Iterator): Iterator of data to evaluate. + + Returns: + The arithmetic mean of the losses computed this batch. + """ + self.eval_return_logits = return_logits + self.module.eval() + + # Curriculum learning could change activation shape + if self.curriculum_enabled_legacy(): + new_difficulty = self.curriculum_scheduler_legacy.update_difficulty( \ + self.global_steps + 1) + if self.global_steps == 0 or self.curriculum_scheduler_legacy.first_step: + self.reset_activation_shape() + self.curriculum_scheduler_legacy.first_step = False + elif new_difficulty != self.curriculum_scheduler_legacy.get_difficulty( \ + self.global_steps): + self.reset_activation_shape() + + eval_output = None + + self._compute_loss = compute_loss + + # Use the provided data iterator + train_iterator = self.data_iterator + self.set_dataiterator(data_iter) + + # set the number micro batches in case the user chose value than training + micro_batches = self.micro_batches if num_micro_batches is None else num_micro_batches + + # Do the work + sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=self.num_stages, stage_id=self.stage_id) + + # prevent dead-lock with multiple evals sequence + dist.barrier() + + with torch.no_grad(): + self._exec_schedule(sched) + + if self.is_last_stage(): + eval_output = self._reduce_outputs(self.fwd_outputs, reduce=reduce_output, micro_batches=micro_batches) + + if compute_loss and (bcast_loss or self.monitor.enabled): + eval_output = self._bcast_pipe_scalar(eval_output) + + if self.global_rank == 0 and self.monitor.enabled: + self.summary_events = [(f'Train/Samples/eval_loss', eval_output.mean().item(), self.global_samples)] + self.monitor.write_events(self.summary_events) + + # Restore the training iterator + self.set_dataiterator(train_iterator) + + # Reset any buffers that may have been populated during the forward passes. + #ds_checkpointing.reset() + self.eval_return_logits = False + if return_logits: + outputs = self.outputs + self.outputs = None + return eval_output, outputs + return eval_output + + def set_train_batch_size(self, train_batch_size): + """Adjust the global batch size by increasing or decreasing the number of + micro-batches (i.e., gradient accumulation steps). The size of each micro-batch + (i.e., ``train_micro_batch_size_per_gpu``) is not changed. + Args: + train_batch_size (int): The new global batch size for training. + Raises: + ValueError: if ``train_batch_size`` is not divisible by the + configured micro-batch size and data parallelism. + """ + super().set_train_batch_size(train_batch_size) + self.micro_batches = self.gradient_accumulation_steps() + + def is_first_stage(self): + """True if this process is in the first stage in the pipeline.""" + return self.stage_id == 0 + + def is_last_stage(self): + """True if this process is in the last stage in the pipeline.""" + return self.stage_id == self.num_stages - 1 + + def _reduce_outputs(self, outputs, reduce='avg', reduce_dp=True, micro_batches=None): + if reduce is None: + return outputs + + if reduce.lower() == 'avg': + # first sum over all microbatches + if torch.is_tensor(outputs[0]): + reduced = sum(outputs) + else: + assert isinstance(outputs, (list, tuple)) + reduced = [torch.zeros_like(o) for o in outputs[0]] + for idx, out in outputs: + reduced[idx] += out + + # Average over the microbatches + reduced = self._scale_loss_by_gas(reduced, eval_micro_batches=micro_batches) + + # Average over DP groups + if reduce_dp and self.is_data_parallel: + if torch.is_tensor(reduced): + dist.all_reduce(reduced, group=self.mpu.get_data_parallel_group()) + reduced /= self.dp_world_size + else: + for idx in range(len(reduced)): + dist.all_reduce(reduced[idx], group=self.mpu.get_data_parallel_group()) + reduced[idx] /= self.dp_world_size + + return reduced + else: + raise NotImplementedError(f'reduction type {reduce} not supported.') + + def _bcast_pipe_scalar(self, data, src_rank=None, dtype=torch.float32): + # Default to last stage (e.g., for broadcasting loss) + if src_rank is None: + src_rank = self.grid.stage_to_global(self.num_stages - 1) + assert src_rank in self.grid.pp_group + + if self.global_rank == src_rank: + result = data.clone().detach().type(dtype).to(self.device) + else: + result = torch.Tensor([0.]).type(dtype).to(self.device) + + dist.broadcast(tensor=result, src=src_rank, group=self.mpu.get_pipe_parallel_group()) + + return result + + def _aggregate_total_loss(self): + # Scale loss, average among DP ranks, and bcast loss to the rest of my DP group + if self.is_last_stage(): + # Scale loss and additional losses, if any + loss = self._scale_loss_by_gas(self.total_loss) + self.agg_additional_losses = self.total_additional_losses + if self.agg_additional_losses is not None: + self.agg_additional_losses = OrderedDict({ + loss_name: self._scale_loss_by_gas(_loss.clone().detach()) + for loss_name, _loss in self.agg_additional_losses.items() + }) + + self.dp_group_loss = loss.clone().detach() + agg_loss = self.dp_group_loss.clone().detach() + #print(f'RANK={self.global_rank} bcast SENDER src={self.global_rank} group={self.grid.pp_group}', flush=True) + + # Average loss across all data-parallel groups + if self.is_data_parallel: + if self.agg_additional_losses is None: + dist.all_reduce(agg_loss, group=self.mpu.get_data_parallel_group()) + agg_loss /= self.dp_world_size + else: + # use a single reduce op for agg_loss and additional losses, if any + assert '__train_loss__' not in self.agg_additional_losses.keys() + tensors = OrderedDict({'__train_loss__': agg_loss}) + tensors.update(self.agg_additional_losses.items()) + flat_tensor = torch.cat([t.clone().reshape(-1).detach() for t in tensors.values()]) + dist.all_reduce(flat_tensor, group=self.mpu.get_data_parallel_group()) + flat_tensor /= self.dp_world_size + offset = 0 + reduced_tensor = {} + for name, t in tensors.items(): + n_elem = t.numel() + reduced_tensor[name] = flat_tensor[offset:offset + n_elem].clone().detach().reshape(t.shape) + offset += n_elem + agg_loss = reduced_tensor['__train_loss__'] + self.agg_additional_losses = OrderedDict( + {name: reduced_tensor[name] + for name in self.agg_additional_losses.keys()}) + + assert self.global_rank in self.grid.pp_group + losses = [self.dp_group_loss, agg_loss] + if self.agg_additional_losses is not None: + losses += list(self.agg_additional_losses.values()) + losses = torch.stack(losses).float() + if self.is_pipe_parallel: + dist.broadcast(tensor=losses, src=self.global_rank, group=self.mpu.get_pipe_parallel_group()) + else: + # Get loss from last stage + src_rank = self.grid.stage_to_global(self.num_stages - 1) + assert src_rank in self.grid.pp_group + # losses to reduce are: dp_group_loss, agg_loss, model additional losses + # therefore: 2 + n_additional_losses + additional_losses = self.module.get_additional_losses() + n_additional_losses = 0 if additional_losses is None else len(additional_losses) + losses = torch.Tensor([0.] * (2 + n_additional_losses)).to(self.device) + dist.broadcast(tensor=losses, src=src_rank, group=self.grid.get_pipe_parallel_group()) + self.dp_group_loss = losses[0].clone().detach() + agg_loss = losses[1].clone().detach() + if additional_losses is not None: + self.agg_additional_losses = OrderedDict({ + name: losses[2 + i].clone().detach() + for i, name in enumerate(additional_losses.keys()) + }) + return agg_loss + + def set_dataloader(self, loader): + """""" + if self.is_first_stage() or self.is_last_stage(): + self.training_dataloader = loader + self.data_iterator = iter(self.training_dataloader) + + def set_dataiterator(self, iterator): + """ Store an iterator to sample for training data. """ + if self.is_first_stage() or self.is_last_stage(): + self.training_dataloader = None + self.data_iterator = iterator + + def set_batch_fn(self, fn): + """Execute a post-processing function on input data. + + Args: + fn (function): The function to run. + """ + self.batch_fn = fn + + def is_gradient_accumulation_boundary(self): + """True if the engine is executing a gradient reduction or optimizer step instruction. + + This is overridden from :class:`DeepSpeedEngine` to force reductions + and steps when the pipeline engine is instructed to do so. + + Returns: + bool: whether reductions and optimizer steps should occur. + """ + return self._force_grad_boundary + + def log_for_device(self, *msg): + if LOG_STAGE == self.stage_id or LOG_STAGE == -1: + if DATA_PARALLEL_ID == self.grid.data_parallel_id or DATA_PARALLEL_ID == -1: + print( + f'RANK={dist.get_rank()} ' + f'PIPE-ID={self.stage_id} ' + f'DATA-ID={self.grid.data_parallel_id} ' + f'MBATCH-ID={self.microbatch_id} ' + f'STEP-ID={self.log_batch_step_id} ' + '::', + *msg, + flush=True) + + def tput_log(self, *msg): + if self.global_rank == 0 and self.global_steps % self.steps_per_print() == 0: + print(*msg) + + def _next_batch(self): + # If using 3D parallelism, only some first-stage ranks may do IO + batch = None + if self.data_iterator is not None: + batch = next(self.data_iterator) + + # Any post-processing, like broadcasting across a slice-parallel group. + if self.batch_fn: + batch = self.batch_fn(batch) + + return batch + + def _exec_forward_pass(self, buffer_id): + self.tput_timer.start() + self.mem_status('BEFORE FWD', reset_max=True) + + if isinstance(self.pipe_buffers['inputs'][buffer_id], tuple): + inputs = tuple(t.clone() for t in self.pipe_buffers['inputs'][buffer_id]) + else: + inputs = self.pipe_buffers['inputs'][buffer_id].clone() + + # collect the partitioned input from the previous stage + if self.is_pipe_partitioned and not self.is_first_stage(): + if self.pipe_partition_input_meta_cache is None: + self.pipe_partition_input_meta_cache = inputs[0].to('cpu') + part_input = PartitionedTensor.from_meta(meta=self.pipe_partition_input_meta_cache, + local_part=inputs[1], + group=self.grid.get_slice_parallel_group()) + + inputs = (part_input.full(), *inputs[2:]) + inputs[0].requires_grad = True + # skip mask + #inputs[1].requires_grad = True + part_input = None + inputs = inputs[0] if len(inputs) == 1 else inputs + self.pipe_buffers['inputs'][buffer_id] = inputs + + # inputs has no gradient because it is from a cloned tensor + outputs = super().forward(inputs) + + # Reset activation checkpointing buffers. + # Need to call this between evaluation iterations + if not self.module.training: + ds_checkpointing.reset() + + # Partition the outputs if we are not the last stage + if self.is_pipe_partitioned and not self.is_last_stage(): + if isinstance(outputs, tuple): + first_output = outputs[0] + # TODO: Improve pipe partitioning to pass multiple tensors that require grads + assert all([torch.is_tensor(elt) and elt.requires_grad is False for elt in outputs[1:]]) + outputs_tail = outputs[1:] + elif torch.is_tensor(outputs): + first_output = outputs + outputs_tail = [] + else: + raise ValueError("expecting a tensor or a tuple of tensors") + part = PartitionedTensor(tensor=first_output, group=self.grid.get_slice_parallel_group()) + # Clear the large output data, but save the computation graph + first_output.data = torch.zeros(1, device=first_output.data.device) + self.pipe_buffers['output_tensors'][buffer_id] = first_output + # Inject the partitioned tensor into the output before sending + outputs = (part.to_meta(), part.data(), *outputs_tail) + part = None + + self.pipe_buffers['outputs'][buffer_id] = outputs + + # Optionally compute loss on the last device + if self.is_last_stage(): + if self._compute_loss and self.module.loss_fn is not None: + labels = self.pipe_buffers['labels'][buffer_id] + self.loss = self.module.loss_fn(outputs, labels) + else: + # Some models just return loss from forward() + self.loss = outputs + if self.eval_return_logits: + self.outputs = outputs + + if isinstance(self.loss, torch.Tensor): + self.fwd_outputs.append(self.loss.detach()) + else: + self.fwd_outputs.append([l.detach() for l in self.loss]) + + def add_to_total_loss(_total_loss, _loss): + if isinstance(_loss, torch.Tensor): + if _total_loss is None: + _total_loss = torch.zeros_like(_loss) + _total_loss += _loss.detach() + else: + if _total_loss is None: + _total_loss = [torch.zeros_like(_l) for _l in _loss] + for _idx, _l in enumerate(_loss): + _total_loss[_idx] += _l.detach() + return _total_loss + + self.total_loss = add_to_total_loss(self.total_loss, self.loss) + + # aggregate additional losses across gradient accumulation steps + additional_losses = self.module.get_additional_losses() + if additional_losses is not None: + if self.total_additional_losses is None: + self.total_additional_losses = OrderedDict() + for name, loss in additional_losses.items(): + total = self.total_additional_losses[name] if name in self.total_additional_losses else None + self.total_additional_losses[name] = add_to_total_loss(total, loss) + + def _exec_backward_pass(self, buffer_id): + assert self.optimizer is not None, "must provide optimizer during " \ + "init in order to use backward" + + self.mem_status('BEFORE BWD', reset_max=True) + + # The last stage just runs backward on the loss using DeepSpeed's typical + # mechanisms. + if self.is_last_stage(): + super().backward(self.loss) + self.mem_status('AFTER BWD') + return + + outputs = self.pipe_buffers['outputs'][buffer_id] + + if self.wall_clock_breakdown(): + self.timers(BACKWARD_MICRO_TIMER).start() + self.timers(BACKWARD_GLOBAL_TIMER).start() + self.timers(BACKWARD_INNER_MICRO_TIMER).start() + self.timers(BACKWARD_INNER_GLOBAL_TIMER).start() + + # Reconstruct if we previously partitioned the output. We must be + # careful to also restore the computational graph of the tensors we partitioned. + if self.is_pipe_partitioned: + if self.is_grad_partitioned: + if self.pipe_partition_output_meta_cache is None: + self.pipe_partition_output_meta_cache = outputs[0].to('cpu') + part_output = PartitionedTensor.from_meta(meta=self.pipe_partition_output_meta_cache, + local_part=outputs[1], + group=self.grid.get_slice_parallel_group()) + self.pipe_buffers['output_tensors'][buffer_id].data = part_output.full() + outputs = (self.pipe_buffers['output_tensors'][buffer_id], *outputs[2:]) + else: + # Already restored from partition + self.pipe_buffers['output_tensors'][buffer_id].data = outputs[0] + outputs = (self.pipe_buffers['output_tensors'][buffer_id], *outputs[1:]) + + grad_tensors = self.grad_layer + if self.is_grad_partitioned: + #print(f'RANK={self.global_rank} BEFORE-BWD restoring grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}') + if self.grad_partition_grad_layer_meta_cache is None: + self.grad_partition_grad_layer_meta_cache = self.grad_layer[0].to('cpu') + part_grad = PartitionedTensor.from_meta(meta=self.grad_partition_grad_layer_meta_cache, + local_part=self.grad_layer[1], + group=self.grid.get_slice_parallel_group()) + grad_tensors = (part_grad.full(), *grad_tensors[2:]) + part_grad = None + #print(f'RANK={self.global_rank} BEFORE-BWD restored grad={self.grad_layer[0].size()} {self.grad_layer[1].size()}') + + if self.using_bf16_optimizer and not self.is_last_stage(): + # manually call because we don't call optimizer.backward() + self.optimizer.clear_lp_grads() + + # This handles either a single tensor or tuple of tensors. + if isinstance(outputs, tuple): + out_tensors = [t for t in outputs if t.is_floating_point()] + assert len(out_tensors) == len(grad_tensors) + torch.autograd.backward(tensors=out_tensors, grad_tensors=grad_tensors) + else: + torch.autograd.backward(tensors=(outputs, ), grad_tensors=(grad_tensors, )) + + if self.using_bf16_optimizer and not self.is_last_stage(): + # manually call because we don't call optimizer.backward() + if not self._config.bfloat16_immediate_grad_update: + self.optimizer.update_hp_grads(clear_lp_grads=False) + + # Free up the memory from the output of forward() + self.pipe_buffers['output_tensors'][buffer_id] = None + self.pipe_buffers['outputs'][buffer_id] = None + grad_tensors = None + + if self.wall_clock_breakdown(): + self.timers(BACKWARD_INNER_MICRO_TIMER).stop() + self.timers(BACKWARD_INNER_GLOBAL_TIMER).stop() + self.timers(BACKWARD_MICRO_TIMER).stop() + self.timers(BACKWARD_GLOBAL_TIMER).stop() + + self.mem_status('AFTER BWD') + + def _exec_load_micro_batch(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(BATCH_INPUT_TIMER).start() + + batch = self._next_batch() + + if self.is_first_stage(): + loaded = None + if torch.is_tensor(batch[0]): + loaded = batch[0].clone().to(self.device).detach() + if self._config.pipeline['activation_checkpoint_interval'] > 0 and self._config.pipeline[ + 'use_reentrant']: + loaded.requires_grad = loaded.is_floating_point() + else: + assert isinstance(batch[0], (tuple, list)) + # Assume list or tuple + loaded = [] + for x in batch[0]: + assert torch.is_tensor(x) + mine = x.clone().detach().to(self.device) + if self._config.pipeline['activation_checkpoint_interval'] > 0 and self._config.pipeline[ + 'use_reentrant']: + mine.requires_grad = mine.is_floating_point() + loaded.append(mine) + loaded = tuple(loaded) + + self.pipe_buffers['inputs'][buffer_id] = loaded + + if self.is_last_stage(): + loaded = batch[1] + if torch.is_tensor(batch[1]): + loaded = batch[1].to(self.device) + # XXX: torch 1.6.0 DataLoader will auto convert tuple to list + elif isinstance(batch[1], (tuple, list)): + loaded = [] + for x in batch[1]: + assert torch.is_tensor(x) + x = x.to(self.device).detach() + loaded.append(x) + loaded = tuple(loaded) + + self.pipe_buffers['labels'][buffer_id] = loaded + + if self.wall_clock_breakdown(): + self.timers(BATCH_INPUT_TIMER).stop() + + def _send_tensor_meta(self, buffer, recv_stage): + """ Communicate metadata about upcoming p2p transfers. + + Metadata is communicated in this order: + * type (0: tensor, 1: list) + * num_tensors if type=list + foreach tensor in buffer: + * ndims + * shape + """ + meta_buffer = torch.empty(TENSOR_META_SIZE, dtype=torch.int32, device=self.device) + if isinstance(buffer, torch.Tensor): + meta_buf_list = [ + 0, # type of data (0: tensor, 1: list (unused), 2: tuple) + self.DTYPE_TO_ID[buffer.dtype], # dtype + len(buffer.size()) # ndims + ] + meta_buf_list.extend(buffer.size()) + assert len( + meta_buf_list + ) <= TENSOR_META_SIZE, f"Buffer for metadata is too small. Current buffer size: {TENSOR_META_SIZE} but required {len(meta_buf_list)}" + meta_buffer[:len(meta_buf_list)].copy_(torch.tensor(meta_buf_list, dtype=torch.int32)) + p2p.send(meta_buffer, recv_stage) + + elif isinstance(buffer, tuple): + meta_buf_list = [ + 2, # type of data (0: tensor, 1: list (unused), 2: tuple) + len(buffer) # num_tensors + ] + + for tensor in buffer: + assert isinstance(tensor, torch.Tensor) + meta_buf_list.append(self.DTYPE_TO_ID[tensor.dtype]) + meta_buf_list.append(len(tensor.size())) + meta_buf_list.extend(tensor.size()) + + assert len( + meta_buf_list + ) <= TENSOR_META_SIZE, f"Buffer for metadata is too small. Current buffer size: {TENSOR_META_SIZE} but required {len(meta_buf_list)}" + meta_buffer[:len(meta_buf_list)].copy_(torch.tensor(meta_buf_list, dtype=torch.int32)) + p2p.send(meta_buffer, recv_stage) + + else: + raise NotImplementedError(f'Could not send meta type {type(buffer)}') + + # Useful for performance debugging. + ''' + if self.grid.data_parallel_id == 0: + print(f'STAGE={self.stage_id} pipe-send-volume: {send_bytes/1024**2:0.2f}MB') + ''' + + def _recv_tensor_meta(self, send_stage): + """Receive metadata about upcoming p2p transfers and return allocated buffers. + + Returns: + Allocated buffer for receiving from send_stage. + """ + buffer = torch.empty(TENSOR_META_SIZE, dtype=torch.int32, device=self.device) + p2p.recv(buffer, send_stage) + + recv_type = buffer[0].item() + + # A single tensor will be sent. + if recv_type == 0: + recv_dtype = self.ID_TO_DTYPE[buffer[1].item()] + recv_ndims = buffer[2].item() + recv_shape = buffer[3:3 + recv_ndims].tolist() + return self._allocate_or_extend_buffers(0, recv_shape, recv_dtype) + + # List or tuple of tensors (recv_type == 1 (list) is currently unused) + elif recv_type == 1 or recv_type == 2: + num_tensors = buffer[1].item() + + buffers = [] + offset = 2 + for idx in range(num_tensors): + recv_dtype = self.ID_TO_DTYPE[buffer[offset].item()] + recv_ndims = buffer[offset + 1].item() + recv_shape = buffer[offset + 2:offset + 2 + recv_ndims].tolist() + offset += 2 + recv_ndims + + buffers.append(self._allocate_or_extend_buffers(idx, recv_shape, recv_dtype)) + + # Convert to tuples if requested. + if recv_type == 2: + buffers = tuple(buffers) + return buffers + + else: + raise NotImplementedError(f'Could not receive type {type(recv_type)}') + + def _exec_send_activations(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(PIPE_SEND_OUTPUT_TIMER).start() + + outputs = self.pipe_buffers['outputs'][buffer_id] + + # NCCL does not like to send torch.BoolTensor types, so cast the mask to half(). + # We could do char, but with half() we can eventually flatten with other fp16 + # messages (TODO) + if self.has_attention_mask or self.has_bool_tensors: + outputs = list(outputs) + outputs[-1] = outputs[-1].half() + outputs = tuple(outputs) + + if self.dynamic_shape or self.first_output_send: + self.first_output_send = False + self._send_tensor_meta(outputs, self.next_stage) + + if isinstance(outputs, torch.Tensor): + p2p.send(outputs, self.next_stage) + elif isinstance(outputs, tuple): + for idx, buffer in enumerate(outputs): + p2p.send(buffer, self.next_stage) + else: + raise NotImplementedError('Could not send output of type ' + f'{type(outputs)}') + + # Restore the boolean tensor + if self.has_attention_mask or self.has_bool_tensors: + outputs = list(outputs) + outputs[-1] = outputs[-1].bool() + outputs = tuple(outputs) + + if self.wall_clock_breakdown(): + self.timers(PIPE_SEND_OUTPUT_TIMER).stop() + + def _exec_send_grads(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(PIPE_SEND_GRAD_TIMER).start() + + inputs = self.pipe_buffers['inputs'][buffer_id] + + # Partition the gradient + if self.is_grad_partitioned: + if isinstance(inputs, tuple): + first_input = inputs[0] + assert all([torch.is_tensor(elt) for elt in inputs[1:]]) + inputs_grad_tail = [elt.grad for elt in inputs[1:]] + elif torch.is_tensor(inputs): + first_input = inputs + inputs_grad_tail = [] + else: + raise ValueError("expecting a tensor or a tuple of tensors") + assert torch.is_tensor(first_input) + part = PartitionedTensor(tensor=first_input.grad, group=self.grid.get_slice_parallel_group()) + + inputs = (part.to_meta(), part.data(), *inputs_grad_tail) + + # XXX Terrible hack + # Drop the attention mask from the input buffer here. It does not have + # a grad that needs to be communicated. We free the buffer immediately + # after, so no need to restore it. The receiver also has a hack that skips + # the recv. This is because NCCL does not let us send torch.BoolTensor :-(. + if self.has_attention_mask or self.has_bool_tensors: + inputs = list(inputs) + inputs.pop() + inputs = tuple(inputs) + + if isinstance(inputs, torch.Tensor): + assert inputs.grad is not None + p2p.send(inputs.grad, self.prev_stage) + else: + # XXX terrible hacky branch + if self.is_grad_partitioned: + # First two sends are partitioned gradient + p2p.send(inputs[0], self.prev_stage) + p2p.send(inputs[1], self.prev_stage) + else: + for idx, buffer in enumerate(inputs): + # Skip tensors that will not produce a grad + if not buffer.is_floating_point(): + assert buffer.grad is None + continue + assert buffer.grad is not None + p2p.send(buffer.grad, self.prev_stage) + + # We can free up the input buffer now + self.pipe_buffers['inputs'][buffer_id] = None + + if self.wall_clock_breakdown(): + self.timers(PIPE_SEND_GRAD_TIMER).stop() + + def _exec_recv_activations(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(PIPE_RECV_INPUT_TIMER).start() + + recvd = None + + # Allocate the buffer if necessary + if self.dynamic_shape or self.pipe_recv_buf is None: + self.pipe_recv_buf = self._recv_tensor_meta(self.prev_stage) + + if isinstance(self.pipe_recv_buf, torch.Tensor): + p2p.recv(self.pipe_recv_buf, self.prev_stage) + recvd = self.pipe_recv_buf.clone().detach() + recvd.requires_grad = recvd.is_floating_point() + else: + assert isinstance(self.pipe_recv_buf, tuple) + recvd = [None] * len(self.pipe_recv_buf) + for idx, buffer in enumerate(self.pipe_recv_buf): + assert torch.is_tensor(buffer) + # XXX hardcode meta type + if self.is_pipe_partitioned and idx == 0 and buffer.dtype != torch.long: + if self.meta_buffer is None: + self.meta_buffer = torch.zeros(buffer.size(), dtype=torch.long, device=self.device) + buffer = self.meta_buffer + + p2p.recv(buffer, self.prev_stage) + recvd[idx] = buffer.clone().detach() + + # NCCL does not like to send torch.BoolTensor types, so un-cast the + # attention mask + if self.has_attention_mask or self.has_bool_tensors: + recvd[-1] = recvd[-1].bool() + + recvd = tuple(recvd) + + for buffer in recvd: + buffer.requires_grad = buffer.is_floating_point() + + self.pipe_buffers['inputs'][buffer_id] = recvd + + if self.wall_clock_breakdown(): + self.timers(PIPE_RECV_INPUT_TIMER).stop() + + def _exec_recv_grads(self, buffer_id): + if self.wall_clock_breakdown(): + self.timers(PIPE_RECV_GRAD_TIMER).start() + + outputs = self.pipe_buffers['outputs'][buffer_id] + # XXX these shapes are hardcoded for Megatron + # Restore partitioned output if it was partitioned and we are sending full gradients + if self.is_pipe_partitioned and not self.is_grad_partitioned: + if self.pipe_partition_grad_meta_cache is None: + self.pipe_partition_grad_meta_cache = outputs[0].to('cpu') + part_output = PartitionedTensor.from_meta(meta=self.pipe_partition_grad_meta_cache, + local_part=outputs[1], + group=self.grid.get_slice_parallel_group()) + outputs[0].data = part_output.full() + outputs = (outputs[0], *outputs[2:]) + # save for backward + self.pipe_buffers['outputs'][buffer_id] = outputs + + # Allocate gradient if necessary + if self.dynamic_shape or self.grad_layer is None: + if isinstance(outputs, torch.Tensor): + self.grad_layer = self._allocate_or_extend_buffers(0, list(outputs.size()), outputs.dtype) + else: + # XXX This is a HACK + # When we exchange activations/gradients, the two pipe stages + # need to issue the send/recv with the same buffer sizes or + # else there is a deadlock. The is_floating_point() filter is + # used to avoid sending gradients for tensors that do not + # produce gradients. When TP>1, we partition the first + # activations/gradients across TP ranks to save communication + # volume and memory. That partitioned tensor is represented as + # two tensors: a 1/TPth chunk of the original data and also a + # small LongTensor storing the metadata used to reconstruct on + # the other side. When combined, the floating point filter also + # filtered out the metadata tensor. This quick (hacky) fix just + # branches on is_grad_partitioned so we don't filter out the + # metadata tensor. + if self.is_grad_partitioned: + sizes_and_dtypes = [(list(t.size()), t.dtype) + for t in outputs[:2]] + [(list(t.size()), t.dtype) + for t in outputs[2:] if t.is_floating_point()] + else: + sizes_and_dtypes = [(list(t.size()), t.dtype) for t in outputs if t.is_floating_point()] + + self.grad_layer = [ + self._allocate_or_extend_buffers(i, size, dtype) + for i, (size, dtype) in enumerate(sizes_and_dtypes) + ] + + if isinstance(self.grad_layer, torch.Tensor): + p2p.recv(self.grad_layer, self.next_stage) + else: + assert isinstance(outputs, tuple) + for idx, buffer in enumerate(self.grad_layer): + # XXX GPT-2 hack + if self.is_grad_partitioned and idx == 0 and buffer.dtype != torch.long: + buffer.data = torch.zeros(buffer.size(), dtype=torch.long, device=self.device) + p2p.recv(buffer, self.next_stage) + + if self.wall_clock_breakdown(): + self.timers(PIPE_RECV_GRAD_TIMER).stop() + + def _exec_optimizer_step(self, lr_kwargs=None): + if self.wall_clock_breakdown(): + self.timers(STEP_MICRO_TIMER).start() + self.timers(STEP_GLOBAL_TIMER).start() + self.mem_status('BEFORE STEP', reset_max=True) + + self._force_grad_boundary = True + self._take_model_step(lr_kwargs) + self._force_grad_boundary = False + + self.mem_status('AFTER STEP') + + if self.global_rank == 0 and self.monitor.enabled: + self.summary_events = [(f'Train/Samples/lr', self.get_lr()[0], self.global_samples)] + if self.fp16_enabled() and hasattr(self.optimizer, 'cur_scale'): + self.summary_events.append( + (f'Train/Samples/loss_scale', self.optimizer.cur_scale, self.global_samples)) + self.monitor.write_events(self.summary_events) + + if self.wall_clock_breakdown(): + self.timers(STEP_MICRO_TIMER).stop() + self.timers(STEP_GLOBAL_TIMER).stop() + if self.global_steps % self.steps_per_print() == 0: + self.timers.log([ + BATCH_INPUT_TIMER, + FORWARD_MICRO_TIMER, + BACKWARD_MICRO_TIMER, + BACKWARD_INNER_MICRO_TIMER, + BACKWARD_REDUCE_MICRO_TIMER, + STEP_MICRO_TIMER, + ]) + if self.global_steps % self.steps_per_print() == 0: + self.timers.log([ + FORWARD_GLOBAL_TIMER, + BACKWARD_GLOBAL_TIMER, + BACKWARD_INNER_GLOBAL_TIMER, + BACKWARD_REDUCE_GLOBAL_TIMER, + STEP_GLOBAL_TIMER, + ]) + + def _allocate_zeros(self, shape, **kwargs): + """ Allocate a tensor of zeros on the engine's device. + + Arguments: + shape: the shape of the tensor to allocate + kwargs: passed to torch.zeros() + + Returns: + A tensor from torch.zeros() allocated on self.device. + """ + if "dtype" not in kwargs: + if self.fp16_enabled(): + kwargs["dtype"] = torch.half + if self.bfloat16_enabled(): + kwargs["dtype"] = torch.bfloat16 + + return torch.zeros(shape, device=self.device, **kwargs) + + def _allocate_buffer(self, shape, num_buffers=-1, **kwargs): + buffers = [] + if num_buffers == -1: + num_buffers = self.num_pipe_buffers + for count in range(num_buffers): + buffers.append(self._allocate_zeros(shape, **kwargs)) + return buffers + + def _allocate_or_extend_buffers(self, idx, shape, dtype): + numel = reduce(mul, shape) if len(shape) > 0 else 1 + if len(self._grad_layer_buf) <= idx or self._grad_layer_buf[idx].numel() < numel: + new_buf = self._allocate_buffer(shape, dtype=dtype, num_buffers=1)[0] + if len(self._grad_layer_buf) <= idx: + self._grad_layer_buf.append(new_buf) + else: + self._grad_layer_buf[idx] = new_buf + return self._grad_layer_buf[idx] + else: + return self._grad_layer_buf[idx].flatten()[:numel].view(shape) + + def forward(self, *args, **kwargs): + """Disabled for pipeline parallel training. See ``train_batch()``. """ + raise PipelineError("Only train_batch() is accessible in pipeline mode.") + + def backward(self, *args, **kwargs): + """Disabled for pipeline parallel training. See ``train_batch()``. """ + raise PipelineError("Only train_batch() is accessible in pipeline mode.") + + def step(self, *args, **kwargs): + """Disabled for pipeline parallel training. See ``train_batch()``. """ + raise PipelineError("Only train_batch() is accessible in pipeline mode.") + + def mem_status(self, msg, print_rank=-1, reset_max=False): + return + global mem_alloced, mem_cached + if not self.global_steps == 0 or not self.global_steps == 9: + #return + pass + if self.mpu.get_data_parallel_rank() != 0: + return + + if self.global_rank != 0: + return + + rank = self.global_rank + if print_rank != -1 and rank != print_rank: + return + + get_accelerator().synchronize() + + if reset_max: + get_accelerator().reset_max_memory_cached() + get_accelerator().reset_max_memory_allocated() + + new_alloced = get_accelerator().memory_allocated() + new_cached = get_accelerator().memory_cached() + + delta_alloced = new_alloced - mem_alloced + delta_cached = new_cached - mem_cached + + mem_cached = new_cached + mem_alloced = new_alloced + + max_alloced = get_accelerator().max_memory_allocated() + max_cached = get_accelerator().max_memory_cached() + + # convert to GB for printing + new_alloced /= 1024**3 + new_cached /= 1024**3 + delta_alloced /= 1024**3 + delta_cached /= 1024**3 + max_alloced /= 1024**3 + max_cached /= 1024**3 + + print( + f'RANK={rank} STAGE={self.stage_id} STEP={self.global_steps} MEMSTATS', msg, + f'current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) ' + f'current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)') + + def module_state_dict(self, exclude_frozen_parameters=False): + """Override hack to save a pipe model and return the directory path of the save. + + This method should only be called by DeepSpeed's ``save_checkpoint()``. The + recommended way of saving a ``PipelineModule`` outside of ``save_checkpoint()`` + is ``save_state_dict()``. + + Returns: + None + """ + assert isinstance(self.module, PipelineModule) + assert self._curr_ckpt_path is not None, \ + "PipelineEngine expects module_state_dict() to be called from save_checkpoint()" + + self.module.save_state_dict(self._curr_ckpt_path, + checkpoint_engine=self.checkpoint_engine, + exclude_frozen_params=exclude_frozen_parameters) + return None + + def load_module_state_dict(self, checkpoint, strict=True, custom_load_fn=None, fetch_z3_params=False): + """Override hack to instead use a directory path. + + This is important because pipeline models checkpoint by layer instead of rank. + + If ``state_dict`` is not ``None`` or a ``str``, we revert to ``super()`` expecting a ``dict``. + + Args: + state_dict (str, None): unused + strict (bool, optional): Strict state loading. Defaults to True. + """ + assert custom_load_fn is None, "custom_load_fn not supported w. pipeline parallelism" + state_dict = checkpoint if self.has_moe_layers else checkpoint['module'] + if (state_dict is not None) and (not isinstance(state_dict, str)): + super().load_module_state_dict(state_dict, strict) + return + + self.module.load_state_dir(load_dir=self._curr_ckpt_path, + strict=strict, + checkpoint_engine=self.checkpoint_engine) + + # A map of PipeInstruction types to methods. Each method will be executed with the + # kwargs provided to the PipeInstruction from the scheduler. + _INSTRUCTION_MAP = { + schedule.OptimizerStep: _exec_optimizer_step, + schedule.ReduceGrads: _exec_reduce_grads, + schedule.ReduceTiedGrads: _exec_reduce_tied_grads, + schedule.LoadMicroBatch: _exec_load_micro_batch, + schedule.ForwardPass: _exec_forward_pass, + schedule.BackwardPass: _exec_backward_pass, + schedule.SendActivation: _exec_send_activations, + schedule.RecvActivation: _exec_recv_activations, + schedule.SendGrad: _exec_send_grads, + schedule.RecvGrad: _exec_recv_grads, + } + + def _exec_schedule(self, pipe_schedule): + # Reserve and reset buffers. + self._reserve_pipe_buffers(pipe_schedule.num_pipe_buffers()) + self.fwd_outputs = [] + + # For each step in the schedule + for step_cmds in pipe_schedule: + # For each instruction in the step + for cmd in step_cmds: + if type(cmd) not in self._INSTRUCTION_MAP: + raise RuntimeError(f'{self.__class__.__name__} does not understand instruction {repr(cmd)}') + + # Equivalent to: self._exec_forward_pass(buffer_id=0) + self._exec_instr = MethodType(self._INSTRUCTION_MAP[type(cmd)], self) + self._exec_instr(**cmd.kwargs) + + def get_additional_losses(self): + return self.agg_additional_losses diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/module.py b/lib/python3.12/site-packages/deepspeed/runtime/pipe/module.py new file mode 100644 index 0000000000000000000000000000000000000000..2bc0c37bffb7f52ea462315476bd8f5962779425 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/pipe/module.py @@ -0,0 +1,693 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import os +import glob + +import re as regex + +from functools import partial + +import torch +import torch.nn as nn +from deepspeed import comm as dist + +from deepspeed.utils import logger +from .. import utils as ds_utils +from ..activation_checkpointing import checkpointing +from .topology import PipeDataParallelTopology, PipelineParallelGrid +from deepspeed.runtime.state_dict_factory import SDLoaderFactory +from deepspeed.accelerator import get_accelerator +from deepspeed.checkpoint.utils import clone_tensors_for_torch_save + + +class PipelineError(Exception): + """Errors related to the use of deepspeed.PipelineModule """ + + +class LayerSpec: + """Building block for specifying pipeline-parallel modules. + + LayerSpec stores the type information and parameters for each stage in a + PipelineModule. For example: + + .. code-block:: python + + nn.Sequence( + torch.nn.Linear(self.in_dim, self.hidden_dim, bias=False), + torch.nn.Linear(self.hidden_hidden, self.out_dim) + ) + + becomes + + .. code-block:: python + + layer_specs = [ + LayerSpec(torch.nn.Linear, self.in_dim, self.hidden_dim, bias=False), + LayerSpec(torch.nn.Linear, self.hidden_hidden, self.out_dim)] + ] + """ + + def __init__(self, typename, *module_args, **module_kwargs): + self.typename = typename + self.module_args = module_args + self.module_kwargs = module_kwargs + + if not issubclass(typename, nn.Module): + raise RuntimeError('LayerSpec only supports torch.nn.Module types.') + + if dist.is_initialized(): + self.global_rank = dist.get_rank() + else: + self.global_rank = -1 + + def __repr__(self): + return ds_utils.call_to_str(self.typename.__name__, self.module_args, self.module_kwargs) + + def build(self, log=False): + """Build the stored specification.""" + if log: + logger.info(f'RANK={self.global_rank} building {repr(self)}') + + return self.typename(*self.module_args, **self.module_kwargs) + + +class TiedLayerSpec(LayerSpec): + + def __init__(self, key, typename, *module_args, forward_fn=None, tied_weight_attr=['weight'], **module_kwargs): + super().__init__(typename, *module_args, **module_kwargs) + self.key = key + self.forward_fn = forward_fn + self.tied_weight_attr = [tied_weight_attr] if type(tied_weight_attr) == str else tied_weight_attr + + +class PipelineModule(nn.Module): + """Modules to be parallelized with pipeline parallelism. + + The key constraint that enables pipeline parallelism is the + representation of the forward pass as a sequence of layers + and the enforcement of a simple interface between them. The + forward pass is implicitly defined by the module ``layers``. The key + assumption is that the output of each layer can be directly fed as + input to the next, like a ``torch.nn.Sequence``. The forward pass is + implicitly: + + .. code-block:: python + + def forward(self, inputs): + x = inputs + for layer in self.layers: + x = layer(x) + return x + + .. note:: + Pipeline parallelism is not compatible with ZeRO-2 and ZeRO-3. + + Args: + layers (Iterable): A sequence of layers defining pipeline structure. Can be a ``torch.nn.Sequential`` module. + num_stages (int, optional): The degree of pipeline parallelism. If not specified, ``topology`` must be provided. + topology (``deepspeed.runtime.pipe.ProcessTopology``, optional): Defines the axes of parallelism axes for training. Must be provided if ``num_stages`` is ``None``. + loss_fn (callable, optional): Loss is computed ``loss = loss_fn(outputs, label)`` + seed_layers(bool, optional): Use a different seed for each layer. Defaults to False. + seed_fn(type, optional): The custom seed generating function. Defaults to random seed generator. + base_seed (int, optional): The starting seed. Defaults to 1234. + partition_method (str, optional): The method upon which the layers are partitioned. Defaults to 'parameters'. + activation_checkpoint_interval (int, optional): The granularity activation checkpointing in terms of number of layers. 0 disables activation checkpointing. + activation_checkpoint_func (callable, optional): The function to use for activation checkpointing. Defaults to ``deepspeed.checkpointing.checkpoint``. + checkpointable_layers (list[str], optional): List of layer class names that are eligible for checkpointing. For GPT models, + ParallelTransformerLayerPipe is always checkpointed regardless of this list. If None, all layers with parameters are + considered checkpointable. Defaults to None. + dynamic_shape: Allows dynamic shapes of inputs. This might have a performance impact. + """ + + def __init__(self, + layers, + num_stages=None, + topology=None, + loss_fn=None, + seed_layers=False, + seed_fn=None, + base_seed=1234, + partition_method='parameters', + activation_checkpoint_interval=0, + activation_checkpoint_func=checkpointing.checkpoint, + checkpointable_layers=None, + dynamic_shape=False): + + super().__init__() + + if num_stages is None and topology is None: + raise RuntimeError('must provide num_stages or topology') + + self.micro_offset = 0 + + self.loss_fn = loss_fn + + self.checkpointable_layers = checkpointable_layers + if checkpointable_layers is not None: + assert isinstance(checkpointable_layers, list), "param `checkpointable_layers` must be type of list." + + self.seed_layers = seed_layers + self.seed_fn = seed_fn + self.base_seed = base_seed + if dist.get_rank() == 0: + try: + seed_str = self.seed_fn.__name__ + except AttributeError: + seed_str = None + print(f'SEED_LAYERS={self.seed_layers} BASE_SEED={self.base_seed} SEED_FN={seed_str}') + + # Setup world info + self.world_group = dist.new_group(ranks=range(dist.get_world_size())) + self.global_rank = dist.get_rank(group=self.world_group) + self.world_size = dist.get_world_size(group=self.world_group) + self.local_rank = int(os.environ.get("LOCAL_RANK", None)) + assert self.local_rank is not None + + if topology: + self._topo = topology + self.num_stages = self._topo.get_dim('pipe') + else: + self.num_stages = num_stages + if topology is None: + if self.world_size % self.num_stages != 0: + raise RuntimeError( + f'num_stages ({self.num_stages}) must divide distributed world size ({self.world_size})') + dp = self.world_size // num_stages + topology = PipeDataParallelTopology(num_pp=num_stages, num_dp=dp) + self._topo = topology + + # Construct communicators for pipeline topology + self._grid = PipelineParallelGrid(process_group=self.world_group, topology=self._topo) + + self.stage_id = self._topo.get_coord(self.global_rank).pipe + + # Initialize partition information + self._layer_specs = list(layers) + self._num_layers = len(self._layer_specs) + self._local_start = 0 + self._local_stop = None + self._partition_layers(method=partition_method) + + self.forward_funcs = [] + self.fwd_map = {} + self.tied_modules = nn.ModuleDict() + self.tied_weight_attrs = {} + + # Offset the random seed by the stage ID. + #newseed = get_accelerator().initial_seed() + self._grid.get_stage_id() + #ds_utils.set_random_seed(newseed) + + self.activation_checkpoint_interval = activation_checkpoint_interval + + self.activation_checkpoint_func = activation_checkpoint_func + + #storage for precomputed checkpointeble results + self.is_checkpointable_results = [] + self.is_checkpointable_results_interval = None + + # if configuration use_reentrant = False, self.activation_checkpoint_func will be set to ``checkpointing.non_reentrant_checkpoint`` + + #with torch.random.fork_rng(devices=[get_accelerator().current_device_name()]): + self._build() + self.to(get_accelerator().device_name(self.local_rank)) + + self.tied_comms = self._index_tied_modules() + self._synchronize_tied_weights() + + self.dynamic_shape = dynamic_shape + + def _precompute_checkpointable_values(self): + if self.activation_checkpoint_interval > 0 and self.is_checkpointable_results_interval != self.activation_checkpoint_interval: + num_layers = len(self.forward_funcs) + self.interval_was_zero = False + for start_idx in range(0, num_layers, self.activation_checkpoint_interval): + end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers) + funcs = self.forward_funcs[start_idx:end_idx] + self.is_checkpointable_results.append(self._is_checkpointable(funcs)) + self.is_checkpointable_results_interval = self.activation_checkpoint_interval + + def _build(self): + specs = self._layer_specs + + for local_idx, layer in enumerate(specs[self._local_start:self._local_stop]): + layer_idx = local_idx + self._local_start + if self.seed_layers: + if self.seed_fn: + self.seed_fn(self.base_seed + layer_idx) + else: + ds_utils.set_random_seed(self.base_seed + layer_idx) + + # Recursively build PipelineModule objects + if isinstance(layer, PipelineModule): + raise NotImplementedError('RECURSIVE BUILD NOT YET IMPLEMENTED') + + # LayerSpec objects contain an nn.Module that should be allocated now. + elif isinstance(layer, nn.Module): + name = str(layer_idx) + self.forward_funcs.append(layer) + self.fwd_map.update({name: len(self.forward_funcs) - 1}) + self.add_module(name, layer) + + # TiedLayerSpec objects contain an nn.Module that should be allocated now. + elif isinstance(layer, TiedLayerSpec): + # Build and register the module if we haven't seen it before. + if layer.key not in self.tied_modules: + self.tied_modules[layer.key] = layer.build() + self.tied_weight_attrs[layer.key] = layer.tied_weight_attr + + if layer.forward_fn is None: + # Just use forward() + self.forward_funcs.append(self.tied_modules[layer.key]) + else: + # User specified fn with args (module, input) + self.forward_funcs.append(partial(layer.forward_fn, self.tied_modules[layer.key])) + + # LayerSpec objects contain an nn.Module that should be allocated now. + elif isinstance(layer, LayerSpec): + module = layer.build() + name = str(layer_idx) + self.forward_funcs.append(module) + self.fwd_map.update({name: len(self.forward_funcs) - 1}) + self.add_module(name, module) + + # Last option: layer may be a functional (e.g., lambda). We do nothing in + # that case and just use it in forward() + else: + self.forward_funcs.append(layer) + + # All pipeline parameters should be considered as model parallel in the context + # of our FP16 optimizer + for p in self.parameters(): + p.ds_pipe_replicated = False + + def _get_frozen_parameter_names(self, layer): + """ Get names of frozen parameters in the layer. + + Returns: + A list of frozen parameter names + """ + if isinstance(layer, LayerSpec): + l = layer.build() + return [n for n, p in l.named_parameters() if not p.requires_grad] + elif isinstance(layer, nn.Module): + return [n for n, p in layer.named_parameters() if not p.requires_grad] + + return [] + + def _count_layer_params(self): + """Count the trainable parameters in individual layers. + + This routine will only build one layer at a time. + + Returns: + A list of the number of parameters in each layer. + """ + param_counts = [0] * len(self._layer_specs) + for idx, layer in enumerate(self._layer_specs): + if isinstance(layer, LayerSpec): + l = layer.build() + params = filter(lambda p: p.requires_grad, l.parameters()) + param_counts[idx] = sum(p.numel() for p in params) + elif isinstance(layer, nn.Module): + params = filter(lambda p: p.requires_grad, layer.parameters()) + param_counts[idx] = sum(p.numel() for p in params) + return param_counts + + def _find_layer_type(self, layername): + idxs = [] + typeregex = regex.compile(layername, regex.IGNORECASE) + for idx, layer in enumerate(self._layer_specs): + name = None + if isinstance(layer, LayerSpec): + name = layer.typename.__name__ + elif isinstance(layer, nn.Module): + name = layer.__class__.__name__ + else: + try: + name = layer.__name__ + except AttributeError: + continue + if typeregex.search(name): + idxs.append(idx) + + if len(idxs) == 0: + raise RuntimeError(f"Partitioning '{layername}' found no valid layers to partition.") + return idxs + + def forward(self, forward_input): + # We need to offset the seed by the microbatch ID. Save it in a local var to + # ensure it is preserved in the closure. Otherwise checkpointed forward funcs + # will see a different offset. + self.micro_offset += 1 + + def exec_range_func(start, end): + ''' Helper function to be used with checkpoint() + Adapted from torch.utils.checkpoint:checkpoint_sequential() + ''' + local_micro_offset = self.micro_offset + 1 + + def exec_func(*inputs): + # Single tensor inputs need to be unwrapped + if len(inputs) == 1: + inputs = inputs[0] + for idx, layer in enumerate(self.forward_funcs[start:end]): + self.curr_layer = idx + self._local_start + if self.seed_layers: + new_seed = (self.base_seed * local_micro_offset) + self.curr_layer + if self.seed_fn: + self.seed_fn(new_seed) + else: + ds_utils.set_random_seed(new_seed) + + inputs = layer(inputs) + return inputs + + return exec_func + + if self.activation_checkpoint_interval == 0: + func = exec_range_func(0, len(self.forward_funcs)) + x = func(forward_input) + else: + num_layers = len(self.forward_funcs) + x = forward_input + for start_idx, is_checkpointable_result in \ + zip(range(0, num_layers, self.activation_checkpoint_interval), self.is_checkpointable_results): + + end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers) + + funcs = self.forward_funcs[start_idx:end_idx] + # Since we either pass tensors or tuples of tensors without unpacking, we + # need to be careful not to double-wrap tensors with tuple. + if not isinstance(x, tuple): + x = (x, ) + + if is_checkpointable_result: + x = self.activation_checkpoint_func(exec_range_func(start_idx, end_idx), *x) + else: + x = exec_range_func(start_idx, end_idx)(*x) + return x + + def _partition_layers(self, method='uniform'): + num_stages = self._topo.get_dim('pipe') + stage_id = self._topo.get_coord(self.global_rank).pipe + + if self.global_rank == 0: + logger.info(f'Partitioning pipeline stages with method {method}') + + method = method.lower() + + # Each stage gets a simple uniform number of layers. + if method == 'uniform': + num_layers = len(self._layer_specs) + self.parts = ds_utils.partition_uniform(num_items=num_layers, num_parts=num_stages) + elif method == 'parameters': + param_counts = self._count_layer_params() + self.parts = ds_utils.partition_balanced(weights=param_counts, num_parts=num_stages) + elif method.startswith('type:'): + layertype = method.split(':')[1] + binary_weights = [0] * len(self._layer_specs) + for idx in self._find_layer_type(layertype): + binary_weights[idx] = 1 + self.parts = ds_utils.partition_balanced(weights=binary_weights, num_parts=num_stages) + elif method == 'profile': + raise NotImplementedError(f'Partitioning method {method} not implemented.') + else: + raise NotImplementedError(f'Partitioning method {method} not implemented.') + + # Print some information on the partitioning. + if self.global_rank == 0: + for stage in range(num_stages): + start = self.parts[stage] + stop = self.parts[stage + 1] + print(f'stage={stage} layers={stop - start}') + for idx, layer in enumerate(self._layer_specs[start:stop]): + name = str(layer) + if isinstance(layer, LayerSpec): + name = layer.typename.__name__ + if isinstance(layer, nn.Module): + name = layer.__class__.__name__ + else: + try: + name = layer.__name__ + except AttributeError: + pass + print(f' {idx+start:2d}: {name}') + if self.loss_fn: + try: + print(f' loss: {self.loss_fn.__name__}') + except AttributeError: + print(f' loss: {self.loss_fn.__class__.__name__}') + + self._set_bounds(start=self.parts[stage_id], stop=self.parts[stage_id + 1]) + + @staticmethod + def _recursive_getattr(module: torch.nn.Module, attr_name: str) -> torch.Tensor: + '''Allow getting an attribute like "linear.weight"''' + weight = module + for item in attr_name.split("."): + weight = getattr(weight, item) + return weight + + def allreduce_tied_weight_gradients(self): + '''All reduce the gradients of the tied weights between tied stages''' + for key, comm in self.tied_comms.items(): + for attr_name in comm['weight_attr']: + weight = self._recursive_getattr(self.tied_modules[key], attr_name) + dist.all_reduce(weight.grad, group=comm['group']) + + def get_tied_weights_and_groups(self): + weight_group_list = [] + for key, comm in self.tied_comms.items(): + for attr_name in comm['weight_attr']: + weight = self._recursive_getattr(self.tied_modules[key], attr_name) + weight_group_list.append((weight, comm['group'])) + return weight_group_list + + def _synchronize_tied_weights(self): + for key, comm in self.tied_comms.items(): + for attr_name in comm['weight_attr']: + dist.broadcast( + self._recursive_getattr(comm['module'], attr_name), + src=min(comm['ranks']), + group=comm['group'], + ) + + def _index_tied_modules(self): + ''' Build communication structures for tied modules. ''' + tied_comms = {} + if self._topo.get_dim('pipe') == 1: + return tied_comms + + specs = self._layer_specs + tie_keys = set(s.key for s in specs if isinstance(s, TiedLayerSpec)) + # Since Python 3.7, "Dictionary order is guaranteed to be insertion order." + # Sort tie_keys here so that orders of self.tied_comms.items() are consistent + # among ranks. + for key in sorted(tie_keys): + # Find the layers that the tied module appears in + tied_layers = [] + for idx, layer in enumerate(specs): + if isinstance(layer, TiedLayerSpec) and layer.key == key: + tied_layers.append(idx) + # Find all stages with this tied module + # TODO: Would be nice to remove the nested data/model parallelism loops and + # TODO: instead generalize in some way, since we really just care about the + # TODO: stage that owns the tied layer. Then loop over each (dp, mp, ...) + # TODO: fiber to generate process groups. + tied_stages = set(self.stage_owner(idx) for idx in tied_layers) + for dp in range(self._grid.data_parallel_size): + for mp in range(self._grid.get_slice_parallel_world_size()): + tied_ranks = [] + for s in sorted(tied_stages): + if self._grid.get_slice_parallel_world_size() > 1: + tied_ranks.append(self._grid.stage_to_global(stage_id=s, data=dp, model=mp)) + else: + tied_ranks.append(self._grid.stage_to_global(stage_id=s, data=dp)) + group = dist.new_group(ranks=tied_ranks) + + # Record this tied module if we own a local copy of it. + if self.global_rank in tied_ranks: + assert key in self.tied_modules + if key in self.tied_modules: + tied_comms[key] = { + 'ranks': tied_ranks, + 'group': group, + 'weight_attr': self.tied_weight_attrs[key], + 'module': self.tied_modules[key], + } + # Only count the tied module once in the eyes of the FP16 optimizer + if self.global_rank != tied_ranks[0]: + for p in self.tied_modules[key].parameters(): + p.ds_pipe_replicated = True + ''' + if len(tied_comms) > 0: + print(f'RANK={self.global_rank} tied_comms={tied_comms}') + ''' + + return tied_comms + + def partitions(self): + return self.parts + + def stage_owner(self, layer_idx): + assert 0 <= layer_idx < self._num_layers + for stage in range(self._topo.get_dim('pipe')): + if self.parts[stage] <= layer_idx < self.parts[stage + 1]: + return stage + raise RuntimeError(f'Layer {layer_idx} not owned? parts={self.parts}') + + def _set_bounds(self, start=None, stop=None): + """Manually define the range of layers that will be built on this process. + + These boundaries are treated as list slices and so start is inclusive and stop is + exclusive. The default of None for both results in all layers being built + locally. + """ + self._local_start = start + self._local_stop = stop + + def set_checkpoint_interval(self, interval): + assert interval >= 0 + self.checkpoint_interval = interval + + def topology(self): + """ ProcessTopology object to query process mappings. """ + return self._topo + + def mpu(self): + return self._grid + + def num_pipeline_stages(self): + return self._topo.get_dim('pipe') + + def ckpt_prefix(self, checkpoints_path, tag): + """Build a prefix for all checkpoint files written by this module. """ + # All checkpoint files start with this + rank_name = 'module' + + # Data parallelism is omitted from the naming convention because we are agnostic + # to this in the checkpoint. + omit_dims = frozenset(['data']) + axes = [a for a in self._grid._topo.get_axis_names() if a not in omit_dims] + for dim in axes: + rank = getattr(self._grid._topo.get_coord(rank=self.global_rank), dim) + rank_name += f'-{dim}_{rank:02d}' + + ckpt_name = os.path.join(checkpoints_path, str(tag), rank_name) + return ckpt_name + + def ckpt_layer_path(self, ckpt_dir, local_layer_idx): + """Customize a prefix for a specific pipeline module layer. """ + idx = local_layer_idx + self._local_start + layer_ckpt_path = os.path.join(ckpt_dir, f'layer_{idx:02d}') + rank_repr = self._grid._topo.get_rank_repr(rank=self.global_rank) + if rank_repr != '': + layer_ckpt_path += f'-{rank_repr}' + layer_ckpt_path += '-model_states.pt' + return layer_ckpt_path + + def ckpt_layer_path_list(self, ckpt_dir, local_layer_idx): + """Get all ckpt file list for a specific pipeline module layer. """ + idx = local_layer_idx + self._local_start + layer_ckpt_path = os.path.join(ckpt_dir, f'layer_{idx:02d}-') + layer_ckpt_path += "*model_states.pt" + ckpt_files = glob.glob(layer_ckpt_path) + ckpt_files.sort() + return ckpt_files + + def save_state_dict(self, save_dir, checkpoint_engine, exclude_frozen_params=False): + # Processes having the same model parallel rank on different data parallel instances + # have identical layer weights. We can distribute the task of saving the layer weights + # among the data parallel ranks. For example, if a pipeline stage has 9 layers and + # if there are 2 data parallel instances, rank 0 will save the first 5 layers and + # rank 1 will save the last 4. + dp_rank = self._grid.data_parallel_id + dp_size = self._grid.data_parallel_size + num_layers = len(self.forward_funcs) + if self.checkpoint_parallel_write_pipeline: + # spread layers evenly across data parallel ranks + offsets = ds_utils.partition_uniform(num_layers, dp_size) + start, end = offsets[dp_rank], offsets[dp_rank + 1] + else: + # data parallel rank 0 writes all layers + if dp_rank != 0: + return + start, end = 0, num_layers + layer_list = self.forward_funcs[start:end] + + checkpoint_engine.makedirs(save_dir, exist_ok=True) + for idx, layer in enumerate(layer_list): + model_ckpt_path = self.ckpt_layer_path(save_dir, start + idx) + if not hasattr(layer, 'state_dict'): + continue + + orig_state_dict = layer.state_dict() + if exclude_frozen_params: + for n in self._get_frozen_parameter_names(layer): + del orig_state_dict[n] + final_state_dict = clone_tensors_for_torch_save(orig_state_dict) + checkpoint_engine.save(final_state_dict, model_ckpt_path) + + def load_state_dir(self, load_dir, checkpoint_engine, strict=True): + for idx, layer in enumerate(self.forward_funcs): + # Functions, etc. will not have state_dicts + if not hasattr(layer, 'load_state_dict'): + continue + + # get all checkpoint files for the layer. + model_ckpt_list = self.ckpt_layer_path_list(load_dir, idx) + mp_rank = self._grid.get_slice_parallel_rank() + mp_world_size = self._grid.get_slice_parallel_world_size() + + sd_loader = SDLoaderFactory.get_sd_loader(model_ckpt_list, + version=2.0, + checkpoint_engine=checkpoint_engine) + load_path, checkpoint, _ = sd_loader.load(mp_world_size, mp_rank, module_key=None, is_pipe_parallel=True) + + layer.load_state_dict(checkpoint, strict=strict) + + # if self._grid.data_parallel_id == 0: + # logger.info( + # f'RANK={self.global_rank} Loaded layer={idx+self._local_start} file={load_path}' + # ) + + self._synchronize_tied_weights() + + def _is_checkpointable(self, funcs): + + if self.activation_checkpoint_func is not checkpointing.non_reentrant_checkpoint: + # This hook excludes the embedding layer + # because only non_reentrant_checkpoint can accept inputs with requires_grad=False + # otherwise, the backward of the embedding layer won't receive gradients. + if self.__class__.__name__ in ('GPTModelPipe', 'GPT2ModelPipe'): + # For GPT models, checkpoint both transformer layers and any additional + # layers specified in checkpointable_layers (if provided) + return all('ParallelTransformerLayerPipe' in f.__class__.__name__ or ( + self.checkpointable_layers is not None and f.__class__.__name__ in self.checkpointable_layers) + for f in funcs) + + if self.checkpointable_layers is not None: + # For non-GPT models, only checkpoint layers specified in checkpointable_layers + return all(f.__class__.__name__ in self.checkpointable_layers for f in funcs) + + # Default behavior: checkpoint any layer that has parameters + params = [f.parameters() for f in funcs if isinstance(f, torch.nn.Module)] + return any(len(list(p)) > 0 for p in params) + + def get_additional_losses(self): + """ Returns model specific additional losses for reporting + + Return a dictionary of {"loss name": loss_value} or None if no additional losses. + """ + return None + + def compile(self, *args, **kwargs): + for idx, layer in enumerate(self.forward_funcs): + if isinstance(layer, nn.Module): + layer.compile(*args, **kwargs) + else: + new_layer = torch.compile(layer, *args, **kwargs) + self.forward_funcs[idx] = new_layer diff --git a/lib/python3.12/site-packages/deepspeed/runtime/pipe/schedule.py b/lib/python3.12/site-packages/deepspeed/runtime/pipe/schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..21bf24dad8edcffd54d05d0c7092ab225e1e39dc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/pipe/schedule.py @@ -0,0 +1,494 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from ..utils import call_to_str + +from abc import ABC, abstractmethod + + +class PipeSchedule(ABC): + """Directs the execution of a pipeline engine by generating sequences of + :class:`PipeInstruction`. + + Schedules are generators that yield sequences of + :class:`PipeInstruction` to process the micro-batches in one batch. + Each yielded step is atomic in the sense that a barrier + synchronization can be placed between successive steps without + deadlock. + + Below is an example schedule that implements data parallelism with gradient accumulation: + + .. code-block:: python + + class DataParallelSchedule(PipeSchedule): + def steps(self): + for step_id in range(self.micro_batches): + cmds = [ + LoadMicroBatch(buffer_id=0), + ForwardPass(buffer_id=0), + BackwardPass(buffer_id=0), + ] + if step_id == self.micro_batches - 1: + cmds.extend([ + ReduceGrads(), + OptimizerStep(), + ]) + yield cmds + + def num_pipe_buffers(self): + return 1 + + Args: + micro_batches (int): The number of micro-batches that comprise a batch. + stages (int): The number of pipeline stages. + stage_id (int): The pipe stage that will execute the generated schedule. + """ + + def __init__(self, micro_batches, stages, stage_id): + super().__init__() + self.micro_batches = micro_batches + self.stages = stages + self.stage_id = stage_id + self.prev_stage = self.stage_id - 1 + self.next_stage = self.stage_id + 1 + + @abstractmethod + def steps(self): + """Yield a list of :class:`PipeInstruction` for each step in the schedule. + + .. note:: + Schedules must implement ``steps()`` to define the schedule. + + Returns: + Instructions to be executed as one step of the pipeline + """ + pass + + def num_pipe_buffers(self): + """The number of pipeline buffers that will be used by this stage. + + .. note:: + Schedules should specialize ``num_pipe_buffers()`` for memory savings at scale. + + Returns: + The number of buffers for the engine to allocate. + """ + return self.micro_batches + + def _valid_micro_batch(self, micro_batch_id): + return 0 <= micro_batch_id < self.micro_batches + + def _valid_stage(self, stage_id): + return 0 <= stage_id < self.stages + + @property + def stage(self): + """Stage index used to configure this schedule.""" + return self.stage_id + + @property + def num_stages(self): + """The number of total pipeline stages used to configure this schedule.""" + return self.stages + + @property + def num_micro_batches(self): + """The number of total micro_batches used to configure this schedule.""" + return self.micro_batches + + @property + def is_first_stage(self): + """True if the configured ``stage_id`` is the first stage in the pipeline.""" + return self.stage_id == 0 + + @property + def is_last_stage(self): + """True if the configured ``stage_id`` is the last stage in the pipeline.""" + return self.stage_id == self.stages - 1 + + def _buffer_idx(self, micro_batch_id): + """Map a micro-batch index to a pipeline buffer index. + + This method uses a cyclic allocation strategy. + + Args: + micro_batch_id (int): The micro-batch index relative to the beginning of the schedule. + + Returns: + int: The index of the buffer that should store data. + """ + assert self._valid_micro_batch(micro_batch_id) + return micro_batch_id % self.num_pipe_buffers() + + def __iter__(self): + self.it = None + return self + + def __next__(self): + if self.it is None: + self.it = self.steps() + return next(self.it) + + +class InferenceSchedule(PipeSchedule): + """A schedule for inferencing batches using pipeline parallelism. + """ + + def steps(self): + """""" + prev_micro_batch_id = -1 + total_steps = self.micro_batches + self.stages - 1 + for step_id in range(total_steps): + cmds = [] + micro_batch_id = step_id - self.stage_id + + # Alternate send/recv buffers + if _is_even(self.stage_id): + recv_buf = step_id % 2 + send_buf = (step_id + 1) % 2 + else: + recv_buf = (step_id + 1) % 2 + send_buf = step_id % 2 + + if self.is_first_stage or self.is_last_stage: + if self._valid_micro_batch(micro_batch_id): + cmds.append(LoadMicroBatch(recv_buf)) + + if _is_even(self.stage_id): + if self._valid_stage(self.next_stage): + if self._valid_micro_batch(micro_batch_id - 1): + cmds.append(SendActivation(send_buf)) + if self._valid_stage(self.prev_stage): + if self._valid_micro_batch(micro_batch_id): + cmds.append(RecvActivation(recv_buf)) + else: + if self._valid_stage(self.prev_stage): + if self._valid_micro_batch(micro_batch_id): + cmds.append(RecvActivation(recv_buf)) + + if self._valid_stage(self.next_stage): + if self._valid_micro_batch(micro_batch_id - 1): + cmds.append(SendActivation(send_buf)) + + if self._valid_micro_batch(micro_batch_id): + cmds.append(ForwardPass(recv_buf)) + + yield cmds + + def num_pipe_buffers(self): + """Only two pipeline buffers are required for inferencing. + + Returns: + ``2`` + """ + return 2 + + +class TrainSchedule(PipeSchedule): + """A schedule for training a batch using hybrid parallelism. + + Pipeline parallelism is extracted through gradient accumulation and thus + convergence follows that of a data parallel approach with the same batch + size. + """ + + def steps(self): + """""" + prev_micro_batch_id = -1 + total_steps = 2 * (self.micro_batches + self.stages - 1) + for step_id in range(total_steps): + # Map the step of the pipeline to the micro-batch id and also whether it is a + # forward or backward pass step. + micro_batch_id, is_forward = self._step_to_micro_batch(step_id) + + if self._valid_micro_batch(prev_micro_batch_id): + prev_buffer = self._buffer_idx(prev_micro_batch_id) + if self._valid_micro_batch(micro_batch_id): + curr_buffer = self._buffer_idx(micro_batch_id) + + cmds = [] + + # Exchange activations + if is_forward: + if self._valid_micro_batch(prev_micro_batch_id) and self._valid_stage(self.prev_stage): + cmds.append(SendGrad(prev_buffer)) + if self._valid_micro_batch(micro_batch_id) and self._valid_stage(self.prev_stage): + cmds.append(RecvActivation(curr_buffer)) + else: + if self._valid_micro_batch(micro_batch_id) and self._valid_stage(self.next_stage): + cmds.append(RecvGrad(curr_buffer)) + if self._valid_micro_batch(prev_micro_batch_id) and self._valid_stage(self.next_stage): + cmds.append(SendActivation(prev_buffer)) + + # First/last stage loads + if self.stage_id == 0 or self.stage_id == self.stages - 1: + if is_forward and self._valid_micro_batch(micro_batch_id): + cmds.append(LoadMicroBatch(curr_buffer)) + + # Computation + if self._valid_micro_batch(micro_batch_id): + if is_forward: + cmds.append(ForwardPass(curr_buffer)) + else: + cmds.append(BackwardPass(curr_buffer)) + + # Model step at the end of the batch + if step_id == total_steps - 1: + cmds.append(ReduceTiedGrads()) + cmds.append(ReduceGrads()) + cmds.append(OptimizerStep()) + + # Prepare state for next time + prev_micro_batch_id = micro_batch_id + yield cmds + + def num_pipe_buffers(self): + """Return the number of pipeline buffers required for this stage. + + This is equivalent to the maximum number of in-flight forward passes, + since we need to remember the activations of forward passes in order + to run backpropagation. For synchronous 1F1B, this is equivalent to + the index difference between this stage and the last stage. + """ + buffers = min(self.stages - self.stage_id, self.micro_batches) + return max(2, buffers) + + def _step_to_micro_batch(self, step_id): + if _is_even(step_id) and _is_even(self.stage_id): + micro_batch_id = self._even_step_forward_id(step_id) + is_forward = True + + elif _is_odd(step_id) and _is_odd(self.stage_id): + micro_batch_id = self._odd_step_forward_id(step_id) + is_forward = True + + elif _is_even(step_id) and _is_odd(self.stage_id): + micro_batch_id = self._even_step_backward_id(step_id) + is_forward = False + + elif _is_odd(step_id) and _is_even(self.stage_id): + micro_batch_id = self._odd_step_backward_id(step_id) + is_forward = False + + else: + assert False + + return micro_batch_id, is_forward + + def _even_step_forward_id(self, step_id): + base = step_id // 2 + micro_batch_id = int(base - self.stage_id // 2) + return micro_batch_id + + def _odd_step_forward_id(self, step_id): + base = (step_id - 1) // 2 + micro_batch_id = int(base - self.stage_id // 2) + return micro_batch_id + + def _even_step_backward_id(self, step_id): + base = step_id // 2 + micro_batch_id = int(base - self.stages + (self.stage_id + 1) // 2) + return micro_batch_id + + def _odd_step_backward_id(self, step_id): + base = ((step_id - 1) // 2) - self.stages + 1 + micro_batch_id = int(base + self.stage_id // 2) + return micro_batch_id + + +class DataParallelSchedule(PipeSchedule): + """An example schedule that trains using traditional data parallelism with gradient + accumulation. + """ + + def steps(self): + """""" + for step_id in range(self.micro_batches): + cmds = [ + LoadMicroBatch(buffer_id=0), + ForwardPass(buffer_id=0), + BackwardPass(buffer_id=0), + ] + if step_id == self.micro_batches - 1: + cmds.extend([ + ReduceGrads(), + OptimizerStep(), + ]) + yield cmds + + def num_pipe_buffers(self): + """Only one pipeline buffer needed. + """ + return 1 + + +class PipeInstruction: + """Base class for all instructions to be executed by the pipeline engine. + + All keyword arguments are stored as members similar to a ``namedtuple``. These are + then accessible to the :class:`PipeEngine` during execution. + + Args: + kwargs (optional): keyword arguments to store as members + """ + + def __init__(self, **kwargs): + self.name = self.__class__.__name__ + self.kwargs = kwargs + for key, val in kwargs.items(): + setattr(self, key, val) + + def __repr__(self): + return call_to_str(self.name, **self.kwargs) + + +class OptimizerStep(PipeInstruction): + """Performs one step with the optimizer and zeros gradients. + + .. note:: Should be issued after :class:`ReduceGrads` and :class:`ReduceTiedGrads`. + + .. note:: Can be a synchronization point among data-parallel ranks. + """ + pass + + +class ReduceGrads(PipeInstruction): + """Reduce the computed gradients among data-parallel processes within the stage. + """ + pass + + +class ReduceTiedGrads(PipeInstruction): + """Reduce the computed gradients of tied modules within a pipeline-parallel group. + + .. warning:: + The stages included in this synchronization point are not known until + the model is partitioned among pipeline stages. In the worst case, it + includes all pipeline stages. This instruction should be scheduled + carefully to avoid deadlocks. + """ + pass + + +class BufferOpInstruction(PipeInstruction): + """A pipeline instruction that operates on pipeline buffer(s). + + Args: + buffer_id (int): the index of the pipeline buffer() to modify. + """ + + def __init__(self, buffer_id, **kwargs): + super().__init__(buffer_id=buffer_id, **kwargs) + + +# IO +class LoadMicroBatch(BufferOpInstruction): + """Load a micro-batch into a buffer. + + Roughly: + + .. code-block:: python + + buffers['inputs'][buffer_id] = next(data_iter) + """ + pass + + +# Compute +class ForwardPass(BufferOpInstruction): + """Compute a forward pass. + + Roughly: + + .. code-block:: python + + buffers['outputs'][buffer_id] = forward(buffers['inputs'][buffer_id]) + """ + pass + + +class BackwardPass(BufferOpInstruction): + """Compute a backward pass and accumulate gradients. + + Roughly: + + .. code-block:: python + + outputs = buffers['outputs'][buffer_id] + gradients = buffers['gradients'][buffer_id] + torch.autograd.backward(tensors=outputs, + grad_tensors=gradients) + """ + pass + + +# Communication +class SendActivation(BufferOpInstruction): + """Send activations to the next stage in the pipeline. + + Roughly: + + .. code-block:: python + + send(buffers['outputs'][buffer_id]) + + .. note:: + The communication is blocking and must be paired with a :class:`RecvActivation` + on the next pipeline stage to avoid deadlock. + """ + pass + + +class RecvActivation(BufferOpInstruction): + """Receive activations from the previous stage in the pipeline. + + Roughly: + + .. code-block:: python + + buffers['inputs'][buffer_id] = recv() + + .. note:: + The communication is blocking and must be paired with a :class:`SendActivation` + on the previous pipeline stage to avoid deadlock. + """ + pass + + +class SendGrad(BufferOpInstruction): + """Send computed gradients to the previous pipeline stage. + with respect to the received activations + + .. note:: + Only received tensors with ``requires_grad==True`` will produce gradients. + Missing gradients will be replaced with ``None`` on the receiving stage. + + .. note:: + The communication is blocking and must be paired with a :class:`RecvGrad` + on the previous pipeline stage to avoid deadlock. + """ + pass + + +class RecvGrad(BufferOpInstruction): + """Receive computed gradients the next pipeline stage. + + .. note:: + Only activations with ``requires_grad==True`` will produce gradients. + Missing gradients will be replaced with ``None``. + + .. note:: + The communication is blocking and must be paired with a :class:`SendGrad` + on the next pipeline stage to avoid deadlock. + """ + pass + + +def _is_even(x): + return x % 2 == 0 + + +def _is_odd(x): + return x % 2 != 0 diff --git a/lib/python3.12/site-packages/deepspeed/runtime/progressive_layer_drop.py b/lib/python3.12/site-packages/deepspeed/runtime/progressive_layer_drop.py new file mode 100644 index 0000000000000000000000000000000000000000..bd3be8f13011a4c9322839ff8a631778e7fb6847 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/progressive_layer_drop.py @@ -0,0 +1,40 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import numpy as np +from deepspeed.utils import log_dist + + +class ProgressiveLayerDrop(object): + r""" Progressive Layer Dropping (PLD) for model training. + This implements the PLD technique for compressed model training + from this paper: https://arxiv.org/pdf/2010.13369.pdf + Args: + theta (float): a hyper-parameter that controls the trade-off between training time and robustness. + The lower the theta value, the faster the training speed. Default value: 0.5. + gamma (float): a hyper-parameter that controls how fast the drop ratio increases. Default value: 0.001. + """ + + def __init__(self, theta=0.5, gamma=0.001): + super().__init__() + + self.theta = theta + self.gamma = gamma + self.current_theta = 1.0 + log_dist(f'Enabled progressive layer dropping (theta = {self.theta})', ranks=[0]) + + def get_state(self): + kwargs = {'progressive_layer_drop': True, 'pld_theta': self.get_theta()} + return kwargs + + def get_theta(self): + return self.current_theta + + def update_state(self, global_step): + + def _prob(x, gamma, p): + return (1. - p) * np.exp(-gamma * x) + p + + self.current_theta = _prob(global_step, self.gamma, self.theta) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/quantize.py b/lib/python3.12/site-packages/deepspeed/runtime/quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..ad553138b2d7c8095aa1568f56a1f1a0b25e16d3 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/quantize.py @@ -0,0 +1,180 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import math +from deepspeed.utils import logger +from deepspeed.ops.quantizer import ds_quantizer + +TWO_D_PARAMS = 6 + + +class Quantizer(object): + + def __init__(self, + q_groups=1, + q_mixed_fp16=False, + q_change_ratio=0.01, + q_type=0, + q_rounding=0, + q_verbose=False, + q_eigenvalue=False, + use_quantizer_kernel=False, + layer_num=0): + + self.q_groups = q_groups + self.q_mixed_fp16 = q_mixed_fp16 + self.q_change_ratio = q_change_ratio + self.q_type = q_type + self.qsteps = 0 + self.quantize_real_ratio = 1.000 + self.q_verbose = q_verbose + self.q_eigenvalue = q_eigenvalue + self.use_quantizer_kernel = use_quantizer_kernel + self.q_rounding = q_rounding + self.layer_num = layer_num + + def any_precision_switch(self): + # Temporary disabled functionality + if self.layer_num == 0: + return True + result = False + for index in range(self.layer_num): + if self.q_start_bits[index] != self.q_target_bits: + next_step = self.qsteps + (TWO_D_PARAMS * (self.layer_num if self.layer_num != 0 else 1)) + if next_step >= self.q_period[index]: + result = True + return result + + def quantize(self, parameter_group, overflow, eigenvalue_enabled, block_eigenvalue={}): + + if overflow and not eigenvalue_enabled: + return + + self.step() + + self.update_fp16_ratio() + + for i in range(len(parameter_group)): + for p in parameter_group[i]: + if len(p.size()) > 1 and hasattr(p, "start_bits") and p.start_bits: + param_id = id(p) + if block_eigenvalue is None: + eigenvalue, layer_id = None, 0 + else: + eigenvalue, layer_id = block_eigenvalue[param_id] if param_id in block_eigenvalue else (None, + 0) + if eigenvalue is not None: + factor = 1 + math.floor(eigenvalue * 4) + p.data = self.compute_quantization(p.data, layer_id, factor) + else: + p.data = self.compute_quantization(p, layer_id) + + def step(self): + self.qsteps += 1 + + def quantize_highbit(self, inputs, num_bits): + + q_range = 2**num_bits + input_flat = inputs.reshape(self.q_groups, -1) + g_min = input_flat.amin(dim=-1, keepdim=True) + g_max = input_flat.amax(dim=-1, keepdim=True) + + # Random number generator (Uniform) + if self.q_rounding == 'nearest': + p = 0. + else: + p = input_flat.new(input_flat.shape).uniform_(-0.5, 0.5) + + if self.q_type == 'symmetric': + scale = 2 * torch.max(torch.abs(g_min), torch.abs(g_max)) / q_range + zero_point = 0. + input_flat = (input_flat / scale + p).round().clamp(-(q_range >> 1), (q_range >> 1) - 1) * scale + elif self.q_type == 'asymmetric': + scale = (g_max - g_min) / q_range + zero_point = (g_min / scale).round() * scale + input_flat = ((input_flat - zero_point) / scale + p).round().clamp(0, (q_range - 1)) * scale + zero_point + output = input_flat.reshape(inputs.shape).contiguous() + return output + + def quantize_tenary(self, inputs): + input_flat = inputs.reshape(self.q_groups, -1) + n = input_flat.shape[1] + m = input_flat.norm(p=1, dim=1).div(n) + thres = (0.7 * m).view(-1, 1) #.expand_as(input_flat) + pos = (input_flat > thres).type(inputs.type()) + neg = (input_flat < -thres).type(inputs.type()) + mask = (input_flat.abs() > thres).type(inputs.type()) + alpha = ((mask * input_flat).abs().sum(dim=1) / mask.sum(dim=1)).view(-1, 1) + output = alpha * pos - alpha * neg + output = output.reshape(inputs.shape).contiguous() + return output + + def quantize_binary(self, inputs): + input_flat = inputs.reshape(self.q_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(inputs.shape).contiguous() + return output + + def mixed_fp16_quantize(self, input, input_q, index): + if self.q_mixed_fp16 and self.q_start_bits[index] >= (self.q_target_bits - 1): + input_q = input * self.quantize_real_ratio + (1 - self.quantize_real_ratio) * input_q + return input_q + return input_q + + def compute_quantization(self, input, index=0, factor=1): + # fixing the quantization bits based on the training steps + # when reducing 1 bit at each period, we increase the period + # to go slowly toward the target quantization bits + # the period and starting bit can be configured + + if input.start_bits != input.target_bits: + if self.qsteps >= input.q_period: + self.quantize_real_ratio = 1.0 + input.q_period <<= 1 + input.q_period *= factor + input.start_bits -= 1 + if self.q_verbose: + logger.info( + f'Quantization settings: current bit-precision = {input.start_bits}, step = {self.qsteps}, quantization period = {input.q_period}, index = {index}' + ) + assert (input.start_bits >= input.target_bits), \ + 'Quantization bit is lower than target precision bits!' + + if self.use_quantizer_kernel: + if input.start_bits <= 2: + raise ValueError('Quantization bit is too low, please do it without quantization kernel!') + input_q = ds_quantizer(input.data.clone(), + self.q_groups, + input.start_bits, + asym=False if self.q_type == 'symmetric' else True, + sr=False if self.q_rounding == 'nearest_neighbor' else True) + else: + if input.start_bits >= 3: + input_flat = self.quantize_highbit(input.data, input.start_bits) + elif input.start_bits == 2: + assert self.q_type == 'symmetric', 'Quantization type is not symmetric!' + assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!' + input_flat = self.quantize_tenary(input.data) + elif input.start_bits == 1: + assert self.q_type == 'symmetric', 'Quantization type is not symmetric!' + assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!' + input_flat = self.quantize_binary(input.data) + if self.use_quantizer_kernel: + return self.mixed_fp16_quantize(input.data, input_q, index) + else: + if self.q_mixed_fp16 and input.start_bits >= input.target_bits - 1: + input_flat = self.quantize_real_ratio * input.data + \ + (1 - self.quantize_real_ratio) * input_flat + return input_flat + + def update_fp16_ratio(self): + if self.q_mixed_fp16: + if self.quantize_real_ratio > 0: + self.quantize_real_ratio -= self.q_change_ratio + else: + self.quantize_real_ratio = 0.000 diff --git a/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d8cb728da375e05de5023b4196d686c25e1c4e63 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) The DeepSpeed Contributors +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team diff --git a/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3cf130dc69ca0d1f539ecf9275cded4bb19b27d9 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b3117f49d7b1e1290a21ec97e3d200983e39640 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/parallel_state_sp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95f8e468df2352f1554f52a81dc5289f50bf888f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/__pycache__/ulysses_sp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py new file mode 100644 index 0000000000000000000000000000000000000000..c16880e15d42467f29f8103746f175b2698f4e6c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/parallel_state_sp.py @@ -0,0 +1,90 @@ +# Copyright (c) The DeepSpeed Contributors +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +This is a slimmed-down version of parallel_state.py (mpu) from Megatron-Deepspeed +""" + +from deepspeed import comm as dist + +# Sequence parallel groups to handle both data and sequence parallelisms. +# These groups are used to reduce gradients and shard parameters and optimizer stages for ZeRO. +_SEQUENCE_PARALLEL_GROUP = None +_SEQUENCE_DATA_PARALLEL_GROUP = None + + +def initialize_sequence_parallel(sequence_parallel_size: int) -> None: + """Initialize sequence parallel groups.""" + + assert dist.is_initialized() + world_size: int = dist.get_world_size() + + if world_size < sequence_parallel_size: + raise RuntimeError(f"world_size ({world_size}) is less than sequence_parallel_size {sequence_parallel_size}") + + if sequence_parallel_size <= 1: + raise ValueError(f"sequence_parallel_size must be greater than 1, got {sequence_parallel_size}") + + if world_size % sequence_parallel_size != 0: + raise RuntimeError( + f"world_size ({world_size}) is not divisible by sequence_parallel_size {sequence_parallel_size})") + + data_parallel_size: int = world_size // sequence_parallel_size + sequence_data_parallel_size: int = sequence_parallel_size * data_parallel_size + num_sequence_parallel_groups: int = world_size // sequence_parallel_size + num_sequence_data_parallel_groups: int = world_size // sequence_parallel_size // data_parallel_size + + rank = dist.get_rank() + + # Build the sequence parallel groups. + global _SEQUENCE_PARALLEL_GROUP + assert _SEQUENCE_PARALLEL_GROUP is None, "sequence parallel group is already initialized" + for i in range(num_sequence_parallel_groups): + ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) + group = dist.new_group(ranks) + if rank in ranks: + _SEQUENCE_PARALLEL_GROUP = group + + # Build the sequence data parallel groups. + global _SEQUENCE_DATA_PARALLEL_GROUP + assert _SEQUENCE_DATA_PARALLEL_GROUP is None, "sequence data parallel group is already initialized" + all_data_sequence_parallel_group_ranks = [] + for i in range(num_sequence_data_parallel_groups): + ranks = range(i * sequence_data_parallel_size, (i + 1) * sequence_data_parallel_size) + group = dist.new_group(ranks) + all_data_sequence_parallel_group_ranks.append(list(ranks)) + if rank in ranks: + _SEQUENCE_DATA_PARALLEL_GROUP = group + + +def get_sequence_parallel_group(): + """Get the sequence parallel group the caller rank belongs to.""" + assert _SEQUENCE_PARALLEL_GROUP is not None, "sequence parallel group is not initialized" + return _SEQUENCE_PARALLEL_GROUP + + +def get_sequence_data_parallel_group(): + """Get the sequence parallel group the caller rank belongs to.""" + assert _SEQUENCE_DATA_PARALLEL_GROUP is not None, "sequence data parallel group is not initialized" + return _SEQUENCE_DATA_PARALLEL_GROUP + + +def get_sequence_parallel_world_size(): + """Return world size for the sequence parallel group.""" + return dist.get_world_size(group=get_sequence_parallel_group()) + + +def get_sequence_data_parallel_world_size(): + """Return world size for the sequence parallel group.""" + return dist.get_world_size(group=get_sequence_data_parallel_group()) + + +def get_sequence_parallel_rank(): + """Return my rank for the sequence parallel group.""" + return dist.get_rank(group=get_sequence_parallel_group()) + + +def get_sequence_data_parallel_rank(): + """Return my rank for the sequence data parallel group.""" + return dist.get_rank(group=get_sequence_data_parallel_group()) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py new file mode 100644 index 0000000000000000000000000000000000000000..f649211d0284f9e654b5e66a6dc2c4bf5073a7f8 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/sequence_parallel/ulysses_sp.py @@ -0,0 +1,1226 @@ +# Copyright (c) The DeepSpeed Contributors +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Ulysses Plus features + +- `UlyssesSPAttentionHF` (port of UlyssesAttention from Megatron-Deepspeed plus modern MHA-variations) +- `UlyssesSPDataLoaderAdapter` - DL adapter to shard the normal DL batches to be used by `UlyssesSPAttentionHF` +- `SequenceTiledCompute` - generic autograd function to perform compute after tiling on the sequence dimension +- `TiledMLP` - a specific autograd function to perform tiled MLP (it's much easier to understand before trying to grok `SequenceTiledCompute`) + +The other UlyssesPlus features live inside https://github.com/snowflakedb/ArcticTraining (XXX: where exactly?) + +""" + +from collections import defaultdict +from deepspeed.runtime.utils import see_memory_usage +from deepspeed.sequence.layer import _DimZeroAllToAll +from einops import rearrange +from packaging import version +from torch import Tensor +from torch.utils.data import DataLoader +from typing import Any +from typing import Tuple +import deepspeed.comm as dist +import importlib.metadata +import math +import torch +import torch.distributed.nn + + +class UlyssesSPAttentionHF(torch.nn.Module): + """Re-Implementation of deepspeed.sequence.layer.DistributedAttention. This implementation enforces the input shape + to be standard [sl, bs, hc, hs] form. Any deviation from this shape will raise an error. + + The primary reason for the re-implementation is to make this less error prone, and remove what seemed like bugs in scenarios where batch size > 1 and when using different versions of + flash attention each of which takes different input shape. Those should be handled by + the actual attn implementation, and not by this module. + + This class then has been further adapted to work with HF Transformers' supported attention mechanism. + + Dimension annotation: + bs = bs + hc = head count + hc_l = head count local + hs = head_size + sl = seqlen + sl_l = seqlen local + ws = world_size + em = embedding (hidden size) + em_l = embedding (hidden size) local + + Arguments: + attn: normal attention implementation from transformers.modeling_utils.ALL_ATTENTION_FUNCTIONS + local_seq_length (int): local sequence length per GPU + global_seq_length (int): actual sequence length + batch_size (int): batch size + attn_head_size (int): size of each attention head + attn_head_count (int): total number of attention heads + kv_head_count (int): total number of kv heads + num_hidden_layers (int): total number of layers + process_group (dist.ProcessGroup): Ulysses process group + seq_length_is_variable (bool): whether global seqlen may change between batches + + + Extras: + - set self.skip_all_but_last_attention_debug_mode to True to enable fast debug which will skip calling all core attn layers but the last one, it will produce garbage of course quality-wise. + """ + + def __init__( + self, + attn, + local_seq_length: int, + global_seq_length: int, + batch_size: int, + attn_head_count: int, + attn_head_size: int, + kv_head_count: int, + num_hidden_layers: int, + process_group: dist.ProcessGroup, + seq_length_is_variable: bool = False, + ) -> None: + super().__init__() + self.attn = attn + self.process_group = process_group + self.world_size = dist.get_world_size(process_group) + self.sp_rank = dist.get_rank(process_group) + + self.local_seq_length = local_seq_length + self.global_seq_length = global_seq_length + self.batch_size = batch_size + self.seq_length_is_variable = seq_length_is_variable + + self.attn_head_size = attn_head_size + self.attn_head_count = attn_head_count + self.global_kv_head_count = kv_head_count + + self.num_hidden_layers = num_hidden_layers + self.skip_all_but_last_attention_debug_mode = False + self.rotating_layer_counter = 0 # used for dev work + + self.local_q_head_count = attn_head_count // self.world_size + + # if we have 4 kv heads and sp 8, we need to replicate kv heads 2x + self.kv_replication_factor = self.world_size // kv_head_count + if self.kv_replication_factor > 1: + self.local_kv_head_count = 1 + else: + self.local_kv_head_count = kv_head_count // self.world_size + + transformers_version_min = "4.51.3" + transformers_version_have = importlib.metadata.version("transformers") + if version.parse(transformers_version_have) < version.parse(transformers_version_min): + raise ValueError( + f"transformers>={transformers_version_min} is required, but you have transformers=={transformers_version_have}" + ) + + if self.attn_head_count % self.world_size != 0: + raise ValueError(f"Attention head count {attn_head_count} is not divisible by SP size {self.world_size}") + if not (self.global_kv_head_count % self.world_size == 0 or self.world_size % self.global_kv_head_count == 0): + raise ValueError( + f"KV attention head count {self.global_kv_head_count} is not divisible by SP size {self.world_size} or" + " vice versa") + + # [sl_l bs hc hs] + self.required_query_shape = torch.Size([local_seq_length, batch_size, attn_head_count, attn_head_size]) + self.required_key_value_shape = torch.Size([local_seq_length, batch_size, kv_head_count, attn_head_size]) + + # [sl bs em_l] + self.required_context_shape = torch.Size( + [global_seq_length, batch_size, attn_head_size * attn_head_count // self.world_size]) + + def _combine_local_sequences(self, query, key, value) -> Tuple[Tensor, Tensor, Tensor]: + + def combine_sequence(input, head_type): + """ + expects inputs in shape: [sl_l bs hc hs] + returns output in shape: [sl bs hc_l hs] + + local_head_count could be different for k,v vs q if it's not an MHA situation + """ + if head_type == "q": + local_head_count = self.local_q_head_count + else: # kv + local_head_count = self.local_kv_head_count + + # MQA and some GQA cases: + if self.kv_replication_factor > 1: + # local_head_count *= self.kv_replication_factor + # replicate heads to the kv_replication_factor on hc dimension [sl_l bs hc hs] - so dim=2 + input = input.repeat_interleave(self.kv_replication_factor, dim=2) + + # [sl_l bs hc hs] -> [sl_l bs ws hc_l hs] + input = input.reshape( + [self.local_seq_length, self.batch_size, self.world_size, local_head_count, self.attn_head_size]) + + input = rearrange(input, "sl_l bs ws hc_l hs -> ws sl_l bs hc_l hs").contiguous() + + output = _DimZeroAllToAll.apply(self.process_group, input) + + # [ws sl_l bs hc_l hs] -> [sl bs hc_l hs] + output = output.reshape([self.global_seq_length, *output.shape[2:]]).contiguous() + + # [sl bs hc_l hs] + return output + + return ( + combine_sequence(query, head_type="q"), + combine_sequence(key, head_type="kv"), + combine_sequence(value, head_type="kv"), + ) + + def _partition_global_sequence(self, input) -> Tensor: + """ + expects input in shape: [sl bs em_l] + returns output in shape: [sl_l bs em] + """ + + # [sl bs em_l] -> [ws sl_l bs em_l] + input = input.reshape([ + self.world_size, + self.local_seq_length, + self.batch_size, + self.attn_head_size * self.attn_head_count // self.world_size, + ]).contiguous() + + output = _DimZeroAllToAll.apply(self.process_group, input) + output = rearrange(output, "ws sl_l bs em_l -> sl_l bs ws em_l") + + # [sl_l bs ws em_l] -> [sl_l bs em] + output = output.reshape([*output.shape[:2], -1]).contiguous() + + # [sl_l bs em] + return output + + def forward( + self, + module: torch.nn.Module, + query: Tensor, + key: Tensor, + value: Tensor, + attention_mask: Tensor, + *args: Any, + **kwargs: Any, + ) -> Tensor: + """forward + + Arguments: + query (Tensor): query input to the layer + key (Tensor): key input to the layer + value (Tensor): value input to the layer + attention_mask (Tensor): Attention mask + args: other args + + Returns: + * output (Tensor): context output + """ + # HF incoming shapes are: + # [batch_size, num_heads, seqlen, head_size] + # UlyssesSPAttentionHF expects: + # [seqlen, batch_size, num_heads, head_size] + # print_rank0(f"{query.shape=}") + # print_rank0(f"{key.shape=}") + # print_rank0(f"{value.shape=}") + # print_rank0(f"{self.required_input_shape=}") + current_local_seq_length = query.shape[2] + if self.seq_length_is_variable and current_local_seq_length != self.required_query_shape[0]: + self.local_seq_length = current_local_seq_length + self.global_seq_length = current_local_seq_length * self.world_size + # update the required seqlen shapes + self.required_query_shape = torch.Size([self.local_seq_length] + list(self.required_query_shape)[1:]) + self.required_key_value_shape = torch.Size([self.local_seq_length] + + list(self.required_key_value_shape)[1:]) + self.required_context_shape = torch.Size([self.global_seq_length] + list(self.required_context_shape)[1:]) + + # make the blocks contiguous as early as possible to minimize fragmentation + query = rearrange(query, "bs hc sl hs -> sl bs hc hs") # .contiguous() + key = rearrange(key, "bs hc sl hs -> sl bs hc hs") # .contiguous() + value = rearrange(value, "bs hc sl hs -> sl bs hc hs") # .contiguous() + + # core attn like FA2 expects an unsharded `position_ids` - without which packed samples + # will return loss=nan. + # + # XXX: need to figure out if we can do the same for SDPA - as it doesn't require this and + # wants an attention mask, so possibly doing this for FA2 only? + # + # Ideally we would passing the original unsharded position_ids - but we have no way to pass + # it here as HF Transformers drops unexpected keys in `batch` - so either we need to stash + # it somewhere in UlyssesSPDataLoaderAdapter and retrieve it here or we could gather it once + # per batch and stash it inside `module` arg - I already have a machinery to figure out + # which layer number is being called below in the skip_all_but_last_attention_debug_mode + # code where rotating_layer_counter is used - so we could calculate it on the first layer + # and re-use on the remaining layers + if "position_ids" in kwargs: + position_ids_list = [torch.empty_like(kwargs["position_ids"]) for _ in range(self.world_size)] + dist.all_gather(position_ids_list, kwargs["position_ids"], group=self.process_group) + kwargs["position_ids"] = torch.cat(position_ids_list, dim=1) + + # please don't remove the white-space vertical alignment in the error message + assert query.shape == self.required_query_shape, ( + f"[{dist.get_rank()}]: query input tensor does not match the required shape\n " + f" {self.required_query_shape}:\n {query.shape=}\n {key.shape=}\n {value.shape=}") + assert key.shape == value.shape == self.required_key_value_shape, ( + f"[{dist.get_rank()}]: key or value input tensor does not match the required shape\n " + f" {self.required_key_value_shape}:\n {query.shape=}\n {key.shape=}\n {value.shape=}") + + # expects: [sl_l bs hc hs] + query_layer, key_layer, value_layer = self._combine_local_sequences(query, key, value) + # returns: [sl bs hc_l hs] + + query_layer = rearrange(query_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() + key_layer = rearrange(key_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() + value_layer = rearrange(value_layer, "sl bs hc_l hs -> bs hc_l sl hs").contiguous() + + # crucial in the case of MQA and some GQA cases we need to fix `module.num_key_value_groups` + # XXX: could move this somewhere to do it only once per run + if self.kv_replication_factor > 1: + module.num_key_value_groups = query_layer.size(-3) // key_layer.size(-3) + + if not self.skip_all_but_last_attention_debug_mode: + # expects: [bs hc_l sl hs] + context_layer, attn_weights = self.attn(module, query_layer, key_layer, value_layer, attention_mask, *args, + **kwargs) + # returns [bs sl hc_l hs] + else: + # we need this hack during development in order to be able to check memory fitting w/o + # waiting for 3h to compute 1.5M seqlen attention, because it's quadratic in dense + # attention, so we skip all but the last core attention call - we want the last one to + # still get the memory usage approximately close to the real memory usage. of course + # the loss will be wrong when we do that. + self.rotating_layer_counter = (self.rotating_layer_counter + 1) % self.num_hidden_layers + # we detect the last layer by module counting since we know how many layers there are + if self.rotating_layer_counter % self.num_hidden_layers == 0: + # do the real pass + context_layer, attn_weights = self.attn(module, query_layer, key_layer, value_layer, attention_mask, + *args, **kwargs) + else: + # this feeds bogus data of the right shape - good enough for quick debug + context_layer = rearrange(query_layer, "bs hc_l sl ... -> bs sl hc_l ...") + attn_weights = None + + # [bs sl hc_l hs] -> [sl bs hc_l hs]' + context_layer = rearrange(context_layer, "bs sl ... -> sl bs ...") + context_layer = context_layer.reshape([*context_layer.shape[:2], -1]) + + assert ( + context_layer.shape == self.required_context_shape + ), f"The context shape {context_layer.shape} is not of the expected shape {self.required_context_shape}" + + # expects: [sl bs em_l] + output = self._partition_global_sequence(context_layer) + # returns: [sl_l bs em] + + output = rearrange(output, "sl_l bs ... -> bs sl_l ...") + + output = output.reshape([*output.shape[:2], -1]) + + # expects [bs sl em] + return output, attn_weights + + @classmethod + def register_with_transformers( + cls, + model_name_or_path, + core_attn_implementation, + sequence_parallel_size, + max_length, + micro_batch_size, + seq_length_is_variable=True, + ): + """ + Register "ulysses" attn_implementation with HF transformers and return mpu (Megatron-LM-style parallel state object). + If sequence_parallel_size==1 do nothng and return None. + + """ + if sequence_parallel_size == 1: + return None + + from transformers import AutoConfig + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + import deepspeed.runtime.sequence_parallel.parallel_state_sp as mpu + + mpu.initialize_sequence_parallel(sequence_parallel_size=sequence_parallel_size) + + # we don't have the model yet at this stage + hf_model_config = AutoConfig.from_pretrained(model_name_or_path) + if core_attn_implementation not in ["flash_attention_2", "sdpa"]: + # notes on the excluded ones: + # - eager: The problem is that `eager` wants an attention_mask and it creates the wrong attention mask it seems if we don't provide one - it's possible that we could somehow solve this, but it's also unlikely someone will want to use the slow eager attention with sequence parallelism + # - flex_attention: haven't tried + + raise ValueError( + f"{core_attn_implementation} attn_implementation isn't currently supported by Ulysses sequence" + " parallelism. Set core_attn_implementation arg to either 'flash_attention_2' or 'sdpa'.") + + if core_attn_implementation not in ALL_ATTENTION_FUNCTIONS: + raise ValueError( + f"{core_attn_implementation} is not a valid attn_implementation. The choices are {ALL_ATTENTION_FUNCTIONS.valid_keys()}" + ) + core_attn_function = ALL_ATTENTION_FUNCTIONS[core_attn_implementation] + uattn = UlyssesSPAttentionHF( + attn=core_attn_function, + local_seq_length=max_length // mpu.get_sequence_parallel_world_size(), + global_seq_length=max_length, + batch_size=micro_batch_size, + attn_head_count=hf_model_config.num_attention_heads, + attn_head_size=getattr(hf_model_config, "head_dim", + hf_model_config.hidden_size // hf_model_config.num_attention_heads), + kv_head_count=hf_model_config.num_key_value_heads, + num_hidden_layers=hf_model_config.num_hidden_layers, + process_group=mpu.get_sequence_parallel_group(), + seq_length_is_variable=seq_length_is_variable, + ) + + def uattn_wrapper( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + *args, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + + # We are relaying on position_ids for SP to work so attention_mask has to be None + # the problem is that HF currently doesn't know anything about ALL_ATTENTION_FUNCTIONS["ulysses"] so it doesn't make a special case like for "flash_attention_2" and "sdpa" and it creates an attention mask on the fly and it breaks things. + attention_mask = None + + attn_output, attn_weights = uattn( + module, + query, + key, + value, + attention_mask, + # XXX: fixme + *args, + **kwargs, + ) + return attn_output, attn_weights + + # We don't do: ALL_ATTENTION_FUNCTIONS.register("ulysses", uattn_wrapper) + # The problem with this approach is that we are missing on all the special use cases in HF Transformers that do things like: if self.config._attn_implementation == "flash_attention_2": ... + # So instead we hack `ALL_ATTENTION_FUNCTIONS` to override all existing keys with our implementation, since it only gets used at the point of calling the attention and that's what we want, all other code branches relying on the original core `attn_implementation` will still be executed. This is what we called "Being John Malkovich" + for key in ALL_ATTENTION_FUNCTIONS.keys(): + ALL_ATTENTION_FUNCTIONS[key] = uattn_wrapper + + return mpu + + +class UlyssesSPDataLoaderAdapter: + + def __init__( + self, + dl: DataLoader, + sp_rank: int, + sp_group, + sp_world_size, + device, + ): + """ + This a DataLoader adapter which wraps around any existing DataLoader. It is used in conjunction with Ulysses to perform batch sharding on the sequence dimension. + + It gathers 1 sample from each participating rank, using the DL it wraps, then shards each of them and sends back to the ranks. So that when dl->iter->next is called, we end up with: + - rank 0: getting batch 0 shard 0 + - rank 1: getting batch 0 shard 1 + ... + - rank n: getting batch 0 shard n + which is used to compute the batch (from rank0) using all SP ranks. + + When the next iteration starts and dl->iter->next is called, we end up with: + - rank 0: getting batch 1 shard 0 + - rank 1: getting batch 1 shard 1 + ... + - rank n: getting batch 1 shard n + which is used to compute a second batch (from rank1) using all SP ranks. + + This continues until SP iterations are performed. At this point we need to get more data and so the above repeats. + + The key thing to understand is that all SP ranks participate in processing a single DL sample. So instead of normal DataParallel we perform a sort of SP over DP. + + When SP number of iterations is completed it's an equivalent of performing a single iteration with normal DP. + + If more tokens need to be consumed per step use the gradient accumulation feature. + + Arguments: + - `dl`: an existing DataLoader object to wrap + - `sp_rank`: SP rank + - `sp_group`: SP group + - `sp_world_size`: SP world size + - `device`: cuda device + + Returns: + Another DataLoader object + + Here are the current assumptions on the inputs fetched by dl->iter->next + - the batch is a dict with at least the keys: `input_ids`, `labels`, `position_ids` - but can have any additional keys necessary. + - the tensor values get sharded, the non-tensor values are passed along as is + """ + + self.dl = dl + self.sp_rank = sp_rank + self.sp_group = sp_group + self.sp_world_size = sp_world_size + self.device = device + + self.iter = iter(dl) + self.micro_batches: list[Any] = [] + + def __len__(self): + return len(self.dl) * self.sp_world_size + + def __iter__(self): + return self + + def __next__(self): + if len(self.micro_batches) == 0: + self.refill() + + return self.micro_batches.pop(0) + + def refill(self): + # this will raise StopIteration when empty + batch = next(self.iter) + micro_batches = defaultdict(dict) + # XXX: replace with more efficient all-to-all? + + # we have batches of variable seqlen so in order to do all_gather on batches - we need to know the exact length of each tensor on each rank + seqlen = torch.tensor(batch["input_ids"].shape[1], dtype=torch.int64, device=self.device) + seqlens = [torch.zeros(1, dtype=torch.int64, device=self.device) for _ in range(self.sp_world_size)] + dist.all_gather(seqlens, seqlen, group=self.sp_group) + seqlens = [x[0].item() for x in seqlens] + + for k in batch.keys(): + if torch.is_tensor(batch[k]): + batch[k] = batch[k].to(self.device) + with torch.no_grad(): + tensor_list = [ + torch.zeros((batch[k].shape[0], seqlens[i]), dtype=batch[k].dtype, device=batch[k].device) + for i in range(self.sp_world_size) + ] + dist.all_gather(tensor_list, batch[k], group=self.sp_group) + else: + tensor_list = [None for _ in range(self.sp_world_size)] + dist.all_gather_object(tensor_list, batch[k], group=self.sp_group) + + for rank, tensor in enumerate(tensor_list): + micro_batches[rank][k] = tensor + + del tensor_list + del batch + + for batch in micro_batches.values(): + seq_length = len(batch["input_ids"][0]) + + if seq_length % self.sp_world_size != 0: + raise ValueError(f"batch's seqlen={seq_length} isn't divisible by sp-size={self.sp_world_size}") + chunk_len = seq_length // self.sp_world_size + + # because we have to gather logits from all sp ranks we have to do the loss function ourselves + # therefore remove labels to avoid an attempt to calculate loss by transformers + labels = batch.pop("labels") + labels = torch.nn.functional.pad(labels, (0, 1), value=-100) + batch["shift_labels"] = labels[..., 1:].contiguous() + # free up temp memory + del labels + + # batch sharding + for k in batch.keys(): + # leave non-tensors alone + if not torch.is_tensor(batch[k]): + continue + # at seqlen>10M and 32+ gpus this can take GBs of memory so keep the prefill buffer on cpu + batch[k] = batch[k][:, chunk_len * self.sp_rank:chunk_len * (self.sp_rank + 1)].cpu() + + self.micro_batches.append(batch) + + +def sequence_tiled_compute( + fn, + seqlen, + shards, + kwargs_to_shard, + kwargs_to_pass, + grad_requiring_tensor_key, + compute_params=None, + output_unshard_dimension=1, + output_reduction="mean", +): + """ + This is a wrapper for SequenceTiledCompute which we need since torch.autograd.Function can't work with dicts of tensors (in backward it has to return a grad value and not a dict that may have a non-None grad value). It's also useful for setting default values which we can't do either in torch.autograd.Function. + + Args: + - `fn`: the function to call on sharded inputs + - `seqlen`: total seqlen of the seqlen dimension + - `shards`: how many shards to use + - `kwargs_to_shard`: this dict will be passed to `fn` as `**kwargs` after sharding on seqlen dimension + - `kwargs_to_pass`: this dict will be passed to `fn` as is, as `**kwargs` + - `grad_requiring_tensor_key`: which main key requires grads + - `compute_params`: a list of weights engaged in the compute. Default: `None` (only needed when using DeepSpeed ZeRO) + - `output_reduction`: None, "mean" or "sum": Default: "mean" + - `output_unshard_dimension`: the dimension to concat the outputs on: Default: 1 (seqlen dim) + + Returns: + - unsharded output with an optional reduction applied, depending on the `output_reduction` value: + `None` - return the unsharded output tensor + `"mean"` - apply mean + `"sum"` - apply sum + + Please note that this implementation doesn't require DeepSpeed and can work without it. `compute_params` can remain `None` in such a case. + + """ + args_to_shard = kwargs_to_shard.values() + keys_to_shard = list(kwargs_to_shard.keys()) + args_to_pass = kwargs_to_pass.values() + keys_to_pass = list(kwargs_to_pass.keys()) + + return SequenceTiledCompute.apply( + fn, + seqlen, + shards, + keys_to_shard, + keys_to_pass, + grad_requiring_tensor_key, + compute_params, + output_unshard_dimension, + output_reduction, + *args_to_shard, + *args_to_pass, + ) + + +class SequenceTiledCompute(torch.autograd.Function): + """ + A generic autograd function to perform a tiled compute. + + Please note that this implementation doesn't require DeepSpeed and can work without it. `compute_params` can remain `None` in such a case. + + For an easier to understand example see TiledMLP - which is the same as this autograd function but without the generalization code. + """ + + @staticmethod + def forward( + ctx, + fn, + seqlen, + shards, + keys_to_shard, + keys_to_pass, + grad_requiring_tensor_key, + compute_params, + output_unshard_dimension, + output_reduction, + *args, + ) -> torch.Tensor: + """ + for args and return values see `sequence_tiled_compute`'s doc + + Currently we assume that all kwargs_to_shard values have a shape of `[bs, seqlen, ...]` and we shard on seqlen dimension + """ + ctx.fn = fn + ctx.seqlen = seqlen + ctx.shards = shards + ctx.grad_requiring_tensor_key = grad_requiring_tensor_key + ctx.compute_params = [p for p in compute_params if p.requires_grad] + ctx.output_unshard_dimension = output_unshard_dimension + + with torch.no_grad(): + args = list(args) + ctx.total_args = len(args) + ctx.grad_requiring_tensor_key_index = (keys_to_shard + keys_to_pass).index(grad_requiring_tensor_key) + + kwargs_to_shard = {k: args.pop(0) for k in keys_to_shard} + kwargs_to_pass = {k: args.pop(0) for k in keys_to_pass} + ctx.kwargs_to_shard = kwargs_to_shard + ctx.kwargs_to_pass = kwargs_to_pass + + with torch.no_grad(): + shard_step = math.ceil(seqlen / shards) + output_shards = [] + + for i in range(shards): + output = fn( + **{ + k: v[:, i * shard_step:(i + 1) * shard_step] + for k, v in kwargs_to_shard.items() + }, + **kwargs_to_pass, + ) + output_shards.append(output) + + if output_unshard_dimension == 0: + # this is just the shape=[1] loss use-case, not sure if it's generic enough + output_unsharded = torch.cat([l.unsqueeze(0) for l in output_shards], dim=output_unshard_dimension) + else: + output_unsharded = torch.cat(output_shards, dim=output_unshard_dimension) # .clone().detach() + + if output_reduction is None: + return output_unsharded + elif output_reduction == "mean": + return output_unsharded.mean() + elif output_reduction == "sum": + return output_unsharded.sum() + else: + raise ValueError(f"unknown value {output_reduction}: valid values are: none/mean/sum") + + @staticmethod + def backward(ctx, *grads) -> torch.Tensor: + fn = ctx.fn + shards = ctx.shards + kwargs_to_shard = ctx.kwargs_to_shard + kwargs_to_pass = ctx.kwargs_to_pass + + grad_requiring_tensor_key = ctx.grad_requiring_tensor_key + grad_requiring_tensor_key_index = ctx.grad_requiring_tensor_key_index + compute_params = ctx.compute_params + output_unshard_dimension = ctx.output_unshard_dimension + grad_requiring_tensor = kwargs_to_shard[grad_requiring_tensor_key] + + grad_requiring_tensor_requires_grad = grad_requiring_tensor.requires_grad + grad_requiring_tensor = grad_requiring_tensor.detach() + # detach() unsets `grad_requiring_tensor.requires_grad`, so restore it + grad_requiring_tensor.requires_grad_(grad_requiring_tensor_requires_grad) + + incoming_grad = grads[0] + grad_requiring_tensor_grad = torch.zeros_like(grad_requiring_tensor) + + kwargs_to_shard_shards = { + k: list(torch.chunk(kwargs_to_shard[k], chunks=shards, dim=1)) + for k in kwargs_to_shard.keys() + } + + # if seqlen is not exactly divisible by shards the last step will be shorter than shard_step + shard_step = kwargs_to_shard_shards[grad_requiring_tensor_key][0].numel() + for i in range(shards): + + # when fn involves one or more model weights deepspeed will normally push a grad to + # reduce per sub-module call, so since we only want it to add a grad for the last + # shard's call , we signal to zero not to add new gradients to reduce until the last + # shard when all gradients have been accumulated. An example for such a call is + # `model.lm_head(hidden_states)` + if compute_params is not None: + if i + 1 < shards: + for param in compute_params: + param.ds_grad_is_ready = False + else: + # last shard, can add the grad + for param in compute_params: + param.ds_grad_is_ready = True + + kwargs_to_shard_shard = {k: kwargs_to_shard_shards[k].pop(0) for k in kwargs_to_shard_shards.keys()} + grad_requiring_tensor_shard = kwargs_to_shard_shard[grad_requiring_tensor_key] + + grad_requiring_tensor_shard.requires_grad_(grad_requiring_tensor_requires_grad) + + shard_offset = i * shard_step + # this will enable gradual population of the pre-allocated + # `grad_requiring_tensor_shard.grad` during `torch.autograd.backward` calls + grad_requiring_tensor_shard.grad = (grad_requiring_tensor_grad.view(-1).narrow( + 0, shard_offset, grad_requiring_tensor_shard.numel()).view_as(grad_requiring_tensor_shard)) + + with torch.enable_grad(): + output = fn(**kwargs_to_shard_shard, **kwargs_to_pass) + + if output_unshard_dimension == 0: + # loss use-case + torch.autograd.backward(output, incoming_grad) + else: + incoming_grad_shard = (incoming_grad.view(-1).narrow( + 0, shard_offset, grad_requiring_tensor_shard.numel()).view_as(grad_requiring_tensor_shard)) + torch.autograd.backward(output, incoming_grad_shard) + + # positional args + grad_outputs = [None] * 9 + # inject the grad for the position of forward input that is grad-requiring + arg_outputs = [None] * ctx.total_args + arg_outputs[grad_requiring_tensor_key_index] = grad_requiring_tensor_grad + + return tuple(grad_outputs + arg_outputs) + + +class TiledMLP(torch.autograd.Function): + """ + Perform a tiled MLP computation to massively reduce memory usage needed to compute MLP when using very long sequence lengths + + For a general tiled compute implementation that can handle any `forward` see `SequenceTiledCompute` + + Args: + - fn: the function to call on sharded inputs + - `self`: the MLP nn.Module object + - `x`: the input to MLP.forward (`hidden_states`) + - `shards`: how many shards to use + - compute_params: a list of weights engaged in the compute Default: `None` (only needed when using DeepSpeed ZeRO) + + Returns: + - the computed `hidden_states` + + Here is an example that monkey patches HF Transformers' LLamaMLP: + + def tiled_mlp_forward(self, x): + bs, seqlen, hidden = x.shape + num_shards = math.ceil(seqlen / hidden) + # to avoid deadlocks get all ranks to agree on the same num_shards by using the max value + tensor = torch.tensor(num_shards, device=x.device) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + num_shards = tensor.item() + compute_params = [self.down_proj.weight, self.gate_proj.weight, self.up_proj.weight] + + def mlp_forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + return TiledMLP.apply( + mlp_forward, + self, + x, + num_shards, + compute_params, + ) + + # this needs to be done before the model is instantiated + from transformers.models.llama import modeling_llama + modeling_llama.LlamaMLP.forward = tiled_mlp_forward + """ + + @staticmethod + def forward( + ctx, + fn, + self, + x, + shards, + compute_params, + ) -> torch.Tensor: + ctx.fn = fn + ctx.self = self + ctx.shards = shards + ctx.compute_params = [p for p in compute_params if p.requires_grad] + ctx.save_for_backward(x) + + x_shards = list(torch.chunk(x, chunks=shards, dim=1)) + with torch.no_grad(): + output_shards = [fn(self, x_shard) for x_shard in x_shards] + output_unsharded = torch.cat(output_shards, dim=1) + + return output_unsharded + + @staticmethod + def backward(ctx, *grads) -> torch.Tensor: + fn = ctx.fn + (x, ) = ctx.saved_tensors + self = ctx.self + shards = ctx.shards + compute_params = ctx.compute_params + + x_requires_grad = x.requires_grad + x = x.detach() + # detach() unsets `x.requires_grad`, so restore it + x.requires_grad_(x_requires_grad) + + incoming_grad = grads[0] + x_grad = torch.zeros_like(x) + x_shards = list(torch.chunk(x, chunks=shards, dim=1)) + + shard_step = x_shards[0].numel() + for i, x_shard in enumerate(x_shards): + + # Tell deepspeed not to add a new grad to its ipg bucket until the last shard is run + if compute_params is not None: + if i + 1 < shards: + for param in compute_params: + param.ds_grad_is_ready = False + else: + # last shard, can add the grad + for param in compute_params: + param.ds_grad_is_ready = True + + x_shard.requires_grad_(x_requires_grad) + + shard_offset = i * shard_step + x_shard.grad = x_grad.view(-1).narrow(0, shard_offset, x_shard.numel()).view_as(x_shard) + incoming_grad_shard = incoming_grad.view(-1).narrow(0, shard_offset, x_shard.numel()).view_as(x_shard) + with torch.enable_grad(): + output = fn(self, x_shard) + torch.autograd.backward(output, incoming_grad_shard) + + return (None, None, x_grad, None, None) + + +class AutogradComputeMLP(torch.autograd.Function): + """ + This is a simplified example to override the normal MLP via an autograd function - then tiling can be added - this simplified version was useful to detect a leak in Deepspeed, so let's keep it. + + Here is an example of performing the monkey patching on LlamaMLP + + def mlp_forward_new(self, x): + + def mlp_forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + return AutogradComputeMLP.apply(mlp_forward, self, x) + + from transformers.models.llama import modeling_llama + modeling_llama.LlamaMLP.forward = mlp_forward_new + """ + + @staticmethod + def forward( + ctx, + fn, + self, + x, + ) -> torch.Tensor: + ctx.fn = fn + ctx.self = self + ctx.save_for_backward(x) + + with torch.no_grad(): + return fn(self, x) + + @staticmethod + def backward(ctx, *grads) -> torch.Tensor: + fn = ctx.fn + (x, ) = ctx.saved_tensors + self = ctx.self + + x1 = x.detach() + x1.requires_grad = x.requires_grad + with torch.enable_grad(): + output = fn(self, x1) + + torch.autograd.backward(output, grads[0]) + return (None, None, x1.grad, None) + + +########################################################### +### below are older versions that some might still want ### +########################################################### + + +class TiledLoss(torch.autograd.Function): + + @staticmethod + def forward(ctx, loss_fn, logits, vocab_size, shift_labels, shards) -> torch.Tensor: + """ + + This is a memory efficient loss autograd function that takes the existing logits and performs loss calculation in shards. + + This one is an SFT-aware version, therefore it takes care of special cases where the whole shard is made of -100 labels and which requires then a special care. + + Note: logits seqlen dimension doesn't have to be divisible by shards, the last shard will be shorter than the rest. The calculating of the number of shards is in the example. + + Here is an example of using it: + + def loss(self, batch) -> torch.Tensor: + batch = to_device(batch, self.device) + shift_labels = batch.pop("shift_labels") + outputs = self.model(**batch, use_cache=False) + logits = outputs.logits + + if all((shift_labels == -100).squeeze()): + # this is the case where all labels in a micro-batch are -100 (very common for SFT if the seqlen is short) - CE returns `nan` in this case, so we don't want to call loss and instead create a differentiable loss `0` which will also set all the grads to `0` in `backward` - the effect of this is akin to a perfect score where the model needs no adjustment since grads will be all zeros. + loss = (logits.sum() * 0.0).float() + + num_shards: Any = "auto" + if num_shards == "auto": + # parameterize to about 1GB fp32 logits shards + slice_size_in_gb = 1 + size_in_gb = logits.numel() * 4 / 2**30 # fp32 + # the sp shard's seqlen sp shard can be easily not divisible by the derived number of chunked loss shards, so we use the uppper ceiling and allow the last chunk to be shorter than the rest + num_shards = math.ceil(size_in_gb / slice_size_in_gb) + # print(f"derived {num_shards} shards for size {size_in_gb}GB") + if num_shards > 1: + # if shards == 1 this will lead to a higher memory usage then calling the normal loss function, so don't do that. + loss = TiledLoss.apply( + self.model_unwrapped.loss_function, + logits, + self.model_unwrapped.config.vocab_size, + shift_labels, + num_shards, + ) + else: + loss = self.model_unwrapped.loss_function( + logits=logits, + labels=None, + vocab_size=self.model_unwrapped.config.vocab_size, + shift_labels=shift_labels, + ) + + return loss + + + """ + ctx.save_for_backward(logits, shift_labels) + ctx.loss_fn = loss_fn + ctx.vocab_size = vocab_size + ctx.shards = shards + + with torch.no_grad(): + seqlen = shift_labels.shape[1] + shard_step = math.ceil(seqlen / shards) + loss_shards = [] + total_good_items = 0 + + # since -100s are ignored we have to perform a weighted average on each loss slice as each slice may contribute a different number of non- -100 labels + # if seqlen / shards != 0 - the last chunk is just shorter than the rest but no data is ignored + for i in range(shards): + # XXX: here and everywhere don't make a copy, pass the slice or perhaps narrow/view? + shift_labels_shard = shift_labels[:, i * shard_step:(i + 1) * shard_step] + if all((shift_labels_shard == -100).squeeze()): + continue # ignore this shard + loss_shard = loss_fn( + logits=logits[:, i * shard_step:(i + 1) * shard_step, :], + labels=None, + vocab_size=vocab_size, + shift_labels=shift_labels_shard, + ) + good_items = sum((shift_labels_shard != -100).squeeze()) + loss_shards.append(loss_shard * good_items) + total_good_items += good_items + total_loss = torch.cat([l.unsqueeze(0) for l in loss_shards], dim=0).sum() + weighted_loss = total_loss / total_good_items + + return weighted_loss + + @staticmethod + def backward(ctx, *grads) -> torch.Tensor: + logits, shift_labels = ctx.saved_tensors + loss_fn = ctx.loss_fn + vocab_size = ctx.vocab_size + shards = ctx.shards + + grad = grads[0] + logits_grad = torch.zeros_like(logits) + logits_shards = list(torch.chunk(logits, chunks=shards, dim=1)) + shift_labels_shards = list(torch.chunk(shift_labels, chunks=shards, dim=1)) + + # if seqlen is not exactly divisible by shards the last step will be shorter than shard_step + shard_step = logits_shards[0].numel() + for i in range(shards): + logits_shard = logits_shards.pop(0) + shift_labels_shard = shift_labels_shards.pop(0) + + shard_offset = i * shard_step + # this will enable gradual population of the pre-allocated `logits_shard.grad` during `torch.autograd.backward` calls + logits_shard.grad = (logits_grad.view(-1).narrow(0, shard_offset, + logits_shard.numel()).view_as(logits_shard)) + + with torch.enable_grad(): + if all((shift_labels_shard == -100).squeeze()): + # fake loss calculation, since CE will return nan, but grads will be set + # a normal loss_fn upcasts logits to float so match it + loss_shard = (logits_shard.sum() * 0.0).float() + else: + loss_shard = loss_fn( + logits=logits_shard.requires_grad_(), + labels=None, + vocab_size=vocab_size, + shift_labels=shift_labels_shard, + ) + + torch.autograd.backward(loss_shard, grad) + + logits_grad /= shards + + # only logits (2nd arg) needs grads + return None, logits_grad, None, None, None + + +# This is the original implementation/integration of UlyssesSP into the training loop, which was superseded by using UlyssesSPDataLoaderAdapter which did all the sharding and pull the shards from the DL +# +# There are 2 issues with this implementation: +# - it's complex and difficult to integrate into various training scenarios +# - it could lead to a huge number of tokens per step - e.g. 32 ranks of 15M seqlen -> 0.5B token step - which is very wasteful +# +# Therefore if you want to use UlyssesSP via UlyssesSPFwdLossBwdWithLogits with its fwd/loss/bwd for those don't want to use UlyssesSPDataLoaderAdapter - here is how it should be installed into the sub-trainer class: +# class SFTTrainer(Trainer): +# def sp_fwd_loss_bwd(self, batch) -> torch.Tensor: +# batch = to_device(batch, self.device) +# +# from arctic_training.trainer.trainer import UlyssesAttentionHFFwdLossBwdWithLogits +# ulysses = UlyssesAttentionHFFwdLossBwdWithLogits( +# model=self.model, +# model_unwrapped=self.model_unwrapped, +# device=self.device, +# num_loss_logit_shards="auto", +# ) +# return ulysses.sp_fwd_loss_bwd(batch) + + +class UlyssesSPFwdLossBwdWithLogits: + + def __init__(self, model, model_unwrapped, device, num_loss_logit_shards="auto", **kwargs): + + self.model = model + self.model_unwrapped = model_unwrapped + self.device = device + self.num_loss_logit_shards = num_loss_logit_shards + self.kwargs = kwargs + + from deepspeed.utils import groups + + self.sp_group = groups._get_sequence_parallel_group() + self.sp_world_size = groups._get_sequence_parallel_world_size() + self.sp_rank = groups._get_sequence_parallel_rank() + + def sp_fwd_loss_bwd(self, batch) -> torch.Tensor: + + see_memory_usage(f"entered sp_fwd_loss_bwd", force=True) + + # ensure shapes are correct + if not (batch["input_ids"].shape == batch["position_ids"].shape == batch["labels"].shape): + raise ValueError( + f'Borked batch {batch["input_ids"].shape=} != {batch["position_ids"].shape=} !=' + f' {batch["labels"].shape=}) in DataLoader->iter->next, cannot continue with Ulysses Sequence' + " parallelism") + + # gather DL batches into super-batches + # Important: DL doesn't always yield max_length batches. Different ranks may have different seqlen and each could be <= max_length (but always divisible by 256) + + micro_batches: list[Any] = defaultdict(dict) + # Efficient gathering of batch inputs across ranks: + # The problem is that our DL doesn't guarantee the same seqlen on all ranks and may give, 3x 1024 and 1x 768 on 4 gpus for max_length 1024. so 3 options we have to be able to gather batches are: + # 1. use all_gather_object - which allows different shapes - but potentially introducing an undesired overhead - 2x pickle calls + # 2. use all_gather and change DL pad to make sure that all ranks always get the same input shape - this creates its own overhead since if we say have ranks with seqlen 512, 768, 1024, 1024 - now we will need to process 4x 1024 seqlens + # 3. use all_gather and post gathering truncate tensors to their intended length - another overhead of allocating and truncating tensors + # using approach (1) for now but might want to benchmark later the other 2 approaches + + # XXX: if using all_gather_object we can gather the whole batch at once and not per-key! so can drop the loop for that approach + + # we have batches of variable seqlen so in order to do all_gather on batches - we need to know the exact length of each tensor on each rank + seqlen = torch.tensor(batch["input_ids"].shape[1], dtype=torch.int64, device=self.device) + # print(seqlen) + seqlens = [torch.zeros(1, dtype=torch.int64, device=self.device) for _ in range(self.sp_world_size)] + dist.all_gather(seqlens, seqlen, group=self.sp_group) + seqlens = [x[0].item() for x in seqlens] + + for k in batch.keys(): + batch[k] = batch[k].to(self.device) + with torch.no_grad(): + tensor_list = [ + torch.zeros((batch[k].shape[0], seqlens[i]), dtype=batch[k].dtype, device=batch[k].device) + for i in range(self.sp_world_size) + ] + dist.all_gather(tensor_list, batch[k], group=self.sp_group) + + # gathering on the data dimension + # will be concatenating and later splitting again for the more general case + # batch[k] = torch.cat(tensor_list, dim=1) + for rank, tensor in enumerate(tensor_list): + micro_batches[rank][k] = tensor + + del tensor_list + del batch + + # we need to chunk twice - each time on SP size level + # - the first time is because we artificially made the seqlen SP-times longer + # - the second time is because of the Ulysses algorithm + + see_memory_usage("after gathering", force=False) + + self.model.set_gradient_accumulation_boundary(False) + + losses = [] + for sub_step_id in range(self.sp_world_size): + batch = micro_batches[sub_step_id] + seq_length = len(batch["input_ids"][0]) + + if seq_length % self.sp_world_size != 0: + raise ValueError( + f"{sub_step_id=}: batch's seqlen={seq_length} isn't divisible by sp-size={self.sp_world_size}") + chunk_len = int(seq_length / self.sp_world_size) + + # to enable the correct mean calculation across shards before sharding the micro batch: + # 1. count the number of non- `-100`` elements per shard + # 2. and subtract one more element because of label shifting + non_skipped_items = {} + for rank in range(self.sp_world_size): + non_skipped = (batch["labels"][:, chunk_len * rank:chunk_len * (rank + 1)] != -100).sum().item() + if non_skipped > 1: + non_skipped -= 1 + non_skipped_items[rank] = non_skipped + + # because we have to gather logits from all sp ranks we have to do the loss function ourselves + # therefore remove labels to avoid an attempt to calculate loss by transformers + labels = batch.pop("labels") + labels = torch.nn.functional.pad(labels, (0, 1), value=-100) + batch["shift_labels"] = labels[..., 1:].contiguous() + # free up temp memory + del labels + + # batch sharding + for k in batch.keys(): + batch[k] = batch[k][:, chunk_len * self.sp_rank:chunk_len * (self.sp_rank + 1)].to(self.device) + + shift_labels = batch.pop("shift_labels") + + outputs = self.forward(batch) + loss = self.compute_loss(labels=None, shift_labels=shift_labels) + + # free up temp mem (e.g. outputs.logits are huge) + del outputs + + # differentiable loss aggregation across ranks + losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=self.sp_group) + + # since each shard may have a variable number of skipped elemented - need to calculate a weighted mean depending on each rank's contribution - this will also take care of loss=0 when all elements are -100 in a shard + # XXX: not expecting a total of 0-non-skipped items for div + loss = sum(losses_per_rank[rank] * non_skipped_items[rank] + for rank in range(self.sp_world_size)) / sum(non_skipped_items.values()) + + self.backward() + + losses.append(loss.detach().item()) + + self.model.set_gradient_accumulation_boundary(True) + + # for per-iteration reporting + if len(losses) == 0: + loss = float("nan") + else: + loss = sum(losses) / len(losses) + + return loss + + def forward(self, batch): + # critical: the labels shouldn't be in batch + outputs = self.model(**batch, use_cache=False) + self.logits = outputs.logits + return outputs + + def compute_loss(self, labels, shift_labels): + if all((shift_labels == -100).squeeze()): + # this is the case where all labels in a micro-batch are -100 (very common for SFT) - CE returns `nan` in this case, so we don't want to call loss and instead create a differentiable loss `0` which will also set all the grads to `0` in `backward` - the effect of this is akin to a perfect score where the model needs no adjustment since grads will be all zeros. + # XXX: should this be float and not the original dtype? + loss = (self.logits.sum() * 0.0).float() + else: + if self.num_loss_logit_shards == "auto": + # parameterize to about 1GB fp32 logits shards + slice_size_in_gb = 1 # XXX: make configurable? + size_in_gb = self.logits.numel() * 4 / 2**30 # fp32 + # the sp shard's seqlen sp shard can be easily not divisible by the derived number of chunked loss shards, so we use the uppper ceiling and allow the last chunk to be shorter than the rest + self.num_loss_logit_shards = math.ceil(size_in_gb / slice_size_in_gb) + # print(f"derived {self.num_loss_logit_shards} shards for size {size_in_gb}GB") + if self.num_loss_logit_shards > 1: + loss = TiledLoss.apply( + self.model_unwrapped.loss_function, + self.logits, + self.model_unwrapped.config.vocab_size, + shift_labels, + self.num_loss_logit_shards, + ) + else: + # XXX: for some reason this fails with zero1 + loss = self.model_unwrapped.loss_function( + logits=self.logits, + labels=None, + vocab_size=self.model_unwrapped.config.vocab_size, + shift_labels=shift_labels, + ) + + self.loss = loss + return loss + + def backward(self): + self.model.backward(self.loss) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/sparse_tensor.py b/lib/python3.12/site-packages/deepspeed/runtime/sparse_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..291ba5f0c78668b793bea12749bd0a188b680ece --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/sparse_tensor.py @@ -0,0 +1,69 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Implementation of a compressed sparse tensor. Similar in +functionality to TensorFlow's IndexedSlices implementation. +""" + +import torch + + +class SparseTensor(object): + """ Compressed Sparse Tensor """ + + def __init__(self, dense_tensor=None): + self.orig_dense_tensor = dense_tensor + self.dtype = self.orig_dense_tensor.dtype + self.is_sparse = dense_tensor.is_sparse + if dense_tensor is not None: + if dense_tensor.is_sparse: + dense_tensor = dense_tensor.coalesce() + self.indices = dense_tensor.indices().flatten() + self.values = dense_tensor.values() + else: + result = torch.sum(dense_tensor, dim=1) + self.indices = result.nonzero().flatten() + self.values = dense_tensor[self.indices] + self.dense_size = list(dense_tensor.size()) + else: + self.indices = None + self.values = None + self.dense_size = None + + def to_coo_tensor(self): + return torch.sparse_coo_tensor(self.indices.unsqueeze(0), self.values, self.dense_size) + + @staticmethod + def type(): + return "deepspeed.SparseTensor" + + def to_dense(self): + it = self.indices.unsqueeze(1) + full_indices = torch.cat([it for _ in range(self.dense_size[1])], dim=1) + return self.values.new_zeros(self.dense_size).scatter_add_(0, full_indices, self.values) + + def sparse_size(self): + index_size = list(self.indices.size()) + index_size = index_size[0] + value_size = list(self.values.size()) + value_size = value_size[0] * value_size[1] + dense_size = self.dense_size[0] * self.dense_size[1] + return index_size + value_size, dense_size + + def add(self, b): + assert self.dense_size == b.dense_size + self.indices = torch.cat([self.indices, b.indices]) + self.values = torch.cat([self.values, b.values]) + + def __str__(self): + sparse_size, dense_size = self.sparse_size() + return "DeepSpeed.SparseTensor(indices_size={}, values_size={}, " \ + "dense_size={}, device={}, reduction_factor={})".format( + self.indices.size(), self.values.size(), self.dense_size, + self.indices.get_device(), dense_size / sparse_size + ) + + def __repr__(self): + return self.__str__() diff --git a/lib/python3.12/site-packages/deepspeed/runtime/state_dict_factory.py b/lib/python3.12/site-packages/deepspeed/runtime/state_dict_factory.py new file mode 100644 index 0000000000000000000000000000000000000000..c2db85d1ba58f16dbe59b1dcd3d3066ab11ff6d2 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/state_dict_factory.py @@ -0,0 +1,427 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import os +import copy +import collections +import json +from abc import ABC, abstractmethod + +from deepspeed.utils import logger +from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine + +from .weight_quantizer import WeightQuantization + +AUTO_MODULE_KEY = 'auto' + + +class SDLoaderFactory: + + @staticmethod + def get_sd_loader_json(json_file, checkpoint_engine): + if isinstance(json_file, str): + with open(json_file) as f: + data = json.load(f) + else: + assert isinstance(json_file, dict) + data = json_file + sd_type = data['type'] + ckpt_list = data['checkpoints'] + version = data['version'] + ckpt_type = data.get('parallelization', 'pp') + mp_size = data.get('mp_size', 0) + if sd_type.lower() in ['bloom', 'ds_model']: + return data + return SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine, sd_type, version) + + @staticmethod + def get_sd_loader(ckpt_list, checkpoint_engine, sd_type='Megatron', version=None): + if sd_type == 'Megatron': + return MegatronSDLoader(ckpt_list, version, checkpoint_engine) + else: + assert False, '{} checkpoint type is not supported'.format(sd_type) + + +class SDLoaderBase(ABC): + + def __init__(self, ckpt_list, version, checkpoint_engine): + self.module_key = None + self.ckpt_list = ckpt_list + self.version = version + self.checkpoint_engine = TorchCheckpointEngine() if checkpoint_engine is None else checkpoint_engine + self.check_ckpt_list() + + def load(self, + mp_world_size, + mp_rank, + module_key=AUTO_MODULE_KEY, + is_pipe_parallel=False, + quantize=False, + quantize_bits=8, + quantize_groups=64, + mlp_extra_grouping=True): + self.module_key = module_key + num_ckpt = len(self.ckpt_list) + idx = mp_rank * num_ckpt // mp_world_size + """ We have multiple cases to handle here for both training and inference: + 1. PipeModule loading mp_rank_*.pt files, is_pipe_parallel=True, module_key is not None + a. if no mp_size/pp_size resizing occurs, for both training & inference, loading + the mp_rank related checkpoint directly. + b. if has mp_size/pp_size resizing, only Megatron model inference is supported, + in this case each mp_rank_*.pt have same content, we will load the first checkpoint + file (idx=0), to avoid idx exceeding file list boundary. + + 2. PipeModule loading layer_*.pt files, is_pipe_parallel=True, module_key is None + a. if no mp_size resizing occurs, for both training & inference, loading + the mp_rank related checkpoint directly. + b. if has mp_size resizing, only Megatron model inference is supported, + checkpoint file(s) will be merged/split according to mp_rank, mp_world_size and + checkpoint file list. + + 3. Non-PipeModule loading mp_rank_*.pt files, is_pipe_parallel=False + Same with case (2). + """ + if is_pipe_parallel and module_key is not None and mp_world_size != num_ckpt: + mp_world_size = num_ckpt + idx = 0 + + load_path = self.ckpt_list[idx] + + merge_count = 1 + if num_ckpt == mp_world_size: + assert os.path.exists(load_path) + #logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}') + sd = self.checkpoint_engine.load(load_path, map_location=lambda storage, \ + loc: storage) + + if quantize: + quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) + sd_module, all_scales = quantizer.sd_quantize_megatron(self.get_module(sd), quantize_bits, + quantize_groups) + self.set_module(sd, sd_module) + else: + all_scales = None + elif num_ckpt > mp_world_size: + sd, all_scales, merge_count = self.merge_state_dict(mp_world_size, mp_rank, quantize, \ + quantize_bits, quantize_groups, mlp_extra_grouping) + else: + sd, all_scales = self.split_state_dict(mp_world_size, mp_rank, quantize, quantize_bits, \ + quantize_groups, mlp_extra_grouping) + return load_path, sd, (all_scales, merge_count) + + def get_merge_state_dicts(self, mp_world_size, mp_rank): + num_ckpt = len(self.ckpt_list) + assert num_ckpt % mp_world_size == 0, 'Invalid checkpoints and world size for sd merge' + + num_to_merge = num_ckpt // mp_world_size + ckpt_list = [self.ckpt_list[i] for i in range(num_to_merge * mp_rank, num_to_merge * (mp_rank + 1))] + + logger.info(f"mp_rank: {mp_rank}, ckpt_list: {ckpt_list}") + sd_list = [self.checkpoint_engine.load(ckpt, map_location=lambda storage, loc: storage) for ckpt in ckpt_list] + return sd_list + + def get_split_state_dict(self, mp_world_size, mp_rank): + num_ckpt = len(self.ckpt_list) + assert mp_world_size % num_ckpt == 0, 'Invalid checkpoints and world size for sd split' + + num_to_split = mp_world_size // num_ckpt + ckpt_index = mp_rank // num_to_split + ckpt_offset = mp_rank % num_to_split + + logger.info(f"mp_rank: {mp_rank}, ckpt_list: {self.ckpt_list[ckpt_index]}, offset: {ckpt_offset}") + + sd = self.checkpoint_engine.load(self.ckpt_list[ckpt_index], map_location=lambda storage, loc: storage) + + return sd, num_to_split, ckpt_offset + + def _choose_module_key(self, sd): + assert not ('module' in sd + and 'model' in sd), "checkpoint has both 'model' and 'module' keys, not sure how to proceed" + assert 'module' in sd or 'model' in sd, "checkpoint contains neither 'model' or 'module' keys, not sure how to proceed" + if 'module' in sd: + return 'module' + elif 'model' in sd: + return 'model' + + def get_module(self, sd): + if self.module_key is None: + return sd + elif self.module_key == AUTO_MODULE_KEY: + return sd[self._choose_module_key(sd)] + else: + return sd[self.module_key] + + def set_module(self, sd, module): + if self.module_key is None: + sd = module + elif self.module_key == AUTO_MODULE_KEY: + sd[self._choose_module_key(sd)] = module + else: + sd[self.module_key] = module + return sd + + def check_ckpt_list(self): + #logger.info(f'checkpoint file list: {self.ckpt_list}') + assert len(self.ckpt_list) > 0 + + sd = self.checkpoint_engine.load(self.ckpt_list[0], map_location=lambda storage, loc: storage) + + # check checkpoint count is same with saved mp_world_size + if 'mp_world_size' in sd.keys(): + assert len(self.ckpt_list) == sd[ + 'mp_world_size'], f"checkpoint count {len(self.ckpt_list)} is different from saved mp_world_size {sd['mp_world_size']}" + + @abstractmethod + def merge_state_dict(self, mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping): + pass + + @abstractmethod + def split_state_dict(self, mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping): + pass + + @abstractmethod + def sanity_check(self, ckpt_file_name): + pass + + +class MegatronSDLoader(SDLoaderBase): + + def __init__(self, ckpt_list, version, checkpoint_engine): + super().__init__(ckpt_list, version, checkpoint_engine) + """ + ## Q/K/V data need special processing + key: transformer.layers.0.attention.query_key_value.weight, shape: torch.Size([3192, 4256]) + key: transformer.layers.0.attention.query_key_value.bias, shape: torch.Size([3192]) + + ## merge or split on axis=0 + key: word_embeddings.weight, shape: torch.Size([12672, 4256]) + key: transformer.layers.0.mlp.dense_h_to_4h.bias, shape: torch.Size([4256]) + key: transformer.layers.0.mlp.dense_h_to_4h.weight, shape: torch.Size([4256, 4256]) + + ## merge or split on axis=1 + key: transformer.layers.0.attention.dense.weight, shape: torch.Size([4256, 1064]) + key: transformer.layers.0.mlp.dense_4h_to_h.weight, shape: torch.Size([4256, 4256]) + + ## no change required + key: transformer.layers.0.mlp.dense_4h_to_h.bias, shape: torch.Size([4256]) + key: transformer.final_layernorm.weight, shape: torch.Size([4256]) + key: transformer.final_layernorm.bias, shape: torch.Size([4256]) + key: transformer.layers.0.attention.dense.bias, shape: torch.Size([4256]) + key: transformer.layers.0.post_attention_layernorm.weight, shape: torch.Size([4256]) + key: transformer.layers.0.post_attention_layernorm.bias, shape: torch.Size([4256]) + key: transformer.layers.0.input_layernorm.weight, shape: torch.Size([4256]) + key: transformer.layers.0.input_layernorm.bias, shape: torch.Size([4256]) + key: position_embeddings.weight, shape: torch.Size([1024, 4256]) + """ + + def merge_query_key_value(self, param_list, ckpt_ver): + """ + Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions: + + 1. version 0, there is no version information saved in checkpoint. + format: [(3 * np * hn), h] + 2. version 1.0 + format: [(np * hn * 3), h] + 3. version 2.0 + format: [(np * 3 * hn), h] + + h: hidden size + n: number of attention heads + p: number of model parallel partitions + np: n/p + hn: h/n + """ + + new_qkv = None + if ckpt_ver == 0: + # [(3 * np * hn), h] + assert param_list[0].shape[0] % 3 == 0 + size_qkv = param_list[0].shape[0] // 3 + split_tensors = [torch.split(param, size_qkv, dim=0) for param in param_list] + + tensors = [] + for i in range(3): + tensor_tuple = [t[i] for t in split_tensors] + tensors.append(torch.cat(tensor_tuple, axis=0)) + new_qkv = torch.cat(tensors, axis=0) + elif ckpt_ver == 1.0 or ckpt_ver == 2.0: + # [(np * hn * 3), h] or [(np * 3 * hn), h] + new_qkv = torch.cat(param_list, axis=0) + else: + assert False, f'checkpoint version: {ckpt_ver} is not supported' + + return new_qkv + + def split_query_key_value(self, param, num_to_split, offset, ckpt_ver): + """ + Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions: + + 1. version 0, there is no version information saved in checkpoint. + format: [(3 * np * hn), h] + 2. version 1.0 + format: [(np * hn * 3), h] + 3. version 2.0 + format: [(np * 3 * hn), h] + + h: hidden size + n: number of attention heads + p: number of model parallel partitions + np: n/p + hn: h/n + """ + + new_qkv = None + if ckpt_ver == 0: + # [(3 * np * hn), h] + assert param.shape[0] % 3 == 0 + size_qkv = param.shape[0] // 3 + split_tensors = torch.split(param, size_qkv, dim=0) + + assert split_tensors[0].shape[0] % num_to_split == 0 + split_size = split_tensors[0].shape[0] // num_to_split + + tensors = [] + for i in range(3): + tensors.append(torch.split(split_tensors[i], split_size, dim=0)[offset]) + new_qkv = torch.cat(tensors, axis=0) + elif ckpt_ver == 1.0 or ckpt_ver == 2.0: + # [(np * hn * 3), h] or [(np * 3 * hn), h] + assert param.shape[0] % num_to_split == 0 + size_qkv = param.shape[0] // num_to_split + split_tensors = torch.split(param, size_qkv, dim=0) + new_qkv = split_tensors[offset] + else: + assert False, f'checkpoint version: {ckpt_ver} is not supported' + + return new_qkv + + def merge_state_dict(self, + mp_world_size, + mp_rank, + quantize=False, + quantize_bits=8, + groups=64, + mlp_extra_grouping=True): + self.sanity_check(self.ckpt_list[0]) + + sd_list = self.get_merge_state_dicts(mp_world_size, mp_rank) + ds_sd = copy.deepcopy(sd_list[0]) + new_client_sd = collections.OrderedDict() + + client_sd_list = [self.get_module(sd) for sd in sd_list] + keys = client_sd_list[0].keys() + + ckpt_ver = self.get_checkpoint_version(ds_sd) + logger.info(f"checkpoint version: {ckpt_ver}") + if quantize: + quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) + + for key in keys: + value_list = [sd[key] for sd in client_sd_list] + + if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key: + if quantize: + value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key, merge_dim=1) + new_client_sd[key] = torch.cat(value_list, axis=1) + elif "attention.query_key_value" in key: + if quantize and "attention.query_key_value.weight" in key: + value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key) + new_client_sd[key] = torch.cat(value_list, axis=0) + else: + if quantize: + new_client_sd[key] = torch.cat(value_list, axis=0) + else: + new_client_sd[key] = self.merge_query_key_value(value_list, ckpt_ver) + elif "mlp.dense_h_to_4h.weight" in key or "word_embeddings.weight" in key or "mlp.dense_h_to_4h.bias" in key: + if quantize and "mlp.dense_h_to_4h.weight" in key: + value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key) + new_client_sd[key] = torch.cat(value_list, axis=0) + else: + new_client_sd[key] = value_list[0] + if quantize: + all_scales = quantizer.merge_scales() + ds_sd = self.set_module(ds_sd, new_client_sd) + + return ds_sd, (all_scales if quantize else None), len(client_sd_list) + + def split_state_dict(self, + mp_world_size, + mp_rank, + quantize=False, + quantize_bits=8, + groups=64, + mlp_extra_grouping=True): + #self.sanity_check(self.ckpt_list[0]) + + sd, num_to_split, ckpt_offset = self.get_split_state_dict(mp_world_size, mp_rank) + ds_sd = copy.deepcopy(sd) + new_client_sd = collections.OrderedDict() + + client_sd = self.get_module(sd) + + ckpt_ver = self.get_checkpoint_version(ds_sd) + logger.info(f"checkpoint version: {ckpt_ver}") + + if quantize: + quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size) + + for key in client_sd.keys(): + value = client_sd[key] + + if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key: + assert value.shape[1] % num_to_split == 0 + split_size = value.shape[1] // num_to_split + if quantize: + q_vals = quantizer.Quantize([value], quantize_bits, groups, key) + value = q_vals[0] + new_client_sd[key] = torch.split(value, split_size, dim=1)[ckpt_offset] + elif "attention.query_key_value" in key: + if quantize and "attention.query_key_value.weight" in key: + q_vals = quantizer.Quantize([value], quantize_bits, groups, key) + value = q_vals[0] + new_client_sd[key] = self.split_query_key_value(value, num_to_split, ckpt_offset, ckpt_ver) + elif "mlp.dense_h_to_4h.weight" in key or "word_embeddings.weight" in key or "mlp.dense_h_to_4h.bias" in key or "final_linear.weight" in key: + assert value.shape[0] % num_to_split == 0 + split_size = value.shape[0] // num_to_split + if quantize and "mlp.dense_h_to_4h.weight" in key: + q_vals = quantizer.Quantize([value], quantize_bits, groups, key) + value = q_vals[0] + new_client_sd[key] = torch.split(value, split_size, dim=0)[ckpt_offset] + else: + new_client_sd[key] = value + + if quantize: + all_scales = quantizer.merge_scales_split(num_to_split) + + ds_sd = self.set_module(ds_sd, new_client_sd) + + return ds_sd, (all_scales if quantize else None) + + def sanity_check(self, ckpt_file_name): + keys_to_check = [ + "attention.dense.weight", "mlp.dense_4h_to_h.weight", "attention.query_key_value", + "mlp.dense_h_to_4h.weight", "mlp.dense_h_to_4h.bias" + ] + + sd = self.checkpoint_engine.load(ckpt_file_name, map_location=lambda storage, loc: storage) + + # partial_key is a sub-string of one key in the sd + def check_key_exist(partial_key, sd): + keys = sd.keys() + found = False + for k in keys: + if partial_key in k: + found = True + break + return found + + for key in keys_to_check: + assert check_key_exist(key, + self.get_module(sd)), f'key: {key} is not found in the checkpoint {ckpt_file_name}' + + def get_checkpoint_version(self, state_dict): + # Use 0 if version info doesn't exist + return self.version if self.version is not None else state_dict.get('checkpoint_version', 0) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__init__.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..006dfd6dcbc6f7675733ba10938552cdc13f7bcb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +from .utils import MIN_SWAPPABLE_BYTES diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2642ba7ffcdf252a3be97f8d35e5a464a0b12bb6 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96ebaac9b1bfb1a0551bff69b5aad3d23259358c Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/aio_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f72d3fe57127661a643af6b7327a1a20d765781 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/async_swapper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f248c810b1b2f71b77f7357e680d8f1a22bda18 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37822f7ddd8bc738ba98233a125e42a871280760 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/optimizer_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65443388c363731ea15cc58d7e64e290e005e382 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_optimizer_swapper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1ab94d83b12bd699c870adedf0e924c0fe47222 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/partitioned_param_swapper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16159782e0198a4d8420faf26e749efb5daa147d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/pipelined_optimizer_swapper.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df7a41aae31682868a1ae5873f6825f782fe69b3 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/aio_config.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/aio_config.py new file mode 100644 index 0000000000000000000000000000000000000000..be6c7d93c86a4f5e268391f717217db1a687d0dc --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/aio_config.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import get_scalar_param +from deepspeed.runtime.swap_tensor.constants import * +from deepspeed.accelerator import get_accelerator + +AIO_DEFAULT_DICT = { + AIO_BLOCK_SIZE: AIO_BLOCK_SIZE_DEFAULT, + AIO_QUEUE_DEPTH: AIO_QUEUE_DEPTH_DEFAULT, + AIO_INTRA_OP_PARALLELISM: AIO_INTRA_OP_PARALLELISM_DEFAULT, + AIO_SINGLE_SUBMIT: AIO_SINGLE_SUBMIT_DEFAULT, + AIO_OVERLAP_EVENTS: AIO_OVERLAP_EVENTS_DEFAULT, + AIO_USE_GDS: AIO_USE_GDS_DEFAULT +} + + +def get_aio_config(param_dict): + if AIO in param_dict.keys() and param_dict[AIO] is not None: + aio_dict = param_dict[AIO] + aio_config = { + AIO_BLOCK_SIZE: + get_scalar_param(aio_dict, AIO_BLOCK_SIZE, AIO_BLOCK_SIZE_DEFAULT), + AIO_QUEUE_DEPTH: + get_scalar_param(aio_dict, AIO_QUEUE_DEPTH, AIO_QUEUE_DEPTH_DEFAULT), + AIO_INTRA_OP_PARALLELISM: + get_scalar_param(aio_dict, AIO_INTRA_OP_PARALLELISM, AIO_INTRA_OP_PARALLELISM_DEFAULT), + AIO_SINGLE_SUBMIT: + get_scalar_param(aio_dict, AIO_SINGLE_SUBMIT, AIO_SINGLE_SUBMIT_DEFAULT), + AIO_OVERLAP_EVENTS: + get_scalar_param(aio_dict, AIO_OVERLAP_EVENTS, AIO_OVERLAP_EVENTS_DEFAULT), + AIO_USE_GDS: + get_scalar_param(aio_dict, AIO_USE_GDS, AIO_USE_GDS_DEFAULT) + } + + if aio_config[AIO_USE_GDS]: + assert get_accelerator().device_name() == 'cuda', 'GDS currently only supported for CUDA accelerator' + + return aio_config + + return AIO_DEFAULT_DICT diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..b808721537fef5e3907bd3b8819b2e2d0b2f6936 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/async_swapper.py @@ -0,0 +1,175 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" +import torch + +from deepspeed import comm as dist +from deepspeed.utils.logging import logger +from deepspeed.runtime.swap_tensor.utils import swap_out_tensors, SwapBuffer +from deepspeed.accelerator import get_accelerator + +INVALID_BUFFER_INDEX = -1 +ASYNC_SWAPPER_WAIT_TIMER = 'async_swap_gradient_wait' + + +class AsyncTensorSwapper(object): + + def __init__(self, aio_handle, numel_alignment, timers): + self.free_buffer_index = [] + self.swapping_buffer_index = [] + self.ready_buffer_index = [] + self.current_buffer_index = INVALID_BUFFER_INDEX + self.all_buffers = [] + self.aio_handle = aio_handle + self.numel_alignment = numel_alignment + self.max_numel = 0 + self.num_pending_swaps = 0 + self.timers = timers + self.timer_names = set() + self.num_elements_swapped = 0 + self.dtype = None + + def has_buffers(self): + return len(self.all_buffers) > 0 + + def add_buffers(self, buffer_list): + assert len(self.all_buffers) == 0 + assert all([get_accelerator().is_pinned(buffer) for buffer in buffer_list]) + dtype = buffer_list[0].dtype + assert all([buffer.dtype == dtype for buffer in buffer_list]) + + self.dtype = dtype + self.all_buffers = [SwapBuffer(buffer) for buffer in buffer_list] + self.free_buffer_index += [i for i in range(len(self.all_buffers))] + self.max_numel = max([buffer.numel() for buffer in buffer_list]) + self.timer_names = set() + + def get_timer_names(self): + return list(self.timer_names) + + def release_buffers(self): + self._report_statistics('Swapped out[Before flush]') + self._flush_buffers_until_complete() + self._report_statistics('Swapped out[After flush]') + + pinned_buffers = [buf.buffer for buf in self.all_buffers] + self.all_buffers = [] + self.free_buffer_index = [] + self.current_buffer_index = INVALID_BUFFER_INDEX + self.num_elements_swapped = 0 + self.dtype = None + + return pinned_buffers + + def swap_out_tensors(self, tensor_list, path_list): + for tensor, swap_path in zip(tensor_list, path_list): + self._swap_out_tensor(tensor, swap_path) + + def _report_statistics(self, message): + if dist.get_rank() == 0: + element_size = torch.tensor([], dtype=self.dtype).element_size() + swapped_GB = (self.num_elements_swapped * element_size) / (1024**3) + logger.debug(f'{message} num_elems = {self.num_elements_swapped}, {swapped_GB:5.2f} GB') + + def _swap_out_tensor(self, tensor, swap_path): + assert len(self.all_buffers) > 0 + + aligned_numel = self._io_aligned_numel(tensor.numel()) + assert aligned_numel <= self.max_numel + + self._make_swap_space(aligned_numel) + assert self.current_buffer_index != INVALID_BUFFER_INDEX + + swap_buffer = self._get_current_buffer() + swap_buffer.insert_tensor(tensor, swap_path, aligned_numel) + + def _make_swap_space(self, numel): + if self.current_buffer_index == INVALID_BUFFER_INDEX: + self._allocate_buffer() + return + + if not self._get_current_buffer().has_space(numel): + if len(self.free_buffer_index) > 0: + self._flush_ready_buffers() + else: + self._flush_buffers_until_complete() + self._allocate_buffer() + + def _io_aligned_numel(self, numel): + remainder = numel % self.numel_alignment + return numel if remainder == 0 else (numel + self.numel_alignment - remainder) + + def _allocate_buffer(self): + assert self.current_buffer_index == INVALID_BUFFER_INDEX + assert len(self.all_buffers) > 0 + assert len(self.free_buffer_index) > 0 + self.current_buffer_index = self.free_buffer_index[-1] + self.free_buffer_index = self.free_buffer_index[:-1] + + def _flush_ready_buffers(self): + if self.current_buffer_index != INVALID_BUFFER_INDEX: + self.ready_buffer_index.append(self.current_buffer_index) + self.current_buffer_index = INVALID_BUFFER_INDEX + + self._swap_out_ready_buffers() + + def _flush_buffers_until_complete(self): + self._flush_ready_buffers() + assert len(self.ready_buffer_index) == 0 + + self._wait_for_swap_complete() + assert len(self.swapping_buffer_index) == 0 + assert len(self.free_buffer_index) == len(self.all_buffers) + + def _swap_out_ready_buffers(self): + for buffer_index in self.ready_buffer_index: + buffer = self._get_buffer(buffer_index) + swap_tensors = buffer.get_swap_tensors() + swap_paths = buffer.get_swap_paths() + self.num_pending_swaps += len(swap_tensors) + swap_out_tensors(self.aio_handle, swap_tensors, swap_paths) + + self.swapping_buffer_index += self.ready_buffer_index + self.ready_buffer_index = [] + + def _wait_for_swap_complete(self): + assert len(self.swapping_buffer_index) > 0 + + self._start_timer(ASYNC_SWAPPER_WAIT_TIMER) + assert self.aio_handle.wait() == self.num_pending_swaps + self._stop_timer(ASYNC_SWAPPER_WAIT_TIMER) + self.timer_names.add(ASYNC_SWAPPER_WAIT_TIMER) + + self.num_pending_swaps = 0 + + for buffer_index in self.swapping_buffer_index: + buffer = self._get_buffer(buffer_index) + self.num_elements_swapped += buffer.get_num_elem() + buffer.reset() + + self.free_buffer_index += self.swapping_buffer_index + assert len(self.free_buffer_index) <= len(self.all_buffers) + self.swapping_buffer_index = [] + + def _get_buffer(self, index): + assert index != INVALID_BUFFER_INDEX + return self.all_buffers[index] + + def _get_current_buffer(self): + return self._get_buffer(self.current_buffer_index) + + def _start_timer(self, name): + if self.timers: + self.timers(name).start() + + def _stop_timer(self, name): + if self.timers: + self.timers(name).stop() + + def _log_timers(self, name_list, force=False): + if self.timers and force: + self.timers.log(name_list) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/constants.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..c1207749eac6915b37c9d3d01f93f2e06e99250e --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/constants.py @@ -0,0 +1,30 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +AIO +""" +AIO_FORMAT = ''' +"aio": { + "block_size": 1048576, + "queue_depth": 8, + "intra_op_parallelism": 1, + "single_submit": false, + "overlap_events": true, + "use_gds": false +} +''' +AIO = "aio" +AIO_BLOCK_SIZE = "block_size" +AIO_BLOCK_SIZE_DEFAULT = 1048576 +AIO_QUEUE_DEPTH = "queue_depth" +AIO_QUEUE_DEPTH_DEFAULT = 8 +AIO_INTRA_OP_PARALLELISM = "intra_op_parallelism" +AIO_INTRA_OP_PARALLELISM_DEFAULT = 1 +AIO_SINGLE_SUBMIT = "single_submit" +AIO_SINGLE_SUBMIT_DEFAULT = False +AIO_OVERLAP_EVENTS = "overlap_events" +AIO_OVERLAP_EVENTS_DEFAULT = True +AIO_USE_GDS = "use_gds" +AIO_USE_GDS_DEFAULT = False diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6729fd28b8fc8acd33dd552d6e70d86b326d2ac2 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/optimizer_utils.py @@ -0,0 +1,527 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" + +import os +import torch + +from deepspeed import comm as dist +from deepspeed.utils.logging import logger +from deepspeed.runtime.swap_tensor.constants import * +from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, \ + MIN_AIO_BYTES, AIO_ALIGNED_BYTES, get_sized_buffers +from deepspeed.runtime.swap_tensor.utils import SwapBufferManager, SwapBufferPool +from deepspeed.accelerator import get_accelerator + + +class FlattenedTensorSwapInfo(object): + + def __init__(self, path, length, offset): + self.path = path + self.offset = offset + self.length = length + + +class SwapTensorContext(object): + + def __init__(self, tensor, swap_folder): + self.compute_tensor = tensor + self.swap_tensor = torch.Tensor() + self.swap_path = os.path.join(swap_folder, f'{OptimizerSwapper.parameter_id(tensor)}.tensor.swp') + + def release_memory(self): + self.compute_tensor.data = torch.Tensor() + self.swap_tensor.data = torch.Tensor() + + def set_buffers(self, compute_buffer, swap_buffer): + self.compute_tensor.data = compute_buffer.data + self.swap_tensor.data = swap_buffer.data + + +class OptimizerStateSwapInfo(object): + + def __init__(self, parameter, numel, base_folder): + self.tensors = [] + self.param_id = OptimizerSwapper.parameter_id(parameter) + self.swap_folder = base_folder + self.swapped_gradients = {} + self.unswapped_gradients = {} + self.tensor_numel = numel + self.tensor_dtype = parameter.dtype + self.tensor_device = parameter.device + self.has_state_tensors = False + self.swap_buffers = [] + self._add_tensors([parameter]) + + def numel(self): + return self.tensor_numel + + def has_gradients(self): + return bool(self.swapped_gradients) or bool(self.unswapped_gradients) + + def _add_tensors(self, tensor_list): + for t in tensor_list: + self.tensors.append(SwapTensorContext(t, self.swap_folder)) + + def add_state_tensors(self, tensor_list): + self.has_state_tensors = True + self._add_tensors(tensor_list) + + def num_tensors(self): + return len(self.tensors) + + def device(self): + return self.tensor_device + + def dtype(self): + return self.tensor_dtype + + def release_memory(self): + for t in self.tensors: + t.release_memory() + + def get_compute_tensors(self): + return [t.compute_tensor for t in self.tensors] + + def get_swap_paths(self): + return [t.swap_path for t in self.tensors] + + def get_swap_buffers_and_paths(self, pinned): + swap_buffers = [] + swap_paths = [] + select_tensors = [t for t in self.tensors if get_accelerator().is_pinned(t.compute_tensor) == pinned] + for t in select_tensors: + swap_buffers.append(t.swap_tensor if pinned else t.compute_tensor) + swap_paths.append(t.swap_path) + return swap_buffers, swap_paths + + def get_or_create_gradient_paths(self, offsets, lengths): + gradient_paths = [] + for offset, length in zip(offsets, lengths): + if not offset in self.swapped_gradients.keys(): + path = os.path.join(self.swap_folder, f'{self.param_id}_gradient_{offset}_{length}.tensor.swp') + self.swapped_gradients[offset] = FlattenedTensorSwapInfo(path, length, offset) + + gradient_paths.append(self.swapped_gradients[offset].path) + + return gradient_paths + + def set_swap_buffers(self, buffers, aligned_numel): + num_tensors = len(self.tensors) + compute_lengths = [self.numel()] * num_tensors + compute_buffers = get_sized_buffers(buffers, compute_lengths) + swap_lengths = [aligned_numel] * num_tensors + swap_buffers = get_sized_buffers(buffers, swap_lengths) + + for i, t in enumerate(self.tensors): + t.set_buffers(compute_buffer=compute_buffers[i], swap_buffer=swap_buffers[i]) + + def get_swap_gradient_buffers(self, swap_buffer): + assert self.numel() <= swap_buffer.numel() + return [swap_buffer.narrow(0, grad.offset, grad.length) for grad in self.swapped_gradients.values()] + + def get_swap_gradient_paths(self): + return [grad.path for grad in self.swapped_gradients.values()] + + def get_unpinned_state_tensors(self): + return [t.compute_tensor for t in self.tensors if not get_accelerator().is_pinned(t.compute_tensor)] + + def read_unswapped_gradients(self, dest_buffer): + num_elem_count = 0 + for offset, grad_partition in self.unswapped_gradients.items(): + dst_tensor = dest_buffer.narrow(0, offset, grad_partition.numel()) + dst_tensor.data.copy_(grad_partition.data) + num_elem_count += grad_partition.numel() + + return num_elem_count + + def write_unswapped_gradients(self, src_buffer): + num_elem_count = 0 + for offset, grad_partition in self.unswapped_gradients.items(): + src_tensor = src_buffer.narrow(0, offset, grad_partition.numel()) + grad_partition.data.copy_(src_tensor.data) + num_elem_count += grad_partition.numel() + + return num_elem_count + + def release_unswapped_gradients(self): + self.unswapped_gradients = {} + + +SWAPPER_DEBUG_MODE = False +SWAP_OUT_GRADIENT_TIMER = 'swap_out_gradient' + + +class OptimizerSwapper(object): + + @staticmethod + def parameter_id(param): + return param.ds_id + + def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): + self.swap_config = swap_config + self.aio_config = aio_config + + # NVMe swap management + self.swap_params_info = {} + self.swap_element_size = torch.tensor([], dtype=dtype).element_size() + self.swap_folder = os.path.join(base_folder, 'optimizer', f'rank{dist.get_rank()}') + os.makedirs(self.swap_folder, exist_ok=True) + + self.optimizer = optimizer + + # Read/Write alignment for each thread during Intra-request parallelism + self.min_aio_bytes = max(MIN_AIO_BYTES, aio_config[AIO_BLOCK_SIZE]) + self.aligned_bytes = AIO_ALIGNED_BYTES * aio_config[AIO_INTRA_OP_PARALLELISM] + self.numel_alignment = self.aligned_bytes // self.swap_element_size + + # Swap buffer management + self.largest_numel = self._io_aligned_numel(largest_numel) + self.dtype = dtype + self.swap_buffer_manager = SwapBufferManager(num_elems=self.largest_numel, + count=swap_config.buffer_count, + dtype=dtype) + + # Timers + self.timers = timers + self.timer_names = set() + + # Print exclusion list + self.print_exclude_list = [ + 'optimizer', + 'swap_buffer_manager', + 'swap_params_info', + 'timers', + 'timer_names', + ] + + def purge_state(self): + for swap_info in self.swap_params_info.values(): + swap_info.tensors = [swap_info.tensors[0]] + swap_info.has_state_tensors = False + + def is_swappable_tensor(self, tensor=None, numel=None): + assert tensor is not None or numel is not None, "Either tensor or numel must be provided" + if tensor is not None: + return self.min_aio_bytes <= (tensor.numel() * self.swap_element_size) + return self.min_aio_bytes <= (numel * self.swap_element_size) + + def init_timers(self): + self.timer_names = set() + + def log_timers(self): + if self.timer_names: + self._log_timers(list(self.timer_names), force=True) + + def pre_backward(self): + self.init_timers() + + def post_backward(self): + pass + + def _flush_gradient_swapper(self, gradient_swapper): + if gradient_swapper.has_buffers(): + self._start_timer(SWAP_OUT_GRADIENT_TIMER) + pinned_buffers = gradient_swapper.release_buffers() + self.swap_buffer_manager.free(pinned_buffers) + self._stop_timer(SWAP_OUT_GRADIENT_TIMER) + self.timer_names.add(SWAP_OUT_GRADIENT_TIMER) + self.timer_names.update(gradient_swapper.get_timer_names()) + + def _swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors, gradient_swapper): + if not OptimizerSwapper.parameter_id(parameter) in self.swap_params_info.keys(): + return + + swap_info = self.swap_params_info[OptimizerSwapper.parameter_id(parameter)] + + swappable_tensors = [] + swappable_offsets = [] + swappable_lengths = [] + + aligned_gradients, aligned_offsets = self._adjust_for_misaligned_lengths(tensors=gradient_tensors, + offsets=gradient_offsets) + + self._start_timer(SWAP_OUT_GRADIENT_TIMER) + for tensor, offset in zip(aligned_gradients, aligned_offsets): + if not self.is_swappable_tensor(tensor=tensor): + swap_info.unswapped_gradients[offset] = tensor + continue + + swappable_tensors.append(tensor) + swappable_offsets.append(offset) + swappable_lengths.append(tensor.numel()) + + if len(swappable_tensors) > 0: + if not gradient_swapper.has_buffers(): + pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) + + gradient_swapper.add_buffers(pinned_buffers) + + swappable_paths = swap_info.get_or_create_gradient_paths(swappable_offsets, swappable_lengths) + + gradient_swapper.swap_out_tensors(tensor_list=swappable_tensors, path_list=swappable_paths) + + self._stop_timer(SWAP_OUT_GRADIENT_TIMER) + self.timer_names.add(SWAP_OUT_GRADIENT_TIMER) + + def _initialize_from_swapped_fp16_params(self, aio_handle, fp16_partitions_info, fp16_num_elems, + fp16_pinned_buffers, fp32_parameters): + assert len(fp32_parameters) == len(fp16_partitions_info) + assert len(fp32_parameters) == len(fp16_num_elems) + assert all([get_accelerator().is_pinned(buffer) for buffer in fp16_pinned_buffers]) + + fp32_swap_paths = self._get_swap_paths(parameters=fp32_parameters, num_elems=fp16_num_elems) + + fp32_pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) + + fp16_buffer_numel = [buf.numel() for buf in fp16_pinned_buffers] + assert all([numel >= self.largest_numel for numel in fp16_buffer_numel]), \ + f"numel of fp16 buffers {fp16_buffer_numel} is too small for initializing fp32 params {self.largest_numel}" + + fp32_swap_buffers = SwapBufferPool(fp32_pinned_buffers) + fp16_swap_buffers = SwapBufferPool(fp16_pinned_buffers) + + curr_index = 0 + while curr_index < len(fp32_parameters): + fp16_pinned_tensors = self._swap_in_fp16_params(aio_handle=aio_handle, + fp16_num_elems=fp16_num_elems[curr_index:], + fp16_partitions_info=fp16_partitions_info[curr_index:], + fp16_swap_buffers=fp16_swap_buffers) + + if dist.get_rank() == 0 and SWAPPER_DEBUG_MODE: + for i, tensor in enumerate(fp16_pinned_tensors): + true_index = curr_index + i + logger.info( + f'swap_in_fp16_param: fp32_id = {OptimizerSwapper.parameter_id(fp32_parameters[true_index])} index = {true_index} orig_num_elem = {fp16_num_elems[true_index]}, swap_num_elem = {fp16_pinned_tensors[i].numel()}' + ) + + swap_out_count = self._swap_out_fp16_params(aio_handle=aio_handle, + fp32_swap_paths=fp32_swap_paths[curr_index:], + fp32_swap_buffers=fp32_swap_buffers, + fp16_pinned_tensors=fp16_pinned_tensors) + assert swap_out_count == len(fp16_pinned_tensors), \ + f"{swap_out_count} does not match {len(fp16_pinned_tensors)}" + + fp16_swap_buffers.reset() + fp32_swap_buffers.reset() + curr_index += swap_out_count + + self.swap_buffer_manager.free(fp32_pinned_buffers) + + def _swap_in_fp16_params(self, aio_handle, fp16_num_elems, fp16_partitions_info, fp16_swap_buffers): + assert len(fp16_num_elems) > 0 + + swapped_fp16_tensors = [] + swap_tensors = [] + swap_paths = [] + unswapped_srcs = [] + unswapped_dsts = [] + + for i, numel in enumerate(fp16_num_elems): + pinned_tensor, _ = fp16_swap_buffers.allocate_tensor(numel, None, numel) + if pinned_tensor is None: + break + + swapped_fp16_tensors.append(pinned_tensor) + offset = 0 + for tensor, partition_numel, partition_path in fp16_partitions_info[i]: + dst_tensor = pinned_tensor.narrow(0, offset, partition_numel) + if partition_path is None: + unswapped_srcs.append(tensor) + unswapped_dsts.append(dst_tensor) + else: + swap_paths.append(partition_path) + swap_tensors.append(dst_tensor) + offset += partition_numel + + assert len(swapped_fp16_tensors) + len(unswapped_srcs) > 0 + ret = swap_in_tensors(aio_handle, swap_tensors, swap_paths) + for src, dst in zip(unswapped_srcs, unswapped_dsts): + dst.data.copy_(src.data) + + assert len(swap_tensors) == aio_handle.wait() + + return swapped_fp16_tensors + + def _swap_out_fp16_params(self, aio_handle, fp32_swap_paths, fp32_swap_buffers, fp16_pinned_tensors): + + assert len(fp16_pinned_tensors) <= len(fp32_swap_paths) + swap_out_count = 0 + for i, fp16_tensor in enumerate(fp16_pinned_tensors): + if not fp32_swap_buffers.has_space(fp16_tensor.numel()): + fp32_swap_buffers.swap_out(aio_handle) + fp32_swap_buffers.reset() + + pinned_tensor, _ = fp32_swap_buffers.insert_tensor(fp16_tensor, fp32_swap_paths[i], + self._io_aligned_numel(fp16_tensor.numel())) + assert pinned_tensor is not None + swap_out_count += 1 + + if len(fp32_swap_buffers.get_swap_tensors()) > 0: + fp32_swap_buffers.swap_out(aio_handle) + + return swap_out_count + + def _initialize_parameters(self, parameters, src_tensors, aio_handle): + assert len(parameters) == len(src_tensors) + + swap_paths = self._get_swap_paths(parameters=parameters, num_elems=[src.numel() for src in src_tensors]) + + SWAP_INIT_TIMER = "swap_init_write" + self._start_timer(SWAP_INIT_TIMER) + + pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) + assert pinned_buffers is not None + + self._swap_out_unpinned_tensors(aio_handle=aio_handle, + unpinned_tensors=src_tensors, + dest_paths=swap_paths, + pinned_buffers=pinned_buffers) + + if dist.get_rank() == 0 and SWAPPER_DEBUG_MODE: + for i, tensor in enumerate(src_tensors): + logger.info( + f'copy_in_fp16_param: fp32_id = {OptimizerSwapper.parameter_id(parameters[i])} index = {i}, swap_num_elem = {src_tensors[i].numel()}' + ) + + self.swap_buffer_manager.free(pinned_buffers) + + self._stop_timer(SWAP_INIT_TIMER) + self._log_timers([SWAP_INIT_TIMER]) + + def _get_swap_paths(self, parameters, num_elems): + swap_info_list = [ + self._create_param_swap_info(parameter=p, + numel=numel) \ + for p, numel in zip(parameters, num_elems) + ] + assert len(swap_info_list) == len(num_elems) + + swap_paths = [info.tensors[0].swap_path for info in swap_info_list] + return swap_paths + + def _swap_out_unpinned_tensors(self, aio_handle, unpinned_tensors, dest_paths, pinned_buffers): + + swap_buffer_count = len(pinned_buffers) + unpinned_tensor_count = len(unpinned_tensors) + + for i in range(0, unpinned_tensor_count, swap_buffer_count): + swap_tensor_count = min((unpinned_tensor_count - i), swap_buffer_count) + + src_tensors = unpinned_tensors[i:(i + swap_tensor_count)] + compute_lengths = [t.numel() for t in src_tensors] + compute_buffers = get_sized_buffers(pinned_buffers, compute_lengths) + + for dst, src in zip(compute_buffers, src_tensors): + dst.data.copy_(src.data) + + swap_lengths = [self._io_aligned_numel(t.numel()) for t in src_tensors] + swap_buffers = get_sized_buffers(pinned_buffers, swap_lengths) + + swap_paths = dest_paths[i:(i + swap_tensor_count)] + swap_out_tensors(aio_handle, swap_buffers, swap_paths) + + assert aio_handle.wait() == swap_tensor_count + + def _adjust_for_misaligned_lengths(self, tensors, offsets): + new_tensors = [] + new_offsets = [] + + for orig_tensor, orig_offset in zip(tensors, offsets): + if not self.is_swappable_tensor(tensor=orig_tensor): + new_tensors.append(orig_tensor) + new_offsets.append(orig_offset) + continue + + remainder = orig_tensor.numel() % self.numel_alignment + if remainder == 0: + new_tensors.append(orig_tensor) + new_offsets.append(orig_offset) + continue + + # Split into two by making remainder a tensor + aligned_length = (orig_tensor.numel() // self.numel_alignment) * self.numel_alignment + new_tensors.append(orig_tensor.narrow(0, 0, aligned_length)) + new_offsets.append(orig_offset) + + # remainder tensor + new_tensors.append(orig_tensor.narrow(0, aligned_length, remainder)) + new_offsets.append(orig_offset + aligned_length) + + return new_tensors, new_offsets + + def _retrieve_unswapped_grad_partitions(self, swap_info, dest_buffer): + UNSWAPPED_READ_GRADIENTS = 'unswapped_read_gradients' + self._start_timer(UNSWAPPED_READ_GRADIENTS) + tensor_count = len(swap_info.unswapped_gradients) + num_elem_count = swap_info.read_unswapped_gradients(dest_buffer) + self._stop_timer(UNSWAPPED_READ_GRADIENTS) + self._log_timers([UNSWAPPED_READ_GRADIENTS]) + + # It should be safe to discard unswapped gradient partitions + swap_info.release_unswapped_gradients() + + if SWAPPER_DEBUG_MODE: + logger.info( + f'optimizer_retrieve_unswapped_gradients: param={swap_info.param_id} tensor_count={tensor_count} elem_count={num_elem_count}' + ) + + def _get_state_tensors(self, parameter): + if not parameter in self.optimizer.state: + return [] + + tensor_list = [] + for state_name, value in self.optimizer.state[parameter].items(): + if torch.is_tensor(value) and self.is_swappable_tensor(tensor=value): + value.ds_id = state_name + '-' + parameter.ds_id + tensor_list.append(value) + + return tensor_list + + def _update_param_state_info(self, swap_info, parameter): + if not swap_info.has_state_tensors: + state_tensors = self._get_state_tensors(parameter) + if state_tensors: + swap_info.add_state_tensors(state_tensors) + + def _create_param_swap_info(self, parameter, numel): + param_id = OptimizerSwapper.parameter_id(parameter) + assert not param_id in self.swap_params_info + + self.swap_params_info[param_id] = OptimizerStateSwapInfo(parameter=parameter, + numel=numel, + base_folder=self.swap_folder) + swap_info = self.swap_params_info[param_id] + + self._update_param_state_info(swap_info, parameter) + + return swap_info + + def _get_param_swap_info(self, parameter): + param_id = OptimizerSwapper.parameter_id(parameter) + swap_info = self.swap_params_info.get(param_id, None) + + if swap_info is not None: + self._update_param_state_info(swap_info, parameter) + + return swap_info + + def _start_timer(self, name): + if self.timers: + self.timers(name).start() + + def _stop_timer(self, name): + if self.timers: + self.timers(name).stop() + + def _log_timers(self, name_list, force=False): + if self.timers and (SWAPPER_DEBUG_MODE or force): + self.timers.log(name_list) + + def _io_aligned_numel(self, numel): + remainder = numel % self.numel_alignment + return numel if remainder == 0 else (numel + self.numel_alignment - remainder) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..52b873ba58a1f7816aae4f79527773eb26870a34 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/partitioned_optimizer_swapper.py @@ -0,0 +1,229 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +from deepspeed.utils.logging import logger +from deepspeed.ops.op_builder import AsyncIOBuilder +from deepspeed import comm as dist + +from deepspeed.runtime.swap_tensor.constants import * +from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object, \ + get_sized_buffers +from deepspeed.runtime.swap_tensor.async_swapper import AsyncTensorSwapper +from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper +from deepspeed.accelerator import get_accelerator + +DEBUG_MODE = False + +SWAP_IN_PARAM_TIMER = 'swap_in_param' +SWAP_OUT_PARAM_TIMER = 'swap_out_param' +SWAP_IN_GRADIENT_TIMER = 'swap_in_gradient' + + +class PartitionedOptimizerSwapper(OptimizerSwapper): + + def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): + super(PartitionedOptimizerSwapper, self).__init__(swap_config, aio_config, base_folder, optimizer, + largest_numel, device, dtype, timers) + + aio_op = AsyncIOBuilder().load() + self.aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], + queue_depth=aio_config[AIO_QUEUE_DEPTH], + single_submit=aio_config[AIO_SINGLE_SUBMIT], + overlap_events=aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) + + # Overlap swapping out + self.gradient_swapper = AsyncTensorSwapper(aio_handle=self.aio_handle, + numel_alignment=self.numel_alignment, + timers=self.timers) + + self.print_exclude_list += ['aio_handle', 'gradient_swapper', 'print_exclude_list'] + + if dist.get_rank() == 0: + print_object(obj=self, name='PartitionedOptimizerSwapper', exclude_list=self.print_exclude_list) + + def initialize_parameters(self, parameters, src_tensors): + self._initialize_parameters(parameters=parameters, src_tensors=src_tensors, aio_handle=self.aio_handle) + + def initialize_from_swapped_fp16_params(self, fp16_partitions_info, fp16_num_elems, fp16_pinned_buffers, + fp32_parameters): + self._initialize_from_swapped_fp16_params(aio_handle=self.aio_handle, + fp16_partitions_info=fp16_partitions_info, + fp16_num_elems=fp16_num_elems, + fp16_pinned_buffers=fp16_pinned_buffers, + fp32_parameters=fp32_parameters) + + def flush_gradients(self): + self._flush_gradient_swapper(self.gradient_swapper) + + def release_swap_buffers(self, parameter): + swap_info = self._get_param_swap_info(parameter) + if swap_info is None: + return + swap_info.release_memory() + + self.swap_buffer_manager.free(swap_info.swap_buffers) + swap_info.swap_buffers = [] + + def swap_in_optimizer_state(self, parameter, async_parameter=None): + swap_info = self._get_param_swap_info(parameter) + if swap_info is None: + return + + self._flush_gradient_swapper(self.gradient_swapper) + + required_buffer_count = swap_info.num_tensors() + (1 if swap_info.has_gradients() else 0) + aligned_numel = self._io_aligned_numel(swap_info.numel()) + pinned_buffers = self.swap_buffer_manager.allocate(num_elems=aligned_numel, + count=required_buffer_count, + dtype=parameter.dtype) + assert pinned_buffers is not None + swap_info.swap_buffers = pinned_buffers.copy() + + self._start_timer(SWAP_IN_PARAM_TIMER) + self._swap_in_parameter(aio_handle=self.aio_handle, + parameter=parameter, + dest_buffers=pinned_buffers[:swap_info.num_tensors()]) + self._stop_timer(SWAP_IN_PARAM_TIMER) + self.timer_names.add(SWAP_IN_PARAM_TIMER) + + if swap_info.has_gradients(): + self._start_timer(SWAP_IN_GRADIENT_TIMER) + self._swap_in_gradients(aio_handle=self.aio_handle, parameter=parameter, dest_buffer=pinned_buffers[-1]) + self._stop_timer(SWAP_IN_GRADIENT_TIMER) + self.timer_names.add(SWAP_IN_GRADIENT_TIMER) + + def _swap_out_optimizer_state(self, swap_info): + pinned_tensors, pinned_paths = swap_info.get_swap_buffers_and_paths(True) + WRITE_TIMER = 'swap_submit_write' + self._start_timer(WRITE_TIMER) + + swap_out_tensors(self.aio_handle, pinned_tensors, pinned_paths) + assert self.aio_handle.wait() == len(pinned_tensors) + + unpinned_tensors, unpinned_paths = swap_info.get_swap_buffers_and_paths(False) + if len(unpinned_tensors) > 0: + pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype) + self._swap_out_unpinned_tensors(aio_handle=self.aio_handle, + unpinned_tensors=unpinned_tensors, + dest_paths=unpinned_paths, + pinned_buffers=pinned_buffers) + swap_info.swap_buffers += pinned_buffers.copy() + + self._stop_timer(WRITE_TIMER) + self._log_timers([WRITE_TIMER]) + + def writeback_optimizer_state_and_gradients(self, parameter, write_opt_state, write_gradients): + swap_info = self._get_param_swap_info(parameter=parameter) + + if swap_info is None: + return + + if write_opt_state: + self._swap_out_optimizer_state(swap_info) + + if write_gradients and swap_info.has_gradients(): + param_gradients = swap_info.swapped_gradients.values() + swap_buffers = [parameter.grad.narrow(0, grad.offset, grad.length) for grad in param_gradients] + swap_paths = [grad.path for grad in param_gradients] + swap_out_tensors(self.aio_handle, swap_buffers, swap_paths) + assert len(swap_buffers) == self.aio_handle.wait() + if swap_info.unswapped_gradients: + swap_info.write_unswapped_gradients(src_buffer=parameter.grad) + + self.release_swap_buffers(parameter) + + def swap_out_optimizer_state(self, parameter, async_swap=False): + swap_info = self._get_param_swap_info(parameter=parameter) + + if swap_info is None: + return + + swap_bytes = sum( + [self._io_aligned_numel(t.numel()) * t.element_size() for t in swap_info.get_compute_tensors()]) + + self._start_timer(SWAP_OUT_PARAM_TIMER) + self._swap_out_optimizer_state(swap_info) + self.release_swap_buffers(parameter) + self._stop_timer(SWAP_OUT_PARAM_TIMER) + self.timer_names.add(SWAP_OUT_PARAM_TIMER) + + if DEBUG_MODE and dist.get_rank() == 0: + logger.info(f'optimizer_param_swap_out: {(swap_bytes/(1024**3)):5.2f} GB') + + def swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors): + self._swap_out_gradients(parameter=parameter, + gradient_offsets=gradient_offsets, + gradient_tensors=gradient_tensors, + gradient_swapper=self.gradient_swapper) + + def _swap_in_parameter(self, aio_handle, parameter, dest_buffers): + swap_info = self._get_param_swap_info(parameter) + if swap_info is None: + return + + num_swap_tensors = swap_info.num_tensors() + assert num_swap_tensors <= len(dest_buffers) + + swap_lengths = [self._io_aligned_numel(swap_info.numel())] * num_swap_tensors + swap_buffers = get_sized_buffers(dest_buffers, swap_lengths) + + compute_lengths = [swap_info.numel()] * num_swap_tensors + compute_buffers = get_sized_buffers(dest_buffers, compute_lengths) + + READ_TIMER = 'swap_submit_read_param' + WAIT_TIMER = 'swap_wait_read_param' + + self._start_timer(READ_TIMER) + swap_in_tensors(aio_handle, swap_buffers, swap_info.get_swap_paths()) + self._stop_timer(READ_TIMER) + + swap_bytes = sum([buffer.numel() * buffer.element_size() for buffer in swap_buffers]) + + self._start_timer(WAIT_TIMER) + aio_handle.wait() + self._stop_timer(WAIT_TIMER) + + swap_info.set_swap_buffers(dest_buffers, self._io_aligned_numel(swap_info.numel())) + + self._log_timers([READ_TIMER, WAIT_TIMER]) + if DEBUG_MODE and dist.get_rank() == 0: + logger.info(f'optimizer_param_swap_in: {(swap_bytes/(1024**3)):5.2f} GB') + + def _swap_in_pinned_gradients(self, aio_handle, parameter, gradient_tensor): + swap_info = self.swap_params_info[OptimizerSwapper.parameter_id(parameter)] + param_gradients = swap_info.swapped_gradients.values() + swap_buffers = [gradient_tensor.narrow(0, grad.offset, grad.length) for grad in param_gradients] + swap_paths = [grad.path for grad in param_gradients] + SWAP_READ_GRADIENTS = 'swap_submit_read_gradient' + SWAP_WAIT_GRADIENTS = 'swap_submit_wait_gradient' + self._start_timer(SWAP_READ_GRADIENTS) + swap_in_tensors(aio_handle, swap_buffers, swap_paths) + self._stop_timer(SWAP_READ_GRADIENTS) + + self._start_timer(SWAP_WAIT_GRADIENTS) + assert len(swap_buffers) == aio_handle.wait() + self._stop_timer(SWAP_WAIT_GRADIENTS) + + self._log_timers([SWAP_READ_GRADIENTS, SWAP_WAIT_GRADIENTS]) + + def _swap_in_gradients(self, aio_handle, parameter, dest_buffer): + swap_info = self.swap_params_info.get(OptimizerSwapper.parameter_id(parameter), None) + if not (swap_info and swap_info.has_gradients()): + return + + assert get_accelerator().is_pinned(dest_buffer) + assert parameter.numel() <= dest_buffer.numel() + + parameter.grad = dest_buffer.narrow(0, 0, parameter.numel()) + + if swap_info.swapped_gradients: + self._swap_in_pinned_gradients(aio_handle, parameter, parameter.grad) + + if swap_info.unswapped_gradients: + self._retrieve_unswapped_grad_partitions(swap_info=swap_info, dest_buffer=parameter.grad) diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..f80fe1501c006249e12d20f7a3f5ce75dec673cb --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py @@ -0,0 +1,422 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" + +import os +import shutil +from enum import Enum +import torch +from deepspeed import comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import AsyncIOBuilder +from deepspeed.ops.op_builder import GDSBuilder +from .constants import * +from .utils import swap_in_tensors, swap_out_tensors, MIN_AIO_BYTES, AIO_ALIGNED_BYTES, print_object, SwapBufferPool + + +def print_rank_0(message, debug=False, force=False): + if dist.get_rank() == 0 and (debug or force): + print(message) + + +class PartitionedParamStatus(Enum): + # Partitioned parameters are present and ready for use + AVAILABLE = 1 + + # partitioned params are in some non-memory device + NOT_AVAILABLE = 2 + + # partitioned params are being read from some non-memory device. + INFLIGHT = 3 + + +class AsyncPartitionedParameterSwapper(object): + + def __init__(self, ds_config, model_dtype): + + self.dtype = model_dtype + + #set swap buffers, create aio handles + self._configure_aio(ds_config) + + #mapping from param id to path + self.id_to_path = {} + + #mapping from pram_id to buffer id + self.param_id_to_buffer_id = {} + + # mapping from param_id to swap buffer + self.param_id_to_swap_buffer = {} + + #number of elements in the param + self.param_id_to_numel = {} + + self.pending_writes = 0 + self.pending_reads = 0 + + #keep track of async swap in params and buffers + self.inflight_params = [] + self.inflight_swap_in_buffers = [] + self.inflight_numel = 0 + + #keep track of available params + self.available_params = set() + self.available_numel = 0 + + # for swapping out from partitioned fp32 params + self.partitioned_swap_buffer = None + self.partitioned_swap_pool = None + + self.invalid_buffer = torch.tensor(1).half() + + if dist.get_rank() == 0: + exclude_list = ['aio_read_handle', 'aio_write_handle', 'buffers'] + print_object(obj=self, name='AsyncPartitionedParameterSwapper', exclude_list=exclude_list) + + def available_swap_in_buffers(self): + return len(self.available_buffer_ids) + + def _configure_aio(self, ds_config): + self.swap_config = ds_config.zero_config.offload_param + torch_dtype_string = str(self.dtype).split(".")[1] + self.swap_folder = os.path.join(self.swap_config.nvme_path, 'zero_stage_3', f'{torch_dtype_string}params', + f'rank{dist.get_rank()}') + shutil.rmtree(self.swap_folder, ignore_errors=True) + os.makedirs(self.swap_folder, exist_ok=True) + + self.swap_element_size = torch.tensor([], dtype=self.dtype).element_size() + + self.aio_config = ds_config.aio_config + + self.use_gds = self.aio_config[AIO_USE_GDS] + self.aio_handle = GDSBuilder().load(verbose=False).gds_handle if self.use_gds else AsyncIOBuilder().load( + verbose=False).aio_handle + + # Read/Write alignment for each thread during Intra-request parallelism + self.min_aio_bytes = max(MIN_AIO_BYTES, self.aio_config[AIO_BLOCK_SIZE]) + self.aligned_bytes = AIO_ALIGNED_BYTES * self.aio_config[AIO_INTRA_OP_PARALLELISM] + self.numel_alignment = self.aligned_bytes // self.swap_element_size + + self.elements_per_buffer = self.swap_config.buffer_size + self.aligned_elements_per_buffer = self._io_aligned_numel(self.elements_per_buffer) + self.param_buffer_count = self.swap_config.buffer_count + + self.available_buffer_ids = [i for i in range(self.param_buffer_count)] + self.reserved_buffer_ids = [] + + self.aio_read_handle = self.aio_handle(block_size=self.aio_config[AIO_BLOCK_SIZE], + queue_depth=self.aio_config[AIO_QUEUE_DEPTH], + single_submit=self.aio_config[AIO_SINGLE_SUBMIT], + overlap_events=self.aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=self.aio_config[AIO_INTRA_OP_PARALLELISM]) + + self.aio_write_handle = self.aio_handle(block_size=self.aio_config[AIO_BLOCK_SIZE], + queue_depth=self.aio_config[AIO_QUEUE_DEPTH], + single_submit=self.aio_config[AIO_SINGLE_SUBMIT], + overlap_events=self.aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=self.aio_config[AIO_INTRA_OP_PARALLELISM]) + + if self.use_gds: + self.buffers = torch.empty(int(self.aligned_elements_per_buffer * self.param_buffer_count), + dtype=self.dtype, + device=get_accelerator().device_name(), + requires_grad=False) + self.aio_read_handle.pin_device_tensor(self.buffers) + else: + self.buffers = get_accelerator().pin_memory(torch.empty(int(self.aligned_elements_per_buffer * + self.param_buffer_count), + dtype=self.dtype, + requires_grad=False), + align_bytes=0) + + self.swap_out_params = [] + + #Check if partitioned param or numel in a tensor is swappable or not + def swappable_tensor(self, param=None, numel=None): + if param is not None: + assert numel is None, "Both parma and numel cannot be provided" + numel = param.ds_tensor.ds_numel + if numel is not None: + return self.min_aio_bytes <= numel * self.swap_element_size + assert False, "Either param or numel must be provided" + + def get_path(self, param, must_exist=False): + paths = self._get_swap_paths([param], must_exist=must_exist) + return paths[0] + + def _get_swap_paths(self, params, must_exist=False): + paths = [] + for param in params: + param_id = param.ds_id + if param_id in self.id_to_path.keys(): + param_path = self.id_to_path[param_id] + else: + assert not must_exist, f"Path for param id {param_id} does not exist" + param_path = os.path.join(self.swap_folder, f'{param_id}_param.tensor.swp') + + self.id_to_path[param_id] = param_path + paths.append(param_path) + + return paths + + def _get_swap_buffers(self, params): + buffers = [] + for param in params: + param_id = param.ds_id + assert param_id in self.param_id_to_swap_buffer.keys(), \ + f'param {param_id} has not been assigned a swap buffer' + buffers.append(self.param_id_to_swap_buffer[param_id]) + + return buffers + + def _track_numel(self, params): + for param in params: + assert param.ds_tensor is not None, "Partitioned tensor is None" + self.param_id_to_numel[param.ds_id] = param.ds_tensor.ds_numel + + def _allocate_and_return_buffers_for_swap_in(self, params): + compute_buffers = [] + swap_buffers = [] + + for param in params: + param_id = param.ds_id + assert param_id in self.param_id_to_numel.keys(), f" Number of elements in param {param_id} is unknown" + assert param_id not in self.param_id_to_buffer_id.keys( + ), f"param {param_id} already assigned swap buffer id {self.param_id_to_buffer_id[param_id]}" + assert param_id not in self.param_id_to_swap_buffer.keys( + ), f"param {param_id} has already been assigned a swap buffer" + + buffer_id = self.available_buffer_ids.pop() + print_rank_0(f"param {param.ds_id} is assigned swap in buffer id {buffer_id} ") + self.param_id_to_buffer_id[param_id] = buffer_id + aligned_swap_numel = self._io_aligned_numel(self.param_id_to_numel[param_id]) + swap_buffer = self.buffers.narrow(0, int(buffer_id * self.aligned_elements_per_buffer), aligned_swap_numel) + + self.param_id_to_swap_buffer[param_id] = swap_buffer + compute_buffer = swap_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) + compute_buffers.append(compute_buffer) + swap_buffers.append(swap_buffer) + + return compute_buffers, swap_buffers + + #waits for inflight nvme write to complete + def synchronize_writes(self): + if self.pending_writes == 0: + return + assert self.pending_writes == self.aio_write_handle.wait() + self.pending_writes = 0 + self.remove_partition_and_release_buffers(self.swap_out_params) + self.swap_out_params = [] + + #waits for inflight nvme reads to complete + def synchronize_reads(self): + if self.pending_reads == 0: + return + + assert self.pending_reads == self.aio_read_handle.wait() + + self.pending_reads = 0 + + for param, swap_in_buffer in zip(self.inflight_params, self.inflight_swap_in_buffers): + param_id = param.ds_id + compute_buffer = swap_in_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) + param.ds_tensor.data = compute_buffer.data + param.ds_tensor.status = PartitionedParamStatus.AVAILABLE + + self.available_params.update([param.ds_id for param in self.inflight_params]) + self.available_numel += self.inflight_numel + + self.inflight_params = [] + self.inflight_swap_in_buffers = [] + self.inflight_numel = 0 + + #Removes the memory assignment and releases the buffers + #Should only be executed after swapping out the tensors + def remove_partition_and_release_buffers(self, params): + for param in params: + param_id = param.ds_id + + if param_id in self.param_id_to_buffer_id.keys(): + + buffer_id = self.param_id_to_buffer_id[param_id] + + assert buffer_id is not None, "Missing buffer id for releasing" + + self.available_buffer_ids.append(buffer_id) + del self.param_id_to_buffer_id[param_id] + del self.param_id_to_swap_buffer[param_id] + print_rank_0(f"param {param.ds_id} releases buffer id {buffer_id} ") + + if param_id in self.available_params: + self.available_params.remove(param_id) + self.available_numel -= self.param_id_to_numel[param_id] + + param.ds_tensor.data = self.invalid_buffer.data + param.ds_tensor.status = PartitionedParamStatus.NOT_AVAILABLE + + #writes from in memory to nvme. Does not release the buffers + def _swap_out(self, params, async_op=True): + + swap_out_paths = self._get_swap_paths(params) + swap_out_params = self._get_swap_buffers(params) + self._track_numel(params) + + swap_out_tensors(self.aio_write_handle, swap_out_params, swap_out_paths) + + self.pending_writes += len(swap_out_params) + self.swap_out_params += params + + if not async_op: + self.synchronize_writes() + + #blocking swap out followed by releasing the memory buffers + def swap_out_and_release(self, params, async_op=False, force_buffer_release=False): + if async_op: + assert force_buffer_release, "Should not release preallocated buffers without completing the swap out. Set force_buffer_release to True to do it anyways" + self._swap_out(params, async_op=async_op) + + # book keeping function for inflight swap in + def _update_inflight_swap_in(self, params, swap_in_buffers, inflight_numel): + self.inflight_params.extend(params) + self.inflight_swap_in_buffers.extend(swap_in_buffers) + self.inflight_numel += inflight_numel + + for param in params: + param.ds_tensor.status = PartitionedParamStatus.INFLIGHT + + self.pending_reads += len(params) + + #assigns an in memory buffer and swaps in from nvme + def swap_in(self, params, async_op=True, swap_in_buffers=None): + + assert all([param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE + for param in params]), "Some params are already available or in flight" + swap_in_paths = self._get_swap_paths(params) + + if swap_in_buffers is None: + if len(self.available_buffer_ids) < len(swap_in_paths): + ids = [p.ds_id for p in params] + print_rank_0( + f'Not enough swap in buffers {len(self.available_buffer_ids)} for {len(swap_in_paths)} params, ids = {ids}', + force=True) + print_rank_0( + f'Num inflight: params {len(self.inflight_params)}, buffers {len(self.inflight_swap_in_buffers)}, numel = {self.inflight_numel}', + force=True) + print_rank_0( + f'Num available params: count = {len(self.available_params)}, ids = {self.available_params}, numel = {self.available_numel}', + force=True) + + assert len(swap_in_paths) <= len( + self.available_buffer_ids + ), f"Not enough buffers {len(self.available_buffer_ids)} for swapping {len(swap_in_paths)}" + compute_buffers, swap_in_buffers = self._allocate_and_return_buffers_for_swap_in(params) + inflight_numel = sum([t.numel() for t in compute_buffers]) + else: + inflight_numel = sum([t.numel() for t in swap_in_buffers]) + + swap_in_tensors(self.aio_read_handle, swap_in_buffers, swap_in_paths) + + self._update_inflight_swap_in(params, swap_in_buffers, inflight_numel) + + if not async_op: + self.synchronize_reads() + + # Enables swapping into buffer that is out the control of swapper. This is always synchronous + def swap_into_buffer(self, param, dest_buffer): + assert param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE, f"param {param.ds_id} is already available or inflight" + + require_swap_buffer = not (get_accelerator().is_pinned(dest_buffer) + and self._is_io_aligned(dest_buffer.numel())) + + if require_swap_buffer: + assert len(self.available_buffer_ids) > 0, f"No buffer available to swap param {param.ds_id}." + compute_buffers, swap_in_buffers = self._allocate_and_return_buffers_for_swap_in([param]) + inflight_numel = compute_buffers[0].numel() + else: + swap_in_buffers = [dest_buffer] + inflight_numel = dest_buffer.numel() + + swap_in_paths = self._get_swap_paths([param]) + + swap_in_tensors(self.aio_read_handle, swap_in_buffers, swap_in_paths) + self._update_inflight_swap_in([param], swap_in_buffers, inflight_numel) + self.synchronize_reads() + + if require_swap_buffer: + dest_buffer.data.copy_(param.ds_tensor.data) + # Release swap buffer memory assignment. Note, this will mark the parameter not available. + self.remove_partition_and_release_buffers([param]) + + #assign a buffer to a param and return the buffer + def get_buffer(self, param, numel): + param_id = param.ds_id + + assert self.available_swap_in_buffers( + ) > 0, f"No swap buffers to allocate for fp16 param {param_id} of numel = {numel}" + assert numel < self.elements_per_buffer, f"More elements {numel} than buffer size {self.elements_per_buffer}" + + self.param_id_to_numel[param_id] = numel + buffer_id = self.available_buffer_ids.pop() + self.param_id_to_buffer_id[param_id] = buffer_id + aligned_swap_numel = self._io_aligned_numel(self.param_id_to_numel[param_id]) + swap_buffer = self.buffers.narrow(0, int(buffer_id * self.aligned_elements_per_buffer), aligned_swap_numel) + + self.param_id_to_swap_buffer[param_id] = swap_buffer + compute_buffer = swap_buffer.narrow(0, 0, self.param_id_to_numel[param_id]) + print_rank_0(f"param {param.ds_id} is assigned swap in buffer id {buffer_id}") + return compute_buffer + + def reserve_available_buffers(self): + buffers = [] + for id in self.available_buffer_ids: + buffers.append( + self.buffers.narrow(0, int(id * self.aligned_elements_per_buffer), + int(self.aligned_elements_per_buffer))) + self.reserved_buffer_ids.append(id) + + self.available_buffer_ids = [] + return buffers + + def release_reserved_buffers(self): + for id in self.reserved_buffer_ids: + self.available_buffer_ids.append(id) + self.reserved_buffer_ids = [] + + def _io_aligned_numel(self, numel): + remainder = numel % self.numel_alignment + return numel if remainder == 0 else (numel + self.numel_alignment - remainder) + + def _is_io_aligned(self, numel): + return (numel % self.numel_alignment) == 0 + + def reserve_partitioned_swap_space(self, partition_num_elems): + aligned_numel = sum([self._io_aligned_numel(numel) for numel in partition_num_elems]) + self.partitioned_swap_buffer = get_accelerator().pin_memory(torch.zeros(aligned_numel, + device='cpu', + dtype=self.dtype), + align_bytes=0) + self.partitioned_swap_pool = SwapBufferPool([self.partitioned_swap_buffer]) + + def swap_out_partitioned_params(self, dst_fp16_params, src_fp32_params): + assert self.partitioned_swap_buffer is not None, f'partitioned swap buffers for fp16 params not initialized' + assert self.partitioned_swap_pool is not None, f'partitioned swap pool for fp16 params not initialized' + assert len(dst_fp16_params) == len(src_fp32_params), \ + f'mismatch in number of fp16 params {len(dst_fp16_params)} and fp32 params {len(src_fp32_params)}' + + fp16_swap_paths = self._get_swap_paths(dst_fp16_params, must_exist=True) + self.synchronize_writes() + self.partitioned_swap_pool.reset() + for i, fp32_tensor in enumerate(src_fp32_params): + swap_tensor, _ = self.partitioned_swap_pool.insert_tensor(fp32_tensor, fp16_swap_paths[i], + self._io_aligned_numel(fp32_tensor.numel())) + assert swap_tensor is not None + dst_fp16_params[i].ds_tensor.status = PartitionedParamStatus.AVAILABLE + + self.partitioned_swap_pool.swap_out(self.aio_write_handle) + + for param in dst_fp16_params: + param.ds_tensor.status = PartitionedParamStatus.NOT_AVAILABLE diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py new file mode 100644 index 0000000000000000000000000000000000000000..1ff570ed3bb9dda06c68de5be8a7c105b98694d7 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/pipelined_optimizer_swapper.py @@ -0,0 +1,242 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping optimizer tensors to/from (NVMe) storage devices. +""" + +from deepspeed.ops.op_builder import AsyncIOBuilder +from deepspeed import comm as dist +import torch + +from deepspeed.runtime.swap_tensor.constants import * +from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object +from deepspeed.runtime.swap_tensor.async_swapper import AsyncTensorSwapper +from deepspeed.runtime.swap_tensor.utils import get_sized_buffer +from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper + + +class OptimizerSwapOp(object): + + def __init__(self, aio_handle, read_op, param_info, allocated_buffers, state_buffers, num_ops): + self.aio_handle = aio_handle + self.read_op = read_op + self.param_info = param_info + self.allocated_buffers = allocated_buffers + self.state_buffers = state_buffers + self.wait_required = True + self.num_ops = num_ops + + def is_parameter(self, parameter): + return OptimizerSwapper.parameter_id(parameter) == self.param_info.param_id + + def wait(self): + assert self.wait_required + assert self.aio_handle.wait() == self.num_ops + self.wait_required = False + + +SYNC_SWAP_IN = 'sync_swap_in' +ASYNC_SWAP_IN = 'async_swap_in' +SYNC_SWAP_OUT = 'sync_swap_out' +ASYNC_SWAP_OUT = 'async_swap_out' + +SWAP_IN_STATE_TIMER = 'swap_in_state' +SWAP_OUT_STATE_TIMER = 'swap_out_state' +SWAP_OUT_GRADIENT_TIMER = 'swap_out_gradient' +ASYNC_SWAP_IN_STATE_TIMER = "async_swap_in_state" +ASYNC_SWAP_OUT_STATE_TIMER = 'async_swap_out_state' + + +class PipelinedOptimizerSwapper(OptimizerSwapper): + + def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers): + super(PipelinedOptimizerSwapper, self).__init__(swap_config, aio_config, base_folder, optimizer, largest_numel, + device, dtype, timers) + + aio_op = AsyncIOBuilder().load() + self.write_aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], + queue_depth=aio_config[AIO_QUEUE_DEPTH], + single_submit=aio_config[AIO_SINGLE_SUBMIT], + overlap_events=aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) + + self.read_aio_handle = aio_op.aio_handle(block_size=aio_config[AIO_BLOCK_SIZE], + queue_depth=aio_config[AIO_QUEUE_DEPTH], + single_submit=aio_config[AIO_SINGLE_SUBMIT], + overlap_events=aio_config[AIO_OVERLAP_EVENTS], + intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM]) + + # Overlap gradient swap out + self.gradient_swapper = AsyncTensorSwapper(aio_handle=self.write_aio_handle, + numel_alignment=self.numel_alignment, + timers=self.timers) + + self.async_swap_in = swap_config.pipeline_read + self.async_swap_out = swap_config.pipeline_write + + self.swap_ops = {SYNC_SWAP_IN: None, ASYNC_SWAP_IN: None, SYNC_SWAP_OUT: None, ASYNC_SWAP_OUT: None} + + self.print_exclude_list += [ + 'gradient_swapper', 'read_aio_handle', 'write_aio_handle', 'swap_ops', 'print_exclude_list' + ] + + if dist.get_rank() == 0: + print_object(obj=self, name='PipelinedOptimizerSwapper', exclude_list=self.print_exclude_list) + + def initialize_parameters(self, parameters, src_tensors): + self._initialize_parameters(parameters=parameters, src_tensors=src_tensors, aio_handle=self.write_aio_handle) + + def initialize_from_swapped_fp16_params(self, fp16_partitions_info, fp16_num_elems, fp16_pinned_buffers, + fp32_parameters): + self._initialize_from_swapped_fp16_params(aio_handle=self.write_aio_handle, + fp16_partitions_info=fp16_partitions_info, + fp16_num_elems=fp16_num_elems, + fp16_pinned_buffers=fp16_pinned_buffers, + fp32_parameters=fp32_parameters) + + def flush_gradients(self): + self._flush_gradient_swapper(self.gradient_swapper) + + def swap_in_optimizer_state(self, parameter, async_parameter): + assert parameter is not None + assert self.swap_ops[SYNC_SWAP_IN] is None + + self._flush_gradient_swapper(self.gradient_swapper) + + self._start_timer(SWAP_IN_STATE_TIMER) + + if self.swap_ops[ASYNC_SWAP_IN]: + assert self.swap_ops[ASYNC_SWAP_IN].is_parameter(parameter) + self.swap_ops[SYNC_SWAP_IN] = self.swap_ops[ASYNC_SWAP_IN] + self.swap_ops[ASYNC_SWAP_IN] = None + else: + self.swap_ops[SYNC_SWAP_IN] = self._swap_in_optimizer_state(aio_handle=self.read_aio_handle, + parameter=parameter) + + if self.swap_ops[SYNC_SWAP_IN]: + self.swap_ops[SYNC_SWAP_IN].wait() + + if self.async_swap_in and async_parameter is not None: + assert self.swap_ops[ASYNC_SWAP_IN] is None + self.swap_ops[ASYNC_SWAP_IN] = self._swap_in_optimizer_state(aio_handle=self.read_aio_handle, + parameter=async_parameter) + + self._stop_timer(SWAP_IN_STATE_TIMER) + self.timer_names.add(SWAP_IN_STATE_TIMER) + + def swap_out_optimizer_state(self, parameter, async_swap): + self._start_timer(SWAP_OUT_STATE_TIMER) + + if self.swap_ops[ASYNC_SWAP_OUT]: + self._start_timer(ASYNC_SWAP_OUT_STATE_TIMER) + self._complete_swap_out(ASYNC_SWAP_OUT) + self._stop_timer(ASYNC_SWAP_OUT_STATE_TIMER) + self.timer_names.add(ASYNC_SWAP_OUT_STATE_TIMER) + + assert self.swap_ops[SYNC_SWAP_IN] is not None + assert not self.swap_ops[SYNC_SWAP_IN].wait_required + swap_op = self._swap_out_optimizer_state(aio_handle=self.write_aio_handle, + parameter=parameter, + swap_in_op=self.swap_ops[SYNC_SWAP_IN]) + self.swap_ops[SYNC_SWAP_IN] = None + + if self.async_swap_out and async_swap: + self.swap_ops[ASYNC_SWAP_OUT] = swap_op + else: + self.swap_ops[SYNC_SWAP_OUT] = swap_op + self._complete_swap_out(SYNC_SWAP_OUT) + + self._stop_timer(SWAP_OUT_STATE_TIMER) + self.timer_names.add(SWAP_OUT_STATE_TIMER) + + def swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors): + self._swap_out_gradients(parameter=parameter, + gradient_offsets=gradient_offsets, + gradient_tensors=gradient_tensors, + gradient_swapper=self.gradient_swapper) + + def _complete_swap_out(self, swap_out_type): + self.swap_ops[swap_out_type].wait() + for buffer in self.swap_ops[swap_out_type].state_buffers: + buffer = torch.Tensor() + self.swap_buffer_manager.free(self.swap_ops[swap_out_type].allocated_buffers) + self.swap_ops[swap_out_type] = None + + def _swap_out_optimizer_state(self, aio_handle, parameter, swap_in_op): + assert swap_in_op.is_parameter(parameter) + + allocated_buffers = swap_in_op.allocated_buffers.copy() + swap_buffers = swap_in_op.state_buffers.copy() + + param_info = swap_in_op.param_info + self._update_param_state_info(param_info, parameter) + unpinned_tensors = param_info.get_unpinned_state_tensors() + + if len(unpinned_tensors) > 0: + new_alloc_buffers = self.swap_buffer_manager.allocate(num_elems=self._io_aligned_numel(param_info.numel()), + count=len(unpinned_tensors), + dtype=param_info.dtype()) + assert new_alloc_buffers is not None + + allocated_buffers += new_alloc_buffers + swap_buffers += new_alloc_buffers + + for pinned_dst, unpinned_src in zip(new_alloc_buffers, unpinned_tensors): + dst = get_sized_buffer(pinned_dst, unpinned_src.numel()) + dst.data.copy_(unpinned_src.data) + + swap_paths = param_info.get_swap_paths() + assert len(swap_paths) == len(swap_buffers) + + swap_out_tensors(aio_handle, swap_buffers, swap_paths) + + swap_out_op = OptimizerSwapOp(aio_handle=aio_handle, + param_info=param_info, + read_op=False, + allocated_buffers=allocated_buffers, + state_buffers=swap_buffers, + num_ops=len(swap_buffers)) + + return swap_out_op + + def _swap_in_optimizer_state(self, aio_handle, parameter): + param_info = self._get_param_swap_info(parameter) + if param_info is None: + return None + + num_swap_tensors = param_info.num_tensors() + required_buffer_count = num_swap_tensors + (1 if param_info.has_gradients() else 0) + aligned_numel = self._io_aligned_numel(param_info.numel()) + allocated_buffers = self.swap_buffer_manager.allocate(num_elems=aligned_numel, + count=required_buffer_count, + dtype=parameter.dtype) + assert allocated_buffers is not None, \ + f"PipelinedOptimizerSwapper ran out of swap buffers, try increasing 'buffer_count'" + + state_buffers = allocated_buffers[:num_swap_tensors] + param_info.set_swap_buffers(state_buffers, aligned_numel) + + swap_buffers = state_buffers.copy() + swap_paths = param_info.get_swap_paths() + + if param_info.has_gradients(): + parameter.grad = allocated_buffers[-1].narrow(0, 0, param_info.numel()) + if param_info.swapped_gradients: + swap_buffers += param_info.get_swap_gradient_buffers(parameter.grad) + swap_paths += param_info.get_swap_gradient_paths() + + swap_in_tensors(aio_handle, swap_buffers, swap_paths) + + if param_info.unswapped_gradients: + self._retrieve_unswapped_grad_partitions(swap_info=param_info, dest_buffer=parameter.grad) + + swap_in_op = OptimizerSwapOp(aio_handle=aio_handle, + param_info=param_info, + read_op=True, + allocated_buffers=allocated_buffers, + state_buffers=state_buffers, + num_ops=len(swap_buffers)) + + return swap_in_op diff --git a/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/utils.py b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..60a4004381d9d9080023bc843c1eb425afc51377 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/swap_tensor/utils.py @@ -0,0 +1,241 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Functionality of swapping tensors to/from (NVMe) storage devices. +""" + +import torch +from deepspeed.utils.logging import logger +from deepspeed.accelerator import get_accelerator + +from deepspeed import comm as dist + +MIN_AIO_BYTES = 1024**2 +AIO_ALIGNED_BYTES = 1024 +MIN_SWAPPABLE_BYTES = MIN_AIO_BYTES + + +def swap_in_tensors(swap_handle, tensor_buffers, swap_paths): + for buffer, path in zip(tensor_buffers, swap_paths): + assert (swap_handle.async_pread(buffer, path, 0) == 0) + + +def swap_out_tensors(swap_handle, tensor_buffers, swap_paths): + for buffer, path in zip(tensor_buffers, swap_paths): + assert (swap_handle.async_pwrite(buffer, path, 0) == 0) + + +def print_object(obj, name, exclude_list=[]): + logger.info('{}:'.format(name)) + for arg in sorted(vars(obj)): + if not arg in exclude_list: + dots = '.' * (29 - len(arg)) + logger.info(' {} {} {}'.format(arg, dots, getattr(obj, arg))) + + +class SwapBuffer(object): + + def __init__(self, buffer): + self.buffer = buffer + self.reset() + + def reset(self): + self.offset = 0 + self.swap_tensors = {} + self.compute_tensors = {} + self.swap_paths = {} + self.num_elem = 0 + + def insert_tensor(self, tensor, swap_path, aligned_numel): + swap_tensor, compute_tensor = self.allocate_tensor(swap_path, tensor.numel(), aligned_numel) + compute_tensor.data.copy_(tensor.data) + return swap_tensor, compute_tensor + + def allocate_tensor(self, swap_path, numel, aligned_numel): + assert self.has_space(aligned_numel) + assert not self.offset in self.swap_tensors + + allocate_offset = self.offset + swap_tensor = self.buffer.narrow(0, allocate_offset, aligned_numel) + dest_tensor = swap_tensor.narrow(0, 0, numel) + + self.swap_tensors[allocate_offset] = swap_tensor + self.compute_tensors[allocate_offset] = dest_tensor + self.swap_paths[allocate_offset] = swap_path + self.offset += aligned_numel + self.num_elem += numel + + return self.swap_tensors[allocate_offset], self.compute_tensors[allocate_offset] + + def has_space(self, numel): + return (self.offset + numel) <= self.buffer.numel() + + def get_swap_tensors(self): + return [tensor for tensor in self.swap_tensors.values()] + + def get_swap_paths(self): + return [path for path in self.swap_paths.values()] + + def get_compute_tensors(self): + return [tensor for tensor in self.compute_tensors.values()] + + def get_num_elem(self): + return self.num_elem + + def get_swap_tensor(self, offset): + return self.swap_tensors.get(offset, None) + + def get_compute_tensor(self, offset): + return self.compute_tensors.get(offset, None) + + def get_swap_path(self, offset): + return self.swap_paths(offset, None) + + +class SwapBufferPool(object): + + def __init__(self, buffers): + assert all([get_accelerator().is_pinned(buf) for buf in buffers]) + self.buffers = [SwapBuffer(buf) for buf in buffers] + self.current_index = 0 + + def reset(self): + self.current_index = 0 + for buffer in self.buffers: + buffer.reset() + + def allocate_tensor(self, numel, swap_path, aligned_numel): + if self.has_space(aligned_numel): + swap_tensor, compute_tensor = self._get_current_buffer().allocate_tensor(swap_path, numel, aligned_numel) + return swap_tensor, compute_tensor + + return None, None + + def insert_tensor(self, tensor, swap_path, aligned_numel): + if self.has_space(aligned_numel): + swap_tensor, compute_tensor = self._get_current_buffer().insert_tensor(tensor, swap_path, aligned_numel) + return swap_tensor, compute_tensor + + return None, None + + def get_swap_tensors(self): + swap_tensors = [] + for buffer in self._get_used_buffers(): + swap_tensors += buffer.get_swap_tensors() + + return swap_tensors + + def get_swap_paths(self): + swap_paths = [] + for buffer in self._get_used_buffers(): + swap_paths += buffer.get_swap_paths() + + return swap_paths + + def get_compute_tensors(self): + compute_tensors = [] + for buffer in self._get_used_buffers(): + compute_tensors += buffer.get_compute_tensors() + + return compute_tensors + + def has_space(self, numel): + if self._get_current_buffer().has_space(numel): + return True + + if self.current_index == len(self.buffers) - 1: + return False + + self.current_index += 1 + return self._get_current_buffer().has_space(numel) + + def swap_out(self, aio_handle, async_op=False): + swap_tensors = self.get_swap_tensors() + swap_paths = self.get_swap_paths() + assert all([p is not None for p in swap_paths]) + + swap_out_tensors(aio_handle, swap_tensors, swap_paths) + + if not async_op: + assert len(swap_tensors) == aio_handle.wait() + + def swap_in(self, aio_handle, async_op=False): + swap_tensors = self.get_swap_tensors() + swap_paths = self.get_swap_paths() + assert all([p is not None for p in swap_paths]) + + swap_in_tensors(aio_handle, swap_tensors, swap_paths) + + if not async_op: + assert len(swap_tensors) == aio_handle.wait() + + def _get_current_buffer(self): + return self.buffers[self.current_index] + + def _get_used_buffers(self): + return self.buffers[:self.current_index + 1] + + +class SwapBufferManager(object): + + def __init__(self, num_elems, count, dtype): + self.num_elems = num_elems + self.count = count + self.dtype = dtype + self.all_buffers = [ + get_accelerator().pin_memory(torch.zeros(num_elems, device='cpu', dtype=dtype), align_bytes=0) + for _ in range(count) + ] + self.free_buffer_index = [i for i in range(count)] + self.used_buffer_index = {} + self.gigabytes = (self.all_buffers[0].element_size() * num_elems * count) / (1024**3) + + if dist.get_rank() == 0: + exclude_list = ['all_buffers'] + print_object(obj=self, name='SwapBufferManager', exclude_list=exclude_list) + + def allocate(self, num_elems, count, dtype): + assert dtype == self.dtype + assert num_elems <= self.num_elems + if count > len(self.free_buffer_index): + return None + + used_indices = self.free_buffer_index[-count:] + self.free_buffer_index = self.free_buffer_index[:-count] + + buffers = [] + for i in used_indices: + tmp_buffer = self.all_buffers[i].narrow(0, 0, num_elems) + buffers.append(tmp_buffer) + self.used_buffer_index[id(tmp_buffer)] = i + return buffers + + def allocate_all(self, num_elems, dtype): + return self.allocate(num_elems=num_elems, count=len(self.free_buffer_index), dtype=dtype) + + def free(self, buffers): + buffer_ids = [] + for buf in buffers: + buffer_ids.append(id(buf)) + + assert all([b_id in self.used_buffer_index for b_id in buffer_ids]) + + for b_id in buffer_ids: + self.free_buffer_index.append(self.used_buffer_index[b_id]) + del (self.used_buffer_index[b_id]) + + +def get_sized_buffer(buffer, num_elems): + assert num_elems <= buffer.numel(), \ + f'num_elems {num_elems} > buffer {buffer.numel()}' + return buffer.narrow(0, 0, num_elems) if num_elems < buffer.numel() else buffer + + +def get_sized_buffers(buffer_list, num_elems_list): + swap_buffers = [ + get_sized_buffer(buffer, num_elems) \ + for buffer, num_elems in zip(buffer_list, num_elems_list) + ] + return swap_buffers diff --git a/lib/python3.12/site-packages/deepspeed/runtime/utils.py b/lib/python3.12/site-packages/deepspeed/runtime/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..edd40a8f145dd97fbbcd69d0acb2063ef3b851e1 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/utils.py @@ -0,0 +1,1154 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +""" +Copyright NVIDIA/Megatron + +Helper functions and classes from multiple sources. +""" + +from collections.abc import Iterable +import os +import psutil +import gc +from math import sqrt + +from numpy import prod + +import torch +from torch.nn import functional as F +try: + from torch._six import inf +except ModuleNotFoundError: + from torch import inf +from typing import Union, List, Dict +from deepspeed import comm as dist +from deepspeed.moe.utils import is_moe_param +from deepspeed.utils import groups, logger +from deepspeed.utils.bwc import (bwc_tensor_model_parallel_rank, bwc_pipeline_parallel_world_size, + bwc_pipeline_parallel_group) +from deepspeed.runtime.constants import PIPE_REPLICATED +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.policy import transpose + +torch_memory_reserved = get_accelerator().memory_reserved +torch_max_memory_reserved = get_accelerator().max_memory_reserved + + +class DummyOptim(): + """ + Dummy optimizer presents model parameters as a param group, this is + primarily used to allow ZeRO-3 without an optimizer + """ + + def __init__(self, params): + self.param_groups = [] + self.param_groups.append({'params': params}) + + +graph_cache = {} + + +def graph_process(replay_first_step, func, *args, **kwargs): + # `func` should only contain operations on the GPU + # Please ensure that the memory address of the data required by 'func' remains constant + if func.__name__ not in graph_cache: + cuda_stream = get_accelerator().Stream() + cuda_stream.wait_stream(get_accelerator().current_stream()) + with get_accelerator().stream(cuda_stream): + func(*args, **kwargs) + get_accelerator().current_stream().wait_stream(cuda_stream) + graph_cache[func.__name__] = get_accelerator().create_graph() + with get_accelerator().capture_to_graph(graph_cache[func.__name__]): + func(*args, **kwargs) + if replay_first_step: + get_accelerator().replay_graph(graph_cache[func.__name__]) + else: + get_accelerator().replay_graph(graph_cache[func.__name__]) + + +def noop_decorator(func): + return func + + +class noop_context(object): + + def __init__(self): + pass + + def __enter__(self): + pass + + def __exit__(self, exc_type, exc_val, exc_tb): + pass + + +def ensure_directory_exists(filename): + """Create the directory path to ``filename`` if it does not already exist. + + Args: + filename (str): A file path. + """ + dirname = os.path.dirname(filename) + os.makedirs(dirname, exist_ok=True) + + +def set_random_seed(seed): + """Set the random seed for common PRNGs used during training: random, numpy, and torch. + + Args: + seed (int): the seed to use + """ + import numpy + import random + random.seed(seed) + numpy.random.seed(seed) + torch.manual_seed(seed) + + +def is_model_parallel_parameter(p) -> bool: + if hasattr(p, 'model_parallel') and p.model_parallel: + return True + + if hasattr(p, 'tensor_model_parallel') and p.tensor_model_parallel: + return True + + return False + + +def copy_to_device(item, device, criterion_func): + """ + Return a copy of tensor on specified device. + Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. + Parameters: + item: tensor to copy or (possibly nested) container of tensors to copy. + device: target device + criterion_func: Function to restrict copy operation to items meet criterion + + Returns: + None + """ + if criterion_func(item): + return item.to(device) + elif isinstance(item, list): + return [copy_to_device(v, device, criterion_func) for v in item] + elif isinstance(item, tuple): + return tuple([copy_to_device(v, device, criterion_func) for v in item]) + elif isinstance(item, dict): + return {k: copy_to_device(v, device, criterion_func) for k, v in item.items()} + else: + return item + + +def move_to_device(item, device, criterion_func=None): + """ + Move tensor on to specified device by changing the storage. + Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. + Parameters: + item: tensor to move or (possibly nested) container of tensors to move. + device: target device + criterion_func: Function to restrict move operation to items meet criterion, defaults to `None` which is an equivalent to always move + + Returns: + None + """ + if (criterion_func is not None and criterion_func(item)): + device_copy = item.to(device) + item.data = device_copy.data + return item + elif isinstance(item, list): + return [move_to_device(v, device, criterion_func) for v in item] + elif isinstance(item, tuple): + return tuple([move_to_device(v, device, criterion_func) for v in item]) + elif isinstance(item, dict): + return {k: move_to_device(v, device, criterion_func) for k, v in item.items()} + else: + return item.to(device) + + +def get_norm_with_moe_layers_fast(all_groups_norm, group): + # This implementation standardizes the grad_norm across ranks. A more precise implementation can be found in 'get_norm_with_moe_layers'. + # Need to allreduce (avg) the norms across different ranks because moe params will not be synced during allreduce + scaled_norm = all_groups_norm * 1.0 / float(dist.get_world_size(group=group)) + scaled_norm_tensor = torch.tensor(scaled_norm, device=get_accelerator().current_device_name(), dtype=torch.float) + dist.all_reduce(scaled_norm_tensor, group=group) + all_groups_norm = scaled_norm_tensor.item() + #print(f"old = {all_groups_norm_old} and new = {all_groups_norm} at rank: {deepspeed.comm.get_rank()}") + return all_groups_norm + + +class CheckOverflow(object): + '''Checks for overflow in gradient across parallel process''' + + def __init__(self, param_groups=None, mpu=None, zero_reduce_scatter=False, deepspeed=None): + self.mpu = mpu + self.params = [] if param_groups else None + self.zero_reduce_scatter = zero_reduce_scatter + self.deepspeed = deepspeed + self.has_moe_params = False + if param_groups: + for group in param_groups: + for param in group: + self.params.append(param) + if is_moe_param(param): + self.has_moe_params = True + + def check_using_norm(self, norm_group, reduce_overflow=True): + # TODO: I don't think reduce_overflow is needed if mpu is None + overflow = -1 in norm_group + overflow_gpu = get_accelerator().FloatTensor([overflow]) + if self.has_moe_params: + # In this case, we need to do an all_reduce across + # the expert_parallel_group, so that if there was + # an overflow due to expert weights, we detect it + + # Only need to check groups.get_largest_expert_parallel_group() + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) + if self.mpu is not None: + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) + elif reduce_overflow: + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX) + dist.barrier() + overflow = overflow_gpu[0].item() + return bool(overflow) + + def check(self, param_groups=None): + params = [] + has_moe_params = False + if param_groups is None: + params = self.params + has_moe_params = self.has_moe_params + else: + assert param_groups is not None, \ + "self.params and param_groups both cannot be none" + + for group in param_groups: + for param in group: + params.append(param) + if is_moe_param(param): + has_moe_params = True + + return self.has_overflow(params, has_moe_params=has_moe_params) + + # `params` is a list / generator of torch.Variable + def has_overflow_serial(self, params): + for i, p in enumerate(params): + if p.grad is not None and self._has_inf_or_nan(p.grad.data, i): + return True + return False + + def has_overflow(self, params, has_moe_params=None): + if has_moe_params is None: + has_moe_params = self.has_moe_params + overflow = self.has_overflow_serial(params) + # Since each model parallel GPU carries only part of the model, + # make sure overflow flag is synced across all the model parallel GPUs + overflow_gpu = get_accelerator().ByteTensor([overflow]) + # deepspeed.comm.all_reduce(overflow_gpu, + # op=deepspeed.comm.ReduceOp.MAX, + # group=mpu.get_model_parallel_group()) + if has_moe_params: + # All reduce this across expert_parallel_group, so that if an expert + # overflows, we detect it here + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) + if self.zero_reduce_scatter: + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group()) + elif self.mpu is not None: + if self.deepspeed is not None: + using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce') + if (using_pipeline and self.deepspeed.pipeline_enable_backward_allreduce + is False) or (not using_pipeline and self.deepspeed.enable_backward_allreduce is False): + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_data_parallel_group()) + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) + elif self.deepspeed is not None and self.deepspeed.enable_backward_allreduce is False: + dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group()) + + overflow = overflow_gpu[0].item() + return bool(overflow) + + # `x` is a torch.Tensor + @staticmethod + def _has_inf_or_nan(x, i): + try: + # if x is half, the .float() incurs an additional deep copy, but it's necessary if + # Pytorch's .sum() creates a one-element tensor of the same type as x + # (which is true for some recent version of pytorch). + cpu_sum = float(x.float().sum()) + # More efficient version that can be used if .sum() returns a Python scalar + # cpu_sum = float(x.sum()) + except RuntimeError as instance: + # We want to check if inst is actually an overflow exception. + # RuntimeError could come from a different error. + # If so, we still want the exception to propagate. + if "value cannot be converted" not in instance.args[0]: + raise + return True + else: + if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum: + return True + return False + + +def _handle_overflow(cpu_sum, x, i): + import math + rank = dist.get_rank() + if rank == 0: + t_i = -1 + for v_i, v in enumerate(x.data.contiguous().view(-1)): + if not math.isfinite(float(v)): + t_i = v_i + break + logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}") + + +def get_global_norm(norm_list): + """ Compute total from a list of norms + """ + total_norm = 0.0 + for norm in norm_list: + total_norm += norm**2.0 + # logger.info(f'norm_list = {norm_list} global = {sqrt(total_norm)}') + return sqrt(total_norm) + + +def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None): + """Clips gradient norm of an iterable of parameters. + + This has been adapted from Nvidia megatron. We add norm averaging + to consider MoE params when calculating norm as they will result + in different norms across different ranks. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + max_norm (float or int): max norm of the gradients + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the parameters (viewed as a single vector). + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = list(filter(lambda p: p.grad is not None, parameters)) + norm_type = float(norm_type) + all_norms = [] + if norm_type == inf: + for p in parameters: + all_norms.append(p.grad.data.abs().max().float()) + total_norm = torch.stack(all_norms).max() + total_norm = total_norm.to(get_accelerator().current_device_name()) + # Take max across all GPUs. + if mpu is not None: + dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) + else: + total_norm = 0 + for p in parameters: + if mpu is not None: + if (mpu.get_model_parallel_rank() == 0) or is_model_parallel_parameter(p): + param_norm = p.grad.data.detach().float().norm(norm_type) + all_norms.append(param_norm) + else: + param_norm = p.grad.data.detach().float().norm(norm_type) + all_norms.append(param_norm) + if len(all_norms) > 0: + total_norm = torch.stack(all_norms).square().sum().float() + else: + total_norm = get_accelerator().FloatTensor([0.0]) + total_norm = total_norm.to(get_accelerator().current_device_name()) + # Sum across all model parallel GPUs. + if mpu is not None: + dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + total_norm = total_norm.pow(1. / norm_type) + + # Need to average total_norm across different GPUs due to the presence of moe params + pg = groups._get_data_parallel_group() + scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg)) + scaled_norm_tensor = scaled_norm + + dist.all_reduce(scaled_norm_tensor, group=pg) + total_norm = scaled_norm_tensor + total_norm = total_norm.to(parameters[0].device) + + max_norm = torch.tensor([float(max_norm)], device=total_norm.device) + clip_coef = max_norm / (total_norm + 1e-6) + tmp_tensor = torch.tensor([1.0], device=clip_coef.device) + clip_coef = torch.min(tmp_tensor, clip_coef) + for p in parameters: + p.grad.data.mul_(clip_coef) + return total_norm + + +def get_flattened_grad_norm(parameters, norm_type=2, mpu=None, grad_norm_mask=None): + """Get grad norm of an iterable of parameters. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. Taken from Nvidia Megatron. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + grad_norm_mask (List[Tensor]): A list of Tensor, where + each Tensor is a 2D Tensor containing ranges of [start_index, end_index]. + Returns: + Total norm of the parameters (viewed as a single vector). + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = list(filter(lambda p: p.grad is not None, parameters)) + + norm_type = float(norm_type) + if norm_type == inf: + total_norm = max(p.grad.data.abs().max() for p in parameters) + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + # Take max across all GPUs. + if mpu is not None: + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) + total_norm = total_norm_cuda[0].item() + else: + total_norm = 0. + for idx, p in enumerate(parameters): + # Use grad_norm_mask to avoid redundant computation of flattened gradient norm + if grad_norm_mask is not None and len(grad_norm_mask[idx]) > 0: + + # A loop-free implementation to create a mask tensor based on a range list + # which is logically equivalent to the following implementation. + # # mask_tensor_ = torch.zeros_like(p, device=p.device, dtype=bool) + # # for mask_idx in grad_norm_mask[idx]: + # # mask_tensor_[mask_idx[0]:mask_idx[1]] = True + cum_sum_pairs = torch.tensor([1, -1], device=get_accelerator().current_device_name(), + dtype=p.dtype).repeat(grad_norm_mask[idx].shape[0], 1) + mask_tensor = torch.zeros(p.shape[0] + 1, + device=get_accelerator().current_device_name(), + dtype=p.dtype) + mask_tensor = mask_tensor.scatter_(0, grad_norm_mask[idx].view(-1), + cum_sum_pairs.view(-1)).cumsum(0).bool()[:-1] + + param_norm = torch.masked_fill(p.grad.data, mask_tensor, 0).float().norm(norm_type) + + else: + param_norm = p.grad.data.float().norm(norm_type) + total_norm += param_norm.item()**norm_type + + # Sum across all model parallel GPUs. + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + if mpu is not None: + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + total_norm = total_norm_cuda[0].item()**(1. / norm_type) + + if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: + total_norm = -1 + + return total_norm + + +def get_grad_zeros(parameters, mpu=None): + """Compute the number of grads with zero values. + + This is adapted from get_grad_norm + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + + Returns: + Total number of params with zero values (viewed as a single vector). + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = list(filter(lambda p: p.grad is not None, parameters)) + + total_zeros = 0. + tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) + for p in parameters: + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + continue + + # Filter to avoid over-counting replicated tensors from tensor + # model parallelism + if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): + continue + + count_zeros = p.grad.numel() - torch.count_nonzero(p.grad) + total_zeros += count_zeros.item() + + # Sum across all model parallel GPUs. + total_zeros_cuda = get_accelerator().FloatTensor([float(total_zeros)]) + if mpu is not None: + dist.all_reduce(total_zeros_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + total_zeros = total_zeros_cuda[0].item() + + return total_zeros + + +def get_weight_norm(parameters, norm_type=2, mpu=None): + """Get norm of an iterable of parameters. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Note that + the gradients are modified in place. Taken from Nvidia Megatron. + + Arguments: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the parameters (viewed as a single vector). + -1 if the norm value is NaN or Inf. + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + + norm_type = float(norm_type) + if norm_type == inf: + total_norm = max(p.data.abs().max() for p in parameters) + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + # Take max across all GPUs. + if mpu is not None: + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) + total_norm = total_norm_cuda[0].item() + else: + total_norm = 0. + tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu) + for p in parameters: + # Pipeline parallelism may replicate parameters. Avoid multi-counting. + if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated: + continue + + # Filter to avoid over-counting replicated tensors from tensor + # model parallelism + if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p): + continue + + param_norm = p.data.float().norm(norm_type) + total_norm += param_norm**norm_type + + # Sum across all model parallel GPUs. + total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]) + if mpu is not None: + dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + total_norm = total_norm_cuda[0].item()**(1. / norm_type) + + if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm: + total_norm = -1 + + return total_norm + + +def prefix_sum_inc(weights): + """ Compute an inclusive prefix sum. + + Example: + >>> prefix_sum_inc([3,4,5]) + [3, 7, 12] + """ + weights_ = [w for w in weights] + for x in range(1, len(weights_)): + weights_[x] += weights_[x - 1] + return weights_ + + +def partition_uniform(num_items, num_parts): + import numpy + parts = [0] * (num_parts + 1) + # First check for the trivial edge case + if num_items <= num_parts: + for p in range(num_parts + 1): + parts[p] = min(p, num_items) + return parts + + chunksize = num_items // num_parts + residual = num_items - (chunksize * num_parts) + + parts = numpy.arange(0, (num_parts + 1) * chunksize, chunksize) + + for i in range(residual): + parts[i + 1:] += 1 + parts = parts.tolist() + + return parts + + +def partition_balanced(weights, num_parts): + """ + use dynamic programming solve `The Linear Partition Problem`. + see https://www8.cs.umu.se/kurser/TDBAfl/VT06/algorithms/BOOK/BOOK2/NODE45.HTM + """ + import numpy as np + n = len(weights) + m = num_parts + + if n <= m: + return partition_uniform(n, m) + + dp_max = np.full((n + 1, m + 1), np.inf) + dp_min = np.full((n + 1, m + 1), np.inf) + dp_cost = np.full((n + 1, m + 1), np.inf) + position = np.zeros((n + 1, m + 1), dtype=int) + prefix_sum = np.zeros((n + 1)) + prefix_sum[1:] = np.cumsum(weights) + + dp_max[0, 0] = 0 + dp_cost[0, 0] = 0 + for i in range(1, n + 1): + for j in range(1, min(i, m) + 1): + for k in range(i): + max_sum = max(dp_max[k, j - 1], prefix_sum[i] - prefix_sum[k]) + min_sum = min(dp_min[k, j - 1], prefix_sum[i] - prefix_sum[k]) + cost = max_sum - min_sum + if dp_cost[i, j] >= cost: + dp_cost[i, j] = cost + dp_max[i, j] = max_sum + dp_min[i, j] = min_sum + position[i, j] = k + + parts = [n] + for i in reversed(range(1, m + 1)): + parts.append(position[parts[-1], i]) + parts.reverse() + + return parts + + +class PartitionedTensor: + + def __init__(self, tensor, group, partition_meta=None): + super().__init__() + + self.group = group + self.num_parts = dist.get_world_size(group=self.group) + self.rank = dist.get_rank(group=self.group) + self.orig_size = list(tensor.size()) + self.orig_device = tensor.device + self.local_data, self.partition = self._partition_tensor(tensor) + self.even_split = tensor.numel() % self.num_parts == 0 + + @classmethod + def from_meta(cls, meta, local_part, group, device=get_accelerator().device_name()): + assert meta.dtype == torch.long + dummy = torch.ones(dist.get_world_size(group=group)) + part_obj = cls(tensor=dummy, group=group) + + meta = meta.tolist() + + # [N, list0, ..., listN-1] + part_obj.orig_size = meta[1:(1 + meta[0])] + meta = meta[1 + meta[0]:] + + part_obj.orig_device = device + part_obj.local_data = local_part.detach() + + part_obj.group = group + + # Partition is encoded like the rowptr of a CSR matrix: + # [num_parts, rank, 0, part_1, ..., part_num_parts] + # TODO: support shuffle between different partition granularities + assert part_obj.num_parts == meta[0] + assert part_obj.rank == meta[1] + part_obj.partition = meta[2:] # length num_parts+1 + + return part_obj + + def _partition_tensor(self, tensor): + partition = partition_uniform(num_items=tensor.numel(), num_parts=self.num_parts) + start = partition[self.rank] + length = partition[self.rank + 1] - start + tensor_part = tensor.detach().contiguous().view(-1).narrow(0, start=start, length=length).clone() + + return tensor_part, partition + + def full(self, device=None): + if device is None: + device = self.orig_device + + # Allocate the full tensor as a flat buffer. + full_numel = prod(self.full_size()) + flat_tensor = torch.zeros([full_numel], dtype=self.local_data.dtype, device=device) + if self.even_split: + # Collect the full tensor + dist.all_gather_into_tensor(flat_tensor, self.local_data, group=self.group) + else: + for part_id in range(self.num_parts): + part_size = self.partition[part_id + 1] - self.partition[part_id] + buf = flat_tensor.narrow(0, start=self.partition[part_id], length=part_size) + if part_id == self.rank: + buf.copy_(self.local_data) + dist.broadcast(buf, part_id, self.group) + return flat_tensor.view(self.full_size()).clone().detach() + + def to_meta(self): + """Returns a torch.LongTensor that encodes partitioning information. + + Can be used along with ``data()`` to serialize a ``PartitionedTensor`` for + communication. + + Returns: + torch.LongTensor: a tensor encoding the meta-information for the partitioning + """ + meta = [] + meta.append(len(self.orig_size)) + meta += list(self.orig_size) + meta.append(self.num_parts) + meta.append(self.rank) + meta += self.partition + return torch.LongTensor(data=meta).to(self.orig_device) + + def data(self): + return self.local_data + + def local_size(self): + return self.local_data.size() + + def full_size(self): + return self.orig_size + + +mem_alloced = 0 +mem_cached = 0 + + +def memory_status(msg, print_rank=-1, reset_max=False): + global mem_alloced, mem_cached + + rank = dist.get_rank() + if print_rank != -1 and rank != print_rank: + return + + get_accelerator().synchronize() + + if reset_max: + get_accelerator().reset_max_memory_cached() + get_accelerator().reset_max_memory_allocated() + + new_alloced = get_accelerator().memory_allocated() + new_cached = get_accelerator().memory_cached() + + delta_alloced = new_alloced - mem_alloced + delta_cached = new_cached - mem_cached + + mem_cached = new_cached + mem_alloced = new_alloced + + max_alloced = get_accelerator().max_memory_allocated() + max_cached = get_accelerator().max_memory_cached() + + # convert to GB for printing + new_alloced /= 1024**3 + new_cached /= 1024**3 + delta_alloced /= 1024**3 + delta_cached /= 1024**3 + max_alloced /= 1024**3 + max_cached /= 1024**3 + + print( + f'RANK={rank} MEMSTATS', msg, f'device={get_accelerator().current_device_name()} ' + f'current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) ' + f'current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)') + + +def get_ma_status(): + if dist.is_initialized() and not dist.get_rank() == 0: + return 0 + return get_accelerator().memory_allocated() + + +def empty_cache(): + get_accelerator().empty_cache() + get_accelerator().reset_peak_memory_stats() + + +def see_memory_usage(message, force=False): + if not force: + return + if dist.is_initialized() and not dist.get_rank() == 0: + return + + # python doesn't do real-time garbage collection so do it explicitly to get the correct RAM reports + gc.collect() + + # Print message except when distributed but not rank 0 + logger.info(message) + logger.info(f"MA {round(get_accelerator().memory_allocated() / (1024 * 1024 * 1024),2 )} GB \ + Max_MA {round(get_accelerator().max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \ + CA {round(torch_memory_reserved() / (1024 * 1024 * 1024),2)} GB \ + Max_CA {round(torch_max_memory_reserved() / (1024 * 1024 * 1024))} GB ") + + vm_stats = psutil.virtual_memory() + used_GB = round(((vm_stats.total - vm_stats.available) / (1024**3)), 2) + logger.info(f'CPU Virtual Memory: used = {used_GB} GB, percent = {vm_stats.percent}%') + + # get the peak memory to report correct data, so reset the counter for the next call + get_accelerator().reset_peak_memory_stats() + + +def call_to_str(base, *args, **kwargs): + """Construct a string representation of a call. + + Args: + base (str): name of the call + args (tuple, optional): args to ``base`` + kwargs (dict, optional): kwargs supplied to ``base`` + + Returns: + str: A string representation of base(*args, **kwargs) + """ + name = f'{base}(' + if args: + name += ', '.join(repr(arg) for arg in args) + if kwargs: + name += ', ' + if kwargs: + name += ', '.join(f'{key}={repr(arg)}' for key, arg in kwargs.items()) + name += ')' + return name + + +def get_only_unique_item(items): + item_set = set(items) + if len(item_set) != 1: + raise RuntimeError(f"expected there to be only one unique element in {items}") + unique_item, = item_set + + return unique_item + + +def mask_nan_or_inf_with_val_inplace(input, device=None, val=-1.): + norm_is_inf = input.isinf() + norm_is_nan = input.isnan() + inf_or_nan = norm_is_nan.logical_or(norm_is_inf) + err = torch.tensor(-1.0, device=device, dtype=torch.float) + input.masked_fill_(inf_or_nan, err) + + +def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, use_graph=False, moe_ep_group=None): + """Get norm of an iterable of tensors. + + This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and + added functionality to handle model parallel parameters. Taken from Nvidia Megatron. + + Arguments: + input_tensors (Iterable[Tensor]): an iterable of Tensors will have norm computed + norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + + Returns: + Total norm of the tensors (viewed as a single vector). + """ + assert isinstance(input_tensors, Iterable), f'expected Iterable type not {type(input_tensors)}' + assert all([torch.is_tensor(t) for t in input_tensors]), f'expected list of only tensors' + + norm_type = float(norm_type) + all_norms = [] + if norm_type == inf: + for t in input_tensors: + all_norms.append(t.data.abs().max().float()) + total_norm = torch.stack(all_norms).max() + device_total_norm = total_norm.to(get_accelerator().current_device_name()) + # Max across model parallel + if mpu is not None: + # For MoE grads, max over model parallel only if MoE-TP is enabled + if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) + # If MoE grads and MoE-TP disabled, max over pipeline parallel + elif bwc_pipeline_parallel_world_size(mpu) > 1: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=bwc_pipeline_parallel_group(mpu)) + + # MoE grads: max across expert parallel group + if moe_ep_group is not None: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.MAX, group=moe_ep_group) + total_norm = device_total_norm.to(input_tensors[0].device) + else: + + if 'norm_tensors_compute_buffer' not in graph_cache or len( + graph_cache['norm_tensors_compute_buffer']) != len(input_tensors): + graph_cache['norm_tensors_compute_buffer'] = [ + torch.empty([], dtype=torch.float, device=get_accelerator().current_device_name()) + for t in input_tensors + ] + compute_buffer = graph_cache['norm_tensors_compute_buffer'] + + def _norm_tensors(tensor_list, _compute_buffer, _norm_type): + for i, t in enumerate(tensor_list): + _compute_buffer[i].data.copy_(t.data.float().norm(_norm_type)**_norm_type) + if i != 0: + _compute_buffer[0].data.add_(_compute_buffer[i].data) + + if use_graph: + graph_process(False, _norm_tensors, input_tensors, compute_buffer, norm_type) + else: + _norm_tensors(input_tensors, compute_buffer, norm_type) + + device_total_norm = compute_buffer[0].float().detach() + + # Sum across model parallel + if mpu is not None: + # For MoE grads, sum over model parallel only if MoE-TP is enabled + if moe_ep_group is None or groups._get_expert_model_parallel_world_size() > 1: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) + # If MoE grads and MoE-TP disabled, sum over pipeline parallel + elif bwc_pipeline_parallel_world_size(mpu) > 1: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=bwc_pipeline_parallel_group(mpu)) + + # MoE grads: sum across expert parallel group + if moe_ep_group is not None: + dist.all_reduce(device_total_norm, op=dist.ReduceOp.SUM, group=moe_ep_group) + total_norm = device_total_norm.to(input_tensors[0].device).pow(1. / norm_type) + + mask_nan_or_inf_with_val_inplace(total_norm, device=total_norm.device) + + return total_norm + + +def clip_tensors_by_global_norm(input_tensors, max_norm=1.0, global_norm=None, mpu=None, eps=1e-6, use_graph=False): + """Clip list of tensors by global norm. + Args: + input_tensors: List of tensors to be clipped + global_norm (float, optional): Precomputed norm. Defaults to None. + mpu (optional): model parallelism unit. Defaults to None. + eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6 + Returns: + float: the global norm + """ + if global_norm is None: + global_norm = get_global_norm_of_tensors(input_tensors, mpu=mpu, use_graph=use_graph) + clip_coef = max_norm / (global_norm + eps) + if clip_coef < 1: + if use_graph: + + def clip_tensors(_tensor_list, _clip_coef_tensor): + for t in _tensor_list: + t.detach().mul_(_clip_coef_tensor) + + if 'clip_coef_tensor' not in graph_cache: + # Alloc memory + graph_cache['clip_coef_tensor'] = torch.tensor(clip_coef, + dtype=torch.float32).to(get_accelerator().device_name()) + clip_coef_tensor = graph_cache['clip_coef_tensor'] + clip_coef_tensor.copy_(torch.tensor(clip_coef, dtype=torch.float32)) + graph_process(False, clip_tensors, input_tensors, clip_coef_tensor) + + else: + for t in input_tensors: + t.detach().mul_(clip_coef) + return global_norm + + +def align_dense_tensors(tensor_list, alignment): + num_elements = sum(t.numel() for t in tensor_list) + remaining = num_elements % alignment + + if remaining: + elements_to_add = alignment - remaining + pad_tensor = torch.zeros(elements_to_add, device=tensor_list[0].device, dtype=tensor_list[0].dtype) + padded_tensor_list = tensor_list + [pad_tensor] + else: + padded_tensor_list = tensor_list + + return padded_tensor_list + + +def all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, dp_process_group): + for group_id, (group_flat, partitioned_params) in enumerate(zip(groups_flat, partitioned_param_groups)): + partition_id = dist.get_rank(group=dp_process_group[group_id]) + dp_world_size = dist.get_world_size(group=dp_process_group[group_id]) + if dp_world_size == 1: + # no groups share optimizer states + # pipeline parallel with bf16 will default call this even if dp size = 1. + continue + dist.all_gather_into_tensor(group_flat, partitioned_params[partition_id], dp_process_group[group_id]) + + +def all_gather_dp_groups(groups_flat, partitioned_param_groups, dp_process_group, start_alignment_factor, + allgather_bucket_size): + if dist.has_all_gather_into_tensor(): + return all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, dp_process_group) + + for group_id, partitioned_params in enumerate(partitioned_param_groups): + # Sequential AllGather Best of both worlds + partition_id = dist.get_rank(group=dp_process_group[group_id]) + dp_world_size = dist.get_world_size(group=dp_process_group[group_id]) + + if dp_world_size == 1: + # no groups share optimizer states + # pipeline parallel with bf16 will default call this even if dp size = 1. + continue + num_shards = max(1, partitioned_params[partition_id].numel() * dp_world_size // allgather_bucket_size) + + shard_size = partitioned_params[partition_id].numel() // num_shards + + # Enforce nccl/rccl alignment of start location of each shard + shard_size = shard_size - (shard_size % start_alignment_factor) + + num_elements = shard_size + + assert shard_size * num_shards <= partitioned_params[partition_id].numel() + + for shard_id in range(num_shards): + + if shard_id == (num_shards - 1): + num_elements = partitioned_params[partition_id].numel() - shard_id * shard_size + + shard_list = [] + for dp_id in range(dp_world_size): + curr_shard = partitioned_params[dp_id].narrow(0, shard_id * shard_size, num_elements).detach() + shard_list.append(curr_shard) + + dist.all_gather(shard_list, shard_list[partition_id], dp_process_group[group_id]) + + +class TLinear(torch.nn.Linear): + + def __init__(self, orig_layer, name=""): + self.name = name + super().__init__(orig_layer.weight.shape[1], orig_layer.weight.shape[0], bias=(orig_layer.bias is not None)) + self.weight.data = transpose(orig_layer.weight.data) + self.bias = orig_layer.bias + self._fwd_func = self._fwd_bias_add if self.bias is not None else self._fwd + + def _fwd(self, input): + return F.linear(input, self.weight) + + def _fwd_bias_add(self, input): + return F.linear(input, self.weight, bias=self.bias) + + def forward(self, input): + return self._fwd_func(input) + + +def get_inactive_params(param_list): + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + return [param for param in param_list if (hasattr(param, 'ds_id') and \ + param.ds_status == ZeroParamStatus.NOT_AVAILABLE)] + + +def get_norm_with_moe_layers(non_expert_norm, mpu, expert_tensors, norm_type=2): + """ Compute the global norm with MoE experts + + Inputs: + non_expert_norm (float) : the calculated norm of the non-expert params + expert_tensors (Dict[ep_name, List[Tensor]): Dictionary of expert group name to list of grad tensors + norm_type (int): the norm to use + + Returns: + if norm is (-/+) inf, returns -1 + otherwise the global norm (float) + """ + + def to_tensor(v): + return get_accelerator().FloatTensor(float(v)).detach() + + group_norms = [non_expert_norm] + for exp_name, tensors in expert_tensors.items(): + group_norm = get_global_norm_of_tensors(input_tensors=tensors, + mpu=mpu, + norm_type=norm_type, + use_graph=False, + moe_ep_group=groups._get_expert_parallel_group(exp_name)) + group_norms.append(group_norm) + + # check if all norms are valid + group_norms = torch.stack([to_tensor(norm) for norm in group_norms]) + if group_norms.eq(-1).any(): + return -1 + + # combine norms + if norm_type == inf: + total_norm = group_norms.max().item() + else: + total_norm = group_norms.pow(norm_type).sum() + total_norm = total_norm.item()**(1. / norm_type) + if total_norm == float('inf') or total_norm == -float('inf'): + total_norm = -1 + + return total_norm + + +def _make_offload_state_key(key): + return f"{key}_offload_buffer" + + +def offload_adam_states(optimizer, device, pin_memory: bool = False, non_blocking: bool = False): + """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + + def move_key(state, key): + offload_buf_key = _make_offload_state_key(key) + if offload_buf_key not in state: + state[offload_buf_key] = torch.empty_like(state[key], device=device) + if pin_memory: + state[offload_buf_key] = get_accelerator().pin_memory(state[offload_buf_key]) + state[offload_buf_key].copy_(state[key], non_blocking=non_blocking) + state[key].data = state[offload_buf_key] + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + move_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_key(state, "exp_avg_sq") + + +def reload_adam_states(optimizer, device, non_blocking: bool = False): + """Move optimizer states to device. Note that this assumes the state structure of DeepSpeed Adam.""" + + def move_back_key(state, key): + state[key].data = state[_make_offload_state_key(key)].to(device, non_blocking=non_blocking) + + for _, state in optimizer.state.items(): + if "exp_avg" in state: + move_back_key(state, "exp_avg") + if "exp_avg_sq" in state: + move_back_key(state, "exp_avg_sq") + + +def compare_tensors_in_structures(inputs1: Union[List, Dict], inputs2: Union[List, Dict]) -> bool: + """ + Compare two lists or dictionaries for equality, including any tensors they may contain. + + Args: + inputs1: First input, either a list or a dictionary. + inputs2: Second input, either a list or a dictionary. + + Returns: + True if inputs1 and inputs2 are equal; False otherwise. + """ + if type(inputs1) != type(inputs2): # Ensure types match + return False + + if isinstance(inputs1, list) and isinstance(inputs2, list): + if len(inputs1) != len(inputs2): + return False + for val1, val2 in zip(inputs1, inputs2): + if isinstance(val1, torch.Tensor) and isinstance(val2, torch.Tensor): + val1 = val1.to(get_accelerator().current_device()) + val2 = val2.to(get_accelerator().current_device()) + if not torch.equal(val1, val2): + return False + elif val1 != val2: + return False + return True + + elif isinstance(inputs1, dict) and isinstance(inputs2, dict): + if inputs1.keys() != inputs2.keys(): + return False + for key in inputs1: + val1, val2 = inputs1[key], inputs2[key] + if isinstance(val1, torch.Tensor) and isinstance(val2, torch.Tensor): + val1 = val1.to(get_accelerator().current_device()) + val2 = val2.to(get_accelerator().current_device()) + if not torch.equal(val1, val2): + return False + elif val1 != val2: + return False + return True + + return False diff --git a/lib/python3.12/site-packages/deepspeed/runtime/weight_quantizer.py b/lib/python3.12/site-packages/deepspeed/runtime/weight_quantizer.py new file mode 100644 index 0000000000000000000000000000000000000000..95d737614e594ee200d13617654714d6ae9d26f6 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/runtime/weight_quantizer.py @@ -0,0 +1,153 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from ..module_inject.replace_policy import HFBertLayerPolicy, replace_policies +from deepspeed.accelerator import get_accelerator + + +class WeightQuantization(object): + + def __init__(self, mlp_extra_grouping=True, mp_size=1): + self.dense_scales = [] + self.qkv_scales = [] + self.mlp4hh_scales = [] + self.mlph4h_scales = [] + self.mlp_extra_grouping = mlp_extra_grouping + self.mp_size = mp_size + + def quantize_data(self, data, quantize_bits, groups, key=None): + data_groups = torch.split(data.float().view(-1), data.numel() // groups) + max_d = [max(g.max(), g.min().abs()) for g in data_groups] + data_scale = [float(1 << quantize_bits) / (2 * mx + 1e-5) for mx in max_d] + data_int = [(g * s) for g, s in zip(data_groups, data_scale)] + data_int = [ + di.round().clamp(-(1 << (quantize_bits - 1)), (((1 << (quantize_bits - 1)) - 1))) for di in data_int + ] + data_int = torch.cat(data_int).reshape(data.shape) + data_int = data_int.to(torch.int8) + data_scale = torch.cat([s.unsqueeze(0).unsqueeze(0) for s in data_scale]) + return data_int, data_scale + + def is_mlp(self, data, merge_count=1): + return ((self.mp_size *data.shape[0] * merge_count) / data.shape[1] == 4 or \ + (self.mp_size *data.shape[1] * merge_count) / data.shape[0] == 4) + + def is_qkv(self, data): + return ((self.mp_size * data.shape[0]) / data.shape[1] == 3 or \ + (self.mp_size * data.shape[1]) / data.shape[0] == 3) + + def Quantize(self, value_list, quantize_bits, groups, key, merge_dim=0): + if self.mlp_extra_grouping and self.is_mlp(value_list[0], merge_count=len(value_list)): + groups *= 2 + q_scale = [] + index = 0 + for data in value_list: + data_int, data_scale = self.quantize_data(data, quantize_bits, groups, key) + q_scale.append(data_scale) + value_list[index] = data_int + index += 1 + q_scale = (1 / + torch.cat(q_scale, dim=merge_dim).to(get_accelerator().current_device_name()).view(-1).unsqueeze(0)) + if "mlp.dense_4h_to_h.weight" in key: + self.mlp4hh_scales.append(q_scale) + elif "mlp.dense_h_to_4h.weight" in key: + self.mlph4h_scales.append(q_scale) + elif "attention.query_key_value.weight" in key: + self.qkv_scales.append(q_scale) + else: + self.dense_scales.append(q_scale) + return value_list + + def merge_layer_scales(self, layer_scales): + max_dim = max([s.shape[-1] for s in layer_scales]) + layer_scales = [ + torch.cat((s, torch.zeros((1, max_dim - s.shape[-1]), device=get_accelerator().current_device_name())), + dim=-1) if s.shape[-1] < max_dim else s for s in layer_scales + ] + return torch.cat(layer_scales).unsqueeze(0) + + def merge_scales(self): + all_scales = [] + for dense_scale, qkv_scale, m4hh_scale, mh4h_scale in \ + zip(self.dense_scales, self.qkv_scales, self.mlp4hh_scales, self.mlph4h_scales): + all_scales.append(self.merge_layer_scales([qkv_scale, dense_scale, mh4h_scale, m4hh_scale])) + return torch.cat(all_scales) + + def merge_scales_split(self, split_count): + all_scales = [[] for _ in range(split_count)] + for dense_scale, qkv_scale, m4hh_scale, mh4h_scale in \ + zip(self.dense_scales, self.qkv_scales, self.mlp4hh_scales, self.mlph4h_scales): + dense_scale = torch.split(dense_scale, dense_scale.numel() // split_count) + qkv_scale = torch.split(qkv_scale, qkv_scale.numel() // split_count) + m4hh_scale = torch.split(m4hh_scale, m4hh_scale.numel() // split_count) + mh4h_scale = torch.split(mh4h_scale, mh4h_scale.numel() // split_count) + for s in range(split_count): + all_scales[s].append( + torch.cat([ + torch.cat((qkv_scale[s], torch.zeros_like(qkv_scale[s])), dim=1), + torch.cat((dense_scale[s], torch.zeros_like(dense_scale[s])), dim=1), mh4h_scale[s], + m4hh_scale[s] + ]).unsqueeze(0)) + for scales_a in all_scales: + torch.cat(scales_a) + return all_scales + + def sd_quantize_megatron(self, sd, quantize_bits, groups): + keys = sd.keys() + for key in keys: + value_list = [sd[key]] + if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key or \ + "mlp.dense_h_to_4h.weight" in key or "attention.query_key_value.weight" in key: + value_list = self.Quantize(value_list, quantize_bits, groups, key=key) + sd[key] = value_list[0] + + all_scales = self.merge_scales() + return sd, all_scales + + def model_quantize(self, model, quantize_policy, quantize_bits, groups): + all_scales = [] + + def quantize_fn(layer, policy_cls): + policy = policy_cls(layer) + + _, qkvw, _, dense_w, _, _ = policy.attention() + _, _h4h_w, _, _4hh_w, _ = policy.mlp() + keys = [qkvw, dense_w, _h4h_w, _4hh_w] + layer_scales = [] + + for key in range(len(keys)): + if self.mlp_extra_grouping and self.is_mlp(keys[key]): + data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups * 2) + elif policy_cls is HFBertLayerPolicy and self.is_qkv(keys[key]): + data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups * 3) + else: + data_quantized, data_scale = self.quantize_data(keys[key], quantize_bits, groups) + keys[key].copy_(data_quantized) + layer_scales.append((1 / data_scale.to(get_accelerator().current_device_name()).view(-1).unsqueeze(0))) + all_scales.append(self.merge_layer_scales(layer_scales)) + return layer + + def _quantize_module(model, policies): + for name, child in model.named_children(): + if child.__class__ in policies: + quantize_fn, replace_policy = policies[child.__class__] + setattr(model, name, quantize_fn(child, replace_policy)) + else: + _quantize_module(child, policies) + + return model + + policy = {} + if quantize_policy is not None: + for layer_name, replace_policy in quantize_policy.items(): + policy.update({layer_name: (quantize_fn, replace_policy)}) + else: + for plcy in replace_policies: + policy.update({plcy._orig_layer_class: (quantize_fn, plcy)}) + + quantized_module = _quantize_module(model, policy) + + return quantized_module, torch.cat(all_scales) diff --git a/lib/python3.12/site-packages/deepspeed/utils/__init__.py b/lib/python3.12/site-packages/deepspeed/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5d902a883f6eeee72b3c9ccbf5a11abf22de8e33 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .logging import logger, log_dist +from .comms_logging import get_caller_func +#from .distributed import init_distributed +from .init_on_device import OnDevice +from .groups import * +from .nvtx import instrument_w_nvtx +# TODO: Move tensor fragment and mixed precision to zero utils +from .tensor_fragment import tensor_fragment, get_full_hp_param, get_hp_fragment_mapping, fragment_address, get_full_hp_grad, map_to_flat_opt_states +from .tensor_fragment import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state +from .tensor_fragment import set_full_hp_param, set_full_hp_grad +from .tensor_fragment import safe_set_full_fp32_param, safe_set_full_optimizer_state, safe_set_full_grad +from .tensor_fragment import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state +from .tensor_fragment import safe_set_local_fp32_param, safe_set_local_grad, safe_set_local_optimizer_state +from .tensor_fragment import safe_update_full_grad_vectorized +from .z3_leaf_module import set_z3_leaf_modules, unset_z3_leaf_modules, get_z3_leaf_modules, z3_leaf_module, z3_leaf_parameter, set_z3_leaf_module +from .mixed_precision_linkage import link_hp_params, lazy_init_hp_params_optimizer_state +from deepspeed.runtime.dataloader import RepeatingLoader +from .numa import get_numactl_cmd diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..481e156806ea703faf1d607e5ede4450d33e3986 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/bwc.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/bwc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f7ba06606df5ed4949b0a2aca81ed866fc11468e Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/bwc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59f46648d9b058a081ef44eb3418caf2b1d11892 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca408077c42a3e27f00afe22daeda2172a9938e1 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/exceptions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/groups.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/groups.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ead669985d0318423613650e2c40226e8b950b4 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/groups.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/logging.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f24c791e81da1258faae719f24f553e6fb2e085f Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00e22065d57c7f9a3a0c88d329b269436b92adee Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/mixed_precision_linkage.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/numa.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/numa.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78c8cce78c6a9a37414e1cf24ca88d0db5b484dd Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/numa.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86db15b998d6e86283b01a1b9e852fcde681e148 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/tensor_fragment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/torch.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/torch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2afc6448336fa7049e1ea4b4dca0a07cc26a565 Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/torch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e7339e8622042918134673b95330d546f19a61d Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/z3_leaf_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-312.pyc b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40748507c279d370c2cbdf7d678291361264a81b Binary files /dev/null and b/lib/python3.12/site-packages/deepspeed/utils/__pycache__/zero_to_fp32.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/deepspeed/utils/bwc.py b/lib/python3.12/site-packages/deepspeed/utils/bwc.py new file mode 100644 index 0000000000000000000000000000000000000000..69fcc251a68429a768892b7cf7c6824de7cb5ba2 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/bwc.py @@ -0,0 +1,104 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + + +def bwc_tensor_model_parallel_rank(mpu=None): + """Backwards-compatible way of querying the tensor model parallel rank from + an ``mpu`` object. + + *Tensor* model parallelism means that tensors are physically split across + processes. This contrasts with *pipeline* model parallelism, in which the + layers are partitioned but tensors left intact. + + The API for tensor model parallelism has changed across versions and this + helper provides a best-effort implementation across versions of ``mpu`` + objects. The preferred mechanism is + ``mpu.get_tensor_model_parallel_rank()``. + + This should "just work" with both Megatron-LM and DeepSpeed's pipeline + parallelism. + + Args: + mpu (model parallel unit, optional): The tensor model parallel rank. + If ``mpu=None``, returns 0. Defaults to ``None``. + + Returns: + int: the rank + """ + if mpu is None: + # No model parallelism in easy :) + return 0 + + if hasattr(mpu, 'get_tensor_model_parallel_rank'): + # New Megatron and DeepSpeed convention (post pipeline-parallelism release) + return mpu.get_tensor_model_parallel_rank() + elif hasattr(mpu, 'get_slice_parallel_rank'): + # Some DeepSpeed + pipeline parallelism versions + return mpu.get_slice_parallel_rank() + else: + # Deprecated Megatron and DeepSpeed convention + return mpu.get_model_parallel_rank() + + +def bwc_tensor_model_parallel_world_size(mpu=None): + """Backwards-compatible way of querying the tensor model parallel world size. + Similar to bwc_tensor_model_parallel_rank. + """ + if mpu is None: + return 1 + + if hasattr(mpu, 'get_tensor_model_parallel_world_size'): + # New Megatron and DeepSpeed convention (post pipeline-parallelism release) + return mpu.get_tensor_model_parallel_world_size() + elif hasattr(mpu, 'get_slice_parallel_world_size'): + # Some DeepSpeed + pipeline parallelism versions + return mpu.get_slice_parallel_world_size() + else: + # Deprecated Megatron and DeepSpeed convention + return mpu.get_model_parallel_world_size() + + +def bwc_tensor_model_parallel_group(mpu=None): + """Backwards-compatible way of querying the tensor model parallel group. + Similar to bwc_tensor_model_parallel_rank. + """ + if mpu is None: + return None + + if hasattr(mpu, 'get_tensor_model_parallel_group'): + # New Megatron and DeepSpeed convention (post pipeline-parallelism release) + return mpu.get_tensor_model_parallel_group() + elif hasattr(mpu, 'get_slice_parallel_group'): + # Some DeepSpeed + pipeline parallelism versions + return mpu.get_slice_parallel_group() + else: + # Deprecated Megatron and DeepSpeed convention + return mpu.get_model_parallel_group() + + +def bwc_pipeline_parallel_world_size(mpu=None): + """Backwards-compatible way of querying the pipeline parallel world size.""" + world_size = 1 + if mpu is not None: + if hasattr(mpu, 'get_pipeline_model_parallel_world_size'): + # New Megatron and DeepSpeed convention (post pipeline-parallelism release) + world_size = mpu.get_pipeline_model_parallel_world_size() + elif hasattr(mpu, 'get_pipe_parallel_world_size'): + # DeepSpeed Topology + world_size = mpu.get_pipe_parallel_world_size() + return world_size + + +def bwc_pipeline_parallel_group(mpu=None): + """Backwards-compatible way of querying the pipeline parallel group.""" + if mpu is None: + return None + if hasattr(mpu, 'get_pipeline_model_parallel_group'): + # Megatron + return mpu.get_pipeline_model_parallel_group() + elif hasattr(mpu, 'get_pipe_parallel_group'): + # DeepSpeed Topology + return mpu.get_pipe_parallel_group() + assert False, 'mpu does not support pipeline parallel group' diff --git a/lib/python3.12/site-packages/deepspeed/utils/comms_logging.py b/lib/python3.12/site-packages/deepspeed/utils/comms_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..0e405a06d23e8925b942e26273a62fc1a6b69743 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/comms_logging.py @@ -0,0 +1,178 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import math +from deepspeed.utils import log_dist + + +def get_caller_func(frame=3): + import sys + return sys._getframe(frame).f_code.co_name + + +def print_rank_0(message): + import deepspeed.comm as dist + if dist.get_rank() == 0: + print(message) + + +# Helper function to pretty-print message sizes +def convert_size(size_bytes): + if size_bytes == 0: + return "0B" + size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB") + i = int(math.floor(math.log(size_bytes, 1024))) + p = math.pow(1024, i) + s = round(size_bytes / p, 2) + return "%s %s" % (s, size_name[i]) + + +# Helper function to calculate algbw and busbw. +# See https://gist.github.com/jeffra/b5e80466b4c86be00ea3b6f130fb7a36 and https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md +def calc_bw_log(comm_op, size, duration): + import deepspeed.comm as dist + + n = dist.get_world_size() + tput = 0 + busbw = 0 + if comm_op == "all_to_all_single": + tput = (size / duration) + busbw = (size / duration) * ((n - 1) / n) + elif comm_op == "all_gather" or comm_op == "all_gather_into_tensor" or comm_op == "reduce_scatter" or comm_op == "reduce_scatter_tensor": + size *= n + tput = (size / duration) + busbw = (size / duration) * ((n - 1) / n) + elif comm_op == "all_reduce" or comm_op == "all_reduce_coalesced" or comm_op == "inference_all_reduce": + tput = (size * 2 / duration) + busbw = (size / duration) * (2 * (n - 1) / n) + elif comm_op == "send" or comm_op == "recv" or comm_op == "isend" or comm_op == "irecv" or comm_op == "broadcast" or comm_op == "reduce" or comm_op == "gather" or comm_op == "scatter" or comm_op == "barrier": + tput = (size / duration) + busbw = tput + else: + print_rank_0("wrong comm_op specified") # noqa: F821 + exit(0) + + # convert to Gbps + tput *= 8 + busbw *= 8 + + tput /= 1e6 + busbw /= 1e6 + + return tput, busbw + + +class CommsLogger: + + def __init__(self): + from deepspeed.comm.constants import COMMS_LOGGER_VERBOSE_DEFAULT, COMMS_LOGGER_DEBUG_DEFAULT, COMMS_LOGGER_PROF_OPS_DEFAULT, COMMS_LOGGER_PROF_ALL_DEFAULT, COMMS_LOGGER_ENABLED_DEFAULT + self.comms_dict = {} + self.verbose = COMMS_LOGGER_VERBOSE_DEFAULT + self.debug = COMMS_LOGGER_DEBUG_DEFAULT + self.prof_ops = COMMS_LOGGER_PROF_OPS_DEFAULT + self.prof_all = COMMS_LOGGER_PROF_ALL_DEFAULT + self.enabled = COMMS_LOGGER_ENABLED_DEFAULT + + def configure(self, comms_config): + self.enabled = comms_config.comms_logger_enabled + if self.enabled: + self.verbose = comms_config.comms_logger.verbose + self.debug = comms_config.comms_logger.debug + self.prof_ops = comms_config.comms_logger.prof_ops + self.prof_all = comms_config.comms_logger.prof_all + + # There are three settings for the op profiler: + # - Global profiling (profile all comms) + # - Op-type profiling (e.g. profile all all_reduce comms) + # - Op profiling (e.g. profile a specific all_reduce op) + def start_profiling_comms(self): + self.prof_all = True + + def stop_profiling_comms(self): + self.prof_all = True + + # E.g. start_profiling_op('all_reduce') + def start_profiling_op(self, op_name_list): + self.prof_ops = list(set(self.prof_ops) | set(op_name_list)) + + def stop_profiling_op(self, op_name_list): + self.prof_ops = [op for op in self.prof_ops if op not in op_name_list] + + # Add log entry + def append(self, raw_name, record_name, latency, msg_size): + algbw, busbw = calc_bw_log(raw_name, msg_size, latency) + if record_name in self.comms_dict.keys(): + # If this comm_op has already been logged with this message size, just add to existing record + if msg_size in self.comms_dict[record_name].keys(): + self.comms_dict[record_name][msg_size][0] += 1 + self.comms_dict[record_name][msg_size][1].append(latency) + self.comms_dict[record_name][msg_size][2].append(algbw) + self.comms_dict[record_name][msg_size][3].append(busbw) + # If this is a new message size for this comm_op, add new record under existing comm_op + else: + self.comms_dict[record_name][msg_size] = [1, [latency], [algbw], [busbw]] + else: + # Create entirely new record + self.comms_dict[record_name] = {msg_size: [1, [latency], [algbw], [busbw]]} + # If verbose, print every comm op + # TODO: Add to tensorboard + if self.verbose: + log_str = f"comm op: {record_name} | time (ms): {latency:.2f} | msg size: {convert_size(msg_size)} | algbw (Gbps): {algbw:.2f} | busbw (Gbps): {busbw:.2f}" + log_dist(log_str, [0]) + + # Print summary at end of iteration, epoch, or training + def log_all(self, print_log=True, show_straggler=False): + import torch + from deepspeed.utils.timer import trim_mean + import deepspeed.comm as dist + from deepspeed.comm.reduce_op import ReduceOp + if print_log: + print( + f"{'Comm. Op': <20}{'Message Size': <20}{'Count': <20}{'Total Latency(ms)': <20}{'Avg Latency(ms)': <20}{'tput_avg (Gbps)': <20}{'busbw_avg (Gbps)': <20}" + ) + for record_name in self.comms_dict.keys(): + if print_log: + print(record_name) + for msg_size, vals in sorted(self.comms_dict[record_name].items()): + # vals[0] is the count for each msg size + count = vals[0] + # vals[1] is a list of latency records for each msg size + total_lat = sum(vals[1]) + # vals[2] and vals[3] are the lists of algbw and busbw, respectively + # Get rid of outliers when we print + avg_lat = trim_mean(vals[1], 0.1) + avg_algbw = trim_mean(vals[2], 0.1) + avg_busbw = trim_mean(vals[3], 0.1) + if print_log: + print( + f"{' ': <20}{convert_size(msg_size): <20}{count: <20}{total_lat: <20.2f}{avg_lat: <20.2f}{avg_algbw: <20.2f}{avg_busbw: <20.2f}" + ) + + if show_straggler: + if print_log: + print("_______________________________") + print("Breakdown with straggler effect") + print("-------------------------------") + print( + f"{'Comm. Op': <20}{'Message Size': <20}{'Count': <20}{'Total comm lat(ms)': <20}{'Total straggler(ms)': <20}{'Avg comm lat(ms)': <20}{'Avg straggler(ms)': <20}" + ) + for record_name in self.comms_dict.keys(): + if print_log: + print(record_name) + for msg_size, vals in sorted(self.comms_dict[record_name].items()): + # vals[0] is the count for each msg size + count = vals[0] + # vals[1] is a list of latency records for each msg size + lats = torch.tensor(vals[1]) + min_lats = torch.tensor(vals[1]) + dist.all_reduce(min_lats, op=ReduceOp.MIN) + total_lat = min_lats.sum().item() + total_straggler = (lats - min_lats).sum().item() + avg_lat = trim_mean(min_lats.tolist(), 0.1) + avg_straggler = trim_mean((lats - min_lats).tolist(), 0.1) + if print_log: + print( + f"{' ': <20}{convert_size(msg_size): <20}{count: <20}{total_lat: <20.2f}{total_straggler: <20.2f}{avg_lat: <20.2f}{avg_straggler: <20.2f}" + ) diff --git a/lib/python3.12/site-packages/deepspeed/utils/config.py b/lib/python3.12/site-packages/deepspeed/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..15f37ca7d8745cccee562d08b0eae3f10c5d9877 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/config.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.runtime.config_utils import DeepSpeedConfigModel + +######################################### +# Timers +######################################### +# Timers. By default, timers are enabled. +# Users can configure in ds_config.json as below example: +TIMERS_FORMAT = ''' +Timers should be enabled as: +"timers": { + "throughput": { + "enabled": true, + "synchronized": true + } +} +''' + +TIMERS = "timers" +TIMERS_THROUGHPUT = "throughput" + + +def get_timers_config(param_dict): + if param_dict and TIMERS in param_dict and TIMERS_THROUGHPUT in param_dict[TIMERS]: + timers_config_dict = param_dict[TIMERS][TIMERS_THROUGHPUT] + else: + timers_config_dict = {} + return DeepSpeedThroughputTimerConfig(**timers_config_dict) + + +class DeepSpeedThroughputTimerConfig(DeepSpeedConfigModel): + """ Configure throughput timers """ + + enabled: bool = True + """ Turn on/off throughput timers """ + + synchronized: bool = True + """ Whether to synchronize a device when measuring the time. + Synchronizing a device is required to produce the most accurate timer measurements. + However, this comes at the expense of performance degradation. The CPU timer provides + sufficient accuracy in many cases. + """ diff --git a/lib/python3.12/site-packages/deepspeed/utils/debug.py b/lib/python3.12/site-packages/deepspeed/utils/debug.py new file mode 100644 index 0000000000000000000000000000000000000000..6174ba45137e1e70dd004593a69ff41015c0212c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/debug.py @@ -0,0 +1,173 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import deepspeed.comm as dist + +# For lazy import with printflock() +fcntl = None + +# for debug purposes map module and param objects to their fully qualified names +module_names = {} +param_names = {} + + +def debug_clear_module_and_param_names(): + global module_names + global param_names + module_names = {} + param_names = {} + + +def debug_extract_module_and_param_names(model): + # extract the fully qualified names as soon as the model is acquired + global module_names + global param_names + # XXX: can probably make a map of param2module and vice-versa + module_names = {module: name for name, module in model.named_modules()} + param_names = {param: name for name, param in model.named_parameters()} + + +def debug_module2name(module): + if module in module_names: + return module_names[module] + else: + return "unknown" + + +def debug_module2name_id(module): + return f"name={debug_module2name(module)}" + + +def debug_module2name_class(module): + return f"name={debug_module2name(module)} {module.__class__.__name__}" + + +def debug_param2name(param): + if param in param_names: + return param_names[param] + else: + return "unknown" + + +def debug_param2name_id(param): + return f"name={debug_param2name(param)} id={param.ds_id}" + + +def debug_param2name_id_shape(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape}" + + +def debug_param2name_id_shape_device(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape} device={param.device}" + + +def debug_param2name_id_numel(param): + return f"name={debug_param2name(param)} id={param.ds_id} numel={param.numel()}" + + +def debug_param2name_id_shape_status(param): + return f"name={debug_param2name(param)} id={param.ds_id} shape={param.ds_shape} status={param.ds_status}" + + +def printflock(*msgs): + """ + + For printing messages for all concurrent gpus w/o getting interleaved text. + + This is useful when debugging issues where multi-gpus don't sync. + + 1. Enable the force debug in say partitioning and zero3 files + 2. Override the usual versions with :: + + def print_rank_0(message, debug=False, force=False): + rank = deepspeed.comm.get_rank() + printflock(f"[{rank}] {message}") + 3. run the program and you get both logs non-interleaved + + But this makes it very difficult to make sense of the output, so the ``log_rank_file`` helper + function might be more useful, as it's easier to send each log stream into a separate file and + then compare those. + + """ + global fcntl + if fcntl is None: + import fcntl + + with open(__file__, "r") as fh: + fcntl.flock(fh, fcntl.LOCK_EX) + try: + print(*msgs) + finally: + fcntl.flock(fh, fcntl.LOCK_UN) + + +fh = None + + +def log_rank_file(rank, *msgs): + """ + Print to a log file of the given rank + + This is useful for debugging hanging in sync processes. Here is a possible workflow: + + 1. Enable the force debug in say partitioning and zero3 files + 2. Override the usual versions of print_rank_0 in those files with :: + + def print_rank_0(message, debug=False, force=False): + rank = deepspeed.comm.get_rank() + log_rank_file(rank, message) + + 3. run the program + 4. fix up the expected differences, e.g. different cuda numbers :: + + perl -pi -e 's|cuda:1|cuda:0|' log_rank_* + + 5. now diff and see where names and ids diverge - you will find where the gpus don't do the same + work (e.g. when some layers get conditionally skipped on one gpu but not all) + + diff -u log_rank_0.txt log_rank_1.txt | less + + """ + global fh + if fh is None: + fh = open(f"log_rank_{rank}.txt", "w") + for m in msgs: + fh.write(f"{m}\n") + fh.flush() + + +def print_backward_tensors(tensor): + + def _print_bwd_tensors(grad_fn): + print(f"Backward tensors in {grad_fn}") + for funcs in grad_fn.next_functions: + if funcs[0]: + try: + tensor = getattr(funcs[0], 'variable') + print(funcs[0]) + print(f"Tensor - id: {id(tensor)}, shape: {tensor.shape}, data: {tensor}, grad: {tensor.grad}") + except AttributeError as e: + _print_bwd_tensors(funcs[0]) + + if hasattr(tensor, 'grad_fn'): + _print_bwd_tensors(tensor.grad_fn) + + +def print_rank(*msg, force=False): + """print something on all global ranks with [rank] prefix. + """ + if not force: + return + global_rank = dist.get_rank() + print(f"[{global_rank}]", *msg) + + +def print_rank0(*msg, force=False): + """print something only on rank 0""" + if not force: + return + global_rank = dist.get_rank() + if global_rank == 0: + print(f"[{global_rank}]", *msg) diff --git a/lib/python3.12/site-packages/deepspeed/utils/exceptions.py b/lib/python3.12/site-packages/deepspeed/utils/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..9dba0933cc9e236ec150bfd5439f7ce8e157394a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/exceptions.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + + +class DeprecatedException(Exception): + pass diff --git a/lib/python3.12/site-packages/deepspeed/utils/groups.py b/lib/python3.12/site-packages/deepspeed/utils/groups.py new file mode 100644 index 0000000000000000000000000000000000000000..3a9108864be2149672f8f922860b1ae0f8fecdfd --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/groups.py @@ -0,0 +1,713 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# The file has been adapted from https://github.com/NVIDIA/Megatron-LM and retains the following license from the original file + +# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" + Support different forms of parallelism in DeepSpeed using multiple process groups. + Given that there are multiple scenarios and use-cases, this file is going to be updated + frequently. For now, the group creation needed for the training scenario is being implemented. + For inference and other new scenarios, the code will be either reused or added to this file. +""" + +from deepspeed import comm as dist +from deepspeed.utils import log_dist +from deepspeed.utils.bwc import bwc_tensor_model_parallel_world_size, bwc_pipeline_parallel_world_size +from deepspeed.utils.exceptions import DeprecatedException +from deepspeed.accelerator import get_accelerator + +# Expert parallel group that the current rank belongs to. +_EXPERT_PARALLEL_GROUP = {} +# Expert data parallel group that the current rank belongs to. +_EXPERT_DATA_PARALLEL_GROUP = {} +# dist world group needs to be cloned for some cases +_WORLD_GROUP = None +# ZeRO parameter partitioning group that the current rank belongs to. +_ZERO_PARAM_INTRA_PARALLEL_GROUP = None +# global object to maintain mpu object if passed by a Megatron client +mpu = None +# global object that stores tensor parallel world size for experts +expert_tensor_parallel_world_size = 1 +# All to All quantized graident communication groups +_ALL_TO_ALL_GROUP = {} + +mesh_device = None + + +# Deprecated groups initialize function. +def initialize(ep_size=1, mpu=None): + """ Deprecated function. Retained to inform the users.""" + raise DeprecatedException( + "Please do not use the groups.initialize() API as it is deprecated. Instead, pass the desired ep_size to deepspeed.moe.layer.MoE(..,ep_size,..)" + ) + + +def _ensure_divisibility(numerator, denominator): + """Ensure that numerator is divisible by the denominator.""" + assert numerator % denominator == 0, '{} is not divisible by {}'.format(numerator, denominator) + + +# ======== Start: Tensor Parallel Group Attributes ======== + +# Intra-layer model parallel group that the current rank belongs to. +_TENSOR_MODEL_PARALLEL_GROUP = None + +# Model parallel group (both intra- and pipeline) that the current rank belongs to. +_MODEL_PARALLEL_GROUP = None +# Data parallel group that the current rank belongs to. +_DATA_PARALLEL_GROUP = None + +# These values enable us to change the mpu sizes on the fly. +_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None +_MPU_TENSOR_MODEL_PARALLEL_RANK = None + + +def _init_tp_mesh_device(tensor_model_parallel_size=1, data_parallel_size=None): + """Initialize model data parallel groups.""" + + global _DATA_PARALLEL_GROUP + global _MODEL_PARALLEL_GROUP + global _TENSOR_MODEL_PARALLEL_GROUP + + if _TENSOR_MODEL_PARALLEL_GROUP is not None: + return + + if data_parallel_size is None: + data_parallel_size = dist.get_world_size() // tensor_model_parallel_size + + mesh_device = dist.initialize_mesh_device((data_parallel_size, tensor_model_parallel_size), + ("data_parallel", "tensor_parallel")) + _TENSOR_MODEL_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="tensor_parallel") + _DATA_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="data_parallel") + + # They are always equal only in 2D (DP + TP) parallelism. + # _MODEL_PARALLEL_GROUP is assigned the same value as _TENSOR_MODEL_PARALLEL_GROUP + # to allow for future potential changes. + _MODEL_PARALLEL_GROUP = _TENSOR_MODEL_PARALLEL_GROUP + + return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP + + +def get_tensor_model_parallel_group(): + """Get the tensor model parallel group the caller rank belongs to.""" + + assert _TENSOR_MODEL_PARALLEL_GROUP is not None, \ + 'intra_layer_model parallel group is not initialized' + return _TENSOR_MODEL_PARALLEL_GROUP + + +def get_model_parallel_group(): + """Get the model parallel group the caller rank belongs to.""" + + assert _MODEL_PARALLEL_GROUP is not None, \ + 'model parallel group is not initialized' + return _MODEL_PARALLEL_GROUP + + +def get_data_parallel_group(): + """Get the data parallel group the caller rank belongs to.""" + assert _DATA_PARALLEL_GROUP is not None, \ + 'data parallel group is not initialized' + return _DATA_PARALLEL_GROUP + + +def set_tensor_model_parallel_world_size(world_size): + """Set the tensor model parallel size""" + global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE + _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size + + +def get_tensor_model_parallel_world_size(): + """Return world size for the tensor model parallel group.""" + global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE + if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None: + return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE + return dist.get_world_size(group=get_tensor_model_parallel_group()) + + +def get_model_parallel_world_size(): + return get_tensor_model_parallel_world_size() + + +def set_tensor_model_parallel_rank(rank): + """Set tensor model parallel rank.""" + global _MPU_TENSOR_MODEL_PARALLEL_RANK + _MPU_TENSOR_MODEL_PARALLEL_RANK = rank + + +def get_tensor_model_parallel_rank(): + """Return my rank for the tensor model parallel group.""" + global _MPU_TENSOR_MODEL_PARALLEL_RANK + if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None: + return _MPU_TENSOR_MODEL_PARALLEL_RANK + return dist.get_rank(group=get_tensor_model_parallel_group()) + + +def get_model_parallel_rank(): + return get_tensor_model_parallel_rank() + + +def get_tensor_model_parallel_src_rank(): + """Calculate the global rank corresponding to the first local rank + in the tensor model parallel group.""" + global_rank = dist.get_rank() + local_world_size = get_tensor_model_parallel_world_size() + return (global_rank // local_world_size) * local_world_size + + +def get_data_parallel_world_size(): + """Return world size for the data parallel group.""" + return dist.get_world_size(group=get_data_parallel_group()) + + +def get_data_parallel_rank(): + """Return my rank for the data parallel group.""" + return dist.get_rank(group=get_data_parallel_group()) + + +# ======== End: Tensor Parallel Group Attributes ======== + + +# Not currently used. Helper function to create a model (tensor) parallel group. +def _create_model_parallel(model_parallel_size_): + """ + Initialize model data parallel groups. + + Arguments: + model_parallel_size: number of GPUs used to parallelize model. + + Returns: + Tuple of data parallel group and model parallel group + + Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we + use 2 GPUs to parallelize the model. The present function will + create 4 model parallel groups and 2 data parallel groups as: + 4 model parallel groups: + [g0, g1], [g2, g3], [g4, g5], [g6, g7] + 2 data parallel groups: + [g0, g2, g4, g6], [g1, g3, g5, g7] + Note that for efficiency, the caller should make sure adjacent ranks + are on the same DGX box. For example if we are using 2 DGX-1 boxes + with a total of 16 GPUs, rank 0 to 7 belong to the first box and + ranks 8 to 15 belong to the second box. + """ + log_dist(f'Creating model parallel group with size {model_parallel_size_}', ranks=[0]) + # Get world size and rank. Ensure some consistencies. + assert dist.is_initialized() + world_size = dist.get_world_size() + model_parallel_size = min(model_parallel_size_, world_size) + _ensure_divisibility(world_size, model_parallel_size) + rank = dist.get_rank() + + _DATA_PARALLEL_GROUP = None + _MODEL_PARALLEL_GROUP = None + # Build the data parallel groups. + for i in range(model_parallel_size): + ranks = range(i, world_size, model_parallel_size) + group = dist.new_group(ranks) + if i == (rank % model_parallel_size): + _DATA_PARALLEL_GROUP = group + + # Build the model parallel groups. + for i in range(world_size // model_parallel_size): + ranks = range(i * model_parallel_size, (i + 1) * model_parallel_size) + group = dist.new_group(ranks) + if i == (rank // model_parallel_size): + _MODEL_PARALLEL_GROUP = group + + return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP + + +def _create_expert_and_data_parallel(expert_parallel_size_, use_data_before_expert_parallel_=False): + """ + Create expert and data parallel groups. + + Note: Caller of this function is responsible to check if the groups already exist. + + Example - E + D parallel + world_size = 16 + expert_parallel_size = 2 # number of experts in same group + expert_data_parallel_group = [0,2,4,6,8,10,12,14], [1,3,5,7,9,11,13,15] - all reduce is only on MoE params + expert_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - no all reduce, but all to all + data_parallel_group = [0,1,...,15] - all reduce is only on non-MoE + use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology + """ + assert dist.is_initialized() + + log_dist(f'Creating expert and data parallel groups with size {expert_parallel_size_}', ranks=[0]) + world_size = dist.get_world_size() + pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu) + rank = dist.get_rank() + + pp_stride = world_size // pp_world_size + _ensure_divisibility(pp_stride, expert_parallel_size_) + + group_name = f"ep_size_{expert_parallel_size_}" + + # Build the expert data parallel groups. + global _EXPERT_DATA_PARALLEL_GROUP + + ep_stride = pp_stride // expert_parallel_size_ + + # Only create group if it does not already exist + if group_name not in _EXPERT_DATA_PARALLEL_GROUP: + for pp_stage_start in range(0, world_size, pp_stride): + for i in range(expert_parallel_size_): + if use_data_before_expert_parallel_: + ranks = range(pp_stage_start + i * ep_stride, pp_stage_start + (i + 1) * ep_stride) + else: + ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, expert_parallel_size_) + group = dist.new_group(ranks) + log_dist( + f'Creating expert data parallel process group named {group_name} ' + f'with ranks: {list(ranks)}', [0]) + if rank in ranks: + _EXPERT_DATA_PARALLEL_GROUP[group_name] = group + + # Build the expert parallel groups. + global _EXPERT_PARALLEL_GROUP + + # Only create group if it does not already exist + if group_name not in _EXPERT_PARALLEL_GROUP: + if use_data_before_expert_parallel_: + for pp_stage_start in range(0, world_size, pp_stride): + for i in range(ep_stride): + ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, ep_stride) + group = dist.new_group(ranks) + log_dist( + f'creating expert parallel process group named {group_name} ' + f'with ranks: {list(ranks)}', [0]) + if rank in ranks: + _EXPERT_PARALLEL_GROUP[group_name] = group + else: + for i in range(world_size // expert_parallel_size_): + ranks = range(i * expert_parallel_size_, (i + 1) * expert_parallel_size_) + group = dist.new_group(ranks) + log_dist(f'creating expert parallel process group named {group_name} ' + f'with ranks: {list(ranks)}', [0]) + if rank in ranks: + _EXPERT_PARALLEL_GROUP[group_name] = group + + +def _get_expert_parallel_ranks(world_size, + tensor_parallel_size_, + expert_parallel_size_, + pipeline_parallel_size_=1, + use_data_before_expert_parallel_=False): + """Generate expert parallel and expert data parallel group ranks list. + + Example - E + M + D parallel + world_size = 16 + model_degree = 2 + expert_degree = 4 # number of experts in same group + mp_group = [0, 1], [2,3], [4,5] ... + data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15] + expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15] + expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15] + + Args: + world_size (int): Distributed world size. + tensor_parallel_size_ (int): Tensor parallel group size. + expert_parallel_size_ (int): Expert parallel group size. + pipeline_parallel_size_ (int): Pipeline parallel group size + use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology + Returns: + Expert parallel group ranks and Expert data parallel group ranks list. + """ + _ensure_divisibility(world_size, tensor_parallel_size_ * pipeline_parallel_size_) + dp_world_size = world_size // (tensor_parallel_size_ * pipeline_parallel_size_) + _ensure_divisibility(dp_world_size, expert_parallel_size_) + + # Generate data parallel groups + data_parallel_groups = [] + dp_group_size = tensor_parallel_size_ + pp_stride = world_size // pipeline_parallel_size_ + + if use_data_before_expert_parallel_: + dp_stride = world_size // expert_parallel_size_ // tensor_parallel_size_ // pipeline_parallel_size_ + for pp_stage_start in range(0, world_size, pp_stride): + pp_stage_next = pp_stage_start + pp_stride + for i in range(dp_group_size): + data_parallel_groups.append(list()) + for ds in range(dp_stride): + # [0, 4, 8, 12, 16, 20, 24, 28, 2, 6, 10, 14, 18, 22, 26, 30] + # [1, 5, 9, 13, 17, 21, 25, 29, 3, 7, 11, 15, 19, 23, 27, 31] + data_parallel_groups[-1].extend( + list( + range(pp_stage_start + i + ds * tensor_parallel_size_, pp_stage_next, + dp_stride * tensor_parallel_size_))) + else: + for pp_stage_start in range(0, world_size, pp_stride): + pp_stage_next = pp_stage_start + pp_stride + for i in range(dp_group_size): + data_parallel_groups.append(list(range(pp_stage_start + i, pp_stage_next, dp_group_size))) + + expert_parallel_groups = [] + expert_data_parallel_groups = [] + for dp_ranks in data_parallel_groups: + # partition of expert parallel groups, e.g. [0,2,4,6], [8,10,12,14] + part_ep_groups = [] + for i in range(0, dp_world_size, expert_parallel_size_): + part_ep_groups.append(dp_ranks[i:i + expert_parallel_size_]) + expert_parallel_groups.extend(part_ep_groups) + + # zip part_ep_groups get expert data parallel ranks, e.g [0,8],[2,10],[4,12],[6,14] + for expert_dp_ranks in zip(*part_ep_groups): + expert_data_parallel_groups.append(list(expert_dp_ranks)) + + return expert_parallel_groups, expert_data_parallel_groups + + +def _create_expert_data_and_model_parallel(expert_parallel_size_, mpu, use_data_before_expert_parallel_=False): + """ + Create expert and data parallel groups based on MPU (model parallel) group. + + Note: Caller of this function is responsible to check if the groups already exist. + + Example - E + M + D parallel + world_size = 16 + model_degree = 2 + expert_degree = 4 # number of experts in same group + mp_group = [0, 1], [2,3], [4,5] ... + data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15] + expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15] + expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15] + """ + assert dist.is_initialized(), "dist is not initialized" + tensor_parallel_size_ = bwc_tensor_model_parallel_world_size(mpu) + + global expert_tensor_parallel_world_size + expert_tensor_parallel_world_size = tensor_parallel_size_ + + world_size = dist.get_world_size() + rank = dist.get_rank() + dp_world_size = mpu.get_data_parallel_world_size() + pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu) + + _ensure_divisibility(world_size, tensor_parallel_size_) + _ensure_divisibility(dp_world_size, expert_parallel_size_) + + log_dist( + f"Creating deepspeed groups with model parallel size {tensor_parallel_size_}, " + f"pipeline parallel size {pp_world_size}, expert parallel size {expert_parallel_size_}, " + f"world size {world_size}, dp world size {dp_world_size}", [0]) + + global _EXPERT_PARALLEL_GROUP, _EXPERT_DATA_PARALLEL_GROUP + + group_name = f"ep_size_{expert_parallel_size_}" + + # Only create groups if they don't already exist + # Need to check conditions outside the group creation loop because of the way torch.dist group creation works + if group_name not in _EXPERT_DATA_PARALLEL_GROUP and group_name not in _EXPERT_PARALLEL_GROUP: + expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks( + world_size, tensor_parallel_size_, expert_parallel_size_, pp_world_size, use_data_before_expert_parallel_) + for ranks in expert_parallel_groups: + group = dist.new_group(ranks) + if rank in list(ranks): + _EXPERT_PARALLEL_GROUP[group_name] = group + + for ranks in expert_data_parallel_groups: + group = dist.new_group(ranks) + if rank in list(ranks): + _EXPERT_DATA_PARALLEL_GROUP[group_name] = group + + +def _get_max_expert_size(): + """Get the maximum ep_size from all the created groups.""" + assert _EXPERT_PARALLEL_GROUP is not None, "Warning! Process group not initialized" + keylist = [] + for key in _EXPERT_PARALLEL_GROUP.keys(): + # index 2 is ep_size in the group name: ep_size_ + index = 2 + keylist.append(int(key.split('_')[index])) + return max(keylist) if len(keylist) > 0 else None + + +def _get_max_expert_size_name(): + """Get the name of the group with max. ep_size""" + return f'ep_size_{_get_max_expert_size()}' + + +def _get_max_expert_parallel_group(): + """Get the max expert parallel size.""" + return _get_expert_parallel_group(_get_max_expert_size_name()) + + +def _get_expert_parallel_group(group_name): + """Get the expert parallel group the caller rank belongs to.""" + assert group_name in _EXPERT_PARALLEL_GROUP, \ + 'expert parallel group is not initialized' + return _EXPERT_PARALLEL_GROUP[group_name] + + +def _get_expert_parallel_group_dict(): + """Get the expert parallel group dict.""" + return _EXPERT_PARALLEL_GROUP + + +def _get_expert_data_parallel_group(group_name): + """Get the expert data parallel group the caller rank belongs to.""" + assert group_name in _EXPERT_DATA_PARALLEL_GROUP, \ + 'expert data parallel group is not initialized' + return _EXPERT_DATA_PARALLEL_GROUP[group_name] + + +def _get_expert_data_parallel_group_dict(): + """Get the expert data parallel group dict.""" + return _EXPERT_DATA_PARALLEL_GROUP + + +def _clone_world_group(): + """Create a clone of the world group + Note: We need to clone the dist world group because we + use dist.get_global_rank() utility function in DeepSpeed at many places. + As that function does not work on dist.group.WORLD, we + need to keep a clone of it. + """ + assert dist.is_initialized(), "dist is not initialized" + global _WORLD_GROUP + if _WORLD_GROUP is None: + # If not cloned already, clone the world group + _WORLD_GROUP = dist.new_group(ranks=range(dist.get_world_size())) + return _WORLD_GROUP + + +def _get_local_all_to_all_group(): + assert dist.is_initialized(), 'dist is not initialized' + global _ALL_TO_ALL_GROUP + device_per_node = get_accelerator().device_count() + num_local = dist.get_world_size() // device_per_node + if num_local == 0 and dist.get_world_size() > 0: + assert dist.get_world_size() >= 1, 'num_gpus must >=1, cannot initialize All-To-All' + cur_rank = [] + for i in range(dist.get_world_size()): + cur_rank.append(i) + _ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=cur_rank) + elif num_local == 1: + assert dist.get_world_size( + ) == device_per_node, 'num_gpus not equal to device per node, cannot initialize All-To-All' + _ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=[i for i in range(device_per_node)]) + else: + assert dist.get_world_size() > device_per_node, 'num_nodes<2 cannot initialize All-To-All' + for i in range(num_local): + local_rank = [j + device_per_node * i for j in range(device_per_node)] + _ALL_TO_ALL_GROUP[f"local_{i}"] = dist.new_group(ranks=local_rank) + + for i in range(device_per_node): + cur_rank = [] + for j in range(num_local): + cur_rank.append(i + j * device_per_node) + _ALL_TO_ALL_GROUP[f"global_{i}"] = dist.new_group(ranks=cur_rank) + return _ALL_TO_ALL_GROUP + + +def _get_data_parallel_group(): + """Get the data parallel group the caller rank belongs to.""" + assert dist.is_initialized(), 'dist is not initialized' + global mpu + if mesh_device is not None: + return mesh_device.get_group(mesh_dim="data_parallel") + if mpu is not None: + if hasattr(mpu, 'initialize_sequence_parallel'): + return None + else: + return mpu.get_data_parallel_group() + + # Return the clone of dist world group + return _clone_world_group() + + +def _get_broadcast_src_rank(): + return dist.get_global_rank(_get_sequence_data_parallel_group(), 0) + + +def _get_expert_broadcast_src_rank(group_name): + return dist.get_global_rank(_get_expert_data_parallel_group(group_name), 0) + + +def _get_expert_parallel_world_size(group_name): + """Return world size for the expert parallel group.""" + return dist.get_world_size(group=_get_expert_parallel_group(group_name)) + + +def _get_expert_data_parallel_world_size(group_name): + """Return world size for the expert data parallel group.""" + return dist.get_world_size(group=_get_expert_data_parallel_group(group_name)) + + +def _get_expert_parallel_rank(group_name): + """Return my rank for the expert parallel group.""" + return dist.get_rank(group=_get_expert_parallel_group(group_name)) + + +def _get_expert_parallel_src_rank(group_name): + """Calculate the global rank corresponding to a local rank zero + in the expert parallel group.""" + global_rank = dist.get_rank() + local_world_size = _get_expert_parallel_world_size(group_name) + return (global_rank // local_world_size) * local_world_size + + +def _get_expert_data_parallel_rank(group_name): + """Return my rank for the expert data parallel group.""" + return dist.get_rank(group=_get_expert_data_parallel_group(group_name)) + + +def _get_data_parallel_world_size(): + """Return world size for the data parallel group.""" + if mesh_device is not None: + return dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel")) + global mpu + if mpu is not None: + if hasattr(mpu, 'initialize_sequence_parallel'): + return None + else: + return mpu.get_data_parallel_world_size() + return dist.get_world_size(group=_get_data_parallel_group()) + + +def _get_model_parallel_world_size(): + """Return world size for the model parallel group.""" + global mpu + if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'): + return 1 + return mpu.get_model_parallel_world_size() + + +def _get_data_parallel_rank(): + """Return my rank for the data parallel group.""" + return dist.get_rank(group=_get_data_parallel_group()) + + +def _get_sequence_parallel_world_size(): + """Return world size for the sequence parallel group.""" + global mpu + if mesh_device is not None: + return dist.get_world_size(mesh_device.get_group(mesh_dim="sequence_parallel")) + if mpu is not None and hasattr(mpu, 'get_sequence_parallel_world_size'): + return mpu.get_sequence_parallel_world_size() + return 1 + + +def _get_sequence_parallel_rank(): + """Return my rank for the sequence parallel group.""" + global mpu + if mpu is not None and hasattr(mpu, 'get_sequence_parallel_rank'): + return mpu.get_sequence_parallel_rank() + if mesh_device is not None: + return dist.get_rank(mesh_device.get_group(mesh_dim="sequence_parallel")) + return 0 + + +def _get_sequence_parallel_group(): + global mpu + if mpu is None or not hasattr(mpu, 'get_sequence_parallel_group'): + if mesh_device is None: + raise KeyError("No sequence parallel group found") + return mesh_device.get_group(mesh_dim="sequence_parallel") + return mpu.get_sequence_parallel_group() + + +def _get_sequence_data_parallel_world_size(): + """Return world size for the model parallel group.""" + global mpu + if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_world_size'): + return mpu.get_sequence_data_parallel_world_size() + return _get_data_parallel_world_size() + + +def _get_sequence_data_parallel_rank(): + """Return my rank for the data parallel group.""" + global mpu + if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_rank'): + return mpu.get_sequence_data_parallel_rank() + return _get_data_parallel_rank() + + +def _get_sequence_data_parallel_group(): + global mpu + # When sequence parallelism is enabled, the process group for zero sharding and + # gradient allreduce must be across both dimensions of data and sequence parallelism. + if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_group'): + return mpu.get_sequence_data_parallel_group() + return _get_data_parallel_group() + + +def _get_expert_model_parallel_world_size(): + global expert_tensor_parallel_world_size + return expert_tensor_parallel_world_size + + +def _create_zero_param_parallel_group(group_size): + """ + Create parameter partitioning group within ZeRO data parallel groups. + + Example - ZP + D parallel + world_size = 16 + zero_hpz_partition_size = 2 # number of ranks with replicated params (dual partitioning) + zero_param_intra_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - segmented (subgroup) with rep partition + data_parallel_group = [0,1,...,15] - all reduce is on ZeRO model + """ + assert dist.is_initialized() + global _ZERO_PARAM_INTRA_PARALLEL_GROUP + # Only create group if it does not already exist + assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is None, \ + 'ZeRO parameter intra parallel group is already initialized' + + world_size = dist.get_world_size() + rank = dist.get_rank() + + zero_param_parallel_size_ = min(group_size, world_size) + _ensure_divisibility(world_size, zero_param_parallel_size_) + + # Build the ZeRO param intra parallel groups. + for i in range(world_size // zero_param_parallel_size_): + ranks = range(i * zero_param_parallel_size_, (i + 1) * zero_param_parallel_size_) + group = dist.new_group(ranks) + if i == (rank // zero_param_parallel_size_): + _ZERO_PARAM_INTRA_PARALLEL_GROUP = group + + +def _get_zero_param_intra_parallel_group(): + """Get the ZeRO parameter partitioning intra parallel group the caller rank belongs to.""" + #assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is not None, \ + # 'ZeRO parameter partitioning group is not initialized' + #TODO: Add warning + return _ZERO_PARAM_INTRA_PARALLEL_GROUP + + +def _zero_param_parallel_is_initialized(): + """Check if ZeRO data parallel with parameter partititioning groups are initialized.""" + ###TODO: assert that MPU is not set + if _ZERO_PARAM_INTRA_PARALLEL_GROUP is None and _DATA_PARALLEL_GROUP is None: + return False + + +def _get_zero_param_intra_parallel_rank_in_mygroup(): + """Return my rank for the ZeRO parameter inter parallel group.""" + return dist.get_rank(group=_get_zero_param_intra_parallel_group()) + + +def _get_zero_param_intra_parallel_group_world_size(): + """Return world size for the ZeRO parameter parallel group.""" + return dist.get_world_size(group=_get_zero_param_intra_parallel_group()) + + +def _get_zero_param_intra_parallel_group_ranks(): + """Return all ranks for the ZeRO parameter intra parallel group.""" + return dist.get_all_ranks_from_group(group=_get_zero_param_intra_parallel_group()) diff --git a/lib/python3.12/site-packages/deepspeed/utils/init_on_device.py b/lib/python3.12/site-packages/deepspeed/utils/init_on_device.py new file mode 100644 index 0000000000000000000000000000000000000000..52dbf71d9562b73d868f71ee7ec7894a187e00f0 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/init_on_device.py @@ -0,0 +1,81 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from typing import Callable +from torch import Tensor +from packaging import version as pkg_version + + +class OnDevice(object): + """ + Create modules/tensors w. specific devices and dtypes. Examples: + + Create MyModule which consists of many different sub-modules and parameters. In this case we can create + MyModule as a collection of 'meta' tensors by passing `device='meta'` or we can create the module _directly_ + on a CUDA device by passing `device=f'cuda:{local_rank}'` (where `local_rank` is the local GPU id. + + with OnDevice(dtype=torch.float16, device='meta'): + model = MyModel() + + with OnDevice(dtype=torch.float16, device=f'cuda:{local_rank}'): + model = MyModel() + + """ + + _orig_torch_empty = torch.empty + _orig_torch_zeros = torch.zeros + _orig_torch_ones = torch.ones + _orig_torch_full = torch.full + + def __init__(self, dtype, device="meta", enabled=True): + self.dtype = dtype + self.enabled = enabled + self.device = device + + if device == "meta": + if pkg_version.parse('1.10') > pkg_version.parse(torch.__version__): + raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") + + def fp_tensor_constructor(self, fn: Callable, target_fp_dtype: torch.dtype) -> Callable: + + def wrapped_fn(*args, **kwargs) -> Tensor: + if kwargs.get("device", None) is None: + kwargs['device'] = self.device + tensor: Tensor = fn(*args, **kwargs) + if tensor.is_floating_point(): + tensor = tensor.to(target_fp_dtype) + return tensor + + return wrapped_fn + + def get_new_tensor_fn_for_dtype(self, dtype: torch.dtype) -> Callable: + + def new_tensor(cls, *args) -> Tensor: + tensor = OnDevice._orig_torch_empty(0, device=self.device).new_empty(*args) + if tensor.is_floating_point(): + tensor = tensor.to(dtype) + return tensor + + return new_tensor + + def __enter__(self): + if not self.enabled: + return + torch.Tensor.__old_new__ = torch.Tensor.__new__ + torch.Tensor.__new__ = self.get_new_tensor_fn_for_dtype(self.dtype) + torch.empty = self.fp_tensor_constructor(self._orig_torch_empty, self.dtype) + torch.zeros = self.fp_tensor_constructor(self._orig_torch_zeros, self.dtype) + torch.ones = self.fp_tensor_constructor(self._orig_torch_ones, self.dtype) + torch.full = self.fp_tensor_constructor(self._orig_torch_full, self.dtype) + + def __exit__(self, exc_type, exc_value, traceback): + if not self.enabled: + return + torch.Tensor.__new__ = torch.Tensor.__old_new__ + torch.empty = self._orig_torch_empty + torch.zeros = self._orig_torch_zeros + torch.ones = self._orig_torch_ones + torch.full = self._orig_torch_full diff --git a/lib/python3.12/site-packages/deepspeed/utils/logging.py b/lib/python3.12/site-packages/deepspeed/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..ed691e3985e1f1ca7360fe8dccb30f886013f411 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/logging.py @@ -0,0 +1,163 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import functools +import logging +import sys +import os +import torch +from deepspeed.utils.torch import required_torch_version + +log_levels = { + "debug": logging.DEBUG, + "info": logging.INFO, + "warning": logging.WARNING, + "error": logging.ERROR, + "critical": logging.CRITICAL, +} + + +class LoggerFactory: + + @staticmethod + def create_logger(name=None, level=logging.INFO): + """create a logger + + Args: + name (str): name of the logger + level: level of logger + + Raises: + ValueError is name is None + """ + + if name is None: + raise ValueError("name for logger cannot be None") + + formatter = logging.Formatter("[%(asctime)s] [%(levelname)s] " + "[%(filename)s:%(lineno)d:%(funcName)s] %(message)s") + + logger_ = logging.getLogger(name) + logger_.setLevel(level) + logger_.propagate = False + ch = logging.StreamHandler(stream=sys.stdout) + ch.setLevel(level) + ch.setFormatter(formatter) + logger_.addHandler(ch) + if required_torch_version(min_version=2.6) and os.getenv("DISABLE_LOGS_WHILE_COMPILING", "0") == "1": + excluded_set = { + item.strip() + for item in os.getenv("LOGGER_METHODS_TO_EXCLUDE_FROM_DISABLE", "").split(",") + } + ignore_set = {'info', 'debug', 'error', 'warning', 'critical', 'exception', 'isEnabledFor'} - excluded_set + for method in ignore_set: + original_logger = getattr(logger_, method) + torch._dynamo.config.ignore_logger_methods.add(original_logger) + return logger_ + + +logger = LoggerFactory.create_logger(name="DeepSpeed", level=logging.INFO) + + +@functools.lru_cache(None) +def warning_once(*args, **kwargs): + """ + This method is identical to `logger.warning()`, but will emit the warning with the same message only once + + Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache. + The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to + another type of cache that includes the caller frame information in the hashing function. + """ + logger.warning(*args, **kwargs) + + +logger.warning_once = warning_once + + +def print_configuration(args, name): + logger.info("{}:".format(name)) + for arg in sorted(vars(args)): + dots = "." * (29 - len(arg)) + logger.info(" {} {} {}".format(arg, dots, getattr(args, arg))) + + +def log_dist(message, ranks=None, level=logging.INFO): + from deepspeed import comm as dist + """Log message when one of following condition meets + + + not dist.is_initialized() + + dist.get_rank() in ranks if ranks is not None or ranks = [-1] + + Args: + message (str) + ranks (list) + level (int) + + """ + should_log = not dist.is_initialized() + ranks = ranks or [] + my_rank = dist.get_rank() if dist.is_initialized() else -1 + if ranks and not should_log: + should_log = ranks[0] == -1 + should_log = should_log or (my_rank in set(ranks)) + if should_log: + final_message = "[Rank {}] {}".format(my_rank, message) + logger.log(level, final_message) + + +def print_json_dist(message, ranks=None, path=None): + from deepspeed import comm as dist + """Print message when one of following condition meets + + + not dist.is_initialized() + + dist.get_rank() in ranks if ranks is not None or ranks = [-1] + + Args: + message (str) + ranks (list) + path (str) + + """ + should_log = not dist.is_initialized() + ranks = ranks or [] + my_rank = dist.get_rank() if dist.is_initialized() else -1 + if ranks and not should_log: + should_log = ranks[0] == -1 + should_log = should_log or (my_rank in set(ranks)) + if should_log: + message['rank'] = my_rank + import json + with open(path, 'w') as outfile: + json.dump(message, outfile) + os.fsync(outfile) + + +def get_current_level(): + """ + Return logger's current log level + """ + return logger.getEffectiveLevel() + + +def should_log_le(max_log_level_str): + """ + Args: + max_log_level_str: maximum log level as a string + + Returns ``True`` if the current log_level is less or equal to the specified log level. Otherwise ``False``. + + Example: + + ``should_log_le("info")`` will return ``True`` if the current log level is either ``logging.INFO`` or ``logging.DEBUG`` + """ + + if not isinstance(max_log_level_str, str): + raise ValueError(f"{max_log_level_str} is not a string") + + max_log_level_str = max_log_level_str.lower() + if max_log_level_str not in log_levels: + raise ValueError(f"{max_log_level_str} is not one of the logging levels") + + return get_current_level() <= log_levels[max_log_level_str] diff --git a/lib/python3.12/site-packages/deepspeed/utils/mixed_precision_linkage.py b/lib/python3.12/site-packages/deepspeed/utils/mixed_precision_linkage.py new file mode 100644 index 0000000000000000000000000000000000000000..c97515ca8fef7477e115914e6410bc0e778e3b8a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/mixed_precision_linkage.py @@ -0,0 +1,51 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import types +from deepspeed.utils import get_full_hp_param, get_full_hp_grad, get_hp_fragment_mapping +from deepspeed.utils import set_full_hp_param, set_full_hp_grad + + +def link_hp_params(lp_param_list, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload, + param_group_index, partition_start, partition_size, dp_group): + local_lp_param_and_offset = _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_group) + + for lp_param, lp_start in local_lp_param_and_offset: + lp_param._hp_mapping = get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict, + offload_gradient_dict, use_offload, param_group_index, + partition_start, partition_size) + + +def lazy_init_hp_params_optimizer_state(lp_param_list, flat_hp_partition, optimizer_state): + for lp in lp_param_list: + if lp._hp_mapping is not None: + lp._hp_mapping.set_optim_state_fragment(flat_hp_partition, optimizer_state[flat_hp_partition]) + + +def _init_lp_to_hp_mapping(lp_param_list, partition_start, partition_size, dp_group): + current_offset = 0 + param_and_offset_list = [] + partition_end = partition_start + partition_size + index_in_param_group = 0 + for i, lp_param in enumerate(lp_param_list): + lp_param._hp_mapping = None + lp_param._dp_group = dp_group + lp_param.get_full_hp_param = types.MethodType(get_full_hp_param, lp_param) + lp_param.get_full_hp_grad = types.MethodType(get_full_hp_grad, lp_param) + lp_param.set_full_hp_param = types.MethodType(set_full_hp_param, lp_param) + lp_param.set_full_hp_grad = types.MethodType(set_full_hp_grad, lp_param) + + # lp_param overlaps with partition if both are true + # 1) current_offset < partition_end, + # 2) current_offset + lp_param.numel() >= partition_start + lp_param_end = current_offset + lp_param.numel() + if current_offset < partition_end and lp_param_end > partition_start: + param_and_offset_list.append((lp_param, current_offset)) + lp_param._index_in_param_group = index_in_param_group + # Indices for params in this partition/GPU + index_in_param_group += 1 + current_offset += lp_param.numel() + + return param_and_offset_list diff --git a/lib/python3.12/site-packages/deepspeed/utils/numa.py b/lib/python3.12/site-packages/deepspeed/utils/numa.py new file mode 100644 index 0000000000000000000000000000000000000000..75f0442a04e53042b9593cb72a7b9aae836ba25c --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/numa.py @@ -0,0 +1,205 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +# return a list of list for cores to numa mapping +# [ +# [ cores for numa 0 ] +# [ cores belong to numa 1 ] +# ... +# ] + +import os +import psutil +import shutil +import subprocess + + +# return a list of list for cores to numa mapping +# [ +# [ cores for numa 0 ] +# [ cores belong to numa 1 ] +# ... +# ] +def get_numa_cores(): + ret = [] + try: + output = subprocess.check_output(['numactl', '--hardware']).decode("utf-8") + except: + return [] + lines = output.split('\n') + for line in lines: + if line.startswith('available:'): + num_numas = int(line.split(' ')[1]) + break + for numa in range(num_numas): + for line in lines: + if line.startswith(f'node {numa} cpus:'): + cores = line.split(' ')[3:] + ret.append([int(core) for core in cores]) + return ret + + +def check_for_numactl_pkg(): + libs = dict( + dpkg=["-l", "numactl", "apt"], + pacman=["-Q", "numactl", "pacman"], + rpm=["-q", "numactl", "yum"], + ) + + found = False + for pkgmgr, data in libs.items(): + flag, lib, tool = data + path = shutil.which(pkgmgr) + if path is not None: + cmd = [pkgmgr, flag, lib] + result = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if result.wait() == 0: + found = True + else: + print(f"please install the {lib} package with {tool}") + break + return found + + +def parse_range(rng): + try: + value = int(rng) + return range(value, value + 1) + except ValueError: + # value is not a single number + parts = rng.split('-') + if len(parts) != 2: + raise ValueError("Bad range: '%s', range must be either a number or two number separated by dash" % + (rng, )) + start = int(parts[0]) + end = int(parts[1]) + if start > end: + raise ValueError("Bad range: '%s', range end must larger than or equal to start" % (rng, )) + return range(start, end + 1) + + +# parse comma and dash separated range list into list +# i.e. "0,2-4,6" --> [0, 2, 3, 4, 6] +# rules: +# 1. Range list number be comma separated, each item are either a single number, +# or a range marked by two numbers (both number are included in the range) +# 2. Sub ranges must be in ascend order and not overlap with each other +# 3. No space in the range expression +def parse_range_list(range_str): + number_list = [] + last = -1 + range_list = range_str.split(',') + for sub_range in range_list: + sub_number_list = parse_range(sub_range) + if sub_number_list[0] <= last: + raise ValueError( + "Bad range: '%s', sub ranges must not overlap with each other and should be in ascend order" % + (range_str, )) + last = sub_number_list[-1] + number_list.extend(sub_number_list) + return number_list + + +def get_numactl_cmd(bind_core_list, num_local_procs, local_rank): + numactl_cmd = [] + check_for_numactl_pkg() + if 'KMP_AFFINITY' in os.environ.keys(): + raise ValueError("Environment variable KMP_AFFINITY conflicts with numactl " + "because it interfere with how many CPU cores numactl can set. " + "Unset KMP_AFFINITY before launching deepspeed.\n\n" + "\t$ unset KMP_AFFINITY\n" + "\t$ deepspeed ") + if bind_core_list is not None: + core_list = parse_range_list(bind_core_list) + total_cores = len(core_list) + else: + total_cores = psutil.cpu_count(logical=False) + core_list = range(total_cores) + cores_per_rank = total_cores // num_local_procs + assert cores_per_rank >= 1, "At least one core needs to be assigned to each rank" + core_list_for_rank = core_list[cores_per_rank * local_rank:cores_per_rank * (local_rank + 1)] + numactl_cmd.append("numactl") + + # check if all cores belong to same numa, if true, bind process to that numa domain with -m parameter + numa_cores = get_numa_cores() + num_numas = len(numa_cores) + + numa_mode = "normal" + + non_empty_numa_list = [] + empty_numa_list = [] + previous_numa_cores = [] + numa_node_list = [] + numa_node_list_list = [] + for i in range(num_numas): + # look for empty numa which is HBM numa + if numa_cores[i] == []: + empty_numa_list.append(i) + else: + non_empty_numa_list.append(i) + + # check for fakenuma + if numa_cores[i] == previous_numa_cores: + if numa_node_list == []: + #first duplication, add previous node into list + numa_node_list.append(i - 1) + numa_node_list.append(i) + else: + if numa_node_list != []: + numa_node_list_list.append(numa_node_list) + numa_node_list = [] + previous_numa_cores = numa_cores[i] + if numa_node_list != []: + numa_node_list_list.append(numa_node_list) + + if empty_numa_list != [] and len(empty_numa_list) == len(non_empty_numa_list): + numa_mode = "flat_hbm" + numa_dict = dict(zip(non_empty_numa_list, empty_numa_list)) + elif numa_node_list_list != []: + numa_mode = "fake" + + if numa_mode == "normal": + for i in range(num_numas): + if set(core_list_for_rank) <= set(numa_cores[i]): + numactl_cmd.append("-m") + numactl_cmd.append(f"{i}") + break + elif numa_mode == "flat_hbm": + for i in range(num_numas): + if set(core_list_for_rank) <= set(numa_cores[i]): + numactl_cmd.append("-p") + numactl_cmd.append(f"{numa_dict[i]}") + break + elif numa_mode == "fake": + for i in range(num_numas): + if set(core_list_for_rank) <= set(numa_cores[i]): + for nodes in numa_node_list_list: + if i in nodes: + numactl_cmd.append("-m") + numactl_cmd.append(f"{','.join(map(str, nodes))}") + break + # the following construct break the outer loop if inner loop breaks + else: + continue + break + + numactl_cmd.append("-C") + last_core = core_list_for_rank[0] + first_core = last_core + core_list_str = f"{last_core}" + for core_id in core_list_for_rank[1:]: + if core_id == last_core + 1: + last_core = core_id + continue + else: + if first_core == last_core: + core_list_str = f"{core_list_str},{core_id}" + else: + core_list_str = f"{core_list_str}-{last_core},{core_id}" + first_core = core_id + last_core = core_id + if first_core != last_core: + core_list_str = f"{core_list_str}-{last_core}" + numactl_cmd.append(f"{core_list_str}") + return cores_per_rank, numactl_cmd diff --git a/lib/python3.12/site-packages/deepspeed/utils/nvtx.py b/lib/python3.12/site-packages/deepspeed/utils/nvtx.py new file mode 100644 index 0000000000000000000000000000000000000000..72d7c863a33f5628fd7c327a6b11652ccba17a01 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/nvtx.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.compiler import is_compiling + +enable_nvtx = True + + +def instrument_w_nvtx(func): + """Decorator that records an NVTX range for the duration of the function call. + Skips NVTX instrumentation when torch.compile is active to avoid graph breaks. + """ + + def wrapped_fn(*args, **kwargs): + if enable_nvtx and not is_compiling(): + get_accelerator().range_push(func.__qualname__) + ret_val = func(*args, **kwargs) + if enable_nvtx and not is_compiling(): + get_accelerator().range_pop() + return ret_val + + return wrapped_fn diff --git a/lib/python3.12/site-packages/deepspeed/utils/tensor_fragment.py b/lib/python3.12/site-packages/deepspeed/utils/tensor_fragment.py new file mode 100644 index 0000000000000000000000000000000000000000..39cf094be5c622dcb3fc9329f9e889e1d601554a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/tensor_fragment.py @@ -0,0 +1,479 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from dataclasses import dataclass +from deepspeed import comm as dist +from typing import Dict, List, Callable + + +@dataclass +class fragment_address: + numel: int + start: int + + +@dataclass +class tensor_fragment: + lp_fragment: torch.Tensor + lp_fragment_address: fragment_address + hp_fragment: torch.Tensor + hp_fragment_address: fragment_address + gradient_dict: Dict + offload_gradient_dict: Dict + use_offload: bool + param_group_index: int + optim_fragment: Dict = None + + def update_hp(self): + self.hp_fragment.data.copy_(self.lp_fragment.data) + + def update_lp(self): + self.lp_fragment.data.copy_(self.hp_fragment.data) + + def get_optim_state_fragment(self, key): + if key in self.optim_fragment: + return self.optim_fragment[key] + else: + raise ValueError(f'{key} not found in optimizer state fragment') + + def set_optim_state_fragment(self, flat_hp_partition, optim_fragment): + self.optim_fragment = { + key: value.narrow(0, self.hp_fragment_address.start, self.hp_fragment_address.numel) + for key, value in optim_fragment.items() + if torch.is_tensor(value) and value.shape == flat_hp_partition.shape + } + + def get_hp_fragment_address(self): + return self.hp_fragment_address + + def get_optim_state_keys(self): + return list(self.optim_fragment.keys()) + + def get_hp_fragment(self, optim_state_key=None): + if optim_state_key is None: + return self.hp_fragment + return self.get_optim_state_fragment(optim_state_key) + + def get_lp_grad_fragment(self, index_in_param_group): + if self.use_offload: + gradient_dict = self.offload_gradient_dict + else: + gradient_dict = self.gradient_dict + + if self.param_group_index not in gradient_dict or gradient_dict[self.param_group_index] is None: + raise ValueError("Gradients are only available immediately after backward and before engine step") + + return gradient_dict[self.param_group_index][index_in_param_group] + + +def map_to_flat_opt_states(flat_hp_tensor, lp_tensors, optim_state, opt_keys): + for key in opt_keys: + hp_param = flat_hp_tensor + buffer = torch.zeros_like(hp_param) + + for lp in lp_tensors: + if lp._hp_mapping is not None: + hp_fragment_address = lp._hp_mapping.get_hp_fragment_address() + hp_fragment = buffer.narrow(0, hp_fragment_address.start, hp_fragment_address.numel) + hp_fragment.data.copy_(lp._hp_mapping.get_hp_fragment(optim_state_key=key).data) + lp._hp_mapping.hp_fragment = hp_fragment + + optim_state[hp_param][key] = buffer + + +def get_full_hp_param(self, optim_state_key=None): + reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten() + if self._hp_mapping is not None: + lp_frag_address = self._hp_mapping.lp_fragment_address + reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel) + hp_fragment = self._hp_mapping.get_hp_fragment(optim_state_key) + reduce_fragment.data.copy_(hp_fragment.data) + dist.all_reduce(reduce_buffer, group=self._dp_group) + return reduce_buffer.reshape_as(self) + + +def set_full_hp_param(self, value, optim_state_key=None): + if self._hp_mapping is not None: + lp_frag_address = self._hp_mapping.lp_fragment_address + value_fragment = torch.narrow(value.flatten(), 0, lp_frag_address.start, lp_frag_address.numel) + hp_fragment = self._hp_mapping.get_hp_fragment(optim_state_key) + hp_fragment.data.copy_(value_fragment.data) + + +def get_full_hp_grad(self): + reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten() + if self._hp_mapping is not None: + lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group) + hp_grad_fragment = lp_grad_fragment.to(torch.float32).flatten() + + lp_frag_address = self._hp_mapping.lp_fragment_address + reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel) + + if self.view(-1).shape == hp_grad_fragment.shape: + reduce_buffer.data.copy_(hp_grad_fragment.data) + else: + reduce_fragment.data.copy_(hp_grad_fragment.data) + + dist.all_reduce(reduce_buffer, group=self._dp_group) + return reduce_buffer.reshape_as(self) + + +def set_full_hp_grad(self, value): + if self._hp_mapping is not None: + lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group) + lp_frag_address = self._hp_mapping.lp_fragment_address + value_fragment = torch.narrow(value.flatten(), 0, lp_frag_address.start, lp_frag_address.numel) + lp_grad_fragment.data.copy_(value_fragment.data.reshape_as(lp_grad_fragment.data)) + + +def safe_get_full_fp32_param(param): + """Assemble and return the fp32 parameter of a low-precision (e.g., fp16) parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + return param._z3_optimizer.get_full_hp_param(param) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + return param.get_full_hp_param() + return None + + +def safe_set_full_fp32_param(param, value): + """Update the partitioned fp32 parameter of a low-precision (e.g., fp16) parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter + value (``torch.Tensor``): New value + """ + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + param._z3_optimizer.set_full_hp_param(value, param) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + param.set_full_hp_param(value) + + +def safe_get_full_optimizer_state(param, optim_state_key): + """Assemble and return the fp32 optimizer state of a low-precision (e.g., fp16) parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter + optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device +""" + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + return param._z3_optimizer.get_full_hp_param(param, optim_state_key) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + return param.get_full_hp_param(optim_state_key) + return None + + +def safe_set_full_optimizer_state(param, value, optim_state_key): + """Update the partitioned fp32 optimizer state of a low-precision (e.g., fp16) parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter + value (``torch.Tensor``): New value + optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) + """ + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + param._z3_optimizer.set_full_hp_param(value, param, optim_state_key) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + param.set_full_hp_param(value, optim_state_key) + + +# TODO: Figure out the correct return dtype +def safe_get_full_grad(param): + """ + Assemble and return the fp32 gradient of a low-precision (e.g., fp16) parameter. + The return data type is that used for gradient accumulation. This is usually the param data type, + but could also be different (e.g., bf16 param training with fp32 gradient accumulation). + + Args: + param (``torch.nn.Parameter``): A model parameter + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + if param.grad is not None: + return param.grad + + # ZeRO stage 3 param + if hasattr(param, 'ds_id'): + return param._z3_optimizer.get_fp32_grad_for_param(param) + + # ZeRO stage 1, 2, and bf16_optimizer params + if hasattr(param, '_hp_mapping'): + return param.get_full_hp_grad() + + return None + + +def safe_set_full_grad(param, value): + """ + Update the partitioned gradient of a low-precision (e.g., fp16) parameter. + To avoid precision issues, the update value should have the data type of + gradient accumulation. + + Args: + param (``torch.nn.Parameter``): A model parameter + value (``torch.Tensor``): The un-partitioned new gradient value. + """ + if param.grad is not None: + param.grad.copy_(value) + elif hasattr(param, 'ds_id'): + # ZeRO stage 3 param + param._z3_optimizer.set_fp32_grad_for_param(value, param) + elif hasattr(param, '_hp_mapping'): + # ZeRO stage 1, 2, and bf16_optimizer params + param.set_full_hp_grad(value) + + +### Local API START ### +def safe_get_local_grad(param): + """ + Get the local gradient partition of a ZeRO-3 partitioned parameter. + The return data type is that used for gradient accumulation. This is usually the param data type, + but could also be different (e.g., bf16 param training with fp32 gradient accumulation). + + Args: + param (``torch.nn.Parameter``): A model parameter + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + return param._z3_optimizer.get_local_fp32_grad_for_param(param) + + +def safe_set_local_grad(param, value): + """ + Update the local gradient partition of a ZeRO-3 partitioned parameter. + To avoid precision issues, the update value should have the data type of + gradient accumulation. + + Args: + param (``torch.nn.Parameter``): A model parameter. + value (``torch.Tensor``): New value of local gradient partition. + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + param._z3_optimizer.set_local_grad_for_param(value, param) + + +def safe_get_local_fp32_param(param): + """Get the local partition of a ZeRO-3 partitioned parameter in fp32 precision. + + Args: + param (``torch.nn.Parameter``): A model parameter. + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + return param._z3_optimizer.get_local_fp32_param(param) + + +def safe_get_local_optimizer_state(param, optim_state_key): + """Get the local optimizer state partition of ZeRO-3 partitioned parameter in fp32 precision. + + Args: + param (``torch.nn.Parameter``): A model parameter + optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer) + + Returns: + Union[torch.Tensor, None]: A tensor on accelerator device + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + return param._z3_optimizer.get_local_fp32_param(param, optim_state_key) + + +def safe_set_local_optimizer_state(param, value, optim_state_key): + """Update the local optimizer state partition of a ZeRO-3 partitioned parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter. + value (``torch.Tensor``): New value of local optimizer state partition. + optim_state_key (``string``): Key value of optimizer state (e.g., `exp_avg` in Adam optimizer). + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + param._z3_optimizer.set_local_hp_param(value, param, optim_state_key) + + +def safe_set_local_fp32_param(param, value): + """Update the local partition of ZeRO-3 partitioned parameter. + + Args: + param (``torch.nn.Parameter``): A model parameter. + value (``torch.Tensor``): New value of local parameter partition. + """ + assert hasattr(param, 'ds_id'), f'This API is only defined for ZeRO-3 partitioned parameters' + param._z3_optimizer.set_local_hp_param(value, param) + + +### Local API END ### + + +### VECTORIZED API BEGIN ### +def safe_update_full_grad_vectorized(param_list: List[torch.nn.Parameter], update_func: Callable): + """ + Vectorized update of the partitioned gradients of a list of low-precision (e.g., fp16) parameters. + To avoid precision issues, the update value should have the data type of + gradient accumulation. + + Args: + param_list (``List[torch.nn.Parameter]``): List of model parameters + update_func (``torch.Tensor``): A function that takes current full gradient value and returns new one. + """ + partitioned_grad_params = [] + for p in param_list: + if p.grad is not None: + p.grad.copy_(update_func(p.grad, p)) + elif p.requires_grad: + partitioned_grad_params.append(p) + + if not partitioned_grad_params: + return + + if hasattr(partitioned_grad_params[0], 'ds_id'): + # ZeRO stage 3 param + partitioned_grad_params[0]._z3_optimizer.update_fp32_grad_for_param_vectorized( + update_func, partitioned_grad_params) + elif hasattr(partitioned_grad_params[0], '_hp_mapping'): + # ZeRO stage 1, 2, and bf16_optimizer params + for p in partitioned_grad_params: + old_grad = safe_get_full_grad(p) + new_grad = update_func(old_grad, p) + p.set_full_hp_grad(new_grad) + + +### VECTORIZED API END ### + + +def get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload, + param_group_index, partition_start, partition_size): + lp_end = lp_param.numel() + lp_start + hp_start = partition_start + hp_end = partition_start + partition_size + + fragment_start = max(lp_start, hp_start) + fragment_end = min(lp_end, hp_end) + assert fragment_start < fragment_end, \ + f'fragment start {fragment_start} should be < fragment_end {fragment_end}' + + fragment_numel = fragment_end - fragment_start + hp_frag_address = fragment_address(start=fragment_start - hp_start, numel=fragment_numel) + hp_fragment_tensor = flat_hp_partition.narrow(0, hp_frag_address.start, hp_frag_address.numel) + + lp_frag_address = fragment_address(start=fragment_start - lp_start, numel=fragment_numel) + lp_fragment_tensor = lp_param.flatten().narrow(0, lp_frag_address.start, lp_frag_address.numel) + + return tensor_fragment(lp_fragment=lp_fragment_tensor, + lp_fragment_address=lp_frag_address, + hp_fragment=hp_fragment_tensor, + hp_fragment_address=hp_frag_address, + gradient_dict=gradient_dict, + offload_gradient_dict=offload_gradient_dict, + use_offload=use_offload, + param_group_index=param_group_index) + + +''' +Logic for lp_param to hp_param mapping + +lp lp0 lp1 lp2 lp3 lp4 <------- indices/names +lp [ ][ ][ ][ ][ ] <-------- tensors +flat_lp [ ] <-------- flat lp params +flat_hp [ ] <------------------ flat hp partition on current rank +full_hp [ ] <------- full flat hp params + + +lp2 + full numel = 16 + lp_frag + numel = 12 + frag_start = 3 + frag_end = 15 + hp_frag + numel = 12 + frag_start = 0 + frag_end = 11 + + hp_frag.copy_(lp_frag) + + +lp3: + full numel = 4 + lp_frag + numel = 4 + start = 0 + end = 3 + hp_frag + numel = 4 + start = 12 + end = 15 + + +lp4: + full numel = 12 + lp_frag + numel = 4 + start = 0 + end = 3 + hp_frag + numel = 4 + start = 16 + end = 19 + + + +Visual depiction of above +lp { } +flat_lp [ ] +flat_hp ( ) + + +flat_lp [ { ( } ) ] + lx hx ly hy + ly-hx + + +lp { } +flat_lp [ ] +flat_hp ( ) + + +flat_lp [ ( { ) } ] + hx lx hy ly + hy-lx + +lp { } +flat_lp [ ] +flat_hp ( ) + + +flat_lp [ ( { } ) ] + hx lx ly hy + ly-lx + +lp -> (lx, hy) +flat_hp -> (hx, hy) +''' diff --git a/lib/python3.12/site-packages/deepspeed/utils/timer.py b/lib/python3.12/site-packages/deepspeed/utils/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..64ae8ac0e5b4f73d3bf36e8fdddc9dd3ad69d760 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/timer.py @@ -0,0 +1,313 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import time +from numpy import mean +from deepspeed.utils.logging import log_dist +from deepspeed.accelerator import get_accelerator + +FORWARD_MICRO_TIMER = 'fwd_microstep' +FORWARD_GLOBAL_TIMER = 'fwd' +BACKWARD_MICRO_TIMER = 'bwd_microstep' +BACKWARD_GLOBAL_TIMER = 'bwd' +BACKWARD_INNER_MICRO_TIMER = 'bwd_inner_microstep' +BACKWARD_INNER_GLOBAL_TIMER = 'bwd_inner' +BACKWARD_REDUCE_MICRO_TIMER = 'bwd_allreduce_microstep' +BACKWARD_REDUCE_GLOBAL_TIMER = 'bwd_allreduce' +STEP_MICRO_TIMER = 'step_microstep' +STEP_GLOBAL_TIMER = 'step' +TIME_EPSILON = 1e-6 + +try: + import psutil + + PSUTILS_INSTALLED = True +except ImportError: + PSUTILS_INSTALLED = False + pass + + +class CudaEventTimer(object): + + def __init__(self, start_event: get_accelerator().Event, end_event: get_accelerator().Event): + self.start_event = start_event + self.end_event = end_event + + def get_elapsed_msec(self): + get_accelerator().current_stream().wait_event(self.end_event) + self.end_event.synchronize() + return self.start_event.elapsed_time(self.end_event) + + +class SynchronizedWallClockTimer: + """Group of timers. Borrowed from Nvidia Megatron code""" + + class Timer: + """Timer.""" + + def __init__(self, name): + self.name_ = name + self.started_ = False + self.event_timers = [] + self.use_host_timer = get_accelerator().use_host_timers() + self.start_event = None + self.elapsed_records = None + self.start_time = 0.0 + self.end_time = 0.0 + + def start(self): + """Start the timer.""" + assert not self.started_, f"{self.name_} timer has already been started" + if self.use_host_timer: + self.start_time = time.time() + else: + event_class = get_accelerator().Event + self.start_event = event_class(enable_timing=True) + self.start_event.record() + self.started_ = True + + def stop(self, reset=False, record=False): + """Stop the timer.""" + assert self.started_, "timer is not started" + event_class = get_accelerator().Event + if self.use_host_timer: + self.end_time = time.time() + self.event_timers.append(self.end_time - self.start_time) + else: + event_class = get_accelerator().Event + end_event = event_class(enable_timing=True) + end_event.record() + self.event_timers.append(CudaEventTimer(self.start_event, end_event)) + self.start_event = None + self.started_ = False + + def _get_elapsed_msec(self): + if self.use_host_timer: + self.elapsed_records = [et * 1000.0 for et in self.event_timers] + else: + self.elapsed_records = [et.get_elapsed_msec() for et in self.event_timers] + self.event_timers.clear() + return sum(self.elapsed_records) + + def reset(self): + """Reset timer.""" + self.started_ = False + self.start_event = None + self.elapsed_records = None + self.event_timers.clear() + + def elapsed(self, reset=True): + """Calculate the elapsed time.""" + started_ = self.started_ + # If the timing in progress, end it first. + if self.started_: + self.stop() + # Get the elapsed time. + elapsed_ = self._get_elapsed_msec() + # Reset the elapsed time + if reset: + self.reset() + # If timing was in progress, set it back. + if started_: + self.start() + return elapsed_ + + def mean(self): + self.elapsed(reset=False) + return trim_mean(self.elapsed_records, 0.1) + + def __init__(self): + self.timers = {} + + def get_timers(self): + return self.timers + + def __call__(self, name): + if name not in self.timers: + self.timers[name] = self.Timer(name) + return self.timers[name] + + @staticmethod + def memory_usage(): + alloc = "mem_allocated: {:.4f} GB".format(get_accelerator().memory_allocated() / (1024 * 1024 * 1024)) + max_alloc = "max_mem_allocated: {:.4f} GB".format(get_accelerator().max_memory_allocated() / + (1024 * 1024 * 1024)) + cache = "cache_allocated: {:.4f} GB".format(get_accelerator().memory_cached() / (1024 * 1024 * 1024)) + max_cache = "max_cache_allocated: {:.4f} GB".format(get_accelerator().max_memory_cached() / + (1024 * 1024 * 1024)) + return " | {} | {} | {} | {}".format(alloc, max_alloc, cache, max_cache) + + def log(self, names, normalizer=1.0, reset=True, memory_breakdown=False, ranks=None): + """Log a group of timers.""" + assert normalizer > 0.0 + string = f"time (ms)" + for name in names: + if name in self.timers: + elapsed_time = (self.timers[name].elapsed(reset=reset) / normalizer) + string += " | {}: {:.2f}".format(name, elapsed_time) + + log_dist(string, ranks=ranks or [0]) + + def get_mean(self, names, normalizer=1.0, reset=True): + """Get the mean of a group of timers.""" + assert normalizer > 0.0 + means = {} + for name in names: + if name in self.timers: + elapsed_time = (self.timers[name].mean() * 1000.0 / normalizer) + means[name] = elapsed_time + return means + + +class NoopTimer: + + class Timer: + + def start(self): + ... + + def reset(self): + ... + + def stop(self, **kwargs): + ... + + def elapsed(self, **kwargs): + return 0 + + def mean(self): + return 0 + + def __init__(self): + self.timer = self.Timer() + + def __call__(self, name): + return self.timer + + def get_timers(self): + return {} + + def log(self, names, normalizer=1.0, reset=True, memory_breakdown=False, ranks=None): + ... + + def get_mean(self, names, normalizer=1.0, reset=True): + ... + + +class ThroughputTimer: + + def __init__(self, config, batch_size, start_step=2, steps_per_output=None, monitor_memory=False, logging_fn=None): + from deepspeed.utils import logger + self.config = config + self.start_time = 0 + self.end_time = 0 + self.started = False + self.batch_size = 1 if batch_size is None else batch_size + self.start_step = start_step + self.epoch_count = 0 + self.micro_step_count = 0 + self.global_step_count = 0 + self.total_elapsed_time = 0 + self.step_elapsed_time = 0 + self.steps_per_output = steps_per_output + self.monitor_memory = monitor_memory + self.logging = logging_fn + if self.logging is None: + self.logging = logger.info + self.initialized = False + + if self.monitor_memory and not PSUTILS_INSTALLED: + raise ImportError("Unable to import 'psutils', please install package") + + def update_epoch_count(self): + self.epoch_count += 1 + self.micro_step_count = 0 + + def _init_timer(self): + self.initialized = True + + def start(self): + if not self.config.enabled: + return + self._init_timer() + self.started = True + if self.global_step_count >= self.start_step: + if self.config.synchronized: + get_accelerator().synchronize() + self.start_time = time.time() + + def _is_report_boundary(self): + if self.steps_per_output is None: + return False + return self.global_step_count % self.steps_per_output == 0 + + def stop(self, global_step=False, report_speed=True): + if not self.config.enabled or not self.started: + return + self.started = False + self.micro_step_count += 1 + if global_step: + self.global_step_count += 1 + + if self.start_time > 0: + if self.config.synchronized: + get_accelerator().synchronize() + self.end_time = time.time() + duration = self.end_time - self.start_time + self.total_elapsed_time += duration + self.step_elapsed_time += duration + + if global_step: + if report_speed and self._is_report_boundary(): + self.logging( + "epoch={}/micro_step={}/global_step={}, RunningAvgSamplesPerSec={}, CurrSamplesPerSec={}, " + "MemAllocated={}GB, MaxMemAllocated={}GB".format( + self.epoch_count, + self.micro_step_count, + self.global_step_count, + self.avg_samples_per_sec(), + self.batch_size / (self.step_elapsed_time + TIME_EPSILON), + round(get_accelerator().memory_allocated() / 1024**3, 2), + round(get_accelerator().max_memory_allocated() / 1024**3, 2), + )) + if self.monitor_memory: + virt_mem = psutil.virtual_memory() + swap = psutil.swap_memory() + self.logging("epoch={}/micro_step={}/global_step={}, vm %: {}, swap %: {}".format( + self.epoch_count, + self.micro_step_count, + self.global_step_count, + virt_mem.percent, + swap.percent, + )) + self.step_elapsed_time = 0 + + def avg_samples_per_sec(self): + if self.global_step_count > 0: + total_step_offset = self.global_step_count - self.start_step + avg_time_per_step = self.total_elapsed_time / total_step_offset + # training samples per second + return self.batch_size / avg_time_per_step + return float("-inf") + + +def trim_mean(data, trim_percent): + """Compute the trimmed mean of a list of numbers. + + Args: + data (list): List of numbers. + trim_percent (float): Percentage of data to trim. + + Returns: + float: Trimmed mean. + """ + assert 0.0 <= trim_percent <= 1.0 + n = len(data) + # Account for edge case of empty list + if len(data) == 0: + return 0 + data.sort() + k = int(round(n * (trim_percent))) + return mean(data[k:n - k]) diff --git a/lib/python3.12/site-packages/deepspeed/utils/torch.py b/lib/python3.12/site-packages/deepspeed/utils/torch.py new file mode 100644 index 0000000000000000000000000000000000000000..1d32775fe64a8e9bf1bff5df6aaa111c2974d53d --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/torch.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from packaging import version as pkg_version + +import torch + + +def required_torch_version(min_version=None, max_version=None): + assert min_version or max_version, "Must provide a min_version or max_version argument" + + torch_version = pkg_version.parse(torch.__version__) + + if min_version and pkg_version.parse(str(min_version)) > torch_version: + return False + + if max_version and pkg_version.parse(str(max_version)) < torch_version: + return False + + return True + + +def register_grad_hook(param, hook): + if required_torch_version(min_version=2.1): + return param.register_post_accumulate_grad_hook(hook) + else: + param_tmp = param.expand_as(param) + grad_acc = param_tmp.grad_fn.next_functions[0][0] + return grad_acc.register_hook(hook) diff --git a/lib/python3.12/site-packages/deepspeed/utils/types.py b/lib/python3.12/site-packages/deepspeed/utils/types.py new file mode 100644 index 0000000000000000000000000000000000000000..96b5df625965401f52f6fd9b60396821943fe177 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/types.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from enum import IntEnum + + +class ActivationFuncType(IntEnum): + UNKNOWN = 0 + GELU = 1 + ReLU = 2 + GATED_GELU = 3 + GATED_SILU = 4 + + +GATED_ACTIVATION_TYPES = [ + ActivationFuncType.GATED_GELU, + ActivationFuncType.GATED_SILU, +] + + +class NormType(IntEnum): + UNKNOWN = 0 + LayerNorm = 1 + GroupNorm = 2 + RMSNorm = 3 diff --git a/lib/python3.12/site-packages/deepspeed/utils/z3_leaf_module.py b/lib/python3.12/site-packages/deepspeed/utils/z3_leaf_module.py new file mode 100644 index 0000000000000000000000000000000000000000..14e8ae2d28235c46016b9030b3492a2eb4e7122a --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/z3_leaf_module.py @@ -0,0 +1,93 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from typing import List, Type, Union + + +def z3_leaf_module(model: torch.nn.Module) -> bool: + """Returns whether a module in `model` has been flagged as a 'leaf' module. + See `set_z3_leaf_modules` for more details. + Args: + model (torch.nn.Module): The model to which the leaf module flag will be applied. + Returns: + bool: Whether the module has been flagged as a 'leaf' module. + """ + return hasattr(model, '_z3_leaf') and model._z3_leaf + + +def z3_leaf_parameter(model: torch.nn.Parameter) -> bool: + """Returns whether a parameter belongs to a leaf module. + See `set_z3_leaf_modules` for more details. + Args: + model (torch.nn.Parameter): The parameter to which the leaf module flag will be applied. + Returns: + bool: Whether the parameter belongs to a leaf module. + """ + return hasattr(model, 'ds_z3_leaf_module') + + +def get_z3_leaf_modules(model: torch.nn.Module) -> List[torch.nn.Module]: + """Returns a list of modules in `model` that have been flagged as 'leaf' modules. + See `set_z3_leaf_modules` for more details. + Args: + model (torch.nn.Module): The model to which the leaf module flag will be applied. + Returns: + List[torch.nn.Module]: A list of modules that have been flagged as 'leaf' modules. + """ + return [module for module in model.modules() if z3_leaf_module(module)] + + +def set_z3_leaf_module(model: torch.nn.Module, flag: bool): + model._z3_leaf = flag + + +def _do_set_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: Union[List[Type], List[str]], + flag: bool) -> List[torch.nn.Module]: + assert all(isinstance(module_class, (type, str) ) for module_class in leaf_module_classes), \ + f'leaf_module_classes must be a list of types or names, got {leaf_module_classes}' + + leaf_modules = [] + + def _set_z3_leaf_flag(model: torch.nn.Module): + nonlocal leaf_modules + for module in leaf_module_classes: + if (isinstance(module, type) and model.__class__ == module) or \ + (isinstance(module, str) and model.__class__.__name__ == module): + model._z3_leaf = flag + leaf_modules.append(model) + + model.apply(_set_z3_leaf_flag) + + if len(leaf_modules) == 0: + raise ValueError(f'No modules of type {leaf_module_classes} found in model {model}') + + return leaf_modules + + +def set_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: Union[List[Type], + List[str]]) -> List[torch.nn.Module]: + """Sets a flag within a module in `model` to instruct ZeRO3 to stop setting hooks recursively when it encounters a module class listed in `leaf_module_classes`. + This is particularly useful in the context of Mixture of Experts (MoE) models. In MoE models, the computation order of experts varies across forward passes. This variability can disrupt ZeRO3's functionality, as ZeRO3 relies on tracking the computation order of modules to prefetch parameters efficiently. By designating a module as a 'leaf' node, ZeRO3 will prefetch parameters for all child modules upon entering the module. + Another scenario where this functionality is beneficial is in models with excessively fine-grained nested modules, where it helps to avoid the overhead associated with hooks. + Args: + model (torch.nn.Module): The model to which the leaf module flag will be applied. + leaf_module_classes (Union[List[Type], List[str]]): A list of module classes that should be flagged as 'leaf' modules. + Returns: + List[torch.nn.Module]: A list of modules that match the module classes in `leaf_module_classes`. + """ + return _do_set_z3_leaf_modules(model, leaf_module_classes, True) + + +def unset_z3_leaf_modules(model: torch.nn.Module, leaf_module_classes: List[Type]) -> List[torch.nn.Module]: + """Unsets a flag within a module in `model` to instruct ZeRO3 to resume setting hooks recursively when it encounters a module class listed in `leaf_module_classes`. + See `set_z3_leaf_modules` for more details. + Args: + model (torch.nn.Module): The model to which the leaf module flag will be applied. + leaf_module_classes (Union[List[Type], List[str]]): A list of module classes that should be flagged as 'leaf' modules. + Returns: + List[torch.nn.Module]: A list of modules that match the module classes in `leaf_module_classes`. + """ + return _do_set_z3_leaf_modules(model, leaf_module_classes, False) diff --git a/lib/python3.12/site-packages/deepspeed/utils/zero_to_fp32.py b/lib/python3.12/site-packages/deepspeed/utils/zero_to_fp32.py new file mode 100644 index 0000000000000000000000000000000000000000..0e759146cadd92ddfefab3680146c2bd6a2b5c04 --- /dev/null +++ b/lib/python3.12/site-packages/deepspeed/utils/zero_to_fp32.py @@ -0,0 +1,760 @@ +#!/usr/bin/env python + +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets +# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in +# the future. Once extracted, the weights don't require DeepSpeed and can be used in any +# application. +# +# example: +# python zero_to_fp32.py . output_dir/ +# or +# python zero_to_fp32.py . output_dir/ --safe_serialization + +import argparse +import torch +import glob +import math +import os +import re +import gc +import json +import numpy as np +from tqdm import tqdm +from collections import OrderedDict +from dataclasses import dataclass + +# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with +# DeepSpeed data structures it has to be available in the current python environment. +from deepspeed.utils import logger +from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS, + FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES, + FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS) + + +@dataclass +class zero_model_state: + buffers: dict() + param_shapes: dict() + shared_params: list + ds_version: int + frozen_param_shapes: dict() + frozen_param_fragments: dict() + + +debug = 0 + +# load to cpu +device = torch.device('cpu') + + +def atoi(text): + return int(text) if text.isdigit() else text + + +def natural_keys(text): + ''' + alist.sort(key=natural_keys) sorts in human order + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + ''' + return [atoi(c) for c in re.split(r'(\d+)', text)] + + +def get_model_state_file(checkpoint_dir, zero_stage): + if not os.path.isdir(checkpoint_dir): + raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") + + # there should be only one file + if zero_stage <= 2: + file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt") + elif zero_stage == 3: + file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt") + + if not os.path.exists(file): + raise FileNotFoundError(f"can't find model states file at '{file}'") + + return file + + +def get_checkpoint_files(checkpoint_dir, glob_pattern): + # XXX: need to test that this simple glob rule works for multi-node setup too + ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys) + + if len(ckpt_files) == 0: + raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'") + + return ckpt_files + + +def get_optim_files(checkpoint_dir): + return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt") + + +def get_model_state_files(checkpoint_dir): + return get_checkpoint_files(checkpoint_dir, "*_model_states.pt") + + +def parse_model_states(files): + zero_model_states = [] + for file in files: + state_dict = torch.load(file, map_location=device, weights_only=False) + + if BUFFER_NAMES not in state_dict: + raise ValueError(f"{file} is not a model state checkpoint") + buffer_names = state_dict[BUFFER_NAMES] + if debug: + print("Found buffers:", buffer_names) + + # recover just the buffers while restoring them to fp32 if they were saved in fp16 + buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names} + param_shapes = state_dict[PARAM_SHAPES] + + # collect parameters that are included in param_shapes + param_names = [] + for s in param_shapes: + for name in s.keys(): + param_names.append(name) + + # update with frozen parameters + frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None) + if frozen_param_shapes is not None: + if debug: + print(f"Found frozen_param_shapes: {frozen_param_shapes}") + param_names += list(frozen_param_shapes.keys()) + + # handle shared params + shared_params = [[k, v] for k, v in state_dict["shared_params"].items()] + + ds_version = state_dict.get(DS_VERSION, None) + + frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None) + + z_model_state = zero_model_state(buffers=buffers, + param_shapes=param_shapes, + shared_params=shared_params, + ds_version=ds_version, + frozen_param_shapes=frozen_param_shapes, + frozen_param_fragments=frozen_param_fragments) + zero_model_states.append(z_model_state) + + return zero_model_states + + +def parse_optim_states(files, ds_checkpoint_dir): + total_files = len(files) + state_dicts = [] + for f in tqdm(files, desc='Loading checkpoint shards'): + state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False) + # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights + # and also handle the case where it was already removed by another helper script + state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None) + state_dicts.append(state_dict) + + if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]: + raise ValueError(f"{files[0]} is not a zero checkpoint") + zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE] + world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT] + + # For ZeRO-2 each param group can have different partition_count as data parallelism for expert + # parameters can be different from data parallelism for non-expert parameters. So we can just + # use the max of the partition_count to get the dp world_size. + + if type(world_size) is list: + world_size = max(world_size) + + if world_size != total_files: + raise ValueError( + f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. " + "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes." + ) + + # the groups are named differently in each stage + if zero_stage <= 2: + fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS + elif zero_stage == 3: + fp32_groups_key = FP32_FLAT_GROUPS + else: + raise ValueError(f"unknown zero stage {zero_stage}") + + fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))] + return zero_stage, world_size, fp32_flat_groups + + +def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters): + """ + Returns fp32 state_dict reconstructed from ds checkpoint + + Args: + - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are) + + """ + print(f"Processing zero checkpoint '{ds_checkpoint_dir}'") + + optim_files = get_optim_files(ds_checkpoint_dir) + zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir) + print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}") + + model_files = get_model_state_files(ds_checkpoint_dir) + + zero_model_states = parse_model_states(model_files) + print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}') + + if zero_stage <= 2: + return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states, + exclude_frozen_parameters) + elif zero_stage == 3: + return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states, + exclude_frozen_parameters) + + +def _zero2_merge_frozen_params(state_dict, zero_model_states): + if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0: + return + + frozen_param_shapes = zero_model_states[0].frozen_param_shapes + frozen_param_fragments = zero_model_states[0].frozen_param_fragments + + if debug: + num_elem = sum(s.numel() for s in frozen_param_shapes.values()) + print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}') + + wanted_params = len(frozen_param_shapes) + wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) + avail_numel = sum([p.numel() for p in frozen_param_fragments.values()]) + print(f'Frozen params: Have {avail_numel} numels to process.') + print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params') + + total_params = 0 + total_numel = 0 + for name, shape in frozen_param_shapes.items(): + total_params += 1 + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + + state_dict[name] = frozen_param_fragments[name] + + if debug: + print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ") + + print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") + + +def _has_callable(obj, fn): + attr = getattr(obj, fn, None) + return callable(attr) + + +def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states): + param_shapes = zero_model_states[0].param_shapes + + # Reconstruction protocol: + # + # XXX: document this + + if debug: + for i in range(world_size): + for j in range(len(fp32_flat_groups[0])): + print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}") + + # XXX: memory usage doubles here (zero2) + num_param_groups = len(fp32_flat_groups[0]) + merged_single_partition_of_fp32_groups = [] + for i in range(num_param_groups): + merged_partitions = [sd[i] for sd in fp32_flat_groups] + full_single_fp32_vector = torch.cat(merged_partitions, 0) + merged_single_partition_of_fp32_groups.append(full_single_fp32_vector) + avail_numel = sum( + [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups]) + + if debug: + wanted_params = sum([len(shapes) for shapes in param_shapes]) + wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes]) + # not asserting if there is a mismatch due to possible padding + print(f"Have {avail_numel} numels to process.") + print(f"Need {wanted_numel} numels in {wanted_params} params.") + + # params + # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support + # out-of-core computing solution + total_numel = 0 + total_params = 0 + for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups): + offset = 0 + avail_numel = full_single_fp32_vector.numel() + for name, shape in shapes.items(): + + unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape) + total_numel += unpartitioned_numel + total_params += 1 + + if debug: + print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ") + state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape) + offset += unpartitioned_numel + + # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and + # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex + # paddings performed in the code it's almost impossible to predict the exact numbers w/o the + # live optimizer object, so we are checking that the numbers are within the right range + align_to = 2 * world_size + + def zero2_align(x): + return align_to * math.ceil(x / align_to) + + if debug: + print(f"original offset={offset}, avail_numel={avail_numel}") + + offset = zero2_align(offset) + avail_numel = zero2_align(avail_numel) + + if debug: + print(f"aligned offset={offset}, avail_numel={avail_numel}") + + # Sanity check + if offset != avail_numel: + raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") + + print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements") + + +def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states, + exclude_frozen_parameters): + state_dict = OrderedDict() + + # buffers + buffers = zero_model_states[0].buffers + state_dict.update(buffers) + if debug: + print(f"added {len(buffers)} buffers") + + if not exclude_frozen_parameters: + _zero2_merge_frozen_params(state_dict, zero_model_states) + + _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) + + # recover shared parameters + for pair in zero_model_states[0].shared_params: + if pair[1] in state_dict: + state_dict[pair[0]] = state_dict[pair[1]] + + return state_dict + + +def zero3_partitioned_param_info(unpartitioned_numel, world_size): + remainder = unpartitioned_numel % world_size + padding_numel = (world_size - remainder) if remainder else 0 + partitioned_numel = math.ceil(unpartitioned_numel / world_size) + return partitioned_numel, padding_numel + + +def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states): + if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0: + return + + if debug: + for i in range(world_size): + num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values()) + print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}') + + frozen_param_shapes = zero_model_states[0].frozen_param_shapes + wanted_params = len(frozen_param_shapes) + wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) + avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size + print(f'Frozen params: Have {avail_numel} numels to process.') + print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params') + + total_params = 0 + total_numel = 0 + for name, shape in zero_model_states[0].frozen_param_shapes.items(): + total_params += 1 + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + + param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states) + state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape) + + partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size) + + if debug: + print( + f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" + ) + + print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") + + +class GatheredTensor: + """ + A pseudo tensor that collects partitioned weights. + It is more memory efficient when there are multiple groups. + """ + + def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape): + self.flat_groups = flat_groups + self.flat_groups_offset = flat_groups_offset + self.offset = offset + self.partitioned_numel = partitioned_numel + self.shape = shape + self.dtype = self.flat_groups[0][0].dtype + + def contiguous(self): + """ + Merge partitioned weights from flat_groups into a single tensor. + """ + end_idx = self.offset + self.partitioned_numel + world_size = len(self.flat_groups) + pad_flat_param_chunks = [] + + for rank_i in range(world_size): + # for each rank, we need to collect weights from related group/groups + flat_groups_at_rank_i = self.flat_groups[rank_i] + start_group_id = None + end_group_id = None + for group_id in range(len(self.flat_groups_offset)): + if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]: + start_group_id = group_id + if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]: + end_group_id = group_id + break + # collect weights from related group/groups + for group_id in range(start_group_id, end_group_id + 1): + flat_tensor = flat_groups_at_rank_i[group_id] + start_offset = self.offset - self.flat_groups_offset[group_id] + end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id] + pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset]) + + # collect weights from all ranks + pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0) + param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous() + return param + + +def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states): + param_shapes = zero_model_states[0].param_shapes + avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size + + # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each + # param, re-consolidating each param, while dealing with padding if any + + # merge list of dicts, preserving order + param_shapes = {k: v for d in param_shapes for k, v in d.items()} + + if debug: + for i in range(world_size): + print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}") + + wanted_params = len(param_shapes) + wanted_numel = sum(shape.numel() for shape in param_shapes.values()) + # not asserting if there is a mismatch due to possible padding + avail_numel = fp32_flat_groups[0].numel() * world_size + print(f"Trainable params: Have {avail_numel} numels to process.") + print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.") + + # params + # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support + # out-of-core computing solution + offset = 0 + total_numel = 0 + total_params = 0 + flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]])) + for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'): + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + total_params += 1 + partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size) + + if debug: + print( + f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" + ) + + # memory efficient tensor + tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape) + state_dict[name] = tensor + offset += partitioned_numel + + offset *= world_size + + # Sanity check + if offset != avail_numel: + raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") + + print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements") + + +def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states, + exclude_frozen_parameters): + state_dict = OrderedDict() + + # buffers + buffers = zero_model_states[0].buffers + state_dict.update(buffers) + if debug: + print(f"added {len(buffers)} buffers") + + if not exclude_frozen_parameters: + _zero3_merge_frozen_params(state_dict, world_size, zero_model_states) + + _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) + + # recover shared parameters + for pair in zero_model_states[0].shared_params: + if pair[1] in state_dict: + state_dict[pair[0]] = state_dict[pair[1]] + + return state_dict + + +def to_torch_tensor(state_dict, return_empty_tensor=False): + """ + Convert state_dict of GatheredTensor to torch tensor + """ + torch_state_dict = {} + converted_tensors = {} + for name, tensor in state_dict.items(): + tensor_id = id(tensor) + if tensor_id in converted_tensors: # shared tensors + shared_tensor = torch_state_dict[converted_tensors[tensor_id]] + torch_state_dict[name] = shared_tensor + else: + converted_tensors[tensor_id] = name + if return_empty_tensor: + torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype) + else: + torch_state_dict[name] = tensor.contiguous() + return torch_state_dict + + +def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, + tag=None, + exclude_frozen_parameters=False, + lazy_mode=False): + """ + Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with + ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example + via a model hub. + + Args: + - ``checkpoint_dir``: path to the desired checkpoint folder + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14`` + - ``exclude_frozen_parameters``: exclude frozen parameters + - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient. + Convert the pesduo tensor to torch tensor by ``.contiguous()`` + + Returns: + - pytorch ``state_dict`` + + A typical usage might be :: + + from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint + # do the training and checkpoint saving + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu + model = model.cpu() # move to cpu + model.load_state_dict(state_dict) + # submit to model hub or save the model to share with others + + In this example the ``model`` will no longer be usable in the deepspeed context of the same + application. i.e. you will need to re-initialize the deepspeed engine, since + ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it. + + If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead. + + Note: the above usage may not work if your application doesn't have sufficient free CPU memory. + You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with + the checkpoint. Or you can load state_dict in lazy mode :: + + from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu + for name, lazy_tensor in state_dict.item(): + tensor = lazy_tensor.contiguous() # to cpu + print(name, tensor) + # del tensor to release memory if it no longer in use + """ + if tag is None: + latest_path = os.path.join(checkpoint_dir, 'latest') + if os.path.isfile(latest_path): + with open(latest_path, 'r') as fd: + tag = fd.read().strip() + else: + raise ValueError(f"Unable to find 'latest' file at {latest_path}") + + ds_checkpoint_dir = os.path.join(checkpoint_dir, tag) + + if not os.path.isdir(ds_checkpoint_dir): + raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist") + + state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters) + if lazy_mode: + return state_dict + else: + return to_torch_tensor(state_dict) + + +def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, + output_dir, + max_shard_size="5GB", + safe_serialization=False, + tag=None, + exclude_frozen_parameters=False): + """ + Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be + loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed. + + Args: + - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``) + - ``output_dir``: directory to the pytorch fp32 state_dict output files + - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB + - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14`` + - ``exclude_frozen_parameters``: exclude frozen parameters + """ + + # Dependency pre-check + if safe_serialization: + try: + from safetensors.torch import save_file + except ImportError: + print('If you want to use `safe_serialization`, please `pip install safetensors`') + raise + if max_shard_size is not None: + try: + from huggingface_hub import split_torch_state_dict_into_shards + except ImportError: + print('If you want to use `max_shard_size`, please `pip install huggingface_hub`') + raise + + # Convert zero checkpoint to state_dict + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, + tag, + exclude_frozen_parameters, + lazy_mode=True) + + # Shard the model if it is too big. + weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin" + if max_shard_size is not None: + filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors") + # an memory-efficient approach for sharding + empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True) + state_dict_split = split_torch_state_dict_into_shards(empty_state_dict, + filename_pattern=filename_pattern, + max_shard_size=max_shard_size) + else: + from collections import namedtuple + StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"]) + state_dict_split = StateDictSplit(is_sharded=False, + filename_to_tensors={weights_name: list(state_dict.keys())}) + + # Save the model by shard + os.makedirs(output_dir, exist_ok=True) + filename_to_tensors = state_dict_split.filename_to_tensors.items() + for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"): + shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors} + shard_state_dict = to_torch_tensor(shard_state_dict) + output_path = os.path.join(output_dir, shard_file) + if safe_serialization: + save_file(shard_state_dict, output_path, metadata={"format": "pt"}) + else: + torch.save(shard_state_dict, output_path) + # release the memory of current shard + for tensor_name in list(shard_state_dict.keys()): + del state_dict[tensor_name] + del shard_state_dict[tensor_name] + del shard_state_dict + gc.collect() + + # Save index if sharded + if state_dict_split.is_sharded: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json" + save_index_file = os.path.join(output_dir, save_index_file) + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + + +def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None): + """ + 1. Put the provided model to cpu + 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` + 3. Load it into the provided model + + Args: + - ``model``: the model object to update + - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``) + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14`` + + Returns: + - ``model`: modified model + + Make sure you have plenty of CPU memory available before you call this function. If you don't + have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it + conveniently placed for you in the checkpoint folder. + + A typical usage might be :: + + from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint + model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir) + # submit to model hub or save the model to share with others + + Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context + of the same application. i.e. you will need to re-initialize the deepspeed engine, since + ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it. + + """ + logger.info(f"Extracting fp32 weights") + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag) + + logger.info(f"Overwriting model with fp32 weights") + model = model.cpu() + model.load_state_dict(state_dict, strict=False) + + return model + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint_dir", + type=str, + help="path to the desired checkpoint folder, e.g., path/checkpoint-12") + parser.add_argument("output_dir", + type=str, + help="directory to the pytorch fp32 state_dict output files" + "(e.g. path/checkpoint-12-output/)") + parser.add_argument( + "--max_shard_size", + type=str, + default="5GB", + help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size" + "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`" + "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances" + "without CPU OOM issues.") + parser.add_argument( + "--safe_serialization", + default=False, + action='store_true', + help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).") + parser.add_argument("-t", + "--tag", + type=str, + default=None, + help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1") + parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters") + parser.add_argument("-d", "--debug", action='store_true', help="enable debug") + args = parser.parse_args() + + debug = args.debug + + convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, + args.output_dir, + max_shard_size=args.max_shard_size, + safe_serialization=args.safe_serialization, + tag=args.tag, + exclude_frozen_parameters=args.exclude_frozen_parameters)