Stockmark Tech Blog

自然言語処理テクノロジーで社会を進化させる ストックマークのテックブログです。

Stockmark-100b: ビジネスドメインに対応した1000億パラメータ規模の大規模言語モデルの事前学習

ストックマークは最近、ビジネスドメインに対応した1000億パラメータ規模の大規模言語モデルであるStockmark-100bの開発を行い、事前学習モデルと指示学習モデルをオープンソース(MITライセンス)として公開しました。この記事では事前学習における弊社の取り組みを紹介させていただきます。

プレスリリース: stockmark.co.jp

事前学習モデル: huggingface.co

指示学習モデル: huggingface.co

背景

2024年2月に国立研究開発法人新エネルギー・産業技術総合開発機構NEDO)が実施する「ポスト5G情報通信システム基盤強化研究開発事業/ポスト5G情報通信システムの開発」に採択されました。これは、経済産業省の国内の生成AIの開発力強化を目的としたGENIACプロジェクトと連携して行われており、国内事業者に対して生成AIの開発に必要な計算資源の確保と利用料補助を行うものです。採択された当社の開発テーマは「厳密さが要求されるビジネス用途におけるハルシネーションを大幅抑止した基盤モデル 」で、ビジネスでも使える正確性の高いモデルの開発を目指しています。

事前学習

今回のプロジェクトでは、1000億パラメータのLLMの事前学習を行いました。モデルの規模は、日本語を主な対象とした公開モデルとしては、これまでにない大きさになっています。ここでは、ビジネスドメインに対応することを目指して、一般によく使われるWikipediaやCommon Crawlといったデータセットだけではなく、当社が独自に収集している特許やビジネスドメインのウェブデータも用いて学習を行いました。最終的に日本語・英語合わせて9100億トークンのデータを学習しました。

データ

事前学習には日本語と英語のデータセットを用いました。日本語のデータセットのサイズは3050億トークンで、構成は以下のようになっています。ビジネスドメインに対応するために当社が独自に収集しているWebページや特許のデータを用いていることが、特徴になっており、これによりビジネスドメインの知識の獲得が期待されます。

Webデータに関しては、ルールベースのクリーニングや、行単位の重複排除や、MinHashを用いたドキュメント単位の重複排除などにより前処理を行いました。

英語はRedPajamaのデータセット(Booksのカテゴリーを除く)からサンプルしました。

モデル

モデルのアーキテクチャーとしては1000億パラメータのLlama-2モデルを用いて、ゼロから事前学習を行いました。モデルの詳細は以下のようになっており、これらのハイパーパラメータはLlama-2の70Bモデルのものを参考にしながら決めました。

ライブラリ

分散学習のライブラリとしてはMegatron-LMを用いました。Megatron-LMは大規模言語モデルの分散学習を効率的に行うためのスタンダードなライブラリで、Llama-2のモデルもサポートされています。

Megatron-LMのTransformeの実装には、古くから使われているレガシーな実装と、最近の世代のGPUでTransformerの演算の効率性を高めるライブラリであるTransformerEngineを用いた実装の二種類があり、今回は後者の実装を採用しました。今回の事前学習の環境での速度比較を行ったところ、後者の実装では10%ほどのスループットの向上が見られました。

Megatron-LMで学習をした後に、Hugging Faceのtransformersのライブラリで使えるようにするためには、モデルのチェックポイント(モデルのWeightを含むファイル)を変換する必要があります。ただ、開発を進めた時点では、TransformerEngineを用いた実装に対しては、利用可能な変換スクリプトが(探した範囲では)見つからなかったため、これに関しては、自分達で作りました。

環境

事前学習はGoogle Cloud(以下GCP)で行いました。GPUインスタンスとしてはH100が8枚搭載されているA3インスタンスを用いました。

今回の事前学習ではモデルサイズ、データサイズともにとても大きいので、複数のGPUノード(インスタンス)を用いた分散学習を行います。複数のノードを用いることにより、一つのGPUに乗らない大きなモデルを複数のノードに分割して処理できるとともに、並列でデータを処理することによりスループットが向上します。

