Is this page helpful?

`core.resharding.copy_services.nccl_copy_service`#

Module Contents#

Classes#

NCCLCopyService

Thin wrapper around torch.distributed batch_isend_irecv to submit and execute a batch of point-to-point sends and recvs.

Data#

API#

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

class core.resharding.copy_services.nccl_copy_service.NCCLCopyService(group=None)#

Bases: core.resharding.copy_services.base.CopyService

Thin wrapper around torch.distributed batch_isend_irecv to submit and execute a batch of point-to-point sends and recvs.

Initialization

supports_multiple_runs_per_plan#: True

TASK_WINDOW#: ‘int(…)’

_ensure_nccl_connected() → None#: Collectively initialize the NCCL communicator on the current device.

set_plan(plan, *, transform=None) → None#: Remember the plan-global sizes that decide whether run() windows its submission.

_should_window( remote_sends: List[core.resharding.copy_services.base.SendOp], remote_recvs: List[core.resharding.copy_services.base.RecvOp], ) → bool#

_run_in_task_windows( remote_sends: List[core.resharding.copy_services.base.SendOp], remote_recvs: List[core.resharding.copy_services.base.RecvOp], ) → None#: Issue the pending remote ops as one batch_isend_irecv per global task-id window.

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, )#

run()#