fix(trainer): contiguous routed_experts slice for glm4_moe under compile

[S1ro1]

Problem

With router replay (trainer.enable_router_replay), the trainer replays inference's expert choices by passing routed_experts (shape [batch, seq, num_layers, top_k]) into the MoE forward, slicing one layer at a time:

routed_experts_layer = routed_experts[:, :, layer_idx, :]

That slice is a non-contiguous view (dim-1 stride = num_layers * top_k). Under torch.compile, the inductor MoE kernel is compiled expecting a contiguous input, so at runtime it asserts and the trainer crashes on every rank:

AssertionError: expected size 8173==8173, stride 8==368 at dim=1
  assert_size_stride(primals_2, (1, 8173, 8), (24061312, 368, 1))
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: prime_rl.trainer.rl.train FAILED

(368 = 46 layers × 8 top_k for GLM-4.5-Air.) Hit on a GLM-4.5-Air (glm4_moe) RL run with enable_router_replay=true + cp=8 + compile.

Fix

Make the per-layer slice contiguous before it enters the compiled kernel:

routed_experts_layer = routed_experts[:, :, layer_idx, :].contiguous()

Validation

GLM-4.5-Air RL run reached step 0 with no stride assert after this change (Mismatch KL 0.0002, Peak Mem 105.6 GiB), and rollouts trained normally. Scoped to glm4_moe; the same slice pattern exists in other MoE models but only glm4_moe has been observed to trip the assert.

🤖 Generated with Claude Code

---

Note

Low Risk
Single-line tensor layout fix on an optional router-replay path; no auth, data, or routing logic changes, with a small per-layer copy cost when replay is on.

Overview
Fixes router replay crashes on GLM-4 MoE when torch.compile is enabled: the per-layer routed_experts[:, :, layer_idx, :] slice is now .contiguous() before each decoder layer’s MoE forward.

Without this, that slice stays a strided view over [batch, seq, layers, top_k] (stride on the token dimension scales with layers × top_k), which trips the compiled inductor MoE kernel’s assert_size_stride and fails training on every rank. The change is limited to modeling_glm4_moe.py; other MoE models still use the same slicing pattern unchanged.

^{Reviewed by Cursor Bugbot for commit 7e50017. Bugbot is set up for automated code reviews on this repo. Configure here.}