Neste tutorial, mostramos como disponibilizar um modelo de linguagem grande (LLM) Gemma usando Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE). Implante um contêiner pré-criado com JetStream e MaxText no GKE. Você também configura o GKE para carregar os pesos do Gemma 7B do Cloud Storage durante a execução.
Este tutorial é destinado a engenheiros de machine learning (ML), administradores e operadores de plataforma e especialistas em dados e IA interessados em usar os recursos de orquestração de contêineres do Kubernetes para veiculação de LLMs. Para saber mais sobre papéis comuns e tarefas de exemplo que mencionamos no conteúdo doGoogle Cloud , consulte Tarefas e funções de usuário comuns do GKE.
Antes de ler esta página, confira se você conhece os seguintes conceitos:
- Modo Autopilot e modo Standard
- Disponibilidade da versão atual da TPU com a arquitetura do sistema do Cloud TPU
- TPUs no GKE
Contexto
Esta seção descreve as principais tecnologias usadas neste tutorial.
Gemma
O Gemma é um conjunto de modelos de inteligência artificial (IA) generativa, leve e abertamente lançados sob licença aberta. Esses modelos de IA estão disponíveis para execução em aplicativos, hardware, dispositivos móveis ou serviços hospedados. É possível usar os modelos Gemma para geração de texto, mas também é possível ajustá-los para tarefas especializadas.
Para saber mais, consulte a documentação do Gemma.
TPUs
TPUs são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Eles são usados para acelerar modelos de machine learning e de IA criados com o uso de frameworks comoTensorFlow , PyTorch eJAX.
Este tutorial aborda a disponibilização do modelo Gemma 7B. O GKE implanta o modelo em nós TPUv5e de host único com topologias de TPU configuradas com base nos requisitos do modelo para exibir prompts com baixa latência.
JetStream
O JetStream é um framework de veiculação de inferência de código aberto desenvolvido pelo Google. O JetStream permite a inferência de alto desempenho, alta capacidade e otimização de memória em TPUs e GPUs. Ele fornece otimizações avançadas de desempenho, incluindo técnicas de quantização e lotes contínuos, para facilitar a implantação de LLMs. O JetStream permite a veiculação de TPU do PyTorch/XLA e do JAX para alcançar o desempenho ideal.
Para saber mais sobre essas otimizações, consulte os repositórios de projetos JetStream PyTorch e JetStream MaxText (links em inglês).
MaxText
MaxText é uma implementação do LLM JAX de alto desempenho, escalonável e adaptável, criada com Linho .Orbax eOptax. A implementação LLM somente decodificador da MaxText é escrita em Python. Ele aproveita muito o compilador XLA para alcançar alto desempenho sem precisar criar kernels personalizados.
Para saber mais sobre os modelos e tamanhos de parâmetros mais recentes compatíveis com a MaxText, consulte o repositório do projeto MaxtText (link em inglês).