# Параллелизм контекста

[Контекстная параллельность](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism) распределяет входную последовательность между GPU, и каждый GPU обрабатывает свой срез. Это снижает потребление памяти на каждом GPU, что особенно полезно для длинных последовательностей, таких как изображения высокого разрешения и видео. Для сравнения с другими стратегиями см. [Распределённый инференс](../training/distributed_inference).

Контекстный параллелизм выполняется внутри бэкенда внимания, и поддерживается только некоторыми из них. Если ваш бэкенд не поддерживает контекстный параллелизм, `enable_parallelism` вызывает ошибку, в которой перечислены совместимые бэкенды. Переключить бэкенд можно с помощью [set_attention_backend()](/docs/diffusers/main/ru/api/models/overview#diffusers.ModelMixin.set_attention_backend), а полный список приведён в разделе [доступные бэкенды](./attention_backends#available-backends).

Настройте каждый из вариантов ниже с помощью [ContextParallelConfig](/docs/diffusers/main/ru/api/parallel#diffusers.ContextParallelConfig).

| Вариант | Конфигурация | Требования | Когда использовать |
|---|---|---|---|
| [Ring Attention](#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) передаёт блоки ключей (K) и значений (V) между GPU, организованными в кольцо. Каждый GPU вычисляет внимание по своему блоку K/V, а затем передаёт этот блок следующему GPU. После одного полного цикла каждый GPU обработал K/V всех токенов, при этом ни один GPU не хранит полную последовательность.

Установите `ring_degree` равным количеству GPU и передайте конфигурацию в метод `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
```

Чтобы попробовать другие варианты, замените вызов `enable_parallelism` в `ring_attention.py` на тот, что приведён в их разделе.

## 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 для перераспределения heads и tokens, затем на полученном результате запускается Ring Attention, после чего all-to-all выполняется в обратном направлении для восстановления исходной раскладки. Ulysses параллелизует по attention heads, а Ring — по последовательности, поэтому вместе они позволяют использовать больше GPU, чем Ulysses в одиночку.

Используйте Unified Attention, если у вас есть как минимум 4 GPU. Установите значения `ulysses_degree` и `ring_degree` больше 1. GPU образуют двумерную сетку из `ulysses_degree * ring_degree` устройств, при этом каждый GPU входит в одну группу Ulysses и одну группу Ring.

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

Бенчмарк ниже сравнивает Ulysses, Ring и Unified Attention на 4 GPU H100 с помощью [этого скрипта](https://github.com/huggingface/diffusers/pull/12693#issuecomment-3694727532). Для каждого изображения используется 50 шагов инференса.

| CP Backend         | Время / Изображение (мс) | Шагов / сек | Пиковая память (ГБ) |
|--------------------|------------------|-------------|------------------|
| 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]
> Добавьте бэкенд `gloo` для тензоров CPU в [init_process_group](https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group). Это позволяет избежать принудительной синхронизации CUDA, вызываемой передачей данных с хоста на устройство и с устройства на хост.

```py
import torch.distributed as dist

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

Бенчмарк ниже запускает FLUX.1-dev на 4 GPU L20 с помощью [этого скрипта](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999).

| CP Backend         | Время / Итерация (мс) | Шагов / сек | Пиковая память (ГБ) | Форма (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) поддерживает любую длину последовательности. Он дополняет локальные K/V на каждом GPU до самой длинной локальной длины, выполняет all-gather для дополненных K/V и вырезает каждый блок обратно до его истинной длины перед вычислением внимания.

Его нельзя использовать вместе с 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 также выигрывает от использования бэкенда `gloo`, показанного в [Ulysses Anything Attention](#ulysses-anything-attention).

> [!NOTE]
> Ring Anything Attention поддерживает только вывод, и `attn_mask` должен быть `None`.

Бенчмарк ниже запускает FLUX.1-dev на 4 GPU RTX 4090 (48 ГБ).

| CP Backend         | Время / Итерация (мс) | Шагов / сек | Пиковая память (ГБ) | Форма (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/main/ru/api/parallel#diffusers.ContextParallelConfig) в аргумент `parallel_config` метода [from_pretrained()](/docs/diffusers/main/ru/api/models/overview#diffusers.ModelMixin.from_pretrained). Сначала инициализируйте группу процессов, так как `from_pretrained` вызывает ошибку, если `torch.distributed` не инициализирован. Запустите скрипт с помощью `torchrun`, как в примере [Ring Attention](#ring-attention).

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