jax.distributed.initialize() fails on jax==0.4.26 about jax

[hawkinsp]

I think this bug has already been fixed at head, but it made it into the last release because we don't have any way to test SLURM clusters in CI.

Can you try with a nightly jaxlib and jax?

[MScherbela]

You are right, it is indeed already fixed on the nightly build.
It works as expected using

jax==0.4.27.dev20240410
jaxlib==0.4.27.dev20240410+cuda12.cudnn89

Thanks!

[samblouir]

Ran into the same issue without using SLURM, but multi-node GPUs.
Same fix worked with these too:

jax==0.4.27.dev20240415
jaxlib==0.4.27.dev20240415+cuda12.cudnn89

Just dropping the install lines here in case anyone needed it like I did (got it from MaxText)

pip install --pre -U jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda12_releases.html

# Maybe Transformer engine for H100+?
pip install "transformer-engine==1.5.0+297459b" \
          --extra-index-url https://us-python.pkg.dev/gce-ai-infra/maxtext-build-support-packages/simple/