הצגת מודלים של טרנספורמרים לדיפוזיה באמצעות קונטיינר xDiT במעבדים גרפיים ל-Cloud

‫xDiT היא ספריית קוד פתוח שמאיצה את ההסקה של מודלים של Diffusion Transformer ‏ (DiT) באמצעות טכניקות של מקביליות ואופטימיזציה. הטכניקות האלה מאפשרות הגדרה של כמה יחידות GPU שניתנת להרחבה, לעומסי עבודה תובעניים. בדף הזה נסביר איך לפרוס מודלים של DiT באמצעות xDiT ויחידות GPU של Cloud בפלטפורמת הסוכנים של Gemini Enterprise.

מידע נוסף על xDiT זמין בפרויקט xDiT ב-GitHub.

יתרונות

ברשימה הבאה מפורטים היתרונות העיקריים של שימוש ב-xDiT להצגת מודלים של DiT בפלטפורמת הסוכנים:

  • יצירה מהירה פי שלושה: יצירת תמונות וסרטונים ברזולוציה גבוהה בחלק קטן מהזמן בהשוואה לפתרונות אחרים להצגת תמונות.
  • תמיכה במעבדי GPU מרובים עם יכולת הרחבה: חלוקת עומסי העבודה בין כמה מעבדי GPU בצורה יעילה לביצועים אופטימליים.
    • מקביליות היברידית: xDiT תומך בגישות שונות לעיבוד מקבילי, כמו מקביליות מאוחדת של רצפים, PipeFusion, מקביליות של CFG ומקביליות של נתונים. אפשר לשלב בין השיטות האלה כדי ליצור מתכון ייחודי לאופטימיזציה של הביצועים.
  • ביצועים אופטימליים של GPU יחיד: xDiT מספק הסקת מסקנות מהירה יותר גם ב-GPU יחיד.
    • האצת GPU: xDiT משלב כמה שיטות להאצת ליבה ומשתמש בטכניקות מ-DiTFastAttn כדי להאיץ את ההסקה ב-GPU יחיד.
  • פריסה קלה: אפשר להתחיל במהירות עם פריסה בלחיצה אחת או עם מחברות Colab Enterprise ב-Gemini Enterprise Agent Platform Model Garden.

מודלים נתמכים

‫xDiT זמין לארכיטקטורות מסוימות של מודלים של DiT ב-Model Garden של Gemini Enterprise Agent Platform, כמו Flux.1 Schnell,‏ CogVideoX-2b ו-Wan2.1, וריאציות של מודלים להמרת טקסט לסרטון. כדי לבדוק אם מודל DiT תומך ב-xDiT ב-Model Garden, אפשר לעיין בכרטיס המודל שלו ב-Model Garden.

מקביליות היברידית לביצועים של כמה מעבדי GPU:

‫xDiT משתמשת בשילוב של טכניקות מקביליות כדי למקסם את הביצועים בהגדרות של כמה מעבדים גרפיים. הטכניקות האלה פועלות יחד כדי לחלק את עומס העבודה ולבצע אופטימיזציה של ניצול המשאבים:

  • מקביליות מאוחדת של רצפים: הטכניקה הזו מפצלת את נתוני הקלט (למשל, פיצול תמונה לטלאים) בין כמה מעבדי GPU, וכך מצמצמת את השימוש בזיכרון ומשפרת את יכולת ההתאמה.
  • ‫PipeFusion: ‏ PipeFusion מחלק את מודל DiT לשלבים, ומשייך כל שלב למעבד GPU אחר, וכך מאפשר עיבוד מקביל של חלקים שונים במודל.
  • מקביליות CFG: הטכניקה הזו מייעלת מודלים באופן ספציפי באמצעות הנחיה ללא סיווג, שיטה נפוצה לשליטה בסגנון ובתוכן של תמונות שנוצרות. הוא מבצע חישובים מקבילים של ענפים מותנים ולא מותנים, וכך מסיק מסקנות מהר יותר.
  • מקביליות נתונים: בשיטה הזו, המודל כולו משוכפל בכל GPU, וכל GPU מעבד קבוצה שונה של נתוני קלט, וכך מגדיל את התפוקה הכוללת של המערכת.

מידע נוסף על שיפורי הביצועים זמין בדוח של xDiT על Flux.1 Schnell או על CogVideoX-2b. ‫Google הצליחה לשחזר את התוצאות האלה ב-Gemini Enterprise Agent Platform Model Garden.

האצה של GPU יחיד

ספריית xDiT מספקת יתרונות להצגת מודלים ב-GPU יחיד באמצעות torch.compile ו-onediff כדי לשפר את מהירות זמן הריצה ב-GPU. אפשר להשתמש בטכניקות האלה גם בשילוב עם מקביליות היברידית.

בנוסף, ל-xDiT יש טכניקה יעילה לחישוב תשומת לב, שנקראת DiTFastAttn, כדי לטפל בצוואר הבקבוק החישובי של DiT. נכון לעכשיו, הטכניקה הזו זמינה לשימוש רק בהגדרות של GPU יחיד או בשילוב עם מקביליות נתונים.

תחילת העבודה ב-Model Garden

קונטיינר ההגשה של xDiT ל-GPU בענן מותאם ומסופק ב-Model Garden של Gemini Enterprise Agent Platform. במודלים נתמכים, הפריסות משתמשות בקונטיינר הזה כשמשתמשים בפריסות בלחיצה אחת או בדוגמאות של מחברות Colab Enterprise.

בדוגמאות הבאות נעשה שימוש במודל Flux.1-schnell כדי להדגים איך פורסים מודל DiT בקונטיינר xDiT.

שימוש בפריסה בלחיצה אחת

אפשר לפרוס נקודת קצה מותאמת אישית של Agent Platform באמצעות קובץ מאפייני מודל של מאגר xDiT.

  1. מנווטים אל דף כרטיס המודל ולוחצים על פריסה.

  2. כדי לבחור את וריאציית המודל שבה רוצים להשתמש, בוחרים סוג מכונה שבו רוצים להשתמש לפריסה.

  3. לוחצים על Deploy (פריסה) כדי להתחיל את תהליך הפריסה. תקבלו שני אימיילים: אחד כשהמודל יועלה ואחד כשהנקודה הסופית תהיה מוכנה.

שימוש ב-notebook של Colab Enterprise

כדי ליהנות מגמישות ומאפשרויות התאמה אישית, אפשר להשתמש בדוגמאות של מחברות Colab Enterprise כדי לפרוס נקודת קצה של פלטפורמת סוכנים עם מאגר xDiT באמצעות Agent Platform SDK ל-Python.

  1. עוברים אל דף כרטיס המודל ולוחצים על פתיחת מחברת.

  2. בוחרים את מסמך ה-notebook של Vertex Serving. מסמך ה-notebook ייפתח ב-Colab Enterprise.

  3. מריצים את המחברת כדי לפרוס מודל באמצעות קונטיינר xDiT ושולחים בקשות לחיזוי לנקודת הקצה. קטע הקוד לפריסה הוא:

import vertexai
from vertexai import model_garden

vertexai.init(project=<YOUR_PROJECT_ID>, location=<REGION>)

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
endpoint = model.deploy()

ארגומנטים של xDiT

‫xDiT מציעה מגוון ארגומנטים של שרת שאפשר להגדיר כדי לבצע אופטימיזציה של הביצועים לתרחישי שימוש ספציפיים. הארגומנטים האלה מוגדרים כמשתני סביבה במהלך הפריסה. הנה רשימה של הארגומנטים העיקריים שאולי תצטרכו להגדיר:

הגדרת המודל
  • ‫MODEL_ID (מחרוזת): מציין את מזהה המודל לטעינה. השם הזה צריך להיות זהה לשם המודל במרשם או בנתיב.
ארגומנטים של אופטימיזציה בזמן ריצה
  • ‫N_GPUS (מספר שלם): מציין את מספר יחידות ה-GPU שבהן יש להשתמש להסקת מסקנות. ערך ברירת המחדל הוא 1.
  • ‫WARMUP_STEPS (מספר שלם): מספר שלבי החימום שנדרשים לפני שמתחילים להסיק מסקנות. זה חשוב במיוחד כש-PipeFusion מופעל, כדי להבטיח ביצועים יציבים. ערך ברירת המחדל הוא 1.
  • ‫USE_PARALLEL_VAE (boolean): מאפשר עיבוד יעיל של תמונות ברזולוציה גבוהה (מעל 2,048 פיקסלים) על ידי הפעלת רכיב ה-VAE במקביל במכשירים. כך נמנעות בעיות של חוסר זיכרון (OOM) בתמונות גדולות. ערך ברירת המחדל הוא false.
  • ‫USE_TORCH_COMPILE (boolean): מפעיל שיפור מהירות באמצעות GPU יחיד דרך torch.compile, ומספק אופטימיזציות ברמת הליבה לשיפור הביצועים. ערך ברירת המחדל הוא false.
  • ‫USE_ONEDIFF (boolean): מפעיל את טכנולוגיית ההידור המואץ של OneDiff כדי לבצע אופטימיזציה של מהירות ההפעלה של ליבת ה-GPU. ערך ברירת המחדל הוא false.
ארגומנטים של מקביליות נתונים
  • ‫DATA_PARALLEL_DEGREE (מספר שלם): מגדיר את רמת המקביליות של הנתונים. אם משאירים את השדה ריק, התכונה מושבתת. אפשר גם להגדיר את מידת המקבילות הרצויה.
  • ‫USE_CFG_PARALLEL (boolean): הפעלת חישוב מקביל להנחיה ללא סיווג (CFG), שנקראת גם Split Batch. כשההגדרה הזו מופעלת, מידת המקבילות הקבועה היא 2. הערך מוגדר כ-true כשמשתמשים ב-CFG כדי לשלוט בסגנון ובתוכן של הפלט. ערך ברירת המחדל הוא False.
ארגומנטים של Parallel Sequence (USP – Unified Sequence Parallelism)
  • ‫ULYSSES_DEGREE (מספר שלם): מגדיר את מידת ה-Ulysses עבור הגישה המקבילית של הרצף המאוחד, שמשלבת בין DeepSpeed-Ulysses לבין Ring-Attention. ההגדרה הזו שולטת בדפוס התקשורת של כולם-לכולם. אם לא מגדירים ערך, נעשה שימוש בערך ברירת המחדל.
  • ‫RING_DEGREE (מספר שלם): מגדיר את מידת הטבעת לתקשורת בין עמיתים במקביל ברצף. פועל בשילוב עם ULYSSES_DEGREE כדי ליצור את רשת התהליכים הדו-ממדית. אם לא מציינים ערך, נעשה שימוש בערך ברירת המחדל.
ארגומנטים של Tensor Parallel
  • ‫TENSOR_PARALLEL_DEGREE (מספר שלם): מגדיר את רמת המקביליות של טנסורים, שמפצלת את פרמטרים של מודלים בין מכשירים לאורך ממדי התכונות כדי להפחית את עלויות הזיכרון לכל מכשיר. כדי להשבית את האפשרות, משאירים את השדה ריק.
  • ‫SPLIT_SCHEME (מחרוזת): מגדיר איך לפצל את טנסורים של המודל בין מכשירים (למשל, לפי ראשי תשומת לב, ממדים מוסתרים). משאירים ריק כדי להשתמש בסכמת הפיצול שמוגדרת כברירת מחדל.
ארגומנטים מבוזרים של Ray
  • ‫USE_RAY (בוליאני): מפעיל את מסגרת הביצוע המבוזרת של Ray לצורך שינוי קנה מידה של חישובים בכמה צמתים. ערך ברירת המחדל הוא false.
  • ‫RAY_WORLD_SIZE (מספר שלם): המספר הכולל של תהליכים באשכול Ray. ערך ברירת המחדל הוא 1.
  • ‫VAE_PARALLEL_SIZE (מספר שלם): מספר התהליכים שמוקדשים לעיבוד מקביל של VAE כשמשתמשים ב-Ray. ערך ברירת המחדל הוא 0.
  • ‫DIT_PARALLEL_SIZE (מספר שלם): מספר התהליכים שמוקדשים לעיבוד מקביל של עמוד השדרה של DiT כשמשתמשים ב-Ray. ערך ברירת המחדל הוא 0.
ארגומנטים מקבילים של PipeFusion
  • ‫PIPEFUSION_PARALLEL_DEGREE (מספר שלם): מגדיר את רמת המקביליות של PipeFusion, מקביליות של צינור עיבוד נתונים ברמת הרצף שמנצלת את מאפייני יתירות הזמנים של מודלים של דיפוזיה. ערכים גבוהים יותר מגדילים את המקביליות, אבל דורשים יותר זיכרון. ערך ברירת המחדל הוא 1.
  • ‫NUM_PIPELINE_PATCH (מספר שלם): מספר הטלאים שאליהם יפוצל הרצף לצורך עיבוד בצינור. אפשר להשאיר את השדה ריק כדי שהמערכת תקבע את הערך באופן אוטומטי.
  • ‫ATTN_LAYER_NUM_FOR_PP (מחרוזת): מציינת באילו שכבות של תשומת לב להשתמש כדי להפעיל מקביליות בצינור. אפשר להפריד בין הערכים באמצעות פסיק (למשל, '10,9') או רווח (למשל, '10 9'). כדי להשתמש בכל השכבות, משאירים את השדה ריק.
ארגומנטים לאופטימיזציה של הזיכרון
  • ‫ENABLE_MODEL_CPU_OFFLOAD (boolean): מעביר את משקלי המודל לזיכרון המעבד כשהם לא בשימוש, וכך מצמצם את השימוש בזיכרון ה-GPU, אבל מגדיל את זמן האחזור. ערך ברירת המחדל הוא False.
  • ‫ENABLE_SEQUENTIAL_CPU_OFFLOAD (boolean): מעביר שכבות של מודלים באופן עקבי ל-CPU במהלך העברה קדימה, וכך מאפשר הסקת מסקנות לגבי מודלים גדולים יותר מזיכרון ה-GPU. ערך ברירת המחדל הוא false.
  • ‫ENABLE_TILING (boolean): מפחית את השימוש בזיכרון ה-GPU על ידי פענוח רכיב ה-VAE בכל פעם משבצת אחת. הארגומנט הזה שימושי לתמונות או לסרטונים גדולים יותר, ולמניעת שגיאות של חוסר זיכרון. ערך ברירת המחדל הוא False.
  • ‫ENABLE_SLICING (boolean): מפחית את השימוש בזיכרון ה-GPU על ידי פיצול טנסור הקלט לפרוסות לצורך פענוח VAE. ערך ברירת המחדל הוא False.
ארגומנטים של DiTFastAttn (אופטימיזציה של תשומת הלב)
  • ‫USE_FAST_ATTN (boolean): מפעיל את ההאצה של DiTFastAttn להסקת מסקנות במעבד גרפי יחיד, באמצעות צמצום זמני של הקלט כדי לצמצם את המורכבות החישובית. ערך ברירת המחדל הוא False.
  • ‫N_CALIB (מספר שלם): מספר דגימות הכיול לאופטימיזציה של DiTFastAttn. ערך ברירת המחדל הוא 8.
  • ‫THRESHOLD (float): מידת הדמיון המינימלית בין המשתמשים לצורך צמצום הדמיון הזמני ב-DiTFastAttn. ערך ברירת המחדל הוא 0.5.
  • ‫WINDOW_SIZE (מספר שלם): גודל החלון של Window Attention עם Residual Caching כדי לצמצם את יתירות המידע המרחבית. ערך ברירת המחדל הוא 64.
  • ‫COCO_PATH (string): הנתיב למערך הנתונים של COCO לכיול DiTFastAttn. נדרש אם USE_FAST_ATTN הוא true. אם לא משתמשים, צריך להשאיר ריק.
ארגומנטים לאופטימיזציה של המטמון
  • ‫USE_CACHE (boolean): הפעלה של מנגנוני שמירה במטמון כדי לצמצם חישובים מיותרים. ערך ברירת המחדל הוא False.
  • ‫USE_TEACACHE (boolean): מפעיל את שיטת האופטימיזציה של TeaCache לשמירת תוצאות ביניים במטמון. ערך ברירת המחדל הוא false.
  • ‫USE_FBCACHE (boolean): הפעלת שיטת האופטימיזציה First-Block-Cache. ערך ברירת המחדל הוא false.
טיעונים לאופטימיזציה של דיוק
  • ‫USE_FP8_T5_ENCODER (boolean): מפעיל דיוק FP8 (נקודה צפה של 8 ביט) עבור מקודד הטקסט T5, ומקטין את השימוש בזיכרון. יכול להיות שההגדרה הזו תשפר את קצב העברת הנתונים עם השפעה מינימלית על האיכות. ערך ברירת המחדל הוא False.

הצגת מוצרים בהתאמה אישית

ב-Model Garden יש הגדרות ברירת מחדל של xDiT parallelization למודלים נתמכים. אפשר לבדוק את הגדרות ברירת המחדל האלה באמצעות Agent Platform SDK ל-Python.

כדי לראות את הגדרת הפריסה שמוגדרת כברירת מחדל למודל, כמו black-forest-labs/FLUX.1-schnell, אפשר להריץ את קטע הקוד הבא:

import vertexai
from vertexai import model_garden

vertexai.init(project=<YOUR_PROJECT_ID>, location=<REGION>)

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")
deploy_options = model.list_deploy_options()


# Example Response
# ['black-forest-labs/flux1-schnell@flux.1-schnell']
# [model_display_name: "Flux1-schnell"
# container_spec {
#   image_uri: "us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-2.ubuntu2204.py310"
#  env {
#    name: "DEPLOY_SOURCE"
#    value: "UI_NATIVE_MODEL"
#  }
#  env {
#    name: "MODEL_ID"
#    value: "gs://vertex-model-garden-restricted-us/black-forest-labs/FLUX.1-schnell"
#  }
#  env {
#    name: "TASK"
#    value: "text-to-image"
#  }
#  env {
#    name: "N_GPUS"
#    value: "2"
#  }
#  env {
#    name: "USE_TORCH_COMPILE"
#    value: "true"
#  }
#  env {
#    name: "RING_DEGREE"
#    value: "2"
#  }
# ..........]

השיטה list_deploy_options() מחזירה את המפרטים של הקונטיינר, כולל משתני הסביבה (env) שמגדירים את התצורה של xDiT.

כדי להתאים אישית את אסטרטגיית ההקבילות, אפשר לשנות את משתני הסביבה האלה כשפורסים את המודל. בדוגמה הבאה מוצג איך לשנות את הערכים של RING_DEGREE ו-ULYSSES_DEGREE בהגדרה של 2 יחידות GPU, כדי לשנות את גישת ההקבלה:

import vertexai
from vertexai import model_garden

# Replace with your project ID and region
vertexai.init(project="<YOUR_PROJECT_ID>", location="<REGION>")

model = model_garden.OpenModel("black-forest-labs/FLUX.1-schnell")

# Custom environment variables to override default settings
# This example sets N_GPUS as 2, so RING_DEGREE * ULYSSES_DEGREE must equal 2
container_env_vars = {
    "N_GPUS": "2",
    "RING_DEGREE": "1",
    "ULYSSES_DEGREE": "2"
    # Add other environment variables to customize here
}

machine_type = "a3-highgpu-2g"
accelerator_type = "NVIDIA_H100_80GB"
accelerator_count = 2

# Deploy the model with the custom environment variables
endpoint = model.deploy(
    machine_type=machine_type,
    accelerator_type=accelerator_type,
    accelerator_count=accelerator_count,
  container_env_vars=container_env_vars
)

כדאי לעיין בקטע 'הסבר על ארגומנטים ספציפיים של xDiT' כדי לקבל פרטים על כל משתנה סביבה. חשוב לוודא שמכפלת דרגות המקביליות (למשל, PIPEFUSION_PARALLEL_DEGREE,‏ ULYSSES_DEGREE,‏ RING_DEGREE ו-USE_CFG_PARALLEL) שווה למספר הכולל של יחידות ה-GPU‏ (N_GPUS).

דוגמאות נוספות למתכונים ולהגדרות של מודלים שונים זמינות במאמרי העזרה הרשמיים של xDiT. מידע נוסף על Model Garden SDK זמין במאמרי העזרה.