Cloud TPU 소개
Tensor Processing Unit(TPU)은 Google에서 커스텀 개발한 ASIC(Application-Specific Integrated Circuit)로서 머신러닝 워크로드를 빠르게 처리하는 데 사용됩니다. TPU 하드웨어에 대한 자세한 내용은 TPU 아키텍처를 참고하세요. Cloud TPU는 Google Cloud에서 TPU를 확장 가능한 컴퓨팅 리소스로 제공하는 웹 서비스입니다.
TPU는 머신러닝 알고리즘에서 자주 쓰이는 대규모 행렬 연산을 수행하도록 설계된 하드웨어를 사용하여 모델을 보다 효율적으로 학습시킵니다. TPU는 온칩 고대역폭 메모리(HBM)를 통해 더 큰 모델과 배치 크기를 사용할 수 있습니다. TPU는 코드 변경 없이 워크로드를 확장하는 슬라이스라는 그룹으로 묶을 수 있습니다.
TPU에서 실행되는 코드는 가속기 선형 대수(XLA) 컴파일러로 컴파일해야 합니다. XLA는 ML 프레임워크 애플리케이션에서 내보낸 그래프를 가져와서 그래프의 선형 대수, 손실, 경사 구성요소를 TPU 기계어 코드로 컴파일하는 적시 컴파일러입니다. 나머지 프로그램은 TPU 호스트 머신에서 실행됩니다. XLA 컴파일러는 TPU 호스트 머신에서 실행되는 TPU VM 이미지의 일부입니다.
TPU 사용 시기
Cloud TPU는 특정 워크로드에 최적화되어 있습니다. 상황에 따라 Compute Engine 인스턴스에서 GPU 또는 CPU를 사용하여 머신러닝 워크로드를 실행해야 할 수 있습니다. 일반적으로 다음 가이드라인을 기반으로 워크로드에 가장 적합한 하드웨어를 결정할 수 있습니다
CPU
- 최대 유연성이 필요한 신속한 프로토타입 생성
- 학습에 오랜 시간이 걸리지 않는 단순한 모델
- 유효 배치 크기가 작은 소형 모델
- C++로 작성된 많은 커스텀 TensorFlow 작업이 포함된 모델
- 호스트 시스템의 사용 가능한 I/O 또는 네트워킹 대역폭으로 제한된 모델
GPU
- 최소한 부분적으로라도 CPU에서 실행해야 하는 커스텀 PyTorch/JAX 작업이 상당히 많은 모델
- Cloud TPU에서 사용할 수 없는 TensorFlow 작업으로 구성된 모델(사용 가능한 TensorFlow 작업 목록 참조)
- 유효 배치 크기가 비교적 큰 중대형 모델
TPU
- 행렬 연산이 주를 이루는 모델
- 기본 학습 루프 내에 커스텀 PyTorch/JAX 작업이 없는 모델
- 몇 주 또는 몇 달간 학습시키는 모델
- 유효 배치 크기가 큰 대형 모델
- 고급 순위 지정 및 맞춤 콘텐츠 워크로드에서 흔히 볼 수 있는 초대규모 임베딩이 있는 모델
Cloud TPU는 다음 워크로드에 적합하지 않습니다.
- 브랜치 생성이 자주 필요하거나 많은 요소별 대수 연산이 포함된 선형 대수 프로그램
- 고정밀 연산이 필요한 워크로드
- 기본 학습 루프에 커스텀 작업이 포함된 신경망 워크로드
Google Cloud의 TPU
Cloud TPU VM, Google Kubernetes Engine, Vertex AI를 통해 TPU를 사용할 수 있습니다. 다음 표에는 각 Google Cloud서비스의 리소스가 나와 있습니다.
| Google Cloud 서비스 | 리소스 |
|---|---|
| Cloud TPU | Cloud TPU VM 시작하기 |
| Google Kubernetes Engine | |
| Vertex AI |
모델 개발 권장사항
연산의 대부분이 비행렬 작업(예: 추가, 구조 변경, 연결)인 프로그램은 높은 MXU 사용률을 달성할 가능성이 별로 없습니다. 아래에는 Cloud TPU에 적합한 모델을 선택하고 빌드하는 데 도움이 되는 몇 가지 가이드라인이 나와 있습니다.
레이아웃
XLA 컴파일러는 행렬 곱셈을 더 작은 단위의 블록 타일로 나누는 작업을 비롯한 코드 변환을 수행하여 MXU(행렬 단위)에서 효과적으로 연산을 실행합니다. MXU 하드웨어의 구조인 128x128 시스톨릭 배열과 TPU 메모리 하위 시스템의 설계는 8배수의 차원을 선호하며 효과적인 타일 생성 작업을 위해 XLA 컴파일러에서 사용됩니다. 따라서 타일 생성에 보다 도움이 되는 특정한 레이아웃도 있지만, 타일을 생성하려면 먼저 형태 변경이 필요한 레이아웃도 있습니다. 형태 변경 오퍼레이션은 Cloud TPU에서 메모리의 제약을 받는 경우가 많습니다.
도형
XLA 컴파일러는 첫 번째 배치에서 적시에 ML 그래프를 컴파일합니다. 후속 배치에 다른 형태가 있으면 모델이 작동하지 않습니다. 형태가 바뀔 때마다 그래프를 다시 컴파일하는 작업이 오래 걸리기 때문입니다. 따라서 동적 형태의 텐서가 포함된 모델은 TPU에 적합하지 않습니다.
패딩
고성능 Cloud TPU 프로그램을 사용하면 조밀 연산을 128x128 조각 타일로 나눌 수 있습니다. 행렬 연산이 전체 MXU를 차지하지 못하는 경우 컴파일러는 0(영)으로 텐서를 패딩합니다. 패딩에는 두 가지 취약점이 있습니다.
- 0(영)으로 패딩된 텐서는 TPU 코어 사용률이 낮습니다.
- 패딩은 텐서에 필요한 온칩 메모리 저장용량을 늘려 최악의 경우 메모리 부족 오류로 이어질 수 있습니다.
패딩은 필요에 따라 XLA 컴파일러에서 자동으로 수행되지만 op_profile 도구의 수단으로 수행되는 패딩의 양을 결정할 수도 있습니다. TPU에 적절하게 맞는 텐서 차원을 선택하여 패딩을 피할 수 있습니다.
측정기준
적절한 텐서 크기를 선택하는 것만으로 TPU 하드웨어, 특히 MXU의 성능을 끌어올리는 데 큰 도움이 됩니다. XLA 컴파일러는 배치 크기 또는 특성 차원을 사용하여 MXU를 최대한 활용하려고 합니다. 따라서 이 중 하나는 128의 배수여야 합니다. 그렇지 않으면 컴파일러가 이 중 하나를 128로 패딩합니다. 이상적으로는 특징 차원과 배치 크기가 모두 8의 배수여야 메모리 하위 시스템에서 높은 성능을 끌어낼 수 있습니다.
Cloud TPU 시작하기
- 계정 설정 Google Cloud
- Cloud TPU API 활성화
- Cloud Storage 버킷에 대한 Cloud TPU 액세스 권한 부여
- TPU에서 기본 계산 실행
- TPU에서 참조 모델 학습
- 모델 분석
도움 요청
Cloud TPU 지원에 문의하세요. 활성 Google Cloud 프로젝트가 있는 경우 다음 정보를 제공할 준비를 합니다.
- Google Cloud 프로젝트 ID
- TPU 이름(있는 경우)
- 기타 제공 정보
다음 단계
Cloud TPU에 대해 더 알고 싶으신가요? 다음 리소스를 참고하시기 바랍니다.