大まかには、以下のようなシステム構成を用いました。基本的にはコードやデータは全て共有のNFSに保存しておき、それぞれのGPUノードは、そこからデータやコードを読み取り、ログやチェックポイントもNFSに保存します。

大規模言語モデルの学習では、より多くのデータを学習した方がモデルの性能が高くなります。そのため処理の効率を上げて、できるだけ多くのデータを学習できるようにすることが重要です。

分散学習では、ノード間でデータをやり取りしながら処理を進めます。多数のインスタンスを用いて学習を行う時には、ノード間の通信の量も増えるため、これが処理のボトルネックとなることがあります。今回は、ノード間の通信をより高速にするために、GCPが提供しているGPUDirect-TCPXという機能を用いました。これはGPUメモリにあるデータをCPUとシステムメモリを経由せずに、直接ネットワークインターフェイスに送信する技術で、私たちの環境では、これにより処理効率が約1.5倍向上しました。

また、学習の途中で定期的にモデルのチェックポイント(モデルの重みやオプティマイザの状態)を定期的に保存しています。これはプロセスが何らかのエラーで止まってしまった場合に、途中から学習を再実行できるようにするためです。今回のような1000億パラメータのモデルとなると、チェックポイントも1.3TBとなり、保存にそれなりの時間がかかってしまいます。チェックポイントの保存をおこなっている間は、学習は止まってしまうので、できるだけ短い時間でチェックポイントの保存を行う必要があります。今回はNFSとしてFilestoreを用いましたが、Filestoreはストレージの容量に応じてディスクスループットが大きくなるので、80TBという大きな容量を設定しました。これにより約8-10分程度でチェックポイントの保存が行えるようになりました。

学習

学習自体はロスのスパイクなどは起こらずに安定して進みました。最終的に日本語・英語合わせて9100億トークンを学習しました。その一方で以下のような課題や苦労などがありました。

今回は日本語と英語を混合したデータでLLMの学習を行なっています。前半では英語3、日本語1の割合で約6000億トークンの学習を行いましたが、学習の進みが想定よりも遅かったです。具体的にはバリデーションロスの値を他のモデルなどと比較してわかりました。そこで、後半では英語1、日本語9の割合で残りの3100億トークンの学習を行いました。日本語のバリデーションロスの値は、これにより、前半よりも早いペースで減少するようになりました。最終的に想定するパフォーマンスを得るために英語と日本語データセットをどのような割合で混合し学習を行うかや、ハイパーパラメータをどのように設定するかなどで課題が残りました。

また、事前学習の最中にプロセスが途中で止まってしまうことがそれなりの頻度で起こりました。事前学習の前半ではエラーが起こる頻度は少なかったのですが、最後の一週間では1日に1から2回ほど止まってしまいヒヤヒヤしました。これらは、コードのバグというよりはシステム起因ではないかと考えています。そのため、途中でプロセスが止まってしまっても、自動でプロセスの停止を検知して、学習を再実行するような仕組みを作ることが安定的な運用には重要でした。またGCPでは不定期にGPUインスタンスのメンテナンスが行われ、メンテナンス時にはインスタンスが停止させられてしまいます。幸いメンテナンスは1時間前から通知が行われるので、通知があった時には自動でチェックポイントの保存を行うような仕組みも導入しました。また、一時的にインスタンス不良により一部のインスタンスが使えないというような状況もあり、このような状況になった時に使用可能な一部のインスタンスのみを使って学習を継続するといったことも想定しました。このように、特に大規模な分散学習では運用面の苦労が多くありました。

最後に

今回は事前学習での取り組みについて簡単に紹介させていただきました。モデルの評価などについては、別の記事で報告させていただければと思います。

謝辞

この成果は、NEDO(国立研究開発法人新エネルギー・産業技術総合開発機構)の助成事業「ポスト5G情報通信システム基盤強化研究開発事業」(JPNP20017)の結果得られたものです。また、当社は産業技術総合研究所とLLMに関する共同研究をおこなっており、今回のプロジェクトの遂行にあたっても助言をいただきました。

(作成)近江崇宏