Assisted Generation: оригинальный способ ускорить генерацию

Прикольный пост на Huggingface про ускорение генерации больших LMок. Реализация уже в 🤗, можно использовать 👍

One sentence summary: маленькая модель генерирует последовательности токенов, большая проверяет

🔹 Как происходит генерация?
- последовательно, токен за токеном
- 1 forward pass на 1 токен
- предыдущие внутренние состояния кешируются
- вычисления скрытого состояния для нового токена идут послойно

🔹 Основной ботлнек и что с этим делать?
Для больших языковых моделей ботлнек это пропускная способность памяти GPU во время матричного умножения. То есть, основное время уходит не на вычисления, а на трансфер весов из памяти к вычислительным ядрам 🤔

Отсюда получается
- INT8 квантизация, Flash Attention и другие техники, учитывающие факт выше, помогают заметно снизить латенси
- Увеличение размера батча увеличивает пропускную способность модели, так как один раз загруженные в вычислительные ядра веса, используются для нескольких примеров (латенси при этом немного возрастает)
- Tensor Parallelism (если есть > 1 GPU) параллелит трансфер весов (и вычисления), что позволяет уменьшить латенси

🔹 Их метод (в посте есть понятная визуализация)
1. Генерируем маленькой assistant моделью несколько последующих токенов (жадно)
2. Большой моделью делаем один forward и получаем логиты для этих сгенерированных токенов
3. Берем от логитов argmax и отбрасываем предложенные assistant моделью токены, начиная с первого несовпадения
4. goto 1.

Для семплинга тоже придумали, но быстрее работает с жадной генерацией

Интересный подход. Получатся разные латенси для разных текстов: легкие (ожидаемые) будут генерироваться быстрее. Типа роутинга запросов, только внутри одного метода генерации 🤔