1000 GPU 規模の非同期 RL 学習フレームワーク「PRIME-RL」— LLM 後学習をスケールさせる

タグ PRIME-RL強化学習LLM後学習分散学習vLLMPyTorch FSDP2オープンソース大規模クラスターcolumnコラムlinuxLinuxGitHubPrimeIntellect-aiprime-rl

ひとことでいうと

PRIME-RL は、大規模な強化学習(AI にタスクを自力で学ばせる手法)を高速・大規模に実行できるオープンソースのフレームワークです。PrimeIntellect 社が開発・公開しており、1,000 枚超の GPU(グラフィックボード)を並べたクラスターでも動くよう設計されています。推論エンジン「vLLM」と分散学習の仕組み「PyTorch FSDP2」を組み合わせることで、1 兆パラメータを超える巨大モデルでも効率よく後学習(Post-Training)を行えます。SFT・強化学習・評価をひとつのフレームワークで完結できる点が最大の魅力です。研究段階の小規模実験から、本番スケールの大規模クラスターまで、同じ設定ファイルで対応できます。


こんな人におすすめ

1. 難問を解ける LLM を自力で育てたい ML 研究者・エンジニア

数学(MATH・AIME)やソフトウェアエンジニアリング(SWE)など、高難度タスクに強い LLM を強化学習で鍛えたい方に最適です。Slurm や Kubernetes(大規模クラスターを管理するソフト)による複数サーバーへの分散学習にも標準対応しており、30B〜100B を超えるモデルも扱えます。

2. 最新 OSS モデルの後学習を再現・改善したいオープンモデル開発者

Qwen3・GLM-4.5・INTELLECT-3 など最新モデルの学習レシピを再現・改良したい方に向いています。ハイパーパラメータ(学習率やバッチサイズなど学習の設定値)はすべて TOML ファイル(テキストベースの設定ファイル形式)で一元管理されており、examples/ にはフルスケールの再現レシピも収録されています。

3. まず小規模で試したい個人研究者

GPU が 1〜8 枚の環境でも、examples/ 以下のチュートリアル(Reverse Text・Wordle・Alphabet Sort など)ですぐに実験を始められます。コンシューマ向け GPU(RTX 3090・4090)での動作も公式が確認済みですので、ゲーミング PC でも試せます。


インストール・使い方

前提条件

PRIME-RL を動かすには NVIDIA 製の GPU が最低 1 基必要です。CPU(パソコンの頭脳)だけの環境では起動できません。公式に動作確認されているハードウェアは RTX 3090 / 4090 / 5090、A100、H100、H200、B200 です。Python(プログラミング言語の一種)は バージョン 3.12 が必要で、パッケージ管理には uv という高速ツールを使います。pip install prime-rl では動作しない点に注意してください。


Step 1: クイックインストール(推奨)

ターミナル(文字を入力してパソコンに命令を送る画面)を開き、以下のコマンドをコピー&ペーストして実行してください。公式が用意したスクリプトが、必要なものをまとめてセットアップしてくれます。

curl -sSL https://raw.githubusercontent.com/PrimeIntellect-ai/prime-rl/main/scripts/install.sh | bash

このコマンドは、インストールスクリプトをダウンロードしてそのまま実行するものです。インターネット接続が必要ですが、手順はこれ 1 行だけで完了します。


Step 2: 手動インストール(細かく確認したい場合)

# リポジトリ(ソースコードの置き場)をクローン(手元にコピー)する
git clone https://github.com/PrimeIntellect-ai/prime-rl.git
cd prime-rl

# uv(高速パッケージマネージャ)をインストールする
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env

# 全依存パッケージを uv.lock から再現性ある形でインストール
uv sync --all-extras

git clone でコードを手元にコピーし、uv sync で必要なライブラリをまとめてインストールします。uv.lock というファイルにバージョンが固定されているため、誰がインストールしても同じ環境が再現されます。


Step 3: 動作確認

# Python バージョン確認(3.12 と表示されれば OK)
uv run python -V

# 高速アテンション計算ライブラリの確認
uv run python -c "import flash_attn"

