diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py index 8a069bd10..e28192b08 100644 --- a/bitsandbytes/autograd/_functions.py +++ b/bitsandbytes/autograd/_functions.py @@ -98,6 +98,11 @@ def reset_grads(self): self.CBt = None +def _to_16bit(t: torch.Tensor) -> torch.Tensor: + """Pass fp16/bf16 tensors through unchanged; cast anything else to fp16 for the int8 quantization kernels.""" + return t if t.dtype in (torch.float16, torch.bfloat16) else t.to(torch.float16) + + class MatMul8bitLt(torch.autograd.Function): @staticmethod def forward( @@ -124,8 +129,8 @@ def forward( input_shape = A.shape - # Cast A to fp16 - if A.dtype != torch.float16 and not _is_compiling(): + # fp16 and bf16 are quantized natively; other dtypes (e.g. fp32) are cast to fp16 for quantization only. + if A.dtype not in (torch.float16, torch.bfloat16) and not _is_compiling(): logger.warning("MatMul8bitLt: inputs will be cast from %s to float16 during quantization", A.dtype) if len(A.shape) == 3: @@ -134,10 +139,10 @@ def forward( # 1. Quantize A. Note that as a side-effect, outliers are suppressed in CA/CAt. if ctx.needs_input_grad[1]: # Slower path - CA, CAt, SCA, SCAt, outlier_cols = F.int8_double_quant(A.to(torch.float16), threshold=state.threshold) + CA, CAt, SCA, SCAt, outlier_cols = F.int8_double_quant(_to_16bit(A), threshold=state.threshold) else: # Fast path - CA, SCA, outlier_cols = F.int8_vectorwise_quant(A.to(torch.float16), threshold=state.threshold) + CA, SCA, outlier_cols = F.int8_vectorwise_quant(_to_16bit(A), threshold=state.threshold) CAt = SCAt = None has_grad = False @@ -152,7 +157,7 @@ def forward( state.reset_grads() # 2. Quantize B - state.CB, state.SCB, _ = F.int8_vectorwise_quant(B.to(torch.float16)) + state.CB, state.SCB, _ = F.int8_vectorwise_quant(_to_16bit(B)) # Handle sparse decomposition if state.threshold > 0.0: @@ -214,19 +219,19 @@ def backward(ctx: torch.autograd.function.FunctionCtx, grad_output: torch.Tensor # compute grad_bias first before changing grad_output dtype grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias) - # Cast grad_output to fp16 if len(grad_output.shape) == 3: grad_output = grad_output.reshape(-1, grad_output.shape[-1]).contiguous() if req_gradB: - Cgrad, _, _, SCgradt, _ = F.int8_double_quant(grad_output.to(torch.float16)) + grad_output_16 = _to_16bit(grad_output) + Cgrad, _, _, SCgradt, _ = F.int8_double_quant(grad_output_16) grad_B = torch.ops.bitsandbytes.int8_scaled_mm.default( Cgrad.t().contiguous(), CAt.t(), SCgradt, SCAt, - dtype=torch.float16, + dtype=grad_output_16.dtype, ) if state.threshold > 0.0 and subA is not None and subA.numel() > 0: @@ -258,7 +263,7 @@ def forward(ctx, A, B, out=None, bias=None, state=MatmulLtState): if (state.is_training and not has_grad) or state.CB is None or state.SCB is None: state.reset_grads() - state.CB, state.SCB, _ = F.int8_vectorwise_quant(B.to(torch.float16)) + state.CB, state.SCB, _ = F.int8_vectorwise_quant(_to_16bit(B)) B = state.CB CB = state.CB.data.to(A.dtype).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0)) diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index a0d9ffe83..603a84db5 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -48,13 +48,13 @@ def _setup_ctypes(names, argtypes, restype=None): # int8 mm dequant: (A, row_stats, col_stats, out, bias, numRows, numCols, stream) _setup_ctypes( - ["cdequant_mm_int32_fp16"], + ["cdequant_mm_int32_fp16", "cdequant_mm_int32_bf16"], [ct.c_void_p] * 5 + [ct.c_int32, ct.c_int32, ct.c_void_p], ) # int8 vectorwise quant: (A, out, row_stats, threshold, rows, cols, stream) _setup_ctypes( - ["cint8_vector_quant"], + ["cint8_vector_quant", "cint8_vector_quant_bf16"], [ct.c_void_p] * 3 + [ct.c_float, ct.c_int32, ct.c_int32, ct.c_void_p], ) @@ -183,37 +183,39 @@ def _( if col_stats.dtype != torch.float32: raise ValueError(f"col_stats must be float32, got {col_stats.dtype}") - # Note: cuda kernel only currently supports fp16 output. - # We'll later cast to desired dtype if needed. - out = torch.empty_like(A, dtype=torch.float16) + # The kernel writes fp16 or bf16 directly. Any other requested dtype (e.g. fp32) is computed + # in fp16 and cast afterwards, matching the previous behavior. + out_dtype = dtype or torch.float16 + kernel_dtype = torch.bfloat16 if out_dtype == torch.bfloat16 else torch.float16 + out = torch.empty_like(A, dtype=kernel_dtype) - # Note: fused bias in the kernel is only supported for fp16 - # TODO(matthewdouglas): Consider supporting bf16 fused bias - bias_ptr = bias.data_ptr() if bias is not None and bias.dtype == torch.float16 else None + # Fuse bias in the kernel when it already matches the kernel dtype; otherwise add it afterward. + fuse_bias = bias is not None and bias.dtype == kernel_dtype + fn = lib.cdequant_mm_int32_bf16 if kernel_dtype == torch.bfloat16 else lib.cdequant_mm_int32_fp16 with _cuda_device_of(A): - lib.cdequant_mm_int32_fp16( + fn( A.data_ptr(), row_stats.data_ptr(), col_stats.data_ptr(), out.data_ptr(), - bias_ptr, + bias.data_ptr() if fuse_bias else None, A.numel() // A.shape[-1], A.shape[-1], _get_raw_stream(A.device.index), ) # Add bias separately if not fused in kernel - if bias is not None and bias.dtype != torch.float16: + if bias is not None and not fuse_bias: out.add_(bias) - return out.to(dtype or torch.float16) + return out.to(out_dtype) @register_kernel("bitsandbytes::int8_vectorwise_quant", "cuda") def _(A: torch.Tensor, threshold=0.0): - if A.dtype != torch.float16: - raise ValueError(f"A must be float16, got {A.dtype}") + if A.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"A must be float16 or bfloat16, got {A.dtype}") if threshold < 0.0: raise ValueError("threshold must be non-negative") @@ -235,8 +237,9 @@ def _(A: torch.Tensor, threshold=0.0): # Needed for torch.compile support. outlier_cols = torch.empty(0, device=A.device, dtype=torch.int64) + fn = lib.cint8_vector_quant_bf16 if A.dtype == torch.bfloat16 else lib.cint8_vector_quant with _cuda_device_of(A): - lib.cint8_vector_quant( + fn( A.data_ptr(), out_row.data_ptr(), row_stats.data_ptr(), diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py index 33b2cd9ac..9f7c43979 100644 --- a/bitsandbytes/functional.py +++ b/bitsandbytes/functional.py @@ -1565,6 +1565,7 @@ def int8_mm_dequant( col_stats: torch.Tensor, out: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, ): """Performs dequantization on the result of a quantized int8 matrix multiplication. @@ -1574,11 +1575,14 @@ def int8_mm_dequant( col_stats (`torch.Tensor`): The column-wise quantization statistics for the rhs operand of the matrix multiplication. out (`torch.Tensor`, *optional*): A pre-allocated tensor to store the output of the operation. bias (`torch.Tensor`, *optional*): An optional bias vector to add to the result. + dtype (`torch.dtype`, *optional*): The desired output dtype. Defaults to `torch.float16`. Returns: - `torch.Tensor`: The dequantized result with an optional bias, with dtype `torch.float16`. + `torch.Tensor`: The dequantized result with an optional bias, with the requested dtype. """ - result = torch.ops.bitsandbytes.int8_mm_dequant.default(A, row_stats, col_stats, dtype=torch.float16, bias=bias) + result = torch.ops.bitsandbytes.int8_mm_dequant.default( + A, row_stats, col_stats, dtype=dtype or torch.float16, bias=bias + ) # TODO(matthewdouglas): Deprecate out kwarg if out is not None: @@ -1607,7 +1611,7 @@ def int8_double_quant( Args: - A (`torch.Tensor` with dtype `torch.float16`): The input matrix. + A (`torch.Tensor` with dtype `torch.float16` or `torch.bfloat16`): The input matrix. col_stats (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the column-wise quantization scales. row_stats (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the row-wise quantization scales. out_col (`torch.Tensor`, *optional*): A pre-allocated tensor to hold the column-wise quantized data. @@ -1653,12 +1657,12 @@ def int8_vectorwise_dequant(A: torch.Tensor, stats: torch.Tensor): def int8_vectorwise_quant(A: torch.Tensor, threshold=0.0): - """Quantizes a tensor with dtype `torch.float16` to `torch.int8` in accordance to the `LLM.int8()` algorithm. + """Quantizes a tensor with dtype `torch.float16` or `torch.bfloat16` to `torch.int8` in accordance to the `LLM.int8()` algorithm. For more information, see the [LLM.int8() paper](https://arxiv.org/abs/2208.07339). Args: - A (`torch.Tensor` with dtype `torch.float16`): The input tensor. + A (`torch.Tensor` with dtype `torch.float16` or `torch.bfloat16`): The input tensor. threshold (`float`, *optional*): An optional threshold for sparse decomposition of outlier features. diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..783ddf170 100644 --- a/bitsandbytes/nn/modules.py +++ b/bitsandbytes/nn/modules.py @@ -739,7 +739,9 @@ def _quantize(self, device): return super().to(device) # We quantize the weight and store in 8bit row-major - B = self.data.contiguous().to(device=device, dtype=torch.float16) + # bf16 weights are quantized natively; everything else goes through fp16. + dtype = torch.bfloat16 if self.data.dtype == torch.bfloat16 else torch.float16 + B = self.data.contiguous().to(device=device, dtype=dtype) CB, SCB, _ = bnb.functional.int8_vectorwise_quant(B) self.data = CB self.CB = CB diff --git a/csrc/kernels.cu b/csrc/kernels.cu index df0f093df..9455f388a 100644 --- a/csrc/kernels.cu +++ b/csrc/kernels.cu @@ -7,6 +7,8 @@ #include "compat_device.cuh" #include "kernels.cuh" +#include + #define HLF_MAX 65504 #define TH 1024 #define NUM 4 @@ -1328,41 +1330,56 @@ __launch_bounds__(256, 3) __global__ void kOptimizerStatic8bit1StateBlockwise( // Outputs: // rowStats [rows] // out [rows, cols] +// Streaming (evict-first) load of a 16-bit float, widened to fp32. +// CUDA < 12.4 only declares the __nv_bfloat16 overload of __ldcs for sm_80+. The half overload is available on +// every architecture, so for bf16 we issue the streaming load through it and reinterpret the 16 raw bits. +template __device__ __forceinline__ float ldcs_as_float(const T* ptr) { + if constexpr (std::is_same_v) { + return __half2float(__ldcs(ptr)); + } else { +#if BNB_HIP + return (float)(*ptr); +#else + const __half_raw bits = __ldcs(reinterpret_cast(ptr)); + __nv_bfloat16_raw raw; + raw.x = bits.x; + return __bfloat162float(__nv_bfloat16(raw)); +#endif + } +} + template __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ void kInt8VectorQuant(T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols) { - - using BlockReduceT = bnb_cub::BlockReduce; - // One block per row. // Threads load column values in a striped arrangement. // e.g. t0 reads row[0], row[0+nthreads], .. // and t1 reads row[1], row[1+nthreads], .. // Each thread will determine its local absmax. // We then do a blockwise reduction to determine the row's absmax. - + using BlockReduceT = bnb_cub::BlockReduce; __shared__ typename BlockReduceT::TempStorage temp_storage; - __shared__ T smem_row_absmax; + __shared__ float smem_row_absmax; const int row_id = blockIdx.x; const T* row_data = A + (row_id * cols); // Threads will read the row values in a striped access pattern and find a local absmax. - T row_local_absmax = -FLT_MIN; + float row_local_absmax = -FLT_MIN; for (int i = threadIdx.x; i < cols; i += THREADS) { - const T absval = fabsf(__ldcs(&(row_data[i]))); + const float absval = fabsf(ldcs_as_float(&row_data[i])); // For sparse decomposition, values outside of the threshold are not to be // included when calculating the row's absmax. if constexpr (SPARSE_DECOMP) { - row_local_absmax = fmaxf(row_local_absmax, absval < T(threshold) ? absval : row_local_absmax); + row_local_absmax = fmaxf(row_local_absmax, absval < threshold ? absval : row_local_absmax); } else { row_local_absmax = fmaxf(row_local_absmax, absval); } } // Reduce thread-local absmax across the block. - const T row_absmax = BlockReduceT(temp_storage).Reduce(row_local_absmax, BNB_MAX_OP, cols); + const float row_absmax = BlockReduceT(temp_storage).Reduce(row_local_absmax, BNB_MAX_OP, cols); if (threadIdx.x == 0) { // Save our block's absmax to shared memory for the quantization step. rowStats[row_id] = smem_row_absmax = row_absmax; @@ -1372,12 +1389,12 @@ __launch_bounds__(1024, BNB_MAX_THREADS_PER_SM / 1024) __global__ // Quantize row-wise. const float scale = __fdividef(127.0f, smem_row_absmax); for (int i = threadIdx.x; i < cols; i += THREADS) { - float val = row_data[i]; + float val = (float)row_data[i]; if constexpr (SPARSE_DECOMP) { // For sparse decomposition, we do not want to quantize the outliers. // Instead they're zeroed out. - out[row_id * cols + i] = fabs(val) < threshold ? __float2int_rn(val * scale) : 0; + out[row_id * cols + i] = fabsf(val) < threshold ? __float2int_rn(val * scale) : 0; } else { out[row_id * cols + i] = __float2int_rn(val * scale); } @@ -1390,13 +1407,19 @@ template __global__ void kInt8VectorQuant( template __global__ void kInt8VectorQuant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols ); +template __global__ void kInt8VectorQuant( + bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols +); +template __global__ void kInt8VectorQuant( + bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols +); #define MM_DEQUANT_CONST 6.200012e-05f // 1.0f/(127.0f*127.0f) -template -__global__ void kdequant_mm_int32_fp16( - int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out, - half* __restrict__ const bias, const int numRows, const int numCols, const int n +template +__global__ void kdequant_mm_int32( + int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, T* out, + T* __restrict__ const bias, const int numRows, const int numCols, const int n ) { const int n_out = numRows * numCols; @@ -1404,7 +1427,7 @@ __global__ void kdequant_mm_int32_fp16( int thread_offset = threadIdx.x * ITEMS_PER_THREAD; int local_values[ITEMS_PER_THREAD]; - half local_output[ITEMS_PER_THREAD]; + T local_output[ITEMS_PER_THREAD]; float local_rowStats[ITEMS_PER_THREAD]; float local_colStats[ITEMS_PER_THREAD]; @@ -1423,7 +1446,7 @@ __global__ void kdequant_mm_int32_fp16( local_colStats[j] = col_idx >= numCols ? 0.0f : __ldg(&colStats[col_idx]); local_rowStats[j] = row_idx >= numRows ? 0.0f : __ldg(&rowStats[row_idx]); - local_biasValue[j] = ((bias == nullptr) || col_idx >= numCols) ? 0.0f : __half2float(bias[col_idx]); + local_biasValue[j] = ((bias == nullptr) || col_idx >= numCols) ? 0.0f : (float)(bias[col_idx]); } // Each block loads THREADS * ITEMS_PER_THREAD values from A @@ -1433,9 +1456,8 @@ __global__ void kdequant_mm_int32_fp16( #pragma unroll ITEMS_PER_THREAD for (int j = 0; j < ITEMS_PER_THREAD; ++j) { - local_output[j] = __float2half( - fmaf(local_values[j] * local_rowStats[j] * local_colStats[j], MM_DEQUANT_CONST, local_biasValue[j]) - ); + local_output[j] = + (T)(fmaf(local_values[j] * local_rowStats[j] * local_colStats[j], MM_DEQUANT_CONST, local_biasValue[j])); } #pragma unroll ITEMS_PER_THREAD @@ -1604,10 +1626,14 @@ template __global__ void kgemm_4bit_inference_naive( float* out, int lda, int ldb, int ldc, int blocksize ); -template __global__ void kdequant_mm_int32_fp16<4, 512>( +template __global__ void kdequant_mm_int32( int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out, half* __restrict__ const bias, const int numRows, const int numCols, const int n ); +template __global__ void kdequant_mm_int32( + int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, + bnb_bfloat16* out, bnb_bfloat16* __restrict__ const bias, const int numRows, const int numCols, const int n +); template __device__ unsigned char dQuantize<0>(float* smem_code, const float rand, float x); template __device__ unsigned char dQuantize<1>(float* smem_code, const float rand, float x); diff --git a/csrc/kernels.cuh b/csrc/kernels.cuh index dc511661b..6ef2fcd67 100644 --- a/csrc/kernels.cuh +++ b/csrc/kernels.cuh @@ -65,10 +65,10 @@ __global__ void kOptimizerStatic8bit1StateBlockwise( const float gnorm_scale, const bool skip_zeros, const int n ); -template -__global__ void kdequant_mm_int32_fp16( - int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, half* out, - half* __restrict__ const bias, const int numRows, const int numCols, const int n +template +__global__ void kdequant_mm_int32( + int* __restrict__ const A, float* __restrict__ const rowStats, float* __restrict__ const colStats, T* out, + T* __restrict__ const bias, const int numRows, const int numCols, const int n ); template diff --git a/csrc/ops.cu b/csrc/ops.cu index f6f890455..74e3d4936 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -495,8 +495,9 @@ int fill_up_to_nearest_multiple(int value, int multiple) { return value + (value % multiple == 0 ? 0 : (multiple - (value % multiple))); } -void dequant_mm_int32_fp16( - int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream +template +void dequant_mm_int32( + int* A, float* rowStats, float* colStats, T* out, T* bias, int numRows, int numCols, bnb_stream_t stream ) { const int threads = 512; const int num_per_thread = 4; @@ -504,18 +505,19 @@ void dequant_mm_int32_fp16( const int n = numRows * numCols; const int num_blocks = (n + num_per_block - 1) / num_per_block; - kdequant_mm_int32_fp16 + kdequant_mm_int32 <<>>(A, rowStats, colStats, out, bias, numRows, numCols, n); BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } +template void int8VectorQuant( - half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream + T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream ) { if (threshold == 0.0) { - kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); + kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); } else { - kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); + kInt8VectorQuant<<>>(A, out, rowStats, threshold, rows, cols); } BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } @@ -569,6 +571,21 @@ template void gemm_4bit_inference_naive( int ldc, int blocksize, bnb_stream_t stream ); +template void int8VectorQuant( + half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream +); +template void int8VectorQuant( + bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream +); + +template void dequant_mm_int32( + int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream +); +template void dequant_mm_int32( + int* A, float* rowStats, float* colStats, bnb_bfloat16* out, bnb_bfloat16* bias, int numRows, int numCols, + bnb_stream_t stream +); + template int igemmlt<32, 0>( bnb_blasLt_handle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda, int ldb, int ldc, bnb_stream_t stream diff --git a/csrc/ops.cuh b/csrc/ops.cuh index c7114bcaa..d0e0dc62a 100644 --- a/csrc/ops.cuh +++ b/csrc/ops.cuh @@ -133,11 +133,15 @@ int igemmlt( void cutlass_igemm( bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc ); -void dequant_mm_int32_fp16( - int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream + +template +void dequant_mm_int32( + int* A, float* rowStats, float* colStats, T* out, T* bias, int numRows, int numCols, bnb_stream_t stream ); + +template void int8VectorQuant( - half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream + T* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, bnb_stream_t stream ); template diff --git a/csrc/pythonInterface.cpp b/csrc/pythonInterface.cpp index 24431603b..9b8a17779 100644 --- a/csrc/pythonInterface.cpp +++ b/csrc/pythonInterface.cpp @@ -545,13 +545,26 @@ int cigemmlt_8_rowscale( void cdequant_mm_int32_fp16( int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, cudaStream_t stream ) { - dequant_mm_int32_fp16(A, rowStats, colStats, out, bias, numRows, numCols, stream); + dequant_mm_int32(A, rowStats, colStats, out, bias, numRows, numCols, stream); +} + +void cdequant_mm_int32_bf16( + int* A, float* rowStats, float* colStats, bnb_bfloat16* out, bnb_bfloat16* bias, int numRows, int numCols, + cudaStream_t stream +) { + dequant_mm_int32(A, rowStats, colStats, out, bias, numRows, numCols, stream); } void cint8_vector_quant( half* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream ) { - int8VectorQuant(A, out, rowStats, threshold, rows, cols, stream); + int8VectorQuant(A, out, rowStats, threshold, rows, cols, stream); +} + +void cint8_vector_quant_bf16( + bnb_bfloat16* __restrict__ A, int8_t* out, float* rowStats, float threshold, int rows, int cols, cudaStream_t stream +) { + int8VectorQuant(A, out, rowStats, threshold, rows, cols, stream); } void* cget_managed_ptr(size_t bytes) { diff --git a/tests/test_autograd.py b/tests/test_autograd.py index 7d273c853..aafb76a4c 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -76,7 +76,7 @@ def test_matmullt( if not transpose[0] and not transpose[1]: B2 = B2.t().contiguous() - state.CB, state.SCB, _ = bnb.functional.int8_vectorwise_quant(B2.to(torch.float16)) + state.CB, state.SCB, _ = bnb.functional.int8_vectorwise_quant(B2) B2 = state.CB if not transpose[0] and transpose[1]: diff --git a/tests/test_functional.py b/tests/test_functional.py index e4cd6a128..4214f050b 100644 --- a/tests/test_functional.py +++ b/tests/test_functional.py @@ -397,22 +397,20 @@ def test_int8_linear_matmul(self, device, dim1, dim2, dim3, dim4, dims, ldb): @pytest.mark.parametrize("dim3", [32], ids=id_formatter("dim3")) @pytest.mark.parametrize("dim4", [32], ids=id_formatter("dim4")) @pytest.mark.parametrize("dims", (2,), ids=id_formatter("dims")) - def test_int8_linear_matmul_half(self, device, dim1, dim2, dim3, dim4, dims): + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) + def test_int8_linear_matmul_half(self, device, dim1, dim2, dim3, dim4, dims, dtype): for i in range(k): if dims == 2: - A = torch.normal(0, 0.5, size=(dim1, dim3), device=device).half() + A = torch.normal(0, 0.5, size=(dim1, dim3), device=device).to(dtype) elif dims == 3: - A = torch.normal(0, 0.5, size=(dim1, dim2, dim3), device=device).half() - B = torch.randn((dim4, dim3), device=device).half() + A = torch.normal(0, 0.5, size=(dim1, dim2, dim3), device=device).to(dtype) + B = torch.randn((dim4, dim3), device=device).to(dtype) torch.nn.init.xavier_uniform_(B) C1 = torch.matmul(A, B.t()) - A = A.view(-1, A.shape[-1]) - CA, statsA, _ = F.int8_vectorwise_quant(A) CB, statsB, _ = F.int8_vectorwise_quant(B) - output = F.int8_mm_dequant(F.int8_linear_matmul(CA, CB), statsA, statsB) - + output = F.int8_mm_dequant(F.int8_linear_matmul(CA, CB), statsA, statsB, dtype=dtype) torch.testing.assert_close(C1.view(-1, C1.shape[-1]), output, atol=0.025, rtol=0.05) @pytest.mark.parametrize("device", get_available_devices()) diff --git a/tests/test_linear8bitlt.py b/tests/test_linear8bitlt.py index b078e82b7..25bc8bd38 100644 --- a/tests/test_linear8bitlt.py +++ b/tests/test_linear8bitlt.py @@ -13,6 +13,7 @@ from bitsandbytes.nn.modules import Linear8bitLt from tests.helpers import ( TRUE_FALSE, + describe_dtype, get_available_devices, id_formatter, torch_load_from_buffer, @@ -364,3 +365,32 @@ def test_linear8bitlt_device_movement(device): # Accelerator outputs should match both times. torch.testing.assert_close(out_accelerator_2, out_accelerator, rtol=1e-8, atol=1e-8) + + +@pytest.mark.parametrize("device", get_available_devices(no_cpu=True)) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) +@pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold")) +def test_linear8bitlt_forward_dtypes(device, dtype, threshold): + # Exercises the activation-quant path end-to-end through Linear8bitLt for both fp16 and bf16. + 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=threshold, + ) + 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.to(device) + + x = torch.randn(4, 32, dtype=dtype, device=device) + out = linear_custom(x) + + assert out.dtype == dtype + assert out.shape == (4, 96) + assert out.isfinite().all() diff --git a/tests/test_modules.py b/tests/test_modules.py index 95f78b6d3..42f38eea7 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -7,7 +7,7 @@ from torch import nn import bitsandbytes as bnb -from tests.helpers import get_available_devices, id_formatter, is_supported_on_hpu +from tests.helpers import describe_dtype, get_available_devices, id_formatter, is_supported_on_hpu @contextlib.contextmanager @@ -61,15 +61,15 @@ def assert_all_approx_close(a, b, atol=1e-8, rtol=1e-5, count=10): @pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("threshold", [0.0, 3.0], ids=id_formatter("threshold")) -def test_linear8bitlt_inference(device, threshold): - l1 = bnb.nn.Linear8bitLt(32, 64, threshold=threshold, has_fp16_weights=False).to(device).half() +def test_linear8bitlt_inference(device, dtype, threshold): + l1 = bnb.nn.Linear8bitLt(32, 64, threshold=threshold, has_fp16_weights=False).to(device).to(dtype) assert l1.weight.device.type == device assert l1.weight.dtype == torch.int8 - l1.eval() for i in range(100): - b1 = torch.randn(16, 8, 32, device=device).half() + b1 = torch.randn(16, 8, 32, device=device, dtype=dtype) o1 = l1(b1) if i == 1: assert l1.state.CB is not None diff --git a/tests/test_ops.py b/tests/test_ops.py index 4ca60f845..aef624f40 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -36,21 +36,19 @@ def test_int8_linear_matmul_out(self, device): opcheck(torch.ops.bitsandbytes.int8_linear_matmul.out, (A, B, out)) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("threshold", [0.0, 6.0]) @pytest.mark.parametrize("device", get_available_devices()) - def test_int8_vectorwise_quant(self, threshold, device): - A = torch.randn(10, 20, dtype=torch.float16, device=device) + def test_int8_vectorwise_quant(self, dtype, threshold, device): + A = torch.randn(10, 20, dtype=dtype, device=device) A[1][0] = 1000.0 - out_row, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant(A, threshold=threshold) - assert out_row.shape == (10, 20) assert out_row.dtype == torch.int8 assert out_row.device == A.device assert row_stats.shape == (10,) assert row_stats.dtype == torch.float32 assert row_stats.device == A.device - if threshold > 0.0: assert outlier_cols is not None assert outlier_cols.dim() == 1 @@ -58,7 +56,6 @@ def test_int8_vectorwise_quant(self, threshold, device): assert outlier_cols.device == A.device else: assert outlier_cols is None - opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A,)) opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold)) @@ -75,6 +72,29 @@ def test_int8_mm_dequant(self, device): opcheck(torch.ops.bitsandbytes.int8_mm_dequant, (A, row_stats, col_stats)) + @pytest.mark.parametrize("device", get_available_devices()) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype) + @pytest.mark.parametrize( + "bias_dtype", [None, torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("bias_dtype") + ) + def test_int8_mm_dequant_dtypes(self, device, dtype, bias_dtype): + # Covers the native fp16/bf16 kernel outputs, the fused-bias path (bias dtype == output dtype), + # the unfused path, and fp32 output (computed in fp16, then cast). + A = torch.randint(-127 * 127, 127 * 127, (64, 96), dtype=torch.int32, device=device) + row_stats = torch.rand(64, dtype=torch.float32, device=device) * 4 + col_stats = torch.rand(96, dtype=torch.float32, device=device) * 4 + bias = torch.randn(96, dtype=bias_dtype, device=device) if bias_dtype is not None else None + + out = torch.ops.bitsandbytes.int8_mm_dequant(A, row_stats, col_stats, dtype=dtype, bias=bias) + + ref = A.float() * row_stats.unsqueeze(1) * col_stats.unsqueeze(0) / (127.0 * 127.0) + if bias is not None: + ref = ref + bias.float() + + assert out.dtype == dtype + assert out.shape == A.shape + torch.testing.assert_close(out.float(), ref, rtol=1e-2, atol=1e-2) + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("has_bias", TRUE_FALSE)