trying vmap for debugging

Author: init-22

algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py

@@ -108,6 +108,7 @@ def init_model_fn(
     initial_params = use_pytorch_weights_inplace(initial_params, file_name="/results/pytorch_base_model_criteo1tb_24_june.pth")
     self._param_shapes = param_utils.jax_param_shapes(initial_params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
+    return initial_params, None
     return jax_utils.replicate(initial_params), None
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:

reference_algorithms/schedule_free/jax/submission.py

@@ -57,15 +57,18 @@ def init_optimizer_state(workload: spec.Workload,
                                    workload.param_shapes)
   optimizer_state = opt_init_fn(params_zeros_like)
 
-  return jax_utils.replicate(optimizer_state), opt_update_fn
+  #return jax_utils.replicate(optimizer_state), opt_update_fn
+  return optimizer_state, opt_update_fn
 
 
 @functools.partial(
-    jax.pmap,
+    jax.vmap,
     axis_name='batch',
-    in_axes=(None, None, 0, 0, 0, 0, 0, None, None),
-    static_broadcasted_argnums=(0, 1),
-    donate_argnums=(2, 3, 4))
+    #in_axes=(None, None, None, 0, 0, 0, 0, None, None))
+    in_axes=(None, None, None, None, None, 0, 0, None, None))
+    #in_axes=(None, None, None, 0, None, 0, 0, None, None))
+    # static_broadcasted_argnums=(0, 1),
+    # donate_argnums=(2, 3, 4))
 def pmapped_train_step(workload,
                        opt_update_fn,
                        model_state,
@@ -143,7 +146,7 @@ def update_params(workload: spec.Workload,
   del eval_results
 
   optimizer_state, opt_update_fn = optimizer_state
-  per_device_rngs = jax.random.split(rng, jax.local_device_count())
+  #per_device_rngs = jax.random.split(rng, jax.local_device_count())
   if hasattr(hyperparameters, 'label_smoothing'):
     label_smoothing = hyperparameters.label_smoothing
   else:
@@ -152,14 +155,22 @@ def update_params(workload: spec.Workload,
     grad_clip = hyperparameters.grad_clip
   else:
     grad_clip = None
+
+  per_example_rngs = jax.random.split(rng, batch['inputs'].shape[0])
+
+  print("Optimizer state: ", jax.tree_map(lambda x: x.shape, optimizer_state))
+  print("Current param container: ", jax.tree_map(lambda x: x.shape, current_param_container))
+  print("model state: ", jax.tree_map(lambda x: x.shape, model_state))
+  print("batch: ", jax.tree_map(lambda x: x.shape, batch))
+  print("rng: ", jax.tree_map(lambda x: x.shape, per_example_rngs))
 
   outputs = pmapped_train_step(workload,
                                opt_update_fn,
                                model_state,
                                optimizer_state,
                                current_param_container,
                                batch,
-                               per_device_rngs,
+                               per_example_rngs,
                                grad_clip,
                                label_smoothing)
   breakpoint()
@@ -186,6 +197,7 @@ def update_params(workload: spec.Workload,
 def get_batch_size(workload_name):
   # Return the global batch size.
   if workload_name == 'criteo1tb':
+    return 8
     return 262_144
   elif workload_name == 'fastmri':
     return 32