TPI-LLM: Sirviendo eficientemente LLMs a escala de 70 mil millones en dispositivos Edge de recursos limitados

TPI-LLM: Serving 70B-scale LLMs Efficiently on Low-resource Edge Devices

October 1, 2024

Autores: Zonghang Li, Wenjiao Feng, Mohsen Guizani, Hongfang Yu

cs.AI

Resumen

La inferencia de modelos grandes está pasando del cloud al edge debido a preocupaciones sobre la privacidad de los datos de interacción del usuario. Sin embargo, los dispositivos edge a menudo tienen problemas con la potencia de cálculo limitada, la memoria y el ancho de banda, lo que requiere colaboración entre varios dispositivos para ejecutar y acelerar la inferencia de modelos grandes. El paralelismo de tuberías, la solución principal, es ineficiente para escenarios de un solo usuario, mientras que el paralelismo de tensores tiene dificultades con las comunicaciones frecuentes. En este documento, argumentamos que el paralelismo de tensores puede ser más efectivo que el de tuberías en dispositivos con recursos limitados, y presentamos un sistema de inferencia paralelo de tensores eficiente en cálculo y memoria, llamado TPI-LLM, para servir modelos a escala de 70 mil millones. TPI-LLM mantiene los datos crudos sensibles locales en los dispositivos de los usuarios e introduce un programador de memoria de ventana deslizante para gestionar dinámicamente los pesos de las capas durante la inferencia, con la superposición de la latencia de E/S de disco con la computación y la comunicación. Esto permite que los modelos más grandes se ejecuten sin problemas en dispositivos con memoria limitada. Analizamos el cuello de botella de comunicación y encontramos que la latencia del enlace, no el ancho de banda, surge como el principal problema, por lo que se implementa un algoritmo de allreduce basado en estrella. A través de experimentos extensos en bancos de pruebas emulados y reales, TPI-LLM demostró más de un 80% menos de tiempo hasta el primer token y latencia de token en comparación con Accelerate, y más del 90% en comparación con Transformers y Galaxy, al tiempo que reduce la huella de memoria máxima de Llama 2-70B en un 90%, requiriendo solo 3.1 GB de memoria para modelos a escala de 70 mil millones.

English

Large model inference is shifting from cloud to edge due to concerns about the privacy of user interaction data. However, edge devices often struggle with limited computing power, memory, and bandwidth, requiring collaboration across multiple devices to run and speed up LLM inference. Pipeline parallelism, the mainstream solution, is inefficient for single-user scenarios, while tensor parallelism struggles with frequent communications. In this paper, we argue that tensor parallelism can be more effective than pipeline on low-resource devices, and present a compute- and memory-efficient tensor parallel inference system, named TPI-LLM, to serve 70B-scale models. TPI-LLM keeps sensitive raw data local in the users' devices and introduces a sliding window memory scheduler to dynamically manage layer weights during inference, with disk I/O latency overlapped with the computation and communication. This allows larger models to run smoothly on memory-limited devices. We analyze the communication bottleneck and find that link latency, not bandwidth, emerges as the main issue, so a star-based allreduce algorithm is implemented. Through extensive experiments on both emulated and real testbeds, TPI-LLM demonstrated over 80% less time-to-first-token and token latency compared to Accelerate, and over 90% compared to Transformers and Galaxy, while cutting the peak memory footprint of Llama 2-70B by 90%, requiring only 3.1 GB of memory for 70B-scale models.