6.5 KiB
Torch-MLU-Ops
目录
简介
Torch-MLU-Ops是寒武纪设计和开发的PyTorch第三方算子库。对于使用PyTorch框架的开发者，通过Torch-MLU-Ops，能够便捷地使用这些自定义算子，进行算子的集成、评测和业务部署。
Torch-MLU-Ops已全量覆盖LLM（Large Language Model）推理场景下的常见算子。作为后端，已支持Cambricon vLLM、Cambricon TGI、Cambricon Stable Diffusion web UI、Cambricon ComfyUI以及Cambricon Diffusers。
安装
环境要求
Torch-MLU-Ops支持的操作系统以及软件依赖如下。
不同平台支持的软件版本如下：
| 操作系统 | 编译器版本 | CMake版本 |
|---|---|---|
| Ubuntu22.04 x86-64 | gcc 9.4.0或者更高版本 | 3.10或者更高版本 |
| Debian10.11 x86-64 | gcc 9.4.0或者更高版本 | 3.10或者更高版本 |
编译、运行Torch-MLU-Ops同时还需要配置以下依赖。
| Torch-MLU-Ops | Torch-MLU | CNNL | CNNL_Extra | CNCL | CNToolkit |
|---|---|---|---|---|---|
| v1.3.2 | v1.24.1 | v1.28.3 | v1.12.3 | v1.24.1 | v3.15.6 |
| v1.3.1 | v1.24.1 | v1.28.3 | v1.12.2 | v1.24.1 | v3.15.5 |
| v1.3.0 | v1.24.0 | v1.28.2 | v1.12.1 | v1.24.0 | v3.15.4 |
| v1.2.3 | v1.24.0 | v1.28.1 | v1.12.1 | v1.24.0 | v3.15.3 |
| v1.2.2 | v1.23.1 | v1.27.4 | v1.11.3 | v1.23.0 | v3.14.3 |
| v1.2.1 | v1.23.1 | v1.27.3 | v1.11.3 | v1.22.3 | v3.14.2 |
| v1.2.0 | v1.23.0 | v1.27.1 | v1.11.1 | v1.22.1 | v3.14.1 |
| v1.1.4 | v1.22.2 | v1.26.7 | v1.10.4 | v1.21.1 | v3.13.7 |
| v1.1.3 | v1.22.2 | v1.26.6 | v1.10.3 | v1.21.1 | v3.13.5 |
此外运行Torch-MLU-Ops还依赖PyTorch环境。Torch-MLU-Ops已支持的PyTorch版本请参考。
| Torch-MLU-Ops | PyTorch |
|---|---|
| v1.3.2 | v2.1, v2.4, v2.5 |
| v1.3.1 | v2.1, v2.4, v2.5 |
| v1.3.0 | v2.1, v2.4, v2.5 |
| v1.2.3 | v2.1, v2.4, v2.5 |
| v1.2.2 | v2.1, v2.4, v2.5 |
| v1.2.1 | v2.1, v2.4, v2.5 |
| v1.2.0 | v2.1, v2.4, v2.5 |
| v1.1.4 | v2.1, v2.3, v2.4 |
| v1.1.3 | v2.1, v2.3 |
通过docker镜像启动
寒武纪提供的解决方案镜像已提供了Torch-MLU-Ops所需要的依赖，请参考Cambricon PyTorch Container Image使用。您可以从寒武纪开发者社区获取该镜像。
编译安装
Torch-MLU-Ops编译脚本依赖于NEUWARE_HOME环境变量，该环境变量指向寒武纪CNToolkit的安装路径，通常为/usr/local/neuware。
在运行示例程序前，您需要将NEUWARE_HOME中的lib64目录添加到LD_LIBRARY_PATH环境变量中，例如：
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$NEUWARE_HOME/lib64
如果您使用寒武纪提供的docker镜像，以上环境变量均已经被设置好了。
设置好NEUWARE_HOME环境变量后，您可以使用以下命令编译Torch-MLU-Ops。
从远端clone仓库，进入项目的主目录：
cd torch_mlu_ops
源码安装：
pip install -e .
wheel包安装：
python setup.py bdist_wheel
pip install dist/torch_mlu_ops*.whl
注意，如果安装过程中遇到权限问题，请获取相关权限或切换到拥有权限的user执行pip install 这步操作。
移除Torch-MLU-Ops（Optional）：
pip uninstall torch_mlu_ops
测试
在测试前，请确保Torch-MLU-Ops已安装，您可以通过 pip list 命令查询Torch-MLU-Ops的安装情况。若有如下类似打印，则表明安装成功。
torch-mlu-ops 1.1.0+pt21
Torch-MLU-Ops模块的导入方法请参考。
import torch_mlu_ops as tmo
在./tests/ops_pytest目录下，提供了ops级别的测试示例程序，您可以参考如下命令进行测试。
cd tests/ops_pytest/
./run_test.sh
# or 单独测试某个测例，例如
python test_flash_attention.py
在./tests/kernels_pytest目录下，提供了kernels级别的测试示例程序，您可以参考如下命令进行测试。
cd tests/kernels_pytest
./build.sh # 需要先编译
./run_test.sh
目录结构
torch_mlu_ops
|-- benchmarks ## 性能benchmarks测试代码目录
|-- csrc ## C以及BangC源码目录
| |-- common
| |-- kernels
| |-- ops
| `-- torch_api
|-- docs ## 文档目录
| |-- release_notes
| `-- user_guide
|-- tests ## 测试代码目录，包含了kernel和ops级别的测试
| |-- kernels_pytest
| `-- ops_pytest
|-- tools ## 工具目录，主要包含CI相关代码
| |-- ci
`-- torch_mlu_ops ## PyTorch第三方算子接口实现相关代码
关键特性
Torch-MLU-Ops提供的自定义算子列表以及接口说明请参考Torch-MLU-Ops用户手册。