Обзор герметичных CUDA, CUDNN, NCCL и NVSHMEM

Обзор и примеры использования приведены в проекте rules_ml_toolchain .

УСТАРЕЛО: Негерметичное использование CUDA/CUDNN

Хотя негерметичное использование CUDA/CUDNN устарело, его можно использовать для некоторых экспериментов, которые в настоящее время официально не поддерживаются (например, сборка колес в Windows с помощью CUDA).

Вот шаги для использования негерметичной CUDA, установленной локально в проектах Google ML:

1. Удалить вызовы правил герметичного репозитория CUDA из файла WORKSPACE проекта, зависящего от XLA.

2. Добавьте вызовы правил негерметичного репозитория CUDA в конец файла WORKSPACE .

   Для XLA и JAX:

   load("@xla//third_party/gpus:cuda_configure.bzl", "cuda_configure")
   cuda_configure(name = "local_config_cuda")
   load("@xla//third_party/nccl:nccl_configure.bzl", "nccl_configure")
   nccl_configure(name = "local_config_nccl")
   

   Для Tensorflow:

   load("@local_xla//third_party/gpus:cuda_configure.bzl", "cuda_configure")
   cuda_configure(name = "local_config_cuda")
   load("@local_xla//third_party/nccl:nccl_configure.bzl", "nccl_configure")
   nccl_configure(name = "local_config_nccl")
   

3. Задайте следующие переменные среды непосредственно в оболочке или в файле .bazelrc , как показано ниже:

   build:cuda --action_env=TF_CUDA_VERSION=<locally installed cuda version>
   build:cuda --action_env=TF_CUDNN_VERSION=<locally installed cudnn version>
   build:cuda --action_env=TF_CUDA_COMPUTE_CAPABILITIES=<CUDA compute capabilities>
   build:cuda --action_env=LD_LIBRARY_PATH=<CUDA/CUDNN libraries folder locations divided by “:” sign>
   build:cuda --action_env=CUDA_TOOLKIT_PATH=<preinstalled CUDA folder location>
   build:cuda --action_env=TF_CUDA_PATHS=<preinstalled CUDA/CUDNN folder locations divided by “,” sign>
   build:cuda --action_env=NCCL_INSTALL_PATH=<preinstalled NCCL library folder location>
   

   Обратите внимание, что TF_CUDA_VERSION и TF_CUDNN_VERSION должны состоять только из основных и дополнительных версий (например, 12.3 для CUDA и 9.1 для CUDNN).

4. Теперь вы можете запустить команду bazel для использования локально установленных CUDA и CUDNN.

   Для XLA никаких изменений в параметрах команды не требуется.

   Для JAX используйте флаг --override_repository=tsl=<tsl_path> в параметрах команды Bazel.

   Для Tensorflow используйте флаг --override_repository=local_tsl=<tsl_path> в параметрах команды Bazel.