[bugfix] Fix distributed update weights for pipeline parallel

[knlnguyen1802]

Summary

This PR fixes a hang in non-colocated full NCCL weight sync when the Megatron actor uses pipeline parallelism, for example --tensor-model-parallel-size 1 --pipeline-model-parallel-size 2.

The root cause is that Vime treated every DP=0, TP=0 pipeline stage as a vLLM weight-sync source rank. With PP>1, multiple trainer ranks concurrently initialized separate trainer-to-vLLM NCCL transfer groups against the same vLLM rollout engines. Current vLLM weight transfer keeps one active receiver communicator, so later initialization could overwrite earlier initialization and leave trainer and receiver waiting on different NCCL groups.

The fix keeps the existing fast path for PP=1. For PP>1, Vime now sends one pipeline stage at a time inside a single vLLM weight-update session. Only the active pipeline stage initializes the vLLM receiver communicator and broadcasts its local weights. Inactive pipeline stages only join the required synchronization barriers.

Problem

The default non-colocated full weight sync path is:

1. MegatronTrainRayActor selects UpdateWeightFromDistributed for --update-weight-mode full --update-weight-transport nccl.
2. UpdateWeightFromDistributed.connect_rollout_engines marks a trainer rank as a source when data_parallel_rank == 0 and tensor_model_parallel_rank == 0.
3. Each source rank creates a trainer-to-vLLM NCCL transfer group through connect_rollout_engines_from_distributed.
4. Each bucket is sent through NCCLWeightTransferEngine.trainer_send_weights, while vLLM receives through its weight-transfer engine.

This works for TP=2, PP=1 because there is only one DP=0, TP=0 source rank.

It fails for TP=1, PP=2 because both pipeline stages have TP=0:

Global rank	TP rank	PP rank	Old source predicate
0	0	0	source
1	0	1	source

The old code named those groups vime-pp_0 and vime-pp_1, but the current vLLM-side wrapper does not actually use group_name to keep multiple receiver communicators. VLLMEngine.init_weights_update_group and VLLMEngine.update_weights_from_distributed pass the request to vLLM's weight-transfer endpoints, whose dense NCCL backend stores a single active model_update_group and receives from source rank 0.

That makes the old multi-PP behavior unsafe:

1. PP stage 0 initializes receiver communicator A.
2. PP stage 1 initializes receiver communicator B.
3. vLLM remembers only the latest active receiver communicator.
4. A trainer rank can broadcast on communicator A while vLLM waits on communicator B.
5. NCCL waits forever because producer and consumer do not join the same group.

The rollout engine lock did not prevent this because it only serialized bucket broadcasts. It did not serialize or isolate communicator initialization.

Goals

• Support non-colocated full NCCL weight sync when Megatron uses PP>1.
• Preserve existing behavior and performance for PP=1.
• Avoid requiring vLLM changes or a new receiver-side communicator map.
• Keep raw Megatron-to-HF and Megatron-Bridge export paths synchronized correctly.
• Make the implementation easy to reason about and test with CPU unit stubs.

Design

1. Keep the PP=1 fast path unchanged

When mpu.get_pipeline_model_parallel_world_size() == 1, the implementation keeps the old persistent transfer group behavior:

1. The single DP=0, TP=0 source rank connects rollout engines during connect_rollout_engines.
2. update_weights starts one vLLM weight-update session.
3. The source rank sends all buckets through the existing _send_weights path.
4. Vime synchronizes CUDA once after all buckets.
5. The vLLM weight-update session is finished.

This means the common non-PP case has the same process-group lifecycle and the same number of NCCL transfer groups as before.

2. Defer vLLM NCCL group creation for PP>1

When PP>1, connect_rollout_engines records the rollout engines, lock, engine GPU counts, PP world size, and local source eligibility, but it does not immediately create a vLLM NCCL transfer group.

This prevents multiple PP source ranks from racing to initialize receiver communicators before any bucket transfer starts.

3. Send one pipeline stage at a time

update_weights still owns the outer vLLM lifecycle:

1. Pause generation.
2. Flush vLLM cache.
3. Start one vLLM weight-update session.
4. Send weights.
5. Finish the vLLM weight-update session.
6. Resume generation.

For PP>1, sending weights is now staged:

for active_pp_rank in range(pp_world_size):
    if local rank is DP=0, TP=0, and PP=active_pp_rank:
        initialize the single vLLM NCCL receiver communicator
        send this PP stage's weight buckets
        cuda synchronize
    else:
        skip local weight iteration and only join matching barriers

Each active PP stage reuses the existing connect_rollout_engines_from_distributed and update_weights_from_distributed helpers. The difference is the timing: only one PP stage owns the vLLM receiver communicator at a time.

4. Preserve barrier symmetry

The raw and bridge weight-export paths have different barrier shapes:

• Raw path has one barrier after dense/non-expert chunks and one barrier after the optional expert pass.
• Bridge path has one barrier after the exported HF chunks.

Inactive PP stages must not iterate local weights for the wrong pipeline stage, but they still need to join the same number of barriers as the active stage. The fix adds an active-stage predicate and mirrors the correct barrier behavior for both paths.

Test Plan

Focused unit coverage should include:

• PP>1 connect defers vLLM NCCL group initialization instead of creating a group immediately.
• The staged send loop initializes only the active PP stage's vLLM group.
• Inactive raw-path PP stages join two barriers without iterating local weights.
• Inactive bridge-path PP stages join one barrier without exporting HF chunks.
• Existing packed and non-packed send behavior still calls NCCLWeightTransferEngine.trainer_send_weights.

Suggested commands:

pytest tests/utils/test_update_weight_from_distributed.py 

Runtime validation should include:

--tensor-model-parallel-size 2
--pipeline-model-parallel-size 1

to confirm the existing non-PP path is unchanged, and:

--tensor-model-parallel-size 1
--pipeline-model-parallel-size 2

to confirm staged vime-pp_0, then vime-pp_1, sync completes without the NCCL hang.

When changing TP/PP topology between runs, use a topology-compatible checkpoint or start from model weights. Resuming a distributed optimizer checkpoint saved with a different (TP, PP) layout can fail before weight sync starts.

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #33499690 | 📁 Comparing 05ff1a8 against latest (2702eef)

  🔍 Preview build  

1 file changed

± examples/qwen3-30B-A3B.html

[gemini-code-assist]

Code Review

This pull request updates the weight synchronization mechanism in UpdateWeightFromDistributed to support pipeline parallel (PP) configurations greater than one by sending weights one pipeline stage at a time. This is achieved by dynamically connecting and disconnecting rollout engines for the active stage and synchronizing across ranks using barriers. The review feedback suggests simplifying the active PP stage loop by only updating self._group_name on the active source rank, which avoids unnecessary state mutation on inactive ranks and simplifies the cleanup logic, along with a corresponding update to the unit tests.

[gemini-code-assist]

We can simplify the loop by only updating self._group_name when self._is_pp_src_rank is True. Since self._group_name is only used on the active PP stage's source rank, there is no need to temporarily set it on inactive or non-source ranks. This also completely eliminates the need to save and restore original_group_name in the finally block, avoiding unnecessary state mutation on non-source ranks.

        pp_rank = mpu.get_pipeline_model_parallel_rank()
        base_is_pp_src_rank = self._is_pp_src_rank
        try:
            for active_pp_rank in range(pp_world_size):
                self._active_weight_sync_pp_rank = active_pp_rank
                self._is_pp_src_rank = base_is_pp_src_rank and pp_rank == active_pp_rank
                if self._is_pp_src_rank:
                    self._group_name = f"vime-pp_{active_pp_rank}"
                    if self._model_update_groups is not None:
                        disconnect_rollout_engines_from_distributed(
                            self.args, self._group_name, self._model_update_groups, self.rollout_engines
                        )
                    self._model_update_groups = connect_rollout_engines_from_distributed(
                        self.args,
                        self._group_name,
                        self.rollout_engines,
                        engine_gpu_counts=self._engine_gpu_counts,
                    )
                    pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0)
                else:
                    pbar = None

                dist.barrier(group=get_gloo_group())
                self._send_weights(pbar)
                if self._is_pp_src_rank:
                    torch.cuda.synchronize()
                dist.barrier(group=get_gloo_group())
        finally:
            self._active_weight_sync_pp_rank = None
            self._is_pp_src_rank = base_is_pp_src_rank

[gemini-code-assist]

Since self._group_name is no longer modified on inactive ranks, we should update the test assertion to expect the original group name ("g") for the inactive PP stage.

Suggested change

	assert send_calls == [
	(0, True, "vime-pp_0", True, DummyGroup("vime-pp_0")),
	(1, False, "vime-pp_1", False, DummyGroup("vime-pp_0")),
	]
	assert send_calls == [
	(0, True, "vime-pp_0", True, DummyGroup("vime-pp_0")),
	(1, False, "g", False, DummyGroup("vime-pp_0")),
	]