RTX 4090/3080上PyTorch Lightning GAN反向传播停滞问题求助

我用PyTorch Lightning写了个GAN模型，在Google Colab的Tesla T4上跑完全正常，但放到本地机房的RTX 4090或者3080上时，训练到判别器的self.manual_backward(discriminator_loss)这一步就彻底卡住，没任何反应。调过batch size，换了同机房的3080也一样，但在CPU上跑完全没问题。核心训练代码如下：

def training_step(self, batch, batch_idx):
    images, future_images = batch
    images = images.float()
    future_images = future_images.float()

    self.global_iteration += 1
    g_opt, d_opt = self.optimizers()
    ##########################
    # 优化判别器 #
    ##########################
    # 每轮生成器训练对应两轮判别器训练

    for _ in range(2):
        d_opt.zero_grad()
        predictions = self(images)

        generated_sequence = torch.cat([images, predictions], dim=1)
        real_sequence = torch.cat([images, future_images], dim=1)

        concatenated_inputs = torch.cat([real_sequence, generated_sequence], dim=0)
        concatenated_outputs = self.discriminator(concatenated_inputs)

        score_real, score_generated = torch.split(
            concatenated_outputs, [real_sequence.shape[0], generated_sequence.shape[0]], dim=0
        )
        score_real_spatial, score_real_temporal = torch.split(score_real, 1, dim=1)
        score_generated_spatial, score_generated_temporal = torch.split(
            score_generated, 1, dim=1
        )

        discriminator_loss = loss_hinge_disc(
            score_generated_spatial, score_real_spatial
        ) + loss_hinge_disc(score_generated_temporal, score_real_temporal)

        self.manual_backward(discriminator_loss)  # 就是这儿卡住
        d_opt.step()

    ######################
    # 优化生成器 #
    ######################
    predictions = [self(images) for _ in range(self.generation_steps)]
    grid_cell_reg = grid_cell_regularizer(torch.stack(predictions, dim=0), future_images)

    generated_sequence = [torch.cat([images, x], dim=1) for x in predictions]
    real_sequence = torch.cat([images, future_images], dim=1)
    
    generated_scores = []
    for g_seq in generated_sequence:
        concatenated_inputs = torch.cat([real_sequence, g_seq], dim=0)
        concatenated_outputs = self.discriminator(concatenated_inputs)

        score_real, score_generated = torch.split(
            concatenated_outputs, [real_sequence.shape[0], g_seq.shape[0]], dim=0
        )
        generated_scores.append(score_generated)
    generator_disc_loss = loss_hinge_gen(torch.cat(generated_scores, dim=0))
    generator_loss = generator_disc_loss + self.grid_lambda * grid_cell_reg
    g_opt.zero_grad()
    self.manual_backward(generator_loss)
    g_opt.step()

1. 先查CUDA和PyTorch版本兼容性 #

RTX 4090/3080属于安培架构，对CUDA版本有明确要求，别用老版本PyTorch搭配新CUDA，或者反过来。

• 先跑命令验证当前环境：

  python -c "import torch; print(torch.version.cuda); print(torch.cuda.is_available())"
  

• 如果版本不匹配，卸载重装对应版本的PyTorch，比如CUDA 11.8搭配PyTorch 2.0+：

  pip uninstall torch torchvision torchaudio -y
  pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
  

2. 检查梯度NaN/Inf导致的停滞 #

反向传播卡住大概率是梯度爆炸或出现NaN/Inf，在卡住的代码前后加检测逻辑：

# 计算完损失后先打印数值
print(f"判别器损失值: {discriminator_loss.item()}")
# 反向传播后检查每个参数的梯度
self.manual_backward(discriminator_loss)
for name, param in self.discriminator.named_parameters():
    if param.grad is not None:
        if torch.isnan(param.grad).any() or torch.isinf(param.grad).any():
            print(f"{name} 参数的梯度存在NaN/Inf！")
            break

如果确实存在异常梯度，修改Hinge损失的实现，添加小偏移保证数值稳定：

def loss_hinge_disc(score_generated, score_real):
    # 原实现可能是torch.mean(torch.relu(1 - score_real)) + torch.mean(torch.relu(1 + score_generated))
    # 修改为带小epsilon的版本
    return torch.mean(torch.nn.functional.relu(1.0 - score_real + 1e-8)) + torch.mean(torch.nn.functional.relu(1.0 + score_generated + 1e-8))

3. 关闭自动混合精度试试 #

安培架构的自动混合精度偶尔会有兼容性问题，强制使用32位浮点精度：
在模型的__init__方法中添加：

def __init__(self, ...):
    super().__init__()
    self.automatic_optimization = False  # 已使用手动反向传播，确认此参数已开启
    self.precision = "32-true"  # 强制32位精度，关闭混合精度

或者在初始化Trainer时指定：

trainer = Trainer(
    precision="32-true",
    # 其他训练参数...
)

4. 排查显存隐性溢出 #

虽然调整了batch size，但RTX 4090显存更大，可能某些张量拼接操作在大显存下分配了过多内存，导致隐性死锁。添加显存监控：

# 反向传播前后打印显存占用
print(f"反向传播前显存: {torch.cuda.memory_allocated()/1024**2:.2f} MB")
self.manual_backward(discriminator_loss)
print(f"反向传播后显存: {torch.cuda.memory_allocated()/1024**2:.2f} MB")

如果显存占用异常，在判别器循环的每次迭代后清理缓存：

for _ in range(2):
    # ... 原有判别器训练代码 ...
    d_opt.step()
    torch.cuda.empty_cache()  # 清理未使用的显存

5. 调整CUDA后端参数 #

安培架构对cudnn的设置比较敏感，尝试关闭benchmark模式并开启确定性计算：

import torch
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

或者禁用TF32（安培架构默认开启，可能导致数值不稳定）：

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

6. 确认手动反向传播逻辑正确性 #

检查手动反向传播的流程是否规范：

• 每次判别器迭代前是否调用了d_opt.zero_grad()
• 反向传播后是否有未释放的计算图，可在d_opt.step()后添加discriminator_loss.detach()

优先排查CUDA版本兼容性和梯度NaN/Inf问题，这是安培架构GPU上反向传播停滞的常见原因。如果以上方法都无效，可尝试降级PyTorch到1.13.x这类稳定版本，或者检查机房的GPU驱动是否为最新版本。

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