Skip to content

feat: add rollout routing replay for MoE RL - #524

Merged
xsuler merged 11 commits into
mainfrom
feat/rollout-routing-replay
Aug 27, 2026
Merged

feat: add rollout routing replay for MoE RL#524
xsuler merged 11 commits into
mainfrom
feat/rollout-routing-replay

Conversation

@xsuler

@xsuler xsuler commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Capture MoE expert ids selected during rollout and replay them in the policy training forward.
  • Recompute routing weights from current router logits, preserving router gradients while removing rollout/train expert-selection drift.
  • Preserve routing metadata through continuous-batching DP merging and agentic multi-turn trajectory assembly.
  • Remove host synchronization from route capture to avoid the earlier rollout and training slowdown.
  • Enable R3 automatically for CUDA sparse-MoE rollout training. Dense models and MLX retain their original paths.

Root cause

Two integration gaps made R3 ineffective for agentic RL:

  1. Async DP output merging discarded routed_experts, so rollout paid the route-capture cost but policy training received no replay metadata.
  2. Multi-turn assembly sliced newly rendered routes using the previous training-row length. Chat-template renders can share only a shorter token prefix, so this misaligned routes across turns.

The fix preserves and reorders route rows with their prompts, then appends routes from the actual shared causal-prefix boundary.

Compatibility

  • CUDA sparse-MoE rollout training: R3 enabled automatically.
  • Dense models: unchanged original path.
  • MLX: unchanged original path.
  • Existing rollout callback mocks without routing metadata remain compatible.

Validation

  • Full manual pre-commit suite passed.
  • CPU tests cover capture/replay, async DP prompt reordering, agentic prefix boundaries, and dense/MoE default resolution.
  • Ling 3.0 Tiny Office agentic E2E produced replay context (2400, 23, 8), logp_abs_diff_mean=0.005596, and logp_diff_mean=0.001522.

Fixes #522

@xsuler xsuler self-assigned this Aug 26, 2026
@xsuler
xsuler requested a review from indevn August 26, 2026 09:05
@xsuler
xsuler deployed to push_check August 26, 2026 09:05 — with GitHub Actions Active
@xsuler
xsuler deployed to push_check August 26, 2026 09:24 — with GitHub Actions Active
@xsuler
xsuler deployed to push_check August 26, 2026 09:37 — with GitHub Actions Active
Comment thread areno/api/trainers/ppo.py
Comment thread areno/api/backend/cuda/training.py Outdated
Comment thread areno/api/agentic.py Outdated
@xsuler
xsuler deployed to push_check August 26, 2026 22:17 — with GitHub Actions Active
@xsuler
xsuler requested a review from indevn August 26, 2026 22:26
@xsuler
xsuler force-pushed the feat/rollout-routing-replay branch from 044114d to 91ed341 Compare August 27, 2026 00:26
@xsuler
xsuler deployed to push_check August 27, 2026 00:26 — with GitHub Actions Active
@xsuler
xsuler deployed to push_check August 27, 2026 01:43 — with GitHub Actions Active
@xsuler
xsuler deployed to push_check August 27, 2026 01:45 — with GitHub Actions Active
@xsuler
xsuler merged commit 78ebad5 into main Aug 27, 2026
6 checks passed
@xsuler
xsuler deleted the feat/rollout-routing-replay branch August 27, 2026 02:16
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.

Support rollout routing replay (R3) for MoE RL training

2 participants