Moonshot AI mã nguồn mở MoonEP: Thư viện tối ưu hóa song song chuyên gia cho mô hình MoE

Moonshot AI vừa mã nguồn mở MoonEP, một thư viện giao tiếp Expert Parallelism (EP) dành cho các khối lượng công việc Mixture-of-Experts (MoE) phân tán. Nhóm phát triển công bố đây là thư viện được xây dựng nhằm giúp việc giao tiếp song song chuyên gia (expert-parallel communication) trở nên hiệu quả hơn ở quy mô lớn. Thư viện này được phát hành theo giấy phép MIT.

MoonEP ra mắt trong khuôn khổ sự kiện Kimi K3 Open Day. Cùng với trọng số mô hình K3 và báo cáo kỹ thuật, Moonshot đã phát hành ba cơ sở mã hạ tầng: MoonEP, FlashKDA và AgentEnv. FlashKDA trước đó đã được mã nguồn mở; còn MoonEP và AgentEnv được công bố cùng với đợt phát hành này. MoonEP là một trong những cải tiến đứng sau tuyên bố cải thiện 2,5 lần hiệu suất mở rộng cho Kimi K3, một mô hình MoE với 2,8 nghìn tỷ tham số, có khả năng thị giác gốc (native vision) và cửa sổ ngữ cảnh 1 triệu token.

Vấn đề mà MoonEP nhắm đến

Trong song song chuyên gia (expert parallelism), một bộ định tuyến (router) sẽ gửi mỗi token đến top-K chuyên gia, vốn nằm trên các rank khác nhau. Các bộ định tuyến hiếm khi đạt được sự cân bằng. Một số chuyên gia nhận được nhiều token hơn hẳn những chuyên gia khác.

Kho lưu trữ định lượng sự lệch lạc (skew) bằng maxvio, được định nghĩa là max_e (T_e / T̄) − 1, trong đó T_e là số token được định tuyến đến chuyên gia e và T̄ là số lượng kỳ vọng trong điều kiện cân bằng hoàn hảo. Chỉ số maxvio bằng 0 nghĩa là cân bằng hoàn hảo.

Chi phí do mất cân bằng mang tính cấu trúc chứ không phải ngẫu nhiên. Độ trễ của một tập hợp được quyết định bởi thành viên chậm nhất, vì vậy rank "nóng" nhất (nhiều tải nhất) sẽ quyết định thời gian lặp. Tệ hơn nữa, số lượng token trên mỗi rank thay đổi sau mỗi bước. Những hình dạng kích hoạt động (dynamic activation shapes) đó làm phân mảnh bộ nhớ GPU và buộc phải đồng bộ hóa host theo từng lớp.

Ý tưởng cốt lõi: các chuyên gia dự phòng động

Điểm đáng chú ý nhất của MoonEP là một bất biến cứng (hard invariant). Mọi rank đều nhận chính xác S × K token, bất kể việc định tuyến bị lệch đến mức nào — trong đó S là số token đầu vào trên mỗi rank và K là số top-k được định tuyến trên mỗi token.

Nó đạt được điều này bằng cách lập kế hoạch cho một số lượng nhỏ các chuyên gia dự phòng trực tuyến (online), trực tiếp từ các đầu ra hiện tại của bộ định tuyến. Các chuyên gia được sao chép đó sẽ được tải trước (prefetch) trước khi quá trình tính toán chuyên gia diễn ra. Trong quá trình truyền ngược (backward pass), gradient của chúng được giảm (reduce) trở lại các rank gốc.

Thiết kế này được thiết lập dựa trên ba đặc tính:

Công cụ giải thích tương tác bên dưới sẽ tính toán kích thước bộ đệm và vùng nhớ tải trước (prefetch-pool) trực tiếp từ cấu hình mà bạn kiểm soát.

Hợp đồng bộ nhớ

