在 A4 Slurm 集群上训练 Qwen2

本教程介绍如何在 Google Cloud上的多节点、多 GPU Slurm 集群上训练大语言模型 (LLM)。本教程中使用的模型基于 Qwen2 15 亿参数模型。Slurm 集群使用两台 a4-highgpu-8g 虚拟机 (VM)，每台虚拟机都有 8 个 NVIDIA B200 GPU。

本教程中介绍的两个主要流程如下：

  1. 使用Google Cloud Cluster Toolkit 部署生产级高性能 Slurm 集群。在此部署过程中，您将创建一个预先安装了必要软件的自定义虚拟机映像。您还可以设置共享 Filestore 实例，并配置高速 RDMA 网络。
  2. 集群部署完毕后，您可以使用本教程随附的一组脚本来运行分布式预训练作业。该作业利用了 Hugging Face Accelerate 库。

本教程适用于机器学习 (ML) 工程师、研究人员、平台管理员和运维人员，以及对在 Google Cloud 上部署高性能 Slurm 集群来训练 LLM 感兴趣的数据和 AI 专家。

目标

  • 使用 Hugging Face 访问 Qwen2 模型。
  • 准备环境。
  • 创建并部署生产级 A4 Slurm 集群。
  • 使用 Accelerate 库训练 Qwen2 模型。
  • 监控作业。
  • 清理。

费用

在本文档中，您将使用 Google Cloud的以下收费组件：

如需根据您的预计使用情况来估算费用，请使用价格计算器。

新 Google Cloud 用户可能有资格申请免费试用。

准备工作

  1. 安装 Google Cloud CLI。

  2. 配置 gcloud CLI 以使用您的联合身份。

    如需了解详情，请参阅使用联合身份登录 gcloud CLI。

  3. 如需初始化 gcloud CLI，请运行以下命令： 

    gcloud init

  4. 创建或选择 Google Cloud 项目。

    选择或创建项目所需的角色

    • 选择项目：选择项目不需要特定的 IAM 角色，您可以选择已获授角色的任何项目。
    • 创建项目：如需创建项目，您需要拥有 Project Creator 角色 (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 权限。如果您创建了项目，则可能已经通过 Owner 角色 (roles/owner) 获得了此权限。否则，您可以通过 Service Usage Admin 角色 (roles/serviceusage.serviceUsageAdmin) 获得此权限。了解如何授予角色。 

    gcloud services enable gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.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

    替换以下内容：

    • PROJECT_ID：您的项目 ID。
    • USER_IDENTIFIER：您的用户 账号。如需查看示例，请参阅 在 IAM 政策中表示员工池用户。
    • ROLE：您向用户账号授予的 IAM 角色。
  8. 为您的 Google Cloud 项目启用默认服务账号： 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@ \
    --project=PROJECT_ID

    将 PROJECT_NUMBER 替换为您的项目编号。如需查看项目编号，请参阅 获取现有项目。

  9. 向默认服务账号授予 Editor 角色 (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 Login： 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  12. 登录或创建 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 存储分区，请运行以下命令：

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

    替换以下内容：

    • BUCKET_NAME：符合存储桶命名要求的 Cloud Storage 存储桶的名称。

    • PROJECT_ID：您要在其中创建 Cloud Storage 存储桶的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 存储桶的名称。

    • PROJECT_ID：您的 Cloud Storage 所在的Google Cloud 项目以及您要创建 Slurm 集群的项目 ID。

    • 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 虚拟机：

    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（集群中每个虚拟机有 8 个 GPU）的使用率在整个训练过程中上升并稳定在特定水平。

  • 作业时长：此作业大约需要 1 小时才能完成。

下载模型

作业成功运行后，训练好的模型会保存在登录节点上的 ~/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

后续步骤