Source: OneFlow
This article is approximately 7732 words long and is recommended to be read in 12 minutes. This article introduces the background of Pathways and provides an in-depth analysis.
01. Why Discuss Pathways?
In the past two years, TensorFlow has been caught off guard by the rise of PyTorch, and the entire industry is looking forward to Jeff Dean turning the tide with a powerful weapon.
Last year, Jeff Dean published a highly anticipated blog post (https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/). On one hand, the author is Jeff Dean, and on the other hand, the title of the blog, “Pathways: The Next Generation AI Architecture,” is sensational. After all, Google and Jeff Dean have led the development of infrastructure in the internet and AI eras. When someone else talks about “next-generation architecture,” it may be boastful, but when Jeff Dean says it, you have to take it seriously. This blog lists Jeff Dean’s judgments on the future evolution of deep learning algorithms: multi-modal, sparse, dynamic routing, etc. These characteristics will indeed impact the underlying architecture of deep learning.
A few days ago, Google first uploaded a paper titled “Pathways: Asynchronous Distributed Dataflow for ML” (https://arxiv.org/abs/2203.12533) on arXiv, introducing the design philosophy of the long-awaited Pathways system. Shortly after, Google also uploaded a large model trained based on the Pathways system, “PaLM: Scaling Language Modeling with Pathways” (https://arxiv.org/abs/2204.02311), which has 540 billion parameters.
To be honest, the Pathways paper is not easy to understand. I believe there are several reasons for this:First, there are very few practitioners worldwide engaged in research and development of distributed deep learning systems. Even among this group, only a few have ever thought about the issues discussed in the paper. Without understanding the WHY, it is difficult to grasp the HOW;Second, this paper is poorly written. I feel that if a paper is easily accepted by a team, its quality cannot be guaranteed. Jeff Dean’s paper falls into this category. I believe that the TensorFlow paper (https://www.usenix.org/system/files/conference/osdi16/osdi16-abadi.pdf) was also quite average, yet it managed to be accepted at OSDI. A few years ago, I gained the confidence to create a better deep learning framework after reading the TensorFlow paper.
I might be one of the few people who can understand Pathways. The issues discussed in Pathways are ones we have considered for years, and we developed OneFlow and even published papers discussing these issues. Unfortunately, when we discussed these issues in our papers, we were rarely understood. Today, when Google discusses these principles in their paper, it becomes an unprovable truth.
As mentioned earlier, this paper is not very easy to understand. I will explain it in hopes of clarifying the confusing aspects of the paper. Especially when we understand the WHY, we may find that the design and implementation of Pathways are still lacking. If we can take it a step further, it happens to be the approach that OneFlow has already implemented. Therefore, I jokingly said in my circle of friends that I want to write a blog titled “Pathways: One Step Forward is OneFlow” or “The End of Pathways is OneFlow.” I plan to discuss Pathways in two articles,the first mainly discussing the background, which is the design motivation of Pathways,and the second discussing the design and implementation of Pathways.
The paper spends a lot of ink discussing the issues of single-controller and multi-controller. Those who have never thought about this issue may find it a bit perplexing: “What do SPMD, MPMD, single-controller (client), and multi-controller (client) mean? Is it really that important? This topic seems quite boring.” Let’s clarify this issue in this blog. The paper first discusses the later-emerging multi-controller and then discusses the earlier-emerging single-controller, like a “flashback.” We will still introduce it according to the timeline.
02. TF v1 and Single-controller
First, let’s introduce some basic concepts in TF v1. Users describe the computational graph (including the structure of the neural network and the available resource information) through a Python client (client), and then this computational graph is handed over to the runtime system through the well-known session.run. The runtime system is divided into a master and several workers, where the master compiles and slices the computational graph and sends each sliced subgraph to the corresponding worker for execution. Here, the client and master grasp the global information, while the worker is only used to execute the assigned subgraph. Therefore, the client and master can be understood as a controller. This structure is referred to as single-client in Google’s previously published papers (such as GShard and GSPMD), but in the Pathways paper, it is called single-controller; the two terms are equivalent.
The above figure shows the mechanism of distributed computing in TF v1. Note that the figure does not distinguish between client and master, collectively referred to as the controller. Each worker consists of a Host represented by CPU and a Device represented by GPU. The arrows in the figure indicate message sending, with downward arrows representing the controller driving the worker (cross-node communication is generally done via Ethernet), while within the worker, the Host drives the Device (intra-node communication is generally done via PCIe). The upward arrows represent the messages sent by the worker to report progress to the controller. Each message sending has a certain latency; the cross symbol in the blue box indicates that cluster communication occurs between two devices (usually occurring in dedicated transmission devices like NVLink or RDMA). The controller sends each subgraph corresponding to each step of the loop to the worker for execution. Once all workers finish executing their respective subgraphs, they inform the controller, which then issues the next step’s computation task.
It is evident that in the TF v1 scheme shown in the figure, due to the latency in message sending, although two adjacent steps on the controller are tightly arranged, there are large idle periods on the devices of both workers, leading to unsaturated device utilization. This is an obvious problem, and it is puzzling why TensorFlow was designed this way from the beginning.
03. Horovod and Multi-controller
The drawbacks of the central scheduler in TF v1 have existed since its inception. In earlier years, even the simplest data parallelism led to TF v1’s performance being “beaten” by other frameworks. Moreover, for a long time, data parallelism was the only requirement for distributed deep learning (so-called model parallelism and pipeline parallelism have only recently become mainstream needs), leading to various optimizations for TF v1’s distributed performance.
The most typical example is Horovod, which was introduced by Uber. It completely abandons the distributed functionality built into various deep learning frameworks and introduces a new plugin that only performs data parallelism. This plugin, when combined with the single-card functionality of frameworks like TF, allows for fast and efficient distributed training, becoming a standard for distributed deep learning. Horovod’s startup and scheduling methods are almost identical to traditional supercomputing MPI programs, where each card starts a process, and all processes run the same code. This is a typical SPMD, where each process is a client or entry point for distributed training, thus also referred to as multi-client or multi-controller. Therefore, the SPMD that frameworks like PyTorch/JAX later adopted actually originated from Horovod.
Let’s see how SPMD addresses the drawbacks of the central scheduler in TF v1. In the multi-controller architecture, each Host executes completely identical programs, only differing in the parameters (rank) passed to each program, indicating that each program processes different slices of data.
Note that there is no central scheduler here; each Host simply sends computation tasks to the Device queues it manages in the same order. Multiple Devices initiate cluster communication at the same time, and once the cluster communication is completed, they can continue executing the computation tasks in their respective queues. It can be seen that the execution queues of each device are filled with computation tasks, avoiding the numerous idle time slots seen in TF v1. In this architecture, the cluster communication between devices is completed through special network connections, and the messages between Host and Device are completed via PCIe, avoiding the high latency caused by the communication between the central controller and worker through Ethernet in TF v1.
Horovod’s “regression” prompts reflection: Has TensorFlow over-designed its framework? In the subsequent versions of TensorFlow and in other next-generation frameworks like PyTorch and JAX, all have embraced the multi-controller architecture.
04. Asymmetric Communication Requirements and MPMD
If distributed deep learning were solely based on symmetrical and regular communication like data parallelism, the sky for distributed deep learning would always be sunny.
In fact, even with model parallelism, communication remains symmetrical and regular, using all-gather, reduce-scatter, and other cluster communications. When pipeline parallelism emerges, the programs executed on processes at different stages of the pipeline become different, which is known as MPMD. With the emergence of dynamic and sparse neural network structures like MoE, such asymmetric and irregular situations are becoming more frequent.
In a previous blog post titled “Combatting Software Complexity: Proper Layering, Not Too Much, Not Too Little,” I discussed the communication patterns in distributed deep learning, which include regular communication patterns represented by cluster communication and irregular communication patterns represented by point-to-point communication.
Thus, the Google Pathways paper states: No, we must return from multi-controller to single-controller.
As shown in the figure above, suppose the workers numbered from top to bottom are 1, 2, and 3. Clearly, their workloads are not entirely the same and are asymmetric. First, an intermediate result produced by 2 needs to be sent to 3’s Recv through Send, and a certain intermediate result from 3 needs to be sent to 1’s Recv through Send. Send and Recv must appear in pairs and have order requirements: Send strictly happens before Recv.
05. Deadlock Risks
Pathways proposes that asymmetric loads (non-SPMD) require centralized scheduling and gang scheduling mechanisms to avoid deadlocks. If one is not a framework developer or has not personally dealt with the challenges posed by such asymmetric and irregular situations, it can be hard to understand. There is a very key statement in the paper regarding this issue:
Gang-scheduling is essential in the case of TPUs, since they are single-threaded and only run non-preemptible kernels, so the system will deadlock if communicating computations are not enqueued in a consistent order. (In the case of TPUs, gang scheduling is essential because the kernels on TPUs are single-threaded and run non-preemptively; if multiple TPUs’ communication operations are not enqueued in the same order, the system will deadlock.)
How to UnderstandThis Statement?
One device is often abstracted as a FIFO task queue, such as the concept of CUDA streams on NVIDIA GPGPU. The kernels in the queue are processed in a FIFO manner. Other hardware resources, including CPU and network, can also be abstracted as task queues, which we collectively refer to as streams.
The term single-threaded means that a TPU has only one queue (similar to allocating only one CUDA stream on a GPU); the term non-preemptive means that once a kernel is launched, it must wait for its execution to finish before the resources it occupies can be reclaimed for use by other kernels. In other words, a launched kernel cannot pause and yield resources to other kernels. To put it more simply, this queue does not allow for “cutting in line.”
Under the conditions of single-threaded and non-preemptive scheduling, if multiple cluster communications are to be executed on two devices, the corresponding kernels for these cluster communications must be enqueued in a consistent order on both devices; otherwise, a deadlock will occur.
As shown in the figure above, there are two cluster communications A and B. If the order of A and B is inconsistent on the two devices, they will get stuck, and neither device will be able to proceed. The upper device first launches the B kernel; it must wait for the corresponding B kernel on the lower device to be launched to complete, but the lower device first launched the A kernel, which also needs to wait for the A kernel on the upper device to complete before it can launch, leading to a deadlock.
The above figure illustrates another scenario that can lead to deadlock (the classic dining philosophers problem). Device 1 and 2 need to perform cluster communication A, while device 2 and device 3 need to perform cluster communication B, and device 3 needs to perform cluster communication C with device 1. If the kernels are launched in the order shown in the figure, a deadlock will also occur.
06. How to Resolve Deadlocks?
Clearly, if a launch order is specified among the queues of each device, deadlocks can be avoided. This essentially requires the use of gang scheduling methods.
Another method is to use preemptive scheduling, which allows kernels to execute out of order (in other words, allows for “cutting in line”), thus avoiding deadlocks.
Another approach is to increase concurrency, allowing each device to create multiple FIFO queues. The Pathways paper emphasizes the necessity of gang scheduling in the TPU scenario due to deadlock issues, but in the GPU scenario, it states:
Even for GPUs or other accelerators that can execute concurrent computations, gang scheduling allows more efficient execution of collectives (Feitelson and Rudolph, 1992).
In other words, the paper believes that if devices support “concurrent” execution, deadlocks may no longer be a problem, thus emphasizing the necessity of gang scheduling from an execution efficiency perspective. How does increasing concurrency help reduce the likelihood of deadlocks?
As shown in the figure below, by allocating several more queues on each device and scheduling A and B to different queues, deadlocks can also be avoided.
However, increasing the number of queues does not fundamentally solve the problem. The number of queues is always limited, and when the number of concurrently executed cluster communications exceeds the number of queues on a device, there is still a risk of deadlock. This issue is real in the GPU scenario and occurs frequently. The deadlock arises because the cluster communications on GPUs all use NCCL developed by NVIDIA, which has a design “flaw”: when the number of cluster communications exceeds the maximum concurrency, deadlocks may occur. This issue is quite interesting, and I believe it deserves some introduction.
07. NCCL’s Deadlock Trap
NCCL’s documentation specifically discusses precautions when executing multiple cluster communications concurrently:
(https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently):
Using multiple NCCL communicators requires careful synchronization, or it can lead to deadlocks. NCCL kernels are blocking (waiting for data to arrive), and any CUDA operation can cause a device synchronization, meaning it will wait for all NCCL kernels to complete. This can quickly lead to deadlocks since NCCL operations perform CUDA calls themselves. Operations on different communicators should therefore be used at different epochs with a locking mechanism, and applications should ensure operations are submitted in the same order across ranks. Launching multiple communication operations (on different streams) might work provided they can fit within the GPU, but could break at any time if NCCL were to use more CUDA blocks per operation, or if some calls used inside NCCL collectives were to perform a device synchronization (e.g., allocate some CUDA memory dynamically).
The above describes the risk of deadlock when the number of concurrently occurring cluster communications exceeds a certain limit. How to Understand this issue?
When a cluster communication operation is started on a GPU, NCCL starts a certain number of threads on the GPU that busy wait for data to arrive. If the data does not arrive, these threads will not exit. In other words, these cluster communication kernels are non-preemptive.
These threads occupy a certain number of GPU cores; for example, 64 cores. Since the total number of cores on a GPU is limited, say 4096 cores, the maximum number of cluster communications that can be started on this GPU is 4096/64 = 64. If the number of concurrently executed cluster communications exceeds 64, there is a risk of deadlock.
08. Summary
To avoid deadlocks caused by irregular communication, the prescription proposed by Pathways is gang scheduling combined with centralized scheduling. Gang scheduling is used within the same group of devices (island of accelerators), while centralized scheduling coordinates the launching of cluster communication operations in a consistent order across different groups of devices.
Does using centralized scheduling solve all possible deadlocks? Not necessarily; centralized dynamic scheduling may still face deadlocks caused by resource dependencies. I discussed a thought experiment in a previous blog post titled “The Curse of Resource Dependency: The Flaws of Existing Deep Learning Frameworks,” and we also discussed it in our paper “OneFlow: Redesign the Distributed Deep Learning Framework from Scratch” (https://arxiv.org/pdf/2110.15032.pdf).
Is centralized scheduling necessary to resolve deadlocks? Not necessarily; if a neural network is a static graph, using control edges to arrange the execution order during static compilation can also prevent deadlocks. In fact, static compilation is equivalent to a form of ahead-of-time centralized scheduler. Once the compiler arranges the execution order, whether there is a centralized scheduler during the runtime phase becomes less important.
Finally, it is worth mentioning that SPMD is suitable for multi-controller, and MPMD is suitable for single-controller,but this is not a one-to-one binding relationship. Using a single-controller can support SPMD (for example, TF v1 supports both data parallelism and model parallelism), and using a multi-controller can also support MPMD (for example, the GSPMD paper supports pipeline parallelism, and the OneFlow multi-client architecture supports pipeline parallelism).
Therefore, arguing for the necessity of using single-controller and centralized scheduling from the perspective of avoiding deadlocks is somewhat forced. Of course, there are other reasons and evidence supporting the necessity of centralized scheduling design. We will discuss this in the next article.
——END——
