尧图网页设计 
DeepSpeed Sparse Attention 深度解析基于 Triton 的块稀疏注意力如何支撑超长序列 Transformer【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeedSparse AttentionSA是 DeepSpeed 为自注意力类模型提供的一套块稀疏block-sparse注意力内核与配置体系目标是把注意力计算随序列长度二次增长的复杂度与显存占用降下来从而支撑更长的输入序列训练与推理。本篇文章基于 DeepSpeed 仓库中 2020 年发布的 Sparse Attention 官方介绍与配套教程结合 deepspeed/ops/sparse_attention 下的实际源码展开讲解读完你可以掌握块稀疏注意力的原理与复杂度优势、SA 内置的多种稀疏结构Fixed / BSLongformer / BigBird / Variable及其全部配置参数、如何在 DeepSpeed 中启用 SA、如何单独复用其 MatMul / Softmax 内核以及如何为任意新稀疏模式编写自定义布局。一、从O(n^2)到O(wn)Sparse Attention 要解决的问题以 Transformer 为代表的注意力模型擅长捕捉长距离 token 之间的依赖关系被广泛用于文本、图像、语音等序列输入场景序列长度常常达到数千 token。然而在实际落地中自注意力的计算量与显存占用随序列长度n呈平方级增长O(n^2)这使得它在长序列上受到严重制约——这是本文要解决的第一个核心问题。DeepSpeed 给出的方案是一套Sparse Attention 内核套件sparse attention kernels通过块稀疏计算block-sparse computation把注意力计算的算力与显存需求降低数个量级。它不仅能缓解注意力计算的内存瓶颈还能让稀疏计算本身保持高效执行。整体而言稀疏注意力把显存占用从O(n^2)降为O(wn)其中1 w n是一个与具体注意力结构相关的参数。Sparse Attention 可以灵活组合多种注意力形态局部注意力local只关注邻近 token全局注意力global通过局部窗口内计算的摘要块summary tokens关注整个序列随机注意力random允许一定数量的随机块被关注。下图展示了三者可任意组合的稀疏结构示例蓝色为局部、橙色为全局、绿色为随机块从实现上看这套库基于 PyTorch其内核通过Triton开发而不是手写 CUDA这为未来支持 CPU/OpenCL/Vulkan 预留了空间详见下文源码佐证。库本身是 DeepSpeed 的扩展既可以跟随 DeepSpeed 使用也可以独立使用。二、技术实现概览块稀疏矩阵与 Triton 内核2.1 前向与反向的块稀疏计算DeepSpeed Sparse Attention 的核心是把注意力分数矩阵按块block切分只对稀疏布局中置 1 的块执行矩阵乘与 softmax。下图中S表示块稀疏矩阵block-sparse matrixD表示稠密矩阵前向大体分两步先在稀疏块上计算S D × D即 Query 与 Key 的注意力分数再对稀疏块做 softmax随后做D S × D注意力分数与 Value 相乘得到 context。反向过程同理只是梯度矩阵同样是稀疏的2.2 模块级源码佐证当前仓库中SA 的全部实现位于 deepspeed/ops/sparse_attention 目录模块文件职责MatMulmatmul.py块稀疏矩阵乘法支持sdd、dsd、dds三种模式Softmaxsoftmax.py在稀疏块上执行 softmax同时处理前向与反向SparseSelfAttentionsparse_self_attention.py组装 MatMul 与 Softmax由 Q/K/V 生成 context 层输出BertSparseSelfAttentionbert_sparse_self_attention.py稀疏版BertSelfAttention可直接替换原稠密层SparseAttentionUtilssparse_attention_utils.py位置编码扩展、tokenizer 更新、padding/unpadding 等适配工具稀疏布局配置sparsity_config.pySparsityConfig基类及全部具体稀疏结构在 sparse_self_attention.py 的forward中可以看到清晰的调用链sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax self.get_ops(num_heads, tgt_len) scaling float(head_dim)**-0.5 # attention scoressparse Dense × Dense attn_output_weights sparse_dot_sdd_nt(query, key) # 块稀疏 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) # outputsDense Sparse × Dense attn_output sparse_dot_dsd_nn(attn_output_weights, value)几个值得注意的实现细节均可从源码确认内核缓存SparseSelfAttention.ops以序列长度L为键缓存(MatMul-sdd, MatMul-dsd, Softmax)三件套get_ops 只在首次遇到该长度时创建内核布局广播稀疏布局作为master_layoutbuffer 注册若分布式环境已初始化且布局未同步会通过dist.broadcast(self.master_layout, src0)从 rank 0 广播到所有进程见源码约束forward中断言query.dtype torch.half即当前只支持 fp16 训练get_layout会校验序列长度必须能被block整除Triton 内核块稀疏 MatMul 的 JIT 内核_kernel直接用triton.jit编写于 matmul.py并通过 LUTlook-up table记录每个非零块的偏移与指针增量避免 GPU 上的分支发散Softmax 内核位于 softmax.py。该目录下的trsrc/子目录存放了对应的 Triton 源文件matmul.tr、softmax_fwd.tr、softmax_bwd.tr印证了内核基于 Triton 而非手写 CUDA的设计。加速器抽象内核通过from deepspeed.accelerator import get_accelerator判断设备类型matmul.py说明该模块已接入 DeepSpeed 的 accelerator 抽象层。三、核心组件逐个拆解3.1 MatMul三种块稀疏矩阵乘模式MatMul类matmul.py接受一个稀疏layout、block尺寸和mode其中mode决定了输入/输出谁是稀疏的mode含义在注意力中的角色sdd输出稀疏 Dense × Dense由 Q、K 生成稀疏注意力分数dsd输出 Dense Sparse × Dense注意力分数 × V 生成 contextdds输出 Dense Dense × Sparse反向传播中计算梯度时使用在反向传播_sparse_matmul.backward中sdd对应的da/db会分别重映射到dds/dsd等模式即梯度计算同样只发生在非零块上。内核会先根据 layout 构造 C、DA、DB 三套 LUT 并缓存make_lut供前向与两次梯度计算复用。3.2 Softmax块内归一化Softmax模块在稀疏块上执行缩放与归一化除常规的scale外还内置了三类掩码的处理用于兼容标准注意力模块rpe相对位置编码relative position embeddingkey_padding_mask对 padding 位屏蔽模式可为add或mulattn_mask自定义注意力掩码当前支持二维Sequence × Sequence模式同样可为add或mul。掩码模式通过SparseSelfAttention构造参数key_padding_mask_modeadd、attn_mask_modemul指定默认值见 sparse_self_attention.py。3.3 SparseAttentionUtils把预训练模型改造成 SA对于想直接改造 Hugging Face BERT / RoBERTa 等已加载模型的用户sparse_attention_utils.py 提供了一组开箱即用的静态方法方法作用replace_model_self_attention_with_sparse_self_attention把 BERT/RoBERTa 编码器中所有 self-attention 层替换为稀疏版保留原有 Q/K/V 权重见 replace_self_attention_layer_with_sparse_self_attention_layerextend_position_embedding将预训练位置编码按倍数复制以适配更长序列例如 128 → 1k 时复制 8 次实验表明这种初始化比从零开始收敛更快见 实现update_tokenizer_model_max_length同步更新 tokenizer 的最大长度pad_to_block_size把输入 token 与注意力掩码在序列维度补齐到 block 的整数倍SA 的硬性要求unpad_sequence_output模型输出前裁剪掉补齐部分四、稀疏结构配置从基类到五种内置模式SA 的稀疏布局由SparsityConfig及其子类描述。基类在 sparsity_config.py 中定义并持有三类共享参数与两个共享方法共享参数num_heads该层注意力头数量必填block块大小当前实现基于块稀疏矩阵每个参与计算的都是block × block的方阵默认16different_layout_per_head是否允许每个注意力头使用不同的稀疏布局默认False所有头共享一份布局。核心方法setup_layout(seq_len)先校验seq_len % block 0再生成(num_heads, num_blocks, num_blocks)的全零布局张量num_blocks seq_len // blockmake_layout(seq_len)由具体子类实现负责画出该稀疏结构的 0/1 矩阵check_and_propagate_first_head_layout当different_layout_per_headFalse时把 head 0 的布局复制给其余所有头。由于布局的构造完全由make_layout决定任何块稀疏注意力模式都可以通过继承SparsityConfig并实现这一个函数来表达。DeepSpeed 官方内置了五种结构下面逐一说明其参数含默认值均与 sparsity_config.py 中类签名一致。4.1 FixedFixedSparsityConfig源自 OpenAI 的 Sparse Transformer 论文Generative Modeling with Sparse Transformers由若干参数固定的局部窗口 全局块构成参数默认值说明num_heads必填注意力头数block16块大小different_layout_per_headFalse每个头是否用不同布局num_local_blocks4每个局部注意力窗口包含的块数窗口内 token 互相可见num_global_blocks1每个局部窗口中作为全局代表的连续块数attentionbidirectionalunidirectional自回归只关注上文注意力矩阵上三角为空或bidirectional如 BERT可双向互看horizontal_global_attentionFalse全局代表块是否也横向关注所有块仅bidirectional下有效num_different_global_patterns1不同全局布局的版本数。例如局部窗口 4 块、全局 1 块时可让第 1/2/3/4 块分别担任全局代表从而有 4 个版本对应约束在构造函数中即被校验源码num_local_blocks必须能被num_global_blocks整除attention只能取unidirectional/bidirectionalhorizontal_global_attention仅限双向注意力num_different_global_patterns 1时必须打开different_layout_per_head且布局版本数不能超过num_local_blocks // num_global_blocks。4.2 BSLongformerBSLongformerSparsityConfigLongformerThe Long-Document Transformer的块稀疏实现——原版 Longformer 是逐 tokenelement-wise稀疏BSLongformer 将其改为块级稀疏参数默认值说明num_sliding_window_blocks3滑动局部窗口包含的块数global_block_indices[0]被视为全局注意力的块索引列表global_block_end_indicesNone全局窗口的结束索引若设置则必须与global_block_indices等长此时对每个iglobal_block_indices[i]到global_block_end_indices[i]不含之间的块都视为全局attentionbidirectional同上4.3 BigBirdBigBirdSparsityConfig源自 Google 的 Big Bird 论文综合了 Fixed/Longformer 的思路并加入随机注意力参数默认值说明num_random_blocks1每行块中随机选取多少个块参与注意力num_sliding_window_blocks3滑动局部窗口块数num_global_blocks1从索引 0 开始的前多少个块作为全局块被所有块关注、也关注所有块attentionbidirectional同上实现上随机块在 set_random_layout 中逐行random.sample若为单向注意力则只在[0, row]范围内采样避免未来信息泄漏。4.4 VariableVariableSparsityConfigVariable是官方提供的模板用来表达局部 全局 随机的任意自定义组合同时支持可变长度的局部窗口参数默认值说明num_random_blocks0每行随机块数0 表示不含随机块local_window_blocks[4]局部窗口块数列表第 k 个数字描述第 k 个窗口最后一个数字会沿用到剩余窗口global_block_indices[0]全局块索引与 end 参数组合时可作为起始索引global_block_end_indicesNone全局块窗口结束索引attentionbidirectional同上horizontal_global_attentionFalse同上4.5 DenseDenseSparsityConfig严格说 Dense 不是稀疏结构——它把 layout 中所有块置 1实现等价于全注意力。官方保留它是为了对比与测试的正确性参照。补充当前 sparsity_config.py 中还额外提供了LocalSlidingWindowSparsityConfig纯局部滑动窗口、默认单向这是教程发布后新增的纯局部模式。五、如何启用 Sparse Attention从改造模型到 JSON 配置官方教程以 BERT 预训练代码为例演示了在 DeepSpeed 中启用 SA 的完整流程。5.1 把稠密注意力替换为稀疏注意力原BertSelfAttention中稠密的打分-掩码-softmax-加权流程attention_scores torch.matmul(query_layer, key_layer) attention_scores attention_scores / math.sqrt(self.attention_head_size) # Apply the attention mask is (precomputed for all layers in BertModel forward() function) attention_scores attention_scores attention_mask pdtype attention_scores.dtype # Normalize the attention scores to probabilities. attention_probs self.softmax(attention_scores) # This is actually dropping out entire tokens to attend to, which might # seem a bit unusual, but is taken from the original Transformer paper. attention_probs self.dropout(attention_probs) context_layer torch.matmul(attention_probs, value_layer)被替换为一次稀疏注意力调用context_layer self.sparse_self_attention( query_layer, key_layer, value_layer, key_padding_maskattention_mask)其中sparse_self_attention是SparseSelfAttention的实例。随后在模型中构造稀疏配置并下传给 encoderBertModel内self.pad_token_id config.pad_token_id if hasattr( config, pad_token_id) and config.pad_token_id is not None else 0 # set sparse_attention_config if it has been selected self.sparse_attention_config get_sparse_attention_config( args, config.num_attention_heads) self.encoder BertEncoder( config, args, sparse_attention_configself.sparse_attention_config)并在 encoder 的每层中当配置存在时启用BertSparseSelfAttentionif sparse_attention_config is not None: from deepspeed.ops.sparse_attention import BertSparseSelfAttention layer.attention.self BertSparseSelfAttention( config, sparsity_configsparse_attention_config)5.2 Padding 与 Unpadding必做由于块稀疏要求序列长度是block的整数倍模型入口处要调用pad_to_block_size出口处再调用unpad_sequence_output还原长度if self.sparse_attention_config is not None: pad_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds SparseAttentionUtils.pad_to_block_size( block_sizeself.sparse_attention_config.block, input_idsinput_ids, attention_maskextended_attention_mask, token_type_idstoken_type_ids, position_idsNone, inputs_embedsNone, pad_token_idself.pad_token_id, model_embeddingsself.embeddings) . . . # If BertEncoder uses sparse attention, and input_ids were padded, sequence output needs to be unpadded to original length if self.sparse_attention_config is not None and pad_len 0: encoded_layers[-1] SparseAttentionUtils.unpad_sequence_output( pad_len, encoded_layers[-1])从 pad_to_block_size 的实现 可以看到padding 会用pad_token_id填充input_ids、以False填充attention_mask保证 padding 位不被关注、以 0 填充token_type_ids与position_ids。5.3 通过启动器启用 SA在 DeepSpeed launcher 脚本中加入开关参数--deepspeed_sparse_attention5.4 DeepSpeed JSON 配置示例稀疏结构通过 DeepSpeed JSON 配置中的sparse_attention字段描述。下面是一个fixed模式的完整示例来自官方教程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 }在运行时DeepSpeed 配置解析器会在 deepspeed/runtime/config.py 的get_sparse_attention中按mode分发到dense/fixed/variable/bigbird/bslongformer五个分支解析结果存入DeepSpeedConfig.sparse_attentionconfig.py。也就是说上层 JSON 只需把模式名与结构参数对应填写即可无需手写任何 Python 配置代码。六、独立使用 Sparse Attention 或单个内核SA 并不强制依赖 DeepSpeed 训练引擎。安装 DeepSpeed 后可以直接把稀疏注意力当作一个普通 PyTorch 模块使用from deepspeed.ops.sparse_attention import SparseSelfAttention同样地MatMul、Softmax、BertSparseSelfAttention、SparseAttentionUtils以及全部*SparsityConfig都可以从deepspeed.ops.sparse_attention顶层导入见init.py。SparseSelfAttention的构造参数与默认值为SparseSelfAttention(sparsity_config..., key_padding_mask_modeadd, attn_mask_modemul, max_seq_length2048)其中max_seq_length决定预生成的master_layout尺寸源码。各模块的使用细节以类 docstring 为准。七、如何支持用户自定义的新稀疏结构由于底层 MatMul 与 Softmax 内核接受任意块布局新增一种稀疏模式不需要改动内核。只需要模仿任一内置结构新建一个继承SparsityConfig的类并按其几何结构实现make_layout必要时重写set_local_layout/set_global_layout/set_random_layout等辅助方法参考 Fixed 与 BigBird 的写法。所需新增参数直接在__init__中声明或用父类默认参数即可。在 sparsity_config.py 中可以看到 layout 的生成契约make_layout(seq_len)先用setup_layout得到全零的(num_heads, num_blocks, num_blocks)张量逐 head 置 1 后通过check_and_propagate_first_head_layout把 head 0 的布局复制到其余 head若共享布局。八、性能表现官方实验数据回顾作为发布于 2020 年的官方实验结果以下数据出自该介绍文章本身用于说明块稀疏方案在长序列场景下的收益实验环境与口径以原文为准如单卡 V100-32GB、batch size 1 等仅作横向参考。8.1 支撑超 10 倍更长的序列在 BERT 预训练实验中官方对比了三种设置稠密dense、稠密 activation checkpoint、稀疏SA activation checkpoint。结果显示 SA 能让 BERT base 与 large 的可训练序列长度分别提升约 10 倍与 16 倍对比稠密基线实验在单张 NVIDIA V100-32GB GPU、batch size 1 下进行8.2 最高 6.3 倍计算加速官方在 BERT base/large 与 Megatron GPT2 上以不同 batch size 与序列长度继续预训练 100 个迭代并记录最后 30 个迭代的平均耗时。结论是相比稠密方案SA 减少了总计算量并提升训练速度序列越长加速越明显提升最高分别达到 BERT base 约6.3×、BERT large 约5.3×、GPT2 约6.1×8.3 精度不低于甚至超过全注意力文章指出Sparse Transformer、Longformer、BigBird 等稀疏注意力相关工作已展示出与全注意力相当或更高的精度DeepSpeed 的实验与其一致在低内存开销与更快计算之外生产场景中还观察到 SA 取得更高精度与更快收敛。下图是长文档阅读理解sequence length 2,048的生产级 BERT 模型精度曲线对比三种设置dense 从头训练、SA 从头训练、以及从 dense512 序列长度checkpoint 继续训练的 SA——SA 从头训练比 dense 收敛更快、精度更高而从预训练 checkpoint 续训的 SA 在时间与精度两方面表现更好8.4 与 Longformer 的系统性对比在与当时 state-of-the-art 稀疏方案 Longformer 的对比中SA 使用Fixed稀疏二者精度相当而系统性能上 SA 在训练与推理均胜出。训练侧Wikitext103 上的 MLM 预训练提速 1.47 倍。该实验沿用 Longformer 提供的 RoBERTa-base checkpoint 转换流程在 8 张 V100-SXM2 GPU 上进行。下表按原文整理RoBERTa Checkpoint 基线 BPC 为 2.5326模型局部窗口大小BPCTrain Step每迭代耗时s时间提升精度提升Longformer5122.65350–1.471.01Sparse Attention5122.63210–Longformer5121.67083k1.62801.01Sparse Attention5121.66133k1.1059Longformer645.78400–1.311.46Sparse Attention643.97370–Longformer642.04663k1.48551.09Sparse Attention641.86933k1.1372说明原文表格在 markdown 转换中行列错位较严重上表将每一步Step 0 / 3k的 BPC 与耗时逐行转录时间/精度提升列按原文放置。可见在相同预算下 SA 的 BPC困惑度类指标更低、每迭代耗时明显更短。推理侧BERT-Base 长文档理解推理最高提速 3.13 倍。在同一长文档理解应用上于 2,048 序列长度、batch size 1 下替换不同窗口大小的 Longformer Attention得到如下结果局部窗口大小时间提升5123.132562.291282.16641.5321.24161.238.5 对任意块稀疏结构的灵活支持DeepSpeed Sparse Attention 并不绑定某一特定稀疏结构而是允许模型研究者借助高效的系统支持自由探索任意块稀疏结构。发布时已内置的流行结构包括Fixed来自 OpenAI Sparse Transformer 论文BigBird来自 GoogleBSLongformerAI2 Longformer 的块稀疏实现见前文配置说明以及可自由定制任意 block-sparse random/local/global 注意力模式的Variable 模板。结合仓库当前的 sparsity_config.py还可看到该体系后来扩展出的DenseSparsityConfig与LocalSlidingWindowSparsityConfig说明其基类 make_layout定制的扩展机制一直延续至今。九、适用前提与注意事项结合官方教程与当前源码使用 Sparse Attention 前请确认以下前提序列长度必须是块大小的整数倍seq_len % block 0不满足时需用pad_to_block_size补齐当前实现仅支持 fp16源码在 sparse_self_attention.py 中以断言强制query.dtype torch.half依赖 Triton内核由 Triton JIT 编译如 matmul.py 的triton.jit内核运行环境需满足对应 Triton 与 PyTorch/CUDA 版本要求官方教程发布时标注为 NVIDIA V100/A100、Torch ≥ 1.6 及特定 CUDA 版本与 DeepSpeed Transformer Kernel 互斥官方教程明确指出 Transformer Kernels 暂不支持 Sparse Attention二者同时使用时需禁用 Transformer Kernels独立使用时无引擎依赖模块可从deepspeed.ops.sparse_attention直接导入与任何 Transformer 模型组合。十、结语与进一步阅读DeepSpeed Sparse Attention 的价值在于把稀疏结构定义与高性能内核执行解耦研究者只需继承SparsityConfig并实现make_layout即可把任意局部 / 全局 / 随机块稀疏想法快速落地为可高效训练与推理的注意力层同时借助 Triton 内核把内存占用从O(n^2)压到O(wn)。如果你想在真实模型中运行可以继续阅读仓库内的配套教程 docs/_tutorials/sparse-attention.md包含本篇文章全部代码示例的原始出处或直接查看 sparsity_config.py、matmul.py 与 softmax.py 的 docstring 与测试用例以掌握每个参数与内核的真实行为。【免费下载链接】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 深度解析：基于 Triton 的块稀疏注意力如何支撑超长序列 Transformer
DeepSpeed Sparse Attention 深度解析基于 Triton 的块稀疏注意力如何支撑超长序列 Transformer【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeedSparse AttentionSA是 DeepSpeed 为自注意力类模型提供的一套块稀疏block-sparse注意力内核与配置体系目标是把注意力计算随序列长度二次增长的复杂度与显存占用降下来从而支撑更长的输入序列训练与推理。本篇文章基于 DeepSpeed 仓库中 2020 年发布的 Sparse Attention 官方介绍与配套教程结合 deepspeed/ops/sparse_attention 下的实际源码展开讲解读完你可以掌握块稀疏注意力的原理与复杂度优势、SA 内置的多种稀疏结构Fixed / BSLongformer / BigBird / Variable及其全部配置参数、如何在 DeepSpeed 中启用 SA、如何单独复用其 MatMul / Softmax 内核以及如何为任意新稀疏模式编写自定义布局。一、从O(n^2)到O(wn)Sparse Attention 要解决的问题以 Transformer 为代表的注意力模型擅长捕捉长距离 token 之间的依赖关系被广泛用于文本、图像、语音等序列输入场景序列长度常常达到数千 token。然而在实际落地中自注意力的计算量与显存占用随序列长度n呈平方级增长O(n^2)这使得它在长序列上受到严重制约——这是本文要解决的第一个核心问题。DeepSpeed 给出的方案是一套Sparse Attention 内核套件sparse attention kernels通过块稀疏计算block-sparse computation把注意力计算的算力与显存需求降低数个量级。它不仅能缓解注意力计算的内存瓶颈还能让稀疏计算本身保持高效执行。整体而言稀疏注意力把显存占用从O(n^2)降为O(wn)其中1 w n是一个与具体注意力结构相关的参数。Sparse Attention 可以灵活组合多种注意力形态局部注意力local只关注邻近 token全局注意力global通过局部窗口内计算的摘要块summary tokens关注整个序列随机注意力random允许一定数量的随机块被关注。下图展示了三者可任意组合的稀疏结构示例蓝色为局部、橙色为全局、绿色为随机块从实现上看这套库基于 PyTorch其内核通过Triton开发而不是手写 CUDA这为未来支持 CPU/OpenCL/Vulkan 预留了空间详见下文源码佐证。库本身是 DeepSpeed 的扩展既可以跟随 DeepSpeed 使用也可以独立使用。二、技术实现概览块稀疏矩阵与 Triton 内核2.1 前向与反向的块稀疏计算DeepSpeed Sparse Attention 的核心是把注意力分数矩阵按块block切分只对稀疏布局中置 1 的块执行矩阵乘与 softmax。下图中S表示块稀疏矩阵block-sparse matrixD表示稠密矩阵前向大体分两步先在稀疏块上计算S D × D即 Query 与 Key 的注意力分数再对稀疏块做 softmax随后做D S × D注意力分数与 Value 相乘得到 context。反向过程同理只是梯度矩阵同样是稀疏的2.2 模块级源码佐证当前仓库中SA 的全部实现位于 deepspeed/ops/sparse_attention 目录模块文件职责MatMulmatmul.py块稀疏矩阵乘法支持sdd、dsd、dds三种模式Softmaxsoftmax.py在稀疏块上执行 softmax同时处理前向与反向SparseSelfAttentionsparse_self_attention.py组装 MatMul 与 Softmax由 Q/K/V 生成 context 层输出BertSparseSelfAttentionbert_sparse_self_attention.py稀疏版BertSelfAttention可直接替换原稠密层SparseAttentionUtilssparse_attention_utils.py位置编码扩展、tokenizer 更新、padding/unpadding 等适配工具稀疏布局配置sparsity_config.pySparsityConfig基类及全部具体稀疏结构在 sparse_self_attention.py 的forward中可以看到清晰的调用链sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax self.get_ops(num_heads, tgt_len) scaling float(head_dim)**-0.5 # attention scoressparse Dense × Dense attn_output_weights sparse_dot_sdd_nt(query, key) # 块稀疏 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) # outputsDense Sparse × Dense attn_output sparse_dot_dsd_nn(attn_output_weights, value)几个值得注意的实现细节均可从源码确认内核缓存SparseSelfAttention.ops以序列长度L为键缓存(MatMul-sdd, MatMul-dsd, Softmax)三件套get_ops 只在首次遇到该长度时创建内核布局广播稀疏布局作为master_layoutbuffer 注册若分布式环境已初始化且布局未同步会通过dist.broadcast(self.master_layout, src0)从 rank 0 广播到所有进程见源码约束forward中断言query.dtype torch.half即当前只支持 fp16 训练get_layout会校验序列长度必须能被block整除Triton 内核块稀疏 MatMul 的 JIT 内核_kernel直接用triton.jit编写于 matmul.py并通过 LUTlook-up table记录每个非零块的偏移与指针增量避免 GPU 上的分支发散Softmax 内核位于 softmax.py。该目录下的trsrc/子目录存放了对应的 Triton 源文件matmul.tr、softmax_fwd.tr、softmax_bwd.tr印证了内核基于 Triton 而非手写 CUDA的设计。加速器抽象内核通过from deepspeed.accelerator import get_accelerator判断设备类型matmul.py说明该模块已接入 DeepSpeed 的 accelerator 抽象层。三、核心组件逐个拆解3.1 MatMul三种块稀疏矩阵乘模式MatMul类matmul.py接受一个稀疏layout、block尺寸和mode其中mode决定了输入/输出谁是稀疏的mode含义在注意力中的角色sdd输出稀疏 Dense × Dense由 Q、K 生成稀疏注意力分数dsd输出 Dense Sparse × Dense注意力分数 × V 生成 contextdds输出 Dense Dense × Sparse反向传播中计算梯度时使用在反向传播_sparse_matmul.backward中sdd对应的da/db会分别重映射到dds/dsd等模式即梯度计算同样只发生在非零块上。内核会先根据 layout 构造 C、DA、DB 三套 LUT 并缓存make_lut供前向与两次梯度计算复用。3.2 Softmax块内归一化Softmax模块在稀疏块上执行缩放与归一化除常规的scale外还内置了三类掩码的处理用于兼容标准注意力模块rpe相对位置编码relative position embeddingkey_padding_mask对 padding 位屏蔽模式可为add或mulattn_mask自定义注意力掩码当前支持二维Sequence × Sequence模式同样可为add或mul。掩码模式通过SparseSelfAttention构造参数key_padding_mask_modeadd、attn_mask_modemul指定默认值见 sparse_self_attention.py。3.3 SparseAttentionUtils把预训练模型改造成 SA对于想直接改造 Hugging Face BERT / RoBERTa 等已加载模型的用户sparse_attention_utils.py 提供了一组开箱即用的静态方法方法作用replace_model_self_attention_with_sparse_self_attention把 BERT/RoBERTa 编码器中所有 self-attention 层替换为稀疏版保留原有 Q/K/V 权重见 replace_self_attention_layer_with_sparse_self_attention_layerextend_position_embedding将预训练位置编码按倍数复制以适配更长序列例如 128 → 1k 时复制 8 次实验表明这种初始化比从零开始收敛更快见 实现update_tokenizer_model_max_length同步更新 tokenizer 的最大长度pad_to_block_size把输入 token 与注意力掩码在序列维度补齐到 block 的整数倍SA 的硬性要求unpad_sequence_output模型输出前裁剪掉补齐部分四、稀疏结构配置从基类到五种内置模式SA 的稀疏布局由SparsityConfig及其子类描述。基类在 sparsity_config.py 中定义并持有三类共享参数与两个共享方法共享参数num_heads该层注意力头数量必填block块大小当前实现基于块稀疏矩阵每个参与计算的都是block × block的方阵默认16different_layout_per_head是否允许每个注意力头使用不同的稀疏布局默认False所有头共享一份布局。核心方法setup_layout(seq_len)先校验seq_len % block 0再生成(num_heads, num_blocks, num_blocks)的全零布局张量num_blocks seq_len // blockmake_layout(seq_len)由具体子类实现负责画出该稀疏结构的 0/1 矩阵check_and_propagate_first_head_layout当different_layout_per_headFalse时把 head 0 的布局复制给其余所有头。由于布局的构造完全由make_layout决定任何块稀疏注意力模式都可以通过继承SparsityConfig并实现这一个函数来表达。DeepSpeed 官方内置了五种结构下面逐一说明其参数含默认值均与 sparsity_config.py 中类签名一致。4.1 FixedFixedSparsityConfig源自 OpenAI 的 Sparse Transformer 论文Generative Modeling with Sparse Transformers由若干参数固定的局部窗口 全局块构成参数默认值说明num_heads必填注意力头数block16块大小different_layout_per_headFalse每个头是否用不同布局num_local_blocks4每个局部注意力窗口包含的块数窗口内 token 互相可见num_global_blocks1每个局部窗口中作为全局代表的连续块数attentionbidirectionalunidirectional自回归只关注上文注意力矩阵上三角为空或bidirectional如 BERT可双向互看horizontal_global_attentionFalse全局代表块是否也横向关注所有块仅bidirectional下有效num_different_global_patterns1不同全局布局的版本数。例如局部窗口 4 块、全局 1 块时可让第 1/2/3/4 块分别担任全局代表从而有 4 个版本对应约束在构造函数中即被校验源码num_local_blocks必须能被num_global_blocks整除attention只能取unidirectional/bidirectionalhorizontal_global_attention仅限双向注意力num_different_global_patterns 1时必须打开different_layout_per_head且布局版本数不能超过num_local_blocks // num_global_blocks。4.2 BSLongformerBSLongformerSparsityConfigLongformerThe Long-Document Transformer的块稀疏实现——原版 Longformer 是逐 tokenelement-wise稀疏BSLongformer 将其改为块级稀疏参数默认值说明num_sliding_window_blocks3滑动局部窗口包含的块数global_block_indices[0]被视为全局注意力的块索引列表global_block_end_indicesNone全局窗口的结束索引若设置则必须与global_block_indices等长此时对每个iglobal_block_indices[i]到global_block_end_indices[i]不含之间的块都视为全局attentionbidirectional同上4.3 BigBirdBigBirdSparsityConfig源自 Google 的 Big Bird 论文综合了 Fixed/Longformer 的思路并加入随机注意力参数默认值说明num_random_blocks1每行块中随机选取多少个块参与注意力num_sliding_window_blocks3滑动局部窗口块数num_global_blocks1从索引 0 开始的前多少个块作为全局块被所有块关注、也关注所有块attentionbidirectional同上实现上随机块在 set_random_layout 中逐行random.sample若为单向注意力则只在[0, row]范围内采样避免未来信息泄漏。4.4 VariableVariableSparsityConfigVariable是官方提供的模板用来表达局部 全局 随机的任意自定义组合同时支持可变长度的局部窗口参数默认值说明num_random_blocks0每行随机块数0 表示不含随机块local_window_blocks[4]局部窗口块数列表第 k 个数字描述第 k 个窗口最后一个数字会沿用到剩余窗口global_block_indices[0]全局块索引与 end 参数组合时可作为起始索引global_block_end_indicesNone全局块窗口结束索引attentionbidirectional同上horizontal_global_attentionFalse同上4.5 DenseDenseSparsityConfig严格说 Dense 不是稀疏结构——它把 layout 中所有块置 1实现等价于全注意力。官方保留它是为了对比与测试的正确性参照。补充当前 sparsity_config.py 中还额外提供了LocalSlidingWindowSparsityConfig纯局部滑动窗口、默认单向这是教程发布后新增的纯局部模式。五、如何启用 Sparse Attention从改造模型到 JSON 配置官方教程以 BERT 预训练代码为例演示了在 DeepSpeed 中启用 SA 的完整流程。5.1 把稠密注意力替换为稀疏注意力原BertSelfAttention中稠密的打分-掩码-softmax-加权流程attention_scores torch.matmul(query_layer, key_layer) attention_scores attention_scores / math.sqrt(self.attention_head_size) # Apply the attention mask is (precomputed for all layers in BertModel forward() function) attention_scores attention_scores attention_mask pdtype attention_scores.dtype # Normalize the attention scores to probabilities. attention_probs self.softmax(attention_scores) # This is actually dropping out entire tokens to attend to, which might # seem a bit unusual, but is taken from the original Transformer paper. attention_probs self.dropout(attention_probs) context_layer torch.matmul(attention_probs, value_layer)被替换为一次稀疏注意力调用context_layer self.sparse_self_attention( query_layer, key_layer, value_layer, key_padding_maskattention_mask)其中sparse_self_attention是SparseSelfAttention的实例。随后在模型中构造稀疏配置并下传给 encoderBertModel内self.pad_token_id config.pad_token_id if hasattr( config, pad_token_id) and config.pad_token_id is not None else 0 # set sparse_attention_config if it has been selected self.sparse_attention_config get_sparse_attention_config( args, config.num_attention_heads) self.encoder BertEncoder( config, args, sparse_attention_configself.sparse_attention_config)并在 encoder 的每层中当配置存在时启用BertSparseSelfAttentionif sparse_attention_config is not None: from deepspeed.ops.sparse_attention import BertSparseSelfAttention layer.attention.self BertSparseSelfAttention( config, sparsity_configsparse_attention_config)5.2 Padding 与 Unpadding必做由于块稀疏要求序列长度是block的整数倍模型入口处要调用pad_to_block_size出口处再调用unpad_sequence_output还原长度if self.sparse_attention_config is not None: pad_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds SparseAttentionUtils.pad_to_block_size( block_sizeself.sparse_attention_config.block, input_idsinput_ids, attention_maskextended_attention_mask, token_type_idstoken_type_ids, position_idsNone, inputs_embedsNone, pad_token_idself.pad_token_id, model_embeddingsself.embeddings) . . . # If BertEncoder uses sparse attention, and input_ids were padded, sequence output needs to be unpadded to original length if self.sparse_attention_config is not None and pad_len 0: encoded_layers[-1] SparseAttentionUtils.unpad_sequence_output( pad_len, encoded_layers[-1])从 pad_to_block_size 的实现 可以看到padding 会用pad_token_id填充input_ids、以False填充attention_mask保证 padding 位不被关注、以 0 填充token_type_ids与position_ids。5.3 通过启动器启用 SA在 DeepSpeed launcher 脚本中加入开关参数--deepspeed_sparse_attention5.4 DeepSpeed JSON 配置示例稀疏结构通过 DeepSpeed JSON 配置中的sparse_attention字段描述。下面是一个fixed模式的完整示例来自官方教程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 }在运行时DeepSpeed 配置解析器会在 deepspeed/runtime/config.py 的get_sparse_attention中按mode分发到dense/fixed/variable/bigbird/bslongformer五个分支解析结果存入DeepSpeedConfig.sparse_attentionconfig.py。也就是说上层 JSON 只需把模式名与结构参数对应填写即可无需手写任何 Python 配置代码。六、独立使用 Sparse Attention 或单个内核SA 并不强制依赖 DeepSpeed 训练引擎。安装 DeepSpeed 后可以直接把稀疏注意力当作一个普通 PyTorch 模块使用from deepspeed.ops.sparse_attention import SparseSelfAttention同样地MatMul、Softmax、BertSparseSelfAttention、SparseAttentionUtils以及全部*SparsityConfig都可以从deepspeed.ops.sparse_attention顶层导入见init.py。SparseSelfAttention的构造参数与默认值为SparseSelfAttention(sparsity_config..., key_padding_mask_modeadd, attn_mask_modemul, max_seq_length2048)其中max_seq_length决定预生成的master_layout尺寸源码。各模块的使用细节以类 docstring 为准。七、如何支持用户自定义的新稀疏结构由于底层 MatMul 与 Softmax 内核接受任意块布局新增一种稀疏模式不需要改动内核。只需要模仿任一内置结构新建一个继承SparsityConfig的类并按其几何结构实现make_layout必要时重写set_local_layout/set_global_layout/set_random_layout等辅助方法参考 Fixed 与 BigBird 的写法。所需新增参数直接在__init__中声明或用父类默认参数即可。在 sparsity_config.py 中可以看到 layout 的生成契约make_layout(seq_len)先用setup_layout得到全零的(num_heads, num_blocks, num_blocks)张量逐 head 置 1 后通过check_and_propagate_first_head_layout把 head 0 的布局复制到其余 head若共享布局。八、性能表现官方实验数据回顾作为发布于 2020 年的官方实验结果以下数据出自该介绍文章本身用于说明块稀疏方案在长序列场景下的收益实验环境与口径以原文为准如单卡 V100-32GB、batch size 1 等仅作横向参考。8.1 支撑超 10 倍更长的序列在 BERT 预训练实验中官方对比了三种设置稠密dense、稠密 activation checkpoint、稀疏SA activation checkpoint。结果显示 SA 能让 BERT base 与 large 的可训练序列长度分别提升约 10 倍与 16 倍对比稠密基线实验在单张 NVIDIA V100-32GB GPU、batch size 1 下进行8.2 最高 6.3 倍计算加速官方在 BERT base/large 与 Megatron GPT2 上以不同 batch size 与序列长度继续预训练 100 个迭代并记录最后 30 个迭代的平均耗时。结论是相比稠密方案SA 减少了总计算量并提升训练速度序列越长加速越明显提升最高分别达到 BERT base 约6.3×、BERT large 约5.3×、GPT2 约6.1×8.3 精度不低于甚至超过全注意力文章指出Sparse Transformer、Longformer、BigBird 等稀疏注意力相关工作已展示出与全注意力相当或更高的精度DeepSpeed 的实验与其一致在低内存开销与更快计算之外生产场景中还观察到 SA 取得更高精度与更快收敛。下图是长文档阅读理解sequence length 2,048的生产级 BERT 模型精度曲线对比三种设置dense 从头训练、SA 从头训练、以及从 dense512 序列长度checkpoint 继续训练的 SA——SA 从头训练比 dense 收敛更快、精度更高而从预训练 checkpoint 续训的 SA 在时间与精度两方面表现更好8.4 与 Longformer 的系统性对比在与当时 state-of-the-art 稀疏方案 Longformer 的对比中SA 使用Fixed稀疏二者精度相当而系统性能上 SA 在训练与推理均胜出。训练侧Wikitext103 上的 MLM 预训练提速 1.47 倍。该实验沿用 Longformer 提供的 RoBERTa-base checkpoint 转换流程在 8 张 V100-SXM2 GPU 上进行。下表按原文整理RoBERTa Checkpoint 基线 BPC 为 2.5326模型局部窗口大小BPCTrain Step每迭代耗时s时间提升精度提升Longformer5122.65350–1.471.01Sparse Attention5122.63210–Longformer5121.67083k1.62801.01Sparse Attention5121.66133k1.1059Longformer645.78400–1.311.46Sparse Attention643.97370–Longformer642.04663k1.48551.09Sparse Attention641.86933k1.1372说明原文表格在 markdown 转换中行列错位较严重上表将每一步Step 0 / 3k的 BPC 与耗时逐行转录时间/精度提升列按原文放置。可见在相同预算下 SA 的 BPC困惑度类指标更低、每迭代耗时明显更短。推理侧BERT-Base 长文档理解推理最高提速 3.13 倍。在同一长文档理解应用上于 2,048 序列长度、batch size 1 下替换不同窗口大小的 Longformer Attention得到如下结果局部窗口大小时间提升5123.132562.291282.16641.5321.24161.238.5 对任意块稀疏结构的灵活支持DeepSpeed Sparse Attention 并不绑定某一特定稀疏结构而是允许模型研究者借助高效的系统支持自由探索任意块稀疏结构。发布时已内置的流行结构包括Fixed来自 OpenAI Sparse Transformer 论文BigBird来自 GoogleBSLongformerAI2 Longformer 的块稀疏实现见前文配置说明以及可自由定制任意 block-sparse random/local/global 注意力模式的Variable 模板。结合仓库当前的 sparsity_config.py还可看到该体系后来扩展出的DenseSparsityConfig与LocalSlidingWindowSparsityConfig说明其基类 make_layout定制的扩展机制一直延续至今。九、适用前提与注意事项结合官方教程与当前源码使用 Sparse Attention 前请确认以下前提序列长度必须是块大小的整数倍seq_len % block 0不满足时需用pad_to_block_size补齐当前实现仅支持 fp16源码在 sparse_self_attention.py 中以断言强制query.dtype torch.half依赖 Triton内核由 Triton JIT 编译如 matmul.py 的triton.jit内核运行环境需满足对应 Triton 与 PyTorch/CUDA 版本要求官方教程发布时标注为 NVIDIA V100/A100、Torch ≥ 1.6 及特定 CUDA 版本与 DeepSpeed Transformer Kernel 互斥官方教程明确指出 Transformer Kernels 暂不支持 Sparse Attention二者同时使用时需禁用 Transformer Kernels独立使用时无引擎依赖模块可从deepspeed.ops.sparse_attention直接导入与任何 Transformer 模型组合。十、结语与进一步阅读DeepSpeed Sparse Attention 的价值在于把稀疏结构定义与高性能内核执行解耦研究者只需继承SparsityConfig并实现make_layout即可把任意局部 / 全局 / 随机块稀疏想法快速落地为可高效训练与推理的注意力层同时借助 Triton 内核把内存占用从O(n^2)压到O(wn)。如果你想在真实模型中运行可以继续阅读仓库内的配套教程 docs/_tutorials/sparse-attention.md包含本篇文章全部代码示例的原始出处或直接查看 sparsity_config.py、matmul.py 与 softmax.py 的 docstring 与测试用例以掌握每个参数与内核的真实行为。【免费下载链接】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），仅供参考