TPU Pod에서 학습

개요

TPU는 TPU Pod로 확장되도록 설계되었습니다. TPU Pod는 전용 고속 네트워크 인터페이스로 연결되는 TPU 기기 모음입니다. TPU Pod를 사용하면 처리 부하를 TPU 여러 개에 분산할 수 있습니다. 각 TPU 보드는 데이터 로드 및 사전 처리와 같은 작업을 위해 고성능 CPU 기반 호스트 머신에 연결됩니다. 더 많은 수의 TPU를 최대한 활용하려면 여러 학습 태스크 매개변수를 조정해야 합니다.

TPU Pod를 사용한 학습 설정은 프레임워크마다 다릅니다. 다음 링크를 사용하면 각 프레임워크로 Pod에서 학습하는 방법에 대한 자세한 정보를 확인할 수 있습니다.

다음 섹션에서는 몇 가지 일반적인 문제, 모델에서 변경해야 하는 사항, 포드 오류를 줄이거나 방지하기 위한 권장사항을 설명합니다.

배치 크기 조정 및 학습 단계

더 큰 TPU 유형에서 선형적 확장을 달성하려면 코어당 배치 크기를 동일하게 유지하세요.

예를 들어 v2-8에서 배치 크기 1,024를 사용하는 경우 v2-32에서는 배치 크기 4,096(4*1,024)을 사용합니다. 이렇게 하면 TPU 하드웨어를 완전히 활용합니다. 더 작은 배치 크기를 사용할 수 있지만 그렇게 하면 학습이 선형적으로 확장되지 않습니다.

일부 모델에는 1단계가 데이터 단일 배치 처리에 해당하는 train_steps 플래그가 포함되어 있습니다. 배치 크기를 늘릴 때 총 학습 예시 수가 동일하게 유지되도록 학습 단계 수를 줄입니다.

예를 들어 100개 단계의 배치 크기가 1,000이면 학습 중 100,000개 예시가 처리됩니다. 작업자 수가 4개이고 유효 배치 크기가 4,000이면 동일하게 100,000개 예시를 처리하도록 단계 수를 25로 조절해야 합니다. 모델에서 epochs 플래그를 사용하면 단계 수를 확장할 필요가 없습니다.

배치 크기가 크면 모델의 수렴 동작이 변경될 수 있으므로 학습률과 같은 일부 초매개변수도 조정할 수 있습니다.

TPU Pod와 동일한 리전에서 리전별 Google Cloud Storage 버킷 사용

일반적으로 TPU 학습 시 권장사항은 항상 같은 리전의 리소스를 사용하는 것입니다. Google Cloud Storage 버킷과 TPU가 같은 리전에 있을 때 데이터 전송 속도가 높아지므로 리소스 리전은 특히 TPU Pod를 사용할 때 중요합니다.

학습 데이터 세트 및 체크포인트에 TPU와 동일한 리전의 리전 Google Cloud Storage 버킷을 사용해야 합니다.

TPU Pod에서 개발 워크플로 권장사항

새 TPU 워크로드를 개발할 때 가장 작은 TPU에서 개발을 시작하고 점진적으로 더 큰 TPU 크기로 반복하는 것이 적합한 경우가 많습니다. 작은 TPU 버전(예: v2-8 또는 v3-8)을 사용하여 시작합니다.

  • 기능을 위해 워크로드 테스트
  • 성능 도구를 사용하여 성능 테스트 및 유효성 검사

워크로드가 작동하고 성능 목표에 도달하면 v2-32 또는 v3-32와 같이 더 큰 TPU 유형으로 수직 확장합니다. 원하는 TPU 크기에 도달할 때까지 확장성(기능 및 성능)의 유효성을 검사하는 동시에 TPU 크기를 점진적으로 반복해서 늘립니다.