TriForce: Без потерь ускорение генерации длинных последовательностей с иерархическим спекулятивным декодированием

TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding

April 18, 2024

Авторы: Hanshi Sun, Zhuoming Chen, Xinyu Yang, Yuandong Tian, Beidi Chen

cs.AI

Аннотация

С широким развертыванием крупных языковых моделей (LLM) для генерации длинного контента в последнее время возникла растущая потребность в эффективной поддержке вывода длинных последовательностей. Однако ключевое-значение (KV) кэш, который хранится для избежания повторного вычисления, стал критическим узким местом, поскольку его размер линейно увеличивается с длиной последовательности. Из-за авторегрессивной природы LLM весь кэш KV будет загружен для каждого сгенерированного токена, что приводит к низкому использованию вычислительных ядер и высокой задержке. Хотя были предложены различные методы сжатия для кэша KV с целью смягчения этой проблемы, они страдают от ухудшения качества генерации. Мы представляем TriForce, иерархическую систему спекулятивного декодирования, которая масштабируется для генерации длинных последовательностей. Этот подход использует исходные веса модели и динамический разреженный кэш KV через извлечение в качестве модели-черновика, которая служит промежуточным уровнем в иерархии и дополнительно спекулируется более маленькой моделью для сокращения времени ее составления. TriForce не только обеспечивает впечатляющее ускорение для Llama2-7B-128K, достигая до 2.31 раза на GPU A100, но также демонстрирует масштабируемость в обработке даже более длинных контекстов. Для сценария выгрузки на двух GPU RTX 4090 TriForce достигает 0.108 с/токен — всего вдвое медленнее, чем авторегрессивная базовая модель на A100, которая достигает 7.78 раз на нашей оптимизированной системе выгрузки. Кроме того, TriForce работает в 4.86 раза быстрее, чем DeepSpeed-Zero-Inference на одном GPU RTX 4090. Устойчивость TriForce подчеркивается его постоянно выдающимися показателями производительности при различных температурах. Код доступен на https://github.com/Infini-AI-Lab/TriForce.

English

With large language models (LLMs) widely deployed in long content generation recently, there has emerged an increasing demand for efficient long-sequence inference support. However, key-value (KV) cache, which is stored to avoid re-computation, has emerged as a critical bottleneck by growing linearly in size with the sequence length. Due to the auto-regressive nature of LLMs, the entire KV cache will be loaded for every generated token, resulting in low utilization of computational cores and high latency. While various compression methods for KV cache have been proposed to alleviate this issue, they suffer from degradation in generation quality. We introduce TriForce, a hierarchical speculative decoding system that is scalable to long sequence generation. This approach leverages the original model weights and dynamic sparse KV cache via retrieval as a draft model, which serves as an intermediate layer in the hierarchy and is further speculated by a smaller model to reduce its drafting latency. TriForce not only facilitates impressive speedups for Llama2-7B-128K, achieving up to 2.31times on an A100 GPU but also showcases scalability in handling even longer contexts. For the offloading setting on two RTX 4090 GPUs, TriForce achieves 0.108s/tokenx2014only half as slow as the auto-regressive baseline on an A100, which attains 7.78times on our optimized offloading system. Additionally, TriForce performs 4.86times than DeepSpeed-Zero-Inference on a single RTX 4090 GPU. TriForce's robustness is highlighted by its consistently outstanding performance across various temperatures. The code is available at https://github.com/Infini-AI-Lab/TriForce.