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

DeepLink-org / DeepLinkExt

readme