Eroga modelli aperti utilizzando il container premium Hex-LLM su Cloud TPU

Hex-LLM, un modello linguistico di grandi dimensioni (LLM) ad alta efficienza erogato con XLA, è il framework di erogazione degli LLM di Gemini Enterprise Agent Platform, progettato e ottimizzato per l'hardware Cloud TPU. Hex-LLM combina tecnologie di erogazione degli LLM come il batching continuo e PagedAttention con ottimizzazioni di Gemini Enterprise Agent Platform personalizzate per XLA e Cloud TPU. È un metodo di erogazione degli LLM ad alta efficienza e dai costi ridotti su Cloud TPU per modelli open source.

Hex-LLM è disponibile in Model Garden tramite l'area di gioco dei modelli, il deployment con un solo clic e il notebook.

Funzionalità

Hex-LLM si basa su progetti open source con ottimizzazioni di Google per XLA e Cloud TPU. Hex-LLM raggiunge un throughput elevato e una bassa latenza durante l'erogazione di LLM utilizzati di frequente.

Hex-LLM include le seguenti ottimizzazioni:

  • Algoritmo di raggruppamento in batch continuo basato su token per garantire che i modelli utilizzino appieno l'hardware con un numero elevato di richieste in parallelo.
  • Una riscrittura completa dei kernel di attenzione ottimizzati per XLA.
  • Strategie di parallelismo dei dati e dei tensori flessibili e componibili con metodi di sharding dei pesi altamente ottimizzati per eseguire in modo efficiente gli LLM su più chip Cloud TPU.

Hex-LLM supporta un'ampia gamma di LLM densi e sparsi:

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

Hex-LLM offre anche una serie di funzionalità, tra cui:

  • Hex-LLM è incluso in un singolo container. Hex-LLM pacchettizza il server API, il motore di inferenza e i modelli supportati in un'unica immagine Docker di cui eseguire il deployment.
  • Compatibile con il formato dei modelli Hugging Face. Hex-LLM può caricare un modello Hugging Face dal disco locale, dall'hub Hugging Face e da un bucket Cloud Storage.
  • Quantizzazione utilizzando bitsandbytes e AWQ.
  • Caricamento dinamico di LoRA. Hex-LLM è in grado di caricare i pesi LoRA leggendo l'argomento della richiesta durante l'erogazione.

Funzionalità avanzate

Hex-LLM supporta le seguenti funzionalità avanzate:

  • Pubblicazione multi-host
  • Pubblicazione disaggregata [sperimentale]
  • Memorizzazione nella cache dei prefissi
  • Supporto della quantizzazione a 4 bit

Pubblicazione multi-host

Hex-LLM ora supporta l'erogazione di modelli con una sezione TPU multi-host. Questa funzionalità ti consente di gestire modelli di grandi dimensioni che non possono essere caricati in una singola VM TPU host, che contiene al massimo otto core v5e.

Per attivare questa funzionalità, imposta --num_hosts negli argomenti del container Hex-LLM e imposta --tpu_topology nella richiesta di caricamento del modello dell'SDK Vertex AI. L'esempio seguente mostra come eseguire il deployment del container Hex-LLM con una topologia TPU 4x4 v5e che gestisce il modello 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,
)

Per un tutorial end-to-end sul deployment del container Hex-LLM con una topologia TPU multi-host, consulta il notebook Agent Platform Model Garden - Llama 3.1 (Deployment).

In genere, le uniche modifiche necessarie per attivare la distribuzione multi-host sono:

  1. Imposta l'argomento --tensor_parallel_size sul numero totale di core all'interno della topologia TPU.
  2. Imposta l'argomento --num_hosts sul numero di host all'interno della topologia TPU.
  3. Imposta --tpu_topology con l'API di caricamento del modello dell'SDK Vertex AI.

Pubblicazione disaggregata [sperimentale]

Hex-LLM ora supporta l'erogazione disaggregata come funzionalità sperimentale. Può essere attivato solo nella configurazione a host singolo e le prestazioni sono in fase di ottimizzazione.

La distribuzione disaggregata è un metodo efficace per bilanciare il tempo al primo token (TTFT) e il tempo per token di output (TPOT) per ogni richiesta, nonché il throughput di distribuzione complessivo. Separa la fase di precompilazione e la fase di decodifica in workload diversi in modo che non interferiscano tra loro. Questo metodo è particolarmente utile per gli scenari che prevedono requisiti di latenza rigorosi.

Per attivare questa funzionalità, imposta --disagg_topo negli argomenti del container Hex-LLM. Di seguito è riportato un esempio che mostra come eseguire il deployment del container Hex-LLM su TPU v5e-8 che gestisce il modello 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,
)

