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로 균등하게 나뉘어야 합니다.
메모리 관리 조정
TPU_PREMAPPED_BUFFER_SIZE환경 변수를 사용하여 하위 수준 런타임 동작을 세부 조정할 수 있습니다.
설명:
TPU_PREMAPPED_BUFFER_SIZE는 TPU 런타임에서 데이터 전송(예: DMA)에 사용하도록 미리 매핑되고 고정된 호스트 메모리 버퍼의 크기(바이트)를 설정합니다. 기본값은 4,294,967,296바이트입니다. 값은 2^12(4KB = 4 * 1,024바이트 = 4,096 = 2^12)의 배수여야 합니다.다음은 유효한 TPU_PRE_MAPPED_BUFFER_SIZE 값의 예시입니다.
17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped). 40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).영향: 특히 텐서가 크거나 호스트-기기 통신이 빈번한 워크로드의 경우 이 크기를 늘리면 호스트와 TPU 기기 간 데이터 전송 성능이 향상될 수 있습니다. 하지만 고정된 호스트 메모리 양도 증가하여 다른 프로세스에 사용할 수 있는 메모리가 줄어듭니다.
버퍼 크기
사전 매핑된 버퍼 리전이 프로그램 런타임 중에 메모리를 할당할 만큼 충분히 크지 않으면 워크로드가 실패하고 다음과 유사한
RESOURCE_EXHAUSTED오류가 반환됩니다.'
RESOURCE_EXHAUSTED로 인해 미리 매핑된 리전에서 버퍼를 할당하지 못했습니다.allocation_size할당을 시도합니다. 할당할 수 없었습니다. 할당되지 않은available_size가 있습니다.'버퍼가 지나치게 크면 TPU가 초기화되는 데 훨씬 오래 걸릴 수 있으며(15초 이상) TPU가 멈춘 것처럼 보일 수 있습니다.
이 문제를 진단하려면 TPU 런타임 로그를 검사합니다. 이러한 로그는 버퍼 사전 매핑을 포함하여 수행되는 작업을 자세히 설명합니다. /tmp/tpu_logs/tpu_driver.INFO에서 로그를 찾거나 환경 변수 TPU_STDERR_LOG_LEVEL=0을 설정하여 콘솔에 직접 출력할 수 있습니다. 이 설정은 다음과 비슷한 출력을 생성합니다.
I0604 12:45:24.926233 62136 tpu_hal.cc:214] Starting premapped memory manager initialization... I0604 12:45:29.411218 62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0 I0604 12:45:29.411244 61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s I0604 12:45:29.411267 61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296 ``` This output will tell you how long it took to initialize the TPU and the size of the premapped buffer.사용: 미리 매핑된 버퍼가 너무 작거나 너무 크면 다음 환경 변수를 사용하여 버퍼 크기를 수동으로 설정할 수 있습니다.
TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the pre-mapped buffer region. TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of a single buffer that can be allocated from the pre-mapped region.예를 들어 다음 작업을 수행할 수 있습니다.
export TPU_PREMAPPED_BUFFER_SIZE=4294967296버퍼 크기를 설정하고 다음 작업을 수행합니다.
export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES ``` to enable it. This export sets the size to the default.안내: 호스트-기기 데이터 전송이 병목 현상이라고 의심되면 TPU_PREMAPPED_BUFFER_SIZE 값을 조정합니다. 호스트 메모리 사용량과 모델 성능을 모니터링하여 최적의 균형을 찾습니다. 기본값은 일반적으로 대부분의 사용 사례에 충분합니다.
XLA 컴파일러 최적화
XLA는 머신러닝 컴파일러로, TPU, CPU, GPU, 기타 플랫폼의 바이너리를 생성할 수 있습니다. XLA는 표준 TensorFlow 코드베이스의 일부이지만 PyTorch 및 JAX 모델에서도 사용 가능합니다. Cloud TPU용 모델은 XLA 그래프로 변환되고 그런 다음 XLA는 TPU 실행 파일로 컴파일됩니다. XLA에 대한 자세한 내용은 XLA: 머신러닝의 컴파일러 최적화를 참조하세요.
패딩
TPU 메모리를 효율적으로 사용하려면 데이터를 128x8개 청크로 분할할 수 있도록 구조화합니다. 행렬 연산용 데이터가 전체 128x8개 청크를 채우지 않으면 XLA 컴파일러에서 텐서를 패딩합니다. 패딩에는 다음과 같은 두 가지 취약점이 있습니다.
- 패딩된 텐서의 TPU 코어 사용률이 낮습니다.
- 패딩은 텐서에 필요한 온칩 메모리 스토리지 양을 늘리므로 메모리 부족 오류가 발생할 수 있습니다.
패딩은 필요한 경우 XLA 컴파일러에서 자동으로 수행되지만 개발자는 메모리 뷰어 도구를 통해 수행되는 패딩 양을 결정할 수 있습니다. TPU에 적합한 텐서 차원을 선택하여 패딩을 방지할 수 있습니다.
텐서 차원
최대 FLOPS를 달성하려면 행렬 곱셈 차원이 사용하는 TPU 버전의 MXU 크기보다 커야 합니다. v6e의 MXU 크기는 256 x 256이고 v6e 이전 버전은 128 x 128입니다. 자세한 내용은 Cloud TPU 시스템 아키텍처를 참조하세요.
배치 크기
XLA 컴파일러는 TPU HBM 메모리에 저장된 텐서의 크기를 반올림하여 계산을 보다 효율적으로 수행합니다. 이 패딩은 하드웨어 수준에서 투명하고 결과에 영향을 미치지 않습니다. 그러나 경우에 따라 패딩으로 인해 메모리 사용량과 실행 시간이 상당히 늘어날 수도 있습니다.
TPU 런타임은 메모리에 텐서를 배치하여 연산 효율성을 극대화하고 패딩을 최소화합니다. 메모리 오버헤드를 최소화하고 연산 효율성을 극대화하려면 다음 중 하나에 해당해야 합니다.
총 배치 크기는 64의 배수(TPU 코어당 8)이고 특징 차원 크기는 128의 배수여야 합니다.
총 배치 크기는 1,024의 배수(TPU 코어당 128)이고 특징 차원 크기는 8의 배수여야 합니다.
배치 크기 1,024 및 128의 배수인 특징 차원을 사용하면 최고의 효율성을 얻을 수 있지만 모든 모델에서 가능하지는 않습니다.
퓨전
퓨전은 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[i] = tmp[i] * z[i];
}
퓨전을 사용하면 배열 액세스가 동시에 발생합니다.
for (i = 0; i < element_count; i++) {
result[i] = (x[i] + y[i]) * z[i];
}
이 예시에서는 메모리 왕복 횟수가 줄어 XLA에서 'tmp'에 공간을 할당하지 않아도 됩니다.
퓨전은 중요한 최적화 방식으로, 다음과 같은 여러 측면에서 Cloud TPU에 도움이 됩니다.
- 기본 메모리에 중간 결과를 저장할 필요성을 없애(느린 속도의 원인) 메모리 전송량을 줄일 수 있습니다.
- 이로써 하드웨어 단위의 활용도가 더욱 높아집니다.
- 동시에 실행해야 하는 버퍼 수가 적기 때문에 모델의 메모리 사용률을 낮출 수 있습니다.
브로드캐스팅
서로 다르지만 호환 가능한 형태의 두 텐서를 조합할 때 암시적으로 브로드캐스팅이 발생합니다.
예를 들어 tf.add(vector, matrix)를 사용하려면 벡터에 행렬 형태 브로드캐스팅을 수행해야 합니다. 작업 결과 형태는 행렬 형태와 동일해집니다. 자세한 내용은 배열 브로드캐스팅 가이드를 참조하세요.
브로드캐스트를 소비자와 퓨전할 수 있는 경우도 있지만 브로드캐스트를 강제 적용하면 성능이 저하되고 메모리 사용량이 증가할 수 있습니다.
다음 예시에서는 벡터 추가 시 브로드캐스트가 포함되며 행렬은 argmax와 퓨전되어 구체화된 브로드캐스트를 만들 수 없습니다.
`tf.argmax(tf.add(vector, zero_matrix), axis=0)`
Ironwood 듀얼 칩렛 아키텍처의 성능 권장사항
Ironwood 프로그래밍 모델을 사용하면 이전 세대(TPU v4 및 v5p)에서 사용된 단일 논리 코어(MegaCore라고도 함) 아키텍처 대신 TPU 기기 2개에 액세스할 수 있습니다. 이러한 변경으로 칩 제조의 비용 효과와 효율성이 향상됩니다. 이는 아키텍처 변화를 나타내지만 새로운 디자인을 통해 기존 소프트웨어 모델을 최소한으로 변경하여 재사용할 수 있습니다.
듀얼 칩렛 아키텍처로 최상의 성능을 얻으려면 다음 방법을 사용하는 것이 좋습니다.
칩렛 간 텐서 동시 로드 사용: 고대역폭 D2D 인터페이스는 효율적인 텐서 동시 로드를 위해 설계되었습니다. 온칩 기기 2개에 텐서를 분할하는 것이 좋습니다.
계층적 집단 활용: 통신 효율성을 극대화하려면 온칩 칩렛 간의 초고속 D2D 링크와 슬라이스 내 고속 ICI 링크라는 2단계 네트워크 계층을 활용합니다. SPMD(단일 프로그램, 다중 데이터)로 자동 동시 로드를 사용할 경우 XLA 컴파일러는 계층적 집단 연산을 자동으로 생성하여 이를 처리합니다. 모델을 수동으로 파티셔닝할 때는 이 계층 구조를 중심으로 통신 패턴도 설계해야 합니다. 다른 칩의 기기와 통신하기 전에 같은 칩에 있는 두 기기 간 통신에 우선순위를 부여합니다.
통신과 계산 중첩: 하드웨어 사용률을 극대화하려면 all-reduce와 같은 집단 통신 작업을 SparseCore로 오프로드합니다. 행렬 곱셈 단위(MXU)에 바인딩되지 않은 이러한 작업은 TensorCore에서 계산을 계속하는 동안 SparseCore에서 동시에 실행될 수 있습니다. 이 기법은 이전 MegaCore 아키텍처의 융합 작업에 내재된 성능 이점 일부를 복구할 수 있습니다.
임베딩을 위해 SparseCore로 오프로드: 듀얼 칩렛 설계에서 임베딩 테이블은 두 칩렛의 HBM에 걸쳐 파티셔닝될 수 있습니다. 이러한 공유 메모리 부족으로 인한 성능 저하를 방지하려면 임베딩 수집 작업을 SparseCore로 오프로드합니다. 이 전략은 고속 D2D 상호 연결을 활용하여 칩렛 간에 임베딩 벡터를 효율적으로 전송합니다. SparseCore 및 임베딩 모델에 대한 자세한 내용은 대규모 임베딩 모델(LEM)을 위한 SparseCore 심층 분석을 참조하세요.
TPU7x의 Ironwood 아키텍처에 대한 자세한 내용은 TPU7x(Ironwood)를 참조하세요.