Skip to content

[Single Controller] Router replay (R3) not supported/validated on the async + TransferQueue path #3327

Description

@terrykong

Summary

Router replay (R3) is partially wired but not supported end-to-end or validated on the Single Controller (SC) async + TransferQueue (TQ) path, and there is no fail-loud guard — so enabling it runs silently with incorrect routing.

The SC path is exactly the "async + TQ data plane" combination that the legacy path explicitly blocks:

RL/nemo_rl/algorithms/grpo.py

Lines 3539 to 3546 in 850b583

wandb_enabled=master_config.logger["wandb_enabled"],
wandb_config=master_config.logger["wandb"],
),
max_rollout_turns=None,
greedy=False,
effort_config=_get_effort_config(master_config),
reward_penalty_config=master_config.reward_penalties,
thinking_tags=get_nemo_gym_thinking_tags(master_config.env),

policy.router_replay.enabled=true with async GRPO is currently supported only
when data_plane.enabled=false. Async + TQ support has not been merged yet.

What is wired (via #3267 and the SC stack)

  • Generation side: configure_vllm_for_router_replay(master_config.policy) in SC _build_generation sets enable_return_routed_experts (single_controller_utils/setup.py:157).
  • NeMo-Gym actor: require_routed_experts=router_replay_enabled(...) (setup.py:364-369).
  • Train side field selection: TQPolicy reads fields_with_optional_routed_experts(DP_TRAIN_FIELDS, enabled=self._router_replay_enabled) (nemo_rl/models/policy/tq_policy.py:121-183, 279-398).

What is missing

  1. No routed_experts preservation in the SC rollout → TQ → train path. routed_experts capture lives only in the legacy sync path nemo_rl/experience/rollouts.py (_attach_routed_experts_to_message_log_prefix, etc.) and in grpo's driver path (_preserve_router_replay_routed_experts, grpo.py:1731-1757). The SC rollout/buffer path (experience/rollout_manager.py, algorithms/async_utils/, DataPlane payload/schema) never writes routed_experts into the TQ records, so TQPolicy reads them missing at train time.
  2. No SC-side guard. No validate_router_replay_config / NotImplementedError is raised anywhere in single_controller.py / single_controller_utils/ / tq_policy.py. Unlike grpo.py:3539, an SC run with policy.router_replay.enabled=true does not fail loud.
  3. No SC R3 validation. The existing R3 async test (tests/unit/algorithms/test_grpo_router_replay_async.py) and recipe (examples/configs/recipes/llm/grpo-qwen3-30ba3b-thinking-swe1-16n8g-megatron-cp2-r3-async-gym.yaml) target the legacy non-TQ path; there is no SC + R3 recipe or functional test.

Effect

An SC run (data_plane.enabled=true, impl=transfer_queue) with policy.router_replay.enabled=true produces generations with routed-expert traces that are then dropped before training, so the trainer computes MoE logprobs without the replayed routing — the exact correctness problem R3 exists to prevent — with no error surfaced.

Ask

Either:

  • (a) Wire routed_experts through the SC rollout → TQReplayBuffer → DataPlane → TQPolicy train path end-to-end (mirroring grpo's _preserve_router_replay_routed_experts) and add an SC + R3 recipe/functional test; or
  • (b) Until (a) lands, add a fail-loud guard at SC setup (mirror grpo.py:3539) so router_replay.enabled=true on the SC/TQ path raises instead of silently mis-training.

Flagged during review of #3267. Tracked as an SC loose end in #2625.

Metadata

Metadata

Assignees

Labels

accuracybugSomething isn't working

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions