基于PyTorch Lightning创建Python包时遇CUDA初始化错误求助

问题场景 #

创建基于PyTorch Lightning的Python包时，若包内源文件包含import pytorch_lightning as pl语句，导入包会触发如下错误：

torch.cuda.DeferredCudaCallError: CUDA call failed lazily at initialization with error: device >= 0 && device < num_gpus INTERNAL ASSERT FAILED at "../aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch.

核心原因 #

PyTorch Lightning导入时会默认尝试初始化CUDA上下文，若当前环境无可用GPU、CUDA驱动与PyTorch版本不兼容，或设备配置异常，就会触发该错误。

解决方法 #

• 延迟导入PyTorch Lightning：避免在包的顶层代码（如__init__.py或模块全局范围）直接导入pl，仅在实际使用的函数/类内部导入：

  # 错误示例（顶层导入）
  import pytorch_lightning as pl
  
  class MyModel:
      def __init__(self):
          self.module = pl.LightningModule()
  
  # 正确示例（延迟导入）
  class MyModel:
      def __init__(self):
          import pytorch_lightning as pl
          self.module = pl.LightningModule()
  

• 强制使用CPU模式：在导入pl前，强制PyTorch跳过CUDA初始化：

  import torch
  # 模拟无CUDA环境
  torch.cuda.is_available = lambda: False
  
  import pytorch_lightning as pl
  

• 检查CUDA环境有效性：若环境有GPU仍报错，先验证基础CUDA配置：

  import torch
  print("CUDA可用:", torch.cuda.is_available())
  print("GPU数量:", torch.cuda.device_count())
  

  若输出False或设备数为0，需重新安装与CUDA驱动匹配的PyTorch版本。

• 修改PL初始化环境变量：通过环境变量禁用自动CUDA初始化：

  import os
  # 指定CPU通信后端
  os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "gloo"
  # 隐藏所有GPU设备
  os.environ["CUDA_VISIBLE_DEVICES"] = ""
  
  import pytorch_lightning as pl
  

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