TensorFlow 문제 해결 - TPU

이 가이드에서는 FAQ와 함께 Cloud TPU에서 TensorFlow 모델을 학습하는 사용자를 위한 문제 해결 도움말을 제공합니다. Pytorch 또는 JAX 학습의 문제를 해결하는 경우 해당 프레임워크의 문제 해결 문서를 참조하세요.

Cloud TPU 사용 방법에 대한 일반적인 가이드는 다음을 참조하세요.

개요

Cloud TPU에서 발생하는 일반적인 문제는 다음과 같은 카테고리로 분류됩니다.

  1. TPU 연결 문제

  2. 일반적인 오류 디버깅

  3. 메모리 사용량 감소

  4. 학습 속도 개선

  5. 모델 정확성 저하 문제 디버깅

TPU 서버 연결 문제

이 섹션에서는 TPU에 연결할 때 TensorFlow가 응답을 중지하거나 오류를 출력하는 상황에서 문제 해결 방법을 설명합니다. 모델이 큰 경우에는 TPU 그래프 컴파일 단계에 시간이 오래 걸릴 수 있으므로 응답을 중단했다고 결론을 내리기 전에 스크립트를 5분 이상 실행해보세요.

첫 번째 단계는 서버 자체의 문제인지, TensorFlow 학습 파이프라인의 문제인지 확인하는 것입니다. 이 작업을 수행하려면 TPU 서버 URL을 사용하여 MNIST 튜토리얼의 작업을 실행하고 올바르게 작동하는지 확인하세요. MNIST 튜토리얼의 작업을 실행할 때에도 연결 문제가 발생하면 TPU 서버 문제인 것입니다. 이 경우 다음 단계를 수행하세요.

  1. 다음 명령어를 실행하여 사용 가능한 TPU를 나열합니다. zone 및 project-id를 영역 및 프로젝트 ID로 바꿉니다.

    (vm)$ gcloud compute tpus list --zone zone --project project-id

    이렇게 하면 다음과 같은 출력이 표시됩니다. 

    NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
demo-tpu   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0  READY

  2. --tpu(위의 예시에서는 demo-tpu)에 올바른 값을 전달하는지, 그리고 이 TPU가 READY로 표시되는지 확인합니다.

  3. TPU가 READY로 표시되지 않거나 연결에 계속 문제가 발생하는 경우 다음을 사용하여 서버를 수동으로 다시 시작하세요.

    (vm)$ gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME

    위의 예시에서 $TPU_SERVER_NAME은 demo-tpu입니다. 완료하는 데 몇 분 정도 걸릴 수 있습니다.

  4. 위의 ... tpus list 명령어를 다시 실행하고 TPU가 READY 상태가 되기를 기다립니다. 이 작업은 몇 분 정도 걸릴 수 있습니다.

  5. MNIST 튜토리얼을 다시 실행해 봅니다.

  6. MNIST 튜토리얼 실행 문제가 지속되는 경우 도움 받기에 설명된 방법 중 하나를 통해 도움을 요청하세요.

MNIST 예시가 올바르게 실행되지만 모델이 여전히 응답하지 않으면 학습 파이프라인에 문제가 있을 가능성이 높습니다. 이를 디버깅하려면 코드에서 TPUStrategy를 기본 전략으로 바꾸는 것으로 시작합니다. 기본 전략을 사용할 경우 strategy.scope() 또는 strategy.run()을 사용할 때마다 모델이 TPU 대신 CPU(또는 GPU)에서 실행됩니다. 모델이 TPU가 아닌 CPU에서 실행될 경우는 TPU 관련 문제가 있는 것입니다. 그래도 실행되지 않을 경우에는 CPU에서 문제를 디버깅하는 것이 좋습니다.

학습 중에 ssh 연결 끊어짐

특히 Cloud Shell을 사용 중일 때는 장기 실행 학습 중 Cloud TPU에 대한 ssh 연결이 시간 초과될 수 있습니다. 이 시점에는 TPU Console 출력이 없고 TPU 학습이 중지된 것처럼 보일 수 있습니다. 이를 방지하려면 터미널 멀티플렉서 또는 tmux 또는 screen과 같은 세션 관리 도구를 사용해서 학습 세션을 실행합니다. 이렇게 하면 학습 길이에 관계없이 ssh 연결이 활성 상태로 유지됩니다.

