mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-20 02:49:32 +00:00
Skip the fast LoRA paths when lora_B has a bias (#12981)
* Skip the fast LoRA paths when lora_B has a bias PEFT's lora_bias=True adds lora_B's bias to the LoRA delta. The fused training kernels and fast_linear_forward only use lora_B.weight, so a trained lora_B bias was dropped in training and in decode. Fall back to PEFT for those layers, as for DoRA. * Handle LoRA bias in installed fused kernels * Trim comments in PR scope --------- Co-authored-by: Maheswar Kumar <110882203+mahiatlinux@users.noreply.github.com> Co-authored-by: mahiatlinux <mahiatlinux@users.noreply.github.com>
This commit is contained in:
Vineeth Sai Varikuntla
GitHub
parent
82bd5f31a1
commit
d8ad0b0806
4 changed files with 95 additions and 17 deletions
|
|
@ -32,13 +32,17 @@ class _Block(torch.nn.Module):
|
|||
self.act_fn = torch.nn.SiLU()
|
||||
|
||||
|
||||
def _peft_block(adapters):
|
||||
def _peft_block(adapters, lora_bias = False):
|
||||
from peft import LoraConfig, get_peft_model
|
||||
|
||||
torch.manual_seed(3407)
|
||||
targets = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
|
||||
cfg = lambda r: LoraConfig(
|
||||
r = r, lora_alpha = 2 * r, target_modules = targets, init_lora_weights = False
|
||||
r = r,
|
||||
lora_alpha = 2 * r,
|
||||
target_modules = targets,
|
||||
init_lora_weights = False,
|
||||
lora_bias = lora_bias,
|
||||
)
|
||||
model = get_peft_model(_Block(), cfg(2), adapter_name = "a")
|
||||
if adapters > 1:
|
||||
|
|
@ -47,6 +51,10 @@ def _peft_block(adapters):
|
|||
model = model.to("cuda", torch.bfloat16)
|
||||
for name, p in model.named_parameters():
|
||||
p.requires_grad_("lora_" in name)
|
||||
if lora_bias:
|
||||
with torch.no_grad():
|
||||
for proj in targets:
|
||||
model.base_model.model.get_submodule(proj).lora_B["a"].bias.normal_()
|
||||
return model, model.base_model.model
|
||||
|
||||
|
||||
|
|
@ -98,10 +106,26 @@ def test_qkv_o_mlp_match_peft(adapters):
|
|||
model,
|
||||
)
|
||||
if adapters > 1:
|
||||
# The second adapter must be in the graph, not just the first.
|
||||
assert any(".b." in n for n in qkv) and any(".b." in n for n in mlp)
|
||||
|
||||
|
||||
def test_qkv_o_mlp_match_peft_with_lora_bias():
|
||||
from unsloth.kernels import apply_lora_mlp_swiglu, apply_lora_o, apply_lora_qkv
|
||||
|
||||
model, block = _peft_block(1, lora_bias = True)
|
||||
_check(
|
||||
lambda X: apply_lora_qkv(block, X),
|
||||
lambda X: (block.q_proj(X), block.k_proj(X), block.v_proj(X)),
|
||||
model,
|
||||
)
|
||||
_check(lambda X: apply_lora_o(block, X), block.o_proj, model)
|
||||
_check(
|
||||
lambda X: apply_lora_mlp_swiglu(block, X),
|
||||
lambda X: block.down_proj(block.act_fn(block.gate_proj(X)) * block.up_proj(X)),
|
||||
model,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("adapters", [1, 2])
|
||||
def test_fast_linear_forward_decode_matches_peft(adapters):
|
||||
from unsloth.kernels import fast_linear_forward
|
||||
|
|
@ -127,8 +151,27 @@ def test_fast_linear_forward_applies_dora_magnitude(q_len):
|
|||
block = get_peft_model(_Block(), cfg).to("cuda", torch.bfloat16).base_model.model
|
||||
X = torch.randn(1, q_len, H, device = "cuda", dtype = torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
# A trained magnitude no longer equals the row norms DoRA starts from.
|
||||
# a trained magnitude differs from the initial DoRA row norms
|
||||
block.q_proj.lora_magnitude_vector["default"].weight.mul_(1.5)
|
||||
torch.testing.assert_close(
|
||||
fast_linear_forward(block.q_proj, X), block.q_proj(X), rtol = 2e-2, atol = 2e-2
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("q_len", [1, 5])
|
||||
def test_fast_linear_forward_adds_lora_bias(q_len):
|
||||
from peft import LoraConfig, get_peft_model
|
||||
from unsloth.kernels import fast_linear_forward
|
||||
|
||||
torch.manual_seed(3407)
|
||||
cfg = LoraConfig(
|
||||
r = 4, lora_alpha = 8, target_modules = ["q_proj"], init_lora_weights = False, lora_bias = True
|
||||
)
|
||||
block = get_peft_model(_Block(), cfg).to("cuda", torch.bfloat16).base_model.model
|
||||
X = torch.randn(1, q_len, H, device = "cuda", dtype = torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
# a trained lora_B bias differs from its zero initialization
|
||||
block.q_proj.lora_B["default"].bias.normal_()
|
||||
torch.testing.assert_close(
|
||||
fast_linear_forward(block.q_proj, X), block.q_proj(X), rtol = 2e-2, atol = 2e-2
|
||||
)
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@
|
|||
import torch
|
||||
from unsloth_zoo.utils import Version
|
||||
from .utils import (
|
||||
_has_active_lora_bias,
|
||||
_has_multiple_active_adapters,
|
||||
addmm_,
|
||||
fast_dequantize,
|
||||
|
|
@ -25,7 +26,7 @@ from .utils import (
|
|||
|
||||
_is_compiling = torch.compiler.is_compiling
|
||||
|
||||
# Inductor before torch 2.11 miscompiles these Functions' traced backward (wrong LoRA gradients).
|
||||
# Inductor before torch 2.11 miscompiles traced backward passes and produces wrong LoRA gradients.
|
||||
TRACE_LORA_FUNCTIONS = Version(torch.__version__) >= Version("2.11.0")
|
||||
|
||||
|
||||
|
|
@ -231,7 +232,7 @@ class LoRA_MLP(torch.autograd.Function):
|
|||
None,
|
||||
None,
|
||||
None,
|
||||
) # _backward and _forward and inplace
|
||||
)
|
||||
|
||||
|
||||
from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
|
||||
|
|
@ -243,7 +244,9 @@ def apply_lora_mlp_swiglu(
|
|||
inplace = True,
|
||||
):
|
||||
if any(
|
||||
_has_multiple_active_adapters(proj) or _has_activation_fake_quantizer(proj)
|
||||
_has_multiple_active_adapters(proj)
|
||||
or _has_active_lora_bias(proj)
|
||||
or _has_activation_fake_quantizer(proj)
|
||||
for proj in (self.gate_proj, self.up_proj, self.down_proj)
|
||||
):
|
||||
return self.down_proj(self.act_fn(self.gate_proj(X)) * self.up_proj(X))
|
||||
|
|
@ -284,7 +287,9 @@ def apply_lora_mlp_geglu_exact(
|
|||
inplace = True,
|
||||
):
|
||||
if any(
|
||||
_has_multiple_active_adapters(proj) or _has_activation_fake_quantizer(proj)
|
||||
_has_multiple_active_adapters(proj)
|
||||
or _has_active_lora_bias(proj)
|
||||
or _has_activation_fake_quantizer(proj)
|
||||
for proj in (self.gate_proj, self.up_proj, self.down_proj)
|
||||
):
|
||||
return self.down_proj(self.act_fn(self.gate_proj(X)) * self.up_proj(X))
|
||||
|
|
@ -321,7 +326,9 @@ from .geglu import geglu_approx_forward_kernel, geglu_approx_backward_kernel
|
|||
|
||||
def apply_lora_mlp_geglu_approx(self, X):
|
||||
if any(
|
||||
_has_multiple_active_adapters(proj) or _has_activation_fake_quantizer(proj)
|
||||
_has_multiple_active_adapters(proj)
|
||||
or _has_active_lora_bias(proj)
|
||||
or _has_activation_fake_quantizer(proj)
|
||||
for proj in (self.gate_proj, self.up_proj, self.down_proj)
|
||||
):
|
||||
return self.down_proj(self.act_fn(self.gate_proj(X)) * self.up_proj(X))
|
||||
|
|
@ -544,7 +551,9 @@ def apply_lora_qkv(
|
|||
inplace = True,
|
||||
):
|
||||
if any(
|
||||
_has_multiple_active_adapters(proj) or _has_activation_fake_quantizer(proj)
|
||||
_has_multiple_active_adapters(proj)
|
||||
or _has_active_lora_bias(proj)
|
||||
or _has_activation_fake_quantizer(proj)
|
||||
for proj in (self.q_proj, self.k_proj, self.v_proj)
|
||||
):
|
||||
return self.q_proj(X), self.k_proj(X), self.v_proj(X)
|
||||
|
|
@ -639,7 +648,6 @@ class LoRA_W(torch.autograd.Function):
|
|||
d_A.addmm_(X.t(), y_dB, alpha = S, beta = 0)
|
||||
d_B.addmm_(A.t() @ X.t(), dY, alpha = S, beta = 0)
|
||||
|
||||
# Get derivative for dX
|
||||
W = fast_dequantize(W.t(), W_quant)
|
||||
dX = dY @ W.t()
|
||||
del W
|
||||
|
|
@ -649,7 +657,11 @@ class LoRA_W(torch.autograd.Function):
|
|||
|
||||
|
||||
def apply_lora_o(self, X):
|
||||
if _has_multiple_active_adapters(self.o_proj) or _has_activation_fake_quantizer(self.o_proj):
|
||||
if (
|
||||
_has_multiple_active_adapters(self.o_proj)
|
||||
or _has_active_lora_bias(self.o_proj)
|
||||
or _has_activation_fake_quantizer(self.o_proj)
|
||||
):
|
||||
return self.o_proj(X)
|
||||
OW, OW_quant, OA, OB, OS = get_lora_parameters(self.o_proj)
|
||||
O = _apply(LoRA_W, X, OW, OW_quant, OA, OB, OS)
|
||||
|
|
|
|||
|
|
@ -361,6 +361,19 @@ def _has_active_dora_adapter(proj):
|
|||
return any(adapter in magnitude for adapter in adapters)
|
||||
|
||||
|
||||
def _has_active_lora_bias(proj):
|
||||
# PEFT lora_bias=True gives lora_B a bias, which the fast LoRA paths never add.
|
||||
lora_B = getattr(proj, "lora_B", None)
|
||||
if not lora_B or getattr(proj, "disable_adapters", True) or getattr(proj, "merged", False):
|
||||
return False
|
||||
adapters = getattr(proj, "active_adapters", None)
|
||||
if adapters is None:
|
||||
adapters = getattr(proj, "active_adapter", ())
|
||||
if isinstance(adapters, str):
|
||||
adapters = (adapters,)
|
||||
return any(adapter in lora_B and lora_B[adapter].bias is not None for adapter in adapters)
|
||||
|
||||
|
||||
def has_mxfp4_base(*projs):
|
||||
"""A packed MXFP4 base that cannot hand its bytes to the fused LoRA kernels (they would then hold a 16-bit
|
||||
weight until backward): use PEFT. Bases with ``mxfp4_quant_state`` pass ``weight_packed`` + that state instead."""
|
||||
|
|
@ -1389,7 +1402,11 @@ def fast_linear_forward(
|
|||
temp_lora = None,
|
||||
out = None,
|
||||
):
|
||||
if _has_multiple_active_adapters(proj) or _has_active_dora_adapter(proj):
|
||||
if (
|
||||
_has_multiple_active_adapters(proj)
|
||||
or _has_active_dora_adapter(proj)
|
||||
or _has_active_lora_bias(proj)
|
||||
):
|
||||
result = proj(X)
|
||||
return result if out is None else out.copy_(result)
|
||||
W, W_quant, lora_A, lora_B, lora_S, bias = get_lora_parameters_bias(proj)
|
||||
|
|
|
|||
|
|
@ -116,7 +116,7 @@ from unsloth.models._attn_mask_compat import (
|
|||
)
|
||||
from transformers.cache_utils import DynamicCache, Cache
|
||||
from ..kernels import *
|
||||
from ..kernels.utils import has_mxfp4_base
|
||||
from ..kernels.utils import _has_active_lora_bias, has_mxfp4_base
|
||||
from ..kernels.bnb_override import install_bnb_nf4_override as _install_bnb_nf4_override
|
||||
from ..tokenizer_utils import *
|
||||
from .vision import FastBaseModel, _is_text_seq2seq_config
|
||||
|
|
@ -4355,10 +4355,13 @@ class FastLlamaModel:
|
|||
and (len(getattr(up_proj, "lora_magnitude_vector", []) or []) == 0)
|
||||
and (len(getattr(down_proj, "lora_magnitude_vector", []) or []) == 0)
|
||||
and not has_mxfp4_base(gate_proj, up_proj, down_proj)
|
||||
and not _has_active_lora_bias(gate_proj)
|
||||
and not _has_active_lora_bias(up_proj)
|
||||
and not _has_active_lora_bias(down_proj)
|
||||
):
|
||||
# See stackoverflow.com/questions/50599045 on replacing a function within a class of a module.
|
||||
# MethodType binds the replacement; see stackoverflow.com/questions/50599045.
|
||||
if hasattr(mlp_module, "_unsloth_forward"):
|
||||
# Then the mlp has been patched to use TiledMLP.
|
||||
# _unsloth_forward identifies an existing TiledMLP patch.
|
||||
mlp_module._unsloth_forward = types.MethodType(
|
||||
_apply_lora_mlp, mlp_module
|
||||
)
|
||||
|
|
@ -4385,6 +4388,9 @@ class FastLlamaModel:
|
|||
and (len(getattr(k_proj, "lora_magnitude_vector", []) or []) == 0)
|
||||
and (len(getattr(v_proj, "lora_magnitude_vector", []) or []) == 0)
|
||||
and not has_mxfp4_base(q_proj, k_proj, v_proj)
|
||||
and not _has_active_lora_bias(q_proj)
|
||||
and not _has_active_lora_bias(k_proj)
|
||||
and not _has_active_lora_bias(v_proj)
|
||||
):
|
||||
layer.self_attn.apply_qkv = apply_lora_qkv
|
||||
n_qkv += 1
|
||||
|
|
@ -4403,6 +4409,7 @@ class FastLlamaModel:
|
|||
and (getattr(o_proj, "base_layer", o_proj).bias is None)
|
||||
and (len(getattr(o_proj, "lora_magnitude_vector", []) or []) == 0)
|
||||
and not has_mxfp4_base(o_proj)
|
||||
and not _has_active_lora_bias(o_proj)
|
||||
):
|
||||
layer.self_attn.apply_o = apply_lora_o
|
||||
n_o += 1
|
||||
|
|
@ -4412,7 +4419,6 @@ class FastLlamaModel:
|
|||
"are not enabled or a bias term (like in Qwen) is used."
|
||||
)
|
||||
|
||||
# A zero count reads as a failure, so say why the fused kernels were skipped.
|
||||
unfused_reason = _fused_lora_skip_reason(
|
||||
lora_dropout, bias, float32_base, fsdp = fused_lora_declined_for_fsdp
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue