Alpa — 把张量/流水/数据并行统一成一道搜索题

是什么

Alpa 是 UC Berkeley RISELab 在 OSDI 2022 发表的自动并行编译器——你给它一段 JAX/PyTorch 模型代码和一组 GPU，它自己算出这模型怎么切才最快，自动生成分布式训练版本。

日常类比：你要把一座大工厂搬到一片厂区，传统做法是工程师手画施工图——哪条产线放哪栋厂房、哪个工序拆给谁。Alpa 的做法是给它一份”工艺流程图 + 厂区平面图”，它自己解一道优化题，吐出施工方案。

它的核心洞察是：长期被分开看的三种并行策略其实是同一棵决策树的不同枝。

数据并行（DP）：每张卡都拿完整模型，分批数据
张量并行（TP）：每张卡拿模型的一片，比如 megatron-lm 把矩阵按列切
流水线并行（PP）：每张卡负责模型的一段层，像 gpipe-2019 pipedream-2019 接力

Alpa 把它们统一成两层搜索：**算子内（intra-op）**决定每个算子怎么 SPMD 切（这一层覆盖 DP+TP），**算子间（inter-op）**决定怎么切流水线段（这一层覆盖 PP）。

为什么重要

不知道 Alpa，下面这些事都解释不通：

为什么 2023 年之后 PyTorch FSDP、Megatron-Core Auto-Parallel、torchtitan 都在做”自动并行”——Alpa 是第一个把它当编译器问题而不是运维问题来解的
为什么 Google GSPMD（内部）也走了类似路线——它和 Alpa 是同一时期、同一思路的两个工程实现
175B 参数的 GPT-3 用 Megatron-LM 训需要专家手调几周；Alpa 在 64 张 A100 上没有任何人工配置就能匹敌
在 Mixture-of-Experts、Wide-ResNet 这些没有公认手调方案的模型上，Alpa 比专家基线快 2.5x 到 9.7x——因为人想不到的切法机器能搜出来

简单说——Alpa 把”怎么并行”从架构师的脑子里抽出来，变成了编译器的输出。

核心要点

Alpa 的整体管线分三层：

IR 层：模型先编译成 XLA HLO 计算图（[[xla-compiler.md|XLA]]），所有决策都在这张图上做
Inter-op pass：用动态规划在图上切流水线段，并给每段分配一块”设备子网格”（device sub-mesh，比如 4×2 张卡）
Intra-op pass：每个子网格内用整数线性规划（ILP）求解每个算子的 SPMD 分片策略

最关键的设计是层级化分解——单层搜索是 NP 难的组合爆炸，分两层后每层都能在分钟级解出来。

Intra-op ILP 在搜什么：每个算子有几种候选分片（复制 / 按行切 / 按列切 / 部分和），相邻算子分片不一致需要插入通信（all-gather / reduce-scatter / all-to-all）。变量是”每个算子选哪种分片”，目标是最小化总通信量 + 计算时间。

Inter-op DP 在搜什么：把图切成 K 段，每段分到一块子网格上，目标是最小化”流水线 bubble + 段内执行时间”的最大值（瓶颈段决定吞吐）。

Auto-resharding 是粘合剂：当上一段的输出分片和下一段的输入分片不一致，编译器自动插入通信算子——这一步让两层搜索能解耦，否则每层都要顾及对方。

为什么是 ILP + DP 而不是单一求解器：算子内的搜索空间是”每个算子选一种分片 spec”，约束是相邻算子的通信代价——这天然适合整数线性规划，目标线性、变量是 0/1 选择。算子间的搜索空间是”在长链上找若干切点”，标准动态规划题型，子问题可以缓存复用。强行合一只会让两边都退化。

实践案例

案例 1：GPT-3 175B 自动匹敌 Megatron

跑 175B GPT 在 64 张 A100 上，Megatron-LM 的官方配方需要专家把 TP=8 / PP=8 / DP=1 写死在配置里。Alpa 不读这份配置，它的搜索器输出的策略几乎一模一样——TP 走层内大矩阵切，PP 切 Transformer 层段，DP 这维不切。吞吐持平。

