FasterTransformer中的Longformer实现详解

FasterTransformer中的Longformer实现详解

FasterTransformer Transformer related optimization, including BERT, GPT 项目地址: https://gitcode.com/gh_mirrors/fa/FasterTransformer

概述

FasterTransformer项目中的Longformer实现是针对长序列处理场景优化的深度学习模型。Longformer模型最初由AllenAI提出，旨在解决传统Transformer模型在处理长序列时面临的内存消耗和计算复杂度问题。FasterTransformer通过高效的CUDA和C++实现，显著提升了Longformer模型的推理性能。

Longformer模型特点

核心创新

Longformer模型的核心创新在于其注意力机制的设计：

1. 局部窗口注意力：每个token只关注固定窗口范围内的邻近token，将计算复杂度从O(n²)降低到O(n)
2. 全局注意力：为特定任务相关的token（如问答中的问题token）保留全局注意力能力
3. 线性复杂度：整体计算复杂度与序列长度呈线性关系，可处理数千token的长文档

技术优势

相比传统Transformer，Longformer具有以下优势：

• 内存消耗显著降低
• 推理延迟大幅减少
• 支持更长的输入序列
• 在长文档任务上性能优于RoBERTa等模型

FasterTransformer实现细节

实现架构

FasterTransformer目前实现了Longformer的Encoder部分，主要包含以下组件：

1. CUDA/C++核心实现：高性能计算内核
2. PyTorch封装：提供与HuggingFace兼容的接口
3. 问答示例：展示实际应用场景

计算流程

Longformer的计算流程如图1所示（见原文档图1），主要包括：

1. 输入嵌入
2. 局部窗口注意力计算
3. 全局注意力计算
4. 前馈网络
5. 层归一化

关键限制

当前实现有以下限制需要注意：

1. 序列长度必须≥2倍局部注意力窗口大小
2. 头维度(head_dim)目前仅支持64
3. 全局token数量必须小于局部注意力窗口大小
4. 全局token必须连续放置在序列开头

快速上手指南

环境准备

硬件要求

• NVIDIA Pascal/Volta/Turing/Ampere架构GPU
• 推荐使用A100等高性能GPU

软件依赖

• CUDA 11.0+
• PyTorch 1.5.0+
• CMake 3.13+
• Python 3.x
• Transformers 4.9.2

安装步骤

1. 获取模型权重：

git lfs install
git clone https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa

2. 编译FasterTransformer：

mkdir build && cd build
cmake -DCMAKE_BUILD_TYPE=Release -DBUILD_PYT=ON ..
make -j12

运行问答示例

python3 examples/pytorch/longformer/longformer_qa.py \
    --ft-longformer-lib build/lib/libth_transformer.so \
    --model-dir path/to/model \
    --passage "Jim Henson was a nice puppet" \
    --question "Who was Jim Henson?" \
    --repeat-test-num 50

常用参数说明

• --data_type fp16/bf16：指定精度模式
• --sequence-length：设置序列长度
• --max-global-attention-num：设置全局token最大数量
• --batch-size：设置批处理大小
• --repeat-test-num：设置测试重复次数

性能对比

测试环境

• 硬件：8×A100-80GB GPU
• 模型配置：
  • 头数：16
  • 头维度：64
  • 中间层维度：4096
  • 局部注意力窗口：512
  • 层数：24

FP32模式性能

| 批大小 | 序列长度 | HuggingFace时延(ms) | FasterTransformer时延(ms) | 加速比 | |--------|----------|----------------------|---------------------------|--------| | 1 | 1024 | 78 | 22 | 3.54× | | 8 | 2048 | 507 | 358 | 1.4× | | 32 | 4096 | 3914 | 2433 | 1.61× |

FP16模式性能

| 批大小 | 序列长度 | HuggingFace时延(ms) | FasterTransformer时延(ms) | 加速比 | |--------|----------|----------------------|---------------------------|--------| | 1 | 1024 | 80 | 14 | 5.71× | | 8 | 4096 | 1000 | 369 | 2.71× | | 32 | 2048 | 1953 | 820 | 2.38× |

从测试数据可以看出，FasterTransformer实现相比原始HuggingFace实现有显著性能提升，特别是在FP16模式下，最大可获得5.7倍的加速。

高级用法

PyTorch OP集成

FasterTransformer提供了与HuggingFace兼容的PyTorch接口，可以无缝替换原有实现：

# 初始化FT Longformer编码器
ft_encoder = FTLongformerEncoder(weights_file, layer_num, head_num, size_per_head,
                               intermediate_size, local_attn_window_size,
                               max_global_token_num, batch_size, seq_len,
                               attn_scaler, ft_longformer_lib)

# 替换HuggingFace编码器
ft_longformer = build_hf_longformer(hf_model_dir)
ft_encoder.set_hf_plugin_mode(True)
ft_longformer.longformer.encoder = ft_encoder

直接推理接口

# 直接调用FT编码器
output = ft_encoder.forward(input_b, local_attn_mask_b, global_attn_mask_b)

# 或使用HuggingFace风格接口
outputs = ft_longformer(input_ids_b,
                       attention_mask=local_attn_mask_b,
                       global_attention_mask=global_attn_mask_b)

应用建议

1. 长文档处理：适合处理法律文档、科研论文等长文本
2. 问答系统：在TriviaQA等长文档问答任务中表现优异
3. 信息检索：可应用于文档检索和关键信息提取
4. 文本摘要：适合生成长文档的摘要

总结

FasterTransformer中的Longformer实现通过高效的CUDA和C++优化，为长序列处理提供了显著的性能提升。其与HuggingFace兼容的接口设计使得现有项目可以轻松集成，特别适合需要处理长文档的实际应用场景。通过利用混合精度计算和优化的注意力机制，该实现在保持模型精度的同时大幅提高了推理速度。

FasterTransformer Transformer related optimization, including BERT, GPT 项目地址: https://gitcode.com/gh_mirrors/fa/FasterTransformer