L'argomento --disagg_topo accetta una stringa nel formato "number_of_prefill_workers,number_of_decode_workers". Nell'esempio precedente, è impostato su "3,1" per configurare tre worker di precompilazione e un worker di decodifica. Ogni worker utilizza due core TPU v5e.

Memorizzazione nella cache dei prefissi

La memorizzazione nella cache del prefisso riduce il tempo al primo token (TTFT) per i prompt che hanno contenuti identici all'inizio, come preamboli a livello aziendale, istruzioni di sistema comuni e cronologia delle conversazioni a più turni. Anziché elaborare ripetutamente gli stessi token di input, Hex-LLM può conservare una cache temporanea dei calcoli dei token di input elaborati per migliorare il TTFT.

Per attivare questa funzionalità, imposta --enable_prefix_cache_hbm negli argomenti del container Hex-LLM. Di seguito è riportato un esempio che mostra come eseguire il deployment del container Hex-LLM su TPU v5e-8 che gestisce il modello 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 utilizza la memorizzazione nella cache dei prefissi per ottimizzare le prestazioni per i prompt che superano una determinata lunghezza (512 token per impostazione predefinita, configurabile utilizzando prefill_len_padding). I successi della cache si verificano in incrementi di questo valore, garantendo che il conteggio dei token memorizzati nella cache sia sempre un multiplo di prefill_len_padding. Il campo cached_tokens di usage.prompt_tokens_details nella risposta dell'API di completamento della chat indica quanti token del prompt hanno generato un successo della cache.

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

Pre-fill suddiviso in chunk

Il pre-riempimento suddiviso in blocchi suddivide un pre-riempimento della richiesta in blocchi più piccoli e combina il pre-riempimento e la decodifica in un unico passaggio del batch. Hex-LLM implementa il pre-riempimento in blocchi per bilanciare il tempo al primo token (TTFT) e il tempo per token di output (TPOT) e migliorare il throughput.

Per attivare questa funzionalità, imposta --enable_chunked_prefill negli argomenti del container Hex-LLM. Di seguito è riportato un esempio che mostra come eseguire il deployment del container Hex-LLM su TPU v5e-8 che gestisce il modello 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,
)

Supporto della quantizzazione a 4 bit

La quantizzazione è una tecnica per ridurre i costi di calcolo e di memoria dell'esecuzione dell'inferenza rappresentando i pesi o le attivazioni con tipi di dati a bassa precisione come INT8 o INT4 anziché i consueti BF16 o FP32.

Hex-LLM supporta la quantizzazione solo per i pesi INT8. Il supporto esteso include modelli con pesi INT4 quantizzati utilizzando la quantizzazione zero-point AWQ. Hex-LLM supporta le varianti INT4 delle famiglie di modelli Mistral, Mixtral e Llama.

Non è richiesto alcun flag aggiuntivo per la distribuzione di modelli quantizzati.

Inizia a utilizzare Model Garden

Il container di distribuzione Hex-LLM Cloud TPU è integrato in Model Garden. Puoi accedere a questa tecnologia di distribuzione tramite il playground, il deployment con un solo clic e gli esempi di notebook Colab Enterprise per una varietà di modelli.

Utilizza Playground

Il playground di Model Garden è un endpoint di Agent Platform di cui è stato eseguito il deployment, raggiungibile inviando richieste nella scheda del modello.

  1. Inserisci un prompt e, facoltativamente, includi gli argomenti per la tua richiesta.

  2. Fai clic su INVIA per ricevere rapidamente la risposta del modello.

Fai una prova con Gemma.

Utilizza il deployment con un solo clic

Puoi eseguire il deployment di un endpoint Agent Platform personalizzato con Hex-LLM utilizzando una scheda del modello.

  1. Vai alla pagina della scheda del modello e fai clic su Esegui il deployment.

  2. Per la variante del modello che vuoi utilizzare, seleziona il tipo di macchina Cloud TPU v5e per il deployment.

  3. Fai clic su Esegui il deployment in basso per avviare il processo di deployment. Riceverai due notifiche via email: una quando il modello viene caricato e un'altra quando l'endpoint è pronto.

Utilizza il notebook Colab Enterprise

Per flessibilità e personalizzazione, puoi utilizzare gli esempi di notebook di Colab Enterprise per eseguire il deployment di un endpoint della piattaforma agente con Hex-LLM utilizzando l'SDK della piattaforma agente per Python.

  1. Vai alla pagina della model card e fai clic su Apri blocco note.

  2. Seleziona il notebook Vertex Serving. Il notebook viene aperto in Colab Enterprise.

  3. Esegui il notebook per eseguire il deployment di un modello utilizzando Hex-LLM e invia richieste di previsione all'endpoint. Lo snippet di codice per il deployment è il seguente:

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

Ecco alcuni esempi di notebook di Colab:

Configura gli argomenti del server e le variabili di ambiente

Puoi impostare i seguenti argomenti per avviare il server Hex-LLM. Puoi personalizzare gli argomenti in modo che si adattino al meglio al caso d'uso e ai requisiti previsti. Tieni presente che gli argomenti sono predefiniti per il deployment con un solo clic, in modo da offrire l'esperienza di deployment più semplice. Per personalizzare gli argomenti, puoi utilizzare gli esempi di notebook come riferimento e impostare gli argomenti di conseguenza.

Modello

  • --model: il modello da caricare. Puoi specificare un ID modello Hugging Face, un percorso del bucket Cloud Storage (gs://my-bucket/my-model) o un percorso locale. Gli artefatti del modello devono seguire il formato Hugging Face e utilizzare i file safetensors per i pesi del modello. Gli artefatti dei modelli quantizzati BitsAndBytes int8 e AWQ sono supportati per Llama, Gemma 2 e Mistral/Mixtral.
  • --tokenizer: il tokenizer da caricare. Può essere un ID modello Hugging Face, un percorso del bucket Cloud Storage (gs://my-bucket/my-model) o un percorso locale. Se questo argomento non è impostato, il valore predefinito è quello di --model.
  • --tokenizer_mode: la modalità del tokenizer. Le scelte possibili sono ["auto", "slow"]. Il valore predefinito è "auto". Se questo valore è impostato su "auto", viene utilizzato il tokenizer rapido, se disponibile. I tokenizer lenti sono scritti in Python e forniti nella libreria Transformers, mentre i tokenizer veloci, che offrono un miglioramento delle prestazioni, sono scritti in Rust e forniti nella libreria Tokenizers. Per saperne di più, consulta la documentazione di Hugging Face.
  • --trust_remote_code: indica se consentire i file di codice remoto definiti nei repository dei modelli Hugging Face. Il valore predefinito è False.
  • --load_format: il formato dei checkpoint del modello da caricare. Le scelte possibili sono ["auto", "dummy"]. Il valore predefinito è "auto". Se questo valore è impostato su "auto", i pesi del modello vengono caricati nel formato safetensors. Se questo valore è impostato su "dummy", i pesi del modello vengono inizializzati in modo casuale. L'impostazione di questo valore su "dummy" è utile per gli esperimenti.
  • --max_model_len: la lunghezza massima della finestra contestuale (lunghezza dell'input più lunghezza dell'output) da utilizzare per il modello. Il valore predefinito viene letto dal file di configurazione del modello in formato Hugging Face: config.json. Una lunghezza massima della finestra contestuale maggiore richiede più memoria TPU.
  • --sliding_window: se impostato, questo argomento sostituisce la dimensione della finestra del modello per l'attenzione della finestra scorrevole. Se imposti questo argomento su un valore più alto, il meccanismo di attenzione include più token e si avvicina all'effetto dell'auto-attenzione standard. Questo argomento è destinato esclusivamente all'utilizzo sperimentale. In casi d'uso generali, consigliamo di utilizzare le dimensioni della finestra originali del modello.
  • --seed: il seed per l'inizializzazione di tutti i generatori di numeri casuali. La modifica di questo argomento potrebbe influire sull'output generato per lo stesso prompt modificando i token campionati come token successivi. Il valore predefinito è 0.

Motore di inferenza

  • --num_hosts: il numero di host da eseguire. Il valore predefinito è 1. Per maggiori dettagli, consulta la documentazione sulla configurazione di TPU v5e.
  • --disagg_topo: definisce il numero di worker di precompilazione e di decodifica con la funzionalità sperimentale di distribuzione disaggregata. Il valore predefinito è None. L'argomento segue il formato: "number_of_prefill_workers,number_of_decode_workers".
  • --data_parallel_size: il numero di repliche parallele dei dati. Il valore predefinito è 1. Se imposti questo valore su N da 1, il throughput migliora di circa N, mantenendo la stessa latenza.
  • --tensor_parallel_size: il numero di repliche parallele del tensore. Il valore predefinito è 1. L'aumento del numero di repliche parallele del tensore generalmente migliora la latenza, perché accelera la moltiplicazione delle matrici riducendo le dimensioni della matrice.
  • --worker_distributed_method: Il metodo distribuito per avviare il worker. Utilizza mp per il modulo multiprocessing o ray per la libreria Ray. Il valore predefinito è mp.
  • --enable_jit: indica se attivare la modalità JIT (Just-In-Time Compilation). Il valore predefinito è True. Se imposti --no-enable_jit, la funzionalità viene disattivata. L'attivazione della modalità JIT migliora le prestazioni di inferenza, ma richiede più tempo per la compilazione iniziale. In generale, i vantaggi in termini di prestazioni dell'inferenza superano l'overhead.
  • --warmup: indica se eseguire il preriscaldamento del server con richieste di esempio durante l'inizializzazione. Il valore predefinito è True. L'impostazione --no-warmup lo disattiva. Il preriscaldamento è consigliato perché le richieste iniziali attivano una compilazione più pesante e, di conseguenza, saranno più lente.
  • --max_prefill_seqs: il numero massimo di sequenze che possono essere pianificate per il pre-riempimento per iterazione. Il valore predefinito è 1. Più alto è questo valore, maggiore è il throughput che il server può raggiungere, ma con potenziali effetti negativi sulla latenza.
  • --prefill_seqs_padding: il server aggiunge un padding alla dimensione del batch di precompilazione per renderla un multiplo di questo valore. Il valore predefinito è 8. L'aumento di questo valore riduce i tempi di ricompilazione del modello, ma aumenta il calcolo sprecato e l'overhead di inferenza. L'impostazione ottimale dipende dal traffico delle richieste.
  • --prefill_len_padding: il server aggiunge spaziatura interna alla lunghezza della sequenza fino a un multiplo di questo valore. Il valore predefinito è 512. L'aumento di questo valore riduce i tempi di ricompilazione del modello, ma aumenta lo spreco di calcolo e l'overhead di inferenza. L'impostazione ottimale dipende dalla distribuzione dei dati delle richieste.
  • --max_decode_seqs/--max_running_seqs: il numero massimo di sequenze che possono essere pianificate per la decodifica per iterazione. Il valore predefinito è 256. Più alto è questo valore, maggiore è il throughput che il server può raggiungere, ma con potenziali effetti negativi sulla latenza.
  • --decode_seqs_padding: il server aggiunge un padding alla dimensione del batch di decodifica per renderla un multiplo di questo valore. Il valore predefinito è 8. L'aumento di questo valore riduce i tempi di ricompilazione del modello, ma aumenta il calcolo sprecato e l'overhead di inferenza. L'impostazione ottimale dipende dal traffico delle richieste.
  • --decode_blocks_padding: il server aggiunge un padding al numero di blocchi di memoria utilizzati per la cache di coppie chiave-valore (cache KV) di una sequenza fino a un multiplo di questo valore durante la decodifica. Il valore predefinito è 128. L'aumento di questo valore riduce i tempi di ricompilazione del modello, ma aumenta lo spreco di calcolo e l'overhead di inferenza. L'impostazione ottimale dipende dalla distribuzione dei dati delle richieste.
  • --enable_prefix_cache_hbm: indica se attivare la memorizzazione nella cache dei prefissi in HBM. Il valore predefinito è False. L'impostazione di questo argomento può migliorare le prestazioni riutilizzando i calcoli dei prefissi condivisi delle richieste precedenti.
  • --enable_chunked_prefill: indica se attivare il pre-riempimento in blocchi. Il valore predefinito è False. L'impostazione di questo argomento può supportare una finestra contestuale maggiore e migliorare le prestazioni.

Gestione della memoria

  • --hbm_utilization_factor: la percentuale di memoria ad alta larghezza di banda (HBM) di Cloud TPU libera che può essere allocata per la cache KV dopo il caricamento dei pesi del modello. Il valore predefinito è 0.9. Se imposti questo argomento su un valore più alto, aumenti la dimensione della cache KV e puoi migliorare la velocità effettiva, ma aumenta il rischio di esaurire la memoria HBM di Cloud TPU durante l'inizializzazione e in fase di runtime.
  • --num_blocks: il numero di blocchi del dispositivo da allocare per la cache KV. Se questo argomento è impostato, il server ignora --hbm_utilization_factor. Se questo argomento non viene impostato, il server profila l'utilizzo dell'HBM e calcola il numero di blocchi di dispositivi da allocare in base a --hbm_utilization_factor. Se imposti questo argomento su un valore più alto, le dimensioni della cache KV aumentano e il throughput può migliorare, ma aumenta il rischio di esaurire la memoria HBM di Cloud TPU durante l'inizializzazione e in fase di runtime.
  • --block_size: il numero di token archiviati in un blocco. Le scelte possibili sono [8, 16, 32, 2048, 8192]. Il valore predefinito è 32. L'impostazione di questo argomento su un valore più alto riduce l'overhead nella gestione dei blocchi, a scapito di un maggiore spreco di memoria. L'impatto esatto sulle prestazioni deve essere determinato empiricamente.

LoRA dinamico

  • --enable_lora: indica se abilitare il caricamento dinamico degli adattatori LoRA da Cloud Storage. Il valore predefinito è False. Questa opzione è supportata per la famiglia di modelli Llama.
  • --max_lora_rank: il ranking LoRA massimo supportato per gli adattatori LoRA definiti nelle richieste. Il valore predefinito è 16. Se imposti questo argomento su un valore più alto, gli adattatori LoRA che possono essere utilizzati con il server sono più flessibili, ma aumenta la quantità di HBM Cloud TPU allocata per i pesi LoRA e diminuisce il throughput.
  • --enable_lora_cache: indica se abilitare la memorizzazione nella cache degli adattatori LoRA dinamici. Il valore predefinito è True. Se imposti --no-enable_lora_cache, la funzionalità viene disattivata. La memorizzazione nella cache migliora le prestazioni perché elimina la necessità di scaricare nuovamente i file dell'adattatore LoRA utilizzati in precedenza.
  • --max_num_mem_cached_lora: il numero massimo di adattatori LoRA archiviati nella memoria cache TPU TPU.Il valore predefinito è 16. Se imposti questo argomento su un valore più alto, aumenti la probabilità di un successo della cache, ma anche l'utilizzo della memoria HBM di Cloud TPU.

Puoi anche configurare il server utilizzando le seguenti variabili di ambiente:

  • HEX_LLM_LOG_LEVEL: controlla la quantità di informazioni di logging generate. Il valore predefinito è INFO. Imposta questo valore su uno dei livelli di logging Python standard definiti nel modulo di logging.
  • HEX_LLM_VERBOSE_LOG: indica se abilitare l'output di logging dettagliato. I valori consentiti sono true o false. Il valore predefinito è false.

Ottimizza gli argomenti del server

Gli argomenti del server sono correlati e hanno un effetto collettivo sulle prestazioni di gestione. Ad esempio, un'impostazione più elevata di --max_model_len=4096 comporta un maggiore utilizzo della memoria della TPU e, di conseguenza, richiede un'allocazione della memoria maggiore e un batching inferiore. Inoltre, alcuni argomenti sono determinati dal caso d'uso, mentre altri possono essere ottimizzati. Ecco un flusso di lavoro per la configurazione del server Hex-LLM.

  1. Determina la famiglia di modelli e la variante di modello di interesse. Ad esempio, Llama 3.1 8B Instruct.
  2. Stima il limite inferiore della memoria TPU necessaria in base alle dimensioni e alla precisione del modello: model_size * (num_bits / 8). Per un modello 8B e una precisione bfloat16, il limite inferiore della memoria TPU necessaria sarebbe 8 * (16 / 8) = 16 GB.
  3. Stima il numero di chip TPU v5e necessari, dove ogni chip v5e offre 16 GB: tpu_memory / 16. Per un modello 8B e una precisione bfloat16, è necessario più di un chip. Tra le configurazioni a 1, 4 e 8 chip, la configurazione più piccola che offre più di 1 chip è quella a 4 chip: ct5lp-hightpu-4t. In seguito puoi impostare --tensor_parallel_size=4.
  4. Determina la lunghezza massima della finestra contestuale (lunghezza dell'input + lunghezza dell'output) per il caso d'uso previsto. Ad esempio, 4096. In seguito puoi impostare --max_model_len=4096.
  5. Regola la quantità di memoria TPU libera allocata per la cache KV al valore massimo raggiungibile in base alle configurazioni del modello, dell'hardware e del server (--hbm_utilization_factor). Inizia con 0.95. Esegui il deployment del server Hex-LLM e testalo con prompt lunghi e un'elevata contemporaneità. Se il server esaurisce la memoria, riduci il fattore di utilizzo di conseguenza.

Un insieme di argomenti di esempio per il deployment di 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

Un insieme di argomenti di esempio per il deployment di Llama 3.1 70B Instruct AWQ su 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

Richiedi la quota di Cloud TPU

In Model Garden, la quota predefinita è di 32 chip Cloud TPU v5e nella regione us-west1. Queste quote si applicano ai deployment con un solo clic e ai deployment dei blocchi note Colab Enterprise. Per richiedere un valore di quota più elevato, consulta Richiedi un aggiustamento della quota.