Llama-Factory 是否支持 FlashAttention 加速
在大模型训练日益普及的今天，一个关键问题始终困扰着开发者：如何在有限的硬件资源下，更快、更稳地完成微调任务？尤其是当处理长文本或高分辨率上下文时，显存溢出、训练缓慢成了家常便饭。这时候，大家自然会问——有没有什么'加速外挂'可以一键开启？
FlashAttention 就是这样一个被广泛寄予厚望的技术。它号称能让注意力计算快上两倍、显存占用直降一个数量级，还不改变模型精度。那么问题来了：像 Llama-Factory 这种主打'开箱即用'的微调框架，能不能顺利接上这个利器？
答案是：能，而且用起来比你想象中更自然。
要理解这一点，得先搞清楚 FlashAttention 到底做了什么。
传统的自注意力机制虽然数学优雅，但实现起来效率堪忧。以 QKᵀ 计算为例，中间结果（比如注意力权重矩阵）必须写入 GPU 的高带宽显存（HBM），后续再读取用于 Softmax 和乘 V 操作。这一来一回，IO 开销巨大，尤其在序列长度超过 1024 后，速度瓶颈和显存压力陡增。
而 FlashAttention 的核心思路非常直接：把整个注意力计算塞进一个 CUDA kernel 里，在 SRAM 中完成所有中间运算，只把最终输出刷回 HBM。这种'算子融合 + 分块处理'（tiling）的方式，让显存访问次数从 $ O(n^2) $ 降到接近 $ O(1) $，实际显存消耗也从 $ O(n^2) $ 趋近于 $ O(n) $。更重要的是，它输出的结果与标准 attention 完全一致——没有近似、没有舍入误差，纯纯的'免费性能提升'。
# 使用 flash-attn 库的典型调用方式
import torch
from flash_attn import flash_attn_func
q, k, v = ... # shape: (batch, seqlen, nheads, headdim), 已转置
out = flash_attn_func(q, k, v)
这段代码看似简单，背后却是对 GPU 内存层级结构的极致利用。只要你的设备是 NVIDIA Ampere 架构及以上（如 A100、RTX 3090/4090），配合 PyTorch 2.0+ 和正确版本的 flash-attn，就能直接享受加速红利。
那 Llama-Factory 呢？它本身并不重新发明轮子，而是站在 Hugging Face Transformers 和 PEFT 的肩膀上构建生态。它的价值不在于从零写模型，而在于把复杂的训练流程封装成可配置、可视化的流水线。用户只需填几个参数，点一下按钮，就能启动 LoRA、QLoRA 或全参微调。
但这是否意味着它无法触达底层优化？恰恰相反。
Llama-Factory 的模型加载阶段实际上是动态注入的过程。当你指定 model_name_or_path 和 finetuning_type: lora 时，框架会通过 Transformers 加载基础模型，并借助 PEFT 插入适配器模块。在这个过程中，如果检测到 flash-attn 可用，很多现代 LLM 实现（如 LLaMA、Qwen、Mistral）都会自动启用 FlashAttention 替代原生 SDPA（Scaled Dot Product Attention）。
换句话说，只要满足以下条件，加速就会悄然生效：
- 安装了兼容版本的
flash-attn（推荐使用pip install flash-attn --no-build-isolation，注意编译依赖） - GPU 支持（NVIDIA 显卡，CUDA ≥ 11.8）
- 模型架构为标准 Transformer 风格（非 GLM、ChatGLM 等特殊结构）
- 使用 fp16 或 bf16 精度训练（FlashAttention 对 fp32 支持较弱）
我们来看一个典型的 LoRA 配置片段：
model_name_or_path: meta-llama/Llama-2-7b-chat-hf
finetuning_type: lora
lora_target:
