PockEngine: Ottimizzazione Sparse ed Efficiente in uno Spazio Ridotto

PockEngine: Sparse and Efficient Fine-tuning in a Pocket

October 26, 2023

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

cs.AI

Abstract

L'apprendimento on-device e il fine-tuning efficiente consentono una personalizzazione continua e rispettosa della privacy (ad esempio, il fine-tuning locale di grandi modelli linguistici su dati personalizzati). Tuttavia, i framework di addestramento esistenti sono progettati per server cloud con acceleratori potenti (ad esempio, GPU, TPU) e mancano delle ottimizzazioni necessarie per l'apprendimento ai margini della rete, che deve affrontare le sfide delle limitazioni delle risorse e della diversità dell'hardware edge. Introduciamo PockEngine: un motore piccolo, sparso ed efficiente per abilitare il fine-tuning su vari dispositivi edge. PockEngine supporta la retropropagazione sparsa: pota il grafo di retropropagazione e aggiorna il modello in modo sparso, ottenendo un risparmio di memoria misurabile e una riduzione della latenza, mantenendo al contempo la qualità del modello. In secondo luogo, PockEngine è basato sulla compilazione: l'intero grafo di addestramento (inclusi i passi di forward, backward e ottimizzazione) viene derivato in fase di compilazione, riducendo l'overhead a runtime e aprendo opportunità per trasformazioni del grafo. PockEngine integra inoltre un ricco set di ottimizzazioni del grafo di addestramento, che possono ulteriormente accelerare i costi di addestramento, tra cui il riordinamento degli operatori e il cambio di backend. PockEngine supporta diverse applicazioni, frontend e backend hardware: compila e ottimizza in modo flessibile modelli definiti in PyTorch/TensorFlow/Jax e distribuisce binari su CPU/GPU/DSP mobili. Abbiamo valutato PockEngine sia su modelli di visione che su grandi modelli linguistici. PockEngine raggiunge un incremento di velocità fino a 15 volte rispetto a TensorFlow preconfigurato (Raspberry Pi) e un risparmio di memoria nella retropropagazione di 5,6 volte (Jetson AGX Orin). In modo notevole, PockEngine abilita il fine-tuning di LLaMav2-7B su NVIDIA Jetson AGX Orin a 550 token/s, 7,9 volte più veloce rispetto a 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.