Параллелизм данных

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

Вы можете настроить параллелизм данных с помощью Accelerate или напрямую с помощью PyTorch Distributed.

Accelerate

Accelerate настраивает процессы и назначает каждому из них GPU.

Установите Accelerate с помощью следующей команды.

uv pip install accelerate

Создайте PartialState в файле run_distributed.py, чтобы настроить распределённую среду. Переместите DiffusionPipeline на PartialState.device, чтобы назначить GPU каждому процессу.

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

В том же файле используйте split_between_processes в качестве контекстного менеджера для распределения промптов между процессами.

with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipeline(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

Запустите скрипт с помощью accelerate launch и установите --num_processes равным количеству GPU.

accelerate launch --num_processes=2 run_distributed.py

Полный пример приведён в этом минимальном скрипте инференса на нескольких GPU. Подробнее — в руководстве Распределённый инференс с Accelerate.

PyTorch Distributed

Вы также можете настроить параллелизм данных с помощью torch.distributed напрямую, без Accelerate.

Вызовите init_process_group, чтобы создать распределённую среду. torchrun задаёт ранг каждого процесса и общее количество процессов (размер мира) через переменные окружения, поэтому передавать их вручную не нужно.

Каждому процессу требуется собственный GPU и собственный промпт. torchrun устанавливает переменную окружения LOCAL_RANK в значение индекса GPU процесса на его узле, поэтому используйте её для выбора GPU. Используйте get_rank, который считает процессы по всем узлам, для выбора промпта.

import os
import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline

def main():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")

    pipeline = DiffusionPipeline.from_pretrained(
        "Qwen/Qwen-Image", dtype=torch.bfloat16,
    ).to(device)

    prompts = ["a dog", "a cat"]
    prompt = prompts[rank]

    image = pipeline(prompt).images[0]
    image.save(f"{prompt.replace(' ', '_')}.png")

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Сохраните скрипт как run_distributed.py и запустите его с помощью torchrun. Установите значение параметра --nproc-per-node равным количеству GPU. В списке prompts содержится один промпт на процесс, поэтому добавьте промпт для каждого используемого GPU.

torchrun --nproc-per-node 2 run_distributed.py