Disponibilizar o Gemma usando TPUs no GKE com o JetStream

Este guia mostra como exibir umaGema um modelo de linguagem grande (LLM) usando Unidades de Processamento de Tensor (TPUs) no Google Kubernetes Engine (GKE) comJetStream através deMaxText. Neste guia, faça o download dos pesos do modelo ajustado por instrução de parâmetro do Gemma 7B no Cloud Storage e implante-os em um Autopilot ou Standard usando um contêiner que executa o JetStream.

Este guia é um bom ponto de partida se você precisar da escalonabilidade, resiliência e economia oferecidos pelos recursos do Kubernetes ao implantar o modelo no JetStream.

Contexto

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

Gemma

O Gemma é um conjunto de modelos de inteligência artificial (IA) generativa, leve e abertamente lançados sob licença aberta. Esses modelos de IA estão disponíveis para execução em aplicativos, hardware, dispositivos móveis ou serviços hospedados. É possível usar os modelos Gemma para geração de texto, mas também é possível ajustá-los para tarefas especializadas.

Para saber mais, consulte a documentação do Gemma.

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 Gemma 7B. O GKE implanta o modelo em nós TPUv5e de host único com topologias de TPU configuradas com base nos requisitos do modelo para exibir prompts com baixa latência.

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

Para saber mais sobre essas otimizações, consulte os repositórios de projetos JetStream PyTorch e JetStream MaxText (links em inglês).

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

Objetivos

Este tutorial é destinado a clientes de IA generativa que usam o JAX, usuários novos ou atuais do GKE, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas interessados em usar os recursos de orquestração de contêineres do Kubernetes para veiculação de LLMs.

Este tutorial abrange as etapas a seguir:

  1. Prepare um cluster do Autopilot ou do GKE Standard com a topologia de TPU recomendada com base nas características do modelo.
  2. Implantar componentes do JetStream no GKE.
  3. Conseguir e publicar o modelo ajustado por instruções do Gemma 7B.
  4. Disponibilizar e interagir com o modelo publicado.

Arquitetura

Nesta seção, descrevemos a arquitetura do GKE usada neste tutorial. A arquitetura compreende um cluster do Autopilot ou padrão do GKE que provisiona TPUs e hospeda componentes do JetStream para implantar e disponibilizar os modelos.

O diagrama a seguir mostra os componentes dessa arquitetura:

Arquitetura do cluster do GKE com pools de nós de TPU de host único que contêm os componentes Maxengine e HTTP.

Essa arquitetura inclui os seguintes componentes:

  • Um cluster regional ou Autopilot do GKE.
  • Dois pools de nós de frações de TPU de host único que hospedam a implantação do JetStream.
  • O componente de serviço distribui o tráfego de entrada para todas as réplicas JetStream HTTP.
  • O JetStream HTTP é um servidor HTTP que aceita solicitações como um wrapper para o formato obrigatório do JetStream e as envia para o cliente GRPC do JetStream (link em inglês).
  • O Maxengine é um servidor JetStream que realiza inferências com lotes contínuos.

Antes de começar

  • Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

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

    Go to project selector

  • Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

  • Enable the required API.

    Enable the API

    •

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

    Go to project selector

  • Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

  • 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

    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 colunn 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. Selecionar um projeto.
    3. Clique em CONCEDER ACESSO.

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

    5. Na lista Selecionar um papel, escolha um.
    6. Para conceder outros papéis, clique em Adicionar outro papel e adicione cada papel adicional.
    7. Clique em Salvar.

Receber acesso ao modelo

Para ter acesso ao modelo Gemma para implantação no GKE, primeiro assine o contrato de consentimento de licença.

É necessário assinar o contrato de consentimento para usar o Gemma. Siga estas instruções:

  1. Acesse a página de consentimento do modelo Gemma em Kaggle.com.
  2. Faça login no Kaggle se ainda não tiver feito isso.
  3. Clique em Solicitar acesso.
  4. Na seção Escolher conta para consentimento, selecione Verificar via conta do Kaggle para usar sua conta do Kaggle para consentimento.
  5. Aceite os Termos e Condições do modelo.

Gerar um token de acesso

Para acessar o modelo pelo Kaggle, você precisa de um token da API do Kaggle.

Siga as etapas abaixo para gerar um novo token, caso ainda não tenha um:

  1. No navegador, acesse as configurações da Kaggle.
  2. Na seção API, clique em Criar novo token.

Um arquivo chamado kaggle.json é transferido por download.

prepare o ambiente

Neste tutorial, você usará o Cloud Shell para gerenciar recursos hospedados no Google Cloud. O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl e gcloud CLI.

Para configurar o ambiente com o Cloud Shell, siga estas etapas:

  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 do Google Cloud. Isso inicia uma sessão no painel inferior do console do Google Cloud.

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

    gcloud config set project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export CLUSTER_NAME=CLUSTER_NAME
export BUCKET_NAME=BUCKET_NAME
export REGION=REGION
export LOCATION=LOCATION

    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://.
    • REGION_NAME: a região em que o cluster do GKE, o bucket do Cloud Storage e os nós da TPU estão localizados. A região contém zonas em que os tipos de máquina TPU v5e estão disponíveis (por exemplo, us-west1, us-west4, us-central1, us-east1, us-east5 ou europe-west4).
    • (Apenas cluster padrão) LOCATION: a zona em que os recursos da TPU estão disponíveis, por exemplo, us-west4-a. Para clusters do Autopilot, não é necessário especificar a zona, apenas a região.

Criar e configurar recursos do Google Cloud

Siga estas instruções para criar os recursos necessários.

Crie um cluster do GKE

É possível disponibilizar o Gemma em TPUs em um cluster do GKE Autopilot ou Standard. Recomendamos que você use um cluster do Autopilot para ter uma experiência totalmente gerenciada do Kubernetes. Para escolher o modo de operação do GKE mais adequado para suas cargas de trabalho, consulte Escolher um modo de operação do GKE.

Piloto automático

No Cloud Shell, execute este comando:

gcloud container clusters create-auto ${CLUSTER_NAME} \
  --project=${PROJECT_ID} \
  --region=${REGION} \
  --release-channel=rapid \
  --cluster-version=1.29

O GKE cria um cluster do Autopilot com nós de CPU e TPU conforme solicitado pelas cargas de trabalho implantadas.

Padrão

  1. Criar um cluster regional do GKE Standard que use a federação de identidade da carga de trabalho para o GKE.

    gcloud container clusters create ${CLUSTER_NAME} \
    --enable-ip-alias \
    --machine-type=e2-standard-4 \
    --num-nodes=2 \
    --release-channel=rapid \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --location=${REGION}

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

  2. Execute o seguinte comando para criar um pool de nós para o cluster:

    gcloud container node-pools create gemma-7b-tpu-nodepool \
  --cluster=${CLUSTER_NAME} \
  --machine-type=ct5lp-hightpu-8t \
  --project=${PROJECT_ID} \
  --num-nodes=2 \
  --region=${REGION} \
  --node-locations=${LOCATION}

    O GKE cria um pool de nós TPU v5e com uma topologia 2x4 e dois nós.

Criar um bucket do Cloud Storage

No Cloud Shell, execute este comando:

gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}

Isso cria um bucket do Cloud Storage para armazenar os arquivos de modelo transferidos por download do Kaggle.

Fazer upload do token de acesso no Cloud Shell

No Cloud Shell, é possível fazer upload do token da API Kaggle para seu projeto do Google Cloud:

  1. No Cloud Shell, clique em Mais > Fazer upload.
  2. Selecione "Arquivo" e clique em Escolher arquivos.
  3. Abra o arquivo kaggle.json.
  4. Clique em Fazer upload.

Crie um Secret do Kubernetes para credenciais do Kaggle

No Cloud Shell, faça o seguinte:

  1. Configure kubectl para se comunicar com o cluster:

    gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}

  2. Crie um Secret para armazenar as credenciais do Kaggle:

    kubectl create secret generic kaggle-secret \
    --from-file=kaggle.json

Configurar o acesso das cargas de trabalho usando a federação de identidade da carga de trabalho do GKE

Atribua uma Kubernetes ServiceAccount ao aplicativo e configure-a para atuar como uma conta de serviço do IAM.

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

    gcloud iam service-accounts create wi-jetstream

  2. Adicione uma vinculação de política do IAM à sua conta de serviço do IAM para gerenciar o Cloud Storage:

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

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

  3. Permita que a conta de serviço do Kubernetes atue como a conta de serviço do IAM adicionando uma vinculação de política do IAM entre as duas. Essa vinculação permite que a Kubernetes ServiceAccount atue como a conta de serviço do Google.

    gcloud iam service-accounts add-iam-policy-binding wi-jetstream@${PROJECT_ID}.iam.gserviceaccount.com \
    --role roles/iam.workloadIdentityUser \
    --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"

  4. Anote a conta de serviço do Kubernetes com o endereço de e-mail da conta de serviço do IAM

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

Converter os checkpoints do modelo

Nesta seção, você cria um job para fazer o seguinte:

  1. Faça o download do checkpoint básico da Orbax pelo Kaggle.
  2. Faça upload do checkpoint em um bucket do Cloud Storage.
  3. Converta o checkpoint em um checkpoint compatível com MaxText.
  4. Cancele a verificação do checkpoint a ser usado para veiculação.

Implantar o job de conversão do checkpoint do modelo

Siga estas instruções para fazer o download e converter os arquivos de checkpoint do modelo do Gemma 7B.

  1. Crie o seguinte manifesto como job-7b.yaml.

    apiVersion: batch/v1
kind: Job
metadata:
  name: data-loader-7b
spec:
  ttlSecondsAfterFinished: 30
  template:
    spec:
      restartPolicy: Never
      containers:
      - name: inference-checkpoint
        image: us-docker.pkg.dev/cloud-tpu-images/inference/inference-checkpoint:v0.2.2
        args:
        - -b=BUCKET_NAME
        - -m=google/gemma/maxtext/7b-it/2
        volumeMounts:
        - mountPath: "/kaggle/"
          name: kaggle-credentials
          readOnly: true
        resources:
          requests:
            google.com/tpu: 8
          limits:
            google.com/tpu: 8
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      volumes:
      - name: kaggle-credentials
        secret:
          defaultMode: 0400
          secretName: kaggle-secret

  2. Aplique o manifesto:

    kubectl apply -f job-7b.yaml

  3. Veja os registros do job:

    kubectl logs -f jobs/data-loader-7b

    Quando o job for concluído, a saída será semelhante a esta:

    Successfully generated decode checkpoint at: gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items
+ echo -e '\nCompleted unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items'

Completed unscanning checkpoint to gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items

Implantar o JetStream

Nesta seção, você vai implantar o contêiner JetStream para exibir o modelo Gemma.

Siga estas instruções para implantar o modelo ajustado por instruções do Gemma 7B.

  1. Crie o seguinte manifesto jetstream-gemma-deployment.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: maxengine-server
spec:
  replicas: 2
  selector:
    matchLabels:
      app: maxengine-server
  template:
    metadata:
      labels:
        app: maxengine-server
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: maxengine-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/maxengine-server:v0.2.2
        args:
        - model_name=gemma-7b
        - tokenizer_path=assets/tokenizer.gemma
        - per_device_batch_size=4
        - max_prefill_predict_length=1024
        - max_target_length=2048
        - async_checkpointing=false
        - ici_fsdp_parallelism=1
        - ici_autoregressive_parallelism=-1
        - ici_tensor_parallelism=1
        - scan_layers=false
        - weight_dtype=bfloat16
        - load_parameters_path=gs://BUCKET_NAME/final/unscanned/gemma_7b-it/0/checkpoints/0/items
        ports:
        - containerPort: 9000
        resources:
          requests:
            google.com/tpu: 8
          limits:
            google.com/tpu: 8
      - name: jetstream-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.2
        ports:
        - containerPort: 8000
---
apiVersion: v1
kind: Service
metadata:
  name: jetstream-svc
spec:
  selector:
    app: maxengine-server
  ports:
  - protocol: TCP
    name: jetstream-http
    port: 8000
    targetPort: 8000
  - protocol: TCP
    name: jetstream-grpc
    port: 9000
    targetPort: 9000

    O manifesto define as seguintes propriedades principais:

    • tokenizer_path: o caminho para o tokenizador do modelo.
    • load_parameters_path: o caminho no bucket do Cloud Storage em que os checkpoints são armazenados.
    • per_device_batch_size: o tamanho do lote de decodificação por dispositivo, em que um chip de TPU é igual a um dispositivo.
    • max_prefill_predict_length: o tamanho máximo do preenchimento prévio ao fazer a regressão automática.
    • max_target_length: o tamanho máximo da sequência.
    • model_name: o nome do modelo (gemma-7b).
    • ici_fsdp_parallelism: o número de fragmentos para paralelismo de dados totalmente fragmentados (FSDP, na sigla em inglês).
    • ici_tensor_parallelism: o número de fragmentos para paralelismo de tensor.
    • ici_autoregressive_parallelism: o número de fragmentos para paralelismo autoregressivo.
    • scan_layers: scan layers boolean flag (boolean).
    • weight_dtype: o tipo de dados de peso (bfloat16).

  2. Aplique o manifesto:

    kubectl apply -f jetstream-gemma-deployment.yaml

  3. Verificar a implantação:

    kubectl get deployment

    O resultado será assim:

    NAME                              READY   UP-TO-DATE   AVAILABLE   AGE
maxengine-server                  2/2     2            2           ##s

  4. Visualize os registros do servidor HTTP para verificar se o modelo foi carregado e compilado. O servidor pode levar alguns minutos para concluir essa operação.

    kubectl logs deploy/maxengine-server -f -c jetstream-http

    O resultado será assim:

    kubectl logs deploy/maxengine-server -f -c jetstream-http

INFO:     Started server process [1]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)

  5. Consulte os registros do MaxEngine e verifique se a compilação foi concluída.

    kubectl logs deploy/maxengine-server -f -c maxengine-server

    O resultado será assim:

    2024-03-29 17:09:08,047 - jax._src.dispatch - DEBUG - Finished XLA compilation of jit(initialize) in 0.26236414909362793 sec
2024-03-29 17:09:08,150 - root - INFO - ---------Generate params 0 loaded.---------

Disponibilizar o modelo

Nesta seção, você vai interagir com o modelo.

Configurar o encaminhamento de portas

Acesse a implantação do JetStream pelo Serviço ClusterIP criado na etapa anterior. Os serviços ClusterIP só podem ser acessados dentro do cluster. Portanto, para acessar o serviço de fora do cluster, conclua as seguintes etapas:

Para estabelecer uma sessão de encaminhamento de portas, execute o seguinte comando:

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

Interagir com o modelo usando curl

  1. Para verificar se é possível acessar o servidor HTTP JetStream, abra um novo terminal e execute o seguinte comando:

    curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

    A solicitação inicial pode levar alguns segundos para ser concluída devido ao aquecimento do modelo. O resultado será assim:

    {
    "response": "\nfor data science in 2023?\n\n**1. Python:**\n- Widely used for data science due to its simplicity, readability, and extensive libraries for data wrangling, analysis, visualization, and machine learning.\n- Popular libraries include pandas, scikit-learn, and matplotlib.\n\n**2. R:**\n- Statistical programming language widely used for data analysis, visualization, and modeling.\n- Popular libraries include ggplot2, dplyr, and caret.\n\n**3. Java:**\n- Enterprise-grade language with strong performance and scalability.\n- Popular libraries include Spark, TensorFlow, and Weka.\n\n**4. C++:**\n- High-performance language often used for data analytics and machine learning models.\n- Popular libraries include TensorFlow, PyTorch, and OpenCV.\n\n**5. SQL:**\n- Relational database language essential for data wrangling and querying large datasets.\n- Popular tools"
}

(Opcional) Interagir com o modelo usando uma interface de chat do GRadio

Nesta seção, você criará um aplicativo de chat na Web que permite interagir com seu modelo ajustado por instruções.

O Gradio é uma biblioteca Python que tem um wrapper ChatInterface que cria interfaces de usuário para chatbots.

Implantar a interface de chat

  1. No Cloud Shell, salve o seguinte manifesto como gradio.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: gradio
  labels:
    app: gradio
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gradio
  template:
    metadata:
      labels:
        app: gradio
    spec:
      containers:
      - name: gradio
        image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.3
        resources:
          requests:
            cpu: "512m"
            memory: "512Mi"
          limits:
            cpu: "1"
            memory: "512Mi"
        env:
        - name: CONTEXT_PATH
          value: "/generate"
        - name: HOST
          value: "http://jetstream-http-svc:8000"
        - name: LLM_ENGINE
          value: "max"
        - name: MODEL_ID
          value: "gemma"
        - name: USER_PROMPT
          value: "<start_of_turn>user\nprompt<end_of_turn>\n"
        - name: SYSTEM_PROMPT
          value: "<start_of_turn>model\nprompt<end_of_turn>\n"
        ports:
        - containerPort: 7860
---
apiVersion: v1
kind: Service
metadata:
  name: gradio
spec:
  selector:
    app: gradio
  ports:
    - protocol: TCP
      port: 8080
      targetPort: 7860
  type: ClusterIP

  2. Aplique o manifesto:

    kubectl apply -f gradio.yaml

  3. Aguarde até que a implantação esteja disponível:

    kubectl wait --for=condition=Available --timeout=300s deployment/gradio

Usar a interface de chat

  1. No Cloud Shell, execute este comando:

    kubectl port-forward service/gradio 8080:8080

    Isso cria um encaminhamento de porta do Cloud Shell para o serviço GRadio.

  2. Clique no botão Ícone de visualização na Web Visualização da Web no canto superior direito da barra de tarefas do Cloud Shell. Clique em Visualizar na porta 8080. Uma nova guia será aberta no navegador.

  3. Interaja com Gemma usando a interface de chat do GRadio. Adicione uma solicitação e clique em Enviar.

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 Connection refused, verifique se o encaminhamento de portas está ativo.

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Excluir os recursos implantados

Para evitar cobranças na sua conta do Google Cloud pelos recursos criados neste guia, execute os comandos a seguir e siga as instruções:

gcloud container clusters delete ${CLUSTER_NAME} --region=${REGION}

gcloud iam service-accounts delete wi-jetstream@PROJECT_ID.iam.gserviceaccount.com

gcloud storage rm --recursive gs://BUCKET_NAME

A seguir