在 A4 Slurm 叢集上訓練 Qwen2

本教學課程說明如何在 Google Cloud的多節點、多 GPU Slurm 叢集上訓練大型語言模型 (LLM)。本教學課程使用的模型是以 Qwen2 15 億參數模型為基礎。Slurm 叢集使用兩個 a4-highgpu-8g 虛擬機器 (VM)，每個 VM 都有 8 個 NVIDIA B200 GPU。

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

  1. 使用Google Cloud Cluster Toolkit 部署高效能的正式版 Slurm 叢集。在這個部署作業中，您會建立已預先安裝必要軟體的自訂 VM 映像檔。您也會設定共用的 Filestore 執行個體，並設定高速 RDMA 網路。
  2. 叢集部署完成後，您可以使用本教學課程隨附的一組指令碼，執行分散式前置訓練工作。這項工作會運用 Hugging Face Accelerate 程式庫。

本教學課程的適用對象為機器學習 (ML) 工程師、研究人員、平台管理員和操作員，以及有興趣在 Google Cloud 部署高效能 Slurm 叢集來訓練 LLM 的資料和 AI 專家。

目標

  • 使用 Hugging Face 存取 Qwen2 模型。
  • 準備環境。
  • 建立並部署正式級 A4 Slurm 叢集。
  • 使用 Accelerate 程式庫訓練 Qwen2 模型。
  • 監控工作。
  • 清除所用資源。

費用

在本文件中，您會使用下列 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，您需要具備服務使用情形管理員 IAM 角色 (roles/serviceusage.serviceUsageAdmin)，其中包含 serviceusage.services.enable 權限。瞭解如何授予角色。

    gcloud services enable gcloud services enable compute.googleapis.com file.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，您需要具備服務使用情形管理員 IAM 角色 (roles/serviceusage.serviceUsageAdmin)，其中包含 serviceusage.services.enable 權限。瞭解如何授予角色。

    gcloud services enable gcloud services enable compute.googleapis.com file.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 帳戶。

使用 Hugging Face 存取 Qwen2

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

  1. 簽署同意聲明協議，即可使用 Qwen 2 1.5B。

  2. 建立 read 存取權杖。

準備環境

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

  1. 複製 Cluster Toolkit GitHub 存放區：

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

  2. 建立 Cloud Storage bucket：

    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：Cloud Storage 所在的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 should be run ONCE on the login node to set up the
# shared Python virtual environment.

set -e
echo "--- Creating Python virtual environment in /home ---"
python3 -m venv ~/.venv
echo "--- Activating virtual environment ---"
source ~/.venv/bin/activate

echo "--- Installing build dependencies ---"
pip install --upgrade pip wheel packaging

echo "--- Installing PyTorch for CUDA 12.8 ---"
pip install torch --index-url https://download.pytorch.org/whl/cu128

echo "--- Installing application requirements ---"
pip install -r requirements.txt

echo "--- Environment setup complete. You can now submit jobs with sbatch. ---"

  2. 如要指定微調工作的設定，請建立 accelerate_config.yaml 檔案，並加入下列內容：

    # Default configuration for a 2-node, 8-GPU-per-node (16 total GPUs) FSDP training job.

compute_environment: "LOCAL_MACHINE"
distributed_type: "FSDP"
downcast_bf16: "no"
machine_rank: 0
main_training_function: "main"
mixed_precision: "bf16"
num_machines: 2
num_processes: 16
rdzv_backend: "static"
same_network: true
tpu_env: []
use_cpu: false

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

    #!/bin/bash
#SBATCH --job-name=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8 # 8 tasks per node
#SBATCH --gpus-per-task=1   # 1 GPU per task
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -e
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
# This command runs once per node.
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c '
  set -e
  echo "Setting up local environment on $(hostname)..."
  LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}"
  LOCAL_CACHE="/mnt/localssd/hf_cache_job_${SLURM_JOB_ID}"
  PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
  rsync -a --info=progress2 ~/./.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export HF_HOME=${LOCAL_CACHE}
  export HF_DATASETS_CACHE=${LOCAL_CACHE}

  # This runs the new preprocessing script. It ensures only ONE process per node
  # downloads and processes the data, avoiding rate limiting and redundant work.
  python preprocess_data.py \
    --dataset_name "HuggingFaceFW/fineweb-edu" \
    --dataset_config "CC-MAIN-2024-10" \
    --tokenizer_id "Qwen/Qwen2-1.5B" \
    --max_seq_length 1024 \
    --output_path ${PROCESSED_DATA_DIR}

  echo "Setup on $(hostname) complete."
'

# --- STAGE 2: Run the Training Job using the Local Environment ---
echo "--- Starting Training ---"

LOCAL_VENV="/mnt/localssd/venv_job_${SLURM_JOB_ID}"
PROCESSED_DATA_DIR="/mnt/localssd/processed_data_${SLURM_JOB_ID}"
LOCAL_OUTPUT_DIR="/mnt/localssd/outputs_${SLURM_JOB_ID}"
mkdir -p ${LOCAL_OUTPUT_DIR}

# This is the main training command. It launches one Python process per GPU.
srun --ntasks=$((SLURM_NNODES * 8)) --gpus-per-task=1 bash -c "
  source ${LOCAL_VENV}/bin/activate

  # The training script now loads the pre-processed data from the local SSD.
  python train.py \
    --model_config_id "Qwen/Qwen2-1.5B" \
    --preprocessed_data_path ${PROCESSED_DATA_DIR} \
    --output_dir ${LOCAL_OUTPUT_DIR} \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --max_steps 10000 \
    --learning_rate 5e-5 \
    --save_strategy steps \
    --save_steps 500
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
# This command runs only on the first node of the job allocation
# and copies the final model back to the persistent shared directory.
srun --nodes=1 --ntasks=1 --ntasks-per-node=1 bash -c "
  rsync -a --info=progress2 ${LOCAL_OUTPUT_DIR}/ ~/qwen2-from-scratch-on-smollm-fineweb/
"

echo "--- Slurm Job Finished ---"

  4. 如要指定微調工作的依附元件，請建立 requirements.txt 檔案，並加入下列內容：

    # Hugging Face Libraries (Pinned to recent, stable versions for reproducibility)
transformers==4.53.3
datasets==4.0.0
accelerate==1.9.0
evaluate==0.4.5
bitsandbytes==0.46.1
trl==0.19.1
peft==0.16.0

# Other dependencies
tensorboard==2.20.0
protobuf==6.31.1
sentencepiece==0.2.0

  5. 如要下載資料集、將資料集標記化並預先處理為可供訓練的格式，請建立 preprocess_data.py 檔案並加入下列內容：

    import argparse
from datasets import load_dataset
from transformers import AutoTokenizer
import os
from itertools import chain

def get_args():
   parser = argparse.ArgumentParser(description="Download and preprocess a dataset.")
   parser.add_argument("--dataset_name", type=str, required=True)
   parser.add_argument("--dataset_config", type=str, required=True)
   parser.add_argument("--tokenizer_id", type=str, required=True)
   parser.add_argument("--max_seq_length", type=int, required=True)
   parser.add_argument("--output_path", type=str, required=True, help="Path to save the processed dataset.")
   return parser.parse_args()

def main():
   args = get_args()

   if os.path.exists(args.output_path) and os.listdir(args.output_path):
       print(f"Processed dataset already exists at {args.output_path}. Skipping.")
       return

   # 1. Load tokenizer
   tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id)

   # 2. Load raw dataset
   print(f"Loading raw dataset {args.dataset_name}...")
   raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

   # 3. Tokenize
   def tokenize_function(examples):
       return tokenizer(examples["text"])

   num_proc = os.cpu_count()
   print(f"Tokenizing dataset using {num_proc} processes...")
   print("Tokenizing dataset...")
   tokenized_dataset = raw_dataset.map(
       tokenize_function,
       batched=True,
       remove_columns=raw_dataset.column_names,
       desc="Running tokenizer on dataset",
       num_proc=num_proc,
   )

   # 4. Group texts
   def group_texts(examples):
       concatenated_examples = {k: list(chain.from_iterable(examples[k])) for k in examples.keys()}
       total_length = len(concatenated_examples[list(examples.keys())[0]])
       if total_length >= args.max_seq_length:
           total_length = (total_length // args.max_seq_length) * args.max_seq_length
       result = {
           k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
           for k, t in concatenated_examples.items()
       }
       result["labels"] = result["input_ids"].copy()
       return result

   print("Grouping texts...")
   lm_dataset = tokenized_dataset.map(
       group_texts,
       batched=True,
       desc=f"Grouping texts in chunks of {args.max_seq_length}",
       num_proc=num_proc,
   )

   # 5. Save to disk
   print(f"Saving processed dataset to {args.output_path}...")
   lm_dataset.save_to_disk(args.output_path)
   print("Preprocessing complete.")

if __name__ == "__main__":
   main()

  6. 如要指定工作指令，請建立包含下列內容的 train.py 檔案：

    import torch
import argparse
from datasets import load_dataset, load_from_disk
import os
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from huggingface_hub import login

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_config_id", type=str, default="Qwen/Qwen2-1.5B", help="Hugging Face model config to use for architecture.")
    # Data arguments - used if preprocessed data is not available
    parser.add_argument("--dataset_name", type=str, default="HuggingFaceFW/fineweb-edu", help="Hugging Face dataset for pre-training.")
    parser.add_argument("--dataset_config", type=str, default="CC-MAIN-2024-10", help="Config for the smollm-corpus dataset, e.g., 'fineweb-edu-dedup'.")
    parser.add_argument("--preprocessed_data_path", type=str, default=None, help="Path to a preprocessed dataset on disk. If provided, skips download and processing.")
    # General arguments
    parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face token for private models/tokenizers")
    parser.add_argument("--output_dir", type=str, default="qwen2-from-scratch-on-olmo", help="Directory to save model checkpoints")

    # TrainingArguments
    parser.add_argument("--max_seq_length", type=int, default=1024, help="Maximum sequence length")
    parser.add_argument("--num_train_epochs", type=int, default=1, help="Number of training epochs")
    parser.add_argument("--max_steps", type=int, default=-1, help="If set to a positive number, it overrides num_train_epochs.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=4, help="Batch size per device during training")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=4, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate")
    parser.add_argument("--logging_steps", type=int, default=10, help="Log every X steps")
    parser.add_argument("--save_strategy", type=str, default="steps", help="Checkpoint save strategy")
    parser.add_argument("--save_steps", type=int, default=500, help="Save checkpoint every X steps")

    return parser.parse_args()

def main():
    args = get_args()

    # --- 1. Setup and Login ---
    if args.hf_token:
        login(args.hf_token)

    # --- 2. Load Tokenizer ---
    # We load the tokenizer from the specified config ID to ensure compatibility
    # with the model architecture (e.g., special tokens).
    tokenizer = AutoTokenizer.from_pretrained(args.model_config_id)

    # --- 4. Initialize Model from Scratch ---
    print(f"Initializing a new model from {args.model_config_id} configuration...")
    config = AutoConfig.from_pretrained(args.model_config_id)
    model = AutoModelForCausalLM.from_config(config)

    print(f"Model has {model.num_parameters():,} parameters.")

    # --- 3. Load or Create and prepare the training dataset ---
    if args.preprocessed_data_path and os.path.exists(args.preprocessed_data_path):
        print(f"Loading preprocessed dataset from {args.preprocessed_data_path}...")
        lm_dataset = load_from_disk(args.preprocessed_data_path)
    else:
        print("No preprocessed dataset found, starting from raw data...")
        raw_dataset = load_dataset(args.dataset_name, name=args.dataset_config, split="train")

        # Tokenization function
        def tokenize_function(examples):
            return tokenizer(examples["text"])

        tokenized_dataset = raw_dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=raw_dataset.column_names,
            desc="Running tokenizer on dataset",
        )

        # Main data processing function that will concatenate all texts from our dataset
        # and generate chunks of max_seq_length.
        def group_texts(examples):
            # Concatenate all texts.
            concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
            total_length = len(concatenated_examples[list(examples.keys())[0]])
            # We drop the small remainder.
            if total_length >= args.max_seq_length:
                total_length = (total_length // args.max_seq_length) * args.max_seq_length
            # Split by chunks of max_len.
            result = {
                k: [t[i : i + args.max_seq_length] for i in range(0, total_length, args.max_seq_length)]
                for k, t in concatenated_examples.items()
            }
            result["labels"] = result["input_ids"].copy()
            return result

        lm_dataset = tokenized_dataset.map(
            group_texts,
            batched=True,
            desc=f"Grouping texts in chunks of {args.max_seq_length}",
        )

    # --- 5. Configure Training Arguments ---
    # Check for bfloat16 support
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    training_args = TrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.num_train_epochs,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        logging_steps=args.logging_steps,
        save_strategy=args.save_strategy,
        save_steps=args.save_steps,
        save_total_limit=2, # Optional: Limit the number of checkpoints
        bf16=use_bf16,
        fp16=not use_bf16,
        optim="adamw_torch",
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        report_to="tensorboard",
        gradient_checkpointing=True,
        # Required for gradient checkpointing with some parallelization strategies
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )

    # --- 6. Create Trainer and Start Training ---
    # Data collator will take care of creating batches for causal language modeling
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=lm_dataset,
        # eval_dataset=... # Optional: if you have a validation set
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    print("Starting training from scratch...")
    trainer.train()
    print("Training finished.")

    # --- 7. Save the final model ---
    print(f"Saving final model to {args.output_dir}")
    trainer.save_model()

if __name__ == "__main__":
    main()

將指令碼上傳至 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 \
  ./train.py \
  ./requirements.txt \
  ./submit.slurm \
  ./install_environment.sh \
  ./accelerate_config.yaml \
  "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 存取權杖建立環境變數：

    export HUGGING_FACE_TOKEN="HUGGING_FACE_TOKEN"

  2. 設定 Python 虛擬環境和所有必要依附元件：

    chmod +x install_environment.sh
./install_environment.sh

開始預先訓練工作負載

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

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

    sbatch submit.slurm

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

    tail -f logs/slurm-qwen2-pretrain-smollm-fineweb.err

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

監控工作負載

您可以監控 Slurm 叢集中的 GPU 使用情形，確認微調工作是否有效率地執行。如要這麼做，請在瀏覽器中開啟下列連結：

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) 的使用率在訓練期間上升並穩定在特定程度。

  • 工作時間：這項工作大約需要一小時才能完成。

下載模型

成功執行工作後，訓練好的模型會儲存在登入節點的 ~/qwen2-from-scratch-on-smollm-fineweb/ 目錄中。由於這個持續性共用目錄會掛接至叢集中的所有節點，因此即使作業完成或運算節點取消分配，模型檢查點仍可供使用。

您可以使用 gcloud compute scp 指令，將儲存的模型從登入節點下載至本機電腦，如下列範例所示：

# From your local machine
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 \
  "${LOGIN_NODE_NAME}":~/qwen2-from-scratch-on-smollm-fineweb/ ./qwen2-trained-model/ --recurse

下載模型後，您可以執行下列操作：

  • 載入模型以進行推論：使用 Hugging Face Transformers 架構載入 qwen2-trained-model/ 目錄，並使用新訓練的 Qwen2 模型執行推論。
  • 額外微調：將儲存的檢查點做為起點，針對更具體的資料集進行額外微調。
  • 將模型推送至 Hugging Face Hub：將訓練好的模型推送至 Hugging Face Hub，即可與他人分享。

清除所用資源

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

刪除專案

刪除 Google Cloud 專案：

gcloud projects delete PROJECT_ID

刪除 Slurm 叢集

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

  1. 前往 cluster-toolkit 目錄。

  2. 刪除 Terraform 檔案和所有已建立的資源：

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

後續步驟