Иногда ошибки в ML-системах выглядят как магия. Модель работает нормально, выдаёт адекватные ответы, а потом вдруг начинает нести полную чушь – и непонятно, что пошло не так. Именно с таким случаем столкнулись инженеры AI21 Labs, когда тестировали свою модель Jamba в популярном фреймворке vLLM.
Что произошло
Команда AI21 Labs обнаружила странное поведение при работе с vLLM – это такой фреймворк для быстрого инференса языковых моделей, который сейчас активно используется в продакшене. Модель Jamba, объединяющая в себе архитектуры Transformer и Mamba, в какой-то момент начинала выдавать полную бессмыслицу вместо нормального текста.
Проблема проявлялась не всегда, а только в определённых ситуациях. Это самый неприятный вид ошибок – когда нельзя просто воспроизвести сбой по щелчку пальцев.
Начало расследования
Первым делом команда проверила очевидные вещи: правильно ли загружены веса модели, корректно ли работает токенизатор, нет ли проблем с памятью. Всё было в порядке. Модель на других фреймворках работала нормально, а на vLLM – нет.
Инженеры начали сужать область поиска. Они запускали модель с разными параметрами, меняли длину контекста, пробовали разные промпты. И постепенно начала вырисовываться закономерность: проблема появлялась, когда модель обрабатывала определённые последовательности токенов.
Что такое Mamba и почему это важно
Чтобы понять суть ошибки, нужно немного разобраться в архитектуре. Jamba использует не только классические блоки Transformer, но и слои Mamba – это более новая архитектура, которая работает как рекуррентная сеть, но обучается эффективнее.
В Mamba есть внутреннее состояние (state), которое обновляется по мере обработки токенов. Это состояние хранит информацию о том, что модель «видела» ранее в тексте. И вот именно с этим состоянием что-то пошло не так.
Момент озарения 💡
Прорыв случился, когда команда начала детально логировать, что происходит с внутренним состоянием Mamba на каждом шаге генерации. Они обнаружили, что в определённый момент состояние вдруг «портилось» – значения становились некорректными, и дальше модель уже не могла нормально работать.
Проблема была в том, как vLLM обрабатывает батчи запросов. Когда несколько запросов обрабатываются параллельно, фреймворк использует разные оптимизации для экономии памяти и ускорения работы. Одна из таких оптимизаций – переиспользование вычислений для одинаковых префиксов в разных запросах.
И вот здесь крылся подвох. В какой-то момент при обработке батча происходила ситуация, когда состояние Mamba от одного запроса «перетекало» в другой запрос. По сути, один токен из чужого контекста попадал не туда, куда нужно, и это ломало всё дальнейшее предсказание.
Технические детали
Если говорить конкретнее, проблема была в механизме кеширования префиксов (prefix caching). vLLM старается не пересчитывать одни и те же токены заново, если они уже встречались. Для моделей Transformer это работает отлично, потому что внимание (attention) – это операция без сохранения состояния (stateless).
Но Mamba – это другая история. У неё есть скрытое состояние, которое зависит от всех предыдущих токенов. И если вы берёте закешированное состояние из одного контекста и применяете его в другом, где последовательность токенов хоть немного отличается, всё ломается.
Представьте, что вы читаете книгу и запоминаете сюжет. Потом кто-то заменяет одну страницу на другую из похожей книги. Вы продолжаете читать дальше, но контекст уже сломан – вы помните что-то, что на самом деле не происходило в этой истории.
Решение
Когда причина стала понятна, исправление оказалось относительно простым. Нужно было скорректировать логику кеширования состояний для слоёв Mamba – убедиться, что состояние никогда не переиспользуется между разными запросами некорректно.
Команда AI21 Labs внесла патч в vLLM, который учитывает особенности гибридных архитектур типа Jamba. Теперь фреймворк правильно отслеживает, когда можно безопасно переиспользовать вычисления, а когда нужно пересчитать состояние с нуля.
Что это значит для индустрии
Эта история показывает несколько важных вещей. Во-первых, новые архитектуры моделей требуют новых подходов к инференсу. То, что хорошо работает для классических Transformer, может сломаться для гибридных моделей или архитектур типа Mamba.
Во-вторых, оптимизации – это всегда компромисс. Кеширование префиксов (prefix caching) даёт огромный прирост производительности для обычных моделей, но для архитектур, сохраняющих состояние (stateful), нужны дополнительные проверки.
В-третьих, отладка ML-систем – это отдельное искусство. Проблема была не в модели и не в данных, а в тонком взаимодействии между архитектурой модели и инфраструктурой инференса. Такие ошибки сложно поймать автоматическими тестами, потому что они проявляются только в специфических сценариях.
Уроки для разработчиков
Если вы работаете с языковыми моделями в продакшене, из этой истории можно вынести несколько практических выводов:
- Всегда тестируйте модели на реальных сценариях использования, а не только на синтетических бенчмарках.
- Обращайте внимание на граничные случаи – когда модель работает в батче с другими запросами, при разной длине контекста, с разными параметрами генерации.
- Логирование промежуточных состояний может быть дорогим в продакшене, но незаменимым при отладке.
- Новые архитектуры могут требовать модификаций инфраструктуры – не стоит предполагать, что всё будет работать «из коробки».
Хорошая новость в том, что vLLM – это проект с открытым исходным кодом (open-source) с активным сообществом. Ошибка была найдена, исправлена, и теперь другие разработчики, которые работают с гибридными моделями, не столкнутся с той же проблемой.
Проще говоря, один неправильный токен действительно может испортить всё – но только если система не учитывает особенности архитектуры. Теперь учитывает.