ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/unsloth/kernels/utils.py
Daniel Han cd9d251f15
 Fix fast inference crash on compressed-tensors FP8 models (#7025) 
* Fix fast_gemv crash on compressed-tensors FP8 models

Loading a compressed-tensors FP8 checkpoint (for example
unsloth/Llama-3.2-1B-Instruct-FP8-Block) with fast_inference=False and
running a forward crashed with 'Parameter object has no attribute absmax'
inside fast_gemv.

A compressed-tensors CompressedLinear exposes an already dequantized bf16
weight at forward time while keeping a weight_scale Parameter. The quant
state resolution in get_lora_parameters/get_lora_parameters_bias fell back
to that weight_scale, so a bf16 weight was routed into the bitsandbytes
fast_gemv/fast_dequantize path, which expects a bitsandbytes QuantState
with an absmax attribute.

Only fall back to weight_scale_inv/weight_scale when the weight is still
fp8. A decompressed bf16 weight then resolves to no quant state and flows
through the normal bf16 path, which already handles bias and the LoRA
backward. Real fp8 and bitsandbytes 4bit weights are unchanged.

* Skip the fast_gemv dispatch test before importing unsloth when bitsandbytes is absent
2026-07-09 04:10:59 -07:00

1129 lines
35 KiB
Python
Raw Blame History

# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import importlib
import triton
import ctypes
MAX_FUSED_SIZE: int = 65536
next_power_of_2 = triton.next_power_of_2
import functools
from typing import Optional
from ..device_type import (
is_hip,
get_device_type,
DEVICE_TYPE,
DEVICE_TYPE_TORCH,
DEVICE_COUNT,
ALLOW_PREQUANTIZED_MODELS,
)
from .fp8 import weight_dequant, fp8_linear
import functools
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
import torch
torch_Tensor = torch.Tensor
from unsloth_zoo.utils import Version
if DEVICE_TYPE == "xpu" and Version(torch.__version__) < Version("2.6.0"):
raise RuntimeError("Intel xpu currently supports unsloth with torch.version >= 2.6.0")
if Version(torch.__version__) < Version("2.4.0"):
torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
else:
torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "cuda")
torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda")
if DEVICE_TYPE == "xpu":
torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "xpu")
torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "xpu")
# tl.math.tanh now is libdevice.tanh
import triton
import triton.language as tl
if Version(triton.__version__) >= Version("3.0.0"):
if DEVICE_TYPE == "xpu":
triton_tanh = tl.extra.intel.libdevice.tanh
else:
from triton.language.extra import libdevice
triton_tanh = libdevice.tanh
triton_cast = tl.cast
else:
triton_tanh = tl.math.tanh
# No casting in old Triton versions
@triton.jit
def triton_cast(x, dtype):
return x.to(dtype)
@functools.lru_cache(1)
def is_cdna():
return is_hip() and triton.runtime.driver.active.get_current_target().arch in (
"gfx940",
"gfx941",
"gfx942",
"gfx950", # CDNA4 (MI350/MI355X)
)
@functools.lru_cache(1)
def is_rdna():
"""Detect ROCm-supported RDNA consumer/workstation GPUs (RDNA2, RDNA3, RDNA3.5, RDNA4)."""
return is_hip() and triton.runtime.driver.active.get_current_target().arch in (
# RDNA2 (Navi 21-24)
"gfx1030",
"gfx1031",
"gfx1032",
"gfx1033",
"gfx1034",
"gfx1035",
"gfx1036",
# RDNA3 (Navi 31-33)
"gfx1100",
"gfx1101",
"gfx1102",
"gfx1103",
# RDNA3.5 (Strix Point / Strix Halo)
"gfx1150",
"gfx1151",
"gfx1152",
# RDNA4 (Navi 48-44)
"gfx1200",
"gfx1201",
)
def calculate_settings(
n: int,
) -> (
int,
int,
):
BLOCK_SIZE: int = next_power_of_2(n)
if BLOCK_SIZE > MAX_FUSED_SIZE:
raise RuntimeError(
f"Cannot launch Triton kernel since n = {n} exceeds "
f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}."
)
num_warps: int = 4
if BLOCK_SIZE >= 32768:
num_warps = 32
elif BLOCK_SIZE >= 8192:
num_warps = 16
elif BLOCK_SIZE >= 2048:
num_warps = 8
return BLOCK_SIZE, num_warps
HAS_CUDA_STREAM = False
import bitsandbytes as bnb
# https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files
HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3")
get_ptr = bnb.functional.get_ptr
if DEVICE_TYPE == "xpu":
HAS_XPU_STREAM = True
if DEVICE_COUNT > 1:
if DEVICE_TYPE in ("cuda", "hip"):
torch_gpu_device = torch.cuda.device
elif DEVICE_TYPE == "xpu":
torch_gpu_device = torch.xpu.device
else:
from contextlib import nullcontext
def torch_gpu_device(device):
return nullcontext()
# INTEL GPU Specific Logic
if DEVICE_TYPE == "xpu":
_gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream
elif DEVICE_TYPE == "mlx":
def _gpu_getCurrentRawStream(_index = 0):
return 0
# NVIDIA GPU Default Logic
elif hasattr(torch._C, "_cuda_getCurrentRawStream"):
_gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream
else:
# CPU-only torch wheel (no compiled CUDA backend). _get_tensor_stream
# is only invoked during real GPU work, so a no-op binding is safe.
def _gpu_getCurrentRawStream(_index = 0):
return 0
c_void_p = ctypes.c_void_p
def _get_tensor_stream(tensor: torch_Tensor) -> c_void_p:
return c_void_p(_gpu_getCurrentRawStream(tensor.device.index))
# Get array of CUDA streams and other buffers
global CUDA_STREAMS
global XPU_STREAMS
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS
# DEVICE_COUNT == 0 = no visible accelerator (e.g. CPU-only CI runner).
# The consumer functions below only index these arrays during real GPU
# work, so empty containers are safe -- they just need to be defined so
# the module imports cleanly.
if DEVICE_TYPE == "xpu":
if DEVICE_COUNT > 0:
_XPU_STREAMS = {
(index := torch.xpu.device(i).idx): ctypes.c_void_p(
torch._C._xpu_getCurrentRawStream(index)
)
for i in range(DEVICE_COUNT)
}
XPU_STREAMS = [None] * (max(_XPU_STREAMS.keys()) + 1)
WEIGHT_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1)
ABSMAX_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1)
for k, v in _XPU_STREAMS.items():
XPU_STREAMS[k] = v
XPU_STREAMS = tuple(XPU_STREAMS)
del _XPU_STREAMS
else:
XPU_STREAMS = ()
WEIGHT_BUFFERS = []
ABSMAX_BUFFERS = []
elif DEVICE_TYPE == "mlx":
CUDA_STREAMS = ()
XPU_STREAMS = ()
WEIGHT_BUFFERS = []
ABSMAX_BUFFERS = []
else:
# NVIDIA GPU Default Logic
if DEVICE_COUNT > 0:
_CUDA_STREAMS = {
(index := torch.cuda.device(i).idx): ctypes.c_void_p(
torch._C._cuda_getCurrentRawStream(index)
)
for i in range(DEVICE_COUNT)
}
CUDA_STREAMS = [None] * (max(_CUDA_STREAMS.keys()) + 1)
WEIGHT_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1)
ABSMAX_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1)
for k, v in _CUDA_STREAMS.items():
CUDA_STREAMS[k] = v
CUDA_STREAMS = tuple(CUDA_STREAMS)
del _CUDA_STREAMS
else:
CUDA_STREAMS = ()
WEIGHT_BUFFERS = []
ABSMAX_BUFFERS = []
# Bitsandbytes operations
ctypes_c_int = ctypes.c_int
ctypes_c_int32 = ctypes.c_int32
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4
cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4
if DEVICE_TYPE == "xpu":
# https://github.com/bitsandbytes-foundation/bitsandbytes/blob/c3b8de268fdb55a88f92feada23fc811a1e6877a/bitsandbytes/backends/xpu/ops.py#L115
# for xpu, inference gemv using above link
cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemv_4bit_inference_fp16
cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemv_4bit_inference_bf16
else:
cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemm_4bit_inference_naive_fp16
cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemm_4bit_inference_naive_bf16
torch_device_stream = (
torch.xpu.current_stream if DEVICE_TYPE == "xpu" else torch.cuda.current_stream
)
torch_mm = torch.mm
torch_mv = torch.mv
torch_matmul = torch.matmul
torch_addmm = torch.addmm
torch_empty = torch.empty
torch_float32 = torch.float32
torch_float16 = torch.float16
torch_bfloat16 = torch.bfloat16
# Check whether torchao can be imported to get Float8Tensor
if importlib.util.find_spec("torchao") is not None:
try:
from torchao.quantization import Float8Tensor
except:
import torchao
if Version(torchao.__version__) >= Version("0.15.0"):
print(
f"Unsloth: `from torchao.quantization import Float8Tensor` failed on version={torchao.__version__}"
)
Float8Tensor = type(None)
else:
Float8Tensor = type(None)
def QUANT_STATE(W):
return getattr(W, "quant_state", None)
# fp8 weight dtypes. A `weight_scale` / `weight_scale_inv` should only be treated as a
# quant state when the weight itself is still fp8. compressed-tensors layers expose an
# already-dequantized bf16 weight at forward time while keeping a `weight_scale` around;
# reading that as a quant state routes a bf16 weight into the bitsandbytes fast_gemv /
# fast_dequantize path, which then reads a missing `absmax` and crashes.
_FP8_WEIGHT_DTYPES = tuple(
dtype
for dtype in (
getattr(torch, "float8_e4m3fn", None),
getattr(torch, "float8_e5m2", None),
)
if dtype is not None
)
def get_lora_parameters(proj):
"""Return (weight, weight quant_state, lora A, lora B, lora scale).
With QAT enabled, also fake-quantizes the base layer and lora weights.
"""
# For DPO or disabled adapters
base_layer = getattr(
proj, "base_layer", proj
) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
W = base_layer.weight
# Optionally apply fake quantization to base layer weights for QAT
if hasattr(base_layer, "weight_fake_quantizer"):
weight_fake_quantizer = getattr(base_layer, "weight_fake_quantizer", None)
if weight_fake_quantizer is not None:
W = weight_fake_quantizer(W)
# Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
# weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
# must not carry a scale as its quant state or fast_gemv will crash on it.
W_quant = getattr(W, "quant_state", None)
if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
W_quant = getattr(base_layer, "weight_scale_inv", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)
if getattr(base_layer, "quant_method", None) == "fp8":
# we need to somehow store and pass this information :)
W.block_size = getattr(base_layer, "block_size", [128, 128])
W_quant.block_size = W.block_size
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
if getattr(proj, "disable_adapters", True) or proj.merged:
return W, W_quant, None, None, None
adapter = getattr(proj, "active_adapters", None)
if adapter is None:
adapter = getattr(proj, "active_adapter", ("default"))
adapter = adapter[0]
# Optionally apply fake quantization to lora weights for QAT
lora_A_linear = proj.lora_A[adapter]
lora_B_linear = proj.lora_B[adapter]
A = lora_A_linear.weight
B = lora_B_linear.weight
if hasattr(lora_A_linear, "weight_fake_quantizer"):
lora_A_fake_quantizer = getattr(lora_A_linear, "weight_fake_quantizer", None)
if lora_A_fake_quantizer is not None:
A = lora_A_fake_quantizer(A)
if hasattr(lora_B_linear, "weight_fake_quantizer"):
lora_B_fake_quantizer = getattr(lora_B_linear, "weight_fake_quantizer", None)
if lora_B_fake_quantizer is not None:
B = lora_B_fake_quantizer(B)
return (
W,
W_quant,
A,
B,
proj.scaling[adapter],
)
def get_lora_parameters_bias(proj):
# For DPO or disabled adapters
base_layer = getattr(
proj, "base_layer", proj
) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
W = base_layer.weight
# Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
# weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
# must not carry a scale as its quant state or fast_gemv will crash on it.
W_quant = getattr(W, "quant_state", None)
if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
W_quant = getattr(base_layer, "weight_scale_inv", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)
# if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
if getattr(proj, "disable_adapters", True) or proj.merged:
return W, W_quant, None, None, None, base_layer.bias
if getattr(base_layer, "quant_method", None) == "fp8":
# we need to somehow store and pass this information :)
W.block_size = getattr(base_layer, "block_size", [128, 128])
W_quant.block_size = W.block_size
adapter = getattr(proj, "active_adapters", None)
if adapter is None:
adapter = getattr(proj, "active_adapter", ("default"))
adapter = adapter[0]
return (
W,
W_quant,
proj.lora_A[adapter].weight,
proj.lora_B[adapter].weight,
proj.scaling[adapter],
base_layer.bias,
)
def _maybe_fake_quantize_activations(X: torch.Tensor, proj: torch.nn.Module) -> torch.Tensor:
"""Fake-quantize input activations if QAT is enabled, else return as-is.
Weights are fake-quantized separately in `get_lora_parameters`.
"""
base_layer = getattr(proj, "base_layer", proj)
activation_fake_quantizer = getattr(base_layer, "activation_fake_quantizer", None)
if activation_fake_quantizer is not None:
X = activation_fake_quantizer(X)
return X
# INTEL GPU Specific Logic
if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM:
@torch.inference_mode
def fast_dequantize(
W,
quant_state = None,
out = None,
use_global_buffer = False,
):
# TODO: After adding XPU BNB support, check this function
if isinstance(W, Float8Tensor):
return W.dequantize()
if quant_state is None:
return W
if W.dtype == torch.float8_e4m3fn:
return weight_dequant(W, quant_state)
if type(quant_state) is not list:
# New quant_state as a class
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
absmax = quant_state.absmax
shape = quant_state.shape
dtype = quant_state.dtype
blocksize = quant_state.blocksize
offset = quant_state.offset
state2 = quant_state.state2
absmax2 = state2.absmax
code2 = state2.code
blocksize2 = state2.blocksize
else:
# Old quant_state as a list of lists
absmax, shape, dtype, blocksize, compressed_stats, _, _ = quant_state
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
global XPU_STREAMS
device = W.device
device_index = device.index
XPU_STREAM = XPU_STREAMS[device_index]
n_elements_absmax = absmax.numel()
# Create weight matrix
if use_global_buffer:
# Use same buffers for faster inference
size = shape[0] * shape[1]
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS
WEIGHT_BUFFER = WEIGHT_BUFFERS[device_index]
ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index]
if WEIGHT_BUFFER is None or WEIGHT_BUFFER.dtype != dtype:
WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch_empty(
size, dtype = dtype, device = device, requires_grad = False
)
ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(
n_elements_absmax,
dtype = torch.float32,
device = device,
requires_grad = False,
)
if size > WEIGHT_BUFFER.numel():
WEIGHT_BUFFER.resize_(size)
if n_elements_absmax > ABSMAX_BUFFER.numel():
ABSMAX_BUFFER.resize_(n_elements_absmax)
out = WEIGHT_BUFFER[:size].view(shape)
out_absmax = ABSMAX_BUFFER[:n_elements_absmax]
else:
if out is None:
out = torch_empty(shape, dtype = dtype, device = device, requires_grad = False)
else:
assert out.shape == shape
assert out.dtype == dtype
out_absmax = torch_empty(
n_elements_absmax,
dtype = torch_float32,
device = device,
requires_grad = False,
)
# NF4 dequantization of statistics
ptr_out_absmax = get_ptr(out_absmax)
with torch_gpu_device(device):
cdequantize_blockwise_fp32(
get_ptr(code2),
get_ptr(absmax),
get_ptr(absmax2),
ptr_out_absmax,
ctypes_c_int(blocksize2),
ctypes_c_int(n_elements_absmax),
XPU_STREAM,
)
out_absmax += offset
# Dequantize W
fx = (
cdequantize_blockwise_fp16_nf4
if dtype == torch_float16
else cdequantize_blockwise_bf16_nf4
)
fx(
get_ptr(None),
get_ptr(W),
ptr_out_absmax,
get_ptr(out),
ctypes_c_int(blocksize),
ctypes_c_int(out.numel()),
XPU_STREAM,
)
# Careful returning transposed data
is_transposed = True if W.shape[0] == 1 else False
return out.t() if is_transposed else out
# NVIDIA GPU Default Logic
elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM:
@torch.inference_mode
def fast_dequantize(
W,
quant_state = None,
out = None,
use_global_buffer = False,
):
if isinstance(W, Float8Tensor):
return W.dequantize()
if quant_state is None:
return W
if W.dtype == torch.float8_e4m3fn:
return weight_dequant(W, quant_state)
if type(quant_state) is not list:
# New quant_state as a class
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
absmax = quant_state.absmax
shape = quant_state.shape
dtype = quant_state.dtype
blocksize = quant_state.blocksize
offset = quant_state.offset
state2 = quant_state.state2
absmax2 = state2.absmax
code2 = state2.code
blocksize2 = state2.blocksize
else:
# Old quant_state as a list of lists
absmax, shape, dtype, blocksize, compressed_stats, _, _ = quant_state
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
pass
global CUDA_STREAMS
device = W.device
device_index = device.index
CUDA_STREAM = CUDA_STREAMS[device_index]
n_elements_absmax = absmax.numel()
# Create weight matrix
if use_global_buffer:
# Use same buffers for faster inference
size = shape[0] * shape[1]
global WEIGHT_BUFFERS
global ABSMAX_BUFFERS
WEIGHT_BUFFER = WEIGHT_BUFFERS[device_index]
ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index]
if WEIGHT_BUFFER is None or WEIGHT_BUFFER.dtype != dtype:
WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch_empty(
size, dtype = dtype, device = device, requires_grad = False
)
ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(
n_elements_absmax,
dtype = torch_float32,
device = device,
requires_grad = False,
)
if size > WEIGHT_BUFFER.numel():
WEIGHT_BUFFER.resize_(size)
if n_elements_absmax > ABSMAX_BUFFER.numel():
ABSMAX_BUFFER.resize_(n_elements_absmax)
out = WEIGHT_BUFFER[:size].view(shape)
out_absmax = ABSMAX_BUFFER[:n_elements_absmax]
else:
if out is None:
out = torch_empty(shape, dtype = dtype, device = device, requires_grad = False)
else:
assert out.shape == shape
assert out.dtype == dtype
out_absmax = torch_empty(
n_elements_absmax,
dtype = torch_float32,
device = device,
requires_grad = False,
)
pass
# NF4 dequantization of statistics
ptr_out_absmax = get_ptr(out_absmax)
with torch_gpu_device(device):
cdequantize_blockwise_fp32(
get_ptr(code2),
get_ptr(absmax),
get_ptr(absmax2),
ptr_out_absmax,
ctypes_c_int(blocksize2),
ctypes_c_int(n_elements_absmax),
CUDA_STREAM,
)
out_absmax += offset
# Dequantize W
fx = (
cdequantize_blockwise_fp16_nf4
if dtype == torch_float16
else cdequantize_blockwise_bf16_nf4
)
fx(
get_ptr(None),
get_ptr(W),
ptr_out_absmax,
get_ptr(out),
ctypes_c_int(blocksize),
ctypes_c_int(out.numel()),
CUDA_STREAM,
)
pass
# Careful returning transposed data
is_transposed = True if W.shape[0] == 1 else False
return out.t() if is_transposed else out
pass
else:
@torch.inference_mode
def fast_dequantize(
W,
quant_state = None,
out = None,
use_global_buffer = False,
):
if isinstance(W, Float8Tensor):
return W.dequantize()
if quant_state is None:
return W
if W.dtype == torch.float8_e4m3fn:
return weight_dequant(W, quant_state)
if type(quant_state) is not list:
# New quant_state as a class
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
absmax = quant_state.absmax
shape = quant_state.shape
dtype = quant_state.dtype
blocksize = quant_state.blocksize
offset = quant_state.offset
state2 = quant_state.state2
absmax2 = state2.absmax
code2 = state2.code
blocksize2 = state2.blocksize
else:
# Old quant_state as a list of lists
absmax, shape, dtype, blocksize, compressed_stats, _, _ = quant_state
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
pass
n_elements_absmax = absmax.numel()
device = W.device
# Create weight matrix
if out is None:
out = torch_empty(shape, dtype = dtype, device = device, requires_grad = False)
else:
assert out.shape == shape
assert out.dtype == dtype
out_absmax = torch_empty(
n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False
)
# Do dequantization
ptr_out_absmax = get_ptr(out_absmax)
cdequantize_blockwise_fp32(
get_ptr(code2),
get_ptr(absmax),
get_ptr(absmax2),
ptr_out_absmax,
ctypes_c_int(blocksize2),
ctypes_c_int(n_elements_absmax),
)
out_absmax += offset
fx = (
cdequantize_blockwise_fp16_nf4
if dtype == torch_float16
else cdequantize_blockwise_bf16_nf4
)
fx(
get_ptr(None),
get_ptr(W),
ptr_out_absmax,
get_ptr(out),
ctypes_c_int(blocksize),
ctypes_c_int(out.numel()),
)
# Careful returning transposed data
is_transposed = True if W.shape[0] == 1 else False
return out.t() if is_transposed else out
pass
# INTEL GPU Specific Logic
if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM:
def fast_gemv(
X,
W,
quant_state,
out = None,
):
if quant_state is None:
return torch_matmul(X, W, out = out)
# For fast X @ W where seq_len == 1
# From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469
_, q_len, hd = X.shape
# assert(q_len == 1)
if type(quant_state) is not list:
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
absmax = quant_state.absmax
shape = quant_state.shape
dtype = quant_state.dtype
blocksize = quant_state.blocksize
stats = quant_state.code
offset = quant_state.offset
state2 = quant_state.state2
absmax2 = state2.absmax
code2 = state2.code
blocksize2 = state2.blocksize
else:
absmax, shape, dtype, blocksize, compressed_stats, quant_type, stats = quant_state
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
global XPU_STREAMS
device = W.device
device_index = device.index
XPU_STREAM = XPU_STREAMS[device_index]
# assert(dtype == X.dtype)
bout = shape[0]
if out is None:
out = torch_empty(
(
1,
1,
bout,
),
dtype = dtype,
device = device,
)
# else:
# assert(out.shape == (1, 1, bout,))
# pass
if DEVICE_TYPE == "xpu":
m = 1
n = shape[0]
else:
n = 1
m = shape[0]
k = shape[1]
lda = shape[0]
ldc = shape[0]
ldb = (hd + 1) // 2
m = ctypes_c_int32(m)
n = ctypes_c_int32(n)
k = ctypes_c_int32(k)
lda = ctypes_c_int32(lda)
ldb = ctypes_c_int32(ldb)
ldc = ctypes_c_int32(ldc)
df = torch_empty(absmax.shape, dtype = torch_float32, device = device)
with torch_gpu_device(device):
cdequantize_blockwise_fp32(
get_ptr(code2),
get_ptr(absmax),
get_ptr(absmax2),
get_ptr(df),
ctypes_c_int(blocksize2),
ctypes_c_int(df.numel()),
XPU_STREAM,
)
df += offset
absmax = df
fx = (
cgemm_4bit_inference_naive_fp16
if dtype == torch_float16
else cgemm_4bit_inference_naive_bf16
)
blocksize = ctypes_c_int32(blocksize)
fx(
m,
n,
k,
get_ptr(X),
get_ptr(W),
get_ptr(absmax),
get_ptr(stats),
get_ptr(out),
lda,
ldb,
ldc,
blocksize,
XPU_STREAM,
)
return out
elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM:
def fast_gemv(
X,
W,
quant_state,
out = None,
):
if quant_state is None:
return torch_matmul(X, W, out = out)
# For fast X @ W where seq_len == 1
# From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469
_, q_len, hd = X.shape
# assert(q_len == 1)
if type(quant_state) is not list:
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
absmax = quant_state.absmax
shape = quant_state.shape
dtype = quant_state.dtype
blocksize = quant_state.blocksize
stats = quant_state.code
offset = quant_state.offset
state2 = quant_state.state2
absmax2 = state2.absmax
code2 = state2.code
blocksize2 = state2.blocksize
else:
absmax, shape, dtype, blocksize, compressed_stats, quant_type, stats = quant_state
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
pass
global CUDA_STREAMS
device = W.device
device_index = device.index
CUDA_STREAM = CUDA_STREAMS[device_index]
# assert(dtype == X.dtype)
bout = shape[0]
if out is None:
out = torch_empty(
(
1,
1,
bout,
),
dtype = dtype,
device = device,
)
# else:
# assert(out.shape == (1, 1, bout,))
# pass
n = 1
m = shape[0]
k = shape[1]
lda = shape[0]
ldc = shape[0]
ldb = (hd + 1) // 2
m = ctypes_c_int32(m)
n = ctypes_c_int32(n)
k = ctypes_c_int32(k)
lda = ctypes_c_int32(lda)
ldb = ctypes_c_int32(ldb)
ldc = ctypes_c_int32(ldc)
df = torch_empty(absmax.shape, dtype = torch_float32, device = device)
with torch_gpu_device(device):
cdequantize_blockwise_fp32(
get_ptr(code2),
get_ptr(absmax),
get_ptr(absmax2),
get_ptr(df),
ctypes_c_int(blocksize2),
ctypes_c_int(df.numel()),
CUDA_STREAM,
)
df += offset
absmax = df
fx = (
cgemm_4bit_inference_naive_fp16
if dtype == torch_float16
else cgemm_4bit_inference_naive_bf16
)
blocksize = ctypes_c_int32(blocksize)
fx(
m,
n,
k,
get_ptr(X),
get_ptr(W),
get_ptr(absmax),
get_ptr(stats),
get_ptr(out),
lda,
ldb,
ldc,
blocksize,
CUDA_STREAM,
)
pass
return out
pass
else:
def fast_gemv(
X,
W,
quant_state,
out = None,
):
if quant_state is None:
return torch_matmul(X, W, out = out)
# For fast X @ W where seq_len == 1
# From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469
_, q_len, hd = X.shape
# assert(q_len == 1)
if type(quant_state) is not list:
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
absmax = quant_state.absmax
shape = quant_state.shape
dtype = quant_state.dtype
blocksize = quant_state.blocksize
stats = quant_state.code
offset = quant_state.offset
state2 = quant_state.state2
absmax2 = state2.absmax
code2 = state2.code
blocksize2 = state2.blocksize
else:
absmax, shape, dtype, blocksize, compressed_stats, quant_type, stats = quant_state
offset, state2 = compressed_stats
absmax2, code2, blocksize2, _, _, _, _ = state2
pass
# assert(dtype == X.dtype)
bout = shape[0]
device = W.device
if out is None:
out = torch_empty(
(
1,
1,
bout,
),
dtype = dtype,
device = device,
)
# else:
# assert(out.shape == (1, 1, bout,))
# pass
n = 1
m = shape[0]
k = shape[1]
lda = shape[0]
ldc = shape[0]
ldb = (hd + 1) // 2
m = ctypes_c_int32(m)
n = ctypes_c_int32(n)
k = ctypes_c_int32(k)
lda = ctypes_c_int32(lda)
ldb = ctypes_c_int32(ldb)
ldc = ctypes_c_int32(ldc)
df = torch_empty(absmax.shape, dtype = torch_float32, device = device)
cdequantize_blockwise_fp32(
get_ptr(code2),
get_ptr(absmax),
get_ptr(absmax2),
get_ptr(df),
ctypes_c_int(blocksize2),
ctypes_c_int(df.numel()),
)
df += offset
absmax = df
fx = (
cgemm_4bit_inference_naive_fp16
if dtype == torch_float16
else cgemm_4bit_inference_naive_bf16
)
blocksize = ctypes_c_int32(blocksize)
fx(
m,
n,
k,
get_ptr(X),
get_ptr(W),
get_ptr(absmax),
get_ptr(stats),
get_ptr(out),
lda,
ldb,
ldc,
blocksize,
)
return out
pass
def fast_linear_forward(
proj,
X,
temp_lora = None,
out = None,
):
W, W_quant, lora_A, lora_B, lora_S, bias = get_lora_parameters_bias(proj)
bsz, q_len, in_dim = X.shape
if q_len != 1:
return matmul_lora(X, W, W_quant, lora_A, lora_B, lora_S)
if W_quant is None:
out = torch_matmul(X, W.t(), out = out)
elif W.dtype == torch.float8_e4m3fn:
out = fp8_linear(X, W, W_quant, bias)
elif bsz == 1 and q_len == 1:
out = fast_gemv(X, W, W_quant, out = out)
else:
W = fast_dequantize(W.t(), W_quant, use_global_buffer = True)
out = torch_matmul(X, W, out = out)
# Add in LoRA weights
if lora_A is not None:
out_dim = out.shape[2]
dtype = X.dtype
if not hasattr(lora_A, "_fast_lora"):
lora_A._fast_lora = lora_A.to(dtype)
lora_B._fast_lora = lora_B.to(dtype)
if bsz == 1:
out = out.view(out_dim)
temp_lora = torch_mv(lora_A._fast_lora, X.ravel(), out = temp_lora)
out.addmv_(lora_B._fast_lora, temp_lora, alpha = lora_S)
else:
out = out.view(bsz, out_dim)
temp_lora = torch_mm(X.view(bsz, in_dim), lora_A._fast_lora.t(), out = temp_lora)
out.addmm_(temp_lora, lora_B._fast_lora.t(), alpha = lora_S)
out = out.view(bsz, 1, out_dim)
if bias is not None:
out += bias
return out
def matmul_lora(
X,
W,
W_quant,
A,
B,
s,
out = None,
):
dtype = X.dtype
if X.dim() == 3:
batch, seq_len, d = X.shape
X = X.view(-1, X.shape[-1])
reshape = True
else:
reshape = False
if isinstance(W, Float8Tensor):
assert W.ndim == 2
if W.block_size[0] == W.shape[0] and W.block_size[1] == 1:
# Rowwise scaling becomes colwise after transpose, so this detects
# the backward pass. TODO: avoid calling matmul_lora in backward.
W = W.dequantize()
else:
W = W.contiguous()
out = torch_matmul(X, W.t(), out = out)
elif W.dtype == torch.float8_e4m3fn:
out = fp8_linear(X, W, W_quant)
else:
W = fast_dequantize(W, W_quant, use_global_buffer = True)
out = torch_matmul(X, W.t(), out = out)
if W_quant is not None:
del W
if A is not None:
# LoRA is enabled
A, B = A.t(), B.t()
XA = torch_matmul(X, A.to(dtype))
out.addmm_(XA, B.to(dtype), alpha = s)
# out += (X @ A.to(dtype)) @ (s * B.to(dtype))
return out.view(batch, seq_len, -1) if reshape else out
Reference in a new issue View git blame Copy permalink