チームでのLLM開発の反省点と解決策の検討

LLMの開発におけるチーム開発で実際に直面した課題をまとめていきます。私は普段はエンジニアとして働いており、エンジニア的な側面からみたLLM開発について、課題とその解決策を検討し、これからLLM開発をチームで行うときの参考となればと思います。

開発期間は約3ヶ月間で、チームの編成としては、学生や社会人で、バックグラウンドも様々なメンバーが集まりました。ほとんどのメンバーはLLMの開発に携わったことはない状態です。私はサブチームのリーダとしてサブチームを管理する役割を担っていました。開発環境は3ノード、24GPUで、10B程度のモデルを学習することが目的です。

LLM開発を終えての課題

LLM開発についての細かい課題はいくつかありましたが、特にチーム開発という点で、以下の課題があったように感じます。

• 環境構築
• 全体像と細かいタスクへの落とし込み
• コード管理の複雑さ

環境構築

CUDA周りのpytorchやdeepspeedなどの複数のライブラリのバージョンを合わせるのが大変でした。 (今回は、さらにPytorchでの混合精度と分散トレーニングを効率化を行うユーティリティであるapexや、flash-attention、TransformerEngineを追加でインストールしました。)

SlurmというJobスケジューラーとanacondaによる環境構築を行いました。これらの環境において、CUDAの環境変数がうまく読み込まれない問題がありました。これにより、異なるバージョンのCUDAを参照したり、インストールされるライブラリのバージョンがわずかに異なるなどの原因により、同じコードでも環境によっては動作しないということが起こりました。

Slurmとは、Jobスケジューラーで、とあるマシンにログインし、サーバー上でJobを実行することで、マシン上でCPUやGPUが扱えるような仕組みです。

全体像と細かいタスクへの落とし込み

チームでのLLM開発の知見があまりない状態、またはチームのメンバーの技術力やバックグラウンドが様々な状況で、実際に手を動かすレベルの作業タスクに落とし込むことが難しかったです。

LLM開発のおおまかな流れとして、データセットの収集し、tokenizerを作成、事前学習、事後学習というような、全体の流れとしては、なんとなくわかるものの、具体的に各フェーズで何をするのか、何を調べるのか、どのくらい時間がかかるのか、がわかりづらい状態でした。チームとして機能するには、各フェーズで行うことを作業タスクに落とし込み、各人が動ける状態にすべきです。

コード管理の複雑さ

Gitでコードを管理していたものの、様々なコードの修正が混ざり合ったことと更新の頻度の多さにより、コード管理と共有に苦労しました。

例えば、学習のパラメタを決定するために何度もコードを書き換え実行するような試行錯誤を行うのですが、この修正と環境依存の修正を分離できず、別の環境でコード動かす際はさらにこの上から環境依存の修正を行うような、その場しのぎ的な修正を行いました。

また、環境構築が行えた人が少なかったことにも起因しますが、コードの共有が難しいことで学習自体を実行できる人数が制限されてしまいました。

環境構築についての改善策

学習コードにエラーがあった際のデバッグや動作検証のため、学習環境と同様の環境を持つメンバーを増やすのは重要です。今回の開発では、簡単な手順書を参考に環境の作成を行いましたが、人によってライブラリのバージョンが異なったり、学習コードがうまく動作しなかったりというような状況となりました。環境構築を共有できる手段の検討は重要であると考えます。

手順のドキュメント化

最もシンプルな手段です。行った手順をまとめるためだけですが、手順のミスなどのヒューマンエラーにより環境がうまく作成できないことがあります。また、わかりやすいドキュメントを作成するには時間がかかるため、後回しになりがちになります。

構成管理ツールの利用

AnsibleやChefなどの構成管理ツールを利用し、インストール手順を記述しておくのも良いかと思います。yamlなどの設定ファイルが手順書として残るので、手順のドキュメント化は不要です。一方で、ツールのインストールやツールを使うための知識などが別途必要となります。

dockerコンテナを使う

nvidiaが提供するdocker imageを使うのが楽かと思います。素早く環境の作成が行え、ポータビリティにも優れています。Dockerが使えない環境であったり、マルチノードで分散学習を行う場合には、コンテナのportやipなど設定が必要となります。

slurm環境であれば、pyxisというslurmのプラグインと、enrootというdockerのランタイムを使うことで、　docker containerの起動が行えます。

検討

起動の速さや起動手順の簡単さからnvidiaのdocker imageを使うのが一番楽かと思います。 ただし、学習を行うホスト上で、sudo権限があるか、dockerが使えるか、新たなライブラリのインストールは可能かなど環境に応じて選択も発生します。いずれの選択肢にしろ、できるだけインフラのコード化を行える手段が良いかと思います。

LLM開発における開発手法

コード管理の複雑さについては、適切な粒度とタイミングでタスクが作成され実施されることでコード修正のまとまりがうまれ、コード管理も行いやすくなるものと思います。

細かいタスクへの落とし込みを行うためにも、特にチームとしてのLLM開発の知見が乏しい状況では、

• データの収集
• tokenizerの作成
• 事前学習

を小さい規模で早い段階でとにかく一度行ってみるというのが1つの解決策のように思えます。

これにより、各フェーズで具体的に何を必要があるか、調べる必要のあること、実装する必要のあることなど、どの程度実施に時間がかかりそうか、などが明らかになります。作成したモデルに一度、言語を生成させてみるというのも、確認作業としてはかなり重要なポイントかと思います。ここで得られた知見により、作業をタスクとして落とし込み、チーム全体として人的リソースがうまく扱えるようになるのではないかと考えます。その後は、モデルサイズ少しづつを大きくしたり、学習手法の検討を行ったりとフィードバックを元にサイクルを繰り返し行うことが重要である思います。アジャイル開発による開発に近いイメージです。

モデルやデータセットを大きくしすぎると1サイクルの実行は遅くなるので、小さいサイズのまま何度かサイクルを回し、学習コードの検証や動作確認を行うのが良いでしょう。このときできるだけ多くの人がサイクルを実行できると理想的です。

振り返ってみると、我々のチームでは、事前学習の検証は行なっていたものの、tokenizerの作成や事前学習、事後学習などは、分離したフェーズと捉え、各フェーズを順番に実行していくようなウォーターフォール開発のようなアプローチとなっていました。経験の少ないチームではとにかく一度試してみるのがよかったという反省があるように思います。

人的リソースについては、各メンバーの得意分野をまとめておくことも重要だと感じいました。我々のチームではこの辺りはもう少し細かく行うべきだったなと思います。チームのマネジメントやコードの修正/管理であったり、論文のキャッチアップ、ドキュメントをまとめる作業など、明確に分担できていない部分がありました。タスクの作成とリソースのアサインはどちらかの解像度が低くてもうまく進まないと感じました。

まとめ

チームでのLLM開発における課題やそれに対する解決策をまとめました。

様々なバックグラウンドを持つ人が集まりLLMを開発するという特殊な状況ではありましたが、このようなLLM開発の知見があまりないような状態での開発は、開発の仕方も試行錯誤して進むような状態でした。そこで、アジャイル開発で見られるような開発サイクルをLLM開発にも取り込むことを検討しました。開発手法については、今後さらに行われていくであろうLLM開発のなかで、より良いものが提案されていくものと思います。

今回のLLMの開発では、期限内に成果物を作成するという開発的な側面と、新規性・性能向上を目指す研究としての側面のバランスをとることの難しさを感じました。

この成果は、NEDO（国立研究開発法人新エネルギー・産業技術総合開発機構）の助成事業「ポスト５Ｇ情報通信システム基盤強化研究開発事業」（JPNP20017）の結果得られたものです。