mcpField
S
slime-rl-training

slime-rl-training

github.com/kalarislabsその他7344 回インストール4 日前に更新
リポジトリ

slime-rl-training について

slimeは、LLMの強化学習によるポストトレーニング向けフレームワークです。Megatron-LMによる学習とSGLangによる生成を組み合わせ、データバッファーや複数ターンの学習にも対応します。

主な機能

  1. 01Megatron-LMによる並列学習
  2. 02SGLangを使った生成
  3. 03プロンプトとサンプルの柔軟な管理
  4. 04非同期・複数ターン学習
  5. 05モデル別の設定スクリプト

活用シーン

  1. 01LLMのGRPO学習
  2. 02生成と学習の並行処理
  3. 03ツール利用や多段階推論を行うエージェントの学習

インストール

コマンドをコピーしてターミナルで実行します。

$npx skills add https://github.com/kalarislabs/research-agent-skills --skill slime-rl-training

インストール後、Claude Code で/slime-rl-trainingを入力して読み込まれたか確認します。

同梱ファイルreferences/api-reference.mdreferences/troubleshooting.md

詳細

slime: RLスケーリング向けLLMポストトレーニングフレームワーク

slimeは、清華大学のTHUDMチームによるLLMポストトレーニングフレームワークで、GLM-4.5、GLM-4.6、GLM-4.7を支えています。学習にはMegatron-LM、高スループットなロールアウト生成にはSGLangを接続します。

slimeを使うべきとき

slimeを選ぶべき場合:

  • SGLangによる推論とともにMegatron-LMネイティブの学習が必要な場合
  • 柔軟なデータバッファを使ったカスタムのデータ生成ワークフローが必要な場合
  • GLM、Qwen3、DeepSeek V3、Llama 3モデルを学習させる場合
  • 本番運用の裏付けがある研究グレードのフレームワークが必要な場合(Z.ai)

代替案を検討すべき場合:

  • エンタープライズグレードの安定性機能が必要な場合 → milesを使用
  • バックエンドを柔軟に切り替えたい場合 → verlを使用
  • PyTorchネイティブの抽象化が必要な場合 → torchforgeを使用

主な機能

  • 学習: 完全な並列化サポート(TP、PP、DP、SP)を備えたMegatron-LM
  • ロールアウト: ルーター付きのSGLangベースの高スループット生成
  • データバッファ: 柔軟なプロンプト管理とサンプル保存
  • モデル: GLM-4.x、Qwen3、DeepSeek V3/R1、Llama 3

アーキテクチャ概要

┌─────────────────────────────────────────────────────────┐
│                    Data Buffer                          │
│ - Prompt initialization and management                  │
│ - Custom data generation and filtering                  │
│ - Rollout sample storage                                │
└─────────────┬───────────────────────────┬───────────────┘
              │                           │
┌─────────────▼───────────┐ ┌─────────────▼───────────────┐
│ Training (Megatron-LM)  │ │ Rollout (SGLang + Router)   │
│ - Actor model training  │ │ - Response generation       │
│ - Critic (optional)     │ │ - Reward/verifier output    │
│ - Weight sync to rollout│ │ - Multi-turn support        │
└─────────────────────────┘ └─────────────────────────────┘

インストール

# Recommended: Docker
docker pull slimerl/slime:latest
docker run --rm --gpus all --ipc=host --shm-size=16g \
  -it slimerl/slime:latest /bin/bash

# Inside container
cd /root/slime && pip install -e . --no-deps
ソースからのインストール
git clone https://github.com/THUDM/slime.git
cd slime
pip install -r requirements.txt
pip install -e .

クイックスタート: GRPO学習

# Source model configuration
source scripts/models/qwen3-4B.sh

# Launch training
python train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 4 \
    --rollout-num-gpus 4 \
    --advantage-estimator grpo \
    --use-kl-loss --kl-loss-coef 0.001 \
    --rollout-batch-size 32 \
    --n-samples-per-prompt 8 \
    --global-batch-size 256 \
    --num-rollout 3000 \
    --prompt-data /path/to/data.jsonl \
    ${MODEL_ARGS[@]} ${CKPT_ARGS[@]}

ワークフロー1: 標準GRPO学習

グループ相対アドバンテージで推論モデルを学習させる場合に、このワークフローを使用します。

前提条件チェックリスト
  • Docker環境、またはMegatron-LMとSGLangのインストール済み環境
  • モデルのチェックポイント(HuggingFace形式またはMegatron形式)
  • JSONL形式の学習データ
ステップ1: データを準備する
# data.jsonl format
{"prompt": "What is 2 + 2?", "label": "4"}
{"prompt": "Solve: 3x = 12", "label": "x = 4"}

チャット形式の場合:

{
    "prompt": [
        {"role": "system", "content": "You are a math tutor."},
        {"role": "user", "content": "What is 15 + 27?"}
    ],
    "label": "42"
}
ステップ2: モデルを設定する

