Usa FSDP para ajustar Mixtral-8x7B en un clúster de Slurm A4

En este instructivo, se muestra cómo ajustar un modelo mistralai/Mixtral-8x7B-v0.1 en un clúster de Slurm de varios nodos y varias GPUs en Google Cloud. El clúster usa dos instancias de máquina virtual (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 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 un sistema de archivos Lustre compartido y redes 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 fragmentados (FSDP) de PyTorch, al que se accede a través de la biblioteca Transformer Reinforcement Learning (TRL) de Hugging Face.

Este instructivo está dirigido a ingenieros e investigadores de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en IA y datos que deseen distribuir una carga de trabajo de IA en varios nodos y GPUs.

Objetivos

  • Accede a Mixtral con Hugging Face
  • Instala Cluster Toolkit
  • Prepara el entorno
  • Crea e implementa un clúster de Slurm de A4 High-GPU apto para la producción.
  • Configura un entorno de varios nodos para el entrenamiento distribuido con FSDP.
  • Ajusta el modelo de Mixtral con la clase 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 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 lustre.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 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 lustre.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. Cuenta de
    • 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 Mixtral con Hugging Face

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

  1. Accede a Hugging Face y explora el modelo de Mixtral.
  2. Crea un token de acceso de Hugging Face read.
  3. Copia y guarda el valor del token. La usarás más adelante en este instructivo.

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 ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export DEPLOYMENT_NAME="YOUR_CLUSTER_NAME"
export BUCKET_NAME="YOUR_GCS_BUCKET"
export HF_TOKEN="YOUR_HF_TOKEN"

    Reemplaza lo siguiente:

    • YOUR_PROJECT_ID: Es el nombre del proyecto Google Cloud en el que deseas crear tu clúster de Slurm.
    • YOUR_ZONE: Es la zona en la que existe tu reserva.
    • YOUR_REGION: Es la región en la que existe tu reserva.
    • YOUR_RESERVATION_NAME: Es la URL o el nombre de la reserva que deseas usar para crear tu clúster de Slurm.
    • YOUR_CLUSTER_NAME: El nombre de tu implementación. Usa un nombre corto que solo contenga letras y números (por ejemplo, a4high). Este nombre también se asigna al clúster de Slurm que crea la implementación.
    • YOUR_GCS_BUCKET: Es el nombre del bucket en el que almacenas los resultados del punto de control del entrenamiento. Especifica un bucket existente o crea uno nuevo. Antes de crearlo, familiarízate con los requisitos de nombres de buckets.
    • YOUR_HF_TOKEN: El token de Hugging Face que creaste en un paso anterior

  2. Crea un bucket de Cloud Storage:

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

Crea un clúster de Slurm A4

Para crear un clúster de Slurm A4, haz lo siguiente:

  1. Ejecuta el siguiente comando para reemplazar el archivo a4high-slurm-deployment.yaml en el directorio examples/machine-learning/a4-highgpu-8g con tus variables de entorno:

    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: ${BUCKET_NAME}

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

  2. Abre el archivo a4high-slurm-blueprint.yaml en el directorio examples/machine-learning/a4-highgpu-8g y edítalo para usar Managed Lustre en el directorio /home compartido de la siguiente manera:

    • Quita el bloque del módulo homefs predeterminado con source: modules/file-system/filestore.
    • Habilita los módulos lustrefs (bloque homefs con remote_mount: lustrefs) y private-service-access.
    • En el bloque vars, configura lo siguiente:
      1. Cambia el valor de install_managed_lustre en /var/tmp/slurm_vars.json a true.
      2. Establece el parámetro per_unit_storage_throughput en 500.
      3. Establece el parámetro lustre_size_gib en 36000.
      4. Quita el comentario de lustre_instance_id: lustre-instance.
      5. Marca como comentario o quita los filestore_ip_range que no se usen.

  3. Implemente el clúster:

    ./gcluster deploy \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  --auto-approve

    El comando ./gcluster deploy inicia un proceso de dos fases, que es el siguiente:

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

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 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 a reliable environment for FSDP training.
# It is meant to be run on a compute node.
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 --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

# --- [STEP 2.4] Build Flash Attention from Source ---
echo "--- Building flash-attn from source... This will take a while. ---"
# Bypass PyTorch CUDA version mismatch check for extension builds
python3 -c "
import torch.utils.cpp_extension as ce
p = ce.__file__
content = open(p).read().replace('raise RuntimeError(CUDA_MISMATCH_MESSAGE', 'pass # raise RuntimeError')
open(p, 'w').write(content)
"

# Use all available CPU cores to speed up the build
MAX_JOBS=$(nproc) pip install flash-attn --no-build-isolation

# --- 3. Download the Model ---
echo "--- [STEP 2.5] Downloading Mixtral 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/hf download mistralai/Mixtral-8x7B-v0.1 --local-dir ~/Mixtral-8x7B-v0.1 --token $HF_TOKEN

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

  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
protobuf==6.31.1

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

    import torch
from torch.distributed.fsdp import MixedPrecision
from datasets import load_dataset
import shutil
import os
import torch.distributed as dist

from peft import LoraConfig, PeftModel, get_peft_model
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    HfArgumentParser,
)

from torch.distributed import get_rank, get_world_size

from transformers.models.mixtral.modeling_mixtral import MixtralDecoderLayer
from trl import SFTTrainer
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ScriptArguments:
    model_id: str = field(default="mistralai/Mixtral-8x7B-v0.1", 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=False, metadata={"help": "When using FSDP activation checkpointing, this must be set to False for Mixtral"})

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": True}

    training_args.optim = "adamw_torch_fused"

    bf16_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    )

    training_args.fsdp = "full_shard"
    training_args.fsdp_config = {
        "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
        "fsdp_transformer_layer_cls_to_wrap": [MixtralDecoderLayer],
        "fsdp_state_dict_type": "SHARDED_STATE_DICT",
        "fsdp_offload_params": False,
        "fsdp_forward_prefetch": True,
        "fsdp_mixed_precision_policy": bf16_policy
    }

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

    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

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

    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"],
    )

    model = get_peft_model(model, peft_config)

    data_splits = load_dataset(script_args.dataset_name)

    dataset = data_splits["train"]
    eval_dataset = data_splits["test"]

    if script_args.dataset_subset_size is not None:
        dataset = dataset.select(range(script_args.dataset_subset_size))

    dataset = dataset.shuffle(seed=training_args.seed)

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        eval_dataset=eval_dataset,
        formatting_func=formatting_prompts_func,
        processing_class=tokenizer,
    )

    trainer.train()

    dist.barrier()
    if trainer.is_world_process_zero():
        best_model_path = trainer.state.best_model_checkpoint

        final_model_dir = os.path.join(training_args.output_dir, "final_best_model")
        print(f"Copying best model to: {final_model_dir}")

        if os.path.exists(final_model_dir):
            shutil.rmtree(final_model_dir)
        shutil.copytree(best_model_path, final_model_dir)

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

def run_post_training_inference(script_args, best_model_path, tokenizer):
    print("\n" + "="*50)
    print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
    print("="*50 + "\n")

    base_model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto"
    )
    model = PeftModel.from_pretrained(base_model, best_model_path)
    model = model.merge_and_unload()
    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."

    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()

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

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

set -e
set -x

echo "--- Slurm Job Started ---"

# --- Define Paths ---
LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
VENV_PATH="${HOME}/.venv/venv-fsdp"
MODEL_PATH="${HOME}/Mixtral-8x7B-v0.1"

# --- 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 ${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}'/train-mixtral.py \
    --model_id="'${LOCAL_SSD_PATH}'/model/" \
    --output_dir="${HOME}/outputs/mixtral_job_${SLURM_JOB_ID}" \
    --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
    --seed=900913 \
    --bf16=True \
    --num_train_epochs=3 \
    --per_device_train_batch_size=32 \
    --gradient_accumulation_steps=4 \
    --learning_rate=4e-5 \
    --logging_steps=3 \
    --lora_r=32 \
    --lora_alpha=32 \
    --lora_dropout=0.05 \
    --eval_strategy=steps \
    --eval_steps=10 \
    --save_strategy=steps \
    --save_steps=10 \
    --load_best_model_at_end=False \
    --metric_for_best_model=eval_loss \
    --run_inference_after_training \
    --dataset_subset_size=67000
'

# --- STAGE 3: 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 las secuencias de comandos a tu 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='${DEPLOYMENT_NAME}' AND labels.slurm_instance_role='login'" \
                --format="value(name)" | head -n 1)"

    La variable LOGIN_NODE almacena un valor similar a ${DEPLOYMENT_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-fsdp.txt \
  ./train-mixtral.py \
  ./train-mixtral.sh \
  "${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, ejecuta la secuencia de comandos de instalación en un nodo de procesamiento para instalar los frameworks y las herramientas requeridos:

# On the login node
srun \
  --job-name=env-setup \
  --nodes=1 \
  --ntasks=1 \
  --gpus-per-node=1 \
  --partition=a4high \
  bash ./install_environment.sh

Este comando configura el entorno virtual, instala todas las dependencias y descarga los pesos del modelo Mixtral en ~/Mixtral-8x7B-v0.1. Este proceso puede demorar más de 30 minutos.

Ejecutar este paso a través de srun delega la secuencia de comandos de instalación del nodo de acceso a un nodo de procesamiento asignado, lo que permite que la secuencia de comandos acceda a los controladores de GPU durante la compilación.

Inicia tu carga de trabajo de ajuste

Para iniciar tu carga de trabajo de entrenamiento, haz lo siguiente:

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

    # On the login node
sbatch train-mixtral.sh

Supervisa tu carga de trabajo

Para supervisar el progreso de tu trabajo de ajuste fino, usa los siguientes métodos:

  1. En el nodo de acceso de tu clúster de Slurm, puedes supervisar el progreso del trabajo verificando los archivos de salida creados en tu directorio home:

    # On the login node:
tail -f mixtral-*.out mixtral-*.err

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

    El trabajo tiene dos fases principales:

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

    Todo el trabajo tarda unos 60 minutos en ejecutarse.

  2. También puedes supervisar el uso de las GPUs en tu clúster de Slurm para verificar que el 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=YOUR_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 la GPU: Para un trabajo de ajuste fino correcto, puedes esperar ver que el uso de las 16 GPU (ocho GPU 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 recursos

Ejecuta los siguientes comandos para borrar los recursos que creaste en este instructivo:

  1. Para borrar tu clúster de Slurm, ve al directorio cluster-toolkit y ejecuta el siguiente comando:

    ./gcluster destroy "${DEPLOYMENT_NAME}" --auto-approve --robust

  2. Borra tu bucket:

    gcloud storage rm --recursive "gs://${BUCKET_NAME}" --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

Borra tu proyecto

Borra un Google Cloud proyecto:

gcloud projects delete PROJECT_ID

¿Qué sigue?