Pubblica modelli aperti utilizzando il contenitore premium Hex-LLM su Cloud TPU

Hex-LLM, un modello linguistico di grandi dimensioni (LLM) ad alta efficienza di cui viene eseguito il servizio con XLA, è il framework di servizio LLM di Vertex AI progettato e ottimizzato per l'hardware Cloud TPU. Hex-LLM combina tecnologie di pubblicazione di modelli LLM come il batching continuo e PagedAttention con ottimizzazioni di Vertex AI personalizzate per XLA e Cloud TPU. Si tratta di un modello LLM ad alta efficienza e a basso costo che viene eseguito su Cloud TPU per i modelli open source.

Hex-LLM è disponibile in Model Garden tramite il laboratorio di modelli, il deployment con un solo clic e il notebook.

Funzionalità

Hex-LLM si basa su progetti open source con le ottimizzazioni di Google per XLA e Cloud TPU. Hex-LLM raggiunge una velocità effettiva elevata e una bassa latenza quando pubblica LLM utilizzati di frequente.

Hex-LLM include le seguenti ottimizzazioni:

  • Algoritmo di raggruppamento continuo basato su token per contribuire a 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 parallelismo dei tensori flessibili e componibili con metodi di suddivisione dei pesi altamente ottimizzati per eseguire in modo efficiente i modelli linguistici di grandi dimensioni su più chip Cloud TPU.

Hex-LLM supporta una vasta 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 Guard 3 1B e 8B
  • Mistral 7B
  • Mixtral 8x7B e 8x22B
  • Phi-3 mini e medium
  • Qwen2 0,5 B, 1,5 B e 7 B
  • Qwen2.5 0,5 miliardi, 1,5 miliardi, 7 miliardi, 14 miliardi e 32 miliardi di AWQ

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

  • Hex-LLM è incluso in un singolo contenitore. Hex-LLM pacchettizza il server API, il motore di inferenza e i modelli supportati in un'unica immagine Docker da eseguire.
  • 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 su più host
  • Pubblicazione disaggregata [sperimentale]
  • Memorizzazione nella cache dei prefissi
  • Supporto della quantizzazione a 4 bit

Pubblicazione su più host

Hex-LLM ora supporta la pubblicazione di modelli con uno slice TPU multi-host. Questa funzionalità ti consente di pubblicare 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 --tpu_topology nella richiesta di caricamento del modello dell'SDK Vertex AI. L'esempio seguente mostra come eseguire il deployment del contenitore Hex-LLM con una topologia TPU 4x4 v5e che serve 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 contenitore 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 su più 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 la pubblicazione disaggregata come funzionalità sperimentale. Può essere attivata solo nella configurazione di un singolo host e il rendimento è in fase di ottimizzazione.

La pubblicazione disaggregata è un metodo efficace per bilanciare il tempo di visualizzazione del primo token (TTFT) e il tempo per token di output (TPOT) per ogni richiesta, nonché il throughput complessivo della pubblicazione. Separa la fase di precompilazione e la fase di decodifica in diversi carichi di lavoro in modo che non interferiscano tra loro. Questo metodo è particolarmente utile per gli scenari che impongono 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 contenitore Hex-LLM su una TPU v5e-8 che serve 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 di risposta al primo token (TTFT) per i prompt che hanno contenuti identici all'inizio, ad esempio preamboli aziendali, istruzioni di sistema comuni e cronologia delle conversazioni con più turni. 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 contenitore Hex-LLM su TPU v5e-8 che serve 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",
    "--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 hit della cache si verificano in incrementi di questo valore, garantendo che il numero di 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 tra i token prompt hanno generato una corrispondenza nella cache.

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

Supporto della quantizzazione a 4 bit

La quantizzazione è una tecnica per ridurre i costi computazionali 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 i modelli con pesi INT4 quantizzati utilizzando la quantizzazione a punto zero AWQ. Hex-LLM supporta le varianti INT4 delle famiglie di modelli Mistral, Mixtral e Llama.

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

