HOSL: Hybrid-Order Split Learning for Memory-Constrained Edge Training
Aakriti, Zhe Li, Dandan Liang, Chao Huang, Rui Li, Haibo Yang
TL;DR
HOSL addresses the memory bottleneck of training large models with Split Learning by assigning zeroth-order gradient estimation to the client and first-order optimization to the server. This hybrid approach preserves SL privacy and model partitioning while dramatically reducing client memory, and it preserves fast convergence through server-side FO updates. The authors prove a convergence rate of $O(\sqrt{d_c/(TQ)})$, showing the benefit of offloading computation to the server and increasing $Q$ to improve convergence, with the client dimension $d_c$ governing the rate rather than the full model size. Empirically, HOSL achieves up to $3.7\times$ lower client memory and accuracy within $0.41\%-4.23\%$ of FO baselines, while outperforming a purely ZO baseline by up to $15.55\%$ across OPT-125M and OPT-1.3B on six tasks, demonstrating practical, memory-efficient edge training for LLMs.
Abstract
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 a $\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.