PockEngine：口袋中的稀疏高效微调引擎
机器学习
2023-10-30 v1
摘要
设备端学习与高效微调实现了持续且保护隐私的定制化（例如，在个性化数据上本地微调大语言模型）。然而，现有训练框架是为配备强大加速器（如 GPU、TPU）的云服务器设计的，缺乏针对边缘学习的优化，而边缘面临资源受限与边缘硬件多样性的挑战。我们介绍 PockEngine：一个微小、稀疏且高效的引擎，用于在各类边缘设备上实现微调。PockEngine 支持稀疏反向传播：它剪枝反向图并以测得的内存节省与延迟降低稀疏更新模型，同时保持模型质量。其次，PockEngine 以编译优先：整个训练图（包括前向、反向和优化步骤）在编译期导出，从而减少运行时开销并带来图变换的机会。PockEngine 还集成了一组丰富的训练图优化，因此可进一步加速训练开销，包括算子重排序与后端切换。PockEngine 支持多样的应用、前端与硬件后端：它灵活编译和调优在 PyTorch/TensorFlow/Jax 中定义的模型，并将二进制部署到移动 CPU/GPU/DSP。我们在视觉模型与大语言模型上评估了 PockEngine。PockEngine 相比现成的 TensorFlow（Raspberry Pi）实现高达 15 加速，5.6 反向传播内存节省（Jetson AGX Orin）。值得注意的是，PockEngine 在 NVIDIA Jetson AGX Orin 上以 550 tokens/s 微调 LLaMav2-7B，比 PyTorch 快 7.9。
引用
@article{arxiv.2310.17752,
title = {PockEngine: Sparse and Efficient Fine-tuning in a Pocket},
author = {Ligeng Zhu and Lanxiang Hu and Ji Lin and Wei-Chen Wang and Wei-Ming Chen and Chuang Gan and Song Han},
journal= {arXiv preprint arXiv:2310.17752},
year = {2023}
}