TPU Pod에서 학습

개요

TPU는 TPU Pod로 설계되었습니다. TPU Pod는 전용 고속 네트워크 인터페이스로 연결되는 TPU 기기 모음입니다. TPU Pod는 최대 2048개의 TPU 코어를 포함할 수 있으며 여러 TPU 간에 처리 부하를 분산할 수 있습니다. 각 TPU 보드는 데이터 로드 및 사전 처리와 같은 작업을 위해 고성능 CPU 기반 호스트 머신에 연결됩니다. 더 많은 수의 TPU를 최대한 활용하려면 여러 학습 태스크 매개변수를 조정해야 합니다. 이 문서에서는 몇 가지 일반적인 문제, 모델에서 변경해야 할 사항, 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 버킷을 사용해야 합니다.

데이터 스토리지에 NFS 사용

TPU Pod를 만들면 각 TPU 노드에 대해 개별 VM이 생성됩니다. 기본적으로 각 TPU VM에는 서로 다른 사용자 ID(UID)가 할당됩니다. 이로 인해 여러 노드에서 동일한 NFS 디렉터리에 액세스하려고 시도할 때 문제가 발생합니다. 동일한 디렉터리의 각 노드마다 소유자가 다르고 표준 Linux 권한이 노드 간에 적용되지 않습니다. 예를 들어 한 노드의 프로세스는 다른 노드에서 생성된 로그 디렉터리에 기록할 수 없습니다.

이 문제를 해결하려면 OS 로그인을 사용 설정할 수 있습니다. 특정 VM 인스턴스 또는 프로젝트에 대해 OS 로그인을 구성할 수 있습니다. 자세한 내용은 OS 로그인 설정을 참조하세요.

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

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

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

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