Skip to content

feat(peft): fused Fast-LoRA SwiGLU MLP autograd (#837) - #1107

Merged
MakazhanAlpamys merged 6 commits into
MakazhanAlpamys:mainfrom
Amix29:codex/fix-837-fast-lora-mlp
Sep 28, 2026
Merged

MakazhanAlpamys merged 6 commits into
MakazhanAlpamys:mainfrom
Amix29:codex/fix-837-fast-lora-mlp

Conversation

@Amix29

@Amix29 Amix29 commented Sep 20, 2026 •

Copy link
Copy Markdown
Contributor

Refs #837

What

  • adds a structural gate/up/down SiLU/swish Fast-LoRA path with a hand-written backward;
  • concatenates gate/up LoRA-A projections into one forward GEMM and one dA GEMM;
  • supports partial adapters and unequal ranks;
  • delegates to PEFT for dropout, merged/disabled adapters, LoRA variants, unsupported 8-bit/custom layers, modules_to_save, unsupported activations and MoE experts;
  • shares dense/NF4 projection-state handling with the merged single-projection kernel;
  • preserves streamed-pool safety by routing every backward-read packed weight and quant-state tensor through save_for_backward.

Review follow-ups

  • rebased onto current main and dropped all Fast-LoRA: the kernel-is-taken pin still misses 4-bit and non-fp32, and gradcheck cannot backstop it #1163/test(peft): assert grad_fn on NF4 and add dtype to fast-LoRA kernel pin #1173 test hunks;
  • preserves the PEFT bf16 result under CPU autocast rather than forcing fp32;
  • uses a shared promote-add-cast helper for mixed bf16-base/fp32-LoRA updates;
  • evaluates the SiLU derivative in fp32 opmath before rounding back to the activation dtype;
  • pins real PEFT Linear8bitLt refusal, parent/module activation gates and modules_to_save ownership;
  • pins kernel execution across all/partial targets, equal/unequal ranks, bias, 2-D/3-D inputs, fp32 and bf16;
  • runs the three-seed NF4 forward/backward parity rule on CPU, prints the measured fast/PEFT error ratios, and retains CUDA rows for the external card;
  • protects CPU bf16 matmuls with the Windows-safe aten_half_matmuls fixture and separately pins the down-projection add;
  • adds a four-layer/two-buffer streamed Llama comparison against resident unpatched PEFT;
  • pins each projection unadapted under float64 gradcheck;
  • measures the Llama-3.1-8B saved-tensor contract at N=2048/r=64 on meta, deduplicated by storage identity, and requires two to three [N, f] tensors saved;
  • unpatching restores class dispatch without leaving an instance forward attribute.

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.

@MakazhanAlpamys

Copy link
Copy Markdown
Owner

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 (1 + G(1-sig)) term from the SiLU derivative fails 8 of 12 tests. The structural gate/up/down fusion, partial adapters, unequal ranks, the seven delegation cases, save_for_backward over every backward-read tensor, and reconstructing the NF4 quant state rather than retaining dense weights — all present and all doing what #837 asks. 12 passed / 1 skipped, ruff clean, merges cleanly with current main, and every non-skipped test fails on main.

What the suite cannot see is whether the kernel ran at all. Measured, each mutation injected at the top of _forward (fast_lora_mlp.py:230):

mutation result
return original_forward(x, *args, **kwargs) — kernel is a total no-op 1 failed / 11 passed
delegate when x.dim() == 2 12 passed — SURVIVED
delegate when the base has no bias 12 passed — SURVIVED
delegate when x.dtype.itemsize != 4 12 passed — SURVIVED
delegate when base.weight.quant_state is not None 12 passed — SURVIVED

The one failure in row 1 is test_saved_tensor_bytes_are_lower_than_peft — it trips only because the fused path happens to save fewer bytes, and only because the assertion is < rather than <=. Nothing in the file asserts what ran. Eight parity tests, the float64 gradcheck, the real-LlamaMLP test and both patching tests all stay green while the kernel does nothing.

The bias-less row is the material one, and it is the same shape as the finding #1087 closed last round: LlamaMLP builds all three projections with bias=False — I confirmed it by execution — so test_real_llama_mlp_matches_peft runs exactly the configuration that survives, and passes, because agreement is free under delegation.

This is #1163, and the fix is the one that worked on #1087: assert type(out.grad_fn).__name__ == "_FastLoraSwiGLUBackward", parametrised over bias=(True, False) × input rank ((4,8), (2,5,8)) × dtype=(fp32, bf16). Lift it from #1087's merged head 1d62c347, not from 89031be9 — this branch is stacked on the pre-widening head, so the class you would copy today is the narrow one.

The dtype and NF4 rows need a test to exist first: grep -c 'nf4\|quant\|bnb\|bitsandbytes' tests/test_issue837_fast_lora_mlp.py is 0, against a module that ships a complete NF4 branch.

A non-SiLU MLP can be patched and silently produce different numbers. patch_fast_lora_mlp (:255-260) gates on the parent model's config.hidden_act and then patches every module that merely has the three child names — it never looks at the module's own act_fn. Demonstrated: a gate/up/down module using nn.GELU, inside a wrapper whose config.hidden_act == "silu", is patched and its output moves by 4.6e-2. I checked the obvious real instance and it is safe today — Gemma3TextConfig has no hidden_act at all, only hidden_activation, so the patcher returns 0 — so this is latent rather than live. It is one getattr(module, "act_fn", None) check away from impossible, and #837's own wording ("detect by structure and by config.hidden_act") reads naturally as a per-module assertion.

Two deviations from #837 that are not in the body. (1) The issue's Math block writes H_gu = X @ cat(A_g, A_u)^T # ONE GEMM and dA_gu = dH_gu^T @ X # ONE GEMM; :67-72 and :137-155 do two separate GEMMs each way, so the concatenation that is the stated point of the fusion is not implemented. Worth either doing or saying why not — #1108 does concatenate, so the pair are inconsistent with each other. (2) #837 scopes MoE experts out, but the patcher matches on child names alone, so a Qwen3-MoE or Mixtral expert block is in scope by accident.

The finding that is cheapest before #1087 merges, and it is really about all three modules. _dense_weight (:28-34) is byte-identical to fast_lora_qkv.py:27-33 and duplicates what fast_lora.py does inline — three copies. _projection_state (:170-224) re-implements, without the explanatory comments, the same seven-guard delegation list that fast_lora.py:243-276 carries with them — and drops the streamed-quant-state repair at fast_lora.py:284-288. It also has no set_compute_type / compute_dtype handling, which fast_lora.py:294-302 does for a real Linear4bit base, so under NF4 with an fp32 input the two paths can disagree on compute precision.

That guard list took #1087 four review rounds to get right. There are now three copies of it and two have already drifted. Pulling _dense_weight and _projection_state into fast_lora.py as shared helpers is the single highest-value thing in this PR, and it is easier to do while #1087's file is the only one merged.

Nits: _mlp_function() caches the Function class in a module-level global — fine, but the sibling carries an identical copy of the idiom; and there are stray blank lines at tests/...:67 and :263 (ruff's E3 is not enabled, so nothing complains).

For the record: the red gate check is not yours. The workflow ran scripts/merge_freshness.py out of your checkout, where a branch cut before #1017 does not have it. Fixed on main in 549ec25d; a rebase clears it.

@MakazhanAlpamys MakazhanAlpamys left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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, one dA GEMM) is implemented now.
  • MoE experts and a per-module act_fn check 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_out or 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

  1. Rebase onto main and drop 513e7f9e..e863bbb7. merge-freshness is 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 names a04cd724, which is not on the branch, and it does not mention that cfa765f2 refactors the merged fast_lora.py.

  2. Please drop the tests/test_issue839_fast_lora_single_projection.py hunks from cfa765f2 and 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 NF4 grad_fn assertion, the dtype parametrisation of test_the_output_carries_the_kernels_own_grad_fn, and #968 to #776 in 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 #968 in the module docstring at :30, which your copy leaves. Two copies would only conflict. Keep the fast_lora.py refactor; that part is exactly what I asked for.

  3. The refactor changes merged #1087 behaviour under autocast. _fast_lora_single_forward now always returns out.to(input_dtype) (fast_lora.py:368). With an fp32 model and input under torch.autocast("cpu", dtype=torch.bfloat16):

    • On main the 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 Linear4bit does. 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 own grad_fn under 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 the grad_fn is the kernel's.

  4. 8-bit layers are patched and give garbage, with no error. patch_fast_lora_mlp matches on child names plus hasattr(lora_A) (fast_lora_mlp.py:266-276), and _dense_weight treats any weight without a quant_state as 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 as lora.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 of patch_fast_lora_single_projection and require it of every adapted child. For unadapted children, refuse a weight that is not floating point and has no quant_state. Add a gpu test asserting that the patch count is 0.

  5. 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_fn test over the five target subsets your parity test already uses, plus one rank_pattern case. Then add NF4 parity (Y, dX and every dA/dB against unpatched peft on the same Linear4bit) in the style of #839's TestNf4Parity.

    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's rtol=atol=1e-2 fails 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 TestStreamedModel for 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 is None. Under peft the same weight gets a gradient with an L1 norm of 1.27. Soup does not expose modules_to_save today, so this is latent. The guard is one line: in the unadapted branch, delegate if the base has any requires_grad parameter.
  • _module_uses_silu fails open when act_fn is missing (:234). A GELU-inline MLP in a silu-config model is patched and its output moves by 1.6e-2. Returning False there is safer, since every in-scope HF MLP has act_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 that unpatch_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 and fan_in_fan_out reasoning, lora_bias and dropout. Please carry them into _projection_state.
  • Smaller gaps:
    • dX is computed even when X needs no grad; #1087 skips it.
    • The gradcheck covers only the all-adapted case; #837 asks for each projection unadapted in turn.
    • The len(active) != 1 guard is unpinned. That one is pre-existing: it also survives on main, so it is for the tracker.

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.py at :70 and :387;
  • the compute-dtype block and the quant-part slicing are duplicated with the QKV module;
  • qmeta["_count"] duplicates len(meta["roles"]).

Amix29 added a commit to Amix29/Soup that referenced this pull request Sep 23, 2026
@Amix29
Amix29 force-pushed the codex/fix-837-fast-lora-mlp branch from 8dca1a5 to ae25111 Compare September 23, 2026 19:52

@MakazhanAlpamys MakazhanAlpamys left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 --check are 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 on main.
  • 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.

  1. The #1163 hunks are still in tests/test_issue839_fast_lora_single_projection.py. All three are there verbatim: the NF4 grad_fn assertion at :858, the dtype parametrisation at :950 and :974, and #968 to #776 at :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 #1163 from the commit title. The fast_lora.py refactor stays.

    This now affects the merge order, not just credit:

    • I merged #1173 onto main and 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 src passes 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.
  2. The autocast regression is still there (fast_lora.py:368, fast_lora_mlp.py:227). Take an fp32 model and input under torch.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.

  3. 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 as lora.bnb.Linear8bitLt), on CUDA. The single-projection patcher takes 0 layers, and the MLP patcher takes 1. The output is NaN, the grad_fn is _FastLoraSwiGLUBackward, and nothing raises.

  4. 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_fn test over the five target subsets plus one rank_pattern case, 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_silu fails open. An inline-GELU MLP in a silu-config model is patched, and its output moves by 1.8e-1.
  • unpatch_fast_lora_mlp leaves forward behind as an instance attribute.
  • The refactor removes 19 comment lines of fast_lora.py rationale 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.

@Amix29

Amix29 commented Sep 24, 2026

Copy link
Copy Markdown
Contributor Author

Implemented every blocking review item and the actionable non-blocking hardening on rebased main (de508b7f):

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.

@MakazhanAlpamys

Copy link
Copy Markdown
Owner

@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 de508b7f with current main (clean; the PR diff is still exactly your three files) and ran everything on CPU:

  • tests/test_issue837_fast_lora_mlp.py + tests/test_issue839_fast_lora_single_projection.py: 103 passed, 6 skipped (the MPS test and the five gpu tests). That is your 104/5 without the Mac-only test. Ruff and git diff --check are clean, and your new file fails on main (66 failed, 4 skipped), as it should.
  • The Fast-LoRA: the kernel-is-taken pin still misses 4-bit and non-fp32, and gradcheck cannot backstop it #1163 hunks are gone. The PR no longer touches any tests/test_issue839_* file, and no commit message mentions Fast-LoRA: the kernel-is-taken pin still misses 4-bit and non-fp32, and gradcheck cannot backstop it #1163. Thank you.
  • The two mutations that survived last round now fail. Delegating unless all three projections are adapted fails 32 tests, and delegating when the ranks differ fails 8. The fast-path matrix is exactly what was needed.
  • The MLP autocast fix, the 8-bit refusal for adapted layers, the modules_to_save refusal and the restore-or-delete unpatch are each pinned: removing any one of them fails its test.
  • The math is still pinned: dropping down's LoRA term from dM (fast_lora_mlp.py:127) fails 7 tests, including the float64 gradcheck.

