FSDP token-based `forward_backward` returns `loss_fn_outputs` in packed order when max_tokens_per_microbatch > 0

[adityasoni9998]

When trainer.max_tokens_per_microbatch > 0, FSDP PolicyWorkerBase.forward_backward uses token-based microbatching and bin-packs samples by token count. The training computation appears internally consistent, but the returned WorkerOutput.loss_fn_outputs are emitted in packed microbatch order rather than original input sample order.

This breaks downstream consumers that assume loss_fn_outputs[i] corresponds to input sample i. For example, the Tinker backend indexes per_sample_outputs[i] by prepared-batch sample index when constructing ForwardBackwardOutput. This looks like an API/output ordering issue rather than a core training-loss issue: the model forward, loss computation, and backward pass happen inside each packed microbatch before loss_fn_outputs are detached and returned.

Expected behavior

forward_backward(...).loss_fn_outputs should be in the same per-sample order as the input batch/datums, matching the existing Tinker API expectation.

Actual behavior

With token-based microbatching enabled, loss_fn_outputs are returned in packed microbatch order.