胡晓光-飞桨框架技术创新之路

飞桨框架技术创新之路 百度杰出架构师/胡晓光 目录 飞桨框架总体设计 动静统一自动并行 飞桨框架关键技术 应用层NewBing Midjourney 百度新搜索…… 模型层 Gemini ChatGPT…… 框架层…… 芯片层 TPU…… 易用的开发接口 •数据抽象：Tensor •操作抽象：Operator •Python，Numpy •避免了解硬件底层开发接口，如CudaC,cuDNN,cuBlas等 自动微分功能 •只需要写神经网络的前向，通过链式法则自动计算Loss关于参数的梯度 •高阶自动微分 •前向自动微分 多硬件 •隔离硬件差异，如CudaC,oneAPI,XTDK,BangC等等 •充分发挥硬件性能 分布式，合理排布计算、访存、通信等操作 可扩展，如硬件后端、算子、Pass、Dialect等 •让深度学习技术的创新与应用更简单 模型 代码规模 Llama3(推理) 2千行 StableDiffusion 1.8万行 AlphaFold 2.4万行 框架 代码规模 Tensorflow(去三方库) 420万行 PyTorch 298万行 PaddlePaddle 307万行 https://github.com/meta-llama/llama3https://github.com/Stability-AI/stablediffusionhttps://github.com/google-deepmind/alphafold统计于2024年6月6日 https://github.com/tensorflow/tensorflowhttps://github.com/pytorch/pytorchhttps://github.com/PaddlePaddle/Paddle PaddlePaddle开源 LEarning) (PArallelDistributedDeep 飞桨2.0发布 飞桨3.0Beta待发布 静态图模式分布式训练 动静统一 千亿模型混合并行 +Compile +shard_tensor Theano开源于2008 Caffe开源 TensorFlow开源 静态图模式 PyTorchv0.1发布 动态图模式 TensorFlowv2.0 +Eager模式 PyTorchv1.3 飞桨CINN探索 Compiler PyTorchv2.0 +Compile 符号微分 Layer Operator 更小的复用粒度 +Graph模式 融合两种开发模式灵活和效率的平衡 大规模分布式极致性能优化 2006 2012 2016 2019 2023 DeepBeliefNets 图像、语音应用 自然语言处理Transformer 预训练模型科学计算、生物计算等 大语言模型、多模态模型 学术创新 科技企业参与 大规模产业应用 大模型时代 SelectiveStateSpaces RWKV使用RNN网络达到跟Transformer可比的效果Mamba:Linear-TimeSequenceModelingwith SwitchTransformers:ScalingtoTrillionParameterModelswithSimpleandEfficientSparsity https://openai.com/research/video-generation-models-as-world-simulators 2021年DeepMind公司提出AlphaFold2模型突破生物学家50年未解蛋白质结构预测难题 2022年英伟达&西门子基于AI和物理的风电场数字孪生以比传统方法快4,000倍 2023年DeepMind公司提出GraphCast气象模型90%气象要素预测精度领先权威气象台 2024年普林斯顿团队用AI提前300毫秒预测核聚变等离子不稳定态 硬件 飞桨框架 产业落地 AI技术发展 开发者 开发体验：动静统一、训推一体、自动并行 MoE、多模态AIforScience 训、压、推全流程 极致性能：广泛硬件适配、软硬一体协同优化 模型层 … 语音 计算机视觉 科学计算 多模态模型 大语言模型 • • 自动微分 反向微分 高阶微分 • • 分布式 并行策略 集合通信 • • 自动并行 张量切分 策略寻优 • • 组合算子 前向规则 反向规则 • • 动转静 AST代码转写 字节码模拟 设备抽象 PHI算子库 神经网络编译器CINN 适配层 参数自动调优 代码自动生成 基础算子 算子内核 计算流管理 存储管理 算子自动融合 动态形状推导 算子定义 Tensor定义 分布式通信 设备运行时 执行层 内存/显存管理优化 静态图新执行器 动态图执行引擎 表示层 PaddleIR Pass Dialect Operation SSAGraph语义 Program语义 接口层 … 分布式 自动微分 优化器 组网 数学计算 张量表示 芯片层 … 寒武纪MLU 昇腾NPU 海光DCU 昆仑XPU 英伟达GPU 大模型自动并行 IR等基础架构统一 高性能编译计算 多硬件 算子库 动态执行引擎 CINN编译器 推理部署 执行器 Pas 组合算子 PIR 自动微分 静态图 （含自动并行） 动转静 动态图 (含自动并行) 动静统一、训推一体广泛硬件适配 自动并行高阶自动微分 动转静拆解对接 动态图静态图框架基础算子编译器 x - Mean(x) x return ReLU BatchNorm Conv2d x Conv w bias deffoo(x): y=conv2d(x,…)z= batch_norm(y,…)returnrelu(z) + return Max(x,0) Std(x) 动态图动转静 新IR自动微分 自动并行组合算子编译器执行器 像写单机代码一样写大规模分布式程序，通过自动并行、自动微分、编译优化等技术，实现极致性能和广泛硬件适配 目录 飞桨框架总体设计动静统一自动并行飞桨框架关键技术 切换到静态图模式 动静统一的组网 创建计算图 （编译期） 每次执行都动态创建计算图 每次执行相同的计算图（执行期） 静态图动态图 to_static 一行命令转换成静态图 飞桨动转静Python语法支持率94%，领先PyTorch的59%飞桨模型整图导出成功率95%，领先PyTorch的62% 通过比较官网文档说明的Python语法34项，以上展示部分语法 自适应图构建机制 opop op 升级 Subgraph opop op •解决深度学习模型下Python语言与框架IR之间差异性，实现转静训练100%的成功率，极致化用户体验 •拓宽接入AI编译器的路径，解锁更多转静加速能力 op动转静 终止 op Subgraph op op 自适应图构建 静态图中间表示运行时 先组网后执行、依赖执行器、全局视野 张量+算子=计算图，即IR中间表示 执行器 硬件设备 推理部署 模型训练仅是深度学习任务链条的一环，更重要地是结合业务落地，部署推理，开放给用户使用 性能优化 IR中包含了计算图的全局信息，如上下游算子的邻接关系等，更有利于进行图优化 推理计算图 常量折叠 算子融合 张量inplace 三段式Pass开发范式，让开发者更聚焦于Pass逻辑的处理，无需关心底层IR的细节 SourcePattern Pass样例：移除冗余连续的Cast计算 ResultPattern DRR应用广泛，在算子Fuse融合类Pass开发上，简洁易用，有效降低60%的开发成本 FlashAtten算子融合Pass Conv2d+Add+Act算子融合Pass 高扩展性IR 支撑大模型性能极致优化 70.00% 60.00% 统计超过900个模型配置PIR升级推理速度提升分布 58.98% 功能完善的Pass体系 PatternRewriter DRR 25.22% 10.23% 4.81% 1.76% 50.00% 多层级的Dialect ShapeDialect ControlFlowDialect CINNDialect KernelDialect OperatorDialect DistributedDialect 40.00% 30.00% 20.00% 10.00% 灵活的基础组件 Interface Trait Type Attribute Value Operation 0.00% 提升30%+ 提升10% 下降10%下降30%持平 通过使用PIR，Pass开发成本降低60%，推理性能提升超过10% •提前静态选Kernel策略，调度更加轻量高效 •常量折叠Pass适用范围更广、执行速度更快 •InplacePass适用范围更广、支持更多算子 计算量、参数量增长大幅超越AI芯片算力、显存增长速度多节点分布式/并行计算带宽瓶颈 计算量vs摩尔定律参数量vs显存算力vs带宽 以NVIDIA近年代表性GPU发展为例 GPU型号 发布时间 TensorPerformance 显存大小 机内互联带宽 A100 2020 312TFLOPS(FP16/BF16) 80GB 600GB/s H100 2022 1000TFLOPS(FP16/BF16) 80/94GB 900GB/s B200 2024 2250TFLOPS(FP16/BF16) 192GB 1800GB/s 图片引自：https://arxiv.org/pdf/2005.14165.pdfhttps://medium.com/riselab/ai-and-memory-wall-2cb4265cb0b8 单设备训练太慢 多设备摊分数据并行训练 多设备优化 “训得快” 面临问题解决思路技术方案 FP32显存占用过高BatchSize=1都跑不起来进一步提高BatchSize 降低单卡显存占用 “可训练” 显存占用过高，导致单卡无法放下整个模型 多设备摊分显存 使用更低精度训练 提前释放Activation节省显存 拆分mini-batch，多次计算前反向，一次参数更新 模型状态切分 算子内部切分 层间切分 参数切片并行 张量并行 流水线并行 数据并行 单设备优化 混合精度训练重计算梯度累加 如何实现高效的混合并行策略以应对大模型的发展？ •不同模型、不同规模、不同设备可能需要使用不同的策略 解决方案 •手动并行：手动管理张量切分和通信，不同并行策略使用不同的API •自动并行：框架自动推导张量切分和通信，用统一API表示不同并行策略 •半自动并行：需要用户添加少量分布式标记信息，框架推导剩余算子和张量的分布式信息 •全自动并行：不需要用户添加分布式标记信息，框架自动推导所有算子和张量的分布式信息 手动并行到自动并行的架构演进 动态图手动并行 动静统一半自动并行 动态图半自动并行 动转静 静态图半自动并行 全自动并行 • 对于百亿甚至千亿规模的大模型，一般需要使用TP+PP+DP/Sharding的混合并行方式进行训练 初始化分布式环境 修改TP组网 修改PP组网 构造分布式model和optimizer运行模型训练 如何用一种方式统一表达各种并行策略？ 如何表示张量在多卡设备的分布状态以及不同的并行策略？ •ProcessMesh每个设备映射一个进程，用进程mesh来表示设备资源 •Placement张量在不同设备的分布（切分）状态 •Replicate:复制的 •Shard(n):dim-n切分的 •Partial:不完整的，需要规约 •sum,mean… Data0 GPU0 GPU1 Data0 GPU4 GPU5 Data1 GPU2 GPU3 Data1 GPU6 GPU7 mesh0/stage0 m