DrJAX: Scalable and Differentiable MapReduce Primitives in JAX

Authors: keith  rush
zachary  charles
zachary  garrett
sean  augenstein
nicole  mitchell
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-06-01 01:07:48

Abstract: We present DrJAX, a JAX-based library designed to support large-scale distributed and parallel machine learning algorithms that use MapReduce-style operations. DrJAX leverages JAX's sharding mechanisms to enable native targeting of TPUs and state-of-the-art JAX runtimes, including Pathways. DrJAX embeds building blocks for MapReduce computations as primitives in JAX. This enables three key benefits. First, DrJAX computations can be translated directly to XLA HLO, enabling flexible integration with a wide array of ML training platforms. Second, DrJAX computations are fully differentiable. Last, DrJAX computations can be interpreted out to existing batch-processing compute systems, including traditional MapReduce systems like Apache Beam and cross-device compute systems like those powering federated learning applications. We show that DrJAX provides an easily programmable, performant, and scalable framework for parallelized algorithm development.

Main paper: 
 introduction: 
 The ability to scale abstractly written compute-intensive programs across large distributed compute environments is a key factor in the success of modern machine learning (ML). This is crucial for ML computations involving large language models, which are often too large to fit on a single compute node. Another key facet of modern ML software is the general ease with which computations can be written and optimized. Techniques such as automatic differentiation (AD) and just-in-time (JIT) compilation have enabled frameworks such as PyTorch (Paszke et al., 2019) , Tensor-Flow (Abadi et al., 2016) , and JAX (Bradbury et al., 2018) to scale to larger and more complex ML workloads.These software frameworks generally focus on enabling parallelism for their most common use case: computation of a function's derivative across a batch of inputs. This computation is typically parallelized in two well-defined manners: across the batch dimension (i.e. data parallelism), and within the computation of the derivative at a single example (i.e. model parallelism). Data parallelism is extremely common across ML frameworks. Model parallelism is more complex, but has seen an explosion of progress in recent years (Gholami et al., 2018; Shazeer et al., 2018; Jia et al., 2019; Lepikhin et al., 2020; Xu et al., 2021) which has enabled the training and deployment of significantly larger models than previously possible (e.g. foundation models).However, many sub-fields of ML do not fit neatly into this description, and instead employ parallelism over higherlevel partitioned structures of data: in meta-learning (where data is partitioned across tasks) (Finn et al., 2017) ; in grouplevel differential privacy (where data is partitioned over discrete groups whose individual contributions to algorithm outputs are information-theoretically bounded) (Dwork, 2010) ; in model merging (Li et al., 2022) or "model soup" algorithms (where data, in the form of hyperparameters, is partitioned across model copies) (Li et al., 2022; Wortsman et al., 2022) ; in federated learning (where data is partitioned across clients who avoid directly sharing data) (McMahan et al., 2017) ; and in optimization with intermittent communication (where data is partitioned across model replicas) (Mangasarian and Solodov, 1993; Zinkevich et al., 2010; Zhang et al., 2016) . The algorithms studied in these nominally different areas share many commonalities. An especially common factor is the use of the MapReduce (Dean and Ghemawat, 2004) programming paradigm, mapping parallel model training steps over partitioned data before invoking a reduction function. Data parallelism is a special case: we simply map and reduce over batches of data. In other applications, MapReduce is applied to coarser groups of data (e.g. multiple batches), which may additionally encode semantic structure of the underlying data.While ML frameworks provide scalability, flexibility, and efficiency for data-and model-parallelism, an algorithm author who wishes to program over partitioned data in a parallel manner finds themselves in an awkward position. For example, frameworks for federated learning (e.g. Ingerman and Ostrowski (2019); Ziller et al. (2021); Ro et al. (2021) ; He et al. (2020) ) offer parallelism over clients, but either do not support other ML use cases discussed above, or do not focus on large-scale datacenter performance. Underlying ML frameworks often offer powerful parallelism primitives (e.g. jax.shard map), but these often assume (as is the case for jax.shard map) that the specification of all model shardings and physical resources is known to the code author (Shazeer et al., 2018) . By contrast, in DrJAX we wish to abstract out MapReduce-style computations, allowing them to be defined in terms of model forwards and backwards passes (for example) that are already sharded. Finally, we note that bridging research and production often requires translating computations (e.g. from JAX) to production platforms.An ideal authoring surface for ML algorithms using MapReduce operations provides several features simultaneously: performant and scalable datacenter performance; the ability to decouple logical partitioning of data (number of groups of data to parallelize over) from physical compute (number of compute nodes); easy and extensible algorithm expression; JIT compilation and AD; and the capacity to translate algorithms to alternative infrastructure.Contributions. We present DrJAX (Differentiable MapReduce JAX) a software library that brings the benefits of modern large-scale machine learning software -sharding, easy-to-use JIT compilation, and AD -to MapReduce-style algorithms operating on partitioned data. DrJAX embeds a simple mapping and reduction programming model, by decomposing computations into differentiable building blocks (Section 2). We fully implement this programming model in JAX by embedding these building blocks via JAX's Primitive mechanism (Section 3). This allows DrJAX to use powerful features like sharding and JIT compilation to XLA. For example, DrJAX can shard computations over data partitions, model, and withindata partitions simultaneously across physical and logical meshes of devices. Because DrJAX is essentially a careful programming of parallel algorithms in XLA, DrJAX can leverage advances in distributed datacenter training like GSPMD (Xu et al., 2021) and Pathways (Barham et al., 2022) . We showcase the scaling and runtime benefits of DrJAX across a suite of large language model training experiments (Section 4). JAX's Primitive mechanism also enables forward-and reverse-mode differentiation which DrJAX leverages to provide full differentiability of its MapReduce-style programs. By implementing the AD framework of Rush et al. (2023) , we ensure that the derivative of a DrJAX program is simply another DrJAX program. This allows DrJAX and its AD system to be interpreted out to other platforms for parallel machine learning, including systems with strong guarantees about data locality and privacy (Section 5).Related work. DrJAX draws inspiration from three main areas of software. First, DrJAX directly utilizes the programming model of MapReduce (Dean and Ghemawat, 2004) and its evolution in Apache Beam (Foundation), highlighting the power of focusing on replications, mappings and reductions on parallelized collections of data.Second, much of DrJAX's treatment of machine learning and automatic differentiation is influenced by modern ML frameworks such as PyTorch (Paszke et al., 2019) and Ten-sorFlow (Abadi et al., 2016 ). DrJAX's design is directly based on the functional-first nature of JAX (Bradbury et al., 2018) . Several libraries implemented in JAX leverage similar (though not identical) mechanisms for ensuring scalability, notably the Praxis library for neural network layers (Google).Finally, DrJAX is inspired by frameworks for federated learning. Without intending to be exhaustive, examples of such frameworks include: PySyft (Ziller et al., 2021 ), FedJAX (Ro et al., 2021) , FedScale (Lai et al., 2022) , FedML (He et al., 2020 ), Flower (Beutel et al., 2020) , FLUTE (Dimitriadis et al., 2022) , FL Pytorch (Burlachenko et al., 2021) , and FATE (Liu et al., 2021) . In particular, DrJAX's programming model was significantly influenced by TensorFlow Federated (Ingerman and Ostrowski, 2019) and its federated computations (Charles et al., 2022) . Moreover, DrJAX uses the AD framework for federated computations proposed by Rush et al. (2023) to extend AD to MapReduce computations more broadly. 
system design: 
 DrJAX is designed with two key ideas in mind. First, the types of parallel computing for machine learning discussed above can all be viewed as applications of MapReduce-style computations that use functions like model forward and backwards passes as black-box subroutines. Second, we can differentiate through MapReduce computations by using the AD techniques proposed by Rush et al. (2023) . While their framework, federated AD, was motivated by federated learning, it can be directly adapted into a mechanism to apply AD to MapReduce computations. By combining standard AD techniques (e.g. computation graphs (Bauer, 1974) ) with proper accounting of how data moves between logical partitions, we can express the derivative of a MapReduce computation as another MapReduce computation.DrJAX operates on partitioned and non-partitioned values. A partition represents the data that we would like to perform MapReduce-style computations over. A non-partitioned value conceptually represents a singleton. We denote a non-partitioned value by v, and a partitioned value as v = [v 1 , . . . , v n ], where all v i are elements of the same space. Note that the value of n depends on the partition, but the ordering within the partition is arbitrary.The inputs and outputs of DrJAX computations can include non-partitioned and partitioned values. Unlike classical MapReduce treatments, we do not assume that a computation will use a reducer. We consider a specific class of computations that can be built from the following computations, which we refer to as building blocks.1. broadcast: Creates a partitioned value in which all groups have the same value (ie. broadcast(v) = [v, v, . . . , v]).2. map fn: Applies a function f across a partition (ie.map fn(f, v) = [f (v 1 ), f (v 2 ), . . . , f (v n )]).3. reduce sum: Sums over a partitioned value, returning a non-partitioned value (ie. reduce sum(v)= n i=1 v i ).This class is sufficiently expressive to include many parallel algorithms of interest, including parallel modelagnostic meta learning (MAML) (Finn et al., 2017) , parallel and local SGD (Mangasarian and Solodov, 1993; Zinkevich et al., 2010) , federated averaging (FedAvg) (McMahan et al., 2017) , Branch-Train-Merge (Li et al., 2022) , DiLoCo (Douillard et al., 2023) , and many others.For our purposes, the key fact about these computations is that they are closed under MapReduce AD. In particular, let f : x → y where x is any collection of partitioned and non-partitioned inputs, and y is non-partitioned. If f can be composed from the building blocks above, then Rush et al. (2023) show that the function ∇f : x → dy /dx can also be expressed in terms of the building blocks above. Moreover, this can be done in a programmatic fashion. We will refer to this as MapReduce AD in the sequel.This leads to our key observation: If we embed the building blocks above into JAX in a suitable manner, then we can (1) lower these computations to data structures accepted by performant data center runtimes, (2) implement MapReduce AD by appropriately delegating to JAX's AD, and (3) preserve partition information to enable translation to other production ML systems. DrJAX does just this, embedding the building blocks into JAX in a JIT-compatible manner.DrJAX also provides implementations of Jacobian-vector and vector-Jacobian products of the building blocks. This allows DrJAX to perform forward-and reverse-mode AD on MapReduce computations by delegating to JAX's forwardand reverse-mode AD.Authoring surface. DrJAX code is almost entirely JAX code, with two general exceptions. First, there are the building blocks above. Second, DrJAX code must specify how many groups are in a partition during the invocation of the computation. To see this, consider the code in Snippet 1, which simply broadcasts a value across a partition, doubles the value in each group, and takes a sum over the partition. To compute the result, DrJAX needs to know the size of the partition. The user has to give this information to the DrJAX programs. For example, Snippet 2 modifies Snippet 1 to include explicit information about the partition size (ie. the number of groups in the partition).  
