Skip to content

[Fix] Transfer routed experts release ownership - #2064

Open
matrix72c wants to merge 1 commit into
InternLM:mainfrom
matrix72c:fix/routed-experts-release-ownership
Open

[Fix] Transfer routed experts release ownership#2064
matrix72c wants to merge 1 commit into
InternLM:mainfrom
matrix72c:fix/routed-experts-release-ownership

Conversation

@matrix72c

@matrix72c matrix72c commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Motivation

When enable_return_routed_experts=True, an RL rollout can attach Ray ObjectRefs to a RolloutState. The same reference may either be created by the rollout worker or be borrowed from a TraceStore session. Before this change, RolloutState did not record which component owned the reference, while several cleanup helpers unconditionally freed references. This made retryable stale samples unsafe and could invalidate references that were still held by the TraceStore.

Failure chain before this PR

The problematic path is easiest to see for a retryable stale sample:

  1. A rollout group becomes stale in ReplayBuffer and is selected for retry.
  2. The old lifecycle calls reset_rollout_response() to keep the prompt and clear the generated response.
  3. reset_rollout_response() also calls free_object_refs() for every routed-experts ObjectRef it sees. It cannot distinguish a locally created rollout reference from a borrowed TraceStore reference.
  4. For a TraceStore sample, the Trie still contains the same reference because the session has not been released yet. The reset therefore frees a shared object too early.
  5. The retry keeps the existing trace/session lifecycle, and a later export or training read tries to fetch the reference from the Trie. The reference is already invalid, which can surface as an object-fetch failure or as routed-experts/sequence-length validation errors.

There were two related versions of the same ownership bug:

  • Partial-rollout post-processing awaited its input routed-experts reference and then freed it unconditionally. A borrowed reference could therefore be destroyed while the TraceStore still owned it.
  • Generic discard recursively freed references without an explicit release decision. Cleanup correctness depended on the caller knowing whether a session had already been released. In addition, the training-batch cleanup ran only after a successful prepare/fit sequence, so an exception could leave a TraceStore session alive.

Finally, replacing a value in the TraceStore trie did not release routed-experts references from the overwritten value, which could leak Ray objects during reroll/overwrite workloads.

What changed

  • Add RolloutState.routed_experts_owner with the values "rollout" and "trace_store".
  • Make reset_rollout_response() a pure state reset: it clears response fields and routed-experts fields, but never calls ray.free.
  • Add release_owned_routed_experts() for explicit caller-owned release, and make discard_rollout_state(..., release_refs=...) opt in to releasing resources.
  • Mark routed-experts references at every direct rollout producer (RolloutWorker, vLLM rollout parsing, and the VERL tool loop) and at every TraceStore export boundary.
  • On retryable stale samples, release only "rollout" references before reset; borrowed "trace_store" references are detached but remain valid until their session is released.
  • Make partial-rollout input release explicit. RolloutWorker.generate() opts in for direct rollout inputs; the handler default is non-releasing for safe reuse by other callers.
  • Keep TraceStore session release as the final release point for its references, detach released routes before generic discard, and ensure the trainer performs session release in a finally block.
  • Release routed-experts references that become unreachable on trie overwrite, while deduplicating references shared by replacement values or multiple trie entries.

Ownership after this change

Direct rollout:
  producer creates ref -> owner = "rollout"
  caller explicitly releases ref -> reset/discard clears state

TraceStore export:
  Trie/session owns ref -> RolloutState borrows it (owner = "trace_store")
  retryable reset only detaches it
  session release -> Trie frees the ref

The LMDeploy generation protocol is unchanged. Ownership is established at XTuner's producer/export boundaries, and old checkpoints remain loadable because the new field defaults to None.

Impact

This prevents premature freeing of shared TraceStore references, makes direct-rollout cleanup explicit, and closes exception/overwrite cleanup gaps without introducing a global reference registry or changing stale reroll/session semantics.

Tests

The following targeted checks pass locally:

  • ruff check for all changed source and test files
  • ruff format --check for all changed source and test files
  • python -m py_compile for all changed source and test files
  • RL state/rollout/producer/trajectory tests: 95 passed in total
  • Staleness ownership tests: 10 passed
  • Retryable replay ownership tests: 2 passed
  • Rollout logic tests: 52 passed
  • Producer tests: 25 passed
  • Trajectory logging tests: 8 passed
  • TraceStore unit tests: 2 passed, 6 skipped because Ray actor initialization is unavailable in the current local environment

The full replay-buffer suite was also inspected; an existing async save/resume test does not complete in this shared test environment, so it was not used as a passing signal.

Related issue

Related to #2025. This PR fixes routed-experts ObjectRef lifetime leaks and premature frees on the XTuner RL path. It does not, by itself, bound learner-side materialization or change sequence-parallel transfer order; those memory-footprint issues remain separate follow-up work. The LMDeploy endpoint and generation protocol are unchanged.

@matrix72c
matrix72c force-pushed the fix/routed-experts-release-ownership branch 2 times, most recently from 47f6556 to 1979eb8 Compare September 7, 2026 03:52
@matrix72c
matrix72c force-pushed the fix/routed-experts-release-ownership branch from 1979eb8 to b68b221 Compare September 7, 2026 05:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant