PockEngine: Разреженная и эффективная тонкая настройка в компактном формате

PockEngine: Sparse and Efficient Fine-tuning in a Pocket

October 26, 2023

Авторы: Ligeng Zhu, Lanxiang Hu, Ji Lin, Wei-Chen Wang, Wei-Ming Chen, Chuang Gan, Song Han

cs.AI

Аннотация

Обучение на устройстве и эффективная тонкая настройка позволяют осуществлять непрерывную и конфиденциальную адаптацию (например, локальную тонкую настройку больших языковых моделей на персональных данных). Однако существующие фреймворки для обучения разработаны для облачных серверов с мощными ускорителями (например, GPU, TPU) и не оптимизированы для обучения на периферийных устройствах, которые сталкиваются с ограниченными ресурсами и разнообразием аппаратного обеспечения. Мы представляем PockEngine: компактный, разреженный и эффективный движок, который позволяет выполнять тонкую настройку на различных периферийных устройствах. PockEngine поддерживает разреженное обратное распространение: он обрезает граф обратного распространения и разреженно обновляет модель, обеспечивая значительное сокращение использования памяти и задержек при сохранении качества модели. Во-вторых, PockEngine ориентирован на компиляцию: весь граф обучения (включая прямой проход, обратное распространение и шаги оптимизации) формируется на этапе компиляции, что снижает накладные расходы во время выполнения и открывает возможности для преобразования графа. PockEngine также интегрирует множество оптимизаций графа обучения, что позволяет дополнительно ускорить процесс обучения, включая переупорядочивание операторов и переключение бэкендов. PockEngine поддерживает разнообразные приложения, интерфейсы и аппаратные бэкенды: он гибко компилирует и настраивает модели, определенные в PyTorch/TensorFlow/Jax, и развертывает бинарные файлы на мобильных CPU/GPU/DSP. Мы оценили PockEngine как на моделях компьютерного зрения, так и на больших языковых моделях. PockEngine демонстрирует ускорение до 15 раз по сравнению с готовым TensorFlow (Raspberry Pi) и сокращение использования памяти при обратном распространении в 5,6 раз (Jetson AGX Orin). Примечательно, что PockEngine позволяет выполнять тонкую настройку модели LLaMav2-7B на NVIDIA Jetson AGX Orin со скоростью 550 токенов в секунду, что в 7,9 раз быстрее, чем в 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.