仅此一文让你掌握OneFlow框架的系统设计(上篇)
OneFlow开源近半年,近期发布了v0.3.2版本,相较于上个大版本,我们又新增了众多算子和功能(如亚线性内存优化、Partial FC、足够灵活易用的新版Checkpoint…),同时完备性(如Serving)、易用性(全新的API)也在快速推进中,敬请期待。
两个月前我们通过DLPerf项目证明了OneFlow是世界上最快的深度学习框架:
成诚:OneFlow是如何做到世界最快深度学习框架的zhuanlan.zhihu.com
其中PK了经NVIDIA深度优化后的各个主流框架实现以及官方实现,对比了ResNet50和BERT这两个CV和NLP领域应用最广的模型在数据并行下的吞吐率和加速比。这两个月我们又做了大量的实验,对比了其他各个应用场景下的第三方框架,如:
- 超大规模人脸识别案例 vs InsightFace
- Wide&Deep vs HugeCTR
- GPT-2 vs Megatron-LM
- SSP vs PipeDream
- Optimizer-Placement-Optimization vs DeepSpeed ZeRO
- Auto Parallelism vs FlexFlow
实验证明了OneFlow这一套简洁的抽象(SBP + Actor)在支持各种模型并行、混合并行、流水并行、自动并行、ZeRO等方面是如此的简单高效。上述的每个特性,都有一个专门的第三方的框架项目对主流框架进行魔改,至少涉及了数千行的项目代码,而这些特性在OneFlow中要么是原生支持的特性,要么是新增一个Actor类型或者是一个图优化的Pass(几十行的代码)就可以支持的非常好。我们认为OneFlow这套设计不仅是性能最快的框架设计,同时也是分布式深度学习训练框架中最简单、最易扩展的框架设计。上述的每个实验,近期都会有专门的技术文章分享,感兴趣的小伙伴可以期待一下~
如果你对OneFlow这套致简致快的框架设计感兴趣,或者对深度学习框架、分布式系统感兴趣的话,本文就会让你完全掌握OneFlow的系统设计。相信读完这篇文章,你就会理解我们是如何看待分布式深度学习训练的,我们为什么要这样设计,这样设计的好处是什么,以及我们为什么相信OneFlow这套设计是分布式深度学习训练框架的最优设计。本文的主要内容如下:
-
深度学习框架原理
-
OneFlow系统架构设计(简略版)
-
OneFlow完整运行流程 与 各模块的交互方式
-
- \1. 分布式集群环境初始化
- \2. Python端搭建计算图
- \3. 编译期: OneFlow(JobSet) -> MergedPlan
- \4. 编译期: Compiler(Job)->Plan
- \5. 运行时: Runtime(Plan)
全文分上、中、下三篇。本文是上篇。
深度学习框架原理
深度学习框架是人工智能领域的“操作系统”,为深度学习相关的算法工程师提供一套简洁易用的用户接口,使之能方便的搭建深度学习模型,进行深度学习模型的训练、验证、测试、调参、迁移、部署、迭代开发等工作。同时深度学习框架作为底层硬件跟算法工程师之间的中间件,要做到设备无关,使得算法工程师可以不用关心具体的计算设备、存储设备的细节就能方便的开发模型。
深度学习框架本质上是一个基于张量(Tensor)之间的计算(Operator)表达式所组成的计算图(Graph)编译执行引擎,提供了一系列张量的定义、一元操作、二元操作等数学原语,并根据反向传播算法(Back Propagation)进行梯度自动求导以及模型更新。在大量数据分批次流入计算图进行模型训练之后,使得模型学习到数据中的内在关联关系,从而获得对应场景中的“智能”感知与判断能力。
OneFlow系统架构设计
OneFlow总体分为3个层次: Python前端、编译期(Compiler)、运行时(Runtime)。
- Python端是用户接口,是OneFlow启动、编译、运行的入口,负责构建逻辑图(Job),且负责运行时跟底层计算图执行引擎交互,包括发送控制指令(运行一个global_function / job)、喂数据(input)、处理输出(output,callback)。
- 编译期(Compiler)负责将前端用户的定义的逻辑上的计算图进行编译,产出实际上的物理计算图 (Plan)
- 运行时(Runtime)负责根据Plan创建真正的执行图——即一个由Actor组成的去中心化流式计算图,每个Actor各司其职,有的Actor负责接收Python端的控制信号,有的Actor负责加载数据,有的Actor负责初始化模型、计算、更新、存储、传输…,有的Actor负责返还给Python端数据,数据在计算图中流动,实现深度学习的模型训练功能。
总体架构图如下图所示: