PyTorch中CPU矩阵乘法快于MPS的原因及MPS训练优势咨询

我正在使用PyTorch学习深度学习，当前正熟悉Tensor（张量）操作。为对比CPU与MPS上的矩阵乘法（matmult）耗时，我使用以下代码测试（芯片为2021款M1 Pro）：

import time
import torch

device = "mps"

torch.manual_seed(1234)
TENSOR_A_CPU = torch.rand(50, 50)
TENSOR_B_CPU = torch.rand(50, 50)

torch.manual_seed(1234)
TENSOR_A_MPS = torch.rand(50, 50).to(device)
TENSOR_B_MPS = torch.rand(50, 50).to(device)

start_time = time.time()
torch.matmul(TENSOR_A_CPU, TENSOR_B_CPU)
print("CPU : --- %s seconds ---" % (time.time() - start_time))

start_time = time.time()
torch.matmul(TENSOR_A_MPS, TENSOR_B_MPS)
print("MPS : --- %s seconds ---" % (time.time() - start_time))

测试结果为：

CPU : --- 7.510185241699219e-05 seconds ---
MPS : --- 0.0003561973571777344 seconds ---

可见CPU计算耗时远低于MPS，与我的预期相反。我想了解MPS是否在深度学习训练的其他流程中更具优化优势？若确实如此则使用MPS更合理，否则我可能全程使用CPU，望解答。

解答 #

• 小张量计算的开销问题：你测试用的是50×50的极小张量，MPS的优势根本没发挥出来。GPU（包括Apple Silicon的MPS后端）启动计算任务时，本身就有调度、数据传输的额外开销，小任务的计算时间甚至赶不上这些开销，而CPU处理小计算的额外成本更低，所以才会出现CPU更快的结果。
• 大张量与批量训练才是MPS的主场：深度学习训练中，我们通常处理的是大尺寸张量（比如2048×2048的权重矩阵）或者大批次的样本数据，此时MPS的并行计算能力会被彻底激活。比如在CNN训练里，大尺寸的卷积操作、大批次的全连接层矩阵乘法，MPS的速度会远超CPU。
• 全流程训练的综合优化：除了核心的矩阵乘法，PyTorch的MPS后端在自动微分、批量归一化、激活函数等训练常用操作上都做了专门优化。当这些操作组合成完整的训练流程时，MPS的整体效率会比CPU高很多，不是单个小操作能比的。
• 统一内存架构的带宽优势：Apple Silicon采用统一内存架构，大张量数据不需要在CPU内存和GPU显存之间来回拷贝，数据传输效率比传统独立GPU架构高很多，这在处理大模型或大尺寸数据集时优势格外明显。

如果你想验证MPS的真实优势，可以把张量尺寸放大到2048×2048，或者构建一个简单的CNN模型用大批次数据训练，此时就能看到MPS的速度碾压CPU的效果。

内容的提问来源于stack exchange，提问作者Kadir