oneflow.comm.broadcast

oneflow.comm.broadcast(tensor, src)

Broadcasts the tensor to the whole group. tensor must have the same number of elements in all processes participating in the collective.

Parameters
  • tensor (Tensor) – Data to be sent if src is the rank of current process, and tensor to be used to save received data otherwise.

  • src (int) – Source rank.

>>> # We have 1 process groups, 2 ranks.
>>> import oneflow as flow
>>> tensor = flow.tensor([[1, 2], [3, 4]], device="cuda") + flow.env.get_local_rank()
>>> # input on rank0
>>> tensor 
tensor([[1, 2],
        [3, 4]], device='cuda:0', dtype=oneflow.int64)
>>> # input on rank1
>>> tensor 
tensor([[2, 3],
        [4, 5]], device='cuda:1', dtype=oneflow.int64)
>>> flow.comm.broadcast(tensor, 0)
>>> # result on rank0
>>> tensor 
tensor([[1, 2],
        [3, 4]], device='cuda:0', dtype=oneflow.int64)