Buckets:

hf-doc-build/doc / diffusers /v0.43.0 /ru /optimization /context_parallelism.md
mishig's picture
|
download
raw
16.5 kB
# Параллелизм контекста
[Контекстная параллельность](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=context_parallelism) распределяет входную последовательность между GPU, и каждый GPU обрабатывает свой срез. Это снижает потребление памяти на каждом GPU, что особенно полезно для длинных последовательностей, таких как изображения высокого разрешения и видео. Для сравнения с другими стратегиями см. [Распределённый инференс](../training/distributed_inference).
Контекстный параллелизм выполняется внутри бэкенда внимания, и поддерживается только некоторыми из них. Если ваш бэкенд не поддерживает контекстный параллелизм, `enable_parallelism` вызывает ошибку, в которой перечислены совместимые бэкенды. Переключить бэкенд можно с помощью [set_attention_backend()](/docs/diffusers/v0.43.0/ru/api/models/overview#diffusers.ModelMixin.set_attention_backend), а полный список приведён в разделе [доступные бэкенды](./attention_backends#available-backends).
Настройте каждый из вариантов ниже с помощью [ContextParallelConfig](/docs/diffusers/v0.43.0/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/v0.43.0/ru/api/parallel#diffusers.ContextParallelConfig) в аргумент `parallel_config` метода [from_pretrained()](/docs/diffusers/v0.43.0/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)
```

Xet Storage Details

Size:
16.5 kB
·
Xet hash:
3045c3960c5037edd0b50e2ea76bc0e8d2f37bde5ba7bad0be793f4b9acb2d6a

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.