diff --git a/examples/on_policy_distillation_trainer/run_dsr1_1.5b_reverse_kl_fsdp.sh b/examples/on_policy_distillation_trainer/run_dsr1_1.5b_reverse_kl_fsdp.sh new file mode 100644 index 00000000000..117bb0f4b33 --- /dev/null +++ b/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[@]}" \ + "$@" diff --git a/examples/on_policy_distillation_trainer/run_qwen3_reverse_kl_fsdp.sh b/examples/on_policy_distillation_trainer/run_qwen3_reverse_kl_fsdp.sh new file mode 100755 index 00000000000..9fa4057ac9c --- /dev/null +++ b/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[@]}" \ + "$@" diff --git a/tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py b/tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py index de5b8a9e026..c56058b9879 100644 --- a/tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py +++ b/tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py @@ -268,6 +268,7 @@ class _DummyWorker: _compute_teacher_logprobs = AgentLoopWorker._compute_teacher_logprobs _pad_token_ids = AgentLoopWorker._pad_token_ids distillation_enabled = False + use_vllm_teacher = False def __init__(self): self.tokenizer = _FakeTokenizer() diff --git a/tests/workers/teacher/test_chunked_gather_on_cpu.py b/tests/workers/teacher/test_chunked_gather_on_cpu.py new file mode 100644 index 00000000000..5f8077a9852 --- /dev/null +++ b/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) diff --git a/tests/workers/test_reverse_kl_topk_on_cpu.py b/tests/workers/test_reverse_kl_topk_on_cpu.py new file mode 100644 index 00000000000..a3a4e3ef52f --- /dev/null +++ b/tests/workers/test_reverse_kl_topk_on_cpu.py @@ -0,0 +1,213 @@ +# 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 the student-top-K reverse-KL distillation loss (verl#6676). + +These tests live next to ``test_distillation_topk_symmetry_on_cpu.py`` and follow +the same CPU-only contract: no FSDP, no CUDA, no flash-attn — just direct calls +into :func:`verl.trainer.distillation.fsdp.losses.compute_reverse_kl_topk` and the +top-level wrapper registered as ``loss_mode='reverse_kl_topk'``. +""" + +import os + +os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from tensordict import TensorDict + +from verl.trainer.distillation.fsdp.losses import compute_reverse_kl_topk as compute_fsdp_reverse_kl_topk +from verl.trainer.distillation.losses import ( + compute_reverse_kl_topk as collect_reverse_kl_topk_metrics, +) +from verl.trainer.distillation.losses import ( + get_distillation_loss_fn, + get_distillation_loss_settings, +) + + +def _config(log_prob_min_clamp=None): + return SimpleNamespace(distillation_loss=SimpleNamespace(log_prob_min_clamp=log_prob_min_clamp)) + + +def _student_topk_inputs(B=2, T=3, V=8, K=3, seed=0): + """Random student logits + the matching student-top-K IDs (stops the gradient + on the indices, like the trainer does in practice).""" + torch.manual_seed(seed) + student_logits = torch.randn(B, T, V, requires_grad=True) + with torch.no_grad(): + topk = torch.topk(F.log_softmax(student_logits, dim=-1), k=K, dim=-1) + return student_logits, topk.indices.detach() + + +def test_reverse_kl_topk_is_nonnegative(): + """KL(Q || P) on the student-top-K support is >= 0 up to numerical noise. + + Use independent random tensors for student logits and teacher log-probs (so + Q and P do not coincide and the KL is materially positive).""" + student_logits, student_topk_ids = _student_topk_inputs() + torch.manual_seed(123) + teacher_logits = torch.randn(*student_logits.shape) + teacher_on_student_logp = torch.gather(F.log_softmax(teacher_logits, dim=-1), dim=-1, index=student_topk_ids) + + out = compute_fsdp_reverse_kl_topk( + student_logits=student_logits, + teacher_on_student_logp=teacher_on_student_logp, + student_topk_ids=student_topk_ids, + config=_config(), + data_format="thd", + ) + assert torch.all(out["distillation_losses"] >= -1e-6), out["distillation_losses"] + # For independent Q and P the KL should be materially positive somewhere. + assert out["distillation_losses"].max().item() > 1e-3 + + +def test_reverse_kl_topk_is_zero_when_distributions_match(): + """KL(Q || Q) == 0 — feeding student log-probs as the teacher log-probs gives a zero loss.""" + student_logits, student_topk_ids = _student_topk_inputs() + student_log_probs = F.log_softmax(student_logits.detach(), dim=-1) + teacher_on_student_logp = torch.gather(student_log_probs, dim=-1, index=student_topk_ids) + + out = compute_fsdp_reverse_kl_topk( + student_logits=student_logits, + teacher_on_student_logp=teacher_on_student_logp, + student_topk_ids=student_topk_ids, + config=_config(), + data_format="thd", + ) + torch.testing.assert_close( + out["distillation_losses"], torch.zeros_like(out["distillation_losses"]), atol=1e-6, rtol=0.0 + ) + + +def test_reverse_kl_topk_gradient_flows_to_student(): + """``loss.sum().backward()`` must populate ``student_logits.grad`` while leaving + the teacher tensor's gradient untouched (teacher is treated as a constant).""" + student_logits, student_topk_ids = _student_topk_inputs() + torch.manual_seed(7) + teacher_logits = torch.randn(*student_logits.shape) + teacher_on_student_logp = torch.gather(F.log_softmax(teacher_logits, dim=-1), dim=-1, index=student_topk_ids) + + out = compute_fsdp_reverse_kl_topk( + student_logits=student_logits, + teacher_on_student_logp=teacher_on_student_logp, + student_topk_ids=student_topk_ids, + config=_config(), + data_format="thd", + ) + out["distillation_losses"].sum().backward() + + assert student_logits.grad is not None + assert torch.isfinite(student_logits.grad).all() + assert student_logits.grad.abs().sum() > 0 + # The teacher tensor is detached input from the worker; we did not set requires_grad on it. + assert teacher_on_student_logp.grad is None + + +def test_reverse_kl_topk_matches_reference_formula(): + """The token-wise loss equals the manual reference reverse-KL computed with ``q~`` / + ``p~``, the student / teacher probabilities renormalized onto the shared + student-top-K support (softmax over the K gathered logits).""" + student_logits, student_topk_ids = _student_topk_inputs(B=1, T=4, V=10, K=4, seed=2) + torch.manual_seed(2) + teacher_logits = torch.randn(*student_logits.shape) + teacher_on_student_logp = torch.gather(F.log_softmax(teacher_logits, dim=-1), dim=-1, index=student_topk_ids) + + out = compute_fsdp_reverse_kl_topk( + student_logits=student_logits, + teacher_on_student_logp=teacher_on_student_logp, + student_topk_ids=student_topk_ids, + config=_config(), + data_format="thd", + ) + student_log_probs = F.log_softmax(student_logits.detach(), dim=-1) + student_topk_logp = torch.gather(student_log_probs, dim=-1, index=student_topk_ids) + # Renormalize both sides onto the student-top-K support (matches the loss). + student_norm = student_topk_logp - torch.logsumexp(student_topk_logp, dim=-1, keepdim=True) + teacher_norm = teacher_on_student_logp - torch.logsumexp(teacher_on_student_logp, dim=-1, keepdim=True) + expected = (student_norm.exp() * (student_norm - teacher_norm)).sum(dim=-1) + torch.testing.assert_close(out["distillation_losses"], expected, atol=1e-6, rtol=1e-6) + + +def test_reverse_kl_topk_invariant_to_mass_outside_topk(): + """Renormalizing onto the student-top-K support makes the loss depend only on + the within-top-K logits, so leaking mass outside the top-K must not change it. + Regression guard for the de-peaking collapse.""" + student_logits, student_topk_ids = _student_topk_inputs(B=2, T=3, V=8, K=3, seed=5) + torch.manual_seed(9) + teacher_logits = torch.randn(*student_logits.shape) + teacher_on_student_logp = torch.gather(F.log_softmax(teacher_logits, dim=-1), dim=-1, index=student_topk_ids) + + def _loss(logits): + return compute_fsdp_reverse_kl_topk( + student_logits=logits, + teacher_on_student_logp=teacher_on_student_logp, + student_topk_ids=student_topk_ids, + config=_config(), + data_format="thd", + )["distillation_losses"] + + base = _loss(student_logits) + + # Raise every non-top-K logit -> mass leaks outside the (fixed) top-K support. + # With renormalization the loss must not move. + bumped = student_logits.detach().clone() + outside_mask = torch.ones_like(bumped, dtype=torch.bool) + outside_mask.scatter_(-1, student_topk_ids, False) + bumped[outside_mask] += 5.0 + bumped.requires_grad_(True) + + torch.testing.assert_close(_loss(bumped), base, atol=1e-5, rtol=1e-5) + + +def test_reverse_kl_topk_registered_and_collects_mass_metrics(): + """``loss_mode='reverse_kl_topk'`` is registered in the distillation loss registry + and its top-level wrapper aggregates the student/teacher mass diagnostics over + the response mask.""" + settings = get_distillation_loss_settings("reverse_kl_topk") + assert settings.use_topk and not settings.use_estimator + assert get_distillation_loss_fn("reverse_kl_topk") is collect_reverse_kl_topk_metrics + + data = TensorDict( + { + "prompts": torch.tensor([[101]]), + "responses": torch.tensor([[11, 12, 0]]), + "attention_mask": torch.tensor([[1, 1, 1, 0]]), + "response_mask": torch.tensor([[1, 1, 0]], dtype=torch.bool), + }, + batch_size=[1], + ) + model_output = { + "distillation_losses": torch.tensor([0.10, 0.20, 0.30]), + "student_mass": torch.tensor([0.97, 0.95, 0.10]), + "teacher_mass": torch.tensor([0.50, 0.70, 0.05]), + } + distillation_config = SimpleNamespace(distillation_loss=SimpleNamespace(topk=8)) + + losses, metrics = collect_reverse_kl_topk_metrics( + config=SimpleNamespace(), + distillation_config=distillation_config, + model_output=model_output, + data=data, + ) + # Loss is clamped non-negative and shape-aligns with the response mask. + assert torch.all(losses >= 0.0) + assert losses.shape == data["response_mask"].shape + + # student_mass / teacher_mass means are computed only over response-masked tokens. + assert metrics["distillation/student_mass"] == pytest.approx((0.97 + 0.95) / 2) + assert metrics["distillation/teacher_mass"] == pytest.approx((0.50 + 0.70) / 2) diff --git a/verl/experimental/agent_loop/agent_loop.py b/verl/experimental/agent_loop/agent_loop.py index 1bd98610dbb..52e836e2062 100644 --- a/verl/experimental/agent_loop/agent_loop.py +++ b/verl/experimental/agent_loop/agent_loop.py @@ -517,7 +517,13 @@ def __init__( # Online policy distillation self.distillation_enabled = is_distillation_enabled(config.distillation) - if self.distillation_enabled: + # Only the vLLM teacher injects logprobs synchronously in the rollout; + # the FSDP teacher is driven by the trainer's main loop after rollout + # (it needs the student top-K first), so agent_loop must not call it here. + self.use_vllm_teacher = ( + self.distillation_enabled and config.distillation.get("teacher_backend", "vllm") == "vllm" + ) + if self.use_vllm_teacher: from verl.experimental.teacher_loop.teacher_manager import AsyncTeacherLLMServerManager self.teacher_key: str = config.distillation.teacher_key @@ -1012,7 +1018,9 @@ async def _compute_teacher_logprobs( sample_kwargs: Optional[dict[str, Any]] = None, ) -> None: """Compute teacher logprobs for single sample.""" - if self.distillation_enabled and not validate: + # Only the vLLM teacher produces logprobs synchronously per trajectory; + # the FSDP teacher is driven by the trainer's main loop after rollout. + if self.use_vllm_teacher and not validate: routing_key = None if sample_kwargs is not None: routing_value = sample_kwargs.get(self.teacher_key) diff --git a/verl/experimental/teacher_loop/teacher_fsdp_manager.py b/verl/experimental/teacher_loop/teacher_fsdp_manager.py new file mode 100644 index 00000000000..6ea6a0bb884 --- /dev/null +++ b/verl/experimental/teacher_loop/teacher_fsdp_manager.py @@ -0,0 +1,148 @@ +# 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. +"""FSDP-teacher counterpart to ``MultiTeacherModelManager``. + +When ``distillation.teacher_backend == 'fsdp'``, the trainer hosts each teacher +as a Ray worker group of :class:`TeacherFSDPWorker` instead of a vLLM rollout +replica. This module provides a thin manager that holds the spawned worker group +handle(s) and dispatches ``compute_logprobs_at_ids`` calls to them. + +Worker spawning happens in ``ray_trainer.py`` (which registers the worker on the +teacher resource pool, spawns, then calls :meth:`set_worker_group`). Single-teacher +is wired; multi-teacher FSDP raises ``NotImplementedError``. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Optional + +from omegaconf import DictConfig +from tensordict import TensorDict + +from verl.utils.config import omega_conf_to_dataclass +from verl.workers.config import DistillationConfig + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +class FSDPTeacherManager: + """Per-teacher wrapper around a Ray worker group of :class:`TeacherFSDPWorker`. + + Mirrors the per-teacher surface of the vLLM ``TeacherModelManager`` so the + trainer's backend branching stays symmetric, but does *not* spawn its own + workers — the trainer spawns the worker group and calls :meth:`set_worker_group`. + """ + + def __init__(self, key: str, distillation_config: DistillationConfig): + self.key = key + self.distillation_config = distillation_config + self.worker_group: Any = None + + def set_worker_group(self, worker_group: Any) -> None: + """Install the spawned :class:`TeacherFSDPWorker` Ray worker group handle.""" + self.worker_group = worker_group + + def compute_logprobs_at_ids(self, data: TensorDict) -> Optional[TensorDict]: + """Forward a teacher-side ``compute_logprobs_at_ids`` call to the worker group. + + See :meth:`verl.workers.teacher.fsdp_teacher_worker.TeacherFSDPWorker.compute_logprobs_at_ids` + for the data contract. + """ + if self.worker_group is None: + raise RuntimeError( + f"FSDPTeacherManager[{self.key!r}].worker_group has not been installed yet — " + "the trainer must call set_worker_group(...) after spawning the unified " + "Ray worker dict in init_workers." + ) + return self.worker_group.compute_logprobs_at_ids(data) + + +class MultiTeacherFSDPManager: + """Manages one :class:`FSDPTeacherManager` per teacher routing key. + + Worker-group spawn is driven by the trainer (so the teacher worker joins the + unified ``RayWorkerGroup`` spawn); the trainer then calls :meth:`set_worker_group` + to install the handle. + """ + + def __init__(self, config: DictConfig, resource_pool: Any): + self.config = config + self.distillation_config: DistillationConfig = omega_conf_to_dataclass(config.distillation) + if self.distillation_config.teacher_backend != "fsdp": + raise ValueError( + f"MultiTeacherFSDPManager requires distillation.teacher_backend='fsdp', " + f"got {self.distillation_config.teacher_backend!r}." + ) + if self.distillation_config.teacher_fsdp_config is None: + raise ValueError("distillation.teacher_fsdp_config must be set when teacher_backend='fsdp'.") + + # Resource pool is held for symmetry with MultiTeacherModelManager (the + # vLLM equivalent), but the FSDP path drives spawn from the trainer. + self.resource_pool = resource_pool + + teacher_models = self.distillation_config.teacher_models + if len(teacher_models) != 1: + raise NotImplementedError( + "Multi-teacher FSDP backend is not yet supported; configure exactly one teacher " + "(teacher_models has 1 entry) when teacher_backend='fsdp'." + ) + + self.teacher_managers: dict[str, FSDPTeacherManager] = { + key: FSDPTeacherManager(key=key, distillation_config=self.distillation_config) for key in teacher_models + } + + @property + def teacher_keys(self) -> list[str]: + return list(self.teacher_managers) + + def set_worker_group(self, worker_group: Any, key: Optional[str] = None) -> None: + """Install the spawned worker group handle for the (single) teacher. + + Args: + worker_group: spawned Ray worker group handle returned by + ``wg_dict.spawn(...)``. + key: routing key of the teacher to attach the worker group to. May + be omitted in single-teacher setups. + """ + if key is None: + if len(self.teacher_managers) != 1: + raise ValueError( + "MultiTeacherFSDPManager.set_worker_group requires an explicit `key` when " + "multiple teachers are configured." + ) + key = next(iter(self.teacher_managers)) + if key not in self.teacher_managers: + raise ValueError(f"Teacher key {key!r} not configured; configured: {self.teacher_keys}.") + self.teacher_managers[key].set_worker_group(worker_group) + + def get(self, routing_key: Optional[str] = None) -> FSDPTeacherManager: + if len(self.teacher_managers) == 1: + return next(iter(self.teacher_managers.values())) + if routing_key is None: + raise ValueError("Routing key is required when multiple FSDP teachers are configured.") + if routing_key not in self.teacher_managers: + raise ValueError( + f"No FSDP teacher configured for routing key {routing_key!r}; " + f"configured teachers: {sorted(self.teacher_managers)}." + ) + return self.teacher_managers[routing_key] + + def compute_logprobs_at_ids(self, data: TensorDict, routing_key: Optional[str] = None) -> Optional[TensorDict]: + return self.get(routing_key).compute_logprobs_at_ids(data) + + +__all__ = ["FSDPTeacherManager", "MultiTeacherFSDPManager"] diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index f636f202e06..eb29f51a0a4 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -899,6 +899,16 @@ distillation: response_length: ${oc.select:actor_rollout_ref.rollout.response_length} temperature: ${oc.select:actor_rollout_ref.rollout.temperature} teacher_key: data_source + teacher_backend: vllm + teacher_chunk_size: 1024 + teacher_fsdp_config: + _target_: verl.workers.config.FSDPEngineConfig + forward_only: true + strategy: fsdp + dtype: bfloat16 + param_offload: false + fsdp_size: -1 + use_dynamic_bsz: true trainer: balance_batch: true total_epochs: 30 diff --git a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml index baccd878ff6..04e5f850f42 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -821,6 +821,16 @@ distillation: response_length: ${oc.select:actor_rollout_ref.rollout.response_length} temperature: ${oc.select:actor_rollout_ref.rollout.temperature} teacher_key: data_source + teacher_backend: vllm + teacher_chunk_size: 1024 + teacher_fsdp_config: + _target_: verl.workers.config.FSDPEngineConfig + forward_only: true + strategy: fsdp + dtype: bfloat16 + param_offload: false + fsdp_size: -1 + use_dynamic_bsz: true trainer: balance_batch: true total_epochs: 30 diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index f91eb3a41f9..b7e03de5dc2 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -837,6 +837,16 @@ distillation: response_length: ${oc.select:actor_rollout_ref.rollout.response_length} temperature: ${oc.select:actor_rollout_ref.rollout.temperature} teacher_key: data_source + teacher_backend: vllm + teacher_chunk_size: 1024 + teacher_fsdp_config: + _target_: verl.workers.config.FSDPEngineConfig + forward_only: true + strategy: fsdp + dtype: bfloat16 + param_offload: false + fsdp_size: -1 + use_dynamic_bsz: true trainer: balance_batch: true total_epochs: 30 diff --git a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml index dd93ec973f6..2ab882c1dd1 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -835,6 +835,16 @@ distillation: response_length: ${oc.select:actor_rollout_ref.rollout.response_length} temperature: ${oc.select:actor_rollout_ref.rollout.temperature} teacher_key: data_source + teacher_backend: vllm + teacher_chunk_size: 1024 + teacher_fsdp_config: + _target_: verl.workers.config.FSDPEngineConfig + forward_only: true + strategy: fsdp + dtype: bfloat16 + param_offload: false + fsdp_size: -1 + use_dynamic_bsz: true trainer: balance_batch: true total_epochs: 30 diff --git a/verl/trainer/config/distillation/distillation.yaml b/verl/trainer/config/distillation/distillation.yaml index 84ee7932ab6..dbc452ab073 100644 --- a/verl/trainer/config/distillation/distillation.yaml +++ b/verl/trainer/config/distillation/distillation.yaml @@ -110,4 +110,34 @@ teacher_models: temperature: ${oc.select:actor_rollout_ref.rollout.temperature} # Key to route examples to the appropriate teacher model in multi-teacher setups. Should correspond to a field in the data proto, e.g., task. -teacher_key: data_source \ No newline at end of file +teacher_key: data_source + +# Backend hosting the teacher: "vllm" (default) or "fsdp" (frozen FSDP worker, for reverse-KL). +teacher_backend: vllm + +# Token-axis chunk size for the FSDP teacher's gather-at-IDs kernel. Only used when teacher_backend == "fsdp". +teacher_chunk_size: 1024 + +# FSDP config for the frozen teacher worker. Only used when teacher_backend == "fsdp". +teacher_fsdp_config: + + # Required for omega_conf_to_dataclass instantiation. + _target_: verl.workers.config.FSDPEngineConfig + + # Teacher is frozen — forward-only. + forward_only: true + + # FSDP backend: "fsdp" (FSDP1) or "fsdp2". + strategy: fsdp + + # Parameter dtype. + dtype: bfloat16 + + # Offload parameters to CPU. + param_offload: false + + # FSDP shard group size; -1 = all GPUs. + fsdp_size: -1 + + # Inference batching. + use_dynamic_bsz: true \ No newline at end of file diff --git a/verl/trainer/distillation/fsdp/losses.py b/verl/trainer/distillation/fsdp/losses.py index 602d3875338..97d3ebf485a 100644 --- a/verl/trainer/distillation/fsdp/losses.py +++ b/verl/trainer/distillation/fsdp/losses.py @@ -1,149 +1,268 @@ -# Copyright 2025 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. - - -import torch -import torch.nn.functional as F - -from verl.utils.ulysses import ( - get_ulysses_sequence_parallel_world_size, - slice_input_tensor, -) -from verl.workers.config import DistillationConfig, DistillationLossConfig - - -def _chunked_topk_log_probs( - logits: torch.Tensor, - topk_ids: torch.Tensor, - chunk_size: int = 4096, -) -> torch.Tensor: - """Compute log_softmax(logits).gather(topk_ids) without materializing [B, T, V]. - - Uses the identity: - log_softmax(x).gather(idx) == x.gather(idx) - logsumexp(x, keepdim=True) - Streams the reduction in chunks of `chunk_size` tokens along (B*T) with fp32 - logsumexp for numerical stability. - - Args: - logits: [B, T, V] student logits. - topk_ids: [B, T, K] indices to gather. - chunk_size: number of tokens per chunk; only affects memory, not numerics. - - Returns: - [B, T, K] tensor with the same dtype as `logits`. - """ - B, T, V = logits.shape - K = topk_ids.shape[-1] - flat_logits = logits.reshape(-1, V) # [N, V] - flat_topk = topk_ids.reshape(-1, K) # [N, K] - N = flat_logits.shape[0] - - # Edge case: empty input (e.g. fully-padded micro-batch). - if N == 0: - return torch.empty((B, T, K), dtype=logits.dtype, device=logits.device) - - out = torch.empty((N, K), dtype=logits.dtype, device=logits.device) - for s in range(0, N, chunk_size): - e = min(s + chunk_size, N) - chunk_logits_fp32 = flat_logits[s:e].float() - log_z = torch.logsumexp(chunk_logits_fp32, dim=-1, keepdim=True) # [c, 1] - chunk_topk_logits = torch.gather(chunk_logits_fp32, dim=-1, index=flat_topk[s:e]) - out[s:e] = (chunk_topk_logits - log_z).to(logits.dtype) - return out.reshape(B, T, K) - - -def kl_divergence(log_q: torch.Tensor, log_p: torch.Tensor) -> torch.Tensor: - """Compute KL divergence between two distributions given their log probabilities.""" - log_p = log_p.float() - log_q = log_q.float() - p = log_p.exp() - kld = p * (log_p - log_q) - return kld.sum(dim=-1) - - -def compute_forward_kl_topk( - student_logits: torch.Tensor, - teacher_topk_log_probs: torch.Tensor, - teacher_topk_ids: torch.Tensor, - config: DistillationConfig, - data_format: str, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute forward KL distillation loss using top-k log probabilities. - - Args: - student_logits: (bsz, seqlen/sp_size, vocab_size). - teacher_topk_log_probs: (bsz, seqlen, topk). - teacher_topk_ids: (bsz, seqlen, topk). - data_format: "thd" or "bshd", models not support THD format, e.g GPT-OSS, Qwen3.5 - - Returns: - - distillation_losses: (bsz, seqlen/sp_size) - - student_mass: (bsz, seqlen/sp_size) - - teacher_mass: (bsz, seqlen/sp_size) - """ - assert teacher_topk_log_probs.is_nested and teacher_topk_ids.is_nested - teacher_topk_log_probs = teacher_topk_log_probs.values().unsqueeze(0) # (1, total_nnz, topk) - teacher_topk_ids = teacher_topk_ids.values().unsqueeze(0) # (1, total_nnz, topk) - - # 1. split across sp groups (bsz, seqlen, topk) => (bsz, seqlen/sp_size, topk) - if get_ulysses_sequence_parallel_world_size() > 1: - teacher_topk_log_probs = slice_input_tensor(teacher_topk_log_probs, dim=1) - teacher_topk_ids = slice_input_tensor(teacher_topk_ids, dim=1) - assert teacher_topk_log_probs.shape[:2] == teacher_topk_ids.shape[:2] == student_logits.shape[:2] - - # 2. compute token-wise KL divergence across sp groups - # ``use_chunked_topk`` (opt-in, default off) trades latency for memory: - # the chunked path streams logsumexp + gather to avoid the [B, T, V] - # log_softmax buffer, enabling long-context (>=64K) where the default - # F.log_softmax path OOMs. See ``DistillationLossConfig.use_chunked_topk`` - # for trade-offs and benchmark numbers. - loss_config: DistillationLossConfig = config.distillation_loss - use_chunked_topk = getattr(loss_config, "use_chunked_topk", False) - if use_chunked_topk: - # log_softmax is monotonic, so topk(logits) == topk(log_softmax(logits)). - student_topk_ids = torch.topk(student_logits, k=teacher_topk_ids.shape[-1], dim=-1).indices - student_topk_log_probs = _chunked_topk_log_probs( - student_logits, - teacher_topk_ids, - chunk_size=getattr(loss_config, "chunked_topk_chunk_size", 4096), - ) - else: - student_log_probs = F.log_softmax(student_logits, dim=-1) - student_topk_ids = torch.topk(student_log_probs, k=teacher_topk_ids.shape[-1], dim=-1).indices - student_topk_log_probs = torch.gather(student_log_probs, dim=-1, index=teacher_topk_ids) - student_mass = student_topk_log_probs.exp().sum(dim=-1) - teacher_mass = teacher_topk_log_probs.exp().sum(dim=-1) - if loss_config.log_prob_min_clamp is not None: - student_topk_log_probs = student_topk_log_probs.clamp_min(loss_config.log_prob_min_clamp) - teacher_topk_log_probs = teacher_topk_log_probs.clamp_min(loss_config.log_prob_min_clamp) - distillation_losses = kl_divergence(log_q=student_topk_log_probs, log_p=teacher_topk_log_probs) - - # Diagnostics for tracking teacher/student top-k overlap in OPD, following - # "Rethinking On-Policy Distillation of Large Language Models" (arXiv:2604.13016). - overlap_mask = (teacher_topk_ids.unsqueeze(-1) == student_topk_ids.unsqueeze(-2)).any(dim=-1) - overlap_count = overlap_mask.sum(dim=-1) - token_kl = teacher_topk_log_probs.exp() * (teacher_topk_log_probs - student_topk_log_probs) - overlap_token_advantage_sum = (-token_kl * overlap_mask).sum(dim=-1) - overlap_token_advantage = overlap_token_advantage_sum / overlap_count.clamp_min(1) - overlap_token_advantage = torch.where( - overlap_count > 0, overlap_token_advantage, torch.zeros_like(overlap_token_advantage) - ) - - return { - "distillation_losses": distillation_losses, - "student_mass": student_mass, - "teacher_mass": teacher_mass, - "overlap_count": overlap_count, - "overlap_token_advantage": overlap_token_advantage, - } +# Copyright 2025 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. + + +import torch +import torch.nn.functional as F + +from verl.utils.ulysses import ( + get_ulysses_sequence_parallel_world_size, + slice_input_tensor, +) +from verl.workers.config import DistillationConfig, DistillationLossConfig + + +def _chunked_topk_log_probs( + logits: torch.Tensor, + topk_ids: torch.Tensor, + chunk_size: int = 4096, +) -> torch.Tensor: + """Compute log_softmax(logits).gather(topk_ids) without materializing [B, T, V]. + + Uses the identity: + log_softmax(x).gather(idx) == x.gather(idx) - logsumexp(x, keepdim=True) + Streams the reduction in chunks of `chunk_size` tokens along (B*T) with fp32 + logsumexp for numerical stability. + + Args: + logits: [B, T, V] student logits. + topk_ids: [B, T, K] indices to gather. + chunk_size: number of tokens per chunk; only affects memory, not numerics. + + Returns: + [B, T, K] tensor with the same dtype as `logits`. + """ + B, T, V = logits.shape + K = topk_ids.shape[-1] + flat_logits = logits.reshape(-1, V) # [N, V] + flat_topk = topk_ids.reshape(-1, K) # [N, K] + N = flat_logits.shape[0] + + # Edge case: empty input (e.g. fully-padded micro-batch). + if N == 0: + return torch.empty((B, T, K), dtype=logits.dtype, device=logits.device) + + out = torch.empty((N, K), dtype=logits.dtype, device=logits.device) + for s in range(0, N, chunk_size): + e = min(s + chunk_size, N) + chunk_logits_fp32 = flat_logits[s:e].float() + log_z = torch.logsumexp(chunk_logits_fp32, dim=-1, keepdim=True) # [c, 1] + chunk_topk_logits = torch.gather(chunk_logits_fp32, dim=-1, index=flat_topk[s:e]) + out[s:e] = (chunk_topk_logits - log_z).to(logits.dtype) + return out.reshape(B, T, K) + + +def kl_divergence(log_q: torch.Tensor, log_p: torch.Tensor) -> torch.Tensor: + """Compute KL divergence between two distributions given their log probabilities.""" + log_p = log_p.float() + log_q = log_q.float() + p = log_p.exp() + kld = p * (log_p - log_q) + return kld.sum(dim=-1) + + +def reverse_kl_divergence(log_q: torch.Tensor, log_p: torch.Tensor) -> torch.Tensor: + """Compute reverse KL divergence on a shared support. + + Computes KL(Q || P) with the expectation taken under Q (the "student" distribution), + which is the mode-seeking direction discussed in + "Rethinking On-Policy Distillation of Large Language Models" (arXiv:2604.13016) + and verl issue #6676. + + Args: + log_q: Student log-probabilities at the shared support, shape ``(..., K)``. + log_p: Teacher log-probabilities at the same support, shape ``(..., K)``. + + Returns: + Token-wise KL(Q || P), shape ``(...,)`` (the trailing K axis is reduced). + """ + log_q = log_q.float() + log_p = log_p.float() + q = log_q.exp() + return (q * (log_q - log_p)).sum(dim=-1) + + +def compute_forward_kl_topk( + student_logits: torch.Tensor, + teacher_topk_log_probs: torch.Tensor, + teacher_topk_ids: torch.Tensor, + config: DistillationConfig, + data_format: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute forward KL distillation loss using top-k log probabilities. + + Args: + student_logits: (bsz, seqlen/sp_size, vocab_size). + teacher_topk_log_probs: (bsz, seqlen, topk). + teacher_topk_ids: (bsz, seqlen, topk). + data_format: "thd" or "bshd", models not support THD format, e.g GPT-OSS, Qwen3.5 + + Returns: + - distillation_losses: (bsz, seqlen/sp_size) + - student_mass: (bsz, seqlen/sp_size) + - teacher_mass: (bsz, seqlen/sp_size) + """ + assert teacher_topk_log_probs.is_nested and teacher_topk_ids.is_nested + teacher_topk_log_probs = teacher_topk_log_probs.values().unsqueeze(0) # (1, total_nnz, topk) + teacher_topk_ids = teacher_topk_ids.values().unsqueeze(0) # (1, total_nnz, topk) + + # 1. split across sp groups (bsz, seqlen, topk) => (bsz, seqlen/sp_size, topk) + if get_ulysses_sequence_parallel_world_size() > 1: + teacher_topk_log_probs = slice_input_tensor(teacher_topk_log_probs, dim=1) + teacher_topk_ids = slice_input_tensor(teacher_topk_ids, dim=1) + assert teacher_topk_log_probs.shape[:2] == teacher_topk_ids.shape[:2] == student_logits.shape[:2] + + # 2. compute token-wise KL divergence across sp groups + # ``use_chunked_topk`` (opt-in, default off) trades latency for memory: + # the chunked path streams logsumexp + gather to avoid the [B, T, V] + # log_softmax buffer, enabling long-context (>=64K) where the default + # F.log_softmax path OOMs. See ``DistillationLossConfig.use_chunked_topk`` + # for trade-offs and benchmark numbers. + loss_config: DistillationLossConfig = config.distillation_loss + use_chunked_topk = getattr(loss_config, "use_chunked_topk", False) + if use_chunked_topk: + # log_softmax is monotonic, so topk(logits) == topk(log_softmax(logits)). + student_topk_ids = torch.topk(student_logits, k=teacher_topk_ids.shape[-1], dim=-1).indices + student_topk_log_probs = _chunked_topk_log_probs( + student_logits, + teacher_topk_ids, + chunk_size=getattr(loss_config, "chunked_topk_chunk_size", 4096), + ) + else: + student_log_probs = F.log_softmax(student_logits, dim=-1) + student_topk_ids = torch.topk(student_log_probs, k=teacher_topk_ids.shape[-1], dim=-1).indices + student_topk_log_probs = torch.gather(student_log_probs, dim=-1, index=teacher_topk_ids) + student_mass = student_topk_log_probs.exp().sum(dim=-1) + teacher_mass = teacher_topk_log_probs.exp().sum(dim=-1) + if loss_config.log_prob_min_clamp is not None: + student_topk_log_probs = student_topk_log_probs.clamp_min(loss_config.log_prob_min_clamp) + teacher_topk_log_probs = teacher_topk_log_probs.clamp_min(loss_config.log_prob_min_clamp) + distillation_losses = kl_divergence(log_q=student_topk_log_probs, log_p=teacher_topk_log_probs) + + # Diagnostics for tracking teacher/student top-k overlap in OPD, following + # "Rethinking On-Policy Distillation of Large Language Models" (arXiv:2604.13016). + overlap_mask = (teacher_topk_ids.unsqueeze(-1) == student_topk_ids.unsqueeze(-2)).any(dim=-1) + overlap_count = overlap_mask.sum(dim=-1) + token_kl = teacher_topk_log_probs.exp() * (teacher_topk_log_probs - student_topk_log_probs) + overlap_token_advantage_sum = (-token_kl * overlap_mask).sum(dim=-1) + overlap_token_advantage = overlap_token_advantage_sum / overlap_count.clamp_min(1) + overlap_token_advantage = torch.where( + overlap_count > 0, overlap_token_advantage, torch.zeros_like(overlap_token_advantage) + ) + + return { + "distillation_losses": distillation_losses, + "student_mass": student_mass, + "teacher_mass": teacher_mass, + "overlap_count": overlap_count, + "overlap_token_advantage": overlap_token_advantage, + } + + +def compute_reverse_kl_topk( + student_logits: torch.Tensor, + teacher_on_student_logp: torch.Tensor, + student_topk_ids: torch.Tensor, + config: DistillationConfig, + data_format: str, + teacher_topk_ids: torch.Tensor = None, +) -> dict[str, torch.Tensor]: + """Compute student-top-K reverse KL distillation loss. + + The support is the student's own top-K; teacher log-probabilities at those IDs + are precomputed by the FSDP teacher worker, so no full ``[B, T, V]`` tensor is + materialized here. + + NOTE: student and teacher are renormalized onto the student-top-K support so the + loss is a proper ``KL(Q_K || P_K)`` (otherwise it could be minimized by leaking + mass outside the top-K). ``student_mass`` stays un-normalized as a diagnostic. + + Args: + student_logits: ``(bsz, seqlen/sp_size, vocab_size)``. + teacher_on_student_logp: ``(bsz, seqlen, K)`` — teacher log-probabilities + gathered at ``student_topk_ids``. + student_topk_ids: ``(bsz, seqlen, K)`` — the student's current-step top-K IDs. + config: Distillation config (used for the ``log_prob_min_clamp`` guard). + data_format: kept for API symmetry with :func:`compute_forward_kl_topk`; + not consumed here. + teacher_topk_ids: ``(bsz, seqlen, K)`` — the teacher's own top-K IDs, used + only for the overlap diagnostics. Overlap metrics are skipped when ``None``. + + Returns: + dict with ``distillation_losses`` and the ``student_mass`` / ``teacher_mass`` + diagnostics; also ``overlap_count`` / ``overlap_token_advantage`` when + ``teacher_topk_ids`` is given. + """ + del data_format # unused; kept for parity with compute_forward_kl_topk's signature + # 1. unwrap nested layouts from the rmpad pipeline + if teacher_on_student_logp.is_nested: + teacher_on_student_logp = teacher_on_student_logp.values().unsqueeze(0) + if student_topk_ids.is_nested: + student_topk_ids = student_topk_ids.values().unsqueeze(0) + if teacher_topk_ids is not None and teacher_topk_ids.is_nested: + teacher_topk_ids = teacher_topk_ids.values().unsqueeze(0) + + # 2. split across sp groups + if get_ulysses_sequence_parallel_world_size() > 1: + teacher_on_student_logp = slice_input_tensor(teacher_on_student_logp, dim=1) + student_topk_ids = slice_input_tensor(student_topk_ids, dim=1) + if teacher_topk_ids is not None: + teacher_topk_ids = slice_input_tensor(teacher_topk_ids, dim=1) + assert teacher_on_student_logp.shape[:2] == student_topk_ids.shape[:2] == student_logits.shape[:2] + assert teacher_on_student_logp.shape[-1] == student_topk_ids.shape[-1] + + # 3. gather student log-probs at the student's own top-K ids + student_log_probs = F.log_softmax(student_logits, dim=-1) + student_topk_log_probs = torch.gather(student_log_probs, dim=-1, index=student_topk_ids) + + # Raw mass on the student-top-K support: diagnostics only, not fed to the loss. + student_mass = student_topk_log_probs.exp().sum(dim=-1) + teacher_mass = teacher_on_student_logp.exp().sum(dim=-1) + + # Renormalize both sides onto the student-top-K support -> proper KL(Q_K || P_K). + student_topk_log_probs = student_topk_log_probs - torch.logsumexp(student_topk_log_probs, dim=-1, keepdim=True) + teacher_on_student_logp = teacher_on_student_logp - torch.logsumexp(teacher_on_student_logp, dim=-1, keepdim=True) + + loss_config: DistillationLossConfig = config.distillation_loss + if loss_config.log_prob_min_clamp is not None: + student_topk_log_probs = student_topk_log_probs.clamp_min(loss_config.log_prob_min_clamp) + teacher_on_student_logp = teacher_on_student_logp.clamp_min(loss_config.log_prob_min_clamp) + + distillation_losses = reverse_kl_divergence( + log_q=student_topk_log_probs, + log_p=teacher_on_student_logp, + ) + + outputs = { + "distillation_losses": distillation_losses, + "student_mass": student_mass, + "teacher_mass": teacher_mass, + } + + # Diagnostics for tracking student/teacher top-k overlap in OPD, following + # "Rethinking On-Policy Distillation of Large Language Models" (arXiv:2604.13016). + if teacher_topk_ids is not None: + # overlap_mask[..., i] = student-top-K token i is also in the teacher top-K. + overlap_mask = (student_topk_ids.unsqueeze(-1) == teacher_topk_ids.unsqueeze(-2)).any(dim=-1) + overlap_count = overlap_mask.sum(dim=-1) + token_rkl = student_topk_log_probs.exp() * (student_topk_log_probs - teacher_on_student_logp) + overlap_token_advantage_sum = (-token_rkl * overlap_mask).sum(dim=-1) + overlap_token_advantage = overlap_token_advantage_sum / overlap_count.clamp_min(1) + overlap_token_advantage = torch.where( + overlap_count > 0, overlap_token_advantage, torch.zeros_like(overlap_token_advantage) + ) + outputs["overlap_count"] = overlap_count + outputs["overlap_token_advantage"] = overlap_token_advantage + + return outputs diff --git a/verl/trainer/distillation/losses.py b/verl/trainer/distillation/losses.py index 3467e59ebaa..533d6fd1f1b 100644 --- a/verl/trainer/distillation/losses.py +++ b/verl/trainer/distillation/losses.py @@ -16,10 +16,12 @@ from typing import Any, Callable, Optional import torch +import torch.nn.functional as F from tensordict import TensorDict from verl.base_config import BaseConfig from verl.trainer.ppo.core_algos import agg_loss, get_policy_loss_fn, kl_penalty +from verl.utils import tensordict_utils as tu from verl.utils.metric import AggregationType, Metric from verl.workers.config import ActorConfig, DistillationConfig, DistillationLossConfig from verl.workers.utils.losses import ppo_loss @@ -134,26 +136,67 @@ def compute_topk_loss( - student_mass: (bsz, seqlen/cp_size) - teacher_mass: (bsz, seqlen/cp_size) """ - match config.strategy: - # VeOmni uses FSDP2 internally, so its loss computation is identical to FSDP. - case "fsdp" | "veomni": - import verl.trainer.distillation.fsdp.losses as fsdp_losses - - distillation_loss_fn = fsdp_losses.compute_forward_kl_topk - case "megatron": - import verl.trainer.distillation.megatron.losses as megatron_losses - - distillation_loss_fn = megatron_losses.compute_forward_kl_topk - case _: - raise NotImplementedError(f"Unsupported strategy: {config.strategy=}") - - outputs = distillation_loss_fn( - student_logits=student_logits, - teacher_topk_log_probs=data["teacher_logprobs"], - teacher_topk_ids=data["teacher_ids"], - config=distillation_config, - data_format=data_format, - ) + # FSDP-teacher first pass: the trainer stamps ``compute_student_topk_only=True`` + # to get the student's own top-K before the teacher forward. Short-circuit and + # return the student top-K log-probs and IDs, with no teacher data dependency. + if tu.get_non_tensor_data(data=data, key="compute_student_topk_only", default=False): + topk = distillation_config.distillation_loss.topk + if topk is None: + raise ValueError( + "distillation_loss.topk must be set when compute_student_topk_only=True; " + "the FSDP teacher pipeline uses it as the student-top-K width." + ) + student_log_probs = F.log_softmax(student_logits, dim=-1) + topk_out = torch.topk(student_log_probs, k=topk, dim=-1) + return { + "student_topk_log_probs": topk_out.values, + "student_topk_ids": topk_out.indices.to(torch.int64), + } + + loss_mode = distillation_config.distillation_loss.loss_mode + + if loss_mode == "forward_kl_topk": + match config.strategy: + # VeOmni uses FSDP2 internally, so its loss computation is identical to FSDP. + case "fsdp" | "veomni": + import verl.trainer.distillation.fsdp.losses as fsdp_losses + + distillation_loss_fn = fsdp_losses.compute_forward_kl_topk + case "megatron": + import verl.trainer.distillation.megatron.losses as megatron_losses + + distillation_loss_fn = megatron_losses.compute_forward_kl_topk + case _: + raise NotImplementedError(f"Unsupported strategy: {config.strategy=}") + + outputs = distillation_loss_fn( + student_logits=student_logits, + teacher_topk_log_probs=data["teacher_logprobs"], + teacher_topk_ids=data["teacher_ids"], + config=distillation_config, + data_format=data_format, + ) + elif loss_mode == "reverse_kl_topk": + match config.strategy: + case "fsdp" | "veomni": + import verl.trainer.distillation.fsdp.losses as fsdp_losses + + distillation_loss_fn = fsdp_losses.compute_reverse_kl_topk + case _: + raise NotImplementedError( + f"reverse_kl_topk is currently only implemented for fsdp/veomni, got {config.strategy=}" + ) + + outputs = distillation_loss_fn( + student_logits=student_logits, + teacher_on_student_logp=data["teacher_on_student_logp"], + student_topk_ids=data["student_topk_ids"], + config=distillation_config, + data_format=data_format, + teacher_topk_ids=data.get("teacher_topk_ids", None), + ) + else: + raise NotImplementedError(f"Unsupported loss_mode for compute_topk_loss: {loss_mode!r}") expected_shape = student_logits.shape[:2] for k, v in outputs.items(): @@ -204,6 +247,10 @@ def distillation_ppo_loss( if student_logits is not None: return compute_topk_loss(config, distillation_config, data, student_logits, data_format) + # FSDP-teacher first pass: only emit student top-K, skip the policy loss. + if tu.get_non_tensor_data(data=data, key="compute_student_topk_only", default=False): + return torch.zeros((), device=model_output["log_probs"].device), {} + # Called as final policy loss distillation_loss_config = distillation_config.distillation_loss distill_loss, distill_metrics = distillation_loss(config, distillation_config, model_output, data) @@ -397,3 +444,75 @@ def compute_distillation_loss_reverse_kl_estimator( "distillation/abs_loss": Metric(AggregationType.MEAN, distillation_losses[response_mask_bool].abs().mean()), } return distillation_losses, metrics + + +@register_distillation_loss(DistillationLossSettings(names=["reverse_kl_topk"], use_topk=True)) # type: ignore[arg-type] +def compute_reverse_kl_topk( + config: ActorConfig, + distillation_config: DistillationConfig, + model_output: dict, + data: TensorDict, +) -> tuple[torch.Tensor, dict[str, Any]]: + """Top-level wrapper for student-top-K reverse KL distillation. + + The token-wise reverse KL ``q * (log q - log p)`` (summed over the K axis) and + the per-token student / teacher mass tensors are produced inside the logits + processor by :func:`verl.trainer.distillation.fsdp.losses.compute_reverse_kl_topk`. + This wrapper only re-pads them, computes diagnostic metrics, and clamps the loss + against minor negative values introduced by ``log_prob_min_clamp``. + + Returns: + - ``distillation_losses``: ``(bsz, resp_len)`` clamped to ``>= 0``. + - ``distillation_metrics``: dict with student/teacher mass diagnostics. The + ``distillation/teacher_mass`` curve is the key mode-seeking signal — it + should rise during training as the student concentrates on tokens the + teacher also assigns high probability to. + """ + distillation_losses = no_padding_2_padding(model_output["distillation_losses"], data) + student_mass = no_padding_2_padding(model_output["student_mass"], data) + teacher_mass = no_padding_2_padding(model_output["teacher_mass"], data) + overlap_count = model_output.get("overlap_count") + overlap_token_advantage = model_output.get("overlap_token_advantage") + if overlap_count is not None and overlap_token_advantage is not None: + overlap_count = no_padding_2_padding(overlap_count, data) + overlap_token_advantage = no_padding_2_padding(overlap_token_advantage, data) + if data["response_mask"].is_nested: + response_mask_bool = data["response_mask"].bool().to_padded_tensor(False) + else: + response_mask_bool = data["response_mask"].bool() + assert distillation_losses.shape == student_mass.shape == teacher_mass.shape == response_mask_bool.shape + + overlap_metrics = {} + if overlap_count is not None and overlap_token_advantage is not None: + assert overlap_count.shape == overlap_token_advantage.shape == response_mask_bool.shape + valid_overlap_count = overlap_count[response_mask_bool] + k = distillation_config.distillation_loss.topk + assert k is not None + # Diagnostics for tracking student/teacher top-k overlap in OPD, following + # "Rethinking On-Policy Distillation of Large Language Models" (arXiv:2604.13016): + # overlap fraction and mean reverse-KL advantage on overlapping tokens. + overlap_metrics["distillation/overlap_ratio"] = (valid_overlap_count.float().mean() / k).item() + overlap_position_mask = response_mask_bool & (overlap_count > 0) + if overlap_position_mask.any(): + overlap_metrics["distillation/overlap_token_advantage"] = ( + overlap_token_advantage[overlap_position_mask].mean().item() + ) + else: + overlap_metrics["distillation/overlap_token_advantage"] = 0.0 + + student_mass_valid = student_mass[response_mask_bool] + teacher_mass_valid = teacher_mass[response_mask_bool] + distillation_metrics = { + "distillation/student_mass": student_mass_valid.mean().item(), + "distillation/student_mass_min": Metric(AggregationType.MIN, student_mass_valid.min()), + "distillation/student_mass_max": Metric(AggregationType.MAX, student_mass_valid.max()), + "distillation/teacher_mass": teacher_mass_valid.mean().item(), + "distillation/teacher_mass_min": Metric(AggregationType.MIN, teacher_mass_valid.min()), + "distillation/teacher_mass_max": Metric(AggregationType.MAX, teacher_mass_valid.max()), + **overlap_metrics, + } + + # clamp small negative drift from the log_prob_min_clamp guard + distillation_losses = distillation_losses.clamp_min(0.0) + + return distillation_losses, distillation_metrics diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index 986a937fe2b..5fa6effe221 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -1,1775 +1,1929 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. 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. -""" -PPO Trainer with Ray-based single controller. -This trainer supports model-agonistic model initialization with huggingface -""" - -import json -import os -import uuid -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor -from pprint import pprint -from typing import Any, Optional - -import numpy as np -import torch -from omegaconf import OmegaConf, open_dict -from torch.utils.data import Dataset, Sampler -from torchdata.stateful_dataloader import StatefulDataLoader -from tqdm import tqdm - -from verl import DataProto -from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto -from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup, ResourcePoolManager -from verl.single_controller.ray.base import create_colocated_worker_cls -from verl.trainer.config import AlgoConfig -from verl.trainer.distillation.losses import is_distillation_enabled -from verl.trainer.ppo import core_algos -from verl.trainer.ppo.core_algos import AdvantageEstimator, agg_loss -from verl.trainer.ppo.metric_utils import ( - compute_data_metrics, - compute_throughout_metrics, - compute_timing_metrics, - compute_variance_proxy_metrics, - process_validation_metrics, -) -from verl.trainer.ppo.reward import extract_reward -from verl.trainer.ppo.utils import ( - Role, - WorkerType, - create_rl_dataset, - create_rl_sampler, - need_critic, - need_reference_policy, - need_reward_model, - need_teacher_policy, -) -from verl.utils import tensordict_utils as tu -from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, should_save_ckpt_esi -from verl.utils.config import omega_conf_to_dataclass -from verl.utils.debug import marked_timer -from verl.utils.import_utils import deprecated, load_class_from_fqn -from verl.utils.metric import reduce_metrics -from verl.utils.py_functional import rename_dict -from verl.utils.seqlen_balancing import calculate_workload, get_seqlen_balanced_partitions, log_seqlen_unbalance -from verl.utils.skip.skip_manager import SkipManager -from verl.utils.torch_functional import masked_mean -from verl.utils.tracking import ValidationGenerationsLogger -from verl.workers.config import DistillationConfig, EngineConfig -from verl.workers.rollout.llm_server import LLMServerManager -from verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding - - -def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty="kl"): - """Apply KL penalty to the token-level rewards. - - This function computes the KL divergence between the reference policy and current policy, - then applies a penalty to the token-level rewards based on this divergence. - - Args: - data (DataProto): The data containing batched model outputs and inputs. - kl_ctrl (core_algos.AdaptiveKLController): Controller for adaptive KL penalty. - kl_penalty (str, optional): Type of KL penalty to apply. Defaults to "kl". - - Returns: - tuple: A tuple containing: - - The updated data with token-level rewards adjusted by KL penalty - - A dictionary of metrics related to the KL penalty - """ - response_mask = data.batch["response_mask"] - token_level_scores = data.batch["token_level_scores"] - batch_size = data.batch.batch_size[0] - - # compute kl between ref_policy and current policy - # When apply_kl_penalty, algorithm.use_kl_in_reward=True, so the reference model has been enabled. - kld = core_algos.kl_penalty( - data.batch["old_log_probs"], data.batch["ref_log_prob"], kl_penalty=kl_penalty - ) # (batch_size, response_length) - kld = kld * response_mask - beta = kl_ctrl.value - - token_level_rewards = token_level_scores - beta * kld - - current_kl = masked_mean(kld, mask=response_mask, axis=-1) # average over sequence - current_kl = torch.mean(current_kl, dim=0).item() - - # according to https://github.com/huggingface/trl/blob/951ca1841f29114b969b57b26c7d3e80a39f75a0/trl/trainer/ppo_trainer.py#L837 - kl_ctrl.update(current_kl=current_kl, n_steps=batch_size) - data.batch["token_level_rewards"] = token_level_rewards - - metrics = {"actor/reward_kl_penalty": current_kl, "actor/reward_kl_penalty_coeff": beta} - - return data, metrics - - -def compute_response_mask(data: DataProto): - """Compute the attention mask for the response part of the sequence. - - This function extracts the portion of the attention mask that corresponds to the model's response, - which is used for masking computations that should only apply to response tokens. - - Args: - data (DataProto): The data containing batched model outputs and inputs. - - Returns: - torch.Tensor: The attention mask for the response tokens. - """ - responses = data.batch["responses"] - response_length = responses.size(1) - attention_mask = data.batch["attention_mask"] - return attention_mask[:, -response_length:] - - -def compute_spec_decode_metrics( - spec_drafts, - spec_accepts, - spec_verifies, - non_padding_mask=None, -) -> dict: - """Aggregate per-request speculative decoding stats. - - Ratios are computed per request and then averaged, so long and short - responses have equal metric weight. - - The three inputs come from the rollout engine (vLLM request spec-decode - stats or sglang ``meta_info["spec_*"]`` keys). Either all three are ``None`` - (caller didn't fetch them, e.g. spec rollout disabled) and the function - is a no-op, or all three are populated; mixed state is a programmer error. - - ``non_padding_mask`` is a numpy bool array used by sync PPO to drop padded - placeholder samples; pass ``None`` for async PPO. - """ - if spec_drafts is None and spec_accepts is None and spec_verifies is None: - return {} - assert spec_drafts is not None and spec_accepts is not None and spec_verifies is not None, ( - "spec_decode metrics require all three of spec_num_draft_tokens / " - "spec_num_accepted_tokens / spec_num_verify_steps; got partial inputs" - ) - - drafts = spec_drafts.tolist() if hasattr(spec_drafts, "tolist") else list(spec_drafts) - accepts = spec_accepts.tolist() if hasattr(spec_accepts, "tolist") else list(spec_accepts) - verifies = spec_verifies.tolist() if hasattr(spec_verifies, "tolist") else list(spec_verifies) - - if non_padding_mask is not None: - drafts = [d for d, keep in zip(drafts, non_padding_mask, strict=True) if keep] - accepts = [a for a, keep in zip(accepts, non_padding_mask, strict=True) if keep] - verifies = [v for v, keep in zip(verifies, non_padding_mask, strict=True) if keep] - - if len(drafts) == 0: - return {} - - # Treat zero-denominator samples as 0.0 and keep them in the mean. - per_sample_accept_rate = [(a / d) if d > 0 else 0.0 for a, d in zip(accepts, drafts, strict=True)] - per_sample_accept_length = [(1.0 + a / v) if v > 0 else 0.0 for a, v in zip(accepts, verifies, strict=True)] - - n = len(drafts) - return { - "rollout/spec_accept_rate": float(sum(per_sample_accept_rate) / n), - "rollout/spec_accept_length": float(sum(per_sample_accept_length) / n), - } - - -def compute_advantage( - data: DataProto, - adv_estimator: AdvantageEstimator, - gamma: float = 1.0, - lam: float = 1.0, - num_repeat: int = 1, - norm_adv_by_std_in_grpo: bool = True, - config: Optional[AlgoConfig] = None, -) -> DataProto: - """Compute advantage estimates for policy optimization. - - This function computes advantage estimates using various estimators like GAE, GRPO, REINFORCE++, etc. - The advantage estimates are used to guide policy optimization in RL algorithms. - - Args: - data (DataProto): The data containing batched model outputs and inputs. - adv_estimator (AdvantageEstimator): The advantage estimator to use (e.g., GAE, GRPO, REINFORCE++). - gamma (float, optional): Discount factor for future rewards. Defaults to 1.0. - lam (float, optional): Lambda parameter for GAE. Defaults to 1.0. - num_repeat (int, optional): Number of times to repeat the computation. Defaults to 1. - norm_adv_by_std_in_grpo (bool, optional): Whether to normalize advantages by standard deviation in - GRPO. Defaults to True. - config (dict, optional): Configuration dictionary for algorithm settings. Defaults to None. - - Returns: - DataProto: The updated data with computed advantages and returns. - """ - # Back-compatible with trainers that do not compute response mask in fit - if "response_mask" not in data.batch.keys(): - data.batch["response_mask"] = compute_response_mask(data) - # prepare response group - if adv_estimator == AdvantageEstimator.GAE: - # Compute advantages and returns using Generalized Advantage Estimation (GAE) - advantages, returns = core_algos.compute_gae_advantage_return( - token_level_rewards=data.batch["token_level_rewards"], - values=data.batch["values"], - response_mask=data.batch["response_mask"], - gamma=gamma, - lam=lam, - ) - data.batch["advantages"] = advantages - data.batch["returns"] = returns - if config.get("use_pf_ppo", False): - data = core_algos.compute_pf_ppo_reweight_data( - data, - config.pf_ppo.get("reweight_method"), - config.pf_ppo.get("weight_pow"), - ) - elif adv_estimator == AdvantageEstimator.GRPO: - # Initialize the mask for GRPO calculation - grpo_calculation_mask = data.batch["response_mask"] - - # Call compute_grpo_outcome_advantage with parameters matching its definition - advantages, returns = core_algos.compute_grpo_outcome_advantage( - token_level_rewards=data.batch["token_level_rewards"], - response_mask=grpo_calculation_mask, - index=data.non_tensor_batch["uid"], - norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo, - ) - data.batch["advantages"] = advantages - data.batch["returns"] = returns - else: - # handle all other adv estimator type other than GAE and GRPO - adv_estimator_fn = core_algos.get_adv_estimator_fn(adv_estimator) - adv_kwargs = { - "token_level_rewards": data.batch["token_level_rewards"], - "response_mask": data.batch["response_mask"], - "config": config, - } - if "uid" in data.non_tensor_batch: # optional - adv_kwargs["index"] = data.non_tensor_batch["uid"] - if "reward_baselines" in data.batch: # optional - adv_kwargs["reward_baselines"] = data.batch["reward_baselines"] - # GDPO: pass raw data for per-dimension reward extraction - if adv_estimator in (AdvantageEstimator.GDPO, "gdpo"): - adv_kwargs["non_tensor_batch"] = data.non_tensor_batch - adv_kwargs["batch"] = data.batch - # Add sum_pi_squared for Optimal Token Baseline - if adv_estimator in (AdvantageEstimator.OPTIMAL_TOKEN_BASELINE, AdvantageEstimator.TIR_OPTIMAL_TOKEN_BASELINE): - # Check if sum_pi_squared is available - assert "sum_pi_squared" in data.batch, ( - "Step-dependent optimal baseline requires sum_pi_squared from actor. " - "Please set actor.calculate_sum_pi_squared=True in config." - ) - adv_kwargs["sum_pi_squared"] = data.batch["sum_pi_squared"] - # old_log_probs needed for path-variance proxy: w_t = 1 - 2*exp(old_log_probs) + sum_pi_squared - adv_kwargs["old_log_probs"] = data.batch["old_log_probs"] - # Get pre-computed rollout IS weights if available - rollout_is_weights = data.batch.get("rollout_is_weights", None) - adv_kwargs["rollout_is_weights"] = rollout_is_weights - - # calculate advantage estimator - advantages, returns = adv_estimator_fn(**adv_kwargs) - data.batch["advantages"] = advantages - data.batch["returns"] = returns - return data - - -@deprecated("Legacy trainer is deprecated, and wil be removed in v0.9.0. Please use `trainer.use_v1=True` instead.") -class RayPPOTrainer: - """Distributed PPO trainer using Ray for scalable reinforcement learning. - - This trainer orchestrates distributed PPO training across multiple nodes and GPUs, - managing actor rollouts, critic training, and reward computation with Ray backend. - Supports various model architectures including FSDP, Megatron, vLLM, and SGLang integration. - """ - - # TODO: support each role have individual ray_worker_group_cls, - # i.e., support different backend of different role - def __init__( - self, - config, - tokenizer, - role_worker_mapping: dict[Role, WorkerType], - resource_pool_manager: ResourcePoolManager, - ray_worker_group_cls: type[RayWorkerGroup] = RayWorkerGroup, - processor=None, - train_dataset: Optional[Dataset] = None, - val_dataset: Optional[Dataset] = None, - collate_fn=None, - train_sampler: Optional[Sampler] = None, - device_name=None, - ): - """ - Initialize distributed PPO trainer with Ray backend. - Note that this trainer runs on the driver process on a single CPU/GPU node. - - Args: - config: Configuration object containing training parameters. - tokenizer: Tokenizer used for encoding and decoding text. - role_worker_mapping (dict[Role, WorkerType]): Mapping from roles to worker classes. - resource_pool_manager (ResourcePoolManager): Manager for Ray resource pools. - ray_worker_group_cls (RayWorkerGroup, optional): Class for Ray worker groups. Defaults to RayWorkerGroup. - processor: Optional data processor, used for multimodal data - train_dataset (Optional[Dataset], optional): Training dataset. Defaults to None. - val_dataset (Optional[Dataset], optional): Validation dataset. Defaults to None. - collate_fn: Function to collate data samples into batches. - train_sampler (Optional[Sampler], optional): Sampler for the training dataset. Defaults to None. - device_name (str, optional): Device name for training (e.g., "cuda", "cpu"). Defaults to None. - """ - - # Store the tokenizer for text processing - self.tokenizer = tokenizer - self.processor = processor - self.config = config - - self.hybrid_engine = config.actor_rollout_ref.hybrid_engine - assert self.hybrid_engine, "Currently, only support hybrid engine" - - if self.hybrid_engine: - assert Role.ActorRollout in role_worker_mapping or Role.ActorRolloutRef in role_worker_mapping, ( - f"{role_worker_mapping.keys()=}" - ) - - self.role_worker_mapping = role_worker_mapping - self.resource_pool_manager = resource_pool_manager - self.use_reference_policy = need_reference_policy(self.config) - self.use_teacher_policy = need_teacher_policy(self.config) - - self.use_rm = need_reward_model(self.config) - - self.use_critic = need_critic(self.config) - self.ray_worker_group_cls = ray_worker_group_cls - self.device_name = device_name if device_name else self.config.trainer.device - self.validation_generations_logger = ValidationGenerationsLogger( - project_name=self.config.trainer.project_name, - experiment_name=self.config.trainer.experiment_name, - ) - - # if ref_in_actor is True, the reference policy will be actor without lora applied - lora_rank = config.actor_rollout_ref.model.get("lora", {}).get("rank", 0) - if lora_rank <= 0: - lora_rank = config.actor_rollout_ref.model.get("lora_rank", 0) - self.ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get("lora_adapter_path") is not None - - # define in-reward KL control - # kl loss control currently not suppoorted - if self.config.algorithm.use_kl_in_reward: - self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl) - - self.use_prefix_grouper = self.config.actor_rollout_ref.actor.get("use_prefix_grouper", False) - - self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler) - - self.checkpoint_manager = None - self._init_dump_executor() - - def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampler: Optional[Sampler]): - """ - Creates the train and validation dataloaders. - """ - if train_dataset is None: - train_dataset = create_rl_dataset( - self.config.data.train_files, - self.config.data, - self.tokenizer, - self.processor, - max_samples=self.config.data.get("train_max_samples", -1), - ) - if val_dataset is None: - val_dataset = create_rl_dataset( - self.config.data.val_files, - self.config.data, - self.tokenizer, - self.processor, - max_samples=self.config.data.get("val_max_samples", -1), - ) - self.train_dataset, self.val_dataset = train_dataset, val_dataset - - if train_sampler is None: - train_sampler = create_rl_sampler(self.config.data, self.train_dataset) - if collate_fn is None: - from verl.utils.dataset.rl_dataset import collate_fn as default_collate_fn - - collate_fn = default_collate_fn - - num_workers = self.config.data["dataloader_num_workers"] - - self.train_dataloader = StatefulDataLoader( - dataset=self.train_dataset, - batch_size=self.config.data.get("gen_batch_size", self.config.data.train_batch_size), - num_workers=num_workers, - drop_last=True, - collate_fn=collate_fn, - sampler=train_sampler, - ) - - val_batch_size = self.config.data.val_batch_size # Prefer config value if set - if val_batch_size is None: - val_batch_size = len(self.val_dataset) - - self.val_dataloader = StatefulDataLoader( - dataset=self.val_dataset, - batch_size=val_batch_size, - num_workers=num_workers, - shuffle=self.config.data.get("validation_shuffle", True), - drop_last=False, - collate_fn=collate_fn, - ) - - assert len(self.train_dataloader) >= 1, "Train dataloader is empty!" - assert len(self.val_dataloader) >= 1, "Validation dataloader is empty!" - - print( - f"Size of train dataloader: {len(self.train_dataloader)}, Size of val dataloader: " - f"{len(self.val_dataloader)}" - ) - - total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs - - if self.config.trainer.total_training_steps is not None: - total_training_steps = self.config.trainer.total_training_steps - - self.total_training_steps = total_training_steps - print(f"Total training steps: {self.total_training_steps}") - - try: - OmegaConf.set_struct(self.config, True) - with open_dict(self.config): - if OmegaConf.select(self.config, "actor_rollout_ref.actor.optim"): - self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps - if OmegaConf.select(self.config, "critic.optim"): - self.config.critic.optim.total_training_steps = total_training_steps - except Exception as e: - print(f"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}") - - @staticmethod - def _write_generations(inputs, outputs, gts, scores, reward_extra_infos_dict, dump_path, global_steps): - """Write generation samples as JSONL (runs in background thread).""" - os.makedirs(dump_path, exist_ok=True) - filename = os.path.join(dump_path, f"{global_steps}.jsonl") - - n = len(inputs) - base_data = { - "input": inputs, - "output": outputs, - "gts": gts, - "score": scores, - "step": [global_steps] * n, - } - - for k, v in reward_extra_infos_dict.items(): - if len(v) == n: - base_data[k] = v - - with open(filename, "w") as f: - for i in range(n): - entry = {k: v[i] for k, v in base_data.items()} - f.write(json.dumps(entry, ensure_ascii=False, default=str) + "\n") - - print(f"Dumped generations to {filename}") - - def _dump_generations(self, inputs, outputs, gts, scores, reward_extra_infos_dict, dump_path): - """Dump rollout/validation samples as JSONL asynchronously.""" - global_steps = self.global_steps - future = self._dump_executor.submit( - self._write_generations, - inputs, - outputs, - gts, - scores, - reward_extra_infos_dict, - dump_path, - global_steps, - ) - self._dump_futures.append(future) - # Clean up completed futures and surface any exceptions early - still_pending = [] - for f in self._dump_futures: - if f.done(): - f.result() # re-raises if the write failed - else: - still_pending.append(f) - self._dump_futures = still_pending - - def _init_dump_executor(self): - """Create or recreate the dump executor and futures list.""" - self._dump_executor = ThreadPoolExecutor(max_workers=1) - self._dump_futures = [] - - def _shutdown_dump_executor(self): - """Drain pending dump futures and shut down the executor.""" - for f in self._dump_futures: - f.result() - self._dump_futures.clear() - self._dump_executor.shutdown(wait=True) - - def _log_rollout_data( - self, batch: DataProto, reward_extra_infos_dict: dict, timing_raw: dict, rollout_data_dir: str - ): - """Log rollout data to disk. - Args: - batch (DataProto): The batch containing rollout data - reward_extra_infos_dict (dict): Additional reward information to log - timing_raw (dict): Timing information for profiling - rollout_data_dir (str): Directory path to save the rollout data - """ - with marked_timer("dump_rollout_generations", timing_raw, color="green"): - inputs = self.tokenizer.batch_decode(batch.batch["prompts"], skip_special_tokens=True) - outputs = self.tokenizer.batch_decode(batch.batch["responses"], skip_special_tokens=True) - scores = batch.batch["token_level_scores"].sum(-1).cpu().tolist() - sample_gts = [item.non_tensor_batch.get("reward_model", {}).get("ground_truth", None) for item in batch] - - reward_extra_infos_to_dump = { - k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in reward_extra_infos_dict.items() - } - if "request_id" in batch.non_tensor_batch: - reward_extra_infos_to_dump.setdefault( - "request_id", - batch.non_tensor_batch["request_id"].tolist(), - ) - - self._dump_generations( - inputs=inputs, - outputs=outputs, - gts=sample_gts, - scores=scores, - reward_extra_infos_dict=reward_extra_infos_to_dump, - dump_path=rollout_data_dir, - ) - - def _maybe_log_val_generations(self, inputs, outputs, scores): - """Log a table of validation samples to the configured logger (wandb or swanlab)""" - - generations_to_log = self.config.trainer.log_val_generations - - if generations_to_log == 0: - return - - import numpy as np - - # Create tuples of (input, output, score) and sort by input text - samples = list(zip(inputs, outputs, scores, strict=True)) - samples.sort(key=lambda x: x[0]) # Sort by input text - - # Use fixed random seed for deterministic shuffling - rng = np.random.RandomState(42) - rng.shuffle(samples) - - # Take first N samples after shuffling - samples = samples[:generations_to_log] - - # Log to each configured logger - self.validation_generations_logger.log(self.config.trainer.logger, samples, self.global_steps) - - def _get_gen_batch(self, batch: DataProto) -> DataProto: - reward_keys = set({"data_source", "reward_model", "extra_info", "uid"}) & batch.non_tensor_batch.keys() - - # pop those keys for generation - batch_keys_to_pop = [] - non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys()) - reward_keys - gen_batch = batch.pop( - batch_keys=batch_keys_to_pop, - non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop), - ) - - # For agent loop, we need reward model keys to compute score. - gen_batch.non_tensor_batch.update(batch.non_tensor_batch) - - return gen_batch - - def _compute_reward_colocate(self, batch: DataProto) -> tuple[torch.Tensor, dict[str, Any]] | torch.Tensor: - """ - compute reward use colocate reward model - """ - assert self.reward_loop_manager is not None, "RewardLoopManager is None" - batch_reward = self.reward_loop_manager.compute_rm_score(batch) - return batch_reward - - def _validate(self, merged: bool = False): - data_source_lst = [] - reward_extra_infos_dict: dict[str, list] = defaultdict(list) - - # Lists to collect samples for the table - sample_inputs = [] - sample_outputs = [] - sample_gts = [] - sample_scores = [] - sample_turns = [] - sample_uids = [] - - for test_data in self.val_dataloader: - test_batch = DataProto.from_single_dict(test_data) - - if "uid" not in test_batch.non_tensor_batch: - test_batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(len(test_batch.batch))], dtype=object - ) - - # repeat test batch - test_batch = test_batch.repeat( - repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, interleave=True - ) - - ground_truths = [ - item.non_tensor_batch.get("reward_model", {}).get("ground_truth", None) for item in test_batch - ] - sample_gts.extend(ground_truths) - - test_gen_batch = self._get_gen_batch(test_batch) - test_gen_batch.meta_info = { - "eos_token_id": self.tokenizer.eos_token_id, - "pad_token_id": self.tokenizer.pad_token_id, - "recompute_log_prob": False, - "do_sample": self.config.actor_rollout_ref.rollout.val_kwargs.do_sample, - "validate": True, - "global_steps": self.global_steps, - } - print(f"test_gen_batch meta info: {test_gen_batch.meta_info}") - - # pad to be divisible by dp_size - size_divisor = self.config.actor_rollout_ref.rollout.agent.num_workers - test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, size_divisor) - test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(test_gen_batch_padded) - - if self.use_rm and "rm_scores" not in test_output_gen_batch_padded.batch.keys(): - # for colocate reward models, we need to sleep rollout model - # to spare GPU memory for reward model - self.checkpoint_manager.sleep_replicas() - batch_reward = self._compute_reward_colocate(test_output_gen_batch_padded) - test_output_gen_batch_padded = test_output_gen_batch_padded.union(batch_reward) - # wake up rollout model - # replace with wake_up method once supported - self.checkpoint_manager.update_weights(self.global_steps) - - # unpad - test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size) - - print("validation generation end") - - # Store generated outputs - output_ids = test_output_gen_batch.batch["responses"] - output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids] - sample_outputs.extend(output_texts) - - test_batch = test_batch.union(test_output_gen_batch) - test_batch.meta_info["validate"] = True - - # Store original inputs - input_ids = test_batch.batch["prompts"] - # TODO: Can we keep special tokens except for padding tokens? - input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids] - sample_inputs.extend(input_texts) - sample_uids.extend(test_batch.non_tensor_batch["uid"]) - - # evaluate using reward_function - reward_tensor, reward_extra_info = extract_reward(test_batch) - - scores = reward_tensor.sum(-1).cpu().tolist() - sample_scores.extend(scores) - - reward_extra_infos_dict["reward"].extend(scores) - for key, values in reward_extra_info.items(): - if key not in reward_extra_infos_dict: - reward_extra_infos_dict[key] = [] - if isinstance(values, np.ndarray): - reward_extra_infos_dict[key].extend(values.tolist()) - else: - reward_extra_infos_dict[key].extend(values if isinstance(values, list) else [values]) - - # collect num_turns of each prompt - if "__num_turns__" in test_batch.non_tensor_batch: - sample_turns.append(test_batch.non_tensor_batch["__num_turns__"]) - - data_source_lst.append(test_batch.non_tensor_batch.get("data_source", ["unknown"] * reward_tensor.shape[0])) - - self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores) - - # dump generations - val_data_dir = self.config.trainer.get("validation_data_dir", None) - if val_data_dir: - self._dump_generations( - inputs=sample_inputs, - outputs=sample_outputs, - gts=sample_gts, - scores=sample_scores, - reward_extra_infos_dict=reward_extra_infos_dict, - dump_path=val_data_dir, - ) - - for key_info, lst in reward_extra_infos_dict.items(): - assert len(lst) == 0 or len(lst) == len(sample_scores), f"{key_info}: {len(lst)=}, {len(sample_scores)=}" - - if merged: - print("_merge_validation_results validate result will be merged") - return { - "data_sources": data_source_lst, - "sample_uids": sample_uids, - "sample_turns": sample_turns, - "reward_extra_infos_dict": reward_extra_infos_dict, - } - data_sources = np.concatenate(data_source_lst, axis=0) - return self._val_metrics_update(data_sources, sample_uids, reward_extra_infos_dict, sample_turns) - - def _val_metrics_update(self, data_sources, sample_uids, reward_extra_infos_dict, sample_turns): - data_src2var2metric2val = process_validation_metrics(data_sources, sample_uids, reward_extra_infos_dict) - metric_dict = {} - for data_source, var2metric2val in data_src2var2metric2val.items(): - core_var = "acc" if "acc" in var2metric2val else "reward" - for var_name, metric2val in var2metric2val.items(): - n_max = max([int(name.split("@")[-1].split("/")[0]) for name in metric2val.keys()]) - for metric_name, metric_val in metric2val.items(): - if ( - (var_name == core_var) - and any(metric_name.startswith(pfx) for pfx in ["mean", "maj", "best"]) - and (f"@{n_max}" in metric_name) - ): - metric_sec = "val-core" - else: - metric_sec = "val-aux" - pfx = f"{metric_sec}/{data_source}/{var_name}/{metric_name}" - metric_dict[pfx] = metric_val - - if len(sample_turns) > 0: - sample_turns = np.concatenate(sample_turns) - metric_dict["val-aux/num_turns/min"] = sample_turns.min() - metric_dict["val-aux/num_turns/max"] = sample_turns.max() - metric_dict["val-aux/num_turns/mean"] = sample_turns.mean() - - return metric_dict - - def _merge_validation_results(self, result_a, result_b): - if result_a is None and result_b is None: - return {} - if result_a is None: - result_a = {"data_sources": [], "sample_uids": [], "sample_turns": [], "reward_extra_infos_dict": {}} - if result_b is None: - result_b = {"data_sources": [], "sample_uids": [], "sample_turns": [], "reward_extra_infos_dict": {}} - - if not result_a.get("data_sources") and not result_b.get("data_sources"): - return {} - - data_sources = np.concatenate(result_a["data_sources"] + result_b["data_sources"], axis=0) - sample_uids = result_a["sample_uids"] + result_b["sample_uids"] - sample_turns = result_a["sample_turns"] + result_b["sample_turns"] - - reward_extra_infos_dict = {} - all_keys = set(result_a["reward_extra_infos_dict"].keys()) | set(result_b["reward_extra_infos_dict"].keys()) - for key in all_keys: - list_a = result_a["reward_extra_infos_dict"].get(key, []) - list_b = result_b["reward_extra_infos_dict"].get(key, []) - reward_extra_infos_dict[key] = list_a + list_b - - return self._val_metrics_update(data_sources, sample_uids, reward_extra_infos_dict, sample_turns) - - def init_workers(self): - """Initialize distributed training workers using Ray backend. - - Creates: - 1. Ray resource pools from configuration - 2. Worker groups for each role (actor, critic, etc.) - """ - self.resource_pool_manager.create_resource_pool() - - self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()} - - # create actor and rollout - actor_role = Role.ActorRolloutRef if Role.ActorRolloutRef in self.role_worker_mapping else Role.ActorRollout - if self.hybrid_engine: - actor_rollout_resource_pool = self.resource_pool_manager.get_resource_pool(actor_role) - actor_rollout_cls = RayClassWithInitArgs( - cls=self.role_worker_mapping[actor_role], - config=self.config.actor_rollout_ref, - distillation_config=self.config.get("distillation"), - role=str(actor_role), - ) - self.resource_pool_to_cls[actor_rollout_resource_pool][str(actor_role)] = actor_rollout_cls - else: - raise NotImplementedError - - # create critic - if self.use_critic: - resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic) - - from verl.workers.config import CriticConfig - - critic_cfg: CriticConfig = omega_conf_to_dataclass(self.config.critic) - - # convert critic_cfg into TrainingWorkerConfig for the unified model engine worker - from verl.workers.engine_workers import TrainingWorkerConfig - - orig_critic_cfg = critic_cfg - engine_config: EngineConfig = orig_critic_cfg.engine - engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu - engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu - - critic_cfg = TrainingWorkerConfig( - model_type="value_model", - model_config=orig_critic_cfg.model, - engine_config=engine_config, - optimizer_config=orig_critic_cfg.optim, - checkpoint_config=orig_critic_cfg.checkpoint, - extra_context=getattr(self, "_critic_extra_context", {}), - ) - - critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg) - self.resource_pool_to_cls[resource_pool][str(Role.Critic)] = critic_cls - - # create reference policy if needed - if self.use_reference_policy and Role.RefPolicy in self.role_worker_mapping: - resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy) - ref_policy_cls = RayClassWithInitArgs( - self.role_worker_mapping[Role.RefPolicy], - config=self.config.actor_rollout_ref, - role=str(Role.RefPolicy), - ) - self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls - - # initialize WorkerGroup - # NOTE: if you want to use a different resource pool for each role, which can support different parallel size, - # you should not use `create_colocated_worker_cls`. - # Instead, directly pass different resource pool to different worker groups. - # See https://github.com/verl-project/verl/blob/master/examples/tutorial/ray/tutorial.ipynb - # for more information. - all_wg = {} - wg_kwargs = {} # Setting up kwargs for RayWorkerGroup - if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None: - wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout - if OmegaConf.select(self.config.global_profiler, "steps") is not None: - wg_kwargs["profile_steps"] = OmegaConf.select(self.config.global_profiler, "steps") - # Only require nsight worker options when tool is nsys - if OmegaConf.select(self.config.global_profiler, "tool") == "nsys": - assert ( - OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options") - is not None - ), "worker_nsight_options must be set when using nsys with profile_steps" - wg_kwargs["worker_nsight_options"] = OmegaConf.to_container( - OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options") - ) - wg_kwargs["device_name"] = self.device_name - - for resource_pool, class_dict in self.resource_pool_to_cls.items(): - if not class_dict: - continue - worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict) - wg_dict = self.ray_worker_group_cls( - resource_pool=resource_pool, - ray_cls_with_init=worker_dict_cls, - **wg_kwargs, - ) - spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys()) - all_wg.update(spawn_wg) - - if self.use_critic: - self.critic_wg = all_wg[str(Role.Critic)] - self.critic_wg.reset() - # assign critic loss - from functools import partial - - from verl.workers.utils.losses import value_loss - - value_loss_ = partial(value_loss, config=orig_critic_cfg) - self.critic_wg.set_loss_fn(value_loss_) - - if self.use_reference_policy and not self.ref_in_actor: - if str(Role.RefPolicy) in all_wg: - self.ref_policy_wg = all_wg[str(Role.RefPolicy)] - self.ref_policy_wg.init_model() - else: - # Model engine: ActorRolloutRefWorker - assert str(Role.ActorRolloutRef) in all_wg, f"{all_wg.keys()=}" - self.ref_policy_wg = all_wg[str(Role.ActorRolloutRef)] - - # we should create rollout at the end so that vllm can have a better estimation of kv cache memory - self.actor_rollout_wg = all_wg[str(actor_role)] - self.actor_rollout_wg.init_model() - - if self.ref_in_actor: - self.ref_policy_wg = self.actor_rollout_wg - - # create reward loop manager - from verl.experimental.reward_loop import RewardLoopManager - - # initalize reward loop manager - # reward model (colocate or standalone): get resource_pool - # no reward model: resource_pool = None - resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) if self.use_rm else None - self.reward_loop_manager = RewardLoopManager( - config=self.config, - rm_resource_pool=resource_pool, - ) - - # create async rollout manager and request scheduler - # Note: mode is always "async" since sync mode is deprecated - self.async_rollout_mode = True - - # initialize teacher loop manager - if self.use_teacher_policy: - from verl.experimental.teacher_loop import MultiTeacherModelManager - - teacher_resource_pool = self.resource_pool_manager.get_resource_pool(Role.TeacherModel) - self.teacher_model_manager = MultiTeacherModelManager( - config=self.config, - resource_pool=teacher_resource_pool, - ) - self.distillation_config: DistillationConfig = omega_conf_to_dataclass(self.config.distillation) - else: - self.teacher_model_manager = None - self.distillation_config = None - - # Support custom AgentLoopManager via config - manager_class_fqn = self.config.actor_rollout_ref.rollout.get("agent", {}).get("agent_loop_manager_class") - if manager_class_fqn: - AgentLoopManager = load_class_from_fqn(manager_class_fqn, "AgentLoopManager") - else: - from verl.experimental.agent_loop import AgentLoopManager - - # infrastructure overview: https://verl.readthedocs.io/en/latest/advance/reward_loop.html#architecture-design - # agent_reward_loop: streaming reward computation with actor rollout - # two conditions satisfied: (1) no reward model, or (2) reward model with extra resource pool - enable_agent_reward_loop = not self.use_rm or self.config.reward.reward_model.enable_resource_pool - - self.llm_server_manager = LLMServerManager.create( - config=self.config, worker_group=self.actor_rollout_wg, rollout_resource_pool=actor_rollout_resource_pool - ) - - # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager - # to stream reward computation with actor rollout - # To stream teacher computation with actor rollout, we instead pass the full manager so that the - # teacher loop workers can sleep/wake together with rollout workers - reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None - self.async_rollout_manager = AgentLoopManager.create( - config=self.config, - llm_client=self.llm_server_manager.get_client(), - teacher_client=self.teacher_model_manager.get_client() if self.use_teacher_policy else None, - reward_loop_worker_handles=reward_loop_worker_handles, - ) - - checkpoint_engine_config = omega_conf_to_dataclass(self.config.actor_rollout_ref.rollout.checkpoint_engine) - # Support custom CheckpointEngineManager via config - checkpoint_manager_class_fqn = self.config.actor_rollout_ref.rollout.get("checkpoint_manager_class") - if checkpoint_manager_class_fqn: - CheckpointEngineManager = load_class_from_fqn(checkpoint_manager_class_fqn, "CheckpointEngineManager") - else: - from verl.checkpoint_engine import CheckpointEngineManager - self.checkpoint_manager = CheckpointEngineManager( - config=checkpoint_engine_config, - actor_wg=self.actor_rollout_wg, - replicas=self.llm_server_manager.get_replicas(), - ) - - # sleep all replicas to load checkpoint - self.checkpoint_manager.sleep_replicas() - - def _save_checkpoint(self): - from verl.utils.fs import local_mkdir_safe - - # path: given_path + `/global_step_{global_steps}` + `/actor` - local_global_step_folder = os.path.join( - self.config.trainer.default_local_dir, f"global_step_{self.global_steps}" - ) - - print(f"local_global_step_folder: {local_global_step_folder}") - actor_local_path = os.path.join(local_global_step_folder, "actor") - - actor_remote_path = ( - None - if self.config.trainer.default_hdfs_dir is None - else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "actor") - ) - - remove_previous_ckpt_in_save = self.config.trainer.get("remove_previous_ckpt_in_save", False) - if remove_previous_ckpt_in_save: - print( - "Warning: remove_previous_ckpt_in_save is deprecated," - + " set max_actor_ckpt_to_keep=1 and max_critic_ckpt_to_keep=1 instead" - ) - max_actor_ckpt_to_keep = ( - self.config.trainer.get("max_actor_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1 - ) - max_critic_ckpt_to_keep = ( - self.config.trainer.get("max_critic_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1 - ) - - self.actor_rollout_wg.save_checkpoint( - actor_local_path, actor_remote_path, self.global_steps, max_ckpt_to_keep=max_actor_ckpt_to_keep - ) - - if self.use_critic: - critic_local_path = os.path.join(local_global_step_folder, str(Role.Critic)) - critic_remote_path = ( - None - if self.config.trainer.default_hdfs_dir is None - else os.path.join( - self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", str(Role.Critic) - ) - ) - self.critic_wg.save_checkpoint( - critic_local_path, critic_remote_path, self.global_steps, max_ckpt_to_keep=max_critic_ckpt_to_keep - ) - - # save dataloader - local_mkdir_safe(local_global_step_folder) - dataloader_local_path = os.path.join(local_global_step_folder, "data.pt") - dataloader_state_dict = self.train_dataloader.state_dict() - torch.save(dataloader_state_dict, dataloader_local_path) - - # latest checkpointed iteration tracker (for atomic usage) - if ( - hasattr(self.config.actor_rollout_ref.actor.checkpoint, "async_save") - and self.config.actor_rollout_ref.actor.checkpoint.async_save - ) or ( - "async_save" in self.config.actor_rollout_ref.actor.checkpoint - and self.config.actor_rollout_ref.actor.checkpoint["async_save"] - ): - print("skip write latest_checkpointed_iteration.txt when async_save is True") - return - local_latest_checkpointed_iteration = os.path.join( - self.config.trainer.default_local_dir, "latest_checkpointed_iteration.txt" - ) - with open(local_latest_checkpointed_iteration, "w") as f: - f.write(str(self.global_steps)) - - def _load_checkpoint(self): - if self.config.trainer.resume_mode == "disable": - return 0 - - # load from hdfs - if self.config.trainer.default_hdfs_dir is not None: - raise NotImplementedError("load from hdfs is not implemented yet") - else: - checkpoint_folder = self.config.trainer.default_local_dir # TODO: check path - if not os.path.isabs(checkpoint_folder): - working_dir = os.getcwd() - checkpoint_folder = os.path.join(working_dir, checkpoint_folder) - global_step_folder = find_latest_ckpt_path(checkpoint_folder) # None if no latest - - # find global_step_folder - if self.config.trainer.resume_mode == "auto": - if global_step_folder is None: - print("Training from scratch") - return 0 - else: - if self.config.trainer.resume_mode == "resume_path": - assert isinstance(self.config.trainer.resume_from_path, str), "resume ckpt must be str type" - assert "global_step_" in self.config.trainer.resume_from_path, ( - "resume ckpt must specify the global_steps" - ) - global_step_folder = self.config.trainer.resume_from_path - if not os.path.isabs(global_step_folder): - working_dir = os.getcwd() - global_step_folder = os.path.join(working_dir, global_step_folder) - print(f"Load from checkpoint folder: {global_step_folder}") - # set global step - self.global_steps = int(global_step_folder.split("global_step_")[-1]) - - print(f"Setting global step to {self.global_steps}") - print(f"Resuming from {global_step_folder}") - - actor_path = os.path.join(global_step_folder, "actor") - critic_path = os.path.join(global_step_folder, str(Role.Critic)) - # load actor - self.actor_rollout_wg.load_checkpoint( - actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load - ) - # load critic - if self.use_critic: - self.critic_wg.load_checkpoint( - critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load - ) - - # load dataloader, - # TODO: from remote not implemented yet - dataloader_local_path = os.path.join(global_step_folder, "data.pt") - if os.path.exists(dataloader_local_path): - steps_per_epoch = len(self.train_dataloader) - at_epoch_boundary = steps_per_epoch > 0 and self.global_steps % steps_per_epoch == 0 - if at_epoch_boundary: - print( - f"Skipping dataloader state restore: global_steps={self.global_steps} " - f"is at an epoch boundary (steps_per_epoch={steps_per_epoch}). " - f"The saved state marks the dataloader as exhausted. " - f"Next epoch will iterate from scratch." - ) - else: - dataloader_state_dict = torch.load(dataloader_local_path, weights_only=False) - self.train_dataloader.load_state_dict(dataloader_state_dict) - else: - print(f"Warning: No dataloader state found at {dataloader_local_path}, will start from scratch") - - def _start_profiling(self, do_profile: bool) -> None: - """Start profiling for all worker groups if profiling is enabled.""" - if do_profile: - self.actor_rollout_wg.start_profile(role="e2e", profile_step=self.global_steps) - if self.use_reference_policy: - self.ref_policy_wg.start_profile(profile_step=self.global_steps) - if self.use_critic: - self.critic_wg.start_profile(profile_step=self.global_steps) - - def _stop_profiling(self, do_profile: bool) -> None: - """Stop profiling for all worker groups if profiling is enabled.""" - if do_profile: - self.actor_rollout_wg.stop_profile() - if self.use_reference_policy: - self.ref_policy_wg.stop_profile() - if self.use_critic: - self.critic_wg.stop_profile() - - def _get_dp_size(self, worker_group, role: str) -> int: - """Get data parallel size from worker group dispatch info. - - This method retrieves the data parallel size by querying the dispatch info - for the specified role. The dispatch info is cached for subsequent calls. - - Args: - worker_group: The worker group to query dispatch info from. - role: The role name (e.g., "actor", "critic") to get DP size for. - - Returns: - The data parallel size (number of DP ranks). - """ - if role not in worker_group._dispatch_info: - dp_rank_mapping = worker_group._query_dispatch_info(role) - worker_group._dispatch_info[role] = dp_rank_mapping - else: - dp_rank_mapping = worker_group._dispatch_info[role] - return max(dp_rank_mapping) + 1 - - def _balance_batch(self, batch: DataProto, metrics, logging_prefix="global_seqlen", keep_minibatch=False): - """Reorder the data on single controller such that each dp rank gets similar total tokens. - - When use_prefix_grouper is enabled, uses group-level balancing to keep samples with - the same uid together on the same rank for prefix sharing optimization. - """ - attention_mask = batch.batch["attention_mask"] - batch_size = attention_mask.shape[0] - global_seqlen_lst = batch.batch["attention_mask"].view(batch_size, -1).sum(-1) # (train_batch_size,) - workload_lst = calculate_workload(global_seqlen_lst) - # Get dp_size from dispatch info to correctly balance across data parallel ranks - # Note: world_size may include tensor/pipeline parallel dimensions, but we only want DP - dp_size = self._get_dp_size(self.actor_rollout_wg, "actor") - - # Use group-level balancing for PrefixGrouper to keep same-uid samples together - if getattr(self, "use_prefix_grouper", False) and "uid" in batch.non_tensor_batch: - from verl.utils.seqlen_balancing import get_group_balanced_partitions - - uid_list = list(batch.non_tensor_batch["uid"]) - seqlen_list = global_seqlen_lst.tolist() - - # Count number of uid groups - num_groups = len(set(uid_list)) - - if num_groups % dp_size != 0: - raise ValueError( - f"PrefixGrouper with balance_batch requires num_uid_groups ({num_groups}) " - f"% dp_size ({dp_size}) == 0. " - f"This ensures each rank gets equal number of groups. " - f"Current batch_size={batch_size}, adjust batch_size to be a multiple of " - f"dp_size * rollout.n." - ) - - global_partition_lst = get_group_balanced_partitions( - seqlen_list=seqlen_list, - uid_list=uid_list, - k_partitions=dp_size, - ) - - elif keep_minibatch: - # Decouple the DP balancing and mini-batching. - minibatch_size = self.config.actor_rollout_ref.actor.get("ppo_mini_batch_size") - minibatch_num = len(workload_lst) // minibatch_size - global_partition_lst = [[] for _ in range(dp_size)] - for i in range(minibatch_num): - rearrange_minibatch_lst = get_seqlen_balanced_partitions( - workload_lst[i * minibatch_size : (i + 1) * minibatch_size], - k_partitions=dp_size, - equal_size=True, - ) - for j, part in enumerate(rearrange_minibatch_lst): - global_partition_lst[j].extend([x + minibatch_size * i for x in part]) - else: - global_partition_lst = get_seqlen_balanced_partitions(workload_lst, k_partitions=dp_size, equal_size=True) - # Place smaller micro-batches at both ends to reduce the bubbles in pipeline parallel. - # Skip reordering within partitions for PrefixGrouper to maintain uid grouping - if not getattr(self, "use_prefix_grouper", False): - for idx, partition in enumerate(global_partition_lst): - partition.sort(key=lambda x: (workload_lst[x], x)) - ordered_partition = partition[::2] + partition[1::2][::-1] - global_partition_lst[idx] = ordered_partition - - # reorder based on index. The data will be automatically equally partitioned by dispatch function - global_idx = torch.tensor([j for partition in global_partition_lst for j in partition]) - batch.reorder(global_idx) - global_balance_stats = log_seqlen_unbalance( - seqlen_list=global_seqlen_lst.tolist(), partitions=global_partition_lst, prefix=logging_prefix - ) - metrics.update(global_balance_stats) - - def _compute_values(self, batch: DataProto) -> DataProto: - batch_td = batch.to_tensordict() - # step 2: convert from padding to nopadding - batch_td = left_right_2_no_padding(batch_td) - # step 3: add meta info - tu.assign_non_tensor(batch_td, compute_loss=False) - output = self.critic_wg.infer_batch(batch_td) - output = output.get() - values = tu.get(output, "values") - values = no_padding_2_padding(values, batch_td) - values = tu.get_tensordict({"values": values.float()}) - values = DataProto.from_tensordict(values) - return values - - def _compute_ref_log_prob(self, batch: DataProto) -> DataProto: - # step 1: convert dataproto to tensordict. - batch_td = batch.to_tensordict() - # step 2: convert from padding to nopadding - batch_td = left_right_2_no_padding(batch_td) - # step 3: add meta info - metadata = {"calculate_entropy": False, "compute_loss": False} - if self.ref_in_actor: - metadata["no_lora_adapter"] = True - tu.assign_non_tensor(batch_td, **metadata) - if self.ref_in_actor: - output = self.actor_rollout_wg.compute_log_prob(batch_td) - else: - output = self.ref_policy_wg.compute_ref_log_prob(batch_td) - # gather output - log_probs = tu.get(output, "log_probs") - # step 4. No padding to padding - log_probs = no_padding_2_padding(log_probs, batch_td) - # step 5: rebuild a tensordict and convert to dataproto - ref_log_prob = tu.get_tensordict({"ref_log_prob": log_probs.float()}) - ref_log_prob = DataProto.from_tensordict(ref_log_prob) - - return ref_log_prob - - def _compute_old_log_prob(self, batch: DataProto): - # TODO: remove step 1, 2, 4 after we make the whole training tensordict and padding free - # step 1: convert dataproto to tensordict. - batch_td = batch.to_tensordict() - # step 2: convert from padding to nopadding - batch_td = left_right_2_no_padding(batch_td) - # step 3: add meta info - calculate_sum_pi_squared = self.config.actor_rollout_ref.actor.get("calculate_sum_pi_squared", False) - tu.assign_non_tensor( - batch_td, - calculate_entropy=True, - calculate_sum_pi_squared=calculate_sum_pi_squared, - compute_loss=False, - ) - output = self.actor_rollout_wg.compute_log_prob(batch_td) - # gather output - entropy = tu.get(output, "entropy") - log_probs = tu.get(output, "log_probs") - routed_experts = tu.get(output, "routed_experts") - sum_pi_squared = tu.get(output, "sum_pi_squared") if calculate_sum_pi_squared else None - - old_log_prob_mfu = tu.get(output, "metrics")["mfu"] - # step 4. No padding to padding - entropy = no_padding_2_padding(entropy, batch_td) - log_probs = no_padding_2_padding(log_probs, batch_td) - if sum_pi_squared is not None: - sum_pi_squared = no_padding_2_padding(sum_pi_squared, batch_td) - # step 5: rebuild a tensordict and convert to dataproto - result = {"old_log_probs": log_probs.float(), "entropys": entropy.float()} - if routed_experts is not None: - result["routed_experts"] = routed_experts - if sum_pi_squared is not None: - result["sum_pi_squared"] = sum_pi_squared.float() - old_log_prob = tu.get_tensordict(result) - old_log_prob = DataProto.from_tensordict(old_log_prob) - return old_log_prob, old_log_prob_mfu - - def _update_actor(self, batch: DataProto) -> DataProto: - rollout_config = self.config.actor_rollout_ref.rollout - batch.meta_info["multi_turn"] = rollout_config.multi_turn.enable - # TODO: Make "temperature" single source of truth from generation. - batch.meta_info["temperature"] = rollout_config.temperature - # update actor - batch_td = batch.to_tensordict() - # step 2: convert from padding to no-padding - batch_td = left_right_2_no_padding(batch_td) - calculate_entropy = self.config.actor_rollout_ref.actor.calculate_entropy or ( - self.config.actor_rollout_ref.actor.entropy_coeff != 0.0 - ) - distillation_use_topk = ( - self.distillation_config.distillation_loss.loss_settings.use_topk - if is_distillation_enabled(self.config.get("distillation")) - else False - ) - distillation_only = False # distillation_only flag means we can skip policy loss and reduce mem footprint - if is_distillation_enabled(self.config.get("distillation")): - distillation_loss_cfg = self.distillation_config.distillation_loss - distillation_only = ( - distillation_use_topk - and not distillation_loss_cfg.use_task_rewards - and not distillation_loss_cfg.use_policy_gradient - ) - ppo_mini_batch_size = self.config.actor_rollout_ref.actor.ppo_mini_batch_size - ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n - ppo_epochs = self.config.actor_rollout_ref.actor.ppo_epochs - seed = self.config.actor_rollout_ref.actor.data_loader_seed - shuffle = self.config.actor_rollout_ref.actor.shuffle - tu.assign_non_tensor( - batch_td, - calculate_entropy=calculate_entropy, - distillation_use_topk=distillation_use_topk, - distillation_only=distillation_only, - global_batch_size=ppo_mini_batch_size, - mini_batch_size=ppo_mini_batch_size, - epochs=ppo_epochs, - seed=seed, - dataloader_kwargs={"shuffle": shuffle}, - compute_loss=True, - ) - actor_output = self.actor_rollout_wg.update_actor(batch_td) - actor_output = tu.get(actor_output, "metrics") - actor_output = rename_dict(actor_output, "actor/") - # modify key name - actor_output["perf/mfu/actor"] = actor_output.pop("actor/mfu") - actor_output = DataProto.from_single_dict(data={}, meta_info={"metrics": actor_output}) - - return actor_output - - def _update_critic(self, batch: DataProto) -> DataProto: - batch_td = batch.to_tensordict() - # step 2: convert from padding to no-padding - batch_td = left_right_2_no_padding(batch_td) - ppo_mini_batch_size = self.config.critic.ppo_mini_batch_size - ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n - ppo_epochs = self.config.critic.ppo_epochs - seed = self.config.critic.data_loader_seed - shuffle = self.config.critic.shuffle - tu.assign_non_tensor( - batch_td, - global_batch_size=ppo_mini_batch_size, - mini_batch_size=ppo_mini_batch_size, - epochs=ppo_epochs, - seed=seed, - dataloader_kwargs={"shuffle": shuffle}, - ) - - output = self.critic_wg.train_mini_batch(batch_td) - output = output.get() - output = tu.get(output, "metrics") - output = rename_dict(output, "critic/") - # modify key name - output["perf/mfu/critic"] = output.pop("critic/mfu") - critic_output = DataProto.from_single_dict(data={}, meta_info={"metrics": output}) - return critic_output - - def fit(self): - """ - The training loop of PPO. - The driver process only need to call the compute functions of the worker group through RPC - to construct the PPO dataflow. - The light-weight advantage computation is done on the driver process. - """ - if self._dump_executor._shutdown: - self._init_dump_executor() - - from omegaconf import OmegaConf - - from verl.utils.tracking import Tracking - - logger = Tracking( - project_name=self.config.trainer.project_name, - experiment_name=self.config.trainer.experiment_name, - default_backend=self.config.trainer.logger, - config=OmegaConf.to_container(self.config, resolve=True), - ) - - self.global_steps = 0 - - # load checkpoint and update weights before doing anything - self._load_checkpoint() - self.checkpoint_manager.update_weights(self.global_steps) - - current_epoch = self.global_steps // len(self.train_dataloader) - - SkipManager.init(self.config) - - # perform validation before training - # currently, we only support validation using the reward_function. - if self.config.trainer.get("val_before_train", True): - val_metrics = self._validate() - assert val_metrics, f"{val_metrics=}" - pprint(f"Initial validation metrics: {val_metrics}") - logger.log(data=val_metrics, step=self.global_steps) - if self.config.trainer.get("val_only", False): - self._shutdown_dump_executor() - return - - # add tqdm - progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress") - - # we start from step 1 - self.global_steps += 1 - last_val_metrics = None - self.max_steps_duration = 0 - - SkipManager.set_step(self.global_steps) - - prev_step_profile = False - curr_step_profile = ( - self.global_steps in self.config.global_profiler.steps - if self.config.global_profiler.steps is not None - else False - ) - next_step_profile = False - - for epoch in range(current_epoch, self.config.trainer.total_epochs): - for batch_dict in self.train_dataloader: - if hasattr(self.actor_rollout_wg, "async_calls_finalize_fn_exec"): - self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=False) - metrics = {} - timing_raw = {} - - with marked_timer("start_profile", timing_raw): - self._start_profiling( - not prev_step_profile and curr_step_profile - if self.config.global_profiler.profile_continuous_steps - else curr_step_profile - ) - batch: DataProto = DataProto.from_single_dict(batch_dict) - batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature - - # add uid to batch - batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object - ) - - gen_batch = self._get_gen_batch(batch) - - # pass global_steps to trace - gen_batch.meta_info["global_steps"] = self.global_steps - rollout_n = self.config.actor_rollout_ref.rollout.n - gen_batch_output = gen_batch.repeat(repeat_times=rollout_n, interleave=True) - - if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX: - # NOTE: REMAX needs one sampled rollout plus one greedy baseline per prompt. - # Keep them in a single agent-loop/vLLM request to avoid sending a second - # rollout after replicas have been put to sleep, which can leave async vLLM - # engines in an invalid state for multi-turn agent workloads. - gen_batch_output.non_tensor_batch["__do_sample__"] = np.ones(len(gen_batch_output), dtype=bool) - gen_baseline_batch = gen_batch.slice(0, None) - gen_baseline_batch.non_tensor_batch["__do_sample__"] = np.zeros(len(gen_baseline_batch), dtype=bool) - combined_gen_batch = DataProto.concat([gen_batch_output, gen_baseline_batch]) - num_sampled_prompts = len(gen_batch_output) - else: - combined_gen_batch = gen_batch_output - num_sampled_prompts = len(gen_batch_output) - - is_last_step = self.global_steps >= self.total_training_steps - with marked_timer("step", timing_raw): - # generate a batch - with marked_timer("gen", timing_raw, color="red"): - if curr_step_profile: - self.llm_server_manager.start_profile() - combined_gen_output = self.async_rollout_manager.generate_sequences(combined_gen_batch) - self.checkpoint_manager.sleep_replicas() - if curr_step_profile: - self.llm_server_manager.stop_profile() - - timing_raw.update(combined_gen_output.meta_info["timing"]) - combined_gen_output.meta_info.pop("timing", None) - - gen_batch_output = combined_gen_output.slice(0, num_sampled_prompts) - if "__do_sample__" in gen_batch_output.non_tensor_batch: - gen_batch_output.pop(non_tensor_batch_keys=["__do_sample__"]) - - if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX: - gen_baseline_output = combined_gen_output.slice(num_sampled_prompts, None) - if "__do_sample__" in gen_baseline_output.non_tensor_batch: - gen_baseline_output.pop(non_tensor_batch_keys=["__do_sample__"]) - - if self.use_rm and "rm_scores" not in gen_baseline_output.batch.keys(): - baseline_reward = self._compute_reward_colocate(gen_baseline_output) - gen_baseline_output = gen_baseline_output.union(baseline_reward) - - reward_baseline_tensor = gen_baseline_output.batch["rm_scores"].sum(dim=-1) - batch.batch["reward_baselines"] = reward_baseline_tensor - - del gen_baseline_output - del combined_gen_batch, combined_gen_output - # repeat to align with repeated responses in rollout - batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) - batch = batch.union(gen_batch_output) - - if "response_mask" not in batch.batch.keys(): - batch.batch["response_mask"] = compute_response_mask(batch) - # Balance the number of valid tokens across DP ranks. - # NOTE: This usually changes the order of data in the `batch`, - # which won't affect the advantage calculation (since it's based on uid), - # but might affect the loss calculation (due to the change of mini-batching). - if self.config.trainer.balance_batch: - self._balance_batch(batch, metrics=metrics) - - # compute global_valid tokens - batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist() - # get images_seqlens - images_seqlens_all = [] - for multi_modal_input in batch.non_tensor_batch["multi_modal_inputs"]: - if "image_grid_thw" not in multi_modal_input.keys(): - continue - images_seqlens_all.extend(multi_modal_input["images_seqlens"].tolist()) - batch.meta_info["images_seqlens"] = images_seqlens_all - with marked_timer("reward", timing_raw, color="yellow"): - # compute reward model score - if self.use_rm and "rm_scores" not in batch.batch.keys(): - batch_reward = self._compute_reward_colocate(batch) - batch = batch.union(batch_reward) - - # extract reward_tensor and reward_extra_infos_dict for training - reward_tensor, reward_extra_infos_dict = extract_reward(batch) - - # Operating Mode Selection: - # - Bypass mode: Sets old_log_probs = rollout_log_probs (2 policies: π_rollout, π_θ) - # - Decoupled mode: Recomputes old_log_probs as proximal anchor (3 policies: π_rollout, π_old, π_θ) - # Note: π_old computed once per data batch, serves as stable reference during mini-batch updates - rollout_corr_config = self.config.algorithm.get("rollout_correction", None) - bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get("bypass_mode", False) - if bypass_recomputing_logprobs: # Use `rollout_log_probs` - from verl.trainer.ppo.rollout_corr_helper import apply_bypass_mode - - apply_bypass_mode( - batch=batch, - rollout_corr_config=rollout_corr_config, - policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss, - ) - else: # Recompute old_log_probs - with marked_timer("old_log_prob", timing_raw, color="blue"): - old_log_prob, old_log_prob_mfu = self._compute_old_log_prob(batch) - entropys = old_log_prob.batch["entropys"] - response_masks = batch.batch["response_mask"] - actor_config = self.config.actor_rollout_ref.actor - entropy_agg = agg_loss( - loss_mat=entropys, - loss_mask=response_masks, - loss_agg_mode=actor_config.loss_agg_mode, - loss_scale_factor=actor_config.loss_scale_factor, - ) - old_log_prob_metrics = { - "actor/entropy": entropy_agg.detach().item(), - "perf/mfu/actor_infer": old_log_prob_mfu, - } - metrics.update(old_log_prob_metrics) - old_log_prob.batch.pop("entropys") - if "routed_experts" in batch.batch and "routed_experts" in old_log_prob.batch: - raise ValueError( - "Detected conflicting router replay configuration: " - "router_replay.mode='R2' and enable_rollout_routing_replay=True " - "cannot be enabled simultaneously. " - "The enable_rollout_routing_replay option is only used in R3 mode; " - "it should not be set when using R2 mode." - ) - batch = batch.union(old_log_prob) - if "rollout_log_probs" in batch.batch.keys(): - # TODO: we may want to add diff of probs too. - from verl.utils.debug.metrics import calculate_debug_metrics - - metrics.update(calculate_debug_metrics(batch)) - - assert "old_log_probs" in batch.batch, f'"old_log_prob" not in {batch.batch.keys()=}' - if self.use_reference_policy: - # compute reference log_prob - with marked_timer(str(Role.RefPolicy), timing_raw, color="olive"): - ref_log_prob = self._compute_ref_log_prob(batch) - batch = batch.union(ref_log_prob) - - # compute values - if self.use_critic: - with marked_timer("values", timing_raw, color="cyan"): - values = self._compute_values(batch) - batch = batch.union(values) - - with marked_timer("adv", timing_raw, color="brown"): - # we combine with rule-based rm - reward_extra_infos_dict: dict[str, list] - batch.batch["token_level_scores"] = reward_tensor - - if reward_extra_infos_dict: - batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()}) - - # compute rewards. apply_kl_penalty if available - if self.config.algorithm.use_kl_in_reward: - batch, kl_metrics = apply_kl_penalty( - batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty - ) - metrics.update(kl_metrics) - else: - batch.batch["token_level_rewards"] = batch.batch["token_level_scores"] - - # Compute rollout correction: IS weights, rejection sampling, and metrics - # Only runs in decoupled mode (computes once per batch using stable π_old) - # In bypass mode, this is skipped - actor computes metrics from evolving π_θ vs π_rollout - if ( - rollout_corr_config is not None - and "rollout_log_probs" in batch.batch - and not bypass_recomputing_logprobs # Only in decoupled mode - ): - from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_add_to_batch - - # Compute IS weights, apply rejection sampling, compute metrics - batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config) - # IS and off-policy metrics already have rollout_corr/ prefix - metrics.update(is_metrics) - - # compute advantages, executed on the driver process - norm_adv_by_std_in_grpo = self.config.algorithm.get( - "norm_adv_by_std_in_grpo", True - ) # GRPO adv normalization factor - - batch = compute_advantage( - batch, - adv_estimator=self.config.algorithm.adv_estimator, - gamma=self.config.algorithm.gamma, - lam=self.config.algorithm.lam, - num_repeat=self.config.actor_rollout_ref.rollout.n, - norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo, - config=self.config.algorithm, - ) - # update critic - if self.use_critic: - with marked_timer("update_critic", timing_raw, color="pink"): - critic_output = self._update_critic(batch) - critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"]) - metrics.update(critic_output_metrics) - - # implement critic warmup - if self.config.trainer.critic_warmup > self.global_steps: - # Still in critic warmup, only update weights to wake up rollout replicas. - self.checkpoint_manager.update_weights(self.global_steps) - else: - # update actor - with marked_timer("update_actor", timing_raw, color="red"): - actor_output = self._update_actor(batch) - - # Check if the ESI (Elastic Server Instance)/training plan is close to expiration. - esi_close_to_expiration = should_save_ckpt_esi( - max_steps_duration=self.max_steps_duration, - redundant_time=self.config.trainer.esi_redundant_time, - ) - # Check if the conditions for saving a checkpoint are met. - # The conditions include a mandatory condition (1) and - # one of the following optional conditions (2/3/4): - # 1. The save frequency is set to a positive value. - # 2. It's the last training step. - # 3. The current step number is a multiple of the save frequency. - # 4. The ESI(Elastic Server Instance)/training plan is close to expiration. - if self.config.trainer.save_freq > 0 and ( - is_last_step - or self.global_steps % self.config.trainer.save_freq == 0 - or esi_close_to_expiration - ): - if esi_close_to_expiration: - print("Force saving checkpoint: ESI instance expiration approaching.") - with marked_timer("save_checkpoint", timing_raw, color="green"): - self._save_checkpoint() - - # update weights from trainer to rollout - with marked_timer("update_weights", timing_raw, color="red"): - self.checkpoint_manager.update_weights(self.global_steps) - - actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) - metrics.update(actor_output_metrics) - - # Log rollout generations if enabled - rollout_data_dir = self.config.trainer.get("rollout_data_dir", None) - if rollout_data_dir: - self._log_rollout_data(batch, reward_extra_infos_dict, timing_raw, rollout_data_dir) - - # validate - if self.config.trainer.test_freq > 0 and ( - is_last_step or self.global_steps % self.config.trainer.test_freq == 0 - ): - with marked_timer("testing", timing_raw, color="green"): - val_metrics: dict = self._validate() - if is_last_step: - last_val_metrics = val_metrics - metrics.update(val_metrics) - - with marked_timer("stop_profile", timing_raw): - next_step_profile = ( - self.global_steps + 1 in self.config.global_profiler.steps - if self.config.global_profiler.steps is not None - else False - ) - self._stop_profiling( - curr_step_profile and not next_step_profile - if self.config.global_profiler.profile_continuous_steps - else curr_step_profile - ) - prev_step_profile = curr_step_profile - curr_step_profile = next_step_profile - - steps_duration = timing_raw["step"] - self.max_steps_duration = max(self.max_steps_duration, steps_duration) - - # training metrics - metrics.update( - { - "training/global_step": self.global_steps, - "training/epoch": epoch, - } - ) - # collect metrics - metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) - # GDPO per-component reward metrics - gdpo_reward_keys = self.config.algorithm.get("gdpo_reward_keys", None) - if gdpo_reward_keys and self.config.algorithm.adv_estimator in ("gdpo", AdvantageEstimator.GDPO): - for key in gdpo_reward_keys: - if key in batch.non_tensor_batch: - vals = np.asarray(batch.non_tensor_batch[key], dtype=np.float32) - metrics[f"gdpo/{key}/mean"] = float(np.mean(vals)) - metrics[f"gdpo/{key}/std"] = float(np.std(vals)) - metrics[f"gdpo/{key}/max"] = float(np.max(vals)) - metrics[f"gdpo/{key}/min"] = float(np.min(vals)) - metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) - # TODO: implement actual tflpo and theoretical tflpo - n_gpus = self.resource_pool_manager.get_n_gpus() - metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus)) - # compute variance proxy metrics - gradient_norm = metrics.get("actor/grad_norm", None) - metrics.update(compute_variance_proxy_metrics(batch=batch, gradient_norm=gradient_norm)) - # Note: mismatch metrics (KL, PPL, etc.) are collected at line 1179 after advantage computation - - # Per-request spec decode metrics. - metrics.update( - compute_spec_decode_metrics( - batch.non_tensor_batch.get("spec_num_draft_tokens", None), - batch.non_tensor_batch.get("spec_num_accepted_tokens", None), - batch.non_tensor_batch.get("spec_num_verify_steps", None), - ) - ) - - # TODO: make a canonical logger that supports various backend - logger.log(data=metrics, step=self.global_steps) - - progress_bar.update(1) - self.global_steps += 1 - SkipManager.set_step(self.global_steps) - - if is_last_step: - if hasattr(self.actor_rollout_wg, "async_calls_finalize_fn_exec"): - self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=True) - self._shutdown_dump_executor() - pprint(f"Final validation metrics: {last_val_metrics}") - progress_bar.close() - return - - # this is experimental and may be changed/removed in the future - # in favor of a general-purpose data buffer pool - if hasattr(self.train_dataset, "on_batch_end"): - # The dataset may be changed after each training batch - self.train_dataset.on_batch_end(batch=batch) - - # Ensure dump executor is shut down when training loop ends without reaching is_last_step - self._shutdown_dump_executor() +# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2024 SGLang Team +# Copyright 2025 ModelBest Inc. 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. +""" +PPO Trainer with Ray-based single controller. +This trainer supports model-agonistic model initialization with huggingface +""" + +import json +import os +import uuid +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from pprint import pprint +from typing import Any, Optional + +import numpy as np +import torch +from omegaconf import OmegaConf, open_dict +from torch.utils.data import Dataset, Sampler +from torchdata.stateful_dataloader import StatefulDataLoader +from tqdm import tqdm + +from verl import DataProto +from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto +from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup, ResourcePoolManager +from verl.single_controller.ray.base import create_colocated_worker_cls +from verl.trainer.config import AlgoConfig +from verl.trainer.distillation.losses import is_distillation_enabled +from verl.trainer.ppo import core_algos +from verl.trainer.ppo.core_algos import AdvantageEstimator, agg_loss +from verl.trainer.ppo.metric_utils import ( + compute_data_metrics, + compute_throughout_metrics, + compute_timing_metrics, + compute_variance_proxy_metrics, + process_validation_metrics, +) +from verl.trainer.ppo.reward import extract_reward +from verl.trainer.ppo.utils import ( + Role, + WorkerType, + create_rl_dataset, + create_rl_sampler, + need_critic, + need_reference_policy, + need_reward_model, + need_teacher_policy, +) +from verl.utils import tensordict_utils as tu +from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, should_save_ckpt_esi +from verl.utils.config import omega_conf_to_dataclass +from verl.utils.debug import marked_timer +from verl.utils.import_utils import deprecated, load_class_from_fqn +from verl.utils.metric import reduce_metrics +from verl.utils.py_functional import rename_dict +from verl.utils.seqlen_balancing import calculate_workload, get_seqlen_balanced_partitions, log_seqlen_unbalance +from verl.utils.skip.skip_manager import SkipManager +from verl.utils.torch_functional import masked_mean +from verl.utils.tracking import ValidationGenerationsLogger +from verl.workers.config import DistillationConfig, EngineConfig +from verl.workers.rollout.llm_server import LLMServerManager +from verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding + + +def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty="kl"): + """Apply KL penalty to the token-level rewards. + + This function computes the KL divergence between the reference policy and current policy, + then applies a penalty to the token-level rewards based on this divergence. + + Args: + data (DataProto): The data containing batched model outputs and inputs. + kl_ctrl (core_algos.AdaptiveKLController): Controller for adaptive KL penalty. + kl_penalty (str, optional): Type of KL penalty to apply. Defaults to "kl". + + Returns: + tuple: A tuple containing: + - The updated data with token-level rewards adjusted by KL penalty + - A dictionary of metrics related to the KL penalty + """ + response_mask = data.batch["response_mask"] + token_level_scores = data.batch["token_level_scores"] + batch_size = data.batch.batch_size[0] + + # compute kl between ref_policy and current policy + # When apply_kl_penalty, algorithm.use_kl_in_reward=True, so the reference model has been enabled. + kld = core_algos.kl_penalty( + data.batch["old_log_probs"], data.batch["ref_log_prob"], kl_penalty=kl_penalty + ) # (batch_size, response_length) + kld = kld * response_mask + beta = kl_ctrl.value + + token_level_rewards = token_level_scores - beta * kld + + current_kl = masked_mean(kld, mask=response_mask, axis=-1) # average over sequence + current_kl = torch.mean(current_kl, dim=0).item() + + # according to https://github.com/huggingface/trl/blob/951ca1841f29114b969b57b26c7d3e80a39f75a0/trl/trainer/ppo_trainer.py#L837 + kl_ctrl.update(current_kl=current_kl, n_steps=batch_size) + data.batch["token_level_rewards"] = token_level_rewards + + metrics = {"actor/reward_kl_penalty": current_kl, "actor/reward_kl_penalty_coeff": beta} + + return data, metrics + + +def compute_response_mask(data: DataProto): + """Compute the attention mask for the response part of the sequence. + + This function extracts the portion of the attention mask that corresponds to the model's response, + which is used for masking computations that should only apply to response tokens. + + Args: + data (DataProto): The data containing batched model outputs and inputs. + + Returns: + torch.Tensor: The attention mask for the response tokens. + """ + responses = data.batch["responses"] + response_length = responses.size(1) + attention_mask = data.batch["attention_mask"] + return attention_mask[:, -response_length:] + + +def compute_spec_decode_metrics( + spec_drafts, + spec_accepts, + spec_verifies, + non_padding_mask=None, +) -> dict: + """Aggregate per-request speculative decoding stats. + + Ratios are computed per request and then averaged, so long and short + responses have equal metric weight. + + The three inputs come from the rollout engine (vLLM request spec-decode + stats or sglang ``meta_info["spec_*"]`` keys). Either all three are ``None`` + (caller didn't fetch them, e.g. spec rollout disabled) and the function + is a no-op, or all three are populated; mixed state is a programmer error. + + ``non_padding_mask`` is a numpy bool array used by sync PPO to drop padded + placeholder samples; pass ``None`` for async PPO. + """ + if spec_drafts is None and spec_accepts is None and spec_verifies is None: + return {} + assert spec_drafts is not None and spec_accepts is not None and spec_verifies is not None, ( + "spec_decode metrics require all three of spec_num_draft_tokens / " + "spec_num_accepted_tokens / spec_num_verify_steps; got partial inputs" + ) + + drafts = spec_drafts.tolist() if hasattr(spec_drafts, "tolist") else list(spec_drafts) + accepts = spec_accepts.tolist() if hasattr(spec_accepts, "tolist") else list(spec_accepts) + verifies = spec_verifies.tolist() if hasattr(spec_verifies, "tolist") else list(spec_verifies) + + if non_padding_mask is not None: + drafts = [d for d, keep in zip(drafts, non_padding_mask, strict=True) if keep] + accepts = [a for a, keep in zip(accepts, non_padding_mask, strict=True) if keep] + verifies = [v for v, keep in zip(verifies, non_padding_mask, strict=True) if keep] + + if len(drafts) == 0: + return {} + + # Treat zero-denominator samples as 0.0 and keep them in the mean. + per_sample_accept_rate = [(a / d) if d > 0 else 0.0 for a, d in zip(accepts, drafts, strict=True)] + per_sample_accept_length = [(1.0 + a / v) if v > 0 else 0.0 for a, v in zip(accepts, verifies, strict=True)] + + n = len(drafts) + return { + "rollout/spec_accept_rate": float(sum(per_sample_accept_rate) / n), + "rollout/spec_accept_length": float(sum(per_sample_accept_length) / n), + } + + +def compute_advantage( + data: DataProto, + adv_estimator: AdvantageEstimator, + gamma: float = 1.0, + lam: float = 1.0, + num_repeat: int = 1, + norm_adv_by_std_in_grpo: bool = True, + config: Optional[AlgoConfig] = None, +) -> DataProto: + """Compute advantage estimates for policy optimization. + + This function computes advantage estimates using various estimators like GAE, GRPO, REINFORCE++, etc. + The advantage estimates are used to guide policy optimization in RL algorithms. + + Args: + data (DataProto): The data containing batched model outputs and inputs. + adv_estimator (AdvantageEstimator): The advantage estimator to use (e.g., GAE, GRPO, REINFORCE++). + gamma (float, optional): Discount factor for future rewards. Defaults to 1.0. + lam (float, optional): Lambda parameter for GAE. Defaults to 1.0. + num_repeat (int, optional): Number of times to repeat the computation. Defaults to 1. + norm_adv_by_std_in_grpo (bool, optional): Whether to normalize advantages by standard deviation in + GRPO. Defaults to True. + config (dict, optional): Configuration dictionary for algorithm settings. Defaults to None. + + Returns: + DataProto: The updated data with computed advantages and returns. + """ + # Back-compatible with trainers that do not compute response mask in fit + if "response_mask" not in data.batch.keys(): + data.batch["response_mask"] = compute_response_mask(data) + # prepare response group + if adv_estimator == AdvantageEstimator.GAE: + # Compute advantages and returns using Generalized Advantage Estimation (GAE) + advantages, returns = core_algos.compute_gae_advantage_return( + token_level_rewards=data.batch["token_level_rewards"], + values=data.batch["values"], + response_mask=data.batch["response_mask"], + gamma=gamma, + lam=lam, + ) + data.batch["advantages"] = advantages + data.batch["returns"] = returns + if config.get("use_pf_ppo", False): + data = core_algos.compute_pf_ppo_reweight_data( + data, + config.pf_ppo.get("reweight_method"), + config.pf_ppo.get("weight_pow"), + ) + elif adv_estimator == AdvantageEstimator.GRPO: + # Initialize the mask for GRPO calculation + grpo_calculation_mask = data.batch["response_mask"] + + # Call compute_grpo_outcome_advantage with parameters matching its definition + advantages, returns = core_algos.compute_grpo_outcome_advantage( + token_level_rewards=data.batch["token_level_rewards"], + response_mask=grpo_calculation_mask, + index=data.non_tensor_batch["uid"], + norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo, + ) + data.batch["advantages"] = advantages + data.batch["returns"] = returns + else: + # handle all other adv estimator type other than GAE and GRPO + adv_estimator_fn = core_algos.get_adv_estimator_fn(adv_estimator) + adv_kwargs = { + "token_level_rewards": data.batch["token_level_rewards"], + "response_mask": data.batch["response_mask"], + "config": config, + } + if "uid" in data.non_tensor_batch: # optional + adv_kwargs["index"] = data.non_tensor_batch["uid"] + if "reward_baselines" in data.batch: # optional + adv_kwargs["reward_baselines"] = data.batch["reward_baselines"] + # GDPO: pass raw data for per-dimension reward extraction + if adv_estimator in (AdvantageEstimator.GDPO, "gdpo"): + adv_kwargs["non_tensor_batch"] = data.non_tensor_batch + adv_kwargs["batch"] = data.batch + # Add sum_pi_squared for Optimal Token Baseline + if adv_estimator in (AdvantageEstimator.OPTIMAL_TOKEN_BASELINE, AdvantageEstimator.TIR_OPTIMAL_TOKEN_BASELINE): + # Check if sum_pi_squared is available + assert "sum_pi_squared" in data.batch, ( + "Step-dependent optimal baseline requires sum_pi_squared from actor. " + "Please set actor.calculate_sum_pi_squared=True in config." + ) + adv_kwargs["sum_pi_squared"] = data.batch["sum_pi_squared"] + # old_log_probs needed for path-variance proxy: w_t = 1 - 2*exp(old_log_probs) + sum_pi_squared + adv_kwargs["old_log_probs"] = data.batch["old_log_probs"] + # Get pre-computed rollout IS weights if available + rollout_is_weights = data.batch.get("rollout_is_weights", None) + adv_kwargs["rollout_is_weights"] = rollout_is_weights + + # calculate advantage estimator + advantages, returns = adv_estimator_fn(**adv_kwargs) + data.batch["advantages"] = advantages + data.batch["returns"] = returns + return data + + +@deprecated("Legacy trainer is deprecated, and wil be removed in v0.9.0. Please use `trainer.use_v1=True` instead.") +class RayPPOTrainer: + """Distributed PPO trainer using Ray for scalable reinforcement learning. + + This trainer orchestrates distributed PPO training across multiple nodes and GPUs, + managing actor rollouts, critic training, and reward computation with Ray backend. + Supports various model architectures including FSDP, Megatron, vLLM, and SGLang integration. + """ + + # TODO: support each role have individual ray_worker_group_cls, + # i.e., support different backend of different role + def __init__( + self, + config, + tokenizer, + role_worker_mapping: dict[Role, WorkerType], + resource_pool_manager: ResourcePoolManager, + ray_worker_group_cls: type[RayWorkerGroup] = RayWorkerGroup, + processor=None, + train_dataset: Optional[Dataset] = None, + val_dataset: Optional[Dataset] = None, + collate_fn=None, + train_sampler: Optional[Sampler] = None, + device_name=None, + ): + """ + Initialize distributed PPO trainer with Ray backend. + Note that this trainer runs on the driver process on a single CPU/GPU node. + + Args: + config: Configuration object containing training parameters. + tokenizer: Tokenizer used for encoding and decoding text. + role_worker_mapping (dict[Role, WorkerType]): Mapping from roles to worker classes. + resource_pool_manager (ResourcePoolManager): Manager for Ray resource pools. + ray_worker_group_cls (RayWorkerGroup, optional): Class for Ray worker groups. Defaults to RayWorkerGroup. + processor: Optional data processor, used for multimodal data + train_dataset (Optional[Dataset], optional): Training dataset. Defaults to None. + val_dataset (Optional[Dataset], optional): Validation dataset. Defaults to None. + collate_fn: Function to collate data samples into batches. + train_sampler (Optional[Sampler], optional): Sampler for the training dataset. Defaults to None. + device_name (str, optional): Device name for training (e.g., "cuda", "cpu"). Defaults to None. + """ + + # Store the tokenizer for text processing + self.tokenizer = tokenizer + self.processor = processor + self.config = config + + self.hybrid_engine = config.actor_rollout_ref.hybrid_engine + assert self.hybrid_engine, "Currently, only support hybrid engine" + + if self.hybrid_engine: + assert Role.ActorRollout in role_worker_mapping or Role.ActorRolloutRef in role_worker_mapping, ( + f"{role_worker_mapping.keys()=}" + ) + + self.role_worker_mapping = role_worker_mapping + self.resource_pool_manager = resource_pool_manager + self.use_reference_policy = need_reference_policy(self.config) + self.use_teacher_policy = need_teacher_policy(self.config) + + self.use_rm = need_reward_model(self.config) + + self.use_critic = need_critic(self.config) + self.ray_worker_group_cls = ray_worker_group_cls + self.device_name = device_name if device_name else self.config.trainer.device + self.validation_generations_logger = ValidationGenerationsLogger( + project_name=self.config.trainer.project_name, + experiment_name=self.config.trainer.experiment_name, + ) + + # if ref_in_actor is True, the reference policy will be actor without lora applied + lora_rank = config.actor_rollout_ref.model.get("lora", {}).get("rank", 0) + if lora_rank <= 0: + lora_rank = config.actor_rollout_ref.model.get("lora_rank", 0) + self.ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get("lora_adapter_path") is not None + + # define in-reward KL control + # kl loss control currently not suppoorted + if self.config.algorithm.use_kl_in_reward: + self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl) + + self.use_prefix_grouper = self.config.actor_rollout_ref.actor.get("use_prefix_grouper", False) + + self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler) + + self.checkpoint_manager = None + self._init_dump_executor() + + def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampler: Optional[Sampler]): + """ + Creates the train and validation dataloaders. + """ + if train_dataset is None: + train_dataset = create_rl_dataset( + self.config.data.train_files, + self.config.data, + self.tokenizer, + self.processor, + max_samples=self.config.data.get("train_max_samples", -1), + ) + if val_dataset is None: + val_dataset = create_rl_dataset( + self.config.data.val_files, + self.config.data, + self.tokenizer, + self.processor, + max_samples=self.config.data.get("val_max_samples", -1), + ) + self.train_dataset, self.val_dataset = train_dataset, val_dataset + + if train_sampler is None: + train_sampler = create_rl_sampler(self.config.data, self.train_dataset) + if collate_fn is None: + from verl.utils.dataset.rl_dataset import collate_fn as default_collate_fn + + collate_fn = default_collate_fn + + num_workers = self.config.data["dataloader_num_workers"] + + self.train_dataloader = StatefulDataLoader( + dataset=self.train_dataset, + batch_size=self.config.data.get("gen_batch_size", self.config.data.train_batch_size), + num_workers=num_workers, + drop_last=True, + collate_fn=collate_fn, + sampler=train_sampler, + ) + + val_batch_size = self.config.data.val_batch_size # Prefer config value if set + if val_batch_size is None: + val_batch_size = len(self.val_dataset) + + self.val_dataloader = StatefulDataLoader( + dataset=self.val_dataset, + batch_size=val_batch_size, + num_workers=num_workers, + shuffle=self.config.data.get("validation_shuffle", True), + drop_last=False, + collate_fn=collate_fn, + ) + + assert len(self.train_dataloader) >= 1, "Train dataloader is empty!" + assert len(self.val_dataloader) >= 1, "Validation dataloader is empty!" + + print( + f"Size of train dataloader: {len(self.train_dataloader)}, Size of val dataloader: " + f"{len(self.val_dataloader)}" + ) + + total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs + + if self.config.trainer.total_training_steps is not None: + total_training_steps = self.config.trainer.total_training_steps + + self.total_training_steps = total_training_steps + print(f"Total training steps: {self.total_training_steps}") + + try: + OmegaConf.set_struct(self.config, True) + with open_dict(self.config): + if OmegaConf.select(self.config, "actor_rollout_ref.actor.optim"): + self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps + if OmegaConf.select(self.config, "critic.optim"): + self.config.critic.optim.total_training_steps = total_training_steps + except Exception as e: + print(f"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}") + + @staticmethod + def _write_generations(inputs, outputs, gts, scores, reward_extra_infos_dict, dump_path, global_steps): + """Write generation samples as JSONL (runs in background thread).""" + os.makedirs(dump_path, exist_ok=True) + filename = os.path.join(dump_path, f"{global_steps}.jsonl") + + n = len(inputs) + base_data = { + "input": inputs, + "output": outputs, + "gts": gts, + "score": scores, + "step": [global_steps] * n, + } + + for k, v in reward_extra_infos_dict.items(): + if len(v) == n: + base_data[k] = v + + with open(filename, "w") as f: + for i in range(n): + entry = {k: v[i] for k, v in base_data.items()} + f.write(json.dumps(entry, ensure_ascii=False, default=str) + "\n") + + print(f"Dumped generations to {filename}") + + def _dump_generations(self, inputs, outputs, gts, scores, reward_extra_infos_dict, dump_path): + """Dump rollout/validation samples as JSONL asynchronously.""" + global_steps = self.global_steps + future = self._dump_executor.submit( + self._write_generations, + inputs, + outputs, + gts, + scores, + reward_extra_infos_dict, + dump_path, + global_steps, + ) + self._dump_futures.append(future) + # Clean up completed futures and surface any exceptions early + still_pending = [] + for f in self._dump_futures: + if f.done(): + f.result() # re-raises if the write failed + else: + still_pending.append(f) + self._dump_futures = still_pending + + def _init_dump_executor(self): + """Create or recreate the dump executor and futures list.""" + self._dump_executor = ThreadPoolExecutor(max_workers=1) + self._dump_futures = [] + + def _shutdown_dump_executor(self): + """Drain pending dump futures and shut down the executor.""" + for f in self._dump_futures: + f.result() + self._dump_futures.clear() + self._dump_executor.shutdown(wait=True) + + def _log_rollout_data( + self, batch: DataProto, reward_extra_infos_dict: dict, timing_raw: dict, rollout_data_dir: str + ): + """Log rollout data to disk. + Args: + batch (DataProto): The batch containing rollout data + reward_extra_infos_dict (dict): Additional reward information to log + timing_raw (dict): Timing information for profiling + rollout_data_dir (str): Directory path to save the rollout data + """ + with marked_timer("dump_rollout_generations", timing_raw, color="green"): + inputs = self.tokenizer.batch_decode(batch.batch["prompts"], skip_special_tokens=True) + outputs = self.tokenizer.batch_decode(batch.batch["responses"], skip_special_tokens=True) + scores = batch.batch["token_level_scores"].sum(-1).cpu().tolist() + sample_gts = [item.non_tensor_batch.get("reward_model", {}).get("ground_truth", None) for item in batch] + + reward_extra_infos_to_dump = { + k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in reward_extra_infos_dict.items() + } + if "request_id" in batch.non_tensor_batch: + reward_extra_infos_to_dump.setdefault( + "request_id", + batch.non_tensor_batch["request_id"].tolist(), + ) + + self._dump_generations( + inputs=inputs, + outputs=outputs, + gts=sample_gts, + scores=scores, + reward_extra_infos_dict=reward_extra_infos_to_dump, + dump_path=rollout_data_dir, + ) + + def _maybe_log_val_generations(self, inputs, outputs, scores): + """Log a table of validation samples to the configured logger (wandb or swanlab)""" + + generations_to_log = self.config.trainer.log_val_generations + + if generations_to_log == 0: + return + + import numpy as np + + # Create tuples of (input, output, score) and sort by input text + samples = list(zip(inputs, outputs, scores, strict=True)) + samples.sort(key=lambda x: x[0]) # Sort by input text + + # Use fixed random seed for deterministic shuffling + rng = np.random.RandomState(42) + rng.shuffle(samples) + + # Take first N samples after shuffling + samples = samples[:generations_to_log] + + # Log to each configured logger + self.validation_generations_logger.log(self.config.trainer.logger, samples, self.global_steps) + + def _get_gen_batch(self, batch: DataProto) -> DataProto: + reward_keys = set({"data_source", "reward_model", "extra_info", "uid"}) & batch.non_tensor_batch.keys() + + # pop those keys for generation + batch_keys_to_pop = [] + non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys()) - reward_keys + gen_batch = batch.pop( + batch_keys=batch_keys_to_pop, + non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop), + ) + + # For agent loop, we need reward model keys to compute score. + gen_batch.non_tensor_batch.update(batch.non_tensor_batch) + + return gen_batch + + def _compute_reward_colocate(self, batch: DataProto) -> tuple[torch.Tensor, dict[str, Any]] | torch.Tensor: + """ + compute reward use colocate reward model + """ + assert self.reward_loop_manager is not None, "RewardLoopManager is None" + batch_reward = self.reward_loop_manager.compute_rm_score(batch) + return batch_reward + + def _validate(self, merged: bool = False): + data_source_lst = [] + reward_extra_infos_dict: dict[str, list] = defaultdict(list) + + # Lists to collect samples for the table + sample_inputs = [] + sample_outputs = [] + sample_gts = [] + sample_scores = [] + sample_turns = [] + sample_uids = [] + + for test_data in self.val_dataloader: + test_batch = DataProto.from_single_dict(test_data) + + if "uid" not in test_batch.non_tensor_batch: + test_batch.non_tensor_batch["uid"] = np.array( + [str(uuid.uuid4()) for _ in range(len(test_batch.batch))], dtype=object + ) + + # repeat test batch + test_batch = test_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, interleave=True + ) + + ground_truths = [ + item.non_tensor_batch.get("reward_model", {}).get("ground_truth", None) for item in test_batch + ] + sample_gts.extend(ground_truths) + + test_gen_batch = self._get_gen_batch(test_batch) + test_gen_batch.meta_info = { + "eos_token_id": self.tokenizer.eos_token_id, + "pad_token_id": self.tokenizer.pad_token_id, + "recompute_log_prob": False, + "do_sample": self.config.actor_rollout_ref.rollout.val_kwargs.do_sample, + "validate": True, + "global_steps": self.global_steps, + } + print(f"test_gen_batch meta info: {test_gen_batch.meta_info}") + + # pad to be divisible by dp_size + size_divisor = self.config.actor_rollout_ref.rollout.agent.num_workers + test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, size_divisor) + test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(test_gen_batch_padded) + + if self.use_rm and "rm_scores" not in test_output_gen_batch_padded.batch.keys(): + # for colocate reward models, we need to sleep rollout model + # to spare GPU memory for reward model + self.checkpoint_manager.sleep_replicas() + batch_reward = self._compute_reward_colocate(test_output_gen_batch_padded) + test_output_gen_batch_padded = test_output_gen_batch_padded.union(batch_reward) + # wake up rollout model + # replace with wake_up method once supported + self.checkpoint_manager.update_weights(self.global_steps) + + # unpad + test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size) + + print("validation generation end") + + # Store generated outputs + output_ids = test_output_gen_batch.batch["responses"] + output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids] + sample_outputs.extend(output_texts) + + test_batch = test_batch.union(test_output_gen_batch) + test_batch.meta_info["validate"] = True + + # Store original inputs + input_ids = test_batch.batch["prompts"] + # TODO: Can we keep special tokens except for padding tokens? + input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids] + sample_inputs.extend(input_texts) + sample_uids.extend(test_batch.non_tensor_batch["uid"]) + + # evaluate using reward_function + reward_tensor, reward_extra_info = extract_reward(test_batch) + + scores = reward_tensor.sum(-1).cpu().tolist() + sample_scores.extend(scores) + + reward_extra_infos_dict["reward"].extend(scores) + for key, values in reward_extra_info.items(): + if key not in reward_extra_infos_dict: + reward_extra_infos_dict[key] = [] + if isinstance(values, np.ndarray): + reward_extra_infos_dict[key].extend(values.tolist()) + else: + reward_extra_infos_dict[key].extend(values if isinstance(values, list) else [values]) + + # collect num_turns of each prompt + if "__num_turns__" in test_batch.non_tensor_batch: + sample_turns.append(test_batch.non_tensor_batch["__num_turns__"]) + + data_source_lst.append(test_batch.non_tensor_batch.get("data_source", ["unknown"] * reward_tensor.shape[0])) + + self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores) + + # dump generations + val_data_dir = self.config.trainer.get("validation_data_dir", None) + if val_data_dir: + self._dump_generations( + inputs=sample_inputs, + outputs=sample_outputs, + gts=sample_gts, + scores=sample_scores, + reward_extra_infos_dict=reward_extra_infos_dict, + dump_path=val_data_dir, + ) + + for key_info, lst in reward_extra_infos_dict.items(): + assert len(lst) == 0 or len(lst) == len(sample_scores), f"{key_info}: {len(lst)=}, {len(sample_scores)=}" + + if merged: + print("_merge_validation_results validate result will be merged") + return { + "data_sources": data_source_lst, + "sample_uids": sample_uids, + "sample_turns": sample_turns, + "reward_extra_infos_dict": reward_extra_infos_dict, + } + data_sources = np.concatenate(data_source_lst, axis=0) + return self._val_metrics_update(data_sources, sample_uids, reward_extra_infos_dict, sample_turns) + + def _val_metrics_update(self, data_sources, sample_uids, reward_extra_infos_dict, sample_turns): + data_src2var2metric2val = process_validation_metrics(data_sources, sample_uids, reward_extra_infos_dict) + metric_dict = {} + for data_source, var2metric2val in data_src2var2metric2val.items(): + core_var = "acc" if "acc" in var2metric2val else "reward" + for var_name, metric2val in var2metric2val.items(): + n_max = max([int(name.split("@")[-1].split("/")[0]) for name in metric2val.keys()]) + for metric_name, metric_val in metric2val.items(): + if ( + (var_name == core_var) + and any(metric_name.startswith(pfx) for pfx in ["mean", "maj", "best"]) + and (f"@{n_max}" in metric_name) + ): + metric_sec = "val-core" + else: + metric_sec = "val-aux" + pfx = f"{metric_sec}/{data_source}/{var_name}/{metric_name}" + metric_dict[pfx] = metric_val + + if len(sample_turns) > 0: + sample_turns = np.concatenate(sample_turns) + metric_dict["val-aux/num_turns/min"] = sample_turns.min() + metric_dict["val-aux/num_turns/max"] = sample_turns.max() + metric_dict["val-aux/num_turns/mean"] = sample_turns.mean() + + return metric_dict + + def _merge_validation_results(self, result_a, result_b): + if result_a is None and result_b is None: + return {} + if result_a is None: + result_a = {"data_sources": [], "sample_uids": [], "sample_turns": [], "reward_extra_infos_dict": {}} + if result_b is None: + result_b = {"data_sources": [], "sample_uids": [], "sample_turns": [], "reward_extra_infos_dict": {}} + + if not result_a.get("data_sources") and not result_b.get("data_sources"): + return {} + + data_sources = np.concatenate(result_a["data_sources"] + result_b["data_sources"], axis=0) + sample_uids = result_a["sample_uids"] + result_b["sample_uids"] + sample_turns = result_a["sample_turns"] + result_b["sample_turns"] + + reward_extra_infos_dict = {} + all_keys = set(result_a["reward_extra_infos_dict"].keys()) | set(result_b["reward_extra_infos_dict"].keys()) + for key in all_keys: + list_a = result_a["reward_extra_infos_dict"].get(key, []) + list_b = result_b["reward_extra_infos_dict"].get(key, []) + reward_extra_infos_dict[key] = list_a + list_b + + return self._val_metrics_update(data_sources, sample_uids, reward_extra_infos_dict, sample_turns) + + def init_workers(self): + """Initialize distributed training workers using Ray backend. + + Creates: + 1. Ray resource pools from configuration + 2. Worker groups for each role (actor, critic, etc.) + """ + self.resource_pool_manager.create_resource_pool() + + self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()} + + # create actor and rollout + actor_role = Role.ActorRolloutRef if Role.ActorRolloutRef in self.role_worker_mapping else Role.ActorRollout + if self.hybrid_engine: + actor_rollout_resource_pool = self.resource_pool_manager.get_resource_pool(actor_role) + actor_rollout_cls = RayClassWithInitArgs( + cls=self.role_worker_mapping[actor_role], + config=self.config.actor_rollout_ref, + distillation_config=self.config.get("distillation"), + role=str(actor_role), + ) + self.resource_pool_to_cls[actor_rollout_resource_pool][str(actor_role)] = actor_rollout_cls + else: + raise NotImplementedError + + # create critic + if self.use_critic: + resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic) + + from verl.workers.config import CriticConfig + + critic_cfg: CriticConfig = omega_conf_to_dataclass(self.config.critic) + + # convert critic_cfg into TrainingWorkerConfig for the unified model engine worker + from verl.workers.engine_workers import TrainingWorkerConfig + + orig_critic_cfg = critic_cfg + engine_config: EngineConfig = orig_critic_cfg.engine + engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu + engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu + + critic_cfg = TrainingWorkerConfig( + model_type="value_model", + model_config=orig_critic_cfg.model, + engine_config=engine_config, + optimizer_config=orig_critic_cfg.optim, + checkpoint_config=orig_critic_cfg.checkpoint, + extra_context=getattr(self, "_critic_extra_context", {}), + ) + + critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg) + self.resource_pool_to_cls[resource_pool][str(Role.Critic)] = critic_cls + + # create reference policy if needed + if self.use_reference_policy and Role.RefPolicy in self.role_worker_mapping: + resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy) + ref_policy_cls = RayClassWithInitArgs( + self.role_worker_mapping[Role.RefPolicy], + config=self.config.actor_rollout_ref, + role=str(Role.RefPolicy), + ) + self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls + + # FSDP teacher (verl#6676): register a TeacherFSDPWorker on the teacher resource + # pool so it joins the unified RayWorkerGroup spawn. The vLLM backend spawns its + # own rollout replica instead. + self._fsdp_teacher_pending = False + if self.use_teacher_policy and self.config.distillation.get("teacher_backend", "vllm") == "fsdp": + teacher_training_config = self._build_fsdp_teacher_training_config() + teacher_resource_pool = self.resource_pool_manager.get_resource_pool(Role.TeacherModel) + + import ray + + from verl.workers.teacher import TeacherFSDPWorker + + teacher_cls = RayClassWithInitArgs(cls=ray.remote(TeacherFSDPWorker), config=teacher_training_config) + self.resource_pool_to_cls[teacher_resource_pool][str(Role.TeacherModel)] = teacher_cls + self._fsdp_teacher_pending = True + + # initialize WorkerGroup + # NOTE: if you want to use a different resource pool for each role, which can support different parallel size, + # you should not use `create_colocated_worker_cls`. + # Instead, directly pass different resource pool to different worker groups. + # See https://github.com/verl-project/verl/blob/master/examples/tutorial/ray/tutorial.ipynb + # for more information. + all_wg = {} + wg_kwargs = {} # Setting up kwargs for RayWorkerGroup + if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None: + wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout + if OmegaConf.select(self.config.global_profiler, "steps") is not None: + wg_kwargs["profile_steps"] = OmegaConf.select(self.config.global_profiler, "steps") + # Only require nsight worker options when tool is nsys + if OmegaConf.select(self.config.global_profiler, "tool") == "nsys": + assert ( + OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options") + is not None + ), "worker_nsight_options must be set when using nsys with profile_steps" + wg_kwargs["worker_nsight_options"] = OmegaConf.to_container( + OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, "worker_nsight_options") + ) + wg_kwargs["device_name"] = self.device_name + + for resource_pool, class_dict in self.resource_pool_to_cls.items(): + if not class_dict: + continue + worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict) + wg_dict = self.ray_worker_group_cls( + resource_pool=resource_pool, + ray_cls_with_init=worker_dict_cls, + **wg_kwargs, + ) + spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys()) + all_wg.update(spawn_wg) + + if self.use_critic: + self.critic_wg = all_wg[str(Role.Critic)] + self.critic_wg.reset() + # assign critic loss + from functools import partial + + from verl.workers.utils.losses import value_loss + + value_loss_ = partial(value_loss, config=orig_critic_cfg) + self.critic_wg.set_loss_fn(value_loss_) + + if self.use_reference_policy and not self.ref_in_actor: + if str(Role.RefPolicy) in all_wg: + self.ref_policy_wg = all_wg[str(Role.RefPolicy)] + self.ref_policy_wg.init_model() + else: + # Model engine: ActorRolloutRefWorker + assert str(Role.ActorRolloutRef) in all_wg, f"{all_wg.keys()=}" + self.ref_policy_wg = all_wg[str(Role.ActorRolloutRef)] + + # we should create rollout at the end so that vllm can have a better estimation of kv cache memory + self.actor_rollout_wg = all_wg[str(actor_role)] + self.actor_rollout_wg.init_model() + + if self.ref_in_actor: + self.ref_policy_wg = self.actor_rollout_wg + + # create reward loop manager + from verl.experimental.reward_loop import RewardLoopManager + + # initalize reward loop manager + # reward model (colocate or standalone): get resource_pool + # no reward model: resource_pool = None + resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) if self.use_rm else None + self.reward_loop_manager = RewardLoopManager( + config=self.config, + rm_resource_pool=resource_pool, + ) + + # create async rollout manager and request scheduler + # Note: mode is always "async" since sync mode is deprecated + self.async_rollout_mode = True + + # initialize teacher loop manager + if self.use_teacher_policy: + teacher_resource_pool = self.resource_pool_manager.get_resource_pool(Role.TeacherModel) + self.distillation_config: DistillationConfig = omega_conf_to_dataclass(self.config.distillation) + teacher_backend = self.distillation_config.teacher_backend + if teacher_backend == "fsdp": + # FSDP teacher (verl#6676): hosted as a frozen FSDP worker group rather + # than a vLLM rollout, because reverse-KL on the student top-K needs + # log-probs at arbitrary token IDs, which vLLM cannot serve. The trainer + # drives the teacher forward after rollout via _inject_fsdp_teacher_data. + from verl.experimental.teacher_loop.teacher_fsdp_manager import MultiTeacherFSDPManager + + self.teacher_model_manager = MultiTeacherFSDPManager( + config=self.config, + resource_pool=teacher_resource_pool, + ) + # Install the spawned worker group handle on the manager and bring up the model. + assert self._fsdp_teacher_pending and str(Role.TeacherModel) in all_wg, ( + f"FSDP teacher worker group not found in spawn output: {all_wg.keys()=}" + ) + self.teacher_wg = all_wg[str(Role.TeacherModel)] + self.teacher_wg.reset() + self.teacher_model_manager.set_worker_group(self.teacher_wg) + elif teacher_backend == "vllm": + from verl.experimental.teacher_loop import MultiTeacherModelManager + + self.teacher_model_manager = MultiTeacherModelManager( + config=self.config, + resource_pool=teacher_resource_pool, + ) + else: + raise ValueError( + f"Unsupported distillation.teacher_backend {teacher_backend!r}; expected 'vllm' or 'fsdp'." + ) + else: + self.teacher_model_manager = None + self.distillation_config = None + + # Support custom AgentLoopManager via config + manager_class_fqn = self.config.actor_rollout_ref.rollout.get("agent", {}).get("agent_loop_manager_class") + if manager_class_fqn: + AgentLoopManager = load_class_from_fqn(manager_class_fqn, "AgentLoopManager") + else: + from verl.experimental.agent_loop import AgentLoopManager + + # infrastructure overview: https://verl.readthedocs.io/en/latest/advance/reward_loop.html#architecture-design + # agent_reward_loop: streaming reward computation with actor rollout + # two conditions satisfied: (1) no reward model, or (2) reward model with extra resource pool + enable_agent_reward_loop = not self.use_rm or self.config.reward.reward_model.enable_resource_pool + + self.llm_server_manager = LLMServerManager.create( + config=self.config, worker_group=self.actor_rollout_wg, rollout_resource_pool=actor_rollout_resource_pool + ) + + # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager + # to stream reward computation with actor rollout + # To stream teacher computation with actor rollout, we instead pass the full manager so that the + # teacher loop workers can sleep/wake together with rollout workers. The FSDP teacher has no + # per-trajectory LLM client, so only the vLLM path forwards a client into the agent loop. + reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None + teacher_client_for_agent_loop = None + if self.use_teacher_policy and self.distillation_config.teacher_backend == "vllm": + teacher_client_for_agent_loop = self.teacher_model_manager.get_client() + self.async_rollout_manager = AgentLoopManager.create( + config=self.config, + llm_client=self.llm_server_manager.get_client(), + teacher_client=teacher_client_for_agent_loop, + reward_loop_worker_handles=reward_loop_worker_handles, + ) + + checkpoint_engine_config = omega_conf_to_dataclass(self.config.actor_rollout_ref.rollout.checkpoint_engine) + # Support custom CheckpointEngineManager via config + checkpoint_manager_class_fqn = self.config.actor_rollout_ref.rollout.get("checkpoint_manager_class") + if checkpoint_manager_class_fqn: + CheckpointEngineManager = load_class_from_fqn(checkpoint_manager_class_fqn, "CheckpointEngineManager") + else: + from verl.checkpoint_engine import CheckpointEngineManager + self.checkpoint_manager = CheckpointEngineManager( + config=checkpoint_engine_config, + actor_wg=self.actor_rollout_wg, + replicas=self.llm_server_manager.get_replicas(), + ) + + # sleep all replicas to load checkpoint + self.checkpoint_manager.sleep_replicas() + + def _save_checkpoint(self): + from verl.utils.fs import local_mkdir_safe + + # path: given_path + `/global_step_{global_steps}` + `/actor` + local_global_step_folder = os.path.join( + self.config.trainer.default_local_dir, f"global_step_{self.global_steps}" + ) + + print(f"local_global_step_folder: {local_global_step_folder}") + actor_local_path = os.path.join(local_global_step_folder, "actor") + + actor_remote_path = ( + None + if self.config.trainer.default_hdfs_dir is None + else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "actor") + ) + + remove_previous_ckpt_in_save = self.config.trainer.get("remove_previous_ckpt_in_save", False) + if remove_previous_ckpt_in_save: + print( + "Warning: remove_previous_ckpt_in_save is deprecated," + + " set max_actor_ckpt_to_keep=1 and max_critic_ckpt_to_keep=1 instead" + ) + max_actor_ckpt_to_keep = ( + self.config.trainer.get("max_actor_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1 + ) + max_critic_ckpt_to_keep = ( + self.config.trainer.get("max_critic_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1 + ) + + self.actor_rollout_wg.save_checkpoint( + actor_local_path, actor_remote_path, self.global_steps, max_ckpt_to_keep=max_actor_ckpt_to_keep + ) + + if self.use_critic: + critic_local_path = os.path.join(local_global_step_folder, str(Role.Critic)) + critic_remote_path = ( + None + if self.config.trainer.default_hdfs_dir is None + else os.path.join( + self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", str(Role.Critic) + ) + ) + self.critic_wg.save_checkpoint( + critic_local_path, critic_remote_path, self.global_steps, max_ckpt_to_keep=max_critic_ckpt_to_keep + ) + + # save dataloader + local_mkdir_safe(local_global_step_folder) + dataloader_local_path = os.path.join(local_global_step_folder, "data.pt") + dataloader_state_dict = self.train_dataloader.state_dict() + torch.save(dataloader_state_dict, dataloader_local_path) + + # latest checkpointed iteration tracker (for atomic usage) + if ( + hasattr(self.config.actor_rollout_ref.actor.checkpoint, "async_save") + and self.config.actor_rollout_ref.actor.checkpoint.async_save + ) or ( + "async_save" in self.config.actor_rollout_ref.actor.checkpoint + and self.config.actor_rollout_ref.actor.checkpoint["async_save"] + ): + print("skip write latest_checkpointed_iteration.txt when async_save is True") + return + local_latest_checkpointed_iteration = os.path.join( + self.config.trainer.default_local_dir, "latest_checkpointed_iteration.txt" + ) + with open(local_latest_checkpointed_iteration, "w") as f: + f.write(str(self.global_steps)) + + def _load_checkpoint(self): + if self.config.trainer.resume_mode == "disable": + return 0 + + # load from hdfs + if self.config.trainer.default_hdfs_dir is not None: + raise NotImplementedError("load from hdfs is not implemented yet") + else: + checkpoint_folder = self.config.trainer.default_local_dir # TODO: check path + if not os.path.isabs(checkpoint_folder): + working_dir = os.getcwd() + checkpoint_folder = os.path.join(working_dir, checkpoint_folder) + global_step_folder = find_latest_ckpt_path(checkpoint_folder) # None if no latest + + # find global_step_folder + if self.config.trainer.resume_mode == "auto": + if global_step_folder is None: + print("Training from scratch") + return 0 + else: + if self.config.trainer.resume_mode == "resume_path": + assert isinstance(self.config.trainer.resume_from_path, str), "resume ckpt must be str type" + assert "global_step_" in self.config.trainer.resume_from_path, ( + "resume ckpt must specify the global_steps" + ) + global_step_folder = self.config.trainer.resume_from_path + if not os.path.isabs(global_step_folder): + working_dir = os.getcwd() + global_step_folder = os.path.join(working_dir, global_step_folder) + print(f"Load from checkpoint folder: {global_step_folder}") + # set global step + self.global_steps = int(global_step_folder.split("global_step_")[-1]) + + print(f"Setting global step to {self.global_steps}") + print(f"Resuming from {global_step_folder}") + + actor_path = os.path.join(global_step_folder, "actor") + critic_path = os.path.join(global_step_folder, str(Role.Critic)) + # load actor + self.actor_rollout_wg.load_checkpoint( + actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load + ) + # load critic + if self.use_critic: + self.critic_wg.load_checkpoint( + critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load + ) + + # load dataloader, + # TODO: from remote not implemented yet + dataloader_local_path = os.path.join(global_step_folder, "data.pt") + if os.path.exists(dataloader_local_path): + steps_per_epoch = len(self.train_dataloader) + at_epoch_boundary = steps_per_epoch > 0 and self.global_steps % steps_per_epoch == 0 + if at_epoch_boundary: + print( + f"Skipping dataloader state restore: global_steps={self.global_steps} " + f"is at an epoch boundary (steps_per_epoch={steps_per_epoch}). " + f"The saved state marks the dataloader as exhausted. " + f"Next epoch will iterate from scratch." + ) + else: + dataloader_state_dict = torch.load(dataloader_local_path, weights_only=False) + self.train_dataloader.load_state_dict(dataloader_state_dict) + else: + print(f"Warning: No dataloader state found at {dataloader_local_path}, will start from scratch") + + def _start_profiling(self, do_profile: bool) -> None: + """Start profiling for all worker groups if profiling is enabled.""" + if do_profile: + self.actor_rollout_wg.start_profile(role="e2e", profile_step=self.global_steps) + if self.use_reference_policy: + self.ref_policy_wg.start_profile(profile_step=self.global_steps) + if self.use_critic: + self.critic_wg.start_profile(profile_step=self.global_steps) + + def _stop_profiling(self, do_profile: bool) -> None: + """Stop profiling for all worker groups if profiling is enabled.""" + if do_profile: + self.actor_rollout_wg.stop_profile() + if self.use_reference_policy: + self.ref_policy_wg.stop_profile() + if self.use_critic: + self.critic_wg.stop_profile() + + def _get_dp_size(self, worker_group, role: str) -> int: + """Get data parallel size from worker group dispatch info. + + This method retrieves the data parallel size by querying the dispatch info + for the specified role. The dispatch info is cached for subsequent calls. + + Args: + worker_group: The worker group to query dispatch info from. + role: The role name (e.g., "actor", "critic") to get DP size for. + + Returns: + The data parallel size (number of DP ranks). + """ + if role not in worker_group._dispatch_info: + dp_rank_mapping = worker_group._query_dispatch_info(role) + worker_group._dispatch_info[role] = dp_rank_mapping + else: + dp_rank_mapping = worker_group._dispatch_info[role] + return max(dp_rank_mapping) + 1 + + def _balance_batch(self, batch: DataProto, metrics, logging_prefix="global_seqlen", keep_minibatch=False): + """Reorder the data on single controller such that each dp rank gets similar total tokens. + + When use_prefix_grouper is enabled, uses group-level balancing to keep samples with + the same uid together on the same rank for prefix sharing optimization. + """ + attention_mask = batch.batch["attention_mask"] + batch_size = attention_mask.shape[0] + global_seqlen_lst = batch.batch["attention_mask"].view(batch_size, -1).sum(-1) # (train_batch_size,) + workload_lst = calculate_workload(global_seqlen_lst) + # Get dp_size from dispatch info to correctly balance across data parallel ranks + # Note: world_size may include tensor/pipeline parallel dimensions, but we only want DP + dp_size = self._get_dp_size(self.actor_rollout_wg, "actor") + + # Use group-level balancing for PrefixGrouper to keep same-uid samples together + if getattr(self, "use_prefix_grouper", False) and "uid" in batch.non_tensor_batch: + from verl.utils.seqlen_balancing import get_group_balanced_partitions + + uid_list = list(batch.non_tensor_batch["uid"]) + seqlen_list = global_seqlen_lst.tolist() + + # Count number of uid groups + num_groups = len(set(uid_list)) + + if num_groups % dp_size != 0: + raise ValueError( + f"PrefixGrouper with balance_batch requires num_uid_groups ({num_groups}) " + f"% dp_size ({dp_size}) == 0. " + f"This ensures each rank gets equal number of groups. " + f"Current batch_size={batch_size}, adjust batch_size to be a multiple of " + f"dp_size * rollout.n." + ) + + global_partition_lst = get_group_balanced_partitions( + seqlen_list=seqlen_list, + uid_list=uid_list, + k_partitions=dp_size, + ) + + elif keep_minibatch: + # Decouple the DP balancing and mini-batching. + minibatch_size = self.config.actor_rollout_ref.actor.get("ppo_mini_batch_size") + minibatch_num = len(workload_lst) // minibatch_size + global_partition_lst = [[] for _ in range(dp_size)] + for i in range(minibatch_num): + rearrange_minibatch_lst = get_seqlen_balanced_partitions( + workload_lst[i * minibatch_size : (i + 1) * minibatch_size], + k_partitions=dp_size, + equal_size=True, + ) + for j, part in enumerate(rearrange_minibatch_lst): + global_partition_lst[j].extend([x + minibatch_size * i for x in part]) + else: + global_partition_lst = get_seqlen_balanced_partitions(workload_lst, k_partitions=dp_size, equal_size=True) + # Place smaller micro-batches at both ends to reduce the bubbles in pipeline parallel. + # Skip reordering within partitions for PrefixGrouper to maintain uid grouping + if not getattr(self, "use_prefix_grouper", False): + for idx, partition in enumerate(global_partition_lst): + partition.sort(key=lambda x: (workload_lst[x], x)) + ordered_partition = partition[::2] + partition[1::2][::-1] + global_partition_lst[idx] = ordered_partition + + # reorder based on index. The data will be automatically equally partitioned by dispatch function + global_idx = torch.tensor([j for partition in global_partition_lst for j in partition]) + batch.reorder(global_idx) + global_balance_stats = log_seqlen_unbalance( + seqlen_list=global_seqlen_lst.tolist(), partitions=global_partition_lst, prefix=logging_prefix + ) + metrics.update(global_balance_stats) + + def _compute_values(self, batch: DataProto) -> DataProto: + batch_td = batch.to_tensordict() + # step 2: convert from padding to nopadding + batch_td = left_right_2_no_padding(batch_td) + # step 3: add meta info + tu.assign_non_tensor(batch_td, compute_loss=False) + output = self.critic_wg.infer_batch(batch_td) + output = output.get() + values = tu.get(output, "values") + values = no_padding_2_padding(values, batch_td) + values = tu.get_tensordict({"values": values.float()}) + values = DataProto.from_tensordict(values) + return values + + def _compute_ref_log_prob(self, batch: DataProto) -> DataProto: + # step 1: convert dataproto to tensordict. + batch_td = batch.to_tensordict() + # step 2: convert from padding to nopadding + batch_td = left_right_2_no_padding(batch_td) + # step 3: add meta info + metadata = {"calculate_entropy": False, "compute_loss": False} + if self.ref_in_actor: + metadata["no_lora_adapter"] = True + tu.assign_non_tensor(batch_td, **metadata) + if self.ref_in_actor: + output = self.actor_rollout_wg.compute_log_prob(batch_td) + else: + output = self.ref_policy_wg.compute_ref_log_prob(batch_td) + # gather output + log_probs = tu.get(output, "log_probs") + # step 4. No padding to padding + log_probs = no_padding_2_padding(log_probs, batch_td) + # step 5: rebuild a tensordict and convert to dataproto + ref_log_prob = tu.get_tensordict({"ref_log_prob": log_probs.float()}) + ref_log_prob = DataProto.from_tensordict(ref_log_prob) + + return ref_log_prob + + def _compute_old_log_prob(self, batch: DataProto): + # TODO: remove step 1, 2, 4 after we make the whole training tensordict and padding free + # step 1: convert dataproto to tensordict. + batch_td = batch.to_tensordict() + # step 2: convert from padding to nopadding + batch_td = left_right_2_no_padding(batch_td) + # step 3: add meta info + calculate_sum_pi_squared = self.config.actor_rollout_ref.actor.get("calculate_sum_pi_squared", False) + tu.assign_non_tensor( + batch_td, + calculate_entropy=True, + calculate_sum_pi_squared=calculate_sum_pi_squared, + compute_loss=False, + ) + output = self.actor_rollout_wg.compute_log_prob(batch_td) + # gather output + entropy = tu.get(output, "entropy") + log_probs = tu.get(output, "log_probs") + routed_experts = tu.get(output, "routed_experts") + sum_pi_squared = tu.get(output, "sum_pi_squared") if calculate_sum_pi_squared else None + + old_log_prob_mfu = tu.get(output, "metrics")["mfu"] + # step 4. No padding to padding + entropy = no_padding_2_padding(entropy, batch_td) + log_probs = no_padding_2_padding(log_probs, batch_td) + if sum_pi_squared is not None: + sum_pi_squared = no_padding_2_padding(sum_pi_squared, batch_td) + # step 5: rebuild a tensordict and convert to dataproto + result = {"old_log_probs": log_probs.float(), "entropys": entropy.float()} + if routed_experts is not None: + result["routed_experts"] = routed_experts + if sum_pi_squared is not None: + result["sum_pi_squared"] = sum_pi_squared.float() + old_log_prob = tu.get_tensordict(result) + old_log_prob = DataProto.from_tensordict(old_log_prob) + return old_log_prob, old_log_prob_mfu + + def _build_fsdp_teacher_training_config(self): + """Construct a TrainingWorkerConfig for the FSDP-served teacher. + + Reused by :meth:`init_workers` to register :class:`TeacherFSDPWorker` on the + teacher resource pool. ``forward_only=True`` is set here as well as in the + worker's ``__init__``. + """ + from verl.workers.config import FSDPEngineConfig, HFModelConfig + from verl.workers.engine_workers import TrainingWorkerConfig + + teacher_models = self.config.distillation.teacher_models + if len(teacher_models) != 1: + raise NotImplementedError( + "Multi-teacher FSDP backend is not yet supported; configure exactly one teacher " + "(teacher_models has 1 entry) when teacher_backend='fsdp'." + ) + teacher_yaml = next(iter(teacher_models.values())) + model_config = HFModelConfig(path=teacher_yaml.model_path) + + fsdp_cfg = self.config.distillation.teacher_fsdp_config + if fsdp_cfg is None: + raise ValueError("distillation.teacher_fsdp_config must be set when teacher_backend='fsdp'.") + engine_config = omega_conf_to_dataclass(fsdp_cfg, dataclass_type=FSDPEngineConfig) + engine_config.forward_only = True + + return TrainingWorkerConfig( + model_type="language_model", + model_config=model_config, + engine_config=engine_config, + optimizer_config=None, + checkpoint_config=None, + ) + + def _inject_fsdp_teacher_data(self, batch: DataProto) -> DataProto: + """Phase-1 student top-K + Phase-2 teacher gather for the FSDP teacher path. + + For ``teacher_backend == 'fsdp'`` (verl#6676), the trainer drives a two-stage + post-rollout step: (1) the actor emits per-token student top-K log-probs and IDs + via the ``compute_student_topk_only`` flag; (2) the frozen FSDP teacher gathers + ``log_softmax(teacher_logits)`` at those IDs. Both are padded to + ``(B, max_response_len, K)`` and unioned into ``batch`` for ``update_actor``. + + No-op when ``teacher_model_manager`` is not the FSDP variant. + """ + from verl.experimental.teacher_loop.teacher_fsdp_manager import MultiTeacherFSDPManager + + if not isinstance(self.teacher_model_manager, MultiTeacherFSDPManager): + return batch + + # Phase 1: actor produces student top-K (one no-grad forward). + batch_td = batch.to_tensordict() + batch_td = left_right_2_no_padding(batch_td) + tu.assign_non_tensor(batch_td, distillation_use_topk=True, compute_loss=True) + student_topk_output = self.actor_rollout_wg.compute_student_topk(batch_td) + student_topk_ids = tu.get(student_topk_output, "student_topk_ids") # nested (B, T_i, K) + + # Phase 2: teacher worker forward + chunked gather at the student top-K IDs. + teacher_input = tu.get_tensordict( + { + "input_ids": batch_td["input_ids"], + "position_ids": batch_td["position_ids"], + "topk_ids": student_topk_ids, + } + ) + tu.assign_non_tensor( + teacher_input, + teacher_chunk_size=self.distillation_config.teacher_chunk_size, + # Match the student's rollout temperature so reverse-KL compares like with like. + temperature=self.config.actor_rollout_ref.rollout.temperature, + ) + teacher_output = self.teacher_model_manager.compute_logprobs_at_ids(teacher_input) + # compute_logprobs_at_ids is non-blocking; materialize the future. + teacher_output = teacher_output.get() + teacher_on_student_logp = tu.get(teacher_output, "teacher_on_student_logp") # nested (B, T_i, K) + teacher_topk_ids = tu.get(teacher_output, "teacher_topk_ids") # nested (B, T_i, K) + + # Phase 3: nested -> left-right padded dense (B, max_seq_len, K), unioned into + # batch so update_actor's left_right_2_no_padding re-nests them aligned to input_ids. + # NOTE: no_padding_2_padding is wrong here — it slices response-only and left-shifts. + indices = tu.get_non_tensor_data(data=batch_td, key="indices", default=None) + max_seq_len = tu.get_non_tensor_data(data=batch_td, key="max_seq_len", default=None) + assert indices is not None and max_seq_len is not None, ( + "left_right_2_no_padding must run before _inject_fsdp_teacher_data Phase 3." + ) + batch_size = student_topk_ids.size(0) + + def _to_left_right_padded(nested: torch.Tensor) -> torch.Tensor: + # Scatter jagged (B, T_i, K) values back to left-right padded dense + # (B, max_seq_len, K); padding slots are dropped again on the next unpad. + values = nested.values() # (total_nnz, K) + k = values.shape[-1] + flat = values.new_zeros((batch_size * max_seq_len, k)) + flat[indices.to(flat.device)] = values + return flat.reshape(batch_size, max_seq_len, k) + + student_topk_ids_padded = _to_left_right_padded(student_topk_ids) + teacher_on_student_logp_padded = _to_left_right_padded(teacher_on_student_logp) + teacher_topk_ids_padded = _to_left_right_padded(teacher_topk_ids) + extras = tu.get_tensordict( + { + "student_topk_ids": student_topk_ids_padded, + "teacher_on_student_logp": teacher_on_student_logp_padded.float(), + "teacher_topk_ids": teacher_topk_ids_padded, + } + ) + return batch.union(DataProto.from_tensordict(extras)) + + def _update_actor(self, batch: DataProto) -> DataProto: + rollout_config = self.config.actor_rollout_ref.rollout + batch.meta_info["multi_turn"] = rollout_config.multi_turn.enable + # TODO: Make "temperature" single source of truth from generation. + batch.meta_info["temperature"] = rollout_config.temperature + # update actor + batch_td = batch.to_tensordict() + # step 2: convert from padding to no-padding + batch_td = left_right_2_no_padding(batch_td) + calculate_entropy = self.config.actor_rollout_ref.actor.calculate_entropy or ( + self.config.actor_rollout_ref.actor.entropy_coeff != 0.0 + ) + distillation_use_topk = ( + self.distillation_config.distillation_loss.loss_settings.use_topk + if is_distillation_enabled(self.config.get("distillation")) + else False + ) + distillation_only = False # distillation_only flag means we can skip policy loss and reduce mem footprint + if is_distillation_enabled(self.config.get("distillation")): + distillation_loss_cfg = self.distillation_config.distillation_loss + distillation_only = ( + distillation_use_topk + and not distillation_loss_cfg.use_task_rewards + and not distillation_loss_cfg.use_policy_gradient + ) + ppo_mini_batch_size = self.config.actor_rollout_ref.actor.ppo_mini_batch_size + ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n + ppo_epochs = self.config.actor_rollout_ref.actor.ppo_epochs + seed = self.config.actor_rollout_ref.actor.data_loader_seed + shuffle = self.config.actor_rollout_ref.actor.shuffle + tu.assign_non_tensor( + batch_td, + calculate_entropy=calculate_entropy, + distillation_use_topk=distillation_use_topk, + distillation_only=distillation_only, + global_batch_size=ppo_mini_batch_size, + mini_batch_size=ppo_mini_batch_size, + epochs=ppo_epochs, + seed=seed, + dataloader_kwargs={"shuffle": shuffle}, + compute_loss=True, + ) + actor_output = self.actor_rollout_wg.update_actor(batch_td) + actor_output = tu.get(actor_output, "metrics") + actor_output = rename_dict(actor_output, "actor/") + # modify key name + actor_output["perf/mfu/actor"] = actor_output.pop("actor/mfu") + actor_output = DataProto.from_single_dict(data={}, meta_info={"metrics": actor_output}) + + return actor_output + + def _update_critic(self, batch: DataProto) -> DataProto: + batch_td = batch.to_tensordict() + # step 2: convert from padding to no-padding + batch_td = left_right_2_no_padding(batch_td) + ppo_mini_batch_size = self.config.critic.ppo_mini_batch_size + ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n + ppo_epochs = self.config.critic.ppo_epochs + seed = self.config.critic.data_loader_seed + shuffle = self.config.critic.shuffle + tu.assign_non_tensor( + batch_td, + global_batch_size=ppo_mini_batch_size, + mini_batch_size=ppo_mini_batch_size, + epochs=ppo_epochs, + seed=seed, + dataloader_kwargs={"shuffle": shuffle}, + ) + + output = self.critic_wg.train_mini_batch(batch_td) + output = output.get() + output = tu.get(output, "metrics") + output = rename_dict(output, "critic/") + # modify key name + output["perf/mfu/critic"] = output.pop("critic/mfu") + critic_output = DataProto.from_single_dict(data={}, meta_info={"metrics": output}) + return critic_output + + def fit(self): + """ + The training loop of PPO. + The driver process only need to call the compute functions of the worker group through RPC + to construct the PPO dataflow. + The light-weight advantage computation is done on the driver process. + """ + if self._dump_executor._shutdown: + self._init_dump_executor() + + from omegaconf import OmegaConf + + from verl.utils.tracking import Tracking + + logger = Tracking( + project_name=self.config.trainer.project_name, + experiment_name=self.config.trainer.experiment_name, + default_backend=self.config.trainer.logger, + config=OmegaConf.to_container(self.config, resolve=True), + ) + + self.global_steps = 0 + + # load checkpoint and update weights before doing anything + self._load_checkpoint() + self.checkpoint_manager.update_weights(self.global_steps) + + current_epoch = self.global_steps // len(self.train_dataloader) + + SkipManager.init(self.config) + + # perform validation before training + # currently, we only support validation using the reward_function. + if self.config.trainer.get("val_before_train", True): + val_metrics = self._validate() + assert val_metrics, f"{val_metrics=}" + pprint(f"Initial validation metrics: {val_metrics}") + logger.log(data=val_metrics, step=self.global_steps) + if self.config.trainer.get("val_only", False): + self._shutdown_dump_executor() + return + + # add tqdm + progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress") + + # we start from step 1 + self.global_steps += 1 + last_val_metrics = None + self.max_steps_duration = 0 + + SkipManager.set_step(self.global_steps) + + prev_step_profile = False + curr_step_profile = ( + self.global_steps in self.config.global_profiler.steps + if self.config.global_profiler.steps is not None + else False + ) + next_step_profile = False + + for epoch in range(current_epoch, self.config.trainer.total_epochs): + for batch_dict in self.train_dataloader: + if hasattr(self.actor_rollout_wg, "async_calls_finalize_fn_exec"): + self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=False) + metrics = {} + timing_raw = {} + + with marked_timer("start_profile", timing_raw): + self._start_profiling( + not prev_step_profile and curr_step_profile + if self.config.global_profiler.profile_continuous_steps + else curr_step_profile + ) + batch: DataProto = DataProto.from_single_dict(batch_dict) + batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature + + # add uid to batch + batch.non_tensor_batch["uid"] = np.array( + [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object + ) + + gen_batch = self._get_gen_batch(batch) + + # pass global_steps to trace + gen_batch.meta_info["global_steps"] = self.global_steps + rollout_n = self.config.actor_rollout_ref.rollout.n + gen_batch_output = gen_batch.repeat(repeat_times=rollout_n, interleave=True) + + if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX: + # NOTE: REMAX needs one sampled rollout plus one greedy baseline per prompt. + # Keep them in a single agent-loop/vLLM request to avoid sending a second + # rollout after replicas have been put to sleep, which can leave async vLLM + # engines in an invalid state for multi-turn agent workloads. + gen_batch_output.non_tensor_batch["__do_sample__"] = np.ones(len(gen_batch_output), dtype=bool) + gen_baseline_batch = gen_batch.slice(0, None) + gen_baseline_batch.non_tensor_batch["__do_sample__"] = np.zeros(len(gen_baseline_batch), dtype=bool) + combined_gen_batch = DataProto.concat([gen_batch_output, gen_baseline_batch]) + num_sampled_prompts = len(gen_batch_output) + else: + combined_gen_batch = gen_batch_output + num_sampled_prompts = len(gen_batch_output) + + is_last_step = self.global_steps >= self.total_training_steps + with marked_timer("step", timing_raw): + # generate a batch + with marked_timer("gen", timing_raw, color="red"): + if curr_step_profile: + self.llm_server_manager.start_profile() + combined_gen_output = self.async_rollout_manager.generate_sequences(combined_gen_batch) + self.checkpoint_manager.sleep_replicas() + if curr_step_profile: + self.llm_server_manager.stop_profile() + + timing_raw.update(combined_gen_output.meta_info["timing"]) + combined_gen_output.meta_info.pop("timing", None) + + gen_batch_output = combined_gen_output.slice(0, num_sampled_prompts) + if "__do_sample__" in gen_batch_output.non_tensor_batch: + gen_batch_output.pop(non_tensor_batch_keys=["__do_sample__"]) + + if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX: + gen_baseline_output = combined_gen_output.slice(num_sampled_prompts, None) + if "__do_sample__" in gen_baseline_output.non_tensor_batch: + gen_baseline_output.pop(non_tensor_batch_keys=["__do_sample__"]) + + if self.use_rm and "rm_scores" not in gen_baseline_output.batch.keys(): + baseline_reward = self._compute_reward_colocate(gen_baseline_output) + gen_baseline_output = gen_baseline_output.union(baseline_reward) + + reward_baseline_tensor = gen_baseline_output.batch["rm_scores"].sum(dim=-1) + batch.batch["reward_baselines"] = reward_baseline_tensor + + del gen_baseline_output + del combined_gen_batch, combined_gen_output + # repeat to align with repeated responses in rollout + batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + batch = batch.union(gen_batch_output) + + if "response_mask" not in batch.batch.keys(): + batch.batch["response_mask"] = compute_response_mask(batch) + # Balance the number of valid tokens across DP ranks. + # NOTE: This usually changes the order of data in the `batch`, + # which won't affect the advantage calculation (since it's based on uid), + # but might affect the loss calculation (due to the change of mini-batching). + if self.config.trainer.balance_batch: + self._balance_batch(batch, metrics=metrics) + + # compute global_valid tokens + batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist() + # get images_seqlens + images_seqlens_all = [] + for multi_modal_input in batch.non_tensor_batch["multi_modal_inputs"]: + if "image_grid_thw" not in multi_modal_input.keys(): + continue + images_seqlens_all.extend(multi_modal_input["images_seqlens"].tolist()) + batch.meta_info["images_seqlens"] = images_seqlens_all + with marked_timer("reward", timing_raw, color="yellow"): + # compute reward model score + if self.use_rm and "rm_scores" not in batch.batch.keys(): + batch_reward = self._compute_reward_colocate(batch) + batch = batch.union(batch_reward) + + # extract reward_tensor and reward_extra_infos_dict for training + reward_tensor, reward_extra_infos_dict = extract_reward(batch) + + # Operating Mode Selection: + # - Bypass mode: Sets old_log_probs = rollout_log_probs (2 policies: π_rollout, π_θ) + # - Decoupled mode: Recomputes old_log_probs as proximal anchor (3 policies: π_rollout, π_old, π_θ) + # Note: π_old computed once per data batch, serves as stable reference during mini-batch updates + rollout_corr_config = self.config.algorithm.get("rollout_correction", None) + bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get("bypass_mode", False) + if bypass_recomputing_logprobs: # Use `rollout_log_probs` + from verl.trainer.ppo.rollout_corr_helper import apply_bypass_mode + + apply_bypass_mode( + batch=batch, + rollout_corr_config=rollout_corr_config, + policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss, + ) + else: # Recompute old_log_probs + with marked_timer("old_log_prob", timing_raw, color="blue"): + old_log_prob, old_log_prob_mfu = self._compute_old_log_prob(batch) + entropys = old_log_prob.batch["entropys"] + response_masks = batch.batch["response_mask"] + actor_config = self.config.actor_rollout_ref.actor + entropy_agg = agg_loss( + loss_mat=entropys, + loss_mask=response_masks, + loss_agg_mode=actor_config.loss_agg_mode, + loss_scale_factor=actor_config.loss_scale_factor, + ) + old_log_prob_metrics = { + "actor/entropy": entropy_agg.detach().item(), + "perf/mfu/actor_infer": old_log_prob_mfu, + } + metrics.update(old_log_prob_metrics) + old_log_prob.batch.pop("entropys") + if "routed_experts" in batch.batch and "routed_experts" in old_log_prob.batch: + raise ValueError( + "Detected conflicting router replay configuration: " + "router_replay.mode='R2' and enable_rollout_routing_replay=True " + "cannot be enabled simultaneously. " + "The enable_rollout_routing_replay option is only used in R3 mode; " + "it should not be set when using R2 mode." + ) + batch = batch.union(old_log_prob) + if "rollout_log_probs" in batch.batch.keys(): + # TODO: we may want to add diff of probs too. + from verl.utils.debug.metrics import calculate_debug_metrics + + metrics.update(calculate_debug_metrics(batch)) + + assert "old_log_probs" in batch.batch, f'"old_log_prob" not in {batch.batch.keys()=}' + if self.use_reference_policy: + # compute reference log_prob + with marked_timer(str(Role.RefPolicy), timing_raw, color="olive"): + ref_log_prob = self._compute_ref_log_prob(batch) + batch = batch.union(ref_log_prob) + + # compute values + if self.use_critic: + with marked_timer("values", timing_raw, color="cyan"): + values = self._compute_values(batch) + batch = batch.union(values) + + with marked_timer("adv", timing_raw, color="brown"): + # we combine with rule-based rm + reward_extra_infos_dict: dict[str, list] + batch.batch["token_level_scores"] = reward_tensor + + if reward_extra_infos_dict: + batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()}) + + # compute rewards. apply_kl_penalty if available + if self.config.algorithm.use_kl_in_reward: + batch, kl_metrics = apply_kl_penalty( + batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty + ) + metrics.update(kl_metrics) + else: + batch.batch["token_level_rewards"] = batch.batch["token_level_scores"] + + # Compute rollout correction: IS weights, rejection sampling, and metrics + # Only runs in decoupled mode (computes once per batch using stable π_old) + # In bypass mode, this is skipped - actor computes metrics from evolving π_θ vs π_rollout + if ( + rollout_corr_config is not None + and "rollout_log_probs" in batch.batch + and not bypass_recomputing_logprobs # Only in decoupled mode + ): + from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_add_to_batch + + # Compute IS weights, apply rejection sampling, compute metrics + batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config) + # IS and off-policy metrics already have rollout_corr/ prefix + metrics.update(is_metrics) + + # compute advantages, executed on the driver process + norm_adv_by_std_in_grpo = self.config.algorithm.get( + "norm_adv_by_std_in_grpo", True + ) # GRPO adv normalization factor + + batch = compute_advantage( + batch, + adv_estimator=self.config.algorithm.adv_estimator, + gamma=self.config.algorithm.gamma, + lam=self.config.algorithm.lam, + num_repeat=self.config.actor_rollout_ref.rollout.n, + norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo, + config=self.config.algorithm, + ) + # update critic + if self.use_critic: + with marked_timer("update_critic", timing_raw, color="pink"): + critic_output = self._update_critic(batch) + critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"]) + metrics.update(critic_output_metrics) + + # implement critic warmup + if self.config.trainer.critic_warmup > self.global_steps: + # Still in critic warmup, only update weights to wake up rollout replicas. + self.checkpoint_manager.update_weights(self.global_steps) + else: + # update actor + with marked_timer("teacher_logprobs", timing_raw, color="purple"): + # No-op unless distillation.teacher_backend == 'fsdp'. + batch = self._inject_fsdp_teacher_data(batch) + with marked_timer("update_actor", timing_raw, color="red"): + actor_output = self._update_actor(batch) + + # Check if the ESI (Elastic Server Instance)/training plan is close to expiration. + esi_close_to_expiration = should_save_ckpt_esi( + max_steps_duration=self.max_steps_duration, + redundant_time=self.config.trainer.esi_redundant_time, + ) + # Check if the conditions for saving a checkpoint are met. + # The conditions include a mandatory condition (1) and + # one of the following optional conditions (2/3/4): + # 1. The save frequency is set to a positive value. + # 2. It's the last training step. + # 3. The current step number is a multiple of the save frequency. + # 4. The ESI(Elastic Server Instance)/training plan is close to expiration. + if self.config.trainer.save_freq > 0 and ( + is_last_step + or self.global_steps % self.config.trainer.save_freq == 0 + or esi_close_to_expiration + ): + if esi_close_to_expiration: + print("Force saving checkpoint: ESI instance expiration approaching.") + with marked_timer("save_checkpoint", timing_raw, color="green"): + self._save_checkpoint() + + # update weights from trainer to rollout + with marked_timer("update_weights", timing_raw, color="red"): + self.checkpoint_manager.update_weights(self.global_steps) + + actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) + metrics.update(actor_output_metrics) + + # Log rollout generations if enabled + rollout_data_dir = self.config.trainer.get("rollout_data_dir", None) + if rollout_data_dir: + self._log_rollout_data(batch, reward_extra_infos_dict, timing_raw, rollout_data_dir) + + # validate + if self.config.trainer.test_freq > 0 and ( + is_last_step or self.global_steps % self.config.trainer.test_freq == 0 + ): + with marked_timer("testing", timing_raw, color="green"): + val_metrics: dict = self._validate() + if is_last_step: + last_val_metrics = val_metrics + metrics.update(val_metrics) + + with marked_timer("stop_profile", timing_raw): + next_step_profile = ( + self.global_steps + 1 in self.config.global_profiler.steps + if self.config.global_profiler.steps is not None + else False + ) + self._stop_profiling( + curr_step_profile and not next_step_profile + if self.config.global_profiler.profile_continuous_steps + else curr_step_profile + ) + prev_step_profile = curr_step_profile + curr_step_profile = next_step_profile + + steps_duration = timing_raw["step"] + self.max_steps_duration = max(self.max_steps_duration, steps_duration) + + # training metrics + metrics.update( + { + "training/global_step": self.global_steps, + "training/epoch": epoch, + } + ) + # collect metrics + metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) + # GDPO per-component reward metrics + gdpo_reward_keys = self.config.algorithm.get("gdpo_reward_keys", None) + if gdpo_reward_keys and self.config.algorithm.adv_estimator in ("gdpo", AdvantageEstimator.GDPO): + for key in gdpo_reward_keys: + if key in batch.non_tensor_batch: + vals = np.asarray(batch.non_tensor_batch[key], dtype=np.float32) + metrics[f"gdpo/{key}/mean"] = float(np.mean(vals)) + metrics[f"gdpo/{key}/std"] = float(np.std(vals)) + metrics[f"gdpo/{key}/max"] = float(np.max(vals)) + metrics[f"gdpo/{key}/min"] = float(np.min(vals)) + metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) + # TODO: implement actual tflpo and theoretical tflpo + n_gpus = self.resource_pool_manager.get_n_gpus() + metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus)) + # compute variance proxy metrics + gradient_norm = metrics.get("actor/grad_norm", None) + metrics.update(compute_variance_proxy_metrics(batch=batch, gradient_norm=gradient_norm)) + # Note: mismatch metrics (KL, PPL, etc.) are collected at line 1179 after advantage computation + + # Per-request spec decode metrics. + metrics.update( + compute_spec_decode_metrics( + batch.non_tensor_batch.get("spec_num_draft_tokens", None), + batch.non_tensor_batch.get("spec_num_accepted_tokens", None), + batch.non_tensor_batch.get("spec_num_verify_steps", None), + ) + ) + + # TODO: make a canonical logger that supports various backend + logger.log(data=metrics, step=self.global_steps) + + progress_bar.update(1) + self.global_steps += 1 + SkipManager.set_step(self.global_steps) + + if is_last_step: + if hasattr(self.actor_rollout_wg, "async_calls_finalize_fn_exec"): + self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=True) + self._shutdown_dump_executor() + pprint(f"Final validation metrics: {last_val_metrics}") + progress_bar.close() + return + + # this is experimental and may be changed/removed in the future + # in favor of a general-purpose data buffer pool + if hasattr(self.train_dataset, "on_batch_end"): + # The dataset may be changed after each training batch + self.train_dataset.on_batch_end(batch=batch) + + # Ensure dump executor is shut down when training loop ends without reaching is_last_step + self._shutdown_dump_executor() diff --git a/verl/workers/config/distillation.py b/verl/workers/config/distillation.py index b58f8bccffd..4c4e7bf69a8 100644 --- a/verl/workers/config/distillation.py +++ b/verl/workers/config/distillation.py @@ -1,320 +1,362 @@ -# Copyright 2024 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. - -import logging -import os -from dataclasses import dataclass, field -from typing import Optional - -from verl.base_config import BaseConfig -from verl.utils.config import omega_conf_to_dataclass - -from .rollout import RolloutConfig - -__all__ = ["DistillationLossConfig", "DistillationTeacherModelConfig", "DistillationConfig"] - -logger = logging.getLogger(__name__) -logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) - - -@dataclass -class DistillationLossConfig(BaseConfig): - """Configuration for distillation loss settings. - - loss_mode (str): - Distillation loss function to use. - topk (int, optional): - Number of top tokens to consider for top-k distillation losses. - use_task_rewards (bool): - Whether to include task rewards alongside distillation loss. - distillation_loss_coef (float): - Coefficient for distillation loss when combined with task rewards. - loss_max_clamp (float, optional): - Maximum value to clamp distillation loss. If None, no clamping is applied. - log_prob_min_clamp (float, optional): - Minimum value to clamp log probabilities for stability, e.g., log q - log p where p or q are - very close to zero. If None, no clamping is applied. - use_policy_gradient (bool): - Whether to incorporate distillation loss as a reward, as done - by https://thinkingmachines.ai/blog/on-policy-distillation/. Recommended to use loss_mode=k1. - Otherwise, distillation loss is directly backpropagated as a supervised loss, - as in https://arxiv.org/abs/2306.13649. Recommended to use loss_mode=k3 or forward_kl_topk. - policy_loss_mode (str): - Name of the policy loss to use when use_policy_gradient is true. - clip_ratio (float): - PPO clipping ratio for policy loss. - clip_ratio_low (float): - Lower bound for PPO clipping ratio. - clip_ratio_high (float): - Upper bound for PPO clipping ratio. - loss_settings (DistillationLossSettings, optional): - Runtime-populated settings based on loss_mode. Not set by user. - """ - - loss_mode: str = "k3" - topk: Optional[int] = 128 - use_task_rewards: bool = True - distillation_loss_coef: float = 1.0 - loss_max_clamp: Optional[float] = 10.0 - log_prob_min_clamp: Optional[float] = -10.0 - - # Chunked top-K log-probs (opt-in, avoids [B, T, V] log_softmax buffer - # at long context). Only consumed by ``loss_mode='forward_kl_topk'``. - # Default ``False`` to preserve short-context performance (chunked path - # has ~6x time overhead at N=14K, V=152K). Set ``True`` when hitting OOM - # at long context (>=64K tokens, V=152K) where the baseline path OOMs. - use_chunked_topk: bool = False - # Tokens per chunk along (B*T) when ``use_chunked_topk=True``. Larger - # chunks reduce kernel-launch overhead but increase per-chunk memory; - # smaller chunks reduce per-chunk memory but increase kernel-launch - # overhead (saved-tensor total stays constant either way). - chunked_topk_chunk_size: int = 4096 - - use_policy_gradient: bool = True - policy_loss_mode: str = "vanilla" - clip_ratio: float = 0.2 - clip_ratio_low: float = 0.2 - clip_ratio_high: float = 0.2 - - # Store global batch info for loss aggregation: - # dp_size: data parallel size - # batch_num_tokens: number of valid tokens in global batch - # global_batch_size: global batch size - global_batch_info: dict = field(default_factory=dict) - - # Store distillation loss settings for computing the specified loss_mode - # Not set by user, populated at runtime - loss_settings: Optional[dict] = None - - def __post_init__(self): - self._mutable_fields.add("loss_settings") - from verl.trainer.distillation.losses import DistillationLossSettings, get_distillation_loss_settings - - self.loss_settings: DistillationLossSettings = get_distillation_loss_settings(self.loss_mode) - - if self.policy_loss_mode != "vanilla": - raise NotImplementedError( - f"Only vanilla policy loss is currently supported when use_policy_gradient is True, " - f"but got {self.policy_loss_mode}." - ) - - if self.use_policy_gradient and self.loss_mode == "forward_kl_topk": - print( - "WARNING: forward_kl_topk is most effective as a supervised distillation loss " - "(use_policy_gradient=False). With policy gradient, the update uses only the sampled" - " token's logprob ∇logπ(a), so the top-k distributional signal (how non-sampled logits " - "should move) is largely unused." - ) - - if not self.use_policy_gradient and self.loss_mode == "k1": - raise ValueError( - "Directly backpropagating k1 loss is incorrect since gradient of k1 loss" - " wrt model weights does not depend on teacher log probabilities." - ) - - -@dataclass -class DistillationTeacherModelConfig(BaseConfig): - """Configuration for on-policy distillation teacher. - - key (str, optional): - Identifier to route examples to the teacher model in multi-teacher setting. - model_path (str, optional): - Model path for the teacher model. Can be a local path or a Hugging Face model - inference (RolloutConfig): - Rollout configuration for the teacher model inference during distillation. - num_replicas (int): - Number of inference replicas of this teacher to launch. Each replica occupies - `per_replica_world_size` GPUs (= inference.data_parallel_size * - inference.tensor_model_parallel_size * inference.pipeline_model_parallel_size), - so the teacher's total GPU footprint is - `num_replicas * per_replica_world_size`. - """ - - _mutable_fields = BaseConfig._mutable_fields | {"num_replicas", "key"} - - key: Optional[str] = None - model_path: Optional[str] = None - inference: RolloutConfig = field(default_factory=RolloutConfig) - num_replicas: Optional[int] = 0 - - @property - def per_replica_world_size(self) -> int: - return ( - self.inference.tensor_model_parallel_size - * self.inference.data_parallel_size - * self.inference.pipeline_model_parallel_size - ) - - @property - def world_size(self) -> int: - return self.num_replicas * self.per_replica_world_size - - def check_configured(self): - if self.model_path is None: - raise ValueError("model_path must be specified for distillation teacher model config.") - if self.key is None: - raise ValueError("key must be specified for distillation teacher model config.") - if self.num_replicas is None: - raise ValueError("num_replicas must be specified for distillation teacher model config.") - - def validate_and_prepare_for_distillation(self, use_topk: bool, topk: Optional[int]) -> None: - # Prompt + Response from student are fed into teacher as context - max_model_len = self.inference.max_model_len - student_prompt_length = self.inference.prompt_length - student_response_length = self.inference.response_length - required_context_len = student_prompt_length + student_response_length + 1 - if max_model_len is not None and required_context_len > max_model_len: - raise ValueError( - "Distillation teacher inference requires room for the student prompt, the full student " - f"response, and one generated token, but got {student_prompt_length=}, " - f"{student_response_length=}, {required_context_len=}, {max_model_len=}." - ) - self.inference.prompt_length = self.inference.prompt_length + self.inference.response_length - self.inference.response_length = 1 - self._validate_topk_logprobs(use_topk=use_topk, topk=topk) - - def _validate_topk_logprobs(self, use_topk: bool, topk: Optional[int]) -> None: - if not use_topk: - return - if topk is None: - raise ValueError("topk must be specified when use_topk is True.") - - engine_name = self.inference.name - engine_kwargs = self.inference.engine_kwargs - match engine_name: - case "vllm": - vllm_engine_kwargs = dict(engine_kwargs.get("vllm", {})) - max_logprobs = vllm_engine_kwargs.get("max_logprobs") - if max_logprobs is None: - vllm_engine_kwargs["max_logprobs"] = topk - max_logprobs = topk - if max_logprobs < topk: - raise ValueError( - f"VLLM max_logprobs ({max_logprobs}) must be >= distillation_loss topk " - f"({topk}) to enable distillation loss computation." - ) - engine_kwargs["vllm"] = vllm_engine_kwargs - case "sglang": - # SGLang's top_logprobs_num is a per-request parameter, so there is no - # engine-boot cap to align (unlike vLLM's max_logprobs). The async - # server translates sampling_params["prompt_logprobs"] into - # return_logprob + logprob_start_len=0 + top_logprobs_num at call time. - pass - case _: - raise NotImplementedError( - f"DistillationTeacherModelConfig does not support inference engine {engine_name}" - ) - - -@dataclass -class DistillationConfig(BaseConfig): - """Configuration for on-policy distillation. - - enabled (bool): - Whether on-policy distillation is enabled. - n_gpus_per_node (int): - Number of GPUs per node in the teacher resource pool. - nnodes (int): - Number of nodes in the teacher resource pool. - teacher_models (dict[str, TeacherModelConfig]): - Configurations for teacher models used for multi-teacher distillation. - teacher_key (str): - Key to route examples to the appropriate teacher model in multi-teacher setups. Should correspond to a field in - the data proto, e.g., data_source. - distillation_loss (DistillationLossConfig): - Configuration for distillation loss settings. - - NOTE: The `teacher_model` entry is in the `teacher_models` dict by default. - Since it is popped when other teacher entries are added, using `teacher_model` as - one of several keys silently drops it. For example, the following CLI overrides result - in ONLY `teacher_model2` being used: - - ```bash - distillation.teacher_models.teacher_model.key=openai/gsm8k - distillation.teacher_models.teacher_model.model_path=Qwen/Qwen3-4B - +distillation.teacher_models.teacher_model2.key=hiyouga/geometry3k - +distillation.teacher_models.teacher_model2.model_path=Qwen/Qwen3-VL-4B-Instruct - ``` - Instead, give the first teacher a different name: - - ```bash - +distillation.teacher_models.teacher_model1.key=openai/gsm8k - +distillation.teacher_models.teacher_model1.model_path=Qwen/Qwen3-4B - +distillation.teacher_models.teacher_model2.key=hiyouga/geometry3k - +distillation.teacher_models.teacher_model2.model_path=Qwen/Qwen3-VL-4B-Instruct - ``` - """ - - _mutable_fields = BaseConfig._mutable_fields | {"teacher_models", "n_gpus_per_node", "nnodes"} - - enabled: bool = False - n_gpus_per_node: int = 0 - nnodes: int = 0 - teacher_models: dict[str, DistillationTeacherModelConfig] = field(default_factory=dict) - teacher_key: str = "data_source" - distillation_loss: DistillationLossConfig = field(default_factory=DistillationLossConfig) - - def __post_init__(self): - if not self.enabled: - return - - self.teacher_models = self._resolve_teacher_models() - teacher_world_size_sum = 0 - for teacher_model in self.teacher_models.values(): - teacher_model.validate_and_prepare_for_distillation( - use_topk=self.distillation_loss.loss_settings.use_topk, - topk=self.distillation_loss.topk, - ) - teacher_world_size_sum += teacher_model.world_size - total_pool_size = self.n_gpus_per_node * self.nnodes - if teacher_world_size_sum != total_pool_size: - raise ValueError( - f"Sum of teacher (num_replicas * per_replica_world_size) ({teacher_world_size_sum}) must match " - f"the distillation resource pool size " - f"({self.n_gpus_per_node=} * {self.nnodes=} = {total_pool_size})." - ) - - def _resolve_teacher_models(self) -> dict[str, DistillationTeacherModelConfig]: - assert "teacher_model" in self.teacher_models - if len(self.teacher_models) == 1: - # Single teacher occupies the entire teacher resource pool. - teacher_model = self.teacher_models["teacher_model"] - inference = teacher_model.inference - per_replica = ( - inference.tensor_model_parallel_size - * inference.data_parallel_size - * inference.pipeline_model_parallel_size - ) - pool_size = self.n_gpus_per_node * self.nnodes - if pool_size % per_replica != 0: - raise ValueError( - f"Single teacher's per_replica_world_size ({per_replica}) must divide the distillation " - f"resource pool size ({self.n_gpus_per_node=} * {self.nnodes=} = {pool_size})." - ) - teacher_model.num_replicas = pool_size // per_replica - teacher_model.key = "default" - else: - # Multiple teachers: remove default single teacher config - self.teacher_models.pop("teacher_model") - - # Teacher models dict is keyed by teacher_key instead of YAML entry name - teacher_models = {} - for teacher_config in self.teacher_models.values(): - teacher_config = omega_conf_to_dataclass(teacher_config, dataclass_type=DistillationTeacherModelConfig) - teacher_config.check_configured() - if teacher_config.key in teacher_models: - raise ValueError(f"Duplicate teacher key {teacher_config.key} found in teacher models.") - teacher_models[teacher_config.key] = teacher_config - return teacher_models +# Copyright 2024 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. + +import logging +import os +from dataclasses import dataclass, field +from typing import Optional + +from verl.base_config import BaseConfig +from verl.utils.config import omega_conf_to_dataclass + +from .engine import FSDPEngineConfig +from .rollout import RolloutConfig + +__all__ = ["DistillationLossConfig", "DistillationTeacherModelConfig", "DistillationConfig"] + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +@dataclass +class DistillationLossConfig(BaseConfig): + """Configuration for distillation loss settings. + + loss_mode (str): + Distillation loss function to use. + topk (int, optional): + Number of top tokens to consider for top-k distillation losses. + use_task_rewards (bool): + Whether to include task rewards alongside distillation loss. + distillation_loss_coef (float): + Coefficient for distillation loss when combined with task rewards. + loss_max_clamp (float, optional): + Maximum value to clamp distillation loss. If None, no clamping is applied. + log_prob_min_clamp (float, optional): + Minimum value to clamp log probabilities for stability, e.g., log q - log p where p or q are + very close to zero. If None, no clamping is applied. + use_policy_gradient (bool): + Whether to incorporate distillation loss as a reward, as done + by https://thinkingmachines.ai/blog/on-policy-distillation/. Recommended to use loss_mode=k1. + Otherwise, distillation loss is directly backpropagated as a supervised loss, + as in https://arxiv.org/abs/2306.13649. Recommended to use loss_mode=k3 or forward_kl_topk. + policy_loss_mode (str): + Name of the policy loss to use when use_policy_gradient is true. + clip_ratio (float): + PPO clipping ratio for policy loss. + clip_ratio_low (float): + Lower bound for PPO clipping ratio. + clip_ratio_high (float): + Upper bound for PPO clipping ratio. + loss_settings (DistillationLossSettings, optional): + Runtime-populated settings based on loss_mode. Not set by user. + """ + + loss_mode: str = "k3" + topk: Optional[int] = 128 + use_task_rewards: bool = True + distillation_loss_coef: float = 1.0 + loss_max_clamp: Optional[float] = 10.0 + log_prob_min_clamp: Optional[float] = -10.0 + + # Chunked top-K log-probs (opt-in, avoids [B, T, V] log_softmax buffer + # at long context). Only consumed by ``loss_mode='forward_kl_topk'``. + # Default ``False`` to preserve short-context performance (chunked path + # has ~6x time overhead at N=14K, V=152K). Set ``True`` when hitting OOM + # at long context (>=64K tokens, V=152K) where the baseline path OOMs. + use_chunked_topk: bool = False + # Tokens per chunk along (B*T) when ``use_chunked_topk=True``. Larger + # chunks reduce kernel-launch overhead but increase per-chunk memory; + # smaller chunks reduce per-chunk memory but increase kernel-launch + # overhead (saved-tensor total stays constant either way). + chunked_topk_chunk_size: int = 4096 + + use_policy_gradient: bool = True + policy_loss_mode: str = "vanilla" + clip_ratio: float = 0.2 + clip_ratio_low: float = 0.2 + clip_ratio_high: float = 0.2 + + # Store global batch info for loss aggregation: + # dp_size: data parallel size + # batch_num_tokens: number of valid tokens in global batch + # global_batch_size: global batch size + global_batch_info: dict = field(default_factory=dict) + + # Store distillation loss settings for computing the specified loss_mode + # Not set by user, populated at runtime + loss_settings: Optional[dict] = None + + def __post_init__(self): + self._mutable_fields.add("loss_settings") + from verl.trainer.distillation.losses import DistillationLossSettings, get_distillation_loss_settings + + self.loss_settings: DistillationLossSettings = get_distillation_loss_settings(self.loss_mode) + + if self.policy_loss_mode != "vanilla": + raise NotImplementedError( + f"Only vanilla policy loss is currently supported when use_policy_gradient is True, " + f"but got {self.policy_loss_mode}." + ) + + if self.use_policy_gradient and self.loss_mode == "forward_kl_topk": + print( + "WARNING: forward_kl_topk is most effective as a supervised distillation loss " + "(use_policy_gradient=False). With policy gradient, the update uses only the sampled" + " token's logprob ∇logπ(a), so the top-k distributional signal (how non-sampled logits " + "should move) is largely unused." + ) + + if self.use_policy_gradient and self.loss_mode == "reverse_kl_topk": + raise ValueError( + "reverse_kl_topk is a supervised distillation loss and requires use_policy_gradient=False." + ) + + if not self.use_policy_gradient and self.loss_mode == "k1": + raise ValueError( + "Directly backpropagating k1 loss is incorrect since gradient of k1 loss" + " wrt model weights does not depend on teacher log probabilities." + ) + + +@dataclass +class DistillationTeacherModelConfig(BaseConfig): + """Configuration for on-policy distillation teacher. + + key (str, optional): + Identifier to route examples to the teacher model in multi-teacher setting. + model_path (str, optional): + Model path for the teacher model. Can be a local path or a Hugging Face model + inference (RolloutConfig): + Rollout configuration for the teacher model inference during distillation. + num_replicas (int): + Number of inference replicas of this teacher to launch. Each replica occupies + `per_replica_world_size` GPUs (= inference.data_parallel_size * + inference.tensor_model_parallel_size * inference.pipeline_model_parallel_size), + so the teacher's total GPU footprint is + `num_replicas * per_replica_world_size`. + """ + + _mutable_fields = BaseConfig._mutable_fields | {"num_replicas", "key"} + + key: Optional[str] = None + model_path: Optional[str] = None + inference: RolloutConfig = field(default_factory=RolloutConfig) + num_replicas: Optional[int] = 0 + + @property + def per_replica_world_size(self) -> int: + return ( + self.inference.tensor_model_parallel_size + * self.inference.data_parallel_size + * self.inference.pipeline_model_parallel_size + ) + + @property + def world_size(self) -> int: + return self.num_replicas * self.per_replica_world_size + + def check_configured(self): + if self.model_path is None: + raise ValueError("model_path must be specified for distillation teacher model config.") + if self.key is None: + raise ValueError("key must be specified for distillation teacher model config.") + if self.num_replicas is None: + raise ValueError("num_replicas must be specified for distillation teacher model config.") + + def validate_and_prepare_for_distillation(self, use_topk: bool, topk: Optional[int]) -> None: + # Prompt + Response from student are fed into teacher as context + max_model_len = self.inference.max_model_len + student_prompt_length = self.inference.prompt_length + student_response_length = self.inference.response_length + required_context_len = student_prompt_length + student_response_length + 1 + if max_model_len is not None and required_context_len > max_model_len: + raise ValueError( + "Distillation teacher inference requires room for the student prompt, the full student " + f"response, and one generated token, but got {student_prompt_length=}, " + f"{student_response_length=}, {required_context_len=}, {max_model_len=}." + ) + self.inference.prompt_length = self.inference.prompt_length + self.inference.response_length + self.inference.response_length = 1 + self._validate_topk_logprobs(use_topk=use_topk, topk=topk) + + def _validate_topk_logprobs(self, use_topk: bool, topk: Optional[int]) -> None: + if not use_topk: + return + if topk is None: + raise ValueError("topk must be specified when use_topk is True.") + + engine_name = self.inference.name + engine_kwargs = self.inference.engine_kwargs + match engine_name: + case "vllm": + vllm_engine_kwargs = dict(engine_kwargs.get("vllm", {})) + max_logprobs = vllm_engine_kwargs.get("max_logprobs") + if max_logprobs is None: + vllm_engine_kwargs["max_logprobs"] = topk + max_logprobs = topk + if max_logprobs < topk: + raise ValueError( + f"VLLM max_logprobs ({max_logprobs}) must be >= distillation_loss topk " + f"({topk}) to enable distillation loss computation." + ) + engine_kwargs["vllm"] = vllm_engine_kwargs + case "sglang": + # SGLang's top_logprobs_num is a per-request parameter, so there is no + # engine-boot cap to align (unlike vLLM's max_logprobs). The async + # server translates sampling_params["prompt_logprobs"] into + # return_logprob + logprob_start_len=0 + top_logprobs_num at call time. + pass + case _: + raise NotImplementedError( + f"DistillationTeacherModelConfig does not support inference engine {engine_name}" + ) + + +@dataclass +class DistillationConfig(BaseConfig): + """Configuration for on-policy distillation. + + enabled (bool): + Whether on-policy distillation is enabled. + n_gpus_per_node (int): + Number of GPUs per node in the teacher resource pool. + nnodes (int): + Number of nodes in the teacher resource pool. + teacher_models (dict[str, TeacherModelConfig]): + Configurations for teacher models used for multi-teacher distillation. + teacher_key (str): + Key to route examples to the appropriate teacher model in multi-teacher setups. Should correspond to a field in + the data proto, e.g., data_source. + distillation_loss (DistillationLossConfig): + Configuration for distillation loss settings. + + NOTE: The `teacher_model` entry is in the `teacher_models` dict by default. + Since it is popped when other teacher entries are added, using `teacher_model` as + one of several keys silently drops it. For example, the following CLI overrides result + in ONLY `teacher_model2` being used: + + ```bash + distillation.teacher_models.teacher_model.key=openai/gsm8k + distillation.teacher_models.teacher_model.model_path=Qwen/Qwen3-4B + +distillation.teacher_models.teacher_model2.key=hiyouga/geometry3k + +distillation.teacher_models.teacher_model2.model_path=Qwen/Qwen3-VL-4B-Instruct + ``` + Instead, give the first teacher a different name: + + ```bash + +distillation.teacher_models.teacher_model1.key=openai/gsm8k + +distillation.teacher_models.teacher_model1.model_path=Qwen/Qwen3-4B + +distillation.teacher_models.teacher_model2.key=hiyouga/geometry3k + +distillation.teacher_models.teacher_model2.model_path=Qwen/Qwen3-VL-4B-Instruct + ``` + + Reverse-KL extension fields (only used when ``teacher_backend == 'fsdp'``): + teacher_backend (str): "vllm" (default, current behavior — teacher serves prompt_logprobs + via vLLM/SGLang rollout) or "fsdp" (teacher hosted as a frozen FSDP worker that + gathers logprobs at student-top-K IDs each train step). + teacher_chunk_size (int): Token-axis chunk size used by the FSDP teacher when gathering + log-probabilities at student-top-K IDs without materializing the full ``[B, T, V]`` + log-softmax tensor. + teacher_fsdp_config (Optional[FSDPEngineConfig]): FSDP wrap / offload / dtype config + applied to the frozen teacher worker. Reuses the same dataclass as actor / ref so the + existing FSDP engine builder can construct the teacher model. + """ + + _mutable_fields = BaseConfig._mutable_fields | {"teacher_models", "n_gpus_per_node", "nnodes"} + + enabled: bool = False + n_gpus_per_node: int = 0 + nnodes: int = 0 + teacher_models: dict[str, DistillationTeacherModelConfig] = field(default_factory=dict) + teacher_key: str = "data_source" + distillation_loss: DistillationLossConfig = field(default_factory=DistillationLossConfig) + + # Reverse-KL extension: vLLM rollout teacher (default) or frozen FSDP worker. + teacher_backend: str = "vllm" + teacher_chunk_size: int = 1024 + teacher_fsdp_config: Optional[FSDPEngineConfig] = None + + def __post_init__(self): + if not self.enabled: + return + + if self.teacher_backend not in ("vllm", "fsdp"): + raise ValueError(f"Unsupported teacher_backend {self.teacher_backend!r}; expected one of 'vllm' or 'fsdp'.") + + if self.teacher_backend == "fsdp": + # FSDP teacher only supports reverse_kl_topk for now; other modes need rollout outputs. + if self.distillation_loss.loss_mode != "reverse_kl_topk": + raise ValueError( + "teacher_backend='fsdp' currently only supports distillation_loss.loss_mode=" + "'reverse_kl_topk', " + f"got loss_mode={self.distillation_loss.loss_mode!r}." + ) + self.teacher_models = self._resolve_teacher_models(skip_inference_validation=True) + return + + self.teacher_models = self._resolve_teacher_models() + teacher_world_size_sum = 0 + for teacher_model in self.teacher_models.values(): + teacher_model.validate_and_prepare_for_distillation( + use_topk=self.distillation_loss.loss_settings.use_topk, + topk=self.distillation_loss.topk, + ) + teacher_world_size_sum += teacher_model.world_size + total_pool_size = self.n_gpus_per_node * self.nnodes + if teacher_world_size_sum != total_pool_size: + raise ValueError( + f"Sum of teacher (num_replicas * per_replica_world_size) ({teacher_world_size_sum}) must match " + f"the distillation resource pool size " + f"({self.n_gpus_per_node=} * {self.nnodes=} = {total_pool_size})." + ) + + def _resolve_teacher_models( + self, skip_inference_validation: bool = False + ) -> dict[str, DistillationTeacherModelConfig]: + assert "teacher_model" in self.teacher_models + if len(self.teacher_models) == 1: + # Single teacher occupies the entire teacher resource pool. + teacher_model = self.teacher_models["teacher_model"] + if skip_inference_validation: + # FSDP teacher does not consume the vLLM/SGLang inference pool sizing fields. + teacher_model.num_replicas = 1 + else: + inference = teacher_model.inference + per_replica = ( + inference.tensor_model_parallel_size + * inference.data_parallel_size + * inference.pipeline_model_parallel_size + ) + pool_size = self.n_gpus_per_node * self.nnodes + if pool_size % per_replica != 0: + raise ValueError( + f"Single teacher's per_replica_world_size ({per_replica}) must divide the distillation " + f"resource pool size ({self.n_gpus_per_node=} * {self.nnodes=} = {pool_size})." + ) + teacher_model.num_replicas = pool_size // per_replica + teacher_model.key = "default" + else: + # Multiple teachers: remove default single teacher config + self.teacher_models.pop("teacher_model") + + # Teacher models dict is keyed by teacher_key instead of YAML entry name + teacher_models = {} + for teacher_config in self.teacher_models.values(): + teacher_config = omega_conf_to_dataclass(teacher_config, dataclass_type=DistillationTeacherModelConfig) + teacher_config.check_configured() + if teacher_config.key in teacher_models: + raise ValueError(f"Duplicate teacher key {teacher_config.key} found in teacher models.") + teacher_models[teacher_config.key] = teacher_config + return teacher_models diff --git a/verl/workers/config/engine.py b/verl/workers/config/engine.py index 0389b68a1f7..289d9e68946 100644 --- a/verl/workers/config/engine.py +++ b/verl/workers/config/engine.py @@ -66,6 +66,9 @@ class EngineRouterReplayConfig(BaseConfig): record_file: Optional[str] = None replay_file: Optional[str] = None + # Allow runtime overrides of mode (e.g., TeacherFSDPWorker disables router replay). + _mutable_fields = BaseConfig._mutable_fields | {"mode"} + def __post_init__(self): """Validate router replay configuration.""" valid_modes = ["disabled", "R2", "R3"] @@ -101,11 +104,11 @@ class EngineConfig(BaseConfig): # whether to use dynamic bsz use_dynamic_bsz: bool = True # for training - max_token_len_per_gpu: int = None - micro_batch_size_per_gpu: int = None + max_token_len_per_gpu: Optional[int] = None + micro_batch_size_per_gpu: Optional[int] = None # for inference - infer_max_token_len_per_gpu: int = None - infer_micro_batch_size_per_gpu: int = None + infer_max_token_len_per_gpu: Optional[int] = None + infer_micro_batch_size_per_gpu: Optional[int] = None # whether use fuse lm head kernel use_fused_kernels: bool = False # TODO (this may conflict with the one in model config) diff --git a/verl/workers/engine/fsdp/transformer_impl.py b/verl/workers/engine/fsdp/transformer_impl.py index 322e78f8141..7d1d70479b7 100644 --- a/verl/workers/engine/fsdp/transformer_impl.py +++ b/verl/workers/engine/fsdp/transformer_impl.py @@ -1158,12 +1158,16 @@ def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict, lo ) # logits_processor_func return tensors with shape (1, total_nnz/sp_size) + # or (1, total_nnz/sp_size, K) when the processor emits per-token, top-K + # outputs (e.g. the FSDP-teacher reverse-KL pipeline's student-top-K stage). if distillation_use_topk: outputs = logits_processor_func(student_logits=logits_rmpad.unsqueeze(0), data=micro_batch) cu_seqlens = input_ids.offsets() for k, v in outputs.items(): v = v.squeeze(0) - assert v.shape == (logits_rmpad.shape[0],), ( + # First-dim only: v may be (total_nnz,) or (total_nnz, K) for + # the top-K path (e.g. student_topk_log_probs / student_topk_ids). + assert v.shape[0] == logits_rmpad.shape[0], ( f"logits_rmpad len: {logits_rmpad.shape[0]}, {k} shape: {v.shape}" ) if self.use_ulysses_sp: @@ -1274,7 +1278,9 @@ def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict, lo outputs = logits_processor_func(student_logits=logits_rmpad.unsqueeze(0), data=micro_batch) for k, v in outputs.items(): v = v.squeeze(0) - assert v.shape == (logits_rmpad.shape[0],), ( + # First-dim only: v may be (total_nnz,) or (total_nnz, K) for + # the top-K path (e.g. student_topk_log_probs / student_topk_ids). + assert v.shape[0] == logits_rmpad.shape[0], ( f"logits_rmpad len: {logits_rmpad.shape[0]}, {k} shape: {v.shape}" ) model_output[k] = torch.nested.nested_tensor_from_jagged(v, cu_seqlens) diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py index bc04c0e6464..8a6e46b02ec 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -650,6 +650,24 @@ def compute_log_prob(self, data: TensorDict) -> TensorDict: return output.cpu() if output is not None else None + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor")) + @DistProfiler.annotate(color="cyan", role="actor_compute_student_topk") + @_with_routing_replay_flag(enabled=False) + def compute_student_topk(self, data: TensorDict) -> TensorDict: + """Run a no-grad student forward and emit the per-token student top-K + log-probs and IDs. + + Used by the FSDP teacher pipeline (verl#6676): the student top-K must be + computed before the teacher forward, since the teacher gathers log-probs at + those IDs. Stamping ``compute_student_topk_only=True`` makes the actor's + distillation logits processor return the top-K instead of the full loss. + """ + tu.assign_non_tensor(data, compute_student_topk_only=True) + # Forward only — never build the teacher-conditional loss tensor. + tu.assign_non_tensor(data, compute_loss=True) + output = self.actor.infer_batch(data) + return output.cpu() if output is not None else None + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor")) @DistProfiler.annotate(color="red", role="actor_update") @_with_routing_replay_flag(enabled=True) diff --git a/verl/workers/teacher/__init__.py b/verl/workers/teacher/__init__.py new file mode 100644 index 00000000000..3ebaf1fd1ba --- /dev/null +++ b/verl/workers/teacher/__init__.py @@ -0,0 +1,25 @@ +# 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. +"""Teacher worker package for FSDP-served distillation teachers. + +The :class:`TeacherFSDPWorker` exposed here is the FSDP-backed alternative to +the vLLM teacher rollout managed by +:class:`verl.experimental.teacher_loop.MultiTeacherModelManager`. It is wired +into the trainer when ``distillation.teacher_backend == 'fsdp'``. +""" + +from .fsdp_teacher_worker import TeacherFSDPWorker +from .utils import chunked_gather_logprobs + +__all__ = ["TeacherFSDPWorker", "chunked_gather_logprobs"] diff --git a/verl/workers/teacher/fsdp_teacher_worker.py b/verl/workers/teacher/fsdp_teacher_worker.py new file mode 100644 index 00000000000..bd03d07fbed --- /dev/null +++ b/verl/workers/teacher/fsdp_teacher_worker.py @@ -0,0 +1,169 @@ +# 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. +"""Frozen teacher hosted as a forward-only FSDP worker. + +The student-top-K reverse-KL pipeline (verl#6676) needs the teacher's +log-probabilities evaluated *at student-top-K token IDs*, which the existing +vLLM-served teacher cannot produce (its prompt_logprobs path is restricted to +the teacher's own top-K). This module hosts the teacher as a Ray FSDP worker +that runs a single forward pass per train step and chunked-gathers +``log_softmax(teacher_logits)`` at IDs supplied by the trainer. + +The worker subclasses :class:`verl.workers.engine_workers.TrainingWorker` to +reuse the standard engine bootstrap (FSDP wrap, dtype, offload, dispatch +collect-info), pinning ``forward_only=True`` so the engine skips optimizer / +gradient setup. +""" + +from __future__ import annotations + +import logging +import os +from typing import Optional + +import torch +from tensordict import TensorDict + +from verl.single_controller.base.decorator import ( + Dispatch, + make_nd_compute_dataproto_dispatch_fn, + register, +) +from verl.utils import tensordict_utils as tu +from verl.utils.profiler import DistProfiler +from verl.workers.config import TrainingWorkerConfig +from verl.workers.engine_workers import TrainingWorker +from verl.workers.teacher.utils import chunked_gather_logprobs + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +class TeacherFSDPWorker(TrainingWorker): + """Forward-only FSDP worker that serves teacher log-probabilities. + + Follows the standard :class:`TrainingWorker` flow but force-flips + ``engine_config.forward_only = True``, so the engine is built without an optimizer. + Exposes :meth:`compute_logprobs_at_ids`, which runs a teacher forward and gathers + ``log_softmax(logits)`` at trainer-supplied token IDs without materializing the + full ``[B, T, V]`` log-softmax tensor. + + Notes: + - Ulysses sequence parallelism on the teacher is not yet supported; the + trainer must set ``teacher_fsdp_config.ulysses_sequence_parallel_size = 1``. + - The trainer must align ``topk_ids`` with the engine's rolled label + convention (``roll(input_ids_rmpad, -1)``); this worker consumes them as-is. + """ + + def __init__(self, config: TrainingWorkerConfig): + if config.engine_config is not None: + # Frozen teacher: no optimizer, no grad, no checkpointing. + config.engine_config.forward_only = True + # Disable router replay (only meaningful for stateful training rollouts). + if hasattr(config.engine_config, "router_replay"): + config.engine_config.router_replay.mode = "disabled" + super().__init__(config) + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="train"), blocking=False) + @DistProfiler.annotate(color="purple", role="teacher_compute_logprobs_at_ids") + def compute_logprobs_at_ids(self, data: TensorDict) -> Optional[TensorDict]: + """Run a teacher forward and gather log-probabilities at given token IDs. + + Expected ``data`` fields (in the rmpad / nested-tensor layout produced by + the trainer's pre-update step): + + * ``input_ids`` / ``position_ids``: nested tensors, jagged ``(bsz, seq_len_i)``. + * ``attention_mask``: nested tensor or ``None``. + * ``topk_ids``: nested ``(bsz, seq_len_i, K)`` int64, *already rolled by 1* + along the sequence axis so position ``i`` carries the IDs for logits ``i``. + * The engine's standard knobs (``use_remove_padding``, ``use_fused_kernels``, + etc.), copied by the trainer from the actor's engine config. + + Returns: + A :class:`TensorDict` with two nested ``(bsz, seq_len_i, K)`` fields: + ``teacher_on_student_logp`` (teacher log-probs at the student top-K) and + ``teacher_topk_ids`` (the teacher's own top-K IDs, for overlap diagnostics). + Returns ``None`` when the worker is not the output-collecting rank + (matches the :meth:`TrainingWorker.infer_batch` contract). + """ + if self.engine.use_ulysses_sp: + raise NotImplementedError( + "Ulysses SP is not yet supported by TeacherFSDPWorker.compute_logprobs_at_ids; " + "set teacher_fsdp_config.ulysses_sequence_parallel_size=1 for now." + ) + + # Inject the same engineering knobs the engine expects (mirrors TrainingWorker.infer_batch). + default_keys = dict( + use_remove_padding=self.model_config.get("use_remove_padding", False), + use_dynamic_bsz=self.engine_config.use_dynamic_bsz, + max_token_len_per_gpu=self.engine_config.infer_max_token_len_per_gpu, + micro_batch_size_per_gpu=self.engine_config.infer_micro_batch_size_per_gpu, + use_fused_kernels=False, # never use fused kernels: they don't expose logits + ) + for key, val in default_keys.items(): + if key not in data.keys(): + tu.assign_non_tensor(data, **{key: val}) + + chunk_size = int(tu.get_non_tensor_data(data=data, key="teacher_chunk_size", default=1024)) + + with self.engine.eval_mode(disable_auto_offload=False), torch.no_grad(): + # Standard input-prep (rmpad / position-id layout), but skip + # prepare_model_outputs which would force label-conditioned log_probs. + model_inputs, _ = self.engine.prepare_model_inputs(micro_batch=data) + output = self.engine.module(**model_inputs) + logits_rmpad = output.logits.squeeze(0) # (total_nnz, V) + + topk_ids_nested = data["topk_ids"] + assert topk_ids_nested.is_nested, "topk_ids must be a nested tensor" + topk_ids_rmpad = topk_ids_nested.values() # (total_nnz, K) + assert topk_ids_rmpad.shape[0] == logits_rmpad.shape[0], ( + f"topk_ids rmpad shape {topk_ids_rmpad.shape} does not match logits " + f"rmpad shape {logits_rmpad.shape}: trainer must align topk_ids with input_ids." + ) + + teacher_logp_rmpad = chunked_gather_logprobs( + logits=logits_rmpad, + topk_ids=topk_ids_rmpad, + chunk_size=chunk_size, + ) + + # Teacher's own top-K IDs at the same width K, for the overlap diagnostics + # only. Indices only — no gradient, no extra V-sized tensor. + k = topk_ids_rmpad.shape[-1] + teacher_topk_ids_rmpad = torch.topk(logits_rmpad, k=k, dim=-1).indices.to(torch.int64) + + cu_seqlens = data["input_ids"].offsets() + # Move offsets to the same device as values for nested_tensor_from_jagged. + cu_seqlens = cu_seqlens.to(teacher_logp_rmpad.device) + teacher_on_student_logp = torch.nested.nested_tensor_from_jagged(teacher_logp_rmpad, cu_seqlens) + teacher_topk_ids = torch.nested.nested_tensor_from_jagged( + teacher_topk_ids_rmpad, cu_seqlens.to(teacher_topk_ids_rmpad.device) + ) + + if not self.engine.is_mp_src_rank_with_outputs(): + return None + + out = TensorDict( + { + "teacher_on_student_logp": teacher_on_student_logp, + "teacher_topk_ids": teacher_topk_ids, + }, + batch_size=data.batch_size, + ) + return out.cpu() + + # Lets the trainer's dispatch helpers check whether a worker handle hosts a teacher. + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def is_teacher_worker(self) -> bool: + return True diff --git a/verl/workers/teacher/utils.py b/verl/workers/teacher/utils.py new file mode 100644 index 00000000000..bad28db8a4b --- /dev/null +++ b/verl/workers/teacher/utils.py @@ -0,0 +1,66 @@ +# 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. +"""Memory-efficient log-probability gather utilities for the FSDP teacher worker.""" + +from __future__ import annotations + +import torch + + +def chunked_gather_logprobs( + logits: torch.Tensor, + topk_ids: torch.Tensor, + chunk_size: int = 1024, +) -> torch.Tensor: + """Gather ``log_softmax(logits)[..., topk_ids]`` without materializing the full + ``log_softmax`` tensor. + + Logits are processed in fp32 in contiguous chunks of ``chunk_size`` leading-dim + rows, so peak memory is ``chunk_size * V * 4`` bytes per side instead of ``N * V * 4``. + + Args: + logits: ``(..., V)``. May be on any device / in any floating dtype. + topk_ids: ``(..., K)``. ``int64``. Leading dims must match ``logits``. + chunk_size: Number of leading-dim rows to process per inner iteration. + Must be ``> 0``. Smaller values reduce peak memory. + + Returns: + Tensor of shape ``(..., K)`` with the same dtype and device as ``logits``. + """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}.") + if logits.shape[:-1] != topk_ids.shape[:-1]: + raise ValueError( + f"Leading shapes of logits {tuple(logits.shape)} and topk_ids " + f"{tuple(topk_ids.shape)} must match (only the trailing dim may differ)." + ) + + leading_shape = logits.shape[:-1] + vocab_size = logits.shape[-1] + k = topk_ids.shape[-1] + + flat_logits = logits.reshape(-1, vocab_size) + # topk_ids may arrive on CPU; gather requires it on logits.device. + flat_topk = topk_ids.reshape(-1, k).to(flat_logits.device) + n = flat_logits.size(0) + out = torch.empty((n, k), dtype=logits.dtype, device=logits.device) + + for start in range(0, n, chunk_size): + end = min(start + chunk_size, n) + chunk_fp32 = flat_logits[start:end].float() + log_z = torch.logsumexp(chunk_fp32, dim=-1, keepdim=True) + gathered = torch.gather(chunk_fp32, dim=-1, index=flat_topk[start:end]) + out[start:end] = (gathered - log_z).to(logits.dtype) + + return out.reshape(*leading_shape, k) diff --git a/verl/workers/utils/padding.py b/verl/workers/utils/padding.py index ed8d1455205..462f4ed6115 100644 --- a/verl/workers/utils/padding.py +++ b/verl/workers/utils/padding.py @@ -93,6 +93,27 @@ def left_right_2_no_padding(data: TensorDict) -> TensorDict: data["teacher_logprobs"] = teacher_logprobs_nested data["teacher_ids"] = teacher_ids_nested + # (bsz, seqlen, K) reverse-KL fields; mirror the teacher_logprobs/teacher_ids handling above + student_topk_ids = data.get("student_topk_ids", None) + teacher_on_student_logp = data.get("teacher_on_student_logp", None) + if student_topk_ids is not None and teacher_on_student_logp is not None: + student_topk_ids_rmpad = index_first_axis(student_topk_ids.unsqueeze(-1).flatten(0, 1), indices) + teacher_on_student_logp_rmpad = index_first_axis(teacher_on_student_logp.unsqueeze(-1).flatten(0, 1), indices) + data["student_topk_ids"] = torch.nested.nested_tensor_from_jagged( + student_topk_ids_rmpad.squeeze(-1), offsets=cu_seqlens + ) + data["teacher_on_student_logp"] = torch.nested.nested_tensor_from_jagged( + teacher_on_student_logp_rmpad.squeeze(-1), offsets=cu_seqlens + ) + + # Teacher's own top-K IDs, carried alongside for the overlap diagnostics only. + teacher_topk_ids = data.get("teacher_topk_ids", None) + if teacher_topk_ids is not None: + teacher_topk_ids_rmpad = index_first_axis(teacher_topk_ids.unsqueeze(-1).flatten(0, 1), indices) + data["teacher_topk_ids"] = torch.nested.nested_tensor_from_jagged( + teacher_topk_ids_rmpad.squeeze(-1), offsets=cu_seqlens + ) + return data