Distributed training of large neural networks

Abstract

Systems and methods for using distributed computing systems to train large neural networks to perform machine learning tasks. A shared set of trainable parameters is maintained in a shared data store, and each set of geographically distributed workers updates their trainable parameters using a sharded training dataset. There are two optimization processes: an external optimization process and an internal optimization loop that runs tens, hundreds, or thousands of times independently and in parallel by each worker. Workers may have different computing capabilities and may be geographically far apart from each other, and the communication bandwidth used by the system may be two or three orders of magnitude smaller than that of other systems.

Description

Technology Field

[0001] Cross-reference regarding related applications

[0002] 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 construed as part of this application and is incorporated by reference in the disclosure of this application. Background Technology

[0003] This specification relates to training neural networks.

[0004] A neural network is a machine learning model that employs one or more layers of non-linear units to predict an output for a received input. The neural network described herein includes one or more hidden layers in addition to an output layer. The output of each hidden layer is used as an input to the next layer of the network, that is, the next hidden layer or output layer. Each layer of the neural network generates an output from a received input according to current values ​​of an individual set of trainable parameters, e.g., weights.

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

[0006] In one embodiment, a distributed computing system is used to provide a neural network, particularly a large neural network, e.g., at least 10 7 A computer implementation method for training a neural network having trainable parameters (e.g., weights) is described. A distributed computing system includes multiple computing devices that are combined by a computer network and operate in parallel.

[0007] The method includes the step of maintaining a shared set of values ​​of trainable parameters in a shared data store (combined to a computer network). A set of workers, typically a set of geographically distributed workers, is implemented on multiple computing devices, and each worker is implemented on one of the devices. Each worker is assigned one of multiple shards of a training dataset (stored in the memory of the computing device). A worker may include a computer, or a computer having one or more machine learning hardware accelerators (e.g., performing matrix multiplication in hardware), or a cluster of computers.

[0008] Each worker acquires a local copy of the trainable parameter values ​​and performs at least H=50, 100, 500, 1000, or 2000 internal optimization steps on them to obtain an updated set of trainable parameter values. The worker acquires the updated set of trainable parameter values ​​by training an instance of a large neural network using shards of the training dataset. Generally, training may involve backpropagating the gradients of the objective function to adjust the values ​​of the local copies of the trainable parameters to obtain an updated set of trainable parameter values ​​for the worker (for any machine learning task being trained to perform the neural network).

[0009] In each of the multiple external optimization steps, a cumulative gradient for some or all of the trainable parameters is determined for each worker, for example, by each worker. The cumulative gradient for the parameters may be determined using the values ​​of the trainable parameters in the local copy obtained by the worker, i.e., the updated trainable parameter values ​​determined by the worker after H internal optimization steps. More specifically, this may involve determining the difference between the values ​​of the trainable parameters in the local copy obtained by the worker from the shared data store and the values ​​of the parameters after H internal optimization steps.

[0010] The method also includes a step of determining an aggregated cumulative gradient in an optimization computer system. The optimization computer system may be, for example, one of computing devices or a shared server connected to a computer network.

[0011] In implementations, the aggregated cumulative gradient is determined for each of some or all of the trainable parameters, for example, by aggregating the cumulative gradients for the trainable parameters determined by each worker. Then, the aggregated cumulative gradient for the trainable parameters can be used to update the values ​​of the trainable parameters in a shared data store.

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

[0013] The techniques described herein are suitable for implementation in environments where some or all of the computing devices are geographically separated, for example, where the computer network includes a Wide Area Network (WAN) and communication bandwidth is relatively limited. However, these techniques may also be implemented on a single worker, for example, for faster convergence and better generalization.

[0014] In a related aspect, the corresponding process is described from the perspective of an optimized computer system.

[0015] The subject matter of the invention described in this specification may be implemented in specific embodiments to realize one or more of the following advantages.

[0016] Conventional approaches for training large neural networks on distributed computing systems require a large number of tightly interconnected hardware accelerators, such as GPUs or TPUs. This is partly because existing algorithms require high-bandwidth communication between devices that need to exchange data with each other during training. For example, a typical synchronous system for training large neural networks can involve thousands of juxtaposed devices interconnected with high-bandwidth cables and careful software engineering to coordinate the devices and achieve high utilization.

[0017] In principle, more widely distributed computing systems may be used, but when training large neural networks, for example, very high-bandwidth communications will be required to transmit many weights or gradients. Implementations of the techniques described herein can reduce the communication required between workers or islands of workers implemented on each computing device by two or three orders of magnitude compared to these conventional approaches.

[0018] The techniques described here also facilitate a much more distributed approach—for example, workers may be physically separated from each other and connected by computer networks such as a local area network (LAN) or a wide area network (WAN). With the described techniques, the islands of workers or locally connected workers may be located in different countries.

[0019] There is no requirement for workers to have the same computing power or memory capacity, and workers may have different computing and / or memory capabilities and / or capabilities that vary over time. Nevertheless, in the embodiments, workers may be maintained to be nearly fully utilized.

[0020] Implementations of the techniques described herein are also highly robust to ranges of different factors. For example, they can resolve situations where parameters need to be rebooted and reloaded from the point where the task interruption ended when a worker stalls or an error occurs. They can also handle varying numbers of workers during the neural network training process and do not require different workers to have the same distribution of training data. These techniques are also robust to interruptions in communications that result in lost packets, for example.

[0021] These techniques can be used, for example, to train a neural network from scratch starting with random initialization or to fine-tune a previously trained neural network. The techniques described herein can also train large neural networks faster than conventional techniques in wall-clock time.

[0022] Neural networks trained using the techniques described here can achieve better performance than some synchronous training techniques, for example, lower complexity or better generalization. At inference time, the resulting model can have the same size and speed as a model trained in fully synchronous mode.

[0023] Details of one or more embodiments of the subject matter of this specification are presented in the accompanying drawings and the following description. Other features, aspects, and advantages of the subject matter of this specification will become apparent from the description, drawings, and claims. Brief explanation of the drawing

[0024] Figure 1 illustrates an exemplary distributed computing system for training a large neural network. Figure 2 is a flowchart of an exemplary process for training a large neural network using a distributed computing system. Figure 3 visualizes an example of the described training process. Figure 4 compares the results of an example of the described training process with the prior art. Figure 5 illustrates the results of an example of a training process described using Nesterov momentum. Figure 6 illustrates the results of an example of a training process described using various numbers of workers. Figure 7 illustrates the results of examples of training processes described using various communication qualities. Figure 8 illustrates the results of an example of the described training process implemented for a single worker. In various drawings, the same reference number and name represent the same element. Specific details for implementing the invention

[0025] Techniques for training neural networks, particularly large neural networks, to perform machine learning tasks are described. Some embodiments of these techniques use distributed computing systems having a number of workers operating in parallel. As used herein, "training" includes fine-tuning a pre-trained model.

[0026] Large neural networks are at least 10 7It can have trainable parameters; in some implementations, at least 10 8 Dog or 10 9 It can have several trainable parameters. Generally, a neural network is configured to process neural network inputs according to the neural network's trainable parameters to generate neural network outputs that define the results of machine learning tasks. For example, a neural network may include so-called foundation models, such as a large language model (LM) or a vision language model (VLM).

[0027] In some implementations, a shared set of trainable parameters is maintained in a shared data store, and each geographically distributed set of workers updates its trainable parameters using a shard training dataset. There are two optimization processes: an external optimization process, and an internal optimization loop that runs tens, hundreds, or thousands of times in parallel and typically independently by each worker. Workers may have different computing capabilities and may be geographically far apart from each other; the communication bandwidth used by the system may be two or three orders of magnitude smaller than that of other systems.

[0028] FIG. 1 illustrates an example of a distributed computing system (100) for training a large neural network.

[0029] A distributed computing system (100) includes a plurality of computing devices (110a..n), each having an associated (local) memory (114a..n). The computing devices may include, as some examples, a general-purpose computer system, a computer equipped with one or more GPUs (Graphics Processing Units), a computer equipped with one or more TPUs (Tensor Processing Units), or a combination thereof.

