【動手學深度學習】第十二章筆記：非同步計算、資料並行

bringlu發表於2023-05-04
深度學習筆記非同步並行

為了更好的閱讀體驗，請點選這裡

12.1 編譯器和直譯器

原書主要關注的是指令式程式設計（imperative programming）。Python 是一種解釋性語言，因此沒有編譯器給程式碼最佳化，程式碼會跑得很慢。

12.1.1 符號式程式設計

考慮另一種選擇符號式程式設計（symbolic programming），即程式碼通常只在完全定義了過程之後才執行計算。這個策略被多個深度學習框架使用，包括 Theano 和 TensorFlow（後者已經獲得了指令式程式設計的擴充套件）。一般包括以下步驟：

  1. 定義計算流程；
  2. 將流程編譯成可執行的程式；
  3. 給定輸入，呼叫編譯好的程式執行。

這將允許進行大量的最佳化。首先，在大多數情況下，我們可以跳過 Python 直譯器。從而消除因為多個更快的 GPU 與單個 CPU 上的單個 Python 執行緒搭配使用時產生的效能瓶頸。其次，編譯器可以將程式碼最佳化和重寫。因為編譯器在將其轉換為機器指令之前可以看到完整的程式碼，所以這種最佳化是可以實現的。例如，只要某個變數不再需要，編譯器就可以釋放記憶體（或者從不分配記憶體），或者將程式碼轉換為一個完全等價的片段。下面，我們將透過模擬指令式程式設計來進一步瞭解符號式程式設計的概念。

def add_():
    return '''
def add(a, b):
    return a + b
'''

def fancy_func_():
    return '''
def fancy_func(a, b, c, d):
    e = add(a, b)
    f = add(c, d)
    g = add(e, f)
    return g
'''

def evoke_():
    return add_() + fancy_func_() + 'print(fancy_func(1, 2, 3, 4))'

prog = evoke_()
print(prog)
y = compile(prog, '', 'exec')
exec(y)

def add(a, b):
    return a + b

def fancy_func(a, b, c, d):
    e = add(a, b)
    f = add(c, d)
    g = add(e, f)
    return g
print(fancy_func(1, 2, 3, 4))
10

裡面出現了神奇的兩個函式 compile() 和 exec()：

  • compile(source, filename, mode[, flags[, dont_inherit]])
    • source：字串或者 AST（Abstract Syntax Trees）物件
    • filename：程式碼檔名稱，如果不是從檔案讀取程式碼則傳遞一些可辨認的值
    • mode：指定編譯程式碼的種類。可以指定為 exec, eval, single
    • flags：變數作用域，區域性名稱空間，如果被提供，可以是任何對映物件
    • flags 和 dont_inherit 是用來控制編譯原始碼時的標誌
  • exec(obj)
    • obj：要執行的表示式。

命令式（解釋型）程式設計和符號式程式設計的區別如下：

  • 指令式程式設計更容易使用。在 Python 中，指令式程式設計的大部分程式碼都是簡單易懂的。指令式程式設計也更容易除錯，這是因為無論是獲取和列印所有的中間變數值，或者使用 Python 的內建除錯工具都更加簡單；
  • 符號式程式設計執行效率更高，更易於移植。符號式程式設計更容易在編譯期間最佳化程式碼，同時還能夠將程式移植到與 Python 無關的格式中，從而允許程式在非 Python 環境中執行，避免了任何潛在的與 Python 直譯器相關的效能問題。

12.1.2 混合式程式設計

PyTorch 是基於指令式程式設計並且使用動態計算圖。為了能夠利用符號式程式設計的可移植性和效率，開發人員思考能否將這兩種程式設計模型的優點結合起來，於是就產生了 TorchScript。TorchScript 允許使用者使用純指令式程式設計進行開發和除錯，同時能夠將大多數程式轉換為符號式程式，以便在需要產品級計算效能和部署時使用。

