能否并行化原生单批次PyTorch模型？附多维度线性模型示例

问题背景 #

原生仅支持单批次输入的PyTorch模型（比如依赖torch.tensordot这类无法直接适配批次维度的算子），能否实现并行化？常规并行化思路是用torch.bmm替代torch.matmul并指定批次维度，但tensordot不支持这种方式。我们需要一种同时支持训练和推理的方案，能并行计算批次中每个样本的梯度。

原模型代码无法直接传入批次维度输入（如torch.randn(batch_size, input_dim, hidden_dim)），示例如下：

import torch
import torch.nn as nn

class LinearMultidimModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(LinearMultidimModel, self).__init__()
        self.weight = nn.Parameter(torch.randn(input_dim, hidden_dim, output_dim))
        self.bias = nn.Parameter(torch.randn(output_dim))

    def forward(self, x):
        # 仅支持(input_dim, hidden_dim)的单样本输入
        out = torch.tensordot(x, self.weight, dims=[[0,1],[0,1]]) + self.bias
        return out

# 单样本调用示例
input_dim = 3
hidden_dim=2
output_dim = 1
model = LinearMultidimModel(input_dim, hidden_dim, output_dim)
x = torch.randn(input_dim, hidden_dim)
output = model(x)
print(output)

解决方案：适配批次维度的并行化实现 #

方法1：调整torch.tensordot的维度匹配参数

修改forward方法中tensordot的维度对应关系，保留批次维度，将输入的后两个维度与权重的前两个维度做张量点积：

class LinearMultidimModelBatch(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(LinearMultidimModelBatch, self).__init__()
        self.weight = nn.Parameter(torch.randn(input_dim, hidden_dim, output_dim))
        self.bias = nn.Parameter(torch.randn(output_dim))

    def forward(self, x):
        # x的形状：(batch_size, input_dim, hidden_dim)
        # 匹配x的第1、2维度与weight的第0、1维度，保留批次维度
        out = torch.tensordot(x, self.weight, dims=[[1,2],[0,1]]) + self.bias
        # 输出形状：(batch_size, output_dim)
        return out

# 批次调用示例
batch_size = 5
model_batch = LinearMultidimModelBatch(input_dim, hidden_dim, output_dim)
x_batch = torch.randn(batch_size, input_dim, hidden_dim)
output_batch = model_batch(x_batch)
print(output_batch.shape)  # 输出: torch.Size([5, 1])

方法2：使用torch.einsum实现灵活维度映射

torch.einsum通过维度符号定义张量运算，完全支持批次维度，且无需固定隐藏维度数量：

class LinearMultidimModelEinsum(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(LinearMultidimModelEinsum, self).__init__()
        self.weight = nn.Parameter(torch.randn(input_dim, hidden_dim, output_dim))
        self.bias = nn.Parameter(torch.randn(output_dim))

    def forward(self, x):
        # 维度符号说明：b=批次, i=输入维度, h=隐藏维度, o=输出维度
        out = torch.einsum('bih, iho -> bo', x, self.weight) + self.bias
        return out

# 批次调用示例
model_einsum = LinearMultidimModelEinsum(input_dim, hidden_dim, output_dim)
output_einsum = model_einsum(x_batch)
print(output_einsum.shape)  # 输出: torch.Size([5, 1])

批次梯度并行计算说明 #

上述两种方案完全支持训练与推理：

• 推理阶段：批次输入会被并行处理，大幅提升计算效率
• 训练阶段：PyTorch会自动对批次中每个样本的梯度进行计算与聚合（根据损失函数的reduction参数，默认求平均），无需额外编写梯度并行逻辑，原生支持反向传播过程的梯度并行计算。

内容的提问来源于stack exchange，提问作者Marko Karbevski