🏋️‍♀️

LiLM 小規模言語モデル TinyLlama 1.1B の日本語追加事前学習(incremental pretrain) を試したメモ

2023/12/01に公開

事前学習

背景

日本語性能のよい軽量 LLM, LiLM, Lightweight Language Model ほしい...
- スマホで動かしたり, ブラウザで動かしたり...
- ドメインは知識応答系. Code 生成とか数学とかのドメインは今は考えない.
Chinese LLaMa https://zenn.dev/syoyo/scraps/6c3e92402e6fd0 でいい感じに incremental pretrain(追加事前学習) するといけるっぽいことがわかった!
ん－, でも 7B or 14 B 規模で試して本当にうまくいくのかわからん...
- (後日たぶん似たような方法で rinna ちゃんが Youri 7B, ELYZA ちゃんが Japanese LLaMa 7B 出してきた! それなりにいい感じになったようだよ)
あと 7B だとやっぱりまだちょっと大きい. 量子化してギリ 4 GB に収まるくらい(WASM はギリだし, iPhone もギリ?). 1 ~ 2B 規模でやりたい...
ちょうどいいタイミングで TinyLlama 1.1B があった! https://github.com/jzhang38/TinyLlama
これに対して日本語データセットで incremental pretrain したらどうか!

2023/10 時点で, 公開されているいい感じの日本語データセット(filtering & dedup 済み)と日本語で追加事前学習するのは前例がなかったため(解説記事みたいなのはあったかもであるが, コード公開しているところは, HojiChar https://github.com/HojiChar/HojiChar のコーパス処理以外無かった), いろいろ手探りでやらねばなりませんでしたが, 以外と最初の一回の学習試行でそこそこいい感じの結果は得られました😸

とりままとめ

日本語性能は獲得できました!
でも日本語しか話せなくなりました...
生成の性能(かしこさ)は微妙でした
- 生成される日本語はまあまあであるが, 構文やコンテキストがおかしい...
- ファインチューンしても間違えたり...
- まあでも 1B 規模なら妥当なのかもしれません
途中でいろいろ問題を見つけたので, 一カ月学習した 30% の時点でいったん打ち切りました.
- いろいろ直して再学習を予定
- ちょうど TinyLLama の 3 T tokens 学習版(2023/12/20 完予定)が出たところを見計らって

途中まで incrmental pretrain したのは公開するかは不明です(リーガルチェックなども必要のため. ただ WASM でブラウザ動作デモできるようにやっぱりどこかの時点で公開するかも)

学習のレシピ

学習を再現するレシピ(スクリプトとコード)の一式は

にあります.

2023/12/01 時点で, 日本語のデータセット(コーパス)構築と日本語 LiLM pretrain のコードをを十分な再現性をもって公開しているのはこれだけだと思います.
(さらにいえば, 日本語で incremental pretrain するコードを公開しているところはどこもない...)

Chinese LLama のやりかたをベースにしました.

どこのご家庭にもある RTX 3090 x 2 + メモリマシマシマシン
1. 日本語データセット CharShu 50B tokens 規模を用意
1. 日本語トークナイザー学習
1. 日本語トークナイザーと日本語データセットで TinyLLama 1.1B を incremental pretain
1. ファインチューンお試し
1. llama.cpp で推論, WASM 対応(ブラウザ動作のため)
1. TODO: モデル評価
1. 学習を振り返って...
その他の話題

ハードウェア構成

x299 + 256 GB CPU mem
RTX 3090 x 2
4 TB NVMe or SSD

追加事前学習では CPU mem は 160 GB くらい使いました. batch size 減らせば 128 GB mem でも動作すると思いますが, 192 GB or 256 GB あると安心です.

データセットが中間表現やらで結構ストレージ領域を使います. 一応 2 TB あればいけると思いますが, 余裕をみて 4 TB 用意しておくとよいでしょう.

日本語データセット準備やトークナイザー学習は CPU only で 128 GB mem でいけます.
今回は別 PC の 128 GB mem + Ryzen 3950X で行いました.

quick review: incremental pretain 最初の試行(2023/10/E ~ 2023/11/E)

incremental pretain スクリプトはこちら.
(Chinese LLaMa 2 の incremental pretrain スクリプトを参考にしています)

学習は 225 W に powerlimit した 3090 x 2 でおよそ 100 日となりました.

学習は TinyLLama 1.1B 1T tokens 学習版(2023/10/2 日あたりリリース)を使い, 10/E(10 月下旬)から 11/E(11 月下旬)の 1 カ月ほど行いました.

学習しはじめ

2% 学習(2 日)の時点で, 日本語性能はそこそこ獲得してくれました.

10 % 程度学習終わった段階での結果

ちょこちょこ checkpoint を評価した時点で, 学習を再度やりなおしたくなるような不都合(主にはデータセット)がいくつか見つかったため, 今回は 30%(概ね 30 日)終了した時点で止めました.

1. 日本語データセット(コーパス)

今回は CharShu(仮プロジェクト名. 文字集)を使いました.

日本語データセットの構築は, RefinedWeb を参考にし, 2023/07 月くらいから行い, 2023/10 月にはおおむねできていました.

いくつか解説はこちら.

LLM 向け日本語データセットの整備メモ

The RefinedWeb Dataset for Falcon LLM のメモ

filtering & dedup して 40 ~ 50B tokens 規模のデータセット(180 GB)となっています.
データセットそのものは huggingface dataset に private にアップロードしていますが, リーガルチェックなどのため公開していません. しかしスクリプトで再現することはできます.

C++ による人類史上最速の dedup を行っていますので, ぜひ手元で動かしてデータセット構築してみてくださいネ.

LLM 日本語データセット向けに C++ で minhash 重複除去を行うメモ

優秀な若人さまによるリコメンドがあります!

いろんなライブラリ/レポジトリ試してきたなかで現実的な時間で終わるの 68 才主ふさんのしかなかったけど

2. 日本語トークナイザー学習

今回は wikipedia データに対して, Unigram で学習しました.
最近のナウは Unigram っぽいようですが, BPE(BytePairEncoding)でもよいかと思います. LLM 用途では, Unigram か BPE かはあんまり性能に違いはないカナと思います.

Unigram では vocab 数少ないのは作れず(Unigram は基本 vocab 数を減らしていくアルゴリズム... のはずで, 指定した vocab 数以上で入力のデータセットを表現できないと学習エラーになる), 21,500 くらいが下限となりました.

これを Chinse LLaMa のトークナイザー merge スクリプトで, merge し, vocab 51470 となりました.

vocab 数は LLM 向けは 30K ~ 60K あたりが標準になっていますが, GPT-4 とかだと 100k(コンテキスト長次第?) になっています.

コンテキスト長が大きい LLM を扱う場合は 100k 規模(llama base 32K + 日本語 vocab 60K とか)でもいいかも... しれません.

学習の準備で見つかった不都合

huggingface datasets の問題

データセットのロードには huggingface datasets(ややこしい名前...)を使いましたが,

キャッシュやらダウンロードファイル保持やら, train 用に事前に token ids を batch 化(して保存)で, 結構データサイズ増えます... デフォだと Apache arrow で非圧縮で再保存したりするので...

元が展開後 180 GB でしたが, 1.6 TB くらい領域が必要となりました.

100B tokens 以上を扱う場合は, 自前でデータセット管理するなどしたほうがよいでしょう.
token ids は整数値なので, zstd でもいいですが lz4 などの整数特化(のはず)の高速圧縮アルゴリズムの利用も検討してもよいでしょう.

LoRA rank?

Chinese LLaMa ですと, 学習させる token 量で rank を決めていました.
(最大 64 っぽい)

このあたり設定どうしたらよいのかはようわかりませんでした.
とりあえずは TinyLLama だと context 2k なのでランク増えると思い, r=256 にしてみました. ただ過剰だったかもしれません.

rank 値はどう決めればよいのか... 事前にわかる方法はないのでしょうか...
学習させてみて, 実際の weight の rank がどうなるか見て判断するしかないのかしらん...

3. 日本語トークナイザーと日本語データセットで incremental pretrain

あとはスクリプト回すだけです!

日本語データセット CharShu は, 文章に対して KenLM で品質スコアリングがされていますが, 今回は低スコアの文章も使いました.
token 数水増しと, ネットスラングみたいたなのも扱えるのを期待してでした.

ただ学習してみて, やっぱり低スコアの文章は使わなかったほうがよかった気がします...

学習サーバ

3090 x 2 学習サーバは, 30 日ほどの学習期間を含め 42 日無問題で動いてくれました.

GPU マイニングで培った運用管理と powerlimit 管理で, いぶし銀のように鍛えられたおかげです!

GPU 温度も, 11 月始めの夏日(25 度)でも 69 度くらいで排熱の問題もナシ!

4. ファインチューンお試し

oasst1-89k-ja https://github.com/kunishou/oasst1-89k-ja でﾌｧｲﾝﾁｭｰﾝしてみました.

書きなぐりですが, 再現コード: https://github.com/lighttransport/japanese-llama-experiment/tree/main/20_finetune

結果はあんまりよくなかったです.
ある程度ファインチューン用データに寄って(アライン)はくれるようですが, 十分ではありませんでした.

pretrain 時点で wikipedia ja データは使っていないので, 知識関連は wikipedia ja のデータセットも含めるとよいかもしれません.

wikipedia データも使って事前学習している rinna-1b はうまく答えてくれました.

5. llama.cpp 変換と WASM 対応

で GGUF 変換して llama.cpp で動かせます. converter に fix がいりますが, それを取り込んでいます.

WASM/WASI

とりま Emscripten で llama.cpp をコンパイルし, node.js で動かしました.
cmake であれば特に変更はありませんが, 最後リンクするときに,
デフォのメモリ設定では足りないので, コンパイルフラッグで TOTAL_MEMORY と STACK_SIZE を適宜増やしておきます.

4bit はいけますがやはり壊れた出力が出たりと微妙なところでした.

WASM でブラウザ動作の場合は 6~8bit で 1GB のモデルサイズくらいがよいでしょう.
(余裕があれば fp16 2GB で)

WASI 動作の場合はスレッド対応を考える必要があります(一応最近(2023/11)だと native threading 対応しているっぽいようであるが)

6. TODO: モデル評価

人目でみてあきらかに ELYZA task 100 などこなせるほどのかしこさはありませんでしたので, モデル評価は行っていません.

7. 追加事前学習してみて

Chinese LLaMa paper & wiki ではあんまり詳しいことは書かれていません.

ソースコードはありますが, ファイル一個しかなく完全再現できる形でもないため, ようわかりません.

Chinese LLama は, 少し前?では

vocab 拡張
英語データ(?)で継続学習
- embedding 再学習
中国語データで継続学習

という手順だったような気がする(最近 wiki だとなんか手順が削除されている)

にありますように, ファインチューン済みのモデルに対して追加事前学習したほうがいいのかもしれません. あと, 日本語しか話せなくなったため, ELYZA ちゃんの記事にありますように, 継続学習(Continual Leaning)や, 追加事前学習の段階でも英語データを混ぜて学習したほうがいいのかも.

今回は Chinse LLaMa にある二段階学習はしなかったため(中国語 120 GB(40B tokens 相当?)ではしていなかったため),vocab 拡張した embedding は日本語データセットだけで学習させることになってしまい, 英語(既存知識)を失念しやすくなってしまったのかもです.

英語を混ぜて学習は, すでに学習しているものに対して, 同じトークンを繰り返し与えることによる学習能力の低下を危惧していましたが,
(e.g. SlimPajama で学習した TinyLLama に, さらに SliPajama のデータセットを与えて学習)

1B だと 100 個くらいまで OK っぽいので, やっぱり問題ないかも.

つまり,

vocab 拡張 tokenizer で英語データセットで (embedding だけ)再学習
次にいくらか英語データセット(SlimPajama)をまぜつつ日本語データセットで学習

計算性能?

225W powerlimit 3090 x 2 で LoRA trainable parameter 66% くらいで 50 B tokens くらいで 1 epoch が 100 日.

一方 TinyLLama 1.1B は A100 40GB x 16 で 3T tokens(3 epoch. 1 epoch およそ 1 T tokens)で 100 日です.

学習スクリプトが違うはいえ, epoch 換算(Tokens/sec 換算)だと 3090 は A100 の 20 倍くらい遅いことになります...
(理論値だと (定格 watt の) 3090 より A100 は 2~3 倍速いくらい)

batch size や, flash attention(一応有効にしてはいるが, うまく聴いているかはわからない)など, いろいろ要因があるのですかね.

学習中に見つかった不都合

vocab size 問題

vocab 数が N の倍数(8, 16 とか)でないと, 効率的に学習できないよーと warning が出ます.

64 の倍数がよいようです.

一応 model.resize_token_embedding の pad_to_multiple_of で内部的に N の倍数にできますが, weight を publish 時は元の tokenizer のサイズにしないといけません.

いろいろめんどいんで, tokenizer を作る段階で 64 の倍数にするとよいでしょう(merge のときになんか使われない token で padding するとか)

loss scale

DeepSpeed で loss scale が 400 万とかで, overflow warning が出ました.
ただ学習自体はうまくいっているので, とりあえずは無視しました😗
(loss 値でしか判断しませんでしたが, 一応 loss は下がっていたので...)

本来ですときちんと学習パラメータをレビューしたほうがいいかなと思います.

CharShu 日本語データセットの問題

日本語データセットの不都合で, 生成されるテキストがおかしくなるのが観測されました.

bunkai https://zenn.dev/syoyo/articles/99c7528863ed5f でのへんな改行がそのままになってしまったミス
- 特に " のクォートが改行入ってしまい, 生成されるテキストがおかしくなりやすい
- その他 https://github.com/lighttransport/japanese-llama-experiment/issues/1
line-wise filtering(詳細は RefinedWeb 参照)していないことによる, 「続きを読むにはクリック」みたいな web 文章の定型文が残ってしまっているミス
- https://github.com/lighttransport/japanese-llama-experiment/issues/3
中国語がまざっている日本語文章がそのまま https://github.com/lighttransport/japanese-llama-experiment/issues/5
- UTF-8 の code 範囲で判定か
あと ASCII 顔文字とか... (^o^)/*
- 入れてもいいが, 入れるならファインチューンの段階のほうがよさそう
NSFW word filtering なし
- このあたりは HojiChar で filter できるか.

あとやはり低スコアの文章は除いたほうがよさそうでした.

句点

英語と親和性がでるかと思い, 句点は「,.」に統一してみましたが, あまり効果は無いようでした.
「、。」にしたほうが, 英語とより区別しやすくなっていいかもしれません.

今回学習したモデルに対して,「、。」の文章を与えるとテキスト生成能力が落ちる感じだった(「,.」にするといい感じになった)ので, 句点はそれなりに学習に影響を及ぼすと考えられます.
(LLM 側で句点をよしなに扱ってくれるというのは少なくとも 1B サイズでは無理っぽい)

あと, Qwen などの中国語のモデルも中国語は「、。」で学習しているようなので, Qwen などの中国語も含まれているようなモデルに対して追加事前学習したい場合も考えるとやはり「、。」がいいかも.

ちなみに, 日本語は原則「、。」(横書きでも)になりました.

https://ja.wikipedia.org/wiki/句読点

公用文作成の考え方

ファインチューン用日本語データセットの問題

TinyLLama は openassistant-guanaco(oasst1 の会話ツリーを一つのやり取り(マルチターン)にしているもの)で chat 用に SFT しています.

日本語ですと oasst1-89k-ja がありますが, 実際使ってみたところ(中身見てみたところ), かなり誤訳や翻訳精度ひくくて(Google 翻訳 API)つらい...
(まあ元のデータも Now explain it to a dog みたいなコンテキストがようわからん文章が多いのであるが...)

希ガス(noble gas)が「高貴なガス」とか...

ちょくちょく issue たてたり PR してみましたが,

手動で対応していると埒が明かず, 100 年はかかりそうなので, 翻訳品質上げるために Qwen14 で翻訳バッチ回し始めました

しばらくは, ファインチューンには日本語ネイティブな llm-japanese-dataset (ややこしい名前であるが, ファインチューン用会話データセット)を使うとよいでしょうか

model.generate の tips

コンテキスト長 2K なので, huggingface transformers の generate を標準的なパラメータで動かすと, 結構繰り返しが出たりします.

カモシカ先生からのありがたいお言葉

ただ no_repeat_ngram_size 少なくするとコンテキスト長が長いと不利っぽい感じでした.

いくらか試した限りでは, llama.cpp のデフォ設定(especially, repeat_penalty 1.1)がちょうどよいような感じがありました.

さらなる高みへ...

Full-parameter continued pretraining?

学習させてからわかったことですが...

8 x A100 40GB で 20B 36 hours だとぉ...???
... Chinese LLaMa ベースで LoRA incremental pretraining より効率がいい?

LoRA とかしないで, TinyLLama のコードで普通に checkpoint から日本語データセット(tokens)で pretrain のほうがいいのかも..
(LoRA rank とか諸々のこと考えなくていいし...)

TinyLLama 3T 事前学習完了版では, フルパラ継続事前学習も試してみたいですね.

日本語データセット(コーパス)を 1T tokens 規模へ拡充

処理自体は基本スケーラブルになっています.

dedup をいくらか効率化する必要はありますが,
CommonCrawl WARC dump(生データ)や自前 crawler でデータ収集し,
filtering&dedup 済みで 1 T tokens 用意するのは比較的楽にできるでしょう.

ただ, web 系のテキストは, LLM 学習にはやっぱり量に対して品質は微妙な気がします.
最近は Phi-1.5 で, (より大きい?) LLM に学習用テキストを生成させて学習させたらいい感じになったとの結果もあるので, 100B 規模あたりの LLM でデータセットを生成というのもよさそうです.

また, Chinchilla scaling laws(1B の LLM には 20B tokens が最適)によれば, 220B(GPT-4 規模)には 4 T tokens(英語とか多言語とか) あればよいことになります.

したがって日本語 only では 1 T くらいあればよく, それ以上は(GPT-4 レヴェルであれば今のところ)不要のような気もします.

幸いにも RedPajama v2 https://www.together.ai/blog/redpajama-data-v2 で英語他言語で 30 T tokens のデータセットがありますので, これに + 日本語 1 T tokens でいい感じになるのではないでしょうか.

大規模言語モデルへの適用... 富岳で LLM したい

基本的にはデータセット(コーパス)も処理はスケーラブルなので, LiLM での知見をそのまま拡張して, トークン量増やしてネットワーク規模上げれば大規模言語モデル(LLM)も (incremental) pretrain できます.

ただいろいろパラメータ変えて pretrain の試行を増やすとか, ネットワーク規模大きくすると, 当然計算量が必要となります(とはいえ基本は O(N^2) ではなくリニア(O(N))に増えるか).

llama.cpp では MPI 対応 + フルパラ train from scratch および LoRA ファインチューンができます.
CPU クラスタで学習させるにはまたいくらか改修がいると思いますが, 富岳で llama.cpp あたりをベースとして C++ で CPU で LLM 学習させたいところですね.

LLM train はメモリ帯域ネックらしいので, ノードあたり HBM 1 TB/s な富岳は, SVE で演算最適化とかしなくとも, 少なくとも A100(40 GB 1.5 TB/s, 80 GB 1.9 TB/s) の半分くらいの性能は出るんじゃないかと思います.
(Flash Attention あたりのメモリアクセス最適化は必要であろう)

理論上は富岳フルノード(15.6 万ノード)で 4 日で Falcon 180B(500 万 GPU hours) の pretrain from scratch が可能なはず.

あとは 10B ~ 220 B 規模を富岳の物量で多量に動かして, 高品位なデータセット生成に使うとかですかね.

マルチモーダル

mini-GPT4 とか LLaVa とかで画像とか Audio とか!

TODO

ブラウザで動かす(ブラウザ側の JS や GUI を用意する)
日本語データセットを再度クリーニングし,TinyLLama 3T tokens 学習版で試す
oasst1-ja 翻訳改善版や japanese-llm-dataset で SFT する
富岳で LLM pretrain from scratch する
マルチモーダル対応
LiLM pretrain 主ふ youtuber に転生し, 「68 才主ふの至高の LiLM pretain レシピ」動画を作り, バズり, 優秀な若人さまが「68 才主ふの至高の LiLM pretain レシピ」を視聴し成長し, 人類史上最速で至高の LiLM pretrain 若人さまへと昇華なされるスキームを確立する旅に出たい