PyTorch Lightning多GPU训练时sync_dist与rank_zero_only设置方法

一、两个参数的核心含义 #

1. rank_zero_only=True #

在4块GPU的DDP训练场景中，每块GPU对应一个独立进程，每个进程有唯一的rank编号，其中rank=0的进程是主进程。设置这个参数后，只有主进程会执行日志记录操作，其他3个进程直接跳过日志步骤。核心目的就是避免多个进程重复输出相同日志，导致日志文件堆满重复内容。

2. sync_dist=True #

DDP训练时，每个GPU进程只处理自己分到的那部分批次数据，计算出的loss、准确率等指标都是当前进程局部批次的结果。设置这个参数后，PyTorch Lightning会自动在4个进程间同步这些指标数据，计算出全局的平均值（可通过sync_dist_fn指定其他聚合方式），再把这个全局结果写入日志。如果不设置，你看到的日志指标只是单GPU的局部值，完全反映不了整个训练/验证集的真实情况。

二、不同场景下的正确设置方式 #

场景1：记录全局聚合后的训练/验证指标 #

比如要记录训练的全局平均loss、验证集的全局准确率，必须让所有进程同步数据得到真实全局指标，同时只让主进程记录一次避免重复：

def training_step(self, batch, batch_idx):
    x, y = batch
    logits = self(x)
    loss = self.loss_fn(logits, y)
    # 同步所有进程的loss，计算全局平均值，仅主进程记录
    self.log("train_loss", loss, sync_dist=True, rank_zero_only=True)
    return loss

def validation_step(self, batch, batch_idx):
    x, y = batch
    logits = self(x)
    acc = self.acc_fn(logits, y)
    # 验证集指标必须同步，否则是单GPU局部准确率
    self.log("val_acc", acc, sync_dist=True, rank_zero_only=True)

场景2：记录进程专属/无需聚合的信息 #

比如只想在主进程打印epoch结束的提示，或者记录单进程的批次处理耗时（不需要全局聚合），只需要开启rank_zero_only=True即可：

def on_train_epoch_end(self):
    # 仅主进程输出epoch完成信息
    self.log("epoch_finish", f"Epoch {self.current_epoch} done!", rank_zero_only=True)

要避开的错误设置 #

• 只开sync_dist=True不开rank_zero_only=True：会导致4个进程都记录同步后的相同指标，日志里同一指标重复出现4次，完全没必要。
• 需要全局指标时不开sync_dist=True：这时候日志里的loss、准确率都是单GPU的局部值，和真实全局指标偏差很大，会误导你判断模型训练情况。

额外注意点 #

• sync_dist会带来少量进程间通信的开销，如果是对全局指标无要求的日志（比如进程专属的调试信息），没必要开启。
• PyTorch Lightning的log方法默认是rank_zero_only=False、sync_dist=False，所以需要手动按需设置。

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