ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/unsloth/kernels/fp8.py
pre-commit-ci[bot] 7f5b7b368c [pre-commit.ci] auto fixes from pre-commit.com hooks 
for more information, see https://pre-commit.ci
2026-03-17 08:24:19 +00:00

800 lines
30 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 os
import torch
import torch.nn as nn
import triton
import triton.language as tl
from torch.nn import functional as F
import math
from unsloth_zoo.utils import Version
from unsloth_zoo.log import logger
from unsloth_zoo.temporary_patches.common import torch_compile
torch_matmul = torch.matmul
try:
from transformers.integrations.finegrained_fp8 import FP8Linear
except:
FP8Linear = None
logger.info(
"Unsloth: FP8 models need importing FP8Linear from `transformers.integrations.finegrained_fp8` but we don't see it."
)
try:
from transformers.integrations.fbgemm_fp8 import FbgemmFp8Linear
except:
FbgemmFp8Linear = None
logger.info(
"Unsloth: FP8 models need importing FbgemmFP8Linear from `transformers.integrations.fbgemm_fp8` but we don't see it."
)
try:
from fbgemm_gpu.experimental.gemm.triton_gemm.fp8_gemm import (
triton_quantize_fp8_block,
)
except:
triton_quantize_fp8_block = None
logger.info(
"Unsloth: Could not find fbgemm_gpu.experimental.gemm.triton_gemm.fp8_gemm.triton_quantize_fp8_block"
)
try:
from torchao.prototype.blockwise_fp8_inference.blockwise_quantization import (
blockwise_fp8_gemm as torchao_blockwise_gemm,
)
except:
torchao_blockwise_gemm = None
logger.info(
"Unsloth: Could not find torchao.prototype.blockwise_fp8_inference.blockwise_quantization.blockwise_fp8_gemm"
)
try:
from compressed_tensors.linear.compressed_linear import CompressedLinear
from compressed_tensors.quantization.quant_args import QuantizationStrategy
from compressed_tensors.quantization.quant_config import QuantizationStatus
except:
CompressedLinear = None
QuantizationStrategy = None
QuantizationStatus = None
@triton.jit
def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
pid_m = tl.program_id(axis = 0)
pid_n = tl.program_id(axis = 1)
n = tl.cdiv(N, BLOCK_SIZE)
offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
offs = offs_m[:, None] * N + offs_n[None, :]
mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
x = tl.load(x_ptr + offs, mask = mask).to(tl.float32)
s = tl.load(s_ptr + pid_m * n + pid_n)
y = x * s
tl.store(y_ptr + offs, y, mask = mask)
def weight_dequant_block(
x: torch.Tensor, s: torch.Tensor, block_size: int = 128, dtype = torch.bfloat16
) -> torch.Tensor:
if not x.is_contiguous():
x = x.contiguous()
if not s.is_contiguous():
s = s.contiguous()
assert x.dim() == 2 and s.dim() == 2
M, N = x.size()
y = torch.empty_like(x, dtype = dtype)
grid = lambda meta: (
triton.cdiv(M, meta["BLOCK_SIZE"]),
triton.cdiv(N, meta["BLOCK_SIZE"]),
)
weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size)
return y
def weight_dequant(x: torch.Tensor, s: torch.Tensor, dtype = torch.bfloat16):
# Per-tensor scale: single value for entire weight matrix
if s.numel() == 1:
return x.to(dtype) * s.view(1, 1).to(dtype)
# Row quantized weight: scale shape is (m, 1) or (n, 1)
elif s.ndim == 2 and s.shape[1] == 1:
if x.shape[0] == s.shape[0]:
y = x.to(dtype) * s.to(dtype)
elif x.shape[1] == s.shape[0]:
# sometimes, this is called with the transpose of the weight. Adjust for that.
y = x.t().to(dtype) * s.to(dtype)
y = y.t()
else:
raise ValueError(f"Incompatible shapes {x.shape = }, {s.shape = }")
return y
# Block quantized weight: scale shape is (ceil(m/block_m), ceil(n/block_n))
else:
return weight_dequant_block(x, s, dtype = dtype)
# Copied from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
@triton.jit
def act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
pid = tl.program_id(axis = 0)
offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
x = tl.load(x_ptr + offs).to(tl.float32)
s = tl.max(tl.abs(x)) / 448.0
# For a row of all zeros, lets return zeros as is
# for LoRA, there are cases where dY has 0 in it and we should not let it be NaN
# this is a deviation from the original implementation.
s = 1.0 if s == 0 else s
y = x / s
y = y.to(y_ptr.dtype.element_ty)
tl.store(y_ptr + offs, y)
tl.store(s_ptr + pid, s)
def act_quant(
x: torch.Tensor, block_size: int = 128
) -> tuple[torch.Tensor, torch.Tensor]:
if not x.is_contiguous():
x = x.contiguous()
assert x.shape[-1] % block_size == 0
y = torch.empty_like(x, dtype = torch.float8_e4m3fn)
s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype = torch.float32)
def grid(meta):
return (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)
act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size)
return y, s
# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py
@triton.jit
def _w8a8_block_fp8_matmul(
# Pointers to inputs and output
A,
B,
C,
As,
Bs,
# Shape for matmul
M,
N,
K,
# Block size for block-wise quantization
group_n,
group_k,
# Stride for inputs and output
stride_am,
stride_ak,
stride_bk,
stride_bn,
stride_cm,
stride_cn,
stride_As_m,
stride_As_k,
stride_Bs_k,
stride_Bs_n,
# Meta-parameters
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
):
"""Triton-accelerated function used to perform linear operations (dot
product) on input tensors `A` and `B` with block-wise quantization, and
store the result in output tensor `C`.
"""
pid = tl.program_id(axis = 0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + (pid % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
As_ptrs = As + offs_am * stride_As_m
offs_bsn = offs_bn // group_n
Bs_ptrs = Bs + offs_bsn * stride_Bs_n
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype = tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = tl.load(a_ptrs, mask = offs_k[None, :] < K - k * BLOCK_SIZE_K, other = 0.0)
b = tl.load(b_ptrs, mask = offs_k[:, None] < K - k * BLOCK_SIZE_K, other = 0.0)
k_start = k * BLOCK_SIZE_K
offs_ks = k_start // group_k
a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
if C.dtype.element_ty == tl.bfloat16:
c = accumulator.to(tl.bfloat16)
elif C.dtype.element_ty == tl.float16:
c = accumulator.to(tl.float16)
else:
c = accumulator.to(tl.float32)
offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
tl.store(c_ptrs, c, mask = c_mask)
def w8a8_block_fp8_matmul_triton(
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
block_size: list[int],
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""Block-wise FP8 matmul."""
if block_size is None:
block_n, block_k = 128, 128
else:
assert len(block_size) == 2
block_n, block_k = block_size[0], block_size[1]
N, K = B.shape
assert A.shape[-1] == B.shape[-1]
assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
assert triton.cdiv(N, block_n) == Bs.shape[0]
assert triton.cdiv(K, block_k) == Bs.shape[1]
M = A.numel() // A.shape[-1]
C_shape = A.shape[:-1] + (N,)
C = A.new_empty(C_shape, dtype = output_dtype)
BLOCK_SIZE_M = 128
if M < BLOCK_SIZE_M:
BLOCK_SIZE_M = max(triton.next_power_of_2(M), 16)
BLOCK_SIZE_K, BLOCK_SIZE_N = block_k, block_n
def grid(META):
return (
triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
)
_w8a8_block_fp8_matmul[grid](
A,
B,
C,
As,
Bs,
M,
N,
K,
block_n,
block_k,
A.stride(-2),
A.stride(-1),
B.stride(1),
B.stride(0),
C.stride(-2),
C.stride(-1),
As.stride(-2),
As.stride(-1),
Bs.stride(1),
Bs.stride(0),
BLOCK_SIZE_M = BLOCK_SIZE_M,
BLOCK_SIZE_N = BLOCK_SIZE_N,
BLOCK_SIZE_K = BLOCK_SIZE_K,
GROUP_SIZE_M = 8,
)
return C
def torchao_block_matmul(
act_q: torch.Tensor,
weight_q: torch.Tensor,
act_scale: torch.Tensor,
weight_scale: torch.Tensor,
block_size: tuple[int, int],
output_dtype: torch.dtype = torch.bfloat16,
):
out = torchao_blockwise_gemm(
act_q.contiguous(),
act_scale.contiguous(),
weight_q.contiguous(),
weight_scale.contiguous(),
block_size = block_size[1],
)
return out.to(output_dtype)
# Note that older versions of fbgemm (<=1.3.0) cause numerical imprecisions resulting in NaNs especially when X has high values in it.
# So our preference order is fbgemm (>=1.4.0) > torchao > triton. All of these have similar outputs/losses. Never use fbgemm (<=1.3.0) for block quantized FP8 matmul.
# This torchao FP8 matmul seems to be ~3x faster than the w8a8_block_fp8_matmul_triton. Though torchao is 15-30% slower than fbgemm implementation (on H100 GPUs).
fp8_block_matmul = (
torchao_block_matmul
if torchao_blockwise_gemm is not None
else w8a8_block_fp8_matmul_triton
)
class FP8BlockQuantLinear(torch.autograd.Function):
@staticmethod
def forward(ctx, X, weight, weight_scale):
m, n = weight.shape
# Save original scale for backward (before any transformation)
original_weight_scale = weight_scale
# Handle per-tensor quantization: expand scalar to block scale shape
if weight_scale.numel() == 1:
block_size = [128, 128]
# Expand scalar to (ceil(m/128), ceil(n/128)) - same value for all blocks
num_blocks_m = triton.cdiv(m, block_size[0])
num_blocks_n = triton.cdiv(n, block_size[1])
weight_scale = weight_scale.expand(num_blocks_m, num_blocks_n).contiguous()
else:
# Block quantization path
p, q = weight_scale.shape
block_size = getattr(weight, "block_size", None) or getattr(
weight_scale, "block_size", [128, 128]
)
assert block_size is not None, "block_size is not set"
if triton.cdiv(m, block_size[0]) != p or triton.cdiv(n, block_size[1]) != q:
if (
triton.cdiv(m, block_size[0]) == q
and triton.cdiv(n, block_size[1]) == p
):
weight_scale = weight_scale.T
original_weight_scale = weight_scale # Update for transposed case
else:
raise ValueError(
f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {block_size}"
)
if not weight.is_contiguous():
weight = weight.contiguous()
# Quantize input and run FP8 matmul
qinput, scale = act_quant(X, block_size[1])
output = fp8_block_matmul(
qinput,
weight,
scale,
weight_scale,
block_size,
output_dtype = X.dtype,
)
ctx.weight = weight
ctx.weight_scale = original_weight_scale # Save original for backward
return output.to(X.dtype)
@staticmethod
def backward(ctx, grad_output):
W_deq = weight_dequant(ctx.weight, ctx.weight_scale)
grad_X = torch_matmul(grad_output, W_deq)
del W_deq
return grad_X, None, None
@torch_compile
def fp8_torch_block_quant_forward(X, weight, weight_scale):
return FP8BlockQuantLinear.apply(X, weight, weight_scale)
class FbgemmFp8Linear_matmul(torch.autograd.Function):
@staticmethod
def forward(ctx, x, weight, weight_scale, bias = None):
if weight.shape[0] == weight_scale.shape[0] and (
weight.shape[0] % 8 == 0 and weight.shape[1] % 8 == 0
):
# Edit: The kernel seems to expect that the weight has dimensions divisible by 8. Otherwise it throws `RuntimeError: cutlass cannot implement`
# One thing we can do is to pad the weight and weight scale to multiple of 8 and perform a F8F8BF16 operation.
# I tried benchmarking that for speed but observed that dequantize+bf16 matmul is significantly faster than padding+f8f8bf16 matmul. So we'll go that route.
# So essentially, f8f8bf16_rowise only happens when shapes are proper (no transposes) and divisible by 8.
# quantize_fp8_per_row will squash the leading dimensions, so save the desired shape here
output_shape = (*x.shape[:-1], -1)
# x_quantized and x_scale are not necessarily on the same device as x, this is an issue.
# https://github.com/pytorch/FBGEMM/blob/e08af8539c391437f447173863df0f3f6f6f1855/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cu#L1237C3-L1237C45
x_quantized, x_scale = torch.ops.fbgemm.quantize_fp8_per_row(
x.view(-1, x.shape[-1]).contiguous(),
scale_ub = getattr(weight, "input_scale_ub", None),
)
# moving x_quantized, x_scale here creates glibberish output ... However, if we move the output, it works
# x_quantized, x_scale = x_quantized.to(x.device), x_scale.to(x.device)
# The computation still happens on the device where self.weight is even if x_quantized is not on the same device as self.weight
weight_scale_float32 = weight_scale.to(torch.float32)
if not weight.is_contiguous():
weight = weight.contiguous()
if not weight_scale.is_contiguous():
weight_scale = weight_scale.contiguous()
output = torch.ops.fbgemm.f8f8bf16_rowwise(
x_quantized, weight, x_scale, weight_scale_float32, use_fast_accum = True
)
output = output + bias if bias is not None else output
# Hacky for now, we have the output to the device of x
output = output.to(x.device, x.dtype)
output = output.reshape(output_shape)
del x_quantized, x_scale
elif (
weight.shape[0] != weight_scale.shape[0]
and weight.shape[1] == weight_scale.shape[0]
) or (weight.shape[0] % 8 != 0 or weight.shape[1] % 8 != 0):
# Either the weight/scale is transposed or its shape is not divisible by 8. Both cases, dequantizing is the preferred way.
# The transpose case is generally noticed in backward pass when we do dY@W instead of @W.T as we do for forward.
# The shape case, I noticed to happen in MLP of Qwen 2.5 VL 7B where the gate proj is of shape (3420, 1280) and 3420/8=427.5
W_deq = weight_dequant(weight, weight_scale).T
output = torch_matmul(x, W_deq)
del W_deq
else:
raise ValueError(
f"Shapes are incompatible {weight.shape = }, {weight_scale.shape = }, {x.shape = }"
)
ctx.weight = weight
ctx.weight_scale = weight_scale
return output
@staticmethod
def backward(ctx, grad_output):
W_deq = weight_dequant(ctx.weight, ctx.weight_scale)
grad_X = torch_matmul(grad_output, W_deq)
del W_deq
return grad_X, None, None, None, None
@torch_compile
def fbgemm_fp8_linear(X, weight, weight_scale, bias = None):
return FbgemmFp8Linear_matmul.apply(X, weight, weight_scale, bias)
class FP8_fbgemm_block_linear(torch.autograd.Function):
@staticmethod
def forward(ctx, X, weight, weight_scale, bias = None):
orig_shape = X.shape
X = X.view(-1, X.shape[-1])
bs_n, bs_k = getattr(weight, "block_size", None) or getattr(
weight_scale, "block_size", [128, 128]
)
bs_m = bs_n
m, n = weight.shape
p, q = weight_scale.shape
if triton.cdiv(m, bs_n) != p or triton.cdiv(n, bs_k) != q:
if triton.cdiv(m, bs_n) == q and triton.cdiv(n, bs_k) == p:
# weights are transposed during backward pass for training :)
# We transpose weight scale to counter that. Note that transposing weight would cause issues with matmul with input X
weight_scale = weight_scale.T
else:
raise ValueError(
f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {bs_n, bs_k}"
)
xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None)
## TODO: Investigate and resolve the high divergence of this output from baseline
# WARNING: This causes the outputs to diverge from expected when X has high values in it.
# That results in the model producing gibberish, especially on longer sequences and training loss starting at high values like 8 instead of <1 ideally
# Please refrain from using this till this issue is resolved. This exists here just for a future headstart.
output = torch.ops.fbgemm.f8f8bf16_blockwise(
xq, weight.contiguous(), xs, weight_scale.contiguous(), bs_m, bs_n, bs_k
)
output = output + bias if bias is not None else output
output = output.view(*orig_shape[:-1], -1)
del xq
del xs
ctx.weight = weight
ctx.weight_scale = weight_scale
ctx.block_size = [bs_m, bs_n, bs_k]
return output
@staticmethod
def backward(ctx, grad_output):
W_deq = weight_dequant(ctx.weight, ctx.weight_scale)
grad_X = torch_matmul(grad_output, W_deq)
del W_deq
return grad_X, None, None, None, None
@torch_compile
def fp8_fbgemm_block_linear(X, weight, weight_scale, bias = None):
return FP8_fbgemm_block_linear.apply(X, weight, weight_scale, bias)
def test_has_fbgemm():
# We must manually check if the faster FBGEMM works on the specific GPU
# For example RTX 5090 and RTX 4090 does not work
# Also SM100 (Blackwell B200/B100) GPUs fail with CUTLASS SM90 kernels
# [TODO] Investigate with TorchAO why FBGEMM fails on consumer GPUs
M, N, K = 128, 128, 128
xq = torch.ones(M, K, dtype = torch.float8_e4m3fn, device = "cuda")
wq = xq
M, K = xq.shape
N, _ = wq.shape
block_scale = torch.ones(M // 128, K // 128, dtype = torch.float32, device = "cuda")
has_fbgemm = False
try:
out = torch.ops.fbgemm.f8f8bf16_blockwise(xq, wq, block_scale, block_scale)
assert torch.unique(out).item() == 128
has_fbgemm = True
del out
except Exception as e:
error_str = str(e).lower()
# Catch any CUTLASS/CUDA errors and disable FBGEMM
# This includes MMA instruction errors, architecture mismatches, kernel launch failures, etc.
cutlass_cuda_errors = (
"cutlass",
"cuda error",
"cuda runtime error",
"no kernel image",
"arch conditional",
"mma instruction",
"compute capability",
"cute_invalid_control_path",
"tma",
)
is_cutlass_cuda_error = any(err in error_str for err in cutlass_cuda_errors)
if is_cutlass_cuda_error:
print(
"Unsloth: FBGEMM on the current GPU cannot load - will switch to Triton kernels"
)
else:
print(
f"Unsloth: FBGEMM on the current GPU cannot load with error = {e} - will switch to Triton kernels"
)
has_fbgemm = False
del block_scale, xq
torch.cuda.empty_cache()
return has_fbgemm
fp8_block_quant_linear = fp8_torch_block_quant_forward
if "UNSLOTH_HAS_FBGEMM" not in os.environ:
os.environ["UNSLOTH_HAS_FBGEMM"] = "0"
try:
import fbgemm_gpu
# Older versions cause numerical imprecisions resulting in NaNs especially when X has high values in it.
# This is both fast and accurate hence preferred.
# This makes it 15% faster than the torchao implementation.
if Version(fbgemm_gpu.__version__) >= Version("1.4.0"):
# We must manually confirm if blockwise FBGEMM works!
# This check is a must for consumer grade GPUs which fail
# Suppress CUDA device printf during probe -- on Blackwell (SM100) GPUs,
# FBGEMM's CUTLASS blockwise kernel (hardcoded SM90) fires thousands of
# "Arch conditional MMA" lines to stdout fd 1 before aborting.
from unsloth.import_fixes import suppress_cuda_printf
with suppress_cuda_printf():
_has_fbgemm = test_has_fbgemm()
if _has_fbgemm:
os.environ["UNSLOTH_HAS_FBGEMM"] = "1"
logger.info(f"Using fbgemm_gpu block quantized FP8 matmul")
fp8_block_quant_linear = fp8_fbgemm_block_linear
else:
os.environ["UNSLOTH_HAS_FBGEMM"] = "0"
except:
pass
HAS_FBGEMM_FP8_OPS = hasattr(torch.ops, "fbgemm") and hasattr(
torch.ops.fbgemm, "quantize_fp8_per_row"
)
@torch_compile
def fp8_linear(X, weight, weight_scale, bias = None):
# Per-tensor quantization: single scalar scale for entire weight
# Block quantized FP8: 2D scale tensor with multiple columns
if weight_scale.numel() == 1 or (
weight_scale.ndim == 2 and weight_scale.shape[1] > 1
):
out = fp8_block_quant_linear(X, weight, weight_scale)
if bias is not None:
out = out + bias
# Row/channel quantized FP8: 2D scale with shape (n, 1)
elif HAS_FBGEMM_FP8_OPS:
out = fbgemm_fp8_linear(X, weight, weight_scale, bias)
else:
# Fallback: dequantize FP8 weight and use standard matmul
W_deq = weight_dequant(weight, weight_scale).T
out = torch_matmul(X, W_deq)
if bias is not None:
out = out + bias
del W_deq
return out
def module_forward_patch(forward_function, scale_attr = "weight_scale"):
def patched_forward(self, X):
return forward_function(X, self.weight, getattr(self, scale_attr))
return patched_forward
def _compressed_linear_supports_unsloth_fp8(self):
if (
CompressedLinear is None
or QuantizationStrategy is None
or QuantizationStatus is None
):
return False
weight = getattr(self, "weight", None)
weight_scale = getattr(self, "weight_scale", None)
quantization_scheme = getattr(self, "quantization_scheme", None)
quantization_args = getattr(quantization_scheme, "weights", None)
if (
weight is None
or weight_scale is None
or quantization_args is None
or weight.dtype != torch.float8_e4m3fn
):
return False
if getattr(quantization_args, "type", None) != "float":
return False
if getattr(quantization_args, "num_bits", None) != 8:
return False
if getattr(quantization_args, "symmetric", None) is not True:
return False
if getattr(self, "weight_zero_point", None) is not None:
return False
strategy = getattr(quantization_args, "strategy", None)
if strategy not in (
QuantizationStrategy.TENSOR,
QuantizationStrategy.CHANNEL,
QuantizationStrategy.BLOCK,
"tensor",
"channel",
"block",
):
return False
if not torch.is_tensor(weight_scale):
return False
if weight_scale.numel() == 0:
return False
return True
def _compressed_linear_forward_fallback(self, input):
if self.quantization_status == QuantizationStatus.COMPRESSED:
weight_data = self.compressor.decompress_module(self)
param = nn.Parameter(weight_data, requires_grad = False)
from compressed_tensors.utils import register_offload_parameter
register_offload_parameter(self, "weight", param)
self.quantization_status = QuantizationStatus.FROZEN
return F.linear(input, self.weight, self.bias)
def compressed_linear_forward_patch(self, input):
if _compressed_linear_supports_unsloth_fp8(self):
return fp8_linear(input, self.weight, self.weight_scale, self.bias)
return _compressed_linear_forward_fallback(self, input)
def _fp8_moe_lora_extractor(wrapper, weight_A, weight_B, scaling, num_experts):
total_rank = weight_A.shape[0]
if num_experts == 0 or total_rank % num_experts != 0:
raise ValueError(
f"LoRA total_rank ({total_rank}) must be divisible by num_experts ({num_experts})"
)
rank_per_expert = total_rank // num_experts
dim_A = weight_A.shape[1]
dim_B = weight_B.shape[0]
hidden_dim = None
intermediate_dim = None
current = wrapper
while hasattr(current, "base_layer"):
current = current.base_layer
if hasattr(current, "hidden_dim"):
hidden_dim = current.hidden_dim
if hasattr(current, "intermediate_dim"):
intermediate_dim = current.intermediate_dim
if hasattr(current, "gate_up_proj") and hasattr(current.gate_up_proj, "shape"):
shape = current.gate_up_proj.shape
if len(shape) == 3:
hidden_dim = shape[2]
intermediate_dim = shape[1] // 2
param_name = getattr(wrapper, "parameter_name", None)
if (
param_name == "down_proj"
and intermediate_dim is not None
and hidden_dim is not None
):
first_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
first_weight = first_weight.permute(1, 0, 2).contiguous()
second_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
return first_weight, second_weight, scaling, num_experts
elif param_name == "gate_up_proj" and hidden_dim is not None:
first_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
first_weight = first_weight.permute(1, 0, 2).contiguous()
second_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
return first_weight, second_weight, scaling, num_experts
if hidden_dim is not None:
if dim_B == hidden_dim:
first_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
first_weight = first_weight.permute(1, 0, 2).contiguous()
second_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
return first_weight, second_weight, scaling, num_experts
elif dim_A == hidden_dim:
first_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
first_weight = first_weight.permute(0, 2, 1).contiguous()
second_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
second_weight = second_weight.permute(1, 2, 0).contiguous()
return first_weight, second_weight, scaling, num_experts
first_weight = weight_A.view(num_experts, rank_per_expert, dim_A)
first_weight = first_weight.permute(0, 2, 1).contiguous()
second_weight = weight_B.view(dim_B, num_experts, rank_per_expert)
second_weight = second_weight.permute(1, 2, 0).contiguous()
return first_weight, second_weight, scaling, num_experts
def _patch_fp8_moe_experts():
try:
from transformers.integrations import finegrained_fp8
from unsloth_zoo.temporary_patches.moe_utils_fp8 import (
forward_moe_backend_fp8,
)
except Exception:
return
experts_interface = getattr(finegrained_fp8, "ALL_FP8_EXPERTS_FUNCTIONS", None)
if experts_interface is not None:
experts_interface["grouped_mm"] = forward_moe_backend_fp8
experts_interface["batched_mm"] = forward_moe_backend_fp8
if hasattr(finegrained_fp8, "FP8Experts"):
finegrained_fp8.FP8Experts._unsloth_lora_extractor_fn = staticmethod(
_fp8_moe_lora_extractor
)
# Patch the forward functions of the layers (for compiled models)
if FbgemmFp8Linear is not None:
FbgemmFp8Linear.forward = module_forward_patch(fbgemm_fp8_linear, "weight_scale")
if FP8Linear is not None:
FP8Linear.forward = module_forward_patch(fp8_block_quant_linear, "weight_scale_inv")
if CompressedLinear is not None:
CompressedLinear.forward = compressed_linear_forward_patch
_patch_fp8_moe_experts()
Reference in a new issue View git blame Copy permalink