0
0

Transformer模型变长序列优化：解析PyTorch上的FlashAttention2与xFormers

随着生成式AI（genAI）模型在应用范围和模型规模方面的持续扩展，其训练和部署所需的计算资源及相关成本也呈现显著增长趋势，模型优化对于提升运行时性能和降低运营成本变得尤为关键。作为现代genAI系统核心组件的Transformer架构及其注意力机制，由于其计算密集型的特性，成为优化的重点对象。

在前面的文章中，我们已经介绍了优化注意力核函数能够显著提升Transformer模型的性能。本文将进一步探讨变长输入序列这一挑战——这是真实世界数据（如文档、代码、时间序列等）的固有特征。

批处理变长输入的技术挑战

在典型的深度学习工作负载中，单个样本在传输至GPU并输入AI模型之前需要进行批处理。批处理不仅能提高计算效率，还能在训练过程中促进模型收敛。通常情况下，批处理操作是通过在新的维度（批次维度）上堆叠所有样本张量来实现的。但是torch.stack操作要求所有张量具有相同的形状，这与变长序列的特性相矛盾。

解决这一挑战的传统方法是将输入序列填充至固定长度后再进行堆叠。这种方法需要在模型中实现适当的掩码机制，以确保输出不受填充元素的影响。在注意力层中，填充掩码用于标识哪些token是填充token，从而在计算注意力时予以忽略（参考PyTorch MultiheadAttention的实现）。这种填充方法会导致GPU资源的显著浪费，增加计算成本并降低开发效率。这一问题在大规模AI模型中表现得尤为突出。

序列连接策略

避免填充的一种替代方案是沿着现有维度连接序列，而非在新维度上堆叠。与torch.stack不同，torch.cat允许处理不同形状的输入张量。连接操作的输出是一个长度等于所有输入序列长度之和的单一序列。为了使这种方案有效，需要为序列配备注意力掩码，确保每个token只关注其原始序列中的其他token，这一过程通常被称为文档掩码。若用N表示所有序列的总长度，采用大O符号表示，则掩码的空间复杂度为*O(N²)，注意力层的计算复杂度也为O(N²)*（因为它需要在计算注意力分数后才应用掩码），这使得该方案的效率极低。

注意力层优化技术

针对上述问题，专门设计的注意力层提供了解决方案。与标准注意力层不同，这类优化的注意力核函数采用了更高效的计算策略。标准注意力层会计算完整的O(N²)注意力分数集合后再应用掩码，而优化后的核函数从设计之初就只计算实际需要的分数。本文将介绍几种具有不同特点的解决方案。

与HuggingFace模型的集成方案

对于使用预训练模型的开发团队来说，迁移至这些优化方案可能存在一定难度。本文将演示如何通过HuggingFace的API简化这一过程，使开发人员能够以最小的代码改动实现这些优化技术的集成。

  • 本文中涉及的平台、库或优化技术的使用并不构成对其的推荐。最适合的技术选择将取决于具体应用场景的要求。
  • 部分讨论的API仍处于原型或测试阶段，其接口可能在未来发生变化。
  • 文中提供的代码示例仅供参考，不保证其在生产环境中的适用性、最优性或稳定性。

实验性LLM模型实现

为了深入讨论这些优化技术，我们首先实现一个简化版的生成模型（部分参考了已有的GPT模型架构）。对于完整的语言模型构建指南，建议参考相关领域的专业教程。

Transformer模块实现

首先构建一个基础的Transformer模块，其设计特别考虑了对不同注意力机制和优化策略的实验需求。虽然该模块执行的计算与标准Transformer模块相同，但我们对传统的运算符选择进行了微调，以支持PyTorch NestedTensor输入的要求。

  1.  # 通用导入  
  2.  importtime, functools  
  3.    
  4.  # torch导入  
  5.  importtorch  
  6.  fromtorch.utils.dataimportDataset, DataLoader  
  7.  importtorch.nnasnn  
  8.    
  9.  # 定义Transformer配置参数  
  10.  BATCH_SIZE=32  
  11.  NUM_HEADS=16  
  12.  HEAD_DIM=64  
  13.  DIM=NUM_HEADS*HEAD_DIM  
  14.  DEPTH=24  
  15.  NUM_TOKENS=1024  
  16.  MAX_SEQ_LEN=1024  
  17.  PAD_ID=0  
  18.  DEVICE='cuda'  
  19.    
  20.  classMyAttentionBlock(nn.Module):  
  21.      def__init__(  
  22.              self,  
  23.              attn_fn,  
  24.              dim,  
  25.              num_heads,  
  26.              format=None,  
  27.              **kwargs  
  28.      ):  
  29.          super().__init__()  
  30.          self.attn_fn=attn_fn  
  31.          self.num_heads=num_heads  
  32.          self.dim=dim  
  33.          self.head_dim=dim//num_heads  
  34.          self.norm1=nn.LayerNorm(dim, bias=False)  
  35.          self.norm2=nn.LayerNorm(dim, bias=False)  
  36.          self.qkv=nn.Linear(dim, dim*3)  
  37.          self.proj=nn.Linear(dim, dim)  
  38.    
  39.          # MLP层实现  
  40.          self.fc1=nn.Linear(dim, dim*4)  
  41.          self.act=nn.GELU()  
  42.          self.fc2=nn.Linear(dim*4, dim)  
  43.    
  44.          self.permute=functools.partial(torch.transpose, dim0=1, dim1=2)  
  45.          ifformat=='bshd':  
  46.              self.permute=nn.Identity()  
  47.    
  48.      defmlp(self, x):  
  49.          x=self.fc1(x)  
  50.          x=self.act(x)  
  51.          x=self.fc2(x)  
  52.          returnx  
  53.    
  54.      defreshape_and_permute(self,x, batch_size):  
  55.          x=x.view(batch_size, -1, self.num_heads, self.head_dim)  
  56.          returnself.permute(x)  
  57.    
  58.      defforward(self, x_in, attn_mask=None):  
  59.          batch_size=x_in.size(0)  
  60.          x=self.norm1(x_in)  
  61.          qkv=self.qkv(x)  
  62.    
  63.          # 为支持PyTorch嵌套张量，采用先分割后重排的策略
  64.          # 而非传统的先重排后分割输入状态的方法
  65.          q, k, v=qkv.chunk(3, -1)  
  66.          q=self.reshape_and_permute(q, batch_size)  
  67.          k=self.reshape_and_permute(k, batch_size)  
  68.          v=self.reshape_and_permute(v, batch_size)  
  69.            
  70.          # 应用注意力函数
  71.          x=self.attn_fn(q, k, v, attn_mask=attn_mask)  
  72.    
  73.          # 输出重排与维度调整 
  74.          x=self.permute(x).reshape(batch_size, -1, self.dim)  
  75.          x=self.proj(x)  
  76.          x=x+x_in  
  77.          x=x+self.mlp(self.norm2(x))  
  78.          returnx

