Migrate JAX workloads from pmap to jit

.github/workflows/CI.yml

@@ -180,7 +180,7 @@ jobs:
         pip install -e .
         python tests/reference_algorithm_tests.py --workload=ogbg     --framework=pytorch     --global_batch_size=8     --submission_path=reference_algorithms/target_setting_algorithms/pytorch_nesterov.py     --tuning_search_space=reference_algorithms/target_setting_algorithms/ogbg/tuning_search_space.json
         python tests/reference_algorithm_tests.py --workload=ogbg     --framework=jax     --global_batch_size=8     --submission_path=reference_algorithms/target_setting_algorithms/jax_nesterov.py     --tuning_search_space=reference_algorithms/target_setting_algorithms/ogbg/tuning_search_space.json
-  pytest:
+  pytest-params:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3

README.md

@@ -57,20 +57,16 @@ You can install this package and dependencies in a [Python virtual environment](
 We recommend using a Docker container (or alternatively, a Singularity/Apptainer container) to ensure a similar environment to our scoring and testing environments.
 Both options are described in detail in the [**Getting Started**](/docs/GETTING_STARTED.md) document.
 
-_TL;DR to install the Jax version for GPU run:_
+*TL;DR to install the Jax version for GPU and all workload dependencies run:*
 
 ```bash
-pip3 install -e '.[pytorch_cpu]'
-pip3 install -e '.[jax_gpu]' -f 'https://storage.googleapis.com/jax-releases/jax_cuda_releases.html'
-pip3 install -e '.[full]'
+pip3 install -e '.[pytorch_cpu,jax_gpu,full]' --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
-_TL;DR to install the PyTorch version for GPU run:_
+*TL;DR to install the PyTorch version for GPU and all workload dependencies run:*
 
 ```bash
-pip3 install -e '.[jax_cpu]'
-pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'
-pip3 install -e '.[full]'
+pip3 install -e '.[jax_cpu,pytorch_gpu,full]'
 ```
 
 ## Getting Started

algoperf/checkpoint_utils.py

@@ -7,7 +7,6 @@
 import os
 from typing import Sequence, Tuple
 
-import jax
 import numpy as np
 import torch
 from absl import logging
@@ -210,10 +209,7 @@ def save_checkpoint(
     train_state, eval_results, global_step, preemption_count).
   """
   if framework == 'jax':
-    model_params = jax.device_get(jax_utils.unreplicate(model_params))
     opt_state, _ = optimizer_state
-    opt_state = jax.device_get(jax_utils.unreplicate(opt_state))
-    model_state = jax.device_get(jax_utils.unreplicate(model_state))
   else:
     if isinstance(
       model_params,

algoperf/data_utils.py

@@ -62,14 +62,19 @@ def _prepare(x):
     if remainder_size != 0 or pad_to_global_batch_size:
       x = pad(x, pad_size, padding_value=padding_value)
 
-    # Reshape (global_batch_size, ...) to
-    # (local_device_count, per_device_batch_size, ...).
-    # Assumes that `global_batch_size % local_device_count == 0`.
-    return x.reshape((local_device_count, -1, *x.shape[1:]))
+    # return x.reshape((local_device_count, -1, *x.shape[1:]))
+    return x
 
   return jax.tree.map(_prepare, batch)
 
 
+def shard(batch):
+  local_device_count = max(torch.cuda.device_count(), jax.local_device_count())
+  return jax.tree.map(
+    lambda x: x.reshape((local_device_count, -1, *x.shape[1:])), batch
+  )
+
+
 def pad(
   tensor: np.ndarray, pad_size: int, padding_value: int = 0
 ) -> np.ndarray:

algoperf/jax_sharding_utils.py

@@ -0,0 +1,40 @@
+"""Utilities for dealing with sharding in JAX."""
+
+import jax
+from jax.sharding import NamedSharding
+from jax.sharding import PartitionSpec as P
+
+
+def get_replicate_sharding():
+  """Returns a sharding spec that replicates data across all devices."""
+  mesh = jax.sharding.Mesh(jax.devices(), ('batch',))
+  return NamedSharding(mesh, P())
+
+
+def get_batch_dim_sharding():
+  """Returns a sharding spec that shards data along the first axis."""
+  mesh = jax.sharding.Mesh(jax.devices(), ('batch',))
+  return NamedSharding(mesh, P('batch'))
+
+
+def shard_along_batch_dim(x):
+  """Shards a tensor across all devices."""
+  mesh = jax.sharding.Mesh(jax.devices(), ('batch',))
+  return jax.tree.map(
+    lambda x: jax.device_put(x, NamedSharding(mesh, P('batch'))), x
+  )
+
+
+def replicate(x):
+  """Replicates tensor across all devices."""
+  mesh = jax.sharding.Mesh(jax.devices(), ('batch',))
+  return jax.tree.map(lambda x: jax.device_put(x, NamedSharding(mesh, P())), x)
+
+
+def display_shard_info(x: jax.Array):
+  """Displays shard info of a jax array."""
+  for shard in x.addressable_shards:
+    print(
+      f'shard.device: {shard.device}, index: {shard.index}, replica_id:'
+      f' {shard.replica_id}.\n'
+    )

algoperf/workloads/cifar/cifar_jax/workload.py

@@ -7,12 +7,11 @@
 import jax.numpy as jnp
 import optax
 import tensorflow_datasets as tfds
-from flax import jax_utils
 from flax import linen as nn
 from flax.core import pop
 from jax import lax
 
-from algoperf import param_utils, spec
+from algoperf import jax_sharding_utils, param_utils, spec
 from algoperf.workloads.cifar.cifar_jax import models
 from algoperf.workloads.cifar.cifar_jax.input_pipeline import create_input_iter
 from algoperf.workloads.cifar.workload import BaseCifarWorkload
@@ -29,6 +28,7 @@ def _build_cifar_dataset(
     repeat_final_dataset: Optional[bool] = None,
   ) -> Iterator[Dict[str, spec.Tensor]]:
     ds_builder = tfds.builder('cifar10:3.0.2', data_dir=data_dir)
+    ds_builder.download_and_prepare()
     train = split == 'train'
     assert self.num_train_examples + self.num_validation_examples == 50000
     if split in ['train', 'eval_train']:
@@ -89,8 +89,8 @@ def init_model_fn(self, rng: spec.RandomState) -> spec.ModelInitState:
     model_state, params = pop(variables, 'params')
     self._param_shapes = param_utils.jax_param_shapes(params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    model_state = jax_utils.replicate(model_state)
-    params = jax_utils.replicate(params)
+    model_state = jax_sharding_utils.replicate(params)
+    params = jax_sharding_utils.replicate(params)
     return params, model_state
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:
@@ -171,15 +171,8 @@ def _compute_metrics(
       'loss': summed_loss,
       'accuracy': accuracy,
     }
-    metrics = lax.psum(metrics, axis_name='batch')
     return metrics
 
-  @functools.partial(
-    jax.pmap,
-    axis_name='batch',
-    in_axes=(None, 0, 0, 0, None),
-    static_broadcasted_argnums=(0,),
-  )
   def _eval_model(
     self,
     params: spec.ParameterContainer,
@@ -188,21 +181,41 @@ def _eval_model(
     rng: spec.RandomState,
   ) -> Dict[spec.Tensor, spec.ModelAuxiliaryState]:
     """Return the mean accuracy and loss as a dict."""
-    logits, _ = self.model_fn(
-      params,
-      batch,
-      model_state,
-      spec.ForwardPassMode.EVAL,
-      rng,
-      update_batch_norm=False,
+
+    @functools.partial(
+      jax.jit,
+      in_shardings=(
+        jax_sharding_utils.get_replicate_sharding(),  # params
+        jax_sharding_utils.get_batch_dim_sharding(),  # batch
+        jax_sharding_utils.get_replicate_sharding(),  # model_state
+        jax_sharding_utils.get_batch_dim_sharding(),  # rng
+      ),
     )
-    weights = batch.get('weights')
-    if weights is None:
-      weights = jnp.ones(len(logits))
-    return self._compute_metrics(logits, batch['targets'], weights)
+    def _eval_model_jitted(
+      params: spec.ParameterContainer,
+      batch: Dict[str, spec.Tensor],
+      model_state: spec.ModelAuxiliaryState,
+      rng: spec.RandomState,
+    ) -> Dict[spec.Tensor, spec.ModelAuxiliaryState]:
+      """Return the mean accuracy and loss as a dict."""
+      logits, _ = self.model_fn(
+        params,
+        batch,
+        model_state,
+        spec.ForwardPassMode.EVAL,
+        rng,
+        update_batch_norm=False,
+      )
+      weights = batch.get('weights')
+      if weights is None:
+        weights = jnp.ones(len(logits))
+      return self._compute_metrics(logits, batch['targets'], weights)
+
+    metrics = _eval_model_jitted(params, batch, model_state, rng)
+    return jax.tree.map(lambda x: x.item(), metrics)
 
   def _normalize_eval_metrics(
     self, num_examples: int, total_metrics: Dict[str, Any]
   ) -> Dict[str, float]:
     """Normalize eval metrics."""
-    return jax.tree.map(lambda x: float(x[0] / num_examples), total_metrics)
+    return jax.tree_map(lambda x: x / num_examples, total_metrics)

algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py

@@ -6,9 +6,8 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from flax import jax_utils
 
-from algoperf import param_utils, spec
+from algoperf import jax_sharding_utils, param_utils, spec
 from algoperf.workloads.criteo1tb.criteo1tb_jax import models
 from algoperf.workloads.criteo1tb.workload import BaseCriteo1TbDlrmSmallWorkload
 
@@ -106,7 +105,7 @@ def init_model_fn(
     initial_params = initial_variables['params']
     self._param_shapes = param_utils.jax_param_shapes(initial_params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    return jax_utils.replicate(initial_params), None
+    return jax_sharding_utils.replicate(initial_params), None
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:
     return param_key == 'Dense_7'
@@ -132,13 +131,40 @@ def model_fn(
     logits_batch = self._model.apply({'params': params}, inputs, **apply_kwargs)
     return logits_batch, None
 
+  def _build_input_queue(
+    self,
+    data_rng: spec.RandomState,
+    split: str,
+    data_dir: str,
+    global_batch_size: int,
+    cache: Optional[bool] = None,
+    repeat_final_dataset: Optional[bool] = None,
+    num_batches: Optional[int] = None,
+  ):
+    it = super()._build_input_queue(
+      data_rng,
+      split,
+      data_dir,
+      global_batch_size,
+      cache,
+      repeat_final_dataset,
+      num_batches,
+    )
+    f = functools.partial(
+      jax.device_put, device=jax_sharding_utils.get_batch_dim_sharding()
+    )
+    return map(f, it)
+
   @functools.partial(
-    jax.pmap,
-    axis_name='batch',
-    in_axes=(None, 0, 0),
-    static_broadcasted_argnums=(0,),
+    jax.jit,
+    in_shardings=(
+      jax_sharding_utils.get_replicate_sharding(),
+      jax_sharding_utils.get_batch_dim_sharding(),
+    ),
+    static_argnums=(0,),
+    out_shardings=jax_sharding_utils.get_replicate_sharding(),
   )
-  def _eval_batch_pmapped(
+  def _eval_batch_jitted(
     self, params: spec.ParameterContainer, batch: Dict[str, spec.Tensor]
   ) -> spec.Tensor:
     logits, _ = self.model_fn(
@@ -162,9 +188,7 @@ def _eval_batch(
   ) -> spec.Tensor:
     # We do NOT psum inside of _eval_batch_pmapped, so the returned tensor of
     # shape (local_device_count,) will all be different values.
-    return np.array(
-      self._eval_batch_pmapped(params, batch).sum(), dtype=np.float64
-    )
+    return np.array(self._eval_batch_jitted(params, batch), dtype=np.float64)
 
 
 class Criteo1TbDlrmSmallTestWorkload(Criteo1TbDlrmSmallWorkload):

algoperf/workloads/criteo1tb/criteo1tb_pytorch/workload.py

@@ -7,7 +7,7 @@
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 
-from algoperf import param_utils, spec
+from algoperf import data_utils, param_utils, spec
 from algoperf.pytorch_utils import pytorch_setup
 from algoperf.workloads.criteo1tb.criteo1tb_pytorch import models
 from algoperf.workloads.criteo1tb.workload import BaseCriteo1TbDlrmSmallWorkload
@@ -152,6 +152,7 @@ def _build_input_queue(
         num_batches=num_batches,
         repeat_final_dataset=repeat_final_dataset,
       )
+      np_iter = map(data_utils.shard, np_iter)
     weights = None
     while True:
       if RANK == 0: