ZeRO — Shard optimizer state, train trillion-param models

Giải pháp bộ nhớ cho model trăm tỷ tham số: ZeRO chia nhỏ optimizer state qua nhiều GPU, loại bỏ sự dư thừa của Data Parallel thông thường.

Training model 100B+ parameters trên GPU consumer tưởng chừng bất khả thi vì bộ nhớ — chỉ riêng trạng thái optimizer (Adam) đã ngốn 600GB. ZeRO (Zero Redundancy Optimizer) phá vỡ giới hạn này bằng cách chia nhỏ trạng thái optimizer qua cluster, giúp bạn train model nghìn tỷ tham số trên phần cứng thông thường mà không cần Model Parallelism phức tạp.

Vấn đề

Data Parallelism thông thường (DDP) là "giải pháp lười biếng" dẫn đến lãng phí bộ nhớ khủng khiếp. Logic đơn giản: nhân bản toàn bộ model lên mỗi GPU, cho mỗi GPU xử lý 1/N batch, rồi average gradients. Nhưng điều này có nghĩa là mỗi GPU giữ một bản sao đầy đủ: weights, gradients, và optimizer states.

Với Adam optimizer, mỗi parameter cần thêm 2 bản sao nữa (momentum và variance) để tính adaptive learning rate. Kết quả: model 175B parameters cần khoảng 320GB chỉ cho optimizer states (fp32), chưa kể 350GB cho weights (fp16) và activations. Không GPU nào đủ VRAM để chứa. Giải pháp truyền thống là Model Parallelism (MP) — chia model theo chiều dọc/layer giữa các GPU — nhưng MP phức tạp, dễ underutilize GPU (pipeline bubble), và khó debug.

Ý tưởng cốt lõi

Chia sẻ tài liệu tham khảo thay vì mỗi người mang một thư viện.

Hãy tưởng tượng bạn và 7 người bạn cùng viết một bách khoa toàn thư. Cách làm cũ (DDP): mỗi người mang theo toàn bộ thư viện từ A-Z, cộng thêm 2 bản sao ghi chú riêng, dù chỉ đang sửa mục "Aardvark". ZeRO hỏi: "Tại sao không chia thư viện ra? Bạn giữ A-H, bạn kia giữ I-P..."

Đây chính là "aha moment": Chúng ta đã replicate dữ liệu một cách vô nghĩa. Trong DDP, dù mỗi GPU chỉ cập nhật 1/N parameters, nó vẫn giữ toàn bộ optimizer states. ZeRO loại bỏ sự dư thừa này bằng cách partition optimizer states (và dần dần cả gradients, weights) qua các GPU.

ZeRO có 3 cấp độ tăng dần:

ZeRO-1 (Optimizer State Partitioning): Chia nhỏ optimizer states. Mỗi GPU chỉ giữ 1/N bản sao. Khi cần update, gather tạm thời, update xong scatter lại.
ZeRO-2 (Add Gradient Partitioning): Chia thêm gradients. Khi backward, reduce-scatter gradients ngay lập tức đến GPU chịu trách nhiệm shard đó, thay vì accumulate đầy đủ rồi mới chia.
ZeRO-3 (Add Parameter Partitioning): Chia cả weights. Mỗi GPU chỉ giữ 1/N parameters thực sự, all-gather khi cần forward/backward, sau đó ngay lập tức free bộ nhớ.

Bạn có thể thấy mình đã bỏ lỡ detail nào đó — nhưng không, chỉ đơn giản là "đừng replicate nữa". Không có ma thuật nào khác.

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

Hiệu quả của ZeRO dựa trên hai trụ cột: tính chất sparse của gradients và cơ chế collective communication hiệu quả.

Gradients là sparse across parameter space nhưng dense across layers. Khi backward qua layer L, gradients chỉ tác động đến một phần parameters. ZeRO-2 khai thác điều này bằng cách reduce-scatter ngay lập tức: thay vì chờ accumulate đầy đủ gradients rồi mới chia (như kế toán gom tiền mặt cuối ngày), ZeRO chia ngay khi có (như mỗi người nộp thuế điện tử ngay lập tức).

Adam states là kẻ ngốn bộ nhớ thực sự. Một model 100B parameters cần 200GB cho weights (fp16), nhưng cần thêm 400GB cho Adam's momentum và variance (fp32). ZeRO-1 chia nhỏ chúng: GPU 0 giữ momentum của shard 0, GPU 1 giữ momentum của shard 1. Việc all-gather chỉ diễn ra trong thời gian update, sau đó bộ nhớ được giải phóng.

Communication được overlap với computation thông qua NCCL (NVIDIA Collective Communications Library). Khi GPU 0 đang tính forward cho layer tiếp theo, nó có thể overlap việc gather parameters với compute. Điều này giấu đi latency của network, đặc biệt khi dùng NVLink hoặc InfiniBand.

Ý nghĩa thực tế

ZeRO biến "impossible" thành "commodity":

Scale: ZeRO-3 cho phép train model 10x lớn hơn baseline. Với 400 GPU, có thể train model nghìn tỷ (trillion) parameters. ZeRO-Infinity mở rộng lên 32 nghìn tỷ.
Throughput: Linear scaling đến 6.8x với 8 GPU (so với 1.4x của DDP thường do memory thrashing).
Dễ dàng: Không cần thiết kế lại model architecture như Model Parallelism. Chỉ cần wrap optimizer với DeepSpeed.

So sánh với các phương pháp khác:

Phương pháp	Memory	Complexity	GPU Utilization	Khi nào dùng
DDP thường	Kém (replicate 100%)	Thấp	Tốt	Model < 10B
ZeRO-1/2/3	Tốt (shard 1/N)	Trung bình	Rất tốt	Model 10B-1T+
Pipeline Parallelism	Trung bình	Cao	Có bubble	Sequential layers
Tensor Parallelism	Tốt	Cao	Cần NVLink	Single node multi-GPU

Trade-offs:

Communication overhead: Tăng lưu lượng mạng (all-gather/reduce-scatter), nhưng thường bị che lấp bởi compute ở batch size lớn.
Activation memory: ZeRO không giảm activation memory (cần activation checkpointing riêng).
ZeRO-3 overhead: Với batch nhỏ, việc liên tục gather/scatter weights có thể chậm hơn DDP thuần.

Ai đang dùng: Microsoft DeepSpeed (nguồn gốc), HuggingFace Transformers (tích hợp sẵn), EleutherAI (GPT-Neo, GPT-J), BigScience (BLOOM 176B), và hầu hết các framework fine-tune lớn hiện nay.

Đào sâu hơn

Paper gốc: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari et al., 2019).
Implementation: DeepSpeed GitHub — thư viện chuẩn với ZeRO-Offload và ZeRO-Infinity.
Bài liên quan TroiSinh:
- Cùng cụm: PEFT Overview — Các phương pháp fine-tune hiệu quả tham số (LoRA, Adapters).
- Cùng cụm: QLoRA — Quantization + LoRA cho fine-tune 65B trên 1 GPU consumer.
- Cùng cụm: Gradient Accumulation — Giả lập batch lớn khi GPU nhỏ.
- Đọc tiếp: Tensor Parallelism — Cắt weight matrix cho nhiều GPU (thường kết hợp với ZeRO).
- Đọc tiếp: Pipeline Parallelism — Mỗi GPU 1 nhóm layer.
- Đọc tiếp: DPO — RLHF không cần reward model (thường train cùng ZeRO cho model alignment lớn).