VeRL

导言

VeRL 作为RL领域趋势最火的开源仓,值得学习。

核心设计

Single/Multi-Controller

Google Pathways^2 在论文中对这两点做了非常详细的阐述。

  • Single-Controller 由一个中心控制器驱动所有 worker 工作,这个中心控制器向所有 worker 发送数据和计算指令,worker 可以运行不同的程序(MPMD)。
    • 优点是灵活、异构,即每个 worker 可以执行不同的算法;
    • 缺点是因为是单点控制,所以会引入额外的通信开销。
    • 典型的 Single-Controller,比如 Spark、TensorFlow1.0 的 non-SPMD 的模式。
  • Multi-Controller 无中心控制器,各个 worker 自驱,各自计算,通过一些通信原语进行同步,worker 运行相同的程序(SPMD)。
    • 优点是高效;
    • 缺点是因为同构所以不够灵活。
    • Multi-Controller 则更普遍,基本上主流的开源 RLHF 分布式训练框架,比如 Megatron-LM、Deepspeed、FSDP,以及 vLLM 都是 Multi-Controller 模式的。这是因为主流 RLHF 框架基本上是从预训练框架演进而来,而预训练基本上都是 Multi-Controller 模式。

HybridFlow

核心思想^1

  • 通过Single-Controller来实现RL算法控制的数据流,并易修改拓展,支持简单单进程调用;
  • 通过Multi-Controller来实现各个推理/训练组件:
    • 基于已有框架
    • 训练:megatron/FSDP1/FSDP2
    • 推理:vllm/SGlang

参数说明

三个batch_size的解释

其余参数参考,link

代码走读

以 GRPO + megatron后端为例,大致逻辑如下:

ray+GRPO算法

title
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class RayPPOTrainer:
def fit(self):
for epoch in range(current_epoch, self.config.trainer.total_epochs):
for batch_dict in self.train_dataloader:
# 推理
gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch_output)

# 可选:old_log_prob
old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)

# 可选:ref_log_prob
ref_log_prob = self.actor_rollout_wg.compute_ref_log_prob(batch)

# update actor
actor_output = self.actor_rollout_wg.update_actor(batch)

推理训练后端

self.actor_rollout_wg的四个子操作,根据后端不同,使用不同worker,megatorn后端使用megatron_worker.py

title
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class AsyncActorRolloutRefWorker(ActorRolloutRefWorker):
class ActorRolloutRefWorker(MegatronWorker, DistProfilerExtension):

def init_model(self):

self.actor = MegatronPPOActor(
config=actor_cfg,
model_config=self.actor_model_config,
hf_config=self.hf_config,
tf_config=self.tf_config,
actor_module=self.actor_module,
actor_optimizer=self.actor_optimizer,
)

self._build_rollout(trust_remote_code=self.config.model.get("trust_remote_code", False))

self.ref_policy = MegatronPPOActor(
config=self.config.ref,
model_config=self.ref_model_config,
hf_config=self.hf_config,
tf_config=self.tf_config,
actor_module=self.ref_module,
actor_optimizer=None,
)

def _build_rollout(self, trust_remote_code=False):
self.rollout = get_rollout_class(rollout_config.name, rollout_config.mode)(
config=rollout_config, model_config=model_config, device_mesh=rollout_device_mesh
)

def generate_sequences(self, prompts: DataProto):
output = self.rollout.generate_sequences(prompts=prompts)

def compute_log_prob(self, data: DataProto):
output, entropys, layers_topk_idx = self.actor.compute_log_prob(data=data, calculate_entropy=True)
# save 2 old_log_probs

def compute_ref_log_prob(self, data: DataProto):
output, _, _ = self.ref_policy.compute_log_prob(data=data, calculate_entropy=False)
# save 2 ref_log_prob

def update_actor(self, data: DataProto):
metrics = self.actor.update_policy(dataloader=dataloader)

推理后端

vllm

训练后端

megatron

待验证选项:

  • use_dynamic_bsz
  • self.use_fused_kernels

最终走 from megatron.core.pipeline_parallel import get_forward_backward_func

loss 计算

这里loss计算的代码又重构了

title
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

def loss_func
# pg_loss 看上去和上图 Token loss 是一个含义
pg_loss, pg_metrics = policy_loss_fn(
old_log_prob=old_log_prob,
log_prob=log_prob,
)
policy_loss = pg_loss
# policy_loss = pg_loss - entropy_coeff * entropy_loss

# kl loss计算如上图
kld = kl_penalty(logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type)
kl_loss = agg_loss(loss_mat=kld, loss_mask=response_mask, loss_agg_mode=self.config.loss_agg_mode)

policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef

old log p

用于重要性采样

vllm优化

FSDP训练优化

特性:multi-turn

swift现已支持multi-turn

VeRL 对 multi-turn的支持

  • 基于SGlang的multi-turn
    • 推荐使用
  • 基于agent_loop设计,参考zhihu,官方wiki
    • 但是没有商发, 评论区和issue都有相当多的bug

精度问题

KL Loss 和 Grad Norm

这是一个在深度强化学习(Deep Reinforcement Learning, DRL),特别是基于策略梯度(Policy Gradient)的方法,如VeRL(可能指的是某种基于[V]ariational或[V]olcano Engine的RL框架,结合RLHF/PPO)中常见的训练不稳定(Training Instability)现象,尤其是涉及到大规模模型(如LLM)的训练。

  1. KL Loss 陡增 ($D_{KL}$ Spike)
  • 含义: 在 VeRL/PPO 这类算法中,KL 散度(Kullback-Leibler Divergence)通常用于衡量新策略 ($\pi_{\text{new}}$)旧策略 ($\pi_{\text{old}}$) 之间的差异。它确保策略更新不会偏离太远。
  • 陡增的意义:
    • 策略剧烈变化: 意味着在某一步更新中,新策略相对于旧策略发生了极大的变化。
    • PPO/KL 约束失效: 这表明用于限制策略更新幅度(如 PPO 中的 ClippingKL 惩罚项)的机制可能失效或不足以应对当前数据批次。
    • 探索过度/错误更新: 模型可能在一个或少数几个批次数据上进行了过度激进的更新,导致策略分布远离了之前收集数据时的策略分布。
  • 合理表现:
    • 在稳定训练中,KL Loss 应该保持在一个较小的、受控的范围内,通常围绕设定的 KL 目标值($D_{KL}^{\text{target}}$)波动。
    • 陡增到非常大的值(如从 0.01 增加到 1.0 甚至更高)不合理的,通常预示着训练即将或已经进入崩溃(Divergence)状态,尽管 Reward 曲线可能由于延迟反馈或采样方差而暂时保持对齐或看似稳定。
  1. Grad Norm 陡增(梯度范数爆炸)
  • 含义: 梯度范数(Gradient Norm, $|g|_2$)是模型所有参数梯度向量的 $L_2$ 范数。它代表了权重更新的总体“力量”或“幅度”
  • 陡增的意义:
    • 梯度爆炸(Exploding Gradients): 这是深度学习训练中一个经典的不稳定问题。它意味着当前批次数据产生了非常大的梯度,如果没有处理(如梯度裁剪),将会导致权重被过度更新
    • 模型参数发散: 极大的梯度更新可能导致模型参数(权重)跳到损失函数的病态(pathological)区域,从而产生更大的梯度,形成恶性循环,最终可能导致参数变为 NaNInf,训练崩溃。
    • 罕见样本或数据问题: 可能是遇到了数据集中极端的难以处理的损坏的样本批次,导致了巨大的梯度。
  • 合理表现:
    • 在稳定训练中,Grad Norm 应该保持在一定范围的波动。如果使用了梯度裁剪(Gradient Clipping),则 Grad Norm 不应超过裁剪阈值。
    • 陡增到远超平时波动范围的值(尤其是在没有梯度裁剪或裁剪阈值太高的情况下)是不合理的,是训练不稳定的明确信号

为什么 Reward 曲线能对齐?

Reward 曲线对齐但 KL/Grad Norm 陡增,可能意味着:

  1. 策略网络和价值网络的独立性: 即使策略更新过度,如果价值函数网络(Value Function)的更新相对稳定,或者早期的 Reward 结果是在策略崩溃前就已经计算并记录下来,Reward 曲线可能短期内不受影响。
  2. 采样方差: 在强化学习中,Reward 曲线本身就具有很高的方差。一两个批次的剧烈更新可能不会立即体现在平滑后的 Reward 曲线中,尤其是在采样间隔较大的情况下。
  3. 延迟效应: 策略的剧烈变化导致模型输出“胡言乱语”或“无效动作”的负面影响,需要一段时间(更多步数、更多环境交互)才会反映在 Reward 的平均表现上。

诊断与解决建议

既然你是在 NPU 上复现 GPU 流程,请重点关注以下可能导致不稳定的因素:

潜在原因 解决建议
梯度爆炸(最常见) 应用或调整梯度裁剪(Gradient Clipping)。 确保 Grad Norm 限制在一个合理的阈值(例如 0.5 或 1.0)。如果已应用,请尝试降低阈值。
KL 惩罚系数/学习率 降低策略网络的学习率(Learning Rate)。 策略更新太快是 KL Loss 陡增的常见原因。尝试使用更小的 $\alpha$。
PPO Clipping 阈值 $\epsilon$ 检查 PPO 的 $\epsilon$ 参数。 尝试降低 clip_range(如从 0.2 降到 0.1),以更严格地限制策略比率 $r_t(\theta)$。
数据精度 检查 NPU 上的数据类型。 如果 NPU 使用 FP16 或其他低精度训练,低精度训练更容易引起梯度和损失的尖峰。尝试使用更高的精度(如 FP32BF16),或使用精度敏感的优化器/归一化
Warmup 阶段 确保在训练初期有适当的 Learning Rate Warmup 阶段,以避免初始阶段就产生巨大的梯度。
批量大小/经验重放 检查批次大小(Batch Size)。 批次太小会增加梯度的方差,可能导致尖峰。
NaN/Inf 问题 检查是否有任何 NaNInf 出现在 Loss 或参数中。KL Loss 的计算(如 $\log(\pi)$)在 $\pi \to 0$ 时可能产生极值,考虑添加 $\epsilon$ 或进行数值稳定处理。

总结: KL Loss 和 Grad Norm 突然陡增是不合理且危险的,是训练不稳定的强烈信号。你需要尽快采取措施(主要是梯度裁剪调整学习率/KL惩罚)来控制更新幅度,否则训练可能会在未来的步骤中崩溃。

参考文献

Author

Shaojie Tan

Posted on

2025-12-02

Updated on

2025-12-17

Licensed under