Pytorch使用BERT时遇CUDA断言失败及CPU索引越界错误如何解决

错误现象 #

在Pytorch环境结合CUDA运行BERT模型时触发如下断言错误：
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:702: indexSelectLargeIndex: block: [234,0,0], thread: [0,0,0] Assertion srcIndex < srcSelectDimSize failed

CUDA环境完整报错日志

Epoch [1/100]
Iter:      0,  Train Loss:   1.1,  Train Acc: 39.06%,  Val Loss:   1.0,  Val Acc: 51.90%,  Time: 0:00:04 *
Iter:     10,  Train Loss:  0.99,  Train Acc: 57.81%,  Val Loss:   1.0,  Val Acc: 52.01%,  Time: 0:00:11 *
Iter:     20,  Train Loss:   1.0,  Train Acc: 42.19%,  Val Loss:  0.99,  Val Acc: 52.01%,  Time: 0:00:17 *
Iter:     30,  Train Loss:   1.0,  Train Acc: 40.62%,  Val Loss:  0.99,  Val Acc: 52.12%,  Time: 0:00:23 *
Iter:     40,  Train Loss:   1.0,  Train Acc: 50.00%,  Val Loss:  0.98,  Val Acc: 52.12%,  Time: 0:00:29 *
Iter:     50,  Train Loss:   1.1,  Train Acc: 43.75%,  Val Loss:  0.98,  Val Acc: 52.12%,  Time: 0:00:35 *
Traceback (most recent call last):
  File "/content/drive/MyDrive/Prediction/run.py", line 38, in <module>
    train(config, model, train_iter, dev_iter, test_iter)
  File "/content/drive/MyDrive/Prediction/train_eval.py", line 50, in train
    outputs = model(trains)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/drive/MyDrive/Prediction/models/BERT+Covid.py", line 68, in forward
    output  = self.bert(context, attention_mask=mask)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py", line 1005, in forward
    return_dict=return_dict,
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py", line 589, in forward
    output_attentions,
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py", line 475, in forward
    past_key_value=self_attn_past_key_value,
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py", line 408, in forward
    output_attentions,
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py", line 323, in forward
    attention_scores = attention_scores / math.sqrt(self.attention_head_size)
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:702: indexSelectLargeIndex: block: [234,0,0], thread: [0,0,0] Assertion `srcIndex < srcSelectDimSize` failed.

#......I SKIPPED SEVERAL LINES DUE TO THE CHARACTER LIMITATION

/pytorch/aten/src/ATen/native/cuda/Indexing.cu:702: indexSelectLargeIndex: block: [235,0,0], thread: [127,0,0] Assertion `srcIndex < srcSelectDimSize` failed

CPU环境复现报错

切换CPU运行代码后，错误收敛为明确的索引越界问题，错误栈如下：

traceback (most recent call last):
  File "/content/drive/MyDrive/Prediction/run.py", line 37, in <module>
    train(config, model, train_iter, dev_iter, test_iter)
  File "/content/drive/MyDrive/Prediction/train_eval.py", line 49, in train
    outputs = model(trains)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/content/drive/MyDrive/Prediction/models/BERT+Covid.py", line 66, in forward
    output  = self.bert(context, attention_mask=mask, )
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py", line 993, in forward
    past_key_values_length=past_key_values_length,
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py", line 215, in forward
    inputs_embeds = self.word_embeddings(input_ids)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/sparse.py", line 160, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py", line 2043, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
IndexError: index out of range in self

已完成排查项 #

• 输入长度未超出模型最大长度：设置的pad size为98，报错行前打印输入形状确实为(batch_size, pad_size)
• 词表大小数值匹配：len(tokenizer)与model.config.vocab_size数值相等，不存在词表大小不匹配问题

模型结构代码 #

class Model(nn.Module):

    def __init__(self, config):
        super(Model, self).__init__()
        self.modelConfig = BertConfig.from_pretrained('./bert_pretrain/config.json')
        self.bert = BertModel.from_pretrained(config.bert_path,config=self.modelConfig)
        for param in self.bert.parameters():
            param.requires_grad = False
        self.cls_fc_layer = FCLayer(config.hidden_size, config.word_size, config.dropout_rate)
        self.label_classifier = FCLayer(
            config.word_size+config.numerical_size,
            config.num_classes,
            config.dropout_rate,
            use_activation=False,
        )

    def forward(self, x):
        context = x[0]  # input token ids
        mask = x[2]  # mask
        numerical=x[3] #size(batch_size,18)
        
        output  = self.bert(context, attention_mask=mask)
        pooled_output=output[1]
        ##size(batch_size,768)
        pooled_output = self.cls_fc_layer(pooled_output)
        ##size(batch_size,18)
        concat_h = torch.cat([pooled_output, numerical], dim=-1)
        ##size(batch_size,36)
        logits = self.label_classifier(concat_h)
        return logits

排查解决思路 #

1. 检查输入token id的数值范围，不要仅验证形状：在调用BERT的代码行前加入调试逻辑，打印当前batch的token id最大值、最小值以及模型词表大小：

print(context.max().item(), context.min().item(), self.bert.config.vocab_size)

该类索引越界90%以上的概率是存在token id数值大于等于词表大小，通常是数据预处理阶段生成了异常id，或者tokenizer处理未知字符时逻辑有误。
2. 统一BERT配置与预训练权重的加载路径：当前代码中BertConfig从./bert_pretrain/config.json加载，而BertModel从config.bert_path加载，两个路径如果指向不同的预训练包，哪怕词表大小数值一致，也可能出现内部索引不匹配问题，建议统一为同一路径加载。
3. 逐样本定位脏数据：将batch size设置为1，逐样本跑训练流程，定位到触发报错的具体样本，查看其分词后的token id是否合法，排查是否存在脏数据混入数据集的情况。
4. 验证tokenizer与预训练模型的匹配性：若你使用的是领域定制版BERT，需确认当前使用的tokenizer与预训练权重是配套发布的版本，避免因词表内容不匹配导致id越界。 #

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