Usa FSDP para ajustar Llama 4 en un clúster de Slurm A4

En este instructivo, se muestra cómo ajustar un modelo de lenguaje grande (LLM) Llama-4-Scout-17 en un clúster de Slurm con varias GPU y varios nodos en Google Cloud. El clúster usa dos instancias de máquina virtual (VM) A4, cada una con 8 GPUs NVIDIA B200.

Los dos procesos principales que se describen en este instructivo son los siguientes:

  1. Implementa un clúster de Slurm de alto rendimiento y apto para la producción con elGoogle Cloud Cluster Toolkit. Como parte de esta implementación, crearás una imagen de VM personalizada con el software necesario preinstalado. También configurarás una instancia de Filestore compartida y redes RDMA de alta velocidad.
  2. 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 el paralelismo de datos completamente fragmentado (FSDP) de PyTorch, al que puedes acceder a través de la biblioteca Transformer Reinforcement Learning de Hugging Face.

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 Llama 4 con Hugging Face

  • Prepara el entorno

  • Crea e implementa un clúster de Slurm de A4 con GPU de alta calidad para producción.

  • Configura un entorno de varios nodos para el entrenamiento distribuido con FSDP.

  • Ajusta el modelo Llama 4 con trl.SFTTrainer de Hugging Face.

  • Transfiere datos a SSD locales.

  • 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.

Es posible que los usuarios nuevos de Google Cloud cumplan con los requisitos para acceder a una prueba gratuita.

Antes de comenzar

  1. 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.

  2. Instala Google Cloud CLI.

  3. Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

  4. Para inicializar gcloud CLI, ejecuta el siguiente comando: 

    gcloud init

  5. 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 permiso resourcemanager.projects.create. Obtén más información para otorgar roles.

    • Crea un proyecto de Google Cloud :

      gcloud projects create PROJECT_ID

      Reemplaza PROJECT_ID por 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_ID por el nombre de tu Google Cloud proyecto.

  6. Verifica que la facturación esté habilitada para tu proyecto de Google Cloud .

  7. 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 permiso serviceusage.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

  8. Instala Google Cloud CLI.

  9. Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

  10. Para inicializar gcloud CLI, ejecuta el siguiente comando: 

    gcloud init

  11. 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 permiso resourcemanager.projects.create. Obtén más información para otorgar roles.

    • Crea un proyecto de Google Cloud :

      gcloud projects create PROJECT_ID

      Reemplaza PROJECT_ID por 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_ID por el nombre de tu Google Cloud proyecto.

  12. Verifica que la facturación esté habilitada para tu proyecto de Google Cloud .

  13. 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 permiso serviceusage.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
    14.

  14. 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.serviceUsageAdmin 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    Reemplaza lo siguiente:

    • PROJECT_ID: ID del proyecto
    • USER_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.
  15. 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.

  16. 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
  17. Crea credenciales de autenticación locales para tu cuenta de usuario: 
    gcloud auth application-default login
  18. Habilita el Acceso al SO para tu proyecto: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. Accede a tu cuenta de Hugging Face o crea una.
  20. Instala las dependencias que necesitas para usar Cluster Toolkit.

Accede a Llama 4 con Hugging Face

Para usar Hugging Face y acceder a Llama 4, haz lo siguiente:

  1. Firma el acuerdo de consentimiento para usar Llama 4.

  2. Crea un token de acceso de Hugging Face read.

    Haz clic en Tu perfil > Configuración > Tokens de acceso > +Crear token nuevo.

  3. 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:

  1. Clona el repositorio de GitHub de Cluster Toolkit:

    git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git

  2. Crea un bucket de Cloud Storage:

    gcloud storage buckets create gs://BUCKET_NAME \
    --project=PROJECT_ID

    Reemplaza 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:

  1. Ve al directorio cluster-toolkit:

    cd cluster-toolkit

  2. Si es la primera vez que usas Cluster Toolkit, compila el objeto binario gcluster:

    make

  3. Ve al directorio examples/machine-learning/a4-highgpu-8g:

    cd examples/machine-learning/a4-highgpu-8g/

  4. Abre el archivo a4high-slurm-deployment.yaml y, 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_URL

    Reemplaza 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_NAME

      • La reserva existe en otro proyecto y tu proyecto puede usarla: projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME

  5. 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-approve

    El comando ./gcluster deploy es 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, haz lo siguiente:

  1. Crea secuencias de comandos de carga de trabajo.

  2. Sube secuencias de comandos al clúster de Slurm.

  3. Conéctate al clúster de Slurm.

  4. Instala frameworks y herramientas.

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:

  1. Para configurar el entorno virtual de Python, crea el archivo install_environment.sh con el siguiente contenido:

    #!/bin/bash
# This script sets up a consistent environment for FSDP training.
# It is meant to be run once on the login node of your Slurm cluster
set -e

# --- 1. Create the Python virtual environment ---
VENV_PATH="$HOME/.venv/venv-fsdp"
if [ ! -d "$VENV_PATH" ]; then
  echo "--- Creating Python virtual environment at $VENV_PATH ---"
  python3 -m venv $VENV_PATH
else
  echo "--- Virtual environment already exists at $VENV_PATH ---"
fi

source $VENV_PATH/bin/activate

# --- 2. Install Dependencies ---
echo "--- [STEP 2.1] Upgrading build toolchain ---"
pip install --upgrade pip wheel packaging

echo "--- [STEP 2.2] Installing PyTorch Nightly ---"
pip install --force-reinstall --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

echo "--- [STEP 2.3] Installing application dependencies ---"
if [ -f "requirements-fsdp.txt" ]; then
    pip install -r requirements-fsdp.txt
else
    echo "ERROR: requirements-fsdp.txt not found!"
    exit 1
fi

# --- 3. Download the Model ---
echo "--- [STEP 2.4] Downloading Llama4 model ---"
if [ -z "$HF_TOKEN" ]; then
  echo "ERROR: The HF_TOKEN environment variable is not set."; exit 1;
fi
pip install huggingface_hub[cli]

# Execute the CLI using its full, explicit path
$VENV_PATH/bin/huggingface-cli download meta-llama/Llama-4-Scout-17B-16E-Instruct --local-dir ~/Llama-4-Scout-17B-16E-Instruct --token $HF_TOKEN

echo "--- Environment setup complete. ---"

    Esta secuencia de comandos configura un entorno virtual de Python confiable, instala una compilación nocturna de PyTorch y descarga el modelo de Llama 4.

  2. Para especificar las dependencias de Python para la secuencia de comandos de entrenamiento, crea un archivo requirements-fsdp.txt con el siguiente contenido:

    transformers==4.55.0
datasets==4.0.0
peft==0.16.0
accelerate==1.9.0
trl==0.21.0

# Other dependencies
sentencepiece==0.2.0

  3. Especifica llama4-train-distributed.py como la secuencia de comandos de entrenamiento principal:

    import torch
from datasets import load_dataset
from peft import LoraConfig, PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    HfArgumentParser,
)

from torch.distributed import get_rank, get_world_size

from transformers.models.llama4.modeling_llama4 import Llama4TextDecoderLayer
from trl import SFTTrainer
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ScriptArguments:
    model_id: str = field(metadata={"help": "Hugging Face model ID from the Hub"})
    dataset_name: str = field(default="philschmid/gretel-synthetic-text-to-sql", metadata={"help": "Dataset from the Hub"})
    run_inference_after_training: bool = field(default=False, metadata={"help": "Run sample inference on rank 0 after training"})
    dataset_subset_size: Optional[int] = field(default=None, metadata={"help": "Number of samples to use from the dataset for training. If None, uses the full dataset."})

@dataclass
class PeftArguments:
    lora_r: int = field(default=16, metadata={"help": "LoRA attention dimension"})
    lora_alpha: int = field(default=32, metadata={"help": "LoRA alpha scaling factor"})
    lora_dropout: float = field(default=0.05, metadata={"help": "LoRA dropout probability"})

@dataclass
class SftTrainingArguments(TrainingArguments):
    max_length: Optional[int] = field(default=2048, metadata={"help": "The maximum sequence length for SFTTrainer"})
    packing: Optional[bool] = field(default=False, metadata={"help": "Enable packing for SFTTrainer"})
    ddp_find_unused_parameters: Optional[bool] = field(default=True, metadata={"help": "When using FSDP activation checkpointing, this must be set to True"})

def formatting_prompts_func(example):
    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 = f"### SCHEMA:\n{example['sql_context']}\n\n### USER QUERY:\n{example['sql_prompt']}"
    response = f"\n\n### SQL QUERY:\n{example['sql']}"
    return f"{system_message}\n\n{user_prompt}{response}"

def main():
    parser = HfArgumentParser((ScriptArguments, PeftArguments, SftTrainingArguments))
    script_args, peft_args, training_args = parser.parse_args_into_dataclasses()

    training_args.gradient_checkpointing = True
    training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}

    training_args.optim = "adamw_torch_fused"

    training_args.fsdp = "full_shard"
    training_args.fsdp_config = {
        "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
        "fsdp_transformer_layer_cls_to_wrap": [Llama4TextDecoderLayer],
        "fsdp_state_dict_type": "FULL_STATE_DICT",
        "fsdp_offload_params": False,
        "fsdp_forward_prefetch": True,
    }

    tokenizer = AutoTokenizer.from_pretrained(script_args.model_id, trust_remote_code=True)

    model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation="sdpa",
    )

    peft_config = LoraConfig(
        r=peft_args.lora_r,
        lora_alpha=peft_args.lora_alpha,
        lora_dropout=peft_args.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    )
    rank = get_rank()
    world_size = get_world_size()

    dataset = load_dataset(script_args.dataset_name, split="train")

    if script_args.dataset_subset_size is not None:
        dataset = dataset.select(range(script_args.dataset_subset_size))
    else:
        print(f"Using the full dataset with {len(dataset)} samples.")

    dataset = dataset.shuffle(seed=training_args.seed)
    print(f"Dataset shuffled with seed: {training_args.seed}.")

    if world_size > 1:
        print(f"Sharding dataset for Rank {rank} of {world_size}.")
        dataset = dataset.shard(num_shards=world_size, index=rank)

    print("Initializing SFTTrainer...")
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        peft_config=peft_config,
        formatting_func=formatting_prompts_func,
        processing_class=tokenizer,
    )

    trainer.train()

    trainer.save_model(training_args.output_dir)

    if script_args.run_inference_after_training and trainer.is_world_process_zero():
        del model
        del trainer
        torch.cuda.empty_cache()
        run_post_training_inference(script_args, training_args, tokenizer)

def run_post_training_inference(script_args, training_args, tokenizer):
    """
    Loads the fine-tuned PEFT adapter from the local output directory and runs inference.
    This should only be called on rank 0 after training is complete.
    """
    print("\n" + "="*50)
    print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
    print("="*50 + "\n")

    # Load the base model and merge the adapter.
    base_model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto"
    )
    # Load the PEFT adapter and merge it into the base model
    model = PeftModel.from_pretrained(base_model, training_args.output_dir)
    model = model.merge_and_unload() # Merge weights for faster inference
    model.eval()

    # Define the test case
    schema = "CREATE TABLE artists (Name TEXT, Country TEXT, Genre TEXT)"
    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."
    question = "Show me all artists from the Country just north of the USA."

    # This must match the formatting_func exactly
    prompt = f"{system_message}\n\n### SCHEMA:\n{schema}\n\n### USER QUERY:\n{question}\n\n### SQL QUERY:\n"

    print(f"Test Prompt:\n{prompt}")

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    print("\n--- Generating SQL... ---")
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=False,
        temperature=None,
        top_p=None,
    )

    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):].strip()

    print(f"\n--- Generated SQL Query ---")
    print(generated_sql)
    print("\n" + "="*50)
    print("=== INFERENCE TEST COMPLETE ===")
    print("="*50 + "\n")

if __name__ == "__main__":
    main()

    Esta secuencia de comandos utiliza el entrenador de ajuste fino supervisado (SFT) de TRL para administrar los bucles de entrenamiento de FSDP, la configuración de adaptación de clasificación baja (LoRA) y el formato de datos.

  4. Para especificar las tareas que ejecutarán los trabajos en tu clúster de Slurm, crea el archivo submit.slurm con el siguiente contenido:

    #!/bin/bash
#SBATCH --job-name=llama4-fsdp-fixed
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --partition=a4high
#SBATCH --output=llama4-%j.out
#SBATCH --error=llama4-%j.err

set -e
set -x

echo "--- Slurm Job Started ---"
echo "Job ID: $SLURM_JOB_ID"
echo "Node List: $SLURM_JOB_NODELIST"

# --- Define Paths ---
LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
VENV_PATH="${HOME}/.venv/venv-fsdp"
MODEL_PATH="${HOME}/Llama-4-Scout-17B-16E-Instruct"

# --- STAGE 1: Stage Data to Local SSD on Each Node ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c "
  echo '--- Staging on node: $(hostname) ---'

  mkdir -p ${LOCAL_SSD_PATH}

  echo 'Copying virtual environment...'
  rsync -a -q ${VENV_PATH}/ ${LOCAL_SSD_PATH}/venv/

  echo 'Copying model weights...'
  rsync -a --info=progress2 ${MODEL_PATH}/ ${LOCAL_SSD_PATH}/model/

  mkdir -p ${LOCAL_SSD_PATH}/hf_cache

  echo '--- Staging on $(hostname) complete ---'
"
echo "--- Staging complete on all nodes ---"

# --- STAGE 2: Run the Training Job ---
echo "--- Launching Distributed Training with GIB NCCL Plugin ---"
nodes=( $( scontrol show hostnames "$SLURM_JOB_NODELIST" ) )
head_node=${nodes[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

export MASTER_ADDR=$head_node_ip
export MASTER_PORT=29500

export NCCL_SOCKET_IFNAME=enp0s19
export NCCL_NET=gIB

# export NCCL_DEBUG=INFO # Un-comment to diagnose NCCL issues if needed

srun --cpu-bind=none --accel-bind=g bash -c '
  # Activate the environment from the local copy
  source '${LOCAL_SSD_PATH}'/venv/bin/activate

  # Point Hugging Face cache to the local SSD
  export HF_HOME='${LOCAL_SSD_PATH}'/hf_cache

  export RANK=$SLURM_PROCID
  export WORLD_SIZE=$SLURM_NTASKS
  export LOCAL_RANK=$SLURM_LOCALID

  export LD_LIBRARY_PATH=/usr/local/gib/lib64:$LD_LIBRARY_PATH
  source /usr/local/gib/scripts/set_nccl_env.sh

  # --- Launch the training ---
  python \
    '${SLURM_SUBMIT_DIR}'/llama4-train-distributed.py \
      --model_id="'${LOCAL_SSD_PATH}'/model/" \
      --output_dir="'${LOCAL_SSD_PATH}'/outputs/" \
      --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
      --seed=900913 \
      --bf16=True \
      --num_train_epochs=1 \
      --per_device_train_batch_size=2 \
      --gradient_accumulation_steps=4 \
      --learning_rate=2e-5 \
      --logging_steps=10 \
      --lora_r=16 \
      --lora_alpha=32 \
      --lora_dropout=0.05 \
      --run_inference_after_training
'

# --- STAGE 3: Copy Final Results Back to Persistent Storage ---
echo "--- Copying final results from local SSD to shared storage ---"
PERSISTENT_OUTPUT_DIR="${HOME}/outputs/llama4_job_${SLURM_JOB_ID}"
mkdir -p "$PERSISTENT_OUTPUT_DIR"

# Only copy from the head node where trl has combined the results
srun --nodes=1 --ntasks=1 -w "$head_node" \
  rsync -a --info=progress2 "${LOCAL_SSD_PATH}/outputs/" "${PERSISTENT_OUTPUT_DIR}/"

# --- STAGE 4: Cleanup ---
echo "--- Cleaning up local SSD on all nodes ---"
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c "rm -rf ${LOCAL_SSD_PATH}"

echo "--- Slurm Job Finished ---"

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:

  1. 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.

  2. Sube tus secuencias de comandos al directorio principal del nodo de acceso:

    gcloud compute scp --project="$PROJECT_ID" --zone="$ZONE" --tunnel-through-iap \
  ./install_environment.sh \
  ./requirements-fsdp.txt \
  ./llama4-train-distributed.py \
  ./submit.slurm \
  "${LOGIN_NODE_NAME}":~/

    Reemplaza LOGIN_NODE_NAME por 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 de la siguiente manera:

  1. Exporta tu token de Hugging Face:

    # On the login node
export HF_TOKEN="hf_..." # Replace with your token

  2. Ejecuta la secuencia de comandos de instalación:

    # On the login node
chmod +x install_environment.sh
./install_environment.sh

    Este comando configura un entorno virtual con todas las dependencias necesarias y descarga los pesos del modelo en el archivo ~/Llama-4-Scout-17B-16E-Instruct.

    Dado que la descarga del modelo es muy grande (alrededor de 200 GB), este proceso tarda alrededor de 30 minutos, según las condiciones de la red.

Inicia tu carga de trabajo de ajuste

Para comenzar a entrenar tu carga de trabajo, haz lo siguiente:

  1. Envía el trabajo al programador de Slurm:

    sbatch submit.slurm

  2. En 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:

    # On the login node
tail -f llama4-*.out

    Si el trabajo se inicia correctamente, el archivo .err mostrará una barra de progreso que se actualizará a medida que avance el trabajo.

    Este trabajo debería tardar un poco más de una hora en completarse en el clúster de Slurm. El trabajo tiene dos fases principales:

    • Copiar el modelo base grande en la SSD local de cada nodo de procesamiento
    • Es el trabajo de entrenamiento, que comienza una vez que se completa la copia del modelo. Este trabajo tarda unos 35 minutos en ejecutarse.

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:

  1. Ve al directorio cluster-toolkit.

  2. Destruye el archivo de Terraform y todos los recursos creados:

    ./gcluster destroy a4-high --auto-approve

Borra tu instancia de Filestore

De forma predeterminada, tu instancia de Filestore tiene el parámetro de configuración deletion_protection establecido en verdadero en el plano cluster-toolkit. Este parámetro de configuración evita la pérdida accidental de datos cuando modificas entornos. Para borrar la instancia de Filestore, debes inhabilitar manualmente la protección contra eliminación.

¿Qué sigue?