Amazon ML Systems：Parallelism 与 Collectives，列切后为何接行切？

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 行/列命名按右乘Y=XW、W[in,out]定义：第一层沿out切列，第二层沿in切行；库中W[out,in]名字可能相反。
• 训练合同为MLP GELU(XW1)W2，X[B,S,H]，W1[H,F]、W2[F,H]，tensor并行p整除F；不假定原模型尺寸。
• 示例理想collective数学语义，通信延迟需按实际拓扑/消息size测量，不承诺线性加速。

解题路线

1. Data Parallel各rank放完整参数处理不同样本，反向同步梯度；batch增加影响优化与学习率合同。参数/梯度/optimizer分片如ZeRO/FSDP改变驻留与all-gather/reduce-scatter调度，不能说DP一定只同步梯度或一定能装下模型。
2. Tensor Parallel切同一层权重/激活：第一层W1=[W1^0...W1^(p−1)]，每rank拿相同X得到Z_i=XW1_i[B,S,F/p]。GELU逐元素，直接本地算，无需先拼全F。
3. 第二层W2沿输入维分块，每rank得U_i=GELU(Z_i)W2_i[B,S,H]。完整U=ΣU_i，forward通常做sum all-reduce。第二层bias应在求和后加一次；每rank先加完整bias再sum会放大p倍。第一层bias按输出维切分可本地加。
4. 反向也要按依赖关系通信：第二层回传到各Z_i的梯度本地可算，第一层回传到共享X的贡献需要sum。第一层输出不应经过跨全F的非局部操作后还假设完全无通信，GELU逐元素这个条件是关键。
5. Pipeline按层放到不同rank，微批流水推进并传边界激活/梯度；bubble、激活驻留、跨stage不均衡与调度一致性是代价。Tensor并行常更频繁小延迟通信，优先放高速互联域；DP/Pipeline可跨更宽拓扑但需实际预算。
6. All-reduce做规约再把结果给全部rank；all-gather拼各rank块；reduce-scatter先规约再分结果块；all-to-all交换不同目的块，常用于expert路由。规约不是拼接，rank要有一致group/次数/shape，否则死锁或错误。
7. 原理参考Megatron：https://arxiv.org/abs/1909.08053；collective语义参考NCCL官方https://docs.nvidia.com/deeplearning/nccl/archives/nccl_2303/user-guide/docs/usage/collectives.html。原Annapurna项目底层实现未知，不把NCCL当Neuron原使用库。

手动走一遍

教学X=[2,3]，W1是2×4矩阵两列组各[1,0],[0,1]，因此Z0=Z1=[2,3]。为易手算本例激活临时取恒等（实际正文为GELU），W2上两行[[1],[2]]、下两行[[3],[4]]，局部U0=8、U1=18，sum=26，与完整Z=[2,3,2,3]乘W2一致。bias=5只加一次得31；每rank加5再sum会错误为36。

规模与取舍

两层MLP算术约按p分摊，但共享X及all-reduce输出尺寸仍约BSH。理想ring all-reduce每rank传输约2(p−1)/p·message_bytes（发送量），时间还含约2(p−1)阶段启动开销；实际算法/拓扑可能不同。参数分片省相应驻留，峰值还含gather buffer与激活。

易错点与追问

容易失分的边界

• 未声明矩阵存储方向。
• GELU之前先sum不同hidden维。
• sum规约当拼接。
• 第二层bias重复p次。
• 把算术除p当总耗时除p。
• 每rank collective调用不一致。

继续深挖

• 列切第一层的backward为什么要同步X梯度？
• 全F LayerNorm插在两层间会增加什么通信？
• sequence parallel怎样改变激活驻留？
• DP/TP/PP怎样按拓扑分配？