Asynchronous distributed dataflow for machine learning workloads

Abstract

Methods, systems, and apparatuses including a computer program encoded on a computer storage medium for distributing a machine learning workload, such as computations for training a neural network or computing inference using a neural network, across a plurality of hardware accelerators, one of the systems comprising a plurality of accelerator islands, each accelerator island comprising a respective plurality of hardware devices including 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 configured to schedule a workload across the plurality of accelerators and corresponding hosts in the accelerator island, the system configured to receive data representing the 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 of the accelerator island.

Description

[Technical Field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims priority to U.S. Provisional Application No. 63 / 186,031, filed May 7, 2021. The disclosure of the prior application is considered part of this application and is incorporated by reference into the disclosure of this application.

[0002] This specification relates to techniques for distributing machine learning workloads, such as computations for training a neural network or for computing inference using a neural network, across multiple hardware accelerators. [Background technology]

[0003] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from the received input according to the current values ​​of its respective set of weights. The inputs to a given neural network layer are referred to herein as the layer's "input activations," and the outputs of a given neural network layer are referred to herein as the layer's output activations. [Prior art documents] [Non-patent literature]

[0004] [Non-Patent Document 1] Abadi, M., et al., TensorFlow: A system for large-scale machine learning. In 12th USENIX Symposium on Operating Systems Design and Implementation(OSDI), Savannah, GA, November 2016. USENIX Association [Non-patent document 2] Paszke, A., et al., PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019 [Non-patent document 3] Bradbury, J., et al., JAX: Composable transformations of Python+NumPy programs. http: / / github.com / google / jax, 2018 [Non-patent document 4] Clarke, L., et al., The MPI message passing interface standard Massively Parallel Distributed Systems, 1994 Summary of the Invention [Means for solving the problem]

[0005] This specification describes a system, implemented as a computer program on one or more computers at one or more locations, that distributes machine learning workloads across multiple hardware accelerators and corresponding hosts.

[0006] A hardware accelerator (or simply "accelerator") is a computing device having specialized hardware configured to perform specialized computations, including, for example, machine learning computations. Examples of accelerators include application-specific integrated circuits ("ASICs"), including graphics processing units ("GPUs"), field programmable gate arrays ("FPGAs"), and tensor processing units ("TPUs"). In some implementations, each accelerator has a separate host, while in other implementations, two or more of the accelerators may share a host.

[0007] In particular, the system maintains data partitioning hardware accelerators and their corresponding hosts into multiple accelerator islands. Each accelerator island includes multiple accelerators and their corresponding hosts. In some implementations, all of the accelerators are the same type, while in other cases, different islands can include accelerators of different types, or a single island can include multiple accelerators of different types. In some implementations, the partitioning is static, while in other implementations, the system dynamically adjusts the partitioning based on current system workload.

[0008] Each accelerator island runs a respective scheduler that schedules the computations assigned to the island across the accelerators and hosts within the island. When the computations assigned to a given island are regular, the island's scheduler can schedule the computations using parallel asynchronous dispatching.

[0009] In some implementations, each scheduler in each island is a single scheduler that directly schedules each operation on a given device, while in other implementations, each scheduler is a collection of schedulers that implement a hierarchical scheduling scheme.

[0010] In general, one innovative aspect of the subject matter described in this specification may be embodied in a system comprising: a plurality of accelerator islands, each accelerator island comprising a respective plurality of hardware devices including 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 configured to schedule a workload across the plurality of accelerators and corresponding hosts in the accelerator island, the system 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 of the accelerator island.

[0011] Other embodiments of this aspect include corresponding methods including operations performed by the system and computer programs recorded on one or more computer storage devices, each configured to perform the actions of the method. One or more computer systems can be configured to perform particular operations or actions by software, firmware, hardware, or any combination thereof installed on the system that, during operation, can cause the system to perform the actions. One or more computer programs can be configured to perform particular operations or actions by including instructions that, when executed by a data processing device, cause the device to perform the actions.

[0012] Implementations may include one or more of the following optional features, either alone or in combination.

[0013] The data representing the machine learning workload may be data representing a sharded dataflow program that includes multiple shards.

[0014] Allocating a respective portion of the machine learning workload to each of the plurality of accelerator islands may include allocating one or more shards of a sharded dataflow program to each of the plurality of accelerator islands.

[0015] Each scheduler may be configured to schedule a portion of the computation using parallel asynchronous dispatching when the respective portion of the machine learning workload assigned to the accelerator island is a regular computation.

[0016] Scheduling a portion of a computation using parallel asynchronous dispatching may include generating a schedule that assigns, to each of a set of hardware accelerators in the accelerator island, a respective set of one or more operations that it receives as inputs and outputs of one or more respective other operations executed by another one of the hardware accelerators in the accelerator island; determining, for each of the set of hardware accelerators, respective sizes of outputs of the one or more respective other operations; and transmitting, in parallel, respective future data for each of the set of hardware accelerators to a corresponding host that specifies the respective sizes of outputs of the one or more respective other operations.

[0017] Each future data may cause the corresponding host to (i) allocate memory on a hardware accelerator for storing the output of one or more respective other operations, and (ii) transmit data to the corresponding host of the accelerator assigned to the one or more respective other operations that identifies the allocated memory.

[0018] The corresponding host of the accelerator assigned to one or more respective other operations may be configured to cause the accelerator assigned to one or more respective other operations to send the output of the respective other operation to the allocated memory.

[0019] The output may be transmitted over an accelerator interconnect network.

[0020] When each portion of the machine learning workload assigned to an accelerator island is a regular computation, assigning each portion of the machine learning workload may include assigning each portion by sending a single message to a respective scheduler of the accelerator island.

[0021] Particular embodiments of the subject matter described in this specification can be implemented to realize one or more of the following advantages.

[0022] The techniques described herein can be used to improve hardware utilization, model throughput, or both by orchestrating distributed computations across multiple hardware accelerators to support machine learning workloads, such as training or running machine learning models. In some examples, using the described techniques can achieve performance comparable to state-of-the-art systems (i.e., near 100% hardware utilization) when running "single program, multiple data" (SPMD) computations across many, potentially thousands, of hardware accelerators, while also providing model throughput, e.g., comparable to the SPMD case of a machine learning model that is pipelined across multiple stages or sharded across two or more islands of accelerators connected via a data center network, in terms of the total number (or size) of outputs that can be generated by the model over a given time period. This enables the training, execution, or both of different machine learning models with more diverse architectures, while being highly efficient, including large models that are scaled up using data parallelism, pipelining, or heterogeneous computation and are not well supported by current hardware systems.

[0023] The described techniques for orchestrating distributed computation across multiple hardware accelerators enable a variety of technical use cases that were previously impossible. For example, instead of training one machine learning model from scratch for each single new machine learning task, the described techniques make it possible to train a single large-scale model across thousands or millions of tasks, allowing the model to leverage and combine its existing capabilities to learn new tasks more quickly and effectively.

[0024] As another example, by dynamically mapping subparts of the overall computation to a collection of smaller islands of accelerators that are more readily available, the described techniques have the flexibility to support the execution of multimodal machine learning models that can process information from more than one modality at a time, such as models configured to simultaneously receive visual, auditory, and language data, and possibly other data in more abstract forms, and compute inferences in a less biased and less error-prone manner by exploiting unique common properties of the received data.

[0025] As yet another example, by using sparse cross-host communication to facilitate efficient data-dependent control flow on accelerators, the described techniques can facilitate the execution of sparse machine learning models in an energy-efficient manner. For example, a sparse model can be a "sparsely" activated neural network, with only small routes through relevant parts of the network operated as needed to perform different tasks. Execution of a sparse model can achieve or even exceed state-of-the-art performance in a range of machine learning tasks, while consuming much less energy, sometimes less than one-tenth the energy, than a similarly sized dense model.

[0026] The details of one or more embodiments of the subject matter herein are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, drawings, and claims. [Brief explanation of the drawings]

[0027] [Figure 1] FIG. 1 illustrates an exemplary system for executing machine learning workloads. [Figure 2] FIG. 1 illustrates a comparison of dispatch overhead and communication patterns between a multi-controller system and a single-controller system. [Figure 3]FIG. 1 is a flow diagram of an example process for executing a machine learning workload. [Figure 4A] FIG. 1 is an exemplary diagram of a sharded dataflow program. [Figure 4B] FIG. 1 illustrates an exemplary diagram of allocating machine learning workloads to multiple accelerator islands. [Figure 4C] FIG. 1 is an exemplary diagram of scheduling allocated computations by a scheduler. [Figure 5] FIG. 1 is a flow diagram of an exemplary process for scheduling portions of a computation using parallel asynchronous dispatching. [Figure 6] FIG. 1 illustrates a comparison of sequential and parallel dispatch for an exemplary sharded dataflow program. DETAILED DESCRIPTION OF THE INVENTION

[0028] Like reference numbers and designations in the various drawings indicate like elements.

[0029] This specification describes a system, implemented as a computer program on one or more computers at one or more locations, that distributes machine learning workloads across multiple hardware accelerators and corresponding hosts.

[0030] FIG. 1 illustrates an exemplary system 100 for executing a machine learning workload 104. The machine learning workload 104 can be specified by a client 102. The system 100 can receive data specifying the machine learning workload 104 from the client 102 and generate output data 154 as a result of executing the machine learning workload 104. In some implementations, the data specifying the machine learning workload 104 can include a source program written in the Python programming language by using a suitable Python programming framework such as TensorFlow and JAX. In other implementations, the data can alternatively include a source program written in another high-level programming language such as the C++ language. In some implementations, the machine learning workload can include computations for training a neural network or for computing inference using a neural network. In some of these implementations, the neural network can be a large-scale neural network. A large-scale neural network is a neural network with many network parameters, for example, 1 billion parameters, 10 billion parameters, 100 billion parameters, or 500 billion or more parameters.

[0031] A neural network may generally be configured to perform a machine learning task on a network input via training to generate a network output for the machine learning task.

[0032] Some examples of machine learning tasks that neural networks can be configured to perform are:

[0033] As an example, the task may be a neural machine translation task. For example, if the input to a neural network is a sequence of text in one language, such as a sequence of words, phrases, characters, or word fragments, the output generated by the neural network may be a translation of the sequence of text into another language, i.e., a sequence of text in the other language that is a translation of the input sequence of text. As a particular example, the task may be a multilingual machine translation task in which a single neural network is configured to translate between multiple different source language-target language pairs. In this example, the source language text may be augmented with an identifier indicating the target language into which the neural network should translate the source language text.

[0034] As another example, the task may be an audio processing task. For example, if the input to a neural network is a sequence representing a spoken utterance, the output generated by the neural network may be a score for each of a set of text fragments, each score representing an estimated likelihood that the text fragment is an accurate transcript of the utterance. As another example, if the input to a neural network is a sequence representing a spoken utterance, the output generated by the neural network may indicate whether a particular word or phrase (a "hot word") was spoken during the utterance. As another example, if the input to a neural network is a sequence representing a spoken utterance, the output generated by the neural network may identify the natural language in which the utterance was spoken.

[0035] As another example, the task may be a natural language processing or understanding task, such as an entailment task, a paraphrase task, a text similarity task, a sentiment task, a sentence completion task, a grammar task, etc., that operates on a sequence of text in some natural language.

[0036] As another example, the task may be a text-to-speech task where the input is natural language text or features of natural language text and the network output is a spectrogram, waveform, or other data that defines the speech of the text spoken in the natural language.

[0037] As another example, the task may be a health prediction task, where the input is a sequence derived from the patient's electronic health record data and the output is a prediction related to the patient's future health, for example, a predicted treatment to be prescribed for the patient, the likelihood of the patient experiencing an adverse health event, or a predicted diagnosis for the patient.

[0038] As another example, the task may be a text generation task, where the input is a sequence of text and the output is another sequence of text, e.g., a completion of the input sequence of text, a response to a question posed during the input sequence, or a sequence of text on a topic specified by the first sequence of text. As another example, the input to a text generation task may be non-textual input, e.g., an image, and the output sequence may be text describing the input.

[0039] As another example, the task may be an image generation task, where the input is a conditioning input and the output is a sequence of intensity value inputs for pixels of an image.

[0040] As another example, the task may be an agent control task, where the inputs are a sequence of observations or other data that characterize the state of the environment, and the outputs define actions to be performed by the agent in response to the most recent data in the sequence. The agent may be, for example, a real-world or simulated robot, a control system for an industrial facility, or a control system that controls different kinds of agents.

[0041] As another example, the task may be a genomics task, where the input is a sequence representing a fragment of a DNA sequence or other molecular sequence, and the output is either an embedding of the fragment for use in a downstream task, or an output for a downstream task, for example, by utilizing unsupervised learning techniques on a dataset of DNA sequence fragments. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting the functional impact of non-coding variants, etc.

[0042] In some cases, the machine learning task is a combination of multiple individual machine learning tasks, i.e., the system is configured to perform multiple different individual machine learning tasks, e.g., two or more of the machine learning tasks described above. For example, the system can be configured to perform multiple individual natural language understanding tasks, and the network input includes identifiers for the individual natural language understanding tasks to be performed on the network input.

[0043] In implementations in which system 100 executes machine learning workload 104 to train a neural network, system 100 can receive architecture data that defines the architecture of the neural network. The architecture defines the number of layers in the neural network, the operations performed by each of the layers, and the connections between layers in the neural network, i.e., which layers receive input from which other layers in the neural network.

[0044] The system 100 can also receive training data for training the neural network to perform one or more of the above-mentioned machine learning tasks. Generally, the training data includes a set of neural network inputs and, for each network input, a respective target output to be generated by the neural network to perform a particular task. In some implementations, a larger set of training data can be randomly divided by the system to generate training data and a validation set for evaluating the performance of the neural network on the task.

[0045] System 100 can receive the architecture data and training data in any of a variety of ways. For example, system 100 can receive the architecture data as an upload from client 102 over a data communications network, e.g., using an application programming interface (API) made available by system 100. As another example, system 100 can receive input from client 102 specifying which data already maintained by system 100 or another cloud storage system accessible by the system should be used to train the neural network.

[0046] Once the system 100 has trained the neural network through execution of the machine learning workload 104, the system can provide data specifying the trained neural network for use in processing new network inputs. That is, the system can output the trained values ​​of the network parameters to the client 102, for example, by outputting to a user device or by storing in memory accessible to the system, for later use in processing inputs using the trained neural network.

[0047] Instead of, or in addition to, outputting trained neural network data, system 100 can instantiate an instance of a neural network having trained values ​​of network parameters, receive inputs to be processed, use the trained neural network to process the received inputs to generate outputs, and then provide the generated outputs related to the received inputs. The system can receive the network inputs via an application programming interface ("API") provided by the system. The trained neural network can be used to process any of the various machine learning tasks described above.

[0048] System 100 is typically hosted within a data center, which may be a distributed cloud-based computing system with hundreds or thousands of hardware accelerators, e.g., Hardware Accelerator A 110A through Hardware Accelerator M 110M, in one or more locations. A hardware accelerator (or "accelerator" for short) is a computing device with specialized hardware configured to perform specialized computations, including, for example, machine learning computations. Examples of accelerators include application-specific integrated circuits ("ASICs"), including graphics processing units ("GPUs"), field programmable gate arrays ("FPGAs"), and tensor processing units ("TPUs").

[0049] Because a hardware accelerator can efficiently perform only a subset of operations for which its hardware is optimized, such as matrix multiplication, the hardware accelerator is connected to a host machine, which may be a CPU-based host machine, such as host A 120A and host B 120B, to perform operations that cannot be efficiently performed on the hardware accelerator. The host machine (or "host" for short) is responsible for operations including loading data from cloud storage, preprocessing the data, sending the data to the hardware accelerator, etc. In some implementations, each accelerator has a separate host, while in other implementations, two or more of the accelerators can share a host.

[0050] Each host manages an object store that can store the inputs and outputs of computations running on the corresponding hardware accelerator. The object store can also keep track of buffers held in the hardware accelerator's memory. For example, a client can use an opaque handle to reference an object in a remote host or accelerator memory, allowing the system to migrate the object as needed. The object store can also store intermediate program values, for example, while the system waits to transfer them between accelerators or to pass them to a subsequent computation.

[0051] Each host instantiates an executor that can dispatch, i.e., schedule, the execution of respective portions of the machine learning workload 104 across the hardware accelerators. Execution is scheduled in parallel when possible, for example, by using multiple CPU cores or GPU streams. For example, the executor may be a CPU-based TensorFlow executor that facilitates serialization of input processing into a dataflow graph representing the machine learning workload.

[0052] Although FIG. 1 illustrates one client 102, system 100 can perform computations on behalf of many clients. In other words, system 100 can receive respective data specifying different machine learning workloads from two or more clients, execute the different workloads with at least some concurrency, and generate respective output data as a result of executing the different machine learning workloads. Each client can be physically adjacent to system 100, e.g., located in the same data center as (some parts of) system 100, or can be a cloud client remote from system 100. In the latter case, system 100 can be at least partially controlled by the cloud client. Each client can run on, e.g., a desktop computer, a laptop computer, a tablet computer, a wearable computer, a mobile phone, a smartphone, a music player, an e-reader, a navigation system, or any other suitable computing device. Each client can communicate with system 100 via a data communications network.

[0053] To submit a machine learning workload 104 to be executed by the system 100, a client 102 can first invoke the system's client library, which first allocates virtual resources to any computations defined in the workload 104 that have not previously been executed, registers the computation with the system's resource manager 142, and triggers compilation of the computation in the background. The client 102 can then build an intermediate representation (IR) for the workload, which may be a hardware accelerator location-independent IR expressed as a custom multilevel intermediate representation (MLIR) dialect. The IR is gradually "degraded" through a series of standard compiler passes, ultimately outputting a low-level representation that includes the locations of physical hardware accelerators. This low-level representation takes into account network connectivity between the physical hardware accelerators and includes operations for forwarding output from source computation shards to their destination computation shard locations, including scatter and gather operations when data exchange is required. This low-level representation is then converted into a sharded dataflow program, which may be represented as a dataflow graph.

[0054] Table 1 below is an example of Python code that may be submitted by a client 102 to the system 100 to run constituent computational shards for a machine learning workload across multiple islands of hardware accelerators.

[0055] [Table 1]

[0056] In this example Python code, a client requests one or more sets of virtual resources, i.e., "virtual TPU devices," and then places specific compiled functions on those resources. When requesting virtual resources, the client can optionally specify constraints on resource type, location, or interconnection topology. The system automatically handles all data movement and resharding between dependent computations.

[0057] Using the JAX library, clients can explicitly wrap standard Python code with decorators to indicate fragments to be compiled into accelerated linear algebra (XLA) computations that may follow the SPMD model. These XLA computations are typically characterized by known input and output types and shapes, bounded loops, and conditionals (if any), making it feasible to estimate the resource requirements of the computation in advance. These computations with known resource requirements are referred to herein as "compiled functions."

[0058] For example, a compiled function may include an addition or multiplication calculation, an operation performed by a certain type of neural network layer, etc. As another example, common types of calculations that may each be a single compiled function within a machine learning training workload may include, for example, forward and / or backward passes of a training step to calculate gradients, a function that accumulates gradients calculated over multiple distributed training steps, etc.

[0059] By default, system 100 converts each compiled function into a standalone program containing only one computational shard, meaning that if a client wants to execute multiple functions sequentially, a separate Python call and remote procedure call from the client to the coordinator is required for each function. Therefore, the system also implements a program tracer that allows clients to wrap blocks of Python code that call multiple compiled functions. The program tracer can generate a single program in which each compiled function is represented by a computational node in a dataflow graph that represents the sharded dataflow program.

[0060] Unlike existing clients that coordinate thousands of individual computation and data buffers corresponding to each shard of computation distributed across many hardware accelerators, which can quickly become a performance bottleneck, the clients 102 of the system 100 use a sharded buffer abstraction to represent logical buffers that can be distributed across multiple hardware accelerators. This abstraction helps clients scale by amortizing the cost of bookkeeping tasks (including reference counting) at the granularity of the logical buffer instead of the individual computation shard.

[0061] Typically, existing distributed computing systems for running machine learning workloads employ either a single-controller architecture (e.g., the TensorFlow v1 system described in more detail in Abadi, M., et al., TensorFlow: A system for large-scale machine learning. In 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI), Savannah, GA, November 2016. USENIX Association) or a multi-controller architecture (e.g., the PyTorch system described in more detail in Paszke, A., et al., PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019, or the JAX system described in more detail in Bradbury, J., et al., JAX: Composable transformations of Python+NumPy programs. http: / / github.com / google / jax, 2018).

