Methods, systems, and media for distributing tensor computations across computing devices

Abstract

Methods, systems, and media for distributing tensor computations across computing devices. One of the methods includes receiving specification data that specifies a distribution of a tensor computation among a plurality of computing devices, where each tensor computation (i) is defined to receive as input one or more respective input tensors each having one or more respective input dimensions, (ii) is defined to generate as output one or more respective output tensors each having one or more respective output dimensions, or both, where the specification data specifies for each input and output tensor a respective layout that assigns each dimension of the input or output tensor to one or more of the plurality of computing devices; assigning to each of the computing devices a respective device-local operation based on the layout of the input and output tensors; and causing the tensor computation to be performed.

Description

[0001] Case Analysis

[0002] This application is a divisional application of Chinese invention patent application 201980017284.6, filed on August 5, 2019.

[0003] Cross-references to related applications

[0004] This application claims priority to U.S. Patent Application No. 62 / 714,586, filed August 3, 2018, entitled “DISTRIBUTING TENSOR COMPUTATIONS ACROSS COMPUTING DEVICES,” pursuant to 35 USC §119(e), which is incorporated herein by reference. Background Technology

[0005] This specification relates to distributing tensor computations across multiple computing devices for execution. For example, tensor computations may be operations used to train neural networks or other machine learning models.

[0006] A neural network is a machine learning model that uses one or more layers of non-linear units to predict the output of a received input. In addition to the output layer, some neural networks also include one or more hidden layers. The output of each hidden layer serves as the input to the next layer in the network (i.e., the next hidden layer or output layer). Each layer of the network generates an output from the received input based on the current values ​​of its corresponding set of parameters. Summary of the Invention

[0007] This specification describes a system implemented as a computer program on one or more computers located at one or more locations, which distributes tensor computation across computing devices in response to a received request. Tensor computation receives one or more input tensors, each having one or more corresponding input dimensions, as input. Alternatively, tensor computation can be defined as generating one or more output tensors, each having one or more corresponding output dimensions, as output. In this specification, tensor computation that receives or generates tensors means that a computing device configured to perform tensor computation performs tensor computation to receive one or more input tensors, generate one or more output tensors, or both, depending on the definition of tensor computation.

[0008] Each computing device includes at least one processor and memory for storing tensors. The computing device may include a processor core, processor, microprocessor, special-purpose logic circuit system (e.g., FPGA (Field-Programmable Gate Array) or ASIC (Application-Specific Integrated Circuit) or any other suitable computing device. In some examples, the computing devices are all of the same type. In other examples, the computing device may include different types of processing units. For example, one device may include one or more CPUs (Central Processing Units), while other devices may include one or more GPUs (Graphics Processing Units) or TPUs (Tensor Processing Units).

[0009] For example, tensor computation can be an operation used to train neural networks or other types of machine learning models, or an operation used to perform inference using neural networks or other machine learning models.

[0010] Specifically, the system allows users to submit specification data that specifies the layout of tensors involved in tensor computations, i.e., how they should be distributed across each dimension of each tensor involved in the computation. For example, a user can submit specification data that specifies for each dimension of each tensor whether the dimension should be split across at least some devices in the device or replicated across all devices in the device.

[0011] Based on the layout, the system compiles tensor computations into device-local operations, and, if necessary, into communication primitive operations that induce communication between devices to combine the outputs of device-local operations.

[0012] The system can then enable tensor computation by having each device perform a corresponding device-local operation (and, if necessary, a communication primitive operation) on the input received by the device. The system can cause the devices to perform the corresponding device-local operations in any of a variety of ways, for example, by sending appropriate instructions or control signals to the devices that cause the devices to perform each device-local operation assigned to them when the input for the device-local operation becomes available.

[0013] When tensor computation is used to train a machine learning model, the user can specify the layout of the tensors processed (i.e., received or generated) during iterations of the machine learning training process performed on a batch of training data. The system can then enable the repeated performance of tensor computations on several different batches of training data to iteratively update the parameters of the machine learning model until it converges or until some other condition is met.

[0014] Specific embodiments of the subject matter described in this specification may be implemented to achieve one or more of the following advantages.

[0015] By allowing users to specify how tensor computations are distributed simply by specifying the layout of the relevant tensors, complex distribution schemes can be defined in a simple way. Therefore, users can effectively implement efficient distribution schemes for performing tensor computations (e.g., model training operations) in a way that leads to one or more efficient uses of the device's computing power, improving performance or reducing computation time.

[0016] The methods described herein for specifying distribution schemes can utilize the computational power of computing devices more efficiently by minimizing communication between computing devices (e.g., by leveraging existing physical connections between computing devices in a grid). Communication between computing devices is typically more expensive than computation and is often a bottleneck in parallel programs, especially in distributed settings. The methods described herein can exhibit finer-grained parallelism compared to other parallelization strategies, such as, for example, master-slave computation strategies.

[0017] By using the techniques described in this specification for specifying distribution schemes, schemes for enabling parallel training of very large machine learning models (i.e., models unsuitable for a single computing device) can be specified and implemented. This allows models to be trained to achieve improved performance on any of a variety of machine learning tasks relative to existing technologies, for example, because models with more parameters (e.g., more hidden units) than conventional models can be trained efficiently. Furthermore, by utilizing the computational power of all available devices, models capable of handling very large inputs can be efficiently implemented for both training and inference. When inference is performed, the system can execute the model with lower latency than conventional methods because operations are distributed across computing devices.

[0018] Furthermore, by specifying the distribution scheme in the manner described in this specification, a given distribution scheme can be easily generalized to other hardware or model types. Therefore, once an effective distribution scheme is found, it can be easily generalized to efficiently distribute other similar tensor computations.

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

[0020] Figure 1 It is a block diagram of a computing system that includes multiple computing devices. Figures 2A to 2D An example layout of a two-dimensional tensor laid out on a two-dimensional grid is shown.

[0021] Figure 3It is a flowchart of an example process for performing operations on multiple computing devices based on canonical data of a distribution calculated from a specified tensor.

[0022] The same reference numerals and designations in the various figures indicate the same elements. Detailed Implementation

[0023] Figure 1 This is a block diagram of a computing system 100, which includes a plurality of computing devices interconnected to form a mesh 110. The computing system 100 may also include a layout engine 120, a computational graph engine 130, and a mesh execution engine 140. For the purposes of this specification, a mesh is an n-dimensional array of interconnected computing devices. For example, mesh 110 is a two-dimensional array of interconnected computing devices. The two-dimensional array of computing devices in mesh 110 may be said to have a shape [4,4], representing a 4x4 array of computing devices.

[0024] Different grids can be defined for the same set of physical devices, and therefore, the shape of the grid does not imply the physical network topology of how the computing devices are interconnected. Instead, the shape of the grid refers to the number of computing devices represented by each dimension of an n-dimensional array, generally denoted as [k1, k2, ..., k]. n ], where each k i (i≤n) represents the length of the i-th dimension in the grid.

