[Tinker API] Cache validated sampling targets

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -1254,11 +1254,13 @@ def forward_backward(
             torch.cuda.empty_cache()
 
         # Aggregate metrics across micro-batches
-        all_loss_fn_outputs = []  # Handle separately from scalar metrics
+        loss_fn_output_batches = []
         for m_batch, metrics in zip(micro_buffer, metrics_list):
             # Extract loss_fn_outputs before reduce_metrics (it's not a scalar metric)
-            if "loss_fn_outputs" in metrics:
-                all_loss_fn_outputs.extend(metrics.pop("loss_fn_outputs"))
+            if metrics is None:
+                loss_fn_output_batches.append([])
+                continue
+            loss_fn_output_batches.append(metrics.pop("loss_fn_outputs", []))
             # Skip fully-padding microbatches: their metrics (clip_ratio=0, policy_entropy=0,
             # ...) are meaningless and would drag down the mean-reduced metrics. Summed
             # metrics (e.g. policy_loss) are unaffected since padding contributes 0, but
@@ -1303,6 +1305,13 @@ def forward_backward(
                 for k, v in moe_metrics.items():
                     status[k] = v
 
+        if not any(loss_fn_output_batches):
+            all_loss_fn_outputs = []
+        elif isinstance(microbatch_iterator, TokenBasedBatchIterator):
+            all_loss_fn_outputs = microbatch_iterator.reorder_and_combine_items(loss_fn_output_batches)
+        else:
+            all_loss_fn_outputs = [item for batch in loss_fn_output_batches for item in batch]
+
         return WorkerOutput(loss_fn_outputs=all_loss_fn_outputs, metrics=status)
 
     def optim_step(self) -> Optional[float]:

skyrl/backends/skyrl_train/workers/worker_utils.py

@@ -424,6 +424,18 @@ def reorder_and_combine_batches(self, batches: List[TensorBatch]) -> TensorBatch
         reordered_batch.metadata = ref_microbatch.metadata
         return reordered_batch
 
+    def reorder_and_combine_items(self, batches: List[List[dict]]) -> List[dict]:
+        """Restore per-sample microbatch outputs to input order."""
+        ordered = [None] * self.data.batch_size
+        for original_indices, items in zip(self._microbatches, batches):
+            if len(items) < len(original_indices):
+                raise ValueError("Microbatch output has fewer items than input samples")
+            for original_idx, item in zip(original_indices, items):
+                ordered[original_idx] = item
+        if any(item is None for item in ordered):
+            raise ValueError("Microbatch outputs do not cover every input sample")
+        return ordered
+
 
 def get_microbatch_iterator(
     data: TrainingInputBatch, micro_batch_size: int, max_tokens_per_microbatch: int

skyrl/backends/skyrl_train_backend.py

@@ -808,6 +808,21 @@ def _normalize_policy_loss_request(
         if role == "critic":
             return loss_fn, loss_fn_config
 
+        if loss_fn == "gspo":
+            normalized_config = dict(loss_fn_config or {})
+            clip_low_threshold = normalized_config.pop("clip_low_threshold", None)
+            clip_high_threshold = normalized_config.pop("clip_high_threshold", None)
+            if clip_low_threshold is None or clip_high_threshold is None:
+                raise ValueError("loss_fn='gspo' requires clip_low_threshold and clip_high_threshold")
+            normalized_config.update(
+                eps_clip_low=1.0 - clip_low_threshold,
+                eps_clip_high=clip_high_threshold - 1.0,
+                loss_reduction="sequence_mean",
+                use_entropy_loss=False,
+                use_kl_loss=False,
+            )
+            return loss_fn, normalized_config
+
         if loss_fn == "dppo":
             # DPPO thresholds live in the nested `algorithm.dppo` sub-config, but
             # Tinker's loss_fn_config is a flat float dict. Re-nest so the
@@ -832,6 +847,47 @@ def _normalize_policy_loss_request(
             normalized_config["eps_clip_high"] = clip_high_threshold - 1.0
         return ("regular" if loss_fn == "ppo" else "gspo"), normalized_config or None
 
+    @staticmethod
+    def _normalize_gspo_batch(batch: TrainingInputBatch) -> None:
+        """Apply Tinker's sequence-mean GSPO reduction before DP sharding."""
+        advantages = batch["advantages"]
+        loss_mask = advantages.ne(0).to(batch["loss_mask"].dtype)
+        token_counts = loss_mask.sum(dim=-1, keepdim=True)
+        trainable_sequences = token_counts.squeeze(-1).gt(0)
+        num_trainable_sequences = trainable_sequences.sum().clamp(min=1)
+        batch["advantages"] = advantages / token_counts.clamp(min=1) / num_trainable_sequences
+        batch["loss_mask"] = loss_mask
+
+    @staticmethod
+    def _select_trainable_gspo_rows(batch: TrainingInputBatch) -> tuple[TrainingInputBatch, list[int]]:
+        """Drop inactive rows while preserving all-zero-batch optimizer semantics."""
+        trainable_rows = batch["loss_mask"].ne(0).any(dim=-1)
+        trainable_indices = trainable_rows.nonzero(as_tuple=False).flatten()
+        if trainable_indices.numel() == 0 or trainable_indices.numel() == batch.batch_size:
+            return batch, list(range(batch.batch_size))
+
+        selected = TrainingInputBatch(
+            {key: None if value is None else value[trainable_indices] for key, value in batch.items()}
+        )
+        selected.metadata = batch.metadata
+        return selected, trainable_indices.tolist()
+
+    @staticmethod
+    def _get_gspo_pruning_metrics(
+        original_token_counts: list[int], submitted_indices: set[int], start_idx: int, end_idx: int
+    ) -> dict[str, float]:
+        """Report one SDK chunk's pruning counts so combined futures sum to the batch total."""
+        request_indices = range(start_idx, end_idx)
+        request_submitted_indices = [index for index in request_indices if index in submitted_indices]
+        submitted_tokens = sum(original_token_counts[index] for index in request_submitted_indices)
+        original_tokens = sum(original_token_counts[index] for index in request_indices)
+        return {
+            "skyrl.ai/submitted_datums:sum": float(len(request_submitted_indices)),
+            "skyrl.ai/skipped_datums:sum": float(end_idx - start_idx - len(request_submitted_indices)),
+            "skyrl.ai/submitted_tokens:sum": float(submitted_tokens),
+            "skyrl.ai/skipped_tokens:sum": float(original_tokens - submitted_tokens),
+        }
+
     def forward_backward(
         self,
         prepared_batch: types.PreparedModelPassBatch,
@@ -860,6 +916,12 @@ def _forward_backward_single_model_batch(
         ):
             raise ValueError("Critic forward_backward requires values and returns for every response token")
         batch = self._to_training_batch(prepared_batch, role)
+        original_batch_size = batch.batch_size
+        original_token_counts = [int(count) for count in batch["attention_mask"].sum(dim=-1).tolist()]
+        submitted_indices = list(range(original_batch_size))
+        if loss_fn == "gspo":
+            self._normalize_gspo_batch(batch)
+            batch, submitted_indices = self._select_trainable_gspo_rows(batch)
         micro_bs = (
             self._cfg.trainer.micro_train_batch_size_per_gpu if self._cfg.trainer.strategy == "megatron" else None
         )
@@ -897,7 +959,14 @@ def _forward_backward_single_model_batch(
         if pad_size > 0 and per_sample_outputs:
             per_sample_outputs = per_sample_outputs[:-pad_size]
 
+        if per_sample_outputs and len(submitted_indices) != original_batch_size:
+            restored_outputs = [{"logprobs": list(logprobs)} for logprobs in prepared_batch.all_sampling_logprobs]
+            for index, output in zip(submitted_indices, per_sample_outputs, strict=True):
+                restored_outputs[index] = output
+            per_sample_outputs = restored_outputs
+
         metrics = self._extract_metrics(data.metrics)
+        submitted_index_set = set(submitted_indices)
 
         results = {}
         for request_id, _, start_idx, end_idx in prepared_batch.request_batch_slices:
@@ -917,10 +986,15 @@ def _forward_backward_single_model_batch(
                     loss_fn_outputs.append(formatted_output)
             else:
                 loss_fn_outputs = [{} for _ in range(end_idx - start_idx)]
+            request_metrics = dict(metrics)
+            if loss_fn == "gspo":
+                request_metrics.update(
+                    self._get_gspo_pruning_metrics(original_token_counts, submitted_index_set, start_idx, end_idx)
+                )
             results[request_id] = types.ForwardBackwardOutput(
                 loss_fn_output_type=data.loss_fn_output_type,
                 loss_fn_outputs=loss_fn_outputs,
-                metrics=metrics,
+                metrics=request_metrics,
             )
         return results
 

skyrl/tinker/api.py

@@ -7,6 +7,7 @@
 import signal
 import threading
 from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from typing import Annotated, Any, AsyncGenerator, ClassVar, Literal
 from uuid import uuid4
@@ -78,6 +79,13 @@
 TERMINAL_STATUSES = (RequestStatus.COMPLETED, RequestStatus.FAILED)
 
 
+@dataclass(frozen=True)
+class SamplingTarget:
+    base_model: str | None
+    model_id: str
+    checkpoint_id: str
+
+
 def raw_json_response(payload: str | None) -> Response:
     """Return already-serialized JSON without routing it through FastAPI's encoder.
 
@@ -260,6 +268,8 @@ async def lifespan(app: FastAPI):
     # SkyRL-Train default is colocate_all=True; only opt into forwarding
     # when the operator explicitly sets it to False.
     is_colocated = bool(backend_cfg.get("trainer.placement.colocate_all", True))
+    app.state.external_sampling_targets = {}
+    app.state.external_sampling_target_locks = {}
     if app.state.engine_config.external_inference_url:
         app.state.external_inference_client = ExternalInferenceClient(app.state.engine_config, app.state.db_engine)
         logger.info(f"External engine configured: {app.state.engine_config.external_inference_url}")
@@ -625,6 +635,8 @@ class ForwardBackwardInput(BaseModel):
     def validate_loss_fn_config_keys(self):
         """Validate loss_fn_config keys based on the selected loss function."""
         if self.loss_fn_config is None:
+            if self.loss_fn == "gspo":
+                raise ValueError("loss_fn='gspo' requires clip_low_threshold and clip_high_threshold.")
             return self
 
         allowed_keys = self._ALLOWED_KEYS_BY_LOSS_FN[self.loss_fn]
@@ -638,6 +650,18 @@ def validate_loss_fn_config_keys(self):
             raise ValueError(
                 f"loss_fn='{self.loss_fn}' does not accept loss_fn_config keys. " f"Received: {invalid_keys}."
             )
+        if self.loss_fn == "gspo":
+            required_keys = {"clip_low_threshold", "clip_high_threshold"}
+            missing_keys = sorted(required_keys - self.loss_fn_config.keys())
+            if missing_keys:
+                raise ValueError(f"loss_fn='gspo' is missing required loss_fn_config keys: {missing_keys}.")
+            clip_low = self.loss_fn_config["clip_low_threshold"]
+            clip_high = self.loss_fn_config["clip_high_threshold"]
+            if not 0.0 <= clip_low <= 1.0 <= clip_high:
+                raise ValueError(
+                    "loss_fn='gspo' requires 0 <= clip_low_threshold <= 1 <= clip_high_threshold; "
+                    f"got {clip_low=} and {clip_high=}."
+                )
         return self
 
     def to_types(self) -> types.ForwardBackwardInput:
@@ -771,12 +795,14 @@ def validate_model_source(self):
 class SaveWeightsRequest(BaseModel):
     model_id: str
     path: str = Field(..., pattern=ID_PATTERN, max_length=ID_MAX_LENGTH)
+    seq_id: int | None = None
     type: Literal["save_weights"] | None = None
 
 
 class LoadWeightsRequest(BaseModel):
     model_id: str
     path: str
+    seq_id: int | None = None
     type: Literal["load_weights"] | None = None
 
 
@@ -859,6 +885,7 @@ class SupportedModel(BaseModel):
 
 class GetServerCapabilitiesResponse(BaseModel):
     supported_models: list[SupportedModel]
+    supported_loss_fns: list[str]
 
 
 class ListCheckpointsResponse(BaseModel):
@@ -1185,6 +1212,7 @@ async def load_weights(request: LoadWeightsRequest, req: Request, session: Async
         request_type=types.RequestType.LOAD_WEIGHTS,
         model_id=request.model_id,
         request_data=types.LoadWeightsInput(source_model_id=source_model_id, checkpoint_id=checkpoint_id),
+        seq_id=request.seq_id,
     )
 
     await session.commit()
@@ -1208,6 +1236,7 @@ async def save_weights(request: SaveWeightsRequest, session: AsyncSession = Depe
         request_type=types.RequestType.SAVE_WEIGHTS,
         model_id=request.model_id,
         request_data=types.SaveWeightsInput(path=request.path),
+        seq_id=request.seq_id,
     )
 
     await session.commit()
@@ -1253,14 +1282,15 @@ async def save_weights_for_sampler(request: SaveWeightsForSamplerRequest, sessio
             seq_id=request.seq_id,
             sampling_session_id=sampling_session_id,
         ),
+        seq_id=request.seq_id,
     )
 
     await session.commit()
 
     return FutureResponse(future_id=str(request_id), status="pending", request_id=str(request_id))
 
 
-async def get_sampling_model(request: SampleRequest, session: AsyncSession) -> (str | None, str | None):
+async def get_sampling_model(request: SampleRequest, session: AsyncSession) -> tuple[str | None, str | None]:
     """Return (base_model, model_path) for a sampling request."""
     # Resolve model/base from sampling_session_id if provided
     if request.sampling_session_id is not None:
@@ -1271,25 +1301,18 @@ async def get_sampling_model(request: SampleRequest, session: AsyncSession) -> (
     return (request.base_model, request.model_path)
 
 
-@app.post("/api/v1/asample", response_model=FutureResponse)
-async def asample(request: SampleRequest, req: Request, session: AsyncSession = Depends(get_session)):
-    """Generates samples from the model (async version)."""
-    if request.sampling_session_id is not None and ":" in request.sampling_session_id:
-        raise HTTPException(
-            status_code=400,
-            detail="sampling_session_id must not contain ':' (the routing-key delimiter)",
-        )
+async def get_sampling_target(request: SampleRequest, req: Request, session: AsyncSession) -> SamplingTarget:
+    """Resolve and validate the model used by a sample request."""
 
-    base_model, model_path = await get_sampling_model(request, session)
+    async def resolve() -> SamplingTarget:
+        base_model, model_path = await get_sampling_model(request, session)
+        if base_model:
+            return SamplingTarget(base_model=base_model, model_id="", checkpoint_id="")
 
-    if base_model:
-        model_id = checkpoint_id = ""
-    else:
         assert model_path is not None
         path = types.TinkerPath.parse(model_path)
         if (
             not path
-            # Accept either tinker://model_id/checkpoint_id or tinker://model_id/sampler_weights/checkpoint_id
             or path.kind not in ("", "sampler_weights")
             or not (model_id := path.primary_id)
             or not (checkpoint_id := path.secondary_id)
@@ -1299,21 +1322,51 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession =
                 detail="model_path must be tinker://model_id/checkpoint_id or tinker://model_id/sampler_weights/checkpoint_id",
             )
         await get_model(session, model_id)
-        # Validate that the checkpoint exists and is ready
         await validate_checkpoint(req, model_id, checkpoint_id, types.CheckpointType.SAMPLER, session)
+        return SamplingTarget(base_model=None, model_id=model_id, checkpoint_id=checkpoint_id)
+
+    sampling_session_id = request.sampling_session_id
+    if sampling_session_id is None or not req.app.state.external_inference_client:
+        return await resolve()
+
+    targets: dict[str, SamplingTarget] = req.app.state.external_sampling_targets
+    target = targets.get(sampling_session_id)
+    if target is not None:
+        return target
+
+    locks: dict[str, asyncio.Lock] = req.app.state.external_sampling_target_locks
+    lock = locks.setdefault(sampling_session_id, asyncio.Lock())
+    async with lock:
+        target = targets.get(sampling_session_id)
+        if target is None:
+            target = await resolve()
+            targets[sampling_session_id] = target
+        return target
+
+
+@app.post("/api/v1/asample", response_model=FutureResponse)
+async def asample(request: SampleRequest, req: Request, session: AsyncSession = Depends(get_session)):
+    """Generates samples from the model (async version)."""
+    if request.sampling_session_id is not None and ":" in request.sampling_session_id:
+        raise HTTPException(
+            status_code=400,
+            detail="sampling_session_id must not contain ':' (the routing-key delimiter)",
+        )
+
+    target = await get_sampling_target(request, req, session)
 
     request_id = await create_future(
         session=session,
         request_type=(
             types.RequestType.EXTERNAL if req.app.state.external_inference_client else types.RequestType.SAMPLE
         ),
-        model_id=model_id,
+        model_id=target.model_id,
         request_data=types.SampleInput(
-            base_model=base_model,
+            base_model=target.base_model,
             prompt=request.prompt.to_types(),
             sampling_params=request.sampling_params.to_types(),
             num_samples=request.num_samples,
-            checkpoint_id=checkpoint_id,
+            checkpoint_id=target.checkpoint_id,
             # A positive topk implies prompt logprobs: both are read off the same
             # prompt forward pass, so asking for one asks for the other.
             prompt_logprobs=bool(request.prompt_logprobs) or request.topk_prompt_logprobs > 0,
@@ -1328,10 +1381,13 @@ async def asample(request: SampleRequest, req: Request, session: AsyncSession =
     if req.app.state.external_inference_client:
         asyncio.create_task(
             req.app.state.external_inference_client.call_and_store_result(
-                request_id, request, model_id, checkpoint_id, base_model=base_model
+                request_id,
+                request,
+                target.model_id,
+                target.checkpoint_id,
+                base_model=target.base_model,
             )
         )
-
     return FutureResponse(future_id=str(request_id), status="pending", request_id=str(request_id))
 
 
@@ -1341,7 +1397,10 @@ async def get_server_capabilities(request: Request):
     supported_models = [
         SupportedModel(model_name=request.app.state.engine_config.base_model),
     ]
-    return GetServerCapabilitiesResponse(supported_models=supported_models)
+    return GetServerCapabilitiesResponse(
+        supported_models=supported_models,
+        supported_loss_fns=sorted(types.SUPPORTED_LOSS_FNS),
+    )
 
 
 class RetrieveFutureRequest(BaseModel):