尧图网页设计 
DeepSpeed Sparse Attention 深度解析块稀疏注意力如何带来 10 倍序列长度与最高 6.3 倍训练加速【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本篇文章系统讲解 DeepSpeed 的块稀疏注意力Sparse Attention技术它通过O(wn)1 w n量级的块稀疏计算把标准 Transformer 自注意力O(n²)的计算与显存压力降为与序列长度近似线性相关从而在 BERT 预训练实验中支撑比稠密注意力长1016 倍的输入序列、获得最高6.3 倍的执行加速并保持可比甚至更高的精度。读完本文你将掌握 DeepSpeed Sparse Attention 的稀疏结构配置Fixed / BigBird / BSLongformer / Variable、基于 Triton 的 MatMul 与 Softmax 内核的调用关系、在任意 Transformer 模型中替换自注意力层的完整步骤以及如何扩展自定义块稀疏模式。本仓库中与本文配套的资料包括官方公告、技术细节博客、使用教程 以及位于 deepspeed/ops/sparse_attention 的完整实现。为什么需要稀疏注意力自注意力的O(n²)瓶颈以 Transformer 为代表的注意力模型能够捕捉输入序列中相距很远 token 之间的依赖关系因此广泛用于文本、图像与语音输入这类输入的特征就是序列长度可达数千 token。然而在实际训练中长序列输入的应用被注意力计算的计算量与显存需求所限制——这两者都随序列长度n以二次方O(n²)增长。DeepSpeed Sparse Attention 的出发点正是解决这一瓶颈它提供一套稀疏注意力内核套件通过块稀疏计算block-sparse computation把注意力计算的计算量和显存需求降低一个甚至多个数量级。它不仅能缓解注意力计算的显存瓶颈还能让稀疏计算本身保持高效执行。核心思想块稀疏、局部 全局 随机三种模式自由组合DeepSpeed Sparse Attention 的设计哲学是不给稀疏结构设限而是让**块block**成为基本计算单元稀疏结构可以设计为只计算邻近 token 之间的局部注意力local attention可以通过局部注意力生成的摘要 token 实现全局注意力global attention也可以引入随机注意力random attention更可以将局部、全局与随机注意力任意组合。由于只计算被选中的块注意力层的显存占用从O(n²)下降为O(wn)其中1 w nw的具体取值取决于所采用的注意力结构。下图中的蓝色、橙色与绿色块分别代表局部、全局与随机注意力位置直观展示了这种可变稀疏结构在具体实现上这套库基于 PyTorch其内核通过Triton平台开发而非手写 CUDA从源码层面为未来支持 CPU / OpenCL / Vulkan 保留了可能性对应源码可见 matmul.py、softmax.py 中的triton.jit内核以及 trsrc 下的.tr源文件。Sparse Attention 是 DeepSpeed 的扩展模块既可以随 DeepSpeed 一起使用也可以脱离 DeepSpeed 独立使用。前向与反向的计算形态块稀疏计算在前向与反向传播中表现为不同的稀疏/稠密矩阵混合形态其中S代表块稀疏矩阵block-sparse matrixD代表稠密矩阵dense matrix前向传播中注意力得分由两个稠密矩阵Query 与 Key 的转置做块稀疏乘得到块稀疏结果Softmax 在该块稀疏矩阵上归一化后再与稠密的 Value 相乘得到稠密的注意力上下文输出。架构解析模块划分与调用关系DeepSpeed Sparse Attention 在 deepspeed/ops/sparse_attention 目录下以清晰的分层结构组织从底层内核到高层模块依次为模块作用源码位置MatMul块稀疏矩阵乘法支持sdd稀疏稠密×稠密、dsd稠密稀疏×稠密、dds稠密稠密×稀疏三种模式matmul.pySoftmax块稀疏 Softmax同时覆盖前向与反向可融合应用 scale、相对位置编码、Key padding mask、attention masksoftmax.pySparseSelfAttention由MatMulSoftmax组装而成给定 Query/Key/Value 直接产出 Context Layer 输出sparse_self_attention.pyBertSparseSelfAttention简化版BertSelfAttention可直接替换原稠密 BERT 自注意力层bert_sparse_self_attention.pySparseAttentionUtils预训练模型适配工具替换注意力层、扩展位置嵌入、padding/unpadding 等sparse_attention_utils.pySparsityConfig及子类稀疏结构抽象类每种结构实现自己的make_layout构造稀疏布局sparsity_config.py三种 MatMul 模式如何拼装注意力从 sparse_self_attention.py#L64-L80 可以看到前向计算本质上是三类稀疏乘法的流水线sparse_dot_sdd_nt MatMul(sparsity_layout, self.sparsity_config.block, sdd, trans_aFalse, trans_bTrue) sparse_dot_dsd_nn MatMul(sparsity_layout, self.sparsity_config.block, dsd, trans_aFalse, trans_bFalse) sparse_softmax Softmax(sparsity_layout, self.sparsity_config.block)实际计算中SparseSelfAttention.forwardsparse_self_attention.py#L97-L146依次执行# 1) 注意力得分Q(稠密) x K^T(稠密) - 块稀疏得分矩阵sdd attn_output_weights sparse_dot_sdd_nt(query, key) # 2) 块稀疏 Softmax内部融合 scale、rpe、key_padding_mask、attn_mask attn_output_weights sparse_softmax(attn_output_weights, scalescaling, rperpe, key_padding_maskkey_padding_mask, attn_maskattn_mask, key_padding_mask_modeself.key_padding_mask_mode, attn_mask_modeself.attn_mask_mode) # 3) 加权求和attn_probs(稀疏) x V(稠密) - 稠密上下文dsd attn_output sparse_dot_dsd_nn(attn_output_weights, value)其中scaling head_dim ** -0.5对应标准多头注意力中的1/sqrt(d_k)缩放。实现还通过master_layout注册为 buffer 的稀疏布局张量与按序列长度L的算子缓存SparseSelfAttention.ops避免重复构建布局与查找表。另外两点值得注意序列长度必须是 block 的整数倍否则get_layout会抛出ValueErrorsparse_self_attention.py#L56-L58当前实现仅支持 fp16 训练forward中显式断言query.dtype torch.halfsparse_self_attention.py#L112。底层的MatMul会依据稀疏布局生成 LUTlook-up table与加载均衡信息Triton 内核按布局头信息调度非零块的乘累加并在需要拆分归约时使用 spin-lock 累加部分结果见 matmul.pySoftmax的前向内核同样按 LUT 只对非零块执行tl.softmax并在同一个内核内完成 scale、相对位置编码与两种 mask 的融合softmax.py#L36-L87。这些内核的数值正确性由 tests/unit/ops/sparse_attention/test_sparse_attention.py 通过与稠密参考实现sparse_to_dense/dense_to_sparse转换后与torch.softmax、稠密 matmul 对比验证fp16 下容差为rtol3e-2, atol2e-3。支持哪些稀疏结构SparsityConfig 参数全解SparsityConfigsparsity_config.py#L10-L60是所有稀疏结构的抽象父类定义了三个共享参数并要求每个子类实现make_layout(seq_len)来构造布局num_heads层中注意力头数block块大小当前实现基于分块稀疏矩阵该参数定义Block × Block的方阵块边长默认 16different_layout_per_head是否允许每个头使用不同的稀疏布局默认False此时会通过check_and_propagate_first_head_layout把第一个头的布局广播到所有头。make_layout生成的布局为(num_heads, num_blocks, num_blocks)的张量其中num_blocks seq_len // block同时基类校验seq_len必须能被block整除。下面是各类具体结构的参数与用途结构类JSON mode设计来源特有参数与含义FixedSparsityConfigfixedOpenAI Sparse Transformernum_local_blocks局部窗口内块数num_global_blocks每个局部窗口中充当全局代表的连续块数attentionunidirectional/bidirectionalhorizontal_global_attention全局代表块是否也横向关注所有块仅双向有效num_different_global_patterns全局代表在局部窗口中的不同取法版本数BSLongformerSparsityConfigbslongformerLongformer 的分块版本块级而非 token 级稀疏num_sliding_window_blocks滑动局部窗口的块数global_block_indices作为全局注意力的块索引列表可为每个全局窗口的起始索引global_block_end_indices各全局窗口的结束索引不含须与起始列表等长BigBirdSparsityConfigbigbirdBigBirdnum_random_blocks每一块行随机关注的块数num_sliding_window_blocks滑动窗口块数num_global_blocks从索引 0 开始的全局块数VariableSparsityConfigvariable自定义模板num_random_blocks每行随机块数local_window_blocks各局部窗口块数组成的列表最后一个值作用于剩余窗口global_block_indices/global_block_end_indices全局窗口索引attention、horizontal_global_attention同上DenseSparsityConfig对照/测试用无额外参数布局全 1等价于稠密注意力LocalSlidingWindowSparsityConfig纯局部滑动窗口源码中额外提供num_sliding_window_blocks、attention以fixed模式为例下图展示了局部窗口与全局代表块的组织方式改编自 Sparse Transformer 原论文这些结构类内置了大量参数一致性校验见 sparsity_config.py可帮助你尽早发现配置错误num_local_blocks必须能被num_global_blocks整除Fixedattention只接受unidirectional与bidirectionalhorizontal_global_attention只对bidirectional有效当num_different_global_patterns 1时强制要求different_layout_per_headTrue且不能超过num_local_blocks // num_global_blocksFixedglobal_block_indices与global_block_end_indices长度必须相同且起始索引必须小于结束索引BSLongformer/Variable滑动窗口、随机块、全局块的数量不能超过序列总块数。在实现层面各类结构将局部local、滑动窗口sliding window、全局global与随机random布局分别以set_local_layout/set_sliding_window_layout/set_global_layout/set_random_layout等私有方法刻画最后在make_layout中按顺序叠加见 sparsity_config.py#L219-L235 的Fixed.make_layout。例如BigBird.make_layout依次叠加随机、滑动窗口与 ITC 形式的全局块unidirectional模式下则通过torch.tril屏蔽未来位置sparsity_config.py#L518-L520。实战在 DeepSpeed 中启用 Sparse Attention将稀疏注意力接入 BERT 这类预训练模型的完整链路分为以下几步以仓库教程 sparse-attention.md 为蓝本第一步在 JSON 配置中声明稀疏注意力在 DeepSpeed 启动所依赖的 JSON 配置文件中加入sparse_attention小节。下面是一个fixedFixedSparsityConfig模式的典型配置sparse_attention: { mode: fixed, block: 16, different_layout_per_head: true, num_local_blocks: 4, num_global_blocks: 1, attention: bidirectional, horizontal_global_attention: false, num_different_global_patterns: 4 }其中mode字段对应上一节各稀疏结构类fixed/bigbird/bslongformer/variable等其余键即该结构的构造参数。第二步用稀疏自注意力替换稠密注意力BertSparseSelfAttentionbert_sparse_self_attention.py复用了原BertSelfAttention的query/key/value线性投影与 head 拆分逻辑只是将中间的稠密计算替换为稀疏版本。在模型构造函数中条件创建即可if sparse_attention_config is not None: from deepspeed.ops.sparse_attention import BertSparseSelfAttention layer.attention.self BertSparseSelfAttention( config, sparsity_configsparse_attention_config)其中sparse_attention_config即根据上文 JSON 解析出的SparsityConfig子类实例。对任意注意力模块替换方式等价于把下面这段稠密计算attention_scores torch.matmul(query_layer, key_layer) attention_scores attention_scores / math.sqrt(self.attention_head_size) attention_scores attention_scores attention_mask attention_probs self.softmax(attention_scores) attention_probs self.dropout(attention_probs) context_layer torch.matmul(attention_probs, value_layer)替换为一次稀疏自注意力调用sparse_self_attention.py#L62-L70 的调用形态context_layer self.sparse_self_attention( query_layer, key_layer, value_layer, key_padding_maskattention_mask)第三步padding / unpadding 输入与输出由于稀疏布局按块组织输入序列长度必须是block的整数倍。SparseAttentionUtilssparse_attention_utils.py提供了一组配套工具pad_to_block_size(block_size, input_ids, attention_mask, ...)把 token 序列、mask 等在序列维度补齐到block的整数倍并返回实际pad_lenunpad_sequence_output(pad_len, sequence_output)在编码器输出后裁掉补齐部分extend_position_embedding(model, max_position)把预训练权重的位置嵌入扩展到更长序列——例如把 128 长度的嵌入扩展到 1k 时会通过复制现有嵌入 8 次来初始化新嵌入。实验表明这种初始化比从零开始收敛更快BERT 直接repeatRoBERTa 因保留 0/1 两个位置需特殊处理update_tokenizer_model_max_length(tokenizer, max_position)同步更新 tokenizer 的最大长度replace_model_self_attention_with_sparse_self_attention(model, ...)加载好模型后一键替换其中所有自注意力层当前支持 BERT / RoBERTa 结构其他模型需仿照扩展。padding 在模型前向入口处进行、unpadding 在输出返回前进行pad_len, input_ids, attention_mask, ... SparseAttentionUtils.pad_to_block_size( block_sizeself.sparse_attention_config.block, input_idsinput_ids, attention_maskextended_attention_mask, ...) # 编码器输出后还原原始序列长度 if pad_len 0: encoded_layers[-1] SparseAttentionUtils.unpad_sequence_output(pad_len, encoded_layers[-1])第四步通过 DeepSpeed launcher 开启在启动训练脚本的命令行参数中加上--deepspeed_sparse_attention开关即可让 DeepSpeed 引擎识别上面的稀疏配置。脱离 DeepSpeed单独使用内核模块由于 Sparse Attention 是独立于 DeepSpeed 主引擎的模块级实现也可以不经过 DeepSpeed 启动器直接使用。安装包含 sparse attention 的 DeepSpeed 后直接导入模块即可例如from deepspeed.ops.sparse_attention import SparseSelfAttention进一步地MatMul与Softmax这两个基础内核也可以分别单独使用。例如构建一个稀疏 Softmaxfrom deepspeed.ops.sparse_attention import Softmax from deepspeed.ops.sparse_attention import FixedSparsityConfig config FixedSparsityConfig(num_heads12, block16) layout config.make_layout(seq_len2048) # (num_heads, 128, 128) 的 0/1 布局 softmax Softmax(layout, block16)Softmax.__call__支持在归一化时融合以下操作见 softmax.py#L260-L296scale归一化前对得分乘的缩放因子SparseSelfAttention中传入head_dim ** -0.5rpe相对位置编码张量加法融合key_padding_mask形状为(BatchSize, SequenceLength)可通过key_padding_mask_modeadd/mul选择加性或乘性融合attn_mask形状为(SequenceLength, SequenceLength)的 2D mask通过attn_mask_modeadd/mul控制融合方式。MatMul.__call__matmul.py#L733-L776则要求稠密输入内维与布局对应维一致且稀疏张量形状为(nnz_blocks, block, block)使用torch.autocast时输入会被自动转成 fp16。扩展你自己的块稀疏结构由于MatMul与Softmax内核接受任意块稀疏布局任何块级别的稀疏模式都可以直接应用。要新增一种稀疏结构只需仿照现有类扩展SparsityConfig并实现make_layoutfrom deepspeed.ops.sparse_attention import SparsityConfig class MySparsityConfig(SparsityConfig): def __init__(self, num_heads, block16, ...): super().__init__(num_heads, block, different_layout_per_headFalse) # 在此声明你的额外参数并做合法性校验 def make_layout(self, seq_len): layout self.setup_layout(seq_len) # 生成全 0 的 (num_heads, num_blocks, num_blocks) # 按你的模式将需要参与计算的块位置置 1 return self.check_and_propagate_first_head_layout(layout)新增子类一旦完成make_layout即可通过 DeepSpeed JSON 配置或直接构造实例接入SparseSelfAttention/BertSparseSelfAttention无需改动任何内核代码。性能表现与实测数据以下数据来自仓库内发布的技术博客所记录的实验注意它们均为特定硬件与实验条件下的结果不应理解为普适性能承诺。支撑更长序列1016 倍在 BERT 预训练实验中对稠密dense、稠密 激活检查点dense activation checkpoint、稀疏 激活检查点SA activation checkpoint三种设置进行对比。在 batch size 为 1、单张 32 GB NVIDIA V100 GPU 上SA 相比稠密设置在BERT base 上支撑长 10 倍、BERT large 上长 16 倍的最大可运行序列长度训练提速最高 6.3 倍同一预训练实验在不同 batch size 与序列长度下持续训练 100 轮、取最后 30 轮平均耗时提升幅度随序列长度增大而增大BERT base 最高 6.3 倍、BERT large 最高 5.3 倍、Megatron GPT2 最高 6.1 倍。这与公告中约一个数量级更长序列、最高约 6 倍更快执行的表述一致。精度不降反升的场景与 Sparse Transformer、Longformer、BigBird 等研究一致稀疏注意力可获得与全注意力可比甚至更高的精度。在基于 BERT 的长文档理解任务序列长度 2048中实验观察到从零预训练时SA 相对稠密收敛更快且精度更高而从稠密 512 序列长度的 checkpoint 继续用 SA 训练则在时间与精度两方面表现更好。与当时 SOTA 稀疏实现 Longformer 的对比在同一稀疏结构Fixed与可比精度下训练MLM 预训练于 Wikitext103基于 RoBERTa-base checkpoint8×V100-SXM2局部窗口 512 时 SA 单轮迭代时间由 1.6280s 降至 1.1059s提速 1.47 倍且 BPC 更低1.6613 vs 1.6708推理BERT-Base序列长度 2048batch size 1替换 Bert Attention 后相对 Longformer 提速最高3.13 倍且窗口越小加速越有限局部窗口 256 为 2.29 倍、128 为 2.16 倍、64 为 1.5 倍、32 为 1.24 倍、16 为 1.23 倍。以上对比再次印证了公告中的论断稀疏实现不仅针对某种固定结构还以高效系统支持帮助研究者自由探索任意块稀疏模式。运行前提与注意事项GPU 环境按教程说明Sparse Attention 最初面向 NVIDIA V100/A100依赖 Torch 1.6 与 CUDA 10.1 / 10.2 / 11.0 / 11.1必须禁用 Transformer Kernels当前 DeepSpeed Transformer Kernels 不支持稀疏注意力两者不可同时开启见 sparse-attention.md 的注意提示序列长度与 dtype序列长度必须能被block整除当前前向仅支持 fp16源码层断言这两个限制均在 sparse_self_attention.py 中有显式校验依赖 Triton内核通过 Triton JIT 编译运行使用前需确保环境具备对应 Triton 支持相关依赖可参考 requirements-triton.txt环境兼容性内核与 DeepSpeed 加速器抽象耦合如matmul.py通过get_accelerator判断设备测试会依据SparseAttnBuilder的兼容性判断自动跳过不兼容环境见 test_sparse_attention.py#L18-L19。总结DeepSpeed Sparse Attention 以分块为单位重新组织了自注意力的计算图用MatMulsdd/dsd/dds与融合型Softmax两个 Triton 内核覆盖前向与反向的全部矩阵运算用SparsityConfig体系把 Fixed、BigBird、BSLongformer、Variable 等稀疏结构统一成局部 全局 随机的组合语言并提供BertSparseSelfAttention与SparseAttentionUtils让既有 BERT/RoBERTa 类模型可以快速接入。它的价值不只在于让长序列训练成为可能BLOG 记录的 BERT 实验中序列长度提升 1016 倍、训练提速最高 6.3 倍更在于把自定义稀疏结构的成本降到只需实现一个make_layout方法从而把稀疏注意力模式的探索权交还给了模型研究者。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 DeepSpeed Sparse Attention 深度解析：块稀疏注意力如何带来 10 倍序列长度与最高 6.3 倍训练加速
DeepSpeed Sparse Attention 深度解析块稀疏注意力如何带来 10 倍序列长度与最高 6.3 倍训练加速【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本篇文章系统讲解 DeepSpeed 的块稀疏注意力Sparse Attention技术它通过O(wn)1 w n量级的块稀疏计算把标准 Transformer 自注意力O(n²)的计算与显存压力降为与序列长度近似线性相关从而在 BERT 预训练实验中支撑比稠密注意力长1016 倍的输入序列、获得最高6.3 倍的执行加速并保持可比甚至更高的精度。读完本文你将掌握 DeepSpeed Sparse Attention 的稀疏结构配置Fixed / BigBird / BSLongformer / Variable、基于 Triton 的 MatMul 与 Softmax 内核的调用关系、在任意 Transformer 模型中替换自注意力层的完整步骤以及如何扩展自定义块稀疏模式。本仓库中与本文配套的资料包括官方公告、技术细节博客、使用教程 以及位于 deepspeed/ops/sparse_attention 的完整实现。为什么需要稀疏注意力自注意力的O(n²)瓶颈以 Transformer 为代表的注意力模型能够捕捉输入序列中相距很远 token 之间的依赖关系因此广泛用于文本、图像与语音输入这类输入的特征就是序列长度可达数千 token。然而在实际训练中长序列输入的应用被注意力计算的计算量与显存需求所限制——这两者都随序列长度n以二次方O(n²)增长。DeepSpeed Sparse Attention 的出发点正是解决这一瓶颈它提供一套稀疏注意力内核套件通过块稀疏计算block-sparse computation把注意力计算的计算量和显存需求降低一个甚至多个数量级。它不仅能缓解注意力计算的显存瓶颈还能让稀疏计算本身保持高效执行。核心思想块稀疏、局部 全局 随机三种模式自由组合DeepSpeed Sparse Attention 的设计哲学是不给稀疏结构设限而是让**块block**成为基本计算单元稀疏结构可以设计为只计算邻近 token 之间的局部注意力local attention可以通过局部注意力生成的摘要 token 实现全局注意力global attention也可以引入随机注意力random attention更可以将局部、全局与随机注意力任意组合。由于只计算被选中的块注意力层的显存占用从O(n²)下降为O(wn)其中1 w nw的具体取值取决于所采用的注意力结构。下图中的蓝色、橙色与绿色块分别代表局部、全局与随机注意力位置直观展示了这种可变稀疏结构在具体实现上这套库基于 PyTorch其内核通过Triton平台开发而非手写 CUDA从源码层面为未来支持 CPU / OpenCL / Vulkan 保留了可能性对应源码可见 matmul.py、softmax.py 中的triton.jit内核以及 trsrc 下的.tr源文件。Sparse Attention 是 DeepSpeed 的扩展模块既可以随 DeepSpeed 一起使用也可以脱离 DeepSpeed 独立使用。前向与反向的计算形态块稀疏计算在前向与反向传播中表现为不同的稀疏/稠密矩阵混合形态其中S代表块稀疏矩阵block-sparse matrixD代表稠密矩阵dense matrix前向传播中注意力得分由两个稠密矩阵Query 与 Key 的转置做块稀疏乘得到块稀疏结果Softmax 在该块稀疏矩阵上归一化后再与稠密的 Value 相乘得到稠密的注意力上下文输出。架构解析模块划分与调用关系DeepSpeed Sparse Attention 在 deepspeed/ops/sparse_attention 目录下以清晰的分层结构组织从底层内核到高层模块依次为模块作用源码位置MatMul块稀疏矩阵乘法支持sdd稀疏稠密×稠密、dsd稠密稀疏×稠密、dds稠密稠密×稀疏三种模式matmul.pySoftmax块稀疏 Softmax同时覆盖前向与反向可融合应用 scale、相对位置编码、Key padding mask、attention masksoftmax.pySparseSelfAttention由MatMulSoftmax组装而成给定 Query/Key/Value 直接产出 Context Layer 输出sparse_self_attention.pyBertSparseSelfAttention简化版BertSelfAttention可直接替换原稠密 BERT 自注意力层bert_sparse_self_attention.pySparseAttentionUtils预训练模型适配工具替换注意力层、扩展位置嵌入、padding/unpadding 等sparse_attention_utils.pySparsityConfig及子类稀疏结构抽象类每种结构实现自己的make_layout构造稀疏布局sparsity_config.py三种 MatMul 模式如何拼装注意力从 sparse_self_attention.py#L64-L80 可以看到前向计算本质上是三类稀疏乘法的流水线sparse_dot_sdd_nt MatMul(sparsity_layout, self.sparsity_config.block, sdd, trans_aFalse, trans_bTrue) sparse_dot_dsd_nn MatMul(sparsity_layout, self.sparsity_config.block, dsd, trans_aFalse, trans_bFalse) sparse_softmax Softmax(sparsity_layout, self.sparsity_config.block)实际计算中SparseSelfAttention.forwardsparse_self_attention.py#L97-L146依次执行# 1) 注意力得分Q(稠密) x K^T(稠密) - 块稀疏得分矩阵sdd attn_output_weights sparse_dot_sdd_nt(query, key) # 2) 块稀疏 Softmax内部融合 scale、rpe、key_padding_mask、attn_mask attn_output_weights sparse_softmax(attn_output_weights, scalescaling, rperpe, key_padding_maskkey_padding_mask, attn_maskattn_mask, key_padding_mask_modeself.key_padding_mask_mode, attn_mask_modeself.attn_mask_mode) # 3) 加权求和attn_probs(稀疏) x V(稠密) - 稠密上下文dsd attn_output sparse_dot_dsd_nn(attn_output_weights, value)其中scaling head_dim ** -0.5对应标准多头注意力中的1/sqrt(d_k)缩放。实现还通过master_layout注册为 buffer 的稀疏布局张量与按序列长度L的算子缓存SparseSelfAttention.ops避免重复构建布局与查找表。另外两点值得注意序列长度必须是 block 的整数倍否则get_layout会抛出ValueErrorsparse_self_attention.py#L56-L58当前实现仅支持 fp16 训练forward中显式断言query.dtype torch.halfsparse_self_attention.py#L112。底层的MatMul会依据稀疏布局生成 LUTlook-up table与加载均衡信息Triton 内核按布局头信息调度非零块的乘累加并在需要拆分归约时使用 spin-lock 累加部分结果见 matmul.pySoftmax的前向内核同样按 LUT 只对非零块执行tl.softmax并在同一个内核内完成 scale、相对位置编码与两种 mask 的融合softmax.py#L36-L87。这些内核的数值正确性由 tests/unit/ops/sparse_attention/test_sparse_attention.py 通过与稠密参考实现sparse_to_dense/dense_to_sparse转换后与torch.softmax、稠密 matmul 对比验证fp16 下容差为rtol3e-2, atol2e-3。支持哪些稀疏结构SparsityConfig 参数全解SparsityConfigsparsity_config.py#L10-L60是所有稀疏结构的抽象父类定义了三个共享参数并要求每个子类实现make_layout(seq_len)来构造布局num_heads层中注意力头数block块大小当前实现基于分块稀疏矩阵该参数定义Block × Block的方阵块边长默认 16different_layout_per_head是否允许每个头使用不同的稀疏布局默认False此时会通过check_and_propagate_first_head_layout把第一个头的布局广播到所有头。make_layout生成的布局为(num_heads, num_blocks, num_blocks)的张量其中num_blocks seq_len // block同时基类校验seq_len必须能被block整除。下面是各类具体结构的参数与用途结构类JSON mode设计来源特有参数与含义FixedSparsityConfigfixedOpenAI Sparse Transformernum_local_blocks局部窗口内块数num_global_blocks每个局部窗口中充当全局代表的连续块数attentionunidirectional/bidirectionalhorizontal_global_attention全局代表块是否也横向关注所有块仅双向有效num_different_global_patterns全局代表在局部窗口中的不同取法版本数BSLongformerSparsityConfigbslongformerLongformer 的分块版本块级而非 token 级稀疏num_sliding_window_blocks滑动局部窗口的块数global_block_indices作为全局注意力的块索引列表可为每个全局窗口的起始索引global_block_end_indices各全局窗口的结束索引不含须与起始列表等长BigBirdSparsityConfigbigbirdBigBirdnum_random_blocks每一块行随机关注的块数num_sliding_window_blocks滑动窗口块数num_global_blocks从索引 0 开始的全局块数VariableSparsityConfigvariable自定义模板num_random_blocks每行随机块数local_window_blocks各局部窗口块数组成的列表最后一个值作用于剩余窗口global_block_indices/global_block_end_indices全局窗口索引attention、horizontal_global_attention同上DenseSparsityConfig对照/测试用无额外参数布局全 1等价于稠密注意力LocalSlidingWindowSparsityConfig纯局部滑动窗口源码中额外提供num_sliding_window_blocks、attention以fixed模式为例下图展示了局部窗口与全局代表块的组织方式改编自 Sparse Transformer 原论文这些结构类内置了大量参数一致性校验见 sparsity_config.py可帮助你尽早发现配置错误num_local_blocks必须能被num_global_blocks整除Fixedattention只接受unidirectional与bidirectionalhorizontal_global_attention只对bidirectional有效当num_different_global_patterns 1时强制要求different_layout_per_headTrue且不能超过num_local_blocks // num_global_blocksFixedglobal_block_indices与global_block_end_indices长度必须相同且起始索引必须小于结束索引BSLongformer/Variable滑动窗口、随机块、全局块的数量不能超过序列总块数。在实现层面各类结构将局部local、滑动窗口sliding window、全局global与随机random布局分别以set_local_layout/set_sliding_window_layout/set_global_layout/set_random_layout等私有方法刻画最后在make_layout中按顺序叠加见 sparsity_config.py#L219-L235 的Fixed.make_layout。例如BigBird.make_layout依次叠加随机、滑动窗口与 ITC 形式的全局块unidirectional模式下则通过torch.tril屏蔽未来位置sparsity_config.py#L518-L520。实战在 DeepSpeed 中启用 Sparse Attention将稀疏注意力接入 BERT 这类预训练模型的完整链路分为以下几步以仓库教程 sparse-attention.md 为蓝本第一步在 JSON 配置中声明稀疏注意力在 DeepSpeed 启动所依赖的 JSON 配置文件中加入sparse_attention小节。下面是一个fixedFixedSparsityConfig模式的典型配置sparse_attention: { mode: fixed, block: 16, different_layout_per_head: true, num_local_blocks: 4, num_global_blocks: 1, attention: bidirectional, horizontal_global_attention: false, num_different_global_patterns: 4 }其中mode字段对应上一节各稀疏结构类fixed/bigbird/bslongformer/variable等其余键即该结构的构造参数。第二步用稀疏自注意力替换稠密注意力BertSparseSelfAttentionbert_sparse_self_attention.py复用了原BertSelfAttention的query/key/value线性投影与 head 拆分逻辑只是将中间的稠密计算替换为稀疏版本。在模型构造函数中条件创建即可if sparse_attention_config is not None: from deepspeed.ops.sparse_attention import BertSparseSelfAttention layer.attention.self BertSparseSelfAttention( config, sparsity_configsparse_attention_config)其中sparse_attention_config即根据上文 JSON 解析出的SparsityConfig子类实例。对任意注意力模块替换方式等价于把下面这段稠密计算attention_scores torch.matmul(query_layer, key_layer) attention_scores attention_scores / math.sqrt(self.attention_head_size) attention_scores attention_scores attention_mask attention_probs self.softmax(attention_scores) attention_probs self.dropout(attention_probs) context_layer torch.matmul(attention_probs, value_layer)替换为一次稀疏自注意力调用sparse_self_attention.py#L62-L70 的调用形态context_layer self.sparse_self_attention( query_layer, key_layer, value_layer, key_padding_maskattention_mask)第三步padding / unpadding 输入与输出由于稀疏布局按块组织输入序列长度必须是block的整数倍。SparseAttentionUtilssparse_attention_utils.py提供了一组配套工具pad_to_block_size(block_size, input_ids, attention_mask, ...)把 token 序列、mask 等在序列维度补齐到block的整数倍并返回实际pad_lenunpad_sequence_output(pad_len, sequence_output)在编码器输出后裁掉补齐部分extend_position_embedding(model, max_position)把预训练权重的位置嵌入扩展到更长序列——例如把 128 长度的嵌入扩展到 1k 时会通过复制现有嵌入 8 次来初始化新嵌入。实验表明这种初始化比从零开始收敛更快BERT 直接repeatRoBERTa 因保留 0/1 两个位置需特殊处理update_tokenizer_model_max_length(tokenizer, max_position)同步更新 tokenizer 的最大长度replace_model_self_attention_with_sparse_self_attention(model, ...)加载好模型后一键替换其中所有自注意力层当前支持 BERT / RoBERTa 结构其他模型需仿照扩展。padding 在模型前向入口处进行、unpadding 在输出返回前进行pad_len, input_ids, attention_mask, ... SparseAttentionUtils.pad_to_block_size( block_sizeself.sparse_attention_config.block, input_idsinput_ids, attention_maskextended_attention_mask, ...) # 编码器输出后还原原始序列长度 if pad_len 0: encoded_layers[-1] SparseAttentionUtils.unpad_sequence_output(pad_len, encoded_layers[-1])第四步通过 DeepSpeed launcher 开启在启动训练脚本的命令行参数中加上--deepspeed_sparse_attention开关即可让 DeepSpeed 引擎识别上面的稀疏配置。脱离 DeepSpeed单独使用内核模块由于 Sparse Attention 是独立于 DeepSpeed 主引擎的模块级实现也可以不经过 DeepSpeed 启动器直接使用。安装包含 sparse attention 的 DeepSpeed 后直接导入模块即可例如from deepspeed.ops.sparse_attention import SparseSelfAttention进一步地MatMul与Softmax这两个基础内核也可以分别单独使用。例如构建一个稀疏 Softmaxfrom deepspeed.ops.sparse_attention import Softmax from deepspeed.ops.sparse_attention import FixedSparsityConfig config FixedSparsityConfig(num_heads12, block16) layout config.make_layout(seq_len2048) # (num_heads, 128, 128) 的 0/1 布局 softmax Softmax(layout, block16)Softmax.__call__支持在归一化时融合以下操作见 softmax.py#L260-L296scale归一化前对得分乘的缩放因子SparseSelfAttention中传入head_dim ** -0.5rpe相对位置编码张量加法融合key_padding_mask形状为(BatchSize, SequenceLength)可通过key_padding_mask_modeadd/mul选择加性或乘性融合attn_mask形状为(SequenceLength, SequenceLength)的 2D mask通过attn_mask_modeadd/mul控制融合方式。MatMul.__call__matmul.py#L733-L776则要求稠密输入内维与布局对应维一致且稀疏张量形状为(nnz_blocks, block, block)使用torch.autocast时输入会被自动转成 fp16。扩展你自己的块稀疏结构由于MatMul与Softmax内核接受任意块稀疏布局任何块级别的稀疏模式都可以直接应用。要新增一种稀疏结构只需仿照现有类扩展SparsityConfig并实现make_layoutfrom deepspeed.ops.sparse_attention import SparsityConfig class MySparsityConfig(SparsityConfig): def __init__(self, num_heads, block16, ...): super().__init__(num_heads, block, different_layout_per_headFalse) # 在此声明你的额外参数并做合法性校验 def make_layout(self, seq_len): layout self.setup_layout(seq_len) # 生成全 0 的 (num_heads, num_blocks, num_blocks) # 按你的模式将需要参与计算的块位置置 1 return self.check_and_propagate_first_head_layout(layout)新增子类一旦完成make_layout即可通过 DeepSpeed JSON 配置或直接构造实例接入SparseSelfAttention/BertSparseSelfAttention无需改动任何内核代码。性能表现与实测数据以下数据来自仓库内发布的技术博客所记录的实验注意它们均为特定硬件与实验条件下的结果不应理解为普适性能承诺。支撑更长序列1016 倍在 BERT 预训练实验中对稠密dense、稠密 激活检查点dense activation checkpoint、稀疏 激活检查点SA activation checkpoint三种设置进行对比。在 batch size 为 1、单张 32 GB NVIDIA V100 GPU 上SA 相比稠密设置在BERT base 上支撑长 10 倍、BERT large 上长 16 倍的最大可运行序列长度训练提速最高 6.3 倍同一预训练实验在不同 batch size 与序列长度下持续训练 100 轮、取最后 30 轮平均耗时提升幅度随序列长度增大而增大BERT base 最高 6.3 倍、BERT large 最高 5.3 倍、Megatron GPT2 最高 6.1 倍。这与公告中约一个数量级更长序列、最高约 6 倍更快执行的表述一致。精度不降反升的场景与 Sparse Transformer、Longformer、BigBird 等研究一致稀疏注意力可获得与全注意力可比甚至更高的精度。在基于 BERT 的长文档理解任务序列长度 2048中实验观察到从零预训练时SA 相对稠密收敛更快且精度更高而从稠密 512 序列长度的 checkpoint 继续用 SA 训练则在时间与精度两方面表现更好。与当时 SOTA 稀疏实现 Longformer 的对比在同一稀疏结构Fixed与可比精度下训练MLM 预训练于 Wikitext103基于 RoBERTa-base checkpoint8×V100-SXM2局部窗口 512 时 SA 单轮迭代时间由 1.6280s 降至 1.1059s提速 1.47 倍且 BPC 更低1.6613 vs 1.6708推理BERT-Base序列长度 2048batch size 1替换 Bert Attention 后相对 Longformer 提速最高3.13 倍且窗口越小加速越有限局部窗口 256 为 2.29 倍、128 为 2.16 倍、64 为 1.5 倍、32 为 1.24 倍、16 为 1.23 倍。以上对比再次印证了公告中的论断稀疏实现不仅针对某种固定结构还以高效系统支持帮助研究者自由探索任意块稀疏模式。运行前提与注意事项GPU 环境按教程说明Sparse Attention 最初面向 NVIDIA V100/A100依赖 Torch 1.6 与 CUDA 10.1 / 10.2 / 11.0 / 11.1必须禁用 Transformer Kernels当前 DeepSpeed Transformer Kernels 不支持稀疏注意力两者不可同时开启见 sparse-attention.md 的注意提示序列长度与 dtype序列长度必须能被block整除当前前向仅支持 fp16源码层断言这两个限制均在 sparse_self_attention.py 中有显式校验依赖 Triton内核通过 Triton JIT 编译运行使用前需确保环境具备对应 Triton 支持相关依赖可参考 requirements-triton.txt环境兼容性内核与 DeepSpeed 加速器抽象耦合如matmul.py通过get_accelerator判断设备测试会依据SparseAttnBuilder的兼容性判断自动跳过不兼容环境见 test_sparse_attention.py#L18-L19。总结DeepSpeed Sparse Attention 以分块为单位重新组织了自注意力的计算图用MatMulsdd/dsd/dds与融合型Softmax两个 Triton 内核覆盖前向与反向的全部矩阵运算用SparsityConfig体系把 Fixed、BigBird、BSLongformer、Variable 等稀疏结构统一成局部 全局 随机的组合语言并提供BertSparseSelfAttention与SparseAttentionUtils让既有 BERT/RoBERTa 类模型可以快速接入。它的价值不只在于让长序列训练成为可能BLOG 记录的 BERT 实验中序列长度提升 1016 倍、训练提速最高 6.3 倍更在于把自定义稀疏结构的成本降到只需实现一个make_layout方法从而把稀疏注意力模式的探索权交还给了模型研究者。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考