test_linear8bitlt.py 7.96 KB
Newer Older
1
from contextlib import nullcontext
2
import copy
Aarni Koskela's avatar
Aarni Koskela committed
3
import os
4
import pickle
5
from tempfile import TemporaryDirectory
6

7
8
9
import pytest
import torch

10
import bitsandbytes as bnb
11
from bitsandbytes.nn.modules import Linear8bitLt
12
13
from tests.helpers import (
    TRUE_FALSE,
14
    get_available_devices,
15
16
17
18
    id_formatter,
    torch_load_from_buffer,
    torch_save_to_buffer,
)
19

20

21
22
# contributed by Alex Borzunov, see:
# https://github.com/bigscience-workshop/petals/blob/main/tests/test_linear8bitlt.py
23
24
@pytest.mark.parametrize("device", get_available_devices())
def test_linear_no_igemmlt(device):
25
26
27
28
29
30
31
32
33
    linear = torch.nn.Linear(1024, 3072)
    x = torch.randn(3, 1024, dtype=torch.half)
    linear_custom = Linear8bitLt(
        linear.in_features,
        linear.out_features,
        linear.bias is not None,
        has_fp16_weights=False,
        threshold=6.0,
    )
34
35

    # TODO: Remove, this is no longer implemented
36
37
38
    linear_custom.state.force_no_igemmlt = True

    linear_custom.weight = bnb.nn.Int8Params(
Ruff's avatar
Ruff committed
39
40
41
        linear.weight.data.clone(),
        requires_grad=False,
        has_fp16_weights=False,
42
43
    ).to(linear.weight.dtype)
    linear_custom.bias = linear.bias
44
45
    linear_custom = linear_custom.to(device)
    linear = linear.half().to(device)
46

47
48
    x_ref = x.clone().to(device).requires_grad_(True)
    x_ours = x.clone().to(device).requires_grad_(True)
49
50
51
52
53
54
    fx_ref = linear(x_ref).float()
    grad_proj = torch.randn_like(fx_ref)
    (fx_ref * grad_proj).mean().backward()

    fx_ours = linear_custom(x_ours).float()
    (fx_ours * grad_proj).mean().backward()
55

56
    assert linear_custom.state.CB is not None
57
58
59
60
61
62
    assert not linear_custom.state.has_fp16_weights

    idx = torch.isclose(fx_ref, fx_ours, atol=0.02, rtol=1e-5)
    assert (idx == 0).sum().item() < fx_ref.numel() * 2.5e-4
    torch.testing.assert_close(fx_ref, fx_ours, atol=0.03, rtol=1e-5)
    torch.testing.assert_close(x_ref.grad, x_ours.grad, atol=0.01, rtol=1e-5)
63
64


65
@pytest.mark.parametrize("device", get_available_devices())
Aarni Koskela's avatar
Aarni Koskela committed
66
@pytest.mark.parametrize("has_fp16_weights", TRUE_FALSE, ids=id_formatter("has_fp16_weights"))
67
@pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold"))
Aarni Koskela's avatar
Aarni Koskela committed
68
69
@pytest.mark.parametrize("serialize_before_forward", TRUE_FALSE, ids=id_formatter("serialize_before_forward"))
@pytest.mark.parametrize("deserialize_before_cuda", TRUE_FALSE, ids=id_formatter("deserialize_before_cuda"))
70
71
@pytest.mark.parametrize("save_before_forward", TRUE_FALSE, ids=id_formatter("save_before_forward"))
@pytest.mark.parametrize("load_before_cuda", TRUE_FALSE, ids=id_formatter("load_before_cuda"))
Ruff's avatar
Ruff committed
72
def test_linear_serialization(
73
    device,
Ruff's avatar
Ruff committed
74
    has_fp16_weights,
75
    threshold,
Ruff's avatar
Ruff committed
76
77
78
79
80
    serialize_before_forward,
    deserialize_before_cuda,
    save_before_forward,
    load_before_cuda,
):
81
82
    if device != "cuda" and has_fp16_weights:
        pytest.skip("has_fp16_weights is only supported on CUDA and is deprecated")
83

84
    linear = torch.nn.Linear(32, 96)
85
86
87
    # TODO: Fallback for bad shapes
    x = torch.randn(4, 32, dtype=torch.half)
    # x = torch.randn(3, 32, dtype=torch.half)
88
89
90
91
92
93

    linear_custom = Linear8bitLt(
        linear.in_features,
        linear.out_features,
        linear.bias is not None,
        has_fp16_weights=has_fp16_weights,
94
        threshold=threshold,
95
    )
96

97
    linear_custom.weight = bnb.nn.Int8Params(
Ruff's avatar
Ruff committed
98
99
100
        linear.weight.data.clone(),
        requires_grad=has_fp16_weights,
        has_fp16_weights=has_fp16_weights,
101
    )
102
    linear_custom.bias = linear.bias
103
    linear_custom = linear_custom.to(device)
104

105
106
107
    if serialize_before_forward:
        state_dict_8bit = linear_custom.state_dict()

108
109
110
    if save_before_forward:
        bytes_8bit = torch_save_to_buffer(linear_custom)

111
    x_first = x.clone().to(device).requires_grad_(True)
112
113
114
115
    fx_first = linear_custom(x_first).float()
    grad_proj = torch.randn_like(fx_first)
    (fx_first * grad_proj).mean().backward()

116
117
118
    if not serialize_before_forward:
        state_dict_8bit = linear_custom.state_dict()

119
120
121
    if not save_before_forward:
        bytes_8bit = torch_save_to_buffer(linear_custom)

122
123
124
125
126
127
128
129
130
131
    with TemporaryDirectory() as tmpdir:
        state_path_8bit = os.path.join(tmpdir, "state_8bit.pth")
        state_path = os.path.join(tmpdir, "state.pth")

        torch.save(linear.state_dict(), state_path)
        torch.save(state_dict_8bit, state_path_8bit)

        if not has_fp16_weights:
            assert os.path.getsize(state_path_8bit) < 0.5 * os.path.getsize(state_path)

132
        new_state_dict = torch.load(state_path_8bit, weights_only=False)
133
134
135
136
137
138

    new_linear_custom = Linear8bitLt(
        linear.in_features,
        linear.out_features,
        linear.bias is not None,
        has_fp16_weights=has_fp16_weights,
139
        threshold=threshold,
140
    )
141
142
143
144
145

    if deserialize_before_cuda:
        with nullcontext() if has_fp16_weights else pytest.raises(RuntimeError):
            new_linear_custom.load_state_dict(new_state_dict, strict=True)

146
147
148
    if load_before_cuda:
        new_linear_custom2 = torch_load_from_buffer(bytes_8bit)

149
    new_linear_custom = new_linear_custom.to(device)
150
151
152

    if not deserialize_before_cuda:
        new_linear_custom.load_state_dict(new_state_dict, strict=True)
153

154
155
156
    if not load_before_cuda:
        new_linear_custom2 = torch_load_from_buffer(bytes_8bit)

157
    x_second = x.clone().to(device).requires_grad_(True)
158
159
160
    fx_second = new_linear_custom(x_second).float()
    (fx_second * grad_proj).mean().backward()

161
    x_third = x.clone().to(device).requires_grad_(True)
162
163
164
    fx_third = new_linear_custom2(x_third).float()
    (fx_third * grad_proj).mean().backward()

165
166
167
168
    # if 8-bit weights were loaded before .cuda, state is incorrect anyway and RuntimeError was raised
    if has_fp16_weights or not deserialize_before_cuda:
        assert torch.allclose(fx_first, fx_second, atol=1e-5)
        assert torch.allclose(x_first.grad, x_second.grad, atol=1e-5)
169
    assert torch.allclose(fx_first, fx_third, atol=1e-5)
170
    assert torch.allclose(x_first.grad, x_third.grad, atol=1e-5)
171
172
173


@pytest.fixture
174
def linear8bit(requires_cuda):
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
    linear = torch.nn.Linear(32, 96)
    linear_custom = Linear8bitLt(
        linear.in_features,
        linear.out_features,
        linear.bias is not None,
        has_fp16_weights=False,
        threshold=6.0,
    )
    linear_custom.weight = bnb.nn.Int8Params(
        linear.weight.data.clone(),
        requires_grad=False,
        has_fp16_weights=False,
    )
    linear_custom.bias = linear.bias
    linear_custom = linear_custom.cuda()
    return linear_custom


def test_linear8bit_copy_param(linear8bit):
    shallow_copy = copy.copy(linear8bit)
    assert linear8bit.weight is shallow_copy.weight
    assert linear8bit.bias is shallow_copy.bias
    assert linear8bit.weight.data.data_ptr() == shallow_copy.weight.data.data_ptr()


def test_linear8bit_deepcopy_param(linear8bit):
    deep_copy = copy.deepcopy(linear8bit)
    assert linear8bit.weight is not deep_copy.weight
    assert linear8bit.bias is not deep_copy.bias
    assert linear8bit.weight.data.data_ptr() != deep_copy.weight.data.data_ptr()
    assert torch.allclose(linear8bit.weight.data, deep_copy.weight.data)
    assert linear8bit.state == deep_copy.state

    # check for a bug where SCB and CB were not copied
    assert deep_copy.weight.SCB is not None
    assert (linear8bit.weight.SCB == deep_copy.weight.SCB).all()
    assert deep_copy.weight.CB is not None
    assert (linear8bit.weight.CB == deep_copy.weight.CB).all()


def test_linear8bit_serialization(linear8bit):
    serialized = pickle.dumps(linear8bit)
    deserialized = pickle.loads(serialized)
    assert linear8bit.weight.data.data_ptr() != deserialized.weight.data.data_ptr()
    assert torch.allclose(linear8bit.weight.data, deserialized.weight.data)
    assert linear8bit.bias.data.data_ptr() != deserialized.bias.data.data_ptr()
    assert torch.allclose(linear8bit.bias.data, deserialized.bias.data)
    assert linear8bit.state == deserialized.state

    # check for a bug where SCB and CB were not copied
    assert (linear8bit.weight.SCB == deserialized.weight.SCB).all()
    assert (linear8bit.weight.CB == deserialized.weight.CB).all()