单卡显存不足时如何进行多卡推理？

Question

在运行VideoChat-7B等预训练大模型时，所需要的显存超出了我单卡的显存，我在服务器上有4张3080-10G显卡，怎样才能同时使用4张卡来运行这个模型？

Answer 1

引自免费微信小程序：皆我百晓生

要解决这个问题，你需要利用NVIDIA的并行计算平台CUDA以及cuDNN库来实现模型在多GPU上的分布式训练和推理。这里是一个基于PyTorch的示例步骤，展示如何将 VideoChat-7B 等大型模型扩展到多GPU上。

首先，确保你的系统已经安装了最新版本的PyTorch、CUDA和cuDNN，并且所有显卡都支持NCCL（NVIDIA Collective Communication Library）用于GPU间的通信。你可以通过以下命令检查NCCL版本：

nvidia-smi --query-gpu=nccl_version,memory.total --format=csv

接下来，我们将分步设置多GPU环境并运行模型：

1. 安装必要的库: 如果尚未安装，可以使用pip安装torch, torchvision, 和 torch_cluster 库：

   pip install torch torchvision torch_cluster
   

2. 编写多GPU运行脚本:

   在Python脚本中，你需要设置设备（这里指GPU），然后创建一个数据加载器，使其能够并行地处理每个GPU的数据。最后，在模型上分布数据，并执行推理。

   import torch
   from torch.nn.parallel import DistributedDataParallel as DDP
   from torch.utils.data.distributed import DistributedSampler
   
   # 设置设备
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   num_gpus = torch.cuda.device_count()
   
   # 假设你有一个名为data_loader的函数，它返回一个数据集迭代器
   data_loader = data_loader(batch_size=your_batch_size)
   
   # 创建分布式数据加载器
   sampler = DistributedSampler(data_loader.dataset)
   data_loader = torch.utils.data.DataLoader(
       data_loader.dataset,
       batch_size=data_loader.batch_size,
       shuffle=False,
       num_workers=data_loader.num_workers,
       pin_memory=True,
       sampler=sampler
   )
   
   # 将模型复制到所有GPU上，并使用DDP封装
   model = YourModelClass(*args, **kwargs).to(device)
   model = DDP(model, device_ids=[i for i in range(num_gpus)])
   
   # 开始推理
   with torch.no_grad():
       for inputs, targets in data_loader:
           inputs, targets = inputs.to(device), targets.to(device)
           outputs = model(inputs)
           # 处理输出...
   

3. 启动多进程: 由于PyTorch不支持在单个进程中跨多个GPU运行，你需要为每个GPU启动一个新的进程。可以使用torch.multiprocessing.spawn来实现这一点。这是一个简单的例子：

   import torch.multiprocessing as mp
   
   def run(rank, world_size):
       # 上述脚本中的代码放在这里，但不需要device设置部分
       ...
   
   if __name__ == "__main__":
       mp.spawn(run, args=(num_gpus,), nprocs=num_gpus, join=True)
   

4. 配置环境变量: 在运行脚本之前，确保设置了以下环境变量，以便正确地分配工作负载：

   export MASTER_ADDR=localhost
   export MASTER_PORT=12345
   CUDA_VISIBLE_DEVICES=0,1,2,3 python your_script.py
   

   这里，MASTER_ADDR和MASTER_PORT用于定义主节点的地址和端口，CUDA_VISIBLE_DEVICES指定参与工作的GPU列表。

注意：上述代码仅作为示例，实际应用中你需要根据具体的模型和数据集调整代码。例如，可能需要修改模型类YourModelClass、数据加载器data_loader和模型的输入输出处理逻辑。如果你使用的是预训练模型，确保该模型支持并行化。对于VideoChat-7B这样的模型，你可能需要找到相应的开源实现，因为它们通常不包含在标准的PyTorch库中。