[feature] Honor max_tokens_per_microbatch in the loss-path forward (forward with loss_fn)

[avigyabb]

Problem

Token-based micro-batching (trainer.max_tokens_per_microbatch, added in #1477) is honored by the loss-less forward and by forward_backward, but the loss-path forward (i.e. forward(data, loss_fn=...), used for eval / loss-only inference such as Tinker FORWARD requests carrying a loss) silently ignores it and always chunks by micro_forward_batch_size_per_gpu:

• Megatron: MegatronPolicyWorker.forward — the loss_fn is None branch goes through _forward_logprobs, which uses get_microbatch_iterator and reorders results back to input order; the loss_fn branch builds micro-batches with a plain BatchIterator.
• FSDP / base worker: PolicyWorkerBase.forward in workers/worker.py — same split: the loss-less branch uses get_microbatch_iterator + reorder_and_combine_batches, the loss branch uses BatchIterator.

Why it matters

• OOM risk: a config sized so that training and loss-less forward only fit thanks to token packing can still OOM on a loss-path forward, since that path falls back to fixed sample-count chunking. This becomes more visible if FORWARD requests are routed through the loss path (as proposed in [tinker] 17/n towards Kimi K2.6: run Tinker API forward requests through the training loss path #2068).
• Consistency: forward, forward(loss_fn=...), and forward_backward should batch the same way so their compute/memory behavior (and micro-batch composition) matches.

Proposal

Mirror forward_backward in both workers' loss-path forward:

1. Build micro-batches with get_microbatch_iterator(data, micro_batch_size=micro_forward_batch_size_per_gpu, max_tokens_per_microbatch=...).
2. For Megatron, pad micro-batches to a uniform size with _pad_microbatch_to_size (Megatron's pipeline schedule requires uniform micro_batch_size) and track num_real_microbatches for metric normalization.
3. Map per-sample loss_fn_outputs back to input order with TokenBasedBatchIterator.reorder_and_combine_items, dropping entries produced for padding — token packing permutes samples, so this step is required for correctness (see [Tinker API] Restore variable-length result order #2043 and [fix] Restore input order for FSDP forward_backward loss_fn_outputs under token batching #2209, which fixed the missing reorder in the two forward_backward paths).
4. Skip padding-microbatch metrics when reducing, as forward_backward does.

Regression coverage can extend the stub-worker tests added in #2209 (tests/backends/skyrl_train/test_token_based_batching_utils.py) to the loss-path forward.

Related

• feat: add max_tokens_per_microbatch config for token-based micro-batching #1477 (introduced token-based micro-batching + the loss-less forward reorder)
• [Tinker API] Restore variable-length result order #2043 (restored input order for megatron forward_backward outputs)
• [fix] Restore input order for FSDP forward_backward loss_fn_outputs under token batching #2209 (same fix for the base/FSDP forward_backward)
• [tinker] 17/n towards Kimi K2.6: run Tinker API forward requests through the training loss path #2068 (proposes routing Tinker FORWARD through the loss path, which makes this gap user-visible)