[0062] Figure 2 compares dispatch overhead and communication patterns between multi-controller and single-controller systems. Specifically, Figure 2 includes a swimlane diagram illustrating actions that may be performed by the host ("Host"), controller ("Ctrlr"), or hardware accelerator ("Dev"). As shown in Figure 2(a), JAX or PyTorch SPMD individually enqueues accelerator computations asynchronously over a high-speed Peripheral Component Interconnect express (PCIe) network, (b) TensorFlow v1 SPMD requires control messages over a slower data center network (DCN), and (c) TensorFlow v1 non-SPMD programs require cross-host coordination or data transfer via explicit send (S) and receive (R) operations (ops).

[0063] As used herein, the term "SPMD" refers to the "single program, multiple data" (SPMD) model, described in more detail in Clarke, L., et al., The MPI message passing interface standard Massively Parallel Distributed Systems, 1994, where all accelerators execute the same program (e.g., the same computation) in lockstep, whereas "MPMD" refers to the "multiple program, multiple data" (MPMD) model, where multiple accelerators execute at least two independent programs (e.g., two different computations) simultaneously.

[0064] Some single-controller systems offer a very general distributed dataflow model, including optimized intra-graph control flow. In a single-controller system, a client—for example, a TensorFlow (TF) Python client—can construct a computational graph and pass it to a coordinator runtime, which then partitions the graph into respective subgraphs, one per worker, and delegates subgraph execution to the worker's local runtime. Coordination between workers is performed using data edges and control edges that pass messages over a data center network (DCN). While single-controller frameworks offer a flexible programming model and resource virtualization, they present implementation challenges, as discussed below. First, clients in a single-controller system are "farther away," and dispatch latency involves communication over the DCN, which is typically an order of magnitude slower than PCIe (see, for example, Figure 2b). Second, to support the concurrent execution of multiple-program, multiple-data (MPMD) programs with SPMD subcomputations, each spanning a subset of accelerators selected from a shared cluster, the runtime must have some mechanism for supporting gang scheduling of accelerator computations. Gang scheduling is essential because some hardware accelerators, such as TPUs, are single-threaded and execute only non-preemptive kernels, resulting in deadlock if communicating computations are not enqueued in a consistent order. Therefore, a single-controller system for executing machine learning workloads requires a distributed scheduling mechanism to order queued computations on behalf of different programs. Finally, systems for modern machine learning workloads must be designed to execute computations distributed across thousands of accelerators, with first-class support for sharded representations and data structures.For example, a naive dataflow graph representing edges between an M-way sharded computation and an N-way sharded computation would require M+N nodes and M×N edges, which quickly becomes unwieldy.

[0065] In contrast, in a multi-controller system, the same client executable runs directly on all hosts in the system and exclusively owns the resources on those hosts for the duration of the program execution. One advantage of a multi-controller architecture is low latency for dispatching accelerator computations because identical copies of the client code run on each accelerator host and dispatching involves only communication over a relatively fast PCIe network (see, for example, Figure 2a). All other communication between hosts occurs exclusively through the collective, using dedicated interconnect networks such as NVLink and Inter-Core Interconnect (ICI), without going through host memory. However, multi-controller architectures are a poor match for modern machine learning workloads that use pipelining, computational sparsity, or both. Any communication beyond the standard collective in a multi-controller system requires users to implement their own coordination primitives. Additionally, multi-controller approaches typically assume exclusive ownership of hardware resources. This not only shifts the responsibility for ensuring high utilization of expensive accelerators onto users, but also complicates the design of features like resource virtualization and multiplexing that are needed to build an efficient cluster-wide machine learning infrastructure.

[0066] At a high level, the system 100 described herein employs a single-controller framework that combines the flexibility of single-controller systems with the performance of multi-controller systems. The architecture of system 100 differs from existing single-controller systems in that system 100 uses asynchronous dispatching to match the performance of multi-controller systems, supports centralized resource management and scheduling with first-class support for gangs of SPMD accelerator computations, and uses a sharded dataflow system for efficient coordination. In addition to matching the functionality and performance of modern systems, the architecture of system 100 also provides the capabilities needed to support a wider range of machine learning workloads, including the ability to support the training or execution of machine learning models that are large, sparse, and / or irregular.

[0067] In particular, system 100 uses resource manager 142 to maintain, i.e., generate or update, data specifying the division of hardware accelerators and their corresponding hosts into multiple accelerator islands. Resource manager 142 is responsible for centralized management of devices, including hardware accelerators, hosts, and schedulers, across all of the islands. Resource manager 142 can track all available devices in system 100, thus allowing underlying computational resources to be dynamically added and removed from the system. Resource manager 142 can also maintain a mapping database that defines a one-to-one mapping of virtual resources that may be requested by client 102 to actual physical devices in system 100. For example, client 102 may request a virtual set of hardware accelerators with a particular 2D or 3D mesh shape that fits its specific communication and / or computation needs for machine learning workload 104. In response, resource manager 142 can dynamically allocate physical hardware accelerators that meet the specific interconnection topology and / or memory capacity for the requested virtual resources.

[0068] In some implementations, resource manager 142 may employ a simple heuristic algorithm that attempts to statically balance the load by distributing computations across all available devices. In other implementations, resource manager 142 may employ a more sophisticated allocation algorithm, for example, that may approximate an optimal allocation of physical devices to computations, taking into account the resource requirements of all client computations and the current state of the system.

[0069] In some implementations, all of the accelerators in system 100 are the same type of accelerator, while in other implementations, different islands can include different types of accelerators, or a single island can include multiple different types of accelerators. In some implementations, the partitioning is static, while in other implementations, resource manager 142 dynamically adjusts the partitioning based on the current system workload.

[0070] Each accelerator island includes multiple accelerators and their corresponding hosts. For example, as shown in FIG. 1, system 100 maintains data partitioning hardware accelerators and their corresponding hosts in two accelerator islands 140A-B. Accelerator island 140A includes hardware accelerators A 110A-H 110H and host A 120A, and accelerator island 140B includes hardware accelerators J 110J-M 110M and host B 120B.

[0071] The hardware accelerators in each island are interconnected with each other via an interconnection network and connected to hardware accelerators in other islands via corresponding hosts via a data center network. For example, the interconnection network can be an Inter-Core Interconnect (ICI) network, and the data center network (DCN) can be an Ethernet network. A scalable, general-purpose data flow engine can be used by the system 100 to manage DCN communication between accelerator islands. The data flow engine can also be used for background housekeeping tasks, including configuration information distribution, program monitoring, error distribution, etc.

[0072] Each accelerator island runs a respective scheduler, e.g., scheduler A 130A for island 140A and scheduler B 130B for island B 140B, which schedules the computations assigned to the island across the accelerators and hosts within the island. Each scheduler may be configured to receive a portion of the machine learning workload and assign operations to hardware accelerators included in the same accelerator island as the scheduler. When the computations assigned to a given island are regular, the island's scheduler schedules the computations using parallel asynchronous dispatching.

[0073] In some implementations, each scheduler in each island is a single scheduler that directly schedules each operation on a given device, while in other implementations, each scheduler is a collection of schedulers that implement a hierarchical scheduling scheme.

[0074] The schedule is configured to schedule computations assigned to an island across accelerators and hosts within the island within strict timing requirements, e.g., on a millisecond timescale, to achieve normal operation of the system. In some implementations, the scheduler may simply enqueue execution of portions of the machine learning workload 104 in first-in-first-out (FIFO) order, while in some other implementations, the scheduler may employ more sophisticated scheduling algorithms, such as reordering computations based on estimated execution times.

[0075] 3 is a flow diagram of an exemplary process 300 for executing a machine learning workload. For convenience, process 300 is described as being performed by one or more computer systems located at one or more locations. For example, a distributed computing system, such as system 100 of FIG. 1, suitably programmed in accordance with this specification, may perform process 300.

[0076] The system maintains data partitioning hardware accelerators and their corresponding hosts into multiple accelerator islands. Each accelerator island includes multiple accelerators and their corresponding hosts. In some implementations, all of the accelerators are the same type, while in other implementations, different islands can include accelerators of different types, or a single island can include multiple accelerators of different types. In some implementations, the partitioning is static, while in other implementations, the system dynamically adjusts the partitioning based on current system workload.

[0077] The system receives data representing a machine learning workload (step 302). In some cases, the system may receive the data from a client over a data communications network. In some cases, the machine learning workload includes computations for training a neural network, and the data representing the machine learning workload includes data representing a sharded dataflow program including multiple computational shards. The machine learning workload may include multiple portions, each corresponding to a respective compiled function.

[0078] 4A is an exemplary diagram of a sharded dataflow program. As shown, the sharded dataflow program can be represented as a dataflow graph, which can be, for example, a directed acyclic graph, where each node represents a computational shard, e.g., an individual compiled function, and edges between nodes represent data flow between computational shards, e.g., between compiled functions. Each computational shard can be represented by a single node in the dataflow graph. For example, the machine learning workload represented by the sharded dataflow program shown in FIG. 4A includes three computational shards, A, B, and C, each of which can be a respective compiled function.

[0079] In response to receiving the data representing the machine learning workload, the system assigns a respective portion of the machine learning workload to each of the plurality of accelerator islands for scheduling by the accelerator island's respective scheduler (step 304). Specifically, one or more computational shards of the sharded dataflow program representing the machine learning workload may be assigned to each of the plurality of accelerator islands.

[0080] 4B is an exemplary diagram of allocating a machine learning workload to multiple accelerator islands. As shown, the system uses a resource manager to allocate a portion of the available computational resources, i.e., a subset of the system's hardware accelerators, to each constituent computational shard of the machine learning workload. The resource manager can determine the allocation by using any suitable load balancing algorithm or technique to load balance the system's hardware accelerators.

[0081] In the example of FIG. 4B, the system assigns compute shard A to a first accelerator island 410, compute shard B to a second accelerator island 412, and compute shard C to a third accelerator island 414. When model parallelism is used to execute a machine learning workload, each hardware accelerator can store a respective portion of the architecture of data (a "model component") that defines the architecture of a corresponding portion of the neural network.

[0082] FIG. 4C illustrates an exemplary diagram of scheduling assigned computations by a scheduler. Each accelerator island runs a respective scheduler that schedules the island's assigned computations across accelerators and hosts within the island. This configuration can support gang scheduling of accelerator computations for concurrent execution of machine learning workloads while avoiding deadlocks. As shown, each island's scheduler gang-schedules its assigned computations to hosts within the island, which are then dispatched by the hosts (by instantiated executors) to hardware accelerators within the island for execution. In FIG. 4C, arrows 420 and 422 indicate control messages, and arrows 424, 426, and 428 indicate datapath transfers. Gang scheduling is a type of co-scheduling that allows groups of computations, including heterogeneous computations, to be scheduled simultaneously on different hardware accelerators and proceed in parallel as a group.

[0083] When each portion of a machine learning workload assigned to an accelerator island is a regular computation, e.g., a regular compiled function with known resource requirements, the accelerator island's scheduler can schedule the portion of the computation using parallel asynchronous dispatching. In these cases, the system can assign each portion of the machine learning workload by sending a single message to each scheduler of the accelerator island, the single message describing a subgraph of the computation graph corresponding to each portion of the machine learning workload. The scheduler is configured to execute all active computation shards in the subgraph sequentially, in order. The use of a single message minimizes network traffic, but does not require the scheduler to actually enqueue all computation shards in the subgraph as a batch, since computations may still be executed interleaved with those submitted by other concurrently running programs.

[0084] 5 is a flow diagram of an exemplary process 500 for scheduling portions of a computation using parallel asynchronous dispatching. For convenience, process 500 is described as being performed by one or more computer systems located at one or more locations. For example, a distributed computing system, such as system 100 of FIG. 1, suitably programmed in accordance with this specification, may perform process 500.

[0085] The system generates a schedule that assigns to each of a set of hardware accelerators in the accelerator island a respective set of one or more operations that it receives as inputs and outputs for one or more respective operations executed by another one of the hardware accelerators in the accelerator island (step 502). The one or more operations may be computations for executing a respective portion of the machine learning workload assigned to each accelerator island. For example, the one or more operations may be constituent operations of a compiled function.

[0086] The system determines, for each of the set of hardware accelerators, the respective sizes of the outputs of one or more respective other operations (step 504).

[0087] The system sends, in parallel, for each of the set of hardware accelerators to the corresponding host, respective future data specifying respective sizes of the outputs of one or more respective other operations (step 506).

[0088] Each future data causes the corresponding host of each of the set of hardware accelerators to (i) allocate memory on the hardware accelerator for storing the output of one or more respective other operations, and (ii) transmit data to the corresponding host of the accelerator assigned to the one or more respective other operations that identifies the allocated memory.

[0089] Also, the corresponding hosts of the accelerators assigned to the one or more respective other operations are configured to cause the accelerators assigned to the one or more respective other operations to transmit outputs of the respective other operations to the allocated memories via the interconnection network.

[0090] Figure 6 illustrates a comparison between sequential and parallel dispatching for an exemplary sharded dataflow program. In Figure 6, the exemplary sharded dataflow program is represented by a dataflow graph that includes three nodes A, B, and C (shown as circles), each corresponding to a respective computational shard. Each computational shard is a regular compiled function. The computational shards corresponding to the three nodes A, B, and C are scheduled (shown as squares) to run on accelerators ("Dev") connected to hosts ("Hosts") A, B, and C.

[0091] As shown on the left side of Figure 6, in the case of sequential dispatching, Host A enqueues compute shard A, receives future data specifying the size of compute shard A's output, and sends the future data to Host B. Host B allocates memory on a hardware accelerator to store compute shard B's input, sends the input buffer address to Host A, and performs some of the preparatory work for launching compute shard B's function. When compute shard A completes, its output is sent directly to compute shard B's input buffer via the interconnection network, and then Host B starts compute shard B. The latency between the completion of one compute shard and the start of the next can be on the order of the data transfer time.

[0092] This typically works well when the computation of the predecessor node takes a reasonable amount of time, but in the case shown in Figure 6, where the computation time is shorter than the time spent on scheduling, resource allocation, and coordination between hosts, the asynchronous pipeline backs up (due to the host-side work in sequential dispatching), and the host-side work becomes a significant bottleneck for executing the overall sequence of computations.

[0093] Therefore, the system overcomes this bottleneck in traditional sequential dispatching by performing host-side work in parallel using parallel asynchronous dispatching, as shown on the right side of Figure 6.

[0094] In particular, if the compiled functions are all regular, the size of the input data of a successor node can actually be pre-computed, i.e., calculated before the preceding computation is enqueued. The parallel asynchronous dispatch mechanism takes advantage of the statically known resource usage of regular compiled functions to perform most of the host-side work for the compute nodes in parallel, rather than serializing such that a node's work occurs after the previous node has been enqueued.

[0095] This specification uses the term "configured" with respect to systems and computer program components. A system of one or more computers configured to perform a particular operation or action means that the system has installed thereon software, firmware, hardware, or a combination thereof that, during operation, causes the system to perform the operation or action. A computer program or programs configured to perform a particular operation or action means that the program or programs contain instructions that, when executed by a data processing apparatus, cause the apparatus to perform the operation or action. Embodiments of the subject matter and functional operations described herein can be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware, including structures disclosed herein and their structural equivalents, or one or more 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 or 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 one or more combinations thereof. Alternatively, or additionally, the program instructions may be encoded on an artificially generated propagated signal, for example, a machine-generated electrical, optical, or electromagnetic signal generated to encode information for transmission to a suitable receiver device for execution by a data processing device.

[0096] The term "data processing apparatus" refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including, for example, a programmable processor, a computer, or multiple processors or computers. An apparatus may also be or include special-purpose logic circuitry, such as, for example, an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). In addition to hardware, an apparatus may optionally include code that creates an execution environment for a computer program, such as code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or one or more combinations thereof.

[0097] A computer program, also referred to or written as a program, software, software application, app, module, software module, script, or code, can be written in any form of programming language, including compiled or interpreted, declarative, or procedural, and can be deployed in any form, such as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program can, but does not necessarily, correspond to a file in a file system. A program can be stored in a single file dedicated to the program, or in multiple coordinated files, e.g., files storing one or more modules, subprograms, or portions of code, or in files holding other programs or data, e.g., one or more scripts stored in a markup language document. A computer program can be deployed to be executed on one computer or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communications network.

[0098] As used herein, the term "database" is used broadly to refer to any collection of data, which need not be structured in any particular way, or at all, and which may be stored on storage devices in one or more locations. Thus, for example, an index database may contain multiple collections of data, each of which may be organized and accessed differently.

[0099] Similarly, the term "engine" is used broadly herein to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine is implemented as one or more software modules or components installed on one or more computers at one or more locations. In some cases, one or more computers are dedicated to a particular engine, and in other cases, multiple engines can be installed and executed on the same one or more computers.

[0100] The processes and logic flows described herein may be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows may also be performed by special purpose logic circuitry, such as an FPGA or ASIC, or a combination of special purpose logic circuitry and one or more programmed computers.

