Disponibilizar LLMs usando TPUs de vários hosts no GKE com o JetStream e o Pathways

Este guia mostra como disponibilizar modelos de linguagem grandes (LLMs) de última geração, como o Llama 3.1 405B no Google Kubernetes Engine (GKE), usando unidades de processamento de tensor (TPUs) em vários nós.

Este guia demonstra como usar tecnologias portáteis de código aberto (Kubernetes, JetStream, Programas de aprendizado no Cloud e a API LeaderWorkerSet (LWS)) para implantar e disponibilizar cargas de trabalho de IA/ML no GKE, aproveitando o controle granular, a escalonabilidade, a resiliência, a portabilidade e a relação custo-benefício do GKE.

Importante: para ter acesso aos programas, entre em contato com o representante da sua conta do Google Cloud.

Contexto

Os modelos de linguagem grandes aumentaram de tamanho e não cabem mais em uma única fração de TPU de host. Para inferência de ML, use os Programas no Cloud (em inglês) para executar inferências em grande escala e vários hosts no GKE em vários nós de TPU interconectados. Neste guia, você vai aprender a provisionar um cluster do GKE com as frações de TPU de vários hosts e usar os binários do Pathways on Cloud, iniciar o servidor JetStream com o framework MaxText e fazer solicitações de inferência de vários hosts.

Ao disponibilizar um LLM usando TPUs no GKE com JetStream, MaxText e Pathways, é possível criar uma solução de exibição robusta e pronta para produção com todos os benefícios do Kubernetes gerenciado, incluindo economia, escalonabilidade e maior disponibilidade. Esta seção descreve as principais tecnologias usadas neste tutorial.

Sobre 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.

Antes de usar TPUs no GKE, recomendamos que você conclua o seguinte programa de aprendizado:

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

Este tutorial aborda a disponibilização do modelo Llama 3.1-405B. O GKE implanta o modelo em nós TPU v6e de vários hosts com topologias de TPU configuradas com base nos requisitos do modelo para exibir prompts com baixa latência.

Programas do Pathways on Cloud

Os Pathways são uma camada de orquestração em grande escala para aceleradores. Os programas são projetados explicitamente para permitir a exploração de novos sistemas e ideias de pesquisa de ML, mantendo o desempenho de ponta para os modelos atuais. Os Pathways permitem que um único processo cliente do JAX coordene a computação em uma ou mais frações grandes de TPU, simplificando as computações de ML que abrangem centenas ou milhares de chips de TPU.

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. O JetStream oferece otimizações avançadas de desempenho, incluindo lotes contínuos, otimizações de cache KV e técnicas de quantização, para facilitar a implantação de LLMs. O JetStream permite a veiculação de TPU do PyTorch/XLA e do JAX para otimizar o desempenho.

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 MaxText (link em inglês).

Llama 3.1 405B

O Llama 3.1 405B é um modelo de linguagem grande da Meta projetado para várias tarefas de processamento de linguagem natural, incluindo geração de texto, tradução e respostas a perguntas. O GKE oferece a infraestrutura necessária para atender às necessidades de treinamento e disponibilização distribuídos de modelos dessa escala.

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

Arquitetura

Nesta seção, descrevemos a arquitetura do GKE usada neste tutorial. A arquitetura inclui um cluster do GKE Standard que provisiona TPUs e hospeda componentes do JetStream e do Pathways para implantar e disponibilizar o modelo.

O diagrama a seguir mostra os componentes dessa arquitetura:

Essa arquitetura inclui os seguintes componentes:

• Um cluster regional do GKE Standard.
• Um pool de nós de fração de TPU de vários hosts que hospeda a implantação do JetStream e os componentes do Pathways.
• O Pathways resource manager gerencia recursos de acelerador e coordena a alocação de aceleradores para jobs do usuário.
• O Pathways client coordena 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, enviando dados de volta para sua carga de trabalho pelo servidor proxy IFRT.
• O IFRT proxy client implementa a API Interim Framework Runtime (IFRT) de código aberto e atua como ponte de comunicação entre sua carga de trabalho e os componentes do Pathways.
• O IFRT proxy server recebe solicitações do IFRT proxy client e as encaminha para o Pathways client, distribuindo o trabalho.
• O contêiner JetStream-Pathways fornece um servidor de inferência baseado em JAX que recebe solicitações de inferência e delega os processos de execução ao Pathways workers.
• O componente de serviço distribui o tráfego de entrada para todas as réplicas JetStream HTTP.
• O JetStream HTTP é um servidor HTTP que aceita solicitações como um wrapper para o formato obrigatório do JetStream e as envia para o cliente GRPC do JetStream (link em inglês).