메모리 효율적인 트랜스포머 학습 및 추론을 위한 텐서와 시퀀스 병렬성
요약
본 논문은 텐서 병렬성(TP)과 시퀀스 병렬성(SP)을 하나의 장치 축에 통합하는 새로운 병렬 실행 전략인 '텐서와 시퀀스 병렬성(TSP)'을 제안합니다. TSP는 모델 가중치와 토큰 모두를 동일한 장치 축을 따라 분산하여, 파라미터 메모리와 활성화 메모리를 동시에 효율적으로 줄입니다. 이를 통해 기존의 다차원 병렬성 방식보다 통신 부하가 적은 메모리 오버헤드로 메모리 제약이 큰 대규모 모델 학습 및 추론에 효과적인 하드웨어 인식 솔루션을 제공합니다.
핵심 포인트
- TSP는 텐서 병렬성과 시퀀스 병렬성을 하나의 장치 축(device axis)에 통합하여 효율을 극대화한 새로운 병렬 실행 전략이다.
- 기존 방식과 달리, TSP는 가중치와 토큰 모두를 동일한 축에서 분산 처리함으로써 파라미터 메모리와 활성화 메모리를 동시에 절감한다.
- 어텐션 및 게이트드 MLP 블록에 대한 구체적인 런타임 스케줄러 구현을 제시하여 이론적 분석과 함께 실제 성능을 검증했다.
- TSP는 긴 컨텍스트와 메모리 제약이 있는 대규모 모델 학습/추론 환경에서 기존 병렬성 기법의 강력한 하드웨어 인식 대안으로 활용 가능하다.
우리는 텐서 병렬성과 시퀀스 병렬성을 하나의 장치 축 (device axis) 위에 접합 (fold)하는 병렬 실행 전략인 텐서와 시퀀스 병렬성 (TSP) 을 제시합니다. 기존의 다차원 병렬성 레이아웃에서는 텐서 병렬성 (TP) 이 모델 가중치를 쉐드하고, 시퀀스 병렬성 (SP) 이 토큰을 쉐드하여 각각 장치당 파라미터 메모리나 활성화 메모리를 줄입니다. 전통적으로 각 방식은 별도의 메시 차원 (mesh dimension) 에 할당됩니다. 반면 TSP 는 랭크 (rank) 에 가중치 쉐드와 시퀀스 쉐드를 모두 할당하여, 동일한 장치 축을 따라 파라미터 메모리와 활성화 메모리 모두를 줄입니다. 우리는 이 설계를 두 가지 런타임 스케줄러로 구현했습니다. 어텐션의 경우 랭크들은 브로드캐스트된 파라미터 쉐드를 반복하며 시퀀스 단위의 키/밸류 교환을 통해 컨텍스트를 재구성합니다. 게이트드 MLP 의 경우 가중치 쉐드가 링 (ring) 을 따라 순환하고 부분 출력들이 로컬로 누적됩니다. 동일한 장치들을 가로지르며 가중치와 활성화 모두를 쉐드함으로써 TSP 는 추가적인 통신 부피를 줄어든 메모리 오버헤드로 교환합니다. 우리는 이론적인 통신 및 메모리 분석을 제공하고, TSP 어텐션 및 게이트드 MLP 블록의 구현을 설명하며, TSP 를 TP, SP, 그리고 TP+SP 와 벤치마크했습니다. 이 결과는 TSP 를 긴 컨텍스트와 메모리 제약이 있는 모델 학습을 위한 하드웨어 인식 (hardware-aware) 대안으로 자리매김시키며, 파이프라인 병렬성과 전문가 병렬성 같은 기존 병렬성 스키마와 함께 밀집형 및 mixture-of-expert 모델의 병렬성 축으로 활용 가능함을 보여줍니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.CL의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기