feat: Add routing replay support for FSDP

[XianlongLi]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

This PR introduces routing replay support for Mixture-of-Experts (MoE) models when using FSDP (Fully Sharded Data Parallel) backend in the twinkle framework. It enables three replay modes (disabled, R2, and R3) to ensure consistent expert routing during reinforcement learning training. The feature is demonstrated through a GRPO training example on math tasks.

Key Changes:

Core Router Replay Infrastructure (src/twinkle/model/transformers/moe/router_replay.py):

• Introduces RouterReplayAction enum defining three replay modes: RECORD, REPLAY_FORWARD, and REPLAY_BACKWARD
• Implements a global registry system for per-MoE-block replay state management
• Provides APIs to record routing decisions, inject pre-computed routings, and retrieve routing data

Expert Parallel Module Enhancements (src/twinkle/model/transformers/moe/expert_parallel.py):

• Adds block-level replay state integration into the router logic
• Implements recording and replay of expert routing indices

TransformersModel Integration (src/twinkle/model/transformers/transformers.py):

• Adds enable_router_replay configuration option
• Implements _router_replay_setup() method for managing replay lifecycle
• Integrates replay into forward(), forward_only(), and backward() methods
• Supports manual cleanup for forward_backward() operations

Data Format Updates:

• Extends InputFeature with routed_experts field (src/twinkle/data_format/input_feature.py)
• Adds routed_experts output support in ModelOutput (src/twinkle/data_format/output.py)
• Adds routed_experts field to SampledSequence (src/twinkle/data_format/sampling.py)

Input Processing Pipeline (src/twinkle/processor/base.py ):

• Adds align_routed_experts() preprocessing step to align vLLM routing data with input sequence lengths
• Handles padding mismatch between routed_experts and input sequences (±1 token tolerance)
• Updates padding and concatenation logic for the new field

vLLM Sampler Integration (src/twinkle/sampler/vllm_sampler/ files):

• Extracts and propagates routed_experts data from vLLM sampling output
• Ensures routing data flows through the sampling pipeline

Usage Example (new files):

• Complete GRPO training cookbook example with routing replay for GSM8K math problems (cookbook/rl/grpo/short_math_grpo_routing_replay.py - 334 lines)
• Accompanying shell script for execution (short_math_grpo_routing_replay.sh)

Replay Modes:

• R2: Records expert routing during a forward-only RECORD pass, then replays during training
• R3: Uses routed_experts data returned by vLLM (v0.14.0+) for direct replay without recording
• Disabled: Standard behavior without routing replay

Usage Tutorial

Refer to cookbook/rl/grpo/short_math_grpo_routing_replay.py

Experiment results

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[tpx818]

/gemini review

Please focus on whether the routing returned in R3 mode and R2 mode is correct, and whether R3 mode correctly uses the vLLM expert routing.

[gemini-code-assist]

Code Review

This pull request introduces GRPO training with routing replay for MoE models under the FSDP backend, supporting three routing replay modes (disabled, R2, and R3). The changes span the training cookbook, CLI arguments, data formats, expert parallel logic, router replay utilities, transformers model class, input processor, and vLLM sampler to handle recording, replaying, and aligning expert routing decisions. The review feedback highlights several critical issues that need to be addressed: potential AttributeError exceptions when labels is None in transformers.py, a guaranteed runtime crash in base.py due to incorrect tensor stacking (concat = False resulting in a 5D tensor instead of 4D), a TypeError when calling self.experts as a nn.ModuleList when EP = 1, potential padding errors if experts_seq_len > input_seq_len, robust handling of string annotations in inspect.signature, and an unused variable recompute_logps in the cookbook.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[XianlongLi]

@tpx818 Please review again.