MindSpore Shard：深度解读算子级并行

经验分享 Tech Blogs 分布式并行-Distributed Parallelsim
1

MindSpore Shard：深度解读算子级并行

如果把训练大模型比作指挥一支庞大的军队，那么 MindSpore 的 shard 接口就是你手中的指挥棒。它允许你跳过繁琐的底层细节，直接告诉系统：“这个关键阵地（算子），我要用这种特定的队形（切分策略）去攻克！”

这篇指南将带你深入理解 MindSpore 的算子级并行（Operator-level Parallelism），不仅知其然，更知其所以然。

1. 为什么我们需要 Shard？

在并行训练的世界里，通常有两种极端：

  • 全自动驾驶 (Auto Parallel)：你什么都不用管，系统帮你决定怎么切分模型。虽然方便，但 Cost Model（代价模型）并不总是完美的，有时候它选的“路”虽然理论代价低，但并不符合你的专家直觉。
  • 纯手动挡 (Data Parallel)：你只能简单地把数据分给各张卡，模型本身不动。这对于超大模型来说，显存根本不够用。

Shard 就像是半自动驾驶。你拥有“专家介入”的权力——对于那些你最了解、最关键的算子（比如巨大的矩阵乘法），你可以手动指定切分方式；而对于剩下的琐碎算子，则交给 MindSpore 自动推导。

一句话总结：shard 让你在“掌控力”和“便利性”之间找到了完美的平衡。

2. 深度揭秘：从策略到执行的旅程

shard 并不是简单地给算子打个标签，它触发了 MindSpore 内部一套复杂的策略传播与图编译机制。让我们拆解一下这个过程。

2.1 第一阶段：锚点植入（Anchor Injection）

当你调用 shard(fn, in_strategy) 时，你实际上是在计算图上打下了几个坚不可摧的锚点。

  • 绝对权威：这些被 shard 标记的算子，其输入和输出的张量布局（Layout）被永久锁定。在 C++ 后端的图优化阶段，系统会识别这些原语（prim::kPrimShard），并通过 SetInputLayout 和 SetOutputLayout 将策略绑定到对应的图节点上。
  • 硬约束：在后续的所有优化过程中，无论系统觉得其他策略多么诱人（比如通信开销更小），它都绝不敢修改你指定的策略。这就像是你在地图上钉下的钉子，路怎么修都可以，但这几个点必须经过。
  • 种子池 (Seed Pool)：这些成功设置了策略的算子会被存入一个特殊的集合（configured_ops），它们将成为下一阶段算法的初始驱动源。

2.2 第二阶段：切分传播（Sharding Propagation）

这是 shard 机制的核心魔法。系统不仅要满足你的要求，还要让整个图“跑通”。MindSpore 使用了一种基于 BFS（广度优先搜索） 的传播算法来实现这一目标。

  • 双向波纹效应：
    算法从种子池中的算子开始，向图的四周扩散（BFS）：

    • 顺流传播（Forward）：遍历当前算子的输出边。如果下游算子 B 未配置策略，系统会根据上游算子 A 的输出布局，为 B 选择一个最匹配的输入策略。
    • 逆流传播（Backward）：遍历当前算子的输入边。如果上游算子 C 未配置策略，系统会根据下游算子 D 的输入需求，反向推导 C 的输出策略。

  • 最小代价决策 (Greedy Cost Minimization)：
    在传播过程中，当系统需要为相邻的未配置算子选择策略时，它遵循一个核心原则：最小化重排布代价。

    1. 零通信优先：如果存在一种策略，使得数据不需要在卡间传输就能直接被下游使用（Layout 完全匹配），那么毫不犹豫地选择它。
    2. 最小通信次之：如果必须传输，则计算所有候选策略的重排布代价（Redistribution Cost），选择通信量最小的那个。

2.3 第三阶段：冲突解决与桥梁搭建（Conflict Resolution）

现实往往不完美。如果你的策略和模型的自然结构发生了冲突，或者你手动指定了两个相邻算子使用完全不同的策略，会发生什么？

  • 自动插入转换器（Redistribution）：
    MindSpore 不会报错，而是会充当“和事佬”。
    假设上游算子 A 产出的是“行切分”数据，而你强行规定下游算子 B 必须接收“列切分”数据。
    系统会自动在 A 和 B 之间插入一组通信算子（如 AllToAll、AllGather、Permute）。这组算子负责在运行时把数据从卡 A 搬运到卡 B，完成布局的转换。

  • 代价权衡：
    虽然系统能解决冲突，但转换是有代价的（时间、带宽）。shard 的艺术就在于：不仅要指定策略，还要尽量减少这种不必要的转换。

3. 实战心法：如何用好 Shard？

理解了原理，我们来看看怎么用。这里不罗列代码，而是讲“心法”。

3.1 心法一：抓住“大鱼”，放过“虾米”

不要试图给每个算子都 shard 一遍。

  • 大鱼：计算量大、参数多的算子（如 MatMul, Conv2D）。这些是性能瓶颈，值得你手动优化（比如做模型并行）。
  • 虾米：激活函数（ReLU）、逐元素操作（Add）。这些算子计算极快，通常跟随上游策略即可（数据并行），不需要你操心。
# [实战示例] 抓大放小
class Net(nn.Cell):
    def __init__(self):
        super().__init__()
        self.dense = nn.Dense(64, 64) # 大鱼：矩阵乘法
        self.relu = nn.ReLU()         # 虾米：激活函数

    def construct(self, x):
        # 1. 只有 Dense 这种重计算算子值得我们手动切分
        # 我们给它配置模型并行策略（假设4卡，参数切4份）
        x = shard(self.dense, in_strategy=((4, 1),), parameter_plan={"self.dense.weight": (1, 4)})(x)
        
        # 2. ReLU 很轻，不要管它。
        # MindSpore 会自动推导：既然上游 Dense 输出了切分后的数据，ReLU 就直接复用这个策略，零通信代价！
        x = self.relu(x) 
        return x

3.2 心法二：顺势而为，减少“搬运”

设计策略时，要顺应数据的流动方向。

  • Bad Case：第一层用模型并行（切参数），第二层突然强行切回数据并行（切 Batch），第三层又切回模型并行。这会导致每层之间都在疯狂通信（AllToAll），训练速度极慢。
  • Good Case：连续的几个层都保持模型并行，直到必须聚合时（比如 Loss 计算前）再统一转回数据并行。
# [实战示例] 顺势而为
# 假设我们构建一个多层感知机 (MLP)
class MLP(nn.Cell):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Dense(128, 128)
        self.fc2 = nn.Dense(128, 128)

    def construct(self, x):
        # Good: 连续使用模型并行，中间不需要转来转去
        # 第一层：输入(4,1) -> 权重(1,4) -> 输出(4,1) [注意：MatMul会自动推导输出策略]
        x = shard(self.fc1, in_strategy=((4, 1),), parameter_plan={"self.fc1.weight": (1, 4)})(x)
        
        # 第二层：继续接收(4,1)的输入。因为上游输出是(4,1)，这里不需要任何通信！
        x = shard(self.fc2, in_strategy=((4, 1),), parameter_plan={"self.fc2.weight": (1, 4)})(x)
        return x

3.3 心法三：利用 Layout 提升可读性

不要在代码里写满 (4, 1), (8, 1) 这种数字天书。使用 Layout 给维度起名字。

  • 把设备维度命名为 dp (Data Parallel) 和 mp (Model Parallel)。
  • 代码里写 layout("dp", "mp")，一眼就能看出这是“数据维走 dp，模型维走 mp”。
# [实战示例] 提升可读性
from mindspore.parallel import Layout

# 定义布局：8卡，4x2
layout = Layout((4, 2), ("dp", "mp"))

def attention_score(q, k):
    return ops.matmul(q, k)

# 不用 Layout: 
# shard(attention_score, in_strategy=((4, 1, 2), (4, 2, 1)))  # 谁知道哪个维度对应什么？

# 使用 Layout: 
# 假设 q: [Batch, Seq, Head], k: [Batch, Head, Seq]
# dp=Batch维度, mp=Head维度
in_strategy = (
    layout("dp", "None", "mp"),  # q: Batch切dp, Head切mp
    layout("dp", "mp", "None")   # k: Batch切dp, Head切mp
)
shard(attention_score, in_strategy=in_strategy) # 清晰明了！

4. 常见误区与避坑

  1. “自动”不是“全能”：
    虽然叫 AUTO_PARALLEL，但如果你 shard 得不合理（比如切分份数不能整除卡数），系统也没法帮你“圆”回来，会直接报错。

  2. PyNative 的遗憾：
    目前 shard 强依赖于静态图编译技术（因为要分析全图做传播），所以在 PyNative 模式（动态图）下暂时无法使用。

  3. 牵一发而动全身：
    你在网络中间改了一个算子的策略，可能会导致整个网络的策略发生“蝴蝶效应”般的剧变。如果不确定，建议先在小规模子网中验证。