הצגת Stable Diffusion XL ‏ (SDXL) באמצעות TPU ב-GKE עם MaxDiffusion

במדריך הזה נראה לכם איך להכניס לשימוש בסביבת הייצור מודל ליצירת תמונות של SDXL באמצעות Tensor Processing Units ‏ (TPUs) ב-Google Kubernetes Engine ‏ (GKE) עם MaxDiffusion. במדריך הזה מורידים את המודל מ-Hugging Face ומפריסים אותו באשכול Autopilot או Standard באמצעות קונטיינר שמריץ MaxDiffusion.

המדריך הזה הוא נקודת התחלה טובה אם אתם צריכים שליטה מפורטת, התאמה אישית, יכולת הרחבה, עמידות, ניידות וחסכוניות של Kubernetes מנוהל כשאתם פורסים ומפעילים את עומסי העבודה של AI/ML. אם אתם צריכים פלטפורמת AI מנוהלת ומאוחדת כדי ליצור ולהכניס לשימוש בסביבת הייצור במהירות מודלים של ML בצורה חסכונית, מומלץ לנסות את פתרון הפריסה שלנו Vertex AI.

רקע

כשמציגים את SDXL באמצעות TPU ב-GKE עם MaxDiffusion, אפשר ליצור פתרון חזק ומוכן לייצור עם כל היתרונות של Kubernetes מנוהל, כולל יעילות בעלויות, יכולת הרחבה וזמינות גבוהה יותר. בקטע הזה מתוארות הטכנולוגיות העיקריות שבהן נעשה שימוש במדריך הזה.

‫Stable Diffusion XL ‏ (SDXL)

Stable Diffusion XL‏ (SDXL) הוא סוג של מודל דיפוזיה לטנטי (LDM) שנתמך על ידי MaxDiffusion להסקת מסקנות. ב-AI גנרטיבי, אפשר להשתמש ב-LDM כדי ליצור תמונות באיכות גבוהה מתיאורי טקסט. מודלים של LDM שימושיים לאפליקציות כמו חיפוש תמונות ויצירת כיתוב לתמונות.

‫SDXL תומך בהסקת מסקנות במארח יחיד או בכמה מארחים עם הערות שיוך. כך אפשר לאמן את SDXL ולהריץ אותו בכמה מכונות, מה שיכול לשפר את היעילות.

מידע נוסף זמין במאגר של מודלים גנרטיביים של Stability AI ובמאמר בנושא SDXL.

TPUs

יחידות TPU הן מעגלים משולבים לאפליקציות ספציפיות (ASIC) שפותחו על ידי Google כדי להאיץ למידת מכונה ומודלים של AI שנבנו באמצעות frameworks כמו TensorFlow,‏ PyTorch ו-JAX.

לפני שמשתמשים ב-TPU ב-GKE, מומלץ להשלים את תוכנית הלימודים הבאה:

  1. מידע על הזמינות של גרסאות TPU עדכניות זמין בארכיטקטורת המערכת של Cloud TPU.
  2. מידע נוסף על TPUs ב-GKE

במדריך הזה נסביר איך להפעיל את מודל SDXL. ‫GKE פורס את המודל בצמתים של TPU v5e עם מארח יחיד, עם טופולוגיות של TPU שמוגדרות על סמך דרישות המודל להצגת הנחיות עם זמן אחזור נמוך. במדריך הזה, המודל משתמש בשבב TPU v5e עם טופולוגיה של 1x1.

MaxDiffusion

‫MaxDiffusion הוא אוסף של יישומי הפניה, שנכתבו ב-Python וב-Jax, של מודלים שונים של דיפוזיה סמויה שפועלים במכשירי XLA, כולל TPU ו-GPU. ‫MaxDiffusion היא נקודת התחלה לפרויקטים של Diffusion למטרות מחקר והפקה.

מידע נוסף זמין במאגר MaxDiffusion.

מטרות

המדריך הזה מיועד ללקוחות של AI גנרטיבי שמשתמשים ב-JAX, למשתמשים חדשים או קיימים ב-SDXL, ולמהנדסי ML, מהנדסי MLOps (DevOps) או אדמינים של פלטפורמות שרוצים להשתמש ביכולות של Kubernetes לניהול קונטיינרים כדי להציג מודלים גדולים של שפה (LLM).

