Ajustar modelos abertos do Gemma usando várias GPUs no GKE

Neste tutorial, mostramos como ajustar o modelo de linguagem grande (LLM) Gemma, família de modelos abertos, usando unidades de processamento gráfico (GPUs) no Google Kubernetes Engine (GKE) com a biblioteca Transformers do Hugging Face. O ajuste fino é um processo de aprendizado supervisionado que melhora a capacidade de um modelo pré-treinado de realizar tarefas específicas atualizando os parâmetros com um novo conjunto de dados. Neste tutorial, você baixa os modelos pré-treinados da família Gemma de parâmetros 2B do Hugging Face e os ajusta em um cluster do GKE Autopilot ou Standard.

Este guia é um bom ponto de partida se você precisar do controle granular, da escalonabilidade, da resiliência, da portabilidade e da economia do Kubernetes gerenciado ao ajustar um LLM.

Prática recomendada:

Teste nossa solução Vertex AI se você precisar de uma plataforma de IA gerenciada unificada para criar e disponibilizar modelos de ML rapidamente e de maneira econômica.

Contexto

Ao disponibilizar o Gemma usando GPUs no GKE com a biblioteca de transformação, é possível implementar uma solução de veiculação de inferência robusta e pronta para produção com todos os benefícios do Kubernetes gerenciado, incluindo escalonabilidade eficiente e maior disponibilidade. Esta seção descreve as principais tecnologias usadas neste guia.

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.

Neste guia, apresentamos o Gemma para geração de texto. Também é possível ajustar esses modelos para se especializar na execução de tarefas específicas.

O conjunto de dados que você usa neste documento é b-mc2/sql-create-context.

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

GPUs

As GPUs permitem acelerar cargas de trabalho específicas em execução nos nós, como machine learning e processamento de dados. O GKE oferece uma gama de opções de tipos de máquina para configuração de nós, incluindo tipos de máquinas com GPUs NVIDIA H100, L4 e A100.

Antes de usar GPUs no GKE, considere concluir o seguinte programa de aprendizado:

  1. Saiba mais sobre a disponibilidade atual da versão da GPU.
  2. Saiba mais sobre GPUs no .

Transformers do Hugging Face

Com a biblioteca Transformers do Hugging Face, você pode acessar modelos pré-treinados de última geração. A biblioteca Transformers permite reduzir o tempo, os recursos e os custos computacionais associados ao treinamento completo de modelo.

Neste tutorial, você vai usar as APIs e ferramentas Hugging Face para fazer o download e ajustar esses modelos pré-treinados.

Receber acesso ao modelo

Para ter acesso aos modelos Gemma para implantação no GKE, primeiro assine o contrato de consentimento de licença e gere um token de acesso do Hugging Face.

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

  1. Acesse a página de consentimento do modelo em Kaggle.com.
  2. Confirme seu consentimento usando sua conta do Hugging Face.
  3. Aceite os termos do modelo.

Gerar um token de acesso

Para acessar o modelo usando Hugging Face, você vai precisar de um token do Hugging Face.

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

  1. Clique em Seu perfil > Configurações > Tokens de acesso.
  2. Selecione Novo token.
  3. Especifique um Nome de sua escolha e um Papel de pelo menos Write.
  4. Selecione Gerar um token.
  5. Copie o token gerado para a área de transferência.

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

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

    gcloud config set project PROJECT_ID
    export PROJECT_ID=$(gcloud config get project)
    export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
    export CLUSTER_NAME=CLUSTER_NAME
    export HF_TOKEN=HF_TOKEN
    export HF_PROFILE=HF_PROFILE
    

    Substitua os seguintes valores:

    • PROJECT_ID: o Google Cloud ID do projeto.
    • CONTROL_PLANE_LOCATION: a região do Compute Engine do plano de controle do cluster. Forneça uma região compatível com o tipo de acelerador que você quer usar, por exemplo, us-central1 para GPUs L4.
    • CLUSTER_NAME: o nome do cluster.
    • HF_TOKEN: o token do Hugging Face gerado anteriormente.
    • HF_PROFILE: o ID do perfil do Hugging Face que você criou antes.
  3. Clone o repositório de exemplo de código do GitHub:

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
    cd kubernetes-engine-samples/ai-ml/llm-finetuning-gemma
    

Criar e configurar recursos Google Cloud

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

Criar um cluster do GKE e um pool de nós

É possível disponibilizar o Gemma em GPUs em um cluster do GKE Autopilot ou Standard. 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.

Prática recomendada:

Use o Autopilot para ter uma experiência totalmente gerenciada do Kubernetes.

Piloto automático

No Cloud Shell, execute este comando:

gcloud container clusters create-auto CLUSTER_NAME \
    --project=PROJECT_ID \
    --location=CONTROL_PLANE_LOCATION \
    --release-channel=rapid \
    --cluster-version=1.29

Substitua os seguintes valores:

  • PROJECT_ID: o Google Cloud ID do projeto.
  • CONTROL_PLANE_LOCATION: a região do Compute Engine do plano de controle do cluster. Forneça uma região compatível com o tipo de acelerador que você quer usar, por exemplo, us-central1 para GPUs L4.
  • CLUSTER_NAME: o nome do cluster.

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

Padrão

  1. No Cloud Shell, execute o seguinte comando para criar um cluster Standard:

    gcloud container clusters create CLUSTER_NAME \
        --project=PROJECT_ID \
        --location=CONTROL_PLANE_LOCATION \
        --workload-pool=PROJECT_ID.svc.id.goog \
        --release-channel=rapid \
        --num-nodes=1
    

    Substitua os seguintes valores:

    • PROJECT_ID: o Google Cloud ID do projeto.
    • CONTROL_PLANE_LOCATION: a região do Compute Engine do plano de controle do cluster. Forneça uma região compatível com o tipo de acelerador que você quer usar, por exemplo, us-central1 para GPUs L4.
    • CLUSTER_NAME: o nome do cluster.

    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 gpupool \
        --accelerator type=nvidia-l4,count=8,gpu-driver-version=latest \
        --project=PROJECT_ID \
        --location=CONTROL_PLANE_LOCATION \
        --node-locations=CONTROL_PLANE_LOCATION-a \
        --cluster=CLUSTER_NAME \
        --machine-type=g2-standard-96 \
        --num-nodes=1
    

    O GKE cria um único pool de nós contendo duas GPUs L4 para cada nó.

Criar um secret do Kubernetes para as credenciais do Hugging Face

No Cloud Shell, faça o seguinte:

  1. Configure kubectl para se comunicar com o cluster:

    gcloud container clusters get-credentials CLUSTER_NAME \
        --location=CONTROL_PLANE_LOCATION
    

    Substitua os seguintes valores:

    • CONTROL_PLANE_LOCATION: a região do Compute Engine do plano de controle do cluster.
    • CLUSTER_NAME: o nome do cluster.
  2. Crie um secret do Kubernetes que contenha o token do Hugging Face:

    kubectl create secret generic hf-secret \
        --from-literal=hf_api_token=$HF_TOKEN \
        --dry-run=client -o yaml | kubectl apply -f -
    

    Substitua $HF_TOKEN pelo token do Hugging Face que você gerou anteriormente ou use a variável de ambiente, se você a definiu.

Criar um contêiner de ajuste fino com o Docker e o Cloud Build

Esse contêiner usa o código dos transformadores PyTorch e Hugging Face para ajustar o modelo Gemma pré-treinado existente.

  1. Crie um repositório Docker do Artifact Registry:

    gcloud artifacts repositories create gemma \
        --project=PROJECT_ID \
        --repository-format=docker \
        --location=us \
        --description="Gemma Repo"
    

    Substitua PROJECT_ID pelo ID do projeto do Google Cloud.

  2. Crie e envie a imagem:

    gcloud builds submit .
    
  3. Exporte o IMAGE_URL para uso posterior neste tutorial.

    export IMAGE_URL=us-docker.pkg.dev/PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0
    

Executar um job de ajuste no GKE

Nesta seção, você vai implantar o job de ajuste fino do Gemma. Um controlador de job no Kubernetes cria um ou mais pods e garante que eles executem uma tarefa específica com sucesso.

  1. Abra o arquivo finetune.yaml.

    apiVersion: batch/v1
    kind: Job
    metadata:
      name: finetune-job
      namespace: default
    spec:
      backoffLimit: 2
      template:
        metadata:
          annotations:
            kubectl.kubernetes.io/default-container: finetuner
        spec:
          terminationGracePeriodSeconds: 600
          containers:
          - name: finetuner
            image: $IMAGE_URL
            resources:
              limits:
                nvidia.com/gpu: "8"
            env:
            - name: MODEL_NAME
              value: "google/gemma-2b"
            - name: NEW_MODEL
              value: "gemma-2b-sql-finetuned"
            - name: LORA_R
              value: "8"
            - name: LORA_ALPHA
              value: "16"
            - name: TRAIN_BATCH_SIZE
              value: "1"
            - name: EVAL_BATCH_SIZE
              value: "2"
            - name: GRADIENT_ACCUMULATION_STEPS
              value: "2"
            - name: DATASET_LIMIT
              value: "1000"
            - name: MAX_SEQ_LENGTH
              value: "512"
            - name: LOGGING_STEPS
              value: "5"
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
          nodeSelector:
            cloud.google.com/gke-accelerator: nvidia-l4
          restartPolicy: OnFailure
  2. Aplique o manifesto para criar o job de ajuste:

    envsubst < finetune.yaml | kubectl apply -f -
    

    Essa instrução substitui o IMAGE_URL pela variável no manifesto.

  3. Implante o job novamente executando o seguinte comando:

    watch kubectl get pods
    
  4. Verifique os registros do job executando o seguinte comando:

    kubectl logs job.batch/finetune-job -f
    

    O recurso de job faz o download dos dados do modelo e, em seguida, ajusta o modelo em todas as oito GPUs. Esse processo pode levar até 20 minutos.

  5. Quando o job for concluído, acesse sua conta do Hugging Face. Um novo modelo chamado HF_PROFILE/gemma-2b-sql-finetuned aparece no seu perfil do Hugging Face.

