En este instructivo, se muestra cómo ajustar el modelo de lenguaje grande (LLM) Gemma 3 en un clúster de Slurm de varios nodos que usa dos instancias de máquina virtual (VM) A4. Como parte de este instructivo, harás lo siguiente:
Crear una imagen personalizada
Configura una red RDMA.
Ejecuta un trabajo de ajuste distribuido. Para un entrenamiento eficiente en varios nodos, usa la biblioteca Accelerate de Hugging Face con el paralelismo de datos completamente fragmentado (FSDP).
Este instructivo está dirigido a ingenieros de aprendizaje automático (AA), administradores y operadores de plataformas, y a especialistas en datos y en IA que estén interesados en usar las capacidades de programación de trabajos de Slurm para controlar cargas de trabajo de ajuste.
Objetivos
Accede a Gemma 3 con Hugging Face.
Prepara tu entorno.
Crea un clúster de Slurm A4.
Prepara tu carga de trabajo.
Ejecuta un trabajo de ajuste.
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 compute.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 compute.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/file.editor, 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, sigue estos pasos:
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, sigue estos pasos:
Clona el repositorio de GitHub de Cluster Toolkit:
git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.gitCrea un bucket de Cloud Storage:
gcloud storage buckets create gs://BUCKET_NAME \ --project=PROJECT_IDReemplaza lo siguiente:
BUCKET_NAME: Es un nombre para tu bucket de Cloud Storage que cumple con los requisitos de nombres de buckets.PROJECT_ID: Es el ID del proyecto deGoogle Cloud en el que deseas crear tu bucket de Cloud Storage.
Crea un clúster de Slurm A4
Para crear un clúster de Slurm A4, sigue estos pasos:
Ve al directorio
cluster-toolkit:cd cluster-toolkitSi es la primera vez que usas Cluster Toolkit, compila el objeto binario
gcluster:makeVe al directorio
examples/machine-learning/a4-highgpu-8g:cd examples/machine-learning/a4-highgpu-8g/Abre el archivo
a4high-slurm-deployment.yamly, luego, edítalo de la siguiente manera:terraform_backend_defaults: type: gcs configuration: bucket: BUCKET_NAME vars: deployment_name: a4-high project_id: PROJECT_ID region: REGION zone: ZONE a4h_cluster_size: 2 a4h_reservation_name: RESERVATION_URLReemplaza lo siguiente:
BUCKET_NAME: Es el nombre del bucket de Cloud Storage que creaste en la sección anterior.PROJECT_ID: Es el ID delGoogle Cloud proyecto en el que existe tu Cloud Storage y en el que deseas crear tu clúster de Slurm.REGION: Es la región en la que existe tu reserva.ZONE: Es la zona en la que existe tu reserva.RESERVATION_URL: Es la URL de la reserva que deseas usar para crear tu clúster de Slurm. Según el proyecto en el que existe la reserva, especifica uno de los siguientes valores:La reserva existe en tu proyecto:
RESERVATION_NAMELa reserva existe en otro proyecto y tu proyecto puede usarla:
projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME
Implemente el clúster:
./gcluster deploy -d examples/machine-learning/a4-highgpu-8g/a4high-slurm-deployment.yaml examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml --auto-approveEl comando
./gcluster deployes un proceso de dos fases, que se describe a continuación:En la primera fase, se compila una imagen personalizada con todo el software preinstalado, lo que puede tardar hasta 35 minutos en completarse.
En la segunda fase, se implementa el clúster con esa imagen personalizada. Este proceso debería completarse más rápido que la primera fase.
Si la primera fase se completa correctamente, pero la segunda falla, puedes intentar implementar el clúster de Slurm nuevamente omitiendo la primera fase:
./gcluster deploy -d examples/machine-learning/a4-highgpu-8g/a4high-slurm-deployment.yaml examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml --auto-approve --skip "image" -w
Prepara tu carga de trabajo
Para preparar tu carga de trabajo, sigue estos pasos:
Crea secuencias de comandos de carga de trabajo
Para crear las secuencias de comandos que usará tu carga de trabajo de ajuste, sigue estos pasos:
Para configurar el entorno virtual de Python, crea el archivo
install_environment.shcon el siguiente contenido:#!/bin/bash # This script should be run ONCE on the login node to set up the # shared Python virtual environment. set -e echo "--- Creating Python virtual environment in /home ---" python3 -m venv ~/.venv echo "--- Activating virtual environment ---" source ~/.venv/bin/activate echo "--- Installing build dependencies ---" pip install --upgrade pip wheel packaging echo "--- Installing PyTorch for CUDA 12.8 ---" pip install torch --index-url https://download.pytorch.org/whl/cu128 echo "--- Installing application requirements ---" pip install -r requirements.txt echo "--- Environment setup complete. You can now submit jobs with sbatch. ---"Para especificar la configuración de tu trabajo de ajuste, crea el archivo
accelerate_config.yamlcon el siguiente contenido:# Default configuration for a 2-node, 8-GPU-per-node (16 total GPUs) FSDP training job. compute_environment: "LOCAL_MACHINE" distributed_type: "FSDP" downcast_bf16: "no" fsdp_config: fsdp_auto_wrap_policy: "TRANSFORMER_BASED_WRAP" fsdp_backward_prefetch: "BACKWARD_PRE" fsdp_cpu_ram_efficient_loading: true fsdp_forward_prefetch: false fsdp_offload_params: false fsdp_sharding_strategy: "FULL_SHARD" fsdp_state_dict_type: "FULL_STATE_DICT" fsdp_transformer_layer_cls_to_wrap: "Gemma3DecoderLayer" fsdp_use_orig_params: true machine_rank: 0 main_training_function: "main" mixed_precision: "bf16" num_machines: 2 num_processes: 16 rdzv_backend: "static" same_network: true tpu_env: [] use_cpu: falsePara especificar las tareas que ejecutarán los trabajos en tu clúster de Slurm, crea el archivo
submit.slurmcon el siguiente contenido:#!/bin/bash #SBATCH --job-name=gemma3-finetune #SBATCH --nodes=2 #SBATCH --ntasks-per-node=8 # 8 tasks per node #SBATCH --gpus-per-task=1 # 1 GPU per task #SBATCH --partition=a4high #SBATCH --output=slurm-%j.out #SBATCH --error=slurm-%j.err set -e echo "--- Slurm Job Started ---" # --- STAGE 1: Copy Environment to Local SSD on all nodes --- srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c ' echo "Setting up local environment on $(hostname)..." LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}" LOCAL_CACHE="/mnt/localssd/hf_cache_job_${SLURM_JOB_ID}" rsync -a --info=progress2 ~/./.venv/ ${LOCAL_VENV}/ mkdir -p ${LOCAL_CACHE} echo "Setup on $(hostname) complete." ' # --- STAGE 2: Run the Training Job using the Local Environment --- echo "--- Starting Training ---" LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}" LOCAL_CACHE="/mnt/localssd/hf_cache_job_${SLURM_JOB_ID}" LOCAL_OUTPUT_DIR="/mnt/localssd/outputs_${SLURM_JOB_ID}" mkdir -p ${LOCAL_OUTPUT_DIR} # This is the main training command. srun --ntasks=$((SLURM_NNODES * 8)) --gpus-per-task=1 bash -c " source ${LOCAL_VENV}/bin/activate export HF_HOME=${LOCAL_CACHE} export HF_DATASETS_CACHE=${LOCAL_CACHE} # Run the Python script directly. # Accelerate will divide the work python ~/train.py \ --model_id google/gemma-3-12b-pt \ --output_dir ${LOCAL_OUTPUT_DIR} \ --per_device_train_batch_size 1 \ --gradient_accumulation_steps 8 \ --num_train_epochs 3 \ --learning_rate 1e-5 \ --save_strategy steps \ --save_steps 100 " # --- STAGE 3: Copy Final Model from Local SSD to Home Directory --- echo "--- Copying final model from local SSD to /home ---" # This command runs only on the first node of the job allocation # and copies the final model back to the persistent shared directory. srun --nodes=1 --ntasks=1 --ntasks-per-node=1 bash -c " rsync -a --info=progress2 ${LOCAL_OUTPUT_DIR}/ ~/gemma-12b-text-to-sql-finetuned/ " echo "--- Slurm Job Finished ---"Para especificar las dependencias de tu trabajo de ajuste, crea el archivo
requirements.txtcon el siguiente contenido:# Hugging Face Libraries (Pinned to recent, stable versions for reproducibility) 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 # Other dependencies tensorboard==2.20.0 protobuf==6.31.1 sentencepiece==0.2.0Para especificar las instrucciones de tu trabajo, crea el archivo
train.pycon el siguiente contenido:import torch import argparse 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("--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") 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 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) 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 Quantized Model and Apply PEFT --- # Define the quantization configuration quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type='nf4', bnb_4bit_compute_dtype=torch_dtype_obj, bnb_4bit_use_double_quant=True, ) config = AutoConfig.from_pretrained(args.model_id) config.use_cache = False # Load the base model with quantization print("Loading base model...") model = AutoModelForCausalLM.from_pretrained( args.model_id, config=config, quantization_config=quantization_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=True, 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=False, 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}") trainer.save_model() if __name__ == "__main__": main()
Sube secuencias de comandos al clúster de Slurm
Para subir las secuencias de comandos que creaste en la sección anterior al clúster de Slurm, sigue estos pasos:
Para identificar tu nodo de acceso, enumera todas las VMs A4 de tu proyecto:
gcloud compute instances list --filter="machineType:a4-highgpu-8g"El nombre del nodo de acceso es similar a
a4-high-login-001.Sube tus secuencias de comandos al directorio principal del nodo de acceso:
gcloud compute scp \ --project=PROJECT_ID \ --zone=ZONE \ --tunnel-through-iap \ ./train.py \ ./requirements.txt \ ./submit.slurm \ ./install_environment.sh \ ./accelerate_config.yaml \ "LOGIN_NODE_NAME":~/Reemplaza
LOGIN_NODE_NAMEpor el nombre del nodo de acceso.
Conéctate al clúster de Slurm
Conéctate al clúster de Slurm conectándote al nodo de acceso a través de SSH:
gcloud compute ssh LOGIN_NODE_NAME \
--project=PROJECT_ID \
--tunnel-through-iap \
--zone=ZONE
Instala frameworks y herramientas
Después de conectarte al nodo de acceso, instala los frameworks y las herramientas siguiendo estos pasos:
Crea una variable de entorno para tu token de acceso de Hugging Face:
export HUGGING_FACE_TOKEN="HUGGING_FACE_TOKEN"Configura un entorno virtual de Python con todas las dependencias necesarias:
chmod +x install_environment.sh ./install_environment.sh
Inicia tu carga de trabajo de ajuste
Para iniciar tu carga de trabajo de ajuste, sigue estos pasos:
Envía el trabajo al programador de Slurm:
sbatch submit.slurmEn el nodo de acceso de tu clúster de Slurm, puedes supervisar el progreso del trabajo. Para ello, verifica los archivos de salida creados en tu directorio
home:tail -f slurm-gemma3-finetune.errSi el trabajo se inicia correctamente, el archivo
.errmostrará una barra de progreso que se actualizará a medida que avance el trabajo.
Supervisa tu carga de trabajo
Puedes supervisar el uso de las GPUs en tu clúster de Slurm 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/monitoring/metrics-explorer?project=PROJECT_ID&pageState=%7B%22xyChart%22%3A%7B%22dataSets%22%3A%5B%7B%22timeSeriesFilter%22%3A%7B%22filter%22%3A%22metric.type%3D%5C%22agent.googleapis.com%2Fgpu%2Futilization%5C%22%20resource.type%3D%5C%22gce_instance%5C%22%22%2C%22perSeriesAligner%22%3A%22ALIGN_MEAN%22%7D%2C%22plotType%22%3A%22LINE%22%7D%5D%7D%7D
Cuando supervisas tu carga de trabajo, puedes ver lo siguiente:
Uso de las GPUs: Para un trabajo de ajuste fino correcto, puedes esperar ver que el uso de las 16 GPUs (ocho GPUs para cada VM del clúster) aumente y se estabilice en un nivel específico durante todo el entrenamiento.
Duración del trabajo: El trabajo debería tomar aproximadamente una hora en completarse.
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 tu clúster de Slurm
Para borrar tu clúster de Slurm, sigue estos pasos:
Ve al directorio
cluster-toolkit.Destruye el archivo de Terraform y todos los recursos creados:
./gcluster destroy a4-high --auto-approve