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-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() 即可

POFO

使用 Rotor[5] 实现。

  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. https://gitlab.inria.fr/hiepacs/rotor/-/tree/offload ↩︎