From e664e0bc384d73fdd500e4da494915332e16b60d Mon Sep 17 00:00:00 2001 From: vx120 <893600387@qq.com> Date: Tue, 28 Jul 2026 20:10:46 +0800 Subject: [PATCH 1/8] add mask replay code and the cookbook Signed-off-by: vx120 <893600387@qq.com> --- cookbook/rl/grpo/grpo_sampling_replay.py | 330 ++++++++++++++++++ cookbook/rl/grpo/grpo_sampling_replay.sh | 19 + src/twinkle/data_format/__init__.py | 2 +- src/twinkle/data_format/sampling.py | 8 + src/twinkle/loss/grpo.py | 16 + src/twinkle/metric/__init__.py | 1 + src/twinkle/metric/grpo.py | 21 +- src/twinkle/metric/rollout.py | 88 +++++ .../model/transformers/transformers.py | 48 ++- .../sampler/vllm_sampler/vllm_engine.py | 68 +++- .../sampler/vllm_sampler/vllm_sampler.py | 1 + src/twinkle/utils/__init__.py | 5 +- src/twinkle/utils/nccl_safe.py | 1 + src/twinkle/utils/torch_utils.py | 116 ++++++ 14 files changed, 708 insertions(+), 16 deletions(-) create mode 100644 cookbook/rl/grpo/grpo_sampling_replay.py create mode 100644 cookbook/rl/grpo/grpo_sampling_replay.sh create mode 100644 src/twinkle/metric/rollout.py diff --git a/cookbook/rl/grpo/grpo_sampling_replay.py b/cookbook/rl/grpo/grpo_sampling_replay.py new file mode 100644 index 000000000..2f8a4a41f --- /dev/null +++ b/cookbook/rl/grpo/grpo_sampling_replay.py @@ -0,0 +1,330 @@ +import os +import time +from typing import List, Tuple, Dict, Any + +from peft import LoraConfig + +import twinkle +from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger +from twinkle.advantage import GRPOAdvantage +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.model import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.reward import GSM8KAccuracyReward, GSM8KFormatReward +from twinkle.sampler import vLLMSampler +from twinkle.metric import ( + CompletionRewardMetric, + compute_grpo_rollout_metrics, +) +from twinkle.preprocessor.llm import GSM8KProcessor + +logger = get_logger() +args = CLI.from_args() + +MODEL_ID = args.model.model_id or 'ms://Qwen/Qwen3.5-4B' +USE_MEGATRON = args.model.strategy != 'native_fsdp' +# This entry point is exclusively for sampling-distribution replay. +ENABLE_SAMPLING_REPLAY = True + +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 200 +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 50 +LOGPROBS_MODE = ( + 'processed_logprobs' + if ENABLE_SAMPLING_REPLAY + else os.getenv('TWINKLE_LOGPROBS_MODE', 'processed_logprobs') +) + +if ENABLE_SAMPLING_REPLAY and USE_MEGATRON: + raise ValueError('Sampling replay currently requires --strategy native_fsdp') + +def create_gsm8k_dataset(): + dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train')) + dataset.set_template('Qwen3_5Template', model_id=MODEL_ID, max_length=400) + dataset.map(GSM8KProcessor()) + dataset.encode(add_generation_prompt=True) + return dataset + +def compute_rewards( + trajectories: List[Dict[str, Any]], +) -> Tuple[List[float], List[float], List[float]]: + accuracy_reward_fn = GSM8KAccuracyReward() + format_reward_fn = GSM8KFormatReward() + + accuracy_rewards = accuracy_reward_fn(trajectories) + format_rewards = format_reward_fn(trajectories) + total_rewards = [a + f for a, f in zip(accuracy_rewards, format_rewards)] + return total_rewards, format_rewards, accuracy_rewards + + +def extract_rollout_batch(sample_responses, *, require_sampling_masks: bool): + """Flatten sampler responses into aligned lists used by reward and training.""" + rollout_batch = { + 'input_data': [], + 'old_logps': [], + 'sampling_masks': [], + 'completion_lengths': [], + 'stop_reasons': [], + } + for sample_response in sample_responses: + for sequence in sample_response.sequences: + if sequence.logprobs is None: + raise RuntimeError('A sampled sequence is missing token log probabilities') + if require_sampling_masks and sequence.sampling_mask is None: + raise RuntimeError( + 'Sampling replay is enabled but a sampled sequence has no sampling mask') + rollout_batch['input_data'].append(sequence.new_input_feature) + rollout_batch['old_logps'].append( + [logprob[0][1] for logprob in sequence.logprobs]) + rollout_batch['sampling_masks'].append(sequence.sampling_mask) + rollout_batch['completion_lengths'].append(len(sequence.tokens)) + rollout_batch['stop_reasons'].append(sequence.stop_reason) + return rollout_batch + + +def main(): + # set sampler and model separate to use different gpus + 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'), + ] + if USE_MEGATRON: + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + else: + 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=32, lora_alpha=64, lora_dropout=0.05) + # Since we are training on text-only data, we avoid using 'all-linear' which would include the ViT layers. + lora_config = LoraConfig( + target_modules=[ + 'q_proj', 'k_proj', 'v_proj', 'o_proj', + 'gate_proj', 'up_proj', 'down_proj', + 'in_proj_qkv', 'in_proj_z', 'in_proj_a', 'in_proj_b', 'out_proj', + ], + r=32, lora_alpha=64, lora_dropout=0.0, + ) + 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') + else: + from transformers import Qwen3_5ForConditionalGeneration + model = TransformersModel( + model_id=MODEL_ID, + model_cls=Qwen3_5ForConditionalGeneration, + device_mesh=model_mesh, + remote_group='model', + ) + + model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=1) + 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, + beta=0.0, + entropy_coef=0.0, + enable_sampling_replay=ENABLE_SAMPLING_REPLAY, + ) + model.set_processor(InputProcessor) + model.set_template('Qwen3_5Template', model_id=MODEL_ID) + + sampler = vLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.8, + 'max_model_len': 4496, + 'max_lora_rank': 32, # save as lora_config + # NOTE: To use enable_lora with qwen3.5, ensure vLLM includes + # PR https://github.com/vllm-project/vllm/pull/36976 + # enable_lora=True used with ckpt_manager.sync_weights(merge_and_sync=False) + # meaning only sync lora weights, if merge_and_sync=True, + # lora will be merged into the base model and sync all weights to vLLM + 'enable_lora': True, + 'enable_sampling_replay': ENABLE_SAMPLING_REPLAY, + 'logprobs_mode': LOGPROBS_MODE, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID) + + 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 if ENABLE_SAMPLING_REPLAY else 1.0, + top_k=-1, + repetition_penalty=1.0, + ) + if ENABLE_SAMPLING_REPLAY: + model.add_metric( + 'GRPOMetric', + is_training=True, + temperature=sampling_params.temperature, + epsilon=0.2, + ) + logger.info( + 'Sampling replay enabled: model_runner=v2, logprobs_mode=processed_logprobs, ' + 'temperature=%s, top_p=%s, top_k=%s', + sampling_params.temperature, + sampling_params.top_p, + sampling_params.top_k, + ) + + optim_step = 0 + sampling_replay_stats_logged = False + logger.info(get_device_placement()) + + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + metrics.reset() + global_prompts = batch if isinstance(batch, list) else [batch] + # enable_lora=True used with ckpt_manager.sync_weights(merge_and_sync=False) + # meaning only sync lora weights, if merge_and_sync=True, + # lora will be merged into the base model and sync all weights to vLLM + weight_sync_started = time.perf_counter() + ckpt_manager.sync_weights(merge_and_sync=False) + weight_sync_seconds = time.perf_counter() - weight_sync_started + sampler.reset_prefix_cache() + def sample_prompt_groups(prompts): + expand_prompts = [] + for prompt in prompts: + expand_prompts.extend([prompt] * NUM_GENERATIONS) + started = time.perf_counter() + responses = sampler.sample(expand_prompts, sampling_params) + elapsed = time.perf_counter() - started + return extract_rollout_batch( + responses, + require_sampling_masks=ENABLE_SAMPLING_REPLAY, + ), elapsed + + rollout_batch, sampling_seconds = sample_prompt_groups(global_prompts) + sampled_tokens_total = sum(rollout_batch['completion_lengths']) + # Match the original GRPO control flow: every sampled rollout is scored, + # logged, and trained. Zero-variance groups keep their zero advantages; + # they are diagnosed below but never resampled, dropped, or skipped. + total_rewards, format_rewards, accuracy_rewards = compute_rewards( + rollout_batch['input_data']) + + all_input_data: List[Dict[str, Any]] = rollout_batch['input_data'] + all_old_logps: List[List[float]] = rollout_batch['old_logps'] + all_sampling_masks = rollout_batch['sampling_masks'] + all_completion_lengths: List[int] = rollout_batch['completion_lengths'] + all_stop_reasons = rollout_batch['stop_reasons'] + metrics.accumulate( + completion_lengths=all_completion_lengths, + rewards={ + 'total': total_rewards, + 'format': format_rewards, + 'accuracy': accuracy_rewards, + }, + ) + rollout_reward_log_dict = metrics.calculate() + + advantages = advantage_fn(total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + rollout_log_dict = compute_grpo_rollout_metrics( + completion_lengths=all_completion_lengths, + stop_reasons=all_stop_reasons, + rewards=total_rewards, + advantages=advantages, + num_generations=NUM_GENERATIONS, + sampling_masks=all_sampling_masks if ENABLE_SAMPLING_REPLAY else None, + ) + num_rollout_tokens = sum(all_completion_lengths) + rollout_log_dict['profiling/weight_sync_seconds'] = weight_sync_seconds + rollout_log_dict['profiling/sampling_seconds'] = sampling_seconds + rollout_log_dict['profiling/sampling_tokens_per_second'] = ( + sampled_tokens_total / sampling_seconds if sampling_seconds else 0.0) + rollout_log_dict['profiling/sampling_generated_tokens'] = sampled_tokens_total + + # Split completions into mini-batches and run one optim step per mini-batch. + total_completions = len(all_input_data) + for mb_start in range(0, total_completions, MINI_BATCH_SIZE): + mb_end = min(mb_start + MINI_BATCH_SIZE, total_completions) + mb_inputs = all_input_data[mb_start:mb_end] + mb_old_logps = all_old_logps[mb_start:mb_end] + mb_advantages = advantages[mb_start:mb_end] + replay_kwargs = {} + if ENABLE_SAMPLING_REPLAY: + replay_kwargs = { + 'sampling_masks': all_sampling_masks[mb_start:mb_end], + 'temperature': sampling_params.temperature, + } + + training_started = time.perf_counter() + model.forward_backward( + inputs=mb_inputs, + old_logps=mb_old_logps, + advantages=mb_advantages, + micro_batch_size=MICRO_BATCH_SIZE, + **replay_kwargs, + ) + model.clip_grad_and_step() + training_seconds = time.perf_counter() - training_started + if ENABLE_SAMPLING_REPLAY and not sampling_replay_stats_logged: + logger.info( + 'Sampling replay active: sequences=%d, tokens=%d, mean_kept_tokens=%.2f', + len(all_sampling_masks), + num_rollout_tokens, + rollout_log_dict['replay/support_size_mean'], + ) + sampling_replay_stats_logged = True + optim_step += 1 + + if optim_step % SAVE_STEPS == 0: + model.save(f'grpo-gsm8k-checkpoint-{optim_step}') + # Copy the rollout reward into every optimizer-step log line. A + # rollout can span multiple mini-batches, but no Step lacks reward. + log_dict = dict(rollout_reward_log_dict) + log_dict.update(model.calculate_metric(is_training=True)) + if mb_start == 0: + log_dict.update(rollout_log_dict) + num_training_tokens = sum(all_completion_lengths[mb_start:mb_end]) + log_dict['profiling/training_seconds'] = training_seconds + log_dict['profiling/training_completion_tokens_per_second'] = ( + num_training_tokens / training_seconds if training_seconds else 0.0) + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') + if optim_step >= MAX_STEPS: + break + + logger.info(f'Training completed. optim_steps={optim_step}') + model.save('grpo-gsm8k-checkpoint') + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/grpo/grpo_sampling_replay.sh b/cookbook/rl/grpo/grpo_sampling_replay.sh new file mode 100644 index 000000000..f1120ecb0 --- /dev/null +++ b/cookbook/rl/grpo/grpo_sampling_replay.sh @@ -0,0 +1,19 @@ +#!/bin/sh +set -eu + +# Sampling-distribution replay example. +python grpo_sampling_replay.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 200 \ + --lr 1e-5 \ + --save-steps 50 \ + --adapter-name default \ + "$@" diff --git a/src/twinkle/data_format/__init__.py b/src/twinkle/data_format/__init__.py index c93bebd2d..1dff273c7 100644 --- a/src/twinkle/data_format/__init__.py +++ b/src/twinkle/data_format/__init__.py @@ -2,5 +2,5 @@ from .input_feature import InputFeature from .message import Message, Tool, ToolCall from .output import LossOutput, ModelOutput -from .sampling import SampledSequence, SampleResponse, SamplingParams +from .sampling import SampledSequence, SampleResponse, SamplingMask, SamplingParams from .trajectory import Trajectory, pack_value, user_data_get diff --git a/src/twinkle/data_format/sampling.py b/src/twinkle/data_format/sampling.py index 05ecdd641..cdd2233a8 100644 --- a/src/twinkle/data_format/sampling.py +++ b/src/twinkle/data_format/sampling.py @@ -166,6 +166,13 @@ def from_dict(cls, d: Dict[str, Any]) -> 'SamplingParams': return cls(**filtered) +@dataclass +class SamplingMask: + """CSR token support sets aligned with sampled sequence tokens.""" + token_ids: List[int] + offsets: List[int] + + @dataclass class SampledSequence: """A single sampled sequence with tokens and logprobs.""" @@ -175,6 +182,7 @@ class SampledSequence: decoded: str = None new_input_feature: InputFeature = None routed_experts: Optional[Any] = None + sampling_mask: Optional[SamplingMask] = None @dataclass diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 781b22060..36970636d 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -32,12 +32,20 @@ def __init__( beta: float = 0.0, entropy_coef: float = 0.0, ignore_index: int = -100, + enable_sampling_replay: bool = False, **kwargs, ): self.epsilon = epsilon self.epsilon_high = epsilon_high if epsilon_high is not None else epsilon self.beta = beta self.entropy_coef = entropy_coef + self.enable_sampling_replay = enable_sampling_replay + if enable_sampling_replay and self.__class__ is not GRPOLoss: + raise ValueError('sampling replay is only supported by GRPOLoss') + if enable_sampling_replay and beta != 0.0: + raise ValueError('sampling replay does not support a GRPO KL penalty (beta must be 0)') + if enable_sampling_replay and entropy_coef != 0.0: + raise ValueError('sampling replay does not support a GRPO entropy bonus') # Gate the expensive entropy compute path in the model forward. self.require_entropy = entropy_coef > 0.0 self.ignore_index = ignore_index @@ -201,6 +209,7 @@ def __call__( old_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, ref_logps: Optional['torch.Tensor'] = None, advantages: Optional[Union['torch.Tensor', List[float], np.ndarray]] = None, + sampling_masks=None, **kwargs, ): """ @@ -222,6 +231,11 @@ def __call__( **kwargs: Additional arguments """ import torch + if self.enable_sampling_replay: + if sampling_masks is None: + raise ValueError('sampling_masks are required when sampling replay is enabled') + if old_logps is None: + raise ValueError('old_logps are required when sampling replay is enabled') labels = inputs.get('labels') assert labels is not None, "inputs must contain 'labels'" if not torch.is_tensor(labels): @@ -230,6 +244,8 @@ def __call__( labels = labels.unsqueeze(0) logps = outputs.get('logps') + if self.enable_sampling_replay and logps is None: + raise RuntimeError('sampling replay logps must be computed by the model forward') loss_mask = (labels != self.ignore_index).bool() if logps is None: logits = outputs.get('logits') diff --git a/src/twinkle/metric/__init__.py b/src/twinkle/metric/__init__.py index baeb6c1c9..cd7d8c99d 100644 --- a/src/twinkle/metric/__init__.py +++ b/src/twinkle/metric/__init__.py @@ -6,4 +6,5 @@ from .embedding import EmbeddingMetric from .grpo import CISPOMetric, GRPOMetric, GSPOMetric from .loss import LossMetric +from .rollout import compute_grpo_rollout_metrics, zero_variance_reward_group_indices from .train_metric import TrainMetric diff --git a/src/twinkle/metric/grpo.py b/src/twinkle/metric/grpo.py index bd85aab67..71fde1de4 100644 --- a/src/twinkle/metric/grpo.py +++ b/src/twinkle/metric/grpo.py @@ -1,6 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import math -from typing import Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from twinkle.data_format import InputFeature, ModelOutput from twinkle.utils import get_logger @@ -9,6 +9,9 @@ logger = get_logger() +if TYPE_CHECKING: + import torch + class GRPOMetric(Metric): @@ -41,6 +44,8 @@ def reset(self): self.sum_new: float = 0.0 self.sum_old: float = 0.0 self.sum_diff: float = 0.0 + self.sum_diff_sq: float = 0.0 + self.sum_ratio: float = 0.0 self.sum_approx_kl: float = 0.0 self.max_token_kl: float = 0.0 self.max_token_ratio: float = 0.0 @@ -185,13 +190,16 @@ def _accumulate_mb( old_f = old_f * scale d = logps_f - old_f # new - old + ratio = torch.exp(d) self.sum_old += float((old_f * mask_f).sum().item()) self.sum_diff += float((d * mask_f).sum().item()) + self.sum_diff_sq += float((d.square() * mask_f).sum().item()) + self.sum_ratio += float((ratio * mask_f).sum().item()) # Schulman K3 estimator of KL(old || new): # samples x ~ old, r(x) = new(x) / old(x), # k3 = r - 1 - log(r) = exp(new - old) - (new - old) - 1. - kl = torch.exp(d) - d - 1.0 + kl = ratio - d - 1.0 kl_masked = kl * mask_f self.sum_approx_kl += float(kl_masked.sum().item()) # Per-token extremes for collapse detection @@ -200,7 +208,7 @@ def _accumulate_mb( if cur_max_kl > self.max_token_kl: self.max_token_kl = cur_max_kl # Track ratio extremes - ratio_masked = torch.exp(d) * mask_f + ratio_masked = ratio * mask_f cur_max_ratio = float(ratio_masked.max().item()) if cur_max_ratio > self.max_token_ratio: self.max_token_ratio = cur_max_ratio @@ -311,11 +319,12 @@ def accumulate( cursor += advanced def calculate(self) -> Dict[str, Any]: - import torch local = [{ 'sum_new': self.sum_new, 'sum_old': self.sum_old, 'sum_diff': self.sum_diff, + 'sum_diff_sq': self.sum_diff_sq, + 'sum_ratio': self.sum_ratio, 'sum_kl': self.sum_approx_kl, 'max_token_kl': self.max_token_kl, 'max_token_ratio': self.max_token_ratio, @@ -344,11 +353,15 @@ def calculate(self) -> Dict[str, Any]: if any(r['has_old'] for r in all_results): mean_old = sum(r['sum_old'] for r in all_results) / n_total mean_diff = sum(r['sum_diff'] for r in all_results) / n_total + mean_diff_sq = sum(r['sum_diff_sq'] for r in all_results) / n_total + mean_ratio = sum(r['sum_ratio'] for r in all_results) / n_total mean_kl = sum(r['sum_kl'] for r in all_results) / n_total global_max_kl = max(r['max_token_kl'] for r in all_results) global_max_ratio = max(r['max_token_ratio'] for r in all_results) results['train/mean_old_logp'] = mean_old results['train/logp_diff_mean'] = mean_diff + results['train/logp_diff_std'] = math.sqrt(max(mean_diff_sq - mean_diff**2, 0.0)) + results['train/importance_ratio_mean'] = mean_ratio results['train/approx_kl'] = mean_kl results['train/token_kl_max'] = global_max_kl results['train/token_ratio_max'] = global_max_ratio diff --git a/src/twinkle/metric/rollout.py b/src/twinkle/metric/rollout.py new file mode 100644 index 000000000..cc261ccd9 --- /dev/null +++ b/src/twinkle/metric/rollout.py @@ -0,0 +1,88 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import Any, Dict, Optional, Sequence + +import numpy as np + + +def zero_variance_reward_group_indices( + rewards: Sequence[float], + num_generations: int, +) -> list[int]: + """Return GRPO group indices that cannot produce a relative advantage.""" + if num_generations <= 0: + raise ValueError('num_generations must be positive') + if len(rewards) % num_generations != 0: + raise ValueError('rewards must form complete num_generations groups') + if len(rewards) == 0: + return [] + + grouped_rewards = np.asarray(rewards, dtype=np.float64).reshape(-1, num_generations) + group_ranges = np.ptp(grouped_rewards, axis=1) + return np.flatnonzero(np.isclose(group_ranges, 0.0)).astype(int).tolist() + + +def compute_grpo_rollout_metrics( + *, + completion_lengths: Sequence[int], + stop_reasons: Sequence[str], + rewards: Sequence[float], + advantages: Sequence[float], + num_generations: int, + sampling_masks: Optional[Sequence[Any]] = None, +) -> Dict[str, float]: + """Reduce one GRPO rollout batch into scalar diagnostics.""" + if len(stop_reasons) != len(completion_lengths): + raise ValueError('stop_reasons must align with completion_lengths') + if len(rewards) != len(completion_lengths): + raise ValueError('rewards must align with completion_lengths') + if num_generations <= 0 or len(rewards) % num_generations != 0: + raise ValueError('rewards must form complete num_generations groups') + if len(advantages) != len(rewards): + raise ValueError('advantages must align with rewards') + + metrics: Dict[str, float] = {} + if len(completion_lengths) > 0: + lengths = np.asarray(completion_lengths, dtype=np.float64) + metrics['rollout/completion_length_p95'] = float(np.percentile(lengths, 95)) + + if len(stop_reasons) > 0: + num_sequences = len(stop_reasons) + metrics['rollout/stop_rate'] = sum(reason == 'stop' for reason in stop_reasons) / num_sequences + metrics['rollout/length_stop_rate'] = ( + sum(reason == 'length' for reason in stop_reasons) / num_sequences) + + if len(rewards) > 0: + grouped_rewards = np.asarray(rewards, dtype=np.float64).reshape(-1, num_generations) + if num_generations > 1: + group_stds = grouped_rewards.std(axis=1, ddof=1) + else: + group_stds = np.zeros(grouped_rewards.shape[0], dtype=np.float64) + metrics['grpo/group_reward_std_mean'] = float(group_stds.mean()) + zero_variance_groups = zero_variance_reward_group_indices(rewards, num_generations) + metrics['grpo/zero_variance_group_fraction'] = ( + len(zero_variance_groups) / grouped_rewards.shape[0]) + metrics['grpo/nonzero_advantage_fraction'] = float( + (~np.isclose(np.asarray(advantages, dtype=np.float64), 0.0)).mean()) + + if sampling_masks is not None: + if len(sampling_masks) != len(completion_lengths): + raise ValueError('sampling_masks must align with completion_lengths') + support_sizes = [] + for sequence_idx, (sampling_mask, completion_length) in enumerate( + zip(sampling_masks, completion_lengths)): + offsets = sampling_mask.offsets + if len(offsets) - 1 != completion_length: + raise ValueError( + f'sampling mask {sequence_idx} has {len(offsets) - 1} rows, ' + f'expected {completion_length}') + support_sizes.extend(end - start for start, end in zip(offsets, offsets[1:])) + + if support_sizes: + sizes = np.asarray(support_sizes, dtype=np.float64) + metrics['replay/support_size_mean'] = float(sizes.mean()) + metrics['replay/support_size_p50'] = float(np.percentile(sizes, 50)) + metrics['replay/support_size_p95'] = float(np.percentile(sizes, 95)) + metrics['replay/support_size_max'] = float(sizes.max()) + metrics['replay/singleton_fraction'] = float((sizes == 1).mean()) + + return metrics diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 017a515b3..aef56f1ff 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -39,7 +39,7 @@ from twinkle.patch import Patch, apply_context, apply_patch from twinkle.processor import InputProcessor from twinkle.template import Template -from twinkle.utils import construct_class, get_logger, selective_log_softmax, torch_util +from twinkle.utils import construct_class, get_logger, replayed_selective_log_softmax, selective_log_softmax, torch_util from twinkle.utils.framework import Torch from twinkle.utils.grad_clip import normalize_and_clip_grad_norm from twinkle.utils.transformers_utils import filter_from_config_kwargs @@ -446,6 +446,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec """ adapter_name = kwargs.pop('adapter_name', self._get_default_group()) temperature = float(kwargs.pop('temperature', 1.0)) + sampling_masks = kwargs.pop('sampling_masks', None) return_logits = kwargs.pop('return_logits', False) task = kwargs.pop('task', 'causal_lm') optimizer_config = self.optimizer_group[adapter_name] @@ -466,6 +467,15 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec loss_require_logits = getattr(loss_instance, 'require_logits', False) loss_require_entropy = getattr(loss_instance, 'require_entropy', False) loss_require_logps = getattr(loss_instance, 'require_logps', True) + enable_sampling_replay = getattr(loss_instance, 'enable_sampling_replay', False) + if enable_sampling_replay: + if sampling_masks is None: + raise ValueError('sampling_masks are required when sampling replay is enabled') + if kwargs.get('old_logps') is None: + raise ValueError('old_logps are required when sampling replay is enabled') + cp_world_size = self.device_mesh.cp_world_size if self.device_mesh is not None else 1 + if getattr(self, '_enable_sp', False) or cp_world_size > 1: + raise ValueError('sampling replay does not support sequence or context parallelism') assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' inputs: Dict[str, Any] = processor( inputs, @@ -497,11 +507,20 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec masked_labels = labels.clone() masked_labels[~loss_mask] = 0 logits = outputs['logits'] - logits.div_(temperature) - if loss_require_entropy: + if enable_sampling_replay: + outputs['logps'] = replayed_selective_log_softmax( + logits=logits, + labels=masked_labels, + loss_mask=loss_mask, + sampling_masks=sampling_masks, + temperature=temperature, + ) + elif loss_require_entropy: + logits.div_(temperature) outputs['logps'], outputs['entropies'] = selective_log_softmax( logits, masked_labels, return_entropy=True) else: + logits.div_(temperature) outputs['logps'] = selective_log_softmax(logits, masked_labels) del logits outputs['past_key_values'] = None @@ -535,6 +554,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T adapter_name = kwargs.pop('adapter_name', self._get_default_group()) disable_lora = kwargs.pop('disable_lora', False) temperature = float(kwargs.pop('temperature', 1.0)) + sampling_masks = kwargs.pop('sampling_masks', None) return_logits = kwargs.pop('return_logits', False) task = kwargs.pop('task', 'causal_lm') optimizer_config = self.optimizer_group[adapter_name] @@ -557,6 +577,15 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T loss_require_logits = getattr(loss_instance, 'require_logits', False) loss_require_entropy = getattr(loss_instance, 'require_entropy', False) loss_require_logps = getattr(loss_instance, 'require_logps', True) + enable_sampling_replay = getattr(loss_instance, 'enable_sampling_replay', False) + if enable_sampling_replay: + if sampling_masks is None: + raise ValueError('sampling_masks are required when sampling replay is enabled') + if kwargs.get('old_logps') is None: + raise ValueError('old_logps are required when sampling replay is enabled') + cp_world_size = self.device_mesh.cp_world_size if self.device_mesh is not None else 1 + if getattr(self, '_enable_sp', False) or cp_world_size > 1: + raise ValueError('sampling replay does not support sequence or context parallelism') inputs: Dict[str, Any] = processor( inputs, sp_strategy=self.sp_strategy, @@ -591,11 +620,20 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T masked_labels = labels.clone() masked_labels[~loss_mask] = 0 logits = outputs['logits'] - logits.div_(temperature) - if loss_require_entropy: + if enable_sampling_replay: + outputs['logps'] = replayed_selective_log_softmax( + logits=logits, + labels=masked_labels, + loss_mask=loss_mask, + sampling_masks=sampling_masks, + temperature=temperature, + ) + elif loss_require_entropy: + logits.div_(temperature) outputs['logps'], outputs['entropies'] = selective_log_softmax( logits, masked_labels, return_entropy=True) else: + logits.div_(temperature) outputs['logps'] = selective_log_softmax(logits, masked_labels) del logits outputs['past_key_values'] = None diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index b1e1790de..5fcf1d844 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -8,7 +8,7 @@ from typing import Any, Dict, List, Optional, Union from twinkle import get_logger -from twinkle.data_format.sampling import SampledSequence, SampleResponse, SamplingParams, StopReason +from twinkle.data_format.sampling import SampledSequence, SampleResponse, SamplingMask, SamplingParams, StopReason from twinkle.sampler.base_engine import BaseSamplerEngine from twinkle.utils import Platform from twinkle.utils.framework import Torch @@ -29,6 +29,48 @@ def _map_finish_reason(reason: str | None) -> StopReason: return _FINISH_REASON_MAP.get(str(reason), 'length') +def _filter_engine_config( + engine_config: Dict[str, Any], + valid_args, + enable_sampling_replay: bool, +): + valid_args = set(valid_args) + invalid_args = set(engine_config) - valid_args + if enable_sampling_replay and 'enable_return_sampling_mask' in invalid_args: + raise RuntimeError( + 'Sampling replay requires a vLLM build whose AsyncEngineArgs accepts ' + 'enable_return_sampling_mask') + filtered_engine_config = {key: value for key, value in engine_config.items() if key in valid_args} + return filtered_engine_config, invalid_args + + +def _copy_sampling_mask(mask, num_tokens: int, required: bool) -> Optional[SamplingMask]: + if mask is None: + if required: + raise RuntimeError('vLLM output is missing sampling mask while sampling replay is enabled') + return None + + token_ids = [int(token_id) for token_id in mask.token_ids] + offsets = [int(offset) for offset in mask.offsets] + num_rows = len(offsets) - 1 + if num_rows != num_tokens: + raise RuntimeError( + f'vLLM sampling mask has {num_rows} rows for {num_tokens} sampled tokens') + if not offsets or offsets[0] != 0 or offsets[-1] != len(token_ids): + raise RuntimeError('vLLM sampling mask has invalid CSR endpoints') + if any(start >= end for start, end in zip(offsets, offsets[1:])): + raise RuntimeError('vLLM sampling mask contains an empty or invalid CSR row') + return SamplingMask(token_ids=token_ids, offsets=offsets) + + +def _set_sampling_replay_output_kind(vllm_params, enable_sampling_replay: bool) -> None: + """Use the only vLLM output mode that carries the full sampling mask.""" + if not enable_sampling_replay: + return + from vllm.sampling_params import RequestOutputKind + vllm_params.output_kind = RequestOutputKind.FINAL_ONLY + + def get_vllm_max_lora_rank(lora_rank: int) -> int: """Get the nearest allowed vLLM LoRA rank.""" from typing import get_args @@ -78,6 +120,7 @@ def __init__( quantization: Optional[str] = None, load_format: str = 'auto', logprobs_mode: Optional[str] = None, + enable_sampling_replay: bool = False, **kwargs, ): from twinkle.hub import HubOperation @@ -97,7 +140,9 @@ def __init__( self.dtype = dtype self.quantization = quantization self.load_format = load_format - self.logprobs_mode = logprobs_mode or 'processed_logprobs' + self.enable_sampling_replay = enable_sampling_replay + self.logprobs_mode = 'processed_logprobs' if enable_sampling_replay else ( + logprobs_mode or 'processed_logprobs') self.engine_kwargs = kwargs or {} self._lora_request_cache: Dict[str, Any] = {} @@ -130,6 +175,8 @@ def __init__( def _create_engine(self): """Create and return the vLLM engine.""" os.environ['VLLM_USE_V1'] = '1' + if self.enable_sampling_replay: + os.environ['VLLM_USE_V2_MODEL_RUNNER'] = '1' from vllm.engine.arg_utils import AsyncEngineArgs from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.async_llm import AsyncLLM @@ -175,9 +222,15 @@ def _create_engine(self): 'twinkle.sampler.vllm_sampler.vllm_worker_extension.TwinkleWorkerExtension') engine_config.update(self.engine_kwargs) + if self.enable_sampling_replay: + engine_config['enable_return_sampling_mask'] = True + engine_config['logprobs_mode'] = 'processed_logprobs' valid_args = inspect.signature(AsyncEngineArgs).parameters.keys() - filtered_engine_config = {k: v for k, v in engine_config.items() if k in valid_args} - invalid_args = set(engine_config.keys()) - set(valid_args) + filtered_engine_config, invalid_args = _filter_engine_config( + engine_config, + valid_args, + self.enable_sampling_replay, + ) if invalid_args: logger.warning(f'VLLMEngine: Filtered out invalid arguments: {invalid_args}') # Create engine using vLLM v1 API @@ -244,6 +297,7 @@ async def sample(self, prompt_logprobs_k = sampling_params.prompt_logprobs or 0 logprobs = sampling_params.logprobs or 0 vllm_params = sampling_params.to_vllm(**kwargs) + _set_sampling_replay_output_kind(vllm_params, self.enable_sampling_replay) # Build request if request_id is None: @@ -291,6 +345,11 @@ async def sample(self, sequences = [] for output in result.outputs: token_ids = list(output.token_ids) + sampling_mask = _copy_sampling_mask( + getattr(output, 'sampling_mask', None), + num_tokens=len(token_ids), + required=self.enable_sampling_replay, + ) # Extract logprobs seq_logprobs = None @@ -319,6 +378,7 @@ async def sample(self, tokens=token_ids, logprobs=seq_logprobs, routed_experts=routed_experts, + sampling_mask=sampling_mask, )) # Extract prompt logprobs if requested diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 3c7b2f686..def44c793 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -269,6 +269,7 @@ async def _sample_single( logprobs=seq.logprobs, decoded=self.template.decode(seq.tokens), new_input_feature=new_input_feature, + sampling_mask=seq.sampling_mask, ) sequences.append(sampled_seq) return SampleResponse( diff --git a/src/twinkle/utils/__init__.py b/src/twinkle/utils/__init__.py index d5d1b698b..53829fa2b 100644 --- a/src/twinkle/utils/__init__.py +++ b/src/twinkle/utils/__init__.py @@ -10,8 +10,9 @@ from .parallel import processing_lock from .platforms import GPU, NPU, Platform, ensure_hccl_socket_env, ensure_npu_backend from .safetensors import LazyTensor, SafetensorLazyLoader, StreamingSafetensorSaver -from .torch_utils import (clone_state_dict_to_cpu, pad_and_stack_tensors, pad_sequence_to_length, selective_log_softmax, - split_cp_inputs, stateless_init_process_group, to_device) +from .torch_utils import (clone_state_dict_to_cpu, pad_and_stack_tensors, pad_sequence_to_length, + replayed_selective_log_softmax, selective_log_softmax, split_cp_inputs, + stateless_init_process_group, to_device) from .transformers_utils import find_all_linears, find_layers, get_modules_to_not_convert from .unsafe import check_unsafe, trust_remote_code from .utils import copy_files_by_pattern, deep_getattr, get_runtime_meta diff --git a/src/twinkle/utils/nccl_safe.py b/src/twinkle/utils/nccl_safe.py index b22b10137..311590b14 100644 --- a/src/twinkle/utils/nccl_safe.py +++ b/src/twinkle/utils/nccl_safe.py @@ -78,6 +78,7 @@ def __init__(self, loss_instance): self.require_logps = getattr(loss_instance, 'require_logps', True) self.require_entropy = getattr(loss_instance, 'require_entropy', False) self.require_logits = getattr(loss_instance, 'require_logits', False) + self.enable_sampling_replay = getattr(loss_instance, 'enable_sampling_replay', False) self.reduction = getattr(loss_instance, 'reduction', 'mean') self._nccl_safe_wrapped = True diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 84a335852..487289e5c 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -136,6 +136,122 @@ def selective_log_softmax(logits, index, return_entropy: bool = False): return per_token_logps +def replayed_selective_log_softmax( + logits: 'torch.Tensor', + labels: 'torch.Tensor', + loss_mask: 'torch.Tensor', + sampling_masks, + temperature: float, +) -> 'torch.Tensor': + """Compute selected log probabilities on rollout-time CSR support sets.""" + import math + import torch + + if not math.isfinite(temperature) or temperature <= 0: + raise ValueError('temperature must be greater than 0 for sampling replay') + if logits.dim() != 3: + raise ValueError(f'logits must have shape [batch, seq_len, vocab], got {tuple(logits.shape)}') + if labels.shape != logits.shape[:2] or loss_mask.shape != labels.shape: + raise ValueError('labels and loss_mask must match the first two logits dimensions') + if len(sampling_masks) != labels.shape[0]: + raise ValueError( + f'sampling mask batch has {len(sampling_masks)} samples, expected {labels.shape[0]}') + + vocab_size = logits.shape[-1] + flat_token_ids = [] + global_offsets = [0] + for batch_idx, sampling_mask in enumerate(sampling_masks): + if sampling_mask is None: + raise ValueError(f'sampling mask is missing for sample {batch_idx}') + token_ids = [int(token_id) for token_id in sampling_mask.token_ids] + offsets = [int(offset) for offset in sampling_mask.offsets] + if not offsets or offsets[0] != 0: + raise ValueError(f'sampling mask offsets for sample {batch_idx} must start at 0') + if offsets[-1] != len(token_ids): + raise ValueError( + f'sampling mask offsets for sample {batch_idx} must end at {len(token_ids)}') + for row_idx, (start, end) in enumerate(zip(offsets, offsets[1:])): + if start > end: + raise ValueError( + f'sampling mask offsets are not monotonic at sample {batch_idx}, row {row_idx}') + if start == end: + raise ValueError(f'sampling mask contains an empty row at sample {batch_idx}, row {row_idx}') + row_token_ids = token_ids[start:end] + if len(set(row_token_ids)) != len(row_token_ids): + raise ValueError( + f'sampling mask contains duplicate token IDs at sample {batch_idx}, row {row_idx}') + + num_rows = len(offsets) - 1 + num_train_tokens = int(loss_mask[batch_idx].sum().item()) + if num_rows != num_train_tokens: + raise ValueError( + f'sampling mask for sample {batch_idx} has {num_rows} rows but ' + f'{num_train_tokens} training tokens') + invalid_token_id = next( + (token_id for token_id in token_ids if token_id < 0 or token_id >= vocab_size), + None, + ) + if invalid_token_id is not None: + raise ValueError( + f'sampling mask token ID {invalid_token_id} is outside vocabulary [0, {vocab_size})') + + base_offset = global_offsets[-1] + flat_token_ids.extend(token_ids) + global_offsets.extend(base_offset + offset for offset in offsets[1:]) + + positions = loss_mask.nonzero(as_tuple=False) + num_rows = positions.shape[0] + if len(global_offsets) != num_rows + 1: + raise ValueError( + f'sampling masks contain {len(global_offsets) - 1} rows for {num_rows} training tokens') + result = torch.zeros(labels.shape, dtype=torch.float32, device=logits.device) + if num_rows == 0: + return result + + offsets_tensor = torch.tensor(global_offsets, dtype=torch.long, device=logits.device) + lengths = offsets_tensor[1:] - offsets_tensor[:-1] + row_ids = torch.repeat_interleave( + torch.arange(num_rows, device=logits.device), + lengths, + ) + kept_token_ids = torch.tensor(flat_token_ids, dtype=torch.long, device=logits.device) + sampled_labels = labels[positions[:, 0], positions[:, 1]].long() + + matches = kept_token_ids == sampled_labels[row_ids] + match_counts = torch.zeros(num_rows, dtype=torch.int32, device=logits.device) + match_counts.scatter_add_(0, row_ids, matches.to(torch.int32)) + missing_rows = (match_counts == 0).nonzero(as_tuple=False) + if missing_rows.numel(): + row_idx = int(missing_rows[0].item()) + raise ValueError( + f'sampled label {int(sampled_labels[row_idx].item())} is absent from ' + f'sampling mask row {row_idx}') + + kept_logits = logits[ + positions[row_ids, 0], + positions[row_ids, 1], + kept_token_ids, + ].float() / temperature + selected_logits = logits[ + positions[:, 0], + positions[:, 1], + sampled_labels, + ].float() / temperature + + row_max = torch.full( + (num_rows,), + -torch.inf, + dtype=torch.float32, + device=logits.device, + ) + row_max.scatter_reduce_(0, row_ids, kept_logits, reduce='amax', include_self=True) + row_exp_sums = torch.zeros(num_rows, dtype=torch.float32, device=logits.device) + row_exp_sums.scatter_add_(0, row_ids, torch.exp(kept_logits - row_max[row_ids])) + flat_logps = selected_logits - (row_max + torch.log(row_exp_sums)) + result[positions[:, 0], positions[:, 1]] = flat_logps + return result + + def _vocab_parallel_selective_log_softmax( logits: 'torch.Tensor', index: 'torch.Tensor', From ef2d5210e443668b490c7d6410147aef7df7d88e Mon Sep 17 00:00:00 2001 From: vx120 <893600387@qq.com> Date: Tue, 28 Jul 2026 22:02:24 +0800 Subject: [PATCH 2/8] remove additional metric Signed-off-by: vx120 <893600387@qq.com> --- cookbook/rl/grpo/grpo_sampling_replay.py | 66 ++---------------- src/twinkle/metric/__init__.py | 1 - src/twinkle/metric/grpo.py | 21 ++---- src/twinkle/metric/rollout.py | 88 ------------------------ 4 files changed, 8 insertions(+), 168 deletions(-) delete mode 100644 src/twinkle/metric/rollout.py diff --git a/cookbook/rl/grpo/grpo_sampling_replay.py b/cookbook/rl/grpo/grpo_sampling_replay.py index 2f8a4a41f..56e215266 100644 --- a/cookbook/rl/grpo/grpo_sampling_replay.py +++ b/cookbook/rl/grpo/grpo_sampling_replay.py @@ -1,5 +1,4 @@ import os -import time from typing import List, Tuple, Dict, Any from peft import LoraConfig @@ -16,10 +15,7 @@ from twinkle.processor import InputProcessor from twinkle.reward import GSM8KAccuracyReward, GSM8KFormatReward from twinkle.sampler import vLLMSampler -from twinkle.metric import ( - CompletionRewardMetric, - compute_grpo_rollout_metrics, -) +from twinkle.metric import CompletionRewardMetric from twinkle.preprocessor.llm import GSM8KProcessor logger = get_logger() @@ -79,7 +75,6 @@ def extract_rollout_batch(sample_responses, *, require_sampling_masks: bool): 'old_logps': [], 'sampling_masks': [], 'completion_lengths': [], - 'stop_reasons': [], } for sample_response in sample_responses: for sequence in sample_response.sequences: @@ -93,7 +88,6 @@ def extract_rollout_batch(sample_responses, *, require_sampling_masks: bool): [logprob[0][1] for logprob in sequence.logprobs]) rollout_batch['sampling_masks'].append(sequence.sampling_mask) rollout_batch['completion_lengths'].append(len(sequence.tokens)) - rollout_batch['stop_reasons'].append(sequence.stop_reason) return rollout_batch @@ -191,23 +185,7 @@ def main(): top_k=-1, repetition_penalty=1.0, ) - if ENABLE_SAMPLING_REPLAY: - model.add_metric( - 'GRPOMetric', - is_training=True, - temperature=sampling_params.temperature, - epsilon=0.2, - ) - logger.info( - 'Sampling replay enabled: model_runner=v2, logprobs_mode=processed_logprobs, ' - 'temperature=%s, top_p=%s, top_k=%s', - sampling_params.temperature, - sampling_params.top_p, - sampling_params.top_k, - ) - optim_step = 0 - sampling_replay_stats_logged = False logger.info(get_device_placement()) for batch in dataloader: @@ -218,27 +196,22 @@ def main(): # enable_lora=True used with ckpt_manager.sync_weights(merge_and_sync=False) # meaning only sync lora weights, if merge_and_sync=True, # lora will be merged into the base model and sync all weights to vLLM - weight_sync_started = time.perf_counter() ckpt_manager.sync_weights(merge_and_sync=False) - weight_sync_seconds = time.perf_counter() - weight_sync_started sampler.reset_prefix_cache() def sample_prompt_groups(prompts): expand_prompts = [] for prompt in prompts: expand_prompts.extend([prompt] * NUM_GENERATIONS) - started = time.perf_counter() responses = sampler.sample(expand_prompts, sampling_params) - elapsed = time.perf_counter() - started return extract_rollout_batch( responses, require_sampling_masks=ENABLE_SAMPLING_REPLAY, - ), elapsed + ) - rollout_batch, sampling_seconds = sample_prompt_groups(global_prompts) - sampled_tokens_total = sum(rollout_batch['completion_lengths']) + rollout_batch = sample_prompt_groups(global_prompts) # Match the original GRPO control flow: every sampled rollout is scored, # logged, and trained. Zero-variance groups keep their zero advantages; - # they are diagnosed below but never resampled, dropped, or skipped. + # they are never resampled, dropped, or skipped. total_rewards, format_rewards, accuracy_rewards = compute_rewards( rollout_batch['input_data']) @@ -246,7 +219,6 @@ def sample_prompt_groups(prompts): all_old_logps: List[List[float]] = rollout_batch['old_logps'] all_sampling_masks = rollout_batch['sampling_masks'] all_completion_lengths: List[int] = rollout_batch['completion_lengths'] - all_stop_reasons = rollout_batch['stop_reasons'] metrics.accumulate( completion_lengths=all_completion_lengths, rewards={ @@ -258,20 +230,6 @@ def sample_prompt_groups(prompts): rollout_reward_log_dict = metrics.calculate() advantages = advantage_fn(total_rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() - rollout_log_dict = compute_grpo_rollout_metrics( - completion_lengths=all_completion_lengths, - stop_reasons=all_stop_reasons, - rewards=total_rewards, - advantages=advantages, - num_generations=NUM_GENERATIONS, - sampling_masks=all_sampling_masks if ENABLE_SAMPLING_REPLAY else None, - ) - num_rollout_tokens = sum(all_completion_lengths) - rollout_log_dict['profiling/weight_sync_seconds'] = weight_sync_seconds - rollout_log_dict['profiling/sampling_seconds'] = sampling_seconds - rollout_log_dict['profiling/sampling_tokens_per_second'] = ( - sampled_tokens_total / sampling_seconds if sampling_seconds else 0.0) - rollout_log_dict['profiling/sampling_generated_tokens'] = sampled_tokens_total # Split completions into mini-batches and run one optim step per mini-batch. total_completions = len(all_input_data) @@ -287,7 +245,6 @@ def sample_prompt_groups(prompts): 'temperature': sampling_params.temperature, } - training_started = time.perf_counter() model.forward_backward( inputs=mb_inputs, old_logps=mb_old_logps, @@ -296,15 +253,6 @@ def sample_prompt_groups(prompts): **replay_kwargs, ) model.clip_grad_and_step() - training_seconds = time.perf_counter() - training_started - if ENABLE_SAMPLING_REPLAY and not sampling_replay_stats_logged: - logger.info( - 'Sampling replay active: sequences=%d, tokens=%d, mean_kept_tokens=%.2f', - len(all_sampling_masks), - num_rollout_tokens, - rollout_log_dict['replay/support_size_mean'], - ) - sampling_replay_stats_logged = True optim_step += 1 if optim_step % SAVE_STEPS == 0: @@ -313,12 +261,6 @@ def sample_prompt_groups(prompts): # rollout can span multiple mini-batches, but no Step lacks reward. log_dict = dict(rollout_reward_log_dict) log_dict.update(model.calculate_metric(is_training=True)) - if mb_start == 0: - log_dict.update(rollout_log_dict) - num_training_tokens = sum(all_completion_lengths[mb_start:mb_end]) - log_dict['profiling/training_seconds'] = training_seconds - log_dict['profiling/training_completion_tokens_per_second'] = ( - num_training_tokens / training_seconds if training_seconds else 0.0) logger.info(f'[Step {optim_step}/{MAX_STEPS}] {log_dict}') if optim_step >= MAX_STEPS: break diff --git a/src/twinkle/metric/__init__.py b/src/twinkle/metric/__init__.py index cd7d8c99d..baeb6c1c9 100644 --- a/src/twinkle/metric/__init__.py +++ b/src/twinkle/metric/__init__.py @@ -6,5 +6,4 @@ from .embedding import EmbeddingMetric from .grpo import CISPOMetric, GRPOMetric, GSPOMetric from .loss import LossMetric -from .rollout import compute_grpo_rollout_metrics, zero_variance_reward_group_indices from .train_metric import TrainMetric diff --git a/src/twinkle/metric/grpo.py b/src/twinkle/metric/grpo.py index 71fde1de4..bd85aab67 100644 --- a/src/twinkle/metric/grpo.py +++ b/src/twinkle/metric/grpo.py @@ -1,6 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import math -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union from twinkle.data_format import InputFeature, ModelOutput from twinkle.utils import get_logger @@ -9,9 +9,6 @@ logger = get_logger() -if TYPE_CHECKING: - import torch - class GRPOMetric(Metric): @@ -44,8 +41,6 @@ def reset(self): self.sum_new: float = 0.0 self.sum_old: float = 0.0 self.sum_diff: float = 0.0 - self.sum_diff_sq: float = 0.0 - self.sum_ratio: float = 0.0 self.sum_approx_kl: float = 0.0 self.max_token_kl: float = 0.0 self.max_token_ratio: float = 0.0 @@ -190,16 +185,13 @@ def _accumulate_mb( old_f = old_f * scale d = logps_f - old_f # new - old - ratio = torch.exp(d) self.sum_old += float((old_f * mask_f).sum().item()) self.sum_diff += float((d * mask_f).sum().item()) - self.sum_diff_sq += float((d.square() * mask_f).sum().item()) - self.sum_ratio += float((ratio * mask_f).sum().item()) # Schulman K3 estimator of KL(old || new): # samples x ~ old, r(x) = new(x) / old(x), # k3 = r - 1 - log(r) = exp(new - old) - (new - old) - 1. - kl = ratio - d - 1.0 + kl = torch.exp(d) - d - 1.0 kl_masked = kl * mask_f self.sum_approx_kl += float(kl_masked.sum().item()) # Per-token extremes for collapse detection @@ -208,7 +200,7 @@ def _accumulate_mb( if cur_max_kl > self.max_token_kl: self.max_token_kl = cur_max_kl # Track ratio extremes - ratio_masked = ratio * mask_f + ratio_masked = torch.exp(d) * mask_f cur_max_ratio = float(ratio_masked.max().item()) if cur_max_ratio > self.max_token_ratio: self.max_token_ratio = cur_max_ratio @@ -319,12 +311,11 @@ def accumulate( cursor += advanced def calculate(self) -> Dict[str, Any]: + import torch local = [{ 'sum_new': self.sum_new, 'sum_old': self.sum_old, 'sum_diff': self.sum_diff, - 'sum_diff_sq': self.sum_diff_sq, - 'sum_ratio': self.sum_ratio, 'sum_kl': self.sum_approx_kl, 'max_token_kl': self.max_token_kl, 'max_token_ratio': self.max_token_ratio, @@ -353,15 +344,11 @@ def calculate(self) -> Dict[str, Any]: if any(r['has_old'] for r in all_results): mean_old = sum(r['sum_old'] for r in all_results) / n_total mean_diff = sum(r['sum_diff'] for r in all_results) / n_total - mean_diff_sq = sum(r['sum_diff_sq'] for r in all_results) / n_total - mean_ratio = sum(r['sum_ratio'] for r in all_results) / n_total mean_kl = sum(r['sum_kl'] for r in all_results) / n_total global_max_kl = max(r['max_token_kl'] for r in all_results) global_max_ratio = max(r['max_token_ratio'] for r in all_results) results['train/mean_old_logp'] = mean_old results['train/logp_diff_mean'] = mean_diff - results['train/logp_diff_std'] = math.sqrt(max(mean_diff_sq - mean_diff**2, 0.0)) - results['train/importance_ratio_mean'] = mean_ratio results['train/approx_kl'] = mean_kl results['train/token_kl_max'] = global_max_kl results['train/token_ratio_max'] = global_max_ratio diff --git a/src/twinkle/metric/rollout.py b/src/twinkle/metric/rollout.py deleted file mode 100644 index cc261ccd9..000000000 --- a/src/twinkle/metric/rollout.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from typing import Any, Dict, Optional, Sequence - -import numpy as np - - -def zero_variance_reward_group_indices( - rewards: Sequence[float], - num_generations: int, -) -> list[int]: - """Return GRPO group indices that cannot produce a relative advantage.""" - if num_generations <= 0: - raise ValueError('num_generations must be positive') - if len(rewards) % num_generations != 0: - raise ValueError('rewards must form complete num_generations groups') - if len(rewards) == 0: - return [] - - grouped_rewards = np.asarray(rewards, dtype=np.float64).reshape(-1, num_generations) - group_ranges = np.ptp(grouped_rewards, axis=1) - return np.flatnonzero(np.isclose(group_ranges, 0.0)).astype(int).tolist() - - -def compute_grpo_rollout_metrics( - *, - completion_lengths: Sequence[int], - stop_reasons: Sequence[str], - rewards: Sequence[float], - advantages: Sequence[float], - num_generations: int, - sampling_masks: Optional[Sequence[Any]] = None, -) -> Dict[str, float]: - """Reduce one GRPO rollout batch into scalar diagnostics.""" - if len(stop_reasons) != len(completion_lengths): - raise ValueError('stop_reasons must align with completion_lengths') - if len(rewards) != len(completion_lengths): - raise ValueError('rewards must align with completion_lengths') - if num_generations <= 0 or len(rewards) % num_generations != 0: - raise ValueError('rewards must form complete num_generations groups') - if len(advantages) != len(rewards): - raise ValueError('advantages must align with rewards') - - metrics: Dict[str, float] = {} - if len(completion_lengths) > 0: - lengths = np.asarray(completion_lengths, dtype=np.float64) - metrics['rollout/completion_length_p95'] = float(np.percentile(lengths, 95)) - - if len(stop_reasons) > 0: - num_sequences = len(stop_reasons) - metrics['rollout/stop_rate'] = sum(reason == 'stop' for reason in stop_reasons) / num_sequences - metrics['rollout/length_stop_rate'] = ( - sum(reason == 'length' for reason in stop_reasons) / num_sequences) - - if len(rewards) > 0: - grouped_rewards = np.asarray(rewards, dtype=np.float64).reshape(-1, num_generations) - if num_generations > 1: - group_stds = grouped_rewards.std(axis=1, ddof=1) - else: - group_stds = np.zeros(grouped_rewards.shape[0], dtype=np.float64) - metrics['grpo/group_reward_std_mean'] = float(group_stds.mean()) - zero_variance_groups = zero_variance_reward_group_indices(rewards, num_generations) - metrics['grpo/zero_variance_group_fraction'] = ( - len(zero_variance_groups) / grouped_rewards.shape[0]) - metrics['grpo/nonzero_advantage_fraction'] = float( - (~np.isclose(np.asarray(advantages, dtype=np.float64), 0.0)).mean()) - - if sampling_masks is not None: - if len(sampling_masks) != len(completion_lengths): - raise ValueError('sampling_masks must align with completion_lengths') - support_sizes = [] - for sequence_idx, (sampling_mask, completion_length) in enumerate( - zip(sampling_masks, completion_lengths)): - offsets = sampling_mask.offsets - if len(offsets) - 1 != completion_length: - raise ValueError( - f'sampling mask {sequence_idx} has {len(offsets) - 1} rows, ' - f'expected {completion_length}') - support_sizes.extend(end - start for start, end in zip(offsets, offsets[1:])) - - if support_sizes: - sizes = np.asarray(support_sizes, dtype=np.float64) - metrics['replay/support_size_mean'] = float(sizes.mean()) - metrics['replay/support_size_p50'] = float(np.percentile(sizes, 50)) - metrics['replay/support_size_p95'] = float(np.percentile(sizes, 95)) - metrics['replay/support_size_max'] = float(sizes.max()) - metrics['replay/singleton_fraction'] = float((sizes == 1).mean()) - - return metrics From cf6ff2fce3cd137b19bdae70b0d631e7517970a2 Mon Sep 17 00:00:00 2001 From: vx120 <893600387@qq.com> Date: Wed, 29 Jul 2026 17:11:30 +0800 Subject: [PATCH 3/8] Added some code comments. Signed-off-by: vx120 <893600387@qq.com> --- src/twinkle/utils/torch_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 487289e5c..34c45c1ec 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -157,6 +157,7 @@ def replayed_selective_log_softmax( raise ValueError( f'sampling mask batch has {len(sampling_masks)} samples, expected {labels.shape[0]}') + # Flatten per-sample CSR rows into one global CSR layout. vocab_size = logits.shape[-1] flat_token_ids = [] global_offsets = [0] @@ -199,6 +200,7 @@ def replayed_selective_log_softmax( flat_token_ids.extend(token_ids) global_offsets.extend(base_offset + offset for offset in offsets[1:]) + # CSR rows are ordered exactly like the masked training-token positions. positions = loss_mask.nonzero(as_tuple=False) num_rows = positions.shape[0] if len(global_offsets) != num_rows + 1: @@ -227,6 +229,7 @@ def replayed_selective_log_softmax( f'sampled label {int(sampled_labels[row_idx].item())} is absent from ' f'sampling mask row {row_idx}') + # Gather only logits retained by the rollout sampler, then normalize per CSR row. kept_logits = logits[ positions[row_ids, 0], positions[row_ids, 1], @@ -238,6 +241,7 @@ def replayed_selective_log_softmax( sampled_labels, ].float() / temperature + # Use max-shifted log-sum-exp for numerically stable restricted softmax. row_max = torch.full( (num_rows,), -torch.inf, From 710021a77e1bd967dbca9702ef0598a962c7a673 Mon Sep 17 00:00:00 2001 From: vx120 <893600387@qq.com> Date: Wed, 29 Jul 2026 17:17:56 +0800 Subject: [PATCH 4/8] add pre-commit code Signed-off-by: vx120 <893600387@qq.com> --- .../sampler/vllm_sampler/vllm_engine.py | 11 +++---- src/twinkle/utils/torch_utils.py | 30 +++++++------------ 2 files changed, 15 insertions(+), 26 deletions(-) diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index 5fcf1d844..d4f3448a6 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -37,9 +37,8 @@ def _filter_engine_config( valid_args = set(valid_args) invalid_args = set(engine_config) - valid_args if enable_sampling_replay and 'enable_return_sampling_mask' in invalid_args: - raise RuntimeError( - 'Sampling replay requires a vLLM build whose AsyncEngineArgs accepts ' - 'enable_return_sampling_mask') + raise RuntimeError('Sampling replay requires a vLLM build whose AsyncEngineArgs accepts ' + 'enable_return_sampling_mask') filtered_engine_config = {key: value for key, value in engine_config.items() if key in valid_args} return filtered_engine_config, invalid_args @@ -54,8 +53,7 @@ def _copy_sampling_mask(mask, num_tokens: int, required: bool) -> Optional[Sampl offsets = [int(offset) for offset in mask.offsets] num_rows = len(offsets) - 1 if num_rows != num_tokens: - raise RuntimeError( - f'vLLM sampling mask has {num_rows} rows for {num_tokens} sampled tokens') + raise RuntimeError(f'vLLM sampling mask has {num_rows} rows for {num_tokens} sampled tokens') if not offsets or offsets[0] != 0 or offsets[-1] != len(token_ids): raise RuntimeError('vLLM sampling mask has invalid CSR endpoints') if any(start >= end for start, end in zip(offsets, offsets[1:])): @@ -141,8 +139,7 @@ def __init__( self.quantization = quantization self.load_format = load_format self.enable_sampling_replay = enable_sampling_replay - self.logprobs_mode = 'processed_logprobs' if enable_sampling_replay else ( - logprobs_mode or 'processed_logprobs') + self.logprobs_mode = 'processed_logprobs' if enable_sampling_replay else (logprobs_mode or 'processed_logprobs') self.engine_kwargs = kwargs or {} self._lora_request_cache: Dict[str, Any] = {} diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 34c45c1ec..9892b79db 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -154,8 +154,7 @@ def replayed_selective_log_softmax( if labels.shape != logits.shape[:2] or loss_mask.shape != labels.shape: raise ValueError('labels and loss_mask must match the first two logits dimensions') if len(sampling_masks) != labels.shape[0]: - raise ValueError( - f'sampling mask batch has {len(sampling_masks)} samples, expected {labels.shape[0]}') + raise ValueError(f'sampling mask batch has {len(sampling_masks)} samples, expected {labels.shape[0]}') # Flatten per-sample CSR rows into one global CSR layout. vocab_size = logits.shape[-1] @@ -169,32 +168,27 @@ def replayed_selective_log_softmax( if not offsets or offsets[0] != 0: raise ValueError(f'sampling mask offsets for sample {batch_idx} must start at 0') if offsets[-1] != len(token_ids): - raise ValueError( - f'sampling mask offsets for sample {batch_idx} must end at {len(token_ids)}') + raise ValueError(f'sampling mask offsets for sample {batch_idx} must end at {len(token_ids)}') for row_idx, (start, end) in enumerate(zip(offsets, offsets[1:])): if start > end: - raise ValueError( - f'sampling mask offsets are not monotonic at sample {batch_idx}, row {row_idx}') + raise ValueError(f'sampling mask offsets are not monotonic at sample {batch_idx}, row {row_idx}') if start == end: raise ValueError(f'sampling mask contains an empty row at sample {batch_idx}, row {row_idx}') row_token_ids = token_ids[start:end] if len(set(row_token_ids)) != len(row_token_ids): - raise ValueError( - f'sampling mask contains duplicate token IDs at sample {batch_idx}, row {row_idx}') + raise ValueError(f'sampling mask contains duplicate token IDs at sample {batch_idx}, row {row_idx}') num_rows = len(offsets) - 1 num_train_tokens = int(loss_mask[batch_idx].sum().item()) if num_rows != num_train_tokens: - raise ValueError( - f'sampling mask for sample {batch_idx} has {num_rows} rows but ' - f'{num_train_tokens} training tokens') + raise ValueError(f'sampling mask for sample {batch_idx} has {num_rows} rows but ' + f'{num_train_tokens} training tokens') invalid_token_id = next( (token_id for token_id in token_ids if token_id < 0 or token_id >= vocab_size), None, ) if invalid_token_id is not None: - raise ValueError( - f'sampling mask token ID {invalid_token_id} is outside vocabulary [0, {vocab_size})') + raise ValueError(f'sampling mask token ID {invalid_token_id} is outside vocabulary [0, {vocab_size})') base_offset = global_offsets[-1] flat_token_ids.extend(token_ids) @@ -204,8 +198,7 @@ def replayed_selective_log_softmax( positions = loss_mask.nonzero(as_tuple=False) num_rows = positions.shape[0] if len(global_offsets) != num_rows + 1: - raise ValueError( - f'sampling masks contain {len(global_offsets) - 1} rows for {num_rows} training tokens') + raise ValueError(f'sampling masks contain {len(global_offsets) - 1} rows for {num_rows} training tokens') result = torch.zeros(labels.shape, dtype=torch.float32, device=logits.device) if num_rows == 0: return result @@ -225,9 +218,8 @@ def replayed_selective_log_softmax( missing_rows = (match_counts == 0).nonzero(as_tuple=False) if missing_rows.numel(): row_idx = int(missing_rows[0].item()) - raise ValueError( - f'sampled label {int(sampled_labels[row_idx].item())} is absent from ' - f'sampling mask row {row_idx}') + raise ValueError(f'sampled label {int(sampled_labels[row_idx].item())} is absent from ' + f'sampling mask row {row_idx}') # Gather only logits retained by the rollout sampler, then normalize per CSR row. kept_logits = logits[ @@ -243,7 +235,7 @@ def replayed_selective_log_softmax( # Use max-shifted log-sum-exp for numerically stable restricted softmax. row_max = torch.full( - (num_rows,), + (num_rows, ), -torch.inf, dtype=torch.float32, device=logits.device, From 8678bd150c24c5d61b5ec05f9c4e2e2dfa17e3ed Mon Sep 17 00:00:00 2001 From: vx120 <893600387@qq.com> Date: Wed, 29 Jul 2026 19:48:40 +0800 Subject: [PATCH 5/8] reduce the judge code Signed-off-by: vx120 <893600387@qq.com> --- cookbook/rl/grpo/grpo_sampling_replay.py | 10 ++-------- src/twinkle/loss/grpo.py | 5 ----- src/twinkle/model/transformers/transformers.py | 4 ---- src/twinkle/utils/torch_utils.py | 17 +++-------------- 4 files changed, 5 insertions(+), 31 deletions(-) diff --git a/cookbook/rl/grpo/grpo_sampling_replay.py b/cookbook/rl/grpo/grpo_sampling_replay.py index 56e215266..3fba27562 100644 --- a/cookbook/rl/grpo/grpo_sampling_replay.py +++ b/cookbook/rl/grpo/grpo_sampling_replay.py @@ -68,7 +68,7 @@ def compute_rewards( return total_rewards, format_rewards, accuracy_rewards -def extract_rollout_batch(sample_responses, *, require_sampling_masks: bool): +def extract_rollout_batch(sample_responses): """Flatten sampler responses into aligned lists used by reward and training.""" rollout_batch = { 'input_data': [], @@ -80,9 +80,6 @@ def extract_rollout_batch(sample_responses, *, require_sampling_masks: bool): for sequence in sample_response.sequences: if sequence.logprobs is None: raise RuntimeError('A sampled sequence is missing token log probabilities') - if require_sampling_masks and sequence.sampling_mask is None: - raise RuntimeError( - 'Sampling replay is enabled but a sampled sequence has no sampling mask') rollout_batch['input_data'].append(sequence.new_input_feature) rollout_batch['old_logps'].append( [logprob[0][1] for logprob in sequence.logprobs]) @@ -203,10 +200,7 @@ def sample_prompt_groups(prompts): for prompt in prompts: expand_prompts.extend([prompt] * NUM_GENERATIONS) responses = sampler.sample(expand_prompts, sampling_params) - return extract_rollout_batch( - responses, - require_sampling_masks=ENABLE_SAMPLING_REPLAY, - ) + return extract_rollout_batch(responses) rollout_batch = sample_prompt_groups(global_prompts) # Match the original GRPO control flow: every sampled rollout is scored, diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 36970636d..52479b92b 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -40,8 +40,6 @@ def __init__( self.beta = beta self.entropy_coef = entropy_coef self.enable_sampling_replay = enable_sampling_replay - if enable_sampling_replay and self.__class__ is not GRPOLoss: - raise ValueError('sampling replay is only supported by GRPOLoss') if enable_sampling_replay and beta != 0.0: raise ValueError('sampling replay does not support a GRPO KL penalty (beta must be 0)') if enable_sampling_replay and entropy_coef != 0.0: @@ -209,7 +207,6 @@ def __call__( old_logps: Optional[Union['torch.Tensor', List[List[float]]]] = None, ref_logps: Optional['torch.Tensor'] = None, advantages: Optional[Union['torch.Tensor', List[float], np.ndarray]] = None, - sampling_masks=None, **kwargs, ): """ @@ -232,8 +229,6 @@ def __call__( """ import torch if self.enable_sampling_replay: - if sampling_masks is None: - raise ValueError('sampling_masks are required when sampling replay is enabled') if old_logps is None: raise ValueError('old_logps are required when sampling replay is enabled') labels = inputs.get('labels') diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index aef56f1ff..5afc46c0a 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -471,8 +471,6 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec if enable_sampling_replay: if sampling_masks is None: raise ValueError('sampling_masks are required when sampling replay is enabled') - if kwargs.get('old_logps') is None: - raise ValueError('old_logps are required when sampling replay is enabled') cp_world_size = self.device_mesh.cp_world_size if self.device_mesh is not None else 1 if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') @@ -581,8 +579,6 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T if enable_sampling_replay: if sampling_masks is None: raise ValueError('sampling_masks are required when sampling replay is enabled') - if kwargs.get('old_logps') is None: - raise ValueError('old_logps are required when sampling replay is enabled') cp_world_size = self.device_mesh.cp_world_size if self.device_mesh is not None else 1 if getattr(self, '_enable_sp', False) or cp_world_size > 1: raise ValueError('sampling replay does not support sequence or context parallelism') diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 9892b79db..f6c7a008e 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -136,6 +136,9 @@ def selective_log_softmax(logits, index, return_entropy: bool = False): return per_token_logps +# Re-normalize trainer logits over each rollout-time top-p/top-k support set +# before reading the sampled token's log probability. Replaying the sampler's +# action space removes the sampling/training distribution mismatch in GRPO. def replayed_selective_log_softmax( logits: 'torch.Tensor', labels: 'torch.Tensor', @@ -165,18 +168,6 @@ def replayed_selective_log_softmax( raise ValueError(f'sampling mask is missing for sample {batch_idx}') token_ids = [int(token_id) for token_id in sampling_mask.token_ids] offsets = [int(offset) for offset in sampling_mask.offsets] - if not offsets or offsets[0] != 0: - raise ValueError(f'sampling mask offsets for sample {batch_idx} must start at 0') - if offsets[-1] != len(token_ids): - raise ValueError(f'sampling mask offsets for sample {batch_idx} must end at {len(token_ids)}') - for row_idx, (start, end) in enumerate(zip(offsets, offsets[1:])): - if start > end: - raise ValueError(f'sampling mask offsets are not monotonic at sample {batch_idx}, row {row_idx}') - if start == end: - raise ValueError(f'sampling mask contains an empty row at sample {batch_idx}, row {row_idx}') - row_token_ids = token_ids[start:end] - if len(set(row_token_ids)) != len(row_token_ids): - raise ValueError(f'sampling mask contains duplicate token IDs at sample {batch_idx}, row {row_idx}') num_rows = len(offsets) - 1 num_train_tokens = int(loss_mask[batch_idx].sum().item()) @@ -197,8 +188,6 @@ def replayed_selective_log_softmax( # CSR rows are ordered exactly like the masked training-token positions. positions = loss_mask.nonzero(as_tuple=False) num_rows = positions.shape[0] - if len(global_offsets) != num_rows + 1: - raise ValueError(f'sampling masks contain {len(global_offsets) - 1} rows for {num_rows} training tokens') result = torch.zeros(labels.shape, dtype=torch.float32, device=logits.device) if num_rows == 0: return result From ce80535391a323b7cce79bc614e721f13c2067d6 Mon Sep 17 00:00:00 2001 From: vx120 <893600387@qq.com> Date: Mon, 10 Aug 2026 16:59:24 +0800 Subject: [PATCH 6/8] add ci test scripts Signed-off-by: vx120 <893600387@qq.com> --- tests/loss/test_sampling_replay.py | 70 ++++++++++ tests/sampler/test_sampling_replay.py | 97 ++++++++++++++ tests/utils/test_sampling_replay.py | 182 ++++++++++++++++++++++++++ 3 files changed, 349 insertions(+) create mode 100644 tests/loss/test_sampling_replay.py create mode 100644 tests/sampler/test_sampling_replay.py create mode 100644 tests/utils/test_sampling_replay.py diff --git a/tests/loss/test_sampling_replay.py b/tests/loss/test_sampling_replay.py new file mode 100644 index 000000000..baf28008e --- /dev/null +++ b/tests/loss/test_sampling_replay.py @@ -0,0 +1,70 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import pytest +import torch + +from twinkle.loss import GRPOLoss + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"beta": 0.1}, "KL penalty"), + ({"entropy_coef": 0.1}, "entropy bonus"), + ], +) +def test_sampling_replay_rejects_incompatible_grpo_options(kwargs, message): + with pytest.raises(ValueError, match=message): + GRPOLoss(enable_sampling_replay=True, **kwargs) + + +def test_sampling_replay_requires_rollout_and_model_logps(): + loss = GRPOLoss(enable_sampling_replay=True) + inputs = {"labels": torch.tensor([[1]])} + + with pytest.raises(ValueError, match="old_logps are required"): + loss(inputs, {"logps": torch.tensor([[-0.5]])}, advantages=[1.0]) + with pytest.raises(RuntimeError, match="must be computed by the model forward"): + loss( + inputs, + {"logits": torch.zeros(1, 1, 2)}, + old_logps=[[-0.5]], + advantages=[1.0], + ) + + +def test_sampling_replay_grpo_uses_replayed_importance_ratio_and_clipping(): + labels = torch.tensor([[-100, 1, 2]]) + replayed_logps = torch.tensor([[0.0, -0.4, -0.6]], requires_grad=True) + old_logps = [[-0.5, -0.5]] + advantages = [[1.0, -2.0]] + + result = GRPOLoss(enable_sampling_replay=True, epsilon=0.2)( + {"labels": labels}, + {"logps": replayed_logps}, + old_logps=old_logps, + advantages=advantages, + ) + + ratio = torch.exp(torch.tensor([0.1, -0.1])) + clipped = ratio.clamp(0.8, 1.2) + expected_tokens = -torch.minimum( + ratio * torch.tensor([1.0, -2.0]), clipped * torch.tensor([1.0, -2.0]) + ) + torch.testing.assert_close(result["loss"], expected_tokens.mean()) + result["loss"].backward() + assert replayed_logps.grad[0, 0] == 0 + assert replayed_logps.grad[0, 1:].abs().sum() > 0 + + +def test_sampling_replay_without_advantages_returns_graph_connected_zero(): + logps = torch.tensor([[-0.2, -0.3]], requires_grad=True) + result = GRPOLoss(enable_sampling_replay=True)( + {"labels": torch.tensor([[1, 2]])}, + {"logps": logps}, + old_logps=[[-0.2, -0.3]], + ) + + assert result["loss"].item() == 0.0 + assert result["num_tokens"] == 0 + result["loss"].backward() + assert torch.equal(logps.grad, torch.zeros_like(logps)) diff --git a/tests/sampler/test_sampling_replay.py b/tests/sampler/test_sampling_replay.py new file mode 100644 index 000000000..b05cbacd3 --- /dev/null +++ b/tests/sampler/test_sampling_replay.py @@ -0,0 +1,97 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from twinkle.data_format import SamplingMask +from twinkle.sampler.vllm_sampler.vllm_engine import ( + _copy_sampling_mask, + _filter_engine_config, + _set_sampling_replay_output_kind, +) + + +def test_copy_sampling_mask_converts_vllm_values_and_detaches_storage(): + source_token_ids = ["1", 3, 4] + source_offsets = [0, 2, 3] + copied = _copy_sampling_mask( + SimpleNamespace(token_ids=source_token_ids, offsets=source_offsets), + num_tokens=2, + required=True, + ) + + assert copied == SamplingMask(token_ids=[1, 3, 4], offsets=[0, 2, 3]) + source_token_ids[0] = 99 + source_offsets[-1] = 99 + assert copied == SamplingMask(token_ids=[1, 3, 4], offsets=[0, 2, 3]) + + +def test_missing_sampling_mask_is_only_allowed_when_replay_is_disabled(): + assert _copy_sampling_mask(None, num_tokens=2, required=False) is None + with pytest.raises(RuntimeError, match="missing sampling mask"): + _copy_sampling_mask(None, num_tokens=2, required=True) + + +@pytest.mark.parametrize( + ("mask", "num_tokens", "message"), + [ + ( + SimpleNamespace(token_ids=[1], offsets=[0, 1]), + 2, + "1 rows for 2 sampled tokens", + ), + (SimpleNamespace(token_ids=[1], offsets=[1, 1]), 1, "invalid CSR endpoints"), + (SimpleNamespace(token_ids=[1], offsets=[0, 0]), 1, "invalid CSR endpoints"), + ( + SimpleNamespace(token_ids=[1, 2], offsets=[0, 2, 1]), + 2, + "invalid CSR endpoints", + ), + ( + SimpleNamespace(token_ids=[1], offsets=[0, 0, 1]), + 2, + "empty or invalid CSR row", + ), + ], +) +def test_copy_sampling_mask_rejects_invalid_csr(mask, num_tokens, message): + with pytest.raises(RuntimeError, match=message): + _copy_sampling_mask(mask, num_tokens=num_tokens, required=True) + + +def test_filter_engine_config_preserves_supported_replay_flag(): + filtered, invalid = _filter_engine_config( + {"dtype": "bfloat16", "enable_return_sampling_mask": True, "unknown": 1}, + {"dtype", "enable_return_sampling_mask"}, + enable_sampling_replay=True, + ) + + assert filtered == {"dtype": "bfloat16", "enable_return_sampling_mask": True} + assert invalid == {"unknown"} + + +def test_filter_engine_config_fails_fast_for_incompatible_vllm(): + with pytest.raises( + RuntimeError, match="AsyncEngineArgs accepts enable_return_sampling_mask" + ): + _filter_engine_config( + {"dtype": "bfloat16", "enable_return_sampling_mask": True}, + {"dtype"}, + enable_sampling_replay=True, + ) + + +def test_replay_forces_final_only_output_kind(monkeypatch): + request_output_kind = SimpleNamespace(FINAL_ONLY=object()) + sampling_params_module = ModuleType("vllm.sampling_params") + sampling_params_module.RequestOutputKind = request_output_kind + vllm_module = ModuleType("vllm") + monkeypatch.setitem(sys.modules, "vllm", vllm_module) + monkeypatch.setitem(sys.modules, "vllm.sampling_params", sampling_params_module) + params = SimpleNamespace(output_kind="unchanged") + + _set_sampling_replay_output_kind(params, enable_sampling_replay=False) + assert params.output_kind == "unchanged" + _set_sampling_replay_output_kind(params, enable_sampling_replay=True) + assert params.output_kind is request_output_kind.FINAL_ONLY diff --git a/tests/utils/test_sampling_replay.py b/tests/utils/test_sampling_replay.py new file mode 100644 index 000000000..2c2a94b91 --- /dev/null +++ b/tests/utils/test_sampling_replay.py @@ -0,0 +1,182 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import math + +import pytest +import torch + +from twinkle.data_format import SamplingMask +from twinkle.utils.torch_utils import replayed_selective_log_softmax + + +def _sampling_mask(*rows): + token_ids = [token_id for row in rows for token_id in row] + offsets = [0] + for row in rows: + offsets.append(offsets[-1] + len(row)) + return SamplingMask(token_ids=token_ids, offsets=offsets) + + +def _reference_replayed_logps(logits, labels, loss_mask, sampling_masks, temperature): + expected = torch.zeros_like(labels, dtype=torch.float32) + for batch_idx, sampling_mask in enumerate(sampling_masks): + row_idx = 0 + for seq_idx in loss_mask[batch_idx].nonzero(as_tuple=True)[0].tolist(): + start = sampling_mask.offsets[row_idx] + end = sampling_mask.offsets[row_idx + 1] + support = sampling_mask.token_ids[start:end] + support_logits = logits[batch_idx, seq_idx, support].float() / temperature + label = int(labels[batch_idx, seq_idx]) + expected[batch_idx, seq_idx] = logits[ + batch_idx, seq_idx, label + ].float() / temperature - torch.logsumexp(support_logits, dim=0) + row_idx += 1 + return expected + + +def test_replayed_logps_match_restricted_softmax_for_ragged_batch(): + logits = torch.tensor( + [ + [ + [0.2, 1.0, -0.5, 2.0, 0.3], + [1.1, -0.2, 0.7, 0.1, 2.4], + [0.3, 0.4, 0.5, 0.6, 0.7], + ], + [ + [2.0, 0.0, 1.0, -1.0, 0.5], + [0.4, 1.4, -0.6, 0.2, 0.8], + [0.9, -0.1, 1.9, 0.3, 0.0], + ], + ], + requires_grad=True, + ) + labels = torch.tensor([[3, 2, -100], [-100, 1, 2]]) + loss_mask = labels != -100 + masks = [ + _sampling_mask([0, 3, 4], [1, 2]), + _sampling_mask([0, 1, 4], [2]), + ] + + actual = replayed_selective_log_softmax( + logits, labels.masked_fill(~loss_mask, 0), loss_mask, masks, 0.7 + ) + expected = _reference_replayed_logps(logits, labels, loss_mask, masks, 0.7) + + torch.testing.assert_close(actual, expected) + assert actual.dtype == torch.float32 + assert torch.equal(actual[~loss_mask], torch.zeros_like(actual[~loss_mask])) + assert actual[1, 2].item() == 0.0 # A singleton support assigns probability one. + + +def test_full_vocab_replay_matches_temperature_scaled_log_softmax(): + torch.manual_seed(7) + logits = torch.randn(2, 3, 6) + labels = torch.tensor([[1, -100, 4], [0, 3, -100]]) + loss_mask = labels != -100 + full_support = list(range(logits.shape[-1])) + masks = [ + _sampling_mask(full_support, full_support), + _sampling_mask(full_support, full_support), + ] + + actual = replayed_selective_log_softmax( + logits, labels.masked_fill(~loss_mask, 0), loss_mask, masks, temperature=1.3 + ) + expected = ( + torch.log_softmax(logits.float() / 1.3, dim=-1) + .gather(-1, labels.masked_fill(~loss_mask, 0).unsqueeze(-1)) + .squeeze(-1) + ) + expected = expected.masked_fill(~loss_mask, 0) + + torch.testing.assert_close(actual, expected) + + +def test_replay_backward_only_touches_retained_support_logits(): + logits = torch.randn(1, 2, 5, requires_grad=True) + labels = torch.tensor([[1, 3]]) + mask = _sampling_mask([0, 1, 4], [2, 3]) + + replayed_selective_log_softmax( + logits, labels, torch.ones_like(labels, dtype=torch.bool), [mask], 1.0 + ).sum().backward() + + assert torch.equal( + logits.grad[0, 0].ne(0), torch.tensor([True, True, False, False, True]) + ) + assert torch.equal( + logits.grad[0, 1].ne(0), torch.tensor([False, False, True, True, False]) + ) + + +def test_empty_training_batch_returns_zeros(): + logits = torch.randn(2, 3, 4) + labels = torch.zeros(2, 3, dtype=torch.long) + loss_mask = torch.zeros_like(labels, dtype=torch.bool) + + result = replayed_selective_log_softmax( + logits, labels, loss_mask, [SamplingMask([], [0]), SamplingMask([], [0])], 1.0 + ) + + assert torch.equal(result, torch.zeros_like(result)) + + +@pytest.mark.parametrize("temperature", [0.0, -1.0, math.inf, math.nan]) +def test_replay_rejects_invalid_temperature(temperature): + with pytest.raises(ValueError, match="temperature"): + replayed_selective_log_softmax( + torch.zeros(1, 1, 2), + torch.zeros(1, 1, dtype=torch.long), + torch.ones(1, 1, dtype=torch.bool), + [_sampling_mask([0])], + temperature, + ) + + +@pytest.mark.parametrize( + ("sampling_mask", "message"), + [ + (None, "missing"), + (SamplingMask([0], [0, 1]), "1 rows but 2 training tokens"), + (SamplingMask([0, 5], [0, 1, 2]), "outside vocabulary"), + (SamplingMask([0, 1], [0, 1, 2]), "absent from sampling mask"), + ], +) +def test_replay_rejects_malformed_or_incompatible_masks(sampling_mask, message): + logits = torch.zeros(1, 2, 3) + labels = torch.tensor([[2, 2]]) + with pytest.raises(ValueError, match=message): + replayed_selective_log_softmax( + logits, + labels, + torch.ones_like(labels, dtype=torch.bool), + [sampling_mask], + 1.0, + ) + + +def test_replay_validates_tensor_and_batch_shapes(): + valid_mask = [_sampling_mask([0])] + with pytest.raises(ValueError, match="logits must have shape"): + replayed_selective_log_softmax( + torch.zeros(1, 2), + torch.zeros(1, 1, dtype=torch.long), + torch.ones(1, 1, dtype=torch.bool), + valid_mask, + 1.0, + ) + with pytest.raises(ValueError, match="labels and loss_mask"): + replayed_selective_log_softmax( + torch.zeros(1, 2, 3), + torch.zeros(1, 1, dtype=torch.long), + torch.ones(1, 1, dtype=torch.bool), + valid_mask, + 1.0, + ) + with pytest.raises(ValueError, match="batch has 0 samples"): + replayed_selective_log_softmax( + torch.zeros(1, 1, 3), + torch.zeros(1, 1, dtype=torch.long), + torch.ones(1, 1, dtype=torch.bool), + [], + 1.0, + ) From b0eee32aa85d4992ed39f80e258d6b9d8570e947 Mon Sep 17 00:00:00 2001 From: vx120 <893600387@qq.com> Date: Tue, 25 Aug 2026 18:01:25 +0800 Subject: [PATCH 7/8] Support sampling replay with packed and padding-free batches Signed-off-by: vx120 <893600387@qq.com> --- .../strategy/sequence_parallel/__init__.py | 8 +- .../model/transformers/transformers.py | 85 +++++++++- src/twinkle/processor/base.py | 8 + .../sampler/vllm_sampler/vllm_engine.py | 103 +++++++++++- src/twinkle/utils/__init__.py | 4 +- src/twinkle/utils/torch_utils.py | 147 +++++++++++++----- src/twinkle/utils/transformers_utils.py | 8 + 7 files changed, 307 insertions(+), 56 deletions(-) diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py index 264503e8e..46ace2c64 100644 --- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py +++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py @@ -11,7 +11,7 @@ from twinkle.patch import apply_patch from twinkle.utils import DeviceMesh -from twinkle.utils.transformers_utils import get_llm_model +from twinkle.utils.transformers_utils import get_llm_model, is_flash_attention_implementation from twinkle.utils.utils import call_with_supported_kwargs, has_signature_parameter from .linear_attention_sp import Qwen3_5GatedDeltaNetUlyssesPatch, _iter_qwen35_gated_delta_net_classes from .utils import (DistributedAttention, GatherLoss, _derive_sequence_parallel_sizes, _get_seq_groups_from_device_mesh, @@ -782,7 +782,7 @@ def pad_and_split_inputs(self, # FlashAttention2 expects a 2D padding mask (or None). Converting it to a 4D causal mask here breaks # the later per-rank sequence split and changes the attention contract relative to the baseline path. if (cache_position is None and hasattr(self, 'causal_mask_func') and self.causal_mask_func is not None - and self.attn_implementation != 'flash_attention_2'): + and not is_flash_attention_implementation(self.attn_implementation)): attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype), local_cache_position, None, None) if extra_split_values is not None: @@ -850,10 +850,10 @@ def prepare_inputs(self, inputs): input_ids = inputs.get('input_ids') position_ids = inputs.get('position_ids') padding_free = bool(inputs.pop('padding_free', False)) - if padding_free and self.attn_implementation not in ('flash_attention_2', 'flash_attention_3'): + if padding_free and not is_flash_attention_implementation(self.attn_implementation): raise RuntimeError('Transformers SequenceParallel does not support padding_free/packed inputs with ' f'attn_implementation={self.attn_implementation!r}. ' - 'Use flash_attention_2 or flash_attention_3, or disable padding_free/packing. ' + 'Use a FlashAttention backend, or disable padding_free/packing. ' 'SDPA/eager attention cannot safely preserve packed sequence boundaries in this path.') real_position_ids = self._extract_real_position_ids(position_ids) if real_position_ids is not None and input_ids is not None and real_position_ids.shape[0] == input_ids.shape[0]: diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index d486b191a..843ff1495 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -39,7 +39,8 @@ from twinkle.patch import Patch, apply_context, apply_patch from twinkle.processor import InputProcessor from twinkle.template import Template -from twinkle.utils import construct_class, get_logger, replayed_selective_log_softmax, selective_log_softmax, torch_util +from twinkle.utils import (construct_class, get_logger, prepare_replayed_selective_log_softmax, + replayed_selective_log_softmax, selective_log_softmax, torch_util) from twinkle.utils.framework import Torch from twinkle.utils.grad_clip import normalize_and_clip_grad_norm from twinkle.utils.transformers_utils import filter_from_config_kwargs @@ -47,6 +48,17 @@ logger = get_logger() +def _get_vocab_size(config: PretrainedConfig) -> int: + """Resolve the LM vocabulary size without running a sharded model forward.""" + vocab_size = getattr(config, 'vocab_size', None) + if vocab_size is None: + text_config = getattr(config, 'text_config', None) + vocab_size = getattr(text_config, 'vocab_size', None) + if not isinstance(vocab_size, int) or vocab_size <= 0: + raise ValueError('sampling replay requires a positive vocab_size in the model config') + return vocab_size + + def _resolve_task_context(model, task): """Return a context manager that applies the right per-forward Patch for ``task``. @@ -223,6 +235,10 @@ def __init__( # Trigger transformers' FSDP-aware loading: meta-device init + rank-0-only weight load. with self.strategy.pretrained_load_context(): self.model = model_cls.from_pretrained(model_id, config=self.hf_config, **kwargs) + # ``from_pretrained``/``from_config`` may deepcopy the supplied config before + # applying runtime overrides such as ``attn_implementation``. Keep Twinkle's + # config reference aligned with the config that the model actually uses. + self.hf_config = self.model.config self.model.gradient_checkpointing_enable() self.sp_strategy = None self._model_wrapped = False @@ -484,6 +500,23 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec enable_sp=getattr(self, '_enable_sp', False), ) labels: torch.Tensor = inputs.pop('labels', None) + replay_metadata = None + replay_loss_mask = None + replay_masked_labels = None + if enable_sampling_replay: + if labels is None: + raise ValueError('labels are required when sampling replay is enabled') + replay_loss_mask = (labels != -100).bool() + replay_masked_labels = labels.masked_fill(~replay_loss_mask, 0) + replay_metadata = prepare_replayed_selective_log_softmax( + labels=labels, + loss_mask=replay_loss_mask, + sampling_masks=sampling_masks, + temperature=temperature, + vocab_size=_get_vocab_size(self.hf_config), + allow_packed_masks=processor.padding_free + or processor._is_packed_position_ids(inputs.get('position_ids')), + ) optimizer_config.accumulate_metrics(True) # Routing replay: respects router_replay_action regardless of caller @@ -502,9 +535,8 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec inputs['labels'] = labels if task != 'embedding' and labels is not None and loss_require_logps: - loss_mask = (labels != -100).bool() - masked_labels = labels.clone() - masked_labels[~loss_mask] = 0 + loss_mask = replay_loss_mask if enable_sampling_replay else (labels != -100).bool() + masked_labels = replay_masked_labels if enable_sampling_replay else labels.masked_fill(~loss_mask, 0) logits = outputs['logits'] if enable_sampling_replay: outputs['logps'] = replayed_selective_log_softmax( @@ -513,6 +545,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec loss_mask=loss_mask, sampling_masks=sampling_masks, temperature=temperature, + metadata=replay_metadata, ) elif loss_require_entropy: logits.div_(temperature) @@ -559,6 +592,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T sampling_masks = kwargs.pop('sampling_masks', None) return_logits = kwargs.pop('return_logits', False) task = kwargs.pop('task', 'causal_lm') + sampling_replay_diagnostics = bool(kwargs.pop('sampling_replay_diagnostics', False)) optimizer_config = self.optimizer_group[adapter_name] self._lazy_wrap_model() if not inputs: @@ -595,6 +629,43 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T enable_sp=getattr(self, '_enable_sp', False), ) labels = inputs.pop('labels', None) + replay_metadata = None + replay_loss_mask = None + replay_masked_labels = None + if enable_sampling_replay: + if labels is None: + raise ValueError('labels are required when sampling replay is enabled') + replay_loss_mask = (labels != -100).bool() + replay_masked_labels = labels.masked_fill(~replay_loss_mask, 0) + replay_metadata = prepare_replayed_selective_log_softmax( + labels=labels, + loss_mask=replay_loss_mask, + sampling_masks=sampling_masks, + temperature=temperature, + vocab_size=_get_vocab_size(self.hf_config), + allow_packed_masks=processor.padding_free + or processor._is_packed_position_ids(inputs.get('position_ids')), + ) + if sampling_replay_diagnostics: + position_ids = inputs.get('position_ids') + logger.info({ + 'sampling_replay/rank': + Platform.get_rank(), + 'sampling_replay/padding_free': + processor.padding_free, + 'sampling_replay/packed_position_ids': + processor._is_packed_position_ids(position_ids), + 'sampling_replay/logical_masks': + len(sampling_masks), + 'sampling_replay/tensor_batch': + labels.shape[0], + 'sampling_replay/tensor_seq_len': + labels.shape[1], + 'sampling_replay/replay_rows': + len(replay_metadata.offsets) - 1, + 'sampling_replay/attn_implementation': + getattr(self.hf_config, '_attn_implementation', None), + }) optimizer_config.accumulate_metrics(False) unwrapped_model = self.strategy.unwrap_model(self.model) @@ -617,9 +688,8 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T inputs['labels'] = labels if task != 'embedding' and labels is not None and loss_require_logps: - loss_mask = (labels != -100).bool() - masked_labels = labels.clone() - masked_labels[~loss_mask] = 0 + loss_mask = replay_loss_mask if enable_sampling_replay else (labels != -100).bool() + masked_labels = replay_masked_labels if enable_sampling_replay else labels.masked_fill(~loss_mask, 0) logits = outputs['logits'] if enable_sampling_replay: outputs['logps'] = replayed_selective_log_softmax( @@ -628,6 +698,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T loss_mask=loss_mask, sampling_masks=sampling_masks, temperature=temperature, + metadata=replay_metadata, ) elif loss_require_entropy: logits.div_(temperature) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 483fb821d..967d35fe2 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -9,6 +9,7 @@ from twinkle import DeviceMesh, Platform, remote_class, remote_function, torch_util from twinkle.data_format import InputFeature +from twinkle.utils.transformers_utils import is_flash_attention_implementation @dataclass @@ -427,6 +428,13 @@ def prepare_transformers_padding_free_patch(self, inputs: List[InputFeature], ** if not padding_free or bool(kwargs.get('enable_sp', False)): return inputs + hf_config = kwargs.get('hf_config') + attn_implementation = getattr(hf_config, '_attn_implementation', None) + if not is_flash_attention_implementation(attn_implementation): + raise RuntimeError('Transformers padding_free/packed batches require a FlashAttention backend; ' + f'got attn_implementation={attn_implementation!r}. SDPA/eager attention cannot isolate ' + 'logical sequences after they are concatenated into one physical row.') + if not getattr(model, '_twinkle_gdn_padding_free_patched', False): from twinkle.patch import apply_patch from twinkle.patch.gdn_padding_free import GatedDeltaNetPaddingFreePatch diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index d4f3448a6..c11e7dd51 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -49,8 +49,61 @@ def _copy_sampling_mask(mask, num_tokens: int, required: bool) -> Optional[Sampl raise RuntimeError('vLLM output is missing sampling mask while sampling replay is enabled') return None - token_ids = [int(token_id) for token_id in mask.token_ids] - offsets = [int(offset) for offset in mask.offsets] + def to_list(value): + if hasattr(value, 'tolist'): + value = value.tolist() + return value + + def flatten_ints(values): + values = to_list(values) + if isinstance(values, (list, tuple)): + flattened = [] + for value in values: + flattened.extend(flatten_ints(value)) + return flattened + return [int(values)] + + raw_token_ids = to_list(mask.token_ids) + token_ids = flatten_ints(raw_token_ids) + + raw_offsets = to_list(getattr(mask, 'offsets', None)) + if raw_offsets is None: + # Newer vLLM sampling-mask variants expose one token-id list per + # generated token instead of a separate CSR offsets field. Preserve + # those row boundaries while converting to Twinkle's CSR format. + if not isinstance(raw_token_ids, (list, tuple)) or len(raw_token_ids) != num_tokens \ + or any(not isinstance(to_list(row), (list, tuple)) for row in raw_token_ids): + raise RuntimeError('vLLM sampling mask without offsets must provide one token-id row ' + 'per sampled token') + offsets = [0] + for row in raw_token_ids: + offsets.append(offsets[-1] + len(flatten_ints(row))) + else: + offsets = flatten_ints(raw_offsets) + + def is_global_csr(candidate_offsets): + return (len(candidate_offsets) == num_tokens + 1 and candidate_offsets[0] == 0 + and candidate_offsets[-1] == len(token_ids) + and all(start < end for start, end in zip(candidate_offsets, candidate_offsets[1:]))) + + # Some vLLM builds preserve per-step/per-chunk tensors as nested lists in + # RequestOutput. A nested offsets field then contains a local CSR per chunk + # (each starts at zero), so compose those local layouts into one global CSR. + if raw_offsets is not None and not is_global_csr(offsets) \ + and isinstance(raw_token_ids, (list, tuple)) \ + and isinstance(raw_offsets, (list, tuple)) and len(raw_token_ids) == len(raw_offsets) \ + and any(isinstance(to_list(offset), (list, tuple)) for offset in raw_offsets): + token_chunks = [flatten_ints(chunk) for chunk in raw_token_ids] + offset_chunks = [flatten_ints(chunk) for chunk in raw_offsets] + composed_offsets = [0] + for chunk_idx, (token_chunk, offset_chunk) in enumerate(zip(token_chunks, offset_chunks)): + if (not offset_chunk or offset_chunk[0] != 0 or offset_chunk[-1] != len(token_chunk) + or any(start >= end for start, end in zip(offset_chunk, offset_chunk[1:]))): + raise RuntimeError(f'vLLM sampling mask chunk {chunk_idx} has invalid CSR offsets') + base_offset = composed_offsets[-1] + composed_offsets.extend(base_offset + offset for offset in offset_chunk[1:]) + offsets = composed_offsets + num_rows = len(offsets) - 1 if num_rows != num_tokens: raise RuntimeError(f'vLLM sampling mask has {num_rows} rows for {num_tokens} sampled tokens') @@ -69,6 +122,50 @@ def _set_sampling_replay_output_kind(vllm_params, enable_sampling_replay: bool) vllm_params.output_kind = RequestOutputKind.FINAL_ONLY +def _validate_sampling_replay_params(sampling_params, extra_kwargs: Dict[str, Any]) -> None: + """Reject score-changing logits processors that a support mask cannot replay.""" + neutral_values = { + 'repetition_penalty': 1.0, + 'presence_penalty': 0.0, + 'frequency_penalty': 0.0, + 'logit_bias': None, + 'logits_processor': None, + 'logits_processors': None, + } + incompatible = [] + for name, neutral in neutral_values.items(): + values = [] + if isinstance(sampling_params, dict): + if name in sampling_params: + values.append(sampling_params[name]) + elif hasattr(sampling_params, name): + values.append(getattr(sampling_params, name)) + if name in extra_kwargs: + values.append(extra_kwargs[name]) + + if neutral is None: + enabled = any(bool(value) for value in values) + else: + enabled = any(value is not None and value != neutral for value in values) + if enabled: + incompatible.append(name) + + if incompatible: + names = ', '.join(incompatible) + raise ValueError('sampling replay does not support score-changing logits processors; ' + f'disable these options: {names}') + + if 'top_k' in extra_kwargs: + top_k = extra_kwargs['top_k'] + elif isinstance(sampling_params, dict): + top_k = sampling_params.get('top_k') + else: + top_k = getattr(sampling_params, 'top_k', None) + if not isinstance(top_k, int) or top_k <= 0: + raise ValueError('sampling distribution replay requires top_k > 0 to bound sampling mask size, ' + 'reduce transfer overhead, and avoid potential OOMs') + + def get_vllm_max_lora_rank(lora_rank: int) -> int: """Get the nearest allowed vLLM LoRA rank.""" from typing import get_args @@ -289,6 +386,8 @@ async def sample(self, from vllm.inputs import TextPrompt, TokensPrompt # Convert to vLLM params + if self.enable_sampling_replay: + _validate_sampling_replay_params(sampling_params, kwargs) if isinstance(sampling_params, dict): sampling_params = SamplingParams.from_dict(sampling_params) prompt_logprobs_k = sampling_params.prompt_logprobs or 0 diff --git a/src/twinkle/utils/__init__.py b/src/twinkle/utils/__init__.py index 53829fa2b..164cd1561 100644 --- a/src/twinkle/utils/__init__.py +++ b/src/twinkle/utils/__init__.py @@ -11,8 +11,8 @@ from .platforms import GPU, NPU, Platform, ensure_hccl_socket_env, ensure_npu_backend from .safetensors import LazyTensor, SafetensorLazyLoader, StreamingSafetensorSaver from .torch_utils import (clone_state_dict_to_cpu, pad_and_stack_tensors, pad_sequence_to_length, - replayed_selective_log_softmax, selective_log_softmax, split_cp_inputs, - stateless_init_process_group, to_device) + prepare_replayed_selective_log_softmax, replayed_selective_log_softmax, selective_log_softmax, + split_cp_inputs, stateless_init_process_group, to_device) from .transformers_utils import find_all_linears, find_layers, get_modules_to_not_convert from .unsafe import check_unsafe, trust_remote_code from .utils import copy_files_by_pattern, deep_getattr, get_runtime_meta diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index f6c7a008e..0d7463b9a 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -1,4 +1,5 @@ import socket +from dataclasses import dataclass from datetime import timedelta from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Union @@ -136,6 +137,99 @@ def selective_log_softmax(logits, index, return_entropy: bool = False): return per_token_logps +@dataclass(frozen=True) +class ReplayedLogSoftmaxMetadata: + """Validated, flattened CSR metadata used by sampling replay.""" + token_ids: tuple[int, ...] + offsets: tuple[int, ...] + + +def prepare_replayed_selective_log_softmax( + labels: 'torch.Tensor', + loss_mask: 'torch.Tensor', + sampling_masks, + temperature: float, + vocab_size: int, + allow_packed_masks: bool = False, +) -> ReplayedLogSoftmaxMetadata: + """Validate and flatten sampling masks before the model forward. + + When ``allow_packed_masks`` is true, a padding-free batch may concatenate + multiple logical masks into one tensor row. Their CSR rows are flattened in + logical-sample order. + """ + import math + import torch + + if not math.isfinite(temperature) or temperature <= 0: + raise ValueError('temperature must be greater than 0 for sampling replay') + if not torch.is_tensor(labels) or labels.dim() != 2: + shape = tuple(labels.shape) if torch.is_tensor(labels) else type(labels) + raise ValueError(f'labels must have shape [batch, seq_len], got {shape}') + if loss_mask.shape != labels.shape: + raise ValueError('loss_mask must have the same shape as labels') + if not isinstance(vocab_size, int) or vocab_size <= 0: + raise ValueError(f'vocab_size must be a positive integer, got {vocab_size!r}') + if not sampling_masks: + raise ValueError('sampling mask batch has 0 samples') + + batch_size = labels.shape[0] + if len(sampling_masks) == batch_size: + mask_groups = [[sampling_mask] for sampling_mask in sampling_masks] + elif batch_size == 1 and allow_packed_masks: + # padding_free/packing collapses multiple logical samples into batch=1. + mask_groups = [list(sampling_masks)] + else: + raise ValueError(f'sampling mask batch has {len(sampling_masks)} samples, expected {batch_size}') + + flat_token_ids = [] + global_offsets = [0] + for batch_idx, mask_group in enumerate(mask_groups): + sampled_labels = labels[batch_idx][loss_mask[batch_idx]].tolist() + parsed_masks = [] + group_rows = 0 + for sampling_mask in mask_group: + if sampling_mask is None: + raise ValueError(f'sampling mask is missing for sample {batch_idx}') + token_ids = [int(token_id) for token_id in sampling_mask.token_ids] + offsets = [int(offset) for offset in sampling_mask.offsets] + if not offsets or offsets[0] != 0 or offsets[-1] != len(token_ids): + raise ValueError(f'sampling mask for sample {batch_idx} has invalid CSR endpoints') + if any(start >= end for start, end in zip(offsets, offsets[1:])): + raise ValueError(f'sampling mask for sample {batch_idx} contains an empty or invalid CSR row') + + invalid_token_id = next( + (token_id for token_id in token_ids if token_id < 0 or token_id >= vocab_size), + None, + ) + if invalid_token_id is not None: + raise ValueError(f'sampling mask token ID {invalid_token_id} is outside vocabulary [0, {vocab_size})') + + base_offset = global_offsets[-1] + flat_token_ids.extend(token_ids) + global_offsets.extend(base_offset + offset for offset in offsets[1:]) + group_rows += len(offsets) - 1 + parsed_masks.append((token_ids, offsets)) + + num_train_tokens = len(sampled_labels) + if group_rows != num_train_tokens: + raise ValueError(f'sampling mask for sample {batch_idx} has {group_rows} rows but ' + f'{num_train_tokens} training tokens') + row_idx = 0 + for token_ids, offsets in parsed_masks: + for start, end in zip(offsets, offsets[1:]): + sampled_label = sampled_labels[row_idx] + if sampled_label < 0 or sampled_label >= vocab_size: + raise ValueError(f'sampled label {sampled_label} is outside vocabulary [0, {vocab_size})') + try: + token_ids.index(sampled_label, start, end) + except ValueError as e: + raise ValueError(f'sampled label {sampled_label} is absent from sampling mask row {row_idx}') from e + row_idx += 1 + + return ReplayedLogSoftmaxMetadata(tuple(flat_token_ids), tuple(global_offsets)) + + # Re-normalize trainer logits over each rollout-time top-p/top-k support set # before reading the sampled token's log probability. Replaying the sampler's # action space removes the sampling/training distribution mismatch in GRPO. @@ -145,45 +239,25 @@ def replayed_selective_log_softmax( loss_mask: 'torch.Tensor', sampling_masks, temperature: float, + metadata: Optional[ReplayedLogSoftmaxMetadata] = None, + allow_packed_masks: bool = False, ) -> 'torch.Tensor': """Compute selected log probabilities on rollout-time CSR support sets.""" - import math import torch - if not math.isfinite(temperature) or temperature <= 0: - raise ValueError('temperature must be greater than 0 for sampling replay') if logits.dim() != 3: raise ValueError(f'logits must have shape [batch, seq_len, vocab], got {tuple(logits.shape)}') if labels.shape != logits.shape[:2] or loss_mask.shape != labels.shape: raise ValueError('labels and loss_mask must match the first two logits dimensions') - if len(sampling_masks) != labels.shape[0]: - raise ValueError(f'sampling mask batch has {len(sampling_masks)} samples, expected {labels.shape[0]}') - - # Flatten per-sample CSR rows into one global CSR layout. - vocab_size = logits.shape[-1] - flat_token_ids = [] - global_offsets = [0] - for batch_idx, sampling_mask in enumerate(sampling_masks): - if sampling_mask is None: - raise ValueError(f'sampling mask is missing for sample {batch_idx}') - token_ids = [int(token_id) for token_id in sampling_mask.token_ids] - offsets = [int(offset) for offset in sampling_mask.offsets] - - num_rows = len(offsets) - 1 - num_train_tokens = int(loss_mask[batch_idx].sum().item()) - if num_rows != num_train_tokens: - raise ValueError(f'sampling mask for sample {batch_idx} has {num_rows} rows but ' - f'{num_train_tokens} training tokens') - invalid_token_id = next( - (token_id for token_id in token_ids if token_id < 0 or token_id >= vocab_size), - None, + if metadata is None: + metadata = prepare_replayed_selective_log_softmax( + labels=labels, + loss_mask=loss_mask, + sampling_masks=sampling_masks, + temperature=temperature, + vocab_size=logits.shape[-1], + allow_packed_masks=allow_packed_masks, ) - if invalid_token_id is not None: - raise ValueError(f'sampling mask token ID {invalid_token_id} is outside vocabulary [0, {vocab_size})') - - base_offset = global_offsets[-1] - flat_token_ids.extend(token_ids) - global_offsets.extend(base_offset + offset for offset in offsets[1:]) # CSR rows are ordered exactly like the masked training-token positions. positions = loss_mask.nonzero(as_tuple=False) @@ -192,24 +266,15 @@ def replayed_selective_log_softmax( if num_rows == 0: return result - offsets_tensor = torch.tensor(global_offsets, dtype=torch.long, device=logits.device) + offsets_tensor = torch.tensor(metadata.offsets, dtype=torch.long, device=logits.device) lengths = offsets_tensor[1:] - offsets_tensor[:-1] row_ids = torch.repeat_interleave( torch.arange(num_rows, device=logits.device), lengths, ) - kept_token_ids = torch.tensor(flat_token_ids, dtype=torch.long, device=logits.device) + kept_token_ids = torch.tensor(metadata.token_ids, dtype=torch.long, device=logits.device) sampled_labels = labels[positions[:, 0], positions[:, 1]].long() - matches = kept_token_ids == sampled_labels[row_ids] - match_counts = torch.zeros(num_rows, dtype=torch.int32, device=logits.device) - match_counts.scatter_add_(0, row_ids, matches.to(torch.int32)) - missing_rows = (match_counts == 0).nonzero(as_tuple=False) - if missing_rows.numel(): - row_idx = int(missing_rows[0].item()) - raise ValueError(f'sampled label {int(sampled_labels[row_idx].item())} is absent from ' - f'sampling mask row {row_idx}') - # Gather only logits retained by the rollout sampler, then normalize per CSR row. kept_logits = logits[ positions[row_ids, 0], diff --git a/src/twinkle/utils/transformers_utils.py b/src/twinkle/utils/transformers_utils.py index 12963c13f..9f997bf70 100644 --- a/src/twinkle/utils/transformers_utils.py +++ b/src/twinkle/utils/transformers_utils.py @@ -9,6 +9,14 @@ import torch.nn as nn +def is_flash_attention_implementation(attn_implementation: Any) -> bool: + """Return whether an attention backend belongs to the FlashAttention family.""" + if not isinstance(attn_implementation, str): + return False + normalized = attn_implementation.strip().lower().replace('-', '_') + return 'flash_attention' in normalized or 'flash_attn' in normalized + + def align_logps_to_mask( ragged: Any, mask: 'torch.Tensor', From 2798fde46b76c5a6f148c168068c0300f37c9226 Mon Sep 17 00:00:00 2001 From: vx120 <893600387@qq.com> Date: Wed, 26 Aug 2026 09:26:13 +0800 Subject: [PATCH 8/8] fix(sampling-replay): normalize vLLM masks and prevalidate packed replay Signed-off-by: vx120 <893600387@qq.com> --- .../model/transformers/transformers.py | 74 +++++----- .../sampler/vllm_sampler/vllm_engine.py | 127 ++++++++++-------- src/twinkle/utils/torch_utils.py | 11 +- 3 files changed, 104 insertions(+), 108 deletions(-) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 648c8fdcc..b561c1664 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -59,6 +59,29 @@ def _get_vocab_size(config: PretrainedConfig) -> int: return vocab_size +def _prepare_sampling_replay( + labels: torch.Tensor, + sampling_masks, + temperature: float, + vocab_size: int, + allow_packed_masks: bool, +): + """Build the labels and metadata shared by training and evaluation replay.""" + if labels is None: + raise ValueError('labels are required when sampling replay is enabled') + loss_mask = (labels != -100).bool() + masked_labels = labels.masked_fill(~loss_mask, 0) + metadata = prepare_replayed_selective_log_softmax( + labels=labels, + loss_mask=loss_mask, + sampling_masks=sampling_masks, + temperature=temperature, + vocab_size=vocab_size, + allow_packed_masks=allow_packed_masks, + ) + return loss_mask, masked_labels, metadata + + def _resolve_task_context(model, task): """Return a context manager that applies the right per-forward Patch for ``task``. @@ -539,22 +562,15 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec enable_sp=getattr(self, '_enable_sp', False), ) labels: torch.Tensor = inputs.pop('labels', None) - replay_metadata = None - replay_loss_mask = None - replay_masked_labels = None + replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: - if labels is None: - raise ValueError('labels are required when sampling replay is enabled') - replay_loss_mask = (labels != -100).bool() - replay_masked_labels = labels.masked_fill(~replay_loss_mask, 0) - replay_metadata = prepare_replayed_selective_log_softmax( + replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay( labels=labels, - loss_mask=replay_loss_mask, sampling_masks=sampling_masks, temperature=temperature, vocab_size=_get_vocab_size(self.hf_config), - allow_packed_masks=processor.padding_free - or processor._is_packed_position_ids(inputs.get('position_ids')), + allow_packed_masks=(processor.padding_free + or processor._is_packed_position_ids(inputs.get('position_ids'))), ) optimizer_config.accumulate_metrics(True) @@ -631,7 +647,6 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T sampling_masks = kwargs.pop('sampling_masks', None) return_logits = kwargs.pop('return_logits', False) task = kwargs.pop('task', 'causal_lm') - sampling_replay_diagnostics = bool(kwargs.pop('sampling_replay_diagnostics', False)) optimizer_config = self.optimizer_group[adapter_name] self._lazy_wrap_model() if not inputs: @@ -668,43 +683,16 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T enable_sp=getattr(self, '_enable_sp', False), ) labels = inputs.pop('labels', None) - replay_metadata = None - replay_loss_mask = None - replay_masked_labels = None + replay_metadata = replay_loss_mask = replay_masked_labels = None if enable_sampling_replay: - if labels is None: - raise ValueError('labels are required when sampling replay is enabled') - replay_loss_mask = (labels != -100).bool() - replay_masked_labels = labels.masked_fill(~replay_loss_mask, 0) - replay_metadata = prepare_replayed_selective_log_softmax( + packed_position_ids = processor._is_packed_position_ids(inputs.get('position_ids')) + replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay( labels=labels, - loss_mask=replay_loss_mask, sampling_masks=sampling_masks, temperature=temperature, vocab_size=_get_vocab_size(self.hf_config), - allow_packed_masks=processor.padding_free - or processor._is_packed_position_ids(inputs.get('position_ids')), + allow_packed_masks=processor.padding_free or packed_position_ids, ) - if sampling_replay_diagnostics: - position_ids = inputs.get('position_ids') - logger.info({ - 'sampling_replay/rank': - Platform.get_rank(), - 'sampling_replay/padding_free': - processor.padding_free, - 'sampling_replay/packed_position_ids': - processor._is_packed_position_ids(position_ids), - 'sampling_replay/logical_masks': - len(sampling_masks), - 'sampling_replay/tensor_batch': - labels.shape[0], - 'sampling_replay/tensor_seq_len': - labels.shape[1], - 'sampling_replay/replay_rows': - len(replay_metadata.offsets) - 1, - 'sampling_replay/attn_implementation': - getattr(self.hf_config, '_attn_implementation', None), - }) optimizer_config.accumulate_metrics(False) unwrapped_model = self.strategy.unwrap_model(self.model) diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index c11e7dd51..51fc4da4a 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -43,67 +43,18 @@ def _filter_engine_config( return filtered_engine_config, invalid_args -def _copy_sampling_mask(mask, num_tokens: int, required: bool) -> Optional[SamplingMask]: - if mask is None: - if required: - raise RuntimeError('vLLM output is missing sampling mask while sampling replay is enabled') - return None +def _to_list(value): + return value.tolist() if hasattr(value, 'tolist') else value - def to_list(value): - if hasattr(value, 'tolist'): - value = value.tolist() - return value - def flatten_ints(values): - values = to_list(values) - if isinstance(values, (list, tuple)): - flattened = [] - for value in values: - flattened.extend(flatten_ints(value)) - return flattened - return [int(values)] +def _as_flat_int_list(values, field_name: str) -> List[int]: + values = _to_list(values) + if not isinstance(values, (list, tuple)) or any(isinstance(_to_list(value), (list, tuple)) for value in values): + raise RuntimeError(f'vLLM sampling mask {field_name} must be a flat integer list') + return [int(value) for value in values] - raw_token_ids = to_list(mask.token_ids) - token_ids = flatten_ints(raw_token_ids) - - raw_offsets = to_list(getattr(mask, 'offsets', None)) - if raw_offsets is None: - # Newer vLLM sampling-mask variants expose one token-id list per - # generated token instead of a separate CSR offsets field. Preserve - # those row boundaries while converting to Twinkle's CSR format. - if not isinstance(raw_token_ids, (list, tuple)) or len(raw_token_ids) != num_tokens \ - or any(not isinstance(to_list(row), (list, tuple)) for row in raw_token_ids): - raise RuntimeError('vLLM sampling mask without offsets must provide one token-id row ' - 'per sampled token') - offsets = [0] - for row in raw_token_ids: - offsets.append(offsets[-1] + len(flatten_ints(row))) - else: - offsets = flatten_ints(raw_offsets) - - def is_global_csr(candidate_offsets): - return (len(candidate_offsets) == num_tokens + 1 and candidate_offsets[0] == 0 - and candidate_offsets[-1] == len(token_ids) - and all(start < end for start, end in zip(candidate_offsets, candidate_offsets[1:]))) - - # Some vLLM builds preserve per-step/per-chunk tensors as nested lists in - # RequestOutput. A nested offsets field then contains a local CSR per chunk - # (each starts at zero), so compose those local layouts into one global CSR. - if raw_offsets is not None and not is_global_csr(offsets) \ - and isinstance(raw_token_ids, (list, tuple)) \ - and isinstance(raw_offsets, (list, tuple)) and len(raw_token_ids) == len(raw_offsets) \ - and any(isinstance(to_list(offset), (list, tuple)) for offset in raw_offsets): - token_chunks = [flatten_ints(chunk) for chunk in raw_token_ids] - offset_chunks = [flatten_ints(chunk) for chunk in raw_offsets] - composed_offsets = [0] - for chunk_idx, (token_chunk, offset_chunk) in enumerate(zip(token_chunks, offset_chunks)): - if (not offset_chunk or offset_chunk[0] != 0 or offset_chunk[-1] != len(token_chunk) - or any(start >= end for start, end in zip(offset_chunk, offset_chunk[1:]))): - raise RuntimeError(f'vLLM sampling mask chunk {chunk_idx} has invalid CSR offsets') - base_offset = composed_offsets[-1] - composed_offsets.extend(base_offset + offset for offset in offset_chunk[1:]) - offsets = composed_offsets +def _validate_sampling_mask_csr(token_ids: List[int], offsets: List[int], num_tokens: int) -> None: num_rows = len(offsets) - 1 if num_rows != num_tokens: raise RuntimeError(f'vLLM sampling mask has {num_rows} rows for {num_tokens} sampled tokens') @@ -111,6 +62,68 @@ def is_global_csr(candidate_offsets): raise RuntimeError('vLLM sampling mask has invalid CSR endpoints') if any(start >= end for start, end in zip(offsets, offsets[1:])): raise RuntimeError('vLLM sampling mask contains an empty or invalid CSR row') + + +def _parse_flat_csr(raw_token_ids, raw_offsets) -> tuple[List[int], List[int]]: + return ( + _as_flat_int_list(raw_token_ids, 'token_ids'), + _as_flat_int_list(raw_offsets, 'offsets'), + ) + + +def _parse_token_rows(raw_token_ids, num_tokens: int) -> tuple[List[int], List[int]]: + rows = _to_list(raw_token_ids) + if not isinstance(rows, (list, tuple)) or len(rows) != num_tokens: + raise RuntimeError('vLLM sampling mask without offsets must provide one token-id row ' + 'per sampled token') + + token_ids = [] + offsets = [0] + for row in rows: + token_row = _as_flat_int_list(row, 'token-id row') + token_ids.extend(token_row) + offsets.append(offsets[-1] + len(token_row)) + return token_ids, offsets + + +def _parse_chunked_csr(raw_token_ids, raw_offsets) -> tuple[List[int], List[int]]: + token_chunks = _to_list(raw_token_ids) + offset_chunks = _to_list(raw_offsets) + if not isinstance(token_chunks, (list, tuple)) or not isinstance(offset_chunks, (list, tuple)) \ + or len(token_chunks) != len(offset_chunks): + raise RuntimeError('vLLM chunked sampling mask must have matching token-id and offset chunks') + + token_ids = [] + offsets = [0] + for chunk_idx, (token_chunk, offset_chunk) in enumerate(zip(token_chunks, offset_chunks)): + chunk_token_ids = _as_flat_int_list(token_chunk, f'token_ids chunk {chunk_idx}') + chunk_offsets = _as_flat_int_list(offset_chunk, f'offsets chunk {chunk_idx}') + if (not chunk_offsets or chunk_offsets[0] != 0 or chunk_offsets[-1] != len(chunk_token_ids) + or any(start >= end for start, end in zip(chunk_offsets, chunk_offsets[1:]))): + raise RuntimeError(f'vLLM sampling mask chunk {chunk_idx} has invalid CSR offsets') + base_offset = offsets[-1] + token_ids.extend(chunk_token_ids) + offsets.extend(base_offset + offset for offset in chunk_offsets[1:]) + return token_ids, offsets + + +def _copy_sampling_mask(mask, num_tokens: int, required: bool) -> Optional[SamplingMask]: + if mask is None: + if required: + raise RuntimeError('vLLM output is missing sampling mask while sampling replay is enabled') + return None + + raw_token_ids = _to_list(mask.token_ids) + raw_offsets = _to_list(getattr(mask, 'offsets', None)) + if raw_offsets is None: + token_ids, offsets = _parse_token_rows(raw_token_ids, num_tokens) + elif any(isinstance(_to_list(offset), (list, tuple)) for offset in raw_offsets): + # Some vLLM builds preserve local CSR tensors for each output chunk. + token_ids, offsets = _parse_chunked_csr(raw_token_ids, raw_offsets) + else: + token_ids, offsets = _parse_flat_csr(raw_token_ids, raw_offsets) + + _validate_sampling_mask_csr(token_ids, offsets, num_tokens) return SamplingMask(token_ids=token_ids, offsets=offsets) diff --git a/src/twinkle/utils/torch_utils.py b/src/twinkle/utils/torch_utils.py index 0d7463b9a..42bdaa0c9 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -152,11 +152,11 @@ def prepare_replayed_selective_log_softmax( vocab_size: int, allow_packed_masks: bool = False, ) -> ReplayedLogSoftmaxMetadata: - """Validate and flatten sampling masks before the model forward. + """Align and flatten sampling masks before the model forward. When ``allow_packed_masks`` is true, a padding-free batch may concatenate - multiple logical masks into one tensor row. Their CSR rows are flattened in - logical-sample order. + multiple logical masks into one tensor row. Their sampler-validated CSR + rows are flattened in logical-sample order. """ import math import torch @@ -193,11 +193,6 @@ def prepare_replayed_selective_log_softmax( raise ValueError(f'sampling mask is missing for sample {batch_idx}') token_ids = [int(token_id) for token_id in sampling_mask.token_ids] offsets = [int(offset) for offset in sampling_mask.offsets] - if not offsets or offsets[0] != 0 or offsets[-1] != len(token_ids): - raise ValueError(f'sampling mask for sample {batch_idx} has invalid CSR endpoints') - if any(start >= end for start, end in zip(offsets, offsets[1:])): - raise ValueError(f'sampling mask for sample {batch_idx} contains an empty or invalid CSR row') - invalid_token_id = next( (token_id for token_id in token_ids if token_id < 0 or token_id >= vocab_size), None,