AReno-OPD: On-Policy Distillation for AReno

为 AReno 开源框架添加 On-Policy Distillation (OPD) 算法支持。OPD 是一种在线策略蒸馏技术，学生模型生成响应（rollout），冻结的教师模型为其打分，学生通过最小化与教师分布之间的 KL 散度来学习。

目录

• 背景介绍
• 算法原理
• 项目结构
• 快速开始
• 配置说明
• API 参考
• 测试
• 算法对比
• License

背景介绍

什么是 On-Policy Distillation?

On-Policy Distillation (OPD) 是一种面向大语言模型后训练的知识蒸馏技术。与传统的离线蒸馏不同，OPD 是在线策略（on-policy） 的——学生模型自己生成响应，然后让教师模型对这些响应进行评分。这使得学生能够在其自身的分布上学习，从而更有效地匹配教师的行为。

为什么在 AReno 中实现 OPD?

AReno 是一个面向大模型后训练的开源框架，致力于将训练框架和推理引擎放在同一套基础设施中，实现训推一体。AReno 已支持 PPO、GRPO、DPO、SFT 等算法，而 OPD 的加入填补了在线策略蒸馏这一重要方向，为用户提供更丰富的模型优化选择。

算法原理

工作流程

┌──────────────┐     ┌──────────────────┐     ┌──────────────────────┐
│   提示数据    │ ──► │  学生模型(可训练)  │ ──► │  Rollout: 每提示生成  │
│   (Prompts)  │     │  (Student)       │     │  n_samples 个响应     │
└──────────────┘     └──────────────────┘     └──────────┬───────────┘
                                                          │
                                                          ▼
                                                ┌──────────────────┐
                                                │  教师模型(冻结)    │
                                                │  (Teacher)        │
                                                │  对响应 token 评分 │
                                                │  (logprobs)       │
                                                └──────────┬───────────┘
                                                          │
                                                          ▼
                                                ┌──────────────────┐
                                                │  KL 散度损失:     │
                                                │  KL(student ||   │
                                                │    teacher)      │
                                                └──────────┬───────────┘
                                                          │
                                                          ▼
                                                ┌──────────────────┐
                                                │  学生模型更新     │
                                                │  (梯度下降)       │
                                                └──────────────────┘

数学公式

对于学生策略 π_s 生成的每个响应 token a：

• 学生 log-probability: log π_s(a | context)
• 教师 log-probability: log π_t(a | context)

蒸馏损失为学生到教师分布的 KL 散度：

Loss = KL(π_s || π_t) = E_{a ~ π_s}[log π_s(a) - log π_t(a)]
                    ≈ mean(log π_s(a) - log π_t(a))   # 在响应 token 上取平均

教师模型的 log-probabilities 被脱离计算图（.detach()），梯度仅流经学生网络：

∇Loss ≈ mean(∇log π_s(a))

温度缩放

通过温度参数 τ 软化两个分布：

Loss_τ = mean((log π_s(a) - log π_t(a)) / τ) = Loss / τ

更高的温度使分布更平滑，降低损失幅度。

与 PPO 的关键区别

项目结构

AReno-OPD/
├── areno/
│   └── api/
│       ├── algorithms.py              # 算法注册表（注册 opd）
│       ├── trainer_config.py          # OPDTrainerConfig 配置类
│       ├── loss_fns/
│       │   ├── __init__.py            # 导出 opd_loss_fn
│       │   └── opd.py                 # OPD 损失函数实现
│       └── trainers/
│           ├── __init__.py            # 导出 OPDTrainer
│           └── opd.py                 # OPD 训练器实现
├── tests/
│   ├── test_opd_loss.py               # 损失函数单元测试（12 个用例）
│   └── test_opd_algorithm.py          # 算法注册与配置测试（6 个用例）
├── docs/
│   └── opd_algorithm.md               # 详细算法文档
├── README.md                          # 本文件
└── .gitignore

核心文件说明

快速开始

前提条件

• Python 3.10+
• PyTorch 2.0+
• AReno 框架（需先安装，见 inclusionAI/AReno）

安装

方式一：从 PyPI 安装（待发布）

pip install areno-opd

