TPU로 파이프라인 실행

이 페이지에서는 TPU를 사용하여 Dataflow에서 Apache Beam 파이프라인을 실행하는 방법을 설명합니다. TPU를 사용하는 작업에는 Dataflow 가격 책정 페이지에 명시된 대로 요금이 청구됩니다.

Dataflow에서 TPU를 사용하는 방법에 대한 자세한 내용은 TPU에 대한 Dataflow 지원을 참고하세요.

선택사항: 가속기를 사용하기 위해 특정 예약하기

TPU를 주문형으로 사용할 수 있지만 구체적으로 타겟팅된 Google Cloud예약과 함께 Dataflow TPU를 사용하는 것이 좋습니다. 이렇게 하면 사용 가능한 액셀러레이터와 빠른 작업자 시작 시간을 이용할 수 있습니다. TPU 예약을 사용하는 파이프라인에는 추가 TPU 할당량이 필요하지 않습니다.

예약을 하지 않고 주문형으로 TPU를 사용하려면 파이프라인을 실행하기 전에 TPU 할당량을 프로비저닝하세요.

선택사항: TPU 할당량 프로비저닝

주문형 용량으로 또는 예약을 사용하여 TPU를 사용할 수 있습니다. 주문형으로 TPU를 사용하려면 먼저 TPU 할당량을 프로비저닝해야 합니다. 구체적으로 타겟팅된 예약을 사용하는 경우 이 섹션을 건너뛰어도 됩니다.

예약 없이 주문형으로 TPU를 사용하려면 다음과 같이 TPU용 Compute Engine API 할당량의 한도 및 현재 사용량을 확인하세요.

콘솔

1. Google Cloud 콘솔에서 할당량 페이지로 이동합니다.

   할당량으로 이동

2. 필터 상자에서 다음을 수행합니다.

   1. 다음 표를 사용하여 TPU 버전 및 머신 유형에 따라 할당량의 속성을 선택하고 복사합니다. 예를 들어 머신 유형이 ct5lp-로 시작하는 주문형 TPU v5e 노드를 만들려면 Name: TPU v5 Lite PodSlice chips를 입력합니다.

      TPU 버전, 머신 유형이 다음으로 시작함	온디맨드 인스턴스의 할당량 속성 및 이름
      TPU v5e, ct5lp-	Name: TPU v5 Lite PodSlice chips
      TPU v5p, ct5p-	Name: TPU v5p chips
      TPU v6e, ct6e-	Dimensions (e.g. location): tpu_family:CT6E

   2. 측정기준 (예: 위치) 속성을 선택하고 region:을 입력한 다음 파이프라인을 시작할 리전의 이름을 입력합니다. 예를 들어 us-west4-a 영역을 사용하려면 region:us-west4를 입력합니다. TPU 할당량은 리전 기준이므로 동일한 리전 내의 모든 영역에서 동일한 TPU 할당량을 사용합니다.

커스텀 컨테이너 이미지 구성

Dataflow 파이프라인에서 TPU와 상호작용하려면 파이프라인 런타임 환경에서 XLA 기기에서 작동할 수 있는 소프트웨어를 제공해야 합니다. 이를 위해서는 파이프라인 요구사항에 따라 TPU 라이브러리를 설치하고 사용하는 TPU 기기에 따라 환경 변수를 구성해야 합니다.

컨테이너 이미지를 맞춤설정하려면 필요한 TPU 라이브러리가 있는 기성 기본 이미지에 Apache Beam을 설치합니다. 또는 Apache Beam SDK 출시와 함께 게시된 이미지에 TPU 소프트웨어를 설치합니다.

커스텀 컨테이너 이미지를 제공하려면 sdk_container_image 파이프라인 옵션을 사용하세요. 자세한 내용은 Dataflow에서 커스텀 컨테이너 사용을 참고하세요.

TPU 가속기를 사용하는 경우 컨테이너 이미지에서 다음 환경 변수를 설정해야 합니다.

ENV TPU_SKIP_MDS_QUERY=1 # Don't query metadata
ENV TPU_HOST_BOUNDS=1,1,1 # There's only one host
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0 # Always 0 for single-host TPUs

사용하는 가속기에 따라 다음 표의 변수도 설정해야 합니다.

유형	토폴로지	Dataflow worker_machine_type 필요	추가 환경 변수
tpu-v5-lite-podslice	1x1	ct5lp-hightpu-1t	TPU_ACCELERATOR_TYPE=v5litepod-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	TPU_ACCELERATOR_TYPE=v5litepod-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	TPU_ACCELERATOR_TYPE=v5litepod-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v6e-slice	1x1	ct6e-standard-1t	TPU_ACCELERATOR_TYPE=v6e-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v6e-slice	2x2	ct6e-standard-4t	TPU_ACCELERATOR_TYPE=v6e-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v6e-slice	2x4	ct6e-standard-8t	TPU_ACCELERATOR_TYPE=v6e-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	TPU_ACCELERATOR_TYPE=v5p-8 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1

커스텀 컨테이너 이미지의 샘플 Dockerfile은 다음 예와 같을 수 있습니다.

FROM python:3.11-slim

COPY --from=apache/beam_python3.11_sdk:2.66.0 /opt/apache/beam /opt/apache/beam

