Asynchronous distributed data flow for machine learning workloads

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for distributing machine learning workloads, e.g., computations for training a neural network or computing an inference using a neural network, across multiple hardware accelerators. One of the systems comprises a plurality of accelerator islands, each hardware accelerator island comprising a respective plurality of hardware devices that include a plurality of hardware accelerators and a corresponding host for each of the plurality of hardware accelerators; and a respective scheduler for each of the accelerator islands that is configured to schedule workloads across the plurality of accelerators and corresponding hosts in the accelerator island, wherein the system is configured to: receive data representing a machine learning workload; and assign a respective portion of the machine learning workload to each of the plurality of accelerator islands for scheduling by the respective scheduler for the accelerator island.

The invention claimed is: 
     
       1. A system comprising:
 a plurality of accelerator islands, each accelerator island comprising a respective plurality of hardware devices that include a plurality of hardware accelerators and a corresponding host for each of the plurality of hardware accelerators, wherein the hardware accelerators within each accelerator island are interconnected with one another over an interconnect network, and are connected to the hardware accelerators within another accelerator island over a data center network through their corresponding hosts; and 
 a respective scheduler for each of the plurality of accelerator islands that is configured to schedule workloads across the plurality of accelerators and corresponding hosts in the accelerator island, wherein the system is configured to:
 receive data representing a machine learning workload; and 
 assign a respective portion of the machine learning workload to each of the plurality of accelerator islands for scheduling by the respective scheduler for the accelerator island, wherein the respective scheduler is configured to, when the respective portion of the machine learning workload assigned to the accelerator island is a regular computation, schedule the respective portion of the machine learning workload using parallel asynchronous dispatch. 
 
 
     
     
       2. The system of  claim 1 , wherein the data representing the machine learning workload is data representing a sharded dataflow program comprising a plurality of shards. 
     
     
       3. The system of  claim 2 , wherein assigning the respective portion of the machine learning workload to each of the plurality of accelerator islands comprises assigning one or more shards of the sharded dataflow program to each of the plurality of accelerator islands. 
     
     
       4. The system of  claim 1 , wherein scheduling the respective portion of the computation using parallel asynchronous dispatch comprises:
 generating a schedule that assigns, to each of a set of the hardware accelerators in the accelerator island, a respective set of one or more operations that takes as input and output of one or more respective other operations that are performed by another one of the hardware accelerators in the accelerator island; 
 determining, for each of the set of hardware accelerators, a respective size of the output of the one or more respective other operations; and 
 transmitting, in parallel and to the corresponding host for each of the set of hardware accelerators, respective future data specifying the respective size of the output of the one or more respective other operations. 
 
     
     
       5. The system of  claim 4 , wherein the respective future data causes the corresponding host to (i) allocate memory on the hardware accelerator for storing the output of the one or more respective other operations and (ii) transmit data to a corresponding host of the accelerator assigned to the one or more respective other operations that identifies the allocated memory. 
     
     
       6. The system of  claim 5 , wherein the corresponding host of the accelerator assigned to the one or more respective other operations is configured to cause the accelerator assigned to the one or more respective other operations to transmit the output of the respective other operations to the allocated memory over the interconnect network. 
     
     
       7. The system of  claim 1 , wherein the interconnect network comprises an Inter-Core Interconnect (ICI) network. 
     
     
       8. The system of  claim 1 , wherein the data center network comprises an Ethernet network. 
     
     
       9. One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the operations comprising:
 receiving data representing a machine learning workload by a plurality of accelerator islands, wherein each accelerator island comprises a respective plurality of hardware devices that include a plurality of hardware accelerators and a corresponding host for each of the plurality of hardware accelerators, wherein the hardware accelerators within each accelerator island are interconnected with one another over an interconnect network, and are connected to the hardware accelerators within another accelerator island over a data center network through their corresponding hosts, and wherein each accelerator island has a respective scheduler that is configured to schedule workloads across the plurality of accelerators and corresponding hosts in the accelerator island; and 
 assigning a respective portion of the machine learning workload to each of the plurality of accelerator islands for scheduling by the respective scheduler for the accelerator island, wherein the respective scheduler is configured to, when the respective portion of the machine learning workload assigned to the accelerator island is a regular computation, schedule the respective portion of the computation using parallel asynchronous dispatch. 
 
     
     
       10. The non-transitory computer-readable storage media of  claim 9 , wherein the interconnect network comprises an Inter-Core Interconnect (ICI) network, and wherein the data center network comprises an Ethernet network. 
     
     
       11. A method comprising:
 receiving data representing a machine learning workload by a plurality of accelerator islands, wherein each accelerator island comprises a respective plurality of hardware devices that include a plurality of hardware accelerators and a corresponding host for each of the plurality of hardware accelerators, wherein the hardware accelerators within each accelerator island are interconnected with one another over an interconnect network, and are connected to the hardware accelerators within another accelerator island over a data center network through their corresponding hosts, and wherein each accelerator island has a respective scheduler that is configured to schedule workloads across the plurality of accelerators and corresponding hosts in the accelerator island; and 
 assigning a respective portion of the machine learning workload to each of the plurality of accelerator islands for scheduling by the respective scheduler for the accelerator island, wherein the respective scheduler is configured to, when the respective portion of the machine learning workload assigned to the accelerator island is a regular computation, schedule the respective portion of the computation using parallel asynchronous dispatch. 
 
     
     
       12. The method of  claim 11 , wherein the data representing the machine learning workload is data representing a sharded dataflow program comprising a plurality of shards. 
     
     
       13. The method of  claim 12 , wherein assigning the respective portion of the machine learning workload to each of the plurality of accelerator islands comprises assigning one or more shards of the sharded dataflow program to each of the plurality of accelerator islands. 
     
     
       14. The method of  claim 11 , wherein scheduling the respective portion of the computation using parallel asynchronous dispatch comprises:
 generating a schedule that assigns, to each of a set of the hardware accelerators in the accelerator island, a respective set of one or more operations that takes as input and output of one or more respective other operations that are performed by another one of the hardware accelerators in the accelerator island; 
 determining, for each of the set of hardware accelerators, a respective size of the output of the one or more respective other operations; and 
 transmitting, in parallel and to the corresponding host for each of the set of hardware accelerators, respective future data specifying the respective size of the output of the one or more respective other operations. 
 
     
     
       15. The method of  claim 14 , wherein the respective future data causes the corresponding host to (i) allocate memory on the hardware accelerator for storing the output of the one or more respective other operations and (ii) transmit data to a corresponding host of the accelerator assigned to the one or more respective other operations that identifies the allocated memory. 
     
     
       16. The method of  claim 15 , wherein the corresponding host of the accelerator assigned to the one or more respective other operations is configured to cause the accelerator assigned to the one or more respective other operations to transmit the output of the respective other operations to the allocated memory. 
     
     
       17. The method of  claim 11 , wherein the interconnect network comprises an Inter-Core Interconnect (ICI) network. 
     
     
       18. The method of  claim 11 , wherein the data center network comprises an Ethernet network.