🔥Hết RAM GPU khi train model? Đây là cách khắc phục!🔥

Nếu bạn mới bắt đầu học Deep Learning với PyTorch, chắc hẳn bạn đã từng gặp phải lỗi đáng sợ này:

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

Đừng lo lắng! Đây là vấn đề cực kỳ phổ biến và có rất nhiều cách để khắc phục. Trong bài viết này, Zootopi sẽ chia sẻ những mẹo từ cơ bản đến nâng cao giúp bạn "sống sót" khi train model với GPU có bộ nhớ hạn chế.

Tại sao lại bị OOM?

Trước khi tìm cách khắc phục, hãy hiểu tại sao GPU của bạn lại "hết RAM". Khi train một model, GPU cần lưu trữ:

  1. Weights: Các trọng số/tham số của mô hình
  2. Activations: Kết quả trung gian của mỗi layer (cần cho backward pass)
  3. Gradients: Đạo hàm để cập nhật weights
  4. Optimizer states: Ví dụ Adam cần lưu momentum và variance cho mỗi tham số

Trong đó, activations thường chiếm nhiều bộ nhớ nhất - đặc biệt với batch size lớn. Đây là lý do tại sao giảm batch size thường là cách nhanh nhất để fix OOM!

Mẹo 1: Giảm Batch Size (Cách đơn giản nhất!)

Đây là bước đầu tiên bạn nên thử. Batch size càng lớn = càng nhiều activations cần lưu = càng nhiều RAM GPU.

from torch.utils.data import DataLoader
 
# Thay vì batch_size=32, thử giảm xuống
train_loader = DataLoader(dataset, batch_size=8)  # hoặc 4, 2, 1

Mẹo nhỏ: Nên chọn batch size là lũy thừa của 2 (8, 16, 32...) để tối ưu hiệu suất GPU.

Mẹo 2: Gradient Accumulation - "Batch size ảo"

"Nhưng giảm batch size thì model train không tốt!" - Đúng vậy, nhưng có cách khắc phục!

Gradient Accumulation cho phép bạn đạt được hiệu quả của batch size lớn mà không cần nhiều RAM. Ý tưởng: tích lũy gradients qua nhiều bước nhỏ, rồi mới cập nhật weights.

accumulation_steps = 8
optimizer.zero_grad()
 
for i, (x, y) in enumerate(train_loader):
    outputs = model(x)
    loss = criterion(outputs, y) / accumulation_steps  # Chia để trung bình
    loss.backward()
 
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Với code trên, nếu batch_size=4 và accumulation_steps=8, bạn đạt hiệu quả tương đương batch_size=32!

Tham khảo thêm: Hugging Face - Gradient Accumulation

Mẹo 3: Mixed Precision Training (AMP) - Giảm 50% bộ nhớ!

Đây là một trong những kỹ thuật hiệu quả nhất. Thay vì dùng float32 (32 bit), AMP sử dụng float16 (16 bit) cho hầu hết các phép tính - giảm gần 50% bộ nhớ và còn tăng tốc độ train!

import torch
from torch.amp import autocast, GradScaler
 
scaler = GradScaler()
 
for x, y in train_loader:
    optimizer.zero_grad()
 
    with autocast(device_type='cuda'):  # Tự động chuyển sang float16
        outputs = model(x)
        loss = criterion(outputs, y)
 
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Lưu ý: AMP hoạt động tốt nhất trên GPU có Tensor Cores (GTX 16xx, RTX series, Tesla V100 trở lên).

Tham khảo: PyTorch AMP Documentation

Mẹo 4: Gradient Checkpointing - Đổi thời gian lấy bộ nhớ

Kỹ thuật này giảm bộ nhớ bằng cách không lưu tất cả activations, mà tính toán lại khi cần trong backward pass.

Theo nghiên cứu của Chen et al. (2016), checkpointing có thể giảm bộ nhớ từ O(n) xuống O(√n) với chỉ khoảng 20-30% thời gian train thêm.

from torch.utils.checkpoint import checkpoint
 
class MyModel(nn.Module):
    def forward(self, x):
        # Wrap các layer tốn nhiều bộ nhớ
        x = checkpoint(self.heavy_layer1, x, use_reentrant=False)
        x = checkpoint(self.heavy_layer2, x, use_reentrant=False)
        return x

Nếu bạn dùng Hugging Face Transformers, chỉ cần bật flag:

from transformers import TrainingArguments
 
training_args = TrainingArguments(
    gradient_checkpointing=True,  # Chỉ cần bật flag này!
    # ... các args khác
)

Tham khảo: PyTorch Checkpoint Utilities

Mẹo 5: Giải phóng bộ nhớ đúng cách

Một số thói quen tốt để tránh "rò rỉ" bộ nhớ:

import torch
 
# 1. Dùng torch.no_grad() khi validation/inference
with torch.no_grad():
    val_outputs = model(val_data)
 
# 2. Xóa các tensor không cần thiết
del intermediate_tensor
torch.cuda.empty_cache()
 
# 3. Detach tensor khi chỉ cần giá trị (không cần gradient)
loss_value = loss.detach().item()
 
# 4. Move tensor về CPU nếu cần lưu lại
saved_output = outputs.detach().cpu()

Mẹo 6: Tối ưu DataLoader

DataLoader cấu hình không tốt có thể gây bottleneck và tốn RAM:

from torch.utils.data import DataLoader
 
train_loader = DataLoader(
    dataset,
    batch_size=16,
    pin_memory=True,         # Tăng tốc transfer CPU -> GPU
    num_workers=4,           # Load data song song (2-8 thường ổn)
    prefetch_factor=2,       # Số batch load trước
    persistent_workers=True  # Tái sử dụng workers giữa các epoch
)

Chú ý:

  • pin_memory=True giúp transfer data lên GPU nhanh hơn nhưng tốn thêm RAM CPU
  • num_workers quá cao có thể gây tràn RAM CPU

Tham khảo: PyTorch DataLoader

Nâng cao: Các công cụ cho model lớn

Nếu các mẹo trên vẫn chưa đủ, đây là một số công cụ mạnh mẽ hơn:

Hugging Face Accelerate

Thư viện giúp quản lý GPU dễ dàng, tích hợp sẵn AMP và distributed training:

from accelerate import Accelerator
 
accelerator = Accelerator(mixed_precision="fp16")
model, optimizer, train_loader = accelerator.prepare(
    model, optimizer, train_loader
)
 
# Train loop
for batch in train_loader:
    outputs = model(batch)
    loss = compute_loss(outputs)
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

Đặc biệt, Accelerate có find_executable_batch_size tự động giảm batch khi gặp OOM:

from accelerate.utils import find_executable_batch_size
 
@find_executable_batch_size(starting_batch_size=32)
def train(batch_size):
    # Code train của bạn - tự động retry với batch nhỏ hơn nếu OOM
    pass

Tham khảo: Accelerate Documentation

DeepSpeed ZeRO (cho multi-GPU)

Chia nhỏ optimizer states, gradients và weights ra nhiều GPU:

{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu" }
  }
}

Tham khảo: DeepSpeed ZeRO Tutorial

PyTorch FSDP (cho multi-GPU)

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
model = FSDP(model)

Tham khảo: PyTorch FSDP Tutorial

Kết luận

OOM là vấn đề ai cũng gặp khi học Deep Learning - đừng nản! Với các kỹ thuật trên, bạn hoàn toàn có thể train model trên GPU có bộ nhớ hạn chế. Hãy bắt đầu từ những cách đơn giản (giảm batch, bật AMP) rồi dần dần thử các kỹ thuật nâng cao hơn.

Chúc các bạn train model thành công!

Tài liệu tham khảo:

Đón xem những trải nghiệm và phân tích cụ thể từ Zootopi tại: