я начинаю хайповать по МоЕ и вам тоже стоит, и ВОТ ПОЧЕМУ

те кто давно на меня подписаны знают что я довольно скептически отношусь к СВЕРХ массивным моделям - да круто, мы обязательно туда придем, но это игрушки для безмерно богатых ребят, на такого размера графах сложно строить продукты - как правило они не влезают в одну ноду(для TPU это не так критично - там сеть заметно быстрее), в итоге на середину 2023 года мы имели примерно такую картинку:

*180b влезает в 8а100 к слову, но на тюн нужно несколько нод по хорошему

ну те вы конечно можете пихнуть 180b модель в прод, но ипотека сама себя не выплатит, а вот этот графичек не наебать:

отсюда https://openai.com/research/instruction-following

Поэтому до выхода хороших откртых моделей увидеть в проде что то большее чем 20b GPT NEO было в диковинку, когда вышли LLAMA13b(по метрикам +- как gpt3 175b) то стало понятно что Scaling сurve не до конца точная - на бОльшем числе бОлее чистых токенов модели становяться умнее, а потом вышла MIstral 7b - которая порвала(~~жопы~~) и показала что в 7b размере можно показывать очень мощные результаты, при этом у mistral супер дешевый инференс - 500tokens/s на 1a100, а значит прод можно скейлить во много раз не закладывая почку БЕЗОСОСУ на gpuшки.

МоЕ - mixture of experts, та самая архитектура из слива про OpenAI, 2т параметров, короче вы знаете.

План примерно такой: я расскажу про switch transfomer, OpenMoE, Megablocks(mixtral)

Cамое просто обьяснение - это бабки у подьезда, учатся они все долго, но за инференс отвечают K самые увереные(выбираются либо Softmax либо сonv либо SelfAttn, в зависимости от архитектур)

Switch Transfomer

это не первая MoE архитектура, MoE пытались делать еще на lstm, но было плохо и не кому не нужно, а гугл сделали ...... нормально и по прежнему никому не нужно, ну серьезно

ладно, там есть модели с 1000+ скачиваний, но суть вы поняли, прием был холодный

Короче, модели и сам подход вышли крутые и вот почему:

в 7раз дешевле претрен чем у DENSE t5 версии используя те же флопсы
можно задистилировать в 1% от оригинального размера sparce модели, при этом сохранив 30% ее качества(waht ever it means, кода я не нашел)

Our work here focuses on TPU architectures, but these class of models may be similarly trained on GPU clusters. In our distributed training setup, our sparsely activated layers split unique weights on different devices. Therefore, the weights of the model increase with the number of devices, all while maintaining a manageable memory and computational footprint on each device.

на самом деле большая часть бенефитов работает именно на TPU с быстрыми шинами, те на GPU скорее всего все будет сильно хуже

Про архитектуру: по сути это все тоже самое т5, но вместо обычных MLP поставили MoE который выбирается через ROuter, на каждого эксперта приходиться по

По метрикам модели вышли довольно посредственными, они не были значительно лучше своих DENSE собратьев, при этом были очень ОЧЕНЬ дорогими для инфернса. В общем то комьюнити забили их допиливать.

собственно да, модель конечно учиться быстрее, но 14b с качеством 700m было перебором по меркам 2021-2022 года, все отложили модель на дальнию полку и забыли о ней.

OpenMoE(ST-MOE)

Прибежали челы из университета сингапура(неожиданно), взяли малоизвестную работу ребят из гугла, прикрутили Z loss отвечающий за выбор эксперта, прикрутиили bf16 и несколько оптимизаций GELU которые мне лениво разбирать, для нас важно другое.

ну те выигрышь был на 10тки процентов, что уже как бы имело смысл))

я конечно все понимаю, но эти цифры меня очень сильно пугают, 1.5t параметров

Собственно авторы обучили на 750B токенах Decoder only модель, которая была при этом ощутимо лучше на бенчах, правда авторы забили сравниваться с LLAMA(иронично), а учили на TPU cloud(см тейк выше про эффективность МоЕ на tpu)

FLAN MOE(2023)

Еще существует FLAN-MoE, он даже получается эффективнее чем просто FLAN, архитектурно совпдает с ST Moe выше

И получается слегка экономичнее Dense версии

MEGABLOCKS: EFFICIENT SPARSE TRAINING WITH MIXTURE-OF-EXPERTS

если вы внимательно читали буковки выше, могли заметить - ключевая эффективность МоЕ провляется на TPU(спецефичный сетап) и на претрене, на инференсе есть разница, но она не столь велика, а с ростом размера модели - уменьшается.

модель орентированная на GPU, сделаны kernels с которыми MoE перестает овощить
ОЧЕНЬ сильная заточенность на удешевление инфернса модели(трен постольку поскольку)

ну начнем с того что теперь у нас эксперты выбираются одновременно, по аналогии с головами трансформера

Переформулировали как mm прокидывание токенов через экспертов, что опять же сильно упрощает жизнь GPUшкам, а еще nvidia kernels быстрее работают со sparse сompute)))

на удивление MoE в 1.38x раз лучше чем ванильный трансформер, авторы обучили только игрушечные модели как proof of concept архитектуры, но по каким то причинам Mistral Ai выбрали именно эту версию, давайте накинем причины почему

инференс и трен требует больше памяти чем dense версия - да, но при этом модели используют меньше флопс на forward - за счет того что используется только 2 эксперта из 8
возможно модели учатся быстрее - я не учил, я хз))))) но авторы обещают что все хорошо будет
ПРИ ЭТОМ инференс MOE крайне выгоден на устройствах с Unified memory - телефоны, M серия чипов, много всякого от AMD. Почему? 40b модель не влезает целиком не в одну из текущих распространенных карт, но в случае с MoE это не нужно, вам нужно использовать только два эксперта, те скорость инференса для МоЕ будет такой же как у 14b модели - а это десятки токенов в секунду

ПАпиры и репы:

https://arxiv.org/pdf/2211.15841 megablock
https://github.com/stanford-futuredata/megablocks megablock code
https://arxiv.org/pdf/2305.14705 flan moe
https://www.notion.so/Aug-2023-OpenMoE-v0-2-Release-43808efc0f5845caa788f2db52021879 open moe
https://arxiv.org/pdf/2202.08906 st moe
https://arxiv.org/abs/2101.03961 switch