Chia trạng thái huấn luyện để giảm bộ nhớ mỗi worker

1Định nghĩa

ZeRO (Zero Redundancy Optimizer) và FSDP (Fully Sharded Data Parallel) là các chiến lược huấn luyện nhiều thiết bị bỏ phần bộ nhớ bị lặp lại giữa các bản sao. Trong huấn luyện song song dữ liệu truyền thống, mỗi GPU giữ một bản đầy đủ của trọng số, gradient (từ bước backward) và trạng thái bộ tối ưu (từ AdamW), nên cả nhóm lãng phí bộ nhớ HBM. ZeRO và FSDP cắt (shard) các trạng thái đó và chia cho N thiết bị, mỗi thiết bị chỉ giữ một phần.

Có ba cấp: ZeRO-1 chỉ chia trạng thái bộ tối ưu; ZeRO-2 chia thêm gradient; ZeRO-3 (tương đương FSDP đầy đủ) chia luôn cả trọng số. Đổi lại việc tiết kiệm, ở mỗi lớp của lượt forward và backward, các worker phải trao đổi qua mạng bằng các phép all-gather và reduce-scatter để ghép tạm phần tham số đang cần tính.

Ví dụ đời thường: ZeRO giống nhóm 4 sinh viên cùng học một bộ sách tham khảo đồ sộ nhưng ba lô mỗi người quá nhỏ. Thay vì cả bốn đều photo trọn 12 cuốn (trùng lặp, tốn chỗ như data parallel chuẩn), mỗi người chỉ giữ 3 cuốn và khi học tới chương nào thì chuyền cuốn đó cho cả nhóm xem (all-gather). Đổi lại là mất công chuyền sách qua lại; nếu đường chuyền chậm, nhóm học chậm hơn dù ba lô nhẹ hơn.

2Luồng

Đây là ô 4A, thuộc cột 4 · Scale & phương pháp. Mũi tên không đi tiếp vào bước kế tiếp mà bọc quanh ô 3B và 3D: nó đổi nơi lưu trữ của gradient và trạng thái bộ tối ưu, còn phép tính toán học của từng bước vẫn như cũ.

Sơ đồ kiến trúc huấn luyện · 22 ô (1A → 5C)

Kiến trúc huấn luyện LLM · Ánh xạ với Bài Lab tổng hợp

Thứ tự đọc: 1A → 1B ở cột 1, 2A → 2D ở cột 2, 3A → 3E ở cột 3; mũi tên next batch từ 3E quay về cột 1 và mũi tên load weights đi xuống băng 5 (5A → 5C). Cột 4 không phải bước kế tiếp: 4A, 4B tiết kiệm bộ nhớ, còn 4C–4G là các phương pháp huấn luyện thay thế nhau. Di chuột hoặc bấm vào từng ô để đọc pop-over.

Tất cả (22)1. Dữ liệu & nhãn2. Forward3. Học từ sai số4. Quy mô & phương pháp5. Suy luận

Ẩn khung overlay

1A1B1C2A2B2C2D3A3B3C3D3E4A4B4C4D4E4F4G5A5B5C

4A4 · Scale & Training Methods

GPU 0..3 · mỗi GPU chỉ giữ một phần của params / grads / optimizer

4A · ZeRO Sharding

Chia trọng số, gradient và trạng thái bộ tối ưu cho nhiều GPU để mỗi GPU chỉ giữ một phần.

Nhận gì

Các trạng thái huấn luyện của mô hình: trọng số, gradient (ô 3B) và trạng thái bộ tối ưu m, v (ô 3D).

Làm gì, vì sao

Khi huấn luyện trên nhiều GPU, mỗi GPU thường giữ đủ mọi thứ nên mô hình lớn không vừa. ZeRO (Zero Redundancy Optimizer) chia các trạng thái đó cho các GPU để mỗi GPU chỉ giữ một phần, đổi bộ nhớ lấy thêm giao tiếp mạng.

Ra gì

Không thêm phép tính mới: gradient và cập nhật vẫn như cũ, nhưng mỗi GPU chỉ giữ và cập nhật phần của mình rồi các phần được gom lại khi cần.

Ví dụ, cách đọc ô

Bốn khối GPU 0 đến GPU 3, mỗi khối ghi params / grads / optimizer. Ví dụ 7 tỉ tham số trên 8 GPU: huấn luyện data parallel thường cần 112 GB mỗi GPU, ZeRO-3 còn khoảng 14 GB (tính tay, chưa gồm activation).

Lưu ý

Con số trên chỉ gồm trạng thái mô hình, chưa gồm activation. ZeRO không cho biết tốc độ vì thời gian phụ thuộc vào mạng.

Đọc bài 4A · ZeRO Sharding↗Giai đoạn 8 trong Bài Lab: Bộ nhớ: ZeRO & recompute↗

flowchart TD
  step1["Chia batch và chia trạng thái"]
  step2["Forward và backward"]
  step3["Reduce-scatter gradient"]
  step4["Cập nhật phần mình giữ"]
  step1 --> step2
  step2 --> step3
  step3 --> step4

3Giải pháp

Vị trí trên sơ đồ: ô 4A "ZeRO" ở cột 4 · Scale & phương pháp. Khác các ô trước, ô này không thêm phép tính mới mà quyết định ai giữ cái gì: nó chạm vào gradient của ô 3B, trạng thái và bước cập nhật của ô 3D (AdamW), và liên quan trực tiếp tới ô 4B (tính lại activation) vì cả hai cùng nhằm tiết kiệm bộ nhớ.

Vấn đề cần giải. Khi huấn luyện, mỗi tham số cần nhiều bộ nhớ hơn bản thân nó, và trong data parallel chuẩn mỗi GPU phải giữ đầy đủ tất cả. Mô hình đủ lớn sẽ không vừa vào một GPU dù dữ liệu đã được chia nhỏ.

Khái niệm tối thiểu. Với số tham số Ψ, huấn luyện hỗn hợp độ chính xác (mixed precision) với AdamW cần cho mỗi tham số: 2 byte trọng số 16 bit, 2 byte gradient 16 bit và 12 byte trạng thái bộ tối ưu (4 byte bản sao trọng số chính 32 bit, 4 byte m, 4 byte v). Tổng 16 byte mỗi tham số, nằm ngoài activation và các bộ đệm. N là số GPU.

Công thức bộ nhớ mỗi GPU cho model states (theo bài báo ZeRO). Data parallel chuẩn: (2 + 2 + 12) × Ψ = 16Ψ. ZeRO-1: 2Ψ + 2Ψ + 12Ψ/N. ZeRO-2: 2Ψ + (2 + 12)Ψ/N. ZeRO-3: 16Ψ/N.

Ví dụ có số (minh họa). Mô hình 7 tỉ tham số (Ψ = 7 × 10⁹), 8 GPU, mỗi GPU 80 GB (giả định). Data parallel chuẩn: 16 × 7 = 112 GB mỗi GPU, không vừa dù chưa tính activation. ZeRO-1: (4 + 12/8) × 7 = 5,5 × 7 = 38,5 GB. ZeRO-2: (2 + 14/8) × 7 = 3,75 × 7 = 26,25 GB. ZeRO-3: 16/8 × 7 = 14 GB.

Cách đọc. Con số là bộ nhớ của riêng model states mỗi GPU; activation, bộ đệm giao tiếp, phân mảnh bộ nhớ còn phải cộng thêm, nên 14 GB không nghĩa là chỉ cần GPU 16 GB. Bước nhảy lớn nhất ở ví dụ là từ 112 xuống 38,5 GB, vì phần trạng thái bộ tối ưu (12 trong 16 byte) là phần lớn nhất.

Cái giá phải trả. Chia nhỏ thì khi cần dùng phải gom lại. Stage 1 và 2 chỉ cần reduce-scatter gradient và all-gather trọng số sau cập nhật, tổng giao tiếp xấp xỉ data parallel chuẩn. Stage 3 còn phải all-gather trọng số của từng lớp trước khi tính forward và backward; bài báo ZeRO ước tính tổng khoảng 1,5 lần giao tiếp của data parallel chuẩn. Vì vậy chọn stage theo cả dung lượng lẫn băng thông mạng chứ không chỉ theo dung lượng.

Điều có thể và không thể kết luận. Công thức cho thấy ZeRO-3 vừa vào bộ nhớ; nó không cho biết thời gian một bước tăng bao nhiêu, vì phụ thuộc mạng, kích thước lớp và mức chồng lấp giao tiếp với tính toán. Có thể kết hợp ZeRO với các kiểu chia khác (tensor, pipeline, expert parallel) theo mô hình, xem chia mô hình hoặc công việc. Đừng áp ngân sách bộ nhớ này cho suy luận: suy luận không có gradient hay trạng thái bộ tối ưu.

Để tự tính bộ nhớ và mô phỏng việc chia cho hai worker từ trạng thái AdamW thật, mở giai đoạn 8 · Bộ nhớ: ZeRO & recompute của Bài Lab tổng hợp. Lưu ý đó là ước tính byte, không phải chạy huấn luyện phân tán; hãy đoán trước tổng byte mỗi worker rồi đối chiếu.

4Profiling

• Bộ nhớ trạng thái mỗi worker

  Lập bảng trọng số, gradient, m, v, bản sao trọng số chính, activation và phần bộ đệm cho từng chiến lược trên cùng mô hình. Kiểm tra bằng công thức tay như ví dụ rồi so với bộ nhớ đỉnh đo được. Công thức không bao gồm phân mảnh bộ nhớ và các bộ đệm, nên con số đo thường cao hơn.

• Giao tiếp và khả năng mở rộng

  Đo thời gian all-gather và reduce-scatter, phần giao tiếp chồng lấp được với tính toán, và thời gian bước khi tăng số worker. Chỉ rút ra khả năng mở rộng khi giữ nguyên kích thước batch mỗi worker hoặc nêu rõ nó thay đổi.

• Tính tương đương

  Trên mô hình nhỏ, so loss và trọng số sau vài bước giữa data parallel chuẩn và ZeRO với cùng dữ liệu, kiểu số và hạt giống. Chúng nên trùng trong sai số số học; nếu không, thứ tự cộng gradient hoặc cắt chuẩn đang khác nhau.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Model states

Ba nhóm tensor phải giữ khi huấn luyện: trọng số (parameters), gradient và trạng thái bộ tối ưu (m, v và bản sao trọng số chính nếu có).

Zero Redundancy Optimizer (ZeRO)

Kỹ thuật của DeepSpeed bỏ phần dư thừa của data parallel: stage 1 chia trạng thái bộ tối ưu, stage 2 thêm gradient, stage 3 thêm cả trọng số.

Distributed Data Parallel (DDP)

Huấn luyện song song dữ liệu: mỗi thiết bị giữ một bản đầy đủ mô hình, xử lý phần batch riêng rồi kết hợp gradient.

Fully Sharded Data Parallel (FSDP)

Cơ chế của PyTorch chia trọng số, gradient và trạng thái bộ tối ưu giữa các thiết bị, về ý tưởng tương đương ZeRO stage 3.

All-gather và reduce-scatter

All-gather: mỗi thiết bị góp phần của mình để mọi thiết bị có bản đầy đủ. Reduce-scatter: cộng gradient từ mọi thiết bị rồi mỗi thiết bị chỉ nhận đúng phần của mình.

7Bài liên quan

8Nguồn