ภาพรวมของ CUDA แบบแยกต่างหาก

Hermetic CUDA ใช้ CUDA เวอร์ชันที่ดาวน์โหลดได้โดยเฉพาะแทน CUDA ที่ผู้ใช้ติดตั้งในเครื่อง Bazel จะดาวน์โหลด CUDA, CUDNN และ NCCL distribution จากนั้นจะใช้ไลบรารีและเครื่องมือ CUDA เป็น Dependency ในเป้าหมาย Bazel ต่างๆ ซึ่งจะช่วยให้สร้างบิลด์ที่ซ้ำกันสำหรับโปรเจ็กต์ Google ML และ CUDA เวอร์ชันที่รองรับได้มากขึ้น

เวอร์ชัน CUDA, CUDNN แบบแยกต่างหากที่รองรับ

เวอร์ชัน CUDA ที่รองรับระบุไว้ในพจนานุกรม CUDA_REDIST_JSON_DICT, third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl

เวอร์ชัน CUDNN ที่รองรับจะระบุไว้ใน CUDNN_REDIST_JSON_DICT dictionary, third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl

ไฟล์ .bazelrc ของโปรเจ็กต์แต่ละรายการมีการตั้งค่าตัวแปรสภาพแวดล้อม HERMETIC_CUDA_VERSION, HERMETIC_CUDNN_VERSION เป็นเวอร์ชันที่ใช้โดยค่าเริ่มต้นเมื่อระบุ --config=cuda ในตัวเลือกคำสั่ง Bazel

ตัวแปรสภาพแวดล้อมที่ควบคุมเวอร์ชัน CUDA/CUDNN แบบรวมกัน

ตัวแปรสภาพแวดล้อม HERMETIC_CUDA_VERSION ควรประกอบด้วยเวอร์ชัน CUDA หลัก รอง และแพตช์ เช่น 12.3.2 ตัวแปรสภาพแวดล้อม HERMETIC_CUDNN_VERSION ควรประกอบด้วยเวอร์ชันหลัก รอง และแพตช์ CUDNN เช่น 9.1.1

3 วิธีในการตั้งค่าตัวแปรสภาพแวดล้อมสําหรับคําสั่ง 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"

หากไม่มี HERMETIC_CUDA_VERSION และ HERMETIC_CUDNN_VERSION กฎที่เก็บข้อมูล CUDA/CUDNN แบบปิดผนึกจะค้นหาค่าตัวแปรสภาพแวดล้อม 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. หากต้องการเปิดใช้ CUDA แบบปิดผนึกระหว่างการทดสอบหรือเมื่อเรียกใช้ไบนารีผ่าน balen ให้เพิ่ม Flag --@local_config_cuda//cuda:include_cuda_libs=true ลงในคำสั่ง bazel คุณสามารถระบุข้อมูลนี้ในเชลล์โดยตรงหรือใน .bazelrc ก็ได้

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

    จำเป็นต้องมีแฟล็กเพื่อตรวจสอบว่าทรัพยากร Dependency ของ CUDA ได้รับการจัดเตรียมอย่างถูกต้องเพื่อทดสอบไฟล์ปฏิบัติการ Flag จะเป็น False โดยค่าเริ่มต้นเพื่อหลีกเลี่ยงการเชื่อมโยงที่ไม่พึงประสงค์ของ Wheel ของ Python ที่ Google เผยแพร่กับไบนารี CUDA

  4. หากต้องการบังคับใช้โหมดความเข้ากันได้แบบย้อนหลังของ CUDA ให้เพิ่ม Flag --@cuda_driver//:enable_forward_compatibility=true ลงในคำสั่ง bazel คุณสามารถระบุข้อมูลนี้ในเชลล์หรือใน .bazelrc โดยตรง ดังนี้

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

    ค่าเริ่มต้นของ Flag คือ false

    เมื่อปิดใช้โหมดความเข้ากันได้กับการส่งต่อของ CUDA เป้าหมาย Bazel จะใช้โหมดผู้ใช้และไดรเวอร์โหมดเคอร์เนลที่ติดตั้งไว้ล่วงหน้าในระบบ

    เมื่อเปิดใช้โหมดความเข้ากันได้กับการส่งต่อของ CUDA เป้าหมาย Bazel จะใช้ไดรเวอร์โหมดผู้ใช้จากการจัดสรรไดรเวอร์ CUDA ใหม่ที่ดาวน์โหลดลงในแคช Bazel และไดรเวอร์โหมดเคอร์เนลที่ติดตั้งไว้ล่วงหน้าในระบบ ซึ่งช่วยให้เปิดใช้ฟีเจอร์ CUDA Toolkit ใหม่ได้ขณะใช้ไดรเวอร์โหมดเคอร์เนลรุ่นเก่า

    คุณควรบังคับใช้โหมดความเข้ากันได้แบบย้อนหลังเฉพาะในกรณีที่เหมาะสมเท่านั้น ดูรายละเอียดได้ในเอกสารประกอบของ NVIDIA

อัปเกรดเวอร์ชัน CUDA/CUDNN แบบแยกต่างหาก

  1. สร้างและส่งคำขอดึงข้อมูลที่มีพจนานุกรม CUDA_REDIST_JSON_DICT ที่อัปเดตแล้วใน CUDA_REDIST_JSON_DICT ใน 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 ใน third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl หากจำเป็น

  2. สำหรับโปรเจ็กต์ Google ML แต่ละรายการ ให้สร้างคำขอพุลแยกต่างหากที่มี HERMETIC_CUDA_VERSION และ HERMETIC_CUDNN_VERSION ที่อัปเดตแล้วในไฟล์ .bazelrc

    การดำเนินการงานก่อนส่ง PR จะเปิดการทดสอบ Bazel และดาวน์โหลดการแจกจ่าย CUDA/CUDNN แบบปิด ตรวจสอบว่างานที่ส่งล่วงหน้าผ่านแล้วก่อนส่ง PR

การชี้ไปยังการเผยแพร่ CUDA/CUDNN/NCCL ซ้ำในระบบไฟล์ในเครื่อง

คุณสามารถใช้ไดเรกทอรี 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 ควรมีลักษณะดังต่อไปนี้ (เหมือนกับว่ามีการแตกไฟล์การแจกจ่ายที่เก็บถาวรไว้ในที่เดียว)

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

โครงสร้างของโฟลเดอร์ภายในไดเรกทอรี CUDNN ควรมีลักษณะดังนี้

<LOCAL_CUDNN_PATH>
    include/
    lib/

โครงสร้างของโฟลเดอร์ภายในไดเรกทอรี NCCL ควรมีลักษณะดังนี้

<LOCAL_NCCL_PATH>
    include/
    lib/

ที่เก็บถาวรของ CUDA/CUDNN ที่กำหนดเองและวงล้อ NCCL

มี 3 ตัวเลือกที่อนุญาตให้ใช้การแจกจ่าย CUDA/CUDNN ที่กําหนดเอง

ไฟล์ JSON ของการแจกจ่าย CUDA/CUDNN ที่กำหนดเอง

ตัวเลือกนี้อนุญาตให้ใช้การกระจายที่กำหนดเองสำหรับทรัพยากร Dependency ของ CUDA/CUDNN ทั้งหมดในโปรเจ็กต์ Google ML

  1. สร้างไฟล์ cuda_redist.json และ/หรือ cudnn_redist.json

    cuda_redist.json รายการต้องเป็นไปตามรูปแบบด้านล่าง

    {
   "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 รายการต้องเป็นไปตามรูปแบบด้านล่าง

    {
   "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 ด้วย full_path สำหรับ URL แบบเต็มและเส้นทางภายในแบบสัมบูรณ์ที่ขึ้นต้นด้วย file:///

  2. ในโปรเจ็กต์ดาวน์สตรีมที่ขึ้นอยู่กับ XLA ให้อัปเดตการเรียกใช้ที่เก็บ JSON ของ cuda แบบปิดผนึกในไฟล์ WORKSPACE เราอนุญาตให้ใช้ทั้งเว็บลิงก์และเส้นทางไฟล์ในเครื่อง ตัวอย่าง

    _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 ที่กําหนดเอง

ตัวเลือกนี้อนุญาตให้ใช้การกระจายที่กำหนดเองสำหรับทรัพยากร Dependency ของ CUDA/CUDNN บางรายการในโปรเจ็กต์ Google ML

  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 ด้วย full_path สำหรับ URL แบบเต็มและเส้นทางภายในแบบสัมบูรณ์ที่ขึ้นต้นด้วย file:///

  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 ผสานกับข้อมูล JSON ที่กำหนดเองใน _CUDA_JSON_DICT และ CUDNN_REDIST_JSON_DICT ผสานกับ _CUDNN_JSON_DICT

ข้อมูลการเผยแพร่ใน _CUDA_DIST_DICT จะลบล้างเนื้อหาของไฟล์ JSON CUDA ที่ได้ และข้อมูลการเผยแพร่ใน _CUDNN_DIST_DICT จะลบล้างเนื้อหาของไฟล์ JSON CUDNN ที่ได้ ระบบจะรวมข้อมูลล้อ 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 แบบไม่ปิดผนึกแล้ว แต่อาจใช้สำหรับการทดสอบบางอย่างที่ยังไม่รองรับอย่างเป็นทางการในขณะนี้ (เช่น การสร้างล้อใน Windows ด้วย CUDA)

ขั้นตอนในการใช้ CUDA แบบไม่ปิดผนึกที่ติดตั้งในเครื่องในโปรเจ็กต์ Google ML มีดังนี้

  1. ลบการเรียกไปยังกฎที่เก็บ CUDA ที่แยกต่างหากออกจากไฟล์ WORKSPACE ของโปรเจ็กต์ที่ขึ้นอยู่กับ XLA

  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 ควรประกอบด้วยเวอร์ชันหลักและเวอร์ชันย่อยเท่านั้น (เช่น 12.3 สำหรับ CUDA และ 9.1 สำหรับ CUDNN)

  4. ตอนนี้คุณเรียกใช้คำสั่ง bazel เพื่อใช้ CUDA และ CUDNN ที่ติดตั้งในเครื่องได้แล้ว

    สำหรับ XLA คุณไม่จำเป็นต้องเปลี่ยนแปลงตัวเลือกคำสั่ง

    สำหรับ JAX ให้ใช้ Flag --override_repository=tsl=<tsl_path> ในตัวเลือกคำสั่ง Bazel

    สำหรับ Tensorflow ให้ใช้ Flag --override_repository=local_tsl=<tsl_path> ในตัวเลือกคำสั่ง Bazel