System and methods of optimizing inference for foundation models

Abstract

A computer-implemented method is disclosed. The method includes: receiving a request to process using a pre-trained machine learning model, the request including an input prompt, a first budget, and a task identifier of a first task; pruning the machine learning model, without storing a separate pruned model in memory, wherein the pruning comprises determining, for each layer of the machine learning model, a layer gate value indicating whether to use the layer for inference based on the first budget and the first task; and utilizing the pruned machine learning model in accordance with the layer gate values to generate an output corresponding to the request.

Description

SYSTEM AND METHODS OF OPTIMIZING INFERENCE FOR FOUNDATION MODELSTECHNICAL FIELD

[0001] The present application relates to machine learning models and, more particularly, to a system and methods of optimizing inference for pre-trained foundation models.BACKGROUND

[0002] Low-Rank Adaptation (LoRA) is a method used to fine-tune large pre-trained machine learning models, particularly in the domains of natural language processing (NLP) and computer vision. Instead of modifying all parameters of a pre-trained model, LoRA adds a small, trainable set of parameters while keeping the original pre-trained weights frozen. The trainable parameters are represented using low-rank matrices that adapt the model’s behavior.

[0003] Pruning is a model optimization technique used to reduce the size and complexity of a machine learning model without significantly affecting its performance. The process involves identifying less important parameters (e.g., weights, neurons, or layers) and removing such parameters from the model. In particular, the parameters that contribute the least to a model’s performance may be set to zero, or entire components of the model may be removed.SUMMARY

[0004] According to an aspect of the present disclosure, there is provided a computer-implemented method. The method includes: loading, in memory, a single pre-trained machine learning model; receiving a plurality of requests to be processed using the machine learning model, each request including an input prompt, a first budget, and a task identifier of a first task; and for each received request: pruning the machine learning model at inference time, wherein the pruning comprises determining, for each layer of the machine learning model, a layer gate value indicating whether to use the layer for inference based on the first budget and the first task; and utilizing the pruned machine learning model in accordance with the layer gate values to generate an output corresponding to the request.

[0005] In some implementations, determining the layer gate value may include: performing lookup of a budget table storing, for each of a plurality of combinations of tasks and budgets, a vector representing layer gate values of all layers of the machine learning model; determining that the budget table contains a first table entry corresponding to the combination of the first budget and the first task; and retrieving layer gate values associated with the first table entry.

[0006] In some implementations, the layer gate values for the layers of the machine learning model may be represented by a Boolean vector in the budget table.

[0007] In some implementations, utilizing the pruned machine learning model may include generating the output using only the layers that have a layer gate value equal to a predefined first value.

[0008] In some implementations, determining the layer gate value may comprise: performing lookup of a budget table storing, for each of a plurality of combinations of tasks and budgets, a vector representing layer gate values of all layers of the machine learning model; determining that the budget table does not contain a table entry corresponding to the combination of the first budget and the first task; and executing a layer-wise pruning method to obtain layer gate values for the first budget and the first task.

[0009] In some implementations, the method may further include storing the obtained layer gate values in the budget table.

[0010] According to another aspect of the present disclosure, there is provide a computer-implemented method. The method includes: loading, in memory, a single pre-trained machine learning model; receiving a plurality of requests to process using the machine learning model with a plurality of Low-Rank Adaptation (LoRA) adapters, each request including an input prompt, a first budget, and an identifier of a first LoRA adapter; and for each received request: pruning the machine learning model at inference time, wherein the pruning comprises determining, for each layer of the machine learning model, a layer gate value indicating whether to use the layer for inference based on the first budget and the first LoRA adapter; and utilizing the pruned machine learning model and corresponding LoRA adapters in accordance with the layer gate values to generate an output corresponding to the request.

[0011] In some implementations, determining the layer gate value may include: performing lookup of a budget table storing, for each of a plurality of combinations of LoRA identifiers and budgets, a vector representing layer gate values of all layers of the machine learning model; determining that the budget table contains a first table entry corresponding to the combination of the first budget and the first LoRA adapter; and retrieving layer gate values associated with the first table entry.

[0012] In some implementations, the layer gate values for the layers of the machine learning model may be represented by a Boolean vector in the budget table.

[0013] In some implementations, utilizing the pruned machine learning model may include generating the output using only the layers that have a layer gate value equal to a predefined first value.

[0014] In some implementations, determining the layer gate value may comprise: performing lookup of a budget table storing, for each of a plurality of combinations of LoRA identifiers and budgets, a vector representing layer gate values of all layers of the machine learning model; determining that the budget table does not contain a table entry corresponding to the combination of the first budget and the first LoRA adapter; and executing a layer-wise pruning method to obtain layer gate values for the first budget and the first LoRA adapter.

[0015] In some implementations, the method may further include storing the obtained layer gate values in the budget table.

[0016] According to another aspect of the present disclosure, there is provide a computer-implemented method. The method includes: receiving a plurality of requests to process using a pre-trained machine learning model, each request including a respective input prompt, a budget, and a task identifier of a task; for each layer of the machine learning model: identifying multiple requests that have a same layer gate value indicating that the layer is to be used for inference based on the respective budget and the task; and in response to identifying the multiple requests, executing the layer for the identified requests simultaneously in batch.

[0017] In some implementations, identifying the multiple requests may comprise: performing lookup of a budget table storing, for each of a plurality of combinations of tasks and budgets, a vector representing layer gate values of all layers of the machine learning model; determining that the budget table contains table entries corresponding to combinations of budgets and tasks associated with multiple ones of the requests; and retrieving layer gate values associated with the table entries.

[0018] In some implementations, the method may further include, for each layer of the machine learning model: identifying only a single request having a layer gate value indicating that the layer is to be used for inference based on the respective budget and the task; and in response to identifying only the single request, executing the layer for the identified request without batching.

[0019] In some implementations, the layer gate values for the layers of the machine learning model may be represented by a Boolean vector in the budget table.

