本教學課程說明如何在多節點、多 GPU 的 Slurm 叢集 Google Cloud上,微調 mistralai/Mixtral-8x7B-v0.1 模型。這個叢集使用兩個a4-highgpu-8g虛擬機器 (VM) 執行個體,每個執行個體都有 8 個 NVIDIA B200 GPU。
本教學課程主要說明以下兩個程序:
- 使用Google Cloud Cluster Toolkit 部署高效能的正式版 Slurm 叢集。在這個部署作業中,您會建立已預先安裝必要軟體的自訂 VM 映像檔。您也會設定共用的 Lustre 檔案系統,並設定高速網路。
- 叢集部署完成後,您可以使用本教學課程隨附的一組指令碼,執行分散式微調工作。這項工作會運用 PyTorch Fully Sharded Data Parallel (FSDP),您可透過 Hugging Face Transformer Reinforcement Learning(TRL) 程式庫存取這項功能。
本教學課程適用於機器學習 (ML) 工程師、研究人員、平台管理員和營運人員,以及對跨多個節點和 GPU 分散 AI 工作負載感興趣的資料和 AI 專家。
目標
- 使用 Hugging Face 存取 Mixtral
- 準備環境
- 建立並部署正式級別的 A4 High-GPU Slurm 叢集。
- 設定多節點環境,透過 FSDP 進行分散式訓練。
- 使用 Hugging Face
trl.SFTTrainer類別微調 Mixtral 模型。 - 將資料暫存到本機 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 lustre.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 lustre.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 存取 Mixtral
如要使用 Hugging Face 存取 Mixtral,請按照下列步驟操作:
- 建立 Hugging Face
read access權杖。 - 複製並儲存
read存取權杖值。您會在稍後的教學課程中用到這項資訊。
準備環境
您會在本地電腦上執行下列步驟,準備部署叢集。
複製 Google Cloud Cluster Toolkit 存放區:
git clone https://github.com/GoogleCloudPlatform/cluster-toolkit.git建立 Cloud Storage bucket:
export BUCKET_NAME="your-unique-bucket-name" gcloud storage buckets create gs://${BUCKET_NAME}
建立 A4 Slurm 叢集
如要建立 A4 Slurm 叢集,請按照下列步驟操作:
前往複製的
cluster-toolkit目錄:cd cluster-toolkit如果是首次使用 Cluster Toolkit,請建構
gcluster二進位檔:make前往
examples/machine-learning/a4-highgpu-8g目錄。開啟
a4high-slurm-deployment.yaml檔案,然後按照下列方式編輯:terraform_backend_defaults: type: gcs configuration: bucket: BUCKET_NAME vars: deployment_name: DEPLOYMENT_NAME project_id: PROJECT_ID region: REGION zone: ZONE a4h_cluster_size: 2 a4h_reservation_name: RESERVATION_NAME更改下列內容:
BUCKET_NAME:您在上一個章節中建立的 Cloud Storage bucket 名稱。PROJECT_ID:Cloud Storage 所在的 Google Cloud 專案 ID,也是您要建立 Slurm 叢集的專案。REGION:預訂項目所在的區域。ZONE:預訂項目所在的可用區。A4h_reservation_name:使用 A4 預留名稱。
開啟
a4high-slurm-blueprint.yaml檔案,然後按照下列方式編輯:- 移除
filestore_homefs模組。 - 啟用
lustrefs和private-service-access模組。 - 在
vars區塊中,設定下列項目:Find slurm_vars,並將install_managed_lustre設為true。- 將
per_unit_storage_throughput參數設為500。 - 將
size_gib參數設為36000。
- 移除
部署叢集:
./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 分鐘才能完成。
- 第二階段會使用該自訂映像檔部署叢集。這個程序應該會比第一階段更快完成。
準備工作負載
如要準備工作負載,請按照下列步驟操作:
建立工作負載指令碼
如要建立微調工作負載使用的指令碼,請按照下列步驟操作:
如要設定 Python 虛擬環境,請建立
install_environment.sh檔案,並加入下列內容:#!/bin/bash # This script sets a reliable environment for FSDP training. # It is meant to be run on a compute node. 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 # --- [STEP 2.4] Build Flash Attention from Source --- echo "--- Building flash-attn from source... This will take a while. ---" # Use all available CPU cores to speed up the build MAX_JOBS=$(nproc) pip install flash-attn --no-build-isolation # --- 3. Download the Model --- echo "--- [STEP 2.5] Downloading Mixtral 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 mistralai/Mixtral-8x7B-v0.1 --local-dir ~/Mixtral-8x7B-v0.1 --token $HF_TOKEN echo "--- Environment setup complete. ---"如要為訓練指令碼指定 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 protobuf==6.31.1將
train-mixtral.py指定為主要訓練指令碼:import torch from torch.distributed.fsdp import MixedPrecision from datasets import load_dataset import shutil import os import torch.distributed as dist from peft import LoraConfig, PeftModel, get_peft_model from transformers import ( AutoModelForCausalLM, AutoTokenizer, TrainingArguments, HfArgumentParser, ) from torch.distributed import get_rank, get_world_size from transformers.models.mixtral.modeling_mixtral import MixtralDecoderLayer from trl import SFTTrainer from dataclasses import dataclass, field from typing import Optional @dataclass class ScriptArguments: model_id: str = field(default="mistralai/Mixtral-8x7B-v0.1", 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=False, metadata={"help": "When using FSDP activation checkpointing, this must be set to False for Mixtral"}) 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": True} training_args.optim = "adamw_torch_fused" bf16_policy = MixedPrecision( param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16, buffer_dtype=torch.bfloat16, ) training_args.fsdp = "full_shard" training_args.fsdp_config = { "fsdp_auto_wrap_policy": "TRANSFORMER_BASED_WRAP", "fsdp_transformer_layer_cls_to_wrap": [MixtralDecoderLayer], "fsdp_state_dict_type": "SHARDED_STATE_DICT", "fsdp_offload_params": False, "fsdp_forward_prefetch": True, "fsdp_mixed_precision_policy": bf16_policy } tokenizer = AutoTokenizer.from_pretrained(script_args.model_id, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" model = AutoModelForCausalLM.from_pretrained( script_args.model_id, torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation="flash_attention_2", ) 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"], ) model = get_peft_model(model, peft_config) data_splits = load_dataset(script_args.dataset_name) dataset = data_splits["train"] eval_dataset = data_splits["test"] if script_args.dataset_subset_size is not None: dataset = dataset.select(range(script_args.dataset_subset_size)) dataset = dataset.shuffle(seed=training_args.seed) trainer = SFTTrainer( model=model, args=training_args, train_dataset=dataset, eval_dataset=eval_dataset, formatting_func=formatting_prompts_func, processing_class=tokenizer, ) trainer.train() dist.barrier() if trainer.is_world_process_zero(): best_model_path = trainer.state.best_model_checkpoint final_model_dir = os.path.join(training_args.output_dir, "final_best_model") print(f"Copying best model to: {final_model_dir}") if os.path.exists(final_model_dir): shutil.rmtree(final_model_dir) shutil.copytree(best_model_path, final_model_dir) if script_args.run_inference_after_training: del model, trainer torch.cuda.empty_cache() run_post_training_inference(script_args, final_model_dir, tokenizer) def run_post_training_inference(script_args, best_model_path, tokenizer): print("\n" + "="*50) print("=== RUNNING POST-TRAINING INFERENCE TEST ===") print("="*50 + "\n") base_model = AutoModelForCausalLM.from_pretrained( script_args.model_id, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto" ) model = PeftModel.from_pretrained(base_model, best_model_path) model = model.merge_and_unload() 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." 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()如要指定工作在 Slurm 叢集上執行的工作,請建立
train-mixtral.sh檔案,並加入下列內容:#!/bin/bash #SBATCH --job-name=mixtral-fsdp #SBATCH --nodes=2 #SBATCH --ntasks-per-node=8 #SBATCH --gpus-per-node=8 #SBATCH --partition=a4high #SBATCH --output=mixtral-%j.out #SBATCH --error=mixtral-%j.err set -e set -x echo "--- Slurm Job Started ---" # --- Define Paths --- LOCAL_SSD_PATH="/mnt/localssd/job_${SLURM_JOB_ID}" VENV_PATH="${HOME}/.venv/venv-fsdp" MODEL_PATH="${HOME}/Mixtral-8x7B-v0.1" # --- 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 ${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}'/train-mixtral.py \ --model_id="'${LOCAL_SSD_PATH}'/model/" \ --output_dir="${HOME}/outputs/mixtral_job_${SLURM_JOB_ID}" \ --dataset_name="philschmid/gretel-synthetic-text-to-sql" \ --seed=900913 \ --bf16=True \ --num_train_epochs=3 \ --per_device_train_batch_size=32 \ --gradient_accumulation_steps=4 \ --learning_rate=4e-5 \ --logging_steps=3 \ --lora_r=32 \ --lora_alpha=32 \ --lora_dropout=0.05 \ --eval_strategy=steps \ --eval_steps=10 \ --save_strategy=steps \ --save_steps=10 \ --load_best_model_at_end=False \ --metric_for_best_model=eval_loss \ --run_inference_after_training \ --dataset_subset_size=67000 ' # --- STAGE 3: 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 叢集,請按照下列步驟操作:
如要找出登入節點,請列出專案中的所有 VM:
gcloud compute instances list登入節點的名稱類似於
a4-high-login-001。將指令碼上傳至登入節點的主目錄:
# Run this from your local machine where you created the files 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 \ ./install_environment.sh \ ./requirements-fsdp.txt \ ./train-mixtral.py \ ./train-mixtral.sh \ "${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 srun \ --job-name=env-setup \ --nodes=1 \ --ntasks=1 \ --gpus-per-node=1 \ --partition=a4high \ bash ./install_environment.sh這個指令會設定虛擬環境、安裝所有依附元件,並將 Mixtral 模型權重下載至
~/Mixtral-8x7B-v0.1。這項程序可能需要 30 分鐘以上才能完成。
開始微調工作負載
如要開始訓練工作負載,請按照下列步驟操作:
將工作提交至 Slurm 排程器:
# On the login node sbatch train-mixtral.sh在 Slurm 叢集的登入節點上,您可以檢查
home目錄中建立的輸出檔案,監控工作進度:# On the login node tail -f mixtral-*.out如果工作順利啟動,
.err檔案會顯示進度列,並隨著工作進度更新。這項工作分為兩個主要階段:
- 將大型基礎模型複製到每個運算節點的本機 SSD。
- 訓練工作,會在模型複製完成後開始。
整個工作大約需要 40 分鐘才能執行完畢。
清除所用資源
為避免因為本教學課程所用資源,導致系統向 Google Cloud 帳戶收取費用,請刪除含有相關資源的專案,或者保留專案但刪除個別資源。
刪除 Slurm 叢集
如要刪除 Slurm 叢集,請按照下列步驟操作:
前往
cluster-toolkit目錄。刪除 Terraform 檔案和所有已建立的資源:
./gcluster destroy DEPLOYMENT_NAME --auto-approve
刪除專案
刪除 Google Cloud 專案:
gcloud projects delete PROJECT_ID
後續步驟
- 重新部署 Slurm 叢集
- 在 Slurm 叢集中測試網路效能
- 監控 Slurm 叢集中的 VM
- 建立服務端點: 微調模型後,您可以使用 GKE 或 Vertex AI 將模型部署至服務端點, 以供推論。