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

[dyurk-lila]

Reviewers: Where to Look

The behavioral core of this PR is small; concentrate review on these sites:

• The gate at the two FSDP sites in skyrl/backends/skyrl_train/workers/worker.py — _forward_backward_micro (train, gate at :939) and _forward_micro_with_loss (eval, gate at :1201), plus the equivalent gate in Megatron's loss_func inside forward_backward_mini_batch (skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py:737). Key things to confirm: (a) loss, the backward pass, and all consumed scalar metrics are computed outside/before the gated block so they are byte-identical regardless of the flag; (b) the else branch yields one empty dict per sequence, and only the cross_entropy branch is gated — the RL (non-cross_entropy) else-branch that builds logprobs must stay ungated.
• The threading of the parameter from the dispatch layer down to each gate: worker_dispatch.py → worker.py forward_backward (:776) / forward (:1079) → the micro methods (:848, :1142), and the Megatron path megatron_worker.py (:992, :1106) → megatron_model_wrapper.py (:474). It is an explicit keyword argument defaulting to True at every level, so any level that forgets to forward it silently reverts to the old behavior rather than erroring — worth a skim that no hop was missed.
• Caller opt-out in skyrl/train/sft_trainer.py — train_step (forward_backward, :1703) and run_eval (forward, :1677) now pass return_per_token_outputs=False. Confirm the SFT trainer truly reads only output.metrics and never output.loss_fn_outputs, so opting out is safe.

Needs less scrutiny: the near-identical per-token-build bodies moved wholesale under the if return_per_token_outputs: guard (unchanged logic, just re-indented); the docstring/comment additions at each gate site; and the test scaffolding.

---

Summary

Threads an explicit return_per_token_outputs keyword argument (default True) from the worker dispatch layer down to the cross_entropy branch of both backends on both the train and eval/forward paths, gating the per-token loss_fn_outputs build. 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 nothing and keep the existing contract (default True).

What changed

• FSDP (workers/worker.py, workers/fsdp/fsdp_worker.py): accept and forward the flag through forward_backward / forward, and gate the per-token build in _forward_backward_micro (train) and _forward_micro_with_loss (eval).
• Megatron (workers/megatron/megatron_worker.py, workers/megatron/megatron_model_wrapper.py): accept and forward the flag, gating the per-token build inside the shared loss_func (covers train + forward_only eval).
• Dispatch (workers/worker_dispatch.py): plumb the flag through and document the reserved argument on forward / forward_backward.
• Caller wiring (train/sft_trainer.py): train_step and run_eval pass return_per_token_outputs=False.

The flag is a plain function argument rather than a loss_fn_config dict key or an AlgorithmConfig field: it is per-call request metadata, not algorithm configuration, and keeping it out of the config avoids both the key-validation issue in build_nested_dataclass and any pop/copy dance around the AlgorithmConfig merge. It is documented in the parameter docstring at each site that reads it — _forward_backward_micro, _forward_micro_with_loss (FSDP), and forward_backward_mini_batch (Megatron).

Numerical equivalence / safety

Byte-identical for all existing callers:

• The default True at every level reproduces the exact prior code paths; no existing caller passes the flag.
• 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. Because the flag is a Python keyword argument and not a loss_fn_config key, it cannot be injected through the Tinker public API's user-supplied config at all.

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, and loss/response_length/lr are identical across the flag. Includes a case where a real eps_clip_low override is passed via loss_fn_config alongside the flag, confirming the two are independent and the config merge still happens. 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.
• tests/train/test_sft_callbacks.py (extended): assert SFTTrainer.train_step and run_eval pass loss_fn="cross_entropy" and return_per_token_outputs=False to the dispatch.
• tests/train/test_trainer.py: the _forward_backward_micro mock in test_forward_backward_batch_calculations accepts the new argument.
• 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:

uv run --isolated --extra skyrl-train --extra dev pytest \
  tests/backends/skyrl_train/workers/test_sft_loss_fn_outputs_gate.py \
  tests/train/test_trainer.py::test_forward_backward_batch_calculations

CPU run: 12 passed. The CPU suite for the new/extended tests runs automatically under the standard CPU CI job (tests/backends/skyrl_train/ + tests/train/); skyrl_train_tests is green on this head.

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, and is not reached by this argument — a possible follow-up, not a regression.
• No config schema field is added — the flag is a per-call argument with a safe default, preferred over a global switch.

Minor cleanup carried here: the loss_config merge guard in the two FSDP micro methods was relaxed from if loss_fn_config is not None: to if loss_fn_config: so an empty override dict skips a no-op OmegaConf.merge. This is a strict no-op for every existing caller (no caller passes {}, and OmegaConf.merge(base, {}) is itself a no-op); it is incidental to the gate rather than required by it, and is easy to drop if reviewers would rather keep the diff minimal.

Relationship to open PRs

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

• SFT loss aggregation consistent with RL path #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.
• [train] VLM SFT support on Megatron backend (Qwen3-VL) #1752 (VLM SFT on Megatron) — disjoint regions in the shared files.
• [skyrl] Preserve staged forward_backward loss_fn_outputs across DP ranks #1534 (preserve staged forward_backward loss_fn_outputs across DP ranks; worker_dispatch.py) — overlaps in worker_dispatch.py; the changes there are additive (a new forwarded argument plus docstring lines), so expect a trivial merge.

The gate logic itself is composable with all of these.

[gemini-code-assist]

Code Review

This pull request introduces a performance optimization to skip the generation of per-token loss function outputs (such as log probabilities and elementwise NLL) when they are not needed by the caller. A new configuration flag return_per_token_outputs is introduced and extracted from loss_fn_config via a utility function pop_return_per_token_outputs. In SFTTrainer, this flag is set to False during training and evaluation steps to avoid unnecessary GPU-to-CPU transfers and Python loops, as only scalar metrics are consumed. The Megatron and FSDP workers are updated to respect this flag, and comprehensive unit and integration tests are added to verify the correctness of this optimization. No review comments were provided, so there is no feedback to address.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[SumanthRH]

Left a comment on the API!

[SumanthRH]

return_per_token_outputs is a control flag that changes the output type

I would not want to smuggle this inside loss_fn_config

Can we just add the flag return_per_token_outputs as an argument to PolicyWorker.forward_backward and PolicyWorker.forward ? This is more clear to users on behaviour.

We should also make a similar change in WorkerDispatch.forward , WorkerDispatch.forward_backward and WorkerDispatch.forward_backward_from_staged -> allow pass through kwargs for worker-type-specific arguments.

[dyurk-lila]

Good point, I did this refactor