接下來假設已經定義好了一個網路比如 net=MLP()，那麼可以使用 net = torch.jit.script(net) 程式碼使用 TorchScript：

  • torch.jit.script(obj, optimize=None, _frames_up=0, _rcb=None, example_inputs=None)：
    • 編寫一個函式或 nn.Module 指令碼將檢查原始碼，使用 TorchScript 編譯器將其編譯為 TorchScript 程式碼，並返回 ScriptModule 或 ScriptFunction。 TorchScript 本身是 Python 語言的一個子集，因此並非 Python 中的所有功能都有效，但我們提供了足夠的功能來計算張量並執行依賴於控制的操作。有關完整指南，請參閱 TorchScript 語言參考。
    • 編寫字典或列表的指令碼會將其中的資料複製到 TorchScript 例項中，隨後可以透過引用在 Python 和 TorchScript 之間以零複製開銷傳遞。
    • torch.jit.script() 可以為模組、函式、字典和列表用作函式，而且還可以被用作裝飾器。
    • 返回：
      • 如果 obj 是 nn.Module，script 會返回一個 ScriptModule。返回的 ScriptModule 將與原來的 nn.Module 有相同的子模組和引數集合。
      • 如果 obj 是獨立的函式，一個 ScriptFunction 將會返回。
      • 如果 obj 是字典，將會返回 torch._C.ScriptDict。
      • 如果 obj 是列表，將會返回 torch._C.ScriptList。

在使用上面轉化成 TorchScript 的程式碼後，一個三層的多層感知機大約增快了 20%。而且，還可以方便地使用 net.save('filepath.pt') 來儲存網路結構。眾所周知，普通的 torch.save()/torch.load() 是不能在沒有原本的模組類定義下讀取模型的。但是在 TorchScript 中，接下來即使我們刪除了原本的多層感知機的類以及衍生的例項，也可以透過 torch.jit.load('filepath.pt') 重新載入模型。當然也不排除是我沒刪乾淨

12.2 非同步計算

PyTorch 使用了 Python 自己的排程器來實現不同的效能權衡。對 PyTorch 來說 GPU操 作在預設情況下是非同步的。當呼叫一個使用 GPU 的函式時，操作會排隊到特定的裝置上，但不一定要等到以後才執行。這允許並行執行更多的計算，包括在 CPU 或其他 GPU 上的操作。

因此，瞭解非同步程式設計是如何工作的，透過主動地減少計算需求和相互依賴，有助於我們開發更高效的程式。這能夠減少記憶體開銷並提高處理器利用率。下面測試一下 numpy(CPU) 和 PyTorch(GPU) 的速度。

# GPU計算熱身
device = d2l.try_gpu()
a = torch.randn(size=(1000, 1000), device=device)
b = torch.mm(a, a)

with d2l.Benchmark('numpy'):
    for _ in range(10):
        a = numpy.random.normal(size=(1000, 1000))
        b = numpy.dot(a, a)

with d2l.Benchmark('torch'):
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)

numpy: 1.0981 sec
torch: 0.0011 sec

預設情況下，GPU 操作在 PyTorch 中是非同步的。強制 PyTorch 在返回之前完成所有計算，這種強制說明了之前發生的情況：計算是由後端執行，而前端將控制權返回給了 Python。

例如下面呼叫 torch.cuda.synchronize(device)，這個函式等待在一個 CUDA 裝置上所有核的所有流都完成。

with d2l.Benchmark():
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)
    torch.cuda.synchronize(device)

Done: 0.0089 sec

廣義上說，PyTorch 有一個用於與使用者直接互動的前端（例如透過 Python），還有一個由系統用來執行計算的後端。使用者可以用各種前端語言編寫 PyTorch 程式，如 Python 和 C++。不管使用的前端程式語言是什麼，PyTorch 程式的執行主要發生在 C++ 實現的後端。由前端語言發出的操作被傳遞到後端執行。後端管理自己的執行緒，這些執行緒不斷收集和執行排隊的任務。請注意，要使其工作，後端必須能夠跟蹤計算圖中各個步驟之間的依賴關係。因此，不可能並行化相互依賴的操作。

當語句的結果需要被列印出來時，Python 前端執行緒將等待 C++ 後端執行緒完成結果計算。這種設計的一個好處是 Python 前端執行緒不需要執行實際的計算。因此，不管 Python 的效能如何，對程式的整體效能幾乎沒有影響。

練習題

（1）在CPU上，對本節中相同的矩陣乘法操作進行基準測試，仍然可以透過後端觀察非同步嗎？

