峰值內(nèi)存消耗是訓(xùn)練深度學(xué)習(xí)模型（如視覺 Transformer 和 LLM）時的常見瓶頸。本文提供了一系列可以在不犧牲建模性能和預(yù)測精度的情況下，將 PyTorch 中的內(nèi)存消耗降低到約 1/20 的技術(shù)。

以 PyTorch 的 Torchvision 庫中的視覺 transformer 為基礎(chǔ)，本文作者編寫了大約 100 行代碼的訓(xùn)練腳本，并且所有代碼示例都可以在 GitHub 上找到。

以下是本文將要介紹的技術(shù)名稱：

(資料圖片)

微調(diào) vision transformer自動混合精度訓(xùn)練低精度訓(xùn)練Reduced Batch Size 訓(xùn)練梯度積累與 Microbatches選擇更精簡的優(yōu)化器在目標設(shè)備上實例化模型分布式訓(xùn)練與張量共享參數(shù)卸載以上九種方法結(jié)合起來，就形成了一種可以用于 llm 的綜合方法，也可以稱之為第十種方法。

這些方法是互相解耦的，可以將它們疊加在一起使用。

本文在實驗中使用的 ViT 為 ViT-L-16 模型。在依次將上述方法添加后，研究者將訓(xùn)練 BigBird-Roberta LLM 來執(zhí)行文本分類任務(wù)。這些技術(shù)使得在消費類硬件上訓(xùn)練這樣的模型成為可能。

微調(diào) vision transformer

為了簡化實驗中的 PyTorch 代碼，本文使用了開源庫 ——Fabric，十幾行代碼就能應(yīng)用各種先進的 PyTorc 技術(shù)（自動混合精度訓(xùn)練、多 GPU 訓(xùn)練、張量分片等）。

原生 PyTorch 代碼和修改后的使用 Fabric 的代碼之間的區(qū)別很微妙，只有較小的修改，如下面的代碼所示：

如上所述，改動雖然不大，但是可以方便的使用 PyTorch 中的高級功能，而無需重新構(gòu)造任何現(xiàn)有代碼。

總結(jié)上圖，將普通 PyTorch 代碼轉(zhuǎn)換為 PyTorch+Fabric 的主要 3 個步驟可以歸納如下：

導(dǎo)入 Fabric 并實例化 Fabric 對象。使用 Fabric 設(shè)置模型、優(yōu)化器和數(shù)據(jù)加載程序。調(diào)用 fabric.backward () 構(gòu)造損失函數(shù)，而不是通常使用的 loss.backward ()

使用普通 PyTorch 和 PyTorch with Fabric 的性能和內(nèi)存消耗幾乎完全相同：

Plain PyTorch (01_pytorch-vit.py):

Time elapsed 17.94 minMemory used: 26.79 GBTest accuracy 95.85%

PyTorch with Fabric (01-2_pytorch-fabric.py)

Time elapsed 17.88 minMemory used: 26.84 GBTest accuracy 96.06%

也可以將下面的代碼：

model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)

替換為：

model = vit_l_16(weights=None)

替換后，將不再是微調(diào)，而是從頭開始訓(xùn)練相同的 ViT 架構(gòu)，預(yù)測準確率會從 96% 以上下降到約 60%：

自動混合精度

上一節(jié)使用 Fabric 修改了 PyTorch 代碼，在此基礎(chǔ)上，使用混合精度和分布式訓(xùn)練，也只需更改一行代碼。

應(yīng)用混合精度訓(xùn)練

應(yīng)用混合精度訓(xùn)練，只需一個小的修改，將下面這行代碼

fabric = Fabric(accelerator="cuda", devices=1)

替換為：

fabric = Fabric(accelerator="cuda", devices=1, precisinotallow="16-mixed")

之后，在不犧牲預(yù)測精度的情況下，內(nèi)存消耗從 26.84GB 減少到 18.21GB，如下所示：

01-2_pytoch-fabric.py 和 02_mixed-precision.py 的結(jié)果對比

此外，混合精確訓(xùn)練不僅減少了內(nèi)存使用，還將運行時間減少了 6 倍（從 17.88 分鐘減少到 3.45 分鐘），這可以說是意外收獲。

什么是混合精度訓(xùn)練？

混合精度訓(xùn)練同時使用 16 位和 32 位精度，以確保不損失精度。16 位表示的梯度計算比 32 位格式快得多，并且節(jié)省了大量的內(nèi)存。這種策略是有益的，尤其是當(dāng)受到內(nèi)存或計算限制時。

之所以被稱為「混合」而不是「低」精度訓(xùn)練的原因是，并不會將所有參數(shù)和操作都轉(zhuǎn)移成 16 位浮點數(shù)。實際上，在訓(xùn)練期間會在 32 位和 16 位運算之間切換。