일반적인 오류 디버깅

TPU를 만들 수 없음

Cloud TPU를 만들 때 다음 오류가 표시될 수 있습니다.

googleapiclient.errors.HttpError: < HttpError 403 when requesting https://content-tpu.googleapis.com/v1/projects/{PROJECT}/locations/{ZONE}/nodes/{TPU_NAME}?alt=json returned "Request had insufficient authentication scopes."

이것은 권한 문제이며 다음 명령어를 실행하여 해결할 수 있습니다.

gcloud auth login --update-adc

이 명령어는 애플리케이션 기본 사용자 인증 정보(ADC)를 업데이트하고 문제를 해결합니다. 자세한 내용은 gcloud auth login을 참조하세요.

로컬 파일 시스템을 사용할 수 없는 경우

오류 메시지

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

영향을 받는 프레임워크 및 구성

이 메시지는 TPU 노드 아키텍처를 사용하여 TensorFlow로 학습할 때 발생할 수 있습니다.

세부정보

모든 입력 파일과 모델 디렉터리는 Cloud Storage 버킷 경로(gs://bucket-name/...)를 사용해야 하며 이 버킷은 TPU 서버에서 액세스할 수 있어야 합니다. 모든 데이터 처리 및 모델 체크포인트는 로컬 머신이 아닌 TPU 서버에서 수행됩니다. TPU에서 사용하기 위해 Cloud Storage를 적절하게 구성하는 방법은 Cloud Storage 버킷에 연결 가이드를 참조하세요.

지원되지 않는 데이터 유형

오류 메시지

TypeError: DataType is not a supported TPU infeed type.

영향을 받는 프레임워크 및 구성

이 메시지는 TPU 노드 아키텍처를 사용하여 TensorFlow로 학습할 때 발생할 수 있습니다.

세부정보

현재 TPU에서는 tf.float32, tf.int32, tf.bfloat16, tf.bool 데이터 유형만 지원됩니다. tf.uint8, tf.string, tf.int64와 같은 다른 일반적인 데이터 유형은 데이터 사전 처리 중에 지원되는 데이터 유형 중 하나로 변환되어야 합니다(즉, tf.data.Dataset 파이프라인).

MNIST 학습에 사용된 decode_image 함수에서 변환 예시를 참조하세요.

지원되지 않는 동적 형태

오류 메시지

ValueError: shape [Shape] must have a fixed size for dimension
d that is known at graph construction time.

영향을 받는 프레임워크 및 구성

이 메시지는 TensorFlow를 사용하는 XLA 컴파일 중에만 발생합니다.

세부정보

TPU에서 모델을 실행하기 위해 TensorFlow는 XLA 컴파일러를 사용하여 모델을 컴파일합니다. 이 컴파일 단계는 학습 속도와 메모리 사용량을 현저하게 개선하지만, 그래프의 모든 텐서 형태(차원 크기)가 그래프 컴파일 시점에 알려져 있어야 합니다. 컴파일 시점에 형태가 결정되지 못한 경우에는 위와 같은 오류와 함께 TPU 컴파일이 실패합니다.

스트림에 남아있는 샘플 수가 배치 크기보다 작을 수 있으므로 동적 형태를 반환하는 일반적인 작업은 dataset.batch(batch_size)입니다. 따라서 TPU에서 학습할 때는 dataset.batch에 대해 drop remainder=True를 설정합니다. 파일에서 잠재적으로 마지막 몇 개의 샘플이 제외되어 모든 배치가 batch_size의 정적 형태를 갖추게 됩니다. 예를 들면 다음과 같습니다.

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)

사용 불가능한 TensorFlow 작업

오류 메시지

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT
devices compatible with node

영향을 받는 프레임워크 및 구성

이 메시지는 TensorFlow를 사용하여 학습할 때 발생할 수 있습니다.

세부정보

모델이 현재 TPU에서 제공되지 않는 텐서플로우 오퍼레이션을 사용합니다.

이후 지원 계획 및 해결 방법 제안과 함께 TPU에서 사용 가능한 오퍼레이션 목록을 살펴보려면 사용 가능한 텐서플로우 오퍼레이션 가이드를 참조하세요.

메모리 부족 오류 메시지

오류 메시지

ResourceExhaustedError: Ran out of memory in memory space hbm; used:
YYY; limit: 7.48G.

영향을 받는 프레임워크 및 구성

이 메시지는 TensorFlow, PyTorch, JAX로 학습할 때 발생할 수 있습니다.

세부정보

각 Cloud TPU는 8개의 TPU 코어로 구성되어 있고, v2 TPU에는 8GB, v3 TPU에는 16GB의 RAM(또는 HBM: 높은 대역폭 메모리)이 있습니다. 이 메모리는 기울기 계산에 필요한 중간 결과 텐서와 가중치(변수) 텐서를 저장하는 데 사용됩니다. TPU RAM에 비해 모델이 너무 큰 경우에는 초기화가 실패하고 위의 오류 메시지가 출력됩니다. 자세한 도움말은 메모리 사용량 감소 섹션을 참조하세요.

메모리 사용량을 줄이기 위한 팁:

실행 중단 문제

TPU 실행 도중 TensorFlow에 오류가 발생하는 경우, 스크립트가 종료되어 셸로 돌아가지 않고 중지되는 것처럼 보이는 경우가 있습니다. 이 경우 키보드에서 CTRL+\를 눌러 SIGQUIT을 트리거하면 즉시 Python이 종료됩니다.

마찬가지로, TPU 실행 중에 CTRL+C를 누르면 TensorFlow가 즉시 종료되지는 않지만 대신 현재 반복 루프의 끝까지 기다렸다가 깔끔하게 종료됩니다.

이런 식으로 종료한 후 TPU에 다시 연결할 때 새 오류가 발생하면 다음 명령어를 사용하여 TPU 서버를 수동으로 재설정하세요.

gcloud compute tpus stop tpu-name --zone=zone
gcloud compute tpus start tpu-name --zone=zone

여기서 tpu-name은 gcloud compute tpus list 명령어로 표시된 첫 번째 열에서 가져오고 zone은 두 번째 열에 표시된 영역입니다.

텐서 패딩이 과도한 경우

메모리 문제의 가능한 원인

TPU 메모리의 텐서에는 패딩이 적용되어 있습니다. 즉, TPU는 메모리에 저장된 텐서의 크기를 반올림하여 연산을 보다 효율적으로 수행합니다. 이 패딩은 하드웨어 수준에서 투명하게 발생하고 결과에 영향을 미치지 않습니다. 그러나 경우에 따라 패딩으로 인해 메모리 사용량 및 실행 시간이 현저하게 늘어날 수도 있습니다.

메모리 사용량을 줄이는 방법

TPU 소프트웨어는 메모리에 텐서를 배치하여 연산 효율성을 최대화하고 패딩을 최소화하려고 시도합니다. 이 메모리 레이아웃 프로세스는 복잡하지만 최적의 결과를 위해서는 모델이 다음과 같은 일반적인 원칙을 따라야 합니다. 메모리 오버헤드를 최소화하고 연산 효율성을 극대화하려면 다음 조건 중 하나를 충족해야 합니다.

  • 총 배치 크기는 64의 배수(TPU 코어당 8)이고 특징 차원은 128의 배수여야 합니다.

    또는

  • 총 배치 크기는 1,024의 배수(TPU 코어당 128)이고 특징 차원은 8의 배수여야 합니다.

1,024의 배치 크기 및 128의 배수인 특징 차원을 사용하면 최고의 효율성을 얻을 수 있지만, 모든 모델에서 가능하지는 않습니다. 명확하게 말해서 '특징 차원'은 완전 연결 레이어의 숨겨진 크기 또는 컨볼루션의 출력 채널 수를 의미합니다. 모든 레이어가 이 규칙을 따르는 것은 아니지만 특히 네트워크의 첫 번째와 마지막 레이어가 여기에 해당됩니다. 이는 정상적인 현상이며 대부분의 모델에는 일정 양의 패딩이 필요합니다.

메모리 사용량 감소

TPU에서 모델을 실행할 때 메모리 부족 오류가 발생하면 모델의 메모리 사용량을 줄이는 단계를 밟아야 합니다.

메모리 사용량을 줄이는 가장 효과적인 방법은 다음과 같습니다.

  • 과도한 텐서 패딩 줄이기
  • 배치 크기 줄이기

배치 크기 또는 모델이 너무 큼

메모리 문제의 가능한 원인

CPU, GPU, TPU에서 신경망을 학습시킬 때 다음과 같은 두 작업에서 메모리 사용이 발생합니다.

  1. 메모리 사용량은 모델의 가중치 수에 비례합니다.
  2. 역방향 전달을 계산하는 데 필요한 정방향 전달의 중간 활성화 함수 저장. 메모리 사용량은 배치 크기, 레이어 크기, 레이어 수에 직접 비례합니다.

따라서 모델에서 필요한 메모리는 주로 배치 크기에 좌우됩니다.

모델에 필요한 메모리는 네트워크에 있는 레이어 수에 따라 달라집니다.

TPU 런타임은 모델을 메모리에 맞추기 위해 연산자를 최적화하려고 시도하지만(경사 체크포인트와 유사하게 재구체화라고 함). 항상 이렇게 할 수 있는 것은 아닙니다.

메모리 사용량을 줄이는 방법

배치 크기를 메모리에 맞을 때까지 천천히 줄이면서 총 배치 크기가 64의 배수인지 확인하세요(코어별 배치 크기는 8의 배수여야 함). 배치 크기가 클수록 TPU에서 더 효율적입니다. 총 배치 크기가 1,024(코어당 128)이면 일반적으로 좋은 출발점입니다.

소형 크기의 배치(예: 64)로도 TPU에서 모델을 실행할 수 없다면 레이어 수 또는 레이어 크기를 줄여보세요.

학습 속도 개선

이 섹션에서는 TPU에서 모델을 성공적으로 실행할 수 있지만 학습 속도가 예상보다 느린 경우 속도를 높일 수 있는 방법을 설명합니다. 학습 성능 향상 방법에 대한 기타 권장사항은 성능 가이드를 참조하세요.

학습 루프당 실행별 단계가 너무 적음

성능 문제 설명

인수 steps_per_execution을 Model.compile에 전달하면 호스트 콜백 간에 실행되는 학습 단계 수를 제어합니다. 각 호스트 콜백은 TPU 서버의 호스트 CPU와 TPU 기기 간에 상당한 통신을 필요로 하므로 steps_per_execution이 너무 작으면 학습 속도가 느려질 수 있습니다.

모델에서 이 문제를 확인하는 방법

TPU 프로필이 TPU 기기 단계 사이에 자주 발생하는 호스트 CPU 콜백을 나타내는 경우 학습 시 더 큰 steps_per_execution 값을 활용할 수 있습니다.

해결 방법

steps_per_execution를 더 큰 값으로 설정합니다. steps_per_execution을 큰 값으로 설정할 수 있지만 메시지 로깅 및 체크포인트 저장은 지정된 수의 단계가 실행된 후에만 발생할 수 있다는 점에 유의하세요.

입력 처리 병목 현상

성능 문제 설명

TPU가 특정한 데이터 청크를 기반으로 학습하는 동안 입력 처리 함수는 CPU에서 다음 데이터 청크를 준비합니다. 입력 함수가 모델 함수보다 시간이 더 오래 걸리면 입력 함수가 데이터를 검색하는 동안 TPU가 유휴 상태로 남습니다.

모델에서 이 문제를 확인하는 방법

Cloud TPU 도구: 입력 파이프라인 분석기의 안내에 따라 텐서보드에서 입력 파이프라인 분석을 살펴보세요.

이미지

입력 파이프라인 분석 페이지에는 입력 처리로 인해 모델에 병목 현상이 발생하는지를 한눈에 볼 수 있는 요약이 표시됩니다. 동일한 페이지에 오퍼레이션별 실행 시간이 표시되므로 문제가 있는 오퍼레이션을 정확히 찾아낼 수 있습니다.

해결 방법

Dataset API를 사용하여 데이터를 로드할 때 다음과 같은 몇 가지 완화 방법을 사용할 수 있습니다.

  1. 데이터를 TFRecord 파일에 tf.train.Example 구조의 모음으로 저장하고 TFRecordDataset를 사용하여 로드합니다. Dataset API 튜토리얼 또는 ResNet 튜토리얼에서 예시를 참조하세요.
  2. dataset.cache() 및/또는 dataset.prefetch()를 사용하여 입력 데이터를 버퍼링합니다. 이렇게 하면 병목 현상으로 인해 파일 액세스 속도가 산발적으로 느려지는 문제가 방지됩니다.
  3. dataset.map() 함수의 num_parallel_calls 매개변수를 지정하여 멀티스레드 map() 작업을 사용 설정합니다. num_parallel_calls 값에 대한 간단한 휴리스틱은 사용 가능한 CPU 코어 수를 사용하는 것입니다.
  4. 모든 학습의 매 세대마다 비용이 발생하도록 하지 말고 오프라인에서 1회 비용으로 값비싼 데이터 사전 처리를 수행하세요.

모든 입력 처리는 로컬 머신이 아니라 TPU 서버에 있는 CPU에서 수행되며, 따라서 로컬 머신의 속도는 문제 요인이 아닙니다.

느린 단계 시간 및 낮은 MXU 사용률

성능 문제 설명

Cloud TPU는 매우 높은 속도로 행렬 곱셈과 컨볼루션을 수행할 수 있습니다. 대부분의 다른 TensorFlow 작업은 TPU에서 효율적으로 구현되지만, 이러한 작업은 다른 하드웨어에 비해 TPU의 주요 장점이 아닙니다. 따라서 TPU를 완전히 활용하려면 모델의 대부분이 행렬 곱셈이나 컨볼루션으로 이루어져 있어야 합니다.

모델에서 이 문제를 확인하는 방법

이 경우 표시되는 단계는 성능 프로파일링 시 표시되는 낮은 MXU 사용률과 결합된 느린 단계 시간입니다.

해결 방법

행렬 곱셈이 아닌 작업 수를 줄여 보세요. 행렬 곱셈 수를 줄인 후 TPU 성능이 적합한지 벤치마크를 다시 확인합니다.

텐서 패딩이 과도한 경우

성능 문제 설명

TPU는 계산 단위를 효율적으로 사용하기 위해 메모리의 텐서에 패딩을 적용합니다. 패딩은 메모리와 메모리 대역폭의 사용량을 늘릴 수 있습니다. 텐서 패딩 문제의 이해 및 해결을 위한 도움말은 텐서 패딩 섹션을 참조하세요.

느린 처리량 및 낮은 메모리 사용량

성능 문제 설명

일반적으로 더 큰 배치를 사용할수록 TPU의 학습 속도(초당 샘플 수)가 빨라집니다.

모델에서 이 문제를 확인하는 방법

TPU가 64 크기에 도달할 때까지 텐서에 패드를 적용하므로 모든 모델의 배치 크기는 항상 64 이상(TPU 코어당 8)이어야 합니다. TPU 학습 시 최적의 배치 크기는 메모리 전송 및 패딩과 관련된 비효율성을 제거할 수 있는 1,024(TPU 코어당 128)입니다.

해결 방법

권장사항은 메모리에 적합하고 64의 배수인 최대 배치 크기를 사용하는 것입니다. 이렇게 하는 가장 쉬운 방법은 1,024부터 시작하는 것입니다. 메모리 부족 오류가 발생하는 경우 모델이 성공적으로 실행될 때까지 배치 크기를 줄여보세요. 모델의 배치 크기를 변경하는 과정에서 동일한 모델 정확성을 유지하기 위해 학습률과 같은 기타 초매개변수를 조정해야 하는 경우도 있으나 상황에 따라 판단해야 합니다.

레이어 크기가 너무 작은 경우

성능 문제 설명

모델의 대부분을 행렬 곱셈이나 컨볼루션이 차지하는 경우에도 입력 텐서가 작으면 TPU가 최대 효율로 실행되지 않을 수 있습니다. 다른 하드웨어와 비교하자면 배치 크기와 레이어 크기가 모두 클 때(예: 차원이 512 이상인 경우) TPU가 가장 효율적으로 실행됩니다.

모델에서 이 문제를 확인하는 방법

일반적으로 레이어 크기가 128보다 작으면 TPU의 효율성이 떨어집니다. 128이 TPU 행렬 곱셈 단위의 기본 차원이기 때문입니다. 완전 연결 레이어의 경우 높은 효율성을 달성하려면 숨겨진 크기가 최소 512인 것이 좋습니다. 동일한 효율성 수준을 달성하기 위해서는 컨볼루션 레이어가 일반적으로 완전 연결형 레이어만큼 클 필요가 없습니다.

해결 방법

모델에 소형 레이어를 사용하는 주요 동기가 학습 속도 때문이라면 TPU에 더 큰 레이어를 사용하는 모델을 다시 벤치마킹하는 것이 좋습니다. 예를 들어 레이어의 출력 크기를 256에서 512로 늘리면 모델의 연산 속도가 2배더라도 늘어나는 학습 시간은 20%에 불과할 수 있습니다.

작업 수준 모델 프로파일링

성능 병목 현상을 식별하기 위해 작업 수준의 실행 시간과 메모리 사용량을 측정하면 유용한 경우가 많습니다. 이렇게 하는 방법에 대한 안내는
Cloud TPU 도구: Trace 뷰어 가이드를 참조하세요.

모델 정확성 저하 문제 디버깅

Cloud TPU 생태계의 목표 중 하나는 CPU 또는 GPU에서 현재 학습되고 있는 모든 모델을 TPU에서 학습시켰을 때 배치 크기 및 학습률과 같은 초매개변수를 약간 조정하는 것만으로 근접한 수준의 정확성을 실현하는 것입니다. 그러나 경우에 따라 TPU에서 모델을 학습시킬 때 정확성이 저하될 수 있습니다. 이러한 문제의 디버깅은 신경망 학습의 무작위적인 속성 때문에 매우 어려울 수 있습니다. 이 섹션에서는 모델을 TPU로 포팅할 때 모델 정확성이 저하되는 근본 원인을 정확히 찾아내는 방법을 제공합니다.

데이터 샤딩 이해(데이터 동시 로드)

TensorFlow의 주요 목표 중 하나는 각 작업이 CPU, GPU, TPU 중 어디에서 실행되든 거의 일치하는 결과를 생성하도록 만드는 것입니다. 단 여기에는 무작위 순서 지정 작업과 같은 예외가 있습니다. 일반적으로 TPU와 CPU에서 임의 이외의 작업 출력을 실행할 때 서로 차이가 크다면 버그로 신고하세요.

그러나 학습 파이프라인 전체로 볼 때는 CPU/GPU 및 TPU 중 어디에서 학습하는지에 따라 큰 차이가 있습니다. TPU에서 학습하는 경우 TensorFlow는 데이터 샤딩을 수행합니다. 각 Cloud TPU에는 독립적인 처리 단위로 작동하는 8개의 TPU 코어가 있습니다. 학습의 각 단계마다 각 TPU 코어에서 데이터 배치 수신, 가중치 경사 계산, 경사를 다른 TPU 코어와 교환 후 가중치 업데이트를 계산합니다. 기본적으로는 여러 코어에서 손실의 평균을 계산하지만, CrossShardOptimizer의 매개변수를 변경하여 합계를 구할 수 있습니다.

모델의 총 손실은 독립된 샘플당 손실의 평균(또는 합계)으로 계산할 수 있고 이 절차는 하나의 대형 배치에 대한 학습과 수학적으로 일치합니다.

샘플별로 독립되지 않은 가장 일반적인 작업은 배치 정규화로, 각 코어별 배치에서 별도로 실행됩니다. 예를 들어 총 배치 크기가 128이면 사전 코어 배치 크기가 16이고 각 8개 코어가 고유한 16개 샘플에 배치 정규화를 수행합니다. 일부 경우에는 32개 미만 등의 적은 배치에 배치 정규화를 수행할 때 정확도가 낮아지는 것으로 확인되었습니다. 이상적인 시나리오에서는 총 배치 크기가 256~1024개 정도로 커야 합니다. 배치 크기가 메모리에 맞추기에 너무 큰 경우 샤딩의 영향을 사례별로 평가해야 합니다.

확정적 학습

모델 정확성의 차이를 디버그하기 어려운 이유 중 하나는 모델이 서로 다른 프레임워크(TensorFlow, PyTorch, JAX)에서 모델이 학습될 때마다 서로 다른 가중치 초기화와 데이터 셔플링을 사용하기 때문입니다. 여러 번 실행했을 때 거의 일치하는 모델이 생성되도록 학습 절차를 확정하여 수정하는 것이 유익합니다. 이 섹션에서는 MNIST 튜토리얼을 확정적으로 실행하는 방법을 보여줍니다.

  1. CPU에서 단일 단계를 실행하여 초기 체크포인트 파일을 생성합니다. 이 단계는 확정적인 가중치 초기화를 달성하기 위해 사용됩니다. 또한 모델에서 랜덤 함수에 대해 고정된 무작위 시드를 사용하는지 확인합니다. 
# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1
  1. 무작위 시드를 사용하려면 입력 함수에서 데이터 셔플링 함수를 수정하세요. 이 작업은 MNIST 튜토리얼에서 이미 수행했을 수 있습니다. 입력 데이터 처리 작업은 항상 CPU에서 실행되기 때문에 이 방식이 효과가 있습니다. 모델 함수의 무작위 순서 지정 작업은 TPU와 CPU에서 확정적이지 않을 수 있습니다. 무작위 순서 지정 작업을 호출할 때마다 고정 시드를 전달하여 실행 간에 동일한 결과가 발생하도록 합니다. 예를 들면 다음과 같습니다. 
# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)
  1. CPU에서 동일한 모델을 두 번 실행하여 확정적인 학습인지 확인합니다. 합리적인 수의 단계(예: 1,000)로 학습을 실행해야 합니다. 하지만 수렴 단계까지 실행할 필요는 없습니다.

    CPU 학습을 단일 코어 TPU 학습과 비교하므로 단일 TPU 코어에 적합한 배치 크기를 사용하세요. 일반적으로 전체 배치 크기를 8로 나누면 됩니다. 텐서플로우는 실행 간에 완벽하게 동일한 확정성을 보장하지 않으나 손실은 거의 같아야 합니다.

초기 가중치 복사

gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

실행 1

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

출력 1

accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

실행 2

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

출력 2

accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

단일 코어 TPU 학습

MNIST 튜토리얼을 확정적으로 실행할 수 있다면 다음 단계는 TPU상의 CPU 학습 결과를 복제하여 단일 TPU 코어를 통해 데이터 분할과 관련된 문제인지, 아니면 TPU 실행 엔진 자체의 문제인지 확인하는 것입니다.

MNIST 가이드에서 단일 코어 학습 및 평가를 실행하는 방법은 다음과 같습니다.

CPU와 동일한 가중치 초기화 사용

gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

1,000 단계로 학습 실행

python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --num_shards=1 \
    --batch_size=128 \
    --train_steps=1000 \
    --eval_steps=10

출력

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

손실이 CPU 학습 모델과 정확히 일치하지 않더라도 비슷합니다. 내 모델과 비슷하지 않다면 TPU 실행 엔진에 버그가 있다는 의미일 수 있습니다. 버그 신고를 제출하기 전에 다음 내용을 다시 확인하세요.

  1. num_shards=1을 TPUConfig에 전달합니다.

  2. 모델 함수에 무작위 순서 지정 작업이 없으며 입력 함수의 모든 무작위 순서 지정 작업은 올바르게 시드됩니다.

  3. CPU 및 TPU 학습에 대해 동일한 초기 체크포인트 파일을 사용하고 있습니다.

다중 코어 TPU 학습 디버그

모델이 CPU 및 단일 코어 TPU에서 동일한 손실을 기록하는 경우 다음 중 하나가 문제일 수 있습니다.

(a) 다른 초기화 값으로 신경 모델을 학습시켰을 때 발생하는 자연적인 무작위 분산으로 인해 정확성이 저하됩니다.

(b) TPU의 데이터 샤딩과 관련된 문제로 인해 정확성이 저하됩니다.

(a) 문제인지 확인하려면 위와 같이 동일한 가중치 초기화 값을 사용하여 CPU/GPU 및 멀티코어 TPU에서 전체 모델을 다시 학습시킵니다.

정확성 저하가 통계적으로 유의미하다는 확신이 들면 데이터 샤딩과 관련해 가장 가능성이 높은 문제는 다음과 같습니다.

  1. 모델에서 배치 정규화를 사용하는 경우 총 배치 크기가 256 미만(예: 코어별 32 미만)이면 정확성이 낮아질 수 있습니다.
  2. 배치 수준의 손실 함수는 샤딩의 영향을 받습니다. 이러한 손실 함수는 일반적으로 특화되어 있습니다. 예를 들어 Karras et al. 2017은 생성적 적대 신경망(GAN)을 학습시킬 때 배치 분류자를 사용합니다.

TPU VM 문제 해결

다음 문제 및 솔루션은 TPU VM 구성에만 적용 가능합니다.

gcloud 설정 문제 해결

문제
gcloud components update가 다음 오류 메시지를 표시합니다.
ERROR: (gcloud.components.update)
You cannot perform this action because the Cloud SDK component manager is
disabled for this installation.
솔루션
TPU VM에 gcloud를 사용하려면 패키지 관리자를 통해 관리되지 않는 gcloud 설치를 사용해야 합니다. 소스 코드에서 gcloud 설치 단계를 따릅니다.
  sudo apt-get remove google-cloud-sdk
  curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  ./google-cloud-sdk/install.sh
  source ~/.bashrc
문제

gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} 명령어가 다음 오류 메시지를 표시합니다.

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.
해결책

SSH 키 전파에 문제가 있을 수 있습니다. gcloud가 다시 만들도록 자동 생성된 키를 백업 위치로 이동해 보세요.

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

디버그 로그

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

libtpu 라이브러리는 디버깅에 유용할 수 있는 정보를 로깅합니다. 기본적으로 이러한 로그는 각 Cloud TPU VM에서 /tmp/tpu_logs에 기록됩니다. 학습 시작 전 다음 환경 변수를 설정하여 로깅 동작을 수정할 수 있습니다.

TPU_LOG_DIR: 로그가 기록되는 디렉터리
디렉터리 위치는 기본적으로 /tmp/tpu_logs입니다. 디렉터리가 없으면 생성되지만 상위 디렉터리는 생성되지 않습니다. 지정된 디렉터리를 찾거나 만들 때 문제가 발생하면 메시지가 stderr에 출력되지만 프로그램이 중지되지 않고 로깅이 사용 중지됩니다. 디렉터리 이름을 "disabled"로 설정하여 디스크 로깅을 함께 사용 중지합니다.
TPU_MIN_LOG_LEVEL: 디스크에 로깅되는 최소 심각도
선택사항은 0(정보), 1(경고), 2(오류), 3(치명적)입니다. 기본값은 0입니다.
TPU_STDERR_LOG_LEVEL: 디스크와 함께 stderr에 로깅되는 최소 심각도입니다(해당하는 경우).
선택 항목은 TPU_MIN_LOG_LEVEL과 동일합니다. 기본값은 3입니다.
TPU_MAX_LOG_SIZE_MB: 각 로그 파일의 최대 크기(MB)
이전 로그 파일이 대략 이 크기에 도달하면 새 로그 파일이 자동으로 시작됩니다. 기본값은 1024입니다.