global_gather

paddle.distributed.utils. global_gather ( x, local_count, global_count, group=None, use_calc_stream=True ) [source]

The global_gather operator gathers the data of x into n_expert * world_size experts according to global_count, and then receives data according to local_count. The expert refers to a user-defined expert network, n_expert refers to the number of expert networks owned by each card, and world_size refers to the number of graphics cards running the network.

As shown below, the value of the world size is 2, n_expert 2, the batch size of the x 4 and local_count is [2, 0, 2, 0]. The global_count of the rank 0 is [2, 0, , ], rank 1 is [2, 0, ,](Due to the limited space, only the data calculated on rank 0 is shown here). In the global_gather operator, the meaning of the global_count and local_count is opposed to global_scatter, global_count[i] represents sending global_count[i] data to the (i % n_expert)th expert of the (i // n_expert)th card, local_count[i] represents receiving local_count[i] data from the (i // n_expert)th card to the (i % n_expert)th expert of this card. The data sent will be arranged according to the experts of each card. The rank in the figure respresent the rank of the current card in all cards.

The process of global_gather sending data is as follows:

The global_count[0] of the 0th card represents sending 2 data to the 0th expert of the 0th card;

The global_count[1] of the 0th card represents sending 0 data to the 1th expert of the 0th card;

The global_count[0] of the 1th card represents sending 2 data to the 0th expert of the 0th card;

The global_count[1] of the 1th card represents sending 0 data to the 1th expert of the 0th card.

global_scatter_gather
Parameters
  • x (Tensor) – Tensor. Tensor whose data type should be float16, float32, float64, int32 or int64.

  • local_count (Tensor) – Tensor which have n_expert * world_size elements that indicates how many data needed to be received. Tensor data type should be int64.

  • global_count (Tensor) – Tensor which have n_expert * world_size elements that indicates how many data needed to be sent. Tensor data type should be int64.

  • group (Group, optional) – The group instance return by new_group or None for global default group. Default: None.

  • use_calc_stream (bool, optional) – Wether to use calculation stream (True) or communication stream. Default: True.

Returns

The data received from all experts.

Return type

out (Tensor)

Examples

# required: distributed
import numpy as np
import paddle
from paddle.distributed import init_parallel_env
init_parallel_env()
n_expert = 2
world_size = 2
d_model = 2
in_feat = d_model
local_input_buf = np.array([[1, 2],[3, 4],[5, 6],[7, 8],[9, 10]],                                        dtype=np.float32)
if paddle.distributed.ParallelEnv().local_rank == 0:
    local_count = np.array([2, 1, 1, 1])
    global_count = np.array([2, 1, 1, 1])
else:
    local_count = np.array([1, 1, 2, 1])
    global_count = np.array([1, 1, 2, 1])
local_input_buf = paddle.to_tensor(local_input_buf, dtype="float32", stop_gradient=False)
local_count = paddle.to_tensor(local_count, dtype="int64")
global_count = paddle.to_tensor(global_count, dtype="int64")
a = paddle.distributed.utils.global_gather(local_input_buf, local_count, global_count)
print(a)
# out for rank 0: [[1, 2], [3, 4], [7, 8], [1, 2], [7, 8]]
# out for rank 1: [[5, 6], [9, 10], [3, 4], [5, 6], [9, 10]]
a.stop_gradient = False
c = a * a
c.backward()
print("local_input_buf.grad", local_input_buf.grad)
# out for rank 0: [[2, 4], [6, 8], [10, 12], [14, 16], [18, 20]]
# out for rank 1: [[2, 4], [6, 8], [10, 12], [14, 16], [18, 20]]