mindspore.ops.communication.all_gather_into_tensor

View Source On AtomGit
mindspore.ops.communication.all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=False)[source]

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

Note

The tensors must have the same shape and format in all processes of the collection.

Parameters

  • output_tensor (Tensor) – The output tensor to be all gathered into tensor. If the number of devices in the group is N, then the shape of output tensor is \((N*x_1, x_2, ..., x_R)\). If the function operates non in-place, this parameter is invalid.

  • input_tensor (Tensor) – The input tensor to be all gathered into tensor. The shape of tensor is \((x_1, x_2, ..., x_R)\).

  • 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

  • If the function operates in-place, return CommHandle.

  • If the function operates non in-place, return Tuple(Tensor, CommHandle). The first element stores the output result, and the second element is CommHandle.

Among them, when async_op is True, then CommHandle is an asynchronous working handle; When async_op is False, CommHandle will return None.

Raises

  • TypeError – If the type of the input_tensor or output_tensor parameter is not Tensor, group is not a str, or async_op is not bool.

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

This example should be run with 2 devices.

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore import ops
>>> from mindspore.ops.communication import init_process_group
>>> from mindspore.ops.communication import all_gather_into_tensor
>>> from mindspore import Tensor
>>>
>>> ms.set_device(device_target="Ascend")
>>> init_process_group()
>>> input_tensor = Tensor(np.ones([2, 8]).astype(np.float32))
>>> out_tensor = Tensor(np.zeros([4, 8]).astype(np.float32))
>>> output = all_gather_into_tensor(out_tensor, input_tensor)
>>> print(out_tensor)
[[1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1. 1.]]