(
device, quant_type, original_dtype, compress_statistics, bias, quant_storage, save_before_forward
)
| 37 | @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) |
| 38 | @pytest.mark.parametrize("save_before_forward", TRUE_FALSE, ids=id_formatter("save_before_forward")) |
| 39 | def test_linear_serialization( |
| 40 | device, quant_type, original_dtype, compress_statistics, bias, quant_storage, save_before_forward |
| 41 | ): |
| 42 | if device == "hpu" and not is_supported_on_hpu(quant_type, original_dtype, storage[quant_storage]): |
| 43 | pytest.skip("This configuration is not supported on HPU.") |
| 44 | |
| 45 | compute_dtype = None |
| 46 | layer_shape = (300, 400) |
| 47 | |
| 48 | linear = torch.nn.Linear(*layer_shape, dtype=original_dtype, device="cpu") # original layer |
| 49 | |
| 50 | # Quantizing original layer |
| 51 | linear_q = bnb.nn.Linear4bit( |
| 52 | linear.in_features, |
| 53 | linear.out_features, |
| 54 | bias=bias, |
| 55 | compute_dtype=compute_dtype, |
| 56 | compress_statistics=compress_statistics, |
| 57 | quant_type=quant_type, |
| 58 | device="meta", |
| 59 | ) |
| 60 | new_weight = bnb.nn.Params4bit(data=linear.weight, quant_type=quant_type, requires_grad=False) |
| 61 | linear_q.weight = new_weight |
| 62 | if bias: |
| 63 | linear_q.bias = torch.nn.Parameter(linear.bias) |
| 64 | linear_q = linear_q.to(device) |
| 65 | |
| 66 | # saving to state_dict: |
| 67 | sd = linear_q.state_dict() |
| 68 | |
| 69 | # restoring from state_dict: |
| 70 | bias_data2 = sd.pop("bias", None) |
| 71 | weight_data2 = sd.pop("weight") |
| 72 | weight2 = bnb.nn.Params4bit.from_prequantized(quantized_stats=sd, data=weight_data2, device=device) |
| 73 | |
| 74 | # creating new layer with same params: |
| 75 | linear_q2 = bnb.nn.Linear4bit( |
| 76 | linear.in_features, |
| 77 | linear.out_features, |
| 78 | bias=bias, |
| 79 | compute_dtype=compute_dtype, |
| 80 | compress_statistics=compress_statistics, |
| 81 | quant_type=quant_type, |
| 82 | device="meta", |
| 83 | ) |
| 84 | # loading weights from state_dict: |
| 85 | linear_q2.weight = weight2 |
| 86 | if bias: |
| 87 | linear_q2.bias = torch.nn.Parameter(bias_data2) |
| 88 | linear_q2 = linear_q2.to(device) |
| 89 | |
| 90 | # MATCHING |
| 91 | a, b = linear_q.weight, linear_q2.weight |
| 92 | |
| 93 | # Quantizing original layer with specified quant_storage type |
| 94 | linear_qs = bnb.nn.Linear4bit( |
| 95 | linear.in_features, |
| 96 | linear.out_features, |
nothing calls this directly
no test coverage detected