深度架构：MindSpore 自动并行实现原理全解析

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

深度架构：MindSpore 自动并行实现原理全解析

MindSpore 自动并行的目标是实现“单机脚本，分布式训练”。开发者只需关注模型逻辑，而框架通过内置的代价模型（Cost Model）与搜索算法，自动寻找计算开销与通信开销平衡的最优切分方案。

一、 自动并行执行流水线

自动并行的核心逻辑起始于 StepAutoParallel 函数。其执行流程如下：

  1. 算子建模（OperatorInfo）：为计算图中的每个 CNode 创建 OperatorInfo 对象，枚举所有可能的切分策略候选集。
  2. 构建代价图（Cost Graph）：在算子间建立 Edge，用于评估因切分策略不一致导致的张量重分布（Redistribution）代价。
  3. 策略搜索：根据用户选择的模式（DP, RP/SAPP, 或 Sharding Propagation）在代价图中寻找最优解。
  4. 策略回写：将确定的策略（如 in_strategy）写回算子属性，并执行张量重排（Resharding）。

二、 三大搜索模式对比

MindSpore 根据集群规模和性能需求提供了三种搜索路径：

模式 核心机制 适用场景
动态规划 (DP) 遍历所有切分组合，通过代价递推寻找全局最优解。 中小规模模型，追求极致性能优化。
切分传播 (SP) 从用户指定的 shard 点出发，自动推导出其余算子的切分方式。 混合并行场景，用户希望部分干预策略。
递归规划 (RP/SAPP) 基于符号化建模与二分法，递归地在设备集群上分配算子策略。 超大规模模型（如万亿参数），搜索速度极快。

三、 SAPP (递归规划) 核心实现机制深度扩展

SAPP（Symbolic Automatic Parallel Planner）是 MindSpore 处理大规模并行（如盘古大模型）的“大杀器”。它通过将并行决策抽象为符号化操作，极大地降低了搜索复杂度。

1. 符号化图消除 (Graph Elimination)

在搜索前，SAPP 会简化计算图，剔除对策略影响较小的“冗余”算子。

  • 算子消除：如 Cast、Identity 等算子在不影响切分逻辑时会被消除。
  • 属性合并：在消除节点时，系统通过 ShapeMappingCombine 将原有的维度变换属性（如 Transpose 的轴变换）合并到相邻边上，确保维度映射关系的完整性。

实现逻辑伪代码：

// 简化自 rec_parse_graph.cc
void SimplifyGraph(graph) {
    for (auto &node : graph.nodes) {
        if (node.is_redundant) {
            // 将当前节点的变换(如Transpose)属性叠加到输出边的 mapping 中
            UpdateEdges(node.prev, node.next, node.transformation_attr); 
            RemoveNode(node); // 消除节点以减小搜索空间
        }
    }
}

2. 权重驱动排序 (SortByWeight)

RP 算法认为计算量大的算子应优先获得切分决策权。系统调用 SortByWeight 对算子排序，确保 MatMul、Conv2D 等高权重算子先于辅助算子决策。

3. 递归二分切分决策 (Recursive Partitioning)

SAPP 的核心循环基于 log_2(Devices) 轮递归。在每一轮中，系统决定是否在算子的某个轴上进行二分。

  • 轴选择决策：针对算子（如 MatMul），评估在 I 轴（行切）、J 轴（列切）或 K 轴（Reduce轴）进行切分的总代价。
  • 分治分配：每一轮递归，设备集群被对半拆分，直到所有设备都被分配到具体的切分轴上。

递归搜索伪代码：

// 简化自 rec_partition.cc
for (int loop = 0; loop < log2(total_devices); loop++) {
    auto sorted_nodes = SortByWeight(graph); //
    for (auto &node : sorted_nodes) {
        // 评估三个轴的切分成本：OpCost (算子开销) + RedisCost (重分布开销)
        Cost cost_I = Calc(node, Axis_I);
        Cost cost_J = Calc(node, Axis_J);
        
        node.strategy = ChooseMin(cost_I, cost_J); // 选择最优切分方向
        ApplyStrategyToTensor(node); // 同步状态到张量符号表达中
    }
}

4. 策略推导与合法性修正

搜索完成后，SAPP 得到的是符号比例（如 0.5 表示二分）。系统需调用 GenerateStrategy 将其转化为具体的 in_strategy。

  • 整除校验：如果切分数无法整除张量维度，系统会自动回退策略以确保执行安全。
  • 动态形状支持：针对 Dynamic Shape，系统会通过 HandleDynamicShapeFix 调整策略，避免因形状变化导致切分非法。

四、 维度变换处理：ReshapeDecompose

对于 Reshape 这种会彻底改变张量语义的算子，普通并行难以追踪维度。MindSpore 通过 ReshapeDecompose 建立映射：

  • 依赖关系映射：计算输出维度 out\_idx 对应输入维度 in\_idx 的缩放比例。
  • 跨算子对齐：确保即便中间插入了 Reshape，系统依然能识别出前后算子是否处于同一个逻辑轴上切分。

五、 显存红线保护：DevicesMemoryControl

显存是自动并行的硬约束。系统在搜索策略时会实时计算：

$$UsedMemory = \sum_{node} (InputSize + OutputSize + Workspace) \times DataType$$

如果预估内存超过物理限制，系统会触发警告并尝试回退策略（如减小模型并行的切分数），从而有效预防 OOM。

总结

MindSpore 的自动并行是编译器技术与并行算法的深度结晶。通过 SAPP 的递归分治，框架实现了在秒级时间内完成大规模集群的最优策略搜索。