ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

PaddleNLP attention_utils 模块深度解析:BigBird 稀疏注意力与 MultiHeadAttention 的实现与实战

PaddleNLP attention_utils 模块深度解析:BigBird 稀疏注意力与 MultiHeadAttention 的实现与实战 人工智能大模型预训练微调LoRARLHF强化学习分布式训练【免费下载链接】PaddleNLPEasy-to-use and powerful LLM and SLM library with awesome model zoo.项目地址https://gitcode.com/gh_mirrors/pa/PaddleNLP点击查看免费下载paddlenlp.transformers.attention_utils是 PaddleNLP 中与 Transformer 注意力机制直接相关的基础工具模块承载了 BigBird 稀疏注意力的完整实现全局注意力 窗口注意力 随机注意力、默认的缩放点积注意力、可插拔的注意力实现注册机制以及一个支持增量推理缓存的多头注意力封装。本文以 docs/zh/source/paddlenlp.transformers.attention_utils.rst 所列 API 为线索深入 attention_utils.py 源码讲解其中每个类与函数的内部逻辑、参数含义并结合 BigBird 模型实现 与测试用例还原完整的调用链帮助读者掌握如何在长序列场景下使用与扩展这套注意力基础设施。一、模块定位PaddleNLP 的注意力基础设施docs/zh/source/paddlenlp.transformers.attention_utils.rst是一份 Sphinx 自动文档页通过automodule指令把 paddlenlp/transformers/attention_utils.py 中所有带文档的公开成员自动渲染成 API 参考。因此该模块的真实内容即其源码核心成员包括成员类型作用Registry/AttentionRegistry类 / 实例注意力实现的注册表通过装饰器按名称注册实现create_bigbird_rand_mask_idx函数生成单层 BigBird 随机注意力块的索引create_bigbird_rand_mask_idx_list函数为每一层生成一份随机块索引列表_convert_param_attr_to_list函数将单个ParamAttr统一展开为长度 n 的列表Linear3D类三维 Q/K/V 线性投影层输出[B, H, T, D]张量Attention抽象基类注意力实现基类定义统一的前向接口DefaultAttention类注册名为default_attention的缩放点积注意力BigBirdSparseAttention类注册名为bigbird的块稀疏注意力MultiHeadAttention类多头注意力封装含 Cache / StaticCache 推理缓存从导入关系看该模块是paddlenlp.transformers的公开基础设施之一init.py 将create_bigbird_rand_mask_idx_list直接导出到顶层命名空间BigBird 模型则同时导入MultiHeadAttention与_convert_param_attr_to_list使用。可以推断该模块是 PaddleNLP 中 BigBird 长序列模型专属的注意力工具箱。二、注册表模式AttentionRegistry 与可插拔注意力模块用一个轻量注册表来解耦注意力实现与使用方class Registry(object): def __init__(self): self.cls_dict {} def register(self, name): def add_item(name, cls): self.cls_dict[name] cls return cls return lambda cls: add_item(name, cls) AttentionRegistry Registry()register(name)返回一个装饰器被装饰的类会以name为键存入AttentionRegistry.cls_dict同时原类被原样返回因此装饰不会改变类的行为attention_utils.py#L26-L38。模块内已有两个注册项default_attention→DefaultAttention、bigbird→BigBirdSparseAttention。消费方通过字符串名取实现AttentionRegistry.cls_dictattention_type见 attention_utils.py#L558-L560。这种设计让MultiHeadAttention无需感知具体算法用户只需在配置中指定attention_typebigbird或original_full即可切换不同注意力实现想要新增算法时也只需实现Attention子类并注册即可无需改动既有模型代码。三、随机块索引生成create_bigbird_rand_mask_idx 与 create_bigbird_rand_mask_idx_listBigBird 稀疏注意力的关键难点在于每个查询块除了关注全局块和窗口块之外还要随机关注若干随机块。为了在多个 batch、多个 head 之间复用同一套随机采样结果模块在 CPU 上用 NumPy 一次性生成好索引再搬运到设备端使用。create_bigbird_rand_mask_idx(num_layers, query_length, key_length, num_heads, block_size, window_size, num_global_blocks, num_rand_blocks, seed)attention_utils.py#L41-L87的核心流程按block_size把序列切块得到num_key_blocks key_length // block_size、num_query_blocks query_length // block_size窗口半宽num_window_blocks window_size // 2。对每个查询块计算其非法块集合左右num_window_blocks范围内的窗口块这些块走窗口注意力不应再出现在随机块中前num_global_blocks个全局块这些块走全局注意力序列边界处会回卷补齐保证头部、尾部查询块的窗口块数量一致。对每个 head 独立做np.random.permutation从合法块中依次取出num_rand_blocks个作为该查询块的随机块。最后把所有 head 的索引堆叠并做一次减num_global_blocks // 2的偏移变换同时把 head 编号与块编号拼成[H*T, 2]形式的 gather 索引列表供后续gather_nd使用。create_bigbird_rand_mask_idx_list(num_layers, ...)attention_utils.py#L90-L108则是对上面函数按层数做列表推导返回形状为[num_layers, H, L, 2]代码中经np.stack堆叠的完整索引保证每一层使用不同的随机采样结果。调用示例与 BigBirdModel 官方示例一致import paddle from paddlenlp.transformers import BigBirdModel, BigBirdTokenizer from paddlenlp.transformers import create_bigbird_rand_mask_idx_list tokenizer BigBirdTokenizer.from_pretrained(bigbird-base-uncased) model BigBirdModel.from_pretrained(bigbird-base-uncased) config model.config max_seq_len 512 text This is a docudrama story on the Lindy Chamberlain case ... input_ids tokenizer.convert_tokens_to_ids(tokenizer(text)) input_ids.extend([0] * (max_seq_len - len(input_ids))) seq_len len(input_ids) input_ids paddle.to_tensor([input_ids]) rand_mask_idx_list create_bigbird_rand_mask_idx_list( config[num_layers], seq_len, seq_len, config[nhead], config[block_size], config[window_size], config[num_global_blocks], config[num_rand_blocks], config[seed]) rand_mask_idx_list [paddle.to_tensor(idx) for idx in rand_mask_idx_list] output model(input_ids, rand_mask_idx_listrand_mask_idx_list)需要注意从 BigBirdModel.forward 的实现看模型内部会依据config里的num_layers / nhead / block_size / window_size / num_global_blocks / num_rand_blocks / seed重新生成rand_mask_idx_list因此调用方即使不显式传入模型也能自洽运行生成逻辑本身仍是create_bigbird_rand_mask_idx_list。随机采样的可复现性由seed参数控制配置中seedNone时则每次生成结果不同。四、参数工具函数_convert_param_attr_to_list_convert_param_attr_to_list(param_attr, n)attention_utils.py#L111-L137用于把用户传入的ParamAttr统一规整为长度n的列表是 PaddleNLP 模型中一个配置同时驱动多个层的常见做法传入list/tuple要求长度必须等于n逐项规整True转为默认ParamAttrFalse表示该层不创建参数记为False否则调用ParamAttr._to_attr归一化。传入单个bool为True时生成n份默认ParamAttr否则生成n个False。传入单个ParamAttr深拷贝为n份若属性带名字则在末尾追加_i后缀以避免参数名冲突。在 BigBird 的 TransformerEncoderLayer 中它被用来一次性为自注意力层 FFN 层展开weight_attr与bias_attr再分别传入MultiHeadAttention和Linear。五、三维线性投影Linear3D常规nn.Linear输出[B, T, D]而多头注意力需要先把投影结果拆成多头。Linear3Dattention_utils.py#L140-L163把这两步合并result paddle.matmul(input, self.weight) # [B, T, D] x [D, D] result paddle.reshape(self.bias, [1, 1, D]) # 加偏置 result paddle.reshape(result, [B, T, H, -1]) # 拆出头维度 result paddle.transpose(result, [0, 2, 1, 3]) # [B, H, T, D]输入形状为[B, T, D]D 即hidden_size权重形状为[hidden_size, hidden_size]输出直接就是多头注意力所需的[B, H, T, D]布局省去了在MultiHeadAttention外层反复reshape transpose的样板代码。MultiHeadAttention的q_proj / k_proj / v_proj全部使用该层实现attention_utils.py#L553-L555。六、注意力基类与默认实现6.1 Attention 基类Attentionattention_utils.py#L166-L182定义了所有注意力实现的统一前向协议def forward(self, query_matrix, key_matrix, value_matrix, d_head, attn_maskNone, rand_mask_idxNone, query_maskNone, key_maskNone, dropoutNone): raise NotImplementedError参数语义query_matrix / key_matrix / value_matrix形状均为[B, H, T, D]d_head单头维度用于缩放点积attn_mask加法式注意力掩码直接加到 logits 上rand_mask_idxBigBird 随机块的 gather 索引query_mask / key_mask形状分别为[B, 1, T, 1]与[B, 1, 1, T]的布尔掩码dropout注意力权重上的 dropout 比例。子类只需实现该协议即可接入MultiHeadAttention。6.2 DefaultAttention缩放点积注意力DefaultAttention注册名default_attentionattention_utils.py#L185-L210是标准的缩放点积注意力计算product Q Kᵀ再乘以d_head ** -0.5做缩放用(1 - Q_mask K_mask) * -1e6把 padding 位置压成很大的负值使 softmax 后权重趋近于 0掩码部分用矩阵乘法生成等价于按行广播的 padding 掩码若传入attn_mask则继续累加支持自定义加法掩码例如单向因果掩码softmax得到权重若dropout非空则以upscale_in_train模式做训练期 dropoutout weights V输出。它对应的就是 BigBird 配置中attention_typeoriginal_full的 O(n²) 全注意力路径见 BigBirdConfig 文档。七、BigBirdSparseAttention全局 窗口 随机三路稀疏注意力BigBirdSparseAttention注册名bigbirdattention_utils.py#L213-L518是整个模块最核心、最复杂的部分。它的目标是把 O(n²) 的全注意力降为近似线性复杂度每个 token 只关注三类 key/value 块——序列首尾的全局块、自身附近的窗口块、以及按预生成索引采样的随机块。7.1 超参与分块策略__init__接收num_heads, block_size, window_size, num_global_blocks, num_rand_blocks, seed并额外计算self.num_global_blocks_back num_global_blocks // 2 self.num_global_blocks_front (num_global_blocks // 2 if num_global_blocks % 2 0 else num_global_blocks // 2 1)即把全局块均分到序列头front与尾back两侧奇数个全局块时前端多分一块。前向开始时输入[B, H, T, D]会被 reshape 成[B, H, L, bs, D]L T // bs为块数query/key/value 与对应的 mask 全部按块切分attention_utils.py#L443-L447。7.2 全局注意力_get_global_out_get_global_out(query_matrix, key_matrix, value_matrix, key_mask, d_head, dropout, is_front)attention_utils.py#L391-L404让序列最前GF * bs个 token或最后GB * bs个 token作为 query对完整序列做标准缩放点积注意力产出全局块输出。从源码结构看它内部同样用(1 - key_mask) * -1e6屏蔽 padding并在 softmax 后做weights V。7.3 带内注意力_get_band_mask 与 _get_band_matrix对于中间的非全局查询块需要同时聚合前端全局块 窗口块 后端全局块对应两个辅助函数_get_band_mask(blocked_query_mask, blocked_key_mask, batch_size, sequence_length)attention_utils.py#L227-L289用 mask 张量拼出形状[B, H, L-G, bs, (GW)*bs]的合法注意力范围掩码。头部/尾部查询块通过zeros_like补零与concat实现窗口块的回卷保证每个查询块都恰好看到W个窗口块。_get_band_matrix(blocked_matrix, B, T)attention_utils.py#L291-L348用同样的回卷逻辑从已分块的 K/V 矩阵中取出对应的(GW)个块重排成[B, H, L-G, (GW)*bs, D]的带内 key/value 矩阵其中全局块部分通过expand广播到每个查询位置。7.4 随机注意力_get_rand_mask 与 _gather_random_key_value_get_rand_mask(blocked_query_mask, blocked_key_mask, rand_mask_idx, batch_size, sequence_length)attention_utils.py#L350-L372依据rand_mask_idx用gather_nd把每个 head 对应的随机 key 掩码抓取出来再与查询掩码做einsum得到形状[B, H, L-G, bs, R*bs]的随机注意力掩码。_gather_random_key_value(blocked_matrix, rand_mask_idx, B, T)attention_utils.py#L374-L389对 K 和 V 分别执行相同的gather_nd得到[B, H, L-G, R*bs, D]的随机 key/value 矩阵。这两个函数都依赖第二节生成的rand_mask_idx形状[H, T]的[head_id, block_id]对因此随机块采样是预计算 查表而非前向中实时采样这也是seed参数能保证可复现的原因。7.5 forward三路结果合并forwardattention_utils.py#L410-L518的完整流程计算前端全局块输出global_front_out与后端全局块输出global_back_out拼接带内掩码与随机掩码得到second_mask拼接带内 K/V 与随机 K/V 得到second_key_matrix / second_value_matrix取中间查询块second_query_matrix blocked_query_matrix[:, :, GF:-GB]通过einsum(bhlqd,bhlkd-bhlqk)计算分数、按d_head**-0.5缩放、加掩码、softmax用_get_splited_matrix把权重与 V 按前端窗口 / 中间 / 后端窗口切成三份分别加权求和——其中中间部分需要额外补上全局块与随机块对应的分值见 attention_utils.py#L496-L508把三部分输出拼回[B, H, (L-G)*bs, D]再与全局前/后输出拼接成完整的[B, H, T, D]最后乘以query_mask屏蔽 padding。整体数据流可概括为Q/K/V [B,H,T,D] --按块 reshape-- [B,H,L,bs,D] ├── 全局前/后块 → 标准全注意力_get_global_out ├── 中间查询块 → 带内块_get_band_* 随机块_gather_random_key_value └── 三路 concat → out [B,H,T,D] × query_mask八、MultiHeadAttention支持推理缓存的多头注意力封装MultiHeadAttentionattention_utils.py#L521-L619把投影、注意力实现、头合并与推理缓存整合成一个可直接使用的nn.Layer。8.1 构造参数参数默认值说明embed_dim必填模型隐藏维度必须是num_heads的整数倍构造时断言num_heads必填注意力头数head_dim embed_dim // num_headsdropout0.0注意力权重 dropout 比例kdim / vdimembed_dimK/V 投影输入维度便于处理 cross-attention 中 K/V 与 Q 维度不同的场景weight_attr / bias_attrNoneQ/K/V 与输出投影的参数属性block_size1BigBird 块大小window_size3BigBird 窗口块数量num_global_blocks1BigBird 全局块数量num_rand_blocks1BigBird 随机块数量seedNone随机块采样种子attention_typebigbird从AttentionRegistry选择实现内部结构q_proj / k_proj / v_proj为Linear3Dout_proj为nn.Linearattn_impl由注册表按attention_type实例化attention_utils.py#L553-L560。8.2 Cache 与 StaticCache模块定义了两个命名元组attention_utils.py#L523-L524Cache collections.namedtuple(Cache, [k, v]) StaticCache collections.namedtuple(StaticCache, [k, v])Cache增量式incremental缓存用于自回归解码的 self-attention。_prepare_qkv在isinstance(cache, self.Cache)时把新计算的 K/V 沿序列维concat到历史缓存上实现只算新 token的逐 token 解码attention_utils.py#L571-L575。StaticCache静态 K/V 缓存用于 encoder-decoder 场景如 UniLM_prepare_qkv检测到该类型时直接复用缓存中的 K/V不再重新投影attention_utils.py#L565-L568。gen_cache(key, valueNone, typeCache)attention_utils.py#L584-L595负责按类型构造缓存StaticCache需要立即对key计算 K/VCache在未提供value时返回形状[-1, num_heads, 0, head_dim]的空缓存传入value时则把初始 K/V 装入缓存注释标明主要用于 UniLM 等场景。8.3 forward 数据流q self.q_proj(query) # [B, H, T, D] if isinstance(cache, self.StaticCache): k, v cache.k, cache.v # 复用静态缓存 else: k, v self.compute_kv(key, value) # k_proj / v_proj if isinstance(cache, self.Cache): k paddle.concat([cache.k, k], axis2) # 增量拼接 v paddle.concat([cache.v, v], axis2) out self.attn_impl(q, k, v, self.head_dim, attn_mask, rand_mask_idx, query_mask, key_mask, self.dropout) out paddle.transpose(out, [0, 2, 1, 3]) # [B, T, H, D] out paddle.reshape(out, [0, 0, out.shape[2] * out.shape[3]]) out self.out_proj(out) # 合并多头并投影attention_typebigbird时attn_impl就是上一节的BigBirdSparseAttention因此MultiHeadAttention既保留了多头注意力的标准外接口投影、多头合并、缓存又完全继承了 BigBird 的稀疏计算。九、在 BigBird 模型中的集成调用链以 paddlenlp/transformers/bigbird/modeling.py 为参考整条调用链为配置层BigBirdConfig定义attention_type默认bigbird、block_size、window_size、num_global_blocks、num_rand_blocks等超参。预训练配置bigbird-base-uncased中block_size16、window_size3、num_global_blocks2、num_rand_blocks3、max_position_embeddings4096见 configuration.py#L23-L45。编码层TransformerEncoderLayer.__init__用_convert_param_attr_to_list展开参数后构造MultiHeadAttention(..., attention_typeconfig.attention_type, block_size..., window_size..., num_global_blocks..., num_rand_blocks..., seed...)modeling.py#L84-L96。模型层BigBirdModel._process_mask依据pad_token_id生成attention_mask / query_mask / key_maskmodeling.py#L380-L394forward内部调用create_bigbird_rand_mask_idx_list为每一层生成随机块索引再交给TransformerEncoder逐层前向modeling.py#L501-L521。任务层BigBirdForSequenceClassification / BigBirdForQuestionAnswering / BigBirdForPretraining等在BigBirdModel之上加输出头复用同一套稀疏注意力。十、测试验证仓库在 tests/transformers/bigbird/test_modeling.py 中对 BigBird 相关实现做了系统测试BigBirdModelTester覆盖batch_size13、seq_length7、hidden_size32、num_attention_heads4、num_hidden_layers5等配置通过BigBirdConfig(...)构造模型并验证各任务模型BigBirdForMultipleChoice、BigBirdForQuestionAnswering、BigBirdForSequenceClassification、BigBirdForTokenClassification、BigBirdForPretraining的输入输出test_modeling.py#L40-L135。测试通过parameterized_class参数化return_dict等选项并复用通用ModelTesterMixin检查attention_utils相关实现与 PaddleNLP 模型基类约定的兼容性test_modeling.py#L20-L37。这些测试一方面印证了attention_utils中注册表、随机索引生成、MultiHeadAttention等组件的正确性另一方面也表明该模块是 PaddleNLP 模型库中可被独立测试与复用的公共组件。十一、小结与实践建议paddlenlp.transformers.attention_utils为 PaddleNLP 提供了一整套即插即用的注意力基础设施需要标准全注意力时使用attention_typeoriginal_fullDefaultAttention需要长序列稀疏注意力时使用attention_typebigbirdBigBirdSparseAttention并配合create_bigbird_rand_mask_idx_list预生成随机块索引注意block_size需能整除序列长度、embed_dim需能被num_heads整除自回归解码时通过MultiHeadAttention.gen_cache初始化Cache实现增量 K/V 缓存想接入自定义注意力算法时继承Attention并实现统一forward协议再用AttentionRegistry.register(your_name)注册即可通过配置字符串启用。理解该模块是深入 PaddleNLP 长序列模型尤其是 BigBird 系列源码、乃至在其基础上做注意力算法二次开发的第一站。赞分享人工智能大模型预训练微调LoRARLHF强化学习分布式训练【免费下载链接】PaddleNLPEasy-to-use and powerful LLM and SLM library with awesome model zoo.项目地址https://gitcode.com/gh_mirrors/pa/PaddleNLP点击查看免费下载相关推荐PaddleNLP BigBird 模型实战指南从块稀疏注意力原理到预训练权重使用PaddleNLP BigBird 模型实战指南从块稀疏注意力原理到预训练权重使用 本篇技术指南围绕 PaddleNLP 中 BigBird 模型的支持情况展人工智能大模型预训练微调LoRARLHF强化学习分布式训练模型推理服务推理引擎模型量化模型压缩本地部署NLPPaddleNLP 中的 BigBird 模型稀疏注意力大模型的使用与原理详解PaddleNLP 中的 BigBird 模型稀疏注意力大模型的使用与原理详解 BigBird 是 Google Research 提出的稀疏注意力 Tran人工智能大模型预训练微调LoRARLHF强化学习分布式训练模型推理服务推理引擎模型量化模型压缩本地部署NLP告别注意力迷雾用BertViz可视化工具轻松破解BigBird稀疏注意力机制告别注意力迷雾用BertViz可视化工具轻松破解BigBird稀疏注意力机制 BertViz是一款强大的NLP模型注意力可视化工具能够帮助开发者和研究者直观AI 可解释性NLP数据可视化上一篇告别9KB冗余bignumber.js生产环境极致优化指南下一篇突破Windows文件系统限制使用WinFsp实现压缩镜像实时挂载的终极指南创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表