本教學課程說明如何在Llama-4-Scout-17上,透過多節點多 GPU 的 Slurm 叢集,微調 Google Cloud大型語言模型 (LLM)。這個叢集使用兩個 A4 虛擬機器 (VM) 執行個體,每個執行個體都有 8 個 NVIDIA B200 GPU。
本教學課程主要說明以下兩個程序:
- 使用Google Cloud Cluster Toolkit 部署高效能的正式版 Slurm 叢集。在這個部署作業中,您會建立已預先安裝必要軟體的自訂 VM 映像檔。您也會設定共用的 Filestore 執行個體,並設定高速 RDMA 網路。
- 叢集部署完成後,您可以使用本教學課程隨附的一組指令碼,執行分散式微調工作。這項工作會運用 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 帳戶。如果您是 Google Cloud新手,歡迎 建立帳戶,親自評估產品在實際工作環境中的成效。新客戶還能獲得價值 $300 美元的免費抵免額,可用於執行、測試及部署工作負載。
-
安裝 Google Cloud CLI。
-
若您採用的是外部識別資訊提供者 (IdP),請先使用聯合身分登入 gcloud CLI。
-
執行下列指令,初始化 gcloud CLI:
gcloud init -
選取或建立專案所需的角色
- 選取專案:選取專案時,不需要具備特定 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 專案名稱。
啟用必要的 API:
啟用 API 時所需的角色
如要啟用 API,您需要具備服務使用情形管理員 IAM 角色 (
roles/serviceusage.serviceUsageAdmin),其中包含serviceusage.services.enable權限。瞭解如何授予角色。gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com
-
安裝 Google Cloud CLI。
-
若您採用的是外部識別資訊提供者 (IdP),請先使用聯合身分登入 gcloud CLI。
-
執行下列指令,初始化 gcloud CLI:
gcloud init -
選取或建立專案所需的角色
- 選取專案:選取專案時,不需要具備特定 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 專案名稱。
啟用必要的 API:
啟用 API 時所需的角色
如要啟用 API,您需要具備服務使用情形管理員 IAM 角色 (
roles/serviceusage.serviceUsageAdmin),其中包含serviceusage.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:使用者帳戶的 ID。 例如:myemail@example.com。ROLE:授予使用者帳戶的 IAM 角色。
- 為 Google Cloud 專案啟用預設服務帳戶:
gcloud iam service-accounts enable PROJECT_NUMBER-compute@developer.gserviceaccount.com
--project=PROJECT_ID將 PROJECT_NUMBER 替換為專案編號。如要查看專案編號,請參閱「 取得現有專案」。
- 將編輯者角色 (
roles/editor) 授予預設服務帳戶:gcloud projects add-iam-policy-binding PROJECT_ID
--member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com"
--role=roles/editor - 為使用者帳戶建立本機驗證憑證:
gcloud auth application-default login
- 為專案啟用 OS 登入功能:
gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
- 登入或建立 Hugging Face 帳戶。
- 安裝使用 Cluster Toolkit 時所需的依附元件。
使用 Hugging Face 存取 Llama 4
如要使用 Hugging Face 存取 Llama 4,請按照下列步驟操作:
-
依序點選「你的個人資料」>「設定」>「存取權杖」>「+ 建立新權杖」
複製並儲存
read access權杖值。本教學課程後續將用到這個位址。
準備環境
如要準備環境,請按照下列步驟操作:
複製 Cluster Toolkit GitHub 存放區:
git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git建立 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 叢集,請按照下列步驟操作:
前往
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 bucket 名稱。PROJECT_ID:Cloud Storage 所在的Google Cloud 專案 ID,也是您要建立 Slurm 叢集的位置。REGION:預訂項目所在的區域。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 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 模型。
如要為訓練指令碼指定 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將
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) 設定和資料格式。
如要指定工作在 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 叢集,請按照下列步驟操作:
如要找出登入節點,請列出專案中的所有 A4 VM:
gcloud compute instances list --filter="machineType:a4-highgpu-8g"登入節點的名稱類似於
a4-high-login-001。將指令碼上傳至登入節點的主目錄:
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
安裝架構和工具
連線至登入節點後,請執行下列操作,安裝架構和工具:
匯出 Hugging Face 權杖:
# On the login node export HF_TOKEN="hf_..." # Replace with your token執行安裝指令碼:
# On the login node chmod +x install_environment.sh ./install_environment.sh這個指令會設定虛擬環境,並納入所有必要依附元件,然後將模型權重下載至
~/Llama-4-Scout-17B-16E-Instruct檔案。由於模型下載檔案非常大 (約 200 GB),視網路狀況而定,這個程序大約需要 30 分鐘。
開始微調工作負載
如要開始訓練工作負載,請按照下列步驟操作:
將工作提交至 Slurm 排程器:
sbatch submit.slurm在 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 叢集,請按照下列步驟操作:
前往
cluster-toolkit目錄。刪除 Terraform 檔案和所有已建立的資源:
./gcluster destroy a4-high --auto-approve
刪除 Filestore 執行個體
根據預設,Filestore 執行個體在 cluster-toolkit 藍圖中會將 deletion_protection 設定設為 true。這項設定可避免修改環境時意外遺失資料。如要刪除 Filestore 執行個體,請手動停用防刪除功能。