Source-linked AI summary
Confidant: Customizing Transformer-based LLMs via Collaborative Edge Training
Yuhao Chen, Yuxuan Yan, Qianqian Yang, Yuanchao Shu, Shibo He, Jiming Chen
TL;DR
Fine-tuning transformer-based LLMs on mobile devices is constrained by limited memory, computing capacity, and energy. Confidant partitions models across devices, uses pipeline parallelism, and schedules attention heads across heterogeneous backends. Preliminary results report up to 45.3% memory reduction and 8.03x speedup versus single-device fine-tuning.
Problem
Mobile devices have limited memory, computing capacity, and energy, making LLM deployment and fine-tuning difficult.
Source: §1 · p. 1Method
Confidant partitions an LLM across mobile devices, applies pipeline-parallel training, and allocates attention heads across heterogeneous CPU and GPU backends.
Source: §1 · p. 2, §2.2 · p. 2, §3.1 · p. 3Results
Up to 45.3% memory reduction and 8.03x speedup are achieved versus fine-tuning on a single mobile device.
Source: §1 · p. 2, §4.2 · p. 5, §4.4 · p. 5Takeaways & Limitations
Confidant demonstrates preliminary feasibility of customizing transformer-based LLMs through collaborative training on commodity mobile devices.
Source: §1 · p. 2, §6 · p. 6The authors identify changing available memory and mobile-device battery power as unresolved reliability challenges for future work.
Source: §5 · pp. 5–6
Abstract
from arXiv · showhide
Transformer-based large language models (LLMs) have demonstrated impressive capabilities in a variety of natural language processing (NLP) tasks. Nonetheless, it is challenging to deploy and fine-tune LLMs on mobile edge devices with limited computing, memory, and energy budgets. In this paper, we propose Confidant, a multi-backend collaborative training framework for customizing state-of-the-art LLMs on commodity mobile devices like smartphones. Confidant partitions an LLM into several sub-models so that each fits into a mobile device's memory. A pipeline parallel training mechanism is further developed to ensure fast and efficient distributed training. In addition, we propose a novel backend scheduler to allocate different attention heads to heterogeneous compute hardware, including mobile CPU and GPUs, to maximize the compute resource utilization on each edge device. Our preliminary experimental results show that Confidant achieves at most 45.3% memory reduction and 8.03x inference speedup in practical settings.
1 INTRODUCTION
Confidant addresses the memory and compute constraints that make LLM fine-tuning difficult on mobile devices by combining collaborative partitioning, pipeline parallelism, and multi-backend scheduling. Preliminary results report up to 45.3% memory reduction and 8.03x speedup versus single-device fine-tuning.
Memory Constraints
- Approximately 5GB of memory is required to fine-tune BERT with batch size 8 on a Redmi K10X Pro.The measurement uses MNN and illustrates why fine-tuning can exceed practical mobile memory availability.Source: §1 · p. 1, Figure 1 · p. 1
- Mobile devices may have more than 5GB of memory, but concurrent applications mean that amount is not consistently available for LLM fine-tuning.The authors therefore identify memory-efficient training as necessary for resource-constrained devices.Source: §1 · p. 1, Table 1 · p. 1
Compute Constraints
- CPU, OpenCL, and Metal show self-attention computation times of the same order of magnitude, with OpenCL sometimes slower than CPU.Mobile frameworks also restrict DNN computation to a single backend per device, leaving available compute capacity underused.Source: Figure 2 · p. 1, §1 · p. 2
Contributions
- 45.3% memory reduction and 8.03x speedup are achieved by Confidant versus fine-tuning on a single mobile device.The framework partitions the LLM across devices, applies pipeline-parallel training, and schedules attention heads across heterogeneous backends.Source: §1 · p. 2
- Confidant partitions an LLM across multiple mobile devices and uses pipeline-parallel training to reduce stalling during collaborative fine-tuning.A backend scheduler allocates different numbers of attention heads to available CPU and GPU backends for concurrent computation.Source: §1 · p. 2
2 BACKGROUND
The background explains how self-attention generates outputs from multiple heads and how pipeline parallelism distributes model stages across devices. Confidant builds on these mechanisms with dynamic partitioning and a scheduler that distributes attention heads across available backends.
2.1 Self Attention Mechanism
- Each self-attention head transforms the input into query, key, and value matrices, computes attention scores, and produces an output later concatenated with other heads.Different head-specific weights capture distinct patterns and dependencies in the input.Source: §2.1 · p. 2
2.2 Pipeline Parallel Training
- With three devices, pipeline training forwards later batches without waiting for earlier backward passes, while versioned weights support the required backward computations.This organization is illustrated as device-to-device collaborative training in Figure 4.Source: §2.2 · p. 2, Figure 4 · p. 3
- Pipeline parallel training divides a model into sub-models on different devices and overlaps forward and backward processing using stale sub-model weights.The 1F1B rule alternates local forwarding and backwarding, while weights are stashed for the corresponding backward pass.Source: §2.2 · p. 2
Backend Scheduler
- The prior dynamic-partitioning approach estimates worker computing capacity and calculates model partition points accordingly.This capacity-aware partitioning forms the basis for adapting pipeline stages to heterogeneous edge devices.Source: §2.2 · pp. 2–3
3 DESIGN OF CONFIDANT
Confidant combines pipeline-parallel model partitioning with multi-backend attention-head scheduling to train transformer LLMs collaboratively across mobile devices. Its scheduler profiles heterogeneous backends and allocates heads to balance their computation time, supported by mobile-framework implementation extensions.
3.1 System Overview
- Confidant dynamically partitions a transformer LLM into sub-models trained on different devices through pipeline parallelism.The design extends prior edge pipeline training from image models to transformer-based LLMs.Source: §3.1 · p. 3
3.2.1 Backend Profiling
- The backend scheduler profiles available devices and measures computation times for different attention-head counts using alternative execution strategies.For each backend, it retains the faster timing between fused large-tensor and individually computed heads.Source: §3.2 · p. 3, §3.2.1 · p. 3
3.2.2 Attention Heads Allocation
- Binary search allocates attention heads so participating backends require roughly equal computation time, bypassing backends whose timing deviates beyond a threshold.The resulting allocation is executed in parallel across the device’s backends.Source: §3.2.2 · pp. 3–4
3.3 Implementation
- The implementation uses MNN with JNI and HTTP communication to connect mobile applications and distributed devices.Pretrained PyTorch weights are converted through ONNX into MNN, while separate computation graphs parallelize each backend’s assigned heads.Source: §3.3 · p. 4
4 EVALUATION
Confidant is evaluated by fine-tuning BERT for CoNLL-2003 classification across mobile devices and comparing memory and latency against single-device and pipeline baselines. The experiments report substantial memory savings and speedups from multi-backend scheduling and pipeline collaboration.
4.1 Evaluation Settings
- The evaluation fine-tunes pretrained BERT with an added classification layer on CoNLL-2003 across the mobile devices listed in Table 2.Training uses the proposed Confidant framework and compares its latency against Single, Single-MBS, and Pipeline baselines.Source: §4.1 · p. 4, §4.4 · p. 5, Table 2 · p. 5
4.2 Memory Usage
- Confidant reduces average memory usage by 45.3% at batch size 8 versus fine-tuning on a single Redmi 10X Pro.The three-device configuration does not increase memory linearly because devices store multiple weight and intermediate-output versions for 1F1B training.Source: §4.2 · pp. 4–5, Figure 5 · p. 5
4.3 Multi-backend Parallel
- On a MacBook Pro, CPU-and-Metal execution accelerates a 12-head self-attention layer compared with either CPU-only or Metal-only execution.The scheduler selects CPU and Metal while omitting slower OpenCL.Source: §4.3 · p. 5, Figure 6 · p. 5
4.4 Training Latency
- The multiple-backend scheduler produces speedups of up to 4.53x, 3.97x, and 4.09x on the three evaluated phones.These results compare single-backend training with scheduler-based multi-backend training on each phone.Source: §4.4 · p. 5
- Confidant achieves an 8.03x maximum speedup over training on a single phone, while pipeline parallelism alone reaches 3.84x.The comparison covers training latency under the evaluated batch sizes and baselines.Source: §4.4 · p. 5
5 DISCUSSION AND FUTURE WORK
The paper reports preliminary progress on collaborative edge training and identifies memory variability, energy reliability, and framework heterogeneity as future-work priorities.
Memory Adaptation
- Varying available memory can interrupt training, motivating techniques that adapt the training process to changing memory budgets.The paper specifically notes that users may run other applications concurrently, causing abrupt memory changes and insufficient-memory interruptions.Source: §5 · pp. 5–6
Training Latency
- Figure 7 compares training latency between single devices and the proposed system, providing a basis for assessing collaborative-training speed.The supplied caption identifies the comparison but does not report individual latency values.Source: Figure 7 · p. 6
Fault Tolerance for Energy-aware Training
- Energy-aware training and fault tolerance remain future goals for improving fine-tuning reliability on battery-powered mobile devices.The proposed direction predicts remaining battery life and relocates workloads to devices with sufficient battery power.Source: §5 · p. 6
Cross Framework Implementation
- Collaborative training across heterogeneous edge devices requires methods that support different DNN frameworks, such as MNN on phones and PyTorch on laptops.The paper identifies cross-framework collaboration as necessary when users possess different device types.Source: §5 · p. 6
6 CONCLUSION
Confidant fine-tunes transformer-based LLMs across mobile devices by partitioning models and combining pipeline parallelism, dynamic partitioning, and multi-backend scheduling. Preliminary evaluations indicate reduced per-device memory usage and faster training.
6 CONCLUSION
- Confidant partitions LLM sub-models across mobile devices, uses pipeline parallel training and dynamic model partitioning, and schedules attention heads across multiple backends.The system was implemented on mobile devices using an industry-scale DNN platform.Source: §6 · p. 6