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

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

목표

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

비용

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

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

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

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

시작하기 전에

  • 이 튜토리얼을 사용하려면 Hugging Face 액세스 토큰이 필요합니다. Hugging Face에서 무료 계정을 등록할 수 있습니다. 계정이 있으면 액세스 토큰을 생성합니다.

    1. Hugging Face 시작 페이지에서 계정 아바타를 클릭하고 액세스 토큰을 선택합니다.
    2. 액세스 토큰 페이지에서 새 토큰 만들기 를 클릭합니다.
    3. 읽기 토큰 유형을 선택하고 토큰 이름을 입력합니다.
    4. 액세스 토큰이 표시됩니다. 토큰을 안전한 곳에 저장합니다.

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

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

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

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

환경 설정

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

export PROJECT="YOUR_PROJECT_ID"
export ZONE="ZONE_NAME"
export RESERVATION="RESERVATION_NAME"
export TPU_NAME="TPU_MACHINE_NAME"

다음을 바꿉니다.

  • YOUR_PROJECT_ID: 프로젝트 ID Google Cloud
  • ZONE_NAME: 사용할 영역
  • RESERVATION_NAME: 용량 예약
  • TPU_MACHINE_NAME: Cloud TPU VM 인스턴스의 이름

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

gcloud auth login

Cloud TPU VM 만들기

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

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project=$PROJECT \
    --accelerator-type=v6e-8 \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

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

gcloud compute tpus tpu-vm 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 python3.12 python3.12-venv

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

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

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

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

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

uv pip install maxtext[tpu-post-train]==0.2.2 --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 워크로드를 시작할 수 있습니다.

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

    # -- MaxText configuration --
export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
export CHIPS_PER_VM=8
export NUM_BATCHES=50
export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items

  2. 학습 작업을 시작합니다. 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 형식으로 다시 변환합니다.

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

    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

  2. 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 alpha compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE --project=$PROJECT --quiet

다음 단계

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