[0030] Computing devices (110) can be combined by a shared computer network (120) and operate in parallel. The computer network (120) may include wired and / or wireless parts and may include, for example, a local area network (LAN), a wide area network (WAN), or the Internet. The advantage of the described techniques is that they are suitable for implementation in environments where some or all of the computing devices are geographically far apart from each other. Accordingly, in some embodiments, the computer network (120) may include a wide area network (WAN).

[0031] Each computing device (110) is configured to implement one or more workers (112a..n), for example, one worker may be implemented on each computing device. Generally, a worker may include code for a software process running on one of the computing devices (110), for example, using multiple GPUs or TPUs of the computing device.

[0032] The distributed computing system (100) also includes a shared data store (130) accessible via a computer network (120), for example, on one of the shared, for example, central server or computing device (110). The shared data store (130) is used to maintain a shared set of values ​​of trainable parameters of a large neural network, as described below.

[0033] The distributed computing system (100) also includes an optimization computer system (150). The optimization computer system may include one of the computing devices (110) or a shared server connected to a computer network.

[0034] FIG. 2 is a flowchart of an exemplary process for training a neural network, particularly a large neural network, using a distributed computing system such as the distributed computing system (100) of FIG. 1. For convenience, this process will be described with reference to the distributed computing system (100).

[0035] A neural network may be randomly initialized, meaning that the initial values ​​of the trainable parameters are randomly selected, or it may include a pre-trained model. Large neural networks can be trained to optimize an objective function. The objective function can be any objective function appropriate for the task to be performed by the (trained) neural network.

[0036] The process includes the step of acquiring a training dataset (step 200). The training dataset comprises a plurality of shards, and each shard is stored in the memory (114) of one of the computing devices (110), for example. Each shard is typically a separate partition of the training dataset, but in principle, some overlap or replication between data within different shards is possible. A sharded training dataset can be acquired by dividing the training dataset into shards. For example, the training dataset can be divided into several parts or shards, each of which holds an individual part of the training dataset and together constitutes a complete training dataset.

[0037] The process maintains a shared set of values ​​of trainable parameters of a large neural network in a shared data store (130). A set of workers (112) is implemented on multiple computing devices (110) (step 202). In particular, each worker is implemented on one of the computing devices, for example, each worker may be implemented on a different individual computing device among the computing devices, and each worker is assigned an individual shard among multiple shards.

[0038] As one example, each worker may be allocated an individual shard among the shards by a shared, for example, central server. As another example, each worker may acquire their own shard without being allocated a shard by any form of central control, which can support privacy protection, for example, by preventing the central server from accessing one or more shards. For example, one or more workers may acquire a portion of the training data without exposing it to any of the other workers or to the central server (if any); this portion of the training data may constitute a shard of the training dataset. There are known techniques to achieve this, such as coherent hashing (a distributed approach where workers can acquire shards based on a hash function), and other hash-based techniques.

[0039] In some embodiments, there is one worker and one shard per computing device; in some embodiments, 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 entail data and / or model parallelism.

[0040] Generally, the present method includes the step of performing a plurality of external optimization steps. In each of the plurality of external optimization steps, each worker (112) performs a plurality of internal optimization steps, generally independently of each other worker and in parallel with each other worker.

[0041] Each external optimization step may include a step in which each worker (112) obtains a local copy of the values ​​of the trainable parameters from the shared data store (130) (step 204). More specifically, a computing device (110) implementing the worker (112) may obtain a local copy of the values ​​of the trainable parameters from the shared data store (130) and store them in the (local) memory (114) of the computing device.

[0042] Each external optimization step may also include a step of performing H internal optimization steps on local copies of the values ​​of the trainable parameters so that each worker (112) obtains an updated set of trainable parameter values ​​for the worker after H internal optimization steps (step 206). Here, H is at least 50, 100, 500, 1000, or 2000.

[0043] In the embodiments, each internal optimization step includes a step of backpropagating the gradients of the objective function evaluated using the shards assigned to the worker to adjust the values ​​of local copies of the trainable parameters (step 208).

[0044] The method also includes the step of determining a cumulative gradient (or "external" gradient) for at least some of the trainable parameters for each of the plurality of external optimization steps and for each worker (112) (step 210).

[0045] The cumulative gradient is determined by determining the 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 H internal optimization steps. That is, the cumulative gradient for the trainable parameter depends on the difference between the value of the parameter before and after H internal optimization steps.

[0046] For example, the cumulative gradient is It can be determined as, where represents the set of updated trainable parameter values ​​for the worker after H internal optimization steps, and represents the values ​​of trainable parameters obtained by the worker from a shared data store, i.e., prior to H internal optimization steps. The cumulative gradient can be determined by the worker, i.e., locally, or (if the updated trainable parameter values ​​are communicated elsewhere), or the cumulative gradient can be determined by one of the other workers or centrally (the term "central" here specifically refers to a function rather than a geographical location).

[0047] Each of the multiple external optimization steps also includes a step for determining an aggregated cumulative gradient (step 212). The aggregated cumulative gradient may be an aggregation of cumulative gradients determined by the workers, for example, an average.

[0048] The aggregated cumulative gradient can be determined in an optimization computer system (150) that may include one of the computing devices (110) or a shared server, as previously mentioned.

[0049] Determining the aggregated cumulative gradient may include aggregating, for example, averaging, the cumulative gradients for each of at least some of the trainable parameters, for the workers (112) or for the trainable parameters determined by the workers. As previously mentioned, the cumulative gradient of the trainable parameters may be determined by the workers (by the workers) or for the workers, for example, by a central server.

[0050] As an example, aggregated cumulative ("external") gradient can be the average cumulative gradient. This average cumulative gradient is, for example, It can be determined by evaluating, and here represents the set of updated trainable parameter values ​​for worker i among k workers after H internal optimization steps, and represents the values ​​of trainable parameters obtained from the shared data store, i.e., prior to H internal optimization steps. That is, the aggregated, for example, average cumulative gradient (per parameter) can be determined from the difference between the trainable parameter values ​​for the current and previous outer loop optimization time steps.

[0051] Optionally, for example, if shards for different workers contain different numbers of training data items, the cumulative gradient for each worker is n, the number of training data items within the shard for the worker. k It can be weighted by. For example, the average cumulative gradient is, for example, a weighted sum It may be the weighted average cumulative gradient determined by evaluating, where w k is n k It represents a weight that can be determined as / n, where n is the total number of training data items in the training dataset.

[0052] Generally, determining the aggregated cumulative gradient for a parameter involves communicating from the worker (112) to an entity performing averaging, e.g., an optimization computer system (150). This may involve communicating 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)).

[0053] As described below, in some embodiments, communication from the workers is pruned so that not all cumulative gradients (or updated parameter values) are transmitted, for example, to the optimization computer system. For example, a threshold may be applied so that only cumulative gradients having magnitudes greater than a threshold are transmitted. Different workers may transmit cumulative gradients for different groups of trainable parameters. Thus, in some embodiments, the aggregated cumulative gradient may be determined for only some of the trainable parameters. For (one of) these trainable parameters, determining the aggregated cumulative gradient may involve aggregating only the cumulative gradients (for the trainable parameter) from some of the workers.

[0054] Each of the multiple external optimization steps also includes a step of updating the values ​​of the trainable parameters in the shared data store, that is, the values ​​of the trainable parameters in the set of shared trainable parameter values ​​(step 214).

[0055] Generally, this involves using an aggregated cumulative gradient for the trainable parameter to determine the value of the trainable parameter. In some implementations, this may include applying a parameter offset that depends on an aggregated, e.g., mean, cumulative gradient for the trainable parameter to a value that depends on the current value of the trainable parameter.

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

[0057] Implementations of the aforementioned techniques can reduce communication between workers, for example, between a central server by a factor of several. Nevertheless, there may be a heavy load on the network, particularly when workers are synchronized, for large models and / or a large number of workers, for example, when H sets of internal optimization steps are completed.

[0058] In some implementations, this can be mitigated by reducing the communication bandwidth from the workers required at these points, and in particular by pruning multiple data items (cumulative gradients or parameter values) sent from the workers to the entity, e.g., the server, and performing cumulative gradient aggregation.

[0059] In some embodiments, the value of H may differ for different workers, and, for example, by making communication to / from the workers and / or to / from the central server time-lag, the workers may be asynchronous and the surge in communication bandwidth may be mitigated.

[0060] In some embodiments, the value of H for a worker may be selected or modified according to the characteristics of the worker or the characteristics of the available communication bandwidth to or from the worker. The characteristics of the worker may be, for example, a metric of the available computing power or memory capacity or its computing speed. The characteristics of communication may be, for example, a metric of the average or peak bandwidth capacity of the network connection to or from the worker. In this way, the distributed computing system may adapt to, for example, changing network or worker capacities or conditions.

[0061] In some implementations, each worker determines the cumulative gradient for each (all) of the trainable parameters. Then, the worker may communicate the cumulative gradient for some or all of the trainable parameters to the optimization computer system. In these implementations, the worker may select a subset to communicate to the optimization computer system, more specifically, an appropriate subset (i.e., less than all) of the cumulative gradients. The selected subset may differ for different workers. When determining the aggregated cumulative gradient, if the cumulative gradient for a trainable parameter has not been communicated by the worker, a value of 0 may be given.

[0062] A subset of cumulative gradients to be communicated can be selected by the worker in a number of ways. For example, the cumulative gradients to be communicated can be selected by comparing each cumulative gradient to a threshold and communicating only those cumulative gradients that have a value greater than the threshold. The threshold is, for example, the top x% (0) of the cumulative gradients. <x<100)를 선택하도록 선택될 수 있다. 또한 또는 대신에, 통신할 누적 그래디언트들이 그들의 부호에 의해 선택될 수도 있다. 보다 구체적으로, 누적 그래디언트들을 각각 포지티브 및 네거티브 부호와 합산함으로써 포지티브 및 네거티브 누적 그래디언트들에 대해 포지티브 및 네거티브 집계된 부호 값이 결정될 수 있고; 그 다음 가장 큰 크기를 갖는 집계된 부호 값의 부호를 갖는 누적 그래디언트들만이 통신될 수 있다.

[0063] In some implementations, instead of communicating cumulative gradients, each worker may communicate some or all of the trainable parameters of an updated set of trainable parameter values ​​for the worker from the worker to the optimization computer system via a computer network. The optimization computer system may then determine the cumulative gradient for the communicated trainable parameters. This may also include, at the worker, selecting a subset of trainable parameters from the updated set of trainable parameter values ​​for the worker, and communicating only the selected subset of trainable parameters. Again, the trainable parameters may be selected in a number of ways, for example, by comparing each to a threshold. Again, the selected subset may differ for different workers.

[0064] In embodiments, each internal optimization step involves the worker (112) adjusting (updating) the values ​​of local copies of the trainable parameters using stochastic gradient descent, particularly stochastic gradient descent using momentum. This may include maintaining a first momentum variable for each trainable parameter locally in the worker. In embodiments, the first momentum variable includes a (damped) moving average of the gradient used to adjust the trainable parameters, e.g., an exponential moving average. Then, the value of the trainable parameter may be adjusted by an offset that depends on the value of the first momentum variable.

[0065] In certain examples, stochastic gradient descent using momentum may include the Adam optimization algorithm or AdamW (Loshchilov et al., 2019). Adam and AdamW are used to adapt the learning rate while maintaining both the first momentum variable and the second momentum variable, where the second momentum variable depends on the square of the gradient.

[0066] In the embodiments, after H > 50 inner optimization steps, aggregating the cumulative gradient at each outer optimization step also uses a momentum-based approach, but differs from that used in the inner optimization steps for stability. This may include maintaining a second momentum variable (different from the first momentum variable) for each trainable parameter in a shared data store (130), and using it to adjust individual trainable parameters, which is effectively achieved by adjusting the relative contributions of previously determined aggregated, e.g., averages, cumulative gradients when updating the values ​​of the trainable parameters.

[0067] In particular, updating the value of a trainable parameter in a shared data store using an aggregated cumulative gradient for the trainable parameter may include maintaining a second momentum variable for the trainable parameter in the shared data store (130). The updated value for the second momentum variable is determined from the aggregated, for example, average cumulative gradient for the trainable parameter. Then, the value of the updated trainable parameter in the shared data store is determined from the updated value of the second momentum variable. In embodiments, the updated value of the second momentum variable is scaled by a factor β that weights the effect of the momentum (e.g., in the range [0, 1]).

[0068] More specifically, the updated value for the second momentum variable can be determined from the current value of the trainable parameter and the aggregated, e.g., mean cumulative gradient for the trainable parameter. Optionally, the aggregated, e.g., mean cumulative gradient can be scaled by an external optimization learning rate parameter (η). As an example, the updated value for the second momentum variable ( )silver It can be determined by summing and can be optionally scaled ( ), for example This can be the case. In some implementations, the external optimization learning rate parameter may remain constant without decay, i.e., only the internal optimization learning rate may decay.

[0069] The value of the updated trainable parameter is the updated value for the second momentum variable ( )(optionally scaled) and the current value of the second momentum variable( It can be determined from the difference between ). Depending on the evaluation method, this may include determining the momentum offset from the difference between the current value of the second momentum variable and the updated value for the second momentum variable. As some examples, the updated trainable parameters are or It can be determined as.

[0070] As merely an example, in some embodiments, the large neural network may include a Transformer neural network. Generally, the Transformer neural network may be characterized by having a series 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 to the attention layer input to generate an attention layer output for each element of the input. In embodiments, the attention mechanism calculates the similarity between a query and a set of key-value pairs. In embodiments, (in the case of self-attention) one or both of the query and the set of key-value pairs are determined from the attention layer input.

[0071] As an example, in a self-attention neural network layer, input embeddings can be used to determine a set of query vector and key-value vector pairs used to generate an updated embedding that includes a weighted sum of values ​​weighted by a similarity function of the query for each individual key. The similarity function may include, for example, a dot product, cosine similarity, or other similarity measures, and the query, key, and value may all be vectors. For example, the attention mechanism may, for the attention layer input for each element of the input sequence X, for example, a matrix Query transformations defined by, for example, a matrix Key transformations defined by, for example, matrices Individual query vectors configured to apply each value transformation defined by to determine the attended sequence for the output , key vector and value vector It can derive.

[0072] In some implementations, large neural networks include large language models (LLMs), e.g., Transformer-based LLMs, or large vision language models (VLMs), e.g., Transformer-based VLMs. Then, the AdamW optimization algorithm can be used to adjust the values ​​of local copies of trainable parameters in the in-house optimization steps, and the Nesterov optimization algorithm (Sutskever et al., "On the importance of initialization and momentum in deep learning", ICML, 2013) can be used to update the values ​​of trainable parameters in a shared data store using aggregated cumulative gradients for the trainable parameters in the out-house optimization steps. The 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 H values, e.g., H≥100, 200, 500, 1000, or 2000.

[0073] In some implementations, the method (and corresponding system) operates synchronously, that is, waits for all workers to report the cumulative gradients for the trainable parameters before determining the aggregated cumulative gradient for the parameters. However, the method (and corresponding system) may also operate asynchronously. That is, determining the aggregated cumulative gradient for the trainable parameters in the optimization computer system may be performed without waiting to obtain updated trainable parameter values ​​or cumulative gradients for the trainable parameters from all workers.

[0074] As described above in some embodiments, the optimization computer system (150) includes a shared server connected to a computer network (120).

[0075] Figure 3 shows a visualization schematically illustrating an example of the training process described above. In this example, The initial pre-trained model (neural network) (300) denoted by is replicated k times (in the example of k=4), and Each worker (112) denoted by trains a replica of the model on its own data shard for H steps independently and in parallel. Afterward, the worker averages the outer gradients, and the outer optimizer copies the global copy of the parameters. Updated model by updating The updated model is redeployed to the workers, and this process is repeated T times (the first two iterations are indicated). The workers train for a total of N=T×H internal steps; communication occurs only during all H steps. As indicated by the flag icons, the replicas may be in different geographical locations, e.g., different countries; and as indicated by the various TPU and GPU designations, the workers may have different hardware capabilities.

[0076] As a specific example of the training process as described above, an initial model that can be randomly or pre-trained ( ) is trained using the training dataset D={(x, y), …}, where x represents the input data item processed by the model and y represents the corresponding target output from the model. For example, if the model includes LLM or VLM, the input data item may include a sequence of tokens, and the target output may include an input sequence shifted by one token. The training dataset is partitioned across multiple shards, and the i-th shard is D i It is denoted as . A training process including k workers and using an inner optimizer (InnerOpt) and an outer optimizer (OuterOpt) can be as follows:

[0077]

[0078] In this example, the T-step external optimization process involves lines 1, 12, and 14, where at each step t, gradients from each worker are collected, averaged, and passed to the OuterOpt to update a shared copy of the parameters. Then, this copy of the shared parameters is redispatched to each local worker (line 3).

[0079] Within each outer optimization step, each worker independently and in parallel performs its own inner optimization (lines 4-9) for H, for example, hundreds of steps using an inner optimizer (InnerOpt). Each worker samples data from its shard (line 5) and updates a local copy of its trainable parameters (line 8).

[0080] In this particular example, the cumulative gradients are averaged once all workers have completed their internal optimization (line 12). Then, the resulting external gradient is used to update a shared copy of the trainable parameters (line 14) prepared for the next round of internal optimizations. Communication occurs only at every H internal optimization steps.

[0081] Adam or variations of Adam can be used for both internal and external optimizers (external optimizer stability is a numerical stability factor (It is aided by increasing to 0.1). However, in implementations using Nesterov momentum in an external optimizer, it is particularly helpful for stable learning, especially when H is large. Generally, updates using Nesterov momentum may involve updating the values ​​of trainable parameters with individual momentum values ​​before determining gradient-based updates, where the momentum values ​​accumulate previous parameter updates.

[0082] As a specific example, the Nesterov optimizer can first update the momentum variable with a moving average of its previous momentum and cumulative or external gradient, e.g., an exponential moving average. The Nesterov optimizer can calculate updates for trainable parameters using another moving average of the previously calculated momentum and cumulative or external gradient, e.g., an exponential moving average. For example, using Nesterov momentum at The update to Ro can be determined as follows (in the order below):

[0083]

[0084]

[0085]

[0086] Here, μ represents momentum, are damping coefficients (which may be the same), Δ represents the aggregated accumulation, or external gradient, and the external learning rate is τ.

[0087] In a variation, the method and system described above may be used to accelerate a single worker. Nevertheless, such a worker may be implemented using, for example, multiple processing threads and may be executed on hardware including multiple hardware accelerators such as GPUs or TPUs.

[0088] Therefore, for example, at least 10 7A computer implementation method and a corresponding system for training a large neural network having 10 trainable parameters are also described. This may include maintaining a set of trainable parameter values ​​to obtain an updated set of trainable parameter values, obtaining a copy of the trainable parameter values ​​at each of a plurality of external optimization steps, and performing H internal optimization steps on the copy of the trainable parameter values, wherein H is at least 50. Each internal optimization step may include a step of backpropagating the gradients of an objective function evaluated using training data from a training dataset to adjust the trainable parameter values.

[0089] The process may also include, for each of a plurality of external optimization steps, a step of determining a cumulative gradient for at least some of the trainable parameters by determining the difference between the value of the trainable parameter in the acquired copy and the value of the parameter after H internal optimization steps, and a step of updating the value of the trainable parameter in a set of trainable parameter values ​​using the cumulative gradient for the trainable parameter.

[0090] Other features and modes of a method or system implemented on such a single worker may be as previously described.

[0091] FIG. 4 illustrates the results 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 complexity 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, and curve (430) uses a batch size eight times larger than curve (420); curve (430) illustrates training the neural network using the techniques described herein using k=8 workers (model replicas) and H=500. While conventional techniques require communicating gradients at each training step, the technique described here communicates 500 times less (and is 8 times faster in wall clock time) and also achieves better generalization performance.

[0092] Other experiments show that performance degradation is gradual as H increases; for example, complexity increases by only 2.9% from H=50 to H=1000 despite 20 times less communication; H=500 appears to be a good trade-off. Additionally, experiments showed that increasing the number of workers improves generalization performance, but a decrease in returns occurs after k=8 (although there was no performance degradation). These techniques appear to work better as the neural network size increases. Even with the use of non-independent homogeneous data, for example, when different training data shards have different data distributions, the use of data appears to have only a small impact on performance when measured by final complexity.

[0093] Figures 5 through 8 relate to training the same Transformer-based large language model neural network ("Chinchilla") with 150 million parameters on the C4 dataset.

[0094] FIG. 5 illustrates the effect of using different outer loop optimizers by comparing the number of training steps on the x-axis and complexity on the y-axis, and SGD (stochastic gradient descent), curve (510); SGDM (SGD with momentum), curve (520); Adam, curve (530); and Nestorb momentum, curve (540) (outer learning rate, τ, = 0.7 and outer momentum = 0.9), which demonstrates that Nestorb momentum can provide improved performance (lower final complexity).

[0095] FIG. 6 shows the number of training steps on the x-axis and complexity on the y-axis, demonstrating that the techniques described herein are robust even when the number of workers changes over time. More specifically, curve (600) illustrates conventional local model training, and curves (610 to 650) illustrate examples of the techniques described herein. Curve (610) illustrates a constant number of workers (k=8); curve (620) illustrates a doubling of the number of workers during training (from 4 to 8); curve (630) illustrates a halving of the number of workers during training (from 8 to 4); curve (640) illustrates ramping up the number of workers (from 1 to 8) during training; and curve (650) illustrates ramping down the number of workers (from 8 to 1) during training. For example, it can be seen that curves (620 and 630) result in final performance as curves (640 and 650), which exemplifies that performance depends on the total number of workers, but does not depend much on how the workers, that is, the available compute, are distributed over time.

[0096] Figure 7 shows the number of training steps on the x-axis and complexity on the y-axis, illustrating that the techniques described herein are robust against bad / asynchronous communications. The curves all show examples of the techniques described herein, and curves (700, 710, 720, and 730) show training for 0%, 10%, 30%, and 50% communication dropout (randomly dropping external gradients), respectively. Even with 50% dropout, i.e., a 50% probability of dropping external gradients, the reduction in complexity is only 2.1% compared to the absence of dropout.

[0097] FIG. 8 shows the number of training steps on the x-axis and complexity on the y-axis, demonstrating that the technique described herein can provide some performance advantages even with a single worker. More specifically, curve (800) illustrates conventional local model training, curve (810) illustrates training as described herein with eight workers, and curve (820) illustrates training as described herein with a single worker (k=1). It can be seen that even with just one worker, the technique described herein provides faster convergence and better final performance.

[0098] Now, some large language models (LLMs) that can be trained or fine-tuned from scratch using the techniques described in this specification will be described.

[0099] Generally, a language model neural network is a neural network trained to generate the next token in a sequence when given a text prompt containing a sequence of tokens in natural language. This process is repeated to expand the text prompt one token at a time to generate natural language output, meaning that natural language output can be generated automatically and regressively on a token-by-token basis. At each time step, the language model neural network processes the current sequence to generate a probability distribution for the vocabulary of tokens. For example, the next token can be selected using the probability distribution by sampling from the distribution using core sampling or other sampling techniques, or by selecting the token with the highest probability. The tokens in the vocabulary can include various arbitrary tokens, such as words, subwords, characters, punctuation and other symbols, and some combinations of numbers. Typically, a language generation neural network is trained on a text corpus consisting of tokens from the vocabulary (and other tokens that can be mapped to tokens outside the specified vocabulary) to predict the next token in a sequence of tokens from the training data.

[0100] Large language model neural networks, for example, 10 n There can be more than 100 trainable parameters, where n = 7, 8, 9, 10, or 11. Language model neural networks can be trained on over 10 billion, 10 billion, or 100 billion words or tokens representing words or other text tokens, for example, subwords (also known as "word fragments"). There are several training datasets in the public domain, one example being the C4 dataset (Raffel et al., arXiv:1910.10683).

[0101] Some large language model neural networks include autoregressive transducer neural networks; transducer-based language models may be encoder-only models, decoder-only models, or encoder-decoder models. Some large language model neural networks may include, for example, diffusion models or expert mixed models.

[0102] Language model neural networks can be configured to perform specific tasks by providing natural language descriptions of desired responses as input or "prompts." Alternatively, or instead, language model neural networks can be "fine-tuned" to perform specific tasks by obtaining a pre-trained language model neural network trained on a large corpus and then further training a portion of the entire network on a relatively small number of examples specialized for the type of task to be performed.

[0103] Now, the techniques described in this specification describe some multimodal models, such as vision language models (VLMs), that can be used to train or fine-tune from scratch.

[0104] Generally, a multimodal machine learning model has a multimodal input configured to receive a first multimodal input and a second multimodal input. As used herein, 'modality' refers to a type of data, and thus a multimodal machine learning model is a model capable of processing a number of different types of data. The multimodal input to the model may include a mixture of inputs of the first and second modalities (at least).

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

[0106] A multimodal machine learning model may be configured to jointly process an encoded version of text and an encoded version of a second multimodal input, for example, an image or video, to generate a model output that defines the result of a machine learning task. Several examples of machine learning tasks that can be performed by a multimodal machine learning model are described below.

[0107] The received text may include text in one or more natural languages, text in computer languages, or both. The computer language may be any form of language used to communicate with a computer, for example, a markup language, a command or configuration language, a data exchange language such as JSON, or a programming language. The text may be received as a series of encoded characters, for example, UTF-8 encoded characters; these 'characters' may include Chinese and other similar characters, as well as ideograms, syllabic characters, etc.

[0108] A multimodal machine learning model may include a text encoder that processes a sequence of text to represent the text as a series of text tokens from a vocabulary of text tokens, each representing, for example, a word, word fragment, or character of a natural language or a computer language.

[0109] If the second multimodal input includes an image or video, it may include image data that defines color or intensity values ​​for pixels of a still or moving image in one, two, or three dimensions. The image may include a LIDAR point cloud, and the image data may define the locations of points of a still or moving point cloud. As another example, the image or video received by the second multimodal input may include, for example, a neural 3D representation representing a 3D scene as a set of latent feature vectors, such as a neural radial field representation. The multimodal machine learning model may include a second multimodal input encoder that processes the second multimodal input using, for example, one or more convolutional, attention, fully connected, or recurrent layers to generate an encoded version of the second multimodal input. Generally, the encoder may implement any form of encoding appropriate for the type of data to be encoded. Just as an example, if the second multimodal input includes an image or video, it may be encoded, for example, as a feature for each of a set of patches tiling the image, or as a sequence of visual tokens selected from a vocabulary of visual tokens, or as a representation of a distinct object in the visual input. These visual tokens may, but do not need to be, interleaved with text tokens processed by the model.

[0110] The model output may include any form of output appropriate for the machine learning task performed by the multimodal machine learning model. In some embodiments, the model output includes text in a natural language or computer language that defines the result of the task, for example, for tasks such as image captioning, visual question answering, or object detection or instance segmentation. Alternatively, the model output may include data that defines an image, video, or audio object in a generative task, for example; or the model output may include non-text 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, for example, a call to a software API for a software tool used to perform the task; and the multimodal input may receive output from a software tool used to generate the final model output that performs the task. Several specific examples of the model output will be provided later.

[0111] These multimodal models can be trained using very large (but potentially noisy) datasets in which text is paired with images and / or one or more other types of data, e.g., audio data, or data regarding the behavior of an agent operating in an environment to perform various tasks. Such models can be trained, for example, using self-supervised learning. The pairings can often be incomplete, and the training dataset may or may not include arbitrary real-world examples of the specific task to be performed; nevertheless, the ability to perform the specific task may still be demonstrated. There are many examples of suitable training datasets available publicly.

[0112] Some exemplary multimodal machine learning models to 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 PaLI-X (Chen et al., arXiv:2305.18565).

[0113] In some embodiments, the second multimodal input may include an observation characterizing the environment of the agent performing the task, e.g., a mechanical agent or a software agent. The observation may characterize the environment at a specific time step, and the model output may define one or more actions to be performed by the agent at that time step. For example, each action may be represented as a sequence of text, e.g., as one or more characters such as letters and numbers representing the action, or as text defining a low-level 'skill' from a set of skills; or the model output may define, e.g., parameters of a probability distribution in which an action is selected. Optionally, the text received by the 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 previous time step. If the agent is a software agent, the model output may include a text output for calling a software API at a time step, and at a subsequent time step, the model input, e.g., the text input, may include a response from the software agent, e.g., from the API.

[0114] Some examples of multimodal machine learning models in which the techniques described herein can be used to control an agent 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).

[0115] Such a multimodal machine learning model may include an audio input, or an agent action input for receiving agent action data representing the actions of an agent performing a task in an environment. Data received in this manner may be processed jointly with data from a text input and a second multimodal input to generate a model output.

[0116] Such multimodal machine learning models have multimodal inputs and, in an embodiment, can perform a range of different tasks. However, in an embodiment, not every task performed by the model requires multimodal input, for example, a task of generating an image from a text description of an image, or an image captioning task. In some embodiments, after training, text inputs may be used to specify a particular task to be performed by the multimodal machine learning model, for example, by providing the model with a "prompt" describing the task to be performed or by providing the model with an example of the task as a prompt. Such prompts may optionally be included in the training data.

[0117] A multimodal machine learning model is configured to process multimodal inputs according to the trainable parameters of the multimodal machine learning model to generate model outputs that define the results of one or more machine learning tasks.

[0118] Generally, each training data item contains multimodal data to be used to train a multimodal machine learning model, for example, using a self-supervised training goal. As another example, such a multimodal machine learning model can also be trained using a reinforcement learning goal, for example, when the model is used to control an agent to perform a task.

[0119] There are many different types of self-supervised objective functions that can be used. As an example, the model can be trained using softmax cross-entropy loss, for example, using a language model-style teacher forced by softmax cross-entropy loss. As another example, the model, for example, one or more preceding second modality inputs such as images or videos Text tokens conditioned and preceding for The lth text token y conditioned for l For a multimodal input comprising a sequence of text encoded as L tokens having, It can be trained with an autoregressive negative log-likelihood (NLL) loss such as the one above. As another example, the model can be trained with a masking loss, for instance, a loss that requires the model to predict masked data, such as masked text tokens. As yet another example, a multimodal machine learning model can be trained using an autosupervised objective function that includes a contrast loss function (which depends on positive examples and one or more negative examples).

[0120] Similarly, there are various types of reinforcement learning objective functions that can be used.

[0121] Each training data item may include, for example, an exemplary sequence of text and an example of a second modality input, for example, an exemplary image or video; generally, these are semantically related to each other (however, this is not always the case because the training dataset may contain noise). As an example, matched text and image or video data may be obtained from web pages, for example, from images or videos and their corresponding alt-text (text from HTML or XHTML alt attributes); or from web pages where images or videos and text are interleaved. An example of such a dataset is WebLI (Web Language Image, Chen et al. arXiv:2305.18565v1). Training datasets for other types of second modality inputs may similarly be obtained from web pages. These training datasets may be large items, for example, 10 7 , 10 8 or 10 9 There could be more items than dogs.

[0122] Alternatively, or instead, smaller but more specialized training datasets may be used to fine-tune models for specific tasks or tasks. Some examples for visual tasks include 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 Image V4 (Kuznetsova et al., arXiv:1811.00982); and the SBU dataset (Ordonez et al., "Im2Text: Describing Images Using 1 Million Captioned Photographs", NeurIPS 2011). Conceptual Captions datasets such as V1 (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 exemplary 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 exemplary task-specific training dataset for agent (robot) control is described in Ebert et al., arXiv:2109.13396.

[0123] Several examples of machine learning tasks that can be performed by multimodal models are now described.

[0124] Generally, a multimodal machine learning model can be trained to perform any type of machine learning task or tasks. After the multimodal machine learning model is trained, it can be deployed for use in performing machine learning task(s). For example, the machine learning model can be deployed in an environment where users can provide requests that cause the machine learning model to process specified multimodal inputs and generate corresponding model outputs. Users can provide requests, for example, through a user interface or through an Application Programming Interface (API). Requests can be transmitted from a user device (e.g., via a data communication network, e.g., the Internet) to one or more computers implementing the machine learning model, for example, in a data center. The machine learning model can process the multimodal inputs specified by the user requests to generate corresponding model outputs, and then transmit the model outputs to user devices (e.g., via a data communication network).

[0125] In some implementations, after training, a specific task to be performed by a multimodal machine learning model may be described by some or all of a sequence of text within the multimodal input to the model. For example, in a multimodal input containing an image, video, or audio item, the prompt may specify "Generate caption," "Generate description," "Answer the following question: [for the image, video, or audio item]," or "Detect person." If the model is used for an agent control task, the prompt may define "Pull a knife out of the drawer" or "Q: What action should the robot take to pull a knife out of the drawer?" Additionally, or instead, these prompts may provide one or more examples of the task to be performed. A multimodal machine learning model may be trained on multiple natural and / or computer languages, and the prompt may then specify the language to use.

[0126] Some examples of machine learning tasks that can be performed by a trained model as described in this specification are as follows.

[0127] For some tasks, the second modality input represents, for example, images or videos as previously described from a camera or other imaging device capturing images or videos from a real-world environment, and / or audio data, such as speech or other sounds captured from a real-world environment. Generally, the tasks described below may be tasks that require spatial recognition or other context from image, video, or audio items. For example, a prompt may ask, "What is the object in the top-left corner?" or "What is the answer to the voice question?"

[0128] As an example, the task may include an object or action detection task. The training data items per task may include 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 may include text providing bounding box coordinates of the object(s) or action(s) (for an image or video). After training, when the model is used for inference, the model output (122) may include or represent text describing or labeling the object(s) or action(s) detected in the second modality image input and bounding box coordinates for the object(s) or action(s) detected (for an image or video), for example, "10 20 90 100 cat 20 30 100 100 dog".

[0129] As another example, the task may include a classification task, for example, an object or action classification task. The training data items per task may include 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 for inference, the model output may include data, for example, text, that classifies the object(s) or action(s) of the second modality input into one of a plurality of classes.

[0130] As another example, a task may include an image, video, or audio item describing the task, for example, a captioning task (including, as used herein, an audio description task to describe what is happening in the video). A task-specific training data item may include 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 for inference, the model output may include data describing the image, video, or audio item in the second modality input, for example, text. For example, the model output may provide a caption or description for the second modality input item, count objects within the second modality input item, or provide some other form of description of the second modality input item.

[0131] As another example, the task may include an image, video, or audio question-and-answer task. Task-specific training data items may include 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 for inference, the model output may include data, e.g., text, that responds to a question regarding a second modality input specified in a sequence of text prompts, e.g., as described above. This may be used to answer questions regarding visual plots and charts or sounds, e.g.

[0132] As another example, the task may include a character or word recognition task, for example, an optical character recognition (OCR) task. Task-specific training data items may include an image, video, or audio item, and a sequence of text including text depicted in the image or video or expressed as speech in the audio item. After training, when the model is used for inference, the model output may include text expressing the characters or words of the second modality input, for example, in natural language.

[0133] As another example, the task may include a still or video or audio generation task. Task-specific training data items may include image, video, or audio items and sequences of text describing the image, video, or audio items. After training, when the model is used for inference, the model output may include data regarding the image, video, or audio items, for example, image data defining values ​​for pixels of a still or video or audio data representing values ​​of an audio waveform, and sequences of text within the multimodal input to the model may describe or characterize the image, video, or audio item to be generated.

[0134] As another example, the task may include a computer language text generation task. Task-specific training data items may include image, video, or audio items and sequences of text in a computer language for generating the image, video, or audio items. After training, when the model is used for inference, the model output may include text in another computer language for generating or rendering image, video, or audio items, for example, web pages, plots, or charts, from a second modality input.

[0135] In another example of a computer language text generation task, the task-specific training data item may include an image, video, or audio item and a sequence of text in a computer language for performing tasks related to the image, video, or audio item, for example, analyzing the content of the image, video, or audio item to provide analysis results, or a data processing task involving a search for retrieving information regarding the content of the image, video, or audio item. The computer language of the model output may include a computer language for calling functions or calling one or more external APIs. As just one example, such output may be formatted as a JSON object. As before, the text sequence of the multimodal input may define the task to be performed, and the second modality input may include a task involving the manipulation of a specific type of data that can benefit from access to an API, such as the image, video, or audio item on which the task is to be performed, for example, mathematical data, date / time related data, scientific data, or recent data that may be post-training of the model (accessible via a search function or API). After training, when the model is used for inference, the model output may include text in another computer language for performing a task in relation to an image, video, or audio item in the second modality input, for example, as described above. Then, the method may include the step of using text in a computer language to perform a task.

[0136] Generally, if the model output includes text, it can be provided as speech representing the text.

[0137] In some embodiments, the machine learning task includes an agent control task in which the agent interacts with the environment to perform an agent control task. In these embodiments, the multimodal input includes observations characterizing the environment. For example, the multimodal input may include a sequence of text defining a task to be performed by the agent, and the second modality input may represent an image, video, audio, or other observation of the environment captured, for example, by a camera or other imaging device or by a microphone from the real-world environment. A task-specific training data item may include 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 for inference, the model output includes an action selection output, for example, containing text used to select one or more actions to be performed by the agent in the environment in response to an observation. As an example, the model output (122) can be converted into a control signal for a mechanical agent such as a robot, such as text like "A: 132 114 128 5 25 156", e.g., " Actions such as " can be defined. As another example, the action selection output can also define one or more low-level skills from a vocabulary of previously learned skills, for example, instead of or. As before, a sequence of text within the multimodal input to the model can describe the task to be performed, for example, "what action the robot must take [to perform the task]."

[0138] In some agent control implementations, the environment is a real-world environment, the agent is a mechanical agent interacting with the real-world environment, e.g., a robot or 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 a task. For example, the agent may be a robot or other mechanical agent interacting with the environment to achieve a specific task, e.g., to position or manipulate an object of interest within the environment, to move an object of interest to a designated location within the environment, or to drive to a designated destination within the environment. In these implementations, observations may include one or more of, e.g., images for capturing observations when the agent interacts with the environment, object location data, and sensor data. Actions may define control signals for controlling the robot or other mechanical agent, e.g., positions, torques, or other control signals for parts of the mechanical agent, or higher-level control commands.

[0139] In some agent control implementations, the agent may be a software agent, i.e., a computer program, configured to perform a task. Some examples where the agent is a software agent are as follows.

[0140] As an 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. Observations may be component locations and / or interconnections, and actions may include component placement or interconnection routing actions. An integrated circuit having interconnection lines routed as determined may subsequently be manufactured.

[0141] As another example, the environment may be a real-world computing environment, and the task may be managing the distribution of tasks or operations across computing resources, for example, on a mobile device and / or in a data center. Observations may include observations of computing resources such as computing or memory capacity, or internet-accessible resources, or may be those related to the behavior of computing resources when processing tasks or operations; and actions may include assigning tasks or operations to specific computing resources.

[0142] As another example, the environment may be a real-world computing environment, and the task is, for example, managing the processing of a queue of sequentially arriving tasks by one or more real-world servers. Observations may include the times of sequential task exits, or the time intervals between sequential task exits, or the time taken for a server to process each task, or the arrival times of sequential tasks, or the time intervals between arrivals, or observations of data characterizing the type of task(s). Actions may include actions that assign specific tasks to specific computing resources.

[0143] As another example, the environment may include a real-world computer system or network, and the task may be maintaining the security of the computer system or network. Observations may include any observations characterizing the behavior of the computer system or network, and actions may include actions to control behavior in order to restrict or modify abnormal or undesirable behavior caused by, for example, the presence of a virus or other security breach.

[0144] As another example, the environment may include a data packet communication network environment, and the task may be routing packets of data through the communication network. Actions may include data packet routing actions, and observations may include observations of routing tables including routing metrics such as routing path length, bandwidth, load, hop count, path cost, delay, maximum transmission unit (MTU), and reliability metrics.

[0145] 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 may be a human user of a digital assistant, such as a smart speaker, a smart display, or some other device used to instruct the user to perform an action. The task may be any real-world task that the user wishes to perform. Observations may be acquired from a monitoring system, such as an observation capture subsystem, such as a video camera or a sound capture system, to capture visual and / or audio observations of the user performing the task. Actions may include commands in the form of audio data, such as text, images, videos, or speech, that guide the user when performing the task.

[0146] In this specification, the term “configured” is used in relation to computer program components as well as computing systems and environments. A computing system or environment is deemed “configured” to perform certain operations or actions when it possesses the necessary software, firmware, hardware, or combination thereof, thereby enabling the performance of said operations or actions during operation. For example, configuring a system may involve installing a software library with specific algorithms, updating firmware with new instructions for handling data, or adding hardware components for enhanced processing capabilities. Similarly, one or more computer programs are “configured” to perform certain operations or actions when, upon execution by a computing device or hardware, they include instructions that cause the device to perform such intended operations or actions.

[0147] The embodiments and functional operations described herein may be implemented in various forms including digital electronic circuits, software, firmware, computer hardware (including the disclosed structures and their structural equivalents), or any combination thereof. The subject matter of the invention may be realized as one or more computer programs, essentially as modules of computer program instructions encoded on a tangible non-transient storage medium for execution by a computing device or hardware or for controlling the operation of a computing device or hardware. The storage medium may be a storage device such as a hard drive or a solid-state drive (SSD), a storage medium, a random or serial access memory device, or a combination thereof. Additionally or alternatively, program instructions may be encoded on a transmitted signal, such as a mechanically 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. Additionally, embodiments may utilize emerging technologies such as quantum computing or neuromorphic computing for specific applications, and components may be deployed in distributed or cloud-based environments residing on different machines or within cloud infrastructure.

[0148] The term “computing device or hardware” refers to 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 may take the form of processor firmware, protocol stacks, database management systems, operating systems, or a combination of these elements. Embodiments may particularly benefit from leveraging the parallel processing capabilities of GPUs in general-purpose computing in the context of graphics processing units (GPGPUs), where code specifically designed for GPU execution—often referred to as kernels or shaders—is utilized. Similarly, TPUs excel at executing optimized tensor operations that are critical to many machine learning algorithms. By utilizing these accelerators and their specialized programming models, systems can achieve significant speed improvements and efficiency gains for tasks involving artificial intelligence and machine learning, particularly in fields such as computer vision, natural language processing, and robotics.

[0149] Computer programs, also referred to as software, applications, modules, scripts, code, or simply programs, may be written in any programming language, including compiled or interpreted languages, and declarative or procedural languages. They may be deployed in various forms, such as standalone programs, modules, components, subroutines, or any other units suitable for use within a computing environment. Programs may or may not correspond to a single file in a file system and may 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). Computer programs may be located at a single site or distributed across multiple sites and executed on a single computer or multiple computers interconnected via a data communication network. Specific embodiments of computer programs may include combinations of traditional programming languages ​​and specialized languages ​​or libraries designed for GPGPU programming or TPU utilization, depending on the selected hardware platform and desired performance characteristics.

[0150] 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 may be located at a single site or distributed across multiple locations. In some cases, one or more dedicated computers may be used for a specific engine, whereas in other cases, multiple engines may operate simultaneously on the same one or more computers. Examples of engine functions within the context of AI and machine learning may include data preprocessing 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 distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors, and on the overall architecture.

[0151] The processes and logic flows described herein may be executed by one or more programmable computers running one or more computer programs to perform functions by operating on input data and generating outputs. Additionally, graphics processing units (GPUs) and tensor processing units (TPUs) may be utilized to significantly accelerate performance by enabling the simultaneous execution of modes of these processes and logic flows. This approach offers significant advantages for computationally intensive tasks often found in AI and machine learning applications, such as matrix multiplications, convolutions, and other operations exhibiting high levels of parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speed and efficiency gains can be achieved compared to relying solely on CPUs. Alternatively, or in combination with programmable computers and specialized processors, these processes and logic flows may also be implemented using specialized processing hardware, such as field programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs), for much greater performance or energy efficiency in specific use cases.

[0152] Computers capable of executing computer programs can be based on general-purpose microprocessors, special-purpose microprocessors, or a combination of both. They may also utilize any other type of central processing unit (CPU). Additionally, graphics processing units (GPUs), tensor processing units (TPUs), and other machine learning accelerators may be used to enhance performance, particularly for tasks involving artificial intelligence and machine learning. These accelerators often work in conjunction with the CPU to handle specialized calculations while the CPU manages overall system operations and other tasks. Typically, the CPU receives instructions and data from read-only memory (ROM), random access memory (RAM), or both. Essential components 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 such as the complexity of the AI ​​model, the volume of data being processed, and desired performance and latency requirements. The 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 such as hard drives, SSDs, or flash memory for permanent data storage.

[0153] Computer-readable media suitable for storing computer program instructions and data include 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 data size, access speed requirements, cost considerations, and the desired level of portability or permanence.

[0154] To facilitate user interaction, embodiments of the subject matter of the invention described herein may 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 a user. Input may be provided by the user through various means including a keyboard, touchscreen, voice command, gesture recognition, or other input modalities, depending on the specific device and application. Additional input methods may include acoustic, voice, or haptic input, while feedback to the user may take the form of visual, auditory, or haptic feedback. Additionally, the computer may interact with the user by exchanging documents with the user's device or application. This may involve sending web content or data in response to requests, or sending and receiving text messages or other forms of messages via 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.

[0155] Machine learning models can be implemented and deployed using machine learning frameworks such as TensorFlow or JAX. These frameworks provide comprehensive tools and libraries that facilitate the development, training, and deployment of machine learning models.

[0156] Embodiments of the subject of the invention described herein may be implemented within a computing system comprising one or more components, depending on specific applications and requirements. These may include back-end components, such as back-end servers or cloud-based infrastructure; optional middleware components, such as middleware servers or application programming interfaces (APIs), to facilitate communication and data exchange; and front-end components, such as a client device having a user interface, a web browser, or an app, to enable a user to interact with the implemented subject of the invention. For example, the described functions may be implemented only 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 may be interconnected using any form or medium of digital data communication, such as a communication network, such as a local area network (LAN) or a wide area network (WAN), including the Internet, where present. The specific system architecture and selection of the 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.

[0157] A computing system may include clients and servers that are geographically separated and can interact via a communication network. Specific types of networks, such as Local Area Networks (LANs), Wide Area Networks (WANs), or the Internet, will depend on the reach and scale of the application. Client-server relationships are established through computer programs running on individual computers and are 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, the server transmits data or commands to a user device, such as a computer, smartphone, or tablet, acting as a client. The client device can then process the received information, display the results to the user, and potentially send the data or feedback back to the server for further processing or storage. This allows for dynamic interaction between the user and the system, enabling a wide range of applications and functions.

[0158] Although this specification contains many specific implementation details, they should not be interpreted as limitations on the scope of any invention or claimables, but rather as descriptions of features that may be assigned to specific embodiments of specific inventions. Certain features described in this specification in the context of separate embodiments may also be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment may also be implemented individually or in any suitable sub-combination in multiple embodiments. Furthermore, while features may be described above as operating in specific combinations and may even be initially claimed as such, one or more features from the claimed combination may be omitted from the combination in some cases, and the claimed combination may relate to a sub-combination or a variation of a sub-combination.

[0159] Similarly, although operations are illustrated in the drawings and cited in a specific order in the claims, this should not be understood as requiring that such operations be performed in the specific order illustrated or in a sequential order, or that all illustrated operations be performed, in order to achieve desired results. In certain situations, multitasking and parallel processing may be advantageous. Furthermore, the separation of the 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 may generally be integrated together into a single software product or packaged into multiple software products.

[0160] Specific embodiments of the subject matter of the invention have been described. Other embodiments are within the scope of the following claims. For example, the actions described in the claims may be performed in different orders and still achieve desirable results. As an example, the processes illustrated in the accompanying drawings do not necessarily require the specific order or sequential order illustrated to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

Claims

Claim 1 A computer implementation method for training a large neural network using a distributed computing system to optimize an objective function, wherein the distributed computing system comprises a plurality of computing devices coupled by a computer network and operating in parallel, each computing device has an associated memory, and the large neural network comprises at least 10 7 The method has trainable parameters, and the method comprises: a training dataset including multiple shards, each of which is 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 multiple computing devices, each worker being implemented on one of the computing devices, each worker being assigned an individual shard among the multiple shards, and each worker, in each of the multiple external optimization steps, Obtain a local copy of the values ​​of trainable parameters from a shared data store; and To obtain an updated set of trainable parameter values ​​for the worker after H internal optimization steps, H internal optimization steps are performed on local copies of the trainable parameter values, where H is at least 50, and each internal optimization step A method configured to include a step of backpropagating the gradient of an objective function evaluated using a shard assigned to a worker to adjust the value of a local copy of a trainable parameter, wherein the method comprises: for each of a plurality of external optimization steps: a step of determining a cumulative gradient for at least some of the trainable parameters by determining, for each worker, the difference between the value of the trainable parameter in the local copy obtained by the worker and the value of the parameter after H internal optimization steps; a step of determining a cumulative gradient aggregated by each of at least some of the trainable parameters in an optimization computer system comprising one of computing devices or a shared server coupled to a computer network; a step of aggregating the cumulative gradient for the trainable parameter determined for the worker; and a step of updating the value of the trainable parameter in a shared data store using the cumulative gradient aggregated for the trainable parameter, wherein the value of the trainable parameter in the shared data store provides a trained parameter for a large neural network. Claim 2 A computer implementation method for training a large neural network using a distributed computing system to optimize an objective function, wherein the distributed computing system comprises a plurality of computing devices coupled by a computer network and operating in parallel, each computing device has an associated memory, and the large neural network comprises at least 10 7 The method has trainable parameters, and the method comprises the steps of: acquiring a training dataset including a plurality of shards, wherein each of the plurality of shards is 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 a plurality of computing devices, wherein each worker is implemented on one of the computing devices, and each worker is assigned an individual shard among the plurality of shards, and each worker, in each of the plurality of external optimization steps, Obtain a local copy of the values ​​of trainable parameters from a shared data store; and To obtain an updated set of trainable parameter values ​​for the worker after H internal optimization steps, H internal optimization steps are performed on local copies of the trainable parameter values, where H is at least 50, and each internal optimization step The method comprises the step of backpropagating the gradient of an objective function evaluated using a shard assigned to a worker to adjust the value of a local copy of a trainable parameter, wherein for each of a plurality of external optimization steps: for each worker, the step of determining a cumulative gradient for at least some of the trainable parameters by determining the difference between the value of the trainable parameter in the local copy obtained by the worker and the value of the parameter after H internal optimization steps; and for each of the optimization computer systems comprising one of the computing devices or a shared server coupled to a computer network, the step of determining a cumulative gradient aggregated by each of at least some of the trainable parameters; the step of aggregating the cumulative gradient for the trainable parameter determined for the worker; and the step of updating the value of the trainable parameter in a shared data store using the aggregated cumulative gradient for the trainable parameter, wherein the value of the trainable parameter in the shared data store provides a trained parameter for a large neural network. Claim 3 A method according to claim 1 or 2, wherein for each worker, the step of determining a cumulative gradient for at least some of the trainable parameters comprises the step of the worker determining a cumulative gradient for each of the trainable parameters, and the method further comprises the step of communicating the cumulative gradient for at least some of the trainable parameters from each worker to an optimization computer system via a computer network. Claim 4 A method according to claim 3, further comprising the steps of: selecting a cumulative gradient subset from the cumulative gradient for each trainable parameter for the worker in the worker; and communicating the cumulative gradient subset from the worker to an optimization computer system through a computer network. Claim 5 A method according to claim 1 or 2, wherein each worker further comprises the step of communicating at least some of the trainable parameters of an updated set of trainable parameter values ​​for the worker from the worker to an optimization computer system via a computer network, and for each worker, the step of determining a cumulative gradient for at least some of the trainable parameters comprises the step of the optimization computer system determining a cumulative gradient for at least some of the trainable parameters. Claim 6 A method according to claim 5, further comprising the steps of selecting a subset of trainable parameters from a set of updated trainable parameter values ​​for a worker, and transmitting the subset of trainable parameters from the worker to an optimization computer system via a computer network. Claim 7 In any one of claims 1 to 6, the step of backpropagating the gradient of an objective function, evaluated using a shard assigned to a worker to adjust the value of a local copy of a trainable parameter, for each trainable parameter: for each trainable parameter, the step of maintaining a first momentum variable for the trainable parameter locally in the worker, wherein the first momentum variable comprises a moving average of the gradient for adjusting the trainable parameter, and the step of adjusting the value of the trainable parameter by an offset dependent on the value of the first momentum variable; and the step of updating the value of the trainable parameter in a shared data store using an aggregated cumulative gradient for the trainable parameter comprises the step of maintaining a second momentum variable for the trainable parameter in a shared data store, the step of determining an updated value for the second momentum variable from the aggregated cumulative gradient for the trainable parameter, and the step of determining an updated value of the trainable parameter in a shared data store from the updated value of the second momentum variable. Claim 8 A method according to claim 7, further comprising the steps of determining an updated value for a second momentum variable from the current value of a trainable parameter and the aggregated cumulative gradient for the trainable parameter, and determining an updated value for a trainable parameter from the difference between the updated value for the second momentum variable and the current value of the second momentum variable. Claim 9 The method according to claim 7 or 8, wherein the step of determining an updated value for a second momentum variable from an aggregated cumulative gradient for the trainable parameter includes the step of scaling the aggregated cumulative gradient by an external optimization learning rate parameter, and the method further includes the step of maintaining the external optimization learning rate parameter constant. Claim 10 A method wherein, in any preceding claim, the large neural network comprises a transformer neural network characterized by having a series of attention neural network layers, each attention neural network layer is configured to have an attention layer input for each element of the input and to apply an attention mechanism to the attention layer input to generate an attention layer output for each element of the input, wherein the attention mechanism calculates a similarity between a query and a set of key-value pairs, and one or both of the query and the set of key-value pairs are determined from the attention layer input. Claim 11 A method wherein, in any preceding term, the large neural network comprises a large language model (LLM) or a large vision language model (VLM) neural network; in the inner optimization step, adjusting the value of a local copy of a trainable parameter uses the AdamW optimization algorithm; and in the outer optimization step, updating the value of a trainable parameter in a shared data store using an aggregated cumulative gradient for the trainable parameter uses the Nesterov optimization algorithm. Claim 12 A method wherein, in any preceding term, the step of aggregating the cumulative gradient for the determined trainable parameters for the worker includes the step of determining a weighted average, wherein the cumulative gradient for each worker is weighted by the number of training data items within the shard for the worker. Claim 13 A method in which, in any preceding term, the step of determining an aggregated cumulative gradient for a trainable parameter in the optimization computer system is performed without waiting to obtain an updated trainable parameter value or cumulative gradient for the trainable parameter from all workers. Claim 14 A method in which, in any preceding claim, the optimization computer system comprises a shared server connected to a computer network. Claim 15 A method in which, in any preceding claim, the computer network includes a wide area network (WAN). Claim 16 As a computer implementation method for training a large neural network, the large neural network comprises at least 10 7 The method having trainable parameters, the method comprises the step of maintaining a set of values ​​of the trainable parameters; and in each of the plurality of external optimization steps: A step of obtaining a copy of the value of a trainable parameter; and The method comprises the step of performing H internal optimization steps on copies of the values ​​of trainable parameters to obtain an updated set of trainable parameter values ​​after H internal optimization steps—where H is at least 50—and each internal optimization step is, A method comprising the step of backpropagating a gradient of an objective function evaluated using training data from a training dataset to adjust the value of a trainable parameter, wherein for each of a plurality of external optimization steps: a step of determining a cumulative gradient for at least some of the trainable parameters by determining the difference between the value of the trainable parameter in the acquired copy and the value of the parameter after H internal optimization steps; and a step of updating the value of the trainable parameter in a set of trainable parameter values ​​using the cumulative gradient for the trainable parameter. Claim 17 One or more non-transient computer storage media storing instructions that, when executed by one or more computers, cause one or more computers to perform the operations of any one of the individual methods of claims 1 through 16. Claim 18 A system comprising: one or more computers; and one or more storage devices coupled to one or more computers so as to be communicable with the one or more computers, wherein the one or more storage devices store instructions that, when executed by one or more computers, cause one or more computers to perform operations of any one of the individual methods of claims 1 through 16.

Images