AIStationで LLM 強化学習はどこまで可能か：verl を用いた 3B/7B モデルの検証

概要

• Fixstars AIStation（以降 AIStation）を利用し、verl による強化学習を試行しました。
• 検証の結果、特定の条件下において、 Qwen2.5-7B-Instruct は OOM（Out of Memory）により学習不可能でしたが、より小さいモデルである Qwen2.5-3B-Instruct は学習可能であることが判明しました。
• このことから、条件に依存するものの、小規模なパラメータを持つモデルであれば、AIStation でも verl による強化学習が実施可能であることが示されました。

はじめに

こんにちは、ソリューション第四事業部ディレクターの森高です。

フィックスターズは AIStation という製品を提供しています。AIStation は NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition（以降 Max-Q）を 2 枚搭載しています。AIStation については [1] を、Max-Q の詳細を知るには過去の投稿 [2] を参照してください。

これまでに、当社は AIStation 上で BLAS ベンチマークの試行 [3] や、LoRA による学習 [4] を実施しました。本稿では、verl [5] という OSS の強化学習フレームワークを使って AIStation 上で強化学習を試します。

本稿では AIStation で verl を動かしたとき、どれくらいの性能が出るのか試した結果を示します。まずは verl で強化学習できるのかを確かめ、そのあと verl のパラメータを振って動作する限界を探します。

verl について

verl は ByteDance Seed チームが開発を始め、現在は verl コミュニティがメンテナンスをしている LLM 向け強化学習のフレームワークで、[6] のオープンソース版実装です。verl は、RL アルゴリズムの変更や既存の LLM との統合も簡単であるなど柔軟性を備え、さらに SoTA クラスのスループットを発揮するフレームワークであると言われています。今回はこの verl を利用した強化学習を AIStation で試してみました。

実験内容概要

verl を使って GSM8K [7] というデータセットの学習 2 種類の実験を行いました。1 つは学習を完走させることです。もう 1 つは短時間の性能計測のみを行い学習を途中 10 ステップで打ち切ります。それぞれの学習のスループットやメモリ消費量、消費電力を測定します。

実験環境

以下に実験環境のハードウェアや OS、フレームワークなどの機種、バージョンを示します。学習実行時に設定する各種パラメータはここには記載せず後述の実験結果と考察セクションに記載します。

表：実験環境

※現在提供中の AIStation とは一部スペックが異なります。ただし GPU のスペック、枚数は同じです。

測定の方法について以下に補足します。

• スループットは verl の出力するログの値 perf/throughput を測定値として利用します。perf/throughput は Tokens processed per second per GPU です（参考 [8]）。
• メモリ消費量、消費電力は nvidia-smi の出力を測定値として利用します。

実験結果

以下に実験結果を示します。

完走実験

まずは以下のパラメータで学習を完走させた結果を示します。ここでいう完走は、後述の性能測定実験とは対照的に、ステップ数を指定するオプション trainer.total_training_steps は何も指定せず --epochs=1 オプションを指定し実行したものです。

表：完走実験のパラメータ

上図は学習を完走させたときの学習曲線です。薄く細い線はステップごとのスコア、太い青と橙の線はデータ全体の 5% の区間で移動平均です。ステップごとの評価だとグラフの通りスコアが乱高下するため移動平均でトレンドを示すようにしています。青い線は 1 GPU での結果、橙の線は 2 GPU での結果です。

このグラフからは、1 GPU でのスコアはほぼ横ばい、2 GPU でのスコアはやや右肩上がりであることがわかります。学習ステップ数が小さいのでこれから学習がうまくいっていることや安定性を論じるのは難しいですが、少なくとも本実験の条件では完走することが確認できました。

上のグラフはメモリの消費量を示します。縦軸がメモリの消費量、横軸が学習の経過時間です。赤い線が 1 GPU での消費量、青い線が 2 GPU での消費量で、黒い実線が 2 GPU の合計メモリサイズ（96 x 2 = 192 GB）、破線が 1 GPU のメモリサイズ（96 GB）を示します。

はじめに、学習時間について見てみます。各折れ線グラフの終端を見てみると、1 GPU は 110 分程度、2 GPU は 70 分程度で学習が終わっています。2 GPU にした分、バッチ数が倍に増えたため学習がより短時間で完了しています。

次にメモリの消費量です。1 GPU での結果に注目すると、黒い破線（96 GB）に触れることはなくぎりぎりのメモリ消費量で学習が完走しています。2 GPU の青い線は少し余裕をもって学習が完走していることを示しています。ただし、GPU 全体のメモリサイズが学習に必要なメモリサイズを許容しても、1 GPU に割り当てられるワークロードに必要なメモリサイズがその GPU のメモリサイズを超えると OOM になってしまうことには注意してください。

最後に、消費電力のグラフを示します。縦軸が消費電力、横軸が学習の経過時間です。赤い線が 1 GPU での消費量、青い線が 2 GPU での消費量です。

[2] によると、Max-Q の TGP は 300 W でした。1 GPU、2 GPU それぞれの場合で、TGP を少し上回る程度で学習が完走しました。

性能測定実験

次に性能測定の結果を示します。前述の通り、10 ステップで学習を打ち切っていることに注意してください。

この実験の結果は、パラメータごとに学習のスループットやそれを元に試算した学習時間、メモリ使用量、消費電力を示します。

表：性能測定実験のパラメータ

完走実験ではモデルは「Qwen2.5-3B-Instruct」のみを使用しましたが、性能測定実験ではそれに加え「Qwen2.5-7B-Instruct」も学習させてみます。また、マイクロバッチサイズも変化させて計測しました。実験のパラメータを増やしたのは、AIStation でどこまで学習させられるかを見たいためです。

ただし、先に断っておくと Qwen2.5-7B-Instruct の学習はいずれのパラメータでも OOM が発生しました。本稿ではグラフなどの実験結果は載せず、OOM してしまったという事実のみを記します。

さて、上図はマイクロバッチサイズを振った時の最大のメモリ消費量を示したグラフです。横軸はマイクロバッチサイズ、縦軸は最大メモリ消費量を示します。2 GPU 分の積み上げ棒グラフであり、青が GPU 0、赤が GPU 1 を示します。

今回の実験では、両 GPU とも 96 GB を上回ることはありませんでした。最初の 10 ステップで学習を打ち切っているのですが、いずれの条件でも OOM せずに完走する可能性は十分ありそうです。

ちなみに、マイクロバッチサイズ 32 と 64 は学習開始後 OOM でエラー終了、128 以上はマイクロバッチサイズがミニバッチサイズより大きくなってしまい学習開始前にアサーションエラーでエラー終了しました。

続いて消費電力のグラフです。上図はマイクロバッチサイズを振った時の最大消費電力を示したグラフです。横軸はマイクロバッチサイズ、縦軸は最大消費電力を示します。2 GPU 分の積み上げ棒グラフであり、青が GPU 0、赤が GPU 1 を示します。いずれの条件でも、600 W を大きく上回ることはありませんでした。学習を完走させていたとしても 600 W を大きく上回ることはまずないと言えそうです。

最後に、下表にスループットを示します。マイクロバッチサイズが大きくなるにつれてスループットも上がっていきますが、4 → 8 → 16 とだんだんサチっていく様子がわかります。

ミニバッチサイズについて

少し話は外れますがミニバッチサイズについて実験を通してわかったことを記しておきます。今回の実験では、学習スクリプトのオプションには 128 を渡しましたが実行時には 64 になっていました。この 64 は、128 を GPU 数の 2 で割ったものだと推測されます。つまり、1 つのミニバッチはさらに GPU の数分のミニバッチに分けられていると考えられます。

本件の発見の経緯についてです。ある実験において、ミニバッチサイズ 128、マイクロバッチサイズ 128 を設定したところ、ミニバッチサイズがマイクロバッチサイズより小さいためにエラーストップすることがありました。ログを読むと上述の通りミニバッチサイズが 128 ではなく 64 になっていたので本件の発見と推測に至りました。

そのときの設定値とエラー内容は以下の通りです。いずれも verl のログより抽出しています。

• 設定

  (TaskRunner pid=143165)                                  'ppo_mini_batch_size': 128,
  

• エラー

  AssertionError: normalized ppo_mini_batch_size 64 should be divisible by ppo_micro_batch_size_per_gpu 128

エラーログを読めばすぐにわかることですが、verl を利用されるときはご注意ください。

まとめ

本検証を通じて、NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition (96GB) を搭載した AIStation における LLM 強化学習（verl）の挙動と特性について以下の知見が得られました。

• 3Bクラスモデルにおける学習の可能性
  • 7B モデルではメモリ消費量がボトルネックとなり OOM が発生しましたが、3B クラスのモデルであれば 96 GB のメモリ内に収まり、一定のマイクロバッチサイズでの動作が確認できました。 このことから、より軽量なモデルを用いたアルゴリズムの動作検証や、小規模な実験サイクルを回す用途においては、本環境が有効に機能する可能性があります。
• オフィス環境への適合性（電力）
  • デュアル GPU 構成で高負荷な学習を実行した際も、消費電力は合計で 600 W 前後で推移しました。 これは一般的な 100 V 電源（1500 W）の容量に対して十分な余裕があり、特別な電源設備を持たないオフィス環境でも、電力面での運用ハードルは低いと考えられます。
• 長時間稼働における安定性
  • 一定の条件下では、エラーによる停止や極端な性能低下は見られず学習プロセスが完走することを確認しました。 ハードウェアおよびソフトウェア（verl）の組み合わせにおいて、一定の安定性が確保されていることが示唆されます。

最後に、AIStation は 3B クラスの軽量モデルを対象とした強化学習タスクにおいて十分なメモリ容量と安定性を備えていることが確認できました。大規模なインフラを必要とせず、身近な環境で学習実験を開始できる点において、ローカル LLM 開発の有力な選択肢の一つとなり得ると言えます。

今回は PPO というアルゴリズムを使った実験を行いましたが、今後は GRPO, GSPO など PPO 以外のアルゴリズムを使った実験も試してみたいと思います。

参考文献

1. Fixstars AIStation – 最新AIをローカルLLMでセキュアに活用 – 株式会社フィックスターズ
2. NVIDIA RTX PRO 6000 Blackwell Max-Q はどのようなGPUなのか？ – Fixstars Tech Blog /proc/cpuinfo
3. NVIDIA RTX PRO 6000 Blackwell Max-Q 上でBLASベンチマークを動かしてみる – Fixstars Tech Blog /proc/cpuinfo
4. NVIDIA Blackwell 上での LoRA チューニングに向けたプロファイリングと最適化 – Fixstars Tech Blog /proc/cpuinfo
5. GitHub – volcengine/verl: verl: Volcano Engine Reinforcement Learning for LLMs
6. HybridFlow: A Flexible and Efficient RLHF Framework
7. openai/gsm8k · Datasets at Hugging Face
8. https://github.com/volcengine/verl/blob/d0997d2a0e14ce8a2a9bb19acbe47ab87ca3a498/verl/trainer/ppo/metric_utils.py#L287
9. [1707.06347] Proximal Policy Optimization Algorithms

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください。