Ulysses Sequence Parallelism: Training with Million-Token Contexts

Context

트랜스포머의 어텐션 메커니즘은 시퀀스 길이 n에 대해 O(n²) 계산량과 메모리를 소비한다. FlashAttention 같은 최적화도 O(n) 메모리만 가능하며 O(n²) 계산 오버헤드는 남아 있어, 32K 토큰 이상의 긴 시퀀스 학습 시 단일 GPU 메모리 한계에 도달한다. 데이터 병렬화는 각 GPU가 전체 시퀀스를 처리해야 하므로 이 문제를 해결하지 못한다.

Technical Solution

• 시퀀스를 P개 GPU에 걸쳐 시퀀스 차원으로 분할: 각 GPU가 토큰 [i·n/P, (i+1)·n/P) 범위 담당
• QKV 프로젝션을 로컬 시퀀스 청크에 대해 병렬 수행: 각 GPU가 자신의 시퀀스 부분에 대한 쿼리, 키, 값 계산
• All-to-All 통신으로 키-값 쌍 재분배: 각 GPU가 모든 시퀀스 위치를 갖되 어텐션 헤드 부분만 담당하도록 변환
• 독립적 어텐션 헤드 계산: 각 GPU가 할당된 헤드 부분에 대해 표준 어텐션(FlashAttention 또는 SDPA) 수행
• All-to-All 통신으로 원래 형태 복원: 시퀀스-샤딩 포맷으로 되돌린 후 아웃풋 프로젝션 수행

Impact

• 메모리 사용량: SP=4 적용 시 96K 토큰에서 피크 메모리 66GB (H100 80GB 범위 내), SP 미적용 시 8K 이상 OOM
• 처리량: SP=4로 64K 토큰 처리 시 13,396 토큰/초 달성 (8K 기준선 대비 3.7배)
• 32K 토큰: 2.1배 처리량 증가
• 동일 시퀀스 길이(8K)에서 SP=4는 단일 GPU 기준선과 비교하여 약 1배 처리량 유지 (All-to-All 통신 오버헤드 최소)

Key Takeaway

어텐션 헤드의 독립성을 활용해 시퀀스 로컬리티를 헤드 로컬리티로 트레이드오프하면, 상대적으로 낮은 통신 오버헤드로 긴 시퀀스 학습을 확장할 수 있다. Ring Attention 대비 O(n·d/P) vs O(n·d) 통신량과 전체 비섹셔널 대역폭 활용으로 더 효율적이다.