Класс ReduceOp

Класс ReduceOp из модуля torch.distributed определяет константы, представляющие различные операции редукции, которые используются в коллективных коммуникациях, таких как all_reduce и reduce. Основные операции включают суммирование (SUM), произведение (PRODUCT), нахождение максимума (MAX), минимума (MIN) и другие.

Класс также предоставляет операцию BAND (побитовое И), BOR (побитовое ИЛИ) и BXOR (побитовое исключающее ИЛИ). Эти операции используются для агрегации данных с нескольких процессов в распределённой среде.

Синтаксис

import torch.distributed as dist

# Использование в коллективных операциях
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

Пример

Давайте выполним редукцию суммирования на двух процессах с использованием all_reduce и ReduceOp.SUM:

import torch
import torch.distributed as dist

# Инициализация процессной группы
dist.init_process_group(backend='gloo')

rank = dist.get_rank()
t = torch.tensor([rank + 1, rank + 2, rank + 3])

# Редукция суммирования
dist.all_reduce(t, op=dist.ReduceOp.SUM)

if rank == 0:
	print(t)

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

tensor([4, 6, 8])

Пример

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

import torch
import torch.distributed as dist

dist.init_process_group(backend='gloo')

rank = dist.get_rank()
t = torch.tensor([rank * 10 + 1, rank * 10 + 5, rank * 10 + 9])

# Редукция максимума
dist.all_reduce(t, op=dist.ReduceOp.MAX)

if rank == 0:
	print(t)

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

tensor([11, 15, 19])

Пример

Давайте используем операцию побитового И с ReduceOp.BAND:

import torch
import torch.distributed as dist

dist.init_process_group(backend='gloo')

rank = dist.get_rank()
t = torch.tensor([rank + 1, rank + 2, rank + 3], dtype=torch.int32)

# Побитовая редукция И
dist.all_reduce(t, op=dist.ReduceOp.BAND)

if rank == 0:
	print(t)

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

tensor([1, 0, 0], dtype=torch.int32)

Пример

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

import torch
import torch.distributed as dist

dist.init_process_group(backend='gloo')

rank = dist.get_rank()
t = torch.tensor([rank + 2, rank + 3, rank + 4])

# Редукция произведения
dist.all_reduce(t, op=dist.ReduceOp.PRODUCT)

if rank == 0:
	print(t)

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

tensor([6, 12, 20])

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

функцию all_reduce,
которая выполняет редукцию данных со всех процессов
функцию reduce,
которая выполняет редукцию данных на указанный процесс
функцию init_process_group,
которая инициализирует распределённую процессную группу
функцию get_world_size,
которая возвращает общее количество процессов в группе