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á 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

  1. 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")
    
  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.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"
    
  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 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

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

  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 los trabajos previos al envío hayan aprobado antes de enviar la solicitud de cambios.

  3. Para la optimización del tiempo, algunas configuraciones de compilación o prueba usan redistribuciones de .tar reflejadas. El archivo json con información sobre las redistribuciones de .tar reflejadas se sube un tiempo después de que se actualizan CUDA_REDIST_JSON_DICT y CUDNN_REDIST_JSON_DICT. Se pueden descargar estos archivos con wget "https://storage.googleapis.com/mirror.tensorflow.org/developer.download.nvidia.com/compute/cuda/redist/redistrib_<cuda_version>_tar.json" para CUDA y wget "https://storage.googleapis.com/mirror.tensorflow.org/developer.download.nvidia.com/compute/cudnn/redist/redistrib_<cudnn_version>_tar.json" para CUDNN. Después, crea y envía una solicitud de extracción con los diccionarios MIRRORED_TARS_CUDA_REDIST_JSON_DICT y MIRRORED_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.

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

  1. 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",
    )
    
  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 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 por full_path para las URLs completas y las rutas de acceso locales absolutas que comienzan 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 distribuciones, el prefijo de ruta de acceso se debe actualizar en las llamadas a 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/"
    )
    

    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:

  1. Borra las llamadas a las reglas del repositorio 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("@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")
    
  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 consistir solo en 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 de comando de Bazel.