Hợp đồng của MoonEP với một framework huấn luyện hoặc suy luận rất cụ thể: một tensor trọng số bộ nhớ đối xứng liên tục (contiguous symmetric-memory weight tensor) cho mỗi phép chiếu chuyên gia (expert projection), cộng với cu_seqlens do bộ lập kế hoạch tạo ra. Nhóm GEMM VM tiêu thụ một tensor trọng số [E+B, H, H'], trong đó E là tổng số chuyên gia được định tuyến, B là các khe tải trước trên mỗi rank, H là kích thước ẩn (hidden size) và H' là kích thước trung gian FFN của chuyên gia. Giá trị cu_seqlens[E+B] được trả về bởi quá trình điều phối sẽ chọn ra các hàng chuyên gia nào đang hoạt động.

Tính liên tục là một yêu cầu bắt buộc, vì nhóm GEMM giải quyết các chuyên gia hoàn toàn bằng chỉ số hàng. Bố cục được phân tách rõ ràng:

Các khe tải trước được lấy từ một vùng nhớ chung toàn cục của tiến trình (process-global pool) được chia sẻ bởi tất cả các lớp. Chi tiết này rất quan trọng: chi phí bộ nhớ bổ sung là B trọng số chuyên gia cho mỗi phép chiếu tính theo tổng thể, chứ không phải trên mỗi lớp.

Cách bạn thiết lập B phụ thuộc vào khối lượng công việc. Huấn luyện phải sử dụng B = E/R, vì bộ lập kế hoạch sao chép các chuyên gia từ tối đa một nhóm gốc từ xa trên mỗi rank. Giới hạn đó đảm bảo mọi chuyên gia mà nhóm GEMM chạm tới đều là cục bộ. Suy luận cho phép B < E/R, và tệp README khuyến nghị B = 3–4. Nếu một rank cần nhiều chuyên gia từ xa khác biệt hơn mức B, nhóm GEMM sẽ đọc các trọng số tràn (overflow weights) trực tiếp từ rank gốc thông qua ánh xạ đối xứng — chậm hơn một chút nhưng không ảnh hưởng đến tính chính xác.

Quá trình huấn luyện phản chiếu bố cục trọng số ở định dạng fp32 với một bộ đệm grad [E+B, H, H'] cho mỗi phép chiếu. Quan trọng là, các hàng [E, E+B) được hỗ trợ bởi một bộ đệm giảm (reduce buffer) riêng biệt, không phải bởi các grad tham số. Gradient của các chuyên gia được sao chép chỉ là tạm thời và phải ẩn đối với quá trình giảm grad của chính framework. Mỗi rank ánh xạ tất cả R bộ đệm giảm thành một view [R, B, H, H'], sau đó reduce_grad đọc các khe chuyên gia của chính nó từ mọi rank qua NVLink, tích lũy vào grad tham số cục bộ và xóa các khe đã tiêu thụ.

Điểm chuẩn so với DeepEP v2

Cả hai điểm chuẩn được công bố đều chạy trên H20 với EP=8, quét qua sự mất cân bằng của bộ định tuyến. Tập lệnh so sánh benchmarks/bench_vs_deepep.py mặc định S=8192, E=384, H=7168, K=8, H'=2048 và 32 SM, với các mục tiêu maxvio là 0.2, 1, 10 và 20. Cả hai thư viện đều nhận được ma trận định tuyến giống hệt nhau từ một seed chung.

Ba phát hiện đã được báo cáo trên trang GitHub của họ. Zero copy giúp giao tiếp thô nhanh hơn bằng cách loại bỏ việc sao chép từ comm-buffer sang user-buffer vốn chiếm ưu thế ở phần kết thúc, vì vậy thời gian giao tiếp của MoonEP luôn thấp hơn DeepEP v2 ở mọi mức độ mất cân bằng. Sự cân bằng hoàn hảo giúp MoonEP gần như miễn nhiễm với sự lệch lạc: thời gian giao tiếp của nó gần như phẳng khi maxvio tăng, trong khi DeepEP v2 — vốn có độ trễ bị quyết định bởi rank nóng nhất — lại suy giảm đều đặn.

Những điểm chính cần lưu ý

Nguồn: MoonshotAI/MoonEP trên GitHub và thông báo từ @Kimi_Moonshot

Michal Sutter là một chuyên gia khoa học dữ liệu với bằng Thạc sĩ Khoa học Dữ liệu từ Đại học Padova. Với nền tảng vững chắc về phân tích thống kê, học máy và kỹ thuật dữ liệu, Michal xuất sắc trong việc chuyển đổi các tập dữ liệu phức tạp thành những thông tin chi tiết có thể hành động được.