Как Nvidia справилась с одной из главных проблем ИИ 😮‍💨

Слышали про KV-кэш? Это одна из самых головоломных проблем при масштабировании контекста в трансформерах, и, да, она действительно может доставить много неприятностей. Проблема в том, что при увеличении длины последовательности KV-кэш растет пропорционально, и для каждой головы и каждого слоя накапливаются колоссальные объемы данных. Всё это не только отнимает уйму памяти, но и сильно тормозит работу моделей.

Например, для модели LLaMA с 65 миллиардами параметров при 128k токенах этот кэш может занять целых 335 ГБ памяти. А это, согласитесь, совсем не шутки! И, кстати, не только память страдает, но и производительность. В общем, проблема довольно масштабная.

Многие стараются бороться с этим, уменьшая размер кэша по слоям или головам. Но в реальности самый большой потенциал для улучшений — это уменьшение по оси токенов. Ведь далеко не все токены действительно важны для модели. В этом и скрывается главный рычаг.

KVzip: попытка решить проблему

Когда появились идеи типа KVzip, оптимизация выглядела многообещающе. В теории можно было сжать кэш до 4 раз без потерь качества. Но на практике метод оказался слишком медленным. Слишком много операций, слишком много вычислений — и вот тут-то на сцену выходит Nvidia.

Что придумали Nvidia?

Nvidia, конечно, не обошли эту проблему стороной. Они взяли концепцию сжатиия KV-кэша, немного доработали её и… вуаля! Теперь получается та же сжатие, но при этом практически без потерь производительности. Всё, что нужно — это обучить маленькую модель, которая будет предсказывать, насколько важен каждый токен для каждого слоя модели. И, внимание, эта модель абсолютно не требует дорогих вычислений, как в случае с KVzip. Просто с помощью линейной модели или двухслойного MLP предсказывается, какие токены можно «отсечь», не потеряв в качестве.

И вот что самое крутое: эффективность увеличивается в 3–4 раза, а при этом производительность практически не страдает. Для вычислений добавляется всего 0,02% дополнительных операций, что — по сути — ничто на фоне квадратичного внимания. Результат: сжатие работает, деградации на бенчмарках практически нет.

А если вам вдруг стало интересно, то всё это решение уже в опенсорсе. Nvidia делится всем этим с миром, так что мы все можем в дальнейшем использовать этот подход 🍑

Nvidia смогла решить проблему, которая стояла перед всеми, кто работал с масштабными моделями, и сделала это почти бесплатно в плане вычислительных затрат.

Data Science