[WIP] JAX.JIT Switch and Sharding

[rka97]

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 CIFAR/MNIST with jit
• Modified AdamW and Nesterov accordingly
• Updated checkpoint and data_utils to support the new approach (mostly removing explicit jax_utils.replicate calls).

Issues

• Prefetching functionality in CIFAR is temporarily disabled (marked with FIXME), not sure how to best support it here.
• I haven't edited any of the PyTorch code, we will need to make sure they still do comparably..

[github-actions]

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

[rka97]

recheck

[priyakasimbeg]

migrating this PR to one that merges from branch on this repo.