PyTorch Lightning多GPU（3卡DDP）训练性能不及单GPU的问题排查

我最近搭了一个PyTorch Lightning的测试模型，专门用来对比3卡DDP多GPU训练（简称3G）和单GPU训练（简称1G）的耗时情况，但结果一直是单GPU训练比多GPU更快。

我尝试过的调整方案 #

• 从较小的批次大小（16）调整到较大的批次大小（256）
• 尝试添加最多8个额外的隐藏层
• 在Trainer参数中调整accumulate_grad_batches的值

可惜这些调整都没带来任何改善。我怀疑可能是我的模型太简单，导致多GPU训练的优势发挥不出来，但我刚接触多GPU训练，不敢轻易下定论，而且我的实验也没得到正向结果。我在Stack Overflow上看到过类似问题，但大多是解决报错的，而我想搞清楚为什么用PyTorch Lightning时性能会这么拉胯。

我的问题 #

这个测试案例里有没有什么关键问题，导致单GPU训练反而比多GPU训练表现更好？

以下是我的测试代码：

import pytorch_lightning as pl
from pytorch_lightning import LightningModule, LightningDataModule
from torch import nn
from torch.utils.data import Dataset, DataLoader
import torch
import time
import wandb
from pytorch_lightning.loggers import WandbLogger
import pynvml

class SimpleDataset(Dataset):
    def __init__(self, size):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        # Generate some random data
        x = torch.randn(1, 10)
        y = torch.randn(1)
        return x, y

class SimpleDataModule(LightningDataModule):
    def __init__(self, batch_size, size):
        super().__init__()
        self.batch_size = batch_size
        self.size = size

    def setup(self, stage=None):
        self.dataset = SimpleDataset(self.size)

    def train_dataloader(self):
        return DataLoader(self.dataset, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.dataset, batch_size=self.batch_size)

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 256)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 256)
        self.fc4 = nn.Linear(256, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = self.fc4(x)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_pred = self(x)
        loss = nn.MSELoss()(y_pred, y)
        self.log('train_loss', loss)

        # Get individual GPU utilization
        pynvml.nvmlInit()
        utilization = []
        for i in range(pynvml.nvmlDeviceGetCount()):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            utilization.append(pynvml.nvmlDeviceGetUtilizationRates(handle).gpu)
        pynvml.nvmlShutdown()

        # Log individual GPU utilization
        for i, util in enumerate(utilization):
            self.log(f'gpu_{i}_utilization', util)

        return loss


    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_pred = self(x)
        loss = nn.MSELoss()(y_pred, y)
        self.log('val_loss', loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.01)

data_module = SimpleDataModule(batch_size=256, size=5000)
model = SimpleModel()

start_time = time.time()

wandb_logger = WandbLogger(project='simple_model')

trainer = pl.Trainer(strategy='ddp', devices=3, max_epochs=200, accumulate_grad_batches=3, logger=wandb_logger)

trainer.fit(model, data_module)
if trainer.is_global_zero:
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Training completed in {elapsed_time:.2f} seconds")

相关指标 #

我用Weights and Biases跟踪了GPU使用率，想从中找到问题所在。结果显示3G训练时GPU使用率波动很大，而1G训练时使用率相对稳定。我觉得这可能是问题的核心，但不知道该怎么解决。下图是相关的监控结果：

有没有什么明显的代码问题，修复后能让3G训练的性能超过1G？

备注：内容来源于stack exchange，提问作者CCranney