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, é a Framework de disponibilização de LLMs da Vertex AI projetado e otimizado para hardware do Cloud TPU. O Hex-LLM combina tecnologias de veiculação de LLM como lotes contínuos e PagedAttentioncom as otimizações da Vertex AI personalizados 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, raciocínio do Phi-4 e raciocínio 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 da API, mecanismo de inferência e modelos com suporte em uma única imagem Docker para ser implantados.
  • 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 à veiculação de modelos com uma fração de TPU de vários hosts. Com esse recurso, é possível veicular modelos grandes que não podem ser carregados em uma única VM de TPU de host, que contém no máximo oito núcleos v5e.

Para ativar esse recurso, defina --num_hosts nos argumentos do contêiner 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 TPU 4x4 v5e que veicula 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 multihost, consulte o notebook Vertex AI Model Garden - Llama 3.1 (implantação).

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 de TPU.
  2. Defina o argumento --num_hosts como o número de hosts na topologia de 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 à veiculação desagregada como um recurso experimental. Ele só pode ser ativado na configuração de host único, e o desempenho está sendo ajustado.

A disponibilização desagregada é um método eficaz para equilibrar o tempo até o primeiro token (TTFT) e o tempo por token de saída (TPOT) de cada solicitação, além da taxa de transferência geral de disponibilização. Ele separa a fase de pré-preenchimento e a fase de decodificação em diferentes cargas de trabalho para que elas não interfiram umas nas outras. 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 Hex-LLM. Confira um exemplo de como implantar o contêiner Hex-LLM na TPU v5e-8, que atende ao 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 está definido como "3,1" para configurar três workers de pré-preenchimento e um worker de decodificação. Cada worker usa dois núcleos de TPU v5e.

Armazenamento em cache de prefixo

O cache de prefixo reduz o tempo para o primeiro token (TTFT) em comandos que têm conteúdo idêntico no início, como preâmbulos em 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. Confira um exemplo que mostra como implantar o contêiner Hex-LLM na TPU v5e-8, que veicula 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 comprimento (512 tokens por padrão, configurável usando prefill_len_padding). Os acertos de 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 quantos dos tokens de solicitação foram um ocorrência em cache.

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

Pré-preenchimento em partes

O pré-preenchimento em partes divide um pré-preenchimento de solicitação em partes menores e mistura pré-preenchimento e decodificação em uma etapa de lote. O Hex-LLM implementa o pré-enchimento em partes para equilibrar o tempo até o primeiro token (TTFT) e o tempo por token de saída (TPOT) e melhora a capacidade de processamento.

Para ativar esse recurso, defina --enable_chunked_prefill nos argumentos do contêiner Hex-LLM. Confira um exemplo de como implantar o contêiner Hex-LLM na TPU v5e-8, que veicula 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. Ela representa os pesos ou as ativações com tipos de dados de baixa precisão, como INT8 ou INT4, em vez dos tipos BF16 ou FP32 comuns.

O Hex-LLM oferece suporte à 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 oferece suporte a variantes INT4 das famílias de modelos Mistral, Mixtral e Llama.

Não é necessário usar outra flag para veicular 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 Model Garden Playground é uma solução pré-implantada endpoint acessível 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 Vertex AI 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 o Colab Enterprise exemplos de notebook para implantar um endpoint da Vertex AI com o Hex-LLM usando o SDK da Vertex AI 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 Enterprise:

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 encaixem melhor no caso de uso e nos requisitos pretendidos. Os argumentos são predefinidos para implantação com um clique, o que facilita a experiência. 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 isso 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 é permitido usar arquivos de código remotos definidos nos repositórios de modelos do Hugging Face. 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 esse valor for definido como "dummy", os pesos do modelo serão inicializados aleatoriamente. Definir como "dummy" é útil para experimentação.
  • --max_model_len: o comprimento máximo do contexto (comprimento da entrada mais o comprimento da saída) para veiculação do modelo. O valor padrão é lido do arquivo de configuração do modelo no formato Hugging Face: config.json. Um tamanho máximo de contexto maior exige mais memória da TPU.
  • --sliding_window: se definido, esse argumento vai substituir o tamanho da janela do modelo para a atenção de janela deslizante. Definir esse argumento com 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 ao uso experimental. Em casos de uso gerais, recomendamos usar o tamanho original da janela 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 ao mudar 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 pré-preenchimento 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 isso como N em vez de 1 melhora aproximadamente a taxa de transferência em 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 de paralelismo de tensores geralmente melhora a latência, porque acelera a multiplicação de matrizes ao reduzir o tamanho delas.
  • --worker_distributed_method: o método distribuído para iniciar o worker. Use mp para o módulo multiprocessing 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. Se definir --no-enable_jit, ela será desativada. Ativar o modo JIT melhora o desempenho da inferência, mas exige mais tempo na compilação inicial. Em geral, os benefícios de desempenho da inferência superam a sobrecarga.
  • --warmup: indica se o servidor deve ser aquecido com solicitações de amostra durante a inicialização. O valor padrão é True. Se definir --no-warmup, ela será desativada. O pré-aquecimento é recomendado 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é-enchimento 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 adiciona padding ao tamanho do lote de pré-preenchimento 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 com um múltiplo desse valor. O valor padrão é 512. 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 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 faz padding do 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 faz padding do número de blocos de memória usados para um 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 o desperdício de computação 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 no HBM. O valor padrão é False. Definir esse argumento pode melhorar o desempenho reutilizando os cálculos de prefixos compartilhados de solicitações anteriores.
  • --enable_chunked_prefill: se a pré-ação em partes está ativada. O valor padrão é False. Definir esse argumento pode aumentar o tamanho do contexto e melhorar a performance.

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 com um valor maior aumenta o tamanho do cache de KV e pode melhorar a capacidade de transmissão, mas aumenta o risco de ficar sem HBM da TPU do Cloud durante a inicialização e o tempo de execução.
  • --num_blocks: número de blocos de dispositivos 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 criar perfis de uso de HBM e calcular o número de blocos de dispositivos a serem alocados com base em --hbm_utilization_factor. Definir esse argumento com um valor mais alto aumenta o tamanho do cache 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 tempo 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 com um valor maior reduz o overhead no gerenciamento de blocos, mas aumenta o desperdício de memória. O impacto exato na performance precisa ser determinado empiricamente.

LoRA dinâmica

  • --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 definidos 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 da Cloud TPU alocada para pesos de LoRA e diminui a capacidade de processamento.
  • --enable_lora_cache: se o armazenamento em cache de adaptadores dinâmicos de LoRA será ativado. O valor padrão é True. Se definir --no-enable_lora_cache, ela será desativada. O armazenamento em cache melhora o desempenho porque remove a necessidade de baixar novamente arquivos de adaptadores LoRA usados anteriormente.
  • --max_num_mem_cached_lora: o número máximo de adaptadores LoRA armazenados no cache de memória da TPU.O valor padrão é 16. Definir esse argumento com um valor maior aumenta a chance de um ocorrência em cache, mas aumenta a quantidade de 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 geração de registros padrão do Python definidos no módulo de geração de registros.
  • HEX_LLM_VERBOSE_LOG: indica se a saída de registro detalhada 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 da veiculação. Por exemplo, uma configuração maior de --max_model_len=4096 leva a um uso maior da memória de TPU e, portanto, exige uma alocação de memória maior e menos loteamento. Além disso, alguns argumentos são determinados pelo caso de uso, enquanto outros podem ser ajustados. Confira um fluxo de trabalho para configurar o servidor Hex-LLM.

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

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