Skip to content

Commit

Permalink
Introduce hermetic CUDA in Google ML projects.
Browse files Browse the repository at this point in the history
1) Hermetic CUDA rules allow building wheels with GPU support on a machine without GPUs, as well as running Bazel GPU tests on a machine with only GPUs and NVIDIA driver installed. When `--config=cuda` is provided in Bazel options, Bazel will download CUDA, CUDNN and NCCL redistributions in the cache, and use them during build and test phases.

    [Default location of CUNN redistributions](https://developer.download.nvidia.com/compute/cudnn/redist/)

    [Default location of CUDA redistributions](https://developer.download.nvidia.com/compute/cuda/redist/)

    [Default location of NCCL redistributions](https://pypi.org/project/nvidia-nccl-cu12/#history)

2) To include hermetic CUDA rules in your project, add the following in the WORKSPACE of the downstream project dependent on XLA.

   Note: use `@local_tsl` instead of `@tsl` in Tensorflow project.

   ```
   load(
      "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
      "cuda_json_init_repository",
   )

   cuda_json_init_repository()

   load(
      "@cuda_redist_json//:distributions.bzl",
      "CUDA_REDISTRIBUTIONS",
      "CUDNN_REDISTRIBUTIONS",
   )
   load(
      "@tsl//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
      "cuda_redist_init_repositories",
      "cudnn_redist_init_repository",
   )

   cuda_redist_init_repositories(
      cuda_redistributions = CUDA_REDISTRIBUTIONS,
   )

   cudnn_redist_init_repository(
      cudnn_redistributions = CUDNN_REDISTRIBUTIONS,
   )

   load(
      "@tsl//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
      "cuda_configure",
   )

   cuda_configure(name = "local_config_cuda")

   load(
      "@tsl//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
      "nccl_redist_init_repository",
   )

   nccl_redist_init_repository()

   load(
      "@tsl//third_party/nccl/hermetic:nccl_configure.bzl",
      "nccl_configure",
   )

   nccl_configure(name = "local_config_nccl")
   ```

PiperOrigin-RevId: 662981325
  • Loading branch information
tensorflower-gardener authored and copybara-github committed Aug 14, 2024
1 parent 5b47d57 commit 7d4187e
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions site/en/install/source.md
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,13 @@ Preconfigured Bazel build configs to DISABLE default on features:

#### GPU support

##### from v.2.18.0
For [GPU support](./pip.md), set `cuda=Y` during configuration and specify the
versions of CUDA and cuDNN if required. Bazel will download CUDA and CUDNN
packages automatically or point to CUDA/CUDNN/NCCL redistributions on local file
system if required.

##### before v.2.18.0
For [GPU support](./pip.md), set `cuda=Y` during configuration and specify the
versions of CUDA and cuDNN. If your system has multiple versions of CUDA or
cuDNN installed, explicitly set the version instead of relying on the default.
Expand Down Expand Up @@ -370,6 +377,9 @@ On your host machine, the TensorFlow *pip* package is in the current directory

### GPU support

Note: Starting from Tensorflow v.2.18.0 the wheels can be built from
source on a machine without GPUs and without NVIDIA driver installed.

Docker is the easiest way to build GPU support for TensorFlow since the *host*
machine only requires the
[NVIDIA® driver](https://github.com/NVIDIA/nvidia-docker/wiki/Frequently-Asked-Questions#how-do-i-install-the-nvidia-driver)
Expand Down

0 comments on commit 7d4187e

Please sign in to comment.