최근 몇 년간 Transformer 아키텍처 기반 대형 언어 모델(LLM)은 뛰어난 성능으로 널리 주목받고 있습니다. 산업용 LLM은 고품질 서비스를 제공하기 위해 긴 시퀀스 입력을 처리해야 합니다. 그러나 시퀀스 길이에 따라 메모리 사용량이 제곱 수준으로 증가하여 긴 시퀀스 학습 확장성에 제한이 있습니다. 기존 병렬 방법은 실행 중에 중복 텐서를 생성하여 메모리 최적화 여지가 있으며, 텐서 병렬화(TP)는 계산과 통신의 효과적인 중첩을 구현할 수 없습니다. 이러한 문제를 해결하기 위해 본문에서는 Transformer 학습의 핵심 계산 단위(연속된 두 행렬 곱셈 및 그 사이에 존재할 수 있는 함수 연산 O=f(AB)C)를 위해 설계된 범용 병렬 방법인 메모리 효율 텐서 병렬(METP)을 제안합니다. METP는 여러 장치에 O 계산의 하위 작업을 분배하고, 집합 통신 대신 점대점 통신(send/recv)을 사용하여 하위 행렬을 교환함으로써 중복 텐서 생성을 방지합니다. 이중 버퍼링 기술로 계산과 통신을 깊이 중첩시키고, 긴 시퀀스 Transformer 학습을 안내하는 완전 중첩 이론 조건을 제시합니다. 이론 분석에 따르면 병렬도 p일 때, FlashAttention을 사용하지 않고 주의(attention)를 계산할 경우 METP의 메모리 비용은 O(1/p³)이며, FlashAttention을 사용해 다중 헤드 자기 주의를 계산할 때는 TP에 비해 최소 41.7% 이상의 메모리를 절감할 수 있습니다. 8대의 A100 GPU 환경을 기반으로 한 실험에서 METP는 다른 방법에 비해 시퀀스 길이를 2.38배에서 2.99배까지 향상시킬 수 있음을 입증했습니다.

분산 학습;대규모 언어 모델;긴 시퀀스;기계 학습 시스템;메모리 효율;텐서 병렬