mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-21 05:11:50 +00:00
5a60350e05
* Fix fused LoRA dtype mismatch on fp32 activations into fp16/bf16 weights
The fused LoRA ops are decorated with @custom_fwd, which disables
autocast. When an fp32 activation (for example the fp32 output of
fast_rms_layernorm on an fp32 hidden/residual stream) reaches a fused
LoRA path whose base weight is fp16/bf16, nothing downcasts the
activation, so torch.matmul raises:
expected mat1 and mat2 to have the same dtype: float != c10::Half
This is not torch-version specific; torch.matmul never auto-promotes a
mixed fp16/fp32 matmul.
Reconcile the dtypes at the matmul_lora choke point: after dequantizing
W (both the plain and the fast_dequantize branches, not the fp8 branch),
downcast X to the base weight compute dtype the way autocast would for a
plain Linear, and keep the LoRA A/B cast dtype consistent with it.
Keep the backward pass dtype-consistent: LoRA_MLP, LoRA_QKV and LoRA_W
now record the original input dtype in ctx.input_dtype, pin the saved
activation to the compute dtype, cast the incoming gradients to that
dtype, and cast the returned dX back to the original input dtype to honor
the autograd grad-dtype contract. Every added branch is a no-op when the
dtypes already match, so normal fp16/bf16/fp32 training is unchanged.
* Tighten fused-LoRA dtype-fix comments
* Avoid recomputing LoRA B-projected gradient in fused-LoRA backward
Each fused-LoRA backward computed dY @ B.t() (the LoRA-B projected gradient,
shape (batch*seq, r)) twice per projection: once for the A weight gradient and
once for the dX accumulation. Hoist it into a single temporary and reuse it in
both, removing 6 small matmuls per layer (3 in LoRA_QKV, 2 in LoRA_MLP, 1 in
LoRA_W). Results are bitwise identical; measured about 2 to 3 percent faster on
the LoRA forward plus backward for qkv and mlp in both 16-bit and 4-bit.
* Follow the produced output dtype in matmul_lora, not the base weight
Pinning the LoRA compute dtype to W.dtype regressed a case that works on
main: under autocast the base matmul produces `out` in the autocast dtype,
not W's, and `out.addmm_(XA, B)` for the LoRA term is in-place and so is
not autocast-eligible. An fp16 activation against a bf16 base weight under
fp16 autocast then raised
RuntimeError: self and mat2 must have the same dtype, but got Half and BFloat16
Derive the LoRA operand dtype from `out` instead, which is the dtype the
base matmul actually produced in both regimes, and skip the activation
pre-cast entirely when autocast is on, since autocast reconciles the matmul
itself and pre-casting would only round the activation through a second
dtype. The autocast query sits behind the existing dtype comparison, so the
ordinary matching-dtype path does not pay for it.
Over the activation x base weight x adapter dtype x inplace x rows x
autocast product (1350 cells), against main: 612 cells that raised now run,
306 are bit-identical, 432 still raise (all of them an autocast dtype that
disagrees with the base weight, which the trainer already rejects at setup),
0 regressions and 0 cells whose values or dtypes moved.
tests/test_fused_lora_activation_dtype.py pins both halves: the fp32
activation case fails on main, and the autocast case fails on this branch
before this commit.
Also corrects three comments that said custom_fwd disables autocast. It does
not with cast_inputs=None; the mechanism is that in-place addmm_ and
matmul(out=) are not autocast-eligible.
* Reconcile base weights in the fused LoRA backward, and dequantize NF4 in fp32
Two defects the forward-side dtype fix exposed, both reproduced through the
public loader with FastLanguageModel.from_pretrained(dtype = torch.float32).
First, the backward passes dequantize the base weight in its own dtype and
then feed it to `torch.matmul(..., out = X)` and `dX.addmm_(...)`. Neither
form is autocast-eligible, so an fp32 base weight against a bf16 compute
dtype raised
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: c10::BFloat16 != float
_dequantize_to reconciles the weight at the six dequantize sites in
LoRA_MLP, LoRA_QKV and LoRA_W. It is a no-op whenever the dtypes already
agree, which is every ordinary fp16/bf16/fp32 run.
Second, with the crash gone, fp32 + load_in_4bit trained to a nonsense loss
(15.1 against 0.75 for the same model in bf16). fast_dequantize allocates
`out` with quant_state.dtype but dispatched the NF4 kernel two ways only,
fp16 or bf16, so an fp32 quant_state wrote bf16 bits into an fp32 buffer.
Measured against bitsandbytes' own dequantize_4bit, the error was 0.127 on
weights of magnitude 0.09. bitsandbytes exports cdequantize_blockwise_fp32_nf4;
binding it and making all three dispatch sites three-way takes that error to
exactly 0, and the loss to 0.7428 against 0.7427 for the same model in fp16.
That second one is a pre-existing bug on main, not something this PR
introduced. It was unreachable only because the crash above got there first.
Over the full kernel-level product (1350 cells) against main: 1044 cells that
raised now run, 306 are bit-identical, 0 regressions, 0 cells whose values or
dtypes moved.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add the fp32 NF4 kernel to the bitsandbytes readiness probe
kernels/utils.py now binds cdequantize_blockwise_fp32_nf4 at module scope, and
tests/python/test_bitsandbytes_kernel_readiness.py asserts the probe list matches
those binds exactly. Without this the probe would also clear a wheel that cannot
serve the fp32 dequantize path.
* Query autocast with the torch device name, and fail open if it is rejected
DEVICE_TYPE is "hip" on ROCm and "mlx" on Apple Silicon, and torch's autocast
APIs accept neither:
torch.is_autocast_enabled("hip")
RuntimeError: unknown device type for autocast in get_autocast_dispatch_key_from_device_type
so the new probe turned the mismatched-dtype case, the one this branch exists
to support, into a crash on every ROCm box. DEVICE_TYPE_TORCH already maps hip
to cuda and mlx to mps.
"npu" is rejected too on a torch build without that backend registered, which
DEVICE_TYPE_TORCH does not remap, so the device name is resolved once at import
and falls back to "assume no ambient autocast" when torch will not take it. That
fails OPEN: matmul_lora reconciles the dtypes itself, which is the behaviour that
makes the mismatched case work at all, rather than erroring on the probe.
Resolving it once also keeps the query off the per-matmul path; what remains
there is a module-global lookup behind the existing dtype comparison.
The 1350-cell matrix is unchanged by this: 1044 fixed, 306 bit-identical, 0
regressions, 0 cells whose values or dtypes moved.
Reported by Codex.
* Let the autocast device-name test run on the CPU-only legs
The file was marked `gpu` module-wide, so every test in it was deselected on any
runner without an accelerator. That is the wrong way round for
test_autocast_probe_uses_the_torch_device_name: it is pure Python, and it is
exactly the check that needs to run on the non-CUDA platforms, since the defect
it pins (torch.is_autocast_enabled("hip") raising on ROCm) cannot show up on the
NVIDIA box the rest of the file needs.
Marking the three matmul tests individually instead: 7 passed / 1 deselected
under -m gpu, 1 passed / 7 deselected without it, where before it was 0 and 8.
* Keep the legacy zero-argument autocast probe for torch 2.1 to 2.3
`torch.is_autocast_enabled` only grew its device-name argument in 2.4, and this
repo still ships extras for older ones (cu118onlytorch211, cu121onlytorch220 and
siblings); this very file already branches on `Version(torch.__version__) <
Version("2.4.0")` for custom_fwd. The single try/except therefore caught the
TypeError on those versions and pinned the probe to "always disabled", so
matmul_lora would pre-cast the activation to the weight dtype and then let
autocast cast it again, changing results by that extra rounding.
Three tiers now, resolved once at import: the device-name form on 2.4+, the
zero-argument form on 2.1-2.3, and only if neither answers, fail open to
"assume no ambient autocast".
The 1350-cell matrix is unchanged: 1044 fixed, 306 bit-identical, 0 regressions,
0 cells whose values or dtypes moved.
Reported by Codex.
* Trim comments
* Cover the fused LoRA backward and fp32 NF4 dequant under mixed dtypes
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Move the dY @ B.t() reuse in the fused LoRA backward to its own PR
* Skip the fused LoRA kernels for float32 base weights instead of reconciling dtypes inside them
fast_lora.py and matmul_lora go back to main. patch_peft_model now leaves float32 models on
PEFT's own forward and says why in the patch summary. Inference keeps two fp32 fixes: NF4
dequantize uses bitsandbytes' fp32 kernel, and single-token decode skips the fp16/bf16-only
4bit gemv.
* Gate the fp32 LoRA fallback GPU tests on has_real_cuda
* Trim comments
---------
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
64 lines
2.4 KiB
Python
64 lines
2.4 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
"""float32 base weights skip the fused LoRA kernels, and fp32 NF4 dequantizes with the fp32 kernel."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import types
|
|
|
|
import pytest
|
|
import torch
|
|
import unsloth # noqa: F401
|
|
|
|
from real_accelerator import has_real_cuda
|
|
from unsloth.models.llama import _base_weight_dtype, _fused_lora_skip_reason
|
|
|
|
|
|
def _proj(weight):
|
|
return types.SimpleNamespace(base_layer = types.SimpleNamespace(weight = weight))
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
|
|
def test_base_weight_dtype_plain(dtype):
|
|
assert _base_weight_dtype(_proj(torch.zeros(2, 2, dtype = dtype))) == dtype
|
|
|
|
|
|
def test_base_weight_dtype_reads_quant_state():
|
|
weight = torch.zeros(2, 2, dtype = torch.uint8)
|
|
weight.quant_state = types.SimpleNamespace(dtype = torch.float32)
|
|
assert _base_weight_dtype(_proj(weight)) == torch.float32
|
|
|
|
|
|
def test_float32_skip_reason():
|
|
assert _fused_lora_skip_reason(0, "none") == ""
|
|
assert "float32" in _fused_lora_skip_reason(0, "none", float32_base = True)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not has_real_cuda(), reason = "needs a real CUDA device")
|
|
def test_fp32_nf4_dequant_matches_bitsandbytes():
|
|
import bitsandbytes as bnb
|
|
from unsloth.kernels.utils import fast_dequantize
|
|
|
|
W = torch.randn(256, 128) * 0.02
|
|
p = bnb.nn.Params4bit(W, requires_grad = False, quant_type = "nf4").cuda()
|
|
assert p.quant_state.dtype == torch.float32
|
|
ref = bnb.functional.dequantize_4bit(p.data, p.quant_state)
|
|
torch.testing.assert_close(fast_dequantize(p, p.quant_state), ref, rtol = 0, atol = 0)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not has_real_cuda(), reason = "needs a real CUDA device")
|
|
def test_fp32_nf4_single_token_decode():
|
|
import bitsandbytes as bnb
|
|
from unsloth.kernels.utils import fast_linear_forward
|
|
|
|
layer = bnb.nn.Linear4bit(128, 256, bias = False, compute_dtype = torch.float32, quant_type = "nf4")
|
|
layer.weight = bnb.nn.Params4bit(
|
|
torch.randn(256, 128) * 0.02, requires_grad = False, quant_type = "nf4"
|
|
)
|
|
layer = layer.cuda()
|
|
X = torch.randn(1, 1, 128, device = "cuda")
|
|
ref = X @ bnb.functional.dequantize_4bit(layer.weight.data, layer.weight.quant_state).t()
|
|
torch.testing.assert_close(fast_linear_forward(layer, X), ref, rtol = 1e-4, atol = 1e-4)
|