[WIP] Migrate JAX workloads from pmap to jit

README.md

@@ -54,16 +54,16 @@ Both options are described in detail in the [**Getting Started**](/docs/GETTING_
 *TL;DR to install the Jax version for GPU 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 '.[pytorch_cpu]' -f https://download.pytorch.org/whl/cpu
+pip3 install -e '.[jax_gpu]'
 pip3 install -e '.[full]'
 ```
 
 *TL;DR to install the PyTorch version for GPU run:*
 
 ```bash
 pip3 install -e '.[jax_cpu]'
-pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'
+pip3 install -e '.[pytorch_gpu]'
 pip3 install -e '.[full]'
 ```
 

algoperf/checkpoint_utils.py

@@ -11,7 +11,6 @@
 from flax import jax_utils
 from flax.training import checkpoints as flax_checkpoints
 from flax.training.checkpoints import latest_checkpoint
-import jax
 import numpy as np
 from tensorflow.io import gfile  # pytype: disable=import-error
 import torch
@@ -193,10 +192,7 @@ def save_checkpoint(framework: str,
     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

@@ -11,6 +11,7 @@
 from torch.utils.data import DistributedSampler
 from torch.utils.data import Sampler
 
+from algoperf import jax_sharding_utils
 from algoperf import spec
 
 
@@ -60,10 +61,7 @@ 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 jax.device_put(x, jax_sharding_utils.get_batch_dim_sharding())
 
   return jax.tree.map(_prepare, batch)
 

algoperf/jax_sharding_utils.py

@@ -0,0 +1,37 @@
+"""Utilities for dealing with sharding in JAX."""
+
+import jax
+from jax.sharding import NamedSharding, 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/input_pipeline.py

@@ -171,5 +171,5 @@ def create_input_iter(
       functools.partial(
           shard_and_maybe_pad_np, global_batch_size=global_batch_size),
       ds)
-  it = jax_utils.prefetch_to_device(it, 2)
+
   return it

algoperf/workloads/cifar/cifar_jax/workload.py

@@ -3,7 +3,6 @@
 import functools
 from typing import Any, Dict, Iterator, Optional, Tuple
 
-from flax import jax_utils
 from flax import linen as nn
 from flax.core import pop
 import jax
@@ -13,6 +12,7 @@
 import tensorflow_datasets as tfds
 
 from algoperf import param_utils
+from algoperf import jax_sharding_utils
 from algoperf import spec
 from algoperf.workloads.cifar.cifar_jax import models
 from algoperf.workloads.cifar.cifar_jax.input_pipeline import create_input_iter
@@ -31,6 +31,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']:
@@ -96,8 +97,8 @@ def init_model_fn(
     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:
@@ -175,35 +176,51 @@ def _compute_metrics(self,
         '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,
-      batch: Dict[str, spec.Tensor],
-      model_state: spec.ModelAuxiliaryState,
-      rng: spec.RandomState) -> Dict[spec.Tensor, spec.ModelAuxiliaryState]:
+    self,
+    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)
+
+    @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
+        ),
+    )
+    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

@@ -11,6 +11,7 @@
 from algoperf import param_utils
 from algoperf import spec
 from algoperf.workloads.criteo1tb.criteo1tb_jax import models
+from algoperf import jax_sharding_utils
 from algoperf.workloads.criteo1tb.workload import \
     BaseCriteo1TbDlrmSmallWorkload
 
@@ -105,7 +106,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'
@@ -129,13 +130,16 @@ def model_fn(
     return logits_batch, None
 
   @functools.partial(
-      jax.pmap,
-      axis_name='batch',
-      in_axes=(None, 0, 0),
-      static_broadcasted_argnums=(0,))
-  def _eval_batch_pmapped(self,
-                          params: spec.ParameterContainer,
-                          batch: Dict[str, spec.Tensor]) -> spec.Tensor:
+      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_jitted(self,
+                         params: spec.ParameterContainer,
+                         batch: Dict[str, spec.Tensor]) -> spec.Tensor:
     logits, _ = self.model_fn(
         params,
         batch,
@@ -156,8 +160,7 @@ def _eval_batch(self,
                   batch: Dict[str, spec.Tensor]) -> 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/fastmri/fastmri_jax/workload.py

@@ -10,6 +10,7 @@
 
 from algoperf import param_utils
 from algoperf import spec
+from algoperf import jax_sharding_utils
 import algoperf.random_utils as prng
 from algoperf.workloads.fastmri.fastmri_jax.models import UNet
 from algoperf.workloads.fastmri.fastmri_jax.ssim import ssim
@@ -39,7 +40,7 @@ def init_model_fn(
     params = variables['params']
     self._param_shapes = param_utils.jax_param_shapes(params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    params = jax_utils.replicate(params)
+    params = jax_sharding_utils.replicate(params)
     return params, None
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:
@@ -94,10 +95,12 @@ def loss_fn(
     }
 
   @functools.partial(
-      jax.pmap,
-      axis_name='batch',
-      in_axes=(None, 0, 0, 0),
-      static_broadcasted_argnums=(0,))
+      jax.jit,
+      in_shardings=(jax_sharding_utils.get_replicate_sharding(),
+                    jax_sharding_utils.get_batch_dim_sharding(),
+                    jax_sharding_utils.get_replicate_sharding()),
+      static_argnums=(0,),
+      out_shardings=jax_sharding_utils.get_replicate_sharding())
   def _eval_model(self,
                   params: spec.Tensor,
                   batch: Dict[str, spec.Tensor],
@@ -126,7 +129,6 @@ def _eval_model(self,
         'ssim': ssim_sum,
         'loss': summed_loss,
     }
-    metrics = jax.lax.psum(metrics, axis_name='batch')
     return metrics
 
   def _eval_model_on_split(self,
@@ -154,13 +156,12 @@ def _eval_model_on_split(self,
           num_batches=num_batches)
 
     total_metrics = {'ssim': 0., 'loss': 0.}
-    eval_rngs = prng.split(model_rng, jax.local_device_count())
     for _ in range(num_batches):
       batch = next(self._eval_iters[split])
       # We already sum these metrics across devices inside _eval_model.
-      synced_metrics = self._eval_model(params, batch, eval_rngs)
+      synced_metrics = self._eval_model(params, batch, model_rng)
       total_metrics = {
-          k: v + synced_metrics[k][0] for k, v in total_metrics.items()
+          k: v + synced_metrics[k] for k, v in total_metrics.items()
       }
     return {k: float(v.item() / num_examples) for k, v in total_metrics.items()}
 

algoperf/workloads/imagenet_resnet/imagenet_jax/input_pipeline.py

@@ -399,6 +399,7 @@ def create_input_iter(split: str,
       ds)
 
   # Note(Dan S): On a Nvidia 2080 Ti GPU, this increased GPU utilization by 10%.
-  it = jax_utils.prefetch_to_device(it, 2)
+  # TODO (kasimbeg): put on device
+  # it = jax_utils.prefetch_to_device(it, 2)
 
   return iter(it)