Esta página foi traduzida pela API Cloud Translation.

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

Padrão

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, o 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.

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:

Saiba mais sobre a disponibilidade da versão atual da TPU com a arquitetura do sistema Cloud TPU.
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 na nuvem

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 TPUs, simplificando as computações de ML que abrangem centenas ou milhares de chips de TPUs.

JetStream

O JetStream é uma framework de fornecimento 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 preparação e publicação distribuídas 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:

Arquitetura do cluster do GKE com um conjunto de nós da TPU de vários anfitriões que contém os componentes JetStream e Pathways.

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.

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.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the API

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

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

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

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 column to see whether the list of roles includes the required roles.

Grant the roles

In the Google Cloud console, go to the IAM page.
Aceder ao IAM
Selecione o projeto.
Clique em Conceder acesso.
No campo Novos responsáveis, introduza o identificador do utilizador. Normalmente, este é o endereço de email de uma Conta Google.
Na lista Selecionar uma função, selecione uma função.
Para conceder funções adicionais, clique em Adicionar outra função e adicione cada função adicional.
Clique em Guardar.

Certifique-se de que tem quota suficiente para dezasseis chips TPU v6e PodSlice Lite. Neste tutorial, vai usar instâncias a pedido.
Certifique-se de que o seu Google Cloud projeto está na lista de autorizações dos Pathways.

Aceda ao modelo

Para aceder ao ponto de verificação do Meta Llama 3.1-405B para implementação no GKE, siga estes passos:

Assine o contrato de consentimento de licença.
Aceda à página de transferências do Meta Llama.
Reveja e aceite os Termos de Utilização do modelo e obtenha o URL necessário para transferir o modelo.
Para transferir o ponto de verificação do modelo, encontre o ID do modelo adequado. Para ver uma lista dos modelos suportados e os respetivos IDs, 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, vai usar o Cloud Shell para gerir recursos alojados no Google Cloud. O Cloud Shell é pré-instalado com o software necessário para este tutorial, incluindo kubectl e a CLI gcloud.

Para configurar o seu ambiente com o Cloud Shell, siga estes passos:

Na Google Cloud consola, inicie uma sessão do Cloud Shell clicando em Ativar Cloud Shell na Google Cloud consola. Esta ação inicia uma sessão no painel inferior da consola. Google Cloud
Defina as variáveis de ambiente predefinidas:
```
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 seu Google Cloud ID do projeto.
- CLUSTER_NAME: o nome do seu cluster do GKE.
- BUCKET_NAME: o nome do seu contentor do Cloud Storage. Não precisa de especificar o prefixo gs://.
- CONTROL_PLANE_LOCATION: a região do Compute Engine onde se encontram o cluster do GKE, o contentor do Cloud Storage e os nós da TPU. A região contém zonas onde os tipos de máquinas TPU v6e estão disponíveis (por exemplo, us-east1, us-east5, europe-west4, asia-northeast1 ou us-south1).
- NODE_LOCATION: a zona onde os recursos da TPU estão disponíveis (por exemplo, us-east1-d).
- CLUSTER_VERSION: a versão do GKE, que tem de suportar o tipo de máquina que quer usar. Tenha em atenção que a versão predefinida do GKE pode não ter disponibilidade para a TPU de destino. Para ver uma lista das versões mínimas do GKE disponíveis por tipo de máquina TPU, consulte o artigo Disponibilidade de TPUs no GKE.
- MACHINE_TYPE: o tipo de máquina v6e.
- TPU_TYPE: um prefixo usado para a nomenclatura de pools de nós (v6e).
- TOPOLOGY: a topologia da TPU v6e.
- WORKERS_PER_SLICE: o número de nós por node pool ou fatia da TPU.

Crie e configure Google Cloud recursos

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

Crie um cluster do GKE

Crie um cluster padrão do GKE regional:

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 demorar vários minutos.

Substitua CLUSTER_VERSION pela versão do cluster adequada.

Crie um node pool da 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

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

Crie uma conta de serviço do IAM para a sua aplicação:

gcloud iam service-accounts create jetstream-pathways

Adicione uma associação de política de IAM para a sua conta de serviço de IAM para gerir o Cloud Storage. Isto destina-se a permitir que a sua conta de serviço de IAM aceda ao contentor de armazenamento onde o ponto de verificação vai 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

Anote a conta de serviço do Kubernetes com o endereço de email da conta de serviço do IAM.

kubectl annotate serviceaccount default \
iam.gke.io/gcp-service-account=jetstream-pathways@${PROJECT}.iam.gserviceaccount.com

Configure o Docker para autenticar no Artifact Registry

Configure o Docker para autenticar no Artifact Registry, de modo que possa extrair as imagens do Pathways na lista de autorizações:

gcloud auth login
gcloud auth configure-docker

Conversão de posto de controlo

Para converter um ponto de verificação do Meta Llama 3.1-405B num ponto de verificação de inferência int8 compatível com o MaxText, conclua os passos em Conversão de pontos de verificação com o Llama3.1-405B. A sua implementação usa o ponto de verificação com a flag load_parameters_path.

Crie um contentor do Cloud Storage para armazenar ficheiros temporários do Pathways

Crie um contentor do Cloud Storage para armazenar os ficheiros temporários dos Pathways, como a cache de compilação:

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

Implemente o JetStream-MaxText e o Pathways

Implemente o servidor do modelo JetStream-MaxText e Pathways.

Ligue-se ao cluster do GKE

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

Implemente a API LeaderWorkerSet (LWS)

O LWS é um recurso personalizado concebido para implementar e gerir aplicações distribuídas com estado, particularmente as que têm uma arquitetura de líder-trabalhador. É especialmente adequado para cargas de trabalho de IA/ML em que um modelo grande é dividido e apresentado 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

O resultado deve ser semelhante ao seguinte:

deployment.apps/lws-controller-manager condition met

Valide se o controlador LeaderWorkerSet está a ser executado no espaço de nomes lws-system:

kubectl get pod -n lws-system

O resultado deve ser semelhante ao seguinte:

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

Implemente o manifesto da carga de trabalho

Guarde 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

Defina o valor do campo load_parameters_path para o caminho do ponto de verificação produzido no processo de conversão do ponto de verificação.
- Para um ponto de verificação bf16, o caminho deve ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
- Para um ponto de verificação int8, deve ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/int8.
Defina o valor do campo gcs_scratch_location para o grupo de caminhos que criou anteriormente.
```
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 das variáveis e das definições.

As descrições seguintes explicam as variáveis e as definições relevantes para o JetStream, o Pathways e o LeaderWorkerSet.
- Jetstream
  - Este contentor é criado a partir do projeto OSS JetStream. Também encapsula a biblioteca cliente baseada em gRPC do IFRT Proxy, que fornece um back-end JAX que pode delegar a execução de programas de TPU para o servidor Pathways Proxy.
  - Para mais informações acerca das flags JetStream-MaxText, consulte as descrições das flags do servidor JetStream-MaxText.
  - Para usar um ponto de verificação quantizado, pode definir as flags quantization, quantize_kvcache ou checkpoint_is_quantized. Se o ponto de verificação for bf16, pode omitir estas flags.
- Caminhos: para uma descrição dos componentes relevantes dos caminhos, consulte Componentes dos caminhos.
- LeaderWorkerSet: para mais informações acerca das definições de LeaderWorkerSet, consulte a definição da API.

Aplique o manifesto de implementação

Aplique o manifesto para implementar o servidor:

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

O servidor de modelos deve ser iniciado.

Valide o início do servidor do modelo

Um modelo 405B pode demorar cerca de 10 a 20 minutos a restaurar o ponto de verificação. 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 é semelhante ao seguinte:

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)

Servir Llama 3.1-405b

Para publicar o modelo Llama 3.1-405b, configure o encaminhamento de porta:

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

O encaminhamento de portas permite-lhe aceder ao serviço a partir de fora do cluster. Pode aceder à implementação do JetStream-Pathways através do serviço ClusterIP do GKE. Os serviços ClusterIP só são acessíveis a partir do cluster.

Interagir com o modelo

Num 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
}'

O pedido inicial pode demorar vários segundos a ser concluído devido ao aquecimento do modelo. O resultado deve ser semelhante ao seguinte:

{
    "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"
}

Concluiu com êxito o seguinte:

Implementou o servidor do modelo JetStream com o MaxText e o Pathways no GKE através de TPUs.
Foi criado um ponto de restauro do Llama 3.1-405B int8 em gs://BUCKET_NAME.
Foram apresentadas e houve interação com o modelo.

Publicação desagregada

A publicação desagregada é uma técnica para publicar LLMs dividindo as fases de preenchimento prévio e descodificação em diferentes anfitriões. Esta abordagem otimiza a utilização de recursos, o que pode levar a um débito e uma latência melhorados.

Preenchimento prévio: uma passagem direta no comando de entrada para inicializar a cache de chaves-valores.
Descodificar: um procedimento que gera tokens de saída de forma incremental, um token por passo e um valor de KV-cache por iteração.

Defina as variáveis de ambiente predefinidas:

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

Crie dois node pools que usem 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 posto de controlo

Para converter um ponto de verificação do Meta Llama 2-70B num ponto de verificação de inferência int8 compatível com o MaxText, conclua os passos em Conversão de pontos de verificação com o Llama2-70B. Selecione o Llama2-70B como modelo quando confirmar os Termos de Utilização da Meta. A sua implementação usa o ponto de verificação com a flag load_parameters_path.

Substitua os seguintes parâmetros no ficheiro 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 vai ser usado na sua implementação com a flag load_parameters_path.

Implemente caminhos do JetStream com o envio desagregado

Guarde 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

Defina o valor do campo load_parameters_path para o caminho do ponto de verificação produzido no processo de conversão do ponto de verificação.
- Para um ponto de verificação bf16, o caminho deve ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
- Para um ponto de verificação int8, deve ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/int8.
Defina o valor do campo gcs_scratch_location para o grupo de caminhos que criou anteriormente.
```
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
```
Aplique o manifesto:
```
kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
```
O servidor do modelo pode demorar algum tempo a restaurar o ponto de verificação, consoante o tamanho do ponto de verificação. Um modelo de 70 mil milhões de parâmetros pode demorar cerca de 8 minutos a restaurar o ponto de verificação, incluindo o aquecimento do modelo. Pode observar ainda mais os registos para determinar o ponto de prontidão verificando o início do servidor do modelo e disponibilizar o modelo configurando o encaminhamento de portas para poder interagir com o modelo.

Concluiu com êxito o seguinte:

Implementou o servidor do modelo JetStream com o MaxText e o Pathways no GKE através de TPUs e da publicação desagregada.
Criou um posto de controlo Llama 2-70B int8 em gs://BUCKET_NAME.
Foram apresentadas e houve interação com o modelo.

Resolver problemas

Se receber a mensagem Empty reply from server, é possível que o contentor não tenha terminado de transferir os dados do modelo. Verifique novamente os registos do Pod para ver a mensagem Connected, que indica que o modelo está pronto para publicação.
Se vir a mensagem Connection refused, verifique se o encaminhamento de portas está ativo.

Limpar

Para evitar incorrer em custos na sua conta do Google Cloud pelos recursos usados neste tutorial, elimine o projeto que contém os recursos ou mantenha o projeto e elimine os recursos individuais.

Elimine os recursos implementados

Para evitar incorrer em custos na sua Google Cloud conta pelos recursos que criou neste guia, execute os seguintes comandos 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

O que se segue?

Descubra como pode executar modelos Gemma no GKE e como executar cargas de trabalho de IA/ML otimizadas com as capacidades de orquestração da plataforma GKE.
Saiba mais acerca das TPUs no GKE.
Explore o repositório do GitHub do JetStream.
Explore o Vertex AI Model Garden.

Publique GMLs com TPUs multi-host no GKE com o JetStream e o Pathways Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

Contexto

Acerca das TPUs

Percursos na nuvem

JetStream

MaxText

Llama 3.1 405B

Arquitetura

Antes de começar

Check for the roles

Grant the roles

Aceda ao modelo

Prepare o ambiente

Crie e configure Google Cloud recursos

Crie um cluster do GKE

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

Configure o Docker para autenticar no Artifact Registry

Conversão de posto de controlo

Crie um contentor do Cloud Storage para armazenar ficheiros temporários do Pathways

Implemente o JetStream-MaxText e o Pathways

Ligue-se ao cluster do GKE

Implemente a API LeaderWorkerSet (LWS)

Implemente o manifesto da carga de trabalho

Aplique o manifesto de implementação

Valide o início do servidor do modelo

Servir Llama 3.1-405b

Interagir com o modelo

Publicação desagregada

Conversão de posto de controlo

Implemente caminhos do JetStream com o envio desagregado

Resolver problemas

Limpar

Elimine os recursos implementados

O que se segue?

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