本地环境Pytorch Lightning无法启动fit方法问题排查求助

问题背景 #

使用PyTorch学习NLP时，相同代码在Google Colab可正常运行，但本地调用trainer.fit()后训练无法启动，DRAM已耗尽却未触发OOM错误。

相关代码 #

trainer.fit(
    task,
    train_dataloaders= train_dataloader,
    val_dataloaders= val_dataloader
)

警告信息 #

Missing logger folder: d:\lightning_logs
C:\Users\Leecarry\AppData\Roaming\Python\Python310\site-packages\pytorch_lightning\callbacks\model_checkpoint.py:608: UserWarning: Checkpoint directory D:\ exists and is not empty.
  rank_zero_warn(f"Checkpoint directory {dirpath} exists and is not empty.")
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
C:\Users\Leecarry\AppData\Roaming\Python\Python310\site-packages\pytorch_lightning\core\optimizer.py:380: RuntimeWarning: Found unsupported keys in the optimizer configuration: {'scheduler'}
  rank_zero_warn(

  | Name  | Type                          | Params
--------------------------------------------------------
0 | model | BertForSequenceClassification | 108 M 
--------------------------------------------------------
108 M     Trainable params
0         Non-trainable params
108 M     Total params
435.680   Total estimated model params size (MB)
Epoch 0:   0%|          | 0/6251 [00:00<?, ?it/s]

原因分析 #

1. 检查点目录配置错误：将检查点目录设为D:\根目录，该目录非空且可能存在权限限制，导致模型无法正常写入检查点，阻塞训练流程。
2. 优化器配置不兼容：优化器配置中包含不被当前PyTorch Lightning版本支持的scheduler键，版本差异导致调度器加载失败，中断训练初始化。
3. 内存负载不匹配：本地DRAM容量小于Colab环境，或数据加载参数（如batch size、num_workers）设置不合理，导致内存耗尽但未触发OOM（可能是系统内存调度机制或PyTorch内存检测延迟）。
4. 环境版本不一致：本地与Colab的PyTorch、PyTorch Lightning、Transformers等依赖版本存在差异，API兼容性问题导致训练无法启动。

解决思路 #

• 修正检查点目录：为ModelCheckpoint指定一个空的、有读写权限的子目录，例如：

  from pytorch_lightning.callbacks import ModelCheckpoint
  checkpoint_callback = ModelCheckpoint(dirpath="./checkpoints", save_top_k=1)
  trainer = Trainer(callbacks=[checkpoint_callback])
  

• 修复优化器配置：按照PyTorch Lightning规范配置学习率调度器，在LightningModule的configure_optimizers()方法中返回优化器与调度器的组合：

  def configure_optimizers(self):
      optimizer = torch.optim.AdamW(self.parameters(), lr=2e-5)
      scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10)
      return [optimizer], [scheduler]
  

• 调整数据加载参数：降低batch size；设置train_dataloader的pin_memory=True（GPU环境下）；根据本地CPU核心数调整num_workers，避免内存过载；检查数据集是否存在重复加载的情况。
• 对齐依赖版本：运行pip list查看Colab中PyTorch、PyTorch Lightning、Transformers的版本，在本地执行对应版本的安装命令，例如：

  pip install torch==2.0.1 pytorch-lightning==2.0.3 transformers==4.30.2
  

• 排查内存细节：使用torch.cuda.memory_summary()查看GPU内存占用，或用系统任务管理器监控DRAM消耗，确认是模型还是数据加载导致的内存耗尽，针对性优化。

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