作者|王益
(資料圖)
OneFlow社區(qū)編譯
翻譯|楊婷
最近,我在處理 PyTorch 分布式和 TorchRec 相關(guān)的工作,為此,我開始學習 PyTorch 2.0。在業(yè)余時間,我也在跟著Alpa作者學習JAX和XLA。如今回顧這些技術(shù),我發(fā)現(xiàn)它們的關(guān)注點似乎都是如下兩個問題:
包含自動求導和并行在內(nèi)的函數(shù)轉(zhuǎn)換,例如 vmap, pmap 和 pjit 等;
異構(gòu)計算,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
函數(shù)轉(zhuǎn)換
“函數(shù)轉(zhuǎn)換”意為將一個程序轉(zhuǎn)變成另一個程序,最常見的例子是自動求導(autograd)。自動求導采用用戶編寫的前向過程并創(chuàng)建后向過程,對于用戶來說,編寫自動求導通常都太過復雜。函數(shù)轉(zhuǎn)換的主要難點在于:在編寫函數(shù)轉(zhuǎn)換算法時以何種方式表示輸入和輸出過程。
Theano:顯式地構(gòu)建 IR
Theano是最早的深度學習工具之一,也就是如今為人們所熟知的Aesara項目。Theano有一個允許用戶在內(nèi)存中將IR構(gòu)建為數(shù)據(jù)結(jié)構(gòu)的API,因此Theano可實現(xiàn)自動求導,并將結(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:用于運行 IR 的虛擬機
TensorFlow 1.x明確保留了構(gòu)建IR的想法。若在TensorFlow中運行上述示例,結(jié)果不會有什么差別;但倘若在TensorFlow 1.x中來運行,最大的差別在于:我們不會將后向 IR 轉(zhuǎn)換為 Python 函數(shù),并使用 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那樣將前向傳播轉(zhuǎn)換為IR。反之,PyTorch 使用 Python 解釋器來運行前向傳播。這樣做的弊端在于會在運行期間生成表示后向傳播的 IR,我們稱之為Eager模式(動態(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增加了一個像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 不會向用戶公開諸如梯度帶等方面的低級別細節(jié)。簡單說來,JAX的思維方式為:將輸入和輸出都用Python函數(shù)來表示。
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)
對于想要自己編寫的函數(shù)轉(zhuǎn)換的高級用戶,他們可以調(diào)用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的函數(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)建了一個 IR,但 PyTorch Eager模式?jīng)]有。IR 不僅可用于自動求導,還可用于其他類型的函數(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高階導數(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動態(tài)控制流
動態(tài)控制流(dynamic control flows)有兩個層級:在 CPU 上運行的粗粒度級別和在 GPU /TPU 上運行的細粒度級別。本部分主要介紹在 CPU 上運行的粗粒度級別的動態(tài)控制流。下面我們將用(if/else)條件語句作為例子檢驗深度學習工具。
TensorFlow 1.x
在 TensorFlow 1.x 中,我們需要將條件語句顯式構(gòu)建到 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 顯式構(gòu)建控制流。此外,實驗項目google/tangent中有AutoGraph功能,它可以將Python控制流轉(zhuǎn)換為tf.cond或tf.while_loop。此功能利用了 Python 解釋器支持的函數(shù)和函數(shù)源代碼。例如下面的g函數(shù)調(diào)用了 Python 的標準庫將源代碼解析為 AST,然后調(diào)用 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經(jīng)常出錯。但如果這種方法很簡單,那么Python開發(fā)者社區(qū)也不會在構(gòu)建Python編譯器時失敗這么多次了。正是由于有這種挑戰(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)
考慮到這一點,你可能會覺得我們可以使用遞歸算法。但是下面用于計算階乘的遞歸無法用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來計算 3!=6。但這會讓遞歸深度超過最大值,因為遞歸不僅依賴于條件,還依賴于函數(shù)定義和調(diào)用。
PyTorch
PyTorch最初是Python-native。正如前文所說,由于多功能調(diào)度機制,grad 和 vamp 的函數(shù)轉(zhuǎn)換都是即時的。值得注意的是:
相比Theano 和 TensorFlow構(gòu)建IR后的函數(shù)轉(zhuǎn)換,即時函數(shù)轉(zhuǎn)換效率更高。
在進行g(shù)rad和vmap 時,JAX也是即時函數(shù)轉(zhuǎn)換。然而像pamp和pjit等更復雜的函數(shù)轉(zhuǎn)換需要對整個計算過程進行概述,在這個過程中IR是必不可少的。
由于IR在pmap 和 pjit中的必要性,PyTorch社區(qū)最近添加了torch.condpytorch/pytorch#83154 ?
4分布式計算
根據(jù)執(zhí)行代碼或 IR 的不同方式,在使用 Python 解釋器或runtime時,有兩種分布式計算方法。
Python-Native
Theano和PyTorch采用了Python-native分布式計算方式。這種分布式訓練工作包含多個Python解釋器進程。這導致出現(xiàn)了以下結(jié)果。
打包和運行(Pack and run)。由于這些 Python 進程在不同的host上運行,因此我們需要打包用戶程序和依賴項,并將它們發(fā)送到這些host上去運行。一直以來TorchX負責了這個打包過程。它支持例如Docker和torch.package等各種打包格式,并且可以與各種集群管理器配合使用,如Kubernetes和SLURM。
單程序多數(shù)據(jù)(SPMD)。由于將用戶程序發(fā)送到各種host上要依賴于打包,與其他權(quán)重較輕的方式(如通過 RPC 發(fā)送代碼)相比,這種方式不太靈活,因此,我們通常只發(fā)送一個程序。當所有這些進程運行同一程序時,這個作業(yè)就變成了單程序多數(shù)據(jù)(SPMD)作業(yè)。
Python-native SPMD
下面是一個簡單的SPMD PyTorch程序,我們可以在相同或不同的host上使用進程運行這個程序。在這個過程中,我們只需要調(diào)用all_gather。真正的分布式訓練程序會調(diào)用更高級別的API,例如torch.nn.parallel.DistributedDataParallel 和 torchrec.DistributedModelParallel, 然后再調(diào)用低級 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 作業(yè)由運行 TensorFlow runtime 程序的進程組成,而不是由 Python 解釋器組成。此分布式運行時作業(yè)執(zhí)行 TensorFlow graph (IR),它是由執(zhí)行用戶程序的 Python 解釋器生成。
用戶程序可以使用低級API(如 tf.device)去指定作業(yè)要運行什么操作、在哪臺設備和主機上運行等等。因為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")
分布式運行時極大地方便了訓練服務的維護,因為我們不再將用戶程序打包到集群上運行。相反,我們打包運行時程序,因為相比用戶程序,運行時程序更加統(tǒng)一。
混合理念
JAX 支持 Python-native 和分布式運行時。
JAX 提供例如vmap、pmap 和 pjit的函數(shù)轉(zhuǎn)換,這可以將 Python 函數(shù)轉(zhuǎn)換為分布式程序。
(本文經(jīng)授權(quán)后由OneFlow社區(qū)編譯,譯文轉(zhuǎn)載請聯(lián)系獲得授權(quán)。原文:https://quip.com/Y8qtAyV4EXRg)
其他人都在看
下載量突破10億,MinIO的開源啟示錄
關(guān)于ChatGPT的一切;CUDA入門之矩陣乘
李白:你的模型權(quán)重很不錯,可惜被我沒收了
單RTX 3090訓練YOLOv5s,時間減少11小時
OpenAI掌門Sam Altman:AI下一個發(fā)展階段
比快更快,開源Stable Diffusion刷新作圖速度
OneEmbedding:單卡訓練TB級推薦模型不是夢
關(guān)鍵詞: