Обёртка нескольких процессов в PyTorch

Обёртка DistributedDataParallel надевается на модуль в каждом процессе группы. После прямого прохода она собирает градиенты с других процессов и усредняет их, чтобы веса шли в одном направлении.

Создать обёртку без живой группы нельзя: конструктор ждёт инициализированную связь. На одной машине без нескольких процессов достаточно обычного модуля. Обёртку подключают, когда обучение идёт параллельно на нескольких картах. Ниже фрагмент, который запускают уже внутри процесса после создания группы:

import torch
import torch.nn as nn

layer = nn.Linear(2, 1)
# после init_process_group в этом процессе:
# layer = nn.parallel.DistributedDataParallel(layer)
print(layer.in_features) # выведет 2

Создайте линейный слой на 3 входа и один выход. В комментарии покажите, как надеть на него обёртку нескольких процессов после создания группы. Выведите число входов слоя.

Соберите слой на 2 входа и 4 выхода. Рядом в комментарии запишите строку обёртки и выведите форму матрицы весов.

Для слоя на 4 входа и один выход выведите число параметров. В комментарии укажите, что обёртку создают только внутри процесса группы.