mindspore.ops.AllGather
- class mindspore.ops.AllGather(group=GlobalComm.WORLD_COMM_GROUP)[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
group (str, optional) – The communication group to work on. Default:
GlobalComm.WORLD_COMM_GROUP
, which means"hccl_world_group"
in Ascend, and"nccl_world_group"
in GPU.
- Inputs:
input_x (Tensor) - The shape of tensor is \((x_1, x_2, ..., x_R)\).
- Outputs:
Tensor. If the number of devices in the group is N, then the shape of output is \((N, x_1, x_2, ..., x_R)\).
- Raises
TypeError – If group is not a str.
ValueError – If the local rank id of the calling process in the group is larger than the group's rank size.
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 numpy as np >>> import mindspore as ms >>> from mindspore import ops >>> import mindspore.nn as nn >>> from mindspore.communication import init >>> from mindspore import Tensor >>> >>> ms.set_context(mode=ms.GRAPH_MODE) >>> init() >>> class Net(nn.Cell): ... def __init__(self): ... super(Net, self).__init__() ... self.allgather = ops.AllGather() ... ... def construct(self, x): ... return self.allgather(x) ... >>> input_x = Tensor(np.ones([2, 8]).astype(np.float32)) >>> net = Net() >>> output = net(input_x) >>> print(output) [[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.]]
- Tutorial Examples: