PockEngine: Ajuste Fino Esparso e Eficiente em um Ambiente Compacto

PockEngine: Sparse and Efficient Fine-tuning in a Pocket

October 26, 2023

Autores: Ligeng Zhu, Lanxiang Hu, Ji Lin, Wei-Chen Wang, Wei-Ming Chen, Chuang Gan, Song Han

cs.AI

Resumo

O aprendizado no dispositivo e o ajuste fino eficiente permitem a personalização contínua e preservadora de privacidade (por exemplo, ajustar localmente grandes modelos de linguagem em dados personalizados). No entanto, os frameworks de treinamento existentes são projetados para servidores em nuvem com aceleradores poderosos (por exemplo, GPUs, TPUs) e carecem de otimizações para o aprendizado na borda, que enfrenta desafios de limitações de recursos e diversidade de hardware de borda. Apresentamos o PockEngine: um motor pequeno, esparso e eficiente para permitir o ajuste fino em diversos dispositivos de borda. O PockEngine suporta retropropagação esparsa: ele poda o grafo de retropropagação e atualiza esparsamente o modelo com economia de memória e redução de latência medidas, mantendo a qualidade do modelo. Em segundo lugar, o PockEngine é orientado para compilação: todo o grafo de treinamento (incluindo passos de propagação direta, retropropagação e otimização) é derivado em tempo de compilação, o que reduz a sobrecarga de tempo de execução e traz oportunidades para transformações de grafo. O PockEngine também integra um conjunto rico de otimizações de grafo de treinamento, podendo assim acelerar ainda mais o custo de treinamento, incluindo reordenação de operadores e troca de backends. O PockEngine suporta diversas aplicações, frontends e backends de hardware: ele compila e ajusta modelos definidos em PyTorch/TensorFlow/Jax de forma flexível e implanta binários em CPUs/GPUs/DSPs móveis. Avaliamos o PockEngine tanto em modelos de visão quanto em grandes modelos de linguagem. O PockEngine alcança até 15 vezes mais velocidade em comparação com o TensorFlow padrão (Raspberry Pi) e 5,6 vezes mais economia de memória na retropropagação (Jetson AGX Orin). Notavelmente, o PockEngine permite ajustar o LLaMav2-7B no NVIDIA Jetson AGX Orin a 550 tokens/s, 7,9 vezes mais rápido que o PyTorch.

English

On-device learning and efficient fine-tuning enable continuous and privacy-preserving customization (e.g., locally fine-tuning large language models on personalized data). However, existing training frameworks are designed for cloud servers with powerful accelerators (e.g., GPUs, TPUs) and lack the optimizations for learning on the edge, which faces challenges of resource limitations and edge hardware diversity. We introduce PockEngine: a tiny, sparse and efficient engine to enable fine-tuning on various edge devices. PockEngine supports sparse backpropagation: it prunes the backward graph and sparsely updates the model with measured memory saving and latency reduction while maintaining the model quality. Secondly, PockEngine is compilation first: the entire training graph (including forward, backward and optimization steps) is derived at compile-time, which reduces the runtime overhead and brings opportunities for graph transformations. PockEngine also integrates a rich set of training graph optimizations, thus can further accelerate the training cost, including operator reordering and backend switching. PockEngine supports diverse applications, frontends and hardware backends: it flexibly compiles and tunes models defined in PyTorch/TensorFlow/Jax and deploys binaries to mobile CPU/GPU/DSPs. We evaluated PockEngine on both vision models and large language models. PockEngine achieves up to 15 times speedup over off-the-shelf TensorFlow (Raspberry Pi), 5.6 times memory saving back-propagation (Jetson AGX Orin). Remarkably, PockEngine enables fine-tuning LLaMav2-7B on NVIDIA Jetson AGX Orin at 550 tokens/s, 7.9times faster than the PyTorch.