Offload routed expert of RL

[hhaAndroid]

当输入给每个 trainworker 过多的数据后，原本不起眼的 routed expert 所占显存将会变得明显。本 PR 提供通过 offload_rollout_routed_experts 配置来将 routed expert 放置于 cpu 上，当需要时候再移到 gpu 上。

当前没有实现 prefetch 功能，原因是训练耗时占比本身不高。

[hhaAndroid]

@claude review

[claude]

Claude: The SFT path resets peak stats before each step (DEVICE_MODULE.reset_peak_memory_stats() at line 831 of the current main), so per-step peaks are meaningful. The RL fit() path has no such reset — these values are monotonically non-decreasing across micro-batches and across fit() calls, so after the first peak they'll be identical in every subsequent log line.

Consider adding DEVICE_MODULE.reset_peak_memory_stats() at the top of the fit() method (or before the training loop) if you want per-rollout peak visibility.

If the intent is to track the overall peak across the entire worker lifetime, this is fine as-is — but then it only needs to be logged once (e.g. at the end of fit()), not per micro-batch.

[claude]

Claude: Nit: since the tensor is unpinned, .to(device) is a synchronous copy regardless. If you later implement prefetch (as mentioned in the PR description), pin_memory() + to(device, non_blocking=True) on a CUDA stream would allow overlapping the transfer with prior-layer compute. Just flagging as a natural next step.

[claude]

Claude: ## Summary

Adds offload_rollout_routed_experts config flag to keep routed-expert tensors on CPU and move each layer's slice to GPU on-the-fly during forward. Also adds per-step max_memory / reserved_memory logging to the RL training loop.

ProduceBatchResult impact: not affected
RoutedExperts impact: ownership unchanged — tensors remain owned by the trainer worker; only device placement changes (CPU vs GPU when flag is True). No new retention paths introduced.

Issues

Warning

• [xtuner/v1/rl/trainer/worker.py:794-795] max_memory_allocated() / max_memory_reserved() are monotonically non-decreasing without a prior reset_peak_memory_stats() call. Unlike the SFT path (which resets per step), the RL path will log identical values after the first peak. See inline comment for details.

Nit

• Commit message "offload r3" doesn't follow the project convention [Tag] Short description (e.g. [Feature] Offload rollout routed experts to CPU).
• [xtuner/v1/module/decoder_layer/moe_decoder_layer.py:643-645] The CPU→GPU transfer is synchronous. pin_memory() + non_blocking=True would be the natural extension when prefetch is implemented.

Verdict

APPROVE — the core logic is correct and well-scoped. The offload flag cleanly integrates with the existing SequenceContext.to() and _pre_moe_forward paths. The memory metric issue is minor (monotonic values are still useful for cross-run comparison).