`nemo_rl.weight_sync.xferdtensor`#

XferDTensor: cross-mesh DTensor reshard for disaggregated refit.

Module Contents#

Classes#

DTensorRef

DTensor-compatible reference for xferdtensor.

Functions#

`_use_golden_api`	Whether `NRL_XFERDTENSOR_GOLDEN` forces the golden reshard path.
`_use_python_api`	Use Python only implementation in for the xferdtensor.
`xferdtensor`	Public XferDTensor entry point used by all external callers.
`_flatten_mesh_ranks`
`_get_tensor_meta`
`_get_mesh_coords`
`_compute_shard_slices`	Return the slice of the global tensor this rank owns.
`xferdtensor_golden`	Broadcast-based reference implementation of XferDTensor.

Data#

_XFERDTENSOR_PATH_LOGGED

API#

nemo_rl.weight_sync.xferdtensor._XFERDTENSOR_PATH_LOGGED#: False

class nemo_rl.weight_sync.xferdtensor.DTensorRef( local_tensor: torch.Tensor, global_shape, dtype=None, device=None, )#

DTensor-compatible reference for xferdtensor.

Provides the interface xferdtensor reads via duck typing:

.shape: global tensor shape (torch.Size)
._local_tensor: local shard (src side) or dst buffer (dst side)
.dtype, .device: tensor metadata

On the src side (train), local_tensor is the TP-local shard from Megatron parameters (no PP broadcast or TP gather needed). global_shape is the full unsharded shape.

On the dst side (gen), local_tensor is either the vLLM local parameter (for direct params) or a temporary buffer (for merged/unmapped params). global_shape is always the full unsharded shape.

Initialization

nemo_rl.weight_sync.xferdtensor._use_golden_api() → bool#: Whether NRL_XFERDTENSOR_GOLDEN forces the golden reshard path.

nemo_rl.weight_sync.xferdtensor._use_python_api() → bool#

Use Python only implementation in for the xferdtensor.

xferdtensor_python_impl is the backup implementation when the nccl.m2n.reshard is not available.

nemo_rl.weight_sync.xferdtensor.xferdtensor( src_tensor, src_mesh, src_placement, dst_tensor, dst_mesh, dst_placement, process_group, stream=None, )#: Public XferDTensor entry point used by all external callers.

nemo_rl.weight_sync.xferdtensor._flatten_mesh_ranks(mesh)#

nemo_rl.weight_sync.xferdtensor._get_tensor_meta(src_tensor, dst_tensor)#

nemo_rl.weight_sync.xferdtensor._get_mesh_coords(mesh, rank)#

nemo_rl.weight_sync.xferdtensor._compute_shard_slices( global_shape, mesh_shape, mesh_coords, placements, )#: Return the slice of the global tensor this rank owns.

nemo_rl.weight_sync.xferdtensor.xferdtensor_golden( src_tensor, src_mesh, src_placement, dst_tensor, dst_mesh, dst_placement, process_group, )#

Broadcast-based reference implementation of XferDTensor.

Reconstructs the full (global) tensor from TP-local shards on the source mesh, then each destination rank extracts its local shard based on dst_placement.

When all src_placement entries are Replicate, a single broadcast from src_ranks[0] suffices. Otherwise, one broadcast per unique shard region reconstructs the full tensor.

This is the canonical 7-argument signature matching the real nccl.m2n.reshard op. Callers must wrap raw tensors in DTensorRef so that .shape reports the global shape and ._local_tensor holds the local shard.

nemo_rl.weight_sync.xferdtensor#