Présentation de CUDA hermétique

CUDA hermétique utilise une version téléchargeable spécifique de CUDA au lieu de la version installée localement par l'utilisateur. Bazel téléchargera les distributions CUDA, CUDNN et NCCL, puis utilisera les bibliothèques et outils CUDA comme dépendances dans diverses cibles Bazel. Cela permet de créer des builds plus reproductibles pour les projets de ML Google et les versions CUDA compatibles.

Versions CUDA et cuDNN hermétiques compatibles

Les versions CUDA compatibles sont spécifiées dans le dictionnaire CUDA_REDIST_JSON_DICT, third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.

Les versions CUDNN compatibles sont spécifiées dans le dictionnaire CUDNN_REDIST_JSON_DICT, third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.

Les fichiers .bazelrc de projets individuels ont des variables d'environnement HERMETIC_CUDA_VERSION et HERMETIC_CUDNN_VERSION définies sur les versions utilisées par défaut lorsque --config=cuda est spécifié dans les options de commande Bazel.

Variables d'environnement contrôlant les versions CUDA/CUDNN hermétiques

La variable d'environnement HERMETIC_CUDA_VERSION doit se composer de la version majeure, mineure et de correctif de CUDA, par exemple 12.3.2. La variable d'environnement HERMETIC_CUDNN_VERSION doit être une version de CUDNN majeure, mineure et de correctif, par exemple 9.1.1.

Il existe trois façons de définir les variables d'environnement pour les commandes 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 et HERMETIC_CUDNN_VERSION ne sont pas présents, les règles du dépôt CUDA/CUDNN hermétique recherchent les valeurs des variables d'environnement TF_CUDA_VERSION et TF_CUDNN_VERSION. Cela permet de garantir la rétrocompatibilité avec les règles de dépôt CUDA/CUDNN non hermétiques.

Le mappage entre la version CUDA et la version de distribution NCCL à télécharger est spécifié dans third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.

Configurer CUDA hermétique

  1. Dans le projet en aval dépendant de XLA, ajoutez les lignes suivantes en bas du fichier 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. Pour sélectionner des versions spécifiques de CUDA et de cuDNN hermétiques, définissez les variables d'environnement HERMETIC_CUDA_VERSION et HERMETIC_CUDNN_VERSION respectivement. N'utilisez que les versions compatibles. Vous pouvez définir les variables d'environnement directement dans votre interface système ou dans le fichier .bazelrc, comme indiqué ci-dessous:

    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. Pour activer CUDA hermétique lors de l'exécution des tests ou lors de l'exécution d'un binaire via bazel, veillez à ajouter l'indicateur --@local_config_cuda//cuda:include_cuda_libs=true à votre commande bazel. Vous pouvez le fournir directement dans un shell ou dans .bazelrc :

    build:cuda --@local_config_cuda//cuda:include_cuda_libs=true
    

    L'indicateur est nécessaire pour s'assurer que les dépendances CUDA sont correctement fournies pour tester les exécutables. L'indicateur est défini sur "false" par défaut pour éviter le couplage indésirable des paquets Python publiés par Google aux binaires CUDA.

  4. Pour appliquer le mode de compatibilité ascendante CUDA, ajoutez l'indicateur --@cuda_driver//:enable_forward_compatibility=true à votre commande bazel. Vous pouvez le fournir directement dans un shell ou dans .bazelrc :

    test:cuda --@cuda_driver//:enable_forward_compatibility=true
    

    La valeur par défaut de l'indicateur est false.

    Lorsque le mode de compatibilité avant CUDA est désactivé, les cibles Bazel utilisent les pilotes en mode utilisateur et en mode noyau préinstallés sur le système.

    Lorsque le mode de compatibilité ascendante CUDA est activé, les cibles Bazel utilisent le pilote en mode utilisateur de la redistribution du pilote CUDA téléchargée dans le cache Bazel et le pilote en mode noyau préinstallé sur le système. Il permet d'activer de nouvelles fonctionnalités du kit d'outils CUDA tout en utilisant un ancien pilote en mode kernel.

    Le mode de compatibilité ascendante ne doit être appliqué que lorsque cela est approprié. Pour en savoir plus, consultez la documentation NVIDIA.

Mettre à niveau la version CUDA/CUDNN hermétique

  1. Créez et envoyez une demande de tirage avec des dictionnaires CUDA_REDIST_JSON_DICT et CUDA_REDIST_JSON_DICT mis à jour dans third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.

    Si nécessaire, mettez à jour CUDA_NCCL_WHEELS dans third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.

    Mettez à jour REDIST_VERSIONS_TO_BUILD_TEMPLATES dans third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl si nécessaire.

  2. Pour chaque projet de ML Google, créez une demande de tirage distincte avec des HERMETIC_CUDA_VERSION et HERMETIC_CUDNN_VERSION mis à jour dans le fichier .bazelrc.

    Les exécutions de tâches de présoumission de PR lancent des tests bazel et téléchargent des distributions CUDA/CUDNN hermétiques. Vérifiez que les tâches de présoumission ont réussi avant d'envoyer la demande de pull.

Pointage vers les redistributions CUDA/CUDNN/NCCL sur le système de fichiers local

Vous pouvez utiliser les répertoires CUDA/CUDNN/NCCL locaux comme source de redistributions. Les variables d'environnement supplémentaires suivantes sont requises :

