PytorchでエラーRuntimeError: CUDA error: Device-side assert triggered

huggingfaceでDeBERTaを使っい、Masked Language Modeling（MLM）をした時に表題のエラー。

  File "/usr/local/lib/python3.8/dist-packages/transformers/models/deberta_v2/modeling_deberta_v2.py", line 341, in forward
    attention_output = self.attention(
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/deberta_v2/modeling_deberta_v2.py", line 272, in forward
    self_output = self.self(
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/deberta_v2/modeling_deberta_v2.py", line 682, in forward
    query_layer = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/linear.py", line 94, in forward
    return F.linear(input, self.weight, self.bias)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py", line 1753, in linear
    return torch._C._nn.linear(input, weight, bias)
RuntimeError: CUDA error: device-side assert triggered

attentionの計算中にエラーを吐いている様子。

input / output の shape が揃っていないとエラーが出ることが多いようだが、自分のデータセットを確認しても shape は合っている。

となるとモデルのどこかで shape が変わったのではないかと予想。

自分のケースでは、tokenizer に special token を追加していたので、token 増えた分だけモデルを拡張しないと駄目だった。

special_tokens_dict = {"additional_special_tokens": my_token_list}
tokenizer.add_special_tokens(special_tokens_dict)

model = AutoModelForMaskedLM.from_pretrained(cfg.model_name)
model.resize_token_embeddings(len(tokenizer))

これでエラーが消え、MLMを実行することが出来た。

参考

https://masatakashiwagi.github.io/portfolio/post/cuda-error-device-side-assert-triggered/