ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Tighten fused-LoRA dtype-fix comments

Browse source
This commit is contained in:
danielhanchen 2026-07-09 07:47:23 +00:00
commit f0705ef6e5
2 changed files with 6 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

10
unsloth/kernels/fast_lora.py
View file

@ -95,12 +95,10 @@ class LoRA_MLP(torch.autograd.Function):
h = _forward_function(e, g)
i = matmul_lora(h, downW, downW_quant, downA, downB, downS)
# custom_fwd disables autocast, so X may arrive in a different dtype than
# the fused-op compute dtype (e.g. fp32 hidden states from
# fast_rms_layernorm meeting fp16/bf16 base weights). matmul_lora computes
# in the base weight dtype (== e.dtype); keep the saved activation in that
# same dtype so the backward pass stays dtype-consistent. The incoming
# dtype is remembered in ctx.input_dtype and restored on the returned dX.
# custom_fwd disables autocast, so X may mismatch the compute dtype (e.g.
# fp32 fast_rms_layernorm output meeting fp16/bf16 base weights). Pin the
# saved activation to the compute dtype (== e.dtype) so backward stays
# consistent; remember the incoming dtype to restore it on the dX grad.
ctx.input_dtype = dtype
X = X.to(e.dtype)

5
unsloth/kernels/utils.py
View file

@ -1092,9 +1092,8 @@ def matmul_lora(
else:
W = W.contiguous()
# custom_fwd disables autocast, so reconcile the activation dtype to the
# weight (compute) dtype the way autocast would for a plain Linear. This
# covers fp32 hidden states (e.g. fp32 fast_rms_layernorm output) meeting
# fp16/bf16 base weights.
# weight (compute) dtype as autocast would (e.g. fp32 fast_rms_layernorm
# output meeting fp16/bf16 base weights).
if X.dtype != W.dtype:
X = X.to(W.dtype)
dtype = W.dtype