PockEngine: Ajuste fino escaso y eficiente en un entorno compacto
PockEngine: Sparse and Efficient Fine-tuning in a Pocket
October 26, 2023
Autores: Ligeng Zhu, Lanxiang Hu, Ji Lin, Wei-Chen Wang, Wei-Ming Chen, Chuang Gan, Song Han
cs.AI
Resumen
El aprendizaje en el dispositivo y el ajuste fino eficiente permiten una personalización continua y que preserva la privacidad (por ejemplo, ajustar localmente modelos de lenguaje grandes con datos personalizados). Sin embargo, los marcos de entrenamiento existentes están diseñados para servidores en la nube con aceleradores potentes (por ejemplo, GPUs, TPUs) y carecen de optimizaciones para el aprendizaje en el edge, que enfrenta desafíos de limitaciones de recursos y diversidad de hardware en el edge. Presentamos PockEngine: un motor pequeño, disperso y eficiente que permite el ajuste fino en diversos dispositivos edge. PockEngine soporta retropropagación dispersa: poda el gráfico de retropropagación y actualiza el modelo de manera dispersa, logrando ahorros de memoria y reducción de latencia medidos, mientras mantiene la calidad del modelo. En segundo lugar, PockEngine es "compilación primero": todo el gráfico de entrenamiento (incluyendo los pasos de avance, retropropagación y optimización) se deriva en tiempo de compilación, lo que reduce la sobrecarga en tiempo de ejecución y abre oportunidades para transformaciones del gráfico. PockEngine también integra un conjunto rico de optimizaciones del gráfico de entrenamiento, lo que permite acelerar aún más el costo del entrenamiento, incluyendo reordenamiento de operadores y cambio de backends. PockEngine soporta diversas aplicaciones, frontends y backends de hardware: compila y ajusta de manera flexible modelos definidos en PyTorch/TensorFlow/Jax y despliega binarios a CPU/GPU/DSPs móviles. Evaluamos PockEngine tanto en modelos de visión como en modelos de lenguaje grandes. PockEngine logra hasta 15 veces más velocidad que TensorFlow estándar (Raspberry Pi) y un ahorro de memoria en retropropagación de 5.6 veces (Jetson AGX Orin). Notablemente, PockEngine permite ajustar fino LLaMav2-7B en NVIDIA Jetson AGX Orin a 550 tokens/s, 7.9 veces más rápido 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.