Offene Modelle mit Hex-LLM-Premium-Container in Cloud TPU bereitstellen

Hex-LLM ist eine hocheffiziente LLM-Bereitstellung (Large Language Model) mit XLA und das für Cloud TPU-Hardware entwickelte und optimierte LLM-Bereitstellungsframework der Gemini Enterprise Agent Platform. Hex-LLM kombiniert LLM-Bereitstellungstechnologien wie die kontinuierliche Batchverarbeitung und PagedAttention mit Optimierungen der Gemini Enterprise Agent Platform, die auf XLA und Cloud TPU zugeschnitten sind. Dies ist eine hocheffiziente und kostengünstige LLM-Bereitstellungsoption für Cloud TPUs für Open-Source-Modellen.

Hex-LLM ist in Model Garden per Modell-Playground, Bereitstellung mit einem Klick und Notebook verfügbar.

Features

Hex-LLM basiert auf Open-Source-Projekten mit Optimierungen von Google für XLA und Cloud TPU. Hex-LLM erreicht bei der Bereitstellung häufig verwendeter LLMs einen hohen Durchsatz und eine niedrige Latenz.

Hex-LLM umfasst die folgenden Optimierungen:

  • Tokenbasierter Algorithmus für kontinuierliches Batching, um sicherzustellen, dass Modelle die Hardware bei einer großen Anzahl gleichzeitiger Anfragen optimal nutzen.
  • Eine vollständige Neuschreibung der Attention-Kernel, für XLA optimiert.
  • Flexible und zusammensetzbare Strategien für Datenparallelität und Tensor-Parallelität die hochoptimierte Methoden zur Gewichtungsfragmentierung nutzen, um LLMs effizient auf mehreren Cloud TPU-Chips auszuführen.

Hex-LLM unterstützt eine Vielzahl von dichten und dünnbesetzten LLMs:

  • Gemma 2B und 7B
  • Gemma-2 9B und 27B
  • Llama-2 7B, 13B und 70B
  • Llama-3 8B und 70B
  • Llama-3.1 8B und 70B
  • Llama-3.2 1B und 3B
  • Llama-3.3 70B
  • Llama-Guard-3 1B und 8B
  • Llama-4 Scout-17B-16E
  • Mistral 7B
  • Mixtral 8x7B und 8x22B
  • Phi-3 Mini und Medium
  • Phi-4, Phi-4 Reasoning und Reasoning Plus
  • Qwen-2 0.5B, 1.5B und 7B
  • Qwen-2.5 0.5B, 1.5B, 7B, 14B und 32B

Hex-LLM bietet außerdem eine Vielzahl von Features, z. B.:

  • Hex-LLM ist in einem einzelnen Container enthalten. Hex-LLM verpackt den API-Server, die Inferenz-Engine und die unterstützten Modelle in einem einzigen, bereitzustellenden Docker-Image.
  • Kompatibel mit dem Hugging Face-Modellformat. Hex-LLM kann ein Hugging Face-Modell von einem lokalen Laufwerk, dem Hugging Face Hub oder einem Cloud Storage-Bucket laden.
  • Quantisierung mit bitsandbytes und AWQ.
  • Dynamisches Laden von LoRA. Hex-LLM kann die LoRA-Gewichtungen laden, indem es das Anfrageargument während der Bereitstellung liest.

Erweiterte Funktionen

Hex-LLM unterstützt die folgenden erweiterten Funktionen:

  • Bereitstellung auf mehreren Hosts
  • Disaggregierte Bereitstellung [experimentell]
  • Präfix-Caching
  • Unterstützung für 4-Bit-Quantisierung

Bereitstellung auf mehreren Hosts

Hex-LLM unterstützt jetzt die Bereitstellung von Modellen mit einem TPU-Slice mit mehreren Hosts. Mit dieser Funktion können Sie große Modelle bereitstellen, die nicht in eine TPU-VM mit einem einzelnen Host geladen werden können, die maximal acht v5e-Kerne enthält.

Wenn Sie diese Funktion aktivieren möchten, legen Sie --num_hosts in den Hex-LLM-Containerargumenten und --tpu_topology in der Modellanfrage zum Hochladen des Vertex AI SDK fest. Im folgenden Beispiel wird gezeigt, wie Sie den Hex-LLM-Container mit einer TPU 4x4 v5e-Topologie bereitstellen, die das Llama 3.1 70B-bfloat16-Modell bereitstellt:

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,
)

Eine End-to-End-Anleitung zum Bereitstellen des Hex-LLM-Containers mit einer TPU-Topologie mit mehreren Hosts finden Sie im Agent Platform Model Garden – Llama 3.1 (Bereitstellung).

Im Allgemeinen sind nur die folgenden Änderungen erforderlich, um die Bereitstellung auf mehreren Hosts zu aktivieren:

  1. Legen Sie das Argument --tensor_parallel_size auf die Gesamtzahl der Kerne in der TPU-Topologie fest.
  2. Setzen Sie das Argument --num_hosts auf die Anzahl der Hosts in der TPU-Topologie.
  3. Legen Sie --tpu_topology mit der Vertex AI SDK-API für den Modellupload fest.

Disaggregierte Bereitstellung [experimentell]

Hex-LLM unterstützt jetzt die disaggregierte Bereitstellung als experimentelles Feature. Sie kann nur für die Einrichtung mit einem einzelnen Host aktiviert werden und die Leistung wird optimiert.

Die disaggregierte Bereitstellung ist eine effektive Methode, um die Zeit bis zum ersten Token (Time to First Token, TTFT) und die Zeit pro Ausgabetoken (Time Per Output Token, TPOT) für jede Anfrage sowie den Gesamtdurchsatz der Bereitstellung auszugleichen. Dabei werden die Vorbefüllungs- und die Decodierungsphase in verschiedene Arbeitslasten aufgeteilt, damit sie sich nicht gegenseitig beeinträchtigen. Diese Methode ist besonders nützlich für Szenarien mit strengen Latenzanforderungen.

Setzen Sie zum Aktivieren dieser Funktion --disagg_topo in den Hex-LLM-Containerargumenten. Das folgende Beispiel zeigt, wie Sie den Hex-LLM-Container auf TPU v5e-8 bereitstellen, der das Llama 3.1 8B bfloat16-Modell bereitstellt:

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,
)

Das Argument --disagg_topo akzeptiert einen String im Format "number_of_prefill_workers,number_of_decode_workers". Im vorherigen Beispiel ist der Wert auf "3,1" festgelegt, um drei Prefill-Worker und einen Decode-Worker zu konfigurieren. Jeder Worker verwendet zwei TPU v5e-Kerne.

Präfix-Caching

Das Präfix-Caching reduziert die Zeit bis zum ersten Token (Time to First Token, TTFT) für Prompts, die am Anfang identischen Inhalt haben, z. B. unternehmensweite Präambeln, allgemeine Systemanweisungen und den Verlauf von Unterhaltungen mit mehreren Runden. Anstatt dieselben Eingabetokens wiederholt zu verarbeiten, kann Hex-LLM einen temporären Cache der verarbeiteten Eingabetoken-Berechnungen beibehalten, um die TTFT zu verbessern.

Wenn Sie diese Funktion aktivieren möchten, legen Sie in den Hex-LLM-Containerargumenten --enable_prefix_cache_hbm fest. Das folgende Beispiel zeigt, wie Sie den Hex-LLM-Container auf einer TPU v5e-8 bereitstellen, die das Llama 3.1 8B bfloat16-Modell bereitstellt:

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 verwendet Präfix-Caching, um die Leistung für Prompts zu optimieren, die eine bestimmte Länge überschreiten (standardmäßig 512 Tokens, konfigurierbar mit prefill_len_padding). Cache-Treffer treten in Schritten dieses Werts auf, sodass die Anzahl der im Cache gespeicherten Tokens immer ein Vielfaches von prefill_len_padding ist. Das Feld cached_tokens von usage.prompt_tokens_details in der API-Antwort für die Chat-Vervollständigung gibt an, wie viele der Prompt-Tokens ein Cache-Treffer waren.

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

Aufgeteilte Vorbefüllung

Beim aufgeteilten Prefill wird ein Prefill einer Anfrage in kleinere Teile aufgeteilt und Prefill und Decodierung werden in einem Batchschritt kombiniert. Hex-LLM implementiert die fragmentierte Vorbefüllung, um die Zeit bis zum ersten Token (Time to First Token, TTFT) und die Zeit pro Ausgabetoken (Time Per Output Token, TPOT) auszugleichen und den Durchsatz zu verbessern.

Wenn Sie diese Funktion aktivieren möchten, legen Sie in den Hex-LLM-Containerargumenten --enable_chunked_prefill fest. Im folgenden Beispiel wird gezeigt, wie Sie den Hex-LLM-Container auf einer TPU v5e-8 bereitstellen, die das Llama 3.1 8B-Modell bereitstellt:

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,
)

Unterstützung für 4-Bit-Quantisierung

Die Quantisierung ist eine Technik, mit der die Rechen- und Speicherkosten für die Ausführung von Inferenzen reduziert werden, indem die Gewichtungen oder Aktivierungen mit Datentypen mit niedriger Genauigkeit wie INT8 oder INT4 anstelle der üblichen BF16 oder FP32 dargestellt werden.

Hex-LLM unterstützt die INT8-Quantisierung nur für Gewichtungen. Der erweiterte Support umfasst Modelle mit INT4-Gewichtungen, die mit der AWQ-Nullpunktquantisierung quantisiert wurden. Hex-LLM unterstützt INT4-Varianten der Modellfamilien Mistral, Mixtral und Llama.

Für die Bereitstellung quantisierter Modelle ist kein zusätzliches Flag erforderlich.

Erste Schritte mit Model Garden

Der Hex-LLM-Cloud-TPU-Bereitstellungscontainer ist in Model Garden eingebunden. Sie können auf diese Bereitstellungstechnologie über die Playground-, die Ein-Klick-Bereitstellung und die Colab Enterprise-Notebook-Beispiele für eine Vielzahl von Modellen zugreifen.

Playground verwenden

Model Garden Playground ist ein vorab bereitgestellter Agent Platform-Endpunkt, der durch das Senden von Anfragen auf der Modellkarte erreichbar ist.

  1. Geben Sie einen Prompt und optional Argumente für Ihre Anfrage ein.

  2. Klicken Sie auf SENDEN, um die Modellantwort schnell zu erhalten.

Probieren Sie es mit Gemma aus.

Bereitstellung mit nur einem Klick verwenden

Sie können einen benutzerdefinierten Agent Platform-Endpunkt mit Hex-LLM mithilfe einer Modellkarte bereitstellen.

  1. Rufen Sie die Seite „Modellkarte“ auf und klicken Sie auf Bereitstellen.

  2. Wählen Sie für die zu verwendende Modellvariante den Cloud TPU v5e-Maschinentyp zur Bereitstellung aus.

  3. Klicken Sie unten auf Bereitstellen, um den Bereitstellungsprozess zu starten. Sie erhalten zwei E-Mail-Benachrichtigungen: eine, wenn das Modell hochgeladen wurde, und eine weitere, wenn der Endpunkt bereit ist.

Colab Enterprise-Notebook verwenden

Für mehr Flexibilität und Anpassungsmöglichkeiten können Sie Colab Enterprise-Notebook-Beispiele verwenden, um einen Agent Platform-Endpunkt mit Hex-LLM mithilfe des Agent Platform SDK für Python bereitzustellen.

  1. Rufen Sie die Seite „Modellkarte“ auf und klicken Sie auf Notebook öffnen.

  2. Wählen Sie das Vertex Serving-Notebook aus. Das Notebook wird in Colab Enterprise geladen.

  3. Führen Sie das Notebook aus, um ein Modell mit Hex-LLM bereitzustellen und Vorhersageanfragen an den Endpunkt zu senden. Das Code-Snippet für die Bereitstellung sieht so aus:

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,
)

Beispiele für Colab-Notebooks:

Serverargumente und Umgebungsvariablen konfigurieren

Sie können die folgenden Argumente festlegen, um den Hex-LLM-Server zu starten. Sie können die Argumente an Ihren beabsichtigten Anwendungsfall und Ihre Anforderungen anpassen. Die Argumente sind für die Bereitstellung mit einem Klick vordefiniert, um die Bereitstellung so einfach wie möglich zu machen. Wenn Sie die Argumente anpassen möchten, können Sie die Notebook-Beispiele als Referenz verwenden und die Argumente entsprechend festlegen.

Modell

  • --model: Das zu ladende Modell. Sie können eine Hugging Face-Modell-ID, einen Cloud Storage-Bucket-Pfad (gs://my-bucket/my-model) oder einen lokalen Pfad angeben. Die Modellartefakte müssen dem Hugging Face-Format entsprechen und safetensors-Dateien für die Modellgewichtungen verwenden. BitsAndBytes-int8- und AWQ-quantisierte Modellartefakte werden für Llama, Gemma 2 und Mistral/Mixtral unterstützt.
  • --tokenizer: Der zu ladende Tokenizer. Dies kann eine Hugging Face-Modell-ID, ein Cloud Storage-Bucket-Pfad (gs://my-bucket/my-model) oder ein lokaler Pfad sein. Wenn dieses Argument nicht festgelegt ist, wird standardmäßig der Wert für --model verwendet.
  • --tokenizer_mode: Der Tokenizer-Modus. Mögliche Optionen sind ["auto", "slow"]. Der Standardwert ist "auto". Wenn diese Option auf "auto" gesetzt ist, wird der schnelle Tokenizer verwendet, falls verfügbar. Die langsamen Tokenizer sind in Python geschrieben und in der Transformers-Bibliothek enthalten. Die schnellen Tokenizer, die eine Leistungsverbesserung bieten, sind in Rust geschrieben und in der Tokenizers-Bibliothek enthalten. Weitere Informationen finden Sie in der Hugging Face-Dokumentation.
  • --trust_remote_code: Gibt an, ob Remote-Codedateien, die in den Hugging Face-Modell-Repositories definiert sind, zugelassen werden sollen. Der Standardwert ist False.
  • --load_format: Format der zu ladenden Modellprüfpunkte. Mögliche Optionen sind ["auto", "dummy"]. Der Standardwert ist "auto". Wenn dieser Wert auf "auto" gesetzt ist, werden die Modellgewichtungen im Safetensors-Format geladen. Wenn dieser Wert auf "dummy" gesetzt ist, werden die Modellgewichtungen zufällig initialisiert. Wenn Sie hier "dummy" festlegen, ist das für Tests nützlich.
  • --max_model_len: Die maximale Kontextlänge (Eingabelänge plus Ausgabelänge), die für das Modell bereitgestellt werden soll. Der Standardwert wird aus der Modellkonfigurationsdatei im Hugging Face-Format gelesen: config.json. Eine größere maximale Kontextlänge erfordert mehr TPU-Arbeitsspeicher.
  • --sliding_window: Wenn dieses Argument festgelegt ist, wird die Fenstergröße des Modells für die Sliding Window Attention überschrieben. Wenn Sie dieses Argument auf einen größeren Wert setzen, werden mehr Tokens in den Aufmerksamkeitsmechanismus einbezogen und der Effekt nähert sich dem der Standard-Self-Attention an. Dieses Argument ist nur für die experimentelle Verwendung gedacht. In allgemeinen Anwendungsfällen empfehlen wir, die ursprüngliche Fenstergröße des Modells zu verwenden.
  • --seed: Der Startwert für die Initialisierung aller Zufallszahlengeneratoren. Wenn Sie dieses Argument ändern, kann sich das auf die generierte Ausgabe für denselben Prompt auswirken, da die Tokens geändert werden, die als nächste Tokens herangezogen werden. Der Standardwert ist 0.

Inferenz-Engine

  • --num_hosts: Die Anzahl der auszuführenden Hosts. Der Standardwert ist 1. Weitere Informationen finden Sie in der Dokumentation zur TPU v5e-Konfiguration.
  • --disagg_topo: Definiert die Anzahl der Prefill- und Decodierungs-Worker mit dem experimentellen Feature der disaggregierten Bereitstellung. Der Standardwert ist None. Das Argument hat das Format "number_of_prefill_workers,number_of_decode_workers".
  • --data_parallel_size: Die Anzahl der datenparallelen Replikate. Der Standardwert ist 1. Wenn Sie diesen Wert von 1 auf N setzen, wird der Durchsatz um etwa N verbessert, während die Latenz gleich bleibt.
  • --tensor_parallel_size: Die Anzahl der tensorparallelen Replikate. Der Standardwert ist 1. Durch Erhöhen der Anzahl der tensorparallelen Replikate wird die Latenz im Allgemeinen verbessert, da die Matrixmultiplikation durch Verringern der Matrixgröße beschleunigt wird.
  • --worker_distributed_method: Die verteilte Methode zum Starten des Workers. Verwenden Sie mp für das Multiprocessing-Modul oder ray für die Ray-Bibliothek. Der Standardwert ist mp.
  • --enable_jit: Gibt an, ob der JIT-Modus (Just-in-Time-Kompilierung) aktiviert werden soll. Der Standardwert ist True. Mit der Einstellung --no-enable_jit wird sie deaktiviert. Durch die Aktivierung des JIT-Modus wird die Inferenzleistung verbessert, allerdings ist für die anfängliche Kompilierung zusätzliche Zeit erforderlich. Im Allgemeinen überwiegen die Vorteile der Inferenzleistung gegenüber dem Overhead.
  • --warmup: Gibt an, ob der Server während der Initialisierung mit Beispielanfragen aufgewärmt werden soll. Der Standardwert ist True. Wenn Sie --no-warmup festlegen, wird sie deaktiviert. Das Aufwärmen wird empfohlen, da die ersten Anfragen eine umfangreichere Kompilierung auslösen und daher langsamer sind.
  • --max_prefill_seqs: Die maximale Anzahl von Sequenzen, die pro Iteration für das Vorbefüllen geplant werden können. Der Standardwert ist 1. Je größer dieser Wert ist, desto höher ist der Durchsatz, den der Server erreichen kann, allerdings mit möglichen negativen Auswirkungen auf die Latenz.
  • --prefill_seqs_padding: Der Server füllt die Batchgröße für die Vorabfüllung auf ein Vielfaches dieses Werts auf. Der Standardwert ist 8. Durch Erhöhen dieses Werts wird die Zeit für die Neukompilierung des Modells verkürzt, aber die verschwendete Rechenleistung und der Inferenz-Overhead werden erhöht. Die optimale Einstellung hängt vom Anfragetraffic ab.
  • --prefill_len_padding: Der Server füllt die Sequenzlänge auf ein Vielfaches dieses Werts auf. Der Standardwert ist 512. Durch Erhöhen dieses Werts wird die Zeit für die Neukompilierung des Modells verkürzt, aber der Rechenaufwand und der Inferenz-Overhead steigen. Die optimale Einstellung hängt von der Datenverteilung der Anfragen ab.
  • --max_decode_seqs/--max_running_seqs: Die maximale Anzahl von Sequenzen, die pro Iteration für die Decodierung geplant werden können. Der Standardwert ist 256. Je größer dieser Wert ist, desto höher ist der Durchsatz, den der Server erreichen kann, allerdings mit möglichen negativen Auswirkungen auf die Latenz.
  • --decode_seqs_padding: Der Server erhöht die Batchgröße für die Dekodierung auf ein Vielfaches dieses Werts. Der Standardwert ist 8. Durch Erhöhen dieses Werts wird die Zeit für die Neukompilierung des Modells verkürzt, aber der Rechenaufwand und der Inferenz-Overhead steigen. Die optimale Einstellung hängt vom Anfragetraffic ab.
  • --decode_blocks_padding: Der Server füllt die Anzahl der Arbeitsspeicherblöcke, die für den Schlüssel/Wert-Cache (KV-Cache) einer Sequenz verwendet werden, während der Dekodierung auf ein Vielfaches dieses Werts auf. Der Standardwert ist 128. Durch Erhöhen dieses Werts wird die Zeit für die Neukompilierung des Modells verkürzt, aber der Rechenaufwand und der Inferenz-Overhead steigen. Die optimale Einstellung hängt von der Datenverteilung der Anfragen ab.
  • --enable_prefix_cache_hbm: Gibt an, ob Präfix-Caching in HBM aktiviert werden soll. Der Standardwert ist False. Wenn Sie dieses Argument festlegen, kann die Leistung verbessert werden, da die Berechnungen von freigegebenen Präfixen früherer Anfragen wiederverwendet werden.
  • --enable_chunked_prefill: Gibt an, ob aufgeteiltes Prefill aktiviert werden soll. Der Standardwert ist False. Wenn Sie dieses Argument festlegen, kann eine längere Kontextlänge unterstützt und die Leistung verbessert werden.

Speicherverwaltung

  • --hbm_utilization_factor: Der Prozentsatz des verfügbaren Cloud TPU High Bandwidth Memory (HBM), der nach dem Laden der Modellgewichtungen für den KV-Cache zugewiesen werden kann. Der Standardwert ist 0.9. Wenn Sie dieses Argument auf einen höheren Wert setzen, wird die Größe des KV-Cache erhöht und der Durchsatz kann verbessert werden. Allerdings steigt dadurch das Risiko, dass während der Initialisierung und zur Laufzeit der Cloud TPU-HBM-Speicher ausgeht.
  • --num_blocks: Anzahl der Geräteblöcke, die für den KV-Cache zugewiesen werden sollen. Wenn dieses Argument festgelegt ist, ignoriert der Server --hbm_utilization_factor. Wenn dieses Argument nicht festgelegt ist, erstellt der Server ein Profil der HBM-Nutzung und berechnet die Anzahl der zuzuweisenden Geräteblöcke basierend auf --hbm_utilization_factor. Wenn Sie dieses Argument auf einen höheren Wert festlegen, wird die KV-Cache-Größe erhöht und der Durchsatz kann verbessert werden. Allerdings steigt das Risiko, dass während der Initialisierung und zur Laufzeit der Cloud TPU-HBM-Speicher ausgeht.
  • --block_size: Anzahl der in einem Block gespeicherten Tokens. Mögliche Optionen sind [8, 16, 32, 2048, 8192]. Der Standardwert ist 32. Wenn Sie dieses Argument auf einen größeren Wert festlegen, wird der Overhead bei der Blockverwaltung reduziert, allerdings auf Kosten einer höheren Speicherverschwendung. Die genauen Auswirkungen auf die Leistung müssen empirisch ermittelt werden.

Dynamisches LoRA

  • --enable_lora: Gibt an, ob das dynamische Laden von LoRA-Adaptern aus Cloud Storage aktiviert werden soll. Der Standardwert ist False. Dies wird für die Llama-Modellfamilie unterstützt.
  • --max_lora_rank: Der maximale LoRA-Rang, der für LoRA-Adapter unterstützt wird, in Anfragen definiert. Der Standardwert ist 16. Wenn Sie dieses Argument auf einen höheren Wert setzen, können Sie die LoRA-Adapter flexibler mit dem Server verwenden. Allerdings wird dadurch mehr Cloud TPU-HBM für LoRA-Gewichtungen zugewiesen und der Durchsatz verringert.
  • --enable_lora_cache: Gibt an, ob das Caching dynamischer LoRA-Adapter aktiviert werden soll. Der Standardwert ist True. Mit der Einstellung --no-enable_lora_cache wird sie deaktiviert. Caching verbessert die Leistung, da zuvor verwendete LoRA-Adapterdateien nicht noch einmal heruntergeladen werden müssen.
  • --max_num_mem_cached_lora: Die maximale Anzahl von LoRA-Adaptern, die im TPU-Speicher-Cache gespeichert werden.Der Standardwert ist 16. Wenn Sie dieses Argument auf einen größeren Wert setzen, erhöht sich die Wahrscheinlichkeit eines Cache-Treffers, aber auch die Cloud TPU HBM-Nutzung.

Sie können den Server auch mit den folgenden Umgebungsvariablen konfigurieren:

  • HEX_LLM_LOG_LEVEL: Hiermit wird festgelegt, wie viele Logging-Informationen generiert werden. Der Standardwert ist INFO. Legen Sie hier eine der standardmäßigen Python-Logging-Ebenen fest, die im Logging-Modul definiert sind.
  • HEX_LLM_VERBOSE_LOG: Gibt an, ob die detaillierte Logging-Ausgabe aktiviert werden soll. Zulässige Werte sind true und false. Der Standardwert ist false.

Serverargumente optimieren

Die Serverargumente hängen zusammen und wirken sich gemeinsam auf die Bereitstellungsleistung aus. Eine größere Einstellung von --max_model_len=4096 führt beispielsweise zu einer höheren TPU-Arbeitsspeichernutzung und erfordert daher eine größere Arbeitsspeicherzuweisung und weniger Batching. Darüber hinaus werden einige Argumente durch den Anwendungsfall bestimmt, während andere optimiert werden können. Hier ist ein Workflow zum Konfigurieren des Hex-LLM-Servers.

  1. Bestimmen Sie die gewünschte Modellfamilie und Modellvariante. Beispiel: Llama 3.1 8B Instruct.
  2. Schätzen Sie die Untergrenze des erforderlichen TPU-Arbeitsspeichers basierend auf der Modellgröße und der Genauigkeit: model_size * (num_bits / 8). Für ein 8B-Modell und bfloat16-Genauigkeit beträgt die Untergrenze des erforderlichen TPU-Arbeitsspeichers 8 * (16 / 8) = 16 GB.
  3. Schätzen Sie die Anzahl der benötigten TPU v5e-Chips, wobei jeder v5e-Chip 16 GB bietet: tpu_memory / 16. Für ein 8B-Modell und bfloat16-Genauigkeit benötigen Sie mehr als einen Chip. Unter den Konfigurationen mit 1, 4 und 8 Chips ist die kleinste Konfiguration mit mehr als 1 Chip die Konfiguration mit 4 Chips: ct5lp-hightpu-4t. Anschließend können Sie --tensor_parallel_size=4 festlegen.
  4. Bestimmen Sie die maximale Kontextlänge (Eingabelänge + Ausgabelänge) für den vorgesehenen Anwendungsfall. Beispiel: 4096. Anschließend können Sie --max_model_len=4096 festlegen.
  5. Passen Sie die Menge des verfügbaren TPU-Arbeitsspeichers, der für den KV-Cache zugewiesen wird, auf den maximalen Wert an, der mit den Modell-, Hardware- und Serverkonfigurationen (--hbm_utilization_factor) erreicht werden kann. Beginnen Sie mit 0.95. Stellen Sie den Hex-LLM-Server bereit und testen Sie ihn mit langen Prompts und hoher Nebenläufigkeit. Wenn der Server nicht mehr genügend Arbeitsspeicher hat, reduzieren Sie den Auslastungsfaktor entsprechend.

Hier ist ein Beispiel für Argumente zum Bereitstellen von 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

Hier ist ein Beispiel für Argumente zum Bereitstellen von Llama 3.1 70B Instruct AWQ auf 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-Kontingent anfordern

In Model Garden beträgt Ihr Standardkontingent 32 Cloud TPU v5e-Chips in der Region us-west1. Diese Kontingente gelten für Bereitstellungen mit nur einem Klick und Colab Enterprise-Notebook-Bereitstellungen. Informationen zum Anfordern eines höheren Kontingentwerts finden Sie unter Kontingentanpassung anfordern.