Merge 864bd67437 into d65ef2c00e

README.md

@@ -911,15 +911,17 @@ response = client.chat.completions.create(
 
 在介绍实现具体算法之前，我先以个人理解的极简视角，阐述所有Policy Optimization (PO)算法的统一共性。
 
-所有RL算法的本质都只是在优化一个期望：
+说到底，这里讨论的 PO 算法都只是在优化一个期望：
 
-$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q), o \sim \pi(O|q)} \left[ \underbrace{f(r_t)}_{\text{策略项}} \cdot \underbrace{g(A_t)}_{\text{优势项}} - \underbrace{h(\text{KL}_t)}_{\text{正则项}} \right]$$
+$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q),\, o \sim \pi_\theta(\cdot \mid q)} \left[ \underbrace{\Phi(r_t, A_t)}_{\text{策略目标}} - \underbrace{h(\text{KL}_t)}_{\text{正则项}} \right]$$
 
-训练时，只需**最小化负目标函数**，即: $\mathcal{L}_{PO} = -\mathcal{J}_{PO}$
+训练时，只需**最小化负目标函数**，即：
+
+$$\mathcal{L}_{PO} = -\mathcal{J}_{PO}$$
 
 这个框架只包含三个核心组件：
-* **策略项** $f(r_t)$: 如何使用概率比 $r_t$? 即告诉模型新旧策略偏差有多大，是否探索到了更好的token
-* **优势项** $g(A_t)$: 如何计算优势 $A_t$, 这很重要！大模型算对定积分也不足为奇，小模型回答对加减法优势通常都是正的
+* **策略项** $\Phi(r_t, A_t)$: 如何结合概率比 $r_t$ 和优势 $A_t$ 更新策略
+* **优势项** $A_t$: 如何计算优势，这很重要！大模型算对定积分也不足为奇，小模型回答对加减法优势通常都是正的
 * **正则项** $h(\text{KL}_t)$: 如何约束变化幅度 $\text{KL}_t$, 既防止跑偏又防止管的太死
 
 <details>
@@ -929,7 +931,7 @@ $$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q), o \sim \pi(O|q)} \left[ \underbrac
 |------|------|------|------|
 | $q$ | 问题/提示词 | 从数据集 $P(Q)$ 中采样 | - |
 | $o$ | 模型输出序列 | 由策略 $\pi$ 生成 | - |
-| $r_t$ | 概率比 | $r_t = \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{ref}(o_t \mid q, o_{<t})}$ | $(0, +\infty)$ |
+| $r_t$ | 概率比 | $r_t = \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{\mathrm{old}}(o_t \mid q, o_{<t})}$ | $(0, +\infty)$ |
 | $A_t$ | 优势函数 | 衡量某个动作相比基线有多好 | $(-\infty, +\infty)$ |
 | $\text{KL}_t$ | KL散度 | 防止策略偏离参考模型太远 | $[0, +\infty)$ |
 
@@ -1085,7 +1087,7 @@ $$\mathcal{L}_{PPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \text{clip}(r_t, 1-\v
 
 其中：
 - **策略项**: $f(r_t) = \min(r_t, \text{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (裁剪概率比防止更新过激)
-- **优势项**: $g(A_t) = R - V(s)$ (通过Critic网络估计价值函数)
+- **优势项**: $A_t$ 通常由Critic网络估计，也可以使用GAE进行计算
 - **正则项**: $h(\text{KL}_t) = \beta \cdot \mathbb{E}[\text{KL}]$ (全局KL散度约束)
 
 对比DPO而言，
@@ -1126,8 +1128,8 @@ python train_ppo.py
 $$\mathcal{L}_{GRPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t) - \beta \cdot \text{KL}_t\right]$$
 
 其中：
-- **策略项**: $f(r_t) = \min(r_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (使用概率比的对称 clip 裁剪)
-- **优势项**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (组内归一化，消除Critic网络)
+- **策略项**: $f(r_t, A_t) = \min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ (对概率比和优势项一起做裁剪，防止更新过激)
+- **优势项**: $g(A_{i,j}) = \frac{R_{i,j} - \mu_i}{\sigma_i + \epsilon}$ (组内归一化，消除Critic网络)
 - **正则项**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token级KL散度约束)
 
 对于同一个问题，模型生成 N 个回答并计算各自奖励，再用组内平均奖励作为 baseline。高于 baseline 的回答被鼓励，低于 baseline 的回答被抑制，因此无需额外训练 critic 网络。
@@ -1160,11 +1162,11 @@ CISPO 的关注点并不是重新设计 group baseline，而是用非常小的 l
 
 **CISPO损失**：
 
-$$\mathcal{L}_{CISPO} = -\mathbb{E}\left[\min(r_t, \varepsilon_{max}) \cdot A_t \cdot \log \pi_\theta(a_t|s) - \beta \cdot \text{KL}_t\right]$$
+$$\mathcal{L}_{CISPO} = -\mathbb{E}\left[\min(r_t, \varepsilon_{\mathrm{high}}) \cdot A_t \cdot \log \pi_\theta(a_t|s) - \beta \cdot \text{KL}_t\right]$$
 
 其中：
-- **策略项**: $f(r_t) = \min(r_t, \varepsilon_{max}) \cdot \log \pi_\theta(a_t|s)$ (ratio 只作为裁剪后的权重)
-- **优势项**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (可直接沿用 GRPO 的组内相对优势)
+- **策略项**: $f(r_t) = \min(r_t, \varepsilon_{\mathrm{high}}) \cdot \log \pi_\theta(a_t|s)$ (ratio 只作为裁剪后的权重)
+- **优势项**: $g(A_{i,j}) = \frac{R_{i,j} - \mu_i}{\sigma_i + \epsilon}$ (可直接沿用 GRPO 的组内相对优势)
 - **正则项**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token级KL散度约束)
 
 CISPO在GRPO基础上，把原本容易被clip成常数的策略项改写成“裁剪权重 × log 概率”的形式。这样ratio即使被截断，也不会把梯度路径一起截断。因此可以直接把CISPO视作GRPO的loss变体来实现，而不是单独维护一套独立脚本。这里不再单列实验。只需在 `train_grpo.py` 把 `loss_type` 配置为 `cispo`，其余训练流程仍沿用 GRPO 的分组采样、奖励计算与优势构造逻辑即可。
@@ -1269,9 +1271,9 @@ python eval_toolcall.py --weight agent
 | 算法 | 策略项 $f(r_t)$ | 优势项 $g(A_t)$ | 正则项 $h(\text{KL}_t)$ | 训练模型数 |
 |------|----------------|----------------|----------------------|----------|
 | **DPO** | $\log r_w - \log r_l$ | 无显式优势项 | 隐含在 $\beta$ 中 | 1 (前向参与 2) | 
-| **PPO** | $\min(r, \text{clip}(r))$ | $R - V(s)$ | $\beta \cdot \mathbb{E}[\text{KL}]$ | 2 | 
-| **GRPO** | $\min(r, \text{clip}(r))$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 |
-| **CISPO** | $\mathrm{clip}(r, 0, \varepsilon_{max}) \cdot A_t \cdot \log \pi_\theta$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 | 
+| **PPO** | $\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ | $A_t$（通常由 Critic 估计，也可以使用 GAE） | $\beta \cdot \mathbb{E}[\text{KL}]$ | 2 |
+| **GRPO** | $\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ | $A_{i,j}=\frac{R_{i,j}-\mu_i}{\sigma_i+\epsilon}$ | $\beta \cdot \text{KL}_t$ | 1 |
+| **CISPO** | $\mathrm{clip}(r, 0, \varepsilon_{\mathrm{high}}) \cdot A_t \cdot \log \pi_\theta$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 | 
 
 **说白了，这些 RL 算法不是割裂独立的，而是在统一优化视角下，对同一目标函数进行不同设计权衡后形成的自然变体，呈现为一种优美自洽的统一。**
 

README_en.md

@@ -911,15 +911,17 @@ However, in specific practice, their differences go beyond just the feedback sou
 
 Before introducing the implementation of specific algorithms, let me first describe the unified commonality of all Policy Optimization (PO) algorithms from my own minimalist perspective.
 
-The essence of all RL algorithms is just optimizing an expectation:
+Put simply, the PO algorithms discussed here are all optimizing an expectation:
 
-$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q), o \sim \pi(O|q)} \left[ \underbrace{f(r_t)}_{\text{policy term}} \cdot \underbrace{g(A_t)}_{\text{advantage term}} - \underbrace{h(\text{KL}_t)}_{\text{regularization term}} \right]$$
+$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q),\, o \sim \pi_\theta(\cdot \mid q)} \left[ \underbrace{\Phi(r_t, A_t)}_{\text{policy objective}} - \underbrace{h(\text{KL}_t)}_{\text{regularization term}} \right]$$
 
-During training, one only needs to **minimize the negative objective function**, i.e.: $\mathcal{L}_{PO} = -\mathcal{J}_{PO}$
+During training, one only needs to **minimize the negative objective function**, i.e.:
+
+$$\mathcal{L}_{PO} = -\mathcal{J}_{PO}$$
 
 This framework contains only three core components:
-* **Policy term** $f(r_t)$: How to use the probability ratio $r_t$? It tells the model how far the new and old policies have deviated, and whether better tokens have been explored
-* **Advantage term** $g(A_t)$: How to compute the advantage $A_t$, this is very important! It's no surprise that large models can solve definite integrals correctly, but for small models, even getting addition and subtraction right usually yields a positive advantage
+* **Policy term** $\Phi(r_t, A_t)$: How to combine the probability ratio $r_t$ and advantage $A_t$ to update the policy
+* **Advantage term** $A_t$: How to compute the advantage, this is very important! It's no surprise that large models can solve definite integrals correctly, but for small models, even getting addition and subtraction right usually yields a positive advantage
 * **Regularization term** $h(\text{KL}_t)$: How to constrain the magnitude of change $\text{KL}_t$, preventing both drifting too far and constraining too tightly
 
 <details>
@@ -929,7 +931,7 @@ This framework contains only three core components:
 |--------|---------|-------------|-------|
 | $q$ | Question/Prompt | Sampled from dataset $P(Q)$ | - |
 | $o$ | Model output sequence | Generated by policy $\pi$ | - |
-| $r_t$ | Probability ratio | $r_t = \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{ref}(o_t \mid q, o_{<t})}$ | $(0, +\infty)$ |
+| $r_t$ | Probability ratio | $r_t = \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{\mathrm{old}}(o_t \mid q, o_{<t})}$ | $(0, +\infty)$ |
 | $A_t$ | Advantage function | Measures how much better a certain action is compared to the baseline | $(-\infty, +\infty)$ |
 | $\text{KL}_t$ | KL divergence | Prevents the policy from drifting too far from the reference model | $[0, +\infty)$ |
 
