PockEngine: Sparse and Efficient Fine-tuning in a Pocket
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.
653f119c8a67c542ee1a755b
Ligeng Zhu
extracted_confirmed
2023-10-30T05:22:40.065Z
62b4b5beb25cb80fcf278354
https://cdn-avatars.huggingface.co/v1/production/uploads/62b4b5beb25cb80fcf278354/7uGNOeE91paZikhK2QFgp.jpeg
Ligeng Zhu
user
Ligeng-Zhu
653f119c8a67c542ee1a755c
Lanxiang Hu
claimed_verified
2024-07-08T08:51:37.425Z
6301d6455e305a35cb0846a7
https://cdn-avatars.huggingface.co/v1/production/uploads/6301d6455e305a35cb0846a7/aT2AtzRMSY_T3y02MIUap.jpeg
Lanxiang Hu
user
Snyhlxde
653f119c8a67c542ee1a755d
Ji Lin
extracted_confirmed
2023-10-30T02:59:29.559Z
62f68e83d278a8f3e784de17
https://cdn-avatars.huggingface.co/v1/production/uploads/1660325537221-62f68e83d278a8f3e784de17.jpeg
Ji Lin
user
jilinmit
653f119c8a67c542ee1a755e
Wei-Chen Wang
extracted_pending
2023-10-30T02:14:53.897Z
6514d59cf95f39fd02acd9dc
https://cdn-avatars.huggingface.co/v1/production/uploads/6514d59cf95f39fd02acd9dc/SSS28-Mjv2UdbGRoKdOZ7.jpeg
Wei-Chen Wang
user
wweichen
653f119c8a67c542ee1a755f
Wei-Ming Chen
admin_assigned
2023-10-30T14:13:20.147Z
650215dd814e13d7fe99baef
/avatars/20e1c93e82c8e44bd1258fabbc8ea427.svg
Wei-Ming Chen
user
meenchen
653f119c8a67c542ee1a7560
Chuang Gan
653f119c8a67c542ee1a7561
Song Han
extracted_pending
2023-10-30T02:14:53.897Z
63797f727df2fefdcaf3ff7e
https://cdn-avatars.huggingface.co/v1/production/uploads/1668906853549-noauth.jpeg
Song
user
songhan
PockEngine: Sparse and Efficient Fine-tuning in a Pocket
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.
PockEngine is an efficient on-device learning engine that supports sparse backpropagation and graph optimizations for fine-tuning models on edge devices with varying hardware and constraints.
sparse backpropagation sparse update model quality compilation first training graph optimizations operator reordering backend switching on-device learning efficient fine-tuning edge devices edge hardware diversity mobile CPU mobile GPU DSPs PyTorch TensorFlow Jax vision models large language models LLaMav2-7B