Proper way to use pmap for DL training

[zhenlan0426]

I am going through the pmap tutorial, it shows a simplified training loop. I can see other ways to implement the update function. Being new to Jax, can someone comment on if one is prefered over another? Thank you!

@functools.partial(jax.pmap, axis_name='num_devices')
def update(params: Params, xs: jnp.ndarray, ys: jnp.ndarray) -> tuple[Params, jnp.ndarray]:
  """original function in Parallel Evaluation in JAX, params are mannual duplicated along the num_devices axis before passing to this function"""

  # Compute the gradients on the given minibatch (individually on each device).
  loss, grads = jax.value_and_grad(loss_fn)(params, xs, ys)

  # Combine the gradient across all devices (by taking their mean).
  grads = jax.lax.pmean(grads, axis_name='num_devices')

  # Also combine the loss. Unnecessary for the update, but useful for logging.
  loss = jax.lax.pmean(loss, axis_name='num_devices')

  # Each device performs its own update, but since we start with the same params
  # and synchronise gradients, the params stay in sync.
  new_params = jax.tree_map(
      lambda param, g: param - g * LEARNING_RATE, params, grads)

  return new_params, loss

@functools.partial(jax.pmap, in_axes=(None,0,0), axis_name='num_devices')
def update2(params: Params, xs: jnp.ndarray, ys: jnp.ndarray) -> tuple[Params, jnp.ndarray]:
  """instead of copying params, just set in_axes to be None"""

  # Compute the gradients on the given minibatch (individually on each device).
  loss, grads = jax.value_and_grad(loss_fn)(params, xs, ys)

  # Combine the gradient across all devices (by taking their mean).
  grads = jax.lax.pmean(grads, axis_name='num_devices')

  # Also combine the loss. Unnecessary for the update, but useful for logging.
  loss = jax.lax.pmean(loss, axis_name='num_devices')

  # Each device performs its own update, but since we start with the same params
  # and synchronise gradients, the params stay in sync.
  new_params = jax.tree_map(
      lambda param, g: param - g * LEARNING_RATE, params, grads)

  return new_params, loss

@jax.jit
def update3(params: Params, xs: jnp.ndarray, ys: jnp.ndarray) -> tuple[Params, jnp.ndarray]:
  """use pmap only for grad_fn. Instead of jax.lax.pmean, mannual average results once they are back to host"""

  # Compute the gradients on the given minibatch (individually on each device).
  grad_fn = jax.pmap(jax.value_and_grad(loss_fn),in_axes=(None,0,0))
  loss, grads = grad_fn(params, xs, ys)
  loss = loss.mean(0)
  grads = jax.tree_map(lambda x:x.mean(0),grads)
  
  # Each device performs its own update, but since we start with the same params
  # and synchronise gradients, the params stay in sync.
  new_params = jax.tree_map(
      lambda param, g: param - g * LEARNING_RATE, params, grads)

  return new_params, loss