PyTorch教程-13.3. 自動并行

深度學(xué)習(xí)框架（例如 MXNet 和 PyTorch）在后端自動構(gòu)建計算圖。使用計算圖，系統(tǒng)了解所有依賴關(guān)系，并可以選擇性地并行執(zhí)行多個非相互依賴的任務(wù)以提高速度。例如，第 13.2 節(jié)中的圖 13.2.2 獨立地初始化了兩個變量。因此，系統(tǒng)可以選擇并行執(zhí)行它們。

通常，單個運算符將使用所有 CPU 或單個 GPU 上的所有計算資源。例如，dot算子將使用所有 CPU 上的所有內(nèi)核（和線程），即使在一臺機器上有多個 CPU 處理器。這同樣適用于單個 GPU。因此，并行化對于單設(shè)備計算機不是很有用。有了多個設(shè)備，事情就更重要了。雖然并行化通常在多個 GPU 之間最相關(guān)，但添加本地 CPU 會略微提高性能。例如，參見 Hadjis等人。( 2016 年)專注于訓(xùn)練結(jié)合 GPU 和 CPU 的計算機視覺模型。借助自動并行化框架的便利，我們可以在幾行 Python 代碼中實現(xiàn)相同的目標(biāo)。更廣泛地說，我們對自動并行計算的討論集中在使用 CPU 和 GPU 的并行計算，以及計算和通信的并行化。

請注意，我們至少需要兩個 GPU 才能運行本節(jié)中的實驗。

import torch
from d2l import torch as d2l

from mxnet import np, npx
from d2l import mxnet as d2l

npx.set_np()

13.3.1。GPU 上的并行計算

讓我們首先定義一個要測試的參考工作負(fù)載：run 下面的函數(shù)使用分配到兩個變量中的數(shù)據(jù)在我們選擇的設(shè)備上執(zhí)行 10 次矩陣-矩陣乘法：x_gpu1和 x_gpu2。

devices = d2l.try_all_gpus()
def run(x):
  return [x.mm(x) for _ in range(50)]

x_gpu1 = torch.rand(size=(4000, 4000), device=devices[0])
x_gpu2 = torch.rand(size=(4000, 4000), device=devices[1])

現(xiàn)在我們將函數(shù)應(yīng)用于數(shù)據(jù)。為了確保緩存不會在結(jié)果中發(fā)揮作用，我們通過在測量之前對其中任何一個執(zhí)行單次傳遞來預(yù)熱設(shè)備。torch.cuda.synchronize() 等待 CUDA 設(shè)備上所有流中的所有內(nèi)核完成。它接受一個device參數(shù)，即我們需要同步的設(shè)備。current_device()如果設(shè)備參數(shù)為（默認(rèn)），則它使用由 給出的當(dāng)前設(shè)備None。

run(x_gpu1)
run(x_gpu2) # Warm-up all devices
torch.cuda.synchronize(devices[0])
torch.cuda.synchronize(devices[1])

with d2l.Benchmark('GPU1 time'):
  run(x_gpu1)
  torch.cuda.synchronize(devices[0])

with d2l.Benchmark('GPU2 time'):
  run(x_gpu2)
  torch.cuda.synchronize(devices[1])

GPU1 time: 0.4967 sec
GPU2 time: 0.5151 sec

如果我們刪除synchronize兩個任務(wù)之間的語句，系統(tǒng)就可以自由地自動在兩個設(shè)備上并行計算。

with d2l.Benchmark('GPU1 & GPU2'):
  run(x_gpu1)
  run(x_gpu2)
  torch.cuda.synchronize()

GPU1 & GPU2: 0.5000 sec

devices = d2l.try_all_gpus()
def run(x):
  return [x.dot(x) for _ in range(50)]

x_gpu1 = np.random.uniform(size=(4000, 4000), ctx=devices[0])
x_gpu2 = np.random.uniform(size=(4000, 4000), ctx=devices[1])

Now we apply the function to the data. To ensure that caching does not play a role in the results we warm up the devices by performing a single pass on either of them prior to measuring.

run(x_gpu1) # Warm-up both devices
run(x_gpu2)
npx.waitall()

with d2l.Benchmark('GPU1 time'):
  run(x_gpu1)
  npx.waitall()

with d2l.Benchmark('GPU2 time'):
  run(x_gpu2)
  npx.waitall()

GPU1 time: 0.5233 sec
GPU2 time: 0.5158 sec

If we remove the waitall statement between both tasks the system is free to parallelize computation on both devices automatically.

with d2l.Benchmark('GPU1 & GPU2'):
  run(x_gpu1)
  run(x_gpu2)
  npx.waitall()

GPU1 & GPU2: 0.5214 sec

在上述情況下，總執(zhí)行時間小于其各部分的總和，因為深度學(xué)習(xí)框架會自動安排兩個 GPU 設(shè)備上的計算，而不需要代表用戶編寫復(fù)雜的代碼。

13.3.2。并行計算與通信

在許多情況下，我們需要在不同設(shè)備之間移動數(shù)據(jù)，比如在 CPU 和 GPU 之間，或者在不同 GPU 之間。例如，當(dāng)我們想要執(zhí)行分布式優(yōu)化時會發(fā)生這種情況，我們需要在多個加速器卡上聚合梯度。讓我們通過在 GPU 上計算然后將結(jié)果復(fù)制回 CPU 來對此進行模擬。

def copy_to_cpu(x, non_blocking=False):
  return [y.to('cpu', non_blocking=non_blocking) for y in x]

with d2l.Benchmark('Run on GPU1'):
  y = run(x_gpu1)
  torch.cuda.synchronize()

with d2l.Benchmark('Copy to CPU'):
  y_cpu = copy_to_cpu(y)
  torch.cuda.synchronize()

Run on GPU1: 0.5019 sec
Copy to CPU: 2.7168 sec

這有點低效。請注意，我們可能已經(jīng)開始將 的部分內(nèi)容復(fù)制y到 CPU，而列表的其余部分仍在計算中。這種情況會發(fā)生，例如，當(dāng)我們計算小批量的（反向傳播）梯度時。一些參數(shù)的梯度將比其他參數(shù)更早可用。因此，在 GPU 仍在運行時開始使用 PCI-Express 總線帶寬對我們有利。在 PyTorch 中，幾個函數(shù)（例如to()和）copy_()承認(rèn)一個顯式non_blocking參數(shù)，它允許調(diào)用者在不需要時繞過同步。設(shè)置non_blocking=True 允許我們模擬這種情況。

with d2l.Benchmark('Run on GPU1 and copy to CPU'):
  y = run(x_gpu1)
  y_cpu = copy_to_cpu(y, True)
  torch.cuda.synchronize()

Run on GPU1 and copy to CPU: 2.4682 sec

def copy_to_cpu(x):
  return [y.copyto(npx.cpu()) for y in x]

with d2l.Benchmark('Run on GPU1'):
  y = run(x_gpu1)
  npx.waitall()

with d2l.Benchmark('Copy to CPU'):
  y_cpu = copy_to_cpu(y)
  npx.waitall()

Run on GPU1: 0.5796 sec
Copy to CPU: 3.0989 sec

This is somewhat inefficient. Note that we could already start copying parts of y to the CPU while the remainder of the list is still being computed. This situation occurs, e.g., when we compute the gradient on a minibatch. The gradients of some of the parameters will be available earlier than that of others. Hence it works to our advantage to start using PCI-Express bus bandwidth while the GPU is still running. Removing waitall between both parts allows us to simulate this scenario.

with d2l.Benchmark('Run on GPU1 and copy to CPU'):
  y = run(x_gpu1)
  y_cpu = copy_to_cpu(y)
  npx.waitall()

Run on GPU1 and copy to CPU: 3.3488 sec

兩個操作所需的總時間（正如預(yù)期的那樣）小于它們各部分的總和。請注意，此任務(wù)不同于并行計算，因為它使用不同的資源：CPU 和 GPU 之間的總線。事實上，我們可以同時在兩個設(shè)備上進行計算和通信。如上所述，計算和通信之間存在依賴關(guān)系：y[i]必須在將其復(fù)制到 CPU 之前進行計算。幸運的是，系統(tǒng)可以y[i-1]邊計算邊 復(fù)制y[i]，以減少總運行時間。

我們以在一個 CPU 和兩個 GPU 上進行訓(xùn)練時簡單的兩層 MLP 的計算圖及其依賴關(guān)系的圖示作為結(jié)尾，如圖13.3.1所示。手動安排由此產(chǎn)生的并行程序?qū)⒎浅Ｍ纯?。這就是擁有基于圖形的計算后端進行優(yōu)化的優(yōu)勢所在。

圖 13.3.1兩層 MLP 在一個 CPU 和兩個 GPU 上的計算圖及其依賴關(guān)系。

13.3.3。概括

現(xiàn)代系統(tǒng)具有多種設(shè)備，例如多個 GPU 和 CPU。它們可以并行、異步使用。

現(xiàn)代系統(tǒng)還具有多種通信資源，例如 PCI Express、存儲（通常是固態(tài)驅(qū)動器或通過網(wǎng)絡(luò)）和網(wǎng)絡(luò)帶寬。它們可以并聯(lián)使用以達(dá)到最高效率。

后端可以通過自動并行計算和通信來提高性能。

13.3.4。練習(xí)

run在本節(jié)定義的函數(shù)中執(zhí)行了八個操作。它們之間沒有依賴關(guān)系。設(shè)計一個實驗，看看深度學(xué)習(xí)框架是否會自動并行執(zhí)行它們。

當(dāng)單個操作員的工作量足夠小時，并行化甚至可以在單個 CPU 或 GPU 上提供幫助。設(shè)計一個實驗來驗證這一點。

設(shè)計一個實驗，在 CPU、GPU 上使用并行計算，并在兩個設(shè)備之間進行通信。

使用 NVIDIA 的Nsight等調(diào)試器 來驗證您的代碼是否有效。

設(shè)計包含更復(fù)雜數(shù)據(jù)依賴關(guān)系的計算任務(wù)，并運行實驗以查看是否可以在提高性能的同時獲得正確的結(jié)果。