הצגת מודלים פתוחים באמצעות קונטיינר פרימיום של Hex-LLM ב-Cloud TPU

‫Hex-LLM הוא מודל שפה גדול (LLM) יעיל במיוחד שמשתמש ב-XLA, והוא מסגרת ה-LLM של Gemini Enterprise Agent Platform שמיועדת לחומרת Cloud TPU ועברה אופטימיזציה בשבילה. ‫Hex-LLM משלב טכנולוגיות להצגת LLM, כמו אצווה רציפה ו-PagedAttention, עם אופטימיזציות של Gemini Enterprise Agent Platform שמותאמות ל-XLA ול-Cloud TPU. זהו מודל LLM יעיל וזול שפועל ב-Cloud TPU למודלים בקוד פתוח.

‫Hex-LLM זמין ב-Model Garden דרך מגרש המשחקים של המודלים, פריסה בלחיצה אחת ו-Notebook.

תכונות

‫Hex-LLM מבוסס על פרויקטים בקוד פתוח עם אופטימיזציות של Google ל-XLA ול-Cloud TPU. ‫Hex-LLM משיג תפוקה גבוהה וזמן אחזור קצר כשמציגים מודלים גדולים של שפה שמשתמשים בהם לעיתים קרובות.

‫Hex-LLM כולל את האופטימיזציות הבאות:

  • אלגוריתם של אצווה רציפה (continuous batching) שמבוסס על טוקנים, כדי להבטיח שהמודלים מנצלים את החומרה באופן מלא עם מספר גדול של בקשות מקבילות.
  • שכתוב מלא של ליבות הקשב (attention kernels) שעברו אופטימיזציה ל-XLA.
  • אסטרטגיות גמישות וקומפוזביליות של מקביליות נתונים ומקביליות טנזור, עם שיטות אופטימיזציה של פיצול משקלים, כדי להריץ ביעילות מודלי LLM בכמה שבבי Cloud TPU.

‫Hex-LLM תומך במגוון רחב של מודלים גדולים של שפה (LLM) דחוסים ודלילים:

  • ‫Gemma 2B ו-7B
  • ‫Gemma-2 9B ו-27B
  • ‫Llama-2 7B, ‏ 13B ו-70B
  • ‫Llama-3 8B ו-70B
  • ‫Llama-3.1 8B ו-70B
  • ‫Llama-3.2 1B ו-3B
  • Llama-3.3 70B
  • ‫Llama-Guard-3 1B ו-8B
  • Llama-4 Scout-17B-16E
  • Mistral 7B
  • ‫Mixtral 8x7B ו-8x22B
  • ‫Phi-3 mini ו-Phi-3 medium
  • ‫Phi-4, ‏ Phi-4 reasoning ו-reasoning plus
  • ‫Qwen-2 0.5B, ‏ 1.5B ו-7B
  • ‫Qwen-2.5 0.5B, ‏ 1.5B, ‏ 7B, ‏ 14B ו-32B

‫Hex-LLM מספק גם מגוון תכונות, כמו:

  • ‫Hex-LLM כלול בקונטיינר אחד. ‫Hex-LLM אורז את שרת ה-API, מנוע ההסקה והמודלים הנתמכים בקובץ אימג' אחד של Docker לפריסה.
  • תואם לפורמט של מודלים של Hugging Face. ‫Hex-LLM יכול לטעון מודל של Hugging Face מדיסק מקומי, מ-Hugging Face Hub ומקטגוריה של Cloud Storage.
  • כימות באמצעות bitsandbytes ו-AWQ.
  • טעינה דינמית של LoRA. ‫Hex-LLM יכול לטעון את משקלי ה-LoRA על ידי קריאת הארגומנט של הבקשה במהלך הצגת המודל.

תכונות מתקדמות

‫Hex-LLM תומך בתכונות המתקדמות הבאות:

  • הצגה בכמה מארחים
  • הצגה מפורקת [ניסיוני]
  • שמירת קידומות במטמון
  • תמיכה בכימות של 4 ביטים

הצגה בכמה מארחים

עכשיו אפשר להשתמש ב-Hex-LLM כדי להציג מודלים עם פרוסת TPU בכמה מארחים. התכונה הזו מאפשרת להכניס לשימוש בסביבת הייצור מודלים גדולים שאי אפשר לטעון ל-TPU VM במארח יחיד, שמכילה לכל היותר שמונה ליבות v5e.

כדי להפעיל את התכונה הזו, מגדירים את --num_hosts בארגומנטים של מאגר Hex-LLM ואת --tpu_topology בבקשה להעלאת מודל של Vertex AI SDK. בדוגמה הבאה מוסבר איך לפרוס את מאגר Hex-LLM עם טופולוגיית TPU 4x4 v5e שמפעילה את המודל Llama 3.1 70B bfloat16:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Meta-Llama-3.1-70B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=16",
    "--num_hosts=4",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=["python", "-m", "hex_llm.server.api_server"],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    tpu_topology="4x4",
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

ב-Agent Platform Model Garden - Llama 3.1 (Deployment) notebook מוסבר איך לפרוס את קונטיינר Hex-LLM עם טופולוגיית TPU בכמה מארחים.

באופן כללי, השינויים היחידים שצריך לבצע כדי להפעיל הצגה בכמה מארחים הם:

  1. מגדירים את הארגומנט --tensor_parallel_size למספר הכולל של הליבות בטופולוגיית ה-TPU.
  2. מגדירים את הארגומנט --num_hosts למספר המארחים בטופולוגיית ה-TPU.
  3. מגדירים את --tpu_topology באמצעות ה-API להעלאת מודלים של Vertex AI SDK.

הצגה מפורקת [ניסיוני]

‫Hex-LLM תומך עכשיו בהצגה מפוצלת כתכונה ניסיונית. אפשר להפעיל אותה רק בהגדרה של מארח יחיד, והביצועים שלה נמצאים בתהליך כוונון.

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

כדי להפעיל את התכונה הזו, מגדירים את --disagg_topo בארגומנטים של מאגר Hex-LLM. בדוגמה הבאה מוסבר איך לפרוס את קונטיינר Hex-LLM ב-TPU v5e-8 שמציג את המודל Llama 3.1 8B bfloat16:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=2",
    "--disagg_topo=3,1",
    "--hbm_utilization_factor=0.9",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=["python", "-m", "hex_llm.server.api_server"],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

הארגומנט --disagg_topo מקבל מחרוזת בפורמט "number_of_prefill_workers,number_of_decode_workers". בדוגמה הקודמת, הוא מוגדר ל-"3,1" כדי להגדיר שלושה עובדי מילוי מראש ועובד פענוח אחד. כל וורקר משתמש בשתי ליבות TPU v5e.

שמירת קידומות במטמון

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

כדי להפעיל את התכונה הזו, מגדירים את --enable_prefix_cache_hbm בארגומנטים של מאגר Hex-LLM. בדוגמה הבאה מוסבר איך לפרוס את קונטיינר Hex-LLM ב-TPU v5e-8 שמציג את המודל Llama 3.1 8B bfloat16:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--hbm_utilization_factor=0.9",
    "--enable_prefix_cache_hbm",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=["python", "-m", "hex_llm.server.api_server"],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

‫Hex-LLM משתמש בשמירת קידומות במטמון כדי לשפר את הביצועים של הנחיות שחורגות מאורך מסוים (512 טוקנים כברירת מחדל, אפשר להגדיר באמצעות prefill_len_padding). מציאות במטמון מתרחשות במרווחים של הערך הזה, כדי להבטיח שמספר הטוקנים שנשמרים במטמון הוא תמיד כפולה של prefill_len_padding. השדה cached_tokens של usage.prompt_tokens_details בתגובה מה-API להשלמת הצ'אט מציין כמה מהאסימונים של ההנחיה נמצאו במטמון.

"usage": {
  "prompt_tokens": 643,
  "total_tokens": 743,
  "completion_tokens": 100,
  "prompt_tokens_details": {
    "cached_tokens": 512
  }
}

מילוי מראש בחלקים

בשיטת Chunked prefill, המילוי מראש של הבקשה מפוצל לחלקים קטנים יותר, והמילוי מראש והפענוח משולבים לשלב אחד באצווה. ‫Hex-LLM מיישם מילוי מראש בחלקים כדי לאזן בין הזמן עד לאסימון הראשון (TTFT) לבין הזמן לכל אסימון פלט (TPOT), ומשפר את התפוקה.

כדי להפעיל את התכונה הזו, מגדירים את --enable_chunked_prefill בארגומנטים של מאגר Hex-LLM. בדוגמה הבאה מוסבר איך לפרוס את הקונטיינר של Hex-LLM ב-TPU v5e-8 שמציג את המודל Llama 3.1 8B:

hexllm_args = [
    "--host=0.0.0.0",
    "--port=7080",
    "--model=meta-llama/Llama-3.1-8B",
    "--data_parallel_size=1",
    "--tensor_parallel_size=4",
    "--hbm_utilization_factor=0.9",
    "--enable_chunked_prefill",
]

model = aiplatform.Model.upload(
    display_name=model_name,
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=["python", "-m", "hex_llm.server.api_server"],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=env_vars,
    serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB
    serving_container_deployment_timeout=7200,
    location=TPU_DEPLOYMENT_REGION,
)

model.deploy(
    endpoint=endpoint,
    machine_type=machine_type,
    deploy_request_timeout=1800,
    service_account=service_account,
    min_replica_count=min_replica_count,
    max_replica_count=max_replica_count,
)

תמיכה בכימות של 4 ביטים

קוונטיזציה היא שיטה להפחתת עלויות החישוב והזיכרון של הרצת הסקה, על ידי ייצוג המשקלים או ההפעלות באמצעות סוגי נתונים ברמת דיוק נמוכה כמו INT8 או INT4, במקום BF16 או FP32 הרגילים.

‫Hex-LLM תומך בכימות של משקולות INT8 בלבד. התמיכה המורחבת כוללת מודלים עם משקלי INT4 שעברו קוונטיזציה באמצעות קוונטיזציה של נקודת אפס של AWQ. ‫Hex-LLM תומך בגרסאות INT4 של משפחות המודלים Mistral, ‏ Mixtral ו-Llama.

לא נדרש דגל נוסף כדי להציג מודלים שעברו קוונטיזציה.

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

מאגר הצגת המודעות של Hex-LLM Cloud TPU משולב ב-Model Garden. אתם יכולים לגשת לטכנולוגיית ההצגה הזו דרך סביבת הניסוי, פריסה בלחיצה אחת ודוגמאות ל-notebook ב-Colab Enterprise למגוון מודלים.

שימוש ב-Playground

‫Model Garden Playground הוא נקודת קצה (endpoint) של Agent Platform שנפרסה מראש, ואפשר להגיע אליה על ידי שליחת בקשות בכרטיס המודל.

  1. מזינים הנחיה ואם רוצים, מוסיפים ארגומנטים לבקשה.

  2. לוחצים על SUBMIT כדי לקבל את התשובה של המודל במהירות.

רוצים לנסות את Gemma?

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

אתם יכולים לפרוס נקודת קצה מותאמת אישית של Agent Platform עם Hex-LLM באמצעות כרטיס מודל.

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

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

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

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

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

  1. עוברים לדף של כרטיס המודל ולוחצים על Open notebook.

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

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

hexllm_args = [
    f"--model=google/gemma-2-9b-it",
    f"--tensor_parallel_size=4",
    f"--hbm_utilization_factor=0.8",
    f"--max_running_seqs=512",
]
hexllm_envs = {
    "PJRT_DEVICE": "TPU",
    "MODEL_ID": "google/gemma-2-9b-it",
    "DEPLOY_SOURCE": "notebook",
}
model = aiplatform.Model.upload(
    display_name="gemma-2-9b-it",
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=[
        "python", "-m", "hex_llm.server.api_server"
    ],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=hexllm_envs,
    serving_container_shared_memory_size_mb=(16 * 1024),
    serving_container_deployment_timeout=7200,
)

endpoint = aiplatform.Endpoint.create(display_name="gemma-2-9b-it-endpoint")
model.deploy(
    endpoint=endpoint,
    machine_type="ct5lp-hightpu-4t",
    deploy_request_timeout=1800,
    service_account="<your-service-account>",
    min_replica_count=1,
    max_replica_count=1,
)

דוגמאות ל-notebooks ב-Colab:

הגדרת ארגומנטים של שרת ומשתני סביבה

כדי להפעיל את שרת Hex-LLM, אפשר להגדיר את הארגומנטים הבאים. אתם יכולים להתאים את הארגומנטים לתרחיש לדוגמה ולדרישות שלכם. שימו לב שהארגומנטים מוגדרים מראש לפריסה בלחיצה אחת, כדי לאפשר את חוויית הפריסה הקלה ביותר. כדי להתאים אישית את הארגומנטים, אתם יכולים להשתמש בדוגמאות ל-Notebook כקובצי עזר ולהגדיר את הארגומנטים בהתאם.

מודל

  • ‫--model: המודל שרוצים לטעון. אפשר לציין מזהה מודל של Hugging Face, נתיב לקטגוריה של Cloud Storage‏ (gs://my-bucket/my-model) או נתיב מקומי. רכיבי ה-Artifact של המודל צריכים להיות בפורמט של Hugging Face, ומשקלי המודל צריכים להיות בקובצי safetensors. יש תמיכה ברכיבי מודל שעברו קוונטיזציה של BitsAndBytes int8 ו-AWQ ב-Llama, ב-Gemma 2 וב-Mistral או Mixtral.
  • ‫--tokenizer: הטוקנייזר שרוצים לטעון. זה יכול להיות מזהה מודל של Hugging Face, נתיב לקטגוריה של Cloud Storage (gs://my-bucket/my-model) או נתיב מקומי. אם לא מגדירים את הארגומנט הזה, ברירת המחדל היא הערך של --model.
  • --tokenizer_mode: מצב ה-tokenizer. האפשרויות הן ["auto", "slow"]. ערך ברירת המחדל הוא "auto". אם המדיניות מוגדרת לערך "auto", המערכת תשתמש ב-tokenizer המהיר, אם הוא זמין. הטוקנייזרים האיטיים כתובים ב-Python ומסופקים בספריית Transformers, והטוקנייזרים המהירים, שמציעים שיפור בביצועים, כתובים ב-Rust ומסופקים בספריית Tokenizers. מידע נוסף זמין במאמרי העזרה של Hugging Face.
  • ‫--trust_remote_code: אם לאפשר קובצי קוד מרוחקים שמוגדרים במאגרי המודלים של Hugging Face. ערך ברירת המחדל הוא False.
  • ‫--load_format: הפורמט של נקודות הבדיקה של המודל שרוצים לטעון. האפשרויות הן ["auto", "dummy"]. ערך ברירת המחדל הוא "auto". אם הערך הוא "auto", משקלי המודל נטענים בפורמט safetensors. אם הערך הוא "dummy", משקלי המודל מאותחלים באופן אקראי. ההגדרה "dummy" שימושית לניסויים.
  • ‫--max_model_len: חלון ההקשר המקסימלי (אורך הקלט בתוספת אורך הפלט) שהמודל יכול להציג. ערך ברירת המחדל נקרא מקובץ התצורה של המודל בפורמט Hugging Face: ‏config.json. ככל שחלון ההקשר המקסימלי גדול יותר, נדרש יותר זיכרון TPU.
  • ‫--sliding_window: אם הארגומנט הזה מוגדר, הוא מבטל את גודל החלון של המודל לקשב של חלון הזזה. אם מגדירים ערך גבוה יותר לארגומנט הזה, מנגנון תשומת הלב כולל יותר אסימונים ומתקרב להשפעה של תשומת לב עצמית רגילה. הארגומנט הזה מיועד לשימוש ניסיוני בלבד. במקרים כלליים, מומלץ להשתמש בגודל החלון המקורי של המודל.
  • ‫--seed: הגרעין לאתחול כל מחוללי המספרים האקראיים. שינוי הארגומנט הזה עשוי להשפיע על הפלט שנוצר לאותה הנחיה, כי הוא משנה את הטוקנים שנדגמים כטוקנים הבאים. ערך ברירת המחדל הוא 0.

מנוע הסקה

  • ‫--num_hosts: מספר המארחים שצריך להפעיל. ערך ברירת המחדל הוא 1. לפרטים נוספים, קראו את המאמר הגדרת TPU v5e.
  • ‫--disagg_topo: מגדירה את מספר ה-workers למילוי מראש ולפענוח בתכונה הניסיונית "הצגה מבוזרת". ערך ברירת המחדל הוא None. הארגומנט הוא בפורמט: "number_of_prefill_workers,number_of_decode_workers".
  • ‫--data_parallel_size: מספר הרפליקות של הנתונים במקביל. ערך ברירת המחדל הוא 1. הגדרה של N במקום 1 משפרת את קצב העברת הנתונים בערך ב-N, תוך שמירה על אותו זמן אחזור.
  • ‫--tensor_parallel_size: מספר הרפליקות המקבילות של הטנזור. ערך ברירת המחדל הוא 1. הגדלת מספר הרפליקות המקבילות של Tensor בדרך כלל משפרת את זמן האחזור, כי היא מאיצה את הכפלת המטריצה על ידי הקטנת גודל המטריצה.
  • ‫--worker_distributed_method: השיטה המבוזרת להפעלת ה-worker. משתמשים ב-mp למודול multiprocessing או ב-ray לספריית Ray. ערך ברירת המחדל הוא mp.
  • ‫--enable_jit: אם להפעיל את מצב JIT (הידור בזמן אמת). ערך ברירת המחדל הוא True. הערך --no-enable_jit משבית את המדיניות. הפעלת מצב JIT משפרת את ביצועי ההסקה, אבל דורשת זמן נוסף להשלמת ההידור הראשוני. באופן כללי, היתרונות של ביצועי ההסקה עולים על התקורה.
  • ‫--warmup: אם להכין את השרת עם בקשות לדוגמה במהלך האתחול. ערך ברירת המחדל הוא True. הערך --no-warmup משבית את המדיניות. מומלץ לבצע חימום, כי בקשות ראשוניות מפעילות הידור כבד יותר ולכן יהיו איטיות יותר.
  • ‫--max_prefill_seqs: המספר המקסימלי של רצפים שאפשר לתזמן למילוי מראש בכל איטרציה. ערך ברירת המחדל הוא 1. ככל שהערך הזה גדול יותר, כך השרת יכול להשיג קצב העברת נתונים גבוה יותר, אבל יכולות להיות לכך השפעות שליליות על זמן האחזור.
  • ‫--prefill_seqs_padding: השרת מרפד את גודל האצווה של המילוי מראש לכפולה של הערך הזה. ערך ברירת המחדל הוא 8. הגדלת הערך הזה מקצרת את זמני הקימפול מחדש של המודל, אבל מגדילה את החישובים המיותרים ואת התקורה של ההסקה. ההגדרה האופטימלית תלויה בתנועת הבקשות.
  • ‫--prefill_len_padding: השרת מוסיף ריפוד לאורך הרצף עד שהוא כפולה של הערך הזה. ערך ברירת המחדל הוא 512. הגדלת הערך הזה מקצרת את זמני הקימפול מחדש של המודל, אבל מגדילה את החישובים המיותרים ואת התקורה של ההסקה. ההגדרה האופטימלית תלויה בהתפלגות הנתונים של הבקשות.
  • ‫--max_decode_seqs/--max_running_seqs: המספר המקסימלי של רצפים שאפשר לתזמן לפענוח בכל איטרציה. ערך ברירת המחדל הוא 256. ככל שהערך הזה גדול יותר, כך השרת יכול להשיג קצב העברת נתונים גבוה יותר, אבל יכולות להיות לכך השפעות שליליות על זמן האחזור.
  • ‫--decode_seqs_padding: השרת מוסיף ריפוד לגודל האצווה של הפענוח עד שהוא מגיע לכפולה של הערך הזה. ערך ברירת המחדל הוא 8. הגדלת הערך הזה מקצרת את זמני הקימפול מחדש של המודל, אבל מגדילה את החישובים המיותרים ואת התקורה של ההסקה. ההגדרה האופטימלית תלויה בתנועת הבקשות.
  • ‫--decode_blocks_padding: השרת מוסיף ריפוד למספר בלוקי הזיכרון שמשמשים למטמון של צמדי מפתח/ערך (KV cache) של רצף, עד לכפולה של הערך הזה במהלך הפענוח. ערך ברירת המחדל הוא 128. הגדלת הערך הזה מקצרת את זמני הקימפול מחדש של המודל, אבל מגדילה את החישובים המיותרים ואת התקורה של ההסקה. ההגדרה האופטימלית תלויה בהתפלגות הנתונים של הבקשות.
  • ‫--enable_prefix_cache_hbm: אם להפעיל שמירת קידומת במטמון ב-HBM. ערך ברירת המחדל הוא False. הגדרת הארגומנט הזה יכולה לשפר את הביצועים על ידי שימוש חוזר בחישובים של קידומות משותפות של בקשות קודמות.
  • ‫--enable_chunked_prefill: אם להפעיל מילוי מראש בחלקים. ערך ברירת המחדל הוא False. הגדרת הארגומנט הזה יכולה לתמוך בחלון הקשר ארוך יותר ולשפר את הביצועים.

ניהול זיכרון

  • ‫--hbm_utilization_factor: אחוז הזיכרון הפנוי ב-Cloud TPU High Bandwidth Memory (HBM) שאפשר להקצות למטמון KV אחרי שטוענים את משקלי המודל. ערך ברירת המחדל הוא 0.9. אם מגדירים ערך גבוה יותר לארגומנט הזה, גודל המטמון של KV גדל וקצב הנתונים יכול להשתפר, אבל יש סיכון גבוה יותר שה-HBM של Cloud TPU ייגמר במהלך האיניציאליזציה ובזמן הריצה.
  • ‫--num_blocks: מספר בלוקי המכשירים שיוקצו למטמון KV. אם הארגומנט הזה מוגדר, השרת מתעלם מ---hbm_utilization_factor. אם הארגומנט הזה לא מוגדר, השרת יוצר פרופיל של השימוש ב-HBM ומחשב את מספר בלוקי המכשירים שצריך להקצות על סמך --hbm_utilization_factor. אם מגדירים לארגומנט הזה ערך גבוה יותר, גודל המטמון של KV גדל וקצב הנתונים יכול להשתפר, אבל הסיכון שייגמר ה-HBM של Cloud TPU במהלך האתחול ובזמן הריצה גדל.
  • ‫--block_size: מספר האסימונים ששמורים בבלוק. האפשרויות הן [8, 16, 32, 2048, 8192]. ערך ברירת המחדל הוא 32. הגדרה של ערך גבוה יותר לארגומנט הזה מפחיתה את התקורה בניהול הבלוקים, אבל גורמת לבזבוז של יותר זיכרון. כדי לקבוע את ההשפעה המדויקת על הביצועים, צריך לבצע בדיקה אמפירית.

Dynamic LoRA

  • ‫--enable_lora: אם להפעיל טעינה דינמית של מתאמי LoRA מ-Cloud Storage. ערך ברירת המחדל הוא False. האפשרות הזו נתמכת במשפחת המודלים של Llama.
  • ‫--max_lora_rank: הדירוג המקסימלי של LoRA שנתמך במתאמי LoRA שמוגדרים בבקשות. ערך ברירת המחדל הוא 16. אם מגדירים ערך גבוה יותר לארגומנט הזה, יש יותר גמישות במתאמי LoRA שאפשר להשתמש בהם עם השרת, אבל כמות ה-HBM של Cloud TPU שמוקצית למשקלים של LoRA גדלה וקצב הנתונים קטן.
  • ‫--enable_lora_cache: אם להפעיל שמירה במטמון של מתאמי LoRA דינמיים. ערך ברירת המחדל הוא True. הערך --no-enable_lora_cache משבית את המדיניות. השמירה במטמון משפרת את הביצועים כי היא מייתרת את הצורך להוריד מחדש קובצי מתאם LoRA שכבר השתמשתם בהם.
  • ‫--max_num_mem_cached_lora: המספר המקסימלי של מתאמי LoRA שמאוחסנים במטמון הזיכרון של ה-TPU.ערך ברירת המחדל הוא 16. הגדרה של ערך גבוה יותר לארגומנט הזה משפרת את הסיכוי למציאה במטמון (cache hit), אבל מגדילה את כמות השימוש ב-Cloud TPU HBM.

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

  • ‫HEX_LLM_LOG_LEVEL: קובעת את כמות המידע שנוצרת ביומן. ערך ברירת המחדל הוא INFO. מגדירים את זה לאחת מרמות הרישום ביומן הרגילות של Python שמוגדרות במודול הרישום ביומן.
  • ‫HEX_LLM_VERBOSE_LOG: אם להפעיל פלט של רישום מפורט ביומן. הערכים האפשריים הם true או false. ערך ברירת המחדל הוא false.

ארגומנטים של שרת Tune

הארגומנטים של השרת קשורים זה לזה ומשפיעים ביחד על ביצועי ההצגה. לדוגמה, הגדרה גבוהה יותר של --max_model_len=4096 מובילה לשימוש רב יותר בזיכרון של ה-TPU, ולכן נדרשת הקצאת זיכרון גדולה יותר ופחות אצווה. בנוסף, חלק מהארגומנטים נקבעים לפי תרחיש השימוש, ואחרים ניתנים לכוונון. כאן מתואר תהליך העבודה להגדרת שרת Hex-LLM.

  1. בוחרים את משפחת המודלים ואת גרסת המודל שרוצים להשתמש בה. לדוגמה, Llama 3.1 8B Instruct.
  2. מעריכים את הגבול התחתון של זיכרון ה-TPU שנדרש על סמך גודל המודל והדיוק: model_size * (num_bits / 8). למודל 8B ודיוק bfloat16, הגבול התחתון של זיכרון ה-TPU הנדרש יהיה 8 * (16 / 8) = 16 GB.
  3. מעריכים את מספר שבבי TPU v5e שצריך, כשכל שבב v5e מציע 16GB: tpu_memory / 16. למודל 8B ודיוק bfloat16, צריך יותר משבב אחד. מבין ההגדרות של שבב אחד, 4 שבבים ו-8 שבבים, ההגדרה הקטנה ביותר שמציעה יותר משבב אחד היא ההגדרה של 4 שבבים: ct5lp-hightpu-4t. אחר כך אפשר להגדיר את --tensor_parallel_size=4.
  4. קובעים את אורך ההקשר המקסימלי (אורך הקלט + אורך הפלט) לתרחיש לדוגמה הרצוי. לדוגמה, 4096. אחר כך אפשר להגדיר את --max_model_len=4096.
  5. מכוונים את כמות הזיכרון הפנוי של ה-TPU שמוקצה למטמון KV לערך המקסימלי שאפשר להשיג בהתאם למודל, לחומרה ולהגדרות השרת (--hbm_utilization_factor). מתחילים עם 0.95. פורסים את שרת Hex-LLM ובודקים אותו עם הנחיות ארוכות ורמת בו-זמניות גבוהה. אם השרת חורג ממגבלת הזיכרון, צריך להקטין את מקדם הניצול בהתאם.

דוגמה לקבוצת ארגומנטים לפריסת Llama 3.1 8B Instruct:

python -m hex_llm.server.api_server \
    --model=meta-llama/Llama-3.1-8B-Instruct \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.95

דוגמה לקבוצת ארגומנטים לפריסת Llama 3.1 70B Instruct AWQ ב-ct5lp-hightpu-4t:

python -m hex_llm.server.api_server \
    --model=hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4 \
    --tensor_parallel_size=4 \
    --max_model_len=4096
    --hbm_utilization_factor=0.45

בקשת מכסה ל-Cloud TPU

ב-Model Garden, מכסת ברירת המחדל היא 32 שבבי Cloud TPU v5e באזור us-west1. המכסות האלה חלות על פריסות בלחיצה אחת ועל פריסות של Notebook ב-Colab Enterprise. כדי לבקש להגדיל את המכסה, קראו את המאמר איך שולחים בקשה לשינוי המכסות.