Moonshot AI مفتوح المصدر MoonEP: مكتبة موازية للخبراء متوازنة تمامًا لتدريب وزارة التربية والتعليم

لدى Moonshot AI مكتبة MoonEP مفتوحة المصدر، وهي مكتبة اتصالات متوازية للخبراء (EP) لأعباء العمل الموزعة لمزيج الخبراء (MoE). أعلن الفريق عن الإصدار كمكتبة مصممة لجعل الاتصالات الموازية بين الخبراء أكثر كفاءة على نطاق واسع. يتم الشحن بموجب ترخيص MIT.

وصل MoonEP كجزء من Kimi K3 Open Day. إلى جانب أوزان طراز K3 والتقرير الفني، أصدرت Moonshot ثلاث قواعد بيانات للبنية التحتية: MoonEP، وFlashKDA، وAgentEnv. كان FlashKDA مفتوح المصدر بالفعل؛ تم نشر MoonEP وAgentEnv مع هذا الإصدار. يعد MoonEP أحد الابتكارات وراء التحسين المزعوم بمقدار 2.5× في كفاءة التوسع لـ Kimi K3، وهو نموذج MoE ذو 2.8 تريليون معلمة مع رؤية أصلية ونافذة سياق رمزية 1M.

المشكلة التي يستهدفها MoonEP

في التوازي الخبير، يرسل جهاز التوجيه كل رمز مميز إلى خبرائه المتميزين، الذين يعيشون في رتب مختلفة. نادراً ما تكون أجهزة التوجيه متوازنة. يحصل بعض الخبراء على رموز أكثر بكثير من غيرهم.

يقوم المستودع بقياس الانحراف باستخدام com.maxvio، تم تعريفها على أنها max_e (T_e / T̄) − 1، أين T_e يتم توجيه الرموز إلى الخبراء e و T̄ هو العدد المتوقع في ظل التوازن المثالي. الحد الأقصى 0 يعني التوازن المثالي.

تكاليف الخلل هي تكاليف هيكلية وليست عرضية. يتم تحديد زمن الوصول الجماعي بواسطة أبطأ مشارك فيه، وبالتالي فإن الترتيب الأكثر سخونة يحدد وقت التكرار. والأسوأ من ذلك أن عدد الرموز المميزة لكل رتبة يتغير في كل خطوة. تعمل أشكال التنشيط الديناميكية هذه على تجزئة ذاكرة GPU وتفرض مزامنة المضيف لكل طبقة.

الفكرة الأساسية: خبراء ديناميكيون زائدون عن الحاجة

الشخصية الرئيسية لـ MoonEP هي شخصية ثابتة. كل رتبة تتلقى بالضبط S × K الرموز المميزة، بغض النظر عن مدى انحراف التوجيه – أين S هي رموز الإدخال لكل رتبة و K يتم توجيهه إلى أعلى k لكل رمز مميز.

ويحقق ذلك من خلال التخطيط لعدد صغير من الخبراء الزائدة عن الحاجة عبر الإنترنت، مباشرة من مخرجات جهاز التوجيه الحالي. يتم جلب هؤلاء الخبراء المكررين مسبقًا قبل حساب الخبراء. في التمريرة الخلفية، يتم تقليص تدرجاتهم إلى صفوفهم الأصلية.

تم إعداد التصميم في ثلاث خصائص:

1. التوازن المثالي: ال S × K الضمان أعلاه، عبر خبراء زائدين عن الحاجة عبر الإنترنت.
2. التخطيط عبر الإنترنت: نواة تخطيط GPU شبه مثالية مع حمل لا يكاد يذكر. يتم تنفيذه في CUTLASS CuTe DSL؛ setup.py دبابيس nvidia-cutlass-dsl==4.4.2.
3. نسخة صفرية وأشكال ثابتة: تنصهر تبديل / غير تبديل. تتم كتابة الرموز مباشرة في مواقعها المجمعة من قبل الخبراء في الرتب البعيدة، ويتم إرجاع طرق العرض المؤقتة إلى الحساب. ثابت فقط S × K هناك حاجة إلى المخزن المؤقت، والأشكال المعروفة بشكل ثابت تقضي على مزامنة مضيف MoE لكل طبقة.

يحسب الشرح التفاعلي أدناه المخزن المؤقت الناتج وأحجام تجمع الجلب المسبق مباشرة من التكوين الذي تتحكم فيه.

عقد الذاكرة

إن عقد MoonEP مع إطار التدريب أو الاستدلال محدد: موتر وزن ذاكرة متماثل واحد متجاور لكل إسقاط خبير، بالإضافة إلى منتج مخطط cu_seqlens. تستهلك مجموعة VM GEMM واحدًا [E+B, H, H'] موتر الوزن، حيث E هو إجمالي الخبراء الموجهين، B هي فتحات الجلب المسبق لكل رتبة، H هو حجم مخفي، و H' هو الحجم المتوسط ​​للخبراء FFN. ال cu_seqlens[E+B] عاد بواسطة dispatch يحدد صفوف الخبراء النشطة.

يعد التواصل متطلبًا صعبًا، لأن مجموعة GEMM تخاطب الخبراء فقط من خلال فهرس الصف. ينقسم التخطيط بشكل نظيف:

• الصفوف [0, E) hold all ranks’ local experts, E/R rows per rank. Each chunk physically is the home rank’s parameter memory, mapped everywhere via symmetric memory.
• Rows [E, E+B) are local prefetch slots, filled by buffer.prefetch_weight.

The prefetch slots draw from a process-global pool shared by all layers. That detail matters: the extra memory cost is B expert weights per projection in total, not per layer.

How you set B depends on the workload. Training must use B = E/R, because the planner duplicates experts from at most one remote home group per rank. That bound guarantees every expert the group GEMM touches is local. Inference allows B < E/R, and the README recommends B = 3–4. If a rank needs more distinct remote experts than B, the group GEMM reads overflow weights straight from the home rank through the symmetric mapping — slightly slower, with no impact on correctness.

Training mirrors the weight layout in fp32 with a [E+B, H, H'] المخزن المؤقت غراد في الإسقاط. بشكل حاسم، الصفوف [E, E+B) are backed by a separate reduce buffer, not by the parameter grads. Duplicated experts’ gradients are temporary and must stay invisible to the framework’s own grad reduce. Each rank maps all R reduce buffers as one [R, B, H, H'] عرض، ثم reduce_grad يقرأ خانات الخبراء الخاصة به من كل رتبة عبر NVLink، ويتراكم في المعلمة المحلية، ويستهلك الأصفار الفتحات.

المعايير ضد DeepEP v2

يعمل كلا المعيارين المنشورين على H20 مع EP=8، مما يؤدي إلى خلل كبير في جهاز التوجيه. السيناريو المقارنة benchmarks/bench_vs_deepep.py الافتراضي ل S=8192, E=384, H=7168, K=8, H'=2048و 32 SM، مع أهداف maxvio تبلغ 0.2 و1 و10 و20. وتتلقى كلتا المكتبتين مصفوفة توجيه متطابقة من بذرة مشتركة.

تم الإبلاغ عن ثلاث نتائج على صفحة جيثب الخاصة بهم. تعمل النسخة الصفرية على جعل الاتصال الأولي أسرع من خلال إزالة المخزن المؤقت للاتصالات → نسخة المخزن المؤقت للمستخدم التي تهيمن على الخاتمة، لذا فإن وقت الاتصال الخاص بـ MoonEP يقع باستمرار تحت DeepEP v2 في كل مستوى من اختلال التوازن. التوازن المثالي يجعل MoonEP محصنًا تقريبًا ضد الانحراف: يظل وقت الاتصال الخاص به ثابتًا تقريبًا مع نمو maxvio، في حين أن DeepEP v2 – الذي يتم ضبط زمن الوصول إليه بواسطة الترتيب الأكثر سخونة – يتدهور بشكل مطرد.

الوجبات السريعة الرئيسية

• تدعي MoonEP أن كل رتبة EP تتلقاها بالضبط S × K الرموز، بغض النظر عن انحراف جهاز التوجيه.
• يأتي التوازن من الخبراء الزائدين الذين تم التخطيط لهم عبر الإنترنت على وحدة معالجة الرسومات، ثم يتم جلبهم مسبقًا قبل حساب الخبراء.
• تقوم الأشكال الثابتة بإزالة مزامنة مضيف MoE لكل طبقة وإيقاف تجزئة الذاكرة التي يقوم بها OOMs DeepEP.
• يتطلب التدريب B = E/R فتحات الجلب المسبق؛ الاستدلال يمكن أن تنخفض إلى B = 3–4 مع عدم وجود تكلفة صحة.
• تم إصداره من قبل معهد ماساتشوستس للتكنولوجيا في إطار Kimi K3 Open Day، جنبًا إلى جنب مع FlashKDA وAgentEnv.

مصادر: إعلان MoonshotAI/MoonEP على GitHub و@Kimi_Moonshot