深入解析 PaddleNLP 的 einsum 算子：从方程语法到底层实现与实战用例

深入解析 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 是什么：一段方程表达多种张量运算

einsum（Einstein 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-L42）：

def einsum(equation, *operands):

两个关键约定值得注意：

• 同一字母的维度必须一致或可广播：源码在回退实现中通过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，可以总结为以下规则：

1. 大小写不敏感：解析时先执行equation.lower()，只接受a-z字母（assert ord("a") <= ord(ch) and ord(ch) <= ord("z")）；
2. 空白被忽略：解析前会移除所有空格，因此'i, j->ij'与'i,j->ij'等价；
3. ->最多出现一次：assert num_eqns_size <= 2，多个->会直接报错；
4. 省略号...只能在每个操作数中出现一次，且要求所有操作数的省略号代表相同维度数（不同操作数间会校验curr_num_ell_idxes == num_ell_idxes）；
5. 输出方程中的字母必须已出现在输入方程中（character {} doesn't exist in input），且同一字母在输出中不能重复出现；
6. 省略输出方程时自动推断：省略号维度与"只出现一次"的字母维度会保留在输出中，重复出现的字母维度自动成为求和维；
7. 每个操作数的显式标签数（含省略号展开）必须等于其真实秩：源码中assert dims_in_terms == operand_rank（paddlenlp/ops/einsum.py#L239），标签数与张量维数不符会报Dimension dismatch错误。

经典范式实战：完整示例与运行结果

以下代码与输出完整继承自einsum的 docstring（paddlenlp/ops/einsum.py#L43-L112），覆盖了该 API 最常用的六种范式。注意原示例在 CUDA 设备上运行，输出中的place=CUDAPlace(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=[], dtype=float64, place=CUDAPlace(0), stop_gradient=True, 2.30369050) # 2) dot：向量点积 print(paddlenlp.ops.einsum('i,i->', x, x)) # Tensor(shape=[], dtype=float64, place=CUDAPlace(0), stop_gradient=True, 1.43773247) # 3) outer：向量外积，得到 [4, 5] 矩阵 print(paddlenlp.ops.einsum("i,j->ij", x, y)), # Tensor(shape=[4, 5], dtype=float64, place=CUDAPlace(0), stop_gradient=True, # [[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], dtype=float64, place=CUDAPlace(0), stop_gradient=True, # [[[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], dtype=float64, place=CUDAPlace(0), stop_gradient=True, # [[[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], dtype=float64, place=CUDAPlace(0), stop_gradient=True, # [[[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], dtype=float64, place=CUDAPlace(0), stop_gradient=True, # [[[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 的版本负责兼容旧环境。

回退实现的完整流程可以分为三个阶段：

1. 解析方程：将方程小写化、去空格，按->分割（最多一次），按逗号切分各操作数；为每个字母分配全局维度索引，并记录每个索引最后出现在哪个操作数（idxes_last_operand），为后续判断求和时机做准备；
2. 预处理操作数：根据输出方程（或自动推断）确定每个全局索引是"输出维"还是"求和维"，对每个操作数执行paddle.transpose调整维度顺序，再用paddle.unsqueeze补齐缺失维度，使所有操作数在逻辑上拥有相同的维度布局（paddlenlp/ops/einsum.py#L291-L323）；
3. 逐对执行 mul_sum：核心辅助函数_mul_sum（paddlenlp/ops/einsum.py#L118-L172）将一对操作数按"批量维 + 左侧输出维 + 求和维 + 右侧输出维"重排为三维，借助paddle.matmul一次完成批量矩阵乘与求和，最后reshape回目标形状并transpose恢复输出顺序；全部操作数处理完后，用paddle.squeeze去掉求和维（paddlenlp/ops/einsum.py#L337-L340）。

从实现细节可以确认几个行为边界：

• 广播维处理：在_mul_sum中，若某一求和维在一个操作数上是大小为 1 的广播维，则先对该操作数单独sum(axis=i, keepdim=True)，再参与矩阵乘法；
• 对角（diagonal）暂不支持：当同一字母在同一操作数中出现两次（即需要取对角线）时，回退实现会抛出NotImplementedError("Can't support diagonal.")，源码中以 TODO 注释注明需要开发paddle.diagonal算子支持（paddlenlp/ops/einsum.py#L308-L311）——这是使用旧版 Paddle 时需要注意的限制；
• 严格形状校验：求和维与非广播维均要求两侧相等，否则断言失败，避免静默产生错误结果。

附带的 transfer_param：FP16/FP32 参数迁移工具

同一模块还导出了transfer_param（paddlenlp/ops/einsum.py#L344-L367），源码注释标明其实现参考自 fast transformers，用于在训练/推理切换时完成参数精度与设备迁移：

def transfer_param(p, is_bias=False, dtype="float16", restore_data=False):

其行为逻辑为：

• 若参数已在目标精度且位于 GPU/CUDA 设备上，直接原样返回（避免重复转换）；
• 当restore_data=True时，动态图模式下通过p.numpy()取出数据并astype(dtype)后写回新参数；静态图模式下则从paddle.static.global_scope()读取全局作用域中的张量数据；
• 最终通过paddle.create_parameter以paddle.nn.initializer.Assign初始化器重建参数。

该工具在仓库中的典型定位是 FP16 训练流程中的参数精度对齐工具。需要注意：tests/transformer/train.py（tests/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）。

测试覆盖的范式清单（含测试类名，便于检索定位）如下：

值得说明的是，该测试用例定义于 2022 年（版权头标注 2022 PaddlePaddle），其测试数据矩阵（tests/ops/test_einsum.py#L22-L36）覆盖了一维到四维张量以及批量维为 1 的广播场景，可作为理解各方程语义的活教材。

在仓库中的真实使用场景

paddle.einsum（以及本模块语义）在 PaddleNLP 内部被广泛使用，以下列举可从源码直接确认的代表性场景：

• 全局指针网络（GlobalPointer）：paddlenlp/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