Skip to content

DeepLink-org/DeepLinkExt

Repository files navigation

DeepLinkExt

该仓库提供了一套在国产硬件上进行大模型训练、微调、推理的解决方案。对下调用DIOPI支持的大模型算子(如flash_attention和rms_norm),对上承接大模型训练、微调、推理框架;并对应提供了对应的大模型算子的torch组合实现。

DIOPI的具体内容请参考DIOPI INTRODUCTION.

目前支持的框架以及对应的算子可以查看框架与算子.

Install

DeepLinkExt依赖deeplink.framework/dipu,需要先完成dipu的编译安装,具体请参考dipu quick_start 完成dipu的编译后,请参考如下代码,设置必要的环境变量。

export PYTHONPATH=$WORKDIR/deeplink.framework/dipu/:$PYTHONPATH
export DIPU_ROOT=$WORKDIR/deeplink.framework/dipu/torch_dipu
export DIOPI_PATH=$WORKDIR/deeplink.framework/dipu/third_party/DIOPI/proto
export VENDOR_INCLUDE_DIRS=${PATH_TO_VENDOR_INCLUDE} # 底层软件栈的include路径,例如/usr/local/Ascend/ascend-toolkit/latest/include

完成上述准备工作后,使用如下命令即可安装DeepLinkExt。提供两种安装实例如下:

1. inplace安装

cd $WORKDIR/DeepLinkExt
python3 setup.py build_ext --inplace

2. 安装到指定目录

cd $WORKDIR/DeepLinkExt
pip install -e . -t $TARGET_INSTALL_DIR

Usage

以InternEvo、LightLLM大模型训练框架为例,参考如下代码,即可实现在训练/推理时使用DeepLinkExt的扩展算子。

InternEvo

DeepLinkExt已完全接入InternEvo,在完成DeepLinkExt的编译安装后,将其添加到PYTHONPATH,使用InternEvo进行训练即可。

LightLLM

对于LightLLM,在启动推理的脚本中,需要添加如下代码,即可实现使用DeepLinkExt扩展算子。

import deeplink_ext.patch_lightllm.py
import lightllm