用對了方法，加速 PyTorch 訓練，有時也不是那么復雜。

近日，深度學習領域知名研究者、Lightning AI 的首席人工智能教育者 Sebastian Raschka 在 CVPR 2023 上發(fā)表了主題演講「Scaling PyTorch Model Training With Minimal Code Changes」。

為了能與更多人分享研究成果，Sebastian Raschka 將演講整理成一篇文章。文章探討了如何在最小代碼更改的情況下擴展 PyTorch 模型訓練，并表明重點是利用混合精度（mixed-precision）方法和多 GPU 訓練模式，而不是低級機器優(yōu)化。

文章使用視覺 Transformer（ViT）作為基礎模型，ViT 模型在一個基本數(shù)據(jù)集上從頭開始，經(jīng)過約 60 分鐘的訓練，在測試集上取得了 62% 的準確率。

GitHub 地址：https://github.com/rasbt/cvpr2023

以下是文章原文：

構建基準

在接下來的部分中，Sebastian 將探討如何在不進行大量代碼重構的情況下改善訓練時間和準確率。

想要注意的是，模型和數(shù)據(jù)集的詳細信息并不是這里的主要關注點（它們只是為了盡可能簡單，以便讀者可以在自己的機器上復現(xiàn)，而不需要下載和安裝太多的依賴）。所有在這里分享的示例都可以在 GitHub 找到，讀者可以探索和重用完整的代碼。

腳本 00_pytorch-vit-random-init.py 的輸出。

不要從頭開始訓練

現(xiàn)如今，從頭開始訓練文本或圖像的深度學習模型通常是低效的。我們通常會利用預訓練模型，并對模型進行微調(diào)，以節(jié)省時間和計算資源，同時獲得更好的建模效果。

如果考慮上面使用的相同 ViT 架構，在另一個數(shù)據(jù)集（ImageNet）上進行預訓練，并對其進行微調(diào)，就可以在更短的時間內(nèi)實現(xiàn)更好的預測性能：20 分鐘（3 個訓練 epoch）內(nèi)達到 95% 的測試準確率。

00_pytorch-vit-random-init.py 和 01_pytorch-vit.py 的對比。

提升計算性能

我們可以看到，相對于從零開始訓練，微調(diào)可以大大提升模型性能。下面的柱狀圖總結(jié)了這一點。

00_pytorch-vit-random-init.py 和 01_pytorch-vit.py 的對比柱狀圖。

當然，模型效果可能因數(shù)據(jù)集或任務的不同而有所差異。但對于許多文本和圖像任務來說，從一個在通用公共數(shù)據(jù)集上預訓練的模型開始是值得的。

接下來的部分將探索各種技巧，以加快訓練時間，同時又不犧牲預測準確性。

開源庫 Fabric

在 PyTorch 中以最小代碼更改來高效擴展訓練的一種方法是使用開源 Fabric 庫，它可以看作是 PyTorch 的一個輕量級包裝庫 / 接口。通過 pip 安裝。

pip install lightning

下面探索的所有技術也可以在純 PyTorch 中實現(xiàn)。Fabric 的目標是使這一過程更加便利。

在探索「加速代碼的高級技術」之前，先介紹一下將 Fabric 集成到 PyTorch 代碼中需要進行的小改動。一旦完成這些改動，只需要改變一行代碼，就可以輕松地使用高級 PyTorch 功能。

PyTorch 代碼和修改后使用 Fabric 的代碼之間的區(qū)別是微小的，只涉及到一些細微的修改，如下面的代碼所示：

普通 PyTorch 代碼（左）和使用 Fabric 的 PyTorch 代碼

總結(jié)一下上圖，就可以得到普通的 PyTorch 代碼轉(zhuǎn)換為 PyTorch+Fabric 的三個步驟：

導入 Fabric 并實例化一個 Fabric 對象。

使用 Fabric 設置模型、優(yōu)化器和 data loader。

損失函數(shù)使用 fabric.backward ()，而不是 loss.backward ()。

這些微小的改動提供了一種利用 PyTorch 高級特性的途徑，而無需對現(xiàn)有代碼進行進一步重構。