Before this leaves draft. Four items, all CPU-sized.

  1. The autocast fix in the refactored single-projection path is not pinned. Last round's regression was in the merged fast_lora.py forward, which your refactor rewrote. The fix is there (fast_lora.py:391), but reverting it to an unconditional out.to(input_dtype) leaves all 103 tests green, because your autocast tests go through the MLP and QKV patchers only. This test passes on this branch and on main, and fails with that revert:

    def test_single_projection_keeps_pefts_bfloat16_output_under_cpu_autocast():
        torch = pytest.importorskip("torch")
        pytest.importorskip("peft")
        import torch.nn as nn
        from peft import LoraConfig, inject_adapter_in_model
    
        from soup_cli.utils.fast_lora import patch_fast_lora_single_projection
    
        class Model(nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = nn.Linear(8, 8)
    
        torch.manual_seed(0)
        model = Model()
        inject_adapter_in_model(
            LoraConfig(r=2, lora_alpha=4, lora_dropout=0.0, target_modules=["o_proj"]), model
        )
        x = torch.randn(2, 5, 8, requires_grad=True)
        with torch.autocast("cpu", dtype=torch.bfloat16):
            reference = model.o_proj(x)
        assert reference.dtype == torch.bfloat16
    
        assert patch_fast_lora_single_projection(model) == 1
        with torch.autocast("cpu", dtype=torch.bfloat16):
            output = model.o_proj(x)
    
        assert output.dtype == reference.dtype
        assert type(output.grad_fn).__name__ == "_FastLoraSingleProjectionBackward"
  2. The second half of the 8-bit item from the 09-23 review is still open: unadapted children with a non-float weight. The allowlist now refuses adapted 8-bit layers, but _projection_state still treats any unadapted weight without a quant_state as dense. I built an MLP whose adapted gate_proj is an nn.Linear and whose unadapted up_proj / down_proj are bnb.nn.Linear8bitLt. (bitsandbytes 0.50.2 quantises Int8Params on CPU, so this runs without a card.) The patcher takes the block, and the output reaches 7.05e4 where peft's is at most 0.312. It carries _FastLoraSwiGLUBackward, and nothing raises. The same block shape in feat(peft): shared-X Fast-LoRA QKV autograd (#838) #1108's QKV kernel is off by 473. One line before compute_dtype = None (fast_lora.py:189) closes it for both kernels:

        if qstate is None and not weight.is_floating_point():
            return None

    With it, both kernels match peft exactly, and all 168 tests across the Fast-LoRA autograd: fused LoRA MLP block (gate/up/down around SwiGLU) with a hand-written backward #837, Fast-LoRA autograd: LoRA on q/k/v with one shared X (concatenated A, one dX accumulation) #838 and Fast-LoRA autograd: single-projection LoRA Function (o_proj, and the fallback for any other adapted Linear) #839 files still pass (I ran them on feat(peft): shared-X Fast-LoRA QKV autograd (#838) #1108's tree, which contains this branch). Soup's own 8-bit loads quantise every projection, so reaching this takes a mixed block. It is still the guard the review asked for, and it is one line.

  3. The NF4 test cannot see a swapped quant state, and its tolerance is the one that failed on a card. Nobody can run the gpu tests on CUDA today. So I copied test_nf4_forward_backward_and_all_adapter_grads_match_peft verbatim and changed only the device from cuda to cpu (bitsandbytes 0.50.2 quantises and dequantises NF4 on CPU). The copy passes for all three seeds. Then I swapped quant states:

    • dequantising gate with up's quant state in the forward (fast_lora_mlp.py:52) fails only seed 7;
    • the same swap in the backward (:138) passes all three seeds.

    The fixture is the reason, not the kernel. Gate and up come from the same init distribution, so their absmax values nearly match, and a swap moves the dequantised gate weight by only 1.3%, 3.5% and 1.1% (seeds 3, 7 and 11). If you scale up_proj.weight by 4 between Model() and .to("cuda") (:435), a swap moves it by about 300%. Both swaps then fail all three seeds, and the unmutated test still passes.

    On tolerance: :466-472 use rtol=atol=1e-2. The 09-23 review measured exactly that tolerance failing on the MLP output on a card (1.1–3.6% of elements over 3 seeds). The forward math has not changed since (git diff ae25111a de508b7f touches no forward op). The CPU copy passing says nothing about CUDA's bf16 kernels. Please judge Y, dX and the adapter grads by Fast-LoRA autograd, correctness half: hand-written backward for the LoRA projections, behind training.fast_lora #792's rule instead: the kernel's error against a float64 reference on the same dequantised weights is at most 2x peft's. That rule also stays valid after you scale up the weights.

  4. The streamed test only checks agreement, so it passes when the kernel never runs. With _forward forced to delegate (fast_lora_mlp.py:211), test_mlp_kernel_matches_resident_peft_across_reused_stream_buffers passed in 6 runs out of 6. The tracker's streamed criterion exists to show the kernel is correct inside a streamed layer, so the test has to show that it ran there. Replacing the streamed_loss = ... line (:564) with this passes as shipped and fails under that mutation:

            seen = []
            hooks = [
                module.register_forward_hook(
                    lambda _m, _i, out: seen.append(type(out.grad_fn).__name__)
                )
                for module in streamed.modules()
                if getattr(module, "_soup_fast_lora_mlp", False)
            ]
            streamed_loss = streamed(input_ids=input_ids, labels=labels).loss
            for hook in hooks:
                hook.remove()
            assert seen == ["_FastLoraSwiGLUBackward"] * 4, seen

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

  • _module_uses_silu now fails closed, but nothing pins it: returning True for a missing act_fn (:242) leaves every test green.
  • The refactor still removes 19 comment lines of rationale from fast_lora.py and adds 3: the Conv1D / fan_in_fan_out reasoning, lora_bias and dropout. They belong on _projection_state.
  • From Fast-LoRA autograd: fused LoRA MLP block (gate/up/down around SwiGLU) with a hand-written backward #837's checklist: the float64 gradcheck still covers only the all-adapted case (the issue asks for each projection unadapted in turn), and bf16 parity and saturation run only on MPS. The saved-bytes item asks for the Llama-3.1-8B MLP shape, which does not need a card.

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 _projection_state (fast_lora.py:170), which the MLP and QKV kernels call too. If it is relaxed for the single path, the other two kernels silently ignore dropout. Deleting it fails one test in your file and one in #839's, but none in #1108's. Whichever of the two lands second should make dropout an explicit opt-in for the single path only.

CI: all 13 required checks on de508b7f are still queued (the run has been waiting since 2026-09-24 16:15Z), which is the state of the Actions queue today, not your change.

What I ran: the two test files above on CPU (CUDA_VISIBLE_DEVICES=-1), on the branch merged with main. Mutations: 7 of 9 on the non-GPU paths were killed. The two survivors were the single-path autocast revert (item 1) and the act_fn fail-open. On the CPU copy of the NF4 test, 1 of 2 swaps was killed, and both were killed after the fixture change. The kernel-never-runs mutation survived the streamed test and is killed by the hook above.

@Amix29

Amix29 commented Sep 25, 2026

Copy link
Copy Markdown
Contributor Author

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.

@MakazhanAlpamys

Copy link
Copy Markdown
Owner

@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 4eea730c merged with today's main (20ab4668), not by reading the reply.

What is closed

  • Single-projection autocast. Reverting fast_lora.py:402 to an unconditional out.to(input_dtype) now fails your new test. That test also passes on main, which is what a guard for a refactor should do.
  • Unadapted non-float sibling. Deleting the guard at fast_lora.py:197-198 fails test_unadapted_int8_sibling_refuses_the_whole_mlp, on a real bnb.nn.Linear8bitLt.
  • NF4 fixture. With up_proj scaled by 4, all four gate/up quant-state swaps now fail 3 of 3 seeds: forward and backward, both directions. I ran them on a CPU copy of the gpu test (bitsandbytes 0.50.2 runs NF4 on CPU).
  • Streamed test. Forcing the MLP _forward to delegate fails it now; the forward-hook record does its job.
  • _module_uses_silu. It fails closed, and that is pinned.
  • Mutations. 20 of 22 on the CPU suite fail at least one test. That includes "delegate unless all three are adapted" (32 failures), "delegate when ranks differ" (8), the SiLU-derivative factor (9) and the LoRA-side dX term (8). The two survivors are below.

What I ran: tests/test_issue837_fast_lora_mlp.py and tests/test_issue839_fast_lora_single_projection.py on CPU (CUDA_VISIBLE_DEVICES=-1). Result: 106 passed, 6 skipped (the MPS test and five gpu tests), which is your 107/5 without the Mac-only test. Ruff and git diff --check are clean. Your new file fails on main: 68 failed, 1 passed (the single-projection autocast guard, as intended), 4 skipped.

One blocking item: the bf16 backward fails your own 2x-vs-float64 rule on one of the three NF4 seeds

I copied test_nf4_forward_backward_and_all_adapter_grads_match_peft verbatim and changed only the device from cuda to cpu. Seeds 3 and 7 pass. Seed 11 fails, the same way on every run: the kernel's dX error is 3.61e-3 against peft's 7.33e-4, which is 4.93x. The CPU copy is not the cause. On CPU both sides use dequantise + GEMM, and changing only the kernel's elementwise SiLU derivative, with peft's reference untouched, brings seed 11 to 1.00x.

There are two independent causes:

  1. fast_lora_mlp.py:132-134 evaluates silu'(G) as a six-op bf16 chain, grad_m * u * sig * (1 + g * (1 - sig)). Autograd's silu_backward, which the unpatched MLP runs, does the same arithmetic in fp32 and rounds once. On 200k random bf16 inputs the chain's mean absolute error is 5.2e-3 against 2.6e-3, and its maximum is 0.315 against 0.137. Near the zero of silu' (G between -1.6 and -1.0) it is 4.6x worse, because 1 + g * (1 - sig) cancels in bf16. This is the cause of seed 11.
  2. :73, :78 and :88 cast the fp32 LoRA term to bf16 before the scaled add. peft promotes, adds and casts once, and so does feat(peft): fast-LoRA single-projection autograd kernel and its tests (#839) #1087's merged kernel (fast_lora.py:324-332). This one shows up without NF4. A pure-bf16 twin of your test (bf16 base, fp32 LoRA, which is Soup's layout since Fix fp16 GradScaler crash on pre-Ampere GPUs (cast bf16 LoRA adapters to fp32) #429) fails the same rule on 1 of 12 seeds: up_proj.lora_B at 2.39x.

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 gpu-only and there is no CPU bf16 parity test. So the CUDA run could fail for this reason rather than for #776's.

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:

  • Your suite is unchanged at 106 passed, 6 skipped, including the float64 gradcheck.
  • The NF4 copy passes 3 of 3 seeds.
  • The 20-seed NF4 sweep passes 19 of 20. The one left is seed 5 at 2.07x, and no variant I tried moved it, so I read it as noise at the rule's margin.
  • The bf16 twin passes 12 of 12; the worst ratio is 1.73x.

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.gpu with:

    @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 device in the two places that say "cuda" (model.to(...) and x_ref). As shipped, the cpu rows fail seed 11; with the fix they pass all three.

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 _deps, _randomise_b, copy and SimpleNamespace). It fails 1 of 12 seeds as shipped, and passes 12 of 12 with the forward fix, with or without the backward fix. So it pins cause (2) specifically:

@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-8

Non-blocking

  • The config.hidden_act gate is unpinned. Deleting it (fast_lora_mlp.py:262-266) leaves all 106 tests green. test_non_silu_activation_falls_back_without_patch also builds a GELU module, so the per-module check refuses first. Fast-LoRA autograd: fused LoRA MLP block (gate/up/down around SwiGLU) with a hand-written backward #837's checklist asks for "falls back with a printed reason". This test passes as shipped and fails with the gate deleted:

    def test_parent_config_hidden_act_gate_refuses_a_silu_module(caplog):
        _deps()
        from soup_cli.utils.fast_lora_mlp import patch_fast_lora_mlp
    
        model = _make_model(("gate_proj",), act="gelu", module_act="silu")
        with caplog.at_level("INFO", logger="soup_cli.utils.fast_lora_mlp"):
            assert patch_fast_lora_mlp(model) == 0
        assert "hidden_act='gelu'" in caplog.text
  • Deleting the allowlist check inside _projection_state (fast_lora.py:146) survives, but only because both patchers check the type first. It is defence in depth, and nothing is needed.

  • Still open from the tracker checklist, all CPU-sized:

    • a gradcheck with each projection unadapted in turn;
    • CPU bf16 parity (the blocking item above is what it finds);
    • saved bytes at the Llama-3.1-8B shape.
  • The PR description still says 104 passed / 5 skipped.

Heads-up on #839's dropout item. The shared dropout guard in _projection_state (fast_lora.py:170-173) is now pinned in all three files. Deleting it fails one test in your MLP file, one in #839's and one in #1108's QKV file. So a shared relaxation cannot slip through silently. Whichever of this and @sloemo01's PR lands second still needs to make dropout an explicit opt-in on the single-projection path only.

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 main moved. A push will re-run it. The PR merges cleanly with today's main.

@Amix29
Amix29 force-pushed the codex/fix-837-fast-lora-mlp branch from 4eea730 to 662a525 Compare September 27, 2026 11:07
@Amix29

Amix29 commented Sep 27, 2026

Copy link
Copy Markdown
Contributor Author

Implemented the blocking bf16 numerics round and the remaining CPU-sized pins at 662a5259, rebased onto current main.

  • shared promote-add-cast LoRA update now matches PEFT for mixed bf16 base / fp32 adapters;
  • SiLU backward uses fp32 opmath and rounds once;
  • NF4 parity now runs on CPU for seeds 3/7/11 while retaining marked CUDA rows;
  • added the 12-seed bf16-vs-float64 rule, parent hidden_act gate test, each-projection-unadapted gradchecks, and a meta-backed Llama-3.1-8B saved-bytes test.

Local M4 Max result: 127 passed, 5 skipped across #837/#839; Ruff, preview E30 and git diff --check clean. The PR description is refreshed. I am not claiming the remaining CUDA NF4 run.

@MakazhanAlpamys

Copy link
Copy Markdown
Owner

@Amix29, thank you, that was fast and it is exactly the fix. I re-checked it by running it on 662a5259 merged with today's main. git range-diff shows your four earlier commits unchanged, so this round is the one new commit.

What is closed

  • The bf16 numerics.
    • Your own NF4 test's new cpu rows pass 3 of 3 seeds. Seed 11 was 4.93x before.
    • Over 20 seeds of the same test, 19 pass. The one left is seed 5 at 2.07x, identical to what I measured with the fix this morning, and no variant moves it.
    • Each half is pinned. Putting the shared helper back to cast-first fails a test (the 12-seed bf16 rule). So does reverting only the gate add (2 failures) or only the up add. Putting the SiLU derivative back to the bf16 chain fails the NF4 cpu row for seed 11.
    • Moving feat(peft): fast-LoRA single-projection autograd kernel and its tests (#839) #1087's single-projection kernel onto _scaled_lora_add is a nice touch: one definition of "add the LoRA term" for all three kernels.
  • The hidden_act gate. Deleting it now fails your new test.
  • The tracker items. The per-projection-unadapted gradchecks pass, and the 8B saved-bytes test is in (one problem with it below).

What I ran: tests/test_issue837_fast_lora_mlp.py and tests/test_issue839_fast_lora_single_projection.py on CPU: 126 passed, 6 skipped (the MPS test, your three cuda rows and #839's two gpu tests), which is your 127/5 without the Mac-only test. Ruff and git diff --check are clean. On main the file gives 88 failed and 1 passed (the single-projection autocast guard, as intended). Of six mutations against the new code, five fail a test; the sixth is below.

Two small test-only fixes before merge

  1. Please add the aten_half_matmuls fixture to test_bf16_mlp_stays_within_twice_pefts_error_against_float64. This one is on me: the test I suggested runs a CPU bf16 F.linear at 10x48x32, which is above oneDNN's 16**3 cutoff. test(prm): run #1235's CPU bf16/fp16 steps on ATen's matmuls, fixing windows-latest 0xc000001d #1314 merged into main today because that exact kind of matmul kills the interpreter on part of GitHub's windows-latest fleet (0xc000001d), and the crash takes every later test in the cell with it. A green run proves nothing, because it depends on which host the job lands on. The fixture keeps the test running everywhere; tests/_windows_ci.py recommends it before any skip:

    def test_bf16_mlp_stays_within_twice_pefts_error_against_float64(seed, aten_half_matmuls):

    With it the test still passes 12 of 12, and putting the helper back to cast-first still fails it.

  2. The 8B saved-bytes test counts peft's X twice. On the meta device every data_ptr() is 0, so the test can't deduplicate storages, and peft's X, saved by both gate's and up's lora_A, is counted twice.

    Something like this:

    seen = set()
    
    def pack(tensor):
        nonlocal total
        # data_ptr() is 0 for every meta tensor; the storage object is what identifies it
        key = (
            tensor.untyped_storage()._cdata,
            tensor.storage_offset(),
            tuple(tensor.shape),
            tuple(tensor.stride()),
        )
        if key not in seen:
            seen.add(key)
            total += tensor.numel() * tensor.element_size()
        return tensor

    Then build with r=64, use x = torch.empty(1, 2048, 4096, ...), and assert against the prediction: n_f = 2048 * 14336 * 2 and assert 2 * n_f <= saved < 3 * n_f. I checked that _cdata is equal for the same tensor and its views and differs between distinct meta tensors.

Non-blocking

  • Reverting only the down projection's add to cast-first leaves every test green. One line after the 2x loop in the bf16 test would pin it. As shipped, your fused forward is bit-exact with peft on all 12 seeds; with that revert it is bit-exact on none. I used 1.05 rather than equality so an FMA difference on another CPU cannot make it flaky. It passes 12 of 12 as shipped and fails 5 of 12 with the revert:

        assert (fast[0] - oracle[0]).abs().max() <= 1.05 * (reference[0] - oracle[0]).abs().max() + 1e-8
  • The new NF4 cpu rows call bitsandbytes' native CPU dequantize_4bit, which picks an AVX-512 path at run time. If a windows-latest cell ever dies with 0xc000001d inside those rows, the fixture can't help a bitsandbytes kernel. The answer then is skip_on_windows_ci from tests/_windows_ci.py on the cpu param only; ubuntu and macOS keep the coverage.

  • Fast-LoRA autograd, correctness half: hand-written backward for the LoRA projections, behind training.fast_lora #792 asks that the measured ratio be reported, not just pass/fail, and the NF4 test prints nothing when it passes. A print of the six ratios would make the CUDA run self-documenting.

After the two fixes, this waits on the required checks (all 13 are still queued for 662a5259) and on the CUDA run of the NF4 tests from the 09-25 review. The PR merges cleanly with today's main.

@Amix29

Amix29 commented Sep 27, 2026

Copy link
Copy Markdown
Contributor Author

Addressed both test-only items and the useful non-blocking pins at 5b2167ce.

  • the 12-seed CPU bf16 test now requests aten_half_matmuls, avoiding the Windows oneDNN illegal-instruction class without dropping coverage;
  • the Llama-3.1-8B saved-byte test now uses N=2048/r=64, deduplicates meta storages by _cdata plus view identity, and pins the expected saving to two-to-three [N, f] tensors;
  • the bf16 rule now separately pins the down-projection add;
  • the NF4 parity test prints the measured fast/PEFT error ratios for the future CUDA acceptance run.

Local M4 Max result remains 127 passed, 5 expected CUDA/platform skips across #837/#839. Ruff, preview E30, and git diff --check pass. No CUDA result is claimed.

@MakazhanAlpamys

Copy link
Copy Markdown
Owner

@Amix29, thank you. This round is exactly what was asked, and nothing more: one commit on top of 662a5259, +32/-7 in the test file, no source change. You also took both optional notes. I checked each item by running it on 5b2167ce merged with today's main (df106781), not by reading the reply. When main moved to 0c95b96c I re-ran the suite, with the same result.

What is closed

  • The Windows fixture. test_bf16_mlp_stays_within_twice_pefts_error_against_float64 now requests aten_half_matmuls. It passes 12 of 12 with it. Putting _scaled_lora_add back to cast-first still fails it, now on 6 seeds (1 last round), so the fixture costs the test none of its teeth. The other CPU bf16 matmuls in the file are all below the 16**3 cutoff (the NF4 rows are 6x12x8), so this was the only test that needed it.
  • The 8B saved bytes. The test now runs at Fast-LoRA autograd: fused LoRA MLP block (gate/up/down around SwiGLU) with a hand-written backward #837's own shape and prints peft=611844096 fast=494141440 saved=117702656. That is exactly two [N, f] tensors (117,440,512 B) plus one [N, r] bf16 tensor (262,144 B), i.e. Fast-LoRA autograd: fused LoRA MLP block (gate/up/down around SwiGLU) with a hand-written backward #837's prediction, measured rather than assumed. I also made the kernel keep one more [N, f] tensor (m added to save_for_backward): your 8B test fails (58,982,400 B saved, under the 2 x [N, f] floor), and it is the only test that does.
  • The down-projection add. Last round, reverting only that add to cast-first left every test green. With your 1.05 line it now fails 5 of 12 seeds.
  • The ratio print. Under -rA, each NF4 cpu row prints its eight ratios. On this box the worst are 1.23 (seed 3, dX), 1.70 (seed 7, gate lora_B) and 1.00 (seed 11). The card run will now document itself.

What I ran: tests/test_issue837_fast_lora_mlp.py and tests/test_issue839_fast_lora_single_projection.py on CPU. Result: 126 passed, 6 skipped (the MPS test, your three cuda rows and #839's two gpu tests), which is your 127/5 without the Mac-only test. Ruff (including preview E30) and git diff --check are clean. Your test file on main gives 88 failed, 1 passed (the single-projection autocast guard, as intended) and 4 skipped. All three mutations above fail a test, and each was reverted afterwards.

Where this stands

From the review side this is ready to merge. Two things remain, and neither is on you:

  1. The 13 required checks. They have been queued on 5b2167ce since 12:14Z. The run for 662a5259 never started; it was cancelled by this push. So the first CI run on this head is also the first time any hosted runner executes the NF4 cpu rows, the 12-seed bf16 test and the new 1.05 line. If a windows-latest cell dies with 0xc000001d inside the NF4 cpu rows, the fix is the one from the last review: skip_on_windows_ci on the cpu param only. Those rows go through bitsandbytes' native CPU dequantise, which the fixture cannot reach.
  2. One run of the three NF4 cuda rows on a card, as the 09-25 review said. With -rA they will print their ratios.

#1108 is not affected. It is stacked on this exact head, and I checked that it lands cleanly on top of it.

@MakazhanAlpamys

Copy link
Copy Markdown
Owner

@Amix29, the CUDA run is done. RTX 5070 (sm_120, torch 2.14 cu130, bitsandbytes 0.50.2), on 5b2167ce merged with main:

pytest tests/test_issue837_fast_lora_mlp.py tests/test_issue839_fast_lora_single_projection.py -m gpu -rA gives 5 passed.

NF4 fast/PEFT error ratios (output, input gradient, then the six adapter grads):

  • seed 3: 1.0, 0.60, all six 1.0
  • seed 7: 1.0, 1.35, gate lora_B 0.80, the other five 1.0
  • seed 11: all eight 1.0

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 5b2167ce. They are still queued, because the Actions queue is long today. The squash message records these numbers.

@MakazhanAlpamys
MakazhanAlpamys merged commit 2ddaa8a into MakazhanAlpamys:main Sep 28, 2026
18 checks passed
@MakazhanAlpamys

Copy link
Copy Markdown
Owner

Merged as 2ddaa8a5. Thank you, @Amix29: when the bf16 backward failed the tracker's own 2x-vs-float64 rule on one NF4 seed (4.93x), you came back within the hour with the fix, one shared promote-add-cast helper for all three kernels, and a CPU test pinning each half. On the final head: 126 passed, 6 skipped on CPU, and the 5 CUDA tests pass on an RTX 5070 with a worst NF4 ratio of 1.35; your test file fails 88 on main. Reverting the shared add helper, the down-projection add, or saving one more [N, f] tensor fails 6, 5 and 1 tests.

MakazhanAlpamys added a commit that referenced this pull request Sep 28, 2026
…#1132); #1296 and #1295 (kok-o), #1107 (Amix29)

Display names from `gh api users/<login>`; sup3dev's is non-Latin, so the
login is used.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
MakazhanAlpamys pushed a commit that referenced this pull request Sep 29, 2026
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.
MakazhanAlpamys added a commit that referenced this pull request Oct 9, 2026
…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants