PyTorch Lightning捕获参数梯度时始终返回None的问题求助

在PyTorch Lightning中需要对梯度执行相关操作，已确认模型权重随步骤更新（每步权重变化、损失逐步下降），但无法捕获梯度：

1. 最初尝试在pl.LightningModule中实现on_after_backward钩子，打印并记录梯度，结果始终得到None：

def on_after_backward(self):
    for p in self.parameters():
        print(p.grad)
    norms = {n: torch.norm(p.grad) for n, p in self.named_parameters()}
    self.log_dict(norms)

2. 关闭自动优化后手动调用backward，training_step代码如下，opt.step()前后权重有变化，但两次打印的梯度仍为None：

def training_step(self, batch, batch_idx):
    if self.trainer.is_global_zero:
        print()
        print("################")
        print(list(self.parameters())[-1].grad)
    args = self.args
    opt = self.optimizers()
    opt.zero_grad()
    idx, targets = batch
    logits = self(idx)
    loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
    loss = L2Wrap.apply(loss, logits)
    self.manual_backward(loss)
    if self.trainer.is_global_zero:
        print(list(self.parameters())[-1].grad)
        print("################")
        print()
    opt.step()
    return loss

1. 排查L2Wrap自定义操作的实现 #

问题大概率出在L2Wrap.apply(loss, logits)这个自定义torch.autograd.Function上：

• 检查L2Wrap的forward方法是否通过save_for_backward保存了反向传播所需的张量（比如logits）
• 确认backward方法是否正确计算并传递梯度到输入张量（尤其是logits），如果反向逻辑没有把梯度传递给logits，模型参数的梯度就无法被累积。

2. 临时移除自定义损失包装验证 #

先注释掉loss = L2Wrap.apply(loss, logits)，直接用原始损失调用self.manual_backward(loss)，再检查梯度是否正常。如果此时梯度能正常打印，说明问题完全来自L2Wrap的实现缺陷。

3. 调整手动优化的梯度检查逻辑 #

在手动优化模式下，确保：

• 调用self.manual_backward(loss)后、opt.step()前立刻检查梯度，不要在这中间执行可能清空梯度的操作
• 避免在training_step开头过早调用opt.zero_grad()（除非你明确需要提前清空）

4. 验证参数的requires_grad状态 #

遍历所有参数确认是否开启梯度追踪：

for name, param in self.named_parameters():
    print(f"{name}: requires_grad={param.requires_grad}")

如果有参数被意外设置为requires_grad=False，梯度会返回None——不过你的权重在更新，这个可能性较低，但可以作为兜底排查项。

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