Menggunakan FSDP untuk menyetel Llama 4 di cluster Slurm A4

Tutorial ini menunjukkan cara melakukan penyesuaian pada model bahasa besar (LLM) Llama-4-Scout-17 di cluster Slurm multi-GPU multi-node di Google Cloud. Cluster ini menggunakan dua instance mesin virtual (VM) A4 yang masing-masing memiliki 8 GPU NVIDIA B200.

Dua proses utama yang dijelaskan dalam tutorial ini adalah sebagai berikut:

  1. Deploy cluster Slurm berperforma tinggi dan siap produksi menggunakan Google Cloud Cluster Toolkit. Sebagai bagian dari deployment ini, Anda akan membuat image VM kustom dengan software yang diperlukan sudah diinstal sebelumnya. Anda juga menyiapkan instance Filestore bersama, dan mengonfigurasi jaringan RDMA berkecepatan tinggi.
  2. Setelah cluster di-deploy, Anda menjalankan tugas penyesuaian terdistribusi dengan menggunakan kumpulan skrip yang disertakan dalam tutorial ini. Tugas ini memanfaatkan PyTorch Fully Sharded Data Parallel (FSDP), yang Anda akses melalui Transformer Reinforcement Learning Hugging Face

Tutorial ini ditujukan untuk engineer machine learning (ML), administrator dan operator platform, serta spesialis data dan AI yang tertarik menggunakan kemampuan penjadwalan tugas Slurm untuk menangani workload penyesuaian.

Tujuan

  • Mengakses Llama 4 menggunakan Hugging Face

  • Menyiapkan lingkungan Anda

  • Buat dan deploy cluster Slurm GPU Tinggi A4 tingkat produksi.

  • Konfigurasi lingkungan multi-node untuk pelatihan terdistribusi dengan FSDP.

  • Sesuaikan model Llama 4 menggunakan trl.SFTTrainer Hugging Face.

  • Siapkan data ke SSD lokal.

  • Pantau tugas Anda.

  • Jalankan pembersihan.

Biaya

Dalam dokumen ini, Anda akan menggunakan komponen Google Cloudyang dapat ditagih berikut:

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga.

Pengguna Google Cloud baru mungkin memenuhi syarat untuk mendapatkan uji coba gratis.

