Infrastructure for fine-tuning neural networks through reinforcement learning

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training a neural network having multiple weights. In one aspect, at each of multiple training iterations, a method comprises loading first model data including current values of the weights into memory of multiple hardware devices in accordance with a first scheme, obtaining multiple training inputs, processing the training inputs using the neural network in accordance with the current values of the weights to sample one or more training outputs, generating a respective reward score for each of the training outputs, loading second model data including the current values of the weights and optimizer data for an optimizer into the memory of the hardware devices in accordance with a second scheme, and training the neural network through reinforcement learning using the reward scores for the training outputs to update the current values of the weights.

Description

[0001] Atorney Docket No. 56113-0618WO1

[0002] INFRASTRUCTURE FOR FINE-TUNING NEURAE NETWORKS THROUGH

[0003] REINFORCEMENT EEARNING

[0004] BACKGROUND

[0005] This specification relates to processing data using machine learning models.

[0006] Machine learning models receive an input and generate an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate the output based on the received input and on values of the parameters of the model.

[0007] Some machine learning models are deep models that employ multiple layers of models to generate an output for a received input. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers that each apply a non-linear transformation to a received input to generate an output.

[0008] SUMMARY

[0009] This specification describes a system implemented as computer programs on one or more computers in one or more locations for training a generative neural network through reinforcement learning. For example, the system can ‘'fine-tune” a pre-trained generative neural network through reinforcement learning using a reward function, e.g., a reward model learned from user preferences.

[0010] In particular, to perform this training, the system performs multiple training iterations. Each training iteration includes a sampling phase and a training phase.

[0011] For the sampling phase, the system loads first model data including current values of the multiple weights into memory' of multiple hardware devices in accordance with a first scheme. After the first model data is loaded in the memory’ of the multiple devices in accordance with the first scheme, the system obtains, by the multiple devices, multiple training inputs, the system processes, using the multiple devices, each of the multiple training inputs using the neural network in accordance with the current values of the multiple weights to sample one or more training outputs for each of the multiple training inputs, and the system generates, by the multiple devices, a respective reward score for each of the training outputs.

[0012] For the training phase, the system loads second model data including current values of the multiple w eights and optimizer data for an optimizer into the memory- of multiple hardware devices in accordance with a second scheme. After the second model data is Atorney Docket No. 56113-0618WO1 loaded in the memory' of the multiple devices in accordance with the second scheme, the system trains the neural network through reinforcement learning using the reward scores for the training outputs to update the current values of the plurality of weights.

[0013] Once the system has trained the neural network, the system can provide the updated values of the weights as the output of the training process..

[0014] Particular embodiments of the subject matter described in this specification can be implemented as to realize one or more of the following advantages.

[0015] Conventional approaches for fine-tuning a neural network (e.g., a generative neural network) through reinforcement learning perform the fine-tuning using a static model sharding scheme that stays the same throughout training.

[0016] A model sharding scheme defines how instances of the neural network are partitioned across sets of hardware devices that are available for the training of the neural network.

[0017] For example, at each training step of multiple training steps, a system can use the same particular sharding scheme for both the ‘"sampling” and "‘training” phases that are carried at the training step.

[0018] However, this fails to efficiently make use of the resources available for training. In particular, using the same sharding strategy' for both the “sampling” and “training” phases that are performed at each training step does not efficiently make use of the resources provided by the set of hardware devices.

[0019] In contrast, the described system can more efficiently train the neural network by performing sampling and training in accordance with different sharding strategies on multiple hardware devices. For example, the system can implement a first sharding scheme for sampling and a second sharding scheme for training based on one or more factors. In particular, for sampling, the first sharding scheme can partition the current model weights across several hardware devices to maximize the speed of generating sampled training outputs.

[0020] For training, the second sharding scheme can partition the model weights across hardware devices in a manner that results in fewer instances of the neural network, but increased data parallelism. That is, the system can more efficiently train the neural network through reinforcement learning using the reward scores for the training outputs to update the values of the weights for the neural network.

[0021] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, Atorney Docket No. 56113-0618WO1 and advantages of the subject matter will become apparent from the description, the drawings, and the claims.

[0022] BRIEF DESCRIPTION OF THE DRAWINGS

[0023] FIG. 1 shows an example neural network training system.

[0024] FIG. 2 shows a block diagram for the example neural network training system during sampling and training.

