Оп оп, что такое MoE в нейронках
Оп оп, что такое MoE в нейронках
Mixture-of-Experts - G.O.A.T. архитектура.
Начнем, пожалуй, самую интересную тему в NLP, по моему мнению - Mixture-of-Experts(MoE).
Наверняка вы слышали что-то, мол, это круче, это эффективнее и тому подобное. Но что значит круче и эффективнее? Сегодня это и разберем. У этого поста будет 2 части: 1 - введение в MoE и 2 - Как учится MoE. Начнем.
1)Что такое MoE, и почему оно не dense?
Все NLP инженеры знают, что базовый декодер - dense модель. Это значит, что у нас в FFN просто огромное количество нейронов. За счет этого как раз наши модели понимают этот мир и делают некие обобщения, усваивая информацию. По сути, FFN отвечает у нас за хранения знаний у LLM. Но есть проблема - память. Точнее, ее катастрофическое малое количество. Забегая вперед, почти все модели учатся через n-d параллелизмы: Tensor Parallelism, Pipeline Parallelism и/или их комбинации. Это нужно для того, чтобы просто уместить модели больше чем условные 100B+ параметров. А обучать... еще больнее... Если мы берем обучение Qwen в BF16(A100/H100, например), то у нас получается: веса модели в FP32 + веса в BF16(AMP) + градиенты в BF16(AMP) + 2 состояния оптимизатора в FP32(AdamW хранит экспоненциально затухающее среднее градиентов и их квадратов). Уже будто 1 карточки на 80GB не хватает... С инференсом, конечно, проще, но есть другая проблема - KV-cache. Эта штука тоже очень много ест видеопамяти. И хоть для него придумали Multi-head Latent Attention, все равно, памяти мало.
И тут появляется MoE… и не решает проблему памяти… Тогда зачем нам MoE?
2)Архитектура MoE.
Вот тут и самое интересное. Вместо того, чтобы инфереть все веса, мы инферем часть экспертов. Эксперт - это маленькая FFN. И таких экспертов много(например в DeepSeek V3 их 257 - 1 shared и 256 routed). И теперь вместо того, чтобы прогонять на инференсе весь FFN в dense модели, мы берем 2-8 экспертов(маленьких FFN) и прогоняем их. В итоге вместо того, чтобы, в случае DeepSeek V3, гонять все 671B параметров, мы используем ~37B. То есть 94% весов "спят". И это намного круче, ведь по сути у нас FLOPS 37B модели, а знания 671B. Погрузимся в то, как работает MoE.
3)Внутри MoE.
Наш токен идет по нашей LLM. Первые несколько слоев(как у DeepSeekMoE) - dense. Но почему? Мы же выяснили, что MoE - круто. Да, MoE действительно крутая штука, но есть у нее проблемы - обобщение информации. У нас все же "набор экспертов", каждый из которых отвечает на свои вопросы. Грубо говоря, один шарит за матан, другой за биологию, а третий - за мемы. Тут на сцену выходит Shared Expert (общий эксперт), который как раз фиксит эту проблему. Он забирает на себя общие знания о языке и всегда обрабатывает токен в фоновом режиме, вне зависимости от выбора роутера. А вот за узкими знаниями наш токен идет уже в роутер(a.k.a Router, Gating или Gating Network). Он выбирает какие эксперты будут работать с токеном. По сути, это Top-k(softmax(nn.Linear(d_token, num_experts))). После чего наш токен проходит эти k экспертов. Но тут проблема... У нас получается k разных векторов? Да. Чтобы сделать final output, мы делаем взвешенную сумму выходов экспертов на "веса" после softmax нашего роутера. Ну и так далее. Да, это все!
4)Беды MoE.
Ну я думаю очевидно было, что все так гладко не будет. MoE хоть и очень прикольная архитектура, но очень капризная. В классическом MoE (как в ранних Switch Transformers) была проблема: роутер - ленивая сущность. Он быстро понимает, что пара экспертов чуть лучше справляется с задачами, и начинает заваливать их всеми токенами. В итоге эти "звездные" эксперты переобучаются, а остальные 90% модели остаются "пустыми" и бесполезными.
Mixture-of-Experts - G.O.A.T. архитектура.
Начнем, пожалуй, самую интересную тему в NLP, по моему мнению - Mixture-of-Experts(MoE).
Наверняка вы слышали что-то, мол, это круче, это эффективнее и тому подобное. Но что значит круче и эффективнее? Сегодня это и разберем. У этого поста будет 2 части: 1 - введение в MoE и 2 - Как учится MoE. Начнем.
1)Что такое MoE, и почему оно не dense?
Все NLP инженеры знают, что базовый декодер - dense модель. Это значит, что у нас в FFN просто огромное количество нейронов. За счет этого как раз наши модели понимают этот мир и делают некие обобщения, усваивая информацию. По сути, FFN отвечает у нас за хранения знаний у LLM. Но есть проблема - память. Точнее, ее катастрофическое малое количество. Забегая вперед, почти все модели учатся через n-d параллелизмы: Tensor Parallelism, Pipeline Parallelism и/или их комбинации. Это нужно для того, чтобы просто уместить модели больше чем условные 100B+ параметров. А обучать... еще больнее... Если мы берем обучение Qwen в BF16(A100/H100, например), то у нас получается: веса модели в FP32 + веса в BF16(AMP) + градиенты в BF16(AMP) + 2 состояния оптимизатора в FP32(AdamW хранит экспоненциально затухающее среднее градиентов и их квадратов). Уже будто 1 карточки на 80GB не хватает... С инференсом, конечно, проще, но есть другая проблема - KV-cache. Эта штука тоже очень много ест видеопамяти. И хоть для него придумали Multi-head Latent Attention, все равно, памяти мало.
И тут появляется MoE… и не решает проблему памяти… Тогда зачем нам MoE?
2)Архитектура MoE.
Вот тут и самое интересное. Вместо того, чтобы инфереть все веса, мы инферем часть экспертов. Эксперт - это маленькая FFN. И таких экспертов много(например в DeepSeek V3 их 257 - 1 shared и 256 routed). И теперь вместо того, чтобы прогонять на инференсе весь FFN в dense модели, мы берем 2-8 экспертов(маленьких FFN) и прогоняем их. В итоге вместо того, чтобы, в случае DeepSeek V3, гонять все 671B параметров, мы используем ~37B. То есть 94% весов "спят". И это намного круче, ведь по сути у нас FLOPS 37B модели, а знания 671B. Погрузимся в то, как работает MoE.
3)Внутри MoE.
Наш токен идет по нашей LLM. Первые несколько слоев(как у DeepSeekMoE) - dense. Но почему? Мы же выяснили, что MoE - круто. Да, MoE действительно крутая штука, но есть у нее проблемы - обобщение информации. У нас все же "набор экспертов", каждый из которых отвечает на свои вопросы. Грубо говоря, один шарит за матан, другой за биологию, а третий - за мемы. Тут на сцену выходит Shared Expert (общий эксперт), который как раз фиксит эту проблему. Он забирает на себя общие знания о языке и всегда обрабатывает токен в фоновом режиме, вне зависимости от выбора роутера. А вот за узкими знаниями наш токен идет уже в роутер(a.k.a Router, Gating или Gating Network). Он выбирает какие эксперты будут работать с токеном. По сути, это Top-k(softmax(nn.Linear(d_token, num_experts))). После чего наш токен проходит эти k экспертов. Но тут проблема... У нас получается k разных векторов? Да. Чтобы сделать final output, мы делаем взвешенную сумму выходов экспертов на "веса" после softmax нашего роутера. Ну и так далее. Да, это все!
4)Беды MoE.
Ну я думаю очевидно было, что все так гладко не будет. MoE хоть и очень прикольная архитектура, но очень капризная. В классическом MoE (как в ранних Switch Transformers) была проблема: роутер - ленивая сущность. Он быстро понимает, что пара экспертов чуть лучше справляется с задачами, и начинает заваливать их всеми токенами. В итоге эти "звездные" эксперты переобучаются, а остальные 90% модели остаются "пустыми" и бесполезными.
Автор: Егор Гречин