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://xn--r1a.website/gonzo_ML/65). Как вариант, вместо него в self-decoder может использоваться RetNet (https://xn--r1a.website/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).
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://xn--r1a.website/gonzo_ML/65). Как вариант, вместо него в self-decoder может использоваться RetNet (https://xn--r1a.website/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
В тестах за основу взяли StableLM-3B-4E1T и сделали сопоставимую YOCO, она даёт результаты сравнимые с другими хорошо затюненными моделями такого же размера. Лосс от размера модели скейлится также как у Llama-optimized трансформера. При этом YOCO с gRet чуть лучше, чем со sliding-window attention (SWA) и обычный трансформер.
Если расширить контекст YOCO-3B до 1M (привет, Gemini!) через продолжение обучения с length schedule 64K, 256K, 1M, то на Needle In A Haystack всё выглядит почти идеально.
В недрах приложений есть сравнение с Mamba, RetNet, Hybrid H3, gRetNet и трансформером. YOCO с трансформером рулят (по перплексии).
Самые интересные результаты в производительности. По памяти улучшение в разы и чем больше длина последовательности, тем больше улучшение. На длине 1М YOCO ест в 9.38x меньше памяти, чем трансформер с GQA, Flash-Decoding и kernel fusion. В основном за счёт KV кеша, но кажется ещё небольшое улучшение у gRet при хранении активаций. По метрике prefilling latency улучшение в десятки раз. По throughput (токены в секунду) на длинных входах ускорение почти до 10 раз (в основном по двум причинам: более быстрый prefill, а также возможность использовать больший батч из-за лучшей работы с памятью). В сочетаниях типа YOCO + BitNet + Groq может быть кумулятивный эффект и вообще бомба.
Хорошая инженерная работа, мне нравится.
Если расширить контекст YOCO-3B до 1M (привет, Gemini!) через продолжение обучения с length schedule 64K, 256K, 1M, то на Needle In A Haystack всё выглядит почти идеально.
В недрах приложений есть сравнение с Mamba, RetNet, Hybrid H3, gRetNet и трансформером. YOCO с трансформером рулят (по перплексии).
Самые интересные результаты в производительности. По памяти улучшение в разы и чем больше длина последовательности, тем больше улучшение. На длине 1М YOCO ест в 9.38x меньше памяти, чем трансформер с GQA, Flash-Decoding и kernel fusion. В основном за счёт KV кеша, но кажется ещё небольшое улучшение у gRet при хранении активаций. По метрике prefilling latency улучшение в десятки раз. По throughput (токены в секунду) на длинных входах ускорение почти до 10 раз (в основном по двум причинам: более быстрый prefill, а также возможность использовать больший батч из-за лучшей работы с памятью). В сочетаниях типа YOCO + BitNet + Groq может быть кумулятивный эффект и вообще бомба.
Хорошая инженерная работа, мне нравится.
arXiv.org
You Only Cache Once: Decoder-Decoder Architectures for Language Models
We introduce a decoder-decoder architecture, YOCO, for large language models, which only caches key-value pairs once. It consists of two components, i.e., a cross-decoder stacked upon a...
🔥22👍6❤1🆒1