Function test_kbit_backprop

tests/test_modules.py:244–309

(device, module, dtype)

@pytest.mark.parametrize("module", module_dict.values(), ids=module_dict.keys())
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
def test_kbit_backprop(device, module, dtype):
    b = 16
    dim1 = 36
    dim2 = 84
    # dim1 = 37
    # dim2 = 83

    ref = nn.Sequential(*[torch.nn.Linear(dim1, dim2), torch.nn.Linear(dim2, 128)])
    torch.nn.init.kaiming_normal_(ref[0].weight)
    torch.nn.init.kaiming_normal_(ref[1].weight)
    ref[1].weight.requires_grad_(False)

    kbit = nn.Sequential(*[torch.nn.Linear(dim1, dim2), module(dim2, 128)])

    if (
        device == "hpu"
        and isinstance(kbit[1], bnb.nn.Linear4bit)
        and not is_supported_on_hpu(kbit[1].weight.quant_type, dtype)
    ):
        pytest.skip("This configuration not supported on HPU")

    kbit[0].weight.detach().copy_(ref[0].weight)
    kbit[1].weight.detach().copy_(ref[1].weight)
    kbit[0].bias.detach().copy_(ref[0].bias)
    kbit[1].bias.detach().copy_(ref[1].bias)
    kbit[1].weight.requires_grad_(False)
    ref = ref.to(device=device, dtype=dtype)
    kbit = kbit.to(device=device, dtype=dtype)
    kbit = kbit.to(device=device, dtype=dtype)

    errs1 = []
    errs2 = []
    relerrs1 = []
    relerrs2 = []
    for i in range(100):
        batch = torch.randn(b, dim1, device=device, dtype=dtype)
        out1 = ref(batch)
        out2 = kbit(batch)
        out1.mean().backward()
        out2.mean().backward()

        grad1 = ref[0].weight.grad
        grad2 = kbit[0].weight.grad
        bgrad1 = ref[0].bias.grad
        bgrad2 = kbit[0].bias.grad

        err1 = (out1 - out2).abs().float()
        err2 = (grad1 - grad2).abs().float()
        relerr1 = err1 / (out1.abs().float() + 1e-9)
        relerr2 = err2 / (grad1.abs().float() + 1e-9)
        errs1.append(err1.mean().item())
        errs2.append(err2.mean().item())
        relerrs1.append(relerr1.mean().item())
        relerrs2.append(relerr2.mean().item())

        if isinstance(module, bnb.nn.Linear8bitLt):
            assert_all_approx_close(grad1, grad2, atol=0.008, rtol=0.05, count=1)
            torch.testing.assert_close(bgrad1, bgrad2, atol=0.008, rtol=0.05)