LLMの学習で RuntimeError: CUDA error: device-side assert triggered と言われた (Transformers)

Transformers の Trainer を使ってこのエラーがでるときはだいたい学習するデータセットの input_ids のトークン ID がモデルの vocab_size を超えるときに発生する気がする。

• モデルの語彙数(vocab_size)がトークナイザーの語彙と合っているか
  既存のモデルにスペシャルトークンを追加して学習するときは特に気をつけるべき。

tokenizer = AutoTokenizer.from_pretrained(
    TOKENIZER_NAME,
    cache_dir=CACHE_DIR,
)

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_NAME,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    cache_dir=CACHE_DIR,
)
model.config.use_cache = False  # type: ignore
# ↓大事
model.resize_token_embeddings(len(tokenizer))  # type: ignore

でモデルを読み込んだあとに resize_token_embeddings(len(tokenizer)) で語彙数を合わせる。

もしあわせてもこのエラーが出る場合は、トークナイザーが壊れているかなにかの理由で最大のトークン ID の値が狂っているので、データセットを確認してそのようなトークン ID が発生していないかを調べるとよさそう。

TransformersのTrainerで学習している状態。

学習自体は普通に始まるのだが、途中のステップで突然

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. compile with `torch_use_cuda_dsa` to enable device-side assertions.

とかいうエラーが出て落ちた。

resume してもこれなのでちょっと謎。

なにかの index エラーっぽいのだが学習開始はするから謎...

https://yuiga.dev/blog/posts/pytorchcuda_error_device-side_assert_triggered_解決の手引き/

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

https://stackoverflow.com/questions/68166721/pytorch-fails-with-cuda-error-device-side-assert-triggered-on-colab

どうやら input_ids がモデルの vocab_size を超えていたようだ

model.resize_token_embeddings(len(tokenizer))

を実行していたのに vocab_size よりも input_ids が出現していた...

データセットを確認してみると確かに len(tokenizer) よりも 1 大きい input_ids が見つかって非常に謎。

学習スクリプトとデータセットのトークナイズを別に分けていたため、そこで誤ったトークナイザーを使ってしまったのか...?

壊れた(?)トークナイザーを使ってトークナイズ、学習に利用していたことが原因だったっぽい。

単語が何か抜けてしまっているのか、len(tokenizer) が最大のトークンIDより小さくなってしまっていた。

自前でトークナイザーを学習して用意して発生していたので、今回はもう一度トークナイザーを作り直し、データセットも再トークナイズすることで解決できた。