安装命令
pip install transformer_engine[pytorch]
或
pip install git+https://github.com/NVIDIA/TransformerEngine.git@main
Cmake报错信息解决
如果出现cmake相关CUDA/CUDNN错误:
(eg.缺少cudnn.h)
更改CUDA_PATH和CUDNN_PATH
版本检查CUDA 12.0+ and cuDNN 8.1+.
首先找到环境下合适的cuda 版本,比如12.1
之后
vim ~/.bashrc
export PATH=/usr