OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2026-01-15 12:15:41 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Enable using custom hermetic NCCL version.

Browse Source 
The NCCL version can be chosen via `HERMETIC_NCCL_VERSION` env var.

See docs [here](https://github.com/google-ml-infra/rules_ml_toolchain/blob/main/gpu/README.md#environment-variables-controlling-the-hermetic-cudacudnnnvshmem-versions).

PiperOrigin-RevId: 846788706
This commit is contained in:
Yulia Baturina
2025-12-19 10:39:43 -08:00
committed by TensorFlower Gardener
parent 6af9ca2d9c
commit 843206ca81
4 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
.bazelrc
View File

@@ -291,10 +291,11 @@ common:mkl_aarch64_threadpool -c opt
# This is an alias for the mkl_aarch64_threadpool build.
common:mkl_aarch64 --config=mkl_aarch64_threadpool
# Default CUDA, CUDNN and NVSHMEM versions.
# Default CUDA, CUDNN, NCCL and NVSHMEM versions.
common:cuda_version --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
common:cuda_version --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
common:cuda_version --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5"
common:cuda_version --repo_env=HERMETIC_NCCL_VERSION="2.27.7"
# CUDA: This config refers to building CUDA op kernels with nvcc.
common:cuda --repo_env TF_NEED_CUDA=1

6
WORKSPACE
View File

@@ -21,10 +21,10 @@ tf_http_archive(
# Details: https://github.com/google-ml-infra/rules_ml_toolchain
tf_http_archive(
name = "rules_ml_toolchain",
sha256 = "53905ede50e3eebc782266e20e9b9ac1d7166ef68b877bea593d3600dcfe03e6",
strip_prefix = "rules_ml_toolchain-a1ff84835e407b41eef5fd1a865a23748c294db6",
sha256 = "1c2c530a054e9e8b3c811ec21ed8a687fc865bec3abbc8ff65beb829b1d67ae4",
strip_prefix = "rules_ml_toolchain-6734d2a174bf29e731d3f473743d1cc1a86100c3",
urls = tf_mirror_urls(
"https://github.com/google-ml-infra/rules_ml_toolchain/archive/a1ff84835e407b41eef5fd1a865a23748c294db6.tar.gz",
"https://github.com/google-ml-infra/rules_ml_toolchain/archive/6734d2a174bf29e731d3f473743d1cc1a86100c3.tar.gz",
),
)

2
ci/official/utilities/code_check_full.bats
View File

@@ -216,6 +216,7 @@ EOF
--@local_config_cuda//cuda:include_cuda_libs=false \
--repo_env=HERMETIC_CUDA_VERSION="12.3.2" \
--repo_env=HERMETIC_CUDNN_VERSION="8.9.7.29" \
--repo_env=HERMETIC_NCCL_VERSION="2.27.7" \
"somepath(//tensorflow/tools/pip_package:wheel, " \
"@local_config_cuda//cuda:cudart + "\
"@local_config_cuda//cuda:cudart + "\
@@ -240,6 +241,7 @@ EOF
--@local_config_cuda//cuda:include_cuda_libs=false \
--repo_env=HERMETIC_CUDA_VERSION="12.3.2" \
--repo_env=HERMETIC_CUDNN_VERSION="8.9.7.29" \
--repo_env=HERMETIC_NCCL_VERSION="2.27.7" \
--define framework_shared_object=false \
"somepath(//tensorflow/tools/pip_package:wheel, " \
"@local_config_cuda//cuda:cudart + "\

6
tensorflow/workspace0.bzl
View File

@@ -108,10 +108,10 @@ def workspace():
# Details: https://github.com/google-ml-infra/rules_ml_toolchain
tf_http_archive(
name = "rules_ml_toolchain",
sha256 = "53905ede50e3eebc782266e20e9b9ac1d7166ef68b877bea593d3600dcfe03e6",
strip_prefix = "rules_ml_toolchain-a1ff84835e407b41eef5fd1a865a23748c294db6",
sha256 = "1c2c530a054e9e8b3c811ec21ed8a687fc865bec3abbc8ff65beb829b1d67ae4",
strip_prefix = "rules_ml_toolchain-6734d2a174bf29e731d3f473743d1cc1a86100c3",
urls = tf_mirror_urls(
"https://github.com/google-ml-infra/rules_ml_toolchain/archive/a1ff84835e407b41eef5fd1a865a23748c294db6.tar.gz",
"https://github.com/google-ml-infra/rules_ml_toolchain/archive/6734d2a174bf29e731d3f473743d1cc1a86100c3.tar.gz",
),
)