深入探討下面的「高級特性」之前，要確保模型的訓練運行時間、預測性能與之前相同。

01_pytorch-vit.py 和 03_fabric-vit.py 的比較結(jié)果。

正如前面柱狀圖中所看到的，訓練運行時間、準確率與之前完全相同，正如預期的那樣。其中，任何波動都可以歸因于隨機性。

在前面的部分中，我們使用 Fabric 修改了 PyTorch 代碼。為什么要費這么大的勁呢？接下來將嘗試高級技術，比如混合精度和分布式訓練，只需更改一行代碼，把下面的代碼

fabric = Fabric(accelerator="cuda")

改為

fabric=Fabric(accelerator="cuda",precision="bf16-mixed")

04_fabric-vit-mixed-precision.py 腳本的比較結(jié)果。腳本地址：https://github.com/rasbt/cvpr2023/blob/main/04_fabric-vit-mixed-precision.py

通過混合精度訓練，我們將訓練時間從 18 分鐘左右縮短到 6 分鐘，同時保持相同的預測性能。這種訓練時間的縮短只需在實例化 Fabric 對象時添加參數(shù)「precision="bf16-mixed"」即可實現(xiàn)。

理解混合精度機制

混合精度訓練實質(zhì)上使用了 16 位和 32 位精度，以確保不會損失準確性。16 位表示中的計算梯度比 32 位格式快得多，并且還節(jié)省了大量內(nèi)存。這種策略在內(nèi)存或計算受限的情況下非常有益。

之所以稱為「混合」而不是「低」精度訓練，是因為不是將所有參數(shù)和操作轉(zhuǎn)換為 16 位浮點數(shù)。相反，在訓練過程中 32 位和 16 位操作之間切換，因此稱為「混合」精度。

如下圖所示，混合精度訓練涉及步驟如下：

將權重轉(zhuǎn)換為較低精度（FP16）以加快計算速度；

計算梯度；

將梯度轉(zhuǎn)換回較高精度（FP32）以保持數(shù)值穩(wěn)定性；

使用縮放后的梯度更新原始權重。

這種方法在保持神經(jīng)網(wǎng)絡準確性和穩(wěn)定性的同時，實現(xiàn)了高效的訓練。

更詳細的步驟如下：

將權重轉(zhuǎn)換為 FP16：在這一步中，神經(jīng)網(wǎng)絡的權重（或參數(shù)）初始時用 FP32 格式表示，將其轉(zhuǎn)換為較低精度的 FP16 格式。這樣可以減少內(nèi)存占用，并且由于 FP16 操作所需的內(nèi)存較少，可以更快地被硬件處理。

計算梯度：使用較低精度的 FP16 權重進行神經(jīng)網(wǎng)絡的前向傳播和反向傳播。這一步計算損失函數(shù)相對于網(wǎng)絡權重的梯度（偏導數(shù)），這些梯度用于在優(yōu)化過程中更新權重。

將梯度轉(zhuǎn)換回 FP32：在計算得到 FP16 格式的梯度后，將其轉(zhuǎn)換回較高精度的 FP32 格式。這種轉(zhuǎn)換對于保持數(shù)值穩(wěn)定性非常重要，避免使用較低精度算術時可能出現(xiàn)的梯度消失或梯度爆炸等問題。

乘學習率并更新權重：以 FP32 格式表示的梯度乘以學習率將用于更新權重（標量值，用于確定優(yōu)化過程中的步長）。

步驟 4 中的乘積用于更新原始的 FP32 神經(jīng)網(wǎng)絡權重。學習率有助于控制優(yōu)化過程的收斂性，對于實現(xiàn)良好的性能非常重要。

Brain Float 16

前面談到了「float 16-bit」精度訓練。需要注意的是，在之前的代碼中，指定了 precision="bf16-mixed"，而不是 precision="16-mixed"。這兩個都是有效的選項。

在這里，"bf16-mixed" 中的「bf16」表示 Brain Floating Point（bfloat16）。谷歌開發(fā)了這種格式，用于機器學習和深度學習應用，尤其是在張量處理單元（TPU）中。Bfloat16 相比傳統(tǒng)的 float16 格式擴展了動態(tài)范圍，但犧牲了一定的精度。

