From 9780cdcca1badf72ef07d9e010ad27c4b2c328c1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 23 Jun 2026 01:29:19 -0700 Subject: [PATCH] Fix FlashAttention fp32 crash with DoRA (use_dora=True) (#6526) * Fix FlashAttention fp32 crash with DoRA (use_dora=True) DoRA upcasts lora_magnitude_vector to fp32 for the optimizer, which promotes the q/k/v_proj output to fp32. FlashAttention only accepts fp16/bf16, so the fp32 q/k/v raised 'FlashAttention only support fp16 and bf16 data type'. Downcast q/k/v to the compute dtype before the flash kernels. Fixes #1013 * Apply kwarg-spacing format hook to DoRA dtype test (pre-commit) * DoRA+FA2: downcast any fp32 among Q/K/V and clamp to a flash-supported dtype * Tighten code comments (no logic change) --------- Co-authored-by: danielhanchen --- .../test_attention_dispatch_dora_dtype.py | 73 +++++++++++++++++++ unsloth/utils/attention_dispatch.py | 19 +++++ 2 files changed, 92 insertions(+) create mode 100644 tests/utils/test_attention_dispatch_dora_dtype.py diff --git a/tests/utils/test_attention_dispatch_dora_dtype.py b/tests/utils/test_attention_dispatch_dora_dtype.py new file mode 100644 index 0000000000..a6586dea8d --- /dev/null +++ b/tests/utils/test_attention_dispatch_dora_dtype.py @@ -0,0 +1,73 @@ +"""Regression test for unslothai/unsloth#1013: run_attention must downcast DoRA's +fp32 q/k/v for FlashAttention while leaving already-16bit tensors untouched.""" + +import torch +import unsloth # noqa: F401 + +from unsloth.utils import attention_dispatch as ad + + +def _run(monkeypatch, qkv_dtype, backend): + captured = {} + + def _check(Q): + captured["dtype"] = Q.dtype + # Mirror the real kernel constraint so an unfixed dispatch fails loudly. + if Q.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError("FlashAttention only support fp16 and bf16 data type") + + def fake_flash_dense(Q, K, V, **kwargs): + _check(Q) + return torch.zeros_like(Q) + + def fake_flash_varlen(Q, K, V, *args, **kwargs): + _check(Q) + return torch.zeros_like(Q) + + monkeypatch.setattr(ad, "flash_attn_func", fake_flash_dense, raising = False) + monkeypatch.setattr(ad, "flash_attn_varlen_func", fake_flash_varlen, raising = False) + + bsz, n_heads, q_len, head_dim = 1, 2, 4, 8 + Q = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype) + K = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype) + V = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype) + + seq_info = None + if backend == ad.FLASH_VARLEN: + cu = torch.tensor([0, q_len], dtype = torch.int32) + seq_info = (None, cu, q_len) + + config = ad.AttentionConfig( + backend = backend, + n_kv_heads = n_heads, + n_groups = 1, + flash_dense_kwargs = {"causal": True}, + flash_varlen_kwargs = {"dropout_p": 0.0, "causal": True}, + ) + context = ad.AttentionContext( + bsz = bsz, + q_len = q_len, + kv_seq_len = q_len, + n_heads = n_heads, + head_dim = head_dim, + requires_grad = False, + seq_info = seq_info, + attention_mask = None, + causal_mask = None, + ) + ad.run_attention(config = config, context = context, Q = Q, K = K, V = V) + return captured["dtype"] + + +def test_dense_flash_downcasts_fp32_qkv(monkeypatch): + # fp32 DoRA output must be downcast to a flash-compatible dtype. + assert _run(monkeypatch, torch.float32, ad.FLASH_DENSE) in (torch.bfloat16, torch.float16) + + +def test_varlen_flash_downcasts_fp32_qkv(monkeypatch): + assert _run(monkeypatch, torch.float32, ad.FLASH_VARLEN) in (torch.bfloat16, torch.float16) + + +def test_bf16_qkv_left_untouched(monkeypatch): + # Standard LoRA path (already bf16) must not be altered. + assert _run(monkeypatch, torch.bfloat16, ad.FLASH_DENSE) == torch.bfloat16 diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 3ff581b7fc..2e984bad0a 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -137,6 +137,25 @@ def run_attention( requires_grad = context.requires_grad sliding_window = context.sliding_window + # DoRA promotes q/k/v_proj outputs to fp32, which FlashAttention rejects, so + # downcast any fp32 Q/K/V to a flash-supported dtype (#1013). + if backend in (FLASH_DENSE, FLASH_VARLEN) and torch.float32 in ( + Q.dtype, + K.dtype, + V.dtype, + ): + # Prefer the autocast dtype, else a non-fp32 input's dtype, then clamp. + if torch.is_autocast_enabled(): + try: + flash_dtype = torch.get_autocast_dtype("cuda") + except (AttributeError, TypeError): + flash_dtype = torch.get_autocast_gpu_dtype() + else: + flash_dtype = next((d for d in (Q.dtype, K.dtype, V.dtype) if d != torch.float32), None) + if flash_dtype not in (torch.float16, torch.bfloat16): + flash_dtype = torch.bfloat16 if SUPPORTS_BFLOAT16 else torch.float16 + Q, K, V = Q.to(flash_dtype), K.to(flash_dtype), V.to(flash_dtype) + if backend == FLASH_VARLEN: Q_f = Q.transpose(1, 2).reshape(bsz * q_len, n_heads, head_dim) K_f = K.transpose(1, 2).reshape(bsz * q_len, config.n_kv_heads, head_dim)