R3 PR: Rollout Routing Replay

skyrl/backends/skyrl_train/inference_engines/base.py

@@ -29,6 +29,7 @@ class InferenceEngineOutput(TypedDict):
     response_ids: List[List[int]]
     stop_reasons: List[str]
     response_logprobs: Optional[List[List[float]]]
+    rollout_inference_indices: Optional[List[List[List[List[int]]]]] # [seq_len, layer_num, topk]
 
 
 class InferenceEngineInterface(ABC):
@@ -63,6 +64,7 @@ async def sample(
         all_responses = []
         all_stop_reasons = []
         all_response_logprobs = []
+        all_rollout_inference_indices = []
 
         for _ in range(num_samples):
             input_batch: InferenceEngineInput = {
@@ -79,12 +81,15 @@ async def sample(
             all_stop_reasons.append(output["stop_reasons"][0])
             if output.get("response_logprobs") is not None:
                 all_response_logprobs.append(output["response_logprobs"][0])
+            if output.get("rollout_inference_indices") is not None:
+                all_rollout_inference_indices.append(output["rollout_inference_indices"][0])
 
         return {
             "response_ids": all_response_ids,
             "responses": all_responses,
             "stop_reasons": all_stop_reasons,
             "response_logprobs": all_response_logprobs if all_response_logprobs else None,
+            "rollout_inference_indices": all_rollout_inference_indices if all_rollout_inference_indices else None,
         }
 
     @abstractmethod

skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py

@@ -135,6 +135,7 @@ def _postprocess_outputs(self, outputs):
         stop_reasons: List[str] = []
         response_ids: List[List[int]] = []
         response_logprobs: Optional[List[List[float]]] = []
+        rollout_inference_indices: Optional[List[List[List[List[int]]]]] = []
 
         for output in outputs:
             # TODO(tgriggs): Support n>1 sampling.
@@ -156,14 +157,25 @@ def _postprocess_outputs(self, outputs):
                     del token_logprobs
             response_logprobs.append(_logprobs)
 
+            if resp.routed_experts is not None:
+                if hasattr(resp.routed_experts, "tolist"):
+                    routed_experts_list = resp.routed_experts.tolist()
+                else:
+                    routed_experts_list = resp.routed_experts
+                rollout_inference_indices.append(routed_experts_list)
+
         if len(response_logprobs) and response_logprobs[0] is None:
             response_logprobs = None  # hack: assume uniform sampling params
 
+        if len(rollout_inference_indices) == 0:
+            rollout_inference_indices = None
+
         return InferenceEngineOutput(
             responses=responses,
             stop_reasons=stop_reasons,
             response_ids=response_ids,
             response_logprobs=response_logprobs,
+            rollout_inference_indices=rollout_inference_indices,
         )
 
     def _get_engine(self):
@@ -320,7 +332,13 @@ def _create_engine(self, *args, **kwargs):
         enable_ray_prometheus_stats = kwargs.pop("enable_ray_prometheus_stats", False)
         enable_log_requests = kwargs.pop("enable_log_requests", False)
         max_log_len = kwargs.pop("max_log_len", None)
-
+
+         # Log if enable_return_routed_experts is being passed
+        if "enable_return_routed_experts" in kwargs:
+            logger.info(f"DEBUG: enable_return_routed_experts={kwargs['enable_return_routed_experts']} is being passed to AsyncEngineArgs")
+        else:
+            logger.warning("DEBUG: enable_return_routed_experts is NOT in kwargs")
+
         if version.parse(vllm.__version__) >= version.parse("0.10.0"):
             engine_args = vllm.AsyncEngineArgs(enable_log_requests=enable_log_requests, **kwargs)
         else:

skyrl/backends/skyrl_train/training_batch.py

@@ -369,6 +369,7 @@ class TrainingInput(TypedDict, total=False):
     kl: Float[torch.Tensor, "batch_size seq_len"]
     rewards: Optional[Float[torch.Tensor, "batch_size seq_len"]]
     rollout_logprobs: Optional[Float[torch.Tensor, "batch_size seq_len"]]
+    rollout_inference_indices: Optional[Integer[torch.Tensor, "batch_size seq_len layer_num topk"]]
 
 
 class TrainingInputBatch(TensorBatch[TrainingInput]):

skyrl/backends/skyrl_train/utils/replay_utils.py

@@ -0,0 +1,62 @@
+"""
+Utility functions for MoE Router Replay.
+"""
+
+import torch
+from typing import Optional, List
+from skyrl.backends.skyrl_train.training_batch import TrainingInputBatch
+
+def _split_replay_indices(rollout_inference_indices: torch.Tensor) -> List[torch.Tensor]:
+    if rollout_inference_indices is None:
+        return None
+    if rollout_inference_indices.dim() != 4:
+        raise ValueError(f"Expected 4D replay indices, got shape {rollout_inference_indices.shape}")
+    per_layer = rollout_inference_indices.permute(2, 0, 1, 3).contiguous()
+    return [per_layer[i] for i in range(per_layer.shape[0])]
+
+def setup_router_replay_forward(data: TrainingInputBatch, enable_router_replay: bool) -> bool:
+    """
+    Set up router replay for forward pass (ref/policy inference).
+    """
+    if not enable_router_replay:
+        return False
+
+    rollout_inference_indices = data.get("rollout_inference_indices")
+    if rollout_inference_indices is None:
+        return False
+
+    from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+
+    RouterReplay.set_replay_data(_split_replay_indices(rollout_inference_indices))
+    RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+
+    return True
+
+
+def setup_router_replay_backward(data: TrainingInputBatch, enable_router_replay: bool) -> bool:
+    """
+    Set up router replay for training forward/backward pass.
+    """
+    if not enable_router_replay:
+        return False
+
+    rollout_inference_indices = data.get("rollout_inference_indices")
+    if rollout_inference_indices is None:
+        return False
+
+    from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+
+    RouterReplay.set_replay_data(_split_replay_indices(rollout_inference_indices))
+    # Use REPLAY_FORWARD - Megatron handles REPLAY_BACKWARD automatically
+    RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+
+    return True
+
+
+def clear_router_replay():
+    """Clear all router replay state."""
+    from megatron.core.transformer.moe.router_replay import RouterReplay
+
+    RouterReplay.clear_global_indices()
+    RouterReplay.clear_global_router_replay_action()
+    RouterReplay.clear_global_router_replay_instances()

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -322,6 +322,7 @@ def init_configs(
 
         self.strategy.hf_config = hf_config
         self.tokenizer = tokenizer
+        self.enable_router_replay = transformer_config_kwargs.get("moe_enable_routing_replay", False)
 
     def configure_lora(self, lora_config, lora_type: Optional[str] = "lora"):
         if lora_type == "lora":
@@ -401,6 +402,10 @@ def forward(self, data: TrainingInputBatch):
         """
         Override `Worker.forward` to support passing the full mini batch to the MegatronModelWrapper.forward method.
         """
+        from skyrl_train.utils.replay_utils import setup_router_replay_forward, clear_router_replay
+
+        setup_router_replay_forward(data, self.enable_router_replay)
+
         # Run in micro batches grouped into a single mini-batch
         micro_bsz = self.cfg.micro_forward_batch_size_per_gpu
         micro_batches = data.chunk(micro_bsz)
@@ -438,6 +443,7 @@ def forward(self, data: TrainingInputBatch):
         log_probs = log_probs.to("cpu")
         output = TrainingOutputBatch({"output": log_probs})
         output.metadata = data.metadata
+        clear_router_replay()
         return output
 
     def save_hf_model(self, export_dir: str, tokenizer):
@@ -593,6 +599,9 @@ def forward_backward(
         Returns:
             Aggregated metrics dict across all micro batches
         """
+        from skyrl_train.utils.replay_utils import setup_router_replay_forward, clear_router_replay
+
+        setup_router_replay_forward(data, self.enable_router_replay)
         self.model.train()
         for chunk in self.actor_module:
             # if use distributed optimizer, zero grad buffer will be handled by optimizer
@@ -665,6 +674,8 @@ def forward_backward(
         # Add loss_fn_outputs back (not reduced, kept as list)
         if all_loss_fn_outputs:
             status["loss_fn_outputs"] = all_loss_fn_outputs
+
+        clear_router_replay()
 
         return status
 

skyrl/train/config/config.py

@@ -440,6 +440,7 @@ class InferenceEngineConfig(BaseConfig):
     """Sets ``VLLM_ENABLE_V1_MULTIPROCESSING=0`` for reproducibility."""
     enable_prefix_caching: bool = True
     enable_chunked_prefill: bool = True
+    enable_return_routed_experts: bool = False
     max_num_batched_tokens: int = 8192
     enforce_eager: bool = True
     """Disable CUDA graphs for stability. Set to ``False`` for higher performance,

skyrl/train/config/megatron_config/policy.yaml

@@ -57,6 +57,7 @@ transformer_config_kwargs:
   recompute_modules: ["core_attn"]
   recompute_method: uniform
   recompute_num_layers: 1
+  moe_enable_routing_replay: ${generator.enable_return_routed_experts}
 
 # flag to manually empty torch's cuda cache between the forward/backward pass and the optimizer step
 # this will free reserved but unallocated memory, and can help avoid OoMs in the optimizer

skyrl/train/config/ppo_base_config.yaml

@@ -286,6 +286,7 @@ generator:
   vllm_v1_disable_multiproc: true
   enable_prefix_caching: true
   enable_chunked_prefill: true
+  enable_return_routed_experts: false
   max_num_batched_tokens: 8192
   # Disable CUDA graphs by default for stability. Set to false for higher performance, but this may affect convergence for long-running and/or long context training jobs.
   enforce_eager: true

skyrl/train/dataset/preprocess.py

@@ -1,7 +1,7 @@
 from typing import List, Tuple, Optional
 import torch
 from transformers import AutoTokenizer
-from jaxtyping import Float
+from jaxtyping import Float, Integer
 
 
 def _verify_inputs(
@@ -32,13 +32,15 @@ def convert_prompts_responses_to_batch_tensors(
     rewards: List[List[float]],
     loss_masks: List[List[int]],
     logprobs: Optional[List[List[float]]] = None,
+    rollout_inference_indices: Optional[List[List[List[List[int]]]]] = None,
 ) -> Tuple[
     Float[torch.Tensor, "batch seq_len"],
     Float[torch.Tensor, "batch seq_len"],
     Float[torch.Tensor, "batch response_len"],
     Float[torch.Tensor, "batch response_len"],
     Float[torch.Tensor, "batch response_len"],
     Optional[Float[torch.Tensor, "batch response_len"]],
+    Optional[Integer[torch.Tensor, "batch seq_len layer_num topk"]],
 ]:
     """
     Convert prompts and responses to batch tensors for training.
@@ -129,4 +131,8 @@ def convert_prompts_responses_to_batch_tensors(
         ]
         logprobs_tensor = torch.tensor(padded_logprobs, dtype=torch.float)
 
-    return sequences, attention_mask, action_mask, ret_rewards, ret_loss_masks, logprobs_tensor
+    rollout_inference_indices_tensor = None
+    if rollout_inference_indices:
+        rollout_inference_indices_tensor = torch.tensor(rollout_inference_indices, dtype=torch.int32)
+
+    return sequences, attention_mask, action_mask, ret_rewards, ret_loss_masks, logprobs_tensor, rollout_inference_indices_tensor

skyrl/train/entrypoints/main_base.py

@@ -68,6 +68,7 @@ def create_ray_wrapped_inference_engines_from_config(
         "backend": ie_cfg.backend,
         "engine_init_kwargs": ie_cfg.engine_init_kwargs,
         "enable_ray_prometheus_stats": ie_cfg.enable_ray_prometheus_stats,
+        "enable_return_routed_experts": ie_cfg.enable_return_routed_experts,
     }
 
     # Conditionally add LoRA parameters if LoRA is enabled

skyrl/train/generators/base.py

@@ -39,6 +39,7 @@ class GeneratorOutput(TypedDict):
     rollout_metrics: Optional[Dict[str, Any]]
     rollout_logprobs: Optional[List[List[float]]]
     trajectory_ids: Optional[List[TrajectoryID]]
+    rollout_inference_indices: Optional[List[List[List[List[int]]]]] # [batch_size, seq_len, layer_num, topk]
     # Applicable only for step-wise training
     is_last_step: Optional[List[bool]]
 

skyrl/train/generators/skyrl_gym_generator.py

@@ -66,6 +66,7 @@ class TurnOutput:
     output_logprobs: Optional[List[float]]
     new_obs: ConversationType
     obs_ids: List[int]
+    rollout_inference_indices: Optional[List[List[List[List[int]]]]] # [seq_len, layer_num, topk]
     reward: Optional[float]
     added_eos: bool = False
 
@@ -300,11 +301,14 @@ async def agent_loop(
             output_ids = engine_output["response_ids"][0]
             stop_reason = engine_output["stop_reasons"][0]
             response_logprobs = engine_output.get("response_logprobs", None)
+            rollout_inference_indices = engine_output.get("rollout_inference_indices", None)
             if response_logprobs is not None:
                 response_logprobs = response_logprobs[0]
                 if self.custom_chat_template is not None:
                     raise ValueError("Response Logprobs bookkeeping is not supported with custom chat template")
 
+            if rollout_inference_indices is not None:
+                rollout_inference_indices = rollout_inference_indices[0]
             # Append eos when sampling_params.stop is not None. Does not affect 3.a as chat templates add eos_token.
             # sampling_params is not None for eval, but None for training (which uses engine.sampling_params which are from cfg)
             stop_strs = current_sampling_params.get("stop", None)
@@ -348,6 +352,7 @@ async def agent_loop(
                 reward=step_reward,
                 obs_ids=obs_ids,
                 added_eos=added_eos,
+                rollout_inference_indices=rollout_inference_indices,
             )
 
             if is_step_wise:

skyrl/train/trainer.py

@@ -604,6 +604,7 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis
         loss_masks: List[List[int]] = generator_output["loss_masks"]
 
         logprobs: Optional[List[List[float]]] = generator_output.get("rollout_logprobs", None)
+        rollout_inference_indices: Optional[List[List[List[List[int]]]]] = generator_output.get("rollout_inference_indices", None)
 
         (
             sequences_tensor,
@@ -612,13 +613,15 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis
             rewards_tensor,
             loss_masks_tensor,
             rollout_logprobs_tensor,
+            rollout_inference_indices_tensor,
         ) = convert_prompts_responses_to_batch_tensors(
             self.tokenizer,
             prompt_ids,
             response_ids,
             rewards,
             loss_masks,
             logprobs,
+            rollout_inference_indices,
         )
 
         # sanity check for off_policy_correction
@@ -639,6 +642,7 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis
                 "rewards": rewards_tensor,
                 "loss_mask": loss_masks_tensor,
                 "rollout_logprobs": rollout_logprobs_tensor,
+                "rollout_inference_indices": rollout_inference_indices_tensor,
                 "is_last_step": (
                     torch.tensor(generator_output["is_last_step"], dtype=torch.bool)
                     if generator_output.get("is_last_step", None) is not None