| 135 | |
| 136 | |
| 137 | class Test4bitBlockwiseQuantOps: |
| 138 | @pytest.mark.parametrize("device", get_available_devices()) |
| 139 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) |
| 140 | @pytest.mark.parametrize("storage_dtype", [torch.uint8, torch.bfloat16], ids=id_formatter("storage_dtype")) |
| 141 | @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) |
| 142 | @pytest.mark.parametrize("blocksize", [32, 64, 128, 256, 512]) |
| 143 | def test_quantize_4bit(self, device, dtype, storage_dtype, quant_type, blocksize): |
| 144 | if device == "hpu" and not is_supported_on_hpu(quant_type, dtype, storage_dtype): |
| 145 | pytest.skip("This configuration is not supported on HPU.") |
| 146 | |
| 147 | A = torch.randn(1024, 1024, dtype=dtype, device=device) |
| 148 | |
| 149 | out, absmax = torch.ops.bitsandbytes.quantize_4bit.default(A, blocksize, quant_type, storage_dtype) |
| 150 | |
| 151 | assert out.device == A.device |
| 152 | assert out.dtype == storage_dtype |
| 153 | |
| 154 | assert absmax.device == A.device |
| 155 | assert absmax.dtype == torch.float32 |
| 156 | |
| 157 | if storage_dtype != torch.uint8: |
| 158 | pytest.xfail("opcheck fails for storage_dtype != torch.uint8") |
| 159 | |
| 160 | opcheck(torch.ops.bitsandbytes.quantize_4bit.default, (A, blocksize, quant_type, storage_dtype)) |
| 161 | |
| 162 | @pytest.mark.parametrize("device", get_available_devices()) |
| 163 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) |
| 164 | @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) |
| 165 | @pytest.mark.parametrize("blocksize", [64, 128, 256]) |
| 166 | def test_quantize_4bit_not_divisible_by_blocksize(self, device, dtype, quant_type, blocksize): |
| 167 | """Test quantize/dequantize roundtrip when n_elements is not divisible by blocksize.""" |
| 168 | # Shape chosen so numel is NOT divisible by blocksize |
| 169 | shape = (7, blocksize - 1) |
| 170 | A = torch.randn(shape, dtype=dtype, device=device) |
| 171 | storage_dtype = torch.uint8 |
| 172 | |
| 173 | # Should not raise |
| 174 | packed, absmax = torch.ops.bitsandbytes.quantize_4bit(A, blocksize, quant_type, storage_dtype) |
| 175 | |
| 176 | assert packed.device == A.device |
| 177 | assert absmax.device == A.device |
| 178 | |
| 179 | # Dequantize back and verify shape is preserved |
| 180 | out = torch.ops.bitsandbytes.dequantize_4bit(packed, absmax, blocksize, quant_type, shape, dtype) |
| 181 | |
| 182 | assert out.shape == shape |
| 183 | assert out.dtype == dtype |
| 184 | |
| 185 | # Verify output is finite (no NaN/Inf) |
| 186 | assert torch.isfinite(out).all(), "Dequantized output contains NaN or Inf" |
| 187 | |
| 188 | @pytest.mark.parametrize("device", get_available_devices()) |
| 189 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) |
| 190 | @pytest.mark.parametrize("storage_dtype", [torch.uint8, torch.bfloat16], ids=id_formatter("storage_dtype")) |
| 191 | @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) |
| 192 | @pytest.mark.parametrize("blocksize", [32, 64, 128, 256, 512]) |
| 193 | def test_dequantize_4bit(self, device, dtype, storage_dtype, quant_type, blocksize): |
| 194 | if device == "hpu" and not is_supported_on_hpu(quant_type, dtype, storage_dtype): |
nothing calls this directly
no outgoing calls
no test coverage detected