Neural networks using distributed path construction

Abstract

To provide a neural network that uses distributed path construction. [Solution] A method, system, and apparatus, including a computer-readable medium, for training a neural network on a training dataset using a distributed computing system including multiple workers is described. The neural network includes multiple components. Each component includes a subset of the multiple parameters of the neural network.

Description

Technical Field

[0001] This specification relates to neural networks.

Background Art

[0002] A neural network is a machine learning model that uses one or more layers of non-linear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from the received input according to the current values of each set of parameters.

[0003] The overall trend of neural networks has been to create larger and more complex networks in order to achieve higher accuracy. As neural networks increase in size and complexity to improve accuracy, the computational cost and communication cost during training of the neural networks also increase.

Prior Art Documents

Non-Patent Documents

[0004]

Non-Patent Document 1

Non-Patent Document 2

[0005] This specification describes a distributed computing system comprising multiple workers in one or more locations that implements, trains, or both implements neural networks to perform one or more machine learning tasks.

[0006] The subjects described herein may be implemented in specific embodiments to achieve one or more of the following advantages:

[0007] By dynamically assembling different components of a neural network into different paths within the neural network for multiple machine learning tasks, the techniques described herein can implement sparsely activated neural networks that use only a smaller subset of all parameters of the neural network to compute inference for a given input for one of the tasks, and therefore consume fewer computational resources, while maintaining performance equivalent to or better than that of a conventional neural network.

[0008] Furthermore, the implementation of the described technology will make training neural networks more practical and less computationally intensive. Conventional methods for training large neural networks on distributed computing systems require a large number of interconnected hardware accelerators, such as GPUs or TPUs. This is partly because many existing data parallelism or model parallelism algorithms require high-bandwidth communication between devices that need to exchange data with each other during training.

[0009] In contrast, the implementation of the described technique does not require the synchronization of parameter values ​​of each local instance of the component across computing devices after each inner optimization step, thus reducing the communication required between computing devices during training by two or three orders of magnitude compared to such conventional methods. This, in turn, enables a much more distributed and flexible technique—for example, computing devices can be heterogeneous, physically separated from each other, and connected by a communication network under poor or suboptimal conditions. In the described technique, the computing devices on which the neural network components are implemented do not need to be manufactured by the same hardware manufacturer and can be located in different countries or even different continents.

[0010] The implementation of the described technology can further facilitate the training of larger neural networks to achieve higher performance in any of a wide range of machine learning tasks. By enabling the automatic scaling of the size of the worker pool according to resource availability, the implementation of the described technology can achieve elastic resource utilization and mitigate the problem of resource overcapacity, thereby improving the efficiency of the use of computing resources by the training process - for example, training a large neural network with over a billion parameters faster (e.g., in wall-clock time) while consuming fewer computing resources (e.g., memory and processing power) than conventional techniques.

[0011] Details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.

Brief Description of the Drawings

[0012] [Figure 1A] FIG. showing an exemplary distributed computing system. [Figure 1B] FIG. showing an example of selecting a path in an exemplary neural network using a router. [Figure 1C] FIG. showing an exemplary path in an exemplary neural network. [Figure 1D] FIG. showing an exemplary worker included in a distributed computing system. [Figure 2] FIG. is a flowchart of an exemplary process for training a neural network. [Figure 3] FIG. is a flowchart of sub-steps of one step of the process of FIG. 2. [Figure 4] FIG. is a flowchart of sub-steps of another step of the process of FIG. 2. [Figure 5]A flowchart of an exemplary process for using a neural network to generate a network output from a network input. **DETAILED DESCRIPTION OF THE INVENTION**

[0013] Like reference numerals and designations in the various drawings indicate like elements.

[0014] FIG. 1A shows an exemplary distributed computing system 100 that implements, trains, or both, a neural network 110 to perform one or more machine learning tasks. The distributed computing system 100 includes a plurality of workers at one or more locations.

[0015] The neural network 110 can perform any kind of machine learning task, i.e., the neural network 110 can receive any kind of network input 102 and, through training, be configured to generate a network output 112 that includes any kind of score, classification, or regression (e.g., generation) output based on the network input.

[0016] In some situations, the neural network 110 can be referred to as an autoregressive neural network when the neural network generates an output sequence of tokens autoregressively as the network output. More specifically, the autoregressively generated output conditions each particular token in the output sequence on the current input sequence, which includes the input sequence included in the network input 102, and any tokens that precede the particular token in the output sequence, i.e., any tokens already generated at any previous position in the output sequence that precedes the particular position of the particular token.

[0017] As a specific example, the neural network 110 may have any of the following transformer-based neural network architectures, such as an encoder-only transformer architecture, an encoder-decoder transformer architecture, a decoder-only transformer architecture, a diffusion transformer architecture, or other attention-based architectures.

[0018] Examples of such transformer-based neural network architectures include: Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu, Exploring the limits of transfer learning with a unified text-to-text transformer, arXiv preprint arXiv:1910.10683, 2019; Daniel Adiwardana, Minh-Thang Luong, David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha, Gaurav Nemade, Yifeng Lu, and Quoc V. Le, Towards a human-like open-domain chatbot, CoRR, abs / 2001.09977, 2020; and Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell et al., Language models are few-shot learners, arXiv preprint arXiv:2005.14165, 2020, Aakanksha Chowdhery et al., PaLM:Scaling Language Modeling with Pathways, arXiv preprint arXiv:2204.02311, Rohan Anil et al., Palm 2 technical report, arXiv preprint arXiv:2305.10403, 2023, and Gemini Team et al., Gemini:a family of highly capable multimodal models, arXiv preprint arXiv:2312.Includes the transformer-based neural network architecture described in 11805 (2023).

[0019] For example, neural network 110 may be a (generative) language model neural network. Examples of generative language model neural networks include Sparrow (Glaese et al., arXiv:2209.14375), Chinchilla (Hoffmann et al., arXiv:2203.15556), and PaLM 2 (Anil et al., arXiv:2305:10403).

