Utiliser FSDP pour affiner Mixtral-8x7B sur un cluster Slurm A4

Ce tutoriel explique comment affiner un modèle mistralai/Mixtral-8x7B-v0.1 sur un cluster Slurm à nœuds et GPU multiples sur Google Cloud. Le cluster utilise deux instances de machine virtuelle a4-highgpu-8g, chacune dotée de huit GPU NVIDIA B200.

Les deux principaux processus décrits dans ce tutoriel sont les suivants :

  1. Déployez un cluster Slurm de production hautes performances à l'aide deGoogle Cloud Cluster Toolkit. Dans le cadre de ce déploiement, vous allez créer une image de VM personnalisée avec les logiciels nécessaires préinstallés. Vous configurez également un système de fichiers Lustre partagé et un réseau à haut débit.
  2. Une fois le cluster déployé, vous exécutez un job d'affinage distribué à l'aide de l'ensemble de scripts qui accompagnent ce tutoriel. Le job utilise PyTorch Fully Sharded Data Parallel (FSDP), auquel vous accédez via la bibliothèque Hugging Face Transformer Reinforcement Learning (TRL).

Ce tutoriel s'adresse aux ingénieurs et chercheurs en machine learning (ML), aux administrateurs et opérateurs de plate-forme, ainsi qu'aux spécialistes des données et de l'IA qui souhaitent distribuer une charge de travail d'IA sur plusieurs nœuds et GPU.

Objectifs

  • Accéder à Mixtral à l'aide de Hugging Face
  • Installer Cluster Toolkit
  • Préparer votre environnement
  • Créez et déployez un cluster Slurm A4 High-GPU de qualité production.
  • Configurez un environnement multinœud pour l'entraînement distribué avec FSDP.
  • Ajustez le modèle Mixtral à l'aide de la classe trl.SFTTrainer de Hugging Face.
  • Transférez les données vers des disques SSD locaux.
  • surveiller votre job ;
  • Effectuer un nettoyage.

Coûts

Dans ce document, vous utilisez les composants facturables suivants de Google Cloud :

Pour obtenir une estimation des coûts en fonction de votre utilisation prévue, utilisez le simulateur de coût.

Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai sans frais.

Avant de commencer

  1. Connectez-vous à votre compte Google Cloud . Si vous débutez sur Google Cloud, créez un compte pour évaluer les performances de nos produits en conditions réelles. Les nouveaux clients bénéficient également de 300 $ de crédits sans frais pour exécuter, tester et déployer des charges de travail.

  2. Installez la Google Cloud CLI.

  3. Si vous utilisez un fournisseur d'identité (IdP) externe, vous devez d'abord vous connecter à la gcloud CLI avec votre identité fédérée.

  4. Pour initialiser la gcloud CLI, exécutez la commande suivante : 

    gcloud init

  5. Créez ou sélectionnez un projet Google Cloud .

    Rôles requis pour sélectionner ou créer un projet

    • Sélectionnez un projet : la sélection d'un projet ne nécessite pas de rôle IAM spécifique. Vous pouvez sélectionner n'importe quel projet pour lequel un rôle vous a été attribué.
    • Créer un projet : pour créer un projet, vous devez disposer du rôle Créateur de projet (roles/resourcemanager.projectCreator), qui contient l'autorisation resourcemanager.projects.create. Découvrez comment attribuer des rôles.

    • Créez un projet Google Cloud  :

      gcloud projects create PROJECT_ID

      Remplacez PROJECT_ID par le nom du projet Google Cloud que vous créez.

    • Sélectionnez le projet Google Cloud que vous avez créé :

      gcloud config set project PROJECT_ID

      Remplacez PROJECT_ID par le nom de votre projet Google Cloud .

  6. Vérifiez que la facturation est activée pour votre projet Google Cloud .

  7. Activez l'API requise :

    Rôles requis pour activer les API

    Pour activer les API, vous devez disposer de l'autorisation serviceusage.services.enable. Si vous avez créé le projet, vous disposez probablement déjà de cette autorisation grâce au rôle Propriétaire (roles/owner). Sinon, vous pouvez obtenir cette autorisation grâce au rôle Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin). Découvrez comment attribuer des rôles. 

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

  8. Installez la Google Cloud CLI.

  9. Si vous utilisez un fournisseur d'identité (IdP) externe, vous devez d'abord vous connecter à la gcloud CLI avec votre identité fédérée.

  10. Pour initialiser la gcloud CLI, exécutez la commande suivante : 

    gcloud init

  11. Créez ou sélectionnez un projet Google Cloud .

    Rôles requis pour sélectionner ou créer un projet

    • Sélectionnez un projet : la sélection d'un projet ne nécessite pas de rôle IAM spécifique. Vous pouvez sélectionner n'importe quel projet pour lequel un rôle vous a été attribué.
    • Créer un projet : pour créer un projet, vous devez disposer du rôle Créateur de projet (roles/resourcemanager.projectCreator), qui contient l'autorisation resourcemanager.projects.create. Découvrez comment attribuer des rôles.

    • Créez un projet Google Cloud  :

      gcloud projects create PROJECT_ID

      Remplacez PROJECT_ID par le nom du projet Google Cloud que vous créez.

    • Sélectionnez le projet Google Cloud que vous avez créé :

      gcloud config set project PROJECT_ID

      Remplacez PROJECT_ID par le nom de votre projet Google Cloud .

  12. Vérifiez que la facturation est activée pour votre projet Google Cloud .

  13. Activez l'API requise :

    Rôles requis pour activer les API

    Pour activer les API, vous devez disposer de l'autorisation serviceusage.services.enable. Si vous avez créé le projet, vous disposez probablement déjà de cette autorisation grâce au rôle Propriétaire (roles/owner). Sinon, vous pouvez obtenir cette autorisation grâce au rôle Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin). Découvrez comment attribuer des rôles. 

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

  14. Attribuez des rôles à votre compte utilisateur. Exécutez la commande suivante une fois pour chacun des rôles IAM suivants : 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

    Remplacez les éléments suivants :

    • PROJECT_ID : ID de votre projet
    • USER_IDENTIFIER : identifiant de votre compte utilisateur . Par exemple, myemail@example.com.
    • ROLE : rôle IAM que vous accordez à votre compte utilisateur.
  15. Activez le compte de service par défaut pour votre projet Google Cloud  : 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID

    Remplacez PROJECT_NUMBER par votre numéro de projet. Pour consulter le numéro de votre projet, consultez Obtenir un projet existant.

  16. Attribuez le rôle Éditeur (roles/editor) au compte de service par défaut : 
    gcloud projects add-iam-policy-binding PROJECT_ID \
  --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
  --role=roles/editor
  17. Créez des identifiants d'authentification locaux pour votre compte utilisateur : 
    gcloud auth application-default login
  18. Activez OS Login pour votre projet : 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. Connectez-vous à votre compte Hugging Face ou créez-en un.
  20. Installez les dépendances dont vous avez besoin pour utiliser Cluster Toolkit.

Accéder à Mixtral à l'aide de Hugging Face

Pour utiliser Hugging Face afin d'accéder à Mixtral, procédez comme suit :

  1. Connectez-vous à Hugging Face et explorez le modèle Mixtral.
  2. Créez un jeton d'accès read Hugging Face.
  3. Copiez et enregistrez la valeur du jeton. Vous en aurez besoin dans la suite de ce tutoriel.

Installer Cluster Toolkit

Cluster Toolkit est un outil Open Source qui simplifie le déploiement des charges de travail de calcul hautes performances (HPC), d'intelligence artificielle (IA) et de machine learning (ML) sur Google Cloud. Pour en savoir plus sur l'utilisation de gcluster et la gestion des clusters, consultez la présentation de Cluster Toolkit.

  1. Préparez la version 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. Téléchargez la version :

    # 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. Définissez le chemin gcluster :

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

Préparer votre environnement

Pour préparer votre environnement, procédez comme suit :

  1. Définissez les variables d'environnement par défaut : 

    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"

    Remplacez les éléments suivants :

    • YOUR_PROJECT_ID : nom du projet Google Cloud dans lequel vous souhaitez créer votre cluster Slurm.
    • YOUR_ZONE : zone où se trouve votre réservation.
    • YOUR_REGION : région où se trouve votre réservation.
    • YOUR_RESERVATION_NAME : URL ou nom de la réservation que vous souhaitez utiliser pour créer votre cluster Slurm.
    • YOUR_CLUSTER_NAME : nom de votre déploiement. Utilisez un nom court ne contenant que des lettres et des chiffres (par exemple, a4high). Ce nom est également attribué au cluster Slurm créé par le déploiement.
    • YOUR_GCS_BUCKET : nom du bucket dans lequel vous stockez les résultats du point de contrôle d'entraînement. Spécifiez un bucket existant ou créez-en un. Avant de le créer, familiarisez-vous avec les règles de dénomination des buckets.
    • YOUR_HF_TOKEN : jeton Hugging Face que vous avez créé à une étape précédente.

  2. Créez un bucket Cloud Storage :

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

Créer un cluster Slurm A4

Pour créer un cluster Slurm A4 :

  1. Écrasez le fichier a4high-slurm-deployment.yaml dans le répertoire examples/machine-learning/a4-highgpu-8g avec vos variables d'environnement en exécutant la commande suivante :

    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. Ouvrez le fichier a4high-slurm-blueprint.yaml dans le répertoire examples/machine-learning/a4-highgpu-8g, puis modifiez-le pour utiliser Managed Lustre pour le répertoire partagé /home comme suit :

    • Supprimez le bloc de module homefs par défaut avec source: modules/file-system/filestore.
    • Activez les modules lustrefs (bloc homefs avec remote_mount: lustrefs) et private-service-access.
    • Dans le bloc vars, configurez les éléments suivants :
      1. Remplacez la valeur de install_managed_lustre dans /var/tmp/slurm_vars.json par true.
      2. Définissez le paramètre per_unit_storage_throughput sur 500.
      3. Définissez le paramètre lustre_size_gib sur 36000.
      4. Décommenter lustre_instance_id: lustre-instance.
      5. Mettez en commentaire ou supprimez les filestore_ip_range inutilisés.

  3. Déployez le cluster :

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

    La commande ./gcluster deploy lance un processus en deux phases :

    • La première phase consiste à créer une image personnalisée avec tous les logiciels préinstallés. Cette opération peut prendre jusqu'à 35 minutes.
    • La deuxième phase déploie le cluster à l'aide de cette image personnalisée. Ce processus devrait se terminer plus rapidement que la première phase.

Préparer votre charge de travail

Pour préparer votre charge de travail, procédez comme suit :

  1. Créez des scripts de charge de travail.

  2. Importez les scripts dans le cluster Slurm.

  3. Connectez-vous au cluster Slurm.

  4. Installez les frameworks et les outils.

Créer des scripts de charge de travail

Pour créer les scripts que votre charge de travail d'affinage utilisera, procédez comme suit :

  1. Pour configurer l'environnement virtuel Python, créez le fichier install_environment.sh avec le contenu suivant :

    #!/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. Pour spécifier les dépendances Python du script d'entraînement, créez un fichier requirements-fsdp.txt avec le contenu suivant :

    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. Spécifiez train-mixtral.py comme script d'entraînement 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. Pour spécifier les tâches à exécuter sur votre cluster Slurm, créez le fichier train-mixtral.sh avec le contenu suivant :

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

Importer les scripts dans votre cluster Slurm

Pour importer les scripts que vous avez créés dans la section précédente vers le cluster Slurm, procédez comme suit :

  1. Définissez la variable LOGIN_NODE en récupérant le nom du nœud de connexion de votre 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 variable LOGIN_NODE stocke une valeur semblable à ${DEPLOYMENT_NAME}-login-001.

  2. Créez une règle de pare-feu :

    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. Importez vos scripts dans le répertoire d'accueil du nœud de connexion :

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

Se connecter au cluster Slurm

Connectez-vous au cluster Slurm en vous connectant au nœud de connexion via SSH :

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

Installer des frameworks et des outils

Une fois connecté au nœud de connexion, installez les frameworks et outils requis en exécutant le script d'installation sur un nœud de calcul :

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

Cette commande configure l'environnement virtuel, installe toutes les dépendances et télécharge les pondérations du modèle Mixtral dans ~/Mixtral-8x7B-v0.1. Ce processus peut prendre jusqu'à 30 minutes.

L'exécution de cette étape via srun délègue le script d'installation du nœud de connexion à un nœud de calcul alloué, ce qui permet au script d'accéder aux pilotes de GPU lors de la compilation.

Démarrer votre charge de travail d'affinage

Pour démarrer votre charge de travail d'entraînement :

  1. Envoyez le job au planificateur Slurm :

    # On the login node
sbatch train-mixtral.sh

Surveiller votre charge de travail

Pour surveiller la progression de votre job d'affinage, utilisez les méthodes suivantes :

  1. Sur le nœud de connexion de votre cluster Slurm, vous pouvez surveiller la progression du job en vérifiant les fichiers de sortie créés dans votre répertoire home :

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

    Si votre job démarre correctement, le fichier .err affiche une barre de progression qui se met à jour à mesure que votre job progresse.

    Le job comporte deux phases principales :

    • Copie du grand modèle de base sur le SSD local de chaque nœud de calcul.
    • Le job d'entraînement, qui commence une fois la copie du modèle terminée.

    L'exécution de l'ensemble du job prend environ 60 minutes.

  2. Vous pouvez également surveiller l'utilisation des GPU dans votre cluster Slurm pour vérifier que votre job d'affinage s'exécute efficacement. Pour ce faire, ouvrez le lien suivant dans votre navigateur :

    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

    Vous pouvez également saisir la commande directement dans le 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"

    Lorsque vous surveillez votre charge de travail, vous pouvez voir les éléments suivants :

    • Utilisation du GPU : pour une tâche d'affinage correcte, vous devriez voir l'utilisation de vos 16 GPU (huit GPU pour chaque VM du cluster) augmenter et se stabiliser à un niveau spécifique tout au long de votre entraînement.

    • Durée du job : l'exécution du job devrait prendre environ une heure.

Effectuer un nettoyage

Pour éviter que les ressources utilisées lors de ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez les ressources individuelles.

Supprimer les ressources

Exécutez les commandes suivantes pour supprimer les ressources que vous avez créées dans ce tutoriel :

  1. Pour supprimer votre cluster Slurm, accédez au répertoire cluster-toolkit et exécutez la commande suivante :

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

  2. Supprimez votre bucket :

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

  3. Pour supprimer une image Packer, ouvrez votre navigateur Web, accédez à la page suivante, recherchez votre image spécifique, puis cliquez sur "Supprimer".

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

Supprimer votre projet

Supprimer un projet Google Cloud  :

gcloud projects delete PROJECT_ID

Étapes suivantes