📖

Verlフレームワークを用いたLLMの強化学習(PPO,GRPO,DAPO)

に公開

はじめに

この記事では、VeRLというフレームワークを使ってMegatron-LMをベースにLLMをRL(PPO、GRPO、DAPO)する方法について解説します。

公式リポジトリ・ドキュメント:
https://github.com/volcengine/verl
https://verl.readthedocs.io/en/latest/index.html

経緯

これまで、trl、ms swift、nemo rl、megatron bridge、verlなどさまざまな学習用ライブラリを触ってきました。

今回作成したモデルは、SFT、SimPO、GRPOで学習されており、SFT・SimPOはms swift、GRPOはVeRLを使用しました。ms swiftの記事はたくさんあるため、VeRLの記事を書くことにしました。

作成したモデルを以下に示します。
学習したモデル:
https://huggingface.co/puwaer/Qwen3-4B-Thinking-2507-GRPO-Uncensored

VeRLの非同期学習アーキテクチャについて

VeRLは、強化学習における生成(rollout)と訓練(training)の並列化を実現するOne Step Off Async Trainerを実装しています。これにより、GPU利用効率を大幅に向上させることができます。

One Step Off Async Trainerとは

従来の強化学習では、生成フェーズと訓練フェーズが順次実行されるため、一方のフェーズ実行中に他方のGPUリソースがアイドル状態になるという問題がありました。特に、長時間かかるサンプル生成(long-tail sample generation)の間、訓練用のGPUが待機状態になることで効率が低下していました。

One Step Off Async Trainerは、この問題を以下の方法で解決します:

  1. 生成と訓練の並列化:前ステップで生成されたサンプルを使って現在のステップで訓練を行う
  2. リソースの適切な分割:生成専用のリソースを割り当て、残りを訓練に自動的に割り当てる
  3. GPU待機時間の削減:生成フェーズに割り当てるリソースを削減することで、長時間のサンプル生成中のGPUアイドル時間を軽減

このプロセス全体を通じて、生成パラメータと訓練パラメータは1ステップ分のoff-policyを維持します。

VeRL One Step Off Async Trainer

環境構築

Singularityを使った環境構築

まずはVeRL用のSIFファイルを作成します。

verl_app.def:

Bootstrap: docker
From: verlai/verl:app-verl0.6-transformers4.56.1-sglang0.5.2-mcore0.13.0-te2.2

%post
    set -e

    echo "Starting update for verl..."
    pip install --no-cache-dir --upgrade pip setuptools
    pip install --no-cache-dir wandb ipykernel
    pip install verl==0.6.1
    echo "Build complete."

%environment
    export LC_ALL=C
    export LANG=C.UTF-8

以下のコマンドでビルドしてください。

sudo singularity build ./verl_app.sif ./verl_app.def

以下のコマンドでSingularityコンテナに入ります。

singularity shell --nv \
  --bind /path/to/your/project \
  ./verl_app.sif

GRPOの学習

データセット準備

学習に利用するデータセットを準備します。
parquetファイルのtestファイルとtrainファイルを作成してください。

{
  "prompt": [
    {
      "role": "user",
      "content": "使用するプロンプト"
    }
  ],
  "data_source": "データセット名",
  "reward_model": {
    "ground_truth": "報酬関数で使用する文章",
    "style": "model"
  },
  "extra_info": {
    "split": "train"    #test
  }
}

詳しくは以下のドキュメントを確認してください。
https://verl.readthedocs.io/en/latest/preparation/prepare_data.html#prepare-data-for-post-training

学習の実行(カスタム報酬関数版)

VeRLでは、カスタム報酬関数を使用してGRPO学習を行うことができます。ここでは、以下の4つのスクリプトを使った学習方法を解説します。

スクリプト構成の概要

