深入解析 PaddleNLP 的 einsum 算子从方程语法到底层实现与实战用例【免费下载链接】PaddleNLPEasy-to-use and powerful LLM and SLM library with awesome model zoo.项目地址: https://gitcode.com/gh_mirrors/pa/PaddleNLP本篇技术指南围绕 PaddleNLP 在 paddlenlp/ops/einsum.py 中提供的einsum与transfer_param两个工具函数展开完整梳理爱因斯坦求和约定的方程语法、全部内置示例、底层实现原理原生委托 自研回退两条路径、配套单元测试以及该算子在 PaddleNLP 各模型中的真实调用场景。读完本文你将能够在自己的 PaddleNLP 项目中用一行方程完成求和、点积、外积、转置、批量矩阵乘法、广播等张量运算并理解其退化到 Paddle 2.3 以下版本时的兼容处理方式。该模块的 API 文档位于 docs/zh/source/paddlenlp.ops.einsum.rst是 PaddleNLP 操作算子文档体系见 docs/zh/source/paddlenlp.ops.rst的组成部分本文内容以该文档对应的源码与测试为准。einsum 是什么一段方程表达多种张量运算einsumEinstein summation convention爱因斯坦求和约定允许用一段紧凑的字符串方程描述先乘后加的张量运算。PaddleNLP 将其封装在paddlenlp.ops命名空间下可以通过以下方式导入import paddlenlp.ops as ops ops.einsum(i-, x) # 求和 ops.einsum(i,i-, x, x) # 点积 ops.einsum(i,j-ij, x, y) # 外积从源码结构看paddlenlp/ops/einsum.py该模块通过__all__ [einsum, transfer_param]对外暴露两个接口并在 paddlenlp/ops/init.py 中以from .einsum import *汇入顶层paddlenlp.ops因此paddlenlp.ops.einsum与paddlenlp.ops.transfer_param均可直接使用。API 签名与参数说明einsum的函数签名与 docstring 定义如下源码见 paddlenlp/ops/einsum.py#L21-L42def einsum(equation, *operands):参数类型说明equationstr用小写字母不区分大小写标注操作数与结果的维度。-左侧为输入方程右侧为输出方程。结果形状可被自动推断因此-与结果字母可省略。多个操作数之间用逗号,分隔例如abc,cde描述两个 3D 操作数。标注为同一字母的维度必须大小相同或为 1可广播。省略号...用于标注广播维度。operandsTensor参与爱因斯坦求和的操作数数量必须与输入方程中描述的操作数一一对应。返回Tensor爱因斯坦求和的结果张量。两个关键约定值得注意同一字母的维度必须一致或可广播源码在回退实现中通过assert left.shape[i] right.shape[i]强制校验非广播维相等paddlenlp/ops/einsum.py#L133操作数个数必须与方程匹配源码中assert len(operand_eqns) len(operands)paddlenlp/ops/einsum.py#L188-L190不匹配会直接报错。方程语法详解方程解析逻辑位于 paddlenlp/ops/einsum.py#L176-L289可以总结为以下规则大小写不敏感解析时先执行equation.lower()只接受a-z字母assert ord(a) ord(ch) and ord(ch) ord(z)空白被忽略解析前会移除所有空格因此i, j-ij与i,j-ij等价-最多出现一次assert num_eqns_size 2多个-会直接报错省略号...只能在每个操作数中出现一次且要求所有操作数的省略号代表相同维度数不同操作数间会校验curr_num_ell_idxes num_ell_idxes输出方程中的字母必须已出现在输入方程中character {} doesnt exist in input且同一字母在输出中不能重复出现省略输出方程时自动推断省略号维度与只出现一次的字母维度会保留在输出中重复出现的字母维度自动成为求和维每个操作数的显式标签数含省略号展开必须等于其真实秩源码中assert dims_in_terms operand_rankpaddlenlp/ops/einsum.py#L239标签数与张量维数不符会报Dimension dismatch错误。经典范式实战完整示例与运行结果以下代码与输出完整继承自einsum的 docstringpaddlenlp/ops/einsum.py#L43-L112覆盖了该 API 最常用的六种范式。注意原示例在 CUDA 设备上运行输出中的placeCUDAPlace(0)即为此环境特征。import numpy as np import paddle import paddlenlp np.random.seed(102) x paddle.to_tensor(np.random.rand(4)) y paddle.to_tensor(np.random.rand(5)) # 1) sum对向量全部元素求和 print(paddlenlp.ops.einsum(i-, x)) # Tensor(shape[], dtypefloat64, placeCUDAPlace(0), stop_gradientTrue, 2.30369050) # 2) dot向量点积 print(paddlenlp.ops.einsum(i,i-, x, x)) # Tensor(shape[], dtypefloat64, placeCUDAPlace(0), stop_gradientTrue, 1.43773247) # 3) outer向量外积得到 [4, 5] 矩阵 print(paddlenlp.ops.einsum(i,j-ij, x, y)), # Tensor(shape[4, 5], dtypefloat64, placeCUDAPlace(0), stop_gradientTrue, # [[0.34590188, 0.48353496, 0.09996135, 0.18656330, 0.21392910], # [0.39122025, 0.54688535, 0.11305780, 0.21100591, 0.24195704], # [0.17320613, 0.24212422, 0.05005442, 0.09341929, 0.10712238], # [0.42290818, 0.59118179, 0.12221522, 0.22809690, 0.26155500]]) A paddle.to_tensor(np.random.rand(2, 3, 2)) B paddle.to_tensor(np.random.rand(2, 2, 3)) # 4) transpose三维转置 print(paddlenlp.ops.einsum(ijk-kji, A)) # Tensor(shape[2, 3, 2], dtypefloat64, placeCUDAPlace(0), stop_gradientTrue, # [[[0.49174730, 0.33344683], # [0.89440989, 0.26162022], # [0.36116209, 0.12241719]], # [[0.49019824, 0.51895050], # [0.18241053, 0.13092809], # [0.81059146, 0.55165734]]]) # 5) batch matrix multiplication批量矩阵乘法 print(paddlenlp.ops.einsum(ijk, ikl-ijl, A, B)) # Tensor(shape[2, 3, 3], dtypefloat64, placeCUDAPlace(0), stop_gradientTrue, # [[[0.13654339, 0.39331432, 0.65059661], # [0.07171420, 0.57518653, 0.77629221], # [0.21250688, 0.37793541, 0.73643411]], # [[0.56925339, 0.65859030, 0.57509818], # [0.30368265, 0.25778348, 0.21630400], # [0.39587265, 0.58031243, 0.51824755]]]) # 6) Ellipsis transpose带省略号的转置 print(paddlenlp.ops.einsum(...jk-...kj, A)) # Tensor(shape[2, 2, 3], dtypefloat64, placeCUDAPlace(0), stop_gradientTrue, # [[[0.49174730, 0.89440989, 0.36116209], # [0.49019824, 0.18241053, 0.81059146]], # [[0.33344683, 0.26162022, 0.12241719], # [0.51895050, 0.13092809, 0.55165734]]]) # 7) Ellipsis batch matrix multiplication带省略号的批量矩阵乘法 print(paddlenlp.ops.einsum(...jk, ...kl-...jl, A, B)) # Tensor(shape[2, 3, 3], dtypefloat64, placeCUDAPlace(0), stop_gradientTrue, # [[[0.13654339, 0.39331432, 0.65059661], # [0.07171420, 0.57518653, 0.77629221], # [0.21250688, 0.37793541, 0.73643411]], # [[0.56925339, 0.65859030, 0.57509818], # [0.30368265, 0.25778348, 0.21630400], # [0.39587265, 0.58031243, 0.51824755]]])示例 6 与 7 展示了省略号的核心价值...会沿未标注的前导维度自动广播使同一段方程既适用于固定形状也适用于任意批量维度的输入。底层实现优先原生算子旧版本自动回退einsum的实现采用版本检测 双路径策略paddlenlp/ops/einsum.py#L114-L116# paddle.einsum can be used in paddle 2.3.0 if hasattr(paddle, einsum): return paddle.einsum(equation, *operands)当运行环境中的 Paddle 版本提供paddle.einsum按源码注释为 2.3.0 及以上版本时直接委托给 Paddle 原生的高性能算子此时 PaddleNLP 的封装只承担统一入口与文档化的职责。只有当paddle.einsum不存在时才会走到自研回退实现。这也意味着在主流新版本 Paddle 上你实际执行的是 Paddle 官方实现PaddleNLP 的版本负责兼容旧环境。回退实现的完整流程可以分为三个阶段解析方程将方程小写化、去空格按-分割最多一次按逗号切分各操作数为每个字母分配全局维度索引并记录每个索引最后出现在哪个操作数idxes_last_operand为后续判断求和时机做准备预处理操作数根据输出方程或自动推断确定每个全局索引是输出维还是求和维对每个操作数执行paddle.transpose调整维度顺序再用paddle.unsqueeze补齐缺失维度使所有操作数在逻辑上拥有相同的维度布局paddlenlp/ops/einsum.py#L291-L323逐对执行 mul_sum核心辅助函数_mul_sumpaddlenlp/ops/einsum.py#L118-L172将一对操作数按批量维 左侧输出维 求和维 右侧输出维重排为三维借助paddle.matmul一次完成批量矩阵乘与求和最后reshape回目标形状并transpose恢复输出顺序全部操作数处理完后用paddle.squeeze去掉求和维paddlenlp/ops/einsum.py#L337-L340。从实现细节可以确认几个行为边界广播维处理在_mul_sum中若某一求和维在一个操作数上是大小为 1 的广播维则先对该操作数单独sum(axisi, keepdimTrue)再参与矩阵乘法对角diagonal暂不支持当同一字母在同一操作数中出现两次即需要取对角线时回退实现会抛出NotImplementedError(Cant support diagonal.)源码中以 TODO 注释注明需要开发paddle.diagonal算子支持paddlenlp/ops/einsum.py#L308-L311——这是使用旧版 Paddle 时需要注意的限制严格形状校验求和维与非广播维均要求两侧相等否则断言失败避免静默产生错误结果。附带的 transfer_paramFP16/FP32 参数迁移工具同一模块还导出了transfer_parampaddlenlp/ops/einsum.py#L344-L367源码注释标明其实现参考自 fast transformers用于在训练/推理切换时完成参数精度与设备迁移def transfer_param(p, is_biasFalse, dtypefloat16, restore_dataFalse):参数默认值作用p—待迁移的原始参数paddle参数对象is_biasFalse是否按偏置参数创建决定is_biasTrue传给底层创建接口dtypefloat16目标精度默认迁移为 FP16可传float32等restore_dataFalse是否将原参数数据拷贝到新参数其行为逻辑为若参数已在目标精度且位于 GPU/CUDA 设备上直接原样返回避免重复转换当restore_dataTrue时动态图模式下通过p.numpy()取出数据并astype(dtype)后写回新参数静态图模式下则从paddle.static.global_scope()读取全局作用域中的张量数据最终通过paddle.create_parameter以paddle.nn.initializer.Assign初始化器重建参数。该工具在仓库中的典型定位是 FP16 训练流程中的参数精度对齐工具。需要注意tests/transformer/train.pytests/transformer/train.py#L116-L119中同名的transfer_param是测试脚本自定义的简单实现仅做paddle.cast到 float32与 paddlenlp/ops/einsum.py 导出的版本并不相同阅读代码时请勿混淆。测试覆盖tests/ops/test_einsum.py 全范式验证仓库在 tests/ops/test_einsum.py 中为paddlenlp.ops.einsum提供了系统性的单元测试。其核心机制是每个测试子类设置一组{paradigm: 方程, data: 操作数名列表}test_forward将 numpy 版np.einsum的结果与 PaddleNLP 版结果逐元素比对tests/ops/test_einsum.py#L43-L51。测试覆盖的范式清单含测试类名便于检索定位如下范式方程测试类i-求和TestEinsumi,i-点积TestEinsumVectorDoti,i-i逐元素乘TestEinsumVectorMuli,j-ij外积TestEinsumVectorOuterij-ji矩阵转置TestEinsumMatrixTransposeij-j/ij-i行列求和TestEinsumMatrixRowSum/TestEinsumMatrixColSumij,ij-ij逐元素乘TestEinsumMatrixEleMulij,j-i矩阵乘向量TestEinsumMatrixVecMulij,kj-ik矩阵乘TestEinsumMatrixMulij,kl-ijkl矩阵外积TestEinsumMatrixOuterbij,bjk-bik批量矩阵乘TestEinsumTensorBMMijk,jk-i/ijk,jk-ik/ijk,jk-ij张量缩并TestEinsumTensorContract1/4/5ijk,lk-ijl/abcd,dfg-abcfg/ijk,lk-ijl多操作数缩并TestEinsumTensorContract2/3ik,ijk-j/ijk,ik-jk混合缩并TestEinsumTensorContract6/7i...-.../ij,...i-j.../k...,jk省略号广播TestEinsumEllipsis1/2/3bn,anm,bm-ba双线性形式TestEinsumTestEinsumBilinearijkl,lmn-ijn高阶缩并TestEinsumTestEinsumOthersblq,bhlk-bhlqk注意力掩码式广播TestEinsumBatch1值得说明的是该测试用例定义于 2022 年版权头标注 2022 PaddlePaddle其测试数据矩阵tests/ops/test_einsum.py#L22-L36覆盖了一维到四维张量以及批量维为 1 的广播场景可作为理解各方程语义的活教材。在仓库中的真实使用场景paddle.einsum以及本模块语义在 PaddleNLP 内部被广泛使用以下列举可从源码直接确认的代表性场景全局指针网络GlobalPointerpaddlenlp/layers/globalpointer.py#L59 使用paddle.einsum(bmd,bnd-bmn, qw, kw)计算头尾实体的打分矩阵配合head_size**0.5缩放这是信息抽取任务中的典型打分结构LoRA 的 NoLA 基展开paddlenlp/peft/lora/lora_layers.py#L217-L219 通过paddle.einsum(k,kir-ir, ...)与paddle.einsum(k,kro-ro, ...)将多个基矩阵按系数线性组合重建低秩适配矩阵注意力掩码与稀疏注意力paddlenlp/transformers/attention_utils.py 中大量使用 einsum 风格方程构造 block 级别注意力掩码如blqd,bmdk-blqk、blkd,bldq-blkq以及融合注意力中的二次乘积bhlqd,bhlkd-bhlqk与输出聚合bhlqk,bhlkd-bhlqd旋转位置编码频率表paddlenlp/experimental/transformers/chatglm/modeling.py#L92 与 paddlenlp/experimental/transformers/deepseek_v2/modeling.py#L120 使用paddle.einsum(i,j-ij, t, inv_freq)生成位置-频率外积矩阵显式替换优化paddlenlp/transformers/auto_utils.py#L29-L59 中提供了一个einsum(rule, a, b)辅助函数对s,se-se、se,sc-sec、se,se-s、sec,sm-ecm、sec,ecm-sm、ks,ksm-sm等常见规则直接用reshape、unsqueeze、bmm、matmul组合替代实现来源注明为 DeepSpeed其余规则才回退到paddle.einsum——这说明在 PaddleNLP 的工程实践中将高频 einsum 方程手工拆解为基元算子是一种常见的性能优化手段也印证了本文所述方程与底层算子之间的对应关系。小结PaddleNLP 的einsum工具以一行字符串方程统一了求和、点积、外积、转置、批量矩阵乘与广播运算对外提供稳定的paddlenlp.ops.einsum入口对内通过版本检测优先委托paddle.einsum、旧环境自研回退的方式保证兼容性配套的transfer_param则服务于 FP16 参数迁移。结合 tests/ops/test_einsum.py 的二十余个范式用例与仓库各模型中的真实调用读者既可以将其直接用于日常张量运算也可以参考 paddlenlp/transformers/auto_utils.py 的拆解思路对热点方程做手写算子级优化。【免费下载链接】PaddleNLPEasy-to-use and powerful LLM and SLM library with awesome model zoo.项目地址: https://gitcode.com/gh_mirrors/pa/PaddleNLP创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
