Parallel Universes: SPMD Parallelism with jax.shard_map and the Evolution Beyond pmap

Parallel Universes: SPMD Parallelism with jax.shard_map and the Evolution Beyond pmap | CodeSignal Learn

Introduction

Welcome to the third lesson of our course "Advanced JAX: Transformations for Speed & Scale"! It's wonderful to continue this journey with you as we explore increasingly sophisticated JAX transformations. In our previous lessons, you mastered JAX's functional approach to randomness through explicit key management and learned how jax.vmap enables elegant, automatic vectorization across batches. These concepts of explicit control and efficient batch processing perfectly prepare us for today's topic: parallel computing across multiple devices . In this lesson, we'll dive into Single Program, Multiple Data (SPMD) parallelism with JAX's modern jax.shard_map transformation, while also examining the now-deprecated jax.pmap for historical context in legacy code. Just as you learned to explicitly manage randomness and batching, SPMD parallelism allows you to explicitly control how computations are distributed across multiple devices — whether CPUs, GPUs, or TPUs. This capability is fundamental for scaling machine learning and scientific computing to handle truly massive datasets and complex models. By the end of this lesson, you'll understand how to create device meshes, partition data across multiple devices, implement collective operations for cross-device communication, and leverage JAX's parallelism for significant performance gains. You'll also appreciate why the JAX ecosystem has evolved from pmap to shard_map, and how these tools integrate with the transformations you've already mastered.

Understanding Parallel Computing Challenges

Before exploring JAX's parallelism tools, let's understand the computational challenges they address. In modern machine learning and scientific computing, we frequently encounter problems that are too large or too computationally intensive for a single device to handle efficiently: Memory Limitations : Large datasets or models may exceed the memory capacity of a single GPU or CPU. For instance, training large language models often requires distributing model parameters across multiple devices. Computational Bottlenecks : Some operations benefit tremendously from parallel execution. Computing statistics across millions of data points, running numerous simulations, or processing large image batches can be dramatically accelerated by distributing work. Scalability Requirements : As datasets grow, serial processing becomes prohibitively slow. Parallel computing allows computational time to scale more favorably with problem size, enabling us to tackle larger and more complex tasks. The traditional approach to parallel computing can involve complex, low-level programming. Single Program, Multiple Data (SPMD) offers a more elegant solution. In SPMD, you write a single program that runs simultaneously on multiple devices, each processing its own portion (or "shard") of the data. When devices need to communicate — to aggregate results or synchronize — they use collective operations . This approach combines the simplicity of writing one program with the power of parallel execution.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

JAX Parallelism API: jax.shard_map

JAX's modern parallelism API centers around jax.shard_map, also known as shmap, which has replaced the older jax.pmap function. The shard_map transformation offers significantly more flexibility and composability, making it the recommended approach for explicit parallelism. It works by taking a function that operates on data shards and mapping it across a device mesh. Key components for jax.shard_map are: The function: Your computation designed to operate on a shard of data. The mesh: Your logical arrangement of devices. Input specs (in_specs): How input arrays should be partitioned across the mesh. Output specs (out_specs): How output arrays should be assembled or structured from device results. Let's start by setting up our environment and examining the available devices. Pythonimport jax import jax.numpy as jnp import jax.lax as lax from jax.sharding import Mesh, PartitionSpec as P # PartitionSpec is often aliased to P # Check available devices n_devices = jax.local_device_count() print(f"Number of JAX local devices: {n_devices}") print(f"Devices: {jax.local_devices()}")import jax import jax.numpy as jnp import jax.lax as lax from jax.sharding import Mesh, PartitionSpec as P # PartitionSpec is often aliased to P # Check available devices n_devices = jax.local_device_count() print(f"Number of JAX local devices: {n_devices}") print(f"Devices: {jax.local_devices()}") This code will show you the available JAX devices. Note that jax.local_devices() returns only the devices accessible to the current process/host, while jax.devices() returns all devices visible to JAX, including remote devices in multi-host distributed setups. For single-host configurations like ours, both functions typically return the same devices. Note: our CodeSignal IDE environment does not yet offer support for accelerators such as GPUs. However, we can set the environment variable XLA_FLAGS="--xla_force_host_platform_device_count=8" to simulate an 8-device configuration, even on CPU-only systems. This simulation is fantastic because it allows you to learn and experiment with the shard_map API and parallelism concepts without needing actual multi-GPU hardware. The principles you learn here apply directly to real multi-device setups. However, note that this simulation doesn't show the actual performance gains you would get from a real multi-device setup. You should see output similar to this, indicating 8 simulated CPU devices: textNumber of JAX local devices: 8 Devices: [CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]Number of JAX local devices: 8 Devices: [CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]

