PockEngine: Sparsame und effiziente Feinabstimmung in der Tasche

PockEngine: Sparse and Efficient Fine-tuning in a Pocket

October 26, 2023

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

cs.AI

papers.abstract

On-Device-Learning und effizientes Fine-Tuning ermöglichen eine kontinuierliche und datenschutzfreundliche Anpassung (z. B. lokales Fine-Tuning großer Sprachmodelle auf personalisierten Daten). Allerdings sind bestehende Trainingsframeworks für Cloud-Server mit leistungsstarken Beschleunigern (z. B. GPUs, TPUs) konzipiert und verfügen nicht über die Optimierungen für das Lernen am Edge, das mit Herausforderungen wie Ressourcenbeschränkungen und der Vielfalt von Edge-Hardware konfrontiert ist. Wir stellen PockEngine vor: eine kompakte, spärliche und effiziente Engine, die Fine-Tuning auf verschiedenen Edge-Geräten ermöglicht. PockEngine unterstützt spärliche Backpropagation: Es beschneidet den Rückwärtsgraphen und aktualisiert das Modell spärlich, wodurch messbare Speichereinsparungen und Latenzreduktionen erzielt werden, während die Modellqualität erhalten bleibt. Zweitens ist PockEngine kompilierungsorientiert: Der gesamte Trainingsgraph (einschließlich Vorwärts-, Rückwärts- und Optimierungsschritte) wird zur Kompilierzeit abgeleitet, was den Laufzeitaufwand reduziert und Möglichkeiten für Graphentransformationen eröffnet. PockEngine integriert zudem eine Vielzahl von Trainingsgraph-Optimierungen, wodurch die Trainingskosten weiter beschleunigt werden können, einschließlich Operator-Neuordnung und Backend-Wechsel. PockEngine unterstützt diverse Anwendungen, Frontends und Hardware-Backends: Es kompiliert und optimiert flexibel Modelle, die in PyTorch/TensorFlow/Jax definiert sind, und stellt Binärdateien für mobile CPUs/GPUs/DSPs bereit. Wir haben PockEngine sowohl an Vision-Modellen als auch an großen Sprachmodellen evaluiert. PockEngine erreicht eine bis zu 15-fache Beschleunigung gegenüber Standard-TensorFlow (Raspberry Pi) und eine 5,6-fache Speichereinsparung bei der Backpropagation (Jetson AGX Orin). Bemerkenswerterweise ermöglicht PockEngine das Fine-Tuning von LLaMav2-7B auf NVIDIA Jetson AGX Orin mit 550 Tokens/s, was 7,9-mal schneller ist als 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.