diff options
| author | Tim Dettmers <TimDettmers@users.noreply.github.com> | 2022-09-19 21:09:25 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-09-19 21:09:25 -0700 |
| commit | 439f2b0c10abd3e9aade386d92810b074c69e9ec (patch) | |
| tree | 75454081c86ba1c96c07e83defc9fc5f4de840cf /tests/test_autograd.py | |
| parent | 9b5f2eda8fbd3f042c4af7ed1b870525d4668f2a (diff) | |
| parent | 76ce9aa6da7d68d2463f0f3e99532ab5b6db58a8 (diff) | |
Merge pull request #33 from dbaranchuk/memory-efficient-backward
Memory efficient backward
Diffstat (limited to 'tests/test_autograd.py')
| -rw-r--r-- | tests/test_autograd.py | 9 |
1 files changed, 6 insertions, 3 deletions
diff --git a/tests/test_autograd.py b/tests/test_autograd.py index bae26de..40bb441 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -253,7 +253,7 @@ for c in req_grad: transpose = [(False, True), (False, False)] str_transpose = ["NT", "NN"] -dtype = [torch.float16] +dtype = [torch.float16, torch.bfloat16, torch.float32] has_fp16_weights = [True, False] has_bias = [True, False] values = list( @@ -354,7 +354,7 @@ def test_matmullt( state.SCB, SCBt, coo_tensorB, - ) = bnb.functional.double_quant(B2) + ) = bnb.functional.double_quant(B2.to(torch.float16)) B2 = state.CB if not transpose[0] and transpose[1]: @@ -367,11 +367,14 @@ def test_matmullt( if has_bias: out_torch += bias + assert out_bnb.dtype == A.dtype, f"bnb matmullt received {A.dtype} but returned {out_bnb.dtype}" + n = out_bnb.numel() err = torch.abs(out_bnb - out_torch).mean().item() # print(f'abs error {err:.4f}') + idx = torch.isclose(out_bnb, out_torch, atol=0.01, rtol=0.1) - assert (idx == 0).sum().item() <= n * 0.0175 + assert (idx == 0).sum().item() <= n * (0.0175 if dtype == torch.float16 else 0.021) idx = torch.isclose(out_bnb, out_torch, atol=0.035, rtol=0.2) assert (idx == 0).sum().item() <= n * 0.001 |