Post-Training on PAI (3): 自研高性能强化学习框架 PAI-ChatLearn
引言
Post-Training（即模型后训练）作为大模型落地的重要一环，能显著优化模型性能，适配特定领域需求。相比于 Pre-Training（即模型预训练），Post-Training 阶段对计算资源和数据资源需求更小，更易迭代，因此备受推崇。
近期，我们将体系化地分享基于阿里云人工智能平台 PAI 平台在强化学习、模型蒸馏、数据预处理、SFT 等方向的技术实践，旨在清晰地展现 PAI 在 Post-Training 各个环节的产品能力和使用方法，欢迎大家随时交流探讨。
作为模型后训练的重点技术领域，强化学习任务具备极高的复杂性，需要数据、算法、计算等诸多方面的高效协作，从而避免模型漂移、泛化能力下降等问题。尤其在企业应用过程中，强化学习阶段还需面临开发易用性和计算性能的双重挑战。
易用性挑战：
对新算法出现时的可扩展性，比如 同步和异步（on policy vs off policy），PPO 和 GRPO（是否需要 critic model）等
新的模型类型是否可以快速支持，比如多模态大模型是否可以快速高效运行起来。
性能挑战：
强化学习的推理和训练过程需要顺序执行，导致的 GPU 资源闲置
分布式推理过程 worker 之间的生成 response 长度不均衡，导致 GPU 空泡，特别是超长序列生成的长尾问题
多个模型（policy，critic，reward 等）的放置，训练和推理的调度，显存管理，对训练效率的挑战
分布式 MOE 训练/推理本身的性能优化
针对上述问题，人工智能平台 PAI 推出了高性能一体化强化学习框架 PAI-ChatLearn，从框架层面解决强化学习在计算性能和易用性方面的挑战。本文将介绍 PAI-ChatLearn 框架特性以及如何在 PAI 平台上使用 PAI-ChatLearn 提交强化学习任务。
PAI-ChatLearn 技术特性
PAI-ChatLearn 是 PAI 自研的灵活、易用、高效的大规模强化学习框架。
灵活易用的框架
框架工程层面，PAI-ChatLearn 支持用户自定义模型、算法和数据流，包括：
使用 Ray 作为基础调度框架，基于高度可定制的模块化设计，通过提前构建计算图的方式，支持定制同步、异步调度，训练/推理 backend 选择，以及高效的模型参数同步。
多 Actor 的设计可以通过简单的配置功能处理 Ray 的资源分配和调度，比如 model colocate 等，支持对 workload 提供更加细粒度的控制
支持 Megatron-core、FSDP 等常用训练框架，支持 vLLM、SGlang 等推理框架，对不同的模型配置不同的并行策略，支持资源调度机制，提升 GPU 资源利用率。
支持 RLHF、DPO、PPO、GRPO 等丰富的强化学习算法，同时也支持用户自定义计算图，可以更方便的定制数据生产和训练流程。
极致的计算性能
针对强化学习场景，PAI-ChatLearn 通过实现 Dynamic Batchsize, Sequence Packing, Sequence Parallel，Partial Rollout 等负载均衡技术，Group GEMM、DeepEP 等 MoE 模型加速的技术，极大提升了 GPU 利用率。
Dynamic Batchsize + Sequence Packing
对于强化学习任务来说，训练样本都来自模型 rollout 生成的结果，因此会出现长短不一的情况。传统的实现方式在处理这种问题时，会将一个 batch 内的所有样本 padding 到相同长度再推理，这个过程中会引入许多无效 token 参与计算，极大降低整体的 tps。 为了更高效的处理 rollout 长度不一致问题，我们采用 Dynamic Batchsize + Sequence Packing 技术，对同一个 minibatch 内的样本进行重组，使得每条样本的长度相近，最大发挥算力能力，提升整体 tps
Sequence Parallel
对于超大模型的超长序列问题，若一张卡内无法放下一整条数据，则会导致 OOM。Sequence Parallel 可以将一条样本的计算均匀分散到多张卡，使得训练可以继续进行。 PAI-ChatLearn 充分借鉴 DeepSpeed 技术方案，通过 Ulyssess Sequence Parallel 来解决超长序列问题。
GroupGemm
对于 MoE 模型，传统的 Sequential MLP 实现在超大 expert 数的情况下，会面临效率挑战。FSDP 依赖的 Transformers 建模没有任何优化，为了让 MoE 模型在 FSDP 训练引擎下可以在合理的时间开销下进行，我们对 MoE 做了重写来支持效率更高的 GroupGemm 算子。
**Partial Rollout： **
为防止 vllm 前向过程中的长尾问题，ChatLearn 对超长序列的生成做截断操作，在下一次 vllm 前向时再继续生成。
实测效果
和其他开源框架相比， PAI-ChatLearn 在性能和规模上有明显优势。
以 Qwen3 模型为例，和开源框架 VeRL 相比， PAI-ChatLearn 端到端加速比在规模和性能都有明显的优势
以 LLaMA2 Dense 模型为例， PAI-ChatLearn 相对业界其他开源框架也有明显的提升
PAI-ChatLearn on PAI
PAI 平台的云原生 AI 训练模块 PAI-DLC（Deep Learning Containers）为企业和开发者提供灵活、稳定、易用和高性能的机器学习训练环境， 支持多种算法框架，能够处理大规模的分布式深度学习任务，在降低成本的同时提升训练效率。
PAI-DLC 支持一键提交 PAI-ChatLearn 框架的强化学习任务。以下我们将以 Qwen3 为例，展示应用 PAI-ChatLearn 的操作步骤。详细步骤可以参考：https://x.sm.cn/5BCZI0l ）
准备 Qwen3 模型
从 ModelScope 下载 Qwen3 模型权重
准备训练数据集
本案例以 MATH-lighteval 数据集为例，介绍 PAI-ChatLearn 强化学习流程。
该数据集是一个数学推理任务数据集，使用固定规则来验证 reward 打分。
如果需要在自定义任务上进行强化学习训练，可参考 PAI-ChatLearn 代码库中
examples/fsdp/models/rule_reward.py实现自定义 reward 打分函数。
提交训练任务
在单机开发调试完成后，您可以在 DLC 环境中配置多机多卡的分布式任务，加快模型的训练速度。具体操作步骤如下：
进入新建任务页面：
登录PAI控制台， 在右侧选择分布式训练( DLC），单击进入。
在分布式训练（DLC）列表页，单击新建任务。
新建任务页面，配置参数：
罗列关键参数，其他参数请参见创建训练任务。
镜像地址：在镜像地址并在文本框中输入：dsw-registry-vpc.cn-wulanchabu.cr.aliyuncs.com/pai-training-algorithm/chatlearn:torch2.5.1-vllm0.6.6-ubuntu22.04-cuda12.6-py310（注：需要根据当前 region 信息来更改镜像地址）
启动命令：配置以下命令，其中 train_grpo_qwen3.sh 脚本输入的启动参数与 DSW 单机预训练模型一致。
如果需要在自定义任务上进行强化学习训练，可参考 ChatLearn 代码库中
examples/fsdp/models/rule_reward.py实现自定义 reward 打分函数。
点击“确定”提交 DLC 任务
还未添加个人签名 2020-10-15 加入
分享阿里云计算平台的大数据和AI方向的技术创新和趋势、实战案例、经验总结。
评论