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

MaxText를 사용하여 TPU VM에서 강화 학습 학습 실행

이 튜토리얼에서는 대규모 언어 모델 (LLM)을 위한 고성능 JAX 기반 학습 스택인 MaxText를 사용하여 단일 v6e-8 Tensor Processing Unit (TPU) 가상 머신 (VM) 인스턴스에서 Google Cloud 강화 학습 (RL) 학습을 실행하는 단계별 안내를 제공합니다.

워크로드 요약

Cloud TPU에서 이 RL 학습 후 워크로드를 실행하려면 다음 구성요소가 필요합니다.

모델: 대화 및 명령어 팔로우 작업에 최적화된 오픈 가중치 LLM인 Llama 3.1 8B Instruct (llama3.1-8b-Instruct). 이 튜토리얼에서는 모델 가중치가 Hugging Face 형식에서 MaxText 호환 형식 (mt-format)으로 변환되고 JAX를 사용하여 v6e-8 TPU에서 분산 실행을 위해 준비됩니다.
강화 학습 (RL) 알고리즘: MaxText의 기본 제공 RL 트레이너를 사용하여 워크로드가 확장 가능한 학습 후 파이프라인을 실행하는 그룹 상대 정책 최적화 (GRPO). GRPO는 동일한 프롬프트에 대해 생성된 출력 그룹을 평가하여 별도의 값 또는 비평가 모델의 필요성을 없애고 상대적 보상을 사용하여 액터 정책을 효율적으로 최적화합니다.
데이터 세트 및 템플릿: 언어적으로 다양한 초등학교 수학 단어 문제로 구성된 고품질 데이터 세트인 Grade School Math 8K (GSM8K). 이 워크로드는 롤아웃 생성, 보상 계산, RL 정렬 단계에서 다단계 추론 기능 개선에 사용합니다.
인프라 및 프레임워크: JAX로 구동되는 고성능 학습 스택인 MaxText(maxtext[tpu-post-train])를 사용하는 단일 Cloud TPU v6e-8 노드 (ct6e-standard-8t).

목표

Cloud TPU VM 인스턴스를 설정합니다.
MaxText 및 종속 항목을 설치합니다.
Hugging Face 모델을 MaxText 형식으로 변환합니다.
TPU에서 RL 그룹 상대 정책 최적화 (GRPO) 워크로드를 실행합니다.
서빙을 위해 학습된 모델을 Hugging Face 형식으로 다시 변환합니다.

비용

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

Cloud TPU

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

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

이 문서에 설명된 태스크를 완료했으면 만든 리소스를 삭제하여 청구가 계속되는 것을 방지할 수 있습니다. 자세한 내용은 삭제를 참조하세요.

시작하기 전에

이 튜토리얼을 사용하려면 Hugging Face 액세스 토큰이 필요합니다. Hugging Face에서 무료 계정에 가입할 수 있습니다. 계정이 있으면 액세스 토큰을 생성합니다.
1. Hugging Face 시작 페이지에서 계정 아바타를 클릭하고 액세스 토큰을 선택합니다.
2. 액세스 토큰 페이지에서 새 토큰 만들기 를 클릭합니다.
3. 읽기 토큰 유형을 선택하고 토큰 이름을 입력합니다.
4. 액세스 토큰이 표시됩니다. 안전한 장소에 토큰을 저장합니다.

Hugging Face 웹사이트에서 학습할 모델의 라이선스 계약에 동의합니다. 이 튜토리얼에서는 llama3.1-8b-Instruct 모델을 사용합니다.

튜토리얼을 완료하는 데 필요한 권한을 얻으려면 관리자에게 프로젝트에 대한 다음 IAM 역할을 부여해 달라고 요청하세요.

TPU 관리자 (roles/tpu.admin)
서비스 계정 사용자 (roles/iam.serviceAccountUser)
Compute 편집자 (roles/compute.editor)

역할 부여에 대한 자세한 내용은 프로젝트, 폴더, 조직에 대한 액세스 관리를 참조하세요.

커스텀 역할이나 다른 사전 정의된 역할을 통해 필요한 권한을 얻을 수도 있습니다.

환경 설정

다음 스크립트를 실행하여 환경 변수를 설정합니다.

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export TPU_NAME="YOUR_TPU_NAME"

다음을 바꿉니다.

YOUR_PROJECT_ID: Google Cloud 프로젝트 ID
YOUR_ZONE: 사용할 영역
YOUR_RESERVATION_NAME: 용량 예약
YOUR_TPU_NAME: Cloud TPU VM 인스턴스의 이름

다음 명령어를 실행하여 인증합니다. Google Cloud

gcloud auth login

Cloud TPU VM 만들기