方式二：从源码安装

git clone https://github.com/Krystal579-max/AReno-OPD.git
cd AReno-OPD
pip install -e .

方式三：合并到 AReno 主仓库

# 克隆 AReno 主仓库
git clone https://github.com/inclusionAI/AReno.git
cd AReno

# 将 OPD 代码合并到 AReno 中
cp -r /path/to/AReno-OPD/areno/* ./areno/

训练脚本

import areno.api
from areno.api.trainer_config import OPDTrainerConfig
from areno.api.trainer_factory import build_trainer

# 配置 OPD 训练
config = OPDTrainerConfig(
    algo="opd",
    ckpt="/path/to/student/model",          # 学生模型（可训练）
    dataset_path="/path/to/training/data",
    ref_ckpt="/path/to/teacher/model",       # 教师模型（冻结，通常更强）
    opd_kl_coef=1.0,                         # KL 损失系数
    opd_temperature=1.0,                     # 温度参数
    n_samples=8,                             # 每提示采样数
    batch_size=32,
    epochs=3,
)

# 构建训练器
trainer = build_trainer(
    config,
    instance=areno.api.Trainer(...),
    dataset=dataset,
    reward_fn=None,      # OPD 不需要奖励函数
    loss_fn=areno.api.loss_fns.opd_loss_fn,
)

# 开始训练
trainer.fit()

CLI 命令行使用

python -m areno train \
    --algo opd \
    --ckpt /path/to/student/model \
    --dataset /path/to/data \
    --ref-ckpt /path/to/teacher/model \
    --opd-kl-coef 1.0 \
    --opd-temperature 1.0 \
    --n-samples 8 \
    --batch-size 32

配置说明

OPDTrainerConfig 参数

从 RolloutTrainerConfig 继承的关键参数：

API 参考

opd_loss_fn(data_pack, logprobs, *, kl_coef=1.0, temperature=1.0)

参数：

• data_pack (dict): 包含响应掩码和 ref_logprobs（教师 log-probabilities）的数据字典
• logprobs (torch.Tensor): 当前学生 log-probabilities（可微）
• kl_coef (float, 可选): KL 损失缩放系数，默认 1.0
• temperature (float, 可选): 温度参数，默认 1.0

返回：

• loss (torch.Tensor): 标量损失值
• stats (dict): 诊断指标，包括 opd_loss, opd_kl, opd_student_logprob_mean, opd_teacher_logprob_mean, opd_kl_coef, opd_temperature

OPDTrainer

继承自 PolicyOnlyTrainer，主要方法：

OPDTrainerConfig

@dataclass(slots=True)
class OPDTrainerConfig(RolloutTrainerConfig):
    ref_ckpt: str | None = None
    opd_kl_coef: float = 1.0
    opd_temperature: float = 1.0

测试

运行所有 OPD 相关测试：

# 运行损失函数测试
python -m pytest tests/test_opd_loss.py -v

# 运行算法注册测试
python -m pytest tests/test_opd_algorithm.py -v

# 运行所有测试
python -m pytest tests/test_opd*.py -v

测试覆盖

损失函数测试（test_opd_loss.py，12 个用例）：

算法注册测试（test_opd_algorithm.py，6 个用例）：

算法对比

何时选择 OPD?

• 您有一个更强的教师模型：如果您可以访问一个更大或更优的模型（如 GPT-4、Qwen-72B 等），OPD 可以有效将其知识蒸馏到小模型中
• 需要在线策略学习：学生模型从其自身的分布中学习，而非静态数据集
• 简化训练流程：不需要奖励模型、critic 网络或优势计算，实现更简单
• 资源受限：相比 PPO，OPD 需要更少的 GPU 显存（无 critic 网络）

License

本项目基于 Apache License 2.0 开源协议。AReno 主仓库同样采用 Apache 2.0 协议。

参考

• AReno 开源框架 — 大模型后训练训推一体框架
• Knowledge Distillation: A Survey — 知识蒸馏综述
• On-Policy Distillation: A Survey — 在线策略蒸馏技术
• PPO: Proximal Policy Optimization — 近端策略优化
• GRPO: Group Relative Policy Optimization — 组相对策略优化