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

Tutorial ini menunjukkan cara melakukan penyesuaian model mistralai/Mixtral-8x7B-v0.1 di cluster Slurm multi-node dan multi-GPU 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:

  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 sistem file Lustre bersama dan mengonfigurasi jaringan berkecepatan tinggi.
  2. 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
  • 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

  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 lustre.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 lustre.googleapis.com
  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 Mixtral menggunakan Hugging Face

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

  1. Buat token Hugging Face read access.
  2. Salin dan simpan nilai token akses read. Anda akan menggunakannya nanti dalam tutorial ini.

Menyiapkan lingkungan Anda

Anda melakukan langkah-langkah berikut di komputer lokal untuk menyiapkan deployment cluster.

  1. Buat clone repositori Google Cloud Cluster Toolkit:

    git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git
    
  2. Membuat bucket Cloud Storage:

    export BUCKET_NAME="your-unique-bucket-name"
    gcloud storage buckets create gs://${BUCKET_NAME}
    

Membuat cluster Slurm A4

Untuk membuat cluster Slurm A4, lakukan hal berikut:

  1. Buka direktori cluster-toolkit yang di-clone:

    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.

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

    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
    

    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.
    • A4h_reservation_name: gunakan nama reservasi A4 Anda.
  4. Buka file a4high-slurm-blueprint.yaml, lalu edit sebagai berikut:

    • Hapus modul filestore_homefs.
    • Aktifkan modul lustrefs dan private-service-access.
    • Di blok vars, konfigurasi hal berikut:
      1. Find slurm_vars dan setel install_managed_lustre ke true.
      2. Tetapkan parameter per_unit_storage_throughput ke 500.
      3. Tetapkan parameter size_gib ke 36000.
  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 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.

Siapkan workload Anda

Untuk menyiapkan beban kerja Anda, ikuti langkah-langkah 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 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 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
    
    # --- [STEP 2.4] Build Flash Attention from Source ---
    echo "--- Building flash-attn from source... This will take a while. ---"
    # 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/huggingface-cli download mistralai/Mixtral-8x7B-v0.1 --local-dir ~/Mixtral-8x7B-v0.1 --token $HF_TOKEN
    
    echo "--- Environment setup complete. ---"
    
  2. 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
    
  3. 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()
    
  4. 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 ---"
    

Mengupload Skrip ke cluster Slurm

Untuk mengupload skrip yang Anda buat di bagian sebelumnya ke cluster Slurm, lakukan hal berikut:

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

    gcloud compute instances list
    

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

  2. Upload skrip Anda ke direktori utama node login:

    # Run this from your local machine where you created the files
    LOGIN_NODE_NAME="your-login-node-name" # e.g., a4high-login-001
    PROJECT_ID="your-gcp-project-id"
    ZONE="your-cluster-zone" # e.g., us-west4-a
    
    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_NAME}":~/
    

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 Anda terhubung ke node login, instal framework dan alat.

  1. Mengekspor token Hugging Face Anda:

    # On the login node
    export HF_TOKEN="hf_..." # Replace with your token
    
  2. Jalankan 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 memakan waktu lebih dari 30 menit hingga selesai.

Mulai workload penyesuaian Anda

Untuk memulai pelatihan workload, lakukan hal berikut:

  1. Kirimkan tugas ke penjadwal Slurm:

    # On the login node
    sbatch train-mixtral.sh
    
  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 mixtral-*.out
    

    Jika tugas Anda berhasil dimulai, file .err akan menampilkan status progres yang diperbarui seiring progres 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 memerlukan waktu sekitar 40 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 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 DEPLOYMENT_NAME --auto-approve
    

Menghapus project Anda

Menghapus Google Cloud project:

gcloud projects delete PROJECT_ID

Langkah berikutnya