Gradient Accumulation — Giả lập batch lớn khi GPU nhỏ

Bạn muốn fine-tune Llama-3-8B với batch size 32 để gradient ổn định, nhưng RTX 4090 (24GB) chỉ chịu được batch=4 trước khi OOM. Mua A100 80GB? Không cần. Thay vào đó, chỉ cần một mẹo algorithm đơn giản: chia batch thành nhiều phần nhỏ, tính gradient từng phần nhưng "cộng dồn" thay vì cập nhật weights ngay. Đây là cách các open-source hacker train model 70B trên phần cứng consumer mà không cần data center.

Vấn đề

Training neural networks là trò chơi của batch size. Batch lớn (32–64) cung cấp gradient ổn định, hội tụ nhanh và generalize tốt hơn. Nhưng backpropagation đòi hỏi lưu trữ activations cho mọi sample trong batch để tính đạo hàm. Với LLM, batch size 64 có thể cần 16GB+ VRAM chỉ cho activations, chưa kể weights và optimizer states — vượt quá giới hạn phần cứng consumer.

Batch nhỏ (< 8) tạo ra gradient nhiễu (high variance), khiến loss dao động mạnh, hội tụ chậm, và dễ rơi vào local minima xấu. Nhưng GPU memory là rào cản cứng: bạn không thể nhét batch 64 vào card 24GB dù muốn. Cách tiếp cận cũ — giảm learning rate để bù — chỉ làm bước đi nhỏ lại theo hướng nhiễu, chứ không sửa được hướng.

Ý tưởng cốt lõi

Insight cốt lõi: gradient là tuyến tính.

Nếu loss là trung bình của các sample loss individual, thì: $\nabla_\theta \mathcal{L} = \nabla_\theta \left( \frac{1}{N} \sum_{i=1}^N \ell(x_i, \theta) \right) = \frac{1}{N} \sum_{i=1}^N \nabla_\theta \ell(x_i, \theta)$

Điều này có nghĩa: gradient của tổng bằng tổng của gradients. Thay vì tính gradient cho 64 sample cùng lúc (cần lưu 64 activations trong VRAM), bạn tính gradient cho 4 sample, cộng vào "hũ" tích lũy, rồi 4 sample tiếp... Sau 16 lần, bạn có tổng của 64 gradients — chính xác như batch 64 thực sự.

Hình dung như đổ nước vào thùng. Thay vì đổ 64 lít một lúc (tràn thùng), bạn đổ 4 lít × 16 lần. Tổng nước như nhau, nhưng thùng chỉ cần chứa 4 lít tại một thời điểm.

Đây chính là "aha moment": Gradient Accumulation không "xấp xỉ" batch lớn — nó tái tạo toán học chính xác, chỉ với chi phí tính toán nhiều hơn (N lần forward-backward) thay vì memory nhiều hơn.

Phá vỡ hiểu lầm phổ biến: Nhiều người nói "Sao không dùng batch nhỏ với learning rate thấp đi 16 lần?" Sai lầm. LR thấp chỉ làm bước đi nhỏ lại theo hướng gradient nhiễu (noisy). GA sửa hướng bằng cách lấy trung bình 16 estimates độc lập — variance giảm 16 lần, hướng gradient ổn định như batch thực sự. Đó là lý do tại sao GA khác biệt về chất lượng so với chỉ giảm LR.

Tại sao nó hoạt động

Về mặt toán học, đạo hàm của hàm tổng bằng tổng đạo hàm. Trong PyTorch, gradient tự động cộng dồn vào .grad mỗi khi gọi backward() trừ khi bạn chủ động gọi zero_grad().

Code pseudocode minh họa:

accumulation_steps = 16
effective_batch = 64
mini_batch = 4

for i in range(accumulation_steps):
    # Forward & loss trên mini-batch nhỏ
    loss = model(mini_batch[i]) / accumulation_steps
    loss.backward()  # Gradient tích lũy vào .grad, chưa update weights
    
# Chỉ update sau khi cộng dồn đủ 16 steps
optimizer.step()
optimizer.zero_grad()

Chia loss cho accumulation_steps để đảm bảo magnitude gradient khớp với true large batch (tránh gradient explosion khi cộng dồn).

Nhưng có chi tiết quan trọng: Floating point không kết hợp (non-associative). sum([grad_i for i in range(4)]) khác biệt nhỏ (~1e-6) so với grad(batch_of_4) do làm tròn số. Điều này tạo ra loss curve khác nhẹ so với true large batch, nhưng thực tế không đáng kể.

BatchNorm và thống kê: BatchNorm cập nhật running mean/variance theo mini-batch hiện tại, không theo effective batch. Dùng GA với BN gây lệch thống kê do mỗi mini-batch nhỏ có phân phối khác nhau. Giải pháp: dùng GroupNorm/LayerNorm (không phụ thuộc batch size), hoặc điều chỉnh BN momentum để ước lượng trên effective batch.

Ý nghĩa thực tế

So sánh thực tế giữa cách tiếp cận truyền thống và Gradient Accumulation:

Benchmarks: Với BLOOM-560M (560M params), batch 64 cần ~16GB VRAM; dùng GA với physical batch 1 chỉ cần ~4GB (theo Lightning AI, 2023). Trên thực tế, bạn có thể train Llama-2-7B trên RTX 4090 (24GB) với effective batch 32 nhờ kết hợp GA + QLoRA.

Ai đang dùng: HuggingFace Transformers (gradient_accumulation_steps), Axolotl (fine-tune LLM open-source), Lightning AI, DeepSpeed. Đặc biệt quan trọng trong QLoRA workflow — quantize xuống 4-bit rồi dùng GA để mô phỏng batch lớn cho ổn định. DPO cũng thường dùng GA vì preference data cần batch lớn để gradient ổn định.

Giới hạn:

• Không giảm wall-clock time (thực tế chậm hơn chút do overhead kernel launch)
• Cần scale learning rate theo sqrt(N) hoặc linear để duy trì convergence dynamics
• BatchNorm layers cần xử lý đặc biệt (dùng GroupNorm hoặc điều chỉnh momentum)
• Sai số số học nhỏ từ float32 non-associativity làm loss curve khác nhẹ

Đào sâu hơn

• Paper gốc: "Accumulated Gradient Normalization" (arXiv:1710.02368, 2017) — cơ sở lý thuyết về distributed training stability.
• Bài liên quan TroiSinh:
  • Cùng cụm:
    • QLoRA — Quantize 4-bit + LoRA, kết hợp hoàn hảo với GA để train model khổng lồ trên 1 GPU consumer
    • PEFT Overview — Các phương pháp efficient fine-tuning từ LoRA đến adapters
    • DPO — Alignment method thường dùng GA vì preference data cần batch lớn để ổn định gradient
    • ZeRO — Shard optimizer states khi GA không đủ, scale lên trillion params
    • Tensor Parallelism — Khi 1 GPU không đủ dù có GA, cắt weight matrix cho nhiều GPU
    • Pipeline Parallelism — Mỗi GPU 1 nhóm layer, pipeline mini-batch qua nhiều GPU
  • Đọc tiếp:
    • Training at Scale — Nền tảng về training model lớn (Level 0)
    • Constitutional AI — Phương pháp alignment tiên tiến (Level 2) thường kết hợp với efficient training
• External:
  • Blog Lightning AI: "Finetuning LLMs on a Single GPU Using Gradient Accumulation" — hướng dẫn thực hành chi tiết về implementation
  • Reddit r/MachineLearning: "Does gradient accumulation achieve anything?" — tranh luận sâu về GA vs small batch + low LR, và numerical precision issues