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

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

@@ -40,6 +40,7 @@
 from skyrl.backends.skyrl_train.utils.torch_utils import masked_mean
 from skyrl.backends.skyrl_train.workers.worker_utils import (
     compute_minibatch_rollout_logprob_diff_metrics,
+    pop_return_per_token_outputs,
 )
 from skyrl.train.config import TrainerConfig
 
@@ -422,6 +423,8 @@ 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 include reserved key ``return_per_token_outputs`` to skip
+                per-token ``loss_fn_outputs`` when callers read only ``metrics``.
             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).
@@ -440,10 +443,12 @@ def forward_backward_mini_batch(
         else:
             current_loss_fn = self.policy_loss_fn
 
+        # Consume the reserved gate before merging AlgorithmConfig overrides.
+        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)
@@ -564,41 +569,43 @@ 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()
+                # Only build per-token outputs for callers that consume them.
+                if return_per_token_outputs:
+                    # Tinker consumes per-token NLL.
+                    with torch.no_grad():
+                        elementwise_loss = -action_log_probs
+                        if loss_mask is not None:
+                            elementwise_loss = elementwise_loss * loss_mask
+
+                    # Trim each sample with one CPU transfer per tensor.
+                    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
+                        )
+
+                    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

@@ -53,6 +53,7 @@
     all_reduce_metrics,
     compute_minibatch_rollout_logprob_diff_metrics,
     get_microbatch_iterator,
+    pop_return_per_token_outputs,
     reduce_metrics,
 )
 from skyrl.env_vars import (
@@ -838,7 +839,9 @@ 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 include reserved key ``return_per_token_outputs`` to skip
+                per-token ``loss_fn_outputs`` when callers read only ``metrics``.
 
         Returns:
             Metrics dict for the worker's local micro batch
@@ -868,9 +871,12 @@ def _forward_backward_micro(
             # Fall back to config default
             current_loss_fn = self.policy_loss_fn
 
+        # Consume the reserved gate before merging AlgorithmConfig overrides.
+        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
@@ -910,40 +916,41 @@ 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()
+            # Only build per-token outputs for callers that consume them.
+            if return_per_token_outputs:
+                # Tinker consumes per-token NLL.
+                with torch.no_grad():
+                    elementwise_loss = -action_log_probs
+                    if loss_mask is not None:
+                        elementwise_loss = elementwise_loss * loss_mask
+
+                # Trim each sample with one CPU transfer per tensor.
+                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)
+
+                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(),
@@ -1109,6 +1116,13 @@ 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 include reserved key ``return_per_token_outputs`` to skip
+                per-token ``loss_fn_outputs`` when callers read only ``metrics``.
         """
         self.model.eval()
         experience.to_device(torch.cuda.current_device())
@@ -1124,9 +1138,12 @@ def _forward_micro_with_loss(
 
         current_loss_fn = PolicyLossRegistry.get(loss_fn)
 
+        # Consume the reserved gate before merging AlgorithmConfig overrides.
+        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))
@@ -1153,36 +1170,39 @@ 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
+            # Only build per-token outputs for callers that consume them.
+            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()
+                # Trim each sample with one CPU transfer per tensor.
+                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)
+
+                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, Optional
+from typing import Any, Dict, Iterator, List, Optional, Tuple
 
 import torch
 import torch.distributed as dist
@@ -49,6 +49,20 @@ def compute_minibatch_rollout_logprob_diff_metrics(
     }
 
 
+# Reserved ``loss_fn_config`` key consumed before AlgorithmConfig validation.
+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]:
+    """Return a copied config and whether per-token outputs should be built."""
+    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.
 

skyrl/train/sft_trainer.py

@@ -44,6 +44,7 @@
 from skyrl.backends.skyrl_train.utils.io import io
 from skyrl.backends.skyrl_train.workers.worker import PPORayActorGroup
 from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch
+from skyrl.backends.skyrl_train.workers.worker_utils import RETURN_PER_TOKEN_OUTPUTS_KEY
 from skyrl.env_vars import SKYRL_RAY_PG_TIMEOUT_IN_S
 from skyrl.train.config import SkyRLTrainConfig
 from skyrl.train.config.sft_config import (
@@ -1546,11 +1547,12 @@ def run_eval(self) -> tuple[dict, int]:
             # was 0/1 before scaling. Recover the count from the batch by counting positive entries.
             # Padded rows have loss_mask=0 so they are excluded here.
             nonpad_tokens = int((batch["loss_mask"] > 0).sum().item())
+            # Eval consumes metrics only; skip per-token loss_fn_outputs.
             output = self.dispatch.forward(
                 "policy",
                 batch,
                 loss_fn="cross_entropy",
-                loss_fn_config=None,
+                loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False},
             )
             batch_loss = float(output.metrics.get("loss", float("nan")))
             total_loss_weighted += batch_loss * nonpad_tokens
@@ -1571,7 +1573,13 @@ def train_step(self, batch: TrainingInputBatch, step: int) -> dict:
         """
         timings: dict[str, float] = {}
         with Timer("forward_backward", timings):
-            output = self.dispatch.forward_backward("policy", batch, loss_fn="cross_entropy")
+            # SFT consumes metrics only; skip per-token loss_fn_outputs.
+            output = self.dispatch.forward_backward(
+                "policy",
+                batch,
+                loss_fn="cross_entropy",
+                loss_fn_config={RETURN_PER_TOKEN_OUTPUTS_KEY: False},
+            )
         with Timer("optim_step", timings):
             grad_norm = self.dispatch.optim_step("policy")