密閉型 CUDA の概要

Hermetic CUDA は、ユーザーがローカルにインストールした CUDA ではなく、特定のダウンロード可能なバージョンの CUDA を使用します。Bazel は、CUDA、CUDNN、NCCL のディストリビューションをダウンロードし、さまざまな Bazel ターゲットで CUDA のライブラリとツールを依存関係として使用します。これにより、Google ML プロジェクトとサポートされている CUDA バージョンのビルドの再現性が向上します。

サポートされているシーリングされた CUDA、CUDNN のバージョン

サポートされている CUDA バージョンは、CUDA_REDIST_JSON_DICT ディクショナリ third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl で指定されています。

サポートされている CUDNN バージョンは、CUDNN_REDIST_JSON_DICT 辞書（third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl）で規定されています。

個々のプロジェクトの .bazelrc ファイルでは、Bazel コマンド オプションで --config=cuda が指定されている場合、HERMETIC_CUDA_VERSION と HERMETIC_CUDNN_VERSION の環境変数がデフォルトで使用されるバージョンに設定されています。

密閉型の CUDA/CUDNN バージョンを制御する環境変数

HERMETIC_CUDA_VERSION 環境変数は、メジャー バージョン、マイナー バージョン、パッチ バージョンの CUDA バージョンで構成する必要があります（例: 12.3.2）。HERMETIC_CUDNN_VERSION 環境変数は、メジャー バージョン、マイナー バージョン、パッチ CUDNN バージョンで構成する必要があります（例: 9.1.1）。

Bazel コマンドの環境変数を設定するには、次の 3 つの方法があります。

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

HERMETIC_CUDA_VERSION と HERMETIC_CUDNN_VERSION が存在しない場合、気密性の高い CUDA / CUDNNN リポジトリ ルールは TF_CUDA_VERSION と TF_CUDNN_VERSION の環境変数値を検索します。これは、非シーリング CUDA / CUDNN リポジトリ ルールとの下位互換性を確保するために行われます。

ダウンロードする CUDA バージョンと NCCL ディストリビューション バージョンのマッピングは、third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl で指定されています。

密閉型の CUDA を設定する

1. XLA に依存するダウンストリーム プロジェクトで、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. 完全な CUDA と CUDNN の特定のバージョンを選択するには、それぞれ HERMETIC_CUDA_VERSION 環境変数と HERMETIC_CUDNN_VERSION 環境変数を設定します。サポートされているバージョンのみを使用してください。次のように、シェルまたは .bazelrc ファイルで環境変数を直接設定できます。

   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. テスト実行中、または bazel 経由でバイナリを実行するときに密閉型の CUDA を有効にするには、bazel コマンドに --@local_config_cuda//cuda:include_cuda_libs=true フラグを追加する必要があります。シェルまたは .bazelrc で直接指定できます。

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

   このフラグは、実行可能ファイルをテストするための CUDA 依存関係が適切に提供されることを確認するために必要です。Google がリリースした Python wheel と CUDA バイナリとの意図しない結合を避けるため、このフラグはデフォルトでは false です。

4. CUDA 上位互換モードを適用するには、bazel コマンドに --@cuda_driver//:enable_forward_compatibility=true フラグを追加します。シェルまたは .bazelrc で直接指定できます。

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

   デフォルトのフラグ値は false です。

   CUDA の将来の互換モードが無効になっている場合、Bazel ターゲットはシステムにプリインストールされているユーザーモード ドライバとカーネルモード ドライバを使用します。

   CUDA 上位互換モードが有効になっている場合、Bazel ターゲットは、Bazel キャッシュにダウンロードされた CUDA ドライバ再配布のユーザーモード ドライバと、システムにプリインストールされているカーネルモード ドライバを使用します。これにより、古いカーネルモード ドライバを使用しながら、新しい CUDA Toolkit 機能を有効にできます。

   上位互換性モードは、適切な場合にのみ適用する必要があります。詳細については、NVIDIA のドキュメントをご覧ください。

気密性の高い CUDA / CUDNN バージョンをアップグレードする

1. third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl にある、更新された CUDA_REDIST_JSON_DICT と CUDA_REDIST_JSON_DICT ディクショナリを使用して pull リクエストを作成して送信します。

   必要に応じて、third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl の CUDA_NCCL_WHEELS を更新します。

   必要に応じて、third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl の REDIST_VERSIONS_TO_BUILD_TEMPLATES を更新します。

2. Google ML プロジェクトごとに、.bazelrc ファイルで HERMETIC_CUDA_VERSION と HERMETIC_CUDNN_VERSION を更新した個別のプル リクエストを作成します。

   PR の presubmit ジョブの実行により、bazel テストが起動され、完全な CUDA / CUDNND ディストリビューションがダウンロードされます。PR を送信する前に、presubmit ジョブに合格したことを確認します。

