[faresobeid]

Merge order: land this AFTER the open PRs that edit trainer/rl/train.py (#3767, #3769, #3793, #3794, #3796, #3819) and #3770 (SFT resume), and after the sibling refactors #3844 (gradpath) and #3843 (moe). I will rebase this PR on top of them.

Summary

The RL and SFT train loops repeated the same setup and teardown code. This PR moves the shared parts into plain helpers in trainer/utils.py. The micro-step loops stay separate, and the gradient-manager calls are unchanged.

• setup_trainer(config, seq_len=None): torch distributed, CPU optimizer threads (full offload), matmul precision, resolve_ep, get_parallel_dims. Replaces setup_full_cpu_optimizer_offload. setup_torch_distributed stays because tools/convert_dcp_to_bf16.py uses it.
• resolve_resume_step(resume, ckpt_dir): the checkpoint step to resume from (same rules as before).
• start_trace(trace_path) / save_trace(prof, trace_path): the torch profiler start and export. I used two functions, not a context manager, so the long loop bodies are not re-indented (that would conflict with every open PR on these files).
• log_step_metrics(step, metrics, ...): adds the shared perf / optim / checkpoint-disk metrics and logs the step in one monitors.log call.
• setup_optimizer(config.optim, model, parallel_dims, config.model): reads the offload settings and builds the full-offload dtype policy itself. Both loops made the same 10-line call before.

Not done: CP rank/group and the final checkpoint stay inline. Each one is 2 to 4 lines that differ a little between RL and SFT, so a helper would not save code.

Behavior change

Metric keys and values are the same. The only change: each trainer step is now one row in metrics.jsonl and one wandb.log call, not 5 rows (RL) or 6 to 7 rows (SFT). The dashboard reads metrics by key, so it is not affected. Also, if a monitor fails to log a step, the whole step row is lost now, not just one group of metrics.

What I ran

• pytest tests/unit/orchestrator tests/unit/test_configs.py -m "not gpu": 291 passed.
• pytest tests/unit/train -m "not gpu" --ignore=tests/unit/train/models: 98 passed. 1 failed: test_cpu_adam.py::test_native_cpu_adamw_matches_fused_torch_and_preserves_gradients needs ninja to build the C++ extension, and this machine does not have it. That test does not touch code changed here.
• ruff check and ruff format --check on the changed files: clean.
• A throwaway equivalence script (not in the repo). It runs the old and the new per-step logging blocks of both loops with a recording monitor. It checks that the union of the old rows equals the new single row, with the same keys and values: RL with/without grad norm and mismatch KL, and SFT with/without MoE stats and with one or several data sources (20 cases). It also checks resolve_resume_step against the old inline code for all resume combinations, and it runs a smoke call of the new setup_optimizer signature (AdamW, with and without optim_cpu_offload).
• Not verified: a real GPU run of either loop.

GPU checks to do

• RL debug run (for example the reverse-text example) with trace_path set: check that the trace file is written and the W&B / metrics.jsonl keys are unchanged.
• SFT debug run with --resume from a checkpoint: check that it resumes at the right step.
• A full-offload run (model.full_offload, e.g. the GLM-5.3 setup) for one or two steps: check that the optimizer is built with the dtype policy and that NUMA binding / thread setup still logs.

Conflicts

• chore: remove Kubernetes deployment path (Helm chart, trainer metrics server) #3767 (removes metrics server, rl/train.py): textual conflict next to the new log_step_metrics call and in the imports. Easy to resolve.
• chore: remove legacy dead code #3769 (resolve_latest_ckpt_step import in rl/train.py): textual conflict. This PR removes that import.
• chore: remove legacy config options #3770 (optim/__init__.py imports, SFT resume load): merges cleanly, but it touches the same files.
• feat: Total Router Recall (replay sampler MoE routing weights) #3793, feat: exact score centering over the sampling mask (arXiv:2609.20807) #3794, feat(trainer): Composer 2-style plausibility filter for router replay #3796, feat: by-handle payload segments: reader, packing, CP-window reads, GC #3819 (rl/train.py): merge cleanly with this branch, but they touch the same file.
• No other listed PR touches these files.
• Sibling refactors: refactor(trainer): one gradient manager for GPU and offload gradients #3844 (refactor/trainer-gradpath) conflicts in optim/__init__.py, utils.py and both loops; refactor(trainer): add iter_moe_blocks for the MoE layer lookup #3843 (refactor/trainer-moe) conflicts in both loops; refactor/trainer-vlm conflicts in rl/train.py.

---

Note

Low Risk
Refactor with equivalent setup and metric semantics; the main operational change is atomic per-step logging (a failed log drops the whole step row).

Overview
Pulls duplicated RL/SFT trainer bootstrap and per-step logging into trainer/utils.py, without changing the main training loops.

setup_trainer replaces inline distributed init, full-offload NUMA/thread setup, matmul precision, resolve_ep, and get_parallel_dims. resolve_resume_step, start_trace / save_trace, and log_step_metrics centralize resume resolution, Chrome trace export, and monitoring.

setup_optimizer now takes (optim config, model, parallel_dims, model_config) and reads CPU/full offload settings plus the full-offload dtype policy via get_full_offload_dtype_policy internally.

Behavior: metric keys/values stay the same, but each step emits one monitors.log row (perf, optim, disk, and loop-specific metrics together) instead of several separate writes per step.

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