xFormers是由Meta开发的一个高性能深度学习库，专门用于优化Transformer架构中的注意力机制和其他组件。它提供了内存高效和计算高效的实现，特别适用于处理长序列和大规模模型。
github地址： xFormers

主要特性

• 内存高效注意力：xFormers的核心功能是提供内存高效的注意力机制实现，可以显著减少GPU内存使用，同时保持计算精度。
• 多种注意力变体：支持标准注意力、Flash Attention、Block-wise attention等多种优化版本。
• 自动优化：根据输入的形状和硬件特性自动选择最优的注意力实现。
• PyTorch集成：与PyTorch深度集成，可以作为drop-in replacement使用。

安装方式

# 要求：torch>=2.7
# 通过pip安装
pip install xformers# 或者从源码安装以获得最新功能
pip install git+https://github.com/facebookresearch/xformers.git

主要优势

内存效率：相比标准注意力机制，xFormers可以节省20-40%的GPU内存，特别是在处理长序列时效果显著。
计算效率：通过优化的CUDA kernel实现，提供更快的计算速度。
易于集成：可以作为现有PyTorch模型的直接替换，无需修改模型架构。
自动优化：根据硬件和输入自动选择最优的实现策略。

使用场景

长序列处理：处理文档级别的文本或长视频序列
大规模语言模型：GPT、BERT等Transformer模型的训练和推理
计算机视觉：Vision Transformer (ViT)等视觉模型
多模态模型：结合文本和图像的大规模模型

注意事项

硬件要求：需要较新的NVIDIA GPU（建议RTX 20系列或更新）
精度：某些情况下可能有轻微的数值差异，但通常可以忽略
调试：由于使用了优化的CUDA kernel，调试可能比标准PyTorch操作稍复杂

代码示例

import torch
import torch.nn as nn
from xformers import ops as xops
import math# 示例1：基础内存高效注意力
def basic_memory_efficient_attention():"""基础的内存高效注意力示例"""batch_size, seq_len, embed_dim = 2, 1024, 512# 创建输入张量query = torch.randn(batch_size, seq_len, embed_dim, device='cuda', dtype=torch.float16)key = torch.randn(batch_size, seq_len, embed_dim, device='cuda', dtype=torch.float16)value = torch.randn(batch_size, seq_len, embed_dim, device='cuda', dtype=torch.float16)# 使用xFormers的内存高效注意力scale = 1.0 / math.sqrt(embed_dim)output = xops.memory_efficient_attention(query, key, value, scale=scale)print(f"Input shape: {query.shape}")print(f"Output shape: {output.shape}")return output# 示例2：多头注意力实现
class MemoryEfficientMultiHeadAttention(nn.Module):def __init__(self, embed_dim, num_heads, dropout=0.0):super().__init__()self.embed_dim = embed_dimself.num_heads = num_headsself.head_dim = embed_dim // num_headsself.scale = 1.0 / math.sqrt(self.head_dim)self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3, bias=False)self.out_proj = nn.Linear(embed_dim, embed_dim)self.dropout = dropoutdef forward(self, x, attn_mask=None):batch_size, seq_len, embed_dim = x.shape# 计算Q, K, Vqkv = self.qkv_proj(x)qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, batch, heads, seq, head_dim]q, k, v = qkv[0], qkv[1], qkv[2]# 重塑为xFormers期望的格式 [batch*heads, seq, head_dim]q = q.reshape(batch_size * self.num_heads, seq_len, self.head_dim)k = k.reshape(batch_size * self.num_heads, seq_len, self.head_dim)v = v.reshape(batch_size * self.num_heads, seq_len, self.head_dim)# 使用内存高效注意力out = xops.memory_efficient_attention(q, k, v, attn_bias=attn_mask,scale=self.scale,p=self.dropout if self.training else 0.0)# 重塑回原始格式out = out.reshape(batch_size, self.num_heads, seq_len, self.head_dim)out = out.transpose(1, 2).reshape(batch_size, seq_len, embed_dim)return self.out_proj(out)# 示例3：带有因果掩码的注意力
def causal_attention_example():"""带有因果掩码的注意力示例（用于decoder）"""batch_size, seq_len, embed_dim = 2, 512, 256query = torch.randn(batch_size, seq_len, embed_dim, device='cuda')key = torch.randn(batch_size, seq_len, embed_dim, device='cuda')value = torch.randn(batch_size, seq_len, embed_dim, device='cuda')# 创建因果掩码（下三角矩阵）causal_mask = torch.tril(torch.ones(seq_len, seq_len, device='cuda'))causal_mask = causal_mask.masked_fill(causal_mask == 0, float('-inf'))causal_mask = causal_mask.masked_fill(causal_mask == 1, 0.0)# 使用带掩码的注意力output = xops.memory_efficient_attention(query, key, value,attn_bias=causal_mask,scale=1.0 / math.sqrt(embed_dim))return output# 示例4：完整的Transformer块
class MemoryEfficientTransformerBlock(nn.Module):def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):super().__init__()self.attention = MemoryEfficientMultiHeadAttention(embed_dim, num_heads, dropout)self.norm1 = nn.LayerNorm(embed_dim)self.norm2 = nn.LayerNorm(embed_dim)# Feed Forward Networkself.ff = nn.Sequential(nn.Linear(embed_dim, ff_dim),nn.GELU(),nn.Dropout(dropout),nn.Linear(ff_dim, embed_dim),nn.Dropout(dropout))def forward(self, x, attn_mask=None):# 注意力 + 残差连接attn_out = self.attention(self.norm1(x), attn_mask)x = x + attn_out# FFN + 残差连接ff_out = self.ff(self.norm2(x))x = x + ff_outreturn x# 示例5：性能对比
def performance_comparison():"""对比标准注意力和内存高效注意力的性能"""batch_size, seq_len, embed_dim = 4, 2048, 768query = torch.randn(batch_size, seq_len, embed_dim, device='cuda', dtype=torch.float16)key = torch.randn(batch_size, seq_len, embed_dim, device='cuda', dtype=torch.float16)value = torch.randn(batch_size, seq_len, embed_dim, device='cuda', dtype=torch.float16)scale = 1.0 / math.sqrt(embed_dim)# 标准注意力实现def standard_attention(q, k, v, scale):scores = torch.matmul(q, k.transpose(-2, -1)) * scaleattn_weights = torch.softmax(scores, dim=-1)return torch.matmul(attn_weights, v)# 测量内存使用（需要在实际环境中运行）print("使用xFormers内存高效注意力...")torch.cuda.reset_peak_memory_stats()xformers_output = xops.memory_efficient_attention(query, key, value, scale=scale)xformers_memory = torch.cuda.max_memory_allocated() / 1024**2  # MBprint("使用标准注意力...")torch.cuda.reset_peak_memory_stats()standard_output = standard_attention(query, key, value, scale)standard_memory = torch.cuda.max_memory_allocated() / 1024**2  # MBprint(f"xFormers峰值内存使用: {xformers_memory:.2f} MB")print(f"标准注意力峰值内存使用: {standard_memory:.2f} MB")print(f"内存节省: {((standard_memory - xformers_memory) / standard_memory * 100):.1f}%")# 示例6：在实际模型中使用
class GPTWithXFormers(nn.Module):def __init__(self, vocab_size, embed_dim, num_heads, num_layers, max_seq_len):super().__init__()self.embed_dim = embed_dimself.token_embedding = nn.Embedding(vocab_size, embed_dim)self.pos_embedding = nn.Embedding(max_seq_len, embed_dim)self.blocks = nn.ModuleList([MemoryEfficientTransformerBlock(embed_dim, num_heads, embed_dim * 4)for _ in range(num_layers)])self.ln_f = nn.LayerNorm(embed_dim)self.head = nn.Linear(embed_dim, vocab_size, bias=False)def forward(self, input_ids):seq_len = input_ids.size(1)pos_ids = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0)# 嵌入x = self.token_embedding(input_ids) + self.pos_embedding(pos_ids)# 创建因果掩码causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device))causal_mask = causal_mask.masked_fill(causal_mask == 0, float('-inf'))causal_mask = causal_mask.masked_fill(causal_mask == 1, 0.0)# Transformer块for block in self.blocks:x = block(x, causal_mask)x = self.ln_f(x)logits = self.head(x)return logits# 使用示例
if __name__ == "__main__":# 检查CUDA是否可用if torch.cuda.is_available():print("CUDA可用，运行示例...")# 运行基础示例output = basic_memory_efficient_attention()print("基础示例完成")# 测试多头注意力mha = MemoryEfficientMultiHeadAttention(512, 8).cuda()x = torch.randn(2, 1024, 512, device='cuda')out = mha(x)print(f"多头注意力输出形状: {out.shape}")# 测试完整模型model = GPTWithXFormers(vocab_size=10000,embed_dim=768,num_heads=12,num_layers=6,max_seq_len=2048).cuda()input_ids = torch.randint(0, 10000, (2, 512), device='cuda')logits = model(input_ids)print(f"模型输出形状: {logits.shape}")else:print("需要CUDA支持才能运行xFormers示例")