[trainer, fsdp, worker, algo, cfg, recipe] feat: add reverse-KL on-policy distillation with FSDP teacher

examples/on_policy_distillation_trainer/run_dsr1_1.5b_reverse_kl_fsdp.sh

@@ -0,0 +1,156 @@
+#!/usr/bin/env bash
+# On-policy distillation | single-teacher | reverse-KL on student-top-K | FSDP teacher | NVIDIA GPUs
+
+# Reproduces the reverse-KL on-policy distillation experiments from
+# "Rethinking On-Policy Distillation of Large Language Models"
+# (https://arxiv.org/abs/2604.13016v1).
+#
+# Student: DeepSeek-R1-Distill-Qwen-1.5B
+# Teacher: JustRL-DeepSeek-1.5B
+# Dataset: DAPO-17k
+# Evaluation: AIME24 / AIME25 / AMC23
+# Teacher backend: frozen FSDP worker
+
+set -xeuo pipefail
+
+# ---- user-adjustable ----
+STUDENT_MODEL=${STUDENT_MODEL:-deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B}
+TEACHER_MODEL=${TEACHER_MODEL:-hbx/JustRL-DeepSeek-1.5B}
+
+# Single 8-GPU node. Student (actor + rollout + ref) and the frozen FSDP teacher
+# are co-located; each worker group is sized independently below.
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+# Teacher runs as its own FSDP worker group; total teacher GPUs = nnodes * gpus_per_node.
+TEACHER_NNODES=${TEACHER_NNODES:-1}
+TEACHER_NGPUS_PER_NODE=${TEACHER_NGPUS_PER_NODE:-8}
+
+# Reverse KL hyperparams. topk is the *student* top-K width (not the teacher top-K).
+distillation_topk=${DISTILLATION_TOPK:-16}
+teacher_chunk_size=${TEACHER_CHUNK_SIZE:-1024}
+
+train_batch_size=${TRAIN_BATCH_SIZE:-64}
+ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-64}
+max_prompt_length=${MAX_PROMPT_LENGTH:-512}
+max_response_length=${MAX_RESPONSE_LENGTH:-7680}
+ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-32768}
+
+actor_lr=${ACTOR_LR:-1e-6}
+
+rollout_tp=${ROLLOUT_TP:-2}
+rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.4}
+
+total_epochs=${TOTAL_EPOCHS:-1}
+save_freq=${SAVE_FREQ:-10}
+test_freq=${TEST_FREQ:-10}
+
+project_name=${PROJECT_NAME:-verl_distill_reverse_kl_fsdp}
+experiment_name=${EXPERIMENT_NAME:-dsr1_1.5b_from_justrl_1.5b_reverse_kl_topk_fsdp}
+# ---- end user-adjustable ----
+
+# Train on DAPO-17k, evaluate on AIME24 / AIME25 / AMC23. Point these at your own
+# parquet files (see examples/data_preprocess for dataset preparation scripts).
+train_data=${TRAIN_DATA:-$HOME/data/dapo/dapo_math_17k.parquet}
+aime24_test=${AIME24_TEST:-$HOME/data/math/aime24.parquet}
+aime25_test=${AIME25_TEST:-$HOME/data/math/aime25.parquet}
+amc23_test=${AMC23_TEST:-$HOME/data/math/amc23.parquet}
+
+train_files="['$train_data']"
+val_files="['$aime24_test','$aime25_test','$amc23_test']"
+
+max_num_tokens=$(( max_prompt_length + max_response_length + 1 ))
+
+########################### parameter arrays ###########################
+
+DATA=(
+    algorithm.adv_estimator=grpo
+    algorithm.use_kl_in_reward=False
+    data.train_files="$train_files"
+    data.val_files="$val_files"
+    data.train_batch_size=${train_batch_size}
+    data.max_prompt_length=${max_prompt_length}
+    data.max_response_length=${max_response_length}
+    data.filter_overlong_prompts=True
+    data.truncation='error'
+    data.shuffle=True
+)
+
+MODEL=(
+    actor_rollout_ref.model.path="$STUDENT_MODEL"
+    actor_rollout_ref.model.use_remove_padding=True
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+)
+
+ACTOR=(
+    actor_rollout_ref.actor.use_torch_compile=True
+    actor_rollout_ref.actor.optim.lr=${actor_lr}
+    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size}
+    actor_rollout_ref.actor.use_dynamic_bsz=True
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu}
+    actor_rollout_ref.actor.fsdp_config.param_offload=True
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.name=vllm
+    actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp}
+    actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util}
+    actor_rollout_ref.rollout.n=1
+    actor_rollout_ref.rollout.max_model_len=${max_num_tokens}
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu}
+    actor_rollout_ref.rollout.val_kwargs.n=16
+    actor_rollout_ref.rollout.val_kwargs.temperature=0.7
+    actor_rollout_ref.rollout.val_kwargs.top_p=0.95
+    actor_rollout_ref.rollout.val_kwargs.do_sample=True
+)
+
+TRAINER=(
+    trainer.balance_batch=True
+    trainer.logger='["console","wandb"]'
+    trainer.project_name=${project_name}
+    trainer.experiment_name=${experiment_name}
+    trainer.n_gpus_per_node=${NGPUS_PER_NODE}
+    trainer.nnodes=${NNODES}
+    trainer.val_before_train=False
+    trainer.save_freq=${save_freq}
+    trainer.test_freq=${test_freq}
+    trainer.total_epochs=${total_epochs}
+)
+
+# Reverse-KL OPD with an FSDP-served teacher: teacher_backend=fsdp hosts the teacher as
+# a frozen FSDP worker instead of a vLLM rollout. loss_mode=reverse_kl_topk uses topk as
+# the *student* top-K width (not the teacher top-K, as in forward_kl_topk).
+EXTRA=(
+    distillation.enabled=True
+    distillation.teacher_backend=fsdp
+    distillation.teacher_chunk_size=${teacher_chunk_size}
+    distillation.n_gpus_per_node=${TEACHER_NGPUS_PER_NODE}
+    distillation.nnodes=${TEACHER_NNODES}
+    distillation.teacher_key=ability
+    distillation.teacher_models.teacher_model.key=MATH
+    distillation.teacher_models.teacher_model.model_path="$TEACHER_MODEL"
+    # --- teacher FSDP engine ---
+    +distillation.teacher_fsdp_config.strategy=fsdp
+    +distillation.teacher_fsdp_config.dtype=bfloat16
+    +distillation.teacher_fsdp_config.param_offload=True
+    +distillation.teacher_fsdp_config.optimizer_offload=False
+    +distillation.teacher_fsdp_config.fsdp_size=-1
+    # --- loss ---
+    distillation.distillation_loss.loss_mode=reverse_kl_topk
+    distillation.distillation_loss.topk=${distillation_topk}
+    distillation.distillation_loss.use_task_rewards=False
+    distillation.distillation_loss.use_policy_gradient=False
+    distillation.distillation_loss.log_prob_min_clamp=-10.0
+)
+
+########################### launch ###########################
+python3 -m verl.trainer.main_ppo \
+    "${DATA[@]}" \
+    "${MODEL[@]}" \
+    "${ACTOR[@]}" \
+    "${ROLLOUT[@]}" \
+    "${TRAINER[@]}" \
+    "${EXTRA[@]}" \
+    "$@"

examples/on_policy_distillation_trainer/run_qwen3_reverse_kl_fsdp.sh

@@ -0,0 +1,136 @@
+#!/usr/bin/env bash
+# On-policy distillation | single-teacher | reverse-KL on student-top-K | FSDP teacher | NVIDIA GPUs
+
+set -xeuo pipefail
+
+# ---- user-adjustable ----
+STUDENT_MODEL=${STUDENT_MODEL:-Qwen/Qwen3-1.7B}
+TEACHER_MODEL=${TEACHER_MODEL:-Qwen/Qwen3-8B}
+
+NNODES=${NNODES:-1}
+NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
+
+# Teacher runs as its own FSDP worker group; total teacher GPUs = nnodes * gpus_per_node.
+TEACHER_NNODES=${TEACHER_NNODES:-1}
+TEACHER_NGPUS_PER_NODE=${TEACHER_NGPUS_PER_NODE:-4}
+
+# Reverse KL hyperparams.
+distillation_topk=${DISTILLATION_TOPK:-32}
+teacher_chunk_size=${TEACHER_CHUNK_SIZE:-1024}
+
+train_batch_size=${TRAIN_BATCH_SIZE:-32}
+ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE:-32}
+max_prompt_length=${MAX_PROMPT_LENGTH:-512}
+max_response_length=${MAX_RESPONSE_LENGTH:-1024}
+ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU:-12288}
+
+actor_lr=${ACTOR_LR:-1e-6}
+
+rollout_tp=${ROLLOUT_TP:-2}
+rollout_gpu_mem_util=${ROLLOUT_GPU_MEM_UTIL:-0.4}
+
+total_epochs=${TOTAL_EPOCHS:-1}
+save_freq=${SAVE_FREQ:-200}
+test_freq=${TEST_FREQ:-50}
+
+project_name=${PROJECT_NAME:-verl_distill_reverse_kl_fsdp}
+experiment_name=${EXPERIMENT_NAME:-qwen3_1.7b_from_qwen3_8b_reverse_kl_topk_fsdp}
+# ---- end user-adjustable ----
+
+gsm8k_train=$HOME/data/gsm8k/train.parquet
+gsm8k_test=$HOME/data/gsm8k/test.parquet
+
+train_files="['$gsm8k_train']"
+val_files="['$gsm8k_test']"
+
+max_num_tokens=$(( max_prompt_length + max_response_length + 1 ))
+
+########################### parameter arrays ###########################
+
+DATA=(
+    algorithm.adv_estimator=grpo
+    algorithm.use_kl_in_reward=False
+    data.train_files="$train_files"
+    data.val_files="$val_files"
+    data.train_batch_size=${train_batch_size}
+    data.max_prompt_length=${max_prompt_length}
+    data.max_response_length=${max_response_length}
+    data.filter_overlong_prompts=True
+    data.truncation='error'
+    data.shuffle=True
+)
+
+MODEL=(
+    actor_rollout_ref.model.path="$STUDENT_MODEL"
+    actor_rollout_ref.model.use_remove_padding=True
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+)
+
+ACTOR=(
+    actor_rollout_ref.actor.use_torch_compile=True
+    actor_rollout_ref.actor.optim.lr=${actor_lr}
+    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size}
+    actor_rollout_ref.actor.use_dynamic_bsz=True
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu}
+    actor_rollout_ref.actor.fsdp_config.param_offload=True
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.name=vllm
+    actor_rollout_ref.rollout.tensor_model_parallel_size=${rollout_tp}
+    actor_rollout_ref.rollout.gpu_memory_utilization=${rollout_gpu_mem_util}
+    actor_rollout_ref.rollout.n=1
+    actor_rollout_ref.rollout.max_model_len=${max_num_tokens}
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ppo_max_token_len_per_gpu}
+)
+
+TRAINER=(
+    trainer.balance_batch=True
+    trainer.logger='["console","wandb"]'
+    trainer.project_name=${project_name}
+    trainer.experiment_name=${experiment_name}
+    trainer.n_gpus_per_node=${NGPUS_PER_NODE}
+    trainer.nnodes=${NNODES}
+    trainer.val_before_train=False
+    trainer.save_freq=${save_freq}
+    trainer.test_freq=${test_freq}
+    trainer.total_epochs=${total_epochs}
+)
+
+# Reverse-KL OPD with an FSDP-served teacher: teacher_backend=fsdp hosts the teacher as
+# a frozen FSDP worker instead of a vLLM rollout. loss_mode=reverse_kl_topk uses topk as
+# the *student* top-K width (not the teacher top-K, as in forward_kl_topk).
+EXTRA=(
+    distillation.enabled=True
+    distillation.teacher_backend=fsdp
+    distillation.teacher_chunk_size=${teacher_chunk_size}
+    distillation.n_gpus_per_node=${TEACHER_NGPUS_PER_NODE}
+    distillation.nnodes=${TEACHER_NNODES}
+    distillation.teacher_key=data_source
+    distillation.teacher_models.teacher_model.key="openai/gsm8k"
+    distillation.teacher_models.teacher_model.model_path="$TEACHER_MODEL"
+    # --- teacher FSDP engine ---
+    +distillation.teacher_fsdp_config.strategy=fsdp
+    +distillation.teacher_fsdp_config.dtype=bfloat16
+    +distillation.teacher_fsdp_config.param_offload=True
+    +distillation.teacher_fsdp_config.optimizer_offload=False
+    +distillation.teacher_fsdp_config.fsdp_size=-1
+    # --- loss ---
+    distillation.distillation_loss.loss_mode=reverse_kl_topk
+    distillation.distillation_loss.topk=${distillation_topk}
+    distillation.distillation_loss.use_task_rewards=False
+    distillation.distillation_loss.use_policy_gradient=False
+    distillation.distillation_loss.log_prob_min_clamp=-10.0
+)
+
+########################### launch ###########################
+python3 -m verl.trainer.main_ppo \
+    "${DATA[@]}" \
+    "${MODEL[@]}" \
+    "${ACTOR[@]}" \
+    "${ROLLOUT[@]}" \
+    "${TRAINER[@]}" \
+    "${EXTRA[@]}" \
+    "$@"

tests/workers/teacher/test_chunked_gather_on_cpu.py

@@ -0,0 +1,96 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for :func:`verl.workers.teacher.utils.chunked_gather_logprobs`.
+
+The chunked gather is the inner kernel of the FSDP teacher worker — it computes
+``log_softmax(logits).gather(idx)`` without materializing the full
+``log_softmax`` tensor. These tests run on CPU with random tensors and do not
+depend on Ray, FSDP, or CUDA.
+"""
+
+import os
+
+os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from verl.workers.teacher.utils import chunked_gather_logprobs
+
+
+def _reference_gather(logits: torch.Tensor, topk_ids: torch.Tensor) -> torch.Tensor:
+    """One-shot reference: full log_softmax then gather. fp32 throughout."""
+    return torch.gather(F.log_softmax(logits.float(), dim=-1), dim=-1, index=topk_ids).to(logits.dtype)
+
+
+@pytest.mark.parametrize("chunk_size", [1, 7, 64, 1_000_000])
+def test_chunked_gather_matches_oneshot(chunk_size):
+    """chunk_size sweep — including a value that exceeds the row count — agrees
+    with the one-shot reference up to fp32 precision."""
+    torch.manual_seed(0)
+    bsz, seq, vocab, k = 2, 5, 13, 3
+    logits = torch.randn(bsz, seq, vocab)
+    topk_ids = torch.randint(0, vocab, (bsz, seq, k))
+
+    out = chunked_gather_logprobs(logits, topk_ids, chunk_size=chunk_size)
+    expected = _reference_gather(logits, topk_ids)
+    torch.testing.assert_close(out, expected, atol=1e-6, rtol=1e-6)
+
+
+def test_chunked_gather_preserves_shape_and_dtype():
+    """Output shape == leading shape of logits + (K,), dtype matches logits."""
+    logits = torch.randn(3, 4, 17, dtype=torch.float32)
+    topk_ids = torch.randint(0, 17, (3, 4, 5), dtype=torch.int64)
+    out = chunked_gather_logprobs(logits, topk_ids, chunk_size=8)
+    assert out.shape == (3, 4, 5)
+    assert out.dtype == torch.float32
+
+
+def test_chunked_gather_works_with_2d_input():
+    """2D inputs (rmpad layout): logits ``(N, V)`` and topk_ids ``(N, K)``."""
+    torch.manual_seed(1)
+    N, V, K = 11, 9, 4
+    logits = torch.randn(N, V)
+    topk_ids = torch.randint(0, V, (N, K))
+    out = chunked_gather_logprobs(logits, topk_ids, chunk_size=3)
+    expected = _reference_gather(logits, topk_ids)
+    torch.testing.assert_close(out, expected, atol=1e-6, rtol=1e-6)
+
+
+def test_chunked_gather_rejects_mismatched_leading_shape():
+    logits = torch.randn(2, 3, 5)
+    topk_ids = torch.randint(0, 5, (2, 4, 2))  # seq dim mismatch
+    with pytest.raises(ValueError, match="Leading shapes"):
+        chunked_gather_logprobs(logits, topk_ids)
+
+
+def test_chunked_gather_rejects_nonpositive_chunk_size():
+    logits = torch.randn(1, 1, 4)
+    topk_ids = torch.zeros(1, 1, 2, dtype=torch.int64)
+    with pytest.raises(ValueError, match="chunk_size must be positive"):
+        chunked_gather_logprobs(logits, topk_ids, chunk_size=0)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_chunked_gather_handles_extreme_logits_in_fp32(dtype):
+    """logsumexp in fp32 keeps ``log_softmax`` numerically stable even when raw
+    logits span large magnitudes (regression for naive ``log(softmax(x))``)."""
+    torch.manual_seed(2)
+    logits = torch.randn(8, 32, dtype=dtype) * 50.0  # magnitude ~ 50
+    topk_ids = torch.randint(0, 32, (8, 4))
+    out = chunked_gather_logprobs(logits, topk_ids, chunk_size=4)
+    assert torch.isfinite(out).all()
+    expected = _reference_gather(logits, topk_ids)
+    torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)