Sharded GEMM与Backprop：行、列、内维三种切分

题意与假设

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

• 训练A为M×K、B为K×N，上游G为M×N；损失L=〈G,AB〉。NumPy模拟各rank计算和聚合，无实际网络collective。
• 代码三种单轴分片允许不均匀与空片；实际MLP starter和五个bug未公开，不将本训练写成其复原。

解题路线

1. 无分片梯度dA=GBᵀ、dB=AᵀG，先用shape推导再讨论collective；loss对每个rank不是默认独立mean，否则会错误多除rank数。
2. 按M行分A/G，各rank Y_r=A_rB，dA_r=G_rBᵀ，局部dB_r=A_rᵀG_r需sum/all-reduce。Y/dA按行concat，dB不是concat。
3. 按N列分B/G，各rank Y_r=AB_r，dB_r=AᵀG_r，局部dA_r=G_rB_rᵀ需sum/all-reduce。Y/dB按列concat，完整G若在每rank需确认切分一致。
4. 按K内维同时分A列/B行，各rank partialY=A_rB_r，再sum得到Y；收到完整G后dA_r=GB_rᵀ、dB_r=A_rᵀG，梯度按分片对应轴concat。forward少一次sum会漏其他分片贡献。
5. 真实分布式实现需要明确哪些结果保留sharded/replicated，all-reduce与reduce-scatter不同；避免无必要all-gather整个张量。MLP多层可组合列并行和行并行，非线性插在partial sum之前通常不等价，GELU(sum partial)≠sum GELU(partial)。

手动走一遍

M3、K5、N4，parts2内维分K为3/2，不整除也应覆盖恰好所有列/行。行分片两个rank的dB必须相加；每rank直接更新各自replicated B却不聚合会逐步分歧。1000组随机尺寸×三轴，用PyTorch autograd核对Y/dA/dB，包含parts大于维度的空片。

Python 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

import numpy as np

def sharded(a,b,gradient,parts,axis):
    a,b,g=map(np.asarray,(a,b,gradient))
    if a.ndim!=2 or b.ndim!=2 or a.shape[1]!=b.shape[0] or g.shape!=(a.shape[0],b.shape[1]):raise ValueError('shape')
    if type(parts) is not int or parts<1 or axis not in ('rows','columns','inner'):raise ValueError('partition')
    if axis=='rows':
        xs=np.array_split(a,parts,axis=0);gs=np.array_split(g,parts,axis=0)
        y=np.concatenate([x@b for x in xs],axis=0)
        da=np.concatenate([chunk@b.T for chunk in gs],axis=0)
        db=sum((x.T@chunk for x,chunk in zip(xs,gs)),np.zeros_like(b,dtype=np.result_type(a,b,g)))
    elif axis=='columns':
        bs=np.array_split(b,parts,axis=1);gs=np.array_split(g,parts,axis=1)
        y=np.concatenate([a@x for x in bs],axis=1)
        da=sum((chunk@x.T for x,chunk in zip(bs,gs)),np.zeros_like(a,dtype=np.result_type(a,b,g)))
        db=np.concatenate([a.T@chunk for chunk in gs],axis=1)
    else:
        xs=np.array_split(a,parts,axis=1);bs=np.array_split(b,parts,axis=0)
        y=sum((x@z for x,z in zip(xs,bs)),np.zeros_like(g,dtype=np.result_type(a,b,g)))
        da=np.concatenate([g@z.T for z in bs],axis=1)
        db=np.concatenate([x.T@g for x in xs],axis=0)
    return y,da,db

复杂度

计算总量仍约O(MKN)，理想均匀每rankO(MKN/P)，通信取决于轴：行分dB K×N聚合、列分dA M×K聚合、内维Y M×N聚合。NumPy模拟串行总工作不代表P倍加速；峰内存与replicated输入另计。

易错点与追问

容易失分的边界

• 梯度该sum却average，引入错误尺度。
• reshape后分片轴改变，shape可能相同但对应元素错位。
• partial matmul后先activation再sum，改变数学模型。

继续深挖

• 二维tensor parallel如何选择通信组与分片布局？
• 不同rank batch mask不均匀，loss分母怎样全局一致？