PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb 是 PyTorch 提供的一項環境變量配置，用于控制 CUDA 顯存分配的行為。通過指定此參數，可以有效管理 GPU 顯存的碎片化，緩解因顯存碎片化而導致的 “CUDA out of memory”（顯存溢出）問題。

一.基本原理

在 PyTorch 中，CUDA 顯存的管理由 PyTorch 自身的 CUDA 緩存分配器負責。當 GPU 顯存中存在多個較小的碎片區域時，即使總體剩余顯存充足，也可能因缺少足夠大的連續顯存塊而無法分配大張量，導致顯存分配失敗（OOM）。

max_split_size_mb 參數作用是：

• 限制 CUDA 緩存分配器在分割較大顯存塊時生成碎片的最小尺寸。

• 更具體而言，當 GPU 需要分配一定大小的顯存時，如果已有的顯存塊太大，分配器將嘗試分割現有顯存塊為兩部分：一部分用于滿足請求，另一部分用于未來請求。

• 通過設置該值，可以避免過度分割顯存導致嚴重碎片化。例如，較高的值意味著較少的碎片，但可能會浪費更多顯存。

簡單理解：設置 max_split_size_mb 的數值越高，意味著更傾向于保留較大的連續顯存塊，顯存碎片化越低，但可能導致顯存空間利用率降低；反之，則利用率更高，但容易碎片化。

二.如何設置

1. 命令行設置

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
python train.py

2. Python代碼中設置

import osos.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"import torch# 后續代碼

注意：環境變量需在導入 PyTorch 前設置。

三.實際示例說明

假設訓練一個較大的 Transformer 模型過程中發生顯存錯誤（OOM）：

RuntimeError: CUDA out of memory. Tried to allocate 128 MiB.

通過配置此變量，嘗試緩解碎片化。

1.命令行示例

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256
python train.py

2.Python代碼示例

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:256"import torch
from torch import nndevice = torch.device('cuda')# 模擬大量顯存分配與釋放過程：
x = [torch.randn(1024, 1024, device=device) for _ in range(10)]
del x[::2]  # 隔一個刪除一個張量，制造顯存碎片# 再次嘗試分配較大的張量
try:large_tensor = torch.randn(4096, 4096, device=device)
except RuntimeError as e:print("OOM:", e)# 顯示當前顯存使用情況
print(torch.cuda.memory_summary(device=device))

四.如何選擇合適的值

• 初始嘗試推薦值：128MB、256MB、512MB

• 經驗法則：

  • 若反復出現OOM但整體顯存還有空余，提高此值以減少碎片。

  • 若顯存非常緊張（使用率很高），則降低此值以提高利用率。

典型應用場景：

• 大模型訓練（Transformer、BERT、GPT模型等）

• 多次頻繁顯存分配/釋放的任務（如頻繁forward/backward）

五.常用命令快速診斷顯存碎片

print(torch.cuda.memory_summary(device='cuda', abbreviated=False))

查看顯存碎片程度：

• 查看 "inactive_split_bytes"和"inactive_split_bytes.all.peak"等指標。

• 根據此診斷是否需要調整 max_split_size_mb。

PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb 提供了一種快速、有效的手段，以緩解 PyTorch CUDA 顯存分配時的碎片化問題，有效避免 GPU 顯存資源的浪費，保障模型訓練穩定進行。