FSDP zum Feinabstimmen von Llama 4 in einem A4-Slurm-Cluster verwenden

In dieser Anleitung wird gezeigt, wie Sie ein Llama-4-Scout-17-LLM (Large Language Model) in einem Slurm-Cluster mit mehreren Knoten und mehreren GPUs auf Google Cloudfeinabstimmen. Der Cluster verwendet zwei A4-VM-Instanzen (virtuelle Maschinen), die jeweils 8 NVIDIA B200-GPUs haben.

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

  1. Stellen Sie einen leistungsstarken Slurm-Cluster für die Produktion 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 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 zugreifen.

Diese Anleitung richtet sich an Entwickler von maschinellem Lernen (ML), Plattformadministratoren und ‑operatoren sowie an Daten- und KI-Spezialisten, die daran interessiert sind, Slurm-Funktionen zur Jobplanung für die Verarbeitung von Arbeitslasten zum Feinabstimmen zu verwenden.

Ziele

  • Über Hugging Face auf Llama 4 zugreifen

  • Umgebung vorbereiten

  • Erstellen und stellen Sie einen A4 High-GPU-Slurm-Cluster für die Produktion bereit.

  • Eine Umgebung mit mehreren Knoten für verteiltes Training mit FSDP konfigurieren

  • Optimieren Sie das Llama 4-Modell mit Hugging Face trl.SFTTrainer.

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

Über Hugging Face auf Llama 4 zugreifen

So verwenden Sie Hugging Face, um auf Llama 4 zuzugreifen:

  1. Einwilligungsvereinbarung unterzeichnen, um Llama 4 zu verwenden

  2. Erstellen Sie ein Hugging Face-Zugriffstoken für read.

    Klicken Sie auf Mein Profil > Einstellungen > Zugriffstokens > +Neuen Token erstellen.

  3. Kopieren und speichern Sie den read access-Tokenwert. Sie benötigen sie später in dieser Anleitung.

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

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

    #!/bin/bash
# This script sets up a consistent environment for FSDP training.
# It is meant to be run once on the login node of your Slurm cluster
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 torchvision torchaudio --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

# --- 3. Download the Model ---
echo "--- [STEP 2.4] Downloading Llama4 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/huggingface-cli download meta-llama/Llama-4-Scout-17B-16E-Instruct --local-dir ~/Llama-4-Scout-17B-16E-Instruct --token $HF_TOKEN

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

    Dieses Skript richtet eine zuverlässige virtuelle Python-Umgebung ein, installiert einen PyTorch-Nightly-Build und lädt das Llama 4-Modell herunter.

  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

  3. Geben Sie llama4-train-distributed.py als Haupttrainingsskript an:

    import torch
from datasets import load_dataset
from peft import LoraConfig, PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    HfArgumentParser,
)

from torch.distributed import get_rank, get_world_size

from transformers.models.llama4.modeling_llama4 import Llama4TextDecoderLayer
from trl import SFTTrainer
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class ScriptArguments:
    model_id: str = field(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=True, metadata={"help": "When using FSDP activation checkpointing, this must be set to True"})

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": False}

    training_args.optim = "adamw_torch_fused"

    training_args.fsdp = "full_shard"
    training_args.fsdp_config = {
        "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
        "fsdp_transformer_layer_cls_to_wrap": [Llama4TextDecoderLayer],
        "fsdp_state_dict_type": "FULL_STATE_DICT",
        "fsdp_offload_params": False,
        "fsdp_forward_prefetch": True,
    }

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

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

    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"],
    )
    rank = get_rank()
    world_size = get_world_size()

    dataset = load_dataset(script_args.dataset_name, split="train")

    if script_args.dataset_subset_size is not None:
        dataset = dataset.select(range(script_args.dataset_subset_size))
    else:
        print(f"Using the full dataset with {len(dataset)} samples.")

    dataset = dataset.shuffle(seed=training_args.seed)
    print(f"Dataset shuffled with seed: {training_args.seed}.")

    if world_size > 1:
        print(f"Sharding dataset for Rank {rank} of {world_size}.")
        dataset = dataset.shard(num_shards=world_size, index=rank)

    print("Initializing SFTTrainer...")
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        peft_config=peft_config,
        formatting_func=formatting_prompts_func,
        processing_class=tokenizer,
    )

    trainer.train()

    trainer.save_model(training_args.output_dir)

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