[0020] According to another aspect of the present disclosure, there is provide a computer-implemented method for gradient-based learning of binary pruning gates. The method includes: parametrizing a binary pruning mask to model hard concrete distributions; and optimizing the minimax objective given by:

[0021] where θ, λ are Lagrange multipliers.

[0022] In some implementations, the hard concrete distribution may be defined as:

[0023] where U is the uniform distribution, β is a temperature parameter, s is a relaxed binary mask that conforms to the hard concrete distribution, and ζ, γ are the bounds of the hard concrete distribution.

[0024] According to another aspect of the present disclosure, there is provided a computer-readable storage medium, comprising one or more instructions, wherein when the one or more instructions are run on a computer, the computer performs any of the methods disclosed herein.

[0025] According to another aspect of the present disclosure, there is provided a non-transitory computer-readable medium storing the instructions that, when executed by a processor, cause the processor to implement any of the methods disclosed herein.

[0026] According to another aspect of the present disclosure, there is provided a device configured to perform any of the methods disclosed herein.

[0027] According to another aspect of the present disclosure, there is provided a processor configured to execute instructions to cause a device to perform any of the methods disclosed herein.

[0028] According to another aspect of the present disclosure, there is provided an integrated circuit configured to perform any of the methods disclosed herein.

[0029] According to another aspect of the present disclosure, there is provided a module comprising one or more circuits for performing any of the methods disclosed herein.

[0030] According to another aspect of the present disclosure, there is provided an apparatus comprising one or more processors functionally connected to a memory for performing any of the methods disclosed herein.

[0031] According to another aspect of the present disclosure, there is provided an apparatus configured to perform any of the methods disclosed herein.

[0032] In some embodiments, the apparatus comprises one or more units configured to perform any of the methods disclosed herein.

[0033] According to another aspect of the present disclosure, there is provided one or more non-transitory, computer-readable storage media comprising computer-executable instructions, wherein the instructions, when executed, cause at least one processing unit, at least one processor, or at least one circuit to perform any of the methods disclosed herein.

[0034] According to another aspect of the present disclosure, there is provided one or more computer-readable storage media storing a computer program, wherein, when the computer program is executed by an apparatus, the apparatus is enabled to implement any of the methods disclosed herein.

[0035] According to another aspect of the present disclosure, there is provided a computer program product including one or more instructions, wherein, when the instructions are executed by an apparatus, the apparatus is enabled to implement any of the methods disclosed herein.

[0036] According to another aspect of the present disclosure, there is provided a computer program, wherein, when the computer program is executed by a computer, an apparatus is enabled to implement any of the methods disclosed herein.

[0037] According to another aspect of the present disclosure, there is provided a system comprising a node for performing any of the methods disclosed herein.BRIEF DESCRIPTION OF THE DRAWINGS

[0038] Embodiments are described in detail below, with reference to the following drawings:

[0039] FIG. 1 illustrates an overall pipeline of an inference-time pruning framework for pre-trained foundation models, in accordance with example embodiments;

[0040] FIG. 2 is a schematic diagram illustrating a budget and task-adaptive inference pipeline for pre-trained foundation models;

[0041] FIG. 3 is a schematic diagram illustrating components of a Gate Update module in the pruning engine of FIG. 2;

[0042] FIG. 4 illustrates an exemplary process flow for batching in budget-adaptive inference;

[0043] FIG. 5 illustrates an overall pipeline of an inference-time pruning framework for pre-trained foundation models with multiple LoRA adapters, in accordance with example embodiments;

[0044] FIG. 6 is a schematic diagram illustrating a budget and task-adaptive inference pipeline for pre-trained foundation models with multiple LoRA adapters;

[0045] FIG. 7 is a schematic diagram illustrating components of a Gate Update module in the pruning engine of FIG. 6;

[0046] FIG. 8 illustrates another exemplary process flow for batching in budget-adaptive inference; and

[0047] FIG. 9 illustrates an example computing system, in accordance with example embodiments.

[0048] Like reference numerals are used in the drawings to denote like elements and features.DETAILED DESCRIPTION OF EMBODIMENTS

[0049] The systems and apparatuses disclosed herein may comprise suitable modules and / or circuitries for executing various procedures.

[0050] As those skilled in the art understand, a “module” is a term of explanation referring to a hardware structure such as a circuitry implemented using technologies such as electrical and / or optical technologies (and with more specific examples of semiconductors) for performing defined operations or processing. A “module” may alternatively refer to the combination of a hardware structure and a software structure, wherein the hardware structure may be implemented using technologies such as electrical and / or optical technologies (and with more specific examples of semiconductors) in a general manner for performing defined operations or processing according to the software structure in the form of a set of instructions stored in one or more non-transitory, computer-readable storage devices or media.

[0051] A module may be a part of a device, an apparatus, a system, and / or the like, wherein the module may be coupled to or integrated with other parts of the device, apparatus, or system such that the combination thereof forms the device, apparatus, or system. Alternatively, the module may be implemented as a standalone device or apparatus.

[0052] The module usually executes a procedure for performing a method. Herein, a procedure has a general meaning equivalent to that of a method. More specifically, a procedure is a defined method implemented using hardware components for processing data. A procedure may comprise or use one or more functions for processing data as designed. Herein, a function is a defined sub-procedure or sub-method for computing, calculating, or otherwise processing input data in a defined manner and generating or otherwise producing output data.

[0053] As those skilled in the art will appreciate, a procedure may be implemented as one or more software and / or firmware programs having necessary computer-executable code or instructions and stored in one or more non-transitory computer-readable storage devices or media which may be any volatile and / or non-volatile, non-removable or removable storage devices such as RAM, ROM, EEPROM, solid-state memory devices, hard disks, CDs, DVDs, flash memory devices, and / or the like. A module may read the computer-executable code from the storage devices and execute the computer-executable code to perform the procedure.

[0054] Alternatively, a procedure may be implemented as one or more hardware structures having necessary electrical and / or optical components, circuits, logic gates, integrated circuit (IC) chips, and / or the like.

[0055] Task-Aware and Budget-Adaptive Inference for Foundation Models

[0056] Foundation models have shown great potential and performance in different data modalities such as language and visual data (e.g., images and videos) . These models have broad understanding of different domains and show strong capabilities to perform multiple tasks. However, running the large models, either for training / fine-tuning or inference, is costly and time-consuming. A variety of techniques have been proposed to facilitate training and inference optimization for foundation models.

[0057] LoRA (Low-Rank Adaptation) is a method for fine-tuning foundation models that significantly reduces the training time while maintaining accuracy. In a neural network, the weights are typically represented as large matrices. LoRA introduces learnable low-rank matrices that adapt the pre-trained model’s weights without modifying them directly. The low-rank matrices act as an additive adjustment to the base model’s frozen weights. Given a weight matrix W, the update of weights during training is approximated by: W′=W+ ΔW,

[0058] where ΔW=A·B, and A and B are low-rank matrices with ranks much smaller than W. The efficiency in the LoRA method comes from decomposing ΔW into the two low-rank matrices. The fine-tuned LoRA weights, i.e., A and B, are referred to as LoRA adapters. During training, A and B are updated while W remains fixed. At inference time, the effective weight W′is used for predictions.

[0059] Various other techniques have been proposed for improving the inference cost of large models (e.g., memory reduction and / or runtime improvements) . One such technique is pruning, where less important parameters, such as weights, neurons, or filters, are removed to reduce the model’s size and computational complexity without significantly affecting its performance. The goal of pruning is to create a smaller, faster, and more efficient version of the base model.

[0060] Pruning methods include unstructured and structured pruning. Unstructured pruning removes individual weights (parameters) from the model without adhering to a specific structural pattern. Individual elements of a weight matrix are removed, leaving the remaining structure intact. This creates sparse matrices, where most values are zero, but the overall structure of the model remains unchanged.

