PrimePar: 대규모 트랜스포머 모델 학습을 위한 효율적인 시공간 텐서 분할

트랜스포머 기반 대규모 언어 모델(LLM)의 급격한 확장과 함께, 이러한 모델들의 학습은 새로운 병렬 학습 기법에 대한 요구가 점점 커지고 있습니다. 텐서 분할은 데이터 병렬성과 모델 병렬성을 포함하는 광범위하게 연구된 병렬 기법으로, LLM 학습 성능에 중요한 영향을 미칩니다. 그러나 기존의 최첨단 병렬 학습 시스템들은 텐서 분할 공간이 불완전하여, 분할된 하위 연산자의 분포가 공간 차원에만 제한되어 있습니다. 우리는 LLM 학습 인스턴스의 텐서 분할에 시간 차원을 도입함으로써, 장치 간 집단 통신을 회피할 추가 기회를 제공하여 메모리 공간을 절약하고, 장치 간 통신을 계산과 중첩할 수 있음을 발견했습니다. 본 논문에서는 하위 연산자를 공간 및 시간 차원 모두에 분배하는 새로운 텐서 분할 원시 연산자를 제안하여, 현재 솔루션 대비 통신 및 메모리 오버헤드 감소를 더욱 탐구합니다. 이 새로운 원시 연산자는 더 넓은 병렬화 공간을 생성하고, 최첨단 기법과 비교하여 더 낮은 최대 메모리 점유율로 더 나은 학습 처리량을 달성하는 병렬 솔루션을 도출합니다. 최적화된 병렬 트랜스포머 모델 학습을 여러 장치에 효율적으로 배포하기 위해, 허용 가능한 탐색 시간 내에 공간-시간 텐서 분할 공간에서 최적의 병렬 솔루션을 찾는 최적화 알고리즘도 제시합니다. 평가 결과, 최적화된 텐서 분할은 LLM 학습 시 최첨단 분산 학습 시스템 대비 최대 1.68배의 학습 처리량과 69%의 최대 메모리 점유율을 달성합니다. 32 GPU로 확장 시, 벤치마크 전반에서 기하 평균 속도 향상은 1.30배에 이르며, 3D 병렬화 적용 시 최대 1.46배의 학습 처리량을 얻을 수 있습니다.