Utiliser FSDP pour affiner Llama 4 sur un cluster Slurm A4

Ce tutoriel explique comment affiner un grand modèle de langage (LLM) Llama-4-Scout-17 sur un cluster Slurm multi-nœuds et multi-GPU sur Google Cloud. Le cluster utilise deux instances de machines virtuelles (VM) A4, chacune disposant 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 allez également configurer une instance Filestore partagée et un réseau RDMA à haut débit.
  2. Une fois le cluster déployé, vous exécutez un job de réglage fin distribué à l'aide de l'ensemble de scripts qui accompagnent ce tutoriel. Le job exploite le parallélisme des données entièrement segmentées (FSDP) de PyTorch, auquel vous accédez via l'apprentissage par renforcement des transformateurs de Hugging Face.

Ce tutoriel s'adresse aux ingénieurs 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 utiliser les fonctionnalités de planification des jobs Slurm pour gérer les charges de travail de réglage fin.

Objectifs

  • Accéder à Llama 4 à l'aide de Hugging Face

  • Préparer votre environnement

  • Créez et déployez un cluster Slurm A4 à GPU élevé de qualité production.

  • Configurez un environnement multinœud pour l'entraînement distribué avec FSDP.

  • Ajustez le modèle Llama 4 à l'aide de trl.SFTTrainer 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 avez besoin du rôle IAM Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin), qui contient l'autorisation serviceusage.services.enable. 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

  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 avez besoin du rôle IAM Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin), qui contient l'autorisation serviceusage.services.enable. 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
    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 d'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 à Llama 4 à l'aide de Hugging Face

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

  1. Signez le contrat de consentement pour utiliser Llama 4.

  2. Créez un jeton d'accès read Hugging Face.

    Cliquez sur Votre profil > Paramètres > Jetons d'accès > + Créer un jeton.

  3. Copiez et enregistrez la valeur du jeton read access. Vous l'utiliserez plus tard dans ce tutoriel.

Préparer votre environnement

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

  1. Clonez le dépôt GitHub Cluster Toolkit :

    git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git

  2. Créez un bucket Cloud Storage :

    gcloud storage buckets create gs://BUCKET_NAME \
    --project=PROJECT_ID

    Remplacez les éléments suivants :

    • BUCKET_NAME : nom de votre bucket Cloud Storage qui respecte les exigences de dénomination des buckets.

    • PROJECT_ID : ID du projetGoogle Cloud dans lequel vous souhaitez créer votre bucket Cloud Storage.

Créer un cluster Slurm A4

Pour créer un cluster Slurm A4, procédez comme suit :

  1. Accédez au répertoire cluster-toolkit :

    cd cluster-toolkit

  2. Si vous utilisez Cluster Toolkit pour la première fois, créez le binaire gcluster :

    make

  3. Accédez au répertoire examples/machine-learning/a4-highgpu-8g :

    cd examples/machine-learning/a4-highgpu-8g/

  4. Ouvrez le fichier a4high-slurm-deployment.yaml, puis modifiez-le comme suit :

    terraform_backend_defaults:
  type: gcs
  configuration:
    bucket: BUCKET_NAME

vars:
  deployment_name: a4-high
  project_id: PROJECT_ID
  region: REGION
  zone: ZONE
  a4h_cluster_size: 2
  a4h_reservation_name: RESERVATION_URL

    Remplacez les éléments suivants :

    • BUCKET_NAME : nom du bucket Cloud Storage que vous avez créé dans la section précédente.

    • PROJECT_ID : ID duGoogle Cloud projet dans lequel existe votre Cloud Storage et dans lequel vous souhaitez créer votre cluster Slurm.

    • REGION : région où se trouve votre réservation.

    • ZONE : zone où se trouve votre réservation.

    • RESERVATION_URL : URL de la réservation que vous souhaitez utiliser pour créer votre cluster Slurm. En fonction du projet dans lequel la réservation existe, spécifiez l'une des valeurs suivantes :

      • La réservation existe dans votre projet : RESERVATION_NAME

      • La réservation existe dans un autre projet et votre projet peut l'utiliser : projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME

  5. Déployez le cluster :

    ./gcluster deploy -d examples/machine-learning/a4-highgpu-8g/a4high-slurm-deployment.yaml examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml --auto-approve

    La commande ./gcluster deploy se déroule 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.

    Si la première phase réussit, mais que la deuxième échoue, vous pouvez essayer de déployer à nouveau le cluster Slurm en ignorant la première phase :

    ./gcluster deploy -d examples/machine-learning/a4-highgpu-8g/a4high-slurm-deployment.yaml examples/machine-learning/a4-highgpu-8g/a4high-slurm-blueprint.yaml --auto-approve --skip "image" -w

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 up a consistent environment for FSDP training.
# It is meant to be run once on the login node of your Slurm cluster
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 torchvision torchaudio --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

