Google utilizza la tecnologia AI per tradurre i contenuti nella tua lingua preferita. Le traduzioni generate dall'AI potrebbero contenere errori.

Utilizza FSDP per il fine tuning di Mixtral-8x7B su un cluster Slurm A4

GPU in cluster

Questo tutorial mostra come eseguire il fine tuning di un modello mistralai/Mixtral-8x7B-v0.1 su un cluster Slurm multi-nodo e multi-GPU su Google Cloud. Il cluster utilizza due istanze di macchine virtuali (VM) a4-highgpu-8g, ognuna delle quali dispone di 8 GPU NVIDIA B200.

Le due procedure principali descritte in questo tutorial sono le seguenti:

Esegui il deployment di un cluster Slurm ad alte prestazioni di livello di produzione utilizzando Google Cloud Cluster Toolkit. Nell'ambito di questo deployment, crei un'immagine VM personalizzata con il software necessario preinstallato. Configuri anche un file system Lustre condiviso e una rete ad alta velocità.
Dopo il deployment del cluster, esegui un job di ottimizzazione distribuita utilizzando il set di script che accompagnano questo tutorial. Il job utilizza PyTorch Fully Sharded Data Parallel (FSDP), a cui puoi accedere tramite la libreria Transformer Reinforcement Learning (TRL) di Hugging Face.

Questo tutorial è rivolto a machine learning engineer, ricercatori, amministratori e operatori di piattaforme e specialisti di dati e AI interessati a distribuire un carico di lavoro di AI su più nodi e GPU.

Obiettivi

Accedere a Mixtral utilizzando Hugging Face
Installa Cluster Toolkit
prepara l'ambiente
Crea ed esegui il deployment di un cluster Slurm A4 High-GPU di livello di produzione.
Configura un ambiente multi-nodo per l'addestramento distribuito con FSDP.
Perfeziona il modello Mixtral utilizzando la classe trl.SFTTrainer di Hugging Face.
Esegui lo staging dei dati sulle unità SSD locali.
Monitorare il job.
Eseguire la pulizia.

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

I nuovi utenti di Google Cloud potrebbero avere diritto a una prova senza costi.

Prima di iniziare

Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei carichi di lavoro.

Installa Google Cloud CLI.

Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

Per inizializzare gcloud CLI, esegui questo comando:

gcloud init

Crea o seleziona un Google Cloud progetto.

Ruoli richiesti per selezionare o creare un progetto

Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

Creare un progetto Google Cloud :
```
gcloud projects create PROJECT_ID
```
Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.
Seleziona il progetto Google Cloud che hai creato:
```
gcloud config set project PROJECT_ID
```
Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

Verifica che la fatturazione sia attivata per il tuo progetto Google Cloud .

Abilita l'API richiesta:

Ruoli richiesti per abilitare le API

Per abilitare le API, devi disporre dell'autorizzazione serviceusage.services.enable. Se hai creato il progetto, probabilmente disponi già di questa autorizzazione tramite il ruolo Proprietario (roles/owner). In caso contrario, puoi ottenere questa autorizzazione tramite il ruolo Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin). Scopri come concedere i ruoli.

gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com lustre.googleapis.com

Installa Google Cloud CLI.

Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

Per inizializzare gcloud CLI, esegui questo comando:

gcloud init

Crea o seleziona un Google Cloud progetto.

Ruoli richiesti per selezionare o creare un progetto

Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

Creare un progetto Google Cloud :
```
gcloud projects create PROJECT_ID
```
Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.
Seleziona il progetto Google Cloud che hai creato:
```
gcloud config set project PROJECT_ID
```
Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

Verifica che la fatturazione sia attivata per il tuo progetto Google Cloud .

Abilita l'API richiesta:

Ruoli richiesti per abilitare le API

gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com lustre.googleapis.com

