Google menggunakan teknologi AI untuk menerjemahkan konten ke dalam bahasa pilihan Anda. Terjemahan AI mungkin mengandung kesalahan.

Menggunakan FSDP untuk menyempurnakan Mixtral-8x7B di cluster Slurm A4

GPU yang dikelompokkan

Tutorial ini menunjukkan cara menyesuaikan model mistralai/Mixtral-8x7B-v0.1 pada cluster Slurm multi-node dan multi-GPU di Google Cloud. Cluster ini menggunakan dua instance virtual machine (VM) a4-highgpu-8g, yang masing-masing memiliki 8 GPU NVIDIA B200.

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

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 sistem file Lustre bersama dan mengonfigurasi jaringan berkecepatan tinggi.
Setelah cluster di-deploy, Anda menjalankan tugas penyesuaian terdistribusi dengan menggunakan kumpulan skrip yang menyertai tutorial ini. Tugas ini memanfaatkan PyTorch Fully Sharded Data Parallel (FSDP), yang Anda akses melalui library Transformer Reinforcement Learning (TRL) Hugging Face.

Tutorial ini ditujukan untuk engineer, peneliti, administrator, dan operator platform machine learning (ML), serta spesialis data dan AI yang tertarik untuk mendistribusikan beban kerja AI di beberapa node dan GPU.

Tujuan

Mengakses Mixtral menggunakan Hugging Face
Instal Cluster Toolkit
Menyiapkan lingkungan Anda
Buat dan deploy cluster Slurm GPU Tinggi A4 tingkat produksi.
Konfigurasi lingkungan multi-node untuk pelatihan terdistribusi dengan FSDP.
Sesuaikan model Mixtral dengan menggunakan class 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

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.

Instal Google Cloud CLI.

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

Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

