def check_compr_random(kv, threshold, nworker):
# set a seed so all workers generate same data. knowing this helps
# calculate expected value after pull
mx.random.seed(123)
rnd.seed(123)
nrepeat = 5
compr_random_keys_shapes = [('2121', shape),('212221',irregular_shape),('21221', big_shape)]
# use new keys so residual is 0 for calculation of expected
for k,s in compr_random_keys_shapes:
kv.init(k, mx.nd.zeros(s))
for k,s in compr_random_keys_shapes:
curr_residual = np.zeros(s)
for l in range(nrepeat):
orig_val = mx.nd.zeros(s)
kv.pull(k, orig_val)
grad = mx.nd.array(rnd.rand(s[0], s[1]))
# creates a copy because push changes grad because of assignment
grad_cpy = mx.nd.array(grad)
kv.push(k, grad)
val = mx.nd.zeros(s)
kv.pull(k, val)
diff = val - orig_val
# compute expected by using simulation of operator
compr, curr_residual, decompr = compute_expected_2bit_quantization(grad_cpy, curr_residual, threshold)
decompr *= nworker * rate
assert_almost_equal(diff.asnumpy(), decompr)
def check_compr_random(kv, threshold, nworker):
# set a seed so all workers generate same data. knowing this helps
# calculate expected value after pull
mx.random.seed(123)
rnd.seed(123)
nrepeat = 5
compr_random_keys_shapes = [('2121', shape),
('212221', irregular_shape),
('21221', big_shape)]
# use new keys so residual is 0 for calculation of expected
for k, s in compr_random_keys_shapes:
kv.init(k, mx.nd.zeros(s))
for k, s in compr_random_keys_shapes:
curr_residual = np.zeros(s)
for l in range(nrepeat):
orig_val = mx.nd.zeros(s)
kv.pull(k, orig_val)
grad = mx.nd.array(rnd.rand(s[0], s[1]))
# creates a copy because push changes grad because of assignment
grad_cpy = mx.nd.array(grad)
kv.push(k, grad)
val = mx.nd.zeros(s)
kv.pull(k, val)
diff = val - orig_val
# compute expected by using simulation of operator
compr, curr_residual, decompr = compute_expected_2bit_quantization(
grad_cpy, curr_residual, threshold)
decompr *= nworker * rate
assert_almost_equal(diff.asnumpy(), decompr)