Closed KuangjuX closed 6 months ago
PockEngine 是一个用于在边缘设备做训练微调的小的稀疏的和有效率的引擎。PockEngine 支持稀疏反向传播，编译优先，全部的训练计算图（包括前向，反向和优化）在编译时被提取，减少运行时开销并且为 graph transformations 带来优化。PockEngine 可以在 Nivida Jetson AGX Orin 上对 LLaMav2-7B 做 fine-tune，比 Pytorch 快 7.9 倍。
On-device 主要有以下困难：
为了应对这些挑战，本篇文章提出了 PockEngine，一个小的和有效率的训练引擎为了 on-device training。它有如下特性：
深度学习的成功建立在很多受欢迎的框架上例如 PyTorch，TensorFlow，MXNet，JAX 等等。这些系统被设计用于灵活开发并依赖于 host language（Python）执行。这带来了昂贵的内存开销并且在低频率的 CPU 例如 ARM Cortex 上执行非常慢。
边缘设备有限制的计算能力。因此，on-device 训练经常需要关注迁移学习。它首先在大规模数据集上预训练模型去学习通用的和丰富的特性，例如 ImageNet 为了 ConvNets 或者 BooksCorpus 为了 BERT。模型随后转移到 downstream 任务重，例如 Visual Wake Words 为了视觉或者 GLUE benchmark 为了语言。在那之后模型可以被定制为小规模个人数据去在相同的任务上做得更好。
但现有的训练框架无法通过稀疏反向传播实现实际的加速或内存节省。PockEngine为此类灵活工作负载提供了系统级支持，以提供更快的程序和高效的运行时。
在推理场景中存在大量的图转换技术。例如，在边缘部署中常用的一种转换是数据布局转换，因为GPU训练中首选的 NCHW 布局在边缘设备上效率不高。另一个常见的优化技术是层融合。IO密集型层（如ReLU）通常可以融合到前面的计算密集型层（如CONV、LINEAR）中。此外，MetaFlow 提出了保持功能的图转换方法来优化DNN架构。TASO 进一步引入了使用形式验证自动生成转换规则的自动化方法。这些技术在推理中已被证明有效，但很少有研究探讨它们在训练中的性能，尽管训练图更加复杂。站在传统智慧的肩膀上，PockEngine是对这些图优化技术在设备端训练中应用的早期探索，并发现了更多潜在的优化。PockEngine显示这些优化可以带来高达1.2倍的加速。
当前存在的这些训练框架（例如 PyTorch，TensorFlow）基于在运行时做自动微分。然而这种方式在边缘设备不使用。因此 PockEngine 是基于编译的 workflow：
在小规模上更新全部模型不一定是一个好的选择，更新模型的一个子集被认为是一个好的解决方案在小的训练开销下。
对于迁移学习来说不是所有权重都同等重要，目标是仅仅 fine-tune 重要的权重去在保证模型精度的情况下减少训练开销。
在我们得到静态的，剪枝后的训练图后，PockEngine 在 unified IR 上进行图优化随后将其翻译到不同的 backend。
PockEngine: Sparse and Efficient Fine-tuning in a Pocket
Abstract
PockEngine 是一个用于在边缘设备做训练微调的小的稀疏的和有效率的引擎。PockEngine 支持稀疏反向传播，编译优先，全部的训练计算图（包括前向，反向和优化）在编译时被提取，减少运行时开销并且为 graph transformations 带来优化。PockEngine 可以在 Nivida Jetson AGX Orin 上对 LLaMav2-7B 做 fine-tune，比 Pytorch 快 7.9 倍。
Introduction
On-device 主要有以下困难：
为了应对这些挑战，本篇文章提出了 PockEngine，一个小的和有效率的训练引擎为了 on-device training。它有如下特性：
Related Work
深度学习的成功建立在很多受欢迎的框架上例如 PyTorch，TensorFlow，MXNet，JAX 等等。这些系统被设计用于灵活开发并依赖于 host language（Python）执行。这带来了昂贵的内存开销并且在低频率的 CPU 例如 ARM Cortex 上执行非常慢。
Edge Deep Learning Systems
Efficient On-Device Learning Algorithms
边缘设备有限制的计算能力。因此，on-device 训练经常需要关注迁移学习。它首先在大规模数据集上预训练模型去学习通用的和丰富的特性，例如 ImageNet 为了 ConvNets 或者 BooksCorpus 为了 BERT。模型随后转移到 downstream 任务重，例如 Visual Wake Words 为了视觉或者 GLUE benchmark 为了语言。在那之后模型可以被定制为小规模个人数据去在相同的任务上做得更好。
但现有的训练框架无法通过稀疏反向传播实现实际的加速或内存节省。PockEngine为此类灵活工作负载提供了系统级支持，以提供更快的程序和高效的运行时。
Computation Graph Transformation and Optimizations
在推理场景中存在大量的图转换技术。例如，在边缘部署中常用的一种转换是数据布局转换，因为GPU训练中首选的 NCHW 布局在边缘设备上效率不高。另一个常见的优化技术是层融合。IO密集型层（如ReLU）通常可以融合到前面的计算密集型层（如CONV、LINEAR）中。此外，MetaFlow 提出了保持功能的图转换方法来优化DNN架构。TASO 进一步引入了使用形式验证自动生成转换规则的自动化方法。这些技术在推理中已被证明有效，但很少有研究探讨它们在训练中的性能，尽管训练图更加复杂。站在传统智慧的肩膀上，PockEngine是对这些图优化技术在设备端训练中应用的早期探索，并发现了更多潜在的优化。PockEngine显示这些优化可以带来高达1.2倍的加速。
Compilation-Based Workflow
当前存在的这些训练框架（例如 PyTorch，TensorFlow）基于在运行时做自动微分。然而这种方式在边缘设备不使用。因此 PockEngine 是基于编译的 workflow：
Sparse Backpropagation and Compution Graph Pruning
在小规模上更新全部模型不一定是一个好的选择，更新模型的一个子集被认为是一个好的解决方案在小的训练开销下。
PockEngine
Searching for Sparse Backpropagation Scheme
对于迁移学习来说不是所有权重都同等重要，目标是仅仅 fine-tune 重要的权重去在保证模型精度的情况下减少训练开销。
Training Graph Optimization
在我们得到静态的，剪枝后的训练图后，PockEngine 在 unified IR 上进行图优化随后将其翻译到不同的 backend。