[0101] A computer suitable for running a computer program can be based on a general-purpose or special-purpose microprocessor, or both, or on another type of central processing unit. Typically, the central processing unit receives instructions and data from a read-only memory or a random-access memory, or both. The essential elements of a computer are a central processing unit for performing 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 in, special-purpose logic circuitry. Typically, a computer also includes one or more mass storage devices for storing data, such as magnetic, magneto-optical, or optical disks, or is operatively coupled to receive data from, transfer data to, or both of, the one or more mass storage devices. However, a computer need not have such devices. Furthermore, a computer may be incorporated in another device, for example, 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 such as a universal serial bus (USB) flash drive, to name just a few.

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

[0103] To provide for user interaction, embodiments of the subject matter described herein can be implemented on a computer having a display device, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user, as well as a keyboard and pointing device, such as a mouse or trackball, through which the user can provide input to the computer. Other types of devices can also be used to provide for user interaction; 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 acoustic, speech, or tactile input. Furthermore, a computer can interact with a user by sending and receiving documents from a device used by the user, such as by sending a web page to a web browser on the user's device in response to a request received from the web browser. A computer can also interact with a user by sending text messages or other forms of messages to a personal device, such as a smartphone running a messaging application, and receiving a reply message from the user in return.

[0104] A data processing apparatus for implementing machine learning models may also include dedicated hardware accelerator units for handling common and computationally intensive parts of the workload, e.g., machine learning training or production, i.e., inference.

[0105] 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.

[0106] Embodiments of the subject matter described herein can be implemented in a computing system that includes a back-end component, e.g., a data server, or includes a middleware component such as an application server, or includes a front-end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described herein, or includes any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), e.g., the Internet.

[0107] A computing system may include clients and servers. Clients and servers are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server sends data, e.g., HTML pages, to a user device, e.g., to display data to and receive user input from a user interacting with the device acting as a client. Data generated at the user device, e.g., the result of a user interaction, can be received at the server from the device.

[0108] While the specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or the scope that may be claimed, but rather as descriptions of features that may be specific to particular embodiments of a particular invention. Some features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Furthermore, while features are described above as acting in some combinations and are initially claimed as such, in some cases, one or more features from a claimed combination can be deleted from the combination, and the claimed combination may be directed to a subcombination or a variation of a subcombination.

[0109] Similarly, although operations are associated with figures and described in a particular order in the claims, this should not be understood as requiring such operations to be performed in the particular order or sequence shown, or that all of the illustrated operations be performed to achieve desirable results. In some situations, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above-described embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems may generally be incorporated together in a single software product or packaged in multiple software products.

[0110] Specific embodiments of the subject matter are described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. By way of example, the processes corresponding to the accompanying figures do not necessarily require the particular order shown or sequential order to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. [Explanation of symbols]

[0111] 100 systems 102 clients 104 Machine Learning Workloads 110A~110H Hardware Accelerator A~Hardware Accelerator H 110J~110M Hardware Accelerator J~Hardware Accelerator M 120A Host A 120B Host B 130A Scheduler A 130B Scheduler B 140A Island A 140B Island B 142 Resource Manager 154 Output Data 300 processes 410 First Accelerator Island 412 Second Accelerator Island 414 The Third Accelerator Island 500 processes

Claims

1. 1. A system comprising: a plurality of accelerator islands, each accelerator island comprising a respective plurality of hardware devices including a plurality of hardware accelerators and a corresponding host for each of the plurality of hardware accelerators; a respective scheduler for each of the accelerator islands configured to schedule workloads across the plurality of hardware accelerators and corresponding hosts within the accelerator islands, the system comprising: receiving data representing a machine learning workload; assigning a respective portion of the machine learning workload to each of the plurality of accelerator islands for scheduling by the respective scheduler of the accelerator island, comprising assigning the respective portion of the machine learning workload to each of the plurality of accelerator islands by sending a single message to the respective scheduler of the accelerator island when the respective portion of the machine learning workload is a regular compiled function with known resource requirements; and each scheduler configured to A system comprising:

2. 10. The system of claim 1, wherein the data representing the machine learning workload is data representing a sharded dataflow program including multiple shards.

3. 3. The system of claim 2, wherein assigning the respective portions 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. 2. The system of claim 1, wherein each scheduler is configured to schedule a portion of computation using parallel asynchronous dispatching when the respective portion of the machine learning workload assigned to the accelerator island is a regular compiled function with known resource requirements.

5. scheduling the portion of the computation using parallel asynchronous dispatching; 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 it receives as inputs and outputs of one or more respective other operations executed by another one of the hardware accelerators in the accelerator island; determining, for each of the set of hardware accelerators, a respective size of an output of each of the one or more other operations; transmitting, in parallel for each of the set of hardware accelerators to the corresponding host, respective future data specifying the respective sizes of the outputs of the one or more respective other operations; The system of claim 4, comprising:

6. 6. The system of claim 5, 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 the corresponding host of the hardware accelerator assigned to the one or more respective other operations that identifies the allocated memory.

7. 7. The system of claim 6, wherein the corresponding host of the hardware accelerator assigned to the one or more respective other operations is configured to cause the hardware accelerator assigned to the one or more respective other operations to send the output of the respective other operation to the allocated memory.

8. The system of claim 7 , wherein the output is transmitted over an accelerator interconnect network.

9. One or more computer-readable storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform the respective operations performed by the system described in any one of claims 1 to 8.

10. A method comprising respective operations performed by the system of any one of claims 1 to 8.

11. One or more computer-readable storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform respective operations performed by the scheduler of any one of claims 1 to 8.

12. A method comprising respective operations performed by the scheduler of any one of claims 1 to 8.

Images