Replace JAX multi-device pmap execution with shard_map and add coverage for sharded training, evaluation, and prediction paths

Author: pbenner

README.md

@@ -655,7 +655,7 @@ print(atoms.get_potential_energy())
 
 ### JAX Backend Multi-Device Notes
 
-- When the JAX backend detects more than one local accelerator, it automatically switches to a multi-device (`pmap`) execution. In that mode the training and evaluation batch size must be divisible by `jax.local_device_count()` so that each device processes an identical number of graphs.
+- When the JAX backend detects more than one local accelerator, it automatically switches to a multi-device (`shard_map`) execution. In that mode the training and evaluation batch size must be divisible by `jax.local_device_count()` so that each device processes an identical number of graphs.
 - On single-device machines no extra configuration is required; the backend falls back to the same single-device behaviour that existing scripts expect.
 
 ---

equitrain/backends/jax_backend.py

@@ -36,12 +36,18 @@
 )
 from equitrain.backends.jax_runtime import ensure_multiprocessing_spawn
 from equitrain.backends.jax_utils import (
+    DEVICE_AXIS_NAME,
+    DEVICE_AXIS_SPEC,
+    REPLICATED_SPEC,
     ModelBundle,
+    add_local_device_axis,
     batched_iterator,
     iter_micro_batches,
     load_model_bundle,
+    remove_local_device_axis,
     replicate_to_local_devices,
     set_jax_platform,
+    shard_map_over_local_devices,
     supports_multiprocessing_workers,
     take_chunk,
     unreplicate_from_local_devices,
@@ -400,7 +406,7 @@ def _unreplicate(tree):
 
 
 def _multi_device_chunk_iterator(loader, device_count: int, *, phase: str, logger):
-    """Group per-device micro-batches to feed into ``jax.pmap`` calls."""
+    """Group per-device micro-batches for mapped local-device calls."""
     micro_iter = iter_micro_batches(loader)
     first_chunk = take_chunk(micro_iter, device_count)
     if len(first_chunk) < device_count:
@@ -493,20 +499,28 @@ def _build_train_functions(
         local_devices = jax.local_devices()
 
         def grad_step(params, batch):
+            params = remove_local_device_axis(params)
+            batch = remove_local_device_axis(batch)
             (loss, aux), grads = jax.value_and_grad(
                 loss_fn, has_aux=True, allow_int=True
             )(params, batch)
             grads = _sanitize_grads(grads, clip_value)
-            grads = jax.lax.pmean(grads, axis_name='device')
-            loss = jax.lax.pmean(loss, axis_name='device')
-            aux = jtu.tree_map(lambda x: jax.lax.pmean(x, axis_name='device'), aux)
-            return loss, aux, grads
+            grads = jax.lax.pmean(grads, axis_name=DEVICE_AXIS_NAME)
+            loss = jax.lax.pmean(loss, axis_name=DEVICE_AXIS_NAME)
+            aux = jtu.tree_map(
+                lambda x: jax.lax.pmean(x, axis_name=DEVICE_AXIS_NAME), aux
+            )
+            return add_local_device_axis((loss, aux, grads))
 
-        grad_step_fn = jax.pmap(
-            grad_step, axis_name='device', in_axes=(0, 0), devices=local_devices
+        grad_step_fn = shard_map_over_local_devices(
+            grad_step,
+            in_specs=(DEVICE_AXIS_SPEC, DEVICE_AXIS_SPEC),
+            devices=local_devices,
         )
 
         def apply_updates(state: TrainState, grads, ema_factor):
+            state = remove_local_device_axis(state)
+            grads = remove_local_device_axis(grads)
             updates, new_opt_state = optimizer.update(
                 grads, state.opt_state, state.params
             )
@@ -520,16 +534,15 @@ def apply_updates(state: TrainState, grads, ema_factor):
                 )
             else:
                 new_ema = state.ema_params
-            return (
+            return add_local_device_axis(
                 TrainState(
                     params=new_params, opt_state=new_opt_state, ema_params=new_ema
-                ),
+                )
             )
 
-        apply_updates_fn = jax.pmap(
+        apply_updates_fn = shard_map_over_local_devices(
             apply_updates,
-            axis_name='device',
-            in_axes=(0, 0, None),
+            in_specs=(DEVICE_AXIS_SPEC, DEVICE_AXIS_SPEC, REPLICATED_SPEC),
             devices=local_devices,
         )
         return grad_step_fn, apply_updates_fn
@@ -565,12 +578,20 @@ def _build_eval_step(loss_fn, *, multi_device: bool):
         local_devices = jax.local_devices()
 
         def step(params, batch):
+            params = remove_local_device_axis(params)
+            batch = remove_local_device_axis(batch)
             loss, aux = loss_fn(params, batch)
-            loss = jax.lax.pmean(loss, axis_name='device')
-            aux = jtu.tree_map(lambda x: jax.lax.pmean(x, axis_name='device'), aux)
-            return loss, aux
+            loss = jax.lax.pmean(loss, axis_name=DEVICE_AXIS_NAME)
+            aux = jtu.tree_map(
+                lambda x: jax.lax.pmean(x, axis_name=DEVICE_AXIS_NAME), aux
+            )
+            return add_local_device_axis((loss, aux))
 
-        return jax.pmap(step, axis_name='device', in_axes=(0, 0), devices=local_devices)
+        return shard_map_over_local_devices(
+            step,
+            in_specs=(DEVICE_AXIS_SPEC, DEVICE_AXIS_SPEC),
+            devices=local_devices,
+        )
 
     return jax.jit(loss_fn)
 

equitrain/backends/jax_predict.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import functools
 import logging
 import pickle
 import threading
@@ -17,9 +16,15 @@
 
 from equitrain.argparser import check_args_complete
 from equitrain.backends.jax_utils import (
+    DEVICE_AXIS_SPEC,
+    REPLICATED_SPEC,
+    add_local_device_axis,
     prepare_sharded_batch,
+    remove_local_device_axis,
     set_jax_platform,
-    split_graphs_for_devices,
+    shard_graphs_for_devices,
+    shard_map_over_local_devices,
+    split_device_outputs,
     stack_or_none,
     supports_multiprocessing_workers,
 )
@@ -48,6 +53,14 @@
 }
 
 
+def _get_global_value(globals_attr, key: str):
+    if globals_attr is None:
+        return None
+    if hasattr(globals_attr, 'items') and key in globals_attr:
+        return globals_attr[key]
+    return getattr(globals_attr, key, None)
+
+
 def _prefetch_to_device(iterator, capacity: int, device_put_fn: Callable[[Any], Any]):
     """Prefetch items onto device(s) using a bounded queue."""
     if capacity is None or capacity <= 0:
@@ -97,7 +110,7 @@ def _split_prediction_outputs(
     outputs: dict[str, Any],
     graph: jraph.GraphsTuple,
 ) -> tuple[list[int], dict[str, list[Any]]]:
-    graph_ids = getattr(graph.globals, 'graph_id', None)
+    graph_ids = _get_global_value(graph.globals, 'graph_id')
     if graph_ids is None:
         raise ValueError(
             'Streaming prediction requires graph.globals.graph_id to be set.'
@@ -197,32 +210,39 @@ def predict_streaming(
             except Exception:  # pragma: no cover - defensive fallback
                 total_hint = 0
 
-    _predict_step = functools.partial(
-        jax.pmap,
-        in_axes=(None, 0),
-        out_axes=0,
-        axis_name='devices',
-        devices=local_devices,
-    )(predictor)
+    def _mapped_predictor(params, batch):
+        batch = remove_local_device_axis(batch)
+        return add_local_device_axis(predictor(params, batch))
+
+    if local_device_count <= 1:
+        _predict_step = jax.jit(_mapped_predictor)
+    else:
+        _predict_step = shard_map_over_local_devices(
+            _mapped_predictor,
+            in_specs=(REPLICATED_SPEC, DEVICE_AXIS_SPEC),
+            devices=local_devices,
+        )
 
     def _prepare_device_graphs(graph):
         if local_device_count <= 1:
             if graph.n_node.ndim == 1:
                 device_batch = jax.tree_util.tree_map(lambda x: x[None, ...], graph)
-            else:
-                device_batch = graph
-            return device_batch, [graph]
+                return device_batch, [graph]
+            device_batch = graph
+            device_graphs = split_device_outputs(graph, local_device_count)
+            return device_batch, device_graphs
         if graph.n_node.ndim == 1:
             device_batch = prepare_sharded_batch(graph, local_device_count)
-        else:
-            if graph.n_node.shape[0] != local_device_count:
-                raise ValueError(
-                    'Expected microbatches with leading axis equal to the number of '
-                    f'local devices ({local_device_count}), got axis size '
-                    f'{graph.n_node.shape[0]}.'
-                )
-            device_batch = graph
-        device_graphs = split_graphs_for_devices(graph, local_device_count)
+            device_graphs = shard_graphs_for_devices(graph, local_device_count)
+            return device_batch, device_graphs
+        if graph.n_node.shape[0] != local_device_count:
+            raise ValueError(
+                'Expected microbatches with leading axis equal to the number of '
+                f'local devices ({local_device_count}), got axis size '
+                f'{graph.n_node.shape[0]}.'
+            )
+        device_batch = graph
+        device_graphs = split_device_outputs(graph, local_device_count)
         return device_batch, device_graphs
 
     host_prefetch_cap = int(getattr(data_loader, '_prefetch_batches', 0) or 0)