Initial commit

2 files changed, 575 insertions, 0 deletions

diff --git a/tests/test_functional.py b/tests/test_functional.py
new file mode 100644
index 0000000..2a7d308
--- /dev/null
+++ b/tests/test_functional.py
@@ -0,0 +1,213 @@
+# Copyright (c) Facebook, Inc. and its affiliates. 
+#   
+# This source code is licensed under the MIT license found in the 
+# LICENSE file in the root directory of this source tree.
+import pytest
+import torch
+import bitsandbytes as bnb
+
+from itertools import product
+
+from bitsandbytes import functional as F
+
+def setup():
+    pass
+
+def teardown():
+    pass
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16], ids=['float', 'half'])
+def test_estimate_quantiles(dtype):
+    A = torch.rand(1024, 1024, device='cuda')
+    A = A.to(dtype)
+    code = F.estimate_quantiles(A)
+
+    percs = torch.linspace(1/512, 511/512, 256, device=A.device)
+    torch.testing.assert_allclose(percs, code, atol=1e-3, rtol=1e-2)
+
+    A = torch.randn(1024, 1024, device='cuda')
+    A = A.to(dtype)
+    code = F.estimate_quantiles(A)
+
+    quantiles = torch.quantile(A.float(), percs)
+    diff = torch.abs(code-quantiles)
+    assert (diff > 5e-02).sum().item() == 0
+
+
+def test_quantile_quantization():
+    for i in range(100):
+        A1 = torch.randn(1024, 1024, device='cuda')
+        code = F.estimate_quantiles(A1)
+        C = F.quantize_no_absmax(A1, code)
+        A2 = F.dequantize_no_absmax(C, code)
+        diff = torch.abs(A1-A2).mean().item()
+        assert diff < 0.0075
+
+        A1 = torch.rand(1024, 1024, device='cuda')
+        code = F.estimate_quantiles(A1)
+        C = F.quantize_no_absmax(A1, code)
+        A2 = F.dequantize_no_absmax(C, code)
+        diff = torch.abs(A1-A2).mean().item()
+        torch.testing.assert_allclose(A1, A2, atol=5e-3, rtol=0)
+        assert diff < 0.001
+
+
+def test_dynamic_quantization():
+    diffs = []
+    reldiffs = []
+    for i in range(100):
+        A1 = torch.randn(1024, 1024, device='cuda')
+        C, S = F.quantize(A1)
+        A2 = F.dequantize(C, S)
+        diff = torch.abs(A1-A2)
+        reldiff = diff/torch.abs(A1+1e-8)
+        diffs.append(diff.mean().item())
+        reldiffs.append(reldiff.mean().item())
+        assert diff.mean().item() < 0.0135
+    print(sum(diffs)/len(diffs))
+    print(sum(reldiffs)/len(reldiffs))
+
+    for i in range(100):
+        A1 = torch.rand(1024, 1024, device='cuda')
+        C, S = F.quantize(A1)
+        A2 = F.dequantize(C, S)
+        diff = torch.abs(A1-A2).mean().item()
+        torch.testing.assert_allclose(A1, A2, atol=1e-2, rtol=0)
+        assert diff < 0.004
+
+
+def test_dynamic_blockwise_quantization():
+    diffs = []
+    reldiffs = []
+    for i in range(100):
+        A1 = torch.randn(1024, 1024, device='cuda')
+        C, S = F.quantize_blockwise(A1)
+        A2 = F.dequantize_blockwise(C, S)
+        diff = torch.abs(A1-A2)
+        reldiff = diff/torch.abs(A1+1e-8)
+        diffs.append(diff.mean().item())
+        reldiffs.append(reldiff.mean().item())
+        assert diffs[-1] < 0.011
+    print(sum(diffs)/len(diffs))
+    print(sum(reldiffs)/len(reldiffs))
+
+    diffs = []
+    for i in range(100):
+        A1 = torch.rand(1024, 1024, device='cuda')
+        C, S = F.quantize_blockwise(A1)
+        A2 = F.dequantize_blockwise(C, S)
+        diff = torch.abs(A1-A2).mean().item()
+        assert diff < 0.0033
+        diffs.append(diff)
+        torch.testing.assert_allclose(A1, A2, atol=1e-2, rtol=0)
+    #print(sum(diffs)/len(diffs))
+
+def test_dynamic_blockwise_stochastic_quantization():
+    diffs = []
+    reldiffs = []
+    rand = torch.rand(1024).cuda()
+    for i in range(100):
+        A1 = torch.randn(1024, 1024, device='cuda')
+        C1, S1 = F.quantize_blockwise(A1, rand=rand)
+        C2, S2 = F.quantize_blockwise(A1)
+        # a maximunm distance of quantized values of 1
+        torch.testing.assert_allclose(C1, C2, atol=1, rtol=0)
+        fraction_smaller = (C1<C2).float().sum()/C1.numel()
+        fraction_larger = (C1>C2).float().sum()/C1.numel()
+        torch.testing.assert_allclose(fraction_larger, fraction_smaller, atol=0.01, rtol=0)
+
+
+
+@pytest.mark.parametrize("gtype", [torch.float32, torch.float16], ids=['float', 'half'])
+def test_percentile_clipping(gtype):
+    gnorm_vec1 = torch.zeros(100, device='cuda')
+    gnorm_vec2 = torch.zeros(100, device='cuda')
+    n = 4
+    step = 0
+    percentile=5
+    for i in range(1000):
+        step += 1
+        g = torch.randn(n, n, dtype=gtype, device='cuda')
+        gnorm1, clip2, gnorm_scale = F.percentile_clipping(g, gnorm_vec2, step, percentile=percentile)
+        assert gnorm_scale == 1.0 if gnorm1 < clip2 else clip2/gnorm1
+
+        gnorm2 = torch.norm(g.float())
+        if step == 1:
+            gnorm_vec1[:] = gnorm2
+        else:
+            gnorm_vec1[step % 100] = gnorm2
+
+        vals, idx = torch.sort(gnorm_vec1)
+        clip1 = vals[percentile]
+
+        torch.testing.assert_allclose(gnorm_vec1, torch.sqrt(gnorm_vec2))
+        torch.testing.assert_allclose(clip1, clip2)
+        torch.testing.assert_allclose(gnorm1, gnorm2)
+
+
+def test_stable_embedding():
+    layer = bnb.nn.StableEmbedding(1024, 1024)
+    layer.reset_parameters()
+
+
+def test_dynamic_blockwise_quantization_cpu():
+    #A1 = torch.randn(1024, 1024, device='cpu')
+    #code = F.create_dynamic_map()
+    #for i in range(1000):
+    #    C, S = F.quantize_blockwise(A1, code=code)
+    #    A2 = F.dequantize_blockwise(C, S)
+
+    for i in range(10):
+        # equivalence with GPU blockwise quantization
+        A1 = torch.randn(1024, 1024, device='cpu')
+        C1, S1 = F.quantize_blockwise(A1)
+        C2, S2 = F.quantize_blockwise(A1.cuda())
+        torch.testing.assert_allclose(S1[0], S2[0].cpu())
+        # there seems to be some issues with precision in CUDA vs CPU
+        # not all elements are usually close, with couple off elements in a million
+        idx = torch.isclose(C1, C2.cpu())
+        assert (idx==0).sum().item() < 15
+
+
+    diffs = []
+    reldiffs = []
+    for i in range(10):
+        A1 = torch.randn(1024, 1024, device='cpu')
+        C, S = F.quantize_blockwise(A1)
+        A2 = F.dequantize_blockwise(C, S)
+        diff = torch.abs(A1-A2)
+        reldiff = diff/torch.abs(A1+1e-8)
+        diffs.append(diff.mean().item())
+        reldiffs.append(reldiff.mean().item())
+        assert diffs[-1] < 0.011
+    #print(sum(diffs)/len(diffs))
+    #print(sum(reldiffs)/len(reldiffs))
+
+    diffs = []
+    for i in range(10):
+        A1 = torch.rand(1024, 1024, device='cpu')
+        C, S = F.quantize_blockwise(A1)
+        A2 = F.dequantize_blockwise(C, S)
+        diff = torch.abs(A1-A2).mean().item()
+        assert diff < 0.0033
+        diffs.append(diff)
+        torch.testing.assert_allclose(A1, A2, atol=1e-2, rtol=0)
+    #print(sum(diffs)/len(diffs))
+
+
+def test_histogram():
+    dim1, dim2 = 32, 32
+    source = torch.rand(dim1, dim2, device='cuda')
+    idx1 = torch.randint(0, 255, size=(dim1, dim2), device='cuda').int()
+    idx2 = torch.randint(0, 255, size=(dim1, dim2), device='cuda').int()
+    histogram1 = torch.zeros((256, 256)).cuda()
+    histogram2 = torch.zeros((256, 256)).cuda()
+
+    F.histogram_scatter_add_2d(histogram2, idx1, idx2, source)
+
+    for i in range(dim1):
+        for j in range(dim2):
+            histogram1[idx1[i, j].item(), idx2[i, j].item()] += source[i, j]
+
+    torch.testing.assert_allclose(histogram1, histogram2)
+    torch.testing.assert_allclose(histogram1.sum(), source.sum())
diff --git a/tests/test_optim.py b/tests/test_optim.py
new file mode 100644
index 0000000..4d67b08
--- /dev/null
+++ b/tests/test_optim.py
@@ -0,0 +1,362 @@
+# Copyright (c) Facebook, Inc. and its affiliates. 
+#   
+# This source code is licensed under the MIT license found in the 
+# LICENSE file in the root directory of this source tree.
+import os
+import time
+import shutil
+import uuid
+import pytest
+import ctypes
+import torch
+import bitsandbytes as bnb
+import bitsandbytes.functional as F
+
+from os.path import join
+from itertools import product
+
+import apex
+
+def get_temp_dir():
+    path = '/tmp/autoswap/{0}'.format(str(uuid.uuid4()))
+    os.makedirs(path, exist_ok=True)
+    return path
+
+def rm_path(path):
+    shutil.rmtree(path)
+
+str2optimizers = {}
+str2optimizers['adam_pytorch'] = (None, torch.optim.Adam, bnb.optim.Adam)
+str2optimizers['adam_apex'] = (None, apex.optimizers.FusedAdam, bnb.optim.Adam)
+str2optimizers['momentum_apex'] = (None, lambda pxx: apex.optimizers.FusedSGD(pxx, 0.01, 0.9), bnb.optim.Adam)
+str2optimizers['momentum_pytorch'] = (None, lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), bnb.optim.Adam)
+str2optimizers['lamb_apex'] = (None, lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.00, use_nvlamb=True), bnb.optim.Adam)
+str2optimizers['lars_apex'] = (None, lambda pxx: apex.parallel.LARC.LARC(apex.optimizers.FusedSGD(pxx, 0.01, 0.9)), bnb.optim.Adam)
+
+str2optimizers['adam'] = (torch.optim.Adam, bnb.optim.Adam)
+str2optimizers['fused_adam'] = (apex.optimizers.FusedAdam, bnb.optim.Adam)
+str2optimizers['momentum'] = (lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), lambda pxx: bnb.optim.SGD(pxx, 0.01, 0.9, block_wise=False))
+str2optimizers['lars'] = (lambda pxx: bnb.optim.PytorchLARS(pxx, 0.01, 0.9), lambda pxx: bnb.optim.LARS(pxx, 0.01, 0.9))
+str2optimizers['lamb'] = (lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.0, max_grad_norm=10000.0, eps=1e-8, use_nvlamb=True), bnb.optim.LAMB)
+str2optimizers['rmsprop'] = (lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9), lambda pxx: bnb.optim.RMSprop(pxx, 0.01, 0.9, block_wise=False))
+str2optimizers['adam8bit'] = (torch.optim.Adam, lambda pxx: bnb.optim.Adam8bit(pxx, block_wise=False))
+str2optimizers['momentum8bit'] = (lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), lambda pxx: bnb.optim.SGD8bit(pxx, 0.01, 0.9, block_wise=False))
+str2optimizers['rmsprop8bit'] = (lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9), lambda pxx: bnb.optim.RMSprop8bit(pxx, 0.01, 0.9, block_wise=False))
+str2optimizers['lamb8bit'] = (lambda pxx: apex.optimizers.FusedLAMB(pxx, weight_decay=0.0, max_grad_norm=10000.0, eps=1e-8, use_nvlamb=True), bnb.optim.LAMB8bit)
+str2optimizers['lars8bit'] = (lambda pxx: bnb.optim.PytorchLARS(pxx, 0.01, 0.9), lambda pxx: bnb.optim.LARS8bit(pxx, 0.01, 0.9))
+
+str2optimizers['adam8bit_blockwise'] = (torch.optim.Adam, lambda pxx: bnb.optim.Adam8bit(pxx, block_wise=True))
+str2optimizers['momentum8bit_blockwise'] = (lambda pxx: torch.optim.SGD(pxx, 0.01, 0.9), lambda pxx: bnb.optim.SGD8bit(pxx, 0.01, 0.9, block_wise=True))
+str2optimizers['rmsprop8bit_blockwise'] = (lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9), lambda pxx: bnb.optim.RMSprop8bit(pxx, 0.01, 0.9, block_wise=True))
+
+str2statenames = {}
+str2statenames['adam'] = [('exp_avg', 'state1'), ('exp_avg_sq', 'state2')]
+str2statenames['momentum'] = [('momentum_buffer', 'state1')]
+str2statenames['lars'] = [('momentum_buffer', 'state1')]
+str2statenames['lamb'] = [('exp_avg', 'state1'), ('exp_avg_sq', 'state2')]
+str2statenames['rmsprop'] = [('square_avg', 'state1')]
+str2statenames['adam8bit'] = [('exp_avg', 'state1', 'qmap1', 'max1'), ('exp_avg_sq', 'state2', 'qmap2', 'max2')]
+str2statenames['lamb8bit'] = [('exp_avg', 'state1', 'qmap1', 'max1'), ('exp_avg_sq', 'state2', 'qmap2', 'max2')]
+str2statenames['adam8bit_blockwise'] = [('exp_avg', 'state1', 'qmap1', 'absmax1'), ('exp_avg_sq', 'state2', 'qmap2', 'absmax2')]
+str2statenames['momentum8bit'] = [('momentum_buffer', 'state1', 'qmap1', 'max1')]
+str2statenames['momentum8bit_blockwise'] = [('momentum_buffer', 'state1', 'qmap1', 'absmax1')]
+str2statenames['lars8bit'] = [('momentum_buffer', 'state1', 'qmap1', 'max1')]
+str2statenames['rmsprop8bit'] = [('square_avg', 'state1', 'qmap1', 'max1')]
+str2statenames['rmsprop8bit_blockwise'] = [('square_avg', 'state1', 'qmap1', 'absmax1')]
+
+dim1 = [1024]
+dim2 = [32, 1024, 4097, 1]
+gtype = [torch.float32, torch.float16]
+optimizer_names = ['adam', 'momentum', 'rmsprop', 'lars', 'lamb']
+values = list(product(dim1,dim2, gtype, optimizer_names))
+names = ['dim1_{0}_dim2_{1}_gtype_{2}_optim_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, gtype, optim_name", values, ids=names)
+def test_optimizer32bit(dim1, dim2, gtype, optim_name):
+    if dim1 == 1 and dim2 == 1: return
+    p1 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1
+    p2 = p1.clone()
+    p1 = p1.float()
+
+
+    torch_optimizer = str2optimizers[optim_name][0]([p1])
+    bnb_optimizer = str2optimizers[optim_name][1]([p2])
+
+    if gtype == torch.float32:
+        atol, rtol = 1e-6, 1e-5
+    else:
+        atol, rtol = 1e-4, 1e-3
+
+
+    for i in range(50):
+        g = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.01
+        p1.grad = g.clone().float()
+        p2.grad = g.clone()
+
+        bnb_optimizer.step()
+        torch_optimizer.step()
+
+        for name1, name2 in str2statenames[optim_name]:
+            torch.testing.assert_allclose(torch_optimizer.state[p1][name1], bnb_optimizer.state[p2][name2], atol=atol, rtol=rtol)
+
+        torch.testing.assert_allclose(p1, p2.float(), atol=atol, rtol=rtol)
+
+        if i % 10 == 0 and i > 0:
+            path = get_temp_dir()
+            torch.save(bnb_optimizer.state_dict(),join(path, 'opt.pt'))
+            del bnb_optimizer
+            bnb_optimizer = None
+            bnb_optimizer = str2optimizers[optim_name][1]([p2])
+            bnb_optimizer.load_state_dict(torch.load(join(path, 'opt.pt')))
+            rm_path(path)
+            torch.testing.assert_allclose(p1, p2.float(), atol=atol, rtol=rtol)
+            for name1, name2 in str2statenames[optim_name]:
+                torch.testing.assert_allclose(torch_optimizer.state[p1][name1], bnb_optimizer.state[p2][name2], atol=atol, rtol=rtol)
+
+        if gtype == torch.float16:
+            # the adam buffers should also be close because they are 32-bit
+            # but the paramters can diverge because they are 16-bit
+            # the difference grow larger and larger with each update
+            # --> copy the state to keep weights close
+            p1.data = p1.data.half().float()
+            p2.copy_(p1.data)
+            torch.testing.assert_allclose(p1.half(), p2)
+        if optim_name in ['lars', 'lamb']:
+            assert bnb_optimizer.state[p2]['unorm_vec'] > 0.0
+
+dim1 = [1024]
+dim2 = [32, 1024, 4097]
+gtype = [torch.float32, torch.float16]
+values = list(product(dim1,dim2, gtype))
+names = ['dim1_{0}_dim2_{1}_gtype_{2}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, gtype", values, ids=names)
+def test_global_config(dim1, dim2, gtype):
+    if dim1 == 1 and dim2 == 1: return
+    p1 = torch.randn(dim1,dim2, device='cpu', dtype=gtype)*0.1
+    p2 = torch.randn(dim1,dim2, device='cpu', dtype=gtype)*0.1
+    p3 = torch.randn(dim1,dim2, device='cpu', dtype=gtype)*0.1
+    mask = torch.rand_like(p2) < 0.1
+    beta1 = 0.9
+    beta2 = 0.999
+    lr = 0.001
+    eps = 1e-8
+
+    bnb.optim.GlobalOptimManager.get_instance().initialize()
+    bnb.optim.GlobalOptimManager.get_instance().override_config(p3, 'optim_bits', 8)
+
+    bnb.optim.GlobalOptimManager.get_instance().register_parameters([p1, p2, p3])
+    p1 = p1.cuda()
+    p2 = p2.cuda()
+    p3 = p3.cuda()
+
+    adam2 = bnb.optim.Adam([p1, p2, p3], lr, (beta1, beta2), eps)
+
+    if gtype == torch.float32:
+        atol, rtol = 1e-6, 1e-5
+    else:
+        atol, rtol = 1e-4, 1e-3
+
+    for i in range(50):
+        g1 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1 + 0.001
+        g2 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1 + 0.001
+        g3 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1 + 0.001
+        p1.grad = g1
+        p2.grad = g2
+        p3.grad = g3
+
+        adam2.step()
+
+        assert adam2.state[p3]['state1'].dtype == torch.uint8
+        assert adam2.state[p3]['state2'].dtype == torch.uint8
+
+
+
+dim1 = [1024]
+dim2 = [32, 1024, 4097]
+gtype = [torch.float32, torch.float16]
+optimizer_names = ['adam8bit', 'momentum8bit', 'rmsprop8bit', 'adam8bit_blockwise', 'lamb8bit', 'lars8bit', 'momentum8bit_blockwise', 'rmsprop8bit_blockwise']
+values = list(product(dim1,dim2, gtype, optimizer_names))
+names = ['dim1_{0}_dim2_{1}_gtype_{2}_optim_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, gtype, optim_name", values, ids=names)
+def test_optimizer8bit(dim1, dim2, gtype, optim_name):
+    if dim1 == 1 and dim2 == 1: return
+    p1 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1
+    p2 = p1.clone()
+    p1 = p1.float()
+    blocksize = 2048
+
+    torch_optimizer = str2optimizers[optim_name][0]([p1])
+    bnb_optimizer = str2optimizers[optim_name][1]([p2])
+
+    if gtype == torch.float32:
+        atol, rtol = 3e-3, 1e-3
+        patol, prtol = 1e-5, 1e-3
+
+    else:
+        atol, rtol = 3e-3, 1e-3
+        patol, prtol = 1e-5, 1e-3
+
+    errors = []
+    relerrors = []
+
+    for i in range(50):
+        g = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.01
+        p1.grad = g.clone().float()
+        p2.grad = g.clone()
+
+        bnb_optimizer.step()
+        torch_optimizer.step()
+
+        torch.testing.assert_allclose(p1, p2.float(), atol=patol, rtol=prtol)
+
+        dequant_states = []
+        for name1, name2, qmap, max_val in str2statenames[optim_name]:
+            #print(bnb_optimizer.state[p2][max_val], name1)
+            if 'blockwise' in optim_name:
+                s1 = F.dequantize_blockwise(code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val], A=bnb_optimizer.state[p2][name2], blocksize=blocksize)
+            else:
+                s1 = F.dequantize(code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val], A=bnb_optimizer.state[p2][name2])
+            num_not_close = torch.isclose(torch_optimizer.state[p1][name1], s1, atol=atol, rtol=rtol)==0
+            assert num_not_close.sum().item() < 20
+            dequant_states.append(s1.clone())
+
+        err  = torch.abs(p1-p2)
+        relerr = err/torch.abs(p1)
+        assert err.mean() < 0.0001
+        assert relerr.mean() < 0.001
+
+        errors.append(err.mean().item())
+        relerrors.append(relerr.mean().item())
+
+        if i % 10 == 0 and i > 0:
+            for (name1, name2, qmap, max_val), s in zip(str2statenames[optim_name], dequant_states):
+                s1cpy = s.clone()
+                raws1cpy = bnb_optimizer.state[p2][name2].clone()
+                qmap1 = bnb_optimizer.state[p2][qmap].clone()
+
+                path = get_temp_dir()
+                torch.save(bnb_optimizer.state_dict(),join(path, 'opt.pt'))
+                del bnb_optimizer
+                bnb_optimizer = None
+                bnb_optimizer = str2optimizers[optim_name][1]([p2])
+                bnb_optimizer.load_state_dict(torch.load(join(path, 'opt.pt')))
+                rm_path(path)
+                torch.testing.assert_allclose(raws1cpy, bnb_optimizer.state[p2][name2])
+                torch.testing.assert_allclose(qmap1, bnb_optimizer.state[p2][qmap])
+
+                if 'blockwise' in optim_name:
+                    s1 = F.dequantize_blockwise(code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val], A=bnb_optimizer.state[p2][name2], blocksize=blocksize)
+                else:
+                    s1 = F.dequantize(code=bnb_optimizer.state[p2][qmap], absmax=bnb_optimizer.state[p2][max_val], A=bnb_optimizer.state[p2][name2])
+                torch.testing.assert_allclose(s1cpy, s1)
+
+                num_not_close = torch.isclose(torch_optimizer.state[p1][name1], s1, atol=atol, rtol=rtol)==0
+                assert num_not_close.sum().item() < 20
+            torch.testing.assert_allclose(p1, p2.float(), atol=patol, rtol=prtol)
+
+        # the parameters diverge quickly. Here we keep them close
+        # together so we can test against the Adam error
+        p1.data = p1.data.to(gtype).float()
+        p2.copy_(p1.data)
+        torch.testing.assert_allclose(p1.to(gtype), p2)
+        for (name1, name2, qmap, max_val), s in zip(str2statenames[optim_name], dequant_states):
+            torch_optimizer.state[p1][name1].copy_(s.data)
+
+    #print(sum(errors)/len(errors))
+    #print(sum(relerrors)/len(relerrors))
+
+
+
+dim1 = [1024]
+dim2 = [32, 1024, 4097]
+gtype = [torch.float32]
+optim_bits = [32, 8]
+values = list(product(dim1,dim2, gtype, optim_bits))
+names = ['dim1_{0}_dim2_{1}_gtype_{2}_optim_bits_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, gtype, optim_bits", values, ids=names)
+def test_adam_percentile_clipping(dim1, dim2, gtype, optim_bits):
+    if dim1 == 1 and dim2 == 1: return
+    p1 = torch.randn(dim1,dim2, device='cpu', dtype=gtype)*0.1
+    beta1 = 0.9
+    beta2 = 0.999
+    lr = 0.001
+    eps = 1e-8
+    p1 = p1.cuda()
+    p2 = p1.clone()
+    adam1 = bnb.optim.Adam([p1], lr, (beta1, beta2), eps, optim_bits=optim_bits)
+    adam2 = bnb.optim.Adam([p2], lr, (beta1, beta2), eps, optim_bits=optim_bits, percentile_clipping=5)
+
+    gnorm_vec = torch.zeros(100).cuda()
+    step = 0
+
+    for i in range(50):
+        step += 1
+        g1 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1 + (0.01*i)
+        g2 = g1.clone()
+        p2.grad = g2
+
+        current_gnorm, clip_val, gnorm_scale = F.percentile_clipping(g1, gnorm_vec, step, 5)
+        g1 = (g1.float()*gnorm_scale).to(gtype)
+        p1.grad = g1
+
+        adam1.step()
+        adam2.step()
+
+        # gnorm_scale is not deterministic (warp reductions), as such there can be slight differences in state
+        if optim_bits == 32:
+            torch.testing.assert_allclose(p1, p2)
+            torch.testing.assert_allclose(adam1.state[p1]['state1'], adam2.state[p2]['state1'], atol=5e-5, rtol=1e-4)
+            torch.testing.assert_allclose(adam1.state[p1]['state2'], adam2.state[p2]['state2'], atol=5e-5, rtol=1e-4)
+        elif optim_bits == 8:
+            torch.testing.assert_allclose(p1, p2, atol=1e-4, rtol=1e-3)
+            torch.testing.assert_allclose(adam1.state[p1]['state1'], adam2.state[p2]['state1'], atol=2, rtol=1e-3)
+            torch.testing.assert_allclose(adam1.state[p1]['state2'], adam2.state[p2]['state2'], atol=2, rtol=1e-3)
+            adam1.state[p1]['state1'].copy_(adam2.state[p2]['state1'])
+            adam1.state[p1]['state2'].copy_(adam2.state[p2]['state2'])
+        if i % 10 == 0 and i > 0:
+            path = get_temp_dir()
+            torch.save(adam2.state_dict(),join(path, 'opt.pt'))
+            del adam2
+            adam2 = None
+            adam2 = bnb.optim.Adam([p2], lr, (beta1, beta2), eps, optim_bits=optim_bits, percentile_clipping=5)
+            adam2.load_state_dict(torch.load(join(path, 'opt.pt')))
+
+
+
+
+dim1 = [4096]
+dim2 = [4096]
+gtype = [torch.float32, torch.float16]
+#optimizer_names = ['adam8bit_blockwise', 'adam8bit', 'lamb8bit']
+#optimizer_names = ['adam8bit_blockwise', 'adam_apex', 'adam8bit', 'adam', 'adam_pytorch']
+#optimizer_names = ['momentum_apex', 'momentum8bit', 'momentum_pytorch']
+#optimizer_names = ['lamb_apex', 'lamb8bit']
+#optimizer_names = ['lars_apex', 'lars8bit']
+optimizer_names = ['adam8bit_blockwise']
+values = list(product(dim1,dim2, gtype, optimizer_names))
+names = ['dim1_{0}_dim2_{1}_gtype_{2}_optim_{3}'.format(*vals) for vals in values]
+@pytest.mark.parametrize("dim1, dim2, gtype, optim_name", values, ids=names)
+def test_benchmark_blockwise(dim1, dim2, gtype, optim_name):
+    if dim1 == 1 and dim2 == 1: return
+    p1 = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.1
+
+
+    bnb_optimizer = str2optimizers[optim_name][1]([p1])
+
+    g = torch.randn(dim1,dim2, device='cuda', dtype=gtype)*0.01
+    p1.grad = g
+    for i in range(5000):
+        if i == 500:
+            # 100 iterations for burn-in
+            torch.cuda.synchronize()
+            t0 = time.time()
+
+        bnb_optimizer.step()
+
+    torch.cuda.synchronize()
+    s = time.time()-t0
+    print('')
+    params = 4500*4096*4096
+    print(optim_name, gtype, s/params)
+    #assert s < 3.9
+
+