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

Standard

---

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 cresceram 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 de treinamentos no 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).

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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Enable the API

• Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin, roles/resourcemanager.projectIamAdmin

  Check for the roles

  1. In the Google Cloud console, go to the IAM page.

     Go to IAM
  2. Select the project.

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

  4. For all rows that specify or include you, check the Role column to see whether the list of roles includes the required roles.

  Grant the roles

  1. In the Google Cloud console, go to the IAM page.

     Acessar o IAM
  2. Selecione o projeto.
  3. Clique em Conceder acesso.

  4. No campo Novos principais, digite seu identificador de usuário. Normalmente, é o endereço de e-mail de uma Conta do Google.

  5. Na lista Selecionar papel, escolha um.
  6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
  7. Clique em Salvar.
  • Verifique se você tem cota suficiente para 16 chips TPU v6e PodSlice Lite. Neste tutorial, você usará as instâncias sob demanda.
  • Verifique se o projeto do Google Cloud está na lista de permissões do programa de aprendizado.

  Receber acesso ao modelo

  Para ter acesso ao ponto de verificação Meta Llama 3.1-405B para implantação no GKE, siga estas etapas:

  1. Assine o contrato de consentimento de licença.
  2. Acesse a página de downloads do Meta Llama.
  3. Leia e aceite os Termos e Condições do modelo e obtenha o URL necessário para fazer o download.
  4. Para fazer o download do checkpoint do modelo, encontre o ID do modelo adequado. Para conferir uma lista dos modelos compatíveis e os IDs deles, consulte a documentação da CLI llama. Por exemplo, use Llama 3.1-405B-Instruct:bf16-mp16 para o modelo Llama 3.1-405B.

  Prepare o ambiente

  Neste tutorial, você vai 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:

  1. No console do Google Cloud , inicie uma sessão do Cloud Shell clicando em Ativar o Cloud Shell no Google Cloud console. Isso inicia uma sessão no painel inferior do console Google Cloud .

  2. Defina as variáveis de ambiente padrão:

     gcloud config set project PROJECT_ID
     gcloud config set billing/quota_project PROJECT_ID
     export PROJECT_ID=$(gcloud config get project)
     export CLUSTER_NAME=CLUSTER_NAME
     export BUCKET_NAME=BUCKET_NAME
     export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
     export NODE_LOCATION=NODE_LOCATION
     export CLUSTER_VERSION=CLUSTER_VERSION
     export MACHINE_TYPE=ct6e-standard-4t
     export TPU_TYPE=v6e
     export TOPOLOGY=4x4
     export WORKERS_PER_SLICE=4
     

     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 prefixo gs://.
     • CONTROL_PLANE_LOCATION: a região do Compute Engine 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 v6e estão disponíveis (por exemplo, us-east1, us-east5, europe-west4, asia-northeast1 ou us-south1).
     • NODE_LOCATION: a zona em que os recursos de TPU estão disponíveis (por exemplo, us-east1-d).
     • 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.
     • MACHINE_TYPE: o tipo de máquina v6e.
     • TPU_TYPE: um prefixo usado para nomear pools de nós (v6e).
     • TOPOLOGY: a topologia da TPU v6e.
     • WORKERS_PER_SLICE: o número de nós por pool de nós ou fração de TPU.

  Criar e configurar recursos Google Cloud

  Para criar os recursos necessários, siga estas instruções:

  Observação: pode ser necessário criar uma reserva de capacidade para usar alguns aceleradores. Para mais informações sobre como reservar e consumir recursos reservados, consulte Como consumir recursos por zona reservados.

  Criar um cluster do GKE

  1. Crie um cluster regional do GKE Standard:

     gcloud container clusters create CLUSTER_NAME \
         --project=PROJECT_ID \
         --cluster-version=CLUSTER_VERSION \
         --location=CONTROL_PLANE_LOCATION \
         --scopes=cloud-platform \
         --machine-type=n2-standard-32
     

     A criação do cluster pode levar vários minutos.

     Substitua CLUSTER_VERSION pela versão do cluster adequada.

  2. Crie um pool de nós de TPU v6e com uma topologia 4x4 e quatro nós cada:

     gcloud container node-pools create multihost-np \
     --project=PROJECT_ID \
     --location=CONTROL_PLANE_LOCATION \
     --node-locations=NODE_LOCATION \
     --cluster=CLUSTER_NAME \
     --machine-type=MACHINE_TYPE \
     --num-nodes=WORKERS_PER_SLICE \
     --tpu-topology=TOPOLOGY \
     --scopes cloud-platform \
     --placement-type=COMPACT \
     --workload-metadata=GCE_METADATA
     

     Observação: o --workload-metadata=GCE_METADATA indica que o pool de nós tem a federação de identidade da carga de trabalho para o GKE desativada. Os pools de nós que executam os programas do Pathways com JetStream e MaxText precisam desativar a federação de identidade da carga de trabalho para o GKE.

  Configurar uma conta de serviço para acesso a objetos do Storage

  Configure uma conta de serviço do Kubernetes para atuar como uma conta de serviço do IAM.

  1. Crie uma conta de serviço do IAM para seu aplicativo:

     gcloud iam service-accounts create jetstream-pathways
     

  2. Adicione uma vinculação de política do IAM à sua conta de serviço do IAM para gerenciar o Cloud Storage. Isso permite que sua conta de serviço do IAM acesse o bucket de armazenamento em que o ponto de verificação será armazenado:

     gcloud projects add-iam-policy-binding ${PROJECT} \
       --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
       --role roles/storage.objectUser
     
     gcloud projects add-iam-policy-binding ${PROJECT} \
       --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
       --role roles/storage.insightsCollectorService
     

  3. 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=jetstream-pathways@${PROJECT}.iam.gserviceaccount.com
     

  Configurar o Docker para autenticação no Artifact Registry

  Configure o Docker para autenticar no Artifact Registry e extrair as imagens dos programas de aprendizado permitidos:

  gcloud auth login
  gcloud auth configure-docker
  

  Conversão de checkpoint

  Para converter um checkpoint do Meta Llama 3.1-405B em um checkpoint de inferência int8 compatível com MaxText, siga as etapas em Conversão de checkpoint com Llama3.1-405B. Sua implantação usa o ponto de verificação com a flag load_parameters_path.

  Crie um bucket do Cloud Storage para armazenar arquivos temporários dos programas

  Crie um bucket do Cloud Storage para armazenar seus arquivos temporários dos programas, como o cache de compilação:

  export PATHWAYS_BUCKET=PATHWAYS_BUCKET
  gcloud storage buckets create gs://$PATHWAYS_BUCKET
  

  Implantar o JetStream-MaxText e o Pathways

  Implante o servidor de modelos JetStream-MaxText e Pathways.

  Conecte-se ao cluster do GKE

  gcloud container clusters get-credentials "${CLUSTER}" --project "${PROJECT}" --location "${ZONE}"
  

  Implantar a API LeaderWorkerSet (LWS)

  O LWS é um recurso personalizado projetado para implantar e gerenciar aplicativos distribuídos com estado, principalmente aqueles com uma arquitetura líder-trabalhador. Ele é especialmente adequado para cargas de trabalho de IA/ML em que um modelo grande é fragmentado e veiculado em vários dispositivos em vários nós.

  VERSION=v0.6.1
  kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml
  

  Aguarde até que o controlador LeaderWorkerSet esteja totalmente disponível:

  kubectl wait deploy/lws-controller-manager -n lws-system --for=condition=available --timeout=5m
  

  A saída será semelhante a esta:

  deployment.apps/lws-controller-manager condition met
  

  Valide se o controlador LeaderWorkerSet está em execução no namespace lws-system:

  kubectl get pod -n lws-system
  

  A saída será semelhante a esta:

  NAME                          READY   STATUS    RESTARTS    AGE
  lws-controller-manager-abcd   1/1     Running   0           40s
  lws-controller-manager-efgh   1/1     Running   0           40s
  

  Implantar o manifesto da carga de trabalho

  1. Salve o seguinte manifesto como jetstream-pathways-llama-3-1-405b-4x4.yaml:

     apiVersion: leaderworkerset.x-k8s.io/v1
     kind: LeaderWorkerSet
     metadata:
       name: jetstream-pathways
       annotations:
         leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
     spec:
       replicas: 1
       leaderWorkerTemplate:
         leaderTemplate:
           metadata:
             labels:
               app: jetstream-pathways
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: pathways-proxy
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
               args:
               imagePullPolicy: Always
               ports:
               - containerPort: 38681
             - name: pathways-rm
               env:
               - name: HOST_ADDRESS
                 value: "$(LWS_LEADER_ADDRESS)"
               - name: TPU_SKIP_MDS_QUERY
                 value: "true"
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               args:
               - --server_port=38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               - --node_type=resource_manager
               - --instance_count=1
               - --instance_type=tpuv6e:4x4
               imagePullPolicy: Always
               ports:
               - containerPort: 38677
             - name: jax-tpu
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
               env:
               - name: LOG_LEVEL
                 value: "INFO"
               args:
               - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
               - model_name=llama3.1-405b
               - load_parameters_path=CHECKPOINT_PATH
               - max_prefill_predict_length=1024
               - max_target_length=2048
               - async_checkpointing=false
               - steps=1
               - ici_fsdp_parallelism=1
               - ici_autoregressive_parallelism=2
               - ici_tensor_parallelism=8
               - scan_layers=false
               - weight_dtype=bfloat16
               - per_device_batch_size=6
               - enable_single_controller=true
               - quantization=int8
               - quantize_kvcache=true
               - checkpoint_is_quantized=true
               - enable_model_warmup=true
               imagePullPolicy: Always
               ports:
               - containerPort: 9000
               startupProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 1
                 initialDelaySeconds: 600
                 failureThreshold: 10000
               livenessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 10
               readinessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 10
             - name: jetstream-http
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
               imagePullPolicy: Always
               ports:
               - containerPort: 8000
         size: 5
         workerTemplate:
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 4x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: worker
               args:
               - --server_port=38679
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               imagePullPolicy: Always
               ports:
               - containerPort: 38679
               resources:
                 limits:
                   google.com/tpu: "4"
     --- 
     apiVersion: v1
     kind: Service
     metadata:
       name: jetstream-svc
     spec:
       selector:
         app: jetstream-pathways
       ports:
       - protocol: TCP
         name: jetstream-http
         port: 8000
         targetPort: 8000

  2. Defina o valor do campo load_parameters_path como o caminho do ponto de verificação produzido no processo de conversão.

     • Para um ponto de verificação bf16, o caminho precisa ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
     • Para um ponto de verificação int8, ele precisa ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/int8.

     Defina o valor do campo gcs_scratch_location como o bucket do programa de aprendizado que você criou antes.

     perl -pi -e 's|CHECKPOINT_PATH|gs://OUTPUT_BUCKET_DIRECTORY/int8|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
     perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
     

     Descrição de variáveis e configurações.

     As descrições a seguir explicam as variáveis e configurações relevantes para JetStream, Pathways e LeaderWorkerSet.

     • Jetstream
       • Esse contêiner foi criado com base no projeto OSS JetStream. Ele também encapsula a biblioteca de cliente baseada em gRPC do proxy IFRT, que fornece um back-end JAX capaz de delegar a execução de programas de TPU ao servidor proxy do Pathways.
       • Para mais informações sobre flags do JetStream-MaxText, consulte as Descrições das flags do servidor JetStream-MaxText.
       • Para usar um ponto de verificação quantizado, defina as flags quantization, quantize_kvcache ou checkpoint_is_quantized. Se o ponto de verificação for bf16, você poderá omitir essas flags.
     • Programas: para uma descrição dos componentes relevantes dos programas, consulte Componentes dos programas.
     • LeaderWorkerSet: para mais informações sobre as configurações do LeaderWorkerSet, consulte a definição da API.

  Aplique o manifesto de implantação.

  Aplique o manifesto para implantar o servidor:

  kubectl apply -f jetstream-pathways-llama-3-1-405b-4x4.yaml
  

  O servidor de modelo será iniciado.

  Verificar a inicialização do servidor de modelo

  Um modelo de 405B pode levar de 10 a 20 minutos para restaurar o checkpoint. Você também pode aguardar mais tempo durante o aquecimento do modelo se tiver ativado a flag enable_model_warmup.

  kubectl logs -f jetstream-pathways-0 -c jax-tpu
  

  O resultado será assim:

  2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
  2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
  2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
  2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
  2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
  2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
  ...
  ...
  ...
  INFO:     Started server process [7]
  INFO:     Waiting for application startup.
  INFO:     Application startup complete.
  INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)
  

  Disponibilizar o Llama 3.1-405b

  Para disponibilizar o modelo Llama 3.1-405b, configure o encaminhamento de portas:

  kubectl port-forward svc/jetstream-svc 8000:8000
  

  Com o encaminhamento de portas, é possível acessar o serviço de fora do cluster. É possível acessar a implantação do JetStream-Pathways pelo serviço ClusterIP do GKE. Os serviços ClusterIP só podem ser acessados dentro do cluster.

  interagir com o modelo

  Em um novo terminal, 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. A saída será semelhante a esta:

  {
      "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
  }
  

  Você fez o seguinte:

  1. Implantamos o servidor de modelos JetStream com MaxText e Pathways no GKE usando TPUs.
  2. Criamos um ponto de verificação Llama 3.1-405B int8 em gs://BUCKET_NAME.
  3. Serviu e interagiu com o modelo.

  Veiculação desagregada

  A veiculação desagregada é uma técnica para veicular LLMs dividindo as etapas de pré-preenchimento e decodificação em diferentes hosts. Essa abordagem otimiza a utilização de recursos, o que pode melhorar a taxa de transferência e a latência.

  • Pré-preenchimento: uma transmissão direta no comando de entrada para inicializar o cache de chave-valor.

  • Decodificação: um procedimento que gera tokens de saída de forma incremental, um token por etapa e um valor de cache KV por iteração.

  1. Defina as variáveis de ambiente padrão:

     export NODE_POOL_NAME=dis-v6e-8
     export NODE_POOL_SIZE=2
     export MACHINE_TYPE=ct6e-standard-4t
     export TOPOLOGY=2x4
     export WORKERS_PER_SLICE=2
     

  2. Crie dois pools de nós que usam nós v6e-8:

     for i in $(seq 1 NODE_POOL_SIZE); do
       gcloud container node-pools create NODE_POOL_NAME-${i}-np \
       --project=PROJECT \
       --cluster=CLUSTER_NAME \
       --location=CONTROL_PLANE_LOCATION \
       --node-locations=NODE_LOCATION \
       --machine-type=MACHINE_TYPE \
       --num-nodes=WORKERS_PER_SLICE \
       --tpu-topology=TOPOLOGY \
       --scopes=cloud-platform \
       --workload-metadata=GCE_METADATA
     done
     

  Conversão de checkpoint

  Para converter um checkpoint do Meta Llama 2-70B em um checkpoint de inferência int8 compatível com MaxText, siga as etapas em Conversão de checkpoint com Llama2-70B. Selecione Llama2-70B como seu modelo ao confirmar os Termos e Condições da Meta. Sua implantação usa o ponto de verificação com a flag load_parameters_path.

  Substitua os seguintes parâmetros no arquivo checkpoint-job.yaml:

  - --meta_url=META_URL
  - --model_name=llama-2
  - --model_path=Llama-2-70b-chat
  - --output_directory=gs://BUCKET_NAME/maxtext/llama-2-70b
  

  O ponto de verificação será usado na sua implantação com a flag load_parameters_path.

  Implantar os programas do JetStream com exibição desagregada

  1. Salve o seguinte manifesto como jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml:

     apiVersion: leaderworkerset.x-k8s.io/v1
     kind: LeaderWorkerSet
     metadata:
       name: jetstream-pathways
       annotations:
         leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
     spec:
       replicas: 1
       leaderWorkerTemplate:
         subGroupPolicy:
           subGroupSize: 2
         leaderTemplate:
           metadata:
             labels:
               app: jetstream-pathways
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 2x4
             tolerations:
             - key: "google.com/tpu"
               operator: "Exists"
               effect: "NoSchedule"
             containers:
             - name: pathways-proxy
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
               args:
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --server_port=38681
               - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
               - --xla_jf_auto_cross_replica_sharding=false
               - --xla_tpu_enable_windowed_einsum_for_reduce_scatter=false
               - --xla_tpu_enable_windowed_einsum_for_all_gather=false
               - --xla_tpu_prefer_latch_optimized_rhs_layouts=true
               - --xla_tpu_enable_experimental_fusion_cost_model=false
               - --xla_tpu_dot_dot_fusion_duplicated=false
               - --xla_tpu_dot_dot_fusion=true
               - --xla_jf_conv_input_fusion=true
               - --xla_jf_conv_output_fusion=true
               - --xla_tpu_rwb_fusion=false
               - --xla_tpu_copy_fusion_pad_unpad_ratio=0
               - --xla_tpu_licm_size_inflation_ratio=1
               - --xla_tpu_copy_elision_analysis_allowance=150000
               - --xla_tpu_copy_insertion_use_region_analysis_limit=10000
               - --xla_tpu_order_dot_after_layout=true
               - --xla_jf_rematerialization_percent_shared_memory_limit=100
               - --xla_tpu_use_repeated_instance_for_preferred_prefetch_time=true
               - --xla_tpu_enforce_prefetch_fifo_order=false
               - --xla_tpu_prefetch_interval_picker_size_override=6000000
               - --xla_tpu_async_copy_bandwidth_scaling_factor=1
               - --xla_tpu_nd_short_transfer_max_chunks=-1
               - --xla_tpu_enable_aggressive_broadcast_priority_update=true
               - --xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers=SQRT
               - --xla_tpu_memory_bound_loop_optimizer_options=enabled:true
               - --xla_tpu_enable_copy_fusion=true
               - --xla_tpu_enable_cross_program_prefetch_freeing=false
               - --xla_tpu_enable_dot_strength_reduction=true
               - --xla_tpu_layout_use_dot_grouping=false
               - --xla_tpu_msa_inefficient_use_to_copy_ratio=0.5
               - --xla_tpu_reduce_loop_fusion_dup_with_unfusable_user=false
               - --xla_tpu_vector_load_fusion_window=1024
               - --xla_tpu_vector_store_fusion_window=256
               - --xla_jf_conv_reshape_fusion=false
               - --xla_tpu_input_conv_multi_users=false
               - --xla_tpu_enable_multi_level_input_dot_dot_fusion=false
               - --xla_tpu_enable_multi_level_output_dot_dot_fusion=false
               - --xla_tpu_dot_dot_fusion_separable_convs_only=false
               - --xla_tpu_enable_multi_level_nested_loop_fusion=true
               - --xla_tpu_nested_dot_fusion=true
               - --xla_tpu_enable_multi_level_nested_dot_fusion=false
               - --xla_jf_enable_multi_output_fusion=true
               - --xla_tpu_use_lp_llo_scheduler_for_dot_dot_fusions=false
               - --xla_tpu_enable_flash_attention=true
               imagePullPolicy: Always
               ports:
               - containerPort: 38681
             - name: pathways-rm
               env:       
               - name: HOST_ADDRESS
                 value: "$(LWS_LEADER_ADDRESS)"
               - name: TPU_SKIP_MDS_QUERY
                 value: "true"
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               args:
               - --server_port=38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               - --node_type=resource_manager
               - --instance_count=2
               - --instance_type=tpuv6e:2x4
               imagePullPolicy: Always
               ports:
               - containerPort: 38677
             - name: jax-tpu
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
               args:
               - MaxText/configs/base.yml
               - tokenizer_path=assets/tokenizer.llama2
               - load_parameters_path=CHECKPOINT_PATH
               - max_prefill_predict_length=1024
               - max_target_length=2048
               - model_name=llama2-70b
               - ici_fsdp_parallelism=1
               - ici_autoregressive_parallelism=1
               - ici_tensor_parallelism=-1
               - scan_layers=false
               - weight_dtype=bfloat16
               - per_device_batch_size=27
               - checkpoint_is_quantized=true 
               - quantization=int8
               - quantize_kvcache=true
               - compute_axis_order=0,2,1,3
               - ar_cache_axis_order=0,2,1,3
               - stack_prefill_result_cache=True
               - inference_server=ExperimentalMaxtextDisaggregatedServer_8
               - inference_benchmark_test=True
               - enable_model_warmup=True
               env:
               - name: LOG_LEVEL
                 value: "INFO"
               imagePullPolicy: Always
               securityContext:
                 capabilities:
                   add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
               ports: 
               - containerPort: 9000
               startupProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 1
                 initialDelaySeconds: 240
                 failureThreshold: 10000
               livenessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 100
               readinessProbe:
                 httpGet:
                   path: /healthcheck
                   port: 8000
                   scheme: HTTP
                 periodSeconds: 60
                 failureThreshold: 100
             - name: jetstream-http
               image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
               imagePullPolicy: Always
               ports:
               - containerPort: 8000
         size: 5
         workerTemplate:
           spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
               cloud.google.com/gke-tpu-topology: 2x4
             containers:
             - name: worker
               args:
               - --server_port=38679
               - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
               - --gcs_scratch_location=PATHWAYS_BUCKET
               image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
               imagePullPolicy: Always
               ports:
               - containerPort: 38679
               resources:
                 limits:
                   google.com/tpu: "4"
     --- 
     apiVersion: v1
     kind: Service
     metadata:
       name: jetstream-svc
     spec:
       selector:
         app: jetstream-pathways
       ports:
       - protocol: TCP
         name: jetstream-http
         port: 8000
         targetPort: 8000

  2. Defina o valor do campo load_parameters_path como o caminho do ponto de verificação produzido no processo de conversão.

     • Para um ponto de verificação bf16, o caminho precisa ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
     • Para um ponto de verificação int8, ele precisa ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/int8.

     Defina o valor do campo gcs_scratch_location como o bucket do programa de aprendizado que você criou antes.

     perl -pi -e 's|CHECKPOINT_PATH|BUCKET_NAME/maxtext/llama-2-70b/int8|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     

  3. Aplique o manifesto:

     kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
     

     O servidor de modelo pode levar algum tempo para restaurar o ponto de verificação, dependendo do tamanho dele. Um modelo de 70 bilhões pode levar cerca de 8 minutos para restaurar o ponto de verificação, incluindo o aquecimento do modelo. Você pode observar ainda mais os registros para determinar o ponto de prontidão verificando o início do servidor de modelo e disponibilizar o modelo configurando o encaminhamento de porta para interagir com ele.

  Você fez o seguinte:

  1. Implantamos o servidor de modelos JetStream com MaxText e Pathways no GKE usando TPUs e disponibilização desagregada.
  2. Criou um checkpoint Llama 2-70B int8 em gs://BUCKET_NAME.
  3. Serviu 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 mensagem Connected, que indica que o modelo está pronto para ser disponibilizado.
  • Se você vir a mensagem 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 --location=CONTROL_PLANE_LOCATION
  
  gcloud iam service-accounts delete jetstream-pathways@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.