OpenRLHFモデルアーキテクチャの深層解析
本稿では、OpenRLHFフレームワークにおけるモデル処理のアーキテクチャについて詳細に解説する。強化学習を組み込んだ大規模言語モデルの訓練において、モデルはどのような役割を果たし、互いにどのように連携するかを焦点を当てる。
強化学習フローの概要
OpenRLHFの設計思想を理解する前に、強化学習ベースのLLM訓練の全体像を把握する必要がある：
参考文献:https://arxiv.org/pdf/2405.11143
処理の大まかな流れは以下の通りである。プロンプトがActorモデルにを入力され、応答が生成される。この応答と元のプロンプトを結合し、CriticモデルやRewardモデルによって評価が行われる。
1、アクターモデルの実装
参照:https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/models/actor.py
アクターモデルはユーザーの指示に対して応答を生成する核心的な役割を担う。
class ActorModel(nn.Module):
def __init__(self, model_path: str, config=None):
super().__init__()
if isinstance(model_path, str):
self.language_model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
attn_implementation="flash_attention_2",
quantization_config=quantization_config,
torch_dtype=torch.bfloat16,
device_map="auto",
)
if use_lora:
self.language_model.enable_input_require_grads()
lora_cfg = LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=lora_rank,
lora_alpha=lora_alpha,
target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
lora_dropout=0.1,
bias="none",
)
self.language_model = get_peft_model(self.language_model, lora_cfg)
else:
self.language_model = model_path
@torch.no_grad()
def generate_text(self, input_tokens: torch.Tensor, **generation_config):
generated_sequences = self.language_model.generate(**generation_config)
eos_token = generation_config["eos_token_id"]
pad_token = generation_config["pad_token_id"]
return self.trim_sequences(generated_sequences, input_tokens.size(1), eos_token, pad_token)
def compute_log_probabilities(self, token_sequences: torch.Tensor, num_actions: int):
model_output = self.language_model(token_sequences)
output_logits = model_output["logits"].to(torch.float32)
log_probabilities = self._calculate_log_probs(
output_logits[:, :-1, :],
token_sequences[:, 1:],
temperature=self.temperature
)
action_log_probs = log_probabilities[:, -num_actions:]
return action_log_probs
この実装の핵となる部分は以下の2点である：
generate_textメソッドは、入力されたトークン列に基づいて新しいテキストを生成する。generation_configには温度パラメータ、top_k、top_pなどの生成制御パラメータが含まれる。
compute_log_probabilitiesメソッドは、生成されたシーケンス中の特定トークン（アクション）の対数尤度を計算する。これはPPOやDPOの訓練において極めて重要である。
強化学習では、単に「正しいシーケンスを生成できるか」ではなく、「生成された各トークンがどの程度"良い"ものであるか」を評価する必要がある。DPOでは以下の目的関数が用いられる：
[L(θ) = E[ min(r(θ) * A, clip(r(θ), 1-ε, 1+ε) * A) ] ]
ここで
[r(\theta)=\pi_{\theta}(a|s)/\pi_{old}(a|s) ]
は確率比を表し、コードでは以下のように実装される：
def compute_log_likelihood_from_logits(logits: torch.Tensor, target_tokens: torch.Tensor, temperature: float = 1.0) -> torch.Tensor:
if temperature != 1.0:
logits.div_(temperature)
batch_size = logits.shape[:-1]
vocab_size = logits.shape[-1]
try:
from flash_attn.ops.triton.cross_entropy import cross_entropy_loss
output = cross_entropy_loss(logits.reshape(-1, vocab_size), target_tokens.reshape(-1))
log_probs = -output[0].view(*batch_size)
except ImportError:
token_logits = torch.gather(logits, dim=-1, index=target_tokens.unsqueeze(-1)).squeeze(-1)
log_sum_exp = compute_logsumexp_chunked(logits.reshape(-1, vocab_size))
log_sum_exp = log_sum_exp.view(*batch_size)
log_probs = token_logits - log_sum_exp
return log_probs
補足説明：
AutoModelForCausalLM.from_pretrainedで読み込んだモデルは以下のようなパラメータをサポートする：
model_output = language_model(
input_ids=input_tokens, # 入力トークン (batch_size, seq_length)
attention_mask=attention_mask, # 有効トークン識別（パディング除外）
position_ids=position_ids, # 位置エンコーディング
past_key_values=past_kv, # KVキャッシュ
inputs_embeds=embedded_inputs, # エンベディング直接入力
use_cache=True, # キャッシュ使用フラグ
labels=target_labels, # ラベル指定で自動Loss計算
)
packing処理の実際：
短いシーケンスを複数束ねてメモリ効率を向上させる場合、OpenRLHFではpacking_samplesパラメータで制御される。packingがない場合はattention_maskから位置IDを再計算するが、packing済みの場合は特別な処理が必要：
if not use_packing:
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 1)
else:
if ring_attention_group is not None:
sequences, attention_mask, position_ids = convert_ring_attn_params(
sequences, attention_mask, packed_lengths, ring_attention_group
)
else:
position_ids = recompute_position_ids(attention_mask)
attention_mask = None
2、価値評価モデルの構築
参照:https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/models/model.py
価値評価モデルは 크게2種類用意される：RewardモデルとCriticモデルである。GRPOやDPOの訓練では、生成されたトークンを評価する必要がある。Rewardモデルは生成テキスト全体の品質を数値化し、Criticモデルは優勢関数（advantage）を計算するための価値推定を行う。
def create_reward_model(base_model_class, base_language_model, value_prefix="score", enable_packing=False):
class RewardModel(base_model_class):
def __init__(self, model_config):
super().__init__(model_config)
self.llm = base_language_model(model_config)
self.value_prefix = value_prefix
setattr(self, value_prefix, nn.Linear(model_config.hidden_size, 1, bias=False))
def forward(self, input_ids, attention_mask=None, return_output=False,
ring_attn_group=None, pad_sequence=False, packed_lengths=None):
outputs = self.llm(input_ids, attention_mask=attention_mask, position_ids=position_ids)
hidden_state = outputs["last_hidden_state"]
value_scores = getattr(self, self.value_prefix)(hidden_state).squeeze(-1)
if not enable_packing:
eos_position = attention_mask.size(1) - 1 - attention_mask.long().fliplr().argmax(dim=1, keepdim=True)
final_reward = value_scores.gather(dim=1, index=eos_position).squeeze(1)
else:
final_reward = extract_packed_rewards(value_scores, packed_lengths)
if not self.training and self.normalize_rewards:
final_reward = (final_reward - self.mean) / self.std
return (final_reward, outputs) if return_output else final_reward
return RewardModel
def create_critic_model(base_model_class, base_language_model, value_prefix="score", enable_packing=False):
class CriticModel(base_model_class):
def __init__(self, model_config):
super().__init__(model_config)
self.llm = base_language_model(model_config)
self.value_prefix = value_prefix
setattr(self, value_prefix, nn.Linear(model_config.hidden_size, 1, bias=False))
def forward(self, input_ids, attention_mask=None, num_actions=None,
return_output=False, packed_lengths=None):
outputs = self.llm(input_ids, attention_mask=attention_mask, position_ids=position_ids)
hidden_state = outputs["last_hidden_state"]
value_scores = getattr(self, self.value_prefix)(hidden_state).squeeze(-1)
if num_actions is None:
assert return_output
return outputs
if not enable_packing:
action_values = value_scores[:, -num_actions:]
else:
assert isinstance(num_actions, list) and len(num_actions) == len(packed_lengths)
action_values = []
offset = 0
for action_count, seq_len in zip(num_actions, packed_lengths):
start_idx = max(0, offset + seq_len - action_count - 1)
end_idx = offset + seq_len - 1
action_values.append(value_scores[:, start_idx:end_idx])
offset += seq_len
action_values = torch.cat(action_values, dim=1)
return (action_values, outputs) if return_output else action_values
return CriticModel
Rewardモデルは、入力されたテキストに対して単一のスコアを付与する。last_hidden_stateから各トークンのスコアを計算し、文章の終わり（EOSトークン位置）のスコアを全体評価として使用する。
Criticモデルは、Actorが生成したアクション（最後のnum_actions個のトークン）に対応する価値推定を出力する。これはPPOにおける優勢関数𝐴(𝑠,𝑎)=𝑄(𝑠,𝑎)−𝑉(𝑠)の計算に直接用いられる。
具体例で理解を深める。プロンプトが「日本の首都は」で、Actorが「東京です」と生成した場合：
結合シーケンス：「日本の首都は東京です」（各トークンのスコアは[0.1, 0.2, 0.3, 0.5, 0.9]など）
RewardモデルはEOSトークン位置のスコア（0.9）を全体の奖励として出力する。Criticモデルは最後の3トークン（ Actions相当）の価値を出力し、これが政策更新の基準となる。
3、損失関数の実装
参照:https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/models/loss.py
補足：
torch.clampは値を指定範囲内に収める関数であり、範囲外の値は境界値に置き換えられる
PolicyLoss（PPO方策損失）：
[\begin{align*} r_t &= \exp(\log \pi(a_t \mid s_t) - \log \pi_{\text{old}}(a_t \mid s_t)) \ \mathcal{L}{\text{clip}}(t) &= \min\left(r_t \cdot A_t,\ \text{clip}(r_t,\ 1 - \epsilon,\ 1 + \epsilon) \cdot A_t\right) \ \mathcal{L}{\text{policy}} &= -\mathbb{E}t \left[ \mathcal{L}{\text{clip}}(t) \right] \end{align*} ]
ValueLoss（価値関数損失）：
[\mathcal{L}{\text{value}} = \frac{1}{2} \cdot \mathbb{E}{t \sim \text{mask}} \left[ \max \left( (V_{\text{clip}, t} - R_t)^2, , (V_t - R_t)^2 \right) \right]\ \text{ただし：}V_{\text{clip}} = V_{\text{old}} + \text{clip}(V - V_{\text{old}}, -\epsilon, \epsilon) ]
まとめ
本稿ではOpenRLHFのモデルアーキテクチャを3つの観点から解説した：
-
ActorModelはプロンプトを受け取り応答を生成する。生成されたトークンの対数尤度を計算し、訓練時のスコアリングに使用される。
-
CriticModelは生成されたアクション（最後の数トークン）の価値を出力する。これによりPPOにおける優勢関数の計算が可能となる。
-
RewardModelはシーケンス全体の品質を単一のスカラ値で評価する。EOSトークン位置のスコアを採用することで文章全体を評価する。
これらのモデルは相互に連携し、強化学習に基づく大規模言語モデルの効率的な訓練を支える。