Function _broadcast

linear_moe/data/utils.py:32–34

(item)

    args = get_args()
    tokenizer = get_tokenizer()
    def _broadcast(item):
        torch.distributed.broadcast(item, mpu.get_tensor_model_parallel_src_rank(),
                                    group=mpu.get_tensor_model_parallel_group())

    if mpu.get_tensor_model_parallel_rank() == 0: