SkewAdam cắt giảm hơn 60% bộ nhớ huấn luyện Mixture-of-Experts và mang lại sự gia tăng đáng kể về độ chính xác, cho phép các nhà phát triển chạy một mô hình MoE 6,78 tỷ tham số trên một GPU 40 GB duy nhất.

Tại sao huấn luyện MoE lại gặp phải rào cản bộ nhớ

Kiến trúc Mixture-of-Experts điều hướng mỗi đầu vào qua một tập hợp con nhỏ các mạng con "chuyên gia" (expert) trong khi vẫn duy trì một khung xương (backbone) dày đặc. Lợi ích là một mô hình có thể mở rộng lên hàng tỷ tham số mà không làm tăng chi phí tính toán theo tỷ lệ thuận. Trong thực tế, bộ tối ưu hóa (optimizer) — cụ thể là biến thể AdamW mà hầu hết các nhóm đều sử dụng — chiếm phần lớn RAM của GPU. Đối với một mô hình 6,78 tỷ tham số, chỉ riêng các tensor momentum và variance của bộ tối ưu hóa đã yêu cầu khoảng 50 GB. Cộng thêm các giá trị kích hoạt (activations) và trọng số mô hình (model weights), bộ nhớ đỉnh sẽ lên tới 81,4 GB, buộc phải sử dụng các thiết lập đa GPU hoặc các lịch trình huấn luyện chậm hơn.

SkewAdam làm điều gì khác biệt

SkewAdam không tạo ra một quy tắc học mới. Nó sắp xếp lại nơi lưu trữ các moment của Adam:

  • Khung xương (backbone) dày đặc giữ lại momentum với độ chính xác đầy đủ (full-precision), giúp duy trì các cập nhật mượt mà mà các lớp dày đặc cần.
  • Ngân hàng chuyên gia (expert bank) lưu trữ một bản xấp xỉ phân tách (factored approximation) của variance, giúp cắt giảm dữ liệu cần lưu trữ cho mỗi chuyên gia.
  • Bộ định tuyến (router), thành phần quyết định kích hoạt chuyên gia nào, sẽ giữ lại một ước tính moment bậc hai (second-moment estimate) chính xác.

Bằng cách coi ba thành phần này là các "tầng" (tiers) riêng biệt, bộ tối ưu hóa sẽ loại bỏ độ chính xác dư thừa ở những nơi ít quan trọng hơn và giữ lại ở những nơi quan trọng nhất.

Tác động có thể đo lường được

Chạy cùng một mô hình MoE 6,78 tỷ tham số với SkewAdam mang lại:

  • Bộ nhớ GPU đỉnh: 31,3 GB (giảm từ 81,4 GB)
  • Dấu chân trạng thái bộ tối ưu hóa (Optimizer-state footprint): 1,29 GB (giảm từ 50 GB)

Trạng thái đã được giảm thiểu này có thể nằm gọn trong một bộ tăng tốc 40 GB duy nhất.

Gia tăng độ chính xác, không chỉ là tiết kiệm

Việc cắt giảm bộ nhớ thường làm giảm chất lượng mô hình, nhưng SkewAdam đã cải thiện độ nhiễu (perplexity) — một chỉ số tiêu chuẩn của mô hình ngôn ngữ — từ 126,8 (AdamW) xuống còn 108,4. Nó cũng vượt trội hơn Muon (120,2) và Lion (393,7) trong cùng một tác vụ. Các tác giả cho biết sự gia tăng này đến từ việc duy trì momentum trên cả ba tầng; các phương pháp loại bỏ hoàn toàn momentum, như Adafactor, đều tụt lại phía sau.

Các câu hỏi còn bỏ ngỏ

Các kết quả của SkewAdam được chứng minh trên mô hình 6,78 tỷ tham số. Vẫn chưa rõ liệu các tỷ lệ trạng thái bộ nhớ tương tự có áp dụng cho các mô hình lớn hơn một bậc hoặc cho các tác vụ ngoài mô hình hóa ngôn ngữ hay không.

Những điều cần theo dõi

  • Các điểm chuẩn (benchmarks) trên các cấu hình MoE lớn hơn (hàng chục tỷ tham số).
  • Việc áp dụng bởi các khung làm việc (frameworks) mã nguồn mở và sự góp mặt trong các kịch bản huấn luyện phổ biến.
  • Phản hồi từ cộng đồng về các đánh đổi tiềm ẩn, chẳng hạn như tốc độ hội tụ hoặc tính ổn định dưới các lịch trình tốc độ học (learning-rate schedules) khác nhau.

Nguồn: bài đăng của nhà phát triển chi tiết về thiết kế và kết quả thực nghiệm của bộ tối ưu hóa.