סקירה כללית על CUDA הרמטית

ב-Hermetic CUDA נעשה שימוש בגרסה ספציפית של CUDA שניתן להורדה, במקום ב-CUDA שמותקן באופן מקומי אצל המשתמש. Bazel יוריד את הפצות CUDA,‏ CUDNN ו-NCCL, ואז ישתמש בספריות ובכלים של CUDA כיחסי תלות ביעדי Bazel שונים. כך ניתן ליצור גרסאות build רבות יותר לפרויקטים של 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 של פרויקטים נפרדים, משתני הסביבה HERMETIC_CUDA_VERSION ו-HERMETIC_CUDNN_VERSION מוגדרים לגרסאות שבהן נעשה שימוש כברירת מחדל כשמציינים את --config=cuda באפשרויות הפקודה של Bazel.

משתני סביבה ששולטים בגרסאות ה-CUDA/CUDNN האטומות

משתנה הסביבה HERMETIC_CUDA_VERSION צריך לכלול גרסה ראשית, גרסת CUDA משנית וגרסת CUDA, למשל 12.3.2. משתנה הסביבה HERMETIC_CUDNN_VERSION צריך לכלול את הגרסה הראשית, המשנית והתיקון של CUDNN, למשל 9.1.1.

יש שלוש דרכים להגדיר את משתני הסביבה לפקודות 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 אטום במהלך ביצוע הבדיקה, או כשמריצים קובץ בינארי באמצעות bazel, צריך להוסיף את הדגל --@local_config_cuda//cuda:include_cuda_libs=true לפקודת bazel. אפשר לספק אותו ישירות במעטפת או ב-.bazelrc:

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

    הדגל הזה נדרש כדי לוודא שהתלות של CUDA מסופקות כראוי כדי לבדוק קובצי הפעלה. הדגל מוגדר כ-false כברירת מחדל כדי למנוע צימוד לא רצוי של חבילות Python שפורסמו על ידי Google לקובצי CUDA בינאריים.

  4. כדי לאכוף את מצב התאימות לעתיד של CUDA, מוסיפים את הדגל --@cuda_driver//:enable_forward_compatibility=true לפקודת bazel. אפשר לספק אותו ישירות במעטפת או ב-.bazelrc:

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

    ערך ברירת המחדל של הדגל הוא false.

    כשמצב התאימות לעתיד של CUDA מושבת, מטרות Bazel ישתמשו במנהלי התקן במצב משתמש ובמנהלי התקן במצב ליבה שמותקנים מראש במערכת.

    כשמפעילים את מצב התאימות לעתיד של CUDA, מטרות Bazel ישתמשו ב-User Mode Driver מההפצה מחדש של מנהל ה-CUDA שהורדתם למטמון של Bazel, וב-Kernel Mode Driver שמותקן מראש במערכת. הוא מאפשר להפעיל תכונות חדשות של CUDA Toolkit תוך שימוש ב-Kernel Mode Driver ישן יותר.

    צריך לאכוף את מצב התאימות לעתיד רק במקרים הרלוונטיים. פרטים נוספים זמינים במסמכי התיעוד של 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

יש שלוש אפשרויות שמאפשרות להשתמש בהפצות מותאמות אישית של 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 ה hermetic של 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,
)

DEPRECATED: Non-hermetic CUDA/CUDNN usage

השימוש ב-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, משתמשים בדגל --override_repository=tsl=<tsl_path> באפשרויות הפקודה של Bazel.

    ב-Tensorflow, משתמשים בדגל --override_repository=local_tsl=<tsl_path> באפשרויות הפקודה של Bazel.