Creating Device Meshes

Now, let's create our device mesh and prepare data for parallel processing. We'll build an example that computes statistics across a dataset, distributing the computation. Pythondata_per_device = 8 # Each device will process 8 data points total_size = n_devices * data_per_device # Create example data key = jax.random.key(42) full_data = jax.random.normal(key, (total_size,)) print(f"\nOriginal data shape: {full_data.shape}") print(f"Data mean: {jnp.mean(full_data):.4f}, variance: {jnp.var(full_data):.4f}") # Create a 1D device mesh # We'll use the first n_devices available devices_array = jax.devices()[:n_devices] mesh = Mesh(devices_array, axis_names=('devices',)) # Name the single axis 'devices' # This print statement is part of the overall script flow print(f"\nModern shard_map approach:") print(f"Mesh: {mesh}")data_per_device = 8 # Each device will process 8 data points total_size = n_devices * data_per_device # Create example data key = jax.random.key(42) full_data = jax.random.normal(key, (total_size,)) print(f"\nOriginal data shape: {full_data.shape}") print(f"Data mean: {jnp.mean(full_data):.4f}, variance: {jnp.var(full_data):.4f}") # Create a 1D device mesh # We'll use the first n_devices available devices_array = jax.devices()[:n_devices] mesh = Mesh(devices_array, axis_names=('devices',)) # Name the single axis 'devices' # This print statement is part of the overall script flow print(f"\nModern shard_map approach:") print(f"Mesh: {mesh}") Here, we generate some random data. Then, we create a Mesh. Our mesh organizes the first n_devices (8 in our simulated case) along a single axis named devices. This mesh will coordinate how our data is distributed. The output of this snippet will be: textOriginal data shape: (64,) Data mean: 0.0827, variance: 0.5778 Modern shard_map approach: Mesh: Mesh('devices': 8, axis_types=(Auto,))Original data shape: (64,) Data mean: 0.0827, variance: 0.5778 Modern shard_map approach: Mesh: Mesh('devices': 8, axis_types=(Auto,))

Creating Sharded Arrays

Next, we create a sharded array — a JAX array that is automatically partitioned across our device mesh. The key to this process is NamedSharding, which serves as the bridge between your device mesh and your data distribution strategy, allowing for: Explicit Control: NamedSharding tells JAX exactly how to distribute your data across the mesh, rather than copying the same data to all devices. Mesh-Data Connection: it links your logical device mesh structure with your data partitioning strategy. Automatic Distribution: JAX handles the complex work of actually splitting and moving data chunks to the right devices. Python# Create sharded array - data is automatically partitioned across devices # The 'with mesh:' context ensures operations know about our device arrangement with mesh: # jax.sharding.NamedSharding specifies how to shard 'full_data' # P('devices') means partition the first axis of 'full_data' along the 'devices' mesh axis sharding_spec = jax.sharding.NamedSharding(mesh, P('devices')) sharded_data = jax.device_put(full_data, sharding_spec) print(f"Sharded data shape: {sharded_data.shape}") # Shape remains the same from a global view print(f"Each device processes chunk of size: {total_size // n_devices}")# Create sharded array - data is automatically partitioned across devices # The 'with mesh:' context ensures operations know about our device arrangement with mesh: # jax.sharding.NamedSharding specifies how to shard 'full_data' # P('devices') means partition the first axis of 'full_data' along the 'devices' mesh axis sharding_spec = jax.sharding.NamedSharding(mesh, P('devices')) sharded_data = jax.device_put(full_data, sharding_spec) print(f"Sharded data shape: {sharded_data.shape}") # Shape remains the same from a global view print(f"Each device processes chunk of size: {total_size // n_devices}") Inside the with mesh: block, jax.device_put with the NamedSharding object distributes full_data according to P('devices') — this PartitionSpec tells JAX to split the first dimension of our data across the devices axis of our mesh. Each of the 8 (simulated) devices receives a contiguous chunk of 8 data points, though the global shape remains unchanged from JAX's perspective: textSharded data shape: (64,) Each device processes chunk of size: 8Sharded data shape: (64,) Each device processes chunk of size: 8

