混合高低阶优化,让边缘设备高效训练大模型
HOSL: Hybrid-Order Split Learning for Memory-Constrained Edge Training
- 客户端用零阶优化避免反向传播,省下存储激活值的内存
- 服务端用一阶优化保证快速收敛,精度仅比传统方法低4.23%
- 适合内存受限的边缘设备做大模型训练,尤其适合资源紧张场景
分割学习(SL)通过在边缘设备与计算能力强的服务器之间分割模型计算,实现大型语言模型的协同训练。然而,现有系统主要依赖一阶(FO)优化,需客户端存储中间激活值,造成显著内存开销。相比之下,零阶(ZO)优化虽可大幅降低内存使用,但收敛慢、性能差。本文提出HOSL框架,将客户端的零阶优化与服务端的一阶优化结合:客户端采用内存高效的零阶梯度估计,消除反向传播和激活存储;服务端则保持一阶优化以确保快速收敛和良好性能。理论上,HOSL的收敛速率可达$/mathcal{O}( ext{√}(d_c/TQ))$,其中$ d_c $为客户端模型维度,优于全模型维度$ d $。在OPT模型(125M和1.3B参数)上共6项任务的实验表明,相比一阶方法,HOSL将客户端显存减少最高达3.7倍,且精度损失在0.20%至4.23%之间;同时相较零阶基线提升高达15.55%,验证了该混合策略在边缘设备上实现高效训练的有效性。
原文摘要 · Abstract (English)
Split learning (SL) enables collaborative training of large language models (LLMs) between resource-constrained edge devices and compute-rich servers by partitioning model computation across the network boundary. However, existing SL systems predominantly rely on first-order (FO) optimization, which requires clients to store intermediate quantities such as activations for backpropagation. This results in substantial memory overhead, largely negating benefits of model partitioning. In contrast, zeroth-order (ZO) optimization eliminates backpropagation and significantly reduces memory usage, but often suffers from slow convergence and degraded performance. In this work, we propose HOSL, a novel Hybrid-Order Split Learning framework that addresses this fundamental trade-off between memory efficiency and optimization effectiveness by strategically integrating ZO optimization on the client side with FO optimization on the server side. By employing memory-efficient ZO gradient estimation at the client, HOSL eliminates backpropagation and activation storage, reducing client memory consumption. Meanwhile, server-side FO optimization ensures fast convergence and competitive performance. Theoretically, we show that HOSL achieves an $\mathcal{O}(\sqrt{d_c/TQ})$ rate, which depends on client-side model dimension $d_c$ rather than the full model dimension $d$, demonstrating that convergence improves as more computation is offloaded to the server. Extensive experiments on OPT models (125M and 1.3B parameters) across 6 tasks demonstrate that HOSL reduces client GPU memory by up to 3.7$\times$ compared to the FO method while achieving accuracy within 0.20%-4.23% of this baseline. Furthermore, HOSL outperforms the ZO baseline by up to 15.55%, validating the effectiveness of our hybrid strategy for memory-efficient training on edge devices.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。