컨텍스트 병렬성

Context parallelism은 입력 시퀀스를 여러 GPU에 분할하여 각 GPU가 자신의 슬라이스를 처리합니다. 이를 통해 GPU당 메모리 사용량이 줄어들며, 고해상도 이미지나 비디오와 같은 긴 시퀀스에서 가장 큰 효과를 발휘합니다. 다른 전략과의 비교는 Distributed inference을 참고하세요.

Context parallelism은 attention backend 내부에서 실행되며, 일부 backend만 이를 지원합니다. 사용 중인 backend가 지원하지 않는 경우, enable_parallelism는 호환되는 backend 목록을 나열하는 오류를 발생시킵니다. set_attention_backend()를 사용하여 backend를 전환할 수 있으며, 전체 목록은 사용 가능한 backend를 참고하세요.

아래의 모든 변형을 ContextParallelConfig로 구성합니다.

Ring Attention

Ring Attention은 링으로 연결된 GPU 사이에서 키(K)와 값(V) 블록을 전달합니다. 각 GPU는 자신이 보유하고 있는 K/V 블록에 대해 어텐션을 계산한 후, 해당 블록을 다음 GPU로 전달합니다. 한 바퀴를 완전히 돌면 모든 GPU가 모든 토큰의 K/V를 확인하게 되며, 어떤 GPU도 전체 시퀀스를 보유하지 않습니다.

ring_degree를 GPU 수로 설정하고, 해당 설정을 transformer의 enable_parallelism() 메서드에 전달합니다.

import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, ContextParallelConfig

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    # torchrun sets LOCAL_RANK to this process's GPU index on its node
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)
    return device

def main():
    device = setup_distributed()
    world_size = dist.get_world_size()

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev", dtype=torch.bfloat16
    ).to(device)
    pipeline.transformer.set_attention_backend("_native_cudnn")

    cp_config = ContextParallelConfig(ring_degree=world_size)
    pipeline.transformer.enable_parallelism(config=cp_config)

    prompt = """
    cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
    highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
    """

    # Must specify generator so all ranks start with same latents (or pass your own)
    generator = torch.Generator().manual_seed(42)
    image = pipeline(
        prompt,
        guidance_scale=3.5,
        num_inference_steps=50,
        generator=generator,
    ).images[0]

    if dist.get_rank() == 0:
        image.save("output.png")

    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

스크립트를 ring_attention.py로 저장한 후 torchrun 또는 다른 PyTorch 호환 런처로 실행합니다. --nproc-per-node를 GPU 수로 설정합니다.

torchrun --nproc-per-node 2 ring_attention.py

다른 변형을 사용해 보려면 ring_attention.py의 enable_parallelism 호출을 해당 섹션의 호출로 교체하세요.

Ulysses Attention

Ulysses Attention uses all-to-all communication, where every GPU exchanges data with every other GPU. The first all-to-all regroups the data so each GPU holds every token for a subset of attention heads. Each GPU computes attention for its heads, and a second all-to-all regroups the results by token for the next layer.

ulysses_degree를 GPU 수로 설정하고, 해당 설정을 enable_parallelism()에 전달합니다.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

통합 어텐션

Unified Sequence Parallelism은 Ring Attention과 Ulysses Attention을 결합합니다. all-to-all을 실행하여 head와 token을 재분배한 후, 그 결과에 Ring Attention을 실행하고, 다시 all-to-all을 역방향으로 실행하여 원래의 레이아웃을 복원합니다. Ulysses는 attention head를 따라 병렬화하고 Ring은 시퀀스를 따라 병렬화하므로, 두 기법을 함께 사용하면 Ulysses 단독 사용 시보다 더 많은 GPU를 활용할 수 있습니다.

GPU가 4개 이상인 경우 Unified Attention을 사용하세요. ulysses_degree와 ring_degree를 모두 1보다 큰 값으로 설정합니다. GPU는 ulysses_degree * ring_degree개의 장치로 구성된 2D 그리드를 형성하며, 각 GPU는 하나의 Ulysses 그룹과 하나의 Ring 그룹에 속합니다.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ring_degree=2))

아래 벤치마크는 이 스크립트를 사용하여 4개의 H100 GPU에서 Ulysses, Ring, Unified Attention을 비교합니다. 각 이미지는 50개의 추론 스텝을 사용합니다.

Ulysses는 가장 높은 처리량을 제공하지만, 모델의 어텐션 헤드 수보다 많은 GPU를 사용할 수 없습니다. Unified Attention은 이러한 제한을 제거합니다.

Ulysses Anything Attention

Ulysses Attention은 시퀀스 길이와 어텐션 헤드 수가 ulysses_degree로 나누어떨어져야 합니다. Ulysses Anything Attention은 이 두 가지 요구사항을 모두 제거합니다. Unified Attention과 함께 사용할 수 없으므로 ring_degree는 1이어야 합니다.

ulysses_degree를 1보다 큰 값으로 설정하고 ulysses_anything=True를 지정하세요.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))

init_process_group에 CPU 텐서를 위한 gloo 백엔드를 추가합니다. 이를 통해 호스트-디바이스 및 디바이스-호스트 전송으로 인한 강제 CUDA 동기화를 피할 수 있습니다.

import torch.distributed as dist

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

아래 벤치마크는 이 스크립트를 사용하여 4개의 L20 GPU에서 FLUX.1-dev를 실행합니다.

1008x1008에서는 시퀀스 길이가 균등하게 나누어지지 않으므로 Ulysses Anything Attention만 실행됩니다. 1024x1024에서는 표준 Ulysses Attention과 동일한 성능을 보입니다.

Ring Anything Attention

Ring Attention은 시퀀스 길이가 ring_degree로 나누어떨어져야 합니다. Ring Anything Attention은 모든 시퀀스 길이를 지원합니다. 각 GPU의 로컬 K/V를 가장 긴 로컬 길이로 패딩하고, 패딩된 K/V를 all-gather한 뒤, attention을 계산하기 전에 각 블록을 원래 길이로 슬라이싱합니다.

Unified Attention과 함께 사용할 수 없으므로 ulysses_degree는 1이어야 합니다. ring_degree를 1보다 큰 값으로 설정하고 ring_anything=True를 지정하세요.

pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2, ring_anything=True))

Ring Anything Attention은 또한 Ulysses Anything Attention에 표시된 gloo 백엔드에서도 이점을 얻습니다.

Ring Anything Attention은 추론만 지원하며, attn_mask는 None이어야 합니다.

아래 벤치마크는 4개의 RTX 4090 (48GB) GPU에서 FLUX.1-dev를 실행합니다.

Ring Anything Attention은 1008x1008에서 실행되며 표준 Ring Attention과 거의 동일한 성능을 보입니다. 동기와 트레이드오프에 대해서는 이 댓글을 참고하세요.

…

enable_parallelism를 로드 후에 호출하는 대신, ContextParallelConfig를 from_pretrained()의 parallel_config 인자로 전달하세요. from_pretrained는 torch.distributed가 초기화되지 않은 경우 오류를 발생시키므로, 먼저 프로세스 그룹을 초기화해야 합니다. Ring Attention 예제와 마찬가지로 스크립트를 torchrun으로 실행하세요.

import os
import torch
from torch import distributed as dist
from diffusers import AutoModel, ContextParallelConfig, DiffusionPipeline

CKPT_ID = "black-forest-labs/FLUX.1-dev"

dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)

cp_config = ContextParallelConfig(ring_degree=dist.get_world_size())
transformer = AutoModel.from_pretrained(
    CKPT_ID,
    subfolder="transformer",
    dtype=torch.bfloat16,
    parallel_config=cp_config
)

pipeline = DiffusionPipeline.from_pretrained(
    CKPT_ID, transformer=transformer, dtype=torch.bfloat16,
).to(device)