torch多机器多卡推理大模型

在PyTorch中，多机推理通常涉及使用DistributedDataParallel模块。以下是一个简化的例子，展示如何在多台机器上进行PyTorch模型的推理。

假设你有两台机器，IP分别为192.168.1.1和192.168.1.2，你想在第一台机器上进行模型的推理。

1. 在每台机器上设置环境变量：

export MASTER_ADDR=192.168.1.1
export MASTER_PORT=12345

2、在第一台机器（master节点）上初始化分布式进程组，并启动PyTorch进程组： 

import torch.distributed as dist
 
dist.init_process_group(backend='tcp', init_method='env://', world_size=2, rank=0)
 
# 在这里加载你的模型和推理代码

3、在其他机器上（slave节点），设置相应的环境变量并初始化进程组，rank为1：

import torch.distributed as dist
 
dist.init_process_group(backend='tcp', init_method='env://', world_size=2, rank=1)
 
# 在这里加载