如下圖所示，混合精度訓(xùn)練可以分解為：將權(quán)重轉(zhuǎn)換為較低精度（如 FP16）以實現(xiàn)更快的計算、計算梯度、將梯度轉(zhuǎn)換回較高精度（FP32）以實現(xiàn)數(shù)值穩(wěn)定性，以及用縮放的梯度更新原始權(quán)重等幾個步驟。

這種方法在保證訓(xùn)練有效的前提下，還能保持神經(jīng)網(wǎng)絡(luò)的準確性和穩(wěn)定性。

感興趣的讀者還可以在本文作者的另一篇文章：《使用混合精度技術(shù)加速大型語言模型》中獲得更多底層概念。

文章地址：https://lightning.ai/pages/community/tutorial/accelerating-large-language-models-with-mixed-precision-techniques/

低精度訓(xùn)練

還可以更進一步，嘗試以「完全」較低的 16 位精度運行，而不是混合精度。

將下面這行代碼：

fabric = Fabric(accelerator="cuda", precisinotallow="16-mixed")

替換為

fabric = Fabric(accelerator="cuda", precisinotallow="16-true")

但需要注意的是，這樣會在訓(xùn)練中產(chǎn)生 NaN 值：

Epoch: 0001/0001 | Batch 0000/0703 | Loss: 2.4105Epoch: 0001/0001 | Batch 0300/0703 | Loss: nanEpoch: 0001/0001 | Batch 0600/0703 | Loss: nan...

這是因為常規(guī)的 16 位浮點只能表示 - 65504 和 65504 之間的數(shù)字：

In [1]: import torchIn [2]: torch.finfo(torch.float16)Out[2]: finfo(resolutinotallow=0.001, min=-65504, max=65504, eps=0.000976562, smallest_normal=6.10352e-05, tiny=6.10352e-05, dtype=float16)

因此，為了避免 NaN 問題，可以將參數(shù)修改為「bf16 true」：

fabric = Fabric(accelerator="cuda", precisinotallow="bf16-true")

可以將內(nèi)存消耗進一步降低到 13.82 GB（同樣，在不犧牲準確性的情況下）：

將 03_bfloat16.py 與之前的代碼的結(jié)果進行比較

什么是 Bfloat16？

「bf16 mixed」中的「bf16」代表 Brain Floating Point（bfloat16）。谷歌為機器學(xué)習(xí)和深度學(xué)習(xí)應(yīng)用程序開發(fā)了這種格式，特別是在其張量處理單元（TPU）中。與傳統(tǒng) float16 格式相比，Bfloat16 以降低精度為代價擴展了動態(tài)范圍。

擴展的動態(tài)范圍有助于 bfloat16 表示非常大和非常小的數(shù)字，使其更適合可能遇到廣泛值的深度學(xué)習(xí)應(yīng)用。然而，較低的精度可能會影響某些計算的準確性，或在某些情況下導(dǎo)致舍入誤差。但在大多數(shù)深度學(xué)習(xí)應(yīng)用中，這種精度的降低對建模性能的影響微乎其微。

雖然 bfloat16 最初是為 TPU 開發(fā)的，但這種格式從 A100 Tensor Core GPU 開始，也得到了其之后的 NVIDIA GPU 的支持。

以下代碼可以檢查 GPU 是否支持 bfloat16：

>>> import torch>>> torch.cuda.is_bf16_supported()True

減少批大小

減少批大小通常是減少內(nèi)存消耗的一個有效方法。然而，它有時會導(dǎo)致較差的預(yù)測性能，因為這樣要改變訓(xùn)練動態(tài)。

無論哪種方式，需要探討減少批量大小對結(jié)果有何影響。事實證明，可以在不犧牲性能的情況下將批大小降低到 16，從而將內(nèi)存消耗降至 5.69 GB：

將 04_lower-batchsize.py 與以前的代碼進行比較。

梯度積累與微批

梯度累積是一種在訓(xùn)練過程中虛擬增加批大小的方法，當(dāng)可用的 GPU 內(nèi)存不足以容納所需的批量大小時，這是非常有用的。并且這種方法只會在運行時產(chǎn)生影響，建模性能并不會受到影響。

梯度累積中，每批計算的量較小，并在多次迭代中累積梯度（通常求和或求平均），而不是在每個批次之后立刻更新模型權(quán)重。一旦累積的梯度達到目標「虛擬」批大小，模型權(quán)重就會用累積的梯度更新。

為了實現(xiàn)梯度積累，只需要對向前和向后傳球進行兩次小的修改：

05_gradient-acum.py 中的代碼修改

本文作者的另一篇文章《使用梯度累積在單個 GPU 上微調(diào) LLM》，更詳細地介紹了梯度累積的細節(jié)。

文章地址：https://lightning.ai/blog/gradient-accumulation/

有效批大小為 16，并且累積步數(shù)為 4，意味著實際批大小為 4（因為 16/4=4）。

05_gradient-acum.py 的結(jié)果

這種技術(shù)的缺點是運行時間從 3.96 分鐘增加到 12.91 分鐘。

值得注意的是，批大小最小可以減少到 1，進一步減少 75% 的內(nèi)存消耗。

使用更精簡的優(yōu)化器

時下流行的 Adam 優(yōu)化器其實附帶了額外的參數(shù)，例如，Adam 為每個模型參數(shù)提供了 2 個額外的優(yōu)化器參數(shù)（平均值和方差）。

因此，通過將 Adam 與 SGD 等無狀態(tài)優(yōu)化器進行交換，可以將參數(shù)數(shù)量減少 2/3，這在使用 ViT 和 LLM 時非常重要。

普通 SGD 的缺點是收斂性較差。因此，Adam 與 SGD 交換后，需要引入余弦衰減學(xué)習(xí)速率調(diào)度器來進行補償。

簡而言之，通過將以下代碼

optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)

替換為：

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)num_steps = NUM_EPOCHS * len(train_loader)scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(    optimizer, T_max=num_steps)

通過這種變化，模型能夠在保持大約 97% 分類準確率的同時減少峰值內(nèi)存消耗：

06_sgd-with-scheduler.py 的結(jié)果

在目標設(shè)備上創(chuàng)建模型

在 PyTorch 中實例化模型時，通常是首先在 CPU 設(shè)備上創(chuàng)建它，然后將它轉(zhuǎn)移到目標設(shè)備上，并將其轉(zhuǎn)換為所需的精度：

model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)model.cuda().float16()

但是 CPU 上生成完整精度的中間模型，是一種低效的方法。所以，可以使用 Fabric 中的 init_module 上下文在目標設(shè)備（例如 GPU）上直接創(chuàng)建所需精度的模型：

import lightning as Lfabric = Fabric(accelerator="cuda", devices=1, precisinotallow="16-true")with fabric.init_module():    model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)

在這種特定情況下（模型），前向通過期間的峰值內(nèi)存大于其全精度表示中的大小。對模型加載本身對 fabric.init_module 方法進行基準測試，結(jié)果如下：

沒有 init_module 的 GPU 峰值內(nèi)存：1.24 GB（07_01_init-module.py）GPU 帶 init_module 的峰值內(nèi)存：0.65 GB（07_03_init-module.py）

可以看到，在這種情況下，init_module 將模型加載的峰值內(nèi)存需求減少了 50%。

有關(guān) init_module 的更多詳細信息，可以參閱這篇關(guān)于大型模型的高效初始化的的文章。

文章地址：https://lightning.ai/pages/community/efficient-initialization-of-large-models/

分布式訓(xùn)練與張量共享

下一個修改是多 GPU 訓(xùn)練。多個 GPU 可供使用是有效的，因為這樣做可以更快地訓(xùn)練模型。

然而，本文探討的是內(nèi)存節(jié)省。因此，需要一種更先進的分布式多 GPU 策略，稱為完全共享數(shù)據(jù)并行（FSDP），該策略利用數(shù)據(jù)并行性和張量并行性在多個設(shè)備上共享大權(quán)重矩陣。

但是如果模型已經(jīng)很小了，例如將此技術(shù)添加到上面第 7 節(jié)的代碼中時，是幾乎看不到任何效果的。因此，為了純粹地關(guān)注分片的效果，可以與第 1 節(jié)中的全精度基線進行比較。

將以下代碼

fabric = Fabric(accelerator="cuda", devices=1)

替換為：

auto_wrap_policy = partial(    transformer_auto_wrap_policy,    transformer_layer_cls={EncoderBlock})    strategy = FSDPStrategy(    auto_wrap_policy=auto_wrap_policy,          activation_checkpointing=EncoderBlock)fabric = Fabric(accelerator="cuda", devices=4, strategy=strategy)

08_fsdp 與 - 01-2.py 的結(jié)果

除了手動定義，請也可以使用以下方法，自動確定要分割哪些層：

fabric = Fabric(accelerator="cuda", devices=4, strategy="fsdp")

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

在數(shù)據(jù)并行中，mini-batch 需要繼續(xù)被劃分，并且每個 GPU 上都有一份模型副本。由于多個 GPU 并行工作，能夠加快模型訓(xùn)練。

以下是工作原理：

在所有 GPU 中復(fù)制相同的模型。然后，每個 GPU 被饋送輸入數(shù)據(jù)的不同子集（不同的小批量）。所有 GPU 獨立地執(zhí)行模型的前向和后向傳遞，計算各自的局部梯度。然后，收集梯度并對所有 GPU 進行平均。然后使用平均梯度來更新模型的參數(shù)。

這種方法的主要優(yōu)點是速度塊。由于每個 GPU 都在與其他 GPU 同時處理一個獨特的小批量數(shù)據(jù)，因此可以在更短的時間內(nèi)在更多數(shù)據(jù)上訓(xùn)練模型。這可以顯著減少訓(xùn)練模型所需的時間，尤其是在使用大型數(shù)據(jù)集時。

然而，數(shù)據(jù)并行性有一些局限性。每個 GPU 必須具有模型及其參數(shù)的完整副本。這限制了可訓(xùn)練模型的大小，因為模型必須適合單個 GPU 的內(nèi)存 —— 這對于現(xiàn)代 ViT 或 LLM 來說是不可行的。

與數(shù)據(jù)并行不同，張量并行將模型本身劃分為 GPU。在數(shù)據(jù)并行中，每個 GPU 都需要適應(yīng)整個模型，這在訓(xùn)練更大的模型時可能會成為一個限制。然而，張量并行性允許通過分解模型并將其分布在多個設(shè)備上來訓(xùn)練對于單個 GPU 來說可能太大的模型。

具體來說，其原理和矩陣乘法相似。按行或按列都可以對模型進行拆解。簡單起見，以按列拆解為例，可以將一個大型矩陣乘法運算分解為單獨的計算，每個計算都可以在不同的 GPU 上執(zhí)行，如下圖所示。然后將結(jié)果連接起來以獲得原始結(jié)果，從而有效地分配了計算負載。

參數(shù)卸載

除了上一節(jié)中解釋的 FSDP 策略之外，還可以將優(yōu)化器參數(shù)卸載到 CPU，可以通過將以下代碼

strategy = FSDPStrategy(    auto_wrap_policy=auto_wrap_policy,          activation_checkpointing=EncoderBlock,)

替換為：

strategy = FSDPStrategy(    auto_wrap_policy=auto_wrap_policy,          activation_checkpointing=EncoderBlock,    cpu_offload=True)

內(nèi)存消耗從 6.59 GB 減少到 6.03 GB：

09_fsdp-cpu-offload-with-01-2.py 的結(jié)果。

美中不足的小缺點是運行時間從 5.5 分鐘增加到了 8.3 分鐘。

將前面幾招連著打出，就成為了最強的降龍十八掌最后一掌！

前幾節(jié)對優(yōu)化 ViT 進行了大量介紹，其實這些技術(shù)也同樣適用于 LLM。

作者在 Lit LLaMA 和 Lit GPT 存儲庫中使用了許多這些技巧，這些存儲庫支持 LLaMA、Falcon、Pythia 和其他流行的模型。盡管如此，為了創(chuàng)建一個更通用的例子，作者從流行的 HF transformers 庫中微調(diào) LLM，用于對 IMDb 電影評論的情緒進行分類。

使用上述技術(shù)，僅使用 1.15 Gb 內(nèi)存（bonus_DistilBERT-after.py）而不是 3.99 Gb（bonus_bigbird-before.py）就可以訓(xùn)練 DistilBERT 分類器。更令人印象深刻的是，通過將這些技術(shù)應(yīng)用于 transformers 庫中的 BigBird 模型，BigBird 僅消耗 4.03 GB（bonus_BigBird-after.py）。

strategy = FSDPStrategy(        cpu_offload=True   )   fabric = Fabric(        accelerator="cuda",        devices=4,        strategy=strategy,         precision="bf16-true"   )   with fabric.init_module():       model = AutoModelForSequenceClassification.from_pretrained(            "google/bigbird-roberta-base", num_labels=2)

結(jié)論

本文展示了 9 種減少 PyTorch 模型內(nèi)存消耗的技術(shù)。當(dāng)將這些技術(shù)應(yīng)用于 ViT 時，單個 GPU 上減少了 20 倍的內(nèi)存消耗?？梢钥吹?，跨 GPU 的張量分片甚至可以降低內(nèi)存消耗。同樣的優(yōu)化還使 BigBird LLM 能夠僅使用 4GB 峰值 GPU RAM 進行訓(xùn)練。

這些技術(shù)都不是特定于模型的，可以與任何 PyTorch 訓(xùn)練腳本一起使用。使用開源 Fabric 庫，大多數(shù)優(yōu)化都可以通過一行代碼實現(xiàn)。

關(guān)鍵詞：