Entrena Qwen2 en un clúster de Slurm con A4

En este instructivo, se muestra cómo entrenar un modelo de lenguaje grande (LLM) en un clúster de Slurm de varios nodos y varias GPUs en Google Cloud. El modelo que usas en este instructivo se basa en un modelo de 1,500 millones de parámetros de Qwen2. El clúster de Slurm usa dos máquinas virtuales (VM) a4-highgpu-8g, 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 conGoogle 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 implementar el clúster, ejecuta un trabajo de entrenamiento previo 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 e investigadores de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en datos y en IA que deseen implementar clústeres de Slurm de alto rendimiento en Google Cloud para entrenar LLMs.

Objetivos

  • Accede al modelo Qwen2 con Hugging Face.
  • Prepara tu entorno.
  • Crea e implementa un clúster de Slurm A4 apto para producción.
  • Entrena el modelo Qwen2 con la biblioteca Accelerate .
  • 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 Google Cloud proyecto:

      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 proyecto de Google Cloud .

  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 APIs, necesitas el permiso serviceusage.services.enable. Si creaste el proyecto, es probable que ya tengas este permiso a través del rol de propietario (roles/owner). De lo contrario, puedes obtener este permiso a través del rol de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin). 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 Google Cloud proyecto:

      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 proyecto de Google Cloud .

  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 APIs, necesitas el permiso serviceusage.services.enable. Si creaste el proyecto, es probable que ya tengas este permiso a través del rol de propietario (roles/owner). De lo contrario, puedes obtener este permiso a través del rol de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin). 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, roles/compute.osAdminLogin, roles/iap.tunnelResourceAccessor 

    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.

Accede a Qwen2 con Hugging Face

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

  1. Firma el acuerdo de consentimiento para usar Qwen 2 1.5B.

  2. Crea un token de acceso de read.

Instala Cluster Toolkit

Cluster Toolkit es una herramienta de código abierto que simplifica la implementación de cargas de trabajo de computación de alto rendimiento (HPC), inteligencia artificial (IA) y aprendizaje automático (AA) en Google Cloud. Para obtener más información sobre el uso de gcluster y la administración de clústeres, consulta la descripción general del Cluster Toolkit.

  1. Prepara la versión de Cluster Toolkit:

    export CLUSTER_TOOLKIT_TAG=v1.97.0

# Detect OS (linux or mac)
case "$(uname -s)" in
  Linux*)     OS="linux" ;;
  Darwin*)    OS="mac" ;;
  *)          echo "Error: Unsupported operating system: $(uname -s)" >&2; exit 1 ;;
esac

# Detect Architecture (amd64 or arm64)
case "$(uname -m)" in
  x86_64)     ARCH="amd64" ;;
  aarch64|arm64) ARCH="arm64" ;;
  *)          echo "Error: Unsupported architecture: $(uname -m)" >&2; exit 1 ;;
esac

  2. Descarga la versión:

    # Download and extract the platform-specific bundle
curl -LO "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_TAG}/gcluster_bundle_${OS}_${ARCH}.zip"
unzip "gcluster_bundle_${OS}_${ARCH}.zip" -d cluster-toolkit/
rm -f "gcluster_bundle_${OS}_${ARCH}.zip"

  3. Define la ruta gcluster:

    export CLUSTER_TOOLKIT_PATH="$(pwd)/cluster-toolkit"
export PATH="${CLUSTER_TOOLKIT_PATH}:${PATH}"
gcluster --version

Prepara el entorno

Para preparar tu entorno, sigue estos pasos:

  1. Configura las variables de entorno predeterminadas: 

    export PROJECT_ID="YOUR_PROJECT_ID"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_URL="RESERVATION_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export HF_TOKEN="HUGGING_FACE_TOKEN"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

    Reemplaza lo siguiente:

    • YOUR_PROJECT_ID: Es el nombre del Google Cloud proyecto en el que deseas crear el clúster de GKE.
    • YOUR_CLUSTER_NAME: Es el nombre del clúster de Slurm que deseas crear.
    • YOUR_ZONE: Es la zona en la que existe tu reserva.
    • YOUR_REGION,: Es la región en la que existe tu reserva.
    • RESERVATION_NAME: Es la URL o el nombre de la reserva que deseas usar para crear tu clúster de Slurm.
    • YOUR_BUCKET_NAME: Es el nombre del bucket en el que almacenas los resultados del punto de control del entrenamiento. Antes de crearlo, familiarízate con los requisitos de nombres de buckets.
    • HUGGING_FACE_TOKEN: El token de Hugging Face que creaste en un paso anterior

  2. Crea un bucket de Cloud Storage:

    gcloud storage buckets create "gs://${GCS_BUCKET}" \
  --project="${PROJECT_ID}"

Crea un clúster de Slurm A4

Para crear un clúster de Slurm A4, sigue estos pasos:

  1. Crea el archivo a4high-slurm-deployment.yaml:

    MANIFEST_PATH="${CLUSTER_TOOLKIT_PATH}/examples/machine-learning/a4-highgpu-8g"
cat <<EOF > "${MANIFEST_PATH}/a4high-slurm-deployment.yaml"
terraform_backend_defaults:
  type: gcs
  configuration:
    bucket: ${GCS_BUCKET}

vars:
  deployment_name: ${CLUSTER_NAME}
  project_id: ${PROJECT_ID}
  region: ${REGION}
  zone: ${ZONE}
  a4h_cluster_size: 2
  a4h_reservation_name: ${RESERVATION_URL}
EOF

  2. Crea los manifiestos de Terraform:

    gcluster create \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  -o "${CLUSTER_NAME}"

  3. Aplica parches a los manifiestos:

    sed -i '/deletion_protection = {/,/}/ { s/enabled = true/enabled = false/; /reason  = "Avoid data loss"/d; }' "${CLUSTER_NAME}/${CLUSTER_NAME}/cluster-env/main.tf"

  4. Implemente el clúster:

    gcluster deploy "${CLUSTER_NAME}/${CLUSTER_NAME}" --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 50 minutos en completarse.

    • En la segunda fase, se implementa el clúster con esa imagen personalizada. Por lo general, este proceso tarda menos en completarse 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 "${CLUSTER_NAME}" --auto-approve --skip "image" -w

Prepara tu carga de trabajo

Para preparar tu carga de trabajo, sigue estos pasos:

  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 entrenamiento, 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 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. ---"

  2. Para especificar la configuración de tu trabajo de ajuste, crea el archivo accelerate_config.yaml con el siguiente contenido:

    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: "SHARDED_STATE_DICT"
  fsdp_transformer_layer_cls_to_wrap: "Qwen2DecoderLayer"
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: false

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

    #SBATCH --job-name=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1  # 1 main srun task on the node that manages accelerate
#SBATCH --gpus-per-node=8    # access to all 8 GPUs on the node
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -euo pipefail
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none bash -c '
  set -e
  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}"
  PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
  LOCAL_TMP="/mnt/localssd/tmp_job_${SLURM_JOB_ID}"

  rsync -a --info=progress2 ~/.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR} ${LOCAL_TMP}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export TMPDIR="${LOCAL_TMP}"
  export TEMP="${LOCAL_TMP}"
  export TMP="${LOCAL_TMP}"
  export HF_HOME="${LOCAL_CACHE}"
  export HF_DATASETS_CACHE="${LOCAL_CACHE}"
  export HF_MODULES_CACHE="${LOCAL_CACHE}/modules"
  export HF_METRICS_CACHE="${LOCAL_CACHE}/metrics"

  python "${HOME}/preprocess_data.py" \
    --dataset_name "HuggingFaceFW/fineweb-edu" \
    --dataset_config "CC-MAIN-2024-10" \
    --tokenizer_id "Qwen/Qwen2-1.5B" \
    --max_seq_length 1024 \
    --output_path ${PROCESSED_DATA_DIR}

  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}"
PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
LOCAL_OUTPUT_DIR="/mnt/localssd/outputs_${SLURM_JOB_ID}"

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29505

# Network and initialization debugging configurations
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=INIT,COLL
export TORCH_DISTRIBUTED_DEBUG=INFO
export NCCL_IB_DISABLE=0

# Launching with full GPU access enabled for accelerate
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none bash -c "
  mkdir -p ${LOCAL_OUTPUT_DIR}
  source ${LOCAL_VENV}/bin/activate

  # Reset the forced Slurm isolation to expose all 8 GPUs to the processes
  unset CUDA_VISIBLE_DEVICES

  # Retrieve the default network device.
  DETECTED_IFACE=\$(ip route show | grep default | awk '{print \$5}' | head -n 1)
  echo \"[INFO] Automatically detected cluster network interface: \$DETECTED_IFACE\"

  # Dynamic injection of the detected interface into the network configuration
  export NCCL_SOCKET_IFNAME=\"\${DETECTED_IFACE},gpu*\"
  export TP_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"
  export GLOO_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"

  accelerate launch \
    --config_file ~/accelerate_config.yaml \
    --num_machines \$SLURM_NNODES \
    --num_processes \$((SLURM_NNODES * 8)) \
    --machine_rank \$SLURM_NODEID \
    --main_process_ip \$MASTER_ADDR \
    --main_process_port \$MASTER_PORT \
    train.py \
      --model_config_id 'Qwen/Qwen2-1.5B' \
      --preprocessed_data_path ${PROCESSED_DATA_DIR} \
      --output_dir ${LOCAL_OUTPUT_DIR} \
      --per_device_train_batch_size 4 \
      --gradient_accumulation_steps 4 \
      --max_steps 10000 \
      --learning_rate 5e-5 \
      --save_strategy steps \
      --save_steps 500 \
      --logging_steps 1
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
mkdir -p ~/qwen2-from-scratch-on-smollm-fineweb/

srun --nodes=1 --ntasks=1 --ntasks-per-node=1 bash -c "
  rsync -a --info=progress2 ${LOCAL_OUTPUT_DIR}/ ~/qwen2-from-scratch-on-smollm-fineweb/
"

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

  4. Para especificar las dependencias de tu trabajo de ajuste, crea un archivo requirements.txt con 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.0
    Nota: En este instructivo, no se usan las versiones más recientes de las dependencias de NVIDIA y PyTorch. Si necesitas dependencias más recientes, consulta la documentación de NVIDIA y la documentación de Pytorch.

  5. Para descargar, tokenizar y preprocesar el conjunto de datos en un formato listo para el entrenamiento, crea un archivo preprocess_data.py con el siguiente contenido:

    import argparse
from datasets import load_dataset
from transformers import AutoTokenizer
import os
from itertools import chain

def get_args():
   parser = argparse.ArgumentParser(description="Download and preprocess a dataset.")
   parser.add_argument("--dataset_name", type=str, required=True)
   parser.add_argument("--dataset_config", type=str, required=True)
   parser.add_argument("--tokenizer_id", type=str, required=True)
   parser.add_argument("--max_seq_length", type=int, required=True)
   parser.add_argument("--output_path", type=str, required=True, help="Path to save the processed dataset.")
   return parser.parse_args()

