On Harnessing Idle Compute at the Edge for Foundation Model Training

TL;DR

Cleave tackles the centralization of foundation-model training by enabling edge devices to participate in large-scale training via a PS-centric framework and fine-grained GEMM partitioning. The selective hybrid tensor parallelism and a cost-model-driven scheduling mechanism enable memory- and communication-efficient training, while explicitly accounting for device heterogeneity and churn. Empirical evaluation shows Cleave can match cloud GPU training times and scale to thousands of edge devices, with significantly faster recovery from failures and up to 8x device scalability compared to baselines. This work offers a practical path to democratize foundation-model development by leveraging idle edge compute without sacrificing accuracy or training speed.

Abstract

The ecosystem behind foundation model development today is highly centralized and limited to large-scale cloud data center operators: training foundation models is costly, needing immense compute resources. Decentralized foundation model training across edge devices, leveraging their spare compute, promises a democratized alternative. However, existing edge-training approaches fall short: they struggle to match cloud-based training performance, exhibit limited scalability with model size, exceed device memory capacity, and have prohibitive communication overhead. They also fail to satisfactorily handle device heterogeneity and dynamism. We introduce a new paradigm, Cleave, which finely partitions training operations through a novel selective hybrid tensor parallelism method. Together with a parameter server centric training framework, Cleave copes with device memory limits and avoids communication bottlenecks, thereby enabling efficient training of large models on par with the cloud. Further, with a cost optimization model to guide device selection and training workload distribution, Cleave effectively accounts for device heterogeneity and churn. Our evaluations show that Cleave matches cloud-based GPU training by scaling efficiently to larger models and thousands of devices, supporting up to 8x more devices than baseline edge-training approaches. It outperforms state-of-the-art edge training methods by up to a factor of 10 in per-batch training time and efficiently handles device failures, achieving at least 100x faster recovery than prior methods.

Figures (10)

• Figure 1: The per-device communication volume when training Lllama2-13B with batch size 128 and sequence length 1024. SOTA approaches -- DTFM dtfm for edge and Alpa alpa for cloud. The additional communication volume of the cloud approach is due to AllReduce and AlltoAll at each layer.
• Figure 2: The workflow of Cleave from model defined in training script to DAG of GEMMs. Edges in the DAG represents the memory dependency. Each GEMM is scheduled selectively across devices with best effort communication and computation overlap.
• Figure 3: Normalized training latency for a batch (lower the better). Cleave shows the ability to match with the cloud-based performance, while others cannot.
• Figure 4: Normalized training latency for a batch with OPT-13B in comparison to multiple GPUs. Cleave can scale down runtime by leveraging more edge devices to match with multiple cloud GPUs.
• Figure 5: Memory consumption per-device using 8192 devices, systems can choose the number of devices. Cleave shows the capability of running large model training, while baselines can experience out-of-memory issue. Red line shows the mobile phone capacity 0.5GB flexnn.