Efficiently Distilling LLM s for Edge Applications (2024)

Abstract: Supernet training of LLMs is of great interest in industrial applications as it confers the ability to produce a palette of smaller models at constant cost, regardless of the number of models (of different size / latency) produced. We propose a new method called Multistage Low-rank Fine-tuning of Super-transformers (MLFS) for parameter-efficient supernet training. We show that it is possible to obtain high-quality encoder models that are suitable for commercial edge applications, and that while decoder-only models are resistant to a comparable degree of compression, decoders can be effectively sliced for a significant reduction in training time.

## introduction given their sizes up to billions of parameters, @xcite , it is challenging for enterprises to fine-tune large language models (llms), and furthermore they are not suitable for deployment on edge devices with limited memory and computational power. we wish to enable llms on edge environments for enterprise use cases. this requires the following two capabilities. (1) accommodating a variety of edge device hardware: a single fine-tuned model is not optimal across the spectrum of devices. for industrial applications, a palette of fine-tuned llms is required for different hardware. (2) dynamically changing resource levels: at run-time, the available resources on edge devices evolve over time, and appropriate model should be dynamically selected based on the available resources of each device. a considerable amount of research has focused on compressing llms @xcite @xcite @xcite . methods that train a single small model guided by a large teacher model such as distilbert @xcite and bert-pkd @xcite , either achieve limited compression or do not scale to a large number of deployment devices. supernet training methods @xcite @xcite @xcite were introduced to address these limitations: multiple smaller subnets within the supernet are trained simultaneously with weight-sharing. this one-time training approach produces a palette of smaller models, helping mitigate the computational cost of fine-tuning a model for each deployment scenario. however, the fullparameter supernet training approach is impractical when fine-tuning of an llm is required for multiple deployment scenarios, limiting its utility for enterprises. parameter-efficient fine-tuning (peft) methods such as low-rank adaptation (lora) reduces the number of trainable parameters by allowing only rank-decomposition matrices to be trained while freezing the pre-trained weights of the model. peft methods, however, are not applicable to supernet training due to the implications on the weight-shared sub-networks. our work bridges this gap to enable efficient fine-tuning of llms for edge devices. our contributions are: classical compression methods have been used for llms including pruning @xcite , low rank approximation @xcite , and quantization @xcite @xcite . knowledge distillation (kd) is adopted in bert-pkd @xcite , tiny-bert @xcite , and distilbert @xcite and @xcite in minillm to distill knowledge from the layers of a large transformer model to a smaller one. see also the survey @xcite . all these existing methods produce a single compressed model, unsuitable for edge scenarios with multiple deployment devices having varying computational capability. neural architecture search (nas) based on reinforcement learning @xcite and evolutionary algorithms @xcite trains every possible architecture and is very slow. weight-sharing nas was thus developed: in @xcite ; @xcite , the building blocks in the same layer are isolated as all architectures are single paths. weight-sharing nas does not scale well to large architecture search spaces, hence, weight-entangled nas, where subnets with common parts share weights, was introduced. for resource-constrained edge deployment, supernet training @xcite @xcite @xcite was developed as a mode of jointly training multiple sub-networks (subnets) with entangled weights: one trains the supernet only once for all deployment scenarios. @xcite introduced an elastic convolutional neural network with "progressive shrinkage", where larger subnets are trained first. recent works have improved sampling strategies, e.g. the sandwich rule with in-place distillation @xcite , attentive sampling @xcite , stochastic nature gradient @xcite , or post-training sampling @xcite . our work is related to supernet training for transformer models @xcite @xcite @xcite . this gradient scaling technique can be used with any of the above supernet methods. parameter-efficient fine-tuning (peft) has been of great benefit in fine tuning llms. bitfit @xcite updates the bias terms in pre-trained models while freezing the remaining parameters. lora @xcite decomposes attention weight gradients into low-rank matrices to reduce the number of trainable parameters. adalora @xcite and qlora @xcite further improve lora @xcite . note that peft allows fine-tuning a base model on a single gpu but does not produce smaller models. none of the peft methods can be used for weight-sharing supernet training. ## solution design for use in enterprise settings, the solution must allow fine-tuning of models on a small gpu footprint. in addition, inference cost in terms of storage must be minimised. we therefore design a solution which does not store the full size model checkpoint for every downstream task but only the frozen weights of the pre-trained base model and the low rank matrices. for inference in commercial edge use cases, we wish to enable storing the desired models locally for a wide variety of edge device resource requirements. we thus develop an approach where storage is minimised, storing only one base model and as many low rank adapter matrices as there are target model size variations, where lowrank adapters are very small. if the model is stored locally on an edge device, our proposed slicing operation takes place where the supernet fine-tuning is performed and the desired model is downloaded for inference. the slicing operation takes place for each model size-task combination and each resulting subnet can be cached for inference. ## problem formulation first, we provide notation. given a transformer model with architectural configuration φ and weights w, we denote its forward-pass mapping by f φ (•; w) : @xmath0. we consider the output space y to be the set of all non-negative vectors in r ν with elements summing to 1, where ν denotes the number of classes / vocabulary size). with slight abuse of notation, we write the forward-pass mapping of an input @xmath1, z, @xmath2; w), where ŷ ∈ y denotes the predicted probability distribution over the (class labels) vocabulary, z denotes the vector of logits, and h represents a tuple of features such as hidden state vectors and attention values from different transformer layers. note that ŷ = σ(z), where σ is the standard soft-max function that maps a vector of logits into a probability vector. given a training data set d train ⊂ x × y, model weights w are learnt by minimizing training loss: next, we introduce the super-transformer and related terminologies. we define three types of networks -teacher network, super-transformer (supernet) and sub-transformer (subnet). the teacher is a fixed network with the same configuration as the pre-trained transformer. a super-transformer is a dynamic model whose architectural dimensions (embedding dimension, number of heads, number of layers, etc.) are configurable at run time. the maxnet (resp. minnet) is the largest (resp. smallest) network in the super-transformer's architecture space. weight entanglement (weight-sharing) allows super-transformer weights to be used across sub-transformers, which are subsets of the supertransformer. pre-trained transformer weights initialise the super-transformer. the dynamic nature of a super-transformer is explicitly specified via a set a, called configuration space, consisting of architectural configurations of all sub-transformer models under consideration. the definition of a super-transformer also includes how the configuration φ ∈ a is to be mapped to a unique transformer model f φ . a weight-sharing super-transformer uses a set of shared weights w sup to define all sub-transformer models' weights. this is done through a weight projection operator π that slices (selects an appropriate subset of) the super-transformer's weights w sup into weights of a sub-transformer model: the aim of a weight-sharing super-transformer is to simultaneously train all the transformer models {f φ (•; π φ (w)) : @xmath3. a typical training objective for super-transformers is the training loss averaged over all model configurations in a: ## mlfs we therefore developed multistage low-rank finetuning of super-transformers (mlfs). given a teacher model with configuration φ tch and pretrained weights w pretrain tch , we assume that its weights (denoted w tch ) can be fine-tuned on the given task by learning low-rank matrices a 0 , b 0 on top of pre-trained weights w pretrain tch : load the sub-transformer model φ j with weights w φ j := π φ j (w sup ). for the case of maxnet (φ 1 ) set the distillation factor α to 0. find the loss: compute gradients (∇ as loss i j , ∇ bs loss i j ) using backward-pass on sub-transformer φ j . update a s , b s using the gradients ( ∇as l sup , ∇bs l sup ) of the super-transformer's loss: stage-s of the fine-tuning process involves learning only the low-rank matrices, a s , b s , by minimizing the super-transform loss as in (3). in stage-1, we sample sub-transformer models by sampling different widths from the super-transformer keeping the depth (number of layers) same as the maxnet. in stage-2, we sample sub-transformer models by sampling different widths as well as depths. we always sample the maxnet model from the supertransformer as the 1 st sub-transformer model, φ 1 , at every iteration. we call this multistage lowrank fine-tuning of super-transformers (mlfs) and present it in algorithm 1. proposition 1 let the individually fine-tuned weights of a subnet, φ, be expressed as w φ = π φ (w pretrain tch ) + ∆w φ . then, mlfs has the following structure on ∆w φ : to illustrate the computational savings, recall w pretrain @xmath4, where d is typically of the order 10 4 -10 6 . for rank r (@xmath5: a @xmath6, b @xmath7, @xmath8, 1, 2, where r ≪ d. then, the number of parameters to be learnt in the mlfs approach is 6rd. in contrast, full fine-tuning requires updating d 2 parameters at every iteration. gradient scaling for faster convergence of the smaller sub-transformers within a super-transformer, we propose a novel weighted-combination of the gradients of the sampled sub-transformers. proposition 2 let 1 st sampled sub-transformer, φ 1 , be the maxnet be in every iteration. then the scaled gradient of the super-transformer training loss, l sup , in algorithm 1 is given by proof: each sub-transformer gradient in (10), grad j , is scaled by (n 1 /n j ), which is obtained from the relative weighting of the loss functions. let l j (w) denote the j-th sub-transformer's loss. using first-order taylor expansion, we get: to define the distillation based losses precisely, let the forward-pass mapping of an input training sample x i through sub-transformer φ j be ŷi j , z i j , h i φ j ← f φ j (x i ; w φ j ), where h i j := (h i,1 j , . . . , h i,l j , . . .) with h i,l j denoting the feature vector from l-th layer of sub-transformer φ j . in super-transformers, the model (maxnet) having the largest configuration, φ 1 , acts as the teacher and knowledge distillation loss for all other subtransformers w.r.t the teacher is defined as ## results on encoder and decoder llms we report performance on encoder tasks using glue @xcite with bert base as the teacher model φ tch . for decoder llms, we use santacoder @xcite and @xcite on a python coding task using bigcode/the-stack data @xcite . we report performance of the sub-transformer models at the end of stage @xmath9. on glue, we use the train set for fine-tuning and the dev set for accuracy evaluation. for santacoder, we evaluate performance using humaneval @xcite and report pass@1 scores. all experiments were conducted using pytorch on a single nvidia a100 (40gb) gpu. additional details on the experiment settings are provided in the appendix. ## perspectives enterprise users require an efficient way to finetune llms for inference on edge devices of many sizes. we developed mlfs for such edge deployment scenarios. we demonstrate its benefits on encoder llms. we show the limitation of compressing decoder llms to a comparable degree; however, mlfs offers significant gains for smaller decoder training/fine-tuning by slicing from a larger pre-trained teacher.