Функция all_gather_into_tensor

Функция distributed.all_gather_into_tensor собирает тензоры от всех процессов распределенной группы в один выходной тензор. Первым параметром функция принимает выходной тензор, в который будут записаны собранные данные. Вторым параметром передается входной тензор текущего процесса. Третьим параметром указывается группа процессов. В отличие от функции all_gather, которая собирает данные в список тензоров, данная функция записывает результаты в один предварительно выделенный тензор.

Синтаксис

import torch.distributed as dist

dist.all_gather_into_tensor(
	tensor_out,
	tensor_in,
	group=None
)

Параметры

Функция принимает следующие параметры:

tensor_out - выходной тензор, в который будут записаны собранные данные. Его размер должен быть достаточным для размещения тензоров от всех процессов.
tensor_in - входной тензор текущего процесса, который будет отправлен другим процессам.
group - группа процессов, участвующих в операции. Если параметр не указан, используется группа по умолчанию.

Пример

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

import torch
import torch.distributed as dist
import os

# Инициализация процесса
dist.init_process_group(backend='gloo', init_method='env://')

# Получение ранга текущего процесса
rank = dist.get_rank()
world_size = dist.get_world_size()

# Создание входного тензора
t_in = torch.tensor([rank + 1, rank + 2])

# Создание выходного тензора
t_out = torch.zeros(world_size * 2)

# Сбор данных
dist.all_gather_into_tensor(t_out, t_in)

if rank == 0:
	print(t_out)

dist.destroy_process_group()

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

tensor([1., 2., 2., 3.])

Пример

Давайте соберем данные в тензор с заданной размерностью:

import torch
import torch.distributed as dist
import os

# Инициализация процесса
dist.init_process_group(backend='gloo', init_method='env://')

# Получение ранга текущего процесса
rank = dist.get_rank()
world_size = dist.get_world_size()

# Создание входного тензора (вектор)
t_in = torch.tensor([rank * 10 + 1, rank * 10 + 2, rank * 10 + 3])

# Создание выходного тензора (матрица)
t_out = torch.zeros(world_size, 3)

# Сбор данных
dist.all_gather_into_tensor(t_out, t_in)

if rank == 0:
	print(t_out)

dist.destroy_process_group()

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

tensor([
	[1., 2., 3.],
	[11., 12., 13.],
])

Пример

Давайте используем функцию с указанием конкретной группы процессов:

import torch
import torch.distributed as dist
import os

# Инициализация процесса
dist.init_process_group(backend='gloo', init_method='env://')

# Получение ранга текущего процесса
rank = dist.get_rank()
world_size = dist.get_world_size()

# Создание подгруппы из первых двух рангов
if rank < 2:
	group = dist.new_group(ranks=[0, 1])
else:
	group = dist.new_group(ranks=[0, 1])

# Проверка принадлежности к группе
if rank in [0, 1]:
	# Создание входного тензора
	t_in = torch.tensor([rank * 100, rank * 100 + 1])
	
	# Создание выходного тензора
	t_out = torch.zeros(4)
	
	# Сбор данных в подгруппе
	dist.all_gather_into_tensor(t_out, t_in, group=group)
	
	if rank == 0:
		print(t_out)

dist.destroy_process_group()

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

tensor([0., 1., 100., 101.])

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

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