Utilizzare FSDP per ottimizzare Llama 4 su un cluster Slurm A4

Questo tutorial mostra come eseguire il fine tuning di un modello linguistico di grandi dimensioni (LLM) Llama-4-Scout-17 su un cluster Slurm multi-nodo e multi-GPU su Google Cloud. Il cluster utilizza due istanze di macchine virtuali (VM) A4, ognuna delle quali ha 8 GPU NVIDIA B200.

Le due procedure principali descritte in questo tutorial sono le seguenti:

  1. Esegui il deployment di un cluster Slurm ad alte prestazioni di livello di produzione utilizzando Google Cloud Cluster Toolkit. Nell'ambito di questo deployment, crei un'immagine VM personalizzata con il software necessario preinstallato. Configura anche un'istanza di Filestore condivisa e la rete RDMA ad alta velocità.
  2. Dopo il deployment del cluster, esegui un job di messa a punto distribuita utilizzando il set di script che accompagnano questo tutorial. Il job utilizza PyTorch Fully Sharded Data Parallel (FSDP), a cui accedi tramite Hugging Face Transformer Reinforcement Learning

Questo tutorial è rivolto a ingegneri di machine learning (ML), amministratori e operatori di piattaforme e specialisti di dati e AI interessati a utilizzare le funzionalità di pianificazione dei job Slurm per gestire i workload di fine-tuning.

Obiettivi

  • Accedere a Llama 4 utilizzando Hugging Face

  • prepara l'ambiente

  • Crea ed esegui il deployment di un cluster Slurm A4 High-GPU di livello di produzione.

  • Configura un ambiente multi-nodo per l'addestramento distribuito con FSDP.

  • Perfeziona il modello Llama 4 utilizzando Hugging Face trl.SFTTrainer.

  • Esegui lo staging dei dati sulle unità SSD locali.

  • Monitorare il job.

  • Eseguire la pulizia.

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

I nuovi utenti di Google Cloud potrebbero avere diritto a una prova senza costi.

Prima di iniziare

  1. Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei workload.

  2. Installa Google Cloud CLI.

  3. Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

  4. Per inizializzare gcloud CLI, esegui questo comando: 

    gcloud init

  5. Crea o seleziona un Google Cloud progetto.

    Ruoli richiesti per selezionare o creare un progetto

    • Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

    • Creare un progetto Google Cloud :

      gcloud projects create PROJECT_ID

      Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

    • Seleziona il progetto Google Cloud che hai creato:

      gcloud config set project PROJECT_ID

      Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

  6. Verifica che la fatturazione sia abilitata per il tuo progetto Google Cloud .

  7. Abilita l'API richiesta:

    Ruoli richiesti per abilitare le API

    Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli. 

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

  8. Installa Google Cloud CLI.

  9. Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

  10. Per inizializzare gcloud CLI, esegui questo comando: 

    gcloud init

  11. Crea o seleziona un Google Cloud progetto.

    Ruoli richiesti per selezionare o creare un progetto

    • Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
    • Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

    • Creare un progetto Google Cloud :

      gcloud projects create PROJECT_ID

      Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

    • Seleziona il progetto Google Cloud che hai creato:

      gcloud config set project PROJECT_ID

      Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

  12. Verifica che la fatturazione sia abilitata per il tuo progetto Google Cloud .

  13. Abilita l'API richiesta:

    Ruoli richiesti per abilitare le API

    Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli. 

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

  14. Concedi ruoli al tuo account utente. Esegui il seguente comando una volta per ciascuno dei seguenti ruoli IAM: 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

    Sostituisci quanto segue:

    • PROJECT_ID: il tuo ID progetto.
    • USER_IDENTIFIER: l'identificatore del tuo account utente . Ad esempio: myemail@example.com.
    • ROLE: il ruolo IAM che concedi al tuo account utente.
  15. Abilita il account di servizio predefinito per il tuo progetto Google Cloud : 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com 
    
    --project=PROJECT_ID

    Sostituisci PROJECT_NUMBER con il numero del progetto. Per rivedere il numero del progetto, consulta Recuperare un progetto esistente.

  16. Concedi il ruolo Editor (roles/editor) al service account predefinito: 
    gcloud projects add-iam-policy-binding PROJECT_ID 
    
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" 
    
    --role=roles/editor
  17. Crea le credenziali di autenticazione locale per il tuo account utente: 
    gcloud auth application-default login
  18. Attiva OS Login per il tuo progetto: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. Accedi o crea un account Hugging Face.
  20. Installa le dipendenze necessarie per utilizzare Cluster Toolkit.

Accedere a Llama 4 utilizzando Hugging Face

Per utilizzare Hugging Face per accedere a Llama 4:

  1. Firma il contratto di consenso per utilizzare Llama 4.

  2. Crea un token di accesso read di Hugging Face.

    Fai clic su Il tuo profilo > Impostazioni > Token di accesso > +Crea nuovo token.

  3. Copia e salva il valore del token read access. Lo utilizzerai più avanti in questo tutorial.

prepara l'ambiente

Per preparare l'ambiente:

  1. Clona il repository GitHub di Cluster Toolkit:

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

  2. Crea un bucket Cloud Storage:

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

    Sostituisci quanto segue:

    • BUCKET_NAME: un nome per il bucket Cloud Storage che rispetti i requisiti di denominazione dei bucket.

    • PROJECT_ID: l'ID del Google Cloud progetto in cui vuoi creare il bucket Cloud Storage.

Crea un cluster Slurm A4

Per creare un cluster Slurm A4:

  1. Vai alla directory cluster-toolkit:

    cd cluster-toolkit

  2. Se è la prima volta che utilizzi Cluster Toolkit, crea il file binario gcluster:

    make

  3. Vai alla directory examples/machine-learning/a4-highgpu-8g:

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

  4. Apri il file a4high-slurm-deployment.yaml e modificalo nel seguente modo:

    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

    Sostituisci quanto segue:

    • BUCKET_NAME: il nome del bucket Cloud Storage creato nella sezione precedente.

    • PROJECT_ID: l'ID del progettoGoogle Cloud in cui esiste Cloud Storage e in cui vuoi creare il cluster Slurm.

    • REGION: la regione in cui esiste la prenotazione.

    • ZONE: la zona in cui esiste la prenotazione.

    • RESERVATION_URL: l'URL della prenotazione che vuoi utilizzare per creare il cluster Slurm. A seconda del progetto in cui esiste la prenotazione, specifica uno dei seguenti valori:

      • La prenotazione esiste nel tuo progetto: RESERVATION_NAME

      • La prenotazione esiste in un progetto diverso e il tuo progetto può utilizzarla: projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME

  5. Esegui il deployment del cluster:

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

    Il comando ./gcluster deploy è un processo in due fasi, che è il seguente:

    • La prima fase crea un'immagine personalizzata con tutto il software preinstallato, il cui completamento può richiedere fino a 35 minuti.

    • La seconda fase esegue il deployment del cluster utilizzando l'immagine personalizzata. Questo processo dovrebbe essere completato più rapidamente rispetto alla prima fase.

    Se la prima fase va a buon fine, ma la seconda no, puoi provare a eseguire nuovamente il deployment del cluster Slurm saltando la prima fase:

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

Prepara il workload

Per preparare il workload:

  1. Crea script di workload.

  2. Carica gli script nel cluster Slurm.

  3. Connettiti al cluster Slurm.

  4. Installa framework e strumenti.

Crea script del workload

Per creare gli script che verranno utilizzati dal tuo workload di perfezionamento:

  1. Per configurare l'ambiente virtuale Python, crea il file install_environment.sh con i seguenti contenuti:

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

    Questo script configura un ambiente virtuale Python affidabile, installa una build notturna di PyTorch e scarica il modello Llama 4.

  2. Per specificare le dipendenze Python per lo script di addestramento, crea un file requirements-fsdp.txt con i seguenti contenuti:

    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. Specifica llama4-train-distributed.py come script di addestramento principale:

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

    Questo script utilizza lo strumento di addestramento TRL Supervised Fine-Tuning (SFT) per gestire i cicli di addestramento FSDP, la configurazione dell'adattamento a basso ranking (LoRA) e la formattazione dei dati.

  4. Per specificare le attività da eseguire per i job sul cluster Slurm, crea il file submit.slurm con i seguenti contenuti:

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

Carica gli script nel cluster Slurm

Per caricare gli script creati nella sezione precedente nel cluster Slurm, segui questi passaggi:

  1. Per identificare il nodo di accesso, elenca tutte le VM A4 nel tuo progetto:

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

    Il nome del nodo di accesso è simile a a4-high-login-001.

  2. Carica gli script nella home directory del nodo di accesso:

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

    Sostituisci LOGIN_NODE_NAME con il nome del nodo di accesso.

Connettiti al cluster Slurm

Connettiti al cluster Slurm connettendoti al nodo di accesso tramite SSH:

gcloud compute ssh LOGIN_NODE_NAME \
    --project=PROJECT_ID \
    --tunnel-through-iap \
    --zone=ZONE

Installare framework e strumenti

Dopo aver eseguito la connessione al nodo di accesso, installa framework e strumenti nel seguente modo:

  1. Esporta il token Hugging Face:

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

  2. Esegui lo script di installazione:

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

    Questo comando configura un ambiente virtuale con tutte le dipendenze richieste e scarica i pesi del modello nel file ~/Llama-4-Scout-17B-16E-Instruct.

    Poiché il download del modello è molto grande (~200 GB), questo processo richiede circa 30 minuti, a seconda delle condizioni di rete.

Avvia il carico di lavoro di perfezionamento

Per iniziare l'addestramento del workload:

  1. Invia il job allo scheduler Slurm:

    sbatch submit.slurm

  2. Sul nodo di accesso del cluster Slurm, puoi monitorare l'avanzamento del job controllando i file di output creati nella directory home:

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

    Se il job viene avviato correttamente, il file .err mostra una barra dei progressi che si aggiorna man mano che il job procede.

    Il completamento di questo job dovrebbe richiedere poco più di un'ora sul cluster Slurm. Il lavoro ha due fasi principali:

    • Copia del modello di base di grandi dimensioni nell'SSD locale di ogni nodo di calcolo.
    • Il job di addestramento, che inizia una volta completata la copia del modello. L'esecuzione di questo job richiede circa 35 minuti.

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina il progetto che contiene le risorse oppure mantieni il progetto ed elimina le singole risorse.

Elimina il progetto

Elimina un progetto Google Cloud :

gcloud projects delete PROJECT_ID

Elimina il cluster Slurm

Per eliminare il cluster Slurm:

  1. Vai alla directory cluster-toolkit.

  2. Elimina il file Terraform e tutte le risorse create:

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

Elimina l'istanza Filestore

Per impostazione predefinita, l'istanza Filestore ha l'impostazione deletion_protection impostata su true nel blueprint cluster-toolkit. Questa impostazione impedisce la perdita accidentale di dati quando modifichi gli ambienti. Per eliminare l'istanza Filestore, devi disattivare manualmente la protezione dall'eliminazione.

Passaggi successivi