[faresobeid]

Summary

Add iter_moe_blocks(model) in trainer/model.py. It yields the MoE MLP of each decoder layer. It replaces the five copies of the decoder-layer MoE lookup in freeze_moe_router, apply_fp32_moe_router, get_full_offload_dtype_policy, apply_force_balanced_routing and get_load_balance_stats.

What changed

• New iter_moe_blocks next to the helpers that use it. The helpers stay in model.py, so no imports change.
• The four layer.mlp if hasattr(...) else layer.feed_forward ... lookups and the getattr(mlp) + hasattr(tokens_per_expert) scan in get_load_balance_stats now call iter_moe_blocks.
• Removed the stale "HF impl" comment in get_expert_load_stats. The HF fallback was removed in chore(trainer): remove HF model-implementation fallback (model.impl) #3773.

Behavior change

None. The feed_forward fallback is gone. Every prime-rl decoder layer uses mlp. get_load_balance_stats already looked only at mlp. Only MoE registers tokens_per_expert, so it selects the same layers as before.

Checks

• CPU equivalence vs main (throwaway script, not in the repo). Tiny Qwen3-MoE (mlp_only_layers=[1]) and GLM4-MoE (first_k_dense_replace=1) models with fixed router-count buffers. Compared is_tt_moe_model, get_load_balance_stats, get_global_moe_stats, compute_expert_load_stats, get_expert_load_stats and get_full_offload_dtype_policy, plus parameter dtype and requires_grad, and router fp32_gate/force_balanced after the three router helpers. All 16 outputs were bit-identical.
• pytest tests/unit/train tests/unit/orchestrator tests/unit/test_configs* -m "not gpu": 450 passed, 1 failed. The failure is test_cpu_adam.py::test_native_cpu_adamw_matches_fused_torch_and_preserves_gradients (ninja build), which fails the same way on main.
• ruff check and ruff format --check.
• GPU: reverse-text MoE RL smoke (samsja/mini-glm-moe, 20 steps, one node), main vs this PR. Both finished. Both froze 23 router params and ran 23 router gates in fp32. Max Vio mean 0.610 (main) vs 0.613 (PR), per-step range 0.57–0.66 in both. Routing Conf. 0.6514–0.6533 vs 0.6512–0.6536. Rollouts are sampled, so per-step values differ by run noise only.

[mikasenghaas]

lgtm