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 il playground del modello, 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 batching continuo basato su token per garantire che i modelli utilizzino completamente l'hardware con un numero elevato di richieste simultanee.
  • 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 le seguenti:

  • Hex-LLM è incluso in un unico container. Hex-LLM raggruppa il server API, il motore di inferenza e i modelli supportati in un'unica immagine Docker da implementare.
  • Compatibile con il formato modelli Hugging Face. Hex-LLM può caricare un modello Hugging Face da un disco locale, da Hugging Face Hub 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 slice TPU multi-host. Questa funzionalità ti consente di erogare 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 contenitore 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 per il deployment del container Hex-LLM con una topologia TPU multi-host, consulta il notebook Agent Platform Model Garden - Llama 3.1 (Deployment).

In generale, le uniche modifiche necessarie per attivare la pubblicazione 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 dei modelli dell'SDK Vertex AI.

Pubblicazione disaggregata [sperimentale]

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

La pubblicazione disaggregata è un metodo efficace per bilanciare il tempo al primo token (TTFT) e il tempo per token di output (TPOT) per ogni richiesta, nonché la velocità effettiva di pubblicazione complessiva. Separa la fase di precompilazione e la fase di decodifica in carichi di lavoro diversi in modo che non interferiscano tra loro. Questo metodo è particolarmente utile per gli scenari che impostano requisiti di latenza rigorosi.

Per attivare questa funzionalità, imposta --disagg_topo negli argomenti del contenitore 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 dei prefissi riduce il tempo al primo token (TTFT) per i prompt che hanno contenuti identici all'inizio del prompt, come i preamboli a livello aziendale, le istruzioni di sistema comuni e la cronologia delle conversazioni multi-turno. Invece di 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 contenitore 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 dei prompt che superano una determinata lunghezza (512 token per impostazione predefinita, configurabile utilizzando prefill_len_padding). I successi della cache si verificano con 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 Chat Completions indica quanti token del prompt sono stati un successo della cache.

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

Precompilazione in blocchi

Il riempimento preliminare in blocchi suddivide un riempimento preliminare della richiesta in blocchi più piccoli e combina il riempimento preliminare e la decodifica in un unico passaggio batch. Hex-LLM implementa il riempimento preliminare in blocchi per bilanciare il Time to First Token (TTFT) e il Time per Output Token (TPOT) e migliora la velocità effettiva.

Per attivare questa funzionalità, imposta --enable_chunked_prefill negli argomenti del contenitore 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 soliti BF16 o FP32.

Hex-LLM supporta la quantizzazione solo dei 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 è necessario alcun flag aggiuntivo per la pubblicazione di modelli quantizzati.

Inizia a utilizzare Model Garden

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

Utilizzare il playground

Il playground di Model Garden è un endpoint di Agent Platform pre-deployment raggiungibile inviando richieste nella scheda del modello.

  1. Inserisci un prompt e, se vuoi, includi argomenti per la tua richiesta.

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

Provalo con Gemma.

Utilizzare il deployment con un solo clic

Puoi eseguire il deployment di un endpoint personalizzato di Agent Platform 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 iniziare la procedura di deployment. Riceverai due notifiche via email: una quando il modello viene caricato e un'altra quando l'endpoint è pronto.

Utilizzare il notebook Colab Enterprise

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

  1. Vai alla pagina della scheda del modello 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 Colab Enterprise:

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 base al caso d'uso e ai requisiti previsti. Tieni presente che gli argomenti sono predefiniti per il deployment con un solo clic per consentire 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 file safetensors per i pesi del modello. Gli artefatti del modello quantizzato BitsAndBytes int8 e AWQ sono supportati per Llama, Gemma 2 e Mistral/Mixtral.
  • --tokenizer: il tokenizer da caricare. Può trattarsi di un ID modello Hugging Face, di un percorso del bucket Cloud Storage (gs://my-bucket/my-model) o di 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 questa opzione è impostata 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 di modelli Hugging Face. Il valore predefinito è False.
  • --load_format: 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 su "dummy" è utile per gli esperimenti.
  • --max_model_len: la massima finestra contestuale (lunghezze di input e output) da erogare per il modello. Il valore predefinito viene letto dal file di configurazione del modello nel formato Hugging Face: config.json. Una finestra contestuale maggiore richiede più memoria TPU.
  • --sliding_window: se impostato, questo argomento sostituisce la dimensione della finestra del modello per l'attenzione a finestra scorrevole. Se imposti questo argomento su un valore più grande, il meccanismo di attenzione include più token e si avvicina all'effetto dell'auto-attenzione standard. Questo argomento è destinato solo all'utilizzo sperimentale. Nei casi d'uso generali, consigliamo di utilizzare la dimensione della finestra originale del modello.
  • --seed: il seme per inizializzare 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 della TPU v5e.
  • --disagg_topo: definisce il numero di worker di precompilazione e di decodifica con la funzionalità sperimentale di pubblicazione 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 anziché 1, la velocità effettiva migliora di circa N, mentre la latenza rimane invariata.
  • --tensor_parallel_size: il numero di repliche parallele dei tensori. Il valore predefinito è 1. Aumentare il numero di repliche parallele di tensori in genere migliora la latenza, perché accelera la moltiplicazione delle matrici riducendone le dimensioni.
  • --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 a scapito di un tempo aggiuntivo dedicato alla compilazione iniziale. In generale, i vantaggi in termini di prestazioni dell'inferenza superano il sovraccarico.
  • --warmup: Indica se riscaldare il 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 quindi saranno più lente.
  • --max_prefill_seqs: Il numero massimo di sequenze che possono essere pianificate per il precompilamento per iterazione. Il valore predefinito è 1. Maggiore è questo valore, maggiore è la velocità effettiva che il server può raggiungere, ma con potenziali effetti negativi sulla latenza.
  • --prefill_seqs_padding: Il server riempie le dimensioni del batch di precompilazione fino a 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 esegue il padding della 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 il sovraccarico di calcolo e inferenza sprecati. 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. Maggiore è questo valore, maggiore è la velocità effettiva che il server può raggiungere, ma con potenziali effetti negativi sulla latenza.
  • --decode_seqs_padding: Il server riempie la dimensione del batch di decodifica fino a un multiplo di questo valore. Il valore predefinito è 8. L'aumento di questo valore riduce i tempi di ricompilazione del modello, ma aumenta il sovraccarico di calcolo e inferenza sprecato. L'impostazione ottimale dipende dal traffico delle richieste.
  • --decode_blocks_padding: il server riempie il numero di blocchi di memoria utilizzati per la cache 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 il sovraccarico di calcolo e inferenza sprecati. L'impostazione ottimale dipende dalla distribuzione dei dati delle richieste.
  • --enable_prefix_cache_hbm: indica se abilitare 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 abilitare il riempimento automatico 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 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 le dimensioni della cache KV e puoi migliorare la velocità effettiva, ma aumenta il rischio di esaurire la HBM di Cloud TPU durante l'inizializzazione e in fase di runtime.
  • --num_blocks: numero di blocchi di dispositivi da allocare per la cache KV. Se questo argomento è impostato, il server ignora --hbm_utilization_factor. Se questo argomento non è impostato, il server profila l'utilizzo di HBM e calcola il numero di blocchi di dispositivi da allocare in base a --hbm_utilization_factor. L'impostazione di questo argomento su un valore più alto aumenta le dimensioni della cache KV e può migliorare il throughput, ma aumenta il rischio di esaurire l'HBM della TPU Cloud durante l'inizializzazione e in fase di runtime.
  • --block_size: Numero di token memorizzati in un blocco. Le scelte possibili sono [8, 16, 32, 2048, 8192]. Il valore predefinito è 32. L'impostazione di questo argomento su un valore maggiore riduce il sovraccarico nella gestione dei blocchi, a scapito di un maggiore spreco di memoria. L'impatto esatto sulle prestazioni deve essere determinato empiricamente.

Dynamic LoRA

  • --enable_lora: indica se attivare il caricamento dinamico degli adattatori LoRA da Cloud Storage. Il valore predefinito è False. Questa funzionalità è 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, puoi utilizzare un numero maggiore di adattatori LoRA con il server, ma aumenta la quantità di HBM Cloud TPU allocata per i pesi LoRA e diminuisce la velocità effettiva.
  • --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.Il valore predefinito è 16. L'impostazione di questo argomento su un valore più grande aumenta la probabilità di un successo della cache, ma aumenta la quantità di utilizzo di 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 attivare l'output di logging dettagliato. I valori consentiti sono true o false. Il valore predefinito è false.

Ottimizzare gli argomenti del server

Gli argomenti del server sono intercorrelati e hanno un effetto collettivo sulle prestazioni di pubblicazione. Ad esempio, un'impostazione più grande di --max_model_len=4096 comporta un maggiore utilizzo della memoria TPU e pertanto richiede un'allocazione di memoria più grande e meno batch. Inoltre, alcuni argomenti sono determinati dal caso d'uso, mentre altri possono essere ottimizzati. Ecco un workflow per configurare il server Hex-LLM.

  1. Determina la famiglia di modelli e la variante del 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 da 8 miliardi 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 da 8 miliardi e una precisione bfloat16, hai bisogno di più di un chip. Tra le configurazioni a 1, 4 e 8 chip, la più piccola che offre più di un chip è quella a 4 chip: ct5lp-hightpu-4t. Successivamente 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. Successivamente puoi impostare --max_model_len=4096.
  5. Ottimizza la quantità di memoria TPU senza costi allocata per la cache KV al valore massimo raggiungibile in base alle configurazioni di modello, hardware e server (--hbm_utilization_factor). Inizia con 0.95. Esegui il deployment del server Hex-LLM e testalo con prompt lunghi e concorrenza elevata. Se il server esaurisce la memoria, riduci di conseguenza il fattore di utilizzo.

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

Richiedere la quota 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, vedi Richiedi un aggiustamento delle quote.