PockEngine : Réglage fin parcimonieux et efficace dans un espace restreint
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
L'apprentissage sur appareil et le réglage fin efficace permettent une personnalisation continue et respectueuse de la vie privée (par exemple, le réglage local de grands modèles de langage sur des données personnalisées). Cependant, les frameworks d'entraînement existants sont conçus pour des serveurs cloud dotés d'accélérateurs puissants (par exemple, GPU, TPU) et manquent d'optimisations pour l'apprentissage en périphérie, qui fait face à des défis liés aux limitations de ressources et à la diversité du matériel périphérique. Nous présentons PockEngine : un moteur compact, épars et efficace permettant le réglage fin sur divers appareils périphériques. PockEngine prend en charge la rétropropagation éparse : il élagage le graphe de rétropropagation et met à jour le modèle de manière éparse avec des économies de mémoire mesurées et une réduction de la latence, tout en maintenant la qualité du modèle. Deuxièmement, PockEngine est axé sur la compilation : l'intégralité du graphe d'entraînement (y compris les étapes de propagation avant, de rétropropagation et d'optimisation) est dérivée au moment de la compilation, ce qui réduit la surcharge d'exécution et offre des opportunités de transformations de graphes. PockEngine intègre également un ensemble riche d'optimisations de graphes d'entraînement, permettant ainsi d'accélérer davantage le coût de l'entraînement, y compris le réordonnancement des opérateurs et le changement de backend. PockEngine prend en charge diverses applications, interfaces et backends matériels : il compile et ajuste de manière flexible des modèles définis dans PyTorch/TensorFlow/Jax et déploie des binaires sur des CPU/GPU/DSP mobiles. Nous avons évalué PockEngine sur des modèles de vision et des grands modèles de langage. PockEngine atteint jusqu'à 15 fois d'accélération par rapport à TensorFlow standard (Raspberry Pi), et une économie de mémoire de 5,6 fois lors de la rétropropagation (Jetson AGX Orin). Remarquablement, PockEngine permet de régler finement LLaMav2-7B sur NVIDIA Jetson AGX Orin à 550 tokens/s, soit 7,9 fois plus rapide que 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.