1. 先把“矩阵运算”这件事在PyTorch张量体系里说清楚PyTorch张量的数学运算里矩阵运算大概是被问得最多的一块。原因很朴素深度学习里几乎所有核心算子——全连接层、注意力里的QKV打分、卷积展开后的乘加、归一化里的协方差矩阵——最后都能归结成几次矩阵乘法。很多人调用API没问题但一遇到shape对不上、速度突然掉一半、反向传播报in-place错误就抓瞎了。问题不在API记不熟而在于脑子里没建立起“形状怎么对齐、内存怎么排布、梯度怎么回传”这三条并行的线。这篇内容面向的是已经能写几行PyTorch、但没系统梳理过矩阵运算的人。不管你是刚跑通MNIST想弄明白到底做了什么还是已经在调模型、想搞清楚为什么加一个transpose就慢了30%下面这些拆解应该都能对上你的痒处。我不会只列API清单重点是解释每个选择背后的代价——为什么这里用mm不用matmul为什么那里宁可多拷一次内存也要contiguous()为什么用inv解方程是新手最容易犯的错之一。1.1 张量不是数学符号是带步长描述的内存视图初学者最容易把张量想成一张二维表格实际上它是一块连续内存加一份“步长stride元数据”。一个(3, 4)的float32张量内存里就是12个连续浮点数步长是(4, 1)意思是“行方向每跳一步跨4个元素列方向每跳一步跨1个元素”。真正决定形状的是这份元数据而不是内存的物理排布。这个认知一旦建立很多诡异行为立刻就不诡异了。比如a.T根本不搬运任何数据它只是把步长从(4, 1)改成(1, 4)告诉后续算子“请按列优先来读”。所以转置是零成本的但也因此转置后的张量不是内存连续的。后面所有关于性能的坑根源都在这里。我习惯在调试时随手打一行print(x.shape, x.stride(), x.is_contiguous())比看什么都管用。尤其是接手别人代码的时候一个看起来平平无奇的(B, S, H, D)四维张量如果stride显示它是从某个permute过来的那它十有八九不是连续的后面接矩阵乘法就要留意开销。注意stride()在排查“为什么这段代码比预期慢”时信息量往往比shape()更大。1.2 元素级运算和矩阵级运算的分界线torch.mul、*、torch.add这类是逐元素运算两个张量的形状要么完全相同要么满足广播规则而torch.mm、torch.matmul、是矩阵乘法遵循的是(m, k) (k, n) - (m, n)这条内维收缩规则。这两套规则长得很像但完全不同*和只差一个符号犯错的成本却很高。一个典型的混淆场景你想算两个向量的内积写了a * b得到的是逐元素乘积还得再.sum()一次写a b一步到位。两者在(n,)形状下的结果一个是张量、一个是标量后续梯度形状也不同。更隐蔽的是a * b在广播下可能悄悄把你的(B, 1)扩成(B, B)代码不报错显存却突然炸了。我的经验是凡是涉及“内维求和”的语义一律用矩阵乘法API凡是涉及“同位对齐”的语义才用逐元素。写代码时强迫自己在注释里标一句“这里做的是内维收缩”能挡掉一大半形状事故。1.3 为什么必须区分mm、bmm、matmul这三个名字这三个API的功能有重叠但约束条件不同选错了轻则报错重则性能打折。API输入维度要求是否支持广播典型用途torch.mm严格2D否单次矩阵乘形状已确定torch.bmm严格3Dbatch数必须一致否批量矩阵乘batch对齐torch.matmul/任意维按后两维做矩阵乘是通用场景含注意力等4D情况torch.einsum任意维按表达式指定视表达式而定复杂收缩、多张量连乘mm和bmm是“我说清楚了我就要这个形状你别帮我猜”所以它们的前置检查更严格、开销更低matmul是“你按广播规则帮我推”灵活但需要多做一轮形状匹配。日常写模型matmul和是主力但在数据加载、损失计算这类形状固定的热点路径上用mm更容易暴露维度错误也少一层推导。2. 矩阵乘法API选型从形状推演到实际写法2.1 用shape对齐法手推广播规则广播规则官方说法是“从尾部维度开始对齐缺失的补1为1的维度可扩展”。这话没错但落到四维张量上容易绕晕。我自己的土办法叫“右对齐补1法”把两个shape右对齐写下来左边不够的补1然后逐位比对某一位要么相等、要么有一个是1否则报错。举个例子注意力打分里最常见的写法import torch B, H, S, D 2, 8, 128, 64 q torch.randn(B, H, S, D) k torch.randn(B, H, S, D) # 右对齐 # q: (2, 8, 128, 64) # kT: (2, 8, 64, 128) scores q k.transpose(-2, -1) # - (2, 8, 128, 128)transpose(-2, -1)只交换最后两维前面(2, 8)原封不动。这里之所以能算是因为matmul只对“最后两维”执行矩阵乘语义前面的维度全部走广播。这个规则一定要刻进肌肉记忆matmul的全部广播发生在batch维度最后两维永远是(m, k) (k, n)。再看一个容易出事的例子a torch.randn(4, 1, 3, 5) b torch.randn(4, 6, 5, 2) c a b # - (4, 6, 3, 2)结果是(4, 6, 3, 2)因为a的第二维1被扩成了6。这看起来“很智能”但如果你的本意是4个batch、每个batch一组实际上多出了6倍计算量和6倍中间缓存。我就见过有人在这里显存翻了好几倍排查了半天以为是模型结构问题。实操心得写四维矩阵乘之前先把两个shape写在纸上右对齐比一遍比在代码里试错快得多。如果推出来某一维是1、结果却是别的数先问自己“这个扩展是我要的吗”。2.2 einsum把推导过程写成可读表达式einsum是我认为最值得花20分钟专门学的API。它的价值不在于性能小规模时略慢和matmul打平而在于它把形状推导变成了显式声明。你写bij,bjk-bik读代码的人一眼就知道谁是内维、谁被保留、谁被求和消掉。A torch.randn(4, 16, 32) B torch.randn(4, 32, 8) # 三种等价写法 c1 torch.bmm(A, B) c2 A B c3 torch.einsum(bij,bjk-bik, A, B)三者在数值上是一致的浮点误差范围内但可读性和调试体验完全不同。当表达式超过两个张量时einsum的优势就压倒性了# 双线性打分q (B,H,S,D), k (B,H,T,D), w (H,D,D) score torch.einsum(bhsd,bhtd,hde-bhst, q, k, w)这一行如果拆成matmul要写三四步中间还要transpose和unsqueeze很容易在某一步把维度搞反。当然代价是einsum对内存布局不敏感它会自己做必要的重排所以在大张量上不一定最优。我的实际做法是先用einsum把逻辑跑通跑对确认shape和数值无误如果这段进了profiler的热点再手工改写成matmul 提前contiguous的版本。顺带提一句einsum还有两个实用变体torch.einsum里加...可以做“省略前导维度”的写法比如...ij,...jk-...ik会自动匹配任意前缀batch维度适合写通用工具函数而ii-i这类表达式能直接取对角线。2.3 multi_dot与加法融合两个容易被忽略的API如果你要连乘三个以上矩阵torch.linalg.multi_dot值得知道。矩阵连乘的结合顺序会显著影响计算量(AB)C的代价是m*n*p m*p*q而A(BC)是n*p*q m*n*q当维度悬殊时能差一个数量级。multi_dot会用动态规划找较优的结合顺序。A torch.randn(1000, 2) B torch.randn(2, 1000) C torch.randn(1000, 2) naive (A B) C # 先算出 1000x1000 的大矩阵 smart torch.linalg.multi_dot([A, B, C]) # 走 1000x2 的低秩路径这个例子里A B会产生一个百万元素的中间结果而multi_dot会先算B C中间量只有2*2。两者最终结果一样但显存和耗时的差距是数量级的。这类低秩结构在手写模型的计算图里其实很常见。另一个是torch.addmm把beta * input alpha * (mat1 mat2)融合成一个kernel。全连接层y x W.T b用它比先matmul再add少一次中间张量读写。在小模型上看不出差别但在大batch的推理路径上是有意义的。同类还有torch.baddbmm处理批量情况。3. 转置与内存布局矩阵运算里最隐蔽的性能暗坑3.1 transpose之后到底发生了什么前面说了transpose只改步长不搬数据。这个设计的直接后果是转置后的张量交给矩阵乘法底层库可能要先做一次隐式拷贝。很多底层实现要求输入是行优先连续的遇到非连续输入会先规整化再计算这次拷贝你看不见但它实实在在吃带宽。a torch.randn(1024, 1024) b torch.randn(1024, 1024) b_t b.T print(b_t.is_contiguous()) # False # 下面这行的实际开销可能比 c1 高出一截 c a b_t在CPU上这个差距有时不明显在GPU上大矩阵时能看出差别。当然底层库也不傻很多情况下它能识别转置模式直接按转置方式读这就是所谓“转置友好的GEMM”。我的建议是别为了这点性能提前优化但要知道这个变量的存在。如果profiler显示某个matmul异常慢而它恰好接在一个permute后面那基本就是这里的问题加一个.contiguous()再测一次对比一下。顺带区分三个容易混的操作transpose交换任意两维permute一次性重排所有维度t()只针对2D转置。三者都不拷贝数据都返回非连续视图。3.2 view为什么报错reshape为什么能救view要求源张量在“按新形状解读时”内存仍然连续否则直接抛错。reshape则聪明一点能view就view不行就悄悄拷贝一份再给你。这就是为什么a.T.view(-1)必报错而a.T.reshape(-1)能过。a torch.arange(12).reshape(3, 4) a.T.view(-1) # RuntimeError a.T.reshape(-1) # 可以但发生了拷贝这个区别在debug时很关键。我见过有人在循环里对转置结果反复reshape每次都在偷偷拷贝循环一跑就是几百次速度自然上不去。正确做法是提前contiguous()一次后面想怎么view怎么view。a_t a.T.contiguous() # 一次拷贝 for _ in range(100): x a_t.view(-1) # 后续零成本3.3 批次维度放在哪一位更划算同样一个四维张量维度顺序不同性能可能差出可观的比例。核心原则是让最内层维度最后一维承载连续访问让需要一起做矩阵乘的两个维度尽量相邻。以多头注意力为例两种常见排布# 排布Abatch优先 q torch.randn(B, H, S, D) scores q q.transpose(-2, -1) # 排布B序列优先 q2 q.transpose(1, 2) # (B, S, H, D)排布A在做matmul时(S, D) (D, S)是标准形状transpose(-2,-1)只影响最后两维前两维连续排布B虽然语义上更贴近“每个token看所有头”但transpose之后中间维度变味了后续再和其他张量做广播时容易出岔子。我踩过的具体坑是把(B, H, S, D)先permute成(B, S, H, D)做了一层运算再permute回去做第二个矩阵乘来回折腾两次每个permute都制造了非连续视图两个matmul都多吃了拷贝开销。后来把维度顺序统一掉从头到尾只在最后两维做文章速度快了一点代码也短了一截。注意不要孤立地看单个shape。整条链路里所有矩阵运算的维度排布如果能保持一致省下的是成串的隐式拷贝。4. 线性代数求解inv、solve、分解怎么选4.1 用inv解方程的精度代价“解Ax b”这件事新手最直觉的写法是x torch.inverse(A) b。这个写法能用但有两个问题一是inverse里藏着一次分解你算出A^{-1}之后再乘b等于把分解的功夫白扔了一半二是显式求逆会放大条件数A越接近奇异误差越夸张。A torch.randn(64, 64) b torch.randn(64) x_bad torch.inverse(A) b x_good torch.linalg.solve(A, b)torch.linalg.solve内部用的是LU分解加前后代入不做显式求逆数值上更稳速度也更快。大批量情况下还有solve_ex可以顺便拿到分解结果复用或者lu_factorlu_solve分离两步。我拿一个故意构造的病态矩阵试过让A的最小奇异值和最大奇异值差出好几个数量级inv路线的相对误差比solve高出不少。练手数据上看不出问题实际模型里一旦遇到接近共线的特征这个误差是会累积的。实操心得代码里出现torch.inverse基本都可以问一句“我是不是在解方程”。如果确实只是要矩阵本身比如某些理论推导、协方差白化那另说如果目的是拿x一律换solve。4.2 SVD、QR、Cholesky的适用场景不同分解对应不同前提选错要么报错要么白算分解前提条件典型用途APILU方阵、可逆不严格要求解线性方程组linalg.lu_factorQR任意形状最小二乘、正交化linalg.qrCholesky对称正定协方差、高斯过程linalg.choleskySVD任意形状降维、伪逆、低秩近似linalg.svd特征值方阵谱分析、稳定性判断linalg.eig/eigh几个实务要点。eigh专门处理对称矩阵比通用eig快而且返回的特征值一定是实数如果你的矩阵理论上对称但浮点累加导致不对称先用(A A.T) / 2对称化再传进去稳定性会好很多。cholesky在遇到数值上不完全正定的矩阵时会失败实践中常加一个极小的对角扰动A eps * Ieps取1e-6量级具体要看数据尺度。svd有个容易忽略的参数full_matrices默认True会返回完整的U和Vh对于高瘦或矮胖矩阵来说是巨大的浪费。算伪逆、做低秩近似时把full_matricesFalse设上省下的显存很可观U, S, Vh torch.linalg.svd(X, full_matricesFalse) k 32 X_lowrank (U[:, :k] * S[:k]) Vh[:k, :]4.3 行列式、迹、秩这类标量量的坑torch.linalg.det在矩阵维度变大时数值上很脆弱因为它本质上是所有特征值之积量级很容易溢出或下溢。真要算对数行列式变分推断里常见用torch.linalg.slogdet返回符号和对数绝对值稳定性好得多sign, logabsdet torch.linalg.slogdet(A)torch.trace只能取主对角线之和想取批次内每个矩阵的迹得用torch.diagonal(A, dim1-2, dim2-1).sum(-1)直接trace会报维度错误。至于秩torch.linalg.matrix_rank内部靠SVD加阈值判断比你手写(A ! 0).sum()靠谱得多但也要理解它是数值意义下的秩——浮点误差会让本该线性相关的行“看起来”独立。检查共线性时我一般直接看奇异值的比值比看秩更直观。5. 矩阵运算与自动求导的配合5.1 矩阵乘的梯度到底怎么回传C A B的反向传播对A的梯度是grad_C B.T对B的梯度是A.T grad_C。这两条公式本身简单但里面藏着一个训练时经常被忽略的事实只要A或B有一个不需要梯度对应的那次矩阵乘就可以跳过。W torch.randn(128, 64, requires_gradTrue) x torch.randn(32, 128) # 不需要梯度 y x W # 反向时不需要算 dA这个机制在冻结骨干网络、做特征提取时非常有用。但要注意torch.no_grad()和requires_gradFalse是两回事前者连中间结果都不记录后者只是标记某个张量不需要梯度图还会照常构建。推理阶段一律套torch.no_grad()训练里冻结参数用后者别混用。另外一个小细节A.T在反向里也会用到如果A本来就是某个转置来的非连续张量两次非连续读会让梯度计算也变慢。这种情况下提前contiguous一次省的是两处开销。5.2 梯度检查小规模数值差分手写自定义算子或者怀疑某段矩阵运算的梯度对不对最直接的办法是数值差分。用double精度、取小维度、中心差分误差能压到1e-6量级import torch def grad_check(fn, x, eps1e-6): x x.double().requires_grad_(True) y fn(x) y.backward(torch.ones_like(y)) analytical x.grad.clone() numerical torch.zeros_like(x) with torch.no_grad(): it torch.ndindex if False else None for idx in torch.cartesian_prod(*[torch.arange(s) for s in x.shape]): idx tuple(idx.tolist()) xp x.clone(); xp[idx] eps xm x.clone(); xm[idx] - eps numerical[idx] (fn(xp).sum() - fn(xm).sum()) / (2 * eps) return (analytical - numerical).abs().max()这里一定要用doublefloat32的差分误差会盖过真实差异。还有一点差分检查只在fn是确定性函数时有效如果里面混了dropout、随机采样得先切到eval模式。我在排查一个问题时发现某个自定义的批量矩阵乘在batch维度上做了广播前向结果完全正确反向却是错的——因为广播维度上的梯度需要sum归并手写实现漏了这一步。数值差分一跑就暴露了最大误差是1e-1量级非常明显。这类bug只靠前向数值对拍是发现不了的。5.3 in-place改动带来的报错复盘矩阵运算和原地操作放在一起最典型的报错是“a leaf Variable that requires grad is being used in an in-place operation”或者反向时的“one of the variables needed for gradient computation has been modified by an inplace operation”。前一种好理解叶子张量不能直接原地改得在no_grad里改或者用.data不推荐。后一种更隐蔽你改的可能不是叶子而是某个中间结果但那个中间结果恰好被反向公式需要。a torch.randn(4, 4, requires_gradTrue) b torch.randn(4, 4, requires_gradTrue) c a b # 反向需要 a 和 b a 1 # 原地修改 a反向会报错正确的改法是a a 1建立一个新节点。如果确实需要在优化器更新里原地改参数那要在torch.no_grad()上下文里做这也是各种优化器内部的标准做法。总结一下我的判断口诀凡是被反向公式引用的张量一律不原地改只有纯前向的临时缓冲才考虑原地省内存。像relu_、add_这种就地版本在激活函数位置用问题不大但在矩阵运算的输入上用风险很高。6. 常见问题与排查速查6.1 形状类报错逐条对照报错信息关键词大概率原因处理办法mat1 and mat2 shapes cannot be multiplied内维不相等打印两个shape右对齐比对内维size mismatch广播不成立某维既不等也不为1补unsqueeze或squeezeExpected 2D or 3D input给mm/bmm传了错误维度换matmul或先整理维度view size is not compatible源张量非连续改reshape或先contiguousexpanded tensor ... in-place广播后的视图被原地改先clone再改CUDA out of memory矩阵乘前后广播把计算量放大了逐维检查是否有意外的1维扩展其中最后一条我想多说两句。显存爆炸不一定是模型太大很多时候是一次不经意的广播。比如(B, 1, S, D) (B, H, D, S)你以为在算H个头结果一个批次里把(B,1,...)扩成了(B,H,...)中间激活直接乘以H。这类问题用torch.cuda.max_memory_allocated()在可疑代码前后各打一次差值能立刻定位到具体哪一行。6.2 NaN与数值异常排查矩阵运算出NaN常见来源就那么几个求逆/求解遇到奇异矩阵检查矩阵条件数用solve替代inv必要时加eps正则。Cholesky对非正定矩阵加eps * I或改用eigh并裁掉负特征值。softmax前打分过大注意力里q k.T没除以sqrt(D)数值直接溢出。这一步缩放不是可选项是必需的。log里出现0或负数行列式、概率密度计算时用slogdet和log_softmax的稳定版本。半精度下的累加float16的表示范围窄大矩阵乘法内维很长时用torch.matmul的累加器通常还是float32但如果你手工分块累加并保持半精度就容易下溢。排查手法上我一般先定位到具体算子用torch.autograd.set_detect_anomaly(True)跑一次它会告诉你反向里哪一步先出现NaN。然后看这个算子的输入torch.isfinite(x).all()逐层往前扫找到第一个出现非有限值的位置。注意set_detect_anomaly会让训练变慢好几倍只在定位问题时开别留在生产代码里。6.3 精度与速度的开关清单矩阵运算的性能除了算法层面还有一组全局开关值得知道# 允许 TF32较新GPU上的张量核心精度模式默认可能已开 torch.backends.cuda.matmul.allow_tf32 True # 半精度矩阵乘的累加精度控制 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction True # 确定性模式复现实验用会变慢 torch.use_deterministic_algorithms(True)这组开关的取舍要看你当前在干什么。做研究要复现结果确定性模式该开就开做推理要吞吐TF32和半精度该开就开。但千万别在同一个项目里一会儿开一会儿关否则你会得到“同一份代码两次跑出不同loss”的灵异现象排查起来极其消耗精力。另外torch.compile对矩阵运算链的融合效果相当明显尤其是连续几个matmul加逐元素操作的组合。我试过一个包含5次中小规模矩阵乘的模块编译后整体耗时下降明显主要收益来自kernel融合和少了中间张量的显存往返。代价是首次编译有等待时间输入shape如果频繁变化还会触发重编译所以更适合形状固定的推理路径。最后分享一个我一直在用的调试习惯任何涉及三次以上矩阵运算的函数写完后先用小维度比如B2, S4, D8跑一遍把每一步的shape和dtype打出来确认整条链路符合预期再换成真实规模。小规模下报错信息更短、更好读而且能在几秒内跑完比在真实数据上等半分钟看一个size mismatch高效得多。这个习惯帮我省下的时间大概比学任何单个API都多。 
长河君 · 2026/10/26 11:31:44 · 阅读 1.2k PyTorch矩阵运算全解析：形状、内存布局与梯度回传
1. 先把“矩阵运算”这件事在PyTorch张量体系里说清楚PyTorch张量的数学运算里矩阵运算大概是被问得最多的一块。原因很朴素深度学习里几乎所有核心算子——全连接层、注意力里的QKV打分、卷积展开后的乘加、归一化里的协方差矩阵——最后都能归结成几次矩阵乘法。很多人调用API没问题但一遇到shape对不上、速度突然掉一半、反向传播报in-place错误就抓瞎了。问题不在API记不熟而在于脑子里没建立起“形状怎么对齐、内存怎么排布、梯度怎么回传”这三条并行的线。这篇内容面向的是已经能写几行PyTorch、但没系统梳理过矩阵运算的人。不管你是刚跑通MNIST想弄明白到底做了什么还是已经在调模型、想搞清楚为什么加一个transpose就慢了30%下面这些拆解应该都能对上你的痒处。我不会只列API清单重点是解释每个选择背后的代价——为什么这里用mm不用matmul为什么那里宁可多拷一次内存也要contiguous()为什么用inv解方程是新手最容易犯的错之一。1.1 张量不是数学符号是带步长描述的内存视图初学者最容易把张量想成一张二维表格实际上它是一块连续内存加一份“步长stride元数据”。一个(3, 4)的float32张量内存里就是12个连续浮点数步长是(4, 1)意思是“行方向每跳一步跨4个元素列方向每跳一步跨1个元素”。真正决定形状的是这份元数据而不是内存的物理排布。这个认知一旦建立很多诡异行为立刻就不诡异了。比如a.T根本不搬运任何数据它只是把步长从(4, 1)改成(1, 4)告诉后续算子“请按列优先来读”。所以转置是零成本的但也因此转置后的张量不是内存连续的。后面所有关于性能的坑根源都在这里。我习惯在调试时随手打一行print(x.shape, x.stride(), x.is_contiguous())比看什么都管用。尤其是接手别人代码的时候一个看起来平平无奇的(B, S, H, D)四维张量如果stride显示它是从某个permute过来的那它十有八九不是连续的后面接矩阵乘法就要留意开销。注意stride()在排查“为什么这段代码比预期慢”时信息量往往比shape()更大。1.2 元素级运算和矩阵级运算的分界线torch.mul、*、torch.add这类是逐元素运算两个张量的形状要么完全相同要么满足广播规则而torch.mm、torch.matmul、是矩阵乘法遵循的是(m, k) (k, n) - (m, n)这条内维收缩规则。这两套规则长得很像但完全不同*和只差一个符号犯错的成本却很高。一个典型的混淆场景你想算两个向量的内积写了a * b得到的是逐元素乘积还得再.sum()一次写a b一步到位。两者在(n,)形状下的结果一个是张量、一个是标量后续梯度形状也不同。更隐蔽的是a * b在广播下可能悄悄把你的(B, 1)扩成(B, B)代码不报错显存却突然炸了。我的经验是凡是涉及“内维求和”的语义一律用矩阵乘法API凡是涉及“同位对齐”的语义才用逐元素。写代码时强迫自己在注释里标一句“这里做的是内维收缩”能挡掉一大半形状事故。1.3 为什么必须区分mm、bmm、matmul这三个名字这三个API的功能有重叠但约束条件不同选错了轻则报错重则性能打折。API输入维度要求是否支持广播典型用途torch.mm严格2D否单次矩阵乘形状已确定torch.bmm严格3Dbatch数必须一致否批量矩阵乘batch对齐torch.matmul/任意维按后两维做矩阵乘是通用场景含注意力等4D情况torch.einsum任意维按表达式指定视表达式而定复杂收缩、多张量连乘mm和bmm是“我说清楚了我就要这个形状你别帮我猜”所以它们的前置检查更严格、开销更低matmul是“你按广播规则帮我推”灵活但需要多做一轮形状匹配。日常写模型matmul和是主力但在数据加载、损失计算这类形状固定的热点路径上用mm更容易暴露维度错误也少一层推导。2. 矩阵乘法API选型从形状推演到实际写法2.1 用shape对齐法手推广播规则广播规则官方说法是“从尾部维度开始对齐缺失的补1为1的维度可扩展”。这话没错但落到四维张量上容易绕晕。我自己的土办法叫“右对齐补1法”把两个shape右对齐写下来左边不够的补1然后逐位比对某一位要么相等、要么有一个是1否则报错。举个例子注意力打分里最常见的写法import torch B, H, S, D 2, 8, 128, 64 q torch.randn(B, H, S, D) k torch.randn(B, H, S, D) # 右对齐 # q: (2, 8, 128, 64) # kT: (2, 8, 64, 128) scores q k.transpose(-2, -1) # - (2, 8, 128, 128)transpose(-2, -1)只交换最后两维前面(2, 8)原封不动。这里之所以能算是因为matmul只对“最后两维”执行矩阵乘语义前面的维度全部走广播。这个规则一定要刻进肌肉记忆matmul的全部广播发生在batch维度最后两维永远是(m, k) (k, n)。再看一个容易出事的例子a torch.randn(4, 1, 3, 5) b torch.randn(4, 6, 5, 2) c a b # - (4, 6, 3, 2)结果是(4, 6, 3, 2)因为a的第二维1被扩成了6。这看起来“很智能”但如果你的本意是4个batch、每个batch一组实际上多出了6倍计算量和6倍中间缓存。我就见过有人在这里显存翻了好几倍排查了半天以为是模型结构问题。实操心得写四维矩阵乘之前先把两个shape写在纸上右对齐比一遍比在代码里试错快得多。如果推出来某一维是1、结果却是别的数先问自己“这个扩展是我要的吗”。2.2 einsum把推导过程写成可读表达式einsum是我认为最值得花20分钟专门学的API。它的价值不在于性能小规模时略慢和matmul打平而在于它把形状推导变成了显式声明。你写bij,bjk-bik读代码的人一眼就知道谁是内维、谁被保留、谁被求和消掉。A torch.randn(4, 16, 32) B torch.randn(4, 32, 8) # 三种等价写法 c1 torch.bmm(A, B) c2 A B c3 torch.einsum(bij,bjk-bik, A, B)三者在数值上是一致的浮点误差范围内但可读性和调试体验完全不同。当表达式超过两个张量时einsum的优势就压倒性了# 双线性打分q (B,H,S,D), k (B,H,T,D), w (H,D,D) score torch.einsum(bhsd,bhtd,hde-bhst, q, k, w)这一行如果拆成matmul要写三四步中间还要transpose和unsqueeze很容易在某一步把维度搞反。当然代价是einsum对内存布局不敏感它会自己做必要的重排所以在大张量上不一定最优。我的实际做法是先用einsum把逻辑跑通跑对确认shape和数值无误如果这段进了profiler的热点再手工改写成matmul 提前contiguous的版本。顺带提一句einsum还有两个实用变体torch.einsum里加...可以做“省略前导维度”的写法比如...ij,...jk-...ik会自动匹配任意前缀batch维度适合写通用工具函数而ii-i这类表达式能直接取对角线。2.3 multi_dot与加法融合两个容易被忽略的API如果你要连乘三个以上矩阵torch.linalg.multi_dot值得知道。矩阵连乘的结合顺序会显著影响计算量(AB)C的代价是m*n*p m*p*q而A(BC)是n*p*q m*n*q当维度悬殊时能差一个数量级。multi_dot会用动态规划找较优的结合顺序。A torch.randn(1000, 2) B torch.randn(2, 1000) C torch.randn(1000, 2) naive (A B) C # 先算出 1000x1000 的大矩阵 smart torch.linalg.multi_dot([A, B, C]) # 走 1000x2 的低秩路径这个例子里A B会产生一个百万元素的中间结果而multi_dot会先算B C中间量只有2*2。两者最终结果一样但显存和耗时的差距是数量级的。这类低秩结构在手写模型的计算图里其实很常见。另一个是torch.addmm把beta * input alpha * (mat1 mat2)融合成一个kernel。全连接层y x W.T b用它比先matmul再add少一次中间张量读写。在小模型上看不出差别但在大batch的推理路径上是有意义的。同类还有torch.baddbmm处理批量情况。3. 转置与内存布局矩阵运算里最隐蔽的性能暗坑3.1 transpose之后到底发生了什么前面说了transpose只改步长不搬数据。这个设计的直接后果是转置后的张量交给矩阵乘法底层库可能要先做一次隐式拷贝。很多底层实现要求输入是行优先连续的遇到非连续输入会先规整化再计算这次拷贝你看不见但它实实在在吃带宽。a torch.randn(1024, 1024) b torch.randn(1024, 1024) b_t b.T print(b_t.is_contiguous()) # False # 下面这行的实际开销可能比 c1 高出一截 c a b_t在CPU上这个差距有时不明显在GPU上大矩阵时能看出差别。当然底层库也不傻很多情况下它能识别转置模式直接按转置方式读这就是所谓“转置友好的GEMM”。我的建议是别为了这点性能提前优化但要知道这个变量的存在。如果profiler显示某个matmul异常慢而它恰好接在一个permute后面那基本就是这里的问题加一个.contiguous()再测一次对比一下。顺带区分三个容易混的操作transpose交换任意两维permute一次性重排所有维度t()只针对2D转置。三者都不拷贝数据都返回非连续视图。3.2 view为什么报错reshape为什么能救view要求源张量在“按新形状解读时”内存仍然连续否则直接抛错。reshape则聪明一点能view就view不行就悄悄拷贝一份再给你。这就是为什么a.T.view(-1)必报错而a.T.reshape(-1)能过。a torch.arange(12).reshape(3, 4) a.T.view(-1) # RuntimeError a.T.reshape(-1) # 可以但发生了拷贝这个区别在debug时很关键。我见过有人在循环里对转置结果反复reshape每次都在偷偷拷贝循环一跑就是几百次速度自然上不去。正确做法是提前contiguous()一次后面想怎么view怎么view。a_t a.T.contiguous() # 一次拷贝 for _ in range(100): x a_t.view(-1) # 后续零成本3.3 批次维度放在哪一位更划算同样一个四维张量维度顺序不同性能可能差出可观的比例。核心原则是让最内层维度最后一维承载连续访问让需要一起做矩阵乘的两个维度尽量相邻。以多头注意力为例两种常见排布# 排布Abatch优先 q torch.randn(B, H, S, D) scores q q.transpose(-2, -1) # 排布B序列优先 q2 q.transpose(1, 2) # (B, S, H, D)排布A在做matmul时(S, D) (D, S)是标准形状transpose(-2,-1)只影响最后两维前两维连续排布B虽然语义上更贴近“每个token看所有头”但transpose之后中间维度变味了后续再和其他张量做广播时容易出岔子。我踩过的具体坑是把(B, H, S, D)先permute成(B, S, H, D)做了一层运算再permute回去做第二个矩阵乘来回折腾两次每个permute都制造了非连续视图两个matmul都多吃了拷贝开销。后来把维度顺序统一掉从头到尾只在最后两维做文章速度快了一点代码也短了一截。注意不要孤立地看单个shape。整条链路里所有矩阵运算的维度排布如果能保持一致省下的是成串的隐式拷贝。4. 线性代数求解inv、solve、分解怎么选4.1 用inv解方程的精度代价“解Ax b”这件事新手最直觉的写法是x torch.inverse(A) b。这个写法能用但有两个问题一是inverse里藏着一次分解你算出A^{-1}之后再乘b等于把分解的功夫白扔了一半二是显式求逆会放大条件数A越接近奇异误差越夸张。A torch.randn(64, 64) b torch.randn(64) x_bad torch.inverse(A) b x_good torch.linalg.solve(A, b)torch.linalg.solve内部用的是LU分解加前后代入不做显式求逆数值上更稳速度也更快。大批量情况下还有solve_ex可以顺便拿到分解结果复用或者lu_factorlu_solve分离两步。我拿一个故意构造的病态矩阵试过让A的最小奇异值和最大奇异值差出好几个数量级inv路线的相对误差比solve高出不少。练手数据上看不出问题实际模型里一旦遇到接近共线的特征这个误差是会累积的。实操心得代码里出现torch.inverse基本都可以问一句“我是不是在解方程”。如果确实只是要矩阵本身比如某些理论推导、协方差白化那另说如果目的是拿x一律换solve。4.2 SVD、QR、Cholesky的适用场景不同分解对应不同前提选错要么报错要么白算分解前提条件典型用途APILU方阵、可逆不严格要求解线性方程组linalg.lu_factorQR任意形状最小二乘、正交化linalg.qrCholesky对称正定协方差、高斯过程linalg.choleskySVD任意形状降维、伪逆、低秩近似linalg.svd特征值方阵谱分析、稳定性判断linalg.eig/eigh几个实务要点。eigh专门处理对称矩阵比通用eig快而且返回的特征值一定是实数如果你的矩阵理论上对称但浮点累加导致不对称先用(A A.T) / 2对称化再传进去稳定性会好很多。cholesky在遇到数值上不完全正定的矩阵时会失败实践中常加一个极小的对角扰动A eps * Ieps取1e-6量级具体要看数据尺度。svd有个容易忽略的参数full_matrices默认True会返回完整的U和Vh对于高瘦或矮胖矩阵来说是巨大的浪费。算伪逆、做低秩近似时把full_matricesFalse设上省下的显存很可观U, S, Vh torch.linalg.svd(X, full_matricesFalse) k 32 X_lowrank (U[:, :k] * S[:k]) Vh[:k, :]4.3 行列式、迹、秩这类标量量的坑torch.linalg.det在矩阵维度变大时数值上很脆弱因为它本质上是所有特征值之积量级很容易溢出或下溢。真要算对数行列式变分推断里常见用torch.linalg.slogdet返回符号和对数绝对值稳定性好得多sign, logabsdet torch.linalg.slogdet(A)torch.trace只能取主对角线之和想取批次内每个矩阵的迹得用torch.diagonal(A, dim1-2, dim2-1).sum(-1)直接trace会报维度错误。至于秩torch.linalg.matrix_rank内部靠SVD加阈值判断比你手写(A ! 0).sum()靠谱得多但也要理解它是数值意义下的秩——浮点误差会让本该线性相关的行“看起来”独立。检查共线性时我一般直接看奇异值的比值比看秩更直观。5. 矩阵运算与自动求导的配合5.1 矩阵乘的梯度到底怎么回传C A B的反向传播对A的梯度是grad_C B.T对B的梯度是A.T grad_C。这两条公式本身简单但里面藏着一个训练时经常被忽略的事实只要A或B有一个不需要梯度对应的那次矩阵乘就可以跳过。W torch.randn(128, 64, requires_gradTrue) x torch.randn(32, 128) # 不需要梯度 y x W # 反向时不需要算 dA这个机制在冻结骨干网络、做特征提取时非常有用。但要注意torch.no_grad()和requires_gradFalse是两回事前者连中间结果都不记录后者只是标记某个张量不需要梯度图还会照常构建。推理阶段一律套torch.no_grad()训练里冻结参数用后者别混用。另外一个小细节A.T在反向里也会用到如果A本来就是某个转置来的非连续张量两次非连续读会让梯度计算也变慢。这种情况下提前contiguous一次省的是两处开销。5.2 梯度检查小规模数值差分手写自定义算子或者怀疑某段矩阵运算的梯度对不对最直接的办法是数值差分。用double精度、取小维度、中心差分误差能压到1e-6量级import torch def grad_check(fn, x, eps1e-6): x x.double().requires_grad_(True) y fn(x) y.backward(torch.ones_like(y)) analytical x.grad.clone() numerical torch.zeros_like(x) with torch.no_grad(): it torch.ndindex if False else None for idx in torch.cartesian_prod(*[torch.arange(s) for s in x.shape]): idx tuple(idx.tolist()) xp x.clone(); xp[idx] eps xm x.clone(); xm[idx] - eps numerical[idx] (fn(xp).sum() - fn(xm).sum()) / (2 * eps) return (analytical - numerical).abs().max()这里一定要用doublefloat32的差分误差会盖过真实差异。还有一点差分检查只在fn是确定性函数时有效如果里面混了dropout、随机采样得先切到eval模式。我在排查一个问题时发现某个自定义的批量矩阵乘在batch维度上做了广播前向结果完全正确反向却是错的——因为广播维度上的梯度需要sum归并手写实现漏了这一步。数值差分一跑就暴露了最大误差是1e-1量级非常明显。这类bug只靠前向数值对拍是发现不了的。5.3 in-place改动带来的报错复盘矩阵运算和原地操作放在一起最典型的报错是“a leaf Variable that requires grad is being used in an in-place operation”或者反向时的“one of the variables needed for gradient computation has been modified by an inplace operation”。前一种好理解叶子张量不能直接原地改得在no_grad里改或者用.data不推荐。后一种更隐蔽你改的可能不是叶子而是某个中间结果但那个中间结果恰好被反向公式需要。a torch.randn(4, 4, requires_gradTrue) b torch.randn(4, 4, requires_gradTrue) c a b # 反向需要 a 和 b a 1 # 原地修改 a反向会报错正确的改法是a a 1建立一个新节点。如果确实需要在优化器更新里原地改参数那要在torch.no_grad()上下文里做这也是各种优化器内部的标准做法。总结一下我的判断口诀凡是被反向公式引用的张量一律不原地改只有纯前向的临时缓冲才考虑原地省内存。像relu_、add_这种就地版本在激活函数位置用问题不大但在矩阵运算的输入上用风险很高。6. 常见问题与排查速查6.1 形状类报错逐条对照报错信息关键词大概率原因处理办法mat1 and mat2 shapes cannot be multiplied内维不相等打印两个shape右对齐比对内维size mismatch广播不成立某维既不等也不为1补unsqueeze或squeezeExpected 2D or 3D input给mm/bmm传了错误维度换matmul或先整理维度view size is not compatible源张量非连续改reshape或先contiguousexpanded tensor ... in-place广播后的视图被原地改先clone再改CUDA out of memory矩阵乘前后广播把计算量放大了逐维检查是否有意外的1维扩展其中最后一条我想多说两句。显存爆炸不一定是模型太大很多时候是一次不经意的广播。比如(B, 1, S, D) (B, H, D, S)你以为在算H个头结果一个批次里把(B,1,...)扩成了(B,H,...)中间激活直接乘以H。这类问题用torch.cuda.max_memory_allocated()在可疑代码前后各打一次差值能立刻定位到具体哪一行。6.2 NaN与数值异常排查矩阵运算出NaN常见来源就那么几个求逆/求解遇到奇异矩阵检查矩阵条件数用solve替代inv必要时加eps正则。Cholesky对非正定矩阵加eps * I或改用eigh并裁掉负特征值。softmax前打分过大注意力里q k.T没除以sqrt(D)数值直接溢出。这一步缩放不是可选项是必需的。log里出现0或负数行列式、概率密度计算时用slogdet和log_softmax的稳定版本。半精度下的累加float16的表示范围窄大矩阵乘法内维很长时用torch.matmul的累加器通常还是float32但如果你手工分块累加并保持半精度就容易下溢。排查手法上我一般先定位到具体算子用torch.autograd.set_detect_anomaly(True)跑一次它会告诉你反向里哪一步先出现NaN。然后看这个算子的输入torch.isfinite(x).all()逐层往前扫找到第一个出现非有限值的位置。注意set_detect_anomaly会让训练变慢好几倍只在定位问题时开别留在生产代码里。6.3 精度与速度的开关清单矩阵运算的性能除了算法层面还有一组全局开关值得知道# 允许 TF32较新GPU上的张量核心精度模式默认可能已开 torch.backends.cuda.matmul.allow_tf32 True # 半精度矩阵乘的累加精度控制 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction True # 确定性模式复现实验用会变慢 torch.use_deterministic_algorithms(True)这组开关的取舍要看你当前在干什么。做研究要复现结果确定性模式该开就开做推理要吞吐TF32和半精度该开就开。但千万别在同一个项目里一会儿开一会儿关否则你会得到“同一份代码两次跑出不同loss”的灵异现象排查起来极其消耗精力。另外torch.compile对矩阵运算链的融合效果相当明显尤其是连续几个matmul加逐元素操作的组合。我试过一个包含5次中小规模矩阵乘的模块编译后整体耗时下降明显主要收益来自kernel融合和少了中间张量的显存往返。代价是首次编译有等待时间输入shape如果频繁变化还会触发重编译所以更适合形状固定的推理路径。最后分享一个我一直在用的调试习惯任何涉及三次以上矩阵运算的函数写完后先用小维度比如B2, S4, D8跑一遍把每一步的shape和dtype打出来确认整条链路符合预期再换成真实规模。小规模下报错信息更短、更好读而且能在几秒内跑完比在真实数据上等半分钟看一个size mismatch高效得多。这个习惯帮我省下的时间大概比学任何单个API都多。