ภาพรวมของ 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.6.3 ตัวแปรสภาพแวดล้อม HERMETIC_CUDNN_VERSION ควรประกอบด้วยเวอร์ชันหลัก รอง และแพตช์ CUDNN เช่น 9.3.0

3 วิธีในการตั้งค่าตัวแปรสภาพแวดล้อมสําหรับคําสั่ง Bazel

# Add an entry to your `.bazelrc` file
build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.6.3"
build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"

# OR pass it directly to your specific build command
bazel build --config=cuda <target> \
--repo_env=HERMETIC_CUDA_VERSION="12.6.3" \
--repo_env=HERMETIC_CUDNN_VERSION="9.3.0"

# 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.6.3"
export HERMETIC_CUDNN_VERSION="9.3.0"

หากไม่มี 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(
   "@xla//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(
   "@xla//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(
   "@xla//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
   "cuda_configure",
)

cuda_configure(name = "local_config_cuda")

load(
   "@xla//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
   "nccl_redist_init_repository",
)

nccl_redist_init_repository()

load(
   "@xla//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.6.3"
build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
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

    จำเป็นต้องใช้ Flag เพื่อให้แน่ใจว่ามีการระบุการพึ่งพา 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 ที่อัปเดตแล้วใน third_party/gpus/cuda/hermetic/cuda_redist_versions.bzlCUDNN_REDIST_JSON_DICT

    อัปเดต 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 หากจำเป็น

    อัปเดต PTX_VERSION_DICT ใน third_party/gpus/cuda/hermetic/cuda_redist_versions.bzl หากจำเป็น

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

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

  3. สำหรับการเพิ่มประสิทธิภาพเวลา การกำหนดค่าบิลด์/การทดสอบบางรายการใช้การแจกจ่าย.tarแบบมิเรอร์ ระบบจะอัปโหลดไฟล์ json ที่มีข้อมูลเกี่ยวกับการแจกจ่าย.tarแบบมิเรอร์ในภายหลังหลังจากที่อัปเดต CUDA_REDIST_JSON_DICT และ CUDNN_REDIST_JSON_DICT แล้ว ผู้ใช้สามารถดาวน์โหลดไฟล์เหล่านี้ได้โดยใช้ wget "https://storage.googleapis.com/mirror.tensorflow.org/developer.download.nvidia.com/compute/cuda/redist/redistrib_<cuda_version>_tar.json" สำหรับ CUDA และ wget "https://storage.googleapis.com/mirror.tensorflow.org/developer.download.nvidia.com/compute/cudnn/redist/redistrib_<cudnn_version>_tar.json" สำหรับ CUDNN หลังจากนั้นให้สร้างและส่งคำขอดึงข้อมูลที่มีพจนานุกรม ที่อัปเดตแล้วใน third_party/gpus/cuda/hermetic/cuda_redist_versions.bzlMIRRORED_TARS_CUDA_REDIST_JSON_DICTMIRRORED_TARS_CUDNN_REDIST_JSON_DICT

ชี้ไปยังการแจกจ่าย 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 ที่กำหนดเอง

ตัวเลือกนี้ช่วยให้ใช้การแจกจ่ายที่กําหนดเองสําหรับทรัพยากร 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 ที่กําหนดเอง

ตัวเลือกนี้ช่วยให้ใช้การแจกจ่ายที่กําหนดเองสําหรับทรัพยากร 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("@xla//third_party/gpus:cuda_configure.bzl", "cuda_configure")
cuda_configure(name = "local_config_cuda")
load("@xla//third_party/nccl:nccl_configure.bzl", "nccl_configure")
nccl_configure(name = "local_config_nccl")

    สำหรับ Tensorflow

    load("@local_xla//third_party/gpus:cuda_configure.bzl", "cuda_configure")
cuda_configure(name = "local_config_cuda")
load("@local_xla//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