`core.resharding.copy_services.nixl_copy_service`#

Module Contents#

Classes#

NixlCopyService

Refit transport over NIXL (UCX/RDMA), for cross-cluster non-collocated refit.

Functions#

_ensure_cuda_ucx_transports

Add cuda transports to UCX_TLS if it’s pinned to a host-only allowlist.

Data#

`logger`
`_CUDA_UCX_TRANSPORTS`
`_MemDesc`

API#

core.resharding.copy_services.nixl_copy_service.logger#: ‘getLogger(…)’

core.resharding.copy_services.nixl_copy_service._CUDA_UCX_TRANSPORTS#: (‘cuda_copy’, ‘cuda_ipc’)

core.resharding.copy_services.nixl_copy_service._MemDesc#: None

core.resharding.copy_services.nixl_copy_service._ensure_cuda_ucx_transports() → None#

Add cuda transports to UCX_TLS if it’s pinned to a host-only allowlist.

UCX reads UCX_TLS once at agent init. Deployments often set it to e.g. “tcp”, which can’t touch GPU memory. Only augment a plain inclusion list; leave an unset value (defaults to “all”) or an exclusion list (“^…”) alone.

class core.resharding.copy_services.nixl_copy_service.NixlCopyService(group=None, agent_name: Optional[str] = None)#

Bases: core.resharding.copy_services.base.CopyService

Refit transport over NIXL (UCX/RDMA), for cross-cluster non-collocated refit.

Each rank runs a NIXL agent. To WRITE into a peer it needs that peer’s agent metadata (connection info) plus a {task_id: (addr, len, dev)} map of the peer’s registered recv buffers. A one-time torch all-gather builds and caches this peer table. Buffers are registered locally, not exchanged.

Every refit after that is pure NIXL and sender-driven: a receiver signals each source that its buffers are free (“ready”); the source waits, syncs its weights, and issues one WRITE per receiver, each carrying a “data” notification. Those two notifications order producer and consumer per refit, so there’s no barrier and no per-refit collective. Notifications are tagged with a per-refit sequence, so stale ones are ignored. Same-rank transfers skip NIXL and copy directly.

Registered buffers are assumed address-stable across refits; if a recv address changes after setup, call clear_service_cache() to rebuild.

Initialization

requires_process_group_barrier#: False

submit_send( src_tensor: torch.Tensor, dest_rank: int, task_id: Optional[int] = None, )#

submit_recv( dest_tensor: torch.Tensor, src_rank: int, task_id: Optional[int] = None, )#

static _mem_desc( tensor: torch.Tensor, ) → core.resharding.copy_services.nixl_copy_service._MemDesc#

_register(which: str, tensors: List[torch.Tensor]) → None#

_handshake( recv_descs: Dict[int, core.resharding.copy_services.nixl_copy_service._MemDesc], ) → None#

_do_local_copies() → None#

_plan_writes( remote_sends: List[core.resharding.copy_services.base.SendOp], )#

static _notif(kind: str, seq: int) → bytes#

static _parse_notif(m: bytes) → Tuple[str, int]#

_await_notifs(kind: str, expected: int, seq: int) → None#

run()#

close() → None#

core.resharding.copy_services.nixl_copy_service#