基于昇腾平台的SAPO训练实践
作者:昇腾实战派
一、算法背景:
强化学习在大语言模型训练中发挥着日益重要的作用,特别是在数学推理、编程和多模态理解等挑战性任务中。基于组的策略优化方法已成为实用方案:每个查询采样多个响应,组内归一化序列级奖励,并通过当前策略与行为策略间的重要性比率加权策略更新。然而,在Mixture-of-Experts(MoE)模型中,token级重要性比率往往表现出高方差——路由异构性和长响应会放大token间的偏差,这种方差增加了不稳定更新的可能性。
在GRPO算法中,在策略更新时,我们通常最大化如下目标
硬截断(Hard Clipping)的引入是为了防止策略更新步幅过大(即 r i , t ( θ ) r_{i,t}(\theta) ri,t(θ)偏离 1 太多)。然而,这种分段常数式的处理方式带来了两个问题:
- 非黑即白的梯度处理:当 r i , t ( θ ) r_{i,t}(\theta) ri,t(θ) 超出$1-\varepsilon ,1+\varepsilon $范围时,梯度直接被置为 0。这意味着稍稍“越界”的高价值样本将完全失去学习信号,降低了样本效率。
- 噪声与稳定性的博弈:如果放宽 ,虽然保留了更多样本,但引入了更多偏离策略分布(Off-policy)的噪声梯度,容易导致训练崩溃(Collapse)
同时GSPO (Group Sequence Policy Optimization) 尝试在序列级别进行截断,即基于序列整体的概率比率 S i ( θ ) S_i(\theta) Si(θ)进行 Clip。虽然这保证了序列层面的一致性,但它缺乏 Token 级的细粒度控制:如果一个序列中仅有几个 Token 严重偏离,GSPO 会丢弃整个序列的梯度,造成极大的浪费。
基于以上现象可以总结出几个问题:
- 高方差问题:token级重要性比率的高方差导致更新不稳定,尤其在MoE模型中更为严重
- 硬截断局限性:现有方法如GRPO和GSPO使用硬截断,难以在稳定性和有效学习间取得平衡
- 样本效率低下:当序列包含少数高度离策略token时,GSPO会抑制该序列的所有梯度,降低样本效率
二、SAPO算法详解
SAPO整体的优化目标:
其中 f i , t f_{i,t} fi,t为最关键的门控函数组件, f i , t f_{i,t} fi,t定义为
其中

为sigmoid函数

右图这个权重函数是一个以 r i , t = 1 r_{i,t}=1 ri,t=1为中心的钟形曲线(类似于高斯分布,但基于 Sigmoid 导数构建):
SAPO 与硬截断(Hard Clipping)在目标函数值与梯度权重上的对比。左图显示 SAPO 的目标函数是光滑的,右图显示 SAPO 的梯度权重随比率偏离 1 而平滑下降,而硬截断则是阶跃式的归零
相较于硬截断的优势: SAPO 实现了一个**连续的信任区域(Continuous Trust Region)**。它不会在某一个阈值突然切断梯度,而是根据偏离程度“软性”地降低权重。即使样本稍微 Off-policy,只要包含有价值的信息,模型依然能以较小的步长进行学习。这种机制在保留学习信号的同时,抑制了过大更新带来的不稳定性。
SAPO 的另一个重要创新是针对正负优势(Advantage)使用了不同的温度参数 τ i , t \tau_{i,t} τi,t:
论文中设置 τ n e g > τ p o s \tau_{neg}>\tau_{pos} τneg>τpos由于 w i , t w_{i,t} wi,t的衰减速度由 τ \tau τ控制( τ \tau τ越大,钟形曲线越窄,衰减越快),这意味着负优势样本的梯度权重会比正优势样本衰减得更快
在 LLM 的大词表(通常 > 100k)场景下,负向更新具有极大的破坏力。因为它会盲目地提升成千上万个未采样 Token 的概率。虽然这能在一定程度上增加熵(探索),但在训练初期或 Off-policy 严重时,这种更新会迅速引入大量噪声,导致 Logit 分布混乱,引发训练崩溃。
通过设置 τ n e g > τ p o s \tau_{neg}>\tau_{pos} τneg>τpos,SAPO 更加激进地抑制那些 Off-policy 的负样本更新,从而显著提升了训练的稳定性。
SAPO 与 GRPO/GSPO 的统一
统一的 Surrogate Objective:

- GRPO 使用分段常数函数作为 f i , t ′ f'_{i,t} fi,t′(导数),即在 [ 1 − ϵ , 1 + ϵ ] [1-\epsilon ,1+\epsilon ] [1−ϵ,1+ϵ]内为 1,否则为 0。
- GSPO 使用序列级比率 s i ( θ ) s_i(\theta) si(θ)进行截断,其 f i , t f_{i,t} fi,t在同一序列内是常数。
- SAPO 使用 s e c h 2 sech^2 sech2形式的软门控。
研究团队使用从Qwen3-30B-A3B-Base微调的冷启动模型在数学推理查询上进行实验。报告了AIME25、HMMT25和BeyondAIME基准上的训练奖励和验证性能(平均Pass@1,16个样本)。
实验结果显示,SAPO在所有基准测试上持续改进模型性能,与GSPO和GRPO-R2相比实现了更高的稳定性和更强的最终性能。

三、调试实践
算法pr:https://github.com/verl-project/verl/pull/4345
基于verl代码分析,可以发现其门控函数相关代码位置位于:verl/trainer/ppo/core_algos.py。
代码中主要新构造了gate function并基于gate function计算loss,同时由于SAPO中不使用clip,将clip相关参数设为0确保兼容。
def compute_policy_loss_sapo(
old_log_prob: torch.Tensor,
log_prob: torch.Tensor,
advantages: torch.Tensor,
response_mask: torch.Tensor,
loss_agg_mode: str = "seq-mean-token-mean",
config: Optional[ActorConfig] = None,
rollout_is_weights: torch.Tensor | None = None,
) -> tuple[torch.Tensor, dict[str, Any]]:
"""
Compute the smoothed policy objective and related metrics for SAPO.
See https://arxiv.org/pdf/2511.20347 for more details.
Args:
old_log_prob (torch.Tensor):
Log-probabilities of actions under the old policy, shape (batch_size, response_length).
log_prob (torch.Tensor):
Log-probabilities of actions under the current policy, shape (batch_size, response_length).
advantages (torch.Tensor):
Advantage estimates for each action, shape (batch_size, response_length).
response_mask (torch.Tensor):
Mask indicating which tokens to include in the loss, shape (batch_size, response_length).
loss_agg_mode (str, optional):
Aggregation mode for `agg_loss`. For SAPO, it is recommended to use "seq-mean-token-mean".
"""
assert config is not None
assert isinstance(config, ActorConfig)
# temperature for positive and negative token updates
tau_pos = torch.as_tensor(config.tau_pos, dtype=advantages.dtype, device=advantages.device)
tau_neg = torch.as_tensor(config.tau_neg, dtype=advantages.dtype, device=advantages.device)
def gate_function(x, tau):
"""The gating function used in SAPO"""
return torch.sigmoid(tau * (x - 1.0)) * (4.0 / tau)
# compute IS at token level:
# r_{i,t}(θ) = π_θ(y_{i,t}|x, y_{i,<t}) / π_θold(y_{i,t}|x, y_{i,<t})]
# In log space: log(r_{i,t}(θ)) = log_prob - ol_log_prob
negative_approx_kl = log_prob - old_log_prob
# Clamp negative_approx_kl for stability
negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)
# finally exp() to remove log and get r_{i,t}(θ)
ratio = torch.exp(negative_approx_kl)
# tau_{i,t} is tau_pos if adv > 0 else tau_neg
taus = torch.where(
condition=advantages > 0,
input=tau_pos, # if A_{i,t} > 0 we set to tau_pos
other=tau_neg, # if A_{i,t} <= 0 we set to tau_neg
)
# compute the gates f_{i,t}(r_{i,t}(θ)) at token level
gates = gate_function(ratio, taus)
# compute policy gradient loss
pg_losses = -gates * advantages
# Apply rollout correction weights if provided
if rollout_is_weights is not None:
pg_losses = pg_losses * rollout_is_weights
# for SAPO, we need to aggregate the loss at the sequence level (seq-mean-token-mean)
pg_loss = agg_loss(
loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode="seq-mean-token-mean", **config.global_batch_info
)
# For compatibility, return zero for both pg_clipfrac and pg_clipfrac_lower (not used in SAPO)
pg_clipfrac = torch.tensor(0.0, device=pg_loss.device)
pg_clipfrac_lower = torch.tensor(0.0, device=pg_loss.device)
# compute KL for metrics tracking
ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)
# return metrics dict
pg_metrics = {
"actor/pg_clipfrac": pg_clipfrac.detach().item(),
"actor/ppo_kl": ppo_kl.detach().item(),
"actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(),
}
return pg_loss, pg_metrics
环境搭建
硬件:Atlas 800T A2
参考每日镜像搭建博客:https://blog.csdn.net/Lumos_Lovegood/article/details/160614580
数据集准备
数据集链接:
训练数据集:https://huggingface.co/datasets/BytedTsinghua-SIA/AIME-2024
验证数据集:https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k
官方GPU脚本链接:https://github.com/verl-project/verl/blob/main/examples/sapo_trainer/run_qwen30b_sapo.sh#
数据集转换,可使用下载好的数据集或在线下载(在线下载需要连通huggingface)
python /verl/examples/data_preprocess/dapo_multiturn_w_tool.py --local_dataset_path /home/c00835985/dataset/dapo
Qwen3-8B打通
基于仓上原有的NPU GRPO Qwen3-8B(FSDP+VLLM)脚本进行修改打通:
set -x
mkdir -p logs
ulimit -n 32768
## Basic Environment Settings
export RAY_DEDUP_LOGS=0
export HYDRA_FULL_ERROR=1
export TASK_QUEUE_ENABLE=1
export HCCL_EXEC_TIMEOUT=3600
export HCCL_CONNECT_TIMEOUT=3600
export HCCL_ASYNC_ERROR_HANDLING=0
export CPU_AFFINITY_CONF=1
export VLLM_USE_V1=1
project_name='SAPO-Qwen3'
exp_name='SAPO-Qwen3-8B-npu'
gen_tp=2
RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/Qwen3-8B"}
CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/dataset/dapo_processed/train.parquet"}
TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/dataset/aime-24_processed/train.parquet"}
# reference policy
use_kl_in_reward=False
kl_coef=0.001
use_kl_loss=False
kl_loss_coef=0.001
# ------Algorithm settings-------
# Positive and negative tau for smoothing function in SAPO (https://arxiv.org/pdf/2511.20347)
# default values used in the paper with Qwen3-30B-A3B-Base
# clipping is not used in SAPO!
loss_mode=sapo # explicitly specify sapo! default is vanilla and is not compatible with SAPO. It uses clipping instead of smoothing.
tau_pos=1.0
tau_neg=1.05
gae_gamma=1.0
gae_lam=0.95
python3 -m verl.trainer.main_ppo \
algorithm.adv_estimator=grpo \
algorithm.use_kl_in_reward=$use_kl_in_reward \
algorithm.kl_ctrl.kl_coef=$kl_coef \
algorithm.gamma=$gae_gamma \
algorithm.lam=$gae_lam \
data.train_files="${TRAIN_FILE}" \
data.val_files="${TEST_FILE}" \
data.train_batch_size=256 \
data.max_prompt_length=512 \
data.max_response_length=1024 \
data.filter_overlong_prompts=True \
data.filter_overlong_prompts_workers=64 \
data.truncation='error' \
actor_rollout_ref.model.path=${MODEL_PATH} \
actor_rollout_ref.actor.optim.lr=1e-6 \
actor_rollout_ref.model.use_remove_padding=True \
actor_rollout_ref.actor.ppo_mini_batch_size=64 \
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \
actor_rollout_ref.actor.tau_pos=$tau_pos \
actor_rollout_ref.actor.tau_neg=$tau_neg \
actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \
actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \
actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \
actor_rollout_ref.actor.use_torch_compile=False \
actor_rollout_ref.ref.use_torch_compile=False \
actor_rollout_ref.model.enable_gradient_checkpointing=True \
actor_rollout_ref.actor.fsdp_config.param_offload=False \
actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \
actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
actor_rollout_ref.rollout.name=vllm \
actor_rollout_ref.rollout.enforce_eager=True \
actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
actor_rollout_ref.rollout.n=5 \
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
actor_rollout_ref.ref.fsdp_config.param_offload=True \
trainer.critic_warmup=0 \
trainer.logger='["console"]' \
trainer.project_name="${project_name}" \
trainer.experiment_name="${exp_name}" \
trainer.n_gpus_per_node=8 \
trainer.nnodes=1 \
trainer.default_local_dir=${CKPTS_DIR} \
trainer.resume_mode=auto \
actor_rollout_ref.actor.fsdp_config.forward_prefetch=True \
actor_rollout_ref.ref.fsdp_config.forward_prefetch=True \
++actor_rollout_ref.actor.entropy_from_logits_with_chunking=True \
++actor_rollout_ref.ref.entropy_from_logits_with_chunking=True \
trainer.val_before_train=True \
trainer.save_freq=-1 \
trainer.test_freq=5 \
trainer.total_epochs=15
主要新增了几个算法相关的参数,如tau_pos,tau_neg,loss_mode指定为sapo等,参数值的选取参照论文中的参考值
同时指定enforce_eager为True来关闭推理的图模式
SAPO Qwen3-8B 打通结果



结论:reward正常上升,val acc分数正常上升
SAPO vs GRPO Qwen3-8B 训练效果对比


在相同训练超参下,可以看到SAPO对比GRPO的收敛速度更快,reward得分也略高于GRPO算法,在Qwen3-8B的小规模dense模型下仍具有优势。
四、打通中遇到的问题
1、使用新版本verl镜像,推理开启图模式报错

通过增加如下参数关闭图模式进行规避
actor_rollout_ref.rollout.enforce_eager=True
2、使用2.4的verl社区镜像拉起报错,同时使用该镜像跑仓上的GRPO脚本也会报同样错误

报错显示rollout.update_weights_bucket_megabytes过小,需要增大该参数,但在脚本中新增该参数后又会报config中无该参数的错误,通过回退verl版本进行规避
原因可能是因为当前verl主线main最近有较多合入导致冲突
3、减层调试Qwen3-30B脚本时,FSDP optimizer offload报错

通过如下参数关闭FSDP卸载规避
actor_rollout_ref.actor.fsdp_config.optimizer_offload=False
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。
更多推荐
所有评论(0)