Publique modelos abertos através do contentor premium Hex-LLM na Cloud TPU

O Hex-LLM, um modelo de linguagem (conteúdo extenso) (MDL/CE) de alta eficiência com o XLA, é a estrutura de serviço de MDL/CE do Vertex AI concebida e otimizada para hardware Cloud TPU. O Hex-LLM combina tecnologias de fornecimento de MDIs/CEs, como o processamento em lote contínuo e o PagedAttention, com otimizações do Vertex AI personalizadas para XLA e Cloud TPU. É um serviço de MDL/CE de alta eficiência e baixo custo nas Cloud TPUs para modelos de código aberto.

O Hex-LLM está disponível no Model Garden através do playground de modelos, da implementação com um clique e do bloco de notas.

Funcionalidades

O Hex-LLM baseia-se em projetos de código aberto com as próprias otimizações da Google para XLA e Cloud TPU. O Hex-LLM alcança um elevado débito e uma baixa latência quando serve MDIs/CEs usados com frequência.

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

  • Algoritmo de processamento em lote contínuo baseado em tokens para ajudar a garantir que os modelos estão a usar totalmente o hardware com um grande número de pedidos simultâneos.
  • Uma reescrita completa dos núcleos de atenção otimizados para XLA.
  • Estratégias de paralelismo de dados e paralelismo de tensores flexíveis e compostas com métodos de divisão de pesos altamente otimizados para executar LLMs de forma eficiente em vários chips de TPU do Google Cloud.

O Hex-LLM suporta uma vasta 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,5 B, 1,5 B e 7 B
  • Qwen-2.5 0.5B, 1.5B, 7B, 14B e 32B

O Hex-LLM também oferece várias funcionalidades, como as seguintes:

  • O Hex-LLM está incluído num único contentor. O Hex-LLM agrupa o servidor da API, o motor de inferência e os modelos suportados numa única imagem do Docker para implementação.
  • Compatível com o formato dos modelos do Hugging Face. O Hex-LLM pode carregar um modelo do Hugging Face a partir do disco local, do Hugging Face Hub e de um contentor do Cloud Storage.
  • Quantização com bitsandbytes e AWQ.
  • Carregamento dinâmico de LoRA. O Hex-LLM consegue carregar os pesos LoRA através da leitura do argumento do pedido durante a publicação.

Funcionalidades avançadas

O Hex-LLM suporta as seguintes funcionalidades avançadas:

  • Publicação em vários anfitriões
  • Publicação desagregada [experimental]
  • Colocação em cache de prefixos
  • Compatibilidade com quantização de 4 bits

Publicação em vários anfitriões

O Hex-LLM suporta agora a publicação de modelos com uma divisão de TPU multi-host. Esta funcionalidade permite-lhe publicar modelos grandes que não podem ser carregados numa única VM de TPU de anfitrião, que contém, no máximo, oito núcleos v5e.

Para ativar esta funcionalidade, defina --num_hosts nos argumentos do contentor Hex-LLM e defina --tpu_topology no pedido de carregamento do modelo do SDK Vertex AI. O exemplo seguinte mostra como implementar o contentor Hex-LLM com uma topologia de TPU 4x4 v5e que publica 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 ponto a ponto sobre a implementação do contentor Hex-LLM com uma topologia de TPU de vários anfitriões, consulte o bloco de notas do Vertex AI Model Garden - Llama 3.1 (implementação).

Em geral, as únicas alterações necessárias para ativar a publicação em vários anfitriões são:

  1. Defina o argumento --tensor_parallel_size para o número total de núcleos na topologia da TPU.
  2. Defina o argumento --num_hosts para o número de anfitriões na topologia da TPU.
  3. Defina --tpu_topology com a API de carregamento de modelos do SDK do Vertex AI.

Publicação desagregada [experimental]

O Hex-LLM suporta agora a publicação desagregada como uma funcionalidade experimental. Só pode ser ativado na configuração de anfitrião único e o desempenho está a ser otimizado.

A publicação desagregada é um método eficaz para equilibrar o tempo até ao primeiro token (TTFT) e o tempo por token de saída (TPOT) para cada pedido, bem como o débito de publicação geral. Separa a fase de preenchimento automático e a fase de descodificação em cargas de trabalho diferentes para que não interfiram umas com as outras. Este método é especialmente útil para cenários que definem requisitos de latência rigorosos.

Para ativar esta funcionalidade, defina --disagg_topo nos argumentos do contentor Hex-LLM. Segue-se um exemplo que mostra como implementar o contentor Hex-LLM na TPU v5e-8 que serve 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, está definido como "3,1" para configurar três trabalhadores de pré-preenchimento e 1 trabalhador de descodificação. Cada trabalhador usa 2 núcleos de TPU v5e.

Colocação em cache de prefixos

A colocação em cache de prefixos reduz o tempo até ao primeiro token (TTFT) para comandos que têm conteúdo idêntico no início do comando, como preâmbulos ao nível da 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 uma cache temporária dos cálculos de tokens de entrada processados para melhorar o TTFT.

Para ativar esta funcionalidade, defina --enable_prefix_cache_hbm nos argumentos do contentor Hex-LLM. Segue-se um exemplo que mostra como implementar o contentor 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 excedam um determinado comprimento (512 tokens por predefinição, configurável através de prefill_len_padding). Os resultados da cache ocorrem em incrementos deste valor, o que garante que a contagem de tokens em cache é 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 comando foram um resultado da cache.

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

Pré-preenchimento segmentado

O preenchimento prévio segmentado divide um preenchimento prévio de pedido em segmentos mais pequenos e mistura o preenchimento prévio e a descodificação num único passo de lote. O Hex-LLM implementa o preenchimento prévio segmentado para equilibrar o tempo até ao primeiro token (TTFT) e o tempo por token de saída (TPOT) e melhora o débito.

Para ativar esta funcionalidade, defina --enable_chunked_prefill nos argumentos do contentor Hex-LLM. Segue-se um exemplo que mostra como implementar o contentor 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,
)

Compatibilidade com quantização de 4 bits

A quantização é uma técnica para reduzir os custos computacionais e de memória da execução da inferência, 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 suporta a quantização apenas de pesos INT8. O apoio técnico alargado inclui modelos com ponderações INT4 quantizadas através da quantização de ponto zero AWQ. O Hex-LLM suporta variantes INT4 das famílias de modelos Mistral, Mixtral e Llama.

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

Comece a usar o Model Garden

O contentor de publicação da Cloud TPU do Hex-LLM está integrado no Model Garden. Pode aceder a esta tecnologia de publicação através do playground, da implementação com um clique e dos exemplos de blocos de notas do Colab Enterprise para uma variedade de modelos.

Use o parque infantil

O playground do Model Garden é um ponto final da Vertex AI pré-implementado que é acessível através do envio de pedidos no cartão do modelo.

  1. Introduza um comando e, opcionalmente, inclua argumentos para o seu pedido.

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

Experimente com o Gemma!

Use a implementação com um clique

Pode implementar um ponto final do Vertex AI personalizado com o Hex-LLM através de um cartão de modelo.

  1. Navegue para a página do cartão do modelo e clique em Implementar.

  2. Para a variação do modelo que quer usar, selecione o tipo de máquina Cloud TPU v5e para implementação.

  3. Clique em Implementar na parte inferior para iniciar o processo de implementação. Recebe duas notificações por email: uma quando o modelo é carregado e outra quando o ponto final está pronto.

Use o bloco de notas do Colab Enterprise

Para flexibilidade e personalização, pode usar exemplos de blocos de notas do Colab Enterprise para implementar um ponto final da Vertex AI com o Hex-LLM através do SDK da Vertex AI para Python.

  1. Navegue para a página do cartão do modelo e clique em Abrir bloco de notas.

  2. Selecione o notebook Vertex Serving. O bloco de notas é aberto no Colab Enterprise.

  3. Execute o bloco de notas para implementar um modelo através do Hex-LLM e envie pedidos de previsão para o ponto final. O fragmento do código para a implementação é o seguinte:

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

Alguns exemplos de blocos de notas do Colab Enterprise:

Configure argumentos do servidor e variáveis de ambiente

Pode definir os seguintes argumentos para iniciar o servidor Hex-LLM. Pode adaptar os argumentos para se adequarem melhor ao exemplo de utilização e aos requisitos pretendidos. Tenha em atenção que os argumentos estão predefinidos para a implementação com um clique, o que permite a experiência de implementação mais fácil. Para personalizar os argumentos, pode basear-se nos exemplos de blocos de notas para referência e definir os argumentos em conformidade.

Modelo

  • --model: o modelo a carregar. Pode especificar um ID do modelo do Hugging Face, um caminho do contentor do Cloud Storage (gs://my-bucket/my-model) ou um caminho local. Espera-se que os artefactos do modelo sigam o formato Hugging Face e usem ficheiros safetensors para os pesos do modelo. Os artefactos de modelos quantizados BitsAndBytes int8 e AWQ são suportados para Llama, Gemma 2 e Mistral/Mixtral.
  • --tokenizer: O tokenizador a carregar. Pode ser um ID do modelo do Hugging Face, um caminho do contentor do Cloud Storage (gs://my-bucket/my-model) ou um caminho local. Se este argumento não estiver definido, é utilizado o valor de --model por predefinição.
  • --tokenizer_mode: o modo de conversão em tokens. As opções possíveis são ["auto", "slow"]. O valor predefinido é "auto". Se esta opção estiver definida como "auto", o tokenizador rápido é usado se estiver disponível. Os tokenizadores lentos são escritos em Python e fornecidos na biblioteca Transformers, enquanto os tokenizadores rápidos que oferecem uma melhoria do 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 deve permitir ficheiros de código remoto definidos nos repositórios de modelos do Hugging Face. O valor predefinido é False.
  • --load_format: formato dos pontos de verificação do modelo a carregar. As opções possíveis são ["auto", "dummy"]. O valor predefinido é "auto". Se esta opção estiver definida como "auto", os pesos do modelo são carregados no formato safetensors. Se esta opção estiver definida como "dummy", os pesos do modelo são inicializados aleatoriamente. Definir esta opção como "dummy" é útil para a experimentação.
  • --max_model_len: o comprimento máximo do contexto (comprimento da entrada mais o comprimento da saída) para servir para o modelo. O valor predefinido é lido a partir do ficheiro de configuração do modelo no formato Hugging Face: config.json. Um comprimento máximo do contexto maior requer mais memória da TPU.
  • --sliding_window: se definido, este argumento substitui o tamanho da janela do modelo para a atenção da janela deslizante. Definir este argumento para um valor maior faz com que o mecanismo de atenção inclua mais tokens e aproxima-se do efeito da autoatenção padrão. Este argumento destina-se apenas a utilização experimental. Nos exemplos de utilização gerais, recomendamos que use o tamanho da janela original do modelo.
  • --seed: a semente para inicializar todos os geradores de números aleatórios. A alteração deste argumento pode afetar o resultado gerado para o mesmo comando através da alteração dos tokens que são amostrados como tokens seguintes. O valor predefinido é 0.

Motor de inferência

  • --num_hosts: o número de anfitriões a executar. O valor predefinido é 1. Para mais detalhes, consulte a documentação sobre a configuração da TPU v5e.
  • --disagg_topo: define o número de trabalhadores de preenchimento e trabalhadores de descodificação com a publicação desagregada da funcionalidade experimental. O valor predefinido é 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 predefinido é 1. A definição deste valor de N para 1 melhora aproximadamente o débito em N, ao mesmo tempo que mantém a mesma latência.
  • --tensor_parallel_size: o número de réplicas paralelas de tensores. O valor predefinido é 1. Aumentar o número de réplicas paralelas de tensores geralmente melhora a latência, porque acelera a multiplicação de matrizes ao reduzir o tamanho da matriz.
  • --worker_distributed_method: o método distribuído para iniciar o trabalhador. Use mp para o módulo multiprocessing ou ray para a biblioteca Ray. O valor predefinido é mp.
  • --enable_jit: se deve ativar o modo JIT (compilação just-in-time). O valor predefinido é True. A definição --no-enable_jit desativa-o. A ativação do modo JIT melhora o desempenho da inferência à custa de exigir tempo adicional gasto na compilação inicial. Em geral, as vantagens de desempenho da inferência superam a sobrecarga.
  • --warmup: indica se o servidor deve ser aquecido com pedidos de amostra durante a inicialização. O valor predefinido é True. A definição --no-warmup desativa-o. Recomendamos o aquecimento, porque os pedidos iniciais acionam uma compilação mais pesada e, por isso, são mais lentos.
  • --max_prefill_seqs: o número máximo de sequências que podem ser agendadas para o preenchimento prévio por iteração. O valor predefinido é 1. Quanto maior for este valor, maior é o débito que o servidor pode alcançar, mas com potenciais efeitos adversos na latência.
  • --prefill_seqs_padding: o servidor preenche o tamanho do lote de pré-preenchimento até um múltiplo deste valor. O valor predefinido é 8. Aumentar este valor reduz os tempos de recompilação do modelo, mas aumenta a computação desperdiçada e a sobrecarga de inferência. A definição ideal depende do tráfego de pedidos.
  • --prefill_len_padding: o servidor preenche o comprimento da sequência até um múltiplo deste valor. O valor predefinido é 512. Aumentar este valor reduz os tempos de recompilação do modelo, mas aumenta a computação desperdiçada e a sobrecarga de inferência. A definição ideal depende da distribuição de dados dos pedidos.
  • --max_decode_seqs/--max_running_seqs: o número máximo de sequências que podem ser agendadas para descodificação por iteração. O valor predefinido é 256. Quanto maior for este valor, maior é o débito que o servidor pode alcançar, mas com potenciais efeitos adversos na latência.
  • --decode_seqs_padding: o servidor preenche o tamanho do lote de descodificação até um múltiplo deste valor. O valor predefinido é 8. Aumentar este valor reduz os tempos de recompilação do modelo, mas aumenta a computação desperdiçada e a sobrecarga de inferência. A definição ideal depende do tráfego de pedidos.
  • --decode_blocks_padding: o servidor preenche o número de blocos de memória usados para uma cache de chave-valor (cache KV) de uma sequência com um múltiplo deste valor durante a descodificação. O valor predefinido é 128. Aumentar este valor reduz os tempos de recompilação do modelo, mas aumenta a computação desperdiçada e a sobrecarga de inferência. A definição ideal depende da distribuição de dados dos pedidos.
  • --enable_prefix_cache_hbm: se deve ativar o armazenamento em cache de prefixos no HBM. O valor predefinido é False. A definição deste argumento pode melhorar o desempenho através da reutilização dos cálculos de prefixos partilhados de pedidos anteriores.
  • --enable_chunked_prefill: se deve ativar o preenchimento prévio segmentado. O valor predefinido é False. A definição deste argumento pode suportar um comprimento do contexto maior e melhorar o desempenho.

Gestão de memória

  • --hbm_utilization_factor: a percentagem de memória de largura de banda elevada (HBM) do Cloud TPU gratuita que pode ser alocada para a cache KV após o carregamento dos pesos do modelo. O valor predefinido é 0.9. Definir este argumento para um valor mais elevado aumenta o tamanho da cache KV e pode melhorar o débito, mas aumenta o risco de ficar sem HBM do Cloud TPU durante a inicialização e o tempo de execução.
  • --num_blocks: número de blocos de dispositivos a atribuir à cache KV. Se este argumento estiver definido, o servidor ignora --hbm_utilization_factor. Se este argumento não estiver definido, o servidor cria perfis de utilização de HBM e calcula o número de blocos de dispositivos a atribuir com base em --hbm_utilization_factor. Definir este argumento para um valor mais elevado aumenta o tamanho da cache KV e pode melhorar o débito, mas aumenta o risco de ficar sem HBM do Cloud TPU durante a inicialização e o tempo de execução.
  • --block_size: número de tokens armazenados num bloco. As opções possíveis são [8, 16, 32, 2048, 8192]. O valor predefinido é 32. Definir este argumento para um valor maior reduz a sobrecarga na gestão de blocos, à custa de um maior desperdício de memória. O impacto exato no desempenho tem de ser determinado empiricamente.

LoRA dinâmico

  • --enable_lora: Indica se deve ativar o carregamento dinâmico de adaptadores LoRA a partir do Cloud Storage. O valor predefinido é False. Isto é suportado para a família de modelos Llama.
  • --max_lora_rank: O limite máximo da classificação LoRA suportado para adaptadores LoRA definidos em pedidos. O valor predefinido é 16. Definir este argumento para um valor mais elevado permite uma maior flexibilidade nos adaptadores LoRA que podem ser usados com o servidor, mas aumenta a quantidade de HBM de TPU na nuvem alocada para ponderações LoRA e diminui o débito.
  • --enable_lora_cache: se deve ativar a colocação em cache de adaptadores LoRA dinâmicos. O valor predefinido é True. A definição --no-enable_lora_cache desativa-o. A colocação em cache melhora o desempenho porque elimina a necessidade de transferir novamente os ficheiros do adaptador LoRA usados anteriormente.
  • --max_num_mem_cached_lora: O número máximo de adaptadores LoRA armazenados na cache de memória da TPU.O valor predefinido é 16. Definir este argumento para um valor maior melhora a probabilidade de um acerto da cache, mas aumenta a quantidade de utilização de HBM do Cloud TPU.

Também pode configurar o servidor através das seguintes variáveis de ambiente:

  • HEX_LLM_LOG_LEVEL: controla a quantidade de informações de registo geradas. O valor predefinido é INFO. Defina esta opção para um dos níveis de registo padrão do Python definidos no módulo de registo.
  • HEX_LLM_VERBOSE_LOG: se deve ativar o resultado do registo detalhado. Os valores permitidos são true ou false. O valor predefinido é false.

Ajuste os argumentos do servidor

Os argumentos do servidor estão interligados e têm um efeito coletivo no desempenho da publicação. Por exemplo, uma definição maior de --max_model_len=4096 resulta numa utilização de memória da TPU mais elevada e, por isso, requer uma atribuição de memória maior e menos processamento em lote. Além disso, alguns argumentos são determinados pelo exemplo de utilização, enquanto outros podem ser ajustados. Segue-se um fluxo de trabalho para configurar o servidor Hex-LLM.

  1. Determine a família de modelos e a variante do 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 do modelo e na precisão: model_size * (num_bits / 8). Para um modelo de 8B e uma precisão bfloat16, o limite inferior da memória da TPU necessária seria de 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 uma precisão bfloat16, precisa de mais de 1 chip. Entre as configurações de 1 chip, 4 chips e 8 chips, a configuração mais pequena que oferece mais de 1 chip é a configuração de 4 chips: ct5lp-hightpu-4t. Posteriormente, pode definir o valor --tensor_parallel_size=4.
  4. Determine o comprimento máximo do contexto (comprimento da entrada + comprimento da saída) para o caso de utilização pretendido. Por exemplo, 4096. Posteriormente, pode definir o valor --max_model_len=4096.
  5. Ajuste a quantidade de memória TPU livre alocada para a cache KV ao valor máximo alcançável, tendo em conta o modelo, o hardware e as configurações do servidor (--hbm_utilization_factor). Comece com 0.95. Implemente o servidor Hex-LLM e teste-o com comandos longos e alta concorrência. Se o servidor ficar sem memória, reduza o fator de utilização em conformidade.

Um conjunto de argumentos de exemplo para implementar 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 exemplo para implementar 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

Peça uma quota do Cloud TPU

No Model Garden, a sua quota predefinida é de 32 chips Cloud TPU v5e na região us-west1. Estas quotas aplicam-se a implementações de um clique e a implementações de blocos de notas do Colab Enterprise. Para pedir um valor de quota mais elevado, consulte o artigo Peça um ajuste de quota.