implementation: 
 We Given the representation of partitioned values above, we can implement the building blocks via straightforward array operations:1. broadcast: Tile an array over its leading axis.2. map fn: Apply a function pointwise over an array's leading axis.3. reduce sum: Sum an array over its leading axis.We extend these to partitioned structures by applying them leaf-wise. DrJAX registers these implementations with JAX lowering logic. This ensures that DrJAX code is entirely replaced by JAX code by the time JAX dispatches logic to an XLA runtime. Other building blocks can be added to DrJAX by registering primitives in a similar fashion, or by defining them in terms of the building blocks above. For example, DrJAX provides a reduce mean symbol which takes an average across groups in a partitioned array, which lowers to two calls to reduce sum. Sharding DrJAX computations. By registering the primitives above, we ensure that compilers like GSPMD (Xu et al., 2021) can shard DrJAX computations across worker nodes. Critically, and distinct from paradigms such as jax.pmap, DrJAX decouples partition size from sharding computations. A partition of size n is purely logical, and can be sharded across any number of m workers as long as m|n. We want to ensure that, no matter how many workers we shard over, DrJAX computations are as efficient as possible. To do so, we only need to focus on how the building blocks above are sharded by compilers. Once this is done, we are free to compose with model-and data-parallelism provided by various JAX libraries.While some DrJAX building blocks are trivially parallelizable (e.g. map fn), compilers may not be able to detect this and generate efficient code. As noted by Xu et al. (2021) and Lepikhin et al. (2020) , internal sharding annotations can dramatically affect the performance of a compiler targeting distributed execution. DrJAX uses static and dynamic sharding constraints to ensure that after compilation, the resulting computation will run efficiently in the data center. As we will see in Section 4, without these annotations, compilers like GSPMD do not optimally shard DrJAX computations, especially as the partition size increases.Derivatives of DrJAX computations. The last benefit of embedding building blocks as JAX primitives is that it gives us a straightforward way to take derivatives of DrJAX computations using AD. We refer to this as MapReduce AD. To do so, we only need to define the action of vector-Jacobian products (VJPs) and Jacobian-vector products (JVPs) on the DrJAX primitives. Rush et al. (2023) discuss how to compute these products, and show that their computation does not require any new building blocks. That is, the JVPs and VJPs of these primitives can be expressed in terms of the same set of primitives. With the JVPs and VJPs, we can now entirely rely on JAX's AD to do forward-and reversemode AD on computations involving these primitives. For more details, see Section 5. 
scalability and efficiency: 
 We now present numerical evidence of the scalability and efficiency of DrJAX, by testing DrJAX in a distributed training setting. We perform multiple rounds of local SGD (Zhang et al., 2016) on transformer language models with 350 million (350M), 1 billion (1B), and 8 billion (8B) parameters. We use a causal language modeling loss and a sequence length of 512. In every round, we parallelize 4 local SGD steps, each with batch size 8, across some number of data groups from a partition before synchronization. We use a partitioned version of CCNews dataset, where news articles are partitioned according to their base URL domain. We use Dataset Grouper (Charles et al., 2023) to iterate over groups of data efficiently. In each round, we sample some number of data groups, which form a partition of some size. We vary the partition size proportionally to the number of workers used in the computation. To describe the scale of the experiments, Table 1 contains the maximum number of tokens processed and model parameters updated per round for each model.For all experiments, we shard the training computation over some number of TPUv2s. The total number of TPU chips, m, is proportional to the partition size, n. For 350M, 1B, and 8B models we use n, 4n, and 8n chips in total, respectively. This means that if we double the partition size, we also double the number of TPU chips used. We fully shard the computations across the workers, and additionally do model parallelism for the 1B and 8B models. For all experiments, our DrJAX computations are first compiled using GSPMD (Xu et al., 2021) and then delegated to the Pathways runtime (Barham et al., 2022) .Weak scaling. The weak scaling of a system refers to how its compute time changes as the workload and compute resources scale simultaneously. Generally, modern ML systems attempt to obtain near-constant weak scaling performance. 1 For DrJAX, we fix the model size and number of local SGD steps computed per data group in the partition, and vary the partition size and number of workers proportionally in order to vary the workload size. As discussed above, we scale the number of TPU chips used in our simulations linearly with respect to the partition size. that a sufficiently advanced compiler could detect this fact, and parallelize worker training when possible (e.g. within resource constraints in the data center environment).This can be a difficult task for a compiler. To illustrate this difficulty, we implemented a double for loop in place of DrJAX-based training (looping over workers, and over each worker's data). For both programs, we JIT-compiled the program, and provided identical input and output shardings to GSPMD and the XLA compiler stack. Though this stack is quite advanced and used to train many of the largest contemporary ML models, it does not recover the performance of DrJAX from this for-loop implementation. Indeed, round runtime scales linearly with the partition size (and therefore, the number of workers), as expected, rather than remaining constant, indicating an inability to use the increased resource scale allocated to the experiment.GSPMD alone is not enough. A better way to parallelize across workers than the for-loop approach above is to implement DrJAX's MapReduce building blocks and use a compiler like GSPMD (Xu et al., 2021) to do automated sharding of the program. This leads to the question: Do we need DrJAX's internal sharding annotations to obtain weak scaling behavior, or can GSPMD alone fully and efficiently parallelize DrJAX computations? Given the relatively simple nature of local SGD's parallel processing patterns (heavily parallelized model training with infrequent synchronization), one might expect that isolating MapReduce building blocks as primitives with specially-designed sharding annotations is unnecessarily complex.To test this, we took a DrJAX-based implementation of local SGD and removed all of DrJAX's internal sharding annotations at function-tracing time, denoting this DrJAX-NS (DrJAX with no sharding). We then re-ran the simulations in Figure 4 . The results in Figure 6 show that at present, these explicit sharding annotations play a crucial role in ensuring DrJAX's weak-scaling behavior. DrJAX-NS computation times increased sublinearly but significantly faster than DrJAX computation times. Moreover, DrJAX-NS exhibited memory footprint scaling issues. We found that for sufficiently large model or worker pool sizes, DrJAX-NS eventually ran out of high-bandwidth memory. In particular, this occurred for the 1B model with 512 workers and for the 8B model with all tested numbers of workers; that is, at 2 workers and beyond.  
interpreting drjax to other platforms: 
 While data center performance is the primary goal of DrJAX, we wish to preserve the optionality to translate DrJAX computations into artifacts interpretable by other systems, such as federated learning systems (Bonawitz et al., 2019; Paulik et al., 2021; Huba et al., 2022) et al. (2017) . This is the loss of a model on some data after some number of steps of (stochastic) gradient descent. For our purposes, we do a single gradient descent step before evaluating the loss, using some user-specified learning rate.In the nomenclature of meta-learning, the data is partitioned across some number of tasks, and we assume we have access to some loss function loss(x, y) where x is the model, and y is the task. To obtain a jaxpr representing this processing pattern, we provide the concrete shape and type of arguments. For brevity, we assume the model and tasks are scalars, and the loss function is the square loss (ie. loss(x, y) = (x -y) 2 ).Given this information, JAX can generate a jaxpr representing Snippet 4. The result is in Snippet 5. The key takeaway is that this jaxpr preserves the DrJAXdefined primitives representing cross-machine communication, broadcast and reduce mean, both of which are primitives registered by DrJAX in JAX. We can trace through the arguments in the jaxpr to see that the computation operates by (1) broadcasting values (the model and learning rate) across the partition, (2) calculating loss using the broadcast values, (3) taking a mean over the partition.{ lambda ; a:f32[] b:f32[] c:f32[3]. let d:f32[1] = broadcast a e:f32[1] = broadcast b f:f32[1] = sub d e _:f32[1] = integer_pow[y=2] f g:f32[1] = integer_pow[y=1] f h:f32[1] = mul 2.0 g i:f32[1] = mul 1.0 h j:f32[1] = mul c i k:f32[1] = sub d j l:f32[1] = sub k e m:f32[1] = integer_pow[y=2] l n:f32[] = reduce_mean m in (n,) }Snippet 5. jaxpr generated for parallel maml loss.Interpreting the jaxpr to a production systems, especially distributed systems, such as TensorFlow Federated (Ingerman and Ostrowski, 2019) is now straightforward: all crossmachine communication is explicit, and the processing inbetween communication is entirely local and can be extracted into standalone functions executed locally by compute nodes in the system.Integrating MapReduce AD. As discussed in Section 2, the Primitive mechanism allows DrJAX to to specify the behavior of building blocks under JAX's functional transformations, including computing forward-and reverse-mode Jacobians (jax.jacfwd and jax.jacrev). This allows DrJAX to apply AD to MapReduce computations via the forward-and reverse-mode algorithms presented in Rush et al. (2023) . For example, forward-and reverse-mode Jacobians of broadcast can be computed via broadcast and reduce sum, respectively. DrJAX can therefore implement MapReduce AD without additional primitives. This means that the jaxpr of DrJAX computations that use MapReduce AD will contain JAX's standard AD symbols, along with DrJAX's primitive set, ensuring that computations using MapReduce AD are still interpretable to other systems.For example, Snippet 6 gives the jaxpr of the reversemode gradient of parallel maml loss) (ie. the deriva- 
discussion: 
 Why MapReduce AD? While features like scalability and efficiency are self-explanatory, the reader may be interested in why we wish to implement MapReduce AD, especially given the care required to interpret to production systems. In short, MapReduce AD makes expressing efficient algorithms easier (Rush et al., 2023) . By way of analogy, AD has made the development of sophisticated neural network architectures significantly easier. Libraries can define the conceptually simpler forward-pass, and rely on AD to perform backpropagation. The result is often faster and less error-prone than hand-implemented gradient computations (Baydin et al., 2018) .Algorithms that operate on partitioned data can see similar benefits. For example, Snippet 4 contains a DrJAX program used to compute the average MAML loss over tasks, as in meta-learning. By simply calling jax.grad(parallel maml loss), we immediately get a DrJAX program that computes the average MAML gradient over tasks. Using this, we can easily write an algorithm that efficiently parallelizes the MAML algorithm. Snippet 7 depicts this, defining implementing a parallel MAML algorithm by simply pairing jax.grad with an SGD update step. Self-tuning algorithms. Another potential use case for MapReduce AD is creating self-tuning algorithms, which use AD to optimize algorithmic hyperparameters (e.g. hypergradient descent). By using MapReduce AD, we can automatically adjust hyperparameters that govern underlying optimization algorithms, but also adjust the MapReduce operations themselves. For example, Rush et al. (2023) show that the weights governing a weighted mean-based reduction can be learned in tandem with performing federated learning. Similarly, Wang et al. (2023) derive formulas for hypergradient descent on federated learning algorithms, but this process can be slow, error-prone, and algorithm-specific. By contrast, Rush et al. (2023) use an AD system (that inspired MapReduce AD, as we discuss in Section 3)) to do the same, without needing to derive algorithm-specific rules. More generally, MapReduce AD opens the door to a wide variety of self-tuning distributed and parallel algorithms.By pairing differentiable MapReduce primitives with an easy-to-use front-end via JAX, performant building block implementations, and useful sharding information, we hope to accelerate research on distributed and parallel ML. Future work includes (1) generalizations of DrJAX to non-linear reductions, (2) extensions of DrJAX to more general types of data, including hierarchically partitioned data, and (3) mature DrJAX interpreters for specific production systems. 
 
 Appendix:  


Decision: We thank the authors for their time and contribution to WANT and we are pleased to share that after the reviewing process the paper has been accepted. Congratulations! We encourage the authors to consider reviewers' feedback for the improvement of the camera-ready version. We hope to see you in person at the workshop and brainstorm on efficient training research together!