[0025] For example, a cluster of 512 core Tensor Processing Units (“TPUs”) with a 16x16x2 ring network interconnect can be represented as: a three-dimensional array of shape [16,16,2]; a two-dimensional array of shape [32,16]; or a one-dimensional array of shape

[512] . A mesh can be defined for computing devices with different types of processing units (e.g., central processing units (CPUs), graphics processing units (GPUs), and TPUs).

[0026] Optionally, the computing system 100 may also receive data for a specified number of operations 180. In some implementations, the computing system 100 is configured to use computing devices in the grid 110 and perform operations 180 on the input data 150.

[0027] In some implementations, computing system 100 is not configured to receive data for a specified number of operations 180. In other implementations, computing system 100 is configured to use computing devices in grid 110 to perform a number of predetermined operations.

[0028] As described in more detail below, the computing system 100 processes the input data 150 according to one or more parallel techniques specified by the layout parameter value 160, and generates output data 170 corresponding to the processed input data.

[0029] The grid execution engine 140 can receive a layout specifying how tensors and operations should be distributed on the grid 110. A tensor is a multidimensional array of numerical or other values ​​(e.g., strings) with a specific order corresponding to the dimensions of the array. For example, scalar values ​​are 0-order tensors, numerical vectors are 1-order tensors, and value matrices are 2-order tensors.

[0030] As described below, layout engine 120 can generate a layout based on layout parameter values ​​160 provided to computing system 100. Mesh execution engine 140 can receive and assign tensor data associated with input data 150, parameter values ​​for different operations (e.g., weights for a neural network), and can assign data and operations according to the layout. The layout generated by layout engine 120 is collectively referred to as "canonical data" because the corresponding layout of the tensors specifies how mesh execution engine 140 will distribute or "lay out" the tensors on mesh 110.

[0031] Depending on the layout, the mesh execution engine 140 can be configured to implement different parallel techniques for processing the input data 150 on the computing devices within the mesh 110. For example, the layout engine 120 can generate a layout and send it to the mesh execution engine 140, which assigns the dimensions of the corresponding tensors and operations to different computing devices to implement data parallelism. Data parallelism refers to a class of techniques in which input data is partitioned and distributed across multiple computing devices, but each computing device performs the same operation on different data.

[0032] In some implementations, layout engine 120 can generate a layout and send it to mesh execution engine 140, which in turn distributes the corresponding tensors and operations to implement model parallelism. In these implementations, the operations performed on mesh 110 are operations used to process input data 150 through a machine learning model. Model parallelism refers to a class of techniques in which input data is copied on each of multiple computing devices, wherein each computing device performs different operations of a machine learning model on a copy of the same data.

[0033] Therefore, in the data parallelism implementation, data is partitioned and distributed across the computing devices in mesh 110; and in the model parallelism implementation, operations are partitioned and distributed across the computing devices. In some implementations, the layout engine 120 can generate and send the corresponding tensors and operations to the mesh execution engine 140 to implement both data parallelism and model parallelism in the layout.

[0034] Turning to the formal details and definitions of the layout generated by the layout engine 120, a tensor can be said to be "laid out" on the grid of a computing device when the mesh execution engine 140 assigns a slice of the tensor to each computing device according to a specific layout. A slice of a tensor is a sub-tensor of the tensor and can be the tensor itself. As described below with examples, each slice of a tensor assigned to a corresponding computing device does not need to be unique. For example, a tensor can be cut into several slices, and a particular slice can be laid out on each computing device in the grid 110.

[0035] The mesh execution engine 140 can slice tensors along one or more dimensions. Slicing along a dimension means dividing the tensor into sub-tensors along the dimension. For example, consider a two-dimensional tensor where dimension 0 passes horizontally through the tensor and dimension 1 passes vertically through the tensor. If the tensor is sliced ​​along dimension 0, the slice will be a sub-tensor whose values ​​cross the tensor horizontally. Similarly, if the tensor is sliced ​​along dimension 1, the sub-tensors will have values ​​that cross the tensor vertically.

[0036] Typically, tensors and meshes can have any number of dimensions, and tensors can be sliced ​​along any tensor dimension and laid out along any mesh dimension. In this specification, the dimensions of a tensor will be referred to as tensor dimensions, and the dimensions of a mesh will be referred to as mesh dimensions.

[0037] The number of slices a tensor is cut into depends on the number of computing devices in the grid dimension on which the tensor is laid out. For example, if the grid is a two-dimensional array of computing devices with two computing devices along a dimension, then if a tensor is laid out in that dimension, the computing system 100 can cut the tensor into two sub-tensors of equal size. Similarly, if the grid is a two-dimensional array of computing devices with six computing devices along a dimension, the computing system 100 can cut the tensor into six sub-tensors of equal size.

[0038] In some implementations, the mesh execution engine 140 can slice and assign corresponding slices of tensors with different shapes. For the purposes of this specification and unless otherwise stated, when a tensor is sliced ​​into multiple slices, each slice has the same shape.

