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

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

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

יתרונות

ברשימה הבאה מפורטים היתרונות העיקריים של שימוש ב-xDiT להפעלת מודלים של DiT ב-Vertex AI:

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

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

‫xDiT זמין לארכיטקטורות מסוימות של מודלים של DiT ב-Vertex AI Model Garden, כמו 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 הצליחה לשחזר את התוצאות האלה ב-Vertex AI Model Garden.

האצה של GPU יחיד

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

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

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

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

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

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

אתם יכולים לפרוס נקודת קצה מותאמת אישית של Vertex AI עם קונטיינר xDiT באמצעות כרטיס מודל.

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

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

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

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

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

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

  2. בוחרים את מחברת 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 (integer): מגדיר את מידת ה-Ulysses לגישה המקבילית של הרצף המאוחד, שמשלבת בין DeepSpeed-Ulysses לבין Ring-Attention. ההגדרה הזו קובעת את דפוס התקשורת בין כל המשתתפים. כדי להשתמש בברירת המחדל, משאירים את השדה ריק.
  • ‫RING_DEGREE (integer): Sets the Ring degree for peer-to-peer communication in sequence parallelism. המאפיין פועל בשילוב עם ULYSSES_DEGREE כדי ליצור את רשת התהליך הדו-ממדית. כדי להשתמש בברירת המחדל, משאירים את השדה ריק.
ארגומנטים של Tensor Parallel
  • ‫TENSOR_PARALLEL_DEGREE (מספר שלם): הגדרת רמת המקביליות של טנסורים, שמפצלת את פרמטרים של המודל בין מכשירים לאורך ממדי התכונות כדי להפחית את עלויות הזיכרון לכל מכשיר. אם משאירים את השדה ריק, התכונה מושבתת.
  • ‫SPLIT_SCHEME (מחרוזת): הגדרה של אופן פיצול טנסורים של המודל בין מכשירים (למשל, לפי ראשי תשומת לב, מימדים מוסתרים). כדי להשתמש בברירת המחדל של תוכנית הפיצול, משאירים את השדה ריק.
ארגומנטים מבוזרים של Ray
  • ‫USE_RAY (boolean): הפעלה של מסגרת ביצוע מבוזרת של Ray לצורך שינוי קנה מידה של חישובים בכמה צמתים. ערך ברירת המחדל הוא false.
  • ‫RAY_WORLD_SIZE (integer): המספר הכולל של תהליכים באשכול 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 להסקת מסקנות ב-GPU יחיד, תוך שימוש ב-Input Temporal Reduction כדי לצמצם את המורכבות החישובית. ערך ברירת המחדל הוא false.
  • ‫N_CALIB (integer): מספר הדגימות לכיול לאופטימיזציה של 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 (בוליאני): הפעלת שיטת האופטימיזציה של TeaCache לשמירת תוצאות ביניים במטמון. ערך ברירת המחדל הוא false.
  • ‫USE_FBCACHE (boolean): הפעלה של שיטת האופטימיזציה First-Block-Cache. ערך ברירת המחדל הוא false.
ארגומנטים לאופטימיזציה של דיוק
  • ‫USE_FP8_T5_ENCODER (boolean): מפעיל דיוק FP8 (נקודה צפה של 8 ביט) עבור מקודד הטקסט T5, ומפחית את השימוש בזיכרון. יכול להיות שיהיה שיפור בנפח הנתונים שעוברים דרך המערכת עם השפעה מינימלית על האיכות. ערך ברירת המחדל הוא false.

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

ב-Model Garden יש הגדרות ברירת מחדל של xDiT parallelization למודלים נתמכים. אפשר לבדוק את הגדרות ברירת המחדל האלה באמצעות Vertex AI 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 זמין במאמרי העזרה.