[0z5a]

Summary

Hardens the weight push from the Megatron trainer to the vLLM rollout engines, for both the IPC (colocated) and the NCCL (distributed) updater. This is the first of the smaller PRs announced in #478 after closing #479, and it is independent of FSDP2.

1. Time out weight-transfer requests.
   • New flags:
     • --update-weight-init-timeout (default 600 s) bounds /init_weight_transfer_engine;
     • --update-weight-request-timeout (default 1800 s) bounds /start_weight_update, /update_weights and /finish_weight_update.
   • Why: vLLM's NCCL trainer_init sends the init request from a side thread and joins it when it leaves its executor. If the trainer-side rendezvous fails while the engine is still waiting, the trainer previously hung instead of raising.
2. Resend a failed push on every rank, check the version before resuming.
   • Before: a failed push raised on the rank that saw it, and generation was neither resumed nor checked.
   • Now: the ranks agree on the outcome over the gloo group. After a failure the trainer closes any session left open on the engines, rebuilds the NCCL channel and sends the whole model once more. A second failure raises on every rank, and generation stays paused.
   • Before continue_generation, rank 0 checks that every engine reports the pushed version (/weight_info); all ranks raise otherwise.
   • Connecting the engines is agreed on the same way.

Why the failure handling differs per transport

A first sketch called send_weights(), then all-reduced a failure flag. That deadlocks, because iterating the weight source is itself a collective: every trainer rank gathers each tensor together.

• NCCL: if the sender raises at start or update, the other ranks wait in the next gather.
• IPC: every chunk ends in an all_gather_object and a barrier over all ranks, so any raise on the sender leaves the others there.

The policy is therefore per transport, entirely in vime's client and source wrappers:

• IPC: the sender's client records a failed request and skips the rest of the session. The trainer then finishes its pass and barriers on every rank before the ranks agree.
• NCCL: the client still raises, so vLLM never broadcasts to engines that are not receiving. ResumableWeightSource then lets the sender finish the pass it had started, or run a full pass if it had not started one, so it joins the gathers the other ranks are in.

Engine session rules relied on

These come from vLLM's worker.

• A failed update_weights ends the session, so a fresh start is accepted.
• start is refused while a session is open, and there is no abort. So before resending, the trainer finishes any open session without a version; the old version stays visible, and the full resend overwrites any partially loaded layers.
• The version is set only by a successful finish.

Not covered

• NCCL update_weights rejected mid-broadcast: vLLM's NCCL trainer runs the engines' update_weights on a side thread. A rejection is caught before the broadcast only if its best-effort future.done() check sees it in time; otherwise the sender enters the broadcast and blocks, which is the data-plane case below. NCCL start and finish failures, and an update rejection caught before the broadcast, are recovered.
• Data-plane failures: a sender blocked inside an NCCL broadcast (for example, an engine that stopped receiving mid-stream), or a gather that fails on one rank, ends that rank's pass while the other ranks wait in the trainer's collectives. Recovering from these needs an NCCL abort and is not attempted here.
• Expert-routing push: it drives start/finish itself and has no resend yet (TODO in _update_rollout_weights).
• Training recovery: after a final failure the trainer raises with generation paused. A restarted trainer reconnects, and the stale-session close lets its first push start cleanly.

Testing (CPU)

• tests/utils/test_weight_sync_resync.py (new): four gloo ranks run fake trainers that follow vLLM's IPC and NCCL call order. A fake engine applies the session rules above. Failures are injected at the sender's requests.
  • One failure, at start, update or finish on either transport: every rank returns after one full resend, the engine reports the new version, and generation resumes. The fake NCCL trainer surfaces an update rejection before the broadcast, so the NCCL update case covers only that path.
  • Every attempt fails: every rank raises, and the version and generation stay unchanged.
  • Negative control: the agree-after-send sketch leaves ranks in the gather until the 10 s gloo timeout.
• Existing updater tests: tests/utils/test_update_weight_from_{distributed,tensor}.py are updated, and new cases cover:
  • the full resend on a new NCCL channel;
  • the draft resend keeping the draft lifecycle;
  • IPC latching;
  • a stale version keeping generation paused.
• tests/utils/test_vllm_engine.py: each request carries its timeout.
• Regression: pre-commit passes, python -m pytest tests/utils passes, and the CPU CI test files in .buildkite/pipeline.yml give the same results as on main.

GPU validation (an NCCL and an IPC push with an injected mid-push failure, plus the init timeout) has not been run yet.

AI assistance

The code, tests and this description were written with Claude Code (AI-assisted). Each commit carries a Co-Authored-By trailer.

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #35254456 | 📁 Comparing 50828d3 against latest (d7e40fe)

  🔍 Preview build  

58 files changed · ± 58 modified

± Modified

• genindex.html
• library.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 48 more...

[gemini-code-assist]

Code Review

This pull request introduces robust error handling and resynchronization mechanisms for weight updates across distributed rollout engines, supporting both NCCL and IPC transports. It adds ResumableWeightSource to manage partial iterations, send_weights_with_resync to coordinate retries across ranks, and configurable timeouts for weight transfer requests. Two valuable review comments were identified: one addressing a potential AttributeError in a finally block that could mask exceptions, and another addressing a potential RuntimeError when dist.all_reduce falls back to NCCL with a CPU tensor.

[gemini-code-assist]

If self._reconnect() fails during send_weights_with_resync, self._trainer is set to None. In the finally block, accessing self._trainer.client will raise an AttributeError, which masks the original RuntimeError raised by send_weights_with_resync. Adding a None check prevents this masking and ensures the original exception is propagated correctly.

Suggested change

	finally:
	self._source.draft = False
	self._trainer.client.draft = False
	finally:
	self._source.draft = False
	if self._trainer is not None:
	self._trainer.client.draft = False

[gemini-code-assist]

If get_gloo_group() returns None (e.g., if Gloo is not initialized or in certain test environments), dist.all_reduce will fallback to the default process group (NCCL). Since NCCL does not support CPU tensors, calling dist.all_reduce on a CPU tensor will raise a RuntimeError: Tensors must be CUDA tensors. Placing the tensor on the current CUDA device when gloo_group is None and CUDA is available prevents this crash.

Suggested change

	flag = torch.tensor([1 if ok else 0], dtype=torch.int32)
	dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=get_gloo_group())
	gloo_group = get_gloo_group()
	device = torch.device("cpu")
	if gloo_group is None and torch.cuda.is_available():
	device = torch.device("cuda", torch.cuda.current_device())
	flag = torch.tensor([1 if ok else 0], dtype=torch.int32, device=device)
	dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=gloo_group)