Code

[WIP] JAX.JIT Switch and Sharding#822

Code

Open

rka97 wants to merge 9 commits intomlcommons:devmlcommons/algorithmic-efficiency:devfrom rka97:jit_switchrka97/algorithmic-efficiency:jit_switchCopy head branch name to clipboard

+424-210

Commits

Commits on Nov 21, 2024

• Use jax.jit for sharding initial steps

  Show description for 89fec8f
  rka97

  committed
  89fec8f

  Copy full SHA for 89fec8f

• Use jax.jit for adamw

  rka97

  committed
  c53729c

  Copy full SHA for c53729c

Commits on Dec 9, 2024

• Pass yapf checks

  rka97

  committed
  ef6af03

  Copy full SHA for ef6af03

• CIFAR workload sharding

  rka97

  committed
  e6037d6

  Copy full SHA for e6037d6

Commits on Jan 9, 2025

• librispeech_conformer now running

  Show description for 0698e34
  rka97

  committed
  0698e34

  Copy full SHA for 0698e34

Commits on Feb 5, 2025

• fix formatting

  rka97

  committed
  6094690

  Copy full SHA for 6094690

• shard default

  rka97

  committed
  8ad564f

  Copy full SHA for 8ad564f

• start imagenet

  rka97

  committed
  1a728dc

  Copy full SHA for 1a728dc

• remove bn sync in imagenet (jit handles it automatically)

  rka97

  committed
  c351975

  Copy full SHA for c351975