“transformer_engine 老给我装 transformer_engine_cu12” | 个人犯蠢记录

一句话：还是官方文档管用

问题描述

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 环境配置 (minimized version)
pip install torch==2.10.0+cu128 flash-attn==2.8.3
pip install "transformer_engine[pytorch]>=1.4"
# 此处输出中可以注意到安装了 transformer_engine_cu13 而不是跟我的 CUDA 版本对应的 transformer_engine_cu12

# 测试语句
$ python -c "import transformer_engine.pytorch"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/opt/conda/lib/python3.10/site-packages/transformer_engine/__init__.py", line 11, in <module>
    import transformer_engine.common
  File "/opt/conda/lib/python3.10/site-packages/transformer_engine/common/__init__.py", line 382, in <module>
    _TE_LIB_CTYPES = _load_core_library()
  File "/opt/conda/lib/python3.10/site-packages/transformer_engine/common/__init__.py", line 360, in
_load_core_library
    return ctypes.CDLL(_get_shared_object_file("core"), mode=ctypes.RTLD_GLOBAL)
  File "/opt/conda/lib/python3.10/ctypes/__init__.py", line 374, in __init__
    self._handle = _dlopen(self._name, mode)
OSError: libcublas.so.13: cannot open shared object file: No such file or directory
(base)

中间也出现过一些别的错误，不过总之一眼能看出来显然是 cuda 相关的版本没对上导致的

Ask AI

提问 AI 已经成为了一个习惯。这次问 AI 的过程很遗憾没有给我提供任何行之有效的解决方案，不过相关的回复也有一定的价值，遂记录一下

Question Prompt

1
2
3
我使用 pip 安装 "transformer_engine[pytorch]>=1.4" 的时候每次都给我安装 transformer_engine_cu13 的版本，但我需要的是 transformer_engine_cu12 的版本。机器上 nvidia-smi 看到的是 12.8，怎么办

具体报错: [Pasted Error Logs]

Result

一些 AI 给出的显然正确的结论：

  1. 是因为给我 resolve 到 _cu13 导致的，说明在解析包的时候产生了一些 mismatch
  2. 我应该想办法去安装 _cu12 的版本

一些有道理但实际上无关的猜测

  1. 我的 nvcc 版本不对 (nvcc -V)
  2. 我的 torch 版本不对 (import torch; python -c "print(torch.version.cuda)")

一些建议与尝试（实际上也没有什么用处）

  1. 强制使用 core_cu12 版本安装 (pip install "transformer_engine[core_cu12,pytorch]>=1.4")
    效果：还真安装了 _cu12，但是 _cu13 也还是会装，最后运行的时候也还是走 _cu13
  2. 不要用最新版（codex 似乎认为最新版可能有 bug）
    效果：无变化
  3. 设置 CUDA_HOME（export CUDA_HOME=/usr/local/cuda-12.8）
    效果：无变化
  4. 设置 CUDA_VERSION（export CUDA_VERSION=12.1）
    效果：无变化

到这里反正能看出来本质上就是解析包的时候不知道为什么给我本地检测成了 cu13 的环境，所以关键点在于找到检测错的核心配置，然后改成 cu12（？

最后的解决方法

看了眼官方文档: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/installation.html
官方安装指令: pip3 install --no-build-isolation transformer_engine[pytorch]

没错，我们需要 --no-build-isolation。否则环境检测可能会错误的检测成 CUDA 13.
That's Over.