作者｜王益

(資料圖)

OneFlow社區編譯

翻譯｜楊婷

最近，我在處理 PyTorch 分布式和 TorchRec 相關的工作，為此，我開始學習 PyTorch 2.0。在業余時間，我也在跟著Alpa作者學習JAX和XLA。如今回顧這些技術，我發現它們的關注點似乎都是如下兩個問題：

包含自動求導和并行在內的函數轉換，例如 vmap, pmap 和 pjit 等；

異構計算，CPU 負責控制流，GPU/TPU 負責張量計算和集合通信。

本文檔中的所有例子都支持在 Colab 中運行：

1

函數轉換

“函數轉換”意為將一個程序轉變成另一個程序，最常見的例子是自動求導（autograd）。自動求導采用用戶編寫的前向過程并創建后向過程，對于用戶來說，編寫自動求導通常都太過復雜。函數轉換的主要難點在于：在編寫函數轉換算法時以何種方式表示輸入和輸出過程。

Theano：顯式地構建 IR

Theano是最早的深度學習工具之一，也就是如今為人們所熟知的Aesara項目。Theano有一個允許用戶在內存中將IR構建為數據結構的API，因此Theano可實現自動求導，并將結果輸出為 Python 函數。 ?

import aesarafrom aesara import tensor as ata = at.dscalar("a") # Define placeholders, which have no values.b = at.dscalar("b")c = a * b # c now contains the IR of an expression.TTdc = aesara.grad(c, a) # Convert the IR in c into another one, dcf_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中來運行，最大的差別在于：我們不會將后向 IR 轉換為 Python 函數，并使用 Python 解釋器來運行。相反，我們會在TensorFlow runtime中來運行。 ?

import tensorflow.compat.v1 as tf # TensorFlow 1.x APIimport numpy as nptf.disable_eager_execution()a = tf.placeholder(tf.float32, shape=())b?=?tf.placeholder(tf.float32,?shape=())c = a * bdc?=?tf.gradients(c,?[a],?stop_gradients=[a,?b])with tf.compat.v1.Session() as sess: # TensorFlow has a runtime to execute the IR, x = np.single(2) # so, no converting it into Python code. y = np.single(3) print(sess.run(dc, feed_dict={a:x, b:y}))

PyTorch 1.x：沒有前向IR

PyTorch不會像Theano或TensorFlow那樣將前向傳播轉換為IR。反之，PyTorch 使用 Python 解釋器來運行前向傳播。這樣做的弊端在于會在運行期間生成表示后向傳播的 IR，我們稱之為Eager模式（動態圖模式）。 ?

import torcha = torch.tensor(1.0, requires_grad=True) # These are not placeholders, but values.b = torch.tensor(2.0)c = a * b # Evaluates c and derives the IR of the backward in c.grad_fn_.c.backward() # Executes c.grad_fn_.print(c.grad)

TensorFlow 2.x: 梯度帶

TensorFlow 2.x增加了一個像PyTorch API的Eager模式API。此 API 追蹤前向傳播如何運行名為梯度帶（GradientTape）的 IR 。TensorFlow 2.x可以從這個跟蹤中找出后向傳播。

import tensorflow as tfa = tf.Variable(1.0) # Like PyTorch, these are values, not placehodlers. b = tf.Variable(2.0)with tf.GradientTape() as tape: c = a * bdcda = tape.gradient(c, a)print(dcda)

JAX

JAX 不會向用戶公開諸如梯度帶等方面的低級別細節。簡單說來，JAX的思維方式為：將輸入和輸出都用Python函數來表示。

import?jax?a = 2.0b = 3.0jax.grad(jax.lax.mul)(a,?b)??# Compute c = a * b w.r.t. a. The result is b=3. jax.jit(jax.grad(jax.lax.mul))(a,b)jax.experimental.pjit(jax.grad(jax.lax.mul), device_mesh(ntpus))(a,b)

對于想要自己編寫的函數轉換的高級用戶，他們可以調用make_jaxpr等低級 API 來訪問 IR，稱為 JAXPR。

jax.make_jaxpr(jax.lax.mul)(2.0, 3.0) # Returns the IR representing jax.lax.mul(2,3)jax.make_jaxpr(jax.grad(jax.lax.mul))(2.0, 3.0) # Returns the IR of grad(mul)(2,3)

FuncTorch

FuncTorch和JAX類似，都是基于PyTorch的函數轉換。

import?torch,?functorcha = torch.tensor([2.0])b = torch.tensor([3.0])functorch.grad(torch.dot)(a, b)

JAX的make_jaxpr類似于functorch的make_fx。

def f(a, b): return torch.dot(a, b) # Have to wrap the builtin function dot into f. # 必須將內置函數dot轉換成f. print(functorch.make_fx(f)(a, b).code)print(functorch.make_fx(functorch.grad(f))(a,?b).code)

TensorFlow 2.x、JAX 和 functorch 都為前向傳遞構建了一個 IR，但 PyTorch Eager模式沒有。IR 不僅可用于自動求導，還可用于其他類型的函數轉換。在下列例子中，functorch.compile.aot_function調用了回調函數print_compile_fn兩次，分別用于前向和后向傳播。

from functorch.compile import aot_functionimport?torch.fx?as?fxdef print_compile_fn(fx_module, args): print(fx_module) return fx_moduleaot_fn = aot_function(torch.dot, print_compile_fn)aot_fn(a, b)

2高階導數

PyTorch

import torchfrom torch import autogradx = torch.tensor(1., requires_grad = True)y = 2*x**3 + 8first_derivative = autograd.grad(y, x, create_graph=True)print(first_derivative)second_derivative = autograd.grad(first_derivative, x)print(second_derivative)

TensorFlow 2.x

import?tensorflow?as?tfx?=?tf.Variable(1.0)with tf.GradientTape() as outer_tape: with tf.GradientTape() as tape: y = 2*x**3 + 8 dy_dx = tape.gradient(y, x) print(dy_dx) d2y_dx2 = outer_tape.gradient(dy_dx, x) print(d2y_dx2)

JAX

def f(a): return 2*a**3 + 8print(jax.grad(f)(1.0))print(jax.grad(jax.grad(f))(1.0))

3動態控制流

動態控制流（dynamic control flows）有兩個層級：在 CPU 上運行的粗粒度級別和在 GPU /TPU 上運行的細粒度級別。本部分主要介紹在 CPU 上運行的粗粒度級別的動態控制流。下面我們將用(if/else)條件語句作為例子檢驗深度學習工具。

TensorFlow 1.x

在 TensorFlow 1.x 中，我們需要將條件語句顯式構建到 IR 中。此時條件語句是一個特殊的運算符 tf.cond。

def f1(): return tf.multiply(a, 17)def f2(): return tf.add(b, 23)r = tf.cond(tf.less(a, b), f1, f2)with tf.compat.v1.Session() as sess: # TensorFlow has a runtime to execute the IR, print(sess.run(r, feed_dict={a:x, b:y}))

TensorFlow 2.x

TensorFlow 2.x 支持使用 tf.cond 和 tf.while_loop 顯式構建控制流。此外，實驗項目google/tangent中有AutoGraph功能，它可以將Python控制流轉換為tf.cond或tf.while_loop。此功能利用了 Python 解釋器支持的函數和函數源代碼。例如下面的g函數調用了 Python 的標準庫將源代碼解析為 AST，然后調用 SSA 表單來理解控制流。

def g(x, y): if tf.reduce_any(x < y): return tf.multiply(x, 17) return tf.add(y, 23) converted_g?=?tf.autograph.to_graph(g)import inspectprint(inspect.getsource(converted_g))

JAX

由于部分Python語法很復雜，所以通過解析源代碼來理解控制流就顯得很困難，這就導致AutoGraph經常出錯。但如果這種方法很簡單，那么Python開發者社區也不會在構建Python編譯器時失敗這么多次了。正是由于有這種挑戰的存在，必須要明確地將控制流構建到 IR 中。為此，JAX 提供了 jax.lax.cond 和 jax.lax.for_loop函數。

jax.lax.cond(a < b, lambda : a*17, lambda: b+23)

考慮到這一點，你可能會覺得我們可以使用遞歸算法。但是下面用于計算階乘的遞歸無法用JAX跟蹤。

def factorial(r, x): return jax.lax.cond(x <= 1.0, lambda: r, lambda: factorial(r*x, x-1))factorial(1.0, 3.0)

可能你還想調用factorial來計算 3！=6。但這會讓遞歸深度超過最大值，因為遞歸不僅依賴于條件，還依賴于函數定義和調用。

PyTorch

PyTorch最初是Python-native。正如前文所說，由于多功能調度機制，grad 和 vamp 的函數轉換都是即時的。值得注意的是：

相比Theano 和 TensorFlow構建IR后的函數轉換，即時函數轉換效率更高。

在進行grad和vmap 時，JAX也是即時函數轉換。然而像pamp和pjit等更復雜的函數轉換需要對整個計算過程進行概述，在這個過程中IR是必不可少的。

由于IR在pmap 和 pjit中的必要性，PyTorch社區最近添加了torch.condpytorch/pytorch#83154 ?

4分布式計算

根據執行代碼或 IR 的不同方式，在使用 Python 解釋器或runtime時，有兩種分布式計算方法。

Python-Native

Theano和PyTorch采用了Python-native分布式計算方式。這種分布式訓練工作包含多個Python解釋器進程。這導致出現了以下結果。

打包和運行（Pack and run）。由于這些 Python 進程在不同的host上運行，因此我們需要打包用戶程序和依賴項，并將它們發送到這些host上去運行。一直以來TorchX負責了這個打包過程。它支持例如Docker和torch.package等各種打包格式，并且可以與各種集群管理器配合使用，如Kubernetes和SLURM。

單程序多數據（SPMD）。由于將用戶程序發送到各種host上要依賴于打包，與其他權重較輕的方式（如通過 RPC 發送代碼）相比，這種方式不太靈活，因此，我們通常只發送一個程序。當所有這些進程運行同一程序時，這個作業就變成了單程序多數據（SPMD）作業。

Python-native SPMD

下面是一個簡單的SPMD PyTorch程序，我們可以在相同或不同的host上使用進程運行這個程序。在這個過程中，我們只需要調用all_gather。真正的分布式訓練程序會調用更高級別的API，例如torch.nn.parallel.DistributedDataParallel 和 torchrec.DistributedModelParallel, 然后再調用低級 API，例如 all_gather 和 all_reduce。

import osimport torchfrom torch import distributed as distdef main(): use_gpu = torch.cuda.is_available() local_rank = int(os.environ.get("LOCAL_RANK", "0")) local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "0")) device = torch.device(f"cuda:{local_rank}" if use_gpu else "cpu") dist.init_distributed(backend="nccl") lst = torch.tensor([local_rank + 100]).to(device) # placeholder rlt_lst = [torch.zeros_like(lst) for _ in range(local_world_size)] dist.all_gather(rlt_lst, lst, async_op=False)????print("After?broadcasting:",?rlt_lst)

Python-native Non-SPMD

PyTorch 不僅限于 SPMD 式的分布式訓練。它還通過torch.distributed.pipeline.sync.Pipe和PiPPy project提供流水并行，其中流水并行的各個階段在不同的設備上運行不同的程序。這些階段常通過 torch.rpc 包來溝通。

分布式運行時機制

分布式 TensorFlow 作業由運行 TensorFlow runtime 程序的進程組成，而不是由 Python 解釋器組成。此分布式運行時作業執行 TensorFlow graph (IR)，它是由執行用戶程序的 Python 解釋器生成。

用戶程序可以使用低級API（如 tf.device）去指定作業要運行什么操作、在哪臺設備和主機上運行等等。因為API有runtime，所以可以做到這一點。

with tf.device("/job:bar/task:0/device:gpu:2"):????#?ops?created?here?have?the?fully?specified?device?above

與PyTorch一樣，TensorFlow也為分布式訓練提供了高級API tf.distributed.strategy，Keras和DTensor。

strategy = tf.distribute.MirroredStrategy() \ if tf.config.list_physical_devices("GPU") \???????????else?tf.distribute.get_strategy()with strategy.scope(): model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])model.compile(loss="mse", optimizer="sgd")

分布式運行時極大地方便了訓練服務的維護，因為我們不再將用戶程序打包到集群上運行。相反，我們打包運行時程序，因為相比用戶程序，運行時程序更加統一。

混合理念

JAX 支持 Python-native 和分布式運行時。

JAX 提供例如vmap、pmap 和 pjit的函數轉換，這可以將 Python 函數轉換為分布式程序。

（本文經授權后由OneFlow社區編譯，譯文轉載請聯系獲得授權。原文：https://quip.com/Y8qtAyV4EXRg）

其他人都在看

下載量突破10億，MinIO的開源啟示錄

關于ChatGPT的一切；CUDA入門之矩陣乘

李白：你的模型權重很不錯，可惜被我沒收了

單RTX 3090訓練YOLOv5s，時間減少11小時

OpenAI掌門Sam Altman：AI下一個發展階段

比快更快，開源Stable Diffusion刷新作圖速度

OneEmbedding:單卡訓練TB級推薦模型不是夢

歡迎Star、試用OneFlow最新版本：GitHub - Oneflow-Inc/oneflow: OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient.OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient. - GitHub - Oneflow-Inc/oneflow: OneFlow is a deep learning framework designed to be user-friendly, scalable and efficient.https://github.com/Oneflow-Inc/oneflow/

關鍵詞：