torch 觀察不到非同步現象，反倒是 numpy 可以觀察到非同步的現象。雖然 torch.cuda.synchronize(torch.device('cpu')) 會彈出報錯，但是仍然可以使用以下兩個程式碼來測試速度：

with d2l.Benchmark('numpy'):
    for _ in range(10):
        a = numpy.random.normal(size=(1000, 1000))
        b = numpy.dot(a, a)
        
# time.sleep(5)
        
with d2l.Benchmark('torch'):
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)

numpy: 0.9737 sec
torch: 0.2859 sec

with d2l.Benchmark('numpy'):
    for _ in range(10):
        a = numpy.random.normal(size=(1000, 1000))
        b = numpy.dot(a, a)
        
time.sleep(5)
        
with d2l.Benchmark('torch'):
    for _ in range(10):
        a = torch.randn(size=(1000, 1000), device=device)
        b = torch.mm(a, a)

numpy: 0.9414 sec
torch: 0.2103 sec

經過多次嘗試，可以發現 torch 的執行時間有明顯差異，這說明有 numpy 有部分仍然佔用裝置的時候，已經開始對 torch 的矩陣乘法計時了。

而如果把這兩個矩陣乘法的順序反過來，numpy 的時間變化不大，因此 torch 幾乎沒有非同步而 numpy 非同步了。

最後，我發現 torch.cuda.synchronize() 直接呼叫不加引數就不會報錯了。如果它的 device 引數為 None，那麼它將使用 current_device 函式找出當前裝置。

12.3 自動並行

深度學習框架會在後端自動構建計算圖。利用計算圖，系統可以瞭解所有依賴關係，並且可以選擇性地並行執行多個不相互依賴的任務以提高速度。

通常情況下單個運算子將使用所有CPU或單個GPU上的所有計算資源。並行化對單裝置計算機來說並不是很有用，而並行化對於多個裝置就很重要了。

請注意，接下來的實驗至少需要兩個GPU來執行。

12.3.1 基於 GPU 的平行計算

測試一下兩個 GPU 序列各執行 10 次矩陣乘法和並行各執行 10 次矩陣乘法的速度。

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])

run(x_gpu1)
run(x_gpu2)  # 預熱裝置
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: 1.5491 sec
GPU2 time: 1.4804 sec

刪除兩個任務之間的 torch.cuda.synchronize() 語句，系統就可以在兩個裝置上自動實現平行計算。

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

GPU1 & GPU2: 1.5745 sec

12.3.2 平行計算與通訊

在許多情況下，我們需要在不同的裝置之間移動資料，比如在CPU和GPU之間，或者在不同的GPU之間。

透過在 GPU 上計算，然後將結果複製回 CPU 來模擬這個過程。

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

with d2l.Benchmark('在GPU1上執行'):
    y = run(x_gpu1)
    torch.cuda.synchronize()

with d2l.Benchmark('複製到CPU'):
    y_cpu = copy_to_cpu(y)
    torch.cuda.synchronize()

在GPU1上執行: 1.6285 sec
複製到CPU: 2.5801 sec

在 GPU 仍在執行時就開始使用 PCI-Express 匯流排頻寬來移動資料是有利的。在 PyTorch 中，to() 和 copy_() 等函式都允許顯式的 non_blocking 引數，這允許在不需要同步時呼叫方可以繞過同步。設定 non_blocking=True 以模擬這個場景。

with d2l.Benchmark('在GPU1上執行並複製到CPU'):
    y = run(x_gpu1)
    y_cpu = copy_to_cpu(y, True)
    torch.cuda.synchronize()

在GPU1上執行並複製到CPU: 1.9456 sec

12.5 多 GPU 訓練

在多個 GPU 上並行總共分為三種：

  1. 網路並行
  2. 按層並行
  3. 資料並行

實際上，資料並行是最常用的方法。原書中也重點討論了資料並行。

12.5.2 資料並行性

假設一臺機器有 \(k\) 個 GPU。 給定需要訓練的模型，雖然每個 GPU 上的引數值都是相同且同步的，但是每個 GPU 都將獨立地維護一組完整的模型引數。

一般來說，\(k\) 個 GPU 並行訓練過程如下：

  • 在任何一次訓練迭代中，給定的隨機的小批次樣本都將被分成 \(k\) 個部分，並均勻地分配到 GPU 上；
  • 每個 GPU 根據分配給它的小批次子集，計算模型引數的損失和梯度；
  • 將 \(k\) 個 GPU 中的區域性梯度聚合，以獲得當前小批次的隨機梯度；
  • 聚合梯度被重新分發到每個 GPU 中；
  • 每個 GPU 使用這個小批次隨機梯度，來更新它所維護的完整的模型引數集。

在實踐中請注意，當在 \(k\) 個 GPU 上訓練時，需要擴大小批次的大小為 \(k\) 的倍數，這樣每個 GPU 都有相同的工作量，就像只在單個 GPU 上訓練一樣。 因此，在 16-GPU 伺服器上可以顯著地增加小批次資料量的大小，同時可能還需要相應地提高學習率。

12.5.4 資料同步

對於高效的多 GPU 訓練，我們需要兩個基本操作。首先，我們需要向多個裝置分發引數並附加梯度（get_params）。如果沒有引數，就不可能在 GPU 上評估網路。第二，需要跨多個裝置對引數求和，也就是說，需要一個 allreduce 函式。

get_params() 函式定義如下：

def get_params(params, device):
    new_params = [p.to(device) for p in params]
    for p in new_params:
        p.requires_grad_()
    return new_params

假設現在有一個向量分佈在多個 GPU 上，下面的 allreduce 函式將所有向量相加，並將結果廣播給所有GPU。請注意，需要將資料複製到累積結果的裝置，才能使函式正常工作。

def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].to(data[0].device)
    for i in range(1, len(data)):
        data[i][:] = data[0].to(data[i].device)

12.5.5 資料分發

nn.parallel.scatter() 是一個簡單的工具函式，將一個小批次資料均勻地分佈在多個 GPU 上。用法如下所示：

data = torch.arange(20).reshape(4, 5)
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
split = nn.parallel.scatter(data, devices)
print('input :', data)
print('load into', devices)
print('output:', split)

input : tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]])
load into [device(type='cuda', index=0), device(type='cuda', index=1)]
output: (tensor([[0, 1, 2, 3, 4],
        [5, 6, 7, 8, 9]], device='cuda:0'), tensor([[10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]], device='cuda:1'))

12.5.6 訓練

def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    # 在每個GPU上分別計算損失
    ls = [loss(lenet(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    for l in ls:  # 反向傳播在每個GPU上分別執行
        l.backward()
    # 將每個GPU的所有梯度相加，並將其廣播到所有GPU
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce(
                [device_params[c][i].grad for c in range(len(devices))])
    # 在每個GPU上分別更新模型引數
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0]) # 在這裡，我們使用全尺寸的小批次

與前幾章中略有不同：訓練函式需要分配 GPU 並將所有模型引數複製到所有裝置。顯然，每個小批次都是使用 train_batch 函式來處理多個 GPU。我們只在一個 GPU 上計算模型的精確度，而讓其他 GPU 保持空閒，儘管這是相對低效的，但是使用方便且程式碼簡潔。

def train(num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    # 將模型引數複製到num_gpus個GPU
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        for X, y in train_iter:
            # 為單個小批次執行多GPU訓練
            train_batch(X, y, device_params, devices, lr)
            torch.cuda.synchronize()
        timer.stop()
        # 在GPU0上評估模型
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
            lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
    print(f'測試精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/輪，'
          f'在{str(devices)}')

12.6 多 GPU 的簡潔實現

12.6.1 DataParallel()

原書出現了一個有趣的函式 net = nn.DataParallel(net, device_ids=devices)，這個函式可以說是本節的重點。

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0) 這個函式：

在模組的層級上實現了資料並行。

這個容器透過在批次維度中分塊將輸入拆分到指定裝置，從而並行化給定模組的應用程式（其他物件將在每個裝置上覆制一次）。在前向傳遞中，模組在每個裝置上被複制，每個副本處理一部分輸入。在向後傳遞期間，來自每個副本的梯度被彙總到原始模組中。

批次大小應大於使用的 GPU 數量。

另外，PyTorch 推薦使用 nn.parallel.DistributedDataParallel() 來代替 nn.Parallel()，原因如下：

大多數涉及批次輸入和多個 GPU 的用例應預設使用 DistributedDataParallel 來利用多個 GPU。

使用具有多處理功能的 CUDA 模型有一些重要的注意事項；除非注意準確地滿足資料處理要求，否則您的程式很可能會出現不正確或未定義的行為。

建議使用 DistributedDataParallel，而不是 DataParallel 進行多 GPU 訓練，即使只有一個裝置。

DistributedDataParallel 和 DataParallel 之間的區別是：DistributedDataParallel 使用多程式，其中為每個 GPU 建立一個程式，而 DataParallel 使用多執行緒。透過使用 multiprocessing，每個 GPU 都有自己的專用程式，這避免了 Python 直譯器的 GIL 帶來的效能開銷。

如果您使用 DistributedDataParallel，您可以使用 torch.distributed.launch 實用程式來啟動您的程式，請參閱第三方後端。

允許將任意位置和關鍵字輸入傳遞到 DataParallel 中，但某些型別需要特殊處理。張量將依託指定的維度被分開（預設為 0）。元組、列表和字典型別將被淺複製。其他型別將在不同的執行緒之間共享，如果寫入模型的正向傳播，則可能會被破壞。

在執行此 DataParallel 模組之前，並行化模組必須在 device_ids[0] 上具有其引數和緩衝區。原因在於：在每個 forward 中，模組在每個裝置上被複制，因此對 forward 中正在執行的模組的任何更新都將丟失。例如，如果模組有一個計數器屬性，在每次轉發時遞增，它將始終保持初始值，因為更新是在轉發後銷燬的副本上完成的。但是，DataParallel 保證 device[0] 上的副本的引數和緩衝區將與基本並行化模組共享儲存。因此，將記錄對裝置 [0] 上的引數或緩衝區的就地更新。例如，BatchNorm2d 和 spectral_norm() 依賴於此行為來更新緩衝區。

當模組在 forward() 中返回一個標量時，此 wrapper 將返回一個長度等於資料並行中使用的裝置數量的向量，其中包含每個裝置的結果。

引數：

  • module (Module) – 要並行的模組
  • device_ids (list of python:int or torch.device) – CUDA 裝置（預設：全部裝置）
  • output_device (int or torch.device) – 輸出的裝置位置（預設：device_ids[0]）

於是，原書中訓練這段程式碼寫成了這樣：

def train(net, num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    def init_weights(m):
        if type(m) in [nn.Linear, nn.Conv2d]:
            nn.init.normal_(m.weight, std=0.01)
    net.apply(init_weights)
    # 在多個GPU上設定模型
    net = nn.DataParallel(net, device_ids=devices)
    trainer = torch.optim.SGD(net.parameters(), lr)
    loss = nn.CrossEntropyLoss()
    timer, num_epochs = d2l.Timer(), 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    for epoch in range(num_epochs):
        net.train()
        timer.start()
        for X, y in train_iter:
            trainer.zero_grad()
            X, y = X.to(devices[0]), y.to(devices[0])
            l = loss(net(X), y)
            l.backward()
            trainer.step()
        timer.stop()
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
    print(f'測試精度：{animator.Y[0][-1]:.2f}，{timer.avg():.1f}秒/輪，'
          f'在{str(devices)}')

注意第 \(19\) 行中是把資料傳到了 \(0\) 號 GPU 上，然後它就會自動切成 GPU 個資料塊然後傳過去了。

執行程式碼測試一下！首先是隻使用 \(1\) 塊 GPU 的程式碼：

train(net, num_gpus=1, batch_size=256, lr=0.1)

測試精度：0.90，222.1秒/輪，在[device(type='cuda', index=0)]

然後是使用 \(2\) 塊 GPU 的程式碼：

train(net, num_gpus=2, batch_size=512, lr=0.2)

測試精度：0.87，111.8秒/輪，在[device(type='cuda', index=0), device(type='cuda', index=1)]

接近一倍的速度提升。原書中跑一輪居然只需要 \(10\) 秒左右，不禁令人感慨。

相關文章