Implementing Collective Operations

Parallelize Collective Operations with shmap

Now, we apply this function using jax.shard_map: Python# Apply shard_map to our function and sharded data shard_map_fn = jax.shard_map( parallel_statistics_shard_map, # The function to parallelize mesh=mesh, # The device mesh to use in_specs=P('devices'), # How the input 'sharded_data' is partitioned out_specs=(P('devices'), P(), P()) # How the outputs should be structured ) # Call the shard_map-transformed function local_means, global_mean, global_var = shard_map_fn(sharded_data) print(f"\nResults from shard_map:") print(f"Local means per device: {local_means}") print(f"Global mean: {global_mean:.4f}") print(f"Global variance: {global_var:.4f}")# Apply shard_map to our function and sharded data shard_map_fn = jax.shard_map( parallel_statistics_shard_map, # The function to parallelize mesh=mesh, # The device mesh to use in_specs=P('devices'), # How the input 'sharded_data' is partitioned out_specs=(P('devices'), P(), P()) # How the outputs should be structured ) # Call the shard_map-transformed function local_means, global_mean, global_var = shard_map_fn(sharded_data) print(f"\nResults from shard_map:") print(f"Local means per device: {local_means}") print(f"Global mean: {global_mean:.4f}") print(f"Global variance: {global_var:.4f}") Let's break down in_specs and out_specs: in_specs=P('devices'): Tells shard_map that the input sharded_data is already partitioned along its first axis, corresponding to the devices mesh axis. out_specs=(P('devices'), P(), P()): This tuple corresponds to the three outputs of parallel_statistics_shard_map: P('devices') for local_mean: The local means (one from each device) will form an array sharded across the devices axis. P() for global_mean: An empty PartitionSpec means this output is replicated; each device will have a copy of the global_mean. P() for global_var: Similarly, global_var is replicated to all devices. The output will be: textResults from shard_map: Local means per device: [ 0.03699342 0.69586265 -0.36393744 0.18055731 -0.10241726 0.12188271 0.12152806 -0.02887268] Global mean: 0.0827 Global variance: 0.5778Results from shard_map: Local means per device: [ 0.03699342 0.69586265 -0.36393744 0.18055731 -0.10241726 0.12188271 0.12152806 -0.02887268] Global mean: 0.0827 Global variance: 0.5778 You can see the local_means array has 8 values (one from each simulated device), while global_mean and global_var are single scalar values, replicated across devices.

Historical Context: The Legacy jax.pmap

To provide context and appreciate JAX's evolution, let's briefly look at the now-deprecated jax.pmap (parallel map). While pmap is no longer recommended for new code, understanding its basics helps highlight the improvements in shard_map and is essential for reading legacy code that may still use this approach. For simplicity, our pmap example will only compute the global mean. Pythondef parallel_statistics_pmap(data_shard): """Compute global mean using pmap - deprecated approach.""" # local_mean = jnp.mean(data_shard) # Not returned, just for illustration local_sum = jnp.sum(data_shard) local_count = data_shard.size # Collective sum across devices global_sum = lax.psum(local_sum, axis_name='devices') # 'devices' matches the axis_name specified in pmap global_count = lax.psum(local_count, axis_name='devices') global_mean = global_sum / global_count return global_mean # pmap replicates outputs by default # Data preparation for pmap: requires explicit reshaping # The first dimension of the input to pmap must be equal to the number of devices pmap_data = full_data.reshape(n_devices, data_per_device) print(f"\nLegacy pmap approach (deprecated):") print(f"Sharded data shape for pmap: {pmap_data.shape}") # Apply pmap # The 'axis_name' here names the mapped axis for use in collectives like psum pmap_fn = jax.pmap(parallel_statistics_pmap, axis_name='devices') pmap_result_mean = pmap_fn(pmap_data) # pmap_data is implicitly sharded across the leading axis print(f"Results from pmap:") print(f"Global mean (replicated): {pmap_result_mean}") # Verification against NumPy and shard_map results print(f"\nVerification:") print(f"shard_map global mean matches numpy: {jnp.allclose(global_mean, jnp.mean(full_data))}") # pmap_result_mean will be an array with the mean replicated n_devices times print(f"pmap global mean matches numpy: {jnp.allclose(pmap_result_mean[0], jnp.mean(full_data))}")def parallel_statistics_pmap(data_shard): """Compute global mean using pmap - deprecated approach.""" # local_mean = jnp.mean(data_shard) # Not returned, just for illustration local_sum = jnp.sum(data_shard) local_count = data_shard.size # Collective sum across devices global_sum = lax.psum(local_sum, axis_name='devices') # 'devices' matches the axis_name specified in pmap global_count = lax.psum(local_count, axis_name='devices') global_mean = global_sum / global_count return global_mean # pmap replicates outputs by default # Data preparation for pmap: requires explicit reshaping # The first dimension of the input to pmap must be equal to the number of devices pmap_data = full_data.reshape(n_devices, data_per_device) print(f"\nLegacy pmap approach (deprecated):") print(f"Sharded data shape for pmap: {pmap_data.shape}") # Apply pmap # The 'axis_name' here names the mapped axis for use in collectives like psum pmap_fn = jax.pmap(parallel_statistics_pmap, axis_name='devices') pmap_result_mean = pmap_fn(pmap_data) # pmap_data is implicitly sharded across the leading axis print(f"Results from pmap:") print(f"Global mean (replicated): {pmap_result_mean}") # Verification against NumPy and shard_map results print(f"\nVerification:") print(f"shard_map global mean matches numpy: {jnp.allclose(global_mean, jnp.mean(full_data))}") # pmap_result_mean will be an array with the mean replicated n_devices times print(f"pmap global mean matches numpy: {jnp.allclose(pmap_result_mean[0], jnp.mean(full_data))}") Key differences from shard_map are illustrated here: Data Preparation: pmap typically requires you to manually reshape your input data so that the leading axis corresponds to the devices. sharded_data for pmap has shape (8, 8). Output Handling: pmap by default replicates all outputs to all devices. pmap_result_mean will be an array of shape (8,), where each element is the same global mean. shard_map offers more control via out_specs. Mesh Explicitness: pmap operates on an implicit 1D array of devices, whereas shard_map uses an explicit Mesh object, allowing for more complex multi-dimensional device arrangements. The output from this section will be: textLegacy pmap approach (deprecated): Sharded data shape for pmap: (8, 8) Results from pmap: Global mean (replicated): [0.0826996 0.0826996 0.0826996 0.0826996 0.0826996 0.0826996 0.0826996 0.0826996] Verification: shard_map global mean matches numpy: True pmap global mean matches numpy: TrueLegacy pmap approach (deprecated): Sharded data shape for pmap: (8, 8) Results from pmap: Global mean (replicated): [0.0826996 0.0826996 0.0826996 0.0826996 0.0826996 0.0826996 0.0826996 0.0826996] Verification: shard_map global mean matches numpy: True pmap global mean matches numpy: True Both methods yield the correct global mean, but shard_map provides a more flexible and modern API for today's complex parallelism needs.

JAX's approach to SPMD parallelism centers around two key concepts: device meshes and data partitioning. Understanding these is crucial for effective parallel programming in JAX.

A device mesh is a logical arrangement of your available computing devices (CPUs, GPUs, or TPUs) into a multi-dimensional grid. Think of it as a coordinate system for your devices. For example, you might arrange 8 devices into a $1 \times 8$ mesh, where devices are organized along a single named axis, perhaps data_parallel_axis. This logical organization allows you to specify precisely how data and computations should be distributed.

Data partitioning determines how JAX arrays are split across the devices in your mesh. JAX uses PartitionSpec objects (often aliased as P) to specify which dimensions of your data should be partitioned along which axes of your device mesh. For instance, you might partition a large batch of training examples along the data_parallel_axis of your mesh, so each device processes a different subset of examples. JAX handles the complex details of data movement and synchronization, allowing you to focus on your parallel algorithm's logic.