Function benchmark_func

colossalai/auto_parallel/offload/solver.py:21–39

benchmark data transfer cost.

(func, number=1, repeat=1, warmup=3)

def benchmark_func(func, number=1, repeat=1, warmup=3):
    """
    benchmark data transfer cost.
    """

    for i in range(warmup):
        func()

    costs = []

    for i in range(repeat):
        torch.cuda.synchronize()
        begin = time.time()
        for i in range(number):
            func()
        torch.cuda.synchronize()
        costs.append((time.time() - begin) / number)

    return sum(costs) / len(costs)


class Solver(ABC):