시스템 아키텍처

Tensor Processing Unit(TPU)은 Google에서 설계된 ML 가속기입니다. Cloud TPU는 TPU를 확장 가능한 GCP 클라우드 리소스로 사용할 수 있게 해줍니다. TensorFlow, Pytorch, JAX와 같은 머신러닝 프레임워크를 사용해서 Cloud TPU에서 머신러닝 워크로드를 실행할 수 있습니다.

MXU는 시스톨릭 배열에서 128 x 128의 곱셈/누산기로 구성됩니다. MXU는 TPU 칩의 대량의 연산 성능을 제공합니다. 각 MXU는 bfloat16 숫자 형식을 사용해서 각 주기당 16,000개의 곱셈-누산 작업을 수행할 수 있습니다.

VPU는 활성화, 소프트맥스 등의 일반 계산에 사용됩니다. 스칼라 단위는 제어 흐름, 메모리 주소 계산, 기타 유지보수 작업에 사용됩니다.

TPU 버전

TPU의 정확한 레이아웃은 사용되는 TPU 버전에 따라 달라집니다. TPU v2 및 v3의 아키텍처 세부정보 및 성능 특성은 심층신경망 학습용 도메인별 슈퍼 컴퓨터에서 확인할 수 있습니다.

TPU v4

가장 작은 TPU v4 구성에는 4개의 TPU 칩과 32GiB의 HBM 및 128MiB의 공유 공통 메모리가 포함됩니다. 각 TPU 칩에는 2개의 코어가 포함됩니다. 각 코어에는 4개의 MXU, 벡터 단위, 스칼라 단위가 있습니다. 다음 다이어그램은 TPU v4 칩을 보여줍니다.

이미지

TPU v3

가장 작은 TPU v3 구성에는 4개의 TPU 칩과 32GiB의 HBM이 포함됩니다. 각 TPU 칩에는 2개의 코어가 포함됩니다. 각 코어에는 2개의 MXU, 벡터 단위, 스칼라 단위가 있습니다. 다음 다이어그램은 TPU v3 칩을 보여줍니다.

이미지

TPU v2

가장 작은 TPU v2 구성에는 4개의 TPU 칩과 16GiB의 HBM이 포함됩니다. 각 TPU 칩에는 2개의 코어가 포함됩니다. 각 코어에는 MXU, 벡터 단위, 스칼라 단위가 있습니다. 다음 다이어그램은 TPU v2 칩을 보여줍니다.

이미지

Cloud TPU는 다음 TPU 구성을 제공합니다.

  • 단일 TPU 기기
  • TPU Pod - 고속 상호 연결로 연결된 TPU 기기 그룹
  • TPU 슬라이스 - TPU Pod의 하위 집합

TPU v4와 v3의 성능 이점 비교

메모리 시스템:

  • v4 TPU 칩은 전체 칩에서 통합된 3232GiB HBM 메모리 공간을 가지므로 두 개의 온칩 TPU 코어 간에 더 효과적으로 조정할 수 있습니다.
  • 최신 메모리 표준 및 속도를 사용하여 HBM 성능이 향상되었습니다.
  • 512B 단위로 고성능 스트라이딩을 기본적으로 지원하는 DMA 성능 프로필이 향상되었습니다.

TPU 코어

  • MXU 수의 두 배이고 클록 속도는 최대 TFLOPS 275개를 제공합니다.
  • 2배 전치 및 치환 대역폭
  • Common Memory(Cmem)의 부하 저장 메모리 액세스 모델
  • 더 빠른 MXU 가중치 로드 대역폭과 8비트 모드가 지원되어 배치 크기를 줄이고 추론 지연 시간을 개선할 수 있습니다.

칩 간 상호 연결:

네트워크 지름이 작은 네트워크 토폴로지를 사용 설정하기 위한 칩당 6개의 상호 연결 링크

기타:

  • 호스팅할 x16 PCIE gen3 인터페이스(직접 연결)입니다.
  • 보안 모델이 향상되었습니다.
  • 에너지 효율이 향상되었습니다.

