如何透過將最佳化器步驟融入反向傳播來節省記憶體

創建於: Oct 02, 2023 | 最後更新於: Jan 16, 2024 | 最後驗證於: Nov 05, 2024

大家好！本教程旨在展示一種透過減少**梯度**佔用的記憶體來降低訓練迴圈記憶體佔用量的方法。假設你有一個模型，並且正在尋找最佳化記憶體的方法，以避免 Out of Memory (OOM) 錯誤，或者只是為了更好地利用 GPU。好吧，你_可能_很幸運（如果梯度佔用了你記憶體的一部分，並且你不需要進行梯度累積）。我們將探討以下內容：

1. 在訓練或微調迴圈中佔用記憶體的因素，

2. 如何捕獲和視覺化記憶體快照以確定瓶頸，

3. 新的 Tensor.register_post_accumulate_grad_hook(hook) API，最後，

4. 如何用 10 行程式碼將所有內容整合在一起以實現記憶體節省。

要執行本教程，你需要：

• 安裝了 torchvision 的 PyTorch 2.1.0 或更高版本

• 如果你想在本地執行記憶體視覺化，需要 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]
  3%|3         | 36.1M/1.13G [00:00<00:03, 379MB/s]
  7%|6         | 79.6M/1.13G [00:00<00:02, 424MB/s]
 11%|#         | 123M/1.13G [00:00<00:02, 436MB/s]
 14%|#4        | 166M/1.13G [00:00<00:02, 445MB/s]
 18%|#8        | 210M/1.13G [00:00<00:02, 450MB/s]
 22%|##1       | 254M/1.13G [00:00<00:02, 452MB/s]
 26%|##5       | 298M/1.13G [00:00<00:01, 454MB/s]
 29%|##9       | 341M/1.13G [00:00<00:01, 455MB/s]
 33%|###3      | 385M/1.13G [00:00<00:01, 456MB/s]
 37%|###6      | 429M/1.13G [00:01<00:01, 457MB/s]
 41%|####      | 472M/1.13G [00:02<00:06, 113MB/s]
 44%|####4     | 515M/1.13G [00:02<00:04, 146MB/s]
 48%|####7     | 552M/1.13G [00:02<00:03, 173MB/s]
 51%|#####     | 587M/1.13G [00:03<00:07, 83.8MB/s]
 54%|#####4    | 631M/1.13G [00:03<00:04, 114MB/s]
 58%|#####8    | 675M/1.13G [00:03<00:03, 150MB/s]
 62%|######1   | 719M/1.13G [00:03<00:02, 190MB/s]
 66%|######5   | 763M/1.13G [00:03<00:01, 232MB/s]
 70%|######9   | 807M/1.13G [00:03<00:01, 274MB/s]
 73%|#######3  | 851M/1.13G [00:03<00:01, 313MB/s]
 77%|#######7  | 895M/1.13G [00:04<00:00, 346MB/s]
 81%|########  | 940M/1.13G [00:04<00:00, 375MB/s]
 85%|########4 | 982M/1.13G [00:04<00:00, 380MB/s]
 88%|########8 | 1.00G/1.13G [00:04<00:00, 395MB/s]
 92%|#########1| 1.04G/1.13G [00:04<00:00, 411MB/s]
 96%|#########5| 1.09G/1.13G [00:04<00:00, 424MB/s]
 99%|#########9| 1.13G/1.13G [00:05<00:00, 112MB/s]
100%|##########| 1.13G/1.13G [00:05<00:00, 214MB/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 檔案，在 CUDA 記憶體視覺化工具 https://pytorch.com.tw/memory_viz 中開啟快照。記憶體時間線是否符合你的預期？

模型引數在訓練步驟之前已經載入到記憶體中，因此我們首先看到一部分記憶體分配給了權重。當我們開始前向傳播時，記憶體會逐漸分配給啟用，即我們為了在反向傳播中計算梯度而儲存的張量。一旦我們開始反向傳播，啟用會逐漸釋放，而梯度記憶體開始累積。

最後，當最佳化器開始工作時，其狀態將被延遲初始化，因此我們應該只在第一個訓練迴圈的最佳化器步驟中看到最佳化器狀態記憶體逐漸增加。在未來的迴圈中，最佳化器記憶體將保留並在原地更新。然後，在每次訓練迴圈結束時呼叫 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 上。

Tensor.register_post_accumulate_grad_hook(hook) API 和我們的技術

我們的技術依賴於在 backward() 期間不必儲存梯度。相反，一旦梯度累積完成，我們將立即將最佳化器應用於相應的引數並完全丟棄該梯度！這消除了在最佳化器步驟之前需要保留一個大的梯度緩衝區的需求。

那麼我們如何才能啟用這種更積極地應用最佳化器的行為呢？在我們的 2.1 版本中，我們添加了一個新的 API torch.Tensor.register_post_accumulate_grad_hook()，它允許我們在 Tensor 的 .grad 欄位累積完成後新增一個鉤子 (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```` 或操作最佳化器配置的使用者。將這個 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 Memory Visualizer 中。

幾項主要觀察

1. 沒有了最佳化器步驟！沒錯……我們把它融入了反向傳播。

2. 同樣，反向傳播持續時間更長，並且中間計算的隨機分配更多。這是預期的，因為最佳化器步驟需要中間計算。

3. 最重要的是！峰值記憶體降低了！現在是約 4GB（希望這與你之前的預期非常接近）。

請注意，與之前相比，現在沒有為梯度分配任何大的記憶體塊，這節省了約 1.2GB 的記憶體。相反，我們將最佳化器步驟儘可能提前，在每個梯度計算完成後非常快速地釋放了它們。太棒了！順便說一下，另外約 1.2GB 的記憶體節省來自將最佳化器分解為按引數劃分的最佳化器，因此中間記憶體也按比例縮小了。這個細節比梯度記憶體節省不那麼重要，因為即使不使用這項技術，你也可以透過關閉 foreach=False 來獲得最佳化器中間記憶體的節省。

你可能正確地想知道：如果我們節省了 2.4GB 的記憶體，為什麼峰值記憶體不是 6GB - 2.4GB = 3.6GB？嗯，峰值位置變了！現在的峰值出現在反向傳播步驟開始附近，那時記憶體中仍然有啟用；而之前，峰值出現在最佳化器步驟期間，那時啟用已經被釋放。因此，約 4.0GB - 約 3.6GB = 約 0.4GB 的差異是由於啟用記憶體。可以想象，這項技術可以與啟用檢查點相結合，以獲得更多的記憶體收益。

結論

在本教程中，我們學習了透過新的 Tensor.register_post_accumulate_grad_hook() API 將最佳化器融入反向傳播以節省記憶體的技術，以及何時應用這項技術（當梯度記憶體佔用顯著時）。在此過程中，我們還學習了記憶體快照，這在記憶體最佳化中通常非常有用。