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

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

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

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 o Pathways on Cloud para executar inferência de vários hosts em grande escala 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 solicitações com baixa latência.

Programas do Pathways on Cloud

Os Pathways são uma camada de orquestração em grande escala para aceleradores. Os programas de treinamentos são projetados explicitamente para permitir a análise detalhada de novos sistemas e ideias de pesquisa de ML, mantendo o desempenho de ponta dos modelos atuais. Com o Pathways, um único processo cliente do JAX coordena 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 oferecer suporte à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:

Arquitetura do cluster do GKE com pool de nós de TPU de vários hosts que contém os componentes JetStream e Pathways.

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

  • Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloud, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

  • 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 role (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 role (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

    •

  • Verifique se você tem os seguintes papéis no projeto: roles/container.admin, roles/iam.serviceAccountAdmin, roles/resourcemanager.projectIamAdmin

    Verificar os papéis

    1. No console do Google Cloud , acesse a página IAM.

      Acessar IAM
    2. Selecione o projeto.

    3. Na coluna Principal, encontre todas as linhas que identificam você ou um grupo no qual você está incluído. Para saber em quais grupos você está incluído, entre em contato com o administrador.

    4. Em todas as linhas que especificam ou incluem você, verifique a coluna Papel para ver se a lista de papéis inclui os papéis necessários.

    Conceder os papéis

    1. No console do Google Cloud , acesse a página IAM.

      Acessar 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. Clique em Selecionar um papel e pesquise o papel.
    6. Para conceder outros papéis, adicione-os clicando em Adicionar outro papel.
    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.

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 de modelos compatíveis e os IDs deles, consulte a documentação da CLI do 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 Ícone de ativação do Cloud Shell Ativar o Cloud Shell no console doGoogle Cloud . 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:

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

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 checkpoint será armazenado:

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

gcloud projects add-iam-policy-binding ${PROJECT_ID} \
  --member "serviceAccount:jetstream-pathways@${PROJECT_ID}.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_ID}.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 treinamentos permitidas:

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 checkpoint com a flag load_parameters_path.

Criar um bucket do Cloud Storage para armazenar arquivos temporários dos programas de treinamento

Crie um bucket do Cloud Storage para armazenar seus arquivos temporários dos programas de treinamentos, 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 do JetStream-MaxText e do Pathways.

Conecte-se ao cluster do GKE

gcloud container clusters get-credentials "${CLUSTER_NAME}" --project "${PROJECT_ID}"  \
    --location "${CONTROL_PLANE_LOCATION}"

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 checkpoint bf16, o caminho precisa ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
    • Para um checkpoint 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://${BUCKET_NAME}/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

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á o 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)

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

O encaminhamento de portas permite 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 checkpoint int8 do Llama 3.1-405B em gs://BUCKET_NAME.
  3. Disponibilizar e interagir 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 capacidade de processamento e a latência.

  • Pré-preenchimento: uma passagem 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_ID} \
  --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 aceitar os Termos e Condições da Meta. Sua implantação usa o checkpoint 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 checkpoint bf16, o caminho precisa ser semelhante a gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
    • Para um checkpoint 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://${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 checkpoint, dependendo do tamanho dele. Um modelo de 70 bilhões de parâmetros pode levar cerca de 8 minutos para restaurar o checkpoint, 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. Disponibilizar e interagir 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