Source-linked AI summary
Pathways: Asynchronous Distributed Dataflow for ML
Paul Barham, Aakanksha Chowdhery, Jeff Dean, Sanjay Ghemawat, Steven Hand, Dan Hurt, Michael Isard, Hyeontaek Lim, Ruoming Pang, Sudip Roy, Brennan Saeta, Parker Schuh, Ryan Sepassi, Laurent El Shafey, Chandramohan A. Thekkath, Yonghui Wu
TL;DR
Current distributed ML systems struggle to support large, sparse, irregular, and heterogeneous workloads while maintaining performance. PATHWAYS addresses this with a single-controller, sharded asynchronous dataflow runtime that coordinates distributed computations and resource management. It matches state-of-the-art multi-controller performance while supporting richer computation patterns and multi-tenant resource sharing, with current resource-management policies and data-dependent sparsity remaining scope boundaries.
Problem
Existing distributed ML systems make large, sparse, irregular, and heterogeneous models difficult to support, while single-controller systems add dispatch and scheduling challenges at scale.
Method
PATHWAYS uses a single-controller, sharded dataflow runtime with asynchronous dispatch, centralized gang-scheduling, resource management, and sparse data exchange.
Results
PATHWAYS matches JAX performance across very large scales for all but the smallest computations while supporting richer computation patterns and multi-tenant sharing.
Takeaways & Limitations
PATHWAYS provides a basis for running current ML models efficiently while exploring future parallelism patterns and sharing accelerator resources.
Takeaways & Limitations
PATHWAYS currently focuses on dynamic time-multiplexing of TPU computations; broader resource types and policies remain planned work.
Abstract
from arXiv · showhide
We present the design of a new large scale orchestration layer for accelerators. Our system, Pathways, is explicitly designed to enable exploration of new systems and ML research ideas, while retaining state of the art performance for current models. Pathways uses a sharded dataflow graph of asynchronous operators that consume and produce futures, and efficiently gang-schedules heterogeneous parallel computations on thousands of accelerators while coordinating data transfers over their dedicated interconnects. Pathways makes use of a novel asynchronous distributed dataflow design that lets the control plane execute in parallel despite dependencies in the data plane. This design, with careful engineering, allows Pathways to adopt a single-controller model that makes it easier to express complex new parallelism patterns. We demonstrate that Pathways can achieve performance parity (~100% accelerator utilization) with state-of-the-art systems when running SPMD computations over 2048 TPUs, while also delivering throughput comparable to the SPMD case for Transformer models that are pipelined across 16 stages, or sharded across two islands of accelerators connected over a data center network.
1 INTRODUCTION
PATHWAYS targets emerging ML workloads that strain specialized SPMD systems, while preserving state-of-the-art functionality and performance for current models.
- SPMD systems run identical computations in lockstep, but pipelining and computational sparsity increasingly require fine-grain control flow and heterogeneous computation.
- Heterogeneous clusters and expensive large accelerator islands motivate MPMD computation across smaller, more available accelerator islands.
- Fine-grained hardware multiplexing can provide workload elasticity, improve fault tolerance, and increase utilization.
- Foundation-model training and inference can multiplex resources across tasks, share state, and combine examples in vectorized batches.
- PATHWAYS uses a client-server runtime, a scalable dataflow model, non-SPMD programming, and centralized resource management to support future workloads.
2 DESIGN MOTIVATION
Existing distributed ML architectures trade dispatch performance against flexibility and centralized coordination. PATHWAYS combines single-controller flexibility with multi-controller performance through asynchronous dispatch, scheduling, and sharded dataflow.
- Multi-controller systems dispatch accelerator computations with low latency over PCIe but poorly support pipelining, sparsity, and communication beyond collectives.
- Single-controller systems incur dispatch communication over DCN, typically an order of magnitude slower than PCIe.
- Gang-scheduling is essential for TPU computations because non-preemptible, single-threaded kernels can deadlock when communicating computations are enqueued inconsistently.
- TensorFlow v1 delays destination dispatch until cross-host transfers complete and materializes full graphs, creating overhead for thousands of shards.
- PATHWAYS uses asynchronous dispatch, centralized gang-scheduling, resource management, and sharded dataflow to combine flexibility with performance.
3 PATHWAYS PROGRAMMING MODEL
PATHWAYS supports compiled, sharded computations from JAX and TensorFlow, with virtual-device placement, automatic data movement, and traced dataflow programs for composing functions.
- PATHWAYS supports TensorFlow and JAX source programs, representing each compiled function as a single sharded computation node.
- PATHWAYS extends JAX across multiple TPU pods by communicating over both ICI and DCN, reaching many thousands of TPU cores.
- Users can request virtual devices with constraints on device type, location, or topology, while PATHWAYS handles data movement and resharding.
- A program tracer combines multiple compiled functions into one PATHWAYS dataflow graph, avoiding separate client calls and RPCs for each function.
- JAX transformations support pipelined models and exploration of data-dependent vectorized control flow.
4 PATHWAYS SYSTEM ARCHITECTURE
PATHWAYS combines a client-server architecture with virtual-device resource management, sharded dataflow coordination, and asynchronous dispatch for distributed accelerator programs.
- PATHWAYS builds on XLA, TensorFlow graphs and executors, JAX, and TensorFlow APIs to run existing ML models with minimal code changes.
- The resource manager allocates virtual slices with requested mesh shapes and dynamically maps virtual devices to physical accelerators.Mappings can satisfy interconnect topology and memory-capacity requirements, while the initial implementation uses a heuristic to spread computations across available devices.
- PATHWAYS lowers a device-location-agnostic IR into programs containing physical locations and transfers between source and destination computation shards.The program can be re-lowered when virtual-to-physical mappings change, while repeated execution is efficient when mappings remain stable.
- Cross-host coordination uses PLAQUE dataflow graphs with compact sharded-computation nodes and sparse exchanges between dynamically selected shards.Sparse communication helps prevent the data-center network from bottlenecking data-dependent accelerator control flow.
- When host-side scheduling and coordination exceed short accelerator execution times, sequential dispatch stalls the asynchronous pipeline.PATHWAYS instead schedules regular compiled subgraphs in parallel using statically known resource usage, while retaining traditional scheduling for data-dependent resource requirements.
- PATHWAYS stores intermediate values in object stores addressed by opaque handles, enabling migration, delayed transfers, garbage collection, and memory back-pressure.Ownership labels support cleanup after program or client failure, and back-pressure can stall computation when HBM allocation is unavailable.
5 EVALUATION
PATHWAYS matches or exceeds JAX performance across dispatch micro-benchmarks, concurrent programs, and large Transformer workloads. It also maintains high throughput for pipelined and cross-island TPU training.
- Single-controller dispatch overheads: PATHWAYS matches JAX performance up to 1000 TPU cores, while its Chained mode outperforms JAX OpByOp up to 256 cores.PATHWAYS executes back-to-back accelerator computations directly from C++, whereas JAX OpByOp transitions to Python for every computation.
- Single-controller dispatch overheads: 2.3 ms reaches JAX throughput on 16 hosts with 128 TPUs, while 35 ms masks PATHWAYS’s overhead on 512 hosts with 2048 TPUs.These thresholds identify computation sizes sufficient to amortize single-controller overhead.
- Multi-tenancy: PATHWAYS achieves at least the same aggregated throughput as JAX when multiple clients concurrently submit programs whose resources fit in HBM.The experiment reports no context-switch overhead between concurrent programs from different clients under this condition.
- Large scale model performance: 133.7k tokens/sec falls to 131.4k tokens/sec as pipeline stages increase from 4 to 16, while throughput scales with TPU cores per stage.PATHWAYS’s throughput also scales linearly with the number of hosts.
- Large scale model performance: 131.4k tokens/sec is achieved on either one island of 128 cores or four islands of 32 cores connected via DCN.Communication time is effectively overlapped with computation in the partitioned configuration.
- Large scale model performance: PATHWAYS achieves ∼97% of single-island throughput when training 64B and 136B Decoder-only Transformers across two accelerator islands.The models use two islands connected over DCN, with ICI reductions within each island followed by cross-island DCN transfers.
6 DISCUSSION
PATHWAYS targets TPU-oriented orchestration while extending beyond current workload patterns and resource-management capabilities. The discussion identifies data-dependent vectorized control flow and broader multi-tenancy support as future work.
- System scope: PATHWAYS’s low-level design is shaped by TPU kernels, which can fuse longer-running computations with control flow and communication primitives.XLA’s support for fused collectives and TPU interconnects makes TPUs a good fit for PATHWAYS.
- System scope: 131.4k tokens/sec is maintained when a 3B Transformer is pipelined across four TPU islands instead of one island.The comparison uses four islands of 32 cores each versus one island of 128 cores.
- Resource management: PATHWAYS plans to support priorities, performance isolation, access control, and resource accounting across larger pools of heterogeneous resources.The initial research focuses on dynamic time-multiplexing of TPU computations; broader resource types remain to be handled.
- Data-dependent vectorized control flow: Data-dependent vectorized control flow remains future work for supporting fine-grain routing and sparsity with a clean programming model and good performance.The motivating examples route examples or sub-examples to accelerators hosting different subsets of model weights.
7 RELATED WORK
Related work covers finer-grained accelerator sharing and computation patterns that depart from SPMD. These lines of work motivate PATHWAYS’s focus on heterogeneous execution and resource efficiency.
- Beyond SPMD: The paper positions its related-work discussion as addressing ML workloads needing capabilities beyond SPMD multi-controllers.The section also uses those prior approaches to validate PATHWAYS’s design choices.
- Accelerator sharing: Conventional resource sharing is coarse-grained, with cloud providers commonly dedicating accelerators to individual users.The passage contrasts this model with newer approaches to finer-grained accelerator sharing.
- Accelerator sharing: Virtualizing accelerators, GPU memory, or DRAM can improve resource efficiency or address memory limits for large models.The cited work includes accelerator virtualization, GPU memory virtualization, and DRAM offload.
- Beyond SPMD: Pipelining, graph neural network training, neural architecture search, and multimodal multitask learning demonstrate workloads that deviate from SPMD.Pipelining partitions models into static heterogeneous computations across accelerators.
8 CONCLUSIONS
PATHWAYS matches state-of-the-art multi-controller performance on current single-tenant SPMD models while enabling richer computation patterns through a single-controller model and centralized resource management.
- The single-controller model provides access to richer computation patterns and centralized resource management between clients and accelerators.
- Careful system design and engineering let PATHWAYS match performance on current ML models while providing features for future models.
A ACCELERATOR DESIGN
Modern deep learning depends critically on hardware accelerators, but achieving high accelerator performance requires non-trivial systems engineering and established optimization techniques.
- Hardware acceleration is critical to modern deep learning.
- Achieving high performance with accelerators is a non-trivial systems exercise.
- Deep learning systems commonly employ established techniques to achieve good accelerator performance.
A.1 Batching
Accelerator systems use asynchronous execution, batching, sharding, and compiled functions to expose parallelism, coordinate communication, and sustain utilization under hardware constraints.
- Asynchronous stream execution masks dispatch latency for small operations when a sufficiently large work pipeline is maintained.
- Neural networks larger than accelerator HBM are sharded across multiple accelerators, making high-speed interconnects critical for performance.
- Dedicated GPU and TPU interconnects expose accelerator communication through MPI primitives such as AllReduce.
- Fine-grained context switching is costly because PCIe bandwidth is much lower than HBM and accelerator-interconnect bandwidth.
- TPUs execute long-running XLA functions with collective communication across their dedicated ICI network and no host interaction.
- Because TPUs lack local pre-emption, PATHWAYS must centrally gang-schedule communicating computations in a consistent order to avoid deadlock.
- Compiled functions have known input and output shapes, bounded loops, and functional conditionals with resources allocated in advance.
- Frameworks enqueue compiled functions asynchronously before predecessors finish, allowing host work to overlap accelerator computation.
C INPUT DATA PROCESSING
PATHWAYS distributes CPU-based input processing across workers while separating CPU resources from TPU-connected hosts; its traces also show concurrent scheduling and relatively small cross-island transfer overhead.
- PATHWAYS runs a CPU-based TensorFlow executor on each host to distribute serialized input-processing graphs across workers.
- A sample TPU trace shows interleaving of gang-scheduled concurrent programs.
- A trace of 64B Transformer training across two 512-TPU islands highlights the relatively small overhead of cross-island transfer over DCN.
- Input processing can run on an independently managed server set, decoupling TPU-connected hosts from available CPU resources.
D EVALUATION WORKLOAD TRACES
Pathways traces show that concurrent clients can saturate otherwise underutilized accelerators, while cross-island Transformer training over DCN retains near-SPMD throughput.
- Multi-tenancy: ~100% device utilization is reached when multiple clients interleave short programs that individually cannot saturate accelerators.Each program has only 0.33 ms of compute time; gang scheduling interleaves programs at millisecond scale or less with little context-switch overhead.
- Cross-island training: 97.2% training throughput is achieved for a 64B Decoder-only Transformer trained across two 512-chip islands using DCN gradient transfers.Each island computes gradients, transfers them to the other island, applies received gradients, and begins the next training step.