文章目錄
1、引言
2、自動混合精度訓練
3、低精度訓練
4、梯度檢查點
5、通過梯度累積減小批量大小
6、張量分片與分布式訓練
7、高效數據加載
8、使用 In-Place 操作
9、Activation and Parameter Offloading
10、使用更精簡的優化器
11、高級策略
12、總結
1、引言
在訓練大型深度學習模型（包括LLM和視覺Transformer）時，最常見的瓶頸之一就是顯存消耗達到峰值。由于大多數人無法使用大規模的GPU集群，因此在本文中將概述一些技術和策略，在不犧牲模型性能和預測準確性的情況下，將顯存消耗降低近20倍。請記住，這些技術中的大多數應用并不互相排斥，可以很容易地結合使用，以提高顯存效率。
2、自動混合精度訓練
混合精度訓練結合了16位（FP16）和32位（FP32）浮點格式。其核心思想是在低精度下執行大部分數學運算，從而降低顯存帶寬和存儲需求，同時在計算的關鍵環節保留必要的精度保障。通過使用FP16存儲激活值和梯度，這些張量的顯存占用量可減少約一半。但需注意，某些網絡層或運算仍需保持FP32精度以避免數值不穩定問題。值得慶幸的是，PyTorch對自動混合精度（AMP）的原生支持極大簡化了這一過程。
注意這里是混合精度訓練而不是低精度訓練
什么是混合精度訓練？
混合精度訓練結合使用16位（FP16）和32位（FP32）浮點格式以保持模型精度。通過使用16位精度計算梯度，相比全32位精度計算，這一過程可大幅加快運算速度并顯著減少顯存占用。這種方法在顯存或計算資源受限的場景下尤為實用。
之所以采用混合精度而非低精度這一表述，是因為并非所有參數或運算都被轉換為16位格式。實際上，訓練過程會在32位與16位運算之間動態切換，這種精度層級的有序交替正是該技術被稱為混合精度的根本原因。
如上述示意圖所示，混合精度訓練流程首先將權重轉換為低精度格式（FP16）以加速計算，隨后梯度計算在低精度環境下完成，但為確保數值穩定性，這些梯度會被重新轉換為高精度格式（FP32），最終經過縮放處理的梯度將用于更新原始權重。因此，通過這種機制既能提升訓練效率，又不會犧牲網絡的整體精度與穩定性。
如前所述，使用 torch.cuda.amp.autocast( ) 可以輕松啟用該功能，一個簡單的代碼示例片段如下：
import?torch
from?torch.cuda.amp?import?autocast, GradScaler# Assume your model and optimizer have been defined elsewhere.
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()
for?data, target?in?data_loader:optimizer.zero_grad()# Enable mixed precisionwith?autocast():output = model(data)loss = loss_fn(output, target)# Scale the loss and backpropagatescaler.scale(loss).backward()scaler.step(optimizer)scaler.update()3、低精度訓練
如原文所述，理論上可以更進一步嘗試完全使用16位低精度（而非混合精度）進行訓練。但此時可能因16位浮點數的固有精度限制出現NaN值異常。為解決這一問題，業界開發了多種新型浮點格式，其中由谷歌專門為此研發的BF16應用較為廣泛。簡而言之，相較于標準的FP16，BF16擁有更大的動態范圍——這種擴展的動態范圍使其能夠更精確地表示極大或極小的數值，從而更適配可能遭遇廣泛數值區間的深度學習場景。雖然其較低的尾數精度在某些情況下可能影響計算準確性或引發舍入誤差，但在大多數實踐中對模型性能的影響微乎其微。
FP16與BF16的動態范圍對比
雖然這種格式最初是為TPU開發的，但在大多數現代GPU（Nvidia Ampere架構及更高版本）也支持這種格式。大家可以使用以下方法檢查您的GPU是否支持這種格式：
import?torch
print(torch.cuda.is_bf16_supported()) ?# should print True4、梯度檢查點
即使采用混合精度與低精度訓練，這些大型模型仍會生成大量中間張量，消耗可觀的顯存。梯度檢查點技術通過在前向傳播過程中選擇性存儲部分中間結果來解決這一問題——未被保存的中間張量將在反向傳播階段重新計算。盡管這會引入額外的計算開銷，卻能顯著節省顯存資源。
通過策略性選擇需設置檢查點的網絡層，大家可通過動態重新計算激活值而非存儲它們來減少顯存使用。這種時間與內存的折中策略對于具有深層架構的模型特別有益，因為中間激活值占用了大部分內存消耗。以下是一個簡單的使用示例：
import?torch
from?torch.utils.checkpoint?import?checkpoint
def?checkpointed_segment(input_tensor):# This function represents a portion of your model# which will be recomputed during the backward pass.# You can create a custom forward pass for this segment.return?model_segment(input_tensor)
# Instead of a conventional forward pass, wrap the segment with checkpoint.
output = checkpoint(checkpointed_segment, input_tensor)采用該方法，在大多數情況下可使激活值的顯存占用量降低40%至50%。盡管反向傳播階段因此增加了額外的計算量，但在GPU顯存成為瓶頸的場景下，這種以時間換空間的策略通常是可接受的。
5、通過梯度累積減小批量大小
通過最初的方法，你可能會問自己：
為什么不干脆減少batchsize大小？
通過減小批量大小的確是減少顯存占用最直接的方法，但需注意的是，這種方式在多數情況下會導致模型預測性弱于使用更大批量訓練的模型。因此需要在顯存限制與模型效果之間謹慎權衡。
那么如何達到平衡呢？
這正是梯度累積技術發揮作用之處！該方法通過在訓練過程中虛擬增大有效批量規模：其核心原理是先在較小的批量上計算梯度，并經過多次迭代的累積（通過采用累加或平均方式），而非在每批次處理后立即更新模型參數。當累積梯度達到目標“虛擬”批量規模時，才使用聚合后的梯度一次性完成模型權重的更新。
這種技術的一個主要缺點是大大增加了訓練時間。
6、張量分片與分布式訓練
對于單個GPU無法容納的龐大訓練模型（即使經過上述優化），完全分片數據并行（FSDP）是不可或缺的。FSDP將模型參數、梯度和優化器狀態分散到多個GPU上。這不僅能將巨大的模型放入顯存，還能通過更好地分配通信開銷提高訓練效率。
FSDP不在每個GPU上維護模型的完整副本，而是在可用設備之間分配模型參數。在執行前向或后向傳遞時，只有相關的分片被加載到顯存中。這種分片機制大大降低了對每臺設備顯存的需求，結合上述技術，在某些情況下甚至可以將顯存需求降低10倍。
Tensor Parallel
樣例如下：
import?torch
from?torch.distributed.fsdp?import?FullyShardedDataParallel?as?FSDP
# Initialize your model and ensure it is on the correct device.
model = MyLargeModel().cuda()
# Wrap the model in FSDP for sharded training across GPUs.
fsdp_model = FSDP(model)7、高效數據加載
在顯存優化實踐中，數據加載環節常被忽視。雖然優化重點通常集中在模型內部結構與計算過程上，但低效的數據處理可能引發不必要的性能瓶頸，同時影響顯存占用與訓練速度。若不確定如何優化數據加載器，可遵循以下經驗法則：優先啟用固定內存（Pinned Memory）與多工作進程（Multiple Workers）配置。
from?torch.utils.data?import?DataLoader# Create your dataset instance and then the DataLoader with pinned memory enabled.
train_loader = DataLoader(dataset,batch_size=64,shuffle=True,num_workers=4, ? ? ?# Adjust based on your CPU capabilitiespin_memory=True? ? ?# Enables faster host-to-device transfers
)8、使用 In-Place 操作
在張量運算中，若未謹慎管理內存，每次操作都可能生成新對象。原地（In-Place）操作通過直接修改現有張量而非創建副本，可有效減少內存碎片化與總體內存占用。這種特性尤其有利于降低迭代訓練循環中的臨時內存分配開銷。例如：
import?torch
x = torch.randn(100,?100, device='cuda')
y = torch.randn(100,?100, device='cuda')
# Using in-place addition
x.add_(y) ?# Here x is modified directly instead of creating a new tensor9、Activation and Parameter Offloading
即便綜合運用前述所有優化技術，在訓練超大規模模型時，仍可能因海量中間激活值的瞬時占用而觸及GPU顯存容量極限。此時，中間數據卸載技術可作為額外的安全閥機制——其核心思路是將部分非即時必需的中間數據臨時轉換至CPU內存，從而為GPU顯存騰出關鍵空間，確保訓練流程持續進行。
我們通過策略性將部分激活值和或模型參數卸載至CPU內存，從而將GPU顯存專用于核心計算任務。雖然如DeepSpeed、Fabric等專業框架已內置管理此類數據遷移的機制，大家仍可通過以下方式自主實現該功能。
def?offload_activation(tensor):# Move tensor to CPU to save GPU memoryreturn?tensor.cpu()def?process_batch(data):# Offload some activations explicitlyintermediate = model.layer1(data)intermediate = offload_activation(intermediate)intermediate = intermediate.cuda() ?# Move back when neededoutput = model.layer2(intermediate)return?output10、使用更精簡的優化器
并非所有優化器對內存的需求均等。以廣泛使用的Adam優化器為例，其針對每個模型參數需額外維護兩個狀態變量（均值與方差），導致內存占用倍增。相比之下，采用無狀態優化器（如SGD）可將參數總量減少近三分之二——這對于訓練大語言模型（LLMs）及其他大規模模型具有顯著意義。
盡管普通SGD優化器存在收斂性能較弱的缺陷，但通過引入余弦衰減學習率調整策略（Cosine Decay Learning Rate Scheduler）可有效補償這一不足。簡而言之：
# instead of this
optimizer?= torch.optim.Adam(model.parameters(), lr=5e-5)
# use this
optimizer?= torch.optim.SGD(model.parameters(), lr=0.01)
num_steps?= NUM_EPOCHS * len(train_loader)
scheduler?= torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps)通過這一調整，大家可以在顯著改變峰值內存占有量的同時（具體取決于實際任務需求），仍能保持模型精度接近97%的水平。
11、高級策略
雖然上面列出的技術確實為我們奠定了堅實的基礎，但我還想列出一些其他高級策略，我們可以考慮將 GPU 提升到極限：
- 內存剖析和高速緩存管理
如果無法測量，就很難優化。PyTorch 提供了一些檢查 GPU 內存使用情況的默認實用程序。使用方法如下：
import?torch
# print a detailed report of current GPU memory usage and fragmentation
print(torch.cuda.memory_summary(device=None, abbreviated=False))
# free up cached memory that’s no longer needed by PyTorch
torch.cuda.empty_cache()- 使用TorchScript進行JIT編譯
PyTorch 的即時（JIT）編譯器使大家使用 TorchScript 將 Python 模型轉換為優化的、可序列化的程序。通過優化內核啟動和減少開銷，這種轉換可以提高內存和性能。您可以通過以下方式輕松訪問它：
import torch
# Suppose `model` is an instance of your PyTorch network.
scripted_model = torch.jit.script(model)
# Now, you can run the scripted model just like before.
output = scripted_model(input_tensor)盡管框架原生方法已能實現基礎功能，但模型編譯技術通常能帶來更深層次的性能優化。
-
自定義內核融合
編譯的另一個主要好處是將多個操作融合到一個內核中。這有助于減少內存讀/寫，提高整體吞吐量。融合后的操作如下：
- 使用torch.compile（）進行動態內存分配
再次從編譯中獲益--使用 JIT 編譯器可通過利用跟蹤和圖形優化技術的編譯時優化來優化動態內存分配，從而進一步壓縮內存并提高性能，尤其是在大型模型和Transformer架構中。
12、總結
隨著 GPU 和云計算變得異常昂貴，只有充分利用現有資源才有意義。這有時可能意味著要在單個 GPU 工作站/筆記本電腦上對 LLM 或視覺Transformer進行訓練/微調。上面列出的技術是研究人員/專業人士在算力緊張的情況下進行訓練所使用的眾多策略中的一部分。
參考資料：AI算法之道
