TPU VM 이미지
TPU 리소스를 만들 때 TPU VM 이미지를 지정하는 --version 또는 --runtime-version 파라미터를 전달합니다. TPU VM 이미지에는 운영체제(Ubuntu)와 TPU에서 코드를 실행하는 데 필요한 기타 소프트웨어(선택사항)가 포함되어 있습니다. 이 문서에서는 Cloud TPU를 만들 때 적합한 TPU VM 이미지를 선택하는 방법을 안내합니다.
PyTorch 및 JAX
PyTorch 및 JAX에 대해 다음과 같은 일반적인 TPU VM 기본 이미지를 사용한 후 사용하려는 프레임워크를 설치합니다.
- tpu-ubuntu2204-base(기본값)
- v2-alpha-tpuv5(TPU v5p)
설치 안내는 PyTorch/XLA 및 JAX의 빠른 시작 문서를 참조하세요.
TensorFlow
각 TensorFlow 버전과 관련된 TPU VM 이미지가 있습니다. Cloud TPU에서는 다음과 같은 TensorFlow 버전이 지원됩니다.
- 2.16.1
- 2.15.1
- 2.15.0
- 2.14.1
- 2.14.0
- 2.13.1
- 2.13.0
- 2.12.1
- 2.12.0
- 2.11.1
- 2.11.0
- 2.10.1
- 2.10.0
- 2.9.3
- 2.9.1
- 2.8.4
- 2.8.3
- 2.8.0
- 2.7.4
- 2.7.3
TensorFlow 패치 버전에 대한 자세한 내용은 지원되는 TensorFlow 패치 버전을 참조하세요.
TensorFlow 버전 2.15.0 이상의 경우 사용 중인 기기 API(PJRT 또는 스트림 실행자)를 기반으로 하는 TPU VM 이미지 변형이 있습니다.
v5p 및 v5e 기반 학습
TPU v5e 및 v5p는 TensorFlow 2.15.0 이상을 지원합니다. tpu-vm-tf-x.y.z-{pod}-pjrt 형식으로 TPU VM 이미지를 지정합니다. 여기서 x는 TensorFlow 주 버전이고, y는 부 버전이고, z는 TensorFlow 패치 버전입니다. 멀티 호스트 TPU를 사용하는 경우 TensorFlow 버전 다음에 pod를 추가합니다. 예를 들어 멀티 호스트 TPU에서 TensorFlow 2.16.0을 사용하는 경우 tpu-vm-tf-2.16.0-pod-pjrt TPU VM 이미지를 사용합니다. 다른 TensorFlow 버전의 경우 2.16.0을 사용 중인 TensorFlow 주 버전 및 패치 버전으로 바꿉니다. 단일 호스트 TPU를 사용하는 경우 pod를 생략합니다.
v5e에서 서빙
TensorFlow, PyTorch, JAX로 서빙하는 데 필요한 모든 소프트웨어 요구사항이 포함된 서빙 Docker 이미지가 있습니다. 자세한 내용은 Cloud TPU v5e 추론 소개를 참조하세요.
TPU v4
TPU v4 및 TensorFlow 2.15.0 이상을 사용하는 경우 v5p 및 v5e 기반 학습 안내를 따르세요. TensorFlow 2.10.0 이하를 사용하는 경우 v4 관련 TPU VM 이미지를 사용합니다.
| TensorFlow 버전 | TPU VM 이미지 버전 |
|---|---|
| 2.10.0 | tpu-vm-tf-2.10.0-v4 tpu-vm-tf-2.10.0-pod-v4 |
| 2.9.3 | tpu-vm-tf-2.9.3-v4 tpu-vm-tf-2.9.3-pod-v4 |
| 2.9.2 | tpu-vm-tf-2.9.2-v4 tpu-vm-tf-2.9.2-pod-v4 |
| 2.9.1 | tpu-vm-tf-2.9.1-v4 tpu-vm-tf-2.9.1-pod-v4 |
TPU v2 및 v3
TPU v2 또는 v3를 사용 중이면 사용 중인 TensorFlow 버전과 일치하는 TPU VM 이미지를 사용합니다. 예를 들어 TensorFlow 2.14.1을 사용 중이면 tpu-vm-tf-2.14.1 TPU 이미지를 사용합니다. 다른 TensorFlow 버전의 경우 2.14.1을 사용 중인 TensorFlow 버전으로 바꿉니다. 멀티 호스트 TPU를 사용 중이면 TPU 이미지 끝에 포드를 추가합니다. 예를 들면 tpu-vm-tf-2.14.1-pod입니다.
TensorFlow 2.15.0부터는 기기 API도 지정해야 합니다.
예를 들어 PJRT API와 함께 TensorFlow 2.16.1을 사용 중이면 TPU 이미지 tpu-vm-tf-2.16.1-pjrt를 사용합니다. 동일한 TensorFlow 버전과 함께 스트림 실행자 API를 사용 중이면 tpu-vm-tf-2.16.1-se TPU 이미지를 사용합니다. 2.15.0 이전의 TensorFlow 버전만 스트림 실행자를 지원합니다.
TensorFlow PJRT 지원
TensorFlow 2.15.0부터는 TPU에서 TensorFlow에 대한 PJRT 인터페이스를 사용할 수 있습니다. PJRT는 자동화된 기기 메모리 조각 모음을 제공하고 하드웨어와 프레임워크의 통합을 간소화합니다. PJRT에 대한 자세한 내용은 PJRT: ML 하드웨어와 프레임워크의 통합 간소화를 참조하세요.
| 가속기 | 특성 | PJRT 지원 | 스트림 실행자 지원 |
|---|---|---|---|
| TPU v2 - v4 | 고밀도 컴퓨팅(TPU 임베딩 API 없음) | 예 | 예 |
| TPU v2 - v4 | 고밀도 컴퓨팅 API + TPU 임베딩 API | 아니요 | 예 |
| TPU v2 - v4 | 소프트웨어 기기 배치가 포함된 tf.summary/tf.print | 아니요 | 예 |
| TPU v5e | 고밀도 컴퓨팅(TPU 임베딩 API 없음) | 예 | 아니요 |
| TPU v5e | TPU 임베딩 API | 해당 사항 없음 | 아니요 |
| TPU v5p | 고밀도 컴퓨팅(TPU 임베딩 API 없음) | 예 | 아니요 |
| TPU v5p | TPU 임베딩 API | 예 | 아니요 |
Libtpu 버전
TPU VM TensorFlow 이미지에는 특정 TensorFlow 버전과 해당 libtpu 라이브러리가 포함됩니다. 자체 VM 이미지를 만드는 경우 다음 TensorFlow TPU 소프트웨어 버전과 해당 libtpu 버전을 사용합니다.
| TensorFlow 버전 | libtpu.so 버전 |
|---|---|
| 2.16.1 | 1.10.1 |
| 2.15.1 | 1.9.0 |
| 2.15.0 | 1.9.0 |
| 2.14.1 | 1.8.1 |
| 2.14.0 | 1.8.0 |
| 2.13.1 | 1.7.1 |
| 2.13.0 | 1.7.0 |
| 2.12.1 | 1.6.1 |
| 2.12.0 | 1.6.0 |
| 2.11.1 | 1.5.1 |
| 2.11.0 | 1.5.0 |
| 2.10.1 | 1.4.1 |
| 2.10.0 | 1.4.0 |
| 2.9.3 | 1.3.2 |
| 2.9.1 | 1.3.0 |
| 2.8.3 | 1.2.3 |
| 2.8.0 | 1.2.0 |
| 2.7.3 | 1.1.2 |