megatron.text_generation.communication.broadcast_tensor

megatron.text_generation.communication.broadcast_tensor(size, dtype, tensor=None, rank=0)

Given size and type of a tensor on all ranks and the tensor value only on a specific rank, broadcast from that rank to all other ranks.