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 781b22060..52479b92b 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 017a515b3..5afc46c0a 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') assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' inputs: Dict[str, Any] = processor( inputs, @@ -497,11 +505,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 +552,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 +575,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') inputs: Dict[str, Any] = processor( inputs, sp_strategy=self.sp_strategy, @@ -591,11 +616,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..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 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..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',