Google은 AI 기술을 사용하여 콘텐츠를 사용자의 기본 언어로 번역합니다. AI 번역에는 오류가 있을 수 있습니다.

A4 Slurm 클러스터에서 FSDP를 사용하여 Mixtral-8x7B 미세 조정

클러스터링된 GPU

이 튜토리얼에서는 Google Cloud의 멀티 노드, 멀티 GPU Slurm 클러스터에서 mistralai/Mixtral-8x7B-v0.1 모델을 파인 튜닝하는 방법을 보여줍니다. 이 클러스터는 각각 8개의 NVIDIA B200 GPU가 있는 두 개의 a4-highgpu-8g 가상 머신 (VM) 인스턴스를 사용합니다.

이 튜토리얼에 설명된 두 가지 주요 프로세스는 다음과 같습니다.

Google Cloud Cluster Toolkit을 사용하여 프로덕션급 고성능 Slurm 클러스터를 배포합니다. 이 배포의 일환으로 필요한 소프트웨어가 사전 설치된 맞춤 VM 이미지를 만듭니다. 공유 Lustre 파일 시스템을 설정하고 고속 네트워킹을 구성합니다.
클러스터가 배포되면 이 튜토리얼에 포함된 스크립트 집합을 사용하여 분산 미세 조정 작업을 실행합니다. 이 작업은 Hugging Face Transformer Reinforcement Learning (TRL) 라이브러리를 통해 액세스하는 PyTorch Fully Sharded Data Parallel(FSDP)을 활용합니다.

이 튜토리얼은 여러 노드와 GPU에 AI 워크로드를 분산하는 데 관심이 있는 머신러닝 (ML) 엔지니어, 연구원, 플랫폼 관리자 및 운영자, 데이터 및 AI 전문가를 대상으로 합니다.

목표

Hugging Face를 사용하여 Mixtral에 액세스
Cluster Toolkit 설치
개발 환경 준비
프로덕션 등급 A4 High-GPU Slurm 클러스터를 만들고 배포합니다.
FSDP를 사용한 분산 학습을 위해 다중 노드 환경을 구성합니다.
Hugging Face trl.SFTTrainer 클래스를 사용하여 Mixtral 모델을 미세 조정합니다.
로컬 SSD에 데이터를 스테이징합니다.
작업 모니터링
삭제

비용

이 문서에서는 비용이 청구될 수 있는 Google Cloud구성요소를 사용합니다.

프로젝트 사용량을 기준으로 예상 비용을 산출하려면 가격 계산기를 사용하세요.

Google Cloud 신규 사용자는 무료 체험판을 사용할 수 있습니다.

시작하기 전에

Google Cloud 계정에 로그인합니다. Google Cloud를 처음 사용하는 경우 계정을 만들고 Google 제품의 실제 성능을 평가해 보세요. 신규 고객에게는 워크로드를 실행, 테스트, 배포하는 데 사용할 수 있는 $300의 무료 크레딧이 제공됩니다.

Google Cloud CLI를 설치합니다.

외부 ID 공급업체(IdP)를 사용하는 경우 먼저 제휴 ID로 gcloud CLI에 로그인해야 합니다.

gcloud CLI를 초기화하려면, 다음 명령어를 실행합니다.

gcloud init

Google Cloud 프로젝트를 만들거나 선택합니다.

프로젝트를 선택하거나 만드는 데 필요한 역할

프로젝트 선택: 프로젝트를 선택하는 데는 특정 IAM 역할이 필요하지 않습니다. 역할이 부여된 프로젝트를 선택하면 됩니다.
프로젝트 만들기: 프로젝트를 만들려면 resourcemanager.projects.create 권한이 포함된 프로젝트 생성자 역할(roles/resourcemanager.projectCreator)이 필요합니다. 역할 부여 방법 알아보기

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를 사용 설정하려면 serviceusage.services.enable 권한이 필요합니다. 프로젝트를 만든 경우 소유자 역할 (roles/owner)을 통해 이 권한이 이미 있을 수 있습니다. 그렇지 않으면 서비스 사용량 관리자 역할 (roles/serviceusage.serviceUsageAdmin)을 통해 이 권한을 얻을 수 있습니다. 역할을 부여하는 방법 알아보기

gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com lustre.googleapis.com

Google Cloud CLI를 설치합니다.

외부 ID 공급업체(IdP)를 사용하는 경우 먼저 제휴 ID로 gcloud CLI에 로그인해야 합니다.

gcloud CLI를 초기화하려면, 다음 명령어를 실행합니다.

gcloud init

Google Cloud 프로젝트를 만들거나 선택합니다.

프로젝트를 선택하거나 만드는 데 필요한 역할

프로젝트 선택: 프로젝트를 선택하는 데는 특정 IAM 역할이 필요하지 않습니다. 역할이 부여된 프로젝트를 선택하면 됩니다.
프로젝트 만들기: 프로젝트를 만들려면 resourcemanager.projects.create 권한이 포함된 프로젝트 생성자 역할(roles/resourcemanager.projectCreator)이 필요합니다. 역할 부여 방법 알아보기

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 사용 설정에 필요한 역할

gcloud services enable compute.googleapis.com file.googleapis.com logging.googleapis.com cloudresourcemanager.googleapis.com servicenetworking.googleapis.com lustre.googleapis.com

사용자 계정에 역할을 부여합니다. 다음 IAM 역할마다 다음 명령어를 1회 실행합니다. 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 역할입니다.
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
```
참고: 로컬 셸과 외부 ID 공급업체(IdP)를 사용하고 위의 명령어를 실행한 후 인증 오류가 발생하면 제휴 ID로 Google Cloud CLI에 로그인하세요.

프로젝트에 OS 로그인을 사용 설정합니다.

gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE

Hugging Face 계정에 로그인하거나 계정을 만듭니다.
클러스터 툴킷을 사용하는 데 필요한 종속 항목을 설치합니다.

Hugging Face를 사용하여 Mixtral에 액세스

Hugging Face를 사용하여 Mixtral에 액세스하려면 다음 단계를 따르세요.

Hugging Face에 로그인하고 Mixtral 모델을 살펴봅니다.
Hugging Face read 액세스 토큰을 만듭니다.
토큰 값을 복사하여 저장합니다. 이 값은 이 튜토리얼의 뒷부분에서 사용합니다.

Cluster Toolkit 설치

Cluster Toolkit은 Google Cloud에서 고성능 컴퓨팅 (HPC), 인공지능 (AI), 머신러닝(ML) 워크로드의 배포를 간소화하는 오픈소스 도구입니다. gcluster 사용 및 클러스터 관리에 대한 자세한 내용은 Cluster Toolkit 개요를 참고하세요.

Cluster Toolkit 버전을 준비합니다.

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

버전을 다운로드합니다.

# 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"

gcluster 경로를 정의합니다.

export CLUSTER_TOOLKIT_PATH="$(pwd)/cluster-toolkit"
export PATH="${CLUSTER_TOOLKIT_PATH}:${PATH}"
gcluster --version

개발 환경 준비

환경을 준비하려면 다음 단계를 수행합니다.

기본 환경 변수를 설정합니다.
```
export PROJECT_ID="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export DEPLOYMENT_NAME="YOUR_CLUSTER_NAME"
export BUCKET_NAME="YOUR_GCS_BUCKET"
export HF_TOKEN="YOUR_HF_TOKEN"
```
다음을 바꿉니다.
- YOUR_PROJECT_ID: Slurm 클러스터를 만들려는 Google Cloud 프로젝트의 이름입니다.
- YOUR_ZONE: 예약이 있는 영역입니다.
- YOUR_REGION: 예약이 있는 리전입니다.
- YOUR_RESERVATION_NAME: Slurm 클러스터를 만드는 데 사용할 예약의 URL 또는 이름입니다.
- YOUR_CLUSTER_NAME: 배포 이름입니다. 영문자와 숫자만 포함된 닉네임을 사용합니다 (예: a4high). 이 이름은 배포에서 생성된 Slurm 클러스터에도 할당됩니다.
- YOUR_GCS_BUCKET: 학습 체크포인트의 결과를 저장하는 버킷의 이름입니다. 기존 버킷을 지정하거나 새 버킷을 만듭니다. 버킷을 만들기 전에 버킷 이름 지정 요구사항을 숙지하세요.
- YOUR_HF_TOKEN: 이전 단계에서 만든 Hugging Face 토큰입니다.

Cloud Storage 버킷을 만듭니다.

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

A4 Slurm 클러스터 만들기

A4 Slurm 클러스터를 만들려면 다음 단계를 따르세요.

다음 명령어를 실행하여 examples/machine-learning/a4-highgpu-8g 디렉터리의 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: ${DEPLOYMENT_NAME}
  project_id: ${PROJECT_ID}
  region: ${REGION}
  zone: ${ZONE}
  a4h_cluster_size: 2
  a4h_reservation_name: ${RESERVATION_NAME}
EOF

examples/machine-learning/a4-highgpu-8g 디렉터리에서 a4high-slurm-blueprint.yaml 파일을 열고 공유 /home 디렉터리에 Managed Lustre를 사용하도록 다음과 같이 수정합니다.
- source: modules/file-system/filestore을 사용하여 기본 homefs 모듈 블록을 삭제합니다.
- lustrefs (remote_mount: lustrefs가 있는 homefs 블록) 및 private-service-access 모듈을 사용 설정합니다.
- vars 블록에서 다음을 구성합니다.
  1. /var/tmp/slurm_vars.json에서 install_managed_lustre 값을 true로 변경합니다.
  2. per_unit_storage_throughput 매개변수를 500로 설정합니다.
  3. lustre_size_gib 매개변수를 36000로 설정합니다.
  4. lustre_instance_id: lustre-instance를 주석 처리 삭제합니다.
  5. 사용하지 않는 filestore_ip_range를 주석 처리하거나 삭제합니다.
클러스터를 배포합니다.
```
./gcluster deploy \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  --auto-approve
```
./gcluster deploy 명령어는 다음과 같은 2단계 프로세스를 시작합니다.
- 첫 번째 단계에서는 모든 소프트웨어가 사전 설치된 맞춤 이미지를 빌드하며, 완료하는 데 최대 35분이 걸릴 수 있습니다.
- 두 번째 단계에서는 해당 맞춤 이미지를 사용하여 클러스터를 배포합니다. 이 프로세스는 첫 번째 단계보다 더 빨리 완료됩니다.

워크로드 준비

워크로드를 준비하려면 다음 단계를 따르세요.

워크로드 스크립트 만들기
Slurm 클러스터에 스크립트를 업로드합니다.
Slurm 클러스터에 연결합니다.
프레임워크 및 도구 설치

워크로드 스크립트 만들기

미세 조정 워크로드에서 사용할 스크립트를 만들려면 다음 단계를 따르세요.

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 --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. ---"
# Bypass PyTorch CUDA version mismatch check for extension builds
python3 -c "
import torch.utils.cpp_extension as ce
p = ce.__file__
content = open(p).read().replace('raise RuntimeError(CUDA_MISMATCH_MESSAGE', 'pass # raise RuntimeError')
open(p, 'w').write(content)
"

# 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/hf 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 클러스터에 업로드하려면 다음 단계를 따르세요.

클러스터의 로그인 노드 이름을 가져와 LOGIN_NODE 변수를 설정합니다.

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

LOGIN_NODE 변수는 ${DEPLOYMENT_NAME}-login-001과 유사한 값을 저장합니다.

방화벽 규칙 만들기

gcloud compute firewall-rules create allow-ssh-ingress-from-iap \
  --project="${PROJECT_ID}" \
  --network="${CLUSTER_NETWORK}" \
  --direction=INGRESS \
  --action=allow \
  --rules=tcp:22 \
  --source-ranges=35.235.240.0/20 \
  --description="Allow SSH ingress from Google Cloud Identity-Aware Proxy (IAP)"

로그인 노드의 홈 디렉터리에 스크립트를 업로드합니다.

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}":~/

Slurm 클러스터에 연결

SSH를 통해 로그인 노드에 연결하여 Slurm 클러스터에 연결합니다.

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

프레임워크 및 도구 설치

로그인 노드에 연결한 후 컴퓨팅 노드에서 설치 스크립트를 실행하여 필요한 프레임워크와 도구를 설치합니다.

# 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분 이상 걸릴 수 있습니다.

srun를 통해 이 단계를 실행하면 로그인 노드에서 할당된 컴퓨팅 노드로 설치 스크립트가 위임되어 스크립트가 컴파일 중에 GPU 드라이버에 액세스할 수 있습니다.

미세 조정 워크로드 시작

학습 워크로드를 시작하려면 다음 단계를 따르세요.

Slurm 스케줄러에 작업을 제출합니다.
```
# On the login node
sbatch train-mixtral.sh
```

워크로드 모니터링

미세 조정 작업의 진행 상황을 모니터링하려면 다음 방법을 사용하세요.

Slurm 클러스터의 로그인 노드에서 home 디렉터리에 생성된 출력 파일을 확인하여 작업 진행 상황을 모니터링할 수 있습니다.
```
# On the login node:
tail -f mixtral-*.out mixtral-*.err
```
작업이 성공적으로 시작되면 .err 파일에 작업이 진행됨에 따라 업데이트되는 진행률 표시줄이 표시됩니다.

작업에는 두 가지 주요 단계가 있습니다.
- 각 컴퓨팅 노드의 로컬 SSD에 대규모 기본 모델을 복사합니다.
- 모델 복사가 완료되면 시작되는 학습 작업입니다.
전체 작업을 실행하는 데 약 60분이 소요됩니다.

Slurm 클러스터에서 GPU 사용량을 모니터링하여 미세 조정 작업이 효율적으로 실행되고 있는지 확인할 수도 있습니다. 이렇게 하려면 브라우저에서 다음 링크를 여세요.

https://console.cloud.google.com/monitoring/metrics-explorer?project=YOUR_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

또는 터미널에 명령어를 직접 입력할 수 있습니다.

open "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 (클러스터의 각 VM에 8개의 GPU)의 사용량이 상승하고 특정 수준으로 안정화되는 것을 확인할 수 있습니다.
작업 기간: 작업이 완료되는 데 약 1시간이 걸립니다.

삭제

이 튜토리얼에서 사용된 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트를 유지하고 개별 리소스를 삭제하세요.

리소스 삭제

다음 명령어를 실행하여 이 튜토리얼에서 만든 리소스를 삭제합니다.

Slurm 클러스터를 삭제하려면 cluster-toolkit 디렉터리로 이동하여 다음 명령어를 실행합니다.
```
./gcluster destroy "${DEPLOYMENT_NAME}" --auto-approve --robust
```

버킷을 삭제합니다.

gcloud storage rm --recursive "gs://${BUCKET_NAME}" --quiet || true

Packer 이미지를 삭제하려면 웹브라우저를 열고 다음 페이지로 이동하여 특정 이미지를 검색하고 삭제를 클릭합니다.
```
http://console.cloud.google.com/compute/images
```

프로젝트 삭제

주의: 프로젝트를 삭제하면 다음과 같은 효과가 발생합니다.

프로젝트의 모든 항목이 삭제됩니다. 이 문서의 태스크에 기존 프로젝트를 사용한 경우 프로젝트를 삭제하면 프로젝트에서 수행한 다른 작업도 삭제됩니다.
커스텀 프로젝트 ID가 손실됩니다. 이 프로젝트를 만들 때 앞으로 사용할 커스텀 프로젝트 ID를 만들었을 수 있습니다. appspot.com URL과 같이 프로젝트 ID를 사용하는 URL을 보존하려면 전체 프로젝트를 삭제하는 대신 프로젝트 내에서 선택한 리소스만 삭제합니다.

여러 아키텍처, 튜토리얼 또는 빠른 시작을 살펴보려는 경우 프로젝트를 재사용하면 프로젝트 할당량 한도 초과를 방지할 수 있습니다.

Google Cloud 프로젝트를 삭제합니다.

gcloud projects delete PROJECT_ID

다음 단계

Slurm 클러스터 재배포
Slurm 클러스터에서 네트워크 성능 테스트
Slurm 클러스터에서 VM 모니터링
제공 엔드포인트 만들기: 미세 조정된 모델이 있으면 Google Kubernetes Engine (GKE) 또는 Vertex AI를 사용하여 제공 엔드포인트에 배포하여 추론에 액세스할 수 있도록 합니다.
AI 및 ML 워크로드에 권장되는 스토리지 서비스 알아보기