diff --git a/cookbook/rl/grpo/group_admission.py b/cookbook/rl/grpo/group_admission.py new file mode 100644 index 000000000..bc9b4d236 --- /dev/null +++ b/cookbook/rl/grpo/group_admission.py @@ -0,0 +1,478 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""End-to-end GSM8K GRPO training with budget-aware group admission. + +This example follows ``short_math_grpo.py`` and inserts complete-group +admission plus bounded replacement sampling between rollout and GRPO. Weight +synchronization happens once per batch, so initial and replacement groups use +the same rollout policy snapshot. +""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from typing import Any, Dict, List, Sequence, Tuple + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceGroup, DeviceMesh, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +from twinkle.advantage.group_admission import ( + GroupAdmissionConfig, + GroupAdmissionDecision, + GroupAdmissionPolicy, + SamplingBudgetConfig, + SamplingBudgetController, + SamplingBudgetState, + group_admission_metrics, +) +from twinkle.checkpoint_engine import CheckpointEngineManager +from twinkle.cli import CLI +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.metric import CompletionRewardMetric +from twinkle.model import TransformersModel +from twinkle.preprocessor.llm import GSM8KProcessor +from twinkle.processor import InputProcessor +from twinkle.reward import GSM8KAccuracyReward +from twinkle.reward.base import Reward +from twinkle.sampler import vLLMSampler + +logger = get_logger() +args = CLI.from_args() + +# ========== Configuration ========== +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.5-4B' +USE_MEGATRON = args.model.strategy != 'native_fsdp' + +MODEL_GPUS = args.infra.model_gpus or 4 +SAMPLER_GPUS = args.infra.sampler_gpus or 4 +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +NUM_GENERATIONS = args.rl.num_generations or 8 +MAX_NEW_TOKENS = args.sampling.max_tokens or 4096 +LEARNING_RATE = args.optimizer.learning_rate or 1e-5 +MAX_STEPS = args.training.max_steps or 1000 +BATCH_SIZE = args.training.batch_size or 8 +MINI_BATCH_SIZE = args.training.mini_batch_size or 8 +MICRO_BATCH_SIZE = args.training.micro_batch_size or 2 +GRADIENT_ACCUMULATION_STEPS = args.training.gradient_accumulation_steps or 1 +ADAPTER_NAME = args.lora.adapter_name or 'default' +SAVE_STEPS = args.training.save_steps or 1000 +LORA_RANK = args.lora.lora_r or 16 + +if MINI_BATCH_SIZE % NUM_GENERATIONS != 0: + raise ValueError( + 'mini_batch_size must be divisible by num_generations to preserve complete group boundaries, ' + f'but got {MINI_BATCH_SIZE} % {NUM_GENERATIONS} != 0') + +# The defaults permit at most one extra complete group per target group. The +# environment variables make the cost envelope adjustable without coupling it +# to a particular sampler implementation. +MAX_EXTRA_GROUPS = int(os.getenv('GROUP_ADMISSION_MAX_EXTRA_GROUPS', BATCH_SIZE)) +MAX_RESAMPLE_ROUNDS = int(os.getenv('GROUP_ADMISSION_MAX_RESAMPLE_ROUNDS', '2')) + +SYSTEM_PROMPT = ('You are a helpful math assistant. Solve the problem with minimal but correct reasoning ' + 'and put your final answer within \\boxed{}.') + + +# ========== Reward Functions ========== +class GSM8KBrevityReward(Reward): + """Reward short completions that contain a recognizable final answer.""" + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for trajectory in trajectories: + messages = trajectory.get('messages', []) + completion = '' + for message in reversed(messages): + if message.get('role') == 'assistant': + completion = message.get('content', '') + break + + has_answer = bool( + re.search(r'\\boxed\{[^}]+\}', completion) + or re.search(r'####\s*[\-\d,\.]+', completion)) + if not has_answer: + rewards.append(0.0) + elif len(completion) <= 300: + rewards.append(1.0) + else: + rewards.append(max(0.0, 1.0 - (len(completion) - 300) / 3000)) + return rewards + + +# ========== Dataset ========== +def create_gsm8k_dataset(): + dataset = Dataset() + dataset.add_dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) + dataset.set_template( + 'Template', + model_id=MODEL_ID, + max_length=4096, + truncation_strategy='delete', + enable_thinking=False, + ) + dataset.map(GSM8KProcessor(system=SYSTEM_PROMPT)) + dataset.encode(add_generation_prompt=True) + return dataset + + +def compute_rewards( + trajectories: List[Dict[str, Any]], +) -> Tuple[List[float], List[float], List[float]]: + accuracy_rewards = GSM8KAccuracyReward()(trajectories) + brevity_rewards = GSM8KBrevityReward()(trajectories) + total_rewards = [accuracy + brevity for accuracy, brevity in zip(accuracy_rewards, brevity_rewards)] + return total_rewards, brevity_rewards, accuracy_rewards + + +@dataclass +class RolloutGroup: + """One prompt and all of its aligned rollout and reward fields.""" + + prompt: Any + input_data: List[Dict[str, Any]] + old_logps: List[List[float]] + completion_lengths: List[int] + completions: List[str] + total_rewards: List[float] + brevity_rewards: List[float] + accuracy_rewards: List[float] + + @property + def generated_tokens(self) -> int: + return sum(self.completion_lengths) + + +def sample_prompt_groups( + sampler: vLLMSampler, + prompts: Sequence[Any], + sampling_params: SamplingParams, +) -> List[RolloutGroup]: + """Sample and score one complete GRPO group for every prompt.""" + expanded_prompts = [prompt for prompt in prompts for _ in range(NUM_GENERATIONS)] + sample_responses = sampler.sample(expanded_prompts, sampling_params) + + input_data: List[Dict[str, Any]] = [] + old_logps: List[List[float]] = [] + completion_lengths: List[int] = [] + completions: List[str] = [] + for response in sample_responses: + for sequence in response.sequences: + if sequence.logprobs is None: + raise RuntimeError('a sampled sequence is missing token log probabilities') + input_data.append(sequence.new_input_feature) + old_logps.append([logprob[0][1] for logprob in sequence.logprobs]) + completion_lengths.append(len(sequence.tokens)) + completions.append(sequence.decoded or '') + + expected = len(prompts) * NUM_GENERATIONS + if len(input_data) != expected: + raise RuntimeError(f'sampler returned {len(input_data)} completions, expected {expected}') + + total_rewards, brevity_rewards, accuracy_rewards = compute_rewards(input_data) + groups = [] + for group_index, prompt in enumerate(prompts): + start = group_index * NUM_GENERATIONS + end = start + NUM_GENERATIONS + groups.append( + RolloutGroup( + prompt=prompt, + input_data=input_data[start:end], + old_logps=old_logps[start:end], + completion_lengths=completion_lengths[start:end], + completions=completions[start:end], + total_rewards=total_rewards[start:end], + brevity_rewards=brevity_rewards[start:end], + accuracy_rewards=accuracy_rewards[start:end], + )) + return groups + + +def flatten_admitted_groups(groups: Sequence[RolloutGroup]) -> Dict[str, list]: + """Flatten groups without changing their order or internal boundaries.""" + return { + 'input_data': [value for group in groups for value in group.input_data], + 'old_logps': [value for group in groups for value in group.old_logps], + 'completion_lengths': [value for group in groups for value in group.completion_lengths], + 'total_rewards': [value for group in groups for value in group.total_rewards], + 'brevity_rewards': [value for group in groups for value in group.brevity_rewards], + 'accuracy_rewards': [value for group in groups for value in group.accuracy_rewards], + } + + +def collect_admitted_groups( + sampler: vLLMSampler, + prompts: Sequence[Any], + sampling_params: SamplingParams, + policy: GroupAdmissionPolicy, +) -> tuple[List[RolloutGroup], List[GroupAdmissionDecision], SamplingBudgetState, bool]: + """Run initial sampling and bounded retries under one policy snapshot.""" + target_groups = len(prompts) + if target_groups == 0: + raise ValueError('at least one prompt is required') + controller = SamplingBudgetController( + SamplingBudgetConfig( + max_extra_groups=MAX_EXTRA_GROUPS, + max_extra_groups_per_round=target_groups, + max_resample_rounds=MAX_RESAMPLE_ROUNDS, + max_total_samples=(target_groups + MAX_EXTRA_GROUPS) * NUM_GENERATIONS, + max_total_tokens=(target_groups + MAX_EXTRA_GROUPS) * NUM_GENERATIONS * MAX_NEW_TOKENS, + )) + state = SamplingBudgetState() + admitted_groups: List[RolloutGroup] = [] + all_decisions: List[GroupAdmissionDecision] = [] + + current_prompts = list(prompts) + retry_queue: List[Any] = [] + round_index = 0 + exhausted = False + + while current_prompts: + groups = sample_prompt_groups(sampler, current_prompts, sampling_params) + decisions = [ + policy.evaluate(group.total_rewards, group.completions) + for group in groups + ] + all_decisions.extend(decisions) + + for group, decision in zip(groups, decisions): + outcome = 'admitted' if decision.admitted else decision.primary_rejection_reason + logger.info( + f'[Group admission] round={round_index} outcome={outcome} ' + f'reward_range={decision.reward_range:.6f} ' + f'nontrivial={decision.nontrivial_advantage_count}') + if decision.admitted: + admitted_groups.append(group) + else: + retry_queue.append(group.prompt) + + controller.observe_round( + state, + decisions, + num_generations=NUM_GENERATIONS, + generated_tokens=sum(group.generated_tokens for group in groups), + ) + plan = controller.plan_resampling( + state, + target_groups=target_groups, + num_generations=NUM_GENERATIONS, + ) + logger.info( + f'[Group admission] plan extra_groups={plan.extra_groups} ' + f'remaining={plan.remaining_target_groups} ' + f'effective_rate={plan.effective_rate:.4f} ' + f'exhausted={plan.exhausted} limited_by={plan.limited_by}') + + if plan.extra_groups == 0: + exhausted = plan.exhausted + break + + # Retry each unresolved prompt at most once per round. When the EMA + # requests more work than there are unresolved prompts, the next round + # can retry them again without admitting duplicate groups for a prompt. + retry_count = min(plan.extra_groups, len(retry_queue)) + current_prompts = retry_queue[:retry_count] + retry_queue = retry_queue[retry_count:] + round_index += 1 + + return admitted_groups, all_decisions, state, exhausted + + +def admission_log_dict( + decisions: Sequence[GroupAdmissionDecision], + state: SamplingBudgetState, + *, + target_groups: int, + admitted_groups: int, + exhausted: bool, +) -> Dict[str, float | int]: + initial_admitted = sum(decision.admitted for decision in decisions[:target_groups]) + initially_rejected = target_groups - initial_admitted + recovered_groups = admitted_groups - initial_admitted + values = { + f'group_admission/{name}': value + for name, value in group_admission_metrics(decisions).items() + } + values.update({ + 'group_admission/target_group_count': target_groups, + 'group_admission/initial_effective_group_rate': initial_admitted / target_groups, + 'group_admission/final_admitted_group_count': admitted_groups, + 'group_admission/final_effective_group_rate': admitted_groups / target_groups, + 'group_admission/recovered_group_count': recovered_groups, + 'group_admission/recovery_rate': recovered_groups / initially_rejected if initially_rejected else 0.0, + 'group_admission/sampling_cost_multiplier': state.generated_groups / target_groups, + 'group_admission/budget_exhausted': int(exhausted), + }) + return values + + +# ========== Main ========== +def main(): + device_groups = [ + DeviceGroup(name='model', ranks=list(range(MODEL_GPUS)), device_type='GPU'), + DeviceGroup(name='sampler', ranks=list(range(MODEL_GPUS, NUM_GPUS)), device_type='GPU'), + ] + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + twinkle.initialize(mode='ray', nproc_per_node=NUM_GPUS, groups=device_groups, lazy_collect=False) + + lora_config = LoraConfig( + target_modules='all-linear', + r=LORA_RANK, + lora_alpha=LORA_RANK * 2, + lora_dropout=0.05, + ) + if USE_MEGATRON: + from twinkle.model.megatron import MegatronModel + model = MegatronModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + mixed_precision='bf16', + variable_seq_lengths=True, + ) + else: + model = TransformersModel( + model_id=MODEL_ID, + device_mesh=model_mesh, + remote_group='model', + ) + + model.add_adapter_to_model( + ADAPTER_NAME, + lora_config, + gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS, + ) + if USE_MEGATRON: + model.set_optimizer('default', lr=LEARNING_RATE) + model.set_lr_scheduler('default', lr_decay_steps=MAX_STEPS, max_lr=LEARNING_RATE) + else: + model.set_optimizer('AdamW', lr=LEARNING_RATE) + model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=0) + + model.set_loss('GRPOLoss', epsilon=0.2) + model.set_processor(InputProcessor, padding_free=True) + model.set_template('Template', model_id=MODEL_ID, enable_thinking=False) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 8192, + 'max_lora_rank': 32, + 'enable_lora': True, + 'enable_tower_connector_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Template', model_id=MODEL_ID, enable_thinking=False) + + ckpt_manager = CheckpointEngineManager(model=model, sampler=sampler) + global_batch_size = BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS + dataloader = DataLoader( + dataset=create_gsm8k_dataset, + batch_size=global_batch_size, + min_batch_size=global_batch_size, + device_mesh=model_mesh, + remote_group='model', + ) + + advantage_fn = GRPOAdvantage() + metrics = CompletionRewardMetric() + sampling_params = SamplingParams( + max_tokens=MAX_NEW_TOKENS, + num_samples=1, + logprobs=1, + temperature=1.0, + top_p=0.95, + ) + admission_policy = GroupAdmissionPolicy( + GroupAdmissionConfig( + min_reward_std=1e-6, + min_reward_range=0.02, + min_nontrivial_advantages=2, + advantage_tolerance=0.01, + )) + + optim_step = 0 + logger.info('Starting GSM8K GRPO training with budget-aware group admission') + logger.info(get_device_placement()) + + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + metrics.reset() + prompts = batch if isinstance(batch, list) else [batch] + + # Keep one rollout-policy snapshot across the initial candidates and + # every replacement round in this batch. + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + admitted_groups, decisions, budget_state, exhausted = collect_admitted_groups( + sampler, + prompts, + sampling_params, + admission_policy, + ) + admission_metrics = admission_log_dict( + decisions, + budget_state, + target_groups=len(prompts), + admitted_groups=len(admitted_groups), + exhausted=exhausted, + ) + + if not admitted_groups: + logger.warning(f'No complete groups admitted; skipping batch. {admission_metrics}') + continue + + rollout_batch = flatten_admitted_groups(admitted_groups) + metrics.accumulate( + completion_lengths=rollout_batch['completion_lengths'], + rewards={ + 'total': rollout_batch['total_rewards'], + 'brevity': rollout_batch['brevity_rewards'], + 'accuracy': rollout_batch['accuracy_rewards'], + }, + ) + advantages = advantage_fn( + rollout_batch['total_rewards'], + num_generations=NUM_GENERATIONS, + scale='group', + ).tolist() + + total_completions = len(rollout_batch['input_data']) + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + model.forward_backward( + inputs=rollout_batch['input_data'][mb_start:mb_end], + old_logps=rollout_batch['old_logps'][mb_start:mb_end], + advantages=advantages[mb_start:mb_end], + micro_batch_size=MICRO_BATCH_SIZE, + ) + model.clip_grad_and_step() + optim_step += 1 + + if optim_step % SAVE_STEPS == 0: + model.save(f'math-grpo-group-admission-checkpoint-{optim_step}') + if optim_step >= MAX_STEPS: + break + + log_dict = metrics.calculate() + log_dict.update(admission_metrics) + log_dict.update(model.calculate_metric(is_training=True)) + metrics.reset() + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('math-grpo-group-admission-final') + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/grpo/group_admission.sh b/cookbook/rl/grpo/group_admission.sh new file mode 100644 index 000000000..e3a0bace3 --- /dev/null +++ b/cookbook/rl/grpo/group_admission.sh @@ -0,0 +1,28 @@ +#!/bin/sh +set -eu + +# End-to-end GSM8K GRPO with complete-group admission and bounded resampling. +# Keep max-steps small for a smoke test, for example: +# CUDA_VISIBLE_DEVICES=0,1 sh group_admission.sh \ +# --model-gpus 1 --sampler-gpus 1 --max-steps 2 + +SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +REPO_ROOT=$(CDPATH= cd -- "$SCRIPT_DIR/../../.." && pwd) +export PYTHONPATH="$REPO_ROOT/src${PYTHONPATH:+:$PYTHONPATH}" + +python "$SCRIPT_DIR/group_admission.py" \ + --model-id ms://Qwen/Qwen3.5-4B \ + --strategy native_fsdp \ + --model-gpus 4 \ + --sampler-gpus 4 \ + --num-generations 8 \ + --max-tokens 4096 \ + --batch-size 8 \ + --mini-batch-size 8 \ + --micro-batch-size 2 \ + --max-steps 1000 \ + --lr 1e-5 \ + --lora-r 16 \ + --save-steps 1000 \ + --adapter-name default \ + "$@" diff --git a/docs/source_en/Components/Advantage/GroupAdmissionPolicy.md b/docs/source_en/Components/Advantage/GroupAdmissionPolicy.md new file mode 100644 index 000000000..665e9fcad --- /dev/null +++ b/docs/source_en/Components/Advantage/GroupAdmissionPolicy.md @@ -0,0 +1,152 @@ +# GroupAdmissionPolicy + +GRPO groups with identical rewards produce zero centered advantages. With +dense or composite rewards, a non-zero difference can also be smaller than the +reward's meaningful resolution. Normalizing that difference may amplify noise +into a strong relative signal. `GroupAdmissionPolicy` separates this quality +decision from sampling infrastructure and always admits or rejects a complete +prompt group. Bounded replacement sampling can then recover useful groups +without an unbounded retry loop. + +This component is intentionally imported from its dedicated submodule rather +than re-exported by `twinkle.advantage`: it applies to group-relative methods +such as GRPO and DAPO-style training, not to every advantage estimator. + +## Signals + +The policy evaluates two optional conditions: + +- The reward condition checks population reward standard deviation, reward + range, and the number of non-trivial centered advantages. + `min_reward_range` is the application-defined reward resolution. +- The optional near-duplicate condition checks whether string representations + are excessively similar. Its token n-gram Jaccard default is a + dependency-free lexical baseline, not a semantic-equivalence test. Callers + can inject a custom scorer when needed. + +All thresholds are disabled by default, so the default policy admits every +valid group and does not change existing training behavior. + +```python +from twinkle.advantage.group_admission import ( + GroupAdmissionConfig, + GroupAdmissionPolicy, +) + +policy = GroupAdmissionPolicy( + GroupAdmissionConfig( + min_reward_std=1e-6, + min_reward_range=0.02, + min_nontrivial_advantages=2, + max_mean_pairwise_similarity=0.95, + ) +) + +decision = policy.evaluate( + rewards=[0.0, 1.0, 0.0, 1.0], + completions=["answer a", "answer b", "answer c", "answer d"], +) +if decision.admitted: + train_complete_group() +``` + +Rejected groups expose one mutually exclusive `primary_rejection_reason`: + +- `exact_dead`: the reward gate rejected an exactly equal-reward group; +- `near_tie`: rewards differ, but fail the configured dispersion or resolution + requirement; +- `redundant`: the reward signal passed, but the near-duplicate gate rejected + the group. + +External precondition failures take precedence and use `external`. Low-level +gate reasons remain available in `decision.reasons`. + +## Budget-aware resampling + +`SamplingBudgetController` tracks the effective-group rate with an exponential +moving average and estimates how many replacement groups are needed. Hard caps +can be set for extra groups, resampling rounds, generated samples, and generated +tokens. + +```python +from twinkle.advantage.group_admission import ( + SamplingBudgetConfig, + SamplingBudgetController, + SamplingBudgetState, +) + +controller = SamplingBudgetController( + SamplingBudgetConfig( + max_extra_groups=16, + max_extra_groups_per_round=4, + max_resample_rounds=2, + max_total_samples=128, + ) +) +state = SamplingBudgetState() + +# `decisions` contains one decision per complete group from the current round. +controller.observe_round( + state, + decisions, + num_generations=4, + generated_tokens=round_output_tokens, +) +plan = controller.plan_resampling( + state, + target_groups=8, + num_generations=4, +) +sample_more_complete_groups(plan.extra_groups) +``` + +Resampling must use the same rollout policy snapshot as the initial candidates. +The controller only plans work; the caller remains responsible for policy +version and staleness checks. + +For a complete GSM8K rollout, admission, bounded-resampling, GRPO advantage, +and optimizer-step loop, run the end-to-end +[`cookbook/rl/grpo/group_admission.py`](https://github.com/modelscope/twinkle/blob/main/cookbook/rl/grpo/group_admission.py) +example. It follows the official `short_math_grpo.py` setup and keeps one +rollout policy snapshot across initial and replacement groups: + +```bash +sh cookbook/rl/grpo/group_admission.sh --max-steps 2 +``` + +Exact-dead and near-tie outcomes are stochastic in a real rollout; their +deterministic branch coverage remains in `tests/advantage/test_group_admission.py`. + +Use `group_admission_metrics(decisions)` to report effective-group rate, the +orthogonal reward/near-duplicate rejection counts, and the mutually exclusive +`exact_dead`, `near_tie`, and `redundant` counts. Derive ratios at the caller +from these raw counts so merged/distributed metrics retain a single +denominator. + +## Challenger integration + +`twinkle_agentic.challenger.Challenger` already refills its batch continuously +when a completed unit contributes no trainable group. Pass the same policy to +`AgenticChallenger` to apply resolution-aware admission at that existing +group-atomic boundary: + +```python +challenger = AgenticChallenger( + backend, + envs=envs, + group_admission_policy=policy, +) +``` + +The original equal-reward filter remains active when no policy is supplied. +When the optional near-duplicate gate is enabled, Challenger uses each +trajectory's final assistant text as its representation. The existing +continuous refill and `max_empty_rounds` behavior are unchanged. The separate +`SamplingBudgetController` remains available to sampling loops that need hard +sample or token budgets; Challenger does not silently relax a gate. + +## Non-goals + +The policy does not relax a configured gate when a budget is exhausted. The +caller should skip or fail the task according to its batch semantics. It also +does not change GRPO loss or normalized-advantage computation. diff --git a/docs/source_en/Components/Advantage/index.rst b/docs/source_en/Components/Advantage/index.rst index 9d38ea8bc..0e4508c50 100644 --- a/docs/source_en/Components/Advantage/index.rst +++ b/docs/source_en/Components/Advantage/index.rst @@ -5,4 +5,5 @@ Advantage Advantage.md GRPOAdvantage.md + GroupAdmissionPolicy.md RLOOAdvantage.md diff --git "a/docs/source_zh/\347\273\204\344\273\266/\344\274\230\345\212\277/GroupAdmissionPolicy.md" "b/docs/source_zh/\347\273\204\344\273\266/\344\274\230\345\212\277/GroupAdmissionPolicy.md" new file mode 100644 index 000000000..259666d8e --- /dev/null +++ "b/docs/source_zh/\347\273\204\344\273\266/\344\274\230\345\212\277/GroupAdmissionPolicy.md" @@ -0,0 +1,132 @@ +# GroupAdmissionPolicy + +GRPO 组内奖励完全相同时，中心化后的 advantage 全为零；对于稠密或组合奖励，非零 +差异也可能低于奖励的有效分辨率，标准化会把这种微小差异放大成较强的相对信号。 +`GroupAdmissionPolicy` 将质量判定与采样基础设施解耦，只对完整 prompt 组执行 +准入或拒绝。调用方可以在硬预算内补采有效组，避免无界重试。 + +本组件有意通过专用子模块导入，而不从 `twinkle.advantage` 根命名空间重新导出：它 +适用于 GRPO、DAPO-style 等 group-relative 方法，并不适用于所有 advantage 估计器。 + +## 判定信号 + +策略检查两类可选条件： + +- reward 条件检查奖励总体标准差、奖励极差，以及非平凡中心化 advantage 的数量； + `min_reward_range` 表示由场景定义的有效奖励分辨率； +- 可选的近重复条件检查字符串表示是否过度相似。默认的 token n-gram + Jaccard 只是零依赖的字面基线，不判断语义等价；有需要的场景可注入自定义 scorer。 + +所有阈值默认关闭，因此默认策略会放行全部合法组，不改变已有训练行为。 + +```python +from twinkle.advantage.group_admission import ( + GroupAdmissionConfig, + GroupAdmissionPolicy, +) + +policy = GroupAdmissionPolicy( + GroupAdmissionConfig( + min_reward_std=1e-6, + min_reward_range=0.02, + min_nontrivial_advantages=2, + max_mean_pairwise_similarity=0.95, + ) +) + +decision = policy.evaluate( + rewards=[0.0, 1.0, 0.0, 1.0], + completions=["答案 A", "答案 B", "答案 C", "答案 D"], +) +if decision.admitted: + train_complete_group() +``` + +被拒绝的组会暴露一个互斥的 `primary_rejection_reason`： + +- `exact_dead`：reward gate 拒绝了完全同分组； +- `near_tie`：reward 虽不完全相同，但未达到配置的离散度或有效分辨率； +- `redundant`：reward 信号有效，但近重复门拒绝该组。 + +外部前置条件失败具有最高优先级，统一标记为 `external`；各 Gate 的底层原因仍保留在 +`decision.reasons` 中。 + +## 预算感知补采 + +`SamplingBudgetController` 使用有效组率的指数移动平均估算下一轮需要补采多少组， +同时支持额外组数、补采轮数、生成样本数和生成 token 数硬上限。 + +```python +from twinkle.advantage.group_admission import ( + SamplingBudgetConfig, + SamplingBudgetController, + SamplingBudgetState, +) + +controller = SamplingBudgetController( + SamplingBudgetConfig( + max_extra_groups=16, + max_extra_groups_per_round=4, + max_resample_rounds=2, + max_total_samples=128, + ) +) +state = SamplingBudgetState() + +# decisions 中每个元素对应本轮的一个完整 GRPO 组。 +controller.observe_round( + state, + decisions, + num_generations=4, + generated_tokens=round_output_tokens, +) +plan = controller.plan_resampling( + state, + target_groups=8, + num_generations=4, +) +sample_more_complete_groups(plan.extra_groups) +``` + +补采必须使用与首批候选相同的 rollout policy 快照。控制器只负责规划补采量， +调用方仍需检查 policy version 和 staleness。 + +从 GSM8K rollout、整组准入、有界补采、GRPO advantage 到优化器更新的 +端到端流程，可参考 +[`cookbook/rl/grpo/group_admission.py`](https://github.com/modelscope/twinkle/blob/main/cookbook/rl/grpo/group_admission.py) +。该示例基于官方 `short_math_grpo.py` 配置，首轮与补采组共用同一个 +rollout policy 快照： + +```bash +sh cookbook/rl/grpo/group_admission.sh --max-steps 2 +``` + +真实 rollout 中 exact-dead 和 near-tie 的出现具有随机性；相关分支的确定性 +覆盖仍由 `tests/advantage/test_group_admission.py` 保证。 + +可以使用 `group_admission_metrics(decisions)` 记录有效组率、reward/近重复条件的独立 +拒绝数量，以及互斥的 `exact_dead`、`near_tie`、`redundant` 数量。比例由调用方 +基于原始计数计算，以便分布式合并时使用统一分母。 + +## Challenger 接入 + +`twinkle_agentic.challenger.Challenger` 已会在已完成单元没有产生可训练组时持续补充新单元。 +将策略传给 `AgenticChallenger`，即可在现有的 group-atomic 边界执行分辨率感知的准入： + +```python +challenger = AgenticChallenger( + backend, + envs=envs, + group_admission_policy=policy, +) +``` + +未传策略时，原有的同分组过滤保持不变。启用可选近重复门后，Challenger 使用每条 +轨迹最后一个 assistant 文本作为表示。现有连续补充和 `max_empty_rounds` 语义均不 +改变。需要样本数或 token 硬预算的其他采样循环仍可使用独立的 +`SamplingBudgetController`；Challenger 不会隐式放宽 Gate。 + +## 非目标 + +预算耗尽时不会隐式放宽任何 Gate；调用方根据自身 batch 语义跳过任务或失败退出。 +本组件也不修改 GRPO loss 和标准化 advantage 的计算。 diff --git "a/docs/source_zh/\347\273\204\344\273\266/\344\274\230\345\212\277/index.rst" "b/docs/source_zh/\347\273\204\344\273\266/\344\274\230\345\212\277/index.rst" index 5938286c4..f1522152d 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\344\274\230\345\212\277/index.rst" +++ "b/docs/source_zh/\347\273\204\344\273\266/\344\274\230\345\212\277/index.rst" @@ -5,4 +5,5 @@ Advantage.md GRPOAdvantage.md + GroupAdmissionPolicy.md RLOOAdvantage.md diff --git a/src/twinkle/advantage/group_admission.py b/src/twinkle/advantage/group_admission.py new file mode 100644 index 000000000..1a0dc0c95 --- /dev/null +++ b/src/twinkle/advantage/group_admission.py @@ -0,0 +1,396 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Infrastructure-agnostic admission policy for dynamic GRPO sampling. + +The policy turns the reward-signal observation used by GRPO into an explicit +sampling decision while preserving complete prompt groups. It intentionally +does not call a sampler or a trainer: synchronous loops, asynchronous workers, +and offline rollout pipelines can share the same deterministic policy. +""" + +from __future__ import annotations + +import math +import re +import statistics +from dataclasses import dataclass +from typing import Callable, Dict, Optional, Sequence, Set, Tuple, Union + +_SimilarityScorer = Callable[[str, str], float] + +__all__ = [ + 'GroupAdmissionConfig', + 'GroupAdmissionDecision', + 'GroupAdmissionPolicy', + 'ResamplePlan', + 'SamplingBudgetConfig', + 'SamplingBudgetController', + 'SamplingBudgetState', + 'group_admission_metrics', +] + +_REWARD_REJECTION_REASONS = frozenset({ + 'reward_std_below_threshold', + 'reward_range_below_threshold', + 'insufficient_nontrivial_advantages', +}) +_NEAR_DUPLICATE_REASON = 'mean_pairwise_similarity_above_threshold' + + +@dataclass(frozen=True) +class GroupAdmissionConfig: + """Thresholds for deciding whether one complete GRPO group is useful. + + ``min_reward_std``, ``min_reward_range``, and + ``min_nontrivial_advantages`` generalize the reward-variance gate + popularized by DAPO-style dynamic sampling. The range threshold lets + dense-reward applications reject differences below their meaningful + reward resolution. The optional similarity threshold adds a + deterministic near-duplicate signal. Defaults are deliberately disabled + so constructing the policy is a no-op. + """ + + min_reward_std: float = 0.0 + min_reward_range: float = 0.0 + min_nontrivial_advantages: int = 0 + advantage_tolerance: float = 1e-8 + max_mean_pairwise_similarity: float | None = None + similarity_ngram_size: int = 3 + + def __post_init__(self) -> None: + if not math.isfinite(self.min_reward_std) or self.min_reward_std < 0: + raise ValueError('min_reward_std must be finite and non-negative') + if not math.isfinite(self.min_reward_range) or self.min_reward_range < 0: + raise ValueError('min_reward_range must be finite and non-negative') + if self.min_nontrivial_advantages < 0: + raise ValueError('min_nontrivial_advantages must be non-negative') + if not math.isfinite(self.advantage_tolerance) or self.advantage_tolerance < 0: + raise ValueError('advantage_tolerance must be finite and non-negative') + threshold = self.max_mean_pairwise_similarity + if threshold is not None and (not math.isfinite(threshold) or not 0 <= threshold <= 1): + raise ValueError('max_mean_pairwise_similarity must be in [0, 1]') + if self.similarity_ngram_size <= 0: + raise ValueError('similarity_ngram_size must be positive') + + +@dataclass(frozen=True) +class GroupAdmissionDecision: + """Decision and auditable observations for one indivisible prompt group.""" + + reasons: tuple[str, ...] + reward_std: float + reward_range: float + nontrivial_advantage_count: int + mean_pairwise_similarity: float | None = None + max_pairwise_similarity: float | None = None + external_reasons: tuple[str, ...] = () + + @property + def admitted(self) -> bool: + return not self.reasons + + @property + def reward_rejected(self) -> bool: + internal_reasons = self.reasons[len(self.external_reasons):] + return any(reason in _REWARD_REJECTION_REASONS for reason in internal_reasons) + + @property + def near_duplicate_rejected(self) -> bool: + internal_reasons = self.reasons[len(self.external_reasons):] + return _NEAR_DUPLICATE_REASON in internal_reasons + + @property + def primary_rejection_reason(self) -> str | None: + """Return one stable, mutually exclusive reason for rejected groups.""" + if self.admitted: + return None + if self.external_reasons: + return 'external' + if self.reward_rejected: + return 'exact_dead' if self.reward_range == 0.0 else 'near_tie' + if self.near_duplicate_rejected: + return 'redundant' + return None + + +class GroupAdmissionPolicy: + """Admit complete groups with useful reward and representation signals.""" + + def __init__( + self, + config: GroupAdmissionConfig | None = None, + *, + scorer: _SimilarityScorer | None = None, + ): + self.config = config or GroupAdmissionConfig() + self._similarity_scorer = scorer + + def evaluate( + self, + rewards: Sequence[float], + completions: Sequence[str] | None = None, + *, + external_reasons: Sequence[str] = (), + ) -> GroupAdmissionDecision: + reward_values = _validated_rewards(rewards) + completion_values = list(completions) if completions is not None else None + if completion_values is not None and len(completion_values) != len(reward_values): + raise ValueError('rewards and completions must describe the same complete group') + + reward_mean = sum(reward_values) / len(reward_values) + reward_std = statistics.pstdev(reward_values) + reward_range = max(reward_values) - min(reward_values) + nontrivial = sum(abs(value - reward_mean) > self.config.advantage_tolerance for value in reward_values) + + reward_reasons = [] + if reward_std < self.config.min_reward_std: + reward_reasons.append('reward_std_below_threshold') + if (reward_range < self.config.min_reward_range and not math.isclose( + reward_range, + self.config.min_reward_range, + rel_tol=1e-12, + abs_tol=1e-12, + )): + reward_reasons.append('reward_range_below_threshold') + if nontrivial < self.config.min_nontrivial_advantages: + reward_reasons.append('insufficient_nontrivial_advantages') + + mean_similarity, max_similarity, similarity_reasons = self._evaluate_similarity(completion_values) + external = tuple(str(reason) for reason in external_reasons if str(reason)) + reasons = (*external, *reward_reasons, *similarity_reasons) + return GroupAdmissionDecision( + reasons=reasons, + reward_std=reward_std, + reward_range=reward_range, + nontrivial_advantage_count=nontrivial, + mean_pairwise_similarity=mean_similarity, + max_pairwise_similarity=max_similarity, + external_reasons=external, + ) + + def _evaluate_similarity( + self, + completions: Sequence[str] | None, + ) -> tuple[float | None, float | None, tuple[str, ...]]: + threshold = self.config.max_mean_pairwise_similarity + if threshold is None: + return None, None, () + if completions is None: + raise ValueError('completions are required when the near-duplicate gate is enabled') + values = list(completions) + if len(values) < 2: + raise ValueError('the near-duplicate gate requires at least two completions') + + scorer = self._similarity_scorer + scores = [] + for left_index, left in enumerate(values): + for right in values[left_index + 1:]: + score = float( + scorer(left, right) if scorer is not None else _ngram_jaccard( + left, + right, + ngram_size=self.config.similarity_ngram_size, + )) + if not math.isfinite(score) or not 0 <= score <= 1: + raise ValueError(f'similarity scorer must return a finite value in [0, 1], got {score}') + scores.append(score) + mean_similarity = sum(scores) / len(scores) + reasons = (_NEAR_DUPLICATE_REASON, ) if mean_similarity > threshold else () + return mean_similarity, max(scores), reasons + + +@dataclass(frozen=True) +class SamplingBudgetConfig: + """Hard caps and EMA parameters for allocating replacement groups.""" + + max_extra_groups: int = 0 + max_extra_groups_per_round: int | None = None + max_resample_rounds: int = 0 + max_total_samples: int | None = None + max_total_tokens: int | None = None + effective_rate_ema_alpha: float = 0.2 + min_effective_rate: float = 0.05 + + def __post_init__(self) -> None: + for name in ('max_extra_groups', 'max_resample_rounds'): + if getattr(self, name) < 0: + raise ValueError(f'{name} must be non-negative') + for name in ('max_extra_groups_per_round', 'max_total_samples', 'max_total_tokens'): + value = getattr(self, name) + if value is not None and value < 0: + raise ValueError(f'{name} must be non-negative when provided') + if not 0 < self.effective_rate_ema_alpha <= 1: + raise ValueError('effective_rate_ema_alpha must be in (0, 1]') + if not 0 < self.min_effective_rate <= 1: + raise ValueError('min_effective_rate must be in (0, 1]') + + +@dataclass +class SamplingBudgetState: + """Mutable, serializable state scoped to one sampling batch/partition.""" + + sampling_rounds: int = 0 + resample_rounds: int = 0 + generated_groups: int = 0 + generated_samples: int = 0 + generated_tokens: int = 0 + admitted_groups: int = 0 + rejected_groups: int = 0 + effective_rate_ema: float | None = None + + +@dataclass(frozen=True) +class ResamplePlan: + extra_groups: int + remaining_target_groups: int + effective_rate: float + exhausted: bool + limited_by: tuple[str, ...] = () + + +class SamplingBudgetController: + """Translate observed group effectiveness into bounded resampling work.""" + + def __init__(self, config: SamplingBudgetConfig | None = None): + self.config = config or SamplingBudgetConfig() + + def observe_round( + self, + state: SamplingBudgetState, + decisions: Sequence[GroupAdmissionDecision], + *, + num_generations: int, + generated_tokens: int = 0, + ) -> None: + if num_generations <= 0: + raise ValueError('num_generations must be positive') + if generated_tokens < 0: + raise ValueError('generated_tokens must be non-negative') + values = list(decisions) + if not values: + raise ValueError('at least one group decision is required') + admitted = sum(decision.admitted for decision in values) + rate = admitted / len(values) + if state.effective_rate_ema is None: + state.effective_rate_ema = rate + else: + alpha = self.config.effective_rate_ema_alpha + state.effective_rate_ema = alpha * rate + (1 - alpha) * state.effective_rate_ema + state.sampling_rounds += 1 + state.generated_groups += len(values) + state.generated_samples += len(values) * num_generations + state.generated_tokens += generated_tokens + state.admitted_groups += admitted + state.rejected_groups += len(values) - admitted + + def plan_resampling( + self, + state: SamplingBudgetState, + *, + target_groups: int, + num_generations: int, + estimated_tokens_per_group: float | None = None, + ) -> ResamplePlan: + if target_groups <= 0 or num_generations <= 0: + raise ValueError('target_groups and num_generations must be positive') + if estimated_tokens_per_group is not None and estimated_tokens_per_group <= 0: + raise ValueError('estimated_tokens_per_group must be positive when provided') + + remaining = max(0, target_groups - state.admitted_groups) + effective_rate = max( + self.config.min_effective_rate, + state.effective_rate_ema if state.effective_rate_ema is not None else 1.0, + ) + if remaining == 0: + return ResamplePlan(0, 0, effective_rate, False, ('target_met', )) + if state.resample_rounds >= self.config.max_resample_rounds: + return ResamplePlan(0, remaining, effective_rate, True, ('resample_round_budget', )) + + requested = max(remaining, math.ceil(remaining / effective_rate)) + limits: list[tuple[str, int]] = [ + ('extra_group_budget', target_groups + self.config.max_extra_groups - state.generated_groups), + ] + if self.config.max_extra_groups_per_round is not None: + limits.append(('per_round_group_budget', self.config.max_extra_groups_per_round)) + if self.config.max_total_samples is not None: + remaining_samples = self.config.max_total_samples - state.generated_samples + limits.append(('sample_budget', remaining_samples // num_generations)) + if self.config.max_total_tokens is not None: + token_estimate = estimated_tokens_per_group + if token_estimate is None and state.generated_groups and state.generated_tokens: + token_estimate = state.generated_tokens / state.generated_groups + if token_estimate is not None: + remaining_tokens = self.config.max_total_tokens - state.generated_tokens + limits.append(('token_budget', math.floor(remaining_tokens / token_estimate))) + + normalized_limits = [(name, max(0, value)) for name, value in limits] + extra_groups = min([requested, *(value for _, value in normalized_limits)]) + limited_by = tuple(name for name, value in normalized_limits if value <= requested and value == extra_groups) + exhausted = extra_groups == 0 + if extra_groups: + state.resample_rounds += 1 + return ResamplePlan(extra_groups, remaining, effective_rate, exhausted, limited_by) + + +def group_admission_metrics(decisions: Sequence[GroupAdmissionDecision]) -> dict[str, float | int]: + """Aggregate decisions using names suitable for Twinkle metric records.""" + values = list(decisions) + if not values: + return { + 'candidate_group_count': 0, + 'admitted_group_count': 0, + 'rejected_group_count': 0, + 'effective_group_rate': 0.0, + 'reward_rejected_group_count': 0, + 'near_duplicate_rejected_group_count': 0, + 'externally_rejected_group_count': 0, + 'exact_dead_group_count': 0, + 'near_tie_group_count': 0, + 'redundant_group_count': 0, + } + admitted = sum(decision.admitted for decision in values) + reward_rejected = sum(decision.reward_rejected for decision in values) + near_duplicate_rejected = sum(decision.near_duplicate_rejected for decision in values) + externally_rejected = sum(bool(decision.external_reasons) for decision in values) + primary_reasons = [decision.primary_rejection_reason for decision in values] + return { + 'candidate_group_count': len(values), + 'admitted_group_count': admitted, + 'rejected_group_count': len(values) - admitted, + 'effective_group_rate': admitted / len(values), + 'reward_rejected_group_count': reward_rejected, + 'near_duplicate_rejected_group_count': near_duplicate_rejected, + 'externally_rejected_group_count': externally_rejected, + 'exact_dead_group_count': primary_reasons.count('exact_dead'), + 'near_tie_group_count': primary_reasons.count('near_tie'), + 'redundant_group_count': primary_reasons.count('redundant'), + } + + +def _validated_rewards(rewards: Sequence[float]) -> list[float]: + values = [float(value) for value in rewards] + if len(values) < 2: + raise ValueError('a GRPO group requires at least two rewards') + if any(not math.isfinite(value) for value in values): + raise ValueError('group rewards must be finite') + return values + + +def _ngram_jaccard(left: str, right: str, *, ngram_size: int) -> float: + """Return lexical overlap for the default near-duplicate detector.""" + left_ngrams = _ngrams(left, ngram_size) + right_ngrams = _ngrams(right, ngram_size) + if not left_ngrams and not right_ngrams: + return 1.0 + union = left_ngrams | right_ngrams + return len(left_ngrams & right_ngrams) / len(union) if union else 1.0 + + +def _ngrams(text: str, ngram_size: int) -> set[tuple[str, ...]]: + if not isinstance(text, str): + raise TypeError(f'completion must be str, got {type(text)!r}') + tokens = re.findall(r'\w+|[^\w\s]', text.lower(), flags=re.UNICODE) + if not tokens: + return set() + if len(tokens) < ngram_size: + return {tuple(tokens)} + return {tuple(tokens[index:index + ngram_size]) for index in range(len(tokens) - ngram_size + 1)} diff --git a/src/twinkle_agentic/challenger/agentic.py b/src/twinkle_agentic/challenger/agentic.py index 23b4ce6dc..669e40019 100644 --- a/src/twinkle_agentic/challenger/agentic.py +++ b/src/twinkle_agentic/challenger/agentic.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple +from twinkle.advantage.group_admission import GroupAdmissionPolicy from twinkle.data_format import SamplingParams, Trajectory, attach_user_data, user_data_get from twinkle.data_format.sampling import SampledSequence, SampleResponse from twinkle.utils import get_logger @@ -191,6 +192,7 @@ def __init__( pass_band: Tuple[float, float] = (1.0, 7.0), pass_rate_width: float = 0.3, max_empty_rounds: int = 0, + group_admission_policy: Optional[GroupAdmissionPolicy] = None, followup_params: Optional[SamplingParams] = None, checker: Optional[Callable[[Trajectory], bool]] = None, save_dir: Optional[str] = None, @@ -203,6 +205,7 @@ def __init__( num_solver_rollouts=num_solver_rollouts, pass_band=pass_band, max_empty_rounds=max_empty_rounds, + group_admission_policy=group_admission_policy, ) if check_retries < 0: raise ValueError(f'check_retries must be >= 0, got {check_retries}') diff --git a/src/twinkle_agentic/challenger/base.py b/src/twinkle_agentic/challenger/base.py index 04b07bf32..3d5815239 100644 --- a/src/twinkle_agentic/challenger/base.py +++ b/src/twinkle_agentic/challenger/base.py @@ -3,13 +3,16 @@ import queue import threading from abc import ABC, abstractmethod +from collections import Counter from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from typing import Callable, Iterator, List, Optional, Sequence, Tuple +from twinkle.advantage.group_admission import GroupAdmissionPolicy from twinkle.data_format import Trajectory from twinkle.utils import get_logger from twinkle_agentic.envs import Env, EnvLeases +from twinkle_agentic.utils.message_utils import assistant_text logger = get_logger() @@ -56,13 +59,14 @@ class Challenger(ABC): """ def __init__( - self, - *, - envs: Sequence[Env], - num_challenger_rollouts: int = 8, - num_solver_rollouts: int = 8, - pass_band: Tuple[float, float] = (1.0, 7.0), - max_empty_rounds: int = 0, + self, + *, + envs: Sequence[Env], + num_challenger_rollouts: int = 8, + num_solver_rollouts: int = 8, + pass_band: Tuple[float, float] = (1.0, 7.0), + max_empty_rounds: int = 0, + group_admission_policy: Optional[GroupAdmissionPolicy] = None, ): if not envs: raise ValueError('envs is empty: a challenger needs a workspace to act in and grade') @@ -85,6 +89,7 @@ def __init__( self.num_solver_rollouts = num_solver_rollouts self.pass_band = pass_band self.max_empty_rounds = max_empty_rounds + self.group_admission_policy = group_admission_policy self.n_proposed = 0 self.n_kept = 0 # One worker per environment, which is what makes a lease never block: a @@ -127,13 +132,47 @@ def _complete(self, challenger: List[List[Trajectory]], solver: List[List[Trajec Called from the job that finished the unit, so it must not block: it drops the groups in a queue and returns to the pool. """ - proposing = [group for group in challenger if self._has_spread(group)] - solving = [group for group in solver if self._has_spread(group)] - flat = len(challenger) - len(proposing) + len(solver) - len(solving) - if flat: - logger.info(f'[{type(self).__name__}] dropped {flat} groups whose rewards were all equal') + reasons: Counter[str] = Counter() + proposing = self._admitted_groups(challenger, reasons) + solving = self._admitted_groups(solver, reasons) + if reasons: + detail = ', '.join(f'{reason}={count}' for reason, count in sorted(reasons.items())) + logger.info(f'[{type(self).__name__}] dropped {sum(reasons.values())} groups: {detail}') self._finished.put((proposing, solving)) + def _admitted_groups( + self, + groups: Sequence[List[Trajectory]], + rejection_reasons: Counter[str], + ) -> List[List[Trajectory]]: + """Apply admission atomically and return only complete accepted groups. + + Challenger historically drops groups with no reward spread. That remains + the default invariant; an optional shared policy can add resolution-aware + reward and near-duplicate checks without changing the refill lifecycle. + """ + admitted = [] + for group in groups: + if not self._has_spread(group): + rejection_reasons['exact_dead'] += 1 + continue + if self.group_admission_policy is None: + admitted.append(group) + continue + + completions = None + if self.group_admission_policy.config.max_mean_pairwise_similarity is not None: + completions = [assistant_text(member) for member in group] + decision = self.group_admission_policy.evaluate( + [float(member.get('rewards') or 0.0) for member in group], + completions=completions, + ) + if decision.admitted: + admitted.append(group) + else: + rejection_reasons[decision.primary_rejection_reason or 'policy'] += 1 + return admitted + @staticmethod def _has_spread(group: List[Trajectory]) -> bool: """True when a group's rewards differ. diff --git a/tests/advantage/test_group_admission.py b/tests/advantage/test_group_admission.py new file mode 100644 index 000000000..b7718a99a --- /dev/null +++ b/tests/advantage/test_group_admission.py @@ -0,0 +1,304 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +import pytest + +from twinkle.advantage import __all__ as advantage_exports +from twinkle.advantage.group_admission import (GroupAdmissionConfig, GroupAdmissionPolicy, SamplingBudgetConfig, + SamplingBudgetController, SamplingBudgetState, + group_admission_metrics) + + +def _decision(admitted: bool): + policy = GroupAdmissionPolicy( + GroupAdmissionConfig( + min_reward_std=0.1, + min_nontrivial_advantages=2, + )) + rewards = [0.0, 1.0] if admitted else [0.0, 0.0] + return policy.evaluate(rewards) + + +def test_group_admission_does_not_expand_the_advantage_root_api() -> None: + specialized_names = { + 'GroupAdmissionConfig', + 'GroupAdmissionPolicy', + 'SamplingBudgetConfig', + 'SamplingBudgetController', + } + assert specialized_names.isdisjoint(advantage_exports) + + +def test_default_policy_preserves_existing_behavior() -> None: + decision = GroupAdmissionPolicy().evaluate( + [0.0, 0.0, 0.0], + ['same completion', 'same completion', 'same completion'], + ) + assert decision.admitted + assert decision.reasons == () + assert decision.primary_rejection_reason is None + assert decision.reward_range == 0.0 + + +def test_reward_gate_rejects_zero_signal_group() -> None: + decision = GroupAdmissionPolicy( + GroupAdmissionConfig( + min_reward_std=1e-6, + min_nontrivial_advantages=2, + )).evaluate([1.0, 1.0, 1.0]) + assert not decision.admitted + assert decision.reward_std == 0.0 + assert decision.reward_range == 0.0 + assert decision.nontrivial_advantage_count == 0 + assert decision.primary_rejection_reason == 'exact_dead' + assert decision.reasons == ( + 'reward_std_below_threshold', + 'insufficient_nontrivial_advantages', + ) + + +def test_reward_range_rejects_nonzero_but_resolution_insignificant_group() -> None: + decision = GroupAdmissionPolicy( + GroupAdmissionConfig(min_reward_range=0.02)).evaluate([0.5, 0.5, 0.51]) + assert not decision.admitted + assert decision.reward_std > 0 + assert decision.reward_range == pytest.approx(0.01) + assert decision.reasons == ('reward_range_below_threshold', ) + assert decision.primary_rejection_reason == 'near_tie' + + +def test_reward_range_admits_difference_at_meaningful_resolution() -> None: + decision = GroupAdmissionPolicy( + GroupAdmissionConfig(min_reward_range=0.02)).evaluate([0.1, 0.1, 0.12]) + assert decision.admitted + assert decision.reward_range == pytest.approx(0.02) + assert decision.primary_rejection_reason is None + + +def test_zero_range_threshold_preserves_std_only_dapo_style_gate() -> None: + policy = GroupAdmissionPolicy( + GroupAdmissionConfig(min_reward_std=1e-6, min_reward_range=0.0)) + assert not policy.evaluate([0.5, 0.5, 0.5]).admitted + assert policy.evaluate([0.5, 0.5, 0.5001]).admitted + + +def test_reward_gate_admits_group_with_nontrivial_advantages() -> None: + decision = GroupAdmissionPolicy( + GroupAdmissionConfig( + min_reward_std=0.4, + min_nontrivial_advantages=2, + )).evaluate([0.0, 1.0, 0.0, 1.0]) + assert decision.admitted + assert decision.reward_std == pytest.approx(0.5) + assert decision.nontrivial_advantage_count == 4 + + +def test_diversity_gate_rejects_collective_near_duplicates() -> None: + policy = GroupAdmissionPolicy( + GroupAdmissionConfig(max_mean_pairwise_similarity=0.8, similarity_ngram_size=1)) + decision = policy.evaluate( + [0.0, 1.0, 0.0], + ['return price + tax', 'return price + tax', 'return price + tax'], + ) + assert not decision.admitted + assert decision.mean_pairwise_similarity == 1.0 + assert decision.reasons == ('mean_pairwise_similarity_above_threshold', ) + assert decision.primary_rejection_reason == 'redundant' + + +def test_diversity_gate_uses_group_mean_instead_of_one_duplicate_pair() -> None: + policy = GroupAdmissionPolicy( + GroupAdmissionConfig(max_mean_pairwise_similarity=0.5, similarity_ngram_size=1)) + decision = policy.evaluate( + [0.0, 1.0, 0.0], + ['alpha beta', 'alpha beta', 'gamma delta'], + ) + assert decision.admitted + assert decision.max_pairwise_similarity == 1.0 + assert decision.mean_pairwise_similarity == pytest.approx(1 / 3) + + +def test_custom_similarity_scorer_supports_domain_semantics() -> None: + + class ExecutionSignatureScorer: + + def __call__(self, left: str, right: str) -> float: + return 1.0 if left.split(':', 1)[0] == right.split(':', 1)[0] else 0.0 + + policy = GroupAdmissionPolicy( + GroupAdmissionConfig(max_mean_pairwise_similarity=0.6), + scorer=ExecutionSignatureScorer(), + ) + decision = policy.evaluate( + [0.0, 1.0, 0.0], + ['success:syntax-a', 'success:syntax-b', 'success:syntax-c'], + ) + assert not decision.admitted + + +def test_diversity_gate_requires_completions_only_when_enabled() -> None: + assert GroupAdmissionPolicy().evaluate([0.0, 1.0]).admitted + policy = GroupAdmissionPolicy(GroupAdmissionConfig(max_mean_pairwise_similarity=0.9)) + with pytest.raises(ValueError, match='completions are required'): + policy.evaluate([0.0, 1.0]) + + +def test_external_precondition_rejects_an_otherwise_useful_group() -> None: + decision = GroupAdmissionPolicy().evaluate( + [0.0, 1.0], + external_reasons=['missing_execution'], + ) + assert not decision.admitted + assert decision.reasons == ('missing_execution', ) + assert decision.primary_rejection_reason == 'external' + + +def test_external_reason_name_cannot_be_misclassified_as_a_reward_rejection() -> None: + decision = GroupAdmissionPolicy().evaluate( + [0.0, 1.0], + external_reasons=['reward_std_below_threshold'], + ) + assert decision.primary_rejection_reason == 'external' + assert not decision.reward_rejected + assert group_admission_metrics([decision])['reward_rejected_group_count'] == 0 + + +def test_default_near_duplicate_scorer_is_normalized_and_handles_empty_text() -> None: + policy = GroupAdmissionPolicy( + GroupAdmissionConfig(max_mean_pairwise_similarity=0.3, similarity_ngram_size=1)) + assert policy.evaluate([0.0, 1.0], ['', '']).mean_pairwise_similarity == 1.0 + assert policy.evaluate([0.0, 1.0], ['', 'value']).mean_pairwise_similarity == 0.0 + decision = policy.evaluate([0.0, 1.0], ['Alpha beta', 'alpha gamma']) + assert decision.mean_pairwise_similarity == pytest.approx(1 / 3) + + +def test_budget_controller_uses_effective_rate_to_plan_replacements() -> None: + controller = SamplingBudgetController( + SamplingBudgetConfig( + max_extra_groups=8, + max_resample_rounds=2, + effective_rate_ema_alpha=0.5, + )) + state = SamplingBudgetState() + controller.observe_round( + state, + [_decision(True), _decision(False), _decision(True), _decision(False)], + num_generations=3, + generated_tokens=120, + ) + plan = controller.plan_resampling(state, target_groups=4, num_generations=3) + assert plan.extra_groups == 4 + assert plan.remaining_target_groups == 2 + assert plan.effective_rate == 0.5 + assert state.resample_rounds == 1 + + +def test_budget_controller_updates_ema_between_rounds() -> None: + controller = SamplingBudgetController( + SamplingBudgetConfig( + max_extra_groups=8, + max_resample_rounds=2, + effective_rate_ema_alpha=0.5, + )) + state = SamplingBudgetState() + controller.observe_round(state, [_decision(True), _decision(False)], num_generations=2) + controller.plan_resampling(state, target_groups=3, num_generations=2) + controller.observe_round(state, [_decision(True), _decision(True)], num_generations=2) + assert state.effective_rate_ema == pytest.approx(0.75) + plan = controller.plan_resampling(state, target_groups=3, num_generations=2) + assert plan.extra_groups == 0 + assert plan.limited_by == ('target_met', ) + + +def test_budget_controller_enforces_sample_cap_at_group_boundary() -> None: + controller = SamplingBudgetController( + SamplingBudgetConfig( + max_extra_groups=10, + max_resample_rounds=3, + max_total_samples=18, + )) + state = SamplingBudgetState() + controller.observe_round(state, [_decision(False)] * 4, num_generations=3) + plan = controller.plan_resampling(state, target_groups=4, num_generations=3) + assert plan.extra_groups == 2 + assert plan.limited_by == ('sample_budget', ) + + +def test_budget_controller_enforces_token_and_per_round_caps() -> None: + controller = SamplingBudgetController( + SamplingBudgetConfig( + max_extra_groups=10, + max_extra_groups_per_round=3, + max_resample_rounds=3, + max_total_tokens=200, + )) + state = SamplingBudgetState() + controller.observe_round( + state, + [_decision(False)] * 4, + num_generations=2, + generated_tokens=120, + ) + plan = controller.plan_resampling( + state, + target_groups=4, + num_generations=2, + estimated_tokens_per_group=40, + ) + assert plan.extra_groups == 2 + assert plan.limited_by == ('token_budget', ) + + +def test_default_budget_disables_resampling() -> None: + controller = SamplingBudgetController() + state = SamplingBudgetState() + controller.observe_round(state, [_decision(False)], num_generations=4) + plan = controller.plan_resampling(state, target_groups=1, num_generations=4) + assert plan.extra_groups == 0 + assert plan.exhausted + assert plan.limited_by == ('resample_round_budget', ) + + +def test_admission_metrics_report_both_rejection_signals() -> None: + reward_reject = _decision(False) + near_tie_reject = GroupAdmissionPolicy( + GroupAdmissionConfig(min_reward_range=0.02)).evaluate([0.5, 0.51]) + diversity_reject = GroupAdmissionPolicy( + GroupAdmissionConfig(max_mean_pairwise_similarity=0.5, similarity_ngram_size=1)).evaluate( + [0.0, 1.0], ['same', 'same']) + metrics = group_admission_metrics( + [_decision(True), reward_reject, near_tie_reject, diversity_reject]) + assert metrics == { + 'candidate_group_count': 4, + 'admitted_group_count': 1, + 'rejected_group_count': 3, + 'effective_group_rate': pytest.approx(1 / 4), + 'reward_rejected_group_count': 2, + 'near_duplicate_rejected_group_count': 1, + 'externally_rejected_group_count': 0, + 'exact_dead_group_count': 1, + 'near_tie_group_count': 1, + 'redundant_group_count': 1, + } + + +def test_empty_admission_metrics_keep_new_counts_zero() -> None: + metrics = group_admission_metrics([]) + assert metrics['candidate_group_count'] == 0 + assert metrics['exact_dead_group_count'] == 0 + assert metrics['near_tie_group_count'] == 0 + assert metrics['redundant_group_count'] == 0 + + +def test_invalid_configs_and_group_inputs_fail_early() -> None: + with pytest.raises(ValueError, match='min_reward_std'): + GroupAdmissionConfig(min_reward_std=-1) + with pytest.raises(ValueError, match='min_reward_range'): + GroupAdmissionConfig(min_reward_range=-1) + with pytest.raises(ValueError, match='max_mean_pairwise_similarity'): + GroupAdmissionConfig(max_mean_pairwise_similarity=1.1) + with pytest.raises(ValueError, match='at least two rewards'): + GroupAdmissionPolicy().evaluate([1.0]) + with pytest.raises(ValueError, match='finite'): + GroupAdmissionPolicy().evaluate([0.0, float('nan')]) + with pytest.raises(ValueError, match='same complete group'): + GroupAdmissionPolicy().evaluate([0.0, 1.0], ['only one completion']) diff --git a/tests/twinkle_agentic/test_challenger_group_admission.py b/tests/twinkle_agentic/test_challenger_group_admission.py new file mode 100644 index 000000000..f1731bc0b --- /dev/null +++ b/tests/twinkle_agentic/test_challenger_group_admission.py @@ -0,0 +1,88 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""CPU-only coverage for group admission in Challenger's refill boundary.""" + +from collections import Counter + +from twinkle.advantage.group_admission import GroupAdmissionConfig, GroupAdmissionPolicy +from twinkle_agentic.challenger.base import Challenger +from twinkle_agentic.envs.base import Env, StepResult + + +class _Env(Env): + + def step(self, tool_name, arguments): + return StepResult() + + +class _Challenger(Challenger): + + def _launch(self): + return False + + +def _trajectory(reward, answer='answer'): + return { + 'rewards': reward, + 'messages': [{'role': 'assistant', 'content': answer}], + } + + +def _admit(challenger, group): + reasons = Counter() + admitted = challenger._admitted_groups([group], reasons) + return admitted, reasons + + +def test_default_keeps_existing_exact_dead_filter(): + challenger = _Challenger(envs=[_Env()]) + try: + admitted, reasons = _admit( + challenger, + [_trajectory(0.5), _trajectory(0.5), _trajectory(0.5)], + ) + assert admitted == [] + assert reasons == {'exact_dead': 1} + finally: + challenger.close() + + +def test_default_accepts_a_group_with_reward_spread(): + challenger = _Challenger(envs=[_Env()]) + group = [_trajectory(0.2), _trajectory(0.4), _trajectory(0.6)] + try: + admitted, reasons = _admit(challenger, group) + assert admitted == [group] + assert not reasons + finally: + challenger.close() + + +def test_resolution_aware_policy_rejects_near_tie_atomically(): + policy = GroupAdmissionPolicy(GroupAdmissionConfig(min_reward_range=0.02)) + challenger = _Challenger(envs=[_Env()], group_admission_policy=policy) + try: + admitted, reasons = _admit( + challenger, + [_trajectory(0.500), _trajectory(0.505), _trajectory(0.510)], + ) + assert admitted == [] + assert reasons == {'near_tie': 1} + finally: + challenger.close() + + +def test_near_duplicate_gate_reads_assistant_completions(): + policy = GroupAdmissionPolicy( + GroupAdmissionConfig(max_mean_pairwise_similarity=0.9), + scorer=lambda left, right: 1.0 if left == right else 0.0, + ) + challenger = _Challenger(envs=[_Env()], group_admission_policy=policy) + try: + admitted, reasons = _admit( + challenger, + [_trajectory(0.0, 'same'), _trajectory(0.5, 'same'), _trajectory(1.0, 'same')], + ) + assert admitted == [] + assert reasons == {'redundant': 1} + finally: + challenger.close()