Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training

Huấn luyện các mô hình Mixture-of-Experts (MoE) ở ngữ cảnh dài thường thất bại khi bất kỳ thành phần nào vượt quá bộ nhớ thiết bị. Nghiên cứu này xác định bốn điểm nghẽn bộ nhớ chưa được giới hạn bởi các kế hoạch song song truyền thống: phân phối chuyên gia, phép chiếu từ vựng, biên độ kiểm tra gradient và trạng thái bộ tối ưu. Tác giả đề xuất bốn lịch trình tối ưu hóa gồm PipelinedLLEP, Ring-DTP, Selective Checkpoint Offload (SCO) và OffloadStreamAdamW để giới hạn các đỉnh bộ nhớ này mà vẫn giữ nguyên giá trị hàm mất mát và gradient. Thử nghiệm trên các mô hình từ 120B đến 667B tham số cho thấy phương pháp này cắt giảm đỉnh phân phối MoE tới 59,3%, đỉnh chiếu từ vựng 86,6%, tối ưu hóa nhanh hơn 2,05 lần, cho phép huấn luyện ngữ cảnh 1M với hiệu suất vượt trội.

In matched component tests, they cut the MoE dispatch peak by up to 59.3% without losing throughput, the vocabulary projection peak by 86.6%, and the offloaded optimizer step by 2.05times faster.

Composed on MoE models from 120B to 667B parameters, they train at 1M context length, 8--32times the reach of a tuned FSDP2 baseline, and up to 10.4times its throughput.