feat: support server-side PPO loss on Tinker-compatible backends

[nashjojo]

Description

The Tinker custom-loss path transfers model logprobs to the client, evaluates the objective, then sends gradients back for another forward/backward request. Add an opt-in model.tinker.server_loss_fn: trinity_ppo path for compatible servers to calculate dual-clipped PPO with optional K2 KL in one server-side forward/backward operation. The default client callback remains unchanged.

The client aligns old logprobs, advantages and optional reference logprobs with shifted target tokens, preserves masked response tokens in the weights, and supplies the full optimizer-batch datum count to every SDK request. This preserves the per-datum masked-token mean when the SDK splits a batch. One optimizer step follows the complete batch; a failed server result does not trigger it. Zero KL skips reference-logprob requests, and diagnostics use additive token counts rather than averaging request means.

Unsupported objectives fail during trainer initialization: asymmetric clipping, sequence masking, fallback policy gradient, non-K2 KL, adaptive KL, entropy loss, alternative reductions, or microbatch rescaling. The server-loss parameters come from the configured algorithm; there is no separate override that can silently change the objective.

Dependencies and scope

• Requires a server that implements the trinity_ppo extension; TuFT support is proposed in agentscope-ai/TuFT#163. This is not a built-in loss on all Tinker services.
• FSDP deployments also need safe handling of SDK requests smaller than the world size (TuFT#164) and bounded, collective-safe microbatching for long sequences. This client PR does not establish an end-to-end memory bound or reproduce a GPU speed benchmark.
• No full-parameter training changes, SDK private-limit monkey patches, experiment-specific paths, or arbitrary request-concurrency settings are included.

Validation

• 13 CPU regression tests and 12 subtests passed: input alignment, masked/empty responses, scalar advantages, required fields, objective validation, additive metrics, default callback behavior, disabled-KL reference requests, optimizer ordering, and server failure handling.
• An additional CPU contract check used Tinker 0.25 protobuf serialization and the actual loss implementation from TuFT#163 (8dc6a68). With KL 0 and 0.001, unequal sequence lengths, a fully masked datum, and uneven 1+2 request splits, loss and gradients matched Trinity's PPO/K2 implementation; maximum gradient difference was zero.
• Black, isort, flake8, mypy, and git diff --check passed for the changed code.
• Mock trainer tests ran with CUDA disabled. The available vLLM installation lacks the vllm.inputs.MultiModalDataDict type alias imported by current main, so the test process supplied that type-only alias before collection.
• No GPU training or new end-to-end experiment was run on this main-based change.

Checklist

• Focused tests and static checks passed as described above.
• Docstrings and the Tinker example documentation describe the option and its limits.
• Code is ready for review.
• Full repository and GPU integration tests.

[nashjojo]

/unittest-module-trainer

[nashjojo]

/unittest-pattern-ServerLoss

[github-actions]

unittest: Run #1884

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Pending ⏳	Other ❓	Flaky 🍂	Duration ⏱️
15	15	0	0	0	0	0	16.6s

🎉 All tests passed!

Github Test Reporter by CTRF 💚

[nashjojo]

/unittest-all

[nashjojo]

/unittest-all

[github-actions]

unittest: Run #1889

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Pending ⏳	Other ❓	Flaky 🍂	Duration ⏱️
491	486	0	5	0	0	0	2h 35m

🎉 All tests passed!

Github Test Reporter by CTRF 💚