Resource Management for Low-latency Cooperative Fine-tuning of Foundation Models at the Network Edge

Abstract

The emergence of large-scale foundation models (FoMo’s) that can perform human-like intelligence motivates their deployment at the network edge for devices to access state-of-the-art artificial intelligence (AI). For better user experiences, the pre-trained FoMo’s need to be adapted to specialized downstream tasks through fine-tuning techniques. To transcend a single device’s memory and computation limitations, we advocate multi-device cooperation within the device-edge cooperative fine-tuning (DEFT) paradigm, where edge devices cooperate to simultaneously optimize different parts of fine-tuning parameters within a FoMo. The edge server is responsible for coordination and gradient aggregation. However, the parameter blocks reside at different depths within a FoMo architecture, leading to varied computation latency-and-memory cost due to gradient backpropagation-based calculations. The heterogeneous on-device computation and memory capacities and channel conditions necessitate an integrated communication-and-computation (C2) allocation of local computation loads and uplink communication resources to achieve low-latency (LoLa) DEFT. To this end, we consider the depth-ware DEFT block allocation problem. The involved optimal block-device matching is tackled by the proposed low-complexity Cutting-RecoUNting-CHecking (CRUNCH) algorithm, which is designed by exploiting the monotone-increasing property between block depth and computation latency-and-memory cost. Next, the joint bandwidthand- block allocation (JBBA) makes the problem more sophisticated, i.e., mathematically NP-hard. We observe a splittable Lagrangian expression through the transformation and analysis of the original problem, where the variables indicating device involvement are introduced to decouple the block and bandwidth allocation. Then, the dual ascent method is employed to tackle the JBBA problem iteratively. Within each iteration, block allocation and bandwidth allocation are optimized concurrently. The optimal block allocation sub-problem is solved efficiently by applying the Hungarian method facilitated by the proposed CRUNCH algorithm. On the other hand, the bandwidth allocation sub-problem is solved in closed form, shedding light on favorable allocations to resource-limited devices. Through extensive experiments conducted on the GLUE benchmark, our results demonstrate significant latency reduction achievable by LoLa DEFT for fine-tuning a RoBERTa model.