[0039] So far, layout has generally been referred to as the format by which the mesh execution engine 140 distributes tensors and operations across computing devices in the mesh 110. More specifically, the layout engine 120 can generate separate layouts for the corresponding tensors, such as tensors representing input data 150, output data 170, or any intermediate input and output data generated by processing operations assigned to the mesh 110.

[0040] The layout engine 120 can also generate separate layouts for operations to be processed on computing devices within the grid 110. For example, an operation can be represented by one or more weight values ​​stored in a corresponding weight tensor. Each computing device can be configured to perform an operation (e.g., matrix multiplication) by multiplying all or part of the input tensor with all or part of the weight tensor. By specifying the layout of the weight tensors, operations can be distributed across computing devices in the grid in the same way as tensors that distribute the storage of input or output data. Formally, a layout is an injective local mapping from the tensor dimensions of a k-dimensional tensor to the grid dimensions of an n-dimensional grid. The layout defines which dimensions of the tensor are split along the grid. A layout can be legal or illegal. Depending on the layout, if the grid execution engine 140 causes some slices of the tensor to be lost when it places the tensor outside the grid 110, the layout is illegal. Otherwise, the layout is considered legal.

[0041] The layout of an n-dimensional tensor laid out on a k-dimensional grid can be represented as an n-tuple: <d0,d1,d2,...d n >, where each d i It is an empty identifier (represented as) ) or an identifier for some integer value between 0 and k. Each d i The identifier indicates which grid dimension is used to lay out the tensor dimension i. If d i Empty, that is, equal to Instead of laying out the tensor along its dimension i, the values ​​stored along dimension i are copied across each computing device.

[0042] For example, the layout of 3D tensors across a 3D mesh layout express:

[0043] • Slice the tensor along tensor dimension 0 and assign the corresponding slice to each computing device along grid dimension 1;

[0044] • Slice the tensor along tensor dimension 1, and assign the corresponding slice to each computing device along grid dimension 0; and

[0045] • Not following empty identifiers The tensor is sliced ​​in the specified tensor dimension 2. Conversely, the tensor is copied completely in the grid dimension 2.

[0046] Figures 2A to 2DExample layouts of two-dimensional tensors laid out on a two-dimensional grid are shown. In these examples, each value in the two-dimensional tensor represents the corresponding pixel of image 210, and the tensor and grid are indexed starting from dimension 0. Dimension 0 is the "horizontal" dimension of the grid in these examples, and dimension 1 is the "vertical" dimension of the grid.

[0047] Figure 2A The tensor of the layout according to the empty layout 230 is shown on the grid 220 of the four computing devices A to D. The empty layout replicates the entire tensor and is composed of tuples. express.

[0048] Figure 2B The diagram shows tensors laid out on a grid 220 of four computing devices A through D, where the tensors are sliced ​​along tensor dimension 0, and the slices are laid out on grid dimension 0. Tensors are not sliced ​​or laid out across grid dimension 1. The tuple representing this layout is...

[0049] Figure 2C Tensors are shown laid out on a grid 220 of four computing devices A to D, wherein the tensors are sliced ​​along tensor dimension 0, and the slices are laid out on grid dimension 0. Tensors are also sliced ​​along tensor dimension 1, wherein the slices are laid out on grid dimension 1. Figure 2C The layout shown is represented as <0, 1>.

[0050] Figure 2D The diagram shows tensors laid out on a grid 220 of four computing devices A to D, where the tensors are sliced ​​along tensor dimension 0, and the slices are laid out on grid dimension 0. The tensors are also sliced ​​again along tensor dimension 0, but the slices are laid out on grid dimension 1. (As shown in...) Figure 2D As shown, after the tensor is laid out across the grid, the right side of image 210 is lost. Therefore, the layout <0, 0> is an illegal layout. The computation system 100 can be configured to receive the layout and the tensor, and to distribute the laid-out tensor to the grid 120 according to the layout.

[0051] The computing system 100 can be configured to receive an identifier for each tensor in the input data 150. <d0,d1,d2,...d n The corresponding n-tuple is taken as part of the layout parameter value 160. Then, the mesh execution engine 140 can lay out each tensor based on the received corresponding n-tuple.

[0052] In some implementations, the computing system 100 is configured to receive an identifier used to implement, for example, a parallelism technique, such as data parallelism, model parallelism, or both. In those implementations, consistent with the parallelism technique identified by the identifier, the layout engine 120 can generate a corresponding layout for each tensor in the input data 150. Therefore, a user specifying a parallelism technique can do so without having to specifically provide parameter values ​​for the layout corresponding to that parallelism technique.

[0053] In some implementations, computing system 100 is configured to receive a combination of layout parameter values ​​160 and one or more identifiers specifying the parallelism techniques to be implemented. For example, a user can provide computing system 100 with an identifier specifying that the system should process input data 150 by implementing data parallelism techniques. Alternatively, the user can provide layout parameter values ​​corresponding to tensors of how the input data 150 should be distributed across grid 110. In this way, the user has greater control over how the grid execution engine 140 allocates data and operations to each computing unit.

[0054] The computing system 100 can be configured to present to a user (e.g., on a display of a user device) the current layout of a tensor represented by an n-tuple of identifiers. The user can modify the values ​​of the n-tuples via a user interface displayed on the user device, and the computing system 100 can then receive and implement these values. The computing system 100 can be configured to first identify whether the layout modified by the user is valid or invalid, and notify the user if the user attempts to provide an invalid layout to the computing system.

[0055] Additionally, the grid execution engine 140 can be configured to execute communication primitive operations that enable computing devices in the grid 110 to communicate with each other. For example, after a computing device executes a corresponding device-local operation assigned by the grid execution engine 140, for example by distributing and storing tensors representing the weights of neural network operations, each computing device generates a corresponding output from the grid execution engine 140 by performing communication primitive operations between computing devices in the grid 110.

[0056] As another example, after the grid execution engine 140 assigns sub-tensors of the input tensor to different computing devices on the grid 110, the grid execution engine 140 can execute communication operation primitives to combine the output sub-tensors corresponding to the assigned input sub-tensors to form output data 170 representing the output of the input data 150.

[0057] The computing system may implement a layout API 190 (“Application Programming Interface”), which defines several functions associated with the layout of a specified tensor and used to allocate and perform operations on mesh 110. For example, the layout API 190 may include functions that can be invoked using appropriate values ​​from layout parameter values ​​160, and, when executed by the computing system 100, cause the computing system 100 to generate a layout based on provided arguments. For example, the arguments may specify whether the generated layout should implement data parallelism, model parallelism, or both.

[0058] As another example, the layout API 190 can implement functions that, when invoked with appropriate input parameter values, cause the computational system 100 to modify the shape of the mesh 110 according to the independent variables. As yet another example, the layout API 190 can implement functions for specifying the machine learning model to be executed on the mesh 110.

[0059] The layout API 190 can also implement several functions related to distributed programming and the distribution of data and operations to different computing devices in the grid 110. For example, the layout API 190 can implement functions specified in the MPI (“messaging interface”).

[0060] Each computing device in the grid 110 of the computing system 100 can implement several communication primitive operations for transferring the output tensor computed at the computing device to other computing devices in the grid 110. For example, the computing system 100 can implement slice-by-slice application of function F across multiple computing devices in the grid 110. Slice-by-slice application means that the operation corresponding to function F can be assigned to each of the multiple computing devices, and each computing device is configured to execute function F on the corresponding assigned tensor slice.

[0061] The computing system 100 can also be configured to process device-local operations assigned to computing devices in the grid 110, which produce output tensors with shapes different from the corresponding input tensors.

[0062] If the input tensor is laid out along the reduced dimension j, then system 100 can be configured to perform Allreduce. j The operation reduces each tensor slice along dimension j. In some implementations, this is similar to grouped Allreduce implemented by MPI. Allreduce refers to the operation of performing a reduction on all output values ​​generated by each computing device in grid 110 and broadcasting that reduction to each computing device. Computing system 100 can implement Allreduce by first dividing the computing devices in grid 110 into groups. jThe operation is performed such that each computing device in the group has coordinates that are different only in the grid dimension j (i.e., the reduced dimension). Then, the tensors assigned to each computing device and for each group can be summed or otherwise reduced, and the resulting reduced values ​​can be assigned to each computing device in the group.

[0063] If the input tensor is not laid out along the reduced dimension j, system 100 can perform piecewise reduction, where there is no communication on grid 110.

[0064] The computation system 100 can also be configured to broadcast tensor slices, which in this specification means inserting one or more new dimensions into a tensor and copying values ​​along those dimensions. In addition to implementing matrix multiplication, the computation system 100 can also be configured to implement other functions from pre-written libraries (such as TensorFlow).

[0065] The computing system 100 can also be configured to transform laid-out tensors from one layout to another. This is useful for quickly modifying the layout to implement different parallelization techniques. Additional slicing of tensors can be performed as a local operation by each computing device, without requiring network communication using primitives like Allreduce. On the other hand, the "un-slicing" dimension is implemented as an Allgarher operation, such as the Allgarher defined in MPI. Allgarher refers to the operation used to collect all output values ​​from each computing device in grid 110 and broadcast the collected output values ​​to each computing device.

[0066] As shown above (refer to the reference) Figure 1 As described, these and other functions can be implemented via Layout API 190. Specification data may include one or more call functions that define how computing devices in computing system 100 and grid 110 should transfer intermediate tensor outputs between each other.

[0067] Computational system 100 can enlarge a computational graph representing a set of interconnected operations and specify the computational layout of the graph. A computational graph is a graph representing a series of related operations. For example, a computational graph can represent operations used to process inputs via a neural network. Each operation can receive zero or more inputs and can generate an output. Inputs and outputs can be, for example, tensors.

[0068] Multiple operations (e.g., operations that the mesh execution engine 140 can distribute across the computing device) can be represented as a directed acyclic computation graph with multiple nodes and one or more edges. Each node in the computation graph represents a corresponding operation among the multiple operations. A node in the computation graph is designated as the final node f. The output of the operation represented by the final node f is the final output of the operation represented by the computation graph.

[0069] For any two nodes u and v in the computation graph, the edge (u, v) is a directed edge and represents a data dependency from u to v. A data dependency from u to v means that the operation represented by node u generates an input to the operation represented by node v, which then generates an output. Therefore, the operation at node u must be performed before the operation at node v. There are no cyclic data dependencies in the computation graph, such as edges (u, v) and (v, u), because the computation graph is acyclic.

[0070] A layout computation graph is a computation graph that represents operations that receive inputs and generate outputs as layout tensors. Computation system 100 can generate a layout-based computation graph from the graph by assigning layouts to each input and output tensor in the computation graph. In this specification, operations used for a layout-based computation graph are referred to as layout-based operations, and the layouts of the corresponding input and output tensors for layout-based operations are referred to as operation layouts. Commonly, the operation layouts corresponding to all layout-based operations used for a layout-based computation graph are referred to as computation layouts.

[0071] The computing system 100 can also receive a regular computation graph and generate a laid-out computation graph via the layout engine 120. The mesh execution engine 140 can then receive the laid-out computation graph and assign the laid-out tensors and laid-out operations to computing devices in the mesh 110. Depending on the layout parameter value 160, the layout engine 120 can generate a computation layout for a computation graph implementing one or more parallelization techniques.

[0072] For example, computing system 100 can receive a computational graph representing a neural network having the following: an input layer, a hidden layer with a non-linear activation function A, and an output layer. The neural network also includes two weight tensors W1 and W2, storing the values ​​of the weights representing the input and hidden layers of the neural network, respectively. Meanwhile, the operation for processing X to obtain Y from the neural network can be represented as:

[0073] γ←A(XW1)W2

[0074] Equation 1

[0075] In this example, it is assumed that the grid 110 of the computing system 100 is a one-dimensional grid of the computing device. The computing system 100 can generate a layouted computation graph that implements data parallelism on the example neural network. The operational layout of the input and output tensors for each operation can be represented as:

[0076]

[0077] The layout of the output tensor Y, the input tensor X, and the activations generated by executing the activation function A(XW1) is as follows: It also specifies slicing each tensor along tensor dimension 0 and laying out each tensor on the grid along grid dimension 0. The layout of weight tensors W1 and W2 is as follows: Furthermore, it specifies that the weight tensor is not sliced, but rather fully replicated on each computing device. Simultaneously, these operational layouts represent a general computational layout for implementing data parallelism. The computing system 100 can allocate appropriate tensor slices to each computing device according to the computational layout, and execute the layout-adjusted computational graph on the grid 120 because the weight tensor is fully replicated, but the data is split, distributed, and manipulated in parallel.

[0078] Equation 3 below represents the computational layout of the computational graph of the neural network shown in Equation 1, which implements model parallelism:

[0079]

[0080] In Equation 3, note: based on the empty layout To lay out the tensors X and Y corresponding to the input and output of the neural network, respectively, and simultaneously based on the layout... We then arrange the weight tensors W1 and W2 on a grid. Since the input and output data are copied, but the weights of the weight tensors are sliced ​​and distributed on the grid, Equation 3 represents the method of model parallelism implementation on the neural network.

[0081] Equation 4 below represents the computational layout of the computational graph of the neural network shown in Equation 1, which implements both data parallelism and model parallelism:

[0082]

[0083] Equation 4 shows the layout Each tensor is laid out on the grid. Since all tensors are sliced ​​and distributed on the grid, Equation 5 represents both data parallelism and model parallelism on the neural network.

[0084] The computation system 100 can receive operations 180 and input data 150, and generate a computation graph. In some implementations, the computation graph engine 130 is configured to receive input data 150 and operations 180, and generate a corresponding computation graph. Then, the layout engine 120 can receive the computation graph and generate a layout-processed computation graph from it, and send the layout-processed computation graph to the mesh execution engine 140.

[0085] The computing system 100 can be configured to display a computing layout on a user device's display and provide an interface for receiving layout parameter values ​​160 and operations 180 from, for example, a user of the computing system 100. For example, a layout API 190 can define one or more functions for receiving, generating, and modifying the computing layout.

[0086] In some implementations, computation system 100 can be configured to generate and output a computation layout corresponding to a computation graph without performing the computation graph on mesh 110. The system can then receive the computation layout as input and is configured to perform operations represented by the computation graph based on the computation layout.

[0087] Computational system 100 can generate a computational layout for a computational graph representing a Transformer sequence-to-sequence model. A sequence-to-sequence model is a machine learning model that receives an input sequence (e.g., an input sequence of words in a language) and generates an output sequence from the input sequence, such as translating the input sequence into another language. The Transformer model includes Encoder and Decoder layers, each containing several sublayers, including Multi-HeadAttention sublayers, Feed-Forward sublayers, and Multi-Head Attention sublayers. A detailed discussion of the Transformer model is provided in "Attention Is All You Need" by Vaswani et al., published in 2017.

[0088] The computational system 100 can generate a computational layout for a computational graph of a Transformer model, which, when executed by the mesh execution engine 140, enables the computational system 100 to execute the Transformer model, implementing model parallelism. Therefore, the computational system 100 can rapidly scale Transformer models to include billions of parameter values ​​and surpass the performance of conventional methods used for these and other types of models.

[0089] Figure 3This is a flowchart of an example process 300 for performing operations on multiple computing devices based on specification data of a distribution calculated from a specified tensor. For convenience, process 300 will be described as being executed by a system of one or more computers located at one or more locations and appropriately programmed according to this specification. For example, an appropriately programmed computing system (e.g., Figure 1 The computing system 100 can execute process 300.

[0090] The computing system receives standardized data on the distribution of the specified tensor computation across multiple computing devices (step 310). See above for reference. Figure 1 As described, the computational system can receive canonical data specifying the appropriate layout for each input and output tensor, which assigns each dimension of the input or output tensor to one or more computational devices among a plurality of computational devices. As mentioned above, each tensor computation can be represented as one or more operations in a computation graph. Each operation can receive one or more input tensors, each having one or more corresponding input dimensions, as input, and generate one or more output tensors, each having one or more corresponding output dimensions, as output. Some operations can both receive input tensors and generate one or more output tensors as output.

[0091] The system assigns appropriate device-local operations to each of the multiple computing devices based on the layout of the input and output tensors (step 320). See above for reference. Figure 1 The computational system described may include a mesh execution engine that distributes tensors according to their respective layouts. Operations can be distributed across multiple computational devices by assigning slices of tensors representing weights to one or more of a plurality of devices.

[0092] The system enables multiple computing devices to perform tensor computations (step 330) by having each of the multiple computing devices perform at least the corresponding device-local operation assigned to it. See above for reference. Figure 1 As described, consistent with how each tensor is laid out according to its corresponding layout, the mesh execution engine of the computational system enables computational devices within the mesh to perform tensor computations. Where necessary, the system can also execute one or more communication primitives to combine individual outputs generated by performing device-local operations on the computational device to produce a combined output, which can be either the output for tensor computation or the input for another computation.

[0093] The embodiments and functional operations of the subject matter described herein can be implemented in digital electronic circuit systems, tangibly embodied computer software or firmware, computer hardware (including the structures disclosed herein and their structural equivalents), or combinations thereof. Embodiments of the subject matter 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-transitory 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 substrate, a random or serial access memory device, or combinations thereof. Alternatively or additionally, the program instructions can be encoded on artificially generated propagated signals, such as machine-generated electrical, optical, or electromagnetic signals, generated to encode information for transmission to a suitable receiver device for execution by the data processing device.

[0094] This specification uses the term "configured to" in conjunction with system and computer program components. For a system of one or more computers to be configured to perform a specific operation or action, this means that software, firmware, hardware, or a combination thereof that causes the system to perform the operation or action is already installed on the system. For one or more computer programs to be configured to perform a specific operation or action, this means that one or more programs include instructions that, when executed by a data processing device, cause the device to perform an operation or action.

[0095] The term "data processing device" refers to data processing hardware and encompasses all kinds of devices, apparatuses, and machines used for processing data, including, for example, programmable processors, computers, or multiple processors or computers. The device may also be or further include special-purpose logic circuit systems, such as FPGAs (Field-Programmable Gate Arrays) or ASICs (Application-Specific Integrated Circuits). In addition to hardware, the device may optionally include code that creates an execution environment for computer programs, such as code constituting processor firmware, protocol stacks, database management systems, operating systems, or combinations thereof.

[0096] Computer programs (also referred to as 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 programming languages), and can be deployed in any form (including as stand-alone programs or modules, components, subroutines, or other units suitable for a computing environment). A program may, but is not required to, correspond to a file in a file system. A program may be stored as part of a file that holds other programs or data (e.g., one or more scripts stored in a markup language document), or in a single file dedicated to the program in question, or in multiple collaborating files (e.g., a file storing one or more modules, subroutines, or portions of code). A computer program may be deployed to execute on a single computer or on multiple computers located at a single site or distributed across multiple sites and interconnected via a data communication network.

[0097] Similarly, throughout this specification, the term "engine" is used extensively to refer to a software-based system, subsystem, or process programmed to perform one or more specific functions. Typically, an engine will be implemented as one or more software modules or components installed on one or more computers located at one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in others, multiple engines may be installed and run on the same one or more computers.

[0098] The processes and logic flows described in this specification can be executed by one or more programmable computers, which execute one or more computer programs to perform functions by manipulating input data and generating output. Processes and logic flows can also be executed by a dedicated logic circuit system (e.g., an FPGA or ASIC) or a combination of a dedicated logic circuit system and one or more programmable computers.

[0099] A computer suitable for executing computer programs can be based on a general-purpose or special-purpose microprocessor, or both, or any other type of central processing unit. Generally, the central processing unit receives instructions and data from read-only memory or random access memory, or both. The essential elements of a computer are the central processing unit for making or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and memory may be supplemented by or incorporated into a special-purpose logic circuit system. Generally, a computer will also include one or more mass storage devices (e.g., disks, magneto-optical disks, or optical disks) for storing data, or the computer may be operatively coupled to receive data from or transfer data to or both from such mass storage devices. However, a computer does not necessarily need to have such devices. Furthermore, a computer may be embedded in another device, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device (e.g., a universal serial bus (USB) flash drive), to name just a few.

[0100] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media, and memory devices, including: for example, 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.

[0101] To provide interaction with a user, embodiments of the subject matter described herein can 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, through which the user provides input to the computer. Other types of devices can also be used to provide interaction with the user; for example, feedback provided to the user can be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form (including sound input, voice input, or tactile input). Additionally, the computer can interact with the user by sending and receiving documents to and from the device used by the user (e.g., by sending a webpage to a web browser on the user's device in response to a request received from a web browser). Moreover, the computer can interact with the user by sending text messages or other forms of messages to a personal device (e.g., a smartphone running a messaging application) and receiving response messages from the user.

[0102] Data processing equipment for implementing machine learning models may also include, for example, dedicated hardware accelerator units for processing common and computationally intensive portions of machine learning training or production (i.e., inference).

[0103] Machine learning frameworks (such as TensorFlow, Microsoft Cognitive Toolkit, Apache Singa, or Apache MXNet) can be used to implement and deploy machine learning models.

[0104] Embodiments of the subject matter described in this specification can be implemented in computing systems, including backend components (e.g., as a data server), middleware components (e.g., an application server), or frontend components (e.g., a client computer having a graphical user interface, web browser, or app through which a user can interact with the implementation of the subject matter described in this invention), or any combination of one or more such backend components, middleware components, or frontend components. The components of the system can be interconnected via digital data communication (e.g., a communication network) of any form or medium. Examples of communication networks include local area networks (“LANs”) and wide area networks (“WANs”), such as the Internet.

[0105] A computing system may include clients and servers. Clients and servers are generally located far apart and typically interact via a communication network. The client-server relationship is created by computer programs running on respective computers that have a client-server relationship with each other. In some embodiments, the server transmits data (e.g., HTML pages) to a user device, for example, to display data to a user interacting with the device (which acts as the client) and to receive user input from that user. Data generated at the user device (e.g., the result of user interaction) may be received at the server from the device.

[0106] In addition to the implementations described in the appended claims and above, the following implementations are also innovative:

[0107] Implementation 1 is a computer-implemented method comprising: receiving canonical data specifying the distribution of tensor computations across a plurality of computing devices, wherein each tensor computation is (i) defined as receiving one or more corresponding input tensors, each having one or more corresponding input dimensions, as input, and (ii) defined as generating one or more corresponding output tensors, each having one or more corresponding output dimensions, as output, or both, and wherein the canonical data specifies a corresponding layout for each input and output tensor, the layout assigning each dimension of the input or output tensor to one or more computing devices among the plurality of computing devices; assigning corresponding device-local operations to each computing device among the plurality of computing devices based on the layout of the input and output tensors; and causing the plurality of computing devices to perform tensor computations by causing each computing device among the plurality of computing devices to perform at least the corresponding device-local operations assigned to the computing devices.

[0108] Implementation 2 is a method comprising: receiving canonical data specifying the distribution of tensor computations across a plurality of computing devices, wherein each tensor computation (i) receives one or more corresponding input tensors, each having one or more corresponding input dimensions, as input, and (ii) generates one or more corresponding output tensors, each having one or more corresponding output dimensions, as output, or both, and wherein the canonical data specifies a corresponding layout for each input and output tensor, the layout assigning each dimension of the input or output tensor to one or more computing devices among the plurality of computing devices; assigning corresponding device-local operations to each computing device among the plurality of computing devices based on the layout of the input and output tensors; and causing the plurality of computing devices to perform tensor computations by causing each computing device among the plurality of computing devices to perform at least the corresponding device-local operations assigned to the computing devices.

[0109] Implementation 3 is based on any one of implementations 1 to 2, wherein tensor computation is an operation used to train a machine learning model on training data using machine learning training techniques, and wherein canonical data specifies a corresponding layout for each tensor processed during iteration of machine learning training techniques on a batch of training data.

[0110] Implementation 4 is based on any one of implementations 1 to 3, wherein tensor computation is represented by a directed graph of nodes connected by directed edges, each node represents the corresponding tensor computation, and each incoming edge of a node represents the input tensor of the tensor computation represented by the node, and each outgoing edge of a node represents the output tensor of the tensor represented by the node, and wherein canonical data specifies the corresponding layout for each tensor among the tensors represented by the edges in the directed graph.

[0111] Implementation 5 is based on any one of implementations 1 to 4, wherein, for each input and output tensor, the layout specifies for each dimension of the tensor (i) copying the data along the dimension on all computing devices in a plurality of computing devices or (ii) splitting the data along the dimension among at least two processors in a plurality of processors.

[0112] Implementation 6 is based on any one of implementations 1 to 5, wherein a plurality of computing devices are arranged in an n-dimensional array, and wherein, for each input and output tensor, the layout specifies for each dimension of the tensor (i) copying data along the dimension on all computing devices in the plurality of devices or (ii) splitting data along the dimension among the computing devices along a dimension of the n-dimensional array.

[0113] Implementation 7 is based on any one of implementations 1 to 6, wherein the shape of each of the input and output tensors is called a tuple of shape symbols, and wherein the canonical data identifies a mapping from the shape symbol to an identifier for one of the n dimensions for at least one shape symbol.

[0114] Implementation 8 is based on any one of implementations 1 to 7, wherein, for each tensor computation, corresponding device-local operations are assigned to multiple computing devices based on the allocation of n dimensions of the input and output tensors: for each input tensor of the tensor computation and each device, a slice of the input tensor assigned to the device by canonical data is identified; and a slice-by-slice instance of the tensor computation is assigned to each device, which applies the tensor computation to the slice of the input tensor assigned to the device.

[0115] Implementation 9 is a method according to any one of implementations 1 to 8, wherein allocating corresponding device-local operations to multiple computing devices based on the allocation of dimensions of input and output tensors includes: for the first tensor computation, determining that a piecewise instance of the first tensor computation requires inter-device communication to generate the correct output for the first tensor computation, and allocating one or more communication operations to the multiple devices, the one or more communication operations causing inter-device communication to combine the output generated by the piecewise instance of the first tensor computation.

[0116] Implementation 10 is a system comprising: one or more computers and one or more storage devices storing instructions operable, when executed by the one or more computers, to cause the one or more computers to perform a method according to any one of implementations 1 to 9.

[0117] Example 11 is a computer storage medium encoded with a computer program including instructions that, when executed by a data processing device, are operable to cause the data processing device to perform a method according to any one of implementations 1 to 9.

[0118] While this specification includes numerous specific implementation details, these details should not be construed as limiting the scope of any invention or what may be claimed, but rather as descriptions of features that may be implemented for particular embodiments of a particular invention. Certain features described in this specification within the context of individual 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 functioning in certain combinations and initially or even equally claimed, in some cases, one or more features from the claimed combination may be removed from the combination. And the claimed combination may refer to a sub-combination or a variation of a sub-combination.

[0119] Similarly, although operations are shown in a specific order in the accompanying drawings, this should not be construed as requiring such operations to be performed in the described specific order or in a sequential order, or as requiring all illustrated operations to achieve the desired result. In some cases, multitasking and parallel processing can be advantageous. Furthermore, the separation of various system modules and components in the above embodiments should not be construed as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0120] Specific embodiments of this subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions set forth in the claims can be performed in a different order and still achieve the desired result. As an example, the processes depicted in the drawings do not necessarily require the specific order or sequential sequence shown to achieve the desired result. In some cases, multitasking and parallel processing can be advantageous.

Claims

1. A computer-implemented method, comprising: For one or more operations from a plurality of operations, and using i) canonical data representing a computation graph of the plurality of operations and ii) specifying the distribution of one or more dimensions of a tensor to multiple components of the system, the operations are assigned to components, which will perform the operations on portions of the tensor corresponding to the one or more dimensions specified by the canonical data. The computation graph includes multiple nodes and multiple edges, wherein each node represents a corresponding operation from the multiple operations, and each edge connects a corresponding first node to a corresponding second node, representing an operation that receives the output of the operation represented by the corresponding first node as input. The canonical data defines the mapping of one or more dimensions of the tensor to corresponding components from the plurality of components of the system; as well as The tensor is used to cause one or more components from the plurality of components to perform the assigned operation on the corresponding data portion of the tensor. The specification data is generated based on the received identifier that identifies at least one parallelism technique used in the implementation of the system.

