ARM64 Build

[Mehdi Amini]

I've been trying to install PAXML on Ubuntu 22.04 ARM64 but I seem to stuck in getting lingvo (mandatory dependency?) running there: I've been struggling to find a recipe for this. Has this been done? Any documentation about it?

[Josh Lipschultz]

Hi Medhi - please see our Dockerfile here: https://github.com/google/paxml/blob/main/paxml/pip_package/Dockerfile#L1. It also uses the Ubuntu 22.04 distribution that you're trying to use, and it's what we use as our public testing environment, so it should work. You can check the other files in that pip_package/ directory for more information. If you could provide more detail about why/how the lingvo dependency is breaking your install, feel free to add it here. Thanks.

[Mehdi Amini]

Is your public testing environment running on ARM64 or just on X86 though? Any link to your CI?

Right now pip is just unhappy about lingvo, here is the verbose output:

$ pip-compile --quiet requirements.in     ../praxis/requirements.in     --output-file $WHEEL_FOLDER/paxml_requirements.txt -vvv
WARNING: the legacy dependency resolver is deprecated and will be removed in future versions of pip-tools. The default resolver will be changed to 'backtracking' in pip-tools 7.0.0. Specify --resolver=backtracking to silence this warning.
Using indexes:
  https://pypi.org/simple

                          ROUND 1                           
Current constraints:
  absl-py (from -r requirements.in (line 4))
  clu (from -r requirements.in (line 5))
  einops (from -r ../praxis/requirements.in (line 6))
  etils (from -r requirements.in (line 6))
  fiddle@ git+https://github.com/google/fiddle from git+https://github.com/google/fiddle (from -r ../praxis/requirements.in (line 8))
  flax (from -r requirements.in (line 7))
  jax@ git+https://github.com/google/jax from git+https://github.com/google/jax (from -r requirements.in (line 8))
  jax-bitempered-loss (from -r ../praxis/requirements.in (line 11))
  lingvo (from -r requirements.in (line 9))
  numpy (from -r requirements.in (line 10))
  optax (from -r ../praxis/requirements.in (line 14))
  optax-shampoo (from -r ../praxis/requirements.in (line 15))
  orbax-checkpoint (from -r requirements.in (line 11))
  protobuf==3.19.6 (from -r requirements.in (line 13))
  pyglove (from -r requirements.in (line 14))
  seqio@ git+https://github.com/google/seqio from git+https://github.com/google/seqio (from -r requirements.in (line 15))
  t5 (from -r requirements.in (line 16))
  tensorflow~=2.9.2 (from -r requirements.in (line 17))
  tensorflow-datasets==4.8.3 (from -r requirements.in (line 20))
  tensorflow-metadata==1.12.0 (from -r requirements.in (line 22))
  tensorflow-text~=2.9.0 (from -r requirements.in (line 18))
  tensorstore (from -r requirements.in (line 19))
  tfds-nightly==4.8.3.dev202303280045 (from -r requirements.in (line 21))

Finding the best candidates:
  found candidate absl-py==1.4.0 (constraint was <any>)
  found candidate clu==0.0.9 (constraint was <any>)
  found candidate einops==0.6.1 (constraint was <any>)
  found candidate etils==1.3.0 (constraint was <any>)
  found candidate fiddle @ git+https://github.com/google/fiddle (constraint was <any>)
  found candidate flax==0.6.10 (constraint was <any>)
  found candidate jax @ git+https://github.com/google/jax (constraint was <any>)
  found candidate jax-bitempered-loss==0.0.2 (constraint was <any>)
Using legacy resolver. Consider using backtracking resolver with `--resolver=backtracking`.
Could not find a version that matches lingvo (from -r requirements.in (line 9))
No versions found
Were https://pypi.org/simple reachable?

[Josh Lipschultz]

Ah right - sorry I misread your question earlier. We currently only provide manylinux2014_x86_64 wheels of lingvo. You can see the available versions here: https://pypi.org/project/lingvo/. The most recent and up to date wheels are compatible with Python 3.8-3.10: e.g. lingvo-0.12.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

If you would like, you could consider following the directions at the following page to build it in your environment. It's not something we've ever tried, however. https://github.com/tensorflow/lingvo/tree/master/pip_package.

[Mehdi Amini]

The most recent and up to date wheels are compatible with Python 3.8-3.10: e.g. lingvo-0.12.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

Yes I saw that: there are only x86 packages unfortunately.

Similarly the Tensorflow tf_sig_build_dockerfiles and scripts are all hardcoded for x86/amd64, so this gets into a deep rabbit hole.

I guess right it stops at "lingvo" (and PAXML) aren't supported on ARM64.