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

  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 da CUDA hermética e da CUDNN, defina as variáveis de ambiente HERMETIC_CUDA_VERSION e HERMETIC_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"
    
  3. 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.

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

  1. Crie e envie uma solicitação de envio com dicionários CUDA_REDIST_JSON_DICT e CUDA_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.

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

  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 os caminhos locais absolutos que começam com file:///.

  2. 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() 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 do CUDA/CUDNN em projetos de ML do Google.

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

  3. 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 chamadas cuda_redist_init_repositories() e 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/"
    )
    

    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:

  1. Exclua chamadas para regras de repositório CUDA herméticas do arquivo WORKSPACE do projeto dependente do XLA.

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