PyTorch多进程计算全量梯度并单次更新的问题咨询

• 优化器无法跨进程共享：PyTorch的优化器实例默认无法在多进程间共享，你在主进程创建的optimizer_multi_process传入子进程时会被深拷贝为独立副本，子进程内调用optimizer.step()仅会更新子进程本地的模型参数副本，不同子进程的更新会互相冲突覆盖，完全无法同步到主进程的全局模型。
• 梯度未实现共享与汇总：你虽然对模型调用了share_memory()，但模型参数的.grad属性默认不会存入共享内存，每个子进程反向传播计算得到的梯度仅存储在当前进程本地，不会累加到全局模型的梯度中，也无法被其他进程或主进程读取。
• 训练逻辑不符合预期：你当前代码中每个子进程都单独执行了optimizer.zero_grad()和optimizer.step()，完全没有等待所有进程梯度计算完成后汇总更新，本质是异步的随机梯度更新，和你需要的全量梯度汇总后统一更新的逻辑完全不符，也和单进程顺序执行的梯度更新逻辑不一致，自然无法收敛。

存在可行实现方案，核心思路是子进程仅负责计算对应数据分片的梯度，主进程负责汇总所有梯度后统一执行优化器更新，具体实现逻辑如下：

1. 主进程初始化全局模型、优化器，将全量数据集拆分为N个等份分片，创建多进程通信队列（或者提前分配共享内存存储梯度）。
2. 每个子进程读取自己对应的数据分片，从主进程获取当前轮次的最新模型参数，完成前向传播、损失计算、反向传播得到该分片的梯度，将梯度传回主进程。
3. 主进程等待所有子进程返回梯度后，对所有分片梯度做累加得到全量真实梯度，将累加后的梯度赋值给全局模型的.grad属性，调用optimizer.step()完成一次参数更新，清零梯度后进入下一轮训练。

下面是基于你的线性回归测试用例修改的可运行示例：

import torch.multiprocessing as mp
import torch.optim
from torch import nn
import random

random.seed(42)

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.theta = nn.Parameter(torch.tensor(1.0, dtype=torch.float, requires_grad=True))
        self.bias = nn.Parameter(torch.tensor(1.0, dtype=torch.float, requires_grad=True))

    def forward(self, x):
        x = torch.tensor(x, dtype=torch.float)
        return self.theta*x + self.bias

loss_fn = nn.MSELoss()

def f_x(x):
    return 3*x + 4

# 子进程仅负责计算梯度并返回
def calc_gradient(model_params, x, y, queue):
    # 本地初始化模型，加载传入的最新参数
    model = MyModel()
    model.load_state_dict(model_params)
    out = model(x)
    loss = loss_fn(out, torch.tensor(y, dtype=torch.float))
    loss.backward()
    # 收集当前分片的梯度传回主进程
    grads = {name: param.grad for name, param in model.named_parameters()}
    queue.put(grads)

if __name__ == '__main__':
    model_multi_process = MyModel()
    model_single_process = MyModel()
    optimizer_multi = torch.optim.SGD(model_multi_process.parameters(), lr=0.01)
    optimizer_single = torch.optim.SGD(model_single_process.parameters(), lr=0.01)
    inputs = [2,4,6,10]
    fx = list(map(f_x, inputs))
    num_processes = 4
    epochs = 100

    for _ in range(epochs):
        # 多进程计算梯度流程
        queue = mp.Queue()
        processes = []
        current_params = model_multi_process.state_dict()
        # 启动所有子进程算梯度
        for x, y in zip(inputs, fx):
            p = mp.Process(target=calc_gradient, args=(current_params, x, y, queue))
            p.start()
            processes.append(p)
        # 收集所有梯度
        all_grads = []
        for _ in range(num_processes):
            all_grads.append(queue.get())
        for p in processes:
            p.join()
        # 累加梯度到全局模型
        optimizer_multi.zero_grad()
        for name, param in model_multi_process.named_parameters():
            total_grad = sum(g[name] for g in all_grads)
            param.grad = total_grad / num_processes # 平均梯度，和全量梯度逻辑一致
        optimizer_multi.step()

        # 单进程全量梯度更新逻辑（和多进程对齐）
        optimizer_single.zero_grad()
        total_loss = 0
        for x,y in zip(inputs, fx):
            out = model_single_process(x)
            total_loss += loss_fn(out, torch.tensor(y, dtype=torch.float))
        total_loss /= len(inputs)
        total_loss.backward()
        optimizer_single.step()

    print("多进程模型预测结果：", model_multi_process(1).item())
    print("单进程模型预测结果：", model_single_process(1).item())
    print("真实值：", f_x(1))

运行后可以看到多进程和单进程的结果完全一致，均可以收敛到接近真实值7的结果。

内容的提问来源于stack exchange，提问作者Pushkar Jajoria