feat(megatron): add bounded sampler-support replay

pyproject.toml

@@ -92,6 +92,8 @@ skyrl-train = [
     "polars",
     "s3fs",
     "fastapi",
+    "orjson>=3.11.9",
+    "pybase64>=1.4.2",
     "uvicorn",
     "vllm-router; sys_platform == 'linux'",
     "pybind11",

skyrl/backends/skyrl_train/inference_servers/base.py

@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Hashable, List, Optional, Tuple, TypedDict
 
+from skyrl.utils.routed_experts import RoutedExpertIndices
+
 if TYPE_CHECKING:
     from skyrl.backends.skyrl_train.weight_sync import WeightUpdateRequest
     from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import (
@@ -32,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):
@@ -47,7 +50,8 @@ class InferenceEngineOutput(TypedDict):
     stop_reasons: List[str]
     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[List[List[int]]]]  # [seq_len, layer_num, topk]
+    rollout_expert_indices: Optional[List[RoutedExpertIndices]]
+    rollout_sample_support: Optional[List[List[List[int]]]]
 
 
 class InferenceEngineInterface(ABC):

skyrl/backends/skyrl_train/inference_servers/routed_experts_wire.py

@@ -0,0 +1,45 @@
+"""Compact routed-expert HTTP payloads."""
+
+import math
+from typing import Any
+
+import numpy as np
+import pybase64
+
+from skyrl.utils.routed_experts import (
+    ROUTED_EXPERT_DTYPES,
+    RoutedExpertIndices,
+    compact_routed_expert_indices,
+)
+
+_DTYPES = {dtype.name: dtype for dtype in ROUTED_EXPERT_DTYPES}
+
+
+def pack_routed_experts(routed_experts: RoutedExpertIndices) -> dict[str, Any]:
+    compact = compact_routed_expert_indices(routed_experts)
+    return {
+        "data": pybase64.b64encode(memoryview(compact)).decode("ascii"),
+        "shape": list(compact.shape),
+        "dtype": compact.dtype.name,
+    }
+
+
+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")
+    try:
+        dtype = _DTYPES[payload["dtype"]]
+        shape = tuple(payload["shape"])
+        data = pybase64.b64decode_as_bytearray(payload["data"], validate=True)
+    except (KeyError, TypeError, ValueError) as exc:
+        raise ValueError("invalid packed routed_experts payload") from exc
+    if len(shape) != 3 or any(type(dim) is not int or dim < 0 for dim in shape):
+        raise ValueError(f"invalid packed routed_experts shape: {shape}")
+    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)
+    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}")
+    return compact

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,

skyrl/backends/skyrl_train/inference_servers/utils.py

@@ -84,6 +84,7 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace:
     args: Namespace = parser.parse_args(args=[])
 
     ie_cfg = cfg.generator.inference_engine
+    sample_support_top_k = cfg.generator.sampling_params.top_k
     overrides = dict(
         model=cfg.trainer.policy.model.path,
         tensor_parallel_size=ie_cfg.tensor_parallel_size,
@@ -118,6 +119,9 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace:
         # Overridable via generator.inference_engine.engine_init_kwargs.trust_remote_code below.
         trust_remote_code=True,
     )
+    if ie_cfg.enable_return_sample_support_set:
+        overrides["max_logprobs"] = sample_support_top_k
+        overrides["logprobs_mode"] = "processed_logprobs"
     for key, value in overrides.items():
         setattr(args, key, value)
 

skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py

@@ -4,15 +4,18 @@
 
 import asyncio
 import logging
+import math
 import os
 import time
 from argparse import Namespace
 from typing import List, Optional, Tuple
 
 import httpx
+import numpy as np
+import orjson
 import uvicorn
 import vllm.envs as envs
-from fastapi import HTTPException, Request
+from fastapi import HTTPException, Request, Response
 from ray.util.placement_group import PlacementGroup
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.engine.async_llm_engine import AsyncLLMEngine
@@ -22,6 +25,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
@@ -34,6 +38,9 @@
     get_node_ip,
 )
 from skyrl.backends.skyrl_train.inference_servers.protocols import ServerActorProtocol
+from skyrl.backends.skyrl_train.inference_servers.routed_experts_wire import (
+    pack_routed_experts,
+)
 from skyrl.env_vars import (
     SKYRL_HTTP_CONNECTION_LIMIT,
     SKYRL_VLLM_DP_PORT_OFFSET,
@@ -43,6 +50,53 @@
 logger = logging.getLogger(__name__)
 
 
+def _sample_support_from_flat_logprobs(
+    logprobs: FlatLogprobs,
+    top_k: int,
+) -> tuple[list[dict[str, float]], list[list[int]]]:
+    """Extract sampled scores and post-filter support from vLLM's flat rows."""
+    # vLLM emits [sampled token, top-1, ..., top-k] for every generated token.
+    row_width = top_k + 1
+    token_ids = np.asarray(logprobs.token_ids, dtype=np.int64).reshape(-1, row_width)
+    processed_logprobs = np.asarray(logprobs.logprobs).reshape(-1, row_width)
+    support_ids = np.where(np.isneginf(processed_logprobs[:, 1:]), -1, token_ids[:, 1:])
+    sampled_logprobs = [{"logprob": value} for value in processed_logprobs[:, 0].tolist()]
+
+    # Repair rows whose sampled token is absent from the captured support.
+    #
+    # The support row is columns 1: (the top_k neighbors); column 0 is the token
+    # vLLM actually sampled. vLLM's approximate Triton top-k/top-p pivot (in
+    # processed_logprobs mode) can leave slightly more than top_k survivors, so the
+    # sampled token can rank just beyond top_k and be absent from cols 1:. When that
+    # happens the downstream sample-support invariant
+    # (assert_sampled_tokens_in_sample_support_sets / build_sample_support_replay)
+    # hard-crashes the whole run. To preserve it we overwrite the row's WEAKEST valid
+    # member (vLLM returns top-k descending, so the trailing valid slot is weakest)
+    # with the sampled id. Overwriting keeps width == top_k, inserts the sampled id
+    # exactly once (so the trainer's renorm denominator is not double-counted), and
+    # leaves any trailing -1 padding intact (the weakest valid col precedes the pad).
+    sampled = token_ids[:, 0]
+    valid = support_ids >= 0
+    present = np.any(support_ids == sampled[:, None], axis=1)
+    # Only repair rows that already have >=1 valid member (matches the downstream
+    # has_support condition); a fully -1 row is left untouched.
+    missing = (~present) & valid.any(axis=1)
+    if np.any(missing):
+        rows = np.flatnonzero(missing)
+        weakest_col = valid.sum(axis=1) - 1  # last valid (weakest) column per row
+        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.tolist()
+
+
 class VLLMServerActor(ServerActorProtocol):
     """
     Ray actor that runs a vLLM OpenAI-compatible API server.
@@ -399,7 +453,10 @@ async def _skyrl_generate(request: Request):
             token_ids = body["token_ids"]
             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:
+                sampling_params_dict["flat_logprobs"] = True
+                sampling_params_dict["logprobs"] = sampling_params_dict["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:
@@ -420,33 +477,42 @@ 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, sample_support = _sample_support_from_flat_logprobs(
+                    resp.logprobs,
+                    sampling_params_dict["top_k"],
+                )
+                logprobs = {"content": content}
+            elif resp.logprobs is not None:
                 content = []
                 for tid, lp_dict in zip(token_ids_out, resp.logprobs):
                     if lp_dict and tid in lp_dict:
-                        content.append({"logprob": lp_dict[tid].logprob})
+                        logprob = lp_dict[tid].logprob
+                        if not math.isfinite(logprob):
+                            raise ValueError("Out of range float values are not JSON compliant")
+                        content.append({"logprob": logprob})
                     else:
                         # -9999.0 is the default in vLLM's ChatCompletionLogProb
                         content.append({"logprob": -9999.0})
                 logprobs = {"content": content}
 
             routed_experts = None
             if resp.routed_experts is not None:
-                if hasattr(resp.routed_experts, "tolist"):
-                    routed_experts = resp.routed_experts.tolist()
-                else:
-                    routed_experts = resp.routed_experts
+                routed_experts = pack_routed_experts(np.asarray(resp.routed_experts))
 
-            return {
+            payload = {
                 "choices": [
                     {
                         "token_ids": token_ids_out,
                         "finish_reason": finish_reason,
                         "logprobs": logprobs,
                         "routed_experts": routed_experts,
+                        "rollout_sample_support": sample_support,
                     }
                 ]
             }
+            return Response(content=orjson.dumps(payload), media_type="application/json")
 
     async def shutdown(self) -> None:
         """Gracefully shutdown the server."""

skyrl/backends/skyrl_train/training_batch.py

@@ -7,7 +7,9 @@
 
 import numpy as np
 import torch
-from jaxtyping import Float, Integer
+from jaxtyping import Bool, Float, Integer
+
+from skyrl.utils.routed_experts import make_replay_padding_indices
 
 DictType = TypeVar("DictType")
 
@@ -476,6 +478,8 @@ class TrainingInput(TypedDict, total=False):
     rewards: Optional[Float[torch.Tensor, "batch_size seq_len"]]
     rollout_logprobs: Optional[Float[torch.Tensor, "batch_size seq_len"]]
     rollout_expert_indices: Optional[Integer[torch.Tensor, "batch_size seq_len layer_num topk"]]
+    router_padding_mask: Optional[Bool[torch.Tensor, "batch_size seq_len"]]
+    sample_support_ids: Optional[Integer[torch.Tensor, "batch_size seq_len topk"]]
     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
 
@@ -524,6 +528,27 @@ def pad_training_input_batch(unpadded_batch: TrainingInputBatch, pad_size: int)
             additional_dims = tensor.shape[1:]
             padding_tensor = torch.zeros(pad_size, *additional_dims, dtype=tensor.dtype, device=tensor.device)
             new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
+        elif key == "rollout_expert_indices":
+            additional_dims = tensor.shape[1:]
+            padding_tensor = make_replay_padding_indices(
+                (pad_size, *additional_dims),
+                dtype=tensor.dtype,
+                device=tensor.device,
+            )
+            new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
+        elif key == "router_padding_mask":
+            additional_dims = tensor.shape[1:]
+            padding_tensor = torch.ones(pad_size, *additional_dims, dtype=torch.bool, device=tensor.device)
+            new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
+        elif key == "sample_support_ids":
+            additional_dims = tensor.shape[1:]
+            padding_tensor = torch.full(
+                (pad_size, *additional_dims),
+                -1,
+                dtype=tensor.dtype,
+                device=tensor.device,
+            )
+            new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
         else:
             # Copy row 0 `pad_size` times. Loss masked so values don't affect the loss. Just need valid shape/dtype.
             assert tensor.shape[0] > 0, f"Cannot pad empty tensor field {key!r}"