A4 Slurm 클러스터에서 Qwen2 학습

이 튜토리얼에서는 Google Cloud의 멀티 노드, 멀티 GPU Slurm 클러스터에서 대규모 언어 모델 (LLM)을 학습시키는 방법을 보여줍니다. 이 튜토리얼에서 사용하는 모델은 Qwen2 15억 파라미터 모델을 기반으로 합니다. Slurm 클러스터는 각각 8개의 NVIDIA B200 GPU를 갖춘 2개의 a4-highgpu-8g 가상 머신 (VM)을 사용합니다.

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

  1. Google Cloud Cluster Toolkit을 사용하여 프로덕션급 고성능 Slurm 클러스터를 배포합니다. 이 배포의 일환으로 필요한 소프트웨어가 사전 설치된 맞춤 VM 이미지를 만듭니다. 공유 Filestore 인스턴스를 설정하고 고속 RDMA 네트워킹을 구성합니다.
  2. 클러스터가 배포된 후 이 튜토리얼에 포함된 스크립트 집합을 사용하여 분산 사전 학습 작업을 실행합니다. 이 작업은 Hugging Face Accelerate 라이브러리를 활용합니다.

이 튜토리얼은 LLM을 학습시키기 위해 Google Cloud 에 고성능 Slurm 클러스터를 배포하는 데 관심이 있는 머신러닝 (ML) 엔지니어, 연구원, 플랫폼 관리자 및 운영자, 데이터 및 AI 전문가를 대상으로 합니다.

목표

  • Hugging Face를 사용하여 Qwen2 모델에 액세스합니다.
  • 환경을 준비합니다.
  • 프로덕션 등급 A4 Slurm 클러스터를 만들고 배포합니다.
  • Accelerate 라이브러리를 사용하여 Qwen2 모델을 학습시킵니다 .
  • 작업 모니터링
  • 삭제

비용

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

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

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

시작하기 전에

  1. Google Cloud CLI를 설치합니다.

  2. gcloud CLI에서 제휴 ID를 사용하도록 구성합니다.

    자세한 내용은 제휴 ID로 gcloud CLI에 로그인을 참고하세요.

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

    gcloud init

  4. 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 프로젝트 이름으로 바꿉니다.

  5. Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

  6. 필요한 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

  7. 사용자 계정에 역할을 부여합니다. 다음 IAM 역할마다 다음 명령어를 1회 실행합니다. roles/compute.admin, roles/iam.serviceAccountUser, roles/file.editor, roles/storage.admin, roles/serviceusage.serviceUsageAdmin, roles/compute.osAdminLogin, roles/iap.tunnelResourceAccessor 

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

    다음을 바꿉니다.

  8. Google Cloud 프로젝트의 기본 서비스 계정을 사용 설정합니다. 
    gcloud iam service-accounts enable PROJECT_NUMBER-compute@ \
    --project=PROJECT_ID

    여기에서 PROJECT_NUMBER를 프로젝트 번호로 바꿉니다. 프로젝트 번호를 검토하려면 기존 프로젝트 가져오기를 참고하세요.

  9. 기본 서비스 계정에 편집자 역할 (roles/editor)을 부여합니다. 
    gcloud projects add-iam-policy-binding PROJECT_ID \
    --member="serviceAccount:PROJECT_NUMBER-compute@" \
    --role=roles/editor
  10. 사용자 계정의 로컬 인증 사용자 인증 정보를 만듭니다. 
    gcloud auth application-default login
  11. 프로젝트에 OS 로그인을 사용 설정합니다. 
    gcloud compute project-info add-metadata --metadata=enable-oslogin=TRUE
  12. Hugging Face 계정에 로그인하거나 계정을 만듭니다.

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

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

  1. Qwen 2 1.5B를 사용하려면 동의 계약에 서명해야 합니다.

  2. read 액세스 토큰을 만듭니다.

Cluster Toolkit 설치

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

  1. 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

  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 CLUSTER_NAME="YOUR_CLUSTER_NAME"
