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

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

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

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

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 шагов инференса.

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 с помощью этого скрипта.

При размере 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 ГБ).

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)