2. The method of claim 1, wherein, For each dimension from the one or more dimensions, the canonical data specifies that (i) the tensor for that dimension is copied across multiple components of the plurality of components or (ii) the tensor for that dimension is split across at least two components from the plurality of components.

3. The method according to claim 1, wherein: The first dimension from the one or more dimensions is called a tuple of shape symbols; and The canonical data refers to the mapping of tuple identifiers for the shape symbols used in the first dimension from tuples of the shape symbols to identifiers for the first dimension.

4. The method of claim 1, wherein, Assigning the operation to the component for one or more operations from the plurality of operations includes: For a first operation from the plurality of operations, determine that a piecewise instance of the first operation requires communication between components to generate the correct output for the first operation; and One or more communication operations are assigned to the plurality of components, the one or more communication operations causing communication between the components to combine the output generated by the piecewise instance of the first operation.

5. The method according to claim 4, wherein, The allocation that leads to communication between components to combine the outputs generated by the slice instances of the first operation includes one or more communication operations such as inserting an Allreduce operation.

6. The method according to claim 1, wherein, The specified data is specified by the user, and the method further includes: It is determined that the operations from the plurality of operations assigned to the component according to the first specification data will result in the loss of data for the tensor corresponding to the first specification data; In response to determining that an operation from the plurality of operations, allocated to a component according to the first specification data, will result in data loss for a tensor corresponding to the first specification data, a notification indicating that the first specification data will result in data loss is provided to the device operated by the user; and After providing a notification indicating that the first specification data will result in data loss, specification data that will not result in data loss is received from the user-operated device.

7. The method according to claim 1, wherein: Each of the plurality of components includes devices from a plurality of devices of the system; as well as Assigning operations to the components includes assigning the operations to devices for one or more operations from a plurality of operations and using i) a computation graph representing the plurality of operations and ii) canonical data specifying the distribution of the one or more dimensions of the tensor to the plurality of devices of the system, wherein the devices will perform the operations on portions of the corresponding dimensions of the one or more dimensions of the tensor specified by the canonical data.

8. The method according to claim 1, wherein, Causing one or more components from the plurality of components to perform the assigned operation on the corresponding data portion of the tensor includes using the tensor to cause components from the plurality of components to perform the assigned operation on the entire tensor.

9. The method according to claim 1, wherein, Causing one or more components from the plurality of components to perform the assigned operation on the corresponding data portion of the tensor includes using the tensor to cause one or more components from the plurality of components to perform the assigned operation on the corresponding portion of the tensor.

10. A system comprising one or more computers and one or more storage devices storing instructions, the instructions being operable, when executed by the one or more computers, to cause the one or more computers to perform operations, the operations including: For one or more operations from a plurality of operations, and using i) canonical data representing a computation graph of the plurality of operations and ii) specifying the distribution of one or more dimensions of a tensor to multiple components of the system, the operations are assigned to components, which will perform the operations on portions of the tensor corresponding to the one or more dimensions specified by the canonical data. The computation graph includes multiple nodes and multiple edges, wherein each node represents a corresponding operation from the multiple operations, and each edge connects a corresponding first node to a corresponding second node, representing an operation that receives the output of the operation represented by the corresponding first node as input. The canonical data defines the mapping of one or more dimensions of the tensor to corresponding components from the plurality of components of the system; as well as The tensor is used to cause one or more components from the plurality of components to perform the assigned operation on the corresponding data portion of the tensor. The specification data is generated based on the received identifier that identifies at least one parallelism technique used in the implementation of the system.

11. The system according to claim 10, wherein, For each dimension from the one or more dimensions, the canonical data specifies that (i) the tensor for that dimension is copied across multiple components of the plurality of components or (ii) the tensor for that dimension is split across at least two components from the plurality of components.

12. The system according to claim 10, wherein: The first dimension from the one or more dimensions is called a tuple of shape symbols; and The canonical data refers to the mapping of tuple identifiers for the shape symbols used in the first dimension from tuples of the shape symbols to identifiers for the first dimension.

13. The system according to claim 10, wherein, Assigning the operation to the component for one or more operations from the plurality of operations includes: For a first operation from the plurality of operations, determine that a piecewise instance of the first operation requires communication between components to generate the correct output for the first operation; and One or more communication operations are assigned to the plurality of components, the one or more communication operations causing communication between the components to combine the output generated by the piecewise instance of the first operation.

14. The system according to claim 13, wherein, The allocation that leads to communication between components to combine the outputs generated by the slice instances of the first operation includes one or more communication operations such as inserting an Allreduce operation.

15. The system according to claim 10, wherein, The specified data is user-defined, and the operation further includes: It is determined that the operations from the plurality of operations assigned to the component according to the first specification data will result in the loss of data for the tensor corresponding to the first specification data; In response to determining that an operation from the plurality of operations, allocated to a component according to the first specification data, will result in data loss for a tensor corresponding to the first specification data, a notification indicating that the first specification data will result in data loss is provided to the device operated by the user; and After providing a notification indicating that the first specification data will result in data loss, specification data that will not result in data loss is received from the user-operated device.

16. The system according to claim 10, wherein: Each of the plurality of components includes devices from a plurality of devices of the system; as well as Assigning operations to the components includes assigning the operations to devices for one or more operations from a plurality of operations and using i) a computation graph representing the plurality of operations and ii) canonical data specifying the distribution of the one or more dimensions of the tensor to the plurality of devices of the system, wherein the devices will perform the operations on portions of the corresponding dimensions of the one or more dimensions of the tensor specified by the canonical data.

17. The system according to claim 10, wherein, Causing one or more components from the plurality of components to perform the assigned operation on the corresponding data portion of the tensor includes using the tensor to cause components from the plurality of components to perform the assigned operation on the entire tensor.

18. The system according to claim 10, wherein, Causing one or more components from the plurality of components to perform the assigned operation on the corresponding data portion of the tensor includes using the tensor to cause one or more components from the plurality of components to perform the assigned operation on the corresponding portion of the tensor.

19. A plurality of non-transitory computer-readable storage media encoded with instructions, said instructions causing said one or more computers to perform operations when executed, said operations including: For one or more operations from a plurality of operations, and using i) canonical data representing a computation graph of the plurality of operations and ii) specifying the distribution of one or more dimensions of a tensor to multiple components of the system, the operations are assigned to components, which will perform the operations on portions of the tensor corresponding to the one or more dimensions specified by the canonical data. The computation graph includes multiple nodes and multiple edges, wherein each node represents a corresponding operation from the multiple operations, and each edge connects a corresponding first node to a corresponding second node, representing an operation that receives the output of the operation represented by the corresponding first node as input. The canonical data defines the mapping of one or more dimensions of the tensor to corresponding components from the plurality of components of the system; as well as The tensor is used to cause one or more components from the plurality of components to perform the assigned operation on the corresponding data portion of the tensor. The specification data is generated based on the received identifier that identifies at least one parallelism technique used in the implementation of the system.

Images