使用 FSDP 在 A4 Slurm 叢集上微調 Mixtral-8x7B

本教學課程說明如何在 Google Cloud的多節點、多 GPU Slurm 叢集上微調 mistralai/Mixtral-8x7B-v0.1 模型。叢集使用兩個 a4-highgpu-8g 虛擬機器 (VM) 執行個體，每個執行個體都有 8 個 NVIDIA B200 GPU。

本教學課程說明的兩項主要程序如下：

  1. 使用Google Cloud Cluster Toolkit 部署正式環境等級的高效能 Slurm 叢集。在這項部署作業中，您會建立自訂 VM 映像檔，並預先安裝必要軟體。您也會設定共用 Lustre 檔案系統和高速網路。
  2. 部署叢集後，您可以使用本教學課程隨附的一組指令碼，執行分散式微調工作。這項工作會運用 PyTorch Fully Sharded Data Parallel (FSDP)，您可以透過 Hugging Face Transformer Reinforcement Learning (TRL) 程式庫存取。

本教學課程適用於機器學習 (ML) 工程師、研究人員、平台管理員和操作人員，以及有興趣在多個節點和 GPU 之間分配 AI 工作負載的資料和 AI 專家。

目標

  • 使用 Hugging Face 存取 Mixtral
  • 安裝 Cluster Toolkit
  • 準備環境
  • 建立及部署正式環境等級的 A4 High-GPU Slurm 叢集。
  • 設定多節點環境，使用 FSDP 進行分散式訓練。
  • 使用 Hugging Face trl.SFTTrainer 類別微調 Mixtral 模型。
  • 將資料暫存至本機 SSD。
  • 監控工作。
  • 清除所用資源。

費用

在本文件中，您會使用下列 Google Cloud計費元件：

您可以使用 Pricing Calculator，根據預測用量估算費用。

初次使用 Google Cloud 的使用者可能符合免費試用期資格。

事前準備

  1. 安裝 Google Cloud CLI。

  2. 設定 gcloud CLI，使用您的聯合身分。

    詳情請參閱「使用聯合身分登入 gcloud CLI」。

  3. 執行下列指令，初始化 gcloud CLI：

    gcloud init

  4. 建立或選取 Google Cloud 專案。

    選取或建立專案所需的角色

    • 選取專案：選取專案時，不需具備特定 IAM 角色，只要您在專案中獲派角色，即可選取該專案。
    • 建立專案：如要建立專案，您必須具備專案建立者角色 (roles/resourcemanager.projectCreator)，其中包含 resourcemanager.projects.create 權限。瞭解如何授予角色。

    • 建立 Google Cloud 專案：

      gcloud projects create PROJECT_ID

      將 PROJECT_ID 替換為您要建立的 Google Cloud 專案名稱。

    • 選取您建立的 Google Cloud 專案：

      gcloud config set project PROJECT_ID

      將 PROJECT_ID 替換為 Google Cloud 專案名稱。

  5. 確認 Google Cloud 專案已啟用計費功能。

  6. 啟用必要的 API (如果尚未啟用)：

    啟用 API 所需的角色

    如要啟用 API，您必須具備 serviceusage.services.enable 權限。如果您建立了專案，可能已透過「擁有者」角色 (roles/owner) 取得這項權限。否則，您可以透過「Service Usage 管理員」角色 (roles/serviceusage.serviceUsageAdmin) 取得這項權限。瞭解如何授予角色。

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

  7. 將角色授予使用者帳戶。針對下列各個 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

    更改下列內容：

  8. 為 Google Cloud 專案啟用預設服務帳戶：
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@ \
    --project=PROJECT_ID

    將 PROJECT_NUMBER 替換為您的專案編號。如要查看專案編號，請參閱「取得現有專案」。

  9. 將編輯者角色 (roles/editor) 授予預設服務帳戶：
    gcloud projects add-iam-policy-binding PROJECT_ID \
  --member="serviceAccount:PROJECT_NUMBER-compute@" \
  --role=roles/editor
  10. 為使用者帳戶建立本機驗證憑證：
    gcloud auth application-default login
  11. 為專案啟用 OS 登入功能： 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  12. 登入或建立 Hugging Face 帳戶。
  13. 安裝使用 Cluster Toolkit 所需的依附元件。

使用 Hugging Face 存取 Mixtral

如要透過 Hugging Face 存取 Mixtral，請按照下列步驟操作：

  1. 登入 Hugging Face 並探索 Mixtral 模型。
  2. 建立 Hugging Face read 存取權杖。
  3. 複製並儲存權杖值。本教學課程稍後會用到。

安裝 Cluster Toolkit

Cluster Toolkit 是開放原始碼工具，可簡化在 Google Cloud部署高效能運算 (HPC)、人工智慧 (AI) 和機器學習 (ML) 工作負載的程序。如要進一步瞭解如何使用 gcluster CLI 和管理叢集，請參閱「Cluster Toolkit 總覽」。

  1. 準備 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

  2. 下載版本：

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

  3. 定義 gcluster 路徑：

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

準備環境

如要準備環境，請按照下列步驟操作：

  1. 設定預設環境變數： 

    export PROJECT_ID="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export LUSTRE_ZONE="YOUR_LUSTRE_ZONE"
export LUSTRE_VPC_NETWORK="YOUR_LUSTRE_VPC_NETWORK"
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"

    更改下列內容：

    • YOUR_PROJECT_ID：要建立 Slurm 叢集的 Google Cloud 專案名稱。
    • YOUR_ZONE：預留項目所在的可用區。
    • YOUR_REGION：預留項目所在的區域。
    • YOUR_RESERVATION_NAME：要用來建立 Slurm 叢集的預留項目網址或名稱。
    • YOUR_CLUSTER_NAME：部署作業的名稱。請使用僅含英文字母和數字的簡稱 (例如 a4high)。這個名稱也會指派給部署作業建立的 Slurm 叢集。
    • YOUR_GCS_BUCKET：儲存訓練檢查點結果的 bucket 名稱。指定現有值區或建立新值區。建立前，請先熟悉 bucket 命名規定。
    • YOUR_HF_TOKEN：您在先前步驟中建立的 Hugging Face 權杖。

  2. 建立 Cloud Storage bucket：

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

建立 A4 Slurm 叢集

如要建立 A4 Slurm 叢集，請按照下列步驟操作：

  1. 執行下列指令，以環境變數覆寫 examples/machine-learning/a4-highgpu-8g 目錄中的 a4high-slurm-deployment.yaml 檔案：

    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}
  lustre_zone: ${LUSTRE_ZONE}
  a4h_cluster_size: 2
  a4h_reservation_name: ${RESERVATION_NAME}
EOF

  2. 開啟 examples/machine-learning/a4-highgpu-8g 目錄中的 a4high-slurm-blueprint.yaml 檔案，並編輯該檔案，將共用 /home 目錄改為使用 Managed Lustre，如下所示：

    • 使用 source: modules/file-system/filestore 移除預設的 homefs 模組區塊。
    • 啟用 lustrefs (homefs 區塊，含 remote_mount: lustrefs) 和 private-service-access 模組。
    • 在 vars 區塊中，設定下列項目：
      1. 將 /var/tmp/slurm_vars.json 中 install_managed_lustre 的值變更為 true。
      2. 將 per_unit_storage_throughput 參數設為 500。
      3. 將 lustre_size_gib 參數設為 36000。
      4. 取消註解 lustre_instance_id: lustre-instance。
      5. 註解排除或移除未使用的 filestore_ip_range。

  3. 部署叢集：

    ./gcluster deploy \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  --auto-approve

    ./gcluster deploy 指令會啟動兩階段程序，如下所示：

    • 第一階段會建構預先安裝所有軟體的自訂映像檔，最多可能需要 35 分鐘才能完成。
    • 第二階段則使用該自訂映像檔部署叢集。這項程序應該會比第一階段更快完成。

準備工作負載

如要準備工作負載，請按照下列步驟操作：

  1. 建立工作負載指令碼。

  2. 將指令碼上傳至 Slurm 叢集。

  3. 連線至 Slurm 叢集。

  4. 安裝框架和工具。

建立工作負載指令碼

如要建立微調工作負載會使用的指令碼，請按照下列步驟操作：

  1. 如要設定 Python 虛擬環境，請建立 install_environment.sh 檔案，並加入下列內容：

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

  2. 如要指定訓練指令碼的 Python 依附元件，請建立 requirements-fsdp.txt 檔案，並加入下列內容：

    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. 將 train-mixtral.py 指定為主要訓練指令碼：

    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. 如要指定要在 Slurm 叢集上執行的工作，請建立 train-mixtral.sh 檔案，並加入下列內容：

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

將指令碼上傳至 Slurm 叢集

如要將上一節建立的指令碼上傳至 Slurm 叢集，請按照下列步驟操作：

  1. 擷取叢集的登入節點名稱，設定 LOGIN_NODE 變數：

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

    LOGIN_NODE 變數會儲存類似 ${DEPLOYMENT_NAME}-login-001 的值。

  2. 建立防火牆規則：

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

  3. 將指令碼上傳至登入節點的主目錄：

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

連線至 Slurm 叢集

透過 SSH 連線至登入節點，連線至 Slurm 叢集：

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

安裝框架和工具

連線至登入節點後，請在運算節點上執行安裝指令碼，安裝必要的架構和工具：

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

這個指令會設定虛擬環境、安裝所有依附元件，並將 Mixtral 模型權重下載至 ~/Mixtral-8x7B-v0.1。設定過程可能需 30 分鐘以上。

透過 srun 執行這個步驟，會將安裝指令碼從登入節點委派給分配的運算節點，讓指令碼在編譯期間存取 GPU 驅動程式。

開始微調工作負載

如要啟動訓練工作負載，請按照下列步驟操作：

  1. 將工作提交至 Slurm 排程器：

    # On the login node
sbatch train-mixtral.sh

監控工作負載

如要監控微調工作的進度，請使用下列方法：

  1. 在 Slurm 叢集的登入節點上，您可以查看 home 目錄中建立的輸出檔案，監控工作進度：

    # On the login node:
tail -f mixtral-*.out mixtral-*.err

    如果工作成功啟動，.err 檔案會顯示進度列，並隨著工作進度更新。

    這項工作有兩個主要階段：

    • 將大型基礎模型複製到各個運算節點的本機 SSD。
    • 模型複製完成後，訓練工作就會開始。

    整個工作大約需要 60 分鐘才能執行完畢。

  2. 您也可以監控 Slurm 叢集中的 GPU 使用情形，確認微調工作是否有效率地執行。方法是在瀏覽器中開啟下列連結：

    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

    或者，您也可以直接在終端機中輸入指令：

    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"

    監控工作負載時，您可以查看下列資訊：

    • GPU 使用率：如果微調工作正常，在整個訓練過程中，所有 16 個 GPU (叢集中每個 VM 有 8 個 GPU) 的使用率應會上升並穩定在特定水準。

    • 工作時間：工作應會在約一小時內完成。

清除所用資源

為避免因為本教學課程所用資源，導致系統向 Google Cloud 帳戶收取費用，請刪除含有相關資源的專案，或者保留專案但刪除個別資源。

刪除資源

執行下列指令，刪除您在本教學課程中建立的資源：

  1. 如要刪除 Slurm 叢集，請前往 cluster-toolkit 目錄，然後執行 gcluster destroy 指令：

    ./gcluster destroy "${DEPLOYMENT_NAME}" --auto-approve --robust

  2. 刪除 bucket：

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

  3. 如要刪除 Packer 映像檔，請開啟網路瀏覽器，前往下列頁面，搜尋特定映像檔，然後按一下「Delete」(刪除)。

    http://console.cloud.google.com/compute/images

刪除專案

刪除 Google Cloud 專案：

gcloud projects delete PROJECT_ID

後續步驟