BMTrain

本项目把 BMTrain 历史上分散在多个分支的硬件后端 合并到 同一个 Python wheel 里。安装时不会做任何针对硬件的编译； 首次 import bmtrain 才会探测当前机器的硬件，对该后端的 csrc/ 目录 跑一次 cmake，把产物 C.*.so 缓存下来。

安装

pip install ./bmtrain_unified 
# 或者打成 wheel
python -m pip wheel ./bmtrain_unified -w dist/

得到的 wheel 是 架构无关 的（只携带源码，不携带二进制）。 也就是说，一个 wheel 就可以分发到所有国产卡 / NVIDIA 集群。

首次 import 会编译原生扩展

# 自动探测（按顺序检查 torch_mlu、MACA_PATH、SDAA_HOME、/usr/local/corex、
# torch.cuda.is_available 等）
python -c "import bmtrain; print(bmtrain.get_backend())"

# 显式指定后端
BMTRAIN_BACKEND=mlu     python -c "import bmtrain; print(bmtrain.__backend__)"
BMTRAIN_BACKEND=tianshu python -c "import bmtrain"
BMTRAIN_BACKEND=muxi    python -c "import bmtrain"
BMTRAIN_BACKEND=sdaa    python -c "import bmtrain"

构建产物缓存在 ~/.cache/bmtrain/<backend>/<py-tag>/， 最终的 .so 会被复制到 bmtrain/_backends/<backend>/__init__.py 旁边，所以第二次 import 完全跳过 cmake。

环境变量

变量	作用
BMTRAIN_BACKEND	强制指定后端（cuda/mlu/muxi/sdaa/tianshu 或别名）
BMTRAIN_BUILD_DIR	缓存根目录（默认 ~/.cache/bmtrain）
BMTRAIN_CMAKE	指定其它的 cmake 二进制
BMTRAIN_VERBOSE	设为 1 时实时输出 cmake / build 日志
BMTRAIN_REBUILD	设为 1 时强制清缓存重新构建
BMTRAIN_BUILD_PARALLEL	作为 -jN 传给 cmake --build
CMAKE_ARGS	追加给 cmake configure 阶段的额外参数

MUXI 后端在检测到 cmake_hpcc 时会自动改用它（与原版 BMTrain_muxi 的构建流程保持一致）。

后端探测顺序

1. 环境变量 BMTRAIN_BACKEND（最高优先级）。
2. import torch_mlu 成功 → mlu。
3. SDAA_HOME 已设置，或 /opt/tecoai 存在，或 import torch_sdaa 成功 → sdaa。
4. MACA_PATH 已设置，或 /opt/maca / /root/cu-bridge 存在 → muxi。
5. COREX_VERSION 已设置，或 /usr/local/corex 存在 → tianshu。
6. torch.cuda.is_available() → cuda。

如果都不命中，import bmtrain 会抛出 RuntimeError，并列出当前 支持的后端。

与既有代码的兼容性

原版 BMTrain 各分支里的所有公开名字和子模块，都会被重新导出到顶层 bmtrain 命名空间，因此用户代码不需要改动：

import bmtrain as bmt
from bmtrain.nn import Linear, ColumnParallelLinear
from bmtrain import optim, distributed, nccl

C 扩展同样会被暴露在传统的 bmtrain.C 路径下。