作者|王益
OneFlow社区编译
翻译|杨婷
最近,我在处理 PyTorch 分布式和 TorchRec 相关的工作,为此,我开始学习 PyTorch 2.0。在业余时间,我也在跟着Alpa作者学习JAX和XLA。如今回顾这些技术,我发现它们的关注点似乎都是如下两个问题:
-
包含自动求导和并行在内的函数转换,例如 vmap, pmap 和 pjit 等;
-
异构计算,CPU 负责控制流,GPU/TPU 负责张量计算和集合通信。
本文档中的所有例子都支持在 Colab 中运行:
Theano/Aesara |
https://colab.research.google.com/drive/1eg7C5WMNokhXgXQ46pNA30dXUCklquPz |
TensorFlow 1.x |
https://colab.research.google.com/drive/1jc0ePg2AAXBihevtoZM_33mmhC70rzqz?usp=sharing |
TensorFlow 2.x |
https://colab.research.google.com/drive/1PbftzJ9E2_FyIiuozTpExMvlFky_G2nv |
PyTorch 1.x |
https://colab.research.google.com/drive/1v4hENL-IJ-C6VT5H9W1NC2te85D8VdJK |
JAX |
https://colab.research.google.com/drive/1PlFijLIzAttIBd3tBjiEbSgPXvq9lVlg |
functorch/PyTorch 2.x |
https://colab.research.google.com/drive/1o-yJ-5g1V084RDaiRw2PqfAjOG7Ty951 |
1
函数转换
“函数转换”意为将一个程序转变成另一个程序,最常见的例子是自动求导(autograd)。自动求导采用用户编写的前向过程并创建后向过程,对于用户来说,编写自动求导通常都太过复杂。函数转换的主要难点在于:在编写函数转换算法时以何种方式表示输入和输出过程。
Theano:显式地构建 IR
Theano是最早的深度学习工具之一,也就是如今为人们所熟知的Aesara项目。Theano有一个允许用户在内存中将IR构建为数据结构的API,因此Theano可实现自动求导,并将结果输出为 Python 函数。
import aesara
from aesara import tensor as at
a = at.dscalar("a") # Define placeholders, which have no values.
b = at.dscalar("b")
c = a * b # c now contains the IR of an expression.TT
dc = aesara.grad(c, a) # Convert the IR in c into another one, dc
f_dc = aesara.function([a, b], dc) # Convert the IR into a Python function,
assert f_dc(1.5, 2.5) == 2.5 # so we can call it.
TensorFlow 1.x:用于运行 IR 的虚拟机
TensorFlow 1.x明确保留了构建IR的想法。若在TensorFlow中运行上述示例,结果不会有什么差别;但倘若在TensorFlow 1.x中来运行,最大的差别在于: