(NLP)こんにちはAIチームの戸田です今回は日本語LLMのOpenCALMの7BモデルをSFTからRLHFまで一通り学習してみたいと思います。一通り動かすことを目的としており、パラメータ調整やデータクレンジングなどのより良い学習を行うための工夫は本記事では行いません。言語モデルの一連の学習については以前、記事で取り上げさせていただきましたのでそちらをご参照いただければと思います。trlxを用いた文書生成モデルの学習1~ILQL編~trlxを用いた文書生成モデルの学習2~PPO編~以前は学習ライブラリにtrlxを使用しましたが、今回はSFT用のTrainerを提供してくれているtrlを使います。データセットJGLUEのMARC-jaを使って、ポジティブな文章を生成するように学習させたいと思います。MARC-jaは多言語のAmazonレビューコーパス(MARC)の日本語部分に基づいて作られたテキスト分類タスクのデータセットです。簡単にするため@shunk031さんがアップロードしてくれたHuggingfaceのDatasetsを使います。Notebookを使って実験してみたいと思います。必要なライブラリは以下の手順で入れることができます。git+https://github.com/huggingface/accelerate.gitまた、trlが提供するTrainerはtransformersのTrainerを継承しており、実験設定や学習ログを簡単にwandbに記録することができます。利用される方は環境変数へのwandbの設定をしてください。{WANDB_PROJECT}SFT学習につかったNotebookはこちらになりますCALM_Training_JGLUE_MARC_SFTコードは公式リポジトリのLlamaにStackExcangeを学習させる例を参考にさせていただきました。2の環境を使用することと、MARC-jaはデータ数が膨大なので学習をstep数で指定しているところがあります。Model学習につかったNotebookはこちらになりますCALM_Training_JGLUE_MARC_RMRLHFで使用する報酬モデルの学習です。従来は人手で評価されたデータセットを用意しますが、今回はMARC-jaに元々付いているラベルを利用します。ポジティブかネガティブかの分類になるので、単純な極性分類のタスクになると思います。難しい問題ではないので、メモリ節約のためcyberagent/open-calm-smallを使用しました。最終的なaccuracyは0.844になりました。LoRAのMergeSFTで学習したLoRAアダプターの重みをMergeします。してこの手順が必要かはこちらのissueでの議論が参考になるかと思われます。こちらのNotebookでMergeを行なっています。Notebookの限られたリソースで動作させるためにコードを分割していますが、メモリが潤沢にあるマシンを使用できる方はこの後のRLHFのコードの中で実行してもらっても問題ないと思います。RLHF学習につかったNotebookはこちらになりますCALM_Training_JGLUE_MARC_RLHFこちらも公式リポジトリのLlamaにStackExcangeを学習させる例を参考にさせていただきました。'test'}}のようにwandbのAPIで指定する際の要素をdictで渡してあげる必要がありました。今回の記事では精度改善のための取り組みは基本的に扱いませんが、一点だけ知見共有としてREWARD_BASELINEパラメータを紹介します。3)で強化学習を少し触った際に、報酬を大きく設定しすぎるとlossが大きく揺れて下がらず、学習がうまくいかなかった経験があったので、この辺りが関わってきているのではないかな、と考えています。(詳しい方がいましたら教えていただきたいです)結果LLMの生成結果の良し悪しを定量的に判断するのは困難ですが、簡単な人手評価として、MARC-jaのvalidationデータから5件テキストをサンプリングして、その最初の五文字に続く文章をSFTのモデルとRLHFのモデルに生成させました。Shift内の5人の多数決でどちらの生成がポジティブなレビュー文章として優れているか評価してみました。結果を以下に添付します。SFT対RLHFが2対3でRLHFの方が若干優っているようにも見えますが、サンプル数が少ないので誤差のようにも思えます。Modelを改善した方がいいような気がします。繰り返しになりますが、本記事では一通り動かすことを目的としており、良い学習を行うための試行錯誤については扱わないので、結果の分析についてはまたの機会に行いたいと思いますおわりに本記事では日本語LLMのOpenCALMの7BモデルをSFTからRLHFまで一通り学習してみました。以前RLHFを試した際にできなかった、SFTや1B以上のモデルでの学習、人手評価などを今回は手を動かして確認できました。Modelの挙動が怪しいので、今後はそこから取り組んでいこうと思います。