O CUDA hermético usa uma versão específica para download do CUDA em vez da versão instalada localmente do usuário. O Bazel vai fazer o download das distribuições CUDA, CUDNN e NCCL e, em seguida, usar as bibliotecas e ferramentas CUDA como dependências em várias metas do Bazel. Isso permite builds mais reproduzíveis para projetos de ML do Google e CUDA compatível. mais recentes.
Versões de CUDA hermética e CUDNN compatíveis
As versões do CUDA compatíveis são especificadas em CUDA_REDIST_JSON_DICT
.
dicionário
third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.
As versões do CUDNN com suporte são especificadas no dicionário
CUDNN_REDIST_JSON_DICT
,
third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.
Os arquivos .bazelrc
de projetos individuais têm as variáveis de ambiente HERMETIC_CUDA_VERSION
e
HERMETIC_CUDNN_VERSION
definidas para as versões usadas por
padrão quando --config=cuda
é especificado nas opções de comando do Bazel.
Variáveis de ambiente que controlam as versões herméticas da CUDA/CUDNN
A variável de ambiente HERMETIC_CUDA_VERSION
precisa consistir na versão principal, secundária e
do patch do CUDA, por exemplo, 12.3.2
.
A variável de ambiente HERMETIC_CUDNN_VERSION
precisa consistir na versão principal, secundária e
de patch da cuDNN, por exemplo, 9.1.1
.
Três maneiras de definir variáveis de ambiente para comandos do 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"
Se HERMETIC_CUDA_VERSION
e HERMETIC_CUDNN_VERSION
não estiverem presentes, as
regras do repositório hermético CUDA/CUDNN vão procurar os valores das variáveis de ambiente TF_CUDA_VERSION
e
TF_CUDNN_VERSION
. Isso é feito para oferecer compatibilidade
com versões anteriores com regras de repositório CUDA/CUDNN não herméticas.
O mapeamento entre a versão do CUDA e a versão da distribuição do NCCL a ser baixada é especificado em third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl
Configurar o CUDA hermético
No projeto downstream dependente da XLA, adicione as linhas abaixo à parte de baixo do arquivo
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 selecionar versões específicas de CUDA e CUDNN herméticas, defina o Variáveis de ambiente
HERMETIC_CUDA_VERSION
eHERMETIC_CUDNN_VERSION
respectivamente. Use apenas versões compatíveis. É possível definir as variáveis de ambiente diretamente no shell ou no arquivo.bazelrc
, conforme mostrado abaixo: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 ativar a CUDA hermética durante a execução do teste ou ao executar um binário pelo bazel, adicione a flag
--@local_config_cuda//cuda:include_cuda_libs=true
ao comando bazel. É possível fornecê-los diretamente em um shell em.bazelrc
:build:cuda --@local_config_cuda//cuda:include_cuda_libs=true
A flag é necessária para garantir que as dependências do CUDA sejam fornecidas corretamente para testar executáveis. A flag é falsa por padrão para evitar acoplamentos indesejados de rodas Python lançadas pelo Google a binários CUDA.
Para aplicar o modo de compatibilidade para frente do CUDA, adicione a flag
--@cuda_driver//:enable_forward_compatibility=true
ao comando bazel. Você pode fornecê-lo diretamente em um shell ou em.bazelrc
:test:cuda --@cuda_driver//:enable_forward_compatibility=true
O modo de compatibilidade futura só deve ser aplicado quando for apropriado. Consulte a documentação da NVIDIA para mais detalhes.
Fazer upgrade da versão hermética do CUDA/CUDNN
Crie e envie uma solicitação de envio com o
CUDA_REDIST_JSON_DICT
atualizado.CUDA_REDIST_JSON_DICT
dicionários em third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl.Atualize
CUDA_NCCL_WHEELS
em third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl se necessário.Atualizar
REDIST_VERSIONS_TO_BUILD_TEMPLATES
em third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl se necessário.Para cada projeto do Google ML, crie uma solicitação de envio separada com atualizações
HERMETIC_CUDA_VERSION
eHERMETIC_CUDNN_VERSION
no arquivo.bazelrc
.As execuções do job de pré-envio de RP vão iniciar testes do Bazel e fazer o download de herméticos. distribuições CUDA/CUDNN. Verifique se os jobs de pré-envio foram aprovados antes de enviar a PR.
Apontar para redistribuições CUDA/CUDNN/NCCL no sistema de arquivos local
É possível usar os diretórios locais do CUDA/CUDNN/NCCL como fonte de redistribuições. As seguintes variáveis de ambiente extras são obrigatórias:
LOCAL_CUDA_PATH
LOCAL_CUDNN_PATH
LOCAL_NCCL_PATH
Exemplo:
# 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"
A estrutura das pastas dentro do diretório CUDA precisa ser a seguinte (como se as redistribuições arquivadas tivessem sido descompactadas em um só lugar):
<LOCAL_CUDA_PATH>/
include/
bin/
lib/
nvvm/
A estrutura das pastas dentro do diretório CUDNN deve ser a seguinte:
<LOCAL_CUDNN_PATH>
include/
lib/
A estrutura das pastas no diretório NCCL deve ser a seguinte:
<LOCAL_NCCL_PATH>
include/
lib/
Arquivos CUDA/CUDNN e rodas NCCL personalizados
Há três opções que permitem o uso de distribuições CUDA/CUDNN personalizadas.
Arquivos JSON de redistribuição CUDA/CUDNN personalizados
Essa opção permite usar distribuições personalizadas para todas as dependências CUDA/CUDNN em projetos de ML do Google.
Crie arquivos
cuda_redist.json
e/oucudnn_redist.json
.O
cuda_redist.json
segue o formato abaixo:{ "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", } }, }
O
cudnn_redist.json
segue o formato abaixo:{ "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", } } } }
O campo
relative_path
pode ser substituído porfull_path
para os URLs completos e caminhos locais absolutos começando comfile:///
.No projeto dependente da XLA, atualize a chamada de repositório JSON hermético do cuda no arquivo
WORKSPACE
. Tanto os links da Web quanto os caminhos de arquivos locais são permitido. Exemplo:_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, )
Se os arquivos JSON tiverem caminhos relativos às distribuições, o prefixo do caminho deve ser atualizados em
cuda_redist_init_repositories()
ecudnn_redist_init_repository()
. Exemplocuda_redist_init_repositories( cuda_redistributions = CUDA_REDISTRIBUTIONS, cuda_redist_path_prefix = "file:///usr/Downloads/dists/", )
Distribuições personalizadas de CUDA/CUDNN
Essa opção permite usar distribuições personalizadas para algumas dependências CUDA/CUDNN em projetos de ML do Google.
No projeto downstream dependente do XLA, remova as linhas abaixo:
<...> "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", )
No mesmo arquivo
WORKSPACE
, crie dicionários com caminhos de distribuição.O dicionário com distribuições CUDA segue o formato abaixo:
_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", } }, }
O dicionário com distribuições CUDNN mostra o seguinte 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", } } } }
O campo
relative_path
pode ser substituído porfull_path
para os URLs completos e caminhos locais absolutos começando comfile:///
.No mesmo arquivo
WORKSPACE
, transmita os dicionários criados para a regra do repositório. Se os dicionários tiverem caminhos relativos às distribuições, o caminho prefixo precisa ser atualizado emcuda_redist_init_repositories()
ecudnn_redist_init_repository()
chamadas.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/" )
Combinação das opções acima
No exemplo abaixo, CUDA_REDIST_JSON_DICT
é mesclado com dados JSON personalizados em
_CUDA_JSON_DICT
e CUDNN_REDIST_JSON_DICT
foram mesclados com
_CUDNN_JSON_DICT
.
Os dados de distribuição em _CUDA_DIST_DICT
substituem o conteúdo do arquivo JSON CUDA resultante, e os dados de distribuição em _CUDNN_DIST_DICT
substituem o conteúdo do arquivo JSON CUDNN resultante. Os dados de rodas NCCL são mesclados de
CUDA_NCCL_WHEELS
e _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 não hermético de CUDA/CUDNN
Embora o uso não hermético da CUDA/CUDNN esteja obsoleto, ele pode ser usado para alguns experimentos não são oficialmente aceitos (por exemplo, rodas de construção no Windows com CUDA).
Confira as etapas para usar o CUDA não hermético instalado localmente em projetos do Google ML:
Exclusão de chamadas para regras de repositório CUDA herméticas do arquivo
WORKSPACE
do projeto dependente da XLA.Adicione as chamadas para regras de repositório CUDA não herméticas na parte inferior da
WORKSPACE
.Para XLA e 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 o 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")
Defina as seguintes variáveis de ambiente diretamente no shell ou no arquivo
.bazelrc
, conforme mostrado abaixo: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>
TF_CUDA_VERSION
eTF_CUDNN_VERSION
precisam consistir apenas de versões principais e secundárias (por exemplo,12.3
para CUDA e9.1
para CUDNN).Agora é possível executar o comando
bazel
para usar o CUDA e o CUDNN instalados localmente.Para XLA, nenhuma alteração nas opções de comando é necessária.
Para JAX, use a flag
--override_repository=tsl=<tsl_path>
nas opções de comando do Bazel.Para o Tensorflow, use a sinalização
--override_repository=local_tsl=<tsl_path>
na Opções de comando do Bazel.