在大規模深度學習模型訓練過程中，GPU內存容量往往成為制約因素，尤其是在訓練大型語言模型(LLM)和視覺Transformer等現代架構時。由于大多數研究者和開發者無法使用配備海量GPU內存的高端計算集群，因此掌握有效的內存優化技術變得尤為關鍵。本文將系統性地介紹多種內存優化策略，這些技術組合應用可使模型訓練的內存消耗降低近20倍，同時不會損害模型性能和預測準確率。以下大部分技術可以相互結合，以獲得更顯著的內存效率提升。

1、自動混合精度訓練

混合精度訓練是降低內存占用的基礎且高效的方法，它充分利用16位(FP16)和32位(FP32)浮點格式的優勢。

混合精度訓練的核心思想是在大部分計算中使用較低精度執行數學運算，從而減少內存帶寬和存儲需求，同時在計算的關鍵環節保持必要的精度。通過對激活值和梯度采用FP16格式，這些張量的內存占用可減少約50%。然而某些特定的層或操作仍需要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 precision
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)

    # Scale the loss and backpropagate
    scaler.scale(loss).backward()
    scaler.step(optimizer)
     scaler.update()

2、低精度訓練

除了混合精度訓練，我們還可以嘗試使用完整的16位低精度格式進行訓練。由于16位浮點數的表示范圍限制，這種方法可能導致NaN值出現。為解決這一問題，研究人員開發了多種專用浮點格式。其中，Brain Floating Point(BF16)是Google為此專門開發的一種廣受歡迎的格式。與標準FP16相比，BF16提供了更大的動態范圍，能夠表示極大和極小的數值，使其更適合于深度學習應用中可能遇到的多樣化數值情況。盡管較低精度可能在某些計算中影響精確度或導致舍入誤差，但在大多數深度學習應用場景中，這種影響對模型性能的影響極小。

雖然BF16最初是為TPU設計的，但現在大多數現代GPU(Nvidia Ampere架構及更高版本)也支持這種格式。可以通過以下方法檢查GPU是否支持BF16格式：

import torch
 print(torch.cuda.is_bf16_supported())  # should print True

3、梯度檢查點

即便采用混合精度和低精度訓練，大型模型在前向傳播過程中產生的大量中間張量仍會消耗大量內存。梯度檢查點(Gradient Checkpointing)技術通過在前向傳播過程中選擇性地僅存儲部分中間結果來解決這一問題。在反向傳播過程中，系統會重新計算缺失的中間值，這雖然增加了計算成本，但可以顯著降低內存需求。

通過戰略性地選擇需要設置檢查點的層，可以通過動態重新計算激活值而非存儲它們來減少內存使用。對于具有深層架構的模型，中間激活值通常占據了內存消耗的主要部分，此時這種權衡尤為有效。以下是梯度檢查點的實現示例：

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內存成為限制因素時，這種權衡通常是合理的。

4、使用梯度累積降低批量大小

在嘗試上述方法后，一個自然的問題是：

為何不直接減小批量大小？

雖然這確實是最直接的方法，但通常使用較小批量大小會導致預測性能下降。簡單減小批量大小雖然能顯著降低內存消耗，但往往會對模型準確率產生不良影響。

如何在這兩者之間取得平衡？

梯度累積(Gradient Accumulation)正是為解決這一問題而設計的技術。它允許在訓練過程中虛擬增加批量大小，其核心原理是為較小的批量計算梯度，并在多次迭代中累積這些梯度(通常通過求和或平均)，而不是在每個批次后立即更新模型權重。一旦累積的梯度達到目標"虛擬"批量大小，才使用這些累積的梯度更新模型參數。

然而需要注意，這種技術的主要缺點是顯著增加了訓練時間。

5、張量分片和分布式訓練

對于即使應用上述優化后仍無法在單個GPU上容納的超大模型，完全分片數據并行(Fully Sharded Data Parallel, FSDP)技術提供了解決方案。FSDP將模型參數、梯度和優化器狀態分片到多個GPU上，這不僅使得訓練超大模型成為可能，還能通過更合理地分配通信開銷提高訓練效率。

FSDP不是在每個GPU上維護完整的模型副本，而是將模型參數分配到多個可用設備上。在執行前向或反向傳播時，系統僅將相關分片加載到內存中。這種分片機制顯著降低了單個設備的內存需求，與前述技術結合使用，在某些情況下可實現高達10倍的內存降低效果。

FSDP可通過以下方式實現：

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)

6、高效的數據加載

內存優化中常被忽視的一個方面是數據加載效率。雖然大部分優化關注點集中在模型內部結構和計算過程，但低效的數據處理同樣可能造成不必要的瓶頸，影響內存利用和計算速度。作為經驗法則，當處理數據加載器時，應始終啟用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 capabilities
    pin_memory=True     # Enables faster host-to-device transfers
 )

7、使用原地操作

在處理張量時，如果不謹慎管理，每個操作都可能創建新的張量對象。原地操作(In-place Operations)通過直接修改現有張量而非分配新張量，有助于減少內存碎片和總體內存占用。這種方式減少了臨時內存分配，在迭代訓練循環中尤為重要。示例如下：

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 tensor

8、激活和參數卸載

對于極大規模模型，即使應用了所有上述技術，由于大量中間激活值的存在，仍可能達到GPU內存限制。激活和參數卸載(Activation and Parameter Offloading)技術通過將部分中間數據移動到CPU內存，為GPU內存提供額外的緩解。

這種方法通過戰略性地將部分激活值和/或參數臨時卸載到主機內存(CPU)，僅在GPU內存中保留關鍵計算所需的數據。雖然DeepSpeed、Fabric等專用框架可自動管理這種數據移動，但也可以按如下方式實現自定義卸載邏輯：

def offload_activation(tensor):
    # Move tensor to CPU to save GPU memory
    return tensor.cpu()

def process_batch(data):
    # Offload some activations explicitly
    intermediate = model.layer1(data)
    intermediate = offload_activation(intermediate)
    intermediate = intermediate.cuda()  # Move back when needed
    output = model.layer2(intermediate)
     return output

9、使用更精簡的優化器

各種優化器在內存消耗方面存在顯著差異。例如，廣泛使用的Adam優化器為每個模型參數維護兩個額外狀態參數(動量和方差)，這意味著更多的內存消耗。將Adam替換為無狀態優化器(如SGD)可將參數數量減少近2/3，這在處理LLM等大型模型時尤為重要。

標準SGD的缺點是收斂特性較差。為彌補這一點，可引入余弦退火學習率調度器以實現更好的收斂效果。實現示例：

# 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%(取決于具體應用)的同時，顯著改善峰值內存消耗。

10、進階優化技術

除上述基礎技術外，以下高級策略可進一步優化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)能夠將Python模型轉換為經過優化的、可序列化的TorchScript程序。這種轉換通過優化內核啟動和減少運行時開銷，可帶來內存和性能的雙重提升：

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架構中實現更顯著的內存和性能優化。

總結

在GPU和云計算資源成本高昂的環境下，最大化利用現有計算資源至關重要。對于希望在有限計算資源條件下訓練或微調大型模型(如LLM或視覺Transformer)的研究者和開發者而言，掌握上述優化技術尤為重要。本文介紹的這些策略代表了研究人員和專業人士在資源受限條件下進行高效模型訓練的常用方法。