意义不是”赶上专家”，而是专家想到的方案，编译器也能想到——下一次换硬件、换模型尺寸，不用重新调。

案例 2：MoE 上 2.5x 反超手调

Mixture-of-Experts 模型的专家层适合按专家维度切（all-to-all 通信），但前后的非专家层适合 TP。这种异质分片人很难调对。Alpa 的 ILP 自动在每个算子选最优分片，比 GShard 风格手调快 2.5x。

教训：当模型结构不是”一根筋的 Transformer”，自动并行的优势立刻显出来——人脑的搜索带宽不够。整个 MoE 层的策略空间大概有几千种组合，人能逐一比较的最多十几种。

案例 4：64 张 A100 上的搜索时间

64 卡集群上跑 GPT 175B，Alpa 的整个 inter-op DP + intra-op ILP 求解约 8-12 分钟完成；相比之下训练任务本身要跑数天，搜索开销可忽略。这是”自动并行能落地”的硬性前提——如果搜索要跑几小时，没人愿意用。

案例 3：Wide-ResNet 上 9.7x

Wide-ResNet 没有公认的”工业最优并行”。Alpa 的搜索找到一种inter-op 切层 + intra-op 按通道切的组合，比朴素 DP 快 9.7x。这种切法人类几乎不会想到——它要求同时考虑”层间瓶颈”和”通道维度通信代价”两个相互拉扯的约束。

踩过的坑

ILP 解得慢：GPT-3 175B 的 HLO 图节点上千，原版 ILP 求解几分钟。后续工作（如 Galvatron）用启发式裁剪状态空间。
静态形状假设：Alpa 假设输入形状编译期已知；动态序列长度（如变长 batch）需要 padding 或重编译。
成本模型靠 profile：Alpa 用 FLOPs + 带宽估通信/计算时间，真实硬件上时常偏 10-20%。在某些拓扑（NVLink 不均匀）上会出错策略。
没有运行时重规划：编译期定下来就不变。生产环境如果某张卡退化（thermal throttle），Alpa 不会重切——这是 Pathways / GSPMD-Runtime 后续要补的洞。
PyTorch 支持靠 torch_xla：原版 Alpa 只对 JAX 生态友好；PyTorch 用户走 torch_xla 桥接，性能损耗 5-10%。这是后来 PyTorch 自己做 FSDP / Auto-Parallel 的直接动因。
inter-op 和 intra-op 的解耦不是免费的：两层独立求解会丢一些跨层联合最优解——比如某种切法只有同时调整段数和段内分片才能逼出来。论文里坦承这是已知妥协，换来的是搜索可解。

适用 vs 不适用场景

适用：

大模型训练超出单卡显存，但没有现成手调配方（新架构、MoE、混合模态）
频繁换硬件 / 换模型尺寸，不想每次重新调并行
教学和研究：理解”并行策略是什么”最干净的视角就是 Alpa 的两层搜索

不适用：

推理（Alpa 只管训练；推理另有 AlpaServe 等工作）
动态形状、变长输入、异构 batch（编译期假设破灭）
极致追 0.5% 性能的工业训练——人工调更稳
单机多卡小模型（搜索的开销 > 收益）
强化学习这种前向反向交错很快变化的训练循环——Alpa 的静态编译模型不匹配

历史小故事（可跳过）

2018-2020：Mesh-TensorFlow / GShard 让用户手写SPMD 标注，已经把”切分作为编译器输入”这个想法埋下
2021：Google 发表 GSPMD（内部用），把 SPMD 编译做完整，但还需要人写 sharding hint
2022 OSDI：Alpa 第一次把搜索器也做了，从此 sharding hint 也不用写
2023 之后：PyTorch FSDP 走”hand-tuned 但简单 API”路线；Megatron-Core / torchtitan 抄 Alpa 的搜索思路；学术界继续做更快的搜索器（Galvatron、AceCoder）

Alpa 的论文和代码都是公开的，仓库地址 github.com/alpa-projects/alpa。

自己写实现的最小思路（如果你想动手）

如果想用 100 行代码理解 Alpa 的核心，可以做这样一个玩具版：

