Descripción general de Hermetic CUDA

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

  1. 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")
    
  2. Para seleccionar versiones específicas de CUDA y CUDNN herméticos, establece las variables de entorno HERMETIC_CUDA_VERSION y HERMETIC_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"
    
  3. 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.

  4. 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

  1. Crea y envía una solicitud de extracción con diccionarios CUDA_REDIST_JSON_DICT y CUDA_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.

  2. Para cada proyecto de Google ML, crea una solicitud de extracción independiente con HERMETIC_CUDA_VERSION y HERMETIC_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.

  1. Crea archivos cuda_redist.json o cudnn_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 por full_path para las URLs completas y las rutas de acceso locales absolutas que comienzan con file:///.

  2. 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() y cudnn_redist_init_repository(). Ejemplo

    cuda_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.

  1. 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",
    )
    
  2. 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 por full_path para las URLs completas y las rutas de acceso locales absolutas que comiencen con file:///.

  3. 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 llamadas cuda_redist_init_repositories() y 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/"
    )
    

    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:

  1. Borra las llamadas a las reglas del repositorio de CUDA hermético del archivo WORKSPACE del proyecto que depende de XLA.

  2. 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")
    
  3. 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 y TF_CUDNN_VERSION deben constar solo de versiones principales y secundarias (p. ej., 12.3 para CUDA y 9.1 para cuDNN).

  4. 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.