标准nn.transformer在长序列上oom或变慢，因其multiheadattention默认使用o(n²)全连接注意力；解决方案包括：①用flash_attn替换（需fp16/bf16、a100/h100）；②改用longformer/linformer等稀疏结构；③分块+重叠缓存。
为什么标准 nn.Transformer 在长序列上会 OOM 或变慢？
因为原生 nn.Transformer 的 nn.MultiheadAttention 默认使用全连接注意力（full attention），计算和存储复杂度都是 O(n²)。当序列长度超过 1024，GPU 显存很快耗尽，训练步数也急剧拉长。
- 典型报错：
RuntimeError: CUDA out of memory或显存占用飙升至 95%+ - 即使能跑，
input_ids长度从 512 增到 2048，单步训练时间可能翻 4 倍以上 -
nn.TransformerEncoderLayer中的src_mask和attn_mask若未正确对齐 shape，会触发RuntimeError: The size of tensor a (2048) must match the size of tensor b (1024)
用 flash_attn 替换默认注意力层
这是目前最直接有效的提速+降显存方案，尤其适合 A100/H100 或支持 FP16/BF16 的卡。它把 self-attention 优化成近似线性复杂度，并原生支持 causal mask 和 packed sequence。
- 安装：
pip install flash-attn --no-build-isolation（注意 CUDA 版本匹配） - 替换时不要动
nn.TransformerEncoder结构，只重写forward中的 attention 调用，或继承nn.MultiheadAttention实现FlashMHA - 必须确保输入是
torch.float16或torch.bfloat16，否则flash_attn会退回到慢路径 - 示例关键片段：
from flash_attn import flash_attn_qkvpacked_func<br>qkv = torch.stack([q, k, v], dim=2) # [b, s, 3, h, d]<br>out = flash_attn_qkvpacked_func(qkv, dropout_p=0.0, causal=True)
改用 Longformer 或 Linformer 的稀疏/低秩注意力
如果不能换显卡或无法编译 flash_attn，就该考虑模型结构层面的降维。PyTorch 生态里，transformers 库已封装好可直接加载的预训练长文本模型。
-
LongformerModel使用滑动窗口 + 全局 token，显存占用接近O(n)；但要注意global_attention_mask必须与input_ids同 shape，且全局位置索引不能越界 -
LinformerModel将 key/value 投影到低维（如 256 维），适合超长文档分类，但会损失局部细节建模能力 - 避免直接用
AutoModel.from_pretrained("allenai/longformer-base-4096")然后塞进自己写的nn.Transformer—— 二者 position embedding 和 layer norm 位置不兼容，容易导致梯度爆炸
分块处理（Chunking）+ 重叠缓存（Recurrent Cache）
对纯自定义 Transformer（比如你手写了一个 CustomTransformerEncoder），最稳妥的 fallback 方案是手动切分序列，再用 cache 复用前序 hidden state。
- 切块大小建议设为 512 或 1024，重叠长度取 64～128（用于缓解边界信息丢失）
- 每次 forward 传入
past_key_values（shape 为[num_layers, 2, b, num_heads, cache_len, head_dim]），注意cache_len要随块推进动态增长 - 别忘了在 loss 计算时 mask 掉重叠部分的 label，否则会重复监督同一 token
- 这个方案 CPU/GPU 通信开销略高，但完全规避了
O(n²)注意力，且兼容所有硬件
实际跑起来你会发现：flash_attn 是“一劳永逸”的最优解，但依赖环境；结构替换（如 Longformer）省心但迁移成本高；而 chunking + cache 最糙却最稳——尤其当你连 torch.compile 都不敢开的时候。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
