[FEATURE] Support Dynamic Sampling with Dynamic Batching · Pull Request !520 · MindSpore/mindrlhf

1. [FEATURE] Support Dynamic Sampling
---
**What type of PR is this?**

/kind feature

**What does this PR do / why do we need it**:

该 PR 在 GRPO 算法基础上，引入了 DAPO 算法中提出的动态采样（Dynamic Sampling）技术，其特性为：

- 在 rollout 过程中，过滤掉 reward 均为 $0$（全错）或 $1$（全对）的 prompt，并循环补充新的 prompt，以保证每个 batch 的样本数达到 batch_size，并且其中每个样本都能提供有效的梯度信息；

![输入图片说明](https://foruda.gitee.com/images/1757901486664042267/7edf049b_9875800.png "屏幕截图")

由于采样（Sampling）效率是 RL 训练效率的最大瓶颈（占比整个训练流程 $70%$ 以上的时间），每个 step 取数次固定大小的 `batch_size` 拼接效率劣于取单次更大的 `batch_size` 效率，因此我们希望能寻找到一个拟合的最优 `batch_size` 数值，使得过滤后的样本数量刚好满足 `train bs`（减少多轮采样，避免样本浪费）。

- 为了尽可能避免“循环补充”这一过程，需要动态调整 `batch_size`，使得每个 step 能进行“过采样”，即在“宁滥毋缺”的前提下最小化“滥”的 prompt 数目。该特性我采用设定滑动窗口 $B$（默认为 $8$），根据历史轮数的实际采样数量和保留数量，预测下一轮的最优采样数。

- 初始（前 $B$ 个 step）采用保守的预设值，如 $3\times \text{target}$；

- 之后我们采用两个预测模型：

- 预测模型一（predict_x1），综合均值和标准差：
$$
\text{predict_x1}=(\mu_X+\alpha \cdot \sigma_X)\cdot \text{ratio}_1
$$
     - 预测模型二（predict_x2），指数加权移动平均：
$$
\text{predict_x2}=\left(\frac{\sum\limits_{i=0}^{B-1} \gamma^i \cdot x_{t-i}}{\sum\limits_{i=0}^{B-1} \gamma^i}\right) \cdot \text{ratio}_2
$$
     - 最终预测结果取两者中的更大者，其中 $q$ 为一安全系数：
$$
\hat{x_{t+1}}=\max(\text{predict_x1},\text{predict_x2})\cdot q
$$

- 将上述流程得到的 batch 生成对应的 rollout 数据（引入了 `rollout_data_manager` 模块进行管理），传入训练模型更新权重。

关于上述预测器的超参配置，我是通过运行实验，从 log 中将 $\text{kept_prompt_size}$ 与 $\text{total_prompt_size}$ 的数据抓取出来，通过调整参数，找到一个“最佳”拟合曲线，经实验，运行效率对比如下：

- baseline 使用默认过采样配置 3 * train_bs，dynamic_gen_bs 为动态过采样。
- num_generation 表示一个问题回答的次数（8/16/32）。
- 实验训推 400 步，统计 100-400 步的平均每步 rollout 用时，记录 rollout 性能的提升：

| num_generation  |  baseline   |  dynamic_gen_bs | 性能提升 |
|---|---|---|---|
| 8  | 123.999s  | 120.727s  | 2.6% |
| 16  | 160.040s  |  129.752s | 18.9% |
| 32  |  252.962s |  165.638s | 34.5% |

可以看出，随着 num_generation 的增大，开启“过采样”方案的性能优势显著。这里顺带记录在不同的 num_generation 下实际采样数与 baseline 的对比：

![输入图片说明](https://foruda.gitee.com/images/1757904650420073050/fde7a7d6_9875800.png "屏幕截图")

**Self-Test result**:

为了验证该 PR 代码修改的正确性，我采用了如下两个方式：

- 将修改的部分（主要集中在 `grpo_experience_maker` 中的 `dynamic_sampling` 函数、`rollout_data_manager` 中的 `filter_by_index` 及 `current_data` 和 `saved_data` 的传递）打印了 log 信息，在 log 文件中监视这些调试信息，确保数据传输完整无误；
- 在 `grpo_config.yaml` 中加上 dynamic sampling 的 `filter_enable` 标签（用于决定是否开启该 dynamic sampling 的开关，默认 False 关闭），跑消融实验，并用脚本将 average reward 信息可视化（见下图）。由于模型训练初期，大部分 prompt 回答全错的概率会更大，因此开启过滤理论上会使得 average reward 偏高，与图中一致，符合预期。

![对比图](https://foruda.gitee.com/images/1757561344670672392/899ad9f7_9875800.png "屏幕截图")

类型	指派人员	状态
审查	yanghaoran	已完成 (0/0人)
测试	yanghaoran	已完成 (0/0人)

MindSpore/mindrlhf暂停 .gitee-modal { width: 500px !important; }

搜索帮助

MindSpore/mindrlhf
暂停