Utilizing multidimensional hardware accelerators to process reduction and broadcast operations on large datasets.

Abstract

Methods, systems, and apparatuses include instructions encoded on a storage medium for performing, for example, reduction of gradient vectors and similarly structured data generated in parallel on nodes organized in a mesh or toroidal topology defined by connections in at least two dimensions between nodes. The methods provide parallel computation and communication between nodes in said topology.

Description

[0001] Cross-references to related applications

[0002] This application claims priority to U.S. Provisional Patent Application No. 62 / 897,239, filed September 6, 2019. The disclosure of the earlier application is considered part of the disclosure of this application and is incorporated herein by reference. Background Technology

[0003] This manual pertains to training machine learning models.

[0004] Machine learning models receive input and generate output based on the received input and the values ​​of the model parameters. These models typically require training with the received input, which can be a very large dataset. Summary of the Invention

[0005] This specification generally describes techniques related to parallel processing of large datasets, specifically involving parallel processing of large datasets across various node topologies to enable faster reduction, mapping, and broadcasting operations on multiple topological nodes. Nodes are networked together and can be executed in hardware via hardware accelerators. Hardware accelerators may include multiple networked processors running in parallel (e.g., in mesh arrays, toroidal configurations, etc.). More specifically, the cross-replication and sum (CRS) operation is configured to sum the gradient contributions of multiple replicas in a data-parallel stochastic gradient descent (SGD) process for such networks. For scalable execution of SGD training on multiple replicas, gradient contributions are merged (e.g., summed) at a given batch size. However, this process can often be a bottleneck for the execution of machine learning models (e.g., neural network models) because the reduction-dispersion and aggregation phases of the CRS operation are performed in a series of one-dimensional (1D) loops within the network. For example, circle algorithms, such as the rotated pincer algorithm, perform global summation operations on a 1D configuration of topological nodes. This configuration is suboptimal because it does not utilize all router links between the topology nodes. More precisely, as the 1D algorithm executes in the second dimension, it executes serially in the first dimension, and during this period, some router links are in a dormant state.

[0006] To overcome this problem, this disclosure describes a process for performing CRS operations in the multidimensional topology using a multidimensional circle algorithm. The multidimensional algorithm is configured to utilize pipelined direct memory access (DMA) transfers for the summation operation to increase the throughput of the CRS operation. This is achieved by simultaneously executing the circle algorithm in the multidimensional topology so that the gradient contribution of the sum can be replicated simultaneously in the multidimensional topology. Because the router links between nodes in the topology are more fully utilized compared to executing the 1D algorithm serially in each dimension of the multidimensional topology, the latency of the CRS operation can be significantly reduced, and in some cases, by more than 200%. Reducing the latency of the CRS operation eliminates a major bottleneck in training machine learning models.

[0007] In one aspect, a method for processing training data is described. A corresponding copy of a machine learning model can be trained on each of multiple nodes organized in a multidimensional topology comprising rows and columns of nodes. Each node can be trained in parallel on a corresponding batch of training data. After training, each node can store the corresponding gradient vector obtained during training. The corresponding gradient vectors in the nodes can be combined by performing operations to generate a final gradient vector. These operations include: performing the first stage of the circle algorithm in parallel for each row of the multidimensional topology by code executed on the nodes to generate different combined data for the portion of the gradient vector at each node of the row; performing the second part of the first stage of the circle algorithm in parallel for each column of the multidimensional topology by code executed on the nodes to generate different column results at each node of the column, which include a portion of the combined data from each row; performing the first part of the second stage of the circle algorithm in parallel for each column of the multidimensional topology by code executed on the nodes to generate a portion of the final gradient vector in each column, which includes each column result; and performing the second part of the second stage of the circle algorithm in parallel for each row of the multidimensional topology by code executed on the nodes to generate a final gradient vector in each row, which includes each portion of the final gradient vector from each column.

[0008] In some variations, one or more of the following may be implemented individually or additionally in any feasible combination. The method may further include: determining the type of the multidimensional topology of the multiple nodes; mapping a first part of the first phase to a first set of nodes and a first set of links between the first set of nodes based on the type; and mapping a second part of the first phase to a second set of nodes and a second set of links between the nodes, the first set of links being different from the second set of links.

[0009] The method may further include: determining a multidimensional topology of multiple nodes, including surrounding links; and configuring each row and column of the multidimensional topology using the surrounding links to perform a circle algorithm.

[0010] Each node in a multidimensional topology can include at least two child nodes. Multidimensional topologies can include toroidal topologies, where, for each row and each column, the child nodes of each node are configured as adjacent nodes in a circle algorithm. Multidimensional topologies can include grid topologies, where, for each row and each column, the child nodes of the grid's edge nodes are configured as adjacent nodes in a circle algorithm.

[0011] In addition to rows and columns, the multidimensional topology can include a set of nodes in the third dimension, wherein the method further includes: executing a third part of the first stage of the circle algorithm in parallel by means of code executed on the node set in the third dimension to generate different third-dimensional results in each node of the column, the results including a portion of the column results from each column; and executing a third part of the first stage of the circle algorithm in parallel by means of code executed on the node set in the third dimension to generate a final gradient vector from each column in each node of the node set in the third dimension.

[0012] The method may further include, during the execution of the first and second phases of the circle algorithm, pipelined data transfer between each node by leveraging the execution of code on each node. The circle algorithm may include a rotation clamp algorithm, wherein the first phase includes a reduction-dispersion phase of the rotation clamp algorithm, and the second phase includes a global aggregation phase of the rotation clamp algorithm. The combined data portion of the gradient vector comprises the sum data from two or more nodes in the row. The portion of the combined data from each row comprises the sum data, which includes the combined data from two or more nodes in the column.

[0013] Details of one or more embodiments of the process and system are described in the accompanying drawings, which are given by way of illustration only and are illustrated in the description below. Other features, aspects, and advantages of the invention will become apparent from the specification, drawings, and claims. The same reference numerals and designations indicate the same elements in the various figures. Attached Figure Description

[0014] Figure 1 An example system of replicas of a machine learning model is illustrated.

[0015] Figure 2 An example of processing units connected in a 1D toroidal topology on the example module is illustrated, showing the execution of the reduction dispersion phase of the circle algorithm.

[0016] Figure 3 An example of processing units connected in a circular topology is illustrated, showing the execution of the reduction and dispersion phase of the circular algorithm.

[0017] Figure 4This is a schematic diagram of an example multidimensional topology including a two-dimensional (2D) torus, illustrating the execution of the reduction dispersion phase of the circle algorithm.

[0018] Figure 5 The illustration shows that Figure 3 An example of a processing unit connected in a circular topology illustrates the execution of the global aggregation phase of the circular algorithm.

[0019] Figure 6 yes Figure 4 A schematic diagram of the multidimensional topology illustrates the execution of the global aggregation phase of the circle algorithm.

[0020] Figure 7 This is a schematic diagram of a multidimensional topology including a two-dimensional (2D) mesh topology, illustrating the execution of the reduction dispersion phase of the circle algorithm.

[0021] Figure 8 This is a schematic diagram of a multidimensional topology including a 2D mesh topology, illustrating the execution of the global aggregation phase of the circle algorithm.

[0022] Figure 9 It is a schematic diagram of a multidimensional topology including a 2D torus topology, which includes multiple processors for each node.

[0023] Figure 10 It is a schematic diagram of a multidimensional topology that includes a 2D mesh topology, which includes multiple processors for each node.

[0024] Figure 11 This is a flowchart illustrating the operations used to perform CRS operations in a multidimensional space.

[0025] Figure 12 This is a flowchart illustrating an example process for combining the corresponding gradient vectors of each node in a multidimensional topology.

[0026] The same reference numerals and markings indicate the same elements in each figure. Detailed Implementation

[0027] A common strategy for training machine learning models is to process batches of training data in parallel across multiple processing units. Figure 1 The diagram illustrates example systems with replicas of machine learning models A through D, each trained on its corresponding processing unit. For convenience, the replicas and the processing units on which they were discovered are referred to using the same labels in the diagram.

[0028] When trained, a machine learning model is defined by the values ​​of its model parameters. Model parameters are typically organized as non-scalar data, such as vectors, two-dimensional (2D) matrices, three-dimensional (3D) matrices, or higher-order matrices, whose elements are typically scalar values, such as integers or floating-point numbers.

[0029] In the system, in each iteration of the model training process, each copy of the model (also called a node) is trained on a unique batch of training data (e.g., from the training dataset). Figure 1 In the diagram, replicas A through D were trained on batches 1 through 4, respectively. When a node has finished processing its batch of training data, it has a set of gradients for the values ​​of the model parameters. Figure 1 In the example, when nodes A through D have completed batches 1 through 4 respectively, nodes A, B, C, and D have gradient values ​​[a1, a2], [b1, b2], [c1, c2], and [d1, d2] respectively. The gradient values ​​in each node have the same structure and generally correspond to the structure of the parameter values. For convenience, these are referred to as vectors.

[0030] In this disclosure, a node generally refers to a part of a system that performs CRS operations and other logical operations. Although a node does not necessarily correspond to a hardware processor, a description specifying that a node is performing a logical operation (e.g., summation) implies that the hardware executing the node is performing a processing operation to perform the logical operation.

[0031] Because the replicas are trained on different data, the gradient vectors of the nodes are combined to generate a final gradient vector, which is used to update the model's parameter values ​​(e.g., the parameter vector). One way to combine the gradient vectors is to generate an element-wise average. The updated parameter values ​​are passed to all machine learning model nodes, and generally another iteration is expected to process a batch of training data and combine (e.g., reduce) the gradient vectors of each replica to a final gradient vector (e.g., the reduced gradient vector), and update the model's parameter values. The portion of the gradient vector passed to each node of the model can be referred to as a fragment of the gradient vector.

[0032] Some machine learning systems have a central parameter computation subsystem that receives gradient vectors from each copy and combines them at a central location. This specification describes techniques that can be used in machine learning systems that do not have a central parameter computation subsystem, but instead reduce gradient vectors and update parameter values ​​in a distributed manner across nodes. The operation that accomplishes this is called the Cross-Replication and Sum (CRS) operation, which typically consists of two phases: a global summation operation (also known as the reduction and dispersion phase) and a global aggregation operation (also known as the aggregation phase). The CRS operation replicates (e.g., aggregates) the final gradient vector across all nodes of the network. In some implementations, element-wise averaging occurs during the reduction and dispersion phase because the contributions from each node are combined.

[0033] Figure 2The diagram illustrates a high-speed connection topology, which connects an example combination of processing units AH (202a to 202h) in a circular topology 200. In some implementations, all processing units reside on a single module. Figure 2 The lines between the processing units shown (such as line 204) represent high-speed data communication links. Processing units are typically fabricated on multiple integrated circuits (“chips”). In some implementations, each chip may include multiple processing units (e.g., two processing units). Links across chip boundaries are called inter-chip network links, while processing units on the same chip communicate via intra-chip interface links. Links may include half-duplex links or full-duplex links, where only one processing unit can transmit data at a time, and where data can be transmitted simultaneously in both directions on a full-duplex link. A full-duplex link may be represented as two arrows between nodes if the link is configured to transmit data for a specific operation in both directions. However, for any diagram referenced in this specification, a link represented by a single arrow does not necessarily refer to a half-duplex link, but rather indicates the direction of data flow for a specific operation shown in the diagram.

[0034] The reference node (e.g., processing unit) AH 202a-h now describes a method for performing reduction on this topology. The reduction has a summation phase and an aggregation phase, which may include a broadcast phase. These are known as circle algorithms.

[0035] A single-path algorithm has one node, for example, node A, that sends data to a neighbor, for example, node B. B is configured to combine (e.g., add) the data from A into its own data and send the sum to C, C repeats this operation and sends it to D, and so on. Data processing, sending, and receiving can be performed in a streaming manner. For example, B is configured to start adding before B receives all the data from A. In this example, the final combination occurs in node H. The final gradient vector is then passed from H to the other nodes AG. This can be achieved by reversing the streaming data path from H back to A. If broadcasting is available to one or more nodes, allowing a node to broadcast to multiple nodes, the passing of the final gradient vector can be accomplished using broadcasting to complete the global aggregation phase of the CRS operation. The single-path circle algorithm consists of N steps in each of the reduction and aggregation phases, where N corresponds to the number of nodes.

[0036] To achieve lower latency than the single-path algorithm, the clamp algorithm sends data across two paths in opposite directions on nodes 202a-h of topology 200. Figure 3 It shows Figure 2 Examples of circular topologies configured to perform clamping algorithms, as shown by arrows (such as arrows 302, 304, etc.). Figure 3In the example, nodes A and B send data to nodes H and C at approximately the same time (step 1). Nodes H and C then send the combined data to nodes G and D (step 2). Nodes G and D then send the combined data to nodes F and E (step 3). Nodes F and E contain the sums across paths AHGF and BCDE, respectively. If the number of nodes is even, the data must be transmitted from F to E (step 4), and vice versa, and there it is combined before being broadcast to all nodes.

[0037] The rotation clamp algorithm can further reduce the latency of the reduction and aggregation phases. For the rotation clamp algorithm, independent instantiation of the clamp algorithm is performed simultaneously, starting with each pair of adjacent nodes. For example, when the clamp algorithm starts running from nodes A and B, the same algorithm runs from nodes B and C, from nodes C and D, and so on. The gradient vector is partitioned into disjoint subsets of equal size (if possible) or essentially equal size, and each instantiation of the algorithm processes one of these subsets. This can potentially change how the stream works, because now, for example, B sends one subset of data to C, and C sends another subset of data to D. In this example, C will not forward data from B until it has finished sending its own data. For this purpose, the data transmitted from B to C is temporarily stored on C. Alternatively, the processing of data arriving at C from B and D can be interleaved.

[0038] If the number of nodes is even, the improved rotation clamping algorithm described below further reduces latency. For the improved rotation clamping algorithm, node A first sends half of its data to node H (as shown by arrow 304) and the other half to node B (then node C, as shown by arrow 302). The clamping algorithm then continues as previously described. Therefore, the data path is symmetrical in each of the two directions from A. If a node is configured to broadcast to every other node, the node sends data in the opposite direction. For example, node E sends data in the opposite direction, and as a final step, nodes B and H send half of their disjoint data to A so that A still receives all the data fragments of the gradient vector. Figure 4 In the image, the arrow is configured to indicate the reduction and dispersion phase of the CRS operation.

[0039] The process can be configured using control messages. For example, control messages can configure which subset of the gradient vector will be processed by the process for each reduction instantiation, whether it is the starting unit, the final unit, or an intermediate unit in the reduction, and in which direction data or combined data and broadcasts should be sent, as described above. Control messages are configured based on the node-specific topology. As described later, control messages can be configured to utilize additional links set up in multidimensional topologies such as 2D meshes and 2D torus topologies.

[0040] The process performed by each intermediate processing unit is as follows. Each intermediate processing unit (e.g., a node) combines its gradient vector, or a portion thereof, with the input gradient vector it received when it received a portion of the input gradient vector from the previous processing unit (e.g., an upstream processing unit). The combination operation performed by each processing unit may be a simple summation of the combined gradient vectors or some other computation. The intermediate processing unit then transmits the combined gradient vector to the next processing unit in the direction (e.g., a downstream processing unit).

[0041] The final processing unit (e.g., a node) combines its gradient vector with the input gradient vector it received when it received the input gradient vector from the preceding processing unit (e.g., an upstream processing unit) in the direction from the data path, and generates the final reduced gradient vector. The final processing unit typically combines the reduced gradient vector with the values ​​of the machine learning model parameters to produce an updated set of parameter values. Assuming a vector of parameters x and a gradient dx, a simple update has the form: x += -learning_rate * dx, where learning_rate is a scalar term. Update rules can be arbitrarily complex; for example, they can depend on previous gradients. After computing the updated parameters, the final processing unit initiates a broadcast operation that reverses the data flow, providing the updated parameter values ​​to all other processing units to send the final output back to the root processing unit (e.g., a node) in a direction from the data path. As a result, each node will have updated parameters to process the next batch of machine learning training data.

[0042] Figure 4 This is a schematic diagram of an example multidimensional topology including a two-dimensional (2D) torus 400, illustrating the execution of the reduction scatter phase of the circle algorithm. The torus topology 400 includes (in this example) 9 replicas and processing nodes A1 to C3 (each of which is a part of a circle consisting of three nodes), horizontal nodes (labeled X), and vertical nodes (labeled Y), and each of them has direct links to 4 other nodes. In some examples, the horizontal circles of three nodes in the 2D topology 400 and the vertical circles of three nodes in the 2D topology 400 correspond to the rows and columns of the 2D topology 400, respectively. While a 3×3 topology is shown for illustration purposes, the size of the topology can be increased or decreased from topology 400. For example, for a topology used to process large machine learning models (e.g., models with 10 million, 25 million, or 50 million parameter values), the appropriate number of processing nodes would be larger, for example, 16 nodes per circle, for a total of 256 nodes. Larger or smaller numbers will be suitable for larger or smaller models.

[0043] The iterative reduction algorithm on the torus simplifies reduction on the torus to a series of circular reductions. The first step is to perform circular reduction along each row of the 2D topology 400. First, perform any of the circular reduction algorithms described above in each row of the 2D topology 400, resulting in each node in the row having a sum vector of gradient vectors for each node in the row. Next, perform circular reduction along each column of the 2D topology 400, at which point each node has the same final gradient vector.

[0044] As mentioned earlier, reduction involves a summation step, a computation step, and a broadcast step. Initial row circle reduction should not perform the complete computation step because even if the computation on the circle has been completed, the gradient has not actually been summed on the torus.

