🤗 accelerate의 시퀀스 병렬

이 가이드에서는 🤗accelerate에서 시퀀스 병렬성을 사용하는 기본 사항을 다룹니다.

매우 관련이 있는 Context Parallelism도 참고하세요.

왜 시퀀스 병렬성인가?

대형 언어 모델, 그리고 최근에는 추론 모델의 등장으로 시퀀스 길이가 빠르게 증가하고 있습니다. 여기에 attention의 2차 메모리 복잡도가 더해지면서, 긴 시퀀스로 모델을 학습하는 더 효율적인 방식에 대한 필요성이 커졌습니다. 시퀀스 길이가 128k일 때, vanilla attention 구현을 가정하면 bf16 정밀도에서 attention matrix의 메모리 요구량은 128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads입니다. 물론, 이러한 attention 가중치를 물리적으로 생성하지 않는 flash attention나 SDPA를 사용하면 메모리 사용량이 크게 줄지만, 메모리 요구량의 증가 폭은 여전히 상당합니다.

Ulysses Sequence parallelism은 attention 계산의 입력을 시퀀스 차원별로 샤딩하고, 각 GPU에서 attention head의 일부만 사용하여 attention을 정상적으로 계산할 수 있게 합니다. 이를 통해 몇 가지 추가 도구를 활용하여 15M+ 시퀀스 길이에 스케일링하여 긴 시퀀스를 가진 모델을 학습할 수 있습니다. Ulysses SP에 TiledMLP, Liger-Kernel, Activation checkpoint offload to cpu 및 기타 몇 가지 트릭을 결합하는 방법에 대해서는 다음 논문을 참고하세요: Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences.

Ulysses SP와 FSDP CP의 차이점은 무엇인가요?

Context Parallelism 문서에서는 시퀀스 차원을 슬라이스하되 헤드 차원 슬라이싱 대신 Ring Attention을 사용하는 또 다른 기술인 Context Parallelism의 배포에 대해 배울 수 있습니다.

다음 글에서는 두 기술의 차이를 매우 상세하게 설명합니다:

• https://insujang.github.io/2024-01-11/tensor-parallelism-and-sequence-parallelism-detailed-analysis/
• https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention

한 기사의 내용을 참고하여 간단히 요약하면 다음과 같습니다:

• Ulysses SP는 통신 오버헤드가 상대적으로 낮지만, Attention Heads의 수에 의해 제한되므로 네트워크 토폴로지(단일 리플리카에 참여하는 GPU 수로 attention heads의 수가 나누어져야 함)에 대해 특정 요구 사항이 있습니다. All-to-all 통신은 지연 시간에 민감하며 Deepspeed가 필요합니다.
• FSDP CP Ring-Attention의 P2P 링 통신은 aforementioned 가분성 요구사항이 없지만, 통신량이 더 높습니다.

마지막으로, USP: A Unified Sequence Parallelism Approach for Long Context Generative AI 논문에서 설명한 대로 SP + CP를 결합하여 더 긴 시퀀스 길이를 지원하는 것이 가능해야 합니다. 다만, 이 기능은 아직 🤗accelerate에 통합되지 않았습니다.

지원되는 시퀀스 병렬성 백엔드

현재 유일한 시퀀스 병렬성 백엔드는 deepspeed입니다. 이는 Arctic Long Sequence Training 기술의 일부인 현대화된 Ulysses SP에서 파생되었습니다. 직접 자체 코드에 통합하려는 경우를 위해 튜토리얼도 제공됩니다.

시퀀스 병렬성(sequence parallelism)을 사용하는 방법

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

+# Example: 4 GPUs with sp_size=4, dp_shard_size=1
+# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
+     sp_backend="deepspeed",
+     sp_size=4,
+     dp_shard_size=1,  # Explicit: no data parallelism
+     sp_handler=DeepSpeedSequenceParallelConfig(
+         sp_seq_length_is_variable: true,
+         sp_attn_implementation="sdpa",
+     ),
+ )

accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

🤗accelerate의 다른 모든 기능과 마찬가지로, accelerate launch에 해당 플래그를 전달하여 시퀀스 병렬성도 활성화할 수 있습니다. 이 경우에도 차이가 없습니다:

accelerate launch --parallelism-config-sp-size 8  ...

sp_size 및 기타 구성을 accelerate config 명령에서 설정할 수도 있습니다. 이렇게 하면 accelerate 구성 파일에 저장되어 스크립트를 실행할 때마다 매번 전달할 필요가 없습니다.

sequence parallelism은 data parallelism과 결합됩니다. 추가 GPU가 필요하지 않습니다. 따라서 8개의 GPU가 있는 경우 --parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8을 사용할 수 있습니다. 또는 ParallelismConfig 클래스를 사용하여 프로그래밍 방식으로 설정할 수도 있습니다.

중요: dp_replicate_size × dp_shard_size × sp_size = num_processes를 반드시 확인해야 합니다. 예를 들어, GPU 8개와 sp_size=8를 사용하는 경우 dp_shard_size=1이 필요합니다(1 × 1 × 8 = 8). GPU 4개와 sp_size=2를 사용하는 경우 2D 병렬화를 위해 dp_shard_size=2를 사용할 수 있습니다(1 × 2 × 2 = 4).

ALST/Ulysses SP 백엔드 구성

ALST/UlyssesSP은 이 논문에서 설명된 대로 어텐션 헤드 병렬성을 사용하여 시퀀스 병렬성을 구현합니다. 단순화를 위해 시퀀스 병렬성의 개념과 설정을 재사용하며, 사용자 관점에서는 동일합니다. 여러 GPU가 단일 배치 처리에 사용됩니다.

ALST가 어떤 가능성을 열어주었는지 감을 잡기 위해, 이 기술 덕분에 단일 H100 GPU에서 500K 토큰, 단일 노드에서 3.7M 토큰, 그리고 Llama-8B 모델의 경우 노드 4개만으로 15M 토큰까지 bf16으로 학습할 수 있었습니다. HF Accelerate의 이 기능은 ALST의 3가지 구성 요소 중 1개만 활성화하므로, 달성 가능한 시퀀스 길이는 더 짧아집니다. ALST의 전체 성능을 발휘하려면 TiledMLP, CPU로의 Activation checkpoint offload 및 기타 몇 가지 설정을 활성화해야 합니다. 자세한 내용은 이 튜토리얼을 참고하세요.

deepspeed 백엔드를 구성하려면:

# Example: 4 GPUs with sp_size=4, dp_shard_size=1
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_shard_size=1,  # Explicit: no data parallelism
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=256,
        sp_seq_length_is_variable=True,
        sp_attn_implementation="sdpa",
    ),
)
accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

• sp_backend: 여기서는 deepspeed로 설정
• sp_size는 시퀀스 병렬도의 크기를 나타내며, 위의 예시에서는 4이므로 단일 배치 처리에 4개의 GPU가 사용됩니다(동일한 GPU에서 DP=4를 수행하는 동안).
• sp_seq_length와 sp_seq_length_is_variable은 시퀀스 길이를 처리하는 데 사용됩니다. sp_seq_length_is_variable=True인 경우 백엔드는 배치 간에 변경될 수 있는 시퀀스 길이를 처리하며, 이때 sp_seq_length 값은 시퀀스 병렬도(sequence parallel degree)로 나누어 떨어지는 값으로 설정하거나 아예 설정하지 않을 수 있습니다. 이 경우 모든 forward에서 시퀀스 변수가 입력으로부터 유도됩니다. False인 경우 seq_length는 배치의 시퀀스 길이 차원과 일치해야 하며, 이는 항상 동일하도록 패딩되어야 합니다. 기본값은 True입니다.
• sp_attn_implementation은 sdpa, flash_attention_2, flash_attention_3 중 하나입니다. 이 시퀀스 병렬 구현은 attention_mask 대신 position_ids를 사용하므로, eager는 position_ids를 지원하는 기능이 추가될 때까지 여기서 작동하지 않습니다. 또한, sdpa는 여러 샘플이 하나로 결합된 경우를 올바르게 처리하지 못하며, 전체 샘플을 하나로 취급하여 어텐션을 수행합니다. 샘플이 결합되지 않은 경우 sdpa는 올바르게 작동합니다. 따라서 항상 작동하는 Flash Attention이 이상적인 선택입니다.

DeepSpeedSequenceParallelConfig 객체에 이러한 값을 설정하는 대신, 동일한 결과를 얻기 위해 환경 변수를 사용할 수도 있습니다. 아래 목록의 끝과 대응하는 환경 변수는 다음과 같습니다.

• PARALLELISM_CONFIG_SP_BACKEND
• PARALLELISM_CONFIG_SP_SEQ_LENGTH
• PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE
• PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION

코드에서 전달하지 않은 경우, sp_size는 --parallelism_config_sp_size CLI 인자를 통해 설정할 수 있습니다. 다른 인자들도 마찬가지입니다. 또한, accelerate config 파일 스타일의 설정도 사용할 수 있습니다. 예를 들어, GPU 2개의 경우:

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 2
parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1  # Must satisfy: 1 × 1 × 2 = 2 num_processes
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: sdpa

앞서 언급했듯이 Ulysses sequence parallelism은 일반적으로 data parallelism과 함께 사용됩니다. 동일한 rank가 고유한 데이터 스트림을 공급하는 데 사용되고, 동시에 Ulysses Sequence Parallelism을 수행합니다. 하지만 다음과 같이 리플리카를 생성할 수도 있습니다:

# Example: 4 GPUs with 2D parallelism (SP=2, DP=2)
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 2 × 1 × 2 = 4 GPUs
parallelism_config = ParallelismConfig(
    dp_replicate_size=2,
    dp_shard_size=1,  # Explicit: no sharding within replicas
    sp_size=2,
    sp_backend="deepspeed",
    sp_handler=DeepSpeedSequenceParallelConfig(...),
)

여기서는 4개의 GPU를 2개의 sequence parallelism 레플리카로 사용합니다. 데이터 병렬성은 Deepspeed-ZeRO가 담당합니다.

UlyssesSPDataLoaderAdapter 내부에는 많은 마법이 숨겨져 있습니다. 이 클래스는 백그라운드에서 동작하며 원래의 DataLoader 객체를 래핑하지만, 문제가 발생할 경우 이를 인지하고 있어야 합니다. 또한, 배치에 참여하는 각 rank로 샤딩이 수행되기 전에 배치 사전에 올바른 shift_labels를 자동으로 주입합니다.

이제 ALST/UlyssesSP를 사용하기 위해 남은 마지막 단계는, 그래디언트를 올바르게 계산하기 위해 다른 가변 all_gather를 사용하여 랭크 간 손실을 집계하는 것입니다. 다음 코드는 이를 수행하며, -100으로 마스킹된 토큰을 제외하여 올바른 평균을 얻습니다:

sp_size = parallelism_config.sp_size if parallelism_config is not None else 1
if sp_size > 1:
    sp_group = accelerator.torch_device_mesh["sp"].get_group()
    sp_world_size = parallelism_config.sp_size

# Normal training loop
for iter, batch in enumerate(dl):
    optimizer.zero_grad()

    batch = move_to_device(batch, model.device)

    # The model automatically receives shift_labels via **kwargs and uses it for loss computation.
    # Both standard transformers models and Liger-patched models handle this correctly.
    outputs = model(**batch)
    loss = outputs.loss
    shift_labels = batch["shift_labels"]

    if sp_size > 1:
        # differentiable weighted per-shard-loss aggregation across ranks
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (shift_labels != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(
            good_tokens, group=sp_group
        )
        # Skip ranks with zero valid tokens to avoid NaN contamination (NaN * 0 = NaN)
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

    if rank == 0: accelerator.print(f"{iter}: {loss=}")
    accelerator.log(dict(train_loss=loss, step=iter))

    accelerator.backward(loss)
    optimizer.step()

shift_labels가 배치에 존재하는 경우 모델이 이를 자동으로 처리한다는 점에 유의하세요. 모델의 forward pass는 **kwargs를 통해 shift_labels를 수신하여 손실 함수로 전달하고, 손실 함수는 시퀀스 병렬성에 대해 손실을 올바르게 계산합니다. Liger Kernel을 사용하는 경우에도 shift_labels를 원활하게 처리하며, 매우 메모리 효율적인 방식으로 손실을 계산합니다. Liger는 융합 연산(예: 메모리에 전체 logits 텐서를 절대 생성하지 않는 융합 logit-loss 계산)을 사용하여 GPU 메모리를 확보하므로 긴 시퀀스 길이에 대해 강력히 권장됩니다.

HF Accelerate가 백그라운드에서 수행한 작업을 확인하려면 이 전체 통합 튜토리얼을 읽어보세요.

ALST/UlyssesSP가 활성화된 Accelerate 학습 루프의 예제는 examples/alst_ulysses_sequence_parallelism를 참고하세요.

[!경고]

이 API는 비교적 신규이며 아직 실험 단계에 있습니다. 안정적인 API를 제공하기 위해 노력하고 있지만, 공개 API의 일부 세부 사항은 향후 변경될 수 있습니다.

이것은 Deepspeed 백엔드이므로 일반적인 Deepspeed 구성이 적용됩니다. 따라서 시퀀스 병렬성을 옵티마이저 상태 및/또는 가중치 오프로딩과 결합하여 더 많은 GPU 메모리를 확보하고 더 긴 시퀀스 길이를 사용할 수 있습니다. 이 기술은 DeepSpeed ZeRO stage 2 및 3과 함께 작동하도록 테스트되었습니다.