[RFC] Shard-level P2P weight sync for non-colocate training

[CalvinXKY]

Summary

This RFC proposes shard-level P2P weight synchronization for slime's non-colocate training path, as an optional weight transfer mode under --update-weight-mode full.

Core idea: each training TP rank converts and sends only its local shard, using dist.send/recv to write directly into the matching SGLang TP rank, instead of gathering full weights on rank 0 and NCCL-broadcasting them.

Three-phase rollout:

Phase	Goal	Relationship to the current path
Phase 1	Land basic functionality, validate the approach, coexist with broadcast	Opt-in via --use-p2p-weight-update; automatic fallback when preconditions fail
Phase 2	Broaden model support; reduce bucket / control-plane overhead	Still opt-in; expand the supported surface area
Phase 3	Enable P2P by default for supported configs; broadcast becomes fallback	Change the default path; keep an escape hatch

Motivation

Current path (UpdateWeightFromDistributed)

For non-colocate + --update-weight-mode full + --update-weight-transport nccl, weight sync roughly follows:

1. all_gather each parameter across the training TP group to reconstruct full tensors
2. PP source rank performs Megatron → HF conversion on full weights
3. rank 0 broadcasts full chunks to SGLang engines over NCCL
4. SGLang load_weights

Main issues:

Issue	Impact
Training TP0 must all_gather and hold broadcast buffers	Higher sync latency; rank-0 memory skew
Non-TP0 ranks are underused during gather/broadcast	Uneven compute and memory utilization
Each inference TP rank receives full weights, then shards internally	High traffic; misaligned with training shard layout
Bucketization further amplifies rank-0 pressure	Especially visible for large models / high TP

Expected benefits

Dimension	P2P expectation
Traffic	Each rank sends only its shard; no full broadcast
Trainer memory	Less rank-0 full-tensor gather buffering
Sync latency	Shorter update_weights when TP sizes align
Correctness	Numerically aligned with broadcast (no logprob regression)

Proposed Architecture

Single sync flow (Phase 1)

1. Vocab params (embed / lm_head): small TP all_gather → strip padding → slice SGLang-aligned shards
2. Other dense params: shard-level Megatron→HF conversion (no all_gather) → send in buckets
3. Expert params (MoE, Phase 2+): send Megatron shards directly (Phase 1 falls back to broadcast for MoE overall)
4. Per bucket: gather per-TP metadata → rank-0 HTTP notify SGLang → NCCL barrier → dist.send/recv → wait for SGLang load to finish

Preconditions (Phase 1)

Requirement	Notes
Non-colocate	Colocate uses CUDA IPC; out of scope
--update-weight-mode full	Mutually exclusive with delta sync
--megatron-to-hf-mode bridge	P2P relies on Bridge weight layout
Megatron TP == SGLang TP	Must hold for every rollout engine
SGLang presharded recv support	Extend SGLang update_weights_from_distributed
Shard-level conversion for the model	Phase 1 target: Qwen2 / Qwen3 dense

When any precondition fails, automatically fall back to NCCL broadcast with a clear rank-0 log line. Existing users are unaffected.

Phase 1 — Basic functionality, validation, coexistence with broadcast

Goal: land P2P with minimal intrusion; do not change default behavior.

Deliverables

• New UpdateWeightFromDistributedP2P class alongside UpdateWeightFromDistributed
• CLI flag: --use-p2p-weight-update
• Startup precondition checks and automatic fallback
• Shard-level Megatron→HF conversion (convert_shard_to_hf; Phase 1: Qwen2/Qwen3 dense)
• SGLang-side patch: presharded recv, per-TP metadata dispatch, barrier sync
• User docs and smoke scripts (Qwen3-4B TP=4 as the reference config)

Acceptance criteria

1. Functionality: Qwen3-4B TP=4 smoke (5 rollout steps) completes without NCCL hang or crash
2. Performance: same config, broadcast vs P2P — steady-state update_weights for P2P ≤ broadcast (TP-aligned, single engine)
3. Correctness: train_rollout_logprob_abs_diff on par with broadcast; no meaningful regression
4. Coexistence: behavior unchanged when the flag is off; fallback path works end-to-end
5. Observability: per-phase / per-bucket timing logs or metrics

Out of scope (Phase 1)

• Shard conversion for MoE / DeepSeek / GLM / VLM, etc.
• Megatron TP ≠ SGLang TP
• Changing the default broadcast path
• Integration with --update-weight-mode delta
• Colocate / vLLM backend

Phase 2 — Broader model support and control-plane optimizations

Goal: expand the set of models where P2P is usable, and reduce fixed per-bucket overhead so more production configs benefit after opt-in.

2.1 Model adapters (shard-level Megatron→HF)

Suggested priority:

Priority	Models	Complexity	Notes
P0	Llama / GLM4 dense	Low	Close to Qwen2 structure
P1	Qwen3 MoE / GLM4 MoE	Medium	Dense via shard convert; experts may be sent as raw Megatron shards
P2	DeepSeek V3 / Qwen3-Next / Qwen3.5	High	MLA, linear attention, MTP, indexer reordering
P3	Qwen3-VL / MiMo / GPT-OSS / MiniMax	Medium–high	Multi-tower / special naming

Each item needs: shard-conversion unit tests + 5-step smoke + weight / logprob alignment vs broadcast.

2.2 Performance and architecture improvements

The broadcast path already uses bucketing. If P2P keeps a similar model, fixed per-bucket cost (HTTP + barrier + load_weights) remains a major residual expense:

Improvement	Expected benefit	Notes
Larger / adaptive update_weight_buffer_size	Fewer buckets	Strong on small dense models; MoE must balance peak memory
Single HTTP request for merged per-TP metadata	Lower lock + Ray overhead	One request carrying per-TP tensor counts
Async dist.isend + batched wait	Less Python loop overhead	Replace per-tensor synchronous send
Convert / send pipelining	Overlap compute and comms	Careful barrier ordering required
Fewer per-bucket loads on SGLang	Less wait time	May need SGLang protocol changes
Multi-engine send optimization	Multi-engine setups	Avoid bucket count × engine count scaling

2.3 Observability

• Promote timing / metrics into formal observability
• Document bucket tuning (buffer size vs TP vs engine count)

Phase 2 acceptance criteria

• At least Llama + Qwen3 MoE usable under P2P with aligned TP
• TP=4, single engine: measurable steady-state gain vs broadcast
• Multi-engine: after bucket optimizations, P2P is not materially slower than broadcast
• Documented fallback matrix (model × TP × engine × bridge/raw)

Phase 3 — P2P enabled by default

Goal: when all preconditions are met, P2P becomes the default path for --update-weight-mode full + nccl; broadcast remains fallback.

Planned changes

1. Default on: non-colocate + bridge + supported model + aligned TP → P2P automatically
2. Explicit opt-out: add --no-use-p2p-weight-update or deprecate the positive flag (TBD)
3. Code cleanup: factor shared bucket / lock / metrics logic between P2P and broadcast
4. Upstream SGLang patch: land presharded recv in SGLang mainline
5. CI: add P2P smoke to CI matrix (at least Qwen3-4B TP4)

Migration and compatibility

User scenario	Phase 3 behavior
Explicit --use-p2p-weight-update	Unchanged
Flag not set, preconditions met	P2P by default
MoE / unsupported model / TP mismatch / raw mode	Broadcast + rank-0 log
delta / disk / colocate	Unchanged

Phase 3 acceptance criteria

• Default supported config on main uses P2P
• Full fallback regression suite passes
• Docs updated from "opt-in" to "default with fallback"
• At least one production-scale MoE 50-step run: P2P vs broadcast

Open Questions

1. Should Phase 3 default-on roll out via a model allowlist (Qwen/Llama first, then MoE), or switch in one step?
2. Should the SGLang patch live in slime, or be contributed upstream?
3. Is P2P needed for --megatron-to-hf-mode raw, or permanent fallback?
4. Should Phase 2 ship an official benchmark script (similar to delta_weight_sync examples)?
5. Are breaking SGLang protocol changes acceptable for bucket optimization (e.g. one HTTP call with full metadata)?

Current Phase 1 implementation

An initial Phase 1 implementation is open for review in PR #2146: feat(p2p): add shard-level weight update with automatic broadcast fallback.

What it includes:

• Opt-in --use-p2p-weight-update (requires --update-weight-mode full; default behavior unchanged)
• UpdateWeightFromDistributedP2P with automatic NCCL broadcast fallback and rank-0 [P2P] logging
• Qwen2 / Qwen3 dense shard Megatron→HF conversion
• sglang_p2p.patch applied at Docker build time (after existing SGLang patches)
• Docs: docs/en/advanced/p2p-weight-sync.md, docs/zh/advanced/p2p-weight-sync.md

Reported validation (see PR comments for details):

• Qwen3-4B TP=4 smoke on 8×A100 (exit_code=0); steady-state update_weights ~0.33–0.40s
• Qwen3-8B TP=4, 50 rollout steps: steady-state perf/update_weights_time ~0.48s (P2P) vs ~0.76s (NCCL broadcast), ~1.56× faster on the weight-sync phase

This PR is intended to land Phase 1 as an opt-in path; Phases 2 and 3 in this RFC remain follow-up work (broader model support, bucket/control-plane optimizations, and optional default-on behavior).