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
  • Mistral 7B e Mixtral 8x7B

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.

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

Você pode modificar os seguintes argumentos de inicialização do servidor Hex-LLM para exibição personalizada:

  • --model: o modelo a ser carregado. É possível especificar um ID de modelo do Hugging Face, um caminho absoluto local ou de bucket do Cloud Storage.
  • --tokenizer: o tokenizador a ser carregado. Pode ser um ID de modelo Hugging Face, um caminho absoluto local ou de bucket do Cloud Storage. O valor padrão é igual ao de --model.
  • --enable_jit: se o modo JIT será ativado. O valor padrão é True.
  • --enable_lora: se o modo de carregamento LoRA será ativado. O valor padrão é False.
  • --max_lora_rank: a classificação máxima de LoRA aceita para adaptadores LoRA. definidas nas solicitações. O valor padrão é 16.
  • --data_parallel_size: o número de réplicas paralelas de dados. O valor padrão é 1.
  • --tensor_parallel_size: o número de réplicas paralelas do tensor. O valor padrão é 1.
  • --max_running_seqs: o número máximo de solicitações que o servidor pode processar concomitantemente. Quanto maior for esse argumento, maior será a capacidade de processamento do servidor poderia alcançar, mas com possíveis efeitos adversos na latência. O valor padrão é 256.
  • --hbm_utilization_factor: a porcentagem da Cloud TPU livre HBM que pode ser alocada para o cache KV após o carregamento dos pesos do modelo. Definir esse argumento com um valor mais baixo pode impedir efetivamente HBM do Cloud TPU sem memória. O valor padrão é 0.9.
  • --seed: a sugestão para inicializar todos os geradores de números aleatórios. Mudança esse argumento pode afetar a saída gerada para o mesmo comando. O valor padrão é 0.

Solicitar cota do Cloud TPU

No Model Garden, sua cota padrão é de quatro 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 mais cotas, consulte Solicitar uma a cota.