在 A4 Slurm 集群上微调 Gemma 3

本教程介绍如何在 Gemma 3 大语言模型 (LLM) 上微调 多节点 Slurm 集群,该集群使用两个 A4 虚拟机 (VM) 实例。在本教程中,您将执行以下操作:

本教程适用于机器学习 (ML) 工程师、平台管理员和运维人员,以及对使用 Slurm 作业调度功能处理微调工作负载感兴趣的数据和 AI 专家。

目标

  1. 使用 Hugging Face 访问 Gemma 3。

  2. 准备环境。

  3. 创建 A4 Slurm 集群。

  4. 准备工作负载。

  5. 运行微调作业。

  6. 监控作业。

  7. 清理。

费用

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

您可使用 价格计算器 根据您的预计使用情况来估算费用。

新 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 角色,您可以选择已获授予角色的任何项目。
    • 创建项目:如需创建项目,您需要 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 项目名称。

  6. 验证是否已为您的 Google Cloud 项目启用结算功能。

  7. 启用必需的 API:

    启用 API 所需的角色

    如需启用 API,您需要拥有 serviceusage.services.enable 权限。如果您已创建项目,则可能已通过 Owner 角色 (roles/owner) 拥有此权限。否则,您可以通过 Service Usage Admin 角色 (roles/serviceusage.serviceUsageAdmin) 获得此权限。 了解如何授予角色

    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 角色,您可以选择已获授予角色的任何项目。
    • 创建项目:如需创建项目,您需要 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 项目名称。

  12. 验证是否已为您的 Google Cloud 项目启用结算功能。

  13. 启用必需的 API:

    启用 API 所需的角色

    如需启用 API,您需要拥有 serviceusage.services.enable 权限。如果您已创建项目,则可能已通过 Owner 角色 (roles/owner) 拥有此权限。否则,您可以通过 Service Usage Admin 角色 (roles/serviceusage.serviceUsageAdmin) 获得此权限。 了解如何授予角色

    gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
  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:您的用户账号的标识符。 例如,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. 向默认服务账号授予 Editor 角色 (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 Login:
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  19. 登录或创建 Hugging Face 账号

使用 Hugging Face 访问 Gemma 3

如需使用 Hugging Face 访问 Gemma 3,请按照以下步骤操作:

  1. 签署同意协议以使用 Gemma 3 12B

  2. 创建 Hugging Face read 访问令牌。 点击您的个人资料 > 设置 > 访问令牌 > +创建新令牌

  3. 复制并保存 read access 令牌值。您将在本教程的后面部分用到该值。

安装 Cluster Toolkit

如需详细了解如何使用 gcluster 和管理集群,请参阅 Cluster Toolkit 概览

  1. 准备 Cluster Toolkit 版本:

    # Find all available releases at: https://github.com/GoogleCloudPlatform/cluster-toolkit/releases
    # Set the desired version TAG (e.g., v1.97.0)
    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 REGION="YOUR_REGION"
    export RESERVATION_URL="YOUR_RESERVATION_NAME"
    export CLUSTER_NAME="YOUR_CLUSTER_NAME"
    export BUCKET_NAME="YOUR_GCS_BUCKET"
    export HF_TOKEN="YOUR_HF_TOKEN"
    
    gcloud config set project "${PROJECT_ID}"
    gcloud config set billing/quota_project "${PROJECT_ID}"

    替换以下内容:

    • YOUR_PROJECT_ID:您要在其中创建 Cloud Storage 存储桶的 Google Cloud 项目的 ID。

    • YOUR_ZONE:预留所在的可用区。

    • YOUR_REGION:预留所在的区域。

    • RESERVATION_NAME:您要用于创建 Slurm 集群的预留的网址或名称。

    • YOUR_CLUSTER_NAME:您要创建的 Slurm 集群的名称。

    • YOUR_GCS_BUCKET:Cloud Storage 存储桶的名称,需遵循 存储桶命名要求

    • YOUR_HF_TOKEN:您在上一个部分中创建的 Hugging Face 访问令牌。

  2. 创建 Cloud Storage 存储桶:

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

创建 A4 Slurm 集群

如需创建 A4 Slurm 集群,请按照以下步骤操作:

  1. 创建 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: ${CLUSTER_NAME}
      project_id: ${PROJECT_ID}
      region: ${REGION}
      zone: ${ZONE}
      a4h_cluster_size: 2
      a4h_reservation_name: ${RESERVATION_URL}
    EOF
  2. 准备清单:

    1. 创建 Terraform 清单:

      gcluster create \
        -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
        "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml"
    2. 修补清单:

      echo "[$(date)] Patching Filestore deletion protection in ${CLUSTER_NAME}..."
      sed -i '/deletion_protection = {/,/}/ { s/enabled = true/enabled = false/; /reason  = "Avoid data loss"/d; }' "${CLUSTER_NAME}/cluster-env/main.tf"
  3. 部署集群:

    gcluster deploy "${CLUSTER_NAME}" --auto-approve

    gcluster deploy 命令是一个包含两个阶段的过程,如下所示:

    • 第一阶段构建预先安装了所有软件的自定义映像,此过程最多可能需要 45 分钟才能完成。

    • 第二阶段使用该自定义映像部署集群。此过程通常比第一阶段花费的时间更少。

    如果第一阶段成功,但第二阶段失败,您可以尝试跳过第一阶段,再次部署 Slurm 集群:

    gcluster deploy "${CLUSTER_NAME}" --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"
    fsdp_config:
      fsdp_auto_wrap_policy: "TRANSFORMER_BASED_WRAP"
      fsdp_backward_prefetch: "BACKWARD_PRE"
      fsdp_cpu_ram_efficient_loading: true
      fsdp_forward_prefetch: false
      fsdp_offload_params: false
      fsdp_sharding_strategy: "FULL_SHARD"
      fsdp_state_dict_type: "FULL_STATE_DICT"
      fsdp_transformer_layer_cls_to_wrap: "Gemma3DecoderLayer"
      fsdp_use_orig_params: true
    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=gemma3-finetune
    #SBATCH --nodes=2
    #SBATCH --ntasks-per-node=1 # 1 task per node
    #SBATCH --gpus-per-node=8   # 8 GPUs per node
    #SBATCH --partition=a4high
    #SBATCH --output=slurm-%j.out
    #SBATCH --error=slurm-%j.err
    
    set -e
    echo "--- Slurm Job Started ---"
    
    # --- STAGE 1: Copy Environment to Local SSD on all nodes ---
    srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c '
      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}"
      rsync -a --info=progress2 ~/./.venv/ ${LOCAL_VENV}/
      mkdir -p ${LOCAL_CACHE}
      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}"
    LOCAL_CACHE="/mnt/localssd/hf_cache_job_${SLURM_JOB_ID}"
    LOCAL_OUTPUT_DIR="/mnt/localssd/outputs_${SLURM_JOB_ID}"
    mkdir -p ${LOCAL_OUTPUT_DIR}
    
    # This is the main training command.
    srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 bash -c '
      # Ensure nccl-gib is at version 1.1.2-1
      INSTALLED_VERSION=$(dpkg-query -W -f="\${Version}" nccl-gib 2>/dev/null || echo "not_installed")
      if [ "${INSTALLED_VERSION}" != "1.1.2-1" ]; then
        echo "Upgrading nccl-gib to 1.1.2-1..."
        sudo apt-get update && sudo apt-get install -y --allow-change-held-packages nccl-gib=1.1.2-1
      fi
    
      source '"${LOCAL_VENV}"'/bin/activate
    
      export HF_HOME='"${LOCAL_CACHE}"'
      export HF_DATASETS_CACHE='"${LOCAL_CACHE}"'
    
      if [ -f /usr/local/gib/scripts/set_nccl_env.sh ]; then
        echo "Sourcing set_nccl_env.sh"
        source /usr/local/gib/scripts/set_nccl_env.sh
      fi
      export LD_LIBRARY_PATH=/usr/local/gib/lib64:${LD_LIBRARY_PATH:-}
    
      export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
      export MASTER_PORT=29500
    
      # Use accelerate launch to properly run the distributed FSDP job
      accelerate launch \
        --config_file ~/accelerate_config.yaml \
        --machine_rank $SLURM_NODEID \
        --main_process_ip $MASTER_ADDR \
        --main_process_port $MASTER_PORT \
        --num_machines $SLURM_NNODES \
        --num_processes $((SLURM_NNODES * 8)) \
        ~/train.py \
        --model_id google/gemma-3-12b-pt \
        --output_dir '"${LOCAL_OUTPUT_DIR}"' \
        --per_device_train_batch_size 1 \
        --gradient_accumulation_steps 8 \
        --num_train_epochs 3 \
        --learning_rate 1e-5 \
        --save_strategy steps \
        --save_steps 100
    '
    
    # --- 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}/ ~/gemma-12b-text-to-sql-finetuned/
    "
    
    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. 如需为作业指定说明,请创建包含以下内容的 train.py 文件:

    import torch
    import argparse
    from datasets import load_dataset
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig
    from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
    from trl import SFTTrainer, SFTConfig
    from huggingface_hub import login
    
    
    def get_args():
        parser = argparse.ArgumentParser()
        parser.add_argument("--model_id", type=str, default="google/gemma-3-12b-pt", help="Hugging Face model ID")
        parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face token for private models")
        parser.add_argument("--dataset_name", type=str, default="philschmid/gretel-synthetic-text-to-sql", help="Hugging Face dataset name")
        parser.add_argument("--output_dir", type=str, default="gemma-12b-text-to-sql", help="Directory to save model checkpoints")
    
        # LoRA arguments
        parser.add_argument("--lora_r", type=int, default=16, help="LoRA attention dimension")
        parser.add_argument("--lora_alpha", type=int, default=16, help="LoRA alpha scaling factor")
        parser.add_argument("--lora_dropout", type=float, default=0.05, help="LoRA dropout probability")
    
        # SFTConfig arguments
        parser.add_argument("--max_seq_length", type=int, default=512, help="Maximum sequence length")
        parser.add_argument("--num_train_epochs", type=int, default=3, help="Number of training epochs")
        parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size per device during training")
        parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
        parser.add_argument("--learning_rate", type=float, default=1e-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=100, 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. Create and prepare the fine-tuning dataset ---
        # The SFTTrainer will use the `formatting_func` to apply the chat template.
        dataset = load_dataset(args.dataset_name, split="train")
        dataset = dataset.shuffle().select(range(12500))
        dataset = dataset.train_test_split(test_size=2500/12500)
    
        # --- 3. Configure Model and Tokenizer ---
        if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8:
            torch_dtype_obj = torch.bfloat16
            torch_dtype_str = "bfloat16"
        else:
            torch_dtype_obj = torch.float16
            torch_dtype_str = "float16"
    
        tokenizer = AutoTokenizer.from_pretrained(args.model_id)
        tokenizer.pad_token = tokenizer.eos_token
    
        gemma_chat_template = (
            "{% for message in messages %}"
            "{% if message['role'] == 'user' %}"
            "{{ '<start_of_turn>user\n' + message['content'] + '<end_of_turn>\n' }}"
            "{% elif message['role'] == 'assistant' %}"
            "{{ '<start_of_turn>model\n' + message['content'] + '<end_of_turn>\n' }}"
            "{% endif %}"
            "{% endfor %}"
            "{% if add_generation_prompt %}"
            "{{ '<start_of_turn>model\n' }}"
            "{% endif %}"
        )
        tokenizer.chat_template = gemma_chat_template
    
        # --- 4. Define the Formatting Function ---
        # This function will be used by the SFTTrainer to format each sample
        # from the dataset into the correct chat template format.
        def formatting_func(example):
            # The create_conversation logic is now implicitly handled by this.
            # We need to construct the messages list here.
            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 = "Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.\n\n<SCHEMA>\n{context}\n</SCHEMA>\n\n<USER_QUERY>\n{question}\n</USER_QUERY>\n"
    
            messages = [
                {"role": "system", "content": system_message},
                {"role": "user", "content": user_prompt.format(question=example["sql_prompt"], context=example["sql_context"])},
                {"role": "assistant", "content": example["sql"]}
            ]
            return tokenizer.apply_chat_template(messages, tokenize=False)
    
        # --- 5. Load Quantized Model and Apply PEFT ---
    
        # Define the quantization configuration
        config = AutoConfig.from_pretrained(args.model_id)
        config.use_cache = False
    
        # Load the base model in native precision
        print("Loading base model...")
        model = AutoModelForCausalLM.from_pretrained(
            args.model_id,
            config=config,
            attn_implementation="sdpa",
            torch_dtype=torch_dtype_obj,
        )
    
        # Configure LoRA.
        peft_config = LoraConfig(
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            r=args.lora_r,
            bias="none",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
            task_type="CAUSAL_LM",
        )
    
        # Apply the PEFT config to the model
        print("Applying PEFT configuration...")
        model = get_peft_model(model, peft_config)
        model.print_trainable_parameters()
    
        # --- 6. Configure Training Arguments ---
        training_args = SFTConfig(
            output_dir=args.output_dir,
            ddp_timeout=300.0,
            max_seq_length=args.max_seq_length,
            num_train_epochs=args.num_train_epochs,
            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,
            packing=False,
            gradient_checkpointing=True,
            gradient_checkpointing_kwargs={"use_reentrant": False},
            optim="adamw_torch",
            fp16=True if torch_dtype_obj == torch.float16 else False,
            bf16=True if torch_dtype_obj == torch.bfloat16 else False,
            max_grad_norm=0.3,
            warmup_ratio=0.03,
            lr_scheduler_type="constant",
            push_to_hub=False,
            report_to="tensorboard",
            dataset_kwargs={
                "add_special_tokens": False,
                "append_concat_token": True,
            }
        )
    
        # --- 7. Create Trainer and Start Training ---
        trainer = SFTTrainer(
            model=model,
            args=training_args,
            train_dataset=dataset["train"],
            eval_dataset=dataset["test"],
            formatting_func=formatting_func,
        )
    
        print("Starting training...")
        trainer.train()
        print("Training finished.")
    
        # --- 8. Save the final model ---
        print(f"Saving final model to {args.output_dir}")
        trainer.save_model()
    
    if __name__ == "__main__":
        main()

将脚本上传到 Slurm 集群

如需将您在上一个部分中创建的脚本上传到 Slurm 集群,请按照以下步骤操作:

  1. 通过检索集群的登录节点的名称来设置 LOGIN_NODE 变量:

    LOGIN_NODE="$(gcloud compute instances list \
                    --project="${PROJECT_ID}" \
                    --filter="labels.ghpc_deployment='${CLUSTER_NAME}' AND labels.slurm_instance_role='login'" \
                    --format="value(name)" | head -n 1)"

    LOGIN_NODE 变量存储的值类似于 ${CLUSTER_NAME}-login-001

  2. 将脚本上传到登录节点的主目录:

    gcloud compute scp \
      --project="${PROJECT_ID}" \
      --zone="${ZONE}" \
      --tunnel-through-iap \
      ./install_environment.sh \
      ./requirements.txt \
      ./submit.slurm \
      ./accelerate_config.yaml \
      ./train.py \
      "${LOGIN_NODE}":~/

连接到 Slurm 集群

连接到登录节点,并将 Hugging Face 令牌传播到新会话:

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

安装框架和工具

连接到登录节点后,设置包含所需依赖项的 Python 虚拟环境:

chmod +x install_environment.sh
./install_environment.sh

启动微调工作负载

如需启动微调工作负载,请按照以下步骤操作:

  1. 将作业提交给 Slurm 调度程序并收集任务 ID:

    JOB_ID="$(sbatch submit.slurm 2>&1 \
      | tee /dev/tty \
      | grep -oP 'Submitted batch job \K\d+')"
  2. 在 Slurm 集群中的登录节点上,您可以通过检查在 home 目录中创建的输出文件来监控作业的进度:

    tail -f "slurm-${JOB_ID}.out" "slurm-${JOB_ID}.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)的使用率都会上升并稳定在特定水平。

  • 作业时长:作业大约需要一小时才能完成。

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用,请删除包含这些资源的项目,或者保留项目但删除各个资源。

删除 Slurm 集群

如需删除 Slurm 集群,请执行以下操作:

gcluster destroy "${CLUSTER_NAME}" --auto-approve

如果您需要删除与项目关联的所有 VPC 网络、防火墙规则、路由器、IP 和子网,请执行以下操作:

echo "========================================================================="
echo " STARTING AUTOMATED NETWORK CLEANUP FOR CLUSTER: ${CLUSTER_NAME}"
echo "========================================================================="

echo "Discovering all VPC networks linked to the cluster..."
NETWORKS=$(gcloud compute networks list --project="${PROJECT_ID}" --format="value(name)" | grep "^${CLUSTER_NAME}" || true)

if [ -z "${NETWORKS}" ]; then
    echo "No VPC networks found starting with ${CLUSTER_NAME}. Everything is already clean!"
    exit 0
fi

echo "Found the following networks to process:"
echo "${NETWORKS}"
echo "-------------------------------------------------------------------------"

echo "=== 1. Wiping Global Firewall Rules ==="
FIREWALL_RULES=$(gcloud compute firewall-rules list \
    --project="${PROJECT_ID}" \
    --filter="network ~ ^${CLUSTER_NAME} OR name ~ ^${CLUSTER_NAME}" \
    --format="value(name)" || echo "")

if [ -n "${FIREWALL_RULES}" ]; then
    echo "Deleting matching firewall rules:"
    echo "${FIREWALL_RULES}"
    echo "${FIREWALL_RULES}" | xargs -r gcloud compute firewall-rules delete --project="${PROJECT_ID}" --quiet
else
    echo "No matching firewall rules found."
fi

echo "=== 2. Tearing Down Network-Specific Infrastructure ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Processing resources for network: ${net_name}"

    ROUTERS=$(gcloud compute routers list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${ROUTERS}" ]; then
        echo "  -> Deleting routers: ${ROUTERS}"
        echo "${ROUTERS}" | xargs -r gcloud compute routers delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    IPS=$(gcloud compute addresses list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="name ~ ^${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${IPS}" ]; then
        echo "  -> Deleting IP reservations: ${IPS}"
        echo "${IPS}" | xargs -r gcloud compute addresses delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi

    SUBNETS=$(gcloud compute networks subnets list \
        --project="${PROJECT_ID}" \
        --regions="${REGION}" \
        --filter="network=${net_name}" \
        --format="value(name)" || echo "")

    if [ -n "${SUBNETS}" ]; then
        echo "  -> Deleting subnetworks:"
        echo "${SUBNETS}"
        echo "${SUBNETS}" | xargs -r gcloud compute networks subnets delete --region="${REGION}" --project="${PROJECT_ID}" --quiet
    fi
done

echo "-------------------------------------------------------------------------"
echo "Waiting 15 seconds for Google Cloud API dependencies to unlock..."
sleep 15

echo "=== 3. Final VPC Networks Destruction ==="
echo "${NETWORKS}" | while read -r net_name; do
    [ -z "${net_name}" ] && continue
    echo "Deleting core VPC network: ${net_name}..."
    gcloud compute networks delete "${net_name}" --project="${PROJECT_ID}" --quiet || \
    echo "Warning: Could not delete ${net_name} yet. If a lock occurred, please rerun in 1 minute."
done

echo "========================================================================="
echo " SUCCESS: All network resources for cluster ${CLUSTER_NAME} have been wiped!"
echo "========================================================================="

删除项目

删除 Google Cloud 项目:

gcloud projects delete PROJECT_ID

后续步骤