울리세스 시퀀스 병렬 처리: 100만 토큰 컨텍스트로의 학습

최근 대규모 언어 모델의 성능 향상을 위해 시퀀스 병렬 처리를 활용한 훈련이 중요해지고 있습니다. 문서 분석, 코드 이해, 복잡한 추론, RAG(Retrieval-Augmented Generation) 워크로드 등 다양한 작업에서 수십만에서 백만 토큰에 이르는 긴 시퀀스를 처리해야 하는 경우가 많아졌습니다. 평균적인 책의 토큰 수는 약 25만 개이므로, 다중 문서 컨텍스트 또는 책 길이의 입력을 처리하려면 단일 GPU 메모리 용량 이상의 긴 시퀀스를 다루어야 합니다.

하지만 이러한 긴 컨텍스트로의 훈련은 상당한 메모리 문제를 야기합니다. 특히, 트랜스포머의 어텐션 연산은 시퀀스 길이에 따라 2차적으로 증가하기 때문에, 컨텍스트 길이가 수만 토큰을 넘어갈 경우 GPU 메모리 용량을 초과할 수 있습니다. 이러한 문제를 해결하기 위해 울리세스 시퀀스 병렬 처리(Snowflake AI Research의 Arctic Long Sequence Training (ALST) 프로토콜의 일부)는 어텐션 연산을 다수의 GPU에 분산시켜 효과적인 해결책을 제시합니다.

The Challenge of Long Sequence Training

트랜스포머의 어텐션 메커니즘은 시퀀스 길이에 따라 2차적으로 스케일됩니다. 길이가 n인 시퀀스의 경우, 표준 어텐션은 O(n^2) FLOPs와 O(n^2) 메모리를 사용하여 어텐션 점수 행렬을 계산하고 저장하는 데 필요합니다. FlashAttention과 같은 최적화된 구현 방식은 타일링된 연산을 통해 O(n) 메모리로 줄이지만, O(n^2) 연산은 여전히 존재합니다. 따라서 매우 긴 시퀀스(32k+ 토큰)의 경우, FlashAttention을 사용하더라도 단일 GPU 메모리 한계를 초과할 수 있습니다.

긴 컨텍스트 훈련이 필수적인 시나리오는 다음과 같습니다:

문서 이해: 책, 법률 문서 또는 연구 논문 전체 처리
코드 분석: 여러 상호 연결된 파일을 포함하는 대규모 코드 기반 이해
추론 작업: 단계별로 ‘생각’하는 모델은 추론 중에 수천 개의 토큰을 생성할 수 있습니다.
검색 증강 생성: 검색된 여러 패시지를 컨텍스트에 통합

전통적인 데이터 병렬 처리 방식은 이러한 문제 해결에 도움이 되지 않습니다. 각 GPU는 여전히 어텐션 블록 내에서 전체 시퀀스를 처리해야 하기 때문입니다. 따라서 시퀀스 자체를 여러 장치에 분산하는 방법이 필요합니다.

How Ulysses Works

울리세스 시퀀스 병렬 처리(DeepSpeed Ulysses 논문에 소개)는 매우 영리한 접근 방식을 취합니다. 시퀀스 차원을 분할하는 것 외에도 어텐션 헤드를 GPU에 분할합니다. 이는 어텐션 헤드를 독립적으로 계산할 수 있다는 점을 활용하여 낮은 통신 오버헤드로 효율적인 병렬화를 가능하게 합니다.

울리세스는 입력 시퀀스를 시퀀스 차원을 따라 분할하고 all-to-all 통신을 사용하여 키-값 쌍을 교환하여 각 GPU가 어텐션 헤드의 부분 집합을 계산할 수 있도록 합니다.

구체적인 작동 방식은 다음과 같습니다:

시퀀스 샤딩: 입력 시퀀스는 P개의 GPU에 따라 시퀀스 차원을 따라 분할됩니다. 각 GPU i는 [i⋅n/P,(i+1)⋅n/P) 토큰을 보유합니다.
QKV 투영: 각 GPU는 로컬 시퀀스 청크에 대한 쿼리, 키, 값 투영을 계산합니다.
All-to-All 통신: all-to-all 집단 연산은 데이터를 재분배하여 각 GPU가 모든 시퀀스 위치에 대한 모든 어텐션 헤드에 대한 데이터를 보유하도록 합니다.
로컬 어텐션: 각 GPU는 할당된 헤드에 대해 표준 어텐션 메커니즘(FlashAttention 또는 SDPA)을 사용하여 어텐션을 계산합니다.
All-to-All 통신: 또 다른 all-to-all 연산은 재분배를 되돌려 시퀀스 샤딩 형식으로 돌아갑니다.
출력 투영: 각 GPU는 로컬 시퀀스 청크에 대한 출력 투영을 계산합니다.

Communication Complexity

울리세스는 각 어텐션 레이어당 두 개의 all-to-all 연산을 필요로 하며, GPU당 총 통신 볼륨은 O(n⋅d/P)입니다. 여기서 n은 시퀀스 길이, d는 은닉 차원, P는 병렬 처리 정도입니다. Ring Attention은 P−1에 이르는 순차적 점대점 전송을 통해 O(n⋅d)를 통신합니다. 울리세스는 all-to-all이 단일 집단 단계에서 전체 분할 대역폭을 활용할 수 있는 반면, Ring Attention은 P−1 호핑을 통해 직렬화되기 때문에 더 낮은 지연 시간을 제공합니다.

Integration with Accelerate

Accelerate는 ParallelismConfig 클래스와 DeepSpeed 통합을 통해 울리세스 시퀀스 병렬 처리를 위한 기반을 제공합니다. 이를 통해 다양한 하드웨어 환경에서 효율적인 학습이 가능합니다.

Best Practices

울리세스 시퀀스 병렬 처리를 최대한 활용하기 위해서는 몇 가지 모범 사례를 따라야 합니다. 첫째, 시퀀스 길이는 sp_size로 나누어 떨어져야 합니다. 둘째, FlashAttention을 사용하여 어텐션 연산 속도를 높여야 합니다. 셋째, DeepSpeed ZeRO와 결합하여 메모리 사용량을 더욱 줄일 수 있습니다. 넷째, PyTorch 메모리 단편화에 적합한 할당기를 사용해야 합니다. 다섯째, 2D 병렬 처리 구성을 조정하여 GPU 활용도를 최적화할 수 있습니다. 여섯째, Liger-Kernel을 활용하여 성능을 향상시킬 수 있습니다. 마지막으로, 토큰을 랭크에 분배하는 방법을 고려하여 효율성을 높일 수 있습니다.

이러한 모범 사례를 따르면, 울리세스 시퀀스 병렬 처리를 사용하여 매우 긴 시퀀스를 효과적으로 학습하고, 대규모 언어 모델의 성능을 향상시킬 수 있습니다.