gonzo-обзоры ML статей 2720

This media is not supported in your browser

5.9K views22:50

❤8😁4

You Only Cache Once: Decoder-Decoder Architectures for Language Models
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, Furu Wei
Статья: https://arxiv.org/abs/2405.05254
Код: https://github.com/microsoft/unilm/tree/master/YOCO

Архитектурные новости. Авторы придумали архитектуру для LLM под названием decoder-decoder.

Напомним, что оригинальный трансформер (и например модели типа T5) был построен на полной архитектуре encoder-decoder, большая часть современных LLM (типа GPT) используют только decoder, и другая популярная ветка недавнего прошлого (модели семейства BERT) состоит только из encoder. Энкодер всегда был двунаправленным (bidirectional) и модели с таким двунаправленным компонентом (то есть encoder и encoder-decoder) имели проблемы с авторегрессионной генерацией — там для генерации нового токена сначала надо было заэнкодить всю последовательность из входа и уже нагенерённой части выхода. Можно конечно использовать только декодерную часть для генерации, но тогда сгенерённые токены не используют на полную мощь параметры энкодера. У decoder тут всё неплохо, при авторегрессионной генерации можно закешировать вектора KV (key и value в блоках внимания) и переиспользовать для генерации нового токена, не надо заново кодировать всю историю.

Но как говорится в сказании о Савитри, “есть один недостаток”. KV-кэш очень пухнет при росте длины генерируемой последовательности, он отжирает кучу памяти GPU и LLM-ки становятся memory-bound. Так для 65B модели (с grouped-query attention и квантизацией KV в 8 бит) для 512k токенов нужно 86Gb памяти, что перекрывает объём памяти H100-80GB. К тому же фаза prefill (см тут или хороший обзор тут), в которой надо обработать все входные токены промпта и вычислить для них значения KV, может занимать сотни секунд для очень длинных входов типа 1М (здесь, кстати, интересно, что Гугл с Gemini 1.5 придумал).

Весь трансформер из L слоёв разделяется поровну и первые L/2 слоёв реализуют self-decoder через efficient self-attention. Размер KV-кеша этой части константен, то есть O(1). Выход последнего слоя self-decoder даёт глобальный KV-кеш, куда ходит вторая половина, cross-decoder, реализованная через оставшиеся L/2 слоёв. Каждый блок получает на вход Q и через cross-attention идёт в этот глобальный KV-кеш. Здесь уже везде стандартное (почти, с GQA, https://arxiv.org/abs/2305.13245) multi-head attention с полным окном.

Под efficient self-attention в self-decoder авторы подразумевают sliding-window attention как в старом добром sparse transformer имени Ильи Суцкевера и ко (https://www.tgoop.com/gonzo_ML/65). Как вариант, вместо него в self-decoder может использоваться RetNet (https://www.tgoop.com/gonzo_ML/1753) под названием gRet (aka gRetNet или RetNet-3) с data-dependent гейтингом. Вроде бы такой же мы и разбирали когда-то давно в оригинальной статье.

В остальном блоки в этих слоях в целом стандартные, чередование внимания и FFN, с использованием pre-RMSNorm, SwiGLU, GQA.

Полученная архитектура называется YOCO (You Only Cache Once, так понимаю тут речь про кеширование в L/2 слое). Это всё похоже на encoder-decoder, но снаружи выглядит как декодер и обе части используют causal masking.

YOCO эффективнее обычного трансформера за счёт меньших требований к памяти, кеш для длинных последовательностей скейлится как O(N) вместо O(NL), то есть можно делать больше инференса и/или с более крупными батчами (что повышает throughput).

Ещё из интересных свойств YOCO есть то, что во время стадии prefill можно сделать early exit и не ходить в cross-decoder, это повышает скорость данной фазы. Поскольку в self-decoder находится половина слоёв, то это уже сокращение вычислений и времени в два раза. К тому же эффективная реализация внимания в self-decoder обычно быстра. Они приводят пример запроса с размером контекста в 512K, на котором prefill latency падает со 180 секунд (трансформер с flash-decoding и kernel fusion) до менее 6 секунд. И даже на длине 32K YOCO всё равно в три раза быстрее (на этой фазе, а не в целом end-to-end).

🔥23❤5👍3👀2

5.96K viewsedited 23:50