wehub-resource-sync
354 lines
15 KiB
Python
354 lines
15 KiB
Python
# 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
|