PockEngine: Sparse and Efficient Fine-tuning in a Pocket
作者: Ligeng Zhu, Lanxiang Hu, Ji Lin, Wei-Chen Wang, Wei-Ming Chen, Chuang Gan, Song Han
分类: cs.LG
发布日期: 2023-10-26
期刊: 56th IEEE/ACM International Symposium on Microarchitecture (MICRO 2023)
💡 一句话要点
提出PockEngine以解决边缘设备上高效微调问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 边缘计算 微调 稀疏反向传播 编译优化 移动设备 个性化学习
📋 核心要点
- 现有的训练框架主要针对云端强大计算资源，无法有效支持边缘设备的资源限制和硬件多样性。
- PockEngine通过稀疏反向传播和编译优先的设计，优化了边缘设备上的微调过程，提升了效率和灵活性。
- 实验结果显示，PockEngine在Raspberry Pi上实现了15倍的速度提升，并在Jetson AGX Orin上实现了5.6倍的内存节省。
📝 摘要（中文）
在设备上学习和高效微调使得持续和隐私保护的个性化定制成为可能。然而，现有的训练框架主要针对云服务器，缺乏针对边缘设备的优化。本文提出PockEngine，一个小型、稀疏且高效的引擎，支持在各种边缘设备上进行微调。PockEngine通过稀疏反向传播技术减少内存占用和延迟，同时保持模型质量。此外，PockEngine采用编译优先的策略，在编译时生成整个训练图，从而降低运行时开销并支持图形转换。经过评估，PockEngine在视觉模型和大型语言模型上表现出色，显著提升了训练速度和内存效率。
🔬 方法详解
问题定义：本文旨在解决现有训练框架在边缘设备上微调大型模型时面临的资源限制和硬件多样性问题。现有方法通常依赖于强大的云计算资源，无法满足边缘设备的需求。
核心思路：PockEngine的核心思路是通过稀疏反向传播和编译优先的策略，优化微调过程，减少内存占用和延迟，同时保持模型的性能。这样的设计使得在资源受限的环境中仍能实现高效的模型训练。
技术框架：PockEngine的整体架构包括三个主要模块：前向传播、反向传播和优化步骤。整个训练图在编译时生成，减少了运行时的开销，并支持多种图形转换和优化。
关键创新：PockEngine的主要创新在于其稀疏反向传播技术和编译优先的设计理念，这与传统方法的动态计算图形成了鲜明对比，显著提升了训练效率。
关键设计：在设计中，PockEngine采用了多种训练图优化技术，如操作符重排序和后端切换，以进一步加速训练过程。同时，支持在PyTorch、TensorFlow和Jax等框架中灵活编译和调优模型。
🖼️ 关键图片
📊 实验亮点
PockEngine在实验中表现出色，在Raspberry Pi上实现了高达15倍的速度提升，并在Jetson AGX Orin上实现了5.6倍的内存节省。特别是在NVIDIA Jetson AGX Orin上，PockEngine以550个tokens/s的速度微调LLaMav2-7B，比PyTorch快7.9倍，显示出其卓越的性能。
🎯 应用场景
PockEngine的研究成果在多个领域具有广泛的应用潜力，包括移动设备上的个性化AI助手、边缘计算中的实时图像处理和智能家居设备的自适应学习等。其高效的微调能力将推动边缘设备智能化的发展，提升用户体验。
📄 摘要（原文）
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.9$\times$ faster than the PyTorch.