使用PyTorch Lightning训练BERT评论分类模型遇RuntimeError问题求助

问题描述 #

训练过程中触发以下错误：

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

背景信息：

• 已通过自定义enable_gradients(model)函数开启所有参数梯度，但错误未解决
• 基于预训练模型aubmindlab/bert-base-arabertv02-twitter构建模型，加载时发现部分BERT权重未正确初始化
• 已确认PyTorch、Transformers、PyTorch Lightning为最新版本
• 尝试先在下游任务预训练BERT再训练目标模型，问题依然存在

相关代码 #

from pytorch_lightning import Trainer

def enable_gradients(model):
    for param in model.parameters():
        param.requires_grad = True

# datamodule
ucc_data_module = UCC_Data_Module(train_path, val_path, test_path, attributes=attributes, batch_size=config['batch_size'])
ucc_data_module.setup()

# model
model = UCC_Comment_Classifier()

enable_gradients(model)

# trainer and fit
# Instantiation of the Lightning Trainer
trainer = Trainer(max_epochs=config['n_epochs'], accelerator='gpu', num_sanity_val_steps=1)

try:
    trainer.fit(model, ucc_data_module)
    torch.save(model.state_dict(), PATH)
except RuntimeError as e:
    print(e)

完整报错栈 #

ProcessRaisedException: 

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/multiprocessing.py", line 
147, in _wrapping_function
    results = function(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 568, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 973, in _run
    results = self._run_stage()
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1016, in _run_stage
    self.fit_loop.run()
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 201, in run
    self.advance()
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 354, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 133, in run
    self.advance(data_fetcher)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 218, in 
advance
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 185, in 
run
    self._optimizer_step(kwargs.get("batch_idx", 0), closure)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 260, in 
_optimizer_step
    call._call_lightning_module_hook(
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 144, in 
_call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1256, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 155, in step
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 256, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 225, in 
optimizer_step
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 114,
in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 69, in wrapper
    return wrapped(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/optim/optimizer.py", line 280, in wrapper
    out = func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/optimization.py", line 439, in step
    loss = closure()
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 101,
in _wrap_closure
    closure_result = closure()
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 140, in 
__call__
    self._result = self.closure(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 135, in 
closure
    self._backward_fn(step_output.closure_loss)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 232, in 
backward_fn
    call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 291, in 
_call_strategy_hook
    output = fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 200, in backward
    self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 67, 
in backward
    model.backward(tensor, *args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1046, in backward
    loss.backward(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_tensor.py", line 487, in backward
    torch.autograd.backward(
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/__init__.py", line 200, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

解决方案 #

1. 修正模型内部的梯度设置逻辑 #

外部调用enable_gradients可能没有正确作用到BERT子模块，建议在UCC_Comment_Classifier的__init__中直接配置梯度：

import pytorch_lightning as pl
from transformers import AutoModel
import torch.nn as nn

class UCC_Comment_Classifier(pl.LightningModule):
    def __init__(self, num_classes):
        super().__init__()
        self.bert = AutoModel.from_pretrained("aubmindlab/bert-base-arabertv02-twitter")
        # 直接开启BERT所有参数的梯度
        for param in self.bert.parameters():
            param.requires_grad = True
        # 定义分类头
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_classes)
        # 损失函数
        self.loss_fn = nn.CrossEntropyLoss()

2. 确保优化器包含所有可训练参数 #

在configure_optimizers方法中，直接传入模型所有参数，避免遗漏：

from transformers import AdamW

def configure_optimizers(self):
    optimizer = AdamW(self.parameters(), lr=2e-5)
    return optimizer

3. 检查损失计算逻辑 #

确保损失张量没有脱离计算图：

• 不要在training_step中使用loss = loss.detach()
• 验证损失的梯度属性，可在training_step中添加打印：print(loss.requires_grad)，确认输出为True

4. 修复DDP模式下的梯度同步问题 #

报错显示使用多GPU DDP模式，可尝试以下调整：

• 移除外部的enable_gradients调用，将梯度开启逻辑放在模型内部
• 初始化Trainer时添加参数：strategy="ddp_find_unused_parameters=False"（需确保所有参数都参与计算）

5. 重新加载预训练权重解决初始化问题 #

针对部分权重未初始化的情况，强制重新下载并加载模型：

self.bert = AutoModel.from_pretrained(
    "aubmindlab/bert-base-arabertv02-twitter",
    force_download=True,
    ignore_mismatched_sizes=True  # 处理层维度不匹配的情况
)

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