ogbg jit migration

Author: priyakasimbeg

algoperf/workloads/ogbg/input_pipeline.py

@@ -148,10 +148,15 @@ def _get_batch_iterator(dataset_iter, global_batch_size, num_shards=None):
     weights_shards.append(weights)
 
     if count == num_shards:
+      # yield {
+      #     'inputs': jraph.batch(graphs_shards),
+      #     'targets': np.vstack(labels_shards),
+      #     'weights': np.vstack(weights_shards)
+      # }
 
       def f(x):
-        return jax.tree.map(lambda *vals: np.stack(vals, axis=0), x[0], *x[1:])
-
+        return jax.tree.map(lambda *vals: np.concatenate(vals, axis=0), x[0], *x[1:])
+      
       graphs_shards = f(graphs_shards)
       labels_shards = f(labels_shards)
       weights_shards = f(weights_shards)

algoperf/workloads/ogbg/ogbg_jax/models.py

@@ -2,6 +2,7 @@
 # https://github.com/google/init2winit/blob/master/init2winit/model_lib/gnn.py.
 from typing import Optional, Tuple
 
+import jax
 from flax import linen as nn
 import jax.numpy as jnp
 import jraph
@@ -78,7 +79,8 @@ def __call__(self, graph, train):
               self.hidden_dims, dropout=dropout, activation_fn=activation_fn),
           update_global_fn=_make_mlp(
               self.hidden_dims, dropout=dropout, activation_fn=activation_fn))
-
+      # jax.debug.print(str(graph))
+    
       graph = net(graph)
 
     # Map globals to represent the final result

algoperf/workloads/ogbg/ogbg_jax/workload.py

@@ -8,6 +8,7 @@
 import jraph
 import optax
 
+from algoperf import sharding_utils
 from algoperf import param_utils
 from algoperf import spec
 from algoperf.workloads.ogbg import metrics
@@ -45,7 +46,8 @@ def init_model_fn(
     params = params['params']
     self._param_shapes = param_utils.jax_param_shapes(params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    return jax_utils.replicate(params), None
+    params = sharding_utils.shard_replicated(params)
+    return params, None
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:
     return param_key == 'Dense_17'
@@ -106,11 +108,20 @@ def _eval_metric(self, labels, logits, masks):
     return metrics.EvalMetrics.single_from_model_output(
         loss=loss['per_example'], logits=logits, labels=labels, mask=masks)
 
+  # @functools.partial(
+  #     jax.pmap,
+  #     axis_name='batch',
+  #     in_axes=(None, 0, 0, 0, None),
+  #     static_broadcasted_argnums=(0,))
   @functools.partial(
-      jax.pmap,
-      axis_name='batch',
-      in_axes=(None, 0, 0, 0, None),
-      static_broadcasted_argnums=(0,))
+    jax.jit,
+    in_shardings=(sharding_utils.get_replicated_sharding(),
+                  sharding_utils.get_naive_sharding_spec(),
+                  sharding_utils.get_replicated_sharding(),
+                  sharding_utils.get_replicated_sharding()),
+    static_argnums=(0,),
+    out_shardings=sharding_utils.get_replicated_sharding(),
+  )
   def _eval_batch(self, params, batch, model_state, rng):
     return super()._eval_batch(params, batch, model_state, rng)
 
@@ -119,7 +130,8 @@ def _normalize_eval_metrics(
                                                    Any]) -> Dict[str, float]:
     """Normalize eval metrics."""
     del num_examples
-    total_metrics = total_metrics.reduce()
+    # total_metrics = total_metrics.reduce()
+    print(total_metrics)
     return {k: float(v) for k, v in total_metrics.compute().items()}
 
 

algoperf/workloads/ogbg/workload.py

@@ -161,6 +161,7 @@ def _eval_batch(self,
         spec.ForwardPassMode.EVAL,
         rng,
         update_batch_norm=False)
+    jax.debug.print(str(logits))
     return self._eval_metric(batch['targets'], logits, batch['weights'])
 
   def _eval_model_on_split(self,

reference_algorithms/paper_baselines/adamw/jax/submission.py

@@ -75,6 +75,7 @@ def _loss_fn(params):
         spec.ForwardPassMode.TRAIN,
         rng,
         update_batch_norm=True,)
+    jax.debug.print("logits: {logits}", logits=logits)
     loss_dict = workload.loss_fn(
         label_batch=batch['targets'],
         logits_batch=logits,
@@ -140,31 +141,29 @@ def update_params(
   replicated = NamedSharding(mesh, P())  # No partitioning
   sharded = NamedSharding(mesh, P('batch'))  # Partition along batch dimension
 
-  # Define input and output shardings
-  arg_shardings = (
-      # workload is static
-      # opt_update_fn is static
-      replicated,  # model_state
-      replicated,  # optimizer_state
-      replicated,  # current_param_container
-      sharded,  # batch
-      replicated,  # rng
-      replicated,  # grad_clip
-      replicated  # label_smoothing
-  )
-  out_shardings = (
-      replicated,  # new_optimizer_state
-      replicated,  # updated_params
-      replicated,  # new_model_state
-      replicated,  # loss
-      replicated  # grad_norm
-  )
   jitted_train_step = jax.jit(
       train_step,
       static_argnums=(0, 1),
       donate_argnums=(2, 3, 4),
-      in_shardings=arg_shardings,
-      out_shardings=out_shardings)
+      in_shardings= (
+                  # workload is static
+                  # opt_update_fn is static
+                  replicated,  # model_state
+                  replicated,  # optimizer_state
+                  replicated,  # current_param_container
+                  sharded,  # batch
+                  replicated,  # rng
+                  replicated,  # grad_clip
+                  replicated  # label_smoothing
+                  ),
+      out_shardings=(
+                  replicated,  # new_optimizer_state
+                  replicated,  # updated_params
+                  replicated,  # new_model_state
+                  replicated,  # loss
+                  replicated  # grad_norm
+  ))
+  # print(batch)
   new_optimizer_state, new_params, new_model_state, loss, grad_norm = jitted_train_step(workload,
                               opt_update_fn,
                               model_state,
@@ -176,7 +175,7 @@ def update_params(
                               label_smoothing)
 
   # Log loss, grad_norm.
-  if global_step % 100 == 0 and workload.metrics_logger is not None:
+  if global_step % 1 == 0 and workload.metrics_logger is not None:
     workload.metrics_logger.append_scalar_metrics(
         {
             'loss': loss.item(),

submission_runner.py

@@ -392,8 +392,9 @@ def train_once(
         train_step_end_time - train_state['last_step_end_time'])
 
     # Check if submission is eligible for an untimed eval.
-    if ((train_step_end_time - train_state['last_eval_time']) >=
-        workload.eval_period_time_sec or train_state['training_complete']):
+    if False:
+    # if ((train_step_end_time - train_state['last_eval_time']) >=
+    #     workload.eval_period_time_sec or train_state['training_complete']):
 
       # Prepare for evaluation (timed).
       if prepare_for_eval is not None: