`nemo_rl.weight_sync.xferdtensor_python`#

Standalone striped-receive Python DTensor resharding.

xferdtensor_python_impl implements exact Shard/Replicate resharding without materializing a global tensor. For compatible replicated destinations:

Split every nonempty local destination shard into equal contiguous stripes.
Receive each stripe directly on a different member of the replica group.
Reconstruct every replica with one in-place NCCL all-gather.

The optimization is selected only when every active destination replica group can split dimension zero evenly, every exact receive can land directly in a contiguous destination view or be sent from a contiguous source view. This avoids striped plans that need both source packing and receive staging for one transfer. Other shapes and scalar tensors use the generic exact-transfer path contained in this module. This preserves arbitrary Shard/Replicate behavior without padding or a full-tensor temporary.

For contiguous local DTensors, the all-gather send stripe and receive buffer are views of the caller-owned destination shard. Explicit temporary storage is limited to noncontiguous source packing and destination receive staging. A noncontiguous local destination uses at most one additional contiguous local-shard buffer; a global-sized buffer is never allocated.

Module Contents#

Functions#

`_mesh_rank_tensor`
`_mesh_ranks`
`_mesh_signature`
`_mesh_coordinates`
`_normalize_shard_dim`
`_placement_signature`
`_validate_layout`
`_compute_shard_slices`	Return DTensor’s sequential sharding slices for one mesh coordinate.
`_rank_regions`
`_intersect`
`_local_slices`
`_region_numel`
`_local_tensor`
`_tensor_metadata`
`_destination_groups`	Group identical destination regions and select one DP0 representative.
`_build_exact_plan`	Return `(source, representative, overlap)` transfers.
`_plan_geometry`
`_destroy_subcommunicator`
`_evict_communicators`
`_finalize_process_group`
`_parent_communicator_key`
`_active_replica_signature`
`_get_replica_subcommunicator`	Collectively split once into all nonempty destination replica groups.
`clear_xferdtensor_python_caches`	Clear generic/striped plans and split communicators.
`_validate_local_inputs`	Validate local buffers before entering communicator split or P2P.
`_stage_rank_operations`
`_exchange_exact_overlaps`
`_broadcast_destination`
`_xferdtensor_python_impl_v1`	Reshard `src_tensor` into `dst_tensor` with exact data movement.
`_ordered_replica_members`	Return members in the rank order used by v1’s split communicator.
`_build_striped_targets`	Return `(destination_region, receiver)` targets and rank stripes.
`_build_striped_transfers`	Build exact source-to-stripe transfers with balanced source replicas.
`_max_p2p_peer_degree`	Return the largest distinct remote-peer count on any rank.
`_region_view_is_contiguous`	Return whether `region` is contiguous inside a row-major parent.
`_remote_transfers_avoid_double_staging`	Reject remote overlaps needing both send packing and receive staging.
`_striped_geometry`	Return cached v1 geometry plus a striped plan, or `None` fallback.
`_destination_buffer`	Return a contiguous local-shard buffer and whether it needs copy-back.
`_stage_striped_operations`	Stage exact transfers into the stripe owned by this destination rank.
`_allgather_destination`	Reconstruct every local destination shard with an in-place all-gather.
`xferdtensor_python_impl`	Reshard exactly, using striped receives for compatible replicas.

Data#

`_PLAN_CACHE`
`_PLAN_CACHE_MAX_SIZE`
`_SUBCOMM_CACHE`
`_INACTIVE_SUBCOMM_CACHE`
`_PROCESS_GROUP_COMM_IDS`
`_PROCESS_GROUP_FINALIZERS`
`_STRIPED_PLAN_CACHE`
`_STRIPED_PLAN_CACHE_MAX_SIZE`
`__all__`

API#

nemo_rl.weight_sync.xferdtensor_python._mesh_rank_tensor(mesh)#

nemo_rl.weight_sync.xferdtensor_python._mesh_ranks(mesh)#

nemo_rl.weight_sync.xferdtensor_python._mesh_signature(mesh)#

nemo_rl.weight_sync.xferdtensor_python._mesh_coordinates(mesh)#

nemo_rl.weight_sync.xferdtensor_python._normalize_shard_dim(dim, tensor_ndim)#

nemo_rl.weight_sync.xferdtensor_python._placement_signature(placements, tensor_ndim)#

nemo_rl.weight_sync.xferdtensor_python._validate_layout(mesh, placements, global_shape, name)#

nemo_rl.weight_sync.xferdtensor_python._compute_shard_slices( global_shape, mesh_shape, coordinates, placements, )#

Return DTensor’s sequential sharding slices for one mesh coordinate.

Sequential semantics matter when multiple mesh dimensions shard the same tensor dimension: each placement shards the local chunk produced by the preceding placement, including its uneven remainder.

nemo_rl.weight_sync.xferdtensor_python._rank_regions(mesh, placements, global_shape)#

nemo_rl.weight_sync.xferdtensor_python._intersect(left, right)#

nemo_rl.weight_sync.xferdtensor_python._local_slices(overlap, owner_region)#

nemo_rl.weight_sync.xferdtensor_python._region_numel(region)#

nemo_rl.weight_sync.xferdtensor_python._local_tensor(tensor)#

nemo_rl.weight_sync.xferdtensor_python._tensor_metadata(src_tensor, dst_tensor)#

nemo_rl.weight_sync.xferdtensor_python._destination_groups(dst_mesh, dst_placements, dst_regions)#: Group identical destination regions and select one DP0 representative.

nemo_rl.weight_sync.xferdtensor_python._build_exact_plan(src_regions, src_ranks, destination_groups)#

Return (source, representative, overlap) transfers.

Replicated source regions are deduplicated. If a destination representative also owns a valid source replica, use it for a local copy; otherwise use the first holder in source-mesh order.

nemo_rl.weight_sync.xferdtensor_python._PLAN_CACHE#: ‘OrderedDict(…)’

nemo_rl.weight_sync.xferdtensor_python._PLAN_CACHE_MAX_SIZE#: 1024

nemo_rl.weight_sync.xferdtensor_python._plan_geometry( src_mesh, src_placements, dst_mesh, dst_placements, global_shape, )#

nemo_rl.weight_sync.xferdtensor_python._SUBCOMM_CACHE#: None

nemo_rl.weight_sync.xferdtensor_python._INACTIVE_SUBCOMM_CACHE#: None

nemo_rl.weight_sync.xferdtensor_python._PROCESS_GROUP_COMM_IDS#: None

nemo_rl.weight_sync.xferdtensor_python._PROCESS_GROUP_FINALIZERS#: None

nemo_rl.weight_sync.xferdtensor_python._destroy_subcommunicator(communicator)#

nemo_rl.weight_sync.xferdtensor_python._evict_communicators(comm_ids)#

nemo_rl.weight_sync.xferdtensor_python._finalize_process_group(process_group_id)#

nemo_rl.weight_sync.xferdtensor_python._parent_communicator_key(process_group)#

nemo_rl.weight_sync.xferdtensor_python._active_replica_signature(destination_groups)#

nemo_rl.weight_sync.xferdtensor_python._get_replica_subcommunicator( process_group, destination_groups, device, )#

Collectively split once into all nonempty destination replica groups.

The nccl4py version used by the target environment can hang if some peers use NCCL_SPLIT_NOCOLOR. Therefore ranks outside an active replica group join one unused leftover color. This preserves one collective split and does not add any data movement.

nemo_rl.weight_sync.xferdtensor_python.clear_xferdtensor_python_caches(process_group=None)#

Clear generic/striped plans and split communicators.

This is part of the distributed call protocol: when a process group may be used again, every rank in that group must call cleanup in the same order after all GPU work has been synchronized. Asymmetric eviction can make one rank enter a future collective split while its peers reuse a cached split. Passing a process group removes only that group’s communicators; omitting it clears every local cache and is intended for coordinated teardown.

Call this explicitly before destroying or replacing a parent process group. The weak-reference finalizer is a best-effort fallback for Python cache bookkeeping, not a substitute for collective NCCL teardown ordering.

nemo_rl.weight_sync.xferdtensor_python._validate_local_inputs( rank, src_tensor, dst_tensor, src_regions, dst_regions, device, dtype, )#: Validate local buffers before entering communicator split or P2P.

nemo_rl.weight_sync.xferdtensor_python._stage_rank_operations( rank, src_tensor, dst_tensor, src_regions, dst_regions, transfers, device, dtype, stream=None, )#

nemo_rl.weight_sync.xferdtensor_python._exchange_exact_overlaps(communicator, sends, receives, stream)#

nemo_rl.weight_sync.xferdtensor_python._broadcast_destination(subcommunicator, dst_local, stream)#

nemo_rl.weight_sync.xferdtensor_python._xferdtensor_python_impl_v1( src_tensor, src_mesh, src_placement, dst_tensor, dst_mesh, dst_placement, process_group, stream=None, ) → None#

Reshard src_tensor into dst_tensor with exact data movement.

This has the same seven-argument signature and in-place effect as xferdtensor_golden. Every process-group rank must call it in the same order. The process group is expected to cover the union of source and destination meshes, as required by the reference implementation.

nemo_rl.weight_sync.xferdtensor_python._STRIPED_PLAN_CACHE#: ‘OrderedDict(…)’

nemo_rl.weight_sync.xferdtensor_python._STRIPED_PLAN_CACHE_MAX_SIZE#: 1024

nemo_rl.weight_sync.xferdtensor_python._ordered_replica_members(representative, members)#: Return members in the rank order used by v1’s split communicator.

nemo_rl.weight_sync.xferdtensor_python._build_striped_targets(destination_groups)#

Return (destination_region, receiver) targets and rank stripes.

NCCL all-gather has an equal-count contract. Splitting dimension zero of a contiguous row-major local shard produces equal, contiguous buffers and makes communicator-rank order match destination memory order exactly. None means that the complete call must use the generic v1 path.

nemo_rl.weight_sync.xferdtensor_python._build_striped_transfers(src_regions, src_ranks, targets)#: Build exact source-to-stripe transfers with balanced source replicas.

nemo_rl.weight_sync.xferdtensor_python._max_p2p_peer_degree(transfers)#: Return the largest distinct remote-peer count on any rank.

nemo_rl.weight_sync.xferdtensor_python._region_view_is_contiguous(region, parent_region)#: Return whether region is contiguous inside a row-major parent.

nemo_rl.weight_sync.xferdtensor_python._remote_transfers_avoid_double_staging( transfers, src_regions, dst_regions, )#: Reject remote overlaps needing both send packing and receive staging.

nemo_rl.weight_sync.xferdtensor_python._striped_geometry( src_mesh, src_placements, dst_mesh, dst_placements, global_shape, )#: Return cached v1 geometry plus a striped plan, or None fallback.

nemo_rl.weight_sync.xferdtensor_python._destination_buffer(dst_local, device, dtype, stream)#: Return a contiguous local-shard buffer and whether it needs copy-back.

nemo_rl.weight_sync.xferdtensor_python._stage_striped_operations( rank, src_tensor, destination_buffer, src_regions, dst_regions, transfers, device, dtype, stream, )#: Stage exact transfers into the stripe owned by this destination rank.

nemo_rl.weight_sync.xferdtensor_python._allgather_destination( subcommunicator, destination_buffer, destination_region, stripe_region, stream, )#: Reconstruct every local destination shard with an in-place all-gather.

nemo_rl.weight_sync.xferdtensor_python.xferdtensor_python_impl( src_tensor, src_mesh, src_placement, dst_tensor, dst_mesh, dst_placement, process_group, stream=None, ) → None#

Reshard exactly, using striped receives for compatible replicas.

stream (optional torch.cuda.Stream) pins every enqueued operation onto the caller’s stream so the reshard is ordered with the caller’s surrounding work (e.g. staging/copy-back of merged params); None uses the current stream.

nemo_rl.weight_sync.xferdtensor_python.__all__#: [‘clear_xferdtensor_python_caches’, ‘xferdtensor_python_impl’]

nemo_rl.weight_sync.xferdtensor_python#

Module Contents#

Functions#

Data#

API#

`nemo_rl.weight_sync.xferdtensor_python`#