# 컨텍스트 병렬성

[Context parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism)은 입력 시퀀스를 여러 GPU에 분할하여 각 GPU가 자신의 슬라이스를 처리합니다. 이를 통해 GPU당 메모리 사용량이 줄어들며, 고해상도 이미지나 비디오와 같은 긴 시퀀스에서 가장 큰 효과를 발휘합니다. 다른 전략과의 비교는 [Distributed inference](../training/distributed_inference)을 참고하세요.

Context parallelism은 attention backend 내부에서 실행되며, 일부 backend만 이를 지원합니다. 사용 중인 backend가 이를 지원하지 않으면 `enable_parallelism`가 호환되는 backend 목록을 나열하는 오류를 발생시킵니다. [set_attention_backend()](/docs/diffusers/v0.43.0/ko/api/models/overview#diffusers.ModelMixin.set_attention_backend)를 사용하여 backend를 전환할 수 있으며, 전체 목록은 [사용 가능한 backend](./attention_backends#available-backends)를 참고하세요.

아래의 모든 변형을 [ContextParallelConfig](/docs/diffusers/v0.43.0/ko/api/parallel#diffusers.ContextParallelConfig)로 구성합니다.

| 변형 | 설정 | 요구 사항 | 사용 시 |
|---|---|---|---|
| [링 어텐션](#ring-attention) | `ring_degree > 1` | 시퀀스 길이는 `ring_degree`로 나누어떨어져야 합니다 | 시퀀스가 길고 인터커넥트 대역폭이 제한되어 있습니다 |
| [Ulysses Attention](#ulysses-attention) | `ulysses_degree > 1` | 시퀀스 길이와 어텐션 헤드 수는 `ulysses_degree`로 나누어떨어져야 합니다 | GPU에 빠른 인터커넥트가 있습니다 |
| [통합 어텐션](#unified-attention) | `ring_degree > 1` 및 `ulysses_degree > 1` | 최소 4개의 GPU | Ulysses Attention 단독으로 사용할 수 있는 장치보다 더 많은 장치를 보유하고 있습니다 |
| [Ulysses Anything Attention](#ulysses-anything-attention) | `ulysses_degree > 1`, `ulysses_anything=True` | `ring_degree`은 1이어야 합니다 | 시퀀스 길이 또는 헤드 수가 균등하게 나누어지지 않는 경우 |
| [Ring Anything Attention](#ring-anything-attention) | `ring_degree > 1`, `ring_anything=True` | `ulysses_degree`는 1이어야 하며, 추론 전용이고, `attn_mask`는 `None`이어야 합니다 | 시퀀스 길이가 균등하게 나누어지지 않는 경우 |

## Ring Attention

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

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

```py
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](https://docs.pytorch.org/docs/stable/elastic/run.html) 또는 다른 PyTorch 호환 런처로 실행합니다. `--nproc-per-node`를 GPU 수로 설정합니다.

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

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

## Ulysses Attention

[Ulysses Attention](https://huggingface.co/papers/2309.14509) 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()`에 전달합니다.

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

## 통합 어텐션

[Unified Sequence Parallelism](https://huggingface.co/papers/2405.07719)은 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 그룹에 속합니다.

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

아래 벤치마크는 [이 스크립트](https://github.com/huggingface/diffusers/pull/12693#issuecomment-3694727532)를 사용하여 4개의 H100 GPU에서 Ulysses, Ring, Unified Attention을 비교합니다. 각 이미지는 50개의 추론 스텝을 사용합니다.

| CP 백엔드         | 시간 / 이미지 (ms) | 초당 스텝 | 피크 메모리 (GB) |
|--------------------|------------------|-------------|------------------|
| ulysses            | 6670.789         | 7.50        | 33.85            |
| 링               | 13076.492        | 3.82        | 56.02            |
| unified_balanced   | 11068.705        | 4.52        | 33.85            |

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

## Ulysses Anything Attention

Ulysses Attention은 시퀀스 길이와 어텐션 헤드 수가 `ulysses_degree`로 나누어떨어져야 합니다. [Ulysses Anything Attention](https://github.com/huggingface/diffusers/pull/12996)은 이 두 가지 요구사항을 모두 제거합니다. Unified Attention과 함께 사용할 수 없으므로 `ring_degree`는 1이어야 합니다.

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

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

> [!TIP]
> [init_process_group](https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group)에 CPU 텐서를 위한 `gloo` 백엔드를 추가합니다. 이를 통해 호스트-디바이스 및 디바이스-호스트 전송으로 인한 강제 CUDA 동기화를 피할 수 있습니다.

```py
import torch.distributed as dist

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

아래 벤치마크는 [이 스크립트](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999)를 사용하여 4개의 L20 GPU에서 FLUX.1-dev를 실행합니다.

| CP 백엔드         | 시간 / 반복 (ms) | 초당 스텝 | 피크 메모리 (GB) | Shape (HxW)|
|--------------------|------------------|-------------|------------------|------------|
| ulysses            |   281.07         |    3.56     |     37.11        | 1024x1024  |
| 링               |   351.34         |    2.85     |     37.01        | 1024x1024  |
| unified_balanced   |   324.37         |    3.08     |     37.16        | 1024x1024  |
| ulysses_anything   |   280.94         |    3.56     |     37.11        | 1024x1024  |
| ulysses            |   실패         |    실패   |     실패       | 1008x1008  |
| 링               |   실패         |    실패   |     실패       | 1008x1008  |
| unified_balanced   |   실패         |    실패   |     실패       | 1008x1008  |
| ulysses_anything   |   278.40         |    3.59     |     36.99        | 1008x1008  |

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

## Ring Anything Attention

Ring Attention은 시퀀스 길이가 `ring_degree`로 나누어떨어져야 합니다. [Ring Anything Attention](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4302195582)은 모든 시퀀스 길이를 지원합니다. 각 GPU의 로컬 K/V를 가장 긴 로컬 길이로 패딩하고, 패딩된 K/V를 all-gather한 뒤, attention을 계산하기 전에 각 블록을 원래 길이로 슬라이싱합니다.

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

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

Ring Anything Attention은 또한 [Ulysses Anything Attention](#ulysses-anything-attention)에 표시된 `gloo` 백엔드에서도 이점을 얻습니다.

> [!NOTE]
> Ring Anything Attention은 추론만 지원하며, `attn_mask`는 `None`이어야 합니다.

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

| CP 백엔드         | 시간 / 반복 (ms) | 초당 스텝 | 피크 메모리 (GB) | Shape (HxW)|
|--------------------|------------------|-------------|------------------|------------|
| ulysses            |   259.07         |    3.86     |     33.83        | 1024x1024  |
| 링               |   338.98         |    2.95     |     33.83        | 1024x1024  |
| unified_balanced   |   321.54         |    3.11     |     33.83        | 1024x1024  |
| ulysses_anything   |   259.07         |    3.86     |     33.83        | 1024x1024  |
| ring_anything      |   340.14         |    2.94     |     33.83        | 1024x1024  |
| ulysses            |   실패         |    실패   |     실패       | 1008x1008  |
| 링               |   실패         |    실패   |     실패       | 1008x1008  |
| unified_balanced   |   실패         |    실패   |     실패       | 1008x1008  |
| ulysses_anything   |   253.16         |    3.95     |     33.75        | 1008x1008  |
| ring_anything      |   335.57         |    2.98     |     33.75        | 1008x1008  |

Ring Anything Attention은 1008x1008에서 실행되며 표준 Ring Attention과 거의 동일한 성능을 보입니다. 동기와 트레이드오프에 대해서는 [이 댓글](https://github.com/huggingface/diffusers/pull/13545#issuecomment-4304104462)을 참고하세요.

## …

`enable_parallelism`를 로드 후에 호출하는 대신, [ContextParallelConfig](/docs/diffusers/v0.43.0/ko/api/parallel#diffusers.ContextParallelConfig)를 [from_pretrained()](/docs/diffusers/v0.43.0/ko/api/models/overview#diffusers.ModelMixin.from_pretrained)의 `parallel_config` 인자로 전달하세요. `from_pretrained`는 `torch.distributed`가 초기화되지 않은 경우 오류를 발생시키므로, 먼저 프로세스 그룹을 초기화해야 합니다. [Ring Attention](#ring-attention) 예제와 마찬가지로 스크립트를 `torchrun`으로 실행하세요.

```py
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)
```