loss.backward()异常致梯度范数为0的模型参数优化问题

import gc
from tqdm.notebook import tqdm
import matplotlib.pyplot as plt

torch.set_printoptions(sci_mode=False)

gradient_norms = []
losses = []

for epoch in tqdm(range(epochs)):
    model_path = f"/kaggle/working/Training/ace_state_dict_{epoch+1}.pth"
    torch.save(model.state_dict(), model_path)
    model.train()
    total_loss = 0
    
    for batch_idx, batch in enumerate(tqdm(train_dataloader, desc=f'Epoch {epoch + 1}/{epochs}')):
        optimizer.zero_grad()
        logits = model(batch["inputs"])
        targets = batch["targets"]
        loss = loss_fn(logits.view(-1, logits.size(-1)), targets.float()) / 1000000000
        loss.backward()        
        # Compute gradient norms
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
        gradient_norms.append(grad_norm)
        optimizer.step()
        scheduler.step()
        total_loss += loss.item()
        
        if batch_idx % 100 == 0:
            print(f'Batch {batch_idx}/{len(train_dataloader)}, Loss: {total_loss/(batch_idx+1)}, Gradient Norms: {grad_norm}')
    
    avg_loss = total_loss / len(train_dataloader)
    losses.append(avg_loss)
    print(f"Train loss: {avg_loss}")

问题描述 #

上述代码中loss.backward()未正常工作，导致grad_norm始终为0.0，阻碍模型参数的迭代优化。已尝试调整超参数、学习率，使用torch.nn.utils.clip_grad_norm_()，且相关张量和模型参数均设置requires_grad=True，问题仍存在。

解决建议 #

• 移除损失值的极端缩放：代码中将损失除以了1e9，这会把梯度直接压缩到几乎为0的量级，是梯度消失的核心原因之一。先删除该除法操作，观察梯度是否恢复正常；若确实需要损失缩放，建议使用PyTorch官方的混合精度训练工具torch.cuda.amp进行合理处理。
• 校验损失计算逻辑：确认loss_fn的输入形状、数据类型是否匹配，比如logits的reshape操作是否正确，targets.float()是否符合损失函数的要求，同时检查损失值本身是否恒为0。
• 检查模型参数的梯度追踪状态：遍历模型所有参数，打印param.requires_grad，确保所有需要优化的层都开启了梯度追踪，避免部分层被意外设置为requires_grad=False。
• 验证优化器配置：确认优化器初始化时正确传入了模型的可训练参数（如optimizer = torch.optim.Adam(model.parameters(), lr=...)），没有遗漏参数组。
• 调整学习率调度器的调用时机：当前代码在每个batch后调用scheduler.step()，若使用的是基于epoch的调度器（如StepLR），会导致学习率被错误更新。建议将scheduler.step()移到epoch循环的末尾，即每个epoch结束后再更新学习率。
• 排查模型前向传播输出：检查model(batch["inputs"])的输出logits是否在训练过程中发生变化，若输出恒为固定值，会导致损失不变、梯度为0，需排查模型的前向传播逻辑。

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