- 日本語データセット
- 日本語トークナイザ
で追加事前学習(incremental pre-training)を行います.
事前に日本語データセットと日本語トークナイザを用意する必要があります.
- tinyllama 1.1B を incremental pretrain
- llama2 7B
- Python 3.10+
- pytorch
- pytorch lightning
$ conda create -n jp-llama-pretrain python=3.10
$ conda activate jp-llama-pretrain
$ python -m pip install -r requirements.txt
$ bash ./run_jp_tinyllama_train.sh
最初にデータセットの tokenize や grouping を行います. 結果はキャッシュされます. CharShu 59B(展開時180GB) では, これらのキャッシュとして概ね 800 GB を必要とします. tokenize 処理には 2~3 時間ほどかかります.
batch=1 で 20 GB ほどメモリを消費します.
transformers に取り込まれた? flash_attn_2 を使います. ROCm の場合は現状(2023/09/15 時点) MI250(gfx90a) 以外対応していないようので, off にします.
checkpokpoint から merge して publish するスクリプト publish_checkpoint.py
model.embedding_size != tokenizer.nvocab のため, llama.cpp などほかのツールではこのままだとうまく扱えないため, うまく利用できるように, model をマージ後, tokenizer の nvocab で model.embedding をリサイズしています. (これにより特定の HW では性能が低下するかもしれません)
TODO: model の embedding と, tokenizer の nvocab を 64 の倍数に合わせる. あとから処理するといろいろめんどいので, Japanese tokenizer を作成の時点で 64 の倍数にしておくのが理想...
- 学習用にトークナイズなどする処理の別 script 化
- libtorch or llama.cpp で CPU クラスタで incremental pretrain できるようにする.
- lit llama: Apache 2.0 https://github.com/Lightning-AI/lit-llama