运行BERT模型执行NER任务时遭遇CUDA内存不足错误求助

我看你已经把batch size调到1还是碰到了CUDA显存不足的问题，结合你的GPU（3.82GiB显存）和提供的代码，整理了几个针对性的解决方案，你可以逐一尝试：

一、先优化验证阶段的冗余计算（最见效的快速调整） #

看你的验证代码，你在同一个batch里两次调用了model做前向传播：

tmp_eval_loss = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)
logits = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask)

这会重复计算一次模型前向，额外占用不少显存。其实BERT的model()调用返回的outputs里同时包含loss和logits，改成一次调用就能搞定：

with torch.no_grad():
    outputs = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)
    tmp_eval_loss = outputs.loss
    logits = outputs.logits

这个改动能立刻减少一次前向传播的显存开销，应该能缓解不少压力。

二、降低序列长度（显存占用和序列长度平方成正比） #

BERT的显存消耗和输入序列长度的平方正相关，如果你的max_seq_len设得比较大（比如默认的512），可以尝试降到128或者256——大部分NER任务的文本都不会太长，这个调整对最终效果影响很小，但能大幅减少显存占用。

三、用梯度累积模拟大batch，不增加单步显存压力 #

虽然你已经用了batch size=1，但可以通过梯度累积来模拟更大的batch效果，同时不增加单步的显存负载：
修改你的训练循环，比如每累积4步的梯度再更新一次参数：

epochs = 10
max_grad_norm = 1.0
accumulation_steps = 4  # 可根据显存情况调整，比如2、4、8

for _ in trange(epochs, desc="Epoch"):
    model.train()
    tr_loss = 0
    nb_tr_examples, nb_tr_steps = 0, 0
    for step, batch in enumerate(train_dataloader):
        batch = tuple(t.to(device) for t in batch)
        b_input_ids, b_input_mask, b_labels = batch
        
        # 前向传播，损失除以累积步数
        loss = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)
        loss = loss / accumulation_steps
        loss.backward()
        
        tr_loss += loss.item() * accumulation_steps
        nb_tr_examples += b_input_ids.size(0)
        nb_tr_steps += 1
        
        # 每accumulation_steps步更新一次参数
        if (step + 1) % accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=max_grad_norm)
            optimizer.step()
            model.zero_grad()
    # 后续验证代码不变...

这样既保证了训练的稳定性，又不会增加单步的显存占用。

四、换成轻量级BERT变体 #

如果以上方法还不够，可以考虑用参数更少的预训练模型，比如：

• 把bert-base-uncased换成distilbert-base-uncased：参数少40%，推理速度快60%，效果和base版BERT几乎一致
• 或者选用bert-small这类超轻量模型
  这些模型的显存占用会显著降低，非常适合小显存GPU。

五、启用自动混合精度训练（AMP） #

PyTorch的自动混合精度（AMP）可以在不损失模型精度的前提下，把部分计算从FP32转为FP16，大幅减少显存占用：
首先导入AMP相关模块：

from torch.cuda.amp import GradScaler, autocast

然后初始化梯度缩放器，修改训练循环的前向和反向传播逻辑：

scaler = GradScaler()  # 放在epochs循环外

for _ in trange(epochs, desc="Epoch"):
    model.train()
    tr_loss = 0
    nb_tr_examples, nb_tr_steps = 0, 0
    for step, batch in enumerate(train_dataloader):
        batch = tuple(t.to(device) for t in batch)
        b_input_ids, b_input_mask, b_labels = batch
        
        with autocast():  # 启用混合精度前向传播
            loss = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)
        loss = loss / accumulation_steps
        scaler.scale(loss).backward()  # 缩放损失避免梯度下溢
        
        tr_loss += loss.item() * accumulation_steps
        nb_tr_examples += b_input_ids.size(0)
        nb_tr_steps += 1
        
        if (step + 1) % accumulation_steps == 0:
            scaler.unscale_(optimizer)  # 反缩放梯度
            torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=max_grad_norm)
            scaler.step(optimizer)  # 更新参数
            scaler.update()  # 更新缩放器
            model.zero_grad()
    # 验证阶段也可以加上with autocast(): 进一步减少显存
    model.eval()
    eval_loss, eval_accuracy = 0, 0
    nb_eval_steps, nb_eval_examples = 0, 0
    predictions , true_labels, true_inputs = [], [],[]
    for batch in valid_dataloader:
        batch = tuple(t.to(device) for t in batch)
        b_input_ids, b_input_mask, b_labels = batch
        with torch.no_grad(), autocast():  # 验证阶段也启用混合精度
            outputs = model(b_input_ids, token_type_ids=None, attention_mask=b_input_mask, labels=b_labels)
            tmp_eval_loss = outputs.loss
            logits = outputs.logits
    # 后续处理不变...

六、手动清理冗余显存 #

在训练和验证循环中，及时把不需要的张量移回CPU或者删除，并手动释放显存：
比如验证阶段处理完数据后，添加：

del logits, label_ids, inputs
torch.cuda.empty_cache()

注意不要太频繁调用empty_cache()，会稍微影响速度，但显存紧张时效果很明显。另外，检查你的nvidia-smi输出，确认没有其他进程占用GPU显存，如果有，关闭这些进程释放资源。

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