[faresobeid]

Router-replay routes arrive from the orchestrator in vLLM's wire dtype (uint8 for ≤256 experts, uint16 above). On main the trainer:

• widens the whole batch's routes to int32 on the host in DataLoader._micro_batch_to_tensor (4x the uint8 bytes, 2x the uint16 bytes);
• copies the full [1, seq, layers, top_k] tensor to the GPU and only then takes the CP chunk. The chunk is a view, so the full tensor stays allocated on the GPU;
• keeps the previous step's micro batches until dataloader.get_batch() returns, so two converted batches are alive while the next one is received.

This PR keeps routes in the wire dtype on the host, takes the CP chunk on the host, copies only that chunk to the GPU and widens it to int32 there. It also drops the step's micro batches after the micro-step loop. The model receives the same int32 tensor as before.

Memory

One trainer rank, 16 micro batches × 65,536 tokens (1M tokens), CP=4, two steps. Routes are the only per-token array that changes. Host peak is RSS above the received frame, sampled after each micro batch is converted. GPU is the peak allocation for the routes in one micro step (H200).

routes	host peak, main	host peak, this PR	GPU per micro step, main	GPU per micro step, this PR
46 layers × top-8, uint8 (GLM-4.5-Air)	3.46 GiB	0.57 GiB	92 MiB	30 MiB
48 layers × top-10, uint16 (Qwen3.8-Flash-Next)	4.91 GiB	1.10 GiB	120 MiB	46 MiB

Host memory is per rank. Every CP rank of a data rank still receives the full payload, so a node holds 8x these numbers.

Overlap

• perf(trainer): zero-copy top-logprob decode, row-copy top logprobs; log load timing #4192 also keeps uint8 routes as a view of the frame and widens them on the GPU. uint16 still widens on the host there, and it does not move the CP shard or free the previous batch. The two PRs touch the same lines in data.py/train.py.
• feat: pass router-replay ids and sampling masks by handle on multi-node runs #3902 already reads only the rank's CP chunk on its by-handle path (multi-node). This PR covers the inline path and feat: keep training arrays with the env servers #4196's fetched arrays, which both go through _micro_batch_to_tensor.

---

Note

Low Risk
Behavior-preserving memory optimization for MoE router replay; model still receives int32 routed experts after CP shard. Minor risk only if CP sharding or dtype handling on CPU diverged from the old GPU path.

Overview
Router-replay routed_experts stay in vLLM wire dtype (uint8/uint16) on the host longer, instead of widening to int32 in DataLoader._micro_batch_to_tensor.

In the training loop, routes are no longer copied to GPU before context-parallel (CP) sharding. CP slicing runs on the host; only the rank’s chunk is moved to CUDA and cast to int32 right before forward—same int32 input to the model, but less host and GPU memory (especially when CP>1, since the full sequence no longer has to live on GPU as a view).

After the micro-batch forward/backward loop, the step del micro_batches, micro_batch drops the converted batch so the next receive/convert does not keep two batches alive at once.

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