[0020] As another example, neural network 110 may be a multimodal model neural network, such as a visual-language model (VLM) neural network. Examples of multimodal neural networks include Flamingo (Alayrac et al., arXiv:2204.14198), PaLI (Chen et al., arXiv:2209.06794), and PaLI-X (Chen et al., arXiv:2305.18565).

[0021] As another example, a neural network 110 may be a foundational model neural network. A foundational model neural network is a large machine learning model trained on a wide range of datasets that can be adapted and fine-tuned for various applications and downstream tasks. Examples of foundational model neural networks include Imagen (Saharia et al., arXiv:2205.11487) and Parti (Yu et al., arXiv:2206.10789).

[0022] To enable the execution of multiple machine learning tasks, the neural network 110 includes multiple “components.” Each component may be considered a modular neural network component, a separate module of the neural network. Thus, each component may alternatively be called a “module” or “expert.”

[0023] Figure 1A shows that the neural network 110 contains only 12 components 121-132, but generally, the neural network 110 can contain a larger number of components. For example, when the neural network 110 is configured as a generative (large-scale) language model (LLM) or a multimodal model having one of the architectures described above, it may contain hundreds of components.

[0024] The neural network 110 includes multiple parameters. The neural network 110 includes multiple layers connected in any appropriate configuration (for example, as a directed graph of layers). The multiple parameters include parameters representing the layer weights and possibly biases for each of the multiple layers of the neural network 110.

[0025] The types and number of layers included in the neural network 110 depend on the task in which the neural network 110 will be used. However, generally speaking, the layers of the neural network 110 may include, to name just a few, convolutional layers, fully connected layers, recurrent layers, attention layers (e.g., self-attention layers or cross-attention layers), embedding layers, and activation layers (e.g., nonlinear activation layers).

[0026] Furthermore, each component contains a non-overlapping subset of the multiple parameters of the neural network 110. Each component is configured to receive its respective component input and process that input according to its respective subset of the multiple parameters of the neural network 110 to produce its respective component output. Each component input can be either a part of the network input 102 or a component output generated by another component.

[0027] For example, in Figure 1A, component 121 is configured to receive a portion of the network input 102 as component input, process the component input according to a subset of multiple parameters of the neural network 110 contained in component 121, and generate a component output. As another example, in Figure 1A, component 124 is configured to receive the component output generated by component 121 as component input, process the component input according to a subset of multiple parameters of the neural network 110 contained in component 124, and generate a component output.

[0028] In some cases, each of the multiple components of the neural network 110 corresponds to a different layer among the multiple layers of the neural network 110. That is, each layer is a component of the neural network 110. For example, in Figure 1A, component 121 is the first layer of the neural network 110, while component 122 is the second layer of the neural network 110.

[0029] In some other cases, each of the multiple layers of the neural network 110 contains multiple components, that is, each component contains a subset of the parameters of the layer of the neural network 110. For example, in Figure 1A, the first layer of the neural network 110 may contain components 121, 122, and 123, while the second layer of the neural network 110 may contain components 124, 125, and 126. In this particular example, component 121 may contain parameters representing some of the weights of the first layer, component 122 may contain parameters representing the rest of the weights of the first layer, while component 123 may contain parameters representing the bias of the first layer.

[0030] In some other cases, each component contains multiple layers of a neural network. Different components may contain the same or different types or numbers of neural network layers. For example, in Figure 1A, component 121 may contain a first subset of the layers of neural network 110, while component 124 may contain a second subset of the layers of neural network 110, which is placed after the first subset of layers of neural network 110.

[0031] In some other cases, each component contains a subset of the parameters of each of the multiple layers of the neural network. For example, in Figure 1A, component 121 may contain a subset of the layer parameters for each layer of the first subset of layers of the neural network 110, while component 122 may contain another subset of the layer parameters for each layer of the first subset of layers of the neural network 110. Thus, components 121 and 122 may contain different subsets of the parameters of multiple layer groups.

[0032] This configuration of neural networks facilitates the composition of various paths within the neural network. Each path is constructed from a sequence of multiple components. Each path represents an input / output function that maps network inputs to corresponding network outputs.

[0033] For example, each path can contain multiple components selected from a large number of components in a neural network and arranged sequentially, with the output of all components except the last one becoming the input to another component in that network. In other words, each path can contain a suitable subset of all components in the neural network.

[0034] The assembly of the sequence of components into a path by the distributed computing system 100 is dynamic. That is, different paths and correspondingly different components of the neural network 110 are used to perform different machine learning tasks from among multiple machine learning tasks in which the neural network 110 can be used.

[0035] In some cases, different paths and correspondingly different components of the neural network 110 may be used to perform the same machine learning task, but the two different paths may partially overlap, that is, they may contain common components to perform the same task.

[0036] For example, an image processing task may have a different path in the neural network 110 than the agent control task. As another example, when performing a data generation task (e.g., a text generation task, an image generation task, a video generation task, or an audio generation task) to generate an output sequence of tokens as network output, the first set of tokens included in the output sequence may be generated by a different path in the neural network 110 than the second set of tokens included in the output sequence.

[0037] Given a given network input 102 on which a machine learning task is performed, the distributed computing system 100 can use a router 105 to determine one or more corresponding paths, that is, based on the network input 102, it can select which components are included in each of the one or more corresponding paths and which other components are not included in each of the one or more corresponding paths.

[0038] Router 105 includes a set of routing parameters. Router 105 is configured to process at least a portion of the network input 102 according to the set of routing parameters of Router 105 to generate a routing output that specifies a path in the neural network 110.

[0039] In practice, the routing output can take many different forms. For example, the routing output may include a weight vector containing the weight scores for each of the possible paths in the neural network 110, i.e., each of the different combinations of sequences of components contained in the neural network 110 that represent input / output functions that map network inputs to corresponding network outputs. For example, in Figure 1A, the routing output may include a weight vector containing the weight scores for each of the nine possible paths in the neural network 110.

[0040] In this example, the weight vector can be a sparse n-dimensional vector containing non-zero weight scores for only a few of the possible paths in the neural network 110, and the router 105 selects only the possible paths that have non-zero weights in the weight vector, for example, by sampling a specific path from the possible paths according to the weight scores contained in the weight vector.

[0041] The number of non-zero weights may be an integer and may be very small compared to the total number of possible paths in the neural network 110. For example, the components included in the neural network 110 may be assembled to form hundreds or thousands of possible paths in the neural network 110, and the weight vector may have fewer than 10, fewer than 5, or fewer than 2 non-zero weight scores.

[0042] As another example, the routing output may include a weight vector containing the respective weight scores of each component included in each of the multiple “levels” within the neural network 110. For example, in Figure 1A, there may be four levels: a first level containing components 121, 122, and 123; a second level containing components 124, 125, and 126; a third level containing components 127, 128, and 129; and a fourth level containing components 130, 131, and 132.

[0043] In this example, for each level, the weight vector can be a sparse n-dimensional vector containing non-zero weight scores for only a few of the possible components included in the level, and the router 105 selects only the components that have non-zero weights in the weight vector, for example, by sampling specific components from the components according to the weight scores included in the weight vector. By selecting specific components for each level, the router 105 can determine the path in the neural network 110 assembled with the selected specific components.

[0044] Figure 1B shows an example of using router 105 to select a path in an exemplary neural network.

[0045] As shown on the left, in some cases, the distributed computing system 100 uses the router 105 to determine a single path for the network input ("prefix").

[0046] Router 105 processes at least a portion of the network input (i.e., at least a portion of the "prefix") according to a set of routing parameters of Router 105 to generate a routing output that specifies a path in the neural network 110. In Figure 1B, the routing output identifies a second path π2 from among the four possible paths π1 to π4 in the neural network 110. Thus, after the second path π2 is determined based on the network input, the second path π2 is used to generate the network output for the network input.

[0047] Figure 1B shows the output as the training output generated during the training time of the neural network 110, although this is not mandatory. That is, the distributed computing system 100 can also use the router 105 to determine a single path for the network input during inference time.

[0048] As shown on the right, in other cases, the system uses router 105 to determine multiple different paths for a network input ("prefix"). For example, a given machine learning task performed on a network input could be a data generation task, and the network output of a given machine learning task could be an output sequence containing a sequence of tokens.

[0049] The network output consists of multiple parts—that is, multiple "chunks"—each part containing a subset of the sequence of tokens included in the output sequence. In these cases, each different path is used to generate a different part of the network output, i.e., a different subset of the sequence of tokens included in the output sequence.

[0050] Router 105 processes at least a portion of the network input (i.e., at least a portion of the "prefixes") according to a set of routing parameters of Router 105 to generate a first weight vector containing the respective weight scores of each of the four possible paths π1 to π4 in the neural network 110. After the first path π1 is determined based on the network input, the first path π1 is used to generate a first chunk of the network output for the network input.

[0051] Subsequently, router 105 processes a first chunk of the network output according to a set of router 105 routing parameters to generate a second weight vector containing the respective weight scores of each of the four possible paths π1 to π4 in the neural network 110. After the fourth path π4 is determined based on the first chunk of the network output, the fourth path π4 is used to generate the second chunk of the network output for the network input.

[0052] In both of the above cases, after determining the path in the neural network 110, the neural network 110 then generates a network output 112 (or a portion of the network output 112) for the network input 102 using only the components within the determined path, that is, without using any components that are not within the determined path.

[0053] In this way, the neural network 110 is sparsely activated during inference. That is, while the neural network 110 processes the network input 102, only a relatively small number of parameters of the neural network 110 are used to generate the corresponding network output 112. For example, although the neural network 110 contains 1 billion parameters, only 500 million, 250 million, or 100 million or fewer parameters, which are included in the components of the determined path within the neural network 110, are used.

[0054] Figure 1C shows an exemplary path in the neural network 110.

[0055] In the example in Figure 1C, the neural network 110 is configured to perform a machine learning task on the network input 102 to produce a network output 112. For example, the task could be a data generation task, where the network input 102 is an input sequence containing a sequence of tokens selected from a token vocabulary, and the output is an output sequence containing another sequence of tokens selected from the token vocabulary.

[0056] In this example, the token vocabulary may include any of the various tokens that represent text symbols or other symbols. For example, the token vocabulary may include one or more characters, subwords, words, punctuation marks, numbers, or other symbols that appear in a corpus of natural language text and / or computer code.

[0057] Additionally or alternatively, a token vocabulary may include tokens that can represent non-textual data. For example, a token vocabulary may include image tokens that represent a discrete set of image patch embeddings for an image that could be generated by an image encoder neural network based on processing image patches of an image. Another example is that a token vocabulary may include audio tokens that represent code vectors in the codebook of a quantizer, such as a residual vector quantizer.

[0058] As part of performing the data generation task, each token included in the input sequence follows a path in the neural network 110.

[0059] In the example in Figure 1C, a path for the data generation task is shown, with components within the path connected by darker colored arrows, while components not within the path are connected by lighter colored arrows. Thus, as seen in the example, to perform the data generation task, the neural network 110 uses components 121, 124, 128, and 131 to process each token contained in the input sequence. The remaining components of the neural network 110 that are not within the path, 122, 123, 125, 126, 127, 129, 130, and 132, are not used.

[0060] Components 121-132 of the neural network 110 are implemented across different workers included in the distributed computing system 100. A worker may consist of either a single computing device or multiple computing devices. Each computing device may include one or more processor cores, one or more processors, one or more microprocessors, dedicated logic circuits, etc.

[0061] For example, each worker may include one or more central processing units (CPUs), one or more graphics processing units (GPUs), one or more field-programmable gate arrays (FPGAs), one or more tensor processing units (TPUs), or any combination thereof. In a specific example, a worker may include a large number of interconnected TPU islands, for example, hundreds or thousands. In another specific example, a worker may include one or more CPUs and one or more GPUs.

[0062] Figure 1D shows an exemplary worker included in a distributed computing system 100.

[0063] The distributed computing system 100 includes a communication network 150 that connects workers to each other. Examples of the communication network 150 include, to name just a few, a local area network (LAN), such as a data center network, and a wide area network (WAN), such as the internet.

[0064] Each worker includes computing and storage resources. In some cases, each worker is substantially independent of other workers, and the computing and storage resources of a worker are separate and different from each other, for example, so that multiple workers can operate in parallel.

[0065] In some cases, each worker implements an instance of one or more components of the neural network 110. For example, a worker may implement an instance of each subset of multiple components of the neural network 110 that are included in the same path within the neural network 110. In some cases, different paths may overlap, meaning that common components may be included in multiple paths, so different local instances of the same component may be implemented in each of multiple workers.

[0066] For example, as shown in Figure 1D, there are a total of four different paths π1 to π4 in the neural network, and worker 152 of the distributed computing system 100 implements a local instance of the component included in path π1, worker 162 of the distributed computing system 100 implements a local instance of the component included in path π2, worker 172 of the distributed computing system 100 implements a local instance of the component included in path π3, and worker 182 of the distributed computing system 100 implements a local instance of the component included in path π4. In Figure 1D, assuming that the four different paths π1 to π4 contain a common component, each of workers 152, 162, 172, and 182 can implement their respective local instances of the common component.

[0067] In some cases, the distributed computing system 100 includes homogeneous computing devices; that is, all workers are of the same type. However, in other cases, the distributed computing system 100 includes heterogeneous computing devices; that is, the distributed computing system 100 includes two or more different types of computing devices.

[0068] For example, as shown in Figure 1D, worker 152 of the distributed computing system 100 may include one or more GPUs, worker 162 of the distributed computing system 100 may include one or more CPUs, while workers 172 and 182 of the distributed computing system 100 may each include one or more TPUs.

[0069] More broadly, heterogeneity of computing devices in a distributed computing system 100 may arise from computing devices having different hardware capabilities, for example, having different hardware versions, or being manufactured by different hardware manufacturers.

[0070] In some cases, multiple workers are located in the same physical location, such as a data center housing multiple computing devices. In some of these cases, the multiple computing devices are managed by a single organization.

[0071] However, in other cases, multiple workers are physically or geographically separated from one another. For example, as shown in Figure 1D, workers 152, 162, 172, and 182 are in different cities, different states, different countries, and even different continents. In some of these other cases, multiple workers are managed by multiple organizations, for example, a community of organizations.

[0072] To facilitate the training of the neural network, each worker holds training data by utilizing storage resources. The training data includes a set of training inputs and, optionally, for each training input, the respective target outputs that should be generated by the neural network 110 to perform a specific task.

[0073] In some cases, the distributed computing system 100 acquires a training dataset and divides it into multiple shards for storage by each of the multiple workers. Thus, each worker stores a local shard of the training dataset.

[0074] In some cases, local shards of a training dataset held by multiple workers will be approximately the same size as each other, but in other cases, the size of local shards of the training dataset will vary from worker to worker. For example, a worker with more storage resources and therefore greater storage capacity than another worker can store larger local shards of the training dataset containing more training inputs than another worker.

[0075] In some cases, local shards of a training dataset held by multiple workers contain different training inputs from each other, while in other cases, local shards of a training dataset held by two or more of the workers contain at least the same training inputs.

[0076] For example, as shown in Figure 1D, worker 152 holds the first local shard 154 of the training dataset, worker 162 holds the second local shard 164 of the training dataset, worker 172 holds the third local shard 174 of the training dataset, and worker 182 holds the fourth local shard 184 of the training dataset.

[0077] There are many ways in which the training dataset can be partitioned after it has been acquired by the distributed computing system 100. For example, the distributed computing system 100 can partition each of the multiple training inputs contained in the training dataset into one of k shards, based on applying a k-means allocation algorithm to at least a portion of the training inputs.

[0078] The k-means allocation algorithm is an unsupervised algorithm that allocates data into k clusters, where k is equal to the number of workers or the number of worker islands. The k-means allocation algorithm is an iterative algorithm, with each iteration consisting of two phases: 1 - Allocation Phase: Each training input in the training dataset is allocated to one of k clusters where the distance between at least a portion of the training input and the cluster centroid is closest, measured by distance, such as Euclidean distance or cosine distance. 2 - Update Phase: Update the centroids of all k clusters by calculating the arithmetic mean of each cluster.

[0079] The allocation and update phases are then repeated until the final allocation of training inputs is complete.

[0080] As another example, a distributed computing system 100 could assign each of the multiple training inputs contained in a training dataset to one of k shards, based on the premise that it processes at least a portion of the training inputs using a classification machine learning model. Examples of classification machine learning models include, to name just a few, logistic regression models, neural networks, support vector machines, and random forest models.

[0081] For each training input, the classification machine learning model processes at least a portion of the training input to generate a classification output that classifies the training input into one of k shards. For example, the classification output could include a score for each of the k shards, where each score represents the likelihood that the training input belongs to that shard.

[0082] Figure 2 is a flowchart of an exemplary process 200 for training a neural network on a training dataset. For convenience, process 200 is described as being performed by a system of one or more computers located in one or more locations. For example, a well-programmed distributed computing system, such as the distributed computing system 100 in Figure 1A, can perform process 200.

[0083] A neural network consists of multiple components. Each component contains a non-overlapping subset of the neural network's multiple parameters. This configuration of a neural network facilitates the assembly of multiple paths within the network.

[0084] Each path is constructed from a subset of multiple components of the neural network. Each path represents an input / output function that maps network inputs to corresponding network outputs. An example of paths in neural network 110, including components 121, 124, 128, and 131 of neural network 110, is shown in Figure 1A.

[0085] Therefore, each subset of the multiple components included in the path includes at least each input component of the neural network that receives network input (or part of network input) to the neural network, and each output component of the neural network that generates network output (or part of network output) of the neural network. Examples of each input component include components 121, 122, and 123 in Figure 1A. Examples of each output component include components 130, 131, and 132 in Figure 1A.

[0086] A distributed computing system includes multiple workers. These workers correspond to multiple paths in a neural network. That is, with respect to a given path among multiple paths, each containing a given subset of components of the neural network, the corresponding worker can implement a local instance of each component within that subset.

[0087] The system updates each local instance of the component in the corresponding path with each worker, both in parallel with and independently of the other workers (step 202). In effect, each worker updates each local instance of each subset of the multiple components included in the path corresponding to the worker, both in parallel with and independently of the other workers.

[0088] Each worker can train its own local instance of a component of its respective subset by performing multiple internal optimization steps using training inputs taken from a local shard of the training dataset stored in the worker.

[0089] The internal optimization steps are described below with reference to Figure 3, which shows a flowchart of substeps 302-308 of step 202. The system can repeatedly perform multiple internal optimization steps to train each local instance of each subset of components by repeatedly executing iterations of steps 302-308 in each worker.

[0090] The system takes a batch of training inputs and, for each training input in the batch, processes the training input according to the parameters of the respective local instances of each component in the respective subset to generate a training output for the training input (step 302). The system may take a batch of training inputs by any preferred method, for example, by sampling. The system refrains from using any other components of the neural network that are not in the respective subset, i.e., not included in the path corresponding to the worker, to process the training inputs.

[0091] The system evaluates an objective function (Step 304) that measures the quality of the training output generated for a batch of training inputs. The objective function can be any suitable differentiable objective function that is appropriate for the training data, i.e., it measures the quality of the training output generated by the neural network for a given training input, relative to the target output for a given training input. Examples of objective functions include the cross-entropy loss function, the squared error loss function, and the negative log-likelihood loss function.

[0092] The system calculates the gradient of the objective function with respect to the parameters contained in each local instance of each component of each subset (step 306). The system can calculate the gradient using gradient-based training techniques, for example, by backpropagation.

[0093] The system determines to update the current values ​​of the parameters of the local instances of each component in each subset based on applying an optimizer to the gradient (step 308). Examples of optimizers include the Adam optimizer, AdamW optimizer, Adafactor optimizer, rmsProp optimizer, and stochastic gradient descent optimizer.

[0094] In some cases, the system performs path-specific early stopping by using fewer training inputs that are removed and not used during the internal optimization step. For each worker, the system decides whether to apply early stopping based on the quality of the training output generated from the removed training inputs according to the parameters of each local instance of the components of each subset. For example, the system may terminate the training of each local instance of a component in a particular path corresponding to a particular worker early, when the quality, determined based on evaluating the objective function against the removed training inputs, meets a quality threshold, i.e., while the training of neural network components in other workers continues.

[0095] In practice, transmitting data between workers over a communication network is expensive, for example, in terms of time, resource usage, or power consumption. Firstly, the performance of the communication network may be poor, for example, the connection between two or more workers may not be established on the first attempt (thus requiring multiple retries) and / or may be lost. Secondly, the physical distance between these multiple working computing devices may be long.

[0096] Therefore, when performing multiple internal optimization steps as part of step 202 and updating local instances of each component of their respective subset in each worker, the multiple workers do not need to communicate with each other via the communication network connecting them, and in some cases do not communicate at all. By not communicating with each other after each internal optimization step during step 202, the system reduces the amount of communication required during neural network training.

[0097] Because two different workers train the same components included in two different paths with different training data, the different local instances of the same components of the neural network implemented in the different workers will have different parameter values ​​as a result of internal optimization steps performed independently in each worker.

[0098] The system determines whether to synchronize the respective local instances of the components of each subset across multiple workers (step 204). Such a decision can be made at any time while step 202 is being executed.

[0099] For example, the system may count the number of internal optimization steps performed in each worker (during the current iteration of process 200), and when the counted number reaches a predetermined number, it may decide that the respective local instances of the respective subset components across multiple workers should be synchronized.

[0100] As another example, the system might count the amount of time that has elapsed since the first internal optimization step was executed in each worker (during the current iteration of process 200), and when the counted amount of time reaches a predetermined amount, it might decide that the respective local instances of the respective subset components across multiple workers should be synchronized.

[0101] In response to the decision to synchronize, the system performs an outer optimization step to update each local instance of the component in each subset implemented on multiple workers so that the different local instances of the same component of the neural network implemented on different workers have the same values ​​as each other (step 206). In practice, the outer optimization step may be performed much less frequently than the internal optimization step. For example, each iteration of process 200 may include 100, 200, 500, or more internal optimization steps for each worker, while each component may include only one outer optimization step.

[0102] The external optimization step is described below with reference to Figure 4, which shows a flowchart of substeps 402-408 of step 206. The system can perform the external optimization step by iterating through steps 402-406 for each of the multiple components of the neural network.

[0103] The system determines, in each of the one or more workers holding each local instance of the component, the respective difference between (i) the previous value of the parameter of each local instance of the component before step 202 (in the current iteration of process 200), for example, before a predetermined number of internal optimization steps, and (ii) the updated value of the parameter of each local instance of the component after step 202 (in the current iteration of process 200), for example, after a predetermined number of internal optimization steps (step 402).

[0104] The system determines an aggregated difference from the individual differences determined for one or more workers (step 404). The aggregated difference can be determined in many different ways, for example, as a weighted or unweighted sum or average of the individual differences.

[0105] As a specific example, the system can determine the aggregated difference by calculating a weighted average of the individual differences determined for one or more workers, where each individual difference determined for each of the one or more workers is weighted by the number of training inputs contained in the shards held by the worker. In this way, each difference for a worker holding a larger number of training inputs is assigned a greater weight than each difference for another worker holding a smaller number of training inputs.

[0106] In some cases, the system may apply norm rescaling to the aggregated difference and use the norm-rescaled aggregated difference as the aggregated difference. Norm rescaling takes into account the variability in the number of paths each component may belong to. For example, the system may scale the norm of the aggregated difference as a function of the number of workers each local instance of the component holds, such as the square root, division, or reciprocal.

[0107] In some cases, the system may apply loss reweighting to the aggregated differences and use the loss-reweighted aggregated differences as the aggregated differences. Loss reweighting takes into account the variability in the size of local shards of the training data held across multiple workers, and can, for example, account for over-sampling of local shards with smaller sizes. For example, the system may weight the aggregated differences in proportion to the size of the local shards held by each worker by calculating the following equation (the definitions of the symbols are consistent with the definitions given further below in Algorithm 1):

number

[0108] In some cases, the system may apply both norm rescaling and loss reweighting to the aggregated difference, and use the norm-rescaled and loss-reweighted aggregated difference as the aggregated difference.

[0109] The system determines synchronized and updated values ​​for the component parameters based on applying an optimizer to the aggregated difference (step 406). Examples of optimizers include the Nesterov Momentum optimizer, Adam optimizer, AdamW optimizer, Adafactor optimizer, rmsProp optimizer, and stochastic gradient descent optimizer.

[0110] In some cases, the optimizer applied in the internal optimization step is the same as the optimizer applied in the external optimization step; however, in other cases, the optimizer applied in the internal optimization step is different from the optimizer applied in the external optimization step. As a specific example, the system applies the AdamW optimizer in the internal optimization step and the Nesterov Momentum optimizer in the external optimization step.

[0111] The system sends data representing the synchronized updated values ​​of the component's parameters to one or more workers so that each of the workers can update the local instance of the component implemented in the worker so that it has the same synchronized updated values ​​(step 408).

[0112] By repeatedly executing process 200 iterations, the system can train a neural network on the training dataset to update the parameter values—and thus all the parameter values—of each of the neural network's multiple components across multiple workers, by using local shards of the training dataset held by each worker.

[0113] An exemplary algorithm for training a neural network on a training dataset is shown below.

[0114] [Table 1]

[0115] Figure 5 is a flowchart of an exemplary process 500 for using a neural network to generate a network output from a network input. For convenience, process 500 is described as being performed by a system of one or more computers located in one or more locations. For example, a well-programmed distributed computing system, e.g., the distributed computing system 100 in Figure 1A, can perform process 500.

[0116] A neural network contains multiple components. A system contains multiple workers. Multiple components are implemented on different workers.

[0117] The system receives a network input which is processed by a neural network to generate a network output (step 502). The network output may consist of multiple parts. For example, the network output may be an output sequence containing a sequence of tokens, where each part contains a subset of the sequence of tokens contained in the output sequence.

[0118] The system selects a path in the neural network from among several paths in the neural network for each portion of the network output (step 504). Each path contains an appropriate subset of each of the multiple components. Each appropriate subset of the multiple components contains at least each input component and each output component. In some cases, the same path is selected for multiple portions of the network output, and in other cases, different paths are selected for multiple portions of the network output.

[0119] For each portion of the network output, each path in the neural network may be selected by the system based on processing either a portion of the network input or a portion of the network output using a router. For example, the portion of the network output processed by the router to select each path could be the portion of the network output immediately preceding the portion of the network output that precedes that portion of the network output. If the portion of the network output is the beginning of the network output, there is no preceding portion, so the router can instead process at least a portion of the network input to select each path.

[0120] With respect to each portion of the network output, the system generates that portion of the network output using appropriate subsets of the multiple components included in each path (step 506). For example, that portion of the network output may be generated by the system using subsets of the multiple components included in each path, based on processing (i) the network input, (ii) one or more preceding portions of already generated network output, or both (i) and (ii).

[0121] In particular, the neural network generates each portion of the network output using only the components of the neural network that are within the appropriate subset included in each path, that is, without using any of the remaining unselected components that are excluded from each path and therefore not included in the appropriate subset.

[0122] Therefore, the network output can be an output sequence containing a sequence of tokens generated using a smaller subset of all the parameters of the neural network. Furthermore, when different paths are selected for multiple parts of the network output, the network output can be an output sequence containing a sequence of tokens generated by different paths of the neural network—and thus, different combinations of components.

[0123] This specification uses the term “configured” in relation to systems and computer program components. One or more computer systems being configured to perform a particular operation or action means that the system has installed software, firmware, hardware, or a combination thereof that causes the system to perform the operation or action while in operation. One or more computer programs being configured to perform a particular operation or action means that one or more programs contain instructions that cause a data processing device to perform the operation or action when executed by the device.

[0124] The embodiments and functional operations described herein can be implemented in digital electronic circuits, tangibly embodied computer software or firmware, computer hardware, or one or more combinations thereof, including the structures disclosed herein and their structural equivalents. The embodiments described herein can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-temporary storage medium for execution by a data processing device or for controlling the operation of a data processing device. The computer storage medium can be a machine-readable storage device, a machine-readable storage board, a random or serial access memory device, or one or more combinations thereof. Alternatively or additionally, the program instructions can be encoded on artificially generated propagating signals, such as machine-generated electrical, optical, or electromagnetic signals, which are generated to encode information for transmission to a receiver device suitable for execution by a data processing device.

[0125] The term “data processing device” refers to data processing hardware and encompasses all types of devices, machines, and equipment for processing data, including, for example, a single programmable processor, a single computer, or multiple processors or computers. A device may also be a dedicated logic circuit, such as an FPGA (Field-Programmable Gate Array) or ASIC (Application-Specific Integrated Circuit), or may further include such dedicated logic circuits. Optionally, in addition to hardware, a device may include code that creates an execution environment for computer programs, such as processor firmware, a protocol stack, a database management system, an operating system, or code that constitutes one or more of these.

[0126] Computer programs, sometimes called or described as programs, software, software applications, apps, modules, software modules, scripts, or code, can be written in any form of programming language, including compiled languages, interpreted languages, declarative languages, or procedural languages, and can be deployed in any form, including as standalone programs or as modules, components, subroutines, or other units suitable for use within a computing environment. A program may, but not necessarily, correspond to a file in a file system. A program may be stored in part of a file that holds other programs or data, for example, one or more scripts stored in a markup language document, a single file dedicated to the program in question, or multiple organized files, for example, one or more modules, subprograms, or files that store parts of code. A computer program can be deployed to run on one computer, located in one place, or distributed across multiple computers connected to each other by a data communication network.

[0127] In this specification, the term “database” is used broadly to refer to any collection of data, that is, data that does not need to be structured in any particular way, or does not need to be structured at all, and can be stored on a storage device in one or more locations. Thus, for example, an index database can contain multiple collections of data, each of which may be organized and accessed in a different way.

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

[0129] The processes and logic flows described herein may be executed by one or more programmable computers running one or more computer programs that perform functions by performing calculations on input data and generating outputs. Alternatively, the processes and logic flows may be executed by dedicated logic circuits, such as FPGAs or ASICs, or by a combination of dedicated logic circuits and one or more programmed computers.

[0130] A computer suitable for running computer programs can be based on a general-purpose microprocessor, a dedicated microprocessor, or both, or any other type of central processing unit. Generally, the central processing unit receives instructions and data from read-only memory, random-access memory, or both. Essential elements of a computer are the central processing unit for executing or running instructions, and one or more memory devices for storing instructions and data. The central processing unit and memory can be complemented by or incorporated into dedicated logic circuits. Also generally, a computer includes one or more mass storage devices for storing data, such as magnetic disks, magneto-optical disks, or optical disks, or is coupled to operate in a manner that it can receive data from such mass storage devices, transfer data to such mass storage devices, or both. However, a computer does not have to have such devices. Furthermore, computers can be integrated into other devices, such as mobile phones, personal digital assistants (PDAs), mobile audio or video players, game consoles, Global Positioning System (GPS) receivers, or portable storage devices, such as Universal Serial Bus (USB) flash drives.

[0131] Computer-readable media suitable for storing computer program instructions and data include, for example, all forms of non-volatile memory, media, and memory devices, such as semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto-optical disks; and CD-ROM and DVD-ROM disks.

[0132] To provide user interaction, the embodiments described herein may be implemented on a computer having a display device for displaying information to the user, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, and a keyboard and pointing device, such as a mouse or trackball, on which the user can provide input to the computer. Other types of devices may also be used to provide user interaction; for example, the feedback provided to the user may be any form of sensory feedback, such as visual feedback, auditory feedback, or haptic feedback, and input from the user may be received in any form, including acoustic, speech, or tactile input. In addition, the computer may interact with the user by sending documents to and receiving documents from a device used by the user, for example, by sending a web page to the user's device's web browser in response to a request received from a web browser. The computer may also interact with the user by sending text messages or other forms of messages to a personal device running a messaging application, such as a smartphone, and receiving a response message from the user in return.

[0133] Data processing devices for implementing machine learning models may include, for example, dedicated hardware accelerator units for handling the computationally intensive parts of the machine learning training or generation, i.e., inference workload.

[0134] Machine learning models can be implemented and deployed using machine learning frameworks, such as the TensorFlow framework or the JAX framework.

[0135] The embodiments of the subject matter described herein can be implemented in a computing system that includes, for example, a data server as a backend component, or a middleware component, for example, an application server, or a client computer having a graphical user interface, a web browser, or an application that allows a user to interact with the subject implementation described herein, or one or more such backend components, middleware components, or frontend components in any combination. The components of the system can be interconnected by digital data communication in any form or medium, for example, a communication network. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), for example, the Internet.

[0136] A computing system may include a client and a server. The client and server are generally geographically separated and typically interact through a communication network. The relationship between the client and the server arises from computer programs running on each computer that are in a client-server relationship with each other. In some embodiments, the server transmits data, such as an HTML page, to the user device for the purpose of displaying data to a user interacting with a device acting as a client and receiving user input from such a user. Data generated on the user device, such as the results of user interactions, can be received from the device by the server.

[0137] This specification includes many specific implementation details, but these should not be considered as limitations on the scope of any invention or claim, but rather as descriptions of features that may be specific to a particular embodiment of a particular invention. Certain features described herein in the context of separate embodiments may be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment may be implemented in multiple embodiments separately or in any preferred partial combination. Furthermore, features may be described above as working in a particular combination, and may even be initially claimed as such, but one or more features of a claimed combination may, in some cases, be removed from the combination, and the claimed combination may cover a partial combination or a variation of a partial combination.

[0138] Similarly, although the operations are depicted in the figures in a specific order and described in the claims, this should not be understood as requiring that such operations be performed in a specific or sequential order as shown, or that all operations shown be performed to achieve the desired result. In certain circumstances, multitasking and parallel processing may be advantageous. Furthermore, the division of various system modules and components in the embodiments described above should not be understood as requiring such division in all embodiments, and it should be understood that the described program components and systems can generally be integrated together into a single software product or packaged into multiple software products.

[0139] A specific embodiment of the subject has been described. Other embodiments are within the scope of the appended claims. For example, the actions described in the claims may be performed in a different order and still achieve the desired results. As an example, the process depicted in the appended diagrams does not necessarily require the specific order or sequential order shown to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous. [Explanation of symbols]

