wuxibin89 · Mar 21, 2026 · Jan 17, 2026 · Jan 17, 2026 · Jan 15, 2026 · Jan 15, 2026
diff --git a/.github/workflows/gpu_unit_tests.yml b/.github/workflows/gpu_unit_tests.yml
@@ -117,6 +117,9 @@ jobs:
       - name: Testing LinearCrossEntropyTP Correctness, Computation Time and Memory Consumption
         run: |
           LOW_MEMORY=True torchrun --standalone --nnodes=1 --nproc-per-node=8 tests/utils/test_special_linear_cross_entropy_tp.py
+      - name: Testing Megatron KL Loss TP Correctness
+        run: |
+          torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/utils/test_special_megatron_kl_loss_tp.py
       - name: Testing FSDP2 actor functionality
         run: |
           torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/workers/actor/test_special_dp_actor.py

diff --git a/examples/on_policy_distillation_trainer/run_qwen3_vl_geo3k.sh b/examples/on_policy_distillation_trainer/run_qwen3_vl_geo3k.sh
@@ -0,0 +1,159 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+############################ Quick Config ############################
+
+ROLLOUT_NAME="vllm" # sglang or vllm
+
+FAMILY="Qwen"
+STUDENT_MODEL=Qwen3-VL-2B-Instruct
+TEACHER_MODEL=Qwen3-VL-4B-Instruct
+
+# USE_POLICY_GRADIENT=False
+# DISTILLATION_LOSS_MODE="k3"
+# DISTILLATION_LOSS_MODE="forward_kl_topk"
+# USE_FUSED_KERNELS=False
+
+USE_POLICY_GRADIENT=True
+DISTILLATION_LOSS_MODE="k1"
+USE_FUSED_KERNELS=True
+
+DISTILLATION_LOSS_MAX_CLAMP=10.0
+DISTILLATION_LOG_PROB_MIN_CLAMP=-10.0
+
+PROJECT_NAME='verl_on_policy_distillation_example_geo3k'
+
+MAX_PROMPT=1024
+MAX_RESPONSE_LENGTH=2048
+MAX_NUM_TOKENS=$(( MAX_PROMPT + MAX_RESPONSE_LENGTH + 1 ))
+TRAIN_PROMPT_BSZ=128
+STUDENT_MICRO_BATCH_SIZE_PER_GPU=1
+STUDENT_MAX_TOKEN_LEN_PER_GPU=$(( STUDENT_MICRO_BATCH_SIZE_PER_GPU * (MAX_PROMPT + MAX_RESPONSE_LENGTH) ))
+USE_DYNAMIC_BSZ=False
+
+STUDENT_WORLD_SIZE=4
+
+TEACHER_RESOURCE_POOL=True
+TEACHER_WORLD_SIZE=4
+
+SP=1
+
+EXP_NAME="fsdp/student-${STUDENT_MODEL}/teacher-${TEACHER_MODEL}/loss-${DISTILLATION_LOSS_MODE}/pg-${USE_POLICY_GRADIENT}"
+
+ENFORCE_EAGER=False # true for faster debugging
+
+############################ Paths ############################
+
+geo3k_train_path=$DATA_PATH/geo3k/train.parquet
+geo3k_test_path=$DATA_PATH/geo3k/test.parquet
+
+TRAIN_FILES="['$geo3k_train_path']"
+TEST_FILES="['$geo3k_test_path']"
+
+############################ Parameter Groups ############################
+
+DATA=(
+    data.train_files="$TRAIN_FILES"
+    data.val_files="$TEST_FILES"
+    data.max_prompt_length=$MAX_PROMPT
+    data.max_response_length=$MAX_RESPONSE_LENGTH
+    data.train_batch_size=$TRAIN_PROMPT_BSZ
+    data.filter_overlong_prompts=True
+    data.truncation='error'
+    data.shuffle=False
+    data.image_key=images
+)
+
+MODEL=(
+    actor_rollout_ref.model.path="${FAMILY}/${STUDENT_MODEL}"
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+    actor_rollout_ref.model.use_remove_padding=True
+    actor_rollout_ref.model.use_fused_kernels=$USE_FUSED_KERNELS
+    actor_rollout_ref.actor.use_torch_compile=True
+    actor_rollout_ref.rollout.enforce_eager=$ENFORCE_EAGER
+)
+
+DISTILLATION=(
+    distillation.enabled=True
+    distillation.num_workers=8
+    distillation.teacher_model.enable_resource_pool=$TEACHER_RESOURCE_POOL
+    distillation.teacher_model.n_gpus_per_node=$TEACHER_WORLD_SIZE
+    distillation.teacher_model.nnodes=1
+    distillation.teacher_model.model_path="${FAMILY}/${TEACHER_MODEL}"
+    distillation.teacher_model.inference.tensor_model_parallel_size=1
+    distillation.teacher_model.inference.name=$ROLLOUT_NAME
+    distillation.teacher_model.inference.gpu_memory_utilization=0.8
+    distillation.teacher_model.inference.enforce_eager=$ENFORCE_EAGER
+    distillation.teacher_model.inference.max_model_len=$MAX_NUM_TOKENS
+    distillation.teacher_model.inference.max_num_batched_tokens=$MAX_NUM_TOKENS
+    distillation.teacher_model.inference.max_num_seqs=$MAX_NUM_TOKENS
+    +distillation.teacher_model.inference.engine_kwargs.vllm.disable_mm_preprocessor_cache=True
+    distillation.distillation_loss.loss_mode=$DISTILLATION_LOSS_MODE
+    distillation.distillation_loss.topk=64
+    distillation.distillation_loss.use_task_rewards=False
+    distillation.distillation_loss.use_policy_gradient=$USE_POLICY_GRADIENT
+    distillation.distillation_loss.loss_max_clamp=$DISTILLATION_LOSS_MAX_CLAMP
+    distillation.distillation_loss.log_prob_min_clamp=$DISTILLATION_LOG_PROB_MIN_CLAMP
+)
+
+STUDENT=(
+    actor_rollout_ref.actor.optim.lr=1e-6
+    actor_rollout_ref.actor.ppo_mini_batch_size=$TRAIN_PROMPT_BSZ
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$STUDENT_MICRO_BATCH_SIZE_PER_GPU
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$STUDENT_MAX_TOKEN_LEN_PER_GPU
+    actor_rollout_ref.actor.use_dynamic_bsz=$USE_DYNAMIC_BSZ
+    actor_rollout_ref.actor.fsdp_config.param_offload=True
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
+    actor_rollout_ref.actor.ulysses_sequence_parallel_size=1
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=$STUDENT_MICRO_BATCH_SIZE_PER_GPU
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=$STUDENT_MAX_TOKEN_LEN_PER_GPU
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=$USE_DYNAMIC_BSZ
+    actor_rollout_ref.rollout.tensor_model_parallel_size=1
+    actor_rollout_ref.rollout.name=$ROLLOUT_NAME
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.8
+    actor_rollout_ref.rollout.calculate_log_probs=False
+    actor_rollout_ref.rollout.max_model_len=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.max_num_batched_tokens=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.max_num_seqs=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.n=1
+    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True
+)
+
+ALGORITHM=(
+    algorithm.adv_estimator=grpo
+    algorithm.use_kl_in_reward=False
+)
+
+TRAINER=(
+    trainer.logger='["console","wandb"]'
+    trainer.project_name=$PROJECT_NAME
+    trainer.experiment_name=$EXP_NAME
+    trainer.n_gpus_per_node=$STUDENT_WORLD_SIZE
+    trainer.nnodes=1
+    trainer.save_freq=200
+    trainer.test_freq=5
+    trainer.total_epochs=15
+    trainer.val_before_train=True
+    trainer.use_legacy_worker_impl=disable
+    trainer.resume_mode=disable
+    trainer.log_val_generations=5
+)
+
+
+
+############################ Launch ############################
+
+python3 -m verl.trainer.main_ppo \
+    --config-path=config \
+    --config-name='ppo_trainer.yaml' \
+    "${DATA[@]}" \
+    "${ALGORITHM[@]}" \
+    "${MODEL[@]}" \
+    "${DISTILLATION[@]}" \
+    "${ROLLOUT[@]}" \
+    "${STUDENT[@]}" \
+    "${TRAINER[@]}" \
+    "$@"
diff --git a/examples/on_policy_distillation_trainer/run_qwen_gsm8k.sh b/examples/on_policy_distillation_trainer/run_qwen_gsm8k.sh
@@ -0,0 +1,156 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+############################ Quick Config ############################
+
+ROLLOUT_NAME="vllm" # sglang or vllm
+
+FAMILY="Qwen"
+STUDENT_MODEL=Qwen2.5-0.5B
+TEACHER_MODEL=Qwen2.5-0.5B
+
+# USE_POLICY_GRADIENT=False
+# DISTILLATION_LOSS_MODE="k3"
+# DISTILLATION_LOSS_MODE="forward_kl_topk"
+# USE_FUSED_KERNELS=False
+
+USE_POLICY_GRADIENT=True
+DISTILLATION_LOSS_MODE="k1"
+USE_FUSED_KERNELS=False
+
+DISTILLATION_LOSS_MAX_CLAMP=10.0
+DISTILLATION_LOG_PROB_MIN_CLAMP=-10.0
+
+PROJECT_NAME='verl_on_policy_distillation_example_gsm8k'
+
+MAX_PROMPT=256
+MAX_RESPONSE_LENGTH=512
+MAX_NUM_TOKENS=$(( MAX_PROMPT + MAX_RESPONSE_LENGTH + 1 ))
+TRAIN_PROMPT_BSZ=128
+STUDENT_MICRO_BATCH_SIZE_PER_GPU=2
+STUDENT_MAX_TOKEN_LEN_PER_GPU=$(( STUDENT_MICRO_BATCH_SIZE_PER_GPU * (MAX_PROMPT + MAX_RESPONSE_LENGTH) ))
+USE_DYNAMIC_BSZ=True
+
+STUDENT_WORLD_SIZE=2
+
+TEACHER_RESOURCE_POOL=False
+TEACHER_WORLD_SIZE=4
+
+SP=1
+
+EXP_NAME="fsdp/student-${STUDENT_MODEL}/teacher-${TEACHER_MODEL}/loss-${DISTILLATION_LOSS_MODE}/pg-${USE_POLICY_GRADIENT}"
+
+ENFORCE_EAGER=True # true for faster debugging
+
+############################ Paths ############################
+
+gsm8k_train_path=$DATA_PATH/gsm8k/train.parquet
+gsm8k_test_path=$DATA_PATH/gsm8k/test.parquet
+
+TRAIN_FILES="['$gsm8k_train_path']"
+TEST_FILES="['$gsm8k_test_path']"
+
+############################ Parameter Groups ############################
+
+DATA=(
+    data.train_files="$TRAIN_FILES"
+    data.val_files="$TEST_FILES"
+    data.max_prompt_length=$MAX_PROMPT
+    data.max_response_length=$MAX_RESPONSE_LENGTH
+    data.train_batch_size=$TRAIN_PROMPT_BSZ
+    data.filter_overlong_prompts=True
+    data.truncation='error'
+    data.shuffle=False
+)
+
+MODEL=(
+    actor_rollout_ref.model.path="${FAMILY}/${STUDENT_MODEL}"
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+    actor_rollout_ref.model.use_remove_padding=True
+    actor_rollout_ref.model.use_fused_kernels=$USE_FUSED_KERNELS
+    actor_rollout_ref.actor.use_torch_compile=True
+    actor_rollout_ref.rollout.enforce_eager=$ENFORCE_EAGER
+)
+
+DISTILLATION=(
+    distillation.enabled=True
+    distillation.num_workers=8
+    distillation.teacher_model.enable_resource_pool=$TEACHER_RESOURCE_POOL
+    distillation.teacher_model.n_gpus_per_node=$TEACHER_WORLD_SIZE
+    distillation.teacher_model.nnodes=1
+    distillation.teacher_model.model_path="${FAMILY}/${TEACHER_MODEL}"
+    distillation.teacher_model.inference.tensor_model_parallel_size=1
+    distillation.teacher_model.inference.name=$ROLLOUT_NAME
+    distillation.teacher_model.inference.gpu_memory_utilization=0.3
+    distillation.teacher_model.inference.enforce_eager=$ENFORCE_EAGER
+    distillation.teacher_model.inference.max_model_len=$MAX_NUM_TOKENS
+    distillation.teacher_model.inference.max_num_batched_tokens=$MAX_NUM_TOKENS
+    distillation.teacher_model.inference.max_num_seqs=$MAX_NUM_TOKENS
+    distillation.distillation_loss.loss_mode=$DISTILLATION_LOSS_MODE
+    distillation.distillation_loss.topk=64
+    distillation.distillation_loss.use_task_rewards=False
+    distillation.distillation_loss.use_policy_gradient=$USE_POLICY_GRADIENT
+    distillation.distillation_loss.loss_max_clamp=$DISTILLATION_LOSS_MAX_CLAMP
+    distillation.distillation_loss.log_prob_min_clamp=$DISTILLATION_LOG_PROB_MIN_CLAMP
+)
+
+STUDENT=(
+    actor_rollout_ref.actor.optim.lr=1e-6
+    actor_rollout_ref.actor.ppo_mini_batch_size=$TRAIN_PROMPT_BSZ
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$STUDENT_MICRO_BATCH_SIZE_PER_GPU
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$STUDENT_MAX_TOKEN_LEN_PER_GPU
+    actor_rollout_ref.actor.use_dynamic_bsz=$USE_DYNAMIC_BSZ
+    actor_rollout_ref.actor.fsdp_config.param_offload=True
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
+    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$SP
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=$STUDENT_MICRO_BATCH_SIZE_PER_GPU
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=$STUDENT_MAX_TOKEN_LEN_PER_GPU
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=$USE_DYNAMIC_BSZ
+    actor_rollout_ref.rollout.tensor_model_parallel_size=1
+    actor_rollout_ref.rollout.name=$ROLLOUT_NAME
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.3
+    actor_rollout_ref.rollout.calculate_log_probs=False
+    actor_rollout_ref.rollout.max_model_len=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.max_num_batched_tokens=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.max_num_seqs=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.n=1
+)
+
+ALGORITHM=(
+    algorithm.adv_estimator=grpo
+    algorithm.use_kl_in_reward=False
+)
+
+TRAINER=(
+    trainer.logger='["console","wandb"]'
+    trainer.project_name=$PROJECT_NAME
+    trainer.experiment_name=$EXP_NAME
+    trainer.n_gpus_per_node=$STUDENT_WORLD_SIZE
+    trainer.nnodes=1
+    trainer.save_freq=200
+    trainer.test_freq=5
+    trainer.total_epochs=15
+    trainer.val_before_train=False
+    trainer.use_legacy_worker_impl=disable
+    trainer.resume_mode=disable
+    trainer.log_val_generations=5
+)
+
+
+
+############################ Launch ############################
+
+python3 -m verl.trainer.main_ppo \
+    --config-path=config \
+    --config-name='ppo_trainer.yaml' \
+    "${DATA[@]}" \
+    "${ALGORITHM[@]}" \
+    "${MODEL[@]}" \
+    "${DISTILLATION[@]}" \
+    "${ROLLOUT[@]}" \
+    "${STUDENT[@]}" \
+    "${TRAINER[@]}" \
+    "$@"