CUDA hermético usa una versión descargable específica de CUDA en lugar de la versión instalada de forma local del usuario. Bazel descargará las distribuciones de CUDA, CUDNN y NCCL, y, luego, usará las bibliotecas y herramientas de CUDA como dependencias en varios destinos de Bazel. Esto permite compilaciones más reproducibles para los proyectos de AA de Google y las versiones compatibles de CUDA.
Versiones compatibles de CUDA y cuDNN herméticas
Las versiones de CUDA compatibles se especifican en el diccionario CUDA_REDIST_JSON_DICT
, third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.
Las versiones de CUDNN compatibles se especifican en el diccionario CUDNN_REDIST_JSON_DICT
, third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.
Los archivos .bazelrc
de proyectos individuales tienen las variables de entorno HERMETIC_CUDA_VERSION
y HERMETIC_CUDNN_VERSION
configuradas en las versiones que se usan de forma predeterminada cuando se especifica --config=cuda
en las opciones de comando de Bazel.
Variables de entorno que controlan las versiones herméticas de CUDA/CUDNN
La variable de entorno HERMETIC_CUDA_VERSION
debe constar de la versión principal, secundaria y de parche de CUDA, p.ej., 12.6.3
.
La variable de entorno HERMETIC_CUDNN_VERSION
debe constar de la versión principal, secundaria y de revisión de cuDNN, p.ej., 9.3.0
.
Existen tres formas de configurar las variables de entorno para los comandos de Bazel:
# Add an entry to your `.bazelrc` file
build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.6.3"
build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
# OR pass it directly to your specific build command
bazel build --config=cuda <target> \
--repo_env=HERMETIC_CUDA_VERSION="12.6.3" \
--repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
# 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.6.3"
export HERMETIC_CUDNN_VERSION="9.3.0"
Si HERMETIC_CUDA_VERSION
y HERMETIC_CUDNN_VERSION
no están presentes, las reglas del repositorio hermético de CUDA/CUDNN buscarán los valores de las variables de entorno TF_CUDA_VERSION
y TF_CUDNN_VERSION
. Esto se hace para la retrocompatibilidad con las reglas del repositorio CUDA/CUDNN no herméticas.
La asignación entre la versión de CUDA y la versión de distribución de NCCL que se descargará se especifica en third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.
Configura CUDA hermético
En el proyecto descendente que depende de XLA, agrega las siguientes líneas al final del archivo
WORKSPACE
:load( "@xla//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( "@xla//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( "@xla//third_party/gpus/cuda/hermetic:cuda_configure.bzl", "cuda_configure", ) cuda_configure(name = "local_config_cuda") load( "@xla//third_party/nccl/hermetic:nccl_redist_init_repository.bzl", "nccl_redist_init_repository", ) nccl_redist_init_repository() load( "@xla//third_party/nccl/hermetic:nccl_configure.bzl", "nccl_configure", ) nccl_configure(name = "local_config_nccl")
Para seleccionar versiones específicas de CUDA y CUDNN herméticos, establece las variables de entorno
HERMETIC_CUDA_VERSION
yHERMETIC_CUDNN_VERSION
, respectivamente. Usa solo versiones compatibles. Puedes establecer las variables de entorno directamente en tu shell o en el archivo.bazelrc
, como se muestra a continuación:build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.6.3" build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.3.0" build:cuda --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90"
Para habilitar CUDA hermético durante la ejecución de pruebas o cuando ejecutes un objeto binario a través de Bazel, asegúrate de agregar la marca
--@local_config_cuda//cuda:include_cuda_libs=true
a tu comando Bazel. Puedes proporcionarlo directamente en un shell o en.bazelrc
:build:cuda --@local_config_cuda//cuda:include_cuda_libs=true
La marca es necesaria para asegurarse de que las dependencias de CUDA se proporcionen correctamente para probar los ejecutables. La marca es falsa de forma predeterminada para evitar la vinculación no deseada de las ruedas de Python publicadas por Google a los objetos binarios de CUDA.
Para aplicar el modo de compatibilidad con versiones posteriores de CUDA, agrega la marca
--@cuda_driver//:enable_forward_compatibility=true
a tu comando bazel. Puedes proporcionarlo directamente en un shell o en.bazelrc
:test:cuda --@cuda_driver//:enable_forward_compatibility=true
El valor predeterminado de la marca es
false
.Cuando se inhabilite el modo de retrocompatibilidad de CUDA, los destinos de Bazel usarán los controladores de modo de usuario y de modo de kernel preinstalados en el sistema.
Cuando se habilita el modo de retrocompatibilidad de CUDA, los destinos de Bazel usarán el controlador de modo de usuario de la redistribución del controlador de CUDA descargado en la caché de Bazel y el controlador de modo de kernel preinstalado en el sistema. Permite habilitar nuevas funciones de CUDA Toolkit mientras se usa un controlador de modo de kernel más antiguo.
El modo de retrocompatibilidad solo debe aplicarse cuando sea apropiado. Consulta la documentación de NVIDIA para obtener detalles.
Actualiza la versión hermética de CUDA/CUDNN
Crea y envía una solicitud de extracción con diccionarios
CUDA_REDIST_JSON_DICT
yCUDNN_REDIST_JSON_DICT
actualizados en third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.Actualiza
CUDA_NCCL_WHEELS
en third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl si es necesario.Actualiza
REDIST_VERSIONS_TO_BUILD_TEMPLATES
en third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl si es necesario.Actualiza
PTX_VERSION_DICT
en third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl si es necesario.Para cada proyecto de Google ML, crea una solicitud de extracción independiente con
HERMETIC_CUDA_VERSION
yHERMETIC_CUDNN_VERSION
actualizados en el archivo.bazelrc
.Las ejecuciones de trabajos previos al envío de PR iniciarán pruebas de bazel y descargarán distribuciones herméticas de CUDA/CUDNN. Verifica que los trabajos previos al envío hayan aprobado antes de enviar la solicitud de cambios.
Para la optimización del tiempo, algunas configuraciones de compilación o prueba usan redistribuciones de
.tar
reflejadas. El archivojson
con información sobre las redistribuciones de.tar
reflejadas se sube un tiempo después de que se actualizanCUDA_REDIST_JSON_DICT
yCUDNN_REDIST_JSON_DICT
. Se pueden descargar estos archivos conwget "https://storage.googleapis.com/mirror.tensorflow.org/developer.download.nvidia.com/compute/cuda/redist/redistrib_<cuda_version>_tar.json"
paraCUDA
ywget "https://storage.googleapis.com/mirror.tensorflow.org/developer.download.nvidia.com/compute/cudnn/redist/redistrib_<cudnn_version>_tar.json"
paraCUDNN
. Después, crea y envía una solicitud de extracción con los diccionariosMIRRORED_TARS_CUDA_REDIST_JSON_DICT
yMIRRORED_TARS_CUDNN_REDIST_JSON_DICT
actualizados en third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.
Cómo apuntar a redistribuciones de CUDA/CUDNN/NCCL en el sistema de archivos local
Puedes usar los directorios locales de CUDA/CUDNN/NCCL como fuente de redistribuciones. Se requieren las siguientes variables de entorno adicionales:
LOCAL_CUDA_PATH
LOCAL_CUDNN_PATH
LOCAL_NCCL_PATH
Ejemplo:
# 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"
La estructura de las carpetas dentro de la carpeta CUDA debe ser la siguiente (como si las redistribuciones archivadas se descomprimieran en un solo lugar):
<LOCAL_CUDA_PATH>/
include/
bin/
lib/
nvvm/
La estructura de las carpetas dentro del directorio CUDNN debe ser la siguiente:
<LOCAL_CUDNN_PATH>
include/
lib/
La estructura de las carpetas dentro del directorio NCCL debe ser la siguiente:
<LOCAL_NCCL_PATH>
include/
lib/
Archivos CUDA/CUDNN personalizados y ruedas de NCCL
Existen tres opciones que permiten el uso de distribuciones personalizadas de CUDA/CUDNN.
Archivos JSON de redistribución personalizados de CUDA/CUDNN
Esta opción permite usar distribuciones personalizadas para todas las dependencias de CUDA/CUDNN en proyectos de Google ML.
Crea archivos
cuda_redist.json
ocudnn_redist.json
.El programa
cuda_redist.json
sigue el siguiente formato:{ "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", } }, }
El programa
cudnn_redist.json
sigue el siguiente formato:{ "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", } } } }
El campo
relative_path
se puede reemplazar porfull_path
para las URLs completas y las rutas de acceso locales absolutas que comienzan confile:///
.En el proyecto descendente que depende de XLA, actualiza la llamada al repositorio JSON de cuda hermético en el archivo
WORKSPACE
. Se permiten vínculos web y rutas de acceso a archivos locales. Ejemplo:_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, )
Si los archivos JSON contienen rutas de acceso relativas a distribuciones, el prefijo de ruta se debe actualizar en las llamadas
cuda_redist_init_repositories()
ycudnn_redist_init_repository()
. Ejemplocuda_redist_init_repositories( cuda_redistributions = CUDA_REDISTRIBUTIONS, cuda_redist_path_prefix = "file:///usr/Downloads/dists/", )
Distribuciones personalizadas de CUDA/CUDNN
Esta opción permite usar distribuciones personalizadas para algunas dependencias de CUDA/CUDNN en proyectos de Google AA.
En el proyecto descendente que depende de XLA, quita las siguientes líneas:
<...> "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", )
En el mismo archivo
WORKSPACE
, crea diccionarios con rutas de distribución.El diccionario con las distribuciones de CUDA sigue el siguiente formato:
_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", } }, }
El diccionario con las distribuciones de CUDNN sigue el siguiente formato:
_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", } } } }
El campo
relative_path
se puede reemplazar porfull_path
para las URLs completas y las rutas de acceso locales absolutas que comienzan confile:///
.En el mismo archivo
WORKSPACE
, pasa los diccionarios creados a la regla del repositorio. Si los diccionarios contienen rutas de acceso relativas a distribuciones, el prefijo de ruta de acceso se debe actualizar en las llamadas acuda_redist_init_repositories()
ycudnn_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/" )
Una combinación de las opciones anteriores
En el siguiente ejemplo, CUDA_REDIST_JSON_DICT
se combina con datos JSON personalizados en _CUDA_JSON_DICT
y CUDNN_REDIST_JSON_DICT
se combina con _CUDNN_JSON_DICT
.
Los datos de las distribuciones en _CUDA_DIST_DICT
anula el contenido del archivo JSON de CUDA resultante, y los datos de las distribuciones en _CUDNN_DIST_DICT
anula el contenido del archivo JSON de CUDNN resultante. Los datos de las ruedas de NCCL se combinan desde CUDA_NCCL_WHEELS
y _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,
)
OBSOLETO: Uso no hermético de CUDA/CUDNN
Aunque el uso no hermético de CUDA/CUDNN dejó de estar disponible, se puede usar para algunos experimentos que actualmente no son compatibles oficialmente (por ejemplo, compilar ruedas en Windows con CUDA).
Estos son los pasos para usar CUDA no hermético instalado de forma local en proyectos de Google ML:
Borra las llamadas a las reglas del repositorio CUDA hermético del archivo
WORKSPACE
del proyecto que depende de XLA.Agrega las llamadas a las reglas del repositorio de CUDA no hermético a la parte inferior del archivo
WORKSPACE
.Para XLA y 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")
Para 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")
Configura las siguientes variables de entorno directamente en tu shell o en el archivo
.bazelrc
como se muestra a continuación: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>
Ten en cuenta que
TF_CUDA_VERSION
yTF_CUDNN_VERSION
deben consistir solo en versiones principales y secundarias (p.ej.,12.3
para CUDA y9.1
para cuDNN).Ahora puedes ejecutar el comando
bazel
para usar CUDA y CUDNN instalados de forma local.Para XLA, no se necesitan cambios en las opciones de comando.
Para JAX, usa la marca
--override_repository=tsl=<tsl_path>
en las opciones del comando de Bazel.Para Tensorflow, usa la marca
--override_repository=local_tsl=<tsl_path>
en las opciones de comando de Bazel.