Efficient Scaling with Pathways

• Efficient scaling with Pathways is a paradigm that decomposes complex computations into asynchronous, sharded dataflows for near-100% resource utilization.
• It orchestrates large-scale ML workloads, such as training PaLM, by employing gang-scheduling and efficient inter-device communication to achieve near-linear scaling.
• The approach extends to molecular simulations, high-dimensional path planning, and metabolic network modeling, demonstrating versatile applicability across scientific domains.

Efficient scaling with Pathways refers to a class of algorithmic, architectural, and systems innovations enabling the rapid expansion of computationally intensive, path-based models and simulations across large-scale distributed infrastructures. Recent advances span machine learning orchestration layers, molecular dynamics rare-event sampling frameworks, numerical PDE solvers for path planning, and formal evolutionary models of metabolic networks. Each of these fields has established distinct methodologies for scaling path-centric computations while ensuring computational, statistical, or biological efficiency.

1. Sharded Asynchronous Dataflow and Distributed ML: The Pathways System

The Pathways system is an orchestration layer designed for large-scale accelerator arrays, such as Google TPU pods, to enable efficient training and inference of modern machine learning workloads. Pathways adopts a sharded dataflow graph of operators, where each node consumes and produces futures corresponding to buffers on device, and schedules these as soon as their inputs are ready. The system architecture divides into a control plane (running the Plaque executor) and a data plane (which moves actual tensor bytes over dedicated interconnects).

Each ML program is represented as a directed acyclic graph $G = (V, E)$, with logical operators as nodes and data-dependency edges carrying futures. When all input-futures for a node $v$ are ready, $v$ is gang-scheduled on a TPU slice. Gang-scheduling is mandatory due to the non-preemptibility and single-streaming execution of TPUs; thus, all shards of an SPMD collective must advance in lockstep to avoid deadlock. Host-side scheduling is aggressively pipelined—resource allocation, buffer preparation, and scheduler RPCs for many downstream nodes are initiated as soon as possible, minimizing end-to-end scheduling overhead relative to compute time and achieving near-100% utilization for both SPMD and non-SPMD workloads.

Pathways efficiently coordinates data transfers using the high-bandwidth, low-latency interconnect (ICI) within pods and the datacenter network (DCN) across pods. Overlapping data transfer with on-device execution and leveraging futures for early scheduling reduces performance bottlenecks. Performance studies demonstrate that token throughput for transformer models is linear up to thousands of TPUs, pipelined models maintain $\leq5\%$ overhead, and scaling out to multiple islands (i.e., pod clusters) incurs less than 3% throughput loss, even under multi-terabyte gradient exchanges (Barham et al., 2022).

2. Scaling Transformer LLMs with Pathways: The PaLM Case

Pathways underlies the training strategy for PaLM, a 540B-parameter transformer LLM. The training leverages two levels of parallelism:

• Intra-pod parallelism exploits a 2D sharding of data and model (weight tensors partitioned across 12 chips, data batches over 256 replicas).
• Inter-pod orchestration: A single Python client asynchronously gang-schedules compute subgraphs on each pod and exchanges gradients via bulk DCN all-reduce.

Empirical scaling results show that throughput scales 1.95$\times$ (out of theoretical 2$\times$) when moving from 3,072 to 6,144 TPUs, with per-step model-FLOPs utilization (MFU) of 46.2%, substantially higher than previous systems (e.g., Gopher, Megatron-Turing, GPT-3). End-to-end training time for PaLM 540B was 2,791 TPU days for 780B tokens, with step throughput not bottlenecked by communication overhead (DCN usage). Observed scaling laws—loss vs. compute, few-shot accuracy vs. parameters, and task performance—reveal smooth log-linear improvements, but with substantial discontinuities (20–50%) in performance on ∼25% of benchmark tasks when scaling from 62B to 540B parameters. A plausible implication is that efficient scaling in the Pathways paradigm not only maximizes hardware utilization but also increases the likelihood of emergent behaviors and rapid capability gains in large models (Chowdhery et al., 2022).

3. Algorithmic Scaling in Molecular Path Sampling

Within biomolecular simulation, the metadynamics of paths (MoP) algorithm implements efficient scaling by decomposing trajectory space into discrete beads, each evolved under a discretized Onsager-Machlup action with friction and bias. Parallelism is achieved by mapping each bead to an independent MPI process/GPU—a pure MPMD architecture—that yields perfect load balance and strictly local communication, only extending to global evaluation of collective variables (CVs) as required by the biasing scheme.

