在Google Cloud TPU上训练时如何显示PyTorch Lightning epoch进度条？

问题根源 #

• nprocs警告：PyTorch Lightning在GCP TPU虚拟机环境下，默认多进程启动逻辑与PJRT runtime适配冲突，识别8个设备时触发不支持提示。
• 进度条消失：TPU分布式训练时，非主进程的日志输出被默认抑制，导致进度条仅在单进程环境可见，多TPU进程下无法正常显示。

修复方案 #

1. 强制主进程显示进度条 #

自定义进度条回调，仅让全局主进程输出进度条，其余进程禁用：

from lightning.pytorch.callbacks import TQDMProgressBar

class MainProcessProgressBar(TQDMProgressBar):
    def init_validation_tqdm(self):
        bar = super().init_validation_tqdm()
        bar.disable = not self.trainer.is_global_zero
        return bar

    def init_predict_tqdm(self):
        bar = super().init_predict_tqdm()
        bar.disable = not self.trainer.is_global_zero
        return bar

    def init_test_tqdm(self):
        bar = super().init_test_tqdm()
        bar.disable = not self.trainer.is_global_zero
        return bar

在初始化Trainer时替换默认进度条，并显式启用进度条：

trainer = pl.Trainer(
    callbacks=[checkpoint_callback, MainProcessProgressBar()],
    max_epochs=N_EPOCHS,
    accelerator='tpu',
    devices=8,
    enable_progress_bar=True,
    logger=True
)

2. 消除nprocs警告 #

在脚本最开头添加环境变量配置，强制TPU使用PJRT runtime的正确进程初始化逻辑：

import os
os.environ['XLA_USE_PJRT'] = '1'
os.environ['TPU_NUM_DEVICES'] = '8'

if __name__ == '__main__':
    # 原有数据模块、模型初始化代码

3. 规范DataModule的分布式逻辑 #

确保数据准备和拆分符合分布式训练要求，避免多进程重复操作：

class IsaDataModule(pl.LightningDataModule):
    def prepare_data(self):
        # 仅主进程执行数据下载/预处理
        if self.trainer.is_global_zero:
            # 你的数据准备代码
            pass

    def setup(self, stage=None):
        # 所有进程执行数据拆分
        if stage in ('fit', None):
            self.train_dataset = ...
            self.val_dataset = ...
        if stage in ('test', None):
            self.test_dataset = ...

4. 版本兼容性优化 #

尝试将PyTorch Lightning升级到2.0.5及以上版本，该版本修复了部分TPU环境下的日志与进度条适配问题，同时保持torch-xla2.0与torch2.0.0的版本匹配。

验证标准 #

启动训练后确认：

• Unsupported nprocs (8), ignoring...警告不再出现
• 终端正常显示epoch和步骤的进度条
• TPU使用率保持稳定，训练流程无异常

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