From 6bcf86bd3392e7030af98edec9e56a2b021e06c4 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 1 Oct 2026 19:05:29 +0000 Subject: [PATCH 1/8] feature(distillation): initial implementation added --- areal/api/cli_args.py | 43 +++++ areal/dataset/__init__.py | 21 +++ areal/dataset/aime.py | 77 ++++++++ areal/trainer/ppo/actor.py | 304 +++++++++++++++++++++++------- areal/utils/logging.py | 1 + areal/workflow/__init__.py | 2 + areal/workflow/self_adaptation.py | 211 +++++++++++++++++++++ docs/en/cli_reference.md | 17 ++ docs/zh/cli_reference.md | 17 ++ examples/distillation/opsa.py | 66 +++++++ 10 files changed, 690 insertions(+), 69 deletions(-) create mode 100644 areal/dataset/aime.py create mode 100644 areal/workflow/self_adaptation.py create mode 100644 examples/distillation/opsa.py diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index cee9dde85f..ce35212683 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -1768,6 +1768,32 @@ def __post_init__(self): ) +@dataclass +class OPSAConfig: + """Configuration for On-Policy Self-Adaptation (OPSA). + + OPSA (Section 4.2 of ``Does On-Policy Distillation Really Distill?``) + suppresses the lowest-probability sampled tokens without task rewards or a + teacher. The selected fraction is evaluated independently for every + response, and its negative advantage is scaled by token entropy. + """ + + lowest_logp_fraction: float = field( + default=0.2, + metadata={ + "help": "Fraction of valid response tokens with the lowest rollout logp to update." + }, + ) + + def __post_init__(self): + """Validate OPSA settings.""" + if not 0 < self.lowest_logp_fraction <= 1: + raise ValueError( + "opsa.lowest_logp_fraction must be in (0, 1], got " + f"{self.lowest_logp_fraction}" + ) + + @dataclass class PPOActorConfig(TrainEngineConfig): """Configuration for PPO actor model, a subclass of a TrainEngine.""" @@ -1867,6 +1893,13 @@ class PPOActorConfig(TrainEngineConfig): adv_norm: NormConfig | None = field( default=None, metadata={"help": "Normalization configuration for advantages."} ) + opsa: OPSAConfig | None = field( + default=None, + metadata={ + "help": "Optional reward-free On-Policy Self-Adaptation configuration." + }, + ) + token_rewards_as_adv: bool = field( default=True, metadata={ @@ -2117,6 +2150,16 @@ def __post_init__(self): " upper: 5.0" ) + if self.opsa is not None: + if self.use_decoupled_loss: + raise ValueError("OPSA requires use_decoupled_loss=False") + if self.use_sapo_loss or self.use_cispo_loss: + raise ValueError("OPSA cannot be combined with SAPO or CISPO") + if self.importance_sampling_level != "token": + raise ValueError("OPSA only supports importance_sampling_level='token'") + if self.m2_threshold is not None or self.rejection_sampling is not None: + raise ValueError("OPSA cannot be combined with token filtering") + # Validate SAPO configuration if self.use_sapo_loss: if self.sapo_tau_pos <= 0 or self.sapo_tau_neg <= 0: diff --git a/areal/dataset/__init__.py b/areal/dataset/__init__.py index 410ee0449a..6ebbcc9147 100644 --- a/areal/dataset/__init__.py +++ b/areal/dataset/__init__.py @@ -31,6 +31,7 @@ "hh-rlhf", "torl_data", "swe_sft", + "aime", ] logger = logging.getLogger("Dataset") @@ -94,6 +95,26 @@ def _get_custom_dataset( max_length=max_length, **kwargs, ) + elif "aime" in path and type == "sft": + from .aime import get_aime_sft_dataset + + return get_aime_sft_dataset( + path=path, + split=split, + tokenizer=tokenizer, + max_length=max_length, + **kwargs, + ) + elif "aime" in path and type == "rl": + from .aime import get_aime_rl_dataset + + return get_aime_rl_dataset( + path=path, + split=split, + tokenizer=tokenizer, + max_length=max_length, + **kwargs, + ) elif "clevr_count_70k" in path and type == "sft": from .clevr_count_70k import get_clevr_count_70k_sft_dataset diff --git a/areal/dataset/aime.py b/areal/dataset/aime.py new file mode 100644 index 0000000000..d378f54190 --- /dev/null +++ b/areal/dataset/aime.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +import os +from datasets import load_dataset + +def get_aime_sft_dataset( + path: str, + split: str, + tokenizer, + max_length: int | None = None, +): + dataset = load_dataset( + "parquet", + data_files={ + "train": os.path.join(path, "train.parquet"), + "test": os.path.join(path, "test.parquet"), + } + ) + + dataset = dataset[split] + + def process(sample): + seq_token = tokenizer.encode( + sample["question"] + sample["answer"] + tokenizer.eos_token + ) + prompt_token = tokenizer.encode(sample["question"]) + loss_mask = [0] * len(prompt_token) + [1] * (len(seq_token) - len(prompt_token)) + return {"input_ids": seq_token, "loss_mask": loss_mask} + + dataset = dataset.map(process).remove_columns(["question", "answer"]) + + if max_length is not None: + # Filter out sequences longer than max_length + dataset = dataset.filter(lambda x: len(x["input_ids"]) <= max_length) + + return dataset + + +def get_aime_rl_dataset( + path: str, + split: str, + tokenizer, + max_length: int | None = None, +): + dataset = load_dataset( + "parquet", + data_files={ + "train": os.path.join(path, "train.parquet"), + "test": os.path.join(path, "test.parquet"), + } + ) + + dataset = dataset[split] + + def process(sample): + messages = [ + { + "role": "user", + "content": sample["question"] + + "\nPlease put your final answer within \\boxed{}.", + } + ] + return {"messages": messages} + + dataset = dataset.map(process).remove_columns(["question"]) + + # Filter out sequences longer than max_length if tokenizer and max_length are provided + if max_length is not None: + + def filter_length(sample): + # Tokenize the user content to check length + content = sample["messages"][0]["content"] + tokens = tokenizer.encode(content) + return len(tokens) <= max_length + + dataset = dataset.filter(filter_length) + + return dataset diff --git a/areal/trainer/ppo/actor.py b/areal/trainer/ppo/actor.py index 3ef1a7bac3..b34cde5cb5 100644 --- a/areal/trainer/ppo/actor.py +++ b/areal/trainer/ppo/actor.py @@ -7,7 +7,12 @@ import torch from areal.api import TrainEngine -from areal.api.cli_args import MOPDLossConfig, PPOActorConfig, RejectionSamplingConfig +from areal.api.cli_args import ( + MOPDLossConfig, + OPSAConfig, + PPOActorConfig, + RejectionSamplingConfig, +) from areal.engine.core import stage_batch_for_engine from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value @@ -56,6 +61,115 @@ logger = logging.getLogger("PPOActor") +def opsa_loss_fn( + logprobs: torch.Tensor, + rollout_logprobs: torch.Tensor, + entropy: torch.Tensor, + loss_mask: torch.Tensor, + lowest_logp_fraction: float, + rewards: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute the reward-free OPSA objective from Eq. 5. + + The rollout log probabilities choose the lowest-logp tokens for each + response. Entropy from the current student forward pass is detached before + forming the adaptive negative advantages, so gradients flow only through + the selected sampled-token log probabilities. + """ + if logprobs.shape != rollout_logprobs.shape or logprobs.shape != entropy.shape: + raise ValueError( + "OPSA logprobs, rollout_logprobs, and entropy must share a shape" + ) + if logprobs.shape != loss_mask.shape or logprobs.ndim not in (1, 2): + raise ValueError("OPSA requires matching 1D packed or 2D padded tensors") + if logprobs.ndim == 2 and cu_seqlens is not None: + raise ValueError("OPSA cu_seqlens is only valid for 1D packed tensors") + if cu_seqlens is not None and cu_seqlens.ndim != 1: + raise ValueError("OPSA cu_seqlens must be a 1D tensor") + + if logprobs.ndim == 2: + response_logprobs = tuple(logprobs.unbind(dim=0)) + response_rollout_logprobs = tuple(rollout_logprobs.unbind(dim=0)) + response_entropies = tuple(entropy.unbind(dim=0)) + response_masks = tuple(loss_mask.bool().unbind(dim=0)) + elif cu_seqlens is None: + response_logprobs = (logprobs,) + response_rollout_logprobs = (rollout_logprobs,) + response_entropies = (entropy,) + response_masks = (loss_mask.bool(),) + else: + # torch.tensor_split requires tensor split indices to be on CPU. + split_indices = cu_seqlens[1:-1].to(device="cpu", dtype=torch.long) + response_logprobs = tuple(torch.tensor_split(logprobs, split_indices)) + response_rollout_logprobs = tuple( + torch.tensor_split(rollout_logprobs, split_indices) + ) + response_entropies = tuple(torch.tensor_split(entropy, split_indices)) + response_masks = tuple(torch.tensor_split(loss_mask.bool(), split_indices)) + + selected_masks = [] + response_advantages = [] + for response_logp, response_rollout_logp, response_entropy, response_mask in zip( + response_logprobs, + response_rollout_logprobs, + response_entropies, + response_masks, + strict=True, + ): + selected_mask = torch.zeros_like(response_mask) + advantages = torch.zeros_like(response_logp) + valid_indices = response_mask.nonzero(as_tuple=True)[0] + if valid_indices.numel() == 0: + selected_masks.append(selected_mask) + response_advantages.append(advantages) + continue + count = max(1, math.ceil(valid_indices.numel() * lowest_logp_fraction)) + selected_indices = valid_indices[ + response_rollout_logp[valid_indices].topk(count, largest=False).indices + ] + selected_entropy = response_entropy[selected_indices].detach() + entropy_range = selected_entropy.max() - selected_entropy.min() + entropy_offset = torch.where( + entropy_range > 0, + (selected_entropy - selected_entropy.min()) / entropy_range, + torch.zeros_like(selected_entropy), + ) + advantages[selected_indices] = -0.5 - 0.5 * entropy_offset + selected_mask[selected_indices] = True + selected_masks.append(selected_mask) + response_advantages.append(advantages) + + if logprobs.ndim == 2: + selected_mask = torch.stack(selected_masks) + advantages = torch.stack(response_advantages) + else: + selected_mask = torch.cat(selected_masks) + advantages = torch.cat(response_advantages) + + if rewards is not None: + rewards = rewards.to(device=advantages.device, dtype=advantages.dtype) + + reward_mask = (rewards > 0).unsqueeze(-1) + + advantages = advantages * reward_mask + selected_mask = selected_mask & reward_mask.bool() + + selected_count = selected_mask.count_nonzero().clamp(min=1) + per_token_loss = -advantages * logprobs + loss = (per_token_loss * selected_mask).sum() / selected_count + zero = torch.zeros_like(selected_mask) + return loss, { + "loss": per_token_loss.detach(), + "importance_weight": torch.ones_like(logprobs), + "approx_kl": (logprobs - rollout_logprobs).detach(), + "clip_mask": zero, + "dual_clip_mask": zero, + "opsa_selected_mask": selected_mask, + "opsa_advantages": advantages.detach(), + } + + def _group_training_metrics( loss_mask: torch.Tensor, group_sizes: list[int], @@ -369,6 +483,8 @@ def _compute_advantages( ) bs = data["input_ids"].shape[0] + if self.config.opsa is not None: + return self._prepare_opsa_batch(data) batch_indices = torch.arange( bs, device=data["input_ids"].device, dtype=torch.long ) @@ -625,6 +741,19 @@ def _compute_advantages( return data + def _prepare_opsa_batch(self, data: dict[str, Any]) -> dict[str, Any]: + """Align rollout log probabilities for the reward-free OPSA objective.""" + loss_mask = torch.roll(data["loss_mask"].float(), shifts=-1, dims=-1) + old_logp = torch.roll(data["logprobs"], shifts=-1, dims=-1) * loss_mask + zeros = torch.zeros_like(old_logp) + data["advantages"] = zeros + data["returns"] = zeros + data["kl_rewards"] = zeros + data["tot_rewards"] = zeros + data["loss_mask"] = loss_mask + data["logprobs"] = old_logp + return data + def _compute_gae_lambda( self, loss_mask: torch.Tensor, @@ -838,6 +967,7 @@ def _ppo_update( sapo_tau_pos=self.config.sapo_tau_pos, sapo_tau_neg=self.config.sapo_tau_neg, use_cispo_loss=self.config.use_cispo_loss, + opsa_config=self.config.opsa, use_decoupled_loss=self.config.use_decoupled_loss, mopd_loss_config=self._mopd_loss_config, ), @@ -1043,6 +1173,7 @@ def grpo_loss_fn( use_sapo_loss: bool = False, sapo_tau_pos: float = 1.0, sapo_tau_neg: float = 1.05, + opsa_config: OPSAConfig | None = None, use_cispo_loss: bool = False, use_decoupled_loss: bool = False, mopd_loss_config: MOPDLossConfig | None = None, @@ -1117,81 +1248,96 @@ def grpo_loss_fn( old_logp = input_data["logprobs"] advantages = input_data["advantages"] - prox_logp_gt = input_data.get("prox_logp") # Could be None if skipped entropy = entropy.detach() + prox_logp_gt = None - if ProxLogpMethod(prox_logp_method) == ProxLogpMethod.REUSE_TRAIN_LOGP: - prox_logp_gt = logprobs.detach() - - # Resolve proximal log-probabilities based on method - prox_logp = _resolve_proximal_logp( - prox_logp_gt=prox_logp_gt, - prox_logp_method=prox_logp_method, - old_logp=old_logp, - logprobs=logprobs.detach(), - versions=input_data.get("versions"), - current_version=current_version, - ) - - # Apply M2PO masking if threshold is set - if m2_threshold is not None: - loss_mask = _apply_m2po_masking(old_logp, prox_logp, loss_mask, m2_threshold) - - # Use CISPO, SAPO, or PPO loss - if use_cispo_loss: - if use_sapo_loss: - raise ValueError( - "CISPO and SAPO are mutually exclusive surrogates. " - "Set at most one of use_cispo_loss / use_sapo_loss." - ) - if importance_sampling_level != "token": - raise ValueError( - "CISPO only supports importance_sampling_level='token'. " - "Sequence-level (GSPO-style) CISPO has no published surrogate." - ) - loss, stat = cispo_loss_fn( + # OPSA directly optimizes selected sampled-token log probabilities rather + # than using PPO's reward-derived surrogate. + if opsa_config is not None: + loss, stat = opsa_loss_fn( logprobs=logprobs, - proximal_logprobs=prox_logp, - advantages=advantages, - eps_clip=eps_clip, - eps_clip_higher=eps_clip_higher, + rollout_logprobs=old_logp, + entropy=entropy, loss_mask=loss_mask, - old_logprobs=old_logp, - rejection_sampling=rejection_sampling, - cu_seqlens=input_data.get("cu_seqlens"), - ) - elif use_sapo_loss: - if use_decoupled_loss: - raise ValueError( - "SAPO is not compatible with `use_decoupled_loss=True`. " - "Please set `actor.use_decoupled_loss=false` in your configuration." - ) - loss, stat = sapo_loss_fn( - logprobs=logprobs, - old_logprobs=old_logp, - advantages=advantages, - tau_pos=sapo_tau_pos, - tau_neg=sapo_tau_neg, - loss_mask=loss_mask, - importance_sampling_level=importance_sampling_level, + lowest_logp_fraction=opsa_config.lowest_logp_fraction, cu_seqlens=input_data.get("cu_seqlens"), ) else: - loss, stat = ppo_actor_loss_fn( - logprobs=logprobs, - old_logprobs=old_logp, - advantages=advantages, - eps_clip=eps_clip, - eps_clip_higher=eps_clip_higher, - loss_mask=loss_mask, - c_clip=c_clip, - proximal_logprobs=prox_logp, - rejection_sampling=rejection_sampling, - importance_sampling_level=importance_sampling_level, - cu_seqlens=input_data.get("cu_seqlens"), + prox_logp_gt = input_data.get("prox_logp") + if ProxLogpMethod(prox_logp_method) == ProxLogpMethod.REUSE_TRAIN_LOGP: + prox_logp_gt = logprobs.detach() + + # Resolve proximal log-probabilities based on method + prox_logp = _resolve_proximal_logp( + prox_logp_gt=prox_logp_gt, + prox_logp_method=prox_logp_method, + old_logp=old_logp, + logprobs=logprobs.detach(), + versions=input_data.get("versions"), + current_version=current_version, ) + # Apply M2PO masking if threshold is set + if m2_threshold is not None: + loss_mask = _apply_m2po_masking( + old_logp, prox_logp, loss_mask, m2_threshold + ) + + # Use CISPO, SAPO, or PPO loss + if use_cispo_loss: + if use_sapo_loss: + raise ValueError( + "CISPO and SAPO are mutually exclusive surrogates. " + "Set at most one of use_cispo_loss / use_sapo_loss." + ) + if importance_sampling_level != "token": + raise ValueError( + "CISPO only supports importance_sampling_level='token'. " + "Sequence-level (GSPO-style) CISPO has no published surrogate." + ) + loss, stat = cispo_loss_fn( + logprobs=logprobs, + proximal_logprobs=prox_logp, + advantages=advantages, + eps_clip=eps_clip, + eps_clip_higher=eps_clip_higher, + loss_mask=loss_mask, + old_logprobs=old_logp, + rejection_sampling=rejection_sampling, + cu_seqlens=input_data.get("cu_seqlens"), + ) + elif use_sapo_loss: + if use_decoupled_loss: + raise ValueError( + "SAPO is not compatible with `use_decoupled_loss=True`. " + "Please set `actor.use_decoupled_loss=false` in your configuration." + ) + loss, stat = sapo_loss_fn( + logprobs=logprobs, + old_logprobs=old_logp, + advantages=advantages, + tau_pos=sapo_tau_pos, + tau_neg=sapo_tau_neg, + loss_mask=loss_mask, + importance_sampling_level=importance_sampling_level, + cu_seqlens=input_data.get("cu_seqlens"), + ) + else: + loss, stat = ppo_actor_loss_fn( + logprobs=logprobs, + old_logprobs=old_logp, + advantages=advantages, + eps_clip=eps_clip, + eps_clip_higher=eps_clip_higher, + loss_mask=loss_mask, + c_clip=c_clip, + proximal_logprobs=prox_logp, + rejection_sampling=rejection_sampling, + importance_sampling_level=importance_sampling_level, + cu_seqlens=input_data.get("cu_seqlens"), + ) + # M2 is part of the shared training-validity contract. Behavioral # rejection may narrow the MOPD numerator further, while its denominator # stays at the pre-rejection count to avoid amplifying accepted tokens. @@ -1283,9 +1429,6 @@ def grpo_loss_fn( stats_tracker.stat( importance_weight=stat["importance_weight"], approx_kl=stat["approx_kl"], - new_logp=logprobs.detach(), - old_logp=old_logp, - entropy=entropy.float(), actor_loss=stat["loss"], clip_ratio=stat["clip_mask"].float(), dual_clip_ratio=stat["dual_clip_mask"].float(), @@ -1294,6 +1437,29 @@ def grpo_loss_fn( denominator="n_valid_tokens", ) + if opsa_config is None: + stats_tracker.stat( + new_logp=logprobs.detach(), + old_logp=old_logp, + entropy=entropy.float(), + denominator="n_valid_tokens", + ) + else: + opsa_selected_mask = stat["opsa_selected_mask"] + opsa_advantages = stat["opsa_advantages"] + + stats_tracker.denominator( + n_opsa_tokens=opsa_selected_mask, + ) + + stats_tracker.stat( + advantages=opsa_advantages, + new_logp=logprobs.detach() * opsa_selected_mask, + old_logp=old_logp * opsa_selected_mask, + entropy=entropy.float() * opsa_selected_mask, + denominator="n_opsa_tokens", + ) + if "behave_imp_weight" in stat: stats_tracker.denominator(unclipped_behave_tokens=stat["behave_mask"]) stats_tracker.stat( diff --git a/areal/utils/logging.py b/areal/utils/logging.py index 349146ac77..e9cc37727e 100644 --- a/areal/utils/logging.py +++ b/areal/utils/logging.py @@ -53,6 +53,7 @@ "VisionRLVRWorkflow": "light_purple", "MultiTurnWorkflow": "light_purple", "MultiTurnV2Workflow": "light_purple", + "SelfAdaptationWorkflow": "light_purple", # Controllers - white "TrainController": "white", "RolloutController": "white", diff --git a/areal/workflow/__init__.py b/areal/workflow/__init__.py index fe4f38417c..aefb212057 100644 --- a/areal/workflow/__init__.py +++ b/areal/workflow/__init__.py @@ -4,12 +4,14 @@ "RLVRWorkflow", "MultiTurnWorkflow", "VisionRLVRWorkflow", + "SelfAdaptationWorkflow", ] _LAZY_IMPORTS = { "RLVRWorkflow": "areal.workflow.rlvr", "MultiTurnWorkflow": "areal.workflow.multi_turn", "VisionRLVRWorkflow": "areal.workflow.vision_rlvr", + "SelfAdaptationWorkflow": "areal.workflow.self_adaptation", } diff --git a/areal/workflow/self_adaptation.py b/areal/workflow/self_adaptation.py new file mode 100644 index 0000000000..c627a46422 --- /dev/null +++ b/areal/workflow/self_adaptation.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Rollout workflow for On-Policy Self-Adaptation. + +The external reward is computed for evaluation/monitoring only. +It is NOT used by the OPSA actor/update path. +""" + +import uuid +from collections.abc import Callable +from typing import Any + +import torch +from transformers import PreTrainedTokenizerFast + +from areal import workflow_context +from areal.api import ( + AsyncRewardWrapper, + InferenceEngine, + ModelRequest, + RolloutWorkflow, +) +from areal.api.cli_args import GenerationHyperparameters +from areal.utils import stats_tracker +from areal.utils.dynamic_import import import_from_string +from areal.utils.hf_utils import apply_chat_template, load_hf_tokenizer +from areal.utils.perf_tracer import ( + atrace_session_phase, + session_context, + trace_session, +) + + +class SelfAdaptationWorkflow(RolloutWorkflow): + """Generate OPSA trajectories without using external rewards for training. + + The reward function is evaluated only for monitoring/evaluation. + The OPSA actor path ignores the reward and uses its own self-adaptation + signal. + """ + + def __init__( + self, + reward_fn: Callable[..., Any] | str, + gconfig: GenerationHyperparameters, + tokenizer: PreTrainedTokenizerFast | str, + enable_thinking: bool = False, + ): + if isinstance(tokenizer, str): + tokenizer = load_hf_tokenizer(tokenizer) + + self.tokenizer = tokenizer + self.gconfig = gconfig.new_with_stop_and_pad_token_ids(tokenizer) + self.enable_thinking = enable_thinking + + # Keep reward function for evaluation/monitoring only. + self.reward_fn = reward_fn + + # Lazily initialize when reward_fn is a string. + if not isinstance(reward_fn, str): + self.async_reward_fn = AsyncRewardWrapper(reward_fn) + + @trace_session("reward") + async def _compute_rewards( + self, + resp, + prompt_str: str, + task_data: dict[str, Any], + ) -> float: + """Compute external reward for evaluation/monitoring only.""" + + completions_str = self.tokenizer.decode(resp.output_tokens) + + reward = await self.async_reward_fn( + prompt_str, + completions_str, + resp.input_tokens, + resp.output_tokens, + **task_data, + ) + + return reward + + @session_context() + async def arun_episode( + self, + engine: InferenceEngine, + data: dict[str, Any], + ) -> dict[str, torch.Tensor]: + """Generate one OPSA trajectory. + + The external reward is computed for evaluation/monitoring only. + The returned `rewards` field remains a zero placeholder because + the OPSA actor path does not consume it. + """ + + # --------------------------------------------------------- + # Lazily load reward function if configured as a string. + # --------------------------------------------------------- + if isinstance(self.reward_fn, str): + self.reward_fn = import_from_string(self.reward_fn) + self.async_reward_fn = AsyncRewardWrapper(self.reward_fn) + + # --------------------------------------------------------- + # Build prompt using the NEW AReaL chat-template helper. + # --------------------------------------------------------- + input_ids = apply_chat_template( + self.tokenizer, + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + + request = ModelRequest( + rid=uuid.uuid4().hex, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + tokenizer=self.tokenizer, + ) + + # Keep this consistent with the generated input IDs. + prompt_str = self.tokenizer.decode(input_ids) + + # --------------------------------------------------------- + # OPSA rollout + # --------------------------------------------------------- + async with atrace_session_phase("generate"): + response = await engine.agenerate(request) + + # --------------------------------------------------------- + # External reward: + # + # IMPORTANT: + # This is ONLY for evaluation/monitoring. + # It does NOT become the OPSA training reward. + # --------------------------------------------------------- + reward = await self._compute_rewards( + response, + prompt_str, + data, + ) + + stats_tracker.get(workflow_context.stat_scope()).scalar( + reward=reward + ) + + # --------------------------------------------------------- + # Build trajectory + # --------------------------------------------------------- + sequence = response.input_tokens + response.output_tokens + + logprobs = ( + [0.0] * response.input_len + + response.output_logprobs + ) + + loss_mask = ( + [0] * response.input_len + + [1] * response.output_len + ) + + versions = ( + [-1] * response.input_len + + response.output_versions + ) + + turn_ids = ( + [-1] * response.input_len + + [0] * response.output_len + ) + + trajectory = { + "input_ids": torch.tensor( + sequence, + dtype=torch.int32, + ), + "loss_mask": torch.tensor( + loss_mask, + dtype=torch.int32, + ), + "logprobs": torch.tensor( + logprobs, + dtype=torch.float32, + ), + "versions": torch.tensor( + versions, + dtype=torch.int32, + ), + "turn_ids": torch.tensor( + turn_ids, + dtype=torch.int32, + ), + "attention_mask": torch.ones( + len(sequence), + dtype=torch.bool, + ), + "rewards": torch.tensor( + reward, + dtype=torch.float32, + ), + + "is_truncated": torch.tensor( + response.stop_reason == "length", + dtype=torch.bool, + ), + } + + return { + key: value.unsqueeze(0) + for key, value in trajectory.items() + } \ No newline at end of file diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index c01479ece1..223f7b0f3c 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -86,6 +86,7 @@ For detailed examples, see the experiment configurations in the `examples/` dire - [MOPDTeacher Specification](section-mopd-teacher) - [MegatronEngine Configuration](section-megatron-engine) - [MemoryProfiler Configuration](section-memory-profiler) +- [OPSA Configuration](section-opsa) - [PRMAdvantageShaping Configuration](section-prm-advantage-shaping) - [PRM Configuration](section-prm) - [PRMScorer Configuration](section-prm-scorer) @@ -419,6 +420,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `gae_lambda_kwargs` | `dict` | `{}` | Keyword arguments passed to a custom gae_lambda function. Ignored when gae_lambda is a float. | | `gae_timestep_unit` | string | `"token"` | Timestep unit used by GAE. 'token' preserves standard token-level GAE; 'turn' applies discount and lambda once per generated turn. **Choices:** `token`, `turn` | | `adv_norm` | [`NormConfig`](section-norm) \| None | `None` | Normalization configuration for advantages. | +| `opsa` | [`OPSAConfig`](section-opsa) \| None | `None` | Optional reward-free On-Policy Self-Adaptation configuration. | | `token_rewards_as_adv` | boolean | `True` | How per-token process rewards in 'token_rewards' enter the objective. True (default): shape advantages after GAE and advantage normalization according to rollout.agent.prm.advantage_shaping, keeping process rewards out of returns. False: add one uniform reward at each turn boundary before GAE, propagating it to preceding tokens and critic returns; this mode is incompatible with GVPO and process-weighted shaping. | | `min_usable_group_size` | integer \| None | `None` | Minimum usable rollout slots a prompt group must keep to stay trainable when some slots fail or are filtered. None derives the minimum from reward_norm/adv_norm: 2 when either uses group statistics (1 for a singleton target group), else 1. | | `kl_ctl` | float | `0.1` | KL divergence coefficient | @@ -1309,6 +1311,21 @@ entries for torch.cuda.memory.\_record_memory_history. | `profile_steps` | list of integer | *list* | List of global steps to capture memory snapshots. | | `max_entries` | integer | `100000` | Max entries for memory history ring buffer. | +(section-opsa)= + +## OPSA Configuration + +Configuration for On-Policy Self-Adaptation (OPSA). + +OPSA (Section 4.2 of `Does On-Policy Distillation Really Distill?`) suppresses the +lowest-probability sampled tokens without task rewards or a teacher. The selected +fraction is evaluated independently for every response, and its negative advantage is +scaled by token entropy. + +| Parameter | Type | Default | Description | +| ---------------------- | ----- | ------- | ------------------------------------------------------------------------- | +| `lowest_logp_fraction` | float | `0.2` | Fraction of valid response tokens with the lowest rollout logp to update. | + (section-prm-advantage-shaping)= ## PRMAdvantageShaping Configuration diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 074145d5ee..193a9c05f6 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -84,6 +84,7 @@ python3 train.py --config path/to/config.yaml actor.lr=1e-4 seed=42 - [MOPDTeacher Specification](section-mopd-teacher) - [MegatronEngine Configuration](section-megatron-engine) - [MemoryProfiler Configuration](section-memory-profiler) +- [OPSA Configuration](section-opsa) - [PRMAdvantageShaping Configuration](section-prm-advantage-shaping) - [PRM Configuration](section-prm) - [PRMScorer Configuration](section-prm-scorer) @@ -417,6 +418,7 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `gae_lambda_kwargs` | `dict` | `{}` | Keyword arguments passed to a custom gae_lambda function. Ignored when gae_lambda is a float. | | `gae_timestep_unit` | string | `"token"` | Timestep unit used by GAE. 'token' preserves standard token-level GAE; 'turn' applies discount and lambda once per generated turn. **Choices:** `token`, `turn` | | `adv_norm` | [`NormConfig`](section-norm) \| None | `None` | Normalization configuration for advantages. | +| `opsa` | [`OPSAConfig`](section-opsa) \| None | `None` | Optional reward-free On-Policy Self-Adaptation configuration. | | `token_rewards_as_adv` | boolean | `True` | How per-token process rewards in 'token_rewards' enter the objective. True (default): shape advantages after GAE and advantage normalization according to rollout.agent.prm.advantage_shaping, keeping process rewards out of returns. False: add one uniform reward at each turn boundary before GAE, propagating it to preceding tokens and critic returns; this mode is incompatible with GVPO and process-weighted shaping. | | `min_usable_group_size` | integer \| None | `None` | Minimum usable rollout slots a prompt group must keep to stay trainable when some slots fail or are filtered. None derives the minimum from reward_norm/adv_norm: 2 when either uses group statistics (1 for a singleton target group), else 1. | | `kl_ctl` | float | `0.1` | KL divergence coefficient | @@ -1307,6 +1309,21 @@ entries for torch.cuda.memory.\_record_memory_history. | `profile_steps` | list of integer | *list* | List of global steps to capture memory snapshots. | | `max_entries` | integer | `100000` | Max entries for memory history ring buffer. | +(section-opsa)= + +## OPSA Configuration + +Configuration for On-Policy Self-Adaptation (OPSA). + +OPSA (Section 4.2 of `Does On-Policy Distillation Really Distill?`) suppresses the +lowest-probability sampled tokens without task rewards or a teacher. The selected +fraction is evaluated independently for every response, and its negative advantage is +scaled by token entropy. + +| Parameter | Type | Default | Description | +| ---------------------- | ----- | ------- | ------------------------------------------------------------------------- | +| `lowest_logp_fraction` | float | `0.2` | Fraction of valid response tokens with the lowest rollout logp to update. | + (section-prm-advantage-shaping)= ## PRMAdvantageShaping Configuration diff --git a/examples/distillation/opsa.py b/examples/distillation/opsa.py new file mode 100644 index 0000000000..a280fb1208 --- /dev/null +++ b/examples/distillation/opsa.py @@ -0,0 +1,66 @@ +"""Train a language model with reward-free On-Policy Self-Adaptation.""" + +import sys + +from areal import PPOTrainer +from areal.api.cli_args import GRPOConfig, load_expr_config +from areal.dataset import get_custom_dataset +from areal.utils.hf_utils import load_hf_tokenizer + + +def main(args: list[str]) -> None: + """Launch OPSA training using prompts only; answer labels are not consumed.""" + config, _ = load_expr_config(args, GRPOConfig) + if config.actor.opsa is None: + raise ValueError("OPSA entrypoint requires actor.opsa to be configured") + tokenizer = load_hf_tokenizer(config.tokenizer_path) + train_dataset = get_custom_dataset( + split="train", + dataset_config=config.train_dataset, + tokenizer=tokenizer, + ) + + valid_dataset = get_custom_dataset( + split="test", + dataset_config=config.valid_dataset, + tokenizer=tokenizer, + ) + + workflow_kwargs = { + "reward_fn": "areal.reward.gsm8k.gsm8k_reward_fn", + "gconfig": config.gconfig, + "tokenizer": tokenizer, + "enable_thinking": False, + } + + eval_workflow_kwargs = workflow_kwargs.copy() + + eval_workflow_kwargs["gconfig"] = config.gconfig.new( + temperature=0.7, + ) + + + with PPOTrainer( + config, + train_dataset=train_dataset, + valid_dataset=valid_dataset, + ) as trainer: + + trainer.train( + workflow=( + "areal.workflow.self_adaptation." + "SelfAdaptationWorkflow" + ), + workflow_kwargs=workflow_kwargs, + + eval_workflow=( + "areal.workflow.self_adaptation." + "SelfAdaptationWorkflow" + ), + eval_workflow_kwargs=eval_workflow_kwargs, + ) + + + +if __name__ == "__main__": + main(sys.argv[1:]) From 024ba745d2c81facc9f9e85f19073ed8d29f516f Mon Sep 17 00:00:00 2001 From: root Date: Thu, 1 Oct 2026 20:39:41 +0000 Subject: [PATCH 2/8] feature(distillation): yaml file added --- examples/distillation/opsa.yaml | 169 ++++++++++++++++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 examples/distillation/opsa.yaml diff --git a/examples/distillation/opsa.yaml b/examples/distillation/opsa.yaml new file mode 100644 index 0000000000..84585a98f2 --- /dev/null +++ b/examples/distillation/opsa.yaml @@ -0,0 +1,169 @@ +experiment_name: gsm8k-opsa +trial_name: trial0 + +seed: 1 +enable_offload: false +total_train_epochs: 10 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 1 + n_gpus_per_node: 8 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + + +scheduler: + type: null + +rollout: + backend: "vllm:d4p1t1" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 2 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: false + agent: + mode: inline + export_style: individual + turn_discount: 1.0 + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 1024 + max_tokens: 2048 + greedy: false + temperature: 1.0 + +actor: + backend: "fsdp:d4p1t1" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen3-1.7B + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + optimizer_dtype: bfloat16 + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + packing_algorithm: ffd + optimizer: + type: adam + lr: 1e-6 + weight_decay: 0.017 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + opsa: + lowest_logp_fraction: 0.2 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: false + use_decoupled_loss: false + rejection_sampling: null + reward_norm: null + adv_norm: null + weight_update_mode: xccl + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: {} + +ref: null + +# SGLang +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + context_length: 32768 + mem_fraction_static: 0.8 + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 32768 + gpu_memory_utilization: 0.15 + +# datasets +train_dataset: + batch_size: 64 + shuffle: true + pin_memory: true + num_workers: 4 + path: /path/to/DAPO-Math-17k + type: rl + max_length: 2048 + +valid_dataset: + batch_size: 30 + pin_memory: true + num_workers: 4 + path: /path/to/AIME24 + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + freq_steps: 100 + freq_secs: null + +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + freq_steps: 100 + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: 20 + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false From 5fa421cbe374d3af7e2f66054bc26b01845506b1 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 1 Oct 2026 21:50:08 +0000 Subject: [PATCH 3/8] docs: add README for OPSA implementation --- docs/en/algorithms/opsa.md | 251 +++++++++++++++++++++++++++++++++++++ 1 file changed, 251 insertions(+) create mode 100644 docs/en/algorithms/opsa.md diff --git a/docs/en/algorithms/opsa.md b/docs/en/algorithms/opsa.md new file mode 100644 index 0000000000..eb911d2839 --- /dev/null +++ b/docs/en/algorithms/opsa.md @@ -0,0 +1,251 @@ +# OPSA Implementation in AReaL + +[OPSA (On-Policy Self-Adaptation)]((https://arxiv.org/abs/2608.31046)) is an RL post-training method that uses the model's own generation behavior to adapt the policy optimization process. In particular, OPSA uses token-level policy information to identify tokens where the model's behavior provides a useful learning signal, allowing the optimization to focus on more informative parts of generated responses. + + +# How OPSA Works + +OPSA (On-Policy Self-Adaptation) is a supervision-free alternative to on-policy distillation (OPD). The key observation behind OPSA is that the improvement from OPD mainly comes from suppressing low-probability ("tail") tokens, rather than from the teacher's token-level supervision. + +The main idea is: + +- Generate responses with the current policy + +The current model generates responses on-policy, just as in standard RL or on-policy distillation. + +- Compute token-level policy information + +For each generated token, OPSA computes the model's token-level probability information, including its entropy. + +- Identify uncertain tokens + +High-entropy positions indicate that the model is more uncertain about which token to generate. OPSA uses this uncertainty to determine where stronger learning signals should be applied. + +- Construct self-adaptive negative advantages + +Instead of obtaining token-level supervision from a teacher model, OPSA constructs negative advantages based on token entropy. + +Higher-entropy positions receive stronger learning signals, while the resulting optimization suppresses low-probability tail tokens. + +- Suppress tail tokens and redistribute probability mass + +The optimization decreases the probability of undesirable low-probability tokens while redistributing probability mass toward the model's higher-probability ("head") tokens. + +Conceptually: + +Current policy + │ + ▼ +Generate response + │ + ▼ +Compute token probabilities / entropy + │ + ▼ +Identify high-entropy positions + │ + ▼ +Construct entropy-adaptive +negative advantages + │ + ▼ +Suppress low-probability tail tokens + │ + ▼ +Redistribute probability mass +toward head tokens + │ + ▼ +Updated policy + +Unlike OPD, OPSA does not require a teacher model or teacher-generated token-level targets. It uses information already available from the policy itself to construct the training signal. + + +This implementation integrates OPSA into the AReaL training pipeline and provides: + +* OPSA-based RL training workflow +* Support for **DAPO-Math-17k** as the training dataset +* Support for **AIME 2024** as a validation/evaluation dataset +* Dataset preprocessing scripts for converting datasets into the format expected by AReaL +* Integration with AReaL's rollout and training infrastructure +* Configurable batch size, sequence length, rollout settings, and optimization parameters + + +# Dataset Preparation + + +The OPSA training pipeline expects datasets in Parquet format with the following columns: + +- question +- answer + +The experiments use: + +- DAPO-Math-17k for training +- AIME 2024 for evaluation + +## DAPO-Math-17k + +The original dataset is available at: + +https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k + +The original DAPO-Math-17k dataset contains the problem prompt in prompt and the ground-truth answer in reward_model. The following script extracts the required fields and saves them as a Parquet file. + +```python +import pandas as pd + + +INPUT_FILE = "dapo-math-17k.parquet" +OUTPUT_FILE = "train.parquet" + + +def main(): + # Load the original dataset + df = pd.read_parquet(INPUT_FILE) + + print(f"Original rows: {len(df)}") + + # Extract the question from the first prompt message + df["question"] = df["prompt"].apply( + lambda prompt: prompt[0]["content"] + if isinstance(prompt, list) and prompt + else None + ) + + # Extract the ground-truth answer + df["answer"] = df["reward_model"].apply( + lambda reward_model: reward_model.get("ground_truth") + if isinstance(reward_model, dict) + else None + ) + + # Keep only the fields required by the AReaL pipeline + df = df[["question", "answer"]].dropna( + subset=["question", "answer"] + ) + + print(f"Final rows: {len(df)}") + print("\nExample rows:") + print(df.head()) + + # Save the processed dataset + df.to_parquet(OUTPUT_FILE, index=False) + + print(f"\nSaved to: {OUTPUT_FILE}") + + +if __name__ == "__main__": + main() +``` + +Run the script with: + +```bash +python prepare_dapo.py +``` + +This produces `train.parquet` with the following schema: + +- question +- answer + +## AIME 2024 + +The AIME 2024 dataset is available at: + +https://huggingface.co/datasets/HuggingFaceH4/aime_2024 + +The original dataset uses problem for the problem statement. The following script renames it to question, keeps the required answer column, and saves the dataset as Parquet. + +```python +import os + +from datasets import load_dataset + + +DATASET_NAME = "HuggingFaceH4/aime_2024" +OUTPUT_DIR = "./aime24" + + +def main(): + # Load the AIME 2024 dataset + dataset = load_dataset(DATASET_NAME) + + for split, split_dataset in dataset.items(): + # Validate the expected columns + required_columns = {"problem", "answer"} + missing_columns = required_columns - set(split_dataset.column_names) + + if missing_columns: + raise ValueError( + f"Missing columns in {split}: {sorted(missing_columns)}" + ) + + # Rename problem -> question + split_dataset = split_dataset.rename_column( + "problem", "question" + ) + + # Keep only the fields required by the AReaL pipeline + columns_to_remove = [ + column + for column in split_dataset.column_names + if column not in {"question", "answer"} + ] + + if columns_to_remove: + split_dataset = split_dataset.remove_columns( + columns_to_remove + ) + + # Save as Parquet + os.makedirs(OUTPUT_DIR, exist_ok=True) + + output_file = os.path.join( + OUTPUT_DIR, + f"{split}.parquet", + ) + + split_dataset.to_parquet(output_file) + + print(f"[OK] Saved {split}: {output_file}") + print(f" Rows: {len(split_dataset)}") + print(f" Columns: {split_dataset.column_names}") + + +if __name__ == "__main__": + main() +``` +Run: + +```bash +python prepare_aime24.py +``` + +# Dataset Configuration + +Update the dataset paths in `examples/distillation/opsa.yaml` to point to the locations of your processed datasets. + +--- + +# Running OPSA Training + +The OPSA workflow can be launched using the AReaL entry point: + +```bash +python examples/distillation/opsa.py \ + --config examples/distillation/opsa.yaml \ + scheduler.type=local +``` + +Depending on the AReaL version and cluster configuration, the scheduler configuration can be changed accordingly. + +For example, for a local setup: + +```bash +scheduler.type=local +``` + +For distributed environments, configure the scheduler and worker resources according to the AReaL deployment. + From 9cf066db5d958d19928fa9228ec33c67d4140aaa Mon Sep 17 00:00:00 2001 From: root Date: Thu, 1 Oct 2026 22:05:59 +0000 Subject: [PATCH 4/8] docs: add README for OPSA implementation --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 45b53d8cdc..081ae4189b 100644 --- a/README.md +++ b/README.md @@ -227,6 +227,7 @@ All RL algorithms support both asynchronous and synchronous versions by setting | **RLHF Reward Modeling** | - | - | [🔗 RLHF Example](examples/alignment/hhrlhf_rw.yaml) | | **SFT** | - | - | [🔗 GSM8K Example](examples/math/gsm8k_sft.py) | | **Distillation** | [📖 Docs](docs/en/algorithms/distillation.md) | [📄 Paper](https://arxiv.org/pdf/2506.02208) | [🔗 GSM8K Example](examples/distillation/gsm8k_grpo_distill_mode_trainEngine.yaml) | +| **On-Policy Self-Adaptation** | [📖 Docs](docs/en/algorithms/opsa.md) | [📄 Paper](https://arxiv.org/abs/2608.31046) | [🔗 DAPO Example](examples/distillation/opsa.yaml) | ### Models From dede5b316da96e8122ad284fd830f7f0bc922946 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 1 Oct 2026 22:32:11 +0000 Subject: [PATCH 5/8] chore: fix formatting --- areal/workflow/self_adaptation.py | 30 ++++++------------------------ 1 file changed, 6 insertions(+), 24 deletions(-) diff --git a/areal/workflow/self_adaptation.py b/areal/workflow/self_adaptation.py index c627a46422..32af0d6b64 100644 --- a/areal/workflow/self_adaptation.py +++ b/areal/workflow/self_adaptation.py @@ -140,34 +140,20 @@ async def arun_episode( data, ) - stats_tracker.get(workflow_context.stat_scope()).scalar( - reward=reward - ) + stats_tracker.get(workflow_context.stat_scope()).scalar(reward=reward) # --------------------------------------------------------- # Build trajectory # --------------------------------------------------------- sequence = response.input_tokens + response.output_tokens - logprobs = ( - [0.0] * response.input_len - + response.output_logprobs - ) + logprobs = [0.0] * response.input_len + response.output_logprobs - loss_mask = ( - [0] * response.input_len - + [1] * response.output_len - ) + loss_mask = [0] * response.input_len + [1] * response.output_len - versions = ( - [-1] * response.input_len - + response.output_versions - ) + versions = [-1] * response.input_len + response.output_versions - turn_ids = ( - [-1] * response.input_len - + [0] * response.output_len - ) + turn_ids = [-1] * response.input_len + [0] * response.output_len trajectory = { "input_ids": torch.tensor( @@ -198,14 +184,10 @@ async def arun_episode( reward, dtype=torch.float32, ), - "is_truncated": torch.tensor( response.stop_reason == "length", dtype=torch.bool, ), } - return { - key: value.unsqueeze(0) - for key, value in trajectory.items() - } \ No newline at end of file + return {key: value.unsqueeze(0) for key, value in trajectory.items()} From a24b5ebf4eeb332e68ff22e758aa025197f67d84 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 2 Oct 2026 17:05:45 +0000 Subject: [PATCH 6/8] chore: fix formatting --- areal/dataset/aime.py | 6 ++++-- docs/en/algorithms/opsa.md | 1 - examples/distillation/opsa.py | 14 ++------------ 3 files changed, 6 insertions(+), 15 deletions(-) diff --git a/areal/dataset/aime.py b/areal/dataset/aime.py index d378f54190..4de8584a35 100644 --- a/areal/dataset/aime.py +++ b/areal/dataset/aime.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 import os + from datasets import load_dataset + def get_aime_sft_dataset( path: str, split: str, @@ -13,7 +15,7 @@ def get_aime_sft_dataset( data_files={ "train": os.path.join(path, "train.parquet"), "test": os.path.join(path, "test.parquet"), - } + }, ) dataset = dataset[split] @@ -46,7 +48,7 @@ def get_aime_rl_dataset( data_files={ "train": os.path.join(path, "train.parquet"), "test": os.path.join(path, "test.parquet"), - } + }, ) dataset = dataset[split] diff --git a/docs/en/algorithms/opsa.md b/docs/en/algorithms/opsa.md index eb911d2839..6758f3f079 100644 --- a/docs/en/algorithms/opsa.md +++ b/docs/en/algorithms/opsa.md @@ -248,4 +248,3 @@ scheduler.type=local ``` For distributed environments, configure the scheduler and worker resources according to the AReaL deployment. - diff --git a/examples/distillation/opsa.py b/examples/distillation/opsa.py index a280fb1208..ae9e7af17c 100644 --- a/examples/distillation/opsa.py +++ b/examples/distillation/opsa.py @@ -39,28 +39,18 @@ def main(args: list[str]) -> None: temperature=0.7, ) - with PPOTrainer( config, train_dataset=train_dataset, valid_dataset=valid_dataset, ) as trainer: - trainer.train( - workflow=( - "areal.workflow.self_adaptation." - "SelfAdaptationWorkflow" - ), + workflow=("areal.workflow.self_adaptation.SelfAdaptationWorkflow"), workflow_kwargs=workflow_kwargs, - - eval_workflow=( - "areal.workflow.self_adaptation." - "SelfAdaptationWorkflow" - ), + eval_workflow=("areal.workflow.self_adaptation.SelfAdaptationWorkflow"), eval_workflow_kwargs=eval_workflow_kwargs, ) - if __name__ == "__main__": main(sys.argv[1:]) From 8f2197da8bb12b471489b719717b2d75f70118ee Mon Sep 17 00:00:00 2001 From: root Date: Fri, 2 Oct 2026 18:23:41 +0000 Subject: [PATCH 7/8] chore: fix formatting --- docs/en/algorithms/opsa.md | 49 ++++++++++++++++++-------------------- 1 file changed, 23 insertions(+), 26 deletions(-) diff --git a/docs/en/algorithms/opsa.md b/docs/en/algorithms/opsa.md index 6758f3f079..241b52dfb6 100644 --- a/docs/en/algorithms/opsa.md +++ b/docs/en/algorithms/opsa.md @@ -1,7 +1,6 @@ # OPSA Implementation in AReaL -[OPSA (On-Policy Self-Adaptation)]((https://arxiv.org/abs/2608.31046)) is an RL post-training method that uses the model's own generation behavior to adapt the policy optimization process. In particular, OPSA uses token-level policy information to identify tokens where the model's behavior provides a useful learning signal, allowing the optimization to focus on more informative parts of generated responses. - +[OPSA (On-Policy Self-Adaptation)](<(https://arxiv.org/abs/2608.31046)>) is an RL post-training method that uses the model's own generation behavior to adapt the policy optimization process. In particular, OPSA uses token-level policy information to identify tokens where the model's behavior provides a useful learning signal, allowing the optimization to focus on more informative parts of generated responses. # How OPSA Works @@ -34,46 +33,43 @@ The optimization decreases the probability of undesirable low-probability tokens Conceptually: Current policy - │ - ▼ +│ +▼ Generate response - │ - ▼ +│ +▼ Compute token probabilities / entropy - │ - ▼ +│ +▼ Identify high-entropy positions - │ - ▼ +│ +▼ Construct entropy-adaptive negative advantages - │ - ▼ +│ +▼ Suppress low-probability tail tokens - │ - ▼ +│ +▼ Redistribute probability mass toward head tokens - │ - ▼ +│ +▼ Updated policy Unlike OPD, OPSA does not require a teacher model or teacher-generated token-level targets. It uses information already available from the policy itself to construct the training signal. - This implementation integrates OPSA into the AReaL training pipeline and provides: -* OPSA-based RL training workflow -* Support for **DAPO-Math-17k** as the training dataset -* Support for **AIME 2024** as a validation/evaluation dataset -* Dataset preprocessing scripts for converting datasets into the format expected by AReaL -* Integration with AReaL's rollout and training infrastructure -* Configurable batch size, sequence length, rollout settings, and optimization parameters - +- OPSA-based RL training workflow +- Support for **DAPO-Math-17k** as the training dataset +- Support for **AIME 2024** as a validation/evaluation dataset +- Dataset preprocessing scripts for converting datasets into the format expected by AReaL +- Integration with AReaL's rollout and training infrastructure +- Configurable batch size, sequence length, rollout settings, and optimization parameters # Dataset Preparation - The OPSA training pipeline expects datasets in Parquet format with the following columns: - question @@ -217,6 +213,7 @@ def main(): if __name__ == "__main__": main() ``` + Run: ```bash @@ -227,7 +224,7 @@ python prepare_aime24.py Update the dataset paths in `examples/distillation/opsa.yaml` to point to the locations of your processed datasets. ---- +______________________________________________________________________ # Running OPSA Training From b9db4a365ef85f314c1928a05f4dd3babac686be Mon Sep 17 00:00:00 2001 From: root Date: Fri, 2 Oct 2026 18:40:55 +0000 Subject: [PATCH 8/8] chore: fix formatting --- docs/en/algorithms/opsa.md | 86 ++++++++++++++++++++------------------ 1 file changed, 46 insertions(+), 40 deletions(-) diff --git a/docs/en/algorithms/opsa.md b/docs/en/algorithms/opsa.md index 241b52dfb6..8cbf8446c1 100644 --- a/docs/en/algorithms/opsa.md +++ b/docs/en/algorithms/opsa.md @@ -1,76 +1,75 @@ # OPSA Implementation in AReaL -[OPSA (On-Policy Self-Adaptation)](<(https://arxiv.org/abs/2608.31046)>) is an RL post-training method that uses the model's own generation behavior to adapt the policy optimization process. In particular, OPSA uses token-level policy information to identify tokens where the model's behavior provides a useful learning signal, allowing the optimization to focus on more informative parts of generated responses. +[OPSA (On-Policy Self-Adaptation)](<(https://arxiv.org/abs/2608.31046)>) is an RL +post-training method that uses the model's own generation behavior to adapt the policy +optimization process. In particular, OPSA uses token-level policy information to +identify tokens where the model's behavior provides a useful learning signal, allowing +the optimization to focus on more informative parts of generated responses. # How OPSA Works -OPSA (On-Policy Self-Adaptation) is a supervision-free alternative to on-policy distillation (OPD). The key observation behind OPSA is that the improvement from OPD mainly comes from suppressing low-probability ("tail") tokens, rather than from the teacher's token-level supervision. +OPSA (On-Policy Self-Adaptation) is a supervision-free alternative to on-policy +distillation (OPD). The key observation behind OPSA is that the improvement from OPD +mainly comes from suppressing low-probability ("tail") tokens, rather than from the +teacher's token-level supervision. The main idea is: - Generate responses with the current policy -The current model generates responses on-policy, just as in standard RL or on-policy distillation. +The current model generates responses on-policy, just as in standard RL or on-policy +distillation. - Compute token-level policy information -For each generated token, OPSA computes the model's token-level probability information, including its entropy. +For each generated token, OPSA computes the model's token-level probability information, +including its entropy. - Identify uncertain tokens -High-entropy positions indicate that the model is more uncertain about which token to generate. OPSA uses this uncertainty to determine where stronger learning signals should be applied. +High-entropy positions indicate that the model is more uncertain about which token to +generate. OPSA uses this uncertainty to determine where stronger learning signals should +be applied. - Construct self-adaptive negative advantages -Instead of obtaining token-level supervision from a teacher model, OPSA constructs negative advantages based on token entropy. +Instead of obtaining token-level supervision from a teacher model, OPSA constructs +negative advantages based on token entropy. -Higher-entropy positions receive stronger learning signals, while the resulting optimization suppresses low-probability tail tokens. +Higher-entropy positions receive stronger learning signals, while the resulting +optimization suppresses low-probability tail tokens. - Suppress tail tokens and redistribute probability mass -The optimization decreases the probability of undesirable low-probability tokens while redistributing probability mass toward the model's higher-probability ("head") tokens. +The optimization decreases the probability of undesirable low-probability tokens while +redistributing probability mass toward the model's higher-probability ("head") tokens. Conceptually: -Current policy -│ -▼ -Generate response -│ -▼ -Compute token probabilities / entropy -│ -▼ -Identify high-entropy positions -│ -▼ -Construct entropy-adaptive -negative advantages -│ -▼ -Suppress low-probability tail tokens -│ -▼ -Redistribute probability mass -toward head tokens -│ -▼ +Current policy => Generate response => Compute token probabilities / entropy => Identify +high-entropy positions => Construct entropy-adaptive negative advantages => Suppress +low-probability tail tokens => Redistribute probability mass toward head tokens => Updated policy -Unlike OPD, OPSA does not require a teacher model or teacher-generated token-level targets. It uses information already available from the policy itself to construct the training signal. +Unlike OPD, OPSA does not require a teacher model or teacher-generated token-level +targets. It uses information already available from the policy itself to construct the +training signal. This implementation integrates OPSA into the AReaL training pipeline and provides: - OPSA-based RL training workflow - Support for **DAPO-Math-17k** as the training dataset - Support for **AIME 2024** as a validation/evaluation dataset -- Dataset preprocessing scripts for converting datasets into the format expected by AReaL +- Dataset preprocessing scripts for converting datasets into the format expected by + AReaL - Integration with AReaL's rollout and training infrastructure -- Configurable batch size, sequence length, rollout settings, and optimization parameters +- Configurable batch size, sequence length, rollout settings, and optimization + parameters # Dataset Preparation -The OPSA training pipeline expects datasets in Parquet format with the following columns: +The OPSA training pipeline expects datasets in Parquet format with the following +columns: - question - answer @@ -86,7 +85,9 @@ The original dataset is available at: https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k -The original DAPO-Math-17k dataset contains the problem prompt in prompt and the ground-truth answer in reward_model. The following script extracts the required fields and saves them as a Parquet file. +The original DAPO-Math-17k dataset contains the problem prompt in prompt and the +ground-truth answer in reward_model. The following script extracts the required fields +and saves them as a Parquet file. ```python import pandas as pd @@ -152,7 +153,9 @@ The AIME 2024 dataset is available at: https://huggingface.co/datasets/HuggingFaceH4/aime_2024 -The original dataset uses problem for the problem statement. The following script renames it to question, keeps the required answer column, and saves the dataset as Parquet. +The original dataset uses problem for the problem statement. The following script +renames it to question, keeps the required answer column, and saves the dataset as +Parquet. ```python import os @@ -222,7 +225,8 @@ python prepare_aime24.py # Dataset Configuration -Update the dataset paths in `examples/distillation/opsa.yaml` to point to the locations of your processed datasets. +Update the dataset paths in `examples/distillation/opsa.yaml` to point to the locations +of your processed datasets. ______________________________________________________________________ @@ -236,7 +240,8 @@ python examples/distillation/opsa.py \ scheduler.type=local ``` -Depending on the AReaL version and cluster configuration, the scheduler configuration can be changed accordingly. +Depending on the AReaL version and cluster configuration, the scheduler configuration +can be changed accordingly. For example, for a local setup: @@ -244,4 +249,5 @@ For example, for a local setup: scheduler.type=local ``` -For distributed environments, configure the scheduler and worker resources according to the AReaL deployment. +For distributed environments, configure the scheduler and worker resources according to +the AReaL deployment.