MindSpeed介绍

MindSpeed 是针对华为昇腾设备的大模型加速库。

大模型训练是一种非常复杂的过程，涉及到许多技术和挑战，其中大模型训练需要大量的显存资源是一个难题，对计算卡提出了不小的挑战。 为了在单个计算卡显存资源不足时，可以通过多张计算卡进行计算，业界出现了类似 Megatron、DeepSpeed 等第三方大模型加速库，对模型、输入数据等进行切分并分配到不同的计算卡上，最后再通过集合通信对结果进行汇总。

昇腾提供 MindSpeed 加速库，使能客户大模型业务快速迁移至昇腾设备，并且支持昇腾专有算法，确保开箱可用。

参考链接：https://gitee.com/ascend/MindSpeed/tree/master

MindSpeed目录介绍：
MindSpeed/

├── ci/ # 门禁执行

├── docs/ # 相关文档

│ ├── features/ #特性相关

│ ├── ops/ #算子相关

├── mindspeed/ #特性放置及使能

│ ├── core/ #性能优化相关特性

│ ├── model/ # 数据集相关

│ ├── moe/ #moe相关

│ ├── ops/ # 自定义算子

│ ├── optimizer/ # 优化器相关特性

│ ├── tokenizer/ # 处理tokenizer相关特性

│ ├── arguments.py #特性使能相关参数

│ ├── megatron_adaptor.py #特性调用

├── tests_extend/ #特性相关UT

│ ├── model_tests/ # 模型测试

│ ├── system_tests/ #模型训练系统测试

│ ├── unit_tests/ #特性UT

├── setup.py #编译安装

MindSpeed特性开发流程

这里以特性fused_swiglu为例，介绍开发流程。

特性分析

在megatron中，MLP部分不同场景下使用了不同的激活函数，swiglu是多数场景下使用的激活函数，以下为megatron部分代码。昇腾提供了swiglu的融合算子，现在要将该融合算子作为特性添加到mindspeed上。

class ParallelMLP(MegatronModule):
    """MLP.

    MLP will take the input with h hidden state, project it to 4*h
    hidden dimension, perform nonlinear transformation, and project the
    state back into h hidden dimension.
    """

    def __init__(self, config, is_expert=False):
        super(ParallelMLP, self).__init__()
        args = get_args()

        self.add_bias = config.add_bias_linear

        ffn_hidden_size = config.ffn_hidden_size
        if config.gated_linear_unit:
            ffn_hidden_size *= 2

        # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
        self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear(
            config.hidden_size,
            ffn_hidden_size,
            config=config,
            init_method=config.init_method,
            bias=self.add_bias,
            gather_output=False,
            skip_bias_add=True,
            is_expert=is_expert,
        )

        self.bias_gelu_fusion = False
        self.activation_func = None
        self.swiglu = args.swiglu

        if args.openai_gelu:
            self.activation_func = openai_gelu
        elif args.onnx_safe:
            self.activation_func = erf_gelu
        elif args.swiglu:
            def swiglu(x):
                x = torch.chunk(x, 2, dim=-1)
                return F.silu(x[0]) * x[1]
            self.activation_func = swiglu
        elif args.squared_relu:
            def squared_relu(x):
                return torch.pow(F.relu(x), 2)
            self.activation_func = squared_relu
        else:
            self.bias_gelu_fusion = args.bias_gelu_fusion
            self.activation_func = F.gelu

特性开发

为了不侵入式修改megatron，这里以patch的方式添加fused_swiglu特性。patch内容如下，位置为mindspeed/model/transformer.py

def parallel_mlp_init_wrapper(fn):
    @wraps(fn)
    def wrapper(self, *args, **kwargs):
        fn(self, *args, **kwargs)
        self.layer_number = None
        args = get_args()
        if args.swiglu and args.use_fused_swiglu:
            self.activation_func = fused_swiglu
    return wrapper

其中fusedz_swiglu定义位置为mindspeed/core/fusions/fused_bias_swiglu.py

def fused_swiglu(x):
    return torch_npu.npu_swiglu(x, dim=-1)

在开发完patch后需要提供参数进行控制，代码片段如下，位置为mindspeed/arguments.py