TPU v3와 v2의 성능 이점 비교

TPU v3 구성의 늘어난 코어당 FLOPS와 메모리 용량은 다음과 같은 방식으로 모델의 성능을 개선할 수 있습니다.

  • TPU v3 구성은 연산 제약 모델의 코어당 성능 측면에서 상당한 이점을 제공합니다. TPU v2 구성의 메모리 제약 모델이 TPU v3 구성에서도 메모리의 제약을 받는 경우 이와 같은 수준의 성능 개선을 달성하지 못할 수 있습니다.

  • TPU v2 구성에서 데이터가 메모리에 들어가지 않는 경우 TPU v3으로 성능을 높이고 중간 값 재계산(재구체화)을 줄일 수 있습니다.

  • TPU v3 구성에서는 TPU v2 구성에서 맞지 않았던 배치 크기로 새 모델을 실행할 수 있습니다. 예를 들어 TPU v3은 더 깊은 ResNets 및 RetinaNet에서 더 큰 이미지를 허용할 수 있습니다.

학습 단계에서 입력을 기다리는 이유로 TPU v2에서 거의 입력에 제약을 받은('인피드') 모델은 Cloud TPU v3에서도 입력에 제약을 받을 수 있습니다. 파이프라인 성능 가이드가 인피드 문제를 해결하는 데 도움이 될 수 있습니다.

TPU v4 구성

TPU v4 포드는 재구성이 가능한 고속 링크로 상호 연결된 4,096개의 칩으로 구성되어 있습니다. TPU v4의 유연한 네트워킹을 사용하면 칩을 동일한 크기의 슬라이스로 다양한 방식으로 연결할 수 있습니다. AcceleratorConfig를 사용하여 v4 TPU를 구성합니다. AcceleratorConfig를 사용하면 토폴로지 필드를 통해 칩 단위로 슬라이스 크기를 지정할 수 있습니다.

TPU v4 토폴로지는 TPU 칩 간의 수, 형태, 상호 연결을 설명하는 3-튜플을 사용하여 지정됩니다. 다음 그림에서는 일반적인 TPU v4 토폴로지를 보여줍니다.

이미지

큰 슬라이스는 하나 이상의 4x4x4 '큐브'의 칩에서 빌드할 수 있습니다.

지정된 칩 수의 TPU 슬라이스는 다양한 방식으로 구성할 수 있습니다. 예를 들어 AcceleratorType이 v4-1024인 TPU 슬라이스는 이제 4x4x32, 4x8x16, 8x8x8로 구성될 수 있습니다. v4-4096의 TPU는 4x4x128, 4x8x64, 4x16x32, 8x16x16, 8x16x16_twisted와 같이 보다 높은 유연성을 제공합니다. 자세한 내용은 TPU 유형 및 토폴로지를 참조하세요.

TPU v2 및 v3 구성

TPU는 다음 구성으로 제공됩니다.

  1. 단일 TPU 보드
  2. TPU Pod
  3. TPU Pod 슬라이스

단일 TPU 보드

단일 보드 TPU 구성은 다른 TPU 보드에 대해 네트워크 연결 없이 4개의 TPU 칩(8개 TPU 코어)이 포함된 독립형 보드입니다. 단일 보드 TPU는 TPU Pod 구성의 일부가 아니며 TPU Pod의 일부를 차지하지 않습니다.

TPU Pod 및 슬라이스

TPU Pod 또는 TPU Pod 슬라이스에서 TPU 칩은 고속 상호 연결을 사용해서 연결되고, 각 TPU 칩은 TPU 기기에서 다른 칩과 직접 통신합니다. TPU 소프트웨어는 포드 또는 슬라이스에서 각 TPU 코어로 데이터 배포를 자동으로 처리합니다. 포드 슬라이스는 32, 128, 512, 1024, 2048 코어로 제공됩니다.

Cloud TPU VM 아키텍처

TPU는 행렬 작업만 수행할 수 있으므로, 각 TPU 보드가 CPU 기반 호스트 머신에 연결되어 TPU에서 실행할 수 없는 작업을 수행합니다. 호스트 머신은 Cloud Storage에서 데이터를 로드하고, 데이터를 사전 처리하고, TPU로 데이터를 전송합니다.

이미지

TPU Pod에는 각 TPU 보드에 대한 TPU 호스트가 있습니다.

이미지

TPU 호스트(및 TPU 보드)와 상호 작용하는 방법은 사용 중인 TPU VM 아키텍처(TPU 노드 또는 TPU VM)에 따라 달라집니다.

TPU 노드

TPU 노드 아키텍처는 gRPC를 통해 TPU 호스트와 통신하는 사용자 VM으로 구성됩니다. 이 아키텍처를 사용할 때는 TPU 호스트에 직접 액세스할 수 없습니다. 이로 인해 학습 및 TPU 오류를 디버깅하기가 어려워질 수 있습니다.

이미지

TPU VM

TPU VM 아키텍처를 사용하면 사용자 VM이 필요하지 않으며 TPU 기기에 물리적으로 연결된 VM에 직접 SSH를 통해 연결할 수 있습니다. 임의 코드를 실행할 수 있도록 VM에 대해 루트 액세스 권한을 갖습니다. 컴파일러 및 런타임 디버그 로그 및 오류 메시지에 액세스할 수 있습니다.

이미지

지원되는 프레임워크

JAX, PyTorch, TensorFlow는 모든 TPU VM에 제공되는 libtpu라는 공유 라이브러리를 통해 TPU에 액세스합니다. 이러한 라이브러리에는 TPU 프로그램을 컴파일하는 데 사용되는 XLA 컴파일러, 컴파일된 프로그램을 실행하는 데 사용되는 TPU 소프트웨어, TPU에 대한 낮은 수준의 액세스를 위해 런타임에 사용되는 TPU 드라이버가 포함되어 있습니다.

이미지

TensorFlow

TPU VM

사용자 VM에서 실행되는 Python 코드 대신 TPU VM을 사용하면 TPU 호스트에서 직접 실행될 수 있습니다.

이미지

TensorFlow 및 Cloud TPU에 대한 자세한 내용은 Cloud TPU에서 TensorFlow 모델 실행을 참조하세요.

TPU 노드

Cloud TPU 노드 시스템 아키텍처는 원래 TensorFlow용으로 빌드되었습니다. TPU 호스트는 사용자에게 액세스되지 않으며 TensorFlow 서버의 헤드리스 복사본을 실행합니다. 이러한 호스트는 TensorFlow 그래프로 표시되지 않는 Python 또는 모든 사용자 코드를 실행하지 않습니다. 사용자 코드는 gRPC 네트워크를 통해 TPU 호스트와 통신하는 별개의 사용자 VM에서 실행됩니다.

이미지

PyTorch

TPU VM

TPU VM을 사용하면 PyTorch 코드가 TPU 호스트에서 직접 실행됩니다.

이미지

PyTorch 및 Cloud TPU에 대한 자세한 내용은 Cloud TPU에서 PyTorch 모델 실행을 참조하세요.

TPU 노드

PyTorch는 XRT라는 라이브러리를 사용하여 Cloud TPU 노드 아키텍처에서 실행되며, TensorFlow gRPC 연결을 통한 XLA 그래프 및 런타임 안내 전송과 TensorFlow 서버에서의 실행을 허용합니다. 각 TPU 호스트에 대해 사용자 VM이 필요합니다.

이미지

JAX

TPU VM

TPU VM에서는 TPU 호스트에서 직접 코드를 실행할 수 있기 때문에 사용자 VM이 필요하지 않습니다.

이미지

Cloud TPU에서 JAX 실행에 대한 자세한 내용은 JAX 빠른 시작을 참조하세요.

TPU 노드

Cloud TPU 노드에서의 JAX는 각 호스트 VM에 개별 사용자 VM이 필요하다는 점에서 PyTorch와 비슷하게 실행됩니다.

이미지

다음 단계