密封 CUDA 使用特定的可下载 CUDA 版本，而不是用户本地安装的 CUDA。Bazel 会下载 CUDA、CUDNN 和 NCCL 发行版，然后在各种 Bazel 目标中使用 CUDA 库和工具作为依赖项。这为 Google ML 项目和受支持的 CUDA 版本提供了更可重现的 build。
支持的封闭 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
在依赖于 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")如需选择封闭的 CUDA 和 CUDNN 的特定版本，请分别设置
HERMETIC_CUDA_VERSION和HERMETIC_CUDNN_VERSION环境变量。仅使用受支持的版本。您可以直接在 shell 或.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"如需在测试执行期间或通过 bazel 运行二进制文件时启用封闭的 CUDA，请务必向 bazel 命令添加
--@local_config_cuda//cuda:include_cuda_libs=true标志。您可以直接在 shell 或.bazelrc中提供它：build:cuda --@local_config_cuda//cuda:include_cuda_libs=true此标志用于确保为测试可执行文件正确提供 CUDA 依赖项。该标志默认为 false，以避免 Google 发布的 Python wheel 与 CUDA 二进制文件之间发生不必要的关联。
如需强制执行 CUDA 向前兼容模式，请在 bazel 命令中添加
--@cuda_driver//:enable_forward_compatibility=true标志。您可以直接在 shell 或.bazelrc中提供该值：test:cuda --@cuda_driver//:enable_forward_compatibility=true默认标志值为
false。停用 CUDA 向前兼容性模式后，Bazel 目标将使用系统上预安装的用户模式和内核模式驱动程序。
启用 CUDA 向前兼容性模式后，Bazel 目标将使用下载到 Bazel 缓存中的 CUDA 驱动程序重新分发中的用户模式驱动程序，以及系统上预安装的内核模式驱动程序。它允许在使用旧版内核模式驱动程序时启用新的 CUDA 工具包功能。
仅应在适当的情况下强制执行向前兼容模式 - 如需了解详情，请参阅 NVIDIA 文档。
升级密封 CUDA/CUDNN 版本
使用 third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl 中的更新后的
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。对于每个 Google 机器学习项目，请在
.bazelrc文件中使用更新后的HERMETIC_CUDA_VERSION和HERMETIC_CUDNN_VERSION创建一个单独的拉取请求。PR 提交前作业执行将启动 bazel 测试并下载密封的 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 目录中的文件夹结构应如下所示（就好像已将归档的重新分发内容解压缩到一个位置一样）：
<LOCAL_CUDA_PATH>/
include/
bin/
lib/
nvvm/
CUDNN 目录内的文件夹结构应如下所示：
<LOCAL_CUDNN_PATH>
include/
lib/
NCCL 目录中的文件夹结构应如下所示：
<LOCAL_NCCL_PATH>
include/
lib/
自定义 CUDA/CUDNN 归档和 NCCL 车轮
有三种方法可以使用自定义 CUDA/CUDNN 发行版。
自定义 CUDA/CUDNN 重新分发 JSON 文件
此选项允许对 Google 机器学习项目中的所有 CUDA/CUDNN 依赖项使用自定义分发。
创建
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", } } } }对于以
file:///开头的完整网址和绝对本地路径，relative_path字段可以替换为full_path。在依赖于 XLA 的下游项目中，更新
WORKSPACE文件中的密封 cuda JSON 代码库调用。允许网页链接和本地文件路径。示例：_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 分发
此选项允许为 Google ML 项目中的某些 CUDA/CUDNN 依赖项使用自定义分发。
在依赖于 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", )在同一
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", } } } }对于完整网址以及以
file:///开头的绝对本地路径，relative_path字段可替换为full_path。在同一
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 与 _CUDA_JSON_DICT 中的自定义 JSON 数据合并，而 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 用法已废弃，但仍可用于目前不受官方支持的一些实验（例如，使用 CUDA 在 Windows 上构建轮）。
如需在 Google ML 项目中使用本地安装的非容器化 CUDA，请按以下步骤操作：
从依赖于 XLA 的项目的
WORKSPACE文件中删除对密封 CUDA 代码库规则的调用。将对非容器化 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")对于 TensorFlow：
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")直接在 shell 或
.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）。现在，您可以运行
bazel命令来使用本地安装的 CUDA 和 CUDNN。对于 XLA，无需更改命令选项。
对于 JAX，请在 Bazel 命令选项中使用
--override_repository=tsl=<tsl_path>标志。对于 Tensorflow，请在 Bazel 命令选项中使用
--override_repository=local_tsl=<tsl_path>标志。