Remove the unreachable torch<2 attention path in the Kolors text encoder - #14620
Remove the unreachable torch<2 attention path in the Kolors text encoder#14620RudraMantri123 wants to merge 2 commits into
Conversation
CoreAttention computes raw attention scores with baddbmm(input=torch.empty(...), beta=0), relying on beta=0 causing input to be ignored. MPS does not honour that contract (pytorch#187521): NaN in the uninitialised buffer propagates into the output. Verified directly at the shapes this code requests -- a freed NaN block of (b*np, sq, sk) is handed back by torch.empty and survives baddbmm. This is the last remaining instance of the idiom in the repository after huggingface#14459 fixed both copies of Attention.get_attention_scores. Use the same buffer-free scaled bmm on MPS, which relies on no contract and skips the scores-sized allocation; all other devices keep the existing baddbmm path unchanged. Adds a CPU-equivalence test that runs on every backend and an MPS-gated test that dirties the allocator before computing scores.
|
Hi @RudraMantri123, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice. Once the PR links an issue (or gets the |
|
Linked: this now closes #14624, which documents the underlying hazard (unsafe |
CoreAttention.forward gates on int(torch.__version__.split(".")[0]) >= 2 and
diffusers requires torch >= 2.6, so the manual attention branch -- the one
carrying the MPS-unsafe baddbmm(torch.empty(...), beta=0) idiom -- cannot
execute on any supported install. Every forward goes through
scaled_dot_product_attention.
That makes the previous commit's approach wrong twice over: the bug it
guarded against is unreachable, and the bmm branch it added was itself dead
code that nothing could execute. Delete the whole torch<2 path instead,
which removes the unsafe idiom truthfully.
Verified the deletion changes nothing observable: outputs are bit-identical
to main across cpu/mps, fp32/fp16, and both mask branches. The test file now
pins CPU-equivalence of the remaining SDPA path (masked and causal) on every
backend; the dirty-allocator test is dropped because it exercised SDPA all
along and could never fail.
RudraMantri123
changed the title
1 participant
Fixes #14624.
What does this PR do?
Removes
CoreAttention's manual attention branch inpipelines/kolors/text_encoder.py— dead code on every supported torch, and the last place in the repository still carrying thebaddbmm(input=torch.empty(...), beta=0)idiom that MPS executes incorrectly (pytorch/pytorch#187521).Correction from the original version of this PR
The first commit here added an MPS-specific
bmmbranch next to thebaddbmmcall, mirroring #14459. That was the wrong fix, and I want to be explicit about why, since I only saw it on re-verification:CoreAttention.forwardgates onint(torch.__version__.split(".")[0]) >= 2and routes everything throughscaled_dot_product_attentionon torch ≥ 2. diffusers requires torch ≥ 2.6, so the manual branch — the one containing the unsafebaddbmmidiom — is unreachable on any supported install. This is also why I could never reproduce end-to-end corruption in this pipeline: the code never runs. The branch my first commit added was itself dead code, and its "dirty allocator" test exercised SDPA all along and could never fail.So the honest change is the second commit: delete the whole torch < 2 path.
What is verified
(16, 1024, 1024)and(16, 2048, 2048)fp16 — a freed NaN block returned bytorch.emptysurvivesbaddbmm(beta=0)on MPS (torch 2.13.0, M-series). Repro in Kolors text encoder carries an unreachable torch<2 attention path built on the MPS-unsafe baddbmm(empty, beta=0) idiom #14624. That is what made the branch worth removing rather than leaving as a trap for anyone who copies it or re-enables it.main(torch.equal, max diff 0.0) across cpu/mps × fp32/fp16 × both mask branches offorward.make quality,check_copies,check_ai, andcheck_forward_call_docstringspass.Tests
tests/pipelines/kolors/test_kolors_text_encoder.pypins CPU-equivalence of the remaining SDPA path — causal (mask=None) and explicit-mask branches — and runs on every backend with no weights or Hub access (CoreAttentionis built directly fromChatGLMConfig).Notes for reviewers
AI disclosure per the contribution policy: AI-assisted debugging and drafting; all experiments were run and verified by me on real hardware.
Who can review?
@yiyixuxu @asomoza