feat(fireworks): fold ECHO env loss into the client loss path (single pass)

[jeffreysijuntan]

What

ECHO's environment-prediction loss is currently applied on the Fireworks backend as a second, gradient-accumulated cross_entropy forward_backward on top of the builtin GRPO kernel pass. The managed builtin kernel can only express the policy loss, so the aux term needed its own pass (extra forward + backward over the same rollout tokens).

This PR folds ECHO into a single pass via the cookbook's client loss path. When an auxiliary loss is active, forward_backward_from_trajectory_groups now runs one forward_backward_custom whose closure computes the policy loss (training.utils.rl.losses.build_loss_fn) plus each aux term, from the single forward's per-token logprobs. Non-ECHO runs keep the fast server-side builtin kernel path. Tinker is unchanged (its cookbook has no client-side policy closure; ECHO isn't wired on Tinker in the cookbooks).

Why it's equivalent

• Env term: reuses the exact same per-token weights the separate pass built (build_aux_ce_datum with scale = coef/(n_seq·|O_i|)). Adding Σ_t weights_t·(−logp_t) inside the closure vs. a separate accumulated cross_entropy backward yields the same gradient.
• Policy term: now computed by the cookbook's client closure (make_grpo_loss_fn → run_loss_loop) instead of the server builtin kernel. The cookbook asserts these are in-lockstep (the client path is its documented always-correct fallback), so any difference is at most fp-level. Advantages are folded and optim_step keeps GradAccNormalization.NONE while aux losses are active, exactly as before — neither term is rescaled.

Cost

ECHO drops from ~2 fwd + 2 bwd to ~2 fwd + 1 bwd per step (forward_backward_custom = one logprobs forward + one linear backward). One fewer backward, no extra rollouts.

Changes

• rllm/trainer/fireworks/fireworks_policy_trainer.py
  • new _forward_backward_client_with_aux (client-path fused pass)
  • new module-level _fold_aux_into_loss (wraps the policy closure to add aux terms)
  • branch in forward_backward_from_trajectory_groups: aux active → client path, else builtin kernel (unchanged)
  • removed now-dead _build_aux_datums; updated optim_step comment
• tests/unified_trainer/test_aux_loss.py: fold reproduces the paper's L_env; _fold_aux_into_loss adds the term with the correct gradient (d/dlogp = −weights) and skips obs-free datums

Notes

• Faithful port: the managed env mask is mask==0, which includes initial-prompt tokens (pre-existing behavior; verl's ECHO excludes them via the response mask). Preserved here to keep equivalence — flagging in case we want to align them separately.
• Verification: py_compile ✓, ruff ✓, trainer + backend import ✓, 19/19 test_aux_loss.py pass.

🤖 Generated with Claude Code