용량 예약에 바인딩된 8개의 v6e TPU 칩으로 Cloud TPU VM 인스턴스를 만듭니다.

gcloud compute instances create "${TPU_NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

VM 인스턴스가 생성된 후 SSH를 사용하여 연결합니다.

gcloud compute ssh "${TPU_NAME}" --zone "${ZONE}" --project "${PROJECT}"

TPU VM 인스턴스 내에서 다음 단계를 완료합니다.

MaxText 설치

TPU VM 인스턴스 내에서 시스템 패키지를 업데이트합니다.

sudo apt update && sudo apt upgrade -y --fix-missing

MaxText에 필요한 Python 3.12와 가상 환경 패키지를 설치합니다.

sudo apt install -y build-essential cmake ninja-build

uv를 사용하여 Python 패키지 설치 속도를 높입니다.

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

maxtext_venv라는 가상 환경을 만들고 활성화합니다.

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

MaxText 및 학습 후 작업에 필요한 종속 항목을 설치합니다.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.3" --resolution=lowest

다음 명령어를 실행하여 나머지 필수 종속 항목을 설치합니다.

install_tpu_post_train_extra_deps

모델을 MaxText 형식으로 변환

MaxText 형식으로 모델을 학습시키려면 Hugging Face 형식에서 MaxText 형식으로 변환해야 합니다.

다음 값을 제공합니다.

Hugging Face 액세스 토큰
사용할 모델의 이름
MaxText 형식으로 모델을 저장할 디렉터리
로드 및 스토리지 옵션

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='llama3.1-8b-Instruct'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

YOUR_HF_TOKEN을 이전에 만든 Hugging Face 액세스 토큰 으로 바꿉니다.

Hugging Face 형식에서 MaxText 형식으로 모델을 변환하려면 다음 스크립트를 실행합니다. 이 변환은 완료하는 데 약 5분이 걸립니다.

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

학습 워크로드 시작

변환 프로세스가 완료되면 RL 워크로드를 시작할 수 있습니다.

RL 워크로드 학습 매개변수를 구성합니다.

# -- MaxText configuration --
export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
export RUN_NAME
export CHIPS_PER_VM=8
export NUM_BATCHES=50
export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
export TPU_ACCELERATOR_TYPE=v6e-8
export TPU_WORKER_ID=0
TPU_NAME=$(hostname)
export TPU_NAME
export TPU_SKIP_MDS_QUERY=1
export TPU_WORKER_HOSTNAMES=localhost
export TPU_TOPOLOGY=2x4
export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
export TPU_HOST_BOUNDS=1,1,1

학습 작업을 시작합니다. v6e-8 VM 인스턴스에서 약 10분이 걸립니다.

python3 -m maxtext.trainers.post_train.rl.train_rl \
    model_name="${MODEL_NAME?}" \
    load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
    run_name="${RUN_NAME?}" \
    base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
    chips_per_vm="${CHIPS_PER_VM?}" \
    num_batches="${NUM_BATCHES?}" \
    num_test_batches=10 \
    rollout_data_parallelism=1 \
    rollout_tensor_parallelism=-1

학습된 모델을 Hugging Face 형식으로 다시 변환

학습 워크로드가 완료되면 모델을 Hugging Face 형식으로 다시 변환합니다.

내보내기 및 학습된 매개변수의 경로를 설정합니다.

export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
export HF_MODEL_NAME=llama3.1-8b
export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/actor/$NUM_BATCHES/model_params

Hugging Face 형식으로 다시 변환을 실행합니다.

python3 -m maxtext.checkpoint_conversion.to_huggingface \
    model_name="${HF_MODEL_NAME?}" \
    load_parameters_path="${POST_TRAIN_PATH?}" \
    base_output_directory="${HF_EXPORT?}" \
    scan_layers=True \
    use_multimodal=False \
    weight_dtype=bfloat16

변환이 완료되면 /dev/shm/$MODEL_NAME/hf-trained에 저장된 조정된 모델을 사용할 수 있습니다. VM이 재부팅되면 /dev/shm 폴더의 콘텐츠에 액세스할 수 없으므로 조정된 모델을 영구 스토리지로 이동하거나 Hugging Face Hub에 업로드해야 합니다.

삭제

추가 비용이 발생하지 않도록 하려면 이 튜토리얼 중에 만든 리소스를 삭제합니다.

TPU VM 인스턴스 삭제

Cloud TPU VM 인스턴스를 삭제합니다.

gcloud compute instances delete "${TPU_NAME}" --zone="${ZONE}" --project="${PROJECT}" --quiet

다음 단계

Cloud TPU에 대한 자세한 내용은 Cloud TPU 소개를 참조하세요.
v6e-8 TPU의 아키텍처 및 구성 세부정보는 TPU v6e를 참조하세요.