# SFT トレーナーのテスト起動(GPU 1 基必要)
uv run sft @ configs/debug/sft/train.toml

# RL トレーナーのテスト起動(GPU 1 基必要)
uv run trainer @ configs/debug/rl/train.toml

# 推論サーバーのテスト起動(GPU 1 基必要)
uv run inference @ configs/debug/infer.toml

まず Python のバージョンと必要ライブラリを確認します。その後、SFT・RL・推論の 3 つをそれぞれデバッグ設定で起動してみることで、環境が正しく整っているかをチェックできます。


Step 4: 外部サービスとの連携(任意)

# W&B(Weights & Biases)でトレーニングの進捗をグラフで可視化する場合
uv run wandb login

# HuggingFace のゲートモデルやプライベートモデルを使う場合
uv run hf auth login

W&B(ダブリューアンドビー)は学習の損失値・精度などをリアルタイムでグラフ表示できる無料サービスです。HuggingFace(ハギングフェイス)はモデルや学習データを共有するプラットフォームで、認証が必要なモデルを使う際にログインします。どちらも任意ですが、学習の進捗確認に役立ちます。


動かしてみた

PRIME-RL は GPU が実行に必須なため、ここでは 試す前に知っておくとよいこと として、README が示す動作要件をまとめます。

項目内容
GPUNVIDIA RTX 3090 以上(最低 1 基)
Python3.12
OSLinux(CUDA 対応環境)
パッケージマネージャuvpip は非対応)
ディスク容量モデルサイズに依存(Qwen3-30B なら 60 GB 以上推奨)

uv.lock ファイルが同梱されているため、uv sync を 1 回実行するだけで依存関係が完全に再現される設計になっています。pip 経由のインストールは現時点でサポートされていないため、必ず uv を使ってください。小規模な動作確認をするなら、まず configs/debug/ 以下の設定で各コンポーネントを順番に起動してみるのがスムーズです。


デモについて

PRIME-RL は NVIDIA GPU の実行が必須であるため、CPU のみの環境やブラウザ上でのデモ実行はできません。モデルのロードだけで数十 GB のメモリを必要とする場合もあり、Gradio などを使った軽量ウェブデモへの変換も困難です。GPU 環境をお持ちの方は、公式の examples/ ディレクトリに用意されている Reverse Text・Wordle・Alphabet Sort・Wiki Search などのチュートリアルから試してみてください。


はじめの一歩:実践のコツ

  • GPU 1 枚から始める: examples/reverse_text/Qwen3-0.6B(6 億パラメータの小型モデル)を使ったチュートリアルが最短ルートです。SFT → RL の流れをコンシューマ GPU 1 枚で数分体験できます。
  • LoRA で省メモリに試す: examples/alphabet_sort/ は LoRA(少ないパラメータでモデルを微調整する手法)を使っており、H100 1 枚で 1 時間強で完了します。フルファインチューニングよりメモリ消費が大幅に少なく、最初の実験に最適です。
  • 設定は TOML ファイルで一元管理: 学習率・バッチサイズ・ロールアウト数はすべて configs/ 以下のテキストファイルを編集するだけで変更できます。コードを書き換える必要がないため、ハイパーパラメータ調整が非常に楽です。
  • デバッグ設定で疎通確認: configs/debug/ 以下の設定は最小限のリソースで動かすためのものです。本番設定に切り替える前に、まずここで環境が正しく動いているか確認する習慣をつけると安心です。
  • W&B で学習の進捗を可視化: uv run wandb login でログインしておくと、損失値や報酬の変化をリアルタイムでグラフ確認できます。学習が意図通りに進んでいるか、早い段階で気づけます。

活用アイデア

  • 数学・論理推論の強化: MATH や AIME などの難問データセットに RL 学習を適用し、CoT(Chain-of-Thought:段階的に考えさせる推論手法)能力を引き上げます。examples/hendrycks_sanity/ が要因比較実験の出発点として最適です。
  • SWE エージェントの開発: examples/qwen30b_swe/ を参考に、コード生成・バグ修正を自動化するソフトウェアエンジニアリングエージェントを学習させられます。verifiers ライブラリと連携すれば、独自のコードテスト環境も定義できます。
  • 独自タスクへの RL 適用: verifiers ライブラリを使って独自のリワード関数(モデルへの「点数」のつけ方)と環境を定義し、特定業務に特化した RL 学習を実施できます。Environments Hub 経由で既製の検証環境も利用可能です。
  • VLM(視覚言語モデル)の後学習: Qwen3-VL など画像とテキストを扱うモデルのマルチモーダル RL トレーニングにも対応しており、画像理解や視覚的推論の能力強化に応用できます。
  • 大規模クラスターへのスケールアウト: examples/ に Slurm ジョブスクリプトや Kubernetes マニフェストが用意されており、32〜2,048 GPU の大規模クラスター実験にそのまま使えます。研究段階から本番スケールへの移行がスムーズです。

用語とポイント解説

FSDP2(エフエスディーピー2) PyTorch の「Fully Sharded Data Parallel v2」の略で、大きなモデルを複数の GPU に分割して学習する技術です。かんたんに言うと、1 台の GPU に収まらない巨大なモデルを複数台で分担して学習するための仕組みです。PRIME-RL はこれを活用することで、1 兆パラメータ超のモデルも扱えます。

vLLM(ブイエルエルエム) 高スループット(大量処理)に特化した LLM 推論エンジンです。PagedAttention という技術で GPU のメモリを効率的に活用します。かんたんに言うと、大量の文章を素早く生成するための「高速エンジン」です。PRIME-RL では学習中の推論フェーズにこれを使います。

MoE(エムオーイー) 「Mixture of Experts(専門家の混合)」の略で、モデルを複数の「専門家」サブネットワークに分割し、入力に応じて使う専門家を切り替える構造です。かんたんに言うと、「問題の種類によって担当者を変える分業体制」をモデルの中に取り入れた設計です。パラメータ数が大きくても計算量を抑えられます。

EP(イーピー) 「Expert Parallelism(専門家並列化)」の略で、MoE の専門家モジュールを複数の GPU に分散して計算する手法です。かんたんに言うと、MoE の「担当者」それぞれを別々の GPU に割り振ることで、大規模な MoE モデルを効率よく動かす仕組みです。

CP(シーピー) 「Context Parallelism(文脈並列化)」の略で、長い入力テキストのシーケンスを複数の GPU に分割して処理する手法です。かんたんに言うと、長い文章を複数の GPU で「分担して読む」ことで、超長文のトレーニングを可能にします。

SFT(エスエフティー) 「Supervised Fine-Tuning(教師ありファインチューニング)」の略で、正解ラベルが付いたデータを使ってモデルを微調整する手法です。かんたんに言うと、「お手本を見て練習する」段階のトレーニングです。RL 学習の前段として、モデルの初期性能を高めるために使います。

GRPO(ジーアールピーオー) 「Group Relative Policy Optimization」の略で、強化学習アルゴリズムの一種です。グループ内の複数サンプルを比べた相対的な報酬でモデルの方策(行動方針)を更新します。かんたんに言うと、「他の回答と比べてどれが良かったか」をもとに学習させる仕組みです。PPO に比べて参照モデルが不要で実装が軽量です。

LoRA(ローラ) 「Low-Rank Adaptation」の略で、モデルの全パラメータではなく小さな追加行列だけを学習することで、フルファインチューニングより大幅にメモリを削減できる手法です。かんたんに言うと、「全部書き直さずに一部だけ改訂する」効率的な微調整方法です。コンシューマ GPU でも大きめのモデルを扱えるようになります。

PD 分離(プリフィル・デコード分離) 「Prefill-Decode 分離推論」の略で、プリフィル(入力を一度に処理するフェーズ)とデコード(トークンを 1 つずつ生成するフェーズ)を別々のノード(サーバー)で処理する手法です。かんたんに言うと、「仕込み係」と「調理係」を分けることで、推論のスループットを大きく向上させます。


ぜひ数学・論理推論の強化学習や SWE エージェントの開発、さらには大規模クラスターを使ったオープンモデルの後学習再現などに活用してみてはいかがでしょうか。