ביצוע הסקה מרובת מארחים באמצעות Pathways

הסקת מסקנות מרובת-מארחים היא שיטה להסקת מסקנות ממודל שמפיצה את המודל על פני כמה מארחים של מאיצים. ההגדרה הזו מאפשרת הסקה של מודלים גדולים שלא נכנסים למארח יחיד. אפשר להשתמש בנתיבי למידה גם בתרחישי שימוש של הסקת מסקנות בכמה מארחים (batch) וגם בתרחישי שימוש של הסקת מסקנות בזמן אמת.

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

חשוב לוודא שיש לכם:

הרצת הסקת מסקנות באצווה באמצעות JetStream

‫JetStream הוא מנוע עם אופטימיזציה של נפח הנתונים וזיכרון להסקת מסקנות של מודלים גדולים של שפה (LLM) במכשירי XLA, בעיקר יחידות לעיבוד טנסורים (TPU) שנכתבו ב-JAX.

אפשר להשתמש בתמונה מוכנה מראש של JetStream Docker כדי להריץ עומס עבודה של הסקת מסקנות באצווה, כמו שמוצג בקובץ ה-YAML הבא. הקונטיינר הזה מבוסס על פרויקט JetStream של OSS. מידע נוסף על דגלים של MaxText-JetStream זמין במאמר דגלים של שרת MaxText-JetStream. בדוגמה הבאה נעשה שימוש בצ'יפים של Trillium‏ (v6e-16) כדי לטעון את נקודת הבדיקה Llama3.1-405b int8 ולבצע עליה הסקה. בדוגמה הזו מניחים שכבר יש לכם אשכול GKE עם לפחות v6e-16 מאגר צמתים אחד בתוכו.

הפעלת שרת מודלים ו-Pathways

  1. מקבלים פרטי כניסה לאשכול ומוסיפים אותם להקשר המקומי של kubectl. 
          gcloud container clusters get-credentials $CLUSTER \
      --zone=$ZONE \
      --project=$PROJECT \
      && kubectl config set-context --current --namespace=default
  2. פריסת LeaderWorkerSet ‏ (LWS) API. 
          VERSION=v0.4.0
      kubectl apply --server-side -f "https://github.com/kubernetes-sigs/lws/releases/download/${VERSION}/manifests.yaml"
  3. מעתיקים את קוד ה-YAML הבא לקובץ בשם pathways-job.yaml: קוד ה-YAML הזה עבר אופטימיזציה לצורת הפרוסה v6e-16. מידע נוסף על המרת נקודת ביקורת של Meta לנקודת ביקורת שתואמת ל-JAX זמין במדריך ליצירת נקודות ביקורת במאמר יצירת נקודות ביקורת של הסקה. לדוגמה, הוראות ל-Llama3.1-405B מפורטות כאן Checkpoint conversion for Llama3.1-405B. 
        apiVersion: leaderworkerset.x-k8s.io/v1
    kind: LeaderWorkerSet
    metadata:
      name: jetstream-pathways
      annotations:
        leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      replicas: 1
      leaderWorkerTemplate:
        leaderTemplate:
          metadata:
            labels:
              app: jetstream-pathways
          spec:
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR_TYPE # Example: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY # Example: 4x4
            tolerations:
            - key: "google.com/tpu"
              operator: "Exists"
              effect: "NoSchedule"
            containers:
            - name: pathways-proxy
              image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
              args:
              - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
              - --server_port=38681
              - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
              imagePullPolicy: Always
              ports:
              - containerPort: 38681
            - name: pathways-rm
              env:
              - name: HOST_ADDRESS
                value: "$(LWS_LEADER_ADDRESS)"
              - name: TPU_SKIP_MDS_QUERY
                value: "true"
              image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
              args:
              - --server_port=38677
              - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
              - --node_type=resource_manager
              - --instance_count=1
              - --instance_type=tpuv6e:TPU_TOPOLOGY # Example: 4x4
              - --temporary_flags_for_debugging=temporary_flag_for_debugging_worker_expected_tpu_chip_config=megachip_tccontrol
              imagePullPolicy: Always
              ports:
              - containerPort: 38677
            - name: jax-tpu
              image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:jax-0.8.0
              # Optimized settings used to serve Llama3.1-405b.
              args:
              - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
              - model_name=llama3.1-405b
              - load_parameters_path=CHECKPOINT_PATH
              - max_prefill_predict_length=1024
              - max_target_length=2048
              - async_checkpointing=false
              - steps=1
              - ici_fsdp_parallelism=1
              - ici_autoregressive_parallelism=2
              - ici_tensor_parallelism=8
              - scan_layers=false
              - weight_dtype=bfloat16
              - per_device_batch_size=10
              - enable_single_controller=true
              - quantization=int8
              - quantize_kvcache=true
              - checkpoint_is_quantized=true
              - enable_model_warmup=true
              imagePullPolicy: Always
              ports:
              - containerPort: 9000
              startupProbe:
                httpGet:
                  path: /healthcheck
                  port: 8000
                  scheme: HTTP
                periodSeconds: 1
                initialDelaySeconds: 900
                failureThreshold: 10000
              livenessProbe:
                httpGet:
                  path: /healthcheck
                  port: 8000
                  scheme: HTTP
                periodSeconds: 60
                failureThreshold: 10
              readinessProbe:
                httpGet:
                  path: /healthcheck
                  port: 8000
                  scheme: HTTP
                periodSeconds: 60
                failureThreshold: 10
            - name: jetstream-http
              image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.4
              imagePullPolicy: Always
              ports:
              - containerPort: 8000
        # The size variable defines the number of worker nodes to be created.
        # It must be equal to the number of hosts + 1 (for the leader node).
        size: 5
        workerTemplate:
          spec:
            nodeSelector:
              cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR_TYPE # Example: tpu-v6e-slice
              cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY # Example: 4x4
            tolerations:
            - key: "google.com/tpu"
              operator: "Exists"
              effect: "NoSchedule"
            containers:
            - name: worker
              args:
              - --server_port=38679
              - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
              - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
              image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
              imagePullPolicy: Always
              ports:
              - containerPort: 38679
              resources:
                limits:
                  google.com/tpu: "4"
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: jetstream-svc
    spec:
      selector:
        app: jetstream-pathways
      ports:
      - protocol: TCP
        name: jetstream-http
        port: 8000
        targetPort: 8000
    מחליפים את מה שכתוב בשדות הבאים:
    • ‫TPU_ACCELERATOR_TYPE: סוג מאיץ ה-TPU. לדוגמה, tpu-v6e-slice.
    • ‫TPU_TOPOLOGY: טופולוגיית ה-TPU. לדוגמה, 2x4.
    • ‫CHECKPOINT_PATH: הנתיב ב-Cloud Storage לנקודת הבדיקה.
    מחילים את ה-YAML הזה. מחכים עד שה-LeaderWorkerSet יתוזמן. אחרי שתזמנו את ההגדרה, יכול להיות שיעבור זמן עד ששרת המודל ישחזר את נקודת הבדיקה. במודל 405B, התהליך נמשך כ-7 דקות.
  4. בודקים ביומני Kubernetes אם שרת המודל JetStream מוכן: עומס העבודה נקרא jetstream-pathways ב-YAML הקודם, ו-0 הוא צומת הראש. 
          kubectl logs -f jetstream-pathways-0 -c jax-tpu
    הפלט אמור להיראות כך, מה שמצביע על כך ששרת מודל JetStream מוכן לטפל בבקשות: 
        2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
    2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
    2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
    2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
    2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
    2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
    ...
    ...
    ...
    INFO:     Started server process [7]
    INFO:     Waiting for application startup.
    INFO:     Application startup complete.
    INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)

התחברות לשרת המודל

