PockEngine: Sparse and Efficient Fine-tuning in a Pocket

Abstract

Translate Article

On-device learning and efficient fine-tuning enable continuous and\nprivacy-preserving customization (e.g., locally fine-tuning large language\nmodels on personalized data). However, existing training frameworks are\ndesigned for cloud servers with powerful accelerators (e.g., GPUs, TPUs) and\nlack the optimizations for learning on the edge, which faces challenges of\nresource limitations and edge hardware diversity. We introduce PockEngine: a\ntiny, sparse and efficient engine to enable fine-tuning on various edge\ndevices. PockEngine supports sparse backpropagation: it prunes the backward\ngraph and sparsely updates the model with measured memory saving and latency\nreduction while maintaining the model quality. Secondly, PockEngine is\ncompilation first: the entire training graph (including forward, backward and\noptimization steps) is derived at compile-time, which reduces the runtime\noverhead and brings opportunities for graph transformations. PockEngine also\nintegrates a rich set of training graph optimizations, thus can further\naccelerate the training cost, including operator reordering and backend\nswitching. PockEngine supports diverse applications, frontends and hardware\nbackends: it flexibly compiles and tunes models defined in\nPyTorch/TensorFlow/Jax and deploys binaries to mobile CPU/GPU/DSPs. We\nevaluated PockEngine on both vision models and large language models.\nPockEngine achieves up to 15 $\\times$ speedup over off-the-shelf TensorFlow\n(Raspberry Pi), 5.6 $\\times$ memory saving back-propagation (Jetson AGX Orin).\nRemarkably, PockEngine enables fine-tuning LLaMav2-7B on NVIDIA Jetson AGX Orin\nat 550 tokens/s, 7.9$\\times$ faster than the PyTorch.\n