במדריך הזה מוסבר איך:

  1. יוצרים אשכול GKE במצב Autopilot או אשכול רגיל עם טופולוגיית TPU מומלצת, על סמך מאפייני המודל.
  2. יוצרים קובץ אימג' של קונטיינר של הסקה (inference) של SDXL.
  3. פריסת שרת ההיקשים של SDXL ב-GKE.
  4. הצגה של המודל ואינטראקציה איתו דרך אפליקציית אינטרנט.

ארכיטקטורה

בקטע הזה מתוארת ארכיטקטורת GKE שבה נעשה שימוש במדריך הזה. הארכיטקטורה מורכבת מאשכול GKE Autopilot או מאשכול רגיל שמקצה TPU ומארח רכיבי MaxDiffusion. ‫GKE משתמש ברכיבים האלה כדי לפרוס את המודלים ולהפעיל אותם.

בתרשים הבא מוצגים רכיבי הארכיטקטורה הזו:

דוגמה לארכיטקטורה של הפעלת MaxDiffusion עם TPU v5e ב-GKE.

הארכיטקטורה הזו כוללת את הרכיבים הבאים:

  • אשכול אזורי של GKE Autopilot או Standard.
  • מאגר צמתים אחד של פרוסת TPU עם מארח יחיד, שמארח את מודל SDXL בהטמעה של MaxDiffusion.
  • רכיב השירות עם מאזן עומסים מסוג ClusterIP. השירות הזה מפזר את התנועה הנכנסת בין כל MaxDiffusion HTTP העותקים.
  • שרת WebApp HTTP עם שירות חיצוני LoadBalancer שמפיץ תנועה נכנסת ומפנה תנועה של מודלים שמשרתים בקשות לשירות ClusterIP.

לפני שמתחילים

  • נכנסים לחשבון Google Cloud . אם אתם משתמשים חדשים ב- Google Cloud, צרו חשבון כדי שתוכלו להעריך את הביצועים של המוצרים שלנו בתרחישים מהעולם האמיתי. לקוחות חדשים מקבלים בחינם גם קרדיט בשווי 300$ להרצה, לבדיקה ולפריסה של עומסי העבודה.
  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required APIs.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the APIs

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required APIs.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the APIs

  • צריך לוודא שיש לכם בפרויקט את התפקיד או התפקידים הבאים: roles/container.admin, roles/iam.serviceAccountAdmin, roles/artifactregistry.admin, roles/cloudbuild.builds.editor

    בדיקת התפקידים

    1. נכנסים לדף IAM במסוף Google Cloud .

      כניסה לדף IAM
    2. בוחרים את הפרויקט.
    3. בעמודה Principal (חשבון המשתמש), מוצאים את כל השורות שבהן מופיע השם שלכם או של קבוצה שאתם נכללים בה. כדי לברר באילו קבוצות אתם נכללים, פנו לאדמין.

    4. בודקים את העמודה Role בכל השורות שבהן מצוין או מופיע השם שלכם, כדי לראות אם רשימת התפקידים כוללת את התפקידים הנדרשים.

    מתן התפקידים

    1. נכנסים לדף IAM במסוף Google Cloud .

      כניסה לדף IAM
    2. בוחרים את הפרויקט.
    3. לוחצים על Grant access.
    4. בשדה New principals, מזינים את מזהה המשתמש. ‫ בדרך כלל מזהה המשתמש הוא כתובת האימייל של חשבון Google.

    5. לוחצים על Select a role ומחפשים את התפקיד.
    6. כדי לתת עוד תפקידים, לוחצים על Add another role ומוסיפים אותם.
    7. לוחצים על Save.
  • בודקים שיש לכם מספיק מכסת שימוש בשבבי TPU v5e PodSlice Lite. במדריך הזה משתמשים במופעים על פי דרישה.

הכנת הסביבה

במדריך הזה משתמשים ב-Cloud Shell כדי לנהל משאבים שמתארחים ב-Google Cloud. ב-Cloud Shell מותקן מראש התוכנה שצריך למדריך הזה, כולל kubectl ו- gcloud CLI.

כדי להגדיר את הסביבה באמצעות Cloud Shell:

  1. ב Google Cloud מסוף, מפעילים סשן של Cloud Shell על ידי לחיצה על Activate Cloud Shell בGoogle Cloud מסוף.סמל ההפעלה של Cloud Shell תופעל סשן בחלונית התחתונה של Google Cloud המסוף.

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

    gcloud config set project PROJECT_ID
    gcloud config set billing/quota_project PROJECT_ID
    export PROJECT_ID=$(gcloud config get project)
    export CLUSTER_NAME=CLUSTER_NAME
    export CLUSTER_VERSION=CLUSTER_VERSION
    export REGION=REGION_NAME
    export ZONE=ZONE
    

    מחליפים את הערכים הבאים:

    • PROJECT_ID: Google Cloud מזהה הפרויקט.
    • CLUSTER_NAME: השם של אשכול GKE.
    • CLUSTER_VERSION : גרסת GKE. צריך לציין גרסת GKE שתומכת ב-TPU Trillium‏ (v6e). מידע נוסף זמין במאמר אימות הזמינות של TPU ב-GKE.
    • REGION_NAME: האזור שבו נמצאים אשכול GKE, הקטגוריה של Cloud Storage וצמתי ה-TPU. האזור מכיל אזורים שבהם זמינים סוגי מכונות TPU v5e (לדוגמה, us-west1,‏ us-west4,‏ us-central1,‏ us-east1,‏ us-east5 או europe-west4).
    • (אשכול רגיל בלבד) ZONE: האזור שבו משאבי ה-TPU זמינים (לדוגמה, us-west4-a). באשכולות במצב Autopilot, לא צריך לציין את האזור, רק את האזור.
  3. משכפלים את מאגר הדוגמאות ופותחים את ספריית ההדרכה:

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
    cd kubernetes-engine-samples/ai-ml/maxdiffusion-tpu
    WORK_DIR=$(pwd)
    gcloud artifacts repositories create gke-llm --repository-format=docker --location=$REGION
    gcloud auth configure-docker $REGION-docker.pkg.dev
    

יצירה והגדרה של Google Cloud משאבים

כדי ליצור את המשאבים הנדרשים, פועלים לפי ההוראות הבאות.

יצירת אשכול GKE

אפשר להפעיל את SDXL ב-TPU באשכול GKE Autopilot או באשכול רגיל. מומלץ להשתמש באשכול Autopilot כדי ליהנות מחוויית Kubernetes מנוהלת באופן מלא. כדי לבחור את מצב הפעולה של GKE שהכי מתאים לעומסי העבודה שלכם, אפשר לעיין במאמר בחירת מצב פעולה של GKE.

טייס אוטומטי

  1. ב-Cloud Shell, מריצים את הפקודה הבאה:

    gcloud container clusters create-auto ${CLUSTER_NAME} \
      --project=${PROJECT_ID} \
      --location=${REGION} \
      --release-channel=rapid \
      --cluster-version=${CLUSTER_VERSION}
    

    ‫GKE יוצר אשכול Autopilot עם צמתי CPU ו-TPU לפי הבקשה של עומסי העבודה שנפרסו.

  2. מגדירים את kubectl לתקשורת עם האשכול:

      gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}
    

רגילה

  1. יוצרים אשכול GKE Standard אזורי שמשתמש ב-Workload Identity Federation for GKE.

    gcloud container clusters create ${CLUSTER_NAME} \
        --enable-ip-alias \
        --machine-type=n2-standard-4 \
        --num-nodes=2 \
        --workload-pool=${PROJECT_ID}.svc.id.goog \
        --location=${REGION}
    

    יצירת האשכול עשויה להימשך כמה דקות.

  2. מריצים את הפקודה הבאה כדי ליצור מאגר צמתים עבור האשכול:

    gcloud container node-pools create maxdiffusion-tpu-nodepool \
      --cluster=${CLUSTER_NAME} \
      --machine-type=ct5lp-hightpu-1t \
      --num-nodes=1 \
      --location=${REGION} \
      --node-locations=${ZONE} \
      --spot
    

    ‫GKE יוצר מאגר צמתים של TPU v5e עם טופולוגיה של 1x1 וצומת אחד.

    כדי ליצור מאגרי צמתים עם טופולוגיות שונות, כדאי לקרוא את המאמר בנושא תכנון ההגדרה של TPU. חשוב לעדכן את ערכי הדוגמה במדריך הזה, כמו cloud.google.com/gke-tpu-topology ו-google.com/tpu.

  3. מגדירים את kubectl לתקשורת עם האשכול:

      gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}
    

יצירת קונטיינר של הסקת מסקנות של SDXL

פועלים לפי ההוראות הבאות כדי ליצור קובץ אימג' של קונטיינר לשרת ההסקה של SDXL.

  1. פותחים את קובץ המניפסט build/server/cloudbuild.yaml:

    steps:
    - name: 'gcr.io/cloud-builders/docker'
      args: [ 'build', '-t', '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion:latest', '.' ]
    images:
    - '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion:latest'
  2. מריצים את ה-build ויוצרים קובץ אימג' של קונטיינר להסקת מסקנות.

    cd $WORK_DIR/build/server
    gcloud builds submit . --region=$REGION
    

    הפלט מכיל את הנתיב של קובץ האימג' של הקונטיינר.

פריסת שרת ההיקשים של SDXL

בקטע הזה, פורסים את שרת ההסקה של SDXL. כדי לפרוס את השרת, במדריך הזה נעשה שימוש בפריסת Kubernetes. פריסה היא אובייקט Kubernetes API שמאפשר להפעיל כמה רפליקות של Pods שמפוזרות בין הצמתים באשכול.

  1. עיון במניפסט של serve_sdxl_v5e.yaml.

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: stable-diffusion-deployment
    spec:
      selector:
        matchLabels:
          app: max-diffusion-server
      replicas: 1  # number of nodes in node-pool
      template:
        metadata:
          labels:
            app: max-diffusion-server
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-topology: 1x1 #  target topology
            cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
            #cloud.google.com/gke-spot: "true"
          volumes:
          - name: dshm
            emptyDir:
                  medium: Memory
          containers:
          - name: serve-stable-diffusion
            image: REGION-docker.pkg.dev/PROJECT_ID/gke-llm/max-diffusion:latest
            env:
            - name: MODEL_NAME
              value: 'stable_diffusion'
            ports:
            - containerPort: 8000
            resources:
              requests:
                google.com/tpu: 1  # TPU chip request
              limits:
                google.com/tpu: 1  # TPU chip request
            volumeMounts:
                - mountPath: /dev/shm
                  name: dshm
    
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: max-diffusion-server
      labels:
        app: max-diffusion-server
    spec:
      type: ClusterIP
      ports:
        - port: 8000
          targetPort: 8000
          name: http-max-diffusion-server
          protocol: TCP
      selector:
        app: max-diffusion-server
  2. מעדכנים את מזהה הפרויקט במניפסט.

    cd $WORK_DIR
    perl -pi -e 's|PROJECT_ID|PROJECT_ID|g' serve_sdxl_v5e.yaml
    perl -pi -e 's|REGION|REGION_NAME|g' serve_sdxl_v5e.yaml
    
  3. החלת המניפסט:

    kubectl apply -f serve_sdxl_v5e.yaml
    

    הפלט אמור להיראות כך:

    deployment.apps/max-diffusion-server created
    
  4. מאמתים את הסטטוס של המודל:

    kubectl get deploy --watch
    

    הפלט אמור להיראות כך:

    NAME                          READY   UP-TO-DATE   AVAILABLE   AGE
    stable-diffusion-deployment   1/1     1            1           8m21s
    
  5. מאחזרים את הכתובת ClusterIP:

    kubectl get service max-diffusion-server
    

    הפלט מכיל את השדה ClusterIP. חשוב לשים לב לערך של CLUSTER-IP.

  6. אימות הפריסה:

     export ClusterIP=CLUSTER_IP
     kubectl run curl --image=curlimages/curl \
        -it --rm --restart=Never \
        -- "$ClusterIP:8000"
    

    מחליפים את CLUSTER_IP בערך CLUSTER-IP שרשמתם קודם. הפלט אמור להיראות כך:

    {"message":"Hello world! From FastAPI running on Uvicorn with Gunicorn."}
    pod "curl" deleted
    
  7. צפייה ביומנים מהפריסה:

    kubectl logs -l app=max-diffusion-server
    

    בסיום הפריסה, הפלט אמור להיראות כך:

    2024-06-12 15:45:45,459 [INFO] __main__: replicate params:
    2024-06-12 15:45:46,175 [INFO] __main__: start initialized compiling
    2024-06-12 15:45:46,175 [INFO] __main__: Compiling ...
    2024-06-12 15:45:46,175 [INFO] __main__: aot compiling:
    2024-06-12 15:45:46,176 [INFO] __main__: tokenize prompts:2024-06-12 15:48:49,093 [INFO] __main__: Compiled in 182.91802048683167
    INFO:     Started server process [1]
    INFO:     Waiting for application startup.
    INFO:     Application startup complete.
    

פריסת לקוח אפליקציית האינטרנט

בקטע הזה, תפרסו את לקוח אפליקציית האינטרנט כדי להציג את מודל SDXL.

  1. עיון במניפסט של build/webapp/cloudbuild.yaml.

    steps:
    - name: 'gcr.io/cloud-builders/docker'
      args: [ 'build', '-t', '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion-web:latest', '.' ]
    images:
    - '$LOCATION-docker.pkg.dev/$PROJECT_ID/gke-llm/max-diffusion-web:latest'
  2. מריצים את ה-build ויוצרים את קובץ האימג' של הקונטיינר של הלקוח בספרייה build/webapp.

    cd $WORK_DIR/build/webapp
    gcloud builds submit . --region=$REGION
    

    הפלט מכיל את הנתיב של קובץ האימג' של הקונטיינר.

  3. פותחים את קובץ המניפסט serve_sdxl_client.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: max-diffusion-client
    spec:
      selector:
        matchLabels:
          app: max-diffusion-client
      template:
        metadata:
          labels:
            app: max-diffusion-client
        spec:
          containers:
          - name: webclient
            image: REGION-docker.pkg.dev/PROJECT_ID/gke-llm/max-diffusion-web:latest
            env:
              - name: SERVER_URL
                value: "http://ClusterIP:8000"
            resources:
              requests:
                memory: "128Mi"
                cpu: "250m"
              limits:
                memory: "256Mi"
                cpu: "500m"
            ports:
            - containerPort: 5000
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: max-diffusion-client-service
    spec:
      type: LoadBalancer
      selector:
        app: max-diffusion-client
      ports:
      - port: 8080
        targetPort: 5000
  4. עורכים את מזהה הפרויקט בקובץ המניפסט:

    cd $WORK_DIR
    perl -pi -e 's|PROJECT_ID|PROJECT_ID|g' serve_sdxl_client.yaml
    perl -pi -e 's|ClusterIP|CLUSTER_IP|g' serve_sdxl_client.yaml
    perl -pi -e 's|REGION|REGION_NAME|g' serve_sdxl_client.yaml
    
  5. החלת המניפסט:

    kubectl apply -f serve_sdxl_client.yaml
    
  6. מאחזרים את כתובת ה-IP‏ LoadBalancer:

    kubectl get service max-diffusion-client-service
    

    הפלט מכיל את השדה LoadBalancer. חשוב לשים לב לערך של EXTERNAL-IP.

אינטראקציה עם המודל באמצעות דף האינטרנט

  1. גישה לכתובת ה-URL הבאה מדפדפן אינטרנט:

    http://EXTERNAL_IP:8080
    

    מחליפים את EXTERNAL_IP בערך EXTERNAL_IP שרשמתם קודם.

  2. לקיים אינטראקציה עם SDXL באמצעות ממשק הצ'אט. מוסיפים הנחיה ולוחצים על שליחה. לדוגמה:

    Create a detailed image of a fictional historical site, capturing its unique architecture and cultural significance
    

הפלט הוא תמונה שנוצרה על ידי מודל, בדומה לדוגמה הבאה:

תמונה שנוצרה על ידי SDXL

הסרת המשאבים

כדי להימנע מחיובים בחשבון Google Cloud בגלל השימוש במשאבים שנעשה במסגרת המדריך הזה, אפשר למחוק את הפרויקט שמכיל את המשאבים, או להשאיר את הפרויקט ולמחוק את המשאבים בנפרד.

מחיקת הפרויקט

  1. במסוף Google Cloud , נכנסים לדף Manage resources.

    כניסה לדף Manage resources

  2. ברשימת הפרויקטים, בוחרים את הפרויקט שרוצים למחוק ולוחצים על Delete.
  3. כדי למחוק את הפרויקט, כותבים את מזהה הפרויקט בתיבת הדו-שיח ולוחצים על Shut down.

מחיקת המשאבים הבודדים

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

gcloud container clusters delete ${CLUSTER_NAME} --location=${REGION}

המאמרים הבאים