# --- 3. Download the Model ---
echo "--- [STEP 2.4] Downloading Llama4 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/huggingface-cli download meta-llama/Llama-4-Scout-17B-16E-Instruct --local-dir ~/Llama-4-Scout-17B-16E-Instruct --token $HF_TOKEN

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

    Ce script configure un environnement virtuel Python fiable, installe une version nightly de PyTorch et télécharge le modèle Llama 4.

  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

  3. Spécifiez llama4-train-distributed.py comme script d'entraînement principal :

    import torch
from datasets import load_dataset
from peft import LoraConfig, PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    HfArgumentParser,
)

from torch.distributed import get_rank, get_world_size

from transformers.models.llama4.modeling_llama4 import Llama4TextDecoderLayer
from trl import SFTTrainer
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ScriptArguments:
    model_id: str = field(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=True, metadata={"help": "When using FSDP activation checkpointing, this must be set to True"})

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

    training_args.optim = "adamw_torch_fused"

    training_args.fsdp = "full_shard"
    training_args.fsdp_config = {
        "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
        "fsdp_transformer_layer_cls_to_wrap": [Llama4TextDecoderLayer],
        "fsdp_state_dict_type": "FULL_STATE_DICT",
        "fsdp_offload_params": False,
        "fsdp_forward_prefetch": True,
    }

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

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

    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"],
    )
    rank = get_rank()
    world_size = get_world_size()

    dataset = load_dataset(script_args.dataset_name, split="train")

    if script_args.dataset_subset_size is not None:
        dataset = dataset.select(range(script_args.dataset_subset_size))
    else:
        print(f"Using the full dataset with {len(dataset)} samples.")

    dataset = dataset.shuffle(seed=training_args.seed)
    print(f"Dataset shuffled with seed: {training_args.seed}.")

    if world_size > 1:
        print(f"Sharding dataset for Rank {rank} of {world_size}.")
        dataset = dataset.shard(num_shards=world_size, index=rank)

    print("Initializing SFTTrainer...")
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        peft_config=peft_config,
        formatting_func=formatting_prompts_func,
        processing_class=tokenizer,
    )

    trainer.train()

    trainer.save_model(training_args.output_dir)

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

def run_post_training_inference(script_args, training_args, tokenizer):
    """
    Loads the fine-tuned PEFT adapter from the local output directory and runs inference.
    This should only be called on rank 0 after training is complete.
    """
    print("\n" + "="*50)
    print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
    print("="*50 + "\n")

    # Load the base model and merge the adapter.
    base_model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto"
    )
    # Load the PEFT adapter and merge it into the base model
    model = PeftModel.from_pretrained(base_model, training_args.output_dir)
    model = model.merge_and_unload() # Merge weights for faster inference
    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."

    # This must match the formatting_func exactly
    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()

    Ce script utilise le TRL Supervised Fine-Tuning (SFT) Trainer pour gérer les boucles d'entraînement FSDP, la configuration LoRA (Low-Rank Adaptation) et le formatage des données.

  4. Pour spécifier les tâches à exécuter sur votre cluster Slurm, créez le fichier submit.slurm avec le contenu suivant :

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

set -e
set -x

echo "--- Slurm Job Started ---"
echo "Job ID: $SLURM_JOB_ID"
echo "Node List: $SLURM_JOB_NODELIST"

# --- Define Paths ---
LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
VENV_PATH="${HOME}/.venv/venv-fsdp"
MODEL_PATH="${HOME}/Llama-4-Scout-17B-16E-Instruct"

# --- 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 --info=progress2 ${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}'/llama4-train-distributed.py \
      --model_id="'${LOCAL_SSD_PATH}'/model/" \
      --output_dir="'${LOCAL_SSD_PATH}'/outputs/" \
      --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
      --seed=900913 \
      --bf16=True \
      --num_train_epochs=1 \
      --per_device_train_batch_size=2 \
      --gradient_accumulation_steps=4 \
      --learning_rate=2e-5 \
      --logging_steps=10 \
      --lora_r=16 \
      --lora_alpha=32 \
      --lora_dropout=0.05 \
      --run_inference_after_training
'

# --- STAGE 3: Copy Final Results Back to Persistent Storage ---
echo "--- Copying final results from local SSD to shared storage ---"
PERSISTENT_OUTPUT_DIR="${HOME}/outputs/llama4_job_${SLURM_JOB_ID}"
mkdir -p "$PERSISTENT_OUTPUT_DIR"

# Only copy from the head node where trl has combined the results
srun --nodes=1 --ntasks=1 -w "$head_node" \
  rsync -a --info=progress2 "${LOCAL_SSD_PATH}/outputs/" "${PERSISTENT_OUTPUT_DIR}/"

# --- STAGE 4: 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 des scripts dans le 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. Pour identifier votre nœud de connexion, listez toutes les VM A4 de votre projet :

    gcloud compute instances list --filter="machineType:a4-highgpu-8g"

    Le nom du nœud de connexion est semblable à a4-high-login-001.

  2. 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 \
  ./llama4-train-distributed.py \
  ./submit.slurm \
  "${LOGIN_NODE_NAME}":~/

    Remplacez LOGIN_NODE_NAME par le nom du nœud de connexion.

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_NAME \
    --project=PROJECT_ID \
    --tunnel-through-iap \
    --zone=ZONE

Installer des frameworks et des outils

Une fois connecté au nœud de connexion, installez les frameworks et les outils en procédant comme suit :

  1. Exportez votre jeton Hugging Face :

    # On the login node
export HF_TOKEN="hf_..." # Replace with your token

  2. Exécutez le script d'installation :

    # On the login node
chmod +x install_environment.sh
./install_environment.sh

    Cette commande configure un environnement virtuel avec toutes les dépendances requises et télécharge les pondérations du modèle dans le fichier ~/Llama-4-Scout-17B-16E-Instruct.

    Comme le téléchargement du modèle est très volumineux (environ 200 Go), ce processus prend environ 30 minutes, en fonction des conditions du réseau.

Démarrer votre charge de travail d'affinage

Pour commencer à entraîner votre charge de travail :

  1. Envoyez le job au planificateur Slurm :

    sbatch submit.slurm

  2. 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 llama4-*.out

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

    Ce job devrait prendre un peu plus d'une heure sur le cluster Slurm. 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 ce job prend environ 35 minutes.

Effectuer un nettoyage

Pour éviter que les ressources utilisées dans ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez chaque ressource individuellement.

Supprimer votre projet

Supprimer un projet Google Cloud  :

gcloud projects delete PROJECT_ID

Supprimer votre cluster Slurm

Pour supprimer votre cluster Slurm, procédez comme suit :

  1. Accédez au répertoire cluster-toolkit.

  2. Détruisez le fichier Terraform et toutes les ressources créées :

    ./gcluster destroy a4-high --auto-approve

Supprimer votre instance Filestore

Par défaut, le paramètre deletion_protection de votre instance Filestore est défini sur "true" dans le blueprint cluster-toolkit. Ce paramètre permet d'éviter toute perte de données accidentelle lorsque vous modifiez des environnements. Pour supprimer l'instance Filestore, vous devez désactiver manuellement la protection contre la suppression.

Étapes suivantes