用 NetworkX 建一张 10 节点的”假计算图”——每个节点带 FLOPs 和输出张量形状
给每个节点定义 3 种候选 sharding：复制 / 按行切 2 路 / 按列切 2 路
用 PuLP（Python 的 ILP wrapper）写一个简化版 intra-op ILP——目标是相邻节点 sharding 不一致时加一个固定的 reshard 代价
用一维 DP 把图切成 K 段，每段独立跑上面的 ILP
比较 K=1（纯 SPMD）和 K=2/3（混合）的总代价

跑起来你会立刻发现——搜索器吐出的策略和你直觉里”应该这么切”几乎一致，只是机器把”为什么是这个”说得清楚。

和 Hindley-Milner 的精神类比

hindley-milner 让编译器替你推类型，不用人手标注；Alpa 让编译器替你推并行，不用人手切模型。两者都把”原本算工程师品味”的事情形式化成可判定的搜索。

不同之处：HM 靠统一（unification）这种语法级算法，Alpa 靠 ILP+DP 这种数值级算法——但思想是一样的：人定问题，机器找解。

学到什么

并行策略是搜索问题，不是模式问题——这个视角转换是 Alpa 最大的贡献
层级化分解是降复杂度的标准武器：单层 NP 难的搜索切两层各自分钟级可解
编译器 IR 是统一抽象：HLO 让 DP/TP/PP 长得一样（都是图上的变换），人才看不出来
auto-resharding 是粘合剂：通信代价从”人算”变成”编译器自动插入”，这是 SPMD 风格能跑通的前提
理论 → 算法 → 工程：和编程语言里的 hindley-milner 异曲同工——把人手做的事变成可判定的搜索

关联

megatron-lm —— Alpa 自动搜出来的策略在 GPT-3 上和 Megatron 手调几乎一致
gpipe-2019 —— Alpa 的 inter-op pass 本质是在搜 GPipe 风格的流水线切法
pipedream-2019 —— Alpa 默认用 1F1B 调度，继承 PipeDream 的 bubble 优化
deepspeed-zero —— ZeRO 切优化器状态，Alpa 切计算；二者可叠加
xla-compiler —— Alpa 整个搜索都在 XLA HLO 图上做
jax —— Alpa 第一公民语言；torch_xla 是二等公民
hindley-milner —— 同一种思想：把人手做的决定变成机器可解的搜索

反向链接

amdahl-law-1967 —— Amdahl 定律 — 串行比例决定并行加速比的上界
aurora-exascale-2024 —— Aurora 2024 — 不用 NVIDIA 也能造 2 EFLOPS 超算
big-little-2011 —— big.LITTLE — 让一颗芯片同时装快核和省电核
blink-2020 —— Blink — 按拓扑动态拼生成树替代 NCCL ring
cudnn-2014 —— cuDNN — 把卷积写成矩阵乘，让所有深度学习框架共享底层加速
deepspeed-zero —— DeepSpeed ZeRO — 微软优化大模型训练显存
gpipe-2019 —— GPipe — micro-batch 流水线让 GPU 排成生产线
hindley-milner —— Hindley-Milner — 编译器自己猜变量类型
jax —— JAX — Google 函数式数值计算
nvlink-nvswitch-2018 —— NVLink 2.0 + NVSwitch — 把 16 块 GPU 拼成一台机器
pipedream-2019 —— PipeDream — 1F1B 调度让流水线工位别空等
ps-li-2014 —— Parameter Server — 多机训练前 AllReduce 时代的工业标准
ring-allreduce-2017 —— Ring All-Reduce — 把 HPC 的环形规约搬进深度学习
sglang-2024 —— SGLang — 把 LLM 程序当成共享前缀的树来跑
taso-2019 —— TASO — 让机器自己发现深度学习图重写规则
tensorflow-osdi-2016 —— TensorFlow — 把神经网络拆成数据流图再跑到任何机器上
xla-compiler —— XLA — 给 TensorFlow / JAX 装一台真正的张量编译器
zero-2020 —— ZeRO 2020 — 把训练状态切成 N 份让万亿参数成为可能