Publique GMLs com TPUs multi-host no GKE com o JetStream e o Pathways

Este guia mostra como publicar modelos de linguagem (conteúdo extenso) (MDI/CE) de última geração, como o Llama 3.1 405B, no Google Kubernetes Engine (GKE) usando unidades de processamento tensorial (TPUs) em vários nós.

Este guia demonstra como usar tecnologias portáteis de código aberto, como o Kubernetes, o JetStream, os Pathways on Cloud e a API LeaderWorkerSet (LWS), para implementar e publicar cargas de trabalho de IA/ML no GKE, tirando partido do controlo detalhado, da escalabilidade, da resiliência, da portabilidade e da rentabilidade do GKE.

Importante: para aceder aos Caminhos, contacte o representante da sua conta do Google Cloud.

Contexto

Os modelos de linguagem (conteúdo extenso) (MDIs/CEs) aumentaram de tamanho e já não cabem numa única fatia de TPU do anfitrião. Para a inferência de ML, pode usar o Pathways on Cloud para executar a inferência multi-anfitrião em grande escala no GKE em vários nós de TPU interligados. Neste guia, vai aprender a aprovisionar um cluster do GKE com as fatias de TPU de vários anfitriões e usar os binários do Pathways on Cloud, iniciar o servidor JetStream com a framework MaxText e fazer pedidos de inferência de vários anfitriões.

Ao publicar um MDG com TPUs no GKE com o JetStream, o MaxText e o Pathways, pode criar uma solução de publicação robusta e pronta para produção com todas as vantagens do Kubernetes gerido, incluindo rentabilidade, escalabilidade e maior disponibilidade. Esta secção descreve as tecnologias principais usadas neste tutorial.

Acerca das TPUs

As TPUs são circuitos integrados específicos da aplicação (ASICs) desenvolvidos pela Google e usados para acelerar a aprendizagem automática e os modelos de IA criados com frameworks como o TensorFlow, o PyTorch e o JAX.

Antes de usar as TPUs no GKE, recomendamos que conclua o seguinte percurso de aprendizagem:

1. Saiba mais sobre a disponibilidade da versão atual da TPU com a arquitetura do sistema Cloud TPU.
2. Saiba mais sobre as TPUs no GKE.

Este tutorial aborda a publicação do modelo Llama 3.1-405B. O GKE implementa o modelo em nós da TPU v6e com vários anfitriões com topologias de TPU configuradas com base nos requisitos do modelo para pedidos de serviço com baixa latência.

Percursos no Google Cloud

O Pathways é uma camada de orquestração em grande escala para aceleradores. O Pathways foi explicitamente concebido para permitir a exploração de novos sistemas e ideias de investigação de ML, mantendo o desempenho de vanguarda dos modelos atuais. O Pathways permite que um único processo de cliente JAX coordene a computação em uma ou mais grandes fatias de TPU, simplificando as computações de ML que abrangem centenas ou milhares de chips de TPU.

JetStream

O JetStream é uma framework de serviço de inferência de código aberto desenvolvida pela Google. O JetStream permite uma inferência de alto desempenho, com elevado débito e otimizada em função da memória em TPUs e GPUs. O JetStream oferece otimizações de desempenho avançadas, incluindo o processamento em lote contínuo, otimizações da cache KV e técnicas de quantização, para facilitar a implementação de LLMs. O JetStream permite o serviço de TPU do PyTorch/XLA e do JAX para otimizar o desempenho.

MaxText

O MaxText é uma implementação de LLM JAX de alto desempenho, escalável e adaptável, criada com base em bibliotecas JAX de código aberto, como o Flax, o Orbax e o Optax. A implementação do LLM apenas com descodificador do MaxText está escrita em Python. Tira partido do compilador XLA para alcançar um elevado desempenho sem precisar de criar kernels personalizados.

Para mais informações sobre os modelos e os tamanhos de parâmetros mais recentes suportados pelo MaxText, consulte o repositório do projeto MaxText.

Llama 3.1 405B

O Llama 3.1 405B é um modelo de linguagem (conteúdo extenso) da Meta concebido para uma variedade de tarefas de processamento de linguagem natural, incluindo geração de texto, tradução e resposta a perguntas. O GKE oferece a infraestrutura necessária para suportar as necessidades de publicação e preparação distribuída de modelos desta escala.

Para mais informações, consulte a documentação do Llama.

Arquitetura

Esta secção descreve a arquitetura do GKE usada neste tutorial. A arquitetura inclui um cluster GKE Standard que aprovisiona TPUs e aloja componentes JetStream e Pathways para implementar e publicar o modelo.

O diagrama seguinte mostra os componentes desta arquitetura:

Esta arquitetura inclui os seguintes componentes:

• Um cluster regional do GKE Standard.
• Um conjunto de nós de segmentação de TPU com vários anfitriões que aloja a implementação do JetStream e os componentes do Pathways.
• O Pathways resource manager gere os recursos do acelerador e coordena a atribuição de aceleradores para tarefas do utilizador.
• As coordenadas Pathways client com o Pathways resource manager para determinar onde os programas compilados são colocados para execução.
• O Pathways worker é executado e realiza cálculos em máquinas aceleradoras, e envia dados de volta para a sua carga de trabalho através do servidor proxy IFRT.
• O IFRT proxy client implementa a API Interim Framework Runtime (IFRT) de SO e atua como a ponte de comunicação entre a sua carga de trabalho e os componentes do Pathways.
• O IFRT proxy server recebe pedidos do IFRT proxy client e encaminha-os para o Pathways client, distribuindo o trabalho.
• O contentor JetStream-Pathways fornece um servidor de inferência baseado em JAX que recebe pedidos de inferência e delega os respetivos processos de execução ao Pathways workers
• O componente Service distribui o tráfego de entrada por todas as réplicas JetStream HTTP.
• JetStream HTTP é um servidor HTTP que aceita pedidos como um wrapper para o formato necessário do JetStream e envia-o para o cliente GRPC do JetStream.