nemo_rl.weight_sync.nccl_reshard_weight_synchronizer#

NCCL-xfer (shard-to-shard) weight synchronizer for non-colocated deployments.

Handles disaggregated Megatron-train -> vLLM/Megatron-gen weight refit via the xferdtensor reshard: bulk FFN/expert params are resharded shard-to-shard between the train and gen parallelism layouts over a dedicated per-PP-stage NCCL communicator, while the remaining “misc” params ride a packed broadcast over the shared model_update_group. Unlike the plain collective synchronizer (which broadcasts every full tensor), this path redistributes each param directly between layouts, avoiding a full gather + broadcast.

Lifecycle: init_communicator(): 1. policy/generation.init_collective() – model_update_group (misc) 2. policy.init_nccl_reshard_comm_group() and generation.rebuild_nccl_reshard_comm_group() – per-PP-stage bulk groups 3. policy.prepare_nccl_reshard_refit_info() -> generation.prepare_nccl_reshard_refit_info() – backend-agnostic metadata sync_weights(): policy.nccl_reshard_refit(kv_scales) + generation.nccl_reshard_refit(); verify.

Like the collective transport, this is a pure data mover. Backend-specific phase transitions are owned by the caller.

Module Contents#

Classes#

NcclReshardWeightSynchronizer

Weight synchronizer using the xferdtensor shard-to-shard reshard.

Functions#

_settle_before_propagating

Let every rank finish unwinding before a refit failure reaches the caller.

API#

nemo_rl.weight_sync.nccl_reshard_weight_synchronizer._settle_before_propagating(futures, budget_s, what: str) → None#

Let every rank finish unwinding before a refit failure reaches the caller.

ray.get raises on the FIRST future that fails and leaves the rest running. That is fine when the caller is going to stop, and wrong when it is going to rebuild: a communicator rebuild is itself a collective, so dispatching init_collective while some ranks are still inside the old refit means they join late or not at all, and the rendezvous times out instead of coming up.

Job 6512153 measured exactly that on the reshard kill variant. Rank 0 gave up on its own deadline, the controller went straight into the recovery, and the rebuild began – line 963 of the log – two lines BEFORE rank 1’s watchdog fired at all. The surviving generation worker then spent 300s twice failing to reach a store that never came up, and the run died at 690s having done everything else right.

Bounded, and swallowing whatever the stragglers raise: they are unwinding from the same failure the caller is already holding, and replacing it with a straggler’s version would lose the diagnosis. If the budget runs out, propagate anyway – a caller stuck here would be a worse wedge than the one being recovered from.

class nemo_rl.weight_sync.nccl_reshard_weight_synchronizer.NcclReshardWeightSynchronizer(
policy: Any,
generation: Any,
train_cluster: Any,
inference_cluster: Any,
refit_timeout_s: Optional[float] = None,
)#

Bases: nemo_rl.weight_sync.interfaces.WeightSynchronizer

Weight synchronizer using the xferdtensor shard-to-shard reshard.

For non-colocated Megatron-train -> vLLM/Megatron-gen deployments where weights are redistributed directly between the two parallelism layouts (bulk path) plus a packed broadcast for the misc params. Mirrors

Class:

CollectiveWeightSynchronizer but additionally bootstraps the per-PP-stage bulk communicators and the nccl_reshard refit metadata.

The train/gen parallelism and per-node GPU count are derived from the policy/generation configs and the clusters, so construction matches the collective synchronizer’s signature.

Parameters:

  • policy – Policy object implementing ColocatablePolicyInterface (Megatron).

  • generation – Generation object implementing GenerationInterface.

  • train_cluster – RayVirtualCluster for the training workers. Only used by init_communicator(); may be None for sync-only instances.

  • inference_cluster – RayVirtualCluster for the inference workers. Only used by init_communicator(); may be None for sync-only instances.

  • refit_timeout_s – Deadline for one refit collective. Each participating worker arms a watchdog and aborts its own communicator when it expires, which is what lets the controller rebuild over the survivors instead of blocking in NCCL forever. None disarms it entirely, so the hang protection is lost.

Initialization

_train_parallelism() → dict[str, int]#
_gen_parallelism() → dict[str, int]#
sync_weights(
*,
timer: Optional[nemo_rl.utils.timer.Timer] = None,
kv_scales: Optional[dict[str, float]] = None,
) → None#
property is_stale: bool#
init_communicator() → None#

Build both communicator families and the refit plan, over the whole fleet.

_build(
membership: nemo_rl.weight_sync.membership.RefitMembership,
) → None#

Build everything this transport needs for the given fleet membership.

Shared by the initial build and by a rebuild after a shard is lost, deliberately. All three pieces below are functions of the inference world size, so keeping two copies of this arithmetic is how the communicators and the refit plan would come to disagree – and that disagreement is silent, because a mesh sized for the old fleet still runs, it just writes the wrong slices. Sharing the path also means every normal run exercises the rebuild code.

_settle_budget_s() → float#

How long to let stragglers unwind: their own deadline, plus a little.

A rank that has not given up yet will do so when its watchdog fires, which is the same refit_timeout_s every rank was armed with. Without a configured deadline there is nothing bounding them, so fall back to a fixed wait rather than blocking the recovery indefinitely.

reconcile_communicator(
absent_shards: collections.abc.Sequence[int],
force: bool = False,
) → bool#

Rebuild both communicator families and regenerate the refit plan.

This transport is harder to recover than the plain broadcast, and the difference is worth stating rather than discovering. Two families must be reconciled: the shared model_update_group, and the per-PP-stage bulk groups whose sub_world_size is itself a function of the inference world size.

More importantly the bulk path is a mesh-to-mesh redistribute, not a broadcast: prepare_nccl_reshard_refit_info derives each parameter’s destination placements from gen_world_size, so every gen rank receives its own slice rather than the same bytes. Dropping a rank therefore does not merely reduce the number of receivers – it orphans the slices that rank owned, and the survivors would come back holding weights that were never written. Resizing the communicators without regenerating the plan would corrupt the refit silently, which is why the plan is regenerated rather than reused.

shutdown() → None#