[aoshen02]

Purpose

Speed up the receiver side of RL weight sync for DeepSeek-V4.1 (MegaMoE, 384 routed experts):

1. [Model Loader] Reload MegaMoE experts in place during layerwise reload. DeepGemmMegaMoEExperts loaders already write the final kernel layout and the layer has no post-load processing, but layerwise reload still put its params on meta, buffered every expert tensor, dry-ran each loader under a copy counter, materialized a new layer and copied it back. A reload_in_place layer opt-in leaves such layers live, so loaders write the live (CUDA-graph-captured) storage directly. Same loaders, same destinations, same bytes.

2. [DSV41] Index the routed-expert mapping in load_weights. Each expert tensor scanned all 1152 (expert, shard) mapping entries with a substring test (~40 µs/tensor, ~1 s per sync per rank at 384 experts). Look up only the entries keyed by the experts.<id>.w<k>. occurrences in the name; fall back to the full scan for other name forms.

3. [DSV41] Load runs of consecutive routed experts with one batched copy. Each routed-expert shard went through its own loader call (name mapping, ownership scan, scale packing, layout views, copy), ~0.2 ms per tensor, ~5 s of every sync per rank at full scale (96 local experts × 40 layers × 6 shards). load_weights now groups consecutive experts of one MegaMoE parameter and hands each run to DeepGemmMegaMoEExperts.load_routed_experts, which stacks the local ones and writes them with one batched copy using the same layout transforms. Bytes are unchanged; NVFP4 global scales and EPLB replica layouts keep the per-expert loader.

4. [DSV41] Collapse Engram huge pages only after loading the Engram tables. load_weights dropped the checkpoint page cache and MADV_COLLAPSEd every Engram table on every call. RL weight sync calls it once per IPC bucket and never sends the frozen tables, so each call paid ~5 ms for nothing. Do both only when the call loaded an Engram table: ~0.9 s per sync on a 1-node proxy.

5. [DSV41] Keep one expert run open per parameter shard. The weight-sync sender emits w1, w2, w3 of one expert before the next, so a single open run was closed by every tensor and (3) never batched (480 calls of 1 expert on the proxy). Keep one run per (parameter, shard) and flush them at the end of the call. This is what makes (3) actually batch.

With (1)–(2) the per-expert-tensor receive cost drops from ~0.6 ms to ~0.115 ms; (3)–(5) remove most of the remaining per-tensor overhead. Together with trainer-side changes (separate repo) a full DeepSeek-V4.1 weight update in colocated verl RL went from ~106 s to ~9 s.

Not a duplicate

Searched open PRs for reload_in_place, "layerwise reload MegaMoE", "deepseek_v41 load_weights expert mapping". Related but different: #57678 / #57733 (model-level post-load finalization on reload), #63119 / #60895 (MegaMoE for other models). None skip the layerwise reload buffering for layers without post-load processing or index the DSV41 expert mapping.

Test Plan

• New unit tests: tests/model_executor/model_loader/test_reload.py::test_reload_in_place_layer_writes_live_storage, tests/models/test_deepseek_v41_expert_mapping.py::test_expert_mapping_index_matches_linear_scan and ::test_routed_expert_shards_load_as_one_run_per_parameter_shard (wired into Basic Models (Other)), tests/kernels/moe/test_deep_gemm_mega_moe.py::test_deep_gemm_mega_moe_batched_load_matches_per_expert_loads (batched run writes the same bytes as per-expert loads, MXFP4 and NVFP4; already in the MegaMoE kernel step).
• pre-commit run on the changed files: passed.

Test Result

On the original base (vLLM nightly image + these files as an overlay, GB200): both new tests pass; test_reload.py 89 passed, 2 failed (Qwen3-0.6B engine start, Llama FP8-KV reload perplexity) — neither uses MegaMoE or reload_in_place, so the new code is a no-op there; treated as image-environment failures, not re-run on a clean baseline. The commits were then cherry-picked onto current main without conflicts in the source files; the unit tests have not been re-run on this exact rebase yet (CI will). Commits (3)–(5) were developed on a downstream DSV41 branch; cherry-picking them here only conflicted in tests (the batched-load test was taken without the unrelated tests next to it, and the run-grouping test moved into test_deepseek_v41_expert_mapping.py). Their two new tests have not been run on this rebase (CI will). pre-commit run on the changed files passed.

End-to-end (DeepSeek-V4.1, colocated verl + Megatron, sleep level 2, 16×GB200):

• Full vLLM parameter/buffer fingerprints after each sync identical to the previous path (4 ranks × 4 sleeps, 0 diffs).
• Train–inference rollout_probs_diff_max = 0 with batch invariance across steps; AIME eval unchanged.
• update_weights per step: ~106 s → 8.8–9.5 s (with the trainer-side changes). All five commits were in the overlay used for those 16-node runs; full-scale update_weights stays ~9 s total with all of them.

AI assistance

This PR was prepared with AI assistance (Claude Code). The human submitter is reviewing every changed line.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.