fix(grpo): make runtime contracts and GLM router replay fail closed

[Alvorecer721]

What does this PR do?

Makes GRPO runtime contracts fail closed and adds a production-scale GLM-5.1 Router Replay (R3) gate with explicit numerical, trace-integrity, learning-signal, and completion checks.

The MoE routing correction itself already exists upstream. The GLM recipe activates it with:

policy:
  router_replay:
    enabled: true

This PR does not introduce a new routing algorithm. It makes the existing integration safe to operate and possible to certify: entrypoints reject incompatible runtime combinations before Ray startup, SingleController rejects unsupported no-op controls, trace validation selects the actual transport contract, and the test harness can no longer report incomplete or stale runs as successful.

Issues

No linked issue.

Why Router Replay is needed

After refit, vLLM and Megatron have the same policy weights, but an MoE model still has a discrete decision inside each layer. Small numerical differences in independently computed router scores can change the final top-k expert when candidates are close. The token then executes a different expert network, creating train-versus-rollout logprob mismatch that is unrelated to a policy update.

R3 records only the expert IDs selected during vLLM generation and replays those IDs in Megatron. Megatron still computes current router scores, gathers current probabilities for the selected experts, executes current expert weights, and computes gradients normally. R3 therefore fixes expert-selection consistency; it does not remove a forward pass or freeze the router.

flowchart LR
    Refit["Refit the same policy weights"]

    subgraph Rollout["Rollout / behavior policy: vLLM"]
        VHidden["Token hidden states"]
        VScores["vLLM router scores"]
        VTopK["vLLM top-k selection"]
        VExperts["Selected experts execute"]
        VLogprobs["generation_logprobs"]
        VRoutes["routed_experts<br/>token x MoE layer x top-k"]

        VHidden --> VScores --> VTopK --> VExperts --> VLogprobs
        VTopK --> VRoutes
    end

    subgraph Transport["NeMo RL rollout payload"]
        Messages["tokens + masks + generation_logprobs"]
        Routes["routed_experts<br/>encoded and batch-aligned"]
        ReplayBuffer["ReplayBuffer or TransferQueue"]

        Messages --> ReplayBuffer
        Routes --> ReplayBuffer
    end

    subgraph Alignment["Megatron route preparation"]
        Decode["Decode route tensor"]
        ParallelMap["Map global MoE layers to local PP layers<br/>and slice token rows for TP/SP/CP"]
        Validate{"Every route complete and valid?"}
        Install["Install expert IDs on each<br/>model-owned RouterReplay instance"]
        Fallback["All -1 row only:<br/>compute Megatron top-k for that row<br/>and emit fallback telemetry"]

        Decode --> ParallelMap --> Validate
        Validate -- "yes" --> Install
        Validate -- "missing sentinel" --> Fallback --> Install
    end

    subgraph Training["Policy backend: Megatron"]
        MWeights["Current Megatron policy weights"]
        MScores["Current Megatron router scores"]
        FixedTopK["Use replayed expert IDs<br/>instead of a new top-k decision"]
        Gather["Gather current scores for<br/>the replayed experts"]
        Prev["Prev-logprob forward"]
        PrevLP["prev_logprobs"]
        Train["Training forward"]
        Loss["GRPO loss / importance ratio"]
        Backward["Backward or activation recomputation<br/>replays the same expert IDs"]
        Update["Update router and expert parameters"]

        MWeights --> MScores --> FixedTopK --> Gather
        Gather --> Prev --> PrevLP
        Gather --> Train --> Loss --> Backward --> Update
        Install --> FixedTopK
        FixedTopK --> Backward
    end

    Refit --> VHidden
    Refit --> MWeights
    VLogprobs --> Messages
    VRoutes --> Routes
    ReplayBuffer --> Decode
    VLogprobs --> Compare["Generation-KL and<br/>token probability-error metrics"]
    PrevLP --> Compare

    Independent["Without R3: Megatron independently<br/>selects top-k; a boundary tie can choose<br/>different experts despite matched weights"]
    MScores -. "R3 disabled" .-> Independent -. "artificial mismatch" .-> Compare

Loading

Runtime evidence

The completed 80-node GLM-5.1 run used TP2/PP18/EP16 Megatron training, TP32/EP32 vLLM generation, legacy async GRPO, and ten training steps.

Measurement	R3 disabled control (3147936)	R3 enabled (3171492)
Generation-KL mean	approximately 0.00250	0.0003876
Generation-KL range	0.0022966-0.0027089	0.0003615-0.0004061
Valid tokens	3,858,221	1,291,712
abs(delta log p) > 0.5	7,873 (0.204%)	4 (0.000310%)
abs(delta log p) > 1.0	570 (0.0148%)	0
Maximum abs(delta log p)	37.7257	0.676

Job 3171492 completed all ten training steps and printed Async GRPO training complete!. Its final wrapper status was a false red: the old trace checker unconditionally required SingleController/TransferQueue events, although this recipe deliberately used legacy async GRPO and its in-memory ReplayBuffer. The 269,952 route records did contain assignments, replay actions, forward verification, and context-parallel identity evidence.

That run is strong routing-correctness evidence but not yet the final learning-quality gate. Only five of ten batches had nonzero reward/advantage/loss because 94.5-100% of responses hit the 1024-token generation cap. The revised recipe increases the envelope to 2048 total / 1536 generated tokens and retains strict requirements instead of weakening them.

Changes

Runtime and configuration contracts

• Reject SingleController schemas passed to the legacy entrypoint and legacy async plus TransferQueue/data-plane mixtures before Ray setup.
• Reject unsupported backends, non-async engines, and SingleController controls that would otherwise parse but have no effect.
• Apply the configured GRPO advantage-clipping range in the SingleController training path.
• Keep the legacy async_grpo block out of SingleController setup.

Trace and completion contracts

• Select an explicit legacy-async or transfer-queue R3 trace schema.
• Require complete per-stage forward verification and context-parallel identity coverage.
• Require TransferQueue producer/fetch integrity only on the TransferQueue path, and reject mixed events on legacy async.
• Centralize test-suite completion handling: a direct or final run must record train/loss through MAX_STEPS; intermediate chained runs remain valid; original nonzero exit codes are preserved.

GLM production gate

• Isolate artifacts by Slurm job ID so stale traces, metrics, or terminal markers cannot certify a new allocation.
• Emit machine-readable terminal failures with the failing phase.
• Require at least 8/10 learning-signal steps and 8/10 nonzero-loss steps.
• Require mean truncation below 0.9, no valid-token error above 1.0, and a fraction below 1e-4 above 0.5.
• Require generation KL below 0.001, valid R3 route traces, and a clean terminal artifact.
• Document the end-to-end Router Replay data and gradient flow.

Usage

Submit the production-scale gate through the checked-in launcher:

bash infra/slurm/cscs/autoresearch/submit_glm51_r3_10step.sh

The launcher preserves the complete GLM checkpoint and conversion cache and uses the existing reservation contract.

Validation completed

• Current HEAD: python -m compileall for changed Python paths.
• Current HEAD: bash -n for changed launch/test shell paths.
• Current HEAD: git diff --check.
• Current HEAD: 13 focused trace/completion tests passed.
• Earlier exact-image component runs passed: 123 entrypoint/config tests, 32 SingleController setup tests, 13 architecture-invariant tests, 7 GLM harness tests, and the focused advantage-clipping test.
• 80-node R3-on training body completed ten steps and materially reduced KL/error tails.

Before this draft is ready for review

• Run the latest dependency-bearing changed-path suite in the exact image; host Python does not include Ray.
• Complete the revised 2048/1536-token 80-node gate with at least 8/10 learning-signal and nonzero-loss steps.
• Attach the final isolated terminal, metrics, and trace artifacts to this PR.
• Read and followed the contributor guidelines.
• Added focused tests and documentation.

Additional information

• Router Replay implementation provenance: NVIDIA-NeMo/RL PR #2590. The newer SingleController + TransferQueue integration followed in NVIDIA-NeMo/RL PR #3378.
• The already-merged optimizer checkpoint memory work in PR fix(checkpointing): make GLM optimizer recovery memory-safe #26 is independent of this routing/config/harness PR.

[Alvorecer721]

/ok to test 1c7d48f

[Alvorecer721]

Superseded by #28, which carries the fail-closed runtime-contract and Router Replay gates on the upstream-a952 line (see its handoff for the merge-loss scan). Closing; the branch is retained.