diff --git a/docs/content/docs/algorithms/off_policy_correction.mdx b/docs/content/docs/algorithms/off_policy_correction.mdx index 49964b0786..5e6756b2d2 100644 --- a/docs/content/docs/algorithms/off_policy_correction.mdx +++ b/docs/content/docs/algorithms/off_policy_correction.mdx @@ -155,6 +155,7 @@ trainer: ``` To enable rollout router replay, set `generator.inference_engine.enable_return_routed_experts=True`, `trainer.policy.megatron_config.moe_enable_routing_replay=True`, and use the `mp` distributed_executor_backend for vLLM. Note that +R3 requires per-token routes that line up with the trained tokens, so `SkyRLGymGenerator` refuses it together with `generator.step_wise_trajectories`, `generator.use_conversation_multi_turn=False`, a custom `generator.chat_template`, and `generator.vision_language_generator` (each breaks that alignment; see the [step-wise training](../tutorials/step-wise-training) page for the step-wise reason). Routes are captured for train batches only. Note that R3 does induce additional training bias when mini-batching, since routing decisions are fixed for all mini-batches in a training batch. However, it has been shown to be important for stabilizing large-scale MoE training, particularly in models adopting a DeepSeek-V3 like architecture (notably the GLM family) due to the use of sigmoid-based affinity scoring instead of softmax for top-k routing. diff --git a/docs/content/docs/tutorials/step-wise-training.mdx b/docs/content/docs/tutorials/step-wise-training.mdx index 83a34ba6f5..d5b053cc2e 100644 --- a/docs/content/docs/tutorials/step-wise-training.mdx +++ b/docs/content/docs/tutorials/step-wise-training.mdx @@ -71,7 +71,10 @@ class GeneratorOutput(TypedDict): rollout_metrics: Optional[Dict[str, Any]] rollout_logprobs: Optional[List[List[float]]] trajectory_ids: Optional[List[TrajectoryID]] - rollout_expert_indices: Optional[List[List[List[List[int]]]]] + trajectory_generation_times: Optional[List[float]] + trajectory_time_splits: Optional[Dict[str, List[float]]] + rollout_expert_indices: Optional[List[RoutedExpertIndices]] + rollout_sample_support: Optional[List[List[List[int]]]] # Applicable only for step-wise training is_last_step: Optional[List[bool]] ``` @@ -85,6 +88,8 @@ When `step_wise_trajectories=True`, some related fields: | `is_last_step` | `List[bool]` | Marks the final step of each trajectory. Must have at least one `True`, and the last element must be `True`. | | `trajectory_ids` | `List[TrajectoryID]` | Associates each step-sample with its parent trajectory. All steps of the same trajectory share the same `TrajectoryID`. | | `rollout_logprobs` | `List[List[float]]` | Per-token logprobs from the inference engine, aligned with `response_ids`. Required for TIS. | +| `rollout_sample_support` | `List[List[List[int]]]` | Per generated token, the bounded top-k vocab IDs the sampler drew from, right-padded with `SAMPLE_SUPPORT_PADDING` (`-1`). One dense `[tokens, top_k]` block per row, aligned with `response_ids`. Every position must be present: an observation token or a synthetic (loop-appended) EOS carries an all-padding row rather than being absent, so the block stays rectangular and aligned. Enabled with `generator.inference_engine.enable_return_sample_support_set`, and requested per request (train batches only — eval's greedy `top_k=-1` params cannot satisfy the capture contract). | +| `rollout_expert_indices` | `Optional[List[RoutedExpertIndices]]` | Per token, the `[layers, topk]` MoE routes vLLM recorded, as one `[tokens, layers, topk]` integer array per trajectory (rollout router replay, R3). **Refused with `step_wise_trajectories=True`**: each step-wise row's prompt is the whole history so far while routes are recorded for that step's generated tokens only, so a step's routes would replay onto the first N prompt tokens of its row with no length mismatch to assert on, silently training against routing that does not match the rollout. | ### Concrete Example diff --git a/skyrl/backends/skyrl_train/distributed/megatron/model_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/model_utils.py index af72391a2c..f985c75f53 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/model_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/model_utils.py @@ -24,6 +24,8 @@ import torch import torch.distributed as dist +from skyrl.backends.skyrl_train.utils.packed_tensor import lengths_from_offsets + @torch.no_grad() def _compute_distributed_log_softmax( @@ -924,7 +926,7 @@ def _packed_sequence_indices( token_indices = torch.arange(total_tokens, device=device) seq_indices = torch.searchsorted(cu_seqlens_padded[1:], token_indices, right=True) seq_offsets = token_indices - cu_seqlens_padded[seq_indices] - seq_lens_padded = cu_seqlens_padded[1:] - cu_seqlens_padded[:-1] + seq_lens_padded = lengths_from_offsets(cu_seqlens_padded) return cu_seqlens_padded, token_indices, seq_indices, seq_offsets, seq_lens_padded diff --git a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py index cb3e526c95..e028441b4b 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py @@ -1,13 +1,16 @@ """Token-aligned metadata layout transforms shared by training features.""" +from collections.abc import Callable, Sequence from dataclasses import dataclass +import numpy as np import torch from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import ( get_packed_seq_align_size, get_unpacked_seq_align_size, ) +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor # Megatron is imported lazily inside functions so that non-Megatron backends can # import this module for its layout dataclass and padding transforms. @@ -99,34 +102,113 @@ def align_token_metadata( f"attention_mask shape {layout.attention_mask.shape}" ) + return _align_token_rows( + lambda row_index: metadata[row_index, layout.attention_mask[row_index]], + metadata, + metadata.shape[2:], + layout, + padding_value, + next_token=next_token, + ) + + +def align_packed_token_metadata( + metadata: PackedTensor, + layout: TokenMetadataLayout, + padding_value: torch.Tensor | bool | int, + *, + next_token: bool = False, + segment_starts: Sequence[int] | None = None, +) -> torch.Tensor: + """Align metadata that already arrives packed as ``[sum(seqlen), *row_shape]``. + + This relies on left padding, which makes each trajectory's real tokens contiguous. + Without ``segment_starts``, segments must match ``layout.sequence_lengths``; + otherwise each segment is placed at its specified real-token offset. + """ + if metadata.device != layout.attention_mask.device: + raise ValueError("Token-aligned metadata and attention_mask must be on the same device") + if len(metadata) != len(layout.sequence_lengths): + raise ValueError( + f"Packed metadata holds {len(metadata)} segments for {len(layout.sequence_lengths)} trajectories" + ) + segment_lengths = metadata.sequence_lengths.tolist() + if segment_starts is None: + if segment_lengths != list(layout.sequence_lengths): + raise ValueError( + f"Packed metadata segments {segment_lengths} do not match " + f"trajectory lengths {list(layout.sequence_lengths)}" + ) + else: + if len(segment_starts) != len(metadata): + raise ValueError(f"Got {len(segment_starts)} segment starts for {len(metadata)} segments") + for row_index, (start, length) in enumerate(zip(segment_starts, segment_lengths, strict=True)): + if start < 0 or start + length > layout.sequence_lengths[row_index]: + raise ValueError( + f"Segment {row_index} spans real tokens [{start}, {start + length}) of a " + f"{layout.sequence_lengths[row_index]}-token trajectory" + ) + + return _align_token_rows( + metadata.segment, + metadata.values, + metadata.row_shape, + layout, + padding_value, + next_token=next_token, + segment_starts=segment_starts, + ) + + +def _align_token_rows( + rows_for: Callable[[int], torch.Tensor], + source: torch.Tensor, + row_shape: tuple[int, ...] | torch.Size, + layout: TokenMetadataLayout, + padding_value: torch.Tensor | bool | int, + *, + next_token: bool = False, + segment_starts: Sequence[int] | None = None, +) -> torch.Tensor: + """Place each trajectory's real-token rows into Megatron's layout and CP-shard them. + + ``rows_for(row_index)`` yields one trajectory's rows. They land at the front of its + padded region unless ``segment_starts`` names a per-trajectory destination offset. + """ if layout.padded_sequence_lengths is None: if next_token: raise ValueError("next-token metadata alignment is only used for packed sequences") aligned = _new_metadata_tensor( - metadata, - (metadata.shape[0], layout.aligned_sequence_length, *metadata.shape[2:]), + source, + (len(layout.sequence_lengths), layout.aligned_sequence_length, *row_shape), padding_value, ) for row_index, sequence_length in enumerate(layout.sequence_lengths): - aligned[row_index, :sequence_length] = metadata[row_index, layout.attention_mask[row_index]] + rows = rows_for(row_index) + start = 0 if segment_starts is None else segment_starts[row_index] + end = sequence_length if segment_starts is None else start + rows.shape[0] + aligned[row_index, start:end] = rows return aligned packed = _new_metadata_tensor( - metadata, - (layout.aligned_sequence_length, *metadata.shape[2:]), + source, + (layout.aligned_sequence_length, *row_shape), padding_value, ) offset = 0 for row_index, (sequence_length, padded_length) in enumerate( zip(layout.sequence_lengths, layout.padded_sequence_lengths, strict=True) ): - packed[offset : offset + sequence_length] = metadata[row_index, layout.attention_mask[row_index]] + rows = rows_for(row_index) + start = offset if segment_starts is None else offset + segment_starts[row_index] + end = offset + sequence_length if segment_starts is None else start + rows.shape[0] + packed[start:end] = rows # Match Megatron's [seq0, pad0, seq1, pad1, ...] microbatch layout. offset += padded_length if next_token: # Each packed logit predicts the next token within its own padded sequence. - shifted = _new_metadata_tensor(metadata, packed.shape, padding_value) + shifted = _new_metadata_tensor(source, packed.shape, padding_value) offset = 0 for padded_length in layout.padded_sequence_lengths: shifted[offset : offset + padded_length - 1] = packed[offset + 1 : offset + padded_length] @@ -135,7 +217,7 @@ def align_token_metadata( if layout.context_parallel_size > 1: out = _new_metadata_tensor( - metadata, + source, (packed.shape[0] // layout.context_parallel_size, *packed.shape[1:]), padding_value, ) @@ -205,3 +287,63 @@ def scatter_packed_token_values_to_batch( ) batch_values[output_mask] = values[packed_mask] return batch_values + + +class TokenMetadataTrace: + """Accumulate arrays whose first dimension is aligned to tokens.""" + + def __init__(self) -> None: + self._chunks: list[np.ndarray] = [] + self._schema: tuple[tuple[int, ...], np.dtype] | None = None + self._num_rows = 0 + self._finalized = False + + @property + def num_rows(self) -> int: + return self._num_rows + + def append(self, rows: np.ndarray, *, expected_rows: int) -> None: + if self._finalized: + raise RuntimeError("token metadata trace is already finalized") + if isinstance(expected_rows, bool) or not isinstance(expected_rows, int) or expected_rows < 0: + raise ValueError(f"expected_rows must be a non-negative integer, got {expected_rows!r}") + if not isinstance(rows, np.ndarray): + raise TypeError("token metadata rows must be a NumPy array") + if rows.ndim < 1: + raise ValueError("token metadata must have a token-row dimension") + if rows.shape[0] != expected_rows: + raise ValueError(f"token metadata has {rows.shape[0]} rows, expected {expected_rows}") + if not rows.flags.c_contiguous: + raise ValueError("token metadata rows must be contiguous") + + schema = (rows.shape[1:], rows.dtype) + if self._schema is None: + self._schema = schema + elif schema != self._schema: + raise ValueError(f"token metadata schema changed from {self._schema} to {schema}") + + self._chunks.append(rows) + self._num_rows += expected_rows + + def append_padding(self, count: int, *, fill: int = -1) -> None: + """Append ``count`` rows of ``fill`` in the schema already established by ``append``.""" + if isinstance(count, bool) or not isinstance(count, int) or count < 0: + raise ValueError(f"padding count must be a non-negative integer, got {count!r}") + if count == 0: + return + if self._schema is None: + raise ValueError("cannot pad token metadata before any rows are captured") + + row_shape, dtype = self._schema + self.append(np.full((count, *row_shape), fill, dtype=dtype, order="C"), expected_rows=count) + + def finalize(self, *, expected_rows: int) -> np.ndarray: + if self._finalized: + raise RuntimeError("token metadata trace is already finalized") + if self._num_rows != expected_rows: + raise ValueError(f"token metadata trace has {self._num_rows} rows, expected {expected_rows}") + if not self._chunks: + raise ValueError("token metadata trace has no chunks") + + self._finalized = True + return self._chunks[0] if len(self._chunks) == 1 else np.concatenate(self._chunks, axis=0) diff --git a/skyrl/backends/skyrl_train/inference_servers/base.py b/skyrl/backends/skyrl_train/inference_servers/base.py index aac9ee6f6d..2c1d9a5ef5 100644 --- a/skyrl/backends/skyrl_train/inference_servers/base.py +++ b/skyrl/backends/skyrl_train/inference_servers/base.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Dict, Hashable, List, Optional, Tuple, TypedDict from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices +from skyrl.backends.skyrl_train.utils.sample_support import SampleSupport if TYPE_CHECKING: from skyrl.backends.skyrl_train.weight_sync import WeightUpdateRequest @@ -34,6 +35,9 @@ class InferenceEngineInput(TypedDict): # Optional prefix-cache salt forwarded to vLLM as the request ``cache_salt`` so cache blocks are # only shared between requests carrying the same salt. See ``GeneratorConfig.use_cache_salt``. cache_salt: Optional[str] + routed_experts_prompt_starts: Optional[List[int]] + # Per-batch opt-in; the engine must enable sample-support capture at startup. + return_sample_support: Optional[bool] class InferenceEngineOutput(TypedDict): @@ -50,6 +54,8 @@ class InferenceEngineOutput(TypedDict): response_logprobs: Optional[List[List[float]]] prompt_logprobs: Optional[List[List[float]]] # per-prompt-token logprobs under the current model rollout_expert_indices: Optional[List[RoutedExpertIndices]] + # One ``[generated_tokens, top_k]`` int32 support array per prompt. + rollout_sample_support: Optional[List[SampleSupport]] class InferenceEngineInterface(ABC): diff --git a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py index a9b8f43a2c..4facecc64d 100644 --- a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py +++ b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py @@ -3,14 +3,21 @@ ``VLLMServerActor`` writes these payloads and ``RemoteInferenceClient`` reads them; nothing else depends on the encoding. Both sides serialize with orjson, which rejects non-finite floats and has no notion of NumPy arrays, so the -helpers here exist to get sampled logprobs and routed-expert IDs across that +helpers here exist to get sampled logprobs and NumPy side channels across that boundary intact. + +Side-channel arrays use ``{data: , shape: [...], dtype: }`` +envelopes. Keeping ``data`` first lets ``load_packed_body`` decode it from the +raw response without materializing a large Python ``str``. """ import math -from typing import Any, Iterable, Mapping, Optional, Tuple +from collections import deque +from enum import StrEnum +from typing import Any, Collection, Iterable, Mapping, Optional, Tuple import numpy as np +import orjson import pybase64 from skyrl.backends.skyrl_train.utils.routed_experts import ( @@ -18,11 +25,42 @@ RoutedExpertIndices, compact_routed_expert_indices, ) +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_DTYPES, + SampleSupport, + validate_sample_support, +) # Matches the floor vLLM applies at its own serving boundaries. CLAMPED_LOGPROB = -9999.0 -_DTYPES = {dtype.name: dtype for dtype in ROUTED_EXPERT_DTYPES} + +class PackedArrayKey(StrEnum): + """Envelope keys, with ``DATA`` first for ``load_packed_body``.""" + + DATA = "data" + SHAPE = "shape" + DTYPE = "dtype" + + +class PackedField(StrEnum): + """Response-body fields whose value is a packed-array envelope.""" + + ROUTED_EXPERTS = "routed_experts" + ROLLOUT_SAMPLE_SUPPORT = "rollout_sample_support" + + +PACKED_SIDE_CHANNEL_FIELDS: tuple[str, ...] = tuple(PackedField) + +_ENVELOPE_KEYS = frozenset(PackedArrayKey) + +_ROUTED_EXPERTS_NDIM = 3 +_SAMPLE_SUPPORT_NDIM = 2 + +_QUOTE = b'"' + +# Base64 cannot contain this scan anchor. +_PACKED_DATA_ANCHOR = f':{{"{PackedArrayKey.DATA}":"'.encode() def build_logprobs_content( @@ -72,35 +110,160 @@ def _to_host_array(routed_experts: Any) -> Any: return routed_experts -def pack_routed_experts(routed_experts: RoutedExpertIndices) -> dict[str, Any]: - compact = compact_routed_expert_indices(_to_host_array(routed_experts)) - return { - "data": pybase64.b64encode(memoryview(compact)).decode("ascii"), - "shape": list(compact.shape), - "dtype": compact.dtype.name, +def pack_ndarray( + arr: np.ndarray, + *, + allowed_dtypes: Collection[np.dtype], + extra: Optional[Mapping[str, Any]] = None, +) -> dict[str, Any]: + """Encode ``arr`` as a base64 envelope carrying ``extra`` as sidecar fields.""" + if not isinstance(arr, np.ndarray): + raise TypeError("packed array must be a NumPy array") + if arr.dtype not in allowed_dtypes: + allowed = sorted(dtype.name for dtype in allowed_dtypes) + raise ValueError(f"packed array {PackedArrayKey.DTYPE} {arr.dtype.name!r} is not one of {allowed}") + if extra is not None: + collisions = sorted(set(extra) & _ENVELOPE_KEYS) + if collisions: + raise ValueError(f"sidecar fields collide with envelope keys: {collisions}") + + contiguous = np.ascontiguousarray(arr) + # `.value` keys: orjson rejects str subclasses as dict keys. + payload = { + PackedArrayKey.DATA.value: pybase64.b64encode(memoryview(contiguous)).decode("ascii"), + PackedArrayKey.SHAPE.value: list(contiguous.shape), + PackedArrayKey.DTYPE.value: contiguous.dtype.name, } + if extra is not None: + payload.update(extra) + return payload -def decode_packed_routed_experts(payload: dict[str, Any]) -> RoutedExpertIndices: - if not isinstance(payload, dict): - raise TypeError("packed routed expert indices must be an object") +def unpack_ndarray( + payload: Mapping[str, Any], + *, + allowed_dtypes: Collection[np.dtype], + ndim: int, +) -> Tuple[np.ndarray, dict[str, Any]]: + """Decode an envelope whose base64 ``data`` may be a string or buffer.""" + if not isinstance(payload, Mapping): + raise TypeError("packed array payload must be an object") try: - dtype = _DTYPES[payload["dtype"]] - shape = tuple(payload["shape"]) - data = pybase64.b64decode_as_bytearray(payload["data"], validate=True) + dtype_name = payload[PackedArrayKey.DTYPE] + shape = tuple(payload[PackedArrayKey.SHAPE]) + data = pybase64.b64decode_as_bytearray(payload[PackedArrayKey.DATA], validate=True) except (KeyError, TypeError, ValueError) as exc: - raise ValueError("invalid packed routed_experts payload") from exc - # bool is a subclass of int, so it needs an explicit rejection; np.integer is - # accepted for in-process callers, since orjson only ever yields plain ints. - if len(shape) != 3 or any( + raise ValueError(f"invalid packed array envelope: {exc}") from exc + + dtypes = {dtype.name: dtype for dtype in allowed_dtypes} + if not isinstance(dtype_name, str) or dtype_name not in dtypes: + raise ValueError(f"packed array {PackedArrayKey.DTYPE} {dtype_name!r} is not one of {sorted(dtypes)}") + dtype = dtypes[dtype_name] + # Reject bool, an int subclass; accept np.integer for in-process callers. + if len(shape) != ndim or any( not isinstance(dim, (int, np.integer)) or isinstance(dim, bool) or dim < 0 for dim in shape ): - raise ValueError(f"invalid packed routed_experts shape: {shape}") + raise ValueError(f"packed array {PackedArrayKey.SHAPE} {shape} is not {ndim} non-negative dimensions") expected_size = math.prod(shape) * dtype.itemsize if len(data) != expected_size: - raise ValueError(f"packed routed_experts has {len(data)} bytes, expected {expected_size}") - decoded = np.frombuffer(data, dtype=dtype).reshape(shape) + raise ValueError( + f"packed array {PackedArrayKey.DATA} has {len(data)} bytes, " + f"expected {expected_size} for {dtype_name}{list(shape)}" + ) + + array = np.frombuffer(data, dtype=dtype).reshape(shape) + sidecar = {key: value for key, value in payload.items() if key not in _ENVELOPE_KEYS} + return array, sidecar + + +def pack_routed_experts(routed_experts: RoutedExpertIndices) -> dict[str, Any]: + compact = compact_routed_expert_indices(_to_host_array(routed_experts)) + return pack_ndarray(compact, allowed_dtypes=ROUTED_EXPERT_DTYPES) + + +def decode_packed_routed_experts(payload: dict[str, Any]) -> RoutedExpertIndices: + decoded, _ = unpack_ndarray(payload, allowed_dtypes=ROUTED_EXPERT_DTYPES, ndim=_ROUTED_EXPERTS_NDIM) compact = compact_routed_expert_indices(decoded) - if compact.dtype != dtype: - raise ValueError(f"packed routed_experts uses non-canonical dtype {dtype.name}; expected {compact.dtype.name}") + if compact.dtype != decoded.dtype: + raise ValueError( + f"packed routed_experts uses non-canonical dtype {decoded.dtype.name}; expected {compact.dtype.name}" + ) return compact + + +def pack_sample_support(sample_support: SampleSupport) -> dict[str, Any]: + return pack_ndarray(validate_sample_support(sample_support), allowed_dtypes=SAMPLE_SUPPORT_DTYPES) + + +def decode_packed_sample_support(payload: dict[str, Any]) -> SampleSupport: + decoded, _ = unpack_ndarray(payload, allowed_dtypes=SAMPLE_SUPPORT_DTYPES, ndim=_SAMPLE_SUPPORT_NDIM) + return validate_sample_support(decoded) + + +def _data_prefix(field: str) -> bytes: + """The bytes an orjson-serialized packed ``field`` opens with.""" + return f'"{field}"'.encode() + _PACKED_DATA_ANCHOR + + +def load_packed_body(raw: bytes, *, fields: tuple[str, ...] = PACKED_SIDE_CHANNEL_FIELDS) -> dict[str, Any]: + """Parse a response after replacing registered base64 blobs with views. + + Null fields pass through. An envelope layout the scan cannot splice raises + instead of falling back to materializing the base64 as a Python string. + """ + prefixes = {field: _data_prefix(field) for field in fields} + blobs: dict[str, deque[memoryview]] = {field: deque() for field in fields} + view = memoryview(raw) + pieces: list[memoryview] = [] + copied = 0 + scan = 0 + while (anchor := raw.find(_PACKED_DATA_ANCHOR, scan)) >= 0: + field = _match_packed_field(raw, anchor, prefixes) + if field is None: + scan = anchor + len(_PACKED_DATA_ANCHOR) + continue + start = anchor + len(_PACKED_DATA_ANCHOR) + end = raw.find(_QUOTE, start) + if end < 0: + raise ValueError(f"unterminated base64 {PackedArrayKey.DATA} for {field} in the response body") + pieces.append(view[copied:start]) + blobs[field].append(view[start:end]) + copied = scan = end + + if pieces: + pieces.append(view[copied:]) + body = orjson.loads(b"".join(pieces)) + else: + body = orjson.loads(raw) + _restore_packed_data(body, blobs) + + unplaced = {field: len(queue) for field, queue in blobs.items() if queue} + if unplaced: + raise ValueError(f"spliced packed blobs found no envelope in the response body: {unplaced}") + return body + + +def _match_packed_field(raw: bytes, anchor: int, prefixes: Mapping[str, bytes]) -> Optional[str]: + """Name the registered field whose prefix ends at ``anchor``, if any.""" + for field, prefix in prefixes.items(): + begin = anchor + len(_PACKED_DATA_ANCHOR) - len(prefix) + if begin >= 0 and raw.startswith(prefix, begin): + return field + return None + + +def _restore_packed_data(node: Any, blobs: Mapping[str, deque[memoryview]]) -> None: + """Put each blob back on its envelope's ``data`` key, in document order.""" + if isinstance(node, dict): + for key, value in node.items(): + queue = blobs.get(key) + if queue is not None and isinstance(value, dict) and PackedArrayKey.DATA in value: + if not queue: + raise ValueError(f"packed {key} survived the scan unspliced; the response-body layout drifted") + value[PackedArrayKey.DATA.value] = queue.popleft() + elif isinstance(value, (dict, list)): + _restore_packed_data(value, blobs) + elif isinstance(node, list): + for item in node: + if isinstance(item, (dict, list)): + _restore_packed_data(item, blobs) diff --git a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py index 7bfbd42a64..d0d1ac937d 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -74,8 +74,13 @@ MultiModalFeatures, ) from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( + PackedField, decode_packed_routed_experts, + decode_packed_sample_support, + load_packed_body, ) +from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices +from skyrl.backends.skyrl_train.utils.sample_support import SampleSupport from skyrl.backends.utils import convert_vllm_prompt_logprobs from skyrl.env_vars import ( SKYRL_GENERATE_CONCURRENCY_PER_ENGINE, @@ -167,6 +172,181 @@ class SampleResponse(TypedDict): topk_prompt_logprobs: Optional[List[Optional[List[Tuple[int, float]]]]] +@dataclass(frozen=True) +class RemoteGenerateResult: + """Raw token generation result returned by ``RemoteGenerateClient``.""" + + raw_response: Dict[str, Any] + response_ids: List[int] + response_logprobs: Optional[List[float]] + stop_reason: str + routed_experts: Optional[RoutedExpertIndices] + sample_support: Optional[SampleSupport] + + +@dataclass +class RemoteGenerateClient: + """Reusable HTTP client for one raw-token generation request.""" + + proxy_url: str + _session: Optional[aiohttp.ClientSession] = field(default=None, init=False, repr=False) + + async def _get_session(self) -> aiohttp.ClientSession: + current_loop = asyncio.get_running_loop() + if self._session is not None and not self._session.closed and self._session.loop != current_loop: + self._session = None + if self._session is None or self._session.closed: + connector = aiohttp.TCPConnector( + limit=SKYRL_HTTP_CONNECTION_LIMIT, + keepalive_timeout=2, + ) + self._session = aiohttp.ClientSession( + connector=connector, + timeout=aiohttp.ClientTimeout(total=None), + ) + return self._session + + async def _post( + self, + url: str, + json: Dict[str, Any], + headers: Optional[Dict[str, str]] = None, + *, + packed_side_channels: bool = False, + ) -> Any: + """POST JSON with retries, optionally splicing packed arrays before parsing.""" + session = await self._get_session() + last_exc: Optional[Exception] = None + for attempt in range(_DATA_PLANE_RETRIES): + try: + async with session.post(url, json=json, headers=headers) as resp: + try: + raw = await resp.read() + body = load_packed_body(raw) if packed_side_channels else orjson.loads(raw) + except orjson.JSONDecodeError as exc: + if 400 <= resp.status < 500: + text = await resp.text() + raise aiohttp.ClientResponseError( + resp.request_info, + resp.history, + status=resp.status, + message=text or resp.reason, + headers=resp.headers, + ) from exc + last_exc = exc + logger.debug(f"retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {exc}") + await asyncio.sleep(1) + continue + raise_for_status(resp, body) + return body + except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as exc: + last_exc = exc + logger.debug(f"POST retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {exc}") + await asyncio.sleep(1) + if last_exc is None: + raise RuntimeError(f"POST failed without an exception for {url=}") + raise last_exc + + async def generate( + self, + *, + prompt_token_ids: List[int], + sampling_params: Dict[str, Any], + session_id: Optional[Any], + model: str, + return_routed_experts: bool = False, + routed_experts_prompt_start: Optional[int] = None, + return_sample_support: bool = False, + mm_features: Optional[MultiModalFeatures] = None, + cache_salt: Optional[str] = None, + ) -> RemoteGenerateResult: + """Generate one raw-token completion with optional per-token replay metadata.""" + if routed_experts_prompt_start is not None: + if not return_routed_experts: + raise ValueError("routed_experts_prompt_start requires return_routed_experts=True") + if ( + isinstance(routed_experts_prompt_start, bool) + or not isinstance(routed_experts_prompt_start, int) + or not 0 <= routed_experts_prompt_start <= len(prompt_token_ids) + ): + raise ValueError("routed_experts_prompt_start must be an integer within the prompt") + + packed_side_channels = return_routed_experts or return_sample_support + path = "/skyrl/v1/generate" if packed_side_channels else "/inference/v1/generate" + request_sampling_params = dict(sampling_params) + if routed_experts_prompt_start is not None: + request_sampling_params["routed_experts_prompt_start"] = routed_experts_prompt_start + payload: Dict[str, Any] = { + "sampling_params": request_sampling_params, + "model": model, + "token_ids": prompt_token_ids, + } + if return_sample_support: + payload["return_sample_support"] = True + if mm_features: + payload["features"] = mm_features + # `cache_salt` is a top-level request field (forwarded to vLLM's TokensPrompt), not a sampling + # param. + if cache_salt is not None: + payload["cache_salt"] = cache_salt + + headers = {"Content-Type": "application/json"} + if session_id: + headers["X-Session-ID"] = str(session_id) + + response = await self._post( + f"{self.proxy_url}{path}", + json=payload, + headers=headers, + packed_side_channels=packed_side_channels, + ) + choice = response["choices"][0] + token_ids = choice["token_ids"] + logprobs = choice.get("logprobs") + response_logprobs = None + if logprobs is not None: + logprobs_content = logprobs.get("content", []) + if logprobs_content: + response_logprobs = [logprob_info["logprob"] for logprob_info in logprobs_content] + + routed_experts = None + if return_routed_experts: + packed_routed_experts = choice.get(PackedField.ROUTED_EXPERTS) + if not isinstance(packed_routed_experts, dict): + raise ValueError("/skyrl/v1/generate must return packed routed_experts") + routed_experts = decode_packed_routed_experts(packed_routed_experts) + + sample_support = None + if return_sample_support: + packed_sample_support = choice.get(PackedField.ROLLOUT_SAMPLE_SUPPORT) + if not isinstance(packed_sample_support, dict): + raise ValueError("/skyrl/v1/generate must return packed rollout_sample_support") + sample_support = decode_packed_sample_support(packed_sample_support) + + return RemoteGenerateResult( + raw_response=response, + response_ids=token_ids, + response_logprobs=response_logprobs, + stop_reason=choice["finish_reason"], + routed_experts=routed_experts, + sample_support=sample_support, + ) + + async def aclose(self) -> None: + if self._session is not None and not self._session.closed: + await self._session.close() + self._session = None + + def __getstate__(self) -> Dict[str, Any]: + state = self.__dict__.copy() + state["_session"] = None + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + self.__dict__.update(state) + self._session = None + + @dataclass class RemoteInferenceClient(InferenceEngineInterface): """ @@ -215,6 +395,10 @@ class RemoteInferenceClient(InferenceEngineInterface): enable_return_routed_experts: bool = False """Whether to return routed expert indices (R3 / rollout router replay).""" + enable_return_sample_support_set: bool = False + """Whether the engine may return the sampler's bounded top-k support per generated token. + Capture is per-request: callers opt a batch in with ``InferenceEngineInput.return_sample_support``.""" + uses_lora_weight_sync: bool = False """True when the trainer syncs LoRA adapters (rather than full/merged weights). When True, `sleep()` is forced to level=1: level=2 discards the base model from VRAM with no CPU backup, @@ -225,7 +409,7 @@ class RemoteInferenceClient(InferenceEngineInterface): """Optional HF tokenizer for local tokenize/detokenize (avoids HTTP round-trips).""" # Private fields excluded from repr for cleaner output - _session: Optional[aiohttp.ClientSession] = field(default=None, repr=False) + _generate_client: Optional[RemoteGenerateClient] = field(default=None, repr=False) _world_size: Optional[Tuple[int, int]] = field(default=None, repr=False) _gen_sem: Optional[asyncio.Semaphore] = field(default=None, repr=False) _detok_sem: Optional[asyncio.Semaphore] = field(default=None, repr=False) @@ -282,66 +466,16 @@ def _get_semaphores(self) -> Tuple[Optional[asyncio.Semaphore], Optional[asyncio self._sem_loop = current_loop return self._gen_sem, self._detok_sem + def _get_generate_client(self) -> RemoteGenerateClient: + if self._generate_client is None: + self._generate_client = RemoteGenerateClient(proxy_url=self.proxy_url) + return self._generate_client + async def _get_session(self) -> aiohttp.ClientSession: - """Get or create the aiohttp session.""" - # Re-use the existing session object if it is not closed. - # Note that we also create a new session object if the event loop has changed, since - # aiohttp.ClientSession is tied to the event loop. - current_loop = asyncio.get_running_loop() - if self._session is not None and not self._session.closed and self._session.loop != current_loop: - # Event loop changed - the old session is unusable (bound to a dead loop). - self._session = None - if self._session is None or self._session.closed: - # keepalive_timeout must be shorter than the server's timeout_keep_alive - # (uvicorn default: 5s). Otherwise aiohttp reuses connections the server - # has already closed, causing ECONNRESET under high concurrency. - connector = aiohttp.TCPConnector( - limit=SKYRL_HTTP_CONNECTION_LIMIT, - keepalive_timeout=2, - ) - self._session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=None)) - return self._session + return await self._get_generate_client()._get_session() async def _post(self, url: str, json: Dict[str, Any], headers: Optional[Dict[str, str]] = None) -> Any: - """POST with retry + backoff on transient connection errors. - - Between generate bursts the pool's keep-alive connections go stale - (server closes them after ``timeout_keep_alive``). An immediate - retry would grab another stale connection from the same pool, so we - sleep briefly to let the connector detect and purge dead sockets - before the next attempt. - """ - session = await self._get_session() - last_exc: Optional[Exception] = None - for attempt in range(_DATA_PLANE_RETRIES): - try: - async with session.post(url, json=json, headers=headers) as resp: - try: - body = orjson.loads(await resp.read()) - except orjson.JSONDecodeError as e: - if 400 <= resp.status < 500: - # Non-JSON client error (e.g. plain text 422 from vllm-router). - # Raise immediately — client errors won't succeed on retry. - text = await resp.text() - raise aiohttp.ClientResponseError( - resp.request_info, - resp.history, - status=resp.status, - message=text or resp.reason, - headers=resp.headers, - ) - last_exc = e - logger.debug(f"retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {e}") - await asyncio.sleep(1) - continue - raise_for_status(resp, body) - return body - except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as e: - last_exc = e - logger.debug(f"POST retry {attempt + 1}/{_DATA_PLANE_RETRIES} for {url=}: {e}") - await asyncio.sleep(1) - continue - raise last_exc # type: ignore[misc] + return await self._get_generate_client()._post(url, json=json, headers=headers) # --------------------------- # Data Plane @@ -406,6 +540,15 @@ async def generate( session_ids = input_batch.get("session_ids") mm_features = input_batch.get("mm_features") cache_salt = input_batch.get("cache_salt") + routed_experts_prompt_starts = input_batch.get("routed_experts_prompt_starts") + if routed_experts_prompt_starts is not None: + if not self.enable_return_routed_experts: + raise ValueError("routed_experts_prompt_starts requires enable_return_routed_experts=True") + if len(routed_experts_prompt_starts) != len(prompt_token_ids): + raise ValueError("routed_experts_prompt_starts must have one entry per prompt") + return_sample_support = self.enable_return_sample_support_set and input_batch.get( + "return_sample_support", False + ) get_logprobs = sampling_params.get("logprobs") is not None # Two semaphores decouple the generate and detokenize stages: @@ -429,6 +572,10 @@ async def _throttled_generate(idx: int) -> Dict[str, Any]: sampling_params=sampling_params, session_id=session_ids[idx] if session_ids and idx < len(session_ids) else None, mm_features=mm_features[idx] if mm_features and idx < len(mm_features) else None, + routed_experts_prompt_start=( + routed_experts_prompt_starts[idx] if routed_experts_prompt_starts is not None else None + ), + return_sample_support=return_sample_support, model=model, cache_salt=cache_salt, ) @@ -438,6 +585,10 @@ async def _throttled_generate(idx: int) -> Dict[str, Any]: sampling_params=sampling_params, session_id=session_ids[idx] if session_ids and idx < len(session_ids) else None, mm_features=mm_features[idx] if mm_features and idx < len(mm_features) else None, + routed_experts_prompt_start=( + routed_experts_prompt_starts[idx] if routed_experts_prompt_starts is not None else None + ), + return_sample_support=return_sample_support, model=model, cache_salt=cache_salt, ) @@ -454,6 +605,9 @@ async def _throttled_detokenize(token_ids: List[int]) -> str: rollout_expert_indices = ( [result["routed_experts"] for result in raw_results] if self.enable_return_routed_experts else None ) + rollout_sample_support = ( + [result[PackedField.ROLLOUT_SAMPLE_SUPPORT] for result in raw_results] if return_sample_support else None + ) return InferenceEngineOutput( responses=responses, @@ -461,6 +615,7 @@ async def _throttled_detokenize(token_ids: List[int]) -> str: response_ids=[r["response_ids"] for r in raw_results], response_logprobs=[r["response_logprobs"] for r in raw_results] if get_logprobs else None, rollout_expert_indices=rollout_expert_indices, + rollout_sample_support=rollout_sample_support, ) async def _generate_single( @@ -471,64 +626,26 @@ async def _generate_single( model: str, mm_features: Optional[MultiModalFeatures] = None, cache_salt: Optional[str] = None, + routed_experts_prompt_start: Optional[int] = None, + return_sample_support: bool = False, ) -> Dict[str, Any]: - """ - Generate completion for a single prompt. - - With keep-mode pause, in-flight requests are frozen by the vLLM - scheduler and resume where they left off after /resume. No retry - logic is needed. - - Returns: - Dict with keys: stop_reason, response_ids, response_logprobs - """ - url = ( - f"{self.proxy_url}/skyrl/v1/generate" - if self.enable_return_routed_experts - else f"{self.proxy_url}/inference/v1/generate" + result = await self._get_generate_client().generate( + prompt_token_ids=prompt_token_ids, + sampling_params=sampling_params, + session_id=session_id, + model=model, + return_routed_experts=self.enable_return_routed_experts, + routed_experts_prompt_start=routed_experts_prompt_start, + return_sample_support=return_sample_support, + mm_features=mm_features, + cache_salt=cache_salt, ) - - payload: dict[str, Any] = { - "sampling_params": sampling_params, - "model": model, - "token_ids": prompt_token_ids, - } - if mm_features: - payload["features"] = mm_features - # `cache_salt` is a top-level request field (forwarded to vLLM's TokensPrompt), not a sampling - # param. - if cache_salt is not None: - payload["cache_salt"] = cache_salt - - headers = {"Content-Type": "application/json"} - if session_id: - headers["X-Session-ID"] = str(session_id) - - response = await self._post(url, json=payload, headers=headers) - - choice = response["choices"][0] - token_ids = choice["token_ids"] - stop_reason = choice["finish_reason"] - - response_logprobs: Optional[List[float]] = None - logprobs = choice.get("logprobs") - if logprobs is not None: - logprobs_content = logprobs.get("content", []) - if logprobs_content: - response_logprobs = [logprob_info["logprob"] for logprob_info in logprobs_content] - - routed_experts = None - if self.enable_return_routed_experts: - packed_routed_experts = choice.get("routed_experts") - if not isinstance(packed_routed_experts, dict): - raise ValueError("/skyrl/v1/generate must return packed routed_experts") - routed_experts = decode_packed_routed_experts(packed_routed_experts) - return { - "stop_reason": stop_reason, - "response_ids": token_ids, - "response_logprobs": response_logprobs, - "routed_experts": routed_experts, + "stop_reason": result.stop_reason, + "response_ids": result.response_ids, + "response_logprobs": result.response_logprobs, + "routed_experts": result.routed_experts, + PackedField.ROLLOUT_SAMPLE_SUPPORT.value: result.sample_support, } async def _render_for_sample( @@ -1405,9 +1522,8 @@ async def get_world_size(self) -> Tuple[int, int]: async def teardown(self) -> None: """Close HTTP session.""" - if self._session and not self._session.closed: - await self._session.close() - self._session = None + if self._generate_client is not None: + await self._generate_client.aclose() async def __aenter__(self) -> "RemoteInferenceClient": """Async context manager entry.""" @@ -1424,7 +1540,6 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: def __getstate__(self) -> dict: """Exclude non-serializable fields from pickle.""" state = self.__dict__.copy() - state["_session"] = None state["_gen_sem"] = None state["_detok_sem"] = None state["_sem_loop"] = None @@ -1433,19 +1548,12 @@ def __getstate__(self) -> dict: def __setstate__(self, state: dict) -> None: """Restore state after unpickling.""" self.__dict__.update(state) - self._session = None self._gen_sem = None self._detok_sem = None self._sem_loop = None - async def aclose(self): - if self._session is not None: - try: - await self._session.close() - except Exception as e: - logger.warning(f"Encountered exception {e} while closing client session") - pass - self._session = None + async def aclose(self) -> None: + await self.teardown() def raise_for_status(resp: aiohttp.ClientResponse, body: Optional[Any] = None) -> None: diff --git a/skyrl/backends/skyrl_train/inference_servers/setup.py b/skyrl/backends/skyrl_train/inference_servers/setup.py index 437d4804f3..84d481d596 100644 --- a/skyrl/backends/skyrl_train/inference_servers/setup.py +++ b/skyrl/backends/skyrl_train/inference_servers/setup.py @@ -293,6 +293,7 @@ def build_new_inference_client( server_urls=server_setup.server_urls, model_name=ie_cfg.served_model_name or cfg.trainer.policy.model.path, enable_return_routed_experts=ie_cfg.enable_return_routed_experts, + enable_return_sample_support_set=ie_cfg.enable_return_sample_support_set, uses_lora_weight_sync=_uses_lora_weight_sync(cfg), data_parallel_size=ie_cfg.data_parallel_size, tokenizer=tokenizer, diff --git a/skyrl/backends/skyrl_train/inference_servers/utils.py b/skyrl/backends/skyrl_train/inference_servers/utils.py index 4772096066..ca56bf0d00 100644 --- a/skyrl/backends/skyrl_train/inference_servers/utils.py +++ b/skyrl/backends/skyrl_train/inference_servers/utils.py @@ -120,6 +120,11 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace: # Overridable via generator.inference_engine.engine_init_kwargs.trust_remote_code below. trust_remote_code=True, ) + # Sample-support capture asks for one logprob per top-k candidate, post-filter, so the + # -inf entries that mark filtered candidates survive to the capture path. + if ie_cfg.enable_return_sample_support_set: + overrides["max_logprobs"] = cfg.generator.sampling_params.top_k + overrides["logprobs_mode"] = "processed_logprobs" for key, value in overrides.items(): setattr(args, key, value) diff --git a/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py b/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py index dd350ea586..84a1ad027a 100644 --- a/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py +++ b/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py @@ -10,6 +10,7 @@ from typing import List, Optional, Tuple import httpx +import numpy as np import orjson import uvicorn import vllm.envs as envs @@ -23,6 +24,7 @@ init_app_state, ) from vllm.inputs import TokensPrompt +from vllm.logprobs import FlatLogprobs from vllm.lora.request import LoRARequest from vllm.sampling_params import SamplingParams as VLLMSamplingParams from vllm.usage.usage_lib import UsageContext @@ -37,10 +39,17 @@ ) from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( CLAMPED_LOGPROB, + PackedField, build_logprobs_content, pack_routed_experts, + pack_sample_support, ) from skyrl.backends.skyrl_train.inference_servers.protocols import ServerActorProtocol +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_DTYPE, + SAMPLE_SUPPORT_PADDING, + SampleSupport, +) from skyrl.env_vars import ( SKYRL_HTTP_CONNECTION_LIMIT, SKYRL_VLLM_DP_PORT_OFFSET, @@ -50,6 +59,46 @@ logger = logging.getLogger(__name__) +def _sample_support_from_flat_logprobs( + logprobs: FlatLogprobs, + top_k: int, +) -> tuple[list[dict[str, float]], SampleSupport]: + """Extract sampled scores and post-filter support from vLLM's flat rows. + + Each row is ``[sampled token, top-1, ..., top-k]``; filtered candidates are ``-inf``. + """ + row_width = top_k + 1 + token_ids = np.asarray(logprobs.token_ids, dtype=SAMPLE_SUPPORT_DTYPE).reshape(-1, row_width) + processed_logprobs = np.asarray(logprobs.logprobs).reshape(-1, row_width) + support_ids = np.where( + np.isneginf(processed_logprobs[:, 1:]), + SAMPLE_SUPPORT_DTYPE.type(SAMPLE_SUPPORT_PADDING), + token_ids[:, 1:], + ) + sampled_logprobs = [{"logprob": value} for value in processed_logprobs[:, 0].tolist()] + + # vLLM's approximate top-k/top-p pivot can omit the sampled token. Replace the + # weakest valid candidate while preserving the support width and trailing padding. + sampled = token_ids[:, 0] + valid = support_ids >= 0 + present = np.any(support_ids == sampled[:, None], axis=1) + missing = (~present) & valid.any(axis=1) + if np.any(missing): + rows = np.flatnonzero(missing) + weakest_col = valid.sum(axis=1) - 1 + support_ids[rows, weakest_col[rows]] = sampled[rows] + logger.warning( + "sample-support repair: %d token(s) had the sampled id absent from top-%d support; " + "overwrote the weakest member to preserve the invariant (vLLM approx top-k/top-p " + "pivot artifact); example: sampled token %d at row %d", + rows.size, + top_k, + int(sampled[rows[0]]), + int(rows[0]), + ) + return sampled_logprobs, support_ids + + class VLLMServerActor(ServerActorProtocol): """ Ray actor that runs a vLLM OpenAI-compatible API server. @@ -430,6 +479,20 @@ async def _skyrl_generate(request: Request): sampling_params_dict = body.get("sampling_params", {}) cache_salt = body.get("cache_salt") + capture_sample_support = body.get("return_sample_support", False) + if capture_sample_support: + # Sample support requires a bounded, non-degenerate top-k set. + top_k = sampling_params_dict.get("top_k") + if not isinstance(top_k, int) or top_k <= 1: + raise HTTPException( + status_code=400, + detail=( + "return_sample_support requires sampling_params.top_k > 1, got " + f"{top_k!r}. Sample-support capture is opt-in per request." + ), + ) + sampling_params_dict["flat_logprobs"] = True + sampling_params_dict["logprobs"] = top_k sampling_params = VLLMSamplingParams(**sampling_params_dict) # `cache_salt` salts vLLM's prefix cache; vLLM rejects an empty salt, so attach only when set. if cache_salt is not None: @@ -450,7 +513,15 @@ async def _skyrl_generate(request: Request): finish_reason = resp.finish_reason logprobs = None - if resp.logprobs is not None: + sample_support = None + if capture_sample_support: + content, support_ids = _sample_support_from_flat_logprobs( + resp.logprobs, + sampling_params_dict["top_k"], + ) + logprobs = {"content": content} + sample_support = pack_sample_support(support_ids) + elif resp.logprobs is not None: content, num_clamped = build_logprobs_content(token_ids_out, resp.logprobs) if num_clamped: logger.warning( @@ -469,7 +540,8 @@ async def _skyrl_generate(request: Request): "token_ids": token_ids_out, "finish_reason": finish_reason, "logprobs": logprobs, - "routed_experts": routed_experts, + PackedField.ROUTED_EXPERTS.value: routed_experts, + PackedField.ROLLOUT_SAMPLE_SUPPORT.value: sample_support, } ] } diff --git a/skyrl/backends/skyrl_train/training_batch.py b/skyrl/backends/skyrl_train/training_batch.py index 1542adeff3..d3b920dc90 100644 --- a/skyrl/backends/skyrl_train/training_batch.py +++ b/skyrl/backends/skyrl_train/training_batch.py @@ -3,44 +3,80 @@ import copy import io import pickle -from typing import Any, Dict, Generic, List, Optional, TypedDict, TypeVar +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from enum import StrEnum +from typing import Any, Dict, Generic, List, Optional, TypedDict, TypeVar, Union import numpy as np import torch from jaxtyping import Bool, Float, Integer -from skyrl.backends.skyrl_train.utils.replay_utils import make_replay_padding_indices +from skyrl.backends.skyrl_train.utils.packed_tensor import ( + PackedTensor, + packed_padding_segments, +) +from skyrl.backends.skyrl_train.utils.replay_utils import replay_padding_row +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_FIELD, + SAMPLE_SUPPORT_PADDING, +) DictType = TypeVar("DictType") -def _serialize_tensor(value: torch.Tensor) -> dict: - """Serialize a single tensor for pickle protocol.""" +class TensorFormat(StrEnum): + """How one serialized batch field is encoded in the pickle stream.""" + + NUMPY = "numpy" + NUMPY_VIEW = "numpy_view" + TORCH = "torch" + TENSOR_LIST = "tensor_list" + PACKED_TENSOR = "packed_tensor" + + +def _serialize_tensor(value: torch.Tensor, *, zero_copy: bool = False) -> dict: + """Serialize a single tensor for pickle protocol. + + With ``zero_copy``, preserve the numpy array so pickle protocol 5 can send its + buffer out of band. The deserialized view may be read-only. + """ try: # Fast path: direct memory copy via numpy (works for most dtypes) arr = value.numpy() - return { - "format": "numpy", - "data": arr.tobytes(), - "shape": arr.shape, - "dtype": str(arr.dtype), - } except TypeError: # Fallback for dtypes not supported by numpy (e.g., bfloat16) buffer = io.BytesIO() torch.save(value, buffer) return { - "format": "torch", + "format": TensorFormat.TORCH, "data": buffer.getvalue(), } + if zero_copy: + # Shape and dtype travel with the array. + return { + "format": TensorFormat.NUMPY_VIEW, + "data": arr, + } + return { + "format": TensorFormat.NUMPY, + "data": arr.tobytes(), + "shape": arr.shape, + "dtype": str(arr.dtype), + } + def _deserialize_tensor(value: dict) -> torch.Tensor: """Deserialize a single tensor from pickle format.""" - if value.get("format") == "torch": + tensor_format = value.get("format") + if tensor_format == TensorFormat.TORCH: # Fallback path: torch.load for unsupported dtypes buffer = io.BytesIO(value["data"]) return torch.load(buffer, weights_only=True) + elif tensor_format == TensorFormat.NUMPY_VIEW: + # Under Ray, this array views a read-only plasma buffer. + return torch.from_numpy(value["data"]) else: # Fast path: reconstruct from numpy bytes # Also handles legacy format without "format" key @@ -110,6 +146,13 @@ def cat(lists: list["TensorList"]) -> "TensorList": return TensorList([t for tl in lists for t in tl.tensors]) +# Value types a batch field may hold: a dense tensor, a ragged list of tensors, or a ragged +# token-aligned field packed to one buffer plus offsets. All three index by batch position. +BATCH_FIELD_TYPES = (torch.Tensor, TensorList, PackedTensor) +BatchField = Union[torch.Tensor, TensorList, PackedTensor] +_BATCH_FIELD_ERROR = f"must be a tensor, {TensorList.__name__}, or {PackedTensor.__name__}" + + def _rebuild_tensor_batch(cls, state: Dict[str, Any]): """Module-level helper for unpickling TensorBatch (must be importable by name).""" obj = dict.__new__(cls) @@ -127,6 +170,9 @@ class TensorBatch(dict, Generic[DictType]): metadata: Optional[Dict[str, Any]] = None + # These fields may be backed by read-only shared memory after deserialization. + ZERO_COPY_KEYS: frozenset[str] = frozenset({"rollout_expert_indices", SAMPLE_SUPPORT_FIELD}) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._batch_size = None @@ -169,8 +215,8 @@ def _check_consistency(self): value = self[key] if value is None: continue - if not isinstance(value, (torch.Tensor, TensorList)): - raise ValueError(f"Field {key} must be a tensor or TensorList, got {type(value)}") + if not isinstance(value, BATCH_FIELD_TYPES): + raise ValueError(f"Field {key} {_BATCH_FIELD_ERROR}, got {type(value)}") self._device = value.device if self._device is None else self._device if len(value) != batch_size: raise ValueError(f"Batch size mismatch in {key}") @@ -185,13 +231,13 @@ def __getitem__(self, index) -> "TensorBatch[DictType]": else: return super().__getitem__(index) - def __setitem__(self, key: str, value: Optional[torch.Tensor | TensorList]) -> None: + def __setitem__(self, key: str, value: Optional[BatchField]) -> None: if value is None: super().__setitem__(key, value) return - if not isinstance(value, (torch.Tensor, TensorList)): - raise ValueError(f"Field {key} must be a tensor or TensorList, got {type(value)}") + if not isinstance(value, BATCH_FIELD_TYPES): + raise ValueError(f"Field {key} {_BATCH_FIELD_ERROR}, got {type(value)}") if hasattr(self, "_batch_size") and self._batch_size is not None and len(value) != self._batch_size: raise ValueError(f"Batch size mismatch in {key}. Expected size {self._batch_size}, got {len(value)}.") @@ -214,9 +260,7 @@ def to( for key, value in self.items(): if value is None: continue - assert isinstance( - value, (torch.Tensor, TensorList) - ), f"Field {key} must be a tensor or TensorList, got {type(value)}" + assert isinstance(value, BATCH_FIELD_TYPES), f"Field {key} {_BATCH_FIELD_ERROR}, got {type(value)}" self[key] = value.to(device=device, dtype=dtype, non_blocking=non_blocking) return self @@ -225,9 +269,7 @@ def contiguous(self) -> "TensorBatch": for key, value in self.items(): if value is None: continue - assert isinstance( - value, (torch.Tensor, TensorList) - ), f"Field {key} must be a tensor or TensorList, got {type(value)}" + assert isinstance(value, BATCH_FIELD_TYPES), f"Field {key} {_BATCH_FIELD_ERROR}, got {type(value)}" self[key] = value.contiguous() return self @@ -256,22 +298,31 @@ def __getstate__(self): """Serialize the `TensorBatch` object for pickle protocol. Uses fast numpy-based serialization when possible, with fallback to torch.save - for dtypes not supported by numpy (e.g., bfloat16). + for dtypes not supported by numpy (e.g., bfloat16). Fields in `ZERO_COPY_KEYS` + skip the intermediate `bytes` copy entirely. """ self.contiguous() if self._device is not None: assert self._device == torch.device("cpu"), "Tensors must be on CPU before serialization" batch_dict = {} for key, value in self.items(): + zero_copy = key in self.ZERO_COPY_KEYS if value is None: batch_dict[key] = None elif isinstance(value, TensorList): batch_dict[key] = { - "format": "tensor_list", + "format": TensorFormat.TENSOR_LIST, "items": [_serialize_tensor(t) for t in value.tensors], } + elif isinstance(value, PackedTensor): + batch_dict[key] = { + "format": TensorFormat.PACKED_TENSOR, + "values": _serialize_tensor(value.values, zero_copy=zero_copy), + # Offsets are too small to benefit from an out-of-band buffer. + "cu_seqlens": _serialize_tensor(value.cu_seqlens), + } else: - batch_dict[key] = _serialize_tensor(value) + batch_dict[key] = _serialize_tensor(value, zero_copy=zero_copy) return { "batch_dict": batch_dict, @@ -288,8 +339,13 @@ def __setstate__(self, state): for key, value in state["batch_dict"].items(): if value is None: self[key] = None - elif value.get("format") == "tensor_list": + elif value.get("format") == TensorFormat.TENSOR_LIST: self[key] = TensorList([_deserialize_tensor(item) for item in value["items"]]) + elif value.get("format") == TensorFormat.PACKED_TENSOR: + self[key] = PackedTensor( + _deserialize_tensor(value["values"]), + _deserialize_tensor(value["cu_seqlens"]), + ) else: self[key] = _deserialize_tensor(value) @@ -314,10 +370,8 @@ def repeat(self, repeats: int) -> "TensorBatch[DictType]": for key, value in self.items(): if value is None: new_batch[key] = value - elif isinstance(value, TensorList): - new_batch[key] = value.repeat(repeats) else: - assert isinstance(value, torch.Tensor), f"Field {key} must be a tensor, got {type(value)}" + assert isinstance(value, BATCH_FIELD_TYPES), f"Field {key} {_BATCH_FIELD_ERROR}, got {type(value)}" new_batch[key] = value.repeat(repeats) new_batch = self.__class__(new_batch) new_batch.metadata = self.metadata @@ -338,10 +392,8 @@ def repeat_interleave(self, repeats: int) -> "TensorBatch[DictType]": for key, value in self.items(): if value is None: new_batch[key] = value - elif isinstance(value, TensorList): - new_batch[key] = value.repeat_interleave(repeats) else: - assert isinstance(value, torch.Tensor), f"Field {key} must be a tensor, got {type(value)}" + assert isinstance(value, BATCH_FIELD_TYPES), f"Field {key} {_BATCH_FIELD_ERROR}, got {type(value)}" new_batch[key] = value.repeat_interleave(repeats) new_batch = self.__class__(new_batch) new_batch.metadata = self.metadata @@ -354,7 +406,7 @@ def chunk(self, chunk_size: int) -> List["TensorBatch[DictType]"]: chunk_data = {} for key, value in self.items(): if value is not None: - if isinstance(value, (torch.Tensor, TensorList)): + if isinstance(value, BATCH_FIELD_TYPES): chunk_data[key] = value[i : i + chunk_size] else: raise ValueError(f"Unsupported type {type(value)} for key {key}") @@ -381,7 +433,7 @@ def slice(self, start: int, end: int, step: int = 1) -> "TensorBatch[DictType]": sliced_data = {} for key, value in self.items(): if value is not None: - if isinstance(value, (torch.Tensor, TensorList)): + if isinstance(value, BATCH_FIELD_TYPES): sliced_data[key] = value[slice_obj] else: raise ValueError(f"Unsupported type {type(value)} for key {key}") @@ -418,6 +470,8 @@ def cat(cls, shards: List["TensorBatch[DictType]"]) -> "TensorBatch[DictType]": if value is not None: if isinstance(value, TensorList): cat_data[key] = TensorList.cat([shard[key] for shard in shards]) + elif isinstance(value, PackedTensor): + cat_data[key] = PackedTensor.cat([shard[key] for shard in shards]) elif isinstance(value, torch.Tensor): cat_data[key] = torch.cat([shard[key] for shard in shards]) else: @@ -483,8 +537,11 @@ class TrainingInput(TypedDict, total=False): kl: Float[torch.Tensor, "batch_size response_len"] # per-token KL, current vs reference policy rewards: Optional[Float[torch.Tensor, "batch_size response_len"]] # env reward, typically only on the last token rollout_logprobs: Optional[Float[torch.Tensor, "batch_size response_len"]] # sampling policy; off-policy corr. - rollout_expert_indices: Optional[Integer[torch.Tensor, "batch_size seq_len layer_num topk"]] # MoE router replay + # MoE router replay, packed to real tokens: values [sum(seq_len_i), layer_num, topk] + cu_seqlens + rollout_expert_indices: Optional[PackedTensor] router_padding_mask: Optional[Bool[torch.Tensor, "batch_size seq_len"]] # True = no captured route (skip in replay) + # Sampler support, packed to RESPONSE tokens: values [sum(response_len_i), top_k] + cu_seqlens + rollout_sample_support: Optional[PackedTensor] pixel_values: Optional[TensorList] # list of `batch_size` [num_patches_i, dim] tensors image_grid_thw: Optional[TensorList] # list of `batch_size` [num_images_i, 3] tensors @@ -501,6 +558,53 @@ class TrainingOutputBatch(TensorBatch[Dict[str, torch.Tensor]]): pass +@dataclass(frozen=True) +class PackedFieldPadding: + """Padding rule for a packed ``TrainingInput`` field. + + ``dummy_row_length`` is its segment length for a synthetic one-token batch row. + """ + + fill: Callable[[PackedTensor], Union[torch.Tensor, int]] + dummy_row_length: int + + +# Every packed batch field needs a padding rule. +PACKED_FIELD_PADDING: Dict[str, PackedFieldPadding] = { + "rollout_expert_indices": PackedFieldPadding( + # Megatron's dropless `tokens * topk` dispatcher needs topk distinct experts per row. + fill=lambda field: replay_padding_row(field.row_shape[-1], dtype=field.dtype, device=field.device), + dummy_row_length=1, + ), + SAMPLE_SUPPORT_FIELD: PackedFieldPadding( + fill=lambda field: SAMPLE_SUPPORT_PADDING, + dummy_row_length=0, + ), +} + + +def _packed_field_padding_rule(key: str) -> PackedFieldPadding: + if key not in PACKED_FIELD_PADDING: + raise ValueError(f"Packed batch field {key!r} has no padding rule") + return PACKED_FIELD_PADDING[key] + + +def make_packed_field_padding(key: str, field: PackedTensor, *, segment_lengths: Sequence[int]) -> PackedTensor: + """Return padding segments for one packed batch field, filled by that field's own rule.""" + rule = _packed_field_padding_rule(key) + return packed_padding_segments(field, segment_lengths=segment_lengths, fill=rule.fill(field)) + + +def append_packed_field_padding(key: str, field: PackedTensor, *, segment_lengths: Sequence[int]) -> PackedTensor: + """Extend ``field`` with one padding segment per appended batch row.""" + return PackedTensor.cat([field, make_packed_field_padding(key, field, segment_lengths=segment_lengths)]) + + +def packed_dummy_row_segments(key: str, count: int) -> List[int]: + """Segment lengths for ``count`` synthetic batch rows, each carrying one attended token.""" + return [_packed_field_padding_rule(key).dummy_row_length] * count + + def pad_training_input_batch(unpadded_batch: TrainingInputBatch, pad_size: int) -> TrainingInputBatch: """Pad `pad_size` entries to `unpadded_batch`, return a newly allocated TrainingInputBatch. If pad_size is 0, return the original batch.""" # TODO(Charlie): This incurs 2x CPU memory usage when pad_size > 0. Optimize when needed. @@ -528,19 +632,16 @@ def pad_training_input_batch(unpadded_batch: TrainingInputBatch, pad_size: int) assert len(tensor) > 0, f"Cannot pad empty TensorList field {key!r}" padding = TensorList([tensor[0].clone() for _ in range(pad_size)]) new_tensors[key] = TensorList.cat([tensor, padding]) + elif isinstance(tensor, PackedTensor): + # Padded rows copy row 0, including its segment length. + new_tensors[key] = append_packed_field_padding( + key, tensor, segment_lengths=[len(tensor.segment(0))] * pad_size + ) elif key == "loss_mask": # Ensures that padding tensors don't count towards the loss additional_dims = tensor.shape[1:] padding_tensor = torch.zeros(pad_size, *additional_dims, dtype=tensor.dtype, device=tensor.device) new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0) - elif key == "rollout_expert_indices": - additional_dims = tensor.shape[1:] - padding_tensor = make_replay_padding_indices( - (pad_size, *additional_dims), - dtype=tensor.dtype, - device=tensor.device, - ) - new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0) elif key == "router_padding_mask": additional_dims = tensor.shape[1:] padding_tensor = torch.ones(pad_size, *additional_dims, dtype=torch.bool, device=tensor.device) diff --git a/skyrl/backends/skyrl_train/utils/packed_tensor.py b/skyrl/backends/skyrl_train/utils/packed_tensor.py new file mode 100644 index 0000000000..55fcd67f92 --- /dev/null +++ b/skyrl/backends/skyrl_train/utils/packed_tensor.py @@ -0,0 +1,190 @@ +"""One packed buffer plus segment offsets for ragged token-aligned batch fields.""" + +from collections.abc import Sequence + +import torch + +# Megatron's own cu_seqlens dtype; a packed global batch stays far inside int32. +CU_SEQLENS_DTYPE = torch.int32 + + +def cu_seqlens_from_lengths( + sequence_lengths: Sequence[int] | torch.Tensor, + *, + device: torch.device | str | int | None = None, +) -> torch.Tensor: + """Return the ``[batch + 1]`` exclusive prefix sum of ``sequence_lengths``.""" + lengths = torch.as_tensor(sequence_lengths, dtype=CU_SEQLENS_DTYPE, device=device) + if lengths.ndim != 1: + raise ValueError(f"sequence lengths must be 1-D, got shape {lengths.shape}") + if lengths.numel() and int(lengths.min()) < 0: + raise ValueError(f"sequence lengths must be non-negative, got {lengths.tolist()}") + offsets = torch.zeros(lengths.numel() + 1, dtype=CU_SEQLENS_DTYPE, device=lengths.device) + # torch.cumsum promotes to int64; accumulate into the target dtype instead. + torch.cumsum(lengths, dim=0, out=offsets[1:]) + return offsets + + +def lengths_from_offsets(cu_seqlens: torch.Tensor) -> torch.Tensor: + """Return the ``[batch]`` segment lengths that ``cu_seqlens`` encodes.""" + return cu_seqlens[1:] - cu_seqlens[:-1] + + +def row_index_from_offsets( + starts: torch.Tensor, + lengths: torch.Tensor, +) -> torch.Tensor: + """Return row indices that lay the requested segments back to back.""" + starts = starts.to(torch.long) + lengths = lengths.to(torch.long) + total_rows = int(lengths.sum()) + # output_size lets repeat_interleave skip its own device-side sum of `lengths`. + destination_starts = torch.repeat_interleave( + cu_seqlens_from_lengths(lengths, device=lengths.device)[:-1].to(torch.long), + lengths, + output_size=total_rows, + ) + within_segment = torch.arange(total_rows, device=lengths.device) - destination_starts + return torch.repeat_interleave(starts, lengths, output_size=total_rows) + within_segment + + +class PackedTensor: + """A ragged batch of token-aligned rows held as one buffer plus ``cu_seqlens``. + + ``values`` is ``[sum(sequence_lengths), *row_shape]`` in canonical batch order and + ``cu_seqlens`` is the ``[batch + 1]`` exclusive prefix sum of the segment lengths. + Indexing and batch operations address segments rather than individual rows. + """ + + def __init__(self, values: torch.Tensor, cu_seqlens: torch.Tensor): + if values.ndim < 1: + raise ValueError("packed values must have a token-row dimension") + if cu_seqlens.ndim != 1 or cu_seqlens.numel() < 2: + raise ValueError(f"cu_seqlens must hold at least two offsets, got shape {cu_seqlens.shape}") + if cu_seqlens.dtype != CU_SEQLENS_DTYPE: + raise ValueError(f"cu_seqlens must be {CU_SEQLENS_DTYPE}, got {cu_seqlens.dtype}") + if cu_seqlens.device != values.device: + raise ValueError( + f"packed values and cu_seqlens must share a device, got {values.device} and {cu_seqlens.device}" + ) + if int(cu_seqlens[0]) != 0 or int(cu_seqlens[-1]) != values.shape[0]: + raise ValueError( + f"cu_seqlens must run from 0 to the {values.shape[0]} packed rows, " + f"got {int(cu_seqlens[0])} to {int(cu_seqlens[-1])}" + ) + self.values = values + self.cu_seqlens = cu_seqlens + + @classmethod + def from_segments(cls, segments: Sequence[torch.Tensor]) -> "PackedTensor": + """Concatenate per-batch-entry row blocks into one packed buffer.""" + if not segments: + raise ValueError("cannot pack an empty list of segments") + cu_seqlens = cu_seqlens_from_lengths([segment.shape[0] for segment in segments], device=segments[0].device) + return cls(torch.cat(segments, dim=0), cu_seqlens) + + @property + def sequence_lengths(self) -> torch.Tensor: + return lengths_from_offsets(self.cu_seqlens) + + @property + def row_shape(self) -> torch.Size: + return self.values.shape[1:] + + @property + def device(self) -> torch.device: + return self.values.device + + @property + def dtype(self) -> torch.dtype: + return self.values.dtype + + def __len__(self) -> int: + return self.cu_seqlens.numel() - 1 + + def __getitem__(self, index) -> "torch.Tensor | PackedTensor": + if isinstance(index, slice): + if index.step in (None, 1): + start, stop, _ = index.indices(len(self)) + stop = max(start, stop) + offsets = self.cu_seqlens[start : stop + 1] + return PackedTensor(self.values[int(offsets[0]) : int(offsets[-1])], offsets - offsets[0]) + return self._gather(range(*index.indices(len(self)))) + if isinstance(index, torch.Tensor): + if index.ndim == 0: + return self.segment(int(index)) + return self._gather(index.tolist()) + if isinstance(index, (list, tuple, range)): + return self._gather(index) + return self.segment(index) + + def segment(self, index: int) -> torch.Tensor: + """Return one batch entry's row block as a view.""" + position = index + len(self) if index < 0 else index + if not 0 <= position < len(self): + raise IndexError(f"segment {index} is out of range for a packed batch of {len(self)}") + return self.values[int(self.cu_seqlens[position]) : int(self.cu_seqlens[position + 1])] + + def _gather(self, indices: Sequence[int]) -> "PackedTensor": + """Select segments in the requested order into a freshly allocated buffer.""" + selected = torch.as_tensor(list(indices), dtype=torch.long, device=self.values.device) + selected_starts = self.cu_seqlens[:-1].to(torch.long)[selected] + selected_lengths = self.sequence_lengths.to(torch.long)[selected] + row_index = row_index_from_offsets(selected_starts, selected_lengths) + cu_seqlens = cu_seqlens_from_lengths(selected_lengths, device=self.values.device) + return PackedTensor(self.values.index_select(0, row_index), cu_seqlens) + + def to(self, device=None, dtype=None, non_blocking: bool = False) -> "PackedTensor": + return PackedTensor( + self.values.to(device=device, dtype=dtype, non_blocking=non_blocking), + self.cu_seqlens.to(device=device, non_blocking=non_blocking), + ) + + def contiguous(self) -> "PackedTensor": + return PackedTensor(self.values.contiguous(), self.cu_seqlens.contiguous()) + + def pin_memory(self) -> "PackedTensor": + return PackedTensor(self.values.pin_memory(), self.cu_seqlens.pin_memory()) + + def repeat(self, repeats: int) -> "PackedTensor": + return self._gather(list(range(len(self))) * repeats) + + def repeat_interleave(self, repeats: int) -> "PackedTensor": + return self._gather([index for index in range(len(self)) for _ in range(repeats)]) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, PackedTensor): + return False + return torch.equal(self.values, other.values) and torch.equal(self.cu_seqlens, other.cu_seqlens) + + def __repr__(self) -> str: + return f"PackedTensor(batch={len(self)}, values={tuple(self.values.shape)}, dtype={self.values.dtype})" + + @staticmethod + def cat(batches: Sequence["PackedTensor"]) -> "PackedTensor": + if not batches: + raise ValueError("cannot cat an empty list of packed batches") + lengths = torch.cat([batch.sequence_lengths for batch in batches]) + return PackedTensor( + torch.cat([batch.values for batch in batches], dim=0), + cu_seqlens_from_lengths(lengths, device=batches[0].device), + ) + + +def packed_padding_segments( + reference: PackedTensor, + *, + segment_lengths: Sequence[int], + fill: torch.Tensor | int | float | bool, +) -> PackedTensor: + """Return ``fill``-valued segments in ``reference``'s row shape, dtype and device. + + ``fill`` broadcasts over the row shape, so it may be a scalar or one whole row. + """ + values = torch.empty( + (sum(segment_lengths), *reference.row_shape), + dtype=reference.dtype, + device=reference.device, + ) + values[...] = fill + return PackedTensor(values, cu_seqlens_from_lengths(segment_lengths, device=reference.device)) diff --git a/skyrl/backends/skyrl_train/utils/replay_utils.py b/skyrl/backends/skyrl_train/utils/replay_utils.py index 400376c65a..47a538be67 100644 --- a/skyrl/backends/skyrl_train/utils/replay_utils.py +++ b/skyrl/backends/skyrl_train/utils/replay_utils.py @@ -4,16 +4,17 @@ from contextlib import contextmanager -import numpy as np import torch from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( TokenMetadataLayout, + align_packed_token_metadata, align_token_metadata, ) +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor -def _replay_padding_row( +def replay_padding_row( topk: int, *, dtype: torch.dtype, @@ -40,27 +41,10 @@ def make_replay_padding_indices( """Return dummy routes with ``topk`` distinct experts in every row.""" if not shape: raise ValueError(f"Replay route padding requires a positive topk dimension, got {shape}") - padding_row = _replay_padding_row(shape[-1], dtype=dtype, device=device) + padding_row = replay_padding_row(shape[-1], dtype=dtype, device=device) return padding_row.expand(shape).clone() -def make_replay_padding_indices_np(shape: tuple[int, ...], *, dtype: np.dtype) -> np.ndarray: - """NumPy sibling of :func:`make_replay_padding_indices`. - - Preprocessing builds the padded route array in NumPy before handing it to - ``torch.from_numpy``, so it needs the same distinct-expert padding rows - without a round trip through torch. - """ - if not shape: - raise ValueError(f"Replay route padding requires a positive topk dimension, got {shape}") - topk = shape[-1] - if topk < 1: - raise ValueError(f"Replay route padding requires a positive topk dimension, got {topk}") - padded = np.empty(shape, dtype=dtype) - padded[...] = np.arange(topk, dtype=dtype) - return padded - - def patch_topk_router_layer_number(): """Monkey-patch TopKRouter.set_layer_number to propagate the global layer number to the RouterReplay instance. @@ -187,7 +171,7 @@ def _get_local_router_layer_indices(model_config, global_num_layers: int, instan def setup_per_microbatch_replay_forward( - rollout_expert_indices: torch.Tensor, + rollout_expert_indices: PackedTensor, router_padding_mask: torch.Tensor | None, attention_mask: torch.Tensor, model, @@ -197,8 +181,9 @@ def setup_per_microbatch_replay_forward( ) -> dict[str, torch.Tensor]: """Set up router replay and return its model-facing keyword arguments. - Replay indices and the router padding mask start in the same batch layout and - undergo matching padding removal or packing and CP sharding. Their destinations + Replay indices arrive packed to their real tokens (``[sum(seqlen), layers, topk]`` plus + ``cu_seqlens``) while the router padding mask arrives batch-major; both undergo matching + padding removal or packing and CP sharding against the shared layout. Their destinations then differ: indices are TP-sliced and installed into per-layer ``RouterReplay`` instances, while the mask follows Megatron's model-specific sequence-parallel path and is passed to the model as ``padding_mask``. @@ -232,8 +217,8 @@ def setup_per_microbatch_replay_forward( if router_padding_mask is None: raise ValueError("router_padding_mask is required with rollout_expert_indices") - if rollout_expert_indices.dim() != 4: - raise ValueError(f"Expected 4D replay indices, got shape {rollout_expert_indices.shape}") + if len(rollout_expert_indices.row_shape) != 2: + raise ValueError(f"Expected [tokens, layers, topk] replay indices, got {rollout_expert_indices!r}") if router_padding_mask.shape != attention_mask.shape: raise ValueError( @@ -243,24 +228,28 @@ def setup_per_microbatch_replay_forward( if router_padding_mask.device != rollout_expert_indices.device: raise ValueError("rollout_expert_indices and router_padding_mask must be on the same device") + num_captured_layers, topk = rollout_expert_indices.row_shape instances = RouterReplay.global_router_replay_instances local_layer_indices = _get_local_router_layer_indices( model_config, - rollout_expert_indices.shape[2], + num_captured_layers, instances, ) layer_index = torch.tensor(local_layer_indices, dtype=torch.long, device=rollout_expert_indices.device) - local_rollout_expert_indices = rollout_expert_indices.index_select(2, layer_index) + local_rollout_expert_indices = PackedTensor( + rollout_expert_indices.values.index_select(1, layer_index), + rollout_expert_indices.cu_seqlens, + ) if (metadata_layout.padded_sequence_lengths is not None) != remove_microbatch_padding: raise ValueError("Shared token metadata layout does not match the model packing mode") aligned_router_padding_mask = align_token_metadata(router_padding_mask.to(torch.bool), metadata_layout, True) - route_padding = _replay_padding_row( - rollout_expert_indices.shape[-1], + route_padding = replay_padding_row( + topk, dtype=rollout_expert_indices.dtype, - device=local_rollout_expert_indices.device, + device=rollout_expert_indices.device, ) - aligned_rollout_expert_indices = align_token_metadata( + aligned_rollout_expert_indices = align_packed_token_metadata( local_rollout_expert_indices, metadata_layout, route_padding, diff --git a/skyrl/backends/skyrl_train/utils/routed_experts.py b/skyrl/backends/skyrl_train/utils/routed_experts.py index b3caec7c8b..cef230728b 100644 --- a/skyrl/backends/skyrl_train/utils/routed_experts.py +++ b/skyrl/backends/skyrl_train/utils/routed_experts.py @@ -1,11 +1,56 @@ +from collections.abc import Sequence from typing import TypeAlias import numpy as np +from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( + TokenMetadataTrace, +) + RoutedExpertIndices: TypeAlias = np.ndarray ROUTED_EXPERT_DTYPES = frozenset({np.dtype(np.uint8), np.dtype(np.int16), np.dtype(np.int32)}) +class RoutedExpertTrace: + """Accumulate routed experts across incremental generation calls.""" + + def __init__(self) -> None: + self._metadata = TokenMetadataTrace() + + @property + def prompt_start(self) -> int: + return self._metadata.num_rows + + def record_generation( + self, + *, + prompt_token_count: int, + generated_token_count: int, + routed_experts: RoutedExpertIndices, + ) -> None: + if prompt_token_count < self.prompt_start: + raise ValueError("routed-expert prompt start exceeds prompt length") + if generated_token_count < 1: + raise ValueError("routed-expert generation must produce at least one token") + + expected_rows = prompt_token_count - self.prompt_start + generated_token_count - 1 + self._metadata.append(compact_routed_expert_indices(routed_experts), expected_rows=expected_rows) + + def finalize(self, *, token_count: int, loss_mask: Sequence[int]) -> RoutedExpertIndices: + """Return the captured route prefix without fabricating rows for its uncovered tail.""" + if len(loss_mask) != token_count: + raise ValueError(f"loss mask has {len(loss_mask)} entries, expected {token_count}") + if self.prompt_start > token_count: + raise ValueError(f"routed-expert trace has {self.prompt_start} rows for {token_count} tokens") + + if any(loss_mask[self.prompt_start + 1 : token_count]): + for source_index in range(self.prompt_start, token_count - 1): + if loss_mask[source_index + 1] != 0: + raise ValueError(f"missing routed-expert row for loss-active target at token {source_index + 1}") + + return self._metadata.finalize(expected_rows=self.prompt_start) + + def compact_routed_expert_indices(routed_experts: RoutedExpertIndices) -> RoutedExpertIndices: """Validate and compact a routed-expert array to the canonical integer dtype.""" if not isinstance(routed_experts, np.ndarray): diff --git a/skyrl/backends/skyrl_train/utils/sample_support.py b/skyrl/backends/skyrl_train/utils/sample_support.py new file mode 100644 index 0000000000..341b7694ee --- /dev/null +++ b/skyrl/backends/skyrl_train/utils/sample_support.py @@ -0,0 +1,109 @@ +"""Per-token bounded sampler support used to renormalize rollout logprobs. + +Rows contain top-k vocabulary IDs and use trailing ``SAMPLE_SUPPORT_PADDING``. +Tokens without captured support use an all-padding row. +""" + +from typing import TypeAlias + +import numpy as np +import torch + +from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( + TokenMetadataLayout, + TokenMetadataTrace, + align_packed_token_metadata, +) +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor + +SampleSupport: TypeAlias = np.ndarray +SAMPLE_SUPPORT_DTYPE = np.dtype(np.int32) +SAMPLE_SUPPORT_TORCH_DTYPE = torch.int32 +SAMPLE_SUPPORT_DTYPES = frozenset({SAMPLE_SUPPORT_DTYPE}) +SAMPLE_SUPPORT_PADDING = -1 +SAMPLE_SUPPORT_FIELD = "rollout_sample_support" +# Sentinel outside the valid packed-row range. +SAMPLE_SUPPORT_NO_ROW = -1 + + +def validate_sample_support(sample_support: SampleSupport) -> SampleSupport: + """Validate vocabulary IDs and trailing padding.""" + if not isinstance(sample_support, np.ndarray): + raise TypeError("sample support must be a NumPy array") + if sample_support.ndim != 2 or not np.issubdtype(sample_support.dtype, np.integer): + raise ValueError( + "sample support must be an integer [tokens, top_k] array, " + f"got shape {sample_support.shape} and dtype {sample_support.dtype}" + ) + if int(sample_support.min(initial=0)) < SAMPLE_SUPPORT_PADDING: + raise ValueError(f"sample support IDs must be {SAMPLE_SUPPORT_PADDING} padding or non-negative vocab IDs") + if np.any((sample_support[:, :-1] == SAMPLE_SUPPORT_PADDING) & (sample_support[:, 1:] >= 0)): + raise ValueError(f"sample support padding must be trailing {SAMPLE_SUPPORT_PADDING} values") + return sample_support + + +def align_sample_support_row_ids( + sample_support: PackedTensor, + layout: TokenMetadataLayout, +) -> torch.Tensor: + """Map model positions to packed support rows. + + Support for response tokens occupies ``[prompt_len - 1, sequence_len - 1)`` because + position ``t`` predicts token ``t + 1``. Derive IDs per micro-batch because slicing and + padding rebase the packed row space. + """ + segment_lengths = sample_support.sequence_lengths.to(torch.long) + if segment_lengths.numel() != len(layout.sequence_lengths): + raise ValueError( + f"Sample support holds {segment_lengths.numel()} segments for " + f"{len(layout.sequence_lengths)} trajectories" + ) + trajectory_lengths = torch.as_tensor( + layout.sequence_lengths, + dtype=torch.long, + device=segment_lengths.device, + ) + # The first response token is predicted at prompt_len - 1. + segment_starts = trajectory_lengths - segment_lengths - 1 + if segment_lengths.numel() and int(segment_starts.min()) < 0: + raise ValueError( + "A trajectory whose support covers all of its real tokens has no position that " + f"predicts its first response token, got lengths {segment_lengths.tolist()} for " + f"trajectories {trajectory_lengths.tolist()}" + ) + row_ids = PackedTensor( + torch.arange(sample_support.values.shape[0], dtype=torch.long, device=sample_support.device), + sample_support.cu_seqlens, + ) + return align_packed_token_metadata( + row_ids, + layout, + SAMPLE_SUPPORT_NO_ROW, + segment_starts=segment_starts.tolist(), + ) + + +class SampleSupportTrace: + """Accumulate sample support across incremental generation calls.""" + + def __init__(self) -> None: + self._metadata = TokenMetadataTrace() + + @property + def num_rows(self) -> int: + return self._metadata.num_rows + + def append(self, sample_support: SampleSupport, *, expected_rows: int) -> None: + self._metadata.append(validate_sample_support(sample_support), expected_rows=expected_rows) + + def append_padding(self, count: int) -> None: + self._metadata.append_padding(count, fill=SAMPLE_SUPPORT_PADDING) + + def finalize(self, *, token_count: int, extra_rows: int) -> SampleSupport: + """Validate the trace length and discard ``extra_rows`` trailing rows.""" + if self.num_rows != token_count + extra_rows: + raise ValueError( + f"sample-support trace has {self.num_rows} rows for {token_count} tokens plus " + f"{extra_rows} trailing rows" + ) + return self._metadata.finalize(expected_rows=self.num_rows)[:token_count] diff --git a/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py b/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py index 41c3a55a64..3a36dbf876 100644 --- a/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py +++ b/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py @@ -32,6 +32,7 @@ decompress_bytes, uint8_tensor_to_bytes, ) +from skyrl.utils.cpu_topology import pool_workers logger = logging.getLogger(__name__) @@ -760,7 +761,8 @@ def apply_one(item: tuple[DeltaTensorRecord, str, bytes]) -> None: mismatches.append(record.name) del region, patch - workers = min(len(payloads), max(1, min(32, os.cpu_count() or 8))) + # Weight sync is a barrier, so it need not reserve cores for colocated work. + workers = min(len(payloads), pool_workers(cap=32, reserved=0)) with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="skyrl-delta-mmap-apply") as executor: list(executor.map(apply_one, payloads)) finally: @@ -1077,7 +1079,8 @@ def _empty_stats() -> dict[str, float]: } def _num_publish_workers(self) -> int: - default = min(8, os.cpu_count() or 1) + # Publishing is a barrier, like the apply path above. + default = pool_workers(cap=8, reserved=0) return self.publish_num_workers or default def _publish_executor_for(self, num_workers: int) -> ThreadPoolExecutor: diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index e66c42841e..0a70b4a45e 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -42,6 +42,7 @@ unpadded_vocab_shard_width, ) from skyrl.backends.skyrl_train.training_batch import TensorList +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor from skyrl.backends.skyrl_train.utils.ppo_utils import ( PolicyLossRegistry, compute_approx_kl, @@ -132,7 +133,7 @@ def _build_packed_valid_mask( def _copy_tensor_tree_to_device(value: Any, device: int) -> Any: """Move all tensors in a nested microbatch to a CUDA device.""" - if torch.is_tensor(value) or isinstance(value, TensorList): + if torch.is_tensor(value) or isinstance(value, (TensorList, PackedTensor)): return value.to(device=device, non_blocking=True) if isinstance(value, dict): return {key: _copy_tensor_tree_to_device(item, device) for key, item in value.items()} diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index d116213abd..808544bc2a 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -45,9 +45,11 @@ from skyrl.backends.skyrl_train.training_batch import ( TrainingInputBatch, TrainingOutputBatch, + append_packed_field_padding, + packed_dummy_row_segments, ) +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor from skyrl.backends.skyrl_train.utils.profiler import build_profiler_from_policy_cfg -from skyrl.backends.skyrl_train.utils.replay_utils import make_replay_padding_indices from skyrl.backends.skyrl_train.weight_sync import ( LoraLoadRequest, WeightChunk, @@ -493,6 +495,15 @@ def init_configs( for k, v in transformer_config_kwargs.items(): setattr(provider, k, v) + # Check the resolved provider because it may supply its own VPP default. Interleaved + # chunks desynchronise each RouterReplay instance's backward FIFO. + vpp_size = provider.virtual_pipeline_model_parallel_size + if provider.moe_enable_routing_replay and vpp_size is not None and vpp_size > 1: + raise ValueError( + f"moe_enable_routing_replay is incompatible with virtual_pipeline_model_parallel_size={vpp_size}: " + "interleaved chunks desync the replay FIFO. Unset virtual_pipeline_model_parallel_size." + ) + # MTP head count: megatron-bridge infers provider.mtp_num_layers from the model's HF config. if not enable_mtp: provider.mtp_num_layers = None @@ -774,6 +785,12 @@ def _pad_microbatch_to_size(self, micro_dict: dict, target_batch_size: int) -> d if value is None: padded[key] = None continue + if isinstance(value, PackedTensor): + # Per-token fields cover the dummy attended token; response fields do not. + padded[key] = append_packed_field_padding( + key, value, segment_lengths=packed_dummy_row_segments(key, pad_count) + ) + continue if isinstance(value, torch.Tensor): if key == "loss_mask": # Pad with zeros so padded samples don't contribute to loss @@ -791,12 +808,6 @@ def _pad_microbatch_to_size(self, micro_dict: dict, target_batch_size: int) -> d pad_tensor = torch.arange(seq_len, device=device).unsqueeze(0).expand(pad_count, -1) elif key == "router_padding_mask": pad_tensor = torch.ones((pad_count, *value.shape[1:]), dtype=torch.bool, device=device) - elif key == "rollout_expert_indices": - pad_tensor = make_replay_padding_indices( - (pad_count, *value.shape[1:]), - dtype=value.dtype, - device=device, - ) elif key == "response_mask": # response_mask should be zeros for padded samples pad_tensor = torch.zeros((pad_count, *value.shape[1:]), dtype=value.dtype, device=device) diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py index 2efa37ad6c..cf6541fefc 100644 --- a/skyrl/backends/skyrl_train/workers/worker_utils.py +++ b/skyrl/backends/skyrl_train/workers/worker_utils.py @@ -5,8 +5,14 @@ import torch.distributed as dist from skyrl.backends.skyrl_train.distributed.strategy import DistributedStrategy -from skyrl.backends.skyrl_train.training_batch import TensorBatch, TrainingInputBatch -from skyrl.backends.skyrl_train.utils.replay_utils import make_replay_padding_indices +from skyrl.backends.skyrl_train.training_batch import ( + PACKED_FIELD_PADDING, + TensorBatch, + TrainingInputBatch, + make_packed_field_padding, + packed_dummy_row_segments, +) +from skyrl.backends.skyrl_train.utils.sample_support import SAMPLE_SUPPORT_FIELD from skyrl.backends.skyrl_train.utils.torch_utils import masked_mean from skyrl.train.dataset.bin_packing import make_seq_packer from skyrl.train.dataset.replay_buffer import Experience @@ -164,6 +170,7 @@ def batch_to_experience(batch: TrainingInputBatch): rollout_logprobs=batch.get("rollout_logprobs"), rollout_expert_indices=batch.get("rollout_expert_indices"), router_padding_mask=batch.get("router_padding_mask"), + rollout_sample_support=batch.get(SAMPLE_SUPPORT_FIELD), # additional info # can be used to log metrics etc for micro-batches in the worker info={}, @@ -321,17 +328,18 @@ def _create_padding_microbatch(self) -> TrainingInputBatch: "response_mask": torch.ones((batch_size, num_actions), dtype=int, device=device), } ) - # Add optional fields such as `rollout_logprobs` and `rollout_expert_indices` to padding batch + # Add optional fields to the padding batch. if self.data.get("rollout_logprobs") is not None: ref_tensor = self.data["rollout_logprobs"] data["rollout_logprobs"] = torch.zeros((batch_size, num_actions), dtype=ref_tensor.dtype, device=device) - if self.data.get("rollout_expert_indices") is not None: - ref_tensor = self.data["rollout_expert_indices"] - data["rollout_expert_indices"] = make_replay_padding_indices( - (batch_size, *ref_tensor.shape[1:]), - dtype=ref_tensor.dtype, - device=device, - ) + for key in PACKED_FIELD_PADDING: + # Per-token fields cover the dummy attended token; response fields do not. + if self.data.get(key) is not None: + data[key] = make_packed_field_padding( + key, + self.data[key], + segment_lengths=packed_dummy_row_segments(key, batch_size), + ) if self.data.get("router_padding_mask") is not None: data["router_padding_mask"] = torch.ones((batch_size, seq_len), dtype=torch.bool, device=device) data.metadata = {} diff --git a/skyrl/benchmarks/bench_packed_route_collation.py b/skyrl/benchmarks/bench_packed_route_collation.py new file mode 100644 index 0000000000..668a9d157d --- /dev/null +++ b/skyrl/benchmarks/bench_packed_route_collation.py @@ -0,0 +1,210 @@ +"""Compare padded and packed route collation with serial and pooled fills. + +Production shapes need ~110 GiB of host RAM, so drive this from a cluster harness:: + + uv run --isolated --extra skyrl-train python -m \ + skyrl.benchmarks.bench_packed_route_collation --num-moe-layers 40 + +Scaled-down smoke:: + + uv run --isolated --extra skyrl-train python -m \ + skyrl.benchmarks.bench_packed_route_collation \ + --num-sequences 64 --max-seqlen 4096 --num-moe-layers 4 --iterations 1 +""" + +import argparse +import functools +import gc +import os +import resource +import statistics +import time + +import numpy as np +import torch + +from skyrl.backends.skyrl_train.utils.packed_tensor import cu_seqlens_from_lengths +from skyrl.backends.skyrl_train.utils.replay_utils import replay_padding_row +from skyrl.train.dataset.parallel_fill import default_fill_workers, fill_batch_rows + +DEFAULT_NUM_MOE_LAYERS = 40 +DEFAULT_TOPK = 22 +DEFAULT_NUM_SEQUENCES = 1024 +DEFAULT_MAX_SEQLEN = 32768 +ROUTE_DTYPE = torch.int16 + +# Fraction of ``max_seqlen`` each distribution draws its shortest sequence from. "uniform" +# has no padding at all; "typical_rl" is the measured production spread. +LENGTH_DISTRIBUTIONS = { + "uniform": 1.0, + "mild_ragged": 0.5, + "typical_rl": 1 / 16, + "heavy_tail": 1 / 64, +} + + +def _sequence_lengths(distribution: str, num_sequences: int, max_seqlen: int, seed: int) -> np.ndarray: + """Draw per-trajectory total lengths, always including one full-length sequence.""" + minimum = max(1, round(max_seqlen * LENGTH_DISTRIBUTIONS[distribution])) + if minimum >= max_seqlen: + return np.full(num_sequences, max_seqlen, dtype=np.int64) + rng = np.random.default_rng(seed) + lengths = rng.integers(minimum, max_seqlen + 1, size=num_sequences).astype(np.int64) + # max_total is set by the longest trajectory, so pin one to the cap for a stable rectangle. + lengths[0] = max_seqlen + return lengths + + +def _make_trajectories(lengths: np.ndarray, num_layers: int, topk: int, seed: int) -> list[np.ndarray]: + """One route array per trajectory, sized to its full sequence length. + + Arrays are views over one template so source allocation does not dominate the benchmark. + """ + rng = np.random.default_rng(seed + 1) + template = rng.integers(0, 128, size=(int(lengths.max()), num_layers, topk), dtype=np.int16) + return [template[: int(length)] for length in lengths] + + +def _write_padded_row( + padded: torch.Tensor, + trajectories: list[np.ndarray], + lengths: np.ndarray, + sample_index: int, +) -> None: + """One trajectory's slot in the rectangle: left dummy rows, routes, trailing dummy rows.""" + padding_row = replay_padding_row(padded.shape[-1], dtype=padded.dtype) + sample_indices = trajectories[sample_index] + left_pad = padded.shape[1] - int(lengths[sample_index]) + route_end = left_pad + sample_indices.shape[0] + padded[sample_index, :left_pad] = padding_row + padded[sample_index, left_pad:route_end] = torch.from_numpy(sample_indices) + padded[sample_index, route_end:] = padding_row + + +def _write_packed_segment( + packed: torch.Tensor, + cu_seqlens: torch.Tensor, + trajectories: list[np.ndarray], + sample_index: int, +) -> None: + """One trajectory's segment of the packed buffer: routes, then any trailing dummy rows.""" + sample_indices = trajectories[sample_index] + segment = packed[int(cu_seqlens[sample_index]) : int(cu_seqlens[sample_index + 1])] + captured = sample_indices.shape[0] + segment[:captured] = torch.from_numpy(sample_indices) + segment[captured:] = replay_padding_row(segment.shape[-1], dtype=packed.dtype) + + +def _make_fill(packed: bool, workers: int): + """Build a fill callable over the trainer's own pool helper.""" + + def fill(buffer: torch.Tensor, trajectories: list[np.ndarray], lengths: np.ndarray) -> None: + if packed: + cu_seqlens = cu_seqlens_from_lengths(lengths) + write = functools.partial(_write_packed_segment, buffer, cu_seqlens, trajectories) + else: + write = functools.partial(_write_padded_row, buffer, trajectories, lengths) + fill_batch_rows(write, len(trajectories), workers=workers) + + return fill + + +def _padded_shape(lengths: np.ndarray, num_layers: int, topk: int) -> tuple[int, ...]: + return (len(lengths), int(lengths.max()), num_layers, topk) + + +def _packed_shape(lengths: np.ndarray, num_layers: int, topk: int) -> tuple[int, ...]: + return (int(lengths.sum()), num_layers, topk) + + +def _peak_rss_bytes() -> int: + """Process high-water RSS. Monotone, so only ever read as a whole-run ceiling.""" + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024 + + +def _time_cold(fill, shape, trajectories, lengths, iterations: int) -> tuple[float, int]: + """Median wall clock over a freshly allocated buffer each iteration. + + Fresh buffers retain the first-touch allocation cost measured in production. + """ + durations = [] + for _ in range(iterations): + gc.collect() + start = time.perf_counter() + buffer = torch.empty(shape, dtype=ROUTE_DTYPE) + fill(buffer, trajectories, lengths) + durations.append(time.perf_counter() - start) + buffer_bytes = buffer.numel() * buffer.element_size() + del buffer + return statistics.median(durations), buffer_bytes + + +def _format_gib(num_bytes: int) -> str: + return f"{num_bytes / 1024**3:8.2f}" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--num-sequences", type=int, default=DEFAULT_NUM_SEQUENCES) + parser.add_argument("--max-seqlen", type=int, default=DEFAULT_MAX_SEQLEN) + parser.add_argument("--num-moe-layers", type=int, default=DEFAULT_NUM_MOE_LAYERS) + parser.add_argument("--topk", type=int, default=DEFAULT_TOPK) + parser.add_argument("--iterations", type=int, default=3) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--workers", type=int, default=default_fill_workers()) + parser.add_argument( + "--distributions", + nargs="+", + default=list(LENGTH_DISTRIBUTIONS), + choices=list(LENGTH_DISTRIBUTIONS), + ) + args = parser.parse_args() + + print( + f"num_sequences={args.num_sequences} max_seqlen={args.max_seqlen} " + f"moe_layers={args.num_moe_layers} topk={args.topk} dtype={ROUTE_DTYPE} " + f"iterations={args.iterations}" + ) + print(f"OMP_NUM_THREADS={os.environ.get('OMP_NUM_THREADS', '')} torch_threads={torch.get_num_threads()}") + print(f"pooled arms use {args.workers} workers (the trainer's autoscaled pool size)") + header = ( + f"{'distribution':<20} {'pad_1t':>8} {'pad_pool':>9} {'pack_1t':>8} {'pack_pool':>10} " + f"{'best_pad':>9} {'vs_best':>8} {'pad_GiB':>9} {'packed_GiB':>11} {'saved':>7}" + ) + print(header) + print("-" * len(header)) + + for distribution in args.distributions: + lengths = _sequence_lengths(distribution, args.num_sequences, args.max_seqlen, args.seed) + trajectories = _make_trajectories(lengths, args.num_moe_layers, args.topk, args.seed) + + padded_shape = _padded_shape(lengths, args.num_moe_layers, args.topk) + packed_shape = _packed_shape(lengths, args.num_moe_layers, args.topk) + arms = { + "pad_1t": (_make_fill(False, 1), padded_shape), + "pad_pool": (_make_fill(False, args.workers), padded_shape), + "pack_1t": (_make_fill(True, 1), packed_shape), + "pack_pool": (_make_fill(True, args.workers), packed_shape), + } + timings = {} + buffer_bytes = {} + for name, (fill, shape) in arms.items(): + timings[name], buffer_bytes[name] = _time_cold(fill, shape, trajectories, lengths, args.iterations) + del trajectories + gc.collect() + + best_padded = min(timings["pad_1t"], timings["pad_pool"]) + best_packed = min(timings["pack_1t"], timings["pack_pool"]) + print( + f"{distribution:<20} {timings['pad_1t'] * 1000:8.1f} {timings['pad_pool'] * 1000:9.1f} " + f"{timings['pack_1t'] * 1000:8.1f} {timings['pack_pool'] * 1000:10.1f} " + f"{best_padded * 1000:9.1f} {best_padded / best_packed:7.2f}x " + f"{_format_gib(buffer_bytes['pad_1t'])} {_format_gib(buffer_bytes['pack_1t']):>11} " + f"{1 - buffer_bytes['pack_1t'] / buffer_bytes['pad_1t']:6.1%}" + ) + + print(f"process peak RSS: {_format_gib(_peak_rss_bytes()).strip()} GiB") + + +if __name__ == "__main__": + main() diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 57c7521910..3aadb26c34 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -1142,6 +1142,8 @@ class InferenceEngineConfig(BaseConfig): enable_return_routed_experts: bool = False """Return per-layer expert routing indices, for rollout router replay (R3) when training an MoE model. Used together with ``trainer.policy.megatron_config.moe_enable_routing_replay``.""" + enable_return_sample_support_set: bool = False + """Return the bounded sampler support used to renormalize rollout logprobs.""" max_num_batched_tokens: int = 8192 """vLLM continuous-batching parameter: maximum number of tokens to pack into a batch.""" enforce_eager: bool = False @@ -1737,6 +1739,24 @@ def __post_init__(self): if self.trainer.algorithm.temperature is None: self.trainer.algorithm.temperature = self.generator.sampling_params.temperature + # Eval requests opt out of capture and do not use these constraints. + if self.generator.inference_engine.enable_return_sample_support_set: + sampling_params = self.generator.sampling_params + if sampling_params.temperature <= 0: + raise ValueError("sample-support capture requires generator.sampling_params.temperature > 0") + if sampling_params.top_k <= 1: + raise ValueError("sample-support capture requires generator.sampling_params.top_k > 1") + if sampling_params.repetition_penalty != 1.0: + raise ValueError("sample-support capture requires repetition_penalty=1.0") + if sampling_params.additional_kwargs: + raise ValueError("sample-support capture does not support sampling_params.additional_kwargs") + if self.generator.vision_language_generator: + raise ValueError("sample-support capture does not support vision_language_generator") + + # The VLM generator does not populate routed-expert indices. + if self.generator.inference_engine.enable_return_routed_experts and self.generator.vision_language_generator: + raise ValueError("rollout router replay (r3) does not support vision_language_generator") + if self.data.dataloader.num_workers is None: self.data.dataloader.num_workers = 8 if self.data.dataloader.persistent_workers and self.data.dataloader.num_workers == 0: diff --git a/skyrl/train/dataset/parallel_fill.py b/skyrl/train/dataset/parallel_fill.py new file mode 100644 index 0000000000..813884604d --- /dev/null +++ b/skyrl/train/dataset/parallel_fill.py @@ -0,0 +1,51 @@ +"""Fill a controller-side batch buffer with a locally sized thread pool. + +The pool parallelises first-touch page faults without changing process-wide torch settings. +Callbacks own disjoint row ranges, and their copies release the GIL. +""" + +import functools +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor + +from skyrl.utils.cpu_topology import pool_workers + +# Extra threads beyond this cap take cores from colocated actors without improving throughput. +MAX_FILL_WORKERS = 32 +# Leave room for Ray services in the same cgroup. +RESERVED_FILL_CORES = 8 + + +@functools.cache +def default_fill_workers() -> int: + return pool_workers(cap=MAX_FILL_WORKERS, reserved=RESERVED_FILL_CORES) + + +def fill_batch_rows( + fill_row: Callable[[int], None], + num_rows: int, + *, + workers: int | None = None, +) -> None: + """Call ``fill_row`` for every row, possibly in parallel. + + Each callback must write to a disjoint row range. + """ + if num_rows < 0: + raise ValueError(f"row count must be non-negative, got {num_rows}") + if num_rows == 0: + return + if workers is None: + workers = default_fill_workers() + if workers < 1: + raise ValueError(f"worker count must be positive, got {workers}") + + workers = min(workers, num_rows) + if workers == 1: + for index in range(num_rows): + fill_row(index) + return + + with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="skyrl-batch-fill") as pool: + # Eagerly consume the map so worker exceptions surface here. + list(pool.map(fill_row, range(num_rows))) diff --git a/skyrl/train/dataset/preprocess.py b/skyrl/train/dataset/preprocess.py index fb7ac6e2f7..19b8c20637 100644 --- a/skyrl/train/dataset/preprocess.py +++ b/skyrl/train/dataset/preprocess.py @@ -1,18 +1,36 @@ +import functools import logging -from typing import List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import numpy as np import torch -from jaxtyping import Bool, Float, Integer +from jaxtyping import Bool, Float -from skyrl.backends.skyrl_train.utils.replay_utils import make_replay_padding_indices_np +from skyrl.backends.skyrl_train.utils.packed_tensor import ( + PackedTensor, + cu_seqlens_from_lengths, +) +from skyrl.backends.skyrl_train.utils.replay_utils import replay_padding_row from skyrl.backends.skyrl_train.utils.routed_experts import ( + ROUTED_EXPERT_DTYPES, RoutedExpertIndices, - compact_routed_expert_indices, ) +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_DTYPES, + SAMPLE_SUPPORT_TORCH_DTYPE, + SampleSupport, +) +from skyrl.train.dataset.parallel_fill import fill_batch_rows logger = logging.getLogger(__name__) +# Torch counterparts of the canonical routed-expert dtypes. +ROUTED_EXPERT_TORCH_DTYPES: Dict[np.dtype, torch.dtype] = { + np.dtype(np.uint8): torch.uint8, + np.dtype(np.int16): torch.int16, + np.dtype(np.int32): torch.int32, +} + def make_router_padding_mask( attention_mask: torch.Tensor, @@ -87,6 +105,148 @@ def _reward_to_numpy(custom_reward: Union[List[float], torch.Tensor]) -> np.ndar return reward_arr +def _fill_routed_expert_segment( + packed: torch.Tensor, + cu_seqlens: torch.Tensor, + rollout_expert_indices: List[RoutedExpertIndices], + sample_index: int, +) -> None: + """Write one route segment, using distinct dummy routes for uncaptured trailing tokens.""" + sample_indices = rollout_expert_indices[sample_index] + flags = sample_indices.flags + # torch.from_numpy refuses a non-writeable buffer, and decoded wire routes may be read-only. + if not flags.c_contiguous or not flags.writeable: + sample_indices = sample_indices.copy(order="C") + segment = packed[int(cu_seqlens[sample_index]) : int(cu_seqlens[sample_index + 1])] + captured = sample_indices.shape[0] + segment[:captured] = torch.from_numpy(sample_indices) + segment[captured:] = replay_padding_row(segment.shape[-1], dtype=packed.dtype) + + +def _collate_rollout_expert_indices( + rollout_expert_indices: List[RoutedExpertIndices], + total_real: np.ndarray, +) -> PackedTensor: + """Pack per-trajectory routes into one ``[sum(seq_len_i), layers, topk]`` buffer. + + Entries already have a canonical dtype and are filled from the trainer's local thread pool. + """ + num_samples = len(rollout_expert_indices) + for sample_index, sample_indices in enumerate(rollout_expert_indices): + if not isinstance(sample_indices, np.ndarray): + raise TypeError( + f"rollout_expert_indices entries must be NumPy arrays, got {type(sample_indices).__name__} " + f"at sample {sample_index}" + ) + if sample_indices.dtype not in ROUTED_EXPERT_DTYPES: + supported = ", ".join( + dtype.name for dtype in sorted(ROUTED_EXPERT_DTYPES, key=lambda dtype: dtype.itemsize) + ) + raise ValueError( + f"rollout_expert_indices entries must use a canonical routed-expert dtype ({supported}), " + f"got {sample_indices.dtype} at sample {sample_index}" + ) + + first_shape = rollout_expert_indices[0].shape + if len(first_shape) != 3 or first_shape[0] == 0: + raise ValueError("rollout_expert_indices must contain routes for every trajectory") + num_layers, topk = first_shape[1:] + if topk < 1: + raise ValueError("rollout_expert_indices must contain at least one expert per layer") + + # Validate serially so an invalid trajectory raises deterministically rather than from a worker. + for sample_index, sample_indices in enumerate(rollout_expert_indices): + if sample_indices.ndim != 3 or sample_indices.shape[1:] != (num_layers, topk): + raise ValueError( + "rollout_expert_indices entries must share [layers, topk], " + f"got shape {sample_indices.shape} at sample {sample_index}" + ) + available = int(total_real[sample_index]) + if sample_indices.shape[0] == 0 or sample_indices.shape[0] > available: + raise ValueError( + f"Trajectory {sample_index} has {sample_indices.shape[0]} route rows for {available} tokens" + ) + + batch_dtype = max((indices.dtype for indices in rollout_expert_indices), key=lambda dtype: dtype.itemsize) + if batch_dtype == np.dtype(np.int32): + logger.warning( + "Collating rollout_expert_indices as int32, which doubles this buffer. No supported expert count " + "needs more than int16, so the inference server is not compacting its routes." + ) + cu_seqlens = cu_seqlens_from_lengths(total_real) + packed = torch.empty( + (int(total_real.sum()), num_layers, topk), + dtype=ROUTED_EXPERT_TORCH_DTYPES[batch_dtype], + ) + fill_batch_rows( + functools.partial(_fill_routed_expert_segment, packed, cu_seqlens, rollout_expert_indices), + num_samples, + ) + return PackedTensor(packed, cu_seqlens) + + +def _fill_sample_support_segment( + packed: torch.Tensor, + cu_seqlens: torch.Tensor, + rollout_sample_support: List[SampleSupport], + sample_index: int, +) -> None: + """Write one trajectory's segment of the packed sample-support buffer.""" + rows = rollout_sample_support[sample_index] + # torch.from_numpy refuses a non-writeable buffer, and decoded wire support may be read-only. + if not rows.flags.c_contiguous or not rows.flags.writeable: + rows = rows.copy(order="C") + packed[int(cu_seqlens[sample_index]) : int(cu_seqlens[sample_index + 1])] = torch.from_numpy(rows) + + +def build_sample_support( + rollout_sample_support: List[SampleSupport], + response_lens: np.ndarray, +) -> PackedTensor: + """Pack one response-token support segment per trajectory.""" + num_samples = len(rollout_sample_support) + for sample_index, rows in enumerate(rollout_sample_support): + if not isinstance(rows, np.ndarray): + raise TypeError( + f"rollout_sample_support entries must be NumPy arrays, got {type(rows).__name__} " + f"at sample {sample_index}" + ) + if rows.dtype not in SAMPLE_SUPPORT_DTYPES: + supported = ", ".join(dtype.name for dtype in SAMPLE_SUPPORT_DTYPES) + raise ValueError( + f"rollout_sample_support entries must use a canonical sample-support dtype ({supported}), " + f"got {rows.dtype} at sample {sample_index}" + ) + + first_shape = rollout_sample_support[0].shape + if len(first_shape) != 2 or first_shape[1] < 1: + raise ValueError( + f"rollout_sample_support must be [response_tokens, top_k] arrays, got shape {first_shape} at sample 0" + ) + top_k = first_shape[1] + + # Validate serially so an invalid trajectory raises deterministically rather than from a worker. + for sample_index, rows in enumerate(rollout_sample_support): + if rows.ndim != 2 or rows.shape[1] != top_k: + raise ValueError( + f"rollout_sample_support entries must share top_k {top_k}, " + f"got shape {rows.shape} at sample {sample_index}" + ) + expected = int(response_lens[sample_index]) + if rows.shape[0] != expected: + raise ValueError( + f"Trajectory {sample_index} has {rows.shape[0]} support rows for {expected} response tokens" + ) + + cu_seqlens = cu_seqlens_from_lengths(response_lens) + packed = torch.empty((int(response_lens.sum()), top_k), dtype=SAMPLE_SUPPORT_TORCH_DTYPE) + fill_batch_rows( + functools.partial(_fill_sample_support_segment, packed, cu_seqlens, rollout_sample_support), + num_samples, + ) + return PackedTensor(packed, cu_seqlens) + + def convert_prompts_responses_to_batch_tensors( pad_token_id: int, prompts: List[List[int]], @@ -95,6 +255,7 @@ def convert_prompts_responses_to_batch_tensors( loss_masks: List[List[int]], logprobs: Optional[List[List[float]]] = None, rollout_expert_indices: Optional[List[RoutedExpertIndices]] = None, + rollout_sample_support: Optional[List[SampleSupport]] = None, max_seq_len: Optional[int] = None, ) -> Tuple[ Float[torch.Tensor, "batch seq_len"], @@ -103,7 +264,8 @@ def convert_prompts_responses_to_batch_tensors( Float[torch.Tensor, "batch response_len"], Float[torch.Tensor, "batch response_len"], Optional[Float[torch.Tensor, "batch response_len"]], - Optional[Integer[torch.Tensor, "batch seq_len layer_num topk"]], + Optional[PackedTensor], + Optional[PackedTensor], ]: """ Convert prompts and responses to batch tensors for training. @@ -160,6 +322,12 @@ def convert_prompts_responses_to_batch_tensors( rewards: ``(batch, max_response)`` — right-aligned. loss_masks: ``(batch, max_response)`` — right-aligned. logprobs: ``(batch, max_response)`` — right-aligned, or ``None``. + rollout_expert_indices: ``PackedTensor`` whose values are + ``(sum(prompt_i + response_i), layers, topk)`` in canonical batch order, with + ``cu_seqlens`` naming each trajectory's segment, or ``None``. + rollout_sample_support: ``PackedTensor`` whose values are + ``(sum(response_i), top_k)`` in canonical batch order, with ``cu_seqlens`` naming + each trajectory's segment, or ``None``. """ _verify_inputs(prompts, responses, rewards, loss_masks) @@ -235,42 +403,16 @@ def convert_prompts_responses_to_batch_tensors( if len(rollout_expert_indices) != num_samples: raise ValueError("rollout_expert_indices must contain routes for every trajectory") - canonical_indices = [] - for sample_index, sample_indices in enumerate(rollout_expert_indices): - if not isinstance(sample_indices, np.ndarray): - raise TypeError( - f"rollout_expert_indices entries must be NumPy arrays, got {type(sample_indices).__name__} " - f"at sample {sample_index}" - ) - canonical_indices.append(compact_routed_expert_indices(sample_indices)) - - first_shape = canonical_indices[0].shape - if len(first_shape) != 3 or first_shape[0] == 0: - raise ValueError("rollout_expert_indices must contain routes for every trajectory") - num_layers, topk = first_shape[1:] - if topk < 1: - raise ValueError("rollout_expert_indices must contain at least one expert per layer") - - batch_dtype = max((indices.dtype for indices in canonical_indices), key=lambda dtype: dtype.itemsize) - padded = make_replay_padding_indices_np( - (num_samples, max_total, num_layers, topk), - dtype=batch_dtype, - ) - for sample_index, sample_indices in enumerate(canonical_indices): - if sample_indices.ndim != 3 or sample_indices.shape[1:] != (num_layers, topk): - raise ValueError( - "rollout_expert_indices entries must share [layers, topk], " - f"got shape {sample_indices.shape} at sample {sample_index}" - ) - left_pad = max_total - (prompt_token_lens[sample_index] + response_token_lens[sample_index]) - available = max_total - left_pad - if sample_indices.shape[0] == 0 or sample_indices.shape[0] > available: - raise ValueError( - f"Trajectory {sample_index} has {sample_indices.shape[0]} route rows for {available} tokens" - ) - route_end = left_pad + sample_indices.shape[0] - padded[sample_index, left_pad:route_end] = sample_indices - rollout_expert_indices_tensor = torch.from_numpy(padded) + rollout_expert_indices_tensor = _collate_rollout_expert_indices(rollout_expert_indices, total_real) + + sample_support_tensor = None + if rollout_sample_support is not None: + if not isinstance(rollout_sample_support, list): + raise TypeError("rollout_sample_support must be a list of NumPy arrays") + if len(rollout_sample_support) != num_samples: + raise ValueError("rollout_sample_support must contain support for every trajectory") + + sample_support_tensor = build_sample_support(rollout_sample_support, response_lens) return ( sequences, @@ -280,6 +422,7 @@ def convert_prompts_responses_to_batch_tensors( ret_loss_masks, logprobs_tensor, rollout_expert_indices_tensor, + sample_support_tensor, ) diff --git a/skyrl/train/dataset/replay_buffer.py b/skyrl/train/dataset/replay_buffer.py index 6f627e31ed..efada7b4dd 100644 --- a/skyrl/train/dataset/replay_buffer.py +++ b/skyrl/train/dataset/replay_buffer.py @@ -15,23 +15,24 @@ from jaxtyping import Bool, Float, Integer from skyrl.backends.skyrl_train.training_batch import TensorList +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor BasicType = Union[int, float, str, bool] -def to(tensor: Union[torch.Tensor, List[torch.Tensor], BasicType], device): +def to(tensor: Union[torch.Tensor, PackedTensor, List[torch.Tensor], BasicType], device): if isinstance(tensor, list): return [to(t, device) for t in tensor] - elif isinstance(tensor, torch.Tensor): + elif isinstance(tensor, (torch.Tensor, PackedTensor)): return tensor.to(device) else: return tensor -def pin_memory(tensor: Union[torch.Tensor, List[torch.Tensor], BasicType]): +def pin_memory(tensor: Union[torch.Tensor, PackedTensor, List[torch.Tensor], BasicType]): if isinstance(tensor, list): return [pin_memory(t) for t in tensor] - elif isinstance(tensor, torch.Tensor): + elif isinstance(tensor, (torch.Tensor, PackedTensor)): return tensor.pin_memory() else: return tensor @@ -67,10 +68,13 @@ class Experience: loss_mask: Optional[Integer[torch.LongTensor, "batch response_len"]] response_mask: Optional[Integer[torch.Tensor, "batch response_len"]] rollout_logprobs: Optional[Float[torch.Tensor, "batch response_len"]] - rollout_expert_indices: Optional[Integer[torch.Tensor, "batch seq_len layer_num topk"]] + # Routes packed to real tokens: values [sum(seq_len_i), layer_num, topk] + cu_seqlens. + rollout_expert_indices: Optional[PackedTensor] num_actions: int info: Optional[dict] router_padding_mask: Optional[Bool[torch.Tensor, "batch seq_len"]] = None + # Sampler support packed to response tokens: values [sum(response_len_i), top_k] + cu_seqlens. + rollout_sample_support: Optional[PackedTensor] = None kl: Optional[Float[torch.Tensor, "batch response_len"]] = None metadata: Optional[Dict[str, Any]] = None pixel_values: Optional[TensorList] = None @@ -104,6 +108,8 @@ def to_device(self, device: torch.device) -> None: self.rollout_expert_indices = to(self.rollout_expert_indices, device) if self.router_padding_mask is not None: self.router_padding_mask = to(self.router_padding_mask, device) + if self.rollout_sample_support is not None: + self.rollout_sample_support = to(self.rollout_sample_support, device) if self.pixel_values is not None: self.pixel_values = self.pixel_values.to(device) if self.image_grid_thw is not None: @@ -135,6 +141,8 @@ def pin_memory(self): self.rollout_expert_indices = self.rollout_expert_indices.pin_memory() if self.router_padding_mask is not None: self.router_padding_mask = self.router_padding_mask.pin_memory() + if self.rollout_sample_support is not None: + self.rollout_sample_support = self.rollout_sample_support.pin_memory() return self diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py index 81792e0a2a..c9724d8805 100644 --- a/skyrl/train/generators/base.py +++ b/skyrl/train/generators/base.py @@ -6,8 +6,11 @@ from skyrl.backends.skyrl_train.inference_servers.base import ConversationType from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices +from skyrl.backends.skyrl_train.utils.sample_support import SampleSupport TrainingPhase = Literal["train", "eval"] +TRAINING_PHASE_TRAIN: TrainingPhase = "train" +TRAINING_PHASE_EVAL: TrainingPhase = "eval" @dataclass @@ -51,7 +54,10 @@ class GeneratorOutput(TypedDict): # e.g. {"llm": [...], "env": [...]}. trajectory_time_splits is None if any trajectory did not # record its split. trajectory_time_splits: Optional[Dict[str, List[float]]] + # Per trajectory, routes for a prefix of its prompt and response tokens. rollout_expert_indices: Optional[List[RoutedExpertIndices]] + # Per trajectory, sampler support for each response token; uncaptured rows are padding. + rollout_sample_support: Optional[List[SampleSupport]] # Applicable only for step-wise training is_last_step: Optional[List[bool]] # Per-row env metrics (one dict per row in the flattened batch). Used by diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py index 26b1102e6c..e740857445 100644 --- a/skyrl/train/generators/skyrl_gym_generator.py +++ b/skyrl/train/generators/skyrl_gym_generator.py @@ -13,6 +13,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union from uuid import uuid4 +import numpy as np import torch from loguru import logger from tqdm.asyncio import tqdm @@ -23,12 +24,24 @@ InferenceEngineInput, InferenceEngineInterface, ) -from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices +from skyrl.backends.skyrl_train.utils.routed_experts import ( + RoutedExpertIndices, + RoutedExpertTrace, +) +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_DTYPE, + SAMPLE_SUPPORT_PADDING, + SampleSupport, + SampleSupportTrace, +) from skyrl.train.config import GeneratorConfig, SkyRLGymConfig from skyrl.train.generators.base import ( + TRAINING_PHASE_EVAL, + TRAINING_PHASE_TRAIN, GeneratorInput, GeneratorInterface, GeneratorOutput, + TrainingPhase, TrajectoryID, ) from skyrl.train.generators.utils import ( @@ -52,6 +65,7 @@ class TrajectoryOutput: rollout_logprobs: Optional[List[float]] env_metrics: Dict[str, Any] rollout_expert_indices: Optional[RoutedExpertIndices] = None + rollout_sample_support: Optional[SampleSupport] = None pixel_values: Optional[torch.Tensor] = None image_grid_thw: Optional[torch.Tensor] = None # End-to-end wall-clock time (seconds) to generate this trajectory. Optional: agent loops may @@ -83,7 +97,10 @@ class AgentLoopState: rollout_logprobs: Optional[List[float]] response_end_idx: Optional[int] done: bool - rollout_expert_indices: Optional[RoutedExpertIndices] = None + routed_expert_trace: Optional[RoutedExpertTrace] = None + sample_support_trace: Optional[SampleSupportTrace] = None + # Support for an EOS sliced from a single-turn response. + dropped_eos_sample_support: Optional[SampleSupport] = None @dataclass @@ -93,13 +110,23 @@ class TurnOutput: output_logprobs: Optional[List[float]] new_obs: ConversationType obs_ids: List[int] - rollout_expert_indices: Optional[RoutedExpertIndices] reward: Optional[float] + rollout_sample_support: Optional[SampleSupport] = None added_eos: bool = False - def get_turn_rollout_expert_indices(self) -> Optional[RoutedExpertIndices]: - """Return only routes that the inference model actually executed.""" - return self.rollout_expert_indices + def get_turn_rollout_sample_support(self) -> Optional[SampleSupport]: + """Return sample support padded over synthetic EOS and observation tokens.""" + if self.rollout_sample_support is None: + return None + padding_count = int(self.added_eos) + len(self.obs_ids) + if not padding_count: + return self.rollout_sample_support + padding = np.full( + (padding_count, self.rollout_sample_support.shape[1]), + SAMPLE_SUPPORT_PADDING, + dtype=self.rollout_sample_support.dtype, + ) + return np.concatenate((self.rollout_sample_support, padding), axis=0) def get_turn_loss_mask(self) -> List[int]: """ @@ -217,12 +244,44 @@ def _validate_cfg(self, generator_cfg: GeneratorConfig): f"`step_wise_trajectories` doesn't support custom chat template, got {generator_cfg.chat_template}" ) - if self.generator_cfg.inference_engine.enable_return_routed_experts: - raise ValueError("`step_wise_trajectories` doesn't support `enable_return_routed_experts=True`") - if not self.use_conversation_multi_turn: raise ValueError("`step_wise_trajectories` doesn't support `use_conversation_multi_turn=False`") + if generator_cfg.inference_engine.enable_return_routed_experts: + raise ValueError( + "`step_wise_trajectories` doesn't support " + "`generator.inference_engine.enable_return_routed_experts=True`. A step's routes are " + "recorded for its generated tokens only, while its row's prompt is the whole history so " + "far, so they would replay onto the first N prompt tokens of the row with no length " + "mismatch to assert on." + ) + + ie_cfg = generator_cfg.inference_engine + + if ie_cfg.enable_return_routed_experts and not self.use_conversation_multi_turn: + raise ValueError( + "`generator.inference_engine.enable_return_routed_experts=True` requires " + "`generator.use_conversation_multi_turn=True`. With `use_conversation_multi_turn=False` the " + "agent loop appends a synthetic EOS that is loss-active but that the inference engine never " + "evaluated, so the routed-expert trace holds no row for it and refuses to dummy-pad a " + "loss-active target." + ) + + if self.custom_chat_template is not None: + if ie_cfg.enable_return_routed_experts: + raise ValueError( + "`generator.inference_engine.enable_return_routed_experts=True` is not compatible with a " + f"custom chat template, got {generator_cfg.chat_template}. Retokenizing the chat history " + "breaks token-in-token-out, so per-token routes no longer align with the response tokens." + ) + if ie_cfg.enable_return_sample_support_set: + raise ValueError( + "`generator.inference_engine.enable_return_sample_support_set=True` is not compatible with a " + f"custom chat template, got {generator_cfg.chat_template}. Retokenizing the chat history " + "breaks token-in-token-out, so per-token support rows no longer align with the response " + "tokens." + ) + async def _run_in_executor_if_available(self, func, *args, **kwargs): if (executor := self.env_executor) is not None: loop = asyncio.get_running_loop() @@ -287,6 +346,7 @@ async def agent_loop( sampling_params: Optional[Dict[str, Any]] = None, trajectory_id: Optional[TrajectoryID] = None, cache_salt: Optional[str] = None, + training_phase: TrainingPhase = TRAINING_PHASE_TRAIN, ) -> Union[TrajectoryOutput, StepWiseOutput]: """ Multi-turn generation loop that executes a single trajectory. @@ -363,9 +423,21 @@ async def agent_loop( current_sampling_params: dict = ( sampling_params if sampling_params is not None else asdict(self.generator_cfg.sampling_params) ) + # Eval uses unbounded top-k sampling and opts out of capture. + capture_sample_support = ( + self.generator_cfg.inference_engine.enable_return_sample_support_set + and training_phase != TRAINING_PHASE_EVAL + ) + sample_support_width = current_sampling_params["top_k"] if capture_sample_support else 0 + # Eval trajectories do not consume routed-expert indices. + capture_routed_experts = ( + self.generator_cfg.inference_engine.enable_return_routed_experts + and training_phase != TRAINING_PHASE_EVAL + ) # Accumulate per-step rewards. Format: (reward, response_end_token_idx) per_step_rewards: List[Tuple[float, Optional[int]]] = [] + final_observation_token_count = 0 is_step_wise = self.generator_cfg.step_wise_trajectories @@ -379,6 +451,8 @@ async def agent_loop( rollout_logprobs=[] if get_logprobs else None, response_end_idx=None, done=False, + routed_expert_trace=RoutedExpertTrace() if capture_routed_experts else None, + sample_support_trace=SampleSupportTrace() if capture_sample_support and not is_step_wise else None, ) while not agent_loop_state.done: @@ -401,11 +475,14 @@ async def agent_loop( agent_loop_state.loss_mask = [] agent_loop_state.rollout_logprobs = None + routed_expert_trace = agent_loop_state.routed_expert_trace engine_input = InferenceEngineInput( prompt_token_ids=[agent_loop_state.input_ids], session_ids=[session_id], sampling_params=sampling_params, cache_salt=cache_salt, + routed_experts_prompt_starts=[routed_expert_trace.prompt_start] if routed_expert_trace else None, + return_sample_support=capture_sample_support, ) llm_call_start_time = time.monotonic() engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name) @@ -422,9 +499,34 @@ async def agent_loop( if rollout_expert_indices is not None: rollout_expert_indices = rollout_expert_indices[0] - if self.custom_chat_template is not None: + assert ( + self.custom_chat_template is None + ), "Rollout expert indices bookkeeping is not supported with custom chat template" + if routed_expert_trace is not None: + if rollout_expert_indices is None: + raise ValueError("R3 generation did not return routed expert indices") + routed_expert_trace.record_generation( + prompt_token_count=len(agent_loop_state.input_ids), + generated_token_count=len(output_ids), + routed_experts=rollout_expert_indices, + ) + + sample_support_rows = None + if capture_sample_support: + assert ( + self.custom_chat_template is None + ), "Sample-support bookkeeping is not supported with custom chat template" + raw_sample_support = engine_output.get("rollout_sample_support", None) + if raw_sample_support is None: + raise ValueError("Sample-support generation did not return a support set") + sample_support_rows = np.asarray( + raw_sample_support[0], + dtype=SAMPLE_SUPPORT_DTYPE, + order="C", + ).reshape(-1, sample_support_width) + if sample_support_rows.shape[0] != len(output_ids): raise ValueError( - "Rollout expert indices bookkeeping is not supported with custom chat template" + f"Sample support has {sample_support_rows.shape[0]} rows for {len(output_ids)} tokens" ) # Append eos when sampling_params.stop is not None. Does not affect 3.a as chat templates add eos_token. # sampling_params is not None for eval, but None for training (which uses engine.sampling_params which are from cfg) @@ -459,6 +561,10 @@ async def agent_loop( ) output = env_step_output["postprocessed_action"] output_ids = self.tokenizer.encode(output, add_special_tokens=False) + if routed_expert_trace is not None: + raise ValueError("R3 bookkeeping is incompatible with postprocessed_action") + if sample_support_rows is not None: + raise ValueError("Sample-support bookkeeping is incompatible with postprocessed_action") obs_ids = self.get_obs_ids_from_obs(new_obs, agent_loop_state.done) @@ -470,8 +576,8 @@ async def agent_loop( new_obs=new_obs, reward=step_reward, obs_ids=obs_ids, + rollout_sample_support=sample_support_rows, added_eos=added_eos, - rollout_expert_indices=rollout_expert_indices, ) if is_step_wise: @@ -482,6 +588,7 @@ async def agent_loop( # agent loop only tracks loss mask and rollout logprobs for this turn with step_wise training turn_loss_mask = turn_output.get_turn_loss_mask() turn_response_logprobs: Optional[List[float]] = turn_output.get_turn_rollout_logprobs() + turn_sample_support = turn_output.get_turn_rollout_sample_support() per_step_output = TrajectoryOutput( response_ids=turn_response_ids, @@ -491,12 +598,15 @@ async def agent_loop( rollout_logprobs=turn_response_logprobs, stop_reason=stop_reason, env_metrics=env.get_metrics() if agent_loop_state.done else {}, - rollout_expert_indices=turn_output.get_turn_rollout_expert_indices(), + rollout_sample_support=turn_sample_support, ) agent_loop_output.step_outputs.append(per_step_output) # 3. Update states: input ids, loss_mask, chat_history, etc. # Three ways of managing input + sample_support_trace = agent_loop_state.sample_support_trace + support_rows_before = sample_support_trace.num_rows if sample_support_trace is not None else 0 + input_length_before = len(agent_loop_state.input_ids) if retokenize_chat_history: # a. custom chat template agent_loop_state = self._update_agent_state_by_retokenizing_chat_history( @@ -513,6 +623,17 @@ async def agent_loop( agent_loop_state, turn_output ) + if sample_support_trace is not None: + # Each appended token must contribute one support row. + support_rows_added = sample_support_trace.num_rows - support_rows_before + tokens_added = len(agent_loop_state.input_ids) - input_length_before + assert support_rows_added == tokens_added, ( + f"sample-support trace advanced {support_rows_added} rows for {tokens_added} tokens " + f"appended by this turn" + ) + # The trace covers this observation, which the response never keeps. + final_observation_token_count = len(turn_output.obs_ids) + per_step_rewards.append((step_reward, agent_loop_state.response_end_idx)) # Get environment-specific metrics after the episode is done @@ -523,6 +644,7 @@ async def agent_loop( prompt_ids = agent_loop_state.input_ids[:initial_prompt_length] rollout_logprobs = None rollout_expert_indices_out = None + rollout_sample_support_out = None response_ids = None # Prepare the final loss_mask, response_ids and rollout_logprobs . @@ -553,10 +675,6 @@ async def agent_loop( rollout_logprobs = agent_loop_state.rollout_logprobs[ : agent_loop_state.response_end_idx - initial_prompt_length + 1 ] - if agent_loop_state.rollout_expert_indices is not None: - rollout_expert_indices_out = agent_loop_state.rollout_expert_indices[ - : agent_loop_state.response_end_idx + 1 - ] # fix index for per_step_rewards per_step_rewards = [(reward, idx - initial_prompt_length) for reward, idx in per_step_rewards] assert len(loss_mask) == len( @@ -571,8 +689,26 @@ async def agent_loop( loss_mask.append(1) if rollout_logprobs is not None: rollout_logprobs.append(0.0) + if agent_loop_state.sample_support_trace is not None: + # Restore support for a sampled EOS; a synthetic EOS remains padding. + dropped_row = agent_loop_state.dropped_eos_sample_support + if dropped_row is not None: + agent_loop_state.sample_support_trace.append(dropped_row, expected_rows=1) + else: + agent_loop_state.sample_support_trace.append_padding(1) appended_eos_token = True + if agent_loop_state.routed_expert_trace is not None and agent_loop_state.routed_expert_trace.prompt_start: + rollout_expert_indices_out = agent_loop_state.routed_expert_trace.finalize( + token_count=len(prompt_ids) + len(response_ids), + loss_mask=[0] * len(prompt_ids) + loss_mask, + ) + if agent_loop_state.sample_support_trace is not None and agent_loop_state.sample_support_trace.num_rows: + rollout_sample_support_out = agent_loop_state.sample_support_trace.finalize( + token_count=len(response_ids), + extra_rows=final_observation_token_count, + ) + if self.generator_cfg.step_wise_trajectories: for per_step_output, (reward, resp_end_idx) in zip(agent_loop_output.step_outputs, per_step_rewards): per_token_reward = [0.0] * len(per_step_output.response_ids) @@ -591,6 +727,7 @@ async def agent_loop( rollout_logprobs=rollout_logprobs, env_metrics=env_metrics, rollout_expert_indices=rollout_expert_indices_out, + rollout_sample_support=rollout_sample_support_out, ) agent_loop_output = self._post_process_agent_loop_output( @@ -718,6 +855,7 @@ async def generate_batched( max_tokens: int, sampling_params: Optional[Dict[str, Any]] = None, cache_salt: Optional[str] = None, + training_phase: TrainingPhase = TRAINING_PHASE_TRAIN, ) -> GeneratorOutput: """ Single-turn batched generation (can use the synchronous offline engine) @@ -749,15 +887,29 @@ async def generate_batched( tokenize=True, return_dict=False, ) + # Eval batches do not capture per-token side channels. + capture_sample_support = ( + self.generator_cfg.inference_engine.enable_return_sample_support_set + and training_phase != TRAINING_PHASE_EVAL + ) + capture_routed_experts = ( + self.generator_cfg.inference_engine.enable_return_routed_experts and training_phase != TRAINING_PHASE_EVAL + ) engine_input = InferenceEngineInput( - prompt_token_ids=prompt_token_ids, sampling_params=sampling_params, cache_salt=cache_salt + prompt_token_ids=prompt_token_ids, + sampling_params=sampling_params, + return_sample_support=capture_sample_support, + cache_salt=cache_salt, ) engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name) outputs = engine_output["responses"] responses = engine_output["response_ids"] stop_reasons = engine_output["stop_reasons"] logprobs = engine_output.get("response_logprobs", None) - raw_rollout_expert_indices = engine_output.get("rollout_expert_indices", None) + raw_rollout_expert_indices = ( + engine_output.get("rollout_expert_indices", None) if capture_routed_experts else None + ) + raw_rollout_sample_support = engine_output.get("rollout_sample_support", None) truncated_responses = [] rewards = [] @@ -765,6 +917,7 @@ async def generate_batched( env_metrics = [] truncated_logprobs: Optional[List[List[float]]] = [] if logprobs is not None else None truncated_indices: Optional[List[RoutedExpertIndices]] = [] if raw_rollout_expert_indices is not None else None + truncated_sample_support: Optional[List[SampleSupport]] = [] if raw_rollout_sample_support is not None else None for i, (output, response, env, env_class) in enumerate(zip(outputs, responses, envs, env_classes)): # step on environment and compute reward @@ -783,6 +936,8 @@ async def generate_batched( sample_indices = raw_rollout_expert_indices[i] prompt_len = len(prompt_token_ids[i]) truncated_indices.append(sample_indices[: prompt_len + len(response)]) + if raw_rollout_sample_support is not None: + truncated_sample_support.append(raw_rollout_sample_support[i][: len(response)]) # Get environment-specific metrics env_metrics.append(env.get_metrics()) @@ -804,6 +959,7 @@ async def generate_batched( "rollout_metrics": rollout_metrics, "rollout_logprobs": truncated_logprobs, "rollout_expert_indices": truncated_indices, + "rollout_sample_support": truncated_sample_support, } return generator_output @@ -833,9 +989,21 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False # every trajectory in this batch shares one salt (the policy version at the start of the batch). cache_salt = self._compute_cache_salt() + # The phase controls per-request sample-support capture. + batch_metadata = input_batch.get("batch_metadata", None) + training_phase: TrainingPhase = ( + batch_metadata.training_phase if batch_metadata is not None else TRAINING_PHASE_TRAIN + ) + if self.batched: return await self.generate_batched( - prompts, env_classes, env_extras, max_tokens, sampling_params, cache_salt=cache_salt + prompts, + env_classes, + env_extras, + max_tokens, + sampling_params, + cache_salt=cache_salt, + training_phase=training_phase, ) # Async agent loop to generate trajectories in parallel. @@ -851,6 +1019,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False sampling_params=sampling_params, trajectory_id=trajectory_ids[i] if trajectory_ids is not None else None, cache_salt=cache_salt, + training_phase=training_phase, ) ) @@ -938,10 +1107,24 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False else: rollout_logprobs = None - if self.generator_cfg.inference_engine.enable_return_routed_experts: - rollout_expert_indices = [output.rollout_expert_indices for output in all_outputs] + if self.generator_cfg.step_wise_trajectories: + expert_indices_values = [None] * len(responses) else: - rollout_expert_indices = None + expert_indices_values = [output.rollout_expert_indices for output in all_outputs] + # Preserve None when no trajectory contains routes. + rollout_expert_indices = ( + expert_indices_values if any(value is not None for value in expert_indices_values) else None + ) + + if self.generator_cfg.step_wise_trajectories: + sample_support_values = [ + step_output.rollout_sample_support for output in all_outputs for step_output in output.step_outputs + ] + else: + sample_support_values = [output.rollout_sample_support for output in all_outputs] + rollout_sample_support = ( + sample_support_values if any(value is not None for value in sample_support_values) else None + ) rollout_metrics = get_rollout_metrics( responses, @@ -976,6 +1159,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False "trajectory_generation_times": out_trajectory_generation_times, "trajectory_time_splits": out_trajectory_time_splits, "rollout_expert_indices": rollout_expert_indices, + "rollout_sample_support": rollout_sample_support, "is_last_step": is_last_step, "env_metrics": env_metrics, } @@ -1037,6 +1221,13 @@ def _update_agent_state_by_retokenizing_chat_history( """ assert self.use_conversation_multi_turn and self.custom_chat_template + if agent_loop_state.routed_expert_trace is not None or agent_loop_state.sample_support_trace is not None: + raise NotImplementedError( + "retokenizing the chat history does not feed the per-token side-channel traces, so routes and " + "sample support would not align with the retokenized response. `generator.chat_template` is " + "refused with `enable_return_routed_experts` and `enable_return_sample_support_set`." + ) + agent_loop_state.chat_history = self._update_chat_history( agent_loop_state.chat_history, turn_output.output, turn_output.new_obs ) @@ -1047,8 +1238,6 @@ def _update_agent_state_by_retokenizing_chat_history( agent_loop_state.response_end_idx = None # `logprobs` are not computed because retokenizing breaks token-in-token-out agent_loop_state.rollout_logprobs = None - # indices are not meaningful when retokenizing - agent_loop_state.rollout_expert_indices = None return agent_loop_state def _update_agent_loop_state_with_multiturn_chat_template( @@ -1100,17 +1289,12 @@ def _update_agent_loop_state_with_multiturn_chat_template( loss_mask_for_turn = turn_output.get_turn_loss_mask() rollout_logprobs_for_turn = turn_output.get_turn_rollout_logprobs() - # use the raw rollout expert indices without any appending of observation tokens - # this will be overwritten each turn, so we don't need to append observation tokens to it - rollout_expert_indices_for_turn = turn_output.rollout_expert_indices - if self.generator_cfg.step_wise_trajectories: # cumulative input_ids is not tracked for step wise training agent_loop_state.response_end_idx = len(turn_output.output_ids) - 1 - # no running loss_mask, `rollout_logprobs`, or `rollout_expert_indices` are tracked for step-wise training + # no running loss_mask or rollout logprobs are tracked for step-wise training agent_loop_state.loss_mask = None agent_loop_state.rollout_logprobs = None - agent_loop_state.rollout_expert_indices = None else: # Directly append turn output turn_ids = turn_output.output_ids + turn_output.obs_ids @@ -1119,11 +1303,9 @@ def _update_agent_loop_state_with_multiturn_chat_template( agent_loop_state.loss_mask += loss_mask_for_turn if agent_loop_state.rollout_logprobs is not None and rollout_logprobs_for_turn is not None: agent_loop_state.rollout_logprobs += rollout_logprobs_for_turn - if rollout_expert_indices_for_turn is not None: - # overwrite the existing rollout inference indices, since the inference engine should - # return the expert indices for the entire sequence including each turn's input - # and the final response should not have an observation appended to it - agent_loop_state.rollout_expert_indices = rollout_expert_indices_for_turn + turn_sample_support = turn_output.get_turn_rollout_sample_support() + if agent_loop_state.sample_support_trace is not None and turn_sample_support is not None: + agent_loop_state.sample_support_trace.append(turn_sample_support, expected_rows=len(turn_ids)) return agent_loop_state @@ -1175,8 +1357,15 @@ def _update_agent_loop_state_with_singleturn_chat_template( # Remove EOS token from response tokens since we are continuing the current assistant message new_resp_tokens = turn_output.output_ids.copy() - if new_resp_tokens and new_resp_tokens[-1] == self.tokenizer.eos_token_id: + dropped_eos = bool(new_resp_tokens) and new_resp_tokens[-1] == self.tokenizer.eos_token_id + if dropped_eos: new_resp_tokens = new_resp_tokens[:-1] + # Keep the sliced EOS token's support row so the trajectory's re-appended EOS can carry it. + agent_loop_state.dropped_eos_sample_support = ( + turn_output.rollout_sample_support[len(new_resp_tokens) : len(new_resp_tokens) + 1] + if dropped_eos and turn_output.rollout_sample_support is not None + else None + ) turn_ids = new_resp_tokens + obs_ids_to_add loss_mask_for_turn = [1] * len(new_resp_tokens) + [0] * len(obs_ids_to_add) @@ -1194,13 +1383,10 @@ def _update_agent_loop_state_with_singleturn_chat_template( agent_loop_state.loss_mask += loss_mask_for_turn if agent_loop_state.rollout_logprobs is not None and rollout_logprobs_for_turn is not None: agent_loop_state.rollout_logprobs += rollout_logprobs_for_turn - if ( - self.generator_cfg.inference_engine.enable_return_routed_experts - and turn_output.rollout_expert_indices is not None - ): - # overwrite the existing rollout inference indices, since the inference engine should - # return the expert indices for the entire sequence including each turn's input and observation tokens - # and the final response should not have an observation appended to it - agent_loop_state.rollout_expert_indices = turn_output.rollout_expert_indices - + if agent_loop_state.sample_support_trace is not None and turn_output.rollout_sample_support is not None: + # A dropped EOS shortens the generated run, so slice rather than reuse the turn's padding. + agent_loop_state.sample_support_trace.append( + turn_output.rollout_sample_support[: len(new_resp_tokens)], expected_rows=len(new_resp_tokens) + ) + agent_loop_state.sample_support_trace.append_padding(len(obs_ids_to_add)) return agent_loop_state diff --git a/skyrl/train/generators/skyrl_vlm_generator.py b/skyrl/train/generators/skyrl_vlm_generator.py index 3793591507..7053bd1a5a 100644 --- a/skyrl/train/generators/skyrl_vlm_generator.py +++ b/skyrl/train/generators/skyrl_vlm_generator.py @@ -21,7 +21,12 @@ RemoteInferenceClient, ) from skyrl.train.config import GeneratorConfig, SkyRLGymConfig -from skyrl.train.generators.base import GeneratorOutput, TrajectoryID +from skyrl.train.generators.base import ( + TRAINING_PHASE_TRAIN, + GeneratorOutput, + TrainingPhase, + TrajectoryID, +) from skyrl.train.generators.skyrl_gym_generator import ( SkyRLGymGenerator, TrajectoryOutput, @@ -57,6 +62,7 @@ def _validate_cfg(self, generator_cfg: GeneratorConfig): "SkyRLVLMGymGenerator requires `use_conversation_multi_turn=True` " "because multi-modal observations must be in separate user messages." ) + super()._validate_cfg(generator_cfg) async def _render_conversation(self, conversation: ConversationType) -> RenderedConversation: rendered = await self.inference_engine_client.render_chat_completion( @@ -74,6 +80,7 @@ async def agent_loop( sampling_params: Optional[Dict[str, Any]] = None, trajectory_id: Optional[TrajectoryID] = None, cache_salt: Optional[str] = None, + training_phase: TrainingPhase = TRAINING_PHASE_TRAIN, ) -> TrajectoryOutput: """Multi-turn VLM generation loop for a single trajectory. The conversation is treated as the source of truth and re-tokenized each step. diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py index c84ea8a5aa..16d1cd2f07 100644 --- a/skyrl/train/generators/utils.py +++ b/skyrl/train/generators/utils.py @@ -8,6 +8,11 @@ from loguru import logger from skyrl.backends.skyrl_train.inference_servers.base import ConversationType +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_DTYPE, + SAMPLE_SUPPORT_PADDING, + SampleSupport, +) from skyrl.train.config import ChatTemplateConfig from skyrl.train.generators.base import ( BatchMetadata, @@ -278,11 +283,14 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step (e.g. `is_last_step`, `trajectory_ids`, contiguous trajectory ordering). """ assert len(generator_outputs) > 0 - has_rollout_logprobs = [output.get("rollout_logprobs") is not None for output in generator_outputs] - if any(has_rollout_logprobs) and not all(has_rollout_logprobs): - raise ValueError( - "generator outputs are expected to all have null rollout_logprobs or all non-null, but received a mix" - ) + # Per-token side channels must be populated consistently across batches. + for all_or_nothing_field in ("rollout_logprobs", "rollout_expert_indices", "rollout_sample_support"): + present = [output.get(all_or_nothing_field) is not None for output in generator_outputs] + if any(present) and not all(present): + raise ValueError( + f"generator outputs are expected to all have null {all_or_nothing_field} or all non-null, " + "but received a mix" + ) first = generator_outputs[0] result: GeneratorOutput = { "prompt_token_ids": _flatten_field(generator_outputs, "prompt_token_ids"), @@ -291,6 +299,8 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step "loss_masks": _flatten_field(generator_outputs, "loss_masks"), "stop_reasons": _concat_optional_field(generator_outputs, "stop_reasons"), "rollout_logprobs": _concat_optional_field(generator_outputs, "rollout_logprobs"), + "rollout_expert_indices": _concat_optional_field(generator_outputs, "rollout_expert_indices"), + "rollout_sample_support": _concat_optional_field(generator_outputs, "rollout_sample_support"), "trajectory_generation_times": _concat_optional_field(generator_outputs, "trajectory_generation_times"), "trajectory_time_splits": _concat_optional_field(generator_outputs, "trajectory_time_splits"), } @@ -783,14 +793,9 @@ def _is_prefix(maybe_prefix: List[int], candidate: List[int]) -> bool: def slice_generator_output( generator_output: GeneratorOutput, indices: List[int], *, preserve_metrics: bool = True ) -> GeneratorOutput: - """Slice a GeneratorOutput to keep only the entries at the given indices. - - Generator-specific per-trajectory fields are sliced without naming them here. - Prefix-aware merging passes entries that all share one ``TrajectoryID``; - dynamic sampling may intentionally select entries from different trajectories. - """ + """Slice list and dict-of-list fields at the given indices.""" assert len(indices) > 0, "indices must be non-empty" - # Every key except `rollout_metrics` is either a per-entry list to slice, or None. + # Every key except `rollout_metrics` is None, a dict of per-entry lists, or a per-entry list. sliced: GeneratorOutput = {} for key, value in generator_output.items(): if key == "rollout_metrics": @@ -798,11 +803,18 @@ def slice_generator_output( sliced[key] = value elif value is None: sliced[key] = None + elif isinstance(value, dict): + sliced[key] = {name: [component[i] for i in indices] for name, component in value.items()} else: sliced[key] = [value[i] for i in indices] return sliced +def _concat_sample_support(blocks: List[SampleSupport]) -> SampleSupport: + """Join one trajectory's per-turn support blocks, without copying an unmerged turn.""" + return blocks[0] if len(blocks) == 1 else np.concatenate(blocks, axis=0) + + def _merge_single_trajectory(gen_out: GeneratorOutput) -> GeneratorOutput: """Greedily merge turns of a single trajectory using prefix matching. @@ -822,6 +834,9 @@ def _merge_single_trajectory(gen_out: GeneratorOutput) -> GeneratorOutput: is_token_level_rewards = isinstance(gen_out["rewards"][0], list) has_logprobs = gen_out.get("rollout_logprobs") is not None has_stop_reasons = gen_out.get("stop_reasons") is not None + has_sample_support = gen_out.get("rollout_sample_support") is not None + # Support rows are dense, so an observation delta contributes full-width padding rows. + sample_support_width = gen_out["rollout_sample_support"][0].shape[1] if has_sample_support else 0 # Per-field output accumulators. # Fields that we take from all the entries in the merge group @@ -829,6 +844,8 @@ def _merge_single_trajectory(gen_out: GeneratorOutput) -> GeneratorOutput: out_response_ids: List[List[int]] = [] out_loss_masks: List[List[int]] = [] out_logprobs: Optional[List[List[float]]] = [] if has_logprobs else None + # Keep one block per turn until the merged trajectory is flushed. + out_sample_support: Optional[List[SampleSupport]] = [] if has_sample_support else None # If per-token rewards, we keep appending. If per-turn rewards, we only take from the last turn. out_rewards: list = [] @@ -842,16 +859,21 @@ def _merge_single_trajectory(gen_out: GeneratorOutput) -> GeneratorOutput: acc_response: List[int] = list(gen_out["response_ids"][0]) acc_loss_mask: List[int] = list(gen_out["loss_masks"][0]) acc_logprobs: Optional[List[float]] = list(gen_out["rollout_logprobs"][0]) if has_logprobs else None + acc_sample_support: Optional[List[SampleSupport]] = ( + [gen_out["rollout_sample_support"][0]] if has_sample_support else None + ) acc_rewards_tokens: Optional[List[float]] = list(gen_out["rewards"][0]) if is_token_level_rewards else None last = 0 def flush(): - nonlocal acc_prompt, acc_response, acc_loss_mask, acc_logprobs, acc_rewards_tokens, last + nonlocal acc_prompt, acc_response, acc_loss_mask, acc_logprobs, acc_sample_support, acc_rewards_tokens, last out_prompt_ids.append(acc_prompt) out_response_ids.append(acc_response) out_loss_masks.append(acc_loss_mask) if has_logprobs: out_logprobs.append(acc_logprobs) + if has_sample_support: + out_sample_support.append(_concat_sample_support(acc_sample_support)) out_rewards.append(acc_rewards_tokens if is_token_level_rewards else gen_out["rewards"][last]) if has_stop_reasons: out_stop_reasons.append(gen_out["stop_reasons"][last]) @@ -869,6 +891,7 @@ def flush(): acc_response = list(gen_out["response_ids"][i]) acc_loss_mask = list(gen_out["loss_masks"][i]) acc_logprobs = list(gen_out["rollout_logprobs"][i]) if has_logprobs else None + acc_sample_support = [gen_out["rollout_sample_support"][i]] if has_sample_support else None acc_rewards_tokens = list(gen_out["rewards"][i]) if is_token_level_rewards else None last = i continue @@ -883,6 +906,10 @@ def flush(): acc_loss_mask.extend([0] * len(obs_delta)) if acc_logprobs is not None: acc_logprobs.extend([0.0] * len(obs_delta)) + if acc_sample_support is not None: + acc_sample_support.append( + np.full((len(obs_delta), sample_support_width), SAMPLE_SUPPORT_PADDING, dtype=SAMPLE_SUPPORT_DTYPE) + ) if acc_rewards_tokens is not None: acc_rewards_tokens.extend([0.0] * len(obs_delta)) @@ -891,6 +918,8 @@ def flush(): acc_loss_mask.extend(gen_out["loss_masks"][i]) if acc_logprobs is not None: acc_logprobs.extend(gen_out["rollout_logprobs"][i]) + if acc_sample_support is not None: + acc_sample_support.append(gen_out["rollout_sample_support"][i]) if acc_rewards_tokens is not None: acc_rewards_tokens.extend(gen_out["rewards"][i]) @@ -905,6 +934,7 @@ def flush(): "loss_masks": out_loss_masks, "stop_reasons": out_stop_reasons, "rollout_logprobs": out_logprobs, + "rollout_sample_support": out_sample_support, "trajectory_ids": out_trajectory_ids, "rollout_expert_indices": None, "is_last_step": out_is_last_step, diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index ae5eec4dfc..3180ca132d 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -45,6 +45,7 @@ compute_approx_kl, get_kl_controller, ) +from skyrl.backends.skyrl_train.utils.sample_support import SAMPLE_SUPPORT_FIELD from skyrl.backends.skyrl_train.utils.torch_utils import masked_mean from skyrl.backends.skyrl_train.workers.worker import PPORayActorGroup from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch @@ -873,6 +874,7 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis logprobs: Optional[List[List[float]]] = generator_output.get("rollout_logprobs", None) rollout_expert_indices = generator_output.get("rollout_expert_indices", None) + rollout_sample_support = generator_output.get("rollout_sample_support", None) pixel_values = generator_output.get("pixel_values", None) image_grid_thw = generator_output.get("image_grid_thw", None) @@ -895,6 +897,7 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis loss_masks_tensor, rollout_logprobs_tensor, rollout_expert_indices_tensor, + rollout_sample_support_tensor, ) = convert_prompts_responses_to_batch_tensors( self.tokenizer.pad_token_id, prompt_ids, @@ -903,6 +906,7 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis loss_masks, logprobs, rollout_expert_indices, + rollout_sample_support, max_seq_len=self.cfg.trainer.algorithm.max_seq_len, ) router_padding_mask = None @@ -933,6 +937,7 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis "rollout_logprobs": rollout_logprobs_tensor, "rollout_expert_indices": rollout_expert_indices_tensor, "router_padding_mask": router_padding_mask, + SAMPLE_SUPPORT_FIELD: rollout_sample_support_tensor, "pixel_values": pixel_values, "image_grid_thw": image_grid_thw, }, @@ -1328,6 +1333,8 @@ def fwd_logprobs_values_reward( fwd_keys.append("rollout_expert_indices") if training_input.get("router_padding_mask") is not None: fwd_keys.append("router_padding_mask") + if training_input.get(SAMPLE_SUPPORT_FIELD) is not None: + fwd_keys.append(SAMPLE_SUPPORT_FIELD) if training_input.get("pixel_values") is not None: fwd_keys.append("pixel_values") if training_input.get("image_grid_thw") is not None: diff --git a/skyrl/train/utils/trainer_utils.py b/skyrl/train/utils/trainer_utils.py index 0d91977ecb..86efc8d634 100644 --- a/skyrl/train/utils/trainer_utils.py +++ b/skyrl/train/utils/trainer_utils.py @@ -15,6 +15,7 @@ from transformers import AutoTokenizer from skyrl.backends.skyrl_train.utils.io import io +from skyrl.backends.skyrl_train.utils.sample_support import SAMPLE_SUPPORT_FIELD from skyrl.backends.skyrl_train.workers.worker import PPORayActorGroup from skyrl.backends.skyrl_train.workers.worker_utils import ( MINIBATCH_ROLLOUT_LOGPROB_DIFF_MEAN_KEY, @@ -702,6 +703,7 @@ def validate_generator_output(num_prompts: int, generator_output: GeneratorOutpu "stop_reasons", "trajectory_ids", "rollout_expert_indices", + "rollout_sample_support", "is_last_step", "pixel_values", "image_grid_thw", @@ -747,10 +749,56 @@ def validate_generator_output(num_prompts: int, generator_output: GeneratorOutpu not isinstance(reward, list) for reward in rewards ), "rewards must be `List[float]` or `List[List[float]]`" + _validate_per_token_side_channels(generator_output, step_wise) + if step_wise: _validate_step_wise_fields(generator_output, num_responses) +def _validate_per_token_side_channels(generator_output: GeneratorOutput, step_wise: bool): + """Validate side-channel row counts against their token domains.""" + rollout_expert_indices = generator_output.get("rollout_expert_indices") + rollout_sample_support = generator_output.get(SAMPLE_SUPPORT_FIELD) + prompt_token_ids = generator_output["prompt_token_ids"] + response_ids = generator_output["response_ids"] + + # Trajectory-aligned routes cannot be replayed against per-turn samples. + assert not (step_wise and rollout_expert_indices is not None), ( + "rollout router replay (r3) is not supported with step-wise training: a route trace is " + "accumulated over one contiguous prompt+response token sequence, so replaying it against " + "per-turn samples would silently route trained tokens by another token's rollout routes" + ) + + if rollout_expert_indices is not None: + loss_masks = generator_output["loss_masks"] + for i, sample_indices in enumerate(rollout_expert_indices): + assert sample_indices is not None, f"rollout_expert_indices[{i}] is None, expected captured routes" + prompt_length = len(prompt_token_ids[i]) + sequence_length = prompt_length + len(response_ids[i]) + captured_rows = len(sample_indices) + assert 0 < captured_rows <= sequence_length, ( + f"rollout_expert_indices[{i}] has {captured_rows} route rows for a " + f"{sequence_length}-token trajectory, expected a non-empty prefix of it" + ) + # Row t covers target t + 1, so every trained target needs a captured row. + trained_positions = np.flatnonzero(np.asarray(loss_masks[i])) + if trained_positions.size: + last_trained_token = prompt_length + int(trained_positions[-1]) + assert captured_rows >= last_trained_token, ( + f"rollout_expert_indices[{i}] captured {captured_rows} route rows, which stops " + f"short of loss-active token {last_trained_token}: replaying that token on a " + "dummy route trains it on a route the rollout never took" + ) + + if rollout_sample_support is not None: + for i, sample_support in enumerate(rollout_sample_support): + assert sample_support is not None, f"{SAMPLE_SUPPORT_FIELD}[{i}] is None, expected captured support" + assert len(sample_support) == len(response_ids[i]), ( + f"{SAMPLE_SUPPORT_FIELD}[{i}] has {len(sample_support)} support rows for " + f"{len(response_ids[i])} response tokens, expected one row per response token" + ) + + def _validate_step_wise_fields(generator_output: GeneratorOutput, num_responses: int): """Validate step-wise specific fields in the generator output. diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index e90c59630e..a19f5e3274 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -217,6 +217,12 @@ def validate_megatron_cfg(cfg: SkyRLTrainConfig): f"{worker_type}.megatron_config: moe_enable_routing_replay is incompatible with " "moe_router_fusion=True -- the fused router bypasses replay. Set moe_router_fusion=False." ) + # Interleaved chunks desynchronise each RouterReplay instance's backward FIFO. + assert not config.megatron_config.transformer_config_kwargs.get("virtual_pipeline_model_parallel_size"), ( + f"{worker_type}.megatron_config: moe_enable_routing_replay is incompatible with " + "virtual_pipeline_model_parallel_size -- interleaved chunks desync the replay FIFO. " + "Unset virtual_pipeline_model_parallel_size." + ) # context, expert, and expert tensor parallel are not yet supported for megatron if config.megatron_config.context_parallel_size > 1: assert ( @@ -369,6 +375,16 @@ def validate_cfg(cfg: SkyRLTrainConfig): "`token_mean_legacy` loss reduction is not supported with step-wise training. Use `token_mean` instead." ) + if cfg.generator.step_wise_trajectories and cfg.generator.inference_engine.enable_return_routed_experts: + raise ValueError( + "`generator.inference_engine.enable_return_routed_experts=True` is not supported with " + "`generator.step_wise_trajectories=True`. Each step-wise row's prompt is the whole history so " + "far, while routes are recorded for that step's generated tokens only. The trainer aligns " + "routes from the start of the sequence, so a step's routes would replay onto the first N prompt " + "tokens of its row with no length mismatch to assert on, silently training against routing that " + "does not match the rollout." + ) + if cfg.generator.merge_stepwise_output and not cfg.generator.step_wise_trajectories: raise ValueError( "`generator.merge_stepwise_output=True` requires `generator.step_wise_trajectories=True`. " diff --git a/skyrl/utils/cpu_topology.py b/skyrl/utils/cpu_topology.py new file mode 100644 index 0000000000..0afb8a6d27 --- /dev/null +++ b/skyrl/utils/cpu_topology.py @@ -0,0 +1,76 @@ +"""Determine usable CPUs from process affinity and cgroup quota.""" + +import os +from typing import Optional, Tuple + +# Container cgroup namespaces expose the current cgroup at these paths. Module-level constants +# let tests replace them with fixtures. +CGROUP_V2_CPU_MAX_PATH = "/sys/fs/cgroup/cpu.max" +CGROUP_V1_CPU_QUOTA_PATH = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us" +CGROUP_V1_CPU_PERIOD_PATH = "/sys/fs/cgroup/cpu/cpu.cfs_period_us" + +# cgroup v2 uses this literal for an unlimited quota. +CGROUP_V2_CPU_MAX_UNLIMITED = "max" + + +def _read_cgroup_file(path: str) -> str: + with open(path, encoding="utf-8") as handle: + return handle.read() + + +def _read_cgroup_v2_cpu_max() -> Optional[Tuple[float, float]]: + """``(quota, period)`` from cgroup v2 ``cpu.max``, or ``None`` if absent or unlimited.""" + try: + quota_text, period_text = _read_cgroup_file(CGROUP_V2_CPU_MAX_PATH).split() + if quota_text == CGROUP_V2_CPU_MAX_UNLIMITED: + return None + return float(quota_text), float(period_text) + except (OSError, ValueError): + return None + + +def _read_cgroup_v1_cpu_max() -> Optional[Tuple[float, float]]: + """``(quota, period)`` from cgroup v1 ``cpu.cfs_*_us``, or ``None`` if absent or unlimited.""" + try: + quota = float(_read_cgroup_file(CGROUP_V1_CPU_QUOTA_PATH).strip()) + period = float(_read_cgroup_file(CGROUP_V1_CPU_PERIOD_PATH).strip()) + except (OSError, ValueError): + return None + if quota < 0: + return None + return quota, period + + +def cgroup_cpu_quota() -> Optional[int]: + """Return whole CPUs permitted by CFS, or ``None`` when no quota applies.""" + limits = _read_cgroup_v2_cpu_max() or _read_cgroup_v1_cpu_max() + if limits is None: + return None + quota, period = limits + if quota <= 0 or period <= 0: + return None + # Floor a fractional allowance, but a sub-CPU quota still gets one worker. + return max(1, int(quota // period)) + + +def permitted_cpu_cores() -> int: + """Return the lesser of the process affinity and cgroup quota.""" + try: + affinity = len(os.sched_getaffinity(0)) + except AttributeError: + affinity = os.cpu_count() or 1 + quota = cgroup_cpu_quota() + if quota is None: + return affinity + return min(affinity, quota) + + +def pool_workers(*, cap: int, reserved: int, cores: Optional[int] = None) -> int: + """Size a pool from permitted cores, a cap, and a reserve for colocated processes.""" + if cap < 1: + raise ValueError(f"pool cap must be positive, got {cap}") + if reserved < 0: + raise ValueError(f"reserved cores must be non-negative, got {reserved}") + if cores is None: + cores = permitted_cpu_cores() + return max(1, min(cap, cores - reserved)) diff --git a/tests/backends/skyrl_train/conftest.py b/tests/backends/skyrl_train/conftest.py index fd7c45ad64..51e1150d70 100644 --- a/tests/backends/skyrl_train/conftest.py +++ b/tests/backends/skyrl_train/conftest.py @@ -1,6 +1,11 @@ +import pickle +from typing import Any + import pytest import ray +OUT_OF_BAND_PICKLE_PROTOCOL = 5 + @pytest.fixture(scope="session", autouse=True) def ray_init(): @@ -10,3 +15,16 @@ def ray_init(): yield if ray.is_initialized(): ray.shutdown() + + +@pytest.fixture +def oob_round_trip(): + """Round trip through protocol-5 buffers, optionally as read-only views.""" + + def round_trip(obj: Any, read_only: bool = False) -> tuple[Any, bytes, list[memoryview]]: + buffers: list[pickle.PickleBuffer] = [] + payload = pickle.dumps(obj, protocol=OUT_OF_BAND_PICKLE_PROTOCOL, buffer_callback=buffers.append) + views = [memoryview(bytes(buffer.raw())) if read_only else buffer.raw() for buffer in buffers] + return pickle.loads(payload, buffers=views), payload, views + + return round_trip diff --git a/tests/backends/skyrl_train/distributed/test_token_metadata.py b/tests/backends/skyrl_train/distributed/test_token_metadata.py index ff8f3dc7de..426ff4372c 100644 --- a/tests/backends/skyrl_train/distributed/test_token_metadata.py +++ b/tests/backends/skyrl_train/distributed/test_token_metadata.py @@ -1,10 +1,16 @@ import sys import types +import numpy as np import pytest import torch from skyrl.backends.skyrl_train.distributed.megatron import token_metadata +from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( + TokenMetadataTrace, +) +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor +from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertTrace @pytest.fixture @@ -86,3 +92,168 @@ def test_packed_layout_aligns_next_token_metadata_and_scatters_rows(monkeypatch, assert aligned.tolist() == [[11, 12, -1, -1, 21, -1, -1, -1]] assert batch_values.tolist() == [[0.0, 1.0, 2.0], [0.0, 0.0, 5.0]] + + +def test_token_metadata_trace_chunks_and_independent_schema() -> None: + trace, other = TokenMetadataTrace(), TokenMetadataTrace() + trace.append(np.ones((2, 3), dtype=np.int32), expected_rows=2) + trace.append(np.zeros((1, 3), dtype=np.int32), expected_rows=1) + other.append(np.empty((0, 4), dtype=np.float32), expected_rows=0) + + with pytest.raises(ValueError, match="expected 4"): + trace.finalize(expected_rows=4) + result = trace.finalize(expected_rows=3) + assert result.shape == (3, 3) + assert other.finalize(expected_rows=0).shape == (0, 4) + with pytest.raises(RuntimeError, match="already finalized"): + trace.finalize(expected_rows=3) + + +@pytest.mark.parametrize( + ("rows", "expected", "match"), + [ + (np.ones((2, 2), dtype=np.int32), 1, "has 2 rows"), + (np.ones((2, 2), dtype=np.int32)[:, ::2], 2, "contiguous"), + (np.ones((1, 3), dtype=np.int32), 1, "schema changed"), + (np.ones((1, 2), dtype=np.int16), 1, "schema changed"), + ], +) +def test_token_metadata_trace_rejects_invalid_chunks(rows, expected, match) -> None: + trace = TokenMetadataTrace() + if rows.shape[0] == 1: + trace.append(np.ones((1, 2), dtype=np.int32), expected_rows=1) + with pytest.raises(ValueError, match=match): + trace.append(rows, expected_rows=expected) + + +def routes(rows: int) -> np.ndarray: + return np.arange(rows * 4, dtype=np.int32).reshape(rows, 2, 2) % 8 + + +def test_routed_expert_trace_tracks_multiturn_suffix_and_terminal_gap() -> None: + trace = RoutedExpertTrace() + trace.record_generation(prompt_token_count=3, generated_token_count=2, routed_experts=routes(4)) + assert trace.prompt_start == 4 + trace.record_generation(prompt_token_count=7, generated_token_count=2, routed_experts=routes(4)) + + # `finalize` returns only the rows the engine actually captured -- `prompt_start`, one short + # of `token_count`, because the last sampled token has no subsequent decode forward. The + # trailing dummy row is built during collation instead, where `make_router_padding_mask` can + # mark it so Megatron excludes it from router accounting. + result = trace.finalize(token_count=9, loss_mask=[0, 0, 0, 1, 1, 0, 0, 1, 1]) + assert trace.prompt_start == 8 + assert result.shape == (8, 2, 2) and result.dtype == np.uint8 + # The last row is a real captured route now, not a dummy `arange(topk)` pad row. + assert np.array_equal(result[-1, 0], [4, 5]) + + +@pytest.mark.parametrize("active", [False, True]) +def test_routed_expert_trace_refuses_a_loss_active_target_without_a_row(active: bool) -> None: + trace = RoutedExpertTrace() + trace.record_generation(prompt_token_count=3, generated_token_count=1, routed_experts=routes(3)) + mask = [0, 0, 0, 0, int(active)] + if active: + with pytest.raises(ValueError, match="loss-active target"): + trace.finalize(token_count=5, loss_mask=mask) + else: + # A masked suffix is no longer dummy-padded up to `token_count`: the captured rows are + # returned verbatim and the gap is filled during collation. + result = trace.finalize(token_count=5, loss_mask=mask) + assert result.shape == (3, 2, 2) + assert np.array_equal(result, routes(3).astype(result.dtype)) + + +@pytest.mark.parametrize("packed", [False, True]) +def test_align_token_rows_places_each_trajectory_from_its_own_row_source(monkeypatch, parallel_state, packed): + """``_align_token_rows`` is the one placement loop both alignment entry points share.""" + monkeypatch.setattr(token_metadata, "get_packed_seq_align_size", lambda *args, **kwargs: 4) + monkeypatch.setattr(token_metadata, "get_unpacked_seq_align_size", lambda *args, **kwargs: 4) + attention_mask = torch.tensor([[0, 1, 1, 1], [0, 0, 1, 1]]) + rows = [torch.tensor([10, 11, 12], dtype=torch.int32), torch.tensor([20, 21], dtype=torch.int32)] + layout = token_metadata.build_token_metadata_layout( + attention_mask, + rows[0].device, + packed=packed, + fp8_enabled=False, + ) + + aligned = token_metadata._align_token_rows( + rows.__getitem__, + rows[0], + (), + layout, + -1, + ) + + if packed: + assert aligned.tolist() == [[10, 11, 12, -1, 20, 21, -1, -1]] + else: + assert aligned.tolist() == [[10, 11, 12, -1], [20, 21, -1, -1]] + + +@pytest.mark.parametrize("packed", [False, True]) +def test_align_packed_token_metadata_honours_per_segment_starts(monkeypatch, parallel_state, packed): + """A response-suffix channel covers part of a trajectory and needs its own start.""" + monkeypatch.setattr(token_metadata, "get_packed_seq_align_size", lambda *args, **kwargs: 4) + monkeypatch.setattr(token_metadata, "get_unpacked_seq_align_size", lambda *args, **kwargs: 4) + attention_mask = torch.tensor([[0, 1, 1, 1], [0, 0, 1, 1]]) + # Trajectory 0 keeps its last 2 of 3 real tokens; trajectory 1 keeps its last 1 of 2. + suffix = PackedTensor.from_segments( + [torch.tensor([11, 12], dtype=torch.int32), torch.tensor([21], dtype=torch.int32)] + ) + layout = token_metadata.build_token_metadata_layout( + attention_mask, + suffix.device, + packed=packed, + fp8_enabled=False, + ) + + aligned = token_metadata.align_packed_token_metadata(suffix, layout, -1, segment_starts=[1, 1]) + + if packed: + assert aligned.tolist() == [[-1, 11, 12, -1, -1, 21, -1, -1]] + else: + assert aligned.tolist() == [[-1, 11, 12, -1], [-1, 21, -1, -1]] + + +def test_align_packed_token_metadata_rejects_segments_that_leave_the_trajectory(monkeypatch, parallel_state): + monkeypatch.setattr(token_metadata, "get_unpacked_seq_align_size", lambda *args, **kwargs: 4) + attention_mask = torch.tensor([[0, 1, 1, 1]]) + suffix = PackedTensor.from_segments([torch.tensor([11, 12], dtype=torch.int32)]) + layout = token_metadata.build_token_metadata_layout( + attention_mask, + suffix.device, + packed=False, + fp8_enabled=False, + ) + + with pytest.raises(ValueError, match="spans real tokens"): + token_metadata.align_packed_token_metadata(suffix, layout, -1, segment_starts=[2]) + with pytest.raises(ValueError, match="do not match"): + token_metadata.align_packed_token_metadata(suffix, layout, -1) + + +def test_append_padding_extends_the_established_schema(): + trace = token_metadata.TokenMetadataTrace() + trace.append(np.array([[7, 8], [9, 10]], dtype=np.int32), expected_rows=2) + + trace.append_padding(0) + assert trace.num_rows == 2 + + trace.append_padding(2) + padded = trace.finalize(expected_rows=4) + + assert padded.dtype == np.int32 + assert padded.flags.c_contiguous + assert padded.tolist() == [[7, 8], [9, 10], [-1, -1], [-1, -1]] + + +def test_append_padding_needs_a_schema_and_a_valid_count(): + with pytest.raises(ValueError, match="before any rows are captured"): + token_metadata.TokenMetadataTrace().append_padding(1) + + trace = token_metadata.TokenMetadataTrace() + trace.append(np.zeros((1, 2), dtype=np.int32), expected_rows=1) + for count in (-1, True, 1.0): + with pytest.raises(ValueError, match="padding count"): + trace.append_padding(count) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py index d2334a8427..5810cc1038 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py @@ -143,7 +143,7 @@ async def generate_with_vllm(generator, client, model_name, tokenizer, return_tr if rewards and not isinstance(rewards[0], list): rewards = [[r] * len(resp) for r, resp in zip(rewards, responses)] - sequences, attention_mask, response_mask, rewards_t, loss_mask_t, logprobs_t, _ = ( + sequences, attention_mask, response_mask, rewards_t, loss_mask_t, logprobs_t, _, _ = ( convert_prompts_responses_to_batch_tensors( pad_token_id=tokenizer.pad_token_id, prompts=generator_output["prompt_token_ids"], diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_router_replay.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_router_replay.py index 3f97e0ed50..10af38c40e 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_router_replay.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_router_replay.py @@ -16,6 +16,7 @@ get_sampling_params_for_backend, ) from skyrl.backends.skyrl_train.training_batch import TrainingInputBatch +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor from skyrl.train.config import SamplingParams, SkyRLTrainConfig from skyrl.train.dataset.preprocess import ( convert_prompts_responses_to_batch_tensors, @@ -36,6 +37,25 @@ NUM_PROMPTS = 10 N_SAMPLES_PER_PROMPT = 4 MAX_GENERATE_LENGTH = 128 +# Moonlight 16B: 27 MoE layers, top_k=6, 64 routed experts. +MOONLIGHT_NUM_LAYERS = 27 +MOONLIGHT_TOPK = 6 +MOONLIGHT_NUM_EXPERTS = 64 + + +def _packed_moonlight_routes(attention_mask: torch.Tensor) -> PackedTensor: + """Moonlight-shaped routes packed to each trajectory's real tokens.""" + route_offsets = torch.arange(MOONLIGHT_TOPK, dtype=torch.int32) + segments = [] + for real_tokens in attention_mask.sum(dim=1).tolist(): + route_start = torch.randint( + 0, + MOONLIGHT_NUM_EXPERTS, + (real_tokens, MOONLIGHT_NUM_LAYERS, 1), + dtype=torch.int32, + ) + segments.append((route_start + route_offsets) % MOONLIGHT_NUM_EXPERTS) + return PackedTensor.from_segments(segments) def _extra_env_vars_for_model(model_name: str) -> dict[str, str] | None: @@ -108,12 +128,14 @@ def build_training_input_from_text_samples( rewards.append([0.0] * len(response_ids)) loss_masks.append([1] * len(response_ids)) - sequences, attention_mask, response_mask, rewards_t, loss_mask_t, _, _ = convert_prompts_responses_to_batch_tensors( - pad_token_id=tokenizer.pad_token_id, - prompts=prompts, - responses=responses, - rewards=rewards, - loss_masks=loss_masks, + sequences, attention_mask, response_mask, rewards_t, loss_mask_t, _, _, _ = ( + convert_prompts_responses_to_batch_tensors( + pad_token_id=tokenizer.pad_token_id, + prompts=prompts, + responses=responses, + rewards=rewards, + loss_masks=loss_masks, + ) ) num_actions = response_mask.shape[1] @@ -217,7 +239,7 @@ async def test_logprobs(tp, pp, cp, ep, etp, extra_tf_kwargs): rewards = generator_output["rewards"] if rewards and not isinstance(rewards[0], list): rewards = [[r] * len(resp) for r, resp in zip(rewards, responses)] - sequences, attention_mask, response_mask, rewards_t, loss_mask_t, logprobs_t, rii_tensor = ( + sequences, attention_mask, response_mask, rewards_t, loss_mask_t, logprobs_t, rii_tensor, _ = ( convert_prompts_responses_to_batch_tensors( pad_token_id=tokenizer.pad_token_id, prompts=generator_output["prompt_token_ids"], @@ -342,7 +364,7 @@ def test_forward_backward(tp, pp, cp, ep, etp, extra_tf_kwargs): rewards.append([1.0] * len(response_ids)) loss_masks.append([1] * len(response_ids)) - sequences, attention_mask, response_mask, rewards_t, loss_mask_t, _, _ = ( + sequences, attention_mask, response_mask, rewards_t, loss_mask_t, _, _, _ = ( convert_prompts_responses_to_batch_tensors( pad_token_id=tokenizer.pad_token_id, prompts=prompts, @@ -353,22 +375,9 @@ def test_forward_backward(tp, pp, cp, ep, etp, extra_tf_kwargs): ) batch_size = sequences.shape[0] - seq_len = sequences.shape[1] num_actions = response_mask.shape[1] - # Moonlight 16B: 27 MoE layers, top_k=6, 64 routed experts - MOONLIGHT_NUM_LAYERS = 27 - MOONLIGHT_TOPK = 6 - MOONLIGHT_NUM_EXPERTS = 64 - route_start = torch.randint( - 0, - MOONLIGHT_NUM_EXPERTS, - (batch_size, seq_len, MOONLIGHT_NUM_LAYERS, 1), - dtype=torch.int32, - ) - route_offsets = torch.arange(MOONLIGHT_TOPK, dtype=torch.int32) - rollout_expert_indices = (route_start + route_offsets) % MOONLIGHT_NUM_EXPERTS - rollout_expert_indices[attention_mask == 0] = route_offsets + rollout_expert_indices = _packed_moonlight_routes(attention_mask) gen = torch.Generator().manual_seed(42) training_input = TrainingInputBatch( @@ -419,3 +428,145 @@ def test_forward_backward(tp, pp, cp, ep, etp, extra_tf_kwargs): for actor in actor_group._actor_handlers: ray.kill(actor) + + +@pytest.mark.h100 +@pytest.mark.parametrize( + "tp,pp,cp,ep,etp,extra_tf_kwargs", + [ + pytest.param(2, 2, 1, 2, 1, {"num_layers_in_first_pipeline_stage": 13}, id="tp2_pp2_ep2"), + ], +) +def test_forward_backward_variable_length_full_recompute(tp, pp, cp, ep, etp, extra_tf_kwargs): + """Replayed routes must stay paired with their own microbatch. + + Each forward microbatch appends its expert routes to a FIFO that + activation-checkpoint recomputation drains once during backward. If a + microbatch queues its routes more than once (or not at all), backward + replays a *different* microbatch's routes. Megatron's MoE all-to-all then + computes split sizes from a token count that doesn't match the tensors in + flight and the collective fails. + + Two conditions are needed to expose that, and both are set here: + + * ``recompute_granularity="full"`` so backward actually recomputes the + MoE layers and consumes the queue (the default only recomputes + ``core_attn``, which never replays routes). + * Sequence lengths that differ *between* microbatches, so a mispaired + replay changes the token count rather than silently reusing a + same-shaped tensor. Uniform lengths can mask the bug entirely. + + Prompts are built with deliberately spread lengths and ``micro_*=2`` over + 8 samples, giving 4 microbatches whose padded widths differ. + """ + with ray_init(extra_env_vars=_extra_env_vars_for_model(MOE_MODEL_NAME)): + cfg = get_test_actor_config(model_name=MOE_MODEL_NAME) + cfg.trainer.strategy = "megatron" + + tokenizer = AutoTokenizer.from_pretrained(MOE_MODEL_NAME, trust_remote_code=True) + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + + # Lengths chosen so consecutive microbatches (pairs, micro_bs=2) have + # different maxima: a stale replay is then shape-visible, not benign. + filler_words = [1, 6, 2, 14, 3, 25, 4, 40] + prompts, responses, rewards, loss_masks = [], [], [], [] + for i, filler in enumerate(filler_words): + prompt_ids = tokenizer.encode( + "Question: " + ("token " * filler) + f"what is {i} plus {i}?", + add_special_tokens=False, + ) + response_ids = tokenizer.encode( + ("because " * filler) + f"the answer is {i + i}.", + add_special_tokens=False, + ) + if tokenizer.eos_token_id is not None and (not response_ids or response_ids[-1] != tokenizer.eos_token_id): + response_ids.append(tokenizer.eos_token_id) + prompts.append(prompt_ids) + responses.append(response_ids) + rewards.append([1.0] * len(response_ids)) + loss_masks.append([1] * len(response_ids)) + + sequences, attention_mask, response_mask, rewards_t, loss_mask_t, _, _, _ = ( + convert_prompts_responses_to_batch_tensors( + tokenizer=tokenizer, + prompts=prompts, + responses=responses, + rewards=rewards, + loss_masks=loss_masks, + ) + ) + + # Guard the premise: if padding collapsed the spread, the test would + # pass for the wrong reason. + real_token_counts = attention_mask.sum(dim=-1) + assert ( + real_token_counts.unique().numel() > 1 + ), f"variable-length premise broken: every sample has {real_token_counts[0].item()} real tokens" + + batch_size = sequences.shape[0] + num_actions = response_mask.shape[1] + + rollout_expert_indices = _packed_moonlight_routes(attention_mask) + + gen = torch.Generator().manual_seed(42) + training_input = TrainingInputBatch( + { + "sequences": sequences, + "attention_mask": attention_mask, + "response_mask": response_mask, + "rewards": rewards_t, + "loss_mask": loss_mask_t, + "rollout_logprobs": -torch.rand((batch_size, num_actions), generator=gen) * 2.0, + "rollout_expert_indices": rollout_expert_indices, + "router_padding_mask": ~attention_mask.bool(), + "action_log_probs": -torch.rand((batch_size, num_actions), generator=gen) * 2.0, + "base_action_log_probs": -torch.rand((batch_size, num_actions), generator=gen) * 2.0, + "advantages": torch.randn((batch_size, num_actions), generator=gen), + "action_mask": response_mask.to(dtype=torch.int64), + } + ) + training_input.metadata = {"response_length": num_actions} + + cfg.trainer.placement.policy_num_gpus_per_node = 4 + if extra_tf_kwargs is not None: + cfg.trainer.policy.megatron_config.transformer_config_kwargs.update(extra_tf_kwargs) + # Recompute whole layers so backward re-runs the MoE routers and drains + # the replay queue; ``core_attn`` alone never replays routes. + cfg.trainer.policy.megatron_config.transformer_config_kwargs.update( + { + "recompute_granularity": "full", + "recompute_method": "uniform", + "recompute_num_layers": 1, + } + ) + cfg.trainer.policy.megatron_config.transformer_config_kwargs.pop("recompute_modules", None) + cfg.trainer.policy.megatron_config.tensor_model_parallel_size = tp + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = pp + cfg.trainer.policy.megatron_config.context_parallel_size = cp + cfg.trainer.policy.megatron_config.expert_model_parallel_size = ep + cfg.trainer.policy.megatron_config.expert_tensor_parallel_size = etp + # 8 samples / 2 per microbatch = 4 microbatches of differing widths. + cfg.trainer.micro_forward_batch_size_per_gpu = 2 + cfg.trainer.micro_train_batch_size_per_gpu = 2 + cfg.trainer.policy.megatron_config.moe_enable_routing_replay = True + + actor_group = init_worker_with_type( + "policy", + num_gpus_per_node=4, + cfg=cfg, + ) + + # Two steps: the first leaves any surplus queue entry behind, so a + # mispairing shows up on the second even if the first survives. + ray.get(actor_group.async_run_ray_method("mesh", "forward_backward", data=training_input)) + ray.get(actor_group.async_run_ray_method("pass_through", "optim_step")) + results = ray.get(actor_group.async_run_ray_method("mesh", "forward_backward", data=training_input)) + + loss = results[0].metrics["policy_loss"] + print(f"Variable-length replay forward_backward - loss: {loss:.6f}") + assert loss is not None and not torch.isnan(torch.tensor(loss)), "Loss should be valid (not NaN)" + assert loss != 0.0, "Loss should be non-zero with non-zero advantages" + + for actor in actor_group._actor_handlers: + ray.kill(actor) diff --git a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py index ad3372e06f..b76166787e 100644 --- a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py +++ b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py @@ -40,6 +40,21 @@ def test_build_vllm_cli_args_succeeds_on_gpu_less_host(monkeypatch): # tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py +@pytest.mark.vllm +def test_sample_support_uses_processed_top_k_logprobs(): + cfg = SkyRLTrainConfig.from_cli_overrides( + [ + "generator.inference_engine.enable_return_sample_support_set=true", + "generator.sampling_params.top_k=8", + ] + ) + + args = build_vllm_cli_args(cfg) + + assert args.max_logprobs == 8 + assert args.logprobs_mode == "processed_logprobs" + + def test_resolve_policy_model_name_uses_served_model_name(): cfg = SkyRLTrainConfig() cfg.trainer.policy.model.path = "base-model" diff --git a/tests/backends/skyrl_train/inference_servers/test_generate_wire.py b/tests/backends/skyrl_train/inference_servers/test_generate_wire.py index ac2a6219ae..35e2c551e1 100644 --- a/tests/backends/skyrl_train/inference_servers/test_generate_wire.py +++ b/tests/backends/skyrl_train/inference_servers/test_generate_wire.py @@ -1,6 +1,7 @@ """Tests for the /skyrl/v1/generate payload contract.""" import base64 +import json import math from dataclasses import dataclass @@ -11,11 +12,29 @@ from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( CLAMPED_LOGPROB, + PackedArrayKey, + PackedField, build_logprobs_content, decode_packed_routed_experts, + decode_packed_sample_support, + load_packed_body, + pack_ndarray, pack_routed_experts, + pack_sample_support, + unpack_ndarray, ) +_FLOAT32 = frozenset({np.dtype(np.float32)}) +_INT16 = frozenset({np.dtype(np.int16)}) + + +def _support_envelope(support: np.ndarray) -> dict: + return pack_ndarray(support, allowed_dtypes=_FLOAT32) + + +def _body(**choice_fields) -> dict: + return {"choices": [{"token_ids": [1, 2, 3], "finish_reason": "stop", **choice_fields}]} + @dataclass class _Logprob: @@ -183,3 +202,215 @@ def test_decode_rejects_noncanonical_dtype(): with pytest.raises(ValueError, match="non-canonical dtype"): decode_packed_routed_experts(payload) + + +def test_packed_sample_support_round_trip(): + support = np.array([[7, 152064, -1, -1], [9, 10, 11, -1], [12, -1, -1, -1]], dtype=np.int32) + + decoded = decode_packed_sample_support(orjson.loads(orjson.dumps(pack_sample_support(support)))) + + assert decoded.dtype == np.int32 + assert decoded.flags.c_contiguous + assert np.array_equal(decoded, support) + + +# Shape, base64, byte-size and dtype allow-listing are covered generically by the +# pack_ndarray/unpack_ndarray tests; these two semantic invariants are not. +@pytest.mark.parametrize( + "support,message", + [ + (np.array([[-2, 1]], dtype=np.int32), "-1 padding"), + (np.array([[1, -1, 2]], dtype=np.int32), "trailing"), + ], +) +def test_sample_support_wire_rejects_bad_padding(support, message): + with pytest.raises(ValueError, match=message): + pack_sample_support(support) + with pytest.raises(ValueError, match=message): + decode_packed_sample_support(pack_ndarray(support, allowed_dtypes=frozenset({np.dtype(np.int32)}))) + + +@pytest.mark.parametrize( + "arr,allowed_dtypes,extra", + [ + (np.arange(6, dtype=np.float32).reshape(2, 3), _FLOAT32, None), + (np.arange(6, dtype=np.float32).reshape(2, 3), _FLOAT32, {"prompt_start": 4, "labels": ["a", "b"]}), + (np.arange(12, dtype=np.int16).reshape(3, 2, 2), _INT16, {"prompt_start": 0}), + (np.empty((0, 3), dtype=np.float32), _FLOAT32, None), + ], +) +def test_ndarray_round_trip_with_sidecar_fields(arr, allowed_dtypes, extra): + payload = pack_ndarray(arr, allowed_dtypes=allowed_dtypes, extra=extra) + decoded, sidecar = unpack_ndarray(payload, allowed_dtypes=allowed_dtypes, ndim=arr.ndim) + + assert np.array_equal(decoded, arr) + assert decoded.dtype == arr.dtype + assert decoded.flags.c_contiguous + assert sidecar == (extra or {}) + + +def test_packed_envelope_leads_with_data(): + payload = pack_ndarray(np.zeros((2, 2), np.float32), allowed_dtypes=_FLOAT32, extra={"prompt_start": 1}) + + assert list(payload) == [PackedArrayKey.DATA, PackedArrayKey.SHAPE, PackedArrayKey.DTYPE, "prompt_start"] + assert orjson.dumps(payload).startswith(b'{"data":"') + + +def test_pack_routed_experts_is_byte_identical_to_the_hand_built_envelope(): + routes = np.arange(12).reshape(3, 2, 2) + + assert orjson.dumps(pack_routed_experts(routes)) == b'{"data":"AAECAwQFBgcICQoL","shape":[3,2,2],"dtype":"uint8"}' + + +@pytest.mark.parametrize( + "wrap", [str, lambda data: memoryview(data.encode("ascii")), lambda data: data.encode("ascii")] +) +def test_unpack_accepts_str_and_buffers(wrap): + support = np.arange(6, dtype=np.float32).reshape(2, 3) + payload = dict(_support_envelope(support)) + payload[PackedArrayKey.DATA.value] = wrap(payload[PackedArrayKey.DATA.value]) + + decoded, _ = unpack_ndarray(payload, allowed_dtypes=_FLOAT32, ndim=2) + assert np.array_equal(decoded, support) + + +def test_pack_rejects_disallowed_dtype(): + with pytest.raises(ValueError, match="dtype"): + pack_ndarray(np.zeros((2, 2), np.float64), allowed_dtypes=_FLOAT32) + + +def test_pack_rejects_sidecar_collision_with_envelope_keys(): + with pytest.raises(ValueError, match="collide"): + pack_ndarray(np.zeros((2, 2), np.float32), allowed_dtypes=_FLOAT32, extra={"dtype": "float64"}) + + +def test_unpack_rejects_disallowed_dtype(): + payload = pack_ndarray(np.zeros((2, 2), np.float32), allowed_dtypes=_FLOAT32) + + with pytest.raises(ValueError, match="dtype"): + unpack_ndarray(payload, allowed_dtypes=_INT16, ndim=2) + + +@pytest.mark.parametrize("ndim", [1, 3]) +def test_unpack_rejects_wrong_ndim(ndim): + payload = pack_ndarray(np.zeros((2, 2), np.float32), allowed_dtypes=_FLOAT32) + + with pytest.raises(ValueError, match="dimensions"): + unpack_ndarray(payload, allowed_dtypes=_FLOAT32, ndim=ndim) + + +def test_unpack_rejects_byte_count_mismatched_with_declared_shape(): + payload = pack_ndarray(np.zeros((2, 3), np.float32), allowed_dtypes=_FLOAT32) + payload[PackedArrayKey.SHAPE.value] = [2, 4] + + with pytest.raises(ValueError, match="24 bytes, expected 32"): + unpack_ndarray(payload, allowed_dtypes=_FLOAT32, ndim=2) + + +def test_load_packed_body_splices_both_blobs_in_one_body(): + routes = np.arange(12).reshape(3, 2, 2) + support = np.arange(6, dtype=np.float32).reshape(2, 3) + raw = orjson.dumps( + _body( + logprobs={"content": [{"logprob": -0.5}]}, + routed_experts=pack_routed_experts(routes), + rollout_sample_support=_support_envelope(support), + ) + ) + + choice = load_packed_body(raw)["choices"][0] + + assert all( + isinstance(choice[field][PackedArrayKey.DATA], memoryview) + for field in (PackedField.ROUTED_EXPERTS, PackedField.ROLLOUT_SAMPLE_SUPPORT) + ) + assert np.array_equal(decode_packed_routed_experts(choice[PackedField.ROUTED_EXPERTS]), routes) + decoded_support, _ = unpack_ndarray(choice[PackedField.ROLLOUT_SAMPLE_SUPPORT], allowed_dtypes=_FLOAT32, ndim=2) + assert np.array_equal(decoded_support, support) + assert choice["logprobs"] == {"content": [{"logprob": -0.5}]} + + +def test_load_packed_body_keeps_sidecar_fields(): + support = np.arange(4, dtype=np.float32).reshape(2, 2) + envelope = pack_ndarray(support, allowed_dtypes=_FLOAT32, extra={"prompt_start": 7}) + raw = orjson.dumps(_body(rollout_sample_support=envelope)) + + decoded, sidecar = unpack_ndarray( + load_packed_body(raw)["choices"][0][PackedField.ROLLOUT_SAMPLE_SUPPORT], + allowed_dtypes=_FLOAT32, + ndim=2, + ) + + assert np.array_equal(decoded, support) + assert sidecar == {"prompt_start": 7} + + +@pytest.mark.parametrize("value", [None, "absent"]) +def test_load_packed_body_passes_through_absent_and_null_fields(value): + fields = {} if value == "absent" else {PackedField.ROUTED_EXPERTS.value: None} + body = _body(logprobs=None, **fields) + + assert load_packed_body(orjson.dumps(body)) == body + + +def test_load_packed_body_rejects_reordered_envelope_keys(): + routes = np.arange(12).reshape(3, 2, 2) + envelope = pack_routed_experts(routes) + reordered = {key: envelope[key] for key in reversed(list(envelope))} + + with pytest.raises(ValueError, match="layout drifted"): + load_packed_body(orjson.dumps(_body(routed_experts=reordered))) + + +def test_load_packed_body_rejects_a_reserialized_body(): + # stdlib json spaces its separators; the blob would silently land in a + # ~121 MiB Python str instead, which is exactly the cost this avoids. + raw = json.dumps(_body(routed_experts=pack_routed_experts(np.arange(12).reshape(3, 2, 2)))).encode() + + with pytest.raises(ValueError, match="layout drifted"): + load_packed_body(raw) + + +def test_load_packed_body_is_not_spoofable_from_a_string_value(): + routes = np.arange(12).reshape(3, 2, 2) + spoof = '"routed_experts":{"data":"AAAA","shape":[1,1,1],"dtype":"uint8"}' + raw = orjson.dumps(_body(note=spoof, routed_experts=pack_routed_experts(routes))) + + choice = load_packed_body(raw)["choices"][0] + + assert choice["note"] == spoof + assert np.array_equal(decode_packed_routed_experts(choice[PackedField.ROUTED_EXPERTS]), routes) + + +def test_load_packed_body_ignores_unregistered_packed_fields(): + envelope = pack_ndarray(np.zeros((2, 2), np.float32), allowed_dtypes=_FLOAT32) + body = _body(some_other_array=envelope) + + assert load_packed_body(orjson.dumps(body)) == body + + +def test_load_packed_body_honours_a_narrowed_field_registry(): + routes = np.arange(12).reshape(3, 2, 2) + raw = orjson.dumps(_body(routed_experts=pack_routed_experts(routes))) + + body = load_packed_body(raw, fields=(PackedField.ROLLOUT_SAMPLE_SUPPORT,)) + + assert isinstance(body["choices"][0][PackedField.ROUTED_EXPERTS][PackedArrayKey.DATA], str) + + +def test_load_packed_body_splices_one_blob_per_choice(): + first, second = np.arange(12).reshape(3, 2, 2), np.arange(12, 24).reshape(3, 2, 2) + raw = orjson.dumps({"choices": [{"routed_experts": pack_routed_experts(routes)} for routes in (first, second)]}) + + choices = load_packed_body(raw)["choices"] + + assert np.array_equal(decode_packed_routed_experts(choices[0][PackedField.ROUTED_EXPERTS]), first) + assert np.array_equal(decode_packed_routed_experts(choices[1][PackedField.ROUTED_EXPERTS]), second) + + +def test_load_packed_body_rejects_an_unterminated_blob(): + raw = orjson.dumps(_body(routed_experts=pack_routed_experts(np.arange(12).reshape(3, 2, 2)))) + truncated = raw[: raw.index(b'"shape"') - 2] + + with pytest.raises(ValueError, match="unterminated"): + load_packed_body(truncated) diff --git a/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py b/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py index 75267428b6..7457d49df2 100644 --- a/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py +++ b/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py @@ -1,6 +1,7 @@ """Tests for RemoteInferenceClient.""" import asyncio +import json import pickle import threading import time @@ -9,19 +10,30 @@ import aiohttp import httpx import numpy as np +import orjson import pytest import pytest_asyncio import uvicorn from fastapi import FastAPI, Query, Request -from fastapi.responses import JSONResponse, PlainTextResponse +from fastapi.responses import JSONResponse, PlainTextResponse, Response +from skyrl.backends.skyrl_train.inference_servers import ( + remote_inference_client as remote_client_module, +) from skyrl.backends.skyrl_train.inference_servers.common import get_open_port from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( + PackedArrayKey, + PackedField, + decode_packed_routed_experts, + pack_ndarray, pack_routed_experts, + pack_sample_support, + unpack_ndarray, ) from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( SKYRL_LORA_ADAPTER_NAME, PauseMode, + RemoteGenerateClient, RemoteInferenceClient, ) from skyrl.backends.skyrl_train.inference_servers.setup import ( @@ -29,12 +41,34 @@ ) from skyrl.train.config import SkyRLTrainConfig +_SUPPORT_DTYPES = frozenset({np.dtype(np.float32)}) +_ROUTES = np.arange(12).reshape(3, 2, 2) + + +async def _fake_detokenize(token_id_lists: List[List[int]]) -> List[str]: + return ["text"] * len(token_id_lists) + + +_SUPPORT = np.arange(6, dtype=np.float32).reshape(2, 3) + + +def _packed_generate_body(*, two_blobs: bool) -> dict: + choice: dict = { + "token_ids": [1, 2, 3], + "finish_reason": "stop", + PackedField.ROUTED_EXPERTS.value: pack_routed_experts(_ROUTES), + } + if two_blobs: + choice[PackedField.ROLLOUT_SAMPLE_SUPPORT.value] = pack_ndarray(_SUPPORT, allowed_dtypes=_SUPPORT_DTYPES) + return {"choices": [choice]} + def create_mock_vllm_server(server_id: int) -> FastAPI: """Create a mock vLLM server with standard endpoints.""" app = FastAPI() app.state.last_generate_features = None app.state.last_generate_model = None + app.state.last_generate_sampling_params = None app.state.last_chat_model = None app.state.last_completion_model = None app.state.last_render_model = None @@ -45,11 +79,45 @@ def create_mock_vllm_server(server_id: int) -> FastAPI: app.state.finished_sessions = [] # Number of /get_world_size hits, used to assert client-side caching. app.state.world_size_calls = 0 + app.state.drifted_body_calls = 0 + app.state.flaky_body_calls = 0 @app.get("/health") async def health(): return {"status": "ok"} + @app.post("/test/packed_body") + async def packed_body(two_blobs: bool = False): + return Response(content=orjson.dumps(_packed_generate_body(two_blobs=two_blobs)), media_type="application/json") + + @app.post("/test/drifted_packed_body") + async def drifted_packed_body(): + app.state.drifted_body_calls += 1 + # stdlib json spaces its separators, so the splice prefix no longer matches. + content = json.dumps(_packed_generate_body(two_blobs=False)).encode() + return Response(content=content, media_type="application/json") + + @app.post("/test/flaky_packed_body") + async def flaky_packed_body(): + app.state.flaky_body_calls += 1 + if app.state.flaky_body_calls == 1: + return Response(content=b"gateway hiccup", media_type="application/json", status_code=502) + return Response(content=orjson.dumps(_packed_generate_body(two_blobs=True)), media_type="application/json") + + @app.post("/test/reset_packed_body_calls") + async def reset_packed_body_calls(): + app.state.drifted_body_calls = 0 + app.state.flaky_body_calls = 0 + return {"status": "ok"} + + @app.get("/test/packed_body_calls") + async def packed_body_calls(): + return {"drifted": app.state.drifted_body_calls, "flaky": app.state.flaky_body_calls} + + @app.post("/test/bad_request_text") + async def bad_request_text(): + return PlainTextResponse("prompt too long", status_code=400) + @app.post("/finish_session") async def finish_session(session_id: str = Query(...)): app.state.finished_sessions.append(session_id) @@ -63,6 +131,10 @@ async def get_finished(): async def get_last_generate_features(): return {"features": app.state.last_generate_features} + @app.get("/test/last_generate_sampling_params") + async def get_last_generate_sampling_params(): + return app.state.last_generate_sampling_params + @app.get("/test/last_models") async def get_last_models(): return { @@ -104,6 +176,7 @@ async def completions(request: Request): async def generate(request: Request): body = await request.json() # Consume body sp = body.get("sampling_params", {}) + app.state.last_generate_sampling_params = sp input_token_ids = body.get("token_ids", []) app.state.last_generate_model = body.get("model") n = sp.get("n", 1) @@ -438,8 +511,7 @@ def test_serialization(self, mock_servers): assert restored.proxy_url == client.proxy_url assert restored.server_urls == client.server_urls assert restored.model_name == client.model_name - # Session should be None after unpickling - assert restored._session is None + assert restored._generate_client is None class TestDataPlane: @@ -480,13 +552,106 @@ async def test_generate_decodes_packed_routed_experts(self, mock_servers): enable_return_routed_experts=True, ) try: - result = await client.generate({"prompt_token_ids": [[1, 2, 3]]}) + result = await client.generate({"prompt_token_ids": [[1, 2, 3]], "routed_experts_prompt_starts": [1]}) + async with httpx.AsyncClient() as http: + captured = (await http.get(f"{mock_servers['proxy_url']}/test/last_generate_sampling_params")).json() finally: await client.teardown() assert len(result["rollout_expert_indices"]) == 1 assert result["rollout_expert_indices"][0].dtype == np.uint8 assert np.array_equal(result["rollout_expert_indices"][0], np.arange(12).reshape(3, 2, 2)) + assert captured["routed_experts_prompt_start"] == 1 + + @pytest.mark.asyncio + async def test_external_generator_requests_sample_support(self, monkeypatch): + generate_client = RemoteGenerateClient(proxy_url="http://unused") + captured = {} + + async def fake_post(url, json, headers, *, packed_side_channels=False): + captured.update(url=url, json=json, headers=headers, packed_side_channels=packed_side_channels) + return { + "choices": [ + { + "token_ids": [7], + "finish_reason": "stop", + "logprobs": {"content": [{"logprob": -0.1}]}, + PackedField.ROLLOUT_SAMPLE_SUPPORT.value: pack_sample_support( + np.array([[7, 8]], dtype=np.int32) + ), + } + ] + } + + monkeypatch.setattr(generate_client, "_post", fake_post) + result = await generate_client.generate( + prompt_token_ids=[1, 2], + sampling_params={}, + session_id=None, + model="default", + return_sample_support=True, + ) + + assert captured["url"].endswith("/skyrl/v1/generate") + assert captured["json"]["return_sample_support"] is True + assert captured["packed_side_channels"] is True + assert np.array_equal(result.sample_support, np.array([[7, 8]], dtype=np.int32)) + + @pytest.mark.asyncio + @pytest.mark.parametrize("input_batch_opts", [{}, {"return_sample_support": False}]) + async def test_sample_support_capture_is_opt_in_per_request(self, monkeypatch, input_batch_opts): + client = RemoteInferenceClient( + proxy_url="http://unused", + server_urls=["http://unused"], + data_parallel_size=1, + enable_return_sample_support_set=True, + ) + captured = {} + + async def fake_post(url, json, headers, *, packed_side_channels=False): + captured.update(url=url, json=json) + return {"choices": [{"token_ids": [1], "finish_reason": "stop"}]} + + monkeypatch.setattr(client._get_generate_client(), "_post", fake_post) + monkeypatch.setattr(client, "detokenize", _fake_detokenize) + + result = await client.generate({"prompt_token_ids": [[1, 2]], **input_batch_opts}) + + assert "return_sample_support" not in captured["json"] + assert captured["url"].endswith("/inference/v1/generate") + assert result["rollout_sample_support"] is None + + @pytest.mark.asyncio + async def test_sample_support_capture_honours_an_explicit_opt_in(self, monkeypatch): + client = RemoteInferenceClient( + proxy_url="http://unused", + server_urls=["http://unused"], + data_parallel_size=1, + enable_return_sample_support_set=True, + ) + captured = {} + + async def fake_post(url, json, headers, *, packed_side_channels=False): + captured.update(url=url, json=json) + return { + "choices": [ + { + "token_ids": [7], + "finish_reason": "stop", + PackedField.ROLLOUT_SAMPLE_SUPPORT.value: pack_sample_support( + np.array([[7, 8]], dtype=np.int32) + ), + } + ] + } + + monkeypatch.setattr(client._get_generate_client(), "_post", fake_post) + monkeypatch.setattr(client, "detokenize", _fake_detokenize) + + result = await client.generate({"prompt_token_ids": [[1, 2]], "return_sample_support": True}) + + assert captured["json"]["return_sample_support"] is True + assert np.array_equal(result["rollout_sample_support"][0], np.array([[7, 8]], dtype=np.int32)) @pytest.mark.asyncio async def test_generate_rejects_list_routed_experts(self, monkeypatch): @@ -508,7 +673,7 @@ async def return_list_routes(*args, **kwargs): ] } - monkeypatch.setattr(client, "_post", return_list_routes) + monkeypatch.setattr(client._get_generate_client(), "_post", return_list_routes) with pytest.raises(ValueError, match="must return packed"): await client._generate_single([1], {}, None, "model") @@ -550,6 +715,90 @@ async def test_detokenize(self, client): assert result[0] == "hello world" # Mock response +class TestPackedSideChannelBodies: + """Tests for response parsing with packed side channels.""" + + async def _post_packed(self, client, mock_servers, path: str, **kwargs): + return await client._get_generate_client()._post( + f"{mock_servers['proxy_url']}{path}", json={}, packed_side_channels=True, **kwargs + ) + + @pytest.mark.asyncio + async def test_splices_both_registered_fields(self, client, mock_servers): + body = await self._post_packed(client, mock_servers, "/test/packed_body?two_blobs=true") + choice = body["choices"][0] + + assert isinstance(choice[PackedField.ROUTED_EXPERTS][PackedArrayKey.DATA], memoryview) + assert isinstance(choice[PackedField.ROLLOUT_SAMPLE_SUPPORT][PackedArrayKey.DATA], memoryview) + assert np.array_equal(decode_packed_routed_experts(choice[PackedField.ROUTED_EXPERTS]), _ROUTES) + support, _ = unpack_ndarray(choice[PackedField.ROLLOUT_SAMPLE_SUPPORT], allowed_dtypes=_SUPPORT_DTYPES, ndim=2) + assert np.array_equal(support, _SUPPORT) + + @pytest.mark.asyncio + async def test_drifted_layout_raises_without_retrying(self, client, mock_servers): + await client._post(f"{mock_servers['proxy_url']}/test/reset_packed_body_calls", json={}) + + with pytest.raises(ValueError, match="layout drifted"): + await self._post_packed(client, mock_servers, "/test/drifted_packed_body") + + async with httpx.AsyncClient() as http: + counts = (await http.get(f"{mock_servers['proxy_url']}/test/packed_body_calls")).json() + assert counts["drifted"] == 1 + + @pytest.mark.asyncio + async def test_undecodable_body_is_retried_then_spliced(self, client, mock_servers): + await client._post(f"{mock_servers['proxy_url']}/test/reset_packed_body_calls", json={}) + + body = await self._post_packed(client, mock_servers, "/test/flaky_packed_body") + + async with httpx.AsyncClient() as http: + counts = (await http.get(f"{mock_servers['proxy_url']}/test/packed_body_calls")).json() + assert counts["flaky"] == 2 + assert np.array_equal( + decode_packed_routed_experts(body["choices"][0][PackedField.ROUTED_EXPERTS]), + _ROUTES, + ) + + @pytest.mark.asyncio + async def test_client_error_with_non_json_body_surfaces_the_text(self, client, mock_servers): + with pytest.raises(aiohttp.ClientResponseError, match="prompt too long"): + await client._post(f"{mock_servers['proxy_url']}/test/bad_request_text", json={}) + + @pytest.mark.asyncio + async def test_non_routed_expert_generate_never_scans_the_body(self, client, monkeypatch): + def fail(*args, **kwargs): + raise AssertionError("the non-R3 path must not scan the response body") + + monkeypatch.setattr(remote_client_module, "load_packed_body", fail) + + result = await client.generate({"prompt_token_ids": [[1, 2, 3]]}) + assert len(result["responses"]) == 1 + + @pytest.mark.asyncio + async def test_routed_expert_generate_goes_through_the_splice(self, mock_servers, monkeypatch): + calls: List[int] = [] + original = remote_client_module.load_packed_body + + def counted(raw, **kwargs): + calls.append(len(raw)) + return original(raw, **kwargs) + + monkeypatch.setattr(remote_client_module, "load_packed_body", counted) + client = RemoteInferenceClient( + proxy_url=mock_servers["proxy_url"], + server_urls=mock_servers["server_urls"], + data_parallel_size=1, + enable_return_routed_experts=True, + ) + try: + result = await client.generate({"prompt_token_ids": [[1, 2, 3]]}) + finally: + await client.teardown() + + assert len(calls) == 1 + assert np.array_equal(result["rollout_expert_indices"][0], _ROUTES) + + class TestControlPlane: """Test control plane methods (fan-out to all servers).""" @@ -1025,8 +1274,7 @@ async def test_async_context_manager(self, mock_servers): result = await client.resume() assert len(result) == 2 - # Session should be closed after exiting context - assert client._session is None or client._session.closed + assert client._generate_client is None or client._generate_client._session is None async def _get_lora_registries(server_urls: List[str]) -> List[Dict[str, str]]: diff --git a/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py b/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py new file mode 100644 index 0000000000..f7f44b5102 --- /dev/null +++ b/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py @@ -0,0 +1,139 @@ +"""Sample-support capture out of vLLM's flat-logprobs rows.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +pytest.importorskip("vllm") + +from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( + PackedField, + decode_packed_sample_support, +) +from skyrl.backends.skyrl_train.inference_servers.vllm_server_actor import ( + VLLMServerActor, + _sample_support_from_flat_logprobs, +) + +pytestmark = pytest.mark.vllm + + +def test_flat_logprobs_extracts_sampled_scores_and_support_rows(): + flat_logprobs = SimpleNamespace( + token_ids=[7, 7, 8, 9, 4, 3, 4, 5], + logprobs=[-0.1, -0.1, -0.2, -0.3, -0.4, -0.2, -0.4, -0.6], + ) + + sampled, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=3) + + assert sampled == [{"logprob": -0.1}, {"logprob": -0.4}] + assert support.dtype == np.int32 + np.testing.assert_array_equal(support, [[7, 8, 9], [3, 4, 5]]) + + +def test_flat_logprobs_replaces_top_p_masked_candidates(): + flat_logprobs = SimpleNamespace( + token_ids=[7, 7, 8, 9], + logprobs=[-0.1, -0.1, -0.2, float("-inf")], + ) + + _, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=3) + + np.testing.assert_array_equal(support, [[7, 8, -1]]) + + +def test_flat_logprobs_repairs_sampled_token_absent_from_support(): + top_k = 3 + flat_logprobs = SimpleNamespace( + token_ids=[100, 8, 9, 10, 7, 7, 8, 9, 5, 6, 7, 8], + logprobs=[ + -0.1, + -0.2, + -0.3, + -0.4, + -0.1, + -0.1, + -0.2, + -0.3, + -0.4, + -0.5, + -0.6, + float("-inf"), + ], + ) + + _, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=top_k) + np.testing.assert_array_equal(support, [[8, 9, 100], [7, 8, 9], [6, 5, -1]]) + + +def test_flat_logprobs_top_k_one_repairs_single_support_column(): + flat_logprobs = SimpleNamespace( + token_ids=[42, 9], + logprobs=[-0.1, -0.2], + ) + + _, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=1) + + np.testing.assert_array_equal(support, [[42]]) + + +class FakeEngine: + sampling_params = None + + async def generate(self, prompt, sampling_params, request_id): + self.sampling_params = sampling_params + yield SimpleNamespace( + outputs=[ + SimpleNamespace( + token_ids=[7], + finish_reason="stop", + logprobs=SimpleNamespace( + token_ids=[7, 7, 8], + logprobs=[-0.1, -0.1, -0.2], + ), + routed_experts=None, + ) + ] + ) + + +@pytest.mark.parametrize("sampling_params", [{"temperature": 1.0}, {"temperature": 0.0, "top_k": -1}, {"top_k": 1}]) +def test_skyrl_generate_rejects_sample_support_without_a_bounded_support(sampling_params): + app = FastAPI() + engine = FakeEngine() + VLLMServerActor._add_custom_endpoints(app, engine, SimpleNamespace(enable_lora=False)) + + with TestClient(app) as client: + response = client.post( + "/skyrl/v1/generate", + json={"token_ids": [1, 2], "sampling_params": sampling_params, "return_sample_support": True}, + ) + + assert response.status_code == 400 + assert "top_k > 1" in response.json()["detail"] + assert engine.sampling_params is None + + +def test_skyrl_generate_returns_packed_sample_support(): + app = FastAPI() + engine = FakeEngine() + VLLMServerActor._add_custom_endpoints(app, engine, SimpleNamespace(enable_lora=False)) + + with TestClient(app) as client: + response = client.post( + "/skyrl/v1/generate", + json={ + "token_ids": [1, 2], + "sampling_params": {"temperature": 1.0, "top_k": 2}, + "return_sample_support": True, + }, + ) + + assert response.status_code == 200 + assert engine.sampling_params.flat_logprobs is True + assert engine.sampling_params.logprobs == 2 + packed = response.json()["choices"][0][PackedField.ROLLOUT_SAMPLE_SUPPORT] + np.testing.assert_array_equal(decode_packed_sample_support(packed), [[7, 8]]) diff --git a/tests/backends/skyrl_train/test_token_based_batching_utils.py b/tests/backends/skyrl_train/test_token_based_batching_utils.py index 4fac8a6679..c04b204eb1 100644 --- a/tests/backends/skyrl_train/test_token_based_batching_utils.py +++ b/tests/backends/skyrl_train/test_token_based_batching_utils.py @@ -12,6 +12,15 @@ import torch from skyrl.backends.skyrl_train.training_batch import TensorList, TrainingInputBatch +from skyrl.backends.skyrl_train.utils.packed_tensor import ( + PackedTensor, + cu_seqlens_from_lengths, +) +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_FIELD, + SAMPLE_SUPPORT_PADDING, + SAMPLE_SUPPORT_TORCH_DTYPE, +) from skyrl.backends.skyrl_train.workers.worker_utils import ( TokenBasedBatchIterator, get_microbatch_iterator, @@ -199,18 +208,73 @@ def test_padding_microbatch_matches_seq_len(self): # Padding rows must not contribute to the loss. assert padding["loss_mask"].sum().item() == 0 + def _add_packed_side_channels(self, batch: TrainingInputBatch) -> None: + """Attach both packed side channels: routes over real tokens, support over responses.""" + batch["rollout_expert_indices"] = PackedTensor( + torch.full((8, 2, 3), 7, dtype=torch.int16), + cu_seqlens_from_lengths([4, 4]), + ) + batch["router_padding_mask"] = torch.zeros((2, 4), dtype=torch.bool) + batch[SAMPLE_SUPPORT_FIELD] = PackedTensor( + torch.full((4, 5), 11, dtype=SAMPLE_SUPPORT_TORCH_DTYPE), + cu_seqlens_from_lengths([2, 2]), + ) + def test_padding_microbatch_uses_unique_dummy_routes(self): batch = self._make_batch([4, 4], num_actions=2) - batch["rollout_expert_indices"] = torch.full((2, 4, 2, 3), 7, dtype=torch.int16) - batch["router_padding_mask"] = torch.zeros((2, 4), dtype=torch.bool) + self._add_packed_side_channels(batch) iterator = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=8) padding = iterator._create_padding_microbatch() - expected = torch.tensor([0, 1, 2], dtype=torch.int16).expand_as(padding["rollout_expert_indices"]) - assert torch.equal(padding["rollout_expert_indices"], expected) + padded_routes = padding["rollout_expert_indices"] + assert padded_routes.sequence_lengths.tolist() == [1] + expected = torch.tensor([0, 1, 2], dtype=torch.int16).expand_as(padded_routes.values) + assert torch.equal(padded_routes.values, expected) assert torch.all(padding["router_padding_mask"]) + def test_padding_microbatch_sample_support_holds_no_response_rows(self): + """A dummy row attends one token but generates no response.""" + batch = self._make_batch([4, 4], num_actions=2) + self._add_packed_side_channels(batch) + iterator = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=8) + + padding = iterator._create_padding_microbatch() + + padded_support = padding[SAMPLE_SUPPORT_FIELD] + assert len(padded_support) == 1 + assert padded_support.sequence_lengths.tolist() == [0] + assert padded_support.values.shape == (0, 5) + assert padded_support.dtype == SAMPLE_SUPPORT_TORCH_DTYPE + assert padding["rollout_expert_indices"].sequence_lengths.tolist() == [1] + + def test_microbatch_selection_gathers_packed_sample_support_segments(self): + batch = self._make_batch([4, 2], num_actions=2) + batch[SAMPLE_SUPPORT_FIELD] = PackedTensor.from_segments( + [ + torch.full((2, 5), 1, dtype=SAMPLE_SUPPORT_TORCH_DTYPE), + torch.full((1, 5), SAMPLE_SUPPORT_PADDING, dtype=SAMPLE_SUPPORT_TORCH_DTYPE), + ] + ) + + microbatch = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=8)._create_microbatch_from_indices([1]) + + support = microbatch[SAMPLE_SUPPORT_FIELD] + assert support.sequence_lengths.tolist() == [1] + assert torch.all(support.segment(0) == SAMPLE_SUPPORT_PADDING) + + def test_microbatch_selection_gathers_packed_route_segments(self): + batch = self._make_batch([4, 2], num_actions=2) + batch["rollout_expert_indices"] = PackedTensor.from_segments( + [torch.full((4, 2, 3), 1, dtype=torch.int16), torch.full((2, 2, 3), 2, dtype=torch.int16)] + ) + + microbatch = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=8)._create_microbatch_from_indices([1]) + + routes = microbatch["rollout_expert_indices"] + assert routes.sequence_lengths.tolist() == [2] + assert torch.equal(routes.segment(0), torch.full((2, 2, 3), 2, dtype=torch.int16)) + def test_multimodal_tensorlist_microbatching(self): """Token-based microbatching must gather TensorList fields (multi-modal pixel_values / image_grid_thw) via the same index gather used for regular tensors.""" diff --git a/tests/backends/skyrl_train/test_train_batch.py b/tests/backends/skyrl_train/test_train_batch.py index e562f5ed48..291dfba5d5 100644 --- a/tests/backends/skyrl_train/test_train_batch.py +++ b/tests/backends/skyrl_train/test_train_batch.py @@ -1,4 +1,6 @@ import pickle +from collections.abc import Callable +from typing import get_args import numpy as np import pytest @@ -6,12 +8,30 @@ import torch from skyrl.backends.skyrl_train.training_batch import ( + PACKED_FIELD_PADDING, + BatchField, TensorBatch, + TensorFormat, TensorList, TrainingInput, TrainingInputBatch, + _deserialize_tensor, + _serialize_tensor, + append_packed_field_padding, + make_packed_field_padding, + packed_dummy_row_segments, pad_training_input_batch, ) +from skyrl.backends.skyrl_train.utils.packed_tensor import ( + PackedTensor, + cu_seqlens_from_lengths, +) +from skyrl.backends.skyrl_train.utils.routed_experts import ROUTED_EXPERT_DTYPES +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_FIELD, + SAMPLE_SUPPORT_PADDING, + SAMPLE_SUPPORT_TORCH_DTYPE, +) def test_train_batch_initialization(): @@ -552,6 +572,7 @@ def test_tensor_batch_none_tensor_list(): "rollout_logprobs", "rollout_expert_indices", "router_padding_mask", + SAMPLE_SUPPORT_FIELD, "pixel_values", "image_grid_thw", } @@ -576,8 +597,17 @@ def _make_full_training_batch(batch_size: int = 4, seq_len: int = 5) -> Training "kl": torch.randn(batch_size, seq_len), "rewards": torch.randn(batch_size, seq_len), "rollout_logprobs": torch.randn(batch_size, seq_len), - "rollout_expert_indices": torch.randint(0, 8, (batch_size, seq_len, 2, 3), dtype=torch.long), + # The fixture is fully attended, so each route segment has seq_len rows. + "rollout_expert_indices": PackedTensor( + torch.randint(0, 8, (batch_size * seq_len, 2, 3), dtype=torch.long), + cu_seqlens_from_lengths([seq_len] * batch_size), + ), "router_padding_mask": torch.zeros((batch_size, seq_len), dtype=torch.bool), + # Support packs to response tokens; this fixture's response spans the whole row. + SAMPLE_SUPPORT_FIELD: PackedTensor( + torch.randint(0, 1000, (batch_size * seq_len, 4), dtype=SAMPLE_SUPPORT_TORCH_DTYPE), + cu_seqlens_from_lengths([seq_len] * batch_size), + ), "pixel_values": TensorList([torch.randn(i + 1, 3) for i in range(batch_size)]), # batch_size * (i + 1) * 3 "image_grid_thw": TensorList([torch.tensor([[1, 2, 3]]) for _ in range(batch_size)]), # batch_size * 1 * 3 } @@ -640,14 +670,23 @@ def test_pad_batch_all_fields(): # padding rows are copies of row 0. assert torch.equal(padded["router_padding_mask"][:batch_size], batch["router_padding_mask"]) assert torch.all(padded["router_padding_mask"][batch_size:]) - assert torch.equal(padded["rollout_expert_indices"][:batch_size], batch["rollout_expert_indices"]) - expected_routes = torch.tensor([0, 1, 2]).expand_as(padded["rollout_expert_indices"][batch_size:]) - assert torch.equal(padded["rollout_expert_indices"][batch_size:], expected_routes) + assert padded["rollout_expert_indices"][:batch_size] == batch["rollout_expert_indices"] + padded_routes = padded["rollout_expert_indices"][batch_size:] + assert padded_routes.sequence_lengths.tolist() == [seq_len] * pad_size + expected_routes = torch.tensor([0, 1, 2]).expand_as(padded_routes.values) + assert torch.equal(padded_routes.values, expected_routes) + + padded_support = padded[SAMPLE_SUPPORT_FIELD][batch_size:] + assert padded[SAMPLE_SUPPORT_FIELD][:batch_size] == batch[SAMPLE_SUPPORT_FIELD] + # Support is indexed over response tokens, and a padded row copies row 0's response. + assert padded_support.sequence_lengths.tolist() == [seq_len] * pad_size + assert torch.all(padded_support.values == SAMPLE_SUPPORT_PADDING) regular_tensor_keys = EXPECTED_TRAINING_INPUT_FIELDS - { "loss_mask", "rollout_expert_indices", "router_padding_mask", + SAMPLE_SUPPORT_FIELD, "pixel_values", "image_grid_thw", } @@ -711,3 +750,205 @@ def test_pad_batch_preserves_none_fields(): padded = pad_training_input_batch(batch, pad_size=1) assert padded["values"] is None assert padded.batch_size == 4 + + +def test_packed_tensor_field_survives_the_ray_pickle_round_trip(): + """Packed route buffers cross to the workers via TensorBatch's custom pickle path.""" + segment_lengths = [3, 1, 4] + routes = PackedTensor( + torch.randint(0, 128, (sum(segment_lengths), 2, 3), dtype=torch.int16), + cu_seqlens_from_lengths(segment_lengths), + ) + data = TensorBatch( + { + "sequences": torch.randn(len(segment_lengths), 4), + "rollout_expert_indices": routes, + } + ) + data.metadata = {"response_length": 4} + + unpickled = pickle.loads(pickle.dumps(data)) + + restored = unpickled["rollout_expert_indices"] + assert isinstance(restored, PackedTensor) + assert restored.values.dtype == torch.int16 + assert restored.cu_seqlens.dtype == routes.cu_seqlens.dtype + assert restored == routes + assert unpickled == data + + +def test_serialized_field_formats_are_stable(): + data = TensorBatch( + { + "sequences": torch.randn(2, 4), + "pixel_values": TensorList([torch.randn(1, 3), torch.randn(2, 3)]), + "rollout_expert_indices": PackedTensor( + torch.zeros((3, 2, 3), dtype=torch.int16), cu_seqlens_from_lengths([2, 1]) + ), + "bf16_logprobs": torch.randn(2, 4, dtype=torch.bfloat16), + } + ) + + state = data.__getstate__()["batch_dict"] + + assert state["sequences"]["format"] == TensorFormat.NUMPY + assert state["bf16_logprobs"]["format"] == TensorFormat.TORCH + assert state["pixel_values"]["format"] == TensorFormat.TENSOR_LIST + assert state["rollout_expert_indices"]["format"] == TensorFormat.PACKED_TENSOR + + +ROUTE_KEY = "rollout_expert_indices" +_ZERO_COPY_SEGMENT_LENGTHS = [512, 256, 256] +_ZERO_COPY_BATCH_SIZE = len(_ZERO_COPY_SEGMENT_LENGTHS) + +# One test payload per opted-in field. +_ZERO_COPY_PAYLOADS: dict[str, Callable[[], BatchField]] = { + ROUTE_KEY: lambda: PackedTensor( + torch.randint(0, 64, (sum(_ZERO_COPY_SEGMENT_LENGTHS), 2, 3), dtype=torch.int16), + cu_seqlens_from_lengths(_ZERO_COPY_SEGMENT_LENGTHS), + ), + SAMPLE_SUPPORT_FIELD: lambda: PackedTensor( + torch.randint(0, 32_000, (sum(_ZERO_COPY_SEGMENT_LENGTHS), 20), dtype=SAMPLE_SUPPORT_TORCH_DTYPE), + cu_seqlens_from_lengths(_ZERO_COPY_SEGMENT_LENGTHS), + ), +} + + +def _zero_copy_buffer(value: BatchField) -> torch.Tensor: + """Return the payload buffer for a zero-copy field.""" + return value.values if isinstance(value, PackedTensor) else value + + +def test_zero_copy_keys_are_declared_and_have_payloads(): + assert TensorBatch.ZERO_COPY_KEYS <= set(TrainingInput.__annotations__) + assert set(_ZERO_COPY_PAYLOADS) == TensorBatch.ZERO_COPY_KEYS + + +@pytest.mark.parametrize("key", sorted(TensorBatch.ZERO_COPY_KEYS)) +def test_zero_copy_field_travels_out_of_band(key, oob_round_trip): + """Only opted-in fields travel out of band.""" + field = _ZERO_COPY_PAYLOADS[key]() + sequences = torch.randint(0, 100, (_ZERO_COPY_BATCH_SIZE, 4)) + batch = TrainingInputBatch({"sequences": sequences, key: field}) + batch.metadata = {"info": "zero-copy"} + buffer = _zero_copy_buffer(field) + + unpickled, payload, views = oob_round_trip(batch) + + assert [view.nbytes for view in views] == [buffer.nbytes], f"{key} is the only out-of-band field" + assert len(payload) < buffer.nbytes, f"{key} must not ALSO sit in the pickle stream" + assert np.shares_memory(_zero_copy_buffer(unpickled[key]).numpy(), buffer.numpy()) + assert not np.shares_memory(unpickled["sequences"].numpy(), sequences.numpy()) + assert unpickled == batch + + +@pytest.mark.parametrize("key", sorted(TensorBatch.ZERO_COPY_KEYS)) +def test_zero_copy_field_tolerates_read_only_plasma_buffer(key, oob_round_trip): + """Zero-copy fields preserve read-only buffers while ordinary fields stay writable.""" + field = _ZERO_COPY_PAYLOADS[key]() + advantages = torch.randn(_ZERO_COPY_BATCH_SIZE, 256) + batch = TrainingInputBatch({"advantages": advantages, key: field}) + batch.metadata = {} + + unpickled, _, views = oob_round_trip(batch, read_only=True) + + assert [view.readonly for view in views] == [True], f"{key} is the only out-of-band field" + restored = _zero_copy_buffer(unpickled[key]).numpy() + assert np.shares_memory(restored, np.frombuffer(views[0], dtype=restored.dtype)), "read-only buffer was copied" + assert unpickled[key] == field + + got = unpickled["advantages"] + assert got.numpy().flags.writeable + got.add_(1.0) + assert torch.allclose(got, advantages + 1.0) + assert torch.allclose(batch["advantages"], advantages), "source must not be aliased" + + +def test_packed_zero_copy_field_ships_only_its_values_buffer(): + """Only a packed field's values use an out-of-band buffer.""" + batch = TrainingInputBatch({ROUTE_KEY: _ZERO_COPY_PAYLOADS[ROUTE_KEY]()}) + + state = batch.__getstate__()["batch_dict"][ROUTE_KEY] + + assert state["format"] == TensorFormat.PACKED_TENSOR + assert state["values"]["format"] == TensorFormat.NUMPY_VIEW + assert state["cu_seqlens"]["format"] == TensorFormat.NUMPY + + +@pytest.mark.parametrize("dtype", sorted(ROUTED_EXPERT_DTYPES, key=str)) +def test_zero_copy_envelope_round_trips_every_route_dtype(dtype): + values = torch.from_numpy(np.arange(48, dtype=dtype).reshape(8, 2, 3)) + + envelope = _serialize_tensor(values, zero_copy=True) + restored = _deserialize_tensor(envelope) + + assert envelope["format"] == TensorFormat.NUMPY_VIEW + assert restored.dtype == values.dtype + assert torch.equal(restored, values) + + +def test_zero_copy_falls_back_for_bfloat16(): + """The numpy TypeError guard must still run before the zero-copy branch.""" + values = torch.randn(3, 4, dtype=torch.bfloat16) + + assert _serialize_tensor(values, zero_copy=True)["format"] == TensorFormat.TORCH + + +# ── packed field padding ───────────────────────────────────────────────────── + +# Expected route and sampler-support padding rows. +_PACKED_PADDING_EXPECTED_ROW: dict[str, Callable[[PackedTensor], torch.Tensor]] = { + ROUTE_KEY: lambda field: torch.arange(field.row_shape[-1], dtype=field.dtype), + SAMPLE_SUPPORT_FIELD: lambda field: torch.full(field.row_shape, SAMPLE_SUPPORT_PADDING, dtype=field.dtype), +} + + +def test_every_packed_training_input_field_has_a_padding_rule(): + """A packed field with no rule cannot be padded, and would raise mid-training-step.""" + packed_fields = { + name for name, annotation in TrainingInput.__annotations__.items() if PackedTensor in get_args(annotation) + } + assert packed_fields == set(PACKED_FIELD_PADDING) + + +@pytest.mark.parametrize("key", sorted(PACKED_FIELD_PADDING)) +def test_packed_field_padding_carries_that_fields_own_fill(key): + field = _ZERO_COPY_PAYLOADS[key]() + segment_lengths = [1, 3] + + padding = make_packed_field_padding(key, field, segment_lengths=segment_lengths) + + assert padding.sequence_lengths.tolist() == segment_lengths + assert padding.row_shape == field.row_shape + assert padding.dtype == field.dtype + assert torch.equal(padding.values, _PACKED_PADDING_EXPECTED_ROW[key](field).expand_as(padding.values)) + + +@pytest.mark.parametrize("key", sorted(PACKED_FIELD_PADDING)) +@pytest.mark.parametrize("pad_count", [1, 3]) +def test_appending_packed_field_padding_keeps_the_real_segments(key, pad_count): + field = _ZERO_COPY_PAYLOADS[key]() + + padded = append_packed_field_padding(key, field, segment_lengths=[2] * pad_count) + + assert padded.sequence_lengths.tolist() == _ZERO_COPY_SEGMENT_LENGTHS + [2] * pad_count + assert padded[: len(field)] == field + appended = padded[len(field) :] + assert torch.equal(appended.values, _PACKED_PADDING_EXPECTED_ROW[key](field).expand_as(appended.values)) + + +@pytest.mark.parametrize(("key", "rows_per_dummy_row"), [(ROUTE_KEY, 1), (SAMPLE_SUPPORT_FIELD, 0)]) +def test_dummy_row_segments_cover_the_single_attended_token(key, rows_per_dummy_row): + field = _ZERO_COPY_PAYLOADS[key]() + + segments = packed_dummy_row_segments(key, 3) + + assert segments == [rows_per_dummy_row] * 3 + padding = make_packed_field_padding(key, field, segment_lengths=segments) + assert len(padding) == 3 + assert padding.values.shape[0] == rows_per_dummy_row * 3 + + +def test_packed_field_padding_refuses_an_unregistered_field(): + with pytest.raises(ValueError, match="no padding rule"): + make_packed_field_padding("unregistered", _ZERO_COPY_PAYLOADS[ROUTE_KEY](), segment_lengths=[1]) diff --git a/tests/backends/skyrl_train/utils/test_packed_tensor.py b/tests/backends/skyrl_train/utils/test_packed_tensor.py new file mode 100644 index 0000000000..f0101cb23a --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_packed_tensor.py @@ -0,0 +1,240 @@ +"""Tests for ``PackedTensor`` batch operations.""" + +import pytest +import torch + +from skyrl.backends.skyrl_train.utils.packed_tensor import ( + CU_SEQLENS_DTYPE, + PackedTensor, + cu_seqlens_from_lengths, + lengths_from_offsets, + row_index_from_offsets, +) + +SEGMENT_LENGTHS = [3, 1, 4, 2] + + +def _segments(lengths=SEGMENT_LENGTHS, *, row_shape=(2, 3)) -> list[torch.Tensor]: + """Distinct rows per segment so any misplacement is visible.""" + segments = [] + next_value = 0 + for length in lengths: + size = length * torch.Size(row_shape).numel() + segments.append(torch.arange(next_value, next_value + size, dtype=torch.int16).reshape(length, *row_shape)) + next_value += size + return segments + + +@pytest.mark.parametrize( + ("lengths", "expected_offsets"), + [ + (SEGMENT_LENGTHS, [0, 3, 4, 8, 10]), + ([0, 2, 0], [0, 0, 2, 2]), + ], +) +def test_offsets_round_trip_lengths_in_int32(lengths, expected_offsets): + offsets = cu_seqlens_from_lengths(lengths) + + assert offsets.tolist() == expected_offsets + assert offsets.dtype == CU_SEQLENS_DTYPE + assert lengths_from_offsets(offsets).tolist() == lengths + assert lengths_from_offsets(offsets).dtype == CU_SEQLENS_DTYPE + + +def test_cu_seqlens_reject_negative_lengths_and_extra_dimensions(): + with pytest.raises(ValueError, match="non-negative"): + cu_seqlens_from_lengths([3, -1]) + with pytest.raises(ValueError, match="must be 1-D"): + cu_seqlens_from_lengths(torch.zeros((2, 2), dtype=torch.int32)) + + +def test_row_index_from_offsets_lays_selected_segments_back_to_back(): + starts = torch.tensor([8, 0]) + lengths = torch.tensor([2, 3]) + + assert row_index_from_offsets(starts, lengths).tolist() == [8, 9, 0, 1, 2] + + +def test_from_segments_round_trips_every_segment(): + segments = _segments() + packed = PackedTensor.from_segments(segments) + + assert len(packed) == len(segments) + assert packed.values.shape == (sum(SEGMENT_LENGTHS), 2, 3) + assert packed.sequence_lengths.tolist() == SEGMENT_LENGTHS + assert packed.row_shape == torch.Size((2, 3)) + assert packed.dtype == torch.int16 + assert packed.device == segments[0].device + for index, segment in enumerate(segments): + assert torch.equal(packed.segment(index), segment) + + +def test_from_segments_rejects_an_empty_batch(): + with pytest.raises(ValueError, match="empty list of segments"): + PackedTensor.from_segments([]) + + +def test_negative_and_out_of_range_segment_indices(): + packed = PackedTensor.from_segments(_segments()) + + assert torch.equal(packed.segment(-1), packed.segment(len(packed) - 1)) + with pytest.raises(IndexError, match="out of range"): + packed.segment(len(packed)) + with pytest.raises(IndexError, match="out of range"): + packed.segment(-len(packed) - 1) + + +def test_integer_index_returns_that_segment(): + segments = _segments() + packed = PackedTensor.from_segments(segments) + + assert torch.equal(packed[2], segments[2]) + assert torch.equal(packed[torch.tensor(2)], segments[2]) + + +@pytest.mark.parametrize("bounds", [(0, 4), (1, 3), (2, 4)]) +def test_contiguous_slice_selects_the_same_segments(bounds): + segments = _segments() + packed = PackedTensor.from_segments(segments) + start, stop = bounds + + sliced = packed[start:stop] + + assert len(sliced) == stop - start + assert sliced.cu_seqlens[0] == 0 + for offset, segment in enumerate(segments[start:stop]): + assert torch.equal(sliced.segment(offset), segment) + + +@pytest.mark.parametrize("indices", [[2, 0], [3, 3, 1], [0, 1, 2, 3], [1]]) +def test_gather_selects_segments_in_the_requested_order(indices): + segments = _segments() + packed = PackedTensor.from_segments(segments) + + for gathered in (packed[torch.tensor(indices)], packed[indices], packed[tuple(indices)]): + assert gathered.sequence_lengths.tolist() == [SEGMENT_LENGTHS[index] for index in indices] + for position, index in enumerate(indices): + assert torch.equal(gathered.segment(position), segments[index]) + + +def test_empty_slice_is_rejected_like_an_empty_tensor_list(): + """``TensorBatch`` fields cannot hold zero batch entries, so neither can a slice.""" + packed = PackedTensor.from_segments(_segments()) + + with pytest.raises(ValueError, match="at least two offsets"): + packed[2:2] + + +def test_strided_slice_falls_back_to_a_gather(): + segments = _segments() + packed = PackedTensor.from_segments(segments) + + strided = packed[::2] + + assert len(strided) == 2 + assert torch.equal(strided.segment(0), segments[0]) + assert torch.equal(strided.segment(1), segments[2]) + + +def test_cat_joins_batches_end_to_end(): + left = PackedTensor.from_segments(_segments([3, 1])) + right = PackedTensor.from_segments(_segments([4, 2])) + + joined = PackedTensor.cat([left, right]) + + assert joined.sequence_lengths.tolist() == [3, 1, 4, 2] + assert torch.equal(joined.segment(0), left.segment(0)) + assert torch.equal(joined.segment(2), right.segment(0)) + + +def test_cat_rejects_an_empty_list(): + with pytest.raises(ValueError, match="empty list of packed batches"): + PackedTensor.cat([]) + + +def test_repeat_tiles_and_repeat_interleave_duplicates(): + packed = PackedTensor.from_segments(_segments([3, 1])) + + tiled = packed.repeat(2) + interleaved = packed.repeat_interleave(2) + + assert tiled.sequence_lengths.tolist() == [3, 1, 3, 1] + assert interleaved.sequence_lengths.tolist() == [3, 3, 1, 1] + assert torch.equal(tiled.segment(2), packed.segment(0)) + assert torch.equal(interleaved.segment(1), packed.segment(0)) + + +def test_to_and_contiguous_preserve_the_batch(): + packed = PackedTensor.from_segments(_segments()) + + widened = packed.to(dtype=torch.int32) + + assert widened.dtype == torch.int32 + assert widened.cu_seqlens.dtype == CU_SEQLENS_DTYPE + assert torch.equal(widened.values, packed.values.to(torch.int32)) + assert packed.contiguous() == packed + + +def test_equality_compares_values_and_offsets(): + packed = PackedTensor.from_segments(_segments()) + + assert packed == PackedTensor.from_segments(_segments()) + assert packed != PackedTensor.from_segments(_segments([3, 1, 4, 2], row_shape=(1, 3))) + assert packed != PackedTensor.from_segments(_segments([4, 4, 2])) + assert packed != packed.values + + +def test_rejects_mismatched_device_or_offset_dtype(): + values = torch.zeros((4, 2), dtype=torch.int16) + + with pytest.raises(ValueError, match="must be torch.int32"): + PackedTensor(values, torch.tensor([0, 4], dtype=torch.int64)) + with pytest.raises(ValueError, match="at least two offsets"): + PackedTensor(values, torch.tensor([0], dtype=CU_SEQLENS_DTYPE)) + with pytest.raises(ValueError, match="at least two offsets"): + PackedTensor(values, torch.zeros((2, 2), dtype=CU_SEQLENS_DTYPE)) + + +def test_rejects_offsets_that_do_not_span_the_buffer(): + values = torch.zeros((4, 2), dtype=torch.int16) + + with pytest.raises(ValueError, match="must run from 0"): + PackedTensor(values, torch.tensor([0, 3], dtype=CU_SEQLENS_DTYPE)) + with pytest.raises(ValueError, match="must run from 0"): + PackedTensor(values, torch.tensor([1, 4], dtype=CU_SEQLENS_DTYPE)) + + +def test_rejects_values_without_a_token_row_dimension(): + with pytest.raises(ValueError, match="token-row dimension"): + PackedTensor(torch.tensor(1), torch.tensor([0, 1], dtype=CU_SEQLENS_DTYPE)) + + +def test_segments_and_contiguous_slices_are_views(): + """Reading a batch entry must not copy: alignment walks every segment per micro-batch.""" + packed = PackedTensor.from_segments(_segments()) + + assert packed.segment(1).data_ptr() == packed.values[3].data_ptr() + assert packed[1:3].values.data_ptr() == packed.values[3].data_ptr() + + +@pytest.mark.parametrize( + "operation", + [ + lambda packed: packed[torch.tensor([1, 0])], + lambda packed: packed[::2], + lambda packed: packed.repeat(2), + lambda packed: packed.repeat_interleave(2), + lambda packed: PackedTensor.cat([packed, packed]), + lambda packed: packed.to(dtype=torch.int32), + ], + ids=["gather", "strided_slice", "repeat", "repeat_interleave", "cat", "to_dtype"], +) +def test_reordering_operations_allocate_rather_than_alias(operation): + """A duplicated or reordered segment must own its rows; the source stays untouched.""" + packed = PackedTensor.from_segments(_segments()) + original = packed.values.clone() + + produced = operation(packed) + produced.values[:] = -1 + + assert torch.equal(packed.values, original) diff --git a/tests/backends/skyrl_train/utils/test_replay_utils.py b/tests/backends/skyrl_train/utils/test_replay_utils.py index 304e11c625..5320da509d 100644 --- a/tests/backends/skyrl_train/utils/test_replay_utils.py +++ b/tests/backends/skyrl_train/utils/test_replay_utils.py @@ -3,7 +3,6 @@ import types from types import SimpleNamespace -import numpy as np import pytest import torch @@ -11,10 +10,13 @@ build_token_metadata_layout, ) from skyrl.backends.skyrl_train.utils import replay_utils -from skyrl.backends.skyrl_train.utils.replay_utils import ( - make_replay_padding_indices, - make_replay_padding_indices_np, -) +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor +from skyrl.backends.skyrl_train.utils.replay_utils import make_replay_padding_indices + + +def _pack_routes(routes: torch.Tensor, attention_mask: torch.Tensor) -> PackedTensor: + """Pack a ``[batch, seq_len, layers, topk]`` fixture to its real tokens.""" + return PackedTensor.from_segments([routes[row][attention_mask[row].bool()] for row in range(routes.shape[0])]) @pytest.fixture @@ -70,22 +72,10 @@ def test_replay_padding_indices_are_unique(dtype): assert torch.equal(padding, torch.tensor([0, 1, 2], dtype=dtype).expand_as(padding)) -@pytest.mark.parametrize("dtype", [np.uint8, np.int16, np.int32]) -def test_numpy_replay_padding_matches_torch(dtype): - torch_dtype = getattr(torch, np.dtype(dtype).name) - padding = make_replay_padding_indices_np((2, 3, 4, 3), dtype=np.dtype(dtype)) - - assert padding.dtype == dtype - assert torch.equal( - torch.from_numpy(padding), - make_replay_padding_indices((2, 3, 4, 3), dtype=torch_dtype), - ) - - @pytest.mark.parametrize("shape", [(), (2, 3, 4, 0)]) -def test_numpy_replay_padding_rejects_missing_topk(shape): +def test_replay_padding_rejects_missing_topk(shape): with pytest.raises(ValueError, match="positive topk"): - make_replay_padding_indices_np(shape, dtype=np.dtype(np.uint8)) + make_replay_padding_indices(shape, dtype=torch.uint8) def test_replay_has_no_dispatcher_specific_patch(): @@ -121,15 +111,14 @@ class RouterReplayAction: "scatter_router_padding_mask_for_model", lambda mask, model, model_config: mask, ) - apply_layout = replay_utils.align_token_metadata + apply_layout = replay_utils.align_packed_token_metadata routed_layer_counts = [] def record_routed_layer_count(metadata, layout, padding_value): - if metadata.ndim == 4: - routed_layer_counts.append(metadata.shape[2]) + routed_layer_counts.append(metadata.row_shape[0]) return apply_layout(metadata, layout, padding_value) - monkeypatch.setattr(replay_utils, "align_token_metadata", record_routed_layer_count) + monkeypatch.setattr(replay_utils, "align_packed_token_metadata", record_routed_layer_count) routes = torch.tensor( [ @@ -152,7 +141,7 @@ def record_routed_layer_count(metadata, layout, padding_value): ) model_kwargs = replay_utils.setup_per_microbatch_replay_forward( - routes, + _pack_routes(routes, attention_mask), router_padding_mask, attention_mask, model=object(), @@ -215,7 +204,7 @@ def run(routes): fp8_enabled=False, ) replay_utils.setup_per_microbatch_replay_forward( - routes, + _pack_routes(routes, attention_mask), router_padding_mask, attention_mask, model=object(), diff --git a/tests/backends/skyrl_train/utils/test_routed_experts.py b/tests/backends/skyrl_train/utils/test_routed_experts.py index 432d9d31c9..fd059eab9b 100644 --- a/tests/backends/skyrl_train/utils/test_routed_experts.py +++ b/tests/backends/skyrl_train/utils/test_routed_experts.py @@ -2,6 +2,7 @@ import pytest from skyrl.backends.skyrl_train.utils.routed_experts import ( + RoutedExpertTrace, compact_routed_expert_indices, ) @@ -61,3 +62,37 @@ def test_compaction_rejects_nested_lists(): def test_compaction_rejects_invalid_routes(routes): with pytest.raises(ValueError): compact_routed_expert_indices(routes) + + +def _turn_routes(num_rows): + return np.arange(num_rows * 4, dtype=np.int16).reshape(num_rows, 2, 2) % 8 + + +def test_trace_returns_only_the_rows_it_captured(): + """The row count identifies where capture stops.""" + trace = RoutedExpertTrace() + trace.record_generation(prompt_token_count=3, generated_token_count=2, routed_experts=_turn_routes(4)) + trace.record_generation(prompt_token_count=7, generated_token_count=2, routed_experts=_turn_routes(4)) + + routes = trace.finalize(token_count=10, loss_mask=[0, 0, 0, 1, 1, 0, 0, 1, 1, 0]) + + assert routes.shape == (8, 2, 2) + assert np.array_equal(routes, np.concatenate((_turn_routes(4), _turn_routes(4)))) + + +def test_trace_keeps_full_coverage_when_every_token_has_a_route(): + trace = RoutedExpertTrace() + trace.record_generation(prompt_token_count=3, generated_token_count=2, routed_experts=_turn_routes(4)) + + routes = trace.finalize(token_count=4, loss_mask=[0, 0, 0, 1]) + + assert routes.shape == (4, 2, 2) + + +def test_trace_rejects_an_uncaptured_loss_active_target(): + """Every loss-active target must have a captured route.""" + trace = RoutedExpertTrace() + trace.record_generation(prompt_token_count=3, generated_token_count=2, routed_experts=_turn_routes(4)) + + with pytest.raises(ValueError, match="missing routed-expert row for loss-active target at token 5"): + trace.finalize(token_count=6, loss_mask=[0, 0, 0, 1, 1, 1]) diff --git a/tests/backends/skyrl_train/utils/test_sample_support.py b/tests/backends/skyrl_train/utils/test_sample_support.py new file mode 100644 index 0000000000..fa1344957f --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_sample_support.py @@ -0,0 +1,50 @@ +import numpy as np +import pytest + +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_DTYPE, + SAMPLE_SUPPORT_PADDING, + SampleSupportTrace, +) + + +def _rows(count: int, first_id: int = 0) -> np.ndarray: + return np.array([[first_id + i, first_id + i + 100] for i in range(count)], dtype=SAMPLE_SUPPORT_DTYPE) + + +def test_finalize_drops_exactly_the_declared_trailing_rows(): + trace = SampleSupportTrace() + trace.append(_rows(3), expected_rows=3) + trace.append_padding(2) + + support = trace.finalize(token_count=3, extra_rows=2) + + np.testing.assert_array_equal(support, _rows(3)) + + +def test_finalize_rejects_a_trace_shorter_than_the_response(): + trace = SampleSupportTrace() + trace.append(_rows(2), expected_rows=2) + + with pytest.raises(ValueError, match="2 rows for 3 tokens plus 0 trailing rows"): + trace.finalize(token_count=3, extra_rows=0) + + +def test_finalize_rejects_an_unexpected_overshoot(): + trace = SampleSupportTrace() + trace.append(_rows(3), expected_rows=3) + trace.append_padding(2) + + with pytest.raises(ValueError, match="5 rows for 3 tokens plus 1 trailing rows"): + trace.finalize(token_count=3, extra_rows=1) + + +def test_padding_rows_are_all_padding_sentinels(): + trace = SampleSupportTrace() + trace.append(_rows(1), expected_rows=1) + trace.append_padding(2) + + support = trace.finalize(token_count=3, extra_rows=0) + + assert support.dtype == SAMPLE_SUPPORT_DTYPE + np.testing.assert_array_equal(support[1:], np.full((2, 2), SAMPLE_SUPPORT_PADDING, dtype=SAMPLE_SUPPORT_DTYPE)) diff --git a/tests/backends/skyrl_train/utils/test_sample_support_row_ids.py b/tests/backends/skyrl_train/utils/test_sample_support_row_ids.py new file mode 100644 index 0000000000..105b938679 --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_sample_support_row_ids.py @@ -0,0 +1,138 @@ +"""Row-id derivation for packed sampler support.""" + +from typing import List, Tuple + +import pytest +import torch + +from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( + TokenMetadataLayout, +) +from skyrl.backends.skyrl_train.training_batch import TrainingInputBatch +from skyrl.backends.skyrl_train.utils.packed_tensor import ( + PackedTensor, + cu_seqlens_from_lengths, +) +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_FIELD, + SAMPLE_SUPPORT_NO_ROW, + SAMPLE_SUPPORT_TORCH_DTYPE, + align_sample_support_row_ids, +) + +TOP_K = 3 +# Anti-correlated lengths exercise different support offsets. +LENGTHS: List[Tuple[int, int]] = [(2, 3), (4, 2)] + + +def _attention_mask(lengths: List[Tuple[int, int]]) -> torch.Tensor: + """Left-padded mask, as ``convert_prompts_responses_to_batch_tensors`` builds it.""" + totals = [prompt + response for prompt, response in lengths] + mask = torch.zeros((len(totals), max(totals)), dtype=torch.bool) + for row, total in enumerate(totals): + mask[row, max(totals) - total :] = True + return mask + + +def _layout(lengths: List[Tuple[int, int]], *, packed: bool = False, align: int = 1) -> TokenMetadataLayout: + mask = _attention_mask(lengths) + totals = [prompt + response for prompt, response in lengths] + if not packed: + return TokenMetadataLayout( + attention_mask=mask, + sequence_lengths=totals, + aligned_sequence_length=mask.shape[1], + ) + padded = [total + (-total % align) for total in totals] + return TokenMetadataLayout( + attention_mask=mask, + sequence_lengths=totals, + aligned_sequence_length=sum(padded), + padded_sequence_lengths=padded, + cu_seqlens_padded=torch.tensor([0, *torch.tensor(padded).cumsum(0).tolist()], dtype=torch.int32), + ) + + +def _support(lengths: List[Tuple[int, int]]) -> PackedTensor: + """Distinct support rows, so a gather by id pins the exact row it landed on.""" + response_lens = [response for _, response in lengths] + total_rows = sum(response_lens) + values = torch.arange(total_rows * TOP_K, dtype=SAMPLE_SUPPORT_TORCH_DTYPE).reshape(total_rows, TOP_K) + return PackedTensor(values, cu_seqlens_from_lengths(response_lens)) + + +def test_row_ids_land_on_the_positions_that_predict_response_tokens(): + """Support includes the last prompt position and excludes the last response position.""" + row_ids = align_sample_support_row_ids(_support(LENGTHS), _layout(LENGTHS)) + + assert row_ids.dtype == torch.int64 + # Trajectory 0 is p=2, r=3, so ids 0..2 sit at real positions 1..3; trajectory 1 is p=4, r=2. + assert row_ids.tolist() == [ + [SAMPLE_SUPPORT_NO_ROW, 0, 1, 2, SAMPLE_SUPPORT_NO_ROW, SAMPLE_SUPPORT_NO_ROW], + [SAMPLE_SUPPORT_NO_ROW, SAMPLE_SUPPORT_NO_ROW, SAMPLE_SUPPORT_NO_ROW, 3, 4, SAMPLE_SUPPORT_NO_ROW], + ] + + +def test_row_ids_gather_the_support_rows_of_each_response_token(): + support = _support(LENGTHS) + + row_ids = align_sample_support_row_ids(support, _layout(LENGTHS)) + + scored = row_ids >= 0 + gathered = support.values[row_ids[scored]] + assert torch.equal(gathered, support.values) + assert scored.sum(dim=1).tolist() == [response for _, response in LENGTHS] + + +@pytest.mark.parametrize("align", [1, 4]) +def test_row_ids_follow_megatron_packed_padding(align): + """Sequence packing interleaves each sequence with its alignment padding.""" + row_ids = align_sample_support_row_ids(_support(LENGTHS), _layout(LENGTHS, packed=True, align=align)) + + padded_lengths = [total + (-total % align) for total in (5, 6)] + assert row_ids.shape == (1, sum(padded_lengths)) + packed = row_ids.squeeze(0).tolist() + assert packed[1:4] == [0, 1, 2] + assert packed[padded_lengths[0] + 3 : padded_lengths[0] + 5] == [3, 4] + + +@pytest.mark.parametrize("selection", ["chunk", "slice"]) +def test_row_ids_rebase_under_batch_selection(selection): + support = _support(LENGTHS) + batch = TrainingInputBatch({"attention_mask": _attention_mask(LENGTHS).long(), SAMPLE_SUPPORT_FIELD: support}) + selected = batch.chunk(1)[1] if selection == "chunk" else batch.slice(1, 2) + + row_ids = align_sample_support_row_ids(selected[SAMPLE_SUPPORT_FIELD], _layout(LENGTHS[1:])) + + assert row_ids[row_ids >= 0].tolist() == [0, 1] + assert torch.equal(selected[SAMPLE_SUPPORT_FIELD].values[row_ids[row_ids >= 0]], support.segment(1)) + + +def test_row_ids_accept_an_empty_padding_segment(): + """A synthetic batch row attends one token and generates nothing, so it scores no position.""" + support = PackedTensor(torch.empty((0, TOP_K), dtype=SAMPLE_SUPPORT_TORCH_DTYPE), cu_seqlens_from_lengths([0])) + mask = torch.zeros((1, 4), dtype=torch.bool) + mask[0, 0] = True + layout = TokenMetadataLayout(attention_mask=mask, sequence_lengths=[1], aligned_sequence_length=4) + + row_ids = align_sample_support_row_ids(support, layout) + + assert torch.all(row_ids == SAMPLE_SUPPORT_NO_ROW) + + +def test_row_ids_reject_a_trajectory_with_no_prompt_token(): + """With no prompt there is no position whose logit predicts the first response token.""" + support = PackedTensor.from_segments([torch.zeros((3, TOP_K), dtype=SAMPLE_SUPPORT_TORCH_DTYPE)]) + layout = TokenMetadataLayout( + attention_mask=torch.ones((1, 3), dtype=torch.bool), + sequence_lengths=[3], + aligned_sequence_length=3, + ) + + with pytest.raises(ValueError, match="predicts its first response token"): + align_sample_support_row_ids(support, layout) + + +def test_row_ids_reject_a_segment_count_mismatch(): + with pytest.raises(ValueError, match="segments for"): + align_sample_support_row_ids(_support(LENGTHS), _layout(LENGTHS[1:])) diff --git a/tests/train/dataset/test_parallel_fill.py b/tests/train/dataset/test_parallel_fill.py new file mode 100644 index 0000000000..1b0e6daedd --- /dev/null +++ b/tests/train/dataset/test_parallel_fill.py @@ -0,0 +1,74 @@ +""" +uv run --isolated --extra dev pytest tests/train/dataset/test_parallel_fill.py +""" + +import threading + +import pytest + +from skyrl.train.dataset.parallel_fill import fill_batch_rows + + +@pytest.mark.parametrize("workers", [None, 1, 4, 64]) +def test_every_index_is_filled_exactly_once(workers): + num_rows = 32 + calls = [0] * num_rows + + def fill_row(index: int) -> None: + calls[index] += 1 + + fill_batch_rows(fill_row, num_rows, workers=workers) + + assert calls == [1] * num_rows + + +def test_single_worker_runs_serially_on_the_calling_thread(): + order = [] + threads = set() + + def fill_row(index: int) -> None: + order.append(index) + threads.add(threading.current_thread()) + + fill_batch_rows(fill_row, 4, workers=1) + + assert order == [0, 1, 2, 3] + assert threads == {threading.current_thread()} + + +def test_multiple_workers_leave_the_calling_thread(): + threads = set() + + def fill_row(index: int) -> None: + threads.add(threading.current_thread()) + + fill_batch_rows(fill_row, 8, workers=4) + + assert threading.current_thread() not in threads + + +def test_zero_rows_is_a_no_op(): + def fill_row(index: int) -> None: + raise AssertionError("fill_row must not be called for an empty batch") + + fill_batch_rows(fill_row, 0) + + +def test_negative_row_count_raises(): + with pytest.raises(ValueError, match="row count must be non-negative"): + fill_batch_rows(lambda index: None, -1) + + +def test_non_positive_worker_count_raises(): + with pytest.raises(ValueError, match="worker count must be positive"): + fill_batch_rows(lambda index: None, 4, workers=0) + + +@pytest.mark.parametrize("workers", [1, 4]) +def test_callback_exception_propagates(workers): + def fill_row(index: int) -> None: + if index == 2: + raise RuntimeError("row 2 failed") + + with pytest.raises(RuntimeError, match="row 2 failed"): + fill_batch_rows(fill_row, 8, workers=workers) diff --git a/tests/train/dataset/test_preprocess.py b/tests/train/dataset/test_preprocess.py index c5e885ebdd..b4fab55172 100644 --- a/tests/train/dataset/test_preprocess.py +++ b/tests/train/dataset/test_preprocess.py @@ -2,13 +2,23 @@ uv run --isolated --extra dev pytest tests/train/dataset/test_preprocess.py """ +import logging +from typing import List from unittest.mock import MagicMock import numpy as np import pytest import torch +from skyrl.backends.skyrl_train.utils.routed_experts import ROUTED_EXPERT_DTYPES +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_DTYPE, + SAMPLE_SUPPORT_PADDING, + SAMPLE_SUPPORT_TORCH_DTYPE, +) +from skyrl.train.dataset import parallel_fill from skyrl.train.dataset.preprocess import ( + ROUTED_EXPERT_TORCH_DTYPES, convert_prompts_responses_to_batch_tensors, make_router_padding_mask, ) @@ -85,7 +95,7 @@ def test_routed_expert_tensor_uses_unique_dummy_routes(tokenizer): ), ] - *_, routed = convert_prompts_responses_to_batch_tensors( + *_, routed, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts=[[10], [20]], responses=[[11, 12], [21, 22]], @@ -94,9 +104,10 @@ def test_routed_expert_tensor_uses_unique_dummy_routes(tokenizer): rollout_expert_indices=routes, ) - assert routed.shape == (2, 3, 2, 2) + assert routed.values.shape == (6, 2, 2) + assert routed.cu_seqlens.tolist() == [0, 3, 6] assert routed.dtype == torch.uint8 - assert routed[0, 2].tolist() == [[0, 1], [0, 1]] + assert routed.segment(0)[2].tolist() == [[0, 1], [0, 1]] @pytest.mark.parametrize( @@ -114,7 +125,7 @@ def test_routed_expert_tensor_promotes_mixed_batch_dtype( np.asarray([[[max_expert_id, max_expert_id + 1]]], dtype=source_dtype), ] - *_, routed = convert_prompts_responses_to_batch_tensors( + *_, routed, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts=[[10], [20]], responses=[[11], [21]], @@ -124,14 +135,14 @@ def test_routed_expert_tensor_promotes_mixed_batch_dtype( ) assert routed.dtype == expected_dtype - assert routed[1, 0].tolist() == [[max_expert_id, max_expert_id + 1]] + assert routed.segment(1)[0].tolist() == [[max_expert_id, max_expert_id + 1]] def test_routed_expert_tensor_accepts_read_only_arrays(tokenizer): routes = np.asarray([[[1, 2]], [[3, 4]]], dtype=np.uint8) routes.flags.writeable = False - *_, routed = convert_prompts_responses_to_batch_tensors( + *_, routed, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts=[[10]], responses=[[11]], @@ -141,7 +152,7 @@ def test_routed_expert_tensor_accepts_read_only_arrays(tokenizer): ) assert routed.dtype == torch.uint8 - assert routed.tolist() == [[[[1, 2]], [[3, 4]]]] + assert routed.segment(0).tolist() == [[[1, 2]], [[3, 4]]] def test_routed_expert_tensor_rejects_nested_lists(tokenizer): @@ -156,17 +167,13 @@ def test_routed_expert_tensor_rejects_nested_lists(tokenizer): ) -@pytest.mark.parametrize("dtype", [np.uint16, np.int64]) -def test_routed_expert_tensor_narrows_wide_dtypes(tokenizer, dtype): - """Wide dtypes are compacted, not rejected. +def test_routed_expert_tensor_accepts_non_contiguous_arrays(tokenizer): + # Every other expert column, which leaves a non-contiguous view. + base = np.asarray([[[1, 9, 2, 9]], [[3, 9, 4, 9]]], dtype=np.uint8) + routes = base[:, :, ::2] + assert not routes.flags.c_contiguous - The wire decoder already restricts routes to uint8/int16/int32, so this path - only sees a wide dtype from a hand-built generator -- narrowing it is more - useful than refusing it. - """ - routes = np.asarray([[[1, 2]], [[3, 4]]], dtype=dtype) - - *_, routed = convert_prompts_responses_to_batch_tensors( + *_, routed, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts=[[10]], responses=[[11]], @@ -176,15 +183,30 @@ def test_routed_expert_tensor_narrows_wide_dtypes(tokenizer, dtype): ) assert routed.dtype == torch.uint8 - assert routed.tolist() == [[[[1, 2]], [[3, 4]]]] + assert routed.segment(0).tolist() == [[[1, 2]], [[3, 4]]] -def test_routed_expert_tensor_retightens_after_truncation(tokenizer): - """A truncated array can fit a narrower dtype than the wire declared.""" - # int16 on the wire because of the trailing 300, which truncation then drops. +@pytest.mark.parametrize("dtype", [np.uint16, np.int64]) +def test_routed_expert_tensor_rejects_non_canonical_dtypes(tokenizer, dtype): + """The sender compacts to the canonical dtype, so collation validates instead of rescanning.""" + routes = np.asarray([[[1, 2]], [[3, 4]]], dtype=dtype) + + with pytest.raises(ValueError, match="canonical routed-expert dtype"): + convert_prompts_responses_to_batch_tensors( + tokenizer.pad_token_id, + prompts=[[10]], + responses=[[11]], + rewards=[[0.0]], + loss_masks=[[1]], + rollout_expert_indices=[routes], + ) + + +def test_routed_expert_tensor_keeps_the_sender_dtype_after_truncation(tokenizer): + # The dropped trailing value required int16 on the sender. routes = np.asarray([[[1, 2]], [[3, 4]], [[300, 5]]], dtype=np.int16) - *_, routed = convert_prompts_responses_to_batch_tensors( + *_, routed, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts=[[10]], responses=[[11]], @@ -193,8 +215,86 @@ def test_routed_expert_tensor_retightens_after_truncation(tokenizer): rollout_expert_indices=[routes[:2]], ) - assert routed.dtype == torch.uint8 - assert routed.tolist() == [[[[1, 2]], [[3, 4]]]] + assert routed.dtype == torch.int16 + assert routed.segment(0).tolist() == [[[1, 2]], [[3, 4]]] + + +@pytest.mark.parametrize( + ("dtype", "expert_id", "expect_warning"), + [(np.int16, 300, False), (np.int32, 2**16, True)], +) +def test_routed_expert_tensor_warns_only_on_an_int32_batch(tokenizer, caplog, dtype, expert_id, expect_warning): + routes = np.asarray([[[expert_id, expert_id + 1]]], dtype=dtype) + + with caplog.at_level(logging.WARNING, logger="skyrl.train.dataset.preprocess"): + *_, routed, _ = convert_prompts_responses_to_batch_tensors( + tokenizer.pad_token_id, + prompts=[[10]], + responses=[[11]], + rewards=[[0.0]], + loss_masks=[[1]], + rollout_expert_indices=[routes], + ) + + assert routed.dtype == ROUTED_EXPERT_TORCH_DTYPES[np.dtype(dtype)] + assert ("not compacting its routes" in caplog.text) is expect_warning + + +def test_routed_expert_torch_dtype_map_covers_the_canonical_dtypes(): + assert set(ROUTED_EXPERT_TORCH_DTYPES) == set(ROUTED_EXPERT_DTYPES) + + +def _numpy_padded_routes( + routes: List[np.ndarray], + prompts: List[List[int]], + responses: List[List[int]], +) -> np.ndarray: + """Reference NumPy implementation of route collation.""" + max_total = max(len(prompt) + len(response) for prompt, response in zip(prompts, responses)) + num_layers, topk = routes[0].shape[1:] + batch_dtype = max((sample.dtype for sample in routes), key=lambda dtype: dtype.itemsize) + padded = np.empty((len(routes), max_total, num_layers, topk), dtype=batch_dtype) + padded[...] = np.arange(topk, dtype=batch_dtype) + for index, sample in enumerate(routes): + left_pad = max_total - (len(prompts[index]) + len(responses[index])) + padded[index, left_pad : left_pad + sample.shape[0]] = sample + return padded + + +def test_routed_expert_tensor_is_bit_identical_to_numpy_collation(tokenizer): + """Packed collation matches the real rows of the padded NumPy reference.""" + prompts = [[1, 2], [3, 4, 5, 6]] + responses = [[10, 11, 12], [20, 21]] + num_layers, topk = 2, 3 + # Sample 0 has 5 tokens but only 4 captured route rows, so its segment pads at the end. + routes = [ + np.arange(4 * num_layers * topk, dtype=np.uint8).reshape(4, num_layers, topk), + (np.arange(6 * num_layers * topk, dtype=np.int16) + 300).reshape(6, num_layers, topk), + ] + + *_, routed, _ = convert_prompts_responses_to_batch_tensors( + tokenizer.pad_token_id, + prompts, + responses, + rewards=[[0.0] * 3, [0.0] * 2], + loss_masks=[[1] * 3, [1] * 2], + rollout_expert_indices=routes, + ) + + padded = _numpy_padded_routes(routes, prompts, responses) + real_rows = np.concatenate( + [ + padded[index, padded.shape[1] - (len(prompt) + len(response)) :] + for index, (prompt, response) in enumerate(zip(prompts, responses)) + ] + ) + assert routed.dtype == torch.int16 + assert routed.cu_seqlens.tolist() == [0, 5, 11] + assert torch.equal(routed.values, torch.from_numpy(real_rows)) + # Padding routes use distinct experts for Megatron's dropless dispatcher. + padding_row = [[0, 1, 2]] * num_layers + assert routed.segment(0)[4].tolist() == padding_row + assert not torch.equal(routed.segment(0)[4], torch.zeros_like(routed.segment(0)[4])) def test_convert_prompts_responses_to_batch_tensors_exact(tokenizer): @@ -215,7 +315,7 @@ def test_convert_prompts_responses_to_batch_tensors_exact(tokenizer): loss_masks = [[1, 1, 0], [1, 1, 1, 0, 0]] rewards = [torch.tensor([0, 1, 0]), torch.tensor([1, 0, 0, 0, 0])] - sequences, attention_mask, response_mask, ret_rewards, ret_loss_masks, ret_log_probs, _ = ( + sequences, attention_mask, response_mask, ret_rewards, ret_loss_masks, ret_log_probs, _, _ = ( convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, @@ -251,7 +351,7 @@ def test_convert_prompts_responses_to_batch_tensors_different_lengths(tokenizer) rewards = [torch.tensor([1.0, 0.5, 0.3]), torch.tensor([0.8])] loss_masks = [[1, 1, 1], [1]] - sequences, attention_mask, response_mask, ret_rewards, ret_loss_masks, ret_log_probs, _ = ( + sequences, attention_mask, response_mask, ret_rewards, ret_loss_masks, ret_log_probs, _, _ = ( convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, @@ -332,7 +432,7 @@ def test_unified_left_padding_layout(tokenizer): rewards = [[0.0] * 3, [0.0] * 2] loss_masks = [[1] * 3, [1] * 2] - seq, attn, action, rew, lm, _, _ = convert_prompts_responses_to_batch_tensors( + seq, attn, action, rew, lm, _, _, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, @@ -362,7 +462,7 @@ def test_right_aligned_response_data(tokenizer): prompts_copy = [p[:] for p in prompts] responses_copy = [r[:] for r in responses] - seq, attn, action, rew, lm, lp, _ = convert_prompts_responses_to_batch_tensors( + seq, attn, action, rew, lm, lp, _, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, @@ -397,7 +497,7 @@ def test_max_seq_len_warns_but_does_not_truncate(tokenizer): rewards = [[0.0] * 10, [0.0] * 50] loss_masks = [[1] * 10, [1] * 50] - seq, _, action, _, _, _, _ = convert_prompts_responses_to_batch_tensors( + seq, _, action, _, _, _, _, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, @@ -411,17 +511,8 @@ def test_max_seq_len_warns_but_does_not_truncate(tokenizer): assert action.shape == (2, 50) -# --------------------------------------------------------------------------- -# R3 (Router Replay) — rollout_expert_indices padding tests -# --------------------------------------------------------------------------- - - def test_rollout_expert_indices_shape_padding_and_alignment(tokenizer): - """rollout_expert_indices tensor should have shape [batch, max_total, layers, topk] - with left-padding aligned to the attention_mask.""" - # Sample 0: prompt=2, response=3 → total=5 - # Sample 1: prompt=4, response=2 → total=6 - # max_total=6 + """Routes pack to [sum(seq_len), layers, topk] with one cu_seqlens segment per trajectory.""" prompts = [[1, 2], [3, 4, 5, 6]] responses = [[10, 11, 12], [20, 21]] rewards = [[0.0] * 3, [0.0] * 2] @@ -429,12 +520,10 @@ def test_rollout_expert_indices_shape_padding_and_alignment(tokenizer): num_layers = 2 topk = 2 - # rollout_expert_indices[i] has shape [prompt_len_i + response_len_i, num_layers, topk] - # Sample 0: 5 tokens, sample 1: 6 tokens - rei_0 = np.asarray([[[1, 2]] * num_layers for _ in range(5)], dtype=np.uint8) # 5 tokens - rei_1 = np.asarray([[[3, 4]] * num_layers for _ in range(6)], dtype=np.uint8) # 6 tokens + rei_0 = np.asarray([[[1, 2]] * num_layers for _ in range(5)], dtype=np.uint8) + rei_1 = np.asarray([[[3, 4]] * num_layers for _ in range(6)], dtype=np.uint8) - seq, attn, action, rew, lm, lp, rei_tensor = convert_prompts_responses_to_batch_tensors( + seq, attn, action, rew, lm, lp, rei_tensor, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, @@ -444,24 +533,11 @@ def test_rollout_expert_indices_shape_padding_and_alignment(tokenizer): ) assert rei_tensor is not None - # Shape: [batch=2, max_total=6, layers=2, topk=2] - assert rei_tensor.shape == (2, 6, num_layers, topk) - - dummy_routes = [[0, 1]] * num_layers - # Sample 0 has total=5, so the first position uses unique dummy routes. - assert rei_tensor[0, 0].tolist() == dummy_routes - assert rei_tensor[0, 1].tolist() == [[1, 2]] * num_layers # first real token - - # Sample 1 has total=6, no padding - assert rei_tensor[1, 0].tolist() == [[3, 4]] * num_layers # first real token - - # Dummy positions in rei_tensor align exactly with attention_mask==0. - for i in range(2): - for pos in range(6): - if attn[i, pos] == 0: - assert rei_tensor[i, pos].tolist() == dummy_routes - else: - assert rei_tensor[i, pos].tolist() != dummy_routes + assert rei_tensor.values.shape == (11, num_layers, topk) + assert rei_tensor.cu_seqlens.tolist() == [0, 5, 11] + assert rei_tensor.sequence_lengths.tolist() == attn.sum(dim=1).tolist() + assert rei_tensor.segment(0).tolist() == [[[1, 2]] * num_layers] * 5 + assert rei_tensor.segment(1).tolist() == [[[3, 4]] * num_layers] * 6 def test_rollout_expert_indices_none_when_not_provided(tokenizer): @@ -471,7 +547,7 @@ def test_rollout_expert_indices_none_when_not_provided(tokenizer): rewards = [[0.0], [0.0]] loss_masks = [[1], [1]] - *_, rei_tensor = convert_prompts_responses_to_batch_tensors( + *_, rei_tensor, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, @@ -491,7 +567,7 @@ def test_stepwise_anti_correlation_no_inflation(tokenizer): rewards = [[0.0] * 90, [0.0] * 10] loss_masks = [[1] * 90, [1] * 10] - seq, attn, action, rew, lm, _, _ = convert_prompts_responses_to_batch_tensors( + seq, attn, action, rew, lm, _, _, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, @@ -508,3 +584,120 @@ def test_stepwise_anti_correlation_no_inflation(tokenizer): # Response data right-aligned: sample 1 has 10 tokens -> [0]*80 + [1]*10 assert action[1].tolist() == [0] * 80 + [1] * 10 + + +SAMPLE_SUPPORT_TOP_K = 3 +# Anti-correlated prompt/response lengths, so a prompt-region rectangle would be mostly filler. +SAMPLE_SUPPORT_LENGTHS = [(2, 3), (9, 2), (5, 5), (1, 1)] + + +def _make_sample_support(lengths: List[tuple], *, seed: int = 0) -> List[np.ndarray]: + """One dense ``[response_len, top_k]`` support block per trajectory, some rows padded.""" + rng = np.random.default_rng(seed) + support = [] + for _, response_len in lengths: + rows = rng.integers(0, 32_000, size=(response_len, SAMPLE_SUPPORT_TOP_K), dtype=np.int64) + # A padded row (nothing captured) and a partially-filled row, as the sampler emits. + rows[0] = SAMPLE_SUPPORT_PADDING + rows[-1, -1] = SAMPLE_SUPPORT_PADDING + support.append(rows.astype(SAMPLE_SUPPORT_DTYPE)) + return support + + +def _convert_with_support(tokenizer, lengths: List[tuple], support: List[np.ndarray]): + prompts = [list(range(1, prompt_len + 1)) for prompt_len, _ in lengths] + responses = [list(range(100, 100 + response_len)) for _, response_len in lengths] + return convert_prompts_responses_to_batch_tensors( + tokenizer.pad_token_id, + prompts, + responses, + rewards=[[0.0] * len(response) for response in responses], + loss_masks=[[1] * len(response) for response in responses], + rollout_sample_support=support, + ) + + +def test_sample_support_packs_to_the_response_tokens(tokenizer): + """One segment per trajectory, holding exactly its response tokens.""" + support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS) + + *_, packed = _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + response_lens = [response_len for _, response_len in SAMPLE_SUPPORT_LENGTHS] + assert packed.values.shape == (sum(response_lens), SAMPLE_SUPPORT_TOP_K) + assert packed.dtype == SAMPLE_SUPPORT_TORCH_DTYPE + assert packed.sequence_lengths.tolist() == response_lens + for index, rows in enumerate(support): + assert torch.equal(packed.segment(index), torch.from_numpy(rows)) + + +def test_sample_support_pooled_fill_equals_serial_fill(monkeypatch, tokenizer): + """Parallel and serial fills must be bit-identical.""" + support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS) + + monkeypatch.setattr(parallel_fill, "default_fill_workers", lambda: len(SAMPLE_SUPPORT_LENGTHS)) + *_, pooled = _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + monkeypatch.setattr(parallel_fill, "default_fill_workers", lambda: 1) + *_, serial = _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + assert pooled == serial + assert torch.equal(pooled.values, torch.from_numpy(np.concatenate(support, axis=0))) + + +def test_sample_support_accepts_read_only_arrays(tokenizer): + lengths = [(2, 2)] + rows = np.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=SAMPLE_SUPPORT_DTYPE) + rows.flags.writeable = False + + *_, packed = _convert_with_support(tokenizer, lengths, [rows[:2]]) + + assert packed.segment(0).tolist() == [[1, 2, 3], [4, 5, 6]] + + +def test_sample_support_rejects_a_row_count_that_is_not_the_response_length(tokenizer): + support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS) + support[1] = support[1][:1] + + with pytest.raises(ValueError, match="support rows for"): + _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + +def test_sample_support_rejects_a_ragged_width(tokenizer): + support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS) + support[2] = support[2][:, :-1] + + with pytest.raises(ValueError, match="must share top_k"): + _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + +def test_sample_support_rejects_nested_lists(tokenizer): + support = [rows.tolist() for rows in _make_sample_support(SAMPLE_SUPPORT_LENGTHS)] + + with pytest.raises(TypeError, match="NumPy arrays"): + _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + +@pytest.mark.parametrize("dtype", [np.int16, np.int64]) +def test_sample_support_rejects_non_canonical_dtypes(tokenizer, dtype): + support = [rows.astype(dtype) for rows in _make_sample_support(SAMPLE_SUPPORT_LENGTHS)] + + with pytest.raises(ValueError, match="canonical sample-support dtype"): + _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + +def test_sample_support_rejects_a_batch_size_mismatch(tokenizer): + support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS)[:-1] + + with pytest.raises(ValueError, match="support for every trajectory"): + _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + +def test_sample_support_none_when_not_provided(tokenizer): + *_, packed = convert_prompts_responses_to_batch_tensors( + tokenizer.pad_token_id, + prompts=[[1, 2]], + responses=[[10]], + rewards=[[0.0]], + loss_masks=[[1]], + ) + assert packed is None diff --git a/tests/train/generators/test_datatypes.py b/tests/train/generators/test_datatypes.py index 12c9efdba2..580cbd97f5 100644 --- a/tests/train/generators/test_datatypes.py +++ b/tests/train/generators/test_datatypes.py @@ -31,7 +31,6 @@ def test_turn_output(output_ids, observation_ids, output_logprobs, added_eos, ex output_logprobs=output_logprobs, new_obs=[], obs_ids=observation_ids, - rollout_expert_indices=None, added_eos=added_eos, reward=1.0, ) diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py index e3ee7a1f5c..fd568276c5 100644 --- a/tests/train/generators/test_generator_output_utils.py +++ b/tests/train/generators/test_generator_output_utils.py @@ -7,6 +7,7 @@ import numpy as np import pytest +from skyrl.backends.skyrl_train.utils.sample_support import SAMPLE_SUPPORT_DTYPE from skyrl.train.generators.base import GeneratorOutput, TrajectoryID from skyrl.train.generators.utils import ( compute_turn_token_counts, @@ -14,6 +15,7 @@ get_metrics_from_generator_output, get_rollout_metrics, merge_stepwise_output, + slice_generator_output, ) from skyrl.train.utils.utils import validate_cfg from tests.train.util import example_dummy_config @@ -30,6 +32,7 @@ def test_generator_output_concatenation(): "rollout_metrics", "rollout_logprobs", "rollout_expert_indices", + "rollout_sample_support", # optional but present in the signature "trajectory_ids", "trajectory_generation_times", @@ -52,6 +55,9 @@ def test_generator_output_concatenation(): "loss_masks": [[1, 1], [1, 1]], "stop_reasons": ["stop", "stop"], "rollout_logprobs": [[0.1, 0.2], [0.3, 0.4]], + # Routes cover every trained token. + "rollout_expert_indices": [np.zeros((3, 1, 2), dtype=np.uint8), np.ones((3, 1, 2), dtype=np.uint8)], + "rollout_sample_support": [[[1, 2], [1, 2]], [[3, 4], [3, 4]]], } generator_output_2: GeneratorOutput = { @@ -61,6 +67,8 @@ def test_generator_output_concatenation(): "loss_masks": [[1, 1, 1], [1]], "stop_reasons": ["stop", "stop"], "rollout_logprobs": [[0.5, 0.6, 0.7], [0.8]], + "rollout_expert_indices": [np.full((5, 1, 2), 2, dtype=np.uint8), np.full((1, 1, 2), 3, dtype=np.uint8)], + "rollout_sample_support": [[[5, 6], [5, 6], [5, 6]], [[7, 8]]], } generator_outputs = [generator_output_1, generator_output_2] @@ -72,6 +80,11 @@ def test_generator_output_concatenation(): assert concatenated_output["loss_masks"] == [[1, 1], [1, 1], [1, 1, 1], [1]] assert concatenated_output["stop_reasons"] == ["stop", "stop", "stop", "stop"] assert concatenated_output["rollout_logprobs"] == [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6, 0.7], [0.8]] + assert [rows[0] for rows in concatenated_output["rollout_sample_support"]] == [[1, 2], [3, 4], [5, 6], [7, 8]] + assert [int(routes.flat[0]) for routes in concatenated_output["rollout_expert_indices"]] == [0, 1, 2, 3] + reversed_output = concatenate_generator_outputs([generator_output_2, generator_output_1]) + assert [rows[0] for rows in reversed_output["rollout_sample_support"]] == [[5, 6], [7, 8], [1, 2], [3, 4]] + assert [int(routes.flat[0]) for routes in reversed_output["rollout_expert_indices"]] == [2, 3, 0, 1] # Validate rollout metrics expected_rollout_metrics = { @@ -95,6 +108,49 @@ def test_generator_output_concatenation(): np.testing.assert_allclose(concatenated_output["rollout_metrics"][key], value) +@pytest.mark.parametrize("side_channel", ["rollout_expert_indices", "rollout_sample_support"]) +def test_side_channel_concatenation_rejects_a_mix(side_channel): + def make_output(value) -> GeneratorOutput: + return { + "prompt_token_ids": [[1]], + "response_ids": [[10]], + "rewards": [1.0], + "loss_masks": [[1]], + "stop_reasons": ["stop"], + "rollout_logprobs": None, + side_channel: value, + } + + populated = make_output([[[1, 2]]] if side_channel == "rollout_sample_support" else [np.zeros((1, 1, 2), np.uint8)]) + missing = make_output(None) + for outputs in ([populated, missing], [missing, populated]): + with pytest.raises(ValueError, match=f"all have null {side_channel}"): + concatenate_generator_outputs(outputs) + + +def test_slice_generator_output_slices_each_component_of_a_dict_field(): + generator_output: GeneratorOutput = { + "prompt_token_ids": [[1], [2], [3]], + "response_ids": [[10], [20], [30]], + "rewards": [1.0, 2.0, 3.0], + "loss_masks": [[1], [1], [1]], + "stop_reasons": ["stop", "stop", "length"], + "rollout_metrics": {"generate/avg_num_tokens": 1.0}, + "rollout_logprobs": None, + "trajectory_generation_times": [10.0, 20.0, 30.0], + "trajectory_time_splits": {"llm": [1.0, 2.0, 3.0], "env": [4.0, 5.0, 6.0]}, + } + + sliced = slice_generator_output(generator_output, [2, 0]) + + assert sliced["trajectory_time_splits"] == {"llm": [3.0, 1.0], "env": [6.0, 4.0]} + assert sliced["response_ids"] == [[30], [10]] + assert sliced["trajectory_generation_times"] == [30.0, 10.0] + assert sliced["rollout_logprobs"] is None + assert sliced["rollout_metrics"] == {"generate/avg_num_tokens": 1.0} + assert "rollout_metrics" not in slice_generator_output(generator_output, [1], preserve_metrics=False) + + def test_time_split_rollout_metrics(): metrics = get_rollout_metrics( responses=[[1, 2]] * 4, @@ -528,6 +584,59 @@ def test_per_trajectory_scalar_rewards_and_overlong_filtering(self): else: assert merged["loss_masks"] == [[1, 0, 1]] + def test_sample_support_observation_deltas_keep_full_width_padding_rows(self): + tid = _make_tid("support") + gen_out: GeneratorOutput = { + "prompt_token_ids": [[10], [10, 20, 30]], + "response_ids": [[20], [40, 41]], + "rewards": [[1.0], [0.0, 5.0]], + "loss_masks": [[1], [1, 1]], + "stop_reasons": ["continue", "eos"], + "rollout_metrics": None, + "rollout_logprobs": None, + "rollout_sample_support": [ + np.array([[20, 21, -1]], dtype=SAMPLE_SUPPORT_DTYPE), + np.array([[40, 44, -1], [41, 45, 46]], dtype=SAMPLE_SUPPORT_DTYPE), + ], + "trajectory_ids": [tid, tid], + "rollout_expert_indices": None, + "is_last_step": [False, True], + } + + merged = merge_stepwise_output(gen_out) + + support = merged["rollout_sample_support"][0] + expected = [[20, 21, -1], [-1, -1, -1], [40, 44, -1], [41, 45, 46]] + np.testing.assert_array_equal(support, np.array(expected, dtype=SAMPLE_SUPPORT_DTYPE)) + assert support.shape == (len(merged["response_ids"][0]), 3) + assert support.dtype == SAMPLE_SUPPORT_DTYPE + + def test_native_output_carrying_dict_valued_time_splits(self): + tid = _make_tid("timed") + gen_out: GeneratorOutput = { + "prompt_token_ids": [[10], [10, 20, 30]], + "response_ids": [[20], [40]], + "rewards": [[0.0], [1.0]], + "loss_masks": [[1], [1]], + "stop_reasons": ["continue", "eos"], + "rollout_metrics": None, + "rollout_logprobs": None, + "trajectory_ids": [tid, tid], + "trajectory_generation_times": [5.0, 5.0], + "trajectory_time_splits": {"llm": [1.0, 2.0], "env": [3.0, 4.0]}, + "rollout_expert_indices": None, + "rollout_sample_support": None, + "is_last_step": [False, True], + "env_metrics": [{}, {}], + } + + merged = merge_stepwise_output(gen_out) + + assert merged["response_ids"] == [[20, 30, 40]] + assert merged["loss_masks"] == [[1, 0, 1]] + assert merged["rewards"] == [[0.0, 0.0, 1.0]] + assert merged["is_last_step"] == [True] + def test_no_logprobs_no_stop_reasons(self): """Works correctly when rollout_logprobs and stop_reasons are None.""" tid = _make_tid("no_lp") @@ -745,6 +854,15 @@ def test_validate_cfg_merge_stepwise_requires_step_wise(self): with pytest.raises(ValueError, match="merge_stepwise_output.*requires.*step_wise_trajectories"): validate_cfg(cfg) + @patch("skyrl.train.utils.utils.validate_batch_sizes", new=lambda cfg: None) + @patch("skyrl.train.utils.utils.validate_generator_cfg", new=lambda cfg: None) + def test_validate_cfg_refuses_step_wise_with_routed_expert_capture(self): + cfg = example_dummy_config() + cfg.generator.step_wise_trajectories = True + cfg.generator.inference_engine.enable_return_routed_experts = True + with pytest.raises(ValueError, match="first N prompt tokens"): + validate_cfg(cfg) + def test_compute_turn_token_counts(): """`compute_turn_token_counts` returns one count per turn (maximal run of non-zero loss-mask diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index 6d03d937b6..1d7e6f85d5 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -8,14 +8,25 @@ import numpy as np import pytest +from skyrl.backends.skyrl_train.utils.sample_support import ( + SAMPLE_SUPPORT_DTYPE, + SAMPLE_SUPPORT_PADDING, + SampleSupportTrace, +) from skyrl.train.config import ChatTemplateConfig, GeneratorConfig from skyrl.train.generators.base import ( + TRAINING_PHASE_EVAL, + TRAINING_PHASE_TRAIN, BatchMetadata, ConversationType, GeneratorInput, GeneratorOutput, ) -from skyrl.train.generators.skyrl_gym_generator import SkyRLGymGenerator, TurnOutput +from skyrl.train.generators.skyrl_gym_generator import ( + AgentLoopState, + SkyRLGymGenerator, + TurnOutput, +) from skyrl_gym.envs.base_text_env import BaseTextEnv, BaseTextEnvStepOutput # Mock constants, where 4 is the eos token id @@ -23,20 +34,22 @@ MOCK_TOKENIZER_ENCODED_IDS = [1, 2, 3, 4] -def test_turn_output_keeps_uncaptured_suffix_out_of_routes(): - routes = np.asarray([[[2, 3]], [[4, 5]]], dtype=np.uint8) +def test_turn_output_masks_uncaptured_suffix(): output = TurnOutput( output="answer", output_ids=[10, 11, 4], output_logprobs=None, new_obs=[], obs_ids=[20, 21], - rollout_expert_indices=routes, reward=1.0, + rollout_sample_support=np.array([[10, 100], [11, 101]], dtype=np.int32), added_eos=True, ) - assert output.get_turn_rollout_expert_indices() is routes + np.testing.assert_array_equal( + output.get_turn_rollout_sample_support(), + np.array([[10, 100], [11, 101], [-1, -1], [-1, -1], [-1, -1]], dtype=np.int32), + ) assert output.get_turn_loss_mask() == [1, 1, 0, 0, 0] @@ -413,6 +426,403 @@ def mock_generate(_, model=None): assert output.stop_reason == "stop" +@pytest.mark.asyncio +@patch("skyrl_gym.make") +async def test_agent_loop_uses_incremental_replay_metadata_traces( + mock_make, + mock_tokenizer, + mock_llm, + mock_env, + generator_cfg, + mock_env_cfg, +): + generator_cfg.batched = False + generator_cfg.max_turns = 2 + generator_cfg.use_conversation_multi_turn = True + generator_cfg.inference_engine.enable_return_routed_experts = True + generator_cfg.inference_engine.enable_return_sample_support_set = True + generator_cfg.sampling_params.top_k = 2 + mock_make.return_value = mock_env + mock_env.init.return_value = ([{"role": "user", "content": "Initial input"}], {}) + + mock_env.step.side_effect = [ + BaseTextEnvStepOutput(observations=[{"role": "user", "content": "next"}], reward=1.0, done=done, metadata={}) + for done in (False, True) + ] + prompt_starts = [] + generation_index = 0 + + def generate(input_batch, model=None): + nonlocal generation_index + assert input_batch["return_sample_support"] is True + prompt_tokens = input_batch["prompt_token_ids"][0] + prompt_start = input_batch["routed_experts_prompt_starts"][0] + prompt_starts.append(prompt_start) + output_ids = [10, 11] + num_route_rows = len(prompt_tokens) - prompt_start + len(output_ids) - 1 + routes = np.arange(num_route_rows * 4, dtype=np.int32).reshape(num_route_rows, 2, 2) % 8 + sample_support = np.array([[10, 100 + generation_index], [11, 110 + generation_index]], dtype=np.int32) + generation_index += 1 + return { + "responses": ["mocked output"], + "response_ids": [output_ids], + "stop_reasons": ["stop"], + "rollout_expert_indices": [routes], + "rollout_sample_support": [sample_support], + } + + mock_llm.generate = AsyncMock(side_effect=generate) + generator = SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + generator.base_conversation_token_ids = [] + + output = await generator.agent_loop( + [{"role": "user", "content": "Start"}], + mock_env_cfg.env_class, + {}, + max_tokens=32, + max_input_length=64, + ) + + assert prompt_starts == [0, 5] + support = output.rollout_sample_support + assert support.dtype == SAMPLE_SUPPORT_DTYPE + np.testing.assert_array_equal(support[:2], np.array([[10, 100], [11, 110]], dtype=SAMPLE_SUPPORT_DTYPE)) + np.testing.assert_array_equal(support[-2:], np.array([[10, 101], [11, 111]], dtype=SAMPLE_SUPPORT_DTYPE)) + assert np.all(support[2:-2] == SAMPLE_SUPPORT_PADDING) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("batched", [True, False]) +@pytest.mark.parametrize("batch_sampling_params", [{"temperature": 1.0, "top_k": 2, "max_tokens": 32}, None]) +@pytest.mark.parametrize("training_phase", [TRAINING_PHASE_TRAIN, TRAINING_PHASE_EVAL]) +@pytest.mark.parametrize("enable_capture", [True, False]) +@patch("skyrl_gym.make") +async def test_generate_requests_sample_support_capture_only_for_the_train_phase( + mock_make, + mock_tokenizer, + mock_llm, + mock_env, + generator_cfg, + mock_env_cfg, + enable_capture, + training_phase, + batch_sampling_params, + batched, +): + generator_cfg.batched = batched + generator_cfg.max_turns = 1 + generator_cfg.inference_engine.enable_return_sample_support_set = enable_capture + generator_cfg.sampling_params.top_k = 2 + mock_make.return_value = mock_env + mock_env.init.return_value = ([{"role": "user", "content": "Initial input"}], {}) + mock_env.step.side_effect = lambda x: BaseTextEnvStepOutput(observations=[], reward=1.0, done=True, metadata={}) + captured = {} + + def generate(input_batch, model=None): + captured.update(input_batch) + num_prompts = len(input_batch["prompt_token_ids"]) + return { + "responses": ["mocked output"] * num_prompts, + "response_ids": [[10, 11]] * num_prompts, + "stop_reasons": ["stop"] * num_prompts, + "rollout_sample_support": ( + [np.array([[10, 100], [11, 110]], dtype=np.int32)] * num_prompts + if input_batch["return_sample_support"] + else None + ), + } + + mock_llm.generate = AsyncMock(side_effect=generate) + generator = SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + generator.base_conversation_token_ids = [] + + input_batch: GeneratorInput = { + "prompts": [[{"role": "user", "content": "What is 3 + 5?"}]], + "env_classes": [mock_env_cfg.env_class], + "env_extras": [{"answer": "8"}], + "sampling_params": batch_sampling_params, + "batch_metadata": BatchMetadata(global_step=1, training_phase=training_phase), + } + + output = await generator.generate(input_batch) + + expected_capture = enable_capture and training_phase == TRAINING_PHASE_TRAIN + assert captured["return_sample_support"] is expected_capture + if expected_capture: + assert output["rollout_sample_support"] is not None + else: + assert output.get("rollout_sample_support", None) is None + + +def test_validate_cfg_refuses_routed_experts_without_conversation_multi_turn( + mock_tokenizer, mock_llm, generator_cfg, mock_env_cfg +): + generator_cfg.batched = False + generator_cfg.use_conversation_multi_turn = False + generator_cfg.inference_engine.enable_return_routed_experts = True + + with pytest.raises(ValueError, match="use_conversation_multi_turn=True"): + SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + + +def test_validate_cfg_refuses_routed_experts_with_step_wise_trajectories( + mock_tokenizer, mock_llm, generator_cfg, mock_env_cfg +): + generator_cfg.batched = False + generator_cfg.step_wise_trajectories = True + generator_cfg.inference_engine.enable_return_routed_experts = True + + with pytest.raises(ValueError, match="first N prompt tokens"): + SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + + +@pytest.mark.parametrize("capture_routed_experts,capture_sample_support", [(True, False), (False, True)]) +def test_validate_cfg_refuses_custom_chat_template_with_side_channel_capture( + capture_routed_experts, capture_sample_support, mock_tokenizer, mock_llm, generator_cfg, mock_env_cfg +): + generator_cfg.batched = False + generator_cfg.chat_template = ChatTemplateConfig(source="name", name_or_path="qwen3_without_thinking") + generator_cfg.inference_engine.enable_return_routed_experts = capture_routed_experts + generator_cfg.inference_engine.enable_return_sample_support_set = capture_sample_support + + with pytest.raises(ValueError, match="custom chat template"): + SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + + +def test_retokenizing_state_update_refuses_a_live_side_channel_trace( + mock_tokenizer, mock_llm, generator_cfg, mock_env_cfg +): + generator_cfg.batched = False + generator_cfg.chat_template = ChatTemplateConfig(source="name", name_or_path="qwen3_without_thinking") + generator = SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + state = AgentLoopState( + chat_history=[{"role": "user", "content": "hi"}], + input_ids=[1, 2], + loss_mask=[], + rollout_logprobs=None, + response_end_idx=None, + done=False, + sample_support_trace=SampleSupportTrace(), + ) + turn_output = TurnOutput( + output="answer", + output_ids=[10, 4], + output_logprobs=None, + new_obs=[], + obs_ids=[], + reward=1.0, + ) + + with pytest.raises(NotImplementedError, match="does not feed the per-token side-channel traces"): + generator._update_agent_state_by_retokenizing_chat_history(state, turn_output) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("batched", [True, False]) +@pytest.mark.parametrize("training_phase", [TRAINING_PHASE_TRAIN, TRAINING_PHASE_EVAL]) +@patch("skyrl_gym.make") +async def test_generate_retains_routed_experts_only_for_the_train_phase( + mock_make, + mock_tokenizer, + mock_llm, + mock_env, + generator_cfg, + mock_env_cfg, + training_phase, + batched, +): + generator_cfg.batched = batched + generator_cfg.max_turns = 1 + generator_cfg.use_conversation_multi_turn = True + generator_cfg.inference_engine.enable_return_routed_experts = True + mock_make.return_value = mock_env + mock_env.init.return_value = ([{"role": "user", "content": "Initial input"}], {}) + mock_env.step.side_effect = lambda x: BaseTextEnvStepOutput(observations=[], reward=1.0, done=True, metadata={}) + captured = {} + + def generate(input_batch, model=None): + captured.update(input_batch) + num_prompts = len(input_batch["prompt_token_ids"]) + prompt_len = len(input_batch["prompt_token_ids"][0]) + return { + "responses": ["mocked output"] * num_prompts, + "response_ids": [[10, 11]] * num_prompts, + "stop_reasons": ["stop"] * num_prompts, + "rollout_expert_indices": [ + np.arange((prompt_len + 1) * 4, dtype=np.int32).reshape(prompt_len + 1, 2, 2) % 8 + ] + * num_prompts, + } + + mock_llm.generate = AsyncMock(side_effect=generate) + generator = SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + generator.base_conversation_token_ids = [] + + input_batch: GeneratorInput = { + "prompts": [[{"role": "user", "content": "What is 3 + 5?"}]], + "env_classes": [mock_env_cfg.env_class], + "env_extras": [{"answer": "8"}], + "sampling_params": None, + "batch_metadata": BatchMetadata(global_step=1, training_phase=training_phase), + } + + output = await generator.generate(input_batch) + + if training_phase == TRAINING_PHASE_TRAIN: + assert output["rollout_expert_indices"] is not None + assert output["rollout_expert_indices"][0] is not None + if not batched: + assert captured["routed_experts_prompt_starts"] == [0] + else: + assert output["rollout_expert_indices"] is None + if not batched: + assert captured["routed_experts_prompt_starts"] is None + + +@pytest.mark.asyncio +@patch("skyrl_gym.make") +async def test_agent_loop_keeps_the_generated_eos_support_row_in_single_turn_mode( + mock_make, + mock_tokenizer, + mock_llm, + mock_env, + generator_cfg, + mock_env_cfg, +): + generator_cfg.batched = False + generator_cfg.max_turns = 1 + generator_cfg.use_conversation_multi_turn = False + generator_cfg.inference_engine.enable_return_sample_support_set = True + generator_cfg.sampling_params.top_k = 2 + mock_make.return_value = mock_env + mock_env.init.return_value = ([{"role": "user", "content": "Initial input"}], {}) + mock_env.step.side_effect = [ + BaseTextEnvStepOutput(observations=[], reward=1.0, done=True, metadata={}), + ] + eos_support_row = [12, 112] + + def generate(input_batch, model=None): + return { + "responses": ["mocked output"], + "response_ids": [[10, 11, 4]], + "stop_reasons": ["stop"], + "rollout_sample_support": [np.array([[10, 110], [11, 111], eos_support_row], dtype=np.int32)], + } + + mock_llm.generate = AsyncMock(side_effect=generate) + generator = SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + generator.base_conversation_token_ids = [] + + output = await generator.agent_loop( + [{"role": "user", "content": "Start"}], + mock_env_cfg.env_class, + {}, + max_tokens=32, + max_input_length=64, + ) + + assert output.response_ids == [10, 11, 4] + assert output.loss_mask == [1, 1, 1] + np.testing.assert_array_equal( + output.rollout_sample_support, + np.array([[10, 110], [11, 111], eos_support_row], dtype=SAMPLE_SUPPORT_DTYPE), + ) + + +@pytest.mark.asyncio +@patch("skyrl_gym.make") +async def test_agent_loop_pads_a_stop_string_eos_support_row_in_single_turn_mode( + mock_make, + mock_tokenizer, + mock_llm, + mock_env, + generator_cfg, + mock_env_cfg, +): + generator_cfg.batched = False + generator_cfg.max_turns = 1 + generator_cfg.use_conversation_multi_turn = False + generator_cfg.inference_engine.enable_return_sample_support_set = True + generator_cfg.sampling_params.top_k = 2 + mock_make.return_value = mock_env + mock_env.init.return_value = ([{"role": "user", "content": "Initial input"}], {}) + mock_env.step.side_effect = [ + BaseTextEnvStepOutput(observations=[], reward=1.0, done=True, metadata={}), + ] + + def generate(input_batch, model=None): + return { + "responses": ["mocked output"], + "response_ids": [[10, 11]], + "stop_reasons": ["stop"], + "rollout_sample_support": [np.array([[10, 110], [11, 111]], dtype=np.int32)], + } + + mock_llm.generate = AsyncMock(side_effect=generate) + generator = SkyRLGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=mock_env_cfg, + inference_engine_client=mock_llm, + tokenizer=mock_tokenizer, + ) + generator.base_conversation_token_ids = [] + + output = await generator.agent_loop( + [{"role": "user", "content": "Start"}], + mock_env_cfg.env_class, + {}, + max_tokens=32, + max_input_length=64, + ) + + assert output.response_ids == [10, 11, 4] + padding_row = [SAMPLE_SUPPORT_PADDING, SAMPLE_SUPPORT_PADDING] + np.testing.assert_array_equal( + output.rollout_sample_support, + np.array([[10, 110], [11, 111], padding_row], dtype=SAMPLE_SUPPORT_DTYPE), + ) + + @pytest.mark.asyncio @patch("skyrl_gym.make") async def test_generate_batched(mock_make, mock_tokenizer, mock_llm, mock_env, generator_cfg, mock_env_cfg): diff --git a/tests/train/generators/test_skyrl_vlm_generator.py b/tests/train/generators/test_skyrl_vlm_generator.py index 0cd504e302..ad3105da61 100644 --- a/tests/train/generators/test_skyrl_vlm_generator.py +++ b/tests/train/generators/test_skyrl_vlm_generator.py @@ -134,6 +134,30 @@ async def mock_generate(input_batch, model=None): # --------------------------------------------------------------------------- +def test_vlm_validate_cfg_still_applies_the_base_refusals(): + tokenizer = MagicMock() + tokenizer.apply_chat_template.side_effect = lambda messages, **kwargs: [1, 2, 3, 4] + tokenizer.eos_token_id = 4 + generator_cfg = GeneratorConfig( + sampling_params=SamplingParams(max_generate_length=200, logprobs=None), + max_input_length=4096, + batched=False, + max_turns=3, + use_conversation_multi_turn=True, + chat_template=ChatTemplateConfig(source="name", name_or_path="qwen3_without_thinking"), + step_wise_trajectories=False, + ) + generator_cfg.inference_engine.enable_return_sample_support_set = True + + with pytest.raises(ValueError, match="custom chat template"): + SkyRLVLMGymGenerator( + generator_cfg=generator_cfg, + skyrl_gym_cfg=SkyRLGymConfig(max_env_workers=0), + inference_engine_client=MagicMock(), + tokenizer=tokenizer, + ) + + @pytest.mark.asyncio @patch("skyrl.train.generators.skyrl_vlm_generator.decode_mm_kwargs") async def test_vlm_obs_offset(mock_decode): diff --git a/tests/train/test_collation_vectorization_equivalence.py b/tests/train/test_collation_vectorization_equivalence.py index af83a424e7..5926b68b6e 100644 --- a/tests/train/test_collation_vectorization_equivalence.py +++ b/tests/train/test_collation_vectorization_equivalence.py @@ -153,7 +153,7 @@ def test_rl_preprocess_bit_identical(seed, with_logprobs): tokenizer = MagicMock() tokenizer.pad_token_id = 0 - seq, attn, action, rew, lm, lp, _ = convert_prompts_responses_to_batch_tensors( + seq, attn, action, rew, lm, lp, _, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, rewards, loss_masks, logprobs ) r_seq, r_attn, r_action, r_rew, r_lm, r_lp = _ref_convert_prompts_responses( @@ -185,7 +185,7 @@ def test_rl_preprocess_accepts_tensor_rewards(): rewards = [torch.tensor([1.0]), torch.tensor([0.5, 0.6, 0.7])] loss_masks = [[1], [1, 0, 1]] - _, _, _, rew, _, _, _ = convert_prompts_responses_to_batch_tensors( + _, _, _, rew, _, _, _, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, rewards, loss_masks ) _, _, _, r_rew, _, _ = _ref_convert_prompts_responses(prompts, responses, rewards, loss_masks, None, pad_token_id=0) @@ -203,7 +203,7 @@ def test_rl_preprocess_accepts_grad_tensor_rewards(): ] loss_masks = [[1], [1, 0, 1]] - _, _, _, rew, _, _, _ = convert_prompts_responses_to_batch_tensors( + _, _, _, rew, _, _, _, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, rewards, loss_masks ) @@ -223,7 +223,7 @@ def test_rl_preprocess_accepts_cuda_tensor_rewards(): ] loss_masks = [[1], [1, 0, 1]] - _, _, _, rew, _, _, _ = convert_prompts_responses_to_batch_tensors( + _, _, _, rew, _, _, _, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, rewards, loss_masks ) diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 5afc726c43..7224613654 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -20,7 +20,11 @@ build_nested_dataclass, overrides_dict_to_dotlist, ) -from skyrl.train.utils.utils import validate_cfg, validate_inference_engine_cfg +from skyrl.train.utils.utils import ( + validate_cfg, + validate_inference_engine_cfg, + validate_megatron_cfg, +) from tests.train.util import example_dummy_config @@ -154,6 +158,62 @@ def test_trainer_config_rejects_invalid_vocab_entropy_chunking(field_name, value TrainerConfig(**{field_name: value}) +@pytest.mark.parametrize( + ("override", "message"), + [ + ("generator.sampling_params.temperature=0", "temperature > 0"), + ("generator.sampling_params.top_k=1", "top_k > 1"), + ("generator.sampling_params.repetition_penalty=1.1", "repetition_penalty=1.0"), + ("generator.sampling_params.additional_kwargs.foo=bar", "additional_kwargs"), + ("generator.vision_language_generator=true", "vision_language_generator"), + ], +) +def test_sample_support_capture_rejects_unsupported_sampling_modifiers(override, message): + with pytest.raises(ValueError, match=message): + SkyRLTrainConfig.from_cli_overrides( + [ + "generator.inference_engine.enable_return_sample_support_set=true", + "generator.sampling_params.top_k=8", + override, + ] + ) + + +def test_routed_expert_capture_rejects_the_vision_language_generator(): + with pytest.raises(ValueError, match="vision_language_generator"): + SkyRLTrainConfig.from_cli_overrides( + [ + "generator.inference_engine.enable_return_routed_experts=true", + "generator.vision_language_generator=true", + ] + ) + + +def test_sample_support_capture_accepts_top_k_top_p_and_min_p(): + cfg = SkyRLTrainConfig.from_cli_overrides( + [ + "generator.inference_engine.enable_return_sample_support_set=true", + "generator.sampling_params.top_k=8", + "generator.sampling_params.top_p=0.9", + "generator.sampling_params.min_p=0.05", + ] + ) + + assert cfg.generator.inference_engine.enable_return_sample_support_set + + +def test_sample_support_capture_leaves_greedy_eval_sampling_params_alone(): + cfg = SkyRLTrainConfig.from_cli_overrides( + [ + "generator.inference_engine.enable_return_sample_support_set=true", + "generator.sampling_params.top_k=8", + ] + ) + + assert cfg.generator.eval_sampling_params.temperature == 0.0 + assert cfg.generator.eval_sampling_params.top_k == -1 + + def test_cli_overrides_plus_prefix_rejected(): with pytest.raises(ValueError, match="The '\\+' prefix"): SkyRLTrainConfig.from_cli_overrides(["+new_field=value"]) @@ -850,3 +910,35 @@ def test_delta_weight_sync_defaults(self): # `publish_staging_dir` and `local_checkpoint_dir` should be constructed based on `sync_dir` assert "my_sync_dir" in cfg.publish_staging_dir assert "my_sync_dir" in cfg.local_checkpoint_dir + + +class TestMegatronRouterReplayValidation: + @staticmethod + def _cfg(): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.generator.inference_engine.enable_return_routed_experts = True + cfg.trainer.policy.megatron_config.moe_enable_routing_replay = True + return cfg + + @pytest.mark.parametrize("vpp_size", [1, 2]) + def test_routing_replay_refuses_virtual_pipeline_parallelism(self, vpp_size): + cfg = self._cfg() + cfg.trainer.policy.megatron_config.transformer_config_kwargs["virtual_pipeline_model_parallel_size"] = vpp_size + + with pytest.raises(AssertionError, match="virtual_pipeline_model_parallel_size"): + validate_megatron_cfg(cfg) + + @pytest.mark.parametrize("vpp_size", [None, 0]) + def test_routing_replay_allows_unset_virtual_pipeline_parallelism(self, vpp_size): + cfg = self._cfg() + cfg.trainer.policy.megatron_config.transformer_config_kwargs["virtual_pipeline_model_parallel_size"] = vpp_size + + validate_megatron_cfg(cfg) + + def test_virtual_pipeline_parallelism_allowed_without_routing_replay(self): + cfg = self._cfg() + cfg.trainer.policy.megatron_config.moe_enable_routing_replay = False + cfg.trainer.policy.megatron_config.transformer_config_kwargs["virtual_pipeline_model_parallel_size"] = 2 + + validate_megatron_cfg(cfg) diff --git a/tests/train/test_packed_route_collation_equivalence.py b/tests/train/test_packed_route_collation_equivalence.py new file mode 100644 index 0000000000..39ff3ba552 --- /dev/null +++ b/tests/train/test_packed_route_collation_equivalence.py @@ -0,0 +1,355 @@ +"""Compare packed route collation with the padded reference path end to end.""" + +import sys +import types +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( + align_token_metadata, + build_token_metadata_layout, +) +from skyrl.backends.skyrl_train.training_batch import ( + TrainingInputBatch, + pad_training_input_batch, +) +from skyrl.backends.skyrl_train.utils import replay_utils +from skyrl.backends.skyrl_train.utils.replay_utils import ( + _split_replay_indices, + make_replay_padding_indices, + replay_padding_row, +) +from skyrl.train.dataset.preprocess import ( + convert_prompts_responses_to_batch_tensors, + make_router_padding_mask, +) + +NUM_LAYERS = 3 +TOPK = 4 +PAD_TOKEN_ID = 0 +# Above 2**8 so the compact dtype is int16, matching the production route width. +MIN_EXPERT_ID = 300 + + +def _reference_padded_routes( + routes: list[np.ndarray], + prompt_lens: list[int], + response_lens: list[int], +) -> torch.Tensor: + """Collate routes into a left-padded batch-major reference tensor.""" + max_total = max(p + r for p, r in zip(prompt_lens, response_lens)) + dtype = max((entry.dtype for entry in routes), key=lambda d: d.itemsize) + torch_dtype = torch.from_numpy(np.empty(0, dtype=dtype)).dtype + padded = torch.empty((len(routes), max_total, NUM_LAYERS, TOPK), dtype=torch_dtype) + padding_row = torch.arange(TOPK, dtype=torch_dtype) + for sample_index, entry in enumerate(routes): + left_pad = max_total - (prompt_lens[sample_index] + response_lens[sample_index]) + route_end = left_pad + entry.shape[0] + padded[sample_index, :left_pad] = padding_row + padded[sample_index, left_pad:route_end] = torch.from_numpy(entry) + padded[sample_index, route_end:] = padding_row + return padded + + +def _reference_replay_data( + padded_routes: torch.Tensor, + attention_mask: torch.Tensor, + local_layers: list[int], + *, + packed: bool, + tp_size: int, + tp_rank: int, +) -> list[torch.Tensor]: + """Gather real-token routes from the padded reference tensor.""" + layout = build_token_metadata_layout( + attention_mask, + padded_routes.device, + packed=packed, + fp8_enabled=False, + ) + local = padded_routes.index_select(2, torch.tensor(local_layers, dtype=torch.long)) + aligned = align_token_metadata( + local, + layout, + replay_padding_row(TOPK, dtype=padded_routes.dtype), + ) + if tp_size > 1: + chunk = aligned.shape[1] // tp_size + aligned = aligned[:, tp_rank * chunk : (tp_rank + 1) * chunk, :, :] + return _split_replay_indices(aligned) + + +@pytest.fixture +def parallel_state(monkeypatch): + try: + import megatron.core.parallel_state as mpu + except ModuleNotFoundError: + megatron = types.ModuleType("megatron") + core = types.ModuleType("megatron.core") + mpu = types.ModuleType("megatron.core.parallel_state") + megatron.core = core + core.parallel_state = mpu + monkeypatch.setitem(sys.modules, "megatron", megatron) + monkeypatch.setitem(sys.modules, "megatron.core", core) + monkeypatch.setitem(sys.modules, "megatron.core.parallel_state", mpu) + + monkeypatch.setattr(mpu, "get_tensor_model_parallel_world_size", lambda: 1, raising=False) + monkeypatch.setattr(mpu, "get_tensor_model_parallel_rank", lambda: 0, raising=False) + monkeypatch.setattr(mpu, "get_context_parallel_world_size", lambda: 1, raising=False) + monkeypatch.setattr(mpu, "get_context_parallel_rank", lambda: 0, raising=False) + return mpu + + +@pytest.fixture +def router_replay(monkeypatch): + """Capture what ``setup_per_microbatch_replay_forward`` hands Megatron.""" + module = types.ModuleType("megatron.core.transformer.moe.router_replay") + + class RouterReplay: + global_router_replay_instances = [object() for _ in range(NUM_LAYERS)] + replay_data: list[torch.Tensor] | None = None + + @classmethod + def set_replay_data(cls, replay_data): + cls.replay_data = replay_data + + @classmethod + def set_global_router_replay_action(cls, action): + pass + + module.RouterReplay = RouterReplay + module.RouterReplayAction = SimpleNamespace(REPLAY_FORWARD="replay_forward") + monkeypatch.setitem(sys.modules, "megatron.core.transformer.moe.router_replay", module) + monkeypatch.setattr( + replay_utils, + "scatter_router_padding_mask_for_model", + lambda mask, model, model_config: mask, + ) + return RouterReplay + + +def _make_batch(lengths: list[tuple[int, int]], *, captured_shortfall: int = 0, seed: int = 0): + """Build one batch of trajectories with the given ``(prompt_len, response_len)`` pairs. + + ``captured_shortfall`` leaves that many trailing tokens of the last trajectory without a + captured route, exercising the dummy-row tail that both paths must fill identically. + """ + rng = np.random.default_rng(seed) + prompts, responses, rewards, loss_masks, routes = [], [], [], [], [] + for index, (prompt_len, response_len) in enumerate(lengths): + prompts.append(list(rng.integers(1, 1000, size=prompt_len))) + responses.append(list(rng.integers(1, 1000, size=response_len))) + rewards.append([0.0] * response_len) + loss_masks.append([1] * response_len) + captured = prompt_len + response_len + if index == len(lengths) - 1: + captured -= captured_shortfall + routes.append( + rng.integers( + MIN_EXPERT_ID, + MIN_EXPERT_ID + 2000, + size=(captured, NUM_LAYERS, TOPK), + dtype=np.int16, + ) + ) + return prompts, responses, rewards, loss_masks, routes + + +def _run_both_paths( + lengths: list[tuple[int, int]], + *, + packed: bool, + tp_size: int, + local_layers: list[int], + captured_shortfall: int = 0, + batch_pad_size: int = 0, + stage_range: tuple[int, int] = (0, NUM_LAYERS), + monkeypatch, + parallel_state, + router_replay, +) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + prompts, responses, rewards, loss_masks, routes = _make_batch(lengths, captured_shortfall=captured_shortfall) + ( + sequences, + attention_mask, + response_mask, + _rewards, + loss_mask, + _logprobs, + packed_routes, + _sample_support, + ) = convert_prompts_responses_to_batch_tensors( + PAD_TOKEN_ID, + prompts, + responses, + rewards, + loss_masks, + rollout_expert_indices=routes, + ) + router_padding_mask = make_router_padding_mask(attention_mask, [entry.shape[0] for entry in routes]) + padded_routes = _reference_padded_routes(routes, [len(p) for p in prompts], [len(r) for r in responses]) + + if batch_pad_size: + batch = TrainingInputBatch( + { + "sequences": sequences, + "attention_mask": attention_mask, + "response_mask": response_mask, + "loss_mask": loss_mask, + "rollout_expert_indices": packed_routes, + "router_padding_mask": router_padding_mask, + } + ) + batch.metadata = {"uids": [f"u{index}" for index in range(len(prompts))]} + batch = pad_training_input_batch(batch, batch_pad_size) + attention_mask = batch["attention_mask"] + router_padding_mask = batch["router_padding_mask"] + packed_routes = batch["rollout_expert_indices"] + # Match the dummy rows added by batch padding in the packed path. + padded_routes = torch.cat( + [ + padded_routes, + make_replay_padding_indices((batch_pad_size, *padded_routes.shape[1:]), dtype=padded_routes.dtype), + ], + dim=0, + ) + + tp_rank = tp_size - 1 + monkeypatch.setattr(parallel_state, "get_tensor_model_parallel_world_size", lambda: tp_size, raising=False) + monkeypatch.setattr(parallel_state, "get_tensor_model_parallel_rank", lambda: tp_rank, raising=False) + router_replay.global_router_replay_instances = [object() for _ in local_layers] + monkeypatch.setattr(replay_utils, "_get_current_pp_stage_layer_range", lambda model_config: stage_range) + + layout = build_token_metadata_layout(attention_mask, packed_routes.device, packed=packed, fp8_enabled=False) + replay_utils.setup_per_microbatch_replay_forward( + packed_routes, + router_padding_mask, + attention_mask, + model=object(), + model_config=SimpleNamespace(fp8=None, sequence_parallel=False), + metadata_layout=layout, + remove_microbatch_padding=packed, + ) + new_replay_data = [tensor.clone() for tensor in router_replay.replay_data] + + reference = _reference_replay_data( + padded_routes, + attention_mask, + local_layers, + packed=packed, + tp_size=tp_size, + tp_rank=tp_rank, + ) + return new_replay_data, reference + + +def _assert_bit_identical(new_data: list[torch.Tensor], reference: list[torch.Tensor]) -> None: + assert len(new_data) == len(reference) + for slot, (produced, expected) in enumerate(zip(new_data, reference, strict=True)): + assert produced.dtype == expected.dtype == torch.int32, slot + assert produced.shape == expected.shape, (slot, produced.shape, expected.shape) + assert torch.equal(produced, expected), slot + + +LENGTH_DISTRIBUTIONS = { + # No padding at all: the packed and padded layouts coincide. + "uniform": [(8, 8), (8, 8), (8, 8), (8, 8)], + "mild_ragged": [(8, 8), (7, 8), (8, 6), (6, 7)], + # Typical RL: an order of magnitude between the shortest and longest trajectory. + "typical_rl": [(2, 2), (8, 24), (4, 6), (1, 31)], + "heavy_tail": [(1, 1), (1, 2), (2, 1), (16, 32)], +} + + +@pytest.mark.parametrize("distribution", sorted(LENGTH_DISTRIBUTIONS)) +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_packed_routes_match_padded_rectangle( + monkeypatch, parallel_state, router_replay, distribution, packed, tp_size +): + new_data, reference = _run_both_paths( + LENGTH_DISTRIBUTIONS[distribution], + packed=packed, + tp_size=tp_size, + local_layers=list(range(NUM_LAYERS)), + monkeypatch=monkeypatch, + parallel_state=parallel_state, + router_replay=router_replay, + ) + _assert_bit_identical(new_data, reference) + + +@pytest.mark.parametrize( + ("distribution", "captured_shortfall", "batch_pad_size", "local_layers", "stage_range"), + [ + pytest.param("typical_rl", 3, 0, list(range(NUM_LAYERS)), (0, NUM_LAYERS), id="uncaptured_suffix"), + pytest.param("mild_ragged", 0, 2, list(range(NUM_LAYERS)), (0, NUM_LAYERS), id="batch_padding"), + pytest.param("typical_rl", 0, 0, [1, 2], (1, 2), id="pipeline_stage_subset"), + ], +) +@pytest.mark.parametrize("packed", [False, True]) +def test_packed_routes_match_edge_cases( + monkeypatch, + parallel_state, + router_replay, + distribution, + captured_shortfall, + batch_pad_size, + local_layers, + stage_range, + packed, +): + new_data, reference = _run_both_paths( + LENGTH_DISTRIBUTIONS[distribution], + packed=packed, + tp_size=1, + local_layers=local_layers, + captured_shortfall=captured_shortfall, + batch_pad_size=batch_pad_size, + stage_range=stage_range, + monkeypatch=monkeypatch, + parallel_state=parallel_state, + router_replay=router_replay, + ) + _assert_bit_identical(new_data, reference) + + +@pytest.mark.parametrize("cp_size", [2, 4]) +def test_packed_routes_match_under_context_parallelism(monkeypatch, parallel_state, router_replay, cp_size): + """CP shards each padded sequence into front/back halves per rank.""" + for cp_rank in range(cp_size): + monkeypatch.setattr(parallel_state, "get_context_parallel_world_size", lambda: cp_size, raising=False) + monkeypatch.setattr(parallel_state, "get_context_parallel_rank", lambda rank=cp_rank: rank, raising=False) + new_data, reference = _run_both_paths( + LENGTH_DISTRIBUTIONS["mild_ragged"], + packed=True, + tp_size=1, + local_layers=list(range(NUM_LAYERS)), + monkeypatch=monkeypatch, + parallel_state=parallel_state, + router_replay=router_replay, + ) + _assert_bit_identical(new_data, reference) + + +@pytest.mark.parametrize("distribution", sorted(LENGTH_DISTRIBUTIONS)) +def test_packed_collation_allocates_no_padded_rectangle(distribution): + """The packed buffer must hold exactly the batch's real tokens.""" + prompts, responses, rewards, loss_masks, routes = _make_batch(LENGTH_DISTRIBUTIONS[distribution]) + *_, packed_routes, _ = convert_prompts_responses_to_batch_tensors( + PAD_TOKEN_ID, + prompts, + responses, + rewards, + loss_masks, + rollout_expert_indices=routes, + ) + + total_real = sum(len(p) + len(r) for p, r in zip(prompts, responses)) + max_total = max(len(p) + len(r) for p, r in zip(prompts, responses)) + assert packed_routes.values.shape == (total_real, NUM_LAYERS, TOPK) + assert packed_routes.values.numel() <= len(prompts) * max_total * NUM_LAYERS * TOPK + assert packed_routes.cu_seqlens.tolist()[-1] == total_real diff --git a/tests/train/test_trainer_utils.py b/tests/train/test_trainer_utils.py index 4fa401f1cf..2d4a39d260 100644 --- a/tests/train/test_trainer_utils.py +++ b/tests/train/test_trainer_utils.py @@ -14,6 +14,8 @@ import pytest import ray +from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertTrace +from skyrl.backends.skyrl_train.utils.sample_support import SAMPLE_SUPPORT_DTYPE from skyrl.train.config import SkyRLTrainConfig from skyrl.train.generators.base import GeneratorInput, GeneratorOutput, TrajectoryID from skyrl.train.utils.trainer_utils import ( @@ -1138,3 +1140,137 @@ def test_validate_stepwise_multiple_is_last_step_true_per_trajectory(): output["is_last_step"] = [True, True, True] with pytest.raises(AssertionError, match="is_last_step.*True.*trajectory continues"): validate_generator_output(num_prompts=1, generator_output=output, step_wise=True) + + +def _make_side_channel_output( + rollout_expert_indices=None, + rollout_sample_support=None, + loss_masks=None, +): + """A two-trajectory GeneratorOutput with 5- and 4-token sequences.""" + return GeneratorOutput( + prompt_token_ids=[[1, 2, 3], [4, 5]], + response_ids=[[10, 11], [12, 13]], + rewards=[0.5, 0.6], + loss_masks=loss_masks if loss_masks is not None else [[1, 1], [1, 1]], + stop_reasons=["stop", "stop"], + rollout_metrics={}, + rollout_logprobs=None, + rollout_expert_indices=rollout_expert_indices, + rollout_sample_support=rollout_sample_support, + ) + + +def _routes(num_rows): + return np.zeros((num_rows, 2, 2), dtype=np.int16) + + +def _support(num_rows): + return np.zeros((num_rows, 3), dtype=SAMPLE_SUPPORT_DTYPE) + + +@pytest.mark.parametrize( + ("route_rows", "loss_masks"), + [ + ((5, 4), [[1, 1], [1, 1]]), + ((4, 3), [[1, 1], [1, 1]]), + ((3, 2), [[1, 0], [1, 0]]), + ], +) +def test_validate_generator_output_accepts_route_under_coverage(route_rows, loss_masks): + """Route prefixes are valid when every trained token is covered.""" + output = _make_side_channel_output( + rollout_expert_indices=[_routes(rows) for rows in route_rows], + loss_masks=loss_masks, + ) + + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_rejects_routes_that_stop_short_of_a_trained_token(): + """A route prefix must cover the last trained token.""" + output = _make_side_channel_output(rollout_expert_indices=[_routes(3), _routes(3)]) + + with pytest.raises(AssertionError, match=r"rollout_expert_indices\[0\] captured 3 route rows.*token 4"): + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_accepts_the_coverage_a_multi_turn_trace_produces(): + """Validation accepts the prefix produced by a multi-turn route trace.""" + trace = RoutedExpertTrace() + trace.record_generation(prompt_token_count=3, generated_token_count=2, routed_experts=_routes(4)) + trace.record_generation(prompt_token_count=7, generated_token_count=2, routed_experts=_routes(4)) + loss_mask = [1, 1, 0, 0, 1, 1] + routes = trace.finalize(token_count=9, loss_mask=[0, 0, 0] + loss_mask) + assert len(routes) == 8 + + output = GeneratorOutput( + prompt_token_ids=[[1, 2, 3]], + response_ids=[[10, 11, 20, 21, 12, 13]], + rewards=[0.5], + loss_masks=[loss_mask], + stop_reasons=["stop"], + rollout_metrics={}, + rollout_logprobs=None, + rollout_expert_indices=[routes], + rollout_sample_support=None, + ) + + validate_generator_output(num_prompts=1, generator_output=output) + + +@pytest.mark.parametrize( + ("invalid_rows", "message"), + [ + (5, r"rollout_expert_indices\[1\] has 5 route rows for a 4-token"), + (0, r"rollout_expert_indices\[1\] has 0 route rows"), + ], +) +def test_validate_generator_output_rejects_invalid_route_bounds(invalid_rows, message): + output = _make_side_channel_output(rollout_expert_indices=[_routes(5), _routes(invalid_rows)]) + + with pytest.raises(AssertionError, match=message): + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_rejects_none_route_entry(): + output = _make_side_channel_output(rollout_expert_indices=[_routes(5), None]) + + with pytest.raises(AssertionError, match=r"rollout_expert_indices\[1\] is None"): + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_accepts_dense_sample_support(): + output = _make_side_channel_output(rollout_sample_support=[_support(2), _support(2)]) + + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_rejects_sample_support_row_shortfall(): + output = _make_side_channel_output(rollout_sample_support=[_support(2), _support(1)]) + + with pytest.raises(AssertionError, match=r"rollout_sample_support\[1\] has 1 support rows for 2 response tokens"): + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_rejects_none_sample_support_entry(): + output = _make_side_channel_output(rollout_sample_support=[None, _support(2)]) + + with pytest.raises(AssertionError, match=r"rollout_sample_support\[0\] is None"): + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_refuses_routes_under_step_wise(): + """Trajectory-aligned routes cannot be replayed against per-turn samples.""" + output = _make_stepwise_output(n_trajectories=1, steps_per_traj=(2,)) + output["rollout_expert_indices"] = [_routes(len(prompt) + 3) for prompt in output["prompt_token_ids"]] + + with pytest.raises(AssertionError, match="not supported with step-wise training"): + validate_generator_output(num_prompts=1, generator_output=output, step_wise=True) + + +def test_validate_generator_output_allows_sample_support_under_step_wise(): + output = _make_stepwise_output(n_trajectories=1, steps_per_traj=(2,)) + output["rollout_sample_support"] = [_support(len(response)) for response in output["response_ids"]] + + validate_generator_output(num_prompts=1, generator_output=output, step_wise=True) diff --git a/tests/utils/test_cpu_topology.py b/tests/utils/test_cpu_topology.py new file mode 100644 index 0000000000..0530273fc2 --- /dev/null +++ b/tests/utils/test_cpu_topology.py @@ -0,0 +1,101 @@ +""" +uv run --isolated --extra dev pytest tests/utils/test_cpu_topology.py +""" + +import os +from pathlib import Path + +import pytest + +from skyrl.utils import cpu_topology +from skyrl.utils.cpu_topology import cgroup_cpu_quota, permitted_cpu_cores, pool_workers + + +@pytest.fixture +def cgroup_paths(monkeypatch, tmp_path: Path): + paths = { + "CGROUP_V2_CPU_MAX_PATH": tmp_path / "cpu.max", + "CGROUP_V1_CPU_QUOTA_PATH": tmp_path / "cpu.cfs_quota_us", + "CGROUP_V1_CPU_PERIOD_PATH": tmp_path / "cpu.cfs_period_us", + } + for name, path in paths.items(): + monkeypatch.setattr(cpu_topology, name, str(path)) + return paths + + +@pytest.mark.parametrize( + ("version", "quota", "expected"), + [(2, "400000", 4), (2, "max", None), (1, "200000", 2), (1, "-1", None)], +) +def test_cgroup_quota(cgroup_paths, version, quota, expected): + if version == 2: + cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text(f"{quota} 100000\n") + else: + cgroup_paths["CGROUP_V1_CPU_QUOTA_PATH"].write_text(f"{quota}\n") + cgroup_paths["CGROUP_V1_CPU_PERIOD_PATH"].write_text("100000\n") + + assert cgroup_cpu_quota() == expected + + +def test_missing_cgroup_files(cgroup_paths): + assert cgroup_cpu_quota() is None + + +def test_unreadable_cgroup_values(cgroup_paths): + cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text("not-a-quota 100000\n") + cgroup_paths["CGROUP_V1_CPU_QUOTA_PATH"].write_text("") + cgroup_paths["CGROUP_V1_CPU_PERIOD_PATH"].write_text("") + + assert cgroup_cpu_quota() is None + + +@pytest.mark.parametrize(("quota", "expected"), [(150000, 1), (50000, 1), (100000, 1), (250000, 2)]) +def test_fractional_quota_floors_to_at_least_one(cgroup_paths, quota, expected): + cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text(f"{quota} 100000\n") + + assert cgroup_cpu_quota() == expected + + +@pytest.mark.parametrize(("affinity", "quota_cpus", "expected"), [(16, 4, 4), (4, 16, 4), (8, 8, 8)]) +def test_permitted_cores_is_the_lesser_of_affinity_and_quota(cgroup_paths, monkeypatch, affinity, quota_cpus, expected): + monkeypatch.setattr(os, "sched_getaffinity", lambda pid: set(range(affinity))) + cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text(f"{quota_cpus * 100000} 100000\n") + + assert permitted_cpu_cores() == expected + + +def test_permitted_cores_without_quota_is_the_affinity_mask(cgroup_paths, monkeypatch): + monkeypatch.setattr(os, "sched_getaffinity", lambda pid: set(range(12))) + + assert permitted_cpu_cores() == 12 + + +@pytest.mark.parametrize( + ("cap", "reserved", "cores", "expected"), + [ + (32, 8, 64, 32), # cap binds + (32, 8, 24, 16), # reserve binds + (32, 8, 8, 1), # reserve would leave nothing, so keep one worker + (1, 0, 64, 1), # cap of one + ], +) +def test_pool_workers(cap, reserved, cores, expected): + assert pool_workers(cap=cap, reserved=reserved, cores=cores) == expected + + +def test_pool_workers_defaults_to_permitted_cores(cgroup_paths, monkeypatch): + monkeypatch.setattr(os, "sched_getaffinity", lambda pid: set(range(64))) + cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text("1200000 100000\n") + + assert pool_workers(cap=32, reserved=4) == 8 + + +@pytest.mark.parametrize("cap", [0, -1]) +def test_pool_workers_rejects_non_positive_cap(cap): + with pytest.raises(ValueError, match="cap must be positive"): + pool_workers(cap=cap, reserved=0, cores=8) + + +def test_pool_workers_rejects_negative_reserve(): + with pytest.raises(ValueError, match="reserved cores must be non-negative"): + pool_workers(cap=8, reserved=-1, cores=8)