Вот что мы делали последние 9 месяцев: сделать обучение MoE примерно в 2 раза быстрее и в 2 раза менее ресурсоемким! Основные моменты: - MoE обычно занимает больше всего времени и памяти в современных моделях. Оказалось, что можно математически переписать обратный проход MoE, чтобы уменьшить объем памяти активации, который нужно хранить в прямом проходе, примерно в 2 раза, что приводит к тем же градиентам без дополнительного пересчета матричных произведений. Мне действительно нравится этот результат, так как он сочетает в себе как алгоритмические, так и системные идеи. - Анализ узких мест в слое MoE приводит к естественной стратегии оптимизации: уменьшить чтение/запись памяти насколько это возможно! Сбор входных данных для прямого прохода и выходных градиентов для обратного прохода иногда занимает столько же времени, сколько и сгруппированные GEMM. Мы объединяем сбор данных с сгруппированным GEMM + перекрываем доступ к памяти и вычисления, чтобы сделать весь слой примерно в 2 раза быстрее. - Вычисление top-k для маршрутизации экспертов может занимать удивительно много времени, около 15-20% всего слоя MoE! Стандартная реализация top-k использует алгоритм radix top-k, который отлично подходит для больших k, но не оптимален для маленьких k. Мы переписали top-k, используя алгоритм bitonic top-k, и он иногда в 20-30 раз быстрее, чем top-k в pytorch! Все основные ядра написаны на Cute-DSL, так что их должно быть легко расширять (и устанавливать :D). Ядра Hopper уже готовы, ядра Blackwell почти готовы. Модели MoE раньше были в 2 раза менее эффективны с точки зрения аппаратного обеспечения для обучения, надеюсь, Sonic-MOE изменит это.

Nvidia продолжает выпускать одни из самых сильных и быстрых открытых моделей. Данные для предварительного и последующего обучения также публикуются, что очень немногие организации сделали.