PyTorch多GPU推理：固定单输入、动态另一输入的实现方案

针对你提出的固定输入b、动态输入a₁~aₙ，用n块GPU各处理一组(a_i, b)并独立保存输出的需求，以下是基于PyTorch DDP的实现方案及相关注意事项：

核心思路 #

利用DDP启动n个进程，每个进程绑定一块GPU并加载一份模型副本；每个进程获取对应的动态输入a_i，并加载固定输入b；独立执行推理后将结果保存到专属文件，无需合并输出。

具体实现步骤与代码示例 #

1. 编写推理脚本 #

import torch
import torch.distributed as dist

# 替换为你的模型类
class YourModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 示例模型结构，根据实际需求替换
        self.fc = torch.nn.Linear(10, 10)
    
    def forward(self, a, b):
        # 模型前向逻辑，接收a和b两个输入
        return self.fc(a + b)

def load_fixed_input_b():
    # 替换为你的固定输入b加载逻辑
    return torch.randn(1, 10)

def load_all_a_inputs():
    # 替换为你的动态输入列表加载逻辑，长度需等于GPU数量
    return [torch.randn(1, 10) for _ in range(dist.get_world_size())]

def main():
    # 初始化分布式环境（由torchrun自动设置环境变量）
    dist.init_process_group(backend='nccl')
    local_rank = dist.get_local_rank()
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # 绑定当前进程到指定GPU
    torch.cuda.set_device(local_rank)

    # 加载模型并转移到当前GPU
    model = YourModel().to(local_rank)
    model.eval()  # 切换到推理模式

    # 加载固定输入b
    # 方案1：每个进程独立加载（适合b体积较小的情况）
    b = load_fixed_input_b().to(local_rank)
    # 方案2：主进程加载后广播给所有进程（适合b体积较大，节省内存）
    # if rank == 0:
    #     b = load_fixed_input_b().to(local_rank)
    # dist.broadcast(b, src=0)  # 主进程（rank=0）将b发送给所有进程

    # 获取当前进程对应的动态输入a_i
    a_list = load_all_a_inputs()
    a_i = a_list[rank].to(local_rank)

    # 执行推理（关闭梯度计算提升效率）
    with torch.no_grad():
        output = model(a_i, b)

    # 保存结果，按进程rank命名避免文件冲突
    save_path = f"output_rank_{rank}.pt"
    torch.save(output.cpu(), save_path)
    print(f"Rank {rank} 推理完成，结果已保存至 {save_path}")

    # 销毁分布式环境
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

2. 启动脚本 #

使用torchrun启动，指定单节点GPU数量（即n）：

torchrun --nproc_per_node=n your_inference_script.py

例如，4块GPU时：

torchrun --nproc_per_node=4 your_inference_script.py

关键注意事项 #

• 动态输入列表a_list的长度必须等于GPU数量（world_size），确保每个进程能获取到唯一的a_i。
• 保存输出时，必须通过进程rank或唯一标识命名文件，避免不同进程的输出互相覆盖。
• 模型加载时，每个进程独立加载副本，无需共享模型权重（DDP会自动处理进程间的权重同步，但推理模式下同步无影响）。
• 若固定输入b体积较大，建议使用主进程广播的方式加载，减少重复内存占用。

替代方案（无需DDP） #

如果不想使用DDP，也可以用torch.multiprocessing手动管理多进程，但DDP在多GPU环境下的稳定性和兼容性更优，推荐优先使用。

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