Migrate JAX workloads from pmap to jit

[priyakasimbeg]

Purpose

The goal of this PR is to allow model parameter and optimizer state sharding, and also to migrate the JAX code from using jax.pmap to using jax.jit.

TODOs:

• Migrate reference optimizers to use jax.jit
  • Nesterov
  • AdamW
  • Others
• Migrate workloads to use jax.jit
  • (Test workload) MNIST
  • (Test workload) CIFAR
  • WMT
  • Criteo1TB
  • FastMRI
  • Librispeech
  • OGBG
  • ImageNet

Changelog

• Added some sharding utilities to handle data distributed
• Replaced pmap code for all workloads with jit
• Modified reference submissions accordingly
• Updated checkpoint and data_utils to support the new approach (mostly removing explicit jax_utils.replicate calls).
• Upgraded Jax version to 0.7.0
• Factored out the sharding into shape (bsz/n_devices, ...) out of shard_and_maybe_pad. PyTorch workloads now call a separate shard function if they used shard_and_maybe_pad before

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[rka97]

We didn't migrate all the optimizers. Should we look into migrating the following?

• adafactor
• lamb
• sam
• shampoo

[priyakasimbeg]

We didn't migrate all the optimizers. Should we look into migrating the following?

• adafactor
• lamb
• sam
• shampoo

The above algorithms are buggy and we have no plans to fix them, so instead of maintaining them we should delete them