Sebelum memulai

  1. Login ke akun Google Cloud Anda. Jika Anda baru menggunakan Google Cloud, buat akun untuk mengevaluasi performa produk kami dalam skenario dunia nyata. Pelanggan baru juga mendapatkan kredit gratis senilai $300 untuk menjalankan, menguji, dan men-deploy workload.

  2. Instal Google Cloud CLI.

  3. Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

  4. Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut: 

    gcloud init

  5. Buat atau pilih Google Cloud project.

    Peran yang diperlukan untuk memilih atau membuat project

    • Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
    • Membuat project: Untuk membuat project, Anda memerlukan peran Pembuat Project (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.

    • Buat Google Cloud project:

      gcloud projects create PROJECT_ID

      Ganti PROJECT_ID dengan nama untuk Google Cloud project yang Anda buat.

    • Pilih project Google Cloud yang Anda buat:

      gcloud config set project PROJECT_ID

      Ganti PROJECT_ID dengan nama project Google Cloud Anda.

  6. Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

  7. Aktifkan API yang diperlukan:

    Peran yang diperlukan untuk mengaktifkan API

    Untuk mengaktifkan API, Anda memerlukan peran IAM Service Usage Admin (roles/serviceusage.serviceUsageAdmin), yang berisi izin serviceusage.services.enable. Pelajari cara memberikan peran. 

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

  8. Instal Google Cloud CLI.

  9. Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

  10. Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut: 

    gcloud init

  11. Buat atau pilih Google Cloud project.

    Peran yang diperlukan untuk memilih atau membuat project

    • Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
    • Membuat project: Untuk membuat project, Anda memerlukan peran Pembuat Project (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.

    • Buat Google Cloud project:

      gcloud projects create PROJECT_ID

      Ganti PROJECT_ID dengan nama untuk Google Cloud project yang Anda buat.

    • Pilih project Google Cloud yang Anda buat:

      gcloud config set project PROJECT_ID

      Ganti PROJECT_ID dengan nama project Google Cloud Anda.

  12. Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

  13. Aktifkan API yang diperlukan:

    Peran yang diperlukan untuk mengaktifkan API

    Untuk mengaktifkan API, Anda memerlukan peran IAM Service Usage Admin (roles/serviceusage.serviceUsageAdmin), yang berisi izin serviceusage.services.enable. Pelajari cara memberikan peran. 

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

  14. Memberikan peran ke akun pengguna Anda. Jalankan perintah berikut satu kali untuk setiap peran IAM berikut: 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

    Ganti kode berikut:

    • PROJECT_ID: Project ID Anda.
    • USER_IDENTIFIER: ID untuk akun pengguna Anda. Misalnya, myemail@example.com.
    • ROLE: Peran IAM yang Anda berikan ke akun pengguna Anda.
  15. Aktifkan akun layanan default untuk project Google Cloud Anda: 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com 
    
    --project=PROJECT_ID

    Ganti PROJECT_NUMBER dengan nomor project Anda. Untuk meninjau nomor project Anda, lihat Mendapatkan project yang sudah ada.

  16. Berikan peran Editor (roles/editor) ke akun layanan default: 
    gcloud projects add-iam-policy-binding PROJECT_ID 
    
    --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" 
    
    --role=roles/editor
  17. Buat kredensial autentikasi lokal untuk akun pengguna Anda: 
    gcloud auth application-default login
  18. Aktifkan Login OS untuk project Anda: 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. Login atau buat akun Hugging Face.
  20. Instal dependensi yang Anda perlukan untuk menggunakan Cluster Toolkit.

Mengakses Llama 4 menggunakan Hugging Face

Untuk menggunakan Hugging Face guna mengakses Llama 4, lakukan hal berikut:

  1. Tandatangani perjanjian izin untuk menggunakan Llama 4.

  2. Buat token akses read Hugging Face.

    Klik Profil Anda > Setelan > Token akses > +Buat token baru

  3. Salin dan simpan nilai token read access. Anda akan menggunakannya nanti dalam tutorial ini.

Menyiapkan lingkungan Anda

Untuk menyiapkan lingkungan Anda, ikuti langkah-langkah berikut:

  1. Buat clone repositori GitHub Cluster Toolkit:

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

  2. Membuat bucket Cloud Storage: 

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

    Ganti kode berikut:

    • BUCKET_NAME: nama untuk bucket Cloud Storage Anda yang mengikuti persyaratan penamaan bucket.

    • PROJECT_ID: ID projectGoogle Cloud tempat Anda ingin membuat bucket Cloud Storage.

Membuat cluster Slurm A4

Untuk membuat cluster Slurm A4, ikuti langkah-langkah berikut:

  1. Buka direktori cluster-toolkit:

    cd cluster-toolkit

  2. Jika Anda baru pertama kali menggunakan Cluster Toolkit, bangun biner gcluster:

    make

  3. Buka direktori examples/machine-learning/a4-highgpu-8g:

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

  4. Buka file a4high-slurm-deployment.yaml, lalu edit sebagai berikut:

    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

    Ganti kode berikut:

    • BUCKET_NAME: nama bucket Cloud Storage yang Anda buat di bagian sebelumnya.

    • PROJECT_ID: ID Google Cloud project tempat Cloud Storage Anda berada dan tempat Anda ingin membuat cluster Slurm.

    • REGION: region tempat pemesanan Anda berada.

    • ZONE: zona tempat pemesanan Anda berada.

    • RESERVATION_URL: URL reservasi yang ingin Anda gunakan untuk membuat cluster Slurm. Berdasarkan project tempat pemesanan berada, tentukan salah satu nilai berikut:

      • Reservasi ada di project Anda: RESERVATION_NAME

      • Pemesanan ada di project lain, dan project Anda dapat menggunakan pemesanan: projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME

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

    Perintah ./gcluster deploy adalah proses dua fase, yang adalah sebagai berikut:

    • Fase pertama membuat image kustom dengan semua software yang sudah diinstal sebelumnya, yang dapat memerlukan waktu hingga 35 menit untuk diselesaikan.

    • Fase kedua men-deploy cluster menggunakan image kustom tersebut. Proses ini akan selesai lebih cepat daripada fase pertama.

    Jika fase pertama berhasil, tetapi fase kedua gagal, Anda dapat mencoba men-deploy cluster Slurm lagi dengan melewati fase pertama:

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

Siapkan workload Anda

Untuk menyiapkan workload, Anda perlu melakukan hal berikut:

  1. Buat skrip beban kerja.

  2. Upload skrip ke cluster Slurm.

  3. Hubungkan ke cluster Slurm.

  4. Instal framework dan alat.

Membuat skrip beban kerja

Untuk membuat skrip yang akan digunakan workload penyesuaian, ikuti langkah-langkah berikut:

  1. Untuk menyiapkan lingkungan virtual Python, buat file install_environment.sh dengan konten berikut:

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

    Skrip ini menyiapkan lingkungan virtual Python yang andal, menginstal build malam PyTorch, dan mendownload model Llama 4.

  2. Untuk menentukan dependensi Python untuk skrip pelatihan, buat file requirements-fsdp.txtdengan konten berikut:

    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. Tentukan llama4-train-distributed.py sebagai skrip pelatihan utama:

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

    Skrip ini menggunakan TRL Supervised Fine-Tuning (SFT) Trainer untuk mengelola loop pelatihan FSDP, konfigurasi Low-Rank Adaptation (LoRA), dan pemformatan data.

  4. Untuk menentukan tugas yang akan dijalankan pada cluster Slurm, buat file submit.slurm dengan konten berikut:

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

Mengupload skrip ke cluster Slurm

Untuk mengupload skrip yang Anda buat di bagian sebelumnya ke cluster Slurm, ikuti langkah-langkah berikut:

  1. Untuk mengidentifikasi node login Anda, cantumkan semua VM A4 di project Anda:

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

    Nama node login mirip dengan a4-high-login-001.

  2. Upload skrip Anda ke direktori utama node login:

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

    Ganti LOGIN_NODE_NAME dengan nama node login.

Menghubungkan ke cluster Slurm

Hubungkan ke cluster Slurm dengan menghubungkan ke node login melalui SSH:

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

Menginstal framework dan alat

Setelah terhubung ke node login, instal framework dan alat dengan melakukan hal berikut:

  1. Mengekspor token Hugging Face Anda:

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

  2. Jalankan skrip penginstalan:

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

    Perintah ini menyiapkan lingkungan virtual dengan semua dependensi yang diperlukan, dan mendownload bobot model ke dalam file ~/Llama-4-Scout-17B-16E-Instruct.

    Karena download model berukuran sangat besar (~200 GB), proses ini memerlukan waktu sekitar 30 menit, bergantung pada kondisi jaringan.

Mulai workload penyesuaian Anda

Untuk mulai melatih workload Anda, lakukan hal berikut:

  1. Kirimkan tugas ke penjadwal Slurm:

    sbatch submit.slurm

  2. Di node login di cluster Slurm, Anda dapat memantau progres tugas dengan memeriksa file output yang dibuat di direktori home:

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

    Jika tugas Anda berhasil dimulai, file .err akan menampilkan status progres yang diperbarui seiring progres tugas Anda.

    Tugas ini akan memerlukan waktu lebih dari satu jam untuk diselesaikan di Slurm Cluster. Pekerjaan ini memiliki dua fase utama:

    • Menyalin model dasar berukuran besar ke SSD lokal setiap node komputasi.
    • Tugas pelatihan, yang dimulai setelah penyalinan model selesai. Tugas ini memerlukan waktu sekitar 35 menit untuk dijalankan.

Pembersihan

Agar tidak perlu membayar biaya pada akun Google Cloud Anda untuk resource yang digunakan dalam tutorial ini, hapus project yang berisi resource tersebut, atau simpan project dan hapus setiap resource.

Menghapus project Anda

Menghapus Google Cloud project:

gcloud projects delete PROJECT_ID

Menghapus cluster Slurm

Untuk menghapus cluster Slurm, ikuti langkah-langkah berikut:

  1. Buka direktori cluster-toolkit.

  2. Hancurkan file Terraform dan semua resource yang dibuat:

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

Menghapus instance Filestore

Secara default, instance Filestore Anda memiliki setelan deletion_protection yang ditetapkan ke benar (true) dalam cetak biru cluster-toolkit. Setelan ini mencegah kehilangan data yang tidak disengaja saat Anda mengubah lingkungan. Untuk menghapus instance Filestore, Anda harus menonaktifkan perlindungan penghapusan secara manual.

Langkah berikutnya