@@ -1085,7 +1087,7 @@ $$\mathcal{L}_{PPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \text{clip}(r_t, 1-\v
 
 Where:
 - **Policy term**: $f(r_t) = \min(r_t, \text{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (clips probability ratio to prevent overly aggressive updates)
-- **Advantage term**: $g(A_t) = R - V(s)$ (estimates value function through Critic network)
+- **Advantage term**: $A_t$ is typically estimated by the Critic network, and GAE can also be used
 - **Regularization term**: $h(\text{KL}_t) = \beta \cdot \mathbb{E}[\text{KL}]$ (global KL divergence constraint)
 
 Compared to DPO,
@@ -1126,8 +1128,8 @@ In early 2025, as DeepSeek-R1 exploded in popularity, GRPO from the DeepSeekMath
 $$\mathcal{L}_{GRPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t) - \beta \cdot \text{KL}_t\right]$$
 
 Where:
-- **Policy term**: $f(r_t) = \min(r_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (uses symmetric clip on the probability ratio)
-- **Advantage term**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (intra-group normalization, eliminating the Critic network)
+- **Policy term**: $f(r_t, A_t) = \min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ (clips the probability ratio and advantage term together)
+- **Advantage term**: $g(A_{i,j}) = \frac{R_{i,j} - \mu_i}{\sigma_i + \epsilon}$ (intra-group normalization, eliminating the Critic network)
 - **Regularization term**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token-level KL divergence constraint)
 
 For the same question, the model generates N responses and computes their respective rewards, then uses the intra-group average reward as a baseline. Responses above the baseline are encouraged, and responses below the baseline are suppressed, thus no additional critic network needs to be trained.
@@ -1160,11 +1162,11 @@ CISPO does not focus on redesigning the group baseline. Instead, it uses a very
 
 **CISPO Loss**:
 
-$$\mathcal{L}_{CISPO} = -\mathbb{E}\left[\min(r_t, \varepsilon_{max}) \cdot A_t \cdot \log \pi_\theta(a_t|s) - \beta \cdot \text{KL}_t\right]$$
+$$\mathcal{L}_{CISPO} = -\mathbb{E}\left[\min(r_t, \varepsilon_{\mathrm{high}}) \cdot A_t \cdot \log \pi_\theta(a_t|s) - \beta \cdot \text{KL}_t\right]$$
 
 Where:
-- **Policy term**: $f(r_t) = \min(r_t, \varepsilon_{max}) \cdot \log \pi_\theta(a_t|s)$ (ratio serves only as a clipped weight)
-- **Advantage term**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (can directly reuse GRPO's intra-group relative advantage)
+- **Policy term**: $f(r_t) = \min(r_t, \varepsilon_{\mathrm{high}}) \cdot \log \pi_\theta(a_t|s)$ (ratio serves only as a clipped weight)
+- **Advantage term**: $g(A_{i,j}) = \frac{R_{i,j} - \mu_i}{\sigma_i + \epsilon}$ (can directly reuse GRPO's intra-group relative advantage)
 - **Regularization term**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token-level KL divergence constraint)
 
 CISPO, building on GRPO, rewrites the policy term that was easily clipped into a constant into the form "clipped weight × log probability". This way, even if the ratio is truncated, the gradient path is not truncated along with it. Therefore, CISPO can be directly viewed as a loss variant of GRPO to implement, rather than maintaining a separate standalone script. No separate experiment is listed here. One only needs to set `loss_type` to `cispo` in `train_grpo.py`; the rest of the training process still follows GRPO's group sampling, reward computation, and advantage construction logic.
@@ -1269,9 +1271,9 @@ Returning to the "**unified framework**", the table below summarizes how differe
 | Algorithm | Policy term $f(r_t)$ | Advantage term $g(A_t)$ | Regularization term $h(\text{KL}_t)$ | Number of training models |
 |-----------|---------------------|------------------------|-------------------------------------|--------------------------|
 | **DPO** | $\log r_w - \log r_l$ | No explicit advantage term | Implicit in $\beta$ | 1 (2 participate in forward) | 
-| **PPO** | $\min(r, \text{clip}(r))$ | $R - V(s)$ | $\beta \cdot \mathbb{E}[\text{KL}]$ | 2 | 
-| **GRPO** | $\min(r, \text{clip}(r))$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 |
-| **CISPO** | $\mathrm{clip}(r, 0, \varepsilon_{max}) \cdot A_t \cdot \log \pi_\theta$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 | 
+| **PPO** | $\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ | $A_t$ (typically estimated by the Critic, or computed with GAE) | $\beta \cdot \mathbb{E}[\text{KL}]$ | 2 |
+| **GRPO** | $\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ | $A_{i,j}=\frac{R_{i,j}-\mu_i}{\sigma_i+\epsilon}$ | $\beta \cdot \text{KL}_t$ | 1 |
+| **CISPO** | $\mathrm{clip}(r, 0, \varepsilon_{\mathrm{high}}) \cdot A_t \cdot \log \pi_\theta$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 | 
 
 **To put it plainly, these RL algorithms are not isolated from one another. From a unified optimization perspective, they are natural variants formed by making different design trade-offs on the same objective function, presenting a beautifully self-consistent unity.**
 

trainer/opd_utils.py

@@ -0,0 +1,148 @@
+"""Core tensor utilities for Generalized Knowledge Distillation (GKD)."""
+
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def build_completion_mask(
+    completion_ids: Tensor,
+    base_mask: Optional[Tensor],
+    eos_token_id: Optional[int],
+    pad_token_id: Optional[int] = None,
+) -> Tensor:
+    """Keep valid completion tokens up to and including the first EOS.
+
+    MiniMind's native ``generate`` pads finished rows with repeated EOS tokens
+    until every row is complete. Those repeated tokens are valid model inputs,
+    but they must not contribute to the distillation objective.
+    """
+    if completion_ids.ndim != 2:
+        raise ValueError(f"completion_ids must be 2D, got {tuple(completion_ids.shape)}")
+    if base_mask is None:
+        mask = torch.ones_like(completion_ids, dtype=torch.bool)
+    else:
+        if base_mask.shape != completion_ids.shape:
+            raise ValueError(
+                f"base_mask shape {tuple(base_mask.shape)} does not match "
+                f"completion_ids shape {tuple(completion_ids.shape)}"
+            )
+        mask = base_mask.bool().clone()
+
+    # Some tokenizers intentionally use EOS as PAD. In that case the first EOS
+    # still belongs to the completion and is handled by the cumulative mask.
+    if pad_token_id is not None and pad_token_id != eos_token_id:
+        mask &= completion_ids.ne(pad_token_id)
+
+    if eos_token_id is not None:
+        eos_hits = completion_ids.eq(eos_token_id) & mask
+        eos_seen_before = eos_hits.long().cumsum(dim=1) - eos_hits.long()
+        mask &= eos_seen_before.eq(0)
+    return mask
+
+
+def _masked_mean(values: Tensor, mask: Tensor) -> Tensor:
+    mask = mask.to(device=values.device, dtype=values.dtype)
+    return (values * mask).sum() / mask.sum().clamp_min(1.0)
+
+
+def generalized_jsd_loss(
+    student_logits: Tensor,
+    teacher_logits: Tensor,
+    completion_mask: Tensor,
+    beta: float = 0.5,
+    temperature: float = 1.0,
+) -> Tuple[Tensor, Dict[str, Tensor]]:
+    """Compute the exact full-vocabulary generalized JSD from GKD.
+
+    With teacher distribution ``p`` and student distribution ``q``, the
+    mixture is ``m = beta * p + (1 - beta) * q`` and the objective is
+    ``beta * KL(p || m) + (1 - beta) * KL(q || m)``. Its limiting cases are
+    forward KL at ``beta=0`` and reverse KL at ``beta=1``.
+
+    The teacher is always treated as a fixed target. Only completion tokens
+    selected by ``completion_mask`` contribute to the token-mean loss.
+    """
+    if student_logits.ndim != 3 or teacher_logits.ndim != 3:
+        raise ValueError("student_logits and teacher_logits must be [batch, time, vocab]")
+    if student_logits.shape != teacher_logits.shape:
+        raise ValueError(
+            "student and teacher logits must have identical shapes, got "
+            f"{tuple(student_logits.shape)} and {tuple(teacher_logits.shape)}"
+        )
+    if completion_mask.shape != student_logits.shape[:2]:
+        raise ValueError(
+            f"completion_mask shape {tuple(completion_mask.shape)} does not match "
+            f"logits shape {tuple(student_logits.shape[:2])}"
+        )
+    if not 0.0 <= beta <= 1.0:
+        raise ValueError(f"beta must be in [0, 1], got {beta}")
+    if temperature <= 0:
+        raise ValueError(f"temperature must be positive, got {temperature}")
+
+    # Float32 log-softmax keeps the divergence stable under fp16/bf16 model
+    # autocast, while preserving gradients back to the student parameters.
+    student_logps = F.log_softmax(student_logits.float() / temperature, dim=-1)
+    teacher_logps = F.log_softmax(
+        teacher_logits.detach().float() / temperature, dim=-1
+    )
+    student_probs = student_logps.exp()
+    teacher_probs = teacher_logps.exp()
+
+    forward_kl = (teacher_probs * (teacher_logps - student_logps)).sum(dim=-1)
+    reverse_kl = (student_probs * (student_logps - teacher_logps)).sum(dim=-1)
+
+    if beta == 0.0:
+        per_token_loss = forward_kl
+    elif beta == 1.0:
+        per_token_loss = reverse_kl
+    else:
+        beta_tensor = student_logps.new_tensor(beta)
+        mixture_logps = torch.logsumexp(
+            torch.stack(
+                (
+                    student_logps + torch.log1p(-beta_tensor),
+                    teacher_logps + torch.log(beta_tensor),
+                )
+            ),
+            dim=0,
+        )
+        teacher_to_mixture = (
+            teacher_probs * (teacher_logps - mixture_logps)
+        ).sum(dim=-1)
+        student_to_mixture = (
+            student_probs * (student_logps - mixture_logps)
+        ).sum(dim=-1)
+        per_token_loss = (
+            beta_tensor * teacher_to_mixture
+            + (1.0 - beta_tensor) * student_to_mixture
+        )
+
+    mask = completion_mask.bool()
+    loss = _masked_mean(per_token_loss, mask)
+    top1_agreement = _masked_mean(
+        student_logps.detach()
+        .argmax(dim=-1)
+        .eq(teacher_logps.argmax(dim=-1))
+        .float(),
+        mask,
+    )
+    metrics = {
+        "divergence": loss.detach(),
+        "forward_kl": _masked_mean(forward_kl.detach(), mask),
+        "reverse_kl": _masked_mean(reverse_kl.detach(), mask),
+        "top1_agreement": top1_agreement.detach(),
+        "valid_tokens": mask.sum().detach(),
+    }
+    return loss, metrics
+
+
+def use_on_policy_batch(lmbda: float, random_value: float) -> bool:
+    """Return the Algorithm 1 branch for a pre-sampled value in ``[0, 1)``."""
+    if not 0.0 <= lmbda <= 1.0:
+        raise ValueError(f"lmbda must be in [0, 1], got {lmbda}")
+    if not 0.0 <= random_value < 1.0:
+        raise ValueError(f"random_value must be in [0, 1), got {random_value}")
+    return random_value < lmbda

trainer/rollout_engine.py

@@ -52,7 +52,10 @@ class RolloutEngine(ABC):
     tokenizer = None
     
     @abstractmethod
-    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int,
+                max_new_tokens: int, temperature: float = 0.8,
+                top_p: Optional[float] = None, top_k: Optional[int] = None,
+                calculate_logps: bool = True) -> RolloutResult:
         pass
     
     @abstractmethod
@@ -68,9 +71,15 @@ class TorchRolloutEngine(RolloutEngine):
         self.device = device
         self.autocast_ctx = autocast_ctx
     
-    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int,
+                max_new_tokens: int, temperature: float = 0.8,
+                top_p: Optional[float] = None, top_k: Optional[int] = None,
+                calculate_logps: bool = True) -> RolloutResult:
         model = self.policy_model.module if isinstance(self.policy_model, DistributedDataParallel) else self.policy_model
         ctx = self.autocast_ctx if self.autocast_ctx else nullcontext()
+        sampling_kwargs = {}
+        if top_p is not None: sampling_kwargs['top_p'] = top_p
+        if top_k is not None: sampling_kwargs['top_k'] = top_k
         with torch.no_grad(), ctx:
             output_ids = model.generate(
                 input_ids=prompt_ids.repeat_interleave(num_generations, dim=0),
@@ -81,11 +90,20 @@ class TorchRolloutEngine(RolloutEngine):
                 num_return_sequences=1,
                 pad_token_id=self.tokenizer.pad_token_id,
                 eos_token_id=self.tokenizer.eos_token_id,
+                **sampling_kwargs,
             ).clone()  # [B*num_gen, P+R]
             prompt_len = prompt_ids.size(1)
             completion_ids = output_ids[:, prompt_len:]  # [B*num_gen, R]
             full_mask = (output_ids != self.tokenizer.pad_token_id).long()
-            per_token_logps = compute_per_token_logps(self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask)
+            if calculate_logps:
+                per_token_logps = compute_per_token_logps(
+                    self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask
+                )
+            else:
+                per_token_logps = torch.zeros(
+                    output_ids.size(0), completion_ids.size(1),
+                    device=output_ids.device, dtype=torch.float32
+                )
         completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
         return RolloutResult(output_ids, completion_ids, per_token_logps, completions,
                              prompt_ids.new_full((output_ids.size(0),), prompt_len),
@@ -104,7 +122,10 @@ class SGLangRolloutEngine(RolloutEngine):
         self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
         self.http = requests
     
-    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int,
+                max_new_tokens: int, temperature: float = 0.8,
+                top_p: Optional[float] = None, top_k: Optional[int] = None,
+                calculate_logps: bool = True) -> RolloutResult:
         # 去除左侧 padding tokens，只保留有效 token
         input_ids_list = []
         for ids, mask in zip(prompt_ids, attention_mask):
@@ -119,8 +140,10 @@ class SGLangRolloutEngine(RolloutEngine):
                 "max_new_tokens": max_new_tokens,
                 "stop_token_ids": [self.tokenizer.eos_token_id] if self.tokenizer.eos_token_id else [],
             },
-            "return_logprob": True,
+            "return_logprob": calculate_logps,
         }