[0140] 100 Distributed Computing Systems 102 Network Input 105 Router 110 Neural Networks 112 Network Output Components 121-132 150 Communication Networks 152 workers 154 First local shard 162 workers 164 Second local shard 172 workers 174 Third local shard 182 workers 184 The 4th local shard 200 processes 500 processes

Claims

1. A method for training a neural network on a training dataset using a distributed computing system that includes multiple workers, A step in which each of the plurality of workers holds a local instance of each component of each subset of the plurality of components of the neural network that represent paths in the neural network, wherein each component of the neural network includes a different subset of the plurality of parameters of the neural network, and each subset of the plurality of components that represent paths in the neural network includes at least an input component of the neural network that receives input to the neural network, and an output component of the neural network that generates the output of the neural network. Each worker updates the respective local instance of the component of the respective subset in parallel with the other workers. The steps include determining whether to synchronize the respective local instances of the components of each subset across the plurality of workers, The steps include updating the respective local instances of the components of the respective subsets implemented in the plurality of workers in response to the decision to synchronize, and Methods that include...

2. Each worker updates the local instance of each component of the neural network in each subset. The steps include: maintaining local shards of the training dataset, The steps include updating the respective local instances of the components of each subset based on the local shards of the training dataset, and The method according to claim 1, including the method described in claim 1.

3. Each worker updates the local instance of each component of the neural network in each subset. The method according to claim 1, comprising the step of allocating each of the multiple training inputs included in the training dataset to one of k shards based on applying a k-means allocation algorithm to at least a portion of the training inputs.

4. Each worker updates the local instance of each component of the neural network in each subset. The method according to claim 1, further comprising the step of assigning each of a plurality of training inputs contained in a training dataset to one of k shards based on processing at least a portion of the training inputs using a logistic regression classifier.

5. The method according to any one of claims 2 to 4, wherein the local shards of the training dataset held by the plurality of workers include different training inputs from each other.

6. The method according to any one of claims 2 to 4, wherein the local shards of the training dataset held by two or more of the plurality of workers include at least the same training input.

7. Each worker updates the local instance of each component of the neural network in each subset, in each of the multiple internal optimization steps. A step of evaluating an objective function that measures the quality of the training output generated according to the parameters of each local instance of each component of each subset, The steps include: calculating the gradient of the objective function with respect to each local instance of each component of each subset; The steps include determining an update to the current value of the parameter of each local instance of each component of each subset, based on applying the AdamW optimizer to the gradient, and The method according to any one of claims 1 to 6, including the method described in any one of claims 1 to 6.

8. The step of determining whether to synchronize the respective local instances of the components of the neural network in each subset across the plurality of workers is: The method according to claim 7, further comprising the step of determining whether a predetermined number of internal optimization steps have been performed by each worker to update the respective local instances of each component of the neural network in each subset.

9. The step of updating the local instance of the component of the neural network implemented in each of the plurality of workers is, with respect to each of the plurality of components of the neural network, In each of the one or more workers holding each local instance of the component, the steps include: (i) determining the difference between the previous value of the parameter of each local instance of the component before a predetermined number of internal optimization steps and (ii) determining the updated value of the parameter of the local instance of the component after the predetermined number of internal optimization steps; A step of determining an aggregated difference from the respective differences determined with respect to one or more workers, The steps include determining the synchronized updated values ​​of the parameters of the component based on applying the Nesterov Momentum Optimizer to the aggregated difference, The steps include: transmitting data representing the synchronized updated values ​​of the parameters of the component to one or more workers; The method according to any one of claims 1 to 8, including the method described in any one of claims 1 to 8.

10. The method according to claim 9, wherein the step of determining the aggregated differences includes the step of determining a weighted average of the respective differences determined with respect to the one or more workers, wherein the respective differences determined with respect to each of the one or more workers are weighted by the number of training inputs contained in the shards held by the workers.

11. The method according to claim 9 or 10, wherein the step of determining the aggregated difference includes scaling the norm of the aggregated difference by the square root of the number of workers each local instance of the component holds.

12. Each worker updates the local instance of each component of the neural network in each subset. The method according to any one of claims 1 to 11, further comprising the step of determining whether to apply early termination based on the quality of the training outputs generated from one or more training inputs according to the parameters of the respective local instances of the components of the respective subset, by each worker.

13. The method according to any one of claims 1 to 12, wherein the neural network includes a transformer neural network.

14. The method according to any one of claims 1 to 13, wherein the neural network includes a large-scale language model (LLM) or a visual language model (VLM) neural network.

15. The method according to any one of claims 1 to 14, wherein the plurality of workers include two or more different types of computing devices.

16. The method according to any one of claims 1 to 15, wherein the plurality of workers are physically, geographically, or physically and geographically separated from one another.

17. The method according to any one of claims 1 to 16, wherein when updating the local instance of each component of the neural network in each subset, the plurality of workers do not communicate with one another via a communication network connecting the plurality of workers.

18. A method performed by one or more computers, A step of receiving a network input to be processed by a neural network to generate a network output for a machine learning task, wherein the neural network comprises a plurality of components, the plurality of components are implemented on different workers, and the network output comprises a plurality of parts. A step of selecting a path in the neural network from a plurality of paths in the neural network for each portion of the network output, wherein each path includes a suitable subset of the plurality of components, and each suitable subset includes at least the respective input component and the respective output component. With respect to each portion of the network output, the step of generating the portion of the network output using only the appropriate subset of the plurality of components included in each of the paths: Methods that include...

19. The method according to claim 18, wherein the step of selecting the respective path in the neural network for each portion of the network output includes the step of selecting the respective path based on using a router to process either a portion of the network input or a portion of the network output.

20. A system comprising one or more computers and one or more storage devices that store instructions causing the one or more computers to perform the operation of each of the methods described in any one of claims 1 to 19 when executed by the one or more computers.

21. One or more computer storage media that store instructions causing one or more computers to perform the operation of each of the methods described in any one of claims 1 to 19 when executed by one or more computers.

Images