訓練、推理 AI 模型的 VRAM 開銷計算筆記

Last Updated on 2024-10-23 by Clay

一直以來，我都只用個大概的公式去推估我的模型量級與我的 GPU VRAM 開銷之間關聯；畢竟這之間牽扯到的變數實在太多了，光是模型架構、層數、注意力機制實現、序列長度、Batch Size、訓練或推理採用的資料精度... 在在都影響我們最後計算的結果。

今天，我找了一些資料（附於文末 References，有一個還滿有趣的），嘗試理論性地計算我所使用的模型的理論值，並且與我訓練所使用的 VRAM 開銷做比對 —— 這之中當然發現過不少公式有些為謬誤之處，迭代修正到 9B、27B 的 Gemma 都與我的實戰經驗符合，這才對理論計算有了一定程度的信心。

當然，我知道不同的訓練框架、硬體架構、是否有做內核融合（Kernel Fusion）... 都會影響到最終的實際結果，然而一般來說不應有太大誤差（比方說突然差個 10GB，你就會知道大概算錯了），並且我們可以控制 LoRA 的一些參數、比如 rank 或是 adapter 要並行哪幾層等等變因，甚至是更頻繁的做資源回收（GC）以此來細微把控 VRAM 使用量。

訓練時的 VRAM 開銷理論值粗略計算

模型的訓練所花費的資源用量，基本可分成全量微調（full fine-tuning）和使用低秩矩陣適配器（Low-Rank Adapter, LoRA）兩大類，使用LoRA微調的情況僅會更新部分參數。

而另一大影響資源用量的情況則為使用的資料精度，現今最主流的精度則為 BFloat16，其指數範圍與 Float32（全精度）相同，僅有較不影響計算之精度有落差。可參考我之前寫過的筆記：深度學習中不同精度表示的差異 float32、float16、float8 和 bfloat16

而以下我計算 AI 模型的 VRAM 開銷，採用有效率之 LoRA 微調與 BFloat16 精度做為配置進行估算，並且在訓練時採用Batch Size = 1的超參數設定，再另行做梯度累積（Gradient Accumulation）以達成較長批次的訓練效果。這也是我習慣壓低 VRAM 的一貫操作方式。

根據 https://vram.asmirnov.xyz 的分析，我們所使用的訓練 VRAM 開銷可分為以下項目：

CUDA Kernels：當 PyTorch 第一次調用 CUDA 時，會使用約 0.977GiB 的 VRAM，此為較為固定的估計值，但最大值為 2GiB
CUDA Kernels VRAM = 2GiB
Parameters（模型參數量）：在 BFloat16 的情況下，一個參數會使用 16bits = 2bytes 進行估計，可由此項目推估模型完整載入記憶體、但不進行計算的使用量
Parameters VRAM（GiB）= Parameters * 2 bytes / (1,024 * 1,024 * 1,024)
Trainable Parameters Ratio：模型要訓練的參數量的比例，訓練參數量的比例越高，VRAM 的使用量也會變大
Gradients（梯度）：模型更新的梯度，會儲存模型要改進參數值，與訓練參數量同等的 VRAM 使用量
Gradients VRAM（GiB）= Parameters * Trainable Parameters Ratio * 2 bytes / (1,024 * 1,024 * 1,024)
First Moments：優化器儲存的梯度移動平均量，與梯度同等的VRAM使用量
First Moments VRAM = Gradients VRAM
Second Moments：優化器儲存的平方梯度移動平均量，與梯度同等的 VRAM 使用量
Second Moments VRAM = Gradients VRAM
Activations（中間激活量）：前向傳播（Forward Propagation）的中間張量大小，VRAM使用量隨著序列長度增長
Activations VRAM（GiB）= Layer Numbers * Batch Size * Sequence Length * Hidden Size * 2 bytes / (1,024 * 1,024 * 1,024)
Output Tensors：模型最後輸出的張量
Output Tensors VRAM（GiB）= Batch Size * Sequence Length * Vocabulary Size * 2 * 2bytes / (1,024 * 1,024 * 1,024)

以下是以一個以 Gemma-2-27b 為範例的計算過程，並有著以下超參數與模型設定（參考自 Gemma-2 技術報告：Gemma 2: Improving Open Language Models at a Practical Size）：

27B 參數量為 27 * (10 ^ 9) 個參數
訓練參數量比例（Trainable Parameters Ratio）按照注意力機制heads（QKV）分配，約為 0.02（2%）
訓練的序列長度（Sequence Length）統一填充為 8,192
詞彙量（Vocabulary）為 256,128 個Tokens
模型層數（Layer Number）為 46 層
隱藏層尺寸（Hidden Size）為 4,068 維
Batch Size設定為 1
資料精度採用 BFloat16（16 bits = 2 bytes）可透過上方公式計算出：
CUDA Kernels VRAM = 2 GiB
Parameters VRAM（GiB）= 27 * (10 ^ 9) * 2 bytes / (1,024 * 1,024 * 1,024) = 50.291 GiB
Gradients VRAM（GiB）= 27 * (10 ^ 9) * 0.02 * 2 bytes / (1,024 * 1,024 * 1,024) =1.006 GiB
First Moments VRAM = Gradients VRAM = 1.006 GiB
Second Moments VRAM = Gradients VRAM = 1.006 GiB
Activations VRAM（GiB）= 46 * 1 * 8,192 * 4,068 * 2 bytes / (1,024 * 1,024 * 1,024) = 2.855 GiB
Output Tensors VRAM（GiB）= 1 * 8,192 * 256,128 * 2 * 2 bytes / (1,024 * 1,024 * 1,024) = 7.816 GiB

總共的VRAM使用量約為 2 + 50.291 + 1.006 + 1.006 + 1.006 + 2.855 + 7.816 = 65.98 GiB，約是一張 80 GB的 H100可供訓練的 VRAM 使用量，還可以將 Batch Size 等超參數調高，進一步提升模型訓練效率呢。

我也寫了一個計算用的腳本：

# Define variables
num_params = 27 * 10**9  # Number of model parameters (27 billion)
seq_length = 8192  # Sequence length
vocab_size = 256128  # Vocabulary size
batch_size = 1  # Batch size
hidden_size = 4068  # Hidden size
num_layers = 46  # Number of model layers
half_precision_bytes = 2  # Half precision (2 bytes)

# LoRA training, represents the proportion of parameters to update
trainable_params_ratio = 0.02  # Assume LoRA updates only 2% of the parameters

# 1. CUDA Kernels
cuda_kernels_vram = 2  # Fixed value (GiB)

# 3. Parameter VRAM (all parameters stored in half precision)
params_vram = (num_params * half_precision_bytes) / (1024**3)  # GiB

# 4. Activation VRAM (using LoRA training, affects only part of the parameters)
lora_activations_vram = (num_layers * batch_size * seq_length * hidden_size * half_precision_bytes) / (1024**3)  # GiB

# 5. Gradient VRAM (using LoRA training, based on partial parameters)
lora_gradients_vram = (num_params * trainable_params_ratio * half_precision_bytes) / (1024**3)  # GiB

# 6. Optimizer's first moment VRAM (using LoRA training, based on partial parameters)
lora_first_moments_vram = (num_params * trainable_params_ratio * half_precision_bytes) / (1024**3)  # GiB

# 7. Optimizer's second moment VRAM (using LoRA training, based on partial parameters)
lora_second_moments_vram = (num_params * trainable_params_ratio * half_precision_bytes) / (1024**3)  # GiB

# 8. Output tensor VRAM
output_tensor_vram = (batch_size * seq_length * vocab_size * half_precision_bytes * 2) / (1024**3)  # GiB

# Calculate total VRAM usage (using LoRA training)
total_lora_vram_usage = (
    cuda_kernels_vram +
    params_vram +
    lora_gradients_vram +
    lora_activations_vram +
    lora_first_moments_vram +
    lora_second_moments_vram +
    output_tensor_vram
)

# Print results
vram_usage_details = {
    "CUDA Kernels VRAM": cuda_kernels_vram,
    "Parameters VRAM": params_vram,
    "LoRA Gradients VRAM": lora_gradients_vram,
    "LoRA Activations VRAM": lora_activations_vram,
    "LoRA First Moments VRAM": lora_first_moments_vram,
    "LoRA Second Moments VRAM": lora_second_moments_vram,
    "Output Tensor VRAM": output_tensor_vram,
    "Total VRAM Usage (LoRA)": total_lora_vram_usage
}

# Output results
for key, value in vram_usage_details.items():
    print(f"{key}: {value:.3f} GiB")

Output:

CUDA Kernels VRAM: 2.000 GiB
Parameters VRAM: 50.291 GiB
LoRA Gradients VRAM: 1.006 GiB
LoRA Activations VRAM: 2.855 GiB
LoRA First Moments VRAM: 1.006 GiB
LoRA Second Moments VRAM: 1.006 GiB
Output Tensor VRAM: 7.816 GiB
Total VRAM Usage (LoRA): 65.981 GiB

推理時的 VRAM 開銷理論值粗略計算

計算完訓練之後，剩下的推理（Inference）就簡單了：在這裡，我們假設我們已經把訓練好的 LoRA 適配器融合回去模型本體，並且移除梯度與優化器的 VRAM 使用量，最後再把本來會儲存的反向傳播（Backward Propagation）VRAM 使用量移除，剩下約是推理的 VRAM 開銷了。

不過，再次強調，這都是理論值，如果我們採用了如 vLLM 等模型的 chunked pre-filling、CPU Offloading 等技術，VRAM 會存在很大的差異，依然是以實戰評估最為精準。

同樣的，我有一個計算用的腳本：

# Define variables
num_params = 27 * 10**9  # Number of model parameters (9 billion)
seq_length = 8192        # Sequence length
vocab_size = 256128      # Vocabulary size
batch_size = 1           # Batch size
hidden_size = 4068       # Hidden size
num_layers = 46          # Number of model layers
half_precision_bytes = 2 # Half precision (2 bytes)

# 1. CUDA Kernels
cuda_kernels_vram = 2  # Fixed value (GiB)

# 2. Parameters VRAM (all parameters stored in half precision)
params_vram = (num_params * half_precision_bytes) / (1024**3)  # GiB

# 3. Activation VRAM (Inference stage)
activations_vram = (num_layers * batch_size * seq_length * hidden_size * half_precision_bytes) / (1024**3)  # GiB

# 4. Output Tensor VRAM
output_tensor_vram = (batch_size * seq_length * vocab_size * half_precision_bytes) / (1024**3)  # GiB

# Calculate total VRAM usage (Inference)
total_inference_vram_usage = (cuda_kernels_vram + params_vram + activations_vram + output_tensor_vram)

# Print results
vram_usage_details = {
    "CUDA Kernels VRAM": cuda_kernels_vram,
    "Parameters VRAM": params_vram,
    "Activations VRAM": activations_vram,
    "Output Tensor VRAM": output_tensor_vram,
    "Total VRAM Usage (Inference)": total_inference_vram_usage
}

# Output results
for key, value in vram_usage_details.items():
    print(f"{key}: {value:.3f} GiB")

Output:

CUDA Kernels VRAM: 2.000 GiB
Parameters VRAM: 50.291 GiB
Activations VRAM: 2.855 GiB
Output Tensor VRAM: 3.908 GiB
Total VRAM Usage (Inference): 59.055 GiB

最後，雖然我認為我再三檢查過了，也與我自己訓練、推理 AI 模型的經驗相符，不過若是我列出的公式或計算有謬誤之處，還請不吝指正！感謝～

一直到今天，我都還在努力學習各式各樣的事情呢。

References

深度學習中不同精度表示的差異 float32、float16、float8 和 bfloat16

使用 vLLM 在 V100 GPU 上加速推理 Gemma-2 疑難雜症排除

訓練、推理 AI 模型的 VRAM 開銷計算筆記

訓練時的 VRAM 開銷理論值粗略計算

推理時的 VRAM 開銷理論值粗略計算

References

Read More

相關

Leave a Reply取消回覆

訓練、推理 AI 模型的 VRAM 開銷計算筆記

訓練時的 VRAM 開銷理論值粗略計算

推理時的 VRAM 開銷理論值粗略計算

References

Read More

分享此文：

相關

Leave a Reply取消回覆