Asynchronous distributed data flow for machine learning workloads

Patent metadata

Field	Value
Publication number	US-12112198-B2
Application number	US-202218082415-A
Country	US
Kind code	B2
Filing date	Dec 15, 2022
Priority date	May 7, 2021
Publication date	Oct 8, 2024
Grant date	Oct 8, 2024

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.

First claim

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 ,

Assignees

• Google Llc

Inventors

• Dean Jeffrey Adgate
• Roy Sudip
• Isard Michael Acheson
• Chowdhery Aakanksha
• Saeta Brennan
• Thekkath Chandramohan Amyangot
• Hurt Daniel William
• Lim Hyeontaek
• El Shafey Laurent
• Schuh Parker Edward
• Barham Paul Ronald
• Pang Ruoming
• Sepassi Ryan
• Ghemawat Sanjay
• Wu Yonghui

Classifications

• G06N3/0495

  Quantised networks; Sparse networks; Compressed networks · CPC title

• G06N3/098

  Distributed learning, e.g. federated learning · CPC title

• G06N3/08

  Learning methods · CPC title

• G06N3/063

  using electronic means · CPC title

• G06F9/5072

  Grid computing · CPC title