在 A4 Slurm 叢集上使用 FSDP 微調 Llama 4

本教學課程會說明如何在 Google Cloud的多節點、多 GPU Slurm 叢集上，微調Llama-4-Scout-17大型語言模型 (LLM)。此叢集使用兩部 A4 虛擬機器 (VM) 執行個體，每部 VM 都配備 8 個 NVIDIA B200 GPU。

本教學課程主要說明以下兩個程序：

  1. 使用 Cluster Toolkit 部署生產環境等級的高效能 Slurm 叢集。在這次部署作業中，您將建立自訂 VM 映像檔，並預先安裝必要軟體。您也會設定高速 RDMA 網路，並建立 Cloud Storage bucket 來託管 Terraform 狀態。
  2. 叢集部署完成後，請使用本教學課程隨附的指令碼集，執行分散式微調工作。這項工作會運用 PyTorch Fully Sharded Data Parallel (FSDP)，您可以透過 Hugging Face Transformer Reinforcement Learning 程式庫存取。

本教學課程適用於機器學習 (ML) 工程師、平台管理員和營運人員，以及有興趣使用 Slurm 工作排程功能處理微調工作負載的資料和 AI 專家。

目標

  • 使用 Hugging Face 存取 Llama 4。

  • 準備環境。

  • 建立及部署正式級 A4 High-GPU Slurm 叢集。

  • 設定多節點環境，以便使用 FSDP 進行分散式訓練。

  • 使用 Hugging Face trl.SFTTrainer 微調 Llama 4 模型。

  • 將資料暫存到本機 SSD。

  • 監控工作。

  • 清除所用資源。

費用

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

如要根據預測用量估算費用，請使用 Pricing Calculator。

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

事前準備

  1. 登入 Google Cloud 帳戶。如果您是 Google Cloud新手，歡迎 建立帳戶，親自體驗產品的實際應用成效。新客戶還能獲得價值 $300 美元的免費抵免額，能用於執行、測試及部署工作負載。

  2. 安裝 Google Cloud CLI。

  3. 如果您使用外部識別資訊提供者 (IdP)，請先 使用聯合身分登入 gcloud CLI。

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

    gcloud init

  5. 建立或選取 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 專案名稱。

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

  7. 如果尚未啟用必要的 API，請按照下列步驟啟用：

    啟用 API 時所需的角色

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

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

  8. 安裝 Google Cloud CLI。

  9. 如果您使用外部識別資訊提供者 (IdP)，請先 使用聯合身分登入 gcloud CLI。

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

    gcloud init

  11. 建立或選取 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 專案名稱。

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

  13. 如果尚未啟用必要的 API，請按照下列步驟啟用：

    啟用 API 時所需的角色

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

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

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

    更改下列內容：

    • PROJECT_ID：專案 ID。
    • USER_IDENTIFIER：使用者帳戶的 ID。 例如：myemail@example.com。
    • ROLE：授予使用者帳戶的 IAM 角色。
  15. 為 Google Cloud 專案啟用預設服務帳戶： 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com 
    
    --project=PROJECT_ID

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

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

使用 Hugging Face 存取 Llama 4

如要使用 Hugging Face 存取 Llama 4，請按照下列步驟操作：

  1. 簽署同意聲明協議，即可使用 Llama 4。

  2. 建立 Hugging Face read 存取權杖。

    依序點選「你的個人資料」>「設定」>「存取權杖」>「+ 建立新權杖」

  3. 複製並儲存 read access 權杖值。本教學課程稍後會用到。

準備環境

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

  1. 複製 Cluster Toolkit GitHub 存放區：

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

  2. 建立 Cloud Storage bucket 來代管 Terraform 狀態：

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

    更改下列內容：

    • BUCKET_NAME：Cloud Storage bucket 的名稱，必須符合 bucket 命名規定。

    • PROJECT_ID：您要建立 Cloud Storage bucket 的Google Cloud 專案 ID。

建立 A4 Slurm 叢集

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

  1. 前往 cluster-toolkit 目錄：

    cd cluster-toolkit

  2. 如果是首次使用 Cluster Toolkit，請建構 gcluster 二進位檔：

    make

  3. 前往 examples/machine-learning/a4-highgpu-8g 目錄：

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

  4. 開啟 a4high-slurm-deployment.yaml 檔案，然後按照下列方式編輯：

    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

    更改下列內容：

    • BUCKET_NAME：您在上一個章節中建立的 Cloud Storage bucket 名稱。

    • PROJECT_ID：bucket 所在的Google Cloud 專案 ID，也是您要建立 Slurm 叢集的專案。

    • REGION：預訂項目所在的區域。

    • ZONE：預訂項目所在的可用區。

    • RESERVATION_URL：您要用來建立 Slurm 叢集的預訂網址。根據預留項目所在的專案，指定下列其中一個值：

      • 專案中已有預留項目： RESERVATION_NAME

      • 預留項目位於其他專案，但您的專案可以使用該預留項目：projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME

  5. 部署叢集：

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

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

    • 第一階段會建立預先安裝所有軟體的自訂映像檔，最多可能需要 35 分鐘才能完成。

    • 第二階段會使用該自訂映像檔部署叢集。這個程序應該會比第一階段更快完成。

    如果第一階段成功，但第二階段失敗，您可以嘗試略過第一階段，再次部署 Slurm 叢集：

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

準備工作負載

如要準備工作負載，請執行下列操作：

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

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

  3. 連線至 Slurm 叢集。

  4. 安裝架構和工具。

建立工作負載指令碼

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

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

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

    這段指令碼會設定可靠的 Python 虛擬環境、安裝 PyTorch 夜間建構版本，並下載 Llama 4 模型。

  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

  3. 將 llama4-train-distributed.py 指定為主要訓練指令碼：

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

    這項指令碼會使用 TRL Supervised Fine-Tuning (SFT) Trainer 管理 FSDP 訓練迴圈、低秩調整 (LoRA) 設定和資料格式。

  4. 如要指定工作在 Slurm 叢集上執行的工作，請建立 submit.slurm 檔案，並加入下列內容：

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

將指令碼上傳至 Slurm 叢集

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

  1. 如要找出登入節點，請列出專案中的所有 A4 VM：

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

    登入節點的名稱類似於 a4-high-login-001。

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

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

    將 LOGIN_NODE_NAME 替換為登入節點的名稱。

連線至 Slurm 叢集

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

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

安裝架構和工具

連線至登入節點後，請執行下列操作，安裝架構和工具：

  1. 匯出 Hugging Face 權杖：

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

  2. 執行安裝指令碼：

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

    這個指令會設定虛擬環境，並下載所有必要的依附元件，然後將模型權重下載至 ~/Llama-4-Scout-17B-16E-Instruct 檔案。

    由於模型下載檔案非常大 (約 200 GB)，因此這個程序大約需要 30 分鐘，實際時間取決於網路狀況。

啟動微調工作負載

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

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

    sbatch submit.slurm

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

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

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

    這項工作在 Slurm 叢集上應需要一小時多一點的時間才能完成。這項工作分為兩個主要階段：

    • 將大型基礎模型複製到每個運算節點的本機 SSD。
    • 訓練工作，會在模型複製完成後開始。這項工作大約需要 35 分鐘才能執行完畢。

清除所用資源

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

刪除專案

刪除 Google Cloud 專案：

gcloud projects delete PROJECT_ID

刪除 Slurm 叢集

如要刪除 Slurm 叢集，請按照下列步驟操作：

  1. 前往 cluster-toolkit 目錄。

  2. 使用 gcluster destroy 指令刪除 Terraform 檔案和所有已建立的資源：

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

後續步驟