简介

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

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

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

安装

1. 安装依赖

在安装AscendSpeed之前，请参考版本配套表，安装最新昇腾软件栈：https://www.hiascend.com/zh/。

2. 安装 AscendSpeed

如有旧版本 AscendSpeed 请先卸载，再进行安装操作。

方式一：从 git 直接安装

pip install git+https://gitee.com/ascend/AscendSpeed.git

方式二：下载源码安装

git clone https://gitee.com/ascend/AscendSpeed.git
pip install -e AscendSpeed

3. 获取 Megatron-LM 并指定 commit id

# 目前版本基于core_r0.6.0的release版本
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout core_r0.6.0

快速上手

以 GPT 模型为例：

方式一：若 Megatron-LM 的 commit id 与上述一致，运行ascendspeed命令

1. 在 Megatron-LM 目录下，运行ascendspeed -P命令。

    ascendspeed -P

2. 在 Megatron-LM 目录下，准备好训练数据，并在示例脚本中填写对应路径，然后执行。

   bash examples/pretrain_gpt_distributed.sh

方式二：修改本地代码

1. 在 Megatron-LM 目录下修改pretrain_gpt.py文件，在import torch下新增一行import ascendspeed.megatron_adaptor

    import os
    import torch
   +import ascendspeed.megatron_adaptor
    from torch import Tensor
    from functools import partial
    from typing import Union

2. 在 Megatron-LM 目录下修改pretrain_gpt.py文件，在model_provider函数中删除assert(args.context_parallel_size == 1), "Context parallelism is only supported with Megatron Core!"。

   else:
   -   assert(args.context_parallel_size == 1), "Context parallelism is only supported with Megatron Core!"
   
       model = megatron.model.GPTModel(
           config,
           num_tokentypes=0,
           parallel_output=True,
           pre_process=pre_process,
           post_process=post_process
       )

3. 在 Megatron-LM 目录下，准备好训练数据，并在示例脚本中填写对应路径，然后执行。

   bash examples/pretrain_gpt_distributed.sh

特性介绍

自定义算子

版本配套表

PyTorch Extension版本号采用{PyTorch版本}-{昇腾版本}命名规则，前者为PyTorch Extension匹配的PyTorch版本，后者用于匹配CANN版本，详细匹配如下：

昇腾辅助软件中有更多关于PyTorch和CANN的版本信息。

安全声明

AscendSpeed 安全声明