diff --git a/docs/en/advanced/miles_server_args.md b/docs/en/advanced/miles_server_args.md index e8d70db59fc..07112b8e47e 100644 --- a/docs/en/advanced/miles_server_args.md +++ b/docs/en/advanced/miles_server_args.md @@ -218,6 +218,7 @@ Arguments for reinforcement learning algorithms and loss calculation. | Argument | Description | Default | Options | Source | | :--- | :--- | :--- | :--- | :--- | | `--advantage-estimator` | Advantage estimator to use. | `"grpo"` | `grpo`, `gspo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`, `on_policy_distillation` | Miles Native | +| `--opd-reward-type` | Token-level reward for on-policy distillation. `logr` is `log(p_teacher/p_student)`; `k3` is the negative k3 KL estimator `1 + log(r) - r`. | `"logr"` | `logr`, `k3` | Miles Native | | `--loss-type` | Type of loss function to use. | `"policy_loss"` | `policy_loss`, `sft_loss`, `custom_loss` | Miles Native | | `--custom-loss-function-path` | Path to a custom loss calculation function (requires `--loss-type custom_loss`). [Ref](../get_started/customization.md#9-custom-loss-function---custom-loss-function-path) | `None` | Type: str | Miles Native | | `--critic-lr` | Learning rate for the Critic. Defaults to `--lr`. | `None` | Type: float | Miles Native | diff --git a/examples/on_policy_distillation/README.md b/examples/on_policy_distillation/README.md index a6b22c5b119..f59318f40de 100644 --- a/examples/on_policy_distillation/README.md +++ b/examples/on_policy_distillation/README.md @@ -35,6 +35,11 @@ PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ bash examples/on_policy_distillation/run-qwen3-8B-opd.sh ``` +The default token-level reward is `logr = log(p_teacher / p_student)`. Pass +`--opd-reward-type k3` to instead use the negative k3 KL estimator +`1 + log(r) - r`, where `r = p_teacher / p_student`. It has the same expected +negative forward-KL objective while using the k3 control variate. + # Preliminary Results Using Qwen3-8B-Base model sfted on part of the [OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) dataset, we performed on-policy distillation with a Qwen3-32B teacher on the remaining data. Evaluation on Math500 shows: @@ -56,4 +61,4 @@ The teacher runs on an independent SGLang server that miles treats as a reward m # References 1. https://thinkingmachines.ai/blog/on-policy-distillation/ 2. https://arxiv.org/abs/2306.13649 -3. https://arxiv.org/abs/2306.08543 \ No newline at end of file +3. https://arxiv.org/abs/2306.08543 diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index bf1eaf75a20..bdcceddc6b0 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -11,6 +11,7 @@ calculate_log_probs_and_entropy, compute_approx_kl, compute_gspo_kl, + compute_opd_reward, compute_opsm_mask, compute_policy_loss, get_advantages_and_returns_batch, @@ -380,7 +381,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) for t_log_prob, response_length in zip(teacher_log_probs, response_lengths, strict=False) ] advantages = [ - teacher_log_prob - student_log_prob + compute_opd_reward(student_log_prob, teacher_log_prob, getattr(args, "opd_reward_type", "logr")) for teacher_log_prob, student_log_prob in zip(teacher_log_probs, student_log_probs, strict=False) ] returns = advantages diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2d6e43e1384..a9af7d4fb7b 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -909,6 +909,17 @@ def add_algo_arguments(parser): ], default="grpo", ) + parser.add_argument( + "--opd-reward-type", + type=str, + choices=["logr", "k3"], + default="logr", + help=( + "Token-level reward for on_policy_distillation. logr uses log(p_teacher/p_student); " + "k3 uses its lower-variance negative-KL counterpart " + "1 + log(p_teacher/p_student) - p_teacher/p_student." + ), + ) parser.add_argument( "--disable-compute-advantages-and-returns", action="store_false", diff --git a/miles/utils/ppo_utils.py b/miles/utils/ppo_utils.py index 634c9d430d5..3b40e4bb094 100644 --- a/miles/utils/ppo_utils.py +++ b/miles/utils/ppo_utils.py @@ -51,6 +51,28 @@ def compute_approx_kl( return kl +def compute_opd_reward( + student_log_prob: torch.Tensor, + teacher_log_prob: torch.Tensor, + reward_type: str, +) -> torch.Tensor: + """Compute the token-level on-policy distillation reward. + + Let ``r = p_teacher / p_student``. ``logr`` is the original OPD reward + ``log(r)``. ``k3`` uses the negative of Schulman's non-negative k3 KL + estimator, ``-(r - 1 - log(r))``, because this value is used as a reward + and OPD should minimize ``KL(p_student || p_teacher)``. + """ + log_r = teacher_log_prob.float() - student_log_prob.float() + + if reward_type == "logr": + return log_r + if reward_type == "k3": + # 1 + log(r) - r, written with expm1 for accuracy near r == 1. + return log_r - torch.expm1(log_r) + raise ValueError(f"Unknown OPD reward type: {reward_type}") + + def compute_opsm_mask( args: Namespace, full_log_probs: list[torch.Tensor], diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py index aa2c35bd311..37ed8046cfc 100644 --- a/tests/fast/utils/test_arguments.py +++ b/tests/fast/utils/test_arguments.py @@ -12,6 +12,16 @@ REQUIRED_ARGS = ["--rollout-batch-size", "64"] +@pytest.mark.parametrize("reward_type", ["logr", "k3"]) +def test_opd_reward_type_is_parsed(reward_type: str) -> None: + with patch.object(sys, "argv", ["test", "--opd-reward-type", reward_type] + REQUIRED_ARGS): + parser = argparse.ArgumentParser() + get_miles_extra_args_provider()(parser) + args, _ = parser.parse_known_args() + + assert args.opd_reward_type == reward_type + + def make_class_with_add_arguments(): class MyFn: @classmethod diff --git a/tests/fast/utils/test_ppo_utils.py b/tests/fast/utils/test_ppo_utils.py new file mode 100644 index 00000000000..4b2e9ce9a11 --- /dev/null +++ b/tests/fast/utils/test_ppo_utils.py @@ -0,0 +1,38 @@ +import pytest +import torch + +from miles.utils.ppo_utils import compute_opd_reward + + +def test_compute_opd_reward_logr() -> None: + student = torch.tensor([-2.0, -1.0]) + teacher = torch.tensor([-1.5, -1.25]) + + actual = compute_opd_reward(student, teacher, "logr") + + torch.testing.assert_close(actual, torch.tensor([0.5, -0.25])) + + +def test_compute_opd_reward_k3_is_negative_kl_estimate() -> None: + student = torch.tensor([-2.0, -1.0]) + teacher = torch.tensor([-1.5, -1.25]) + log_r = teacher - student + expected = 1 + log_r - torch.exp(log_r) + + actual = compute_opd_reward(student, teacher, "k3") + + torch.testing.assert_close(actual, expected) + assert torch.all(actual <= 0) + + +def test_compute_opd_reward_k3_is_zero_when_distributions_match() -> None: + log_probs = torch.tensor([-10.0, -1.0, 0.0]) + + actual = compute_opd_reward(log_probs, log_probs, "k3") + + torch.testing.assert_close(actual, torch.zeros_like(log_probs)) + + +def test_compute_opd_reward_rejects_unknown_type() -> None: + with pytest.raises(ValueError, match="Unknown OPD reward type"): + compute_opd_reward(torch.tensor([0.0]), torch.tensor([0.0]), "unknown")