백만 토큰 컨텍스트 학습을 위한 Ulysses 시퀀스 병렬 처리

Hugging Face는 Ulysses 시퀀스 병렬 처리( Snowflake AI Research의 Arctic Long Sequence Training 프로토콜의 일부)를 Accelerate, Transformers, 그리고 TRL 라이브러리에 통합했습니다. 이 통합을 통해 개발자는 어텐션 연산을 여러 GPU에 분산시켜 어텐션 메커니즘의 2차 메모리 스케일링을 극복하고 수십만에서 수백만 토큰에 이르는 시퀀스로 대형 언어 모델(LLM)을 훈련할 수 있습니다.

긴 시퀀스 훈련의 도전 과제

표준 트랜스포머 어텐션은 시퀀스 길이 $n$에 대해 FLOPs와 메모리 모두 $O(n^2)$로 2차적으로 스케일합니다. FlashAttention은 메모리 사용량을 $O(n)$으로 줄이지만, $O(n^2)$ 연산 요구량은 그대로 남아 있습니다. 32k 토큰을 초과하는 시퀀스의 경우, 훈련은 일반적으로 단일 GPU의 메모리 용량을 초과하므로, 데이터 병렬에만 의존하지 않고 시퀀스 자체를 여러 장치에 나누는 방법이 필요합니다.

Ulysses 시퀀스 병렬 처리 작동 방식

Ulysses 시퀀스 병렬 처리(SP)는 시퀀스 차원과 어텐션 헤드를 모두 GPU에 분할하여 어텐션 연산을 분산합니다. 과정은 다음 단계로 진행됩니다:

  1. Sequence Sharding: 입력 시퀀스가 $P$개의 GPU에 걸쳐 분할되며, 각 GPU는 토큰의 로컬 청크를 보유합니다.
  2. QKV Projection: 각 GPU는 로컬 청크에 대해 query, key, value 프로젝션을 계산합니다.
  3. All-to-All Communication: all-to-all 집합 연산이 데이터를 재분배하여 각 GPU가 모든 시퀀스 위치를 보유하지만 어텐션 헤드의 일부만 보유하도록 합니다.
  4. Local Attention: GPU는 할당된 헤드에 대해 FlashAttention 또는 SDPA를 사용하여 어텐션을 계산합니다.
  5. All-to-All Communication: 두 번째 all-to-all 연산이 데이터를 시퀀스-샤딩 형식으로 되돌립니다.
  6. Output Projection: 각 GPU는 로컬 시퀀스 청크에 대한 출력 프로젝션을 계산합니다.

통신 복잡도

Ulysses는 어텐션 레이어당 두 번의 all-to-all 연산을 필요로 하며, GPU당 통신량은 $O(n · d / P)$입니다($n$은 시퀀스 길이, $d$는 은닉 차원, $P$는 병렬도). 이는 GPU당 $O(n · d)$를 통신하고 $P-1$ 홉에 걸쳐 전송을 직렬화하는 Ring Attention보다 효율적입니다.

생태계 통합

Accelerate

Accelerate는 ParallelismConfig 클래스와 DeepSpeed 통합을 통해 Ulysses를 구현합니다. 주요 매개변수로는 sp_size(시퀀스 병렬에 사용할 GPU 수)와 sp_backend("deepspeed"로 설정해야 함)가 있습니다. accelerator.prepare()가 호출되면 시스템은 모델을 자동으로 UlyssesSPAttentionHF에 등록하고 데이터 로더를 UlyssesSPDataLoaderAdapter로 감싸습니다.

Transformers Trainer

Transformers Trainer는 TrainingArguments.parallelism_config를 통해 Ulysses 통합을 처리합니다. 데이터 로더 감싸기, 시퀀스 샤딩, 가중 손실 집계 등을 자동화하여 토큰이 랭크 간에 고르게 분배되지 않아도 그래디언트가 올바르게 유지됩니다.

TRL SFTTrainer

TRL의 SFTTrainer는 packing 기능과 같이 패딩 낭비를 줄이는 감독 파인튜닝 최적화를 추가합니다. 시퀀스 나눔을 보장하기 위해 pad_to_multiple_of가 sp_size와 동일해야 합니다. 또한 SFTTrainer는 Ulysses가 활성화된 경우 사전 이동된 레이블을 자동으로 관리합니다.

Ulysses와 Ring Attention 비교

항목 Ulysses (DeepSpeed) Ring Attention (FSDP2)
병렬 방식 어텐션 헤드 분할 링 기반 KV 교환
백엔드 DeepSpeed ZeRO PyTorch FSDP2
어텐션 지원 FlashAttention 2/3, SDPA SDPA만
통신 레이어당 두 개의 all-to-all P2P 링 통신
GPU당 통신량 $O(\text{total_seq} \times \text{hidden} / \text{sp_size})$ $O(\text{total_seq} \times \text{hidden})$
헤드 수 제약 num_heads >= sp_size 없음

성능 벤치마크

Hugging Face는 H100 80GB GPU를 사용해 Gutenberg 영어 데이터셋에서 Qwen3-4B를 대상으로 Ulysses SP를 벤치마크했습니다.

메모리 감소

SP=4를 사용하면 동일한 시퀀스 길이에서 GPU당 메모리가 3.3배 감소합니다. 이를 통해 기본 8K 토큰(DP=4)에서 96K 토큰(SP=4)으로 확장할 수 있으며 80GB 메모리 한도 내에 머무릅니다. 128K 토큰에서는 구성에서 메모리 초과(OOM) 상태에 도달했습니다.

처리량

시퀀스 길이가 증가함에 따라 처리량이 증가하는데, 이는 2차 어텐션 연산이 통신 오버헤드를 지배하기 때문입니다. 64K 토큰에서 SP=4는 초당 13,396 토큰을 달성했으며, 이는 8K 기준 대비 3.7배의 처리량입니다.

구현을 위한 모범 사례

  • Sequence Divisibility: pad_to_multiple_of를 사용하여 시퀀스 길이가 sp_size로 나누어 떨어지도록 합니다.
  • Attention Backends: Ampere GPU에는 FlashAttention 2를, Hopper GPU에는 FlashAttention 3을 사용합니다.
  • Memory Optimization: Ulysses를 DeepSpeed ZeRO Stage 3과 결합하고 환경 변수 PYTORCH_ALLOC_CONF=expandable_segments:True를 사용하여 단편화를 줄입니다.
  • 2D Parallelism: GPU 수에 따라 sp_size와 dp_shard_size를 균형 맞춰 최대 시퀀스 길이 또는 높은 처리량 중 하나를 최적화합니다.
  • Additional Kernels: 손실 계산 및 대규모 행렬 연산 중 작업 메모리를 추가로 줄이기 위해 Liger-Kernel의 FusedLinearCrossEntropy와 TiledMLP를 사용합니다.

Sources