Cómo ajustar un LLM con trabajos de Cloud Run

1. Introducción

Descripción general

En este codelab, usarás trabajos de Cloud Run para ajustar un modelo de Gemma y, luego, entregar el resultado en Cloud Run con vLLM.

Para los fines de este codelab, usarás un conjunto de datos de texto a SQL, cuyo objetivo es hacer que el LLM responda con una consulta de SQL cuando se le haga una pregunta en lenguaje natural.

Qué aprenderás

• Cómo realizar un ajuste fino con la GPU de trabajos de Cloud Run
• Cómo entregar un modelo con Cloud Run con vLLM
• Cómo usar la configuración de VPC directa para un trabajo de GPU para subir y entregar el modelo más rápido

2. Antes de comenzar

Habilita las APIs

Antes de comenzar a usar este codelab, habilita las siguientes APIs ejecutando el siguiente comando:

gcloud services enable run.googleapis.com \
    compute.googleapis.com \
    run.googleapis.com \
    cloudbuild.googleapis.com \
    secretmanager.googleapis.com \
    artifactregistry.googleapis.com

Cuota de GPU

Solicita un aumento de la cuota para una región compatible. La cuota es nvidia_l4_gpu_allocation_no_zonal_redundancy, en la API de Cloud Run Admin.

Nota: Si usas un proyecto nuevo, es posible que debas esperar unos minutos entre la habilitación de la API y la aparición de las cuotas en esta página.

Hugging Face

En este codelab, se usa un modelo alojado en Hugging Face. Para obtener este modelo, solicita el token de acceso de usuario de Hugging Face con el permiso de "Leer". Más adelante, harás referencia a esto como YOUR_HF_TOKEN.

También deberás aceptar las condiciones de uso para usar el modelo: https://huggingface.co/google/gemma-2b

3. Configuración y requisitos

Configura los siguientes recursos:

• Cuenta de servicio de IAM y permisos de IAM asociados
• Secreto de Secret Manager para almacenar tu token de Hugging Face
• Bucket de Cloud Storage para almacenar tu modelo ajustado
• Repositorio de Artifact Registry para almacenar la imagen que compilarás para ajustar tu modelo.

1. Configura las variables de entorno para este codelab. Completamos previamente una serie de variables por ti. Especifica el ID, la región y el token de Hugging Face de tu proyecto.

   export PROJECT_ID=<YOUR_PROJECT_ID>
   export REGION=<YOUR_REGION>
   export HF_TOKEN=<YOUR_HF_TOKEN>
   
   export AR_REPO=codelab-finetuning-jobs
   export IMAGE_NAME=finetune-to-gcs
   export JOB_NAME=finetuning-to-gcs-job
   export BUCKET_NAME=$PROJECT_ID-codelab-finetuning-jobs
   export SECRET_ID=HF_TOKEN
   export SERVICE_ACCOUNT="finetune-job-sa"
   export SERVICE_ACCOUNT_ADDRESS=$SERVICE_ACCOUNT@$PROJECT_ID.iam.gserviceaccount.com
   

2. Ejecuta este comando para crear la cuenta de servicio:

   gcloud iam service-accounts create $SERVICE_ACCOUNT \
     --display-name="Service account for fine-tuning codelab"
   

3. Usa Secret Manager para almacenar el token de acceso de Hugging Face:

   gcloud secrets create $SECRET_ID \
         --replication-policy="automatic"
   
   printf $HF_TOKEN | gcloud secrets versions add $SECRET_ID --data-file=-
   

4. Otorga a tu cuenta de servicio el rol de descriptor de acceso a secretos de Secret Manager:

   gcloud secrets add-iam-policy-binding $SECRET_ID \
     --member serviceAccount:$SERVICE_ACCOUNT_ADDRESS \
     --role='roles/secretmanager.secretAccessor'
   

5. Crea un bucket que alojará tu modelo ajustado:

   gcloud storage buckets create -l $REGION gs://$BUCKET_NAME
   

6. Otorga a tu cuenta de servicio acceso al bucket:

   gcloud storage buckets add-iam-policy-binding gs://$BUCKET_NAME \
     --member=serviceAccount:$SERVICE_ACCOUNT_ADDRESS \
     --role=roles/storage.objectAdmin
   

7. Crea un repositorio de Artifact Registry para almacenar la imagen del contenedor:

   gcloud artifacts repositories create $AR_REPO \
       --repository-format=docker \
       --location=$REGION \
       --description="codelab for finetuning using CR jobs" \
       --project=$PROJECT_ID
   

4. Crea la imagen del trabajo de Cloud Run

En el siguiente paso, crearás el código que hará lo siguiente:

• Importa el modelo de Gemma desde Hugging Face.
• Realiza un ajuste fino del modelo con el conjunto de datos de Hugging Face. La tarea usa una sola GPU L4 para el ajuste fino.
• Sube el modelo ajustado llamado new_model a tu bucket de Cloud Storage.

1. Crea un directorio para el código de trabajo de ajuste fino.

   mkdir codelab-finetuning-job
   cd codelab-finetuning-job
   

2. Crea un archivo llamado finetune.py.

   # Copyright 2024 Google LLC
   #
   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at
   #
   #      http://www.apache.org/licenses/LICENSE-2.0
   #
   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.
   
   import os
   import torch
   from datasets import load_dataset
   from transformers import (
       AutoModelForCausalLM,
       AutoTokenizer,
       BitsAndBytesConfig,
       TrainingArguments,
   
   )
   from peft import LoraConfig, PeftModel
   
   from trl import SFTTrainer
   
   # Cloud Storage bucket to upload the model
   bucket_name = os.getenv("BUCKET_NAME", "YOUR_BUCKET_NAME")
   
   # The model that you want to train from the Hugging Face hub
   model_name = os.getenv("MODEL_NAME", "google/gemma-2b")
   
   # The instruction dataset to use
   dataset_name = "b-mc2/sql-create-context"
   
   # Fine-tuned model name
   new_model = os.getenv("NEW_MODEL", "gemma-2b-sql")
   
   ################################################################################
   # QLoRA parameters
   ################################################################################
   
   # LoRA attention dimension
   lora_r = int(os.getenv("LORA_R", "4"))
   
   # Alpha parameter for LoRA scaling
   lora_alpha = int(os.getenv("LORA_ALPHA", "8"))
   
   # Dropout probability for LoRA layers
   lora_dropout = 0.1
   
   ################################################################################
   # bitsandbytes parameters
   ################################################################################
   
   # Activate 4-bit precision base model loading
   use_4bit = True
   
   # Compute dtype for 4-bit base models
   bnb_4bit_compute_dtype = "float16"
   
   # Quantization type (fp4 or nf4)
   bnb_4bit_quant_type = "nf4"
   
   # Activate nested quantization for 4-bit base models (double quantization)
   use_nested_quant = False
   
   ################################################################################
   # TrainingArguments parameters
   ################################################################################
   
   # Output directory where the model predictions and checkpoints will be stored
   output_dir = "./results"
   
   # Number of training epochs
   num_train_epochs = 1
   
   # Enable fp16/bf16 training (set bf16 to True with an A100)
   fp16 = True
   bf16 = False
   
   # Batch size per GPU for training
   per_device_train_batch_size = int(os.getenv("TRAIN_BATCH_SIZE", "1"))
   
   # Batch size per GPU for evaluation
   per_device_eval_batch_size = int(os.getenv("EVAL_BATCH_SIZE", "2"))
   
   # Number of update steps to accumulate the gradients for
   gradient_accumulation_steps = int(os.getenv("GRADIENT_ACCUMULATION_STEPS", "1"))
   
   # Enable gradient checkpointing
   gradient_checkpointing = True
   
   # Maximum gradient normal (gradient clipping)
   max_grad_norm = 0.3
   
   # Initial learning rate (AdamW optimizer)
   learning_rate = 2e-4
   
   # Weight decay to apply to all layers except bias/LayerNorm weights
   weight_decay = 0.001
   
   # Optimizer to use
   optim = "paged_adamw_32bit"
   
   # Learning rate schedule
   lr_scheduler_type = "cosine"
   
   # Number of training steps (overrides num_train_epochs)
   max_steps = -1
   
   # Ratio of steps for a linear warmup (from 0 to learning rate)
   warmup_ratio = 0.03
   
   # Group sequences into batches with same length
   # Saves memory and speeds up training considerably
   group_by_length = True
   
   # Save checkpoint every X updates steps
   save_steps = 0
   
   # Log every X updates steps
   logging_steps = int(os.getenv("LOGGING_STEPS", "50"))
   
   ################################################################################
   # SFT parameters
   ################################################################################
   
   # Maximum sequence length to use
   max_seq_length = int(os.getenv("MAX_SEQ_LENGTH", "512"))
   
   # Pack multiple short examples in the same input sequence to increase efficiency
   packing = False
   
   # Load the entire model on the GPU 0
   device_map = {'':torch.cuda.current_device()}
   
   # Set limit to a positive number
   limit = int(os.getenv("DATASET_LIMIT", "5000"))
   
   dataset = load_dataset(dataset_name, split="train")
   if limit != -1:
       dataset = dataset.shuffle(seed=42).select(range(limit))
   
   
   def transform(data):
       question = data['question']
       context = data['context']
       answer = data['answer']
       template = "Question: {question}\nContext: {context}\nAnswer: {answer}"
       return {'text': template.format(question=question, context=context, answer=answer)}
   
   
   transformed = dataset.map(transform)
   
   # Load tokenizer and model with QLoRA configuration
   compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
   
   bnb_config = BitsAndBytesConfig(
       load_in_4bit=use_4bit,
       bnb_4bit_quant_type=bnb_4bit_quant_type,
       bnb_4bit_compute_dtype=compute_dtype,
       bnb_4bit_use_double_quant=use_nested_quant,
   )
   
   # Check GPU compatibility with bfloat16
   if compute_dtype == torch.float16 and use_4bit:
       major, _ = torch.cuda.get_device_capability()
       if major >= 8:
           print("=" * 80)
           print("Your GPU supports bfloat16")
           print("=" * 80)
   
   # Load base model
   model = AutoModelForCausalLM.from_pretrained(
       model_name,
       quantization_config=bnb_config,
       device_map=device_map,
       torch_dtype=torch.float16,
   )
   model.config.use_cache = False
   model.config.pretraining_tp = 1
   
   # Load LLaMA tokenizer
   tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
   tokenizer.pad_token = tokenizer.eos_token
   tokenizer.padding_side = "right"
   
   # Load LoRA configuration
   peft_config = LoraConfig(
       lora_alpha=lora_alpha,
       lora_dropout=lora_dropout,
       r=lora_r,
       bias="none",
       task_type="CAUSAL_LM",
       target_modules=["q_proj", "v_proj"]
   )
   
   # Set training parameters
   training_arguments = TrainingArguments(
       output_dir=output_dir,
       num_train_epochs=num_train_epochs,
       per_device_train_batch_size=per_device_train_batch_size,
       gradient_accumulation_steps=gradient_accumulation_steps,
       optim=optim,
       save_steps=save_steps,
       logging_steps=logging_steps,
       learning_rate=learning_rate,
       weight_decay=weight_decay,
       fp16=fp16,
       bf16=bf16,
       max_grad_norm=max_grad_norm,
       max_steps=max_steps,
       warmup_ratio=warmup_ratio,
       group_by_length=group_by_length,
       lr_scheduler_type=lr_scheduler_type,
   )
   
   trainer = SFTTrainer(
       model=model,
       train_dataset=transformed,
       peft_config=peft_config,
       dataset_text_field="text",
       max_seq_length=max_seq_length,
       tokenizer=tokenizer,
       args=training_arguments,
       packing=packing,
   )
   
   trainer.train()
   
   trainer.model.save_pretrained(new_model)
   
   # Reload model in FP16 and merge it with LoRA weights
   base_model = AutoModelForCausalLM.from_pretrained(
       model_name,
       low_cpu_mem_usage=True,
       return_dict=True,
       torch_dtype=torch.float16,
       device_map=device_map,
   )
   model = PeftModel.from_pretrained(base_model, new_model)
   model = model.merge_and_unload()
   
   # push to Cloud Storage
   
   file_path_to_save_the_model = '/finetune/new_model'
   model.save_pretrained(file_path_to_save_the_model)
   tokenizer.save_pretrained(file_path_to_save_the_model)
   

3. Crea un archivo requirements.txt:

   accelerate==0.34.2
   bitsandbytes==0.45.5
   datasets==2.19.1
   transformers==4.51.3
   peft==0.11.1
   trl==0.8.6
   torch==2.3.0
   

4. Crear un Dockerfile:

   FROM nvidia/cuda:12.6.2-runtime-ubuntu22.04
   
   RUN apt-get update && \
       apt-get -y --no-install-recommends install python3-dev gcc python3-pip git && \
       rm -rf /var/lib/apt/lists/*
   
   COPY requirements.txt /requirements.txt
   
   RUN pip3 install -r requirements.txt --no-cache-dir
   
   COPY finetune.py /finetune.py
   
   ENV PYTHONUNBUFFERED 1
   
   CMD python3 /finetune.py --device cuda
   

5. Compila el contenedor en tu repositorio de Artifact Registry:

   gcloud builds submit \
     --tag $REGION-docker.pkg.dev/$PROJECT_ID/$AR_REPO/$IMAGE_NAME \
     --region $REGION
   

5. Implementa y ejecuta el trabajo

En este paso, crearás la configuración de YAML para tu trabajo con salida de VPC directa para subir archivos a Google Cloud Storage más rápido.

Ten en cuenta que este archivo contiene variables que actualizarás en un paso posterior.

1. Crea un archivo llamado finetune-job.yaml.tmpl:

   apiVersion: run.googleapis.com/v1
   kind: Job
   metadata:
     name: $JOB_NAME
     labels:
       cloud.googleapis.com/location: $REGION
     annotations:
       run.googleapis.com/launch-stage: ALPHA
   spec:
     template:
       metadata:
         annotations:
           run.googleapis.com/execution-environment: gen2
           run.googleapis.com/network-interfaces: '[{"network":"default","subnetwork":"default"}]'
       spec:
         parallelism: 1
         taskCount: 1
         template:
           spec:
             serviceAccountName: $SERVICE_ACCOUNT_ADDRESS
             containers:
             - name: $IMAGE_NAME
               image: $REGION-docker.pkg.dev/$PROJECT_ID/$AR_REPO/$IMAGE_NAME
               env:
               - name: MODEL_NAME
                 value: "google/gemma-2b"
               - name: NEW_MODEL
                 value: "gemma-2b-sql-finetuned"
               - name: BUCKET_NAME
                 value: "$BUCKET_NAME"
               - name: LORA_R
                 value: "8"
               - name: LORA_ALPHA
                 value: "16"
               - name: GRADIENT_ACCUMULATION_STEPS
                 value: "2"
               - name: DATASET_LIMIT
                 value: "1000"
               - name: LOGGING_STEPS
                 value: "5"
               - name: HF_TOKEN
                 valueFrom:
                   secretKeyRef:
                     key: 'latest'
                     name: HF_TOKEN
               resources:
                 limits:
                   cpu: 8000m
                   nvidia.com/gpu: '1'
                   memory: 32Gi
               volumeMounts:
               - mountPath: /finetune/new_model
                 name: finetuned_model
             volumes:
             - name: finetuned_model
               csi:
                 driver: gcsfuse.run.googleapis.com
                 readOnly: false
                 volumeAttributes:
                   bucketName: $BUCKET_NAME
             maxRetries: 3
             timeoutSeconds: '3600'
             nodeSelector:
               run.googleapis.com/accelerator: nvidia-l4
   

2. Ejecuta el siguiente comando para reemplazar las variables del YAML por tus variables de entorno:

   envsubst < finetune-job.yaml.tmpl > finetune-job.yaml
   

3. Crea el trabajo de Cloud Run:

   gcloud alpha run jobs replace finetune-job.yaml
   

4. Ejecuta el trabajo:

   gcloud alpha run jobs execute $JOB_NAME --region $REGION --async
   

La tarea tardará unos 10 minutos en completarse. Puedes verificar el estado con el vínculo que se proporciona en el resultado del último comando.

6. Usa un servicio de Cloud Run para publicar tu modelo ajustado con vLLM

En este paso, implementarás un servicio de Cloud Run. Esta configuración usa una VPC directa para acceder al bucket de Cloud Storage a través de una red privada para obtener descargas más rápidas.

Ten en cuenta que este archivo contiene variables que actualizarás en un paso posterior.

1. Crea un archivo service.yaml.tmpl:

   apiVersion: serving.knative.dev/v1
   kind: Service
   metadata:
     name: serve-gemma-sql
     labels:
       cloud.googleapis.com/location: $REGION
     annotations:
       run.googleapis.com/launch-stage: BETA
       run.googleapis.com/ingress: all
       run.googleapis.com/ingress-status: all
   spec:
     template:
       metadata:
         labels:
         annotations:
           autoscaling.knative.dev/maxScale: '1'
           run.googleapis.com/cpu-throttling: 'false'
           run.googleapis.com/gpu-zonal-redundancy-disabled: 'true'
           run.googleapis.com/network-interfaces: '[{"network":"default","subnetwork":"default"}]'
       spec:
         containers:
         - name: serve-finetuned
           image: us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250505_0916_RC00
           ports:
           - name: http1
             containerPort: 8000
           resources:
             limits:
               cpu: 8000m
               nvidia.com/gpu: '1'
               memory: 32Gi
           volumeMounts:
           - name: fuse
             mountPath: /finetune/new_model
           command: ["python3", "-m", "vllm.entrypoints.api_server"]
           args:
           - --model=/finetune/new_model
           - --tensor-parallel-size=1
           env:
           - name: MODEL_ID
             value: 'new_model'
           - name: HF_HUB_OFFLINE
             value: '1'
         volumes:
         - name: fuse
           csi:
             driver: gcsfuse.run.googleapis.com
             volumeAttributes:
               bucketName: $BUCKET_NAME
         nodeSelector:
           run.googleapis.com/accelerator: nvidia-l4
   

2. Actualiza el archivo service.yaml con el nombre de tu bucket.

   envsubst < service.yaml.tmpl > service.yaml
   

3. Implementa tu servicio de Cloud Run:

   gcloud alpha run services replace service.yaml
   

7. Prueba tu modelo ajustado

En este paso, le pedirás al modelo que pruebe el ajuste fino.

1. Obtén la URL de servicio de tu servicio de Cloud Run:

   SERVICE_URL=$(gcloud run services describe serve-gemma-sql --platform managed --region $REGION --format 'value(status.url)')
   

2. Crea la instrucción para tu modelo.

   USER_PROMPT="Question: What are the first name and last name of all candidates? Context: CREATE TABLE candidates (candidate_id VARCHAR); CREATE TABLE people (first_name VARCHAR, last_name VARCHAR, person_id VARCHAR)"
   

3. Llama a tu servicio con CURL para solicitarle al modelo que realice una acción:

   curl -X POST $SERVICE_URL/generate \
     -H "Content-Type: application/json" \
     -H "Authorization: bearer $(gcloud auth print-identity-token)" \
     -d @- <<EOF
   {
       "prompt": "${USER_PROMPT}"
   }
   EOF
   

Deberías ver una respuesta similar a la siguiente:

{"predictions":["Prompt:\nQuestion: What are the first name and last name of all candidates? Context: CREATE TABLE candidates (candidate_id VARCHAR); CREATE TABLE people (first_name VARCHAR, last_name VARCHAR, person_id VARCHAR)\nOutput:\n CREATE TABLE people_to_candidates (candidate_id VARCHAR, person_id VARCHAR) CREATE TABLE people_to_people (person_id VARCHAR, person_id VARCHAR) CREATE TABLE people_to_people_to_candidates (person_id VARCHAR, candidate_id"]}

8. ¡Felicitaciones!

¡Felicitaciones por completar el codelab!

Te recomendamos que revises la documentación de Cloud Run.

Temas abordados

• Cómo realizar un ajuste fino con la GPU de trabajos de Cloud Run
• Cómo entregar un modelo con Cloud Run con vLLM
• Cómo usar la configuración de VPC directa para un trabajo de GPU para subir y entregar el modelo más rápido

9. Limpia

Para evitar cargos imprevistos, por ejemplo, si los servicios de Cloud Run se invocan por error más veces que tu asignación mensual de invocaciones de Cloud Run en el nivel gratuito, puedes borrar el servicio de Cloud Run que creaste en el paso 6.

Para borrar el servicio de Cloud Run, ve a la consola de Cloud Run en https://console.cloud.google.com/run y borra el servicio serve-gemma-sql.

Para borrar todo el proyecto, ve a Administrar recursos, selecciona el proyecto que creaste en el paso 2 y elige Borrar. Si borras el proyecto, deberás cambiar los proyectos en tu SDK de Cloud. Para ver la lista de todos los proyectos disponibles, ejecuta gcloud projects list.