Ulysses 시퀀스 병렬성

Ulysses sequence parallelism(SP)은 매우 긴 시퀀스를 여러 GPU에 분할하여 학습합니다. attention을 올바르게 계산하기 위해, all-to-all collective이 sharding 차원을 sequence에서 attention heads로 교환합니다. 각 GPU는 전체 sequence를 보유하며, heads의 부분 집합에 대해 로컬에서 attention을 계산합니다. 두 번째 all-to-all은 sequence-sharded 레이아웃으로 복귀하여, forward pass의 나머지 부분이 각 chunk에서 로컬로 계속 진행되도록 합니다.

                        GPU 0                       GPU 1
                   ┌───────────────┐          ┌───────────────┐
  forward          │ tokens 0..N/2 │          │ tokens N/2..N │  ← each GPU holds half the sequence
  (seq-sharded)    │  all H heads  │          │  all H heads  │
                   └───────┬───────┘          └───────┬───────┘
                           └───────── all-to-all ──────┘
                   ┌───────────────┐          ┌───────────────┐
  attention        │ all N tokens  │          │ all N tokens  │  ← now each GPU has the full sequence
  (head-sharded)   │ heads 0..H/2  │          │ heads H/2..H  │  ← but only half the heads
                   └───────┬───────┘          └───────┬───────┘
                           └───────── all-to-all ──────┘
                   ┌───────────────┐          ┌───────────────┐
  forward          │ tokens 0..N/2 │          │ tokens N/2..N │  ← back to seq-sharded
  (seq-sharded)    │  all H heads  │          │  all H heads  │
                   └───────────────┘          └───────────────┘

이 가이드는 ALST(Arctic Long Sequence Training)의 Ulysses sequence parallelism 구성 요소를 다룹니다. 전체 ALST 시스템에는 Transformers에서 사용할 수 없는 TiledMLP와 activation checkpoint offloading도 포함되어 있습니다. 전체 시스템에 대해서는 DeepSpeed ALST 튜토리얼을 참고하세요.

구성

Sequence parallelism은 Accelerate v1.12.0 이상과 GPU 2개 이상이 필요합니다. Accelerate의 ParallelismConfig에서 sequence parallelism을 구성하고, parallelism_config 또는 Accelerate config file에 전달하세요.

parallelism_config

Accelerate 설정 파일

다음 필드는 시퀀스 병렬성을 구성하는 데 중요합니다.

Trainer는 DataLoader 샤딩, position_ids 생성, 레이블 시프트, SP 랭크 간 손실 집계 등을 자동으로 처리합니다. 커스텀 학습 루프를 작성하는 경우, 대신 Accelerate의 Sequence Parallelism 가이드를 참고하세요.

Ulysses 시퀀스 병렬성을 사용하려면 sp_backend을 "deepspeed"로 설정해야 합니다.
sp_size는 단일 시퀀스를 병렬로 처리하는 GPU의 수입니다. 텐서 병렬성에서는 모든 랭크가 동일한 데이터를 수신하는 반면, 각 SP 랭크는 DataLoader에서 고유한 데이터 스트림을 수신합니다. 유효한 dp_world_size = world_size / sp_size이므로, 4개의 GPU와 sp_size=4를 사용하는 경우 배치 크기 계산에 사용되는 dp_world_size=1이 됩니다. 시퀀스는 또한 sp_size의 배수로 패딩되어야 합니다. 데이터 콜레이터에서 pad_to_multiple_of를 이에 맞게 설정하세요.

Attention head의 수는 sp_size로 나누어떨어져야 합니다. 32개의 head를 가진 모델은 sp_size가 1, 2, 4, 8, 16 또는 32인 것을 지원합니다.
```
from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=sp_size,
)
```
sp_seq_length_is_variable은 가변 시퀀스 길이를 처리하는 방식을 제어합니다. 배치 간 길이가 달라지는 경우 True(권장)로 설정하세요. 모든 시퀀스가 sp_seq_length로 지정된 고정 길이로 패딩되는 경우 False로 설정하세요.
sp_attn_implementation은 attention 백엔드를 설정합니다. 지원되는 값은 "sdpa", "flash_attention_2", 또는 "flash_attention_3"입니다. 특히 배치에 여러 샘플을 패킹하는 경우 FlashAttention을 권장합니다. 샘플이 패킹된 경우 SDPA는 샘플 경계를 넘어 잘못 attend할 수 있습니다. 메모리 및 확장성 문제로 4D attention_mask가 폐기되므로 Eager attention은 지원되지 않습니다.

데이터 병렬성과의 결합

Sequence parallelism과 data parallelism은 동일한 GPU를 사용하며, SP는 추가 하드웨어가 필요하지 않습니다. 두 가지를 모두 실행하려면 dp_replicate_size × dp_shard_size × sp_size가 GPU 총 개수와 같도록 dp_replicate_size 또는 dp_shard_size를 설정하세요.

예를 들어, GPU 8개와 sp_size=4를 사용하는 경우 dp_replicate_size=2로 설정합니다 (2 × 1 × 4 = 8).

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_replicate_size=2,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",
    ),
)

다음 단계

FSDP2의 컨텍스트 병렬성(cp_size)은 Ulysses(sp_size)와 별개입니다. 컨텍스트 병렬성으로 평가하기를 참조하세요.
Accelerate Sequence Parallelism 가이드에서는 Ulysses 구현을 더 깊이 다루며, 커스텀 학습 루프를 작성하는 방법을 보여줍니다.
DeepSpeed ALST 튜토리얼은 TiledMLP와 activation checkpoint offloading을 포함한 전체 ALST 시스템을 다룹니다.
병렬화 방법 가이드는 시퀀스 병렬화를 ZeRO와 같은 다른 전략과 결합하는 방법을 보여줍니다.
Ulysses Sequence Parallelism: Training with Million-Token Contexts 블로그 포스트는 Ulysses가 어떻게 작동하는지, 그리고 Accelerate, Trainer, SFTTrainer에 어떻게 통합되는지를 설명합니다.

Update on GitHub