Функция all_reduce

Функция distributed.all_reduce выполняет коллективную операцию редукции данных между всеми процессами в группе. Эта операция суммирует, перемножает или выполняет другую агрегацию значений тензоров от всех процессов и записывает результат обратно в тот же тензор каждого процесса.

Первый параметр функции - это тензор tensor, который будет изменён. Второй параметр - операция редукции op (по умолчанию ReduceOp.SUM). Третий параметр - группа процессов group, в которой выполняется операция. Функция работает синхронно и блокирует выполнение до завершения операции.

Синтаксис

import torch.distributed as dist

dist.all_reduce(tensor, op=ReduceOp.SUM, group=None)

Пример

Базовый пример использования all_reduce для суммирования тензоров между двумя процессами:

<~show~ file="example_1.py"> <+python+> import torch import torch.distributed as dist dist.init_process_group(backend='gloo', init_method='env://') rank = dist.get_rank() t = torch.tensor([rank + 1, rank + 2], dtype=torch.float32) print(f"Process {rank} before all_reduce: {t}") dist.all_reduce(t, op=dist.ReduceOp.SUM) print(f"Process {rank} after all_reduce: {t}") dist.destroy_process_group() <-python-> <-show->

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

Process 0 before all_reduce: tensor([1., 2.])
Process 1 before all_reduce: tensor([2., 3.])
Process 0 after all_reduce: tensor([3., 5.])
Process 1 after all_reduce: tensor([3., 5.])

Пример

Операция all_reduce с умножением:

<~show~ file="example_2.py"> <+python+> import torch import torch.distributed as dist dist.init_process_group(backend='gloo', init_method='env://') rank = dist.get_rank() t = torch.tensor([rank + 1, rank + 2], dtype=torch.float32) print(f"Process {rank} before all_reduce: {t}") dist.all_reduce(t, op=dist.ReduceOp.PRODUCT) print(f"Process {rank} after all_reduce: {t}") dist.destroy_process_group() <-python-> <-show->

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

Process 0 before all_reduce: tensor([1., 2.])
Process 1 before all_reduce: tensor([2., 3.])
Process 0 after all_reduce: tensor([2., 6.])
Process 1 after all_reduce: tensor([2., 6.])

Пример

Использование all_reduce с пользовательской группой процессов и операцией нахождения максимума:

<~show~ file="example_3.py"> <+python+> import torch import torch.distributed as dist dist.init_process_group(backend='gloo', init_method='env://') rank = dist.get_rank() group = dist.new_group([0, 2]) if rank in [0, 2]: t = torch.tensor([rank + 1, rank + 2], dtype=torch.float32) print(f"Process {rank} before all_reduce: {t}") dist.all_reduce(t, op=dist.ReduceOp.MAX, group=group) print(f"Process {rank} after all_reduce: {t}") dist.destroy_process_group() <-python-> <-show->

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

Process 0 before all_reduce: tensor([1., 2.])
Process 2 before all_reduce: tensor([3., 4.])
Process 0 after all_reduce: tensor([3., 4.])
Process 2 after all_reduce: tensor([3., 4.])

Пример

Использование all_reduce для вычисления средней точности модели на всех устройствах при распределённом обучении:

<~show~ file="example_4.py"> <+python+> import torch import torch.distributed as dist dist.init_process_group(backend='gloo', init_method='env://') rank = dist.get_rank() world_size = dist.get_world_size() torch.manual_seed(rank) accuracy = torch.tensor([float(torch.randint(70, 100, (1,)).item()) / 100.0]) print(f"Process {rank} local accuracy: {accuracy.item():.2f}") dist.all_reduce(accuracy, op=dist.ReduceOp.SUM) avg_accuracy = accuracy / world_size print(f"Process {rank} global average accuracy: {avg_accuracy.item():.2f}") dist.destroy_process_group() <-python-> <-show->

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

Process 0 local accuracy: 0.83
Process 1 local accuracy: 0.95
Process 2 local accuracy: 0.71
Process 3 local accuracy: 0.88
Process 0 global average accuracy: 0.84
Process 1 global average accuracy: 0.84
Process 2 global average accuracy: 0.84
Process 3 global average accuracy: 0.84

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

функцию init_process_group,
которая инициализирует распределённую среду
функцию all_gather,
которая собирает тензоры со всех процессов
функцию broadcast,
которая отправляет тензор от одного процесса всем остальным
функцию reduce,
которая выполняет редукцию данных на одном процессе