def _add_network_size_args(parser):
    group = parser.add_argument_group(title='network size')
    group.add_argument("--use-fused-rmsnorm", action='store_true',
                       help="Use fused rmsnorm.")
    group.add_argument("--use-fused-swiglu", action='store_true',
                       help="Use fused swiglu.")
    group.add_argument("--use-fused-rotary-pos-emb", action='store_true',
                       help="Use fused rotary-pos-emb.")

特性使能

在完成特性开发后，需要使能特性，这里以调用patch的方式实现，代码片段如下，位置为

def megatron_legacy_adaptation(aspm):
    from .core.fusions.fused_layer_norm import FusedLayerNormAffineFunction, FastLayerNormFN, fused_layer_norm_affine
    from .core.fusions.fused_softmax import is_kernel_available, ScaledUpperTriangMaskedSoftmax, ScaledMaskedSoftmax, \
        ScaledSoftmax, forward_fused_softmax
    from .core.fusions.rms_norm import rms_norm_init_wrapper, rms_norm_forward_wrapper
    from .model.transformer import parallel_mlp_init_wrapper, flash_self_attention_forward, flash_self_attention_init_wrapper, parallel_mlp_forward_wrapper, parallel_transformer_init_wrapper, \
        parallel_transformer_forward_wrapper


aspm.register_patch('megatron.legacy.model.transformer.ParallelMLP.__init__', parallel_mlp_init_wrapper)

特性UT

为了验证，看护该特性，需要提供该特性UT，以下为代码，位置为tests_extend/unit_tests/ops/cann/test_npu_swiglu.py

import pytest
import torch
import torch_npu
from mindspeed import megatron_adaptor
from megatron.legacy.model.transformer import ParallelMLP
from unit_tests.common import DistributedTest
from commons import initialize_model_parallel
from megatron.training.global_vars import set_args
from megatron.training.arguments import parse_args
from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
from megatron.core.transformer.transformer_config import TransformerConfig

DEVICE_NAME = torch_npu.npu.get_device_name(0)[:10]


def create_test_args(swiglu=True, use_fused_swiglu=False):
    args = parse_args(None, True)
    args.swiglu = swiglu
    args.use_fused_swiglu = use_fused_swiglu
    return args


class TestNpuSwiglu(DistributedTest):
    world_size = 1

    def patch_ori_op_exec(self, hidden_states):
        args = create_test_args(True, False)
        set_args(args)
        model_parallel_cuda_manual_seed(123)
        self.patch_ori_mlp = ParallelMLP(self.transformer_config).npu().to(hidden_states.dtype)
        hidden_states, output_bias = self.patch_ori_mlp(hidden_states)
        return hidden_states

    def patch_fused_op_exec(self, hidden_states):
        args = create_test_args(True, True)
        set_args(args)
        model_parallel_cuda_manual_seed(123)
        self.patch_fused_mlp = ParallelMLP(self.transformer_config).npu().to(hidden_states.dtype)
        hidden_states, output_bias = self.patch_fused_mlp(hidden_states)
        return hidden_states

    @pytest.mark.skipif(DEVICE_NAME != 'Ascend910B', reason='device type is not supported, skip this UT!')
    @pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16])
    def test_npu_swiglu(self, dtype):
        args = parse_args(None, True)
        set_args(args)
        sequence_length, batch_size, hidden_size = 32, 2, 12
        initialize_model_parallel(1, 1)
        self.transformer_config = TransformerConfig(num_layers=2, hidden_size=hidden_size, num_attention_heads=4,
                                                    use_cpu_initialization=False, gated_linear_unit=True)
        hidden_states = torch.randn((sequence_length, batch_size, hidden_size)).npu().to(dtype)
        output_patch_ori = self.patch_ori_op_exec(hidden_states)
        output_patch_fused = self.patch_fused_op_exec(hidden_states)
        tol = 0.004 if dtype == torch.bfloat16 else 0.001
        assert torch.allclose(output_patch_ori, output_patch_fused, rtol=tol, atol=tol)

特性上库

完成自验后就可以提交上库了，按照上文中特性开发对应的位置进行上库，由mindspeed的commiter进行代码检视、合入。注意在PR中将自己的自验结果展示。