[tinker] 17/n towards Kimi K2.6: run Tinker API forward requests through the training loss path - #2068
Conversation
The FORWARD endpoint dispatched to the loss-less worker forward, which takes Megatron's inference route and bypasses training-only transforms: fake-INT4 QAT weight quantization and the fused LM-head logprob. Its logprobs therefore diverge from what forward_backward computes gradients against — on our Kimi K2.7 INT4-QAT runs the mismatch showed up as a persistent rollout/train logprob gap in importance-sampling metrics and skewed eval losses, because forward() silently scored the *unquantized* weights. The tinker FORWARD request always carries a loss_fn; use it (policy roles): validate + normalize exactly like forward_backward, then call the dispatch's loss-path forward (forward_only). Critic and loss-less callers keep the plain forward. Co-authored-by: Cursor <cursoragent@cursor.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the _forward_single_model_batch method in skyrl_train_backend.py to ensure that the forward pass runs the same pipeline as the forward-backward pass for non-critic roles by utilizing the loss function and its configuration. This prevents logprobs from diverging due to bypassed training-only transforms. The reviewer suggests adding a warning log when mixed loss functions are detected in a single batch, as only the first loss function is currently applied to the entire batch.
| loss_fn = prepared_batch.all_loss_fns[0] if prepared_batch.all_loss_fns else None | ||
| loss_fn_config = next((c for c in prepared_batch.all_loss_fn_configs if c is not None), None) | ||
| if role != "critic" and loss_fn is not None: | ||
| self._validate_batch_role_and_loss(role, loss_fn) |
There was a problem hiding this comment.
When a batch contains mixed loss functions (or a mix of loss-bearing and loss-less requests), only the first loss function is used for the entire batch. To prevent silent issues and maintain consistency with _forward_backward_single_model_batch, we should log a warning when mixed loss functions are detected in the batch.
| loss_fn = prepared_batch.all_loss_fns[0] if prepared_batch.all_loss_fns else None | |
| loss_fn_config = next((c for c in prepared_batch.all_loss_fn_configs if c is not None), None) | |
| if role != "critic" and loss_fn is not None: | |
| self._validate_batch_role_and_loss(role, loss_fn) | |
| loss_fn = prepared_batch.all_loss_fns[0] if prepared_batch.all_loss_fns else None | |
| if loss_fn is not None and len(set(prepared_batch.all_loss_fns)) > 1: | |
| logger.warning( | |
| "SkyRL backend received mixed loss functions %s in one batch; using '%s' for all", | |
| set(prepared_batch.all_loss_fns), | |
| loss_fn, | |
| ) | |
| loss_fn_config = next((c for c in prepared_batch.all_loss_fn_configs if c is not None), None) | |
| if role != "critic" and loss_fn is not None: | |
| self._validate_batch_role_and_loss(role, loss_fn) |
|
Hey Casper, thanks again for this PR! Was wondering if you could explain how you got the loss-less worker forward to bypass QAT? Looking at the code, it seems like both loss/loss-less paths should use QAT? Here are some of my notes on the function calls: MegatronPolicyWorker.init_model() [megatron_worker.py]
_forward_single_model_batch() [skyrl_train_backend.py]
(From my previous comment on #2030) Maybe there’s another reason why the logprobs diverge (and if so, we should probably update the comment)? What do you think? |
2 participants
Part of the Kimi K2.6/K2.7 series. Standalone. This is the second piece salvaged from #2030 (closed as superseded by #2021), as flagged there.
What
_forward_single_model_batchin the SkyRL-Train backend now routes FORWARD requests through the same loss pipeline asforward_backward(the dispatch's loss-path forward withforward_onlysemantics) whenever the request carries aloss_fnand the role is not critic — the tinker FORWARD request always carries one. Validation and loss normalization are shared withforward_backward. Critic and loss-less callers keep the plain worker forward.Why
The loss-less worker forward takes Megatron's inference route, which bypasses training-only transforms — fake-INT4 QAT weight quantization (#1862) and the fused LM-head logprob (#1841). Its logprobs therefore diverge from what gradients are computed against.
On our Kimi K2.7-Code INT4 QAT runs through the Tinker API this was not cosmetic:
forward()silently scored the unquantized weights, so client-side importance-sampling corrections computed from FORWARD logprobs disagreed with the training-side logprobs (a persistent rollout/train logprob gap in the #1880 metrics), and eval losses measured through FORWARD were skewed relative to training loss. With FORWARD on the loss path, the same batch scored viaforward()and viaforward_backward()produces matching logprobs.Made with Cursor