Este guia mostra como exibir um modelo de linguagem grande (LLM) usando Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE) com JetStream através do PyTorch. Neste guia, você baixa pesos de modelo para o Cloud Storage e os implanta em um cluster do GKE Autopilot ou Standard usando um contêiner que executa o JetStream.
Este guia é um bom ponto de partida se você precisar da escalonabilidade, resiliência e economia oferecidos pelos recursos do Kubernetes ao implantar o modelo no JetStream.
Este guia é destinado a clientes de IA generativa que usam o PyTorch, usuários novos ou atuais do GKE, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas interessados em usar os recursos de orquestração de contêineres do Kubernetes para veiculação de LLMs.
Contexto
Ao disponibilizar o LLM usando TPUs no GKE com o JetStream, é 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:
- Saiba mais sobre a disponibilidade atual da versão da TPU com a arquitetura do sistema do Cloud TPU.
- Saiba mais sobre TPUs no GKE.
Este tutorial aborda a disponibilização de vários modelos de LLM. 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.
Sobre o 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 alcançar o desempenho ideal.
Lotes contínuos
O agrupamento contínuo é uma técnica que agrupa dinamicamente as solicitações de inferência recebidas em lotes, reduzindo a latência e aumentando a capacidade de processamento.
Quantização de cache KV
A quantização do cache KV envolve a compactação do cache de valor-chave usado em mecanismos de atenção, reduzindo os requisitos de memória.
Quantização de peso Int8
A quantização de peso Int8 reduz a precisão dos pesos do modelo de ponto flutuante de 32 bits para inteiros de 8 bits, resultando em computação mais rápida e uso reduzido da memória.
Para saber mais sobre essas otimizações, consulte os repositórios de projetos JetStream PyTorch e JetStream MaxText (links em inglês).
Sobre o PyTorch
O PyTorch é um framework de aprendizado de máquina de código aberto desenvolvido pela Meta e agora faz parte da Linux Foundation. O PyTorch fornece recursos de alto nível, como computação de tensor e redes neurais profundas.
Objetivos
- Prepare um cluster do Autopilot ou do GKE Standard com a topologia de TPU recomendada com base nas características do modelo.
- Implantar componentes do JetStream no GKE.
- Receba e publique seu modelo.
- Disponibilizar e interagir com o modelo publicado.
Arquitetura
Nesta seção, descrevemos a arquitetura do GKE usada neste tutorial. A arquitetura inclui um cluster do Autopilot ou padrão do GKE que provisiona TPUs e hospeda componentes do JetStream para implantar e disponibilizar os modelos.
O diagrama a seguir mostra os componentes dessa arquitetura:
Essa arquitetura inclui os seguintes componentes:
- Um cluster regional ou Autopilot do GKE.
- Dois pools de nós de frações de TPU de host único que hospedam a implantação do JetStream.
- 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). - O
JetStream-PyTorch
é um servidor JetStream que realiza inferências com lotes contínuos.
Antes de começar
- Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the required API.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the required API.
-
Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin
Check for the roles
-
In the Google Cloud console, go to the IAM page.
Go to IAM - Select the project.
-
In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.
- For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.
Grant the roles
-
In the Google Cloud console, go to the IAM page.
Acessar o IAM - Selecionar um projeto.
- Clique em CONCEDER ACESSO.
-
No campo Novos principais, insira seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.
- Na lista Selecionar um papel, escolha um.
- Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
- Clique em Salvar.
-
- Verifique se você tem cota suficiente para oito chips TPU v5e PodSlice Lite. Neste tutorial, você usará as instâncias sob demanda.
- Crie uma conta do token do Hugging Face caso ainda não tenha uma.
Receber acesso ao modelo
Tenha acesso a vários modelos no Hugging Face para implantação no GKE
Gemma 7B-it
Para ter acesso ao modelo Gemma para implantação no GKE, primeiro assine o contrato de consentimento de licença.
- Acesse a página de consentimento do modelo Gemma no Hugging Face.
- Faça login no Hugging Face, caso ainda não tenha feito isso.
- Leia e aceite os Termos e Condições do modelo.
Llama 3 8B
Para ter acesso ao modelo Llama 3 para implantação no GKE, primeiro assine o contrato de consentimento de licença.
- Acesse a página de consentimento do modelo Llama 3 no Hugging Face.
- Faça login no Hugging Face, caso ainda não tenha feito isso.
- Leia e aceite os Termos e Condições do modelo.
Prepare o ambiente
Neste tutorial, você usará o Cloud Shell para gerenciar recursos hospedados no
Google Cloud. O Cloud Shell vem pré-instalado com o software necessário
para este tutorial, incluindo
kubectl
e
gcloud CLI.
Para configurar o ambiente com o Cloud Shell, siga estas etapas:
No console do Google Cloud, inicie uma sessão do Cloud Shell clicando em Ativar o Cloud Shell no console do Google Cloud. Isso inicia uma sessão no painel inferior do console do Google Cloud.
Defina as variáveis de ambiente padrão:
gcloud config set project PROJECT_ID export PROJECT_ID=$(gcloud config get project) export CLUSTER_NAME=CLUSTER_NAME export BUCKET_NAME=BUCKET_NAME export REGION=REGION export LOCATION=LOCATION export CLUSTER_VERSION=CLUSTER_VERSION
Substitua os seguintes valores:
PROJECT_ID
: é o ID do projeto do Google Cloud.CLUSTER_NAME
: o nome do cluster do GKE.BUCKET_NAME
: o nome do bucket do Cloud Storage. Não é necessário especificar o prefixogs://
.REGION
: a região em que o cluster do GKE, o bucket do Cloud Storage e os nós da TPU estão localizados. A região contém zonas em que os tipos de máquina TPU v5e estão disponíveis (por exemplo,us-west1
,us-west4
,us-central1
,us-east1
,us-east5
oueurope-west4
). Para clusters do Autopilot, verifique se você tem recursos zonais de TPU v5e suficientes para a região escolhida.- (Apenas cluster padrão)
LOCATION
: a zona em que os recursos da TPU estão disponíveis, por exemplo,us-west4-a
. Para clusters do Autopilot, não é necessário especificar a zona, apenas a região. CLUSTER_VERSION
: a versão do GKE, que precisa ser compatível com o tipo de máquina que você quer usar. A versão padrão do GKE pode não ter disponibilidade para sua TPU de destino. Para conferir uma lista de versões mínimas do GKE disponíveis por tipo de máquina de TPU, consulte Disponibilidade de TPU no GKE.
Criar e configurar recursos do Google Cloud
Siga estas instruções para criar os recursos necessários.
Crie um cluster do GKE
É possível disponibilizar o Gemma em TPUs em um cluster do GKE Autopilot ou Standard. Recomendamos que você use um cluster do Autopilot para ter uma experiência totalmente gerenciada do Kubernetes. Para escolher o modo de operação do GKE mais adequado para suas cargas de trabalho, consulte Escolher um modo de operação do GKE.
Piloto automático
Crie um cluster do GKE Autopilot:
gcloud container clusters create-auto CLUSTER_NAME \
--project=PROJECT_ID \
--region=REGION \
--cluster-version=CLUSTER_VERSION
Padrão
Criar um cluster regional do GKE Standard que use a federação de identidade da carga de trabalho para o GKE:
gcloud container clusters create CLUSTER_NAME \ --enable-ip-alias \ --machine-type=e2-standard-4 \ --num-nodes=2 \ --cluster-version=CLUSTER_VERSION \ --workload-pool=PROJECT_ID.svc.id.goog \ --location=REGION
A criação do cluster pode levar vários minutos.
Crie um pool de nós de TPU v5e com uma topologia
2x4
e dois nós:gcloud container node-pools create tpu-nodepool \ --cluster=CLUSTER_NAME \ --machine-type=ct5lp-hightpu-8t \ --project=PROJECT_ID \ --num-nodes=2 \ --region=REGION \ --node-locations=LOCATION
Criar um bucket do Cloud Storage
Crie um bucket do Cloud Storage para armazenar o checkpoint convertido:
gcloud storage buckets create gs://BUCKET_NAME --location=REGION
Gerar o token da CLI do Hugging Face no Cloud Shell
Gere um novo token do Huggin' Face, caso ainda não tenha um:
- Clique em Seu perfil > Configurações > Tokens de acesso.
- Clique em Novo token.
- Especifique um Nome de sua escolha e um Papel de pelo menos
Read
. - Clique em Generate a token.
- Edite as permissões do seu token de acesso para ter acesso de leitura ao repositório Hugging Face do seu modelo.
- Copie o token gerado para a área de transferência.
Criar um Secret do Kubernetes para as credenciais do Hugging Face
No Cloud Shell, faça o seguinte:
Configure
kubectl
para se comunicar com o cluster:gcloud container clusters get-credentials CLUSTER_NAME --location=REGION
Crie um secret para armazenar as credenciais do Hugging Face:
kubectl create secret generic huggingface-secret \ --from-literal=HUGGINGFACE_TOKEN=HUGGINGFACE_TOKEN
Substitua
HUGGINGFACE_TOKEN
pelo seu token do Hugging Face.
Configurar o acesso das cargas de trabalho usando a federação de identidade da carga de trabalho do GKE
Atribua uma Kubernetes ServiceAccount ao aplicativo e configure-a para atuar como uma conta de serviço do IAM.
Crie uma conta de serviço do IAM para seu aplicativo:
gcloud iam service-accounts create wi-jetstream
Adicione uma vinculação de política do IAM à sua conta de serviço do IAM para gerenciar o Cloud Storage:
gcloud projects add-iam-policy-binding PROJECT_ID \ --member "serviceAccount:wi-jetstream@PROJECT_ID.iam.gserviceaccount.com" \ --role roles/storage.objectUser gcloud projects add-iam-policy-binding PROJECT_ID \ --member "serviceAccount:wi-jetstream@PROJECT_ID.iam.gserviceaccount.com" \ --role roles/storage.insightsCollectorService
Permita que a conta de serviço do Kubernetes atue como a conta de serviço do IAM adicionando uma vinculação de política do IAM entre as duas. Essa vinculação permite que a Kubernetes ServiceAccount atue como a conta de serviço do Google.
gcloud iam service-accounts add-iam-policy-binding wi-jetstream@PROJECT_ID.iam.gserviceaccount.com \ --role roles/iam.workloadIdentityUser \ --member "serviceAccount:PROJECT_ID.svc.id.goog[default/default]"
Anote a conta de serviço do Kubernetes com o endereço de e-mail da conta de serviço do IAM
kubectl annotate serviceaccount default \ iam.gke.io/gcp-service-account=wi-jetstream@PROJECT_ID.iam.gserviceaccount.com
Implantar o JetStream
Implante o contêiner JetStream para exibir o modelo:
Salve o seguinte manifesto como jetstream-pytorch-deployment.yaml
:
Gemma 7B-it
Llama 3 8B
O manifesto define as seguintes propriedades principais:
model_id
: o nome do modelo do Hugging Face (google/gemma-7b-it
,meta-llama/Meta-Llama-3-8B
). Consulte os modelos compatíveis.override_batch_size
: o tamanho do lote de decodificação por dispositivo, em que um chip de TPU é igual a um dispositivo. O padrão é32
.working_dir
: o diretório de trabalho em que o ponto de verificação do modelo está ou será armazenado. Como o bucket do Cloud Storage está montado no seu nó, a carga de trabalho pode acessar os pontos de verificação emgs://BUCKET_NAME/pytorch/org/repo
. O servidor JetStream-PyTorch faz o download dos pesos do modelo do Hugging Face para o diretório de trabalho no primeiro uso. O carregamento de um modelo de 7B ou 8B leva alguns minutos. Para implantações subsequentes, o servidor JetStream-PyTorch carrega os pesos do modelo do diretório de trabalho, levando menos tempo para carregar. Para pontos de verificação já existentes, verifique se o diretório de trabalho usa a mesma convenção de caminho.enable_model_warmup
: essa configuração ativa o aquecimento do modelo depois que o servidor do modelo é iniciado. O padrão éFalse
.
Você pode definir estas propriedades:
max_input_length
: o tamanho máximo da sequência de entrada. O padrão é1024
.max_output_length
: o comprimento máximo de decodificação de saída. Esse valor é padrão para1024
.quantize_weights
: se o checkpoint é quantizado. O padrão é0
. Defina como 1 para ativar a quantizaçãoint8
.internal_jax_compilation_cache
: o diretório do cache de compilação do JAX. O padrão é~/jax_cache
. Defina comogs://BUCKET_NAME/jax_cache
para armazenamento em cache remoto.
No manifesto, uma sondagem de inicialização
é configurada para garantir que o servidor de modelo seja rotulado como Ready
depois que o
modelo for carregado e o aquecimento for concluído. As sondagens de atividade e prontidão são configuradas para garantir a integridade do servidor de modelo.
Substitua
BUCKET_NAME
pelo GSBucket criado anteriormente:sed -i "s|BUCKET_NAME|BUCKET_NAME|g" jetstream-pytorch-deployment.yaml
Aplique o manifesto:
kubectl apply -f jetstream-pytorch-deployment.yaml
Verificar a implantação:
kubectl get deployment
O resultado será assim:
NAME READY UP-TO-DATE AVAILABLE AGE jetstream-pytorch-server 0/2 2 0 ##s
Para clusters do Autopilot, pode levar alguns minutos para provisionar os recursos de TPU necessários.
Acesse os registros do servidor JetStream-PyTorch para verificar se os pesos do modelo foram carregados e se o aquecimento do modelo foi concluído. O servidor pode levar alguns minutos para concluir essa operação.
kubectl logs deploy/jetstream-pytorch-server -f -c jetstream-pytorch-server
O resultado será assim:
Started jetstream_server.... 2024-04-12 04:33:37,128 - root - INFO - ---------Generate params 0 loaded.---------
Verifique se a implantação está pronta:
kubectl get deployment
O resultado será assim:
NAME READY UP-TO-DATE AVAILABLE AGE jetstream-pytorch-server 2/2 2 2 ##s
Pode levar alguns minutos para o endpoint
healthcheck
ser registrado.
Disponibilizar o modelo
Nesta seção, você vai interagir com o modelo.
Configurar o encaminhamento de portas
Acesse a implantação do JetStream pelo Serviço ClusterIP criado na etapa anterior. Os serviços ClusterIP só podem ser acessados dentro do cluster. Portanto, para acessar o serviço de fora do cluster, conclua as seguintes etapas:
Para estabelecer uma sessão de encaminhamento de portas, execute o seguinte comando:
kubectl port-forward svc/jetstream-svc 8000:8000
Interagir com o modelo usando curl
Para verificar se é possível acessar o servidor HTTP JetStream, abra um novo terminal e execute o seguinte comando:
curl --request POST \ --header "Content-type: application/json" \ -s \ localhost:8000/generate \ --data \ '{ "prompt": "What are the top 5 programming languages", "max_tokens": 200 }'
A solicitação inicial pode levar alguns segundos para ser concluída devido ao aquecimento do modelo. O resultado será assim:
{ "response": " for data science in 2023?\n\n**1. Python:**\n- Widely used for data science due to its readability, extensive libraries (pandas, scikit-learn), and integration with other tools.\n- High demand for Python programmers in data science roles.\n\n**2. R:**\n- Popular choice for data analysis and visualization, particularly in academia and research.\n- Extensive libraries for statistical modeling and data wrangling.\n\n**3. Java:**\n- Enterprise-grade platform for data science, with strong performance and scalability.\n- Widely used in data mining and big data analytics.\n\n**4. SQL:**\n- Essential for data querying and manipulation, especially in relational databases.\n- Used for data analysis and visualization in various industries.\n\n**5. Scala:**\n- Scalable and efficient for big data processing and machine learning models.\n- Popular in data science for its parallelism and integration with Spark and Spark MLlib." }
Você fez o seguinte:
- Implantei o servidor de modelo JetStream-PyTorch no GKE usando TPUs.
- Criação de um ponto de controle em
gs://BUCKET_NAME/pytorch/org/repo
. - Exibiu e interagiu com o modelo.
Resolver problemas
- Se você receber a mensagem
Empty reply from server
, é possível que o contêiner não tenha concluído o download dos dados do modelo. Verifique os registros do pod novamente para ver a mensagemConnected
, que indica que o modelo está pronto para ser disponibilizado. - Se você vir
Connection refused
, verifique se o encaminhamento de portas está ativo.
Limpar
Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.
Excluir os recursos implantados
Para evitar cobranças na sua conta do Google Cloud pelos recursos criados neste guia, execute os comandos a seguir e siga as instruções:
gcloud container clusters delete CLUSTER_NAME --region=REGION
gcloud iam service-accounts delete wi-jetstream@PROJECT_ID.iam.gserviceaccount.com
gcloud storage rm --recursive gs://BUCKET_NAME
A seguir
- Descubra como executar modelos Gemma no GKE e como executar cargas de trabalho otimizadas de IA/ML com os recursos de orquestração da plataforma do GKE.
- Saiba mais sobre TPUs no GKE.
- Explore o repositório do JetStream no GitHub (link em inglês).
- Conhecer o Model Garden da Vertex AI.