Gradient Accumulation

Gradient accumulation là kỹ thuật huấn luyện chạy một số batch nhỏ, gọi là micro-batch, qua lượt truyền xuôi và lan truyền ngược, cộng dồn gradient rồi mới cập nhật trọng số model. Lần cập nhật này xấp xỉ một bước huấn luyện trên batch lớn bằng tổng các micro-batch, trong khi GPU chỉ cần lưu activation của một micro-batch tại một thời điểm. Kỹ thuật này cho phép thiết bị bị giới hạn bộ nhớ huấn luyện với batch size mà cấu hình yêu cầu, mà không cần giảm độ phân giải ảnh hoặc thu nhỏ model. Batch size hiệu dụng bằng kích thước micro-batch nhân với số bước tích lũy rồi nhân với số thiết bị; do đó, micro-batch 4 ảnh được tích lũy qua 8 bước trên 4 GPU hoạt động tương tự batch 128.

Cách tích lũy gradient hoạt động#

Trong PyTorch và hầu hết framework khác, mỗi lần gọi backward() sẽ cộng gradient mới vào gradient đã lưu trên các parameter thay vì ghi đè chúng. Gradient accumulation tận dụng trực tiếp cơ chế này: chạy lượt truyền xuôi và ngược cho từng micro-batch, gọi optimizer khi đạt số micro-batch đã đặt, sau đó xóa gradient và bắt đầu cửa sổ tiếp theo:

Hầu hết hàm mất mát lấy trung bình trên các mẫu trong batch, vì vậy loss của mỗi micro-batch được chia cho số bước để biến tổng gradient thành giá trị trung bình. Với cách chia tỷ lệ này, bốn micro-batch gồm 8 mẫu tạo ra gradient giống với một batch gồm 32 mẫu:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(16, 1)
x, y = torch.randn(32, 16), torch.randn(32, 1)
loss_fn = nn.MSELoss()

# Reference: one batch of 32 samples
loss_fn(model(x), y).backward()
batch_grad = model.weight.grad.clone()
model.zero_grad()

# Four micro-batches of 8: backward() adds into .grad, and dividing each loss by 4 turns the sum into a mean
steps = 4
for x_micro, y_micro in zip(x.chunk(steps), y.chunk(steps)):
    (loss_fn(model(x_micro), y_micro) / steps).backward()

print("Matches the full-batch gradient:", torch.allclose(model.weight.grad, batch_grad, atol=1e-6))
print(f"Largest difference: {(model.weight.grad - batch_grad).abs().max():.1e}")

Lệnh in ra Matches the full-batch gradient: True với sai khác lớn nhất là 1.6e-07; đây là sai số làm tròn dấu phẩy động do cộng cùng các số theo thứ tự khác nhau.

Accumulation khác batch lớn ở điểm nào#

Kết quả tương đương đúng với loss lấy trung bình trên các mẫu độc lập, nhưng không còn đúng trong ba trường hợp phổ biến:

• Batch normalization: Tính trung bình và phương sai trên từng micro-batch, vì vậy thống kê từ 4 ảnh nhiễu hơn thống kê từ 128 ảnh.
• Chuỗi có độ dài thay đổi: Khi model ngôn ngữ lấy trung bình loss trên token trong mỗi micro-batch, micro-batch có ít token lại có trọng số ngang với micro-batch dài.
• Gradient clipping: Chỉ thực hiện một lần trên gradient đã tích lũy, ngay trước bước optimizer; cắt gradient ở từng micro-batch sẽ làm thay đổi kết quả.

Hugging Face và Unsloth đã khắc phục sai lệch số lượng token vào tháng 10 năm 2024 đối với quá trình tinh chỉnh mô hình ngôn ngữ lớn (LLM) bằng cách chia tổng loss cho tổng số token của toàn bộ cửa sổ tích lũy. Accumulation cũng không tăng tốc: cùng số mẫu vẫn đi qua mạng lần lượt theo từng micro-batch, và micro-batch nhỏ sử dụng GPU kém hiệu quả hơn. Trong huấn luyện phân tán, DistributedDataParallel của PyTorch cung cấp ngữ cảnh no_sync(), bỏ qua bước all-reduce gradient ở mọi micro-batch trừ micro-batch cuối, giúp giảm chi phí truyền thông. Batch hiệu dụng lớn hơn thường cần learning rate cao hơn: các nhà nghiên cứu Facebook huấn luyện ResNet-50 trên ImageNet trong một giờ với batch 8.192 ảnh trên 256 GPU bằng cách tăng learning rate tuyến tính theo batch size và tăng dần learning rate trong giai đoạn warmup.

Các kỹ thuật liên quan#

• Gradient checkpointing: Tính lại activation trong quá trình lan truyền ngược thay vì lưu chúng, giúp tiết kiệm bộ nhớ cho mỗi micro-batch nhưng tiêu tốn thêm compute.
• Độ chính xác hỗn hợp: Lưu activation ở định dạng 16 bit, tạo thêm không gian cho micro-batch lớn hơn. Loss scaler của kỹ thuật này chỉ unscale gradient một lần sau micro-batch cuối cùng.
• Song song pipeline: Chia mỗi batch thành các micro-batch, luân chuyển chúng qua các giai đoạn của model trên những GPU khác nhau rồi tích lũy gradient trước một bước optimizer duy nhất.

Quá trình huấn luyện Ultralytics YOLO26 mặc định tích lũy gradient: trong chế độ huấn luyện, trainer cập nhật optimizer sau mỗi round(nbs / batch) batch; vì vậy, giá trị mặc định batch=16 cùng batch size danh nghĩa nbs=64 tạo ra bốn batch cho mỗi lần cập nhật, còn batch=8 tạo ra tám batch. Số bước tăng dần từ một trong giai đoạn warmup, còn weight decay được điều chỉnh theo batch × accumulation / nbs, nên hyperparameter mặc định dành cho batch 64 ảnh vẫn áp dụng được trên GPU chỉ chứa vài ảnh tại một thời điểm.