mindspore.mint.distributed.batch_isend_irecv

View Source On Gitee
mindspore.mint.distributed.batch_isend_irecv(p2p_op_list)[source]

Batch send and recv tensors asynchronously.

Note

  • The 'isend' and 'irecv' of P2POp in p2p_op_list between ranks need to match each other.

  • P2POp in p2p_op_list can only use the same communication group.

  • tag of P2POp in p2p_op_list is not support yet.

  • tensor of P2POp in p2p_op_list will not be modified by result inplace.

  • Only support PyNative mode, Graph mode is not currently supported.

Parameters

p2p_op_list (list[P2POp]) – list contains P2POp. P2POp is type of mindspore.mint.distributed.P2POp

Returns

list[CommHandle], CommHandle is an async work handle, Currently only one packaging handle is supported.

Raises

  • TypeError – If p2p_op_list is empty or p2p_op_list are not all type of P2POp.

  • TypeError – The group name in p2p_op_list are not consistent.

  • TypeError – The tensor in p2p_op_list are not Tensor.

  • TypeError – The op in p2p_op_list are not isend or irecv.

Supported Platforms:

Ascend

Examples

Note

Before running the following examples, you need to configure the communication environment variables.

For Ascend devices, it is recommended to use the msrun startup method without any third-party or configuration file dependencies. Please see the msrun start up for more details.

This example should be run with 2 devices.

>>> import numpy as np
>>> import mindspore
>>> from mindspore.mint.distributed import init_process_group, get_rank, get_world_size
>>> from mindspore.mint.distributed import batch_isend_irecv, P2POp
>>> from mindspore import Tensor
>>>
>>> init_process_group()
>>> this_rank = get_rank()
>>> world_size = get_world_size()
>>> next_rank = (this_rank + 1) % world_size
>>> prev_rank = (this_rank + world_size - 1) % world_size
>>>
>>> send_tensor = Tensor(this_rank + 1, dtype=mindspore.float32)
>>> recv_tensor = Tensor(0., dtype=mindspore.float32)
>>>
>>> send_op = P2POp('isend', send_tensor, next_rank)
>>> recv_op = P2POp('irecv', recv_tensor, prev_rank)
>>>
>>> p2p_op_list = [send_op, recv_op]
>>> output = batch_isend_irecv(p2p_op_list)
>>> print(recv_tensor)
rank 0:
2.0
rank 1:
1.0