diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 5bb4600c57a..a138e98576c 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -781,7 +781,7 @@ def compute_log_prob( num_microbatches: list[int], store_prefix: str = "", ) -> dict[str, list[torch.Tensor]]: - + use_rollout_sampling_masks = "rollout_sampling_masks" in data_iterator[0].rollout_data with timer(f"{store_prefix}log_probs"): return forward_only( get_log_probs_and_entropy, @@ -790,6 +790,7 @@ def compute_log_prob( data_iterator, num_microbatches, store_prefix=store_prefix, + use_rollout_sampling_masks=use_rollout_sampling_masks, ) def train(self, rollout_id: int, rollout_data_ref: Box) -> None: diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 85097426ded..316a1d2975c 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -194,6 +194,7 @@ def forward_only( data_iterator: Sequence[DataIterator], num_microbatches: Sequence[int], store_prefix: str = "", + use_rollout_sampling_masks: bool = False, ) -> dict[str, list[torch.Tensor]]: """Run forward passes only and collect non-loss outputs (e.g., logprobs). @@ -239,16 +240,18 @@ def forward_step( assert not return_schedule_plan, "forward_only step should never return schedule plan" # Get the batch. + keys = [ + "tokens", + "loss_masks", + "multimodal_train_inputs", + "total_lengths", + "response_lengths", + "max_seq_lens", + "rollout_sampling_masks", + ] batch = get_batch( data_iterator, - [ - "tokens", - "loss_masks", - "multimodal_train_inputs", - "total_lengths", - "response_lengths", - "max_seq_lens", - ], + keys, args.data_pad_size_multiplier, args.qkv_format, allgather_cp=args.allgather_cp, @@ -268,15 +271,17 @@ def forward_step( **(batch["multimodal_train_inputs"] if batch["multimodal_train_inputs"] is not None else {}), ) - return output_tensor, partial( - f, - args=args, - unconcat_tokens=unconcat_tokens, - total_lengths=total_lengths, - response_lengths=response_lengths, - with_entropy=args.use_rollout_entropy, - max_seq_lens=batch.get("max_seq_lens", None), - ) + callback_kwargs = { + "args": args, + "unconcat_tokens": unconcat_tokens, + "total_lengths": total_lengths, + "response_lengths": response_lengths, + "with_entropy": args.use_rollout_entropy, + "max_seq_lens": batch.get("max_seq_lens", None), + } + if use_rollout_sampling_masks: + callback_kwargs["rollout_sampling_masks"] = batch["rollout_sampling_masks"] + return output_tensor, partial(f, **callback_kwargs) # Turn on evaluation mode which disables dropout. for model_module in model: @@ -381,23 +386,25 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p """ # Get the batch. + keys = [ + "tokens", + "multimodal_train_inputs", + "packed_seq_params", + "total_lengths", + "response_lengths", + "loss_masks", + "log_probs", + "ref_log_probs", + "values", + "advantages", + "returns", + "rollout_log_probs", + "max_seq_lens", + "rollout_sampling_masks", + ] batch = get_batch( data_iterator, - [ - "tokens", - "multimodal_train_inputs", - "packed_seq_params", - "total_lengths", - "response_lengths", - "loss_masks", - "log_probs", - "ref_log_probs", - "values", - "advantages", - "returns", - "rollout_log_probs", - "max_seq_lens", - ], + keys, args.data_pad_size_multiplier, args.qkv_format, allgather_cp=args.allgather_cp, diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 1c215c7b62a..0a670868c6f 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -680,6 +680,9 @@ def _compute_server_args( for key in unused_keys: kwargs.pop(key) + if args.num_rollout != 0 and args.rollout_top_p < 1 and kwargs.get("disaggregation_mode", "null") != "null": + raise ValueError("sampling-mask replay is incompatible with prefill/decode disaggregation") + return kwargs, external_engine_need_check_fields diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index 305c1bdc668..625bcd33e6c 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -21,7 +21,7 @@ get_reinforce_plus_plus_baseline_advantages, get_reinforce_plus_plus_returns, ) -from miles.utils.types import RolloutBatch +from miles.utils.types import RolloutBatch, RolloutSamplingMask from .cp_utils import ( _allgather_cp_redistribute, @@ -129,8 +129,8 @@ def get_responses( total_lengths: list[int], response_lengths: list[int], max_seq_lens: list[int] | None = None, -) -> Iterator[tuple[torch.Tensor, torch.Tensor]]: - """Yield response-aligned `(logits_chunk, tokens_chunk)` pairs per sample. +) -> Iterator[tuple[torch.Tensor, torch.Tensor, list[tuple[int, int]]]]: + """Yield response logits, targets, and response-row spans per sample. After squeezing batch dimension and applying temperature scaling, this function extracts the logits and tokens corresponding to response segments @@ -149,9 +149,10 @@ def get_responses( response_lengths: Response segment lengths per sample. Yields: - Tuple of `(logits_chunk, tokens_chunk)` where `logits_chunk` is shape + Tuple of `(logits_chunk, tokens_chunk, response_spans)` where `logits_chunk` is shape `[R, V]` (policy) or `[R, 1]` (value) and `tokens_chunk` is shape `[R]` (1D int64), both aligned to response tokens for one sample. + `response_spans` contains ordered half-open row ranges in the full response. """ parallel_state = get_parallel_state() qkv_format = args.qkv_format @@ -183,7 +184,8 @@ def get_responses( end += total_length start = end - response_length logits_chunk = logits[start - 1 : end - 1] - tokens_chunk = tokens[-response_length:] + tokens_chunk = tokens[total_length - response_length : total_length] + response_spans = [(0, response_length)] if response_length else [] elif args.allgather_cp: # DSA: global concat then contiguous CP split. Each rank owns logits for # global positions [chunk_start, chunk_end). @@ -203,9 +205,11 @@ def get_responses( if e <= s: logits_chunk = logits[0:0] tokens_chunk = tokens[0:0] + response_spans = [] else: logits_chunk = logits[s - chunk_start : e - chunk_start] tokens_chunk = tokens[(s + 1) - seq_start : (e + 1) - seq_start] + response_spans = [(s - logit_global_start, e - logit_global_start)] assert logits_chunk.size(0) == tokens_chunk.size(0), f"{logits_chunk.size(0)} vs {tokens_chunk.size(0)}" else: # TODO: this is super ugly... do better abstraction. @@ -213,6 +217,11 @@ def get_responses( total_length, response_length, qkv_format, max_seq_len ) + prompt_length = total_length - response_length + response_spans = [ + (start - prompt_length, stop - prompt_length) for start, stop in tokens_offset if start < stop + ] + logits_0, logits_1 = logits[end : end + chunk_size], logits[end + chunk_size : end + 2 * chunk_size] end += 2 * chunk_size @@ -230,7 +239,36 @@ def get_responses( seq_start += total_length - yield logits_chunk, tokens_chunk + yield logits_chunk, tokens_chunk, response_spans + + +def _build_tp_sampling_mask( + logits: torch.Tensor, + sampling_mask: RolloutSamplingMask, + response_spans: list[tuple[int, int]], + tp_rank: int, +) -> torch.Tensor: + mask = torch.zeros(logits.shape, dtype=torch.bool, device=logits.device) + if not response_spans: + return mask + + offsets = sampling_mask.offsets + sizes = torch.cat([offsets[start + 1 : stop + 1] - offsets[start:stop] for start, stop in response_spans]) + parts = [sampling_mask.token_ids[offsets[start].item() : offsets[stop].item()] for start, stop in response_spans] + token_ids = parts[0] if len(parts) == 1 else torch.cat(parts) + + vocab_size = logits.size(-1) + vocab_start = tp_rank * vocab_size + owned_entries = ((token_ids >= vocab_start) & (token_ids < vocab_start + vocab_size)).nonzero(as_tuple=True)[0] + if owned_entries.numel() == 0: + return mask + + rows = torch.repeat_interleave(sizes, output_size=token_ids.numel())[owned_entries] + cols = token_ids[owned_entries].long() - vocab_start + rows = rows.to(logits.device, non_blocking=True) + cols = cols.to(logits.device, non_blocking=True) + mask[rows, cols] = True + return mask def get_log_probs_and_entropy( @@ -243,6 +281,7 @@ def get_log_probs_and_entropy( with_entropy: bool = False, non_loss_data: bool = True, max_seq_lens: list[int] | None = None, + rollout_sampling_masks: list[RolloutSamplingMask] | None = None, ) -> dict[str, list[torch.Tensor]]: """Compute per-token log-probabilities (and optionally entropy) on responses. @@ -260,6 +299,7 @@ def get_log_probs_and_entropy( response_lengths: Response segment lengths per sample. with_entropy: If True, include "entropy" key in result. non_loss_data: Unused; kept for API compatibility. + rollout_sampling_masks: Fixed rollout support per full response; entropy remains unrestricted. Returns: Dict with key "log_probs" mapping to a list of `[R]` tensors per @@ -268,16 +308,24 @@ def get_log_probs_and_entropy( """ parallel_state = get_parallel_state() assert non_loss_data + log_probs_list = [] entropy_list = [] - for logits_chunk, tokens_chunk in get_responses( - logits, - args=args, - unconcat_tokens=unconcat_tokens, - total_lengths=total_lengths, - response_lengths=response_lengths, - max_seq_lens=max_seq_lens, + for i, (logits_chunk, tokens_chunk, response_spans) in enumerate( + get_responses( + logits, + args=args, + unconcat_tokens=unconcat_tokens, + total_lengths=total_lengths, + response_lengths=response_lengths, + max_seq_lens=max_seq_lens, + ) ): + sampling_mask = None + if rollout_sampling_masks is not None: + sampling_mask = _build_tp_sampling_mask( + logits_chunk, rollout_sampling_masks[i], response_spans, parallel_state.tp.rank + ) log_prob, entropy = calculate_log_probs_and_entropy( logits_chunk, tokens_chunk, @@ -285,9 +333,10 @@ def get_log_probs_and_entropy( with_entropy=with_entropy, chunk_size=args.log_probs_chunk_size, true_on_policy=args.true_on_policy_mode, + sampling_mask=sampling_mask, ) - log_probs_list.append(log_prob.squeeze(-1)) + log_probs_list.append(log_prob.reshape(-1)) entropy_list.append(entropy) res = { @@ -341,7 +390,7 @@ def get_values( per sample. """ value_list = [] - for logits_chunk, _ in get_responses( + for logits_chunk, _, _ in get_responses( logits, args=args, unconcat_tokens=unconcat_tokens, @@ -649,6 +698,7 @@ def policy_loss_function( response_lengths=response_lengths, with_entropy=True, max_seq_lens=max_seq_lens, + rollout_sampling_masks=batch.get("rollout_sampling_masks"), ) log_probs = log_probs_and_entropy["log_probs"] diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 0d462d46cdf..d6dfae1141b 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -824,6 +824,9 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl if samples[0].rollout_log_probs is not None: train_data["rollout_log_probs"] = [sample.rollout_log_probs for sample in samples] + if self.args.rollout_top_p < 1: + train_data["rollout_sampling_masks"] = [sample.rollout_sampling_mask for sample in samples] + if samples[0].rollout_routed_experts is not None: train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples] @@ -953,6 +956,7 @@ def _stat(xs): "round_number", "sample_indices", "rollout_log_probs", + "rollout_sampling_masks", "rollout_routed_experts", "prompt", "teacher_log_probs", diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index cfefeccd0e1..68ec9174cc2 100644 --- a/miles/rollout/generate_hub/agentic_tool_call.py +++ b/miles/rollout/generate_hub/agentic_tool_call.py @@ -49,7 +49,8 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: "agentic_tool_call.generate requires session_server_ip/session_server_port. " "Pass --use-session-server to start the session server." ) - tracer = await OpenAIEndpointTracer.create(input.args) + capture = not input.evaluation and input.sampling_params["top_p"] < 1 + tracer = await OpenAIEndpointTracer.create(input.args, capture_sampling_mask=capture) custom_agent_function: Callable = load_function(input.args.custom_agent_function_path) assert ( @@ -114,6 +115,8 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: sample.status = Sample.Status.ABORTED return GenerateFnOutput(samples=sample) + if capture and merged_sample.rollout_sampling_mask is None: + raise ValueError("a replay trajectory must carry a sampling mask") sample = apply_merged_session_sample(input.args, input.sample, merged_sample) sample.metadata.update(agent_metadata or {}) sample.metadata.update(session_metadata) diff --git a/miles/rollout/generate_utils/openai_endpoint_utils.py b/miles/rollout/generate_utils/openai_endpoint_utils.py index f1cb76f1565..42449ae25c6 100644 --- a/miles/rollout/generate_utils/openai_endpoint_utils.py +++ b/miles/rollout/generate_utils/openai_endpoint_utils.py @@ -19,7 +19,7 @@ MergedSessionSample, SessionRecord, ) -from miles.utils.types import Sample +from miles.utils.types import RolloutSamplingMask, Sample logger = logging.getLogger(__name__) @@ -288,7 +288,7 @@ async def _request( raise RuntimeError(f"request failed without exception: phase={phase} method={method} url={url}") @staticmethod - async def create(args: Namespace): + async def create(args: Namespace, *, capture_sampling_mask: bool = False): backends = getattr(args, "session_server_backends", None) if backends: session_url = random.choice(backends).rstrip("/") @@ -336,7 +336,7 @@ async def create(args: Namespace): "POST", f"{session_url}/sessions", phase="create_session", - payload={}, + payload={"capture_sampling_mask": capture_sampling_mask}, max_retries=_CREATE_RETRIES, ) @@ -629,6 +629,11 @@ def apply_merged_session_sample( sample.response_length = merged.response_length sample.loss_mask = merged.loss_mask sample.rollout_log_probs = merged.rollout_log_probs + sample.rollout_sampling_mask = ( + RolloutSamplingMask.from_dict(merged.rollout_sampling_mask) + if merged.rollout_sampling_mask is not None + else None + ) sample.metadata = {**(sample.metadata or {}), **(merged.metadata or {})} sample.status = Sample.Status(merged.status) sample.weight_versions.extend(merged.weight_versions) @@ -743,6 +748,11 @@ def _compute_sample_from_openai_record( output_token_ids = [item[1] for item in choice["meta_info"]["output_token_logprobs"]] output_log_probs = [item[0] for item in choice["meta_info"]["output_token_logprobs"]] + rollout_sampling_mask = None + if record.rollout_sampling_mask is not None: + rollout_sampling_mask = RolloutSamplingMask.from_dict(record.rollout_sampling_mask) + output_log_probs = record.rollout_sampling_log_probs + sample = copy(input_sample) sample.metadata = dict(input_sample.metadata) sample.weight_versions = list(input_sample.weight_versions) @@ -755,6 +765,7 @@ def _compute_sample_from_openai_record( sample.tokens = prompt_token_ids + output_token_ids sample.rollout_log_probs = output_log_probs + sample.rollout_sampling_mask = rollout_sampling_mask sample.response = "" sample.response_length = len(output_token_ids) sample.loss_mask = [1] * len(output_token_ids) @@ -800,6 +811,8 @@ def _strip_last_output_tokens_without_decode(sample: Sample, trim_count: int) -> if sample.rollout_log_probs is not None: sample.rollout_log_probs = sample.rollout_log_probs[:keep_tokens] + if sample.rollout_sampling_mask is not None: + sample.rollout_sampling_mask = sample.rollout_sampling_mask.prefix(keep_tokens) if sample.loss_mask is not None: sample.loss_mask = sample.loss_mask[:keep_tokens] if sample.rollout_routed_experts is not None: @@ -849,6 +862,8 @@ def _truncate_sample_output(sample: Sample, keep_tokens: int, tokenizer) -> None if sample.rollout_log_probs is not None: sample.rollout_log_probs = sample.rollout_log_probs[:keep_tokens] + if sample.rollout_sampling_mask is not None: + sample.rollout_sampling_mask = sample.rollout_sampling_mask.prefix(keep_tokens) if sample.loss_mask is not None: sample.loss_mask = sample.loss_mask[:keep_tokens] if sample.rollout_routed_experts is not None: diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 5844f09b91b..0c7b9913342 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -1,7 +1,7 @@ from copy import copy from dataclasses import fields -from miles.utils.types import Sample +from miles.utils.types import RolloutSamplingMask, Sample def drop_samples_after_first_non_completed(samples: list[Sample]) -> tuple[list[Sample], int]: @@ -23,8 +23,17 @@ def drop_samples_after_first_non_completed(samples: list[Sample]) -> tuple[list[ def merge_samples(samples: list[Sample], tokenizer) -> Sample: acc = samples[0] + # Concatenate masks once after all turns. + mask_parts = [] + if len(samples) > 1 and acc.rollout_sampling_mask is not None: + mask_parts.append(acc.rollout_sampling_mask) for sample in samples[1:]: + if mask_parts: + obs_tokens = sample.tokens[len(acc.tokens) : len(sample.tokens) - sample.response_length] + mask_parts.extend((RolloutSamplingMask.singletons(obs_tokens), sample.rollout_sampling_mask)) acc = _merge_sample_pair(acc, sample, tokenizer=tokenizer) + if mask_parts: + acc.rollout_sampling_mask = RolloutSamplingMask.concatenate(*mask_parts) return acc @@ -80,6 +89,7 @@ def _fill_defaults(sample: Sample): loss_mask=a.loss_mask + [0] * obs_len + b.loss_mask, weight_versions=a.weight_versions + b.weight_versions, rollout_log_probs=a.rollout_log_probs + [0.0] * obs_len + b.rollout_log_probs, + rollout_sampling_mask=None, rollout_routed_experts=b.rollout_routed_experts, remove_sample=_merge_equal_value("remove_sample"), status=b.status, diff --git a/miles/rollout/session/linear_trajectory.py b/miles/rollout/session/linear_trajectory.py index 881d1304d26..9c5afa5e686 100644 --- a/miles/rollout/session/linear_trajectory.py +++ b/miles/rollout/session/linear_trajectory.py @@ -54,6 +54,7 @@ class LinearTrajectory: records: list[SessionRecord] = field(default_factory=list) trajectory_token_ids: list[list[int]] = field(default_factory=list) num_assistant: int = 0 + capture_sampling_mask: bool = False @property def token_ids(self) -> list[int]: @@ -312,9 +313,9 @@ def __init__(self, args, tokenizer: Any, *, tito_tokenizer: TITOTokenizer): self.tito_tokenizer = tito_tokenizer self.comparator = tito_tokenizer.create_comparator() - def create_session(self) -> str: + def create_session(self, *, capture_sampling_mask: bool = False) -> str: session_id = uuid.uuid4().hex - self.sessions[session_id] = LinearTrajectory() + self.sessions[session_id] = LinearTrajectory(capture_sampling_mask=capture_sampling_mask) return session_id def get_session(self, session_id: str) -> LinearTrajectory: diff --git a/miles/rollout/session/session_types.py b/miles/rollout/session/session_types.py index bff5da91b9f..53bf43eccea 100644 --- a/miles/rollout/session/session_types.py +++ b/miles/rollout/session/session_types.py @@ -8,6 +8,8 @@ class SessionRecord(BaseModel): request: dict response: dict status_code: int + rollout_sampling_mask: dict[str, str] | None = None + rollout_sampling_log_probs: list[float] | None = None class GetSessionResponse(BaseModel): @@ -22,6 +24,7 @@ class MergedSessionSample(BaseModel): response_length: int loss_mask: list[int] rollout_log_probs: list[float] + rollout_sampling_mask: dict[str, str] | None = None status: str metadata: dict = Field(default_factory=lambda: {"response_decoded": False}) weight_versions: list[str] = Field(default_factory=list) diff --git a/miles/rollout/session/sessions.py b/miles/rollout/session/sessions.py index 27dd94615dd..7c90c2ab4b2 100644 --- a/miles/rollout/session/sessions.py +++ b/miles/rollout/session/sessions.py @@ -5,7 +5,7 @@ import orjson import pybase64 -from fastapi import Request +from fastapi import Body, Request from fastapi.responses import JSONResponse from starlette.responses import Response @@ -19,6 +19,7 @@ ) from miles.utils.chat_template_utils import get_tito_tokenizer from miles.utils.processing_utils import load_tokenizer +from miles.utils.types import RolloutSamplingMask logger = logging.getLogger(__name__) @@ -31,6 +32,19 @@ _worker_stats: dict = {} +def _compact_sampling_replay(meta_info: dict) -> tuple[dict[str, str], list[float]]: + try: + masks = meta_info.pop("output_token_sampling_mask") + log_probs = meta_info.pop("output_token_sampling_logprobs") + if not (len(masks) == len(log_probs) == meta_info["completion_tokens"] and all(masks)): + raise ValueError("expected one nonempty mask and logprob per output token") + mask = RolloutSamplingMask.from_rows(masks).to_dict() + meta_info.pop("output_token_sampling_mask_length", None) + return mask, [float(value) for value in log_probs] + except (KeyError, TypeError, ValueError, RuntimeError, OverflowError) as exc: + raise UpstreamResponseError(f"invalid sampling-mask metadata: {exc}") from exc + + def get_worker_stats(port): """Return the live stats dict for the given worker port. @@ -101,8 +115,8 @@ async def session_error_handler(request: Request, exc: SessionError): return JSONResponse(status_code=exc.status_code, content={"error": str(exc)}) @app.post("/sessions") - async def create_session(): - session_id = registry.create_session() + async def create_session(capture_sampling_mask: bool = Body(False, embed=True)): + session_id = registry.create_session(capture_sampling_mask=capture_sampling_mask) return {"session_id": session_id} def _compact_output_token_logprobs(output_token_logprobs): @@ -176,6 +190,8 @@ def _compact_session_record(record: SessionRecord) -> SessionRecord: response={ "choices": [compact_choice], }, + rollout_sampling_mask=record.rollout_sampling_mask, + rollout_sampling_log_probs=record.rollout_sampling_log_probs, ) def _status_from_finish_reason(finish_reason: str | None) -> str: @@ -231,6 +247,7 @@ def _merge_session_records_to_sample( records: list[SessionRecord], accumulated_token_ids: list[int], max_trim_tokens: int, + capture_sampling_mask: bool, ) -> MergedSessionSample | None: if not records: return None @@ -241,6 +258,7 @@ def _merge_session_records_to_sample( first_prompt_len: int | None = None loss_mask_full = [0] * len(tokens) rollout_log_probs_full = [0.0] * len(tokens) + sampling_mask_parts = [] weight_versions: list[str] = [] prefix_cache_meta_infos: list[dict] = [] routed_experts = None @@ -302,6 +320,17 @@ def _merge_session_records_to_sample( f"is_last={is_last}, max_trim_tokens={max_trim_tokens}" ) + if capture_sampling_mask: + mask = RolloutSamplingMask.from_dict(record.rollout_sampling_mask) + obs_start = len(prev_checkpoint) if prev_checkpoint is not None else first_prompt_len + sampling_mask_parts.extend( + ( + RolloutSamplingMask.singletons(tokens[obs_start:cursor]), + mask.prefix(matched) if trim_count else mask, + ) + ) + output_log_probs = record.rollout_sampling_log_probs + for j in range(matched): idx = cursor + j assert idx < len(tokens), f"record {i}: output token index {idx} outside accumulated tokens" @@ -364,6 +393,9 @@ def _merge_session_records_to_sample( response_length=response_length, loss_mask=loss_mask_full[first_prompt_len:], rollout_log_probs=rollout_log_probs_full[first_prompt_len:], + rollout_sampling_mask=( + RolloutSamplingMask.concatenate(*sampling_mask_parts).to_dict() if capture_sampling_mask else None + ), status=final_status, metadata={"response_decoded": False}, weight_versions=weight_versions, @@ -417,6 +449,7 @@ async def get_merged_session(session_id: str): records_to_merge, accumulated_token_ids, max_trim_tokens, + session.capture_sampling_mask, ) metadata = { "records_total": len(records), @@ -547,6 +580,7 @@ async def chat_completions(request: Request, session_id: str): request_body["logprobs"] = True request_body["return_prompt_token_ids"] = True request_body["return_meta_info"] = True + request_body["return_sampling_mask"] = session.capture_sampling_mask if getattr(args, "use_rollout_routing_replay", False): request_body["return_routed_experts"] = True # Must be False so stop tokens are trimmed from output: otherwise the @@ -656,6 +690,14 @@ async def chat_completions(request: Request, session_id: str): completion_token_ids = [t[1] for t in output_token_logprobs] + rollout_sampling_mask = None + rollout_sampling_log_probs = None + if session.capture_sampling_mask: + rollout_sampling_mask, rollout_sampling_log_probs = await asyncio.to_thread( + _compact_sampling_replay, meta_info + ) + result["response_body"] = orjson.dumps(response) + async with session.lock: if session.closing: logger.warning(f"Session {session_id} closed during proxy, skipping state update") @@ -703,6 +745,8 @@ async def chat_completions(request: Request, session_id: str): status_code=result["status_code"], request=request_body, response=response, + rollout_sampling_mask=rollout_sampling_mask, + rollout_sampling_log_probs=rollout_sampling_log_probs, ) session.append_record(record) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index a9af7d4fb7b..fcaf18c1887 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1742,6 +1742,23 @@ def dump_args_to_file(args, main_log_dir=None): f.write("-------------------- end of arguments ---------------------\n") +def _validate_sampling_mask_replay(args) -> None: + if args.rollout_top_p >= 1 or args.num_rollout == 0: + return + if args.train_backend != "megatron": + raise ValueError("rollout_top_p < 1 sampling-mask replay requires the Megatron backend") + if not args.use_session_server: + raise ValueError("rollout_top_p < 1 sampling-mask replay requires the session server") + if args.custom_generate_function_path != "miles.rollout.generate_hub.agentic_tool_call.generate": + raise ValueError("sampling-mask replay requires miles.rollout.generate_hub.agentic_tool_call.generate") + if args.true_on_policy_mode: + raise ValueError("sampling-mask replay is incompatible with --true-on-policy-mode") + if args.sglang_pp_size > 1: + raise ValueError("sampling-mask replay requires SGLang pipeline parallel size 1") + if args.sglang_speculative_algorithm is not None: + raise ValueError("sampling-mask replay is incompatible with speculative decoding") + + def parse_args(add_custom_arguments=None): # Users may call `parse_args` very early, thus we ensure logger is configured here configure_logger() @@ -1795,6 +1812,7 @@ def parse_args(add_custom_arguments=None): ), "decoder_first_pipeline_num_layers and decoder_last_pipeline_num_layers should be None when pipeline_model_parallel_size is 1." sglang_validate_args(args) + _validate_sampling_mask_replay(args) dump_args_to_file(args) return args diff --git a/miles/utils/ppo_utils.py b/miles/utils/ppo_utils.py index 3b40e4bb094..9e04c142aac 100644 --- a/miles/utils/ppo_utils.py +++ b/miles/utils/ppo_utils.py @@ -668,7 +668,13 @@ def chunked_gae( def calculate_log_probs_and_entropy( - logits, tokens, tp_group, with_entropy: bool = False, chunk_size: int = -1, true_on_policy: bool = False + logits, + tokens, + tp_group, + with_entropy: bool = False, + chunk_size: int = -1, + true_on_policy: bool = False, + sampling_mask: torch.Tensor | None = None, ): if true_on_policy: return _calculate_log_probs_and_entropy_true_on_policy( @@ -687,11 +693,22 @@ def calculate_log_probs_and_entropy( num_chunks = (logits.size(0) - 1) // chunk_size + 1 tokens_chunks = tokens.chunk(num_chunks, dim=0) logits_chunks = logits.chunk(num_chunks, dim=0) + sampling_mask_chunks = ( + sampling_mask.chunk(num_chunks, dim=0) if sampling_mask is not None else [None] * num_chunks + ) log_probs = [] - for tokens_chunk, logits_chunk in zip(tokens_chunks, logits_chunks, strict=True): + for tokens_chunk, logits_chunk, mask_chunk in zip( + tokens_chunks, logits_chunks, sampling_mask_chunks, strict=True + ): if logits_chunk.dtype != torch.float32: logits_chunk = logits_chunk.float() - log_prob = compute_log_probs(logits_chunk.clone(), tokens_chunk, tp_group) + log_prob_logits = logits_chunk.clone() + if mask_chunk is not None: + # CE gradients are already zero outside the recorded support; + # no_grad avoids saving the dense mask for backward. + with torch.no_grad(): + log_prob_logits.masked_fill_(~mask_chunk, float("-inf")) + log_prob = compute_log_probs(log_prob_logits, tokens_chunk, tp_group) log_probs.append(log_prob) log_prob = torch.cat(log_probs, dim=0) if with_entropy: @@ -705,7 +722,11 @@ def calculate_log_probs_and_entropy( else: if logits.dtype != torch.float32: logits = logits.float() - log_prob = compute_log_probs(logits.clone(), tokens, tp_group) + log_prob_logits = logits.clone() + if sampling_mask is not None: + with torch.no_grad(): + log_prob_logits.masked_fill_(~sampling_mask, float("-inf")) + log_prob = compute_log_probs(log_prob_logits, tokens, tp_group) if with_entropy: entropy = compute_entropy_from_logits(logits.clone(), tp_group) else: diff --git a/miles/utils/types.py b/miles/utils/types.py index 4d7d6ef9b2f..2c470704b5e 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -1,3 +1,5 @@ +import base64 +from collections.abc import Sequence from dataclasses import dataclass, field from enum import Enum from typing import Any @@ -6,6 +8,67 @@ import torch +@dataclass(frozen=True, eq=False) +class RolloutSamplingMask: + """CPU CSR (int32 IDs, int64 offsets). Treat tensors as immutable; inputs may share storage.""" + + token_ids: torch.Tensor + offsets: torch.Tensor + + def __len__(self) -> int: + return self.offsets.numel() - 1 + + @classmethod + def from_rows(cls, rows: Sequence[Sequence[int]]) -> "RolloutSamplingMask": + token_ids: list[int] = [] + offsets = [0] + for row in rows: + token_ids.extend(row) + offsets.append(len(token_ids)) + return cls( + torch.tensor(token_ids, dtype=torch.int32, device="cpu"), + torch.tensor(offsets, dtype=torch.int64, device="cpu"), + ) + + def to_dict(self) -> dict[str, str]: + ids = self.token_ids.numpy().astype(" "RolloutSamplingMask": + ids = numpy.frombuffer(base64.b64decode(data["ids"], validate=True), dtype=" "RolloutSamplingMask": + token_ids = torch.cat([mask.token_ids for mask in masks]) + shifted_offsets = [] + cursor = 0 + for mask in masks: + shifted_offsets.append(mask.offsets[1:] + cursor) + cursor += mask.token_ids.numel() + offsets = torch.cat([torch.zeros(1, dtype=torch.int64, device="cpu"), *shifted_offsets]) + return cls(token_ids, offsets) + + @classmethod + def singletons(cls, token_ids: Sequence[int] | torch.Tensor) -> "RolloutSamplingMask": + token_ids = torch.as_tensor(token_ids, dtype=torch.int32, device="cpu") + return cls(token_ids, torch.arange(len(token_ids) + 1, dtype=torch.int64, device="cpu")) + + def prefix(self, rows: int) -> "RolloutSamplingMask": + end = self.offsets[rows].item() + return type(self)(self.token_ids[:end].clone(), self.offsets[: rows + 1].clone()) + + @dataclass class Sample: """The sample generated""" @@ -25,6 +88,7 @@ class Sample: loss_mask: list[int] | None = None weight_versions: list[str] = field(default_factory=list) rollout_log_probs: list[float] | None = None # Log probabilities from rollout engine + rollout_sampling_mask: RolloutSamplingMask | None = None rollout_routed_experts: numpy.ndarray | None = ( None # Routed experts from rollout engine. shape: (num_tokens-1, num_layers, moe_router_topk), dtype=int32 ) @@ -124,6 +188,9 @@ def from_dict(data: dict): def to_dict(self): value = self.__dict__.copy() + value["rollout_sampling_mask"] = ( + self.rollout_sampling_mask.to_dict() if self.rollout_sampling_mask is not None else None + ) value["status"] = self.status.value value["spec_info"] = self.spec_info.to_dict() value["prefix_cache_info"] = self.prefix_cache_info.to_dict() @@ -132,6 +199,8 @@ def to_dict(self): @staticmethod def from_dict(data: dict): data = dict(data) + if data.get("rollout_sampling_mask") is not None: + data["rollout_sampling_mask"] = RolloutSamplingMask.from_dict(data["rollout_sampling_mask"]) data["status"] = Sample.Status(data["status"]) data["spec_info"] = Sample.SpecInfo.from_dict(data.get("spec_info", {})) data["prefix_cache_info"] = Sample.PrefixCacheInfo.from_dict(data.get("prefix_cache_info", {})) @@ -180,6 +249,8 @@ def strip_last_output_tokens(self, n: int, tokenizer) -> None: ), f"cannot strip {n} tokens: only {self.response_length} output tokens available" self.tokens = self.tokens[:-n] self.response_length -= n + if self.rollout_sampling_mask is not None: + self.rollout_sampling_mask = self.rollout_sampling_mask.prefix(self.response_length) if self.rollout_log_probs is not None: self.rollout_log_probs = self.rollout_log_probs[:-n] if self.loss_mask is not None: @@ -203,6 +274,7 @@ def reset_for_retry(self) -> None: self.loss_mask = None self.weight_versions = [] self.rollout_log_probs = None + self.rollout_sampling_mask = None self.rollout_routed_experts = None self.status = Sample.Status.ABORTED self.non_generation_time = 0.0 @@ -254,7 +326,10 @@ class ParamInfo: # A dict-based batch produced along the rollout -> training path # In Megatron backend, several fields are converted to torch.Tensor lists on GPU # before being consumed by data iterators (see megatron_utils.actor._get_rollout_data). -RolloutBatch = dict[str, list[torch.Tensor] | list[int] | list[float] | list[str]] +RolloutBatch = dict[ + str, + list[torch.Tensor] | list[RolloutSamplingMask] | list[int] | list[float] | list[str], +] @dataclass diff --git a/tests/e2e/precision/test_top_p_sampling_replay.py b/tests/e2e/precision/test_top_p_sampling_replay.py new file mode 100644 index 00000000000..fb3f2d8b74a --- /dev/null +++ b/tests/e2e/precision/test_top_p_sampling_replay.py @@ -0,0 +1,128 @@ +import os +from argparse import Namespace +from itertools import product +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.distributed as dist + +from miles.backends.training_utils.loss import get_log_probs_and_entropy +from miles.utils.types import RolloutSamplingMask + + +def verify_top_p_replay() -> None: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + assert world_size == 4 + + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", device_id=device) + + tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])] + cp_groups = [dist.new_group([0, 2]), dist.new_group([1, 3])] + cp_rank, tp_rank = divmod(rank, 2) + parallel_state = SimpleNamespace( + tp=SimpleNamespace(size=2, rank=tp_rank, group=tp_groups[cp_rank]), + cp=SimpleNamespace(size=2, rank=cp_rank, group=cp_groups[tp_rank]), + ) + + full_logits = torch.randn(16, 8, generator=torch.Generator().manual_seed(7)).to(device).requires_grad_() + tokens = [torch.arange(8, dtype=torch.long, device=device), torch.tensor([0, 1, 2, 3, 4, 5, 6, 3], device=device)] + supports = [ + [[0, 2, 5], [3, 4], [1, 4, 7], [0, 5], [2, 6, 7], [1, 3, 7]], + [[0, 3]], + ] + # Two independently packed sequences. The second response's sole logit is + # position 6, so CP rank 1 has no response rows for that sample. + cp_positions = [0, 1, 6, 7, 8, 9, 14, 15] if cp_rank == 0 else [2, 3, 4, 5, 10, 11, 12, 13] + vocab_start = tp_rank * 4 + local_logits = ( + full_logits[cp_positions, vocab_start : vocab_start + 4].unsqueeze(0).contiguous().detach().requires_grad_() + ) + args = Namespace( + qkv_format="thd", + rollout_temperature=0.7, + allgather_cp=False, + log_probs_chunk_size=1, + true_on_policy_mode=False, + ) + + saved_dense_masks = [] + + def pack(tensor): + if tensor.dtype == torch.bool and tensor.ndim == 2 and tensor.size(-1) == local_logits.size(-1): + saved_dense_masks.append(tuple(tensor.shape)) + return tensor + + # Also cover singleton observations and support wholly owned by the other TP rank. + singleton_supports = [[[2], [0, 3], [4, 7], [5], [6], [7]], [[3]]] + for rows, chunk_size in product((None, supports, singleton_supports), (-1, 1, 3)): + sampling_masks = [RolloutSamplingMask.from_rows(sample_rows) for sample_rows in rows] if rows else None + args.log_probs_chunk_size = chunk_size + saved_dense_masks.clear() + + with ( + torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor), + patch("miles.backends.training_utils.loss.get_parallel_state", return_value=parallel_state), + patch("miles.backends.training_utils.cp_utils.get_parallel_state", return_value=parallel_state), + ): + result = get_log_probs_and_entropy( + local_logits, + args=args, + unconcat_tokens=tokens, + total_lengths=[8, 8], + response_lengths=[6, 1], + with_entropy=True, + rollout_sampling_masks=sampling_masks, + ) + + assert not saved_dense_masks, f"dense sampling masks retained for backward: {saved_dense_masks}" + + owned_rows = [[0, 5], [0]] if cp_rank == 0 else [[1, 2, 3, 4], []] + scaled_logits = full_logits / args.rollout_temperature + expected = {"log_probs": [], "entropy": []} + for sample_idx, (response_length, logit_start) in enumerate(((6, 1), (1, 14))): + log_probs, entropy = [], [] + for response_row in owned_rows[sample_idx]: + support = torch.tensor( + rows[sample_idx][response_row] if rows is not None else range(full_logits.size(-1)), + dtype=torch.long, + device=device, + ) + full_row = scaled_logits[logit_start + response_row] + selected = torch.nonzero(support == tokens[sample_idx][-response_length + response_row]).item() + log_probs.append(torch.log_softmax(full_row[support], dim=-1)[selected]) + entropy.append(-(torch.softmax(full_row, dim=-1) * torch.log_softmax(full_row, dim=-1)).sum()) + + for key, values in (("log_probs", log_probs), ("entropy", entropy)): + expected[key].append(torch.stack(values) if values else full_logits.new_empty(0)) + torch.testing.assert_close(result[key][sample_idx], expected[key][-1], atol=1e-5, rtol=1e-5) + + for key in expected: + actual_gradient = torch.autograd.grad(torch.cat(result[key]).sum(), local_logits, retain_graph=True)[0] + full_gradient = torch.autograd.grad(torch.cat(expected[key]).sum(), full_logits, retain_graph=True)[0] + expected_gradient = full_gradient[cp_positions, vocab_start : vocab_start + 4].unsqueeze(0) + # The existing fused CE returns BF16 gradients; entropy gradients stay FP32. + tolerance = 4e-3 if key == "log_probs" else 1e-5 + torch.testing.assert_close(actual_gradient, expected_gradient, atol=tolerance, rtol=tolerance) + if key == "log_probs": + assert torch.count_nonzero(actual_gradient[expected_gradient == 0]).item() == 0 + + dist.barrier() + if rank == 0: + print("TP2 x CP2 top-p replay PASS") + + +def main() -> None: + try: + verify_top_p_replay() + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/fast/backends/megatron_utils/test_top_p_replay.py b/tests/fast/backends/megatron_utils/test_top_p_replay.py new file mode 100644 index 00000000000..d3c24e8c599 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_top_p_replay.py @@ -0,0 +1,140 @@ +from argparse import Namespace +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import torch + +from miles.backends.megatron_utils.actor import MegatronTrainRayActor +from miles.backends.megatron_utils.model import train_one_step +from miles.backends.training_utils.data import DataIterator, get_rollout_data +from miles.backends.training_utils.parallel import GroupInfo, ParallelState +from miles.ray.rollout import RolloutManager +from miles.utils.types import RolloutSamplingMask, Sample + + +def test_replay_mask_reaches_scoring_and_training(): + samples = [ + Sample( + index=i, + tokens=[i + 1] + response, + response_length=len(response), + rollout_log_probs=[-0.1 * (i + 1)] * len(response), + rollout_sampling_mask=RolloutSamplingMask.from_rows(support), + ) + for i, (response, support) in enumerate( + [([4], [[4]]), ([5], [[1, 5]]), ([7, 8], [[7, 9], [2, 8]]), ([9], [[0, 9]])] + ) + ] + args = Namespace( + loss_type="policy_loss", + rollout_top_p=0.9, + balance_data=False, + data_pad_size_multiplier=1, + qkv_format="bshd", + allgather_cp=False, + use_rollout_entropy=False, + use_rollout_logprobs=True, + use_dynamic_batch_size=True, + use_dynamic_global_batch_size=False, + calculate_per_token_loss=False, + recompute_loss_function=False, + enable_mtp_training=False, + fp16=False, + bf16=False, + dumper_enable=False, + dumper_fwd_only=[], + dumper_fwd_bwd=[], + custom_megatron_before_log_prob_hook_path=None, + custom_megatron_before_train_step_hook_path=None, + seq_length=3, + micro_batch_size=1, + global_batch_size=4, + decoder_seq_length=None, + ci_test=False, + use_opsm=False, + advantage_estimator="grpo", + eps_clip=0.2, + eps_clip_high=0.2, + get_mismatch_metrics=False, + use_tis=False, + entropy_coef=0.0, + use_kl_loss=False, + ) + manager_cls = RolloutManager.__ray_metadata__.modified_class + manager = object.__new__(manager_cls) + manager.args = args + manager.custom_convert_samples_to_train_data_func = None + manager._post_process_rewards = lambda _: ([0.0] * 4, [0.0] * 4) + group = GroupInfo(rank=0, size=1, group=None) + dp_group = GroupInfo(rank=0, size=2, group=None) + parallel = ParallelState(intra_dp=dp_group, intra_dp_cp=dp_group, cp=group, tp=group) + scored = [] + + def forward(*, input_ids, **kwargs): + return torch.zeros(*input_ids.shape, 10, requires_grad=True) + + def schedule(*, forward_step_func, data_iterator, model, num_microbatches, forward_only, **kwargs): + results = [] + for _ in range(num_microbatches): + output, callback = forward_step_func(data_iterator[0], model[0]) + result = callback(output) + if not forward_only: + loss, _, result = result + loss.backward() + assert output.grad is not None and output.grad.abs().sum() > 0 + results.append(result) + return results + + def score(logits, *, unconcat_tokens, response_lengths, rollout_sampling_masks=None, **kwargs): + assert len(unconcat_tokens) == 1 + sample = samples[unconcat_tokens[0][0].item() - 1] + assert unconcat_tokens[0].tolist() == sample.tokens + assert response_lengths == [sample.response_length] + assert rollout_sampling_masks == [sample.rollout_sampling_mask] + scored.append(sample.index) + log_probs = logits[0, : sample.response_length, 0] - 0.1 * (sample.index + 1) + return {"log_probs": [log_probs], "entropy": [log_probs * 0]} + + model = MagicMock(config=SimpleNamespace(), side_effect=forward) + optimizer = MagicMock() + optimizer.step.return_value = (True, 1.0, 0) + with ( + patch("miles.backends.training_utils.parallel._parallel_state", parallel), + patch("torch.cuda.current_device", return_value="cpu"), + patch("miles.ray.rollout.ray.put", side_effect=lambda data: SimpleNamespace(data=data, hex=lambda: "test")), + patch("miles.utils.data.ray.get", side_effect=lambda ref: ref.data), + patch("miles.backends.megatron_utils.model.get_args", return_value=args), + patch("miles.backends.megatron_utils.model.get_forward_backward_func", return_value=schedule), + patch("miles.backends.megatron_utils.model.mpu.is_pipeline_last_stage", return_value=True), + patch("miles.backends.megatron_utils.model.clear_memory"), + patch("miles.backends.training_utils.log_utils.dist.all_reduce"), + patch("miles.backends.megatron_utils.actor.get_log_probs_and_entropy", side_effect=score), + patch("miles.backends.training_utils.loss.get_log_probs_and_entropy", side_effect=score), + ): + converted = manager._convert_samples_to_train_data(samples) + shards = manager._split_train_data_by_dp(converted, dp_size=2) + for rank, shard in enumerate(shards): + assert shard.inner.data["sample_indices"] == [rank, rank + 2] + assert shard.inner.data["rollout_sampling_masks"] == [ + samples[i].rollout_sampling_mask for i in (rank, rank + 2) + ] + data = get_rollout_data(args, shards) + assert data["sample_indices"] == [0, 2] + assert data["total_lengths"] == [2, 3] + data["advantages"] = [torch.ones(length) for length in data["response_lengths"]] + iterator = DataIterator(data, micro_batch_indices=[[1], [0]]) + actor = object.__new__(MegatronTrainRayActor) + actor.args, actor.model = args, [model] + for prefix in ("", "ref_"): + scored.clear() + result = actor.compute_log_prob([iterator], [2], store_prefix=prefix) + assert scored == [2, 0] + assert f"{prefix}log_probs" in result + for actual, expected in zip(result[f"{prefix}log_probs"], data["rollout_log_probs"], strict=True): + torch.testing.assert_close(actual, expected) + data.update(result) + + scored.clear() + train_one_step(args, 0, 0, [iterator.reset()], [model], optimizer, MagicMock(), 2) + assert scored == [2, 0] + optimizer.step.assert_called_once() diff --git a/tests/fast/fixtures/generation_fixtures.py b/tests/fast/fixtures/generation_fixtures.py index 7e70dcd187a..3ab994f9963 100644 --- a/tests/fast/fixtures/generation_fixtures.py +++ b/tests/fast/fixtures/generation_fixtures.py @@ -26,7 +26,7 @@ MODEL_NAME = "Qwen/Qwen3-0.6B" RESPONSE_TEXT = "\\boxed{8}" -DEFAULT_SAMPLING_PARAMS = {"max_new_tokens": 64, "temperature": 0.7} +DEFAULT_SAMPLING_PARAMS = {"max_new_tokens": 64, "temperature": 0.7, "top_p": 1.0} VARIANT_TO_GENERATE_FN_PATH = { "old_sglang_rollout": "miles.rollout.sglang_rollout.generate", diff --git a/tests/fast/rollout/generate_hub/test_agentic_turn_budget.py b/tests/fast/rollout/generate_hub/test_agentic_turn_budget.py index 61e870edd65..f76f44e5aad 100644 --- a/tests/fast/rollout/generate_hub/test_agentic_turn_budget.py +++ b/tests/fast/rollout/generate_hub/test_agentic_turn_budget.py @@ -70,7 +70,7 @@ def test_generate_preserves_server_merged_sample(monkeypatch, exit_status, empty GenerateFnInput( state=SimpleNamespace(args=args, tokenizer=None), sample=original, - sampling_params={}, + sampling_params={"top_p": 1.0}, evaluation=False, ) ) diff --git a/tests/fast/rollout/generate_hub/test_multi_turn.py b/tests/fast/rollout/generate_hub/test_multi_turn.py index 96d01e7efda..dd237e35b70 100644 --- a/tests/fast/rollout/generate_hub/test_multi_turn.py +++ b/tests/fast/rollout/generate_hub/test_multi_turn.py @@ -6,7 +6,14 @@ import numpy as np import pybase64 import pytest -from tests.fast.fixtures.generation_fixtures import GenerateEnv, generation_env, listify, make_sample, run_generate +from tests.fast.fixtures.generation_fixtures import ( + DEFAULT_SAMPLING_PARAMS, + GenerateEnv, + generation_env, + listify, + make_sample, + run_generate, +) from miles.utils.processing_utils import load_tokenizer from miles.utils.test_utils.mock_sglang_server import ProcessResult, ProcessResultMetaInfo @@ -24,7 +31,6 @@ def is_agentic_variant(variant: str) -> bool: MODEL_NAME = "Qwen/Qwen3-0.6B" -DEFAULT_SAMPLING_PARAMS = {"max_new_tokens": 64, "temperature": 0.7} TOKENIZER = load_tokenizer(MODEL_NAME, trust_remote_code=True) @@ -157,6 +163,7 @@ def expected_openai_request(messages: list[dict], **extra) -> dict: "return_prompt_token_ids": True, "return_meta_info": True, "no_stop_trim": False, + "top_p": DEFAULT_SAMPLING_PARAMS["top_p"], **extra, } diff --git a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py index 6d90fdd38b9..00f854a3c32 100644 --- a/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py +++ b/tests/fast/rollout/generate_utils/test_openai_endpoint_utils.py @@ -5,17 +5,20 @@ """ from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest +from miles.rollout.base_types import GenerateFnInput +from miles.rollout.generate_hub.agentic_tool_call import generate from miles.rollout.generate_utils.openai_endpoint_utils import ( OpenAIEndpointTracer, compute_samples_from_openai_records, + truncate_samples_by_total_tokens, ) from miles.rollout.generate_utils.sample_utils import merge_samples -from miles.rollout.session.session_types import SessionRecord -from miles.utils.types import Sample +from miles.rollout.session.session_types import MergedSessionSample, SessionRecord +from miles.utils.types import RolloutSamplingMask, Sample # ── helpers ────────────────────────────────────────────────────────── @@ -92,30 +95,149 @@ def _make_record( ) +def _make_replay_record( + prompt_token_ids: list[int], + output_token_ids: list[int], + sampling_rows: list[list[int]], + sampling_log_probs: list[float], +) -> SessionRecord: + record = _make_record( + prompt_token_ids=prompt_token_ids, + output_token_ids=output_token_ids, + output_log_probs=[-9.0] * len(output_token_ids), + ) + record.rollout_sampling_mask = RolloutSamplingMask.from_rows(sampling_rows).to_dict() + record.rollout_sampling_log_probs = sampling_log_probs + return record + + +def _mask_rows(mask: RolloutSamplingMask) -> list[list[int]]: + return [mask.token_ids[mask.offsets[i] : mask.offsets[i + 1]].tolist() for i in range(len(mask))] + + +def test_sampling_replay_trajectory_integrity(): + tokenizer = _mock_tokenizer() + records = [ + _make_replay_record([1, 2], [10, 11, 99], [[10, 12], [11], [99]], [-0.4, 0.0, -1.0]), + _make_replay_record( + [1, 2, 10, 11, 20, 21], + [30, 31], + [[30, 40], [31, 41]], + [-0.7, -0.8], + ), + ] + + turns = compute_samples_from_openai_records( + _ARGS, + _make_input_sample(), + records, + tokenizer, + accumulated_token_ids=[1, 2, 10, 11, 20, 21, 30, 31], + max_trim_tokens=1, + ) + assert turns[0].rollout_log_probs == [-0.4, 0.0] + assert _mask_rows(turns[0].rollout_sampling_mask) == [[10, 12], [11]] + + merged = merge_samples(turns, tokenizer) + assert merged.loss_mask == [1, 1, 0, 0, 1, 1] + assert merged.rollout_log_probs == [-0.4, 0.0, 0.0, 0.0, -0.7, -0.8] + assert _mask_rows(merged.rollout_sampling_mask) == [ + [10, 12], + [11], + [20], + [21], + [30, 40], + [31, 41], + ] + + round_trip = Sample.from_dict(merged.to_dict()) + round_trip.validate() + truncated = truncate_samples_by_total_tokens( + [round_trip], + max_seq_len=len(round_trip.tokens) - 1, + tokenizer=tokenizer, + )[0] + assert truncated.rollout_log_probs == [-0.4, 0.0, 0.0, 0.0, -0.7] + assert _mask_rows(truncated.rollout_sampling_mask) == [[10, 12], [11], [20], [21], [30, 40]] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("evaluation", [False, True]) +@pytest.mark.parametrize("malformed", [False, True], ids=["merged", "malformed"]) +async def test_generate_sampling_capture_and_collection_errors(monkeypatch, evaluation, malformed): + merged = MergedSessionSample( + tokens=[1, 2, 10], + response_length=1, + loss_mask=[1], + rollout_log_probs=[0.0], + status="truncated", + ) + + async def agent(**kwargs): + return None + + tracer = OpenAIEndpointTracer("http://session", "test") + tracer._request = AsyncMock( + return_value={"session_id": "test", "sample": {"response_length": 1} if malformed else merged.model_dump()} + ) + tracer.delete_session = AsyncMock() + captures = [] + + async def create(args, *, capture_sampling_mask=False): + captures.append(capture_sampling_mask) + return tracer + + monkeypatch.setattr(OpenAIEndpointTracer, "create", create) + monkeypatch.setattr("miles.rollout.generate_hub.agentic_tool_call.load_function", lambda _: agent) + args = SimpleNamespace( + session_server_ip="session", + session_server_port=80, + custom_agent_function_path="agent", + generate_multi_samples=False, + max_seq_len=3, + ) + input = GenerateFnInput( + SimpleNamespace(args=args, tokenizer=_mock_tokenizer()), + _make_input_sample(), + {"top_p": 0.9}, + evaluation, + ) + + if malformed: + assert (await generate(input)).samples.status == Sample.Status.ABORTED + elif evaluation: + assert (await generate(input)).samples.tokens == [1, 2, 10] + else: + with pytest.raises(ValueError, match="sampling mask"): + await generate(input) + assert captures == [not evaluation] + tracer.delete_session.assert_awaited_once() + + @pytest.mark.asyncio async def test_create_fetches_session_server_instance_id(monkeypatch): - calls: list[tuple[str, str]] = [] + calls: list[tuple[str, str, dict]] = [] - async def fake_post(url: str, payload: dict, action: str = "post"): - calls.append((action, url)) - if action == "get": + async def fake_request(method, url, *, payload=None, **kwargs): + calls.append((method, url, payload)) + if method == "GET": assert url == "http://127.0.0.1:12345/health" return {"status": "ok", "session_server_instance_id": "server-instance-123"} - assert action == "post" + assert method == "POST" assert url == "http://127.0.0.1:12345/sessions" return {"session_id": "session-123"} - monkeypatch.setattr("miles.rollout.generate_utils.openai_endpoint_utils.post", fake_post) + monkeypatch.setattr(OpenAIEndpointTracer, "_request", fake_request) args = SimpleNamespace(session_server_ip="127.0.0.1", session_server_port=12345) - tracer = await OpenAIEndpointTracer.create(args) + tracer = await OpenAIEndpointTracer.create(args, capture_sampling_mask=True) assert tracer.base_url == "http://127.0.0.1:12345/sessions/session-123" assert tracer.session_server_instance_id == "server-instance-123" assert args.session_server_instance_id == "server-instance-123" assert calls == [ - ("get", "http://127.0.0.1:12345/health"), - ("post", "http://127.0.0.1:12345/sessions"), + ("GET", "http://127.0.0.1:12345/health", None), + ("POST", "http://127.0.0.1:12345/sessions", {"capture_sampling_mask": True}), ] diff --git a/tests/fast/router/test_sessions.py b/tests/fast/router/test_sessions.py index 875f5116830..f37dacb85b9 100644 --- a/tests/fast/router/test_sessions.py +++ b/tests/fast/router/test_sessions.py @@ -10,11 +10,17 @@ import pytest import requests +from miles.rollout.generate_utils.openai_endpoint_utils import ( + apply_merged_session_sample, + truncate_samples_by_total_tokens, +) from miles.rollout.session.linear_trajectory import LinearTrajectory from miles.rollout.session.session_server import SessionServer +from miles.rollout.session.session_types import MergedSessionSample from miles.utils.http_utils import find_available_port from miles.utils.test_utils.mock_sglang_server import MockSGLangServer, ProcessResult, with_mock_server from miles.utils.test_utils.uvicorn_thread_server import UvicornThreadServer +from miles.utils.types import Sample @pytest.fixture(scope="class") @@ -37,6 +43,13 @@ def patched_chat_response(self, payload: dict) -> dict: "output_token_logprobs": output_token_logprobs, "completion_tokens": len(output_token_logprobs), } + if payload.get("return_sampling_mask") and payload.get("model") != "omit-sampling-mask": + output_ids = [token_id for _, token_id in output_token_logprobs] + choice["meta_info"].update( + output_token_sampling_mask=[[token_id] for token_id in output_ids], + output_token_sampling_logprobs=[0.0] * len(output_ids), + output_token_sampling_mask_length=len(output_ids), + ) return response with patch.object(MockSGLangServer, "_compute_chat_completions_response", new=patched_chat_response): @@ -47,6 +60,8 @@ def patched_chat_response(self, payload: dict) -> dict: chat_template_path=None, trajectory_manager="linear_trajectory", session_server_instance_id=uuid.uuid4().hex, + rollout_temperature=1.0, + tito_allowed_append_roles=["system"], ) server_obj = SessionServer(args, backend_url=backend.url) @@ -57,7 +72,7 @@ def patched_chat_response(self, payload: dict) -> dict: url = f"http://127.0.0.1:{port}" try: - yield SimpleNamespace(url=url) + yield SimpleNamespace(url=url, backend=backend) finally: server.stop() @@ -151,6 +166,68 @@ def test_proxy_chat_appends_record(self, router_env): assert record["path"] == "/v1/chat/completions" assert record["status_code"] == 200 + def test_top_p_sampling_replay_contract(self, router_env): + capture_fields = { + "output_token_sampling_mask", + "output_token_sampling_logprobs", + "output_token_sampling_mask_length", + } + session_id = requests.post( + f"{router_env.url}/sessions", json={"capture_sampling_mask": True}, timeout=5.0 + ).json()["session_id"] + + def chat(**overrides): + payload = { + "messages": [{"role": "user", "content": "What is 1+2?"}], + "temperature": 1.0, + "top_p": 0.9, + } + payload.update(overrides) + return requests.post( + f"{router_env.url}/sessions/{session_id}/v1/chat/completions", + json=payload, + timeout=10.0, + ) + + response = chat(return_sampling_mask=False) + assert response.status_code == 200 + assert router_env.backend.request_log[-1]["top_p"] == 0.9 + assert router_env.backend.request_log[-1]["return_sampling_mask"] is True + assert capture_fields.isdisjoint(response.json()["choices"][0]["meta_info"]) + before_failure = requests.get(f"{router_env.url}/sessions/{session_id}", timeout=5.0).json() + first_record = before_failure["records"][0] + assert set(first_record["rollout_sampling_mask"]) == {"ids", "offsets"} + assert first_record["rollout_sampling_log_probs"] + assert capture_fields.isdisjoint(first_record["response"]["choices"][0]["meta_info"]) + + second_messages = [ + {"role": "user", "content": "What is 1+2?"}, + response.json()["choices"][0]["message"], + {"role": "system", "content": "second turn"}, + ] + assert chat(messages=second_messages, model="omit-sampling-mask").status_code == 502 + after_failure = requests.get(f"{router_env.url}/sessions/{session_id}", timeout=5.0).json() + assert after_failure == before_failure + + assert chat(messages=second_messages).status_code == 200 + second_state = requests.get(f"{router_env.url}/sessions/{session_id}", timeout=5.0).json() + assert len(second_state["records"]) == 2 + assert all(record["rollout_sampling_mask"] is not None for record in second_state["records"]) + + merged_response = requests.get(f"{router_env.url}/sessions/{session_id}/merged", timeout=5.0) + assert merged_response.status_code == 200 + merged = MergedSessionSample.model_validate(merged_response.json()["sample"]) + sample = apply_merged_session_sample(SimpleNamespace(), Sample(), merged) + assert 0 in sample.loss_mask and 1 in sample.loss_mask + assert sample.rollout_log_probs == [0.0] * sample.response_length + mask = sample.rollout_sampling_mask + assert mask.offsets.tolist() == list(range(sample.response_length + 1)) + assert mask.token_ids.tolist() == sample.tokens[-sample.response_length :] + + sample = truncate_samples_by_total_tokens([sample], len(sample.tokens) - 1, None)[0] + assert len(sample.rollout_sampling_mask) == sample.response_length + assert sample.rollout_sampling_mask.token_ids.tolist() == sample.tokens[-sample.response_length :] + class TestTokenizationOffload: """TITO tokenization is CPU work. It must run off the event loop so one slow