[train] Skip building unused per-token loss_fn_outputs when the caller does not consume them

## Summary

Threads a per-request `return_per_token_outputs` flag (default `True`) carried on the already-plumbed `loss_fn_config` dict to gate the per-token `loss_fn_outputs` build on the `cross_entropy` branch of **both** backends and **both** the train and eval/forward paths. When the flag is `False`, the per-token NLL, the two detached `[mb, seq]` D2H copies (logprobs + elementwise loss), and the `.tolist()` loop are skipped; each sequence gets an empty dict instead. The `loss` / `response_length` metrics and the `loss_fn_output_type` tag are unchanged.

SkyRL's own `SFTTrainer` reads only `output.metrics` (loss / response_length), never `output.loss_fn_outputs`, so it now opts out — eliminating dead work on the SFT train + eval hot path. RL and Tinker callers pass no flag and keep the existing contract (default `True`).

## What changed

- **`worker_utils.py`**: new module-level helper `pop_return_per_token_outputs(loss_fn_config) -> (Optional[dict], bool)` that copies the dict before popping the flag (callers' dicts are never mutated) and returns the default `True` for `None`. Its docstring is the single authoritative source for the flag's contract and the why-popped-before-merge rationale.
- **Megatron** (`megatron/megatron_model_wrapper.py`): pop the flag once in `forward_backward_mini_batch` and gate the per-token build inside the shared `loss_func` (covers train + `forward_only` eval).
- **FSDP** (`workers/worker.py`): pop + gate in `_forward_backward_micro` (train) and `_forward_micro_with_loss` (eval).
- **Caller wiring** (`train/sft_trainer.py`): `train_step` (forward_backward) and `run_eval` (forward) now pass `loss_fn_config={"return_per_token_outputs": False}`.

The flag is popped from a per-call **copy** of `loss_fn_config` **before** the `AlgorithmConfig` merge — it is not an `AlgorithmConfig` field, so leaving it in would trip the key validation in `build_nested_dataclass` (reached via `from_dict_config`). The merge guard was changed from `if loss_fn_config is not None:` to `if loss_fn_config:` so an empty dict after the pop skips the merge exactly as today.

`return_per_token_outputs` is documented in the `pop_return_per_token_outputs` helper docstring and at each pop/gate site; the `loss_fn_config` docstrings on the three consuming worker methods that read the flag — `_forward_backward_micro`, `_forward_micro_with_loss` (FSDP), and `forward_backward_mini_batch` (Megatron) — also note the reserved key and its default. The per-build-site rationale comments were collapsed to a single line each (pointing at the helper docstring) to avoid hand-sync drift across the three near-identical sites.

## Numerical equivalence / safety

Byte-identical for all existing callers:

- When the flag is absent/`None`, the pop never runs and the full prior per-token build executes unchanged. Default `True` reproduces the exact prior code paths.
- The merge-guard change (`is not None` → truthy) is a strict no-op: no existing caller ever passed `{}`, a non-empty override dict is still truthy and still merges, and `OmegaConf.merge(base, {})` is itself a no-op.
- `loss`, the backward pass, and all consumed scalar metrics (`loss`, `response_length`, `lr`) are computed before/independent of the gated block, so they are identical whether per-token outputs are kept or skipped.
- `loss_fn_output_type` is a `WorkerOutput` field that always defaults to `"scalar"` (never set explicitly), so the type tag survives automatically — only the arrays become empty. The empty-dict-with-`"scalar"`-tag combination is only reachable behind the explicit opt-out whose sole caller ignores the payload.
- RL's separate (non-`cross_entropy`) `loss_fn_outputs` else-branch is untouched; the RL trainer and Tinker backend pass no flag, so their contracts hold. The Tinker public API whitelists `loss_fn_config` keys (empty allowed-key set for `cross_entropy`), so the flag cannot be injected by users; and even if present it is popped before any merge.

Note on the eval path: `run_eval` already iterates eval batches serially and reads only `output.metrics["loss"]`, so opting out there removes per-token work without changing any reported eval metric.

## Test plan

- `tests/backends/skyrl_train/workers/test_sft_loss_fn_outputs_gate.py` (new, CPU): drives the real FSDP `_forward_backward_micro` / `_forward_micro_with_loss` `cross_entropy` builds on CPU; asserts default/explicit-True populate `logprobs` + `elementwise_loss`, `False` yields empty dicts, `loss`/`response_length`/`lr` are identical across the flag, the flag is popped before the `AlgorithmConfig` merge (a real `eps_clip_low` override alongside the flag still merges without raising), and the caller dict is not mutated. Adds an RL-path test confirming the non-`cross_entropy` else-branch is ungated (logprobs still built; outputs + loss identical across the flag); that test disables `use_kl_loss`/`use_entropy_loss` explicitly (rather than relying on a default) so it isolates the gate from the KL/entropy terms. CPU run: 12 passed.
- `tests/backends/skyrl_train/workers/test_worker_utils.py` (extended): unit tests for `pop_return_per_token_outputs` — `None`→`(None, True)`, absent-flag→`True` with config preserved, explicit `False`/`True` popped leaving legitimate overrides, and no caller-dict mutation.
- `tests/train/test_sft_callbacks.py` (extended): assert `SFTTrainer.train_step` and `run_eval` pass `loss_fn="cross_entropy"` and `loss_fn_config={"return_per_token_outputs": False}` to the dispatch.
- `tests/backends/skyrl_train/gpu/gpu_ci/test_training_step.py` (extended, GPU): parametrized over FSDP + Megatron (the Megatron leg carries `@pytest.mark.megatron` like the sibling tests), DP=2; runs the real worker `forward_backward`/`forward` twice on the same dummy batch (flag default-True vs explicit-False) and asserts `loss` + `response_length` identical and `loss_fn_output_type == "scalar"` in both, with per-token outputs populated when kept vs empty when skipped.

Run locally:

```bash
uv run --isolated --extra skyrl-train --extra dev pytest tests/backends/skyrl_train/workers/test_sft_loss_fn_outputs_gate.py tests/backends/skyrl_train/workers/test_worker_utils.py tests/train/test_sft_callbacks.py
```

The CPU suite for the new/extended tests runs automatically under the standard CPU CI job (`tests/backends/skyrl_train/` + `tests/train/`); confirm that check is green before merge.

## Generality & follow-ups

Covered: both backends (Megatron `loss_func`; FSDP micro methods) and both the train (`forward_backward`) and eval/forward (`forward(loss_fn="cross_entropy")`) paths. RL and Tinker contracts preserved (default `True`).

Intentionally out of scope:
- The RL `loss_fn_outputs` build (separate else-branch, not `cross_entropy`) is untouched; the `forward(loss_fn=None)` pure-inference path is untouched.
- The JAX backend builds `loss_fn_outputs` in a jit-traced path driven by a structured `LossFnConfig` dataclass rather than the runtime `loss_fn_config` dict, so the dict-borne flag does not reach it — a possible follow-up, not a regression.
- No config schema field is added — the flag rides the runtime `loss_fn_config` dict only, preferring a function-arg/flag with a safe default over a global switch.

Doc follow-up: the caller-facing dispatch docstrings `WorkerDispatch.forward` / `WorkerDispatch.forward_backward` (the API the SFT trainer / Tinker call into) are outside this PR's touched files and so are left unchanged here; a one-line note about the reserved key could be added there in a follow-up so future callers can discover the opt-out from the dispatch layer.

## Relationship to open PRs

This gate is orthogonal/complementary to several in-flight efforts touching the same files:

- **#1513** (SFT loss-aggregation rewrite of the same FSDP `cross_entropy` branch) — note it renames the SFT status key `loss`→`sft_loss`, so merge ordering matters; only the test coupling to the literal `loss` metric key would need a touch-up if it lands first.
- **#1752** (VLM SFT on Megatron) — disjoint regions in the shared files.
- **#1534** (preserve staged `forward_backward` loss_fn_outputs across DP ranks; `worker_dispatch.py`) — no overlapping file.

The gate logic itself is composable with all of these.

Author: dyurk-lila

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

@@ -32,6 +32,7 @@
     setup_per_microbatch_replay_forward,
 )
 from skyrl.backends.skyrl_train.utils.torch_utils import masked_mean
+from skyrl.backends.skyrl_train.workers.worker_utils import pop_return_per_token_outputs
 from skyrl.train.config import TrainerConfig
 
 
@@ -286,6 +287,9 @@ def forward_backward_mini_batch(
             loss_fn: Optional loss function name (e.g., "cross_entropy", "ppo").
                      If provided, overrides the config's policy_loss_type.
             loss_fn_config: Optional config overrides for the loss function.
+                May also carry the reserved key ``return_per_token_outputs`` (bool, default
+                ``True``); set ``False`` to skip the per-token ``loss_fn_outputs`` build when
+                the caller reads only ``metrics``. See :func:`pop_return_per_token_outputs`.
             forward_only: If True, run the forward pass without backward (no gradients).
                           Useful for evaluation / loss-only inference paths (e.g., SFT
                           ``forward(loss_fn=...)`` codepath).
@@ -302,10 +306,12 @@ def forward_backward_mini_batch(
         else:
             current_loss_fn = self.policy_loss_fn
 
+        # Pop the per-request flag that gates the per-token loss_fn_outputs build.
+        loss_fn_config, return_per_token_outputs = pop_return_per_token_outputs(loss_fn_config)
+
         # Build config for loss function, applying any overrides
         loss_config = self.cfg.algorithm
-        if loss_fn_config is not None:
-
+        if loss_fn_config:
             new_loss_config = OmegaConf.merge(OmegaConf.create(asdict(loss_config)), OmegaConf.create(loss_fn_config))
             # NOTE: users can provide a custom loss config class, so we need to use the same class after applying overrides
             loss_config = type(loss_config).from_dict_config(new_loss_config)
@@ -379,41 +385,47 @@ def loss_func(logits, data):
             if resolved_loss_name == "cross_entropy":
                 loss = policy_loss
 
-                # Compute elementwise loss for Tinker API (per-token NLL)
-                with torch.no_grad():
-                    elementwise_loss = -action_log_probs
-                    if loss_mask is not None:
-                        elementwise_loss = elementwise_loss * loss_mask
-
-                # Build per-sequence loss_fn_outputs.
-                # Compute valid_lens vectorized on GPU, then move tensors to CPU
-                # exactly once before iterating in Python — avoids ~3N GPU->CPU
-                # syncs per micro-batch (item()/cpu()/tolist() inside the loop).
-                batch_size = action_log_probs.shape[0]
-                seq_len = action_log_probs.shape[1]
-                if action_mask is not None:
-                    valid_lens_t = action_mask.sum(dim=-1).long()
-                elif loss_mask is not None:
-                    valid_lens_t = loss_mask.sum(dim=-1).long()
+                # Build per-token loss_fn_outputs unless the caller opted out (see
+                # pop_return_per_token_outputs); skipping returns one empty dict per sequence.
+                if return_per_token_outputs:
+                    # Compute elementwise loss for Tinker API (per-token NLL)
+                    with torch.no_grad():
+                        elementwise_loss = -action_log_probs
+                        if loss_mask is not None:
+                            elementwise_loss = elementwise_loss * loss_mask
+
+                    # Compute valid_lens vectorized on GPU, then move tensors to CPU
+                    # exactly once before iterating in Python — avoids ~3N GPU->CPU
+                    # syncs per micro-batch (item()/cpu()/tolist() inside the loop).
+                    batch_size = action_log_probs.shape[0]
+                    seq_len = action_log_probs.shape[1]
+                    if action_mask is not None:
+                        valid_lens_t = action_mask.sum(dim=-1).long()
+                    elif loss_mask is not None:
+                        valid_lens_t = loss_mask.sum(dim=-1).long()
+                    else:
+                        valid_lens_t = torch.full(
+                            (batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long
+                        )
+
+                    # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens.
+                    action_log_probs_cpu = action_log_probs.detach().cpu()
+                    elementwise_loss_cpu = elementwise_loss.detach().cpu()
+                    valid_lens = valid_lens_t.cpu().tolist()
+
+                    loss_fn_outputs = []
+                    for i in range(batch_size):
+                        valid_len = valid_lens[i]
+                        loss_fn_outputs.append(
+                            {
+                                "logprobs": (action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else []),
+                                "elementwise_loss": (
+                                    elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else []
+                                ),
+                            }
+                        )
                 else:
-                    valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long)
-
-                # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens.
-                action_log_probs_cpu = action_log_probs.detach().cpu()
-                elementwise_loss_cpu = elementwise_loss.detach().cpu()
-                valid_lens = valid_lens_t.cpu().tolist()
-
-                loss_fn_outputs = []
-                for i in range(batch_size):
-                    valid_len = valid_lens[i]
-                    loss_fn_outputs.append(
-                        {
-                            "logprobs": (action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else []),
-                            "elementwise_loss": (
-                                elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else []
-                            ),
-                        }
-                    )
+                    loss_fn_outputs = [{} for _ in range(action_log_probs.shape[0])]
 
                 metrics = {
                     "loss": loss.item(),

skyrl/backends/skyrl_train/workers/worker.py

@@ -55,6 +55,7 @@
     BatchIterator,
     all_reduce_metrics,
     get_microbatch_iterator,
+    pop_return_per_token_outputs,
     reduce_metrics,
 )
 from skyrl.env_vars import (
@@ -819,7 +820,10 @@ def _forward_backward_micro(
             loss_fn: Optional train loss function name to use instead of config default.
                 Public Tinker aliases such as ``ppo`` should be normalized by the backend
                 before reaching the worker.
-            loss_fn_config: Optional config overrides for the resolved train loss function
+            loss_fn_config: Optional config overrides for the resolved train loss function.
+                May also carry the reserved key ``return_per_token_outputs`` (bool, default
+                ``True``); set ``False`` to skip the per-token ``loss_fn_outputs`` build when
+                the caller reads only ``metrics``. See :func:`pop_return_per_token_outputs`.
 
         Returns:
             Metrics dict for the worker's local micro batch
@@ -849,9 +853,12 @@ def _forward_backward_micro(
             # Fall back to config default
             current_loss_fn = self.policy_loss_fn
 
+        # Pop the per-request flag that gates the per-token loss_fn_outputs build.
+        loss_fn_config, return_per_token_outputs = pop_return_per_token_outputs(loss_fn_config)
+
         # Build config for loss function, applying any overrides
         loss_config = self.cfg.algorithm
-        if loss_fn_config is not None:
+        if loss_fn_config:
             # Create a copy of the config and apply overrides
             # TODO: Fix nested overrides
             from dataclasses import asdict
@@ -891,40 +898,47 @@ def _forward_backward_micro(
             loss = unscaled_loss * microbatch_weight
             self.strategy.backward(loss, self.model, self.optimizer)
 
-            # Compute elementwise loss for Tinker API (per-token NLL)
-            with torch.no_grad():
-                elementwise_loss = -action_log_probs
-                if loss_mask is not None:
-                    elementwise_loss = elementwise_loss * loss_mask
-
-            # Build per-sequence loss_fn_outputs (matches Tinker's ForwardBackwardOutput structure)
-            # Trim to actual response length per sample (Tinker expects variable-length arrays
-            # that align with the input weights, not padded to batch max).
-            # Compute valid_lens vectorized on GPU, then move tensors to CPU exactly
-            # once before iterating in Python — avoids ~3N GPU->CPU syncs per micro-batch.
-            batch_size = action_log_probs.shape[0]
-            seq_len = action_log_probs.shape[1]
-            if action_mask is not None:
-                valid_lens_t = action_mask.sum(dim=-1).long()
-            elif loss_mask is not None:
-                valid_lens_t = loss_mask.sum(dim=-1).long()
+            # Build per-token loss_fn_outputs unless the caller opted out (see
+            # pop_return_per_token_outputs); skipping returns one empty dict per sequence.
+            if return_per_token_outputs:
+                # Compute elementwise loss for Tinker API (per-token NLL)
+                with torch.no_grad():
+                    elementwise_loss = -action_log_probs
+                    if loss_mask is not None:
+                        elementwise_loss = elementwise_loss * loss_mask
+
+                # Build per-sequence loss_fn_outputs (matches Tinker's ForwardBackwardOutput structure)
+                # Trim to actual response length per sample (Tinker expects variable-length arrays
+                # that align with the input weights, not padded to batch max).
+                # Compute valid_lens vectorized on GPU, then move tensors to CPU exactly
+                # once before iterating in Python — avoids ~3N GPU->CPU syncs per micro-batch.
+                batch_size = action_log_probs.shape[0]
+                seq_len = action_log_probs.shape[1]
+                if action_mask is not None:
+                    valid_lens_t = action_mask.sum(dim=-1).long()
+                elif loss_mask is not None:
+                    valid_lens_t = loss_mask.sum(dim=-1).long()
+                else:
+                    valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long)
+
+                # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens.
+                action_log_probs_cpu = action_log_probs.detach().cpu()
+                elementwise_loss_cpu = elementwise_loss.detach().cpu()
+                valid_lens = valid_lens_t.cpu().tolist()
+
+                loss_fn_outputs = []
+                for i in range(batch_size):
+                    valid_len = valid_lens[i]
+                    loss_fn_outputs.append(
+                        {
+                            "logprobs": action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else [],
+                            "elementwise_loss": (
+                                elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else []
+                            ),
+                        }
+                    )
             else:
-                valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long)
-
-            # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens.
-            action_log_probs_cpu = action_log_probs.detach().cpu()
-            elementwise_loss_cpu = elementwise_loss.detach().cpu()
-            valid_lens = valid_lens_t.cpu().tolist()
-
-            loss_fn_outputs = []
-            for i in range(batch_size):
-                valid_len = valid_lens[i]
-                loss_fn_outputs.append(
-                    {
-                        "logprobs": action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else [],
-                        "elementwise_loss": (elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else []),
-                    }
-                )
+                loss_fn_outputs = [{} for _ in range(action_log_probs.shape[0])]
 
             status = {
                 "loss": loss.item(),
@@ -1087,6 +1101,14 @@ def _forward_micro_with_loss(
         Runs the model + loss under ``torch.no_grad()`` (no backward, no KL/entropy terms),
         and returns the same metrics shape as the SFT branch of ``_forward_backward_micro``,
         minus ``lr`` (no optimizer state involved).
+
+        Args:
+            experience: Experience object for one micro batch.
+            loss_fn: Eval loss function name (e.g., "cross_entropy").
+            loss_fn_config: Optional config overrides for the resolved loss function.
+                May also carry the reserved key ``return_per_token_outputs`` (bool, default
+                ``True``); set ``False`` to skip the per-token ``loss_fn_outputs`` build when
+                the caller reads only ``metrics``. See :func:`pop_return_per_token_outputs`.
         """
         self.model.eval()
         experience.to_device(torch.cuda.current_device())
@@ -1102,9 +1124,12 @@ def _forward_micro_with_loss(
 
         current_loss_fn = PolicyLossRegistry.get(loss_fn)
 
+        # Pop the per-request flag that gates the per-token loss_fn_outputs build.
+        loss_fn_config, return_per_token_outputs = pop_return_per_token_outputs(loss_fn_config)
+
         # Build config for loss function, applying any overrides
         loss_config = self.cfg.algorithm
-        if loss_fn_config is not None:
+        if loss_fn_config:
             from dataclasses import asdict
 
             new_loss_config = OmegaConf.merge(OmegaConf.create(asdict(loss_config)), OmegaConf.create(loss_fn_config))
@@ -1131,36 +1156,43 @@ def _forward_micro_with_loss(
                 rollout_logprobs=rollout_action_logprobs,
             )
 
-            elementwise_loss = -action_log_probs
-            if loss_mask is not None:
-                elementwise_loss = elementwise_loss * loss_mask
+            # Build per-token loss_fn_outputs unless the caller opted out (see
+            # pop_return_per_token_outputs); skipping returns one empty dict per sequence.
+            if return_per_token_outputs:
+                elementwise_loss = -action_log_probs
+                if loss_mask is not None:
+                    elementwise_loss = elementwise_loss * loss_mask
 
-            # Compute valid_lens vectorized on GPU, then move tensors to CPU
-            # exactly once before iterating in Python. Avoids ~3N GPU->CPU syncs
-            # per micro-batch (item()/cpu()/tolist() inside the per-sample loop).
-            batch_size = action_log_probs.shape[0]
-            seq_len = action_log_probs.shape[1]
-            if action_mask is not None:
-                valid_lens_t = action_mask.sum(dim=-1).long()
-            elif loss_mask is not None:
-                valid_lens_t = loss_mask.sum(dim=-1).long()
+                # Compute valid_lens vectorized on GPU, then move tensors to CPU
+                # exactly once before iterating in Python. Avoids ~3N GPU->CPU syncs
+                # per micro-batch (item()/cpu()/tolist() inside the per-sample loop).
+                batch_size = action_log_probs.shape[0]
+                seq_len = action_log_probs.shape[1]
+                if action_mask is not None:
+                    valid_lens_t = action_mask.sum(dim=-1).long()
+                elif loss_mask is not None:
+                    valid_lens_t = loss_mask.sum(dim=-1).long()
+                else:
+                    valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long)
+
+                # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens.
+                action_log_probs_cpu = action_log_probs.detach().cpu()
+                elementwise_loss_cpu = elementwise_loss.detach().cpu()
+                valid_lens = valid_lens_t.cpu().tolist()
+
+                loss_fn_outputs = []
+                for i in range(batch_size):
+                    valid_len = valid_lens[i]
+                    loss_fn_outputs.append(
+                        {
+                            "logprobs": action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else [],
+                            "elementwise_loss": (
+                                elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else []
+                            ),
+                        }
+                    )
             else:
-                valid_lens_t = torch.full((batch_size,), seq_len, device=action_log_probs.device, dtype=torch.long)
-
-            # Bulk GPU->CPU sync: one transfer for logprobs, elementwise_loss, and valid_lens.
-            action_log_probs_cpu = action_log_probs.detach().cpu()
-            elementwise_loss_cpu = elementwise_loss.detach().cpu()
-            valid_lens = valid_lens_t.cpu().tolist()
-
-            loss_fn_outputs = []
-            for i in range(batch_size):
-                valid_len = valid_lens[i]
-                loss_fn_outputs.append(
-                    {
-                        "logprobs": action_log_probs_cpu[i, -valid_len:].tolist() if valid_len > 0 else [],
-                        "elementwise_loss": (elementwise_loss_cpu[i, -valid_len:].tolist() if valid_len > 0 else []),
-                    }
-                )
+                loss_fn_outputs = [{} for _ in range(action_log_probs.shape[0])]
 
         return {
             "loss": policy_loss.item(),

skyrl/backends/skyrl_train/workers/worker_utils.py

@@ -1,5 +1,5 @@
 import math
-from typing import Dict, Iterator, List
+from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import torch
 import torch.distributed as dist
@@ -9,6 +9,38 @@
 from skyrl.train.dataset.bin_packing import make_seq_packer
 from skyrl.train.dataset.replay_buffer import Experience
 
+# Reserved ``loss_fn_config`` key that gates the per-token ``loss_fn_outputs`` build.
+RETURN_PER_TOKEN_OUTPUTS_KEY = "return_per_token_outputs"
+
+
+def pop_return_per_token_outputs(
+    loss_fn_config: Optional[Dict[str, Any]],
+) -> Tuple[Optional[Dict[str, Any]], bool]:
+    """Extract the per-request ``return_per_token_outputs`` flag from ``loss_fn_config``.
+
+    Default ``True`` keeps the existing contract (Tinker / RL consume the per-token
+    ``loss_fn_outputs``); SkyRL's own SFTTrainer sets it ``False`` since it only reads
+    ``metrics``. The flag is popped here because it is not an ``AlgorithmConfig`` field:
+    leaving it in would trip the ``AlgorithmConfig`` key validation in
+    ``build_nested_dataclass`` (reached via ``from_dict_config``).
+
+    The pop runs on a shallow copy so the caller's ``loss_fn_config`` dict is never
+    mutated (callers may reuse it across micro-batches). Returns the (possibly copied)
+    config alongside the resolved flag.
+
+    Args:
+        loss_fn_config: Optional per-call loss-function config overrides, or ``None``.
+
+    Returns:
+        ``(loss_fn_config, return_per_token_outputs)`` where ``loss_fn_config`` is a
+        fresh copy with ``return_per_token_outputs`` removed when the input was not
+        ``None``, otherwise the original ``None``.
+    """
+    if loss_fn_config is None:
+        return None, True
+    loss_fn_config = dict(loss_fn_config)
+    return loss_fn_config, loss_fn_config.pop(RETURN_PER_TOKEN_OUTPUTS_KEY, True)
+
 
 def reduce_metrics(metrics: Dict[str, List[float]], sum_loss_metrics: bool = False) -> Dict[str, float]:
     """Reduce scalar metrics from a list of entries per key with the appropriate reduction.