시스템 아키텍처

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

단일 TPU 기기에는 4개의 칩이 있으며, 각 칩에는 2개의 TPU 코어가 포함되어 있습니다. TPU 코어에는 하나 이상의 행렬 곱셈 단위(MXU), 벡터 처리 단위(VPU), 스칼라 단위가 포함됩니다.

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

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

TPU 버전

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

TPU v2

TPU v2 보드에는 TPU 칩 4개와 HBM 16GiB가 포함됩니다. 각 TPU 칩에는 2개의 코어가 포함됩니다. 각 코어에는 MXU, 벡터 단위, 스칼라 단위가 있습니다.

이미지

TPU v3

TPU v3 보드에는 TPU 칩 4개와 HBM 32GiB가 포함됩니다. 각 TPU 칩에는 2개의 코어가 포함됩니다. 각 코어에는 MXU, 벡터 단위, 스칼라 단위가 있습니다.

이미지

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

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

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 구성

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 런타임은 Pod 또는 슬라이스의 각 TPU 코어에 데이터 배포를 자동으로 처리합니다. Pod 슬라이스는 32개, 128개, 512개, 1,024개 또는 2,048개 코어와 함께 사용할 수 있습니다.

Cloud TPU VM 아키텍처

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

이미지

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

이미지

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

TPU 노드

TPU 노드는 원래 TPU 환경입니다. TPU 노드는 gRPC를 통해 TPU 호스트와 통신하는 추가적인 사용자 VM이 필요합니다. TPU 호스트에 대한 직접 액세스는 없습니다.

이미지

TPU VM

TPU VM을 사용할 때는 TPU 기기에 물리적으로 연결되는 Google Compute Engine 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와 비슷하게 실행됩니다.

이미지

다음 단계