mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[JAX] Obtain NCCL via a stub, rather than linking it statically or dynamically.

Browse Source 
This shrinks the CUDA jaxlib wheel size by around 80MB.

PiperOrigin-RevId: 570554454
This commit is contained in:
Peter Hawkins 2023-10-03 18:33:09 -07:00 committed by jax authors
parent 816ebf2fae
commit efc18e4147
3 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.bazelrc
View File

@ -62,6 +62,7 @@ build:native_arch_posix --host_copt=-march=native
build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=1
build:cuda --repo_env TF_NEED_CUDA=1
build:cuda --repo_env TF_NCCL_USE_STUB=1
# "sm" means we emit only cubin, which is forward compatible within a GPU generation.
# "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations.
build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="sm_52,sm_60,sm_70,sm_80,compute_90"

7
CHANGELOG.md
View File

@ -10,6 +10,11 @@ Remember to align the itemized text with the first line of an item within a list
# jaxlib 0.4.18
* Changes
* CUDA jaxlibs now depend on the user to install a compatible NCCL version.
If using the recommended `cuda12_pip` installation, NCCL should be installed
automatically.
# jax 0.4.17 (Oct 3, 2023)
* New features
@ -42,8 +47,6 @@ Remember to align the itemized text with the first line of an item within a list
* Changes:
* Python 3.12 wheels were added in this release.
* The CUDA 12 wheels now require CUDA 12.2 or newer and cuDNN 8.9.4 or newer.
* This is likely the last release of jaxlib that supports CUDA 11. We intend
to drop CUDA 11 support in the next release.
* Bug fixes:
* Fixed log spam from ABSL when the JAX CPU backend was initialized.

2
setup.py
View File

@ -119,6 +119,7 @@ setup(
"nvidia-cufft-cu11>=10.9",
"nvidia-cusolver-cu11>=11.4",
"nvidia-cusparse-cu11>=11.7",
"nvidia-nccl-cu11>=2.18.3",
],
'cuda12_pip': [
@ -131,6 +132,7 @@ setup(
"nvidia-cufft-cu12>=11.0.8.103",
"nvidia-cusolver-cu12>=11.5.2",
"nvidia-cusparse-cu12>=12.1.2.141",
"nvidia-nccl-cu12>=2.18.3",
],
# Target that does not depend on the CUDA pip wheels, for those who want