Pubblica modelli aperti utilizzando il contenitore 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 Vertex AI, 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 Vertex AI 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 la pubblicazione 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:

  • 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 dei 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 uno slice 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 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 Vertex AI Model Garden - Llama 3.1 (Deployment).

In generale, le uniche modifiche necessarie per attivare la pubblicazione multihost 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.

Il servizio disaggregato è un metodo efficace per bilanciare il tempo al primo token (Time to First Token, TTFT) e il tempo per token di output (Time Per Output Token, TPOT) per ogni richiesta, nonché la velocità effettiva complessiva del servizio. 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 risultati 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 trovati nella cache.

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

Precompilazione a blocchi

Il riempimento automatico suddiviso in blocchi divide un riempimento automatico di una richiesta in blocchi più piccoli e combina il riempimento automatico e la decodifica in un unico passaggio batch. Hex-LLM implementa il prefill in blocchi per bilanciare il tempo al primo token (Time to First Token, TTFT) e il tempo per token di output (Time per Output Token, TPOT) e migliora il throughput.

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 erogazione Cloud TPU Hex-LLM è integrato in Model Garden. Puoi accedere a questa tecnologia di pubblicazione tramite il playground, il deployment con un solo clic e gli esempi di notebook Colab Enterprise per una varietà di modelli.

Utilizzare il playground

Il playground di Model Garden è un endpoint Vertex AI 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.

Prova Gemma.

Utilizzare il deployment con un solo clic

Puoi eseguire il deployment di un endpoint Vertex AI 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 iniziare la procedura di deployment. Ricevi 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 Vertex AI con Hex-LLM utilizzando l'SDK Vertex AI 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ò 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, viene utilizzato il valore predefinito 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 di 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. Impostare questo valore su "dummy" è utile per gli esperimenti.
  • --max_model_len: la lunghezza massima del contesto (lunghezza dell'input più la lunghezza dell'output) da utilizzare per il modello. Il valore predefinito viene letto dal file di configurazione del modello nel formato Hugging Face: config.json. Una lunghezza del contesto massima 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ù 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 le dimensioni della finestra originali 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 del tensore. Il valore predefinito è 1. Aumentare il numero di repliche parallele dei tensori in genere 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 a scapito di un tempo aggiuntivo dedicato alla compilazione iniziale. In generale, i vantaggi in termini di prestazioni di 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 precompilamento in blocchi. Il valore predefinito è False. L'impostazione di questo argomento può supportare una lunghezza del contesto maggiore e migliorare le prestazioni.

Gestione della memoria

  • --hbm_utilization_factor: la percentuale di memoria ad alta larghezza di banda (HBM) 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 il throughput, ma aumenti il rischio di esaurire la HBM della 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. Se imposti questo argomento su un valore più alto, aumenti le dimensioni della cache KV e puoi migliorare il throughput, ma aumenti il rischio di esaurire la HBM di Cloud TPU durante l'inizializzazione e l'esecuzione.
  • --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 rango 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 cache di memoria TPU TPU.Il valore predefinito è 16. Se imposti questo argomento su un valore maggiore, aumentano le 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 abilitare 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 del contesto (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 libera 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 della quota.