Qwen2 in einem A4-Slurm-Cluster trainieren

In dieser Anleitung wird beschrieben, wie Sie ein Large Language Model (LLM) in einem Slurm-Cluster mit mehreren Knoten und mehreren GPUs in Google Cloudtrainieren. Das Modell, das Sie in dieser Anleitung verwenden, basiert auf einem Qwen2-Modell mit 1,5 Milliarden Parametern. Der Slurm-Cluster verwendet zwei virtuelle Maschinen (VMs) vom Typ „a4-highgpu-8g“ mit jeweils 8 NVIDIA B200-GPUs.

Die beiden Hauptprozesse, die in dieser Anleitung beschrieben werden, sind:

  1. Stellen Sie mit demGoogle Cloud Cluster Toolkit einen leistungsstarken Slurm-Cluster für die Produktion bereit. Im Rahmen dieser Bereitstellung erstellen Sie ein benutzerdefiniertes VM-Image mit der erforderlichen vorinstallierten Software. Außerdem richten Sie eine freigegebene Filestore-Instanz ein und konfigurieren ein RDMA-Netzwerk mit hoher Geschwindigkeit.
  2. Nachdem der Cluster bereitgestellt wurde, führen Sie einen verteilten Vortrainingsjob mit den Skripts aus, die zu dieser Anleitung gehören. Der Job nutzt die Hugging Face Accelerate-Bibliothek.

Diese Anleitung richtet sich an Entwickler für maschinelles Lernen (ML), Forscher, Plattformadministratoren und -operatoren sowie an Daten- und KI-Spezialisten, die daran interessiert sind, leistungsstarke Slurm-Cluster auf Google Cloud zum Trainieren von LLMs bereitzustellen.

Ziele

  • Greifen Sie über Hugging Face auf das Qwen2-Modell zu.
  • Bereiten Sie Ihre Umgebung vor.
  • Produktionsreifen A4-Slurm-Cluster erstellen und bereitstellen
  • Trainieren Sie das Qwen2-Modell mit der Accelerate-Bibliothek .
  • den Job überwachen
  • bereinigen.

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen.

Neuen Nutzern von Google Cloud steht möglicherweise eine kostenlose Testversion zur Verfügung.

Hinweis

  1. Installieren Sie die Google Cloud CLI.

  2. Konfigurieren Sie die gcloud CLI für die Verwendung Ihrer föderierten Identität.

    Weitere Informationen finden Sie unter Mit Ihrer föderierten Identität in der gcloud CLI anmelden.

  3. Führen Sie den folgenden Befehl aus, um die gcloud CLI zu initialisieren: 

    gcloud init

  4. Erstellen oder wählen Sie ein Google Cloud -Projekt aus.

    Erforderliche Rollen zum Auswählen oder Erstellen eines Projekts

    • Projekt auswählen: Für die Auswahl eines Projekts ist keine bestimmte IAM-Rolle erforderlich. Sie können jedes Projekt auswählen, für das Ihnen eine Rolle zugewiesen wurde.
    • Projekt erstellen: Zum Erstellen eines Projekts benötigen Sie die Rolle „Project Creator“ (roles/resourcemanager.projectCreator), die die Berechtigung resourcemanager.projects.create enthält. Informationen zum Zuweisen von Rollen

    • Erstellen Sie ein Projekt in Google Cloud :

      gcloud projects create PROJECT_ID

      Ersetzen Sie PROJECT_ID durch einen Namen für das Projekt in Google Cloud , das Sie erstellen.

    • Wählen Sie das von Ihnen erstellte Projekt in Google Cloud aus:

      gcloud config set project PROJECT_ID

      Ersetzen Sie PROJECT_ID durch den Namen Ihres Projekts in Google Cloud .

  5. Prüfen Sie, ob die Abrechnung für Ihr Google Cloud -Projekt aktiviert ist.

  6. Aktivieren Sie die erforderliche API, falls sie noch nicht aktiviert ist:

    Rollen, die zum Aktivieren von APIs erforderlich sind

    Zum Aktivieren von APIs benötigen Sie die Berechtigung serviceusage.services.enable. Wenn Sie das Projekt erstellt haben, haben Sie diese Berechtigung wahrscheinlich bereits über die Inhaberrolle (roles/owner). Andernfalls können Sie diese Berechtigung über die Rolle „Service Usage-Administrator“ (roles/serviceusage.serviceUsageAdmin) erhalten. Informationen zum Zuweisen von Rollen 

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

  7. Weisen Sie Ihrem Nutzerkonto Rollen zu. Führen Sie den folgenden Befehl für jede der folgenden IAM-Rollen einmal aus: roles/compute.admin, roles/iam.serviceAccountUser, roles/file.editor, roles/storage.admin, roles/serviceusage.serviceUsageAdmin, roles/compute.osAdminLogin, roles/iap.tunnelResourceAccessor 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    Ersetzen Sie Folgendes:

  8. Aktivieren Sie das Standarddienstkonto für Ihr Google Cloud -Projekt: 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@ \
    --project=PROJECT_ID

    Ersetzen Sie PROJECT_NUMBER durch die Projekt-ID. Informationen zum Ermitteln Ihrer Projektnummer finden Sie unter Vorhandenes Projekt abrufen.

  9. Weisen Sie dem Standarddienstkonto die Rolle „Bearbeiter“ (roles/editor) zu: 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@" \
    --role=roles/editor
  10. Erstellen Sie lokale Authentifizierungsdaten für Ihr Nutzerkonto: 
    gcloud auth application-default login
  11. Aktivieren Sie OS Login für Ihr Projekt: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  12. Melden Sie sich in einem Hugging Face-Konto an oder erstellen Sie eines.

Über Hugging Face auf Qwen2 zugreifen

So greifen Sie mit Hugging Face auf Qwen2 zu:

  1. Unterzeichnen Sie die Einwilligungsvereinbarung, um Qwen 2 1.5B zu verwenden.

  2. Erstellen Sie ein read-Zugriffstoken.

Cluster Toolkit installieren

Cluster Toolkit ist ein Open-Source-Tool, mit dem sich die Bereitstellung von Arbeitslasten für Hochleistungs-Computing (HPC), künstliche Intelligenz (KI) und maschinelles Lernen (ML) in Google Cloudvereinfachen lässt. Weitere Informationen zur Verwendung der gcluster-Befehlszeile und zur Verwaltung von Clustern finden Sie in der Übersicht über das Cluster Toolkit.

  1. Bereiten Sie die Cluster Toolkit-Version vor:

    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. Laden Sie den Release herunter:

    # 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. Definieren Sie den Pfad gcluster:

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

Umgebung vorbereiten

So bereiten Sie die Umgebung vor:

  1. Legen Sie die Standardumgebungsvariablen fest: 

    export PROJECT_ID="YOUR_PROJECT_ID"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_URL="YOUR_RESERVATION_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export HF_TOKEN="YOUR_HF_TOKEN"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

    Ersetzen Sie Folgendes:

    • YOUR_PROJECT_ID: der Name des Google Cloud -Projekts, in dem Sie den GKE-Cluster erstellen möchten.
    • YOUR_CLUSTER_NAME: der Name des Slurm-Clusters, den Sie erstellen möchten.
    • YOUR_ZONE: die Zone, in der sich Ihre Reservierung befindet.
    • YOUR_REGION,: die Region, in der sich Ihre Reservierung befindet.
    • RESERVATION_NAME: die URL oder der Name der Reservierung, die Sie zum Erstellen des Slurm-Clusters verwenden möchten.
    • YOUR_GCS_BUCKET: der Name des Buckets, in dem Sie die Ergebnisse des Trainingsprüfpunkts speichern. Machen Sie sich vor dem Erstellen mit den Anforderungen für Bucket-Namen vertraut.
    • YOUR_HF_TOKEN: das Hugging Face-Token, das Sie in einem vorherigen Schritt erstellt haben.

  2. Erstellen Sie einen Cloud Storage-Bucket:

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

A4-Slurm-Cluster erstellen

So erstellen Sie einen A4-Slurm-Cluster:

  1. Erstellen Sie die Datei a4high-slurm-deployment.yaml.

    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: ${GCS_BUCKET}

vars:
  deployment_name: ${CLUSTER_NAME}
  project_id: ${PROJECT_ID}
  region: ${REGION}
  zone: ${ZONE}
  a4h_cluster_size: 2
  a4h_reservation_name: ${RESERVATION_URL}
EOF

  2. Verwenden Sie den Befehl gcluster create, um die Terraform-Manifeste zu erstellen:

    gcluster create \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  -o "${CLUSTER_NAME}"

  3. Patchen Sie die Manifeste:

    sed -i '/deletion_protection = {/,/}/ { s/enabled = true/enabled = false/; /reason  = "Avoid data loss"/d; }' "${CLUSTER_NAME}/${CLUSTER_NAME}/cluster-env/main.tf"

  4. Stellen Sie den Cluster bereit:

    gcluster deploy "${CLUSTER_NAME}/${CLUSTER_NAME}" --auto-approve

    Der gcluster deploy-Befehl ist ein zweiphasiger Prozess, der so aussieht:

    • In der ersten Phase wird ein benutzerdefiniertes Image mit der gesamten vorinstallierten Software erstellt. Dies kann bis zu 50 Minuten dauern.

    • In der zweiten Phase wird der Cluster mit diesem benutzerdefinierten Image bereitgestellt. Dieser Vorgang dauert in der Regel weniger lange als die erste Phase.

    Wenn die erste Phase erfolgreich ist, die zweite Phase jedoch fehlschlägt, können Sie versuchen, den Slurm-Cluster noch einmal bereitzustellen, indem Sie die erste Phase überspringen:

    gcluster deploy "${CLUSTER_NAME}" --auto-approve --skip "image" -w

Arbeitslast vorbereiten

So bereiten Sie Ihre Arbeitslast vor:

  1. Arbeitslastskripts erstellen

  2. Skripts in den Slurm-Cluster hochladen

  3. Stellen Sie eine Verbindung zum Slurm-Cluster her.

  4. Frameworks und Tools installieren

Arbeitslastskripts erstellen

So erstellen Sie die Skripts, die von Ihrer Trainingsarbeitslast verwendet werden:

  1. Erstellen Sie die Datei install_environment.sh mit folgendem Inhalt, um die virtuelle Python-Umgebung einzurichten:

    #!/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. Erstellen Sie die Datei accelerate_config.yaml mit folgendem Inhalt, um die Konfigurationen für Ihren Feinabstimmungsjob anzugeben:

    compute_environment: "LOCAL_MACHINE"
distributed_type: "FSDP"
downcast_bf16: "no"
fsdp_config:
  fsdp_auto_wrap_policy: "TRANSFORMER_BASED_WRAP"
  fsdp_backward_prefetch: "BACKWARD_PRE"
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: "FULL_SHARD"
  fsdp_state_dict_type: "SHARDED_STATE_DICT"
  fsdp_transformer_layer_cls_to_wrap: "Qwen2DecoderLayer"
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. Erstellen Sie die Datei submit.slurm mit folgendem Inhalt, um die Aufgaben für die Jobs anzugeben, die in Ihrem Slurm-Cluster ausgeführt werden sollen:

    #SBATCH --job-name=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1  # 1 main srun task on the node that manages accelerate
#SBATCH --gpus-per-node=8    # access to all 8 GPUs on the node
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -euo pipefail
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none 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}"
  LOCAL_TMP="/mnt/localssd/tmp_job_${SLURM_JOB_ID}"

  rsync -a --info=progress2 ~/.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR} ${LOCAL_TMP}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export TMPDIR="${LOCAL_TMP}"
  export TEMP="${LOCAL_TMP}"
  export TMP="${LOCAL_TMP}"
  export HF_HOME="${LOCAL_CACHE}"
  export HF_DATASETS_CACHE="${LOCAL_CACHE}"
  export HF_MODULES_CACHE="${LOCAL_CACHE}/modules"
  export HF_METRICS_CACHE="${LOCAL_CACHE}/metrics"

  python "${HOME}/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}"

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29505

# Network and initialization debugging configurations
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=INIT,COLL
export TORCH_DISTRIBUTED_DEBUG=INFO
export NCCL_IB_DISABLE=0

# Launching with full GPU access enabled for accelerate
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none bash -c "
  mkdir -p ${LOCAL_OUTPUT_DIR}
  source ${LOCAL_VENV}/bin/activate

  # Reset the forced Slurm isolation to expose all 8 GPUs to the processes
  unset CUDA_VISIBLE_DEVICES

  # Retrieve the default network device.
  DETECTED_IFACE=\$(ip route show | grep default | awk '{print \$5}' | head -n 1)
  echo \"[INFO] Automatically detected cluster network interface: \$DETECTED_IFACE\"

  # Dynamic injection of the detected interface into the network configuration
  export NCCL_SOCKET_IFNAME=\"\${DETECTED_IFACE},gpu*\"
  export TP_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"
  export GLOO_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"

  accelerate launch \
    --config_file ~/accelerate_config.yaml \
    --num_machines \$SLURM_NNODES \
    --num_processes \$((SLURM_NNODES * 8)) \
    --machine_rank \$SLURM_NODEID \
    --main_process_ip \$MASTER_ADDR \
    --main_process_port \$MASTER_PORT \
    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 \
      --logging_steps 1
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
mkdir -p ~/qwen2-from-scratch-on-smollm-fineweb/

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. Wenn Sie die Abhängigkeiten für Ihren Feinabstimmungsjob angeben möchten, erstellen Sie eine requirements.txt-Datei mit folgendem Inhalt:

    # 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
    Hinweis: In dieser Anleitung werden nicht die aktuellen Versionen der NVIDIA- und PyTorch-Abhängigkeiten verwendet. Wenn Sie neuere Abhängigkeiten benötigen, lesen Sie die NVIDIA-Dokumentation und die PyTorch-Dokumentation.

  5. Wenn Sie das Dataset herunterladen, tokenisieren und in ein für das Training geeignetes Format vorverarbeiten möchten, erstellen Sie eine preprocess_data.py-Datei mit folgendem Inhalt:

    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. Erstellen Sie eine train.py-Datei mit folgendem Inhalt, um die Anweisungen für Ihren Job anzugeben:

    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 fineweb-edu dataset, e.g., 'CC-MAIN-2024-10'.")
    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-smollm-fineweb", 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)

    # --- 3. 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.")

    # --- 4. 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}...")

        # Synchronization of distributed processes
        local_rank = int(os.environ.get("LOCAL_RANK", -1))
        if local_rank != -1:
            # Introducing a minimal time offset per GPU to avoid I/O collisions.
            import time
            time.sleep(local_rank * 0.2)

        lm_dataset = load_from_disk(args.preprocessed_data_path, keep_in_memory=False)

    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: [item for sublist in examples[k] for item in sublist] 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()

Skripts in den Slurm-Cluster hochladen

So laden Sie die im vorherigen Abschnitt erstellten Skripts in den Slurm-Cluster hoch:

  1. Legen Sie die Variable LOGIN_NODE fest, indem Sie den Namen des Login-Knotens für Ihren Cluster abrufen:

    LOGIN_NODE="$(gcloud compute instances list \
                --project="${PROJECT_ID}" \
                --filter="labels.ghpc_deployment='${CLUSTER_NAME}' AND labels.slurm_instance_role='login'" \
                --format="value(name)" | head -n 1)"

    In der Variable LOGIN_NODE wird ein Wert wie ${CLUSTER_NAME}-login-001 gespeichert.

  2. Firewallregel erstellen

    # Identify the VPC network used by the Slurm cluster's login node
VPC_NETWORK_URI=$(gcloud compute instances describe "${LOGIN_NODE}" \
  --project="${PROJECT_ID}" \
  --zone="${ZONE}" \
  --format="value(networkInterfaces[0].network)")
CLUSTER_NETWORK=$(basename "${VPC_NETWORK_URI}")

# Create the IAP firewall rule to allow SSH access to the login node
gcloud compute firewall-rules create "allow-ssh-from-iap-${CLUSTER_NETWORK}" \
  --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. Laden Sie Ihre Skripts in das Basisverzeichnis des Login-Knotens hoch:

    gcloud compute scp \
  --project="${PROJECT_ID}" \
  --zone="${ZONE}" \
  --tunnel-through-iap \
  ./install_environment.sh \
  ./requirements.txt \
  ./submit.slurm \
  ./accelerate_config.yaml \
  ./train.py \
  ./preprocess_data.py \
  "${LOGIN_NODE}":~/

Verbindung zum Slurm-Cluster herstellen

Stellen Sie eine Verbindung zum Slurm-Cluster her, indem Sie über SSH eine Verbindung zum Login-Knoten herstellen:

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

Frameworks und Tools installieren

Nachdem Sie eine Verbindung zum Login-Knoten hergestellt haben, installieren Sie Frameworks und Tools so:

  1. Richten Sie eine virtuelle Python-Umgebung mit allen erforderlichen Abhängigkeiten ein:

    chmod +x install_environment.sh
./install_environment.sh

Vortraining der Arbeitslast starten

So trainieren Sie Ihre Arbeitslast:

  1. Senden Sie den Job an den Slurm-Planer:

    sbatch submit.slurm

  2. Auf dem Login-Knoten in Ihrem Slurm-Cluster können Sie den Fortschritt des Jobs anhand der Ausgabedateien im Verzeichnis home verfolgen:

    tail -f ~/logs/slurm-1.err # (or .out, depending on where the script is currently sending logs)

    Wenn der Job erfolgreich gestartet wird, wird in der Datei .err eine Fortschrittsanzeige angezeigt, die sich im Verlauf des Jobs aktualisiert.

Arbeitslast überwachen

Sie können die Nutzung der GPUs in Ihrem Slurm-Cluster überwachen, um zu prüfen, ob Ihr Feinabstimmungsjob effizient ausgeführt wird. Öffnen Sie dazu den folgenden Link in Ihrem Browser:

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

Alternativ können Sie den Befehl direkt im Terminal eingeben:

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"

Wenn Sie Ihre Arbeitslast überwachen, können Sie Folgendes sehen:

  • GPU-Nutzung: Bei einem fehlerfreien Feinabstimmungsjob sollte die Nutzung aller 16 GPUs (acht GPUs für jede VM im Cluster) während des Trainings steigen und sich auf einem bestimmten Niveau stabilisieren. TEST

  • Jobdauer: Der Job sollte etwa eine Stunde dauern.

Modell herunterladen

Nachdem Sie den Job erfolgreich ausgeführt haben, wird Ihr trainiertes Modell im Verzeichnis ~/qwen2-from-scratch-on-smollm-fineweb/ auf dem Login-Knoten gespeichert. Da dieses nichtflüchtige freigegebene Verzeichnis auf allen Knoten im Cluster bereitgestellt wird, bleiben Ihre Modellprüfpunkte auch nach Abschluss des Jobs oder nach der Aufhebung der Zuweisung der Compute-Knoten verfügbar.

Sie können das gespeicherte Modell mit dem Befehl gcloud compute scp vom Login-Knoten auf Ihren lokalen Computer herunterladen, wie im folgenden Beispiel gezeigt:

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

Nachdem Sie Ihr Modell heruntergeladen haben, können Sie Folgendes tun:

  • Modell für die Inferenz laden: Verwenden Sie das Hugging Face Transformers-Framework, um das Verzeichnis qwen2-trained-model/ zu laden und die Inferenz mit Ihrem neu trainierten Qwen2-Modell durchzuführen.
  • Zusätzliche Feinabstimmung: Verwenden Sie den gespeicherten Prüfpunkt als Ausgangspunkt für eine zusätzliche Feinabstimmung mit einem spezifischeren Dataset.
  • Modell in den Hugging Face-Hub übertragen: Teilen Sie Ihr trainiertes Modell, indem Sie es in den Hugging Face-Hub übertragen.

Bereinigen

Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.

Ressourcen löschen

  1. Führen Sie den Befehl gcluster destroy aus, um den Slurm-Cluster zu löschen:

    ./gcluster destroy "${CLUSTER_NAME}" --auto-approve

  2. So löschen Sie Ihren Cloud Storage-Bucket:

    gcloud storage buckets delete "gs://${GCS_BUCKET}" --quiet || true

  3. Wenn Sie ein Packer-Image löschen möchten, öffnen Sie Ihren Webbrowser, rufen Sie die folgende Seite auf, suchen Sie nach dem gewünschten Image und klicken Sie auf „Löschen“. 

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

  4. So löschen Sie alle VPC-Netzwerke, Firewallregeln, Router, IP-Adressen und Subnetze, die mit dem Projekt verknüpft sind:

    echo "========================================================================="
echo " STARTING AUTOMATED NETWORK CLEANUP FOR CLUSTER: ${CLUSTER_NAME}"
echo "========================================================================="

echo "Discovering all VPC networks linked to the cluster..."
NETWORKS=$(gcloud compute networks list --project="${PROJECT_ID}" --format="value(name)" | grep "^${CLUSTER_NAME}" || true)

if [ -z "${NETWORKS}" ]; then
    echo "No VPC networks found starting with ${CLUSTER_NAME}. Everything is already clean!"
    exit 0
fi

echo "Found the following networks to process:"
echo "${NETWORKS}"
echo "-------------------------------------------------------------------------"

echo "=== 1. Wiping Global Firewall Rules ==="
FIREWALL_RULES=$(gcloud compute firewall-rules list \
    --project="${PROJECT_ID}" \
    --filter="network ~ ^${CLUSTER_NAME} OR name ~ ^${CLUSTER_NAME}" \
    --format="value(name)" || echo "")

if [ -n "${FIREWALL_RULES}" ]; then
    echo "Deleting matching firewall rules:"
    echo "${FIREWALL_RULES}"
    echo "${FIREWALL_RULES}" | xargs -r gcloud compute firewall-rules delete --project="${PROJECT_ID}" --quiet
else
    echo "No matching firewall rules found."
fi

echo "=== 2. Tearing Down Network-Specific Infrastructure ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Processing resources for network: ${net_name}"

    ROUTERS=$(gcloud compute routers list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${ROUTERS}" ]; then
        echo "  -> Deleting routers: ${ROUTERS}"
        echo "${ROUTERS}" | xargs -r gcloud compute routers delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    IPS=$(gcloud compute addresses list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="name ~ ^${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${IPS}" ]; then
        echo "  -> Deleting IP reservations: ${IPS}"
        echo "${IPS}" | xargs -r gcloud compute addresses delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    SUBNETS=$(gcloud compute networks subnets list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${SUBNETS}" ]; then
        echo "  -> Deleting subnetworks:"
        echo "${SUBNETS}"
        echo "${SUBNETS}" | xargs -r gcloud compute networks subnets delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi
done

echo "-------------------------------------------------------------------------"
echo "Waiting 15 seconds for Google Cloud API dependencies to unlock..."
sleep 15

echo "=== 3. Final VPC Networks Destruction ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Deleting core VPC network: ${net_name}..."
    gcloud compute networks delete "${net_name}" --project="${PROJECT_ID}" --quiet || \
    echo "Warning: Could not delete ${net_name} yet. If a lock occurred, please rerun in 1 minute."
done

echo "========================================================================="
echo " SUCCESS: All network resources for cluster ${CLUSTER_NAME} have been wiped!"
echo "========================================================================="

Projekt löschen

Google Cloud -Projekt löschen:

gcloud projects delete PROJECT_ID

Nächste Schritte