אפשר לגשת לפריסת JetStream Pathways באמצעות שירות ClusterIP של GKE. אפשר להגיע לשירות ClusterIP רק מתוך האשכול. לכן, כדי לגשת לשירות מחוץ לאשכול, צריך קודם להריץ את הפקודה הבאה כדי ליצור סשן של העברת נתונים דרך יציאה:

kubectl port-forward pod/${HEAD_POD} 8000:8000

כדי לוודא שיש לכם גישה לשרת ה-HTTP של JetStream, פותחים טרמינל חדש ומריצים את הפקודה הבאה:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

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

{
    "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
}

היקש מפורק

הגשה מפורקת היא טכניקה להפעלת מודלים גדולים של שפה (LLM) שמפרידה בין שלבי ההשלמה והפענוח לתהליכים שונים, שיכולים להתבצע במכונות שונות. כך אפשר לנצל טוב יותר את המשאבים ולשפר את הביצועים והיעילות, במיוחד במודלים גדולים.

  • מילוי מראש: בשלב הזה, ההנחיה לקלט מעובדת ונוצר ייצוג ביניים (כמו מטמון של זוגות מפתח/ערך). לרוב, מדובר בתהליך שדורש הרבה משאבי מחשוב.
  • ‫Decode: בשלב הזה נוצרים טוקנים של הפלט, אחד אחרי השני, באמצעות הייצוג של ה-prefill. בדרך כלל הביצועים מוגבלים על ידי רוחב הפס של הזיכרון.

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

כדי להפעיל את ההצגה המפוצלת, משנים את קובץ ה-YAML הבא כך שישתמש בשני v6e-8 חלקי נתונים: אחד למילוי מראש ואחד ליצירה. לפני שממשיכים, צריך לוודא שבאשכול GKE מוגדרות לפחות שתי קבוצות של צמתים עם טופולוגיה v6e-8 כזו. כדי להשיג ביצועים אופטימליים, הוגדרו דגלים ספציפיים של XLA.

יוצרים נקודת ביקורת של llama2-70b באותו תהליך שבו יוצרים נקודת ביקורת של llama3.1-405b, שמפורט בקטע הקודם.

  1. כדי להפעיל את שרת JetStream במצב מפורק באמצעות Pathways, מעתיקים את קוד ה-YAML הבא ומדביקים אותו בקובץ בשם pathways-job.yaml: 
    apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: jetstream-pathways
  annotations:
    leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
spec:
  replicas: 1
  leaderWorkerTemplate:
    subGroupPolicy:
      subGroupSize: 2
    leaderTemplate:
      metadata:
        labels:
          app: jetstream-pathways
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR_TYPE # Example: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY # Example: 2x4
        tolerations:
        - key: "google.com/tpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: pathways-proxy
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
          args:
          - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
          - --server_port=38681
          - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
          imagePullPolicy: Always
          ports:
          - containerPort: 38681
        - name: pathways-rm
          env:
          - name: HOST_ADDRESS
            value: "$(LWS_LEADER_ADDRESS)"
          - name: TPU_SKIP_MDS_QUERY
            value: "true"
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
          args:
          - --server_port=38677
          - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
          - --node_type=resource_manager
          - --instance_count=2
          - --instance_type=tpuv6e:TPU_TOPOLOGY # Example: 2x4
          imagePullPolicy: Always
          ports:
          - containerPort: 38677
        - name: jax-tpu
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:jax-0.8.0
          # Optimized settings used to serve Llama2-70b.
          args:
          - MaxText/configs/base.yml
          - tokenizer_path=assets/tokenizer.llama2
          - load_parameters_path=CHECKPOINT_PATH
          - max_prefill_predict_length=1024
          - max_target_length=2048
          - model_name=llama2-70b
          - ici_fsdp_parallelism=1
          - ici_autoregressive_parallelism=1
          - ici_tensor_parallelism=-1
          - scan_layers=false
          - weight_dtype=bfloat16
          - per_device_batch_size=1
          - checkpoint_is_quantized=true
          - quantization=int8
          - quantize_kvcache=true
          - compute_axis_order=0,2,1,3
          - ar_cache_axis_order=0,2,1,3
          - stack_prefill_result_cache=True
          # Specify disaggregated mode to run Jetstream
          - inference_server=ExperimentalMaxtextDisaggregatedServer_8
          - inference_benchmark_test=True
          - enable_model_warmup=True
          env:
          - name: LOG_LEVEL
            value: "INFO"
          imagePullPolicy: Always
          securityContext:
            capabilities:
              add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
          ports:
          - containerPort: 9000
          startupProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 1
            initialDelaySeconds: 240
            failureThreshold: 10000
          livenessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 100
          readinessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 100
        - name: jetstream-http
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.4
          imagePullPolicy: Always
          ports:
          - containerPort: 8000
    # The size variable defines the number of worker nodes to be created.
    # It must be equal to the number of hosts + 1 (for the leader node).
    size: 5
    workerTemplate:
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR_TYPE # Example: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY # Example: 2x4
        containers:
        - name: worker
          args:
          - --server_port=38679
          - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
          - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
          imagePullPolicy: Always
          ports:
          - containerPort: 38679
          resources:
            limits:
              google.com/tpu: "4"
---
apiVersion: v1
kind: Service
metadata:
  name: jetstream-svc
spec:
  selector:
    app: jetstream-pathways
  ports:
  - protocol: TCP
    name: jetstream-http
    port: 8000
    targetPort: 8000
    מחליפים את מה שכתוב בשדות הבאים:
    • ‫TPU_ACCELERATOR_TYPE: סוג מאיץ ה-TPU. לדוגמה, tpu-v6e-slice.
    • ‫TPU_TOPOLOGY: טופולוגיית ה-TPU. לדוגמה, 2x4.
    • ‫CHECKPOINT_PATH: הנתיב ב-Cloud Storage לנקודת הבדיקה.
  2. אחרי שמחילים את ה-YAML הזה, לשרת המודל ייקח זמן לשחזר את נקודת הבדיקה. במודל 70B, התהליך עשוי להימשך כ-2 דקות. 
      kubectl apply -f pathways-job.yaml
  3. בודקים ביומנים של Kubernetes אם שרת המודלים של JetStream מוכן: 
        kubectl logs -f jetstream-pathways-0 -c jax-tpu
    יוצג פלט שדומה לזה שבהמשך, שמציין ששרת המודל JetStream מוכן לטפל בבקשות: 
        2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
    2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
    2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
    2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
    2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
    2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
    ...
    ...
    ...
    INFO:     Started server process [7]
    INFO:     Waiting for application startup.
    INFO:     Application startup complete.
    INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)

התחברות לשרת המודל

אפשר לגשת לפריסת JetStream Pathways דרך שירות ClusterIP של GKE. אפשר להגיע לשירות ClusterIP רק מתוך האשכול. לכן, כדי לגשת לשירות מחוץ לאשכול, צריך להקים סשן של העברת נתונים דרך יציאה על ידי הפעלת הפקודה הבאה:

kubectl port-forward pod/${HEAD_POD} 8000:8000

כדי לוודא שיש לכם גישה לשרת ה-HTTP של JetStream, פותחים טרמינל חדש ומריצים את הפקודה הבאה:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

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

{
    "response": " used in software development?\nThe top 5 programming languages used in software development are:\n\n1. Java: Java is a popular programming language used for developing enterprise-level applications, Android apps, and web applications. Its platform independence and ability to run on any device that has a Java Virtual Machine (JVM) installed make it a favorite among developers.\n2. Python: Python is a versatile language that is widely used in software development, data analysis, artificial intelligence, and machine learning. Its simplicity, readability, and ease of use make it a popular choice among developers.\n3. JavaScript: JavaScript is a widely used programming language for web development, allowing developers to create interactive client-side functionality for web applications. It is also used for server-side programming, desktop and mobile application development, and game development.\n4. C++: C++ is a high-performance programming language used for developing operating systems, games, and other high-performance applications."
}

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