TPU v6e를 사용한 모델 학습
이 문서에서는 Cloud TPU v6e (Trillium이라고도 함)에서 모델을 학습하는 방법을 안내하며, 환경 설정, 성능 최적화, JAX 및 PyTorch/XLA를 사용한 실제 학습 예시를 다룹니다.
TPU v6e(Trillium이라고도 함)는 Google의 6세대 TPU입니다. API 및 로그와 같은 모든 기술적 측면과 이 문서 전체에서 Trillium은 v6e로 지칭됩니다. 포드당 256개의 칩이 있는 TPU v6e의 아키텍처는 v5e와 유사한 점이 많습니다. TPU v6e는 변환기, 텍스트-이미지, 컨볼루셔널 신경망 (CNN) 학습, 미세 조정, 서빙에 최적화되어 있습니다. TPU v6e 시스템 아키텍처 및 구성에 관한 자세한 내용은 TPU v6e를 참고하세요.
Cloud TPU v6e에서 추론을 실행하는 방법에 대한 자세한 내용은 다음 튜토리얼을 참고하세요.
- v6e에서 JetStream MaxText 추론
- v6e에서 JetStream MaxText 추론
- v6e에서 MaxDiffusion 추론
- v6e에서 vLLM 추론
- 경로를 사용하여 멀티 호스트 추론 실행
시작하기 전에
시작하기 전에 다음을 수행해야 합니다.
- 결제가 사용 설정된 Google Cloud 계정 및 프로젝트 만들기
- Google Cloud CLI 알파 구성요소 설치
- Cloud TPU API 사용 설정
- Cloud TPU 서비스 에이전트 만들기
- Cloud TPU 서비스 계정 만들기 및 권한 부여
자세한 내용은 Cloud TPU 환경 설정을 참조하세요.
할당량 및 권한 확인
프로젝트에 다음 할당량이 있는지 확인합니다.
XPK와 함께 GKE를 사용하는 경우 Google Cloud 콘솔에 추가 권한이 필요합니다. 자세한 내용은 Google Cloud 콘솔에 필요한 권한 을 참고하세요.
TPU 프로비저닝
다음 방법을 사용하여 TPU v6e를 프로비저닝하고 관리할 수 있습니다.
- GKE: GKE를 사용하여 컨테이너화된 머신러닝 워크로드용 가속기 풀로 TPU를 프로비저닝하고 관리할 수 있습니다. 자세한 내용은 GKE의 TPU 정보를 참고하세요.
- GKE 및 XPK: XPK는 GKE에서 클러스터 생성 및 워크로드 실행을 간소화하는 명령줄 도구입니다. ML 실무자가 심도 있는 Kubernetes 전문 지식 없이도 TPU를 프로비저닝하고 학습 작업을 실행할 수 있도록 설계되었습니다. 자세한 내용은 XPK GitHub 저장소를 참고하세요.
- Cloud TPU 큐에 추가된 리소스: 큐에 추가된 리소스를 사용하면 TPU 용량을 요청할 수 있으며, 이 용량은 사용 가능해지면 프로비저닝됩니다. 큐에서 대기할 수 있는 일괄 작업 및 내결함성 워크로드에 적합합니다. 요청의 시간 범위를 지정할 수 있습니다. 자세한 내용은 큐에 추가된 리소스 관리를 참고하세요.
GKE 및 XPK로 v6e Cloud TPU 프로비저닝
v6e와 함께 GKE 명령어를 사용하는 경우 Kubernetes 명령어 또는 XPK를 사용하여 Cloud TPU를 프로비저닝하고 모델을 학습시키거나 제공할 수 있습니다. GKE 클러스터에서 Cloud TPU 구성을 계획하는 방법은 GKE에서 Cloud TPU 계획을 참고하세요. 다음 섹션에서는 단일 NIC 지원 및 멀티 NIC 지원이 있는 XPK 클러스터를 만드는 명령어를 제공합니다.
단일 NIC를 지원하는 XPK 클러스터 만들기
export CLUSTER_NAME=xpk-cluster-name export ZONE=us-east1-d export PROJECT_ID=your-project-id export TPU_TYPE=v6e-256 export NUM_SLICES=2 export NETWORK_NAME=${CLUSTER_NAME}-mtu9k export NETWORK_FW_NAME=${NETWORK_NAME}-fw
gcloud compute networks create ${NETWORK_NAME} \ --mtu=8896 \ --project=${PROJECT_ID} \ --subnet-mode=auto \ --bgp-routing-mode=regional
gcloud compute firewall-rules create ${NETWORK_FW_NAME} \ --network=${NETWORK_NAME} \ --allow tcp,icmp,udp \ --project=${PROJECT_ID}
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
python3 xpk.py cluster create --cluster=${CLUSTER_NAME} \ --cluster-cpu-machine-type=e2-standard-8 \ --num-slices=${NUM_SLICES} \ --tpu-type=${TPU_TYPE} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --on-demand \ --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \ --create-vertex-tensorboard
명령어 플래그 설명
| 변수 | 설명 |
| CLUSTER_NAME | XPK 클러스터의 사용자 할당 이름입니다. |
| PROJECT_ID | Google Cloud 프로젝트 이름. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다. 자세한 내용은 Google Cloud 프로젝트 설정을 참조하세요. |
| ZONE | 지원되는 영역에 대해서는 Cloud TPU 리전 및 영역 문서를 참조하세요. |
| TPU_TYPE | 가속기 유형을 참조하세요. |
| NUM_SLICES | 만들려는 슬라이스의 수입니다. |
| CLUSTER_ARGUMENTS | 사용할 네트워크 및 서브네트워크입니다.
예: |
| NUM_SLICES | 생성할 슬라이스 수입니다. |
| NETWORK_NAME | 사용할 보조 네트워크의 이름입니다. |
| NETWORK_FW_NAME | 사용할 보조 네트워크 방화벽의 이름입니다. |
멀티 NIC를 지원하는 XPK 클러스터 만들기
export CLUSTER_NAME=xpk-cluster-name export REGION=your-region export ZONE=us-east1-d export PROJECT_ID=your-project-id export TPU_TYPE=v6e-256 export NUM_SLICES=2 export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE} export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE} export NETWORK_FW_NAME_1=${NETWORK_NAME_1}-fw-1-${ZONE} export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-1-${ZONE} export ROUTER_NAME=${CLUSTER_NAME}-network-1-${ZONE} export NAT_CONFIG=${CLUSTER_NAME}-natconfig-1-${ZONE}
gcloud compute networks create ${NETWORK_NAME_1} \ --mtu=8896 \ --bgp-routing-mode=regional \ --subnet-mode=custom \ --project=${PROJECT_ID}
gcloud compute networks subnets create ${SUBNET_NAME_1} \ --network=${NETWORK_NAME_1} \ --range=10.11.0.0/18 \ --region=${REGION} \ --project=${PROJECT_ID}
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \ --network=${NETWORK_NAME_1} \ --allow tcp,icmp,udp \ --project=${PROJECT_ID}
gcloud compute routers create ${ROUTER_NAME} \ --project=${PROJECT_ID} \ --network=${NETWORK_NAME_1} \ --region=${REGION}
gcloud compute routers nats create ${NAT_CONFIG} \ --router=${ROUTER_NAME} \ --region=${REGION} \ --auto-allocate-nat-external-ips \ --nat-all-subnet-ip-ranges \ --project=${PROJECT_ID} \ --enable-logging
# Secondary subnet for multi-nic experience.
# Need custom IP routing to be different from the first network's subnet.
export NETWORK_NAME_2=${CLUSTER_NAME}-privatenetwork-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-2-${ZONE}
gcloud compute networks create ${NETWORK_NAME_2} \ --mtu=8896 \ --bgp-routing-mode=regional \ --subnet-mode=custom \ --project=${PROJECT_ID}
gcloud compute networks subnets create ${SUBNET_NAME_2} \ --network=${NETWORK_NAME_2} \ --range=10.10.0.0/18 \ --region=${REGION} \ --project=${PROJECT_ID}
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \ --network=${NETWORK_NAME_2} \ --allow tcp,icmp,udp \ --project=${PROJECT_ID}
gcloud compute routers create ${ROUTER_NAME} \ --project=${PROJECT_ID} \ --network=${NETWORK_NAME_2} \ --region=${REGION}
gcloud compute routers nats create ${NAT_CONFIG} \ --router=${ROUTER_NAME} \ --region=${REGION} \ --auto-allocate-nat-external-ips \ --nat-all-subnet-ip-ranges \ --project=${PROJECT_ID} \ --enable-logging
export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"
python3 xpk.py cluster create \ --cluster=${CLUSTER_NAME} \ --cluster-cpu-machine-type=e2-standard-8 \ --num-slices=${NUM_SLICES} \ --tpu-type=${TPU_TYPE} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --on-demand \ --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \ --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \ --create-vertex-tensorboard
명령어 플래그 설명
| 변수 | 설명 |
| CLUSTER_NAME | XPK 클러스터의 사용자 할당 이름입니다. |
| PROJECT_ID | Google Cloud 프로젝트 이름. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다. 자세한 내용은 Google Cloud 프로젝트 설정을 참조하세요. |
| ZONE | 지원되는 영역에 대해서는 Cloud TPU 리전 및 영역 문서를 참조하세요. |
| TPU_TYPE | 가속기 유형을 참조하세요. |
| NUM_SLICES | 만들려는 슬라이스의 수입니다. |
| CLUSTER_ARGUMENTS | 사용할 네트워크 및 서브네트워크입니다.
예: |
| NODE_POOL_ARGUMENTS | 사용할 추가 노드 네트워크입니다.
예: |
| NUM_SLICES | 생성할 슬라이스 수입니다(멀티슬라이스에만 필요). |
| NETWORK_NAME | 사용할 보조 네트워크의 이름입니다. |
| NETWORK_FW_NAME | 사용할 보조 네트워크 방화벽의 이름입니다. |
JAX 또는 PyTorch 설정
다음 리소스에서는 사용하는 프로비저닝 및 관리 방법에 따라 Cloud TPU에서 JAX 또는 PyTorch를 설정하는 방법을 보여줍니다.
- GKE Autopilot: TPU 애플리케이션 준비
- GKE Standard: 워크로드 준비
- GKE 및 XPK: XPK 리드미
- JAX를 사용하는 단일 호스트 Cloud TPU: JAX를 사용하여 Cloud TPU VM에서 계산 실행
- JAX를 사용하는 멀티 호스트 Cloud TPU: TPU 슬라이스에서 JAX 코드 실행
- PyTorch를 사용하는 단일 호스트 Cloud TPU: PyTorch를 사용하여 Cloud TPU VM에서 계산 실행
- PyTorch를 사용하는 멀티 호스트 Cloud TPU: TPU 슬라이스에서 PyTorch 코드 실행
MaxText로 XPK를 설정하고 실행하려면 XPK를 사용하여 대규모로 MaxText 실행 을 참고하세요.
네트워크 성능 최적화
이 섹션에서는 최대 전송 단위 (MTU)를 구성하고, 멀티 슬라이스 환경에 다중 NIC를 사용하고, TCP 설정을 개선하여 네트워크 성능을 최적화하는 방법을 설명합니다.
MTU 구성
최적의 네트워크 성능을 위해 MTU (최대 전송 단위)가 8,896인 네트워크를 사용하세요.
기본적으로 가상 프라이빗 클라우드 (VPC)는 1,460바이트의 MTU만 제공하므로 최적의 네트워크 성능을 제공하지 않습니다. VPC 네트워크의 MTU를 1,300바이트에서 8,896바이트(포함) 사이의 값으로 설정할 수 있습니다. 일반적인 커스텀 MTU 크기는 1,500바이트(표준 이더넷) 또는 8,896바이트(최대)입니다. 자세한 내용은 유효한 VPC 네트워크 MTU 크기를 참조하세요.
기존 또는 기본 네트워크의 MTU 설정을 변경하는 방법에 관한 자세한 내용은 VPC 네트워크의 MTU 설정 변경을 참고하세요.
다음 예시에서는 MTU가 8,896인 네트워크와 네트워크 내에서 TCP, ICMP, UDP 트래픽을 허용하는 해당 방화벽 규칙을 만듭니다.
export RESOURCE_NAME=your-resource-name export NETWORK_NAME=${RESOURCE_NAME}-privatenetwork export NETWORK_FW_NAME=${RESOURCE_NAME}-privatefirewall gcloud compute networks create ${NETWORK_NAME} --mtu=8896 --project=${PROJECT_ID} \ --subnet-mode=auto --bgp-routing-mode=regional gcloud compute firewall-rules create ${NETWORK_FW_NAME} --network=${NETWORK_NAME} \ --allow tcp,icmp,udp --project=${PROJECT_ID}
your-resource-name을 네트워크 및 방화벽의 기본 이름으로 바꿉니다.
멀티슬라이스에 멀티 NIC 옵션 사용
멀티슬라이스 환경을 사용하는 경우 보조 서브넷에 필요한 다음 환경 변수를 설정합니다.
export NETWORK_NAME_2=${RESOURCE_NAME} export SUBNET_NAME_2=${RESOURCE_NAME} export FIREWALL_RULE_NAME=${RESOURCE_NAME} export ROUTER_NAME=${RESOURCE_NAME}-network-2 export NAT_CONFIG=${RESOURCE_NAME}-natconfig-2 export REGION=your-region
다음 명령어를 사용하여 네트워크 및 서브넷에 대한 커스텀 IP 라우팅을 만듭니다.
보조 네트워크를 만듭니다.
gcloud compute networks create ${NETWORK_NAME_2} --mtu=8896 \ --bgp-routing-mode=regional --subnet-mode=custom --project=${PROJECT_ID}보조 네트워크의 서브네트워크를 만듭니다.
gcloud compute networks subnets create ${SUBNET_NAME_2} \ --network=${NETWORK_NAME_2} \ --range=10.10.0.0/18 --region=${REGION} \ --project=${PROJECT_ID}새 서브네트워크 내에서 트래픽을 허용하는 방화벽 규칙을 만듭니다.
gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \ --network=${NETWORK_NAME_2} --allow tcp,icmp,udp \ --source-ranges 10.10.0.0/18 --project=${PROJECT_ID}보조 네트워크의 Cloud Router를 만듭니다.
gcloud compute routers create ${ROUTER_NAME} \ --project=${PROJECT_ID} \ --network=${NETWORK_NAME_2} \ --region=${REGION}Cloud Router의 NAT 구성을 만듭니다.
gcloud compute routers nats create ${NAT_CONFIG} \ --router=${ROUTER_NAME} \ --region=${REGION} \ --auto-allocate-nat-external-ips \ --nat-all-subnet-ip-ranges \ --project=${PROJECT_ID} \ --enable-logging
멀티 네트워크 슬라이스를 만든 후 XPK 클러스터를 설정하고 XPK 워크로드 생성 명령어에 --command ifconfig 플래그를 추가하여 두 네트워크 인터페이스 카드 (NIC)가 모두 사용되고 있는지 확인할 수 있습니다.
다음
workload create명령어를 사용하여 Google Cloud 콘솔 로그에ifconfig명령어의 출력을 표시하고 eth0과 eth1 모두 MTU가 8,896으로 설정되어 있는지 확인합니다.python3 xpk.py workload create \ --cluster CLUSTER_NAME \ {--base-docker-image maxtext_base_image | --docker-image your-cloud-image-name} \ --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \ --tpu-type=${ACCELERATOR_TYPE} \ --num-slices=${NUM_SLICES} \ --on-demand \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --command "ifconfig"
디버그 로그를 사용 설정하거나 Vertex AI TensorBoard를 사용하려면 다음 선택적 인수를 명령어에 추가합니다.
--enable-debug-logs \ --use-vertex-tensorboard
Google Cloud 콘솔 로그에서 XPK 워크로드의 출력을 확인하여 eth0과 eth1 모두 MTU가 8,896으로 설정되어 있는지 확인합니다.
TCP 설정 향상
큐에 추가된 리소스를 사용하여 Cloud TPU를 프로비저닝한 경우 다음 명령어를 실행하여 TCP 수신 버퍼 제한을 늘려 네트워크 성능을 개선할 수 있습니다.
gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \ --project "${PROJECT_ID}" \ --zone "${ZONE}" \ --node=all \ --worker=all \ --command=' sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"'
메모리 할당 성능 최적화
tcmalloc 라이브러리는 상당한 크기의 메모리 할당이 빈번한 모델의 성능을 개선하기 위해 Cloud TPU VM에서 기본적으로 사용됩니다. LD_PRELOAD 환경 변수를 통해 구성됩니다.
하지만 일부 워크로드 (예: 임베딩 테이블 할당이 매우 큰 DLRM)의 경우 tcmalloc을 사용하면 느려질 수 있습니다. 이러한 경우 학습 스크립트를 실행하기 전에 셸 세션에서 LD_PRELOAD 변수를 설정 해제하여 표준 malloc 함수로 되돌릴 수 있습니다.
unset LD_PRELOAD
SkyPilot 사용
SkyPilot과 함께 Cloud TPU v6e를 사용할 수 있습니다. SkyPilot은 AI 워크로드의 실행, 관리, 확장을 간소화하는 오픈소스 프레임워크입니다. v6e 관련 위치 및 가격 책정 정보를 SkyPilot에 추가할 수 있습니다. 자세한 내용은 SkyPilot TPU v6e 예시를 참고하세요.
학습 예시
다음 섹션에서는 Cloud TPU v6e에서 MaxText, MaxDiffusion, PyTorch 모델을 학습시키는 예를 제공합니다.
이 예는 다음 소프트웨어 버전으로 테스트되었습니다.
- Python
3.10이상 - Nightly 소프트웨어 버전:
- 야간 JAX
0.4.32.dev20240912 - Nightly LibTPU
0.1.dev20240912+nightly
- 야간 JAX
- 안정적인 소프트웨어 버전:
- JAX + JAX Lib(v0.4.37)
Cloud TPU v6e에서 MaxText 및 MaxDiffusion 학습
다음 섹션에서는 MaxText 및 MaxDiffusion 모델의 학습 수명 주기를 다룹니다.
일반적으로 대략적인 단계는 다음과 같습니다.
- 워크로드 기본 이미지를 빌드합니다.
- XPK를 사용하여 워크로드를 실행합니다.
- 워크로드에 대한 학습 명령어를 빌드합니다.
- 워크로드를 배포합니다.
- 워크로드를 모니터링하고 측정항목을 확인합니다.
- 필요하지 않은 경우 XPK 워크로드를 삭제합니다.
- 더 이상 필요하지 않은 경우 XPK 클러스터를 삭제합니다.
기본 이미지 빌드
MaxText 또는 MaxDiffusion을 설치하고 Docker 이미지를 빌드합니다.
사용하려는 저장소를 클론하고 저장소의 디렉터리로 변경합니다.
MaxText:
git clone https://github.com/google/maxtext.git && cd maxtextMaxDiffusion:
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103Google Cloud CLI를 사용하도록 Docker를 구성합니다.
gcloud auth configure-docker다음 명령어를 사용하거나 JAX AI 이미지를 사용하여 Docker 이미지를 빌드합니다. JAX AI 이미지에 대한 자세한 내용은 JAX AI 이미지를 참고하세요.
MaxText:
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35MaxDiffusion:
bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest활성 gcloud CLI 구성에서 프로젝트 ID를 설정합니다.
gcloud config set project ${PROJECT_ID}로컬에 이미지가 빌드되어 있지 않은 머신에서 워크로드를 실행하려는 경우 해당 이미지를 업로드합니다.
CLOUD_IMAGE_NAME환경 변수를 설정합니다.export CLOUD_IMAGE_NAME=${USER}_runner이미지를 업로드합니다.
bash docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
XPK를 사용하여 워크로드 실행
MaxText에서 설정한 기본값 또는 MaxDiffusion을 사용하지 않는 경우 다음 환경 변수를 설정합니다.
export BASE_OUTPUT_DIR=gs://YOUR_BUCKET export PER_DEVICE_BATCH_SIZE=2 export NUM_STEPS=30 export MAX_TARGET_LENGTH=8192
모델 스크립트를 빌드합니다. 이 스크립트는 나중에 학습 명령어로 복사됩니다.
아직 모델 스크립트를 실행하지 마세요.
MaxText
MaxText는 순수 Python 및 JAX로 작성된 고성능의 확장성이 뛰어난 오픈소스 LLM으로, 학습 및 추론을 위해 Google Cloud TPU 및 GPU를 타겟팅합니다.
JAX_PLATFORMS=tpu,cpu \ ENABLE_PJRT_COMPATIBILITY=true \ TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \ TPU_SLICE_BUILDER_DUMP_ICI=true && \ python3 -m MaxText.train MaxText/configs/base.yml \ base_output_directory=${BASE_OUTPUT_DIR} \ dataset_type=synthetic \ per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ enable_checkpointing=false \ gcs_metrics=true \ profiler=xplane \ skip_first_n_steps_for_profiler=5 \ steps=${NUM_STEPS} # attention='dot_product'"Gemma2
Gemma는 Gemini 연구 및 기술을 기반으로 Google DeepMind에서 개발한 오픈 가중치 LLM 제품군입니다.
python3 -m MaxText.train MaxText/configs/base.yml \ model_name=gemma2-27b \ run_name=gemma2-27b-run \ base_output_directory=${BASE_OUTPUT_DIR} \ max_target_length=${MAX_TARGET_LENGTH} \ per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ steps=${NUM_STEPS} \ enable_checkpointing=false \ use_iota_embed=true \ gcs_metrics=true \ dataset_type=synthetic \ profiler=xplane \ attention=flashMixtral 8x7b
Mixtral은 Mistral AI에서 개발한 최신 AI 모델로, 희소 전문가 망(MoE) 아키텍처를 활용합니다.
python3 -m MaxText.train MaxText/configs/base.yml \ base_output_directory=${BASE_OUTPUT_DIR} \ per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \ model_name=mixtral-8x7b \ steps=${NUM_STEPS} \ max_target_length=${MAX_TARGET_LENGTH} \ tokenizer_path=assets/tokenizer.mistral-v1 \ attention=flash \ dtype=bfloat16 \ dataset_type=synthetic \ profiler=xplaneLlama3-8b
Llama는 Meta에서 개발한 개방형 가중치 LLM 제품군입니다.
PyTorch에서 Llama3를 실행하는 방법에 대한 예시는 torchprime 저장소의 torch_xla 모델을 참조하세요.
MaxDiffusion
MaxDiffusion은 Cloud TPU 및 GPU를 비롯한 XLA 기기에서 실행되는 다양한 잠재 확산 모델의 참조 구현 모음으로, 순수 Python과 JAX로 작성되었습니다. Stable Diffusion은 모든 텍스트 입력에서 실사 이미지를 생성하는 잠재 텍스트 이미지 변환 모델입니다.
다음 학습 스크립트에 표시된 대로 MaxDiffusion을 실행하려면 특정 Git 브랜치를 설치해야 합니다.
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103 && pip install -r requirements.txt && pip install . && pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR} && python src/maxdiffusion/train_sdxl.py \ src/maxdiffusion/configs/base_xl.yml \ revision=refs/pr/95 \ activations_dtype=bfloat16 \ weights_dtype=bfloat16 \ resolution=1024 \ per_device_batch_size=1 \ output_dir=${OUT_DIR} \ jax_cache_dir=${OUT_DIR}/cache_dir/ \ max_train_steps=200 \ attention=flash \ run_name=sdxl-ddp-v6e다음 변수를 내보냅니다.
export CLUSTER_NAME=CLUSTER_NAME export ACCELERATOR_TYPE=ACCELERATOR_TYPE export NUM_SLICES=NUM_SLICES export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT
환경 변수 설명
변수 설명 CLUSTER_NAMEXPK 클러스터의 이름입니다. ACCELERATOR_TYPE가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참고하세요. NUM_SLICESTPU 슬라이스 수입니다. YOUR_MODEL_SCRIPT학습 명령어로 실행할 모델 스크립트입니다. 이전 단계에서 만든 스크립트를 사용하여 모델을 실행합니다. MaxText 기본 이미지를 사용하려면
--base-docker-image플래그를 지정하거나--docker-image플래그와 사용할 이미지를 지정해야 합니다.다음 선택적 플래그를 추가할 수 있습니다.
--enable-debug-logs플래그를 포함하여 디버그 로깅을 사용 설정할 수 있습니다. 자세한 내용은 MaxText에서 JAX 디버그를 참고하세요.--use-vertex-tensorboard플래그를 포함하여 Vertex AI Experiments를 만들어 Vertex AI TensorBoard에 데이터를 업로드할 수 있습니다. 자세한 내용은 Vertex AI를 사용하여 MaxText에서 JAX 모니터링을 참고하세요.
python3 xpk.py workload create \ --cluster ${CLUSTER_NAME} \ {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \ --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \ --tpu-type=${ACCELERATOR_TYPE} \ --num-slices=${NUM_SLICES} \ --on-demand \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --command="${YOUR_MODEL_SCRIPT}"
출력에는 워크로드를 추적할 수 있는 링크가 포함됩니다. 링크를 열고 로그 탭을 클릭하여 워크로드를 실시간으로 추적합니다.
MaxText에서 JAX 디버그
추가 XPK 명령어를 사용하여 클러스터 또는 워크로드가 실행되지 않는 이유를 진단합니다.
- XPK 워크로드 목록
- XPK 검사기
- XPK 워크로드를 만들 때
--enable-debug-logs플래그를 사용하여 워크로드 로그에서 상세 로깅을 사용 설정합니다.
Vertex AI를 사용하여 MaxText에서 JAX 모니터링
TensorBoard를 사용하려면 Google Cloud 사용자 계정에 aiplatform.user 역할이 있어야 합니다. 다음 명령어를 실행하여 이 역할을 부여합니다.
gcloud projects add-iam-policy-binding your-project-id \ --member='user:your-email' \ --role='roles/aiplatform.user'
Vertex AI 관리형 TensorBoard를 통해 스칼라 및 프로필 데이터를 확인합니다.
사용 중인 영역에 대한 리소스 관리(CRUD) 요청을 600에서 5,000으로 늘립니다. 16개 미만의 VM을 사용하는 소규모 워크로드의 경우 문제가 되지 않을 수 있습니다.
Vertex AI용
cloud-accelerator-diagnostics와 같은 종속 항목을 설치합니다.# xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI cd ~/xpk pip install .
Vertex AI TensorBoard 만들기에 설명된 대로
--create-vertex-tensorboard플래그를 사용하여 XPK 클러스터를 만듭니다. 기존 클러스터에서도 이 명령어를 실행할 수 있습니다.--use-vertex-tensorboard플래그와 선택적--experiment-name플래그를 사용하여 XPK 워크로드 실행 시 Vertex AI Experiments를 만듭니다. 전체 단계의 목록은 Vertex AI Experiments를 만들어 Vertex AI TensorBoard에 데이터 업로드를 참조하세요.
로그에는 다음과 유사한 Vertex AI TensorBoard 링크가 포함됩니다.
View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name
Vertex AI TensorBoard 링크는 Google Cloud 콘솔에서도 확인할 수 있습니다. Google Cloud 콘솔에서 Vertex AI Experiments로 이동합니다. 드롭다운에서 적절한 리전을 선택합니다.
TensorBoard 디렉터리도 ${BASE_OUTPUT_DIR}로 지정한 Cloud Storage 버킷에 기록합니다.
XPK 워크로드 삭제
xpk workload delete 명령어를 사용하여 작업 접두사 또는 작업 상태를 기반으로 하나 이상의 워크로드를 삭제합니다. 이 명령어는 더 이상 실행할 필요가 없는 XPK 워크로드를 전송했거나 큐에 멈춰 있는 작업이 있는 경우에 유용할 수 있습니다.
XPK 클러스터 삭제
xpk cluster delete 명령어를 사용하여 클러스터를 삭제합니다.
python3 xpk.py cluster delete --cluster ${CLUSTER_NAME} \ --zone=${ZONE} --project=${PROJECT_ID}
MaxDiffusion 벤치마킹 결과
v6e-4, v6e-16, 2개의 v6e-16에서 MaxDiffusion용 학습 스크립트를 실행했습니다. 다음 표에서는 측정된 처리량을 보여줍니다.
| v6e-4 | v6e-16 | v6e-16 2개 | |
|---|---|---|---|
| 학습 단계 | 0.069 | 0.073 | 0.13 |
| 전역 배치 크기 | 8 | 32 | 64 |
| 처리량(예시/초) | 115.9 | 438.4 | 492.3 |
Cloud TPU v6e에서 PyTorch/XLA를 사용하여 Llama 모델 학습
이 섹션에서는 WikiText 데이터 세트를 사용하여 Cloud TPU v6e에서 PyTorch/XLA를 사용하여 Llama 모델을 학습시키는 방법을 설명합니다.
Hugging Face 및 Llama 3 모델에 액세스
이 예시에는 Hugging Face 사용자 액세스 토큰이 필요합니다. 사용자 액세스 토큰 생성에 관한 자세한 내용은 사용자 액세스 토큰에 관한 Hugging Face 문서를 참조하세요.
Hugging Face의 Llama-3-8B 모델에 액세스할 수 있는 권한도 필요합니다. 액세스 권한을 얻으려면 HuggingFace의 Meta-Llama-3-8B 모델로 이동하여 액세스를 요청하세요.
Cloud TPU VM 만들기
이 예에서는 칩 8개가 있는 Cloud TPU v6e를 만듭니다.
환경 변수를 설정합니다.
export PROJECT_ID=your-project-id export TPU_NAME=your-tpu-name export ZONE=us-east1-d export ACCELERATOR_TYPE=v6e-8 export RUNTIME_VERSION=v2-alpha-tpuv6e
환경 변수 설명
변수 설명 PROJECT_IDGoogle Cloud 프로젝트 ID입니다. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다. TPU_NAMETPU의 이름입니다. ZONETPU VM을 만들 영역입니다. 지원되는 영역에 대한 자세한 내용은 TPU 리전 및 영역을 참조하세요. ACCELERATOR_TYPE가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요. RUNTIME_VERSIONCloud TPU 소프트웨어 버전입니다. Cloud TPU VM을 만듭니다.
gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${RUNTIME_VERSION} \ --accelerator-type=${ACCELERATOR_TYPE} \ --zone=${ZONE} \ --project=${PROJECT_ID}
설치
Hugging Face Transformer 및 종속 항목의 pytorch-tpu/transformers 포크를 설치합니다. 이 예시는 다음 종속 항목 버전으로 테스트되었습니다.
torch: 2.5.0과 호환torch_xla[tpu]: 2.5.0과 호환jax: 0.4.33jaxlib: 0.4.33
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git cd transformers sudo pip3 install -e . pip3 install datasets pip3 install evaluate pip3 install scikit-learn pip3 install accelerate pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/'
모델 구성 파일 설정
다음 섹션의 학습 명령어인 모델 실행은 두 개의 JSON 구성 파일을 사용하여 모델 매개변수와 완전 샤딩된 데이터 동시 로드 (FSDP) 구성을 정의합니다. FSDP 샤딩을 사용하면 여러 TPU에 모델 가중치를 샤딩하여 학습하는 동안 더 큰 배치 크기를 사용할 수 있습니다. 소형 모델로 학습할 때는 데이터 동시 로드를 사용하고 각 기기에서 가중치를 복제하는 것으로 충분할 수 있습니다. PyTorch/XLA에서 기기 간에 텐서를 샤딩하는 방법에 관한 자세한 내용은 PyTorch/XLA SPMD 사용자 가이드를 참조하세요.
모델 매개변수 구성 파일을 만듭니다. 다음은 Llama-3-8B의 모델 매개변수 구성입니다. 다른 모델의 경우 Hugging Face에서 구성 파일을 찾으세요. 예를 들어 Llama-2-7B 구성을 참고하세요.
cat > llama-config.json << EOF { "architectures": [ "LlamaForCausalLM" ], "attention_bias": false, "attention_dropout": 0.0, "bos_token_id": 128000, "eos_token_id": 128001, "hidden_act": "silu", "hidden_size": 4096, "initializer_range": 0.02, "intermediate_size": 14336, "max_position_embeddings": 8192, "model_type": "llama", "num_attention_heads": 32, "num_hidden_layers": 32, "num_key_value_heads": 8, "pretraining_tp": 1, "rms_norm_eps": 1e-05, "rope_scaling": null, "rope_theta": 500000.0, "tie_word_embeddings": false, "torch_dtype": "bfloat16", "transformers_version": "4.40.0.dev0", "use_cache": false, "vocab_size": 128256 } EOFFSDP 구성 파일을 만듭니다.
cat > fsdp-config.json << EOF { "fsdp_transformer_layer_cls_to_wrap": [ "LlamaDecoderLayer" ], "xla": true, "xla_fsdp_v2": true, "xla_fsdp_grad_ckpt": true } EOFFSDP에 대한 자세한 내용은 SPMD를 사용하는 완전 샤딩된 데이터 동시 로드 를 참고하세요.
다음 명령어를 사용하여 구성 파일을 Cloud TPU VM에 업로드합니다.
gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \ --worker=all \ --project=${PROJECT_ID} \ --zone=${ZONE}
모델 실행
이전 섹션에서 만든 구성 파일을 사용하여 run_clm.py 스크립트를 실행하여 WikiText 데이터 세트에서 Llama-3-8B 모델을 학습시킵니다. 학습 스크립트는 Cloud TPU v6e-8에서 실행하는 데 약 10분이 걸립니다.
다음 명령어를 사용하여 Cloud TPU에서 Hugging Face에 로그인합니다.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command=' pip3 install "huggingface_hub[cli]" huggingface-cli login --token HUGGING_FACE_TOKEN'
모델 학습을 실행합니다.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone ${ZONE} \ --worker=all \ --command=' export PJRT_DEVICE=TPU export XLA_USE_SPMD=1 export ENABLE_PJRT_COMPATIBILITY=true # Optional variables for debugging: export XLA_IR_DEBUG=1 export XLA_HLO_DEBUG=1 export PROFILE_EPOCH=0 export PROFILE_STEP=3 export PROFILE_DURATION_MS=100000 # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path export PROFILE_LOGDIR=PROFILE_PATH python3 transformers/examples/pytorch/language-modeling/run_clm.py \ --dataset_name wikitext \ --dataset_config_name wikitext-2-raw-v1 \ --per_device_train_batch_size 16 \ --do_train \ --output_dir /home/$USER/tmp/test-clm \ --overwrite_output_dir \ --config_name /home/$USER/llama-config.json \ --cache_dir /home/$USER/cache \ --tokenizer_name meta-llama/Meta-Llama-3-8B \ --block_size 8192 \ --optim adafactor \ --save_strategy no \ --logging_strategy no \ --fsdp "full_shard" \ --fsdp_config /home/$USER/fsdp-config.json \ --torch_dtype bfloat16 \ --dataloader_drop_last yes \ --flash_attention \ --max_steps 20'
PyTorch/XLA 문제 해결
이전 섹션에서 디버깅을 위한 선택적 변수를 설정하면 모델에 대한 프로필이 PROFILE_LOGDIR 변수로 지정된 위치에 저장됩니다. 이 위치에 저장된 xplane.pb 파일을 추출하고 tensorboard를 사용하여 TensorBoard 안내에 따라 브라우저에서 프로필을 볼 수 있습니다.
PyTorch/XLA가 예상대로 작동하지 않으면 모델 디버깅, 프로파일링, 최적화에 관한 제안사항이 포함된 문제 해결 가이드를 참고하세요.