[0061] Structured pruning removes entire groups of parameters (e.g., neurons, filters, layers, etc. ) based on their importance. This produces a smaller, simpler model with fewer parameters and reduced computation. Structured pruning can be classified into two types: width (i.e., removing entire rows or columns of the model's weights) and depth (i.e., complete removal of a layer, which could include one or more weights, in the model) pruning.

[0062] While pruned models offer computational efficiency and reduced size, they also come with several challenges that can impact their performance and deployment. First, aggressive pruning can remove critical parameters, leading to a significant drop in model accuracy. Second, in order to run inference with different budgets (i.e., pruning ratio) , multiple pruned original models are needed, which results in large memory requirements. Third, LoRA models trained with the original models are not compatible with pruned models, as the fundamental techniques they rely on are in conflict. LoRA assumes the base model’s architecture, and weights remain intact and unaltered. On the other hand, pruning modifies the structure of the base model by, for example, removing weights, neurons, or even layers, altering the computational graph. This structural alteration can disrupt the integration of LoRA’s additional low-rank matrices into the original model. Moreover, pruning often introduces sparsity in the weight matrices. LoRA relies on dense matrix multiplication for low-rank adaptation. Sparse pruned weights may result in inefficiencies and incompatibility with LoRA’s computational requirements.

[0063] Existing solutions for combining LoRA and model pruning techniques fall into three general categories. The first category includes general pruning techniques in which the model is pruned first, and LoRA is subsequently applied to compensate the performance loss of the pruned model. LLM-Pruner is a method that involves a structured width pruning technique for large language models, where the performance of pruned models can be recovered using LoRA. Parameter-Efficient adaPtation for the Speech foundatIon model (PEPSI) adopts an unstructured pruning strategy for multilingual speech foundation models on low-resourced languages. In this method, neurons that are irrelevant to a downstream task are removed in an unstructured way, which can reduce computation overhead. The pruned model performance can be recovered using LoRA.

[0064] The existing solutions that fall into this first category are generally not equipped to handle multi-LoRA inferencing scenarios. Each pruned model only works with its own specific LoRA. Furthermore, these solutions are not budget adaptive by default; in particular, they cannot perform on-the-air budget / task-adaptive inference-time pruning. Multiple initializations of a pruned model are required for different tasks, and the base model and LoRA adapters need to be pruned separately.

[0065] A second category includes techniques for joint pruning of a base model and LoRA adapters. In LoRAPrune, structured pruning meets LoRA fine-tuning, where LoRA-guided pruning criteria based on the weights and gradients of LoRA are used to perform model pruning and LoRA fine-tuning simultaneously. Similar to the solutions of the first category, joint pruning techniques cannot handle multi-LoRA inferencing scenarios, and each pruned model only works with its own specific LoRA. They also cannot perform on-the-air budget / task-adaptive inference-time pruning.

[0066] A third category includes methods for efficient multi-LoRA serving. Punica is a multi-tenant LoRA serving technique that holds only a single copy of the underlying pre-trained model when serving multiple different LoRA models. Punica allows batching of GPU operations for different LoRA models. S-LoRA is another method that falls into this category. Thousands of concurrent LoRA adapters can be served by S-LoRA, which is designed for scalable serving of many LoRA adapters that is useful for large-scale customized fine-tuning services. These multi-LoRA serving techniques do not improve runtime as no pruning is used and generally result in increased latency due to the extra overhead of LoRA. Furthermore, these techniques generally require complicated kernel-level modifications, and also cannot perform on-the-air budget / task-adaptive inference.

[0067] The present disclosure addresses certain technical drawbacks of existing techniques for inference optimization. Existing pruning methods can adjust a model’s computational load based on a given budget, and obtain a pruned model accordingly. However, such methods obtain a different model for each given budget. As such, a separate model needs to be loaded for each pruned model at inference time. The memory requirement for budget-adaptive inference with different pruned models increases significantly when the number of pruned models increases.

[0068] Pruned models are typically incompatible with multi-LoRA adapters. Existing LoRA adapters are trained using the weights of a pre-trained model that is not pruned. After the model is pruned, the existing LoRA adapters are no longer compatible with the pruned model. More particularly, if some weights in the base model are pruned, the low-rank matrices A and B lose their alignment and relevance, leading to unpredictable behavior or degraded performance.

[0069] When multiple input prompts are provided to a pre-trained model to be processed, each of one or more pruned models may have different inference steps. Therefore, the input prompts are processed individually in a sequence. This results in slower inference, which increases the latency of the model. Batching allows the model to process multiple inputs simultaneously, leveraging parallel computation capabilities to improve throughput. Furthermore, memory access patterns are optimized when data is processed in batches, and batching amortizes computational overhead (i.e., overhead of loading data, running operations, transferring results) across multiple inputs. It is desired to leverage batching when using pruned models to process multiple requests.

[0070] The embodiments of the present disclosure relate to inference optimization for pre-trained foundation models (e.g., GPT, BERT, LLaMA, DALL-E, etc. ) . The methods and systems proposed herein are modality-and architecture-agnostic, and can be applied to different models regardless of their architecture and data type, or modalities, they work on.

[0071] The present application describes a solution for on-the-fly (online) budget and task-adaptive pruning. Given a pre-trained model and a user-specified budget and / or desired task, the proposed solution dynamically processes the budget / task and prunes the model accordingly at inference time in order to satisfy the user’s requirements. Importantly, the solution can be applied to multi-LoRA scenarios, when the base pre-trained model is served along with LoRA adapters at inference time.

[0072] The present application also describes a solution for multi-budget / task inference with a single pre-trained model. Only a single model is required to be loaded in memory for on-the-fly serving of requests with multiple different budgets and / or desired tasks. The proposed solution allows for saving significant serving memory compared to loading multiple different models.

[0073] The present application also describes budget-adaptive batching of inputs. The proposed models can provide simultaneous serving of a large number of users or inference requests with different input budgets in an efficient manner.

[0074] Reference is made to FIG. 1, which illustrates an overview of the proposed inference-time pruning framework for pre-trained models. The inference engine 100 is a software framework or system designed to execute a trained machine learning model to perform inference, applying its learned patterns to generate outputs for given inputs. The inference engine 100 is configured to load a pre-trained model 150, by reading the model’s parameters and architecture into memory, and run the model 150 to produce predictions. The inputs provided by the user to the inference pipeline may include, without limitation, user’s prompts (e.g., a natural language query) , user's desired budget (e.g., memory, speed, price, accuracy, etc. ) for processing their request, and optional task name (s) . The available resources, as specified by the user, may be mapped to a pruning budget (e.g., pruning ratio, expressed as a percentage) .

[0075] As shown in FIG. 1, the inference engine 100 includes a pruning engine 110, a pre-processing module 120, and a post-processing module 130. The pre-processing module 120 transforms raw input data, such as scaling images, tokenizing text, and the like, into a format that is compatible with the model (e.g., embedding text into vectors) . For example, a user’s prompt may be processed and mapped to a machine-readable format that the model can interpret (e.g., converting text of a question into numerical tokens) . In the pruning engine 110, the least important components of the model may be identified and pruned on-the-fly, based on the pruning budget and tasks specified by the user. The model then performs inference using the un-pruned layers, i.e., skipping the pruned layers, to satisfy the user’s demand in terms of their given budget. The inference engine 100 executes the model’s computational graph on the inputs, producing raw outputs, such as logits, scores, etc. The model's raw outputs are transformed into user-friendly results by the post-processing module 130.

[0076] When LoRA weights corresponding to different tasks are defined, the LoRA weights may be added to the pruning engine 110. The task name which maps to corresponding LoRA weights is provided as input. An overview of the inference-time pruning framework for pre-trained foundation models with multi-LoRA adapters, as proposed in the present disclosure, is illustrated in FIG. 5. The multi-LoRA system shown in FIG. 5, as contrasted with the framework of FIG. 1, includes input of task data that indicates which LoRA is used during inference (e.g., by means of LoRA identifiers) . The pruning engine 510 prunes the pre-trained model 550 and the corresponding LoRA weights. The pruned model and the pruned LoRA weights may then be used for inference. The two pruning frameworks will be described in greater detail below with reference to FIGS. 1 to 8.

[0077] Reference is now made to FIG. 2 which illustrates a budget and task-adaptive inference pipeline for pre-trained foundation models. As shown in FIG. 2, the pre-trained model 150 consists of N layers of nodes / neurons. For example, the layers may comprise convolutional layers of a neural network. As another example, the layers may comprise embedding, positional encodings, and attention layers of a Transformer model. In some implementations, the pre-trained model 150 may include a "final head" which maps an output of the last layer N to a word in the model's vocabulary.

[0078] Each layer in the pre-trained model 150 is associated with a respective “gate” , or layer gate. A layer gate is a Boolean function whose value represents whether the corresponding layer is to be used for inference. In particular, the value of a layer gate is set to either 1 (set to “ON” ) or 0 (set to “OFF” ) during inference. When the pre-trained model 150 is first loaded in memory, all of the layer gates may be initialized to a value of 1. The budget, the values of the layer gates, and the optional task details (e.g. task name or a related dataset (DS) for a given task) may all be passed to the “Gate Update” module 170.

[0079] The Gate Update module 170 is configured to update the values of the layer gates based on the user-specified budget and task. Exemplary components of the Gate Update module 170 are shown in FIG. 3. A budget table 180 of the Gate Update module 170 stores layer gate values for the layers of the pre-trained model 150. For each of a plurality of combinations of tasks and budgets, the budget table 180 stores a corresponding vector that represents the layer gate values corresponding to the model’s layers. The vector may, for example, be a Boolean vector or binary sequence formed by concatenating layer gate values. The length of a stored vector for any task / budget is equal to the number of layers in the pre-trained model 150.

[0080] If a given task / budget combination is included in the budget table 180, the corresponding Boolean vector can be retrieved from the budget 180. Otherwise, a layer-wise pruning method is used to prune the pre-trained model 150 and obtain values of the layer gates for the given task / budget combination. A suitable off-the-shelf layer-wise pruning method may be used to determine which layers are to be pruned. An example of an off-the-shelf layer-wise pruning method is Shortened LLaMA, which uses gradients over a calibration dataset to identify the pruned layers. The layer gate values are stored in the budget table 180, and then provided to the pruning engine 110.

[0081] After the layer gate values are obtained from Gate Update module 170, the layer gates in the pre-trained model 150 are updated accordingly. Since the layer gates are updated according to the provided task and budget, the end-to-end inference process is budget / task-adaptive. When another request with a different task and / or budget is received at the model, the gates may be updated accordingly.

[0082] An exemplary model inference method based on the budget and task-adaptive inference pipeline is now described. The method may be performed by a computing system that is configured to execute pre-trained foundation models for processing inference requests from users. More particularly, the computing system may serve multiple requests from users with different input budgets and / or tasks in accordance with the method. In at least some implementations, an inference engine (such as the engine 100 of FIG. 1) implemented by a computing system may perform all or part of the operations of the method.

[0083] The computing system may load, in memory, a single pre-trained machine learning model. A plurality of requests from users to be processed using the machine learning model may be served by the computing system. Each request includes, at least, an input prompt, a first budget, and a task identifier of a first task provided by a user. The machine learning model may be used for processing multiple different requests, without having to load additional models in memory.

[0084] For each received request, the computing system may be configured to prune the machine learning model at inference time. In particular, for each layer of the machine learning model, the computing system may determine a layer gate value. This value indicates whether to use the layer for inference based on the first budget and the first task.

[0085] In at least some implementations, the computing system may determine the layer gate value for a layer of the machine learning model by performing lookup of a budget table stored in memory. The budget table stores, for each of a plurality of combinations of tasks and budgets, a vector representing layer gate values of all layers of the model. The layer gate values for the layers of the model may be represented, for example, by a Boolean vector in the budget table. In response to determining that the budget table contains a table entry corresponding to the combination of the first budget and the first task, the computing system retrieves layer gate values associated with the table entry. That is, if the budget table already contains layer gate values for a specific task / budget input, said values (e.g., a vector representing the values) are retrieved.

[0086] If, on the other, the computing system determines, upon performing lookup of the budget table, that the budget table does not contain a table entry corresponding to the first task / budget combination, the computing system is configured to execute a layer-wise pruning method to obtain layer gate value for the first task and the first budget. The obtained layer gate values may then be stored in the budget table in association with the first task / budget.

[0087] The computing system then utilizes the pruned machine learning model in accordance with the layer gate values to generate an output corresponding to the request. In particular, an output for the request is generated using only the layers that have a layer gate value equal to a predefined first value (e.g., gate value of 1) . The first value is a numerical value representing the layer gate being “on” .

[0088] Reference is now made to FIG. 5 which shows an inference-time pruning framework for pre-trained foundation models with multiple LoRA adapters. As will be described in greater detail, the inference pipeline of FIGS. 5 and 6 is budget and task-adaptive. The inference engine of FIG. 5 enables LoRA weights corresponding to different tasks to be added to the pruning engine 510. The layer gates in this framework are associated with the pre-trained layers and the corresponding layers in the LoRA adapters. The layer gates are initialized to 1, meaning that all of the pre-trained model’s weights and the corresponding LoRA weights are active in the beginning.

[0089] The Gate Update module 570 is configured to update the values of the layer gates according to the user-specified budget and task. Exemplary components of the Gate Update module 570 are shown in FIG. 7. The inputs to the Gate Update module 570 include task information (i.e., LoRA IDs) , budget, and the layer gate values corresponding to the pre-trained foundation model and the LoRA adapters. The Gate Update Module 570 outputs updated layer gate values. A budget table 580 of the Gate Update module 570 stores layer gate values for the layers of the pre-trained model 550 and corresponding LoRAs.

[0090] For each of a plurality of combinations of tasks (i.e., LoRA IDs) and budgets, the budget table 580 stores a corresponding Boolean vector that represents the layer gate values corresponding to the model’s layers. Otherwise, a layer-wise pruning method that is capable of pruning the pre-trained model 550 and the LoRA weights is used to obtain the values of layer gates for the given task and budget. A suitable off-the-shelf layer-wise pruning method that can prune the pre-trained model 550 and the LoRA weights may be used to determine which layers are pruned. By way of example, the LoRAprune method can be modified to apply layer-wise pruning to the pre-trained model 550 and the LoRA weights. The layer gate values are stored in the budget table 580, and then provided to the pruning engine 510.

[0091] Once the updated layer gate values are returned from Gate Update module 570, the layer gates in the pre-trained model 550 are updated accordingly. Since the layer gates that are controlling the pre-trained model's layers and the LoRA adapters are updated according to the provided task and budget, the end-to-end inference process is budget / task-adaptive. When another request with a different LoRA ID and / or budget is given, the layer gates may be updated accordingly to the new task and budget.

[0092] An exemplary model inference method based on the budget and task-adaptive inference pipeline with LoRA adapters is now described. The method may be performed by a computing system that is configured to execute pre-trained foundation models for processing inference requests from users. More particularly, the computing system may serve multiple requests from users with different input budgets and / or tasks in accordance with the method. In at least some implementations, an inference engine (such as the engine 500 of FIG. 5) implemented by a computing system may perform all or part of the operations of the method.

[0093] The computing system may load, in memory, a single pre-trained machine learning model. A plurality of requests from users to be processed using the machine learning model with a plurality of Low-Rank Adaptation (LoRA) adapters may be served by the computing system. Each request includes, at least, an input prompt, a first budget, and an identifier of a first LoRA adapter provided by a user. The machine learning model may be used for processing multiple different requests, without having to load additional models in memory.

[0094] For each received request, the computing system may be configured to prune the machine learning model at inference time. In particular, for each layer of the machine learning model, the computing system may determine a layer gate value. This value indicates whether to use the layer for inference based on the first budget and the first LoRA adapter.

[0095] In at least some implementations, the computing system may determine the layer gate value for a layer of the machine learning model by performing lookup of a budget table stored in memory. The budget table stores, for each of a plurality of combinations of LoRA identifiers and budgets, a vector representing layer gate values of all layers of the model. The layer gate values for the layers of the model may be represented, for example, by a Boolean vector in the budget table. In response to determining that the budget table contains a table entry corresponding to the combination of the first budget and the first LoRA adapter, the computing system retrieves layer gate values associated with the table entry. That is, if the budget table already contains layer gate values for a specific LoRA adapter / budget input, said values (e.g., a vector representing the values) are retrieved.

[0096] If, on the other, the computing system determines, upon performing lookup of the budget table, that the budget table does not contain a table entry corresponding to the first LoRA adapter / budget combination, the computing system is configured to execute a layer-wise pruning method to obtain layer gate value for the first LoRA adapter and the first budget. The obtained layer gate values may then be stored in the budget table in association with the first LoRA adapter / budget.

[0097] The computing system then utilizes the pruned machine learning model and corresponding LoRA adapters in accordance with the layer gate values to generate an output corresponding to the request. In particular, an output for the request is generated using only the layers that have a layer gate value equal to a predefined first value (e.g., gate value of 1) . The first value is a numerical value representing the layer gate being “on” .

[0098] Reference is now made to FIGS. 4 and 8 which illustrate examples of batching requests during inference using pre-trained foundation models. Batching for inference refers to processing multiple requests together as a single batch during the inference phase of a model. Instead of making predictions one at a time, the model processes a group of inputs simultaneously. This technique can offer several benefits, particularly in terms of efficiency and resource utilization.

[0099] The present application proposes a batching technique for budget-adaptive inference scenarios. Each of a plurality of requests to the model is associated with a respective input prompt, task, and budget. For each request, a Gate Update module (such as the module 170 of FIG. 2) is used to determine the values of the layer gates. The inference process starts from the first layer of the pre-trained model and proceeds through the subsequent layers. If, at a given layer, the layer gate values for multiple prompts are set to 1, the corresponding layers are run in batch simultaneously. That is, if two or more requests have corresponding layer gate values of 1 at a given layer, the layers corresponding to the requests are processed together as a single batch during inference. If the layer gate value is set to 1 for only one of the prompts, the corresponding layer is run without batching. If none of the layer gates is active in a layer, that layer is skipped altogether.

[0100] An example for batching three prompts is illustrated in FIG. 4. Given prompts, p1, p2 and p3, the corresponding tasks are t1, t2, and t3 and corresponding budgets are b1, b2 and b3, respectively. The gates 401 represent a layer gate value of 1 (i.e., set to ON) and the gates 402 represent a layer gate value of 0 (i.e., set to OFF) . In layer 1, all the layer gate values for the three prompts are set to 1. Therefore, layer 1 of p1, p2 and p3 will be processed in batch. For layer 2, layer gates corresponding to p2 and p3 (but not p1) are active, i.e., set to 1. Hence, layer 2 for these two prompts are processed in batch. In layer N-1, only the layer gate for p1 is 1. Therefore, this layer is run individually, and not in batch.

[0101] An exemplary model inference method that leverages the batching techniques of the present disclosure is now described. The method may be performed by a computing system that is configured to execute pre-trained foundation models for processing inference requests from users. More particularly, the computing system may serve multiple requests from users with different input budgets and / or tasks in accordance with the method. In at least some implementations, an inference engine (such as the engine 100 of FIG. 1) implemented by a computing system may perform all or part of the operations of the method.

[0102] The computing system may load, in memory, a single pre-trained machine learning model. A plurality of requests from users to be processed using the machine learning model may be served by the computing system. Each request includes, at least, an input prompt, a first budget, and a task identifier of a first task provided by a user. The machine learning model may be used for processing multiple different requests, without having to load additional models in memory.

[0103] For each layer of the machine learning model, the computing system identifies multiple requests that have a same layer gate value indicating that the layer is to be used for inference based on the respective budget and task. In at least some implementations, identifying the multiple requests may include first determining layer gate values for the plurality of user requests. The computing system may perform lookup of a budget table storing, for each of a plurality of combinations of tasks and budgets associated with user requests, a vector representing layer gate values of layers of the machine learning model. As described above, the layer gate values for the layers of the model may be represented by a Boolean vector in the budget table. In response to determining that the budget table contains table entries corresponding to combinations of budgets and tasks associated with multiple ones of the requests, the computing system retrieves the layer gate values associated with those table entries.

[0104] In response to identifying the multiple requests, the computing system executes the layer for the identified requests simultaneously in batch. However, if for a layer of the model, only a single request is identified having a layer gate value indicating that the layer is to be used for inference based on the respective budget and task, the computing system executes the layer for the identified request without batching.

[0105] The proposed batching method can be extended to scenarios where LoRA adapters are also available in the inference framework. For the multi-LoRA case, when the layer gates are activated for multiple prompts, the layer and the corresponding LoRA weights are processed in batch. An example for the proposed batching with multiple LoRA adapters is given in FIG. 8. The  is the LoRA layer corresponding to the ith task and the jth layer.

[0106] Joint learning of pruning gates and LoRA adapters can be achieved via gradient-based learning, through the use of Lagrange multipliers. Given a pre-trained language model (PLM) and a target layer sparsity level LT, the proposed solution outputs a learned binary pruning gate vector z*and LoRA weights θ*LORA. By jointly learning the pruning gate simultaneously with the LoRA adapters, the adapters will learn to compensate for performance lose incurred during the pruning.

[0107] The binary pruning gate zi is parametrized using the hard-concrete distribution as represented in:

[0108] The variable u is drawn from a uniform distribution, and β, ζ, γ are hyperparameters which determine the shape of the distribution. More specifically, U is the uniform distribution, β is a temperature parameter, s is a relaxed binary mask that conforms to the hard concrete distribution, and ζ, γ are the bounds of the hard concrete distribution. This distribution has support on [0, 1] , but concentrates most of its mass at {0, 1} , and importantly, is amenable to gradient-based learning. Standard hyperparameter values from the literature are used and thus do not require additional tuning by the user. In particular, the standard values of β=0.83, ζ=1.1, γ=-0.1 may be used.

[0109] Given this parameterization, the following minimax objective is optimized, where θ, λ are Lagrange multipliers:

[0110] The benefits of this approach are four-fold. First, the above equation allows the user to specify a desired sparsity level through the use of the LT parameter. Second, this approach allows for learning task-specific LoRA weights and pruning masks, which are low memory, and can be easily exchanged over-the-air, depending on the type of user query. Third, as previously mentioned, by learning LoRA weights concurrently with the pruning masks, LoRA weights learn to recover the performance loss incurred during model pruning, and conversely, the masks are learned to prune the layers that are least amenable to said compensation. Finally, this approach can be easily extended to allow for width pruning (as well as layer-wise pruning) , to reduce memory requirements of large models.

[0111] FIG. 9 illustrates an example computing system 900, which may be used to implement examples of the present disclosure. In particular, the computing system 900 may be configured to perform budget-adaptive inference for pre-trained foundation models in accordance with various methods described herein.

[0112] The computing system 900 includes at least one processing unit, such as a processor 902, and at least one physical memory 904. The processor 902 may be, for example, a central processing unit, a microprocessor, a digital signal processor, an application-specific integrated circuit (ASIC) , a field-programmable gate array (FPGA) , a dedicated logic circuitry, a dedicated artificial intelligence processor unit, a graphics processing unit (GPU) , a tensor processing unit (TPU) , a neural processing unit (NPU) , a hardware accelerator, or combinations thereof. The memory 904 may include a volatile or non-volatile memory (e.g., a flash memory, a random-access memory (RAM) , and / or a read-only memory (ROM) ) . The memory 904 may store instructions for execution by the processor 902, to the computing system 900 to carry out examples of the methods, functionalities, systems and modules disclosed herein.

[0113] The computing system 900 may also include at least one network interface 906 for wired and / or wireless communications with an external system and / or network (e.g., an intranet, the Internet, a P2P network, a WAN and / or a LAN) . A network interface may enable the computing system 900 to carry out communications (e.g., wireless communications) with systems external to the computing system 900, such as a language model residing on a remote system.

[0114] The computing system 900 may optionally include at least one input / output (I / O) interface 908, which may interface with optional input device (s) 910 and / or optional output device (s) 912. Input device (s) 910 may include, for example, buttons, a microphone, a touchscreen, a keyboard, etc. Output device (s) 912 may include, for example, a display, a speaker, etc. In this example, optional input device (s) 910 and optional output device (s) 912 are shown external to the computing system 900. In other examples, one or more of the input devices 910, and / or output device (s) 912 may be an internal component of the computing system 900.

[0115] Herein, use of language such as “at least one of X, Y, and Z, ” “at least one of X, Y, or Z, ” “at least one or more of X, Y, and Z, ” “at least one or more of X, Y, and / or Z, ” or “at least one of X, Y, and / or Z, ” is intended to be inclusive of both a single item (e.g., just X, or just Y, or just Z) and multiple items (e.g., {X and Y} , {X and Z} , {Y and Z} , or {X, Y, and Z} ) . The phrase “at least one of” and similar phrases are not intended to convey a requirement that each possible item must be present, although each possible item may be present.

[0116] In some implementations, the methods disclosed herein may be implemented as computer-executable instructions stored in one or more non-transitory computer-readable storage devices (in the form of software, firmware, or a combination thereof) such that, the instructions, when executed, may cause one or more physical components such as one or more circuits to perform the methods disclosed herein.

[0117] For example, in some implementations, an apparatus comprising one or more processors functionally connected to one or more non-transitory computer-readable storage devices or media may be used to perform the methods disclosed herein, wherein the one or more non-transitory computer-readable storage devices or media store the computer-executable instructions of the methods disclosed herein, and the one or more processors may read the computer-executable instructions from the one or more non-transitory computer-readable storage devices or media, and executes the instructions to perform the methods disclosed herein.

[0118] In some implementations, an apparatus may not have any processors or computer-readable storage devices or media. Rather, the apparatus may comprise any other suitable physical or virtual components for implementing the methods disclosed herein.

[0119] In some implementations, the computer-executable instructions that implement the methods disclosed herein may be one or more computer programs, one or more program products, or a combination thereof.

[0120] In some implementations, the methods disclosed herein may be implemented as one or more circuits, one or more components, one or more units, one or more modules, one or more integrated-circuit (IC) chips, one or more chipsets, one or more devices, one or more apparatuses, one or more systems, and / or the like.

[0121] The one or more circuits, one or more components, one or more units, one or more modules, one or more IC chips, one or more chipsets, one or more devices, one or more apparatuses, or one or more systems may be physical, virtual, or a combination thereof. Herein, the term “virtual” (such as a “virtual apparatus” ) refers to a circuit, component, unit, module, chipset, device, apparatus, system, or the like that is simulated or emulated or otherwise formed using suitable software or firmware such that it appears as if it is “real” or physical) .

[0122] The present disclosure encompasses various embodiments, including not only method embodiments, but also other embodiments such as apparatus embodiments and embodiments related to non-transitory computer readable storage media. Embodiments may incorporate, individually or in combinations, the features disclosed herein.

[0123] Although this disclosure refers to illustrative embodiments, this is not intended to be construed in a limiting sense. Various modifications and combinations of the illustrative embodiments, as well as other embodiments of the disclosure, will be apparent to persons skilled in the art upon reference to the description.

[0124] Features disclosed herein in the context of any particular embodiments may also or instead be implemented in other embodiments. Method embodiments, for example, may also or instead be implemented in apparatus, system, and / or computer program product embodiments. In addition, although embodiments are described primarily in the context of methods and apparatus, other implementations are also contemplated, as instructions stored on one or more non-transitory computer-readable media, for example. Such media could store programming or instructions to perform any of various methods consistent with the present disclosure.

[0125] Those skilled in the art will appreciate that the above-described embodiments and / or features thereof may be customized, separated, and / or combined as needed or desired. Moreover, although embodiments have been described above with reference to the accompanying drawings, those of skill in the art will appreciate that variations and modifications may be made without departing from the scope thereof as defined by the appended claims.

Claims

1.A computer-implemented method, comprising:loading, in memory, a single pre-trained machine learning model;receiving a plurality of requests to be processed using the machine learning model, each request including an input prompt, a first budget, and a task identifier of a first task; andfor each received request:pruning the machine learning model at inference time, wherein the pruning comprises determining, for each layer of the machine learning model, a layer gate value indicating whether to use the layer for inference based on the first budget and the first task; andutilizing the pruned machine learning model in accordance with the layer gate values to generate an output corresponding to the request.2.The method of claim 1, wherein determining the layer gate value comprises:performing lookup of a budget table storing, for each of a plurality of combinations of tasks and budgets, a vector representing layer gate values of all layers of the machine learning model;determining that the budget table contains a first table entry corresponding to the combination of the first budget and the first task; andretrieving layer gate values associated with the first table entry.3.The method of claim 2, wherein the layer gate values for the layers of the machine learning model are represented by a Boolean vector in the budget table.4.The method of claim 1, wherein utilizing the pruned machine learning model comprises generating the output using only the layers that have a layer gate value equal to a predefined first value.5.The method of claim 1, wherein determining the layer gate value comprises:performing lookup of a budget table storing, for each of a plurality of combinations of tasks and budgets, a vector representing layer gate values of all layers of the machine learning model;determining that the budget table does not contain a table entry corresponding to the combination of the first budget and the first task; andexecuting a layer-wise pruning method to obtain layer gate values for the first budget and the first task.6.The method of claim 5, further comprising storing the obtained layer gate values in the budget table.7.A computer-implemented method, comprising:loading, in memory, a single pre-trained machine learning model;receiving a plurality of requests to process using the machine learning model with a plurality of Low-Rank Adaptation (LoRA) adapters, each request including an input prompt, a first budget, and an identifier of a first LoRA adapter; andfor each received request:pruning the machine learning model at inference time, wherein the pruning comprises determining, for each layer of the machine learning model, a layer gate value indicating whether to use the layer for inference based on the first budget and the first LoRA adapter; andutilizing the pruned machine learning model and corresponding LoRA adapters in accordance with the layer gate values to generate an output corresponding to the request.8.The method of claim 7, wherein determining the layer gate value comprises:performing lookup of a budget table storing, for each of a plurality of combinations of LoRA identifiers and budgets, a vector representing layer gate values of all layers of the machine learning model;determining that the budget table contains a first table entry corresponding to the combination of the first budget and the first LoRA adapter; andretrieving layer gate values associated with the first table entry.9.The method of claim 8, wherein the layer gate values for the layers of the machine learning model are represented by a Boolean vector in the budget table.10.The method of claim 7, wherein utilizing the pruned machine learning model comprises generating the output using only the layers that have a layer gate value equal to a predefined first value.11.The method of claim 7, wherein determining the layer gate value comprises:performing lookup of a budget table storing, for each of a plurality of combinations of LoRA identifiers and budgets, a vector representing layer gate values of all layers of the machine learning model;determining that the budget table does not contain a table entry corresponding to the combination of the first budget and the first LoRA adapter; andexecuting a layer-wise pruning method to obtain layer gate values for the first budget and the first LoRA adapter.12.The method of claim 11, further comprising storing the obtained layer gate values in the budget table.13.A computer-implemented method, comprising:receiving a plurality of requests to process using a pre-trained machine learning model, each request including a respective input prompt, a budget, and a task identifier of a task;for each layer of the machine learning model:identifying multiple requests that have a same layer gate value indicating that the layer is to be used for inference based on the respective budget and the task; andin response to identifying the multiple requests, executing the layer for the identified requests simultaneously in batch.14.The method of claim 13, wherein identifying the multiple requests comprises:performing lookup of a budget table storing, for each of a plurality of combinations of tasks and budgets, a vector representing layer gate values of all layers of the machine learning model;determining that the budget table contains table entries corresponding to combinations of budgets and tasks associated with multiple ones of the requests; andretrieving layer gate values associated with the table entries.15.The method of claim 13, further comprising, for each layer of the machine learning model:identifying only a single request having a layer gate value indicating that the layer is to be used for inference based on the respective budget and the task; andin response to identifying only the single request, executing the layer for the identified request without batching.16.The method of claim 14, wherein the layer gate values for the layers of the machine learning model are represented by a Boolean vector in the budget table.17.A computer-implemented method for gradient-based learning of binary pruning gates, the method comprising:parametrizing a binary pruning mask to model hard concrete distributions; andoptimizing the minimax objective given by:where θ, λ are Lagrange multipliers.18.The method of claim 17, wherein the hard concrete distribution is defined as: where U is the uniform distribution, β is a temperature parameter, s is a relaxed binary mask that conforms to the hard concrete distribution, and ζ, γ are the bounds of the hard concrete distribution.19.A computing system, comprising:a processor;memory coupled to the processor, the memory storing computer-executable instructions that, when executed by the processor, configure the processor to perform any of the methods of claims 1 to 18.20.A non-transitory computer-readable medium storing instructions that, when executed by a processor, configure the processor to perform any of the methods of claims 1 to 18.

Images