test_triton.py 2.55 KB
Newer Older
Tim Dettmers's avatar
Tim Dettmers committed
1
2
3
import pytest
import torch

Mitchell Wortsman's avatar
Mitchell Wortsman committed
4
from bitsandbytes.nn import Linear8bitLt
Aarni Koskela's avatar
Aarni Koskela committed
5
6
from bitsandbytes.nn.triton_based_modules import SwitchBackLinear
from bitsandbytes.triton.triton_utils import is_triton_available
Aarni Koskela's avatar
Aarni Koskela committed
7
from tests.helpers import TRUE_FALSE
Aarni Koskela's avatar
Aarni Koskela committed
8

Tim Dettmers's avatar
Tim Dettmers committed
9

Ruff's avatar
Ruff committed
10
11
12
13
@pytest.mark.skipif(
    not is_triton_available() or not torch.cuda.is_available() or not torch.cuda.get_device_capability()[0] >= 8,
    reason="This test requires triton and a GPU with compute capability 8.0 or higher.",
)
14
@pytest.mark.deprecated
Aarni Koskela's avatar
Aarni Koskela committed
15
@pytest.mark.parametrize("vector_wise_quantization", TRUE_FALSE)
16
17
18
def test_switchback(vector_wise_quantization):
    for dim in [83]:
        for batch in [13]:
Tim Dettmers's avatar
Tim Dettmers committed
19
            standard = torch.nn.Linear(dim, 4 * dim).cuda().half()
Ruff's avatar
Ruff committed
20
21
22
            switchback = (
                SwitchBackLinear(dim, 4 * dim, vector_wise_quantization=vector_wise_quantization).cuda().half()
            )
Mitchell Wortsman's avatar
Mitchell Wortsman committed
23
            baseline = Linear8bitLt(dim, 4 * dim).cuda().half()
Tim Dettmers's avatar
Tim Dettmers committed
24
25
            switchback.weight.data.copy_(standard.weight)
            switchback.bias.data.copy_(standard.bias)
Mitchell Wortsman's avatar
Mitchell Wortsman committed
26
27
28
29
30
31
            baseline.weight.data.copy_(standard.weight)
            baseline.bias.data.copy_(standard.bias)

            x1 = torch.randn(batch, dim).cuda().half().requires_grad_(True)
            x2 = x1.clone().detach().requires_grad_(True)
            x3 = x1.clone().detach().requires_grad_(True)
Tim Dettmers's avatar
Tim Dettmers committed
32

Mitchell Wortsman's avatar
Mitchell Wortsman committed
33
34
            out_standard = standard(x1)
            (2**10 * out_standard.abs().mean()).backward()
Tim Dettmers's avatar
Tim Dettmers committed
35

36
            print(x2.dtype)
Mitchell Wortsman's avatar
Mitchell Wortsman committed
37
38
            out_sb = switchback(x2)
            (2**10 * out_sb.abs().mean()).backward()
Tim Dettmers's avatar
Tim Dettmers committed
39

Mitchell Wortsman's avatar
Mitchell Wortsman committed
40
41
            out_baseline = baseline(x3)
            (2**10 * out_baseline.abs().mean()).backward()
Tim Dettmers's avatar
Tim Dettmers committed
42

Mitchell Wortsman's avatar
Mitchell Wortsman committed
43
44
            err_sb = (out_standard - out_sb).abs().mean()
            err_baseline = (out_standard - out_baseline).abs().mean()
Ruff's avatar
Ruff committed
45
            print("OUT", err_sb, err_baseline)
Mitchell Wortsman's avatar
Mitchell Wortsman committed
46
            assert err_sb < 2 * err_baseline
Tim Dettmers's avatar
Tim Dettmers committed
47

Mitchell Wortsman's avatar
Mitchell Wortsman committed
48
49
            err_sb = (standard.bias.grad - switchback.bias.grad).abs().mean()
            err_baseline = (standard.bias.grad - baseline.bias.grad).abs().mean()
Tim Dettmers's avatar
Tim Dettmers committed
50

Ruff's avatar
Ruff committed
51
            print("GW2", err_sb, err_baseline)
Mitchell Wortsman's avatar
Mitchell Wortsman committed
52
            assert err_sb < 2 * err_baseline
Tim Dettmers's avatar
Tim Dettmers committed
53

Mitchell Wortsman's avatar
Mitchell Wortsman committed
54
55
            err_sb = (standard.weight.grad - switchback.weight.grad).abs().mean()
            err_baseline = (standard.weight.grad - baseline.weight.grad).abs().mean()
Tim Dettmers's avatar
Tim Dettmers committed
56

Ruff's avatar
Ruff committed
57
            print("GW1", err_sb, err_baseline)
Mitchell Wortsman's avatar
Mitchell Wortsman committed
58
            assert err_sb < 2 * err_baseline
Tim Dettmers's avatar
Tim Dettmers committed
59

Mitchell Wortsman's avatar
Mitchell Wortsman committed
60
61
            err_sb = (x1.grad - x2.grad).abs().mean()
            err_baseline = (x1.grad - x3.grad).abs().mean()
Tim Dettmers's avatar
Tim Dettmers committed
62

Ruff's avatar
Ruff committed
63
            print("GX1", err_sb, err_baseline)
Mitchell Wortsman's avatar
Mitchell Wortsman committed
64
            assert err_sb < 2 * err_baseline