优化移动端大模型微调内存布局,提速超5倍
FBLayout: Optimizing Memory Layout for Efficient LLM Finetuning on Mobile GPUs

- 统一R-Tile布局支持前向与反向传播的多维运算
- 通过分块索引转换避免数据物理移动,减少开销
- 根据激活情况全局选择高效布局,适合移动设备开发者
基于Transformer的模型在语言、视觉和多模态任务中表现出色。在设备端微调这些模型可保护隐私,但受限于移动端GPU的严重内存约束及注意力机制训练中频繁的布局变换,效率低下。现有移动端训练框架要么采用前后向统一布局,导致反向传播时内存访问碎片化、GPU利用率低;要么依赖显式布局转换,引入显著转换开销。为此,我们提出FBLayout,一种面向移动端GPU平台的布局感知框架。FBLayout引入:(1) 面向前后向多维归约的统一R-Tile布局;(2) 基于分块的索引转换,消除物理数据移动;(3) 激活引导的布局选择,实现全局高效布局传播。在多种手机(含ARM Mali与Qualcomm Adreno GPU)上的七种Transformer模型评估显示,相比MNN、TFLite和TVM,FBLayout实现2.2–5.7倍加速,显著提升缓存效率并降低内存占用,使实用的设备端大模型微调成为可能。
原文摘要 · Abstract (English)
Transformer-based models have enabled unprecedented capabilities across language, vision, and multimodal tasks. On-device fine-tuning of transformer models offers a privacy-preserving path to personalized AI, yet remains inefficient on mobile GPUs due to severe memory constraints and frequent layout transformations in attention mechanism during training. Existing mobile training frameworks either use unified layouts for forward and backward passes -- leading to fragmented memory access and poor GPU utilization during backpropagation -- or rely on explicit layout conversions, which introduce significant transformation overhead. To overcome this, we propose FBLayout, a layout-aware framework that co-designs tensor organization with mobile GPU platforms. FBLayout introduces: (1) a unified R-Tile layout for multi-dimensional reductions across forward/backward passes; (2) tile-based index transformation to eliminate physical data movement; and (3) activation-guided layout selection to propagate efficient layouts globally. Evaluations on seven transformer models across different mobile phones (including ARM Mali and Qualcomm Adreno GPUs) show that FBLayout achieves 2.2-5.7x speedup over MNN, TFLite, and TVM, while significantly improving cache efficiency and reducing memory footprint, enabling practical on-device large model fine-tuning.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。