diffusers/main/ru/optimization/context_parallelism.md

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

Контекстная параллельность распределяет входную последовательность между GPU, и каждый GPU обрабатывает свой срез. Это снижает потребление памяти на каждом GPU, что особенно полезно для длинных последовательностей, таких как изображения высокого разрешения и видео. Для сравнения с другими стратегиями см. Распределённый инференс.

Контекстный параллелизм выполняется внутри бэкенда внимания, и поддерживается только некоторыми из них. Если ваш бэкенд не поддерживает контекстный параллелизм, enable_parallelism вызывает ошибку, в которой перечислены совместимые бэкенды. Переключить бэкенд можно с помощью set_attention_backend(), а полный список приведён в разделе доступные бэкенды.

Настройте каждый из вариантов ниже с помощью ContextParallelConfig.

Вариант	Конфигурация	Требования	Когда использовать
Ring Attention	ring_degree > 1	длина последовательности должна делиться нацело на ring_degree	последовательности длинные, а пропускная способность межсоединений ограничена
Ulysses Attention	ulysses_degree > 1	длина последовательности и количество головок внимания должны делиться нацело на ulysses_degree	ваши GPU имеют высокоскоростное соединение
Единое внимание	ring_degree > 1 и ulysses_degree > 1	не менее 4 GPU	у вас больше устройств, чем может использовать Ulysses Attention в одиночку
Ulysses Anything Attention	ulysses_degree > 1, ulysses_anything=True	ring_degree должен быть равен 1	длина последовательности или количество головок не делятся нацело
Ring Anything Attention	ring_degree > 1, ring_anything=True	ulysses_degree должен быть равен 1, только для инференса, attn_mask должен быть None	длина последовательности не делится нацело

Ring Attention

Ring Attention передаёт блоки ключей (K) и значений (V) между GPU, организованными в кольцо. Каждый GPU вычисляет внимание по своему блоку K/V, а затем передаёт этот блок следующему GPU. После одного полного цикла каждый GPU обработал K/V всех токенов, при этом ни один GPU не хранит полную последовательность.

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

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

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 для перераспределения 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.

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

Бенчмарк ниже сравнивает Ulysses, Ring и Unified Attention на 4 GPU H100 с помощью этого скрипта. Для каждого изображения используется 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 снимает оба этих требования. Его нельзя использовать вместе с Unified Attention, поэтому ring_degree должен быть равен 1.

Установите ulysses_degree в значение больше 1 и ulysses_anything=True.

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

Добавьте бэкенд gloo для тензоров CPU в init_process_group. Это позволяет избежать принудительной синхронизации CUDA, вызываемой передачей данных с хоста на устройство и с устройства на хост.

import torch.distributed as dist

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

Бенчмарк ниже запускает FLUX.1-dev на 4 GPU L20 с помощью этого скрипта.

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 поддерживает любую длину последовательности. Он дополняет локальные K/V на каждом GPU до самой длинной локальной длины, выполняет all-gather для дополненных K/V и вырезает каждый блок обратно до его истинной длины перед вычислением внимания.

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

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. Подробнее о мотивации и компромиссах см. в этом комментарии.

Включить параллелизм контекста при загрузке

Вместо вызова enable_parallelism после загрузки передайте ContextParallelConfig в аргумент parallel_config метода from_pretrained(). Сначала инициализируйте группу процессов, так как from_pretrained вызывает ошибку, если torch.distributed не инициализирован. Запустите скрипт с помощью torchrun, как в примере Ring Attention.

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)