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)