2026/5/17 1:33:49
网站建设
项目流程
怎么做类似清风dj网站,网站建设必须提供良好的用户体验,如何制作短视频教程,app模板免费下载张量并行 (Tensor Parallelism, TP) 深度解析
1. TP 只能用于 Transformer 吗#xff1f;
答案是#xff1a;不#xff0c;但它在 Transformer 上用得最多#xff0c;也最有效。
咱们从 CV 的角度来类比。TP 的核心思想是“拆分矩阵乘法”。任何包含巨大矩阵乘法 (YX⋅WY …张量并行 (Tensor Parallelism, TP) 深度解析1. TP 只能用于 Transformer 吗答案是不但它在 Transformer 上用得最多也最有效。咱们从 CV 的角度来类比。TP 的核心思想是“拆分矩阵乘法”。任何包含巨大矩阵乘法 (YX⋅WY X \cdot WYX⋅W) 的模型理论上都可以用 TP。比如一个超大的全连接网络 (MLP)你完全可以把它那一层巨大的权重矩阵WWW拆开存到不同卡上。为什么大家都在说 TP Transformer这是因为 Transformer 的结构Multi-Head Attention 和 FFN简直是为 TP天造地设的。Transformer 的每一层都包含两个主要的子模块Multi-Head Attention (MHA)Feed-Forward Network (FFN)这两个模块都有一个共同的神奇特性可以把它们切成两半让两张卡各自独立跑完大部分计算只在最后一步进行一次通信就能得到正确结果。这种特性极大地减少了通信频率让 TP 变得实用。如果一个模型结构导致每做一步小计算都要通信一次那通信开销会大到让并行失去意义。所以TP 是一种通用的技术但 Transformer 的结构完美契合了 TP 的优势成为了它的最佳拍档。2. TP 的详细数据流是怎样的你问到了资料中没有展开的细节这部分是理解 TP 的关键。资料里提到的“切分矩阵”、“算前16个头”、“立刻通信拼起来”具体是怎么发生的呢我们以2 张显卡 (GPU 0, GPU 1)为例分别看看 FFN 和 MHA 的数据流。场景一Feed-Forward Network (FFN) 的数据流FFN 通常是一个两层的 MLP先放大维度再缩小回原维度。假设输入是XXX两个权重矩阵是W1W_1W1和W2W_2W2。第一步列切分 (Column Parallelism) 处理W1W_1W1复制输入GPU 0 和 GPU 1 拿到完全一样的输入数据XXX。切分矩阵巨大的W1W_1W1被竖着切成两半。GPU 0 拿左半边W1,LW_{1,L}W1,L。GPU 1 拿右半边W1,RW_{1,R}W1,R。独立计算GPU 0 算YLReLU(X⋅W1,L)Y_L \text{ReLU}(X \cdot W_{1,L})YLReLU(X⋅W1,L)。GPU 1 算YRReLU(X⋅W1,R)Y_R \text{ReLU}(X \cdot W_{1,R})YRReLU(X⋅W1,R)。关键点到这里两张卡各自拿着输出向量的一半不需要通信。第二步行切分 (Row Parallelism) 处理W2W_2W2为了让最后的结果能加起来第二个矩阵W2W_2W2必须横着切。输入GPU 0 的输入是YLY_LYLGPU 1 的输入是YRY_RYR。切分矩阵W2W_2W2被横着切成两半。GPU 0 拿上半边W2,UpW_{2,Up}W2,Up。GPU 1 拿下半边W2,DownW_{2,Down}W2,Down。独立计算GPU 0 算ZLYL⋅W2,UpZ_L Y_L \cdot W_{2,Up}ZLYL⋅W2,Up。GPU 1 算ZRYR⋅W2,DownZ_R Y_R \cdot W_{2,Down}ZRYR⋅W2,Down。关键通信 (All-Reduce)现在每张卡算出了最终结果的一部分。根据矩阵乘法规则最终结果ZZLZRZ Z_L Z_RZZLZR。动作两张卡进行一次All-Reduce (Sum)通信把各自的ZLZ_LZL和ZRZ_RZR加起来。结果通信结束后GPU 0 和 GPU 1 都得到了完整的、正确的结果ZZZ。FFN 总结输入复制 - 第一层列切分 - 第二层行切分 -最后做一次 All-Reduce。场景二Multi-Head Attention (MHA) 的数据流这个更直观就是资料里说的“CV人的类比…撕成两半”。假设有 32 个头。“头”的分配GPU 0 负责计算第 1-16 个头。GPU 1 负责计算第 17-32 个头。切分 Q, K, V 投影层 (列切分)计算 Attention 需要把输入XXX投影成 Query, Key, Value。对应的权重WQ,WK,WVW_Q, W_K, W_VWQ,WK,WV也是巨大的矩阵。我们像 FFN 第一层一样把它们按列切分。GPU 0 只拿属于前 16 个头的权重GPU 1 拿后 16 个头的。两张卡各自算出自己那 16 个头的Q,K,VQ, K, VQ,K,V。独立计算 AttentionGPU 0 用自己的Q,K,VQ, K, VQ,K,V算它那 16 个头的注意力结果。GPU 1 同理算它那 16 个头的结果。关键点这一步完全独立不需要通信。切分输出投影层WOW_OWO(行切分)Attention 的最后一步是把所有头的结果拼接起来乘一个输出矩阵WOW_OWO。为了配合前面的切分这个WOW_OWO必须按行切分。GPU 0 用它那 16 个头的结果乘以WOW_OWO的上半部分。GPU 1 用它那 16 个头的结果乘以WOW_OWO的下半部分。关键通信 (All-Reduce)和 FFN 一样最后需要把两部分结果加起来。动作进行一次All-Reduce (Sum)通信。结果两张卡都得到了完整的 MHA 输出。MHA 总结按“头”把 Q/K/V 权重列切分 - 独立算 Attention - 把输出权重WOW_OWO行切分 -最后做一次 All-Reduce。导师总结 (数据流)现在回看资料里的那句话“每一层计算结束必须立刻通信将两部分结果拼起来或相加”是不是清晰多了在 Transformer 的每一层里包含一个 MHA 块和一个 FFN 块。TP 的魔法流程是进入 MHA 块大家分头干活最后通信一次 (All-Reduce)得到完整结果。拿着这个结果进入 FFN 块再分头干活最后再通信一次 (All-Reduce)得到完整结果。这一层就算跑完了带着完整的结果进入下一层。这就是为什么 TP 的通信量巨大因为它在一个 Transformer 层内部就要进行两次高频的 All-Reduce 通信。这也是为什么资料强调 TP “通常仅限于同一台机器内部通过 NVLink 连接使用” 的原因。3. 模型是放在哪张卡上边的呢这是一个非常关键的问题理解了这一点你就真正掌握了张量并行TP的精髓。在张量并行TP的模式下完整的模型并不存在于任何单一的显卡上。模型的权重参数是被“撕碎”了分散存储在参与 TP 的所有显卡上的。为了让你彻底明白我们还是用你熟悉的 CV 领域和之前讲过的并行方式做对比3.1 对比三种方式的模型位置数据并行 (DP)模型位置每张卡上都有一份完整的模型复制。CV 类比就像你和你的同学不同的 GPU每人都有一本一模一样的《新华字典》完整的 ResNet 模型你们各自查不同的字处理不同的数据。流水线并行 (PP)模型位置每张卡上存储模型的一部分层 (Layers)。CV 类比就像工厂流水线你负责装轮胎前几层他负责装发动机中间层最后一个人负责喷漆最后几层。大家手里都只有自己负责的那部分图纸。张量并行 (TP)你现在问的这个模型位置每一层的权重矩阵都被切分分散在不同卡上。没有任何一张卡拥有某一层完整的权重。CV 类比资料里那个比喻非常精准——“把一个巨大的卷积核撕成了两半两张卡各拿一半去算”。具体来说假设一个线性层的权重矩阵WWW大小是4096×40964096 \times 40964096×4096。如果你用 2 张卡做 TPGPU 0可能只存了这个矩阵的左半边4096×20484096 \times 20484096×2048。GPU 1存了右半边4096×20484096 \times 20484096×2048。3.2 导师总结 (模型位置)所以回答你的问题“模型是放在哪张卡上边的呢”答案是模型被打散了每个人手里都只有碎片。GPU 0持有第 1 层矩阵的碎片 A第 2 层矩阵的碎片 C第 3 层矩阵的碎片 E…GPU 1持有第 1 层矩阵的碎片 B第 2 层矩阵的碎片 D第 3 层矩阵的碎片 F…这就是为什么资料的总结表里提到 TP 的显存节省是“极高 (单层参数被拆解)”。因为每张卡只需要存储原模型1/N1/N1/N的参数量NNN是 TP 的并行度。但也正因为大家手里都只有碎片所以每次计算完必须立刻通信“将两部分结果拼起来”才能凑出完整的结果往下传。