[Tinker API] Keep forwarded sample futures in memory

docs/content/docs/tinker/architecture.mdx

@@ -8,7 +8,7 @@ This page describes how SkyRL implements the Tinker API, including the system ar
 
 The integration is organized in three high-level layers:
 
-1. **API Layer** (`skyrl.tinker.api`) - FastAPI HTTP server that accepts Tinker API requests, stores them in a database, and returns future IDs for async polling
+1. **API Layer** (`skyrl.tinker.api`) - FastAPI HTTP server that accepts Tinker API requests and returns future IDs for async polling. Engine-owned operations use the database; API-forwarded samples use in-memory futures.
 2. **Engine Layer** (`skyrl.tinker.engine`) - Background subprocess that polls the database, batches pending requests, and dispatches them to the backend
 3. **Backend Layer** (`skyrl.backends`) - Translates Tinker operations into training and inference calls, managing Ray workers, FSDP/Megatron training, and vLLM inference
 
@@ -71,6 +71,9 @@ Sampling requests (`sample`) go through the following lifecycle:
 Client calls sample()
     │
     ▼
+API Server (FastAPI)
+    │  Creates an in-memory future
+    ▼
 SkyRL-Train Backend
     │  Converts Tinker SamplingParams → vLLM params
     ▼
@@ -83,10 +86,10 @@ RemoteInferenceClient
 Inference Workers (vLLM)
     │  Generate tokens with logprobs
     ▼
-Results aggregated → GeneratedSequence objects returned
+Results resolve the in-memory future → GeneratedSequence objects returned
 ```
 
-The `sample()` call is a fairly lightweight wrapper around SkyRL-Train's vLLM inference engines. The backend translates Tinker `SamplingParams` to vLLM format (e.g., `stop_strings` → `stop`, `stop_tokens` → `stop_token_ids`) and delegates prompts to the `RemoteInferenceClient`. `RemoteInferenceClient` forwards the requests to an instance of [vLLM Router](https://github.com/vllm-project/router), which handles load balancing and sticky routing across vLLM workers.
+The `sample()` call is a fairly lightweight wrapper around SkyRL-Train's vLLM inference engines. Non-colocated and external inference requests bypass the engine database: the API owns their futures and signals completion with `asyncio.Event`. The backend translates Tinker `SamplingParams` to vLLM format (e.g., `stop_strings` → `stop`, `stop_tokens` → `stop_token_ids`) and delegates prompts to the `RemoteInferenceClient`. `RemoteInferenceClient` forwards the requests to an instance of [vLLM Router](https://github.com/vllm-project/router), which handles load balancing and sticky routing across vLLM workers.
 
 
 ## Weight Sync

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