Entraîner Qwen2 sur un cluster Slurm A4

Ce tutoriel explique comment entraîner un grand modèle de langage (LLM) sur un cluster Slurm multi-nœuds et multi-GPU sur Google Cloud. Le modèle que vous utilisez dans ce tutoriel est basé sur un modèle Qwen2 à 1,5 milliard de paramètres. Le cluster Slurm utilise deux machines virtuelles a4-highgpu-8g, 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 haute performance adapté à la production à 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 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 pré-entraînement distribué à l'aide de l'ensemble de scripts fournis avec ce tutoriel. Le job utilise la bibliothèque Hugging Face Accelerate.

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 déployer des clusters Slurm hautes performances sur Google Cloud pour entraîner des LLM.

Objectifs

  • Accédez au modèle Qwen2 à l'aide de Hugging Face.
  • Préparez votre environnement.
  • Créez et déployez un cluster Slurm A4 de niveau production.
  • Entraînez le modèle Qwen2 à l'aide de la bibliothèque Accelerate .
  • 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 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 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.

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

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

  1. Signez le contrat de consentement pour utiliser Qwen2 1.5B.

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

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'entraînement 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 should be run ONCE on the login node to set up the
# shared Python virtual environment.

set -e
echo "--- Creating Python virtual environment in /home ---"
python3 -m venv ~/.venv
echo "--- Activating virtual environment ---"
source ~/.venv/bin/activate

echo "--- Installing build dependencies ---"
pip install --upgrade pip wheel packaging

echo "--- Installing PyTorch for CUDA 12.8 ---"
pip install torch --index-url https://download.pytorch.org/whl/cu128

echo "--- Installing application requirements ---"
pip install -r requirements.txt

echo "--- Environment setup complete. You can now submit jobs with sbatch. ---"

  2. Pour spécifier les configurations de votre tâche d'affinage, créez le fichier accelerate_config.yaml avec le contenu suivant :

    # Default configuration for a 2-node, 8-GPU-per-node (16 total GPUs) FSDP training job.

compute_environment: "LOCAL_MACHINE"
distributed_type: "FSDP"
downcast_bf16: "no"
machine_rank: 0
main_training_function: "main"
mixed_precision: "bf16"
num_machines: 2
num_processes: 16
rdzv_backend: "static"
same_network: true
tpu_env: []
use_cpu: false

  3. 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=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8 # 8 tasks per node
#SBATCH --gpus-per-task=1   # 1 GPU per task
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -e
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
# This command runs once per node.
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c '
  set -e
  echo "Setting up local environment on $(hostname)..."
  LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}"
  LOCAL_CACHE="/mnt/localssd/hf_cache_job_${SLURM_JOB_ID}"
  PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
  rsync -a --info=progress2 ~/./.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export HF_HOME=${LOCAL_CACHE}
  export HF_DATASETS_CACHE=${LOCAL_CACHE}

  # This runs the new preprocessing script. It ensures only ONE process per node
  # downloads and processes the data, avoiding rate limiting and redundant work.
  python preprocess_data.py \
    --dataset_name "HuggingFaceFW/fineweb-edu" \
    --dataset_config "CC-MAIN-2024-10" \
    --tokenizer_id "Qwen/Qwen2-1.5B" \
    --max_seq_length 1024 \
    --output_path ${PROCESSED_DATA_DIR}

  echo "Setup on $(hostname) complete."
'

# --- STAGE 2: Run the Training Job using the Local Environment ---
echo "--- Starting Training ---"

LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}"
PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
LOCAL_OUTPUT_DIR="/mnt/localssd/outputs_${SLURM_JOB_ID}"
mkdir -p ${LOCAL_OUTPUT_DIR}

# This is the main training command. It launches one Python process per GPU.
srun --ntasks=$((SLURM_NNODES * 8)) --gpus-per-task=1 bash -c "
  source ${LOCAL_VENV}/bin/activate

  # The training script now loads the pre-processed data from the local SSD.
  python train.py \
    --model_config_id "Qwen/Qwen2-1.5B" \
    --preprocessed_data_path ${PROCESSED_DATA_DIR} \
    --output_dir ${LOCAL_OUTPUT_DIR} \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --max_steps 10000 \
    --learning_rate 5e-5 \
    --save_strategy steps \
    --save_steps 500
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
# This command runs only on the first node of the job allocation
# and copies the final model back to the persistent shared directory.
srun --nodes=1 --ntasks=1 --ntasks-per-node=1 bash -c "
  rsync -a --info=progress2 ${LOCAL_OUTPUT_DIR}/ ~/qwen2-from-scratch-on-smollm-fineweb/
"

echo "--- Slurm Job Finished ---"

  4. Pour spécifier les dépendances de votre job d'affinage, créez un fichier requirements.txt avec le contenu suivant :

    # Hugging Face Libraries (Pinned to recent, stable versions for reproducibility)
transformers==4.53.3
datasets==4.0.0
accelerate==1.9.0
evaluate==0.4.5
bitsandbytes==0.46.1
trl==0.19.1
peft==0.16.0

# Other dependencies
tensorboard==2.20.0
protobuf==6.31.1
sentencepiece==0.2.0

  5. Pour télécharger, tokeniser et prétraiter l'ensemble de données dans un format prêt pour l'entraînement, créez un fichier preprocess_data.py avec le contenu suivant :

    import argparse
from datasets import load_dataset
from transformers import AutoTokenizer
import os
from itertools import chain

def get_args():
   parser = argparse.ArgumentParser(description="Download and preprocess a dataset.")
   parser.add_argument("--dataset_name", type=str, required=True)
   parser.add_argument("--dataset_config", type=str, required=True)
   parser.add_argument("--tokenizer_id", type=str, required=True)
   parser.add_argument("--max_seq_length", type=int, required=True)
   parser.add_argument("--output_path", type=str, required=True, help="Path to save the processed dataset.")
   return parser.parse_args()

def main():
   args = get_args()

   if os.path.exists(args.output_path) and os.listdir(args.output_path):
       print(f"Processed dataset already exists at {args.output_path}. Skipping.")
       return

   # 1. Load tokenizer
   tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id)

   # 2. Load raw dataset
   print(f"Loading raw dataset {args.dataset_name}...")
   raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

   # 3. Tokenize
   def tokenize_function(examples):
       return tokenizer(examples["text"])

   num_proc = os.cpu_count()
   print(f"Tokenizing dataset using {num_proc} processes...")
   print("Tokenizing dataset...")
   tokenized_dataset = raw_dataset.map(
       tokenize_function,
       batched=True,
       remove_columns=raw_dataset.column_names,
       desc="Running tokenizer on dataset",
       num_proc=num_proc,
   )

   # 4. Group texts
   def group_texts(examples):
       concatenated_examples = {k: list(chain.from_iterable(examples[k])) for k in examples.keys()}
       total_length = len(concatenated_examples[list(examples.keys())[0]])
       if total_length >= args.max_seq_length:
           total_length = (total_length // args.max_seq_length) * args.max_seq_length
       result = {
           k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
           for k, t in concatenated_examples.items()
       }
       result["labels"] = result["input_ids"].copy()
       return result

   print("Grouping texts...")
   lm_dataset = tokenized_dataset.map(
       group_texts,
       batched=True,
       desc=f"Grouping texts in chunks of {args.max_seq_length}",
       num_proc=num_proc,
   )

   # 5. Save to disk
   print(f"Saving processed dataset to {args.output_path}...")
   lm_dataset.save_to_disk(args.output_path)
   print("Preprocessing complete.")

if __name__ == "__main__":
   main()

  6. Pour spécifier les instructions de votre job, créez un fichier train.py avec le contenu suivant :

    import torch
import argparse
from datasets import load_dataset, load_from_disk
import os
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from huggingface_hub import login

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_config_id", type=str, default="Qwen/Qwen2-1.5B", help="Hugging Face model config to use for architecture.")
    # Data arguments - used if preprocessed data is not available
    parser.add_argument("--dataset_name", type=str, default="HuggingFaceFW/fineweb-edu", help="Hugging Face dataset for pre-training.")
    parser.add_argument("--dataset_config", type=str, default="CC-MAIN-2024-10", help="Config for the smollm-corpus dataset, e.g., 'fineweb-edu-dedup'.")
    parser.add_argument("--preprocessed_data_path", type=str, default=None, help="Path to a preprocessed dataset on disk. If provided, skips download and processing.")
    # General arguments
    parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face token for private models/tokenizers")
    parser.add_argument("--output_dir", type=str, default="qwen2-from-scratch-on-olmo", help="Directory to save model checkpoints")

    # TrainingArguments
    parser.add_argument("--max_seq_length", type=int, default=1024, help="Maximum sequence length")
    parser.add_argument("--num_train_epochs", type=int, default=1, help="Number of training epochs")
    parser.add_argument("--max_steps", type=int, default=-1, help="If set to a positive number, it overrides num_train_epochs.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=4, help="Batch size per device during training")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate")
    parser.add_argument("--logging_steps", type=int, default=10, help="Log every X steps")
    parser.add_argument("--save_strategy", type=str, default="steps", help="Checkpoint save strategy")
    parser.add_argument("--save_steps", type=int, default=500, help="Save checkpoint every X steps")

    return parser.parse_args()

def main():
    args = get_args()

    # --- 1. Setup and Login ---
    if args.hf_token:
        login(args.hf_token)

    # --- 2. Load Tokenizer ---
    # We load the tokenizer from the specified config ID to ensure compatibility
    # with the model architecture (e.g., special tokens).
    tokenizer = AutoTokenizer.from_pretrained(args.model_config_id)

    # --- 4. Initialize Model from Scratch ---
    print(f"Initializing a new model from {args.model_config_id} configuration...")
    config = AutoConfig.from_pretrained(args.model_config_id)
    model = AutoModelForCausalLM.from_config(config)

    print(f"Model has {model.num_parameters():,} parameters.")

    # --- 3. Load or Create and prepare the training dataset ---
    if args.preprocessed_data_path and os.path.exists(args.preprocessed_data_path):
        print(f"Loading preprocessed dataset from {args.preprocessed_data_path}...")
        lm_dataset = load_from_disk(args.preprocessed_data_path)
    else:
        print("No preprocessed dataset found, starting from raw data...")
        raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

        # Tokenization function
        def tokenize_function(examples):
            return tokenizer(examples["text"])

        tokenized_dataset = raw_dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=raw_dataset.column_names,
            desc="Running tokenizer on dataset",
        )

        # Main data processing function that will concatenate all texts from our dataset
        # and generate chunks of max_seq_length.
        def group_texts(examples):
            # Concatenate all texts.
            concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
            total_length = len(concatenated_examples[list(examples.keys())[0]])
            # We drop the small remainder.
            if total_length >= args.max_seq_length:
                total_length = (total_length // args.max_seq_length) * args.max_seq_length
            # Split by chunks of max_len.
            result = {
                k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
                for k, t in concatenated_examples.items()
            }
            result["labels"] = result["input_ids"].copy()
            return result

        lm_dataset = tokenized_dataset.map(
            group_texts,
            batched=True,
            desc=f"Grouping texts in chunks of {args.max_seq_length}",
        )

    # --- 5. Configure Training Arguments ---
    # Check for bfloat16 support
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    training_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.num_train_epochs,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        logging_steps=args.logging_steps,
        save_strategy=args.save_strategy,
        save_steps=args.save_steps,
        save_total_limit=2, # Optional: Limit the number of checkpoints
        bf16=use_bf16,
        fp16=not use_bf16,
        optim="adamw_torch",
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        report_to="tensorboard",
        gradient_checkpointing=True,
        # Required for gradient checkpointing with some parallelization strategies
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )

    # --- 6. Create Trainer and Start Training ---
    # Data collator will take care of creating batches for causal language modeling
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=lm_dataset,
        # eval_dataset=... # Optional: if you have a validation set
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    print("Starting training from scratch...")
    trainer.train()
    print("Training finished.")

    # --- 7. Save the final model ---
    print(f"Saving final model to {args.output_dir}")
    trainer.save_model()

if __name__ == "__main__":
    main()

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 \
  ./train.py \
  ./requirements.txt \
  ./submit.slurm \
  ./install_environment.sh \
  ./accelerate_config.yaml \
  "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. Créez une variable d'environnement pour votre jeton d'accès Hugging Face :

    export HUGGING_FACE_TOKEN="HUGGING_FACE_TOKEN"

  2. Configurez un environnement virtuel Python avec toutes les dépendances requises :

    chmod +x install_environment.sh
./install_environment.sh

Commencer le pré-entraînement de votre charge de travail

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 :

    tail -f logs/slurm-qwen2-pretrain-smollm-fineweb.err

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

Surveiller votre charge de travail

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

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 des GPU : pour une tâche d'affinage saine, vous pouvez vous attendre à ce que l'utilisation de vos 16 GPU (huit GPU pour chaque VM du cluster) augmente et se stabilise à 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.

Télécharger votre modèle

Une fois votre job exécuté, votre modèle entraîné est enregistré dans le répertoire ~/qwen2-from-scratch-on-smollm-fineweb/ du nœud de connexion. Étant donné que ce répertoire partagé persistant est installé sur tous les nœuds de votre cluster, vos points de contrôle de modèle restent disponibles même une fois le job terminé ou les nœuds de calcul désalloués.

Vous pouvez télécharger le modèle enregistré à partir du nœud de connexion sur votre machine locale à l'aide de la commande gcloud compute scp, comme illustré dans l'exemple suivant :

# From your local machine
LOGIN_NODE_NAME="your-login-node-name" # e.g., a4high-login-001
PROJECT_ID="your-gcp-project-id"
ZONE="your-cluster-zone" # e.g., us-west4-a

gcloud compute scp --project="$PROJECT_ID" --zone="$ZONE" --tunnel-through-iap \
  "${LOGIN_NODE_NAME}":~/qwen2-from-scratch-on-smollm-fineweb/ ./qwen2-trained-model/ --recurse

Une fois votre modèle téléchargé, vous pouvez effectuer les opérations suivantes :

  • Chargez le modèle pour l'inférence : utilisez le framework Hugging Face Transformers pour charger le répertoire qwen2-trained-model/ et effectuer l'inférence avec votre nouveau modèle Qwen2 entraîné.
  • Affinage supplémentaire : utilisez le point de contrôle enregistré comme point de départ pour un affinage supplémentaire sur un ensemble de données plus spécifique.
  • Transférer le modèle vers le Hub Hugging Face : partagez votre modèle entraîné en le transférant vers le Hub Hugging Face.

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

Étapes suivantes