LOCAL_CUDA_PATH
LOCAL_CUDNN_PATH
LOCAL_NCCL_PATH

Exemple :

# 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 structure des dossiers dans le répertoire CUDA doit être la suivante (comme si les redistributions archivées étaient dégroupées en un seul endroit):

<LOCAL_CUDA_PATH>/
    include/
    bin/
    lib/
    nvvm/

La structure des dossiers dans le répertoire CUDNN doit être la suivante :

<LOCAL_CUDNN_PATH>
    include/
    lib/

La structure des dossiers dans le répertoire NCCL doit être la suivante :

<LOCAL_NCCL_PATH>
    include/
    lib/

Archives CUDA/CUDNN personnalisées et roues NCCL

Trois options permettent d'utiliser des distributions CUDA/CUDNN personnalisées.

Fichiers JSON de redistribution CUDA/CUDNN personnalisés

Cette option permet d'utiliser des distributions personnalisées pour toutes les dépendances CUDA/CUDNN dans les projets de ML Google.

  1. Créez des fichiers cuda_redist.json et/ou cudnn_redist.json.

    Les émissions cuda_redist.json suivent le format ci-dessous :

    {
       "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",
          }
       },
    }
    

    Les émissions cudnn_redist.json suivent le format ci-dessous :

    {
       "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",
             }
          }
       }
    }
    

    Le champ relative_path peut être remplacé par full_path pour les URL complètes et les chemins d'accès locaux absolus commençant par file:///.

  2. Dans le projet en aval dépendant de XLA, mettez à jour l'appel du dépôt JSON cuda hermétique dans le fichier WORKSPACE. Les liens Web et les chemins d'accès aux fichiers locaux sont autorisés. Exemple :

    _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 les fichiers JSON contiennent des chemins d'accès relatifs aux distributions, le préfixe de chemin d'accès doit être mis à jour dans les appels cuda_redist_init_repositories() et cudnn_redist_init_repository(). Exemple

    cuda_redist_init_repositories(
       cuda_redistributions = CUDA_REDISTRIBUTIONS,
       cuda_redist_path_prefix = "file:///usr/Downloads/dists/",
    )
    

Distributions CUDA/CUDNN personnalisées

Cette option permet d'utiliser des distributions personnalisées pour certaines dépendances CUDA/CUDNN dans les projets de ML Google.

  1. Dans le projet en aval qui dépend de XLA, supprimez les lignes ci-dessous:

    <...>
       "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. Dans le même fichier WORKSPACE, créez des dictionnaires avec des chemins de distribution.

    Le dictionnaire des distributions CUDA suit le format ci-dessous :

    _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",
          }
       },
    }
    

    Le dictionnaire avec les distributions CUDNN suit le format ci-dessous :

    _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",
             }
          }
       }
    }
    

    Le champ relative_path peut être remplacé par full_path pour les URL complètes et les chemins d'accès locaux absolus commençant par file:///.

  3. Dans le même fichier WORKSPACE, transmettez les dictionnaires créés à la règle du dépôt. Si les dictionnaires contiennent des chemins d'accès relatifs aux distributions, le préfixe de chemin d'accès doit être mis à jour dans les appels cuda_redist_init_repositories() et 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/"
    )
    

    Combinaison des options ci-dessus

Dans l'exemple ci-dessous, CUDA_REDIST_JSON_DICT est fusionné avec les données JSON personnalisées dans _CUDA_JSON_DICT, et CUDNN_REDIST_JSON_DICT est fusionné avec _CUDNN_JSON_DICT.

Les données de distribution dans _CUDA_DIST_DICT remplacent le contenu du fichier JSON CUDA généré, et les données de distribution dans _CUDNN_DIST_DICT remplacent le contenu du fichier JSON CUDNN généré. Les données sur les roues NCCL sont fusionnées à partir de CUDA_NCCL_WHEELS et _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,
)

OBSOLÈTE: Utilisation non hermétique de CUDA/CUDNN

Bien que l'utilisation de CUDA/CUDNN non hermétique soit obsolète, elle peut être utilisée pour certains tests qui ne sont actuellement pas officiellement compatibles (par exemple, créer des roues sous Windows avec CUDA).

Voici les étapes à suivre pour utiliser un CUDA non hermétique installé localement dans des projets Google ML:

  1. Supprimez les appels aux règles de dépôt CUDA hermétique à partir du fichier WORKSPACE du projet qui dépend de XLA.

  2. Ajoutez les appels aux règles de dépôt CUDA non hermétiques en bas du fichier WORKSPACE.

    Pour XLA et 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")
    

    Pour 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. Définissez les variables d'environnement suivantes directement dans votre interface système ou dans le fichier .bazelrc, comme indiqué ci-dessous:

    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>
    

    Notez que TF_CUDA_VERSION et TF_CUDNN_VERSION ne doivent contenir que des versions majeures et mineures (par exemple, 12.3 pour CUDA et 9.1 pour cuDNN).

  4. Vous pouvez maintenant exécuter la commande bazel pour utiliser CUDA et CUDNN installés en local.

    Pour XLA, aucune modification des options de commande n'est nécessaire.

    Pour JAX, utilisez l'option --override_repository=tsl=<tsl_path> dans les options de commande Bazel.

    Pour Tensorflow, utilisez l'option --override_repository=local_tsl=<tsl_path> dans les options de commande Bazel.