大模型训推一体、多硬件适配,LLM时代的飞桨3.0新特性全面剖析
深度学习框架,可以说是整个AI大厦的地基。地基打好了,上面的建筑才能盖得高、盖得稳。它把硬件底层那些复杂操作封装成通用接口,让开发者能专心搞算法,不用去管显存分配、算子调度这些头疼的事。自动微分、动态图这些功能的加入,更是把开发效率拉到了一个全新的高度。
飞桨,作为国内首个自主研发的开源深度学习平台,从1.0版本静态图起步,到2.0版本默认动态图并实现动静统一、训推一体,再到如今为大模型时代量身打造的
3.0版本正式发布
01 设计思想
深度学习框架的设计,终极目标就是降低技术创新和应用的门槛。怎么做到?得从两个维度发力。
首先是用户端。一个好的框架必须给开发者极致的体验——不是说“好用”就行,而是要大幅度削减学习成本和时间成本。飞桨提出的“动静统一、训推一体、自动并行”,就是冲着这个目标去的。其次是硬件端。现代模型跑在各种各样的芯片上,框架必须能屏蔽硬件差异,做到广泛适配,还要能跟硬件协同优化,榨干每一分性能。
另一方面,框架的演进必须紧跟技术趋势和产业需求。大语言模型、MOE、多模态、科学计算……这些前沿方向对分布式训练、存储、通信都提出了更高要求。框架不仅要支持这些新模型,还得具备训练、压缩、推理一体化的全流程能力。只有跟得上趋势、经得住打磨的框架,才能为产学研各界开发者提供持续稳定的支持。
飞桨框架3.0的设计理念和主要特色
基于这些思考,飞桨框架3.0在2.x版本的基础上,
重点推出了动静统一自动并行、编译器自动优化、大模型训推一体、大模型多硬件适配四大新特性
02 框架架构一览
要实现上面这些特性,架构设计必须下硬功夫。下面这张架构图,展示了飞桨框架3.0的核心模块和它们之间的协作关系。
飞桨框架 3.0 架构图
从最上层往下看,飞桨对外提供了丰富的开发接口——张量运算、数学计算、模型组网、优化策略等等。再往下,框架分为4个层次:
- :负责计算图的表达与转换。这里重点引入了高可扩展中间表示PIR,为动转静、自动微分、自动并行、算子组合等核心功能提供了统一的底座。
表示层
- :对代码或计算图进行智能编排和调度,同时做显存、内存优化。无论动态图还是静态图,都能高效执行。
调度层
- :由神经网络编译器CINN和算子库PHI两部分构成,涵盖张量定义、算子内核、自动融合等关键能力。
算子层
- :对接底层硬件,包括设备管理、算子适配、通信适配、编译接入等。
适配层
这次3.0版本的架构升级,主要集中在三大块:
- :打破框架各模块之间的壁垒,提升在科学计算、编译优化、大模型领域的潜力。
高扩展中间表示PIR
- :通过自动融合和策略调优,大幅提升模型端到端性能。
神经网络编译器自动优化
- :降低大模型场景下的开发和性能优化成本,提升用户体验。
动静统一自动并行
03 高扩展中间表示 PIR
计算图的中间表示(IR)是深度学习框架性能优化、推理部署、编译器等方向的基石。在大模型时代,对IR的灵活性、扩展性、完备性要求更高了。飞桨3.0在基础架构层面统一了IR定义,实现了全架构统一表示,让上下游各个方向共享开发成果。
PIR的核心理念是高度灵活和高扩展性。它通过Type、Attribute、Op、Trait、Interface这些基础组件,以及Dialect概念,让开发者能灵活扩展和定制语义表达。在模型表示层,通过多Dialect的模块化管理和统一多端表示,实现了训练与推理一体化的全架构统一表示,无缝对接编译器和多硬件。在图变换层,统一底层模块,简化概念,提供了低成本、易用、高性能的Pass开发机制,支持丰富且可插拔的优化策略。
PIR坚守静态单赋值(SSA)原则,模型等价于一个DAG。它用Operation表示节点,Value表示边。Operation可以包含Region、Block,嵌套构建任意复杂语法。Value连接Operation,描述Use-Define链。飞桨提供了PatternRewriter和Declarative Rewrite Rule(DRR)两种Pass开发机制,兼顾灵活性与易用性。采用三段式Pass开发,开发者只需关注逻辑,不用管底层IR细节。实际效果:
Pass开发成本降低58%;在推理场景,超过84%的模型推理加速超10%
04 神经网络编译器自动优化
为什么要把编译器技术引入深度学习框架?三个核心原因:
第一,硬件发展速度严重不均衡——算力跑得飞快,但访存性能、CPU性能、总线带宽都跟不上。许多算子(如norm、activation)都受访存瓶颈限制。基于编译器的自动融合技术,可以把多个小算子融合成一个大算子,减少访存量和算子数量,大幅提升性能。这已经是深度学习框架的标配能力。
第二,模型结构越来越多样性。手动优化不可能覆盖所有模型,编译器提供的通用优化路径是最现实的解法。
第三,市面上有大量不同架构的硬件,每个都需要针对性优化,人力投入巨大。编译器可以大幅降低这类优化成本。
拿Llama里常用的RMS Normalization来举例。用Python接口组合实现很简单:
class RMSNorm(paddle.nn.Layer):
def __init__(self):
super().__init__()
self.variance_epsilon = 1e-6
self.size = 768
self.weight = paddle.create_parameter(
shape=[self.size],
dtype=paddle.get_default_dtype(),
default_initializer=nn.initializer.Constant(1.0),
)
def forward(self, x):
variance = x.pow(2).mean(-1, keepdim=True)
x = paddle.rsqrt(variance + self.variance_epsilon) * x
return x * self.weight
这段代码开发简单,但性能差、显存占用高。如果手写FusedRMSNorm,性能好但开发成本高。借助编译器自动优化,在A100平台上测试:
编译后的算子比普通实现快了4倍,比手动融合还提升14%
飞桨编译器的整体架构如下图。在表示层,CINN前端利用PIR扩展能力,完成算子拆分、重计算、子图划分、维度推导等操作,生成多个可被编译器后端优化的子图。后端会把子图转换为AST形式的低层IR,在此基础上做循环融合,最终生成一个kernel。底层IR还会进行性能调优,找到最优配置,再生成具体代码。
在Llama和Stable Diffusion上的测试结果显示,
启用编译器优化后,相比未手动优化的基础版本,推理速度分别提升了36%和30%
05 动静统一自动并行
大模型训练为什么要做自动并行?因为现在主流的手动并行方式实在太痛苦了。开发者需要同时精通模型结构、并行策略和框架调度逻辑,才能手工处理切分、通信、显存优化、调度优化这些事。结果就是大模型创新迭代被这些工程细节拖慢了。
举一个简单例子:在手动并行中,因为并行策略会导致Tensor shape变化,那些跟shape相关的算子(比如reshape)都得根据切分策略手动调整参数。稍有不慎就出错。
飞桨的自动并行方案是:
开发者只需少量张量切分标注,框架自动推导所有张量和算子的分布式切分状态,添加正确的通信算子,再根据模型结构和集群信息,结合显存和调度优化,自动寻找最高效的混合并行策略
在具体设计上,我们引入了两个关键概念:ProcessMesh和Placements。ProcessMesh把GPU卡映射为进程,多卡组成一维或多维数组。Placements是由Replicate、Shard、Partial三种标记组成的列表,分别表示张量在不同设备上复制、按维度切分、或是不完整需要做reduce操作。
开发者通过paddle.distributed.shard_tensor()接口标注张量切分,框架自动推导。下图展示了一个数据并行、张量并行、流水线并行混合的例子。
对应代码示例:
import paddle
import paddle.distributed as dist
...
mesh0 = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
mesh1 = dist.ProcessMesh([[4, 5], [6, 7]], dim_names=['x', 'y'])
class MlpModel(paddle.nn.Layer):
def __init__(self):
super().__init__()
self.w0 = dist.shard_tensor(
self.create_parameter(shape=[1024, 4096]),
mesh0, [dist.Replicate(), dist.Shard(1)])
self.w1 = dist.shard_tensor(
self.create_parameter(shape=[4096, 1024]),
mesh1, [dist.Replicate(), dist.Shard(0)])
def forward(self, x):
dist.shard_tensor(x, mesh0, [dist.Shard(0), dist.Replicate()])
y = paddle.matmul(x, self.w0)
y = dist.reshard(y, mesh1, [dist.Shard(0), dist.Shard(2)])
z = paddle.matmul(y, self.w1)
return z
...
model = MlpModel()
opt = paddle.optimizer.AdamW(...)
dist_model, dist_loader = dist.to_static(model, opt, ...)
for step, data in enumerate(dist_loader()):
loss = dist_model(data)
采用自动并行后,以Llama为例,
分布式训练核心代码量减少了50%
借助全局分析优化,性能也优于手动的动态图并行
06 产业优势
飞桨框架3.0-Beta是面向大模型和异构多芯的专属设计。向下适配多种硬件,释放潜能;向上一体化支撑大模型训练、推理。四大核心能力全面提升了服务产业的能力:
- :用户只需少量切分标注,框架自动推导并行策略,混合并行开发成本大幅降低,开发者可以更专注于算法创新。
动静统一自动并行
- :与框架一体化设计,支持生成式模型、科学计算等的高效训练与可变形状推理,Llama2和Stable Diffusion推理性能提升超30%。
编译器自动优化
- :训练和推理能力互相复用,动转静无缝衔接。RLHF训练中生成计算复用推理优化加速2.1倍,推理量化场景复用自动并行策略效率提升3.8倍。
大模型训推一体
- :提供简洁高效的抽象接口和基础算子体系,降低适配成本;优化调度编排和存储共享;编译器提供自动融合调优方案。已吸引硬件厂商贡献3,456个PR,25,000多个commits。
大模型多硬件适配
这就是飞桨新一代框架3.0。目前3.0-Beta版本已面向开发者开放,所有接口与2.x完全兼容,欢迎广大开发者使用和反馈。
-
- 关于宇宙的好的网名有哪些
- 角色扮演 | 1
- 网名