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 auf Google Cloudtrainieren. Das in dieser Anleitung verwendete Modell basiert auf einem Qwen2-Modell mit 1,5 Milliarden Parametern. Der Slurm-Cluster verwendet zwei virtuelle Maschinen (VMs) vom Typ „a4-highgpu-8g“, die jeweils 8 NVIDIA B200-GPUs haben.

Die beiden in dieser Anleitung beschriebenen Hauptprozesse sind:

  1. Stellen Sie einen leistungsstarken Slurm-Cluster in Produktionsqualität mit demGoogle Cloud Cluster Toolkit bereit. Im Rahmen dieser Bereitstellung erstellen Sie ein benutzerdefiniertes VM-Image mit der erforderlichen vorinstallierten Software. Sie richten außerdem eine freigegebene Filestore-Instanz ein und konfigurieren das RDMA-Hochgeschwindigkeitsnetzwerk.
  2. Nachdem der Cluster bereitgestellt wurde, führen Sie einen verteilten Vorabtrainingsjob mit den Skripts aus, die in dieser Anleitung enthalten sind. Für den Job wird die Hugging Face Accelerate-Bibliothek verwendet.

Diese Anleitung richtet sich an Entwickler von maschinellem Lernen (ML), Forscher, Plattformadministratoren und ‑betreiber sowie an Daten- und KI-Spezialisten, die daran interessiert sind, leistungsstarke Slurm-Cluster auf Google Cloud bereitzustellen, um LLMs zu trainieren.

Ziele

  • Über Hugging Face auf das Qwen2-Modell zugreifen
  • Bereiten Sie Ihre Umgebung vor.
  • Erstellen und stellen Sie einen produktionsreifen A4-Slurm-Cluster bereit.
  • 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 Sie ein Google Cloud Projekt oder wählen Sie eines aus.

    Rollen, die zum Auswählen oder Erstellen eines Projekts erforderlich sind

    • 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 „Projektersteller“ (roles/resourcemanager.projectCreator), die die Berechtigung resourcemanager.projects.create enthält. Weitere Informationen zum Zuweisen von Rollen

    • So erstellen Sie ein Google Cloud Projekt:

      gcloud projects create PROJECT_ID

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

    • Wählen Sie das von Ihnen erstellte Google Cloud Projekt 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:

    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 Rolle „Inhaber“ (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 Abrufen 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 Anmeldedaten zur Authentifizierung 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 ein Konto.

Über Hugging Face auf Qwen2 zugreifen

So greifen Sie über Hugging Face auf Qwen2 zu:

  1. Einwilligungsvereinbarung unterzeichnen, um Qwen 2 1.5B zu verwenden

  2. Erstellen Sie ein read-Zugriffstoken.

Cluster Toolkit installieren

Cluster Toolkit ist ein Open-Source-Tool, das die Bereitstellung von Arbeitslasten aus den Bereichen Hochleistungs-Computing (HPC), künstliche Intelligenz (KI) und maschinelles Lernen (ML) auf Google Cloudvereinfacht. Weitere Informationen zur Verwendung von gcluster und zur Verwaltung von Clustern finden Sie in der Übersicht zum Cluster Toolkit.

  1. Cluster Toolkit-Version vorbereiten:

    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 das 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="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 ProjektsYOUR_PROJECT_ID, 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 Ihre Reservierung vorhanden ist.
    • RESERVATION_NAME: Die URL oder der Name der Reservierung, die Sie zum Erstellen Ihres Slurm-Clusters verwenden möchten.
    • YOUR_GCS_BUCKET: Der Name des Buckets, in dem Sie die Ergebnisse des Trainings-Checkpoints speichern. Bevor Sie einen Bucket erstellen, sollten Sie sich mit den Anforderungen für Bucket-Namen vertraut machen.
    • YOUR_HF_TOKEN: das Hugging Face-Token, das Sie in einem früheren 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. Erstellen Sie die Terraform-Manifeste:

    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 Befehl gcluster deploy ist ein zweiphasiger Prozess:

    • In der ersten Phase wird ein benutzerdefiniertes Image mit vorinstallierter Software erstellt. Das 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 jedoch fehlschlägt, können Sie versuchen, den Slurm-Cluster noch einmal bereitzustellen und dabei die erste Phase zu ü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 dem folgenden Inhalt, um die Konfigurationen für den 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 dem folgenden Inhalt, um die Aufgaben anzugeben, die für die Jobs 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 neuesten 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 Format vorverarbeiten möchten, das für das Training geeignet ist, 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 Anmeldenknotens 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)"

    Die Variable LOGIN_NODE speichert einen Wert ähnlich wie ${CLUSTER_NAME}-login-001.

  2. Firewallregel erstellen

    gcloud compute firewall-rules create allow-ssh-ingress-from-iap \
  --project="${PROJECT_ID}" \
  --network="${CLUSTER_NETWORK}" \
  --direction=INGRESS \
  --action=allow \
  --rules=tcp:22 \
  --source-ranges=35.235.240.0/20 \
  --description="Allow SSH ingress from Google Cloud Identity-Aware Proxy (IAP)"

  3. Laden Sie Ihre Skripts in das Basisverzeichnis des Anmeldeknotens 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 Anmeldeknoten 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 Anmeldeknoten hergestellt haben, installieren Sie Frameworks und Tools, indem Sie Folgendes tun:

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

    chmod +x install_environment.sh
./install_environment.sh

Arbeitslast vortrainieren

So starten Sie das Training Ihrer Arbeitslast:

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

    sbatch submit.slurm

  2. Auf dem Anmeldeknoten in Ihrem Slurm-Cluster können Sie den Fortschritt des Jobs überwachen, indem Sie die Ausgabedateien prüfen, die in Ihrem home-Verzeichnis erstellt wurden:

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

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

Arbeitslast überwachen

Sie können die Nutzung der GPUs in Ihrem Slurm-Cluster überwachen, um zu prüfen, ob Ihr Fine-Tuning-Job 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, sehen Sie Folgendes:

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

  • Jobdauer: Der Job sollte etwa eine Stunde dauern.

Modell herunterladen

Nachdem Sie Ihren Job erfolgreich ausgeführt haben, wird Ihr trainiertes Modell im Verzeichnis ~/qwen2-from-scratch-on-smollm-fineweb/ auf dem Log-in-Knoten gespeichert. Da dieses persistente freigegebene Verzeichnis auf allen Knoten in Ihrem Cluster bereitgestellt wird, bleiben Ihre Modellprüfpunkte auch nach Abschluss des Jobs oder nach der Freigabe der Rechenknoten verfügbar.

Sie können das gespeicherte Modell mit dem Befehl gcloud compute scp vom Anmeldenode 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, haben Sie folgende Möglichkeiten:

  • 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 auszuführen.
  • Zusätzliches Fine-Tuning: Verwenden Sie den gespeicherten Prüfpunkt als Ausgangspunkt für zusätzliches Fine-Tuning mit einem spezifischeren Dataset.
  • Modell in den Hugging Face-Hub hochladen: Geben Sie Ihr trainiertes Modell frei, indem Sie es in den Hugging Face-Hub hochladen.

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. So löschen Sie Ihren Slurm-Cluster:

    ./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