This commit is contained in:
wehub-resource-sync
@@ -0,0 +1 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
@@ -0,0 +1,339 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from swift.utils import nanstd
|
||||
|
||||
|
||||
def compute_advantages(
|
||||
rewards_per_func: torch.Tensor,
|
||||
reward_weights: torch.Tensor,
|
||||
num_generations: int,
|
||||
advantage_estimator: str = 'grpo',
|
||||
scale_rewards: str = 'group',
|
||||
kl_in_reward: bool = False,
|
||||
beta: float = 0.0,
|
||||
kl_values: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Compute advantages from per-function rewards.
|
||||
|
||||
This is a pure tensor function suitable for all backends (HF, Megatron, Ray).
|
||||
Input tensors should already be gathered across all processes.
|
||||
|
||||
Produces the **per-sequence** base advantage from rewards. The OPD-RL teacher signal
|
||||
is *not* injected here: it is a per-token signal applied later (when the base
|
||||
advantage is broadcast to ``[B, T]`` while writing it onto the batch). See
|
||||
:func:`compute_teacher_logratio` (the advantage signal) and
|
||||
:func:`compute_teacher_kl_per_token` (the k3 monitoring metric).
|
||||
|
||||
Ref model KL (``kl_in_reward``) is subtracted from rewards **before** advantage
|
||||
normalization — standard GRPO/PPO regularization that prevents the policy from
|
||||
drifting too far from the reference model.
|
||||
|
||||
Args:
|
||||
rewards_per_func: ``[N, n_funcs]`` per-function reward matrix.
|
||||
reward_weights: ``[n_funcs]`` weighting tensor.
|
||||
num_generations: ``K`` — completions per prompt.
|
||||
advantage_estimator: ``'grpo'``, ``'rloo'``, or ``'reinforce_plus_plus'``.
|
||||
scale_rewards: ``'batch'``, ``'group'``, ``'none'``, or ``'gdpo'``.
|
||||
kl_in_reward: Subtract ref model KL from rewards (pre-normalization).
|
||||
beta: Ref model KL penalty coefficient.
|
||||
kl_values: ``[N]`` ref model KL values (required when ``kl_in_reward=True``).
|
||||
|
||||
Returns:
|
||||
``(advantages, rewards)`` both ``[N]``.
|
||||
"""
|
||||
rewards = (rewards_per_func * reward_weights.unsqueeze(0)).nansum(dim=1)
|
||||
|
||||
if kl_in_reward and beta != 0.0 and kl_values is not None:
|
||||
rewards = rewards - beta * kl_values
|
||||
|
||||
K = num_generations
|
||||
grouped = rewards.view(-1, K)
|
||||
group_mean = grouped.mean(dim=1).repeat_interleave(K)
|
||||
|
||||
if advantage_estimator == 'rloo' and K > 1:
|
||||
advantages = rewards * K / (K - 1) - group_mean * K / (K - 1)
|
||||
else:
|
||||
advantages = rewards - group_mean
|
||||
|
||||
std: Optional[torch.Tensor] = None
|
||||
if advantage_estimator == 'reinforce_plus_plus':
|
||||
if scale_rewards == 'batch':
|
||||
std = advantages.std().expand_as(advantages) if advantages.numel() > 1 else torch.zeros_like(advantages)
|
||||
elif scale_rewards == 'group':
|
||||
std = (advantages.view(-1, K).std(dim=1).repeat_interleave(K) if K > 1 else torch.zeros_like(advantages))
|
||||
else:
|
||||
if scale_rewards == 'batch':
|
||||
std = rewards.std().expand_as(rewards) if rewards.numel() > 1 else torch.zeros_like(rewards)
|
||||
elif scale_rewards == 'group':
|
||||
std = grouped.std(dim=1).repeat_interleave(K) if K > 1 else torch.zeros_like(rewards)
|
||||
elif scale_rewards == 'gdpo':
|
||||
n_funcs = rewards_per_func.shape[1]
|
||||
normalized_list = []
|
||||
for i in range(n_funcs):
|
||||
r_i = rewards_per_func[:, i].view(-1, K)
|
||||
g_mean = torch.nanmean(r_i, dim=1, keepdim=True)
|
||||
g_std = nanstd(r_i, dim=1, keepdim=True) + 1e-8
|
||||
norm_i = torch.nan_to_num((r_i - g_mean) / g_std, nan=0.0)
|
||||
normalized_list.append(reward_weights[i] * norm_i.view(-1))
|
||||
summed = sum(normalized_list)
|
||||
advantages = (summed - summed.mean()) / (summed.std() + 1e-8)
|
||||
std = None
|
||||
|
||||
if std is not None and scale_rewards != 'none':
|
||||
advantages = advantages / (std + 1e-4)
|
||||
|
||||
return advantages, rewards
|
||||
|
||||
|
||||
def compute_advantages_dynamic(
|
||||
rewards_per_func: torch.Tensor,
|
||||
reward_weights: torch.Tensor,
|
||||
prompt_ids: List[str],
|
||||
request_ids: List[str],
|
||||
advantage_estimator: str = 'grpo',
|
||||
scale_rewards: str = 'group',
|
||||
kl_in_reward: bool = False,
|
||||
beta: float = 0.0,
|
||||
kl_values: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Request-aware advantage computation for dynamic sample counts.
|
||||
|
||||
Groups rewards by ``prompt_id`` (via ``request_id`` deduplication) and
|
||||
computes advantages within each group. Supports variable numbers of
|
||||
completions per prompt (multi-turn scenarios).
|
||||
|
||||
Like :func:`compute_advantages`, this returns only the per-sequence base advantage;
|
||||
the OPD-RL teacher signal is applied per-token later (see
|
||||
:func:`compute_teacher_logratio`).
|
||||
|
||||
Input tensors should already be gathered across all processes.
|
||||
|
||||
Args:
|
||||
rewards_per_func: ``[N, n_funcs]`` per-function reward matrix.
|
||||
reward_weights: ``[n_funcs]`` weighting tensor.
|
||||
prompt_ids: Length-N list of prompt identifiers.
|
||||
request_ids: Length-N list of request identifiers (may have duplicates).
|
||||
advantage_estimator: ``'grpo'``, ``'rloo'``, or ``'reinforce_plus_plus'``.
|
||||
scale_rewards: ``'batch'``, ``'group'``, ``'none'``.
|
||||
kl_in_reward: Subtract ref model KL from rewards (pre-normalization).
|
||||
beta: Ref model KL penalty coefficient.
|
||||
kl_values: ``[N]`` ref model KL values.
|
||||
|
||||
Returns:
|
||||
``(advantages, rewards)`` both ``[N]`` (with duplicate entries for repeated request_ids).
|
||||
"""
|
||||
device = rewards_per_func.device
|
||||
rewards = (rewards_per_func * reward_weights.unsqueeze(0)).nansum(dim=1)
|
||||
|
||||
if kl_in_reward and beta != 0.0 and kl_values is not None:
|
||||
rewards = rewards - beta * kl_values
|
||||
|
||||
# Deduplicate by request_id (keep last occurrence)
|
||||
seen = {}
|
||||
for idx, rid in enumerate(request_ids):
|
||||
seen[rid] = idx
|
||||
unique_indices = torch.tensor(sorted(seen.values()), device=device)
|
||||
unique_request_ids = [request_ids[i] for i in unique_indices.cpu()]
|
||||
unique_prompt_ids = [prompt_ids[i] for i in unique_indices.cpu()]
|
||||
unique_rewards = rewards[unique_indices]
|
||||
|
||||
# Group by prompt_id
|
||||
prompt_to_indices: Dict[str, List[int]] = {}
|
||||
for idx, pid in enumerate(unique_prompt_ids):
|
||||
prompt_to_indices.setdefault(pid, []).append(idx)
|
||||
|
||||
prompt_means = torch.zeros(len(unique_rewards), device=device)
|
||||
for pid, idxs in prompt_to_indices.items():
|
||||
idx_t = torch.tensor(idxs, device=device)
|
||||
prompt_means[idx_t] = unique_rewards[idx_t].mean()
|
||||
|
||||
if advantage_estimator == 'rloo':
|
||||
request_advantages = torch.zeros_like(unique_rewards)
|
||||
for pid, idxs in prompt_to_indices.items():
|
||||
K = len(idxs)
|
||||
idx_t = torch.tensor(idxs, device=device)
|
||||
r_group = unique_rewards[idx_t]
|
||||
if K > 1:
|
||||
request_advantages[idx_t] = r_group * K / (K - 1) - r_group.mean() * K / (K - 1)
|
||||
else:
|
||||
request_advantages[idx_t] = r_group - r_group.mean()
|
||||
else:
|
||||
request_advantages = unique_rewards - prompt_means
|
||||
|
||||
# Normalize
|
||||
if advantage_estimator == 'reinforce_plus_plus':
|
||||
if scale_rewards == 'batch':
|
||||
adv_std = (request_advantages.std() if request_advantages.numel() > 1 else torch.tensor(0.0, device=device))
|
||||
prompt_stds = torch.full_like(request_advantages, adv_std)
|
||||
elif scale_rewards == 'group':
|
||||
prompt_stds = torch.zeros(len(unique_rewards), device=device)
|
||||
for pid, idxs in prompt_to_indices.items():
|
||||
idx_t = torch.tensor(idxs, device=device)
|
||||
adv_group = request_advantages[idx_t]
|
||||
prompt_stds[idx_t] = adv_group.std() if len(idxs) > 1 else 0.0
|
||||
else:
|
||||
prompt_stds = None
|
||||
else:
|
||||
if scale_rewards == 'batch':
|
||||
r_std = unique_rewards.std() if unique_rewards.numel() > 1 else torch.tensor(0.0, device=device)
|
||||
prompt_stds = torch.full_like(unique_rewards, r_std)
|
||||
elif scale_rewards == 'group':
|
||||
prompt_stds = torch.zeros(len(unique_rewards), device=device)
|
||||
for pid, idxs in prompt_to_indices.items():
|
||||
idx_t = torch.tensor(idxs, device=device)
|
||||
r_group = unique_rewards[idx_t]
|
||||
prompt_stds[idx_t] = r_group.std() if len(idxs) > 1 else 0.0
|
||||
else:
|
||||
prompt_stds = None
|
||||
|
||||
if prompt_stds is not None and scale_rewards != 'none':
|
||||
request_advantages = request_advantages / (prompt_stds + 1e-4)
|
||||
|
||||
# Map back to original order
|
||||
rid_to_idx = {rid: idx for idx, rid in enumerate(unique_request_ids)}
|
||||
indices_in_unique = torch.tensor([rid_to_idx[r] for r in request_ids], device=device)
|
||||
advantages = request_advantages[indices_in_unique]
|
||||
|
||||
return advantages, rewards
|
||||
|
||||
|
||||
def compute_teacher_kl_per_token(
|
||||
teacher_per_token_logps: torch.Tensor,
|
||||
policy_per_token_logps: torch.Tensor,
|
||||
completion_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Per-token teacher KL (OPD-RL) via the k3 estimator -- **monitoring only**.
|
||||
|
||||
``teacher`` and ``policy`` logps are token-in-token-out on the *same* sampled tokens
|
||||
(the teacher logp on the student-sampled token).
|
||||
It is the magnitude of the reverse KL between student and teacher and a good "how far
|
||||
is the student from the teacher" gauge -- it should *decrease* over training.
|
||||
|
||||
Args:
|
||||
teacher_per_token_logps: ``[B, T]`` teacher logp on sampled tokens.
|
||||
policy_per_token_logps: ``[B, T]`` student (old) logp on the same tokens.
|
||||
completion_mask: ``[B, T]`` response-token mask.
|
||||
|
||||
Returns:
|
||||
``[B, T]`` per-token teacher KL (masked outside the response).
|
||||
"""
|
||||
d = teacher_per_token_logps - policy_per_token_logps
|
||||
per_token = torch.exp(d) - d - 1
|
||||
return per_token * completion_mask
|
||||
|
||||
|
||||
def compute_teacher_logratio(
|
||||
teacher_per_token_logps: torch.Tensor,
|
||||
policy_per_token_logps: torch.Tensor,
|
||||
completion_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Per-token signed teacher log-ratio (OPD-RL) -- the k1 reverse-KL estimator.
|
||||
|
||||
This is the correct OPD-RL policy-gradient signal (PG OPD): the negative single-sample
|
||||
reverse-KL estimate used as a reward, ``r_t = teacher_logp(y_t) - student_logp(y_t)``
|
||||
|
||||
Args:
|
||||
teacher_per_token_logps: ``[B, T]`` teacher logp on sampled tokens.
|
||||
policy_per_token_logps: ``[B, T]`` student (old) logp on the same tokens.
|
||||
completion_mask: ``[B, T]`` response-token mask.
|
||||
|
||||
Returns:
|
||||
``[B, T]`` per-token signed log-ratio (masked outside the response).
|
||||
"""
|
||||
d = teacher_per_token_logps - policy_per_token_logps
|
||||
return d * completion_mask
|
||||
|
||||
|
||||
def expand_advantage_to_per_token(
|
||||
advantages: torch.Tensor,
|
||||
completion_mask: torch.Tensor,
|
||||
teacher_per_token_logps: Optional[torch.Tensor] = None,
|
||||
policy_per_token_logps: Optional[torch.Tensor] = None,
|
||||
teacher_kl_coef: float = 0.0,
|
||||
) -> torch.Tensor:
|
||||
"""Expand the per-sequence base advantage ``[B]`` to per-token ``[B, T]``.
|
||||
|
||||
Broadcasting the per-sequence advantage to per-token happens *here* (at batch
|
||||
construction) rather than in the loss, so the OPD-RL teacher signal can be added
|
||||
per token: ``adv_t = base_adv + coef * (teacher_logp_t - student_logp_t)`` .
|
||||
Without a teacher this is a plain broadcast.
|
||||
|
||||
Args:
|
||||
advantages: ``[B]`` per-sequence base advantage.
|
||||
completion_mask: ``[B, T]`` response-token mask (defines the token frame).
|
||||
teacher_per_token_logps: ``[B, T]`` teacher logp on sampled tokens (OPD-RL).
|
||||
policy_per_token_logps: ``[B, T]`` student (old) logp on the same tokens.
|
||||
teacher_kl_coef: Coefficient for the per-token teacher signal.
|
||||
|
||||
Returns:
|
||||
``[B, T]`` per-token advantage.
|
||||
"""
|
||||
per_token_adv = advantages.unsqueeze(1).expand_as(completion_mask).clone()
|
||||
if teacher_per_token_logps is not None and teacher_kl_coef != 0.0:
|
||||
signed = compute_teacher_logratio(teacher_per_token_logps, policy_per_token_logps, completion_mask)
|
||||
per_token_adv = per_token_adv + teacher_kl_coef * signed
|
||||
return per_token_adv
|
||||
|
||||
|
||||
@dataclass
|
||||
class RewardMetrics:
|
||||
"""Reward statistics for logging."""
|
||||
reward_mean: float
|
||||
reward_std: float
|
||||
frac_reward_zero_std: float
|
||||
per_func_mean: Dict[str, float]
|
||||
per_func_std: Dict[str, float]
|
||||
|
||||
|
||||
def compute_reward_metrics(
|
||||
rewards: torch.Tensor,
|
||||
rewards_per_func: torch.Tensor,
|
||||
reward_func_names: List[str],
|
||||
num_generations: int,
|
||||
scale_rewards: str,
|
||||
) -> RewardMetrics:
|
||||
"""Compute reward statistics for monitoring.
|
||||
Args:
|
||||
rewards: ``[N]`` scalar rewards (already weighted).
|
||||
rewards_per_func: ``[N, n_funcs]`` per-function rewards.
|
||||
reward_func_names: Names of reward functions.
|
||||
num_generations: ``K``.
|
||||
scale_rewards: Scaling strategy (affects std computation).
|
||||
Returns:
|
||||
:class:`RewardMetrics` with all statistics.
|
||||
"""
|
||||
group_rewards = rewards.view(-1, num_generations)
|
||||
reward_mean = group_rewards.mean(-1).mean().item()
|
||||
|
||||
if scale_rewards in ('group', 'none', 'gdpo'):
|
||||
reward_std = group_rewards.std(-1).mean().item() if num_generations > 1 else 0.0
|
||||
elif scale_rewards == 'batch':
|
||||
reward_std = rewards.std().item() if rewards.numel() > 1 else 0.0
|
||||
else:
|
||||
reward_std = 0.0
|
||||
|
||||
if num_generations > 1:
|
||||
is_std_zero = torch.isclose(group_rewards.std(dim=1), torch.zeros_like(group_rewards.std(dim=1)))
|
||||
else:
|
||||
is_std_zero = torch.ones(group_rewards.size(0), dtype=torch.bool, device=group_rewards.device)
|
||||
frac_zero_std = is_std_zero.float().mean().item()
|
||||
|
||||
per_func_mean = {}
|
||||
per_func_std = {}
|
||||
for i, name in enumerate(reward_func_names):
|
||||
col = rewards_per_func[:, i]
|
||||
per_func_mean[name] = torch.nanmean(col).item()
|
||||
valid = col[~torch.isnan(col)]
|
||||
per_func_std[name] = valid.std().item() if valid.numel() > 1 else 0.0
|
||||
|
||||
return RewardMetrics(
|
||||
reward_mean=reward_mean,
|
||||
reward_std=reward_std,
|
||||
frac_reward_zero_std=frac_zero_std,
|
||||
per_func_mean=per_func_mean,
|
||||
per_func_std=per_func_std,
|
||||
)
|
||||
@@ -0,0 +1,353 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import torch
|
||||
from dacite import from_dict
|
||||
from dataclasses import dataclass, field, fields
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
from swift.infer_engine.protocol import RolloutOutput
|
||||
from swift.template import Messages
|
||||
from swift.utils import get_logger, remove_response
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from swift.infer_engine.protocol import RolloutInferRequest
|
||||
from swift.rlhf_trainers.gkd_loss import DataSource
|
||||
|
||||
from swift.dataset.preprocessor.core import _pair_keys as StandardKeys
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
# Multimodal keys that a scheduler may override via ``rollout_infos``.
|
||||
_MULTIMODAL_KEYS = ('images', 'videos', 'audios')
|
||||
|
||||
|
||||
@dataclass
|
||||
class OnPolicySample:
|
||||
"""A single on-policy rollout trajectory (pre-collation, per-sample).
|
||||
|
||||
Lifecycle of one sample::
|
||||
|
||||
1. dataset row -> messages + extra
|
||||
2. rollout -> response_token_ids / rollout_logprobs / finish_reason
|
||||
3. rebuild messages -> replace_assistant_response_with_ids
|
||||
4. encode -> encoded = template.encode(self) (per-sample dict)
|
||||
|
||||
The collated model-forward inputs (input_ids[B,T], ...) are produced later
|
||||
at batch level from ``[s.encoded for s in samples]`` by
|
||||
``collate_to_grpo_micro_batch`` (returns ``(model_inputs, grpo_batch)``), never on
|
||||
the sample.
|
||||
"""
|
||||
# --- standard keys ---
|
||||
messages: List[Dict]
|
||||
images: List[Any] = field(default_factory=list)
|
||||
videos: List[Any] = field(default_factory=list)
|
||||
audios: List[Any] = field(default_factory=list)
|
||||
tools: Optional[List[Any]] = None
|
||||
objects: Dict[str, Any] = field(default_factory=dict)
|
||||
extra: Dict[str, Any] = field(default_factory=dict) # dataset passthrough columns (flattened for reward)
|
||||
|
||||
# --- id ---
|
||||
prompt_id: str = ''
|
||||
request_id: str = ''
|
||||
|
||||
# --- rollout output ---
|
||||
response_token_ids: List[List[int]] = field(default_factory=list)
|
||||
response_loss_mask: List[List[int]] = field(default_factory=list)
|
||||
rollout_logprobs: List[List[float]] = field(default_factory=list)
|
||||
finish_reason: Optional[str] = None
|
||||
add_eos: bool = False
|
||||
rollout_infos: Dict[str, Any] = field(default_factory=dict)
|
||||
routed_experts: Optional[Any] = None # R3 router replay (per-sample, pre-collation)
|
||||
|
||||
# --- per-sample template.encode output, used for model forward kwargs ---
|
||||
encoded: Optional[Dict[str, Any]] = None
|
||||
|
||||
# --- OPSD (On-Policy Self-Distillation), shared by GKD and OPD-RL ---
|
||||
teacher_prompt: Optional[Any] = None # OPSD: dataset ``teacher_prompt`` column (pre-collation)
|
||||
teacher_messages: Optional[Messages] = None # OPSD: messages with teacher_prompt replacing last user
|
||||
|
||||
@property
|
||||
def is_truncated(self) -> bool:
|
||||
return self.finish_reason == 'length'
|
||||
|
||||
def get_tag(self, tag_key: str = 'dataset') -> Optional[str]:
|
||||
"""Return the multi-teacher routing tag from ``extra[tag_key]`` (``None`` if unset).
|
||||
|
||||
Single source of truth for where a sample's routing tag lives (``--teacher_tag_key``,
|
||||
default ``dataset``); routing keys off exactly this, with no fallback to other columns.
|
||||
"""
|
||||
val = self.extra.get(tag_key)
|
||||
return str(val) if val is not None else None
|
||||
|
||||
def build_teacher_view(self) -> bool:
|
||||
"""Populate the OPSD teacher view from ``teacher_prompt`` + the on-policy response.
|
||||
|
||||
OPSD scores the teacher on its OWN (teacher_prompt + same on-policy response)
|
||||
sequence: replace the last user message with ``teacher_prompt`` and keep the
|
||||
assistant response. Teacher and student share ``response_token_ids`` (identical
|
||||
response tokens, only the prompt differs). Returns ``True`` when an OPSD view
|
||||
exists (idempotent) and ``False`` when ``teacher_prompt`` is unset (non-OPSD).
|
||||
"""
|
||||
if self.teacher_messages is not None:
|
||||
return True
|
||||
|
||||
if not self.teacher_prompt:
|
||||
return False
|
||||
|
||||
messages = [dict(m) for m in self.messages]
|
||||
for msg in reversed(messages):
|
||||
if msg['role'] == 'user':
|
||||
msg['content'] = self.teacher_prompt
|
||||
break
|
||||
|
||||
self.teacher_messages = messages
|
||||
return True
|
||||
|
||||
def to_teacher_template_dict(self) -> Dict[str, Any]:
|
||||
"""Reconstruct the teacher-side dict consumed by ``template.encode`` (OPSD).
|
||||
|
||||
Uses ``teacher_messages`` (teacher_prompt-replaced) + the shared
|
||||
``response_token_ids`` (same on-policy response as the student).
|
||||
"""
|
||||
d = self._standard_fields()
|
||||
d['messages'] = self.teacher_messages
|
||||
chat_template_kwargs = self.extra.get('chat_template_kwargs')
|
||||
if chat_template_kwargs is not None:
|
||||
d['chat_template_kwargs'] = chat_template_kwargs
|
||||
if self.response_token_ids:
|
||||
d['response_token_ids'] = self.response_token_ids
|
||||
d['add_eos'] = False
|
||||
return d
|
||||
|
||||
def _standard_fields(self) -> Dict[str, Any]:
|
||||
"""Non-empty StandardKeys fields from this sample (replaces _multimodal_columns)."""
|
||||
return {k: getattr(self, k) for k in StandardKeys if getattr(self, k)}
|
||||
|
||||
def to_reward_row(self) -> Dict[str, Any]:
|
||||
"""Build the dict consumed by reward functions.
|
||||
|
||||
Flattens ``extra`` to top level so dataset columns (``solution``,
|
||||
``target``, ...) remain accessible as keyword args after
|
||||
``RowPreprocessor.rows_to_batched``. ``encoded`` is excluded (heavy,
|
||||
model-internal). Keeps the legacy column contract intact so existing
|
||||
reward functions need no change.
|
||||
"""
|
||||
row: Dict[str, Any] = {
|
||||
'messages': self.messages,
|
||||
'prompt_id': self.prompt_id,
|
||||
'request_id': self.request_id,
|
||||
'finish_reason': self.finish_reason,
|
||||
'is_truncated': self.is_truncated,
|
||||
'rollout_infos': self.rollout_infos,
|
||||
'response_token_ids': self.response_token_ids,
|
||||
}
|
||||
row.update(self._standard_fields())
|
||||
row.update(self.extra)
|
||||
return row
|
||||
|
||||
@classmethod
|
||||
def from_row(cls, row: Dict[str, Any]) -> OnPolicySample:
|
||||
"""Build a sample from a dataloader / rollout dict row.
|
||||
|
||||
Known keys are mapped to explicit dataclass fields via field-name
|
||||
introspection (no hand-maintained key list); every other column
|
||||
(dataset passthrough like ``solution`` / ``chat_template_kwargs``) goes
|
||||
to ``extra``. ``is_truncated`` is a derived property and is dropped.
|
||||
|
||||
``row`` values are deep-copied so the sample owns its data: dataloaders
|
||||
(e.g. RepeatSampler with steps_per_generation) cache and re-yield the same
|
||||
row dict, and the rollout/encode pipeline mutates messages in place
|
||||
(remove_response / response injection). Sharing references would corrupt
|
||||
the cached dataset rows across micro-steps.
|
||||
"""
|
||||
field_names = {f.name for f in fields(cls)} - {'extra', 'encoded'}
|
||||
standard: Dict[str, Any] = {}
|
||||
extra: Dict[str, Any] = {}
|
||||
for key, value in row.items():
|
||||
if key in field_names:
|
||||
standard[key] = copy.deepcopy(value)
|
||||
elif key == 'is_truncated':
|
||||
continue # derived from finish_reason
|
||||
else:
|
||||
extra[key] = copy.deepcopy(value)
|
||||
return cls(extra=extra, **standard)
|
||||
|
||||
def to_template_dict(self) -> Dict[str, Any]:
|
||||
"""Reconstruct the dict consumed by ``template.encode``.
|
||||
|
||||
messages + StandardKeys (images/videos/audios/tools/objects) +
|
||||
``chat_template_kwargs`` (the only dataset-passthrough column encode
|
||||
consumes — drives enable_thinking / max_pixels / reasoning_effort) +
|
||||
add_eos. Other ``extra`` columns (solution/target/...) are reward-only
|
||||
and intentionally excluded from encode. Response tokens are already
|
||||
injected into ``messages`` via ``replace_assistant_response_with_ids``
|
||||
before encoding.
|
||||
"""
|
||||
d = self._standard_fields()
|
||||
chat_template_kwargs = self.extra.get('chat_template_kwargs')
|
||||
if chat_template_kwargs is not None:
|
||||
d['chat_template_kwargs'] = chat_template_kwargs
|
||||
d['add_eos'] = self.add_eos
|
||||
return d
|
||||
|
||||
def apply_rollout_output(self, *, rollout_output: RolloutOutput) -> None:
|
||||
"""Merge one rollout output back onto this sample (used post-rollout).
|
||||
|
||||
Single- vs multi-turn is inferred from ``rollout_output`` itself, not
|
||||
from external scheduler state: a multi-turn scheduler (colocate *or*
|
||||
server) always returns the full ``messages`` history, while single-turn
|
||||
leaves ``messages`` empty. ``choice.token_ids`` is the last single-shot
|
||||
generation, so it is only a safe fallback in the single-turn case.
|
||||
"""
|
||||
choice = rollout_output.response.choices[0]
|
||||
return_by_scheduler = rollout_output.messages is not None
|
||||
|
||||
# messages: multi-turn returns full history; single-turn appends response
|
||||
if return_by_scheduler:
|
||||
self.messages = rollout_output.messages
|
||||
else:
|
||||
remove_response(self.messages)
|
||||
self.messages.append({'role': 'assistant', 'content': choice.message.content})
|
||||
|
||||
# response token ids: prefer explicit TITO / scheduler output; only fall
|
||||
# back to single-shot choice.token_ids in genuine single-turn rollout.
|
||||
if rollout_output.response_token_ids:
|
||||
self.response_token_ids = rollout_output.response_token_ids
|
||||
if rollout_output.response_loss_mask:
|
||||
self.response_loss_mask = rollout_output.response_loss_mask
|
||||
elif not return_by_scheduler and choice.token_ids:
|
||||
self.response_token_ids = choice.token_ids
|
||||
|
||||
# rollout logprobs (for importance sampling); keep nested [turn][token]
|
||||
if rollout_output.rollout_logprobs:
|
||||
self.rollout_logprobs = rollout_output.rollout_logprobs
|
||||
elif choice.logprobs is not None and 'content' in choice.logprobs:
|
||||
self.rollout_logprobs = [[item['logprob'] for item in choice.logprobs['content']]]
|
||||
|
||||
self.finish_reason = choice.finish_reason
|
||||
self.add_eos = False
|
||||
|
||||
# R3 router replay: transfer routed_experts from the rollout choice
|
||||
routed_experts = getattr(choice, 'routed_experts', None)
|
||||
if routed_experts is not None:
|
||||
self.routed_experts = routed_experts
|
||||
|
||||
# rollout_infos may carry scheduler-overridden multi-modal data
|
||||
if rollout_output.rollout_infos:
|
||||
self.rollout_infos = rollout_output.rollout_infos
|
||||
for key in _MULTIMODAL_KEYS:
|
||||
if key in rollout_output.rollout_infos:
|
||||
setattr(self, key, rollout_output.rollout_infos[key])
|
||||
|
||||
def to_infer_request(self, include_extra: bool = False) -> RolloutInferRequest:
|
||||
"""Build the ``RolloutInferRequest`` consumed by the rollout engine.
|
||||
|
||||
Maps messages + multimodal/standard columns (images/videos/audios/
|
||||
tools/objects) + ``uuid`` (defaults to ``request_id``). Images given as
|
||||
``{'bytes': ...}`` / ``{'path': ...}`` dicts are normalized to base64 /
|
||||
path strings. ``tools`` given as a JSON string is parsed.
|
||||
|
||||
When ``include_extra`` is True, dataset passthrough columns (``extra``)
|
||||
are forwarded via ``data_dict`` (used for server mode / multi-turn).
|
||||
"""
|
||||
import base64
|
||||
import json
|
||||
|
||||
from swift.infer_engine.protocol import RolloutInferRequest
|
||||
|
||||
def _process_image_data(image_data):
|
||||
if isinstance(image_data, dict):
|
||||
if image_data.get('bytes'):
|
||||
return base64.b64encode(image_data['bytes']).decode('utf-8')
|
||||
if image_data.get('path'):
|
||||
return image_data['path']
|
||||
return image_data
|
||||
|
||||
request_data: Dict[str, Any] = {'uuid': self.request_id}
|
||||
request_data.update(self._standard_fields())
|
||||
|
||||
chat_template_kwargs = self.extra.get('chat_template_kwargs')
|
||||
if chat_template_kwargs:
|
||||
request_data['chat_template_kwargs'] = chat_template_kwargs
|
||||
|
||||
if request_data.get('images'):
|
||||
imgs = request_data['images']
|
||||
if not isinstance(imgs, list):
|
||||
imgs = [imgs]
|
||||
request_data['images'] = [_process_image_data(img) for img in imgs]
|
||||
|
||||
if isinstance(request_data.get('tools'), str):
|
||||
try:
|
||||
request_data['tools'] = json.loads(request_data['tools'])
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
if include_extra and self.extra:
|
||||
base_data_dict = self.extra.get('data_dict')
|
||||
if base_data_dict is not None and not isinstance(base_data_dict, dict):
|
||||
raise ValueError('data_dict exists but is not a dictionary')
|
||||
extra_data = {k: v for k, v in self.extra.items() if k != 'data_dict' and v is not None}
|
||||
request_data['data_dict'] = {**extra_data, **(base_data_dict or {})}
|
||||
|
||||
return from_dict(RolloutInferRequest, request_data)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GRPOSample(OnPolicySample):
|
||||
"""On-policy sample with GRPO reward/advantage signals."""
|
||||
rewards: Optional[List[Optional[float]]] = None # optional mirror; main path uses rewards_per_func tensor
|
||||
advantages: Optional[torch.Tensor] = None # filled after _compute_advantages (0-dim tensor per sample)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GRPOBatch:
|
||||
"""Batch data for GRPO loss computation (post-collation, batch-level).
|
||||
|
||||
1. ``completion_mask``, ``truncated_mask``, ``seq_lengths`` — derived from
|
||||
collated ``labels`` right after ``data_collator``.
|
||||
2. ``old_per_token_logps``, ``ref_per_token_logps`` — computed via
|
||||
model/ref forward on the collated batch.
|
||||
3. ``advantages`` — computed from gathered rewards.
|
||||
4. ``rollout_per_token_logps``, ``num_items_in_batch`` — optional,
|
||||
filled when rollout IS / DAPO is enabled.
|
||||
"""
|
||||
completion_mask: torch.Tensor # [B, T]
|
||||
truncated_mask: torch.Tensor # [B]
|
||||
seq_lengths: torch.Tensor # [B] or [B+n] for padding_free
|
||||
|
||||
old_per_token_logps: Optional[torch.Tensor] = None # [B, T]
|
||||
ref_per_token_logps: Optional[torch.Tensor] = None # [B, T]
|
||||
rollout_per_token_logps: Optional[torch.Tensor] = None # [B, T]
|
||||
teacher_per_token_logps: Optional[torch.Tensor] = None # [B, T], OPD-RL teacher logp on sampled tokens
|
||||
advantages: Optional[torch.Tensor] = None # [B, T] per-token (base broadcast minus per-token teacher KL)
|
||||
num_items_in_batch: Optional[torch.Tensor] = None # scalar
|
||||
logits_to_keep: Optional[int] = None
|
||||
|
||||
def to_device(self, device) -> 'GRPOBatch':
|
||||
"""Move all tensor fields to ``device`` in place (Ray: collated on the CPU
|
||||
driver, moved to the GPU worker before forward). Returns self."""
|
||||
for f in fields(self):
|
||||
v = getattr(self, f.name)
|
||||
if isinstance(v, torch.Tensor):
|
||||
setattr(self, f.name, v.to(device))
|
||||
return self
|
||||
|
||||
|
||||
@dataclass
|
||||
class GKDSample(OnPolicySample):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class GKDBatch:
|
||||
"""Batch-level GKD signals (post-collation), symmetric to :class:`GRPOBatch`.
|
||||
|
||||
- ``data_source``: STUDENT / TEACHER / DATASET for this micro-batch.
|
||||
- ``teacher_topk_logprobs`` / ``teacher_topk_indices``: assembled teacher
|
||||
top-k (teacher-API mode), batch tensors aligned to student tokens.
|
||||
"""
|
||||
data_source: 'DataSource'
|
||||
teacher_topk_logprobs: Optional[torch.Tensor] = None
|
||||
teacher_topk_indices: Optional[torch.Tensor] = None
|
||||
@@ -0,0 +1,144 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from swift.rl_core.data import GRPOSample
|
||||
from swift.utils import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _is_async_reward(func: Callable) -> bool:
|
||||
import asyncio
|
||||
return asyncio.iscoroutinefunction(func) or asyncio.iscoroutinefunction(getattr(func, '__call__', None))
|
||||
|
||||
|
||||
def compute_rewards_per_func(
|
||||
samples: List[GRPOSample],
|
||||
reward_funcs: List[Callable],
|
||||
reward_model_plugins: List[Optional[Any]],
|
||||
device: torch.device,
|
||||
trainer_state: Optional[Any] = None,
|
||||
extra_reward_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Compute per-function rewards for ``samples``.
|
||||
|
||||
Supports sync reward callables, async reward callables (auto-detected and
|
||||
run via ``asyncio.run``), and reward models (``nn.Module`` instances) via
|
||||
``reward_model_plugins``.
|
||||
|
||||
Args:
|
||||
samples: On-policy samples carrying completions in ``messages[-1]``.
|
||||
reward_funcs: Reward callables / models.
|
||||
reward_model_plugins: Optional model plugins aligned with ``reward_funcs``.
|
||||
device: Target device for the returned tensor.
|
||||
trainer_state: Passed to reward functions as ``trainer_state`` kwarg.
|
||||
|
||||
Returns:
|
||||
``[N, n_funcs]`` tensor of rewards.
|
||||
"""
|
||||
if reward_model_plugins is None:
|
||||
reward_model_plugins = [None] * len(reward_funcs)
|
||||
async_indices = [i for i, func in enumerate(reward_funcs) if _is_async_reward(func)]
|
||||
|
||||
rewards_per_func = torch.zeros((len(samples), len(reward_funcs)), device=device)
|
||||
completions = [s.messages[-1]['content'] for s in samples]
|
||||
|
||||
reward_kwargs: Dict[str, Any] = {'trainer_state': trainer_state}
|
||||
if extra_reward_kwargs:
|
||||
reward_kwargs.update(extra_reward_kwargs)
|
||||
reward_rows = [s.to_reward_row() for s in samples]
|
||||
if reward_rows:
|
||||
from swift.dataset import RowPreprocessor
|
||||
reward_kwargs.update(RowPreprocessor.rows_to_batched(reward_rows))
|
||||
|
||||
for i, (reward_func, reward_model_plugin) in enumerate(zip(reward_funcs, reward_model_plugins)):
|
||||
if isinstance(reward_func, nn.Module):
|
||||
output = reward_model_plugin(inputs=reward_rows, **reward_kwargs)
|
||||
output = [reward if reward is not None else torch.nan for reward in output]
|
||||
rewards_per_func[:, i] = torch.tensor(output, dtype=torch.float32, device=device)
|
||||
elif i in async_indices:
|
||||
# Async rewards are executed below.
|
||||
pass
|
||||
else:
|
||||
output = reward_func(completions, **reward_kwargs)
|
||||
output = [reward if reward is not None else torch.nan for reward in output]
|
||||
rewards_per_func[:, i] = torch.tensor(output, dtype=torch.float32, device=device)
|
||||
|
||||
if async_indices:
|
||||
import asyncio
|
||||
|
||||
async def _run_async_funcs():
|
||||
coros = [reward_funcs[idx](completions, **reward_kwargs) for idx in async_indices]
|
||||
return await asyncio.gather(*coros)
|
||||
|
||||
for idx, output in zip(async_indices, asyncio.run(_run_async_funcs())):
|
||||
output = [r if r is not None else torch.nan for r in output]
|
||||
rewards_per_func[:, idx] = torch.tensor(output, dtype=torch.float32, device=device)
|
||||
|
||||
if rewards_per_func.shape[1] > 0 and torch.isnan(rewards_per_func).all(dim=1).any():
|
||||
nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0]
|
||||
row_reward_kwargs = {key: value[nan_row_idx] for key, value in reward_kwargs.items() if key != 'trainer_state'}
|
||||
row_reward_kwargs['completion'] = completions[nan_row_idx]
|
||||
logger.warning(f'All reward functions returned None for kwargs: {row_reward_kwargs}. '
|
||||
'Please ensure that at least one reward function returns a valid reward.')
|
||||
|
||||
return rewards_per_func
|
||||
|
||||
|
||||
def score_completions(
|
||||
samples: List[GRPOSample],
|
||||
reward_funcs: List[Callable],
|
||||
reward_model_plugins: List[Optional[Any]],
|
||||
use_gym_env: bool,
|
||||
device: torch.device,
|
||||
trainer_state: Optional[Any] = None,
|
||||
extra_reward_kwargs: Optional[Dict[str, Any]] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Score completions and return per-function rewards.
|
||||
|
||||
When ``use_gym_env`` is set, the ``total_reward`` stored in
|
||||
``sample.rollout_infos`` is appended as an extra reward column.
|
||||
"""
|
||||
if use_gym_env:
|
||||
gym_reward = torch.tensor([s.rollout_infos['total_reward'] for s in samples],
|
||||
dtype=torch.float32,
|
||||
device=device).unsqueeze(1)
|
||||
if not reward_funcs:
|
||||
return gym_reward
|
||||
func_rewards = compute_rewards_per_func(
|
||||
samples,
|
||||
reward_funcs,
|
||||
reward_model_plugins,
|
||||
device=device,
|
||||
trainer_state=trainer_state,
|
||||
extra_reward_kwargs=extra_reward_kwargs,
|
||||
)
|
||||
return torch.cat([func_rewards, gym_reward], dim=1)
|
||||
|
||||
return compute_rewards_per_func(
|
||||
samples,
|
||||
reward_funcs,
|
||||
reward_model_plugins,
|
||||
device=device,
|
||||
trainer_state=trainer_state,
|
||||
extra_reward_kwargs=extra_reward_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def compute_std_for_dynamic_sampling(
|
||||
rewards_per_func: torch.Tensor,
|
||||
reward_weights: torch.Tensor,
|
||||
num_generations: int,
|
||||
) -> torch.Tensor:
|
||||
"""Compute per-sample reward std used by dynamic sampling (DAPO).
|
||||
|
||||
Returns a ``[N]`` tensor; callers are expected to pass already-global rewards.
|
||||
"""
|
||||
rewards = (rewards_per_func * reward_weights.unsqueeze(0)).nansum(dim=1)
|
||||
|
||||
if num_generations > 1:
|
||||
grouped = rewards.view(-1, num_generations)
|
||||
return grouped.std(dim=1).repeat_interleave(num_generations)
|
||||
return torch.zeros_like(rewards)
|
||||
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
"""Shared resample logic for HF / Megatron / Megatron-Ray trainers.
|
||||
|
||||
When ``truncation_strategy='delete'`` (or dynamic sampling), samples whose
|
||||
``template.encode`` fails (e.g. exceed ``max_length``, or multimodal processing
|
||||
errors) must be replaced with fresh ones drawn from a backup iterator until we
|
||||
have ``len(inputs)`` valid samples. The backends previously each carried a
|
||||
near-identical copy of this loop; they only differed in the iterator they pull
|
||||
from and whether the assistant response is stripped before encoding (prompt-only
|
||||
algorithms like GRPO strip; GKD off-policy distillation keeps it).
|
||||
"""
|
||||
from typing import Iterator, List
|
||||
|
||||
from swift.template.base import Template
|
||||
from swift.utils import get_logger, remove_response
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def resample_encode_failed_inputs(
|
||||
template: Template,
|
||||
data_iterator: Iterator,
|
||||
inputs: List[dict],
|
||||
max_resample_rounds: int = 10,
|
||||
strip_response: bool = True,
|
||||
) -> List[dict]:
|
||||
"""Replace samples whose encode fails with fresh ones from ``data_iterator``.
|
||||
|
||||
Caps the TOTAL encode attempts (fail-fast): a systematic failure
|
||||
(e.g. ``max_length`` too small so every prompt is over-length) raises quickly
|
||||
instead of churning through the iterator, and an empty batch breaks the loop
|
||||
instead of spinning forever.
|
||||
|
||||
Args:
|
||||
template: Template used to encode (and thereby validate) a sample.
|
||||
data_iterator: Backup iterator yielding batches (lists) of fresh samples.
|
||||
inputs: The current batch; its length is the required valid count.
|
||||
max_resample_rounds: Resample budget; total attempts == required * (rounds + 1).
|
||||
strip_response: Remove the assistant response (in place) before encoding.
|
||||
|
||||
Returns:
|
||||
A list of valid samples with the same length as ``inputs``.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If not enough valid samples are collected after the budget.
|
||||
"""
|
||||
required = len(inputs)
|
||||
max_attempts = required * (max_resample_rounds + 1)
|
||||
valid, pending = [], list(inputs)
|
||||
attempts = n_dropped = 0
|
||||
|
||||
while len(valid) < required and attempts < max_attempts:
|
||||
if not pending:
|
||||
batch = list(next(data_iterator))
|
||||
if not batch: # guard: an empty batch would otherwise spin forever
|
||||
break
|
||||
pending.extend(batch)
|
||||
item = pending.pop(0)
|
||||
attempts += 1
|
||||
try:
|
||||
if strip_response:
|
||||
remove_response(item['messages'])
|
||||
template.encode(item)
|
||||
valid.append(item)
|
||||
except Exception as e:
|
||||
n_dropped += 1
|
||||
logger.info(f'Encoding failed for one sample; will resample. {e}')
|
||||
|
||||
if len(valid) < required:
|
||||
raise RuntimeError(f'resample: only collected {len(valid)}/{required} valid samples after {attempts} '
|
||||
f'attempts ({n_dropped} failed). Increase `max_length` or adjust `truncation_strategy`.')
|
||||
return valid[:required]
|
||||
Reference in New Issue
Block a user