PockEngine: Sparse en efficiënte fine-tuning in een pocket

PockEngine: Sparse and Efficient Fine-tuning in a Pocket

October 26, 2023

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

cs.AI

Samenvatting

On-device learning en efficiënte fine-tuning maken continue en privacy-beschermende aanpassing mogelijk (bijvoorbeeld het lokaal fine-tunen van grote taalmodellen op gepersonaliseerde data). Bestaande trainingsframeworks zijn echter ontworpen voor cloudservers met krachtige accelerators (zoals GPU's en TPU's) en missen de optimalisaties voor leren aan de edge, wat te maken heeft met uitdagingen zoals beperkte middelen en diversiteit in edge-hardware. Wij introduceren PockEngine: een compacte, sparse en efficiënte engine om fine-tuning op diverse edge-apparaten mogelijk te maken. PockEngine ondersteunt sparse backpropagatie: het snoeit de backward-grafiek en werkt het model spaarzaam bij, wat resulteert in gemeten geheugenbesparing en latentievermindering terwijl de modelkwaliteit behouden blijft. Ten tweede is PockEngine compilation-first: de volledige trainingsgrafiek (inclusief forward, backward en optimalisatiestappen) wordt afgeleid tijdens compilatie, wat de runtime-overhead vermindert en mogelijkheden biedt voor grafiektransformaties. PockEngine integreert ook een uitgebreide set van trainingsgrafiekoptimalisaties, waardoor de trainingskosten verder kunnen worden versneld, waaronder operatorherordening en backend-switching. PockEngine ondersteunt diverse toepassingen, frontends en hardware-backends: het compileert en optimaliseert flexibel modellen die zijn gedefinieerd in PyTorch/TensorFlow/Jax en implementeert binaire bestanden op mobiele CPU's/GPU's/DSP's. We hebben PockEngine geëvalueerd op zowel vision-modellen als grote taalmodellen. PockEngine behaalt een versnelling tot 15 keer ten opzichte van standaard TensorFlow (Raspberry Pi) en een geheugenbesparing van 5,6 keer bij backpropagatie (Jetson AGX Orin). Opmerkelijk is dat PockEngine fine-tuning van LLaMav2-7B op NVIDIA Jetson AGX Orin mogelijk maakt met 550 tokens/s, 7,9 keer sneller dan 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.