Disponibilize modelos abertos usando o contêiner premium Hex-LLM no Cloud TPU

O Hex-LLM, um modelo de linguagem grande (LLM) de alta eficiência que veicula com XLA, é o framework de disponibilização de LLMs da Gemini Enterprise Agent Platform projetado e otimizado para hardware do Cloud TPU. O Hex-LLM combina tecnologias de veiculação de LLM, como lotes contínuos e PagedAttention, com as otimizações da Gemini Enterprise Agent Platform personalizadas para XLA e Cloud TPU. É um LLM de alta eficiência e baixo custo que é veiculado no Cloud TPU para os modelos de origem.

O Hex-LLM está disponível no Model Garden por meio de modelo Playground, implantação com um clique e notebook.

Recursos

O Hex-LLM é baseado em projetos de código aberto com as otimizações do próprio Google para XLA e o Cloud TPU. O Hex-LLM atinge alta capacidade de processamento e baixa latência na disponibilização os LLMs usados com frequência.

O Hex-LLM inclui as seguintes otimizações:

  • Algoritmo de lote contínuo baseado em token para ajudar a garantir que os modelos sejam totalmente utilizando o hardware com um grande número de solicitações simultâneas.
  • Uma reescrita completa dos kernels de atenção otimizados para XLA.
  • Estratégias de paralelismo de dados flexíveis e combináveis e de paralelismo de tensores com métodos de fragmentação de peso altamente otimizados para executar LLMs com eficiência vários chips do Cloud TPU.

O Hex-LLM oferece suporte a uma ampla gama de LLMs densos e esparsos:

  • 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

O Hex-LLM também oferece vários recursos, entre eles:

  • Hex-LLM está incluído em um único contêiner. O Hex-LLM empacota o servidor de API, o mecanismo de inferência e os modelos compatíveis em uma única imagem Docker para ser implantada.
  • Compatível com os modelos Hugging Face . O Hex-LLM pode carregar um modelo Hugging Face do disco local, o Hugging o Face Hub e um bucket do Cloud Storage.
  • Quantização usando bitsandbytes e AWQ:
  • Carregamento dinâmico LoRA (em inglês). O Hex-LLM é capaz de carregar os pesos LoRA lendo o argumento da solicitação durante a disponibilização.

Recursos avançados

O Hex-LLM é compatível com os seguintes recursos avançados:

  • Disponibilidade de vários hosts
  • Exibição desagregada [experimental]
  • Armazenamento em cache de prefixo
  • Suporte à quantização de 4 bits

Disponibilidade de vários hosts

O Hex-LLM agora oferece suporte à disponibilização de modelos com uma fração de TPU de vários hosts. Esse recurso permite disponibilizar modelos grandes que não podem ser carregados em uma VM de TPU de host único, que contém no máximo oito núcleos v5e.

Para ativar esse recurso, defina --num_hosts nos argumentos do contêiner do Hex-LLM e defina --tpu_topology na solicitação de upload do modelo do SDK da Vertex AI. O exemplo a seguir mostra como implantar o contêiner Hex-LLM com uma topologia de TPU 4x4 v5e que disponibiliza o modelo 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,
)

Para um tutorial completo sobre como implantar o contêiner Hex-LLM com uma topologia de TPU de vários hosts, consulte o notebook Agent Platform Model Garden - Llama 3.1 (Deployment).

Em geral, as únicas mudanças necessárias para ativar a veiculação em vários hosts são:

  1. Defina o argumento --tensor_parallel_size como o número total de núcleos na topologia da TPU.
  2. Defina o argumento --num_hosts como o número de hosts na topologia da TPU.
  3. Defina --tpu_topology com a API de upload de modelo do SDK da Vertex AI.

Exibição desagregada [experimental]

O Hex-LLM agora oferece suporte à disponibilização desagregada como um recurso experimental. Ele só pode ser ativado na configuração de host único, e o desempenho está em ajuste.

A veiculação desagregada é um método eficaz para equilibrar o tempo até o primeiro token (TTFT, na sigla em inglês) e o tempo por token de saída (TPOT, na sigla em inglês) para cada solicitação, além da capacidade de veiculação geral. Ele separa as fases de preenchimento e decodificação em cargas de trabalho diferentes para que não interfiram uma na outra. Esse método é especialmente útil para cenários que definem requisitos de latência rigorosos.

Para ativar esse recurso, defina --disagg_topo nos argumentos do contêiner do Hex-LLM. Veja um exemplo de como implantar o contêiner Hex-LLM na TPU v5e-8 que disponibiliza o modelo 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,
)

O argumento --disagg_topo aceita uma string no formato "number_of_prefill_workers,number_of_decode_workers". No exemplo anterior, ele é definido como "3,1" para configurar três workers de preenchimento prévio e um worker de decodificação. Cada worker usa dois núcleos de TPU v5e.

Armazenamento em cache de prefixo

O armazenamento em cache de prefixos reduz o tempo para o primeiro token (TTFT, na sigla em inglês) em comandos que têm conteúdo idêntico no início, como preâmbulos de toda a empresa, instruções comuns do sistema e histórico de conversas de várias interações. Em vez de processar os mesmos tokens de entrada repetidamente, o Hex-LLM pode reter um cache temporário dos cálculos de tokens de entrada processados para melhorar o TTFT.

Para ativar esse recurso, defina --enable_prefix_cache_hbm nos argumentos do contêiner Hex-LLM. Veja a seguir um exemplo que mostra como implantar o contêiner Hex-LLM na TPU v5e-8 que disponibiliza o modelo 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,
)

O Hex-LLM usa o armazenamento em cache de prefixos para otimizar o desempenho de comandos que excedem um determinado tamanho (512 tokens por padrão, configurável usando prefill_len_padding). As ocorrências em cache ocorrem em incrementos desse valor, garantindo que a contagem de tokens em cache seja sempre um múltiplo de prefill_len_padding. O campo cached_tokens de usage.prompt_tokens_details na resposta da API de conclusão de chat indica quantas ocorrência em cache os tokens de comando tiveram.

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

Preenchimento em blocos

O pré-preenchimento em blocos divide um pré-preenchimento de solicitação em blocos menores e mistura o pré-preenchimento e a decodificação em uma etapa de lote. O Hex-LLM implementa o preenchimento prévio em blocos para equilibrar o tempo até o primeiro token (TTFT, na sigla em inglês) e o tempo por token de saída (TPOT, na sigla em inglês) e melhorar a capacidade de processamento.

Para ativar esse recurso, defina --enable_chunked_prefill nos argumentos do contêiner Hex-LLM. Confira a seguir um exemplo que mostra como implantar o contêiner Hex-LLM na TPU v5e-8 que disponibiliza o modelo 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,
)

Suporte à quantização de 4 bits

A quantização é uma técnica para reduzir os custos computacionais e de memória da execução de inferências, representando os pesos ou as ativações com tipos de dados de baixa precisão, como INT8 ou INT4, em vez dos habituais BF16 ou FP32.

O Hex-LLM é compatível com a quantização somente de peso INT8. O suporte estendido inclui modelos com pesos INT4 quantizados usando a quantização de ponto zero AWQ. O Hex-LLM é compatível com as variantes INT4 das famílias de modelos Mistral, Mixtral e Llama.

Não é necessária nenhuma flag adicional para disponibilizar modelos quantizados.

Introdução ao Model Garden

O contêiner de veiculação Hex-LLM do Cloud TPU está integrado o Grupo de modelos. Você pode acessar essa tecnologia de veiculação por meio do Playground, implantação com um clique e bloco do Colab Enterprise exemplos de uma variedade de modelos.

Usar o playground

O playground do Model Garden é um endpoint pré-implantado da Agent Platform que pode ser acessado enviando solicitações no card de modelo.

  1. Insira um comando e, se quiser, inclua argumentos para sua solicitação.

  2. Clique em ENVIAR para receber a resposta do modelo rapidamente.

Faça um teste com o Gemma!

Usar a implantação com um clique

É possível implantar um endpoint personalizado da Agent Platform com o Hex-LLM usando um card de modelo.

  1. Navegue até o card de modelo. página e clique em Implantar.

  2. Para a variação de modelo que você quer usar, selecione a Cloud TPU Tipo de máquina v5e para implantação.

  3. Clique em Implantar na parte de baixo para iniciar o processo. Você recebe duas notificações por e-mail: um quando o modelo é carregado e outro quando endpoint esteja pronto.

Usar o bloco do Colab Enterprise

Para flexibilidade e personalização, use os exemplos de notebook do Colab Enterprise para implantar um endpoint da Agent Platform com o Hex-LLM usando o SDK da Agent Platform para Python.

  1. Acesse a página do card de modelo e clique em Abrir notebook.

  2. Selecione o notebook da Vertex Serving. O notebook está aberto no Colab Enterprise

  3. Executar pelo notebook para implantar um modelo usando o Hex-LLM e enviar solicitações de previsão ao endpoint. O snippet de código da implantação é da seguinte forma:

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

Exemplos de notebooks do Colab:

Configurar argumentos do servidor e variáveis de ambiente

É possível definir os seguintes argumentos para iniciar o servidor Hex-LLM. É possível adaptar os argumentos para que se ajustem melhor ao caso de uso e aos requisitos pretendidos. Os argumentos são predefinidos para implantação com um clique, o que facilita o processo. Para personalizar os argumentos, use os exemplos de notebook como referência e defina os argumentos de acordo.

Modelo

  • --model: o modelo a ser carregado. É possível especificar um ID de modelo do Hugging Face, um caminho de bucket do Cloud Storage (gs://my-bucket/my-model) ou um caminho local. Os artefatos do modelo precisam seguir o formato do Hugging Face e usar arquivos safetensors para os pesos do modelo. Os artefatos de modelo quantizados BitsAndBytes int8 e AWQ são compatíveis com Llama, Gemma 2 e Mistral/Mixtral.
  • --tokenizer: o tokenizador a ser carregado. Pode ser um ID de modelo do Hugging Face, um caminho de bucket do Cloud Storage (gs://my-bucket/my-model) ou um caminho local. Se esse argumento não for definido, o padrão será o valor de --model.
  • --tokenizer_mode: o modo do tokenizador. As opções possíveis são ["auto", "slow"]. O valor padrão é "auto". Se for definido como "auto", o tokenizador rápido será usado, se disponível. Os tokenizadores lentos são escritos em Python e fornecidos na biblioteca Transformers, enquanto os tokenizadores rápidos que oferecem melhoria de desempenho são escritos em Rust e fornecidos na biblioteca Tokenizers. Para mais informações, consulte a documentação do Hugging Face.
  • --trust_remote_code: se os arquivos de código remoto definidos nos repositórios de modelos do Hugging Face devem ser permitidos. O valor padrão é False.
  • --load_format: formato dos checkpoints do modelo a serem carregados. As opções possíveis são ["auto", "dummy"]. O valor padrão é "auto". Se definido como "auto", os pesos do modelo serão carregados no formato safetensors. Se for definido como "dummy", os pesos do modelo serão inicializados aleatoriamente. Definir como "dummy" é útil para experimentação.
  • --max_model_len: o tamanho máximo do contexto (tamanho da entrada mais o tamanho da saída) a ser disponibilizado para o modelo. O valor padrão é lido no arquivo de configuração do modelo no formato Hugging Face: config.json. Um tamanho do contexto máximo maior requer mais memória de TPU.
  • --sliding_window: se definido, esse argumento substitui o tamanho da janela do modelo para a atenção da janela deslizante. Definir esse argumento como um valor maior faz com que o mecanismo de atenção inclua mais tokens e se aproxime do efeito da autoatenção padrão. Esse argumento é destinado apenas para uso experimental. Em casos de uso gerais, recomendamos usar o tamanho de janela original do modelo.
  • --seed: a sugestão para inicializar todos os geradores de números aleatórios. Mudar esse argumento pode afetar a saída gerada para o mesmo comando, alterando os tokens que são amostrados como os próximos tokens. O valor padrão é 0.

Mecanismo de inferência

  • --num_hosts: o número de hosts a serem executados. O valor padrão é 1. Para mais detalhes, consulte a documentação sobre a configuração da TPU v5e.
  • --disagg_topo: define o número de workers de preenchimento prévio e de decodificação com o recurso experimental de veiculação desagregada. O valor padrão é None. O argumento segue o formato: "number_of_prefill_workers,number_of_decode_workers".
  • --data_parallel_size: o número de réplicas paralelas de dados. O valor padrão é 1. Definir esse valor como N em vez de 1 melhora a capacidade de processamento em aproximadamente N, mantendo a mesma latência.
  • --tensor_parallel_size: o número de réplicas paralelas do tensor. O valor padrão é 1. Aumentar o número de réplicas paralelas de tensor geralmente melhora a latência, porque acelera a multiplicação de matrizes reduzindo o tamanho da matriz.
  • --worker_distributed_method: o método distribuído para iniciar o worker. Use mp para o módulo de multiprocessamento ou ray para a biblioteca Ray. O valor padrão é mp.
  • --enable_jit: se o modo JIT (compilação Just-In-Time) será ativado. O valor padrão é True. A configuração --no-enable_jit o desativa. Ativar o modo JIT melhora o desempenho da inferência, mas exige mais tempo para a compilação inicial. Em geral, os benefícios de desempenho de inferência superam a sobrecarga.
  • --warmup: se o servidor deve ser aquecido com solicitações de amostra durante a inicialização. O valor padrão é True. Definir --no-warmup desativa o recurso. A inicialização é recomendada porque as solicitações iniciais acionam uma compilação mais pesada e, portanto, são mais lentas.
  • --max_prefill_seqs: o número máximo de sequências que podem ser programadas para pré-preenchimento por iteração. O valor padrão é 1. Quanto maior for esse valor, maior será a capacidade de processamento do servidor, mas com possíveis efeitos adversos na latência.
  • --prefill_seqs_padding: o servidor preenche o tamanho do lote de preenchimento prévio para um múltiplo desse valor. O valor padrão é 8. Aumentar esse valor reduz os tempos de recompilação do modelo, mas aumenta o desperdício de computação e a sobrecarga de inferência. A configuração ideal depende do tráfego de solicitações.
  • --prefill_len_padding: o servidor preenche o tamanho da sequência para um múltiplo desse valor. O valor padrão é 512. Aumentar esse valor reduz os tempos de recompilação do modelo, mas aumenta a computação desperdiçada e a sobrecarga de inferência. A configuração ideal depende da distribuição de dados das solicitações.
  • --max_decode_seqs/--max_running_seqs: o número máximo de sequências que podem ser programadas para decodificação por iteração. O valor padrão é 256. Quanto maior for esse valor, maior será a capacidade de processamento do servidor, mas com possíveis efeitos adversos na latência.
  • --decode_seqs_padding: o servidor preenche o tamanho do lote de decodificação para um múltiplo desse valor. O valor padrão é 8. Aumentar esse valor reduz os tempos de recompilação do modelo, mas aumenta o desperdício de computação e a sobrecarga de inferência. A configuração ideal depende do tráfego de solicitações.
  • --decode_blocks_padding: o servidor preenche o número de blocos de memória usados para o cache de chave-valor (cache KV) de uma sequência até um múltiplo desse valor durante a decodificação. O valor padrão é 128. Aumentar esse valor reduz os tempos de recompilação do modelo, mas aumenta a computação desperdiçada e a sobrecarga de inferência. A configuração ideal depende da distribuição de dados das solicitações.
  • --enable_prefix_cache_hbm: se o armazenamento em cache de prefixo será ativado na HBM. O valor padrão é False. Definir esse argumento pode melhorar o desempenho ao reutilizar os cálculos de prefixos compartilhados de solicitações anteriores.
  • --enable_chunked_prefill: se o preenchimento prévio em blocos será ativado. O valor padrão é False. Definir esse argumento pode permitir um tamanho de contexto maior e melhorar o desempenho.

Gerenciamento de memória

  • --hbm_utilization_factor: a porcentagem de memória de alta largura de banda (HBM) do Cloud TPU livre que pode ser alocada para o cache KV após o carregamento dos pesos do modelo. O valor padrão é 0.9. Definir esse argumento como um valor mais alto aumenta o tamanho do cache de KV e pode melhorar a capacidade de processamento, mas aumenta o risco de ficar sem HBM do Cloud TPU durante a inicialização e no ambiente de execução.
  • --num_blocks: número de blocos de dispositivo a serem alocados para o cache de KV. Se esse argumento for definido, o servidor vai ignorar --hbm_utilization_factor. Se esse argumento não for definido, o servidor vai analisar o uso de HBM e calcular o número de blocos de dispositivo a serem alocados com base em --hbm_utilization_factor. Definir esse argumento como um valor mais alto aumenta o tamanho do cache de KV e pode melhorar a capacidade de processamento, mas aumenta o risco de ficar sem HBM da Cloud TPU durante a inicialização e no ambiente de execução.
  • --block_size: número de tokens armazenados em um bloco. As opções possíveis são [8, 16, 32, 2048, 8192]. O valor padrão é 32. Definir esse argumento como um valor maior reduz a sobrecarga no gerenciamento de blocos, mas aumenta o desperdício de memória. O impacto exato no desempenho precisa ser determinado empiricamente.

LoRA dinâmico

  • --enable_lora: se o carregamento dinâmico de adaptadores LoRA do Cloud Storage será ativado. O valor padrão é False. Isso é compatível com a família de modelos Llama.
  • --max_lora_rank: a classificação máxima de LoRA aceita para adaptadores LoRA. definidas nas solicitações. O valor padrão é 16. Definir esse argumento com um valor mais alto permite maior flexibilidade nos adaptadores LoRA que podem ser usados com o servidor, mas aumenta a quantidade de HBM do Cloud TPU alocada para pesos LoRA e diminui a capacidade de processamento.
  • --enable_lora_cache: se o armazenamento em cache de adaptadores LoRA dinâmicos será ativado. O valor padrão é True. A configuração --no-enable_lora_cache o desativa. O armazenamento em cache melhora o desempenho porque elimina a necessidade de fazer o download novamente dos arquivos do adaptador LoRA usados anteriormente.
  • --max_num_mem_cached_lora: o número máximo de adaptadores LoRA armazenados no cache em memória da TPU.O valor padrão é 16. Definir esse argumento com um valor maior melhora a chance de uma ocorrência em cache, mas aumenta o uso do HBM do Cloud TPU.

Também é possível configurar o servidor usando as seguintes variáveis de ambiente:

  • HEX_LLM_LOG_LEVEL: controla a quantidade de informações de registro geradas. O valor padrão é INFO. Defina como um dos níveis de registro padrão do Python definidos no módulo de registro.
  • HEX_LLM_VERBOSE_LOG: se a saída de registro detalhada deve ser ativada. Os valores permitidos são true ou false. O valor padrão é false.

Ajustar argumentos do servidor

Os argumentos do servidor estão inter-relacionados e têm um efeito coletivo no desempenho de exibição. Por exemplo, uma configuração maior de --max_model_len=4096 leva a um maior uso da memória da TPU e, portanto, requer uma alocação de memória maior e menos lotes. Além disso, alguns argumentos são determinados pelo caso de uso, enquanto outros podem ser ajustados. Veja um fluxo de trabalho para configurar o servidor Hex-LLM.

  1. Determine a família e a variante do modelo de interesse. Por exemplo, Llama 3.1 8B Instruct.
  2. Estime o limite inferior da memória de TPU necessária com base no tamanho e na precisão do modelo: model_size * (num_bits / 8). Para um modelo 8B e precisão bfloat16, o limite inferior de memória de TPU necessária seria 8 * (16 / 8) = 16 GB.
  3. Estime o número de chips TPU v5e necessários, em que cada chip v5e oferece 16 GB: tpu_memory / 16. Para um modelo 8B e precisão bfloat16, você precisa de mais de um chip. Entre as configurações de 1, 4 e 8 chips, a menor configuração que oferece mais de 1 chip é a de 4 chips: ct5lp-hightpu-4t. Em seguida, é possível definir --tensor_parallel_size=4.
  4. Determine o tamanho máximo do contexto (tamanho da entrada + tamanho da saída) para o caso de uso pretendido. Por exemplo, 4096. Em seguida, é possível definir --max_model_len=4096.
  5. Ajuste a quantidade de memória de TPU livre alocada para o cache KV para o valor máximo alcançável, considerando as configurações do modelo, do hardware e do servidor (--hbm_utilization_factor). Comece com 0.95. Implante o servidor Hex-LLM e teste-o com comandos longos e alta simultaneidade. Se o servidor ficar sem memória, reduza o fator de utilização de acordo.

Um conjunto de argumentos de amostra para implantar o 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

Um conjunto de argumentos de amostra para implantar o Llama 3.1 70B Instruct AWQ no 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

Solicitar cota do Cloud TPU

No Model Garden, sua cota padrão é de 32 chips do Cloud TPU v5e na região us-west1. Essas cotas se aplicam a implantações com um clique e Implantações de notebooks do Colab Enterprise. Para solicitar um valor de cota maior, consulte Solicitar um ajuste de cota.