Skip to content

[Fix] Use Ray reference counting for routed-expert lifetime - #2243

Open
matrix72c wants to merge 1 commit into
InternLM:stack/rl-ppofrom
matrix72c:fix/routed-experts-reference-lifetime
Open

matrix72c wants to merge 1 commit into
InternLM:stack/rl-ppofrom
matrix72c:fix/routed-experts-reference-lifetime

Conversation

@matrix72c

Copy link
Copy Markdown
Contributor

Stacked on #2186 (stack/rl-ppo).

Root cause

Routed-expert ObjectRefs are shared: partial rollouts concatenate them, the trace store keeps them per session, sibling samples and trainable segments can point at the same object, and the actor reads each pack several times (old/ref/teacher/train). Several holders still call ray.internal.free on them: reset_rollout_response, trace-store release, _resolve_routed_experts after the first read, and the actor's free_rollout_routed_expert_refs at the end of fit. A force-free by one holder invalidates the object for every other holder. free_rollout_routed_experts_in_worker / _agent_loop_manager_uses_trace_store only switch off one of those paths.

Fix

Use Ray reference counting instead of explicit frees. Each holder drops its own reference when its lifetime ends (reset/discard, session release, after the training batch), and Ray reclaims the object after the last one.

  • Remove the explicit frees above, the free_rollout_routed_experts_in_worker option and the trace-store special case in the trainer.
  • free_rollout_state_refs detaches routed_experts before freeing the payloads the state owns exclusively.
  • The trainer drops the batch's routed_experts after fit, together with releasing its trace sessions.
  • vLLM HTTP responses must carry routed-expert arrays. A serialized ObjectRef received over HTTP is rejected, because an out-of-band ref cannot take part in reference counting.

Test plan

tests/rl/test_routed_experts_lifetime.py runs against a local Ray cluster. It covers reset/discard keeping sibling routes alive, partial-rollout concatenation keeping borrowed inputs, trace-store overwrite and batch release keeping other consumers, repeated batches leaving no object references in the owner, and the trained batch releasing its refs while an external borrower keeps the data. test_trace_store.py and test_rollout_logic.py are updated for the new lifetime. With the trainer and PPO role suites (test_rl_colocate_trainer.py, test_rl_disaggregated_trainer.py, test_ppo_role_wiring.py, test_routed_experts_dtype.py, test_training_worker_rank.py): 156 passed on CPU.

This branch has not been deployed

No deployments
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