Обзор герметичного CUDA

Hermetic CUDA использует специальную загружаемую версию CUDA вместо локально установленной пользователем CUDA. Bazel загрузит дистрибутивы CUDA, CUDNN и NCCL, а затем будет использовать библиотеки и инструменты CUDA в качестве зависимостей в различных целях Bazel. Это обеспечивает более воспроизводимые сборки для проектов Google ML и поддерживаемых версий CUDA.

Поддерживаемые герметичные версии CUDA, CUDNN.

Поддерживаемые версии CUDA указаны в словаре CUDA_REDIST_JSON_DICT , Third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl .

Поддерживаемые версии CUDNN указаны в словаре CUDNN_REDIST_JSON_DICT , Third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl .

В файлах .bazelrc отдельных проектов переменные среды HERMETIC_CUDA_VERSION , HERMETIC_CUDNN_VERSION установлены в версии, используемые по умолчанию, когда в параметрах команды Bazel указан --config=cuda .

Переменные среды, управляющие герметичными версиями CUDA/CUDNN.

Переменная среды HERMETIC_CUDA_VERSION должна состоять из основной, дополнительной и исправленной версии CUDA, например 12.3.2 . Переменная среды HERMETIC_CUDNN_VERSION должна состоять из основной, дополнительной и исправленной версии CUDNN, например 9.1.1 .

Три способа установки переменных среды для команд Bazel:

# Add an entry to your `.bazelrc` file
build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.3.2"
build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.1.1"

# OR pass it directly to your specific build command
bazel build --config=cuda <target> \
--repo_env=HERMETIC_CUDA_VERSION="12.3.2" \
--repo_env=HERMETIC_CUDNN_VERSION="9.1.1"

# If .bazelrc doesn't have corresponding entries and the environment variables
# are not passed to bazel command, you can set them globally in your shell:
export HERMETIC_CUDA_VERSION="12.3.2"
export HERMETIC_CUDNN_VERSION="9.1.1"

Если HERMETIC_CUDA_VERSION и HERMETIC_CUDNN_VERSION отсутствуют, правила герметичного репозитория CUDA/CUDNN будут искать значения переменных среды TF_CUDA_VERSION и TF_CUDNN_VERSION . Это сделано для обратной совместимости с правилами негерметичного репозитория CUDA/CUDNN.

Сопоставление между версией CUDA и загружаемой версией дистрибутива NCCL указано в файле Third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.

Настроить герметичный CUDA

  1. В нижестоящем проекте, зависящем от XLA, добавьте следующие строки в конец файла WORKSPACE :

    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")
    
  2. Чтобы выбрать конкретные версии герметичных CUDA и CUDNN, установите переменные среды HERMETIC_CUDA_VERSION и HERMETIC_CUDNN_VERSION соответственно. Используйте только поддерживаемые версии. Вы можете установить переменные среды непосредственно в вашей оболочке или в файле .bazelrc , как показано ниже:

    build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.3.2"
    build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.1.1"
    build:cuda --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90"
    
  3. Чтобы включить герметичный CUDA во время выполнения теста или при запуске двоичного файла через bazel, обязательно добавьте флаг --@local_config_cuda//cuda:include_cuda_libs=true в свою команду bazel. Вы можете предоставить его либо непосредственно в оболочке, либо в .bazelrc :

    build:cuda --@local_config_cuda//cuda:include_cuda_libs=true
    

    Флаг необходим, чтобы убедиться, что зависимости CUDA правильно предоставлены для тестирования исполняемых файлов. По умолчанию этот флаг имеет значение false, чтобы избежать нежелательного связывания выпущенных Google колес Python с двоичными файлами CUDA.

  4. Чтобы обеспечить режим прямой совместимости CUDA, добавьте флаг --@cuda_driver//:enable_forward_compatibility=true к вашей команде bazel. Вы можете предоставить его либо непосредственно в оболочке, либо в .bazelrc :

    test:cuda --@cuda_driver//:enable_forward_compatibility=true
    

    Значение флага по умолчанию — false .

    Если режим прямой совместимости CUDA отключен, целевые объекты Bazel будут использовать драйверы пользовательского режима и режима ядра, предварительно установленные в системе.

    Когда включен режим прямой совместимости CUDA, целевые объекты Bazel будут использовать драйвер пользовательского режима из перераспределенного драйвера CUDA, загруженного в кэш Bazel, и драйвер режима ядра, предварительно установленный в системе. Это позволяет включить новые функции CUDA Toolkit при использовании старого драйвера режима ядра.

    Режим прямой совместимости следует включать только тогда, когда это необходимо — подробности см. в документации NVIDIA .

Обновление герметичной версии CUDA/CUDNN

  1. Создайте и отправьте запрос на включение с обновленными словарями CUDA_REDIST_JSON_DICT , CUDA_REDIST_JSON_DICT в файле Third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl .

    При необходимости обновите CUDA_NCCL_WHEELS в файле Third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl .

    При необходимости обновите REDIST_VERSIONS_TO_BUILD_TEMPLATES в файле Third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl .

  2. Для каждого проекта Google ML создайте отдельный запрос на включение с обновленными HERMETIC_CUDA_VERSION и HERMETIC_CUDNN_VERSION в файле .bazelrc .

    Выполнение заданий PR перед отправкой запустит базельские тесты и загрузит герметичные дистрибутивы CUDA/CUDNN. Перед отправкой PR убедитесь, что задания предварительной отправки прошли успешно.

Указание на перераспределения CUDA/CUDNN/NCCL в локальной файловой системе.

Вы можете использовать локальные каталоги CUDA/CUDNN/NCCL в качестве источника перераспределения. Требуются следующие дополнительные переменные среды:

LOCAL_CUDA_PATH
LOCAL_CUDNN_PATH
LOCAL_NCCL_PATH

Пример:

# Add an entry to your `.bazelrc` file
build:cuda --repo_env=LOCAL_CUDA_PATH="/foo/bar/nvidia/cuda"
build:cuda --repo_env=LOCAL_CUDNN_PATH="/foo/bar/nvidia/cudnn"
build:cuda --repo_env=LOCAL_NCCL_PATH="/foo/bar/nvidia/nccl"

# OR pass it directly to your specific build command
bazel build --config=cuda <target> \
--repo_env=LOCAL_CUDA_PATH="/foo/bar/nvidia/cuda" \
--repo_env=LOCAL_CUDNN_PATH="/foo/bar/nvidia/cudnn" \
--repo_env=LOCAL_NCCL_PATH="/foo/bar/nvidia/nccl"

# If .bazelrc doesn't have corresponding entries and the environment variables
# are not passed to bazel command, you can set them globally in your shell:
export LOCAL_CUDA_PATH="/foo/bar/nvidia/cuda"
export LOCAL_CUDNN_PATH="/foo/bar/nvidia/cudnn"
export LOCAL_NCCL_PATH="/foo/bar/nvidia/nccl"

Структура папок внутри CUDA dir должна быть следующей (как если бы архивные редистрибутивы были распакованы в одно место):

<LOCAL_CUDA_PATH>/
    include/
    bin/
    lib/
    nvvm/

Структура папок внутри каталога CUDNN должна быть следующей:

<LOCAL_CUDNN_PATH>
    include/
    lib/

Структура папок внутри каталога NCCL должна быть следующей:

<LOCAL_NCCL_PATH>
    include/
    lib/

Пользовательские архивы CUDA/CUDNN и колеса NCCL.

Существует три варианта, которые позволяют использовать собственные дистрибутивы CUDA/CUDNN.

Пользовательские JSON-файлы перераспределения CUDA/CUDNN

Эта опция позволяет использовать собственные дистрибутивы для всех зависимостей CUDA/CUDNN в проектах Google ML.

  1. Создайте файлы cuda_redist.json и/или cudnn_redist.json .

    cuda_redist.json показывает следующий формат:

    {
       "cuda_cccl": {
          "linux-x86_64": {
             "relative_path": "cuda_cccl-linux-x86_64-12.4.99-archive.tar.xz",
          },
          "linux-sbsa": {
             "relative_path": "cuda_cccl-linux-sbsa-12.4.99-archive.tar.xz",
          }
       },
    }
    

    cudnn_redist.json показывает следующий формат:

    {
       "cudnn": {
          "linux-x86_64": {
             "cuda12": {
             "relative_path": "cudnn/linux-x86_64/cudnn-linux-x86_64-9.0.0.312_cuda12-archive.tar.xz",
             }
          },
          "linux-sbsa": {
             "cuda12": {
             "relative_path": "cudnn/linux-sbsa/cudnn-linux-sbsa-9.0.0.312_cuda12-archive.tar.xz",
             }
          }
       }
    }
    

    Поле relative_path можно заменить на full_path для полных URL-адресов и абсолютных локальных путей, начинающихся с file:/// .

  2. В нижестоящем проекте, зависящем от XLA, обновите вызов герметичного хранилища JSON cuda в файле WORKSPACE . Разрешены как веб-ссылки, так и пути к локальным файлам. Пример:

    _CUDA_JSON_DICT = {
       "12.4.0": [
          "file:///home/user/Downloads/redistrib_12.4.0_updated.json",
       ],
    }
    
    _CUDNN_JSON_DICT = {
       "9.0.0": [
          "https://developer.download.nvidia.com/compute/cudnn/redist/redistrib_9.0.0.json",
       ],
    }
    
    cuda_json_init_repository(
       cuda_json_dict = _CUDA_JSON_DICT,
       cudnn_json_dict = _CUDNN_JSON_DICT,
    )
    

    Если файлы JSON содержат относительные пути к дистрибутивам, префикс пути следует обновить в вызовах cuda_redist_init_repositories() и cudnn_redist_init_repository() . Пример

    cuda_redist_init_repositories(
       cuda_redistributions = CUDA_REDISTRIBUTIONS,
       cuda_redist_path_prefix = "file:///usr/Downloads/dists/",
    )
    

Пользовательские дистрибутивы CUDA/CUDNN

Эта опция позволяет использовать собственные дистрибутивы для некоторых зависимостей CUDA/CUDNN в проектах Google ML.

  1. В нижестоящем проекте, зависящем от XLA, удалите строки ниже:

    <...>
       "CUDA_REDIST_JSON_DICT",
    <...>
       "CUDNN_REDIST_JSON_DICT",
    <...>
    
    cuda_json_init_repository(
       cuda_json_dict = CUDA_REDIST_JSON_DICT,
       cudnn_json_dict = CUDNN_REDIST_JSON_DICT,
    )
    
    load(
       "@cuda_redist_json//:distributions.bzl",
       "CUDA_REDISTRIBUTIONS",
       "CUDNN_REDISTRIBUTIONS",
    )
    
  2. В этом же файле WORKSPACE создайте словари с путями распространения.

    Словарь с дистрибутивами CUDA имеет следующий формат:

    _CUSTOM_CUDA_REDISTRIBUTIONS = {
       "cuda_cccl": {
          "linux-x86_64": {
             "relative_path": "cuda_cccl-linux-x86_64-12.4.99-archive.tar.xz",
          },
          "linux-sbsa": {
             "relative_path": "cuda_cccl-linux-sbsa-12.4.99-archive.tar.xz",
          }
       },
    }
    

    Словарь с дистрибутивами CUDNN имеет следующий формат:

    _CUSTOM_CUDNN_REDISTRIBUTIONS = {
       "cudnn": {
          "linux-x86_64": {
             "cuda12": {
             "relative_path": "cudnn/linux-x86_64/cudnn-linux-x86_64-9.0.0.312_cuda12-archive.tar.xz",
             }
          },
          "linux-sbsa": {
             "cuda12": {
             "relative_path": "cudnn/linux-sbsa/cudnn-linux-sbsa-9.0.0.312_cuda12-archive.tar.xz",
             }
          }
       }
    }
    

    Поле relative_path можно заменить на full_path для полных URL-адресов и абсолютных локальных путей, начинающихся с file:/// .

  3. В том же файле WORKSPACE передайте созданные словари в правило репозитория. Если словари содержат относительные пути к дистрибутивам, префикс пути должен быть обновлен в вызовах cuda_redist_init_repositories() и cudnn_redist_init_repository() .

    cuda_redist_init_repositories(
       cuda_redistributions = _CUSTOM_CUDA_REDISTRIBUTIONS,
       cuda_redist_path_prefix = "file:///home/usr/Downloads/dists/",
    )
    
    cudnn_redist_init_repository(
       cudnn_redistributions = _CUSTOM_CUDNN_REDISTRIBUTIONS,
       cudnn_redist_path_prefix = "file:///home/usr/Downloads/dists/cudnn/"
    )
    

    Комбинация вариантов выше

В приведенном ниже примере CUDA_REDIST_JSON_DICT объединяется с пользовательскими данными JSON в _CUDA_JSON_DICT , а CUDNN_REDIST_JSON_DICT объединяется с _CUDNN_JSON_DICT .

Данные о распределениях в _CUDA_DIST_DICT переопределяют содержимое результирующего файла CUDA JSON, а данные о распределениях в _CUDNN_DIST_DICT переопределяют содержимое результирующего файла CUDNN JSON. Данные колес NCCL объединяются из CUDA_NCCL_WHEELS и _NCCL_WHEEL_DICT .

load(
    //third_party/gpus/cuda/hermetic:cuda_redist_versions.bzl",
    "CUDA_REDIST_PATH_PREFIX",
    "CUDA_NCCL_WHEELS",
    "CUDA_REDIST_JSON_DICT",
    "CUDNN_REDIST_PATH_PREFIX",
    "CUDNN_REDIST_JSON_DICT",
)

_CUDA_JSON_DICT = {
   "12.4.0": [
      "file:///usr/Downloads/redistrib_12.4.0_updated.json",
   ],
}

_CUDNN_JSON_DICT = {
   "9.0.0": [
      "https://developer.download.nvidia.com/compute/cudnn/redist/redistrib_9.0.0.json",
   ],
}

cuda_json_init_repository(
   cuda_json_dict = CUDA_REDIST_JSON_DICT | _CUDA_JSON_DICT,
   cudnn_json_dict = CUDNN_REDIST_JSON_DICT | _CUDNN_JSON_DICT,
)

load(
   "@cuda_redist_json//:distributions.bzl",
   "CUDA_REDISTRIBUTIONS",
   "CUDNN_REDISTRIBUTIONS",
)

load(
   "//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
   "cuda_redist_init_repositories",
   "cudnn_redist_init_repository",
)

_CUDA_DIST_DICT = {
   "cuda_cccl": {
      "linux-x86_64": {
            "relative_path": "cuda_cccl-linux-x86_64-12.4.99-archive.tar.xz",
      },
      "linux-sbsa": {
            "relative_path": "cuda_cccl-linux-sbsa-12.4.99-archive.tar.xz",
      },
   },
   "libcusolver": {
      "linux-x86_64": {
            "full_path": "file:///usr/Downloads/dists/libcusolver-linux-x86_64-11.6.0.99-archive.tar.xz",
      },
      "linux-sbsa": {
         "relative_path": "libcusolver-linux-sbsa-11.6.0.99-archive.tar.xz",
      },
   },
}

_CUDNN_DIST_DICT = {
   "cudnn": {
      "linux-x86_64": {
            "cuda12": {
               "relative_path": "cudnn-linux-x86_64-9.0.0.312_cuda12-archive.tar.xz",
            },
      },
      "linux-sbsa": {
            "cuda12": {
               "relative_path": "cudnn-linux-sbsa-9.0.0.312_cuda12-archive.tar.xz",
            },
      },
   },
}

cudnn_redist_init_repositories(
   cuda_redistributions = CUDA_REDISTRIBUTIONS | _CUDA_DIST_DICT,
   cuda_redist_path_prefix = "file:///usr/Downloads/dists/",
)

cudnn_redist_init_repository(
   cudnn_redistributions = CUDNN_REDISTRIBUTIONS | _CUDNN_DIST_DICT,
   cudnn_redist_path_prefix = "file:///usr/Downloads/dists/cudnn/"
)

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

_NCCL_WHEEL_DICT = {
   "12.4.0": {
      "x86_64-unknown-linux-gnu": {
            "url": "https://files.pythonhosted.org/packages/38/00/d0d4e48aef772ad5aebcf70b73028f88db6e5640b36c38e90445b7a57c45/nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl",
      },
   },
}

nccl_redist_init_repository(
   cuda_nccl_wheels = CUDA_NCCL_WHEELS | _NCCL_WHEEL_DICT,
)

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

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

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

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

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

    Для XLA и JAX:

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

    Для Тензорного потока:

    load("@local_tsl//third_party/gpus:cuda_configure.bzl", "cuda_configure")
    cuda_configure(name = "local_config_cuda")
    load("@local_tsl//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.