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 的操作步骤。详细步骤可以参考： ）

准备 Qwen3 模型

从 ModelScope下载 Qwen3 模型权重

modelscope download --model Qwen/Qwen3-8B --local_dir Qwen3-8B

准备训练数据集

本案例以 MATH-lighteval 数据集为例，介绍 PAI-ChatLearn 强化学习流程。

• 该数据集是一个数学推理任务数据集，使用固定规则来验证 reward 打分。
• 如果需要在自定义任务上进行强化学习训练，可参考 PAI-ChatLearn 代码库中examples/fsdp/models/rule_reward.py实现自定义 reward 打分函数。

# 下载数据集
mkdir -p dataset
modelscope download --dataset AI-ModelScope/MATH-lighteval --local_dir dataset/MATH-lighteval
# 数据集预处理
python examples/fsdp/data/data_preprocess/math_lighteval.py --input_dir dataset/MATH-lighteval --local_dir dataset/MATH-lighteval

提交训练任务

在单机开发调试完成后，您可以在DLC环境中配置多机多卡的分布式任务，加快模型的训练速度。具体操作步骤如下：

1. 进入新建任务页面：
   1. 登录PAI控制台， 在右侧选择分布式训练( DLC），单击进入。
   2. 在分布式训练（DLC）列表页，单击新建任务。
2. 新建任务页面，配置参数：

罗列关键参数，其他参数请参见创建训练任务。

• **镜像地址：**在镜像地址并在文本框中输入：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单机预训练模型一致。

cd /mnt/data/ChatLearn && bash examples/fsdp/scripts/train_grpo_qwen3.sh

• 如果需要在自定义任务上进行强化学习训练，可参考ChatLearn代码库中examples/fsdp/models/rule_reward.py实现自定义reward打分函数。

# 下载数据集
mkdir -p dataset
modelscope download --dataset AI-ModelScope/MATH-lighteval --local_dir dataset/MATH-lighteval
# 数据集预处理
python examples/fsdp/data/data_preprocess/math_lighteval.py --input_dir dataset/MATH-lighteval --local_dir dataset/MATH-lighteval

3. 点击"确定"提交DLC任务