Новый проект
kyegomez/Latent-MoE
Однофайловый слой PyTorch, который переносит экспертов MoE в латентное пространство, снижая коммуникационные и параметрические затраты.
- Stars
- ★ 18 ↗ +1 за 7 дней
- Forks
- ⑂ 0
- Язык
- Python
- Статус
- В канале 1 янв. 1 г.
AI-разбор
Проект реализует LatentMoE из работы Elango и других авторов NVIDIA для построения более эффективных слоёв Mixture of Experts в PyTorch. Каждый токен сначала проецируется из размерности скрытого состояния d в уменьшенную латентную размерность l = d / alpha, после чего маршрутизируемые эксперты работают в этом пространстве, а результат возвращается в d. Такой подход снижает объём all-to-all-коммуникаций и памяти для загрузки весов экспертов в alpha раз, что позволяет увеличить общее число экспертов. Вариант acc увеличивает top-k для сохранения стоимости и потенциального повышения точности, а вариант eff сохраняет top-k и ориентирован на снижение вычислительных затрат. В репозитории есть слой LatentMoE, конфигурация, оценка асимптотических затрат, пример MoETransformer с GQA, генерацией и обучением на потоковом корпусе через HuggingFace datasets; для обучения требуются отдельные зависимости, а Wikimedia Wikipedia может загружать целые parquet-фрагменты даже в потоковом режиме.
Исследователи и инженеры, работающие с PyTorch, Transformer и Mixture of Experts, которым нужен компактный экспериментальный слой и готовый пример обучения.
Что внутри
- Однофайловый слой LatentMoE, совместимый по форме входа и выхода со стандартным MoE FFN.
- Сжатие токенов в латентную размерность l = d / alpha перед вычислением маршрутизируемых экспертов.
- Два режима работы: acc для издержек, сопоставимых с базовой моделью, и eff для более дешёвого вывода.
- Поддержка общих постоянно активных экспертов, маршрутизации top-k и масштабирования числа экспертов.
- Метод cost_summary() для просмотра асимптотических показателей из таблицы 1 исходной работы.
- Пример MoETransformer с GQA, вычислением вспомогательной потери, генерацией, обратным распространением и скриптом обучения с контрольными точками.
Где применять
- Замена стандартного MoE FFN в исследовательских Transformer-моделях на PyTorch.
- Эксперименты с компромиссом между точностью, количеством экспертов, FLOP и памятью.
- Обучение небольших MoE-трансформеров на Wikipedia или более лёгком корпусе Wikitext.
- Проверка влияния латентной маршрутизации на коммуникационные затраты и загрузку весов экспертов.
- Прототипирование архитектур MoE с общими экспертами и настраиваемым top-k.
Темы