export ZONE="YOUR_ZONE"
export REGION="YOUR_REGION"
export RESERVATION_URL="RESERVATION_NAME"
export GCS_BUCKET="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: GKE 클러스터를 만들려는 Google Cloud 프로젝트의 이름입니다.
    • YOUR_CLUSTER_NAME: 만들려는 Slurm 클러스터의 이름입니다.
    • YOUR_ZONE: 예약이 있는 영역입니다.
    • YOUR_REGION,: 예약이 있는 리전입니다.
    • RESERVATION_NAME: Slurm 클러스터를 만드는 데 사용할 예약의 URL 또는 이름입니다.
    • YOUR_GCS_BUCKET: 학습 체크포인트의 결과를 저장하는 버킷의 이름입니다. 버킷을 만들기 전에 버킷 이름 지정 요구사항을 숙지하세요.
    • YOUR_HF_TOKEN: 이전 단계에서 만든 Hugging Face 토큰입니다.

  2. Cloud Storage 버킷을 만듭니다.

    gcloud storage buckets create "gs://${GCS_BUCKET}" \
  --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: ${GCS_BUCKET}

vars:
  deployment_name: ${CLUSTER_NAME}
  project_id: ${PROJECT_ID}
  region: ${REGION}
  zone: ${ZONE}
  a4h_cluster_size: 2
  a4h_reservation_name: ${RESERVATION_URL}
EOF

  2. Terraform 매니페스트를 만듭니다.

    gcluster create \
  -d "${MANIFEST_PATH}/a4high-slurm-deployment.yaml" \
  "${MANIFEST_PATH}/a4high-slurm-blueprint.yaml" \
  -o "${CLUSTER_NAME}"

  3. 매니페스트를 패치합니다.

    sed -i '/deletion_protection = {/,/}/ { s/enabled = true/enabled = false/; /reason  = "Avoid data loss"/d; }' "${CLUSTER_NAME}/${CLUSTER_NAME}/cluster-env/main.tf"

  4. 클러스터를 배포합니다.

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

    gcluster deploy 명령어는 다음과 같은 2단계 프로세스입니다.

    • 첫 번째 단계에서는 모든 소프트웨어가 사전 설치된 맞춤 이미지를 빌드하며, 완료하는 데 최대 50분이 걸릴 수 있습니다.

    • 두 번째 단계에서는 해당 맞춤 이미지를 사용하여 클러스터를 배포합니다. 이 프로세스는 일반적으로 첫 번째 단계보다 완료하는 데 시간이 덜 걸립니다.

    첫 번째 단계는 성공했지만 두 번째 단계가 실패한 경우 첫 번째 단계를 건너뛰고 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 파일을 만듭니다.

    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: "SHARDED_STATE_DICT"
  fsdp_transformer_layer_cls_to_wrap: "Qwen2DecoderLayer"
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 파일을 만듭니다.

    #SBATCH --job-name=qwen2-pretrain-smollm-fineweb
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1  # 1 main srun task on the node that manages accelerate
#SBATCH --gpus-per-node=8    # access to all 8 GPUs on the node
#SBATCH --partition=a4high
#SBATCH --output=logs/slurm-%j.out
#SBATCH --error=logs/slurm-%j.err

set -euo pipefail
echo "--- Slurm Job Started ---"

# --- STAGE 1: Setup environment and pre-process data on each node's local SSD ---
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none 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}"
  LOCAL_TMP="/mnt/localssd/tmp_job_${SLURM_JOB_ID}"

  rsync -a --info=progress2 ~/.venv/ ${LOCAL_VENV}/
  mkdir -p ${LOCAL_CACHE} ${PROCESSED_DATA_DIR} ${LOCAL_TMP}

  echo "Pre-processing data on $(hostname)..."
  source ${LOCAL_VENV}/bin/activate
  export TMPDIR="${LOCAL_TMP}"
  export TEMP="${LOCAL_TMP}"
  export TMP="${LOCAL_TMP}"
  export HF_HOME="${LOCAL_CACHE}"
  export HF_DATASETS_CACHE="${LOCAL_CACHE}"
  export HF_MODULES_CACHE="${LOCAL_CACHE}/modules"
  export HF_METRICS_CACHE="${LOCAL_CACHE}/metrics"

  python "${HOME}/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}"

export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29505

# Network and initialization debugging configurations
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=INIT,COLL
export TORCH_DISTRIBUTED_DEBUG=INFO
export NCCL_IB_DISABLE=0

# Launching with full GPU access enabled for accelerate
srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 --gpu-bind=none bash -c "
  mkdir -p ${LOCAL_OUTPUT_DIR}
  source ${LOCAL_VENV}/bin/activate

  # Reset the forced Slurm isolation to expose all 8 GPUs to the processes
  unset CUDA_VISIBLE_DEVICES

  # Retrieve the default network device.
  DETECTED_IFACE=\$(ip route show | grep default | awk '{print \$5}' | head -n 1)
  echo \"[INFO] Automatically detected cluster network interface: \$DETECTED_IFACE\"

  # Dynamic injection of the detected interface into the network configuration
  export NCCL_SOCKET_IFNAME=\"\${DETECTED_IFACE},gpu*\"
  export TP_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"
  export GLOO_SOCKET_IFNAME=\"\${DETECTED_IFACE}\"

  accelerate launch \
    --config_file ~/accelerate_config.yaml \
    --num_machines \$SLURM_NNODES \
    --num_processes \$((SLURM_NNODES * 8)) \
    --machine_rank \$SLURM_NODEID \
    --main_process_ip \$MASTER_ADDR \
    --main_process_port \$MASTER_PORT \
    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 \
      --logging_steps 1
"

# --- STAGE 3: Copy Final Model from Local SSD to Home Directory ---
echo "--- Copying final model from local SSD to /home ---"
mkdir -p ~/qwen2-from-scratch-on-smollm-fineweb/

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

  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
    참고: 이 튜토리얼에서는 최신 버전의 NVIDIA 및 Pytorch 종속 항목을 사용하지 않습니다. 최신 종속 항목이 필요한 경우 NVIDIA 문서 및 Pytorch 문서를 참고하세요.

  5. 학습에 적합한 형식으로 데이터 세트를 다운로드하고, 토큰화하고, 전처리하려면 다음 콘텐츠로 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()

  6. 작업의 안내를 지정하려면 다음 콘텐츠가 포함된 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 fineweb-edu dataset, e.g., 'CC-MAIN-2024-10'.")
    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-smollm-fineweb", 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)

    # --- 3. 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.")

    # --- 4. 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}...")

        # Synchronization of distributed processes
        local_rank = int(os.environ.get("LOCAL_RANK", -1))
        if local_rank != -1:
            # Introducing a minimal time offset per GPU to avoid I/O collisions.
            import time
            time.sleep(local_rank * 0.2)

        lm_dataset = load_from_disk(args.preprocessed_data_path, keep_in_memory=False)

    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: [item for sublist in examples[k] for item in sublist] 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 클러스터에 업로드하려면 다음 단계를 따르세요.

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

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

    gcloud compute scp \
  --project="${PROJECT_ID}" \
  --zone="${ZONE}" \
  --tunnel-through-iap \
  ./install_environment.sh \
  ./requirements.txt \
  ./submit.slurm \
  ./accelerate_config.yaml \
  ./train.py \
  ./preprocess_data.py \
  "${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"

프레임워크 및 도구 설치

로그인 노드에 연결한 후 다음을 실행하여 프레임워크와 도구를 설치합니다.

  1. 필요한 모든 종속 항목을 사용하여 Python 가상 환경을 설정합니다.

    chmod +x install_environment.sh
./install_environment.sh

워크로드 사전 학습 시작

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

  1. Slurm 스케줄러에 작업을 제출합니다.

    sbatch submit.slurm

  2. Slurm 클러스터의 로그인 노드에서 home 디렉터리에 생성된 출력 파일을 확인하여 작업의 진행 상황을 모니터링할 수 있습니다.

    tail -f ~/logs/slurm-1.err # (or .out, depending on where the script is currently sending logs)

    작업이 성공적으로 시작되면 .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

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

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시간이 걸립니다.

모델 다운로드

작업을 성공적으로 실행하면 학습된 모델이 로그인 노드의 ~/qwen2-from-scratch-on-smollm-fineweb/ 디렉터리에 저장됩니다. 이 영구 공유 디렉터리는 클러스터의 모든 노드에 마운트되므로 작업이 완료되거나 컴퓨팅 노드가 할당 해제된 후에도 모델 체크포인트를 계속 사용할 수 있습니다.

다음 예와 같이 gcloud compute scp 명령어를 사용하여 로그인 노드에서 로컬 머신으로 저장된 모델을 다운로드할 수 있습니다.

# From your local machine
gcloud compute scp --project="${PROJECT_ID}" --zone="${ZONE}" --tunnel-through-iap \
  "${LOGIN_NODE}":~/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 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트를 유지하고 개별 리소스를 삭제하세요.

리소스 삭제

  1. Slurm 클러스터를 삭제하려면 다음 단계를 따르세요.

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

  2. Cloud Storage 버킷을 삭제하려면 다음 단계를 따르세요.

    gcloud storage buckets delete "gs://${GCS_BUCKET}" --quiet || true

  3. Packer 이미지를 삭제하려면 웹브라우저를 열고 다음 페이지로 이동하여 특정 이미지를 검색하고 삭제를 클릭합니다. 

    http://console.cloud.google.com/compute/images

  4. 프로젝트와 연결된 모든 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

다음 단계