Function func

tests/microbenchmarks/all_reduce.py:68–71

(x, loop_num=inner_loop)

    allreduce = AllReduce(mapping=mapping, strategy=strategy)

    def func(x, loop_num=inner_loop):
        for _ in range(loop_num):
            output = allreduce(x, all_reduce_params=allreduce_params)
        return output if fusion == AllReduceFusionOp.NONE else output[0]

    start = [torch.cuda.Event(enable_timing=True) for _ in range(outer_loop)]
    stop = [torch.cuda.Event(enable_timing=True) for _ in range(outer_loop)]