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 de nivel de producción con el Google Cloud Cluster Toolkit. Como parte de esta implementación, crearás una imagen de VM personalizada con el software necesario preinstalado. También configuras un sistema de archivos Lustre compartido y una red de alta velocidad.
  2. Después de implementar el clúster, ejecuta 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 accedes a través de la biblioteca Transformer aprendizaje por refuerzo (TRL) de Hugging Face.

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

Objetivos

  • Accede a Mixtral con Hugging Face
  • Instala Cluster Toolkit
  • Prepara el entorno
  • Crear e implementar un clúster de Slurm de GPU alta A4 de nivel de producción
  • Configurar un entorno de varios nodos para el entrenamiento distribuido con FSDP
  • Ajusta el modelo de Mixtral con la clase trl.SFTTrainer de Hugging Face.
  • Almacena datos de forma temporal en 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. Instala Google Cloud CLI.

  2. Configura gcloud CLI para usar tu identidad federada.

    Para obtener más información, consulta Accede a gcloud CLI con tu identidad federada.

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

    gcloud init

  4. Crea o selecciona un proyecto de Google Cloud .

    Roles necesarios para seleccionar o crear un proyecto

    • Seleccionar 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.
    • Crea 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 de Google Cloud que estás creando.

    • Selecciona el proyecto de Google Cloud que creaste:

      gcloud config set project PROJECT_ID

      Reemplaza PROJECT_ID por el nombre de tu proyecto de Google Cloud .

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

  6. Habilita la API requerida, si aún no está habilitada:

    Roles necesarios para habilitar las APIs

    Para habilitar las 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 obtenerlo 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

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

  8. Habilita la cuenta de servicio predeterminada para tu proyecto de Google Cloud : 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@ \
    --project=PROJECT_ID

    Reemplaza PROJECT_NUMBER por el número del proyecto. Para revisar el número de tu proyecto, consulta Obtén un proyecto existente.

  9. 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@" \
  --role=roles/editor
  10. Crea credenciales de autenticación locales para tu cuenta de usuario: 
    gcloud auth application-default login
  11. Habilita el Acceso al SO para tu proyecto: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  12. Accede a una cuenta de Hugging Face o crea una.
  13. 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. Lo 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 la CLI de gcluster y la administración de clústeres, consulta la descripción general de 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 LUSTRE_ZONE="YOUR_LUSTRE_ZONE"
export LUSTRE_VPC_NETWORK="YOUR_LUSTRE_VPC_NETWORK"
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 de Google Cloud en el que quieres 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: Es 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 de 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. Reemplaza el archivo a4high-slurm-deployment.yaml en el directorio examples/machine-learning/a4-highgpu-8g por tus variables de entorno. Para ello, ejecuta el siguiente comando:

    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}
  lustre_zone: ${LUSTRE_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 de 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 de 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 de los trabajos que se ejecutarán 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. Establece la variable LOGIN_NODE recuperando el nombre del nodo de acceso de tu clúster:

    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, instala los frameworks y las herramientas necesarios ejecutando la secuencia de comandos de instalación en un nodo de procesamiento:

# 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 de Mixtral en ~/Mixtral-8x7B-v0.1. Este proceso puede tardar 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 muestra una barra de progreso que se actualiza a medida que avanza el trabajo.

    El trabajo tiene dos fases principales:

    • Copiar el modelo base grande en la SSD local de cada nodo de procesamiento
    • 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 GPU en tu clúster de Slurm para verificar que tu trabajo de ajuste 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 GPU: Para un trabajo de ajuste correcto, puedes esperar 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 el entrenamiento.

    • Duración del trabajo: El trabajo debería tardar 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 comando gcluster destroy:

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

  2. Borra el 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 proyecto de Google Cloud :

gcloud projects delete PROJECT_ID

¿Qué sigue?