[0045] To reduce the latency introduced by the iterative reduction algorithm described above, topology 400 can be configured to perform the rotational clamping algorithm in multiple dimensions. For example, performing the rotational clamping algorithm in two phases of a 2D topology 400... Figure 3 The described reduction and dispersion operation. In the first phase (in... Figure 4 In stage 0), a reduction dispersion operation is performed along the first dimension (e.g., the horizontal dimension or the X direction). For example, for each circle, such as circle 200a including A1, B1, and C1, circle 200b including A2, B2, and C2, and circle 200c including A3, B3, and C3, a reduction dispersion operation is performed along the first dimension (e.g., the horizontal dimension or the X direction). Figure 3 The reduction dispersion described herein. A first reduction dispersion operation is followed by a second reduction dispersion operation in the second dimension (e.g., the vertical dimension or the Y dimension). For example, each circle, such as circle 200d including A1, A2, and A3, circle 200e including B1, B2, and B3, and circle 200f including C1, C2, and C3, is executed in parallel. Figure 3 The reduction and scatter operation is shown in the diagram. The reduction and scatter operation is performed in two parts because topology 400 has two dimensions. An additional part of the reduction and scatter operation is performed on the additional dimensions of topology 400.

[0046] Turning Figure 5The diagram illustrates an example of processing units connected in a circular topology 200, demonstrating the execution of the global aggregation phase of the circular algorithm. Node E, including the sum of each node in nodes AD and FH, sends its data to nodes D and F (step 1). Nodes F and D store this data and also send the sum data to nodes G and C, respectively (step 2). Subsequently, nodes G and C repeat the process, respectively, to store the total gradient vector data and send the data to nodes H and B, respectively (step 3). H and B then send the data to node A, respectively (step 4). Alternatively, in step 1, node E can send the sum data to node F. Then, in step 2, nodes F and E can simultaneously send the data to nodes G and D, and so on, to propagate the data through nodes H and C, respectively, and terminate at nodes A and B.

[0047] Turning Figure 6 , showed Figure 4 A schematic diagram of the multidimensional topology 400 illustrates the execution of the global aggregation phase of the aforementioned circle algorithm. For each dimension of the multidimensional topology... Figure 4 The described reduction and dispersion operation is followed by a global aggregation phase for each dimension of the network. Figure 6 Phase 1 is shown. The aggregation phase of topology 400 consists of two parts, one for each dimension. The two aggregation parts are arranged relative to each other. Figure 4 The reduction and dispersion operations described are performed in the reverse order of the dimensions described. For example, in loop (or circle) 200d including A1, A2, and A3, loop 200e including B1, B2, and B3, and loop 200f including C1, C2, and C3, the first aggregation portion occurs along the Y dimension. The second portion of the aggregation phase is performed along the horizontal (X) dimension, including: loop 200a including A1, B1, and C1, loop 200b including A2, B2, and C2, and loop 200c including A3, B3, and C3.

[0048] The data payload in the vertical (Y) dimension is the sum of the payloads scaled down by the size of the horizontal (X) dimension. In other words, the data fragments in the X dimension are the entire reduced span in the Y dimension. In the 2D rotation clamp algorithm, the time cost of short cross-copying and calling is O(2*(size(X)+size(Y))). On large 2D meshes or tori, the time cost in the 1D algorithm is less than O(size(X)*size(Y)).

