En este instructivo, se muestra cómo ajustar un modelo de lenguaje grande (LLM) Gemma 3 en un clúster de GKE con varios nodos y varias GPUs en Google Cloud. Este clúster usa una instancia de máquina virtual (VM) A4 que tiene 8 GPUs NVIDIA B200.
Los dos procesos principales que se describen en este instructivo son los siguientes:
- Implementa un clúster de GKE de alto rendimiento con GKE Autopilot. Como parte de esta implementación, crearás una imagen de VM personalizada con el software necesario preinstalado.
- Después de que se implemente el clúster, ejecutarás un trabajo de ajuste distribuido con el conjunto de secuencias de comandos que acompañan a este instructivo. El trabajo aprovecha la biblioteca Accelerate de Hugging Face.
Este instructivo está dirigido a ingenieros, investigadores, administradores y operadores de plataformas de aprendizaje automático (AA), y a especialistas en datos y en IA que deseen implementar clústeres de GKE en Google Cloud para entrenar LLMs.
Objetivos
Accede al modelo de Gemma 3 con Hugging Face.
Prepara tu entorno.
Crea e implementa un clúster de GKE de A4.
Ajusta el modelo Gemma 3 con la biblioteca Accelerate de Hugging Face y el paralelismo de datos completamente fragmentado (FSDP).
Supervisar tu trabajo
Realizar una limpieza
Costos
En este documento, usarás los siguientes componentes facturables de Google Cloud:
Para generar una estimación de costos en función del uso previsto,
usa la calculadora de precios.
Antes de comenzar
- Accede a tu cuenta de Google Cloud . Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.
-
Instala Google Cloud CLI.
-
Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.
-
Para inicializar gcloud CLI, ejecuta el siguiente comando:
gcloud init -
Crea o selecciona un Google Cloud proyecto.
Roles necesarios para seleccionar o crear un proyecto
- Selecciona un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
-
Crear un proyecto: Para crear un proyecto, necesitas el rol de Creador de proyectos (
roles/resourcemanager.projectCreator), que contiene el permisoresourcemanager.projects.create. Obtén más información para otorgar roles.
-
Crea un proyecto de Google Cloud :
gcloud projects create PROJECT_ID
Reemplaza
PROJECT_IDpor un nombre para el proyecto Google Cloud que estás creando. -
Selecciona el proyecto Google Cloud que creaste:
gcloud config set project PROJECT_ID
Reemplaza
PROJECT_IDpor el nombre de tu Google Cloud proyecto.
-
Verifica que la facturación esté habilitada para tu proyecto de Google Cloud .
Habilita la API necesaria:
Roles necesarios para habilitar las APIs
Para habilitar las APIs, necesitas el rol de IAM de administrador de Service Usage (
roles/serviceusage.serviceUsageAdmin), que contiene el permisoserviceusage.services.enable. Obtén más información para otorgar roles.gcloud services enable gcloud services enable compute.googleapis.com container.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
-
Instala Google Cloud CLI.
-
Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.
-
Para inicializar gcloud CLI, ejecuta el siguiente comando:
gcloud init -
Crea o selecciona un Google Cloud proyecto.
Roles necesarios para seleccionar o crear un proyecto
- Selecciona un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
-
Crear un proyecto: Para crear un proyecto, necesitas el rol de Creador de proyectos (
roles/resourcemanager.projectCreator), que contiene el permisoresourcemanager.projects.create. Obtén más información para otorgar roles.
-
Crea un proyecto de Google Cloud :
gcloud projects create PROJECT_ID
Reemplaza
PROJECT_IDpor un nombre para el proyecto Google Cloud que estás creando. -
Selecciona el proyecto Google Cloud que creaste:
gcloud config set project PROJECT_ID
Reemplaza
PROJECT_IDpor el nombre de tu Google Cloud proyecto.
-
Verifica que la facturación esté habilitada para tu proyecto de Google Cloud .
Habilita la API necesaria:
Roles necesarios para habilitar las APIs
Para habilitar las APIs, necesitas el rol de IAM de administrador de Service Usage (
roles/serviceusage.serviceUsageAdmin), que contiene el permisoserviceusage.services.enable. Obtén más información para otorgar roles.gcloud services enable gcloud services enable compute.googleapis.com container.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
-
Otorga roles a tu cuenta de usuario. Ejecuta el siguiente comando una vez para cada uno de los siguientes roles de IAM:
roles/compute.admin, roles/iam.serviceAccountUser, roles/cloudbuild.builds.editor, roles/artifactregistry.admin, roles/storage.admin, roles/serviceusage.serviceUsageAdmingcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
Reemplaza lo siguiente:
PROJECT_ID: ID del proyectoUSER_IDENTIFIER: Es el identificador de tu cuenta de usuario de . Por ejemplo,myemail@example.com.ROLE: Es el rol de IAM que otorgas a tu cuenta de usuario.
- Habilita la cuenta de servicio predeterminada para tu proyecto Google Cloud :
gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \ --project=PROJECT_ID
Reemplaza PROJECT_NUMBER por el número del proyecto. Para revisar el número de tu proyecto, consulta Cómo obtener un proyecto existente.
- Otorga el rol de editor (
roles/editor) a la cuenta de servicio predeterminada:gcloud projects add-iam-policy-binding PROJECT_ID \ --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \ --role=roles/editor
- Crea credenciales de autenticación locales para tu cuenta de usuario:
gcloud auth application-default login
- Habilita el Acceso al SO para tu proyecto:
gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
- Accede a tu cuenta de Hugging Face o crea una.
Accede a Gemma 3 con Hugging Face
Para usar Hugging Face y acceder a Gemma 3, haz lo siguiente:
- Accede a Hugging Face
- Crea un token de acceso de Hugging Face
read.
Haz clic en Tu perfil > Configuración > Tokens de acceso > +Crear token nuevo. - Copia y guarda el valor del token
read access. La usarás más adelante en este instructivo.
Prepara el entorno
Para preparar tu entorno, configura lo siguiente:
gcloud config set project PROJECT_NAME
gcloud config set billing/quota_project PROJECT_NAME
export RESERVATION=YOUR_RESERVATION_ID
export PROJECT_ID=$(gcloud config get project)
export REGION=CLUSTER_REGION
export CLUSTER_NAME=CLUSTER_NAME
export HF_TOKEN=YOUR_TOKEN
export NETWORK=default
Reemplaza lo siguiente:
PROJECT_NAME: Es el nombre del Google Cloud proyecto en el que deseas crear el clúster de GKE.YOUR_RESERVATION_ID: Es el identificador de tu capacidad reservada.CLUSTER_REGION: Es la región en la que deseas crear tu clúster de GKE. Solo puedes crear el clúster en la región en la que existe tu reserva.CLUSTER_NAME: Es el nombre del clúster de GKE que se creará.HF_TOKEN: El token de acceso de Hugging Face que creaste en la sección anterior.
Crea un clúster de GKE en modo Autopilot
Para crear un clúster de GKE en modo Autopilot, ejecuta el siguiente comando:
gcloud container clusters create-auto ${CLUSTER_NAME} \
--project=${PROJECT_ID} \
--location=${REGION} \
--release-channel=rapid
La creación del clúster de GKE puede tardar un tiempo en completarse. Para verificar que Google Cloud haya terminado de crear tu clúster, ve aClústeres de Kubernetesen la consola de Google Cloud .
Crea un secreto de Kubernetes para las credenciales de Hugging Face
Para crear un secreto de Kubernetes para las credenciales de Hugging Face, sigue estos pasos:
Configura
kubectlpara comunicarse con tu clúster de GKE:gcloud container clusters get-credentials $CLUSTER_NAME \ --location=$REGIONCrea un secreto de Kubernetes para almacenar tu token de Hugging Face:
gcloud container clusters get-credentials ${CLUSTER_NAME} \ --location=${REGION} kubectl create secret generic hf-secret \ --from-literal=hf_api_token=${HF_TOKEN} \ --dry-run=client -o yaml | kubectl apply -f -
Prepara tu carga de trabajo
Para preparar tu carga de trabajo, haz lo siguiente:
Crea secuencias de comandos de carga de trabajo
Para crear las secuencias de comandos que usa tu carga de trabajo de ajuste, haz lo siguiente:
Crea un directorio para las secuencias de comandos de la carga de trabajo. Usa este directorio como tu directorio de trabajo.
mkdir llm-finetuning-gemma cd llm-finetuning-gemmaCrea el archivo
cloudbuild.yamlpara usar Google Cloud Build. Este archivo crea tu contenedor de cargas de trabajo y lo almacena en Artifact Registry:steps: - name: 'gcr.io/cloud-builders/docker' args: [ 'build', '-t', 'us-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0', '.' ] images: - 'us-docker.pkg.dev/$PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0'Crea un archivo
Dockerfilepara ejecutar el trabajo de ajuste:FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 RUN apt-get update && \ apt-get -y install python3 python3-dev gcc python3-pip python3-venv git curl vim RUN python3 -m venv /opt/venv ENV PATH="/opt/venv/bin:/usr/local/nvidia/bin:$PATH" ENV LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$LD_LIBRARY_PATH" RUN pip3 install setuptools wheel packaging ninja RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 RUN pip3 install \ transformers==4.53.3 \ datasets==4.0.0 \ accelerate==1.9.0 \ evaluate==0.4.5 \ bitsandbytes==0.46.1 \ trl==0.19.1 \ peft==0.16.0 \ tensorboard==2.20.0 \ protobuf==6.31.1 \ sentencepiece==0.2.0 COPY finetune.py /finetune.py COPY accel_fsdp_gemma3_config.yaml /accel_fsdp_gemma3_config.yaml CMD accelerate launch --config_file accel_fsdp_gemma3_config.yaml finetune.pyCrea el archivo
accel_fsdp_gemma3_config.yaml. Este archivo de configuración dirige Hugging Face Accelerate para dividir el trabajo de ajuste en varias GPUs.compute_environment: LOCAL_MACHINE debug: false distributed_type: FSDP downcast_bf16: 'no' enable_cpu_affinity: false fsdp_config: fsdp_activation_checkpointing: false fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP fsdp_cpu_ram_efficient_loading: true fsdp_offload_params: false fsdp_reshard_after_forward: true fsdp_state_dict_type: FULL_STATE_DICT fsdp_transformer_layer_cls_to_wrap: Gemma3DecoderLayer fsdp_version: 2 machine_rank: 0 main_training_function: main mixed_precision: bf16 num_machines: 1 num_processes: 8 rdzv_backend: static same_network: true tpu_env: [] tpu_use_cluster: false tpu_use_sudo: false use_cpu: falseCrea el archivo
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 command: ["accelerate","launch"] args: - "--config_file" - "accel_fsdp_gemma3_config.yaml" - "finetune.py" - "--model_id" - "google/gemma-3-12b-pt" - "--output_dir" - "gemma-12b-text-to-sql" - "--per_device_train_batch_size" - "8" - "--gradient_accumulation_steps" - "8" - "--num_train_epochs" - "3" - "--learning_rate" - "1e-5" - "--save_strategy" - "steps" - "--save_steps" - "100" resources: limits: nvidia.com/gpu: "8" env: - 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-b200 cloud.google.com/reservation-name: $RESERVATION cloud.google.com/reservation-affinity: "specific" cloud.google.com/gke-gpu-driver-version: latest restartPolicy: OnFailureCrea el archivo
finetune.py:import torch import argparse import subprocess from datasets import load_dataset from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model from trl import SFTTrainer, SFTConfig from huggingface_hub import login def get_args(): parser = argparse.ArgumentParser() parser.add_argument("--model_id", type=str, default="google/gemma-3-12b-pt", help="Hugging Face model ID") parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face token for private models") parser.add_argument("--trust_remote", type=bool, default="False", help="Trust remote code when loading tokenizer") parser.add_argument("--use_fast", type=bool, default="True", help="Determines if a fast Rust-based tokenizer should be used") parser.add_argument("--dataset_name", type=str, default="philschmid/gretel-synthetic-text-to-sql", help="Hugging Face dataset name") parser.add_argument("--output_dir", type=str, default="gemma-12b-text-to-sql", help="Directory to save model checkpoints") # LoRA arguments parser.add_argument("--lora_r", type=int, default=16, help="LoRA attention dimension") parser.add_argument("--lora_alpha", type=int, default=16, help="LoRA alpha scaling factor") parser.add_argument("--lora_dropout", type=float, default=0.05, help="LoRA dropout probability") # SFTConfig arguments parser.add_argument("--max_seq_length", type=int, default=512, help="Maximum sequence length") parser.add_argument("--num_train_epochs", type=int, default=3, help="Number of training epochs") parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size per device during training") parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps") parser.add_argument("--learning_rate", type=float, default=1e-5, help="Learning rate") parser.add_argument("--logging_steps", type=int, default=10, help="Log every X steps") parser.add_argument("--save_strategy", type=str, default="steps", help="Checkpoint save strategy") parser.add_argument("--save_steps", type=int, default=100, help="Save checkpoint every X steps") parser.add_argument("--push_to_hub", action='store_true', help="Push model back up to HF") parser.add_argument("--hub_private_repo", type=bool, default="True", help="Push to a private repo") return parser.parse_args() def main(): args = get_args() # --- 1. Setup and Login --- if args.hf_token: login(args.hf_token) # --- 2. Create and prepare the fine-tuning dataset --- # The `create_conversation` function is no longer needed. # The SFTTrainer will use the `formatting_func` to apply the chat template. dataset = load_dataset(args.dataset_name, split="train") dataset = dataset.shuffle().select(range(12500)) dataset = dataset.train_test_split(test_size=2500/12500) # --- 3. Configure Model and Tokenizer --- if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8: torch_dtype_obj = torch.bfloat16 torch_dtype_str = "bfloat16" else: torch_dtype_obj = torch.float16 torch_dtype_str = "float16" tokenizer = AutoTokenizer.from_pretrained(args.model_id, trust_remote_code=args.trust_remote, use_fast=args.use_fast) tokenizer.pad_token = tokenizer.eos_token gemma_chat_template = ( "" "" ) tokenizer.chat_template = gemma_chat_template # --- 4. Define the Formatting Function --- # This function will be used by the SFTTrainer to format each sample # from the dataset into the correct chat template format. def formatting_func(example): # The create_conversation logic is now implicitly handled by this. # We need to construct the messages list here. system_message = "You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA." user_prompt = "Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.\n\n<SCHEMA>\n{context}\n</SCHEMA>\n\n<USER_QUERY>\n{question}\n</USER_QUERY>\n" messages = [ {"role": "user", "content": user_prompt.format(question=example["sql_prompt"][0], context=example["sql_context"][0])}, {"role": "assistant", "content": example["sql"][0]} ] return tokenizer.apply_chat_template(messages, tokenize=False) # --- 5. Load Model and Apply PEFT --- config = AutoConfig.from_pretrained(args.model_id) config.use_cache = False # We'll be loading this model full precision because we're planning to do FSDP # Load the base model with quantization print("Loading base model...") model = AutoModelForCausalLM.from_pretrained( args.model_id, config=config, attn_implementation="eager", torch_dtype=torch_dtype_obj, ) # Prepare the model for k-bit training model = prepare_model_for_kbit_training(model) # Configure LoRA. peft_config = LoraConfig( lora_alpha=args.lora_alpha, lora_dropout=args.lora_dropout, r=args.lora_r, bias="none", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], task_type="CAUSAL_LM", ) # Apply the PEFT config to the model print("Applying PEFT configuration...") model = get_peft_model(model, peft_config) model.print_trainable_parameters() # --- 6. Configure Training Arguments --- training_args = SFTConfig( output_dir=args.output_dir, max_seq_length=args.max_seq_length, num_train_epochs=args.num_train_epochs, per_device_train_batch_size=args.per_device_train_batch_size, gradient_accumulation_steps=args.gradient_accumulation_steps, learning_rate=args.learning_rate, logging_steps=args.logging_steps, save_strategy=args.save_strategy, save_steps=args.save_steps, packing=False, label_names=["domain"], gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, optim="adamw_torch", fp16=True if torch_dtype_obj == torch.float16 else False, bf16=True if torch_dtype_obj == torch.bfloat16 else False, max_grad_norm=0.3, warmup_ratio=0.03, lr_scheduler_type="constant", push_to_hub=True, report_to="tensorboard", dataset_kwargs={ "add_special_tokens": False, "append_concat_token": True, } ) # --- 7. Create Trainer and Start Training --- trainer = SFTTrainer( model=model, args=training_args, train_dataset=dataset["train"], eval_dataset=dataset["test"], formatting_func=formatting_func, ) print("Starting training...") trainer.train() print("Training finished.") # --- 8. Save the final model --- print(f"Saving final model to {args.output_dir}") model.cpu() trainer.save_model(args.output_dir) torch.distributed.destroy_process_group() if __name__ == "__main__": main()
Usa Docker y Cloud Build para crear un contenedor de ajuste
Crea un repositorio de Docker de Artifact Registry:
gcloud artifacts repositories create gemma \ --project=${PROJECT_ID} \ --repository-format=docker \ --location=us \ --description="Gemma Repo"En el directorio
llm-finetuning-gemmaque creaste en un paso anterior, ejecuta el siguiente comando para crear el contenedor de ajuste y enviarlo a Artifact Registry.gcloud builds submit .Exporta la URL de la imagen. Lo usarás en un paso posterior de este instructivo:
export IMAGE_URL=us-docker.pkg.dev/${PROJECT_ID}/gemma/finetune-gemma-gpu:1.0.0
Inicia tu carga de trabajo de ajuste
Para iniciar tu carga de trabajo de ajuste, haz lo siguiente:
Aplica el manifiesto de ajuste para crear el trabajo de ajuste:
envsubst < finetune.yaml | kubectl apply -f -Como usas clústeres en el modo Autopilot de GKE, es posible que el inicio del nodo habilitado para GPU tarde unos minutos.
Supervisa el trabajo ejecutando el siguiente comando:
ewatch kubectl get podsEjecuta el siguiente comando para verificar los registros del trabajo:
kubectl logs job.batch/finetune-job -fEl recurso de trabajo descarga los datos del modelo y, luego, ajusta el modelo en las ocho GPUs. La descarga tarda alrededor de cinco minutos en completarse. Una vez que se completa la descarga, el proceso de ajuste fino tarda aproximadamente dos horas y 30 minutos en completarse.
Supervisa tu carga de trabajo
Puedes supervisar el uso de las GPUs en tu clúster de GKE para verificar que tu trabajo de ajuste fino se ejecute de manera eficiente. Para ello, abre el siguiente vínculo en tu navegador:
https://console.cloud.google.com/kubernetes/clusters/details/us-central1/[CLUSTER_NAME]/observability?mods=monitoring_api_prod&project=[YOUR_PROJECT_ID]]&pageState=("timeRange":("duration":"PT1H"),"nav":("section":"gpu"),"groupBy":("groupByType":"namespacesTop5"))
Cuando supervisas tu carga de trabajo, puedes ver lo siguiente:
- Uso de GPU: Para un trabajo de ajuste fino correcto, puedes esperar ver que el uso de las 8 GPU aumente y se estabilice en un nivel alto durante todo el entrenamiento.
- Duración del trabajo: El trabajo debería tardar aproximadamente 10 minutos en completarse en el clúster A4 especificado.
Realiza una limpieza
Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.
Borra tu proyecto
Borra un Google Cloud proyecto:
gcloud projects delete PROJECT_ID
Borra tus recursos
Para borrar tu trabajo de ajuste, ejecuta el siguiente comando:
kubectl delete job finetune-jobPara borrar tu clúster de GKE, ejecuta el siguiente comando:
gcloud container clusters delete $CLUSTER_NAME \ --region=$REGION