Repository navigation
feat(peft): fused Fast-LoRA SwiGLU MLP autograd (#837) - #1107
Conversation
a04cd72 to
d0115a2
Compare
|
Advisory review while this is still a draft — nothing here blocks anything, and I am posting it now rather than when you mark it ready because one finding is cheapest to act on before #1087 merges. The math is genuinely pinned, and that is the hard half. Dropping the What the suite cannot see is whether the kernel ran at all. Measured, each mutation injected at the top of
The one failure in row 1 is The bias-less row is the material one, and it is the same shape as the finding #1087 closed last round: This is #1163, and the fix is the one that worked on #1087: assert The dtype and NF4 rows need a test to exist first: A non-SiLU MLP can be patched and silently produce different numbers. Two deviations from #837 that are not in the body. (1) The issue's Math block writes The finding that is cheapest before #1087 merges, and it is really about all three modules. That guard list took #1087 four review rounds to get right. There are now three copies of it and two have already drifted. Pulling Nits: For the record: the red |
d0115a2 to
cfa765f
Compare
MakazhanAlpamys
left a comment
There was a problem hiding this comment.
Advisory review of the draft at 8dca1a53.
#1087 has landed (squash 7b78e528), so I reviewed this merged with current main (0235759e). The merge conflicts (add/add in fast_lora.py and tests/test_issue839_fast_lora_single_projection.py) because the branch still carries #1087's four pre-squash commits. The conflict is mechanical: main's copy of both files is byte-identical to e863bbb7. I resolved it by taking your side, and the merged tree differs from main by exactly your delta: 4 files, +875/-64. Everything below was run on that tree, with the card visible so the gpu tests ran: 54 passed, 1 skipped (the MPS test) across the #837 and #839 files, and ruff is clean.
What is closed. I checked each item by execution.
Every mutation from my last round is now killed. Each one is injected at the top of _forward (fast_lora_mlp.py:199):
| mutation | last round | now |
|---|---|---|
| kernel never runs | 1 failed (a < in the bytes test) |
10 failed |
| delegate on 2-D / on 3-D input | 2-D survived | 4 / 6 failed |
| delegate on a bias-less base | survived | 5 failed |
delegate when x.dtype.itemsize != 4 |
survived | 5 failed, including the bf16 CPU rows, so CI sees it |
delegate when quant_state is not None |
survived | 1 failed: the new gpu NF4 test, so only on a card, like #1087 |
- The math is still pinned: dropping the
(1 + G(1-sig))term fails 8 tests. - The gate/up concatenation (one
X @ A_gu^T, onedAGEMM) is implemented now. - MoE experts and a per-module
act_fncheck are both in, and the mutation that removes either one fails its test. - The shared-helper extraction is the change I most wanted, and it holds. #839's delegation tests still pin every guard through
_projection_state: removing the disabled/merged, dropout,lora_bias,fan_in_fan_outor variant guard fails 1 or 2 tests each.
I also measured four things that hold today but are not in the file:
disable_adapter()(DPO's reference pass) gives exactly the base-only logits, max diff 0.- An fp32-LoRA / bf16-base block (Soup's layout since #429) takes the fused path. Against a float64 reference its error is 1.00x / 0.71x / 1.47x peft's (output / dX / LoRA grads), inside #792's 2x rule.
- NF4 parity: the numbers are under item 5 below.
- A streamed model: the numbers are under the first non-blocking item.
Blocking before this leaves draft
-
Rebase onto
mainand drop513e7f9e..e863bbb7.merge-freshnessis red ("live merge has conflicts"). The 13 required checks are green, but they ran 12:06–14:25Z on 2026-09-23, before #1087 merged at 17:43Z. Please also refresh the description. It still namesa04cd724, which is not on the branch, and it does not mention thatcfa765f2refactors the mergedfast_lora.py. -
Please drop the
tests/test_issue839_fast_lora_single_projection.pyhunks fromcfa765f2and build on #1173 instead. Thank you for tightening #1087's pin while you were in there. The trouble is that #1163 was claimed by @kok-o on 2026-09-22 at 19:59Z, and #1173 implements it. Your commit is dated the next morning and carries the same three edits: the NF4grad_fnassertion, thedtypeparametrisation oftest_the_output_carries_the_kernels_own_grad_fn, and#968to#776in the NF4 comment. #1173's copy is the complete one. It also adds the vacuity precondition on the NF4 reference, which #1163's acceptance criteria ask for, and it fixes the#968in the module docstring at:30, which your copy leaves. Two copies would only conflict. Keep thefast_lora.pyrefactor; that part is exactly what I asked for. -
The refactor changes merged #1087 behaviour under autocast.
_fast_lora_single_forwardnow always returnsout.to(input_dtype)(fast_lora.py:368). With an fp32 model and input undertorch.autocast("cpu", dtype=torch.bfloat16):- On
mainthe kernel returns bf16 with_FastLoraSingleProjectionBackward, and peft also returns bf16. - On this branch it returns fp32 via
ToCopyBackward0. - The MLP kernel (
fast_lora_mlp.py:227) and #1108's QKV kernel (fast_lora_qkv.py:218) do the same. The block's output is fp32 where peft's is bf16.
Only the 4-bit path should cast back, because bnb's
Linear4bitdoes. I checked this fix:return out if work_x is x else out.to(input_dtype)
With the equivalent
if work_x is not x:in QKV, all three return bf16 with their owngrad_fnunder autocast. All 74 tests across the three files still pass, and NF4 with fp32 input still casts back. Please add an autocast test that asserts the dtype equals peft's and thegrad_fnis the kernel's. - On
-
8-bit layers are patched and give garbage, with no error.
patch_fast_lora_mlpmatches on child names plushasattr(lora_A)(fast_lora_mlp.py:266-276), and_dense_weighttreats any weight without aquant_stateas dense (fast_lora.py:225). I built the layers the way a real 8-bit load does (is_loaded_in_8bit=True, so peft wraps them aslora.bnb.Linear8bitLt). The single-projection patcher takes 0 of them, because of its type allowlist (fast_lora.py:384-396). The MLP patcher takes the block, and its output is NaN. #792 scopes 8-bit out. Please reuse the allowlist: factor it out ofpatch_fast_lora_single_projectionand require it of every adapted child. For unadapted children, refuse a weight that is not floating point and has noquant_state. Add agputest asserting that the patch count is 0. -
Four mutations still survive (23 passed each):
- delegate unless all three projections are adapted;
- delegate when the ranks differ;
- NF4 backward dequantises gate with up's quant-state tensors (
:134); - the same swap in the NF4 forward (
:50).
The partial-adapter and unequal-rank cases are covered only by agreement tests, and the NF4 test has no parity check and no backward. Please parametrise the
grad_fntest over the five target subsets your parity test already uses, plus onerank_patterncase. Then add NF4 parity (Y, dX and every dA/dB against unpatched peft on the sameLinear4bit) in the style of #839'sTestNf4Parity.The shipped NF4 math is correct. Against unpatched peft with
lora.bnb.Linear4bit, the output differs by at most 6.25e-2 at |y| up to 4.7, and dX by 2.0e-3. Choose the tolerance with care, though: #839'srtol=atol=1e-2fails on the MLP output (1.1–3.6% of elements over 3 seeds), and the kernel is not at fault. Against a float64 reference on the same dequantised weights, your error is 1.66x / 1.30x / 0.78x peft's. Three chained bf16 NF4 matmuls need #792's 2x-versus-float64 criterion rather than #839's fixed tolerance.
Non-blocking
- There is no streamed test (a tracker criterion), and it would pass today. I ran a 4-layer tiny Llama with 2 stream buffers, and the MLP and QKV kernels patched every layer. The loss matched the resident unpatched twin to 7 digits and all 40 LoRA grads matched (worst 3.4e-7). You can lift #1087's
TestStreamedModelfor it. - bf16 parity and the saturation test run only on Apple MPS, so they skip on CPU, CUDA and CI. CPU bf16 would cover them.
- A trainable unadapted projection silently stops training. With
modules_to_save=["down_proj"], the forward uses the trained copy, but its gradient isNone. Under peft the same weight gets a gradient with an L1 norm of 1.27. Soup does not exposemodules_to_savetoday, so this is latent. The guard is one line: in the unadapted branch, delegate if the base has anyrequires_gradparameter. _module_uses_silufails open whenact_fnis missing (:234). A GELU-inline MLP in a silu-config model is patched and its output moves by 1.6e-2. ReturningFalsethere is safer, since every in-scope HF MLP hasact_fn.- #1165 against the merged patcher is benign for this PR. MLP always wins over single-projection in either order, and all four patch/unpatch orders end with the never-patched
grad_fns. The one leftover is thatunpatch_fast_lora_mlp(:291) writes the bound method back as an instance attribute. #1108's restore-or-delete is the pattern to copy. - The refactor dropped 16 comment lines of hard-won rationale from
fast_lora.py: the Conv1D andfan_in_fan_outreasoning,lora_biasand dropout. Please carry them into_projection_state. - Smaller gaps:
No changelog fragment or docs are needed yet: nothing is wired, the same as #1087.
Nits:
- stray blank lines in
tests/test_issue837_fast_lora_mlp.pyat:70and:387; - the compute-dtype block and the quant-part slicing are duplicated with the QKV module;
qmeta["_count"]duplicateslen(meta["roles"]).
8dca1a5 to
ae25111
Compare
MakazhanAlpamys
left a comment
There was a problem hiding this comment.
Advisory review of the draft at ae25111a.
Thank you for the rebase. It is exactly the clean-up I asked for in item 1. git range-diff shows #1087's four pre-squash commits gone, and your two commits carried over unchanged (cfa765f2 = 8e1c27b0, 8dca1a53 = ae25111a, identical patch-ids). merge-freshness is green, and the PR diff is now just your four files, +875/-64.
That also means the code did not change. All four files have the same blobs as at 8dca1a53, so my review of that head still applies in full. Since then, main has moved three commits, none of which touch these files. I re-ran everything on this head merged with current main (1c053a8c), with the card visible so the gpu tests ran:
- 54 passed, 1 skipped (the Apple MPS test) across the #837 and #839 files. Ruff and
git diff --checkare clean. - The mutation results are identical to last round: 9 of 13 killed on the MLP kernel, 5 of 6 on the shared helper. The one helper survivor,
len(active) != 1, also survives onmain. - Your new test file still fails on
main(23 failed, 1 skipped), as it should.
Blocking before this leaves draft. These are unchanged from last round. I re-checked each one by running it.
-
The #1163 hunks are still in
tests/test_issue839_fast_lora_single_projection.py. All three are there verbatim: the NF4grad_fnassertion at:858, thedtypeparametrisation at:950and:974, and#968to#776at:861. The commit title also still says(#837, #1163). #1163 is @kok-o's issue, and #1173 implements it. Please drop those three hunks and remove#1163from the commit title. Thefast_lora.pyrefactor stays.This now affects the merge order, not just credit:
- I merged #1173 onto
mainand then this branch on top of it. That gives a content conflict in that test file, so whichever of the two lands second has to be reworked. - With the hunks dropped, your
srcpasses main's copy of the file (50 passed, 1 skipped, together with your #837 file) and #1173's copy (62 passed, 1 skipped). The two PRs then merge cleanly in either order.
- I merged #1173 onto
-
The autocast regression is still there (
fast_lora.py:368,fast_lora_mlp.py:227). Take an fp32 model and input undertorch.autocast("cpu", dtype=torch.bfloat16):- On
main, the single-projection kernel returns bf16 with_FastLoraSingleProjectionBackward. - On this branch, both the single-projection kernel and the MLP block return fp32 via
ToCopyBackward0, where peft returns bf16.
The one-line fix from last round still applies:
return out if work_x is x else out.to(input_dtype). Please add an autocast test with it. - On
-
8-bit is still patched, and the output is garbage. I dispatched the layers the way a real 8-bit load does (
is_loaded_in_8bit=True, so peft wraps them aslora.bnb.Linear8bitLt), on CUDA. The single-projection patcher takes 0 layers, and the MLP patcher takes 1. The output is NaN, thegrad_fnis_FastLoraSwiGLUBackward, and nothing raises. -
The same four mutations survive, with all 23 tests passing each time:
- delegate unless all three projections are adapted;
- delegate when the ranks differ;
- the NF4 gate/up quant-state swap in the backward (
:134); - the same swap in the forward (
:50).
The fix path is unchanged: parametrise the
grad_fntest over the five target subsets plus onerank_patterncase, then add NF4 parity with a backward, judged by #792's 2x-versus-float64 rule.
Please also refresh the description. It still says "Stacked on #1087", names a04cd724, and quotes 33 passed and 18,016 bytes.
Non-blocking. These are unchanged too. I re-ran the ones I could:
- With
modules_to_save=["down_proj"], the trained copy gets no gradient. Under peft its gradients have L1 norms of 5.69 (weight) and 5.61 (bias). _module_uses_silufails open. An inline-GELU MLP in a silu-config model is patched, and its output moves by 1.8e-1.unpatch_fast_lora_mlpleavesforwardbehind as an instance attribute.- The refactor removes 19 comment lines of
fast_lora.pyrationale and adds back 3. - There is still no streamed test, and the bf16 tests still run only on MPS.
- The MLP backward computes dX even when X needs no grad.
The details and fixes for each are in my previous review.
No changelog fragment or docs are needed while nothing is wired.
Nit: the stray blank lines at tests/test_issue837_fast_lora_mlp.py:70 and :387 are still there.
ae25111 to
de508b7
Compare
|
Implemented every blocking review item and the actionable non-blocking hardening on rebased
Local validation: 104 passed, 5 expected GPU/platform skips, Ruff and diff-check clean. The CUDA-marked NF4 tests are present for the maintainer's card; no CUDA result is claimed here. |
|
@Amix29, thank you for this round. It closes most of what was open, and I checked each item by running it rather than by reading the reply. What is closed. I merged
Before this leaves draft. Four items, all CPU-sized.
After these, please mark it ready for review. Before it can merge, it will also need one run of the three NF4 tests on a CUDA card. That cannot happen today, and I would rather say so than guess a date. Non-blocking
A heads-up on the dropout work in #839. @sloemo01 has taken the dropout stretch item on the single-projection path. After your refactor, that path's dropout check lives in the shared CI: all 13 required checks on What I ran: the two test files above on CPU ( |
de508b7 to
4eea730
Compare
|
Addressed the 2026-09-25 CPU review at head 4eea730, rebased onto current main. Added the single-projection autocast pin, refused non-floating unadapted siblings, scaled the NF4 fixture and replaced fixed tolerance with the <=2x-peft-vs-float64 oracle rule, proved the MLP kernel runs in the streamed test, pinned missing act_fn fail-closed, restored shared guard rationale, and added a real mixed Linear8bitLt regression. Validation: 107 passed, 5 expected GPU/MPS skips; Ruff and git diff --check clean. Marked ready for review; the requested CUDA NF4 run remains external and is not claimed here. |
|
@Amix29, thank you for this round. Everything the 09-25 review asked for is in and pinned. I checked each item by running it on What is closed
What I ran: One blocking item: the bf16 backward fails your own 2x-vs-float64 rule on one of the three NF4 seeds I copied There are two independent causes:
Over 20 seeds, the NF4 copy fails 3 (2.07x, 2.71x and 4.93x). Nothing in CI can see this, because the NF4 test is The fix below is what I verified: # forward, fast_lora_mlp.py:73, :78, :88 (same pattern for u with hu/bul/su and y with hd/bdl/sd)
promoted = torch.promote_types(g.dtype, hg.dtype)
g = torch.add(g.to(promoted), torch.matmul(hg, bgl.t()).to(promoted), alpha=float(sg)).to(g.dtype)
# backward, fast_lora_mlp.py:132-134
opmath = torch.promote_types(g.dtype, torch.float32) # float64 stays float64 for gradcheck
g_hi = g.to(opmath)
sig = torch.sigmoid(g_hi)
grad_u = grad_m * silu_g
grad_g = (grad_m.to(opmath) * u.to(opmath) * sig * (1 + g_hi * (1 - sig))).to(g.dtype)With it:
Each half fixes its own case: (1) alone takes seed 11 from 4.93x to 1.00x, and (2) alone takes the bf16 seed from 2.39x to 1.09x. To keep it pinned in CI, let the NF4 test also run on CPU. macOS CI already installs bitsandbytes 0.50.2 and runs your Int8-on-CPU tests (they PASS there). Replace the method's @pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)])
@pytest.mark.parametrize("seed", [3, 7, 11])
def test_nf4_forward_backward_and_all_adapter_grads_match_peft(self, seed, device):Then use That row pins cause (1). Cause (2) needs a bf16 case, and this one needs no bitsandbytes. Paste it at the end of your file (it uses your @pytest.mark.parametrize("seed", range(12))
def test_bf16_mlp_stays_within_twice_pefts_error_against_float64(seed):
torch = _deps()
import torch.nn as nn
from peft import LoraConfig, inject_adapter_in_model
from soup_cli.utils.fast_lora_mlp import patch_fast_lora_mlp
class TinyMLP(nn.Module):
def __init__(self):
super().__init__()
self.gate_proj = nn.Linear(32, 48, bias=False)
self.up_proj = nn.Linear(32, 48, bias=False)
self.down_proj = nn.Linear(48, 32, bias=False)
self.act_fn = nn.SiLU()
def forward(self, x):
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
class Model(nn.Module):
def __init__(self):
super().__init__()
self.config = SimpleNamespace(hidden_act="silu")
self.mlp = TinyMLP()
def forward(self, x):
return self.mlp(x)
torch.manual_seed(seed)
model = Model()
with torch.no_grad():
model.mlp.up_proj.weight.mul_(4)
inject_adapter_in_model(
LoraConfig(
r=4, lora_alpha=8, lora_dropout=0.0,
target_modules=["gate_proj", "up_proj", "down_proj"],
),
model,
)
_randomise_b(model)
for name, param in model.named_parameters(): # bf16 base, fp32 LoRA (#429 layout)
param.data = param.data.float() if "lora_" in name else param.data.to(torch.bfloat16)
x0 = torch.randn(2, 5, 32, dtype=torch.bfloat16)
def run(module, x):
x = x.clone().requires_grad_(True)
out = module(x)
out.double().square().mean().backward()
grads = [p.grad for n, p in sorted(module.named_parameters()) if "lora_" in n]
result = [t.detach().double().clone() for t in (out, x.grad, *grads)]
module.zero_grad(set_to_none=True)
return result
oracle = run(copy.deepcopy(model).double(), x0.double())
reference = run(model, x0)
assert patch_fast_lora_mlp(model) == 1
fast = run(model, x0)
for got, ref, exact in zip(fast, reference, oracle):
assert (got - exact).abs().max() <= 2 * (ref - exact).abs().max() + 1e-8Non-blocking
Heads-up on #839's dropout item. The shared dropout guard in Once the fix is in, the one run of the NF4 tests on a CUDA card from the 09-25 review is what remains before merge. Required CI is green on this head, but it ran on 09-25, before today's |
4eea730 to
662a525
Compare
|
Implemented the blocking bf16 numerics round and the remaining CPU-sized pins at
Local M4 Max result: 127 passed, 5 skipped across #837/#839; Ruff, preview E30 and |
|
@Amix29, thank you, that was fast and it is exactly the fix. I re-checked it by running it on What is closed
What I ran: Two small test-only fixes before merge
Non-blocking
After the two fixes, this waits on the required checks (all 13 are still queued for |
|
Addressed both test-only items and the useful non-blocking pins at
Local M4 Max result remains 127 passed, 5 expected CUDA/platform skips across #837/#839. Ruff, preview E30, and |
|
@Amix29, thank you. This round is exactly what was asked, and nothing more: one commit on top of What is closed
What I ran: Where this stands From the review side this is ready to merge. Two things remain, and neither is on you:
#1108 is not affected. It is stacked on this exact head, and I checked that it lands cleanly on top of it. |
|
@Amix29, the CUDA run is done. RTX 5070 (sm_120, torch 2.14 cu130, bitsandbytes 0.50.2), on
NF4 fast/PEFT error ratios (output, input gradient, then the six adapter grads):
The worst is 1.35, inside #792's 2x rule, and #839's NF4 parity test passes too, so no re-run at the larger shape was needed. That was the last open item: this merges as soon as the 13 required checks are green on |
|
Merged as |
feat(peft): shared-X Fast-LoRA QKV autograd (#838) With peft LoRA on `q_proj`/`k_proj`/`v_proj`, each projection runs its own `X @ A^T` in forward and its own `dA` and A-side `dX` GEMMs in backward, and `dX` is assembled from six contributions. - `soup_cli.utils.fast_lora_qkv` adds a shared-X Q/K/V autograd Function: `A_cat = cat(A_q, A_k, A_v)`, one `X @ A_cat^T` GEMM in forward, one `dA` GEMM and one accumulated `dX` in backward. The base GEMMs stay separate, since packed NF4 weights cannot be concatenated. - The hook: the patched `q_proj` computes Q, K and V together and leaves K and V in a short-lived cache on the attention module, keyed by the identity of the input tensor, and `k_proj`/`v_proj` consume it. A call with a different tensor (cross-attention) or extra arguments delegates to peft and drops the cache. Modules whose q/k/v input widths differ (TrOCR, Dia, T5Gemma and Idefics cross-attention) are not patched. - It handles any adapted subset (Soup's default on Llama is q and v), unequal ranks, GQA widths, a frozen base bias, 2-D/3-D and non-contiguous inputs, CPU autocast and NF4 bases. It delegates to peft for dropout, merged or disabled adapters, LoRA variants, 8-bit layers (adapted, or an unadapted non-float sibling) and `modules_to_save`. The LoRA adds use #1107's shared promote-add-cast helper. - Ownership with the single-projection patcher holds in both patch orders: q/k/v are marked group-owned, each patcher skips the other's children, and unpatch restores only the forwards QKV installed. Nothing is user-reachable yet. No flag or call site exists until #792 wires one, so there is no changelog fragment or docs change. Verified at a7077a8 merged with main (dccc868, and again on cb5bd02): - CPU: 210 passed, 10 skipped across the #837, #838 and #839 test files. The new file fails 84 tests (4 skipped) on main. - float64 gradchecks on GQA shapes with ranks (2, 0, 3) and (2, 2, 2); fp32 parity against a real Llama attention with GQA and q/v adapters; a 12-seed bf16-base/fp32-LoRA rule, within 2x of peft's error against float64; NF4 parity on CPU for three seeds; a streamed 4-layer Llama. - CUDA (RTX 5070 Laptop, sm_120): 8 passed with `-m gpu`, including the three QKV NF4 rows. A 40-seed NF4 sweep at 6 and 34 rows passed every check; the worst fast/peft error ratio against a float64 oracle was 1.43. - Mutations: 23 of 24 fail a test. The kernel never running fails 57, `k` or `v` never consuming the cache 56 each, either `dX` term dropped 24 each, and a `k`/`v` quant-state swap in the NF4 forward or backward 3 of 3. The survivor is the refusal of children that another group patcher owns; today only the QKV patcher sets that marker. Refs #838.
…ated 6 (#1735) `docs/issue-difficulty.md` cited #792 as the example of a level-10 issue. The three hand-written kernels that made it a 10 are merged (#1087, #1107, #1108); what is left on #792 is the `training.fast_lora` flag and its wiring, re-rated to `difficulty:6` with the reason in its thread. No open issue carries level 10 at the moment, so the example cell of that row says so. One line, docs only; the table's definitions are unchanged. Refs #792.
Refs #837
What
Review follow-ups
Validation on Apple M4 Max
No CUDA performance or parity claim is made from this Mac. The requested final CUDA NF4 run remains external validation before merge. No changelog fragment or user docs are added because this kernel is not wired into a user-facing training path yet.