def test_optimize_local_bounds(perf_params, grad_params):
fun = GoodputFunction(perf_params, grad_params, 128)
goodput, bsz, steps = fun.optimize(1, 1, atomic_bsz_range=(64, 256))
assert (bsz == 128), "expected bsz = 128, got {}".format(bsz)
assert (isinstance(goodput, float))
replicas = np.asarray(range(1, 100))
# single-node
goodput, bsz, steps = fun.optimize(np.ones_like(replicas),
replicas,
atomic_bsz_range=(64, 256))
assert (np.all(bsz >= np.ceil(128 / replicas).astype(int)))
assert (np.all(np.logical_or(bsz >= (64), goodput == 0.0)))
assert (np.all(bsz <= (256)))
assert (np.all(bsz * replicas <= 100 * 128))
assert (bsz[0] == 128)
assert (np.all(steps == 0))
# multi-node
goodput, bsz, steps = fun.optimize(replicas,
replicas,
atomic_bsz_range=(64, 256))
assert (np.all(bsz >= np.ceil(128 / replicas).astype(int)))
assert (np.all(np.logical_or(bsz >= (64), goodput == 0.0)))
assert (np.all(bsz <= (256)))
assert (np.all(bsz * replicas <= 100 * 128))
assert (bsz[0] == 128)
assert (np.all(steps == 0))
def test_optimize_accumulation(perf_params, grad_params):
fun = GoodputFunction(perf_params, grad_params, 128)
goodput, bsz, steps = fun.optimize(1,
1,
max_batch_size=1280,
atomic_bsz_range=(64, 256),
accumulation=True)
assert (isinstance(goodput, float))
replicas = np.asarray(range(1, 20))
# single-node
goodput, bsz, steps = fun.optimize(np.ones_like(replicas),
replicas,
max_batch_size=1280,
atomic_bsz_range=(64, 256),
accumulation=True)
assert (np.all(
np.logical_or(bsz >= np.ceil(128 / replicas).astype(int),
goodput == 0.0)))
assert (np.all(np.logical_or(bsz >= (64), goodput == 0.0)))
assert (np.all(bsz <= (256)))
assert (np.all(
np.logical_or(
bsz * replicas * (steps + 1) < 1280 + replicas * (steps + 1),
goodput == 0.0)))
assert (np.all(steps <= 15))
assert (np.all(steps >= 0))
assert (np.all(
np.logical_or(replicas > 1, np.logical_or(bsz == 128, steps > 0))))
# multi-node
goodput, bsz, steps = fun.optimize(replicas,
replicas,
max_batch_size=1280,
atomic_bsz_range=(64, 256),
accumulation=True)
assert (np.all(
np.logical_or(bsz >= np.ceil(128 / replicas).astype(int),
goodput == 0.0)))
assert (np.all(np.logical_or(bsz >= (64), goodput == 0.0)))
assert (np.all(bsz <= (256)))
assert (np.all(
np.logical_or(
bsz * replicas * (steps + 1) < 1280 + replicas * (steps + 1),
goodput == 0.0)))
assert (np.all(steps <= 15))
assert (np.all(steps >= 0))
assert (np.all(np.logical_or(np.multiply(steps, bsz) >= 256, steps == 0)))
def test_optimize_no_bounds(perf_params, grad_params):
goodput_fn = GoodputFunction(perf_params, grad_params, 128)
goodput, bsz, steps = goodput_fn.optimize(1, 3)
assert (bsz == 128 // 3 + 1), "expected bsz = 43, got {}".format(bsz)
assert (isinstance(goodput, float))
replicas = np.asarray([1, 2, 3, 4, 5])
# single-node
goodput, bsz, steps = goodput_fn.optimize(np.ones_like(replicas), replicas)
assert (bsz.shape == (5, ))
assert (np.all(bsz == np.ceil(128 / replicas).astype(int)))
assert (goodput.shape == (5, ))
assert (bsz[0] == 128)
assert (np.all(steps == 0))
# multi-node
goodput, bsz, steps = goodput_fn.optimize(replicas, replicas)
assert (bsz.shape == (5, ))
assert (np.all(bsz == np.ceil(128 / replicas).astype(int)))
assert (goodput.shape == (5, ))
assert (bsz[0] == 128)
assert (np.all(steps == 0))
def test_optimize_max_bounds(perf_params, grad_params):
fun = GoodputFunction(perf_params, grad_params, 128)
goodput, bsz, steps = fun.optimize(1, 1, max_batch_size=1280)
assert (bsz == 128), "expected bsz = 128, got {}".format(bsz)
assert (isinstance(goodput, float))
replicas = np.asarray(range(1, 100))
# single-node
goodput, bsz, steps = fun.optimize(np.ones_like(replicas),
replicas,
max_batch_size=1280)
assert (np.all(bsz >= np.ceil(128 / replicas).astype(int)))
assert (np.all(bsz * replicas <= 1280 + replicas))
assert (bsz[0] == 128)
assert (np.all(steps == 0))
# multi-node
goodput, bsz, steps = fun.optimize(replicas, replicas, max_batch_size=1280)
assert (np.all(bsz >= np.ceil(128 / replicas).astype(int)))
assert (np.all(bsz * replicas <= 1280 + replicas))
assert (bsz[0] == 128)
assert (np.all(steps == 0))
def test_one_replica_accumulation(perf_params, grad_params):
fun = GoodputFunction(perf_params, grad_params, 128)
replicas = np.asarray([1])
max_batch_sizes = np.asarray(range(128, 128 * 20, 128))
# single-node
for max_batch_size in max_batch_sizes:
goodput, bsz, steps = fun.optimize(np.ones_like(replicas),
replicas,
max_batch_size=1280,
atomic_bsz_range=(64, 256),
accumulation=True)
assert (np.all(np.logical_or(bsz >= (64), goodput == 0.0)))
assert (np.all(bsz <= (256)))
assert (np.all(
np.logical_or(bsz * (steps + 1) <= max_batch_size,
goodput == 0.0)))
assert (np.all(
np.logical_or(bsz >= np.ceil(128 / replicas).astype(int),
goodput == 0.0)))
assert (np.all(np.logical_or(bsz * (steps + 1) != 128, steps == 0)))