시스템 아키텍처

이 문서에서는 Cloud TPU 시스템의 하드웨어 및 소프트웨어 구성요소의 아키텍처를 설명합니다.

Tensor Processing Unit(TPU)은 Google에서 맞춤 개발한 ASIC(Application-Specific Integrated Circuit)로, 머신러닝(ML) 워크로드를 빠르게 처리하는 데 사용됩니다. TPU는 Google이 머신러닝 분야에서 쌓은 심화된 경험과 경쟁력을 바탕으로 설계되었습니다.

단일 TPU 보드에는 4개의 TPU 칩이 포함됩니다. 각 칩에는 2개의 TPU 코어가 포함됩니다. TPU 코어에서 사용할 수 있는 리소스는 버전에 따라 다릅니다. 각 TPU 코어에는 스칼라, 벡터, 행렬 곱셈 단위(MXU)가 있습니다. MXU는 TPU 칩의 대량의 컴퓨팅 성능을 제공합니다. 각 MXU는 낮은 bfloat16 정확성으로 각 주기에서 16,000개의 곱셈-누산 작업을 수행할 수 있습니다. bfloat16은 IEEE 절반의 정밀도 표현보다 높은 학습 및 모델 정확성을 제공하는 16비트 부동 소수점 표현입니다. TPU 보드의 각 코어는 사용자 계산을 독립적으로 실행할 수 있습니다. 높은 대역폭 상호 연결은 칩이 서로 직접 통신하도록 허용합니다. 더 자세한 기술 정보는 심층신경망 학습용 도메인별 슈퍼컴퓨터를 참조하세요.

TPU는 TPU Pod로 확장되도록 설계되었습니다. TPU Pod는 TPU 코어가 최대 2,048개까지 있을 수 있는 슈퍼컴퓨터로, 여러 TPU 보드에 처리 부하를 분산할 수 있습니다. TPU Pod 구성에서 전용 고속 네트워크 인터페이스는 여러 TPU 기기를 연결하여 머신러닝 워크로드를 위한 더 많은 수의 TPU 코어와 더 큰 TPU 메모리 풀을 제공합니다.

각 TPU 보드는 TPU에 피드하기 위한 데이터 로드 및 사전 처리와 같은 고성능 CPU 기반 호스트 머신에 연결되어 있습니다.

Cloud TPU

Cloud TPU는 Google Cloud Platform(GCP)을 통해 TPU에 대한 액세스 권한을 부여하는 서비스입니다. Cloud TPU를 사용하여 Google의 TPU 가속기 하드웨어에서 머신러닝 워크로드를 실행할 수 있습니다. Cloud TPU는 연구자, 개발자, 기업이 ML 워크로드를 학습시키는 데 도움이 되는 최대의 성능과 유연성을 발휘하도록 설계되었습니다.

Cloud TPU API를 사용하여 TPU의 TPU 관리를 자동화할 수 있습니다. 이렇게 하면 손쉽게 대규모 연산 클러스터로 확장하고 워크로드를 실행하고 워크로드가 완료된 후 클러스터를 다시 축소할 수 있습니다. 칩에 내장된 하드웨어 지원 덕분에 다양한 딥 러닝 워크로드에 걸쳐 선형적인 성능 확장이 실현됩니다.

TPU 구성

Cloud TPU는 다른 TPU 구성에 대한 액세스를 제공합니다.

  • TPU Pod
  • TPU 슬라이스
  • TPU 보드

TPU Pod

TPU Pod에서 TPU Pod의 모든 TPU 칩은 CPU 호스트를 통한 통신 지연을 건너뛰는 고속 상호 연결을 통해 서로 직접 연결됩니다. 칩은 2D 토러스에 연결되어 각 칩이 4개의 이웃과 직접 통신합니다. 이 아키텍처는 all-reduce와 같은 ML 워크로드의 공통 통신 패턴에 특히 고성능을 발휘합니다. TPU v2 및 v3의 아키텍처 세부정보 및 성능 특성은 심층신경망 학습용 도메인별 슈퍼컴퓨터에서 확인할 수 있습니다.

TPU 슬라이스

TPU Slice는 TPU Pod의 일부입니다. 전체 pod의 리소스가 필요하지 않으면 pod의 일부를 사용할 수 있습니다. Cloud TPU에서 Cloud TPU 가격 책정 섹션을 참조하시면 자세한 내용을 확인할 수 있습니다.

단일 TPU 보드

단일 보드 TPU 구성은 다른 TPU 보드에 대한 네트워크 연결이 없는 TPU 칩 4개(TPU 코어 8개)가 있는 독립형 보드입니다. 단일 보드 TPU는 TPU Pod 구성에 속하지 않으므로 TPU Pod의 한 부분을 점유하지 않습니다. TPU 유형 페이지에서 사용 가능한 단일 보드 TPU 구성을 확인하세요.

TPU 버전

TPU 버전은 각 TPU 코어의 아키텍처, 각 TPU 코어를 위한 고대역폭 메모리(HBM)의 크기, 각 TPU 보드의 코어 간 상호 연결, 기기 간 통신에 사용할 수 있는 네트워킹 인터페이스를 정의합니다. 사용 가능한 TPU 버전은 v1과 v2입니다.

TPU v2

TPU v2 보드에는 각각 2개의 코어가 있는 4개의 TPU 칩이 포함됩니다. 각 TPU 코어는 HBM 8GiB와 코어당 1개의 MXU가 있습니다. TPU v2 Pod에는 최대 512개의 TPU 코어와 4TiB의 메모리가 있습니다.

이미지

TPU v3

TPU v3 보드에는 각각 2개의 코어가 있는 4개의 TPU 칩이 포함됩니다. TPU 코어당 16GiB의 HBM과 각 코어에 2개의 MXU가 있습니다. TPU v3 Pod에는 최대 2,048개의 TPU 코어와 32TiB의 메모리가 포함됩니다.

이미지

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에서도 입력에 제약을 받을 수 있습니다. 파이프라인 성능 가이드가 인피드 문제를 해결하는 데 도움이 될 수 있습니다.

Cloud TPU VM 아키텍처

각 TPU 보드는 물리적으로 호스트 머신(TPU 호스트)에 연결됩니다.

이미지

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

이미지

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

TPU 노드

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

이미지

TPU VM

TPU VM을 사용하면 TPU 호스트에서 실행되는 Google Compute Engine VM에 직접 SSH를 통해 연결됩니다. 원하는 코드를 실행할 수 있도록 머신에 대한 루트 액세스 권한을 얻습니다. TPU 컴파일러 및 런타임에서 직접 디버그 로그 및 오류 메시지에 액세스할 수 있습니다. TPU VM은 TPU 노드에서 불가능한 새로운 사용 사례를 지원합니다. 예를 들어 입력 파이프라인에서 커스텀 작업을 실행하고 로컬 스토리지를 사용할 수 있습니다.

이미지

사용자 VM이 없으므로 코드와 TPU VM 사이에 네트워크, Virtual Private Cloud, 방화벽이 없어도 입력 파이프라인의 성능이 향상됩니다. 또한 TPU VM은 사용자 VM에 대한 비용을 지불할 필요가 없으므로 저렴한 비용으로 제공됩니다.

지원되는 프레임워크

JAX, PyTorch, TensorFlow와 같은 프레임워크는 모든 TPU VM에 있는 공유 라이브러리libtpu를 통해 TPU에 액세스합니다. 이 라이브러리에는 TPU 프로그램을 컴파일하는 데 사용되는 XLA 컴파일러, 컴파일된 프로그램을 실행하는 데 사용되는 TPU 런타임, 런타임에서 TPU에 액세스하기 위해 사용하는 TPU 드라이버가 포함되어 있습니다.

이미지

TensorFlow

TPU VM

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

이미지

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

TPU 노드

Cloud TPU 노드 시스템 아키텍처는 원래 TensorFlow 및 분산 프로그래밍 모델을 위해 빌드되었습니다. TPU 호스트는 사용자가 액세스할 수 없으며 TensorFlow 서버의 헤드리스 사본만 실행합니다. Python 또는 TensorFlow 그래프로 표현되지 않은 사용자 코드를 실행하지 않습니다. 사용자 코드는 네트워크를 통해 TPU 호스트와 통신하는 별도의 원격 VM에서 실행됩니다.

이미지

PyTorch

TPU VM

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

이미지

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

TPU 노드

PyTorch는 XRT라는 라이브러리를 사용하여 Cloud TPU 노드 아키텍처에서 실행되므로 TensorFlow gRPC 연결을 통해 XLA 그래프 및 런타임 안내를 전송하고 이를 TensorFlow 서버에서 실행할 수 있습니다. 사용자 VM은 각 TPU 호스트에 필요합니다.

이미지

JAX

TPU VM

TPU VM의 경우 TPU 호스트에서 직접 코드를 실행할 수 있으므로 사용자 VM이 필요하지 않습니다.

이미지

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

TPU 노드

Cloud TPU 노드의 JAX는 각 호스트 VM에 별도의 사용자 VM이 필요하다는 점에서 PyTorch와 유사합니다.

이미지

다음 단계