作者:昇腾实战派
知识地图https://blog.csdn.net/Lumos_Lovegood/article/details/161455142

一、背景概述

随着大语言模型后训练技术的发展,SFT-RL-SFT 已成为当前主流的训练流程范式。VeRL 作为一款高性能强化学习训练框架,凭借其优异的扩展性和实用性,已在多项大模型训练任务中广泛应用。

本文基于实际项目经验,重点分析了在NPU环境中使用VeRL对Qwen2.5-32B模型进行DAPO训练时遇到的性能瓶颈,并提出一套系统性的优化方案,最终显著提升训练效率,为同类场景下的模型调优提供参考。

二、DAPO算法概述

DAPO(Dynamic Advantage Policy Optimization)是在GRPO基础上引入多项技巧的强化学习算法,据论文报告显示,其性能优于原版GRPO。为便于后续优化方案的理解,本节简要介绍DAPO的核心机制。

算法流程

DAPO 的算法流程与 GRPO 基本相同:

请添加图片描述

  1. Rollout阶段:采集满足数量要求的训练样本;
  2. 动态采样筛选:过滤掉奖励为1或0的样本,提升训练有效性;
  3. 奖励计算:对剩余样本进行奖励打分;
  4. 旧策略概率计算:计算样本在旧策略下的对数概率;
  5. 优势估计:基于奖励计算组间优势(Adv);
  6. 策略更新:利用优势函数和旧策略概率计算目标函数,更新模型参数。

核心特性

对比 GRPO,DAPO 主要作出了以下修改:

  • 更高裁剪上限(Clip-Higher):将GRPO中的裁剪上下阈值分离,并提高上界以增强模型探索能力;
  • 动态采样机制:对采样样本设定筛选条件,过滤奖励过高或过低的样本,提升训练效率;
  • Token级损失计算:将序列级平均损失改为全局Token平均损失,增强长序列中Token对损失的贡献;
  • 长序列奖励塑造:重新设计长序列生成中的长度奖励机制;
  • 取消KL散度约束:无需依赖参考模型来约束策略更新,简化训练流程。

三、性能优化方案

本章主要分通用优化、训练阶段优化和推理阶段优化三个部分介绍:

请添加图片描述

以Atlas 800T A2训练服务器上的Qwen2.5-32B模型训练优化为例,其中各优化项性能收益如下,其中无收益的部分优化方法在其他场景下可能会取得一定性能提升,也一并列出供参考:
请添加图片描述

性能采集

VeRL Profiling 采集

由于VeRL使用ray进行资源调度,相比torchrun无法在主循环内直接打点采集,许多人会碰到采不下来的情况,因此本节专门针对VeRL给出profiling采集方法

原仓工具采集

VeRL原仓合入了昇腾设备采集 profiling 工具,可直接参考使用说明

打点采集

由于强化学习分推理阶段和训练阶段,如果按照通常在主循环进行打点,推理阶段会导致profiling文件过大无法解析,建议训练推理阶段分开采。

VeRL中训推流程分别单独包装在 generate_sequenceupdate_actor 两个函数中,位于 verl/worker/fsdp_worker.py
推理阶段会逐token进行采集,如果不减少输出 token 的数量,单个profiling文件可能达到几百个G,采集时推荐修改的关键参数调整如下:

#  减少输出token数量,实际输出token数应该是actor_rollout_ref.rollout.n * data.gen_batch_size * data.max_response_length
data.max_response_length=16

#  减小 batchsize
data.gen_batch_size = $((train_prompt_bsz * 1))

#  关闭训前评估
trainer.val_before_train=False

#  关闭动态采样,减小token数容易触发跳过采样
algorithm.filter_groups.enable=False

#  关闭过长惩罚,该参数的一个关联参数与max_response耦合, 不关会报错
reward_model.overlong_buffer.enable=False

打点位置:

def generate_sequences(self, prompts: DataProto):
    ······
        with simple_timer("generate_sequences", timing_generate):
            rollout_prof_flag = False
            import torch_npu
            if torch.distributed.get_rank() == 0 and rollout_prof_flag:
                prof_save_path = "/path/to/prof_result/NPU_qwen25_32B_without_stack_rollout"
                experimental_config = torch_npu.profiler._ExperimentalConfig(
                                export_type=torch_npu.profiler.ExportType.Text,
                                aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, 
                                profiler_level=torch_npu.profiler.ProfilerLevel.Level1, 
                                l2_cache=False, 
                                data_simplification=False)
                prof = torch_npu.profiler.profile(
                                activities=[
                                    torch_npu.profiler.ProfilerActivity.CPU,    #  采集框架侧数据开关
                                    torch_npu.profiler.ProfilerActivity.NPU],    #  采集NPU数据开关
                                schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1, skip_first=0),    #  设置不同step的行
                                on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(prof_save_path),    #  将采集到的性能数据导出为TensorBoard工具支持的格式
                                record_shapes=False,    #  算子的InputShapes和InputTypes,Bool类
                                profile_memory=False,    #  算子的内存占用情况,Bool类型
                                with_stack=False,    #  算子调用栈,Bool类型
                                experimental_config=experimental_config)
                print(f"{'<' * 20} prof start {'>' * 20}")
                prof.start()
            output = self.rollout.generate_sequences(prompts=prompts)
            if torch.distributed.get_rank() == 0 and rollout_prof_flag:
                prof.stop()
                rollout_prof_flag = False
                print(f"{'<' * 20} rollout prof stop {'>' * 20}")

        log_gpu_memory_usage("After rollout generation", logger=logger)

       ······

训练阶段则无需更改参数,如果采集数据过大可以使用离线解析或减小 data.train_batch_size

def update_actor(self, data: DataProto):
actor_prof_flag = False
import torch_npu
if torch.distributed.get_rank() == 0 and actor_prof_flag:
     prof_save_path = "/path/to/prof_result/NPU_qwen25_32B_without_stack_actor"
     experimental_config = torch_npu.profiler._ExperimentalConfig(
                    export_type=torch_npu.profiler.ExportType.Text,
                    aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, 
                    profiler_level=torch_npu.profiler.ProfilerLevel.Level1, 
                    l2_cache=False, 
                    data_simplification=False)
    prof = torch_npu.profiler.profile(
                    activities=[
                        torch_npu.profiler.ProfilerActivity.CPU,    #  采集框架侧数据开关
                        torch_npu.profiler.ProfilerActivity.NPU],    #  采集NPU数据开关
                    schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1, skip_first=0),    #  设置不同step的行
                    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(prof_save_path),    #  将采集到的性能数据导出为TensorBoard工具支持的格式
                    record_shapes=False,    #  算子的InputShapes和InputTypes,Bool类
                    profile_memory=False,    #  算子的内存占用情况,Bool类型
                    with_stack=False,    #  算子调用栈,Bool类型
                    experimental_config=experimental_config)
    print(f"{'<' * 20} prof start {'>' * 20}")
    prof.start()
#  Support all hardwares
data = data.to(get_device_id())

······

if torch.distributed.get_rank() == 0 and actor_prof_flag:
    prof.stop()
    actor_prof_flag = False
    print(f"{'<' * 20} actor prof stop {'>' * 20}")

return output

上述方法可以完成0号卡在每个step的 profiling 采集
注意:VeRL采用RAY进行资源调度,可以关注打屏日志中打印prof信息的机器ip来找到rank0的所在的机器,采集数据只会在该机器中落盘(也可以把路径设为共享存储路径)

单推理采集

在vllm的启动脚本里手动加入启动位置,代码位置vllm/vllm/entrypoints/llm.py,在**_run_engine**函数修改为如下代码:

from vllm import LLM, SamplingParams
import json
import os
······
#  Create an LLM.
llm = LLM(
    model="/home/wxx/Qwen2.5-32B",
    #  model="",
    tensor_parallel_size=2,
    distributed_executor_backend="ray",
    trust_remote_code=True,
    load_format = "dummy"
)

#  Generate texts from the prompts.
#  llm.start_profile()
outputs = llm.generate(prompts, sampling_params)
#  llm.stop_profile()
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

注:在VeRL框架中,该单推采集方案产生的profiling每单步保存单个profiling文件,无法连续查看,仅建议在定位vllm后端版本问题时使用

离线解析

NPU上 torch profiler 会在解析后会默认对采集信息进行解析,但如果文件过大或权限问题等因素有时会解析失败,这时可以使用离线解析代码:

from torch_npu.profiler.profiler import analyse

if __name__ == "__main__":
analyse(profiler_path="./result_data", max_process_number=想使用的cpu核数)

性能瓶颈分析

基于性能采集与算子分析,梳理出多项可优化方向:

  1. 模型 ROPE 逻辑当前调用torch_npu._npu_rotary_embedding ATB 算子,算子内部存在index_select逻辑,会重复从cos_sin_cache读取正余弦参数,具备优化空间。
  2. 现有allreducematmul算子,可替换为 mc2 融合算子进一步提效。
  3. 当前未启用图模式,建议组合开启taskqueue level 1与图模式开展调优。
  4. 推理过程因样本长度动态变化,出现任务收尾不均的现象:大部分卡完成推理后,仍有少量卡持续执行,对性能会有一定影响,显存与执行节奏受样本随机性影响较大。

通用性能优化手段

本节主要介绍 NPU 训练的一些通用优化能力,除了目标模型其他大多数场景也可以获得可观收益

融合算子替换

首先将常用的 RotaryMul & RotaryMulGrad、RmsNorm & RmsNormGrad、Swiglu 三个融合算子替换到对应模型的 models 文件目录中,融合算子的具体设计及实现方式见昇腾社区性能优化部分

在 VeRL 开源代码仓中有专门的 npu_patch 文件(verl/models/transformers/npu_patch.py),对transformers 的 Qwen2 Model

部分进行替换:

import torch
import torch_npu
from torch_npu import npu_rotary_mul as apply_rotary_emb
from transformers.models.qwen2 import modeling_qwen2
from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm as Qwen25RMSNorm
from transformers.models.qwen2.modeling_qwen2 import Qwen2MLP

def apply_rotary_pos_emb_flashatt_npu(
 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
 cos = cos.chunk(2, dim=-1)[0].contiguous()
 sin = sin.chunk(2, dim=-1)[0].contiguous()
 cos = cos.repeat(1, 2)
 sin = sin.repeat(1, 2)
 q_embed = apply_rotary_emb(
 q.float(), cos.unsqueeze(0).unsqueeze(2).float(), sin.unsqueeze(0).unsqueeze(2).float()
 ).type_as(q)
 k_embed = apply_rotary_emb(
 k.float(), cos.unsqueeze(0).unsqueeze(2).float(), sin.unsqueeze(0).unsqueeze(2).float()
 ).type_as(k)
 return q_embed, k_embed


#  This api can improve performance on ASCEND NPU
def rms_norm_forward(self, x):
 return torch_npu.npu_rms_norm(x, self.weight, epsilon=self.variance_epsilon)[0]

def fused_apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 cos = cos.unsqueeze(unsqueeze_dim)
 sin = sin.unsqueeze(unsqueeze_dim)
 q_embed = torch_npu.npu_rotary_mul(q.contiguous(), cos, sin).to(q.dtype)
 k_embed = torch_npu.npu_rotary_mul(k.contiguous(), cos, sin).to(k.dtype)
 return q_embed, k_embed 

def qwen2_mlp_forward(self, x):
 gate_up_output = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1)
 return self.down_proj(torch_npu.npu_swiglu(gate_up_output, dim=-1))

Qwen25RMSNorm.forward = rms_norm_forward
Qwen2MLP.forward = qwen2_mlp_forward
modeling_qwen2.apply_rotary_pos_emb = fused_apply_rotary_pos_emb

替换融合算子后,端到端性能可提升8%+

rope 计算 sin/cos 前置

当前rotary embedding 走的是 atb算子torch_npu._npu_rotary_embedding,该算子包含index_select操作,每一次执行都会从con_sin_cache中重复获取cos, sin
请添加图片描述

profiling体现如下:
请添加图片描述

由于同一位置的token,在不同层中的位置编码相同,因此可以把cos, sin的获取提前至模型脚本中的layer loop之外,直接向rope_forward_oot中传入cos, sin,并走相对耗时更短的aclnn算子torch_npu.npu_apply_rotary_pos_emb,消除index_select操作
请添加图片描述

完成修改后,profiling对应部分如下:

请添加图片描述

训练部分性能优化

强化学习训练中单训练过程的时间占比较小,比竞品略高的算力在计算密集场景下只要切分得当往往不会存在太大的GAP,本节主要就切分给出建议,并提供一些VeRL支持的训练优化方法。

切分优化

fsdp size:
通过参数 actor_rollout_ref.actor.fsdp_config.fsdp_size 在fsdp后端下可以调整 fsdp 切分的范围,由于 32B 模型无法在单机(8卡)范围内进行 fsdp 切分来减少跨机通信,共进行了fsdpsize=32、64、128三个case,在平均输出长度约4K时,测试数据如下:

fsdp_size 32 64 128
throughput 40 42 47

实测 fsdp_size=128 时性能最优。

SP(ulysses):
长序列场景下通常需要开启sp来优化训练性能,如果不开sp会造成OOM。在训练早期(response length mean ~ 500),相比sp8,sp4会得到更好的性能(约20%的性能收益),但随着response length越来越长,这样的切分可能会导致OOM训练中断和性能劣化,因此可以采取在response length mean < 2500 的情况下采取sp4,在达到该长度后合适的断点位置中止任务切换为sp8,从而缩短整体的训练时间。

权重预取

可以通过参数 actor_rollout_ref.actor.fsdp_config.forward_prefetch=True 开启前向的参数预取,在当前前向计算完成前,提前发起下一个前向传播所需的 all-gather 操作,通过通算掩盖提升训练效率,详细实现见FSDP前向预取

反向预取由于梯度计算时序可能动态变化,导致 BACKWARD_POST无法准确预测下一次预取的时机,引发数据依赖错误,详情见 FSDP documentation;目前VeRL最新代码已禁用反向预取接口

Chunk logits

模型处理长序列张量时(如大语言模型的token序列),完整张量会一次性占用大量内存,VeRL中在训练阶段提供了参数actor_rollout_ref.actor.entropy_from_logits_with_chunking=True,在前向传播过程中,将张量按 [chunk_size, voc] 形状分块处理,而非直接处理完整序列长度的张量,从而降低显存峰值。

推理部分性能优化

V1与Chunked Prefill

简述:
V1是 vllm 对推理架构的一次整体更新,主要解决 V0 在高并发、长序列场景和资源效率上的局限性,在长序列场景下有大幅性能优化;Chunked Prefill 为 V1 对长序列场景的一种优化手段,将长短不一的 prompts 拆分为长短一致的 chunks 进行 prefill,这些 chunks 间的气泡再插入其他完成了 prefill 的 prompts 的 decode 需求,一方面降低了峰值显存,避免长序列场景下过长的prompt造成oom,另一方面大幅提升了并行计算的效率。

Qwen2.5-32B 在128卡、2K -> 20K场景下,开启 V1 和 Chunked Prefill 后端到端实测性能提升15%+

使能方式:

开启 V1 ,需要设置环境变量 VLLM_USE_V1=1,在 VeRL 中则是在runtime_env.yaml中添加对应参数值

VeRL 中开启 chunked_prefill,需要在启动脚本中设置参数:actor_rollout_ref.rollout.enable_chunked_prefill=True

V1 的支持度如下,根据实际需求决定是否开启:
请添加图片描述

vLLM V1 的详细实现与功能:vLLM V1: A Major Upgrade to vLLM’s Core Architecture

Chunked Prefill 的详细实现与技术细节:图解大模型计算加速系列:分离式推理架构2,模糊分离与合并边界的chunked-prefills

mc2

简述:

在vLLM的RowParallelLinear前向函数中,原本会分别执行allreduce和matmul操作。该功能通过使用torch_npu.npu_mm_all_reduce_base实现了allreduce与matmul的融合内核操作实现性能优化。

使能方式

PR链接

注意:mc2 优化与 flashcom 优化当前无法同时获取收益

当前该优化手段仅合入 vllm-ascend >= v0.10.0 版本,VeRL 使能方法如下:

  1. 在 runtime_env.yaml 添加环境变量 VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE=1
  2. 由于当前 VeRL 版本中 vllm-ascend 与 vllm 的导入顺序存在问题,会导致 vllm-ascend 对 vllm patch 无法生效,包括上述的mc2优化,因此需要额外在 verl/workers/fsdp_workers.py 中添加如下导入代码:
# # #  In verl/workers/fsdp_workers.py
import datetime
import json
import logging
import os
import warnings
from dataclasses import asdict
from typing import Any, Optional

import numpy as np
import psutil
import torch
# ------------------Add------------------------# 
import vllm
import vllm_ascend.patch.worker
# ---------------------------------------------# 
import torch.distributed
import torch.distributed as dist

该问题由于需要导入vllm-ascend,暂时难以合入VeRL原仓代码,当前只能通过上述手段进行使能。

Taskqueue level2

简述:

由于profiling中拖尾阶段存在一定host bound,因此使能taskqueue进行流水优化

Level 1优化:使能task_queue算子下发队列优化,将算子下发任务分为两段,一部分任务(主要是aclnn算子的调用)放在新增的二级流水上,一、二级流水通过算子队列传递任务,相互并行,通过部分掩盖减少整体的下发耗时,提升端到端性能。
请添加图片描述

Level 2优化:包含Level 1的优化并进一步平衡了一、二级流水的任务负载,主要是将workspace相关任务迁移至二级流水,掩盖效果更好,性能收益更大。
请添加图片描述

使能taskqueue level2 之后,host bound明显缓解,吞吐提升10%
请添加图片描述

使能方式

level1:添加环境变量 TASK_QUEUE_ENABLE=1

level2:添加环境变量 TASK_QUEUE_ENABLE=2

Flash Attention V2

简述

核心思想:解决memory bound,不减少计算量,优化读取逻辑,从而提高性能

FAV2原理解析:猛猿:图解大模型计算加速系列:Flash Attention V2,从原理到并行计算

使能方式

PR链接

torch_npu._npu_flash_attention接口替换为torch_npu.atb._npu_flash_attention_v2

请添加图片描述
请添加图片描述

开启FAV2无法与Chunked Prefill共同使用,长序列单推实测数据如下(throughput):

case test1 test2 test3 avg
chunked-prefill 300 302 331 311
FAV2 261 258 293 271

因此使能FAV2,关闭Chunked Prefill 在长序列场景下较开启 Chunked Prefill性能存在劣化,最终性能优化未使能FAV2

Flashcomm

简述

Flashcomm = 拆分AllReduce + 更改通信算子位置

  1. AllReduce 等价于 ReduceScatter + AllGather ,直接使用 AllReduce 算子可以减少通信次数,小并发场景下可以获得一定收益。但在大并发场景下,通信量增大,使用AllReduce 带来的启动开销收益并不明显。而且拆分后的两个算子具有更高的灵活性,可以在数学等价的前提下调整通信算子的位置,例如先对数据进行降维或者量化,再进行通信,可以有效减少通信的数据量;
  2. Transformer 结构中,在 AllReduce 算子之后,往往会有一些其他计算操作,如 RMSNorm、以及 MLA 中的降维计算等。这些计算过程会在不同卡上执行相同的计算操作,在小并发场景下可能耗时不高,但在大并发场景下会带来较大代价
    请添加图片描述

该特性与mc2存在原理性冲突,实测在32B长序列场景下mc2性能收益更大,因此暂未使用Flashcomm V1
使能方式

PR链接

加入环境变量VLLM_ASCEND_ENABLE_FLASHCOMM=1

Graph mode

简述:
开启图模式后,32B场景实测性能无明显提升,故此优化暂未开启;而其在别的场景下可能会有性能收益,NPU上图模式主要对下发进行优化,可以分别使用【taskqueue level2】和【taskqueue level 1 + 图模式】两种方案择优使用,taskqueue level2和图模式当前无法同时获取收益。

使能方式:
VeRL 中,通过设定参数 actor_rollout_ref.rollout.enforce_eager=False 来开启静态图模式,注意这里 eager mode 为动态图模式,因此置为 False 为开启静态图。

注:在最新的Verl代码中该参数默认值为False,如期望关闭图模式需要手动加入该参数并置为True

模块化切分

简述

dense模型的一个decoderlayer由attntion、layernorm、mlp三个模块组成,常规的切分策略是每个模块都用相同的并行策略切分,如qwen2.5的最优并行策略是dp1tp4,这种情况下每个attn结束后都需要有一个allreduce通信算子,导致通信等待较为严重,因此可以考虑分离并行策略,对attntion模块不切tp(matmul shape较小),只对mlp模块切tp。

平均生成长度4K场景下,PageAttention与MLP时间基本相当,PA耗时随长度增加持续增长,因此在此阶段Attention更应该进行切分,造成该优化手段在平均生成长度较长的场景下无明显性能优化

使能方式

PR链接

添加环境变量VLLM_ASCEND_ENABLE_MLP_OPTIMIZE=1

四、结语

在强化学习后训练中,短推长场景下显存还是主要限制点,针对显存进行优化是一大方向,另外性能低的另一个重要原因是负载不均,观察显存可以发现推理过程中大多数卡都推完的情况下,只在等少数卡完成推理,也会对性能造成一定影响,后续对强化学习的优化手段也主要从这两个方向出发。

目前主要的性能优化方向集中在训推异步上,开源方案有one step off、full async等,核心思想是让训推并行进行,这类方案通常会导致精度无法对齐,如何并行的同时保证训练效果是当前异步方案攻关的重点。

请添加图片描述

Logo

脑启社区是一个专注类脑智能领域的开发者社区。欢迎加入社区,共建类脑智能生态。社区为开发者提供了丰富的开源类脑工具软件、类脑算法模型及数据集、类脑知识库、类脑技术培训课程以及类脑应用案例等资源。

更多推荐