def main():
   args = get_args()

   if os.path.exists(args.output_path) and os.listdir(args.output_path):
       print(f"Processed dataset already exists at {args.output_path}. Skipping.")
       return

   # 1. Load tokenizer
   tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id)

   # 2. Load raw dataset
   print(f"Loading raw dataset {args.dataset_name}...")
   raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

   # 3. Tokenize
   def tokenize_function(examples):
       return tokenizer(examples["text"])

   num_proc = os.cpu_count()
   print(f"Tokenizing dataset using {num_proc} processes...")
   print("Tokenizing dataset...")
   tokenized_dataset = raw_dataset.map(
       tokenize_function,
       batched=True,
       remove_columns=raw_dataset.column_names,
       desc="Running tokenizer on dataset",
       num_proc=num_proc,
   )

   # 4. Group texts
   def group_texts(examples):
       concatenated_examples = {k: list(chain.from_iterable(examples[k])) for k in examples.keys()}
       total_length = len(concatenated_examples[list(examples.keys())[0]])
       if total_length >= args.max_seq_length:
           total_length = (total_length // args.max_seq_length) * args.max_seq_length
       result = {
           k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
           for k, t in concatenated_examples.items()
       }
       result["labels"] = result["input_ids"].copy()
       return result

   print("Grouping texts...")
   lm_dataset = tokenized_dataset.map(
       group_texts,
       batched=True,
       desc=f"Grouping texts in chunks of {args.max_seq_length}",
       num_proc=num_proc,
   )

   # 5. Save to disk
   print(f"Saving processed dataset to {args.output_path}...")
   lm_dataset.save_to_disk(args.output_path)
   print("Preprocessing complete.")

if __name__ == "__main__":
   main()

  6. Para especificar las instrucciones de tu trabajo, crea un archivo train.py con el siguiente contenido:

    import torch
import argparse
from datasets import load_dataset, load_from_disk
import os
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from huggingface_hub import login

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_config_id", type=str, default="Qwen/Qwen2-1.5B", help="Hugging Face model config to use for architecture.")
    # Data arguments - used if preprocessed data is not available
    parser.add_argument("--dataset_name", type=str, default="HuggingFaceFW/fineweb-edu", help="Hugging Face dataset for pre-training.")
    parser.add_argument("--dataset_config", type=str, default="CC-MAIN-2024-10", help="Config for the fineweb-edu dataset, e.g., 'CC-MAIN-2024-10'.")
    parser.add_argument("--preprocessed_data_path", type=str, default=None, help="Path to a preprocessed dataset on disk. If provided, skips download and processing.")
    # General arguments
    parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face token for private models/tokenizers")
    parser.add_argument("--output_dir", type=str, default="qwen2-from-scratch-on-smollm-fineweb", help="Directory to save model checkpoints")

    # TrainingArguments
    parser.add_argument("--max_seq_length", type=int, default=1024, help="Maximum sequence length")
    parser.add_argument("--num_train_epochs", type=int, default=1, help="Number of training epochs")
    parser.add_argument("--max_steps", type=int, default=-1, help="If set to a positive number, it overrides num_train_epochs.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=4, help="Batch size per device during training")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=5e-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=500, 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. Load Tokenizer ---
    # We load the tokenizer from the specified config ID to ensure compatibility
    # with the model architecture (e.g., special tokens).
    tokenizer = AutoTokenizer.from_pretrained(args.model_config_id)

    # --- 3. Initialize Model from Scratch ---
    print(f"Initializing a new model from {args.model_config_id} configuration...")
    config = AutoConfig.from_pretrained(args.model_config_id)
    model = AutoModelForCausalLM.from_config(config)

    print(f"Model has {model.num_parameters():,} parameters.")

    # --- 4. Load or Create and prepare the training dataset ---
    if args.preprocessed_data_path and os.path.exists(args.preprocessed_data_path):
        print(f"Loading preprocessed dataset from {args.preprocessed_data_path}...")

        # Synchronization of distributed processes
        local_rank = int(os.environ.get("LOCAL_RANK", -1))
        if local_rank != -1:
            # Introducing a minimal time offset per GPU to avoid I/O collisions.
            import time
            time.sleep(local_rank * 0.2)

        lm_dataset = load_from_disk(args.preprocessed_data_path, keep_in_memory=False)

    else:
        print("No preprocessed dataset found, starting from raw data...")
        raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

        # Tokenization function
        def tokenize_function(examples):
            return tokenizer(examples["text"])

        tokenized_dataset = raw_dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=raw_dataset.column_names,
            desc="Running tokenizer on dataset",
        )

        # Main data processing function that will concatenate all texts from our dataset
        # and generate chunks of max_seq_length.
        def group_texts(examples):
            # Concatenate all texts.
            concatenated_examples = {k: [item for sublist in examples[k] for item in sublist] for k in examples.keys()}
            total_length = len(concatenated_examples[list(examples.keys())[0]])
            # We drop the small remainder.
            if total_length >= args.max_seq_length:
                total_length = (total_length // args.max_seq_length) * args.max_seq_length
            # Split by chunks of max_len.
            result = {
                k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
                for k, t in concatenated_examples.items()
            }
            result["labels"] = result["input_ids"].copy()
            return result

        lm_dataset = tokenized_dataset.map(
            group_texts,
            batched=True,
            desc=f"Grouping texts in chunks of {args.max_seq_length}",
        )


    # --- 5. Configure Training Arguments ---
    # Check for bfloat16 support
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    training_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.num_train_epochs,
        max_steps=args.max_steps,
        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,
        save_total_limit=2, # Optional: Limit the number of checkpoints
        bf16=use_bf16,
        fp16=not use_bf16,
        optim="adamw_torch",
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        report_to="tensorboard",
        gradient_checkpointing=True,
        # Required for gradient checkpointing with some parallelization strategies
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )

    # --- 6. Create Trainer and Start Training ---
    # Data collator will take care of creating batches for causal language modeling
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=lm_dataset,
        # eval_dataset=... # Optional: if you have a validation set
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    print("Starting training from scratch...")
    trainer.train()
    print("Training finished.")

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

  1. Recupera el nombre del nodo de acceso de tu clúster para establecer la variable LOGIN_NODE:

    LOGIN_NODE="$(gcloud compute instances list \
                --project="${PROJECT_ID}" \
                --filter="labels.ghpc_deployment='${CLUSTER_NAME}' AND labels.slurm_instance_role='login'" \
                --format="value(name)" | head -n 1)"

    La variable LOGIN_NODE almacena un valor similar a ${CLUSTER_NAME}-login-001.

  2. Crea una regla de firewall:

    gcloud compute firewall-rules create allow-ssh-ingress-from-iap \
  --project="${PROJECT_ID}" \
  --network="${CLUSTER_NETWORK}" \
  --direction=INGRESS \
  --action=allow \
  --rules=tcp:22 \
  --source-ranges=35.235.240.0/20 \
  --description="Allow SSH ingress from Google Cloud Identity-Aware Proxy (IAP)"

  3. 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.txt \
  ./submit.slurm \
  ./accelerate_config.yaml \
  ./train.py \
  ./preprocess_data.py \
  "${LOGIN_NODE}":~/

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}" \
    --project="${PROJECT_ID}" \
    --tunnel-through-iap \
    --zone="${ZONE}"
    -- -t "export HF_TOKEN='${HF_TOKEN}'; bash -l"

Instala frameworks y herramientas

Después de conectarte al nodo de acceso, instala los frameworks y las herramientas de la siguiente manera:

  1. Configura un entorno virtual de Python con todas las dependencias necesarias:

    chmod +x install_environment.sh
./install_environment.sh

Comienza a entrenar previamente tu carga de trabajo

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 resultado creados en tu directorio home:

    tail -f ~/logs/slurm-1.err # (or .out, depending on where the script is currently sending logs)

    Si el trabajo se inicia correctamente, el archivo .err mostrará 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

Como alternativa, puedes ingresar el comando directamente en la terminal:

open "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. PROBAR

  • Duración del trabajo: El trabajo debería tomar aproximadamente una hora en completarse.

Descarga tu modelo

Después de ejecutar el trabajo correctamente, el modelo entrenado se guarda en el directorio ~/qwen2-from-scratch-on-smollm-fineweb/ del nodo de acceso. Como este directorio compartido persistente se activa en todos los nodos de tu clúster, los puntos de control del modelo permanecen disponibles incluso después de que se completa el trabajo o se desasignan los nodos de procesamiento.

Puedes descargar el modelo guardado desde el nodo de acceso a tu máquina local con el comando gcloud compute scp, como se muestra en el siguiente ejemplo:

# From your local machine
gcloud compute scp --project="${PROJECT_ID}" --zone="${ZONE}" --tunnel-through-iap \
  "${LOGIN_NODE}":~/qwen2-from-scratch-on-smollm-fineweb/ ./qwen2-trained-model/ --recurse

Después de descargar tu modelo, puedes hacer lo siguiente:

  • Carga el modelo para la inferencia: Usa el framework de Hugging Face Transformers para cargar el directorio qwen2-trained-model/ y realizar inferencias con tu modelo de Qwen2 recién entrenado.
  • Ajuste adicional: Usa el punto de control guardado como punto de partida para realizar un ajuste adicional en un conjunto de datos más específico.
  • Envía el modelo a Hugging Face Hub: Comparte tu modelo entrenado enviándolo a Hugging Face Hub.

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 tus recursos

  1. Para borrar tu clúster de Slurm, sigue estos pasos:

    ./gcluster destroy "${CLUSTER_NAME}" --auto-approve

  2. Para borrar tu bucket de Cloud Storage, sigue estos pasos:

    gcloud storage buckets delete "gs://${GCS_BUCKET}" --quiet || true

  3. Para borrar una imagen de Packer, abre tu navegador web, navega a la siguiente página, busca tu imagen específica y haz clic en Borrar. 

    http://console.cloud.google.com/compute/images

  4. Para borrar todas las redes de VPC, las reglas de firewall, los routers, las IPs y las subredes asociadas con el proyecto, haz lo siguiente:

    echo "========================================================================="
echo " STARTING AUTOMATED NETWORK CLEANUP FOR CLUSTER: ${CLUSTER_NAME}"
echo "========================================================================="

echo "Discovering all VPC networks linked to the cluster..."
NETWORKS=$(gcloud compute networks list --project="${PROJECT_ID}" --format="value(name)" | grep "^${CLUSTER_NAME}" || true)

if [ -z "${NETWORKS}" ]; then
    echo "No VPC networks found starting with ${CLUSTER_NAME}. Everything is already clean!"
    exit 0
fi

echo "Found the following networks to process:"
echo "${NETWORKS}"
echo "-------------------------------------------------------------------------"

echo "=== 1. Wiping Global Firewall Rules ==="
FIREWALL_RULES=$(gcloud compute firewall-rules list \
    --project="${PROJECT_ID}" \
    --filter="network ~ ^${CLUSTER_NAME} OR name ~ ^${CLUSTER_NAME}" \
    --format="value(name)" || echo "")

if [ -n "${FIREWALL_RULES}" ]; then
    echo "Deleting matching firewall rules:"
    echo "${FIREWALL_RULES}"
    echo "${FIREWALL_RULES}" | xargs -r gcloud compute firewall-rules delete --project="${PROJECT_ID}" --quiet
else
    echo "No matching firewall rules found."
fi

echo "=== 2. Tearing Down Network-Specific Infrastructure ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Processing resources for network: ${net_name}"

    ROUTERS=$(gcloud compute routers list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${ROUTERS}" ]; then
        echo "  -> Deleting routers: ${ROUTERS}"
        echo "${ROUTERS}" | xargs -r gcloud compute routers delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    IPS=$(gcloud compute addresses list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="name ~ ^${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${IPS}" ]; then
        echo "  -> Deleting IP reservations: ${IPS}"
        echo "${IPS}" | xargs -r gcloud compute addresses delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    SUBNETS=$(gcloud compute networks subnets list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${SUBNETS}" ]; then
        echo "  -> Deleting subnetworks:"
        echo "${SUBNETS}"
        echo "${SUBNETS}" | xargs -r gcloud compute networks subnets delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi
done

echo "-------------------------------------------------------------------------"
echo "Waiting 15 seconds for Google Cloud API dependencies to unlock..."
sleep 15

echo "=== 3. Final VPC Networks Destruction ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Deleting core VPC network: ${net_name}..."
    gcloud compute networks delete "${net_name}" --project="${PROJECT_ID}" --quiet || \
    echo "Warning: Could not delete ${net_name} yet. If a lock occurred, please rerun in 1 minute."
done

echo "========================================================================="
echo " SUCCESS: All network resources for cluster ${CLUSTER_NAME} have been wiped!"
echo "========================================================================="

Borra tu proyecto

Borra un Google Cloud proyecto:

gcloud projects delete PROJECT_ID

¿Qué sigue?