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

Hex-LLM, eine hocheffizientes LLM-Bereitstellung (Large Language Model) mit XLA, ist das Vertex AI LLM-Bereitstellungs-Framework, das für Cloud TPU-Hardware entwickelt und optimiert wurde. Hex-LLM kombiniert LLM-Bereitstellungstechnologien wie die kontinuierliche Batchverarbeitung und PagedAttention mit Vertex AI-Optimierungen, 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 Googles eigenen Optimierungen 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, die für XLA optimiert sind.
  • 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 spärlichen 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 Format Hugging Face-Modelle. 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-Gewichte laden, indem es das Anfrageargument während der Bereitstellung liest.

Erweiterte Funktionen

Hex-LLM unterstützt die folgenden erweiterten Funktionen:

  • Bereitstellung über mehrere Hosts
  • Aufgeschlüsselte Auslieferung [experimentell]
  • Präfix-Caching
  • Unterstützung für 4-Bit-Quantisierung

Bereitstellung über mehrere 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 einzelne TPU-VM mit einem Host geladen werden können, die maximal acht v5e-Kerne enthält.

Um diese Funktion zu aktivieren, legen Sie --num_hosts in den Hex-LLM-Containerargumenten und --tpu_topology in der Vertex AI SDK-Modell-Upload-Anfrage fest. Das folgende Beispiel zeigt, wie Sie den Hex-LLM-Container mit einer TPU 4x4 v5e-Topologie bereitstellen, die das Llama 3.1 70B-Modell mit bfloat16-Datentyp 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,
)

Ein End-to-End-Tutorial zum Bereitstellen des Hex-LLM-Containers mit einer TPU-Topologie mit mehreren Hosts finden Sie im Vertex AI Model Garden-Notebook – Llama 3.1 (Deployment).

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. Legen Sie das Argument --num_hosts auf die Anzahl der Hosts in der TPU-Topologie fest.
  3. Legen Sie --tpu_topology mit der API zum Hochladen von Modellen im Vertex AI SDK fest.

Aufgeschlüsselte Auslieferung [experimentell]

Hex-LLM unterstützt jetzt die disaggregierte Bereitstellung als experimentelles Feature. Sie kann nur auf dem Einzelhost-Setup aktiviert werden und die Leistung wird noch 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 gesamten Bereitstellungsdurchsatz auszugleichen. Dabei werden die Vorabfüllphase und die Decodierungsphase in separate Arbeitslasten unterteilt, damit sie sich nicht gegenseitig beeinträchtigen. Diese Methode ist besonders nützlich für Szenarien mit strengen Latenzanforderungen.

Wenn Sie diese Funktion aktivieren möchten, legen Sie --disagg_topo in den Hex-LLM-Containerargumenten fest. Das folgende Beispiel zeigt, wie Sie den Hex-LLM-Container auf TPU v5e-8 bereitstellen, der das Llama 3.1 8B-Modell mit bfloat16-Datentyp 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 es auf "3,1" festgelegt, um drei Prefill-Worker und einen Decode-Worker zu konfigurieren. Jeder Worker verwendet zwei TPU v5e-Kerne.

Präfix-Caching

Durch das Zwischenspeichern von Präfixen wird die Zeit bis zum ersten Token (Time to First Token, TTFT) für Prompts mit identischen Inhalten am Anfang des Prompts verkürzt, z. B. unternehmensweite Präambeln, allgemeine Systemanweisungen und der Verlauf von Unterhaltungen mit mehreren Zügen. Anstatt dieselben Eingabetokens wiederholt zu verarbeiten, kann Hex-LLM einen temporären Cache der verarbeiteten Eingabetokenberechnungen beibehalten, um die TTFT zu verbessern.

Wenn Sie diese Funktion aktivieren möchten, legen Sie --enable_prefix_cache_hbm in den Hex-LLM-Containerargumenten fest. Das folgende Beispiel zeigt, wie Sie den Hex-LLM-Container auf TPU v5e-8 bereitstellen, der das Llama 3.1-Modell mit 8 Milliarden Parametern in bfloat16 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 erfolgen in Schritten dieses Werts. So wird sichergestellt, dass 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
  }
}

Blockweise Vorgabe

Beim Chunked Prefill wird der Prefill einer Anfrage in kleinere Chunks aufgeteilt und Prefill und Dekodierung werden in einem Batch-Schritt kombiniert. Hex-LLM implementiert Chunked Prefill, 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 --enable_chunked_prefill in den Hex-LLM-Containerargumenten fest. Das folgende Beispiel zeigt, wie Sie den Hex-LLM-Container auf TPU v5e-8 bereitstellen, der das Llama 3.1-Modell mit 8 Milliarden Parametern 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 zur Reduzierung der Rechen- und Speicherkosten für die Ausführung von Inferenzen, indem die Gewichte 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 Gewichte. Der erweiterte Support umfasst Modelle mit INT4-Gewichten, 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 integriert. 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 Vertex AI-Endpunkt, der durch das Senden von Anfragen auf der Modellkarte erreichbar ist.

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

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

Jetzt mit Gemma ausprobieren

Bereitstellung mit nur einem Klick verwenden

Sie können einen benutzerdefinierten Vertex AI-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 wird, 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 Vertex AI-Endpunkt mit Hex-LLM mithilfe des Vertex AI 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 durch, 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 Enterprise-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 gestalten. Um die Argumente anzupassen, 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 Modellgewichte verwenden. BitsAndBytes-int8- und AWQ-quantisierte Modellartefakte werden für Llama, Gemma 2 und Mistral/Mixtral unterstützt.
  • --tokenizer: Der Tokenizer, der geladen werden soll. 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, sofern verfügbar. Die langsamen Tokenizer sind in Python geschrieben und in der Transformers-Bibliothek enthalten, während die schnellen Tokenizer, die eine Leistungssteigerung bieten, in Rust geschrieben und in der Tokenizers-Bibliothek enthalten sind. Weitere Informationen finden Sie in der Hugging Face-Dokumentation.
  • --trust_remote_code: Gibt an, ob Remotecodedateien, die in den Hugging Face-Modell-Repositories definiert sind, zulässig sind. Der Standardwert ist False.
  • --load_format: Format der zu ladenden Modell-Checkpoints. Mögliche Optionen sind ["auto", "dummy"]. Der Standardwert ist "auto". Wenn dieser Wert auf "auto" festgelegt ist, werden die Modellgewichte im Safetensors-Format geladen. Wenn dieser Wert auf "dummy" festgelegt ist, werden die Modellgewichte zufällig initialisiert. Wenn Sie diese Option auf "dummy" setzen, 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 festgelegt, überschreibt dieses Argument die Fenstergröße des Modells für Sliding Window Attention. Wenn Sie dieses Argument auf einen größeren Wert setzen, werden mehr Tokens in den Aufmerksamkeitsmechanismus einbezogen, was dem Effekt der standardmäßigen Selbstaufmerksamkeit entspricht. Dieses Argument ist nur für experimentelle Zwecke vorgesehen. 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 sich die Tokens ändern, die als nächste Tokens ausgewählt 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-Worker und Decode-Worker mit der experimentellen Funktion „Disaggregated Serving“. Der Standardwert ist None. Das Argument hat das Format "number_of_prefill_workers,number_of_decode_workers".
  • --data_parallel_size: Die Anzahl der Datenparallelreplikate. Der Standardwert ist 1. Wenn Sie diesen Wert von 1 auf N ändern, wird der Durchsatz um N verbessert, während die Latenz gleich bleibt.
  • --tensor_parallel_size: Die Anzahl der Tensor-Parallel-Replikate. Der Standardwert ist 1. Durch Erhöhen der Anzahl der Tensor-parallelen Replikate wird die Latenz in der Regel 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 Modul multiprocessing 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. Wenn Sie --no-enable_jit festlegen, wird sie deaktiviert. Durch Aktivieren des JIT-Modus wird die Inferenzleistung verbessert, allerdings ist für die erste Kompilierung mehr Zeit erforderlich. Im Allgemeinen überwiegen die Vorteile der Inferenzleistung den Mehraufwand.
  • --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 die Funktion deaktiviert. Ein Warm-up wird empfohlen, da bei den ersten Anfragen eine stärkere Kompilierung ausgelöst wird und sie daher langsamer sind.
  • --max_prefill_seqs: Die maximale Anzahl von Sequenzen, die pro Iteration für das Vorabfü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 Prefill-Batchgröße 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 der Rechenaufwand und der Inferenz-Overhead steigen. 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 werden die Zeiten für die Neukompilierung des Modells verkürzt, aber der Rechenaufwand und der Overhead für die Inferenz 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 Dekodierung 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 füllt die Batchgröße für die Dekodierung 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 der Rechenaufwand und der Inferenz-Overhead steigen. Die optimale Einstellung hängt vom Anfragetraffic ab.
  • --decode_blocks_padding: Der Server erhöht die Anzahl der Arbeitsspeicherblöcke, die für den Key-Value-Cache (KV-Cache) einer Sequenz verwendet werden, während der Dekodierung auf ein Vielfaches dieses Werts. Der Standardwert ist 128. Durch Erhöhen dieses Werts werden die Zeiten für die Neukompilierung des Modells verkürzt, aber der Rechenaufwand und der Overhead für die Inferenz 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 gemeinsamer Präfixe vorheriger Anfragen wiederverwendet werden.
  • --enable_chunked_prefill: Gibt an, ob Chunked Prefill aktiviert werden soll. Der Standardwert ist False. Wenn Sie dieses Argument festlegen, können Sie eine längere Kontextlänge unterstützen und die Leistung verbessern.

Arbeitsspeicherverwaltung

  • --hbm_utilization_factor: Der Prozentsatz des freien Cloud TPU High Bandwidth Memory (HBM), der nach dem Laden der Modellgewichte 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. Dadurch steigt jedoch das Risiko, dass während der Initialisierung und zur Laufzeit der Cloud TPU-HBM-Speicherplatz nicht ausreicht.
  • --num_blocks: Anzahl der Geräteblöcke, die für den KV-Cache zugewiesen werden sollen. Wenn dieses Argument festgelegt ist, wird --hbm_utilization_factor vom Server ignoriert. Wenn dieses Argument nicht festgelegt ist, wird die HBM-Nutzung der Serverprofile berechnet und die Anzahl der zuzuweisenden Geräteblöcke basierend auf --hbm_utilization_factor berechnet. 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 auch das Risiko, dass während der Initialisierung und zur Laufzeit der HBM-Speicher der Cloud TPU nicht ausreicht.
  • --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 Aufwand für die Blockverwaltung verringert, allerdings auf Kosten von mehr Speicherplatzverschwendung. Die genauen Auswirkungen auf die Leistung müssen empirisch ermittelt werden.

Dynamische 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 flexibler LoRA-Adapter verwenden, die mit dem Server kompatibel sind. Dadurch wird jedoch die Menge an Cloud TPU-HBM erhöht, die für LoRA-Gewichte zugewiesen wird, und der Durchsatz sinkt.
  • --enable_lora_cache: Gibt an, ob das Caching dynamischer LoRA-Adapter aktiviert werden soll. Der Standardwert ist True. Wenn Sie --no-enable_lora_cache festlegen, wird sie deaktiviert. Durch das Caching wird die Leistung verbessert, 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-Arbeitsspeicher-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 Menge der 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 abstimmen

Die Serverargumente sind miteinander verknüpft 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. Außerdem werden einige Argumente durch den Anwendungsfall bestimmt, während andere angepasst 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 benötigten TPU-Speichers basierend auf der Modellgröße und ‑genauigkeit: model_size * (num_bits / 8). Für ein 8B-Modell und bfloat16-Präzision liegt die Untergrenze des benötigten TPU-Arbeitsspeichers bei 8 * (16 / 8) = 16 GB.
  3. Schätzen Sie die Anzahl der benötigten TPU v5e-Chips. Jeder v5e-Chip bietet 16 GB: tpu_memory / 16. Für ein 8B-Modell und bfloat16-Präzision benötigen Sie mehr als einen Chip. Von den Konfigurationen mit 1, 4 und 8 Chips ist die Konfiguration mit 4 Chips die kleinste Konfiguration mit mehr als einem Chip: ct5lp-hightpu-4t. Sie können --tensor_parallel_size=4 später festlegen.
  4. Bestimmen Sie die maximale Kontextlänge (Eingabelänge + Ausgabelänge) für den vorgesehenen Anwendungsfall. Beispiel: 4096. Sie können --max_model_len=4096 später festlegen.
  5. Stimmen Sie die Menge an freiem TPU-Speicher, der für den KV-Cache zugewiesen wird, auf den maximalen Wert ab, der bei den Modell-, Hardware- und Serverkonfigurationen möglich ist (--hbm_utilization_factor). Beginnen Sie mit 0.95. Stellen Sie den Hex-LLM-Server bereit und testen Sie ihn mit langen Prompts und hoher Parallelität. Wenn der Server nicht mehr genügend Arbeitsspeicher hat, reduzieren Sie den Auslastungsfaktor entsprechend.

Hier ist ein Beispiel für eine Reihe von Argumenten für die Bereitstellung 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

Ein Beispielsatz von Argumenten für die Bereitstellung von Llama 3.1 70B Instruct AWQ auf ct5lp-hightpu-4t ist:

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

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