CHIPLET ARCHITECTURE FOR INFERENCE, FINE-TUNING TRAINING, AND TRANSFER LEARNING

A method for training and fine-tuning an artificial intelligence model is disclosed. In one embodiment, such a method distributes, across multiple chiplets of a package, functionality associated with a deep neural network. The method implements, within a first set of chiplets, frozen layers of the deep neural network. By contrast, the method implements, within a second set of chiplets, trainable layers of the deep neural network. The number of chiplets in the second set may be smaller than the number of chiplets in the first set and may consist of a single chiplet in some embodiments. In certain embodiments, the second set of chiplets has one or more of additional memory capacity and additional processing capacity compared to the first set of chiplets in order to train and fine tune the trainable layers. A corresponding apparatus is also disclosed.

BACKGROUND

Field of the Invention

This invention relates generally to artificial intelligence, and more particularly to systems and methods for more efficiently training and fine-tuning artificial intelligence models.

Background of the Invention

Training artificial intelligence models and deploying them for inference typically have different requirements in terms of computing power and memory. For example, training an artificial intelligence model typically involves feeding large amounts of labeled data into the model and adjusting the model's parameters iteratively to optimize its performance. This training is computationally intensive and requires significant computing and memory resources especially for complex models and large datasets. Training typically occurs offline, with the goal being to create an optimized model that captures patterns and relationships in the training data.

On the other hand, inference deployment involves using the trained model to make predictions or decisions on new, unseen data. The computing power required for inference is generally lower than that needed for training. This is due to the reduced complexity of performing inference, where there is no need for gradient computations, backpropagation, or parameter updates required for training. This may also be due to the fact that inference is often implemented on dedicated hardware accelerators, such as GPUs or specialized chips like TPUs (Tensor Processing Units) which can enhance computational efficiency. Inference is often performed on edge devices or in cloud environments, with the goal being to achieve real-time or near-real-time predictions based on the trained model.

SUMMARY

The invention has been developed in response to the present state of the art and, in particular, in response to the problems and needs in the art that have not yet been fully solved by currently available systems and methods. Accordingly, systems and methods have been developed for training and fine-tuning artificial intelligence models. The features and advantages of the invention will become more fully apparent from the following description and appended claims, or may be learned by practice of the invention as set forth hereinafter.

Consistent with the foregoing, a method for training and fine-tuning an artificial intelligence model is disclosed. In one embodiment, such a method distributes, across multiple chiplets of a package, functionality associated with a deep neural network. The method implements, within a first set of chiplets, frozen layers of the deep neural network. By contrast, the method implements, within a second set of chiplets, trainable layers of the deep neural network. The number of chiplets in the second set may be smaller than the number of chiplets in the first set and may consist of a single chiplet in some embodiments. In certain embodiments, the second set of chiplets has one or more of additional memory capacity and additional processing capacity compared to the first set of chiplets in order to train and fine tune the trainable layers.

A corresponding apparatus is also disclosed and claimed herein.

DETAILED DESCRIPTION

Computing environment100contains an example of an environment for the execution of at least some of the computer code involved in performing the inventive methods, such as code150(i.e., a “fine tuning module150”) for more efficiently training and fine-tuning artificial intelligence models. In addition to block150, computing environment100includes, for example, computer101, wide area network (WAN)102, end user device (EUD)103, remote server104, public cloud105, and private cloud106. In this embodiment, computer101includes processor set110(including processing circuitry120and cache121), communication fabric111, volatile memory112, persistent storage113(including operating system122and block150, as identified above), peripheral device set114(including user interface (UI) device set123, storage124, and Internet of Things (IoT) sensor set125), and network module115. Remote server104includes remote database130. Public cloud105includes gateway140, cloud orchestration module141, host physical machine set142, virtual machine set143, and container set144.

Referring toFIG.2, as previously mentioned, training artificial intelligence models and deploying them for inference typically have different requirements in terms of computing power and memory. For example, training an artificial intelligence model typically involves feeding large amounts of labeled data into the model and adjusting the model's parameters iteratively to optimize its performance. This training is computationally intensive and requires significant computing and memory resources especially for complex models and large datasets. Training typically occurs offline, with the goal being to create an optimized model that captures patterns and relationships in the training data.

On the other hand, inference deployment involves using the trained model to make predictions or decisions on new, unseen data. The computing power required for inference is generally lower than that needed for training. This is due to the reduced complexity of performing inference, where there is no need for gradient computations, backpropagation, or parameter updates required for training. This may also be due to the fact that inference is often implemented on dedicated hardware accelerators, such as GPUs or specialized chips like TPUs (Tensor Processing Units), which can enhance computational efficiency. Inference is often performed on edge devices or in cloud environments, with the goal being to achieve real-time or near-real-time predictions based on the trained model.

Nevertheless, even after a deployment, a trained artificial intelligence model may require further fine tuning and/or adjustments. This fine-tuning may need to be performed on a system with substantial computing and memory requirements, similar to the initial training, as opposed to on edge devices or cloud environments with lower computing resources. The fine-tuned model may then be re-deployed on an edge device or cloud environment to perform inference operations. This process may need to be repeated each time the artificial intelligence model requires further tuning and/or adjustments. This recurring exchange between training and deployment systems can be both cumbersome and time-consuming. In some cases, the training and deployment systems may not even be owned or operated by the same entities, further increasing the difficulty in transitioning between the two.

To address the issues described above, in certain embodiments, instead of repeatedly transitioning between a higher performance training system and a lower performance deployment system, a deployment system may be provided that enables some fine-tuning and/or adjustments to an artificial intelligence model. In general, as shown inFIG.2, for a pre-trained base deep neural network200that includes multiple hidden layers202between an input layer204and output layer206, fine-tuning and transfer learning (i.e., training or fine tuning the base deep neural network200to perform a new task) are generally performed on a relatively small number of layers202bnear the output layer206. That is, only the weights of a last number of layers202btypically need to be modified to fine tune or optimize a deep neural network200. By contrast, the weights of the remaining layers202agenerally remain frozen. For the purposes of this disclosure, the layers202bmay be referred to as “trainable layers202b” and the layers202amay be referred to as “frozen layers202a.” The deep neural network200shown inFIG.2is shown in very simplified form for explanation purposes and may include more or fewer layers202and/or neurons208within the layers202than those that are illustrated.

Referring toFIG.3, as previously mentioned, inference operations associated with an artificial intelligence model may be implemented on dedicated hardware such as GPUs or specialized chips like TPUs to enhance computational efficiency. However, conventional hardware devices are typically not designed to efficiently train or fine tune the artificial intelligence model.FIG.3shows one embodiment of a hardware component (i.e., a package300comprising multiple chiplets302) in accordance with the invention that may be designed to perform inference operations associated with an artificial intelligence model. Beneficially, unlike conventional hardware, the package300may also be configured to efficiently perform some fine-tuning and training of the artificial intelligence model deployed thereon.

As shown inFIG.3, in certain embodiments, a package300in accordance with the invention may include multiple chiplets302in communication with one another. The illustrated chiplets302are arranged in series although in other embodiments the chiplets302may be arranged in parallel or even a configuration comprising both series and parallel connections. An artificial intelligence model comprising a deep neural network200is distributed across the chiplets302. The package300implements, within a first set of chiplets302a-f, frozen layers202aof the deep neural network200. Similarly, the package300implements, within a second set chiplets302g, trainable layers202bof the deep neural network200. In certain embodiments, the second set of chiplets302gis a single chiplet302galthough in other embodiments the second set of chiplets302gmay include multiple chiplets302g. It should also be recognized that although the chiplets302are shown in a single package300the chiplets302may also be distributed across several packages300. For example, chiplets302with frozen layers202acould be on a separate package300from chiplets302with trainable layers202b.

As shown inFIG.3, each chiplet302in the chain may implement some number of layers202(or portions of layers202) of the deep neural network200. The output of the layers202implemented by a chiplet302may be passed to a next chiplet302in the chain, where it will be input to the layers202implemented by that chiplet302. The full set of chiplets302may, in certain embodiments, implement all layers202of the deep neural network200. As was previously mentioned, for a pre-trained base deep neural network200that includes multiple hidden layers202between an input layer204and output layer206, fine-tuning and transfer learning are generally performed on a relatively small number of layers202bnear the output layer206. Thus, in certain embodiments, only a single or smaller number of chiplets302gin the chain, namely those at the end of the chain, may host layers202that are used for fine-tuning and/or transfer learning. The weights or other parameters for the layers202bimplemented by these chiplet(s)302gmay be adjustable. By contrast, weights or other parameters for the layers202aimplemented by remaining chiplets302a-fmay be frozen. Thus, only chiplet(s)302gmay be associated with fine-tuning and/or training the deep neural network200.

