菜单
Prevents MoE router instability during RL training by recording and replaying inference-phase routing distributions back into training. Reduces training-inference routing divergence and KL divergence, enabling stable MoE RL scaling without sacrificing training speed.
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/ADu2021/skillXiv --skill moe-routing-alignment命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
正在显示 SKILL.md
基于 SOC 职业分类
| name | moe-routing-alignment |
| title | Stabilizing MoE RL by Aligning Training and Inference Routers |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2510.11370 |
| keywords | ["MoE","Routing","Reinforcement Learning","Stability","Training"] |
| description | Prevents MoE router instability during RL training by recording and replaying inference-phase routing distributions back into training. Reduces training-inference routing divergence and KL divergence, enabling stable MoE RL scaling without sacrificing training speed. |
Mixture-of-Experts models with RL training suffer from routing instability: routers diverge between training and inference phases, causing catastrophic collapse. Rollout Routing Replay (R3) synchronizes routing behavior across both phases.
By recording inference routing and replaying those decisions during training, R3 prevents divergence and enables stable scaling.
The core problem: routers make different decisions during training vs. inference, even with identical inputs. This divergence causes:
R3 solves this by:
Record routing distributions during inference rollouts. When the model makes decisions in the environment, capture which experts were selected:
class RoutingRecorder:
def __init__(self, num_experts):
self.num_experts = num_experts
self.routing_history = []
def record_routing(self, inputs, routing_logits):
"""Capture routing decisions during inference."""
# Get expert selection distribution
routing_probs = torch.softmax(routing_logits, dim=-1)
# For each token, record which expert was selected
selected_experts = torch.argmax(routing_probs, dim=-1)
record = {
'inputs': inputs.detach().cpu(),
'routing_logits': routing_logits.detach().cpu(),
: routing_probs.detach().cpu(),
: selected_experts.detach().cpu()
}
.routing_history.append(record)
record
():
.routing_history
():
.routing_history = []
Implement training with enforced routing replay. During policy gradient updates, use recorded routing to constrain router behavior:
class RoutingReplayPG:
def __init__(self, model, recorder, learning_rate=1e-4):
self.model = model
self.recorder = recorder
self.optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
def compute_loss_with_routing_replay(self, batch_inputs, recorded_routing, rewards):
"""Compute RL loss while enforcing routing consistency."""
# Forward pass through model
logits, routing_logits = self.model(batch_inputs)
# Get current routing probabilities
current_probs = torch.softmax(routing_logits, dim=-1)
recorded_probs = recorded_routing['routing_probs']
# Policy gradient loss
action_log_probs = torch.log(current_probs.gather(-1, recorded_routing['selected_experts'].unsqueeze(-1)))
pg_loss = -(action_log_probs * rewards).mean()
# Routing consistency loss: penalize divergence from recorded routing
kl_divergence = torch.nn.functional.kl_div(
torch.log(current_probs + 1e-10),
recorded_probs,
reduction='batchmean'
)
# Total loss: PG + routing alignment
total_loss = pg_loss + 0.5 * kl_divergence
return total_loss, {
'pg_loss': pg_loss.item(),
'kl_div': kl_divergence.item(),
'total_loss': total_loss.item()
}
def training_step(self, batch_inputs, recorded_routing, rewards):
"""Perform one gradient update with routing replay."""
loss, metrics = self.compute_loss_with_routing_replay(
batch_inputs, recorded_routing, rewards
)
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
self.optimizer.step()
return metrics
Implement the full training loop that alternates between inference (recording) and training (replaying):
def train_moe_with_routing_replay(model, env, num_episodes=1000):
"""Main training loop with routing replay."""
recorder = RoutingRecorder(num_experts=model.num_experts)
pg_trainer = RoutingReplayPG(model, recorder)
for episode in range(num_episodes):
# Phase 1: Inference rollouts with routing recording
episode_transitions = []
for _ in range(batch_size):
state = env.reset()
done = False
while not done:
# Forward pass records routing
with torch.no_grad():
action, routing_logits = model(state)
recorder.record_routing(state, routing_logits)
state, reward, done, _ = env.step(action)
episode_transitions.append((state, reward))
# Phase 2: Training with recorded routing
recorded_routing = recorder.get_routing_distribution()
episode_reward = sum(r for _, r in episode_transitions)
# Train on batch with enforced routing
metrics = pg_trainer.training_step(
batch_inputs=torch.stack([s for s, _ in episode_transitions]),
recorded_routing=recorded_routing[0], # Use first recorded distribution
rewards=torch.tensor([r for _, r in episode_transitions])
)
recorder.clear()
if episode % 100 == 0:
print(f"Episode {episode}, Reward: {episode_reward:.2f}, "
f"KL Div: {metrics['kl_div']:.4f}")
| Parameter | Recommendation |
|---|---|
| KL divergence weight | 0.5-1.0 (balance PG loss and routing consistency) |
| Routing history retention | Single batch (memory-efficient) |
| Recording overhead | ~5-10% inference time |
| Batch size for replay | Same as inference batch |
When to use:
When NOT to use:
Common pitfalls:
Reference: Stabilizing MoE RL on arXiv