search

young-geng / EasyLM

Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
Apache License 2.0
2.33k stars 247 forks source link

TPU Installation broken because of change in Orbax #84

Open akhilkedia opened 11 months ago

akhilkedia commented 11 months ago

Hi!

Running the scripts/tpu_vm_setup.sh file currently errors out because of a change in the orbax package, as discussed here https://github.com/google/orbax/issues/436

Collecting orbax
  Using cached orbax-0.1.8.tar.gz (1.6 kB)
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [3 lines of output]

      *** Orbax is a namespace, and not a standalone package. For model checkpointing and exporting utilities, please install `orbax-checkpoint` and `orbax-export` respectively (instead of `orbax`). ***

      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

The orbax dependency originates from flax, and flax==0.6.8 requires orbax

While latest flax has changed this to orbax-checkpoint (as can be seen by running curl -L 'https://pypi.python.org/pypi/flax/json' | jq '.info.requires_dist'), the version of flax in scripts/tpu_vm_setup.sh directly depends on orbax.

So perhaps the flax dependency version should be changed in scripts/tpu_vm_setup.sh, or orbax==0.1.7 should be added to fix the dependency.

young-geng commented 11 months ago

We've also noticed that recently, and we think that this would be a good time for us to upgrade the JAX and Flax version. We are still testing some code to make sure everything is compatible with the newest version of JAX and Flax, and will release it as soon as possible.