TPI-LLM: Servindo LLMs em Escala de 70 bilhões Eficientemente em Dispositivos de Borda com 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

Resumo

A inferência de modelos grandes está migrando do cloud para a borda devido a preocupações com a privacidade dos dados de interação do usuário. No entanto, os dispositivos de borda frequentemente enfrentam desafios com limitações de potência computacional, memória e largura de banda, exigindo colaboração entre vários dispositivos para executar e acelerar a inferência de modelos grandes. O paralelismo de pipeline, a solução convencional, é ineficiente para cenários de usuário único, enquanto o paralelismo de tensor enfrenta dificuldades com comunicações frequentes. Neste artigo, argumentamos que o paralelismo de tensor pode ser mais eficaz do que o pipeline em dispositivos com recursos limitados, e apresentamos um sistema de inferência de paralelismo de tensor eficiente em computação e memória, chamado TPI-LLM, para atender modelos em escala de 70 bilhões. O TPI-LLM mantém dados brutos sensíveis localmente nos dispositivos dos usuários e introduz um agendador de memória de janela deslizante para gerenciar dinamicamente os pesos das camadas durante a inferência, com sobreposição da latência de E/S de disco com a computação e comunicação. Isso permite que modelos maiores sejam executados suavemente em dispositivos com limitações de memória. Analisamos o gargalo de comunicação e descobrimos que a latência de link, não a largura de banda, surge como o principal problema, então um algoritmo de allreduce baseado em estrela é implementado. Através de experimentos extensivos em bancadas de teste emuladas e reais, o TPI-LLM demonstrou mais de 80% menos tempo para o primeiro token e latência de token em comparação com o Accelerate, e mais de 90% em comparação com Transformers e Galaxy, enquanto reduz a pegada de memória de pico do Llama 2-70B em 90%, exigindo apenas 3,1 GB de memória para modelos em escala de 70 bilhões.

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.