FlexFlow Train：自动搜索分布式训练的最优并行策略

FlexFlow Train 在 GitHub 上拿到了 1,886 Star。

这是一个深度学习框架，专门解决分布式 DNN 训练里的并行策略问题。它不让你手动调数据并行还是模型并行，而是自动搜索最高效的组合方案。开发团队来自 CMU、Facebook、Los Alamos National Lab、MIT、Stanford 和 UCSD。

1、这玩意儿是干嘛的

分布式训练里，并行策略的选择直接决定性能和资源利用率。数据并行、模型并行、流水线并行，每种都有适用场景，组合起来更是复杂。网络层数一多，设备数量一大，手动设计策略几乎不可能覆盖所有可能性。

FlexFlow Train 的核心思路是把并行策略的搜索自动化。框架基于 MCMC 搜索算法，在给定的硬件配置下找出最优的并行化方案，覆盖算子划分、设备映射、通信调度等多个维度。搜索空间比传统方法更大，因为它同时考虑代数变换和并行化的联合优化。

它跟 PyTorch 兼容。你可以把现有的 PyTorch 模型导出成 FlexFlow 格式，然后直接调用自动调优功能来优化并行性能。不用重写模型定义，迁移成本很低。

2、为什么要用它

手动设计并行策略需要大量经验。模型结构一变，或者硬件配置不同，之前的方案可能就失效了。很多团队在这上面消耗了大量时间，结果拿到的还不是最优解。

FlexFlow Train 把这一层抽象出来。你只需要定义模型和硬件拓扑，框架自己跑搜索，输出一个执行策略。搜出来的结果可以导出保存，下次直接复用，不用重复计算。

它背后的论文发表在 OSDI 2022、MLSys 2019 和 ICML 2018，作者来自 CMU、Stanford、MIT 等机构。核心思路是联合优化代数变换和并行化，搜索空间覆盖传统数据并行和模型并行之外的更多维度，包括参数并行和属性并行。

3、支持什么

除了 PyTorch，FlexFlow Train 也支持 TensorFlow Keras 和 ONNX 模型导入。喜欢用 C++ 的开发者可以直接用 C++ API，功能跟 Python 版本对等。

运行时通过命令行参数控制训练配置，比如 epoch 数、batch size、GPU 数量、设备内存大小等。搜索过程也有独立的参数可以调，包括搜索迭代次数、参数并行开关、属性并行开关等。

命令行参数分两块：训练配置和 Legion 运行时配置。训练参数管 epoch 数、batch size、打印频率、数据集路径这些。运行时参数管 GPU 数量、每个 GPU 的显存大小、数据加载 worker 数这些。

4、适合谁用

• 在大规模集群上训练大模型，需要自动优化并行策略的团队
• 研究分布式系统并行化算法的学者
• 想把现有 PyTorch 模型迁移到分布式环境、又不想重写代码的开发者

项目采用 Apache 2.0 协议。FlexFlow 仓库已经拆分，训练相关的内容在这个仓库，推理和 serving 相关的功能在 flexflow-serve 仓库里。

仓库，推理和 serving 相关的功能在 flexflow-serve 仓库里。