Comparasion Schemes

Sys 工作的对比实验做起来十分麻烦。MAGIS 对比了 TVM、XLA 等多个后端。虽是体力劳动，但也努力许久，若不作记录，实在可惜。故将本文作为复现文章之附录，单开一篇，逐一记录各个后端的配置及试错过程。

TVM

参考官方文档[1]安装的是 apache-tvm v0.11，这个版本是跑不起来。是要带 --pre 参数安装。

1
2
pip install apache-tvm --pre
Successfully installed apache-tvm-0.14.dev273 ml-dtypes-0.5.0

补装以下库

1
2
pip install tlcpack-nightly-cu116 --pre -f https://tlcpack.ai/wheels
pip install pytest

XLA

PyTorch 与 TensorFlow 都有支持。当然是先尝试 PyTorch-XLA，一步一坑，无比麻烦。相关链接 [1:1] [2] [3] [4]，这里只粗略记录。

  1. 为兼容低版本 torch 而只能装旧版 xla，无教程可循。
  2. 跑样例，ImportError: libpython3.10.so.1.0: cannot open shared object file
  3. 检查 /usr/local/lib，只有 libpython3.10.a，原来是编译时只生成了静态库
  4. 重新编译安装 Python
  5. 更新动态链接器缓存 sudo ldconfig
  6. 重跑提示 Undefined Symbol
  7. 至暗时刻。尝试换 xla 版本，但是根本找不到仓库，无奈去翻 xla 的 GCS Bucket 自行拼链接下载，反复测试数个版本不得行
  8. 最后还是会报各种 Attribute Error 等

无奈转投 Tensorflow-XLA 门下。

其他后端

使用 Triton 算子库加速

一般使用，通过 torch.compile() 即可。[5]

POFO

冷门东西，怀疑是拉来凑数的，使用 Rotor[6] 实现。

  1. (Homepage) openxla/xla: A machine learning compiler for GPUs, CPUs, and ML accelerators ↩︎ ↩︎

  2. How to run with PyTorch/XLA:GPU — PyTorch/XLA master documentation ↩︎

  3. pytorch/xla: Enabling PyTorch on XLA Devices (e.g. Google TPU) ↩︎

  4. Releases · pytorch/pytorch ↩︎

  5. Triton 并不是与其他几位同类的概念，但是论文里确实这样写了，这里直接沿用。 ↩︎

  6. https://gitlab.inria.fr/hiepacs/rotor/-/tree/offload ↩︎