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 versões CUDA compatíveis.
Versões herméticas do CUDA e da cuDNN com suporte
As versões do CUDA compatíveis são especificadas no dicionário CUDA_REDIST_JSON_DICT
, 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 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
de 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
.
Há três maneiras de definir as 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 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 da CUDA hermética e da CUDNN, defina as variáveis de ambiente
HERMETIC_CUDA_VERSION
eHERMETIC_CUDNN_VERSION
, respectivamente. Use apenas versões compatíveis. Você pode 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 o CUDA hermético durante a execução de testes ou ao executar um binário via bazel, adicione a sinalização
--@local_config_cuda//cuda:include_cuda_libs=true
ao comando do Bazel. Você pode fornecê-lo diretamente em um shell ou 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 o acoplamento indesejado de rodas Python lançadas pelo Google com 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 fornecer o valor diretamente em um shell ou em.bazelrc
:test:cuda --@cuda_driver//:enable_forward_compatibility=true
O valor padrão da flag é
false
.Quando o modo de compatibilidade com encaminhamento de CUDA estiver desativado, os destinos do Bazel vão usar os drivers do modo do usuário e do modo kernel pré-instalados no sistema.
Quando o modo de compatibilidade futura do CUDA está ativado, as metas do Bazel usam o driver do modo de usuário da redistribuição do driver CUDA transferido para o cache do Bazel e o driver do modo Kernel pré-instalado no sistema. Ele permite ativar novos recursos do Toolkit CUDA ao usar drivers de modo de kernel mais antigos.
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 dicionários
CUDA_REDIST_JSON_DICT
eCUDA_REDIST_JSON_DICT
atualizados 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.Atualize
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
HERMETIC_CUDA_VERSION
eHERMETIC_CUDNN_VERSION
atualizados no arquivo.bazelrc
.As execuções de job de pré-envio de PR vão iniciar testes do Bazel e fazer o download de distribuições herméticas do CUDA/CUDNN. Verifique se os jobs de pré-envio foram aprovados antes de enviar a PR.
Como apontar para redistribuições CUDA/CUDNN/NCCL no sistema de arquivos local
É possível usar as pastas locais CUDA/CUDNN/NCCL como uma 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 no diretório CUDNN precisa 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 do 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 os caminhos locais absolutos que começam comfile:///
.No projeto dependente da XLA, atualize a chamada de repositório JSON hermético do cuda no arquivo
WORKSPACE
. Links da Web e caminhos de arquivos locais são permitidos. 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 contiverem caminhos relativos para distribuições, o prefixo do caminho precisará ser atualizado nas chamadas
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 do CUDA/CUDNN em projetos de ML do Google.
No projeto downstream dependente da 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 formato abaixo:
_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 que começam comfile:///
.No mesmo arquivo
WORKSPACE
, transmita os dicionários criados para a regra do repositório. Se os dicionários contiverem caminhos relativos às distribuições, o prefixo do caminho precisará ser atualizado em chamadascuda_redist_init_repositories()
ecudnn_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/" )
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
é mesclado 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 tenha sido descontinuado, ele pode ser usado para alguns experimentos atualmente sem suporte (por exemplo, criação de rodas no Windows com CUDA).
Estas são as etapas para usar o CUDA não hermético instalado localmente em projetos do Google ML:
Exclua chamadas para regras de repositório CUDA herméticas do arquivo
WORKSPACE
do projeto dependente do XLA.Adicione as chamadas às regras do repositório CUDA não herméticas na parte de baixo do arquivo
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 a XLA, não é necessário fazer mudanças nas opções de comando.
Para JAX, use a sinalização
--override_repository=tsl=<tsl_path>
nas opções de comando do Bazel.Para o Tensorflow, use a flag
--override_repository=local_tsl=<tsl_path>
nas opções de comando do Bazel.