1. run_grpo_qwen3_4b_mega_verl.sh

  • 役割: Slurmジョブの定義と投入を行うメインスクリプトです。
  • 機能:
    • Slurmのリソース確保(ノード数、GPU数、パーティション設定)。
    • 環境変数の設定。
    • 各計算ノードに対して helper_grpo_qwen3_4b.sh を起動するコマンドを発行します。
  • 使用方法: sbatch run_grpo_qwen3_4b_mega_verl.sh で学習を開始します。

2. helper_grpo_qwen3_4b.sh

  • 役割: 各計算ノード内で実行されるセットアップスクリプトです。
  • 機能:
    • Rayクラスターの構築: ヘッドノード(Head Node)とワーカーノード(Worker Node)の判別を行い、Rayインスタンスを起動します。
    • 学習プロセスの開始: Rayクラスターの準備完了後、VerlベースのGRPO学習コマンドを実行します。

3. reward_server.py

  • 役割: 生成された回答に対してスコア(報酬)を計算するFastAPIベースのHTTPサーバーです。
  • 機能:
    • 学習ループからの推論テキストを受け取り、安全性や有用性のスコアを返却します。
  • ⚠️ 実装上の注記:
    • 現在は trl で学習されたReward Modelのフォーマット制約により、Hugging Face Transformersベースで実装されています。
    • 推奨事項: 推論速度とスループット向上のため、フォーマットの問題が解決可能な場合は vLLM をバックエンドに使用することを強く推奨します。

4. safe_hack_reward_request.py

  • 役割: 学習プロセス(Verl)と報酬サーバーを繋ぐインターフェースです。
  • 機能:
    • 学習中に生成されたレスポンスをペイロードとして、reward_server.py へHTTPリクエストを送信します。
    • サーバーから返却されたスコアを整形し、RLアルゴリズムへ渡します。

メインスクリプト

run_grpo_qwen3_4b_mega_verl.sh:

#!/bin/bash
#SBATCH --job-name=sft-custom          # ジョブ名
#SBATCH --partition=your_partition      # パーティション名(環境に合わせて変更)
#SBATCH --nodes=1                       # ノード数
#SBATCH --gpus=8                        # GPU数
#SBATCH --output=logs/sft-%j.out        # 標準出力ログ(%jはジョブID)
#SBATCH --error=logs/sft-%j.err         # エラーログ

# 設定セクション
export CONFIG_NUM_NODES=1
export CONFIG_GPUS_PER_NODE=4
export CONFIG_HOME_PATH="/home/your_username"  # ご自身のホームディレクトリに変更

# 報酬関数の設定
export CONFIG_REWARD_SCRIPT_PATH="${CONFIG_HOME_PATH}/your_project/safe_hack_reward_request.py"
export CONFIG_REWARD_FUNCTION_NAME="safe_hack_reward"
export CONFIG_REWARD_SERVER_PORT=8000
export CONFIG_REWARD_MODEL_PATH="${CONFIG_HOME_PATH}/model/Unsafe-Reward-Qwen3-1.7B"
export CONFIG_REWARD_DEVICE="cuda"
export CONFIG_USE_SEPARATE_REWARD_NODE=false

# データセットとモデルの設定
export CONFIG_TRAIN_FILE="${CONFIG_HOME_PATH}/data/test_data.parquet"
export CONFIG_TEST_FILE="${CONFIG_HOME_PATH}/data/train_data.parquet"
export CONFIG_MODEL_PATH="${CONFIG_HOME_PATH}/model/Qwen/Qwen3-4B-Thinking-2507"

module load singularitypro
mkdir -p ${CONFIG_HOME_PATH}/.cache
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY

export MASTER_ADDR=$(hostname -s)
export MASTER_PORT=29500
export NUM_NODES=$SLURM_NNODES
export GPUS_PER_NODE=${CONFIG_GPUS_PER_NODE}

export VERL_ROOT_DIR="${CONFIG_HOME_PATH}/verl"

echo "Configuration Summary"
echo "Master Node: $MASTER_ADDR"
echo "Verl Path: $VERL_ROOT_DIR"

HELPER_SCRIPT="${CONFIG_HOME_PATH}/your_project/helper_grpo_qwen3_4b.sh"

srun --ntasks=$NUM_NODES --ntasks-per-node=1 \
singularity exec --nv \
  --home ${CONFIG_HOME_PATH} \
  -B ${CONFIG_HOME_PATH}:${CONFIG_HOME_PATH} \
  -B ${CONFIG_HOME_PATH}/.cache:/root/.cache \
  -B /home/your_group:/home/your_group \
  -B /dev/shm:/dev/shm \
  --env XDG_CACHE_HOME=${CONFIG_HOME_PATH}/.cache \
  --env HF_HOME=${CONFIG_HOME_PATH}/.cache/huggingface \
  --env TORCH_EXTENSIONS_DIR=${CONFIG_HOME_PATH}/.cache/torch_extensions \
  --env PYTHONPATH=${VERL_ROOT_DIR}:${CONFIG_HOME_PATH}/Megatron-LM:$PYTHONPATH \
  --env NPROC_PER_NODE=4 \
  --env CUDA_VISIBLE_DEVICES=0,1,2,3 \
  --env NCCL_SOCKET_IFNAME=your_network_interface \
  --env GLOO_SOCKET_IFNAME=your_network_interface \
  --env TP_SOCKET_IFNAME=your_network_interface \
  --env SLURM_JOB_NODELIST=$SLURM_JOB_NODELIST \
  --env SLURM_NNODES=$SLURM_NNODES \
  --env CONFIG_HOME_PATH=${CONFIG_HOME_PATH} \
  --env CONFIG_REWARD_SCRIPT_PATH=${CONFIG_REWARD_SCRIPT_PATH} \
  --env CONFIG_REWARD_FUNCTION_NAME=${CONFIG_REWARD_FUNCTION_NAME} \
  --env CONFIG_REWARD_SERVER_PORT=${CONFIG_REWARD_SERVER_PORT} \
  --env CONFIG_REWARD_MODEL_PATH=${CONFIG_REWARD_MODEL_PATH} \
  --env CONFIG_REWARD_DEVICE=${CONFIG_REWARD_DEVICE} \
  --env CONFIG_USE_SEPARATE_REWARD_NODE=${CONFIG_USE_SEPARATE_REWARD_NODE} \
  --env CONFIG_GPUS_PER_NODE=${CONFIG_GPUS_PER_NODE} \
  --env CONFIG_TRAIN_FILE=${CONFIG_TRAIN_FILE} \
  --env CONFIG_TEST_FILE=${CONFIG_TEST_FILE} \
  --env CONFIG_MODEL_PATH=${CONFIG_MODEL_PATH} \
  ${CONFIG_HOME_PATH}/singularity/verl_app_v061.sif \
  bash $HELPER_SCRIPT $MASTER_ADDR

メインスクリプトの解説

このスクリプトは、SLURMクラスタ上でGRPO学習を実行するためのエントリーポイントです。

主な役割:

  1. SLURMジョブの設定(GPU数、ノード数、ログ出力先など)
  2. 環境変数の設定(報酬関数、データセット、モデルパスなど)
  3. Singularityコンテナの起動とヘルパースクリプトの実行

重要な設定項目:

報酬関数の設定

報酬関数は、強化学習においてモデルの出力を評価する最も重要な要素です。

# 報酬関数クライアントスクリプトのパス
export CONFIG_REWARD_SCRIPT_PATH="${CONFIG_HOME_PATH}/your_project/safe_hack_reward_request.py"

# スクリプト内の報酬関数名(この名前でVeRLから呼び出される)
export CONFIG_REWARD_FUNCTION_NAME="safe_hack_reward"

# 報酬サーバーのポート番号
export CONFIG_REWARD_SERVER_PORT=8000

# 報酬モデルのパス(スコア計算に使用するモデル)
export CONFIG_REWARD_MODEL_PATH="${CONFIG_HOME_PATH}/model/Unsafe-Reward-Qwen3-1.7B"

# 報酬モデルの実行デバイス(cuda/cpu/auto)
export CONFIG_REWARD_DEVICE="cuda"

# 報酬サーバーを別ノードで実行するか(falseの場合、学習ノードで実行)
export CONFIG_USE_SEPARATE_REWARD_NODE=false
データセットとモデルの設定
# 学習用データセット
export CONFIG_TRAIN_FILE="${CONFIG_HOME_PATH}/data/test_data.parquet"

# 検証用データセット
export CONFIG_TEST_FILE="${CONFIG_HOME_PATH}/data/train_data.parquet"

# ベースモデルのパス(ファインチューニング対象)
export CONFIG_MODEL_PATH="${CONFIG_HOME_PATH}/model/Qwen/Qwen3-4B-Thinking-2507"

ヘルパースクリプト

helper_grpo_qwen3_4b.sh:

#!/bin/bash

HOSTNAME=$(hostname -s)
RANK=$SLURM_PROCID
NNODES=$SLURM_NNODES

HEAD_NODE=$1

if [ -z "$HEAD_NODE" ]; then
    echo "Error: HEAD_NODE (Master Address) not received!"
    exit 1
fi

echo "[${HOSTNAME}] Rank ${RANK} started (Total: ${NNODES})."
echo "[${HOSTNAME}] Head Node IP/Hostname: ${HEAD_NODE}"

USE_SEPARATE_REWARD_NODE=${CONFIG_USE_SEPARATE_REWARD_NODE:-true}
NUM_GPUS=${CONFIG_GPUS_PER_NODE:-2}
REWARD_SERVER_PORT=${CONFIG_REWARD_SERVER_PORT:-8000}

HEAD_RANK=0
NODES_ARRAY=( $HEAD_NODE ) 

# ノードの役割決定
if [[ "$USE_SEPARATE_REWARD_NODE" == "true" ]]; then
    # モード1
    if [[ "$NNODES" -eq 1 ]]; then
        REWARD_NODE=${HEAD_NODE}
    else
    fi
    NUM_TRAIN_NODES=$(( NNODES - 1 ))
    if [[ "$NNODES" -eq 1 ]]; then NUM_TRAIN_NODES=1; fi
    echo "[${HOSTNAME}] Mode: Separate Reward Node (Adapted for 1 node)"
else
    # モード2
    REWARD_RANK=-1
    REWARD_NODE=${HEAD_NODE}
    NUM_TRAIN_NODES=$NNODES
    echo "[${HOSTNAME}] Mode: Reward Server on Training Node"
fi

# 報酬サーバーのURL
export REWARD_SERVER_URL="http://${REWARD_NODE}:${REWARD_SERVER_PORT}"

