Merge branch 'main' into nightly
This commit is contained in:
Daniel Han
commit
f9ea2cb9f2
7 changed files with 271 additions and 54 deletions
|
|
@ -20,7 +20,7 @@ from .utils import (
|
|||
MAX_FUSED_SIZE,
|
||||
triton_tanh,
|
||||
triton_cast,
|
||||
torch_cuda_device,
|
||||
torch_gpu_device,
|
||||
)
|
||||
from transformers.models.llama.modeling_llama import logger
|
||||
from packaging.version import Version
|
||||
|
|
@ -301,7 +301,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
|
|||
BLOCK_SIZE, num_warps = calculate_settings(vocab_size)
|
||||
logsumexp = torch.empty(n_rows, dtype = torch.float32, device = device)
|
||||
|
||||
with torch_cuda_device(device):
|
||||
with torch_gpu_device(device):
|
||||
_cross_entropy_forward[(n_rows,)](
|
||||
logits, logits.stride(0),
|
||||
losses,
|
||||
|
|
@ -319,7 +319,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
|
|||
# For large vocabs > 65336 like Gemma 256K
|
||||
logsumexp = torch.empty((n_rows, n_chunks,), dtype = torch.float32, device = device)
|
||||
|
||||
with torch_cuda_device(device):
|
||||
with torch_gpu_device(device):
|
||||
_chunked_cross_entropy_forward[(n_rows, n_chunks,)](
|
||||
logits, logits.stride(0),
|
||||
losses,
|
||||
|
|
@ -363,7 +363,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
|
|||
div, mod = divmod(vocab_size, BLOCK_SIZE)
|
||||
n_blocks : int = div + (mod != 0)
|
||||
|
||||
with torch_cuda_device(dlosses.device):
|
||||
with torch_gpu_device(dlosses.device):
|
||||
_cross_entropy_backward[(n_rows, n_blocks,)](
|
||||
logits, logits.stride(0),
|
||||
dlosses, dlosses.stride(0),
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import torch
|
|||
from .utils import (
|
||||
calculate_settings,
|
||||
triton_tanh,
|
||||
torch_cuda_device,
|
||||
torch_gpu_device,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -48,7 +48,7 @@ def geglu_exact_forward_kernel(gate, up):
|
|||
device = gate.device
|
||||
out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = device)
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
with torch_cuda_device(device):
|
||||
with torch_gpu_device(device):
|
||||
_exact_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,)
|
||||
return out
|
||||
pass
|
||||
|
|
@ -105,7 +105,7 @@ def geglu_exact_backward_kernel(DW, e, g):
|
|||
batch_seq_len, hd = e.shape
|
||||
n_elements = e.numel()
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
with torch_cuda_device(e.device):
|
||||
with torch_gpu_device(e.device):
|
||||
_exact_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,)
|
||||
return DW, e, g
|
||||
pass
|
||||
|
|
@ -143,7 +143,7 @@ def geglu_approx_forward_kernel(gate, up):
|
|||
device = gate.device
|
||||
out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = device)
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
with torch_cuda_device(device):
|
||||
with torch_gpu_device(device):
|
||||
_approx_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,)
|
||||
return out
|
||||
pass
|
||||
|
|
@ -207,7 +207,7 @@ def geglu_approx_backward_kernel(DW, e, g):
|
|||
batch_seq_len, hd = e.shape
|
||||
n_elements = e.numel()
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
with torch_cuda_device(e.device):
|
||||
with torch_gpu_device(e.device):
|
||||
_approx_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,)
|
||||
return DW, e, g
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@
|
|||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
from .utils import calculate_settings, torch_cuda_device
|
||||
from .utils import calculate_settings, torch_gpu_device
|
||||
from unsloth_zoo.patching_utils import (
|
||||
patch_layernorm,
|
||||
)
|
||||
|
|
@ -113,7 +113,7 @@ class Fast_Layernorm(torch.autograd.Function):
|
|||
r = torch.empty(n_rows, dtype = torch.float32, device = device)
|
||||
mu = torch.empty(n_rows, dtype = torch.float32, device = device)
|
||||
|
||||
with torch_cuda_device(device):
|
||||
with torch_gpu_device(device):
|
||||
layernorm_forward[(n_rows,)](
|
||||
Y, Y.stride(0),
|
||||
X, X.stride(0),
|
||||
|
|
@ -140,7 +140,7 @@ class Fast_Layernorm(torch.autograd.Function):
|
|||
X, W, b, r, mu = ctx.saved_tensors
|
||||
n_rows, n_cols = dY.shape
|
||||
|
||||
with torch_cuda_device(dY.device):
|
||||
with torch_gpu_device(dY.device):
|
||||
layernorm_backward[(n_rows,)](
|
||||
dY, dY.stride(0),
|
||||
X, X .stride(0),
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@
|
|||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
from .utils import calculate_settings, torch_cuda_device
|
||||
from .utils import calculate_settings, torch_gpu_device
|
||||
|
||||
@triton.jit
|
||||
def _rms_layernorm_forward(
|
||||
|
|
@ -156,7 +156,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
|
|||
r = torch.empty(n_rows, dtype = torch.float32, device = device)
|
||||
|
||||
fx = _gemma_rms_layernorm_forward if gemma else _rms_layernorm_forward
|
||||
with torch_cuda_device(device):
|
||||
with torch_gpu_device(device):
|
||||
fx[(n_rows,)](
|
||||
Y, Y.stride(0),
|
||||
X, X.stride(0),
|
||||
|
|
@ -186,7 +186,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
|
|||
# dW = X
|
||||
dX = torch.empty_like(dY) if ctx.GEMMA else dY
|
||||
|
||||
with torch_cuda_device(dY.device):
|
||||
with torch_gpu_device(dY.device):
|
||||
_rms_layernorm_backward[(n_rows,)](
|
||||
dY, dY.stride(0),
|
||||
dX, dX.stride(0),
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@
|
|||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
from .utils import calculate_settings, torch_cuda_device
|
||||
from .utils import calculate_settings, torch_gpu_device
|
||||
ROPE_GROUP_SIZE : int = 4
|
||||
|
||||
def _rope_embedding(
|
||||
|
|
@ -100,7 +100,7 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
div, mod = divmod(n_heads, ROPE_GROUP_SIZE)
|
||||
n_groups : int = div + (mod != 0)
|
||||
|
||||
with torch_cuda_device(Q.device):
|
||||
with torch_gpu_device(Q.device):
|
||||
_rope_embedding[(n_rows, n_groups, )](
|
||||
Q, Q.stride(0),
|
||||
cos, cos.stride(0),
|
||||
|
|
@ -135,7 +135,7 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
cos = ctx.cos
|
||||
sin = ctx.sin
|
||||
|
||||
with torch_cuda_device(dY.device):
|
||||
with torch_gpu_device(dY.device):
|
||||
_rope_embedding[(n_rows, ctx.n_groups, )](
|
||||
dY, dY .stride(0),
|
||||
cos, cos.stride(0),
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@
|
|||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
from .utils import calculate_settings, torch_cuda_device
|
||||
from .utils import calculate_settings, torch_gpu_device
|
||||
|
||||
|
||||
@triton.jit
|
||||
|
|
@ -43,7 +43,7 @@ def swiglu_fg_kernel(e, g):
|
|||
n_elements = e.numel()
|
||||
h = torch.empty((batch, seq_len, hd), dtype = e.dtype, device = e.device)
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
with torch_cuda_device(e.device):
|
||||
with torch_gpu_device(e.device):
|
||||
_fg_kernel[grid](e, g, h, n_elements, BLOCK_SIZE = 1024,)
|
||||
return h
|
||||
pass
|
||||
|
|
@ -95,7 +95,7 @@ def swiglu_DWf_DW_dfg_kernel(DW, e, g):
|
|||
batch_seq_len, hd = e.shape
|
||||
n_elements = e.numel()
|
||||
grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
|
||||
with torch_cuda_device(e.device):
|
||||
with torch_gpu_device(e.device):
|
||||
_DWf_DW_dfg_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,)
|
||||
return DW, e, g
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -13,14 +13,21 @@
|
|||
# limitations under the License.
|
||||
|
||||
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 unsloth import DEVICE_TYPE
|
||||
|
||||
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
|
||||
import torch
|
||||
torch_Tensor = torch.Tensor
|
||||
from packaging.version 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
|
||||
|
|
@ -29,14 +36,21 @@ else:
|
|||
torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda")
|
||||
pass
|
||||
|
||||
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
|
||||
from packaging.version import Version
|
||||
import triton
|
||||
import triton.language as tl
|
||||
if Version(triton.__version__) >= Version("3.0.0"):
|
||||
from triton.language.extra import libdevice
|
||||
triton_tanh = libdevice.tanh
|
||||
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
|
||||
|
|
@ -60,50 +74,104 @@ def calculate_settings(n : int) -> (int, int,):
|
|||
return BLOCK_SIZE, num_warps
|
||||
pass
|
||||
|
||||
HAS_CUDA_STREAM = False
|
||||
# INTEL GPU specific logic
|
||||
if DEVICE_TYPE == "xpu":
|
||||
# TODO: Changed here after adding XPU BNB support
|
||||
HAS_XPU_STREAM = False
|
||||
def get_ptr(x: Optional[torch.Tensor]):
|
||||
raise RuntimeError("XPU BNB support is not implemented yet. This function should not be called.")
|
||||
else:
|
||||
# NVIDIA-GPU logic here as default
|
||||
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
|
||||
|
||||
import bitsandbytes as bnb
|
||||
import ctypes
|
||||
|
||||
# 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 torch.cuda.device_count() > 1:
|
||||
torch_cuda_device = torch.cuda.device
|
||||
if DEVICE_TYPE == "cuda" and torch.cuda.device_count() > 1:
|
||||
torch_gpu_device = torch.cuda.device
|
||||
elif DEVICE_TYPE == "xpu" and torch.xpu.device_count() > 1:
|
||||
torch_gpu_device = torch.xpu.device
|
||||
else:
|
||||
from contextlib import nullcontext
|
||||
def torch_cuda_device(device): return nullcontext()
|
||||
pass
|
||||
_cuda_getCurrentRawStream = torch._C._cuda_getCurrentRawStream
|
||||
def torch_gpu_device(device): return nullcontext()
|
||||
pass
|
||||
|
||||
# INTEL GPU Specific Logic
|
||||
if DEVICE_TYPE == "xpu":
|
||||
_gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream
|
||||
# NVIDIA GPU Default Logic
|
||||
else:
|
||||
_gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream
|
||||
|
||||
c_void_p = ctypes.c_void_p
|
||||
def _get_tensor_stream(tensor: torch_Tensor) -> c_void_p:
|
||||
return c_void_p(_cuda_getCurrentRawStream(tensor.device.index))
|
||||
return c_void_p(_gpu_getCurrentRawStream(tensor.device.index))
|
||||
pass
|
||||
|
||||
|
||||
# Get array of CUDA streams and other buffers
|
||||
global CUDA_STREAMS
|
||||
global XPU_STREAMS
|
||||
global WEIGHT_BUFFERS
|
||||
global ABSMAX_BUFFERS
|
||||
|
||||
_CUDA_STREAMS = {
|
||||
(index := torch.cuda.device(i).idx) : ctypes.c_void_p(torch._C._cuda_getCurrentRawStream(index))
|
||||
for i in range(torch.cuda.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
|
||||
# INTEL GPU Specific Logic
|
||||
if DEVICE_TYPE == "xpu":
|
||||
_XPU_STREAMS = {
|
||||
(index := torch.xpu.device(i).idx) : ctypes.c_void_p(torch._C._xpu_getCurrentRawStream(index))
|
||||
for i in range(torch.xpu.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:
|
||||
# NVIDIA GPU Default Logic
|
||||
_CUDA_STREAMS = {
|
||||
(index := torch.cuda.device(i).idx) : ctypes.c_void_p(torch._C._cuda_getCurrentRawStream(index))
|
||||
for i in range(torch.cuda.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
|
||||
|
||||
|
||||
# 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
|
||||
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
|
||||
# INTEL GPU Specific Logic
|
||||
if DEVICE_TYPE == "xpu":
|
||||
# TODO: After adding XPU BNB support, this function should be implemented
|
||||
def cdequantize_blockwise_fp32(*args, **kwargs):
|
||||
raise RuntimeError("XPU BNB support is not implemented yet. cdequantize_blockwise_fp32 should not be called now.")
|
||||
|
||||
def cdequantize_blockwise_fp16_nf4(*args, **kwargs):
|
||||
raise RuntimeError("XPU BNB support is not implemented yet. cdequantize_blockwise_fp16_nf4 should not be called now.")
|
||||
|
||||
def cdequantize_blockwise_bf16_nf4(*args, **kwargs):
|
||||
raise RuntimeError("XPU BNB support is not implemented yet. cdequantize_blockwise_bf16_nf4 should not be called now.")
|
||||
|
||||
def cgemm_4bit_inference_naive_fp16(*args, **kwargs):
|
||||
raise RuntimeError("XPU BNB support is not implemented yet. cgemm_4bit_inference_naive_fp16 should not be called now.")
|
||||
|
||||
def cgemm_4bit_inference_naive_bf16(*args, **kwargs):
|
||||
raise RuntimeError("XPU BNB support is not implemented yet. cgemm_4bit_inference_naive_bf16 should not be called now.")
|
||||
else:
|
||||
# NVIDIA GPU Default Logic
|
||||
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
|
||||
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_mm = torch.mm
|
||||
torch_mv = torch.mv
|
||||
torch_matmul = torch.matmul
|
||||
|
|
@ -160,7 +228,84 @@ def get_lora_parameters_bias(proj):
|
|||
)
|
||||
pass
|
||||
|
||||
if HAS_CUDA_STREAM:
|
||||
# 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 quant_state is None: return W
|
||||
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 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:
|
||||
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), 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,)
|
||||
pass
|
||||
# Careful returning transposed data
|
||||
is_transposed = (True if W.shape[0] == 1 else False)
|
||||
return out.t() if is_transposed else out
|
||||
pass
|
||||
# NVIDIA GPU Default Logic
|
||||
elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM:
|
||||
@torch.inference_mode
|
||||
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
|
||||
if quant_state is None: return W
|
||||
|
|
@ -218,7 +363,7 @@ if HAS_CUDA_STREAM:
|
|||
|
||||
# NF4 dequantization of statistics
|
||||
ptr_out_absmax = get_ptr(out_absmax)
|
||||
with torch_cuda_device(device):
|
||||
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
|
||||
|
|
@ -289,7 +434,79 @@ else:
|
|||
pass
|
||||
|
||||
|
||||
if HAS_CUDA_STREAM:
|
||||
# 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
|
||||
pass
|
||||
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
|
||||
|
||||
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,)
|
||||
pass
|
||||
|
||||
return out
|
||||
pass
|
||||
elif DEVICE_TYPE == "cuda" 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
|
||||
|
|
@ -342,7 +559,7 @@ if HAS_CUDA_STREAM:
|
|||
ldc = ctypes_c_int32(ldc)
|
||||
|
||||
df = torch_empty(absmax.shape, dtype = torch.float32, device = device)
|
||||
with torch_cuda_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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue