引言
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 模型权重
plain
modelscope download --model Qwen/Qwen3-8B --local_dir Qwen3-8B
准备训练数据集
本案例以 MATH-lighteval 数据集为例,介绍 PAI-ChatLearn 强化学习流程。
- 该数据集是一个数学推理任务数据集,使用固定规则来验证 reward 打分。
- 如果需要在自定义任务上进行强化学习训练,可参考 PAI-ChatLearn 代码库中
examples/fsdp/models/rule_reward.py
实现自定义 reward 打分函数。
plain
# 下载数据集
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环境中配置多机多卡的分布式任务,加快模型的训练速度。具体操作步骤如下:
- 进入新建任务页面:
- 登录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单机预训练模型一致。
plain
cd /mnt/data/ChatLearn && bash examples/fsdp/scripts/train_grpo_qwen3.sh
- 如果需要在自定义任务上进行强化学习训练,可参考ChatLearn代码库中
examples/fsdp/models/rule_reward.py
实现自定义reward打分函数。
plain
# 下载数据集
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任务
