nemo_rl.weight_sync.nccl_reshard_weight_synchronizer#

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

Handles disaggregated Megatron-train -> vLLM-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/generation.init_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: policy and generation run on separate GPU clusters, so the phase transitions (offload / restore) are owned by the orchestrator, not here.

Module Contents#

Classes#

NcclReshardWeightSynchronizer

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

API#

class nemo_rl.weight_sync.nccl_reshard_weight_synchronizer.NcclReshardWeightSynchronizer(
policy: Any,
generation: Any,
train_cluster: Any,
inference_cluster: Any,
)#

Bases: nemo_rl.weight_sync.interfaces.WeightSynchronizer

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

For non-colocated Megatron-train -> vLLM-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 (vLLM).

  • 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.

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#
shutdown() → None#