[0025] FIG. 3 is a flow diagram of an example process for fine-tuning weights of a neural network by performing a sampling phase and a training phase according to respective strategies.

[0026] Like reference numbers and designations in the various drawings indicate like elements.

[0027] DETAILED DESCRIPTION

[0028] FIG. 1 shows an example neural network training system 100. The neural network training system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

[0029] The neural network training system 100 is a system that obtains data specifying an initial, pre-trained generative neural network 110 and further trains (“fine-tunes”) the pretrained neural network 110 using one or more reward functions 108 to generate fine-tuned weight values 126 for the neural network 110. That is, the system 100 fine-tunes current weight values 120 of the neural network 110 to generate the fine-tuned weight values 126.

[0030] The neural network training system 100 can obtain data specifying the current weight values 120 of the pre-trained neural network 110 from an off-chip memory 104. The off-chip memory' 104 can be any appropriate memory' that is not located on any of the hardw are devices (e.g., not in the on-device memory of the hardware devices). For example, the off-chip memory 104 can be random access memory (RAM) allocated on one or more storage devices. The obtained current weight values 120 can then be loaded into the on-device memory of the plurality of hardware devices for use in training as described below.

[0031] The system 100 performs the training of the neural network 110 over multiple training steps. Atorney Docket No. 56113-0618WO1

[0032] At each training step, the system 100 performs “sampling’' of training inputs 118. That is, the system performs sampling by using the neural network 110 to generate one or more output examples (e.g., sampled training outputs 122) for each of the multiple context inputs (e.g., sampled training inputs 118).

[0033] The system then determines reward scores 124 of the sampled training outputs 122 using one or more reward functions 108.

[0034] The system 100 then performs “training” in order to train the neural network 110 using the sampled training outputs 122 and the reward scores 124. As part of the training, the system 100 uses the reward scores 124 to train the neural network 110 through reinforcement learning, as described in further detail below with reference to FIG. 2.

[0035] Conventionally, the system can use a particular weight partitioning (e.g., sharding) scheme for both sampling and training, i.e., use the same sharding scheme for both the sampling and training phases that are carried at each training step.

[0036] However, performing sampling and training using the same sharding scheme can result in inefficiencies in processing, as the preferred sharding strategy (e.g., sharding scheme) for sampling may not be ideal or efficient for training. In contrast, the reinforcement learning system 102 can more efficiently train the neural network 110 by performing sampling and training using different sharding strategies on multiple hardware devices using reward functions, as described in further detail below with reference to FIG. 2.

[0037] The neural network 1 10 is referred to as a “generative” neural network because the neural network 110 generates a new output example conditioned on a context input, i.e., instead of discriminating between existing output examples.

[0038] For example, the neural network 110 can be a language model neural network, i.e., an auto-regressive neural network that generates output sequences of tokens from a vocabulary, e.g., conditioned on a context sequence.

[0039] As a particular example, the tokens in the vocabulary can be text tokens, such that the language model neural network maps a context sequence of text tokens to an output sequence of text tokens. As another particular example, the language model neural network can be a multi-modal neural network, with the tokens in the context sequence including tokens representing another data modality, e.g., images, videos, or audio, instead of or in addition to text tokens.

[0040] The neural network 110 is referred to as an auto-regressive neural network because the neural network 110 auto-regressively generates an output sequence of tokens by Atorney Docket No. 56113-0618WO1 generating each particular token in the output sequence conditioned on a current input sequence that includes any tokens that precede the particular text token in the output sequence, i.e., the tokens that have for already been generated for any previous positions in the output sequence that precede the particular position of the particular token, and a context input that provides context for the output sequence. For example, the current input sequence when generating a token at any given position in the output sequence can include the context sequence and the tokens at any preceding positions that precede the given position in the output sequence. As a particular example, the current input sequence can include the context sequence followed by the tokens at any preceding positions that precede the given position in the output sequence. Optionally, the context and the current output sequence can be separated by one or more predetermined tokens within the current input sequence.

[0041] More specifically, to generate a particular token at a particular position within an output sequence, the neural network 110 can process the current input sequence to generate a score distribution, e.g., a probability distribution, that assigns a respective score, e.g., a respective probability, to each token in the vocabulary of tokens. The neural network 110 can then select, as the particular token, a token from the vocabulary using the score distribution. For example, the neural network 110 can greedily select the highest-scoring token or can sample, e.g., using nucleus sampling or another sampling technique, a token from the distribution.

[0042] As a particular example, the language model neural network 1 10 can be an autoregressive Transformer-based neural network that includes (i) a plurality of attention blocks that each apply a self-attention operation and (ii) an output subnetwork that processes an output of the last attention block to generate the score distribution.

[0043] The neural network 1 10 can have any of a variety of Transformer-based neural network architectures. Generally, however, the Transformer-based neural network includes a sequence of attention blocks (a block that applies an attention mechanism over a block input to generate a block output), and, during the processing of a given input sequence, each attention block in the sequence receives a respective input hidden state for each input token in the given input sequence. The attention block then updates each of the hidden states at least in part by applying self-attention to generate a respective output hidden state for each of the input tokens. The input hidden states for the first attention block are embeddings of the input tokens in the input sequence and the input hidden states for each subsequent attention block are the output hidden states generated by the preceding attention block. Atorney Docket No. 56113-0618WO1

[0044] In this example, the output subnetwork processes the output hidden state generated by the last attention block in the sequence for the last input token in the input sequence to generate the score distribution.

[0045] As described above, the language model neural network 110 has been pre-trained to determine current weight values 120, i.e., current values of the weights of the neural network 110, e.g. of the weights and. optionally, biases of the layers of the neural network 110. For example, the system 100 or another training system can have pre-trained the language model neural network 110 on a language modeling task, e.g., a task that requires predicting, given a current sequence of text tokens, the next token that follows the current sequence in the training data. As a particular example, the language model neural network 110 can be pre-trained on a next token prediction objective, i.e.. a maximum-likelihood objective on a large dataset of text, e.g., text that is publicly available from the Internet or another text corpus. Optionally, the pre-training on the language modeling task can have been followed by one or more additional sets of training, e.g., supervised fine-tuning, instruction tuning, and so on.

[0046] The neural network can have been trained by another system and the system 100 can receive the pre-trained parameter values for the neural network from the other system. As a particular example, the generative neural network can be a pre-trained “foundation” model that is available for use by the system.

[0047] Generally, because the neural network 110 is auto-regressive, the system 100 can use the same neural network 1 10 to generate multiple different candidate output sequences in response to the same context sequence, e.g., by using beam search decoding from score distributions generated by the neural network 110, using a Sample-and-Rank decoding scheme, by using different random seeds for the pseudo-random number generator that’s used in sampling for different runs through the neural network 110 or using another decoding scheme that leverages the auto-regressive nature of the neural network.

[0048] Other examples of such generative neural networks 110 include image or video generation neural networks, e.g., diffusion models, that generate images, audio or videos conditioned on context inputs, e.g.. text, audio, categorical variables, or other images.

[0049] The training output generated by the generative neural network can comprise image, video and / or audio content. The training input from which the training output is generated can comprise an instruction in natural language that instructs the generative neural network as to what content to generate. For example, the training input can specify the objects / entities that are to be present and / or the actions that are to be performed by the entities in the Atorney Docket No. 56113-0618WO1 image / video / audio. For audio, the training input can specify the sounds and any speech to be present in the generated audio.

[0050] More generally, the generative neural network 110 can be any appropriate type of generative model that can be used to generate multiple different output examples from any given context input and for which one or more reward functions 108 are available to the system 100.

[0051] As one example, the one or more reward functions 108 can include one or more reward models. A reward model is a machine learning model, e.g., a neural network, that has been trained to process an input that includes a sampled training output 122 generated by the generative neural network 110 for a training input 118 and, in some cases, the training input 118 to generate as output a reward score 124 that measures the quality of the sampled training output 122, e.g., the qualify according to an absolute measure or the quality given the training input 118.

[0052] As another example, instead of being a learned model, the reward function 108 can be a hard-coded function that evaluates the quality of an output example generated by the generative neural network 110. For example, when ground truth outputs for the context inputs are available, the reward function can be, e.g., a BLEU score function, an edit distance function, and so on. When ground truth output for context inputs are not available, the reward function can be. e.g., a textual coherence measure when the output examples are text or a non-reference image quality measure when the output examples are images.

[0053] To improve the performance of the neural network 1 10, the system 100 fine-tunes, i.e., further trains, the generative neural network 110 on a training data set 106 to generate the fine-tuned weight values 126. In particular, the system uses multiple hardware devices to perform sampling and training of the neural network 110 by using multiple data partitioning and model partitioning strategies, as described in further detail below with reference to FIG. 2.

