Ulysses 시퀀스 병렬

Ulysses 시퀀스 병렬 (Ulysses sequence parallelism)

Ulysses sequence parallelism (SP)은 매우 긴 시퀀스를 여러 GPU로 분할해서 학습하는 기법이에요. attention을 올바르게 계산하려면 all-to-all 집단 통신(collective)이 샤딩 차원을 시퀀스에서 attention 헤드로 바꿔줘요. 그러면 각 GPU가 전체 시퀀스를 갖게 되고, 헤드의 일부에 대해 로컬로 attention을 계산해요. 두 번째 all-to-all이 다시 시퀀스-샤딩 레이아웃으로 되돌려서, 나머지 forward pass가 각 청크에서 로컬로 계속되게 해줘요.

출처: 문서

본문

                        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  │
                   └───────────────┘          └───────────────┘

[!NOTE] 이 가이드는 ALST (Arctic Long Sequence Training)의 Ulysses sequence parallelism 구성 요소를 다룹니다. 전체 ALST 시스템에는 TiledMLP와 activation checkpoint offloading도 포함되지만, 이들은 Transformers에서는 사용할 수 없습니다. 전체 시스템은 DeepSpeed ALST tutorial을 참고해 주세요.

본문

구성하기

시퀀스 병렬은 Accelerate v1.12.0과 GPU 2개 이상이 필요해요. Accelerate의 ParallelismConfig에서 시퀀스 병렬을 구성하고, 이를 parallelism_config 또는 Accelerate config file에 전달해요.

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

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

training_args = TrainingArguments(
    ...,
    deepspeed="path/to/deepspeed_config.json",
    parallelism_config=parallelism_config,
)

Trainer 기반 스크립트로 accelerate launch를 실행해요.

accelerate launch --num_processes 4 train.py \
--output_dir output_dir \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 1

accelerate config 명령을 실행하고 하드웨어와 학습 설정에 대한 질문에 답해서 캐시에 default_config.yaml 파일을 만들어요.

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 4
parallelism_config:
  parallelism_config_sp_size: 4
  parallelism_config_dp_replicate_size: 1
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: flash_attention_2

Trainer 기반 스크립트로 accelerate launch를 실행해요.

accelerate launch --config_file alst_config.yaml train.py \
--output_dir output_dir \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 1

시퀀스 병렬을 구성할 때 중요한 필드는 다음과 같아요.

[!TIP] Trainer는 DataLoader 샤딩, position_ids 생성, 라벨 shift, SP rank 간 loss 집계를 자동으로 처리해요. 커스텀 학습 루프를 작성한다면 Accelerate Sequence Parallelism 가이드를 대신 참고해 주세요.

  • sp_backend는 Ulysses sequence parallelism을 사용하려면 "deepspeed"로 설정해야 해요.

  • sp_size는 하나의 시퀀스를 병렬로 처리하는 GPU 수예요. 각 SP rank는 DataLoader로부터 고유한 데이터 스트림을 받는데, 모든 rank가 동일한 데이터를 받는 텐서 병렬과는 달라요. 유효 dp_world_size = world_size / sp_size이므로, GPU 4개에 sp_size=4라면 배치 크기 계산에서 dp_world_size=1이 돼요. 시퀀스도 sp_size의 배수로 패딩해야 해요. 데이터 콜레이터의 pad_to_multiple_of를 그에 맞게 설정하세요.

    [!WARNING] attention 헤드 수는 sp_size로 나누어 떨어져야 합니다. 헤드 32개 모델은 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가 샘플 경계를 가로질러 attention을 잘못 계산할 수 있어요. Eager attention은 4D attention_mask가 메모리와 스케일링 이유로 버려지기 때문에 지원되지 않아요.

데이터 병렬과 결합하기

시퀀스 병렬과 데이터 병렬은 동일한 GPU를 사용하므로, SP에 추가 하드웨어가 필요하진 않아요. 둘 다 실행하려면 dp_replicate_size 또는 dp_shard_size를 설정해서 dp_replicate_size × dp_shard_size × sp_size가 총 GPU 수와 같게 해요.

예를 들어 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",
    ),
)

더 알아보기 (Learn more)