xFormers 教程：用 Packed Sequences、GQA、ALiBi、SwiGLU 构建内存高效 Transformer

一句话看懂

xFormers 是一个 GPU 上的 Transformer 加速工具包，通过内存高效注意力、GQA、ALiBi 等技术，显著降低显存占用并提升训练速度。

详细发生了什么

xFormers 是 Meta 开源的 Transformer 优化库，核心是 memory_efficient_attention 算子。它避免了标准注意力中 [B, H, M, M] 分数矩阵的显式存储，将显存占用从 O(M²) 降为 O(M)。

教程首先验证了 xFormers 注意力与标准注意力的数值一致性（最大绝对差在 fp16 精度范围内）。然后对比了不同序列长度（512 到 4096）下的显存和速度：标准注意力显存随序列长度平方增长，而 xFormers 近似线性增长，且速度更快。

接着演示了因果掩码（LowerTriangularMask）、变长序列打包（BlockDiagonalMask）和分组查询注意力（GQA）。GQA 让多个查询头共享更少的键值头（如 8 个查询头对应 2 个 KV 头），大幅减小 KV-cache，这正是 Llama/Mistral 等模型使用的技术。

最后，教程展示了如何添加自定义 ALiBi 位置偏置，并将所有技术组合成一个可训练的 GPT 风格模型，包含 SwiGLU 前馈层和自动混合精度训练。

中文圈视角

xFormers 对国内开发者非常实用，尤其是需要训练或微调大模型的场景。虽然 xFormers 是 Meta 的项目，但完全开源，无需梯子即可从 GitHub 或 PyPI 安装。

国内同类工具有 DeepSeek 的 FlashMLA（针对 MoE 优化）和华为昇腾的 MindSpore 注意力优化，但 xFormers 生态更成熟，与 Hugging Face Transformers 集成良好。对于使用 PyTorch 的中文开发者，xFormers 是最直接的内存优化方案。

具体场景：

• 长文本训练：处理 8K+ 上下文时，xFormers 可节省 50% 以上显存。
• 推理加速：GQA 和 packed sequences 能降低 KV-cache 大小，适合部署多轮对话模型。
• 国产 GPU 适配：xFormers 主要针对 NVIDIA GPU，国产 GPU（如昇腾、寒武纪）需关注各自厂商的优化库，但思路可借鉴。

一个中文圈尚未广泛讨论的盲点：xFormers 的 BlockDiagonalMask 可以高效处理变长序列，这对中文 NLP 中常见的文档级任务（如长文档摘要、多轮对话）非常有用，但很多教程仍在使用 padding 方式，浪费显存。

几条值得记住的细节

• xFormers 的 memory_efficient_attention 在 fp16 下与标准注意力数值误差小于 1e-5。
• 序列长度从 512 到 4096，标准注意力显存从约 100 MB 飙升至 1600 MB，而 xFormers 仅从 80 MB 增至 300 MB。
• GQA 支持 5-D 张量布局 [B, M, G, Hq, K]，其中 G 是 KV 头数，Hq 是每个 KV 头对应的查询头数。
• 变长序列打包使用 BlockDiagonalMask.from_seqlens，无需填充，零显存浪费。
• 教程最终模型包含 SwiGLU 和 AMP，可在单 GPU 上训练 1B 参数级别的模型。

一句话总结

如果你用 PyTorch 训练 Transformer，xFormers 是降低显存、加速训练的最简单方案，尤其适合长序列和 GQA 场景。