[0049] The 2D circle algorithm is configured to utilize more router links in a 2D topology. Return Figure 4 and Figure 6The 2D circle algorithm can be executed in parallel along the X and Y dimensions. Each of these two parallel steps (each step involves executing the 2D circle along one of the multiple dimensions) is called a color. Multicolor full reduction occurs during the reduction dispersion phase ( Figure 4 During phase 0 (as shown), execution is performed in parallel in stages along different dimensions of the multidimensional topology. For the global aggregation phase (… Figure 6 As shown in stage 1), the order of dimensions is changed so that each node of the topology ends with a final gradient vector. For 2D topologies, such as... Figure 4 and Figure 6 As shown, the order of the X and Y dimensions is flipped from stage 0 to stage 1. For example, a two-color full reduction can be configured to perform along the X dimension following the Y dimension of color 1, while it performs along the Y dimension following the X dimension of color 2. This two-color scheme can be extended to n colors on an n-dimensional grid or torus. Higher throughput can be achieved with multiple colors when more links are utilized simultaneously during each stage of the reduction dispersion and aggregation phases of the CRS operation.

[0050] The algorithms described in this specification can be executed on other topologies. (Turn) Figures 7 to 8 The diagram illustrates a 2D mesh 700 topology. Mesh 700 is configured in a manner similar to toroidal topology 400 for performing the reduction dispersion phase. Figure 7 Phase 0) and aggregation phase ( Figure 8 Each stage in stage 1). For grid 700, full-duplex links can be used to create loops 200a through 200f, instead of the looping links in topology 400. Similar to topology 400, multicolor full reduction occurs in the reduction dispersion stage ( Figure 7 During phase 0 (as shown), execution is performed in parallel in stages along different dimensions of the multidimensional topology. For the global aggregation phase (… Figure 8 As shown in stage 1), the order of dimensions is changed so that each node of the topology ends with a final gradient vector. For 2D topologies, such as... Figures 7 to 8 As shown, the order of the X and Y dimensions is flipped from stage 0 to stage 1. For example, a two-color full reduction can be configured to perform along the X dimension following the Y dimension of color 1, while it performs along the Y dimension following the X dimension of color 2. This two-color scheme can be extended to n colors on an n-dimensional grid or torus. Higher throughput can be achieved with multiple colors when more links are utilized simultaneously during each stage of the reduction dispersion and aggregation phases of the CRS operation.

[0051] In some implementations, to create cycles 200a-f in mesh 700, as described later, each node from A1 to C3 can include multiple cores, allowing each node to use a different core to process the received data. For example, each node from A1 to C3 can include two or more cores, each acting in a similar manner to a node with a single core. Cores can be internally connected to nodes to form cycles 200a-200f in mesh 700, as described later. Figure 10 As described.

[0052] Other topologies are also possible. For example, a 16×4 topology can perform CRS operations. This topology could include a loop link in a dimension with 16 nodes, instead of a loop link in another dimension with 4 nodes. The algorithm would be modified to use circular reduction in dimensions with loop links and some form of line reduction, such as the single-path algorithm described earlier, in dimensions without loop links. Each algorithm can be implemented in a distributed manner via processes or threads running on processing units.

[0053] Turning Figure 9 The diagram illustrates a multidimensional topology 910 including a torus, where each node comprises multiple processing cores. Loops 900a to 900f are formed by inter-node links and intra-node links, and are connected to... Figure 4 and Figure 6 Loops 200a through 200f are similar. For example, node 902 includes processing cores A1 and A2, each of which acts as a node, as previously described. Thus, loop 900a includes six cores A1 through A6, loop 900b includes six cores B1 through B6, loop 900c includes cores C1 through C6, loop 900d includes cores A1 through A2, B1 through B2, and C1 through C2, loop 900e includes cores A3 through A4, cores B3 through B4, and nodes C3 through C4, and loop 900f includes cores A5 through A6, B5 through B6, and C5 through C6. When this topology is mapped to a physical computing system, each node (such as node 902) represents a chip, and the cores, such as cores A1 through A2 of node 902, represent the processing cores of the chip. While each chip is shown as having two cores, any number of cores physically located on the chip can be included as cores in loops 900a through 900f.

[0054] The physical links that map loops to chips are now described. Loop neighbors 900a through 900f are mapped to physical toroidal links. Additionally, local contributions between processing cores on a chip (e.g., chip 902) are summed. This summation operation is pipelined using summation operations across peer chips on the ICI network.

[0055] As previously mentioned, in the rotating clamp algorithm, a bidirectional loop with full-duplex links is used to reduce latency in stages (e.g., the reduce-scatter and aggregate stages). Compared to a unidirectional loop, a bidirectional loop reduces synchronization and buffering overhead because the loop neighbors are no more than one step apart. Therefore, a 1-bit adder can be used to specify the address of the target buffer on the downstream neighbor for sending fragments. Conversely, with a unidirectional loop, the receiver must send an explicit flow control packet back to the transmitter when the buffer on the receiver is available to receive packets.

[0056] In a toroidal network (such as the toroidal network shown in topology 910), a bidirectional ring is directly mapped to the physical links of the toroid. For example... Figure 9 As shown, the cores of each chip can then be linked to form two-dimensional loops 900a to 900f. This scheme utilizes toroidal links in a toroidal network. Here, loops 900a to 900f are constructed along the dimensions of the toroidal network along the nodes (e.g., chips). Within each chip, data fragments arrive at a core (e.g., A1 of chip 902). The packets are summed and then sent downstream to a peer tensor core on the same chip (e.g., A2 of chip 902). In the next step, data fragments (e.g., data packets) are sent from a second core on a chip node to the first core of a downstream chip. This zigzag transmission ensures that half of the communication from any tensor core is within the chip.

[0057] Since any core's peer (e.g., a neighbor) is within the same chip, this scheme reduces the communication load on the inter-chip toroidal links from network packets by a factor of two compared to a mesh scheme.

[0058] Although toroidal links map directly to hardware links, in a mesh topology, some transmissions on a bidirectional ring will share links with transmissions at other peers. (Turn around) Figure 10 The diagram illustrates a 2D mesh topology 1010 with cycles 1000a to 1000f. In this scheme, links of a single mesh dimension are used in both directions to construct cycles 1000a to 1000f. At mesh corners (such as corner 1002), data fragments are routed between two cores (e.g., cores A5 and A6) of the corner chip. In this method, neighbors in the bidirectional rotation clamping algorithm are physical neighbors on the mesh network or peers on the same chip at the corner. Therefore, as... Figure 10As shown, loop 1000a includes nodes A1, A2, B2, C1, C2, and B1; loop 1000b includes nodes A3, A4, B4, C4, C3, and B3, and so on. Similar to grid 700, different colors 1 and 2 can be executed in parallel on topology 910. The links shared by both colors 1 and 2 are shown in bold and dashed lines. Additional buffers can be used to ensure that colors can be executed in parallel in this scheme.

[0059] Figure 9 and Figure 10 The connections between the cores of each chip in topologies 910 and 1010 are determined based on the availability of toroidal links. In some implementations, if toroidal links are available, the hardware is configured according to topology 910 because toroidal loops 900a to 900f require fewer inter-chip links compared to the aforementioned mesh loops 1000a to 1000f. However, when toroidal links are unavailable, the system can use mesh topology 1010 to configure the loops and still form loops for performing circular algorithms in each dimension.

[0060] When performing two-color operations in parallel on both the 2D mesh and the torus, the first communication is actually performed along the rows of the first color, then along the columns of the first color, and first along the columns of the second color, then along the rows of the second color. This has proven optimal when each chip has two cores and results in the use of all network links in each reduction scattering and aggregation phase. The rotation clamp algorithm execution is generally determined by the first phase. In the second phase, the amount of data transmitted is scaled down to the size of the first dimension. Therefore, the second phase does not add significant overhead to the multidimensional rotation clamp algorithm because the payload transmitted over the network is significantly smaller than that in the first phase.

[0061] To achieve high utilization on the toroidal link (e.g., highly parallel DMA transfers), a single descriptor can be injected along all four directions into the clockwise and counterclockwise neighbors of each of the two colors. This can be done for each fragmentation step of the rotation clamp algorithm. Propagating descriptors in this way prevents head-end blocking, where a descriptor along one direction can block descriptors along other directions. The number of injection cycles is determined by the fragment payload size, which depends on the cross-copy size and the toroidal dimension.

[0062] The floating-point summation operation of cross-copy and concatenation can be performed by a vector operation with a payload in the vector memory. To have a payload in the vector memory (VMEM), it must be transferred from high-bandwidth memory (HBM) to the VMEM via a DMA operation. The summed output is sent downstream in a loop. The size of the fragments in the VMEM determines the throughput achieved on the toroidal link. If the fragment size is small, the execution of the transfer will be determined by the latency of the toroidal link. Each transfer of color (clockwise or counterclockwise) may require five buffers of fragment size, which can be used in parallel at any time as cross-copy and concatenation are performed on the chip core. These fragment-sized buffers are used to transfer data from HBM to VMEM and to receive fragments from upstream replicas. A 1-bit counter can be used to synchronize the transmitter and receiver. Fragment-sized buffers can be used to add the network contribution to the corresponding local contribution in the VMEM. Fragment-sized buffers can be used to transfer the summed result to downstream replicas (e.g., nodes).

[0063] Figures 11 to 12 This is a flowchart illustrating the operations used to perform CRS operations in a multidimensional space. Figure 11 In this process 1100, process 1102 involves training (1102) a corresponding copy of the machine learning model on each of a plurality of nodes organized in a multidimensional topology comprising rows and columns of nodes. Each node is trained in parallel on a corresponding batch of training data. After training, each node can save the corresponding gradient vector obtained during training. Process 1100 includes combining (1104) the corresponding gradient vectors in the nodes to generate a final gradient vector. Figure 12 The process 1200 describes the combination of the corresponding gradient vectors.

[0064] Figure 12 This is a flowchart illustrating an example process 1200 for combining the corresponding gradient vectors of each node in a multidimensional topology. Process 1200 includes performing a first part of the first stage of the circle algorithm (1202) in parallel on each row of the multidimensional topology via code executed on the nodes to generate combined data for each corresponding field of the gradient vector for that row. Process 1200 includes performing a second part of the first stage of the circle algorithm (1204) in parallel on each column of the multidimensional topology via code executed on the nodes to generate a column result in each column, which includes a portion of the combined data for each corresponding field of the gradient vectors of one or more rows. In some implementations, the first and second parts (which may be referred to as colors) may be executed in parallel with each other to fully utilize the DMA transfer of the multidimensional topology. In some implementations, the first stage is referred to as the reduction-dispersion stage.

[0065] Process 1200 includes executing, in parallel, the first part of the second stage of the (1206) circle algorithm on each column of the multidimensional topology via code executed on the node, to generate a portion of the final gradient vector in each column that includes the result of each column. Process 1200 includes executing, in parallel, the second part of the second stage of the (1208) circle algorithm on each row of the multidimensional topology via code executed on the node, to generate a final gradient vector in each row that includes each portion of the final gradient vector from each column. In some implementations, the first and second parts of the second stage can be executed in parallel to fully utilize the DMA transfer of the multidimensional topology. The second stage may be referred to as the aggregation stage.

[0066] For each of processes 1100 and 1200, each node in the multidimensional topology may include at least two child nodes (or even more child nodes). Process 1200 may further include, during the first and second phases of executing the circle algorithm, utilizing code executed on each node to pipeline data transfer between each node so that summation and transfer operations can be performed in parallel on one or more nodes.

[0067] In some implementations, the circle algorithm includes the aforementioned rotation clamping algorithm. The multidimensional topology can include toroidal topology or mesh topology, as previously described. For a toroidal topology, the child nodes of each node (which can represent chips in the toroidal network) can then be linked within the node. In some implementations, when the multidimensional topology includes a mesh topology, and when each node in the mesh includes at least two child nodes, each node can be linked via bidirectional links (e.g., full-duplex links).

[0068] Processes 1100 and 1200 may include mapping the topology of the nodes to a physical computing architecture by determining the type of the multidimensional topology of the multiple nodes, and mapping a first part of the first phase to a first set of nodes and a first set of links between the first set of nodes based on the type. Additionally, the process includes mapping a second part of the first phase to a second set of nodes and a second set of links between the nodes, where the first set of links differs from the second set of links.

[0069] In some implementations, in addition to rows and columns, the multidimensional topology includes a third-dimensional set of nodes to form a three-dimensional topology. Here, the first phase of the circle algorithm includes a corresponding third part, executed by the third-dimensional set of nodes (e.g., adding a third color). Additionally, the second phase of the circle algorithm includes a corresponding third part, executed by the third-dimensional set of nodes (to complete the execution of the third color).

[0070] In some implementations, the combined data for each corresponding field of the gradient vector in each row includes the sum of the data for each corresponding field of the gradient vector in that row. Similarly, the combined data for each corresponding field of the gradient vector in each column includes the sum of the data for each corresponding field of the gradient vector in that column.

[0071] This specification uses the term "configuration" in relation to system and computer program components. For a system of one or more computers to be configured to perform a particular operation or action, this means that the system has software, firmware, hardware, or a combination thereof installed thereon that causes the system to perform the operation or action during operation. For one or more computer programs to be configured to perform a particular operation or action, this means that one or more programs include instructions that, when executed by a data processing device, cause that device to perform the operation or action.

[0072] Embodiments of the subject matter and functional operation described in this specification can be implemented in digital electronic circuit systems, in tangibly embodied computer software or firmware, in computer hardware including the structures disclosed in this specification and their equivalents, or in combinations thereof. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium for execution by a data processing apparatus or for controlling the operation of a data processing apparatus. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination thereof. Alternatively or additionally, the program instructions can be encoded on artificially generated propagation signals (e.g., machine-generated electrical, optical, or electromagnetic signals) generated for encoding information to be transmitted to a suitable receiver device for execution by the data processing apparatus.

[0073] The term "data processing apparatus" 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 apparatus may also be or further include a dedicated logic circuit system, such as a FPGA (Field-Programmable Gate Array) or an ASIC (Application-Specific Integrated Circuit). In addition to hardware, the apparatus may optionally include code that creates an execution environment for computer programs, such as code constituting processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of these.

[0074] The data processing apparatus used to implement machine learning models may also include, for example, dedicated hardware accelerator units for handling the general and computationally intensive portions of machine learning training or production, i.e., inference, and workloads.

[0075] Machine learning models can be implemented and deployed using machine learning frameworks, such as the TensorFlow framework, the Microsoft Cognitive Toolkit framework, the Apache Singa framework, or the Apache MXNet framework.

[0076] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes backend components (e.g., as a data server), middleware components (e.g., an application server), or frontend components, such as a client computer having a graphical user interface, a web browser, or an application through which a user can interact with the implementation of the subject matter described in this specification, or any combination of one or more such backend, middleware, 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.

[0077] A computing system may include clients and servers. Clients and servers are typically geographically separated and usually interact via a communication network. The client-server relationship is established by means of computer programs running on corresponding computers and having a client-server relationship with each other. In some embodiments, the server transmits data (e.g., HTML pages) to a user device (e.g., to display data to a user interacting with the device (which acts as a 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 from the device at the server.

[0078] While this specification contains numerous details of specific implementations, these details should not be construed as limiting the scope of any invention or the scope of what may be claimed, but rather as descriptions of features specific to particular embodiments of a particular invention. In this specification, the terms “vector,” “vector data,” and “vector element” are used broadly to refer to any non-scalar data. Examples of non-scalar data besides vectors include strings, arrays, structures, matrices, and tensors.

[0079] The reduction and update functions described in this specification are merely examples and are entirely independent of the invention itself. Additionally, the invention is described for machine learning, but can be used for any purpose involving reducing data distributed across a network.

[0080] Certain features described in this specification in 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 embodiments, or even initially claimed in this way, one or more features from the claimed combination may be removed from the combination in some cases, and the claimed combination may involve sub-combinations or variations thereof.

[0081] Similarly, although the operations are depicted in a specific order in the diagrams, this should not be construed as requiring such operations to be performed in the specific order shown or in a sequential order, or as performing all the operations shown to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous. Furthermore, the separation of the 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.

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

Claims

1. A method of processing training data using processing hardware, the method comprising: training a respective copy of a machine learning model on each of a plurality of nodes organized in a multi-dimensional topology comprising rows and columns of nodes, wherein each node is trained in parallel on a respective batch of training data, and where each node holds a respective gradient vector resulting from the training after the training; combining respective gradient vectors in the nodes to generate a final gradient vector by performing operations comprising: performing, by code executing on the nodes, a first phase of a round algorithm in parallel on each row of the multi-dimensional topology to generate different combination data for a portion of the final gradient vector in each node of the row; performing, by code executing on the nodes in parallel while performing the first phase of the round algorithm in parallel on each row of the multi-dimensional topology, a second portion of the first phase of the round algorithm in parallel on each column of the multi-dimensional topology to generate different column results in each node of the column, the column results comprising portions of the combination data from each row; performing, by code executing on the nodes, a first portion of a second phase of the round algorithm in parallel on each column of the multi-dimensional topology to generate a portion of the final gradient vector in each column comprising each column result; and performing, by code executing on the nodes, a second portion of the second phase of the round algorithm in parallel on each row of the multi-dimensional topology to generate the final gradient vector in each row, the final gradient vector comprising each portion of the final gradient vector from each column; determining a type of the multi-dimensional topology of the plurality of nodes; based on the type, mapping the first portion of the first phase to a first set of nodes and a first set of links between the first set of nodes; and mapping the second portion of the first phase to a second set of nodes and a second set of links between the nodes, the first set of links being different from the second set of links.

2. The method of claim 1, the method further comprising: determining that the multi-dimensional topology of the plurality of nodes comprises wrap-around links; and configuring each row and each column of the multi-dimensional topology to perform the round algorithm using the wrap-around links.

3. The method of claim 1, wherein, each node of the multi-dimensional topology comprises at least two sub-nodes.

4. The method of claim 3, wherein, the multi-dimensional topology comprises a torus topology, and wherein for each row and each column, the sub-nodes of each node are configured as adjacent nodes of the round algorithm.

5. The method of claim 3, wherein, the multi-dimensional topology comprises a grid topology, and wherein for each row and each column, the sub-nodes of edge nodes of the grid are configured as adjacent nodes of the round algorithm.

6. The method of claim 1, wherein, in addition to the rows and the columns, the multi-dimensional topology comprises a set of nodes of a third dimension, wherein the method further comprises: performing, by code executing on the set of nodes of the third dimension, a third portion of the first phase of the round algorithm in parallel to generate different third dimension results in each node of the column, the third dimension results comprising portions of the column results from each column; and performing, by code executing on the set of nodes of the third dimension, a fourth portion of the second phase of the round algorithm in parallel to generate the final gradient vector in each node of the third dimension, the final gradient vector comprising each portion of the final gradient vector from each column. a third portion of the first phase of the round algorithm is performed in parallel by code executing on the set of nodes of the third dimension to generate the final gradient vector from each column in each node of the set of nodes of the third dimension.

7. The method of claim 1, further comprising pipelining data transfers between each node with execution of code on each node during performance of the first phase of the round algorithm and performance of the second phase of the round algorithm.

8. The method of claim 1, wherein, the round algorithm comprises a rotation pincer algorithm, wherein the first phase comprises a reduction-scatter phase of the rotation pincer algorithm, and wherein the second phase comprises a global gather phase of the rotation pincer algorithm.

9. The method of any one of claims 1-8, wherein, the portion of the final gradient vector that includes combined data from each row comprises combined data from two or more nodes of the column.

10. A data processing apparatus comprising a plurality of processing units each corresponding to a respective node, wherein the data processing apparatus is configured to perform operations comprising: training respective copies of a machine learning model on each node of a plurality of nodes organized in a multi-dimensional topology comprising rows and columns of nodes, wherein each node is trained in parallel on a respective batch of training data to thereby save, after the training, a respective gradient vector resulting from the training; combining the respective gradient vectors in the nodes to generate a final gradient vector by: performing a first phase of a round algorithm on each row of the multi-dimensional topology in parallel by code executing on the nodes to generate different combined data for a portion of the final gradient vector in each node of the row; performing a second portion of the first phase of the round algorithm on each column of the multi-dimensional topology in parallel by code executing on the nodes in parallel with performing the first phase of the round algorithm on each row of the multi-dimensional topology to generate different column results in each node of the column, the column results comprising a portion of combined data from each row; performing a first portion of a second phase of the round algorithm on each column of the multi-dimensional topology in parallel by code executing on the nodes to generate a portion of the final gradient vector in each column that includes each column result; and performing a second portion of the second phase of the round algorithm on each row of the multi-dimensional topology in parallel by code executing on the nodes to generate the final gradient vector in each row, the final gradient vector comprising each portion of the final gradient vector from each column; determining a type of the multi-dimensional topology of the plurality of nodes; based on the type, mapping the first portion of the first phase to a first set of nodes and a first set of links between the first set of nodes; and based on the type, mapping the first portion of the first phase to a first set of nodes and a first set of links between the first set of nodes; and mapping a second portion of the first phase to a second set of nodes and a second set of links between the nodes, the first set of links being different from the second set of links.

11. The data processing apparatus of claim 10, the operations further comprising: determining that the multi-dimensional topology of the plurality of nodes includes a wrap-around link; and configuring each row and each column of the multi-dimensional topology using the wrap-around link to perform the round algorithm.

12. The data processing apparatus of claim 10, wherein, Each node of the multi-dimensional topology includes at least two child nodes.

13. The data processing apparatus of claim 12, wherein, The multi-dimensional topology includes a torus topology, and wherein for each row and each column, the child nodes of each node are configured as neighboring nodes of the round algorithm.

14. The data processing apparatus of claim 12, wherein, The multi-dimensional topology includes a mesh topology, and wherein for each row and each column, the child nodes of edge nodes of the mesh are configured as neighboring nodes of the round algorithm.

15. The data processing apparatus of claim 10, wherein, In addition to the rows and the columns, the multi-dimensional topology includes a third set of nodes, wherein the operations further comprise: in parallel by code executing on the third set of nodes, executing a third portion of the first phase of the round algorithm to generate a different third dimensional result in each node of the columns, the third dimensional result including a portion of the column result from each column; and in parallel by code executing on the third set of nodes, executing a third portion of the first phase of the round algorithm to generate the final gradient vector from each column in each node of the third set of nodes.

16. The data processing apparatus of claim 10, the operations further comprising, during execution of the first phase of the round algorithm and execution of the second phase of the round algorithm, pipelining data transfers between each node with execution of code on each node.

17. The data processing apparatus of claim 10, wherein, The round algorithm includes a rotation pincer algorithm, wherein the first phase includes a reduce-scatter phase of the rotation pincer algorithm, and wherein the second phase includes a global gather phase of the rotation pincer algorithm.

18. The data processing apparatus of any one of claims 10-17, wherein, The combined data of the portion of the final gradient vector includes sum data from two or more nodes of the rows, and wherein the portion of the combined data from each row that includes sum data includes combined data from two or more nodes of the columns.

19. One or more non-transitory computer-readable storage media encoded with instructions that, when executed by processing hardware, cause the processing hardware to perform operations comprising: training a respective copy of a machine learning model on each node of a plurality of nodes, the plurality of nodes being organized in a multi-dimensional topology including rows and columns of nodes, wherein each node is trained in parallel on a respective batch of training data to thereby save, after the training, a respective gradient vector resulting from the training; combining the respective gradient vectors in the nodes to generate a final gradient vector by: in parallel by code executing on the nodes, performing a first phase of a round algorithm on each row of the multi-dimensional topology to generate different combined data for a portion of the final gradient vector in each node of the row; and in parallel by code executing on the nodes, performing a second phase of the round algorithm on each column of the multi-dimensional topology to generate different combined data for a portion of the final gradient vector in each node of the column. a second portion of the first phase of the round algorithm is executed in parallel by code executing on the nodes for each column of the multi-dimensional topology to generate a different column result in each node of the column, the column result including a portion of the combined data from each row; a first portion of a second phase of the round algorithm is executed in parallel by code executing on the nodes for each column of the multi-dimensional topology to generate in each column a portion of the final gradient vector including each column result; and a second portion of the second phase of the round algorithm is executed in parallel by code executing on the nodes for each row of the multi-dimensional topology to generate in each row the final gradient vector, the final gradient vector including each portion of the final gradient vector from each column; determining a type of a multi-dimensional topology of the plurality of nodes; based on the type, mapping a first portion of the first phase to a first set of nodes and a first set of links between the first set of nodes; and mapping a second portion of the first phase to a second set of nodes and a second set of links between the nodes, the first set of links being different from the second set of links.

20. The non-transitory computer-readable storage medium of claim 19, the operations further comprising: determining that the multi-dimensional topology of the plurality of nodes includes wrap-around links; and configuring each row and each column of the multi-dimensional topology to execute the round algorithm using the wrap-around links.

21. The non-transitory computer-readable storage medium of claim 19, wherein, each node of the multi-dimensional topology includes at least two sub-nodes.

22. The non-transitory computer-readable storage medium of claim 21, wherein, the multi-dimensional topology includes a torus topology, and wherein for each row and each column, the sub-nodes of each node are configured as adjacent nodes of the round algorithm.

23. The non-transitory computer-readable storage medium of claim 21, wherein, the multi-dimensional topology includes a mesh topology, and wherein for each row and each column, the sub-nodes of edge nodes of the mesh are configured as adjacent nodes of the round algorithm.

24. The non-transitory computer-readable storage medium of claim 19, wherein, in addition to the rows and the columns, the multi-dimensional topology includes a third dimension of nodes, wherein the operations further comprise: a third portion of the first phase of the round algorithm is executed in parallel by code executing on the third dimension of nodes to generate a different third dimension result in each node of the column, the third dimension result including a portion of the column result from each column; and a third portion of the first phase of the round algorithm is executed in parallel by code executing on the third dimension of nodes to generate in each node of the third dimension of nodes the final gradient vector from each column.

25. The non-transitory computer-readable storage medium of claim 19, the operations further comprising pipelining data transfers between each node with execution of code on each node during execution of the first phase of the round algorithm and execution of the second phase of the round algorithm.

26. The non-transitory computer-readable storage medium of claim 19, wherein, The circular algorithm includes a rotate-pincer algorithm, wherein the first stage includes a reduction-spreading stage of the rotate-pincer algorithm, and wherein the second stage includes a global gathering stage of the rotate-pincer algorithm.

27. The non-transitory computer-readable storage medium of any of claims 19-26, wherein, The combined data of the portion of the final gradient vector includes sum data from two or more nodes of the row, and wherein the portion of the combined data from each row that includes sum data includes combined data from two or more nodes of the column.

Images