[0054] In particular, at each training step, the system 100 performs sampling to generate a set of sampled training outputs 122. That is, the system samples a set of training inputs 118 and uses the neural network 110 to process each of the training inputs 118 to generate one or more training outputs 122 for each training input 118. That is, as described above, the same generative neural network can be used to “sample” multiple different training outputs 122 for a given training input 118.

[0055] After the sampling, the system 100 uses the reward functions 108 to train the generative neural network 110 by scoring the sampled training outputs 122 to generate Atorney Docket No. 56113-0618WO1 reward scores 124 and fine-tuning the neural network 110, i.e., updating the weight values of the neural network 110, based on the sampled training outputs 122 and the reward scores 124.

[0056] Fine-tuning the neural network 110 is described below with reference to FIGS. 2-3.

[0057] After training, the system 100 or another inference system deploys the fine-tuned language model neural network 110 for performing the task by storing the fine-tuned weight values 126 in the off-chip memory 104.

[0058] FIG. 2 shows a block diagram 200 showing the sampling and training phases of a given training step. For example, the system 102 can perform the sampling and training phases shown in FIG. 2. .

[0059] In particular, the reinforcement learning system 102 can more efficiently train the neural network 110 by parallelizing sampling and training across multiple different hardw are devices 202 to generate updated w eight values of the neural network 110.

[0060] Each set of hardware devices 202 can include one or more hardw are computing devices, e.g., one or more hardware accelerators, and optionally a general purpose processor. The hardware accelerators can each be computer chips that perform certain operations, e.g., matrix multiplication, in hardware. For example, the hardware accelerators can be tensor processing units (TPUs), graphics processing units (GPUs), or other machine learning accelerators that perform machine learning operations in hardw are.

[0061] To train the neural network, the system can maintain instances of the generative neural netw ork distributed across multiple sets of one or more hardware devices 202. That is, each of multiple instances of the generative neural network can be maintained by a set of one or more hardware devices.

[0062] Importantly, the system can perform sampling and training of the neural network 110 across the hardware devices 202 according to multiple different sharding schemes. A sharding scheme defines how many instances of the generative neural network are distributed across the sets of hardware devices and how7each instance is distributed across a given set of one or more devices.

[0063] That is, the system can use a preferred first sharding scheme for sampling that is different than a preferred second sharding scheme for training. A sharding scheme is a strategy7for dividing (e.g., partitioning) weights of the neural network 110 and, optionally, other data, into partitions or shards across one or more hardware devices 202.

[0064] The system can implement the first sharding scheme for sampling and the second sharding scheme for training based on one or more factors. Atorney Docket No. 56113-0618WO1

[0065] In general, for any given phase of a given training step, the system implements a corresponding sharding scheme to divide the weights across multiple hardware devices for optimizing overall system throughput according to a memory layout for a particular computation performed by each set of hardware devices during the phase. Because the sampling and training phases require the sets of hardware devices to perform different computation, the system implements different sharding schemes for each of the phases.

[0066] That is. for sampling, the system can use a first sharding scheme to partition the weights according to a memon layout that reduces communication across hardware devices to increase processing speed, while for training, the system can use a second sharding scheme to partition the weights according to a relatively lower memory usage scheme to increase memory availability for optimizer operations, gradient computation, or a combination thereof.

[0067] In particular, for sampling, the first sharding scheme can partition the current model weights across several sets of hardware devices 202 to maximize a number of instances of the neural network 110 deployed on the hardware devices. By maximizing the number of instances of the neural network 110, the system can maximize a speed of generating sampled training outputs 122.

[0068] For training, the second sharding scheme can partition the current model weights across hardware devices 202 in a manner that results in fewer instances of the neural network 110, but increased data parallelism. By implementing a scheme with fewer instances of the neural network 1 10 and increased data parallelism, the system can increase the efficiency of gradient computation, optimizer computations, or both.

[0069] In particular, for sampling, the system can assign a batch of one or more training inputs 118 to each set of hardware devices 202. The system can perform sampling by generating, in parallel across the sets of hardware devices 202, one or more sampled training outputs 122 for each training input 118 using the respective instance of the neural network 110 maintained by each of the sets of hardware devices 202.

[0070] The system then uses the reward function to generate a respective reward score 124 for each sampled training output 122.

[0071] As described above, the reward function can be a machine learning model, e.g., a neural network, that has been trained to process an input that includes a training input 118 and the sampled training output 122 to generate the respective reward score 124. In some implementations, each set of hardware devices 202 maintains a respective instance of the reward model and uses the respective instance of the reward model to score the sampled Atorney Docket No. 56113-0618WO1 training outputs 122 that are generated by the set of hardware devices 202 in parallel. In some other implementations, one or more respective instances of the reward model can be deployed off-chip. In some other implementations, the reward model can be a function that is evaluated off-chip.

[0072] By parallelizing the sampling in this manner using the first data sharding scheme, the system can significantly decrease the time required for generating sampled training outputs and scoring the sampled training outputs to generate the reward scores.

[0073] The sampling stage is followed by a training (e.g., updating) stage that uses the second sharding scheme to update the weight values. During training, the system uses the sampled training outputs 122 and the reward scores 124 to compute gradients and to generate updated model weight values 212 based on the second sharding scheme. In this way, the system can iteratively use the updated model weight values 212 to perform sampling based on the first data sharding scheme, and the system can generate the sampled training outputs 122 and the reward scores 124 for training based on the second sharding scheme.

[0074] In particular, the system includes a gradient engine 204 configured to generate gradients 208 from sampled outputs and corresponding reward scores and an optimizer 206 configured to optimize (e.g., fine-tune) the model weight values based on the gradients 208. During training, the system performs multiple training stages to improve the quality of the outputs that are generated by the neural network 110 by updating the model weight values 212. In particular, when the current training stage is the first training stage, the current model weight values are the pre-trained values of the weights. When the current training stage is not the first training stage, the current values are the values of the weights after being updated at the preceding training stage.

[0075] At each training stage, the system uses the gradient engine 204 to compute the gradients 208 by performing backpropagation using the sampled training outputs 122 and the reward scores 124. In particular, the system provides the hardware devices with a particular state of the neural network 110 for training. For example the system can provide the final neural network 110 trained at previous training stages (e.g., the final trained "‘supervised” neural network) to the hardware devices. In another example, the system can provide a randomly initialized neural network 110 to the hardware devices.

[0076] The system can compute the gradients 208 with respect to a reinforcement learning loss function, as shown by Equation 1 : Atorney Docket No. 56113-0618WO1

[0077] (1) Loss = (1 - a)VEn[r x,y ] - aEsEpK[VKL(n s>)\\psuperviseds>))] where (1 — a) is a reward tradeoff parameter, VEn[r(x, y)] is the reward function, and is a constraint term that is based on the weights of the pre-trained neural network 110 prior to fine-tuning. The constraint term measures a divergence between probabilities assigned to a given output by the generative neural network in accordance with the current weight values 120 and probabilities assigned to the given output by generative neural network in accordance with the pre-trained supervised neural network weights.

[0078] More generally, the system can perform the training phase using any appropriate reinforcement learning loss function that is based on the reward scores for the training outputs.

[0079] The system uses the loss function to compute a respective gradient 208 corresponding to each weight value of the neural network 110.

[0080] To perform the training phase, the system also obtains optimizer data 210 for the optimizer 206. Generally, the optimizer 206 can map the updated model weight values to the gradients 208. For example, the optimizer 206 can be a stochastic gradient descent (SGD) optimizer, an Adam optimizer, an Adafactor optimizer, or an Adamax optimizer.

[0081] For example, in the case that the system includes an Adam optimizer or an Adafactor optimizer, the optimizer data 210 can be one or more momentum values associated with the optimizer 206. In particular, the one or more momentum values can be based on one or more moving averages of the gradients 208. The optimizer 206 is configured to map the updated model weight values 212 to each of the gradients 208. In some cases, the optimizer 206 is configured to use the momentum values to map the updated model weight values 212 according to a particular algorithm.

[0082] As described above, the system then provides the updated model weight values 212 to the hardware devices 202 to perform sampling of the outputs generated by the neural network 110 using the updated model weight values 212, i.e., as part of the next training step.

[0083] FIG. 3 is a flow diagram of an example process for fine-tuning weights of a neural network by performing a sampling phase and a training phase according to respective strategies. For convenience, the process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural Atorney Docket No. 56113-0618WO1 network training system, e.g., the neural network training system 100 of FIG. 1, appropriately programmed, can perform the process 300.

[0084] The system trains the neural network by fine-tuning weight values over multiple training steps. Each training step includes the sampling phase in accordance with a first scheme and the training phase in accordance with a second scheme.

[0085] At the first training step, the system trains the neural network starting from the values of the weights determined from the pre-trained values of the weights. For any subsequent training steps, the system trains the neural network starting from the values of the weights determined at the end of the preceding training step.

[0086] During the sampling phase, the system loads first model data including current values of the multiple weights into memory of multiple hardware devices in accordance with the first scheme (302). In particular, the system loads the first model data in accordance with a first sharding scheme that partitions the multiple weights among the multiple hardware devices.

[0087] Each of the sets of one or more hardware devices obtain multiple training inputs (304). In particular, each set of hardware devices obtains the multiple training inputs while the current values of the multiple weights are loaded (e.g., with the current values of the multiple weights being stored in accordance with the first scheme).

[0088] At each set of hardware devices, the system processes each of the multiple training inputs using the neural network according to the current values of the multiple weights to sample one or more training outputs for each training output (306).

[0089] The system then generates a respective reward score for each of the training outputs generated by each of the hardware devices (308). The system can use one or more reward functions to generate the reward score for the sampled training outputs. In particular, for each set of one or more hardware devices, the system can use an instance of a reward model deployed on the set of the one or more hardware devices to generate the respective reward score.

[0090] In some implementations, the system can use a third, different sharding scheme to generate the reward scores. For example, for each set of one or more hardware devices, the system can move the weights of the neural network to off-chip memory, e g., to random access memory (RAM) located off-chip, e.g., on a host computer system, and load the weights of the reward model into on-chip memory, and then score the training outputs generated by the set of hardware devices using the reward model. Atorney Docket No. 56113-0618WO1

[0091] To perform the training phase, the system loads second model data including current values of the multiple weights of the neural network and optimizer data for an optimizer into memory of multiple hardware devices according to the second scheme (310). In particular, the system loads the second model data in accordance with a second sharding scheme that partitions the weights among the multiple hardware devices. The second sharding scheme is generally different than the first sharding scheme. For example, the second sharding scheme may deploy fewer instances of the neural network among the multiple hardware devices than the first sharding scheme in order to maximize gradient computation efficiency.

[0092] The system then trains the neural network through reinforcement learning using the reward scores to update the current values of the multiple weights (312). In particular, the system computes gradients using the sampled training outputs and the reward scores, and the system generates updated model weight values based on the second sharding scheme.

[0093] To update the model weight values, the system can implement an optimizer configured to update the model weight values based on the gradients. Generally, the optimizer can map the gradients to the updated model weight values to the gradients.

[0094] The system can then provide the updated model weight values to the multiple hardware devices to perform sampling using the updated model weight values for the subsequent training step.

[0095] In general, for each subsequent training step, the system performs a new sampling phase and anew training phase by optimizing the on-device memory space of the one or more hardware devices according to the first scheme and according to the second scheme, respectively. Initially, the system can attempt to perform the phases while keeping data in the on-device memory of the hardware device (e.g., without moving the data to the off-chip memory 104. If the system determines that there is not enough memory availability of the one or more hardware devices, the system can move the data from the on-device memory to the off-chip memory7104 (e.g., data that is not required for performing the particular stage). That is, the system can optimize memory layout for the hardware devices by determining whether to remove data associated with particular training step components, such as the gradient engine 204 or the optimizer 206, based on whether the data is required during a particular phase of training step.

[0096] In particular, prior to sampling, the system can remove data (e g., model weights) associated with the optimizer 206, the reward model for scoring 114, or any other models from the on-device memory of the hardware devices to the off-chip memory' 104. During Atorney Docket No. 56113-0618WO1 sampling, the system performs scoring of the outputs using the updated model weight values 212. Once the system has performed scoring of the sampled training outputs 122 using the reward model, the system can move the updated weight values 212 from the on-device memory to the off-chip memory 104, and the system can load the reward model onto the hardware devices 202. During training, the system then partitions the model weights (e.g., re-shards the model weights) according to the second scheme to reduce memory’ usage and to load the optimizer data onto the hardware devices for training.

[0097] As one example, prior to performing any given sampling phase, to implement the first sharding scheme, the system can move out the optimiser state, the reward model state and optionally other models if they are involved in the training process to RAM in order to make more space for sampling using the neural network. The neural network is then re- sharded (optimised) for sampling, which requires more device memory (e.g., high bandwidth memory’ (HBM) or other on-chip memory) than the training phase.

[0098] The system then preforms the sampling phase to sample from the neural network.

[0099] For scoring, the system can move the neural network to RAM. while the reward model is loaded onto devices.

[0100] For the training phase, the neural network is again re-sharded to reduce its memory usage for the training step and optimiser is loaded onto devices.

[0101] Generally, moves to RAM happen only if required, e.g., the system can determine whether performing operations given the data currently stored in on-chip memory is possible and only if not, moves out model states to RAM which are not involved in subsequent computations within the current phase.

[0102] After performing the multiple training steps, the system can output the updated model weight values (e.g., the fine-tuned weight values). For example, the system can output the fine-tuned weight values to one or more of the multiple hardware devices, to the memory', off-chip, or a combination thereof.

[0103] This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions. Atorney Docket No. 56113-0618WO1

[0104] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry. in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by. or to control the operation of, data processing apparatus. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially-generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0105] The term "‘data processing apparatus"’ refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry , e.g., an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). The apparatus can optionally’ include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0106] A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may. but need not. correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub-programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one Atorney Docket No. 56113-0618WO1 site or distributed across multiple sites and interconnected by a data communication network.

[0107] In this specification the term “engine’’ is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.

[0108] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry . e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0109] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read-only memory' or a random access memory' or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto-optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

[0110] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory' devices; magnetic disks, e.g., internal hard disks or removable disks; magnetooptical disks; and CD-ROM and DVD-ROM disks. Atorney Docket No. 56113-0618WO1

[0111] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g.. visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

[0112] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads.

[0113] Machine learning models can be implemented and deployed using a machine learning framework, e.g.. a TensorFlow framework.

[0114] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back-end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front-end component, e g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0115] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some Atorney Docket No. 56113-0618WO1 embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device.

[0116] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0117] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0118] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

[0119] What is claimed is:

Claims

Atorney Docket No. 56113-0618WO1CLAIMS1. A method of training a neural network having a plurality of weights, the method comprising, at each of a plurality of training iterations: loading first model data comprising current values of the plurality of weights into memory of a plurality' of hardware devices in accordance with a first scheme; after the first model data has been loaded in the memory of the plurality of hardware devices in accordance with the first scheme: obtaining, by the plurality of hardware devices, a plurality of training inputs; processing, using the plurality’ of hardware devices, each of the plurality of training inputs using the neural network in accordance with the current values of the plurality of weights to sample one or more training outputs for each of the plurality’ of training inputs; generating, by the plurality of hardware devices, a respective reward score for each of the training outputs; loading second model data comprising the current values of the plurality’ of weights and optimizer data for an optimizer into the memory of the plurality of hardware devices in accordance with a second scheme; after the second model data has been loaded in the memory of the plurality of hardware devices in accordance with the second scheme: training the neural network through reinforcement learning using the reward scores for the training outputs to update the current values of the plurality of weights.

2. The method of claim 1, wherein generating, by the plurality of hardware devices, a respective reward score for each of the training outputs comprises: determining the respective reward score for each of the training outputs by using a reward function.

3. The method of claim 2, wherein the reward function is a machine learning model pre-trained to process an input comprising the training output to generate the respective reward score for the training output.Attorney Docket No. 56113-0618WO14. The method of claim 1, wherein training the neural network through reinforcement learning using the reward scores for the training outputs to update the cunent values of the plurality of weights comprises: computing gradients with respect to the plurality of weights of a reinforcement loss function based on the reward scores; and updating the current values using the gradients.

5. The method of claim 4, wherein updating the current values using the gradients comprises: applying, to the gradients, the optimizer in accordance with the optimizer data to update the current values.

6. The method of claim 1, further comprising: after training the neural network, providing the updated values of the plurality of weights to one or more hardware devices of the plurality of hardware devices.

7. The method of any preceding claim, wherein the first scheme distributes a first number of instances of the generative neural network across the plurality of hardware devices.

8. The method of claim 7, wherein the second scheme distributes a second number of instances of the generative neural network across the plurality of hardware devices.

9. The method of claim 8, wherein the first number is greater than the second number.

10. A system comprising one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform the operations of the respective method of any one of claims 1-9.

11. One or more computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations comprising the operations of the respective system of any one of claims 1 -9.

Images