Iniziare a utilizzare Model Garden

Il contenitore di servizio Hex-LLM Cloud TPU è 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 serie di modelli.

Utilizzare il parco giochi

Il playground di Model Garden è un endpoint Vertex AI pre-implementato che è raggiungibile inviando richieste nella scheda del modello.

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

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

Prova con 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 avviare 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 di 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 notebook.

  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 è come segue:

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 alle esigenze e al caso d'uso previsti. Tieni presente che gli argomenti sono predefiniti per il deployment con un solo clic per abilitare l'esperienza di deployment più semplice. Per personalizzare gli argomenti, puoi utilizzare gli esempi del notebook come riferimento e impostarli 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 elementi del modello devono seguire il formato Hugging Face e utilizzare i file safetensors per i pesi del modello. Gli elementi del modello quantizzato BitsAndBytes int8 e AWQ sono supportati per Llama, Gemma 2 e Mistral/Mixtral.
  • --tokenizer: il tokenizzatore 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 per impostazione predefinita il valore di --model.
  • --tokenizer_mode: la modalità del tokenizzatore. Le possibili scelte sono ["auto", "slow"]. Il valore predefinito è "auto". Se questo valore è impostato su "auto", viene utilizzato lo tokenizer rapido, se disponibile. I tokenizzatori lenti sono scritti in Python e forniti nella libreria Transformers, mentre i tokenizzatori rapidi 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 di Hugging Face. Il valore predefinito è False.
  • --load_format: formato dei checkpoint del modello da caricare. Le possibili scelte 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 la sperimentazione.
  • --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 in formato Hugging Face: config.json. Una lunghezza massima del contesto maggiore richiede più memoria TPU.
  • --sliding_window: se impostato, questo argomento sostituisce le dimensioni della finestra del modello per l'attenzione con finestra scorrevole. Impostando questo argomento su un valore maggiore, il meccanismo di attenzione include più token e si avvicina all'effetto dell'attenzione autoreferenziale standard. Questo argomento è pensato solo per l'utilizzo sperimentale. Per i casi d'uso generali, consigliamo di utilizzare le dimensioni della finestra originale 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 pubblicazione disaggregata della funzionalità sperimentale. 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. Impostando questo valore su N da 1, la velocità effettiva migliorerà di circa N, mantenendo la stessa latenza.
  • --tensor_parallel_size: il numero di repliche parallele di tensori. Il valore predefinito è 1. L'aumento del numero di repliche parallele del tensore in genere migliora la latenza, perché accelera la moltiplicazione della matrice 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 (compilazione Just-in-Time). Il valore predefinito è True. L'impostazione --no-enable_jit la disattiva. L'attivazione della modalità JIT migliora le prestazioni di inferenza, ma richiede un tempo aggiuntivo per la compilazione iniziale. In generale, i vantaggi in termini di prestazioni dell'inferenza superano l'overhead.
  • --warmup: indica se eseguire l'inizializzazione del server con richieste di esempio. Il valore predefinito è True. L'impostazione --no-warmup la disattiva. Il riscaldamento è 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 è il throughput che il server può raggiungere, ma con potenziali effetti negativi sulla latenza.
  • --prefill_seqs_padding: il server aumenta la dimensione del batch di precompilazione a un moltiplo di questo valore. Il valore predefinito è 8. L'aumento di questo valore riduce i tempi di ricompilamento del modello, ma aumenta il calcolo sprecato e il carico aggiuntivo dell'inferenza. L'impostazione ottimale dipende dal traffico delle richieste.
  • --prefill_len_padding: il server aggiunge spazi alla lunghezza della sequenza fino a un multiplo di questo valore. Il valore predefinito è 512. L'aumento di questo valore riduce i tempi di ricompilamento del modello, ma aumenta il calcolo sprecato e il carico di lavoro 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 ogni 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 aggiunge spazi alla dimensione del batch di decodifica per ottenere un multiplo di questo valore. Il valore predefinito è 8. L'aumento di questo valore riduce i tempi di ricompilarazione del modello, ma aumenta il sovraccarico di calcolo e di inferenza sprecato. L'impostazione ottimale dipende dal traffico delle richieste.
  • --decode_blocks_padding: il server aggiunge il numero di blocchi di memoria utilizzati per la cache chiave-valore (cache KV) di una sequenza a un multiplo di questo valore durante la decodifica. Il valore predefinito è 128. L'aumento di questo valore riduce i tempi di ricompilamento del modello, ma aumenta il calcolo sprecato e il carico di lavoro di inferenza. L'impostazione ottimale dipende dalla distribuzione dei dati delle richieste.
  • --enable_prefix_cache_hbm: indica se abilitare la cache dei prefissi in HBM. Il valore predefinito è False. L'impostazione di questo argomento può migliorare il rendimento riutilizzando i calcoli dei prefissi condivisi delle richieste precedenti.

