[kok-o]

Resolves #1163.

Summary

Follow-up to #839 / PR #1087 (by @sloemo01). TestTheFastPathIsActuallyTaken in tests/test_issue839_fast_lora_single_projection.py pinned that the fast-LoRA kernel actually runs over bias x rank in fp32, but missed non-fp32 dtypes and 4-bit base weights.

Neither agreement parity nor gradcheck can distinguish kernel execution from silent delegation. This PR widens the pin to assert on grad_fn across dtypes and on NF4.

What it does

1. Dtype Parameterization in TestTheFastPathIsActuallyTaken:

   • Parametrised test_the_output_carries_the_kernels_own_grad_fn over dtype (fp32, bf16).
   • Verified that inserting mutation if x.dtype.itemsize != 4: return original_forward(x) fails 4 tests (all bf16 variants) with AssertionError: the kernel did not run: grad_fn is AddBackward0.

2. NF4 grad_fn Assertion:

   • Added unpatched vacuity precondition assertion on ref.grad_fn: assert type(ref.grad_fn).__name__ != "_FastLoraSingleProjectionBackward".
   • Asserted type(out.grad_fn).__name__ == "_FastLoraSingleProjectionBackward" after patching.

3. Citation Cleanup:

   • Replaced closed issue citation #968 with root-cause issue #776.

Note on base

This branch is built on top of PR #1087 (1d62c347). Once #1087 merges or is lifted into #837/#838, this commit can be integrated cleanly.

Verification

• pytest --no-cov tests/test_issue839_fast_lora_single_projection.py: 29 passed, 2 skipped (gpu) in 14.0s
• ruff check tests/test_issue839_fast_lora_single_projection.py: 100% clean

[MakazhanAlpamys]

Review of 5b6b180b.

What holds, verified by execution. Your commit is exactly what #1163 asked for, and both of the issue's mutations now die. I ran the file on a card so that both gpu tests execute, on main plus your change:

mutation at fast_lora.py:278	on main	with this PR (card)	with this PR (CPU, what CI sees)
none	27 passed	31 passed	29 passed, 2 skipped
if x.dtype.itemsize != 4: return original_forward(x)	27 passed (survives)	5 failed (4 bf16 arms + NF4)	4 failed (the bf16 arms)
delegate when the base weight has a quant_state	27 passed (survives)	1 failed: :861, "grad_fn is AddBackward0"	29 passed, 2 skipped

The NF4 vacuity precondition is in place (:844), both #968 citations now point at #776, nothing was deleted or loosened (the 8 original fp32 IDs still collect), no new skip, and ruff is clean. A test-only change needs no fragment. Your body reports only the dtype mutation. The NF4 one needs a card, so the second row above is my measurement.

Blocking

1. The branch conflicts with main. It still carries #1087's four pre-squash commits, and main has the squash (7b78e528), so the test file is an add/add conflict. GitHub reports CONFLICTING and merge-freshness fails. The 13 required checks are green, but they ran on a merge ref from before #1087 landed. Fix: git rebase --onto origin/main 1d62c347. I simulated it: it applies cleanly as 1 file, +18/-13, and it is identical to the tree I tested. CI then runs on the new head.

2. The dtype arm never builds the dtype layout training actually uses, so a narrowing of the same shape still survives, even on a card. get_peft_model defaults to autocast_adapter_dtype=True, which leaves fp32 adapters on a bf16 base. I measured that on a dense bf16 base and on an NF4 base through get_peft_model: lora_A comes back float32 on both. So every real bf16 LoRA and QLoRA run takes the kernel's mixed-dtype branch. The bf16 arm casts the whole wrapper with model.to(bf16) (:976), and the NF4 fixture uses inject_adapter_in_model, which gives bf16 adapters. Both are uniform-dtype. Results with your file:

   • if lora_a.dtype != x.dtype: return original_forward(x) gives 31 passed on the card.
   • if x.dtype == torch.float16: return original_forward(x) also gives 31 passed. fp16 is the pre-Ampere training dtype.

   Without the mutation, the kernel does run in both layouts (_FastLoraSingleProjectionBackward on CPU, and for mixed on NF4 too). So the pin can stay on CPU:

   @pytest.mark.parametrize("dtype", ["fp32", "bf16", "fp16", "mixed"], ids=["fp32", "bf16", "fp16", "mixed"])
   ...
   torch_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16,
                  "fp16": torch.float16, "mixed": torch.bfloat16}[dtype]
   model, layer = _make_adapted_linear(bias=bias)
   model.to(torch_dtype)
   if dtype == "mixed":  # what get_peft_model leaves by default: fp32 adapters on a bf16 base
       for name, param in model.named_parameters():
           if "lora_" in name:
               param.data = param.data.float()
   x = torch.randn(*shape, dtype=torch_dtype, requires_grad=True)

   I ran exactly this on CPU: 37 passed and 2 skipped with no mutation. The mixed-dtype mutation fails 4 tests (the mixed arms) and the fp16 one fails 4 (the fp16 arms). I know this goes past the issue's literal text. It is still the same pattern the issue describes, one round later, and mixed dtype is the configuration this kernel exists for.

Non-blocking

• The NF4 kill only happens on a card. TestNf4Parity is correctly gpu-marked, so under the quant-state mutation every CI cell stays green (29 passed, 2 skipped). A CPU twin is possible: bitsandbytes 0.50.2 quantizes NF4 on CPU, and the patched forward runs there (AddBackward0 unpatched, the kernel's node patched, backward fine). CI's test matrix installs .[dev], which pulls bitsandbytes. Guard it with importorskip, plus a skip when CPU quantization raises on an older bitsandbytes. Optional, since the issue put the assertion in the gpu test, but it would let CI catch that mutation.
• Please keep the provenance paragraph. The docstring paragraph you replaced (old :957-965) was the file's only record of the three survivors at :278 and of why the o_proj shape (3-D, no bias) matters. Keep it, and add the rows this PR adds, with their counts.

Nits

• The module docstring's gpu bullet for NF4 still says "parity" only; it now also pins that the kernel ran. The "Asserted on CPU" list doesn't mention the dtype arms.
• :975-977: nn.Module.to works in place, so both reassignments are no-ops.

[sloemo01]

Confirming blocking item 1, and pinning the exact rebase so it is a copy-paste rather than a re-derivation.

The conflict is mine. #1087 merged as a squash (7b78e52), while this branch is built on the four pre-squash commits (295882d → ea55686 → 89031be → 1d62c34), so both sides add tests/test_issue839_fast_lora_single_projection.py and git cannot choose. Measured on 5b6b180, mergeable: false, mergeable_state: dirty.

I ran the rebase you proposed, on a scratch branch, with nothing pushed:

git fetch upstream
git rebase --onto upstream/main 1d62c347

It applies cleanly as one commit and the result is byte-identical to what you tested: 1 file changed, 18 insertions(+), 13 deletions(-), the diff against main being only this branch's change to the test file. So CI will re-run on a head that carries the squashed kernel rather than the four-commit chain.

One correction to the issue body while you are here: 5b6b180 does not conflict with main in the sense of a textual clash inside the file. It is an add/add on the whole path, which is why merge-freshness fails with CONFLICTING rather than a marker you could resolve by hand. The rebase removes the second side entirely, and no manual resolution is needed.

For what it is worth, the failure mode is the one the kernel is about: a green suite that cannot see the difference. The 13 required checks were green on a merge ref from before #1087 landed, so they agreed with the old tree. Same shape as the delegation problem, one level up.

Blocking item 2 is yours to decide, and I have no stake in it either way. Mixed dtype being the configuration get_peft_model actually leaves (autocast_adapter_dtype=True giving fp32 adapters on a bf16 base) matches what I saw while building #1087's fixtures, so I read that pin as the more valuable half of this PR, whatever #1163's text said literally.

[kok-o]

Thank you @MakazhanAlpamys and @sloemo01!

The branch is cleanly rebased onto upstream/main (aa39479e directly on 0235759e), resolving the add/add merge-freshness conflict with #1087:

1. Rebase & Conflict Resolution:
   • Cleanly rebased on top of main after feat(peft): fast-LoRA single-projection autograd kernel and its tests (#839) #1087 squashed, carrying only the 1 commit with +36/-7 on tests/test_issue839_fast_lora_single_projection.py.
2. Dtype Parameterization (including mixed dtype):
   • Parametrised test_the_output_carries_the_kernels_own_grad_fn over ["fp32", "bf16", "fp16", "mixed"].
   • The mixed arm casts the base to bfloat16 and keeps lora_* parameters in float32, matching the real-world layout left by get_peft_model(..., autocast_adapter_dtype=True).
3. Docstrings & Provenance:
   • Preserved and expanded the provenance paragraphs and module docstrings, detailing both CPU dtype arms and NF4 kernel execution assertions.
   • Cleaned up citations from Pin the bitsandbytes fused-vs-dequant NF4 divergence that #857 rests on #968 to root-cause Streamed NF4 is not bit-exact vs resident NF4 on Blackwell (sm_120): 4.9e-04 / 3.9e-03, while quantization=none stays exact #776.
4. Verification:
   • All 37 CPU tests pass (2 GPU tests skipped) in 7.8s.
   • ruff check 100% clean.

Ready for merge!

[MakazhanAlpamys]

Review of aa39479e.

Thank you for the rebase and for taking the mixed/fp16 arms as written. Both blocking items from my last review are closed, and I checked each by execution.

1. The conflict is gone. The branch is now a single commit on main. It merges cleanly onto today's main (1c053a8c), and the only file that differs is this test file. GitHub now reports it clean.

2. The narrowings that survived last round now die. I ran the file on the merged tree, both CPU-only (what CI sees) and with the card visible, so both gpu tests execute. Each mutation is inserted at fast_lora.py:278:

mutation	main's file (card)	this PR, CPU	this PR, card
none	27 passed	37 passed, 2 skipped	39 passed
if x.dtype.itemsize != 4: return original_forward(x)	27 passed	12 failed	13 failed
delegate when the base weight has a quant_state	27 passed	37 passed, 2 skipped	1 failed (NF4 grad_fn)
if lora_a.dtype != x.dtype: return original_forward(x)	27 passed	4 failed (the mixed arms)	4 failed
delegate on fp16	27 passed	4 failed (the fp16 arms)	4 failed
delegate on bf16	—	8 failed	9 failed

The mixed-dtype and fp16 rows are the two that gave 31/31 at 5b6b180b.

Other checks:

• Nothing was deleted or loosened. Collection goes from 27 to 39, and the four original IDs survive as the fp32 arm. There is no new skip.
• Ruff is clean, and the file is format-clean.
• The provenance paragraph is back, and both docstring nits are fixed.
• This is test-only, so no fragment is needed.
• CI is 13/13 required checks green on this head. That merge ref predates three later merges to main, but none of them touch fast_lora.py or this file. My local merge gives the numbers above.

This does not depend on the two SwiGLU/QKV drafts. When this lands, their copy of the #1163 hunk has to go; I'll take that up on those PRs.

Non-blocking

• The quant-state row still dies only on a card, so every CI cell stays green under it. A CPU twin (bitsandbytes 0.50.2 quantizes NF4 on CPU) would let CI see it. That's fine as a follow-up; the issue put the assertion in the gpu test, and you did exactly that.
• The sentence you appended to the provenance paragraph has no counts, while the paragraph above it records counts for the earlier rounds. Something like: "itemsize: 12 CPU / 13 card; quant_state: 1, card only; mixed: 4; fp16: 4".

Nit

• The PR body still describes the previous head: fp32/bf16 only, "built on top of PR #1087 (1d62c347)", and 29 passed. Please refresh it before the squash.

Approving.

[MakazhanAlpamys]

Merged as 92bd2207. Thank you, @kok-o.

Both mutations that survived my last round (mixed dtype and fp16) now fail four tests each, on CPU, which is where CI can see them. All six mutations I ran were killed with the card visible, and the file merged cleanly onto current main. #1107 and #1108 will drop their copy of the #1163 hunk now that this has landed.