简介

本文详细对比分析了torchtitan、megatron、xdit、vllm等9个主流大模型训练/推理/RL框架,包括其核心特点、优势与局限性。文章从分布式训练、并行策略、资源管理等角度深入剖析各框架设计理念,并结合实际应用场景提供选择建议,为大模型开发者提供全面的技术参考和实践指导。

我把主流训练/推理/RL 框架写成了一份真实简评,先叠个甲,欢迎交流,轻喷~🐶

一、torchtitan

torch native 的训练打包解决方案,其实一开始就是推 torch native 的 tp 解决方案,dtensor 三种分布式语义:shard、replicate、partial。

torch dtensor 也是借鉴了 oneflow 当时 sbp 的思想,oneflow 当时在设计上两个非常大的贡献就是 actor 和 sbp 机制。

torchtitan 作为一个 tutorial 教给你 torch native 的 tp 应该怎么用,现在又集成进了 fsdp、基于 dtensor 构建的各种并行 plan、checkpointing、dcp。

其实严格意义上来说 torchtitan 不太像一个训练的框架,更像是一个手册和字典,等需要用的时候过来查一下。

二、megatron

抓两条主线,mpu/parallel_state 分布式环境初始化 initialize_megatron,构建通信进程大组和并行进程小组;

模型并行,著名的就是 megatron 的三篇论文:mlp/attention tp、interleaved 1f1b、tp-sp。

在 mcore 的代码实现里 model_provider 提供一个简单的 cpu 模型,没有 fp16 或 ddp,但是已经被 megatron 改造为并行的版本。

get_language_model 会获取一个 TransformerLanguageModel,由 ParallelTransformer 组成,如果需要数据并行,再配置 ddp/fsdp。

现在 megatron 的代码里也支持了 fsdp,既集成了 pytorch fsdp,也有 megatron 实现的版本。

megatron 被诟病的比较多的主要有两个点:框架和模型代码耦合,hf 参数转换。

一个事情都有它的一体两面,megatron 的风格是直给,把细节展现给你,不像 torch tp 的 parallel_plan 和模型分离,对于算法这部分是不可见的。

torch native 的 tp 解决方案更适合新手,写一个 parallel 的配置就好,但当需要定制化就不方便。

megatron 直接把 rowparallellinear、columnparallellinear 给到你,切分逻辑和通信都会比较直观,但是对于新模型的集成需要自己实现。

参数的转换,有 mbridge。

三、xdit

xdit 是较早开出来做 dit 推理的框架,核心理念主打 hf+并行,对标的是视觉生成的 vllm,创新性的提出了 ulysses+ring attn 结合的 usp,但是 wrap 和 decorator 设计,使框架层变得比较复杂。

我们开出来的 dax,做视频推理的框架,并行的部分就是参考的 xdit。

目前这个时间点 llm 推理框架也开始做视觉的模态了,出发点又有些区别:

  • sglang-diffusin:纯粹的 llm 开始做视觉生成,引入了 fastvideo
  • vllm-omni:多模态 omni,支持多 modal 的 encoder 和 decoder,视觉的 decoder 还是 diffusion 范式,顺带把视觉推理给做了,这就是社区强大之后的降维打击了
  • sglang-dlm:ar 范式下的 block diffusion,block-wise 的 kvcache,支持 llada

四、vllm

背靠 uc,从 continuous batching 和 pagedattention 起家,接触 vllm 源码主要是业务中出现了问题,才去源码里找答案。

排查 llm 循环输出问题,把 llm 生成的推理流程看了一遍,对于 llm 的推理生成流程是,输入 prompt 通过 tokenizer 得到对应的 token_ids。

然后经过 embedding 得到 hidden_states,作为 decode layer 层的输入,经过 Wq、Wk、Wv 权重矩阵的 projection,分别得到 Q、K、V,用于计算 sdpa。

再经过 out projection 得到输出的 hidden_states,对于类 transformer 架构来说 decode layer 是多层的,前一层 decode layer 输出的 hidden_states 会被送入下一层 decode layer。

作为当前层 decode layer 的输入,直到经过所有的 decode layer 层之后得到最终输出的 hidden_states。

再经过 llm_head 计算得到最后一个 token 的概率,进行采样得到下一个输出的 token_id,最后解决方案其实很简单,加一个 max_length 就好了。

排查 ray 进程启动问题,看了一下 vllm 的 backend_executor,支持 uni、multi、ray、external 等。

executor 可以有多个 worker,通过 collective_rpc 在所有 worker 上执行对应的函数。

五、openrlhf

应该是国内比较早做 rl 并开出来的框架,使用 ray pg group 做资源管理,group 内部复用已有的推理/训练框架走 spmd,后续的 rl 框架基本上都延续了这一设定,是作为 rl 入门比较好的框架。

目前社区更新已经有点力不从心,当时的大佬也基本都去了各个地方,可以说这波人现在基本撑起了国内 rl 框架开发的半壁江山,功不可没。

六、trl

hf 生态,易用性较好,好上手,适合单步 debug rl 流程,代码实现上更像是在 hf trainer 的基础上插入了 rl 流程的 hook,在 prepare_inputs 里面插入了参数更新和 rollout。

七、verl

verl 的出发点是 rl 的计算流程可以被建模为数据流 dataflow,每个节点都是一个分布式的 spmd 程序。

在 inter-node 为了灵活性使用 ray 作为胶水层实现了中心化的 single-controller。

在 intra-node 的 model 层面为了通信的效率使用了 multi-controller,可以复用现有推理/训练系统的 spmd 并行编程模型。

各个组件之间的连接基于 dataproto 的数据协议抽象,更像是预训练的框架打了一个 rl 的 patch,从内部系统 clean 出来的代码,感觉更多的是面向 mlsyser,而不是 rler 的框架。

这个感觉有点类似 megatron,可以问一个算法同学喜不喜欢 megatron。

八、slime

代码比较干净,基于 sglang、megatron 作为推理/训练后端提供高性能训练支持,支持共置/分离式架构,将 rl 流程拆分为 rollout/databuffer/train 三个独立模块,有较高的自由度,rl 框架的冉冉新星。

九、ray

也是 uc 的项目,出发点是解决 rl 环境交互复杂编程问题,我们从 22、23 年开始使用 ray 作为离线大规模数据处理和模型服务化的框架,发展到现在,ray 已经进入了 pytorch 开源基金会,有点 ai 领域的 k8s 那么个意思了。

task 和 actor 抽象的足够干净,强依赖于 grpc,ray core 提供了进程级资源管理和远程 rpc 调用的能力,可以很好的作为其他框架之间的胶水层,现在的 rl 框架基本都是选择用 ray 做的资源管理和控制层。

pytorch 也推出了一个类似的东西:monarch,看社区发展了。

十、如何学习AI大模型?

我在一线互联网企业工作十余年里,指导过不少同行后辈,帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家,也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑,所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限,很多互联网行业朋友无法获得正确的资料得到学习提升,故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

这份完整版的大模型 AI 学习和面试资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】
在这里插入图片描述

在这里插入图片描述

第一阶段: 从大模型系统设计入手,讲解大模型的主要方法;

第二阶段: 在通过大模型提示词工程从Prompts角度入手更好发挥模型的作用;

第三阶段: 大模型平台应用开发借助阿里云PAI平台构建电商领域虚拟试衣系统;

第四阶段: 大模型知识库应用开发以LangChain框架为例,构建物流行业咨询智能问答系统;

第五阶段: 大模型微调开发借助以大健康、新零售、新媒体领域构建适合当前领域大模型;

第六阶段: 以SD多模态大模型为主,搭建了文生图小程序案例;

第七阶段: 以大模型平台应用与开发为主,通过星火大模型,文心大模型等成熟大模型构建大模型行业应用。

在这里插入图片描述

👉学会后的收获:👈

• 基于大模型全栈工程实现(前端、后端、产品经理、设计、数据分析等),通过这门课可获得不同能力;

• 能够利用大模型解决相关实际项目需求: 大数据时代,越来越多的企业和机构需要处理海量数据,利用大模型技术可以更好地处理这些数据,提高数据分析和决策的准确性。因此,掌握大模型应用开发技能,可以让程序员更好地应对实际项目需求;

• 基于大模型和企业数据AI应用开发,实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能, 学会Fine-tuning垂直训练大模型(数据准备、数据蒸馏、大模型部署)一站式掌握;

• 能够完成时下热门大模型垂直领域模型训练能力,提高程序员的编码能力: 大模型应用开发需要掌握机器学习算法、深度学习框架等技术,这些技术的掌握可以提高程序员的编码能力和分析能力,让程序员更加熟练地编写高质量的代码。

在这里插入图片描述

1.AI大模型学习路线图
2.100套AI大模型商业化落地方案
3.100集大模型视频教程
4.200本大模型PDF书籍
5.LLM面试题合集
6.AI产品经理资源合集

👉获取方式:
😝有需要的小伙伴,可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓

在这里插入图片描述

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