refactor(unified_trainer): extract step merging into a shared backend-agnostic module

[listar2000]

Summary

Both the Tinker and Verl transforms used to implement near-identical cumulative-prefix step merging in backend-specific code, with the same delta-extension loop accessing slightly different fields. This PR pulls the merge into one place behind a small typed protocol, and turns the backend transforms into thin Datum / DataProto adapters over a common MergedSegment.

What changed

New module rllm/experimental/common/step_merge.py:

• MergedSegment — post-merge intermediate (prompt_ids, response_ids, plus per-flat-token response_mask / response_logprobs / response_advantages and an extras dict).
• merge_trajectory_steps(trajectory, *, token_ops, ...) — walks Trajectory.steps once and emits one MergedSegment per cumulative-prefix run.
• TokenOps Protocol typed against TokenInput from rllm/experimental/rollout/types.py. DefaultTokenOps covers list[int] backends (Verl) and is the default; Tinker ships its own _TinkerTokenOps that handles EncodedTextChunk unwrapping and chunk .length counting.
• PerTokenExtras (Tinker router_replay) and per_segment_extras (Verl multi_modal_inputs) are the two extension hooks the existing backends needed.

Backend transforms become thin adapters:

• rllm/trainer/tinker/transform.py: trajectory_to_datums calls the merge with _TINKER_OPS, then _segment_to_datum builds the Tinker Datum (concat prompt+response, prepend prompt-region zeros, right-shift).
• rllm/experimental/verl/transform.py: _process_trajectory calls the merge with the default ops, pad_short_logprobs=True, skip_steps_without_model_output=True, then materialises tensors and feeds accumulated.add_step.

Net diff in existing files: -262 / +136 lines (~−40%).

Tests

New tests/unified_trainer/test_step_merge.py (33 tests) covers the merge directly with both the default int ops and a mock chunk-aware TokenOps mirroring what Tinker ships:

• Single-step + per-token advantage
• Multi-step prefix merging (int and chunks)
• Prefix breaks and mixed prefix/break
• ImageChunk equality on prefix matches and splits; delta_obs containing chunks expands the per-token arrays correctly
• Strict (require_logprobs / require_advantage) vs tolerant (pad_short_logprobs, missing→zeros) modes
• skip_steps_without_model_output filter
• PerTokenExtras (routing matrices) with pad behavior over delta_obs and steps lacking the field
• per_segment_extras (multimodal) taken from each segment's first step
• Per-flat-token length invariant in both ops

The Tinker test file drops TestIsPrefix and TestFlattenTokenInput (now subsumed by test_step_merge.py) and keeps the Datum-builder tests; one assertion-message regex was relaxed to match the new wording.

Test plan

• ./venv/bin/python -m pytest tests/unified_trainer/ — 70 passed (33 step_merge + 9 tinker + 6 verl + 22 others)
• No regressions in broader tests/ suite (the same 13 pre-existing failures on main reproduce, unrelated to this change)

🤖 Generated with Claude Code

[claude]

⚠️ Code review skipped — your organization's overage spend limit has been reached.

Code review is billed via overage credits. To resume reviews, an organization admin can raise the monthly limit at claude.ai/admin-settings/claude-code.

Once credits are available, reopen this pull request to trigger a review.

[luyuzhe111]

would like to see some e2e tests on this one to ensure no regression : )

[listar2000]

would like to see some e2e tests on this one to ensure no regression : )

For sure. Would do the tests soon.