Функция reduce_scatter

Функция distributed.reduce_scatter выполняет две операции одновременно: редуцирование (суммирование, умножение и т.д.) тензоров со всех процессов и разброс (scatter) результата по процессам. Первым параметром функция принимает выходной тензор output, вторым - список входных тензоров input_list. Третьим параметром можно указать операцию редуцирования (по умолчанию ReduceOp.SUM). Функция работает асинхронно с текущим потоком, но блокирует выполнение до завершения операции.

Синтаксис

torch.distributed.reduce_scatter(output, input_list, op=ReduceOp.SUM, group=None, async_op=False)

Параметры

output - тензор, в который будет записан результат для текущего процесса.

input_list - список тензоров, которые будут редуцированы и разбросаны. Каждый процесс должен предоставить список размером равным количеству процессов в группе.

op - операция редуцирования (по умолчанию ReduceOp.SUM). Может принимать значения ReduceOp.SUM, ReduceOp.PRODUCT, ReduceOp.MAX, ReduceOp.MIN, ReduceOp.BAND, ReduceOp.BOR, ReduceOp.BXOR.

group - группа процессов (по умолчанию используется группа по умолчанию).

async_op - если установлено в True, функция возвращает объект Work для асинхронного выполнения.

Пример

Давайте выполним редуцирование и разброс с суммированием для двух процессов:

import torch
import torch.distributed as dist

dist.init_process_group(backend='gloo')
rank = dist.get_rank()
world_size = dist.get_world_size()

t = torch.tensor([rank + 1, rank + 2])
input_list = [t + i for i in range(world_size)]
output = torch.zeros(2)

dist.reduce_scatter(output, input_list)

print(f"Rank {rank}: output {output}")

Результат выполнения кода для первого процесса:

"Rank 0: output tensor([3., 3.])"

Результат выполнения кода для второго процесса:

"Rank 1: output tensor([3., 3.])"

Пример

Давайте используем операцию MAX для редуцирования:

import torch
import torch.distributed as dist

dist.init_process_group(backend='gloo')
rank = dist.get_rank()
world_size = dist.get_world_size()

t = torch.tensor([rank * 2, rank * 2 + 1])
input_list = [t + i for i in range(world_size)]
output = torch.zeros(2)

dist.reduce_scatter(output, input_list, op=dist.ReduceOp.MAX)

print(f"Rank {rank}: output {output}")

Результат выполнения кода для двух процессов:

"Rank 0: output tensor([2., 3.])"
"Rank 1: output tensor([4., 5.])"

Пример

Давайте выполним асинхронную операцию с помощью параметра async_op:

import torch
import torch.distributed as dist

dist.init_process_group(backend='gloo')
rank = dist.get_rank()
world_size = dist.get_world_size()

t = torch.tensor([rank + 1, rank + 2])
input_list = [t + i for i in range(world_size)]
output = torch.zeros(2)

work = dist.reduce_scatter(output, input_list, async_op=True)
work.wait()

print(f"Rank {rank}: output {output}")

Результат выполнения кода:

"Rank 0: output tensor([3., 3.])"
"Rank 1: output tensor([3., 3.])"

Смотрите также

функцию all_reduce,
которая выполняет редуцирование данных со всех процессов
функцию all_gather,
которая собирает данные со всех процессов
класс ReduceOp,
который определяет операции редуцирования
функцию init_process_group,
которая инициализирует распределенную среду