def run_post_training_inference(script_args, training_args, tokenizer):
    """
    Loads the fine-tuned PEFT adapter from the local output directory and runs inference.
    This should only be called on rank 0 after training is complete.
    """
    print("\n" + "="*50)
    print("=== RUNNING POST-TRAINING INFERENCE TEST ===")
    print("="*50 + "\n")

    # Load the base model and merge the adapter.
    base_model = AutoModelForCausalLM.from_pretrained(
        script_args.model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="auto"
    )
    # Load the PEFT adapter and merge it into the base model
    model = PeftModel.from_pretrained(base_model, training_args.output_dir)
    model = model.merge_and_unload() # Merge weights for faster inference
    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."

    # This must match the formatting_func exactly
    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()

    In diesem Skript wird der TRL Supervised Fine-Tuning (SFT) Trainer verwendet, um FSDP-Trainingsschleifen, die LoRA-Konfiguration (Low-Rank Adaptation) und die Datenformatierung zu verwalten.

  4. 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=llama4-fsdp-fixed
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8
#SBATCH --partition=a4high
#SBATCH --output=llama4-%j.out
#SBATCH --error=llama4-%j.err

set -e
set -x

echo "--- Slurm Job Started ---"
echo "Job ID: $SLURM_JOB_ID"
echo "Node List: $SLURM_JOB_NODELIST"

# --- Define Paths ---
LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}"
VENV_PATH="${HOME}/.venv/venv-fsdp"
MODEL_PATH="${HOME}/Llama-4-Scout-17B-16E-Instruct"

# --- 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 --info=progress2 ${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}'/llama4-train-distributed.py \
      --model_id="'${LOCAL_SSD_PATH}'/model/" \
      --output_dir="'${LOCAL_SSD_PATH}'/outputs/" \
      --dataset_name="philschmid/gretel-synthetic-text-to-sql" \
      --seed=900913 \
      --bf16=True \
      --num_train_epochs=1 \
      --per_device_train_batch_size=2 \
      --gradient_accumulation_steps=4 \
      --learning_rate=2e-5 \
      --logging_steps=10 \
      --lora_r=16 \
      --lora_alpha=32 \
      --lora_dropout=0.05 \
      --run_inference_after_training
'

# --- STAGE 3: Copy Final Results Back to Persistent Storage ---
echo "--- Copying final results from local SSD to shared storage ---"
PERSISTENT_OUTPUT_DIR="${HOME}/outputs/llama4_job_${SLURM_JOB_ID}"
mkdir -p "$PERSISTENT_OUTPUT_DIR"

# Only copy from the head node where trl has combined the results
srun --nodes=1 --ntasks=1 -w "$head_node" \
  rsync -a --info=progress2 "${LOCAL_SSD_PATH}/outputs/" "${PERSISTENT_OUTPUT_DIR}/"

# --- STAGE 4: 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 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 \
  ./install_environment.sh \
  ./requirements-fsdp.txt \
  ./llama4-train-distributed.py \
  ./submit.slurm \
  "${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. Hugging Face-Token exportieren:

    # On the login node
export HF_TOKEN="hf_..." # Replace with your token

  2. Führen Sie das Installationsskript aus:

    # On the login node
chmod +x install_environment.sh
./install_environment.sh

    Mit diesem Befehl wird eine virtuelle Umgebung mit allen erforderlichen Abhängigkeiten eingerichtet und die Modellgewichte werden in die Datei ~/Llama-4-Scout-17B-16E-Instruct heruntergeladen.

    Da der Model-Download sehr groß ist (ca. 200 GB), dauert dieser Vorgang je nach Netzwerkbedingungen etwa 30 Minuten.

Arbeitslast für das Feinabstimmen starten

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:

    # On the login node
tail -f llama4-*.out

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

    Dieser Job sollte auf dem Slurm-Cluster etwas mehr als eine Stunde dauern. 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 dieses Jobs dauert etwa 35 Minuten.

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

Filestore-Instanz löschen

Standardmäßig ist für Ihre Filestore-Instanz in der cluster-toolkit-Blaupause die Einstellung deletion_protection auf „true“ gesetzt. Diese Einstellung verhindert versehentlichen Datenverlust, wenn Sie Umgebungen ändern. Wenn Sie die Filestore-Instanz löschen möchten, müssen Sie den Löschschutz manuell deaktivieren.

Nächste Schritte