Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update conda-forge installation docs after drop of support for CUDA 11 #25374

Merged
merged 1 commit into from
Dec 10, 2024

Conversation

traversaro
Copy link
Contributor

@traversaro traversaro commented Dec 10, 2024

Since April 2024 jax does not support anymore CUDA 11, see https://github.com/jax-ml/jax/blob/jax-v0.4.37/CHANGELOG.md#jaxlib-0426-april-3-2024 .

If only CUDA 12 is supported, the installation instructions for the conda-forge build of jax can be simplified:

  • Since CUDA 12, the ptxas executable is now available directly in conda-forge and is automatically installed as a dependency of cuda-enabled jaxlib builds (see Add run dependency on cuda-nvcc when using CUDA 12 conda-forge/jaxlib-feedstock#241), so there is no need to manually add the nvidia channel or install additional packages, or explain what this additional packages are
  • Unless the latest build of CUDA-enabled jaxlib gets marked as broken (i.e. yanked) for any reason, on a machine with a CUDA-enabled GPU conda install jax -c conda-forge will automatically install the CUDA-enabled build of jaxlib. However, this will silently fail if for any reason nvidia-smi does not work on the machine, so I also left the documentation on the command to explicitly request the installation of a CUDA-enabled build of jaxlib.

@traversaro traversaro changed the title Update conda-forge installation docs after CUDA 12 upgrade Update conda-forge installation docs after drop of support for CUDA 11 Dec 10, 2024
@traversaro
Copy link
Contributor Author

fyi @ericmjl @ngam @xhochy

Copy link
Collaborator

@dfm dfm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Dec 10, 2024
@dfm dfm self-assigned this Dec 10, 2024
@copybara-service copybara-service bot merged commit 8e7aaa7 into jax-ml:main Dec 10, 2024
15 of 17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants