密封式 CUDA 會使用特定可下載的 CUDA 版本，而非使用者在本機安裝的 CUDA。Bazel 會下載 CUDA、CUDNN 和 NCCL 發布項目，然後使用 CUDA 程式庫和工具做為各種 Bazel 目標的依附元件。如此一來，Google 機器學習專案和支援的 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 密封
在依附 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環境變數。請只使用支援的版本。您可以直接在殼層或.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標記。您可以直接在殼層或.bazelrc中提供：build:cuda --@local_config_cuda//cuda:include_cuda_libs=true這個標記可確保系統正確提供 CUDA 依附元件，以便測試可執行檔。根據預設，這個旗標為 false，可避免 Google 發布的 Python 輪子與 CUDA 二進位檔產生不必要的耦合。
如要強制執行 CUDA 前瞻相容性模式，請在 bazel 指令中新增
--@cuda_driver//:enable_forward_compatibility=true標記。您可以直接在殼層或.bazelrc中提供：test:cuda --@cuda_driver//:enable_forward_compatibility=true預設標記值為
false。停用 CUDA 向前相容性模式後，Bazel 目標會使用系統預先安裝的使用者模式和核心模式驅動程式。
啟用 CUDA 前瞻相容性模式時，Bazel 目標會使用從 CUDA 驅動程式重新分配下載的使用者模式驅動程式 (已下載到 Bazel 快取，並預先安裝在核心模式驅動程式)。可在使用舊版核心模式驅動程式的同時，啟用新的 CUDA 工具包功能。
只有在適當情況下，才應強制執行向前相容性模式。詳情請參閱 NVIDIA 說明文件。
升級密封的 CUDA/CUDNN 版本
建立並提交包含更新版
CUDA_REDIST_JSON_DICT的拉取要求，CUDA_REDIST_JSON_DICT字典位於 third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl。視需要更新 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 ML 專案，在
.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 dir 內的資料夾結構應如下所示：
<LOCAL_CUDNN_PATH>
include/
lib/
NCCL 目錄內的資料夾結構應如下所示：
<LOCAL_NCCL_PATH>
include/
lib/
自訂 CUDA/CUDNN 封存檔和 NCCL 輪
有三種選項可讓您使用自訂 CUDA/CUDNN 發布版本。
自訂 CUDA/CUDNN 重新分配 JSON 檔案
這個選項可讓您為 Google ML 專案中的所有 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 存放區規則的呼叫。在
WORKSPACE檔案底部新增對非密封 CUDA 存放區規則的呼叫。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")如要直接在殼層或
.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>旗標。