本教程介绍如何在 Google Cloud 上的 多节点、多 GPU Slurm 集群上训练大语言模型 (LLM) Google Cloud。本教程中使用的模型基于 Qwen2 15 亿参数模型。Slurm 集群使用两台 a4-highgpu-8g 虚拟机 (VM),每台虚拟机都有 8 个 NVIDIA B200 GPU。
本教程中介绍的两个主要流程如下:
- 使用 Google Cloud Cluster Toolkit 部署生产级高性能 Slurm 集群。在此部署过程中,您将创建一个预先安装了必要软件的自定义虚拟机映像。您还将设置共享 Filestore 实例,并配置高速 RDMA 网络。
- 集群部署完毕后,您可以使用本教程附带的一组脚本运行分布式预训练作业。该作业利用了 the Hugging Face Accelerate 库。
本教程适用于机器学习 (ML) 工程师、研究人员、平台管理员和运维人员,以及对在 Google Cloud 上部署高性能 Slurm 集群以训练 LLM 感兴趣的数据和 AI 专家。
目标
- 使用 Hugging Face 访问 Qwen2 模型。
- 准备环境。
- 创建和部署生产级 A4 Slurm 集群。
- 使用 Accelerate 库训练 Qwen2 模型。
- 监控作业。
- 清理。
费用
在本文档中,您将使用的以下收费组件: Google Cloud
您可使用 价格计算器 根据您的预计使用情况来估算费用。
准备工作
-
安装 Google Cloud CLI。
-
配置 gcloud CLI 以使用您的联合身份。
如需了解详情,请参阅使用联合身份登录 gcloud CLI。
-
如需初始化 gcloud CLI,请运行以下命令:
gcloud init -
选择或创建项目所需角色
- 选择项目:选择项目不需要特定的 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 项目名称。
-
验证是否已为您的 Google Cloud 项目启用结算功能。
启用必需的 API:
启用 API 所需的角色
如需启用 API,您需要拥有 Service Usage Admin 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
-
将角色授予您的用户账号。对以下每个 IAM 角色运行以下命令一次:
roles/compute.admin, roles/iam.serviceAccountUser, roles/file.editor, roles/storage.admin, roles/serviceusage.serviceUsageAdmingcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
替换以下内容:
PROJECT_ID:您的项目 ID。USER_IDENTIFIER:您的用户账号的标识符。 账号。如需查看示例,请参阅 在 IAM 政策中表示员工池用户。ROLE:您授予用户账号的 IAM 角色。
- 为您的 Google Cloud 项目启用默认服务帐号:
gcloud iam service-accounts enable PROJECT_NUMBER-compute@ \ --project=PROJECT_ID
将 PROJECT_NUMBER 替换为您的项目编号。如需查看您的 项目编号,请参阅 获取现有项目。
- 向默认服务账号授予 Editor 角色 (
roles/editor):gcloud projects add-iam-policy-binding PROJECT_ID \ --member="serviceAccount:PROJECT_NUMBER-compute@" \ --role=roles/editor
- 为您的用户账号创建本地身份验证凭据:
gcloud auth application-default login
- 为您的项目启用 OS Login:
gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
- 登录或创建 Hugging Face 账号。
使用 Hugging Face 访问 Qwen2
如需使用 Hugging Face 访问 Qwen2,请执行以下操作:
准备环境
如需准备环境,请按照以下步骤操作:
克隆 Cluster Toolkit GitHub 代码库:
git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git创建 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 集群,请按照以下步骤操作:
转到
cluster-toolkit目录:cd cluster-toolkit如果您是首次使用 Cluster Toolkit,请构建
gcluster二进制文件:make转到
examples/machine-learning/a4-highgpu-8g目录:cd examples/machine-learning/a4-highgpu-8g/打开
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 所在的项目的 ID,也是您要在其中创建 Slurm 集群的项目。Google CloudREGION:预留所在的区域。ZONE:预留所在的可用区。RESERVATION_URL:您要用于创建 Slurm 集群的预留的网址。根据预留所在的项目的不同,指定以下某个值:预留存在于您的项目中:
RESERVATION_NAME预留存在于其他项目中,并且您的项目可以使用该预留:
projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME
部署集群:
./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
准备工作负载
如需准备工作负载,请按照以下步骤操作:
创建工作负载脚本
如需创建训练工作负载将使用的脚本,请按照以下步骤操作:
如需设置 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. ---"如需为微调作业指定配置,请创建
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如需为要在 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 ---"如需为微调作业指定依赖项,请创建
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如需下载数据集、对其进行分词并将其预处理为可用于训练的格式,请创建
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()如需为作业指定说明,请创建
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 集群,请按照以下步骤操作:
如需确定登录节点,请列出项目中的所有 A4 虚拟机:
gcloud compute instances list --filter="machineType:a4-highgpu-8g"登录节点的名称类似于
a4-high-login-001。将脚本上传到登录节点的主目录:
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
安装框架和工具
连接到登录节点后,请执行以下操作来安装框架和工具:
为您的 Hugging Face 访问令牌创建环境变量:
export HUGGING_FACE_TOKEN="HUGGING_FACE_TOKEN"设置包含所有必需依赖项的 Python 虚拟环境:
chmod +x install_environment.sh ./install_environment.sh
开始预训练工作负载
如需开始训练工作负载,请执行以下操作:
将作业提交给 Slurm 调度程序:
sbatch submit.slurm在 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)的 使用率在整个训练过程中上升并 稳定在特定水平。
作业时长:作业大约需要一小时才能完成。
下载模型
成功运行作业后,训练的模型会保存在登录节点上的 ~/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 集群,请按照以下步骤操作:
转到
cluster-toolkit目录。销毁 Terraform 文件和所有已创建的资源:
./gcluster destroy a4-high --auto-approve