Visão geral do CUDA hermético

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

  1. 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")
    
  2. Para selecionar versões específicas de CUDA e CUDNN herméticas, defina o Variáveis de ambiente HERMETIC_CUDA_VERSION e HERMETIC_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"
    
  3. 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.

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

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

  2. Para cada projeto do Google ML, crie uma solicitação de envio separada com atualizações HERMETIC_CUDA_VERSION e HERMETIC_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.

  1. Crie arquivos cuda_redist.json e/ou cudnn_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 por full_path para os URLs completos e caminhos locais absolutos começando com file:///.

  2. 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() e cudnn_redist_init_repository(). Exemplo

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

  1. 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",
    )
    
  2. 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 por full_path para os URLs completos e caminhos locais absolutos começando com file:///.

  3. 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 em cuda_redist_init_repositories() e cudnn_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:

  1. Exclusão de chamadas para regras de repositório CUDA herméticas do arquivo WORKSPACE do projeto dependente da XLA.

  2. 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")
    
  3. 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 e TF_CUDNN_VERSION precisam consistir apenas de versões principais e secundárias (por exemplo, 12.3 para CUDA e 9.1 para CUDNN).

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