mindspore.communication.comm_func.scatter_tensor

View Source On Gitee
mindspore.communication.comm_func.scatter_tensor(tensor, src=0, group=GlobalComm.WORLD_COMM_GROUP)[source]

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

Note

The interface behavior only support Tensor input and scatter evenly, which is different from that of pytoch.distributed.scatter. 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 input tensor to be scattered. The shape of tensor is \((x_1, x_2, ..., x_R)\).

  • 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. Default: "GlobalComm.WORLD_COMM_GROUP".

Returns

Tensor, the shape of output is \((x_1/src\_rank, x_2, ..., x_R)\). The dimension 0 of data is equal to the dimension of input tensor divided by src, and the other dimension keep the same.

Raises

  • TypeError – If the type of the first input parameter is not Tensor, or any of op and group is not a str.

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

Supported Platforms:

Ascend GPU

Examples

Note

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

For Ascend/GPU/CPU 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 mindspore as ms
>>> from mindspore.communication import init
>>> from mindspore.communication.comm_func import scatter_tensor
>>> import numpy as np
>>> # Launch 2 processes.
>>>
>>> init()
>>> input = ms.Tensor(np.arange(8).reshape([4, 2]).astype(np.float32))
>>> out = scatter_tensor(tensor=data, src=0)
>>> print(out)
# rank_0
[[0. 1.]
 [2. 3.]]
# rank_1
[[4. 5.]
 [6. 7.]]