From ee54593eeb6097cc202acda5691f9afa10024643 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Thu, 16 Jul 2026 21:44:33 +0000 Subject: [PATCH 01/17] refactor(inference): extract reusable RemoteGenerateClient Extract the single-request HTTP generation path out of RemoteInferenceClient into a standalone RemoteGenerateClient plus a RemoteGenerateResult dataclass, so routed-expert results can be obtained without constructing the full inference/control-plane client. RemoteInferenceClient now owns an internal RemoteGenerateClient and delegates session management, _post, and _generate_single to it. Endpoint routing, retry/backoff, cache_salt handling, serialization, and lifecycle behavior are unchanged. --- .../remote_inference_client.py | 286 ++++++++++-------- .../test_remote_inference_client.py | 8 +- 2 files changed, 163 insertions(+), 131 deletions(-) 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..e3d8b21620 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -76,6 +76,7 @@ from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( decode_packed_routed_experts, ) +from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices from skyrl.backends.utils import convert_vllm_prompt_logprobs from skyrl.env_vars import ( SKYRL_GENERATE_CONCURRENCY_PER_ENGINE, @@ -167,6 +168,141 @@ 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] + + +@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) -> Any: + """POST JSON with retry on transient connection and response-decoding failures.""" + 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 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, + mm_features: Optional[MultiModalFeatures] = None, + cache_salt: Optional[str] = None, + ) -> RemoteGenerateResult: + """Generate one raw-token completion, optionally returning R3 routes.""" + path = "/skyrl/v1/generate" if return_routed_experts else "/inference/v1/generate" + 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(f"{self.proxy_url}{path}", json=payload, headers=headers) + 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("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 RemoteGenerateResult( + raw_response=response, + response_ids=token_ids, + response_logprobs=response_logprobs, + stop_reason=choice["finish_reason"], + routed_experts=routed_experts, + ) + + 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): """ @@ -225,7 +361,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 +418,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 @@ -472,63 +558,20 @@ async def _generate_single( mm_features: Optional[MultiModalFeatures] = None, cache_salt: Optional[str] = None, ) -> 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, + 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, } async def _render_for_sample( @@ -1405,9 +1448,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 +1466,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 +1474,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/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..23f3db0ef9 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 @@ -438,8 +438,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: @@ -508,7 +507,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") @@ -1025,8 +1024,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]]: From d881ecd5c58ab6627ca21c17ef9d87b7ea8805de Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Thu, 16 Jul 2026 21:48:07 +0000 Subject: [PATCH 02/17] refactor(generators): assemble routed-expert traces incrementally For multi-turn rollouts, build routed-expert (R3) trace data incrementally as the conversation grows instead of re-gathering the whole conversation's routes on every turn. - Add `routed_experts_prompt_starts` to `InferenceEngineInput` and thread a per-request `routed_experts_prompt_start` through `RemoteInferenceClient` and `RemoteInferenceGenerator.generate()`, forwarded as a sampling param so the engine only returns routes for the newly generated suffix. - Introduce `TokenMetadataTrace` (token-aligned array accumulator) and `RoutedExpertTrace`, which records each generation's routes and finalizes a full per-token routed-expert array with loss-mask-aware terminal padding. - Track a `RoutedExpertTrace` on `AgentLoopState` and record routes per turn, replacing the previous whole-conversation re-gather in `SkyRLGymGenerator.agent_loop`. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../distributed/megatron/token_metadata.py | 49 +++++++++++++++ .../skyrl_train/inference_servers/base.py | 1 + .../remote_inference_client.py | 30 ++++++++- .../skyrl_train/utils/routed_experts.py | 54 ++++++++++++++++ skyrl/train/generators/skyrl_gym_generator.py | 59 +++++++---------- .../distributed/test_token_metadata.py | 63 +++++++++++++++++++ .../test_remote_inference_client.py | 11 +++- .../generators/test_skyrl_gym_generator.py | 62 ++++++++++++++++-- 8 files changed, 288 insertions(+), 41 deletions(-) diff --git a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py index cb3e526c95..98059023f1 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py @@ -2,6 +2,7 @@ from dataclasses import dataclass +import numpy as np import torch from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import ( @@ -205,3 +206,51 @@ 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 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..4a8de4cdf4 100644 --- a/skyrl/backends/skyrl_train/inference_servers/base.py +++ b/skyrl/backends/skyrl_train/inference_servers/base.py @@ -34,6 +34,7 @@ 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]] class InferenceEngineOutput(TypedDict): 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 e3d8b21620..f6f22aa768 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -242,13 +242,27 @@ async def generate( session_id: Optional[Any], model: str, return_routed_experts: bool = False, + routed_experts_prompt_start: Optional[int] = None, mm_features: Optional[MultiModalFeatures] = None, cache_salt: Optional[str] = None, ) -> RemoteGenerateResult: """Generate one raw-token completion, optionally returning R3 routes.""" + 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") + path = "/skyrl/v1/generate" if return_routed_experts 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": sampling_params, + "sampling_params": request_sampling_params, "model": model, "token_ids": prompt_token_ids, } @@ -492,6 +506,12 @@ 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") get_logprobs = sampling_params.get("logprobs") is not None # Two semaphores decouple the generate and detokenize stages: @@ -515,6 +535,9 @@ 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 + ), model=model, cache_salt=cache_salt, ) @@ -524,6 +547,9 @@ 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 + ), model=model, cache_salt=cache_salt, ) @@ -557,6 +583,7 @@ async def _generate_single( model: str, mm_features: Optional[MultiModalFeatures] = None, cache_salt: Optional[str] = None, + routed_experts_prompt_start: Optional[int] = None, ) -> Dict[str, Any]: result = await self._get_generate_client().generate( prompt_token_ids=prompt_token_ids, @@ -564,6 +591,7 @@ async def _generate_single( session_id=session_id, model=model, return_routed_experts=self.enable_return_routed_experts, + routed_experts_prompt_start=routed_experts_prompt_start, mm_features=mm_features, cache_salt=cache_salt, ) diff --git a/skyrl/backends/skyrl_train/utils/routed_experts.py b/skyrl/backends/skyrl_train/utils/routed_experts.py index b3caec7c8b..53b011bcab 100644 --- a/skyrl/backends/skyrl_train/utils/routed_experts.py +++ b/skyrl/backends/skyrl_train/utils/routed_experts.py @@ -1,11 +1,65 @@ +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() + self._schema: tuple[int, int, np.dtype] | None = None + + @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 + compact = compact_routed_expert_indices(routed_experts) + if self._schema is None: + self._schema = (*compact.shape[1:], compact.dtype) + self._metadata.append(compact, expected_rows=expected_rows) + + def finalize(self, *, token_count: int, loss_mask: Sequence[int]) -> RoutedExpertIndices: + 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") + + 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}") + + padding_count = token_count - self.prompt_start + if padding_count: + if self._schema is None: + raise ValueError("cannot pad routed-expert trace before any routes are captured") + num_layers, topk, dtype = self._schema + padding_row = np.arange(topk, dtype=dtype) + padding = np.broadcast_to(padding_row, (padding_count, num_layers, topk)).copy() + self._metadata.append(padding, expected_rows=padding_count) + + return self._metadata.finalize(expected_rows=token_count) + + 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/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py index 26b1102e6c..7871f5cbe0 100644 --- a/skyrl/train/generators/skyrl_gym_generator.py +++ b/skyrl/train/generators/skyrl_gym_generator.py @@ -23,7 +23,7 @@ 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.train.config import GeneratorConfig, SkyRLGymConfig from skyrl.train.generators.base import ( GeneratorInput, @@ -83,7 +83,7 @@ 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 @dataclass @@ -93,14 +93,9 @@ class TurnOutput: output_logprobs: Optional[List[float]] new_obs: ConversationType obs_ids: List[int] - rollout_expert_indices: Optional[RoutedExpertIndices] reward: Optional[float] 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_loss_mask(self) -> List[int]: """ Get loss mask for this turn's tokens. @@ -379,6 +374,9 @@ async def agent_loop( rollout_logprobs=[] if get_logprobs else None, response_end_idx=None, done=False, + routed_expert_trace=( + RoutedExpertTrace() if self.generator_cfg.inference_engine.enable_return_routed_experts else None + ), ) while not agent_loop_state.done: @@ -401,11 +399,13 @@ 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, ) llm_call_start_time = time.monotonic() engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name) @@ -426,6 +426,14 @@ async def agent_loop( raise ValueError( "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, + ) # 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) stop_strs = current_sampling_params.get("stop", None) @@ -459,6 +467,8 @@ 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") obs_ids = self.get_obs_ids_from_obs(new_obs, agent_loop_state.done) @@ -471,7 +481,6 @@ async def agent_loop( reward=step_reward, obs_ids=obs_ids, added_eos=added_eos, - rollout_expert_indices=rollout_expert_indices, ) if is_step_wise: @@ -491,7 +500,6 @@ 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(), ) agent_loop_output.step_outputs.append(per_step_output) @@ -553,10 +561,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( @@ -573,6 +577,12 @@ async def agent_loop( rollout_logprobs.append(0.0) 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 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) @@ -1047,8 +1057,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 +1108,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 +1122,6 @@ 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 return agent_loop_state @@ -1194,13 +1192,4 @@ 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 - return agent_loop_state diff --git a/tests/backends/skyrl_train/distributed/test_token_metadata.py b/tests/backends/skyrl_train/distributed/test_token_metadata.py index ff8f3dc7de..e8b9bf7c77 100644 --- a/tests/backends/skyrl_train/distributed/test_token_metadata.py +++ b/tests/backends/skyrl_train/distributed/test_token_metadata.py @@ -1,10 +1,13 @@ 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.routed_experts import RoutedExpertTrace @pytest.fixture @@ -86,3 +89,63 @@ 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)) + + result = trace.finalize(token_count=9, loss_mask=[0, 0, 0, 1, 1, 0, 0, 1, 1]) + assert result.shape == (9, 2, 2) and result.dtype == np.uint8 + assert np.array_equal(result[-1, 0], [0, 1]) + + +@pytest.mark.parametrize("active", [False, True]) +def test_routed_expert_trace_only_pads_masked_suffix(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: + result = trace.finalize(token_count=5, loss_mask=mask) + assert np.array_equal(result[-2:, 0], [[0, 1], [0, 1]]) 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 23f3db0ef9..9c19654406 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 @@ -35,6 +35,7 @@ def create_mock_vllm_server(server_id: int) -> FastAPI: 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 @@ -63,6 +64,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 +109,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) @@ -479,13 +485,16 @@ 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_generate_rejects_list_routed_experts(self, monkeypatch): diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index 6d03d937b6..b253f02433 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -23,20 +23,17 @@ 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, added_eos=True, ) - assert output.get_turn_rollout_expert_indices() is routes assert output.get_turn_loss_mask() == [1, 1, 0, 0, 0] @@ -413,6 +410,63 @@ def mock_generate(_, model=None): assert output.stop_reason == "stop" +@pytest.mark.asyncio +@patch("skyrl_gym.make") +async def test_agent_loop_uses_incremental_routed_expert_trace( + 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 + 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 = [] + + def generate(input_batch, model=None): + 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 + return { + "responses": ["mocked output"], + "response_ids": [output_ids], + "stop_reasons": ["stop"], + "rollout_expert_indices": [routes], + } + + 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 = [] + + 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] + + @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): From 707724c43e511c4c4ffbe332e3f7f963009f6e04 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Mon, 20 Jul 2026 20:08:51 +0000 Subject: [PATCH 03/17] test(generators): drop obsolete turn route argument --- skyrl/backends/skyrl_train/utils/routed_experts.py | 4 +++- skyrl/train/generators/skyrl_gym_generator.py | 5 ++++- .../backends/skyrl_train/distributed/test_token_metadata.py | 4 +++- tests/train/generators/test_datatypes.py | 1 - 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/skyrl/backends/skyrl_train/utils/routed_experts.py b/skyrl/backends/skyrl_train/utils/routed_experts.py index 53b011bcab..6b213ef84e 100644 --- a/skyrl/backends/skyrl_train/utils/routed_experts.py +++ b/skyrl/backends/skyrl_train/utils/routed_experts.py @@ -3,7 +3,9 @@ import numpy as np -from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import TokenMetadataTrace +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)}) diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py index 7871f5cbe0..75c5d07671 100644 --- a/skyrl/train/generators/skyrl_gym_generator.py +++ b/skyrl/train/generators/skyrl_gym_generator.py @@ -23,7 +23,10 @@ InferenceEngineInput, InferenceEngineInterface, ) -from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices, RoutedExpertTrace +from skyrl.backends.skyrl_train.utils.routed_experts import ( + RoutedExpertIndices, + RoutedExpertTrace, +) from skyrl.train.config import GeneratorConfig, SkyRLGymConfig from skyrl.train.generators.base import ( GeneratorInput, diff --git a/tests/backends/skyrl_train/distributed/test_token_metadata.py b/tests/backends/skyrl_train/distributed/test_token_metadata.py index e8b9bf7c77..73ba160579 100644 --- a/tests/backends/skyrl_train/distributed/test_token_metadata.py +++ b/tests/backends/skyrl_train/distributed/test_token_metadata.py @@ -6,7 +6,9 @@ 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.distributed.megatron.token_metadata import ( + TokenMetadataTrace, +) from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertTrace 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, ) From 4bb8674956eb75fc53d8cd397ffdbe3980539012 Mon Sep 17 00:00:00 2001 From: Eric Tang <46737979+erictang000@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:19:43 -0700 Subject: [PATCH 04/17] Update skyrl/backends/skyrl_train/utils/routed_experts.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- skyrl/backends/skyrl_train/utils/routed_experts.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/skyrl/backends/skyrl_train/utils/routed_experts.py b/skyrl/backends/skyrl_train/utils/routed_experts.py index 6b213ef84e..a35bcc543a 100644 --- a/skyrl/backends/skyrl_train/utils/routed_experts.py +++ b/skyrl/backends/skyrl_train/utils/routed_experts.py @@ -46,9 +46,10 @@ def finalize(self, *, token_count: int, loss_mask: Sequence[int]) -> RoutedExper if self.prompt_start > token_count: raise ValueError(f"routed-expert trace has {self.prompt_start} rows for {token_count} tokens") - 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}") + 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}") padding_count = token_count - self.prompt_start if padding_count: From 31d6b3c1a3bbe3453a8843d59a2b3115af6eef89 Mon Sep 17 00:00:00 2001 From: lila-sync-bot Date: Fri, 14 Aug 2026 20:34:01 +0000 Subject: [PATCH 05/17] perf(r3): collate routes through a container-aware pool, and refuse replay under VPP --- .../skyrl_train/utils/replay_utils.py | 24 +--- .../weight_sync/delta_checkpoint.py | 8 +- .../workers/megatron/megatron_worker.py | 11 ++ skyrl/train/dataset/parallel_fill.py | 61 ++++++++ skyrl/train/dataset/preprocess.py | 134 +++++++++++++----- skyrl/train/utils/utils.py | 8 ++ skyrl/utils/cpu_topology.py | 97 +++++++++++++ .../skyrl_train/utils/test_replay_utils.py | 22 +-- tests/train/dataset/test_parallel_fill.py | 87 ++++++++++++ tests/train/dataset/test_preprocess.py | 112 +++++++++++++-- tests/train/test_config.py | 40 +++++- tests/utils/test_cpu_topology.py | 114 +++++++++++++++ 12 files changed, 625 insertions(+), 93 deletions(-) create mode 100644 skyrl/train/dataset/parallel_fill.py create mode 100644 skyrl/utils/cpu_topology.py create mode 100644 tests/train/dataset/test_parallel_fill.py create mode 100644 tests/utils/test_cpu_topology.py diff --git a/skyrl/backends/skyrl_train/utils/replay_utils.py b/skyrl/backends/skyrl_train/utils/replay_utils.py index 400376c65a..94696ad8d3 100644 --- a/skyrl/backends/skyrl_train/utils/replay_utils.py +++ b/skyrl/backends/skyrl_train/utils/replay_utils.py @@ -4,7 +4,6 @@ from contextlib import contextmanager -import numpy as np import torch from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( @@ -13,7 +12,7 @@ ) -def _replay_padding_row( +def replay_padding_row( topk: int, *, dtype: torch.dtype, @@ -40,27 +39,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. @@ -255,7 +237,7 @@ def setup_per_microbatch_replay_forward( 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( + route_padding = replay_padding_row( rollout_expert_indices.shape[-1], dtype=rollout_expert_indices.dtype, device=local_rollout_expert_indices.device, diff --git a/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py b/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py index 41c3a55a64..e8338aa006 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,9 @@ 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))) + # reserved=0: weight sync is a barrier, so no colocated work needs a core, and an + # unconstrained host keeps the pool size it had before the quota became visible. + 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 +1080,8 @@ def _empty_stats() -> dict[str, float]: } def _num_publish_workers(self) -> int: - default = min(8, os.cpu_count() or 1) + # reserved=0: 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_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index d116213abd..e9be351918 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -493,6 +493,17 @@ def init_configs( for k, v in transformer_config_kwargs.items(): setattr(provider, k, v) + # Checked on the resolved provider because to_megatron_provider() can supply its own VPP + # default. Every forward appends its routes to all local RouterReplay instances, but under + # VPP only the chunk being forwarded consumes them, so each instance's backward FIFO + # desyncs by the chunk count. + 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 diff --git a/skyrl/train/dataset/parallel_fill.py b/skyrl/train/dataset/parallel_fill.py new file mode 100644 index 0000000000..3762aee7a1 --- /dev/null +++ b/skyrl/train/dataset/parallel_fill.py @@ -0,0 +1,61 @@ +"""Fill a large controller-side batch buffer from a locally-sized thread pool. + +Controller-side collation runs single-threaded on the whole global batch before DP sharding, so +its cost lands on every training step and does not shard away. For a multi-GiB buffer the fill is +dominated by first-touch page faults on a fresh mapping rather than by the copies themselves, and +the trainer driver runs under ``@ray.remote(num_cpus=1)`` with ``OMP_NUM_THREADS=1``, so torch +cannot fault those pages in parallel on its own. + +``torch.set_num_threads`` is the wrong lever here: it is process-global, so on the fully-async +path it would also reach the generation coroutines running concurrently with collation. A local +pool keeps the effect scoped to this fill. Each callback owns a disjoint row range and the copies +release the GIL. +""" + +import functools +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor + +from skyrl.utils.cpu_topology import pool_workers + +# Copy throughput keeps scaling to the cap; past it the extra threads only take cores from +# colocated actors. +MAX_FILL_WORKERS = 32 +# Leave room for raylet, the GCS, the dashboard and the log monitor 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(index)`` for every index in ``range(num_rows)``, in parallel. + + ``fill_row`` must write a row range that no other index touches; nothing here serialises + overlapping writes. + """ + 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: + # list() forces the map so an exception in a worker surfaces here rather than being dropped. + list(pool.map(fill_row, range(num_rows))) diff --git a/skyrl/train/dataset/preprocess.py b/skyrl/train/dataset/preprocess.py index fb7ac6e2f7..4f876c4623 100644 --- a/skyrl/train/dataset/preprocess.py +++ b/skyrl/train/dataset/preprocess.py @@ -1,18 +1,26 @@ 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 skyrl.backends.skyrl_train.utils.replay_utils import make_replay_padding_indices_np +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.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 +95,87 @@ def _reward_to_numpy(custom_reward: Union[List[float], torch.Tensor]) -> np.ndar return reward_arr +def _collate_rollout_expert_indices( + rollout_expert_indices: List[RoutedExpertIndices], + pad_lens: np.ndarray, + max_total: int, +) -> Integer[torch.Tensor, "batch seq_len layer_num topk"]: + """Pack per-trajectory routes into one left-padded ``[batch, seq_len, layers, topk]`` buffer. + + ``pack_routed_experts`` establishes the canonical dtype on the sending side, so entries are + validated rather than rescanned here. Every region of the buffer is written exactly once, from + a locally sized thread pool: this fill runs on the whole global batch before DP sharding, on + every training step, and is bound by first-touch page faults on a fresh multi-GiB mapping. + """ + 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. + route_ends = [] + 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}" + ) + left_pad = int(pad_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_ends.append(left_pad + sample_indices.shape[0]) + + 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." + ) + padded = torch.empty( + (num_samples, max_total, num_layers, topk), + dtype=ROUTED_EXPERT_TORCH_DTYPES[batch_dtype], + ) + # Distinct experts per padding row, as Megatron's dropless dispatcher requires. + padding_row = replay_padding_row(topk, dtype=padded.dtype) + + def fill_sample(sample_index: int) -> None: + 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") + left_pad = int(pad_lens[sample_index]) + route_end = route_ends[sample_index] + sample_rows = padded[sample_index] + sample_rows[:left_pad] = padding_row + sample_rows[left_pad:route_end] = torch.from_numpy(sample_indices) + sample_rows[route_end:] = padding_row + + fill_batch_rows(fill_sample, num_samples) + return padded + + def convert_prompts_responses_to_batch_tensors( pad_token_id: int, prompts: List[List[int]], @@ -235,42 +324,11 @@ 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, + rollout_expert_indices_tensor = _collate_rollout_expert_indices( + rollout_expert_indices, + pad_lens, + max_total, ) - 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) return ( sequences, diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index e90c59630e..005a8163e5 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -217,6 +217,14 @@ 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." ) + # Every forward appends its routes to all local RouterReplay instances, but under VPP + # only the chunk being forwarded consumes them, so each instance's backward FIFO + # desyncs by the chunk count. + 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 ( diff --git a/skyrl/utils/cpu_topology.py b/skyrl/utils/cpu_topology.py new file mode 100644 index 0000000000..9ce1e51970 --- /dev/null +++ b/skyrl/utils/cpu_topology.py @@ -0,0 +1,97 @@ +"""How many CPUs this process may actually keep busy, and how to size a pool from it. + +``os.cpu_count()`` reports the machine, not the container. Under Ray a worker runs with +``OMP_NUM_THREADS=1`` and ``torch.get_num_threads() == 1`` without the process itself being +restricted, so neither of those can size a pool either. The two limits that do bind are the +affinity mask (cpuset / taskset pinning) and the CFS quota, and they are independent. +""" + +import os +from typing import Optional, Tuple + +# A container's own cgroup appears at the root of its cgroup namespace, so these paths already +# describe this process (``/proc/self/cgroup`` reads ``0::/``) and need no prefix join. Module +# level so tests can point them at 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 spells "no quota" as this literal; v1 spells it as a negative 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]: + """Whole CPUs this process' CFS quota permits, or ``None`` when no quota applies. + + Kubernetes ``limits.cpu`` becomes this quota, and flytekit's ``pod_spec_from_resources`` ends + with ``limits = limits or requests``, so a pod declaring only ``requests.cpu`` still carries + one. ``sched_getaffinity`` cannot see it: a quota caps CPU *time*, not which CPUs are runnable. + """ + 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: + """CPUs this process can actually keep busy: the lesser of its affinity mask and its quota. + + ``sched_getaffinity`` honours cpuset/taskset pinning and is Linux-only; elsewhere + ``cpu_count`` is the closest read available. + """ + 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: + """Pool size for a CPU-bound thread pool: capped, otherwise permitted cores minus a reserve. + + ``reserved`` leaves room for the colocated processes sharing this cgroup -- under Ray that is + raylet, the GCS, the dashboard and the log monitor. Oversubscribing a CFS quota buys a few + percent of throughput for constant throttling. + """ + 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/utils/test_replay_utils.py b/tests/backends/skyrl_train/utils/test_replay_utils.py index 304e11c625..beb3a59b0f 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,7 @@ 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.replay_utils import make_replay_padding_indices @pytest.fixture @@ -70,22 +66,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(): diff --git a/tests/train/dataset/test_parallel_fill.py b/tests/train/dataset/test_parallel_fill.py new file mode 100644 index 0000000000..36bbef2ab0 --- /dev/null +++ b/tests/train/dataset/test_parallel_fill.py @@ -0,0 +1,87 @@ +""" +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", [1, 4, 64]) +def test_every_index_is_filled_exactly_once(workers): + """Includes ``workers`` above ``num_rows``, which clamps rather than raising.""" + 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_default_worker_count_fills_every_index(): + num_rows = 8 + calls = [0] * num_rows + + def fill_row(index: int) -> None: + calls[index] += 1 + + fill_batch_rows(fill_row, num_rows) + + 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..ea5782a546 100644 --- a/tests/train/dataset/test_preprocess.py +++ b/tests/train/dataset/test_preprocess.py @@ -2,13 +2,17 @@ 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.train.dataset.preprocess import ( + ROUTED_EXPERT_TORCH_DTYPES, convert_prompts_responses_to_batch_tensors, make_router_padding_mask, ) @@ -156,15 +160,11 @@ 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. - - 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) +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 *_, routed = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, @@ -179,8 +179,24 @@ def test_routed_expert_tensor_narrows_wide_dtypes(tokenizer, dtype): assert routed.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.""" +@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): + """A truncated array keeps the dtype the wire declared; nothing rescans it to retighten.""" # int16 on the wire because of the trailing 300, which truncation then drops. routes = np.asarray([[[1, 2]], [[3, 4]], [[300, 5]]], dtype=np.int16) @@ -193,10 +209,82 @@ def test_routed_expert_tensor_retightens_after_truncation(tokenizer): rollout_expert_indices=[routes[:2]], ) - assert routed.dtype == torch.uint8 + assert routed.dtype == torch.int16 assert routed.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: + """The route collation as NumPy expressed it: broadcast ``arange(topk)``, then write each sample.""" + 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): + """Left padding, a short route prefix and trailing padding, across a mixed-dtype batch.""" + 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 it pads on both sides. + 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, + ) + + assert routed.dtype == torch.int16 + assert torch.equal(routed, torch.from_numpy(_numpy_padded_routes(routes, prompts, responses))) + # Padding routes are topk distinct experts, which Megatron's dropless dispatcher requires; + # zeros would collapse them onto one expert. + padding_row = [[0, 1, 2]] * num_layers + assert routed[0, 0].tolist() == padding_row + assert routed[0, 5].tolist() == padding_row + assert not torch.equal(routed[0, 0], torch.zeros_like(routed[0, 0])) + + def test_convert_prompts_responses_to_batch_tensors_exact(tokenizer): """ Test with inputs of exact lengths. diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 5afc726c43..a25ba56063 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 @@ -850,3 +854,37 @@ 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: + """Router replay (R3) incompatibilities refused at submission time.""" + + @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/utils/test_cpu_topology.py b/tests/utils/test_cpu_topology.py new file mode 100644 index 0000000000..1e52f6763f --- /dev/null +++ b/tests/utils/test_cpu_topology.py @@ -0,0 +1,114 @@ +""" +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): + """Point every cgroup path at a fixture directory; each file is absent until written.""" + 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 + + +def test_cgroup_v2_quota(cgroup_paths): + cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text("400000 100000\n") + + assert cgroup_cpu_quota() == 4 + + +def test_cgroup_v2_max_literal_is_unlimited(cgroup_paths): + cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text("max 100000\n") + + assert cgroup_cpu_quota() is None + + +def test_cgroup_v1_quota(cgroup_paths): + cgroup_paths["CGROUP_V1_CPU_QUOTA_PATH"].write_text("200000\n") + cgroup_paths["CGROUP_V1_CPU_PERIOD_PATH"].write_text("100000\n") + + assert cgroup_cpu_quota() == 2 + + +def test_cgroup_v1_negative_quota_is_unlimited(cgroup_paths): + cgroup_paths["CGROUP_V1_CPU_QUOTA_PATH"].write_text("-1\n") + cgroup_paths["CGROUP_V1_CPU_PERIOD_PATH"].write_text("100000\n") + + assert cgroup_cpu_quota() is None + + +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) From e080afea8a9ad2842ff43c72e2ba428d4b9d7942 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Sat, 22 Aug 2026 00:05:46 +0000 Subject: [PATCH 06/17] style: tighten route collation comments and tests --- .../weight_sync/delta_checkpoint.py | 5 ++- .../workers/megatron/megatron_worker.py | 6 ++-- skyrl/train/dataset/parallel_fill.py | 28 +++++---------- skyrl/train/dataset/preprocess.py | 11 +++--- skyrl/train/utils/utils.py | 4 +-- skyrl/utils/cpu_topology.py | 35 ++++--------------- tests/train/dataset/test_parallel_fill.py | 15 +------- tests/train/dataset/test_preprocess.py | 12 +++---- tests/train/test_config.py | 2 -- tests/utils/test_cpu_topology.py | 35 ++++++------------- 10 files changed, 41 insertions(+), 112 deletions(-) diff --git a/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py b/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py index e8338aa006..3a36dbf876 100644 --- a/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py +++ b/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py @@ -761,8 +761,7 @@ def apply_one(item: tuple[DeltaTensorRecord, str, bytes]) -> None: mismatches.append(record.name) del region, patch - # reserved=0: weight sync is a barrier, so no colocated work needs a core, and an - # unconstrained host keeps the pool size it had before the quota became visible. + # 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)) @@ -1080,7 +1079,7 @@ def _empty_stats() -> dict[str, float]: } def _num_publish_workers(self) -> int: - # reserved=0: publishing is a barrier like the apply path above. + # Publishing is a barrier, like the apply path above. default = pool_workers(cap=8, reserved=0) return self.publish_num_workers or default diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index e9be351918..7c009260df 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -493,10 +493,8 @@ def init_configs( for k, v in transformer_config_kwargs.items(): setattr(provider, k, v) - # Checked on the resolved provider because to_megatron_provider() can supply its own VPP - # default. Every forward appends its routes to all local RouterReplay instances, but under - # VPP only the chunk being forwarded consumes them, so each instance's backward FIFO - # desyncs by the chunk count. + # 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( diff --git a/skyrl/train/dataset/parallel_fill.py b/skyrl/train/dataset/parallel_fill.py index 3762aee7a1..813884604d 100644 --- a/skyrl/train/dataset/parallel_fill.py +++ b/skyrl/train/dataset/parallel_fill.py @@ -1,15 +1,7 @@ -"""Fill a large controller-side batch buffer from a locally-sized thread pool. - -Controller-side collation runs single-threaded on the whole global batch before DP sharding, so -its cost lands on every training step and does not shard away. For a multi-GiB buffer the fill is -dominated by first-touch page faults on a fresh mapping rather than by the copies themselves, and -the trainer driver runs under ``@ray.remote(num_cpus=1)`` with ``OMP_NUM_THREADS=1``, so torch -cannot fault those pages in parallel on its own. - -``torch.set_num_threads`` is the wrong lever here: it is process-global, so on the fully-async -path it would also reach the generation coroutines running concurrently with collation. A local -pool keeps the effect scoped to this fill. Each callback owns a disjoint row range and the copies -release the GIL. +"""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 @@ -18,10 +10,9 @@ from skyrl.utils.cpu_topology import pool_workers -# Copy throughput keeps scaling to the cap; past it the extra threads only take cores from -# colocated actors. +# Extra threads beyond this cap take cores from colocated actors without improving throughput. MAX_FILL_WORKERS = 32 -# Leave room for raylet, the GCS, the dashboard and the log monitor in the same cgroup. +# Leave room for Ray services in the same cgroup. RESERVED_FILL_CORES = 8 @@ -36,10 +27,9 @@ def fill_batch_rows( *, workers: int | None = None, ) -> None: - """Call ``fill_row(index)`` for every index in ``range(num_rows)``, in parallel. + """Call ``fill_row`` for every row, possibly in parallel. - ``fill_row`` must write a row range that no other index touches; nothing here serialises - overlapping writes. + 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}") @@ -57,5 +47,5 @@ def fill_batch_rows( return with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="skyrl-batch-fill") as pool: - # list() forces the map so an exception in a worker surfaces here rather than being dropped. + # 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 4f876c4623..f7b0b73e1a 100644 --- a/skyrl/train/dataset/preprocess.py +++ b/skyrl/train/dataset/preprocess.py @@ -100,12 +100,9 @@ def _collate_rollout_expert_indices( pad_lens: np.ndarray, max_total: int, ) -> Integer[torch.Tensor, "batch seq_len layer_num topk"]: - """Pack per-trajectory routes into one left-padded ``[batch, seq_len, layers, topk]`` buffer. + """Pack routes into a left-padded ``[batch, seq_len, layers, topk]`` buffer. - ``pack_routed_experts`` establishes the canonical dtype on the sending side, so entries are - validated rather than rescanned here. Every region of the buffer is written exactly once, from - a locally sized thread pool: this fill runs on the whole global batch before DP sharding, on - every training step, and is bound by first-touch page faults on a fresh multi-GiB mapping. + The sender establishes canonical dtypes, so this path validates rather than rescans entries. """ num_samples = len(rollout_expert_indices) for sample_index, sample_indices in enumerate(rollout_expert_indices): @@ -130,7 +127,7 @@ def _collate_rollout_expert_indices( 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. + # Validate before dispatch so errors are deterministic. route_ends = [] for sample_index, sample_indices in enumerate(rollout_expert_indices): if sample_indices.ndim != 3 or sample_indices.shape[1:] != (num_layers, topk): @@ -162,7 +159,7 @@ def _collate_rollout_expert_indices( def fill_sample(sample_index: int) -> None: 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. + # torch.from_numpy requires a writable buffer. if not flags.c_contiguous or not flags.writeable: sample_indices = sample_indices.copy(order="C") left_pad = int(pad_lens[sample_index]) diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 005a8163e5..a1961e5a0c 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -217,9 +217,7 @@ 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." ) - # Every forward appends its routes to all local RouterReplay instances, but under VPP - # only the chunk being forwarded consumes them, so each instance's backward FIFO - # desyncs by the chunk count. + # 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. " diff --git a/skyrl/utils/cpu_topology.py b/skyrl/utils/cpu_topology.py index 9ce1e51970..0afb8a6d27 100644 --- a/skyrl/utils/cpu_topology.py +++ b/skyrl/utils/cpu_topology.py @@ -1,22 +1,15 @@ -"""How many CPUs this process may actually keep busy, and how to size a pool from it. - -``os.cpu_count()`` reports the machine, not the container. Under Ray a worker runs with -``OMP_NUM_THREADS=1`` and ``torch.get_num_threads() == 1`` without the process itself being -restricted, so neither of those can size a pool either. The two limits that do bind are the -affinity mask (cpuset / taskset pinning) and the CFS quota, and they are independent. -""" +"""Determine usable CPUs from process affinity and cgroup quota.""" import os from typing import Optional, Tuple -# A container's own cgroup appears at the root of its cgroup namespace, so these paths already -# describe this process (``/proc/self/cgroup`` reads ``0::/``) and need no prefix join. Module -# level so tests can point them at fixtures. +# 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 spells "no quota" as this literal; v1 spells it as a negative quota. +# cgroup v2 uses this literal for an unlimited quota. CGROUP_V2_CPU_MAX_UNLIMITED = "max" @@ -49,12 +42,7 @@ def _read_cgroup_v1_cpu_max() -> Optional[Tuple[float, float]]: def cgroup_cpu_quota() -> Optional[int]: - """Whole CPUs this process' CFS quota permits, or ``None`` when no quota applies. - - Kubernetes ``limits.cpu`` becomes this quota, and flytekit's ``pod_spec_from_resources`` ends - with ``limits = limits or requests``, so a pod declaring only ``requests.cpu`` still carries - one. ``sched_getaffinity`` cannot see it: a quota caps CPU *time*, not which CPUs are runnable. - """ + """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 @@ -66,11 +54,7 @@ def cgroup_cpu_quota() -> Optional[int]: def permitted_cpu_cores() -> int: - """CPUs this process can actually keep busy: the lesser of its affinity mask and its quota. - - ``sched_getaffinity`` honours cpuset/taskset pinning and is Linux-only; elsewhere - ``cpu_count`` is the closest read available. - """ + """Return the lesser of the process affinity and cgroup quota.""" try: affinity = len(os.sched_getaffinity(0)) except AttributeError: @@ -82,12 +66,7 @@ def permitted_cpu_cores() -> int: def pool_workers(*, cap: int, reserved: int, cores: Optional[int] = None) -> int: - """Pool size for a CPU-bound thread pool: capped, otherwise permitted cores minus a reserve. - - ``reserved`` leaves room for the colocated processes sharing this cgroup -- under Ray that is - raylet, the GCS, the dashboard and the log monitor. Oversubscribing a CFS quota buys a few - percent of throughput for constant throttling. - """ + """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: diff --git a/tests/train/dataset/test_parallel_fill.py b/tests/train/dataset/test_parallel_fill.py index 36bbef2ab0..1b0e6daedd 100644 --- a/tests/train/dataset/test_parallel_fill.py +++ b/tests/train/dataset/test_parallel_fill.py @@ -9,9 +9,8 @@ from skyrl.train.dataset.parallel_fill import fill_batch_rows -@pytest.mark.parametrize("workers", [1, 4, 64]) +@pytest.mark.parametrize("workers", [None, 1, 4, 64]) def test_every_index_is_filled_exactly_once(workers): - """Includes ``workers`` above ``num_rows``, which clamps rather than raising.""" num_rows = 32 calls = [0] * num_rows @@ -23,18 +22,6 @@ def fill_row(index: int) -> None: assert calls == [1] * num_rows -def test_default_worker_count_fills_every_index(): - num_rows = 8 - calls = [0] * num_rows - - def fill_row(index: int) -> None: - calls[index] += 1 - - fill_batch_rows(fill_row, num_rows) - - assert calls == [1] * num_rows - - def test_single_worker_runs_serially_on_the_calling_thread(): order = [] threads = set() diff --git a/tests/train/dataset/test_preprocess.py b/tests/train/dataset/test_preprocess.py index ea5782a546..89545b85e4 100644 --- a/tests/train/dataset/test_preprocess.py +++ b/tests/train/dataset/test_preprocess.py @@ -196,8 +196,7 @@ def test_routed_expert_tensor_rejects_non_canonical_dtypes(tokenizer, dtype): def test_routed_expert_tensor_keeps_the_sender_dtype_after_truncation(tokenizer): - """A truncated array keeps the dtype the wire declared; nothing rescans it to retighten.""" - # int16 on the wire because of the trailing 300, which truncation then drops. + # 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( @@ -243,7 +242,7 @@ def _numpy_padded_routes( prompts: List[List[int]], responses: List[List[int]], ) -> np.ndarray: - """The route collation as NumPy expressed it: broadcast ``arange(topk)``, then write each sample.""" + """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) @@ -256,11 +255,10 @@ def _numpy_padded_routes( def test_routed_expert_tensor_is_bit_identical_to_numpy_collation(tokenizer): - """Left padding, a short route prefix and trailing padding, across a mixed-dtype batch.""" 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 it pads on both sides. + # Sample 0 has fewer route rows than tokens and needs padding on both sides. 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), @@ -277,12 +275,10 @@ def test_routed_expert_tensor_is_bit_identical_to_numpy_collation(tokenizer): assert routed.dtype == torch.int16 assert torch.equal(routed, torch.from_numpy(_numpy_padded_routes(routes, prompts, responses))) - # Padding routes are topk distinct experts, which Megatron's dropless dispatcher requires; - # zeros would collapse them onto one expert. + # Padding routes must select distinct experts for the dropless dispatcher. padding_row = [[0, 1, 2]] * num_layers assert routed[0, 0].tolist() == padding_row assert routed[0, 5].tolist() == padding_row - assert not torch.equal(routed[0, 0], torch.zeros_like(routed[0, 0])) def test_convert_prompts_responses_to_batch_tensors_exact(tokenizer): diff --git a/tests/train/test_config.py b/tests/train/test_config.py index a25ba56063..7e4dc1a493 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -857,8 +857,6 @@ def test_delta_weight_sync_defaults(self): class TestMegatronRouterReplayValidation: - """Router replay (R3) incompatibilities refused at submission time.""" - @staticmethod def _cfg(): cfg = _make_validated_test_config() diff --git a/tests/utils/test_cpu_topology.py b/tests/utils/test_cpu_topology.py index 1e52f6763f..0530273fc2 100644 --- a/tests/utils/test_cpu_topology.py +++ b/tests/utils/test_cpu_topology.py @@ -13,7 +13,6 @@ @pytest.fixture def cgroup_paths(monkeypatch, tmp_path: Path): - """Point every cgroup path at a fixture directory; each file is absent until written.""" paths = { "CGROUP_V2_CPU_MAX_PATH": tmp_path / "cpu.max", "CGROUP_V1_CPU_QUOTA_PATH": tmp_path / "cpu.cfs_quota_us", @@ -24,30 +23,18 @@ def cgroup_paths(monkeypatch, tmp_path: Path): return paths -def test_cgroup_v2_quota(cgroup_paths): - cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text("400000 100000\n") - - assert cgroup_cpu_quota() == 4 - - -def test_cgroup_v2_max_literal_is_unlimited(cgroup_paths): - cgroup_paths["CGROUP_V2_CPU_MAX_PATH"].write_text("max 100000\n") - - assert cgroup_cpu_quota() is None - - -def test_cgroup_v1_quota(cgroup_paths): - cgroup_paths["CGROUP_V1_CPU_QUOTA_PATH"].write_text("200000\n") - cgroup_paths["CGROUP_V1_CPU_PERIOD_PATH"].write_text("100000\n") - - assert cgroup_cpu_quota() == 2 - - -def test_cgroup_v1_negative_quota_is_unlimited(cgroup_paths): - cgroup_paths["CGROUP_V1_CPU_QUOTA_PATH"].write_text("-1\n") - cgroup_paths["CGROUP_V1_CPU_PERIOD_PATH"].write_text("100000\n") +@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() is None + assert cgroup_cpu_quota() == expected def test_missing_cgroup_files(cgroup_paths): From 5dccdc0d58787e6138d1a2a4f5618699f4da3315 Mon Sep 17 00:00:00 2001 From: lila-sync-bot Date: Fri, 14 Aug 2026 21:14:53 +0000 Subject: [PATCH 07/17] refactor(wire): generic packed-ndarray codec and an N-field response-body splice --- .../inference_servers/generate_wire.py | 211 +++++++++++++++-- .../remote_inference_client.py | 30 ++- .../inference_servers/vllm_server_actor.py | 3 +- .../inference_servers/test_generate_wire.py | 220 ++++++++++++++++++ .../test_remote_inference_client.py | 152 +++++++++++- 5 files changed, 587 insertions(+), 29 deletions(-) diff --git a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py index a9b8f43a2c..4e33778e77 100644 --- a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py +++ b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py @@ -3,14 +3,23 @@ ``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. + +A side-channel array travels as a ``{data: , shape: [...], dtype: +}`` envelope plus any sidecar fields. ``data`` is emitted first so +``load_packed_body`` can cut the base64 straight out of the raw response bytes +and hand the decoder a ``memoryview``, sparing orjson the cost of materializing +a multi-hundred-megabyte 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 ( @@ -22,7 +31,38 @@ # 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 of a packed NumPy array. + + ``DATA`` is emitted first: ``load_packed_body`` locates a blob by byte + prefix, so any other order defeats the splice. + """ + + 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) +"""Fields ``load_packed_body`` splices out of a raw response body.""" + +_ENVELOPE_KEYS = frozenset(PackedArrayKey) + +_ROUTED_EXPERTS_NDIM = 3 + +_QUOTE = b'"' + +# Bytes shared by every envelope, and the anchor for the single scan in +# ``load_packed_body``: base64 contains none of them, so the search skips blobs. +_PACKED_DATA_ANCHOR = f':{{"{PackedArrayKey.DATA}":"'.encode() def build_logprobs_content( @@ -72,35 +112,162 @@ 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 a packed envelope into its array and its sidecar fields. + + ``data`` may be base64 in a ``str`` or in any buffer, so ``load_packed_body`` + can pass a ``memoryview`` into the raw response body. + """ + 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 + 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] # 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( + 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 _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 body, cutting every registered packed blob out first. + + One scan of ``raw`` finds each envelope, replaces its base64 with an empty + string for orjson, and keeps the blob as a ``memoryview`` that + ``unpack_ndarray`` decodes in place of a Python ``str``. + + A ``null`` field passes through -- the server sends that when a request + captured nothing. Any other layout the scan cannot cut, such as a body + re-serialized with a different key order or spacing, raises: parsing the + blob as a ``str`` would silently forfeit the whole point of the splice. + """ + 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 f6f22aa768..719540489e 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -74,7 +74,9 @@ MultiModalFeatures, ) from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( + PackedField, decode_packed_routed_experts, + load_packed_body, ) from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices from skyrl.backends.utils import convert_vllm_prompt_logprobs @@ -201,15 +203,28 @@ async def _get_session(self) -> aiohttp.ClientSession: ) return self._session - async def _post(self, url: str, json: Dict[str, Any], headers: Optional[Dict[str, str]] = None) -> Any: - """POST JSON with retry on transient connection and response-decoding failures.""" + 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 retry on transient connection and response-decoding failures. + + ``packed_side_channels`` splices packed arrays out of the raw bytes before + parsing; only ``/skyrl/v1/generate`` returns them, and no other caller + should pay for the scan. + """ 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()) + 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() @@ -277,7 +292,12 @@ async def generate( if session_id: headers["X-Session-ID"] = str(session_id) - response = await self._post(f"{self.proxy_url}{path}", json=payload, headers=headers) + response = await self._post( + f"{self.proxy_url}{path}", + json=payload, + headers=headers, + packed_side_channels=return_routed_experts, + ) choice = response["choices"][0] token_ids = choice["token_ids"] logprobs = choice.get("logprobs") @@ -289,7 +309,7 @@ async def generate( routed_experts = None if return_routed_experts: - packed_routed_experts = choice.get("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) 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..dab33322e5 100644 --- a/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py +++ b/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py @@ -37,6 +37,7 @@ ) from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( CLAMPED_LOGPROB, + PackedField, build_logprobs_content, pack_routed_experts, ) @@ -469,7 +470,7 @@ 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, } ] } 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..ece60a57f3 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,28 @@ from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( CLAMPED_LOGPROB, + PACKED_SIDE_CHANNEL_FIELDS, + PackedArrayKey, + PackedField, build_logprobs_content, decode_packed_routed_experts, + load_packed_body, + pack_ndarray, pack_routed_experts, + 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 +201,205 @@ def test_decode_rejects_noncanonical_dtype(): with pytest.raises(ValueError, match="non-canonical dtype"): decode_packed_routed_experts(payload) + + +@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(): + # load_packed_body finds a blob by byte prefix, so `data` must stay first + # even when sidecar fields are present. + 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_registry_seeds_both_side_channels(): + # A later branch adds the rollout_sample_support producer; the splice must + # already know the name so that branch stays a pure addition. + assert PACKED_SIDE_CHANNEL_FIELDS == ("routed_experts", "rollout_sample_support") + + +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] + + # Both blobs are handed over as memoryviews into `raw`, never as Python strs. + 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_leaves_a_spoofed_prefix_alone_when_no_field_is_present(): + body = _body(note='"routed_experts":{"data":"AAAA"') + + assert load_packed_body(orjson.dumps(body)) == body + + +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"] + + # Blobs are matched to envelopes in document order. + 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 9c19654406..04320323b6 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,15 +10,24 @@ 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, + unpack_ndarray, ) from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( SKYRL_LORA_ADAPTER_NAME, @@ -29,6 +39,21 @@ ) from skyrl.train.config import SkyRLTrainConfig +_SUPPORT_DTYPES = frozenset({np.dtype(np.float32)}) +_ROUTES = np.arange(12).reshape(3, 2, 2) +_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.""" @@ -46,11 +71,46 @@ 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 + # Hits on the packed-body endpoints below, used to assert retry behaviour. + 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) @@ -558,6 +618,96 @@ async def test_detokenize(self, client): assert result[0] == "hello world" # Mock response +class TestPackedSideChannelBodies: + """``_post(packed_side_channels=True)`` splices packed blobs out of the raw bytes. + + The splice is what keeps a ~121 MiB base64 blob from becoming a Python + ``str``; a layout it cannot cut must fail instead of quietly falling back to + whole-body parsing. + """ + + 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] + + # Both blobs arrive as memoryviews into the raw response, never as strs. + 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).""" From a604e5ccc876f69eeb973ac118044cb79c4ea0ae Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Sat, 22 Aug 2026 00:06:17 +0000 Subject: [PATCH 08/17] style: tighten wire codec comments and tests --- .../inference_servers/generate_wire.py | 39 +++++-------------- .../remote_inference_client.py | 7 +--- .../inference_servers/test_generate_wire.py | 17 -------- .../test_remote_inference_client.py | 9 +---- 4 files changed, 12 insertions(+), 60 deletions(-) diff --git a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py index 4e33778e77..a5f7aac23c 100644 --- a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py +++ b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py @@ -6,11 +6,9 @@ helpers here exist to get sampled logprobs and NumPy side channels across that boundary intact. -A side-channel array travels as a ``{data: , shape: [...], dtype: -}`` envelope plus any sidecar fields. ``data`` is emitted first so -``load_packed_body`` can cut the base64 straight out of the raw response bytes -and hand the decoder a ``memoryview``, sparing orjson the cost of materializing -a multi-hundred-megabyte Python ``str``. +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 @@ -33,11 +31,7 @@ class PackedArrayKey(StrEnum): - """Envelope keys of a packed NumPy array. - - ``DATA`` is emitted first: ``load_packed_body`` locates a blob by byte - prefix, so any other order defeats the splice. - """ + """Envelope keys, with ``DATA`` first for ``load_packed_body``.""" DATA = "data" SHAPE = "shape" @@ -52,7 +46,6 @@ class PackedField(StrEnum): PACKED_SIDE_CHANNEL_FIELDS: tuple[str, ...] = tuple(PackedField) -"""Fields ``load_packed_body`` splices out of a raw response body.""" _ENVELOPE_KEYS = frozenset(PackedArrayKey) @@ -60,8 +53,7 @@ class PackedField(StrEnum): _QUOTE = b'"' -# Bytes shared by every envelope, and the anchor for the single scan in -# ``load_packed_body``: base64 contains none of them, so the search skips blobs. +# Base64 cannot contain this scan anchor. _PACKED_DATA_ANCHOR = f':{{"{PackedArrayKey.DATA}":"'.encode() @@ -147,11 +139,7 @@ def unpack_ndarray( allowed_dtypes: Collection[np.dtype], ndim: int, ) -> Tuple[np.ndarray, dict[str, Any]]: - """Decode a packed envelope into its array and its sidecar fields. - - ``data`` may be base64 in a ``str`` or in any buffer, so ``load_packed_body`` - can pass a ``memoryview`` into the raw response body. - """ + """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: @@ -165,8 +153,7 @@ def unpack_ndarray( 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] - # 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. + # 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 ): @@ -204,16 +191,10 @@ def _data_prefix(field: str) -> bytes: def load_packed_body(raw: bytes, *, fields: tuple[str, ...] = PACKED_SIDE_CHANNEL_FIELDS) -> dict[str, Any]: - """Parse a response body, cutting every registered packed blob out first. - - One scan of ``raw`` finds each envelope, replaces its base64 with an empty - string for orjson, and keeps the blob as a ``memoryview`` that - ``unpack_ndarray`` decodes in place of a Python ``str``. + """Parse a response after replacing registered base64 blobs with views. - A ``null`` field passes through -- the server sends that when a request - captured nothing. Any other layout the scan cannot cut, such as a body - re-serialized with a different key order or spacing, raises: parsing the - blob as a ``str`` would silently forfeit the whole point of the splice. + 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} 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 719540489e..f6137b8e49 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -211,12 +211,7 @@ async def _post( *, packed_side_channels: bool = False, ) -> Any: - """POST JSON with retry on transient connection and response-decoding failures. - - ``packed_side_channels`` splices packed arrays out of the raw bytes before - parsing; only ``/skyrl/v1/generate`` returns them, and no other caller - should pay for the scan. - """ + """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): 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 ece60a57f3..c93a4d4cfd 100644 --- a/tests/backends/skyrl_train/inference_servers/test_generate_wire.py +++ b/tests/backends/skyrl_train/inference_servers/test_generate_wire.py @@ -12,7 +12,6 @@ from skyrl.backends.skyrl_train.inference_servers.generate_wire import ( CLAMPED_LOGPROB, - PACKED_SIDE_CHANNEL_FIELDS, PackedArrayKey, PackedField, build_logprobs_content, @@ -223,8 +222,6 @@ def test_ndarray_round_trip_with_sidecar_fields(arr, allowed_dtypes, extra): def test_packed_envelope_leads_with_data(): - # load_packed_body finds a blob by byte prefix, so `data` must stay first - # even when sidecar fields are present. 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"] @@ -282,12 +279,6 @@ def test_unpack_rejects_byte_count_mismatched_with_declared_shape(): unpack_ndarray(payload, allowed_dtypes=_FLOAT32, ndim=2) -def test_registry_seeds_both_side_channels(): - # A later branch adds the rollout_sample_support producer; the splice must - # already know the name so that branch stays a pure addition. - assert PACKED_SIDE_CHANNEL_FIELDS == ("routed_experts", "rollout_sample_support") - - 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) @@ -301,7 +292,6 @@ def test_load_packed_body_splices_both_blobs_in_one_body(): choice = load_packed_body(raw)["choices"][0] - # Both blobs are handed over as memoryviews into `raw`, never as Python strs. assert all( isinstance(choice[field][PackedArrayKey.DATA], memoryview) for field in (PackedField.ROUTED_EXPERTS, PackedField.ROLLOUT_SAMPLE_SUPPORT) @@ -364,12 +354,6 @@ def test_load_packed_body_is_not_spoofable_from_a_string_value(): assert np.array_equal(decode_packed_routed_experts(choice[PackedField.ROUTED_EXPERTS]), routes) -def test_load_packed_body_leaves_a_spoofed_prefix_alone_when_no_field_is_present(): - body = _body(note='"routed_experts":{"data":"AAAA"') - - assert load_packed_body(orjson.dumps(body)) == body - - 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) @@ -392,7 +376,6 @@ def test_load_packed_body_splices_one_blob_per_choice(): choices = load_packed_body(raw)["choices"] - # Blobs are matched to envelopes in document order. 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) 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 04320323b6..85b4f09151 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 @@ -71,7 +71,6 @@ 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 - # Hits on the packed-body endpoints below, used to assert retry behaviour. app.state.drifted_body_calls = 0 app.state.flaky_body_calls = 0 @@ -619,12 +618,7 @@ async def test_detokenize(self, client): class TestPackedSideChannelBodies: - """``_post(packed_side_channels=True)`` splices packed blobs out of the raw bytes. - - The splice is what keeps a ~121 MiB base64 blob from becoming a Python - ``str``; a layout it cannot cut must fail instead of quietly falling back to - whole-body parsing. - """ + """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( @@ -636,7 +630,6 @@ 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] - # Both blobs arrive as memoryviews into the raw response, never as strs. 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) From 58ed80b2e73c1838cf880219c70fa6787e00b9f7 Mon Sep 17 00:00:00 2001 From: lila-sync-bot Date: Fri, 14 Aug 2026 21:53:08 +0000 Subject: [PATCH 09/17] perf(r3): route rollout experts to the trainer packed with cu_seqlens --- .../distributed/megatron/model_utils.py | 4 +- .../distributed/megatron/token_metadata.py | 103 ++++- skyrl/backends/skyrl_train/training_batch.py | 88 ++-- .../skyrl_train/utils/packed_tensor.py | 180 ++++++++ .../skyrl_train/utils/replay_utils.py | 58 ++- .../megatron/megatron_model_wrapper.py | 3 +- .../workers/megatron/megatron_worker.py | 16 +- .../skyrl_train/workers/worker_utils.py | 11 +- .../bench_packed_route_collation.py | 225 ++++++++++ skyrl/train/dataset/preprocess.py | 94 +++-- skyrl/train/dataset/replay_buffer.py | 12 +- .../distributed/test_token_metadata.py | 71 ++++ .../gpu/gpu_ci/megatron/test_router_replay.py | 177 +++++++- .../test_token_based_batching_utils.py | 29 +- .../backends/skyrl_train/test_train_batch.py | 69 ++- .../skyrl_train/utils/test_packed_tensor.py | 274 ++++++++++++ .../skyrl_train/utils/test_replay_utils.py | 51 ++- tests/train/dataset/test_preprocess.py | 70 ++-- ...test_packed_route_collation_equivalence.py | 393 ++++++++++++++++++ 19 files changed, 1756 insertions(+), 172 deletions(-) create mode 100644 skyrl/backends/skyrl_train/utils/packed_tensor.py create mode 100644 skyrl/benchmarks/bench_packed_route_collation.py create mode 100644 tests/backends/skyrl_train/utils/test_packed_tensor.py create mode 100644 tests/train/test_packed_route_collation_equivalence.py 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 98059023f1..c70ccdad5e 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py @@ -1,5 +1,6 @@ """Token-aligned metadata layout transforms shared by training features.""" +from collections.abc import Callable, Sequence from dataclasses import dataclass import numpy as np @@ -9,6 +10,7 @@ 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. @@ -100,34 +102,119 @@ 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]``. + + Equivalent to ``align_token_metadata`` on the batch-major padded rectangle the same + rows would occupy, without ever materializing it. This relies on the batch being + purely LEFT padded, so a trajectory's real tokens are one contiguous run and + ``cu_seqlens`` alone locates them -- no per-row mask is needed or consulted. + + Without ``segment_starts`` each segment must hold exactly its trajectory's real tokens, + which ``layout.sequence_lengths`` states independently. With it, segment ``i`` covers + only part of its trajectory and lands at ``segment_starts[i]`` real tokens in, which is + what a response-suffix channel needs. + """ + 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] @@ -136,7 +223,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, ) diff --git a/skyrl/backends/skyrl_train/training_batch.py b/skyrl/backends/skyrl_train/training_batch.py index 1542adeff3..db6c85d88f 100644 --- a/skyrl/backends/skyrl_train/training_batch.py +++ b/skyrl/backends/skyrl_train/training_batch.py @@ -3,24 +3,35 @@ import copy import io import pickle -from typing import Any, Dict, Generic, List, Optional, TypedDict, TypeVar +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 +from skyrl.backends.skyrl_train.utils.replay_utils import append_packed_replay_padding DictType = TypeVar("DictType") +class TensorFormat(StrEnum): + """How one serialized batch field is encoded in the pickle stream.""" + + NUMPY = "numpy" + TORCH = "torch" + TENSOR_LIST = "tensor_list" + PACKED_TENSOR = "packed_tensor" + + def _serialize_tensor(value: torch.Tensor) -> dict: """Serialize a single tensor for pickle protocol.""" try: # Fast path: direct memory copy via numpy (works for most dtypes) arr = value.numpy() return { - "format": "numpy", + "format": TensorFormat.NUMPY, "data": arr.tobytes(), "shape": arr.shape, "dtype": str(arr.dtype), @@ -30,14 +41,14 @@ def _serialize_tensor(value: torch.Tensor) -> dict: buffer = io.BytesIO() torch.save(value, buffer) return { - "format": "torch", + "format": TensorFormat.TORCH, "data": buffer.getvalue(), } def _deserialize_tensor(value: dict) -> torch.Tensor: """Deserialize a single tensor from pickle format.""" - if value.get("format") == "torch": + if value.get("format") == TensorFormat.TORCH: # Fallback path: torch.load for unsupported dtypes buffer = io.BytesIO(value["data"]) return torch.load(buffer, weights_only=True) @@ -110,6 +121,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) @@ -169,8 +187,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 +203,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 +232,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 +241,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 @@ -267,9 +281,15 @@ def __getstate__(self): 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), + "cu_seqlens": _serialize_tensor(value.cu_seqlens), + } else: batch_dict[key] = _serialize_tensor(value) @@ -288,8 +308,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 +339,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 +361,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 +375,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 +402,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 +439,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,7 +506,8 @@ 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) 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 @@ -534,13 +558,9 @@ def pad_training_input_batch(unpadded_batch: TrainingInputBatch, pad_size: int) 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) + # Every other field copies row 0 into the padding rows, so each padded row holds + # as many real tokens as row 0 and needs a route segment of that length. + new_tensors[key] = append_packed_replay_padding(tensor, segment_lengths=[len(tensor.segment(0))] * pad_size) 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..8f560fc5b2 --- /dev/null +++ b/skyrl/backends/skyrl_train/utils/packed_tensor.py @@ -0,0 +1,180 @@ +"""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 the row ids of the segments at ``starts``, laid out back to back. + + ``starts`` and ``lengths`` name source segments in any order; the result gathers them + into one contiguous buffer, so it is the row index a packed gather selects with. + """ + 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. + Every batch operation -- indexing, chunking, concatenation, device transfer -- + addresses segments, so this drops into a ``TensorBatch`` field wherever a + ``[batch, seq_len, ...]`` tensor would otherwise carry per-row padding. + + Segments are contiguous and consecutive, so this describes exactly one ragged level: + a row belongs to the segment whose offset range contains it, and nothing else. + """ + + 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), + ) diff --git a/skyrl/backends/skyrl_train/utils/replay_utils.py b/skyrl/backends/skyrl_train/utils/replay_utils.py index 94696ad8d3..3e22c7afd7 100644 --- a/skyrl/backends/skyrl_train/utils/replay_utils.py +++ b/skyrl/backends/skyrl_train/utils/replay_utils.py @@ -2,14 +2,20 @@ Utility functions for MoE Router Replay. """ +from collections.abc import Sequence from contextlib import contextmanager 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, + cu_seqlens_from_lengths, +) def replay_padding_row( @@ -43,6 +49,33 @@ def make_replay_padding_indices( return padding_row.expand(shape).clone() +def make_packed_replay_padding( + reference: PackedTensor, + *, + segment_lengths: Sequence[int], +) -> PackedTensor: + """Return dummy-route segments matching ``reference``'s row shape and dtype. + + Batch padding rows exist only to give Megatron a uniform micro-batch size; their + tokens are loss-masked, so one dummy route per token is all they need. + """ + padding = make_replay_padding_indices( + (sum(segment_lengths), *reference.row_shape), + dtype=reference.dtype, + device=reference.device, + ) + return PackedTensor(padding, cu_seqlens_from_lengths(segment_lengths, device=reference.device)) + + +def append_packed_replay_padding( + routes: PackedTensor, + *, + segment_lengths: Sequence[int], +) -> PackedTensor: + """Extend ``routes`` with one dummy-route segment per batch padding row.""" + return PackedTensor.cat([routes, make_packed_replay_padding(routes, segment_lengths=segment_lengths)]) + + def patch_topk_router_layer_number(): """Monkey-patch TopKRouter.set_layer_number to propagate the global layer number to the RouterReplay instance. @@ -169,7 +202,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, @@ -179,8 +212,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``. @@ -214,8 +248,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( @@ -225,24 +259,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], + 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/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 7c009260df..3b99902d56 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -46,8 +46,9 @@ TrainingInputBatch, TrainingOutputBatch, ) +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.utils.replay_utils import append_packed_replay_padding from skyrl.backends.skyrl_train.weight_sync import ( LoraLoadRequest, WeightChunk, @@ -783,6 +784,13 @@ def _pad_microbatch_to_size(self, micro_dict: dict, target_batch_size: int) -> d if value is None: padded[key] = None continue + if key == "rollout_expert_indices": + # The dummy attention_mask row below marks one valid token, so each padded + # row's route segment holds one row. + padded[key] = append_packed_replay_padding(value, segment_lengths=[1] * pad_count) + continue + if isinstance(value, PackedTensor): + raise ValueError(f"Micro-batch field {key!r} is packed and has no padding rule") if isinstance(value, torch.Tensor): if key == "loss_mask": # Pad with zeros so padded samples don't contribute to loss @@ -800,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..d3f04b5e79 100644 --- a/skyrl/backends/skyrl_train/workers/worker_utils.py +++ b/skyrl/backends/skyrl_train/workers/worker_utils.py @@ -6,7 +6,7 @@ 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.utils.replay_utils import make_packed_replay_padding 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 @@ -326,11 +326,10 @@ def _create_padding_microbatch(self) -> TrainingInputBatch: 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, + # The dummy attention_mask row marks one valid token, so its route segment holds one row. + data["rollout_expert_indices"] = make_packed_replay_padding( + self.data["rollout_expert_indices"], + segment_lengths=[1] * 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) diff --git a/skyrl/benchmarks/bench_packed_route_collation.py b/skyrl/benchmarks/bench_packed_route_collation.py new file mode 100644 index 0000000000..0b0b922230 --- /dev/null +++ b/skyrl/benchmarks/bench_packed_route_collation.py @@ -0,0 +1,225 @@ +"""Padded-rectangle versus packed R3 route collation. + +The old path allocated ``[batch, max_total, layers, topk]`` and filled the left padding with +dummy routes; the trainer's Megatron layout then compacted it straight back to ragged. This +measures what that rectangle costs versus writing only the real tokens. + +Both fills are carried in a serial and a pooled arm. The pooled padded arm is the real +baseline -- the trainer already pools -- so a serial packed fill has to beat that, not just +the serial padded one, to be a win. The pooled arms default to whatever pool size the +trainer would size on the measuring host; a fixed default measures an arm no trainer runs. + +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. + + ``RoutedExpertTrace.finalize`` self-pads to ``token_count``, so a trajectory's route + array always covers every real token -- the batch's only padding is left padding. + + The arrays are views over one template buffer. Materializing a distinct 55 GiB of source + routes would dwarf the collation under test, and a leading-axis slice is already + C-contiguous, so the copies being measured see exactly the layout production hands them. + """ + 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. + + First-touch page faults on a fresh multi-GiB mapping dominate the fill, so the buffer + must be new every time; a reused one measures only the copies. + """ + 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/dataset/preprocess.py b/skyrl/train/dataset/preprocess.py index f7b0b73e1a..1a59b5e948 100644 --- a/skyrl/train/dataset/preprocess.py +++ b/skyrl/train/dataset/preprocess.py @@ -1,10 +1,15 @@ +import functools import logging 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.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, @@ -95,14 +100,40 @@ def _reward_to_numpy(custom_reward: Union[List[float], torch.Tensor]) -> np.ndar return reward_arr -def _collate_rollout_expert_indices( +def _fill_routed_expert_segment( + packed: torch.Tensor, + cu_seqlens: torch.Tensor, rollout_expert_indices: List[RoutedExpertIndices], - pad_lens: np.ndarray, - max_total: int, -) -> Integer[torch.Tensor, "batch seq_len layer_num topk"]: - """Pack routes into a left-padded ``[batch, seq_len, layers, topk]`` buffer. + sample_index: int, +) -> None: + """Write one trajectory's segment of the packed route buffer. - The sender establishes canonical dtypes, so this path validates rather than rescans entries. + vLLM may capture no route for a trailing token, so a segment can end with dummy rows. + Those hold ``topk`` distinct experts to keep Megatron's dropless dispatcher supplied with + one row per token without collapsing its router accounting. + """ + 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. + + ``pack_routed_experts`` establishes the canonical dtype on the sending side, so entries are + validated rather than rescanned here. Every region of the buffer is written exactly once, from + a locally sized thread pool: this fill runs on the whole global batch before DP sharding, on + every training step, and is bound by first-touch page faults on a fresh multi-GiB mapping. + Packing shrinks that mapping but does not make the fill cheap, so the pool has to survive it. """ num_samples = len(rollout_expert_indices) for sample_index, sample_indices in enumerate(rollout_expert_indices): @@ -127,21 +158,18 @@ def _collate_rollout_expert_indices( if topk < 1: raise ValueError("rollout_expert_indices must contain at least one expert per layer") - # Validate before dispatch so errors are deterministic. - route_ends = [] + # 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}" ) - left_pad = int(pad_lens[sample_index]) - available = max_total - left_pad + 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" ) - route_ends.append(left_pad + sample_indices.shape[0]) batch_dtype = max((indices.dtype for indices in rollout_expert_indices), key=lambda dtype: dtype.itemsize) if batch_dtype == np.dtype(np.int32): @@ -149,28 +177,19 @@ def _collate_rollout_expert_indices( "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." ) - padded = torch.empty( - (num_samples, max_total, num_layers, topk), + # Routes stay packed to the real tokens they describe. A [batch, max_total, layers, topk] + # rectangle instead reaches ~55 GiB per global batch at a 120B route shape, and the trainer's + # Megatron layout compacts it straight back to ragged. + 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], ) - # Distinct experts per padding row, as Megatron's dropless dispatcher requires. - padding_row = replay_padding_row(topk, dtype=padded.dtype) - - def fill_sample(sample_index: int) -> None: - sample_indices = rollout_expert_indices[sample_index] - flags = sample_indices.flags - # torch.from_numpy requires a writable buffer. - if not flags.c_contiguous or not flags.writeable: - sample_indices = sample_indices.copy(order="C") - left_pad = int(pad_lens[sample_index]) - route_end = route_ends[sample_index] - sample_rows = padded[sample_index] - sample_rows[:left_pad] = padding_row - sample_rows[left_pad:route_end] = torch.from_numpy(sample_indices) - sample_rows[route_end:] = padding_row - - fill_batch_rows(fill_sample, num_samples) - return padded + fill_batch_rows( + functools.partial(_fill_routed_expert_segment, packed, cu_seqlens, rollout_expert_indices), + num_samples, + ) + return PackedTensor(packed, cu_seqlens) def convert_prompts_responses_to_batch_tensors( @@ -189,7 +208,7 @@ 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], ]: """ Convert prompts and responses to batch tensors for training. @@ -246,6 +265,9 @@ 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``. """ _verify_inputs(prompts, responses, rewards, loss_masks) @@ -321,11 +343,7 @@ 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") - rollout_expert_indices_tensor = _collate_rollout_expert_indices( - rollout_expert_indices, - pad_lens, - max_total, - ) + rollout_expert_indices_tensor = _collate_rollout_expert_indices(rollout_expert_indices, total_real) return ( sequences, diff --git a/skyrl/train/dataset/replay_buffer.py b/skyrl/train/dataset/replay_buffer.py index 6f627e31ed..35668a310a 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,7 +68,8 @@ 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 diff --git a/tests/backends/skyrl_train/distributed/test_token_metadata.py b/tests/backends/skyrl_train/distributed/test_token_metadata.py index 73ba160579..e3e7660567 100644 --- a/tests/backends/skyrl_train/distributed/test_token_metadata.py +++ b/tests/backends/skyrl_train/distributed/test_token_metadata.py @@ -9,6 +9,7 @@ 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 @@ -151,3 +152,73 @@ def test_routed_expert_trace_only_pads_masked_suffix(active: bool) -> None: else: result = trace.finalize(token_count=5, loss_mask=mask) assert np.array_equal(result[-2:, 0], [[0, 1], [0, 1]]) + + +@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) 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..89bfd85515 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: @@ -353,22 +373,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 +426,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/test_token_based_batching_utils.py b/tests/backends/skyrl_train/test_token_based_batching_utils.py index 4fac8a6679..e68fc499a6 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,10 @@ 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.workers.worker_utils import ( TokenBasedBatchIterator, get_microbatch_iterator, @@ -201,16 +205,35 @@ def test_padding_microbatch_matches_seq_len(self): 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["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) 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"] + # The dummy attention_mask row marks one valid token, so one dummy route per row. + 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_microbatch_selection_gathers_packed_route_segments(self): + """Token-based microbatching must select route segments alongside the dense rows.""" + 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..4786a1d55a 100644 --- a/tests/backends/skyrl_train/test_train_batch.py +++ b/tests/backends/skyrl_train/test_train_batch.py @@ -7,11 +7,16 @@ from skyrl.backends.skyrl_train.training_batch import ( TensorBatch, + TensorFormat, TensorList, TrainingInput, TrainingInputBatch, pad_training_input_batch, ) +from skyrl.backends.skyrl_train.utils.packed_tensor import ( + PackedTensor, + cu_seqlens_from_lengths, +) def test_train_batch_initialization(): @@ -576,7 +581,12 @@ 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), + # Routes arrive packed to real tokens; this fixture is fully attended, so every + # segment holds 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), "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,9 +650,12 @@ 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:] + # Each padded row copies row 0, so its route segment matches row 0's length. + 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) regular_tensor_keys = EXPECTED_TRAINING_INPUT_FIELDS - { "loss_mask", @@ -711,3 +724,51 @@ 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_named_by_the_tensor_format_enum(): + """Every branch of ``__setstate__`` keys off a ``TensorFormat`` member, not a bare string.""" + 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 + # Legacy pickles carry the plain strings; StrEnum members must keep matching them. + assert [format.value for format in TensorFormat] == ["numpy", "torch", "tensor_list", "packed_tensor"] 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..80ad4db7fd --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_packed_tensor.py @@ -0,0 +1,274 @@ +"""Batch operations on ``PackedTensor`` must match the equivalent list-of-segments result. + +uv run --isolated --extra dev pytest tests/backends/skyrl_train/utils/test_packed_tensor.py +""" + +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 + + +# --------------------------------------------------------------------------- +# Offsets algebra +# --------------------------------------------------------------------------- + + +def test_cu_seqlens_stay_int32_through_the_prefix_sum(): + offsets = cu_seqlens_from_lengths(SEGMENT_LENGTHS) + + assert offsets.dtype == CU_SEQLENS_DTYPE + assert offsets.tolist() == [0, 3, 4, 8, 10] + + +def test_offsets_algebra_round_trips_lengths(): + offsets = cu_seqlens_from_lengths(SEGMENT_LENGTHS) + + assert lengths_from_offsets(offsets).tolist() == SEGMENT_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_cu_seqlens_tolerate_zero_length_segments(): + offsets = cu_seqlens_from_lengths([0, 2, 0]) + + assert offsets.tolist() == [0, 0, 2, 2] + assert lengths_from_offsets(offsets).tolist() == [0, 2, 0] + + +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] + + +# --------------------------------------------------------------------------- +# Container surface +# --------------------------------------------------------------------------- + + +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_repr_names_the_batch_and_row_shape(): + assert ( + repr(PackedTensor.from_segments(_segments())) == "PackedTensor(batch=4, values=(10, 2, 3), dtype=torch.int16)" + ) + + +# --------------------------------------------------------------------------- +# Construction guards +# --------------------------------------------------------------------------- + + +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)) + + +# --------------------------------------------------------------------------- +# Aliasing contract +# --------------------------------------------------------------------------- + + +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 beb3a59b0f..20aa077605 100644 --- a/tests/backends/skyrl_train/utils/test_replay_utils.py +++ b/tests/backends/skyrl_train/utils/test_replay_utils.py @@ -10,7 +10,17 @@ 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 +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor +from skyrl.backends.skyrl_train.utils.replay_utils import ( + append_packed_replay_padding, + make_packed_replay_padding, + 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 @@ -72,6 +82,34 @@ def test_replay_padding_rejects_missing_topk(shape): make_replay_padding_indices(shape, dtype=torch.uint8) +@pytest.mark.parametrize("segment_lengths", [[1, 1], [3], [2, 5, 1]]) +def test_packed_replay_padding_matches_the_reference_row_shape(segment_lengths): + reference = PackedTensor.from_segments([torch.full((4, 2, 3), 9, dtype=torch.int16)]) + + padding = make_packed_replay_padding(reference, segment_lengths=segment_lengths) + + assert padding.sequence_lengths.tolist() == segment_lengths + assert padding.row_shape == reference.row_shape + assert padding.dtype == reference.dtype + # Distinct experts per dummy row, as Megatron's dropless dispatcher requires. + assert torch.equal(padding.values, torch.tensor([0, 1, 2], dtype=torch.int16).expand_as(padding.values)) + + +@pytest.mark.parametrize("pad_count", [1, 3]) +def test_appending_replay_padding_keeps_the_real_segments_and_the_arange_invariant(pad_count): + """Both batch-padding sites append through here, so the round trip is asserted once.""" + routes = PackedTensor.from_segments( + [torch.full((4, 2, 3), 9, dtype=torch.int16), torch.full((2, 2, 3), 8, dtype=torch.int16)] + ) + + padded = append_packed_replay_padding(routes, segment_lengths=[1] * pad_count) + + assert padded.sequence_lengths.tolist() == [4, 2] + [1] * pad_count + assert padded[: len(routes)] == routes + appended = padded[len(routes) :] + assert torch.equal(appended.values, torch.tensor([0, 1, 2], dtype=torch.int16).expand_as(appended.values)) + + def test_replay_has_no_dispatcher_specific_patch(): assert "TokenDispatcher" not in inspect.getsource(replay_utils) @@ -105,15 +143,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( [ @@ -136,7 +173,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(), @@ -199,7 +236,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/train/dataset/test_preprocess.py b/tests/train/dataset/test_preprocess.py index 89545b85e4..f2399c1333 100644 --- a/tests/train/dataset/test_preprocess.py +++ b/tests/train/dataset/test_preprocess.py @@ -98,9 +98,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( @@ -128,7 +129,7 @@ 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): @@ -145,7 +146,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): @@ -176,7 +177,7 @@ def test_routed_expert_tensor_accepts_non_contiguous_arrays(tokenizer): ) assert routed.dtype == torch.uint8 - assert routed.tolist() == [[[[1, 2]], [[3, 4]]]] + assert routed.segment(0).tolist() == [[[1, 2]], [[3, 4]]] @pytest.mark.parametrize("dtype", [np.uint16, np.int64]) @@ -209,7 +210,7 @@ def test_routed_expert_tensor_keeps_the_sender_dtype_after_truncation(tokenizer) ) assert routed.dtype == torch.int16 - assert routed.tolist() == [[[[1, 2]], [[3, 4]]]] + assert routed.segment(0).tolist() == [[[1, 2]], [[3, 4]]] @pytest.mark.parametrize( @@ -255,10 +256,15 @@ def _numpy_padded_routes( def test_routed_expert_tensor_is_bit_identical_to_numpy_collation(tokenizer): + """A short route prefix and trailing padding, across a mixed-dtype batch. + + The packed buffer must hold exactly the real-token rows of the rectangle NumPy built, + left padding dropped rather than written and read back. + """ prompts = [[1, 2], [3, 4, 5, 6]] responses = [[10, 11, 12], [20, 21]] num_layers, topk = 2, 3 - # Sample 0 has fewer route rows than tokens and needs padding on both sides. + # 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), @@ -273,12 +279,21 @@ def test_routed_expert_tensor_is_bit_identical_to_numpy_collation(tokenizer): 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 torch.equal(routed, torch.from_numpy(_numpy_padded_routes(routes, prompts, responses))) - # Padding routes must select distinct experts for the dropless dispatcher. + assert routed.cu_seqlens.tolist() == [0, 5, 11] + assert torch.equal(routed.values, torch.from_numpy(real_rows)) + # Padding routes are topk distinct experts, which Megatron's dropless dispatcher requires; + # zeros would collapse them onto one expert. padding_row = [[0, 1, 2]] * num_layers - assert routed[0, 0].tolist() == padding_row - assert routed[0, 5].tolist() == padding_row + 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): @@ -496,16 +511,14 @@ def test_max_seq_len_warns_but_does_not_truncate(tokenizer): # --------------------------------------------------------------------------- -# R3 (Router Replay) — rollout_expert_indices padding tests +# R3 (Router Replay) — packed rollout_expert_indices 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.""" + """Routes pack to [sum(seq_len), layers, topk] with one cu_seqlens segment per trajectory.""" # Sample 0: prompt=2, response=3 → total=5 # Sample 1: prompt=4, response=2 → total=6 - # max_total=6 prompts = [[1, 2], [3, 4, 5, 6]] responses = [[10, 11, 12], [20, 21]] rewards = [[0.0] * 3, [0.0] * 2] @@ -514,7 +527,6 @@ 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 @@ -528,24 +540,14 @@ 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 + # Packed to the batch's real tokens: 5 + 6, never the 2 x 6 rectangle. + assert rei_tensor.values.shape == (11, num_layers, topk) + assert rei_tensor.cu_seqlens.tolist() == [0, 5, 11] + + # Segment lengths match each trajectory's real-token count in the attention mask. + 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): 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..96337a35d3 --- /dev/null +++ b/tests/train/test_packed_route_collation_equivalence.py @@ -0,0 +1,393 @@ +"""Bit-identity of packed R3 route collation against the padded-rectangle path. + +The packed buffer plus ``cu_seqlens`` replaces a ``[batch, max_total, layers, topk]`` +rectangle that ``align_token_metadata`` immediately compacted back to ragged. What +Megatron receives must not change, so both paths are run end to end here -- collation, +micro-batch padding, layer selection, packing/CP layout, and the TP slice -- and the +per-layer tensors handed to ``RouterReplay.set_replay_data`` compared. + +Run with: + uv run --isolated --extra dev pytest tests/train/test_packed_route_collation_equivalence.py +""" + +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 + + +# --------------------------------------------------------------------------- +# Reference: the padded [batch, max_total, layers, topk] rectangle +# --------------------------------------------------------------------------- + + +def _reference_padded_routes( + routes: list[np.ndarray], + prompt_lens: list[int], + response_lens: list[int], +) -> torch.Tensor: + """The pre-change collation: a left-padded batch-major rectangle.""" + 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]: + """The pre-change trainer path: gather real tokens out of the padded rectangle.""" + 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) + + +# --------------------------------------------------------------------------- +# Fixtures and harness +# --------------------------------------------------------------------------- + + +@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, + ) = 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"] + # The old path copied row 0 into each padding row and filled its routes with + # dummies; reproduce exactly that rectangle for the reference. + 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 + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +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("packed", [False, True]) +def test_packed_routes_match_with_uncaptured_suffix(monkeypatch, parallel_state, router_replay, packed): + """A trajectory whose trailing tokens have no captured route needs identical dummy rows.""" + new_data, reference = _run_both_paths( + LENGTH_DISTRIBUTIONS["typical_rl"], + packed=packed, + tp_size=1, + local_layers=list(range(NUM_LAYERS)), + captured_shortfall=3, + monkeypatch=monkeypatch, + parallel_state=parallel_state, + router_replay=router_replay, + ) + _assert_bit_identical(new_data, reference) + + +@pytest.mark.parametrize("packed", [False, True]) +def test_packed_routes_match_under_batch_padding(monkeypatch, parallel_state, router_replay, packed): + """``pad_training_input_batch`` rows must land on the same dummy routes.""" + new_data, reference = _run_both_paths( + LENGTH_DISTRIBUTIONS["mild_ragged"], + packed=packed, + tp_size=1, + local_layers=list(range(NUM_LAYERS)), + batch_pad_size=2, + monkeypatch=monkeypatch, + parallel_state=parallel_state, + router_replay=router_replay, + ) + _assert_bit_identical(new_data, reference) + + +@pytest.mark.parametrize("packed", [False, True]) +def test_packed_routes_match_on_a_pipeline_stage_subset(monkeypatch, parallel_state, router_replay, packed): + """Under PP a stage owns a subset of routers, so layer selection must agree too.""" + # A stage starting at layer 1 with 2 routers owns captured layers 1 and 2, not 0. + new_data, reference = _run_both_paths( + LENGTH_DISTRIBUTIONS["typical_rl"], + packed=packed, + tp_size=1, + local_layers=[1, 2], + stage_range=(1, 2), + 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 From c1e71a0b8e319e47347b7764595e778558436a09 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Sat, 22 Aug 2026 00:06:58 +0000 Subject: [PATCH 10/17] style: tighten packed-route comments and tests --- .../distributed/megatron/token_metadata.py | 12 +- .../skyrl_train/utils/packed_tensor.py | 13 +- .../bench_packed_route_collation.py | 21 +--- skyrl/train/dataset/preprocess.py | 16 +-- .../test_token_based_batching_utils.py | 2 - .../backends/skyrl_train/test_train_batch.py | 9 +- .../skyrl_train/utils/test_packed_tensor.py | 58 ++------- .../skyrl_train/utils/test_replay_utils.py | 2 - tests/train/dataset/test_preprocess.py | 24 +--- ...test_packed_route_collation_equivalence.py | 115 +++++------------- 10 files changed, 56 insertions(+), 216 deletions(-) diff --git a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py index c70ccdad5e..3257bd57e8 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py @@ -122,15 +122,9 @@ def align_packed_token_metadata( ) -> torch.Tensor: """Align metadata that already arrives packed as ``[sum(seqlen), *row_shape]``. - Equivalent to ``align_token_metadata`` on the batch-major padded rectangle the same - rows would occupy, without ever materializing it. This relies on the batch being - purely LEFT padded, so a trajectory's real tokens are one contiguous run and - ``cu_seqlens`` alone locates them -- no per-row mask is needed or consulted. - - Without ``segment_starts`` each segment must hold exactly its trajectory's real tokens, - which ``layout.sequence_lengths`` states independently. With it, segment ``i`` covers - only part of its trajectory and lands at ``segment_starts[i]`` real tokens in, which is - what a response-suffix channel needs. + 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") diff --git a/skyrl/backends/skyrl_train/utils/packed_tensor.py b/skyrl/backends/skyrl_train/utils/packed_tensor.py index 8f560fc5b2..c5a0c20983 100644 --- a/skyrl/backends/skyrl_train/utils/packed_tensor.py +++ b/skyrl/backends/skyrl_train/utils/packed_tensor.py @@ -34,11 +34,7 @@ def row_index_from_offsets( starts: torch.Tensor, lengths: torch.Tensor, ) -> torch.Tensor: - """Return the row ids of the segments at ``starts``, laid out back to back. - - ``starts`` and ``lengths`` name source segments in any order; the result gathers them - into one contiguous buffer, so it is the row index a packed gather selects with. - """ + """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()) @@ -57,12 +53,7 @@ class PackedTensor: ``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. - Every batch operation -- indexing, chunking, concatenation, device transfer -- - addresses segments, so this drops into a ``TensorBatch`` field wherever a - ``[batch, seq_len, ...]`` tensor would otherwise carry per-row padding. - - Segments are contiguous and consecutive, so this describes exactly one ragged level: - a row belongs to the segment whose offset range contains it, and nothing else. + Indexing and batch operations address segments rather than individual rows. """ def __init__(self, values: torch.Tensor, cu_seqlens: torch.Tensor): diff --git a/skyrl/benchmarks/bench_packed_route_collation.py b/skyrl/benchmarks/bench_packed_route_collation.py index 0b0b922230..668a9d157d 100644 --- a/skyrl/benchmarks/bench_packed_route_collation.py +++ b/skyrl/benchmarks/bench_packed_route_collation.py @@ -1,13 +1,4 @@ -"""Padded-rectangle versus packed R3 route collation. - -The old path allocated ``[batch, max_total, layers, topk]`` and filled the left padding with -dummy routes; the trainer's Megatron layout then compacted it straight back to ragged. This -measures what that rectangle costs versus writing only the real tokens. - -Both fills are carried in a serial and a pooled arm. The pooled padded arm is the real -baseline -- the trainer already pools -- so a serial packed fill has to beat that, not just -the serial padded one, to be a win. The pooled arms default to whatever pool size the -trainer would size on the measuring host; a fixed default measures an arm no trainer runs. +"""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:: @@ -67,12 +58,7 @@ def _sequence_lengths(distribution: str, num_sequences: int, max_seqlen: int, se 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. - ``RoutedExpertTrace.finalize`` self-pads to ``token_count``, so a trajectory's route - array always covers every real token -- the batch's only padding is left padding. - - The arrays are views over one template buffer. Materializing a distinct 55 GiB of source - routes would dwarf the collation under test, and a leading-axis slice is already - C-contiguous, so the copies being measured see exactly the layout production hands them. + 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) @@ -139,8 +125,7 @@ def _peak_rss_bytes() -> int: def _time_cold(fill, shape, trajectories, lengths, iterations: int) -> tuple[float, int]: """Median wall clock over a freshly allocated buffer each iteration. - First-touch page faults on a fresh multi-GiB mapping dominate the fill, so the buffer - must be new every time; a reused one measures only the copies. + Fresh buffers retain the first-touch allocation cost measured in production. """ durations = [] for _ in range(iterations): diff --git a/skyrl/train/dataset/preprocess.py b/skyrl/train/dataset/preprocess.py index 1a59b5e948..23aa0a123c 100644 --- a/skyrl/train/dataset/preprocess.py +++ b/skyrl/train/dataset/preprocess.py @@ -106,12 +106,7 @@ def _fill_routed_expert_segment( rollout_expert_indices: List[RoutedExpertIndices], sample_index: int, ) -> None: - """Write one trajectory's segment of the packed route buffer. - - vLLM may capture no route for a trailing token, so a segment can end with dummy rows. - Those hold ``topk`` distinct experts to keep Megatron's dropless dispatcher supplied with - one row per token without collapsing its router accounting. - """ + """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. @@ -129,11 +124,7 @@ def _collate_rollout_expert_indices( ) -> PackedTensor: """Pack per-trajectory routes into one ``[sum(seq_len_i), layers, topk]`` buffer. - ``pack_routed_experts`` establishes the canonical dtype on the sending side, so entries are - validated rather than rescanned here. Every region of the buffer is written exactly once, from - a locally sized thread pool: this fill runs on the whole global batch before DP sharding, on - every training step, and is bound by first-touch page faults on a fresh multi-GiB mapping. - Packing shrinks that mapping but does not make the fill cheap, so the pool has to survive it. + 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): @@ -177,9 +168,6 @@ def _collate_rollout_expert_indices( "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." ) - # Routes stay packed to the real tokens they describe. A [batch, max_total, layers, topk] - # rectangle instead reaches ~55 GiB per global batch at a 120B route shape, and the trainer's - # Megatron layout compacts it straight back to ragged. cu_seqlens = cu_seqlens_from_lengths(total_real) packed = torch.empty( (int(total_real.sum()), num_layers, topk), 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 e68fc499a6..c78d0fee80 100644 --- a/tests/backends/skyrl_train/test_token_based_batching_utils.py +++ b/tests/backends/skyrl_train/test_token_based_batching_utils.py @@ -215,14 +215,12 @@ def test_padding_microbatch_uses_unique_dummy_routes(self): padding = iterator._create_padding_microbatch() padded_routes = padding["rollout_expert_indices"] - # The dummy attention_mask row marks one valid token, so one dummy route per row. 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_microbatch_selection_gathers_packed_route_segments(self): - """Token-based microbatching must select route segments alongside the dense rows.""" 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)] diff --git a/tests/backends/skyrl_train/test_train_batch.py b/tests/backends/skyrl_train/test_train_batch.py index 4786a1d55a..9aae1c5bb7 100644 --- a/tests/backends/skyrl_train/test_train_batch.py +++ b/tests/backends/skyrl_train/test_train_batch.py @@ -581,8 +581,7 @@ 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), - # Routes arrive packed to real tokens; this fixture is fully attended, so every - # segment holds seq_len rows. + # 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), @@ -652,7 +651,6 @@ def test_pad_batch_all_fields(): assert torch.all(padded["router_padding_mask"][batch_size:]) assert padded["rollout_expert_indices"][:batch_size] == batch["rollout_expert_indices"] padded_routes = padded["rollout_expert_indices"][batch_size:] - # Each padded row copies row 0, so its route segment matches row 0's length. 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) @@ -751,8 +749,7 @@ def test_packed_tensor_field_survives_the_ray_pickle_round_trip(): assert unpickled == data -def test_serialized_field_formats_are_named_by_the_tensor_format_enum(): - """Every branch of ``__setstate__`` keys off a ``TensorFormat`` member, not a bare string.""" +def test_serialized_field_formats_are_stable(): data = TensorBatch( { "sequences": torch.randn(2, 4), @@ -770,5 +767,3 @@ def test_serialized_field_formats_are_named_by_the_tensor_format_enum(): assert state["bf16_logprobs"]["format"] == TensorFormat.TORCH assert state["pixel_values"]["format"] == TensorFormat.TENSOR_LIST assert state["rollout_expert_indices"]["format"] == TensorFormat.PACKED_TENSOR - # Legacy pickles carry the plain strings; StrEnum members must keep matching them. - assert [format.value for format in TensorFormat] == ["numpy", "torch", "tensor_list", "packed_tensor"] diff --git a/tests/backends/skyrl_train/utils/test_packed_tensor.py b/tests/backends/skyrl_train/utils/test_packed_tensor.py index 80ad4db7fd..f0101cb23a 100644 --- a/tests/backends/skyrl_train/utils/test_packed_tensor.py +++ b/tests/backends/skyrl_train/utils/test_packed_tensor.py @@ -1,7 +1,4 @@ -"""Batch operations on ``PackedTensor`` must match the equivalent list-of-segments result. - -uv run --isolated --extra dev pytest tests/backends/skyrl_train/utils/test_packed_tensor.py -""" +"""Tests for ``PackedTensor`` batch operations.""" import pytest import torch @@ -28,22 +25,19 @@ def _segments(lengths=SEGMENT_LENGTHS, *, row_shape=(2, 3)) -> list[torch.Tensor return segments -# --------------------------------------------------------------------------- -# Offsets algebra -# --------------------------------------------------------------------------- - - -def test_cu_seqlens_stay_int32_through_the_prefix_sum(): - offsets = cu_seqlens_from_lengths(SEGMENT_LENGTHS) +@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 offsets.tolist() == [0, 3, 4, 8, 10] - - -def test_offsets_algebra_round_trips_lengths(): - offsets = cu_seqlens_from_lengths(SEGMENT_LENGTHS) - - assert lengths_from_offsets(offsets).tolist() == SEGMENT_LENGTHS + assert lengths_from_offsets(offsets).tolist() == lengths assert lengths_from_offsets(offsets).dtype == CU_SEQLENS_DTYPE @@ -54,13 +48,6 @@ def test_cu_seqlens_reject_negative_lengths_and_extra_dimensions(): cu_seqlens_from_lengths(torch.zeros((2, 2), dtype=torch.int32)) -def test_cu_seqlens_tolerate_zero_length_segments(): - offsets = cu_seqlens_from_lengths([0, 2, 0]) - - assert offsets.tolist() == [0, 0, 2, 2] - assert lengths_from_offsets(offsets).tolist() == [0, 2, 0] - - def test_row_index_from_offsets_lays_selected_segments_back_to_back(): starts = torch.tensor([8, 0]) lengths = torch.tensor([2, 3]) @@ -68,11 +55,6 @@ def test_row_index_from_offsets_lays_selected_segments_back_to_back(): assert row_index_from_offsets(starts, lengths).tolist() == [8, 9, 0, 1, 2] -# --------------------------------------------------------------------------- -# Container surface -# --------------------------------------------------------------------------- - - def test_from_segments_round_trips_every_segment(): segments = _segments() packed = PackedTensor.from_segments(segments) @@ -202,17 +184,6 @@ def test_equality_compares_values_and_offsets(): assert packed != packed.values -def test_repr_names_the_batch_and_row_shape(): - assert ( - repr(PackedTensor.from_segments(_segments())) == "PackedTensor(batch=4, values=(10, 2, 3), dtype=torch.int16)" - ) - - -# --------------------------------------------------------------------------- -# Construction guards -# --------------------------------------------------------------------------- - - def test_rejects_mismatched_device_or_offset_dtype(): values = torch.zeros((4, 2), dtype=torch.int16) @@ -238,11 +209,6 @@ def test_rejects_values_without_a_token_row_dimension(): PackedTensor(torch.tensor(1), torch.tensor([0, 1], dtype=CU_SEQLENS_DTYPE)) -# --------------------------------------------------------------------------- -# Aliasing contract -# --------------------------------------------------------------------------- - - 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()) diff --git a/tests/backends/skyrl_train/utils/test_replay_utils.py b/tests/backends/skyrl_train/utils/test_replay_utils.py index 20aa077605..f713908a92 100644 --- a/tests/backends/skyrl_train/utils/test_replay_utils.py +++ b/tests/backends/skyrl_train/utils/test_replay_utils.py @@ -91,13 +91,11 @@ def test_packed_replay_padding_matches_the_reference_row_shape(segment_lengths): assert padding.sequence_lengths.tolist() == segment_lengths assert padding.row_shape == reference.row_shape assert padding.dtype == reference.dtype - # Distinct experts per dummy row, as Megatron's dropless dispatcher requires. assert torch.equal(padding.values, torch.tensor([0, 1, 2], dtype=torch.int16).expand_as(padding.values)) @pytest.mark.parametrize("pad_count", [1, 3]) def test_appending_replay_padding_keeps_the_real_segments_and_the_arange_invariant(pad_count): - """Both batch-padding sites append through here, so the round trip is asserted once.""" routes = PackedTensor.from_segments( [torch.full((4, 2, 3), 9, dtype=torch.int16), torch.full((2, 2, 3), 8, dtype=torch.int16)] ) diff --git a/tests/train/dataset/test_preprocess.py b/tests/train/dataset/test_preprocess.py index f2399c1333..91e7db7b1e 100644 --- a/tests/train/dataset/test_preprocess.py +++ b/tests/train/dataset/test_preprocess.py @@ -256,11 +256,7 @@ def _numpy_padded_routes( def test_routed_expert_tensor_is_bit_identical_to_numpy_collation(tokenizer): - """A short route prefix and trailing padding, across a mixed-dtype batch. - - The packed buffer must hold exactly the real-token rows of the rectangle NumPy built, - left padding dropped rather than written and read back. - """ + """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 @@ -289,8 +285,7 @@ def test_routed_expert_tensor_is_bit_identical_to_numpy_collation(tokenizer): 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 are topk distinct experts, which Megatron's dropless dispatcher requires; - # zeros would collapse them onto one expert. + # 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])) @@ -510,15 +505,8 @@ def test_max_seq_len_warns_but_does_not_truncate(tokenizer): assert action.shape == (2, 50) -# --------------------------------------------------------------------------- -# R3 (Router Replay) — packed rollout_expert_indices tests -# --------------------------------------------------------------------------- - - def test_rollout_expert_indices_shape_padding_and_alignment(tokenizer): """Routes pack to [sum(seq_len), layers, topk] with one cu_seqlens segment per trajectory.""" - # Sample 0: prompt=2, response=3 → total=5 - # Sample 1: prompt=4, response=2 → total=6 prompts = [[1, 2], [3, 4, 5, 6]] responses = [[10, 11, 12], [20, 21]] rewards = [[0.0] * 3, [0.0] * 2] @@ -526,9 +514,8 @@ 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] - 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( tokenizer.pad_token_id, @@ -540,11 +527,8 @@ def test_rollout_expert_indices_shape_padding_and_alignment(tokenizer): ) assert rei_tensor is not None - # Packed to the batch's real tokens: 5 + 6, never the 2 x 6 rectangle. assert rei_tensor.values.shape == (11, num_layers, topk) assert rei_tensor.cu_seqlens.tolist() == [0, 5, 11] - - # Segment lengths match each trajectory's real-token count in the attention mask. 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 diff --git a/tests/train/test_packed_route_collation_equivalence.py b/tests/train/test_packed_route_collation_equivalence.py index 96337a35d3..5d8c9ac679 100644 --- a/tests/train/test_packed_route_collation_equivalence.py +++ b/tests/train/test_packed_route_collation_equivalence.py @@ -1,14 +1,4 @@ -"""Bit-identity of packed R3 route collation against the padded-rectangle path. - -The packed buffer plus ``cu_seqlens`` replaces a ``[batch, max_total, layers, topk]`` -rectangle that ``align_token_metadata`` immediately compacted back to ragged. What -Megatron receives must not change, so both paths are run end to end here -- collation, -micro-batch padding, layer selection, packing/CP layout, and the TP slice -- and the -per-layer tensors handed to ``RouterReplay.set_replay_data`` compared. - -Run with: - uv run --isolated --extra dev pytest tests/train/test_packed_route_collation_equivalence.py -""" +"""Compare packed route collation with the padded reference path end to end.""" import sys import types @@ -44,17 +34,12 @@ MIN_EXPERT_ID = 300 -# --------------------------------------------------------------------------- -# Reference: the padded [batch, max_total, layers, topk] rectangle -# --------------------------------------------------------------------------- - - def _reference_padded_routes( routes: list[np.ndarray], prompt_lens: list[int], response_lens: list[int], ) -> torch.Tensor: - """The pre-change collation: a left-padded batch-major rectangle.""" + """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 @@ -78,7 +63,7 @@ def _reference_replay_data( tp_size: int, tp_rank: int, ) -> list[torch.Tensor]: - """The pre-change trainer path: gather real tokens out of the padded rectangle.""" + """Gather real-token routes from the padded reference tensor.""" layout = build_token_metadata_layout( attention_mask, padded_routes.device, @@ -97,11 +82,6 @@ def _reference_replay_data( return _split_replay_indices(aligned) -# --------------------------------------------------------------------------- -# Fixtures and harness -# --------------------------------------------------------------------------- - - @pytest.fixture def parallel_state(monkeypatch): try: @@ -227,8 +207,7 @@ def _run_both_paths( attention_mask = batch["attention_mask"] router_padding_mask = batch["router_padding_mask"] packed_routes = batch["rollout_expert_indices"] - # The old path copied row 0 into each padding row and filled its routes with - # dummies; reproduce exactly that rectangle for the reference. + # Match the dummy rows added by batch padding in the packed path. padded_routes = torch.cat( [ padded_routes, @@ -274,10 +253,6 @@ def _assert_bit_identical(new_data: list[torch.Tensor], reference: list[torch.Te assert torch.equal(produced, expected), slot -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - LENGTH_DISTRIBUTIONS = { # No padding at all: the packed and padded layouts coincide. "uniform": [(8, 8), (8, 8), (8, 8), (8, 8)], @@ -306,48 +281,34 @@ def test_packed_routes_match_padded_rectangle( _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_with_uncaptured_suffix(monkeypatch, parallel_state, router_replay, packed): - """A trajectory whose trailing tokens have no captured route needs identical dummy rows.""" - new_data, reference = _run_both_paths( - LENGTH_DISTRIBUTIONS["typical_rl"], - packed=packed, - tp_size=1, - local_layers=list(range(NUM_LAYERS)), - captured_shortfall=3, - monkeypatch=monkeypatch, - parallel_state=parallel_state, - router_replay=router_replay, - ) - _assert_bit_identical(new_data, reference) - - -@pytest.mark.parametrize("packed", [False, True]) -def test_packed_routes_match_under_batch_padding(monkeypatch, parallel_state, router_replay, packed): - """``pad_training_input_batch`` rows must land on the same dummy routes.""" - new_data, reference = _run_both_paths( - LENGTH_DISTRIBUTIONS["mild_ragged"], - packed=packed, - tp_size=1, - local_layers=list(range(NUM_LAYERS)), - batch_pad_size=2, - monkeypatch=monkeypatch, - parallel_state=parallel_state, - router_replay=router_replay, - ) - _assert_bit_identical(new_data, reference) - - -@pytest.mark.parametrize("packed", [False, True]) -def test_packed_routes_match_on_a_pipeline_stage_subset(monkeypatch, parallel_state, router_replay, packed): - """Under PP a stage owns a subset of routers, so layer selection must agree too.""" - # A stage starting at layer 1 with 2 routers owns captured layers 1 and 2, not 0. +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["typical_rl"], + LENGTH_DISTRIBUTIONS[distribution], packed=packed, tp_size=1, - local_layers=[1, 2], - stage_range=(1, 2), + 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, @@ -371,23 +332,3 @@ def test_packed_routes_match_under_context_parallelism(monkeypatch, parallel_sta 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 From 646f87f02b78f429b0b4d96c8e0bd51a152ed7ce Mon Sep 17 00:00:00 2001 From: lila-sync-bot Date: Fri, 14 Aug 2026 22:11:49 +0000 Subject: [PATCH 11/17] perf(batch): ship large per-token side channels to the trainer zero-copy --- skyrl/backends/skyrl_train/training_batch.py | 54 +++++++-- tests/backends/skyrl_train/conftest.py | 24 ++++ .../backends/skyrl_train/test_train_batch.py | 106 ++++++++++++++++++ 3 files changed, 172 insertions(+), 12 deletions(-) diff --git a/skyrl/backends/skyrl_train/training_batch.py b/skyrl/backends/skyrl_train/training_batch.py index db6c85d88f..ef6f95951b 100644 --- a/skyrl/backends/skyrl_train/training_batch.py +++ b/skyrl/backends/skyrl_train/training_batch.py @@ -20,22 +20,25 @@ 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) -> dict: - """Serialize a single tensor for pickle protocol.""" +def _serialize_tensor(value: torch.Tensor, *, zero_copy: bool = False) -> dict: + """Serialize a single tensor for pickle protocol. + + With ``zero_copy`` the payload carries the numpy array itself rather than a fresh + ``bytes`` copy of it. Pickle protocol 5 hands a payload to Ray's plasma out-of-band + path only if it reduces to a ``PickleBuffer``, which a C-contiguous array does and + ``bytes`` does not; the reader then rebuilds a view onto shared memory instead of + copying. That view is read-only, so only fields in ``TensorBatch.ZERO_COPY_KEYS`` + may take this path. + """ try: # Fast path: direct memory copy via numpy (works for most dtypes) arr = value.numpy() - return { - "format": TensorFormat.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() @@ -45,13 +48,32 @@ def _serialize_tensor(value: torch.Tensor) -> dict: "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") == TensorFormat.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: + # Zero-copy path: `data` is already an array, and under Ray it views the plasma + # buffer. `torch.from_numpy` warns once per process when that buffer is read-only; + # the returned tensor must not be mutated in place. + return torch.from_numpy(value["data"]) else: # Fast path: reconstruct from numpy bytes # Also handles legacy format without "format" key @@ -145,6 +167,11 @@ class TensorBatch(dict, Generic[DictType]): metadata: Optional[Dict[str, Any]] = None + # Fields serialized as a zero-copy numpy view rather than a copied `bytes` blob (see + # `_serialize_tensor`). Deserialized tensors for these keys can be backed by read-only + # shared memory, so a field qualifies only if no consumer mutates it in place. + ZERO_COPY_KEYS: frozenset[str] = frozenset({"rollout_expert_indices"}) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._batch_size = None @@ -270,13 +297,15 @@ 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): @@ -287,11 +316,12 @@ def __getstate__(self): elif isinstance(value, PackedTensor): batch_dict[key] = { "format": TensorFormat.PACKED_TENSOR, - "values": _serialize_tensor(value.values), + "values": _serialize_tensor(value.values, zero_copy=zero_copy), + # `cu_seqlens` is [batch + 1] offsets: too small to be worth a plasma 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, diff --git a/tests/backends/skyrl_train/conftest.py b/tests/backends/skyrl_train/conftest.py index fd7c45ad64..4fb663effc 100644 --- a/tests/backends/skyrl_train/conftest.py +++ b/tests/backends/skyrl_train/conftest.py @@ -1,6 +1,12 @@ +import pickle +from typing import Any + import pytest import ray +# The protocol that carries buffers out of band, i.e. the one Ray pickles with. +OUT_OF_BAND_PICKLE_PROTOCOL = 5 + @pytest.fixture(scope="session", autouse=True) def ray_init(): @@ -10,3 +16,21 @@ def ray_init(): yield if ray.is_initialized(): ray.shutdown() + + +@pytest.fixture +def oob_round_trip(): + """Round trip an object through protocol-5 out-of-band buffers, the way Ray does. + + Returns the rebuilt object, the in-band payload, and the out-of-band buffer views. + With ``read_only`` the buffers are re-wrapped as immutable memoryviews, to stand in + for the plasma memory Ray maps read-only in the reader. + """ + + 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/test_train_batch.py b/tests/backends/skyrl_train/test_train_batch.py index 9aae1c5bb7..79825391cd 100644 --- a/tests/backends/skyrl_train/test_train_batch.py +++ b/tests/backends/skyrl_train/test_train_batch.py @@ -1,4 +1,5 @@ import pickle +from collections.abc import Callable import numpy as np import pytest @@ -6,17 +7,21 @@ import torch from skyrl.backends.skyrl_train.training_batch import ( + BatchField, TensorBatch, TensorFormat, TensorList, TrainingInput, TrainingInputBatch, + _deserialize_tensor, + _serialize_tensor, 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 def test_train_batch_initialization(): @@ -767,3 +772,104 @@ def test_serialized_field_formats_are_stable(): assert state["bf16_logprobs"]["format"] == TensorFormat.TORCH assert state["pixel_values"]["format"] == TensorFormat.TENSOR_LIST assert state["rollout_expert_indices"]["format"] == TensorFormat.PACKED_TENSOR +# ── zero-copy field transport ──────────────────────────────────────────────── + +ROUTE_KEY = "rollout_expert_indices" +_ZERO_COPY_SEGMENT_LENGTHS = [512, 256, 256] +_ZERO_COPY_BATCH_SIZE = len(_ZERO_COPY_SEGMENT_LENGTHS) + +# One payload per opted-in field, each holding `_ZERO_COPY_BATCH_SIZE` batch entries, so the +# mechanism tests cover whatever `ZERO_COPY_KEYS` holds. A key without an entry fails +# `test_zero_copy_keys_are_live_and_exclude_mutated_fields`. +_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), + ), +} + + +def _zero_copy_buffer(value: BatchField) -> torch.Tensor: + """The one buffer a zero-copy field ships out of band.""" + return value.values if isinstance(value, PackedTensor) else value + + +def test_zero_copy_keys_are_live_and_exclude_mutated_fields(): + """A key that silently stops naming a field makes the path inert, with nothing failing.""" + assert TensorBatch.ZERO_COPY_KEYS <= set(TrainingInput.__annotations__), "stale key would be inert" + assert set(_ZERO_COPY_PAYLOADS) == TensorBatch.ZERO_COPY_KEYS, "every zero-copy field needs a test payload" + # The trainer writes through these: `collators.py` scales `loss_mask` in place, and the + # advantage pipeline normalizes its own tensors. + for mutated in ("sequences", "attention_mask", "loss_mask", "response_mask", "advantages", "returns", "values"): + assert mutated not in 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): + """An opted-in field reaches Ray's plasma path as a buffer; every other field does not.""" + 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): + """Plasma is read-only in the reader, so an opted-in field must not copy -- while a field + the trainer writes through must stay off that path and 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(): + """A packed field's offsets are a `[batch + 1]` array; only its token buffer earns plasma.""" + 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 From bafe55360face69e9a07f8b0441faac26d3c41df Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Sat, 22 Aug 2026 00:08:09 +0000 Subject: [PATCH 12/17] style: tighten zero-copy transport comments and tests --- skyrl/backends/skyrl_train/training_batch.py | 18 ++++--------- tests/backends/skyrl_train/conftest.py | 8 +----- .../backends/skyrl_train/test_train_batch.py | 26 ++++++------------- 3 files changed, 14 insertions(+), 38 deletions(-) diff --git a/skyrl/backends/skyrl_train/training_batch.py b/skyrl/backends/skyrl_train/training_batch.py index ef6f95951b..9981fb62f7 100644 --- a/skyrl/backends/skyrl_train/training_batch.py +++ b/skyrl/backends/skyrl_train/training_batch.py @@ -29,12 +29,8 @@ class TensorFormat(StrEnum): def _serialize_tensor(value: torch.Tensor, *, zero_copy: bool = False) -> dict: """Serialize a single tensor for pickle protocol. - With ``zero_copy`` the payload carries the numpy array itself rather than a fresh - ``bytes`` copy of it. Pickle protocol 5 hands a payload to Ray's plasma out-of-band - path only if it reduces to a ``PickleBuffer``, which a C-contiguous array does and - ``bytes`` does not; the reader then rebuilds a view onto shared memory instead of - copying. That view is read-only, so only fields in ``TensorBatch.ZERO_COPY_KEYS`` - may take this path. + 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) @@ -70,9 +66,7 @@ def _deserialize_tensor(value: dict) -> torch.Tensor: buffer = io.BytesIO(value["data"]) return torch.load(buffer, weights_only=True) elif tensor_format == TensorFormat.NUMPY_VIEW: - # Zero-copy path: `data` is already an array, and under Ray it views the plasma - # buffer. `torch.from_numpy` warns once per process when that buffer is read-only; - # the returned tensor must not be mutated in place. + # Under Ray, this array views a read-only plasma buffer. return torch.from_numpy(value["data"]) else: # Fast path: reconstruct from numpy bytes @@ -167,9 +161,7 @@ class TensorBatch(dict, Generic[DictType]): metadata: Optional[Dict[str, Any]] = None - # Fields serialized as a zero-copy numpy view rather than a copied `bytes` blob (see - # `_serialize_tensor`). Deserialized tensors for these keys can be backed by read-only - # shared memory, so a field qualifies only if no consumer mutates it in place. + # These fields may be backed by read-only shared memory after deserialization. ZERO_COPY_KEYS: frozenset[str] = frozenset({"rollout_expert_indices"}) def __init__(self, *args, **kwargs): @@ -317,7 +309,7 @@ def __getstate__(self): batch_dict[key] = { "format": TensorFormat.PACKED_TENSOR, "values": _serialize_tensor(value.values, zero_copy=zero_copy), - # `cu_seqlens` is [batch + 1] offsets: too small to be worth a plasma buffer. + # Offsets are too small to benefit from an out-of-band buffer. "cu_seqlens": _serialize_tensor(value.cu_seqlens), } else: diff --git a/tests/backends/skyrl_train/conftest.py b/tests/backends/skyrl_train/conftest.py index 4fb663effc..51e1150d70 100644 --- a/tests/backends/skyrl_train/conftest.py +++ b/tests/backends/skyrl_train/conftest.py @@ -4,7 +4,6 @@ import pytest import ray -# The protocol that carries buffers out of band, i.e. the one Ray pickles with. OUT_OF_BAND_PICKLE_PROTOCOL = 5 @@ -20,12 +19,7 @@ def ray_init(): @pytest.fixture def oob_round_trip(): - """Round trip an object through protocol-5 out-of-band buffers, the way Ray does. - - Returns the rebuilt object, the in-band payload, and the out-of-band buffer views. - With ``read_only`` the buffers are re-wrapped as immutable memoryviews, to stand in - for the plasma memory Ray maps read-only in the reader. - """ + """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] = [] diff --git a/tests/backends/skyrl_train/test_train_batch.py b/tests/backends/skyrl_train/test_train_batch.py index 79825391cd..c4fbbb5c94 100644 --- a/tests/backends/skyrl_train/test_train_batch.py +++ b/tests/backends/skyrl_train/test_train_batch.py @@ -772,15 +772,11 @@ def test_serialized_field_formats_are_stable(): assert state["bf16_logprobs"]["format"] == TensorFormat.TORCH assert state["pixel_values"]["format"] == TensorFormat.TENSOR_LIST assert state["rollout_expert_indices"]["format"] == TensorFormat.PACKED_TENSOR -# ── zero-copy field transport ──────────────────────────────────────────────── - ROUTE_KEY = "rollout_expert_indices" _ZERO_COPY_SEGMENT_LENGTHS = [512, 256, 256] _ZERO_COPY_BATCH_SIZE = len(_ZERO_COPY_SEGMENT_LENGTHS) -# One payload per opted-in field, each holding `_ZERO_COPY_BATCH_SIZE` batch entries, so the -# mechanism tests cover whatever `ZERO_COPY_KEYS` holds. A key without an entry fails -# `test_zero_copy_keys_are_live_and_exclude_mutated_fields`. +# 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), @@ -790,23 +786,18 @@ def test_serialized_field_formats_are_stable(): def _zero_copy_buffer(value: BatchField) -> torch.Tensor: - """The one buffer a zero-copy field ships out of band.""" + """Return the payload buffer for a zero-copy field.""" return value.values if isinstance(value, PackedTensor) else value -def test_zero_copy_keys_are_live_and_exclude_mutated_fields(): - """A key that silently stops naming a field makes the path inert, with nothing failing.""" - assert TensorBatch.ZERO_COPY_KEYS <= set(TrainingInput.__annotations__), "stale key would be inert" - assert set(_ZERO_COPY_PAYLOADS) == TensorBatch.ZERO_COPY_KEYS, "every zero-copy field needs a test payload" - # The trainer writes through these: `collators.py` scales `loss_mask` in place, and the - # advantage pipeline normalizes its own tensors. - for mutated in ("sequences", "attention_mask", "loss_mask", "response_mask", "advantages", "returns", "values"): - assert mutated not in TensorBatch.ZERO_COPY_KEYS +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): - """An opted-in field reaches Ray's plasma path as a buffer; every other field does not.""" + """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}) @@ -824,8 +815,7 @@ def test_zero_copy_field_travels_out_of_band(key, oob_round_trip): @pytest.mark.parametrize("key", sorted(TensorBatch.ZERO_COPY_KEYS)) def test_zero_copy_field_tolerates_read_only_plasma_buffer(key, oob_round_trip): - """Plasma is read-only in the reader, so an opted-in field must not copy -- while a field - the trainer writes through must stay off that path and stay writable.""" + """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}) @@ -846,7 +836,7 @@ def test_zero_copy_field_tolerates_read_only_plasma_buffer(key, oob_round_trip): def test_packed_zero_copy_field_ships_only_its_values_buffer(): - """A packed field's offsets are a `[batch + 1]` array; only its token buffer earns plasma.""" + """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] From 0caad2f26088092bd98992e82bf2269df1914e92 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Sat, 22 Aug 2026 00:50:05 +0000 Subject: [PATCH 13/17] style: apply Black formatting --- tests/backends/skyrl_train/test_train_batch.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/backends/skyrl_train/test_train_batch.py b/tests/backends/skyrl_train/test_train_batch.py index c4fbbb5c94..ff7fb8286b 100644 --- a/tests/backends/skyrl_train/test_train_batch.py +++ b/tests/backends/skyrl_train/test_train_batch.py @@ -772,6 +772,8 @@ def test_serialized_field_formats_are_stable(): 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) From 1c3510a4cbf96b957346ccd7527e61247aa39a24 Mon Sep 17 00:00:00 2001 From: lila-sync-bot Date: Fri, 14 Aug 2026 22:41:00 +0000 Subject: [PATCH 14/17] feat(sample-support): capture the sampler's bounded support and trace it across turns --- .../docs/algorithms/off_policy_correction.mdx | 1 + .../docs/tutorials/step-wise-training.mdx | 7 +- .../distributed/megatron/token_metadata.py | 12 + .../skyrl_train/inference_servers/base.py | 8 + .../inference_servers/generate_wire.py | 15 + .../remote_inference_client.py | 37 +- .../skyrl_train/inference_servers/setup.py | 1 + .../skyrl_train/inference_servers/utils.py | 5 + .../inference_servers/vllm_server_actor.py | 81 +++- .../skyrl_train/utils/sample_support.py | 69 +++ skyrl/train/config/config.py | 25 + skyrl/train/generators/base.py | 6 + skyrl/train/generators/skyrl_gym_generator.py | 251 +++++++++- skyrl/train/generators/skyrl_vlm_generator.py | 11 +- skyrl/train/generators/utils.py | 43 +- skyrl/train/utils/trainer_utils.py | 1 + skyrl/train/utils/utils.py | 10 + .../distributed/test_token_metadata.py | 26 ++ .../test_build_vllm_cli_args.py | 15 + .../inference_servers/test_generate_wire.py | 28 ++ .../test_remote_inference_client.py | 101 ++++ .../test_vllm_sample_support.py | 167 +++++++ .../skyrl_train/utils/test_sample_support.py | 52 +++ .../generators/test_generator_output_utils.py | 126 +++++ .../generators/test_skyrl_gym_generator.py | 430 +++++++++++++++++- .../generators/test_skyrl_vlm_generator.py | 26 ++ tests/train/test_config.py | 60 +++ 27 files changed, 1581 insertions(+), 33 deletions(-) create mode 100644 skyrl/backends/skyrl_train/utils/sample_support.py create mode 100644 tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py create mode 100644 tests/backends/skyrl_train/utils/test_sample_support.py 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/token_metadata.py b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py index 3257bd57e8..e028441b4b 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py @@ -325,6 +325,18 @@ def append(self, rows: np.ndarray, *, expected_rows: int) -> None: 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") diff --git a/skyrl/backends/skyrl_train/inference_servers/base.py b/skyrl/backends/skyrl_train/inference_servers/base.py index 4a8de4cdf4..b2343c64e7 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 @@ -35,6 +36,11 @@ class InferenceEngineInput(TypedDict): # only shared between requests carrying the same salt. See ``GeneratorConfig.use_cache_salt``. cache_salt: Optional[str] routed_experts_prompt_starts: Optional[List[int]] + # Opt a single batch into sample-support capture; requires the engine to have been started with + # ``InferenceEngineConfig.enable_return_sample_support_set``. Defaults to not capturing, so + # requests that do not consume the support (in-agent tool calls, eval requests whose greedy + # ``top_k=-1`` params cannot satisfy the capture contract) do not pay for it. + return_sample_support: Optional[bool] class InferenceEngineOutput(TypedDict): @@ -51,6 +57,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 a5f7aac23c..4facecc64d 100644 --- a/skyrl/backends/skyrl_train/inference_servers/generate_wire.py +++ b/skyrl/backends/skyrl_train/inference_servers/generate_wire.py @@ -25,6 +25,11 @@ 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 @@ -50,6 +55,7 @@ class PackedField(StrEnum): _ENVELOPE_KEYS = frozenset(PackedArrayKey) _ROUTED_EXPERTS_NDIM = 3 +_SAMPLE_SUPPORT_NDIM = 2 _QUOTE = b'"' @@ -185,6 +191,15 @@ def decode_packed_routed_experts(payload: dict[str, Any]) -> RoutedExpertIndices 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 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 f6137b8e49..d0d1ac937d 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -76,9 +76,11 @@ 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, @@ -179,6 +181,7 @@ class RemoteGenerateResult: response_logprobs: Optional[List[float]] stop_reason: str routed_experts: Optional[RoutedExpertIndices] + sample_support: Optional[SampleSupport] @dataclass @@ -253,10 +256,11 @@ async def generate( 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, optionally returning R3 routes.""" + """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") @@ -267,7 +271,8 @@ async def generate( ): raise ValueError("routed_experts_prompt_start must be an integer within the prompt") - path = "/skyrl/v1/generate" if return_routed_experts else "/inference/v1/generate" + 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 @@ -276,6 +281,8 @@ async def generate( "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 @@ -291,7 +298,7 @@ async def generate( f"{self.proxy_url}{path}", json=payload, headers=headers, - packed_side_channels=return_routed_experts, + packed_side_channels=packed_side_channels, ) choice = response["choices"][0] token_ids = choice["token_ids"] @@ -309,12 +316,20 @@ async def generate( 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: @@ -380,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, @@ -527,6 +546,9 @@ async def generate( 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: @@ -553,6 +575,7 @@ async def _throttled_generate(idx: int) -> Dict[str, Any]: 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, ) @@ -565,6 +588,7 @@ async def _throttled_generate(idx: int) -> Dict[str, Any]: 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, ) @@ -581,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, @@ -588,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( @@ -599,6 +627,7 @@ async def _generate_single( 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]: result = await self._get_generate_client().generate( prompt_token_ids=prompt_token_ids, @@ -607,6 +636,7 @@ async def _generate_single( 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, ) @@ -615,6 +645,7 @@ async def _generate_single( "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( 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 dab33322e5..9326aa9b78 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 @@ -40,8 +42,14 @@ 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, @@ -51,6 +59,52 @@ 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. + + vLLM emits ``[sampled token, top-1, ..., top-k]`` per generated token, so column 0 + carries the sampled-token logprob and columns ``1:`` are the support. Candidates the + top-p/min-p filters removed come back at ``-inf`` and become padding. + """ + 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 Triton top-k/top-p pivot can leave slightly more than top_k + # survivors, so the sampled token can rank just past top_k and be absent from its own + # support. Overwrite the row's weakest valid member (vLLM returns top-k descending, so + # that is the trailing valid slot) with the sampled id: width stays top_k, the sampled + # id appears exactly once so the trainer's renorm denominator is not double-counted, + # and trailing padding is untouched. Rows with no valid member at all are left alone. + 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. @@ -431,6 +485,22 @@ 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 is the sampler's bounded top-k set, so an unbounded or degenerate + # top_k has no support to return. Reject it here rather than forwarding + # `logprobs=top_k` to vLLM, which rejects a negative value with an opaque 500. + 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: @@ -451,7 +521,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( @@ -471,6 +549,7 @@ async def _skyrl_generate(request: Request): "finish_reason": finish_reason, "logprobs": logprobs, PackedField.ROUTED_EXPERTS.value: routed_experts, + PackedField.ROLLOUT_SAMPLE_SUPPORT.value: sample_support, } ] } 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..258bf72b92 --- /dev/null +++ b/skyrl/backends/skyrl_train/utils/sample_support.py @@ -0,0 +1,69 @@ +"""Per-token sampler support: the bounded top-k set vLLM actually sampled from. + +A support row is ``[top-1, ..., top-k]`` vocab IDs for one generated token, so the +trainer can renormalize its logprobs over the same bounded set the rollout sampler +drew from instead of the full vocabulary. Rows are dense and right-padded with +``SAMPLE_SUPPORT_PADDING``; a token with no captured support (prompt tokens, +observation tokens, a synthetic EOS) is an all-padding row. +""" + +from typing import TypeAlias + +import numpy as np + +from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( + TokenMetadataTrace, +) + +SampleSupport: TypeAlias = np.ndarray +SAMPLE_SUPPORT_DTYPE = np.dtype(np.int32) +SAMPLE_SUPPORT_DTYPES = frozenset({SAMPLE_SUPPORT_DTYPE}) +SAMPLE_SUPPORT_PADDING = -1 + + +def validate_sample_support(sample_support: SampleSupport) -> SampleSupport: + """Check the two invariants the generic packed-array codec cannot: no + negatives other than the padding sentinel, and padding only ever trailing.""" + 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 + + +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: + """Concatenate the trace and cut it to ``token_count`` rows. + + ``extra_rows`` is the trailing row count the response never keeps (the final observation), + so the total is checked exactly in both directions and an unexpected overshoot raises + instead of being silently truncated away. + """ + 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/train/config/config.py b/skyrl/train/config/config.py index 57c7521910..0164549470 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -1142,6 +1142,10 @@ 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 post-filter sampler support for each generated token, so the trainer can + renormalize logprobs over the same support the rollout sampler drew from. Constrains + ``generator.sampling_params``; eval requests opt out per-request instead.""" 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 +1741,27 @@ def __post_init__(self): if self.trainer.algorithm.temperature is None: self.trainer.algorithm.temperature = self.generator.sampling_params.temperature + # Capture guards apply to generator.sampling_params only. generator.eval_sampling_params is + # greedy with top_k=-1 by default and cannot satisfy them, so the generator opts eval + # requests out of capture instead. + 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 re-renders the conversation each turn and never populates + # ``rollout_expert_indices``, so the pair would generate a full batch and then fail collation. + 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/generators/base.py b/skyrl/train/generators/base.py index 81792e0a2a..24180fd463 100644 --- a/skyrl/train/generators/base.py +++ b/skyrl/train/generators/base.py @@ -8,6 +8,8 @@ from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices TrainingPhase = Literal["train", "eval"] +TRAINING_PHASE_TRAIN: TrainingPhase = "train" +TRAINING_PHASE_EVAL: TrainingPhase = "eval" @dataclass @@ -52,6 +54,10 @@ class GeneratorOutput(TypedDict): # record its split. trajectory_time_splits: Optional[Dict[str, List[float]]] rollout_expert_indices: Optional[List[RoutedExpertIndices]] + # Per trajectory, one dense ``[response_tokens, top_k]`` row block of the sampler support each + # response token was drawn from, right-padded with ``SAMPLE_SUPPORT_PADDING``. Tokens with no + # captured support (observations, a synthetic EOS) are all-padding rows. + rollout_sample_support: Optional[List[List[List[int]]]] # 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 75c5d07671..2548c3040c 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 @@ -27,11 +28,20 @@ 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 ( @@ -55,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[List[List[int]]] = 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 @@ -87,6 +98,10 @@ class AgentLoopState: response_end_idx: Optional[int] done: bool routed_expert_trace: Optional[RoutedExpertTrace] = None + sample_support_trace: Optional[SampleSupportTrace] = None + # The support row of the latest turn's own EOS token, which the ``use_conversation_multi_turn=False`` + # convention slices off the response. ``None`` when the turn did not end in a generated EOS. + dropped_eos_sample_support: Optional[SampleSupport] = None @dataclass @@ -97,8 +112,27 @@ class TurnOutput: new_obs: ConversationType obs_ids: List[int] reward: Optional[float] + rollout_sample_support: Optional[SampleSupport] = None added_eos: bool = False + def get_turn_rollout_sample_support(self) -> Optional[SampleSupport]: + """Sample support for this turn's generated tokens, padded over the suffix. + + Mirrors ``get_turn_loss_mask``: a synthetic EOS and the observation tokens were + never sampled, so they get all-padding rows. + """ + 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]: """ Get loss mask for this turn's tokens. @@ -215,12 +249,46 @@ 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." + ) + + # Keyed on `custom_chat_template` rather than the narrower `retokenize_chat_history`, so the inert + # `use_conversation_multi_turn=False` + custom-template pair is refused too. + 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() @@ -285,6 +353,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. @@ -361,9 +430,24 @@ async def agent_loop( current_sampling_params: dict = ( sampling_params if sampling_params is not None else asdict(self.generator_cfg.sampling_params) ) + # The eval phase's greedy/`top_k=-1` params cannot satisfy the capture contract that + # ``SkyRLTrainConfig`` enforces on ``generator.sampling_params``, so capture is requested + # per-request keyed on the phase, not by the engine-level flag alone. + 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 + # Routes are gated on the same phase: no consumer reads an eval trajectory's routes, and the + # driver holds every trajectory of an eval batch at once. + 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 @@ -377,9 +461,8 @@ async def agent_loop( rollout_logprobs=[] if get_logprobs else None, response_end_idx=None, done=False, - routed_expert_trace=( - RoutedExpertTrace() if self.generator_cfg.inference_engine.enable_return_routed_experts else None - ), + 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: @@ -409,6 +492,7 @@ async def agent_loop( 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) @@ -425,10 +509,10 @@ 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: - raise ValueError( - "Rollout expert indices bookkeeping is not supported with custom chat template" - ) + # `_validate_cfg` refuses this pair; subclasses that replace it must too. + 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") @@ -437,6 +521,25 @@ async def agent_loop( generated_token_count=len(output_ids), routed_experts=rollout_expert_indices, ) + + sample_support_rows = None + if capture_sample_support: + # `_validate_cfg` refuses this pair; subclasses that replace it must too. + 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( + 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) stop_strs = current_sampling_params.get("stop", None) @@ -472,6 +575,8 @@ async def agent_loop( 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) @@ -483,6 +588,7 @@ 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, ) @@ -494,6 +600,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, @@ -503,11 +610,17 @@ 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_sample_support=( + turn_sample_support.tolist() if turn_sample_support is not None else None + ), ) 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( @@ -524,6 +637,18 @@ async def agent_loop( agent_loop_state, turn_output ) + if sample_support_trace is not None: + # Support rows are 1:1 with the tokens the turn appended to the trajectory, so a state + # update that does not feed the trace is caught here instead of silently emitting None. + 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 @@ -534,6 +659,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 . @@ -578,6 +704,15 @@ 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: + # The last turn's own EOS was sampled from a real support set, which the + # single-turn convention sliced off; give it back to the token that carries + # the terminal reward. A stop-string EOS was never sampled, so it stays 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: @@ -585,6 +720,11 @@ async def agent_loop( 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, + ).tolist() 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): @@ -604,6 +744,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( @@ -731,6 +872,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) @@ -762,15 +904,30 @@ async def generate_batched( tokenize=True, return_dict=False, ) + # Both side channels are captured only for train batches: eval params cannot satisfy the + # sample-support contract, and no consumer reads an eval batch's routes. + 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 = [] @@ -778,6 +935,9 @@ 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[List[List[int]]]] = ( + [] 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 @@ -796,6 +956,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)].tolist()) # Get environment-specific metrics env_metrics.append(env.get_metrics()) @@ -817,6 +979,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 @@ -846,9 +1009,22 @@ 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() + # Drives the per-request sample-support capture opt-out. Both phases carry non-None + # `sampling_params`, so the phase is the only discriminator available here. + 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. @@ -864,6 +1040,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, ) ) @@ -951,10 +1128,26 @@ 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: + # Step-wise rows carry no routes: `_validate_cfg` refuses R3 with `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] + # Keyed on value presence rather than the engine flag, so an eval batch (which captures no + # routes) and generator subclasses that leave the field unpopulated emit None. + 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, @@ -989,6 +1182,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, } @@ -1050,6 +1244,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 ) @@ -1125,6 +1326,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 + 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 @@ -1176,8 +1380,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) @@ -1195,4 +1406,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 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..8edc390827 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,9 @@ def _validate_cfg(self, generator_cfg: GeneratorConfig): "SkyRLVLMGymGenerator requires `use_conversation_multi_turn=True` " "because multi-modal observations must be in separate user messages." ) + # The base refusals still apply to this generator; the VLM-specific ones above take precedence + # where they overlap. + super()._validate_cfg(generator_cfg) async def _render_conversation(self, conversation: ConversationType) -> RenderedConversation: rendered = await self.inference_engine_client.render_chat_completion( @@ -74,6 +82,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..7cdbf13b13 100644 --- a/skyrl/train/generators/utils.py +++ b/skyrl/train/generators/utils.py @@ -8,6 +8,7 @@ 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_PADDING from skyrl.train.config import ChatTemplateConfig from skyrl.train.generators.base import ( BatchMetadata, @@ -278,11 +279,15 @@ 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 are all-or-nothing across the batches: a partially populated field would + # be dropped (or raise a bare TypeError) depending on which batch happened to come first. + 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 +296,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"), } @@ -788,9 +795,10 @@ def slice_generator_output( 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. + Handles a flat list or a dict-of-lists, slicing each component. """ 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,6 +806,8 @@ 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 @@ -822,6 +832,11 @@ 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 = ( + next((len(row) for rows in gen_out["rollout_sample_support"] for row in rows), 0) 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,7 @@ 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 + out_sample_support: Optional[List[List[List[int]]]] = [] 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 +858,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[List[int]]] = ( + [list(row) for row in 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(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 +890,9 @@ 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 = ( + [list(row) for row in 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 +907,8 @@ 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.extend([SAMPLE_SUPPORT_PADDING] * sample_support_width for _ in obs_delta) if acc_rewards_tokens is not None: acc_rewards_tokens.extend([0.0] * len(obs_delta)) @@ -891,6 +917,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.extend(gen_out["rollout_sample_support"][i]) if acc_rewards_tokens is not None: acc_rewards_tokens.extend(gen_out["rewards"][i]) @@ -905,6 +933,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/utils/trainer_utils.py b/skyrl/train/utils/trainer_utils.py index 0d91977ecb..05c258bd14 100644 --- a/skyrl/train/utils/trainer_utils.py +++ b/skyrl/train/utils/trainer_utils.py @@ -702,6 +702,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", diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index a1961e5a0c..a19f5e3274 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -375,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/tests/backends/skyrl_train/distributed/test_token_metadata.py b/tests/backends/skyrl_train/distributed/test_token_metadata.py index e3e7660567..3d518fc782 100644 --- a/tests/backends/skyrl_train/distributed/test_token_metadata.py +++ b/tests/backends/skyrl_train/distributed/test_token_metadata.py @@ -222,3 +222,29 @@ def test_align_packed_token_metadata_rejects_segments_that_leave_the_trajectory( 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/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 c93a4d4cfd..35e2c551e1 100644 --- a/tests/backends/skyrl_train/inference_servers/test_generate_wire.py +++ b/tests/backends/skyrl_train/inference_servers/test_generate_wire.py @@ -16,9 +16,11 @@ 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, ) @@ -202,6 +204,32 @@ def test_decode_rejects_noncanonical_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", [ 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 85b4f09151..f232aca062 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 @@ -27,11 +27,13 @@ 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 ( @@ -41,6 +43,12 @@ _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) @@ -555,6 +563,99 @@ async def test_generate_decodes_packed_routed_experts(self, mock_servers): 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): + """A request that does not ask for support must not pay for it: callers that never consume it + (in-agent tool calls, eval batches) omit the key, and eval's `top_k=-1` would make the server + request negative logprobs.""" + 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): client = RemoteInferenceClient( 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..fa7e47aa7b --- /dev/null +++ b/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py @@ -0,0 +1,167 @@ +"""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(): + # Three rows, top_k=3 (row_width=4): + # Row A: sampled id (100) absent from a fully-valid support row -> repair. + # Row B: sampled id (7) already present -> unchanged. + # Row C: sampled id (5) absent from a support row that has trailing -1 padding. + 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, # row A: all valid + -0.1, + -0.1, + -0.2, + -0.3, # row B: all valid + -0.4, + -0.5, + -0.6, + float("-inf"), # row C: last col filtered -> padding + ], + ) + + _, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=top_k) + sampled_ids = [100, 7, 5] + + # (b) each row keeps width == top_k + assert all(row.size == top_k for row in support) + + # (a) every row's support now contains its sampled id + for sampled_id, row in zip(sampled_ids, support): + assert sampled_id in row + + # (c) the sampled id appears exactly once per repaired row (no duplicate) + assert np.count_nonzero(support[0] == 100) == 1 + assert np.count_nonzero(support[2] == 5) == 1 + + # (d) trailing -1 padding preserved on the padded row + assert support[2][-1] == -1 + + # (e) the unaffected row (sampled already present) is unchanged + np.testing.assert_array_equal(support[1], [7, 8, 9]) + + # Concrete expected repair: weakest (trailing) valid member overwritten. + np.testing.assert_array_equal(support[0], [8, 9, 100]) + np.testing.assert_array_equal(support[2], [6, 5, -1]) + + +def test_flat_logprobs_top_k_one_repairs_single_support_column(): + # top_k == 1 (row_width == 2): a single support column that must hold the sampled id. + 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): + """Support is the sampler's bounded top-k set, so an unbounded or degenerate ``top_k`` has none. + Forwarding it as ``logprobs`` instead yields an opaque 500 from vLLM.""" + 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/utils/test_sample_support.py b/tests/backends/skyrl_train/utils/test_sample_support.py new file mode 100644 index 0000000000..3bc5c87e00 --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_sample_support.py @@ -0,0 +1,52 @@ +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(): + """A trailing row count the caller did not declare means the trace and the response disagree + about which tokens were sampled, which silent truncation would hide.""" + 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/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py index e3ee7a1f5c..043a3b5392 100644 --- a/tests/train/generators/test_generator_output_utils.py +++ b/tests/train/generators/test_generator_output_utils.py @@ -14,6 +14,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 +31,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 +54,8 @@ 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]], + "rollout_expert_indices": [np.zeros((2, 1, 2), dtype=np.uint8), np.ones((2, 1, 2), dtype=np.uint8)], + "rollout_sample_support": [[[1, 2], [1, 2]], [[3, 4], [3, 4]]], } generator_output_2: GeneratorOutput = { @@ -61,6 +65,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((3, 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 +78,13 @@ 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]] + # Both per-token side channels are named fields, so the input order cannot decide whether they + # survive concatenation. + 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,53 @@ 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): + """A batch that captured a side channel cannot be concatenated with one that did not: the + consumer packs the field for every trajectory in the batch or for none of them.""" + + 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(): + """``trajectory_time_splits`` is a dict of per-entry lists rather than a per-entry list.""" + 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 +588,61 @@ 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): + """Every other producer emits dense width-``top_k`` rows, so an observation delta + must too -- an empty row here would make the trajectory ragged.""" + 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": [[[20, 21, -1]], [[40, 44, -1], [41, 45, 46]]], + "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] + assert support == [[20, 21, -1], [-1, -1, -1], [40, 44, -1], [41, 45, 46]] + assert len(support) == len(merged["response_ids"][0]) + assert {len(row) for row in support} == {3} + + def test_native_output_carrying_dict_valued_time_splits(self): + """The native step-wise ``GeneratorOutput`` carries ``trajectory_time_splits`` as a dict of + per-entry lists; the other fixtures in this file omit the key entirely.""" + 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) + + # `_merge_single_trajectory` returns a fixed key set that drops the timing fields, so only the + # merge itself is asserted here. + 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 +860,17 @@ 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): + """Refused in `validate_cfg`, which runs before Ray and the engines start, and the message + must state why: a step's routes would replay onto the first N prompt tokens of its row.""" + 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 b253f02433..daf73b3115 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -8,14 +8,21 @@ import numpy as np import pytest +from skyrl.backends.skyrl_train.utils.sample_support import 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 @@ -31,9 +38,14 @@ def test_turn_output_masks_uncaptured_suffix(): new_obs=[], obs_ids=[20, 21], reward=1.0, + rollout_sample_support=np.array([[10, 100], [11, 101]], dtype=np.int32), added_eos=True, ) + 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] @@ -412,7 +424,7 @@ def mock_generate(_, model=None): @pytest.mark.asyncio @patch("skyrl_gym.make") -async def test_agent_loop_uses_incremental_routed_expert_trace( +async def test_agent_loop_uses_incremental_replay_metadata_traces( mock_make, mock_tokenizer, mock_llm, @@ -424,6 +436,8 @@ async def test_agent_loop_uses_incremental_routed_expert_trace( 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"}], {}) @@ -432,19 +446,25 @@ async def test_agent_loop_uses_incremental_routed_expert_trace( 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) @@ -456,7 +476,7 @@ def generate(input_batch, model=None): ) generator.base_conversation_token_ids = [] - await generator.agent_loop( + output = await generator.agent_loop( [{"role": "user", "content": "Start"}], mock_env_cfg.env_class, {}, @@ -465,6 +485,410 @@ def generate(input_batch, model=None): ) assert prompt_starts == [0, 5] + assert output.rollout_sample_support[:2] == [[10, 100], [11, 110]] + assert output.rollout_sample_support[-2:] == [[10, 101], [11, 111]] + assert all(row == [-1, -1] for row in output.rollout_sample_support[2:-2]) + + +@pytest.mark.asyncio +@patch("skyrl_gym.make") +async def test_agent_loop_skips_sample_support_capture_on_the_eval_path( + mock_make, + mock_tokenizer, + mock_llm, + mock_env, + generator_cfg, + mock_env_cfg, +): + """The eval phase's sampling params (greedy, top_k=-1) cannot satisfy the capture + contract; the engine-level flag must not force capture on them.""" + generator_cfg.batched = False + generator_cfg.max_turns = 1 + 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={}), + ] + captured = {} + + def generate(input_batch, model=None): + captured.update(input_batch) + return { + "responses": ["mocked output"], + "response_ids": [[10, 11]], + "stop_reasons": ["stop"], + } + + 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, + sampling_params={"temperature": 0.0, "top_k": -1, "max_tokens": 32}, + training_phase=TRAINING_PHASE_EVAL, + ) + + assert captured["return_sample_support"] is False + assert output.rollout_sample_support is None + + +@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, +): + """Capture is requested iff the engine flag is on and the batch is a train-phase batch. + + Both phases supply non-None ``sampling_params`` (train from ``generator.sampling_params``, + eval from ``generator.eval_sampling_params``), so ``batch_metadata.training_phase`` is the + only usable discriminator -- the presence of ``sampling_params`` is not one. Driven through + ``generate`` so both the batched and agent-loop routes are covered. + """ + 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 = [] + + # Mirrors `prepare_generator_input`: both phases carry explicit sampling params. + 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 +): + """The single-turn convention re-appends a loss-active EOS the inference engine never evaluated, + so no routed-expert row exists for it and the trace refuses to dummy-pad a loss-active target.""" + 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 +): + """A step-wise `generate` has no field to put routes in, so without this refusal a generator built + outside the entrypoint (which is where `validate_cfg` runs) emits `rollout_expert_indices=None` + and trains with routing replay silently disabled.""" + 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 +): + """Refused at config time rather than on the first train batch, after a large model has loaded.""" + 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 +): + """The retokenizing state update is the one helper that never feeds a trace, so a trace reaching + it would finalize empty (support) or misaligned (routes) instead of failing.""" + 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, +): + """Nothing reads an eval trajectory's routes, and the driver holds the whole eval batch at once + (~3.9 KB per token at 120B), so eval routes must not be retained.""" + 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: + # A captured trace also asks the engine where its prompt prefix already ended. + 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, +): + """With `use_conversation_multi_turn=False` the generated EOS is sliced off the turn and + re-appended to the trajectory as the loss-active token carrying the terminal reward. It keeps + the support set it was actually sampled from rather than an all-padding row, which a replay + consumer would read as a synthetic EOS and score over the full vocabulary.""" + 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"], + # The engine's own EOS (id 4) terminates the turn. + "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] + # The re-appended EOS is loss-active and carries the terminal reward. + assert output.loss_mask == [1, 1, 1] + assert output.rollout_sample_support == [[10, 110], [11, 111], eos_support_row] + + +@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, +): + """An EOS the loop appends because generation stopped on a stop string was never sampled, so it + gets an all-padding row.""" + 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] + assert output.rollout_sample_support == [[10, 110], [11, 111], [-1, -1]] @pytest.mark.asyncio diff --git a/tests/train/generators/test_skyrl_vlm_generator.py b/tests/train/generators/test_skyrl_vlm_generator.py index 0cd504e302..9758a524ab 100644 --- a/tests/train/generators/test_skyrl_vlm_generator.py +++ b/tests/train/generators/test_skyrl_vlm_generator.py @@ -134,6 +134,32 @@ async def mock_generate(input_batch, model=None): # --------------------------------------------------------------------------- +def test_vlm_validate_cfg_still_applies_the_base_refusals(): + """``SkyRLVLMGymGenerator._validate_cfg`` replaces the base method, so it must chain to it; + otherwise every base refusal is silently void for this generator.""" + 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_config.py b/tests/train/test_config.py index 7e4dc1a493..1e136be3e2 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -158,6 +158,66 @@ 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(): + """The VLM generator never populates ``rollout_expert_indices``, so without this guard the pair + generates a whole batch and then dies in collation on a None route array.""" + 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(): + # The five guards deliberately do not reach eval_sampling_params: forcing + # temperature > 0 there would change eval semantics. Capture is opted out per request. + 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"]) From 9a8327d7806e8a1ab3c62da545919558b4432ac9 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Sat, 22 Aug 2026 00:08:52 +0000 Subject: [PATCH 15/17] style: tighten sample-support capture comments and tests --- .../skyrl_train/inference_servers/base.py | 5 +- .../inference_servers/vllm_server_actor.py | 16 +--- .../skyrl_train/utils/sample_support.py | 18 +---- skyrl/train/config/config.py | 11 +-- skyrl/train/generators/skyrl_gym_generator.py | 37 +++------ skyrl/train/generators/skyrl_vlm_generator.py | 2 - skyrl/train/generators/utils.py | 11 +-- .../test_remote_inference_client.py | 3 - .../test_vllm_sample_support.py | 36 +-------- .../skyrl_train/utils/test_sample_support.py | 2 - .../generators/test_generator_output_utils.py | 14 ---- .../generators/test_skyrl_gym_generator.py | 81 ------------------- .../generators/test_skyrl_vlm_generator.py | 2 - tests/train/test_config.py | 4 - 14 files changed, 27 insertions(+), 215 deletions(-) diff --git a/skyrl/backends/skyrl_train/inference_servers/base.py b/skyrl/backends/skyrl_train/inference_servers/base.py index b2343c64e7..2c1d9a5ef5 100644 --- a/skyrl/backends/skyrl_train/inference_servers/base.py +++ b/skyrl/backends/skyrl_train/inference_servers/base.py @@ -36,10 +36,7 @@ class InferenceEngineInput(TypedDict): # only shared between requests carrying the same salt. See ``GeneratorConfig.use_cache_salt``. cache_salt: Optional[str] routed_experts_prompt_starts: Optional[List[int]] - # Opt a single batch into sample-support capture; requires the engine to have been started with - # ``InferenceEngineConfig.enable_return_sample_support_set``. Defaults to not capturing, so - # requests that do not consume the support (in-agent tool calls, eval requests whose greedy - # ``top_k=-1`` params cannot satisfy the capture contract) do not pay for it. + # Per-batch opt-in; the engine must enable sample-support capture at startup. return_sample_support: Optional[bool] 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 9326aa9b78..84a1ad027a 100644 --- a/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py +++ b/skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py @@ -65,9 +65,7 @@ def _sample_support_from_flat_logprobs( ) -> tuple[list[dict[str, float]], SampleSupport]: """Extract sampled scores and post-filter support from vLLM's flat rows. - vLLM emits ``[sampled token, top-1, ..., top-k]`` per generated token, so column 0 - carries the sampled-token logprob and columns ``1:`` are the support. Candidates the - top-p/min-p filters removed come back at ``-inf`` and become padding. + 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) @@ -79,12 +77,8 @@ def _sample_support_from_flat_logprobs( ) sampled_logprobs = [{"logprob": value} for value in processed_logprobs[:, 0].tolist()] - # vLLM's approximate Triton top-k/top-p pivot can leave slightly more than top_k - # survivors, so the sampled token can rank just past top_k and be absent from its own - # support. Overwrite the row's weakest valid member (vLLM returns top-k descending, so - # that is the trailing valid slot) with the sampled id: width stays top_k, the sampled - # id appears exactly once so the trainer's renorm denominator is not double-counted, - # and trailing padding is untouched. Rows with no valid member at all are left alone. + # 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) @@ -487,9 +481,7 @@ async def _skyrl_generate(request: Request): capture_sample_support = body.get("return_sample_support", False) if capture_sample_support: - # Sample support is the sampler's bounded top-k set, so an unbounded or degenerate - # top_k has no support to return. Reject it here rather than forwarding - # `logprobs=top_k` to vLLM, which rejects a negative value with an opaque 500. + # 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( diff --git a/skyrl/backends/skyrl_train/utils/sample_support.py b/skyrl/backends/skyrl_train/utils/sample_support.py index 258bf72b92..5fb78367fd 100644 --- a/skyrl/backends/skyrl_train/utils/sample_support.py +++ b/skyrl/backends/skyrl_train/utils/sample_support.py @@ -1,10 +1,6 @@ -"""Per-token sampler support: the bounded top-k set vLLM actually sampled from. +"""Per-token bounded sampler support used to renormalize rollout logprobs. -A support row is ``[top-1, ..., top-k]`` vocab IDs for one generated token, so the -trainer can renormalize its logprobs over the same bounded set the rollout sampler -drew from instead of the full vocabulary. Rows are dense and right-padded with -``SAMPLE_SUPPORT_PADDING``; a token with no captured support (prompt tokens, -observation tokens, a synthetic EOS) is an all-padding row. +Rows contain top-k vocabulary IDs and use trailing ``SAMPLE_SUPPORT_PADDING``. """ from typing import TypeAlias @@ -22,8 +18,7 @@ def validate_sample_support(sample_support: SampleSupport) -> SampleSupport: - """Check the two invariants the generic packed-array codec cannot: no - negatives other than the padding sentinel, and padding only ever trailing.""" + """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): @@ -55,12 +50,7 @@ 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: - """Concatenate the trace and cut it to ``token_count`` rows. - - ``extra_rows`` is the trailing row count the response never keeps (the final observation), - so the total is checked exactly in both directions and an unexpected overshoot raises - instead of being silently truncated away. - """ + """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 " diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 0164549470..3aadb26c34 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -1143,9 +1143,7 @@ class InferenceEngineConfig(BaseConfig): """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 post-filter sampler support for each generated token, so the trainer can - renormalize logprobs over the same support the rollout sampler drew from. Constrains - ``generator.sampling_params``; eval requests opt out per-request instead.""" + """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 @@ -1741,9 +1739,7 @@ def __post_init__(self): if self.trainer.algorithm.temperature is None: self.trainer.algorithm.temperature = self.generator.sampling_params.temperature - # Capture guards apply to generator.sampling_params only. generator.eval_sampling_params is - # greedy with top_k=-1 by default and cannot satisfy them, so the generator opts eval - # requests out of capture instead. + # 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: @@ -1757,8 +1753,7 @@ def __post_init__(self): if self.generator.vision_language_generator: raise ValueError("sample-support capture does not support vision_language_generator") - # The VLM generator re-renders the conversation each turn and never populates - # ``rollout_expert_indices``, so the pair would generate a full batch and then fail collation. + # 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") diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py index 2548c3040c..9f44f53491 100644 --- a/skyrl/train/generators/skyrl_gym_generator.py +++ b/skyrl/train/generators/skyrl_gym_generator.py @@ -99,8 +99,7 @@ class AgentLoopState: done: bool routed_expert_trace: Optional[RoutedExpertTrace] = None sample_support_trace: Optional[SampleSupportTrace] = None - # The support row of the latest turn's own EOS token, which the ``use_conversation_multi_turn=False`` - # convention slices off the response. ``None`` when the turn did not end in a generated EOS. + # Support for an EOS sliced from a single-turn response. dropped_eos_sample_support: Optional[SampleSupport] = None @@ -116,11 +115,7 @@ class TurnOutput: added_eos: bool = False def get_turn_rollout_sample_support(self) -> Optional[SampleSupport]: - """Sample support for this turn's generated tokens, padded over the suffix. - - Mirrors ``get_turn_loss_mask``: a synthetic EOS and the observation tokens were - never sampled, so they get all-padding rows. - """ + """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) @@ -272,8 +267,6 @@ def _validate_cfg(self, generator_cfg: GeneratorConfig): "loss-active target." ) - # Keyed on `custom_chat_template` rather than the narrower `retokenize_chat_history`, so the inert - # `use_conversation_multi_turn=False` + custom-template pair is refused too. if self.custom_chat_template is not None: if ie_cfg.enable_return_routed_experts: raise ValueError( @@ -430,16 +423,13 @@ async def agent_loop( current_sampling_params: dict = ( sampling_params if sampling_params is not None else asdict(self.generator_cfg.sampling_params) ) - # The eval phase's greedy/`top_k=-1` params cannot satisfy the capture contract that - # ``SkyRLTrainConfig`` enforces on ``generator.sampling_params``, so capture is requested - # per-request keyed on the phase, not by the engine-level flag alone. + # 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 - # Routes are gated on the same phase: no consumer reads an eval trajectory's routes, and the - # driver holds every trajectory of an eval batch at once. + # 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 @@ -509,7 +499,6 @@ async def agent_loop( if rollout_expert_indices is not None: rollout_expert_indices = rollout_expert_indices[0] - # `_validate_cfg` refuses this pair; subclasses that replace it must too. assert ( self.custom_chat_template is None ), "Rollout expert indices bookkeeping is not supported with custom chat template" @@ -524,7 +513,6 @@ async def agent_loop( sample_support_rows = None if capture_sample_support: - # `_validate_cfg` refuses this pair; subclasses that replace it must too. assert ( self.custom_chat_template is None ), "Sample-support bookkeeping is not supported with custom chat template" @@ -638,8 +626,7 @@ async def agent_loop( ) if sample_support_trace is not None: - # Support rows are 1:1 with the tokens the turn appended to the trajectory, so a state - # update that does not feed the trace is caught here instead of silently emitting 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, ( @@ -705,9 +692,7 @@ async def agent_loop( if rollout_logprobs is not None: rollout_logprobs.append(0.0) if agent_loop_state.sample_support_trace is not None: - # The last turn's own EOS was sampled from a real support set, which the - # single-turn convention sliced off; give it back to the token that carries - # the terminal reward. A stop-string EOS was never sampled, so it stays padding. + # 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) @@ -904,8 +889,7 @@ async def generate_batched( tokenize=True, return_dict=False, ) - # Both side channels are captured only for train batches: eval params cannot satisfy the - # sample-support contract, and no consumer reads an eval batch's routes. + # 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 @@ -1009,8 +993,7 @@ 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() - # Drives the per-request sample-support capture opt-out. Both phases carry non-None - # `sampling_params`, so the phase is the only discriminator available here. + # 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 @@ -1129,12 +1112,10 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False rollout_logprobs = None if self.generator_cfg.step_wise_trajectories: - # Step-wise rows carry no routes: `_validate_cfg` refuses R3 with `step_wise_trajectories`. expert_indices_values = [None] * len(responses) else: expert_indices_values = [output.rollout_expert_indices for output in all_outputs] - # Keyed on value presence rather than the engine flag, so an eval batch (which captures no - # routes) and generator subclasses that leave the field unpopulated emit None. + # 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 ) diff --git a/skyrl/train/generators/skyrl_vlm_generator.py b/skyrl/train/generators/skyrl_vlm_generator.py index 8edc390827..7053bd1a5a 100644 --- a/skyrl/train/generators/skyrl_vlm_generator.py +++ b/skyrl/train/generators/skyrl_vlm_generator.py @@ -62,8 +62,6 @@ def _validate_cfg(self, generator_cfg: GeneratorConfig): "SkyRLVLMGymGenerator requires `use_conversation_multi_turn=True` " "because multi-modal observations must be in separate user messages." ) - # The base refusals still apply to this generator; the VLM-specific ones above take precedence - # where they overlap. super()._validate_cfg(generator_cfg) async def _render_conversation(self, conversation: ConversationType) -> RenderedConversation: diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py index 7cdbf13b13..18ab425ee5 100644 --- a/skyrl/train/generators/utils.py +++ b/skyrl/train/generators/utils.py @@ -279,8 +279,7 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step (e.g. `is_last_step`, `trajectory_ids`, contiguous trajectory ordering). """ assert len(generator_outputs) > 0 - # Per-token side channels are all-or-nothing across the batches: a partially populated field would - # be dropped (or raise a bare TypeError) depending on which batch happened to come first. + # 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): @@ -790,13 +789,7 @@ 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. - Handles a flat list or a dict-of-lists, slicing each component. - """ + """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 None, a dict of per-entry lists, or a per-entry list. sliced: GeneratorOutput = {} 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 f232aca062..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 @@ -600,9 +600,6 @@ async def fake_post(url, json, headers, *, packed_side_channels=False): @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): - """A request that does not ask for support must not pay for it: callers that never consume it - (in-agent tool calls, eval batches) omit the key, and eval's `top_k=-1` would make the server - request negative logprobs.""" client = RemoteInferenceClient( proxy_url="http://unused", server_urls=["http://unused"], 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 index fa7e47aa7b..f7f44b5102 100644 --- a/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py +++ b/tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py @@ -46,10 +46,6 @@ def test_flat_logprobs_replaces_top_p_masked_candidates(): def test_flat_logprobs_repairs_sampled_token_absent_from_support(): - # Three rows, top_k=3 (row_width=4): - # Row A: sampled id (100) absent from a fully-valid support row -> repair. - # Row B: sampled id (7) already present -> unchanged. - # Row C: sampled id (5) absent from a support row that has trailing -1 padding. top_k = 3 flat_logprobs = SimpleNamespace( token_ids=[100, 8, 9, 10, 7, 7, 8, 9, 5, 6, 7, 8], @@ -57,45 +53,23 @@ def test_flat_logprobs_repairs_sampled_token_absent_from_support(): -0.1, -0.2, -0.3, - -0.4, # row A: all valid + -0.4, -0.1, -0.1, -0.2, - -0.3, # row B: all valid + -0.3, -0.4, -0.5, -0.6, - float("-inf"), # row C: last col filtered -> padding + float("-inf"), ], ) _, support = _sample_support_from_flat_logprobs(flat_logprobs, top_k=top_k) - sampled_ids = [100, 7, 5] - - # (b) each row keeps width == top_k - assert all(row.size == top_k for row in support) - - # (a) every row's support now contains its sampled id - for sampled_id, row in zip(sampled_ids, support): - assert sampled_id in row - - # (c) the sampled id appears exactly once per repaired row (no duplicate) - assert np.count_nonzero(support[0] == 100) == 1 - assert np.count_nonzero(support[2] == 5) == 1 - - # (d) trailing -1 padding preserved on the padded row - assert support[2][-1] == -1 - - # (e) the unaffected row (sampled already present) is unchanged - np.testing.assert_array_equal(support[1], [7, 8, 9]) - - # Concrete expected repair: weakest (trailing) valid member overwritten. - np.testing.assert_array_equal(support[0], [8, 9, 100]) - np.testing.assert_array_equal(support[2], [6, 5, -1]) + 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(): - # top_k == 1 (row_width == 2): a single support column that must hold the sampled id. flat_logprobs = SimpleNamespace( token_ids=[42, 9], logprobs=[-0.1, -0.2], @@ -128,8 +102,6 @@ async def generate(self, prompt, sampling_params, request_id): @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): - """Support is the sampler's bounded top-k set, so an unbounded or degenerate ``top_k`` has none. - Forwarding it as ``logprobs`` instead yields an opaque 500 from vLLM.""" app = FastAPI() engine = FakeEngine() VLLMServerActor._add_custom_endpoints(app, engine, SimpleNamespace(enable_lora=False)) diff --git a/tests/backends/skyrl_train/utils/test_sample_support.py b/tests/backends/skyrl_train/utils/test_sample_support.py index 3bc5c87e00..fa1344957f 100644 --- a/tests/backends/skyrl_train/utils/test_sample_support.py +++ b/tests/backends/skyrl_train/utils/test_sample_support.py @@ -31,8 +31,6 @@ def test_finalize_rejects_a_trace_shorter_than_the_response(): def test_finalize_rejects_an_unexpected_overshoot(): - """A trailing row count the caller did not declare means the trace and the response disagree - about which tokens were sampled, which silent truncation would hide.""" trace = SampleSupportTrace() trace.append(_rows(3), expected_rows=3) trace.append_padding(2) diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py index 043a3b5392..cc6c82c5af 100644 --- a/tests/train/generators/test_generator_output_utils.py +++ b/tests/train/generators/test_generator_output_utils.py @@ -78,8 +78,6 @@ 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]] - # Both per-token side channels are named fields, so the input order cannot decide whether they - # survive concatenation. 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]) @@ -110,9 +108,6 @@ def test_generator_output_concatenation(): @pytest.mark.parametrize("side_channel", ["rollout_expert_indices", "rollout_sample_support"]) def test_side_channel_concatenation_rejects_a_mix(side_channel): - """A batch that captured a side channel cannot be concatenated with one that did not: the - consumer packs the field for every trajectory in the batch or for none of them.""" - def make_output(value) -> GeneratorOutput: return { "prompt_token_ids": [[1]], @@ -132,7 +127,6 @@ def make_output(value) -> GeneratorOutput: def test_slice_generator_output_slices_each_component_of_a_dict_field(): - """``trajectory_time_splits`` is a dict of per-entry lists rather than a per-entry list.""" generator_output: GeneratorOutput = { "prompt_token_ids": [[1], [2], [3]], "response_ids": [[10], [20], [30]], @@ -589,8 +583,6 @@ def test_per_trajectory_scalar_rewards_and_overlong_filtering(self): assert merged["loss_masks"] == [[1, 0, 1]] def test_sample_support_observation_deltas_keep_full_width_padding_rows(self): - """Every other producer emits dense width-``top_k`` rows, so an observation delta - must too -- an empty row here would make the trajectory ragged.""" tid = _make_tid("support") gen_out: GeneratorOutput = { "prompt_token_ids": [[10], [10, 20, 30]], @@ -614,8 +606,6 @@ def test_sample_support_observation_deltas_keep_full_width_padding_rows(self): assert {len(row) for row in support} == {3} def test_native_output_carrying_dict_valued_time_splits(self): - """The native step-wise ``GeneratorOutput`` carries ``trajectory_time_splits`` as a dict of - per-entry lists; the other fixtures in this file omit the key entirely.""" tid = _make_tid("timed") gen_out: GeneratorOutput = { "prompt_token_ids": [[10], [10, 20, 30]], @@ -636,8 +626,6 @@ def test_native_output_carrying_dict_valued_time_splits(self): merged = merge_stepwise_output(gen_out) - # `_merge_single_trajectory` returns a fixed key set that drops the timing fields, so only the - # merge itself is asserted here. assert merged["response_ids"] == [[20, 30, 40]] assert merged["loss_masks"] == [[1, 0, 1]] assert merged["rewards"] == [[0.0, 0.0, 1.0]] @@ -863,8 +851,6 @@ def test_validate_cfg_merge_stepwise_requires_step_wise(self): @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): - """Refused in `validate_cfg`, which runs before Ray and the engines start, and the message - must state why: a step's routes would replay onto the first N prompt tokens of its row.""" cfg = example_dummy_config() cfg.generator.step_wise_trajectories = True cfg.generator.inference_engine.enable_return_routed_experts = True diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index daf73b3115..ec623096a0 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -490,60 +490,6 @@ def generate(input_batch, model=None): assert all(row == [-1, -1] for row in output.rollout_sample_support[2:-2]) -@pytest.mark.asyncio -@patch("skyrl_gym.make") -async def test_agent_loop_skips_sample_support_capture_on_the_eval_path( - mock_make, - mock_tokenizer, - mock_llm, - mock_env, - generator_cfg, - mock_env_cfg, -): - """The eval phase's sampling params (greedy, top_k=-1) cannot satisfy the capture - contract; the engine-level flag must not force capture on them.""" - generator_cfg.batched = False - generator_cfg.max_turns = 1 - 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={}), - ] - captured = {} - - def generate(input_batch, model=None): - captured.update(input_batch) - return { - "responses": ["mocked output"], - "response_ids": [[10, 11]], - "stop_reasons": ["stop"], - } - - 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, - sampling_params={"temperature": 0.0, "top_k": -1, "max_tokens": 32}, - training_phase=TRAINING_PHASE_EVAL, - ) - - assert captured["return_sample_support"] is False - assert output.rollout_sample_support is None - - @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]) @@ -562,13 +508,6 @@ async def test_generate_requests_sample_support_capture_only_for_the_train_phase batch_sampling_params, batched, ): - """Capture is requested iff the engine flag is on and the batch is a train-phase batch. - - Both phases supply non-None ``sampling_params`` (train from ``generator.sampling_params``, - eval from ``generator.eval_sampling_params``), so ``batch_metadata.training_phase`` is the - only usable discriminator -- the presence of ``sampling_params`` is not one. Driven through - ``generate`` so both the batched and agent-loop routes are covered. - """ generator_cfg.batched = batched generator_cfg.max_turns = 1 generator_cfg.inference_engine.enable_return_sample_support_set = enable_capture @@ -601,7 +540,6 @@ def generate(input_batch, model=None): ) generator.base_conversation_token_ids = [] - # Mirrors `prepare_generator_input`: both phases carry explicit sampling params. input_batch: GeneratorInput = { "prompts": [[{"role": "user", "content": "What is 3 + 5?"}]], "env_classes": [mock_env_cfg.env_class], @@ -623,8 +561,6 @@ def generate(input_batch, model=None): def test_validate_cfg_refuses_routed_experts_without_conversation_multi_turn( mock_tokenizer, mock_llm, generator_cfg, mock_env_cfg ): - """The single-turn convention re-appends a loss-active EOS the inference engine never evaluated, - so no routed-expert row exists for it and the trace refuses to dummy-pad a loss-active target.""" generator_cfg.batched = False generator_cfg.use_conversation_multi_turn = False generator_cfg.inference_engine.enable_return_routed_experts = True @@ -641,9 +577,6 @@ def test_validate_cfg_refuses_routed_experts_without_conversation_multi_turn( def test_validate_cfg_refuses_routed_experts_with_step_wise_trajectories( mock_tokenizer, mock_llm, generator_cfg, mock_env_cfg ): - """A step-wise `generate` has no field to put routes in, so without this refusal a generator built - outside the entrypoint (which is where `validate_cfg` runs) emits `rollout_expert_indices=None` - and trains with routing replay silently disabled.""" generator_cfg.batched = False generator_cfg.step_wise_trajectories = True generator_cfg.inference_engine.enable_return_routed_experts = True @@ -661,7 +594,6 @@ def test_validate_cfg_refuses_routed_experts_with_step_wise_trajectories( 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 ): - """Refused at config time rather than on the first train batch, after a large model has loaded.""" 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 @@ -679,8 +611,6 @@ def test_validate_cfg_refuses_custom_chat_template_with_side_channel_capture( def test_retokenizing_state_update_refuses_a_live_side_channel_trace( mock_tokenizer, mock_llm, generator_cfg, mock_env_cfg ): - """The retokenizing state update is the one helper that never feeds a trace, so a trace reaching - it would finalize empty (support) or misaligned (routes) instead of failing.""" generator_cfg.batched = False generator_cfg.chat_template = ChatTemplateConfig(source="name", name_or_path="qwen3_without_thinking") generator = SkyRLGymGenerator( @@ -725,8 +655,6 @@ async def test_generate_retains_routed_experts_only_for_the_train_phase( training_phase, batched, ): - """Nothing reads an eval trajectory's routes, and the driver holds the whole eval batch at once - (~3.9 KB per token at 120B), so eval routes must not be retained.""" generator_cfg.batched = batched generator_cfg.max_turns = 1 generator_cfg.use_conversation_multi_turn = True @@ -773,7 +701,6 @@ def generate(input_batch, model=None): assert output["rollout_expert_indices"] is not None assert output["rollout_expert_indices"][0] is not None if not batched: - # A captured trace also asks the engine where its prompt prefix already ended. assert captured["routed_experts_prompt_starts"] == [0] else: assert output["rollout_expert_indices"] is None @@ -791,10 +718,6 @@ async def test_agent_loop_keeps_the_generated_eos_support_row_in_single_turn_mod generator_cfg, mock_env_cfg, ): - """With `use_conversation_multi_turn=False` the generated EOS is sliced off the turn and - re-appended to the trajectory as the loss-active token carrying the terminal reward. It keeps - the support set it was actually sampled from rather than an all-padding row, which a replay - consumer would read as a synthetic EOS and score over the full vocabulary.""" generator_cfg.batched = False generator_cfg.max_turns = 1 generator_cfg.use_conversation_multi_turn = False @@ -810,7 +733,6 @@ async def test_agent_loop_keeps_the_generated_eos_support_row_in_single_turn_mod def generate(input_batch, model=None): return { "responses": ["mocked output"], - # The engine's own EOS (id 4) terminates the turn. "response_ids": [[10, 11, 4]], "stop_reasons": ["stop"], "rollout_sample_support": [np.array([[10, 110], [11, 111], eos_support_row], dtype=np.int32)], @@ -834,7 +756,6 @@ def generate(input_batch, model=None): ) assert output.response_ids == [10, 11, 4] - # The re-appended EOS is loss-active and carries the terminal reward. assert output.loss_mask == [1, 1, 1] assert output.rollout_sample_support == [[10, 110], [11, 111], eos_support_row] @@ -849,8 +770,6 @@ async def test_agent_loop_pads_a_stop_string_eos_support_row_in_single_turn_mode generator_cfg, mock_env_cfg, ): - """An EOS the loop appends because generation stopped on a stop string was never sampled, so it - gets an all-padding row.""" generator_cfg.batched = False generator_cfg.max_turns = 1 generator_cfg.use_conversation_multi_turn = False diff --git a/tests/train/generators/test_skyrl_vlm_generator.py b/tests/train/generators/test_skyrl_vlm_generator.py index 9758a524ab..ad3105da61 100644 --- a/tests/train/generators/test_skyrl_vlm_generator.py +++ b/tests/train/generators/test_skyrl_vlm_generator.py @@ -135,8 +135,6 @@ async def mock_generate(input_batch, model=None): def test_vlm_validate_cfg_still_applies_the_base_refusals(): - """``SkyRLVLMGymGenerator._validate_cfg`` replaces the base method, so it must chain to it; - otherwise every base refusal is silently void for this generator.""" tokenizer = MagicMock() tokenizer.apply_chat_template.side_effect = lambda messages, **kwargs: [1, 2, 3, 4] tokenizer.eos_token_id = 4 diff --git a/tests/train/test_config.py b/tests/train/test_config.py index 1e136be3e2..7224613654 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -180,8 +180,6 @@ def test_sample_support_capture_rejects_unsupported_sampling_modifiers(override, def test_routed_expert_capture_rejects_the_vision_language_generator(): - """The VLM generator never populates ``rollout_expert_indices``, so without this guard the pair - generates a whole batch and then dies in collation on a None route array.""" with pytest.raises(ValueError, match="vision_language_generator"): SkyRLTrainConfig.from_cli_overrides( [ @@ -205,8 +203,6 @@ def test_sample_support_capture_accepts_top_k_top_p_and_min_p(): def test_sample_support_capture_leaves_greedy_eval_sampling_params_alone(): - # The five guards deliberately do not reach eval_sampling_params: forcing - # temperature > 0 there would change eval semantics. Capture is opted out per request. cfg = SkyRLTrainConfig.from_cli_overrides( [ "generator.inference_engine.enable_return_sample_support_set=true", From 098017bff4d995f2f7c0adab673e632a8a74dfe2 Mon Sep 17 00:00:00 2001 From: lila-sync-bot Date: Fri, 14 Aug 2026 23:14:56 +0000 Subject: [PATCH 16/17] perf(sample-support): carry the captured support to the trainer as one packed field --- skyrl/backends/skyrl_train/training_batch.py | 78 +++++- .../skyrl_train/utils/packed_tensor.py | 19 ++ .../skyrl_train/utils/replay_utils.py | 33 +-- .../skyrl_train/utils/routed_experts.py | 26 +- .../skyrl_train/utils/sample_support.py | 58 +++++ .../workers/megatron/megatron_worker.py | 16 +- .../skyrl_train/workers/worker_utils.py | 28 ++- skyrl/train/dataset/preprocess.py | 93 +++++++ skyrl/train/dataset/replay_buffer.py | 6 + skyrl/train/generators/base.py | 13 +- skyrl/train/generators/skyrl_gym_generator.py | 14 +- skyrl/train/generators/utils.py | 34 ++- skyrl/train/trainer.py | 7 + skyrl/train/utils/trainer_utils.py | 67 +++++ .../distributed/test_token_metadata.py | 17 +- .../gpu_ci/megatron/test_megatron_models.py | 2 +- .../gpu/gpu_ci/megatron/test_router_replay.py | 20 +- .../test_token_based_batching_utils.py | 48 +++- .../backends/skyrl_train/test_train_batch.py | 92 +++++++ .../skyrl_train/utils/test_replay_utils.py | 32 +-- .../skyrl_train/utils/test_routed_experts.py | 40 +++ .../utils/test_sample_support_row_ids.py | 163 ++++++++++++ tests/train/dataset/test_preprocess.py | 232 ++++++++++++++++-- .../generators/test_generator_output_utils.py | 20 +- .../generators/test_skyrl_gym_generator.py | 26 +- ...est_collation_vectorization_equivalence.py | 8 +- ...test_packed_route_collation_equivalence.py | 21 ++ tests/train/test_trainer_utils.py | 156 ++++++++++++ 28 files changed, 1195 insertions(+), 174 deletions(-) create mode 100644 tests/backends/skyrl_train/utils/test_sample_support_row_ids.py diff --git a/skyrl/backends/skyrl_train/training_batch.py b/skyrl/backends/skyrl_train/training_batch.py index 9981fb62f7..61294bffff 100644 --- a/skyrl/backends/skyrl_train/training_batch.py +++ b/skyrl/backends/skyrl_train/training_batch.py @@ -3,6 +3,8 @@ import copy import io import pickle +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 @@ -10,8 +12,15 @@ import torch from jaxtyping import Bool, Float, Integer -from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor -from skyrl.backends.skyrl_train.utils.replay_utils import append_packed_replay_padding +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") @@ -162,7 +171,7 @@ 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"}) + ZERO_COPY_KEYS: frozenset[str] = frozenset({"rollout_expert_indices", SAMPLE_SUPPORT_FIELD}) def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -531,6 +540,8 @@ class TrainingInput(TypedDict, total=False): # 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 @@ -547,6 +558,57 @@ class TrainingOutputBatch(TensorBatch[Dict[str, torch.Tensor]]): pass +@dataclass(frozen=True) +class PackedFieldPadding: + """How one packed ``TrainingInput`` field fills the segments batch padding appends. + + ``dummy_row_length`` is the segment length for a synthetic batch row, which carries a + single attended token: a field indexed over every real token needs one row for it, a + field indexed over response tokens needs none. + """ + + fill: Callable[[PackedTensor], Union[torch.Tensor, int]] + dummy_row_length: int + + +# Every packed batch field needs an entry here: the three batch padding sites look its rule +# up by name, so a field without one raises instead of reaching the trainer short a segment +# (or, in `_pad_microbatch_to_size`, being skipped as a non-Tensor). +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. @@ -574,15 +636,17 @@ 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): + # Every other field copies row 0 into the padding rows, so each padded row spans as + # many tokens as row 0 and needs a segment of row 0's 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": - # Every other field copies row 0 into the padding rows, so each padded row holds - # as many real tokens as row 0 and needs a route segment of that length. - new_tensors[key] = append_packed_replay_padding(tensor, segment_lengths=[len(tensor.segment(0))] * pad_size) 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 index c5a0c20983..55fcd67f92 100644 --- a/skyrl/backends/skyrl_train/utils/packed_tensor.py +++ b/skyrl/backends/skyrl_train/utils/packed_tensor.py @@ -169,3 +169,22 @@ def cat(batches: Sequence["PackedTensor"]) -> "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 3e22c7afd7..47a538be67 100644 --- a/skyrl/backends/skyrl_train/utils/replay_utils.py +++ b/skyrl/backends/skyrl_train/utils/replay_utils.py @@ -2,7 +2,6 @@ Utility functions for MoE Router Replay. """ -from collections.abc import Sequence from contextlib import contextmanager import torch @@ -12,10 +11,7 @@ align_packed_token_metadata, align_token_metadata, ) -from skyrl.backends.skyrl_train.utils.packed_tensor import ( - PackedTensor, - cu_seqlens_from_lengths, -) +from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor def replay_padding_row( @@ -49,33 +45,6 @@ def make_replay_padding_indices( return padding_row.expand(shape).clone() -def make_packed_replay_padding( - reference: PackedTensor, - *, - segment_lengths: Sequence[int], -) -> PackedTensor: - """Return dummy-route segments matching ``reference``'s row shape and dtype. - - Batch padding rows exist only to give Megatron a uniform micro-batch size; their - tokens are loss-masked, so one dummy route per token is all they need. - """ - padding = make_replay_padding_indices( - (sum(segment_lengths), *reference.row_shape), - dtype=reference.dtype, - device=reference.device, - ) - return PackedTensor(padding, cu_seqlens_from_lengths(segment_lengths, device=reference.device)) - - -def append_packed_replay_padding( - routes: PackedTensor, - *, - segment_lengths: Sequence[int], -) -> PackedTensor: - """Extend ``routes`` with one dummy-route segment per batch padding row.""" - return PackedTensor.cat([routes, make_packed_replay_padding(routes, segment_lengths=segment_lengths)]) - - def patch_topk_router_layer_number(): """Monkey-patch TopKRouter.set_layer_number to propagate the global layer number to the RouterReplay instance. diff --git a/skyrl/backends/skyrl_train/utils/routed_experts.py b/skyrl/backends/skyrl_train/utils/routed_experts.py index a35bcc543a..1ec3048626 100644 --- a/skyrl/backends/skyrl_train/utils/routed_experts.py +++ b/skyrl/backends/skyrl_train/utils/routed_experts.py @@ -16,7 +16,6 @@ class RoutedExpertTrace: def __init__(self) -> None: self._metadata = TokenMetadataTrace() - self._schema: tuple[int, int, np.dtype] | None = None @property def prompt_start(self) -> int: @@ -35,12 +34,18 @@ def record_generation( raise ValueError("routed-expert generation must produce at least one token") expected_rows = prompt_token_count - self.prompt_start + generated_token_count - 1 - compact = compact_routed_expert_indices(routed_experts) - if self._schema is None: - self._schema = (*compact.shape[1:], compact.dtype) - self._metadata.append(compact, expected_rows=expected_rows) + 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 routes, which cover a prefix of the sequence's real tokens. + + The capture ends short of ``token_count``: the last sampled token has no subsequent + decode forward to record its route, and a synthetic EOS is never evaluated at all. + The row count is the trace's only report of where the capture stops -- collation + dummy-fills the uncovered tail and ``make_router_padding_mask`` excludes exactly those + rows from router accounting -- so the trace must not pad the tail itself, which would + report fabricated routes as captured ones. + """ 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: @@ -51,16 +56,7 @@ def finalize(self, *, token_count: int, loss_mask: Sequence[int]) -> RoutedExper if loss_mask[source_index + 1] != 0: raise ValueError(f"missing routed-expert row for loss-active target at token {source_index + 1}") - padding_count = token_count - self.prompt_start - if padding_count: - if self._schema is None: - raise ValueError("cannot pad routed-expert trace before any routes are captured") - num_layers, topk, dtype = self._schema - padding_row = np.arange(topk, dtype=dtype) - padding = np.broadcast_to(padding_row, (padding_count, num_layers, topk)).copy() - self._metadata.append(padding, expected_rows=padding_count) - - return self._metadata.finalize(expected_rows=token_count) + return self._metadata.finalize(expected_rows=self.prompt_start) def compact_routed_expert_indices(routed_experts: RoutedExpertIndices) -> RoutedExpertIndices: diff --git a/skyrl/backends/skyrl_train/utils/sample_support.py b/skyrl/backends/skyrl_train/utils/sample_support.py index 5fb78367fd..ba5a3be155 100644 --- a/skyrl/backends/skyrl_train/utils/sample_support.py +++ b/skyrl/backends/skyrl_train/utils/sample_support.py @@ -6,15 +6,25 @@ 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 +# Names the support field in ``GeneratorOutput``, ``TrainingInput`` and ``Experience``. +SAMPLE_SUPPORT_FIELD = "rollout_sample_support" +# Row-id channel value for a model position no support row scores. Out of range for any +# packed row index, so a gather by id cannot silently pick up a real row. +SAMPLE_SUPPORT_NO_ROW = -1 def validate_sample_support(sample_support: SampleSupport) -> SampleSupport: @@ -33,6 +43,54 @@ def validate_sample_support(sample_support: SampleSupport) -> SampleSupport: return sample_support +def align_sample_support_row_ids( + sample_support: PackedTensor, + layout: TokenMetadataLayout, +) -> torch.Tensor: + """Return the per-token channel naming which packed support row scores each model position. + + The payload itself must not go through ``align_packed_token_metadata``: that places a + segment at a fixed offset inside its trajectory's padded region, while support scoring + happens one position to the left of the token it describes -- the logit at position ``t`` + predicts token ``t + 1``. A trajectory's support therefore covers real tokens + ``[p_i - 1, p_i + r_i - 1)``, which includes the last prompt token and excludes the last + response token. Aligning only these int64 row ids keeps that placement in one place, and + the scorer gathers ``[top_k]`` rows by id. + + Row ids index ``sample_support.values``, so they must be derived per micro-batch: + ``chunk``, ``slice`` and batch padding all 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, + ) + # p_i = L_i - r_i, and the position predicting the first response token is p_i - 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.""" diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 3b99902d56..a3adb1f06e 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -45,10 +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 append_packed_replay_padding from skyrl.backends.skyrl_train.weight_sync import ( LoraLoadRequest, WeightChunk, @@ -784,13 +785,14 @@ def _pad_microbatch_to_size(self, micro_dict: dict, target_batch_size: int) -> d if value is None: padded[key] = None continue - if key == "rollout_expert_indices": - # The dummy attention_mask row below marks one valid token, so each padded - # row's route segment holds one row. - padded[key] = append_packed_replay_padding(value, segment_lengths=[1] * pad_count) - continue if isinstance(value, PackedTensor): - raise ValueError(f"Micro-batch field {key!r} is packed and has no padding rule") + # The dummy attention_mask row below marks one valid token, so a per-token field + # gets one dummy row for it while a response-token field gets none. A packed field + # with no rule raises here rather than falling through as a non-Tensor. + 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 diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py index d3f04b5e79..9b466c6b9e 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_packed_replay_padding +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,16 +328,19 @@ 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 such as `rollout_logprobs` and the packed side channels to 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: - # The dummy attention_mask row marks one valid token, so its route segment holds one row. - data["rollout_expert_indices"] = make_packed_replay_padding( - self.data["rollout_expert_indices"], - segment_lengths=[1] * batch_size, - ) + for key in PACKED_FIELD_PADDING: + # The dummy attention_mask row marks one valid token, so a per-token field gets one + # dummy row for it while a response-token field gets none. + 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/train/dataset/preprocess.py b/skyrl/train/dataset/preprocess.py index 23aa0a123c..10d785184c 100644 --- a/skyrl/train/dataset/preprocess.py +++ b/skyrl/train/dataset/preprocess.py @@ -15,6 +15,11 @@ ROUTED_EXPERT_DTYPES, RoutedExpertIndices, ) +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__) @@ -180,6 +185,79 @@ def _collate_rollout_expert_indices( 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 per-trajectory sampler support into one ``[sum(response_len_i), top_k]`` buffer. + + Support describes generated tokens only, so it packs to the response tokens rather than to + a ``[batch, seq_len, top_k]`` rectangle whose whole prompt region would be padding written + on the driver and read back only to be discarded. The outer ragged level is one segment per + trajectory, exactly as for packed routes, so the trainer indexes both by segment. + + The wire side establishes the canonical dtype and the trailing-padding invariant, so entries + are validated rather than rescanned here. The fill runs from a locally sized thread pool for + the same reason route collation does: it touches the whole global batch before DP sharding, + on every training step, and is bound by first-touch page faults on a fresh mapping. + """ + 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]], @@ -188,6 +266,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"], @@ -197,6 +276,7 @@ def convert_prompts_responses_to_batch_tensors( Float[torch.Tensor, "batch response_len"], Optional[Float[torch.Tensor, "batch response_len"]], Optional[PackedTensor], + Optional[PackedTensor], ]: """ Convert prompts and responses to batch tensors for training. @@ -256,6 +336,9 @@ def convert_prompts_responses_to_batch_tensors( 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) @@ -333,6 +416,15 @@ def convert_prompts_responses_to_batch_tensors( 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, attention_mask, @@ -341,6 +433,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 35668a310a..efada7b4dd 100644 --- a/skyrl/train/dataset/replay_buffer.py +++ b/skyrl/train/dataset/replay_buffer.py @@ -73,6 +73,8 @@ class Experience: 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 @@ -106,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: @@ -137,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 24180fd463..e77c29d273 100644 --- a/skyrl/train/generators/base.py +++ b/skyrl/train/generators/base.py @@ -6,6 +6,7 @@ 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" @@ -53,11 +54,17 @@ 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, one ``[tokens, layers, topk]`` array of the routes the rollout took over a + # prefix of its ``prompt + response`` tokens: no decode forward follows the last sampled token, + # and a multi-turn trace ends further short of a synthetic EOS. Collation + # dummy-fills the uncovered tail and the router padding mask keeps it out of router accounting, + # so the row count is what states where the capture stops. rollout_expert_indices: Optional[List[RoutedExpertIndices]] - # Per trajectory, one dense ``[response_tokens, top_k]`` row block of the sampler support each + # Per trajectory, one dense ``[response_tokens, top_k]`` array of the sampler support each # response token was drawn from, right-padded with ``SAMPLE_SUPPORT_PADDING``. Tokens with no - # captured support (observations, a synthetic EOS) are all-padding rows. - rollout_sample_support: Optional[List[List[List[int]]]] + # captured support (observations, a synthetic EOS) are all-padding rows. Stays an ndarray from + # the wire to the packed trainer field: nested lists of it cost ~50x the int32 buffer. + 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 9f44f53491..e740857445 100644 --- a/skyrl/train/generators/skyrl_gym_generator.py +++ b/skyrl/train/generators/skyrl_gym_generator.py @@ -65,7 +65,7 @@ class TrajectoryOutput: rollout_logprobs: Optional[List[float]] env_metrics: Dict[str, Any] rollout_expert_indices: Optional[RoutedExpertIndices] = None - rollout_sample_support: Optional[List[List[int]]] = 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 @@ -598,9 +598,7 @@ 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_sample_support=( - turn_sample_support.tolist() if turn_sample_support is not None else None - ), + rollout_sample_support=turn_sample_support, ) agent_loop_output.step_outputs.append(per_step_output) @@ -709,7 +707,7 @@ async def agent_loop( rollout_sample_support_out = agent_loop_state.sample_support_trace.finalize( token_count=len(response_ids), extra_rows=final_observation_token_count, - ).tolist() + ) 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): @@ -919,9 +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[List[List[int]]]] = ( - [] if raw_rollout_sample_support 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 @@ -941,7 +937,7 @@ async def generate_batched( 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)].tolist()) + truncated_sample_support.append(raw_rollout_sample_support[i][: len(response)]) # Get environment-specific metrics env_metrics.append(env.get_metrics()) diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py index 18ab425ee5..590e23d1ea 100644 --- a/skyrl/train/generators/utils.py +++ b/skyrl/train/generators/utils.py @@ -8,7 +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_PADDING +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, @@ -806,6 +810,11 @@ def slice_generator_output( 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. @@ -827,9 +836,7 @@ def _merge_single_trajectory(gen_out: GeneratorOutput) -> GeneratorOutput: 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 = ( - next((len(row) for rows in gen_out["rollout_sample_support"] for row in rows), 0) if has_sample_support else 0 - ) + 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 @@ -837,7 +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 - out_sample_support: Optional[List[List[List[int]]]] = [] if has_sample_support else None + # One row block per merged turn; concatenated at flush rather than extended row by row. + 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 = [] @@ -851,8 +859,8 @@ 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[List[int]]] = ( - [list(row) for row in gen_out["rollout_sample_support"][0]] if has_sample_support 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 @@ -865,7 +873,7 @@ def flush(): if has_logprobs: out_logprobs.append(acc_logprobs) if has_sample_support: - out_sample_support.append(acc_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]) @@ -883,9 +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 = ( - [list(row) for row in gen_out["rollout_sample_support"][i]] if has_sample_support 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 @@ -901,7 +907,9 @@ def flush(): if acc_logprobs is not None: acc_logprobs.extend([0.0] * len(obs_delta)) if acc_sample_support is not None: - acc_sample_support.extend([SAMPLE_SUPPORT_PADDING] * sample_support_width for _ in obs_delta) + 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)) @@ -911,7 +919,7 @@ def flush(): if acc_logprobs is not None: acc_logprobs.extend(gen_out["rollout_logprobs"][i]) if acc_sample_support is not None: - acc_sample_support.extend(gen_out["rollout_sample_support"][i]) + 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]) 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 05c258bd14..0599b57bf5 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, @@ -748,10 +749,76 @@ 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 the per-generated-token side channels against what the trainer consumes. + + The outer length checks above cover only the per-trajectory list; a ``None`` entry or a + row count the trainer cannot place still reaches ``convert_prompts_responses_to_batch_tensors`` + (or, for routes, ``make_router_padding_mask`` one step earlier) as an opaque ``TypeError``. + + Routes cover a non-empty *prefix* of a trajectory's ``prompt + response`` tokens: vLLM + records no route for the last sampled token, and a multi-turn trace ends further short of + a synthetic EOS. Collation dummy-fills the uncovered tail and the router padding mask + excludes it, so the row count is bounded, not fixed -- bounded above by the sequence and + below by the last token the loss trains. Sample support is dense over the response, so its + row count is exact. + """ + 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"] + + # Route rows are aligned to one contiguous prompt+response token sequence, while step-wise + # splits that trajectory into per-turn samples whose prompts re-cover earlier turns. The + # trajectory-aligned rows would land on the wrong tokens of every step but the first, which + # the row-count bounds below cannot see: they check coverage, not which token a row describes. + 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" + ) + # The row at source position ``t`` holds the route that produced token ``t + 1``, so + # ``captured_rows`` rows cover targets ``[1, captured_rows]``. Beyond that the trainer + # replays dummy routes, and the router padding mask only keeps those out of router + # accounting -- it cannot stop a trained token from being replayed on a route the + # rollout never took. ``RoutedExpertTrace.finalize`` proves the same bound for traces + # it builds; this is the boundary check for routes from any other producer. + 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/tests/backends/skyrl_train/distributed/test_token_metadata.py b/tests/backends/skyrl_train/distributed/test_token_metadata.py index 3d518fc782..426ff4372c 100644 --- a/tests/backends/skyrl_train/distributed/test_token_metadata.py +++ b/tests/backends/skyrl_train/distributed/test_token_metadata.py @@ -136,13 +136,19 @@ def test_routed_expert_trace_tracks_multiturn_suffix_and_terminal_gap() -> None: 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 result.shape == (9, 2, 2) and result.dtype == np.uint8 - assert np.array_equal(result[-1, 0], [0, 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_only_pads_masked_suffix(active: bool) -> None: +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)] @@ -150,8 +156,11 @@ def test_routed_expert_trace_only_pads_masked_suffix(active: bool) -> None: 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 np.array_equal(result[-2:, 0], [[0, 1], [0, 1]]) + assert result.shape == (3, 2, 2) + assert np.array_equal(result, routes(3).astype(result.dtype)) @pytest.mark.parametrize("packed", [False, True]) 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 89bfd85515..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 @@ -128,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] @@ -237,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"], @@ -362,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, @@ -485,7 +487,7 @@ def test_forward_backward_variable_length_full_recompute(tp, pp, cp, ep, etp, ex 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( tokenizer=tokenizer, prompts=prompts, 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 c78d0fee80..b2a7e25f0a 100644 --- a/tests/backends/skyrl_train/test_token_based_batching_utils.py +++ b/tests/backends/skyrl_train/test_token_based_batching_utils.py @@ -16,6 +16,11 @@ 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, @@ -203,13 +208,21 @@ def test_padding_microbatch_matches_seq_len(self): # Padding rows must not contribute to the loss. assert padding["loss_mask"].sum().item() == 0 - def test_padding_microbatch_uses_unique_dummy_routes(self): - batch = self._make_batch([4, 4], num_actions=2) + 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) + self._add_packed_side_channels(batch) iterator = TokenBasedBatchIterator(batch, max_tokens_per_microbatch=8) padding = iterator._create_padding_microbatch() @@ -220,6 +233,37 @@ def test_padding_microbatch_uses_unique_dummy_routes(self): 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): + """The dummy row attends one prompt-side token and generates nothing, so its support + segment is empty -- and the route field keeps its own one-row dummy segment.""" + 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( diff --git a/tests/backends/skyrl_train/test_train_batch.py b/tests/backends/skyrl_train/test_train_batch.py index ff7fb8286b..353539d5ed 100644 --- a/tests/backends/skyrl_train/test_train_batch.py +++ b/tests/backends/skyrl_train/test_train_batch.py @@ -1,5 +1,6 @@ import pickle from collections.abc import Callable +from typing import get_args import numpy as np import pytest @@ -7,6 +8,7 @@ import torch from skyrl.backends.skyrl_train.training_batch import ( + PACKED_FIELD_PADDING, BatchField, TensorBatch, TensorFormat, @@ -15,6 +17,9 @@ 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 ( @@ -22,6 +27,11 @@ 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(): @@ -562,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", } @@ -592,6 +603,11 @@ def _make_full_training_batch(batch_size: int = 4, seq_len: int = 5) -> Training 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 } @@ -660,10 +676,17 @@ def test_pad_batch_all_fields(): 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", } @@ -784,6 +807,10 @@ def test_serialized_field_formats_are_stable(): 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), + ), } @@ -865,3 +892,68 @@ def test_zero_copy_falls_back_for_bfloat16(): values = torch.randn(3, 4, dtype=torch.bfloat16) assert _serialize_tensor(values, zero_copy=True)["format"] == TensorFormat.TORCH + + +# ── packed field padding ───────────────────────────────────────────────────── + +# The row each field's padding segments must carry, stated independently of the production +# rule: distinct experts for Megatron's dropless dispatcher, "no support" for the sampler +# support. Padding a field with the other's fill would mis-train silently. +_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): + """All three batch padding sites append through here, so the round trip is asserted once.""" + 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): + """A synthetic batch row attends one token: a per-token field needs one row for it, a + response-token field needs none.""" + 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_replay_utils.py b/tests/backends/skyrl_train/utils/test_replay_utils.py index f713908a92..5320da509d 100644 --- a/tests/backends/skyrl_train/utils/test_replay_utils.py +++ b/tests/backends/skyrl_train/utils/test_replay_utils.py @@ -11,11 +11,7 @@ ) from skyrl.backends.skyrl_train.utils import replay_utils from skyrl.backends.skyrl_train.utils.packed_tensor import PackedTensor -from skyrl.backends.skyrl_train.utils.replay_utils import ( - append_packed_replay_padding, - make_packed_replay_padding, - make_replay_padding_indices, -) +from skyrl.backends.skyrl_train.utils.replay_utils import make_replay_padding_indices def _pack_routes(routes: torch.Tensor, attention_mask: torch.Tensor) -> PackedTensor: @@ -82,32 +78,6 @@ def test_replay_padding_rejects_missing_topk(shape): make_replay_padding_indices(shape, dtype=torch.uint8) -@pytest.mark.parametrize("segment_lengths", [[1, 1], [3], [2, 5, 1]]) -def test_packed_replay_padding_matches_the_reference_row_shape(segment_lengths): - reference = PackedTensor.from_segments([torch.full((4, 2, 3), 9, dtype=torch.int16)]) - - padding = make_packed_replay_padding(reference, segment_lengths=segment_lengths) - - assert padding.sequence_lengths.tolist() == segment_lengths - assert padding.row_shape == reference.row_shape - assert padding.dtype == reference.dtype - assert torch.equal(padding.values, torch.tensor([0, 1, 2], dtype=torch.int16).expand_as(padding.values)) - - -@pytest.mark.parametrize("pad_count", [1, 3]) -def test_appending_replay_padding_keeps_the_real_segments_and_the_arange_invariant(pad_count): - routes = PackedTensor.from_segments( - [torch.full((4, 2, 3), 9, dtype=torch.int16), torch.full((2, 2, 3), 8, dtype=torch.int16)] - ) - - padded = append_packed_replay_padding(routes, segment_lengths=[1] * pad_count) - - assert padded.sequence_lengths.tolist() == [4, 2] + [1] * pad_count - assert padded[: len(routes)] == routes - appended = padded[len(routes) :] - assert torch.equal(appended.values, torch.tensor([0, 1, 2], dtype=torch.int16).expand_as(appended.values)) - - def test_replay_has_no_dispatcher_specific_patch(): assert "TokenDispatcher" not in inspect.getsource(replay_utils) diff --git a/tests/backends/skyrl_train/utils/test_routed_experts.py b/tests/backends/skyrl_train/utils/test_routed_experts.py index 432d9d31c9..1d7a957449 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,42 @@ 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 is what tells the trainer where the capture stops, so the trace never + fabricates rows for the tail it did not cover: collation dummy-fills that tail and the + router padding mask keeps those rows out of router accounting. + """ + 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(): + """The uncovered tail is only safe because every loss-active target inside it is masked: + a forced route at a masked position can perturb nothing but later masked positions. + """ + 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_row_ids.py b/tests/backends/skyrl_train/utils/test_sample_support_row_ids.py new file mode 100644 index 0000000000..4f5d452905 --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_sample_support_row_ids.py @@ -0,0 +1,163 @@ +"""Row-id derivation for packed sampler support. + +The support payload is never pushed through ``align_packed_token_metadata``: its domain is a +response suffix shifted one position left, which segment placement alone cannot express. Only +this int64 row-id channel is aligned, and the scorer gathers ``[top_k]`` rows by id. + +Run with: +uv run --isolated --extra dev --extra skyrl-train pytest tests/backends/skyrl_train/utils/test_sample_support.py +""" + +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 +# (prompt_len, response_len) per trajectory: anti-correlated, so the support suffixes start at +# different offsets and a mis-shifted placement cannot pass by coincidence. +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(): + """A trajectory's support covers real tokens ``[p - 1, p + r - 1)``: it includes the last + prompt token and excludes the last response token.""" + 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): + """Under sequence packing the channel is one row of ``[seq0, pad0, seq1, pad1, ...]``.""" + 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] + + +def test_row_ids_rebase_under_chunk(): + """``chunk`` rebases the packed row space, so ids carried in the batch would be wrong.""" + support = _support(LENGTHS) + batch = TrainingInputBatch({"attention_mask": _attention_mask(LENGTHS).long(), SAMPLE_SUPPORT_FIELD: support}) + full_batch_ids = align_sample_support_row_ids(support, _layout(LENGTHS)) + + chunk = batch.chunk(1)[1] + chunk_ids = align_sample_support_row_ids(chunk[SAMPLE_SUPPORT_FIELD], _layout(LENGTHS[1:])) + + # The chunk's own row space starts at 0 again, so its ids differ from the batch's. + assert chunk_ids[chunk_ids >= 0].tolist() == [0, 1] + assert full_batch_ids[1][full_batch_ids[1] >= 0].tolist() == [3, 4] + # Gathering with the rebased ids still lands on the same support rows. + gathered = chunk[SAMPLE_SUPPORT_FIELD].values[chunk_ids[chunk_ids >= 0]] + assert torch.equal(gathered, support.segment(1)) + + +def test_row_ids_rebase_under_slice(): + support = _support(LENGTHS) + batch = TrainingInputBatch({"attention_mask": _attention_mask(LENGTHS).long(), SAMPLE_SUPPORT_FIELD: support}) + + sliced = batch.slice(1, 2) + row_ids = align_sample_support_row_ids(sliced[SAMPLE_SUPPORT_FIELD], _layout(LENGTHS[1:])) + + assert torch.equal(sliced[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_preprocess.py b/tests/train/dataset/test_preprocess.py index 91e7db7b1e..556b3c6210 100644 --- a/tests/train/dataset/test_preprocess.py +++ b/tests/train/dataset/test_preprocess.py @@ -10,9 +10,19 @@ import pytest import torch -from skyrl.backends.skyrl_train.utils.routed_experts import ROUTED_EXPERT_DTYPES +from skyrl.backends.skyrl_train.utils.routed_experts import ( + ROUTED_EXPERT_DTYPES, + RoutedExpertTrace, +) +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, + build_sample_support, convert_prompts_responses_to_batch_tensors, make_router_padding_mask, ) @@ -70,6 +80,43 @@ def test_router_padding_mask_marks_left_padding_and_uncaptured_suffix(): assert mask.tolist() == [[True, False, False, True], [False, False, False, False]] +def test_router_padding_mask_marks_the_tail_a_multi_turn_trace_never_captured(): + """A trace reports only the rows it captured, so the mask covers exactly the tail that + collation dummy-fills. Two turns over a 3-token prompt and a 2-token observation capture + 8 rows of a 10-token sequence: no decode forward followed the last sampled token, and the + synthetic EOS was never evaluated at all. + """ + trace = RoutedExpertTrace() + trace.record_generation( + prompt_token_count=3, + generated_token_count=2, + routed_experts=np.zeros((4, 2, 2), dtype=np.int16), + ) + trace.record_generation( + prompt_token_count=7, + generated_token_count=2, + routed_experts=np.zeros((4, 2, 2), dtype=np.int16), + ) + # The last two tokens are loss-masked, which is what the trace's own guard enforces. + routes = trace.finalize(token_count=10, loss_mask=[0, 0, 0, 1, 1, 0, 0, 1, 1, 0]) + assert routes.shape[0] == 8 + + attention_mask = torch.tensor([[0, 0] + [1] * 10]) + + mask = make_router_padding_mask(attention_mask, [routes.shape[0]]) + + assert mask.tolist() == [[True, True] + [False] * 8 + [True, True]] + + +def test_router_padding_mask_marks_the_last_token_for_batched_routes(): + """``generate_batched`` returns ``seq_len - 1`` rows, one short of its sequence.""" + attention_mask = torch.tensor([[0, 0, 1, 1, 1, 1, 1]]) + + mask = make_router_padding_mask(attention_mask, [4]) + + assert mask.tolist() == [[True, True, False, False, False, False, True]] + + def test_routed_expert_tensor_uses_unique_dummy_routes(tokenizer): routes = [ np.asarray( @@ -89,7 +136,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]], @@ -119,7 +166,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]], @@ -136,7 +183,7 @@ 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]], @@ -167,7 +214,7 @@ def test_routed_expert_tensor_accepts_non_contiguous_arrays(tokenizer): routes = base[:, :, ::2] assert not routes.flags.c_contiguous - *_, routed = convert_prompts_responses_to_batch_tensors( + *_, routed, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts=[[10]], responses=[[11]], @@ -200,7 +247,7 @@ 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]], @@ -221,7 +268,7 @@ def test_routed_expert_tensor_warns_only_on_an_int32_batch(tokenizer, caplog, dt 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( + *_, routed, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts=[[10]], responses=[[11]], @@ -266,7 +313,7 @@ def test_routed_expert_tensor_is_bit_identical_to_numpy_collation(tokenizer): (np.arange(6 * num_layers * topk, dtype=np.int16) + 300).reshape(6, num_layers, topk), ] - *_, routed = convert_prompts_responses_to_batch_tensors( + *_, routed, _ = convert_prompts_responses_to_batch_tensors( tokenizer.pad_token_id, prompts, responses, @@ -309,7 +356,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, @@ -345,7 +392,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, @@ -426,7 +473,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, @@ -456,7 +503,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, @@ -491,7 +538,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, @@ -517,7 +564,7 @@ def test_rollout_expert_indices_shape_padding_and_alignment(tokenizer): 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, @@ -541,7 +588,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, @@ -561,7 +608,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, @@ -578,3 +625,156 @@ 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 — packed to response tokens +# --------------------------------------------------------------------------- + +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_allocates_no_prompt_region_rectangle(tokenizer): + """A ``[batch, seq_len, top_k]`` rectangle would be mostly prompt-region filler.""" + support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS) + + *_, packed = _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + max_total = max(prompt_len + response_len for prompt_len, response_len in SAMPLE_SUPPORT_LENGTHS) + rectangle_rows = len(SAMPLE_SUPPORT_LENGTHS) * max_total + assert packed.values.shape[0] == sum(response_len for _, response_len in SAMPLE_SUPPORT_LENGTHS) + assert packed.values.shape[0] < rectangle_rows + assert packed.values.numel() == packed.values.shape[0] * SAMPLE_SUPPORT_TOP_K + + +def test_sample_support_preserves_padding_rows(tokenizer): + """Padding stays ``-1``: not zeros, and not the route field's ``arange(topk)``.""" + support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS) + + *_, packed = _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) + + first_rows = torch.stack([packed.segment(index)[0] for index in range(len(support))]) + assert torch.all(first_rows == SAMPLE_SUPPORT_PADDING) + assert not torch.any(packed.values[packed.values >= 0] == SAMPLE_SUPPORT_PADDING) + + +def test_sample_support_pooled_fill_equals_serial_fill(monkeypatch, tokenizer): + """The pooled fill writes disjoint segments, so it must be bit-identical to a serial one.""" + 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_and_sliced_arrays(tokenizer): + """Wire arrays arrive read-only, and the single-turn path hands over a row slice.""" + 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): + """The sender establishes int32, so collation validates instead of rescanning.""" + 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 + + +def test_build_sample_support_torch_dtype_matches_the_wire_dtype(): + assert torch.empty(0, dtype=SAMPLE_SUPPORT_TORCH_DTYPE).numpy().dtype == SAMPLE_SUPPORT_DTYPE + packed = build_sample_support([np.zeros((2, 3), dtype=SAMPLE_SUPPORT_DTYPE)], np.asarray([2])) + assert packed.dtype == SAMPLE_SUPPORT_TORCH_DTYPE diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py index cc6c82c5af..e2418070ac 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, @@ -54,7 +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]], - "rollout_expert_indices": [np.zeros((2, 1, 2), dtype=np.uint8), np.ones((2, 1, 2), dtype=np.uint8)], + # Routes cover every token the loss trains: one row short of the sequence, since the last + # sampled token has no subsequent forward to record a route for. + "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]]], } @@ -65,7 +68,7 @@ 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((3, 1, 2), 2, dtype=np.uint8), np.full((1, 1, 2), 3, dtype=np.uint8)], + "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]]], } @@ -592,7 +595,10 @@ def test_sample_support_observation_deltas_keep_full_width_padding_rows(self): "stop_reasons": ["continue", "eos"], "rollout_metrics": None, "rollout_logprobs": None, - "rollout_sample_support": [[[20, 21, -1]], [[40, 44, -1], [41, 45, 46]]], + "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], @@ -601,9 +607,11 @@ def test_sample_support_observation_deltas_keep_full_width_padding_rows(self): merged = merge_stepwise_output(gen_out) support = merged["rollout_sample_support"][0] - assert support == [[20, 21, -1], [-1, -1, -1], [40, 44, -1], [41, 45, 46]] - assert len(support) == len(merged["response_ids"][0]) - assert {len(row) for row in support} == {3} + 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)) + # One dense ndarray, so the trainer can pack it without a nested-list round trip. + 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") diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index ec623096a0..42f51b1835 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -8,7 +8,11 @@ import numpy as np import pytest -from skyrl.backends.skyrl_train.utils.sample_support import SampleSupportTrace +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, @@ -485,9 +489,12 @@ def generate(input_batch, model=None): ) assert prompt_starts == [0, 5] - assert output.rollout_sample_support[:2] == [[10, 100], [11, 110]] - assert output.rollout_sample_support[-2:] == [[10, 101], [11, 111]] - assert all(row == [-1, -1] for row in output.rollout_sample_support[2:-2]) + # One dense int32 block, never a nested list: the payload stays an ndarray to the trainer. + 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 @@ -757,7 +764,10 @@ def generate(input_batch, model=None): assert output.response_ids == [10, 11, 4] assert output.loss_mask == [1, 1, 1] - assert output.rollout_sample_support == [[10, 110], [11, 111], eos_support_row] + 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 @@ -807,7 +817,11 @@ def generate(input_batch, model=None): ) assert output.response_ids == [10, 11, 4] - assert output.rollout_sample_support == [[10, 110], [11, 111], [-1, -1]] + 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 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_packed_route_collation_equivalence.py b/tests/train/test_packed_route_collation_equivalence.py index 5d8c9ac679..39ff3ba552 100644 --- a/tests/train/test_packed_route_collation_equivalence.py +++ b/tests/train/test_packed_route_collation_equivalence.py @@ -180,6 +180,7 @@ def _run_both_paths( loss_mask, _logprobs, packed_routes, + _sample_support, ) = convert_prompts_responses_to_batch_tensors( PAD_TOKEN_ID, prompts, @@ -332,3 +333,23 @@ def test_packed_routes_match_under_context_parallelism(monkeypatch, parallel_sta 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..10ef8da11b 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,157 @@ 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) + + +# ============================================================ +# Per-generated-token side-channel validation tests +# ============================================================ + + +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): + """Routes cover a prefix: full coverage, ``generate_batched``'s seq_len - 1, and the deeper + shortfall a multi-turn trace leaves are all legal -- the deeper one only because the tokens + past the capture are loss-masked, which is what makes their dummy routes harmless.""" + 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(): + """The row count is bounded below by the last loss-active token: past the capture the + trainer replays dummy routes, and the router padding mask only excludes those rows from + router accounting -- it does not stop a trained token from being replayed on them.""" + 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(): + """The lower bound is exactly what ``RoutedExpertTrace`` captures, so the boundary check + cannot refuse this repo's own multi-turn producer: two turns over a 3-token prompt and a + 2-token observation capture 8 rows of a 9-token trajectory whose last token is trained. + """ + 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) + + +def test_validate_generator_output_rejects_routes_past_the_sequence(): + """More route rows than tokens means the trainer cannot place them: collation would + reject it, and the router padding mask has no position to mark them at.""" + output = _make_side_channel_output(rollout_expert_indices=[_routes(5), _routes(5)]) + + with pytest.raises(AssertionError, match=r"rollout_expert_indices\[1\] has 5 route rows for a 4-token"): + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_rejects_empty_routes(): + output = _make_side_channel_output(rollout_expert_indices=[_routes(5), _routes(0)]) + + with pytest.raises(AssertionError, match=r"rollout_expert_indices\[1\] has 0 route rows"): + validate_generator_output(num_prompts=2, generator_output=output) + + +def test_validate_generator_output_rejects_none_route_entry(): + """``None`` survives the outer length check and reaches the router-mask build as a + ``TypeError`` on ``len(None)``.""" + 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(): + """``build_sample_support`` requires exactly one row per response token.""" + 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(): + """Any step-wise generator is refused, not just ``SkyRLGymGenerator``: route rows are + aligned to the whole trajectory, so per-turn samples would replay another token's routes. + """ + 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(): + """Support is per-step dense, so step-wise keeps it -- only routes are refused.""" + 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) From ec57e6f4179069ce23147a51bf3c97e6b3877209 Mon Sep 17 00:00:00 2001 From: dyurk-lila Date: Sat, 22 Aug 2026 00:09:44 +0000 Subject: [PATCH 17/17] style: tighten packed sample-support comments and tests --- skyrl/backends/skyrl_train/training_batch.py | 13 +-- .../skyrl_train/utils/routed_experts.py | 10 +-- .../skyrl_train/utils/sample_support.py | 24 ++---- .../workers/megatron/megatron_worker.py | 4 +- .../skyrl_train/workers/worker_utils.py | 5 +- skyrl/train/dataset/preprocess.py | 13 +-- skyrl/train/generators/base.py | 11 +-- skyrl/train/generators/utils.py | 2 +- skyrl/train/utils/trainer_utils.py | 26 +----- .../test_token_based_batching_utils.py | 3 +- .../backends/skyrl_train/test_train_batch.py | 7 +- .../skyrl_train/utils/test_routed_experts.py | 9 +- .../utils/test_sample_support_row_ids.py | 45 +++------- tests/train/dataset/test_preprocess.py | 83 +------------------ .../generators/test_generator_output_utils.py | 4 +- .../generators/test_skyrl_gym_generator.py | 1 - tests/train/test_trainer_utils.py | 48 ++++------- 17 files changed, 55 insertions(+), 253 deletions(-) diff --git a/skyrl/backends/skyrl_train/training_batch.py b/skyrl/backends/skyrl_train/training_batch.py index 61294bffff..d3b920dc90 100644 --- a/skyrl/backends/skyrl_train/training_batch.py +++ b/skyrl/backends/skyrl_train/training_batch.py @@ -560,20 +560,16 @@ class TrainingOutputBatch(TensorBatch[Dict[str, torch.Tensor]]): @dataclass(frozen=True) class PackedFieldPadding: - """How one packed ``TrainingInput`` field fills the segments batch padding appends. + """Padding rule for a packed ``TrainingInput`` field. - ``dummy_row_length`` is the segment length for a synthetic batch row, which carries a - single attended token: a field indexed over every real token needs one row for it, a - field indexed over response tokens needs none. + ``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 an entry here: the three batch padding sites look its rule -# up by name, so a field without one raises instead of reaching the trainer short a segment -# (or, in `_pad_microbatch_to_size`, being skipped as a non-Tensor). +# 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. @@ -637,8 +633,7 @@ def pad_training_input_batch(unpadded_batch: TrainingInputBatch, pad_size: int) padding = TensorList([tensor[0].clone() for _ in range(pad_size)]) new_tensors[key] = TensorList.cat([tensor, padding]) elif isinstance(tensor, PackedTensor): - # Every other field copies row 0 into the padding rows, so each padded row spans as - # many tokens as row 0 and needs a segment of row 0's length. + # 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 ) diff --git a/skyrl/backends/skyrl_train/utils/routed_experts.py b/skyrl/backends/skyrl_train/utils/routed_experts.py index 1ec3048626..cef230728b 100644 --- a/skyrl/backends/skyrl_train/utils/routed_experts.py +++ b/skyrl/backends/skyrl_train/utils/routed_experts.py @@ -37,15 +37,7 @@ def record_generation( 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 routes, which cover a prefix of the sequence's real tokens. - - The capture ends short of ``token_count``: the last sampled token has no subsequent - decode forward to record its route, and a synthetic EOS is never evaluated at all. - The row count is the trace's only report of where the capture stops -- collation - dummy-fills the uncovered tail and ``make_router_padding_mask`` excludes exactly those - rows from router accounting -- so the trace must not pad the tail itself, which would - report fabricated routes as captured ones. - """ + """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: diff --git a/skyrl/backends/skyrl_train/utils/sample_support.py b/skyrl/backends/skyrl_train/utils/sample_support.py index ba5a3be155..341b7694ee 100644 --- a/skyrl/backends/skyrl_train/utils/sample_support.py +++ b/skyrl/backends/skyrl_train/utils/sample_support.py @@ -1,6 +1,7 @@ """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 @@ -20,10 +21,8 @@ SAMPLE_SUPPORT_TORCH_DTYPE = torch.int32 SAMPLE_SUPPORT_DTYPES = frozenset({SAMPLE_SUPPORT_DTYPE}) SAMPLE_SUPPORT_PADDING = -1 -# Names the support field in ``GeneratorOutput``, ``TrainingInput`` and ``Experience``. SAMPLE_SUPPORT_FIELD = "rollout_sample_support" -# Row-id channel value for a model position no support row scores. Out of range for any -# packed row index, so a gather by id cannot silently pick up a real row. +# Sentinel outside the valid packed-row range. SAMPLE_SUPPORT_NO_ROW = -1 @@ -47,18 +46,11 @@ def align_sample_support_row_ids( sample_support: PackedTensor, layout: TokenMetadataLayout, ) -> torch.Tensor: - """Return the per-token channel naming which packed support row scores each model position. - - The payload itself must not go through ``align_packed_token_metadata``: that places a - segment at a fixed offset inside its trajectory's padded region, while support scoring - happens one position to the left of the token it describes -- the logit at position ``t`` - predicts token ``t + 1``. A trajectory's support therefore covers real tokens - ``[p_i - 1, p_i + r_i - 1)``, which includes the last prompt token and excludes the last - response token. Aligning only these int64 row ids keeps that placement in one place, and - the scorer gathers ``[top_k]`` rows by id. - - Row ids index ``sample_support.values``, so they must be derived per micro-batch: - ``chunk``, ``slice`` and batch padding all rebase the packed row space. + """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): @@ -71,7 +63,7 @@ def align_sample_support_row_ids( dtype=torch.long, device=segment_lengths.device, ) - # p_i = L_i - r_i, and the position predicting the first response token is p_i - 1. + # 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( diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index a3adb1f06e..808544bc2a 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -786,9 +786,7 @@ def _pad_microbatch_to_size(self, micro_dict: dict, target_batch_size: int) -> d padded[key] = None continue if isinstance(value, PackedTensor): - # The dummy attention_mask row below marks one valid token, so a per-token field - # gets one dummy row for it while a response-token field gets none. A packed field - # with no rule raises here rather than falling through as a non-Tensor. + # 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) ) diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py index 9b466c6b9e..cf6541fefc 100644 --- a/skyrl/backends/skyrl_train/workers/worker_utils.py +++ b/skyrl/backends/skyrl_train/workers/worker_utils.py @@ -328,13 +328,12 @@ 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 the packed side channels 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) for key in PACKED_FIELD_PADDING: - # The dummy attention_mask row marks one valid token, so a per-token field gets one - # dummy row for it while a response-token field gets none. + # 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, diff --git a/skyrl/train/dataset/preprocess.py b/skyrl/train/dataset/preprocess.py index 10d785184c..19b8c20637 100644 --- a/skyrl/train/dataset/preprocess.py +++ b/skyrl/train/dataset/preprocess.py @@ -203,18 +203,7 @@ def build_sample_support( rollout_sample_support: List[SampleSupport], response_lens: np.ndarray, ) -> PackedTensor: - """Pack per-trajectory sampler support into one ``[sum(response_len_i), top_k]`` buffer. - - Support describes generated tokens only, so it packs to the response tokens rather than to - a ``[batch, seq_len, top_k]`` rectangle whose whole prompt region would be padding written - on the driver and read back only to be discarded. The outer ragged level is one segment per - trajectory, exactly as for packed routes, so the trainer indexes both by segment. - - The wire side establishes the canonical dtype and the trailing-padding invariant, so entries - are validated rather than rescanned here. The fill runs from a locally sized thread pool for - the same reason route collation does: it touches the whole global batch before DP sharding, - on every training step, and is bound by first-touch page faults on a fresh mapping. - """ + """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): diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py index e77c29d273..c9724d8805 100644 --- a/skyrl/train/generators/base.py +++ b/skyrl/train/generators/base.py @@ -54,16 +54,9 @@ 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, one ``[tokens, layers, topk]`` array of the routes the rollout took over a - # prefix of its ``prompt + response`` tokens: no decode forward follows the last sampled token, - # and a multi-turn trace ends further short of a synthetic EOS. Collation - # dummy-fills the uncovered tail and the router padding mask keeps it out of router accounting, - # so the row count is what states where the capture stops. + # Per trajectory, routes for a prefix of its prompt and response tokens. rollout_expert_indices: Optional[List[RoutedExpertIndices]] - # Per trajectory, one dense ``[response_tokens, top_k]`` array of the sampler support each - # response token was drawn from, right-padded with ``SAMPLE_SUPPORT_PADDING``. Tokens with no - # captured support (observations, a synthetic EOS) are all-padding rows. Stays an ndarray from - # the wire to the packed trainer field: nested lists of it cost ~50x the int32 buffer. + # 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]] diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py index 590e23d1ea..16d1cd2f07 100644 --- a/skyrl/train/generators/utils.py +++ b/skyrl/train/generators/utils.py @@ -844,7 +844,7 @@ 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 - # One row block per merged turn; concatenated at flush rather than extended row by row. + # 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 = [] diff --git a/skyrl/train/utils/trainer_utils.py b/skyrl/train/utils/trainer_utils.py index 0599b57bf5..86efc8d634 100644 --- a/skyrl/train/utils/trainer_utils.py +++ b/skyrl/train/utils/trainer_utils.py @@ -756,28 +756,13 @@ def validate_generator_output(num_prompts: int, generator_output: GeneratorOutpu def _validate_per_token_side_channels(generator_output: GeneratorOutput, step_wise: bool): - """Validate the per-generated-token side channels against what the trainer consumes. - - The outer length checks above cover only the per-trajectory list; a ``None`` entry or a - row count the trainer cannot place still reaches ``convert_prompts_responses_to_batch_tensors`` - (or, for routes, ``make_router_padding_mask`` one step earlier) as an opaque ``TypeError``. - - Routes cover a non-empty *prefix* of a trajectory's ``prompt + response`` tokens: vLLM - records no route for the last sampled token, and a multi-turn trace ends further short of - a synthetic EOS. Collation dummy-fills the uncovered tail and the router padding mask - excludes it, so the row count is bounded, not fixed -- bounded above by the sequence and - below by the last token the loss trains. Sample support is dense over the response, so its - row count is exact. - """ + """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"] - # Route rows are aligned to one contiguous prompt+response token sequence, while step-wise - # splits that trajectory into per-turn samples whose prompts re-cover earlier turns. The - # trajectory-aligned rows would land on the wrong tokens of every step but the first, which - # the row-count bounds below cannot see: they check coverage, not which token a row describes. + # 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 " @@ -795,12 +780,7 @@ def _validate_per_token_side_channels(generator_output: GeneratorOutput, step_wi f"rollout_expert_indices[{i}] has {captured_rows} route rows for a " f"{sequence_length}-token trajectory, expected a non-empty prefix of it" ) - # The row at source position ``t`` holds the route that produced token ``t + 1``, so - # ``captured_rows`` rows cover targets ``[1, captured_rows]``. Beyond that the trainer - # replays dummy routes, and the router padding mask only keeps those out of router - # accounting -- it cannot stop a trained token from being replayed on a route the - # rollout never took. ``RoutedExpertTrace.finalize`` proves the same bound for traces - # it builds; this is the boundary check for routes from any other producer. + # 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]) 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 b2a7e25f0a..c04b204eb1 100644 --- a/tests/backends/skyrl_train/test_token_based_batching_utils.py +++ b/tests/backends/skyrl_train/test_token_based_batching_utils.py @@ -234,8 +234,7 @@ def test_padding_microbatch_uses_unique_dummy_routes(self): assert torch.all(padding["router_padding_mask"]) def test_padding_microbatch_sample_support_holds_no_response_rows(self): - """The dummy row attends one prompt-side token and generates nothing, so its support - segment is empty -- and the route field keeps its own one-row dummy segment.""" + """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) diff --git a/tests/backends/skyrl_train/test_train_batch.py b/tests/backends/skyrl_train/test_train_batch.py index 353539d5ed..291dfba5d5 100644 --- a/tests/backends/skyrl_train/test_train_batch.py +++ b/tests/backends/skyrl_train/test_train_batch.py @@ -896,9 +896,7 @@ def test_zero_copy_falls_back_for_bfloat16(): # ── packed field padding ───────────────────────────────────────────────────── -# The row each field's padding segments must carry, stated independently of the production -# rule: distinct experts for Megatron's dropless dispatcher, "no support" for the sampler -# support. Padding a field with the other's fill would mis-train silently. +# 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), @@ -929,7 +927,6 @@ def test_packed_field_padding_carries_that_fields_own_fill(key): @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): - """All three batch padding sites append through here, so the round trip is asserted once.""" field = _ZERO_COPY_PAYLOADS[key]() padded = append_packed_field_padding(key, field, segment_lengths=[2] * pad_count) @@ -942,8 +939,6 @@ def test_appending_packed_field_padding_keeps_the_real_segments(key, pad_count): @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): - """A synthetic batch row attends one token: a per-token field needs one row for it, a - response-token field needs none.""" field = _ZERO_COPY_PAYLOADS[key]() segments = packed_dummy_row_segments(key, 3) diff --git a/tests/backends/skyrl_train/utils/test_routed_experts.py b/tests/backends/skyrl_train/utils/test_routed_experts.py index 1d7a957449..fd059eab9b 100644 --- a/tests/backends/skyrl_train/utils/test_routed_experts.py +++ b/tests/backends/skyrl_train/utils/test_routed_experts.py @@ -69,10 +69,7 @@ def _turn_routes(num_rows): def test_trace_returns_only_the_rows_it_captured(): - """The row count is what tells the trainer where the capture stops, so the trace never - fabricates rows for the tail it did not cover: collation dummy-fills that tail and the - router padding mask keeps those rows out of router accounting. - """ + """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)) @@ -93,9 +90,7 @@ def test_trace_keeps_full_coverage_when_every_token_has_a_route(): def test_trace_rejects_an_uncaptured_loss_active_target(): - """The uncovered tail is only safe because every loss-active target inside it is masked: - a forced route at a masked position can perturb nothing but later masked positions. - """ + """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)) 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 index 4f5d452905..105b938679 100644 --- a/tests/backends/skyrl_train/utils/test_sample_support_row_ids.py +++ b/tests/backends/skyrl_train/utils/test_sample_support_row_ids.py @@ -1,12 +1,4 @@ -"""Row-id derivation for packed sampler support. - -The support payload is never pushed through ``align_packed_token_metadata``: its domain is a -response suffix shifted one position left, which segment placement alone cannot express. Only -this int64 row-id channel is aligned, and the scorer gathers ``[top_k]`` rows by id. - -Run with: -uv run --isolated --extra dev --extra skyrl-train pytest tests/backends/skyrl_train/utils/test_sample_support.py -""" +"""Row-id derivation for packed sampler support.""" from typing import List, Tuple @@ -29,8 +21,7 @@ ) TOP_K = 3 -# (prompt_len, response_len) per trajectory: anti-correlated, so the support suffixes start at -# different offsets and a mis-shifted placement cannot pass by coincidence. +# Anti-correlated lengths exercise different support offsets. LENGTHS: List[Tuple[int, int]] = [(2, 3), (4, 2)] @@ -71,8 +62,7 @@ def _support(lengths: List[Tuple[int, int]]) -> PackedTensor: def test_row_ids_land_on_the_positions_that_predict_response_tokens(): - """A trajectory's support covers real tokens ``[p - 1, p + r - 1)``: it includes the last - prompt token and excludes the last response token.""" + """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 @@ -96,7 +86,7 @@ def test_row_ids_gather_the_support_rows_of_each_response_token(): @pytest.mark.parametrize("align", [1, 4]) def test_row_ids_follow_megatron_packed_padding(align): - """Under sequence packing the channel is one row of ``[seq0, pad0, seq1, pad1, ...]``.""" + """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)] @@ -106,31 +96,16 @@ def test_row_ids_follow_megatron_packed_padding(align): assert packed[padded_lengths[0] + 3 : padded_lengths[0] + 5] == [3, 4] -def test_row_ids_rebase_under_chunk(): - """``chunk`` rebases the packed row space, so ids carried in the batch would be wrong.""" - support = _support(LENGTHS) - batch = TrainingInputBatch({"attention_mask": _attention_mask(LENGTHS).long(), SAMPLE_SUPPORT_FIELD: support}) - full_batch_ids = align_sample_support_row_ids(support, _layout(LENGTHS)) - - chunk = batch.chunk(1)[1] - chunk_ids = align_sample_support_row_ids(chunk[SAMPLE_SUPPORT_FIELD], _layout(LENGTHS[1:])) - - # The chunk's own row space starts at 0 again, so its ids differ from the batch's. - assert chunk_ids[chunk_ids >= 0].tolist() == [0, 1] - assert full_batch_ids[1][full_batch_ids[1] >= 0].tolist() == [3, 4] - # Gathering with the rebased ids still lands on the same support rows. - gathered = chunk[SAMPLE_SUPPORT_FIELD].values[chunk_ids[chunk_ids >= 0]] - assert torch.equal(gathered, support.segment(1)) - - -def test_row_ids_rebase_under_slice(): +@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) - sliced = batch.slice(1, 2) - row_ids = align_sample_support_row_ids(sliced[SAMPLE_SUPPORT_FIELD], _layout(LENGTHS[1:])) + row_ids = align_sample_support_row_ids(selected[SAMPLE_SUPPORT_FIELD], _layout(LENGTHS[1:])) - assert torch.equal(sliced[SAMPLE_SUPPORT_FIELD].values[row_ids[row_ids >= 0]], support.segment(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(): diff --git a/tests/train/dataset/test_preprocess.py b/tests/train/dataset/test_preprocess.py index 556b3c6210..b4fab55172 100644 --- a/tests/train/dataset/test_preprocess.py +++ b/tests/train/dataset/test_preprocess.py @@ -10,10 +10,7 @@ import pytest import torch -from skyrl.backends.skyrl_train.utils.routed_experts import ( - ROUTED_EXPERT_DTYPES, - RoutedExpertTrace, -) +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, @@ -22,7 +19,6 @@ from skyrl.train.dataset import parallel_fill from skyrl.train.dataset.preprocess import ( ROUTED_EXPERT_TORCH_DTYPES, - build_sample_support, convert_prompts_responses_to_batch_tensors, make_router_padding_mask, ) @@ -80,43 +76,6 @@ def test_router_padding_mask_marks_left_padding_and_uncaptured_suffix(): assert mask.tolist() == [[True, False, False, True], [False, False, False, False]] -def test_router_padding_mask_marks_the_tail_a_multi_turn_trace_never_captured(): - """A trace reports only the rows it captured, so the mask covers exactly the tail that - collation dummy-fills. Two turns over a 3-token prompt and a 2-token observation capture - 8 rows of a 10-token sequence: no decode forward followed the last sampled token, and the - synthetic EOS was never evaluated at all. - """ - trace = RoutedExpertTrace() - trace.record_generation( - prompt_token_count=3, - generated_token_count=2, - routed_experts=np.zeros((4, 2, 2), dtype=np.int16), - ) - trace.record_generation( - prompt_token_count=7, - generated_token_count=2, - routed_experts=np.zeros((4, 2, 2), dtype=np.int16), - ) - # The last two tokens are loss-masked, which is what the trace's own guard enforces. - routes = trace.finalize(token_count=10, loss_mask=[0, 0, 0, 1, 1, 0, 0, 1, 1, 0]) - assert routes.shape[0] == 8 - - attention_mask = torch.tensor([[0, 0] + [1] * 10]) - - mask = make_router_padding_mask(attention_mask, [routes.shape[0]]) - - assert mask.tolist() == [[True, True] + [False] * 8 + [True, True]] - - -def test_router_padding_mask_marks_the_last_token_for_batched_routes(): - """``generate_batched`` returns ``seq_len - 1`` rows, one short of its sequence.""" - attention_mask = torch.tensor([[0, 0, 1, 1, 1, 1, 1]]) - - mask = make_router_padding_mask(attention_mask, [4]) - - assert mask.tolist() == [[True, True, False, False, False, False, True]] - - def test_routed_expert_tensor_uses_unique_dummy_routes(tokenizer): routes = [ np.asarray( @@ -627,10 +586,6 @@ def test_stepwise_anti_correlation_no_inflation(tokenizer): assert action[1].tolist() == [0] * 80 + [1] * 10 -# --------------------------------------------------------------------------- -# Sample support — packed to response tokens -# --------------------------------------------------------------------------- - 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)] @@ -676,32 +631,8 @@ def test_sample_support_packs_to_the_response_tokens(tokenizer): assert torch.equal(packed.segment(index), torch.from_numpy(rows)) -def test_sample_support_allocates_no_prompt_region_rectangle(tokenizer): - """A ``[batch, seq_len, top_k]`` rectangle would be mostly prompt-region filler.""" - support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS) - - *_, packed = _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) - - max_total = max(prompt_len + response_len for prompt_len, response_len in SAMPLE_SUPPORT_LENGTHS) - rectangle_rows = len(SAMPLE_SUPPORT_LENGTHS) * max_total - assert packed.values.shape[0] == sum(response_len for _, response_len in SAMPLE_SUPPORT_LENGTHS) - assert packed.values.shape[0] < rectangle_rows - assert packed.values.numel() == packed.values.shape[0] * SAMPLE_SUPPORT_TOP_K - - -def test_sample_support_preserves_padding_rows(tokenizer): - """Padding stays ``-1``: not zeros, and not the route field's ``arange(topk)``.""" - support = _make_sample_support(SAMPLE_SUPPORT_LENGTHS) - - *_, packed = _convert_with_support(tokenizer, SAMPLE_SUPPORT_LENGTHS, support) - - first_rows = torch.stack([packed.segment(index)[0] for index in range(len(support))]) - assert torch.all(first_rows == SAMPLE_SUPPORT_PADDING) - assert not torch.any(packed.values[packed.values >= 0] == SAMPLE_SUPPORT_PADDING) - - def test_sample_support_pooled_fill_equals_serial_fill(monkeypatch, tokenizer): - """The pooled fill writes disjoint segments, so it must be bit-identical to a serial one.""" + """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)) @@ -713,8 +644,7 @@ def test_sample_support_pooled_fill_equals_serial_fill(monkeypatch, tokenizer): assert torch.equal(pooled.values, torch.from_numpy(np.concatenate(support, axis=0))) -def test_sample_support_accepts_read_only_and_sliced_arrays(tokenizer): - """Wire arrays arrive read-only, and the single-turn path hands over a row slice.""" +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 @@ -749,7 +679,6 @@ def test_sample_support_rejects_nested_lists(tokenizer): @pytest.mark.parametrize("dtype", [np.int16, np.int64]) def test_sample_support_rejects_non_canonical_dtypes(tokenizer, dtype): - """The sender establishes int32, so collation validates instead of rescanning.""" support = [rows.astype(dtype) for rows in _make_sample_support(SAMPLE_SUPPORT_LENGTHS)] with pytest.raises(ValueError, match="canonical sample-support dtype"): @@ -772,9 +701,3 @@ def test_sample_support_none_when_not_provided(tokenizer): loss_masks=[[1]], ) assert packed is None - - -def test_build_sample_support_torch_dtype_matches_the_wire_dtype(): - assert torch.empty(0, dtype=SAMPLE_SUPPORT_TORCH_DTYPE).numpy().dtype == SAMPLE_SUPPORT_DTYPE - packed = build_sample_support([np.zeros((2, 3), dtype=SAMPLE_SUPPORT_DTYPE)], np.asarray([2])) - assert packed.dtype == SAMPLE_SUPPORT_TORCH_DTYPE diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py index e2418070ac..fd568276c5 100644 --- a/tests/train/generators/test_generator_output_utils.py +++ b/tests/train/generators/test_generator_output_utils.py @@ -55,8 +55,7 @@ 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 token the loss trains: one row short of the sequence, since the last - # sampled token has no subsequent forward to record a route for. + # 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]]], } @@ -609,7 +608,6 @@ def test_sample_support_observation_deltas_keep_full_width_padding_rows(self): 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)) - # One dense ndarray, so the trainer can pack it without a nested-list round trip. assert support.shape == (len(merged["response_ids"][0]), 3) assert support.dtype == SAMPLE_SUPPORT_DTYPE diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index 42f51b1835..1d7e6f85d5 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -489,7 +489,6 @@ def generate(input_batch, model=None): ) assert prompt_starts == [0, 5] - # One dense int32 block, never a nested list: the payload stays an ndarray to the trainer. 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)) diff --git a/tests/train/test_trainer_utils.py b/tests/train/test_trainer_utils.py index 10ef8da11b..2d4a39d260 100644 --- a/tests/train/test_trainer_utils.py +++ b/tests/train/test_trainer_utils.py @@ -1142,11 +1142,6 @@ def test_validate_stepwise_multiple_is_last_step_true_per_trajectory(): validate_generator_output(num_prompts=1, generator_output=output, step_wise=True) -# ============================================================ -# Per-generated-token side-channel validation tests -# ============================================================ - - def _make_side_channel_output( rollout_expert_indices=None, rollout_sample_support=None, @@ -1183,9 +1178,7 @@ def _support(num_rows): ], ) def test_validate_generator_output_accepts_route_under_coverage(route_rows, loss_masks): - """Routes cover a prefix: full coverage, ``generate_batched``'s seq_len - 1, and the deeper - shortfall a multi-turn trace leaves are all legal -- the deeper one only because the tokens - past the capture are loss-masked, which is what makes their dummy routes harmless.""" + """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, @@ -1195,9 +1188,7 @@ def test_validate_generator_output_accepts_route_under_coverage(route_rows, loss def test_validate_generator_output_rejects_routes_that_stop_short_of_a_trained_token(): - """The row count is bounded below by the last loss-active token: past the capture the - trainer replays dummy routes, and the router padding mask only excludes those rows from - router accounting -- it does not stop a trained token from being replayed on them.""" + """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"): @@ -1205,10 +1196,7 @@ def test_validate_generator_output_rejects_routes_that_stop_short_of_a_trained_t def test_validate_generator_output_accepts_the_coverage_a_multi_turn_trace_produces(): - """The lower bound is exactly what ``RoutedExpertTrace`` captures, so the boundary check - cannot refuse this repo's own multi-turn producer: two turns over a 3-token prompt and a - 2-token observation capture 8 rows of a 9-token trajectory whose last token is trained. - """ + """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)) @@ -1231,25 +1219,21 @@ def test_validate_generator_output_accepts_the_coverage_a_multi_turn_trace_produ validate_generator_output(num_prompts=1, generator_output=output) -def test_validate_generator_output_rejects_routes_past_the_sequence(): - """More route rows than tokens means the trainer cannot place them: collation would - reject it, and the router padding mask has no position to mark them at.""" - output = _make_side_channel_output(rollout_expert_indices=[_routes(5), _routes(5)]) - - with pytest.raises(AssertionError, match=r"rollout_expert_indices\[1\] has 5 route rows for a 4-token"): - validate_generator_output(num_prompts=2, generator_output=output) - - -def test_validate_generator_output_rejects_empty_routes(): - output = _make_side_channel_output(rollout_expert_indices=[_routes(5), _routes(0)]) +@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=r"rollout_expert_indices\[1\] has 0 route 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(): - """``None`` survives the outer length check and reaches the router-mask build as a - ``TypeError`` on ``len(None)``.""" output = _make_side_channel_output(rollout_expert_indices=[_routes(5), None]) with pytest.raises(AssertionError, match=r"rollout_expert_indices\[1\] is None"): @@ -1263,7 +1247,6 @@ def test_validate_generator_output_accepts_dense_sample_support(): def test_validate_generator_output_rejects_sample_support_row_shortfall(): - """``build_sample_support`` requires exactly one row per response token.""" 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"): @@ -1278,9 +1261,7 @@ def test_validate_generator_output_rejects_none_sample_support_entry(): def test_validate_generator_output_refuses_routes_under_step_wise(): - """Any step-wise generator is refused, not just ``SkyRLGymGenerator``: route rows are - aligned to the whole trajectory, so per-turn samples would replay another token's routes. - """ + """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"]] @@ -1289,7 +1270,6 @@ def test_validate_generator_output_refuses_routes_under_step_wise(): def test_validate_generator_output_allows_sample_support_under_step_wise(): - """Support is per-step dense, so step-wise keeps it -- only routes are refused.""" output = _make_stepwise_output(n_trajectories=1, steps_per_traj=(2,)) output["rollout_sample_support"] = [_support(len(response)) for response in output["response_ids"]]