Распределённый вывод

Распределённый вывод можно разделить на две категории:

1. Загрузка всей модели на каждый GPU и последовательная передача частей батча через копию модели на каждом GPU
2. Загрузка частей модели на каждый GPU и обработка одного входа за раз

Это руководство охватывает первую скобку. Во вторую см. Big Model Inference.

Автоматическая отправка фрагментов батча каждой загруженной модели

Это наиболее ресурсоемкое по памяти решение, поскольку в каждый момент времени каждому GPU требуется хранить в памяти полную копию модели.

Обычно при выполнении этой операции пользователи отправляют модель на конкретное устройство, чтобы загрузить её с CPU, а затем перемещают каждый промпт на другое устройство.

Базовый конвейер с использованием библиотеки diffusers может выглядеть примерно так:

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

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)

Затем выполняется инференс на основе конкретного промпта:

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    pipe.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    result = pipe(prompt).images[0]
    result.save(f"result_{rank}.png")

Можно заметить, что нам приходится проверять ранг, чтобы понять, какой промпт отправить, что может быть немного утомительно.

Пользователь может также подумать, что с помощью Accelerate использование Accelerator для подготовки загрузчика данных для такой задачи может быть простым способом управления этим. (Чтобы узнать больше, ознакомьтесь с соответствующим разделом в Quick Tour)

Сможет ли он справиться? Да. Однако добавляет ли он ненужный лишний код? Тоже да.

С помощью Accelerate мы можем упростить этот процесс, используя контекстный менеджер Accelerator.split_between_processes() (он также доступен в PartialState и AcceleratorState). Эта функция автоматически распределяет переданные ей данные (будь то промпт, набор тензоров, словарь с предыдущими данными и т. д.) по всем процессам (с возможной дополнительной отступной частью), чтобы вы могли использовать их сразу.

Перепишем приведённый выше пример, используя этот контекстный менеджер:

import torch
from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

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

А затем, чтобы запустить код, мы можем использовать Accelerate:

Если вы сгенерировали файл конфигурации для использования с помощью accelerate config:

accelerate launch distributed_inference.py

Если у вас есть конкретный файл конфигурации, который вы хотите использовать:

accelerate launch --config_file my_config.json distributed_inference.py

Или, если вы не хотите создавать файлы конфигурации и хотите запустить на двух GPU:

Примечание: вы получите предупреждения о том, что значения были угаданы на основе вашей системы. Чтобы их убрать, выполните accelerate config default или пройдите accelerate config, чтобы создать файл конфигурации.

accelerate launch --num_processes 2 distributed_inference.py

Теперь шаблонный код, необходимый для разделения этих данных, легко сводится к нескольким строкам.

Но что, если распределение промптов по GPU будет неравномерным? Например, если у нас 3 промпта, но только 2 GPU?

В рамках контекстного менеджера первая GPU получает первые два запроса, а вторая GPU — третий, что обеспечивает разбиение всех запросов без каких-либо накладных расходов.

Однако, что делать, если нам нужно выполнить какие-то действия с результатами со всех GPU? (Например, собрать их все и выполнить какую-либо постобработку) Вы можете передать apply_padding=True, чтобы убедиться, что списки промптов дополнены до одинаковой длины, при этом дополнительные данные берутся из последнего примера. Таким образом, на всех GPU будет одинаковое количество промптов, и вы сможете собрать результаты.

Это требуется только при выполнении действий, например, при сборе результатов, когда данные на каждом устройстве должны иметь одинаковую длину. Для базового инференса это не нужно.

Например:

import torch
from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat", "a chicken"], apply_padding=True) as prompt:
    result = pipe(prompt).images

На первом GPU промпты будут ["a dog", "a cat"], а на втором GPU — ["a chicken", "a chicken"]. Убедитесь, что вы отбрасываете последний образец, так как он будет дублировать предыдущий.

Более сложные примеры, например, по использованию с LLM, можно найти здесь.