ローカル ファイル システム上の CUDA / CUDNN / NCC 再配布を指す

ローカルの CUDA/CUDNN/NCCL ディレクトリを再配布のソースとして使用できます。次の追加の環境変数が必要です。

LOCAL_CUDA_PATH
LOCAL_CUDNN_PATH
LOCAL_NCCL_PATH

例:

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

CUDA ディレクトリ内のフォルダの構造は、アーカイブされた再配布が 1 か所に解凍された場合と同じになります。

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

CUDNN ディレクトリ内のフォルダの構造は次のようになります。

<LOCAL_CUDNN_PATH>
    include/
    lib/

NCCL ディレクトリ内のフォルダの構造は次のようになります。

<LOCAL_NCCL_PATH>
    include/
    lib/

カスタム CUDA / cuDNN アーカイブと NCCL ホイール

カスタム CUDA / CUDNN ディストリビューションを使用するには、次の 3 つの方法があります。

カスタム CUDA/CUDNN 再配布 JSON ファイル

このオプションを使用すると、Google ML プロジェクトのすべての CUDA / CUDNN 依存関係にカスタム ディストリビューションを使用できます。

1. cuda_redist.json ファイルまたは cudnn_redist.json ファイルを作成します。

   cuda_redist.json show の形式は次のとおりです。

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

   cudnn_redist.json show の形式は次のとおりです。

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

   完全な URL と file:/// で始まる絶対ローカルパスでは、relative_path フィールドを full_path に置き換えることができます。

2. XLA に依存するダウンストリーム プロジェクトで、WORKSPACE ファイルの完全な cuda JSON リポジトリ呼び出しを更新します。ウェブリンクとローカル ファイルパスの両方を使用できます。例:

   _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,
   )
   

   JSON ファイルにディストリビューションへの相対パスが含まれている場合は、cuda_redist_init_repositories() 呼び出しと cudnn_redist_init_repository() 呼び出しでパス接頭辞を更新する必要があります。例

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

カスタム CUDA / CUDNN ディストリビューション

このオプションを使用すると、Google ML プロジェクトの一部の CUDA / CUDNN 依存関係にカスタム ディストリビューションを使用できます。

1. XLA に依存するダウンストリーム プロジェクトで、次の行を削除します。

   <...>
      "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. 同じ WORKSPACE ファイルで、ディストリビューション パスを含むディクショナリを作成します。

   CUDA ディストリビューションを含むディクショナリは、次の形式で表示されます。

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

   CUDNN ディストリビューションを含む辞書は、次の形式で表示されます。

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

   relative_path フィールドは、file:/// で始まる完全な URL と絶対ローカルパスの場合は full_path に置き換えることができます。

3. 同じ WORKSPACE ファイルで、作成した辞書をリポジトリ ルールに渡します。辞書にディストリビューションへの相対パスが含まれている場合は、cuda_redist_init_repositories() と 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/"
   )
   

   上記のオプションの組み合わせ

次の例では、CUDA_REDIST_JSON_DICT は _CUDA_JSON_DICT のカスタム JSON データと統合され、CUDNN_REDIST_JSON_DICT は _CUDNN_JSON_DICT と統合されます。

_CUDA_DIST_DICT の分布データは、生成される CUDA JSON ファイルの内容をオーバーライドし、_CUDNN_DIST_DICT の分布データは、生成される CUDNN JSON ファイルの内容をオーバーライドします。NCCL ホイールデータは、CUDA_NCCL_WHEELS と _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,
)

非推奨: 非密閉型の CUDA / CUDNN の使用

非完全な CUDA / CUDNN の使用は非推奨ですが、現在公式にサポートされていない一部のテスト（CUDA で Windows にホイールをビルドするなど）で使用される場合があります。

Google ML プロジェクトにローカルにインストールされた非完全な CUDA を使用する手順は次のとおりです。

1. XLA に依存するプロジェクトの WORKSPACE ファイルから、密閉型の CUDA リポジトリ ルールの呼び出しを削除します。

2. 非完全な CUDA リポジトリ ルールの呼び出しを WORKSPACE ファイルの末尾に追加します。

   XLA と 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")
   

   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. 次のように、シェルまたは .bazelrc ファイルで次の環境変数を直接設定します。

   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 と TF_CUDNN_VERSION は、メジャー バージョンとマイナー バージョンのみで構成する必要があります（例: CUDA の場合は 12.3、CUDNN の場合は 9.1）。

4. これで、bazel コマンドを実行して、ローカルにインストールされた CUDA と CUDNN を使用できるようになりました。

   XLA の場合、コマンド オプションを変更する必要はありません。

   JAX の場合は、Bazel コマンド オプションで --override_repository=tsl=<tsl_path> フラグを使用します。

   TensorFlow の場合は、Bazel コマンド オプションで --override_repository=local_tsl=<tsl_path> フラグを使用します。