| #!/usr/bin/env python3 |
| # -*- coding: utf-8 -*- |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import sys |
| sys.path.insert(0, "../../python/") |
| import mxnet as mx |
| import numpy as np |
| import numpy.random as rnd |
| import copy |
| |
| from mxnet.test_utils import assert_almost_equal |
| |
| |
| def check_diff_to_scalar(A, x, rank=None): |
| """ assert A == x""" |
| assert(np.sum(np.abs((A - x).asnumpy())) == 0), (rank, A.asnumpy(), x) |
| |
| def compute_1bit(arr, curr_residual, threshold): |
| str_quant = "" |
| new_residual = [] |
| decompr = [] |
| |
| for idx, val in np.ndenumerate(arr): |
| val += curr_residual[idx] |
| if val > threshold: |
| str_quant += "1" |
| new_residual.append(val - 1) |
| decompr.append(1) |
| else: |
| str_quant += "0" |
| new_residual.append(val + 1) |
| decompr.append(-1) |
| |
| # append extra bits when size of array not a factor of 32 |
| if len(str_quant) != 32: |
| str_quant += "0" * (32 - len(str_quant) % 32) |
| return str_quant, new_residual, decompr |
| |
| def compute_2bit(arr, curr_residual, threshold): |
| str_quant = "" |
| new_residual = [] |
| decompr = [] |
| |
| for idx, val in np.ndenumerate(arr): |
| val += curr_residual[idx] |
| if val >= threshold: |
| str_quant += "11" |
| new_residual.append(val - threshold) |
| decompr.append(threshold) |
| elif val <= -threshold: |
| str_quant += "10" |
| new_residual.append(val + threshold) |
| decompr.append(-threshold) |
| else: |
| str_quant += "00" |
| new_residual.append(val) |
| decompr.append(0) |
| |
| # append extra bits when size of array not a factor of 16 |
| if len(str_quant) % 16 != 0: |
| str_quant += "0" * (16 - len(str_quant) % 16) |
| return str_quant, new_residual, decompr |
| |
| def compute_expected_quantization(arr, curr_residual, threshold, quantize_func): |
| |
| from struct import pack,unpack |
| def as_float32(s): |
| return unpack("f",pack("I", int(s, 2)))[0] |
| |
| arr_npy = arr.asnumpy() |
| # str_quant stores the quantized representation as a sequence of bits |
| str_quant, new_residual, decompr = quantize_func(arr_npy, curr_residual, threshold) |
| |
| compr = [] |
| # converts the string generated into integers 32chars at a time |
| for i in range(0, len(str_quant), 32): |
| cur_float = str_quant[i+24:i+32] + str_quant[i+16:i+24] + str_quant[i+8:i+16] + str_quant[i:i+8] |
| compr.append(as_float32(cur_float)) |
| |
| return np.array(compr), np.array(new_residual).reshape(arr.shape), np.array(decompr).reshape(arr.shape) |
| |
| ## individual key interface |
| def test_kvstore(kv_type, stype): |
| print(kv_type) |
| kv = mx.kv.create(kv_type) |
| kv.set_optimizer(mx.optimizer.create('test', learning_rate=-lr)) |
| for k, s in zip(keys, shapes): |
| kv.init(k, mx.nd.zeros(s)) |
| |
| res = [np.zeros(s) for s in shapes] |
| for i in range(nrepeat): |
| for j in range(len(keys)): |
| kv.push(keys[j], [mx.nd.array( |
| data[i][j][g], mx.gpu(g)).tostype(stype) for g in range(nworker)]) |
| |
| res = [a + b * lr for a, b in zip(res, [sum(d) for d in data[i]])] |
| for j in range(len(keys)): |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| err = [np.sum(np.abs(o.asnumpy() - res[j])) for o in out] |
| err = sum(err) / np.sum(np.abs(res[j])) |
| assert(err < 1e-6), (err, shapes[j]) |
| |
| def test_compress_kvstore(kv_type, compression='2bit', threshold=0.5): |
| print(kv_type + ' with ' + compression + ' compression and threshold is ' + str(threshold)) |
| rate = 2 |
| quantize_func = None |
| if compression == '1bit': |
| quantize_func = compute_1bit |
| elif compression == '2bit': |
| quantize_func = compute_2bit |
| else: |
| raise RuntimeError("Unknown gradient compression type!") |
| kv = mx.kv.create(kv_type) |
| kv.set_gradient_compression({'type':compression, 'threshold':threshold}) |
| kv.set_optimizer(mx.optimizer.create('test', learning_rate=-rate)) |
| for k, s in zip(keys, shapes): |
| kv.init(k, mx.nd.zeros(s)) |
| # init one key with 1s so we can check if it was compressed during init |
| kv.init(gc_init_test_key, mx.nd.ones(shapes[0])) |
| # use different keys for random tests so that |
| # we can track residual from start |
| random_keys = [13, 15, 17] |
| for k, s in zip(random_keys, shapes): |
| kv.init(k, mx.nd.zeros(s)) |
| |
| def pull_init_test(kv): |
| # checks that compression is not applied to init of key |
| out = [mx.nd.zeros(shapes[0], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(gc_init_test_key, out=out) |
| exp = np.ones_like(out[0].asnumpy()) |
| for o in out: |
| assert_almost_equal(o.asnumpy(), exp) |
| |
| def pull_before_push(kv): |
| for _ in range(nrepeat): |
| for j in range(len(keys)): |
| out = [mx.nd.ones(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| exp = np.zeros_like(out[0].asnumpy()) |
| for o in out: |
| assert_almost_equal(o.asnumpy(), exp) |
| |
| def push_ones(kv, sign=1): |
| for i in range(nrepeat): |
| for j in range(len(keys)): |
| kv.push(keys[j], [sign * mx.nd.ones(shapes[j], mx.gpu(g)) for g in range(nworker)]) |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| if sign == 1: |
| exp = (i + 1) * rate * nworker * np.ones_like(out[0].asnumpy()) |
| else: |
| exp = (nrepeat - i - 1) * rate * nworker * np.ones_like(out[0].asnumpy()) |
| for o in out: |
| assert_almost_equal(o.asnumpy(), exp) |
| |
| def verify_residual_1bit(kv, threshold, rate): |
| # current values must equal to zero |
| for j in range(len(keys)): |
| out = [mx.nd.ones(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| exp = np.zeros_like(out[0].asnumpy()) |
| for o in out: |
| assert_almost_equal(o.asnumpy(), exp) |
| |
| curr_residual = 0 |
| curr_val = rate * nworker if 2 > threshold else -rate * nworker |
| for j in range(len(keys)): |
| kv.push(keys[j], [2 * mx.nd.ones(shapes[j], mx.gpu(g)) for g in range(nworker)]) |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| |
| for o in out: |
| check_diff_to_scalar(o, curr_val) |
| |
| curr_residual = 1 if 2 > threshold else 3 |
| curr_val += rate * nworker if 0 + curr_residual > threshold else -rate * nworker |
| for j in range(len(keys)): |
| kv.push(keys[j], [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)]) |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| for o in out: |
| check_diff_to_scalar(o, curr_val) |
| |
| curr_residual += -1 if curr_residual > threshold else +1 |
| curr_val += rate * nworker if -2 + curr_residual > threshold else -rate * nworker |
| for j in range(len(keys)): |
| kv.push(keys[j], [-2 * mx.nd.ones(shapes[j], mx.gpu(g)) for g in range(nworker)]) |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| for o in out: |
| check_diff_to_scalar(o, curr_val) |
| |
| def push_zeros(kv): |
| for _ in range(nrepeat): |
| for j in range(len(keys)): |
| kv.push(keys[j], [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)]) |
| out = [mx.nd.ones(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| exp = np.zeros_like(out[0].asnumpy()) |
| for o in out: |
| assert_almost_equal(o.asnumpy(), exp) |
| |
| def verify_residual_2bit(kv, threshold, rate): |
| for j in range(len(keys)): |
| kv.push(keys[j], [mx.nd.ones(shapes[j], mx.gpu(g))*0.4 for g in range(nworker)]) |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j],out=out) |
| for o in out: |
| check_diff_to_scalar(o, 0) |
| |
| kv.push(keys[j], [mx.nd.ones(shapes[j], mx.gpu(g))*(threshold-0.3) for g in range(nworker)]) |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j],out=out) |
| curval = threshold * rate * nworker |
| for o in out: |
| check_diff_to_scalar(o, curval) |
| |
| kv.push(keys[j], [mx.nd.ones(shapes[j], mx.gpu(g))*(0.2) for g in range(nworker)]) |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j],out=out) |
| for o in out: |
| check_diff_to_scalar(o, curval) |
| |
| kv.push(keys[j], [mx.nd.ones(shapes[j], mx.gpu(g))*(threshold-0.3) for g in range(nworker)]) |
| out = [mx.nd.zeros(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j],out=out) |
| curval += threshold*rate*nworker |
| for o in out: |
| check_diff_to_scalar(o, curval) |
| # residual would be 0 now |
| return curval |
| |
| def check_neg(kv, neg, rate, curval): |
| for _ in range(nrepeat): |
| curval = curval + rate*nworker*neg |
| for j in range(len(keys)): |
| kv.push(keys[j], [mx.nd.ones(shapes[j], mx.gpu(g))*neg for g in range(nworker)]) |
| out = [mx.nd.ones(shapes[j], mx.gpu(g)) for g in range(nworker)] |
| kv.pull(keys[j], out=out) |
| for o in out: |
| check_diff_to_scalar(o, curval) |
| # residual would be 0 again |
| |
| def check_compr_random(kv, threshold): |
| for k, s in zip(random_keys, shapes): |
| curr_residual = [np.zeros(s) for g in range(nworker)] |
| orig_val = [mx.nd.zeros(s, mx.gpu(g)) for g in range(nworker)] |
| kv.pull(k, out=orig_val) |
| grads = [mx.nd.random_uniform(-0.6, 0.6, shape=s, ctx=mx.gpu(g)) for g in range(nworker)] |
| grads_cpy = copy.deepcopy(grads) |
| kv.push(k, grads) |
| val = [mx.nd.zeros(s, mx.gpu(g)) for g in range(nworker)] |
| kv.pull(k, out=val) |
| diffs = [val[g] - orig_val[g] for g in range(nworker)] |
| # compute expected by using simulation of operator |
| # on cpu |
| sum_dequantized_vals = np.zeros(s) |
| for g in range(nworker): |
| compr, curr_residual[g], decompr = compute_expected_quantization( |
| grads_cpy[g], curr_residual[g], threshold, quantize_func) |
| sum_dequantized_vals += (decompr * rate) |
| |
| for g in range(nworker): |
| assert_almost_equal(diffs[g].asnumpy(), sum_dequantized_vals) |
| |
| pull_init_test(kv) |
| pull_before_push(kv) |
| if compression == '1bit': |
| push_ones(kv, sign=1) |
| push_ones(kv, sign=-1) |
| verify_residual_1bit(kv, threshold, rate) |
| elif compression == '2bit': |
| push_zeros(kv) |
| curval = verify_residual_2bit(kv, threshold, rate) |
| check_neg(kv, -1*threshold, rate, curval) |
| check_compr_random(kv, threshold) |
| |
| ## group keys interface |
| def test_group_kvstore(kv_type, stype): |
| print(kv_type) |
| kv = mx.kv.create(kv_type) |
| kv.set_optimizer(mx.optimizer.create('test', learning_rate=-lr)) |
| kv.init(keys, [mx.nd.zeros(s) for s in shapes]) |
| res = [np.zeros(s) for s in shapes] |
| out = [[mx.nd.zeros(s, mx.gpu(g)) for g in range(nworker)] for s in shapes] |
| for i in range(nrepeat): |
| kv.push(keys, [[ |
| mx.nd.array(data[i][j][g], mx.gpu(g)).tostype(stype) for g in range(nworker)] |
| for j in range(len(keys))]) |
| |
| kv.pull(keys, out=out) |
| res = [a + b * lr for a, b in zip(res, [sum(d) for d in data[i]])] |
| for a, b in zip(res, out): |
| err = [np.sum(np.abs(o.asnumpy() - a)) for o in b] |
| err = sum(err) / np.sum(np.abs(a)) |
| assert(err < 1e-6), (err, a.shape) |
| |
| if __name__ == "__main__": |
| keys = [3, 5, 7] |
| # let the last shape exceed MXNET_KVSTORE_BIGARRAY_BOUND |
| shapes = [(4, 4), (100, 100), (2000, 2000)] |
| stypes = ['default', 'row_sparse'] |
| |
| gc_init_test_key = 9 |
| |
| lr = .1 |
| nworker = 4 |
| nrepeat = 10 |
| |
| # generate data |
| data = [[[np.random.random(s)*2-1 for i in range(nworker)] for s in shapes] for j in range(nrepeat)] |
| |
| for stype in stypes: |
| test_kvstore('local_update_cpu', stype) |
| test_kvstore('local_allreduce_cpu', stype) |
| test_kvstore('local_allreduce_device', stype) |
| |
| ## compression for local kvstore happens only when reduce is on device |
| test_compress_kvstore('local_allreduce_device', '1bit', -.5) |
| test_compress_kvstore('local_allreduce_device', '1bit', 0) |
| test_compress_kvstore('local_allreduce_device', '1bit', .5) |
| test_compress_kvstore('local_allreduce_device', '2bit', .5) |
| for stype in stypes: |
| test_group_kvstore('local_update_cpu', stype) |
| test_group_kvstore('local_allreduce_cpu', stype) |
| test_group_kvstore('local_allreduce_device', stype) |
