diff --git a/cookbook/rl/grpo/grpo_sampling_replay.py b/cookbook/rl/grpo/grpo_sampling_replay.py new file mode 100644 index 000000000..3fba27562 --- /dev/null +++ b/cookbook/rl/grpo/grpo_sampling_replay.py @@ -0,0 +1,266 @@ +import os +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 +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): + """Flatten sampler responses into aligned lists used by reward and training.""" + rollout_batch = { + 'input_data': [], + 'old_logps': [], + 'sampling_masks': [], + 'completion_lengths': [], + } + 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') + 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)) + 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, + ) + optim_step = 0 + 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 + ckpt_manager.sync_weights(merge_and_sync=False) + sampler.reset_prefix_cache() + def sample_prompt_groups(prompts): + expand_prompts = [] + for prompt in prompts: + expand_prompts.extend([prompt] * NUM_GENERATIONS) + responses = sampler.sample(expand_prompts, sampling_params) + return extract_rollout_batch(responses) + + 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 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'] + 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() + + # 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, + } + + 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() + 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)) + 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 81e0b9208..96199f71f 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -32,12 +32,18 @@ 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 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 @@ -222,6 +228,9 @@ def __call__( **kwargs: Additional arguments """ import torch + if self.enable_sampling_replay: + 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 +239,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/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index ffbcebccd..d486b191a 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,13 @@ 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') + 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') loss_require_values = getattr(loss_instance, 'require_values', False) assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' inputs: Dict[str, Any] = processor( @@ -498,11 +506,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 if loss_require_values: @@ -539,6 +556,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] @@ -561,6 +579,13 @@ 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') + 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') loss_require_values = getattr(loss_instance, 'require_values', False) inputs: Dict[str, Any] = processor( inputs, @@ -596,11 +621,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 if loss_require_values: diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index b1e1790de..d4f3448a6 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,46 @@ 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 +118,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 +138,8 @@ 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 +172,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 +219,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 +294,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 +342,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 +375,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 f1e6d4095..d0bdd2e13 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.require_values = getattr(loss_instance, 'require_values', 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..f6c7a008e 100644 --- a/src/twinkle/utils/torch_utils.py +++ b/src/twinkle/utils/torch_utils.py @@ -136,6 +136,107 @@ 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', + 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]}') + + # 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 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) + num_rows = positions.shape[0] + 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}') + + # Gather only logits retained by the rollout sampler, then normalize per CSR row. + 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 + + # Use max-shifted log-sum-exp for numerically stable restricted softmax. + 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', 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, + )