事前設定済みのモデルスクリプトを選択します:

# List available models
ls scripts/models/
# glm4-9B.sh, qwen3-4B.sh, qwen3-30B-A3B.sh, deepseek-v3.sh, llama3-8B.sh, ...

# Source your model
source scripts/models/qwen3-4B.sh
ステップ3: 学習を開始する
python train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 8 \
    --rollout-num-gpus 8 \
    --advantage-estimator grpo \
    --use-kl-loss \
    --kl-loss-coef 0.001 \
    --prompt-data /path/to/train.jsonl \
    --input-key prompt \
    --label-key label \
    --apply-chat-template \
    --rollout-batch-size 32 \
    --n-samples-per-prompt 8 \
    --global-batch-size 256 \
    --num-rollout 3000 \
    --save-interval 100 \
    --eval-interval 50 \
    ${MODEL_ARGS[@]}
ステップ4: 学習を監視する
  • TensorBoardを確認する: tensorboard --logdir outputs/
  • 報酬カーブが上昇していることを確認する
  • ノード全体のGPU使用率を監視する

ワークフロー2: 非同期学習

ロールアウトと学習を重ね合わせることで、より高いスループットを得たい場合に非同期モードを使用します。

非同期モードを使うべきとき
  • 生成時間が長い大規模モデルの場合
  • 同期モードでGPUのアイドル時間が多い場合
  • バッファリングに十分なメモリがある場合
非同期学習の開始
python train_async.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 8 \
    --rollout-num-gpus 8 \
    --advantage-estimator grpo \
    --async-buffer-size 4 \
    --prompt-data /path/to/train.jsonl \
    ${MODEL_ARGS[@]}
非同期固有のパラメータ
--async-buffer-size 4        # Number of rollouts to buffer
--update-weights-interval 2  # Sync weights every N rollouts

ワークフロー3: マルチターンのエージェント学習

ツール利用や多段階の推論を行うエージェントを学習させる場合に、このワークフローを使用します。

前提条件
  • マルチターンのロジックを実装したカスタム生成関数
  • ツール/環境のインターフェース
ステップ1: カスタム生成関数を定義する
# custom_generate.py
async def custom_generate(args, samples, evaluation=False):
    """Multi-turn generation with tool calling."""
    for sample in samples:
        conversation = sample.prompt

        for turn in range(args.max_turns):
            # Generate response
            response = await generate_single(conversation)

            # Check for tool call
            tool_call = extract_tool_call(response)
            if tool_call:
                tool_result = execute_tool(tool_call)
                conversation.append({"role": "assistant", "content": response})
                conversation.append({"role": "tool", "content": tool_result})
            else:
                break

        sample.response = response
        sample.reward = compute_reward(sample)

    return samples
ステップ2: カスタム関数を指定して起動する
python train.py \
    --custom-generate-function-path custom_generate.py \
    --max-turns 5 \
    --prompt-data /path/to/agent_data.jsonl \
    ${MODEL_ARGS[@]}

完全なマルチターン検索の例については、examples/search-r1/を参照してください。


設定リファレンス

3種類の引数カテゴリ

slimeでは3種類の引数を使用します:

1. Megatron引数(そのまま渡す):

--tensor-model-parallel-size 2
--pipeline-model-parallel-size 1
--num-layers 32
--hidden-size 4096

2. SGLang引数(--sglang-のプレフィックス付き):

--sglang-mem-fraction-static 0.8
--sglang-context-length 8192
--sglang-log-level INFO

3. slime引数:

# Resource allocation
--actor-num-nodes 1
--actor-num-gpus-per-node 8
--rollout-num-gpus 8
--colocate  # Share GPUs between training/inference

# Data
--prompt-data /path/to/data.jsonl
--input-key prompt
--label-key label

# Training loop
--num-rollout 3000
--rollout-batch-size 32
--n-samples-per-prompt 8
--global-batch-size 256

# Algorithm
--advantage-estimator grpo  # or: gspo, ppo, reinforce_plus_plus
--use-kl-loss
--kl-loss-coef 0.001
主な制約
rollout_batch_size × n_samples_per_prompt = global_batch_size × num_steps_per_rollout

例: 32 × 8 = 256 × 1


データバッファシステム

slimeのデータバッファにより、柔軟なデータ管理が可能になります:

基本的なデータソース
class RolloutDataSource:
    def get_samples(self, num_samples):
        """Fetch prompts from dataset."""
        return self.dataset.sample(num_samples)

    def add_samples(self, samples):
        """Called after generation (no-op by default)."""
        pass
バッファ付きデータソース(オフポリシー)
class RolloutDataSourceWithBuffer(RolloutDataSource):
    def __init__(self):
        self.buffer = []

    def add_samples(self, samples):
        """Store generated samples for reuse."""
        self.buffer.extend(samples)

    def buffer_filter(self, args, buffer, num_samples):
        """Custom selection logic (prioritized, stratified, etc.)."""
        return select_best(buffer, num_samples)

よくある問題と解決策

問題: SGLangエンジンのクラッシュ

症状: 学習の途中で推論エンジンが停止する

解決策:

# Enable fault tolerance
--use-fault-tolerance

# Increase memory allocation
--sglang-mem-fraction-static 0.85

# Reduce batch size
--rollout-batch-size 16
問題: 重み同期のタイムアウト

症状: ロールアウトの後で学習が停止する

解決策:

# Increase sync interval
--update-weights-interval 5

# Use colocated mode (no network transfer)
--colocate
問題: 学習中のOOM

症状: バックワードパスでCUDA OOMが発生する

解決策:

# Enable gradient checkpointing
--recompute-activations

# Reduce micro-batch size
--micro-batch-size 1

# Enable sequence parallelism
--sequence-parallel
問題: データ読み込みの低速化

症状: データ取得中にGPUがアイドル状態になる

解決策:

# Increase data workers
--num-data-workers 4

# Use streaming dataset
--streaming-data

対応モデル

モデルファミリー構成
GLMGLM-4.5、GLM-4.6、GLM-4.7、GLM-Z1-9B
QwenQwen3(4B、8B、30B-A3B)、Qwen3-MoE、Qwen2.5
DeepSeekV3、V3.1、R1
LlamaLlama 3(8B、70B)
その他Kimi K2、Moonlight-16B

各モデルには、scripts/models/に事前設定済みのスクリプトがあります。


応用トピック

コロケーションモード

学習と推論でGPUを共有し、メモリ使用量を抑えます:

python train.py \
    --colocate \
    --actor-num-gpus-per-node 8 \
    --sglang-mem-fraction-static 0.4 \
    ${MODEL_ARGS[@]}
カスタム報酬モデル
# custom_rm.py
class CustomRewardModel:
    def __init__(self, model_path):
        self.model = load_model(model_path)

    def compute_reward(self, prompts, responses):
        inputs = self.tokenize(prompts, responses)
        scores = self.model(inputs)
        return scores.tolist()
--custom-rm-path custom_rm.py
評価のマルチタスク
--eval-prompt-data aime /path/to/aime.jsonl \
--eval-prompt-data gsm8k /path/to/gsm8k.jsonl \
--n-samples-per-eval-prompt 16

リソース

エージェントの作業手順

  1. 環境を確認します。 GPU の種類、メモリ、ドライバー/CUDA のバージョン(nvidia-smi)、フレームワークのバージョン、データセットの場所とサイズを確認します。
  2. 入力を明確にします。 データまたはユーザーから形式、識別子、パラメーターを確認します。結果が変わる値は推測せず、質問してください。
  3. まず小規模に実行します。 スモークテストを行います。小さなモデルまたはサブセットを使い、少ないステップ数で実行して、損失が減少し、チェックポイントが保存されることを確認します。
  4. 上記の手順と参考資料に従って、タスク全体を実行します。
  5. 結果を検証します。 保留データで指標を追跡し、ベースラインと比較して、シード、設定、ハードウェアを記録します。
  6. 報告します。 実行内容(バージョン、コマンド、パラメーター)、確認した内容、未解決の不確実な点を明記します。
次の状況が発生した場合対応
CUDA のメモリ不足バッチサイズを小さくする、勾配蓄積/チェックポインティングまたは混合精度を有効にする、あるいはモデルをシャーディングします。
損失が NaN になる、または発散する学習率を下げ、データに無効な値がないか確認し、勾配クリッピングを有効にします。
この手順に記載された関数、フラグ、エンドポイントがインストール済みバージョンにないインストール済みバージョンのドキュメント(help()、--help、公式ドキュメント)を確認して適宜対応し、ユーザーに伝えます。API を捏造してはいけません。
必須の入力、識別子、パラメーターが曖昧ユーザーに質問するか、実行前に前提を明示します。

誠実性に関するルール

  • 結果、パラメーター、識別子、引用、統計を決して捏造しないでください。実行または検証できない場合は、その旨を明確に伝えてください。
  • ログなしで学習結果を主張してはいけません。大規模なジョブを開始する前に、計算コストを見積もってください。
  • ここに記載されているバージョン固有の詳細は古くなっている可能性があるため、インストール済みバージョンの公式ドキュメントで確認してください。
  • 費用が発生する、共有 GPU またはクラウドのクォータを消費する、個人データや患者データを扱う、または元に戻せない操作を行う前に、ユーザーに確認してください。

関連スキル

  • verl-rl-training:verl(Volcano Engine RL)を使用した LLM の強化学習について説明します。
  • miles-rl-training:slime の本番環境向けフォークである miles を使用した、エンタープライズグレードの強化学習について説明します。
  • nanogpt:Karpathy による最小限の PyTorch GPT 実装(model.py と train.py)を提供し、文字レベルの Shakes… の学習ワークフローを説明します。

コメント