Concedi ruoli al tuo account utente. Esegui il seguente comando una volta per ciascuno dei seguenti ruoli 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
```
Sostituisci quanto segue:
- PROJECT_ID: il tuo ID progetto.
- USER_IDENTIFIER: l'identificatore del tuo account utente . Ad esempio: myemail@example.com.
- ROLE: il ruolo IAM che concedi al tuo account utente.
Abilita il account di servizio predefinito per il tuo progetto Google Cloud :
```
gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID
```
Sostituisci PROJECT_NUMBER con il numero del progetto. Per rivedere il numero del progetto, consulta Ottenere un progetto esistente.

Concedi il ruolo Editor (roles/editor) al service account predefinito:

gcloud projects add-iam-policy-binding PROJECT_ID \
  --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
  --role=roles/editor

Crea le credenziali di autenticazione locale per il tuo account utente:
```
gcloud auth application-default login
```
Nota: se utilizzi una shell locale e un provider di identità (IdP) esterno e si verifica un errore di autenticazione dopo l'esecuzione del comando precedente, accedi a Google Cloud CLI con la tua identità federata.

Attiva OS Login per il tuo progetto:

gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE

Accedi o crea un account Hugging Face.
Installa le dipendenze necessarie per utilizzare Cluster Toolkit.

Accedere a Mixtral utilizzando Hugging Face

Per utilizzare Hugging Face per accedere a Mixtral:

Accedi a Hugging Face ed esplora il modello Mixtral.
Crea un token di accesso read di Hugging Face.
Copia e salva il valore del token. Lo utilizzerai più avanti in questo tutorial.

Installa Cluster Toolkit

Cluster Toolkit è uno strumento open source che semplifica il deployment di carichi di lavoro di computing ad alte prestazioni (HPC), intelligenza artificiale (AI) e machine learning (ML) su Google Cloud. Per saperne di più sull'utilizzo di gcluster e sulla gestione dei cluster, consulta la panoramica di Cluster Toolkit.

Prepara la versione di 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

Scarica la release:

# 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"

Definisci il percorso gcluster:

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

prepara l'ambiente

Per preparare l'ambiente:

Imposta le variabili di ambiente predefinite:
```
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"
```
Sostituisci quanto segue:
- YOUR_PROJECT_ID: il nome del progetto Google Cloud in cui vuoi creare il cluster Slurm.
- YOUR_ZONE: la zona in cui esiste la prenotazione.
- YOUR_REGION: la regione in cui esiste la prenotazione.
- YOUR_RESERVATION_NAME: l'URL o il nome della prenotazione che vuoi utilizzare per creare il cluster Slurm.
- YOUR_CLUSTER_NAME: il nome del deployment. Utilizza un nome breve contenente solo lettere e numeri (ad esempio a4high). Questo nome viene assegnato anche al cluster Slurm creato dal deployment.
- YOUR_GCS_BUCKET: il nome del bucket in cui memorizzi i risultati del checkpoint di addestramento. Specifica un bucket esistente o creane uno nuovo. Prima di crearlo, acquisisci familiarità con i requisiti per la denominazione dei bucket.
- YOUR_HF_TOKEN: il token Hugging Face creato in un passaggio precedente.

Crea un bucket Cloud Storage:

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

Crea un cluster Slurm A4

Per creare un cluster A4 Slurm:

Sovrascrivi il file a4high-slurm-deployment.yaml nella directory examples/machine-learning/a4-highgpu-8g con le tue variabili di ambiente eseguendo il seguente 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}
  a4h_cluster_size: 2
  a4h_reservation_name: ${RESERVATION_NAME}
EOF

Apri il file a4high-slurm-blueprint.yaml nella directory examples/machine-learning/a4-highgpu-8g e modificalo per utilizzare Managed Lustre per la directory /home condivisa come segue:
- Rimuovi il blocco del modulo homefs predefinito con source: modules/file-system/filestore.
- Attiva i moduli lustrefs (blocco homefs con remote_mount: lustrefs) e private-service-access.
- Nel blocco vars, configura quanto segue:
  1. Modifica il valore di install_managed_lustre in /var/tmp/slurm_vars.json impostandolo su true.
  2. Imposta il parametro per_unit_storage_throughput su 500.
  3. Imposta il parametro lustre_size_gib su 36000.
  4. Rimuovi commento da lustre_instance_id: lustre-instance.
  5. Commenta o rimuovi filestore_ip_range inutilizzati.
Esegui il deployment del cluster:
```
./gcluster deploy \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  --auto-approve
```
Il comando ./gcluster deploy avvia un processo in due fasi, che è il seguente:
- La prima fase crea un'immagine personalizzata con tutto il software preinstallato, il cui completamento può richiedere fino a 35 minuti.
- La seconda fase esegue il deployment del cluster utilizzando l'immagine personalizzata. Questo processo dovrebbe essere completato più rapidamente rispetto alla prima fase.

Prepara il workload

Per preparare il workload:

Crea script di workload.
Carica gli script nel cluster Slurm.
Connettiti al cluster Slurm.
Installa framework e strumenti.

Crea script del workload

Per creare gli script che verranno utilizzati dal tuo workload di messa a punto:

Per configurare l'ambiente virtuale Python, crea il file install_environment.sh con i seguenti contenuti:

#!/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. ---"

Per specificare le dipendenze Python per lo script di addestramento, crea un file requirements-fsdp.txt con i seguenti contenuti:

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

Specifica train-mixtral.py come script di addestramento principale:

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

Per specificare le attività da eseguire nei job sul cluster Slurm, crea il file train-mixtral.sh con i seguenti contenuti:

#!/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 ---"

Carica gli script nel cluster Slurm

Per caricare gli script creati nella sezione precedente nel cluster Slurm, segui questi passaggi:

Imposta la variabile LOGIN_NODE recuperando il nome del nodo di accesso per il cluster:

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 variabile LOGIN_NODE memorizza un valore simile a ${DEPLOYMENT_NAME}-login-001.

Crea una regola 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)"

Carica gli script nella home directory del nodo di accesso:

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}":~/

Connettiti al cluster Slurm

Connettiti al cluster Slurm connettendoti al nodo di accesso tramite SSH:

gcloud compute ssh "${LOGIN_NODE}" \
    --project="${PROJECT_ID}" \
    --tunnel-through-iap \
    --zone="${ZONE}"
    -- -t "export HF_TOKEN='${HF_TOKEN}'; bash -l"

Installare framework e strumenti

Dopo aver eseguito la connessione al nodo di accesso, installa i framework e gli strumenti richiesti eseguendo lo script di installazione su un nodo di calcolo:

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

Questo comando configura l'ambiente virtuale, installa tutte le dipendenze e scarica i pesi del modello Mixtral in ~/Mixtral-8x7B-v0.1. Questo processo può richiedere più di 30 minuti.

L'esecuzione di questo passaggio tramite srun delega lo script di installazione dal nodo di accesso a un nodo di calcolo allocato, consentendo allo script di accedere ai driver GPU durante la compilazione.

Avvia il carico di lavoro di fine-tuning

Per avviare il workload di addestramento:

Invia il job allo scheduler Slurm:

# On the login node
sbatch train-mixtral.sh

Monitorare il workload

Per monitorare l'avanzamento del job di perfezionamento, utilizza i seguenti metodi:

Sul nodo di accesso nel cluster Slurm, puoi monitorare l'avanzamento del job controllando i file di output creati nella directory home:
```
# On the login node:
tail -f mixtral-*.out mixtral-*.err
```
Se il job viene avviato correttamente, il file .err mostra una barra di avanzamento che si aggiorna man mano che il job procede.

Il job ha due fasi principali:
- Copia del modello di base di grandi dimensioni sull'SSD locale di ogni nodo di calcolo.
- Il job di addestramento, che inizia una volta completata la copia del modello.
L'esecuzione dell'intero job richiede circa 60 minuti.

Puoi anche monitorare l'utilizzo delle GPU nel cluster Slurm per verificare che il job di ottimizzazione venga eseguito in modo efficiente. Per farlo, apri il seguente link nel browser:

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

In alternativa, puoi inserire il comando direttamente nel terminale:

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"

Quando monitori il carico di lavoro, puoi visualizzare quanto segue:

Utilizzo GPU: per un job di perfezionamento corretto, puoi aspettarti di vedere l'utilizzo di tutte le 16 GPU (otto GPU per ogni VM nel cluster) aumentare e stabilizzarsi a un livello specifico durante l'addestramento.
Durata del job: il completamento del job dovrebbe richiedere circa un'ora.

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

Elimina le risorse

Esegui i seguenti comandi per eliminare le risorse create in questo tutorial:

Per eliminare il cluster Slurm, vai alla directory cluster-toolkit ed esegui questo comando:
```
./gcluster destroy "${DEPLOYMENT_NAME}" --auto-approve --robust
```

Elimina il bucket:

gcloud storage rm --recursive "gs://${BUCKET_NAME}" --quiet || true

Per eliminare un'immagine Packer, apri il browser web, vai alla pagina seguente, cerca l'immagine specifica e fai clic su Elimina.
```
http://console.cloud.google.com/compute/images
```

Elimina il progetto

Attenzione: l'eliminazione di un progetto ha i seguenti effetti:

L'intero contenuto del progetto viene eliminato. Se hai utilizzato un progetto esistente per, le attività descritte in questo documento, eliminerai anche tutto il lavoro che hai svolto nel progetto.
Gli ID progetto personalizzati non sono più disponibili. Quando hai creato questo progetto, potresti aver creato un ID progetto personalizzato che vuoi utilizzare in futuro. Per conservare gli URL che utilizzano l'ID progetto, ad esempio un URL appspot.com, elimina le risorse selezionate all'interno del progetto anziché eliminare l'intero progetto.

Se intendi esplorare più architetture, tutorial o guide rapide, puoi riutilizzare i progetti ed evitare così di superare i limiti di quota.

Elimina un progetto Google Cloud :

gcloud projects delete PROJECT_ID

Passaggi successivi

Riesegui il deployment di un cluster Slurm
Testare le prestazioni di rete su un cluster Slurm
Monitora le VM in un cluster Slurm
Crea un endpoint di pubblicazione: una volta ottimizzato il modello, puoi eseguirne il deployment in un endpoint di pubblicazione utilizzando Google Kubernetes Engine (GKE) o Vertex AI per renderlo accessibile per l'inferenza.
Scopri di più sui servizi di archiviazione consigliati per i workload di AI e ML