作者｜王益

(資料圖)

OneFlow社區(qū)編譯

翻譯｜楊婷

最近，我在處理 PyTorch 分布式和 TorchRec 相關(guān)的工作，為此，我開始學(xué)習(xí) PyTorch 2.0。在業(yè)余時(shí)間，我也在跟著Alpa作者學(xué)習(xí)JAX和XLA。如今回顧這些技術(shù)，我發(fā)現(xiàn)它們的關(guān)注點(diǎn)似乎都是如下兩個(gè)問(wèn)題：

包含自動(dòng)求導(dǎo)和并行在內(nèi)的函數(shù)轉(zhuǎn)換，例如 vmap, pmap 和 pjit 等；

異構(gòu)計(jì)算，CPU 負(fù)責(zé)控制流，GPU/TPU 負(fù)責(zé)張量計(jì)算和集合通信。

本文檔中的所有例子都支持在 Colab 中運(yùn)行：

1

函數(shù)轉(zhuǎn)換

“函數(shù)轉(zhuǎn)換”意為將一個(gè)程序轉(zhuǎn)變成另一個(gè)程序，最常見(jiàn)的例子是自動(dòng)求導(dǎo)（autograd）。自動(dòng)求導(dǎo)采用用戶編寫的前向過(guò)程并創(chuàng)建后向過(guò)程，對(duì)于用戶來(lái)說(shuō)，編寫自動(dòng)求導(dǎo)通常都太過(guò)復(fù)雜。函數(shù)轉(zhuǎn)換的主要難點(diǎn)在于：在編寫函數(shù)轉(zhuǎn)換算法時(shí)以何種方式表示輸入和輸出過(guò)程。

Theano：顯式地構(gòu)建 IR

Theano是最早的深度學(xué)習(xí)工具之一，也就是如今為人們所熟知的Aesara項(xiàng)目。Theano有一個(gè)允許用戶在內(nèi)存中將IR構(gòu)建為數(shù)據(jù)結(jié)構(gòu)的API，因此Theano可實(shí)現(xiàn)自動(dòng)求導(dǎo)，并將結(jié)果輸出為 Python 函數(shù)。 ?

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：用于運(yùn)行 IR 的虛擬機(jī)

TensorFlow 1.x明確保留了構(gòu)建IR的想法。若在TensorFlow中運(yùn)行上述示例，結(jié)果不會(huì)有什么差別；但倘若在TensorFlow 1.x中來(lái)運(yùn)行，最大的差別在于：我們不會(huì)將后向 IR 轉(zhuǎn)換為 Python 函數(shù)，并使用 Python 解釋器來(lái)運(yùn)行。相反，我們會(huì)在TensorFlow runtime中來(lái)運(yùn)行。 ?

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：沒(méi)有前向IR

PyTorch不會(huì)像Theano或TensorFlow那樣將前向傳播轉(zhuǎn)換為IR。反之，PyTorch 使用 Python 解釋器來(lái)運(yùn)行前向傳播。這樣做的弊端在于會(huì)在運(yùn)行期間生成表示后向傳播的 IR，我們稱之為Eager模式（動(dòng)態(tài)圖模式）。 ?

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增加了一個(gè)像PyTorch API的Eager模式API。此 API 追蹤前向傳播如何運(yùn)行名為梯度帶（GradientTape）的 IR 。TensorFlow 2.x可以從這個(gè)跟蹤中找出后向傳播。

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 不會(huì)向用戶公開諸如梯度帶等方面的低級(jí)別細(xì)節(jié)。簡(jiǎn)單說(shuō)來(lái)，JAX的思維方式為：將輸入和輸出都用Python函數(shù)來(lái)表示。

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)

對(duì)于想要自己編寫的函數(shù)轉(zhuǎn)換的高級(jí)用戶，他們可以調(diào)用make_jaxpr等低級(jí) API 來(lái)訪問(wèn) 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的函數(shù)轉(zhuǎn)換。

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. # 必須將內(nèi)置函數(shù)dot轉(zhuǎn)換成f. print(functorch.make_fx(f)(a, b).code)print(functorch.make_fx(functorch.grad(f))(a,?b).code)

TensorFlow 2.x、JAX 和 functorch 都為前向傳遞構(gòu)建了一個(gè) IR，但 PyTorch Eager模式?jīng)]有。IR 不僅可用于自動(dòng)求導(dǎo)，還可用于其他類型的函數(shù)轉(zhuǎn)換。在下列例子中，functorch.compile.aot_function調(diào)用了回調(diào)函數(shù)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高階導(dǎo)數(shù)

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動(dòng)態(tài)控制流

動(dòng)態(tài)控制流（dynamic control flows）有兩個(gè)層級(jí)：在 CPU 上運(yùn)行的粗粒度級(jí)別和在 GPU /TPU 上運(yùn)行的細(xì)粒度級(jí)別。本部分主要介紹在 CPU 上運(yùn)行的粗粒度級(jí)別的動(dòng)態(tài)控制流。下面我們將用(if/else)條件語(yǔ)句作為例子檢驗(yàn)深度學(xué)習(xí)工具。

TensorFlow 1.x

在 TensorFlow 1.x 中，我們需要將條件語(yǔ)句顯式構(gòu)建到 IR 中。此時(shí)條件語(yǔ)句是一個(gè)特殊的運(yùn)算符 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 顯式構(gòu)建控制流。此外，實(shí)驗(yàn)項(xiàng)目google/tangent中有AutoGraph功能，它可以將Python控制流轉(zhuǎn)換為tf.cond或tf.while_loop。此功能利用了 Python 解釋器支持的函數(shù)和函數(shù)源代碼。例如下面的g函數(shù)調(diào)用了 Python 的標(biāo)準(zhǔn)庫(kù)將源代碼解析為 AST，然后調(diào)用 SSA 表單來(lái)理解控制流。

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語(yǔ)法很復(fù)雜，所以通過(guò)解析源代碼來(lái)理解控制流就顯得很困難，這就導(dǎo)致AutoGraph經(jīng)常出錯(cuò)。但如果這種方法很簡(jiǎn)單，那么Python開發(fā)者社區(qū)也不會(huì)在構(gòu)建Python編譯器時(shí)失敗這么多次了。正是由于有這種挑戰(zhàn)的存在，必須要明確地將控制流構(gòu)建到 IR 中。為此，JAX 提供了 jax.lax.cond 和 jax.lax.for_loop函數(shù)。

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

考慮到這一點(diǎn)，你可能會(huì)覺(jué)得我們可以使用遞歸算法。但是下面用于計(jì)算階乘的遞歸無(wú)法用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)

可能你還想調(diào)用factorial來(lái)計(jì)算 3！=6。但這會(huì)讓遞歸深度超過(guò)最大值，因?yàn)檫f歸不僅依賴于條件，還依賴于函數(shù)定義和調(diào)用。

PyTorch

PyTorch最初是Python-native。正如前文所說(shuō)，由于多功能調(diào)度機(jī)制，grad 和 vamp 的函數(shù)轉(zhuǎn)換都是即時(shí)的。值得注意的是：

相比Theano 和 TensorFlow構(gòu)建IR后的函數(shù)轉(zhuǎn)換，即時(shí)函數(shù)轉(zhuǎn)換效率更高。

在進(jìn)行g(shù)rad和vmap 時(shí)，JAX也是即時(shí)函數(shù)轉(zhuǎn)換。然而像pamp和pjit等更復(fù)雜的函數(shù)轉(zhuǎn)換需要對(duì)整個(gè)計(jì)算過(guò)程進(jìn)行概述，在這個(gè)過(guò)程中IR是必不可少的。

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

4分布式計(jì)算

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

Python-Native

Theano和PyTorch采用了Python-native分布式計(jì)算方式。這種分布式訓(xùn)練工作包含多個(gè)Python解釋器進(jìn)程。這導(dǎo)致出現(xiàn)了以下結(jié)果。

打包和運(yùn)行（Pack and run）。由于這些 Python 進(jìn)程在不同的host上運(yùn)行，因此我們需要打包用戶程序和依賴項(xiàng)，并將它們發(fā)送到這些host上去運(yùn)行。一直以來(lái)TorchX負(fù)責(zé)了這個(gè)打包過(guò)程。它支持例如Docker和torch.package等各種打包格式，并且可以與各種集群管理器配合使用，如Kubernetes和SLURM。

單程序多數(shù)據(jù)（SPMD）。由于將用戶程序發(fā)送到各種host上要依賴于打包，與其他權(quán)重較輕的方式（如通過(guò) RPC 發(fā)送代碼）相比，這種方式不太靈活，因此，我們通常只發(fā)送一個(gè)程序。當(dāng)所有這些進(jìn)程運(yùn)行同一程序時(shí)，這個(gè)作業(yè)就變成了單程序多數(shù)據(jù)（SPMD）作業(yè)。

Python-native SPMD

下面是一個(gè)簡(jiǎn)單的SPMD PyTorch程序，我們可以在相同或不同的host上使用進(jìn)程運(yùn)行這個(gè)程序。在這個(gè)過(guò)程中，我們只需要調(diào)用all_gather。真正的分布式訓(xùn)練程序會(huì)調(diào)用更高級(jí)別的API，例如torch.nn.parallel.DistributedDataParallel 和 torchrec.DistributedModelParallel, 然后再調(diào)用低級(jí) 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 式的分布式訓(xùn)練。它還通過(guò)torch.distributed.pipeline.sync.Pipe和PiPPy project提供流水并行，其中流水并行的各個(gè)階段在不同的設(shè)備上運(yùn)行不同的程序。這些階段常通過(guò) torch.rpc 包來(lái)溝通。

分布式運(yùn)行時(shí)機(jī)制

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

用戶程序可以使用低級(jí)API（如 tf.device）去指定作業(yè)要運(yùn)行什么操作、在哪臺(tái)設(shè)備和主機(jī)上運(yùn)行等等。因?yàn)锳PI有runtime，所以可以做到這一點(diǎn)。

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

與PyTorch一樣，TensorFlow也為分布式訓(xùn)練提供了高級(jí)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")

分布式運(yùn)行時(shí)極大地方便了訓(xùn)練服務(wù)的維護(hù)，因?yàn)槲覀儾辉賹⒂脩舫绦虼虬郊荷线\(yùn)行。相反，我們打包運(yùn)行時(shí)程序，因?yàn)橄啾扔脩舫绦颍\(yùn)行時(shí)程序更加統(tǒng)一。

混合理念

JAX 支持 Python-native 和分布式運(yùn)行時(shí)。

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

（本文經(jīng)授權(quán)后由OneFlow社區(qū)編譯，譯文轉(zhuǎn)載請(qǐng)聯(lián)系獲得授權(quán)。原文：https://quip.com/Y8qtAyV4EXRg）

其他人都在看

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

關(guān)于ChatGPT的一切；CUDA入門之矩陣乘

李白：你的模型權(quán)重很不錯(cuò)，可惜被我沒(méi)收了

單RTX 3090訓(xùn)練YOLOv5s，時(shí)間減少11小時(shí)

OpenAI掌門Sam Altman：AI下一個(gè)發(fā)展階段

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

OneEmbedding:單卡訓(xùn)練TB級(jí)推薦模型不是夢(mèng)

歡迎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/

關(guān)鍵詞：