[Not maintained · vLLM 0.16.0 only · kept for reference] Fix trainer GPU0 memory skew: replace NcclBridge with in-process trainer_init

[CalvinXKY]

Status: Legacy / reference only — targets vLLM 0.16.0. The project has moved to newer vLLM; this PR is not planned for merge or further iteration on main, but is kept open for anyone still on 0.16.0 who needs the in-process trainer_init weight-sync path (see test results below).

Purpose

Follow-up to vLLM native weight sync integration (PR #3). Issue #4 reports that the previous _NcclBridge subprocess on trainer rank 0 / GPU 0 pins ~10 GiB extra memory vs other trainer ranks and adds ~30 s to the first weight sync (subprocess spawn + second CUDA context).

This PR is step 1/N of the performance work for #4: call vLLM’s NCCLWeightTransferEngine.trainer_init() in the Megatron train actor process (SkyRL-style), and send weights via trainer_send_weights() on the existing packed bucket path—no torch.multiprocessing child process.

Later steps (not in this PR) may further shrink the remaining rank-0 vs rank-1 gap (packed buffer sizing, optional teardown/recreate of PyNccl after each sync, docs).

What’s included

• UpdateWeightFromDistributed (vLLM path)
  • connect_rollout_engines_from_distributed: NCCLWeightTransferEngine.trainer_init(...) in-process on trainer rank 0.
  • update_weights_from_distributed: NCCLWeightTransferEngine.trainer_send_weights(..., packed=...) (unchanged packed bucket semantics; MoE expert buckets still packed=False).
  • disconnect_rollout_engines_from_distributed: notify rollout engines only; no bridge shutdown().
• Removed: _NcclBridge, _nccl_bridge_worker, and related multiprocessing spawn/teardown.
• Unchanged: SGLang path still uses init_process_group + dist.broadcast; _is_vllm_backend() routing; rollout HTTP metadata + engine init_weight_transfer_engine flow.

Test plan

Same environment as #4 and PR #3:

• Model: Qwen3-4B
• GPUs: 8× NVIDIA A100-SXM4-80GB (4 Megatron actor + 4 vLLM rollout, tensor-model-parallel-size 2 on actor)
• Flags: --rollout-backend vllm --vllm-weight-sync-mode native
• Compare: this branch vs previous NcclBridge subprocess implementation on the same script and hardware.

Docker / install (same as PR #3)

docker pull slimerl/slime:latest

docker run -itd --rm --gpus all --network=host --ipc=host --ulimit memlock=-1 \
  --ulimit stack=67108864 -v /data/nfs_87:/data/nfs_87 \
  --name slime-dev slimerl/slime:latest bash

pip install vllm-router ray
pip uninstall slime -y

cd vime
pip install -e . --no-deps

E2E: Megatron + vLLM native sync (this change)

From repo root:

bash run_scripts/qwen_4b.sh

Key flags (also in the script):

• --rollout-backend vllm
• --vllm-weight-sync-mode native
• 4 actor GPUs + 4 rollout GPUs, GRPO, global-batch-size 256, etc.

Optional: minimal coexistence check (no full training)

bash tools/verify_inprocess_trainer_init/run.sh

Validates trainer_init + trainer_send_weights in-process alongside an existing torch.distributed NCCL group (coexist mode).

Test results

Hardware: 8× A100-80GB, Qwen3-4B, run_scripts/qwen_4b.sh, native vLLM weight sync.

Weight sync latency (perf/update_weights_time)

Step	Before (NcclBridge subprocess)	After (in-process trainer_init)
0 (first sync, includes init)	34.4 s	11.8 s
1	2.7 s	1.9 s
2	3.0 s	1.9 s
3	2.2 s	2.1 s
4	3.1 s	1.1 s
5	6.3 s	1.4 s

• First sync ~3× faster (no subprocess spawn / second CUDA context on rank 0).
• Steady-state sync mostly ~1.1–2.1 s over 6+ steps (one outlier ~4.4 s when sync overlapped with rollout start; completed successfully).

Trainer GPU memory (rank 0, used_GB after update_weights)

Phase	Before	After
First sync (post-init)	~28.7	~28.1
Steady state (steps 1–5)	~39.6 (plateau)	~38.2 (flat across steps)

• No upward drift in post-sync steady state across 6 training steps on the new path.
• Training peak before clear_memory (~50.6 GiB on rank 0 after step ≥1) is in line with the old run; clear_memory still returns rank 0 to ~30.7 GiB before the next sync.

Correctness / stability

• 6+ full train → sync → rollout cycles with no OOM, NCCL errors, or weight-sync failures.
• Logs show vLLM in-process weight transfer on connect; no NcclBridge ready.
• Rollout throughput ~1280–1440 tokens/GPU/s; training metrics and rollout continue normally.

Rank imbalance (remaining)

	Rank 0 (sender)	Other trainer ranks	Gap
After sync (steady)	~38.2 GiB	~29.3 GiB	~9 GiB

Removing the bridge subprocess eliminates the ~10 GiB child-process component from #4; the remaining gap is expected from rank-0 HF conversion buckets + in-process PyNccl / packed buffers—target for a follow-up [2/N] PR.

Closes #4 partially (subprocess + first-sync latency); rank-0 gap reduction tracked as follow-ups above.

[gemini-code-assist]

Code Review

This pull request simplifies the vLLM weight transfer process by removing the _NcclBridge subprocess and implementing in-process transfer using NCCLWeightTransferEngine. This change reduces architectural complexity by running the transfer logic directly in the Megatron trainer. A review comment suggests using a generator expression when preparing tensors for transfer to avoid a significant memory spike on rank 0, ensuring that only one extra contiguous copy of a tensor exists in memory at a time.

[gemini-code-assist]

Building a list of contiguous tensors for the entire bucket creates a significant memory spike on rank 0, as it effectively doubles the memory required for the weights being transferred in that bucket. Since trainer_send_weights accepts an iterator, you can use a generator expression to yield contiguous tensors one by one. This ensures that only one extra contiguous copy exists in memory at any given time, which aligns better with the goal of reducing GPU 0 memory skew.

        named_gpu_iter = ((name, (param.data if hasattr(param, "data") else param).contiguous())
                          for name, param in converted_named_tensors)
        NCCLWeightTransferEngine.trainer_send_weights(
            iterator=named_gpu_iter,

[CalvinXKY]

[CalvinXKY]