Cloud TPU 성능 가이드

TPU 성능 문제를 해결하는 첫 번째 단계는 모델을 프로파일링하는 것입니다. 성능 프로필을 캡처하는 방법에 대한 자세한 내용은 Cloud TPU에서 모델 프로파일링을 참조하세요.

TPU 모델 성능

이 섹션에서는 모델 성능을 저하시킬 수 있는 일반적인 문제와 이를 해결하는 방법을 설명합니다.

모델이 입력 제약적

TPU는 매우 빠르게 계산을 수행합니다. TPU가 유휴 상태가 아니게 하려면 TPU에 안정적인 데이터 스트림이 로드되고 있는지 확인하는 것이 중요합니다. 이 방법은 데이터 세트를 로드하고 사전 처리하는 방법에 따라 다릅니다. 예를 들어 tf.data.TFRecordset() 및 num_parallel_reads 매개변수를 사용하여 데이터 파일을 동시에 읽을 수 있습니다.
샤딩으로 인해 배치 크기가 너무 작음(코어간 배치 분할)

TPU 런타임은 TPU 기기의 모든 8개 코어에 배치를 분할합니다(예: v2-8 또는 v3-8). 전역 배치 크기를 128로 지정하면 각 코어의 배치 크기는 16(128/8)입니다.

메모리 사용량을 최적화하려면 TPU 메모리에 적합한 가장 큰 배치 크기를 사용합니다. TPU 코어마다 행렬 곱셈을 처리하는 데 2차원의 8X128 벡터 레지스터를 사용합니다. 일반적으로 배치 크기는 8 또는 128로 균등하게 나눠져야 합니다.

XLA 컴파일러 최적화

XLA는 TPU, CPU, GPU, 기타 플랫폼의 바이너리를 생성할 수 있는 머신러닝의 컴파일러입니다. XLA는 표준 TensorFlow 코드베이스의 일부이지만 PyTorch 및 JAX 모델에서도 사용 가능합니다. Cloud TPU 모델은 XLA 그래프로 변환된 후 XLA는 TPU 실행 파일로 컴파일됩니다. XLA에 대한 자세한 내용은 XLA: 머신러닝의 컴파일러 최적화를 참조하세요.

패딩

TPU 메모리를 효율적으로 사용하려면 데이터를 128x8개의 청크로 분할할 수 있도록 구조화합니다. 행렬 연산용 데이터가 전체 128x8개의 청크를 채우지 않으면 XLA 컴파일러가 텐서를 패딩합니다. 패딩에는 두 가지 취약점이 있습니다.

패딩된 텐서는 TPU 코어 사용률이 낮습니다.
패딩은 텐서에 필요한 온칩 메모리 스토리지 양을 늘리므로 메모리 부족 오류가 발생할 수 있습니다.

패딩은 필요한 경우 XLA 컴파일러에서 자동으로 수행되지만 메모리 뷰어 도구를 사용하여 수행되는 패딩 양을 결정할 수 있습니다. TPU에 적합한 텐서 측정기준을 선택하여 패딩을 방지할 수 있습니다.

텐서 측정기준

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

TPU 런타임은 메모리에 텐서를 배치하여 연산 효율성을 극대화하고 패딩을 최소화합니다. 메모리 오버헤드를 최소화하고 연산 효율성을 극대화하려면 다음 조건 중 하나를 충족해야 합니다.

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

1,024의 배치 크기 및 128의 배수인 특징 차원을 사용하면 최고의 효율성을 얻을 수 있지만, 모든 모델에서 가능하지는 않습니다.

Fusion

Fusion은 XLA 컴파일러가 프로그램을 최적화하기 위해 사용하는 일반적인 기술입니다. 퓨전 오퍼레이션이란 실행할 여러 개의 오퍼레이션으로 구성된 조합을 의미합니다.

예를 들어 다음과 같이 연속된 오퍼레이션을 생각해 보세요.

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

이 코드는 다음 유사 코드와 거의 동일합니다.

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result = tmp[i] * z[i];
    }

퓨전을 사용하면 배열 액세스가 동시에 발생합니다.

    for (i = 0; i < element_count; i++) {
      result = (x[i] + y[i]) * z[i];
    }

이 예시에서는 메모리 왕복 횟수가 줄어 XLA가 'tmp'에 공간을 할당할 필요가 없습니다.

퓨전은 중요한 최적화 방식으로, 다음과 같은 여러 측면에서 Cloud TPU에 도움이 됩니다.

기본 메모리에 중간 결과를 저장할 필요성을 없애(느린 속도의 원인) 메모리 전송량을 줄일 수 있습니다.
이로써 하드웨어 단위의 활용도가 더욱 높아집니다.
동시에 실행해야 하는 버퍼 수가 적기 때문에 모델의 메모리 사용률을 낮출 수 있습니다.

브로드캐스팅

서로 다르지만 호환 가능한 형태의 두 텐서를 조합할 때 암시적으로 브로드캐스팅이 발생합니다.

예를 들어 tf.add(vector, matrix)에서는 벡터에 행렬 형태 브로드캐스팅을 수행해야 합니다. 오퍼레이션의 결과, 행렬과 형태가 동일해집니다. 자세한 내용은 배열 브로드캐스팅 가이드를 참조하세요.

브로드캐스트를 소비 함수와 퓨전할 수 있는 경우도 있지만 브로드캐스트를 강제 적용하면 성능이 저하되고 메모리 사용량이 증가할 수 있습니다.

다음 예에서는 벡터 추가 시 브로드캐스트가 포함되며 행렬은 argmax와 퓨전되어 구체화된 브로드캐스트로 만들 수 없습니다.

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`