注意
點擊這裡以下載完整的範例程式碼
如何透過將優化器步驟融合到反向傳播中來節省記憶體¶
建立於：2023 年 10 月 02 日 | 最後更新：2024 年 1 月 16 日 | 最後驗證：2024 年 11 月 05 日
您好！本教學旨在展示一種透過減少梯度所佔用的記憶體，來減少訓練迴圈記憶體佔用量的方法。假設您有一個模型，並且您對最佳化記憶體以避免 Out of Memory (OOM) 錯誤，或僅僅是從您的 GPU 中榨取更多效能感興趣。那麼，您_可能_很幸運（如果梯度佔用您記憶體的一部分，並且您不需要進行梯度累積）。我們將探索以下內容：
在您的訓練或微調迴圈期間，什麼佔用記憶體，
如何捕獲和視覺化記憶體快照以確定瓶頸，
新的
Tensor.register_post_accumulate_grad_hook(hook)API，以及，所有內容如何以 10 行程式碼組合在一起以實現記憶體節省。
要執行本教學，您需要
PyTorch 2.1.0 或更新版本，且包含
torchvision如果您想在本機執行記憶體視覺化，則需要 1 個 CUDA GPU。否則，這種技術在任何裝置上都會有類似的好處。
讓我們從匯入所需的模組和模型開始。我們將使用 torchvision 中的 vision transformer 模型，但您可以隨意替換為您自己的模型。我們也將使用 torch.optim.Adam 作為我們的優化器，但同樣地，您可以隨意替換為您自己的優化器。
import torch
from torchvision import models
from pickle import dump
model = models.vit_l_16(weights='DEFAULT').cuda()
optimizer = torch.optim.Adam(model.parameters())
Downloading: "https://download.pytorch.org/models/vit_l_16-852ce7e3.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/vit_l_16-852ce7e3.pth
0%| | 0.00/1.13G [00:00<?, ?B/s]
1%| | 7.88M/1.13G [00:00<00:14, 82.2MB/s]
2%|1 | 22.2M/1.13G [00:00<00:09, 122MB/s]
3%|3 | 39.8M/1.13G [00:00<00:07, 150MB/s]
5%|4 | 57.4M/1.13G [00:00<00:07, 163MB/s]
6%|6 | 75.0M/1.13G [00:00<00:06, 171MB/s]
8%|7 | 92.6M/1.13G [00:00<00:06, 175MB/s]
9%|9 | 110M/1.13G [00:00<00:06, 178MB/s]
11%|#1 | 128M/1.13G [00:00<00:06, 180MB/s]
13%|#2 | 145M/1.13G [00:00<00:05, 180MB/s]
14%|#3 | 162M/1.13G [00:01<00:05, 179MB/s]
15%|#5 | 180M/1.13G [00:01<00:05, 179MB/s]
17%|#6 | 197M/1.13G [00:01<00:05, 179MB/s]
18%|#8 | 214M/1.13G [00:01<00:06, 160MB/s]
20%|#9 | 231M/1.13G [00:01<00:05, 165MB/s]
21%|##1 | 248M/1.13G [00:01<00:05, 170MB/s]
23%|##2 | 266M/1.13G [00:01<00:05, 173MB/s]
24%|##4 | 283M/1.13G [00:01<00:05, 175MB/s]
26%|##5 | 300M/1.13G [00:01<00:05, 177MB/s]
27%|##7 | 318M/1.13G [00:01<00:04, 181MB/s]
29%|##8 | 336M/1.13G [00:02<00:04, 183MB/s]
31%|### | 354M/1.13G [00:02<00:04, 184MB/s]
32%|###2 | 372M/1.13G [00:02<00:04, 186MB/s]
34%|###3 | 390M/1.13G [00:02<00:04, 186MB/s]
35%|###5 | 408M/1.13G [00:02<00:04, 187MB/s]
37%|###6 | 426M/1.13G [00:02<00:04, 187MB/s]
38%|###8 | 444M/1.13G [00:02<00:04, 187MB/s]
40%|###9 | 462M/1.13G [00:02<00:03, 188MB/s]
41%|####1 | 480M/1.13G [00:02<00:03, 188MB/s]
43%|####2 | 498M/1.13G [00:02<00:03, 188MB/s]
44%|####4 | 516M/1.13G [00:03<00:03, 188MB/s]
46%|####6 | 534M/1.13G [00:03<00:03, 188MB/s]
48%|####7 | 552M/1.13G [00:03<00:03, 188MB/s]
49%|####9 | 570M/1.13G [00:03<00:03, 188MB/s]
51%|##### | 588M/1.13G [00:03<00:03, 188MB/s]
52%|#####2 | 607M/1.13G [00:03<00:03, 189MB/s]
54%|#####3 | 625M/1.13G [00:03<00:02, 189MB/s]
55%|#####5 | 643M/1.13G [00:03<00:02, 189MB/s]
57%|#####6 | 661M/1.13G [00:03<00:02, 188MB/s]
58%|#####8 | 679M/1.13G [00:03<00:02, 186MB/s]
60%|#####9 | 696M/1.13G [00:04<00:02, 185MB/s]
61%|######1 | 714M/1.13G [00:04<00:02, 184MB/s]
63%|######3 | 732M/1.13G [00:04<00:02, 183MB/s]
65%|######4 | 749M/1.13G [00:04<00:02, 182MB/s]
66%|######6 | 767M/1.13G [00:04<00:02, 182MB/s]
68%|######7 | 784M/1.13G [00:04<00:02, 182MB/s]
69%|######9 | 801M/1.13G [00:04<00:02, 182MB/s]
71%|####### | 819M/1.13G [00:04<00:01, 181MB/s]
72%|#######2 | 836M/1.13G [00:04<00:01, 181MB/s]
74%|#######3 | 854M/1.13G [00:04<00:01, 182MB/s]
75%|#######5 | 871M/1.13G [00:05<00:01, 181MB/s]
77%|#######6 | 888M/1.13G [00:05<00:01, 182MB/s]
78%|#######8 | 906M/1.13G [00:05<00:01, 182MB/s]
79%|#######9 | 923M/1.13G [00:05<00:01, 181MB/s]
81%|######## | 940M/1.13G [00:05<00:01, 181MB/s]
82%|########2 | 958M/1.13G [00:05<00:01, 181MB/s]
84%|########3 | 975M/1.13G [00:05<00:01, 181MB/s]
85%|########5 | 992M/1.13G [00:05<00:00, 181MB/s]
87%|########6 | 0.99G/1.13G [00:05<00:00, 181MB/s]
88%|########8 | 1.00G/1.13G [00:05<00:00, 181MB/s]
90%|########9 | 1.02G/1.13G [00:06<00:00, 181MB/s]
92%|#########1| 1.04G/1.13G [00:06<00:00, 183MB/s]
93%|#########3| 1.06G/1.13G [00:06<00:00, 185MB/s]
95%|#########4| 1.07G/1.13G [00:06<00:00, 186MB/s]
96%|#########6| 1.09G/1.13G [00:06<00:00, 186MB/s]
98%|#########7| 1.11G/1.13G [00:06<00:00, 187MB/s]
99%|#########9| 1.13G/1.13G [00:06<00:00, 187MB/s]
100%|##########| 1.13G/1.13G [00:06<00:00, 181MB/s]
現在讓我們定義典型的訓練迴圈。您應該在訓練時使用真實圖像，但為了本教學的目的，我們將傳入虛假的輸入，而不擔心載入任何實際資料。
IMAGE_SIZE = 224
def train(model, optimizer):
# create our fake image input: tensor shape is batch_size, channels, height, width
fake_image = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE).cuda()
# call our forward and backward
loss = model.forward(fake_image)
loss.sum().backward()
# optimizer update
optimizer.step()
optimizer.zero_grad()
訓練期間的記憶體使用情況¶
我們即將查看一些記憶體快照，因此我們應該準備好正確分析它們。通常，訓練記憶體包含：
模型參數（大小 P）
為反向傳播儲存的激活值（大小 A）
梯度，其大小與模型參數相同，因此大小 G = P。
優化器狀態，其與參數的大小成正比。在本例中，Adam 的狀態需要 2 倍的模型參數，因此大小 O = 2P。
中間張量，其在整個計算過程中分配。我們暫時不擔心它們，因為它們通常很小且是暫時的。
捕獲和視覺化記憶體快照¶
讓我們取得一個記憶體快照！當您的程式碼執行時，考慮一下您期望 CUDA 記憶體時間軸的樣子。
# tell CUDA to start recording memory allocations
torch.cuda.memory._record_memory_history(enabled='all')
# train 3 steps
for _ in range(3):
train(model, optimizer)
# save a snapshot of the memory allocations
s = torch.cuda.memory._snapshot()
with open(f"snapshot.pickle", "wb") as f:
dump(s, f)
# tell CUDA to stop recording memory allocations now
torch.cuda.memory._record_memory_history(enabled=None)
現在，透過拖放 snapshot.pickle 檔案，在 https://pytorch.dev.org.tw/memory_viz 的 CUDA 記憶體視覺化工具中開啟快照。記憶體時間軸是否符合您的期望？
模型參數已經在訓練步驟之前載入到記憶體中，因此我們在一開始就看到一塊記憶體專用於權重。當我們開始正向傳播時，會逐漸為激活值分配記憶體，或者我們儲存的張量以便能夠在反向傳播中計算梯度。一旦我們開始反向傳播，激活值會逐漸釋放，同時梯度的記憶體開始建立。
最後，當優化器啟動時，其狀態將被延遲初始化，因此我們應該只在第一個訓練迴圈的優化器步驟中看到優化器狀態記憶體逐漸增加。在未來的迴圈中，優化器記憶體將保持不變並就地更新。然後，在每次訓練迴圈結束時呼叫 zero_grad 時，梯度的記憶體會相應地釋放。
此訓練迴圈中的記憶體瓶頸在哪裡？或者，換句話說，峰值記憶體在哪裡？
記憶體使用量的高峰出現在優化器步驟中！請注意，此時的記憶體組成如預期：約 1.2GB 的參數、約 1.2GB 的梯度，以及約 2.4GB = 2 * 1.2GB 的優化器狀態。最後的約 1.2GB 來自 Adam 優化器，需要記憶體來儲存中間值，總共達到約 6GB 的記憶體使用高峰。從技術上講，如果您設定 Adam(model.parameters(), foreach=False)，則可以消除對最後 1.2GB 優化器中間值的需求，這會犧牲執行時間來換取記憶體。如果關閉 foreach 執行時間優化足以節省記憶體，那很好，但如果您好奇本教學課程如何幫助您做得更好，請繼續閱讀！透過我們即將介紹的技術，我們將透過消除對約 1.2GB 的梯度記憶體以及優化器中間值記憶體的需求來降低記憶體使用高峰。現在，您預期新的記憶體使用高峰會是多少？答案將在下一個快照中揭曉。
免責聲明：此技術不適用於所有人¶
在我們過於興奮之前，我們必須考慮此技術是否適用於您的使用案例。這不是萬靈丹！將優化器步驟融合到反向傳播中的技術，僅針對減少梯度記憶體（並且作為副作用也減少優化器中間值記憶體）。因此，梯度佔用的記憶體越大，記憶體減少就越重要。在上面的例子中，梯度佔據了 20% 的記憶體空間，這相當可觀！
對於您來說，情況可能並非如此。例如，如果您的權重已經很小（例如，由於應用了 LoRa），那麼梯度在您的訓練迴圈中不會佔用太多空間，並且收益會少得多。在這種情況下，您應該首先嘗試其他技術，例如啟用檢查點、分散式訓練、量化或減少批次大小。然後，當梯度再次成為瓶頸時，再回到本教學課程！
還在這裡嗎？太好了，讓我們介紹 Tensor 上新的 register_post_accumulate_grad_hook(hook) API。
Tensor.register_post_accumulate_grad_hook(hook) API 和我們的技術¶
我們的技術依賴於在 backward() 期間不必儲存梯度。相反，一旦梯度被累積，我們將立即將優化器應用於相應的參數並完全丟棄該梯度！這消除了在優化器步驟之前保留大量梯度緩衝區的需求。
那麼我們如何解鎖更積極地應用優化器的行為呢？在我們的 2.1 版本中，我們新增了一個新的 API torch.Tensor.register_post_accumulate_grad_hook()，它允許我們在 Tensor 的 .grad 欄位被累積後，將一個 hook 添加到該 Tensor 上。我們將把優化器步驟封裝到這個 hook 中。怎麼做？
如何在 10 行程式碼中將所有內容組合在一起¶
還記得我們從一開始的模型和優化器設定嗎？我將把它們註釋掉，這樣我們就不會浪費資源重新執行程式碼。
model = models.vit_l_16(weights='DEFAULT').cuda()
optimizer = torch.optim.Adam(model.parameters())
# Instead of having just *one* optimizer, we will have a ``dict`` of optimizers
# for every parameter so we could reference them in our hook.
optimizer_dict = {p: torch.optim.Adam([p], foreach=False) for p in model.parameters()}
# Define our hook, which will call the optimizer ``step()`` and ``zero_grad()``
def optimizer_hook(parameter) -> None:
optimizer_dict[parameter].step()
optimizer_dict[parameter].zero_grad()
# Register the hook onto every parameter
for p in model.parameters():
p.register_post_accumulate_grad_hook(optimizer_hook)
# Now remember our previous ``train()`` function? Since the optimizer has been
# fused into the backward, we can remove the optimizer step and zero_grad calls.
def train(model):
# create our fake image input: tensor shape is batch_size, channels, height, width
fake_image = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE).cuda()
# call our forward and backward
loss = model.forward(fake_image)
loss.sum().backward()
# optimizer update --> no longer needed!
# optimizer.step()
# optimizer.zero_grad()
這在我們的範例模型中只花了約 10 行程式碼的修改，這很棒。然而，對於真實模型，將優化器替換為優化器字典可能是一個相當侵入性的變更，特別是對於那些使用``LRScheduler``s 或在整個訓練週期中操作優化器配置的人來說。使用這些變更設計此 API 將更加複雜，並且可能需要將更多配置移動到全域狀態，但應該不是不可能的。也就是說，PyTorch 的下一步是使此 API 更容易與 LRScheduler 和您已經習慣的其他功能一起採用。
但讓我回到說服您這項技術值得的理由。我們將諮詢我們的朋友，記憶體快照。
# delete optimizer memory from before to get a clean slate for the next
# memory snapshot
del optimizer
# tell CUDA to start recording memory allocations
torch.cuda.memory._record_memory_history(enabled='all')
# train 3 steps. note that we no longer pass the optimizer into train()
for _ in range(3):
train(model)
# save a snapshot of the memory allocations
s = torch.cuda.memory._snapshot()
with open(f"snapshot-opt-in-bwd.pickle", "wb") as f:
dump(s, f)
# tell CUDA to stop recording memory allocations now
torch.cuda.memory._record_memory_history(enabled=None)
是的，花一些時間將您的快照拖曳到 CUDA 記憶體視覺化工具中。
- 幾個主要的觀察結果
不再有優化器步驟了！是的...我們將其融合到反向傳播中。
同樣，反向傳播拖得更長，並且有更多隨機分配用於中間值。這是預期的，因為優化器步驟需要中間值。
最重要的是！記憶體使用高峰更低了！現在約為 4GB（我希望這與您之前的預期非常接近）。
請注意，與之前相比，不再有任何大塊記憶體分配給梯度，從而節省了約 1.2GB 的記憶體。相反，我們在計算出每個梯度後，透過將優化器步驟盡可能提前來非常快速地釋放了它們。萬歲！順便說一句，另一個約 1.2GB 的記憶體節省來自於將優化器分解為每個參數的優化器，因此中間值也相應地縮小了。這個細節不如梯度記憶體節省重要，因為您可以透過僅關閉 foreach=False 來獲得優化器中間值節省，而無需使用此技術。
您可能會正確地想知道：如果我們節省了 2.4GB 的記憶體，為什麼記憶體使用高峰不是 6GB - 2.4GB = 3.6GB？嗯，高峰已經移動了！現在，高峰出現在反向傳播步驟的開始附近，當我們仍然在記憶體中有激活時，而在之前，高峰出現在激活已被釋放的優化器步驟中。約 0.4GB 的差異，即約 4.0GB - 約 3.6GB，因此是由於激活記憶體。因此，人們可以想像，此技術可以與啟用檢查點相結合，以獲得更多的記憶體優勢。
結論¶
在本教學課程中，我們學習了透過新的 Tensor.register_post_accumulate_grad_hook() API 將優化器融合到反向傳播步驟中的記憶體節省技術，以及何時應用此技術（當梯度記憶體非常重要時）。一路上，我們還學習了記憶體快照，它通常在記憶體最佳化中很有用。
腳本的總執行時間：（0 分鐘 16.185 秒）
