FlexFlow Train

[!WARNING] FlexFlow 仓库已被拆分为独立的 flexflow-train 和 flexflow-serve 仓库。 您当前正在查看的是 flexflow-train。 如需了解与推理/服务相关的内容，请访问 flexflow-serve。

FlexFlow Train 是一个深度学习框架，通过自动搜索高效的并行化策略来加速分布式 DNN 训练。

贡献

如果您遇到任何错误或有任何建议，请通过 提交问题告知我们。

有关如何向 FlexFlow Train 贡献代码的说明，请参阅 CONTRIBUTING.md。

我们欢迎所有对 FlexFlow Train 的贡献，无论是修复 bug 还是添加新功能和扩展。

引用文献

• Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandeep Baines, Carlos Efrain Quintero Narvaez, Vinay Ramakrishnaiah, Nirmal Prajapati, Pat McCormick, Jamaludin Mohd-Yusof, Xi Luo, Dheevatsa Mudigere, Jongsoo Park, Misha Smelyanskiy, and Alex Aiken. Unity: 通过联合优化代数变换与并行化加速 DNN 训练。载于 2022 年 7 月举行的操作系统设计与实现研讨会 (OSDI) 论文集。

• Zhihao Jia, Matei Zaharia, 和 Alex Aiken. 超越深度神经网络的数据与模型并行化。载于 2019 年 4 月举行的第二届机器学习与系统会议 (MLSys) 论文集。

• Zhihao Jia, Sina Lin, Charles R. Qi, 和 Alex Aiken. 探索卷积神经网络并行化的隐藏维度。载于 2018 年 7 月举行的国际机器学习大会 (ICML) 论文集。

团队

FlexFlow Train 由卡内基梅隆大学、Facebook、洛斯阿拉莫斯国家实验室、麻省理工学院、斯坦福大学和加州大学圣迭戈分校的团队开发并维护（按字母顺序排列）。

许可证

FlexFlow Train 采用 Apache License 2.0 许可证。

FlexFlow Train 快速上手指南

FlexFlow Train 是一个深度学习框架，旨在通过自动搜索高效的并行化策略来加速分布式 DNN 训练。它支持 PyTorch、TensorFlow Keras 和 ONNX 模型，并提供了 C++ 接口。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04/22.04)
• GPU: NVIDIA GPU (需安装 CUDA) 或 AMD GPU (需安装 ROCm/HIP)
  • 注意: 预构建的 CUDA Docker 镜像目前完全兼容宿主机安装 CUDA 11.7 的环境。
• 依赖项:
  • Python 3.x
  • CMake, GCC/G++
  • CUDA Toolkit 或 ROCm SDK
  • (可选) Docker: 用于快速部署预构建环境

安装步骤

您可以选择使用预构建的 Docker 镜像（推荐快速体验）或从源代码编译。

方法一：使用 Docker (推荐)

FlexFlow 提供了包含所有依赖项的预构建 Docker 镜像。

1. 拉取镜像

根据您的硬件后端选择对应的镜像：

• NVIDIA CUDA 后端:

  docker pull ghcr.io/flexflow/flexflow-cuda:latest
  

• AMD HIP/ROCM 后端:

  docker pull ghcr.io/flexflow/flexflow-hip_rocm:latest
  

2. 运行容器

启动容器并挂载当前目录以便代码开发（以 CUDA 为例）：

docker run --gpus all -it --rm -v $(pwd):/workspace ghcr.io/flexflow/flexflow-cuda:latest

方法二：从源代码编译

如需自定义编译或最新特性，请参考官方 INSTALL.md 文档进行源码编译。基本流程如下：

git clone https://github.com/flexflow/flexflow-train.git
cd flexflow-train
# 遵循 INSTALL.md 中的具体 cmake 和 make 指令进行构建

(注：Conda 安装包即将推出，敬请期待)

基本使用

FlexFlow 的核心优势在于能够优化现有 PyTorch 模型的并行性能。以下是基于 PyTorch 的最简使用流程。

第一步：导出 PyTorch 模型

首先，将现有的 PyTorch 模型转换为 FlexFlow 格式 (.ff)。

import torch
import flexflow.torch.fx as fx

# 假设您有一个已定义的 PyTorch 模型
model = MyPyTorchModule()

# 导出为 FlexFlow 模型文件
fx.torch_to_flexflow(model, "mymodel.ff")

第二步：加载模型并自动调优

在 FlexFlow 程序中导入模型，配置优化器，并利用其自动调优功能在给定机器上寻找最佳并行策略。

from flexflow.pytorch.model import PyTorchModel
# 假设 ffmodel 已初始化，input_tensor 已定义

def top_level_task():
  # 加载之前保存的 PyTorch 模型
  torch_model = PyTorchModel("mymodel.ff")
  
  # 应用模型
  output_tensor = torch_model.apply(ffmodel, input_tensor)
  
  ## 模型编译
  ffmodel.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
  
  ## 模型训练 (此处以 CIFAR-10 为例)
  # (x_train, y_train) = cifar10.load_data() 
  # 实际使用时请替换为您的数据加载逻辑
  ffmodel.fit(x_train, y_train, epochs=30)

第三步：运行训练脚本

使用命令行参数启动训练。FlexFlow 支持多种运行时配置和自动调优参数。

基础训练命令：

python your_script.py -e 30 -b 64 -ll:gpu 4

常用参数说明：

• 训练控制:
  • -e, --epochs: 总训练轮数 (默认: 1)
  • -b, --batch-size: 全局批次大小 (默认: 64)
  • -d, --dataset: 数据集路径 (若不设置则使用合成数据)
• 资源分配 (Legion 运行时):
  • -ll:gpu: 每个节点使用的 GPU 数量 (默认: 0)
  • -ll:cpu: 数据加载工作进程数 (默认: 4)
• 性能自动调优:
  • --search-budget: MCMC 搜索的迭代次数 (设为 0 禁用搜索，直接使用默认策略)
  • --export-strategy: 导出最佳发现策略的路径
  • --import-strategy: 导入之前保存的策略路径
  • --enable-parameter-parallel: 允许探索参数并行 (默认仅考虑数据和模型并行)

更多高级示例请参考项目仓库中的 examples/python/pytorch 文件夹。