mindspore.ops.communication.all_gather

View Source On AtomGit
mindspore.ops.communication.all_gather(tensor_list, tensor, group=None, async_op=False)[source]

Gathers tensors from the specified communication group and returns the tensor list which is all gathered.

Note

In the MCCL communication group scenario, elements in tensor_list with different shapes are not supported.

Parameters:

  • tensor_list (list[Tensor]) – Output list.

  • tensor (Tensor) – The input tensor to be all gathered into tensor.

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

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

Returns:

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

Raises:

  • TypeError – If the type of input tensor is not Tensor, tensor_list is not Tensor List, group is not a str or async_op is not bool.

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

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

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

Supported Platforms:

Ascend CPU

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 startup for more details.

This example should be run with 2 devices.

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore.ops.communication import init_process_group
>>> from mindspore.ops.communication import all_gather
>>> from mindspore import Tensor
>>>
>>> init_process_group()
>>> input_tensor = Tensor(np.ones([2, 8]).astype(np.float32))
>>> out_tensors = [Tensor(np.zeros([2, 8]).astype(np.float32)), Tensor(np.zeros([2, 8]).astype(np.float32))]
>>> output = all_gather(out_tensors, input_tensor)
>>> print(out_tensors)
[Tensor(shape=[2, 8], dtype=Float32, value=
[[ 1.00000000e+00,  1.00000000e+00,  1.00000000e+00 ...  1.00000000e+00,  1.00000000e+00,  1.00000000e+00],
 [ 1.00000000e+00,  1.00000000e+00,  1.00000000e+00 ...  1.00000000e+00,  1.00000000e+00,  1.00000000e+00]]),
Tensor(shape=[2, 8], dtype=Float32, value=
[[ 1.00000000e+00,  1.00000000e+00,  1.00000000e+00 ...  1.00000000e+00,  1.00000000e+00,  1.00000000e+00],
 [ 1.00000000e+00,  1.00000000e+00,  1.00000000e+00 ...  1.00000000e+00,  1.00000000e+00,  1.00000000e+00]])]