Ось що ми працювали останні 9 місяців: зробити так, щоб навчання MoEs працювало ~2 рази швидше і ~2 рази менше пам'яті! Основні моменти: - MoE зазвичай займає найбільше часу та пам'яті в сучасних моделях. Виявляється, можна математично переписати зворотний пропуск MoE, щоб зменшити активаційний мем, який потрібно зберігати у передньому привіді, на ~2x, що дає ті ж градієнти без додаткового перерахунку matmul. Мені дуже подобається цей результат, адже він поєднує як алгоритмічну, так і системну інсайти. - Аналіз вузьких місць у шарі MoE веде до природної стратегії оптимізації: максимально зменшити читання/запис у мемі! Збір вхідних даних для переднього приводу та вихідного градута для BWD іноді може займати стільки ж часу, скільки згруповані GEMM. Ми збираємо злиття з групованим GEMM + overlap mem доступом і обчислюємо, щоб увесь шар працював ~2 рази швидше. - Обчислення top-k для експертної маршрутизації може займати дивовижно багато часу — ~15-20% усього шару MoE! Стандартний Top-K impl використовує Radix Top-K алгоритм, що чудово підходить для великого k, але не оптимальний для малого k. Ми переписували топ-к за допомогою бітонічного топ-к алгоритму, і іноді він у 20-30 разів швидший за топ-к у pytorch! Всі основні ядра написані на Cute-DSL, тому їх має бути легко розширити (і встановити :D). Зерна Hopper вийшли, зерна Blackwell майже готові. Раніше моделі MoE були вдвічі менш ефективними для навчання, сподіваюся, Sonic-MOE це змінить.