Furthermore, although the illustrated trainable chiplets302gare located at or near an end of the chain of chiplets302, in other embodiments trainable chiplets302gmay be placed at other locations within the deep neural network200to perform other types of fine-tuning. This is because different methods of fine tuning may modify different portions of the deep neural network200. Within fine tuning, some current techniques include prompt tuning, head tuning, and LORA (Low-Rank Adaptation of Large Language Models). Prompt tuning modifies a first layer202of the deep neural network200. Head tuning, by contrast, modifies the last layer202. LORA may add trainable parameters at every layer202or transformer block. Thus, in certain embodiments trainable chiplets302gmay be provided at various and even several locations within the deep neural network200, anywhere between the input layer204and the output layer206.

Referring toFIG.4, because fine-tuning and/or training requires more computing resources (e.g., processing power, memory, etc.) than simply performing inference operations, it follows that the chiplet(s)302ginvolved in fine-tuning and/or training the deep neural network200may require more computing resources than other chiplets302a-fof the package300. Thus, in certain embodiments, the chiplet(s)302gmay be designed to include sufficient on-chip memory (or high bandwidth and/or high capacity memory in close proximity to the chiplet(s)302g) to accommodate training of the changeable layers202b, including optimizer states and activations needed for backpropagation. For a set of changeable weights requiring memory capacity wc, the memory must be designed to accommodate αwc, where α is a factor accounting for the training and is typically on the order of three to eight depending on the optimizer. The chiplet(s)302gmay also be designed with sufficient computing performance such that forward/backward propagation of the changeable layers202bcan be performed in a time similar to the forward propagation time of the fixed layers202ain the other chiplets302a-f. This may be done to avoid a bottleneck when employing pipeline parallelism. Because changeable weights may be implemented on a single chiplet302g, demanding bandwidth requirements for training may be accommodated by high on-chip bandwidth.

FIG.4shows a conceptual design wherein 3D stacked memory400of capacity c is coupled to a chiplet302g(e.g., an artificial intelligence accelerator) to provide sufficient on-chip memory to accommodate fine-tuning and training of the changeable layers202b. Storing weight parameters on chip reduces or eliminates off-chip memory accesses, which may be costly in energy and latency. The chiplets302a-fhosting the frozen layers202amay, in certain embodiments, not include this additional memory400since no fine-tuning or training is performed on these devices, but may otherwise have an architecture that is the same as or similar to the chiplet302g. These chiplets302a-fmay also, in certain embodiments, store weight parameters on chip to reduce or eliminate off-chip memory accesses.

FIG.5is a high-level block diagram showing how the chiplets302on a package300perform inference operations. When in inference mode, a first chiplet302ain the chain may receive inference inputs500and propagate these inference inputs500through the fixed layers202aimplemented by the chiplet302a. The output of the first chiplet302amay then be propagated in a forward direction through the layers202of the other chiplets302in the chain to produce inference predictions502at an output thereof. Each chiplet302in the chain may have enough scratchpad memory to store the activations that are generated within the chiplet302. The link bandwidth between the chiplets302may be designed so that communication latency is small compared to the compute time within the chiplets302.

FIG.6is a high-level block diagram showing a pipeline diagram600for fine-tuning the deep neural network200. Each box of the pipeline diagram600represents a microbatch of data processed by one of the chiplets302(four chiplets302in this example) of the package300when performing fine-tune training on the deep neural network200. The bottom row of boxes represents the microbatches processed by chiplet302a, the second to the bottom row of boxes represents the microbatches processed by chiplet302b, the third from the bottom row of boxes represents the microbatches processed by chiplet302c, and the top row of boxes represents the microbatches processed by chiplet302g(i.e., the trainable chiplet302g). The pipeline diagram600(as well as the pipeline diagram700shown inFIG.7) use the following notation: Fmodel,μbatch=Forward Propagation; Bmodel,μbatch=Backward Propagation; FBmodel,μbatch=Forward+Backward Propagation; and Umodel=Update. The microbatch size may be selected to be compatible with the batch size when performing inference processing so that training the deep neural network200does not impose additional capacity requirements.

As shown inFIG.6, because weights are only modified in the trainable chiplet302g, data only needs to propagate in the forward direction starting with chiplet302aand ending with chiplet302g. Forward and backward propagation would only occur in the trainable chiplet302gin order to update the weights in the trainable chiplet302g. By contrast if all the chiplets302were trainable, both forward and backward propagation would be needed through all the chiplets302, as show inFIG.7in order to update the weights in all of the chiplets302. Thus, using only a single trainable chiplet302gwith the other chiplets302having fixed weights significantly reduces processing requirements needed to train the deep neural network200compared to having changeable weights in all the chiplets302.

Referring toFIG.8, an alternative embodiment of a technique for fine-tuning or training the trainable chiplet302gis illustrated. In this embodiment, during inference operations, activations from the frozen chiplets302a-c(i.e., frozen layers202a) are stored in an external memory800or storage device800either on or off chip. This may be performed for a significant period of time. In order to fine-tune the deep neural network200and/or facilitate transfer learning, these activations may be read from the external memory800and processed on the trainable chiplet302g, thereby eliminating the need to regenerate the activations on the chiplets302a-cduring the training process. Only forward and backward propagation is needed on the trainable chiplet302gin order to train or fine tune the deep neural network200as indicated by the pipeline diagram802ofFIG.8.

Referring toFIG.9, as previously mentioned, in certain embodiments a package300may include more than one trainable chiplet302g, such as to accommodate a greater memory requirement. In certain cases, chiplets302gcontaining trainable layers202bmay have higher bandwidth requirements in order to satisfy training requirements. This can be accomplished by way of greater serialization and/or advanced packaging (e.g., a bridge between the trainable chiplets302g).FIG.9shows a plurality of chiplets302with the same architecture that are arranged in a way to satisfy high speed communication requirements between trainable chiplets302g, while allowing lower speed communication between frozen chiplets302, which may be advantageous to reduce power consumption or provide other benefits. More specifically,FIG.9shows a plurality of chiplets302with a combination of higher and lower speed interfaces, where the higher speed interfaces of the trainable chiplets302gare aligned and connected to enable higher speed communication therebetween.

Referring toFIG.10, as previously discussed in association withFIG.8, in certain embodiments, activations from the frozen chiplets302a-c(i.e., frozen layers202a) may be stored in an external memory800or storage device800on or off chip over a significant period of time. These activations may then be read from the external memory800and processed on the trainable chiplet302gin order to fine-tune the deep neural network200and/or facilitate transfer learning.

In an alternative embodiment, instead of processing all of the activations on the trainable chiplet302g, the processing may be spread across all (or at least more) of the chiplets302on the package300to increase the efficiency of the training process. In order to accomplish this, the frozen chiplets302a-cmay be temporarily reprogrammed to assist in training the deep neural network200, which as mentioned above is a computationally intensive task. This may provide more “workers” in performing the training and fine-tuning process.

For example, in a first step, weights and other parameters associated with the frozen layers202aon the frozen chiplets302may be saved to external memory, which may include on-chip or off-chip memory. The activations that were stored in the external memory800(as was described in association withFIG.8) may then be divided between the chiplets302(labelled as activation sub-components Act′(a), Act′(b), Act′(c), and Act′(d)). In certain embodiments these activation sub-components may correspond to different inference tasks, such as sentiment analysis, question and answer, instruction following, image classification, image segmentation, or the like.

In a second step, weights (labelled as W(a), W(b), W(c), and W(d)) may be copied to each of the chiplets302to perform the fine-tuning task. In a third step, fine-tuning training tasks may be concurrently run on each the chiplets302to adjust and optimize the weights (labelled as W′(a), W′(b), W′(c), and W′(d)). These optimized weights may then be saved to external memory. The chiplets302a-cmay then be reprogrammed again to re-implement the frozen layers202ausing the weights and other parameters that were previously saved to external memory. The weights (i.e., W′(a), W′(b), W′(c), and W′(d)) may be copied into the trainable chiplet302gto finish the fine-tuning of the deep neural network200.

The flowcharts and block diagrams in the Figures illustrate the architecture, functionality, and operation of possible implementations of systems, methods, and computer program products according to various embodiments of the present invention. In this regard, each block in the flowcharts or block diagrams may represent a module, segment, or portion of code, which comprises one or more executable instructions for implementing the specified logical function(s). It should also be noted that, in some alternative implementations, the functions noted in the block may occur out of the order noted in the Figures. For example, two blocks shown in succession may, in fact, be executed substantially concurrently, or the blocks may sometimes be executed in the reverse order, depending upon the functionality involved. Other implementations may not require all of the disclosed steps to achieve the desired functionality. It will also be noted that each block of the block diagrams and/or flowchart illustrations, and combinations of blocks in the block diagrams and/or flowchart illustrations, may be implemented by special purpose hardware-based systems that perform the specified functions or acts, or combinations of special purpose hardware and computer instructions.