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에서 추론을 실행하는 방법에 대한 자세한 내용은 다음 튜토리얼을 참조하세요.
시작하기 전에
시작하기 전에 다음을 수행해야 합니다.
- 결제가 사용 설정된 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 클러스터를 설정하고 --command ifconfig 플래그를 XPK 워크로드 생성 명령어에 추가하여 두 네트워크 인터페이스 카드(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='