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示例")