0.1-1B程度の日本語・大規模言語モデルを作る練習日記

Kan Hatakeyama

2024年3月11日 11:52

はじめに

最近は大規模言語モデルを作っています。

最終的には１０B級のモデルを作るのですが、その前に、0.1-1b程度で練習をします。
本記事は、その取り組みのメモ・日記です。（逐次更新します）

注意点
・日記なので、内容はかなり雑多です。
・練習を兼ねたcode実装なので、先端知見が含まれているとも限りません。
・どんなことをやっているか、時系列を追って理解する参考にはなるかもしれません。

3/6頃コード確認

・「標準コード」を支給頂いたので、それをトレースした。
・諸々の処理（pipe）が繋がっていないので、それをつなげる必要有り

3/7 パイプライン化

・初心者向けに、shコマンドを一発叩くだけで、諸々の処理が終わるパイプラインの構築が必要である。
・事前学習程度までのパイプは繋がった。
https://github.com/hatakeyama-llm-team/llm/tree/dev_hatakeyama/codes

・練習でWikipediaの日英記事、それぞれ１００万件を0.1Bモデルに事前学習させた。
・わかってはいたが、0.１ｂモデルは流石のクオリティの低さであった。

つまづき（？）点
学習の際、train_samplesと,train_tokensを設定する項目があった。
train_tokensを、実際のdataset sizeにしておけば、そのタイミングで学習は終わる。
train_samplesは適当に大きなサイズにしておけば良いのだが、これをあまりに大きくしすぎると、memory errorが発生した。
（以下は、30000B token相当に設定した例）
train_samples=$(( 30000 * 1000 * 1000 * 1000 * 2 / ${seq_len} ))

小さすぎると、train_samples側の基準でtrainが終わってしまうので、1000Bくらいにしておくのが良いか。

3/8 0.3bの学習

0.1bがアホすぎたので､0.3bで回してみた
- A100 (80GB) x2で計算｡
- バッチサイズは､deepspeed 0､global batch size 64で､ちょいど良い感じであった

wikipedia(200件､日英)は3.6b token｡
- 学習速度から推定すると､37000 sec (10時間程度)で1epoch
学習後は、0.1bよりも知性を感じる結果となった｡
- (1 epoch = 27466 step後)

つまづき（？）点
・本teamでは、特定の順番でデータを学習させることを計画している。
・まずはデフォルトで、データがどの順番に学習されているかどうかの確認が必要であった。
・Megatraon-DeepSpeed内のpretrain_gpt.pyを一時的にhackして、t.txtに実際に学習させる文章を出力する処理を暫定で追加した。



def get_batch_pipe(data):
	"""Modification of `get_batch` to work on `next(data_iterator)` instead of `data_iterator`"""
	args = get_args()
	tokenizer = get_tokenizer()


	#実際の学習データを確認するscript
	import sentencepiece as spm
	model_path = "../../data/tokenizer/tokenizer.model"
	sp = spm.SentencePieceProcessor(model_file=model_path)
	with open("t.txt", "a", encoding='utf-8') as f:
	    token = (data['text'])
	    shape = token.shape
	    f.write(str(shape) + "\\n")
	    txt = sp.decode(token[0].tolist())
	    f.write(txt)

・その結果、学習直前にdataをindexするタイミングで、データのshuffle処理が自動実行されていることがわかった。
・https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/data/gpt_dataset.pyの中で、np_rng.shuffleという処理が３回行われるので、そこをコメントアウトすることで、shuffleを防ぐことができた。

3/8 deep speedのzero stage チェック

deepspeedのzeroを増やすとどういうメリットがあるかを調べることにした
ベース: 0.3bモデル
stage 0
- global batch: 64
- (2gbほど､別プロセスで使ってたので注意)

stage 1
- ちょっとvramの使用量が減った?
- 速度はほぼ変化なし

・Stage 2以降は、エラーが出ました。追加のconfig設定が必要なようです。

3/9 データの学習順序がモデルに与える影響の確認

・0.3bモデル
・Wikipedia 英語１００万件、日本語１００万件の順序で学習
・1 epochで24k step

・英語のみを学習させた場合（@ 10k step）

・英語→日本語の順に学習させた場合（@ 24k step）

・日英がshuffleされた場合　(以前の検討モデル。penaltyを一部変更)

つまづき（？）点
・想定とは異なり、完全に英→日の順番で学習させてしまうと、LLMが適切に日本語能力を獲得できないことが判明
・初期の学習時、英語用にモデル重みが固定してしまうことが原因？
・学習順序に、グラデーションをかけていくことが重要かもしれません。

・Perplexityも計算してみました。

def perplexity(model, tokenizer, text) -> torch.Tensor:
    tokenized_input = tokenizer.encode(
        text, add_special_tokens=False, return_tensors="pt"
    ).to(model.device)
    with torch.inference_mode():
        output = model(tokenized_input, labels=tokenized_input)
    ppl = torch.exp(output.loss)
    return ppl.item()

英→日の場合、、

今日は: 354
富士山は：　422
I have a : 357
Functional polymers: 11

でした。やはり、日本語に対する困惑度が高いようです。
学習初期はlearning rateが高いこともあり（？）、英語だけを学習させると、日本語の能力が身につきにくいようです。

3/11 日英Wikipediaの学習順序を最適化する

0.3bモデル
Wikipedia 英語１００万件、日本語１００万件を学習
- 1 epochで24k step
3/9の検討で、完全に英語→日本語の順番では良くないことが判明
今回は、3 Stageに分けてデータセットを配備
- １：　日本語９万件、英語９万件
- ２：　日本語９万件、英語８１万件
- ３：　日本語８１万件、英語９万件
一晩待って､計算結果が出ました｡
ちゃんと日本語を喋れるようになりました!
- perplexityも､どちらかというと日本語が下がり､英語で上がる傾向があるような気がしました｡

個人的な手応え: 学習初期に､様々なトークンに触れさせておかないと､局所解に陥るかも(?)

3/12 tokenizerをデカくしてみる

背景

日本語を喋るLLMのtokenizerの語彙サイズは､60kくらいが一つの目安になっています｡

一方､2024年に登場したGemmaの語彙は脅威の256kでした｡
- 中国のQwen系も150k
語彙が多いほど､学習token数を減らせる

一方､下手に増やせば良いというものでもないようです
- BLOOMは250k
  - embedding層のサイズが大きすぎて大変だった模様
  - あんまり評判は高くない印象
- どこのサイトか忘れましたが､Gemmaも､日本語の圧縮率は一般的なモデルの1.5倍程度(?)と､そこまで高くない

今回のタスク

これまではvocab sizeを5kでやってたので､ものは試しに､65k, 250kにしてみます｡
学習データは3/11と同じです｡
語彙の比較
- 流石に語彙を増やしすぎた気がします｡少なくとも､もっと多くの文章でtokenizerを学習させた方が良さそうです｡
- sentencepiece (unigram)を使用しました｡

学習結果
- いまいちでした｡ (計算は途中で止めました)
- 65k付近に最適解がありそうです｡

ポイント
- vocabを増やすと､データセットのtoken数を減らせる
  - しかし､減る速度はどんどん落ちていく
  - 少なくとも今回のdatasetでは､65k程度のvocabで十分
- vocabを増やすと､必要なVRAM量が上昇し､batch sizeを小さくする必要がある
  - そのせいか､250kではlearning timeが5k条件の二倍程度に｡
- 65kでもglobal batchは小さくなったが､学習tokenが減ったこともあり､1 epochあたりのlearning timeを削減できた

3/13 tokenizerの語彙数で出力を比較する

語彙数が65kのモデル学習(0.3b)が終わったので､出力を語彙数5k, 65kで比較してみます｡
日英Wikipediaを学習: 3 Stageに分けてデータセットを配置
- １：　日本語９万件、英語９万件
- ２：　日本語９万件、英語８１万件
- ３：　日本語８１万件、英語９万件
結果: perplexityは増えるが､体感の精度はあまり変わらない

語彙数が多い ≒ 予測すべきtokenの選択肢が増えるので､perplexityが増大するのは､想定内の結果かもしれません｡
- このあたりは結局､学習トークン数を稼ぐしかなさそうです
1文字あたりの出力時間も､ 0.005 sec/文字 → 0.004 sec/文字程度まで削減できました｡
今回はwikipediaのようなきれいな文章のみでtokenizerを学習しましたが､mc4など､種々の文章を混ぜたときの語彙の最適値は微妙に変わるかもしれません｡ (今後の検討課題)

3/13 雑多なweb文章を学習させる

wikipediaではなく､雑多なweb文章を学習させてみます｡
- mc4-jaを100万件､wikipedia-enを100万件
  - 全くクリーニングしていません
- 学習順序やモデルサイズは､3/11と同じです｡
tokenizerは今回のデータセットで再学習しました｡

mc4-jaは､そのままでは目も当てられないほど､汚いテキストが多量に含まれています(こちら)｡
それをそのまま学習させるという作業をしてみました｡

結果: わりと日本語を喋れました｡
- ただし､perplexityは高めの印象です｡
  - ゴミ情報が多いからなのか､それとも､mc4の情報が多彩だからなのかは､なんともわかりません
- 広告っぽい単語の羅列をいれると､広告サイト風のテキストが出てきます｡

感想: 大規模言語モデルは賢いと思いました｡

3/18 ジャンルごとに継続学習してMoEとして推論する

以下のステージにわけてモデルを学習させました(件数は概算, 約6b tokens)｡

Stage 0: ランダム
- 日本語Wiki 3万件､mc4-ja 3万件､英語Wiki 3万件
Stage 1: 英語Wiki
- 日本語Wiki 3万件､mc4-ja 3万件､英語Wiki 80万件
Stage 2: 日本語雑多(mc4)
- 日本語Wiki 3万件､mc4-ja 80万件､英語Wiki 3万件
Stage 3: 日本語Wiki
- 日本語Wiki 80万件､mc4-ja 3万件､英語Wiki 3万件

モデル
- 1 bのモデルにしました
  - 某先生のサーバーをお借りしました｡ありがとうございます!
  - A100 x8 で1日弱でおわりました｡
- 少しズレはありますが､stage 2,3,4に対応するcheckpointを保存して､独立したモデルとして使います
推論
- Stage 2,3,4のモデルの中から､最も予測性能が高いものを選んで回答を生成させます
- 入力文章に対するperplexityを計算し､一番小さなモデルを使うアルゴリズムを適用します
  - perplexityは､入力文章そのものをモデルに予測させた際の「予測誤差」です｡
  - perplexityが小さい ≒ その文章と親和性が高い ≒ 回答精度も期待できるというロジックです
- 実装メモはこちら

結果

いくつかの入力に対して､推論させました｡
- model id
  - 0: 英語特化(stage 1)
  - 1: 雑多な日本語特化(stage 2)
  - 2: きれいな日本語特化(stage 3)
- PPL List
  - 各モデルのinputに対するppl

ポイント

「今日は晴れてるから」､「ブログを書きました」のような､wikipediaへの記載がなさそうな文章については､Stage 2のモデル(雑多な日本語特化)が選ばれました｡
「地球温暖化」のような､お硬いトピックについては､Wikipediaを読んだ直後のStage3のモデルが選ばれました｡
予想外だったのは､「global warming」のような英語の入力に対しては､英語を学んだ直後(stage 1)ではなく､日本語を重点的に学んだstage 3以降のモデルが選ばれたことです
- stage 2,3でも英語は少しだけ学ばせてますが､perplexityがStage 1,2,3の順に顕著に下がっていることを感が見えると､
- 日本語を学ぶことで英語の能力も一応は向上しているのかもしれない?と思いました｡
  - ただ､この規模のモデルで､知識転移が起きる印象はあまりないです
  - 1 bモデルに対して､6b tokenしか学習させていない点にも注意が必要そうです(Chinchilla 論文以降は､少なくとも20倍は学ばせるのが､定石です)｡
  - 文法の共通規則のようなものを学んでいる可能性はあるかもしれません｡

3/19 カリキュラム学習&ジャンルごとのモデル切り替え練習

・1bモデルを以下のステージに分けて学習させました｡
・前回に比べ､英語や日本語の件数を数倍に増やしています(全部で10b)｡

・Stage 2,3,4に時点のモデルをそれぞれid 0,1,2として､入力に対するperplexityを計算しながら､出力を見ました｡
・各入力に対する､出力は以下の通り｡(boldが､perplexityが最も低い､ベストと判断されたモデルの出力です)

------
input: hello, I'm John
model id 0: hello, I'm John. 2016年3月
model id 1: hello, I'm John. I am not sure if you have any
model id 2: hello, I'm John." "I Love You Like That (
------
input: global warming
model id 0: global warming 2016年3月【お知らせ
model id 1: global warming 2018-5 Posted by
model id 2: global warming The climate of the United States is characterized by
------
input: 英語: He is a good man. 日本語:
model id 0: 英語: He is a good man. 日本語: 「彼は私のことを愛している」 20
model id 1: 英語: He is a good man. 日本語: Thank you for your kindness and help in finding
model id 2: 英語: He is a good man. 日本語: 私は、私の夫は私を愛している。
------
input: 今日は晴れてるから
model id 0: 今日は晴れてるから、ちょっとだけ雨降ってた。でも
model id 1: 今日は晴れてるから、ちょっとだけ散歩してきます。 20
model id 2: 今日は晴れてるから、また行こうかな。 2017
------
input: ブログを書きました
model id 0: ブログを書きました。 2017年3月9
model id 1: ブログを書きました。 2014年5月3
model id 2: ブログを書きました。 2017年6月3
------
input: 地球温暖化を防ぐに
model id 0: 地球温暖化を防ぐにはどうしたらいいか? 2017
model id 1: 地球温暖化を防ぐにはどうすればいいか? 2018
model id 2: 地球温暖化を防ぐにはどうしたらよいか。 2017
------
input: ガリレオによるピサの斜塔実験とは
model id 0: ガリレオによるピサの斜塔実験とは? 2017年3月6
model id 1: ガリレオによるピサの斜塔実験とは? 2018年5月3
model id 2: ガリレオによるピサの斜塔実験とは、1960年代後半にイタリアで
------
input: スーパーコンピューターの富岳は
model id 0: スーパーコンピューターの富岳は、1980年に噴火した火山灰
model id 1: スーパーコンピューターの富岳は、2016年3月5日
model id 2: スーパーコンピューターの富岳は、2017年3月で稼働

感想:

わりといい感じにモデルを切り替えられました｡
基本的には､stageが上がるほど､(当然ながら?)モデルは賢くなるようです｡
- 英語を中心的に学習させた､最初のモデルは､ほぼ使われませんでした｡
- ただし､mc4にいかにも入っていそうなblog系の内容などでは､model id =1 が選ばれる結果となり､ドメイン特化の有効性を改めて確認できました｡

3/24 DeepSpeedのZero Stage 0,1の速度比較

共通条件

パラメータ
- model_size: 1.3
- num_layers: 24
- num_attn_heads: 16
- hidden_size: 2048
GPU
- A100(80gb) x8

結果

Stage 0: global batchは88が限界
Stage 1: global batchは112が限界
どちらも約1.7 samples/secで、速度にほぼ違いなし。

感想

Stage 2,3では、よりVRAM使用量を削減できるとの噂です。
が、速度面では、 0>1>2>3>3+offload との噂なので、十分にGPUを確保できる場合は、敢えて使う必要はない感じでしょうか、、。

4/2 2.7bの事前学習が終了

概ね､英語wikipedia(360万件)→ CommonCrawl系日本語(1000万件)→日本語wikipedia(120万件)がメインになるような順番でカリキュラム学習させました｡
合計で30 btoken程度です｡

0.1bに比べて､「雑多な日本語ゾーン」でのlossの下がりが改善されました｡もっとモデルサイズを大きくすれば､雑多な日本語を喋れるようになるのでしょうか｡

評価結果

以下のcheckpointで保存された3つのモデルの出力を見ることにしました｡

model id 0: 24000/199305 step → 英語メイン
model id 1: 180000/199305 → 雑多な日本語メイン
model id 2: 199305/199305 → きれいな日本語メイン

------
input: hello, I'm John
model id 0: hello, I'm John. 1970s: The Last of the Seven Stars (The Four Star) - A short film about a group called "the
model id 1: hello, I'm John and then a little bit of my life." 2018: "They are not really in love with me but have
model id 2: hello, I'm John Doe." 1973: "They are the ones who have been in charge of this country for so long and

------
input: global warming
model id 0: global warming 10–25 °C, and then to a temperature of about-38.4° C in August; it was reported that
model id 1: global warming 2018年7月3日【動画】中国で「日本製品不買運動」が活発化、日本の食品メーカー
model id 2: global warming climatic conditions. The climate change scenarios are based on the assumption that a large part of Earth'need is being lost due to human

------
input: 英語: He is a good man. 日本語:
model id 0: 英語: He is a good man. 日本語: 'I am not the best of my friends, but I have always been in love with you forever and it will be better than that which would
model id 1: 英語: He is a good man. 日本語: 私は良い人です。【楽天市場】購入者さんの【送料無料】【2ケースセット】(北海道・沖縄は別途送料必要)ポッカサッポロ生ビール
model id 2: 英語: He is a good man. 日本語: 彼は素晴らしい人です。「He's the one who made me feel like I was in love with him/But he doesn’t

------
input: 今日は晴れてるから
model id 0: 今日は晴れてるから、雨が降らないと困る。でもね～・・・この前は雨が降ってなかったのにな～。【楽天市場】10%OFF
model id 1: 今日は晴れてるから、ちょっと散歩にでも行こうかな。【楽天市場】購入者さんの【送料無料】【2ケースセット】(北海道・沖縄は別途送料必要)サントリー天然
model id 2: 今日は晴れてるから、雨の心配はなさそう。 The Battle of the River Plate was a naval battle fought in April and May, during World War I. It

------
input: ブログを書きました
model id 0: ブログを書きました。 2017/3月5日(水)の朝、東京・新宿で「第6回『日本と世界の平和
model id 1: ブログを書きました。 2018/9月3日(土)に、東京・渋谷のNHKホールで「第九」演奏会が開催されました
model id 2: ブログを書きました。 The following is a list of the Los Premios MTV Latinoamérica winners and nominees for Best New Artist. The awards were established
------
input: 地球温暖化を防ぐに
model id 0: 地球温暖化を防ぐにはどうしたらいいのでしょうか。「宇宙は、人類が作り出したものだ」という考えを広める活動を行っています。「2017年4
model id 1: 地球温暖化を防ぐにはどうしたらいいか、という問題は人類が解決しなければならない大きな課題です。【楽天市場】購入者さんの【送料無料】【2ケースセット】(北海道
model id 2: 地球温暖化を防ぐにはどうすればいいか、という議論が展開された。 2017年3月8日～9日に、第4回「気候
------
input: ガリレオによるピサの斜塔実験とは
model id 0: ガリレオによるピサの斜塔実験とは、1960年代後半にアメリカで開発された世界初の超高速ロケットである。 2.5トン級(約3,784kg
model id 1: ガリレオによるピサの斜塔実験とは、1960年8月24日にイタリア・トリノで行われた「第3回国際科学技術博覧会」でのこと。【楽天市場】
model id 2: ガリレオによるピサの斜塔実験とは、1907年6月25日にイタリアのローマで実施された。背景と目的は、ピサーノ・デッレが
------
input: スーパーコンピューターの富岳は
model id 0: スーパーコンピューターの富岳は、2017年4月3日(土)から5月上旬にかけて雨天が予想されます。 ■ 「第6回東京
model id 1: スーパーコンピューターの富岳は、パソコンやスマートフォンなどの電子機器をインターネットに接続するための無線LAN(Wi-Fi)機能を搭載しています。「このたびの大雪により
model id 2: スーパーコンピューターの富岳は、1980年代後半から23世紀の初めにヨーロッパで流行したコンピュータ・ミュージック(電子音楽)を発展させた。背景