[dependencies] Upgrade to cuda 13 by default - #2040
Conversation
There was a problem hiding this comment.
Code Review
This pull request upgrades the project's CUDA dependency from version 12.8 to 13.0 across Dockerfiles, CI configurations, package requirements, and documentation. However, two critical package resolution issues were identified: JAX does not publish a cuda13 extra on PyPI, and NVIDIA does not publish a nvidia-cudnn-cu13 package. Both of these changes will cause installation and Docker build failures and should be reverted to their cu12 equivalents.
| [project.optional-dependencies] | ||
| gpu = [ | ||
| "jax[cuda12]>=0.7.2; sys_platform == 'linux'", | ||
| "jax[cuda13]>=0.7.2; sys_platform == 'linux'", |
There was a problem hiding this comment.
JAX does not publish a cuda13 extra on PyPI. The correct extra name for installing CUDA-enabled JAX is cuda12 (or cuda12_pip), which is compatible with CUDA 13 runtimes via CUDA's forward compatibility. Using cuda13 will cause installation of the gpu extra to fail.
"jax[cuda12]>=0.7.2; sys_platform == 'linux'",
| # ---------- PyTorch + cuDNN (setup for transformer-engine) ---------- | ||
| # PyTorch + cuDNN | ||
| RUN uv pip install --system --no-cache-dir "torch==2.11.0" "nvidia-cudnn-cu12>=9.3" && \ | ||
| RUN uv pip install --system --no-cache-dir "torch==2.11.0" "nvidia-cudnn-cu13>=9.3" && \ |
There was a problem hiding this comment.
NVIDIA does not publish a nvidia-cudnn-cu13 package on PyPI. The cuDNN 9.x libraries for CUDA 12 and 13 are distributed under the nvidia-cudnn-cu12 package name. Attempting to install nvidia-cudnn-cu13 will result in a package resolution error and break the Docker build.
RUN uv pip install --system --no-cache-dir "torch==2.11.0" "nvidia-cudnn-cu12>=9.3" && \
Upgrades SkyRL to use cuda 13 wheels (upgraded from cuda 12.8) for all dependencies.
Prereqs: