EP4609322A1 - Distributed training of large neural networks - Google Patents

Abstract

Systems and methods for using a distributed computing system to train a large neural network to perform a machine learning task. A shared set of trainable parameters is maintained in a shared data store, and each of a geographically distributed set of workers updates their trainable parameters using a shard training dataset. There are two optimization processes: an outer optimization process, and an inner optimization loop that is executed by each worker independently and in parallel tens, hundreds, or thousands of times. The workers can have different computing capabilities and can be geographically distant from one another, and the communications bandwidth used by the system can be two or three orders of magnitude less than that of other systems.

Description

DISTRIBUTED TRAINING OF LARGE NEURAL NETWORKS

CROSS-REFERENCE TO RELATED APPLICATION

[0001] This application claims priority to U.S. Provisional Application No. 63/598,463, filed on November 13, 2023. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application.

BACKGROUND

[0002] This specification relates to training neural networks.

[0003] Neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. The neural networks described herein include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of trainable parameters, e.g. weights.

SUMMARY

[0004] This specification generally describes systems and methods implemented as computer programs on one or more computers, generally in multiple locations, for using a distributed computing system to train a large neural network to perform a machine learning task.

[0005] In one aspect there is described a computer-implemented method of training a neural network, in particular a large neural network, e.g. one having at least 10⁷ trainable parameters (such as weights), using a distributed computing system. The distributed computing system comprises a plurality of computing devices coupled by a computer network and operating in parallel.

[0006] The method involves maintaining a shared set of values of the trainable parameters in a shared data store (coupled to the computer network). A set of workers, typically a geographically distributed set of workers, is implemented on the plurality of computing devices, each worker implemented on one of the devices. Each worker is allocated one of a plurality of shards of a training dataset (stored in the memory of the computing device). A worker may comprise a computer, or a computer with one or more machine learning hardware accelerators (e.g. that perform matrix multiplication in hardware), or a cluster of computers.

[0007] Each worker obtains a local copy of the values of the trainable parameters and performs at least H = 50, 100, 500, 1000, or 2000 inner optimization steps on these, to obtain an updated set of trainable parameter values. A worker obtains the updated set of trainable parameter values by training an instance of the large neural network using its shard of the training dataset. In general the training can involve backpropagating gradients of an objective function (for whatever machine learning task the neural network is being trained to perform) to adjust the values of the local copy of the trainable parameters to obtain the updated set of trainable parameter values for the worker.

[0008] At each of a plurality of outer optimization steps a cumulative gradient for some or all of the trainable parameters is determined for each worker, e.g. by each worker. The cumulative gradient for a parameter can be determined using the value of the trainable parameter in the local copy obtained by the worker, i.e. the updated trainable parameter value determined by the worker after the H inner optimization steps. More particularly this can involve determining a difference between the value of the trainable parameter in the local copy obtained by the worker from the shared data store, and the value of the parameter after the H inner optimization steps. [0009] The method also involves determining an aggregated cumulative gradient at an optimization computer system. The optimization computer system can be, e.g., one of the computing devices, or a shared server coupled to the computer network.

[0010] In implementations the aggregated cumulative gradient is determined by, for each of some or all of the trainable parameters, aggregating the cumulative gradients for the trainable parameter determined for, e.g. by, the respective workers. The aggregated cumulative gradient for the trainable parameter can then be used to update the value of the trainable parameter in the shared data store.

[0011] At the end of the training process the values of the trainable parameters in the shared data store provide the trained parameters for the large neural network.

[0012] The described techniques are suitable for implementation in an environment in which some or all of the computing devices are geographically distant from one another, e.g. where the computer network comprises a wide area network (WAN), and in which the communications bandwidth is relatively limited. However the techniques can also be implemented on a single worker, e.g. for faster convergence and better generalization.

[0013] In a related aspect there is described a corresponding process, from the perspective of the optimization computer system.

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

[0015] Conventional approaches to training large neural networks on a distributed computing system require a large number of tightly interconnected hardware accelerators such as GPUs or TPUs. This is in part because existing algorithms require high bandwidth communications between devices that need to exchange data with one another during the training. For example a typical synchronous system for training a large neural network can involve thousands of colocated devices interconnected with high-bandwidth cables, and with careful software engineering to orchestrate the devices and achieve high utilization.

[0016] In principle a more widely distributed computing system could be used, but when training a large neural network very high bandwidth communications would be needed to transport, e.g., the many weights or gradients. Implementations of the described techniques can reduce the communication required between workers, or islands of workers, implemented on respective computing devices by two or three orders of magnitude compared with such conventional approaches.

[0017] The described techniques also facilitate a much more distributed approach - for example the workers can be physically spaced apart from one another and connected by a computer network such as a LAN (local area network) or WAN (wide area network). With the described techniques workers, or islands of locally connected workers, can even be located in different countries.

[0018] There is no requirement for the workers to have the same computing capability or memory capacity, and the workers can have different computing and/or memory capabilities, and/or capabilities that vary over time. In implementations the workers can nonetheless be kept almost fully utilized.

[0019] Implementations of the described techniques are also very robust to a range of different factors. For example they can cope with a worker stalling or failing and having to reboot and reload parameters from a point at which it stopped. They can also cope with varying numbers of workers over the course of training a neural network, and do not require different workers to have the same distribution of training data. The techniques are also robust to disruption of communications, e.g. resulting in lost packets.

[0020] The techniques can be used to train a neural network from scratch, e.g. from random initialization, or to fine tune a previously trained neural network. The described techniques can also train a large neural network faster than conventional techniques in wall-clock time.

[0021] A neural network trained using the described techniques can achieve better performance, e.g. lower perplexity or better generalization, than some synchronous training techniques. At inference time the resulting model can have the same size and speed as a model trained in fully synchronous mode. [0022] 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, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.

BRIEF DESCRIPTION OF THE DRAWINGS

[0023] FIG. 1 shows an example distributed computing system for training a large neural network.

[0024] FIG. 2 is a flow diagram of an example process for training a large neural network using a distributed computing system.

[0025] FIG. 3 is a visualization of an example of the described training process.

[0026] FIG. 4 compares results of an example of the described training process with conventional techniques.

[0027] FIG. 5 illustrates a result of an example of the described training process using Nesterov momentum.

[0028] FIG. 6 illustrates results of examples of the described training process with varying numbers of workers.

[0029] FIG. 7 illustrates results of examples of the described training process with varying communications quality.

[0030] FIG. 8 illustrates a result of an example of the described training process implemented on a single worker.

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

DETAILED DESCRIPTION

[0032] Techniques are described for training a neural network, in particular a large neural network to perform a machine learning task. Some implementations of the techniques use a distributed computing system with multiple workers that work in parallel. As used herein, “training” includes fine tuning a pre-trained model.

[0033] The large neural network can have at least 10⁷ trainable parameters; in some implementations at least 10⁸ or 10⁹ trainable parameters. In general the neural network is configured to process a neural network input in accordance with the trainable parameters of the neural network, to generate a neural network output that defines a result of a machine learning task. For example the neural network may comprise a so-called foundation model such as a large language model (LLM) or vision language model (VLM).

[0034] In some implementations a shared set of the trainable parameters is maintained in a shared data store, and each of a geographically distributed set of workers updates their trainable parameters using a shard training dataset. There are two optimization processes: an outer optimization process, and an inner optimization loop that is executed by each worker in parallel, and in general independently, tens, hundreds or thousands of times. The workers can have different computing capabilities and can be geographically distant from one another; the communications bandwidth used by the system can be two or three orders of magnitude less than that of other systems.

[0035] FIG. 1 shows an example of a distributed computing system 100 for training a large neural network.

[0036] The distributed computing system 100 comprises a plurality of computing devices 110a..n each having associated (local) memory 114a..n. A computing device may comprise, as some examples, a general purpose computer system, or a computer equipped with one or more GPUs (Graphics Processing Units), or a computer equipped with one or more TPUs (Tensor Processing Units), or a combination of these.

[0037] The computing devices 110 are coupled by a shared computer network 120 and can operate in parallel. The computer network 120 may include wired and/or wireless parts and may comprise, e.g. a local area network (LAN), a wide area network (WAN), or the Internet. An advantage of the described techniques is that they are suitable for implementation in an environment in which some or all of the computing devices are geographically distant from one another. In some implementations the computer network 120 may therefore include a wide area network (WAN).

[0038] Each of the computing devices 110 is configured to implement one or more workers 112a.. n, e.g. one worker can be implemented on each computing device. In general a worker can comprise code for a software process that operates on one of the computing devices 110, e.g. using multiple GPUs or TPUs of the computing device.

[0039] The distributed computing system 100 also includes a shared data store 130, accessible over the computer network 120, e.g. on a shared, e.g. central, server or on one of the computing devices 110. The shared data store 130 is used to maintain shared set of values of the trainable parameters of the large neural network, as described later. [0040] The distributed computing system 100 also includes an optimization computer system 150. The optimization computer system may comprise one of the computing devices 110 or it may comprise a shared server coupled to the computer network.

[0041] FIG. 2 is a flow diagram of an example process for training a neural network, in particular a large neural network using a distributed computing system, such as the distributed computing system 100 of FIG. 1. For convenience the process will be described with reference to the distributed computing system 100.

[0042] The neural network can be one that has been initialized at random, i.e. for which initial values of the trainable parameters have been chosen at random, or it can comprise a pre-trained model. The large neural network can be trained to optimize an objective function. The objective function can be any objective function appropriate to a task to be performed by the (trained) neural network.

[0043] The process involves obtaining a training dataset (step 200). The training dataset comprises a plurality of shards, each stored in the memory 114 of one, e.g. a respective one, of the computing devices 110. Each shard is typically a disjoint partition of the training dataset, though in principle some overlap or duplication between data in different shards is possible. A sharded training dataset can be obtained by partitioning a training dataset into shards. For example the training dataset can be divided into portions, or shards, each of which holds a respective portion of the training dataset, and that together make up the complete training dataset.

[0044] The process maintains a shared set of values of the trainable parameters of the large neural network in the shared data store 130. A set of workers 112 is implemented on the plurality of computing devices 110 (step 202). In particular each worker is implemented on one of the computing devices, e.g. each worker can be implemented on a different respective one of the computing devices, and each worker is allocated a respective one of the plurality of shards.

[0045] As one example, each worker may be allocated a respective one of the shards by the shared, e.g. central, server. As another example, each worker may obtain its shard without the shard being allocated by any form of central control, e.g. to avoid a central server having access to one or more of the shards, e.g. to assist in preserving privacy. For example, one or more of the workers may obtain a portion of training data without exposing the portion of training data to any of the other workers or to a central server (if present); and this portion of training data may constitute a shard of the training dataset. There are known techniques for achieving this, such as consistent hashing (a decentralized approach in which a worker can obtain a shard based on a hash function), and other hash-based techniques.

[0046] In some implementations there is one worker and one shard per computing device; in some implementations one or more of the computing devices may implement more than one worker and store more than one shard. Optionally, implementing a worker 110 on a computing device 112 may itself involve data and/or model parallelism.

[0047] In general the method involves performing a plurality of outer optimization steps. At each of the plurality of outer optimization steps each worker 112 performs a number of inner optimization steps, in general independently of each other worker, and in parallel with each other worker.

[0048] Each outer optimization step can involve each worker 112 obtaining a local copy of the values of the trainable parameters from the shared data store 130 (step 204). More particularly the computing device 110 implementing the worker 112 can obtain the local copy of the values of the trainable parameters from the shared data store 130 and store it in the (local) memory 114 of the computing device.

[0049] Each outer optimization step can also involve each worker 112 performing H inner optimization steps on the local copy of the values of the trainable parameters, to obtain an updated set of trainable parameter values for the worker after the H inner optimization steps (step 206). Here H is at least 50, 100, 500, 1000, or 2000.

[0050] In implementations each inner optimization step involves backpropagating gradients of the objective function, evaluated using the shard allocated to the worker, to adjust the values of the local copy of the trainable parameters (step 208).

[0051] The method also involves, for each of the plurality of outer optimization steps, and for each worker 112, determining a cumulative gradient (or “outer” gradient) for at least some of the trainable parameters (step 210).

[0052] The cumulative gradient is determined by determining a difference between the (updated) value of the trainable parameter in the local copy obtained by the worker computing device and the value of the parameter after the H inner optimization steps. That is, the cumulative gradient for a trainable parameter is dependent upon a difference between a value of the parameter before and after the H inner optimization steps

[0053] For example the cumulative gradient can be determined as denotes the updated set of trainable parameter values for the worker after the H inner optimization steps, and denotes the values of the trainable parameters obtained by the worker from the shared data store, i.e. before the H inner optimization steps. The cumulative gradient may be determined by the worker, i.e. locally, or (if the updated trainable parameter values are communicated elsewhere), the cumulative gradient may be determined by one of the other workers, or centrally (as used here “central” refers to function rather than particularly to geography).

[0054] Each of the plurality of outer optimization steps also involves determining an aggregated cumulative gradient (step 212). The aggregated cumulative gradient can be an aggregate, e.g. an average, of the cumulative gradients determined by the workers.

[0055] The aggregated cumulative gradient can be determined at the optimization computer system 150 which, as previously mentioned, may comprise one of the computing devices 110 or a shared server.

[0056] Determining the aggregated cumulative gradient can involve, for each trainable parameter of at least some of the trainable parameters, aggregating, e.g. averaging, the cumulative gradients for the trainable parameter determined for or by the workers 112. As previously mentioned the cumulative gradient of a trainable parameter may be determined at (by) a worker, or for a worker, e.g. by a central server.

[0057] As an example the aggregated cumulative (“outer”) gradient, A®, can be an average cumulative gradient. Such an average cumulative gradient can, e.g., be determined by evaluating denotes the updated set of trainable parameter values for worker i of k workers after the H inner optimization steps, and denotes the values of the trainable parameters obtained from the shared data store, i.e. before the H inner optimization steps. That is, the aggregated, e.g. average, cumulative gradient (per parameter) can be determined from a difference between trainable parameter values for a current and previous outer loop optimization time step.

[0058] Optionally, for example where the shards for different workers comprise different numbers of training data items, the cumulative gradient for each worker can be weighted by a number of training data items, n_k, in the shard for the worker. For example the average cumulative gradient can be a weighted average cumulative gradient, e.g. determined by evaluating a weighted sum — 0^f) where w_k denotes a weight, which may be determined as n_k/n where n is a total number of training data items in the training dataset.

[0059] In general determining the aggregated cumulative gradient for a parameter involves communicating from a worker 112 to an entity performing the averaging, e.g. the optimization computer system 150. The can involve communicating either the cumulative gradient for the parameter from the inner loop, or the updated parameter value from the inner loop (the original parameter value will typically still be available from the shared data store 130).

[0060] As described later, in some implementations the communication from a worker is pruned so that not all the cumulative gradients (or updated parameter values) are transmitted, e.g., to the optimization computer system. For example a threshold may be applied so that only cumulative gradients that have magnitudes greater than a threshold value are sent. Different workers may send cumulative gradients for different groups of trainable parameters. Thus in some implementations the aggregated cumulative gradient may only be determined for some of the trainable parameters. For (one of) these trainable parameters determining the aggregated cumulative gradient may only involve aggregating the cumulative gradients (for the trainable parameter) from some of the workers.

[0061] Each of the plurality of outer optimization steps also involves updating the value of the trainable parameter in the shared data store (step 214), i.e. the value of the trainable parameter in the shared set of values of the trainable parameters.

[0062] In general this involves using the aggregated cumulative gradient for the trainable parameter to determine the value of the trainable parameter. In some implementations this can involve applying, to a value dependent on a current value of the trainable parameter, a parameter offset dependent on the aggregated, e.g. average, cumulative gradient for the trainable parameter.

[0063] In implementations, after the training, the shared set of values of the trainable parameters in the shared data store 130 provides the trained parameters for the large neural network (step 216).

[0064] Implementations of the above described techniques can reduce communication between the workers and, e.g. a central server, by some orders of magnitude. Nonetheless, particularly when workers are synchronized, for large models and/or a large number of workers there can be a heavy load on the network e.g. when a set of H inner optimization steps is completed.

[0065] In some implementations this can be mitigated by reducing the communications bandwidth from the workers necessary at these points, in particular by pruning a number of data items (cumulative gradients or parameter values) sent from a worker to an entity, e.g. server, performing the cumulative gradient aggregating.

[0066] In some implementations the value of H may be different for different workers, e.g. to desynchronized the workers and mitigate spikes in the communications bandwidth, as the communication to/from the workers and/or to/from a central server, is then staggered. [0067] In some implementations the value of H for a worker may be selected or modified dependent on a characteristic of the worker or on a characteristic of available communications bandwidth to or from the worker. The characteristic of the worker may, e.g., be a metric of its available computing capability or memory capacity or of its computing speed. The characteristic of the communications may, e.g., be a metric of average or peak bandwidth capacity of a network connection to or from the worker. In this way the distributed computing system can adapt to, e.g., changing network or worker capacity or conditions.

[0068] In some implementations each worker determines the cumulative gradient for each (all) of the trainable parameters. The worker can then communicate the cumulative gradient for some or all of the trainable parameters to the optimization computer system. In such implementations the worker may select a subset, more particularly a proper subset (i.e. less than all) of the cumulative gradients to communicate to the optimization computer system. The selected subset may be different for different workers. When determining the aggregated cumulative gradient, if the cumulative gradient for a trainable parameter has not been communicated by a worker it may be given a value of zero.

[0069] A subset of cumulative gradients to communicate may be selected by a worker in many ways. For example the cumulative gradients to communicate may be selected by comparing each cumulative gradient with a threshold value, communicating only cumulative gradients that have a value greater than the threshold. The threshold may be chosen, e.g. to select a largest x% (0 < x < 100) of the cumulative gradients. Also or instead cumulative gradients to communicate may be selected by their sign. More particularly a positive and a negative aggregated sign value may be determined for positive and for negative cumulative gradients, by summing the cumulative gradients with, respectively, a positive and a negative sign; and then only cumulative gradients having the sign of the aggregated sign value with the largest magnitude may be communicated.

[0070] In some implementations, rather than communicating the cumulative gradients, each worker may communicate some or all of the trainable parameters of the updated set of trainable parameter values for the worker, from the worker to optimization computer system over the computer network. Then the optimization computer system can determine the cumulative gradient for the communicated trainable parameters. This can also involve selecting, at a worker, a subset of trainable parameters from the updated set of trainable parameter values for the worker, and communicating only the selected subset of the trainable parameters. Again the trainable parameters may be selected in many ways, e.g. comparing each with a threshold value. Again the selected subset may be different for different workers. [0071] In implementations each inner optimization step involves a worker 112 using stochastic gradient descent to adjust (update) the values of the local copy of the trainable parameters, in particular stochastic gradient descent with momentum. This can involve maintaining a first momentum variable for each trainable parameter, locally at the worker. In implementations the first momentum variable comprises a (decaying) moving average, e.g. an exponential moving average, of a gradient used for adjusting the trainable parameter. The value of the trainable parameter can then be adjusted by an offset dependent on a value of the first momentum variable.

[0072] In a particular example the stochastic gradient descent with momentum can comprise the Adam optimization algorithm, or AdamW (Loshchilov et al., 2019). Adam and AdamW maintain both first and second momentum variables where the second momentum variable depends on the square of the gradient and is used to adapt the learning rate.

[0073] In implementations, aggregating the cumulative gradients at each outer optimization step, after H > 50 inner optimization steps, also uses a momentum-based approach, but different to that used for the inner optimization step, for stability. This can involve maintaining a second momentum variable (different to the first momentum variable) for each trainable parameter in the shared data store 130, and using this to adjust the respective trainable parameter, in effect by adjusting a relative contribution of previously determined aggregated, e.g. average, cumulative gradients when updating the value of the trainable parameter.

[0074] In particular updating the value of the trainable parameter in the shared data store using the aggregated cumulative gradient for the trainable parameter can involve maintaining, in the shared data store 130, a second momentum variable for the trainable parameter. An updated value for the second momentum variable is determined from the aggregated, e.g. average, cumulative gradient for the trainable parameter. Then an updated value of the trainable parameter in the shared data store is determined from the updated value of the second momentum variable. In implementations the updated value of the second momentum variable is scaled by a factor, e.g. in the range [0,1], that weights an effect of the momentum.

[0075] More particularly the updated value for the second momentum variable can be determined from a current value of the trainable parameter and the aggregated, e.g. average, cumulative gradient for the trainable parameter. Optionally the aggregated, e.g. average, cumulative gradient can be scaled by an outer optimization learning rate parameter, rj. As an example the updated value for the second momentum variable, can be determined by summing and, optionally scaled, j_n some implementations the outer optimization learning rate parameter may be kept constant rather than decayed; i.e. just a learning rate of the inner optimization may be decayed.

[0076] The updated value of the trainable parameter can be determined from a difference between the updated value for the second momentum variable, (optionally scaled), and a current value of the second momentum variable, Depending on how it is evaluated this may involve determining a momentum offset from a difference between a current value of the second momentum variable and the updated value for the second momentum variable. As some examples updated trainable parameters may be determined 0^ = —

[0077] Merely as an example, in some implementations the large neural network can comprise a Transformer neural network. In general a Transformer neural network can be characterized by having a succession of attention, e.g. self-attention, neural network layers. Each attention neural network layer has an attention layer input for each element of the input and is configured to apply an attention mechanism over the attention layer input to generate an attention layer output for each element of the input. In implementations the attention mechanism computes a similarity between a query and a set of key-value pairs. In implementations one or both (in the case of self-attention) of the query and the set of key-value pairs are determined from the attention layer input.

[0078] As an example, in a self-attention neural network layer an input embedding may be used to determine a query vector and a set of key -value vector pairs, that are used to generate an updated embedding comprising a weighted sum of the values, weighted by a similarity function of the query to each respective key. The similarity function may comprise, e.g., a dot product, cosine similarity, or other similarity measure; the query, keys, and values may all be vectors. For example the attention mechanism may be configured to apply each of a query transformation e.g. defined by a matrix W^Q , a key transformation e.g. defined by a matrix W^K , and a value transformation e.g. defined by a matrix W^v, to the attention layer input for each element of an input sequence Xto derive a respective query vector Q = XW^Q, key vector K = XW^K , and value vector V = XW^V which are used determine an attended sequence for the output.

[0079] In some implementations the large neural network comprises a large language model (LLM), e.g. a transformer-based LLM or a large vision language model (VLM), e.g. a transformer-based VLM. Then adjusting the values of the local copy of the trainable parameters in the inner optimization steps can use the AdamW optimization algorithm and, in the outer optimization steps, updating the value of the trainable parameter in the shared data store using the aggregated cumulative gradient for the trainable parameter can use the Nesterov optimization algorithm (Sutskever et al., “On the importance of initialization and momentum in deep learning”, ICML, 2013). Use of the Nesterov optimization algorithm, also referred to as Nesterov’s Accelerated Gradient (NAG), has been found to be particularly stable and effective, especially for larger values of H, e.g. H > 100, 200, 500, 1000 or 2000.

[0080] In some implementations the method (and corresponding system) are operated synchronously, i.e. waiting for all the workers to report the cumulative gradients for a trainable parameter before determining the aggregated cumulative gradient for the parameter. However the method (and corresponding system) may also be operated asynchronously. That is, determining the aggregated cumulative gradient for a trainable parameter at the optimization computer system may be performed without waiting to obtain an updated trainable parameter value or cumulative gradient for the trainable parameter from all of the workers.

[0081] As previously described in some implementations the optimization computer system 150 comprises a shared server coupled to the computer network 120.

[0082] FIG. 3 shows a visualization that outlines an example of a training process as described above. In this example an initial, pre-trained model (neural network) 300, denoted 0^\ is replicated k times (in the example k = 4), and each worker 112, denoted 0 , trains a model replica on its own shard of data for H steps independently and in parallel. Afterwards the workers average their outer gradients, and an outer optimizer updates the global copy of the parameters, 0^\ to obtain an updated model, 0^² The updated model is re-dispatched to the workers, and the process repeats T times (the first two iterations are indicated). A worker trains for a total of N = T X H inner steps; communication only happens every H inner steps. As indicated by the flag icons, the replicas can be in different geographical locations, e.g. different countries; and as indicated by the various TPU and GPU designations, the workers can have different hardware capabilities.

[0083] As a particular example of a training process as described above, an initial model, 0^° which can be initialized at random or pre-trained, is trained using a training dataset T> = {(x, y), ■■■ }, where x denotes an input data item processed by the model and y denotes a corresponding target output from the model. For example, where the model comprises an LLM or VLM the input data item can comprise a sequence of tokens and the target output can comprise the input sequence shifted by one token. The training dataset is split across multiple shards, and the ith shard is denoted A training process that involves k workers, and that uses an inner optimizer InnerOpt and an outer optimizer OuterOpt, can then be as follows:

1 for outer step t = 1 . . . T do

2 for worker i = 1 . .. fc do

9 end for

10 end for

11 Averaging outer grade

[0084] In this example the T-step outer optimization process involves lines 1, 12 and 14, and at each step t gradients from each worker are collected, averaged, and passed to the outer optimizer OuterOpt to update the shared copy of the parameters. Afterwards this shared copy of the parameters is re-dispatched to each local worker (line 3).

[0085] Within each outer optimization phase each worker independently and in parallel performs its own inner optimization (lines 4-9) for H, e.g. several hundred, steps, using the inner optimizer InnerOpt. Each worker samples data from its own shard (line 5), and updates its own local copy of the trainable parameters (line 8).

[0086] In this particular example the cumulative gradients are averaged (line 12) once all workers have completed their inner optimization. The resulting outer gradient is then used to update the shared copy of the trainable parameters (line 14), ready for the next round of inner optimizations. Communication need only happen every H inner optimization steps.

[0087] Adam, or a variant of Adam, can be used for both the inner and the outer optimizer (in the outer optimizer stability is helped by increasing the numerical stability factor e to 0.1). However in implementations using Nesterov momentum in the outer optimizer is particularly helpful for stable learning, especially when H is large. In general an update using Nesterov momentum can involve updating the value of the trainable parameter with respective momentum values before determining a gradient-based update, where the momentum values accumulate previous parameter updates.

[0088] As a particular example, a Nesterov optimizer can first update a momentum variable with a moving average, e.g. an exponential moving average, of its previous momentum and the cumulative or outer gradient. The Nesterov optimizer can then compute an update for a trainable parameter using another moving average, e.g. an exponential moving average, of the previously computed momentum and the cumulative or outer gradient. For example an update 0_oid to 0_new using Nesterov momentum can be determined as (in the order below):

// = (1 — <z)/z + a ■ A outer_update = (1 — a')/! + a' ■ A d_new = Gold + T ■ outer_update where denotes momentum, a, a' E [0,1) are decay coefficients (that may be the same), A denotes the aggregated cumulative, or outer, gradient, and the outer learning rate is T.

[0089] In a variant, a method and system as described above may be used to accelerate a single worker. Nonetheless such a worker can be implemented using, e.g., multiple processing threads, and can run on hardware comprising multiple hardware accelerators such as GPUs or TPUs.

[0090] Thus there is also described a computer-implemented method, and a corresponding system, for training a large neural network, e.g. having at least 10⁷ trainable parameters. This can involve maintaining a set of values of the trainable parameters and, at each of a plurality of outer optimization steps, obtaining a copy of the values of the trainable parameters and performing H inner optimization steps on the copy of the values of the trainable parameters, where H is at least 50, to obtain an updated set of trainable parameter values. Each inner optimization step can comprise backpropagating gradients of an objective function, evaluated using training data from a training dataset, to adjust the values of the trainable parameters.

[0091] The process can also involve, for each of the plurality of outer optimization steps, determining a cumulative gradient for at least some of the trainable parameters by determining a difference between the value of the trainable parameter in the copy obtained and the value of the parameter after the H inner optimization steps, and updating the value of the trainable parameter in the set of values of the trainable parameters using the cumulative gradient for the trainable parameter.

[0092] Other features and aspects of a method or system implemented on such a single worker may be as previously described. [0093] FIG. 4 illustrates a result of training a transformer-based large language model neural network (“Chinchilla”) with 150 million parameters on the C4 dataset, for 64,000 training steps, showing the number of training steps on the x-axis and perplexity on the y-axis. Curve 410 illustrates training the neural network from scratch, and curves 420, 430, and 440 illustrate fine tuning the neural network after pre-training for 24,000 steps. Curves 410, 420, and 430 illustrate training the neural network using conventional techniques, curve 430 using an 8 times larger batch size than curve 420; curve 430. Curve 440 training the neural network using the techniques described herein, with k = 8 workers (model replicas) and H = 500. The conventional techniques need to communicate gradients at each training step whereas the described technique communicates 500 times less (and is 8 times faster in wall clock time), and also finally reaches better generalization performance.

[0094] Other experiments show that the performance degradation is mild with increasing H, e.g. the perplexity increasing only 2.9% from H = 50 to H = 1000 despite communicating 20 times less; H = 500 appears to be a good compromise value. Experiments also show that increasing the number of workers improved generalization performance, but with diminishing returns (though no performance degradation) after k = 8. The techniques appear to work better as the neural network size increases. Use of non-i.i.d. data, e.g. where different training data shards have different data distributions, appears to have only a small impact on performance as measured by final perplexity.

[0095] FIGS. 5 to 8 relate to training the same transformer-based large language model neural network (“Chinchilla”) with 150 million parameters, on the C4 dataset.

[0096] FIG. 5 shows the number of training steps on the x-axis and perplexity on the y-axis, and illustrates the effect of the use of different outer loop optimizers, comparing SGD (stochastic gradient descent), curve 510; SGDM (SGD with momentum), curve 520; Adam, curve 530, and an optimizer using Nesterov momentum, curve 540 (with outer learning rate, T, = 0.7 and outer momentum = 0.9), illustrating that Nesterov momentum can provide improved performance (lower final perplexity).

[0097] FIG. 6 shows the number of training steps on the x-axis and perplexity on the y-axis, and illustrates that the described techniques are robust to varying the number of workers over time. In more detail, curve 600 shows conventional, local model training, and curves 610 to 650 shows examples of the techniques described herein. Curve 610 illustrates a constant number of workers (fc = 8); curve 620 illustrates doubling the number of workers (from 4 to 8) halfway through the training; curve 630 illustrates halving the number of workers (from 8 to 4) halfway through the training; curve 640 illustrates ramping up the number of workers (from 1 to 8) during the training; and curve 650 illustrates ramping down the number of workers (from 8 to 1) during the training. It can be seen, for example, that curves 620 and 630 result in final performance, as do curves 640 and 650, illustrating that while the performance depends on the total number of workers, it does not depend much on how the workers, i.e. available compute, is distributed over time.

[0098] FIG. 7 shows the number of training steps on the x-axis and perplexity on the y-axis, and illustrates that the described techniques are robust to poor/asynchronous communications. The curves all show examples of the techniques described herein, with curves 700, 710, 720 and 730 showing training for, respectively 0%, 10%, 30%, and 50% communications dropout (randomly dropping outer gradients). Even with 50% dropout, i.e. a 50% probability of dropping an outer gradient, the degradation in perplexity is only 2.1% relative to no dropout. [0099] FIG. 8 shows the number of training steps on the x-axis and perplexity on the y-axis, and illustrates that the described techniques can provide some performance benefit, even with a single worker. In more detail, curve 800 shows conventional, local model training, curves 810 shows training as described herein with 8 workers, and curve 820 shows training as described herein with a single worker (k = 1). It can be seen that even with a single worker the described technique provides faster convergence and better final performance.

[0100] There now follows a description of some large language models (LLMs) that the techniques described herein can be used to either train from scratch or fine-tune.

[0101] In general, a language model neural network is a neural network that has been trained so that, given a text prompt that includes a sequence of tokens in a natural language, the neural network can generate the next token in the sequence. This process can be repeated to extend the text prompt one token at a time to generate a natural language output, i.e., to generate the natural language output auto-regressively token by token. At each time step the language model neural network processes the current sequence to generate a probability distribution over a vocabulary of tokens. The next token can then be selected using the probability distribution, e.g., by sampling from the distribution using nucleus sampling or another sampling technique or by selecting the highest-probability token. The tokens in the vocabulary can include any of a variety of tokens, e.g., some combination of words, sub-words, characters, punctuation and other symbols, and numbers. In general, the language model neural network is trained on a corpus of text made up of tokens from the vocabulary (and optionally other tokens that can be mapped to a designated out-of-vocabulary token), to predict the next token in a sequence of tokens from the training data. [0102] A large language model neural network may have, e.g., greater than 10ⁿ trainable parameters where, e.g., n = 7, 8, 9, 10, or 11. The language model neural network can be trained on greater than 10 billion, 100 billion or 1000 billion words or tokens representing words or other text tokens, e.g., sub-words (also known as “word pieces”). There are several training datasets in the public domain; an example is the C4 dataset (Raffel et al., arXiv: 1910.10683).

[0103] Some large language model neural networks comprise an autoregressive transformer neural network; a transformer-based language model can be an encoder-only model, a decoder- only model or an encoder-decoder model. Some large language model neural network may comprise, e.g. a diffusion model or a mixture-of-experts model.

[0104] A language model neural network can be made to perform a particular task by providing a natural language description of the desired response as an input or “prompt”. Also or instead a language model neural network may be “fine-tuned” to perform a particular task, by obtaining a pre-trained language model neural network trained on a large corpus of examples and then further training part of all of the language model neural network on a relatively small number of examples particular to the type of task that is to be performed.

[0105] There now follows a description of some multimodal models, such as vision language models (VLMs), that the techniques described herein can be used to either train from scratch or fine-tune.

[0106] In general a multimodal machine learning model has a multimodal input configured to receive a first multimodal input and a second multimodal input. As used herein a “modality” refers to a type of data, and thus a multimodal machine learning model is one that can process multiple different types of data. The multimodal input to the model may comprise a mix of inputs of (at least) the first and second modality.

[0107] The first multimodal input may comprise a text input to receive a sequence of text and the second multimodal input may be configured to receive a different type of input data. As one example it may comprise a visual input to receive an image or video. As some other examples the second multimodal input may be configured to receive audio data representing values of an audio waveform, e.g. instantaneous amplitude data or time-frequency domain data; or data representing observations (not necessarily visual) of an environment with which an agent controlled by the multimodal machine learning model interacts.

[0108] The multimodal machine learning model may be configured to jointly process an encoded version of the text and an encoded version of the second multimodal input, e.g. of the image or video, to generate a model output that defines a result of a machine learning task. A few examples of machine learning tasks that can be performed by such a multimodal machine learning model are described later.

[0109] The text received may comprise text in one or more natural languages, or text in a computer language, or both. The computer language may be any formal language used to communicate with a computer, e.g. a markup language, or a command or configuration language, or a data exchange language such as JSON, or a programming language. The text may be received, e.g., as a series of encoded characters, e.g. UTF-8 encoded characters; such “characters” can include Chinese and other similar characters, as well as logograms, syllabograms and the like.

[0110] The multimodal machine learning model can include a text encoder that processes the sequence of text to represent the text as a series of text tokens from a vocabulary of text tokens, e.g. that each represent words, wordpieces or characters in a natural or computer language.

[0111] Where the second multimodal input comprises an image or video it may comprise image data defining color or intensity values for pixels of a still or moving image in one, two, or three dimensions. An image can include a LIDAR point cloud, and the image data can define the locations of points of a still or moving point cloud. As another example, an image or video received by the second multimodal input may comprise a neural 3D representation, e.g. that represents a 3D scene as a set of latent feature vectors, e.g. a neural radiance field representation. The multimodal machine learning model can include second multimodal input encoder that processes the second multimodal input, e.g. using one or more convolutional, attention, fully connected, or recurrent layers, to generate the encoded version of the second multimodal input. In general such an encoder may implement any form of encoding appropriate for the type of data to be encoded. Merely as an example, where the second multimodal input comprises an image or video this may be encoded, e.g., as features for each of a set of patches that tile the image, or as a sequence of visual tokens selected from a vocabulary of visual tokens, or as a representation of distinct objects in the visual input. Such visual tokens may, but need not be, interleaved with text tokens processed by the model.

[0112] The model output may comprise any form of output appropriate to the machine learning task performed by the multimodal machine learning model. For example the model output may comprises text in a natural or computer language that defines a result of the task, e.g. for tasks such as image captioning, visual question answering, or object detection or instance segmentation. Also or instead the model output may comprise data defining an image, video, or audio object, e.g. in a generative task; or the model output may comprise non-textual action selection data for selecting an action to be performed by an agent controlled by the model. As another example the model output may also or instead define an intermediate step to be performed during the task, e.g. a call to a software API for a software tool that is used when performing the task; the multimodal input may then receive an output from the software tool that is used to generate a final model output that performs the task. A few particular examples of model output are given later.

[0113] Such a multimodal model can be trained using very large (but possibly noisy) datasets in which text is paired with an image and/or with one or more other types of data, e.g. audio data, or data relating to the operation of an agent acting in an environment to perform a variety of tasks. Such a model is can be trained, e.g., using self-supervised learning. The pairing can often be imperfect, and the training dataset can, but may not, include any actual examples of a particular task to be performed, but nonetheless an ability to perform a particular task can emerge. There are many examples of suitable, publicly available training datasets.

[0114] Some example multimodal machine learning models with which the techniques described herein may be used include: Flamingo (Alayrac et al. arXiv:2204.14198); ALIGN (Jia et al., arXiv:2102.05918); PaLI (Chen et al. arXiv:2209.06794); and PaLLX (Chen et al. arXiv:2305.18565).

[0115] In some implementations, the second multimodal input can include an observation characterizing an environment of an agent performing a task, e.g. a mechanical agent or software agent. The observation may characterize the environment at a particular time step and the model output may define one or more actions to be performed by the agent at the time step. For example each action may be expressed as a sequence of text, e.g. as one or more characters such as letters and numbers, that represents the action, or as text that defines a low- level “skill” from a set of skills; or the model output may, e.g., define parameters of a probability distribution from which an action is selected. Optionally the text received by a text input may include text describing the task to be performed. Optionally the text input may include a description of one or more actions performed at a preceding time step. Where the agent is a software agent the model output may comprise a text output for calling a software API at a time step, and the model input at a subsequent time step, e.g. the text input, may comprise a response from the software agent, e.g. from the API.

[0116] Some examples of multimodal machine learning models controlling an agent, and with which the techniques described herein may be used, are described in: PaLM-E (Driess et al. arXiv:2303.03378); RT-1 (Brohan et al. arXiv:2212.06817); and RT-2 (Brohan et al. arXiv:2307.15818). [0117] Such a multimodal machine learning model can have an audio input, or an agent action input to receive agent action data representing an action of an agent performing a task in an environment. Data received in this way may be jointly processed with data from a text input and from a second multimodal input to generate the model output.

[0118] Such a multimodal machine learning model has a multimodal input and can, in implementations, perform a range of different tasks. However in implementations not every task that the model performs requires a multimodal input, e.g. a task to generate an image from a text description of the image, or an image captioning task. In some implementations, after training, the text input can be used to specify a particular task that is to be performed by the multimodal machine learning model, e.g. by providing a “prompt” to the model describing the task to be performed or giving an example of the task as a prompt to the model. Such prompts may optionally be included in the training data.

[0119] The multimodal machine learning model is configured to process the multimodal input in accordance with the trainable parameters of the multimodal machine learning model, to generate a model output that defines a result of one or more machine learning tasks.

[0120] In general each training data item comprises multimodal data for use in training the multimodal machine learning model, e.g. using a self-supervised training objective. As another example such a multimodal machine learning model can also or instead be trained using a reinforcement learning objective, e.g. when the model is used to control an agent to perform a task.

[0121] There are many different types of self-supervised objective function that may be used. As one example the model may be trained using a softmax cross entropy loss, e.g. using language model style teacher forcing with a softmax cross entropy loss. As another example the model may be trained with an autoregressive negative log likelihood (NLL) loss, such as — Sf=i l°gp(yj |y<j' ^x<j) for ^a multimodal input comprising a sequence of text encoded as L tokens with the Zth text token y_t conditioned on preceding second modality inputs such as one or more images or videos, and conditioned on preceding text tokens y_<t. As another example the model may be trained with a masking loss, e.g. a loss that requires the model to predict masked-out data such as masked out text tokens. As another example the multimodal machine learning model can be trained using a self-supervised objective function that comprises a contrastive loss function (one that is dependent upon a positive example and one or more negative examples). [0122] There are, similarly, many different types of reinforcement learning objective function that may be used.

[0123] Each training data item may comprise, e.g., an example sequence of text and an example of the second modality input, e.g. an example image or video; in general these are semantically related to one another (but not always, as the training dataset may be noisy). As an example, matched, text and image or video data and may be obtained from web pages, e.g. from images or videos and their corresponding alt-text (text from the HTML or XHTML alt attribute); or from web pages where images or video and text are interleaved with one another. One example of such a dataset is WebLI (Web Language Image, Chen et al. arXiv:2305.18565vl). Training datasets for other types of second modality input can similarly be obtained from web pages. Such training datasets can be large, e.g. > 10⁷, 10⁸ or 10⁹ items. [0124] Also or instead, smaller but more specialized training datasets can be used, e.g. to fine tune a model for a particular task or tasks. A few examples for visual tasks are the Visual Genome dataset for Visual Question Answering (Krishna et al., arXiv: 1602.07332); Objects365 (Shao et al., “Objects365: A large-scale, high-quality dataset for object detection”, IEEE/CVF international conference on computer vision, pages 8430-8439); Open Images V4 (Kuznetsova et al., arXiv: 1811.00982); the SBU dataset (Ordonez et al. “Im2Text: Describing Images Using 1 Million Captioned Photographs”, NeurlPS 2011); the Conceptual Captions datasets, e.g. VI (2M images) or V2 (10M images) (Sharma et al., “Conceptual Captions: A Cleaned, Hypernymed, Image Alt-text Dataset For Automatic Image Captioning”, ACL 2018); and Kinetics for video (Kay et al., arXiv: 1705.06950). An example task-specific training dataset for audio data is AudioSet (Gemmeke et al., “Audio set: An ontology and human- labeled dataset for audio events,” ICASSP, IEEE, 2017, pp. 776-780). An example taskspecific training dataset for agent (robot) control is described in Ebert et al., arXiv:2109.13396. [0125] A few examples of machine learning tasks that can be performed by a multimodal model are now described.

[0126] In general a multimodal machine learning model can be trained to perform any sort of machine learning task or tasks. After the multimodal machine learning model has been trained it can be deployed for use in performing the machine learning task(s). For instance, the machine learning model can be deployed in an environment that enables users to provide requests for the machine learning model to process specified multimodal inputs to generate corresponding model outputs. Users can provide the requests, e.g., by way of a user interface or through an application programming interface (API). The requests can be transmitted from a user device (e.g., over a data communication network, e.g., the internet) to one or more computers implementing the machine learning model, e.g., in a data center. The machine learning model can process multimodal inputs specified by user requests to generate corresponding model outputs, and then transmit the model outputs to user devices (e.g., over a data communication network).

[0127] In some implementations, after training, a particular task that is to be performed by the multimodal machine learning model can be described by part or all of the sequence of text in the multimodal input to the model. For example in a multimodal input that includes an image, video, or audio item such a prompt might specify “Generate a caption”, “Generate a description”, “Answer the following question: [about the image, video, or audio item]”, or “Detect a person”. Where the model is used for an agent control task a prompt may define “Take the knife out of the drawer”, or “Q: What action should the robot take to take the knife out of the drawer?”. Also or instead such a prompt may give one or more examples of a task to be performed. A multimodal machine learning model can be trained on multiple natural and/or computer languages and the prompt may then specify a language to use.

[0128] A few examples of some machine learning tasks that can be performed by a model trained as described herein follow.

[0129] For some tasks the second modality input represents an image or video as previously described, e.g. from a camera or other imaging device that captures the image or video from a real-world environment, and/or audio, e.g. audio data such as speech or other sounds captured from a real-world environment. In general the tasks described below may be tasks that require spatial awareness or other context from the image, video, or audio item. For example, a prompt may ask “What is the object in the top left corner?”, or “What was the answer to the spoken question?”.

[0130] As one example the task may comprise an object or action detection task. A taskspecific training data item may comprise an image, video, or audio item containing one or more objects or actions, and a sequence of text. The sequence of text may describe or otherwise label the object(s) or action(s) and (for an image or video) may include text giving bounding box coordinates for the object(s) or action(s). After training, when the model is used in inference, the model output 122 may comprise or represent text that describes or otherwise labels detected object(s) or action(s) in the second modality input, and may (for an image or video) include bounding-box coordinates for the detected object(s) or action(s), e.g. " 10 20 90 100 cat 20 30 100 100 dog”.

[0131] As another example the task may comprise a classification task, e.g. an object or action classification task. A task-specific training data item may comprise an image, video, or audio item containing one or more objects or actions and a sequence of text. The sequence of text may describe or otherwise classify the object(s) or action(s). After training, when the model is used in inference, the model output may comprise data, e.g. text, that classifies the object(s) or action(s) in the second modality input into one of a plurality of classes.

[0132] As another example the task may comprise an image, video, or audio item describing task, e.g. a captioning task (which, as used here, includes an audio description task to explain what is happening in a video). A task-specific training data item may comprise an image, video, or audio item and a sequence of text describing the image, video, or audio item. After training, when the model is used in inference, the model output may comprise data, e.g. text, describing an image, video, or audio item in the second modality input. For example the model output may provide a caption or description for a second modality input item, or it may count objects in the second modality input item, or it may provide some other form of description of the second modality input item.

[0133] As another example the task may comprise an image, video, or audio questionanswering task. A task-specific training data item may comprise an image, video, or audio item and a sequence of text that describes the image, video, or audio item. After training, when the model is used in inference, the model output may comprise data, e.g. text, that answers a question about the second modality input specified in a prompt sequence of text, e.g. as described above. This may be used, e.g., to answer questions about visual plots and charts or about sounds.

[0134] As another example the task may comprise a character or word recognition task, e.g. an OCR (optical character recognition) task. A task-specific training data item may comprise an image, video, or audio item and a sequence of text that includes text that is depicted in the image or video, or that is represented as speech in the audio item. After training, when the model is used in inference, the model output may comprise text that represents characters or words in the second modality input, e.g. in a natural language.

[0135] As another example the task may comprise a still or moving image or audio generation task. A task-specific training data item may comprise an image, video, or audio item and a sequence of text that describes the image, video, or audio item. After training, when the model is used in inference, the model output may comprise data for an image, video, or audio item, e.g. image data defining values for pixels of a still or moving image or audio data representing values of an audio waveform, and the sequence of text in the multimodal input to the model may describe or characterize the image, video, or audio item to be generated. [0136] As another example the task may comprise a computer language text generation task. A task-specific training data item may comprise an image, video, or audio item and a sequence of text in a computer language for generating the image, video, or audio item. After training, when the model is used in inference, the model output may comprise text in the or another computer language for generating or rendering an image, video, or audio item in the second modality input, e.g. a web page, plot, or chart.

[0137] In another example of a computer language text generation task a task-specific training data item may comprise an image, video, or audio item and a sequence of text in a computer language for performing a task in relation to the image, video, or audio item, e.g. a data processing task that involves analyzing the content of the image, video, or audio item to provide a result of the analysis or, e.g., a search to search for information relating to the content of the image, video, or audio item. The computer language in the model output may comprise computer language for invoking a function or calling one or more external APIs. Merely as one example, such an output may be formatted as a JSON object. As previously, the sequence of text in the multimodal input may define the task to be performed and the second modality input may comprise, e.g. an image, video, or audio item in relation to which the task is to be performed, e.g. a task that involves manipulation of particular types of data that may benefit from access to an API such as mathematical data, date/time related data, scientific data, recent data that may post-date training of the model (that may be accessed by a search function or API), and so forth. After training, when the model is used in inference, the model output may comprise text in the or another computer language for performing a task, e.g. as described above, in relation to an image, video, or audio item in the second modality input. The method may then include using the text in the computer language to perform the task.

[0138] In general where the model output comprises text this may be provided as speech representing the text.

[0139] In some implementations the machine learning task comprises an agent control task in which the agent interacts with an environment to perform the agent control task. In these implementations the multimodal input includes an observation characterizing the environment. For example the multimodal input can include a sequence of text that defines the task to be performed by the agent and the second modality input can represents an image, video, audio, or other observation of the environment, e.g. captured by a camera or other imaging device, or by a microphone, from a real-world environment. A task-specific training data item may comprise a sequence of text representing one or more actions of the agent, and a second modality input representing an observation of the environment. After training, when the model is used in inference, the model output comprises an action selection output, e.g. including text, that is used to select one or more actions to be performed by the agent in the environment in response to the observation. As an illustration the model output 122 may define an action as text such as “A: 132 114 128 5 25 156”, that can be converted into a control signal for a mechanical agent, such as a robot, e.g. “AT = [0.1, —0.2,0] A/? = [10°, 25°, —7°]”. As another example the action selection output may also or instead define one or more low-level skills, e.g. from a vocabulary of previously learnt skills. As before, the sequence of text in the multimodal input to the model may describe the task to be performed, e.g. “What action should the robot take to [perform task]”.

[0140] In some agent control implementations, the environment is a real-world environment and the agent is a mechanical agent interacting with the real-world environment, e.g., a robot or an autonomous or semi-autonomous land, air, or sea vehicle operating in or navigating through the environment, and the actions are actions taken by the mechanical agent in the real- world environment to perform the task. For example, the agent may be a robot or other mechanical agent interacting with the environment to accomplish a specific task, e.g., to locate or manipulate an object of interest in the environment or to move an object of interest to a specified location in the environment or to navigate to a specified destination in the environment. In these implementations, the observations may include, e.g., one or more of: images, object position data, and sensor data to capture observations as the agent interacts with the environment. The actions may define control signals to control the robot or other mechanical agent, e.g., positions, torques, or other control signals for the parts of the mechanical agent, or higher-level control commands.

[0141] In some agent control implementations the agent can be a software agent, i.e. a computer program, configured to perform a task. Some examples where the agent is a software agent now follow.

[0142] As one example the environment may be an integrated circuit design and the task may be a routing task for routing interconnection lines of the integrated circuit. The observations may be of component positions and/or interconnections, and the actions may comprise component placing or interconnect routing actions. An integrated circuit with interconnection lines routed as determined may then be fabricated.

[0143] As another example the environment may be a real-world computing environment and the task may be to manage the distribution of jobs or tasks across computing resources e.g. on a mobile device and/or in a data center. The observations may include observations of computing resources such as compute or memory capacity, or Internet-accessible resources, or that relate to the operation of the computing resources in processing the jobs or tasks; and the actions may include assigning jobs or tasks to particular computing resources.

[0144] As another example the environment may be a real-world computing environment and the task is to manage the processing, e.g. by one or more real-world servers, of a queue of continuously arriving jobs. The observations may comprise observations of the times of departures of successive jobs, or the time intervals between the departures of successive jobs, or the time a server takes to process each job, or the arrival times, or time intervals between the arrivals, of successive jobs, or data characterizing the type of job(s). The actions may comprise actions that allocate particular jobs to particular computing resources.

[0145] As another example the environment may comprise a real-world computer system or network and the task may be to maintain security of the computer system or network. The observations may comprise any observations characterizing operation of the computer system or network, and the actions may comprise actions to control the operation e.g. to limit or correct abnormal or undesired operation e.g. because of the presence of a virus or other security breach. [0146] As another example the environment may comprise a data packet communications network environment, and the task may be to route packets of data over the communications network. The actions may comprise data packet routing actions and the observations may comprise, e.g., observations of a routing table which includes routing metrics such as a metric of routing path length, bandwidth, load, hop count, path cost, delay, maximum transmission unit (MTU), and reliability.

[0147] In some agent control implementations the agent may be a human agent and the environment may be a real -world environment. For example the agent can be a human user of a digital assistant such as a smart speaker, smart display, or some other device that is used to instruct the user to perform actions. The task may be any real-world task that the user wishes to perform. The observations may be obtained from an observation capture subsystem, e.g. a monitoring system such as a video camera or sound capture system, to capture visual and/or audio observations of the user performing the task. The actions may comprise instructions in the form of, e.g., text, image, video, or audio data such as speech, that guide the user in performing the task.

[0148] In this specification, the term "configured" is used in relation to computing systems and environments, as well as computer program components. A computing system or environment is considered "configured" to perform specific operations or actions when it possesses the necessary software, firmware, hardware, or a combination thereof, enabling it to carry out those operations or actions during operation. For instance, configuring a system might involve installing a software library with specific algorithms, updating firmware with new instructions for handling data, or adding a hardware component for enhanced processing capabilities. Similarly, one or more computer programs are "configured" to perform particular operations or actions when they contain instructions that, upon execution by a computing device or hardware, cause the device to perform those intended operations or actions.

[0149] The embodiments and functional operations described in this specification can be implemented in various forms, including digital electronic circuitry, software, firmware, computer hardware (encompassing the disclosed structures and their structural equivalents), or any combination thereof. The subject matter can be realized as one or more computer programs, essentially modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by or to control the operation of a computing device or hardware. The storage medium can be a storage device such as a hard drive or solid-state drive (SSD), a storage medium, a random or serial access memory device, or a combination of these. Additionally or alternatively, the program instructions can be encoded on a transmitted signal, such as a machine-generated electrical, optical, or electromagnetic signal, designed to carry information for transmission to a receiving device or system for execution by a computing device or hardware. Furthermore, implementations may leverage emerging technologies like quantum computing or neuromorphic computing for specific applications, and may be deployed in distributed or cloud-based environments where components reside on different machines or within a cloud infrastructure.

[0150] The term "computing device or hardware" refers to the physical components involved in data processing and encompasses all types of devices and machines used for this purpose. Examples include processors or processing units, computers, multiple processors or computers working together, graphics processing units (GPUs), tensor processing units (TPUs), and specialized processing hardware such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs). In addition to hardware, a computing device or hardware may also include code that creates an execution environment for computer programs. This code can take the form of processor firmware, a protocol stack, a database management system, an operating system, or a combination of these elements. Embodiments may particularly benefit from utilizing the parallel processing capabilities of GPUs, in a General-Purpose computing on Graphics Processing Units (GPGPU) context, where code specifically designed for GPU execution, often called kernels or shaders, is employed. Similarly, TPUs excel at running optimized tensor operations crucial for many machine learning algorithms. By leveraging these accelerators and their specialized programming models, the system can achieve significant speedups and efficiency gains for tasks involving artificial intelligence and machine learning, particularly in areas such as computer vision, natural language processing, and robotics.

[0151] A computer program, also referred to as software, an application, a module, a script, code, or simply a program, can be written in any programming language, including compiled or interpreted languages, and declarative or procedural languages. It can be deployed in various forms, such as a standalone program, a module, a component, a subroutine, or any other unit suitable for use within a computing environment. A program may or may not correspond to a single file in a file system and can be stored in various ways. This includes being embedded within a file containing other programs or data (e.g., scripts within a markup language document), residing in a dedicated file, or distributed across multiple coordinated files (e.g., files storing modules, subprograms, or code segments). A computer program can be executed on a single computer or across multiple computers, whether located at a single site or distributed across multiple sites and interconnected through a data communication network. The specific implementation of the computer programs may involve a combination of traditional programming languages and specialized languages or libraries designed for GPGPU programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics.

[0152] In this specification, the term "engine" broadly refers to a software-based system, subsystem, or process designed to perform one or more specific functions. An engine is typically implemented as one or more software modules or components installed on one or more computers, which can be located at a single site or distributed across multiple locations. In some instances, one or more dedicated computers may be used for a particular engine, while in other cases, multiple engines may operate concurrently on the same one or more computers. Examples of engine functions within the context of Al and machine learning could include data pre-processing and cleaning, feature engineering and extraction, model training and optimization, inference and prediction generation, and post-processing of results. The specific design and implementation of engines will depend on the overall architecture and the distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors.

[0153] The processes and logic flows described in this specification can be executed by one or more programmable computers running one or more computer programs to perform functions by operating on input data and generating output. Additionally, graphics processing units (GPUs) and tensor processing units (TPUs) can be utilized to enable concurrent execution of aspects of these processes and logic flows, significantly accelerating performance. This approach offers significant advantages for computationally intensive tasks often found in Al and machine learning applications, such as matrix multiplications, convolutions, and other operations that exhibit a high degree of parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speedups and efficiency gains compared to relying solely on CPUs can be achieved. Alternatively or in combination with programmable computers and specialized processors, these processes and logic flows can also be implemented using specialized processing hardware, such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs), for even greater performance or energy efficiency in specific use cases.

[0154] Computers capable of executing a computer program can be based on general-purpose microprocessors, special-purpose microprocessors, or a combination of both. They can also utilize any other type of central processing unit (CPU). Additionally, graphics processing units (GPUs), tensor processing units (TPUs), and other machine learning accelerators can be employed to enhance performance, particularly for tasks involving artificial intelligence and machine learning. These accelerators often work in conjunction with CPUs, handling specialized computations while the CPU manages overall system operations and other tasks. Typically, a CPU receives instructions and data from read-only memory (ROM), random access memory (RAM), or both. The essential elements of a computer include a CPU for executing instructions and one or more memory devices for storing instructions and data. The specific configuration of processing units and memory will depend on factors like the complexity of the Al model, the volume of data being processed, and the desired performance and latency requirements. Embodiments can be implemented on a wide range of computing platforms, from small embedded devices with limited resources to large-scale data center systems with high-performance computing capabilities. The system may include storage devices like hard drives, SSDs, or flash memory for persistent data storage.

[0155] Computer-readable media suitable for storing computer program instructions and data encompass all forms of non-volatile memory, media, and memory devices. Examples include semiconductor memory devices such as read-only memory (ROM), solid-state drives (SSDs), and flash memory devices; hard disk drives (HDDs); optical media; and optical discs such as CDs, DVDs, and Blu-ray discs. The specific type of computer-readable media used will depend on factors such as the size of the data, access speed requirements, cost considerations, and the desired level of portability or permanence. [0156] To facilitate user interaction, embodiments of the subject matter described in this specification can be implemented on a computing device equipped with a display device, such as a liquid crystal display (LCD) or an organic light-emitting diode (OLED) display, for presenting information to the user. Input can be provided by the user through various means, including a keyboard), touchscreens, voice commands, gesture recognition, or other input modalities depending on the specific device and application. Additional input methods can include acoustic, speech, or tactile input, while feedback to the user can take the form of visual, auditory, or tactile feedback. Furthermore, computers can interact with users by exchanging documents with a user's device or application. This can involve sending web content or data in response to requests or sending and receiving text messages or other forms of messages through mobile devices or messaging platforms. The selection of input and output modalities will depend on the specific application and the desired form of user interaction.

[0157] Machine learning models can be implemented and deployed using machine learning frameworks, such as TensorFlow or JAX. These frameworks offer comprehensive tools and libraries that facilitate the development, training, and deployment of machine learning models. [0158] Embodiments of the subject matter described in this specification can be implemented within a computing system comprising one or more components, depending on the specific application and requirements. These may include a back-end component, such as a back-end server or cloud-based infrastructure; an optional middleware component, such as a middleware server or application programming interface (API), to facilitate communication and data exchange; and a front-end component, such as a client device with a user interface, a web browser, or an app, through which a user can interact with the implemented subject matter. For instance, the described functionality could be implemented solely on a client device (e.g., for on-device machine learning) or deployed as a combination of front-end and back-end components for more complex applications. These components, when present, can be interconnected using any form or medium of digital data communication, such as a communication network like a local area network (LAN) or a wide area network (WAN) including the Internet. The specific system architecture and choice of components will depend on factors such as the scale of the application, the need for real-time processing, data security requirements, and the desired user experience.

[0159] The computing system can include clients and servers that may be geographically separated and interact through a communication network. The specific type of network, such as a local area network (LAN), a wide area network (WAN), or the Internet, will depend on the reach and scale of the application. The client-server relationship is established through computer programs running on the respective computers and designed to communicate with each other using appropriate protocols. These protocols may include HTTP, TCP/IP, or other specialized protocols depending on the nature of the data being exchanged and the security requirements of the system. In certain embodiments, a server transmits data or instructions to a user's device, such as a computer, smartphone, or tablet, acting as a client. The client device can then process the received information, display results to the user, and potentially send data or feedback back to the server for further processing or storage. This allows for dynamic interactions between the user and the system, enabling a wide range of applications and functionalities.

[0160] 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.

[0161] 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.

[0162] 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.

Claims

1. A computer-implemented method of training a large neural network, using a distributed computing system, to optimize an objective function; the distributed computing system comprising a plurality of computing devices being coupled by a computer network and operating in parallel, each computing device having associated memory; the large neural network having at least 10⁷ trainable parameters; the method using: a training dataset comprising a plurality of shards, each stored in the memory of one of the computing devices; a shared set of values of the trainable parameters maintained in a shared data store; and a set of workers implemented on the plurality of computing devices, each worker being implemented on one of the computing devices, each worker being allocated a respective one of the plurality of shards, and each worker configured to, at each of a plurality of outer optimization steps: obtain a local copy of the values of the trainable parameters from the shared data store; and perform H inner optimization steps on the local copy of the values of the trainable parameters, where H is at least 50, to obtain an updated set of trainable parameter values for the worker after the H inner optimization steps, each inner optimization step comprising backpropagating gradients of the objective function, evaluated using the shard allocated to the worker, to adjust the values of the local copy of the trainable parameters; the method comprising, for each of the plurality of outer optimization steps: for each worker, determining a cumulative gradient for at least some of the trainable parameters by determining a difference between the value of the trainable parameter in the local copy obtained by the worker and the value of the parameter after the H inner optimization steps; and determining, at an optimization computer system comprising one of the computing devices or a shared server coupled to the computer network, an aggregated cumulative gradient by, for each trainable parameter of at least some of the trainable parameters aggregating the cumulative gradients for the trainable parameter determined for the workers; and updating the value of the trainable parameter in the shared data store using the aggregated cumulative gradient for the trainable parameter; wherein the values of the trainable parameters in the shared data store provide the trained parameters for the large neural network.

2. A computer-implemented method of training a large neural network, using a distributed computing system, to optimize an objective function; the distributed computing system comprising a plurality of computing devices being coupled by a computer network and operating in parallel, each computing device having associated memory; the large neural network having at least 10⁷ trainable parameters; the method comprising: obtaining a training dataset comprising a plurality of shards, each stored in the memory of one of the computing devices; maintaining a shared set of values of the trainable parameters in a shared data store; and implementing a set of workers on the plurality of computing devices, each worker being implemented on one of the computing devices, each worker being allocated a respective one of the plurality of shards, and each worker, at each of a plurality of outer optimization steps: obtaining a local copy of the values of the trainable parameters from the shared data store; and performing H inner optimization steps on the local copy of the values of the trainable parameters, where H is at least 50, to obtain an updated set of trainable parameter values for the worker after the H inner optimization steps, each inner optimization step comprising backpropagating gradients of the objective function, evaluated using the shard allocated to the worker, to adjust the values of the local copy of the trainable parameters; the method further comprising, for each of the plurality of outer optimization steps: for each worker, determining a cumulative gradient for at least some of the trainable parameters by determining a difference between the value of the trainable parameter in the local copy obtained by the worker and the value of the parameter after the H inner optimization steps; and determining, at an optimization computer system comprising one of the computing devices or a shared server coupled to the computer network, an aggregated cumulative gradient by, for each trainable parameter of at least some of the trainable parameters aggregating the cumulative gradients for the trainable parameter determined for the workers; and updating the value of the trainable parameter in the shared data store using the aggregated cumulative gradient for the trainable parameter; wherein the values of the trainable parameters in the shared data store provide the trained parameters for the large neural network.

3. The method of claim 1 or 2, wherein for each worker, determining the cumulative gradient for at least some of the trainable parameters, comprises the worker determining the cumulative gradient for each of the trainable parameters; the method further comprising: communicating the cumulative gradient for the at least some of the trainable parameters, over the computer network, from each worker to the optimization computer system.

4. The method of claim 3, further comprising: selecting, at a worker, a subset of the cumulative gradients from the cumulative gradient for each of the trainable parameters for the worker, and communicating the subset of the cumulative gradients over the computer network, from the worker to the optimization computer system.

5. The method of claim 1 or 2, comprising: each worker communicating at least some of the trainable parameters of the updated set of trainable parameter values for the worker, from the worker to optimization computer system over the computer network; and wherein for each worker, determining the cumulative gradient for at least some of the trainable parameters, comprises the optimization computer system determining the cumulative gradient for at least some of the trainable parameters.

6. The method of claim 5, further comprising: selecting, at a worker, a subset of trainable parameters from the updated set of trainable parameter values for the worker, and communicating the subset of the trainable parameters over the computer network, from the worker to the optimization computer system.

7. The method of any one of claims 1-6, wherein each worker backpropagating gradients of the objective function, evaluated using the shard allocated to the worker, to adjust the values of the local copy of the trainable parameters comprises, for each of the trainable parameters: maintaining a first momentum variable for the trainable parameter locally at the worker, wherein the first momentum variable comprises a moving average of a gradient for adjusting the trainable parameter, and adjusting the value of the trainable parameter by an offset dependent on a value of the first momentum variable; and wherein updating the value of the trainable parameter in the shared data store using the aggregated cumulative gradient for the trainable parameter comprises: maintaining a second momentum variable for the trainable parameter in the shared data store, determining an updated value for the second momentum variable from the aggregated cumulative gradient for the trainable parameter, and determining an updated value of the trainable parameter in the shared data store from the updated value of the second momentum variable.

8. The method of claim 7, further comprising determining the updated value for the second momentum variable from a current value of the trainable parameter and the aggregated cumulative gradient for the trainable parameter, and determining the updated value of the trainable parameter from a difference between the updated value for the second momentum variable and a current value of the second momentum variable.

9. The method of claim 7 or 8, wherein determining the updated value for the second momentum variable from the aggregated cumulative gradient for the trainable parameter comprises scaling the aggregated cumulative gradient by an outer optimization learning rate parameter, the method further comprising keeping the outer optimization learning rate parameter constant.

10. The method of any preceding claim, wherein the large neural network comprises a Transformer neural network characterized by having a succession of attention neural network layers, wherein each attention neural network layer has an attention layer input for each element of the input and is configured to apply an attention mechanism over the attention layer input to generate an attention layer output for each element of the input, wherein the attention mechanism computes a similarity between a query and a set of key -value pairs, wherein one or both of the query and the set of key -value pairs are determined from the attention layer input.

11. The method of any preceding claim, wherein the large neural network comprises a large language model (LLM) or a large vision language model (VLM) neural network; wherein adjusting the values of the local copy of the trainable parameters in the inner optimization steps uses the AdamW optimization algorithm; and wherein, in the outer optimization steps, updating the value of the trainable parameter in the shared data store using the aggregated cumulative gradient for the trainable parameter uses the Nesterov optimization algorithm.

12. The method of any preceding claim, wherein aggregating the cumulative gradients for the trainable parameter determined for the workers comprises determining a weighted average wherein the cumulative gradient for each worker is weighted by a number of training data items in the shard for the worker.

13. The method of any preceding claim, wherein determining the aggregated cumulative gradient for a trainable parameter at the optimization computer system is performed without waiting to obtain an updated trainable parameter value or cumulative gradient for the trainable parameter from all of the workers.

14. The method of any preceding claim, wherein the optimization computer system comprises a shared server coupled to the computer network.

15. The method of any preceding claim, wherein the computer network comprises a wide area network (WAN).

16. A computer-implemented method of training a large neural network, the large neural network having at least 10⁷ trainable parameters; the method comprising: maintaining a set of values of the trainable parameters; and at each of a plurality of outer optimization steps: obtaining a copy of the values of the trainable parameters; and performing H inner optimization steps on the copy of the values of the trainable parameters, where H is at least 50, to obtain an updated set of trainable parameter values after the H inner optimization steps, each inner optimization step comprising: backpropagating gradients of an objective function, evaluated using training data from a training dataset, to adjust the values of the trainable parameters; the method further comprising, for each of the plurality of outer optimization steps: determining a cumulative gradient for at least some of the trainable parameters by determining a difference between the value of the trainable parameter in the copy obtained and the value of the parameter after the H inner optimization steps; and updating the value of the trainable parameter in the set of values of the trainable parameters using the cumulative gradient for the trainable parameter.

17. One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations of the respective method of any one of claims 1-16.

18. A system comprising: one or more computers; and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective method of any one of claims 1-16.