Benchmarks on GROMACS show that such decomposition attains $\geq$70% weak-scaling parallel efficiency for unbiased path molecular dynamics, and 60–65% for MoP simulations up to 3,200 GPUs. Communication overhead (from bead-to-bead, and CV evaluation) stays under 30% even at massive scale, suggesting that the method's inherent one-dimensional path topology is ideally matched to distributed parallelization. Algorithmic optimizations—including vectorized kernels, nonblocking MPI communication, and GPU-offload—are essential to sustain scaling. This strategy generalizes to other path-sampling frameworks, provided communication is localized and global reductions are minimized (Malapally et al., 21 Jan 2025).

4. High-Dimensional Path Planning on Manifolds

Optimal path planning on manifolds, as formalized by a time-optimal control problem constrained to the graph of a smooth function, confronts challenges of dimensionality and differential geometry. Traditional methods reliant on grid-based discretizations of the corresponding Hamilton-Jacobi-Bellman (HJB) PDE become intractable at moderate $n$ due to the curse of dimensionality.

Recent work applies a generalized Hopf-Lax formula, discretizing time (not state), and reduces the HJB-value function to a finite-dimensional saddle-point problem. A primal-dual hybrid gradient (PDHG)-inspired splitting algorithm alternately updates state and co-state variables, leveraging closed-form proximal steps and momentum (inertial acceleration) for rapid convergence. The core innovation is that, since only time is discretized, per-iteration cost is $O(nJ)$: no state-space grid is needed, eliminating the exponential scaling bottleneck.

Empirical studies demonstrate sub-exponential runtime growth up to $n=30$, with solution times remaining on the order of a few seconds for all tested $v$0. The geometry of the manifold is incorporated through a local positive-definite matrix $v$1 (derived from the manifold's inducing function $v$2), with the Cholesky factorization critical for efficient dual updates. This approach provides efficient scaling in both runtime and memory, even as $v$3 increases, provided the cost of repeated Cholesky decompositions is controlled (Huynh et al., 2024).

5. Scaling Laws in Biological Networks: The Toolbox Model

The scaling of metabolic and regulatory networks is governed by evolutionary and topological constraints. The toolbox model posits that metabolic networks evolve by acquiring new pathways via horizontal transfer from a universal metabolic network. Each acquired pathway corresponds to a new transcriptional regulator. On a critical (mean upstream degree $v$4) rooted tree, the model predicts and analytically proves a quadratic scaling law: total number of regulators $v$5 scales as the square of the number of enzymes $v$6 ($v$7). This property—robust to network details and generalizable to multi-substrate, branched, and cyclic metabolic pathways—depends fundamentally on maintaining a “large-world” (critical) network topology where the distribution of pathway lengths is broad.

Numerical simulations on KEGG-based universal networks confirm approximately quadratic scaling, with minimal pathways typically spanning 30–40 reaction layers and exhibiting high yield ($v$897% with zero byproducts). This scaling is destroyed if network topology becomes “small-world,” as in naive graph representations that ignore biochemical constraints, reducing to (at best) linear scaling with logarithmic corrections. Thus, efficient scaling in real metabolic networks is both an evolutionary and a topological phenomenon (Pang et al., 2010).

6. Design Principles and Cross-Domain Insights

Comparison of scaling strategies across domains highlights convergent themes:

Domain	Key Principle	Scaling Mechanism
ML Systems	Asynchronous dispatch	Sharded dataflow, futures
Molecular Sim.	Path decomposition	MPMD, local comms
Path Planning	Time discretization	PDHG/Hopf-Lax, no grids
Metabolic Evo.	Critical topology	Quadratic regulatory law

Efficient scaling with Pathways, as the unifying concept, emphasizes:

• Structural decomposition (across devices, trajectories, or graphs),
• Local communication and coordination (scheduling or data exchange),
• Minimal global synchronization (asynchronous execution where possible),
• Numerical formulations that avoid the curse of dimensionality (no spatial grids, only time or path discretization),
• Exploitation of domain-specific topological or evolutionary constraints to maintain or enhance statistical or biological scaling laws.

This synthesis suggests that robust, efficient scaling combines: architectural abstraction (shards, futures), algorithmic design (splitting, saddle-point, decomposition), and careful mapping to the topology of the computational/network domain. The diversity of applications—ranging from exascale biomolecular path sampling, half-trillion parameter LLMs, high-dimensional path planning, to prokaryotic metabolic networks—underscores the versatility and foundational nature of these scaling paradigms across scientific computing and biological modeling (Barham et al., 2022, Chowdhery et al., 2022, Malapally et al., 21 Jan 2025, Huynh et al., 2024, Pang et al., 2010).