gcloud init

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 Project Creator (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.

Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

Aktifkan API yang diperlukan:

Peran yang diperlukan untuk mengaktifkan API

Untuk mengaktifkan API, Anda memerlukan izin serviceusage.services.enable. Jika Anda membuat project, kemungkinan Anda sudah memiliki izin ini melalui peran Pemilik (roles/owner). Jika tidak, Anda bisa mendapatkan izin ini melalui peran Admin Penggunaan Layanan (roles/serviceusage.serviceUsageAdmin). Pelajari cara memberikan peran.

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

Instal Google Cloud CLI.

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

Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

gcloud init

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 Project Creator (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.

Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

Aktifkan API yang diperlukan:

Peran yang diperlukan untuk mengaktifkan API

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

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

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

Buat kredensial autentikasi lokal untuk akun pengguna Anda:
```
gcloud auth application-default login
```
Catatan: Jika Anda menggunakan shell lokal dan penyedia identitas (IdP) eksternal, dan Anda mengalami error autentikasi setelah menjalankan perintah sebelumnya, maka login ke Google Cloud CLI dengan identitas gabungan Anda.

Aktifkan Login OS untuk project Anda:

gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE

Login atau buat akun Hugging Face.
Instal dependensi yang Anda perlukan untuk menggunakan Cluster Toolkit.

Mengakses Mixtral menggunakan Hugging Face

Untuk menggunakan Hugging Face guna mengakses Mixtral, lakukan hal berikut:

Login ke Hugging Face dan jelajahi model Mixtral.
Buat token akses read Hugging Face.
Salin dan simpan nilai token. Anda akan menggunakannya nanti dalam tutorial ini.

Instal Cluster Toolkit

Cluster Toolkit adalah alat open source yang menyederhanakan deployment workload komputasi berperforma tinggi (HPC), kecerdasan buatan (AI), dan machine learning (ML) di Google Cloud. Untuk mengetahui informasi selengkapnya tentang cara menggunakan gcluster dan mengelola cluster, lihat Ringkasan Cluster Toolkit.

Siapkan versi Cluster Toolkit:

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

Download rilis:

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

Tentukan jalur gcluster:

export CLUSTER_TOOLKIT_PATH="$(pwd)/cluster-toolkit"
export PATH="${CLUSTER_TOOLKIT_PATH}:${PATH}"
gcluster --version

Menyiapkan lingkungan Anda

Untuk menyiapkan lingkungan Anda, ikuti langkah-langkah berikut:

Tetapkan variabel lingkungan default:
```
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"
```
Ganti kode berikut:
- YOUR_PROJECT_ID: nama Google Cloud project tempat Anda ingin membuat cluster Slurm.
- YOUR_ZONE: zona tempat pemesanan Anda berada.
- YOUR_REGION: region tempat pemesanan Anda berada.
- YOUR_RESERVATION_NAME: URL atau nama reservasi yang ingin Anda gunakan untuk membuat cluster Slurm.
- YOUR_CLUSTER_NAME: nama deployment Anda. Gunakan nama pendek yang hanya berisi huruf dan angka (misalnya, a4high). Nama ini juga ditetapkan ke cluster Slurm yang dibuat oleh deployment.
- YOUR_GCS_BUCKET: nama bucket tempat Anda menyimpan hasil dari titik pemeriksaan pelatihan. Tentukan bucket yang ada atau buat bucket baru. Sebelum membuatnya, pahami persyaratan penamaan bucket.
- YOUR_HF_TOKEN: token Hugging Face yang Anda buat di langkah sebelumnya.

Membuat bucket Cloud Storage:

gcloud storage buckets create "gs://${BUCKET_NAME}" \
  --project="${PROJECT_ID}"

Membuat cluster Slurm A4

Untuk membuat cluster Slurm A4, lakukan hal berikut:

Timpa file a4high-slurm-deployment.yaml di direktori examples/machine-learning/a4-highgpu-8g dengan variabel lingkungan Anda dengan menjalankan perintah berikut:

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

Buka file a4high-slurm-blueprint.yaml di direktori examples/machine-learning/a4-highgpu-8g, lalu edit untuk menggunakan Managed Lustre untuk direktori /home bersama sebagai berikut:
- Hapus blok modul homefs default dengan source: modules/file-system/filestore.
- Aktifkan modul lustrefs (blok homefs dengan remote_mount: lustrefs) dan private-service-access.
- Di blok vars, konfigurasi hal berikut:
  1. Ubah nilai install_managed_lustre di /var/tmp/slurm_vars.json menjadi true.
  2. Tetapkan parameter per_unit_storage_throughput ke 500.
  3. Tetapkan parameter lustre_size_gib ke 36000.
  4. Hapus tanda komentar lustre_instance_id: lustre-instance.
  5. Menambahkan komentar atau menghapus filestore_ip_range yang tidak digunakan.
Deploy cluster:
```
./gcluster deploy \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  --auto-approve
```
Perintah ./gcluster deploy memulai 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.

Menyiapkan workload Anda

Untuk menyiapkan beban kerja, ikuti langkah-langkah berikut:

Buat skrip beban kerja.
Upload skrip ke cluster Slurm.
Hubungkan ke cluster Slurm .
Instal framework dan alat.

Membuat skrip beban kerja

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

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

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

Untuk menentukan dependensi Python untuk skrip pelatihan, buat file requirements-fsdp.txt dengan 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
protobuf==6.31.1

Tentukan train-mixtral.py sebagai skrip pelatihan utama:

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

Untuk menentukan tugas yang akan dijalankan pada cluster Slurm, buat file train-mixtral.sh dengan konten berikut:

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

Upload skrip ke cluster Slurm Anda

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

Tetapkan variabel LOGIN_NODE dengan mengambil nama node login untuk cluster Anda:

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

Variabel LOGIN_NODE menyimpan nilai yang mirip dengan ${DEPLOYMENT_NAME}-login-001.

Buat aturan firewall:

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

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 \
  ./train-mixtral.py \
  ./train-mixtral.sh \
  "${LOGIN_NODE}":~/

Menghubungkan ke cluster Slurm

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

gcloud compute ssh "${LOGIN_NODE}" \
    --project="${PROJECT_ID}" \
    --tunnel-through-iap \
    --zone="${ZONE}"
    -- -t "export HF_TOKEN='${HF_TOKEN}'; bash -l"

Menginstal framework dan alat

Setelah Anda terhubung ke node login, instal framework dan alat yang diperlukan dengan menjalankan skrip penginstalan di node komputasi:

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

Perintah ini menyiapkan lingkungan virtual, menginstal semua dependensi, dan mendownload bobot model Mixtral ke ~/Mixtral-8x7B-v0.1. Proses ini bisa memerlukan waktu lebih dari 30 menit hingga selesai.

Menjalankan langkah ini melalui srun akan mendelegasikan skrip penginstalan dari node login ke node komputasi yang dialokasikan, sehingga memungkinkan skrip mengakses driver GPU selama kompilasi.

Mulai workload penyesuaian

Untuk memulai workload pelatihan, lakukan hal berikut:

Kirimkan tugas ke penjadwal Slurm:

# On the login node
sbatch train-mixtral.sh

Memantau workload Anda

Untuk memantau progres tugas penyesuaian, gunakan metode berikut:

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 mixtral-*.out mixtral-*.err
```
Jika tugas Anda berhasil dimulai, file .err akan menampilkan status progres yang diperbarui seiring berjalannya tugas Anda.

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.
Seluruh tugas akan memerlukan waktu sekitar 60 menit untuk dijalankan.

Anda juga dapat memantau penggunaan GPU di cluster Slurm untuk memverifikasi bahwa tugas penyesuaian Anda berjalan secara efisien. Untuk melakukannya, buka link berikut di browser Anda:

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

Atau, Anda dapat memasukkan perintah langsung di terminal:

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"

Saat memantau workload, Anda dapat melihat hal berikut:

Penggunaan GPU: untuk tugas penyesuaian yang berjalan lancar, Anda dapat melihat penggunaan semua 16 GPU Anda (delapan GPU untuk setiap VM dalam cluster) meningkat dan stabil ke tingkat tertentu selama pelatihan.
Durasi tugas: tugas akan memerlukan waktu sekitar satu jam untuk diselesaikan.

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 resource

Jalankan perintah berikut untuk menghapus resource yang Anda buat dalam tutorial ini:

Untuk menghapus cluster Slurm, buka direktori cluster-toolkit dan jalankan perintah berikut:
```
./gcluster destroy "${DEPLOYMENT_NAME}" --auto-approve --robust
```

Menghapus bucket Anda:

gcloud storage rm --recursive "gs://${BUCKET_NAME}" --quiet || true

Untuk menghapus gambar Packer, buka browser web Anda, buka halaman berikut, telusuri gambar spesifik Anda, lalu klik Hapus.
```
http://console.cloud.google.com/compute/images
```

Menghapus project Anda

Perhatian: Menghapus project akan mengakibatkan:

Semua hal dalam project akan dihapus. Jika Anda menggunakan project yang ada untuk mengerjakan tugas di dokumen ini, saat Anda menghapusnya, pekerjaan lain dalam project tersebut juga akan dihapus.
Project ID kustom akan hilang. Saat membuat project ini, Anda mungkin juga membuat project ID kustom yang masih ingin digunakan pada masa mendatang. Agar tidak kehilangan URL yang menggunakan project ID tersebut, seperti URL appspot.com, hapus resource yang dipilih di dalam project, bukan menghapus seluruh project.

Jika Anda berencana mempelajari beberapa arsitektur, tutorial atau panduan memulai, dengan menggunakan kembali project dapat membantu Anda agar tidak melampaui batas kuota project.

Menghapus Google Cloud project:

gcloud projects delete PROJECT_ID

Langkah berikutnya

Men-deploy ulang cluster Slurm
Menguji performa jaringan di cluster Slurm
Memantau VM di cluster Slurm
Buat endpoint penyaluran: Setelah memiliki model yang di-fine-tune, Anda dapat men-deploy-nya ke endpoint penyaluran menggunakan Google Kubernetes Engine (GKE) atau Vertex AI agar dapat diakses untuk inferensi.
Pelajari layanan penyimpanan yang direkomendasikan untuk workload AI dan ML