Gestione della memoria

  • --hbm_utilization_factor: la percentuale di memoria HBM (High Bandwidth Memory) Cloud TPU libera che può essere allocata per la cache KV dopo il caricamento dei pesi del modello. Il valore predefinito è 0.9. L'impostazione di questo argomento su un valore più elevato aumenta la dimensione della cache KV e può migliorare il throughput, ma aumenta il rischio di esaurire la memoria HBM di Cloud TPU durante l'inizializzazione e in fase di esecuzione.
  • --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 dell'HBM e calcola il numero di blocchi del dispositivo da allocare in base a --hbm_utilization_factor. Impostare questo argomento su un valore più alto aumenta le dimensioni della cache KV e può migliorare il throughput, ma aumenta il rischio di esaurire la memoria HBM di Cloud TPU durante l'inizializzazione e in fase di esecuzione.
  • --block_size: numero di token memorizzati in un blocco. Le possibili scelte 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 costo di un maggiore spreco di memoria. L'impatto esatto sul rendimento deve essere determinato empiricamente.

LoRa dinamico

  • --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. L'impostazione di un valore più elevato per questo argomento consente una maggiore flessibilità negli adattatori LoRa che possono essere utilizzati con il server, 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. L'impostazione --no-enable_lora_cache la disattiva. 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 memorizzati nella cache della memoria TPU.Il valore predefinito è 16. L'impostazione di questo argomento su un valore più elevato migliora le probabilità di una successo della cache, ma aumenta 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 generate durante la registrazione. Il valore predefinito è INFO. Imposta questo valore su uno dei livelli di logging Python standard definiti nel modulo logging.
  • HEX_LLM_VERBOSE_LOG: indica se attivare l'output del 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 sul funzionamento del servizio. Ad esempio, un'impostazione più elevata di --max_model_len=4096 porta a un maggiore utilizzo della memoria TPU e, di conseguenza, richiede un'allocazione della memoria più elevata e meno batching. Inoltre, alcuni argomenti sono determinati dal caso d'uso, mentre altri possono essere ottimizzati. Di seguito è riportato un flusso di lavoro per la configurazione del server Hex-LLM.

  1. Determina la famiglia e la variante di modello di tuo interesse. Ad esempio, Llama 3.1 8B Instruct.
  2. Stimare 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 di byte e una precisione bfloat16, il limite inferiore della memoria TPU necessaria è 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 di bit e una precisione bfloat16, sono necessari più di un chip. Tra le configurazioni a 1 chip, 4 chip e 8 chip, la configurazione più piccola che offre più di 1 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 possibile in base al modello, all'hardware e alle configurazioni del server (--hbm_utilization_factor). Inizia con 0.95. Esegui il deployment del server Hex-LLM e testa il server con prompt lunghi e alta concorrenza. Se il server non dispone di memoria sufficiente, 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

Richiedi la quota Cloud TPU

In Model Garden, la quota predefinita è di 4 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 quote aggiuntive, consulta Richiedere una quota più alta.