这种实现为我们提供了一个灵活的实验平台，可以用于测试各种注意力机制的性能表现。

Transformer解码器架构实现

基于前面的可配置Transformer模块，我们构建了一个标准的Transformer解码器架构。这一实现遵循了典型的Transformer架构设计范式，同时保持了足够的灵活性以适应不同的实验需求。

  1.  classMyDecoder(nn.Module):  
  2.      def__init__(  
  3.              self,  
  4.              block_fn,  
  5.              num_tokens,  
  6.              dim,  
  7.              num_heads,  
  8.              num_layers,  
  9.              max_seq_len,  
  10.              pad_idx=None  
  11.      ):  
  12.          super().__init__()  
  13.          self.num_heads=num_heads  
  14.          self.pad_idx=pad_idx  
  15.          # 构建token嵌入层，支持填充标记处理
  16.          self.embedding=nn.Embedding(num_tokens, dim, padding_idx=pad_idx)  
  17.          # 位置编码嵌入层
  18.          self.positional_embedding=nn.Embedding(max_seq_len, dim)  
  19.          # 构建多层Transformer块
  20.          self.blocks=nn.ModuleList([  
  21.              block_fn(  
  22.                  dim=dim,  
  23.                  num_heads=num_heads  
  24.              )  
  25.              for_inrange(num_layers)])  
  26.          # 输出投影层
  27.          self.output=nn.Linear(dim, num_tokens)  
  28.    
  29.      defembed_tokens(self, input_ids, position_ids=None):  
  30.          # token嵌入与位置编码的结合
  31.          x=self.embedding(input_ids)  
  32.          ifposition_idsisNone:  
  33.              position_ids=torch.arange(input_ids.shape[1],  
  34.                                          device=x.device)  
  35.          x=x+self.positional_embedding(position_ids)  
  36.          returnx  
  37.    
  38.      defforward(self, input_ids, position_ids=None, attn_mask=None):  
  39.          # Token嵌入与位置编码的整合
  40.          x=self.embed_tokens(input_ids, position_ids)  
  41.          
  42.          # 自动生成填充掩码（如果需要）
  43.          ifself.pad_idxisnotNone:  
  44.              assertattn_maskisNone  
  45.              # 生成布尔类型的填充掩码
  46.              attn_mask= (input_ids!=self.pad_idx)  
  47.              attn_mask=attn_mask.view(BATCH_SIZE, 1, 1, -1) \  
  48.                  .expand(-1, self.num_heads, -1, -1)  
  49.    
  50.          # 依次通过所有Transformer层
  51.          forbinself.blocks:  
  52.              x=b(x, attn_mask)  
  53.    
  54.          # 生成最终输出logits
  55.          logits=self.output(x)  
  56.          returnlogits

变长序列数据集实现

为了系统评估模型性能，还要实现了一个包含变长序列的模拟数据集。为简化实验设置，我们采用了一个固定的序列长度分布。在实际应用场景中，序列长度分布通常由数据本身的特性决定，如文档长度或音频片段持续时间。值得注意的是，序列长度分布会直接影响填充策略导致的计算效率损失。

  1.  # 随机数据集实现
  2.  classFakeDataset(Dataset):  
  3.      def__len__(self):  
  4.          return1000000  
  5.    
  6.      def__getitem__(self, index):  
  7.          # 生成随机长度的序列
  8.          length=torch.randint(1, MAX_SEQ_LEN, (1,))  
  9.          sequence=torch.randint(1, NUM_TOKENS, (length+1,))  
  10.          # 准备输入输出对
  11.          inputs=sequence[:-1]  
  12.          targets=sequence[1:]  
  13.          returninputs, targets  
  14.    
  15.  defpad_sequence(sequence, length, pad_val):  
  16.      """实现序列填充功能"""
  17.      returntorch.nn.functional.pad(  
  18.          sequence,  
  19.          (0, length-sequence.shape[0]),  
  20.          value=pad_val  
  21.      )  
  22.    
  23.  defcollate_with_padding(batch):  
  24.      """批处理数据整理函数，包含填充操作"""
  25.      padded_inputs= []  
  26.      padded_targets= []  
  27.      forbinbatch:  
  28.          padded_inputs.append(pad_sequence(b[0], MAX_SEQ_LEN, PAD_ID))  
  29.          padded_targets.append(pad_sequence(b[1], MAX_SEQ_LEN, PAD_ID))  
  30.      padded_inputs=torch.stack(padded_inputs, dim=0)  
  31.      padded_targets=torch.stack(padded_targets, dim=0)  
  32.      return {  
  33.          'inputs': padded_inputs,  
  34.          'targets': padded_targets  
  35.      }  
  36.    
  37.  defdata_to_device(data, device):  
  38.      """数据传输至指定设备的工具函数"""
  39.      ifisinstance(data, dict):  
  40.          return {  
  41.              key: data_to_device(val,device)  
  42.              forkey, valindata.items()  
  43.          }  
  44.      elifisinstance(data, (list, tuple)):  
  45.          returntype(data)(  
  46.              data_to_device(val, device) forvalindata  
  47.          )  
  48.      elifisinstance(data, torch.Tensor):  
  49.          returndata.to(device=device, non_blocking=True)  
  50.      else:  
  51.          returndata.to(device=device)

这些实现为我们提供了一个完整的实验框架，可以系统地评估不同优化策略的效果。通过这个框架可以精确地量化各种优化方法对计算效率的影响。

训练与评估流程设计

最后还要实现了一个完整的main函数，用于对变长序列数据进行训练和评估。这个实现包含了现代深度学习训练流程的关键要素，包括混合精度训练、梯度计算和性能监控。

  1.  defmain(  
  2.      block_fn,   
  3.      data_collate_fn=collate_with_padding,  
  4.      pad_idx=None,  
  5.      train=True,  
  6.      compile=False  
  7.  ):  
  8.      # 初始化随机种子，确保实验可重现
  9.      torch.random.manual_seed(0)  
  10.      device=torch.device(DEVICE)  
  11.      # 启用高精度矩阵乘法，提升数值稳定性
  12.      torch.set_float32_matmul_precision("high")  
  13.    
  14.      # 配置数据加载器，优化数据传输效率
  15.      data_set=FakeDataset()  
  16.      data_loader=DataLoader(  
  17.          data_set,  
  18.          batch_size=BATCH_SIZE,  
  19.          collate_fn=data_collate_fn,  
  20.          num_workers=12,  # 并行数据加载
  21.          pin_memory=True,  # 启用内存钉扎，加速GPU传输
  22.          drop_last=True   # 确保批次大小一致
  23.      )  
  24.    
  25.      # 模型实例化与设备配置
  26.      model=MyDecoder(  
  27.          block_fn=block_fn,  
  28.          num_tokens=NUM_TOKENS,  
  29.          dim=DIM,  
  30.          num_heads=NUM_HEADS,  
  31.          num_layers=DEPTH,  
  32.          max_seq_len=MAX_SEQ_LEN,  
  33.          pad_idx=pad_idx  
  34.      ).to(device)  
  35.    
  36.      # 启用torch编译优化（如果指定）
  37.      ifcompile:  
  38.          model=torch.compile(model)  
  39.    
  40.      # 损失函数与优化器配置
  41.      criterion=torch.nn.CrossEntropyLoss(ignore_index=PAD_ID)  
  42.      optimizer=torch.optim.SGD(model.parameters())  
  43.    
  44.      # 定义训练步骤
  45.      deftrain_step(model, inputs, targets,   
  46.                     position_ids=None, attn_mask=None):  
  47.          # 使用自动混合精度训练
  48.          withtorch.amp.autocast(DEVICE, dtype=torch.bfloat16):  
  49.              outputs=model(inputs, position_ids, attn_mask)  
  50.              outputs=outputs.view(-1, NUM_TOKENS)  
  51.              targets=targets.flatten()  
  52.              loss=criterion(outputs, targets)  
  53.          # 梯度清零、反向传播与参数更新
  54.          optimizer.zero_grad(set_to_none=True)  
  55.          loss.backward()  
  56.          optimizer.step()  
  57.    
  58.      # 定义评估步骤
  59.      @torch.no_grad()  
  60.      defeval_step(model, inputs, targets,   
  61.                    position_ids=None, attn_mask=None):  
  62.          withtorch.amp.autocast(DEVICE, dtype=torch.bfloat16):  
  63.              outputs=model(inputs, position_ids, attn_mask)  
  64.              # 处理嵌套张量的特殊情况
  65.              ifoutputs.is_nested:  
  66.                  outputs=outputs.data._values  
  67.                  targets=targets.data._values  
  68.              else:
  69.                  outputs=outputs.view(-1, NUM_TOKENS)  
  70.                  targets=targets.flatten()  
  71.              loss=criterion(outputs, targets)  
  72.          returnloss  
  73.    
  74.      # 根据模式选择执行函数
  75.      iftrain:  
  76.          model.train()  
  77.          step_fn=train_step  
  78.      else:  
  79.          model.eval()  
  80.          step_fn=eval_step  
  81.    
  82.      # 性能监控初始化
  83.      t0=time.perf_counter()  
  84.      summ=0  
  85.      count=0  
  86.    
  87.      # 主训练循环
  88.      forstep, datainenumerate(data_loader):  
  89.          # 数据传输至GPU
  90.          data=data_to_device(data, device=device)  
  91.          step_fn(model, data['inputs'], data['targets'],  
  92.                         position_ids=data.get('indices'),  
  93.                         attn_mask=data.get('attn_mask'))  
  94.    
  95.          # 性能统计计算
  96.          batch_time=time.perf_counter() -t0  
  97.          ifstep>20:  # 跳过预热阶段
  98.              summ+=batch_time  
  99.              count+=1  
  100.          t0=time.perf_counter()  
  101.          ifstep>=100:  # 限制评估步数
  102.              break  
  103.      # 输出平均步骤时间
  104.      print(f'average step time: {summ/count}')

PyTorch SDPA基准测试

在基准测试中，我们使用PyTorch的SDPA（Scaled Dot-Product Attention）机制作为baseline。在NVIDIA H100 GPU上进行测试，环境配置为CUDA 12.4和PyTorch 2.5.1，同时评估了启用和禁用torch.compile的场景。

  1.  # 配置SDPA基准测试
  2.  fromtorch.nn.functionalimportscaled_dot_product_attentionassdpa  
  3.  block_fn=functools.partial(MyAttentionBlock, attn_fn=sdpa)  
  4.  causal_block_fn=functools.partial(  
  5.      MyAttentionBlock,  
  6.      attn_fn=functools.partial(sdpa, is_causal=True)  
  7.  )  
  8.  
  9.  # 执行评估与训练测试
  10.  formodein ['eval', 'train']:  
  11.      forcompilein [False, True]:  
  12.          block_func=causal_block_fn\  
  13.              ifmode=='train'elseblock_fn  
  14.          print(f'{mode} with {collate}, '  
  15.                f'{"compiled"ifcompileelse"uncompiled"}')  
  16.          main(block_fn=block_func,  
  17.               pad_idx=PAD_ID,  
  18.               train=mode=='train',  
  19.               compile=compile)

基准测试结果显示：

  • 评估模式- 未编译：132毫秒(ms)- 编译后：130 ms
  • 训练模式- 未编译：342 ms- 编译后：299 ms

这些基准数据为后续优化方案提供了重要的参考点。通过这些数据可以客观评估不同优化策略的效果。

变长输入序列的优化策略

下面我们就要详细探讨几种针对Transformer模型中变长输入序列处理的优化策略。每种策略都有其独特的优势和应用场景，将通过实验数据来评估它们的效果。

动态填充优化

第一个优化策略关注填充机制本身。不同于传统方法将每个批次的序列填充到固定长度，我们采用了一种动态填充策略：将序列填充到当前批次中最长序列的长度。这种方法可以显著减少不必要的计算开销。以下是具体实现：

  1.  defcollate_pad_to_longest(batch):  
  2.      padded_inputs= []  
  3.      padded_targets= []  
  4.      # 计算当前批次中的最大序列长度
  5.      max_length=max([b[0].shape[0] forbinbatch])  
  6.      # 只填充到最大长度，而不是预设的固定长度
  7.      forbinbatch:  
  8.          padded_inputs.append(pad_sequence(b[0], max_length, PAD_ID))  
  9.          padded_targets.append(pad_sequence(b[1], max_length, PAD_ID))  
  10.      # 堆叠处理后的序列
  11.      padded_inputs=torch.stack(padded_inputs, dim=0)  
  12.      padded_targets=torch.stack(padded_targets, dim=0)  
  13.      return {  
  14.          'inputs': padded_inputs,  
  15.          'targets': padded_targets  
  16.      }  
  17.  
  18.  # 执行动态填充优化的性能测试
  19.  formodein ['eval', 'train']:  
  20.      forcompilein [False, True]:  
  21.          block_func=causal_block_fn\  
  22.              ifmode=='train'elseblock_fn  
  23.          print(f'{mode} with {collate}, '  
  24.                f'{"compiled"ifcompileelse"uncompiled"}')  
  25.          main(block_fn=block_func,  
  26.               data_collate_fn=collate_pad_to_longest,  
  27.               pad_idx=PAD_ID,  
  28.               train=mode=='train',  
  29.               compile=compile)

实验结果表明，动态填充策略带来了可观的性能提升：

  • 评估模式- 未编译：129 ms（相比基准提升2.3%）- 编译后：116 ms（相比基准提升10.8%）
  • 训练模式- 未编译：337 ms（相比基准提升1.5%）- 编译后：294 ms（相比基准提升1.7%）

PyTorch NestedTensors优化方案

接下来，评估在PyTorch NestedTensors下的应用。这是一个目前处于原型阶段的特性，它允许我们直接处理不同长度的张量，这些张量被称为"jagged"或"ragged"张量。这种方法避免了显式填充的需求，但需要特别注意张量操作的兼容性。

以下代码展示了如何使用NestedTensors处理变长序列：

  1.  defnested_tensor_collate(batch):  
  2.      # 创建嵌套张量结构
  3.      inputs=torch.nested.as_nested_tensor([b[0] forbinbatch],  
  4.                                             layout=torch.jagged)  
  5.      targets=torch.nested.as_nested_tensor([b[1] forbinbatch],  
  6.                                              layout=torch.jagged)  
  7.      # 生成位置索引
  8.      indices=torch.concat([torch.arange(b[0].shape[0]) forbinbatch])  
  9.    
  10.      # 特别处理：创建与输入具有相同"jagged"形状的索引张量
  11.      # 这是由于NestedTensors对张量操作的限制
  12.      xx=torch.empty_like(inputs)  
  13.      xx.data._values[:] =indices  
  14.    
  15.      return {  
  16.          'inputs': inputs,  
  17.          'targets': targets,  
  18.          'indices': xx  
  19.      }  
  20.  
  21.  # 执行NestedTensors性能评估
  22.  forcompilein [False, True]:  
  23.      print(f'eval with nested tensors, '  
  24.            f'{"compiled"ifcompileelse"uncompiled"}')  
  25.      main(  
  26.          block_fn=block_fn,  
  27.          data_collate_fn=nested_tensor_collate,  
  28.          train=False,  
  29.          compile=compile  
  30.      )

NestedTensors方案的性能表现令人瞩目：

  • 未编译模式下：131 ms（与基准相当）
  • 编译模式下：42 ms（性能提升约3倍）

这个显著的性能提升主要得益于两个因素：

  1. 避免了填充带来的冗余计算
  2. 编译优化能够更好地利用NestedTensors的特性进行计算优化

但是需要注意的是，由于NestedTensors仍处于原型阶段，在实际应用中需要谨慎评估其稳定性和兼容性。

FlashAttention2优化实现

前面的文章我们已经探讨了FlashAttention对Transformer模型性能的影响。本节将重点介绍flash-attn 2.7.0版本中的flash_attn_varlen_func，这是一个专门为处理可变长度输入设计的API。这个优化方案的核心思想是将批次中的所有序列连接成一个连续序列，同时使用一个特殊的索引张量（cu_seqlens）来追踪各个原始序列的边界位置。

以下是这种方法的详细实现：

  1.  defcollate_concat(batch):  
  2.      # 将所有序列连接为单一序列，并添加batch维度
  3.      inputs=torch.concat([b[0] forbinbatch]).unsqueeze(0)  
  4.      targets=torch.concat([b[1] forbinbatch]).unsqueeze(0)  
  5.      # 生成位置索引
  6.      indices=torch.concat([torch.arange(b[0].shape[0]) forbinbatch])  
  7.      # 计算累积序列长度，用于定位序列边界
  8.      seqlens=torch.tensor([b[0].shape[0] forbinbatch])  
  9.      seqlens=torch.cumsum(seqlens, dim=0, dtype=torch.int32)  
  10.      # 添加起始位置0，形成完整的区间索引
  11.      cu_seqlens=torch.nn.functional.pad(seqlens, (1, 0))  
  12.    
  13.      return {  
  14.          'inputs': inputs,  
  15.          'targets': targets,  
  16.          'indices': indices,  
  17.          'attn_mask': cu_seqlens  
  18.      }  
  19.  
  20.  # 配置FlashAttention变长序列处理函数
  21.  fromflash_attnimportflash_attn_varlen_func  
  22.  # 标准版本：用于评估模式
  23.  fa_varlen=lambdaq, k, v, attn_mask: flash_attn_varlen_func(  
  24.      q.squeeze(0),  
  25.      k.squeeze(0),  
  26.      v.squeeze(0),  
  27.      cu_seqlens_q=attn_mask,  
  28.      cu_seqlens_k=attn_mask,  
  29.      max_seqlen_q=MAX_SEQ_LEN,  
  30.      max_seqlen_k=MAX_SEQ_LEN  
  31.  ).unsqueeze(0)  
  32.  
  33.  # 因果版本：用于训练模式，确保注意力的因果性
  34.  fa_varlen_causal=lambdaq, k, v, attn_mask: flash_attn_varlen_func(  
  35.      q.squeeze(0),  
  36.      k.squeeze(0),  
  37.      v.squeeze(0),  
  38.      cu_seqlens_q=attn_mask,  
  39.      cu_seqlens_k=attn_mask,  
  40.      max_seqlen_q=MAX_SEQ_LEN,  
  41.      max_seqlen_k=MAX_SEQ_LEN,  
  42.      causal=True  
  43.  ).unsqueeze(0)  
  44.  
  45.  # 配置使用FlashAttention的Transformer块
  46.  block_fn=functools.partial(MyAttentionBlock,  
  47.                               attn_fn=fa_varlen,  
  48.                               format='bshd')  
  49.  causal_block_fn=functools.partial(MyAttentionBlock,  
  50.                                      attn_fn=fa_varlen_causal,  
  51.                                      format='bshd')  
  52.  
  53.  # 执行性能评估
  54.  print('flash-attn eval')  
  55.  main(  
  56.      block_fn=block_fn,  
  57.      data_collate_fn=collate_concat,  
  58.      train=False  
  59.  )  
  60.  
  61.  print('flash-attn train')  
  62.  main(  
  63.      block_fn=causal_block_fn,  
  64.      data_collate_fn=collate_concat,  
  65.      train=True,  
  66.  )

FlashAttention2的性能提升非常显著：

  • 评估模式：51 ms（比基准快2.6倍）
  • 训练模式：160 ms（比基准快2.1倍）

这种显著的性能提升主要源于以下技术创新：

  1. 高效的内存访问模式：通过分块计算和重排访问模式，减少了内存带宽需求
  2. 序列边界的精确追踪：使用cu_seqlens避免了填充带来的计算浪费
  3. 优化的CUDA核函数实现：专门针对变长序列场景进行了优化

目前flash_attn_varlen_func还不支持torch.compile。这是因为其实现包含了复杂的CUDA核函数，这些函数目前还无法被PyTorch的编译器正确处理。在实际应用中，即使没有编译优化，其性能仍然远超基准实现。

XFormers内存高效注意力机制的实现

前面的文章我们也介绍了xFormers (0.0.28)中的memory_efficient_attention操作符。下面我们将重点探讨BlockDiagonalMask的应用，这是一个专门为处理任意长度输入序列设计的掩码机制。这种方法的独特之处在于它能够在保持高计算效率的同时精确处理序列边界。

让我们详细分析这个实现：

  1.  fromxformers.opsimportfmha  
  2.  fromxformers.opsimportmemory_efficient_attentionasmea  
  3.  
  4.  defcollate_xformer(batch):  
  5.      # 序列连接与维度处理
  6.      inputs=torch.concat([b[0] forbinbatch]).unsqueeze(0)  
  7.      targets=torch.concat([b[1] forbinbatch]).unsqueeze(0)  
  8.      indices=torch.concat([torch.arange(b[0].shape[0]) forbinbatch])  
  9.      
  10.      # 提取每个序列的长度信息
  11.      seqlens= [b[0].shape[0] forbinbatch]  
  12.      # 设置批次大小，用于掩码生成
  13.      batch_sizes= [1forbinbatch]  
  14.      # 创建块对角掩码，处理序列边界
  15.      block_diag=fmha.BlockDiagonalMask.from_seqlens(seqlens, device='cpu')  
  16.      block_diag._batch_sizes=batch_sizes  
  17.    
  18.      return {  
  19.          'inputs': inputs,  
  20.          'targets': targets,  
  21.          'indices': indices,  
  22.          'attn_mask': block_diag  
  23.      }  
  24.  
  25.  # 配置评估模式的注意力计算
  26.  mea_eval=lambdaq, k, v, attn_mask: mea(  
  27.      q,k,v, attn_bias=attn_mask)  
  28.  
  29.  # 配置训练模式的因果注意力计算
  30.  mea_train=lambdaq, k, v, attn_mask: mea(  
  31.      q,k,v, attn_bias=attn_mask.make_causal())  
  32.  
  33.  # 构建评估模式的Transformer块配置
  34.  block_fn=functools.partial(MyAttentionBlock,  
  35.                               attn_fn=mea_eval,  
  36.                               format='bshd')  
  37.  
  38.  # 构建训练模式的Transformer块配置
  39.  causal_block_fn=functools.partial(MyAttentionBlock,  
  40.                               attn_fn=mea_train,  
  41.                               format='bshd')  
  42.  
  43.  # 执行性能评估测试
  44.  print(f'xFormer Attention ')  
  45.  forcompilein [False, True]:  
  46.      print(f'eval with xFormer Attention, '  
  47.            f'{"compiled"ifcompileelse"uncompiled"}')  
  48.      main(block_fn=block_fn,  
  49.           train=False,  
  50.           data_collate_fn=collate_xformer,  
  51.           compile=compile)  
  52.  
  53.  print(f'train with xFormer Attention')  
  54.  main(block_fn=causal_block_fn,  
  55.       train=True,  
  56.       data_collate_fn=collate_xformer)

XFormers的性能表现令人印象深刻：

  • 评估模式（未编译）：50 ms
  • 评估模式（已编译）：42 ms
  • 训练模式：159 ms

这些性能数据反映了XFormers优化策略的几个关键优势：

高效的内存管理：

  • BlockDiagonalMask提供了一种内存友好的方式来处理序列边界
  • 避免了传统填充方法带来的内存浪费

优化的计算模式：

  • 通过块对角矩阵的形式组织注意力计算
  • 有效减少了不必要的计算操作

灵活的掩码机制：

  • make_causal方法可以方便地转换为因果注意力模式
  • 支持动态序列长度而不损失性能

编译优化兼容性：

  • 在评估模式下，能够充分利用torch.compile带来的优化
  • 虽然训练模式下的编译支持仍有限制，但基础性能已经相当优秀

但是在实际应用中还需要注意以下几点：

  1. 训练模式下的编译支持目前仍有限制
  2. 性能提升的程度可能会随具体的硬件环境和模型配置而变化
  3. 内存使用模式与标准实现有所不同，可能需要相应调整内存预算

优化结果综合分析

通过上述一系列实验，我们获得了不同优化策略的性能数据。让我们通过一张对比图来直观地理解这些结果：

不同优化方法的步骤时间对比（数值越低表示性能越好），xFormer的memory_efficient_attention表现最为出色，在评估时实现了约3倍的性能提升，在训练时获得了约2倍的加速。这里需要特别说明的是，这些性能数据不应被视为普适性结论。在实际应用中，不同注意力计算方法的性能表现会因具体的模型架构、硬件配置和应用场景而呈现显著差异。

HuggingFace模型的变长输入优化

相比从零开始构建模型，现代机器学习开发更多地依赖于预训练模型的微调。虽然前文描述的优化技术可以在不改变模型权重和行为的前提下集成到现有模型中，但如何高效实施这种集成仍然是一个重要问题。在本节中，我们将探讨如何在HuggingFace生态系统中实现这些优化。

GPT2LMHeadModel实验设计

为了演示优化过程，我们选择了GPT2LMHeadModel作为实验对象。首先要调整数据集和数据处理流程以适配HuggingFace的标准接口：

  1.  fromtransformersimportGPT2Config, GPT2LMHeadModel  
  2.  
  3.  classHuggingFaceFakeDataset(Dataset):  
  4.      """用于HuggingFace模型实验的模拟数据集"""
  5.      def__len__(self):  
  6.          return1000000  
  7.    
  8.      def__getitem__(self, index):  
  9.          # 生成随机长度序列
  10.          length=torch.randint(1, MAX_SEQ_LEN, (1,))  
  11.          input_ids=torch.randint(1, NUM_TOKENS, (length,))  
  12.          # 创建训练标签，将第一个token标记为填充
  13.          labels=input_ids.clone()  
  14.          labels[0] =PAD_ID  
  15.          return {  
  16.              'input_ids': input_ids,  
  17.              'labels': labels  
  18.          }  
  19.  
  20.  defhf_collate_with_padding(batch):  
  21.      """HuggingFace模型的批处理数据整理函数"""
  22.      padded_inputs= []  
  23.      padded_labels= []  
  24.      forbinbatch:  
  25.          input_ids=b['input_ids']  
  26.          labels=b['labels']  
  27.          # 执行序列填充
  28.          padded_inputs.append(pad_sequence(input_ids, MAX_SEQ_LEN, PAD_ID))  
  29.          padded_labels.append(pad_sequence(labels, MAX_SEQ_LEN, PAD_ID))  
  30.      # 转换为批次张量
  31.      padded_inputs=torch.stack(padded_inputs, dim=0)  
  32.      padded_labels=torch.stack(padded_labels, dim=0)  
  33.      return {  
  34.          'input_ids': padded_inputs,  
  35.          'labels': padded_labels,  
  36.          'attention_mask': (padded_inputs!=PAD_ID)  # 生成注意力掩码
  37.      }

在这个实现中，特别注意以下几个方面：

  1. 数据格式的一致性：确保生成的数据符合HuggingFace的预期格式
  2. 填充处理：实现了高效的填充策略，同时保持数据的语义完整性
  3. 注意力掩码：自动生成适当的注意力掩码，用于处理填充token

HuggingFace模型训练流程的实现

为了系统评估优化效果，还需要实现了一个完整的训练流程。这个实现不仅包含了标准的训练循环，还集成了各种现代深度学习的优化技术：

  1.  defhf_main(  
  2.      config,  
  3.      collate_fn=hf_collate_with_padding,  
  4.      compile=False  
  5.  ):  
  6.      # 初始化环境配置
  7.      torch.random.manual_seed(0)  
  8.      device=torch.device(DEVICE)  
  9.      # 启用高精度矩阵计算，提升数值稳定性
  10.      torch.set_float32_matmul_precision("high")  
  11.    
  12.      # 配置数据加载器
  13.      data_set=HuggingFaceFakeDataset()  
  14.      data_loader=DataLoader(  
  15.          data_set,  
  16.          batch_size=BATCH_SIZE,  
  17.          collate_fn=collate_fn,  
  18.          num_workers=12ifDEVICE=="CUDA"else0,  # 根据设备类型调整并行加载
  19.          pin_memory=True,  # 启用内存钉扎加速数据传输
  20.          drop_last=True   # 保持批次大小一致
  21.      )  
  22.    
  23.      # 模型初始化和设备迁移
  24.      model=GPT2LMHeadModel(config).to(device)  
  25.    
  26.      # 条件编译优化
  27.      ifcompile:  
  28.          model=torch.compile(model)  
  29.    
  30.      # 配置训练组件
  31.      criterion=torch.nn.CrossEntropyLoss(ignore_index=PAD_ID)  
  32.      optimizer=torch.optim.SGD(model.parameters())  
  33.    
  34.      model.train()  # 设置训练模式
  35.    
  36.      # 性能监控初始化
  37.      t0=time.perf_counter()  
  38.      summ=0  
  39.      count=0  
  40.    
  41.      # 主训练循环
  42.      forstep, datainenumerate(data_loader):  
  43.          # 数据迁移到目标设备
  44.          data=data_to_device(data, device=device)  
  45.          input_ids=data['input_ids']  
  46.          labels=data['labels']  
  47.          position_ids=data.get('position_ids')  
  48.          attn_mask=data.get('attention_mask')  
  49.          
  50.          # 使用混合精度训练
  51.          withtorch.amp.autocast(DEVICE, dtype=torch.bfloat16):  
  52.              outputs=model(input_ids=input_ids,  
  53.                            position_ids=position_ids,  
  54.                            attention_mask=attn_mask)  
  55.              # 处理序列偏移，确保正确的预测目标
  56.              logits=outputs.logits[..., :-1, :].contiguous()  
  57.              labels=labels[..., 1:].contiguous()  
  58.              # 计算损失
  59.              loss=criterion(logits.view(-1, NUM_TOKENS), labels.flatten())  
  60.          
  61.          # 梯度更新
  62.          optimizer.zero_grad(set_to_none=True)  
  63.          loss.backward()  
  64.          optimizer.step()  
  65.    
  66.          # 记录和更新性能统计
  67.          batch_time=time.perf_counter() -t0  
  68.          ifstep>20:  # 跳过初始预热步骤
  69.              summ+=batch_time  
  70.              count+=1  
  71.          t0=time.perf_counter()  
  72.          ifstep>=100:  # 限制评估步数
  73.              break  
  74.      # 输出平均步骤时间
  75.      print(f'average step time: {summ/count}')

这个实现中包含了几个关键的技术考虑：

  1. 混合精度训练：- 使用torch.amp.autocast实现自动混合精度- 选择bfloat16数据类型，在保持数值稳定性的同时提升计算效率
  2. 内存优化：- 使用pin_memory和non_blocking数据传输- 通过set_to_none=True优化梯度清零操作
  3. 性能监控：- 实现了精确的性能统计- 考虑了预热阶段的影响
  4. 数据处理：- 处理序列的因果关系，确保预测目标的正确性- 优化了张量操作的内存布局（通过contiguous()调用）

这种实现为我们提供了一个可靠的基准测试平台，使我们能够准确评估不同优化策略的效果。

基于SDPA的基准测试实现

使用标准SDPA（Scaled Dot-Product Attention）的基准测试。这为后续的优化策略提供了一个参考点：

  1.  config=GPT2Config(  
  2.          n_layer=DEPTH,     # 设置模型深度
  3.          n_embd=DIM,       # 嵌入维度
  4.          n_head=NUM_HEADS, # 注意力头数量
  5.          vocab_size=NUM_TOKENS,  # 词表大小
  6.      )  
  7.  
  8.  # 执行不同配置的基准测试
  9.  forcompilein [False, True]:  
  10.      print(f"HF GPT2 train with SDPA, compile={compile}")  
  11.      hf_main(config=config, compile=compile)

基准测试的结果显示，在不同配置下的性能表现存在显著差异：

  • 未启用编译优化时：815毫秒
  • 启用编译优化后：440毫秒

这个基准测试揭示了编译优化能带来接近1.85倍的性能提升，这主要得益于PyTorch的动态编译技术对计算图的优化。

FlashAttention2的集成实现

接下来需要通过配置HuggingFace的内置支持来启用FlashAttention2。这种方法的优势在于实现简单，只需要修改配置参数即可：

  1.  flash_config=GPT2Config(  
  2.          n_layer=DEPTH,  
  3.          n_embd=DIM,  
  4.          n_head=NUM_HEADS,  
  5.          vocab_size=NUM_TOKENS,  
  6.          attn_implementation='flash_attention_2'  # 启用FlashAttention2
  7.      )  
  8.  
  9.  print(f"HF GPT2 train with flash")  
  10.  hf_main(config=flash_config)

这个简单的配置修改产生了显著的性能提升：

  • 步骤时间降至620毫秒
  • 相比未优化版本提升了约31.9%的性能

这种改进的实现原理是：HuggingFace框架会在内部自动将填充后的输入数据进行"解填充"（unpadding）操作，然后将处理后的数据传递给优化过的flash_attn_varlen_func函数。这个过程虽然引入了一些数据预处理开销，但整体性能仍然获得了显著提升。

FlashAttention2的无填充优化

虽然上述实现已经带来了可观的性能提升，但在数据处理流程中仍存在一些冗余操作：先对序列进行填充，然后又在内部进行解填充。为了进一步优化性能，可以直接使用未填充的输入数据。

最近的HuggingFace更新增加了对连接序列（未填充）输入的支持，但这个功能目前仅限于特定模型。要在GPT2模型中启用这个功能，需要对modeling_gpt2.py文件进行少量修改。以下是完整的补丁内容：

  1.  @@-370,0+371@@  
  2.  +        position_ids=None  
  3.  @@-444,0+446@@  
  4.  +            position_ids=position_ids  
  5.  @@-611,0+614@@  
  6.  +        position_ids=None  
  7.  @@-621,0+625@@  
  8.  +            position_ids=position_ids  
  9.  @@-1140,0+1145@@  
  10.  +                    position_ids=position_ids

这些修改的主要目的是在attention计算过程中正确传递位置编码信息，确保模型能够准确处理未填充的序列数据。

在了解了必要的修改后，我们可以实现一个专门用于处理未填充序列的数据处理函数：

  1.  defcollate_flatten(batch):  
  2.      # 直接连接所有序列，避免填充操作
  3.      input_ids=torch.concat([b['input_ids'] forbinbatch]).unsqueeze(0)  
  4.      labels=torch.concat([b['labels'] forbinbatch]).unsqueeze(0)  
  5.      # 为每个序列生成对应的位置编码
  6.      position_ids= [torch.arange(b['input_ids'].shape[0]) forbinbatch]  
  7.      position_ids=torch.concat(position_ids)  
  8.    
  9.      return {  
  10.          'input_ids': input_ids,  
  11.          'labels': labels,  
  12.          'position_ids': position_ids  
  13.      }  
  14.  
  15.  # 使用无填充配置进行性能测试
  16.  print(f"HF GPT2 train with flash, no padding")  
  17.  hf_main(config=flash_config, collate_fn=collate_flatten)

这种优化方案的性能表现令人瞩目：

  • 步骤时间降至323毫秒
  • 与使用填充数据的FlashAttention2相比，性能提升了约90%
  • 相比原始基准测试，性能提升了约2.5倍

综合性能分析

通过一张图表来直观地比较不同优化策略的效果：

不同优化方法的步骤时间对比（数值越低表示性能越好）通过系统性的优化，我们实现了显著的性能提升：

  1. 相比未编译的基准版本：性能提升了约2.5倍
  2. 相比启用编译的版本：性能提升了约36%

这些优化成果充分展示了HuggingFace API的灵活性，它使我们能够轻松集成高效的注意力计算核心，显著提升模型在处理变长序列时的训练性能。

总结

随着AI模型在复杂度和应用范围上的持续扩展，性能优化变得越来越重要。本文着重探讨了注意力层的优化策略，并提供了一系列实用的工具和技术来提升Transformer模型的性能。主要的技术贡献包括：

  1. 数据处理优化：- 提出了动态填充和无填充策略- 优化了序列连接和位置编码的处理方式
  2. 计算效率提升：- 利用FlashAttention2实现高效的注意力计算- 通过编译优化提升执行效率
  3. 框架集成：- 展示了如何在HuggingFace生态系统中实现这些优化- 提供了具体的代码修改指南

这些优化策略不仅提供了显著的性能提升，还保持了实现的简洁性和可维护性。对于希望优化自己模型性能的开发者来说，这些方法提供了实用的参考方案。

为了深入了解更多AI模型优化相关内容，建议查阅本系列的第一篇文章以及其他相关技术文档。随着深度学习技术的不断发展，我们期待看到更多创新的优化方案出现。

作者：Chaim Rand

标签： 深度学习 Transformer PyTorch
登录后发布评论

“Transformer模型变长序列优化：解析PyTorch上的FlashAttention2与xFormers”的评论:

还没有评论
关于作者
...
Deephub
更多文章请关注公众号：Deephub-IMBA
相关阅读

标签噪声下的模型评估：如何准确评估AI模型的真实性能，提高模型性能测量的可信度

VideoMind：Chain-of-LoRA突破时间盲区让AI真正看懂长视频

9个主流GAN损失函数的数学原理和Pytorch代码实现：从经典模型到现代变体

DAPO: 面向开源大语言模型的解耦裁剪与动态采样策略优化系统

多模态AI核心技术：CLIP与SigLIP技术原理与应用进展

SWEET-RL：基于训练时信息的多轮LLM代理强化学习框架

时间序列异常检测：MSET-SPRT组合方法的原理和Python代码实现