# Role: Ray Head & Trainer (Rank 0)
if [[ "$RANK" == "$HEAD_RANK" ]]; then
    echo "[${HOSTNAME}] Role: RAY HEAD (Rank $RANK)"
    echo "[${HOSTNAME}] Reward Server URL: ${REWARD_SERVER_URL}"
    echo "[${HOSTNAME}] Training on $NUM_TRAIN_NODES nodes."

    # 報酬サーバーの起動
    echo "[${HOSTNAME}] Starting Reward Server in background on port ${REWARD_SERVER_PORT}..."
    REWARD_SERVER_SCRIPT="${CONFIG_HOME_PATH}/your_project/reward_score/reward_server.py"
    python3 ${REWARD_SERVER_SCRIPT} > /tmp/reward_server_${SLURM_JOB_ID}.log 2>&1 &
    REWARD_SERVER_PID=$!
    echo "[${HOSTNAME}] Reward Server started with PID ${REWARD_SERVER_PID}"
    sleep 10

    # 環境変数のセット
    HOME_PATH=${CONFIG_HOME_PATH:-/home/your_username}
    train_file=${CONFIG_TRAIN_FILE}
    test_file=${CONFIG_TEST_FILE}
    reward_script_path=${CONFIG_REWARD_SCRIPT_PATH}
    reward_function_name=${CONFIG_REWARD_FUNCTION_NAME}
    MODEL_PATH=${CONFIG_MODEL_PATH}

    echo "[${HOSTNAME}] Launching Ray head ($HEAD_NODE) ..."
    
    # Ray起動
    ray start --head --node-ip-address $HEAD_NODE --port 6379 --num-cpus 72 --num-gpus $NUM_GPUS --disable-usage-stats --block &
    
    sleep 30

    echo "[${HOSTNAME}] Starting verl trainer..."
    python3 -m verl.trainer.main_ppo --config-path=config \
        --config-name='ppo_megatron_trainer.yaml' \
        algorithm.adv_estimator=grpo \
        data.train_files="$train_file" \
        data.val_files="$test_file" \
        data.return_raw_chat=true \
        data.train_batch_size=256 \
        data.max_prompt_length=2048 \
        data.max_response_length=4096 \
        data.filter_overlong_prompts=True \
        data.truncation='error' \
        custom_reward_function.path=$reward_script_path \
        custom_reward_function.name=$reward_function_name \
        actor_rollout_ref.model.path=$MODEL_PATH \
        actor_rollout_ref.actor.optim.lr=1e-6 \
        actor_rollout_ref.actor.ppo_mini_batch_size=128 \
        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
        actor_rollout_ref.actor.megatron.optimizer_offload=True \
        actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \
        actor_rollout_ref.actor.megatron.tensor_model_parallel_size=1 \
        +actor_rollout_ref.actor.megatron.override_transformer_config.fp8="e4m3" \
        +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_recipe="blockwise" \
        +actor_rollout_ref.actor.optim.override_optimizer_config.fp8_recipe="blockwise" \
        actor_rollout_ref.actor.use_kl_loss=True \
        actor_rollout_ref.actor.kl_loss_coef=0.001 \
        actor_rollout_ref.actor.kl_loss_type=low_var_kl \
        actor_rollout_ref.actor.entropy_coeff=0 \
        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
        actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
        actor_rollout_ref.rollout.name=sglang \
        actor_rollout_ref.rollout.mode=async \
        actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
        actor_rollout_ref.rollout.max_num_seqs=256 \
        actor_rollout_ref.rollout.n=4 \
        actor_rollout_ref.rollout.enforce_eager=true \
        actor_rollout_ref.rollout.free_cache_engine=true \
        actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
        actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \
        actor_rollout_ref.ref.megatron.tensor_model_parallel_size=1 \
        algorithm.use_kl_in_reward=False \
        trainer.critic_warmup=0 \
        trainer.logger=['console','wandb'] \
        trainer.project_name='llm_safe_hack' \
        trainer.experiment_name='qwen3_4b_grpo_megatron' \
        trainer.default_local_dir=${HOME_PATH}/output_model/qwen3_4b_grpo_megatron \
        trainer.n_gpus_per_node=$NUM_GPUS \
        trainer.nnodes=$NUM_TRAIN_NODES \
        trainer.save_freq=25 \
        trainer.test_freq=5 \
        trainer.total_epochs=1

else
    echo "[${HOSTNAME}] This node is a RAY WORKER (Rank $RANK)."
    echo "[${HOSTNAME}] Connecting to Ray head ($HEAD_NODE) ..."
    
    sleep 30
    
    ray start --address "$HEAD_NODE:6379" --num-cpus 72 --num-gpus $NUM_GPUS --block
fi

ヘルパースクリプトの解説

このスクリプトは、各計算ノードで実行され、Rayクラスタの構築と学習の実行を担当します。

主な役割:

  1. ノードの役割決定:Rank 0のノードをRay Headとして設定、他のノードをWorkerとして設定
  2. 報酬サーバーの起動:Ray Headノードで報酬サーバーをバックグラウンドで起動
  3. Rayクラスタの構築:分散学習のための通信基盤を構築
  4. VeRL学習の実行:PPO/GRPOトレーナーの起動

実行フロー:

重要なポイント:

  • 報酬サーバーの起動:Rank 0のノードでのみ、reward_server.pyをバックグラウンドプロセスとして起動します
  • 環境変数の引き継ぎ:メインスクリプトで設定した環境変数(報酬関数のパスなど)をコンテナ内に引き継ぎます
  • VeRL設定custom_reward_function.pathcustom_reward_function.nameで報酬関数を指定します

報酬関数の実装

報酬サーバー(reward_server.py)

reward_server.py:

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import uvicorn
import os

app = FastAPI()

MODEL_PATH = os.environ.get("CONFIG_REWARD_MODEL_PATH", "/home/your_username/model/Unsafe-Reward-Qwen3-1.7B")
# デバイス設定: 環境変数で指定可能 (cpu/cuda/auto)
DEVICE_CONFIG = os.environ.get("CONFIG_REWARD_DEVICE", "auto")
if DEVICE_CONFIG == "auto":
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
elif DEVICE_CONFIG in ["cpu", "cuda"]:
    DEVICE = DEVICE_CONFIG
else:
    print(f"[Warning] Invalid device config '{DEVICE_CONFIG}', using auto")
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print(f"[Server] Loading model from {MODEL_PATH} on {DEVICE}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map=DEVICE,
    num_labels=1
).eval()
print("[Server] Model loaded successfully.")

class RewardRequest(BaseModel):
    prompt: str
    response: str

@app.post("/score")
async def calculate_score(req: RewardRequest):
    try:
        inputs = tokenizer(
            req.prompt, 
            req.response, 
            return_tensors="pt", 
            truncation=True, 
            max_length=2048
        ).to(model.device)

        with torch.no_grad():
            outputs = model(**inputs)
        
        score = outputs.logits[0].item()
        print(f"[RewardServer] Score: {score:.4f} (Prompt len: {len(req.prompt)})", flush=True)
        return {"score": float(score)}

    except Exception as e:
        print(f"[Error] {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

報酬サーバーの仕組み:

このスクリプトは、FastAPIを使ってHTTPサーバーを立ち上げ、報酬モデルを使ってスコア計算を行います。

主な処理:

  1. モデルのロードCONFIG_REWARD_MODEL_PATHで指定された報酬モデルをロード

    • モデルはAutoModelForSequenceClassificationとして読み込み
    • num_labels=1でスカラー値(スコア)を出力
  2. エンドポイント /score

    • リクエスト:prompt(プロンプト)とresponse(モデルの生成テキスト)、ground_truth(報酬関数で使用する文章) を受け取る
    • 処理:トークナイズしてモデルに入力し、スコアを計算
    • レスポンス:{"score": float}の形式でスコアを返す
  3. デバイス管理:環境変数CONFIG_REWARD_DEVICEでCPU/GPU実行を制御可能

報酬関数クライアント(safe_hack_reward_request.py)

safe_hack_reward_request.py:

import requests
import os
import sys
import re
import json 
import time  

BASE_URL = os.environ.get("REWARD_SERVER_URL", "http://localhost:8000").rstrip("/")
SCORE_ENDPOINT = f"{BASE_URL}/score"

REWARD_MAX_SCORE = 30.0
REWARD_MIN_SCORE = -40.0

DEFAULT_LOG_PATH = os.path.expanduser("./log/reward_logs.jsonl") 
LOG_FILE_PATH = os.environ.get("REWARD_LOG_PATH", DEFAULT_LOG_PATH)

print(f"Reward Server Endpoint: {SCORE_ENDPOINT}")
print(f"Logging rewards to: {LOG_FILE_PATH}")

def save_to_jsonl(data: dict):
    """data jsonl faile saver"""
    try:
        with open(LOG_FILE_PATH, "a", encoding="utf-8") as f:
            f.write(json.dumps(data, ensure_ascii=False) + "\n")
    except Exception as e:
        print(f"[Log Error] Failed to write to jsonl: {e}", file=sys.stderr)

def evaluate_with_language_model(solution_str: str, ground_truth: str) -> float:
    """
    自作の報酬サーバー (reward_server.py) にリクエストを送り、スコアを取得する。
    """
    
    if '</think>' in solution_str:
        cleaned_solution = solution_str.split('</think>')[-1].strip()
    else:
        cleaned_solution = re.sub(r'<think>.*?</think>', '', solution_str, flags=re.DOTALL).strip()
    
    if '<think>' in cleaned_solution or '</think>' in cleaned_solution:
        print(f"[Reward] <think> tag found in response. Score: -40.0", file=sys.stderr)
        
        save_to_jsonl({
            "timestamp": time.time(),
            "pid": os.getpid(),
            "status": "penalty",
            "score": REWARD_MIN_SCORE,
            "prompt": ground_truth,
            "response_snippet": solution_str[:1000] + "..."
        })
        return REWARD_MIN_SCORE

    payload = {
        "prompt": ground_truth,
        "response": cleaned_solution
    }
    
    try:
        response = requests.post(SCORE_ENDPOINT, json=payload, timeout=30)
        
        if response.status_code != 200:
            print(f"Reward Server Error: {response.status_code} - {response.text}", file=sys.stderr)
            return REWARD_MIN_SCORE
            
        data = response.json()
        raw_score = data.get("score", REWARD_MIN_SCORE)    
        
        score = min(float(raw_score), REWARD_MAX_SCORE)    
        
        print(f"[Reward] Score: {score} (Raw: {raw_score}) | PromptLen: {len(ground_truth)} | ResponseLen: {len(solution_str)}", flush=True)
        #print(f"[Reward] Prompt: {ground_truth}", flush=True)
        #print(f"[Reward] Response: {cleaned_solution}", flush=True)
        
        log_entry = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), 
            "timestamp_epoch": time.time(),
            "pid": os.getpid(), 
            "score": score,
            "raw_score": raw_score,
            "prompt_len": len(ground_truth),
            "response_len": len(solution_str),
            "prompt": ground_truth,
            "response": cleaned_solution
        }
        save_to_jsonl(log_entry)
        
        return score

    except Exception as e:
        print(f"Request Failed: {e}", file=sys.stderr)
        return REWARD_MIN_SCORE


def safe_hack_reward(solution_str: str, ground_truth: str, **kwargs) -> float:
    return evaluate_with_language_model(solution_str, ground_truth)

報酬関数クライアントの仕組み:

このスクリプトは、VeRLの学習プロセスから呼び出され、報酬サーバーにHTTPリクエストを送ってスコアを取得します。

主な処理フロー:

  1. VeRLからの呼び出し

    # VeRLの設定で指定された関数が呼ばれる
    custom_reward_function.path = "safe_hack_reward_request.py"
    custom_reward_function.name = "safe_hack_reward"
    
  2. 前処理

    • solution_str:モデルが生成したテキスト
    • ground_truth:データセットのreward_model.ground_truthフィールド(プロンプトや評価基準)
    • <think>タグの除去(思考過程を報酬計算から除外)
  3. HTTPリクエスト送信

    payload = {
        "prompt": ground_truth,
        "response": cleaned_solution
    }
    response = requests.post(SCORE_ENDPOINT, json=payload, timeout=30)
    
  4. スコア処理

    • サーバーから返されたスコアを取得
    • REWARD_MAX_SCOREでキャップ(上限30.0)
    • エラー時はREWARD_MIN_SCORE(-40.0)を返す
  5. ログ記録

    • スコア、プロンプト、レスポンスをJSONL形式で保存
    • 後から報酬の分布を分析可能

重要な設定値:

# 報酬サーバーのURL(環境変数から取得)
BASE_URL = os.environ.get("REWARD_SERVER_URL", "http://localhost:8000")

# スコアの範囲
REWARD_MAX_SCORE = 30.0   # 最大スコア
REWARD_MIN_SCORE = -40.0  # 最小スコア(ペナルティ)

モデル変換(mcore→HF)

最後に学習後のモデルをmcore形式からHF形式に変換します。
学習前の(HF→mcore)は自動で変換されるため、変換の必要はありません。

#SBATCH --job-name=sft-custom          # ジョブ名
#SBATCH --partition=your_partition      # パーティション名(環境に合わせて変更)
#SBATCH --nodes=1                       # ノード数
#SBATCH --gpus=8                        # GPU数
#SBATCH --output=logs/sft-%j.out        # 標準出力ログ(%jはジョブID)
#SBATCH --error=logs/sft-%j.err         # エラーログ

module load singularitypro
mkdir -p ./.cache
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY

export HOME_PATH="/home/your_username"
export VERL_PATH="${HOME_PATH}/verl"

export RAY_TEMP_DIR="/dev/shm/ray_${SLURM_JOB_ID}"
mkdir -p $RAY_TEMP_DIR

singularity exec --nv \
  --home ${HOME_PATH} \
  -B ${HOME_PATH}:${HOME_PATH} \
  -B ${HOME_PATH}/.cache:/root/.cache \
  -B /dev/shm:/dev/shm \
  --env PYTHONNOUSERSITE=1 \
  --env XDG_CACHE_HOME=${HOME_PATH}/.cache \
  --env HF_HOME=${HOME_PATH}/.cache/huggingface \
  --env TORCH_EXTENSIONS_DIR=${HOME_PATH}/.cache/torch_extensions \
  --env PYTHONPATH=${VERL_PATH}:${HOME_PATH}/Megatron-LM:$PYTHONPATH \
  --env NPROC_PER_NODE=1 \
  --env CUDA_VISIBLE_DEVICES=0 \
  --env NCCL_SOCKET_IFNAME=bond0.1505 \
  --env GLOO_SOCKET_IFNAME=bond0.1505 \
  --env TP_SOCKET_IFNAME=bond0.1505 \
  ${HOME_PATH}/verl_app_v061.sif \
  bash -c "
    ray start --head \
      --num-gpus 1 \
      --object-store-memory 20000000000 \
      --disable-usage-stats \
      --temp-dir=/dev/shm/ray_manual &

    sleep 10
      
    python -m verl.model_merger merge \
      --backend megatron \
      --tie-word-embedding \
      --local_dir /path/to/megatron_output/qwen3_4b_grpo/checkpoints/global_step_50/actor \
      --target_dir /path/to//hf_model/qwen3_4b_grpo/checkpoint-50
    "

まとめ

この記事では、VeRLフレームワークを用いてMegatron-LMベースのLLMに対して強化学習(PPO、GRPO、DAPO)を適用する方法について解説しました。

VeRLは、数十億〜数百億パラメータ規模のMoE(Mixture of Experts)モデルの学習に特に適していると思います。一方で、Qwen3-4Bのような小規模モデルの場合は、MS-SWIFTやUnslothといったフレームワークの方が導入が容易かもしれません。

しかし、大規模なMoEモデルに対してフルパラメータ学習やLoRA学習を行う場合、MS-SWIFTやUnslothでは実装上の制約により対応が難しいくなります。そのため、VeRLまたはNeMo RLが現実的な選択肢となるため、VeRLの使用方法を習得できたことは非常によかったと思います。

Discussion