+        if top_p is not None: payload['sampling_params']['top_p'] = top_p
+        if top_k is not None: payload['sampling_params']['top_k'] = top_k if top_k > 0 else -1
         
         resp = self.http.post(f"{self.base_url}/generate", json=payload, timeout=self.timeout)
         resp.raise_for_status()

trainer/train_opd.py

@@ -0,0 +1,755 @@
+import os
+import sys
+
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import argparse
+import gc
+import json
+import math
+import random
+import warnings
+from contextlib import nullcontext
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from datasets import load_dataset
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+from transformers import AutoTokenizer
+
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from trainer.opd_utils import (
+    build_completion_mask,
+    generalized_jsd_loss,
+    use_on_policy_batch,
+)
+from trainer.rollout_engine import create_rollout_engine
+from trainer.trainer_utils import (
+    Logger,
+    SkipBatchSampler,
+    init_distributed_mode,
+    is_main_process,
+    lm_checkpoint,
+    setup_seed,
+)
+
+warnings.filterwarnings("ignore")
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+class OPDDataset(Dataset):
+    """Raw conversational samples used by both GKD data branches."""
+
+    def __init__(self, data_path):
+        self.samples = load_dataset("json", data_files=str(data_path), split="train")
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        conversations = sample.get("conversations", sample.get("messages"))
+        if not isinstance(conversations, list) or not conversations:
+            raise ValueError(
+                "OPD data requires a non-empty 'conversations' or 'messages' list"
+            )
+        return {"conversations": [dict(message) for message in conversations]}
+
+
+def collate_conversations(batch):
+    return {"conversations": [item["conversations"] for item in batch]}
+
+
+def resolve_device(device):
+    if device != "auto":
+        return device
+    if torch.cuda.is_available():
+        return "cuda:0"
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def resolve_path(value):
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path.resolve()
+    return (REPO_ROOT / path).resolve()
+
+
+def validate_args(parsed_args):
+    if not 0.0 <= parsed_args.lmbda <= 1.0:
+        raise ValueError("--lmbda must be in [0, 1]")
+    if not 0.0 <= parsed_args.beta <= 1.0:
+        raise ValueError("--beta must be in [0, 1]")
+    if parsed_args.temperature <= 0 or parsed_args.distill_temperature <= 0:
+        raise ValueError("sampling and distillation temperatures must be positive")
+    if not 0.0 < parsed_args.rollout_top_p <= 1.0:
+        raise ValueError("--rollout_top_p must be in (0, 1]")
+    if parsed_args.rollout_top_k < 0:
+        raise ValueError("--rollout_top_k must be zero (disabled) or positive")
+    if parsed_args.num_generations < 1 or parsed_args.max_gen_len < 1:
+        raise ValueError("--num_generations and --max_gen_len must be positive")
+    if parsed_args.max_seq_len < 1 or parsed_args.accumulation_steps < 1:
+        raise ValueError("--max_seq_len and --accumulation_steps must be positive")
+    if parsed_args.max_train_steps < 0:
+        raise ValueError("--max_train_steps must be zero or positive")
+    if parsed_args.log_interval < 1 or parsed_args.save_interval < 1:
+        raise ValueError("log/save intervals must be positive")
+    if parsed_args.rollout_top_p < 1.0 or parsed_args.rollout_top_k > 0:
+        Logger(
+            "[OPD WARNING] top-p/top-k truncation changes the behavior policy; "
+            "top_p=1 and top_k=0 match the GKD baseline"
+        )
+
+
+def _torch_load(path):
+    try:
+        return torch.load(path, map_location="cpu", weights_only=True)
+    except TypeError:
+        return torch.load(path, map_location="cpu")
+
+
+def _load_transformers_state_dict(model_dir):
+    single_safetensors = model_dir / "model.safetensors"
+    safetensors_index = model_dir / "model.safetensors.index.json"
+    pytorch_bin = model_dir / "pytorch_model.bin"
+    if single_safetensors.exists() or safetensors_index.exists():
+        try:
+            from safetensors.torch import load_file
+        except ImportError as exc:
+            raise ImportError(
+                "Loading a Transformers safetensors model requires safetensors"
+            ) from exc
+    if single_safetensors.exists():
+        return load_file(str(single_safetensors), device="cpu")
+    if safetensors_index.exists():
+        index = json.loads(safetensors_index.read_text(encoding="utf-8"))
+        state_dict = {}
+        for shard_name in sorted(set(index["weight_map"].values())):
+            state_dict.update(load_file(str(model_dir / shard_name), device="cpu"))
+        return state_dict
+    if pytorch_bin.exists():
+        return _torch_load(pytorch_bin)
+    raise FileNotFoundError(
+        f"No model.safetensors, model.safetensors.index.json, or pytorch_model.bin in {model_dir}"
+    )
+
+
+def resolve_model_source(source, config):
+    direct = Path(source).expanduser()
+    if not direct.is_absolute():
+        direct = REPO_ROOT / direct
+    if direct.exists():
+        return direct.resolve()
+    moe_suffix = "_moe" if config.use_moe else ""
+    weight_path = REPO_ROOT / "out" / f"{source}_{config.hidden_size}{moe_suffix}.pth"
+    if weight_path.exists():
+        return weight_path.resolve()
+    raise FileNotFoundError(
+        f"Model source {source!r} was not found directly or as {weight_path}"
+    )
+
+
+def load_state_dict(source, config):
+    source = resolve_model_source(source, config)
+    state_dict = (
+        _load_transformers_state_dict(source) if source.is_dir() else _torch_load(source)
+    )
+    if isinstance(state_dict, dict) and isinstance(state_dict.get("model"), dict):
+        state_dict = state_dict["model"]
+    if not isinstance(state_dict, dict) or not state_dict:
+        raise ValueError(f"Invalid or empty state dict from {source}")
+    for prefix in ("module.", "_orig_mod."):
+        if all(key.startswith(prefix) for key in state_dict):
+            state_dict = {key[len(prefix):]: value for key, value in state_dict.items()}
+    if "lm_head.weight" not in state_dict and "model.embed_tokens.weight" in state_dict:
+        state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"]
+    if "model.embed_tokens.weight" not in state_dict and "lm_head.weight" in state_dict:
+        state_dict["model.embed_tokens.weight"] = state_dict["lm_head.weight"]
+    return state_dict, source
+
+
+def validate_directory_config(source, config, label):
+    source_path = Path(source).expanduser()
+    if not source_path.is_absolute():
+        source_path = REPO_ROOT / source_path
+    config_path = source_path / "config.json" if source_path.is_dir() else None
+    if not config_path or not config_path.exists():
+        return
+    saved = json.loads(config_path.read_text(encoding="utf-8"))
+    expected = {
+        "hidden_size": config.hidden_size,
+        "num_hidden_layers": config.num_hidden_layers,
+        "vocab_size": config.vocab_size,
+        "num_attention_heads": config.num_attention_heads,
+        "num_key_value_heads": config.num_key_value_heads,
+    }
+    mismatches = {
+        key: (saved.get(key), value)
+        for key, value in expected.items()
+        if saved.get(key) is not None and saved.get(key) != value
+    }
+    if mismatches:
+        raise ValueError(f"{label} config does not match CLI architecture: {mismatches}")
+
+
+def load_native_model(config, source, device, label, model_dtype=None):
+    validate_directory_config(source, config, label)
+    state_dict, resolved_source = load_state_dict(source, config)
+    loaded_model = MiniMindForCausalLM(config)
+    loaded_model.load_state_dict(state_dict, strict=True)
+    del state_dict
+    gc.collect()
+    if model_dtype is None:
+        loaded_model = loaded_model.to(device)
+    else:
+        loaded_model = loaded_model.to(device=device, dtype=model_dtype)
+    Logger(f"Loaded {label} strictly from {resolved_source}")
+    return loaded_model
+
+
+def load_tokenizer(tokenizer_path, vocab_size):
+    tokenizer = AutoTokenizer.from_pretrained(
+        str(resolve_path(tokenizer_path)), trust_remote_code=True
+    )
+    if len(tokenizer) != vocab_size:
+        raise ValueError(
+            f"Tokenizer size {len(tokenizer)} does not match model vocab {vocab_size}"
+        )
+    if tokenizer.pad_token_id is None or tokenizer.eos_token_id is None:
+        raise ValueError("OPD requires tokenizer.pad_token_id and tokenizer.eos_token_id")
+    tokenizer.padding_side = "left"
+    tokenizer.truncation_side = "left"
+    return tokenizer
+
+
+def _parse_json_field(value):
+    if not isinstance(value, str):
+        return value
+    try:
+        return json.loads(value)
+    except json.JSONDecodeError:
+        return value
+
+
+def render_sample(tokenizer, conversations):
+    """Render one prompt and its optional fixed completion."""
+    messages = []
+    tools = None
+    for raw_message in conversations:
+        message = {key: value for key, value in dict(raw_message).items() if value is not None}
+        if message.get("tools"):
+            parsed_tools = _parse_json_field(message["tools"])
+            if isinstance(parsed_tools, list):
+                tools = parsed_tools
+        if message.get("tool_calls"):
+            message["tool_calls"] = _parse_json_field(message["tool_calls"])
+        messages.append(message)
+
+    has_completion = messages[-1].get("role") == "assistant" and any(
+        messages[-1].get(field) for field in ("content", "reasoning_content", "tool_calls")
+    )
+    prompt_messages = messages[:-1] if messages[-1].get("role") == "assistant" else messages
+    template_kwargs = {"tools": tools} if tools else {}
+    prompt_text = tokenizer.apply_chat_template(
+        prompt_messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        **template_kwargs,
+    )
+    full_text = None
+    if has_completion:
+        full_text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+            **template_kwargs,
+        )
+    return prompt_text, full_text
+
+
+def tokenize_prompts(tokenizer, prompt_texts, max_seq_len, device):
+    encoded = tokenizer(
+        prompt_texts,
+        add_special_tokens=False,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=max_seq_len,
+        return_token_type_ids=False,
+    )
+    return {key: value.to(device) for key, value in encoded.items()}
+
+
+def build_off_policy_batch(
+    tokenizer, prompt_texts, full_texts, max_seq_len, max_gen_len, device
+):
+    """Tokenize fixed prompt-completion pairs and mask all prompt tokens."""
+    sequences = []
+    label_masks = []
+    for prompt_text, full_text in zip(prompt_texts, full_texts):
+        if full_text is None:
+            raise ValueError(
+                "Off-policy GKD selected a sample without a fixed assistant completion; "
+                "use --lmbda 1 for prompt-only data"
+            )
+        prompt_ids = tokenizer(prompt_text, add_special_tokens=False).input_ids
+        full_ids = tokenizer(full_text, add_special_tokens=False).input_ids
+        if full_ids[:len(prompt_ids)] != prompt_ids:
+            raise ValueError(
+                "The rendered prompt is not a token prefix of the fixed conversation; "
+                "the tokenizer/chat template is incompatible with prompt-completion GKD"
+            )
+        completion_ids = full_ids[len(prompt_ids):]
+        prompt_ids = prompt_ids[-max_seq_len:]
+        completion_ids = completion_ids[:max_gen_len]
+        if not completion_ids:
+            raise ValueError("A fixed completion became empty after tokenization/truncation")
+        sequences.append(prompt_ids + completion_ids)
+        label_masks.append([False] * len(prompt_ids) + [True] * len(completion_ids))
+
+    max_length = max(len(sequence) for sequence in sequences)
+    input_ids = []
+    attention_mask = []
+    labels = []
+    for sequence, label_mask in zip(sequences, label_masks):
+        padding = max_length - len(sequence)
+        input_ids.append(sequence + [tokenizer.pad_token_id] * padding)
+        attention_mask.append([1] * len(sequence) + [0] * padding)
+        labels.append(
+            [token if keep else -100 for token, keep in zip(sequence, label_mask)]
+            + [-100] * padding
+        )
+    return (
+        torch.tensor(input_ids, dtype=torch.long, device=device),
+        torch.tensor(attention_mask, dtype=torch.long, device=device),
+        torch.tensor(labels, dtype=torch.long, device=device),
+    )
+
+
+def distributed_mean(value):
+    result = value.detach().float().clone()
+    if dist.is_initialized():
+        dist.all_reduce(result, op=dist.ReduceOp.SUM)
+        result /= dist.get_world_size()
+    return result.item()
+
+
+def choose_on_policy(lmbda, device):
+    if dist.is_initialized():
+        draw = torch.zeros((), device=device)
+        if dist.get_rank() == 0:
+            draw.fill_(random.random())
+        dist.broadcast(draw, src=0)
+        random_value = draw.item()
+    else:
+        random_value = random.random()
+    return use_on_policy_batch(lmbda, random_value)
+
+
+def save_student(lm_config, epoch, step, wandb):
+    if not is_main_process():
+        return
+    model.eval()
+    raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+    raw_model = getattr(raw_model, "_orig_mod", raw_model)
+    moe_suffix = "_moe" if lm_config.use_moe else ""
+    output_path = resolve_path(args.save_dir) / (
+        f"{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth"
+    )
+    state_dict = raw_model.state_dict()
+    torch.save({key: value.half().cpu() for key, value in state_dict.items()}, output_path)
+    lm_checkpoint(
+        lm_config,
+        weight=args.save_weight,
+        model=model,
+        optimizer=optimizer,
+        epoch=epoch,
+        step=step,
+        wandb=wandb,
+        save_dir=str(resolve_path(args.checkpoint_dir)),
+        scheduler=scheduler,
+        scaler=scaler,
+    )
+    model.train()
+    del state_dict
+
+
+def gkd_train_epoch(
+    epoch, loader, iters, rollout_engine, teacher_model, lm_config,
+    start_step=0, wandb=None
+):
+    teacher_model.eval().requires_grad_(False)
+    model.train()
+    on_policy_batches = 0
+    processed_batches = 0
+
+    for step, batch in enumerate(loader, start=start_step + 1):
+        if args.max_train_steps > 0 and processed_batches >= args.max_train_steps:
+            break
+        processed_batches += 1
+        on_policy = choose_on_policy(args.lmbda, args.device)
+        on_policy_batches += int(on_policy)
+        rendered = [render_sample(tokenizer, item) for item in batch["conversations"]]
+        prompt_texts = [item[0] for item in rendered]
+        full_texts = [item[1] for item in rendered]
+
+        if on_policy:
+            prompt_inputs = tokenize_prompts(
+                tokenizer, prompt_texts, args.max_seq_len, args.device
+            )
+            model.eval()
+            rollout_result = rollout_engine.rollout(
+                prompt_ids=prompt_inputs["input_ids"],
+                attention_mask=prompt_inputs["attention_mask"],
+                num_generations=args.num_generations,
+                max_new_tokens=args.max_gen_len,
+                temperature=args.temperature,
+                top_p=args.rollout_top_p,
+                top_k=args.rollout_top_k,
+                calculate_logps=False,
+            )
+            model.train()
+            input_ids = rollout_result.output_ids.to(args.device)
+            completion_ids = rollout_result.completion_ids.to(args.device)
+            completion_mask = build_completion_mask(
+                completion_ids,
+                rollout_result.completion_mask.to(args.device),
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+            prompt_mask = prompt_inputs["attention_mask"].repeat_interleave(
+                args.num_generations, dim=0
+            )
+            attention_mask = torch.cat(
+                (prompt_mask, completion_mask.to(prompt_mask.dtype)), dim=1
+            )
+            with autocast_ctx:
+                student_output = model(
+                    input_ids,
+                    attention_mask=attention_mask,
+                    logits_to_keep=completion_ids.size(1) + 1,
+                )
+                student_logits = student_output.logits[:, :-1, :]
+            with torch.no_grad(), autocast_ctx:
+                teacher_logits = teacher_model(
+                    input_ids,
+                    attention_mask=attention_mask,
+                    logits_to_keep=completion_ids.size(1) + 1,
+                ).logits[:, :-1, :]
+            avg_completion_length = completion_mask.sum(dim=1).float().mean()
+            debug_completion = rollout_result.completions[0]
+        else:
+            input_ids, attention_mask, labels = build_off_policy_batch(
+                tokenizer,
+                prompt_texts,
+                full_texts,
+                args.max_seq_len,
+                args.max_gen_len,
+                args.device,
+            )
+            completion_mask = labels[:, 1:].ne(-100)
+            with autocast_ctx:
+                student_output = model(input_ids, attention_mask=attention_mask)
+                student_logits = student_output.logits[:, :-1, :]
+            with torch.no_grad(), autocast_ctx:
+                teacher_logits = teacher_model(
+                    input_ids, attention_mask=attention_mask
+                ).logits[:, :-1, :]
+            avg_completion_length = completion_mask.sum(dim=1).float().mean()
+            debug_completion = full_texts[0]
+
+        gkd_loss, gkd_metrics = generalized_jsd_loss(
+            student_logits,
+            teacher_logits,
+            completion_mask,
+            beta=args.beta,
+            temperature=args.distill_temperature,
+        )
+        aux_loss = student_output.aux_loss if lm_config.use_moe else gkd_loss.new_zeros(())
+        total_loss = gkd_loss + aux_loss
+        scaler.scale(total_loss / args.accumulation_steps).backward()
+
+        should_step = step % args.accumulation_steps == 0 or processed_batches == iters
+        if should_step:
+            scaler.unscale_(optimizer)
+            if args.grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+            scheduler.step()
+            optimizer.zero_grad(set_to_none=True)
+
+        if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
+            Logger(f"[OPD DEBUG] on_policy={on_policy}, prompt={prompt_texts[0]!r}")
+            Logger(f"[OPD DEBUG] completion={debug_completion!r}")
+
+        if step % args.log_interval == 0 or processed_batches == iters:
+            values = {
+                "loss/gkd": distributed_mean(gkd_loss),
+                "loss/aux": distributed_mean(aux_loss),
+                "loss/total": distributed_mean(total_loss),
+                "distill/forward_kl": distributed_mean(gkd_metrics["forward_kl"]),
+                "distill/reverse_kl": distributed_mean(gkd_metrics["reverse_kl"]),
+                "distill/top1_agreement": distributed_mean(gkd_metrics["top1_agreement"]),
+                "distill/valid_tokens": distributed_mean(gkd_metrics["valid_tokens"]),
+                "distill/on_policy": float(on_policy),
+                "distill/on_policy_fraction": on_policy_batches / processed_batches,
+                "rollout/avg_completion_length": distributed_mean(avg_completion_length),
+                "train/learning_rate": optimizer.param_groups[0]["lr"],
+            }
+            Logger(
+                f"Epoch:[{epoch + 1}/{args.epochs}]({step}/{start_step + iters}), "
+                f"GKD Loss: {values['loss/gkd']:.4f}, "
+                f"Forward KL: {values['distill/forward_kl']:.4f}, "
+                f"Reverse KL: {values['distill/reverse_kl']:.4f}, "
+                f"Top1 Agree: {values['distill/top1_agreement']:.4f}, "
+                f"On Policy: {int(on_policy)}, LR: {values['train/learning_rate']:.8f}"
+            )
+            if wandb and is_main_process():
+                wandb.log(values, step=epoch * iters + step)
+
+        if should_step and (step % args.save_interval == 0 or processed_batches == iters):
+            save_student(lm_config, epoch, step, wandb)
+
+        del input_ids, attention_mask, completion_mask
+        del student_output, student_logits, teacher_logits
+        del gkd_loss, gkd_metrics, aux_loss, total_loss
+        if on_policy:
+            del prompt_inputs, completion_ids, prompt_mask, rollout_result
+        else:
+            del labels
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="MiniMind Generalized/On-Policy Knowledge Distillation"
+    )
+    parser.add_argument("--data_path", type=str, default="dataset/rlaif.jsonl")
+    parser.add_argument("--student_model", type=str, default="full_sft")
+    parser.add_argument("--teacher_model", type=str, default="full_sft")
+    parser.add_argument("--tokenizer_path", type=str, default="model")
+    parser.add_argument("--save_dir", type=str, default="out")
+    parser.add_argument("--checkpoint_dir", type=str, default="checkpoints")
+    parser.add_argument("--save_weight", type=str, default="opd")
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=2)
+    parser.add_argument("--learning_rate", type=float, default=1e-6)
+    parser.add_argument("--device", type=str, default="auto")
+    parser.add_argument(
+        "--dtype", type=str, default="bfloat16",
+        choices=["float32", "float16", "bfloat16"]
+    )
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--accumulation_steps", type=int, default=1)
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--log_interval", type=int, default=1)
+    parser.add_argument("--save_interval", type=int, default=100)
+    parser.add_argument("--max_seq_len", type=int, default=768)
+    parser.add_argument("--max_gen_len", type=int, default=256)
+    parser.add_argument("--max_train_steps", type=int, default=0)
+    parser.add_argument(
+        "--lmbda", type=float, default=1.0,
+        help="fraction of student-generated batches; 1.0 is pure OPD"
+    )
+    parser.add_argument(
+        "--beta", type=float, default=0.5,
+        help="generalized JSD coefficient; 0=forward KL, 1=reverse KL"
+    )
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--distill_temperature", type=float, default=1.0)
+    parser.add_argument("--rollout_top_p", type=float, default=1.0)
+    parser.add_argument("--rollout_top_k", type=int, default=0)
+    parser.add_argument("--num_generations", type=int, default=1)
+    parser.add_argument("--student_hidden_size", type=int, default=768)
+    parser.add_argument("--student_num_layers", type=int, default=8)
+    parser.add_argument("--student_use_moe", type=int, default=0, choices=[0, 1])
+    parser.add_argument("--teacher_hidden_size", type=int, default=768)
+    parser.add_argument("--teacher_num_layers", type=int, default=8)
+    parser.add_argument("--teacher_use_moe", type=int, default=1, choices=[0, 1])
+    parser.add_argument("--from_resume", type=int, default=0, choices=[0, 1])
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-OPD")
+    parser.add_argument("--wandb_run_name", type=str, default="")
+    parser.add_argument(
+        "--wandb_mode", type=str, default="cloud",
+        choices=["cloud", "local", "disabled"]
+    )
+    parser.add_argument("--wandb_logdir", type=str, default="swanlog")
+    parser.add_argument("--use_compile", type=int, default=0, choices=[0, 1])
+    parser.add_argument("--debug_mode", action="store_true")
+    parser.add_argument("--debug_interval", type=int, default=20)
+    args = parser.parse_args()
+    validate_args(args)
+
+    local_rank = init_distributed_mode()
+    if dist.is_initialized():
+        args.device = f"cuda:{local_rank}"
+    else:
+        args.device = resolve_device(args.device)
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    setup_seed(args.seed + rank)
+    resolve_path(args.save_dir).mkdir(parents=True, exist_ok=True)
+    resolve_path(args.checkpoint_dir).mkdir(parents=True, exist_ok=True)
+
+    max_position_embeddings = max(32768, args.max_seq_len + args.max_gen_len)
+    student_config = MiniMindConfig(
+        hidden_size=args.student_hidden_size,
+        num_hidden_layers=args.student_num_layers,
+        use_moe=bool(args.student_use_moe),
+        max_position_embeddings=max_position_embeddings,
+    )
+    teacher_config = MiniMindConfig(
+        hidden_size=args.teacher_hidden_size,
+        num_hidden_layers=args.teacher_num_layers,
+        use_moe=bool(args.teacher_use_moe),
+        max_position_embeddings=max_position_embeddings,
+    )
+    if student_config.vocab_size != teacher_config.vocab_size:
+        raise ValueError("GKD requires student and teacher to share the same vocabulary")
+    tokenizer = load_tokenizer(args.tokenizer_path, student_config.vocab_size)
+
+    device_type = args.device.split(":", 1)[0]
+    mixed_dtype = {
+        "float32": torch.float32,
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+    }[args.dtype]
+    if device_type == "mps" and mixed_dtype == torch.bfloat16:
+        Logger("[OPD WARNING] MPS uses float16 for the frozen teacher")
+        mixed_dtype = torch.float16
+    teacher_dtype = mixed_dtype if device_type in {"cuda", "mps"} else torch.float32
+    autocast_ctx = (
+        torch.cuda.amp.autocast(dtype=mixed_dtype)
+        if device_type == "cuda" and mixed_dtype != torch.float32
+        else nullcontext()
+    )
+
+    model = load_native_model(
+        student_config, args.student_model, args.device, "student"
+    )
+    teacher_model = load_native_model(
+        teacher_config,
+        args.teacher_model,
+        args.device,
+        "teacher",
+        model_dtype=teacher_dtype,
+    ).eval().requires_grad_(False)
+    Logger(f"Student params: {sum(p.numel() for p in model.parameters()) / 1e6:.3f}M")
+    Logger(f"Teacher params: {sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f}M")
+
+    checkpoint = (
+        lm_checkpoint(
+            student_config,
+            weight=args.save_weight,
+            save_dir=str(resolve_path(args.checkpoint_dir)),
+        )
+        if args.from_resume == 1 else None
+    )
+    train_ds = OPDDataset(resolve_path(args.data_path))
+    train_sampler = (
+        DistributedSampler(train_ds, shuffle=True, seed=args.seed)
+        if dist.is_initialized() else None
+    )
+    count_loader = DataLoader(
+        train_ds,
+        batch_size=args.batch_size,
+        sampler=train_sampler,
+        collate_fn=collate_conversations,
+    )
+    full_iters = len(count_loader)
+    iters = min(full_iters, args.max_train_steps) if args.max_train_steps > 0 else full_iters
+    optimizer_steps_per_epoch = math.ceil(iters / args.accumulation_steps)
+    total_optimizer_steps = max(1, optimizer_steps_per_epoch * args.epochs)
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    scheduler = CosineAnnealingLR(
+        optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10
+    )
+    scaler = torch.cuda.amp.GradScaler(
+        enabled=(device_type == "cuda" and mixed_dtype == torch.float16)
+    )
+
+    start_epoch, start_step = 0, 0
+    if checkpoint:
+        model.load_state_dict(checkpoint["model"], strict=True)
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        if "scheduler" in checkpoint:
+            scheduler.load_state_dict(checkpoint["scheduler"])
+        if "scaler" in checkpoint:
+            scaler.load_state_dict(checkpoint["scaler"])
+        start_epoch = checkpoint["epoch"]
+        start_step = checkpoint.get("step", 0)
+        Logger(f"Resuming GKD from epoch={start_epoch}, step={start_step}")
+
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger("torch.compile enabled")
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    rollout_engine = create_rollout_engine(
+        engine_type="torch",
+        policy_model=model,
+        tokenizer=tokenizer,
+        device=args.device,
+        autocast_ctx=autocast_ctx,
+    )
+
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = checkpoint.get("wandb_id") if checkpoint else None
+        wandb.init(
+            project=args.wandb_project,
+            name=args.wandb_run_name or (
+                f"GKD-L{args.lmbda}-B{args.beta}-BS{args.batch_size}x"
+                f"{args.accumulation_steps}-LR{args.learning_rate}"
+            ),
+            config=vars(args),
+            mode=args.wandb_mode,
+            logdir=str(resolve_path(args.wandb_logdir)),
+            id=wandb_id if args.wandb_mode == "cloud" else None,
+            resume="must" if wandb_id and args.wandb_mode == "cloud" else None,
+        )
+
+    optimizer.zero_grad(set_to_none=True)
+    for epoch in range(start_epoch, args.epochs):
+        if train_sampler:
+            train_sampler.set_epoch(epoch)
+        setup_seed(args.seed + epoch + rank)
+        indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if epoch == start_epoch and start_step > 0 else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(
+            train_ds,
+            batch_sampler=batch_sampler,
+            num_workers=args.num_workers,
+            pin_memory=(device_type == "cuda"),
+            collate_fn=collate_conversations,
+        )
+        if skip > 0:
+            Logger(f"Epoch [{epoch + 1}/{args.epochs}]: skipping {skip} completed batches")
+        gkd_train_epoch(
+            epoch,
+            loader,
+            iters,
+            rollout_engine,
+            teacher_model,
+            student_config,
+            start_step=skip,
+            wandb=wandb,
+        )
+        start_step = 0
+
+    if wandb and is_main_process() and hasattr(wandb, "finish"):
+        wandb.finish()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()