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á distribuciones CUDA, CUDNN y NCCL y, luego, usará bibliotecas y herramientas CUDA como dependencias en varios destinos de Bazel. Esto permite compilaciones más reproducibles para proyectos de AA de Google y versiones de CUDA compatibles.
Versiones herméticas compatibles con CUDA y CUDNN
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 revisión de CUDA, p. ej., 12.3.2
.
La variable de entorno HERMETIC_CUDNN_VERSION
debe constar de la versión principal, secundaria y de revisión de cuDNN, p. ej., 9.1.1
.
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.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"
Si HERMETIC_CUDA_VERSION
y HERMETIC_CUDNN_VERSION
no están presentes, las reglas herméticas del repositorio 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.
Configurar CUDA hermético
En el proyecto downstream que depende de XLA, agrega las siguientes líneas al final del archivo
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")
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.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"
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 anterior.
El modo de retrocompatibilidad solo se debe aplicar 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
yCUDA_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.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 las tareas de envío previo se hayan aprobado antes de enviar la solicitud de cambios.
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 CUDA/CUDNN personalizados
Esta opción permite usar distribuciones personalizadas para todas las dependencias de CUDA/CUDNN en proyectos de AA de Google.
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 ML.
En el proyecto downstream 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 distribuciones de CUDNN muestra que 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 comiencen confile:///
.En el mismo archivo
WORKSPACE
, pasa los diccionarios creados a la regla del repositorio. Si los diccionarios contienen rutas de acceso relativas a las distribuciones, el prefijo de ruta de acceso debe actualizarse en las llamadascuda_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/" )
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 ruedas de NCCL se combinan a partir de 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 de 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("@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")
Para 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")
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 constar solo de 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 del comando de Bazel.