TPU 런타임
이 문서에서는 사용되는 프레임워크 및 TPU 아키텍처(VM 또는 노드)에 따라 Cloud TPU에 적합한 런타임 버전을 선택하는 방법을 안내합니다.
TPU VM 런타임 버전
이 섹션에서는 TPU VM 아키텍처를 갖춘 TPU에 사용해야 하는 TPU 소프트웨어 버전을 설명합니다. TPU 노드 아키텍처의 경우 TPU 노드 소프트웨어 버전을 참조하세요.
TPU 소프트웨어 버전은 TensorFlow, PyTorch, JAX 프레임워크에서 사용할 수 있습니다.
TensorFlow
모델이 작성된 TensorFlow 버전과 일치하는 TPU 소프트웨어 버전을 사용합니다.
TensorFlow 2.15.0부터는 스트림 실행자(SE) 런타임 또는 PJRT 런타임도 지정해야 합니다. 예를 들어 PJRT 런타임과 함께 TensorFlow 2.16.1을 사용하는 경우 tpu-vm-tf-2.16.1-pjrt TPU 소프트웨어 버전을 사용합니다. TensorFlow 2.15.0 이전 버전은 스트림 실행자만 지원합니다. PJRT에 대한 자세한 내용은 TensorFlow PJRT 지원을 참조하세요.
현재 지원되는 TensorFlow TPU VM 소프트웨어 버전은 다음과 같습니다.
- tpu-vm-tf-2.16.1-pjrt
- tpu-vm-tf-2.16.1-se
- tpu-vm-tf-2.15.1-pjrt
- tpu-vm-tf-2.15.1-se
- tpu-vm-tf-2.15.0-pjrt
- tpu-vm-tf-2.15.0-se
- tpu-vm-tf-2.14.1
- tpu-vm-tf-2.14.0
- tpu-vm-tf-2.13.1
- tpu-vm-tf-2.13.0
- tpu-vm-tf-2.12.1
- tpu-vm-tf-2.12.0
- tpu-vm-tf-2.11.1
- tpu-vm-tf-2.11.0
- tpu-vm-tf-2.10.1
- tpu-vm-tf-2.10.0
- tpu-vm-tf-2.9.3
- tpu-vm-tf-2.9.1
- tpu-vm-tf-2.8.4
- tpu-vm-tf-2.8.3
- tpu-vm-tf-2.8.0
- tpu-vm-tf-2.7.4
- tpu-vm-tf-2.7.3
TensorFlow 패치 버전에 대한 자세한 내용은 지원되는 TensorFlow 패치 버전을 참조하세요.
TensorFlow PJRT 지원
TensorFlow 2.15.0부터는 TPU에서 TensorFlow에 PJRT 인터페이스를 사용할 수 있습니다. PJRT는 자동화된 기기 메모리 조각 모음을 제공하고 하드웨어와 프레임워크의 통합을 간소화합니다. PJRT에 대한 자세한 내용은 Google 오픈소스 블로그의 PJRT: ML 하드웨어와 프레임워크의 통합 간소화를 참조하세요.
TPU v2, v3, v4의 모든 기능이 PJRT 런타임으로 마이그레이션된 것은 아닙니다. 다음 표에서는 PJRT 또는 스트림 실행자에서 지원되는 기능을 설명합니다.
| 가속기 | 기능 | 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 v5e는 TPU Embedding API를 지원하지 않음 | 해당 사항 없음 |
| TPU v5p | 고밀도 컴퓨팅(TPU 임베딩 API 없음) | 예 | No |
| TPU v5p | TPU 임베딩 API | 예 | No |
TensorFlow 버전 2.10.0 이하를 사용하는 TPU v4
TensorFlow를 사용하여 TPU v4에서 모델을 학습시키는 경우 TensorFlow 버전 2.10.0 이하에서는 다음 표에 표시된 v4 특정 버전을 사용하세요. 사용 중인 TensorFlow 버전이 테이블에 없는 경우 TensorFlow 섹션의 안내를 따르세요.
| TensorFlow 버전 | TPU 소프트웨어 버전 |
|---|---|
| 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 |
Libtpu 버전
TPU VM은 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.* | 1.2.0 |
| 2.7.3 | 1.1.2 |
PyTorch 및 JAX
PyTorch 및 JAX에 다음 공통 기본 이미지를 사용한 후 pip와 함께 필요한 프레임워크 버전을 설치합니다.
tpu-ubuntu2204-base(기본)v2-alpha-tpuv5(TPU v5p)
설치 안내는 PyTorch/XLA 및 JAX에 대한 빠른 시작 문서를 참조하세요.
TPU 노드 런타임 버전
이 섹션에서는 TPU 노드 아키텍처를 갖춘 TPU에 사용해야 하는 TPU 소프트웨어 버전을 설명합니다. TPU VM 아키텍처의 경우 TPU VM 소프트웨어 버전을 참조하세요.
TPU 소프트웨어 버전은 TensorFlow, PyTorch, JAX 프레임워크에서 사용할 수 있습니다.
TensorFlow
모델이 작성된 TensorFlow 버전과 일치하는 TPU 소프트웨어 버전을 사용합니다. 예를 들어 TensorFlow 2.12.0을 사용하는 경우 2.12.0 TPU 소프트웨어 버전을 사용합니다. TensorFlow 특정 TPU 소프트웨어 버전은 다음과 같습니다.
- 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.2
- 2.7.3
TensorFlow 패치 버전에 대한 자세한 내용은 지원되는 TensorFlow 패치 버전을 참조하세요.
TPU 노드를 만들면 최신 버전의 TensorFlow가 TPU 노드에 사전 설치됩니다.