擴展的動態(tài)范圍使得 bfloat16 能夠表示非常大和非常小的數(shù)字，使其更適用于深度學習應用中可能遇到的數(shù)值范圍。然而，較低的精度可能會影響某些計算的準確性，或在某些情況下導致舍入誤差。但在大多數(shù)深度學習應用中，這種降低的精度對建模性能的影響很小。

雖然 bfloat16 最初是為 TPU 開發(fā)的，但從 NVIDIA Ampere 架構的 A100 Tensor Core GPU 開始，已經(jīng)有幾種 NVIDIA GPU 開始支持 bfloat16。

我們可以使用下面的代碼檢查 GPU 是否支持 bfloat16：

>>> torch.cuda.is_bf16_supported() True

如果你的 GPU 不支持 bfloat16，可以將 precision="bf16-mixed" 更改為 precision="16-mixed"。

多 GPU 訓練和完全分片數(shù)據(jù)并行

接下來要嘗試修改多 GPU 訓練。如果我們有多個 GPU 可供使用，這會帶來好處，因為它可以讓我們的模型訓練速度更快。

這里介紹一種更先進的技術 — 完全分片數(shù)據(jù)并行（Fully Sharded Data Parallelism (FSDP)），它同時利用了數(shù)據(jù)并行性和張量并行性。

在 Fabric 中，我們可以通過下面的方式利用 FSDP 添加設備數(shù)量和多 GPU 訓練策略：

fabric = Fabric( accelerator="cuda", precision="bf16-mixed", devices=4, strategy="FSDP" # new! )

06_fabric-vit-mixed-fsdp.py 腳本的輸出。

現(xiàn)在使用 4 個 GPU，我們的代碼運行時間大約為 2 分鐘，是之前僅使用混合精度訓練時的近 3 倍。

理解數(shù)據(jù)并行和張量并行

在數(shù)據(jù)并行中，小批量數(shù)據(jù)被分割，并且每個 GPU 上都有模型的副本。這個過程通過多個 GPU 的并行工作來加速模型的訓練速度。

如下簡要概述了數(shù)據(jù)并行的工作原理：

同一個模型被復制到所有的 GPU 上。

每個 GPU 分別接收不同的輸入數(shù)據(jù)子集（不同的小批量數(shù)據(jù)）。

所有的 GPU 獨立地對模型進行前向傳播和反向傳播，計算各自的局部梯度。

收集并對所有 GPU 的梯度求平均值。

平均梯度被用于更新模型的參數(shù)。

每個 GPU 都在并行地處理不同的數(shù)據(jù)子集，通過梯度的平均化和參數(shù)的更新，整個模型的訓練過程得以加速。

這種方法的主要優(yōu)勢是速度。由于每個 GPU 同時處理不同的小批量數(shù)據(jù)，模型可以在更短的時間內(nèi)處理更多的數(shù)據(jù)。這可以顯著減少訓練模型所需的時間，特別是在處理大型數(shù)據(jù)集時。

然而，數(shù)據(jù)并行也有一些限制。最重要的是，每個 GPU 必須具有完整的模型和參數(shù)副本。這限制了可以訓練的模型大小，因為模型必須適應單個 GPU 的內(nèi)存。這對于現(xiàn)代的 ViTs 或 LLMs 來說這是不可行的。

與數(shù)據(jù)并行不同，張量并行將模型本身劃分到多個 GPU 上。并且在數(shù)據(jù)并行中，每個 GPU 都需要適 應整個模型，這在訓練較大的模型時可能成為一個限制。而張量并行允許訓練那些對單個 GPU 而言可能過大的模型，通過將模型分解并分布到多個設備上進行訓練。

張量并行是如何工作的呢？想象一下矩陣乘法，有兩種方式可以進行分布計算 —— 按行或按列。為了簡單起見，考慮按列進行分布計算。例如，我們可以將一個大型矩陣乘法操作分解為多個獨立的計算，每個計算可以在不同的 GPU 上進行，如下圖所示。然后將結(jié)果連接起來以獲取結(jié)果，這有效地分攤了計算負載。