FSDP zum Feinabstimmen von Mixtral-8x7B in einem A4-Slurm-Cluster verwenden

In dieser Anleitung wird beschrieben, wie Sie ein mistralai/Mixtral-8x7B-v0.1-Modell in einem Slurm-Cluster mit mehreren Knoten und mehreren GPUs auf Google Cloudabstimmen. Der Cluster verwendet zwei a4-highgpu-8g-VM-Instanzen mit jeweils 8 NVIDIA B200-GPUs.

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. Außerdem richten Sie ein freigegebenes Lustre-Dateisystem ein und konfigurieren Hochgeschwindigkeitsnetzwerke.
  2. Nachdem der Cluster bereitgestellt wurde, führen Sie einen verteilten Feinabstimmungsjob mit den Skripts aus, die in dieser Anleitung enthalten sind. Für den Job wird PyTorch Fully Sharded Data Parallel (FSDP) verwendet, auf das Sie über die Hugging Face-Bibliothek Transformer Reinforcement Learning (TRL) zugreifen.

Diese Anleitung richtet sich an Entwickler von maschinellem Lernen (ML), Forscher, Plattformadministratoren und ‑betreiber sowie Daten- und KI-Spezialisten, die daran interessiert sind, eine KI-Arbeitslast auf mehrere Knoten und GPUs zu verteilen.

Ziele

  • Über Hugging Face auf Mixtral zugreifen
  • Cluster Toolkit installieren
  • Umgebung vorbereiten
  • Einen produktionsreifen A4-High-GPU-Slurm-Cluster erstellen und bereitstellen
  • Eine Umgebung mit mehreren Knoten für verteiltes Training mit FSDP konfigurieren
  • Das Mixtral-Modell mit der Hugging Face-Klasse trl.SFTTrainer optimieren.
  • Daten auf lokalen SSDs bereitstellen.
  • 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 die Abrechnung für Ihr Google Cloud Projekt aktiviert ist.

  7. 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 lustre.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 die Abrechnung für Ihr Google Cloud Projekt aktiviert ist.

  13. 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 lustre.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.
  20. Installieren Sie die Abhängigkeiten, die Sie für die Verwendung des Cluster Toolkits benötigen.

Über Hugging Face auf Mixtral zugreifen

So greifen Sie über Hugging Face auf Mixtral zu:

  1. Melden Sie sich bei Hugging Face an und sehen Sie sich das Mixtral-Modell an.
  2. Erstellen Sie ein Hugging Face-Zugriffstoken für read.
  3. Kopieren und speichern Sie den Tokenwert. Sie benötigen sie später in dieser Anleitung.

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 ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export DEPLOYMENT_NAME="YOUR_CLUSTER_NAME"
export BUCKET_NAME="YOUR_GCS_BUCKET"
export HF_TOKEN="YOUR_HF_TOKEN"

    Ersetzen Sie Folgendes:

    • YOUR_PROJECT_ID: der Name des Google Cloud -Projekts, in dem Sie den Slurm-Cluster erstellen möchten.
    • YOUR_ZONE: die Zone, in der sich Ihre Reservierung befindet.
    • YOUR_REGION: die Region, in der Ihre Reservierung vorhanden ist.
    • YOUR_RESERVATION_NAME: Die URL oder der Name der Reservierung, die Sie zum Erstellen Ihres Slurm-Clusters verwenden möchten.
    • YOUR_CLUSTER_NAME: Der Name der Bereitstellung. Verwenden Sie einen kurzen Namen, der nur Buchstaben und Ziffern enthält (z. B. a4high). Dieser Name wird auch dem vom Deployment erstellten Slurm-Cluster zugewiesen.
    • YOUR_GCS_BUCKET: Der Name des Buckets, in dem Sie die Ergebnisse des Trainings-Checkpoints speichern. Geben Sie einen vorhandenen Bucket an oder erstellen Sie einen neuen. 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://${BUCKET_NAME}" \
  --project="${PROJECT_ID}"

A4-Slurm-Cluster erstellen

So erstellen Sie einen A4-Slurm-Cluster:

  1. Überschreiben Sie die Datei a4high-slurm-deployment.yaml im Verzeichnis examples/machine-learning/a4-highgpu-8g mit Ihren Umgebungsvariablen, indem Sie den folgenden Befehl ausführen:

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

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

  2. Öffnen Sie die Datei a4high-slurm-blueprint.yaml im Verzeichnis examples/machine-learning/a4-highgpu-8g und bearbeiten Sie sie so, dass Managed Lustre für das freigegebene Verzeichnis /home verwendet wird:

    • Entfernen Sie den Standardblock für das homefs-Modul mit source: modules/file-system/filestore.
    • Aktivieren Sie die Module lustrefs (homefs-Block mit remote_mount: lustrefs) und private-service-access.
    • Konfigurieren Sie im Block vars Folgendes:
      1. Ändern Sie den Wert für install_managed_lustre in /var/tmp/slurm_vars.json in true.
      2. Setzen Sie den Parameter per_unit_storage_throughput auf 500.
      3. Setzen Sie den Parameter lustre_size_gib auf 36000.
      4. Entfernen Sie die Auskommentierung von lustre_instance_id: lustre-instance.
      5. Kommentieren Sie nicht verwendete filestore_ip_range aus oder entfernen Sie sie.

  3. Stellen Sie den Cluster bereit:

    ./gcluster deploy \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  --auto-approve

    Mit dem Befehl ./gcluster deploy wird ein zweistufiger Prozess gestartet:

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

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 Arbeitslast zum Feinabstimmen verwendet werden:

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

    #!/bin/bash
# This script sets a reliable environment for FSDP training.
# It is meant to be run on a compute node.
set -e

# --- 1. Create the Python virtual environment ---
VENV_PATH="$HOME/.venv/venv-fsdp"
if [ ! -d "$VENV_PATH" ]; then
  echo "--- Creating Python virtual environment at $VENV_PATH ---"
  python3 -m venv $VENV_PATH
else
  echo "--- Virtual environment already exists at $VENV_PATH ---"
fi

source $VENV_PATH/bin/activate

# --- 2. Install Dependencies ---
echo "--- [STEP 2.1] Upgrading build toolchain ---"
pip install --upgrade pip wheel packaging

echo "--- [STEP 2.2] Installing PyTorch Nightly ---"
pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128

echo "--- [STEP 2.3] Installing application dependencies ---"
if [ -f "requirements-fsdp.txt" ]; then
    pip install -r requirements-fsdp.txt
else
    echo "ERROR: requirements-fsdp.txt not found!"
    exit 1
fi

# --- [STEP 2.4] Build Flash Attention from Source ---
echo "--- Building flash-attn from source... This will take a while. ---"
# Bypass PyTorch CUDA version mismatch check for extension builds
python3 -c "
import torch.utils.cpp_extension as ce
p = ce.__file__
content = open(p).read().replace('raise RuntimeError(CUDA_MISMATCH_MESSAGE', 'pass # raise RuntimeError')
open(p, 'w').write(content)
"

# Use all available CPU cores to speed up the build
MAX_JOBS=$(nproc) pip install flash-attn --no-build-isolation

# --- 3. Download the Model ---
echo "--- [STEP 2.5] Downloading Mixtral model ---"
if [ -z "$HF_TOKEN" ]; then
  echo "ERROR: The HF_TOKEN environment variable is not set."; exit 1;
fi
pip install huggingface_hub[cli]

# Execute the CLI using its full, explicit path
$VENV_PATH/bin/hf download mistralai/Mixtral-8x7B-v0.1 --local-dir ~/Mixtral-8x7B-v0.1 --token $HF_TOKEN

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

  2. Erstellen Sie eine requirements-fsdp.txt-Datei mit folgendem Inhalt, um die Python-Abhängigkeiten für das Trainingsskript anzugeben:

    transformers==4.55.0
datasets==4.0.0
peft==0.16.0
accelerate==1.9.0
trl==0.21.0

# Other dependencies
sentencepiece==0.2.0
protobuf==6.31.1

  3. Geben Sie train-mixtral.py als Haupttrainingsskript an:

    import torch
from torch.distributed.fsdp import MixedPrecision
from datasets import load_dataset
import shutil
import os
import torch.distributed as dist

from peft import LoraConfig, PeftModel, get_peft_model
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    HfArgumentParser,
)

from torch.distributed import get_rank, get_world_size

from transformers.models.mixtral.modeling_mixtral import MixtralDecoderLayer
from trl import SFTTrainer
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ScriptArguments:
    model_id: str = field(default="mistralai/Mixtral-8x7B-v0.1", metadata={"help": "Hugging Face model ID from the Hub"})
    dataset_name: str = field(default="philschmid/gretel-synthetic-text-to-sql", metadata={"help": "Dataset from the Hub"})
    run_inference_after_training: bool = field(default=False, metadata={"help": "Run sample inference on rank 0 after training"})
    dataset_subset_size: Optional[int] = field(default=None, metadata={"help": "Number of samples to use from the dataset for training. If None, uses the full dataset."})

@dataclass
class PeftArguments:
    lora_r: int = field(default=16, metadata={"help": "LoRA attention dimension"})
    lora_alpha: int = field(default=32, metadata={"help": "LoRA alpha scaling factor"})
    lora_dropout: float = field(default=0.05, metadata={"help": "LoRA dropout probability"})

@dataclass
class SftTrainingArguments(TrainingArguments):
    max_length: Optional[int] = field(default=2048, metadata={"help": "The maximum sequence length for SFTTrainer"})
    packing: Optional[bool] = field(default=False, metadata={"help": "Enable packing for SFTTrainer"})
    ddp_find_unused_parameters: Optional[bool] = field(default=False, metadata={"help": "When using FSDP activation checkpointing, this must be set to False for Mixtral"})

def formatting_prompts_func(example):
    system_message = "You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."
    user_prompt = f"### SCHEMA:\n{example['sql_context']}\n\n### USER QUERY:\n{example['sql_prompt']}"
    response = f"\n\n### SQL QUERY:\n{example['sql']}"
    return f"{system_message}\n\n{user_prompt}{response}"

def main():
    parser = HfArgumentParser((ScriptArguments, PeftArguments, SftTrainingArguments))
    script_args, peft_args, training_args = parser.parse_args_into_dataclasses()

    training_args.gradient_checkpointing = True
    training_args.gradient_checkpointing_kwargs = {"use_reentrant": True}

    training_args.optim = "adamw_torch_fused"

    bf16_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    )

    training_args.fsdp = "full_shard"
    training_args.fsdp_config = {
        "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
        "fsdp_transformer_layer_cls_to_wrap": [MixtralDecoderLayer],
        "fsdp_state_dict_type": "SHARDED_STATE_DICT",
        "fsdp_offload_params": False,
        "fsdp_forward_prefetch": True,
        "fsdp_mixed_precision_policy": bf16_policy
    }

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

    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

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

    peft_config = LoraConfig(
        r=peft_args.lora_r,
        lora_alpha=peft_args.lora_alpha,
        lora_dropout=peft_args.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    )

    model = get_peft_model(model, peft_config)

    data_splits = load_dataset(script_args.dataset_name)

    dataset = data_splits["train"]
    eval_dataset = data_splits["test"]

    if script_args.dataset_subset_size is not None:
        dataset = dataset.select(range(script_args.dataset_subset_size))

    dataset = dataset.shuffle(seed=training_args.seed)

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        eval_dataset=eval_dataset,
        formatting_func=formatting_prompts_func,
        processing_class=tokenizer,
    )

    trainer.train()

    dist.barrier()
    if trainer.is_world_process_zero():
        best_model_path = trainer.state.best_model_checkpoint

        final_model_dir = os.path.join(training_args.output_dir, "final_best_model")
        print(f"Copying best model to: {final_model_dir}")

        if os.path.exists(final_model_dir):
            shutil.rmtree(final_model_dir)
        shutil.copytree(best_model_path, final_model_dir)

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

def run_post_training_inference(script_args, best_model_path, tokenizer):
    print("\n" + "="*50)
    print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
    print("="*50 + "\n")

    base_model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto"
    )
    model = PeftModel.from_pretrained(base_model, best_model_path)
    model = model.merge_and_unload()
    model.eval()

    # Define the test case
    schema = "CREATE TABLE artists (Name TEXT, Country TEXT, Genre TEXT)"
    system_message = "You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."
    question = "Show me all artists from the Country just north of the USA."

    prompt = f"{system_message}\n\n### SCHEMA:\n{schema}\n\n### USER QUERY:\n{question}\n\n### SQL QUERY:\n"

    print(f"Test Prompt:\n{prompt}")

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    print("\n--- Generating SQL... ---")
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=False,
        temperature=None,
        top_p=None,
    )

    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):].strip()

    print(f"\n--- Generated SQL Query ---")
    print(generated_sql)
    print("\n" + "="*50)
    print("=== INFERENCE TEST COMPLETE ===")
    print("="*50 + "\n")

if __name__ == "__main__":
    main()

  4. Erstellen Sie die Datei train-mixtral.sh 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=mixtral-fsdp
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --partition=a4high
#SBATCH --output=mixtral-%j.out
#SBATCH --error=mixtral-%j.err

set -e
set -x

echo "--- Slurm Job Started ---"

# --- Define Paths ---
LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
VENV_PATH="${HOME}/.venv/venv-fsdp"
MODEL_PATH="${HOME}/Mixtral-8x7B-v0.1"

# --- STAGE 1: Stage Data to Local SSD on Each Node ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c "
echo '--- Staging on node: $(hostname) ---'
mkdir -p ${LOCAL_SSD_PATH}

echo 'Copying virtual environment...'
rsync -a -q ${VENV_PATH}/ ${LOCAL_SSD_PATH}/venv/

echo 'Copying model weights...'
rsync -a ${MODEL_PATH}/ ${LOCAL_SSD_PATH}/model/

mkdir -p ${LOCAL_SSD_PATH}/hf_cache

echo '--- Staging on $(hostname) complete ---'
"
echo "--- Staging complete on all nodes ---"

# --- STAGE 2: Run the Training Job ---
echo "--- Launching Distributed Training with GIB NCCL Plugin ---"
nodes=( $( scontrol show hostnames "$SLURM_JOB_NODELIST" ) )
head_node=${nodes[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

export MASTER_ADDR=$head_node_ip
export MASTER_PORT=29500

export NCCL_SOCKET_IFNAME=enp0s19

export NCCL_NET=gIB

# export NCCL_DEBUG=INFO # Un-comment to diagnose NCCL issues if needed

srun --cpu-bind=none --accel-bind=g bash -c '
# Activate the environment from the local copy
source '${LOCAL_SSD_PATH}'/venv/bin/activate

# Point Hugging Face cache to the local SSD
export HF_HOME='${LOCAL_SSD_PATH}'/hf_cache

export RANK=$SLURM_PROCID
export WORLD_SIZE=$SLURM_NTASKS
export LOCAL_RANK=$SLURM_LOCALID

export LD_LIBRARY_PATH=/usr/local/gib/lib64:$LD_LIBRARY_PATH
source /usr/local/gib/scripts/set_nccl_env.sh

# --- Launch the training ---
python \
    '${SLURM_SUBMIT_DIR}'/train-mixtral.py \
    --model_id="'${LOCAL_SSD_PATH}'/model/" \
    --output_dir="${HOME}/outputs/mixtral_job_${SLURM_JOB_ID}" \
    --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
    --seed=900913 \
    --bf16=True \
    --num_train_epochs=3 \
    --per_device_train_batch_size=32 \
    --gradient_accumulation_steps=4 \
    --learning_rate=4e-5 \
    --logging_steps=3 \
    --lora_r=32 \
    --lora_alpha=32 \
    --lora_dropout=0.05 \
    --eval_strategy=steps \
    --eval_steps=10 \
    --save_strategy=steps \
    --save_steps=10 \
    --load_best_model_at_end=False \
    --metric_for_best_model=eval_loss \
    --run_inference_after_training \
    --dataset_subset_size=67000
'

# --- STAGE 3: Cleanup ---
echo "--- Cleaning up local SSD on all nodes ---"
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c "rm -rf ${LOCAL_SSD_PATH}"

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

Skripts in Ihren 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='${DEPLOYMENT_NAME}' AND labels.slurm_instance_role='login'" \
                --format="value(name)" | head -n 1)"

    Die Variable LOGIN_NODE speichert einen Wert ähnlich wie ${DEPLOYMENT_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-fsdp.txt \
  ./train-mixtral.py \
  ./train-mixtral.sh \
  "${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 die erforderlichen Frameworks und Tools, indem Sie das Installationsskript auf einem Rechenknoten ausführen:

# On the login node
srun \
  --job-name=env-setup \
  --nodes=1 \
  --ntasks=1 \
  --gpus-per-node=1 \
  --partition=a4high \
  bash ./install_environment.sh

Mit diesem Befehl wird die virtuelle Umgebung eingerichtet, alle Abhängigkeiten werden installiert und die Mixtral-Modellgewichte werden in ~/Mixtral-8x7B-v0.1 heruntergeladen. Das kann über 30 Minuten dauern.

Wenn Sie diesen Schritt über srun ausführen, wird das Installationsskript vom Anmeldeknoten an einen zugewiesenen Rechenknoten delegiert. So kann das Skript während der Kompilierung auf GPU-Treiber zugreifen.

Arbeitslast für das Feinabstimmen starten

So starten Sie Ihre Trainingsarbeitslast:

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

    # On the login node
sbatch train-mixtral.sh

Arbeitslast überwachen

So können Sie den Fortschritt Ihres Feinabstimmungsjobs überwachen:

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

    # On the login node:
tail -f mixtral-*.out mixtral-*.err

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

    Der Job besteht aus zwei Hauptphasen:

    • Das große Basismodell wird auf die lokale SSD jedes Rechenknotens kopiert.
    • Der Trainingsjob, der beginnt, sobald das Kopieren des Modells abgeschlossen ist.

    Die Ausführung des gesamten Jobs dauert etwa 60 Minuten.

  2. Sie können auch 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=YOUR_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.

    • Jobdauer: Der Job sollte etwa eine Stunde dauern.

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

Führen Sie die folgenden Befehle aus, um die Ressourcen zu löschen, die Sie in dieser Anleitung erstellt haben:

  1. Wenn Sie Ihren Slurm-Cluster löschen möchten, wechseln Sie zum Verzeichnis cluster-toolkit und führen Sie den folgenden Befehl aus:

    ./gcluster destroy "${DEPLOYMENT_NAME}" --auto-approve --robust

  2. Löschen Sie den Bucket:

    gcloud storage rm --recursive "gs://${BUCKET_NAME}" --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

Projekt löschen

Google Cloud -Projekt löschen:

gcloud projects delete PROJECT_ID

Nächste Schritte