# Configure the environment to access TPU device

ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

# Configure the environment for the chosen accelerator.
# Adjust according to the accelerator you use.
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1

# Install TPU software stack.
RUN pip install jax[tpu] apache-beam[gcp]==2.66.0 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

ENTRYPOINT ["/opt/apache/beam/boot"]

TPU로 작업 실행

TPU로 Dataflow 작업을 실행할 때는 다음과 같은 사항을 고려해야 합니다.

• TPU 컨테이너가 클 수 있으므로 디스크 공간 부족을 방지하려면 --disk_size_gb 파이프라인 옵션을 사용하여 기본 부팅 디스크 크기를 50GB 또는 컨테이너 이미지에 필요한 적절한 크기로 늘리세요.
• 작업자 내 동시 로드를 제한합니다.

TPU 및 작업자 동시 로드

기본 구성에서 Dataflow Python 파이프라인은 VM 코어당 하나의 Apache Beam SDK 프로세스를 실행합니다. TPU 머신 유형에는 vCPU 코어가 많이 있지만 하나의 프로세스만 TPU 기기에서 계산을 실행할 수 있습니다. 또한 TPU 기기는 프로세스의 수명 동안 프로세스에 의해 예약될 수 있습니다. 따라서 Dataflow TPU 파이프라인을 실행할 때는 작업자 내 병렬 처리를 제한해야 합니다. 작업자 동시 로드를 제한하려면 다음 안내를 따르세요.

• 사용 사례에 모델에 대한 추론 실행이 포함되는 경우 Beam RunInference API를 사용합니다. 자세한 내용은 Beam의 대규모 언어 모델 추론을 참고하세요.
• Beam RunInference API를 사용할 수 없는 경우 Beam의 다중 프로세스 공유 객체를 사용하여 특정 작업을 단일 프로세스로 제한합니다.
• 앞의 권장사항을 사용할 수 없고 작업자당 하나의 Python 프로세스만 실행하려면 --experiments=no_use_multiple_sdk_containers 파이프라인 옵션을 설정하세요.
• vCPU가 100개를 초과하는 작업자의 경우 --number_of_worker_harness_threads 파이프라인 옵션을 사용하여 스레드 수를 줄입니다. 다음 표를 사용하여 TPU 유형이 100개가 넘는 vCPU를 사용하는지 확인하세요.

다음 표에는 각 TPU 구성의 작업자당 총 컴퓨팅 리소스가 나와 있습니다.

TPU 유형	토폴로지	머신 유형	TPU 칩	vCPU	RAM(GB)
tpu-v5-lite-podslice	1x1	ct5lp-hightpu-1t	1	24	48
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	4	112	192
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	8	224	384
tpu-v6e-slice	1x1	ct6e-standard-1t	1	44	176
tpu-v6e-slice	2x2	ct6e-standard-4t	4	180	720
tpu-v6e-slice	2x4	ct6e-standard-8t	8	360	1440
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	4	208	448

TPU로 파이프라인 실행

TPU로 Dataflow 작업을 실행하려면 다음 명령어를 사용합니다.

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --dataflow_service_options "worker_accelerator=type:TPU_TYPE;topology:TPU_TOPOLOGY" \
  --worker_machine_type "MACHINE_TYPE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --sdk_container_image "IMAGE" \
  --number_of_worker_harness_threads NUMBER_OF_THREADS

다음을 바꿉니다.

• PIPELINE: 파이프라인 소스 코드 파일입니다.
• PROJECT: Google Cloud 프로젝트 이름
• BUCKET: Cloud Storage 버킷입니다.
• REGION: Dataflow 리전입니다(예: us-central1).
• TPU_TYPE: 지원되는 TPU 유형입니다(예: tpu-v5-lite-podslice). 유형 및 토폴로지의 전체 목록은 지원되는 TPU 가속기를 참고하세요.
• TPU_TOPOLOGY: TPU 토폴로지(예: 1x1)
• MACHINE_TYPE: 해당 머신 유형(예: ct5lp-hightpu-1t)
• DISK_SIZE_GB: 각 작업자 VM의 부팅 디스크 크기입니다(예: 100).
• IMAGE: Docker 이미지의 Artifact Registry 경로입니다.
• NUMBER_OF_THREADS: 선택사항. 작업자 하네스 스레드 수입니다.

Dataflow 작업 확인

작업에서 TPU가 있는 작업자 VM을 사용하는지 확인하려면 다음 단계를 따르세요.

1. Google Cloud 콘솔에서 Dataflow > 작업 페이지로 이동합니다.

   작업으로 이동

2. 작업을 선택합니다.

3. 작업 측정항목 탭을 클릭합니다.

4. 자동 확장 섹션에서 현재 작업자 VM이 하나 이상 있는지 확인합니다.

5. 측면 작업 정보 창에서 machine_type가 ct로 시작하는지 확인합니다. 예를 들면 ct6e-standard-1t입니다. TPU 사용량을 나타냅니다.

Dataflow 작업 문제 해결

TPU로 Dataflow 작업을 실행하는 데 문제가 발생하면 Dataflow TPU 작업 문제 해결을 참고하세요.

다음 단계

• Dataflow의 TPU 지원 자세히 알아보기
• Beam의 대규모 모델 추론에 대해 알아봅니다.