Disponibilizar o modelo ajustado no GKE

Nesta seção, você vai implantar o contêiner vLLM para exibir o modelo Gemma. Este tutorial usa uma implantação do Kubernetes para implantar o contêiner vLLM. Uma implantação é um objeto da API Kubernetes que permite executar várias réplicas de pods distribuídos entre os nós de um cluster.

  1. Crie o seguinte manifesto serve-gemma.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: vllm-gemma-deployment
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: gemma-server
      template:
        metadata:
          labels:
            app: gemma-server
            ai.gke.io/model: gemma-2b
            ai.gke.io/inference-server: vllm
            examples.ai.gke.io/source: user-guide
        spec:
          containers:
          - name: inference-server
            image: docker.io/vllm/vllm-openai:v0.10.0
            resources:
              requests:
                cpu: "2"
                memory: "7Gi"
                ephemeral-storage: "10Gi"
                nvidia.com/gpu: 1
              limits:
                cpu: "2"
                memory: "7Gi"
                ephemeral-storage: "10Gi"
                nvidia.com/gpu: 1
            command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
            args:
            - --model=$(MODEL_ID)
            - --tensor-parallel-size=1
            env:
            - name: LD_LIBRARY_PATH
              value: ${LD_LIBRARY_PATH}:/usr/local/nvidia/lib64
            - name: MODEL_ID
              value: google/gemma-2b
            - name: HUGGING_FACE_HUB_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
                medium: Memory
          nodeSelector:
            cloud.google.com/gke-accelerator: nvidia-l4
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: llm-service
    spec:
      selector:
        app: gemma-server
      type: ClusterIP
      ports:
        - protocol: TCP
          port: 8000
          targetPort: 8000
  2. Crie a variável de ambiente para o novo MODEL_ID:

    export MODEL_ID=HF_PROFILE/gemma-2b-sql-finetuned
    

    Substitua HF_PROFILE pelo ID do perfil do Hugging Face que você criou anteriormente.

  3. Substitua MODEL_ID no manifesto:

    sed -i "s|google/gemma-2b|$MODEL_ID|g" serve-gemma.yaml
    
  4. Aplique o manifesto:

    kubectl apply -f serve-gemma.yaml
    

    Um pod no cluster faz o download dos pesos do modelo do Hugging Face e inicia o mecanismo de exibição.

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

    kubectl wait --for=condition=Available --timeout=700s deployment/vllm-gemma-deployment
    
  6. Confira os registros da implantação em execução:

    kubectl logs -f -l app=gemma-server
    

O recurso de implantação faz o download dos dados do modelo. O processo pode levar alguns minutos. O resultado será assim:

INFO 01-26 19:02:54 model_runner.py:689] Graph capturing finished in 4 secs.
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)

Verifique se o download do modelo foi concluído antes de prosseguir para a próxima seção.

Disponibilizar o modelo

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

Configurar o encaminhamento de portas

Depois que o modelo for implantado, execute o comando abaixo para configurar o encaminhamento de portas para o modelo:

kubectl port-forward service/llm-service 8000:8000

O resultado será assim:

Forwarding from 127.0.0.1:8000 -> 8000

Interagir com o modelo usando curl

Em uma nova sessão do terminal, use curl para conversar com seu modelo:

O exemplo de comando a seguir é para TGI:

USER_PROMPT="Question: What is the total number of attendees with age over 30 at kubecon eu? Context: CREATE TABLE attendees (name VARCHAR, age INTEGER, kubecon VARCHAR)"

curl -X POST http://localhost:8000/generate \
  -H "Content-Type: application/json" \
  -d @- <<EOF
{
    "prompt": "${USER_PROMPT}",
    "temperature": 0.1,
    "top_p": 1.0,
    "max_tokens": 24
}
EOF

A saída a seguir mostra um exemplo da resposta do modelo:

{"generated_text":" Answer: SELECT COUNT(age) FROM attendees WHERE age > 30 AND kubecon = 'eu'\n"}

Dependendo da consulta, talvez seja necessário alterar o max_token para conseguir um resultado melhor. Também é possível usar o modelo de orientação personalizada para uma melhor experiência de chat.