mindspore.mint.distributed.scatter

View Source On Gitee
mindspore.mint.distributed.scatter(tensor, scatter_list, src=0, group=None, async_op=False)[source]

Scatter tensor evently across the processes in the specified communication group.

Note

  • The interface behavior only support Tensor List input and scatter evenly.

  • Only the tensor in process src (global rank) will do scatter.

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

Parameters

  • tensor (Tensor) – the output tensor.

  • scatter_list (list[Tensor]) – List of same-sized tensors to scatter. default is None, must be specified on the source rank.

  • src (int, optional) – Specifies the rank(global rank) of the process that send the tensor. And only process src will send the tensor.

  • group (str, optional) – The communication group to work on. If None, which means "hccl_world_group" in Ascend. Default: None.

  • async_op (bool, optional) – Whether this operator should be an async operator. Default: False .

Returns

CommHandle, CommHandle is an async work handle, if async_op is set to True. CommHandle will be None, when async_op is False.

Raises

  • TypeError – If the type of tensor parameter is not Tensor, scatter_list is not Tensor List.

  • TypeError – If any of op and group is not a str. async_op is not bool or 'op' is invalid.

  • TypeError – If size of scatter_list is not equal to group size.

  • TypeError – If the type or shape of tensor not equal to the member of scatter_list.

  • RuntimeError – If device target is invalid, or backend is invalid, or distributed initialization fails.

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.

>>> from mindspore import Tensor
>>> from mindspore.mint.distributed import init_process_group, scatter
>>> import numpy as np
>>> # Launch 2 processes.
>>>
>>> init_process_group()
>>> inputs = [Tensor(np.ones([2, 2]).astype(np.float32)), Tensor(np.ones([2, 2]).astype(np.float32))]
>>> output = Tensor(np.zeros([2, 2]).astype(np.float32))
>>> scatter(output, inputs, src=0)
>>> print(output)
# rank_0
[[1. 1.]
 [1. 1.]]
# rank_1
[[1. 1.]
 [1. 1.]]