Qwen3-235B 长序列强化学习训练性能优化实践
作者昇腾实战派知识地图链接强化学习知识地图一、背景概述随着大模型后训练范式从SFT向SFT-RL-SFT演进强化学习在大模型对齐与能力提升中扮演关键角色。基于昇腾NPU平台的Verl框架已成为主流训练工具之一尤其在长序列推理场景下对性能与显存效率提出更高要求。本文基于Atlas 800T A2服务器聚焦于Qwen3-235B规模模型在2k输入、30k输出长度下的强化学习训练性能优化针对推理阶段耗时过长、显存压力大、训练中断风险高等问题系统性地开展性能分析与调优。版本信息性能优化版本部分非版本号为master分支commit id组件版本HDK25.2.1CANN8.3.RC1VeRL21271aavllmrelease/v0.11.0vllm-ascend15b2e5cMindSpeed35da6acMegatron-LMcore_v0.12.1torch2.7.1toch-npu2.7.1-0919MindSpeed-RL 2.2.0商发配套版本链接组件版本HDK25.2.1CANN8.3.RC1VeRL796871d7d092f7cbc6a64e7f4a3796f7a2217f5evllm38217877aa70041c0115ee367b75197af9cbc5advllm-ascend1de16ead8eecfec8903ec1b330b27a4fa2593c35MindSpeed1cdd0abMegatron-LMcore_v0.12.1torch2.7.1toch-npu2.7.1三、性能瓶颈分析在初始配置下经profiling拆解发现generate阶段耗时较大成为核心瓶颈。主要问题集中在以下三方面输入配置不合理batch_size × n_samples 384 × 16 6144远超卡数整除能力导致数据分发效率低下推理性能未充分优化部分关键特性未启用如aclgraph图模式、融合算子等训练阶段配置缺失update阶段占比过高缺乏分布式优化器、ETP、TND子序列Batch Rope等关键优化手段。三、性能优化方案3.1 通用性能优化3.1.1 二级流水Task Queue Level 2通过设置环境变量export TASK_QUEUE_ENABLE2启用二级流水算子下发机制。该优化将算子任务拆分为一、二级流水并行执行尤其将workspace相关任务迁移至二级流水显著掩盖Host调度延迟提升整体端到端性能。该配置仅在二进制场景生效建议优先使用Level 2。3.1.2 高性能内存管理jemalloc在ARM架构环境下启用jemalloc可有效降低内存碎片提升内存分配效率。安装方式如下aptinstalllibjemalloc2exportLD_PRELOAD$LD_PRELOAD:/usr/lib/aarch64-linux-gnu/libjemalloc.so.2注x86平台无明显收益仅在ARM环境推荐使用。3.2 推理性能优化3.2.1 ACLGraph图模式PIECEWISEVLLM后端支持eager与图模式切换。通过将enforce_eagerFalse启用ACLGraph可将多个小算子融合为图执行减少Host下发开销。实测端到端性能提升达15%-20%。关键修复早期版本因MindSpeed patch导致图捕获异常已通过上下文管理器contextmanager机制修复确保在rollout阶段动态切换torch.compile与dummy_compile避免通信异常。3.2.2 大专家并行EP配置优化开启enable_expert_parallelTrue可提升专家层并行度但需注意依赖CANN 8.3.RC1及以上版本A2平台通信带宽有限EP过大如128反而导致通信开销激增性能劣化建议根据实际通信能力合理设置EP大小避免“过切”。3.3 训练性能优化3.3.1 融合算子优化通过启用多个融合算子减少算子下发次数提升计算效率算子类型配置项说明RMSNormuse_fused_rmsnormTrue替换原始分步归一化SwiGLUuse_fused_swigluTrue融合GELU与线性变换RoPEapply_rope_fusionTrueuse_fused_rotary_pos_embTrue将7个算子融合为1个耗时从86μs降至24μs端到端收益约0.5%GMMmoe_grouped_gemmTrue融合多专家计算提升GMM算子效率3.3.2 分布式优化器启用分布式优化器可显著降低显存占用与通信压力actor_rollout_ref.actor.megatron.use_distributed_optimizerTrue actor_rollout_ref.ref.megatron.use_distributed_optimizerTrue actor_rollout_ref.actor.megatron.override_transformer_config.use_distributed_optimizerTrue3.3.3 ETPExpert Tensor Parallel优化设置expert_tensor_parallel_size1避免专家参数在行列维度切分提升小专家场景下的计算效率尤其适用于MoE模型。3.3.4 TND场景Batch Rope优化在TNDToken-Nested-Data格式下原始方案对每条子序列独立执行RoPE存在循环开销。新方案通过生成统一频率矩阵实现一次RoPE完成全部编码RoPE阶段耗时减少50%端到端性能提升1%-2%且序列越长收益越显著。3.3.5 训练并行策略调优原始配置中CPContext Parallel过大PP通信占比高达56.6%成为瓶颈。通过以下调整优化优化项原配置新配置效果CP84降低上下文通信PP128显著减少PP通信EP816缓解显存压力ETP未开启1提升专家计算效率优化后PP通信占比从56.6%降至29.55%整体通信效率显著提升。四、关键问题与解决方案问题原因解决方案长序列推理超时HCCL通信超时将输入与输出batch的interleave均设为False避免长尾问题精度异常reward跌0仅修改输入batch interleave未同步输出batch保持输入与输出batch排布一致开启CP报错缺少override_transformer_config.context_parallel_size配置补全配置项group_type报错torch-npu与mindspeed版本不匹配升级至兼容版本找不到index_first_axisrl-plugin未正确patch手动修改verl/__init__.py文件训练拉起慢校验函数阻塞临时返回空值加速启动DataProto报错batch_size × n_samples 不能被卡数整除调整为可整除组合如384×16tensordict精度问题.to(cpu)为非阻塞操作升级tensordict至0.10.0及以上版本开启PP报错megatron.training模块缺失升级mindspeed至9.5后版本或修复optimize_p2p_comm.pylog_prob阶段OOM30k序列下显存压力过大降低log_prob_max_token_len_per_gpu至16k平衡显存与性能建议log_prob_max_token_len不宜过小否则输入shape过细计算效率下降。五、新模型适配与复用目标模型与Qwen3-MoE结构高度一致仅在Attention部分省略两个Norm层其余结构、权重格式与配置基本兼容。适配过程参考已有Qwen3-MoE方案所有优化配置可直接迁移。适配后性能与原模型基本一致验证了优化方案的通用性与可复用性。六、总结与展望本项目通过系统性分析与多维度调优核心优化路径包括推理侧启用ACLGraph图模式、融合算子、Batch Rope训练侧启用分布式优化器、ETP、合理调整并行策略配置侧优化batch排布、显存控制、版本兼容性。未来可进一步探索训推异步方案如VeRL Fully Async、更高效的图模式如Torch-AIR以及动态切分策略持续提升大规模模型训练的效率与稳定性。