Mesh-TensorFlow: 슈퍼컴퓨터를 위한 딥 러닝

배치 분할(data-parallelism)은 메모리 제약 조건으로 인해 매우 큰 모델을 훈련시키는 데 어려움이 있으며, 높은 지연 시간과 작은 배치 크기에서의 비효율성 등의 문제를 겪고 있지만, 이 모든 문제가 더 일반적인 분산 전략(model-parallelism)을 통해 해결될 수 있습니다. 그러나 효율적인 모델 병렬 알고리즘은 특히 대규모 클러스터에서 발견, 설명 및 구현하기가 복잡합니다. 우리는 이러한 일반적인 분산 텐서 계산을 지정하기 위한 언어인 Mesh-TensorFlow를 소개합니다. 데이터 병렬 처리는 텐서와 연산을 "배치" 차원에 따라 분할하는 것으로 볼 수 있는 반면, Mesh-TensorFlow에서는 사용자가 프로세서의 다차원 메시(mesh)의 어떤 차원에도 텐서의 어떤 차원을 분할하도록 지정할 수 있습니다. Mesh-TensorFlow 그래프는 Allreduce와 같은 집합적 통신 원시형을 포함한 병렬 연산으로 구성된 SPMD(Single-Program-Multiple-Data) 프로그램으로 컴파일됩니다. 우리는 Mesh-TensorFlow를 사용하여 Transformer 시퀀스-투-시퀀스(sequence-to-sequence) 모델의 효율적인 데이터 병렬 및 모델 병렬 버전을 구현했습니다. 최대 512개 코어까지 구성된 TPU 메시를 사용하여 최대 50억 개 매개변수를 가진 Transformer 모델을 훈련시키며, WMT'14 영어-프랑스어 번역 작업과 10억 단어 언어 모델링 벤치마크에서 기존 최고 성능 결과를 초월하였습니다. Mesh-TensorFlow는 https://github.com/tensorflow/mesh 에서 이용 가능합니다.