FasterTransformer Longformer 推理加速实践:PyTorch 集成、输入约束与端到端性能解析
推理引擎算子库大模型【免费下载链接】FasterTransformerTransformer related optimization, including BERT, GPT项目地址https://gitcode.com/gh_mirrors/fa/FasterTransformer点击查看免费下载Longformer 通过局部滑动窗口注意力 任务驱动的全局注意力将自注意力的复杂度从 O(n²) 降为 O(n)是长文档建模数千 token 甚至更长的主流方案之一。本文以 docs/longformer_guide.md 为核心完整梳理 FasterTransformer 中 Longformer Encoder 的实现范围、输入约束、计算流程给出基于 PyTorch OP 的 QA 推理 Demo 全流程环境搭建、构建、运行与参数说明并结合 src/fastertransformer/models/longformer/LongformerEncoder.cc 等源码与官方 FP32/FP16 基准解析其加速原理与实战要点。读完本文你将能够在 FasterTransformer 仓库内复现 Longformer QA 加速推理理解其与 HuggingFace 实现的替换方式及性能差异。Longformer 模型与 FasterTransformer 的实现定位Longformer 出自论文《Longformer: The Long-Document Transformer》。标准 Transformer 的自注意力操作复杂度随序列长度呈平方增长导致无法处理长序列Longformer 引入一种随序列长度线性扩展的注意力机制易于处理数千 token 乃至更长的文档。其注意力机制是标准自注意力的即插即用替代品组合了局部窗口注意力local windowed attention与任务驱动的全局注意力task motivated global attention在字符级语言建模text8、enwik8以及 WikiHop、TriviaQA 等长文档任务上取得了当时领先的结果。FasterTransformer 中的 Longformer 实现与 HuggingFace 的 Longformer 对齐longformer-large-4096-finetuned-triviaqa等模型权重可直接加载。从仓库支持矩阵见 README.md看Longformer 在 FasterTransformer 中仅在PyTorch框架下提供FP16支持底层 CUDA/C 实现的精度还包括 FP32 与 BF16见下文。从 src/fastertransformer/models/longformer/LongformerEncoder.h 的模板实例化可知LongformerEncoder支持float、half以及在ENABLE_BF16编译开关下__nv_bfloat16三种数据类型template class LongformerEncoderfloat; template class LongformerEncoderhalf; #ifdef ENABLE_BF16 template class LongformerEncoder__nv_bfloat16; #endif目前仓库中只实现了 Longformer 的 Encoder 部分并配套一个真实场景的问答Question and Answering示例来同时演示正确性与加速效果。LongformerEncoder以 CUDA 和 C 实现并通过 PyTorch 自定义 OP 封装src/fastertransformer/th_op/longformer/LongformerEncoderOp.h供用户平滑集成。使用前提四条硬性约束在动手前必须先满足原文档列出的 4 条约束否则运行会失败或结果错误seq_len必须 2 * local_attn_window_size且seq_len % local_attn_window_size 0。这条约束在 LongformerAttentionLayer.cc 中通过FT_CHECK(seq_len % local_attn_window_size_ 0)强制校验在 model.py 的 Python 侧也有对应断言。head_dim即size_per_head目前只支持64。同样在 LongformerAttentionLayer.cc 中由FT_CHECK(size_per_head_ 64)保证。max_global_token_num应 local_attn_window_size。注意源码侧的实际检查为FT_CHECK(max_global_token_num_ local_attn_window_size_)见 LongformerAttentionLayer.cc即小于等于也能通过校验原文档建议严格小于以留出安全余量。输入序列中全局 tokenglobal_tokens必须连续地放在序列开头。这是 FasterTransformer 侧实现的前提假设与 HuggingFace 中任意位置标记全局 token不同使用时需要自行把全局 token 移动到序列头部。Longformer Encoder 计算流程原文档给出的计算流程如下图 1Longformer 计算流程图来自 docs/longformer_guide.md从源码层面可以还原该流程的每一步。LongformerEncoder::forwardLongformerEncoder.cc接收三类输入input形状batch_size x seq_len x in_dimlocal_attn_mask形状batch_size x seq_len0.0表示不做局部注意力1.0表示做局部注意力global_attn_mask形状batch_size x seq_len-10000.0表示非全局 token0.0表示全局 token。前向过程分两个预处理 kernel 与逐层循环预处理longformer_kernels.cuinvokeInitLongformerIdx基于全局注意力掩码通过cub::DevicePartition::Flagged提取每个 batch 中全局 token 的索引global_idx与数量global_token_nums供后续全局注意力按索引取值invokeLocalAttnMaskShift将局部注意力掩码经fma(mask, 10000.f, -10000.f)变换为注意力打分可加的偏移量对应 Python 侧extended_mask_b (global_attn_mask_b dummy_local_attn_mask_b) * 10000. - 10000.的做法。每个 Transformer 层forwardLayerLongformerEncoder.cc用cublasMMWrapper::Gemm计算 Q、K、V以及全局路径的kgglobal K、vgglobal V五个投影全局 Qqg因只对max_global_token_num_个 token 计算用stridedBatchedGemm单独处理invokeAddBiasTransposeToMultiHead把 Q/K/V/global-K/global-V 五组结果加 bias 并转置为多头的[batch, head, seq, dim]布局LongformerAttentionLayer::forward完成局部滑动窗口注意力序列被切成 head/tail 与 middle 三段对应不同 stride 的窗口 GEMM见 LongformerAttentionLayer.cc 中的buf_sizes/buf_strides设计与全局注意力invokeLongformerMHASoftmaxlongformer_kernels.cu 中的longformerMHASoftmaxKernel对局部 全局注意力打分做带 mask 的 softmax局部 token 只 softmax 窗口内及其前方的全局 token全局 token 则对全序列 softmax且当mask_blk[row_idx] ! 0时按 Python 实现语义整体置零invokeTransposeMultiHeadToSingle把多头注意力输出还原再经attention_output_weights线性层、AddBiasResidualLayerNormGeluFfnLayer完成 FFNintermediate GEMM bias GELU output GEMM最后再一次AddBiasResidualLayerNorm输出该层结果。整体呈现出局部窗口 GEMM 复用共享内存缓冲 全局 token 索引化 gather 融合 softmax的 CUDA 优化形态这也是相对 HuggingFace 朴素实现获得加速的核心所在。Demo以 PyTorch OP 运行 Longformer QA仓库提供了问答示例脚本 examples/pytorch/longformer/longformer_qa.py。它会先用随机输入对比 HuggingFace 与 FasterTransformer 的 Longformer Encoder 延迟再分别解码输出答案并打印用于同时验证性能与正确性。环境要求CMake 3.13用于编译 PyTorch 扩展CUDA 11.0 或更高版本Python 3仅验证过 Python 3PyTorch验证版本 1.8.0 1.5.0应可工作。推荐直接使用 NVIDIA NGC 的 PyTorch Docker 镜像原文档示例使用nvcr.io/nvidia/pytorch:22.09-py3并确保宿主机具备基于 Pascal、Volta、Turing 或 Ampere 架构的 GPU。环境搭建启动 Docker 容器并挂载项目目录docker run \ -it \ --shm-size 5g \ --rm \ --ipchost \ --gpusall \ -v {YOUR_FASTER_TRANSFORMER_PROJECT_DIR_ON_HOST}:/workspace/FasterTransformer \ --workdir /workspace/FasterTransformer \ nvcr.io/nvidia/pytorch:22.09-py3 bash也可以换成其他带 CUDA 的 PyTorch NGC 容器但需满足上述版本要求。安装 HuggingFace Transformers指定版本以对齐依赖pip install transformers4.9.2以 PyTorch 模式构建 FasterTransformermkdir build cd build cmake -DCMAKE_BUILD_TYPERelease -DBUILD_PYTON .. make -j12下载 HuggingFace Longformer QA 模型需 git-lfsapt-get update apt-get install git-lfs git lfs install git clone https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa ../examples/pytorch/longformer/longformer-large-4096-finetuned-triviaqa运行cd /workspace/FasterTransformer python3 examples/pytorch/longformer/longformer_qa.py \ --ft-longformer-lib build/lib/libth_transformer.so \ --model-dir examples/pytorch/longformer/longformer-large-4096-finetuned-triviaqa \ --passage Jim Henson was a nice puppet \ --question Who was Jim Henson? \ --repeat-test-num 50脚本会先对比 HuggingFace PyTorch 与 FasterTransformer 的 Longformer Encoder 性能随后打印问答答案。更多命令行参数以python3 longformer_qa.py -h查看核心参数如下参数说明--data_type fp16/bf16指定 FP16 或 BF16 模式即以 FP16/BF16 输入并产出 FP16/BF16 输出HuggingFace 侧同样切换对应精度。注意脚本中对 HF 的 BF16 会回退到 FP32HF Longformer Encoder 不支持 BF16--sequence-length选择序列长度。必须满足sequence_length 2 * local_attention_window_size且sequence_length % local_attention_window_size 0--max-global-attention-num全局 token 数量上限。实际全局 token 数不能超过该值FT 仅支持全局 token 位于序列开头。QA 示例中问题question的所有 token 会被放在序列开头并标记为全局 token--batch-size批大小。QA 示例中会把同一份 question passage 复制batch_size份堆叠成 batch仅用于性能测试--repeat-test-num推理对比的重复运行次数用于计算平均延迟脚本其余参数还包括--passage/--passage-file正文文本二选一与--model-dir、--ft-longformer-lib等见 longformer_qa.py。FasterTransformer Longformer Encoder PyTorch OPPyTorch 侧的封装类FTLongformerEncoder定义在 model.py它通过torch.classes.load_library(ft_longformer_lib)加载libth_transformer.so再实例化torch.classes.FasterTransformer.LongformerEncoder其底层对应 LongformerEncoderOp.h 中的FasterTransformerLongformerEncoder。模型权重由from_hf_longformer_weight_to_ft从 HuggingFace 的pytorch_model.bin中按层抽取 q/k/v、global q/k/v、attention output、两处 LayerNorm、intermediate/output 等全部参数并转置、展平、拼接FP32/FP16/BF16 按需转换形成连续权重张量——这也正是 LongformerEncoderOp.h 中q k v kg vg 权重与 bias 必须连续存放的要求。构建 FT Longformer Encoder# 传入必要配置与参数 weights_file os.path.join(hf_model_dir, pytorch_model.bin) 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 Longformer Encoder ft_longformer build_hf_longformer(hf_model_dir) ft_encoder.set_hf_plugin_mode(True) ft_longformer.longformer.encoder ft_encoder推理# 方式一直接向 FT Encoder 传准备好的输入 ft_encoder.set_hf_plugin_mode(False) output ft_encoder.forward(input_b, local_attn_mask_b, global_attn_mask_b) # local_attn_mask_b 与 global_attn_mask_b 需与输入等长1.0 表示有效注意力0.0 表示无注意力 # 方式二按 HuggingFace Longformer 的常规用法直接调用 outputs ft_longformer(input_ids_b, attention_masklocal_attn_mask_b, global_attention_maskglobal_attn_mask_b)hf_plugin_mode为True时forward会把 HuggingFace 传入的扩展注意力掩码extended_attention_mask转换回 FT 需要的local_attn_mask与global_attn_mask并返回LongformerBaseModelOutput结构从而对 HuggingFace 代码保持透明见 model.py。性能基准FP32 与 FP16官方在 A100 上对比了 HuggingFace Longformer Encoder不含 embedding 与最终qa_output线性层与 FT Longformer Encoder 的延迟。硬件环境为 8×A100-80GBmclk 1593MHzpclk 1410MHz AMD EPYC 7742 64 核处理器。模型超参数head_num16、size_per_head64、intermediate_size4096、local_attn_window_size512、num_layers24。FP32Batch_sizeSeq_lenPrecisionHuggingFace Latency (ms)FasterTransformers Latency (ms)Speedup11024FP3278223.5412048FP32104531.9614096FP32160841.9081024FP322731222.2382048FP325073581.484096FP329886071.62321024FP329784692.08322048FP32195514071.39324096FP32391424331.61FP16Batch_sizeSeq_lenPrecisionHuggingFace Latency (ms)FasterTransformers Latency (ms)Speedup11024FP1680145.7112048FP16106313.414096FP16157503.481024FP16276793.4982048FP165182082.4984096FP1610003692.71321024FP169953053.26322048FP1619538202.38324096FP1619668232.39数据解读基于上表事实FP16 在 batch1、seq1024 时加速比最高达 5.71×FP32 在相同配置下为 3.54×。整体上小 batch、短序列的加速更明显大 batch 下受 GEMM 效率与访存带宽影响加速比有所回落但所有测试配置均快于 HuggingFace 基线。该数据为官方原文档在特定软硬件环境下的测试结果实际部署中应根据自身 GPU 型号与序列长度重新基准。小结FasterTransformer 的 Longformer 实现把长文档 Transformer 中最昂贵的自注意力拆解为局部滑动窗口 GEMM 全局 token 索引化注意力并通过融合 softmax、bias-transpose 等 CUDA kernel 与 PyTorch OP 封装实现了对 HuggingFace Longformer Encoder 的透明替换与显著加速。使用时的核心要点是遵守四条输入约束序列长度与窗口的整除关系、head_dim64、全局 token 数量上限、全局 token 置于序列开头并按照构建 lib → 权重转换 → 替换 encoder的流程集成。需要进一步探索时可继续阅读 LongformerAttentionLayer.cc 的窗口 GEMM 缓冲布局或参考仓库通用的模型接入模板 templates/adding_a_new_model/README.md。赞分享推理引擎算子库大模型【免费下载链接】FasterTransformerTransformer related optimization, including BERT, GPT项目地址https://gitcode.com/gh_mirrors/fa/FasterTransformer点击查看免费下载相关推荐FasterTransformer BART 推理指南架构、优化原理与 PyTorch 端到端部署FasterTransformer BART 推理指南架构、优化原理与 PyTorch 端到端部署 导读 本文以 FasterTransformer 仓库中的推理引擎算子库大模型FasterTransformer中的Longformer实现详解FasterTransformer中的Longformer实现详解 概述 FasterTransformer项目中的Longformer实现是针对长序列处理场景推理引擎算子库大模型终极文档解析指南如何用AnythingLLM打破格式壁垒构建智能知识库终极文档解析指南如何用AnythingLLM打破格式壁垒构建智能知识库 你是否曾被海量文档淹没PDF报告、Word文档、Excel表格、Markdown笔人工智能AI 应用RAGAI Agent后端前端上一篇Android权限请求终极指南深入解析PermissionRequest生命周期管理下一篇 10分钟极速上手v3-admin-viteVue3后台框架新革命创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考