在 PyTorch 中視覺化並理解 GPU 記憶體

你一定很熟悉這條訊息 🤬

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 7.93 GiB total capacity; 6.00 GiB already allocated; 14.88 MiB free; 6.00 GiB reserved in total by PyTorch)

雖然很容易看到 GPU 記憶體已滿，但理解原因以及如何解決它可能更具挑戰性。在本教程中，我們將逐步介紹如何在訓練期間視覺化並理解 PyTorch 中的 GPU 記憶體使用情況。我們還將瞭解如何估算記憶體需求並最佳化 GPU 記憶體使用情況。

🔎 PyTorch 視覺化工具

PyTorch 提供了一個方便的工具用於視覺化 GPU 記憶體使用情況

import torch
from torch import nn

# Start recording memory snapshot history
torch.cuda.memory._record_memory_history(max_entries=100000)

model = nn.Linear(10_000, 50_000, device ="cuda")
for _ in range(3):
    inputs = torch.randn(5_000, 10_000, device="cuda")
    outputs = model(inputs)

# Dump memory snapshot history to a file and stop recording
torch.cuda.memory._dump_snapshot("profile.pkl")
torch.cuda.memory._record_memory_history(enabled=None)

執行此程式碼會生成一個包含 GPU 記憶體使用歷史記錄的 profile.pkl 檔案。你可以在以下網址視覺化此歷史記錄：https://pytorch.org/memory_viz。

透過拖放 profile.pkl 檔案，你將看到一個類似這樣的圖表

讓我們將此圖表分解為關鍵部分

1. 模型建立：記憶體增加 2 GB，對應於模型的大小

   $10{,}000 \times 50{,}000 \text{ weights} + 50{,}000 \text{ biases in } \texttt{float32 }\text{(4 bytes)} \implies (5 \times 10^8) \times 4 \, \text{bytes} = 2 \, \text{GB}.$

   此記憶體（藍色部分）在整個執行過程中持續存在。

2. 輸入張量建立（第一次迴圈）：記憶體增加 200 MB，與輸入張量大小匹配

   $5{,}000 \times 10{,}000 \text{ elements in } \texttt{float32 }\text{(4 bytes)} \implies (5 \times 10^7) \times 4 \, \text{bytes} = 0.2 \, \text{GB}.$

3. 正向傳播（第一次迴圈）：輸出張量記憶體增加 1 GB

   $5{,}000 \times 50{,}000 \text{ elements in } \texttt{float32 }\text{(4 bytes)} \implies (25 \times 10^7) \times 4 \, \text{bytes} = 1 \, \text{GB}.$

4. 輸入張量建立（第二次迴圈）：記憶體再次增加 200 MB，用於新的輸入張量。此時，你可能期望步驟 2 中的輸入張量被釋放。然而它並沒有被釋放：模型保留了其啟用值，因此即使該張量不再賦值給變數 inputs，它仍然被模型的正向傳播計算所引用。模型保留其啟用值是因為在神經網路中反向傳播過程需要這些張量。嘗試使用 torch.no_grad() 檢視差異。

5. 正向傳播（第二次迴圈）：記憶體增加 1 GB，用於新的輸出張量，計算方式與步驟 3 相同。

6. 釋放第一次迴圈啟用值：在第二次迴圈的正向傳播之後，第一次迴圈（步驟 2）的輸入張量可以被釋放。模型保留了第一個輸入張量的啟用值，這些啟用值被第二次迴圈的輸入覆蓋。一旦第二次迴圈完成，第一個張量不再被引用，其記憶體可以被釋放。

7. 更新 output：步驟 3 的輸出張量重新賦值給變數 output。之前的張量不再被引用並被刪除，釋放了其記憶體。

8. 輸入張量建立（第三次迴圈）：與步驟 4 相同。

9. 正向傳播（第三次迴圈）：與步驟 5 相同。

10. 釋放第二次迴圈啟用值：步驟 4 的輸入張量被釋放。

11. 再次更新 output：步驟 5 的輸出張量重新賦值給變數 output，釋放了之前的張量。

12. 程式碼執行結束：所有記憶體被釋放。

📊 訓練期間的記憶體視覺化

前面的例子被簡化了。在實際場景中，我們通常訓練複雜的模型，而不是單個線性層。此外，前面的例子沒有包括訓練過程。在這裡，我們將研究在一個真實的大型語言模型（LLM）的完整訓練迴圈中 GPU 記憶體的行為。

import torch
from transformers import AutoModelForCausalLM

# Start recording memory snapshot history
torch.cuda.memory._record_memory_history(max_entries=100000)

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for _ in range(3):
    inputs = torch.randint(0, 100, (16, 256), device="cuda")  # Dummy input
    loss = torch.mean(model(inputs).logits)  # Dummy loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# Dump memory snapshot history to a file and stop recording
torch.cuda.memory._dump_snapshot("profile.pkl")
torch.cuda.memory._record_memory_history(enabled=None)

💡 提示： 在進行效能分析時，限制步數。每個 GPU 記憶體事件都會被記錄，檔案可能會變得非常大。例如，上述程式碼生成了一個 8 MB 的檔案。

這是此示例的記憶體配置檔案

此圖表比前一個示例更復雜，但我們仍然可以逐步分解它。注意三個峰值，每個峰值對應於訓練迴圈的一次迭代。讓我們簡化圖表，使其更容易解釋

1. 模型初始化 (model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda"))
   第一步涉及將模型載入到 GPU。模型引數（藍色）佔用記憶體並保持在那裡直到訓練結束。

2. 正向傳播 (model(inputs))
   在正向傳播過程中，啟用值（每層的中間輸出）被計算並存儲在記憶體中用於反向傳播。這些啟用值（橙色表示）逐層增長，直到最後一層。損失在橙色區域的峰值處計算。

3. 反向傳播 (loss.backward())
   梯度（黃色）在此階段計算並存儲。同時，啟用值被丟棄，因為它們不再需要，導致橙色區域縮小。黃色區域表示梯度計算的記憶體使用情況。

4. 最佳化器步驟 (optimizer.step())
   梯度用於更新模型的引數。最初，最佳化器本身被初始化（綠色區域）。此初始化只進行一次。之後，最佳化器使用梯度更新模型的引數。為了更新引數，最佳化器會臨時儲存中間值（紅色區域）。更新後，梯度（黃色）和最佳化器中間值（紅色）都被丟棄，釋放記憶體。

至此，一次訓練迭代完成。此過程對剩餘的迭代重複，從而在圖表中產生可見的三個記憶體峰值。

像這樣的訓練配置檔案通常遵循一致的模式，這使得它們對於估算給定模型和訓練迴圈的 GPU 記憶體需求非常有用。

📐 估算記憶體需求

從上面一節來看，估算 GPU 記憶體需求似乎很簡單。所需的總記憶體應對應於記憶體配置檔案中的最高峰值，這發生在正向傳播期間。在這種情況下，記憶體需求是（藍色 + 綠色 + 橙色）：$\text{Model Parameters} + \text{Optimizer State} + \text{Activations}$

就這麼簡單嗎？實際上，這裡有一個陷阱。配置檔案可能因訓練設定而異。例如，將批次大小從 16 減少到 2 會改變情況

- inputs = torch.randint(0, 100, (16, 256), device="cuda")  # Dummy input
+ inputs = torch.randint(0, 100, (2, 256), device="cuda")  # Dummy input

現在，最高峰值出現在最佳化器步驟期間，而不是正向傳播期間。在這種情況下，記憶體需求變為（藍色 + 綠色 + 黃色 + 紅色）：$\text{Model Parameters} + \text{Optimizer State} + \text{Gradients} + \text{Optimizer Intermediates}$

為了概括記憶體估算，我們需要考慮所有可能的峰值，無論它們發生在正向傳播還是最佳化器步驟期間。$\text{Model Parameters} + \text{Optimizer State} + \max(\text{Gradients} + {\text{Optimizer Intermediates}, \text{Activations}})$

現在我們有了公式，讓我們看看如何估算每個元件。

模型引數

模型引數最容易估算。$\text{Model Memory} = N \times P$

其中：

• $N$ 是引數數量。
• $P$ 是精度（以位元組為單位，例如 float32 為 4）。

例如，一個具有 15 億引數且精度為 4 位元組的模型需要

在上述示例中，模型大小為：$\text{Model Memory} = 1.5 \times 10^9 \times 4 \, \text{bytes} = 6 \, \text{GB}$

最佳化器狀態

最佳化器狀態所需的記憶體取決於最佳化器型別和模型引數。例如，AdamW 最佳化器為每個引數儲存兩個動量（一階和二階）。這使得最佳化器狀態大小為：$\text{Optimizer State Size} = 2 \times N \times P$

啟用值

啟用值所需的記憶體更難估算，因為它包括在正向傳播過程中計算的所有中間值。要計算啟用記憶體，我們可以使用正向鉤子來測量輸出的大小

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")

activation_sizes = []

def forward_hook(module, input, output):
    """
    Hook to calculate activation size for each module.
    """
    if isinstance(output, torch.Tensor):
        activation_sizes.append(output.numel() * output.element_size())
    elif isinstance(output, (tuple, list)):
        for tensor in output:
            if isinstance(tensor, torch.Tensor):
                activation_sizes.append(tensor.numel() * tensor.element_size())

# Register hooks for each submodule
hooks = []
for submodule in model.modules():
    hooks.append(submodule.register_forward_hook(forward_hook))

# Perform a forward pass with a dummy input
dummy_input = torch.zeros((1, 1), dtype=torch.int64, device="cuda")
model.eval()  # No gradients needed for memory measurement
with torch.no_grad():
    model(dummy_input)

# Clean up hooks
for hook in hooks:
    hook.remove()

print(sum(activation_sizes))  # Output: 5065216

對於 Qwen2.5-1.5B 模型，這會給出每個輸入 token 5,065,216 個啟用值。要估算輸入張量的總啟用記憶體，請使用：$\text{Activation Memory} = A \times B \times L \times P$

其中：

• $A$ 是每個 token 的啟用值數量。
• $B$ 是批次大小。
• $L$ 是序列長度。

然而，直接使用這種方法並不總是實用的。理想情況下，我們希望有一種啟發式方法來估算啟用記憶體，而無需執行模型。此外，我們可以直觀地看到更大的模型具有更多的啟用。這引出了一個問題：模型引數數量與啟用數量之間是否存在關聯？

並非直接相關，因為每個 token 的啟用數量取決於模型架構。然而，LLM 傾向於具有相似的結構。透過分析不同的模型，我們觀察到引數數量與啟用數量之間存在大致的線性關係

這種線性關係使我們能夠使用啟發式方法估算啟用值：$A = 4.6894 \times 10^{-4} \times N + 1.8494 \times 10^{6}$

儘管這只是一個近似值，但它提供了一種無需為每個模型執行復雜計算即可估算啟用記憶體的實用方法。

梯度

梯度更容易估算。梯度所需的記憶體與模型引數相同：$\text{Gradients Memory} = N \times P$

最佳化器中間值

在更新模型引數時，最佳化器會儲存中間值。這些值所需的記憶體與模型引數相同：$\text{Optimizer Intermediates Memory} = N \times P$

總記憶體

總而言之，訓練模型所需的總記憶體為：$\text{Total Memory} = \text{Model Memory} + \text{Optimizer State} + \max(\text{Gradients}, \text{Optimizer Intermediates}, \text{Activations})$

元件如下

• 模型記憶體：$N \times P$
• 最佳化器狀態：$2 \times N \times P$
• 梯度：$N \times P$
• 最佳化器中間值：$N \times P$
• 啟用值：$A \times B \times L \times P$，使用啟發式方法估算：$A = 4.6894 \times 10^{-4} \times N + 1.8494 \times 10^{6}$

為了讓這個計算更容易，我為你建立了一個小工具

🚀 後續步驟

你最初理解記憶體使用情況的動機很可能是因為有一天你記憶體不足了。這篇部落格為你提供了直接的解決方案來解決這個問題嗎？可能沒有。然而，現在你對記憶體使用情況有了更好的理解，並且知道了如何分析它，你就能更好地找到減少記憶體使用的方法了。

有關最佳化 TRL 記憶體使用的具體技巧列表，你可以檢視文件的減少記憶體使用部分。不過，這些技巧不限於 TRL，可以應用於任何基於 PyTorch 的訓練過程。

🤝 致謝

感謝 Kashif Rasul 對本部落格文章提出的寶貴反饋和建議。