Функция init_process_group

Функция distributed.init_process_group инициализирует распределенную среду PyTorch для обучения моделей на нескольких устройствах или узлах. Первым параметром передается бэкенд ('gloo', 'nccl' или 'mpi'), вторым параметром можно указать URL мастер-узла через init_method. Обязательными параметрами являются rank (номер текущего процесса) и world_size (общее количество процессов).

Синтаксис

import torch.distributed as dist

dist.init_process_group(
	backend,
	init_method='env://',
	rank=-1,
	world_size=-1,
	timeout=datetime.timedelta(minutes=30),
	group_name=''
)

Пример

Базовая инициализация группы с использованием переменных окружения:

import torch.distributed as dist

# Инициализация с бэкендом 'gloo'
dist.init_process_group(
	backend='gloo',
	init_method='env://',
	rank=0,
	world_size=2
)

print(f"Rank: {dist.get_rank()}")
print(f"World size: {dist.get_world_size()}")
print(f"Backend: {dist.get_backend()}")

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

"Rank: 0"
"World size: 2"
"Backend: gloo"

Пример

Инициализация группы с использованием TCP-адреса мастер-узла:

import torch.distributed as dist

# Инициализация через TCP
dist.init_process_group(
	backend='nccl',
	init_method='tcp://127.0.0.1:23456',
	rank=1,
	world_size=4
)

res = f"Process {dist.get_rank()} of {dist.get_world_size()}"
print(res)

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

"Process 1 of 4"

Пример

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

import torch.distributed as dist

# Проверка перед инициализацией
if not dist.is_initialized():
	dist.init_process_group(
		backend='gloo',
		init_method='env://',
		rank=0,
		world_size=1
	)

print(f"Is initialized: {dist.is_initialized()}")

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

"Is initialized: True"

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

функцию destroy_process_group,
которая завершает работу распределенной группы
функцию get_rank,
которая возвращает номер текущего процесса
функцию get_world_size,
которая возвращает общее количество процессов
функцию is_initialized,
которая проверяет, инициализирована ли группа