Qwen2 in einem A4-Slurm-Cluster trainieren

In dieser Anleitung erfahren Sie, wie Sie ein Large Language Model (LLM) in einem Slurm-Cluster mit mehreren Knoten und mehreren GPUs auf Google Cloudtrainieren. Das Modell, das Sie in diesem Tutorial verwenden, 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 Hauptprozesse, die in dieser Anleitung beschrieben werden, 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 auch eine freigegebene Filestore-Instanz ein und konfigurieren das RDMA-Hochgeschwindigkeitsnetzwerk.
  2. Nachdem der Cluster bereitgestellt wurde, führen Sie einen verteilten Pre-Training-Job 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 für maschinelles Lernen (ML), Forscher, Plattformadministratoren und ‑operatoren sowie an Daten- und KI-Spezialisten, die daran interessiert sind, leistungsstarke Slurm-Cluster in 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. Melden Sie sich in Ihrem Google Cloud -Konto an. Wenn Sie mit Google Cloudnoch nicht vertraut sind, erstellen Sie ein Konto, um die Leistungsfähigkeit unserer Produkte in der Praxis sehen und bewerten zu können. Neukunden erhalten außerdem ein Guthaben von 300 $, um Arbeitslasten auszuführen, zu testen und bereitzustellen.

  2. Installieren Sie die Google Cloud CLI.

  3. Wenn Sie einen externen Identitätsanbieter (IdP) verwenden, müssen Sie sich zuerst mit Ihrer föderierten Identität in der gcloud CLI anmelden.

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

    gcloud init

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

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

  7. Aktivieren Sie die erforderliche API:

    Rollen, die zum Aktivieren von APIs erforderlich sind

    Zum Aktivieren von APIs benötigen Sie die IAM-Rolle „Service Usage-Administrator“ (roles/serviceusage.serviceUsageAdmin), die die Berechtigung serviceusage.services.enable enthält. Weitere Informationen zum Zuweisen von Rollen 

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

  8. Installieren Sie die Google Cloud CLI.

  9. Wenn Sie einen externen Identitätsanbieter (IdP) verwenden, müssen Sie sich zuerst mit Ihrer föderierten Identität in der gcloud CLI anmelden.

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

    gcloud init

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

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

  13. Aktivieren Sie die erforderliche API:

    Rollen, die zum Aktivieren von APIs erforderlich sind

    Zum Aktivieren von APIs benötigen Sie die IAM-Rolle „Service Usage-Administrator“ (roles/serviceusage.serviceUsageAdmin), die die Berechtigung serviceusage.services.enable enthält. Weitere Informationen zum Zuweisen von Rollen 

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

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

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

    Ersetzen Sie Folgendes:

    • PROJECT_ID: Ihre Projekt-ID.
    • USER_IDENTIFIER: Die Kennung für Ihr Nutzerkonto . Beispiel: myemail@example.com
    • ROLE: Die IAM-Rolle, die Sie Ihrem Nutzerkonto zuweisen.
  15. Aktivieren Sie das Standarddienstkonto für Ihr Google Cloud Projekt: 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --project=PROJECT_ID

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

  16. Weisen Sie dem Standarddienstkonto die Rolle „Bearbeiter“ (roles/editor) zu: 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \
    --role=roles/editor
  17. Erstellen Sie lokale Anmeldedaten zur Authentifizierung für Ihr Nutzerkonto: 
    gcloud auth application-default login
  18. Aktivieren Sie OS Login für Ihr Projekt: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. 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.

Umgebung vorbereiten

So bereiten Sie die Umgebung vor:

  1. Klonen Sie das GitHub-Repository für Cluster-Toolkit:

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

  2. Erstellen Sie einen Cloud Storage-Bucket:

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

    Ersetzen Sie Folgendes:

    • BUCKET_NAME: Ein Name für Ihren Cloud Storage-Bucket, der den Anforderungen für Bucket-Namen entspricht.

    • PROJECT_ID: Die ID desGoogle Cloud -Projekts, in dem Sie Ihren Cloud Storage-Bucket erstellen möchten.

A4-Slurm-Cluster erstellen

So erstellen Sie einen A4-Slurm-Cluster:

  1. Wechseln Sie in das Verzeichnis cluster-toolkit:

    cd cluster-toolkit

  2. Wenn Sie das Cluster-Toolkit zum ersten Mal verwenden, erstellen Sie das Binärprogramm gcluster:

    make

  3. Wechseln Sie in das Verzeichnis examples/machine-learning/a4-highgpu-8g:

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

  4. Öffnen Sie die Datei a4high-slurm-deployment.yaml und bearbeiten Sie sie so:

    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

    Ersetzen Sie Folgendes:

    • BUCKET_NAME: Der Name des Cloud Storage-Bucket, den Sie im vorherigen Abschnitt erstellt haben.

    • PROJECT_ID: Die ID desGoogle Cloud -Projekts, in dem sich Ihr Cloud Storage befindet und in dem Sie Ihren Slurm-Cluster erstellen möchten.

    • REGION: Die Region, in der Ihre Reservierung vorhanden ist.

    • ZONE: die Zone, in der sich Ihre Reservierung befindet.

    • RESERVATION_URL: die URL der Reservierung, die Sie zum Erstellen Ihres Slurm-Clusters verwenden möchten. Geben Sie je nach Projekt, in dem die Reservierung vorhanden ist, einen der folgenden Werte an:

      • Die Reservierung ist in Ihrem Projekt vorhanden: RESERVATION_NAME

      • Die Reservierung ist in einem anderen Projekt vorhanden und Ihr Projekt kann sie nutzen: projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME

  5. Stellen Sie den Cluster bereit:

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

    Der Befehl ./gcluster deploy umfasst zwei Phasen:

    • In der ersten Phase wird ein benutzerdefiniertes Image mit vorinstallierter Software erstellt. Das kann bis zu 35 Minuten dauern.

    • In der zweiten Phase wird der Cluster mit diesem benutzerdefinierten Image bereitgestellt. Dieser Vorgang sollte schneller abgeschlossen sein 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 -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

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. Um die virtuelle Python-Umgebung einzurichten, erstellen Sie die Datei install_environment.sh mit folgendem Inhalt:

    #!/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 Ihren Feinabstimmungsjob anzugeben:

    # 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. 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:

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

  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 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()

Skripts in den Slurm-Cluster hochladen

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

  1. So ermitteln Sie Ihren Anmeldeknoten: Listen Sie alle A4-VMs in Ihrem Projekt auf:

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

    Der Name des Anmeldeknotens ähnelt a4-high-login-001.

  2. Laden Sie Ihre Skripts in das Basisverzeichnis des Anmeldeknotens hoch:

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

    Ersetzen Sie LOGIN_NODE_NAME durch den Namen des Anmeldeknotens.

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

Frameworks und Tools installieren

Nachdem Sie eine Verbindung zum Anmeldeknoten hergestellt haben, installieren Sie Frameworks und Tools, indem Sie Folgendes tun:

  1. Erstellen Sie eine Umgebungsvariable für Ihr Hugging Face-Zugriffstoken:

    export HUGGING_FACE_TOKEN="HUGGING_FACE_TOKEN"

  2. 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-qwen2-pretrain-smollm-fineweb.err

    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

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.

  • 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 Anmeldeknoten auf Ihren lokalen Computer herunterladen, wie im folgenden Beispiel gezeigt:

# 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

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.

Projekt löschen

Google Cloud -Projekt löschen:

gcloud projects delete PROJECT_ID

Slurm-Cluster löschen

So löschen Sie Ihren Slurm-Cluster:

  1. Wechseln Sie zum Verzeichnis cluster-toolkit.

  2. Löschen Sie die Terraform-Datei und alle erstellten Ressourcen:

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

Nächste Schritte