From 1c0dc0c80ce61325e27deae7099e58b9573006f1 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Wed, 24 Jun 2026 17:50:40 +0800 Subject: [PATCH 01/14] [distillation] feat: add fsdp teacher backend hooks to DistillationConfig (verl#6676) Introduce three opt-in fields on DistillationConfig that scaffold the upcoming student-top-K reverse-KL pipeline (verl issue #6676): - teacher_backend ("vllm" | "fsdp"): selects whether the teacher is served by the existing vLLM/SGLang rollout (default, unchanged) or a frozen FSDP worker that gathers logprobs at student-top-K IDs each step. - teacher_chunk_size: token-axis chunk size for the FSDP gather kernel. - teacher_fsdp_config: reuses FSDPEngineConfig so the existing engine builder can construct the teacher model. When teacher_backend='fsdp' the post-init now skips the vLLM-only prompt_logprobs / max_logprobs / resource-pool-sizing checks (those assume a rollout deployment) and pins loss_mode to 'reverse_kl_topk' so misuse is caught early. The default vllm path is unchanged. Co-Authored-By: Claude Opus 4.8 --- verl/workers/config/distillation.py | 65 +++++++++++++++++++++++------ 1 file changed, 52 insertions(+), 13 deletions(-) diff --git a/verl/workers/config/distillation.py b/verl/workers/config/distillation.py index cc075d41f75..95592d5c635 100644 --- a/verl/workers/config/distillation.py +++ b/verl/workers/config/distillation.py @@ -20,6 +20,7 @@ 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"] @@ -243,6 +244,17 @@ class DistillationConfig(BaseConfig): +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"} @@ -254,10 +266,31 @@ class DistillationConfig(BaseConfig): 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(): @@ -274,24 +307,30 @@ def __post_init__(self): f"({self.n_gpus_per_node=} * {self.nnodes=} = {total_pool_size})." ) - def _resolve_teacher_models(self) -> dict[str, DistillationTeacherModelConfig]: + 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"] - 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})." + 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 ) - teacher_model.num_replicas = pool_size // per_replica + 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 From 9387fda064a7b2f65352866ffc155c4abc823cd8 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Wed, 24 Jun 2026 18:52:23 +0800 Subject: [PATCH 02/14] [distillation] feat: add student-top-K reverse-KL loss (verl#6676) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement the mode-seeking objective requested in verl issue #6676: L_t = sum_{i in TopK_student} q_i * (log q_i - log p_i) with the support taken on the *student*'s own top-K and the expectation under the student distribution. The teacher log-probabilities at those student-top-K IDs are precomputed by the upcoming FSDP teacher worker (P3), so this loss only consumes the gathered (B, T, K) tensor and does not need to materialize a full [B, T, V] log-softmax. Changes: - verl/trainer/distillation/fsdp/losses.py: add reverse_kl_divergence helper and compute_reverse_kl_topk, mirroring compute_forward_kl_topk's nested-tensor / ulysses-SP unwrap path. Returns distillation_losses, student_mass (≈ 1, sanity), and teacher_mass (the key mode-seeking diagnostic — should rise during training). - verl/trainer/distillation/losses.py: dispatch compute_topk_loss by loss_mode; route 'reverse_kl_topk' to the fsdp/veomni implementation (megatron raises NotImplementedError pending a vocab-parallel port); register the top-level wrapper that re-pads outputs, aggregates the mass diagnostics over the response mask, and clamps minor negative drift introduced by log_prob_min_clamp. - tests/workers/test_reverse_kl_topk_on_cpu.py: CPU-only unit tests for non-negativity, zero at distribution match, gradient flow into the student logits, agreement with the manual reference formula, and registry/metric-aggregation contracts. Co-Authored-By: Claude Opus 4.8 --- tests/workers/test_reverse_kl_topk_on_cpu.py | 182 +++++++++++++++++++ verl/trainer/distillation/fsdp/losses.py | 87 +++++++++ verl/trainer/distillation/losses.py | 111 +++++++++-- 3 files changed, 360 insertions(+), 20 deletions(-) create mode 100644 tests/workers/test_reverse_kl_topk_on_cpu.py 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..0bd81151189 --- /dev/null +++ b/tests/workers/test_reverse_kl_topk_on_cpu.py @@ -0,0 +1,182 @@ +# 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 + + sum_{i in TopK_student} q_i * (log q_i - log p_i) + + where ``q`` is the student top-K probability gathered from + ``softmax(student_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) + expected = (student_topk_logp.exp() * (student_topk_logp - teacher_on_student_logp)).sum(dim=-1) + torch.testing.assert_close(out["distillation_losses"], expected, atol=1e-6, rtol=1e-6) + + +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/trainer/distillation/fsdp/losses.py b/verl/trainer/distillation/fsdp/losses.py index c900286a875..d5f61f806ce 100644 --- a/verl/trainer/distillation/fsdp/losses.py +++ b/verl/trainer/distillation/fsdp/losses.py @@ -32,6 +32,27 @@ def kl_divergence(log_q: torch.Tensor, log_p: torch.Tensor) -> torch.Tensor: 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, @@ -92,3 +113,69 @@ def compute_forward_kl_topk( "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, +) -> dict[str, torch.Tensor]: + """Compute student-top-K reverse KL distillation loss. + + The support is the student's own top-K and the expectation is taken under the + student distribution Q; ``L_t = sum_{i in TopK_student} q_i * (log q_i - log p_i)``. + The teacher log-probabilities at those IDs are precomputed by the FSDP teacher + worker, so no full ``[B, T, V]`` tensor is materialized here. + + 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. + + Returns: + dict with ``distillation_losses`` (per-token KL(Q || P)) and the + ``student_mass`` / ``teacher_mass`` diagnostics (student- and teacher-side + probability on the student-top-K; ``teacher_mass`` should rise in training). + """ + 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) + + # 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) + 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) + + student_mass = student_topk_log_probs.exp().sum(dim=-1) + teacher_mass = teacher_on_student_logp.exp().sum(dim=-1) + + 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, + ) + + return { + "distillation_losses": distillation_losses, + "student_mass": student_mass, + "teacher_mass": teacher_mass, + } diff --git a/verl/trainer/distillation/losses.py b/verl/trainer/distillation/losses.py index 07df0850e35..edd9104dc92 100644 --- a/verl/trainer/distillation/losses.py +++ b/verl/trainer/distillation/losses.py @@ -134,26 +134,49 @@ 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, - ) + 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, + ) + 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(): @@ -392,3 +415,51 @@ 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) + 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 + + 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()), + } + + # clamp small negative drift from the log_prob_min_clamp guard + distillation_losses = distillation_losses.clamp_min(0.0) + + return distillation_losses, distillation_metrics From 7a4ae8fd6648c28fd8e00d73baac07311219daf0 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Wed, 24 Jun 2026 19:17:56 +0800 Subject: [PATCH 03/14] [distillation] feat: add forward-only FSDP teacher worker (verl#6676) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce the worker-side scaffold for the FSDP teacher backend: - verl/workers/teacher/utils.py: chunked_gather_logprobs — kernel that computes log_softmax(logits).gather(idx) without materializing the full log_softmax tensor, by using log_softmax(x).gather(idx) == x.gather(idx) - logsumexp(x) in fp32 over chunks of the leading axis. Peak memory becomes chunk_size * V * 4 bytes per side instead of N * V * 4. - verl/workers/teacher/fsdp_teacher_worker.py: TeacherFSDPWorker — subclasses TrainingWorker and pins engine_config.forward_only=True so FSDPEngine builds without optimizer/grad. Adds compute_logprobs_at_ids which reuses FSDPEngineWithLMHead.prepare_model_inputs (matching the actor's rmpad / nested-tensor layout), runs a single forward, and delegates to chunked_gather_logprobs at trainer-supplied topk_ids. The trainer is responsible for the roll-by-1 alignment; Ulysses SP on the teacher is gated off until a follow-up. - verl/experimental/teacher_loop/teacher_fsdp_manager.py: FSDPTeacherManager and MultiTeacherFSDPManager — mirror the public surface of TeacherModelManager / MultiTeacherModelManager so the trainer's resource-pool plumbing can switch on teacher_backend without diverging downstream consumers. Worker-group spawning and main-loop dispatch are explicitly NotImplementedError-marked for P4. - tests/workers/teacher/test_chunked_gather_on_cpu.py: CPU-only unit tests covering chunk_size sweep including > N rows, 2D rmpad input, shape/dtype preservation, validation errors, and fp32-stable behavior on large-magnitude logits. Co-Authored-By: Claude Opus 4.8 --- .../teacher/test_chunked_gather_on_cpu.py | 96 +++++++++++ .../teacher_loop/teacher_fsdp_manager.py | 158 ++++++++++++++++++ verl/workers/teacher/__init__.py | 25 +++ verl/workers/teacher/fsdp_teacher_worker.py | 156 +++++++++++++++++ verl/workers/teacher/utils.py | 65 +++++++ 5 files changed, 500 insertions(+) create mode 100644 tests/workers/teacher/test_chunked_gather_on_cpu.py create mode 100644 verl/experimental/teacher_loop/teacher_fsdp_manager.py create mode 100644 verl/workers/teacher/__init__.py create mode 100644 verl/workers/teacher/fsdp_teacher_worker.py create mode 100644 verl/workers/teacher/utils.py 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/verl/experimental/teacher_loop/teacher_fsdp_manager.py b/verl/experimental/teacher_loop/teacher_fsdp_manager.py new file mode 100644 index 00000000000..9950599ee7a --- /dev/null +++ b/verl/experimental/teacher_loop/teacher_fsdp_manager.py @@ -0,0 +1,158 @@ +# 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 exposes the matching multi-teacher manager +(``MultiTeacherFSDPManager``) and a per-teacher wrapper (``FSDPTeacherManager``). + +Worker-group spawning is finalized alongside the trainer wiring in +``verl/trainer/ppo/ray_trainer.py``; the not-yet-wired entrypoints raise +``NotImplementedError``. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Optional + +from omegaconf import DictConfig +from tensordict import TensorDict + +from verl.single_controller.ray.base import RayResourcePool, split_resource_pool +from verl.utils.config import omega_conf_to_dataclass +from verl.workers.config import DistillationConfig, DistillationTeacherModelConfig + +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 public surface of the vLLM equivalent ``TeacherModelManager`` so the + trainer's resource-pool plumbing can switch on ``teacher_backend`` transparently. + + Args: + distillation_config: Resolved distillation config (must have + ``teacher_backend == 'fsdp'`` and a non-None ``teacher_fsdp_config``). + teacher_model_config: Per-teacher entry from + ``distillation_config.teacher_models``. + resource_pool: Dedicated Ray resource pool for this teacher. + """ + + def __init__( + self, + distillation_config: DistillationConfig, + teacher_model_config: DistillationTeacherModelConfig, + resource_pool: RayResourcePool, + ): + self.distillation_config = distillation_config + self.teacher_model_config = teacher_model_config + self.resource_pool = resource_pool + self.worker_group: Any = None # populated by _spawn_worker_group + + def _spawn_worker_group(self) -> Any: + """Build a Ray worker group of :class:`TeacherFSDPWorker` over this + manager's resource pool. + + Deferred to the trainer integration, which reuses the existing + ``RayWorkerGroup`` plumbing in ``ray_trainer.py``. + """ + raise NotImplementedError( + "FSDPTeacherManager._spawn_worker_group is wired alongside the " + "trainer's resource-pool dispatch refactor." + ) + + def compute_logprobs_at_ids(self, data: TensorDict) -> Optional[TensorDict]: + """Dispatch a forward + chunked gather to the teacher worker group. + + Returns the trainer-facing tensor dict with ``teacher_on_student_logp`` + populated as a nested tensor of shape ``(bsz, seq_len_i, K)``. + """ + if self.worker_group is None: + raise RuntimeError( + "FSDPTeacherManager.compute_logprobs_at_ids called before _spawn_worker_group. " + "Wire this manager into the trainer's main loop before invoking." + ) + # Will become ``return self.worker_group.compute_logprobs_at_ids(data)`` + # once the worker group dispatch decorators are aligned with the trainer. + raise NotImplementedError( + "FSDPTeacherManager.compute_logprobs_at_ids is wired alongside the " + "trainer's main-loop refactor." + ) + + +class MultiTeacherFSDPManager: + """Manages one :class:`FSDPTeacherManager` per teacher routing key. + + Mirrors the vLLM ``MultiTeacherModelManager`` so the trainer can branch on + ``distillation.teacher_backend`` while keeping the rest of the glue identical. + """ + + def __init__(self, config: DictConfig, resource_pool: RayResourcePool): + 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'." + ) + + self.resource_pool = resource_pool + self.teacher_managers: dict[str, FSDPTeacherManager] = {} + self._initialize_teacher_managers() + + def _initialize_teacher_managers(self) -> None: + teacher_models = self.distillation_config.teacher_models + # FSDP teachers ignore vllm/sglang inference parallelism; num_replicas=1 is + # pinned in config post-init. Multi-teacher splits the pool per teacher. + split_sizes = [self.resource_pool.world_size] if len(teacher_models) == 1 else None + if split_sizes is None: + raise NotImplementedError( + "Multi-teacher FSDP pool sizing is wired once the trainer-side " + "teacher resource pool layout is finalized." + ) + split_pools = split_resource_pool(self.resource_pool, split_size=split_sizes) + for (key, teacher_model_config), teacher_pool in zip( + teacher_models.items(), split_pools, strict=True + ): + self.teacher_managers[key] = FSDPTeacherManager( + distillation_config=self.distillation_config, + teacher_model_config=teacher_model_config, + resource_pool=teacher_pool, + ) + + 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] + + +__all__ = ["FSDPTeacherManager", "MultiTeacherFSDPManager"] 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..82081d589bf --- /dev/null +++ b/verl/workers/teacher/fsdp_teacher_worker.py @@ -0,0 +1,156 @@ +# 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 the field ``teacher_on_student_logp`` + (nested ``(bsz, seq_len_i, K)``), or ``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, + ) + + cu_seqlens = data["input_ids"].offsets() + teacher_on_student_logp = torch.nested.nested_tensor_from_jagged( + teacher_logp_rmpad, cu_seqlens + ) + + if not self.engine.is_mp_src_rank_with_outputs(): + return None + + out = TensorDict( + {"teacher_on_student_logp": teacher_on_student_logp}, + 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..2326777afd4 --- /dev/null +++ b/verl/workers/teacher/utils.py @@ -0,0 +1,65 @@ +# 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) + flat_topk = topk_ids.reshape(-1, k) + 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) From f69157a8237fa3dd013b0791f4c36c226138c687 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Wed, 24 Jun 2026 19:47:23 +0800 Subject: [PATCH 04/14] [distillation] feat: wire FSDP teacher into trainer main loop (verl#6676) Connect the post-rollout student-top-K -> teacher-gather pipeline into the trainer's RayPPOTrainer fit loop. With teacher_backend='fsdp': - agent_loop.py: gate the per-trajectory vLLM teacher logprob hook on use_vllm_teacher (False under the FSDP backend) so rollouts no longer attempt to reach a non-existent LLMServerClient. - workers/engine/fsdp/transformer_impl.py: relax the per-key shape assertion in prepare_model_outputs from `v.shape == log_probs.shape` to `v.shape[0] == log_probs.shape[0]`, allowing the logits processor to emit per-token, top-K outputs (e.g. (total_nnz, K)) which the reverse-KL pipeline needs for student_topk_log_probs / student_topk_ids. - trainer/distillation/losses.py: add a `compute_student_topk_only` data-flag fast path on compute_topk_loss that returns the student's own top-K log-probs and IDs, used by the actor's first-pass forward before the teacher is queried. - workers/engine_workers.py: add ActorRolloutRefWorker.compute_student_topk which stamps the data flag above and runs a no-grad infer_batch on the actor engine, returning the per-token student top-K as a TensorDict the trainer can hand to the teacher worker. - experimental/teacher_loop/teacher_fsdp_manager.py: rewrite as a thin manager (single-teacher only for now) that holds a Ray worker group handle installed by the trainer after the unified spawn. Multi-teacher raises NotImplementedError until the per-teacher resource-pool split is finalized. - trainer/ppo/ray_trainer.py: branch teacher manager creation by teacher_backend; pre-spawn, register TeacherFSDPWorker on the teacher resource pool with a TrainingWorkerConfig assembled from distillation.teacher_fsdp_config; post-spawn, install the spawned handle on MultiTeacherFSDPManager and reset the worker. Add helpers _build_fsdp_teacher_training_config and _inject_fsdp_teacher_data, and call the latter inside the fit loop's update_actor block (under a dedicated marked_timer) so update_actor sees student_topk_ids / teacher_on_student_logp on every step. The vLLM path is unchanged. Co-Authored-By: Claude Opus 4.8 --- verl/experimental/agent_loop/agent_loop.py | 13 +- .../teacher_loop/teacher_fsdp_manager.py | 133 ++++++++-------- verl/trainer/distillation/losses.py | 19 +++ verl/trainer/ppo/ray_trainer.py | 150 +++++++++++++++++- verl/workers/engine/fsdp/transformer_impl.py | 8 +- verl/workers/engine_workers.py | 18 +++ 6 files changed, 262 insertions(+), 79 deletions(-) diff --git a/verl/experimental/agent_loop/agent_loop.py b/verl/experimental/agent_loop/agent_loop.py index b9e8e20ed8b..ea8fce44194 100644 --- a/verl/experimental/agent_loop/agent_loop.py +++ b/verl/experimental/agent_loop/agent_loop.py @@ -424,7 +424,14 @@ 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 @@ -910,7 +917,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 index 9950599ee7a..7c0f162158c 100644 --- a/verl/experimental/teacher_loop/teacher_fsdp_manager.py +++ b/verl/experimental/teacher_loop/teacher_fsdp_manager.py @@ -15,12 +15,12 @@ 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 exposes the matching multi-teacher manager -(``MultiTeacherFSDPManager``) and a per-teacher wrapper (``FSDPTeacherManager``). +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-group spawning is finalized alongside the trainer wiring in -``verl/trainer/ppo/ray_trainer.py``; the not-yet-wired entrypoints raise -``NotImplementedError``. +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 @@ -32,9 +32,8 @@ from omegaconf import DictConfig from tensordict import TensorDict -from verl.single_controller.ray.base import RayResourcePool, split_resource_pool from verl.utils.config import omega_conf_to_dataclass -from verl.workers.config import DistillationConfig, DistillationTeacherModelConfig +from verl.workers.config import DistillationConfig logger = logging.getLogger(__file__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) @@ -43,67 +42,44 @@ class FSDPTeacherManager: """Per-teacher wrapper around a Ray worker group of :class:`TeacherFSDPWorker`. - Mirrors the public surface of the vLLM equivalent ``TeacherModelManager`` so the - trainer's resource-pool plumbing can switch on ``teacher_backend`` transparently. - - Args: - distillation_config: Resolved distillation config (must have - ``teacher_backend == 'fsdp'`` and a non-None ``teacher_fsdp_config``). - teacher_model_config: Per-teacher entry from - ``distillation_config.teacher_models``. - resource_pool: Dedicated Ray resource pool for this teacher. + 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, - distillation_config: DistillationConfig, - teacher_model_config: DistillationTeacherModelConfig, - resource_pool: RayResourcePool, - ): + def __init__(self, key: str, distillation_config: DistillationConfig): + self.key = key self.distillation_config = distillation_config - self.teacher_model_config = teacher_model_config - self.resource_pool = resource_pool - self.worker_group: Any = None # populated by _spawn_worker_group - - def _spawn_worker_group(self) -> Any: - """Build a Ray worker group of :class:`TeacherFSDPWorker` over this - manager's resource pool. + self.worker_group: Any = None - Deferred to the trainer integration, which reuses the existing - ``RayWorkerGroup`` plumbing in ``ray_trainer.py``. - """ - raise NotImplementedError( - "FSDPTeacherManager._spawn_worker_group is wired alongside the " - "trainer's resource-pool dispatch refactor." - ) + 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]: - """Dispatch a forward + chunked gather to the teacher worker group. + """Forward a teacher-side ``compute_logprobs_at_ids`` call to the worker group. - Returns the trainer-facing tensor dict with ``teacher_on_student_logp`` - populated as a nested tensor of shape ``(bsz, seq_len_i, K)``. + 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( - "FSDPTeacherManager.compute_logprobs_at_ids called before _spawn_worker_group. " - "Wire this manager into the trainer's main loop before invoking." + 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." ) - # Will become ``return self.worker_group.compute_logprobs_at_ids(data)`` - # once the worker group dispatch decorators are aligned with the trainer. - raise NotImplementedError( - "FSDPTeacherManager.compute_logprobs_at_ids is wired alongside the " - "trainer's main-loop refactor." - ) + return self.worker_group.compute_logprobs_at_ids(data) class MultiTeacherFSDPManager: """Manages one :class:`FSDPTeacherManager` per teacher routing key. - Mirrors the vLLM ``MultiTeacherModelManager`` so the trainer can branch on - ``distillation.teacher_backend`` while keeping the rest of the glue identical. + 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: RayResourcePool): + 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": @@ -116,29 +92,47 @@ def __init__(self, config: DictConfig, resource_pool: RayResourcePool): "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 - self.teacher_managers: dict[str, FSDPTeacherManager] = {} - self._initialize_teacher_managers() - def _initialize_teacher_managers(self) -> None: teacher_models = self.distillation_config.teacher_models - # FSDP teachers ignore vllm/sglang inference parallelism; num_replicas=1 is - # pinned in config post-init. Multi-teacher splits the pool per teacher. - split_sizes = [self.resource_pool.world_size] if len(teacher_models) == 1 else None - if split_sizes is None: + if len(teacher_models) != 1: raise NotImplementedError( - "Multi-teacher FSDP pool sizing is wired once the trainer-side " - "teacher resource pool layout is finalized." + "Multi-teacher FSDP backend is not yet supported; configure exactly one teacher " + "(teacher_models has 1 entry) when teacher_backend='fsdp'." ) - split_pools = split_resource_pool(self.resource_pool, split_size=split_sizes) - for (key, teacher_model_config), teacher_pool in zip( - teacher_models.items(), split_pools, strict=True - ): - self.teacher_managers[key] = FSDPTeacherManager( - distillation_config=self.distillation_config, - teacher_model_config=teacher_model_config, - resource_pool=teacher_pool, + + 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: @@ -154,5 +148,10 @@ def get(self, routing_key: Optional[str] = None) -> FSDPTeacherManager: ) 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/distillation/losses.py b/verl/trainer/distillation/losses.py index edd9104dc92..39863c86ff8 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,6 +136,23 @@ def compute_topk_loss( - student_mass: (bsz, seqlen/cp_size) - teacher_mass: (bsz, seqlen/cp_size) """ + # 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": diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index 439731f86a9..256fd56aa88 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -835,6 +835,23 @@ def init_workers(self): ) 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) + + from verl.workers.teacher import TeacherFSDPWorker + + teacher_cls = RayClassWithInitArgs(cls=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`. @@ -915,14 +932,39 @@ def init_workers(self): # 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) + 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 @@ -946,12 +988,16 @@ def init_workers(self): # 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 + # 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=self.teacher_model_manager.get_client() if self.use_teacher_policy else None, + teacher_client=teacher_client_for_agent_loop, reward_loop_worker_handles=reward_loop_worker_handles, ) @@ -1290,6 +1336,91 @@ def _compute_old_log_prob(self, batch: DataProto): 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, + ) + teacher_output = self.teacher_model_manager.compute_logprobs_at_ids(teacher_input) + teacher_on_student_logp = tu.get(teacher_output, "teacher_on_student_logp") # nested (B, T_i, K) + + # Phase 3: nested -> dense (B, max_response_len, K), union into batch so the engine + # left_right_2_no_padding'es them back to nested when update_actor micro-batches run. + student_topk_ids_padded = no_padding_2_padding(student_topk_ids, batch_td) + teacher_on_student_logp_padded = no_padding_2_padding(teacher_on_student_logp, batch_td) + extras = tu.get_tensordict( + { + "student_topk_ids": student_topk_ids_padded, + "teacher_on_student_logp": teacher_on_student_logp_padded.float(), + } + ) + 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 @@ -1645,6 +1776,9 @@ def fit(self): 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) diff --git a/verl/workers/engine/fsdp/transformer_impl.py b/verl/workers/engine/fsdp/transformer_impl.py index 2c66c6a3fad..68218b71ad2 100644 --- a/verl/workers/engine/fsdp/transformer_impl.py +++ b/verl/workers/engine/fsdp/transformer_impl.py @@ -1135,12 +1135,16 @@ def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict, lo sum_pi_squared_rmpad = verl_F.calculate_sum_pi_squared_from_logits(logits_rmpad) # 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 == log_probs.shape, f"log_probs shape: {log_probs.shape}, {k} shape: {v.shape}" + assert v.shape[0] == log_probs.shape[0], ( + f"log_probs shape: {log_probs.shape}, {k} shape: {v.shape}" + ) if self.use_ulysses_sp: pad_size = output_args["pad_size"] v = gather_outputs_and_unpad(v, gather_dim=0, unpad_dim=0, padding_size=pad_size) @@ -1222,7 +1226,7 @@ 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 == log_probs.shape, ( + assert v.shape[0] == log_probs.shape[0], ( f"log_probs shape: {log_probs.shape}, {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 bf379b49c3d..508c62f4a4c 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -646,6 +646,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) From be035df042d694d9042f8c625a7035ea60a4c74d Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Wed, 24 Jun 2026 19:56:43 +0800 Subject: [PATCH 05/14] [examples] feat: add reverse-KL FSDP-teacher distillation script (verl#6676) Single-teacher OPD recipe that runs the verl#6676 student-top-K reverse-KL pipeline end-to-end: - distillation.teacher_backend=fsdp routes the teacher through the TeacherFSDPWorker instead of a vLLM rollout replica. - distillation.distillation_loss.loss_mode=reverse_kl_topk drives the student-top-K reverse-KL loss, with topk repurposed as the *student*-top-K width. - teacher_fsdp_config exposes the FSDPEngineConfig defaults the trainer now wires through (param_offload + bf16 + fsdp_size=-1). Defaults are sized for a single 8-GPU node with a 1.7B student and an 8B FSDP teacher; the script's user-adjustable block exposes the knobs needed to retarget at larger pairs. Co-Authored-By: Claude Opus 4.8 --- .../run_qwen3_reverse_kl_fsdp.sh | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100755 examples/on_policy_distillation_trainer/run_qwen3_reverse_kl_fsdp.sh 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[@]}" \ + "$@" From 0c3a90fc76d5ba4ab80cdaabd20926f2030f78eb Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Fri, 26 Jun 2026 17:06:57 +0800 Subject: [PATCH 06/14] [distillation] fix: FSDP teacher pipeline glue and config knobs (verl#6676) - losses.py: short-circuit policy loss on the FSDP-teacher first pass (compute_student_topk_only) so the student top-K emit step does not trigger the final loss path. - ray_trainer.py: wrap TeacherFSDPWorker with ray.remote, materialize the non-blocking compute_logprobs_at_ids future, and pass temperature=1.0 so the FSDP engine emits the teacher's raw distribution. - engine.py: allow runtime override of EngineRouterReplayConfig.mode and type max/micro batch fields as Optional[int]. - fsdp_teacher_worker.py: move cu_seqlens to the values device before nested_tensor_from_jagged. - teacher/utils.py: move topk_ids to logits.device for torch.gather. - distillation.yaml: add teacher_backend / teacher_chunk_size / teacher_fsdp_config knobs for the FSDP teacher path; regenerate the trainer yaml variants. Co-Authored-By: Claude Opus 4.8 --- .../_generated_ppo_megatron_trainer.yaml | 10 ++++++ .../_generated_ppo_torchtitan_trainer.yaml | 10 ++++++ .../config/_generated_ppo_trainer.yaml | 10 ++++++ .../config/_generated_ppo_veomni_trainer.yaml | 10 ++++++ .../config/distillation/distillation.yaml | 32 ++++++++++++++++++- verl/trainer/distillation/losses.py | 4 +++ verl/trainer/ppo/ray_trainer.py | 8 ++++- verl/workers/config/engine.py | 11 ++++--- verl/workers/teacher/fsdp_teacher_worker.py | 2 ++ verl/workers/teacher/utils.py | 3 +- 10 files changed, 93 insertions(+), 7 deletions(-) diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index 45359349dfa..c6e7d4a18bc 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -866,6 +866,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 52bfb78a0a8..9124e7bdc49 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -800,6 +800,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 a300875cc7b..63e79dbeddb 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -820,6 +820,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 b48c2fb5652..b6e7e3e78d2 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -805,6 +805,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/losses.py b/verl/trainer/distillation/losses.py index 39863c86ff8..614739c9cd0 100644 --- a/verl/trainer/distillation/losses.py +++ b/verl/trainer/distillation/losses.py @@ -246,6 +246,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) diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index 256fd56aa88..260322fdff7 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -846,9 +846,11 @@ def init_workers(self): 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=TeacherFSDPWorker, config=teacher_training_config) + 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 @@ -1405,8 +1407,12 @@ def _inject_fsdp_teacher_data(self, batch: DataProto) -> DataProto: tu.assign_non_tensor( teacher_input, teacher_chunk_size=self.distillation_config.teacher_chunk_size, + # FSDP engine requires temperature; use T=1.0 for the raw teacher distribution. + temperature=1.0, ) 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) # Phase 3: nested -> dense (B, max_response_len, K), union into batch so the engine diff --git a/verl/workers/config/engine.py b/verl/workers/config/engine.py index b5ad8e9cfea..6f0c696a9ad 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/teacher/fsdp_teacher_worker.py b/verl/workers/teacher/fsdp_teacher_worker.py index 82081d589bf..9d0c59e6d99 100644 --- a/verl/workers/teacher/fsdp_teacher_worker.py +++ b/verl/workers/teacher/fsdp_teacher_worker.py @@ -137,6 +137,8 @@ def compute_logprobs_at_ids(self, data: TensorDict) -> Optional[TensorDict]: ) 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 ) diff --git a/verl/workers/teacher/utils.py b/verl/workers/teacher/utils.py index 2326777afd4..bad28db8a4b 100644 --- a/verl/workers/teacher/utils.py +++ b/verl/workers/teacher/utils.py @@ -51,7 +51,8 @@ def chunked_gather_logprobs( k = topk_ids.shape[-1] flat_logits = logits.reshape(-1, vocab_size) - flat_topk = topk_ids.reshape(-1, k) + # 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) From 3859df00a18121d092c938c0bef527b6b02d4bcb Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Fri, 26 Jun 2026 17:48:57 +0800 Subject: [PATCH 07/14] [distillation] fix: align FSDP teacher topk fields through left_right padding (verl#6676) - ray_trainer.py: Phase 3 now scatters student_topk_ids / teacher_on_student_logp back to left-right padded dense (B, max_seq_len, K) using the Phase-1 ``indices`` / ``max_seq_len``, instead of no_padding_2_padding (which sliced response-only and left-shifted, breaking input_ids alignment). - padding.py: left_right_2_no_padding now also re-nests student_topk_ids and teacher_on_student_logp aligned to input_ids, matching the teacher_logprobs / teacher_ids handling so the reverse_kl_topk logits processor sees aligned jagged tensors. --- verl/trainer/ppo/ray_trainer.py | 25 +++++++++++++++++++++---- verl/workers/utils/padding.py | 13 +++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index 260322fdff7..e760919bb0f 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -1415,10 +1415,27 @@ def _inject_fsdp_teacher_data(self, batch: DataProto) -> DataProto: teacher_output = teacher_output.get() teacher_on_student_logp = tu.get(teacher_output, "teacher_on_student_logp") # nested (B, T_i, K) - # Phase 3: nested -> dense (B, max_response_len, K), union into batch so the engine - # left_right_2_no_padding'es them back to nested when update_actor micro-batches run. - student_topk_ids_padded = no_padding_2_padding(student_topk_ids, batch_td) - teacher_on_student_logp_padded = no_padding_2_padding(teacher_on_student_logp, batch_td) + # 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) extras = tu.get_tensordict( { "student_topk_ids": student_topk_ids_padded, diff --git a/verl/workers/utils/padding.py b/verl/workers/utils/padding.py index ed8d1455205..f19f576748b 100644 --- a/verl/workers/utils/padding.py +++ b/verl/workers/utils/padding.py @@ -93,6 +93,19 @@ 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 + ) + return data From a70ad119906485f85785b7a4bf182486abac24af Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Mon, 29 Jun 2026 17:13:07 +0800 Subject: [PATCH 08/14] [distillation] feat: add student/teacher top-K overlap diagnostics (verl#6676) Track the student/teacher top-K intersection during reverse-KL OPD, following "Rethinking On-Policy Distillation of Large Language Models". The FSDP teacher now also emits its own top-K ids, carried through the left_right padding glue, so compute_reverse_kl_topk can report overlap_ratio and overlap_token_advantage. Also default use_policy_gradient to False. Co-Authored-By: Claude Opus 4.8 --- verl/trainer/distillation/fsdp/losses.py | 33 ++++++++++++++++++--- verl/trainer/distillation/losses.py | 25 ++++++++++++++++ verl/trainer/ppo/ray_trainer.py | 3 ++ verl/workers/config/distillation.py | 2 +- verl/workers/teacher/fsdp_teacher_worker.py | 21 ++++++++++--- verl/workers/utils/padding.py | 8 +++++ 6 files changed, 83 insertions(+), 9 deletions(-) diff --git a/verl/trainer/distillation/fsdp/losses.py b/verl/trainer/distillation/fsdp/losses.py index d5f61f806ce..192d17e72c9 100644 --- a/verl/trainer/distillation/fsdp/losses.py +++ b/verl/trainer/distillation/fsdp/losses.py @@ -121,13 +121,13 @@ def compute_reverse_kl_topk( 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 and the expectation is taken under the - student distribution Q; ``L_t = sum_{i in TopK_student} q_i * (log q_i - log p_i)``. - The teacher log-probabilities at those IDs are precomputed by the FSDP teacher - worker, so no full ``[B, T, V]`` tensor is materialized here. + student distribution Q. The teacher log-probabilities at those IDs are precomputed + by the FSDP teacher worker, so no full ``[B, T, V]`` tensor is materialized here. Args: student_logits: ``(bsz, seqlen/sp_size, vocab_size)``. @@ -137,11 +137,15 @@ def compute_reverse_kl_topk( 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`` (per-token KL(Q || P)) and the ``student_mass`` / ``teacher_mass`` diagnostics (student- and teacher-side probability on the student-top-K; ``teacher_mass`` should rise in training). + When ``teacher_topk_ids`` is given, also ``overlap_count`` / + ``overlap_token_advantage`` (student/teacher top-K intersection per token). """ del data_format # unused; kept for parity with compute_forward_kl_topk's signature # 1. unwrap nested layouts from the rmpad pipeline @@ -149,11 +153,15 @@ def compute_reverse_kl_topk( 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] @@ -174,8 +182,25 @@ def compute_reverse_kl_topk( log_p=teacher_on_student_logp, ) - return { + 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 614739c9cd0..9a3a137af8e 100644 --- a/verl/trainer/distillation/losses.py +++ b/verl/trainer/distillation/losses.py @@ -193,6 +193,7 @@ def compute_topk_loss( 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}") @@ -465,12 +466,35 @@ def compute_reverse_kl_topk( 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 = { @@ -480,6 +504,7 @@ def compute_reverse_kl_topk( "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 diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index e760919bb0f..efca2e064a7 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -1414,6 +1414,7 @@ def _inject_fsdp_teacher_data(self, batch: DataProto) -> DataProto: # 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. @@ -1436,10 +1437,12 @@ def _to_left_right_padded(nested: torch.Tensor) -> torch.Tensor: 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)) diff --git a/verl/workers/config/distillation.py b/verl/workers/config/distillation.py index 95592d5c635..31b3dbe4baf 100644 --- a/verl/workers/config/distillation.py +++ b/verl/workers/config/distillation.py @@ -70,7 +70,7 @@ class DistillationLossConfig(BaseConfig): loss_max_clamp: Optional[float] = 10.0 log_prob_min_clamp: Optional[float] = -10.0 - use_policy_gradient: bool = True + use_policy_gradient: bool = False policy_loss_mode: str = "vanilla" clip_ratio: float = 0.2 clip_ratio_low: float = 0.2 diff --git a/verl/workers/teacher/fsdp_teacher_worker.py b/verl/workers/teacher/fsdp_teacher_worker.py index 9d0c59e6d99..b55809998f1 100644 --- a/verl/workers/teacher/fsdp_teacher_worker.py +++ b/verl/workers/teacher/fsdp_teacher_worker.py @@ -91,9 +91,11 @@ def compute_logprobs_at_ids(self, data: TensorDict) -> Optional[TensorDict]: etc.), copied by the trainer from the actor's engine config. Returns: - A :class:`TensorDict` with the field ``teacher_on_student_logp`` - (nested ``(bsz, seq_len_i, K)``), or ``None`` when the worker is not the - output-collecting rank (matches the :meth:`TrainingWorker.infer_batch` contract). + 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( @@ -136,18 +138,29 @@ def compute_logprobs_at_ids(self, data: TensorDict) -> Optional[TensorDict]: 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_on_student_logp": teacher_on_student_logp, + "teacher_topk_ids": teacher_topk_ids, + }, batch_size=data.batch_size, ) return out.cpu() diff --git a/verl/workers/utils/padding.py b/verl/workers/utils/padding.py index f19f576748b..462f4ed6115 100644 --- a/verl/workers/utils/padding.py +++ b/verl/workers/utils/padding.py @@ -106,6 +106,14 @@ def left_right_2_no_padding(data: TensorDict) -> TensorDict: 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 From 6a9ecb62a0619cdbea4e56d9e584ff675c348377 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Tue, 30 Jun 2026 20:47:32 +0800 Subject: [PATCH 09/14] [distillation] fix: renormalize reverse-KL onto student top-K support (verl#6676) The student-top-K reverse-KL loss summed q_i*(log q_i - log p_i) over the student's own top-K without renormalizing. Once the student starts to de-peak, that truncated sum can be minimized by leaking probability mass *outside* the top-K (every q_i shrinks, so every term shrinks) instead of by matching the teacher. This drives a self-reinforcing collapse: entropy explodes, student_mass falls from ~0.97 to ~0.4 and overlap_ratio drops sharply, while the reported loss keeps decreasing. The onset is governed by the deterministic optimization drift, not batch noise, so it reproduces at the same step regardless of train batch size. Fix: renormalize both student and teacher log-probs onto the shared student-top-K support before the divergence, so the per-token objective is a proper KL(Q_K || P_K). student_mass / teacher_mass remain un-normalized as de-peaking diagnostics. Add a regression test asserting the loss is invariant to mass leaked outside the top-K, and update the reference-formula test. --- tests/workers/test_reverse_kl_topk_on_cpu.py | 46 +++++++++++++++++--- verl/trainer/distillation/fsdp/losses.py | 27 ++++++++---- 2 files changed, 58 insertions(+), 15 deletions(-) diff --git a/tests/workers/test_reverse_kl_topk_on_cpu.py b/tests/workers/test_reverse_kl_topk_on_cpu.py index 0bd81151189..71cf41ec37b 100644 --- a/tests/workers/test_reverse_kl_topk_on_cpu.py +++ b/tests/workers/test_reverse_kl_topk_on_cpu.py @@ -119,12 +119,9 @@ def test_reverse_kl_topk_gradient_flows_to_student(): def test_reverse_kl_topk_matches_reference_formula(): - """The token-wise loss equals the manual reference - - sum_{i in TopK_student} q_i * (log q_i - log p_i) - - where ``q`` is the student top-K probability gathered from - ``softmax(student_logits)``.""" + """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) @@ -139,10 +136,45 @@ def test_reverse_kl_topk_matches_reference_formula(): ) 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) - expected = (student_topk_logp.exp() * (student_topk_logp - teacher_on_student_logp)).sum(dim=-1) + # 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 diff --git a/verl/trainer/distillation/fsdp/losses.py b/verl/trainer/distillation/fsdp/losses.py index 192d17e72c9..499d2957647 100644 --- a/verl/trainer/distillation/fsdp/losses.py +++ b/verl/trainer/distillation/fsdp/losses.py @@ -125,9 +125,13 @@ def compute_reverse_kl_topk( ) -> dict[str, torch.Tensor]: """Compute student-top-K reverse KL distillation loss. - The support is the student's own top-K and the expectation is taken under the - student distribution Q. The teacher log-probabilities at those IDs are precomputed - by the FSDP teacher worker, so no full ``[B, T, V]`` tensor is materialized here. + 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)``. @@ -141,11 +145,9 @@ def compute_reverse_kl_topk( only for the overlap diagnostics. Overlap metrics are skipped when ``None``. Returns: - dict with ``distillation_losses`` (per-token KL(Q || P)) and the - ``student_mass`` / ``teacher_mass`` diagnostics (student- and teacher-side - probability on the student-top-K; ``teacher_mass`` should rise in training). - When ``teacher_topk_ids`` is given, also ``overlap_count`` / - ``overlap_token_advantage`` (student/teacher top-K intersection per token). + 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 @@ -169,9 +171,18 @@ def compute_reverse_kl_topk( 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) From 2cbe9f46e64508e3dbd267d91e37e1fd1f7492a5 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Tue, 30 Jun 2026 20:49:54 +0800 Subject: [PATCH 10/14] [distillation] fix: query FSDP teacher at the student rollout temperature (verl#6676) The reverse-KL OPD loss compares the student and teacher distributions, so both must be taken at the same temperature. The FSDP teacher was queried at a fixed T=1.0 while the student rollout uses rollout.temperature, so the divergence was computed between two distributions at mismatched temperatures. Pass rollout.temperature through to the teacher forward to align them. --- verl/trainer/ppo/ray_trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index efca2e064a7..e31a0a25fa0 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -1407,8 +1407,8 @@ def _inject_fsdp_teacher_data(self, batch: DataProto) -> DataProto: tu.assign_non_tensor( teacher_input, teacher_chunk_size=self.distillation_config.teacher_chunk_size, - # FSDP engine requires temperature; use T=1.0 for the raw teacher distribution. - temperature=1.0, + # 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. From e215dc0d76b50c9661312f757789513482ca798f Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Thu, 2 Jul 2026 22:13:56 +0800 Subject: [PATCH 11/14] fix: relax shape assert to check only first dimension in top-K path The top-K distillation path may produce tensors of shape (total_nnz,) or (total_nnz, K) for student_topk_log_probs / student_topk_ids. The old strict assert required v.shape == (total_nnz,), which would fail for 2D outputs. Change it to only compare the first dimension. --- verl/workers/engine/fsdp/transformer_impl.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/verl/workers/engine/fsdp/transformer_impl.py b/verl/workers/engine/fsdp/transformer_impl.py index 93c66ce8b83..7d1d70479b7 100644 --- a/verl/workers/engine/fsdp/transformer_impl.py +++ b/verl/workers/engine/fsdp/transformer_impl.py @@ -1165,7 +1165,9 @@ def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict, lo 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: @@ -1276,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) From 497f0dfeeab27967f6f874577286c31f56e7a506 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Thu, 2 Jul 2026 22:43:53 +0800 Subject: [PATCH 12/14] feat: add reverse-KL FSDP example for DSR1 1.5B reproducing Rethinking OPD --- .../run_dsr1_1.5b_reverse_kl_fsdp.sh | 156 ++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 examples/on_policy_distillation_trainer/run_dsr1_1.5b_reverse_kl_fsdp.sh 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[@]}" \ + "$@" From 8017adeca9a8f21813de8852ca6d3f8753769d28 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Thu, 2 Jul 2026 22:45:52 +0800 Subject: [PATCH 13/14] style: apply pre-commit formatting fixes --- tests/workers/test_reverse_kl_topk_on_cpu.py | 427 +- verl/experimental/agent_loop/agent_loop.py | 2543 ++++++----- .../teacher_loop/teacher_fsdp_manager.py | 305 +- verl/trainer/distillation/fsdp/losses.py | 540 ++- verl/trainer/ppo/ray_trainer.py | 3864 ++++++++--------- verl/workers/config/distillation.py | 726 ++-- verl/workers/teacher/fsdp_teacher_worker.py | 340 +- 7 files changed, 4360 insertions(+), 4385 deletions(-) diff --git a/tests/workers/test_reverse_kl_topk_on_cpu.py b/tests/workers/test_reverse_kl_topk_on_cpu.py index 71cf41ec37b..a3a4e3ef52f 100644 --- a/tests/workers/test_reverse_kl_topk_on_cpu.py +++ b/tests/workers/test_reverse_kl_topk_on_cpu.py @@ -1,214 +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) +# 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 b6f79c06355..08a09d0d4c7 100644 --- a/verl/experimental/agent_loop/agent_loop.py +++ b/verl/experimental/agent_loop/agent_loop.py @@ -1,1272 +1,1271 @@ -# 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. -""" -Agent framework for multi-turn rollout and agentic reinforcement learning. -- AgentLoopBase: coroutine based abstract base class for agent loop. - - SingleTurnAgentLoop: single turn agent loop. - - ToolAgentLoop: ReAct agent loop with tool calling, with user defined tools. -- AgentLoopWorker: worker class for running agent loop coroutines in parallel. -- AgentLoopManager: manager class for running agent loop workers in parallel. - -AgentLoopManager is one specific agent-framework implementation in verl, -and is designed to be fully replaceable by other agent frameworks such as: -- NVIDIA Nemo-Gym -- AWS Bedrock AgentCore -- SWE-agent -- ... -""" - -import asyncio -import logging -import os -import random -from abc import ABC, abstractmethod -from typing import Any, Optional -from uuid import uuid4 - -import hydra -import numpy as np -import ray -import torch -from omegaconf import DictConfig, OmegaConf -from PIL import Image -from pydantic import BaseModel, ConfigDict -from tensordict import TensorDict -from transformers import AutoProcessor, AutoTokenizer - -from verl.experimental.agent_loop.utils import resolve_config_path -from verl.protocol import DataProto -from verl.tools.tool_registry import load_all_tools -from verl.trainer.distillation import is_distillation_enabled -from verl.utils.config import omega_conf_to_dataclass -from verl.utils.dataset.rl_dataset import RLHFDataset, get_dataset_class -from verl.utils.model import compute_position_id_with_mask -from verl.utils.profiler import simple_timer -from verl.utils.ray_utils import auto_await, get_event_loop -from verl.utils.rollout_trace import ( - RolloutTraceConfig, - rollout_trace_attr, -) -from verl.utils.skip import SkipManager -from verl.utils.tokenizer import ( - build_multimodal_processor_inputs, - get_processor_token_id, - normalize_token_ids, -) -from verl.utils.tokenizer.chat_template import apply_chat_template, initialize_system_prompt -from verl.utils.tokenizer.continuous_token_wiring import create_continuous_token_builder -from verl.workers.config import ( - HFModelConfig, - RolloutConfig, -) -from verl.workers.rollout.llm_server import LLMServerClient - -logger = logging.getLogger(__file__) -logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) - -DEFAULT_ROUTING_CACHE_SIZE = 10000 - - -class AgentLoopMetrics(BaseModel): - """Agent loop performance metrics.""" - - generate_sequences: float = 0.0 - tool_calls: float = 0.0 - compute_score: float = 0.0 - num_preempted: int = -1 # -1 means not available - - -class AgentLoopOutput(BaseModel): - """Agent loop output.""" - - prompt_ids: list[int] - """Prompt token ids.""" - response_ids: list[int] - """Response token ids including LLM generated token, tool response token.""" - response_mask: list[int] - """Response mask, 1 for LLM generated token, 0 for tool response token.""" - response_logprobs: Optional[list[float]] = None - """Log probabilities for the response tokens.""" - routed_experts: Optional[Any] = None - """Routed experts for the total tokens.""" - multi_modal_data: Optional[dict[str, Any]] = None - """Multi-modal data for multi-modal tools.""" - reward_score: Optional[float] = None - """Reward score for the trajectory.""" - num_turns: int = 0 - """Number of chat turns, including user, assistant, tool.""" - metrics: AgentLoopMetrics - """Auxiliary performance metrics""" - extra_fields: dict[str, Any] = {} - """Extra fields for dynamic addition.""" - mm_processor_kwargs: Optional[dict[str, Any]] = None - """Processor/backend kwargs that must stay aligned across rollout and training paths.""" - - def as_dict(self) -> dict[str, Any]: - """Convert agent loop output to a dictionary.""" - output = self.model_dump(exclude_unset=True) - - output["prompts"] = torch.tensor(output.pop("prompt_ids"), dtype=torch.int64) - output["responses"] = torch.tensor(output.pop("response_ids"), dtype=torch.int64) - output["response_mask"] = torch.tensor(output.pop("response_mask"), dtype=torch.int64) - - response_logprobs = output.pop("response_logprobs", None) - if response_logprobs is not None: - output["rollout_log_probs"] = torch.tensor(response_logprobs, dtype=torch.float32) - - routed_experts = output.pop("routed_experts", None) - if routed_experts is not None: - output["routed_experts"] = torch.tensor(routed_experts, dtype=torch.int64) - - # rm_scores: reward score for each token - reward_score = output.pop("reward_score", None) - if reward_score is not None: - rm_scores = torch.zeros_like(output["response_mask"], dtype=torch.float32) - rm_scores[-1] = reward_score - output["rm_scores"] = rm_scores - - teacher_ids, teacher_logprobs = ( - output["extra_fields"].pop("teacher_ids", None), - output["extra_fields"].pop("teacher_logprobs", None), - ) - if teacher_ids is not None: - output["teacher_ids"] = teacher_ids - if teacher_logprobs is not None: - output["teacher_logprobs"] = teacher_logprobs - return output - - -class _InternalAgentLoopOutput(AgentLoopOutput): - """Internal agent loop output with padded sequences.""" - - model_config = ConfigDict(arbitrary_types_allowed=True) - - prompt_ids: torch.Tensor - """Padded prompt token ids.""" - response_ids: torch.Tensor - """Padded response token ids.""" - input_ids: torch.Tensor - """Padded input ids(prompt_ids + response_ids).""" - position_ids: torch.Tensor - """Padded position ids.""" - response_mask: torch.Tensor - """Padded response mask.""" - attention_mask: torch.Tensor - """Padded attention mask.""" - response_logprobs: Optional[torch.Tensor] = None - """Padded log probabilities for the response tokens.""" - teacher_logprobs: Optional[torch.Tensor] = None - """Padded log probabilities from teacher model for prompt/response tokens.""" - teacher_ids: Optional[torch.Tensor] = None - """Padded token ids corresponding to the teacher log probabilities.""" - routed_experts: Optional[torch.Tensor] = None - """Padded routed experts for the total tokens.""" - multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None - """Multi-modal inputs for processors (e.g. pixel_values, image_grid_thw, video_grid_thw).""" - extra_fields: dict[str, Any] = {} - """Extra fields for dynamic addition.""" - - -class DictConfigWrap: - """Wrapper for DictConfig to avoid hydra.utils.instantiate recursive resolve.""" - - def __init__(self, config: DictConfig): - self.config = config - - -class ToolListWrap: - """Wraps a tool list so ``hydra.utils.instantiate`` doesn't recursively - resolve its elements (which would demote them to ``DictConfig``).""" - - def __init__(self, tools: list): - self.tools = tools - - -class AgentLoopBase(ABC): - """An agent loop takes an input message, chat with OpenAI compatible LLM server and interact with various - environments. - - Args: - trainer_config (DictConfig): whole config for main entrypoint. - server_manager (LLMServerClient): OpenAI compatible LLM server manager. - tokenizer (AutoTokenizer): Tokenizer for tokenize messages. - processor (AutoProcessor): Processor for process messages. - dataset_cls (type[Dataset]): Dataset class for creating dataset, Defaults to RLHFDataset. - data_config (DictConfigWrap): Dataset config. - """ - - def __init__( - self, - trainer_config: DictConfigWrap, - server_manager: LLMServerClient, - tokenizer: AutoTokenizer, - processor: AutoProcessor, - dataset_cls: type[RLHFDataset], - data_config: DictConfigWrap, - **kwargs, - ): - self.config = trainer_config.config - self.rollout_config = self.config.actor_rollout_ref.rollout - self.server_manager = server_manager - self.tokenizer = tokenizer - self.processor = processor - self.dataset_cls = dataset_cls - self.data_config = data_config.config - self.apply_chat_template_kwargs = self.data_config.get("apply_chat_template_kwargs", {}) - self.mm_processor_kwargs = self.data_config.get("mm_processor_kwargs", {}) - self.continuous_token_builder = None - self.enable_continuous_token = False - continuous_token_config = self.data_config.continuous_token - if continuous_token_config.enable and self.processor is None: - model_config = self.config.actor_rollout_ref.model - self.continuous_token_builder = create_continuous_token_builder( - self.tokenizer, - model_family=continuous_token_config.model_family, - model_path=model_config.path, - tokenizer_name_or_path=model_config.tokenizer_path, - chat_template_kwargs=self.apply_chat_template_kwargs, - ) - self.enable_continuous_token = True - # Continuous Token doesn't use the legacy removable system prompt. - self.system_prompt = None - else: - if continuous_token_config.enable and self.processor is not None: - logger.warning( - "Continuous Token is enabled but processor is set; falling back to legacy multimodal path." - ) - processing_class = self.processor if self.processor is not None else self.tokenizer - self.system_prompt = initialize_system_prompt(processing_class, **self.apply_chat_template_kwargs) - self.loop = get_event_loop() - - def _get_mm_processor_kwargs(self, audio_data: Optional[list[Any]] = None) -> dict[str, Any]: - mm_processor_kwargs = dict(self.mm_processor_kwargs or {}) - if audio_data is not None and "sampling_rate" not in mm_processor_kwargs: - sampling_rate = getattr(getattr(self.processor, "feature_extractor", None), "sampling_rate", None) - if sampling_rate is not None: - mm_processor_kwargs["sampling_rate"] = int(sampling_rate) - return mm_processor_kwargs - - async def process_vision_info(self, messages: list[dict]) -> dict: - """Backward-compatible wrapper for multi-modal extraction.""" - return await self.process_multi_modal_info(messages) - - async def process_multi_modal_info(self, messages: list[dict]) -> dict: - """Extract images, videos and audios from messages. - - Args: - messages (list[dict]): Input messages. - - Returns: - dict: Multi-modal data with keys like "images", "videos" and "audios". - """ - multi_modal_data = {} - if self.processor is not None: - image_patch_size = getattr(getattr(self.processor, "image_processor", None), "patch_size", 14) - if hasattr(self.dataset_cls, "process_multi_modal_info"): - images, videos, audios = await self.dataset_cls.process_multi_modal_info( - messages, image_patch_size=image_patch_size, config=self.data_config - ) - else: - images, videos = await self.dataset_cls.process_vision_info( - messages, image_patch_size=image_patch_size, config=self.data_config - ) - audios = None - if images is not None: - multi_modal_data["images"] = images - if videos is not None: - multi_modal_data["videos"] = videos - if audios is not None: - multi_modal_data["audios"] = audios - - return multi_modal_data - - async def ct_build_initial_tokens( - self, - messages: list[dict], - tools: list[dict] = None, - ) -> list[int]: - """Build the initial prompt token ids with Continuous Token.""" - prompt_ids = await self.loop.run_in_executor( - None, - lambda: self.continuous_token_builder.build_initial_tokens(messages, tools=tools), - ) - return self._cap_text_prompt_length(prompt_ids) - - async def ct_merge_non_assistant_msg( - self, - previous_messages: list[dict], - updated_messages: list[dict], - runtime_token_ids: list[int], - response_mask: list[int], - response_logprobs: Optional[list[float]] = None, - tools: list[dict] = None, - ): - """Merge appended non-assistant messages into runtime tokens and metadata.""" - merge_result = await self.loop.run_in_executor( - None, - lambda: self.continuous_token_builder.merge_non_assistant_tokens( - previous_messages, - updated_messages, - runtime_token_ids, - tools=tools, - ), - ) - aligned_response_mask, aligned_response_logprobs = self.continuous_token_builder.align_response_metadata( - merge_result, response_mask, response_logprobs - ) - return merge_result, aligned_response_mask, aligned_response_logprobs - - async def ct_merge_assistant_token( - self, - runtime_token_ids: list[int], - assistant_token_ids: list[int], - response_mask: list[int], - response_logprobs: Optional[list[float]] = None, - assistant_logprobs: Optional[list[float]] = None, - ): - """Merge assistant-generated tokens and align response metadata.""" - merge_result = await self.loop.run_in_executor( - None, - lambda: self.continuous_token_builder.merge_assistant_tokens( - runtime_token_ids, - assistant_token_ids, - ), - ) - aligned_response_mask, aligned_response_logprobs = self.continuous_token_builder.align_response_metadata( - merge_result, - response_mask, - response_logprobs, - assistant_logprobs=assistant_logprobs, - ) - return merge_result, aligned_response_mask, aligned_response_logprobs - - def _cap_text_prompt_length(self, prompt_ids: list[int]) -> list[int]: - prompt_length = self.rollout_config.prompt_length - if len(prompt_ids) > prompt_length: - logger.warning( - "Prompt of %d tokens exceeds rollout.prompt_length=%d; left-truncating.", - len(prompt_ids), - prompt_length, - ) - return prompt_ids[-prompt_length:] - return prompt_ids - - async def apply_chat_template( - self, - messages: list[dict], - tools: list[dict] = None, - images: list[Image.Image] = None, - videos: list[tuple[torch.Tensor, dict]] = None, - audios: list[Any] = None, - mm_processor_kwargs: Optional[dict[str, Any]] = None, - remove_system_prompt: bool = False, - ): - """Apply chat template to messages with optional tools, images, and videos. - - Args: - messages (list[dict]): Input messages. - tools (list[dict], optional): Tools schemas. Defaults to None. - images (list[Image.Image], optional): Input images. Defaults to None. - videos (list[tuple[torch.Tensor, dict]], optional): Input videos. Defaults to None. - remove_system_prompt (bool, optional): Whether to remove system prompt. Defaults to False. - - Returns: - list[int]: Prompt token ids. - """ - if self.processor is not None: - raw_prompt = await self.loop.run_in_executor( - None, - lambda: apply_chat_template( - self.processor, - messages, - tools=tools, - add_generation_prompt=True, - tokenize=False, - **self.apply_chat_template_kwargs, - ), - ) - - model_inputs = build_multimodal_processor_inputs( - self.processor, - text=[raw_prompt], - images=images, - videos=videos, - audio=audios, - mm_processor_kwargs=mm_processor_kwargs - if mm_processor_kwargs is not None - else self._get_mm_processor_kwargs(audios), - ) - prompt_ids = normalize_token_ids(model_inputs.pop("input_ids")) - else: - tokenized_prompt = await self.loop.run_in_executor( - None, - lambda: apply_chat_template( - self.tokenizer, - messages, - tools=tools, - add_generation_prompt=True, - tokenize=True, - **self.apply_chat_template_kwargs, - ), - ) - prompt_ids = normalize_token_ids(tokenized_prompt) - - if remove_system_prompt: - prompt_ids = prompt_ids[len(self.system_prompt) :] - - # Mirror the response-side ``response_ids[:response_length]`` cap on the prompt side: - # every prompt produced by the agent loop must fit in ``rollout.prompt_length`` so that - # ``_pad_token_ids`` (and downstream ``torch.cat``) can rely on uniform shapes. - # Multimodal prompts cannot be sliced here because placeholder tokens must remain - # aligned 1:1 with ``multi_modal_inputs`` features, so we fail loudly instead. - prompt_length = self.rollout_config.prompt_length - if len(prompt_ids) > prompt_length: - if images or videos or audios: - raise ValueError( - f"Multimodal prompt produced {len(prompt_ids)} tokens, exceeding " - f"rollout.prompt_length={prompt_length}. Truncating multimodal token " - f"sequences corrupts vision/audio feature alignment, so this is treated " - f"as a configuration error. Reduce the multimodal input size " - f"(e.g. ``total_pixels`` / ``max_pixels`` / fps / number of frames) or " - f"increase ``rollout.prompt_length``." - ) - prompt_ids = self._cap_text_prompt_length(prompt_ids) - - return prompt_ids - - @abstractmethod - async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput: - """Run agent loop to interact with LLM server and environment. - - Args: - sampling_params (Dict[str, Any]): LLM sampling params. - **kwargs: dataset fields from `verl.utils.dataset.RLHFDataset`. - - Returns: - AgentLoopOutput: Agent loop output. - """ - raise NotImplementedError - - -"""Agent loop registry: key is agent_name, value is a dict of agent loop config -used by hydra.utils.instantiate to initialize agent loop instance. - -https://hydra.cc/docs/advanced/instantiate_objects/overview/ -""" -_agent_loop_registry: dict[str, dict] = {} - - -def register(agent_name: str): - """Register agent loop class.""" - - def decorator(subclass: type[AgentLoopBase]) -> type[AgentLoopBase]: - fqdn = f"{subclass.__module__}.{subclass.__qualname__}" - _agent_loop_registry[agent_name] = {"_target_": fqdn} - return subclass - - return decorator - - -class AgentLoopWorker: - """Agent loop worker takes a batch of messages and run each message in an agent loop. - - Args: - config (DictConfig): whole config for main entrypoint. - llm_client (LLMServerClient): Client for the LLM server. - teacher_client (dict[str, LLMServerClient]): Client for multiple teacher servers. - reward_loop_worker_handles (List[ray.actor.ActorHandle]): Actor handles for streaming reward computation. - """ - - def __init__( - self, - config: DictConfig, - llm_client: LLMServerClient, - teacher_client: dict[str, LLMServerClient] = None, - reward_loop_worker_handles: list[ray.actor.ActorHandle] = None, - ): - self.config = config - self.llm_client = llm_client - self.teacher_client = teacher_client - self.reward_loop_worker_handles = reward_loop_worker_handles - - rollout_config, model_config = config.actor_rollout_ref.rollout, config.actor_rollout_ref.model - self.rollout_config: RolloutConfig = omega_conf_to_dataclass(rollout_config) - self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config) - - self.dataset_cls = get_dataset_class(config.data) - self.tokenizer = self.model_config.tokenizer - self.processor = self.model_config.processor - self.mm_processor_kwargs = config.data.get("mm_processor_kwargs", {}) - - # Online policy distillation - self.distillation_enabled = is_distillation_enabled(config.distillation) - # 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 - self.teacher_server_manager = AsyncTeacherLLMServerManager( - config=config, - teacher_client=teacher_client, - ) - - # Load tools once per worker; each trajectory just reuses self.tools. - tool_config_path = self.rollout_config.multi_turn.tool_config_path - function_tool_path = self.rollout_config.multi_turn.function_tool_path - self.tools = load_all_tools( - tool_config_path=resolve_config_path(tool_config_path) if tool_config_path else None, - function_tool_path=resolve_config_path(function_tool_path) if function_tool_path else None, - ) - - # Load custom agent loop implementations from config path - agent_loop_config_path = self.rollout_config.agent.agent_loop_config_path - if agent_loop_config_path: - resolved_path = resolve_config_path(agent_loop_config_path) - agent_loop_configs = OmegaConf.load(resolved_path) - for agent_loop_config in agent_loop_configs: - _agent_loop_registry[agent_loop_config.name] = agent_loop_config - if self.model_config.get("custom_chat_template", None) is not None: - if self.model_config.processor is not None: - self.model_config.processor.chat_template = self.model_config.custom_chat_template - self.model_config.tokenizer.chat_template = self.model_config.custom_chat_template - - trace_config = self.rollout_config.trace - RolloutTraceConfig.init( - self.rollout_config.trace.project_name, - self.rollout_config.trace.experiment_name, - trace_config.get("backend"), - trace_config.get("token2text", False), - trace_config.get("max_samples_per_step_per_worker", None), - ) - - def _get_mm_processor_kwargs(self, audio_data: Optional[list[Any]] = None) -> dict[str, Any]: - """Return multimodal processor kwargs with audio sampling-rate defaults.""" - mm_processor_kwargs = dict(self.mm_processor_kwargs or {}) - if audio_data is not None and "sampling_rate" not in mm_processor_kwargs: - sampling_rate = getattr(getattr(self.processor, "feature_extractor", None), "sampling_rate", None) - if sampling_rate is not None: - mm_processor_kwargs["sampling_rate"] = int(sampling_rate) - return mm_processor_kwargs - - async def generate_sequences(self, batch: DataProto) -> DataProto: - """Generate sequences from agent loop. - - Args: - batch (DataProto): Input batch. - - Returns: - DataProto: Output batch. - - prompts: [bsz, prompt_length], prompt token ids from dataset. - - responses: [bsz, response_length], output token ids include response tokens - from LLM generation and observation tokens from tool_calls. - - response_mask: [bsz, response_length], 1 for LLM generated tokens, 0 for observation/padding tokens. - - input_ids: [bsz, prompt_length + response_length], whole sequence token ids, including prompt tokens - and response tokens. - - attention_mask: [bsz, prompt_length + response_length], 0 for padding tokens, 1 for other tokens. - - position_ids: [bsz, prompt_length + response_length], incremental position ids. - - For multi-turn conversations: - responses: |<- LLM generation ->|<- tool_calls ->|<- LLM generation ->|<- padding ->| - response_mask: | 1, 1, 1, ..., 1, 1 | 0, 0, .., 0, 0 | 1, 1, 1, ..., 1, 1 | 0, 0, ..., 0| - """ - config = self.rollout_config - validate = batch.meta_info.get("validate", False) - sampling_params = dict( - temperature=config.temperature, - top_p=config.top_p, - top_k=config.top_k, - repetition_penalty=1.0, - logprobs=config.calculate_log_probs, - ) - - def apply_greedy_sampling_params(params: dict[str, Any]) -> None: - params["top_p"] = 1.0 - params["top_k"] = -1 - params["temperature"] = 0 - - # override sampling params for validation - if validate: - sampling_params["top_p"] = config.val_kwargs.top_p - sampling_params["top_k"] = config.val_kwargs.top_k - sampling_params["temperature"] = config.val_kwargs.temperature - - # by default, we assume it's a single turn agent - if "agent_name" not in batch.non_tensor_batch: - default_agent_loop = config.agent.default_agent_loop - batch.non_tensor_batch["agent_name"] = np.array([default_agent_loop] * len(batch), dtype=object) - - if "index" in batch.non_tensor_batch: - index = batch.non_tensor_batch["index"] - else: - index = np.arange(len(batch)) - - max_samples_per_worker = RolloutTraceConfig.get_instance().max_samples_per_step_per_worker - - # For n rollouts per sample, we trace all n rollouts for selected samples - # Note: This sampling happens per-worker, so total traces = max_samples_per_worker * num_workers * n - if max_samples_per_worker is not None: - unique_sample_indices = np.unique(index) - if max_samples_per_worker < len(unique_sample_indices): - selected_samples = set( - np.random.choice(unique_sample_indices, max_samples_per_worker, replace=False).tolist() - ) - traced_indices = set(i for i in range(len(batch)) if index[i] in selected_samples) - else: - traced_indices = set(range(len(batch))) - else: - traced_indices = set(range(len(batch))) - - trajectory_info = await get_trajectory_info( - batch.meta_info.get("global_steps", -1), index.tolist(), batch.meta_info.get("validate", False) - ) - - # NOTE: __do_sample__ is an internal per-sample override used by REMAX combined rollout. - # Do not forward it to concrete agent loops, which may reject unknown kwargs. - per_sample_do_sample = batch.non_tensor_batch.get("__do_sample__") - tasks = [] - for i in range(len(batch)): - trace_this_sample = i in traced_indices - kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items() if k != "__do_sample__"} - sample_sampling_params = dict(sampling_params) - if not validate and per_sample_do_sample is not None and not bool(per_sample_do_sample[i]): - apply_greedy_sampling_params(sample_sampling_params) - tasks.append( - asyncio.create_task( - self._run_agent_loop(sample_sampling_params, trajectory_info[i], trace=trace_this_sample, **kwargs) - ) - ) - outputs = await asyncio.gather(*tasks) - - output = self._postprocess( - outputs, input_non_tensor_batch=batch.non_tensor_batch, validate=batch.meta_info.get("validate", False) - ) - return output - - async def _run_agent_loop( - self, - sampling_params: dict[str, Any], - trajectory: dict[str, Any], - *, - agent_name: str, - trace: bool = True, - **kwargs, - ) -> _InternalAgentLoopOutput: - with rollout_trace_attr( - step=trajectory["step"], - sample_index=trajectory["sample_index"], - rollout_n=trajectory["rollout_n"], - validate=trajectory["validate"], - name="agent_loop", - trace=trace, - ): - assert agent_name in _agent_loop_registry, ( - f"Agent loop {agent_name} not registered, registered agent loops: {_agent_loop_registry.keys()}" - ) - - agent_loop_config = _agent_loop_registry[agent_name] - agent_loop = hydra.utils.instantiate( - config=agent_loop_config, - trainer_config=DictConfigWrap(config=self.config), - server_manager=self.llm_client, - tokenizer=self.tokenizer, - processor=self.processor, - dataset_cls=self.dataset_cls, - data_config=DictConfigWrap(self.config.data), - tools=ToolListWrap(self.tools), - ) - output: AgentLoopOutput = await agent_loop.run(sampling_params, **kwargs) - return await self._agent_loop_postprocess(output, trajectory["validate"], **kwargs) - - def _pad_token_ids( - self, - tokens: list[int], - *, - max_length: int, - padding_side: str, - return_attention_mask: bool, - ) -> dict[str, torch.Tensor]: - """Right/left pad a flat list of token ids to a ``(1, max_length)`` tensor.""" - # tokenizer.pad() with empty input returns dict with list values - # instead of tensors, which breaks downstream .dim() calls. - if not tokens: - pad_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0 - result = {"input_ids": torch.full((1, max_length), pad_id, dtype=torch.long)} - if return_attention_mask: - result["attention_mask"] = torch.zeros((1, max_length), dtype=torch.long) - return result - self.tokenizer.padding_side = padding_side - padded = self.tokenizer.pad( - {"input_ids": tokens}, - padding="max_length", - max_length=max_length, - return_tensors="pt", - return_attention_mask=return_attention_mask, - ) - if padded["input_ids"].dim() == 1: - padded["input_ids"] = padded["input_ids"].unsqueeze(0) - if return_attention_mask: - padded["attention_mask"] = padded["attention_mask"].unsqueeze(0) - return padded - - async def _agent_loop_postprocess(self, output, validate, **kwargs) -> _InternalAgentLoopOutput: - """Perform post-processing operations on the output of each individual agent loop.""" - output.extra_fields["raw_prompt"] = kwargs["raw_prompt"] - - # Some AgentLoop may have already computed the reward score, e.g SWE-agent. - - # NOTE: consistent with the legacy batch version of generate_sequences that existed in the - # deprecated vLLM SPMD rollout implementation. - # prompt_ids: left padded with zeros (e.g., [0,0,0,0,1,2,3,4]) - # response_ids: right padded with zeros (e.g., [5,6,7,8,0,0,0,0]) - # input_ids: concatenation of prompt + response - # Mask: - # For example, if the prompt is [1,2,3,4] and the response is [5,6,7,(tool start)8,9(tool end),10,11,12] - # - prompt_attention_mask: 0s for padding, 1s for tokens - # e.g., [0,0,0,0,1,1,1,1] - # - response_attention_mask: 0s for padding, 1s for tokens - # e.g., [1,1,1,1,1,1,1,1,1,1,1,0,0,0,0] - # attention_mask: concatenation of prompt_attention_mask and response_attention_mask - # e.g., [0,0,0,0,1,1,1,1(prompt),1,1,1,1,1,1,1,1,1,1,1,0,0,0,0(response)] - # - response_mask: 1s for LLM generated tokens, 0 for tool response/padding tokens - # e.g., [1,1,1,1,1,1,1,(tool start),0,0(tool end),1,1,0,0,0,0] - # - position_ids: sequential positions for tokens, starting at 0 - # e.g., [0,0,0,0,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,0,0,0,0] - - # TODO(wuxibin): remove padding and use tensordict. - prompt_output = self._pad_token_ids( - output.prompt_ids, - max_length=self.rollout_config.prompt_length, - padding_side="left", - return_attention_mask=True, - ) - - response_output = self._pad_token_ids( - output.response_ids, - max_length=self.rollout_config.response_length, - padding_side="right", - return_attention_mask=True, - ) - - response_mask_output = self._pad_token_ids( - output.response_mask, - max_length=self.rollout_config.response_length, - padding_side="right", - return_attention_mask=False, - ) - - response_logprobs = None - if output.response_logprobs is not None: - pad_size = self.rollout_config.response_length - len(output.response_logprobs) - response_logprobs = torch.tensor(output.response_logprobs + [0.0] * pad_size).unsqueeze(0) - - response_mask = response_mask_output["input_ids"] * response_output["attention_mask"] - attention_mask = torch.cat([prompt_output["attention_mask"], response_output["attention_mask"]], dim=1) - input_ids = torch.cat([prompt_output["input_ids"], response_output["input_ids"]], dim=1) - - routed_experts = None - if output.routed_experts is not None: - total_length = input_ids.shape[1] - length, layer_num, topk_num = output.routed_experts.shape - if isinstance(output.routed_experts, np.ndarray): - routed_experts_array = output.routed_experts - if not routed_experts_array.flags.writeable: - routed_experts_array = routed_experts_array.copy() - experts_tensor = torch.from_numpy(routed_experts_array) - elif isinstance(output.routed_experts, torch.Tensor): - experts_tensor = output.routed_experts - else: - raise TypeError(f"Unsupported type for routed_experts: {type(output.routed_experts)}") - routed_experts = torch.zeros(1, total_length, layer_num, topk_num, dtype=experts_tensor.dtype) - - # Calculate start position: left padding means original prompt starts at the end - start_pos = prompt_output["input_ids"].shape[1] - len(output.prompt_ids) - end_pos = min(start_pos + length, total_length) - - # Add boundary checks for robustness - if start_pos < 0 or end_pos > total_length: - raise ValueError( - f"Invalid position range: start_pos={start_pos}, end_pos={end_pos}, total_length={total_length}" - ) - - routed_experts[:, start_pos:end_pos] = experts_tensor.unsqueeze(0) - - multi_modal_inputs = self._compute_multi_modal_inputs(output, input_ids) - position_ids = self._compute_position_ids( - input_ids, - attention_mask, - multi_modal_inputs, - output.mm_processor_kwargs - if output.mm_processor_kwargs is not None - else self._get_mm_processor_kwargs( - output.multi_modal_data.get("audios") if output.multi_modal_data else None - ), - ) - await self._compute_score([output], kwargs=kwargs) - await self._compute_teacher_logprobs( - output, - prompt_ids=output.prompt_ids, - response_ids=output.response_ids, - validate=validate, - sample_kwargs=kwargs, - ) - teacher_ids, teacher_logprobs = ( - output.extra_fields.pop("teacher_ids", None), - output.extra_fields.pop("teacher_logprobs", None), - ) - if teacher_ids is not None and teacher_logprobs is not None: - # TODO(wuxibin): remove padding and use tensordict. - from verl.experimental.teacher_loop.teacher_manager import _pad_teacher_outputs - - teacher_ids, teacher_logprobs = _pad_teacher_outputs( - teacher_ids, - teacher_logprobs, - prompt_width=prompt_output["input_ids"].shape[1], - response_width=response_output["input_ids"].shape[1], - prompt_length=len(output.prompt_ids), - response_length=len(output.response_ids), - pad_token_id=self.tokenizer.pad_token_id, - ) - - return _InternalAgentLoopOutput( - prompt_ids=prompt_output["input_ids"], - response_ids=response_output["input_ids"], - input_ids=input_ids, - position_ids=position_ids, - response_mask=response_mask, - attention_mask=attention_mask, - response_logprobs=response_logprobs, - routed_experts=routed_experts, - multi_modal_inputs=multi_modal_inputs, - multi_modal_data=output.multi_modal_data, - mm_processor_kwargs=output.mm_processor_kwargs, - teacher_logprobs=teacher_logprobs, - teacher_ids=teacher_ids, - reward_score=output.reward_score, - num_turns=output.num_turns, - metrics=output.metrics, - extra_fields=output.extra_fields, - ) - - def _compute_multi_modal_inputs(self, output, input_ids) -> dict[str, torch.Tensor]: - """Compute multi-modal inputs with image, video and audio.""" - multi_modal_inputs = {} - if self.processor is None: - return multi_modal_inputs - - multi_modal_data = output.multi_modal_data or {} - images = multi_modal_data.get("images") - videos = multi_modal_data.get("videos") - audios = multi_modal_data.get("audios") - current_text = self.tokenizer.decode(input_ids.squeeze(0), skip_special_tokens=True) - - multi_modal_inputs = build_multimodal_processor_inputs( - self.processor, - text=[current_text], - images=images, - videos=videos, - audio=audios, - mm_processor_kwargs=output.mm_processor_kwargs - if output.mm_processor_kwargs is not None - else self._get_mm_processor_kwargs(audios), - ) - multi_modal_inputs.pop("input_ids", None) - multi_modal_inputs.pop("attention_mask", None) - - # We must use dict(multi_modal_inputs) to convert BatchFeature values to a new dict - # because np.array() only keeps the keys for BatchFeature. - multi_modal_inputs = dict(multi_modal_inputs.convert_to_tensors("pt")) - image_grid_thw = multi_modal_inputs.get("image_grid_thw") - if image_grid_thw is not None: - images_seqlens = torch.repeat_interleave(image_grid_thw[:, 1] * image_grid_thw[:, 2], image_grid_thw[:, 0]) - multi_modal_inputs["images_seqlens"] = images_seqlens - return multi_modal_inputs - - def _compute_position_ids( - self, - input_ids, - attention_mask, - multi_modal_inputs, - mm_processor_kwargs: Optional[dict[str, Any]] = None, - ) -> torch.Tensor: - """Compute position ids for multi-modal inputs.""" - # text-only OR non-M-RoPE multimodal (e.g. Gemma4) -> standard 1D positions - if self.processor is None or not hasattr(self.processor, "get_rope_index"): - return compute_position_id_with_mask(attention_mask) # (1, seq_len) - - multi_modal_kwargs = { - "image_grid_thw": multi_modal_inputs.get("image_grid_thw"), - "video_grid_thw": multi_modal_inputs.get("video_grid_thw"), - } - # For transformers>=5.3.0, mm_token_type_ids is only used to calculate position ids. - if multi_modal_inputs.pop("mm_token_type_ids", None) is not None: - mm_token_type_ids = torch.zeros_like(input_ids) - image_token_id = get_processor_token_id(self.processor, "image") - video_token_id = get_processor_token_id(self.processor, "video") - if image_token_id is not None: - mm_token_type_ids[0][input_ids[0] == image_token_id] = 1 - if video_token_id is not None: - mm_token_type_ids[0][input_ids[0] == video_token_id] = 2 - multi_modal_kwargs["mm_token_type_ids"] = mm_token_type_ids - - # Model's get_rope_index has been dynamically bind to the processor. - vision_position_ids, _ = self.processor.get_rope_index( - input_ids=input_ids, - attention_mask=attention_mask, - **multi_modal_kwargs, - ) - vision_position_ids = vision_position_ids.transpose(0, 1) # (3, 1, seq_len) => (1, 3, seq_len) - - valid_mask = attention_mask[0].bool() - text_position_ids = torch.ones((1, len(input_ids[0])), dtype=torch.long) - text_position_ids[0, valid_mask] = torch.arange(valid_mask.sum().item()) - text_position_ids = text_position_ids.unsqueeze(0) - position_ids = torch.cat((text_position_ids, vision_position_ids), dim=1) # (1, 4, seq_length) - return position_ids - - async def _compute_score(self, outputs: list[AgentLoopOutput], kwargs: dict) -> None: - """Compute reward score for all outputs in a trajectory; assigns result to outputs[-1].""" - enable_async_reward = self.reward_loop_worker_handles is not None - - final_output = outputs[-1] - if final_output.reward_score is None and enable_async_reward: - timing = {} - with simple_timer("compute_score", timing): - all_prompts, all_responses, all_input_ids, all_attention_mask, all_position_ids = [], [], [], [], [] - for output in outputs: - prompts = torch.tensor(output.prompt_ids, dtype=torch.int64) - responses = torch.tensor(output.response_ids, dtype=torch.int64) - input_ids = torch.cat([prompts, responses], dim=0) - attention_mask = torch.ones_like(input_ids, dtype=torch.int64) - multi_modal_inputs = self._compute_multi_modal_inputs(output, input_ids) - position_ids = self._compute_position_ids( - input_ids.unsqueeze(0), - attention_mask.unsqueeze(0), - multi_modal_inputs, - output.mm_processor_kwargs - if output.mm_processor_kwargs is not None - else self._get_mm_processor_kwargs( - output.multi_modal_data.get("audios") if output.multi_modal_data else None - ), - ).squeeze(0) - all_prompts.append(prompts) - all_responses.append(responses) - all_input_ids.append(input_ids) - all_attention_mask.append(attention_mask) - all_position_ids.append(position_ids) - - n = len(outputs) - batch = TensorDict( - { - "prompts": torch.nn.utils.rnn.pad_sequence(all_prompts, batch_first=True, padding_value=0), - "responses": torch.nn.utils.rnn.pad_sequence(all_responses, batch_first=True, padding_value=0), - "attention_mask": torch.nn.utils.rnn.pad_sequence( - all_attention_mask, batch_first=True, padding_value=0 - ), - "input_ids": torch.nn.utils.rnn.pad_sequence(all_input_ids, batch_first=True, padding_value=0), - "position_ids": torch.nn.utils.rnn.pad_sequence( - all_position_ids, batch_first=True, padding_value=0 - ), - }, - batch_size=n, - ) - non_tensor_batch = { - **{k: np.array([v] * n) for k, v in kwargs.items()}, - "__num_turns__": np.array([o.num_turns for o in outputs]), - "tool_extra_fields": np.array([o.extra_fields for o in outputs], dtype=object), - "prompt_len": np.array([len(o.prompt_ids) for o in outputs]), - "response_len": np.array([len(o.response_ids) for o in outputs]), - } - - data = DataProto( - batch=batch, - non_tensor_batch=non_tensor_batch, - ) - selected_reward_loop_worker_handle = random.choice(self.reward_loop_worker_handles) - result = await selected_reward_loop_worker_handle.compute_score.remote(data) - final_output.reward_score = result["reward_score"] - final_output.extra_fields["reward_extra_info"] = result["reward_extra_info"] - final_output.metrics.compute_score = timing["compute_score"] - - async def _compute_teacher_logprobs( - self, - output: AgentLoopOutput, - prompt_ids: list[int], - response_ids: list[int], - validate: bool, - sample_kwargs: Optional[dict[str, Any]] = None, - ) -> None: - """Compute teacher logprobs for single sample.""" - # 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) - if routing_value is not None: - # Non-tensor batch values arrive as 0-d numpy objects / arrays; normalize to Python. - routing_key = routing_value.item() if hasattr(routing_value, "item") else routing_value - teacher_ids, teacher_logprobs = await self.teacher_server_manager.compute_teacher_logprobs_single( - sequence_ids=prompt_ids + response_ids, - multi_modal_data=output.multi_modal_data, - mm_processor_kwargs=output.mm_processor_kwargs, - routing_key=routing_key, - ) - output.extra_fields["teacher_ids"] = teacher_ids - output.extra_fields["teacher_logprobs"] = teacher_logprobs - - def _postprocess( - self, - inputs: list[_InternalAgentLoopOutput], - input_non_tensor_batch: dict | None = None, - validate: bool = False, - ) -> DataProto: - """Process the padded outputs from _run_agent_loop and combine them into a batch.""" - # Convert lists back to tensors and stack them to create a batch. - prompt_ids = torch.cat([input.prompt_ids for input in inputs], dim=0) - response_ids = torch.cat([input.response_ids for input in inputs], dim=0) - response_mask = torch.cat([input.response_mask for input in inputs], dim=0) - attention_mask = torch.cat([input.attention_mask for input in inputs], dim=0) - input_ids = torch.cat([input.input_ids for input in inputs], dim=0) - position_ids = torch.cat([input.position_ids for input in inputs], dim=0) - optional_outputs = {} - if inputs[0].response_logprobs is not None: - optional_outputs["rollout_log_probs"] = torch.cat([input.response_logprobs for input in inputs], dim=0) - if inputs[0].routed_experts is not None: - optional_outputs["routed_experts"] = torch.cat([input.routed_experts for input in inputs], dim=0) - if inputs[0].teacher_logprobs is not None and inputs[0].teacher_ids is not None: - optional_outputs["teacher_logprobs"] = torch.cat([input.teacher_logprobs for input in inputs], dim=0) - optional_outputs["teacher_ids"] = torch.cat([input.teacher_ids for input in inputs], dim=0) - batch = TensorDict( - { - "prompts": prompt_ids, # [bsz, prompt_length] - "responses": response_ids, # [bsz, response_length] - "response_mask": response_mask, # [bsz, response_length] - "input_ids": input_ids, # [bsz, prompt_length + response_length] - "attention_mask": attention_mask, # [bsz, prompt_length + response_length] - # position_ids: [bsz, 3, prompt_length + response_length] or [bsz, prompt_length + response_length] - "position_ids": position_ids, - **optional_outputs, - }, - batch_size=len(inputs), - ) - - scores = [input.reward_score for input in inputs] - if all(score is not None for score in scores): - prompt_length = prompt_ids.size(1) - response_length = attention_mask[:, prompt_length:].sum(dim=1) - 1 - rm_scores = torch.zeros_like(response_mask, dtype=torch.float32) - rm_scores[torch.arange(response_mask.size(0)), response_length] = torch.tensor(scores, dtype=torch.float32) - batch["rm_scores"] = rm_scores - - non_tensor_batch = { - "__num_turns__": np.array([input.num_turns for input in inputs], dtype=np.int32), - } - if self.reward_loop_worker_handles is None and input_non_tensor_batch: - non_tensor_batch.update(input_non_tensor_batch) - - # add reward_extra_info to non_tensor_batch - reward_extra_infos = [input.extra_fields.get("reward_extra_info", {}) for input in inputs] - reward_extra_keys = list(reward_extra_infos[0].keys()) - for key in reward_extra_keys: - non_tensor_batch[key] = np.array([info[key] for info in reward_extra_infos]) - - # Add multi_modal_inputs to non_tensor_batch if any samples have them - multi_modal_inputs_list = [input.multi_modal_inputs for input in inputs] - if any(mmi is not None for mmi in multi_modal_inputs_list): - non_tensor_batch["multi_modal_inputs"] = np.array(multi_modal_inputs_list, dtype=object) - - metrics = [input.metrics.model_dump() for input in inputs] - # Collect extra fields from all inputs and convert them to np.ndarray - # Keep a stable set of keys so downstream batch concat stays consistent across agent loops. - extra_fields = {} - default_extra_keys = { - "turn_scores", - "tool_rewards", - "min_global_steps", - "max_global_steps", - "extras", - } - all_keys = set(key for input_item in inputs for key in input_item.extra_fields) | default_extra_keys - for key in all_keys: - temp_arr = np.empty(len(inputs), dtype=object) - temp_arr[:] = [input.extra_fields.get(key) for input in inputs] - extra_fields[key] = temp_arr - - non_tensor_batch.update(extra_fields) - - # Only include reward_extra_keys in meta_info if rm_scores is in batch - # This avoids conflicts when reward_tensor is merged later in ray_trainer.py - if "rm_scores" in batch.keys(): - meta_info = {"metrics": metrics, "reward_extra_keys": reward_extra_keys} - else: - meta_info = {"metrics": metrics} - - return DataProto( - batch=batch, - non_tensor_batch=non_tensor_batch, - meta_info=meta_info, - ) - - -async def get_trajectory_info(step, index, validate): - """Get trajectory info. - - Args: - step (int): global steps in the trainer. - index (list): form datastore extra_info.index column. - validate (bool): whether is a validate step. - - Returns: - list: trajectory. - """ - trajectory_info = [] - rollout_n = 0 - for i in range(len(index)): - if i > 0 and index[i - 1] == index[i]: - rollout_n += 1 - else: - rollout_n = 0 - trajectory_info.append({"step": step, "sample_index": index[i], "rollout_n": rollout_n, "validate": validate}) - return trajectory_info - - -class AgentLoopManager: - """Agent loop manager that manages a group of agent loop workers. - - Args: - config (DictConfig): whole config for main entrypoint. - llm_client (LLMServerClient): Client for the LLM server. - teacher_client (dict[str, LLMServerClient]): Client for multiple teacher servers. - reward_loop_worker_handles (List[ray.actor.ActorHandle]): Actor handles for streaming reward computation. - """ - - def __init__( - self, - config: DictConfig, - llm_client: LLMServerClient, - teacher_client: dict[str, LLMServerClient] = None, - reward_loop_worker_handles: list[ray.actor.ActorHandle] = None, - ): - self.config = config - self.rollout_config = config.actor_rollout_ref.rollout - self.model_config = config.actor_rollout_ref.model - self.llm_client = llm_client - self.teacher_client = teacher_client - self.reward_loop_worker_handles = reward_loop_worker_handles - - if not hasattr(self, "agent_loop_workers_class"): - self.agent_loop_workers_class = ray.remote(AgentLoopWorker) - - @classmethod - @auto_await - async def create(cls, *args, **kwargs): - """Create agent loop manager.""" - instance = cls(*args, **kwargs) - await instance._init_agent_loop_workers() - return instance - - async def _init_agent_loop_workers(self): - self.agent_loop_workers = [] - num_workers = self.rollout_config.agent.num_workers - - node_ids = [node["NodeID"] for node in ray.nodes() if node["Alive"] and node["Resources"].get("CPU", 0) > 0] - for i in range(num_workers): - # Round-robin scheduling over the all nodes - node_id = node_ids[i % len(node_ids)] - self.agent_loop_workers.append( - self.agent_loop_workers_class.options( - name=f"agent_loop_worker_{i}" + f"_{uuid4().hex[:8]}", - scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( - node_id=node_id, soft=True - ), - ).remote( - self.config, - self.llm_client, - self.teacher_client, - self.reward_loop_worker_handles, - ) - ) - - @auto_await - @SkipManager.annotate(role="rollout") - async def generate_sequences(self, prompts: DataProto) -> DataProto: - """Split input batch and dispatch to agent loop workers. - - Args: - prompts (DataProto): Input batch. - - Returns: - DataProto: Output batch. - """ - # Attach per-sample priority to the batch (like ``uid``) so each sample gets - # a globally-unique priority that flows to vLLM request scheduling. Assigned - # before chunking so chunks own disjoint ranges without per-worker offsets. - if "priority" not in prompts.non_tensor_batch: - prompts.non_tensor_batch["priority"] = np.arange(len(prompts), dtype=np.int64) - - chunkes = prompts.chunk(len(self.agent_loop_workers)) - outputs = await asyncio.gather( - *[ - worker.generate_sequences.remote(chunk) - for worker, chunk in zip(self.agent_loop_workers, chunkes, strict=True) - ] - ) - output = DataProto.concat(outputs) - - # calculate performance metrics - metrics = [output.meta_info.pop("metrics") for output in outputs] # List[List[Dict[str, str]]] - timing = self._performance_metrics(metrics, output) - - output.meta_info = {"timing": timing, **outputs[0].meta_info} - return output - - def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: DataProto) -> dict[str, float]: - timing = {} - t_generate_sequences = np.array([metric["generate_sequences"] for chunk in metrics for metric in chunk]) - t_tool_calls = np.array([metric["tool_calls"] for chunk in metrics for metric in chunk]) - t_compute_score = np.array([metric["compute_score"] for chunk in metrics for metric in chunk]) - num_preempted = np.array([metric["num_preempted"] for chunk in metrics for metric in chunk]) - timing["agent_loop/num_preempted/min"] = num_preempted.min() - timing["agent_loop/num_preempted/max"] = num_preempted.max() - timing["agent_loop/num_preempted/mean"] = num_preempted.mean() - timing["agent_loop/generate_sequences/min"] = t_generate_sequences.min() - timing["agent_loop/generate_sequences/max"] = t_generate_sequences.max() - timing["agent_loop/generate_sequences/mean"] = t_generate_sequences.mean() - timing["agent_loop/tool_calls/min"] = t_tool_calls.min() - timing["agent_loop/tool_calls/max"] = t_tool_calls.max() - timing["agent_loop/tool_calls/mean"] = t_tool_calls.mean() - timing["agent_loop/compute_score/min"] = t_compute_score.min() - timing["agent_loop/compute_score/max"] = t_compute_score.max() - timing["agent_loop/compute_score/mean"] = t_compute_score.mean() - - # batch sequence generation is bounded by the slowest sample - slowest = np.argmax(t_generate_sequences + t_tool_calls + t_compute_score) - prompt_length = output.batch["prompts"].shape[1] - timing["agent_loop/slowest/generate_sequences"] = t_generate_sequences[slowest] - timing["agent_loop/slowest/tool_calls"] = t_tool_calls[slowest] - timing["agent_loop/slowest/compute_score"] = t_compute_score[slowest] - timing["agent_loop/slowest/num_preempted"] = num_preempted[slowest] - - if "attention_mask" in output.batch: - attention_mask = output.batch["attention_mask"][slowest] - timing["agent_loop/slowest/prompt_length"] = attention_mask[:prompt_length].sum().item() - timing["agent_loop/slowest/response_length"] = attention_mask[prompt_length:].sum().item() - - return timing +# 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. +""" +Agent framework for multi-turn rollout and agentic reinforcement learning. +- AgentLoopBase: coroutine based abstract base class for agent loop. + - SingleTurnAgentLoop: single turn agent loop. + - ToolAgentLoop: ReAct agent loop with tool calling, with user defined tools. +- AgentLoopWorker: worker class for running agent loop coroutines in parallel. +- AgentLoopManager: manager class for running agent loop workers in parallel. + +AgentLoopManager is one specific agent-framework implementation in verl, +and is designed to be fully replaceable by other agent frameworks such as: +- NVIDIA Nemo-Gym +- AWS Bedrock AgentCore +- SWE-agent +- ... +""" + +import asyncio +import logging +import os +import random +from abc import ABC, abstractmethod +from typing import Any, Optional +from uuid import uuid4 + +import hydra +import numpy as np +import ray +import torch +from omegaconf import DictConfig, OmegaConf +from PIL import Image +from pydantic import BaseModel, ConfigDict +from tensordict import TensorDict +from transformers import AutoProcessor, AutoTokenizer + +from verl.experimental.agent_loop.utils import resolve_config_path +from verl.protocol import DataProto +from verl.tools.tool_registry import load_all_tools +from verl.trainer.distillation import is_distillation_enabled +from verl.utils.config import omega_conf_to_dataclass +from verl.utils.dataset.rl_dataset import RLHFDataset, get_dataset_class +from verl.utils.model import compute_position_id_with_mask +from verl.utils.profiler import simple_timer +from verl.utils.ray_utils import auto_await, get_event_loop +from verl.utils.rollout_trace import ( + RolloutTraceConfig, + rollout_trace_attr, +) +from verl.utils.skip import SkipManager +from verl.utils.tokenizer import ( + build_multimodal_processor_inputs, + get_processor_token_id, + normalize_token_ids, +) +from verl.utils.tokenizer.chat_template import apply_chat_template, initialize_system_prompt +from verl.utils.tokenizer.continuous_token_wiring import create_continuous_token_builder +from verl.workers.config import ( + HFModelConfig, + RolloutConfig, +) +from verl.workers.rollout.llm_server import LLMServerClient + +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +DEFAULT_ROUTING_CACHE_SIZE = 10000 + + +class AgentLoopMetrics(BaseModel): + """Agent loop performance metrics.""" + + generate_sequences: float = 0.0 + tool_calls: float = 0.0 + compute_score: float = 0.0 + num_preempted: int = -1 # -1 means not available + + +class AgentLoopOutput(BaseModel): + """Agent loop output.""" + + prompt_ids: list[int] + """Prompt token ids.""" + response_ids: list[int] + """Response token ids including LLM generated token, tool response token.""" + response_mask: list[int] + """Response mask, 1 for LLM generated token, 0 for tool response token.""" + response_logprobs: Optional[list[float]] = None + """Log probabilities for the response tokens.""" + routed_experts: Optional[Any] = None + """Routed experts for the total tokens.""" + multi_modal_data: Optional[dict[str, Any]] = None + """Multi-modal data for multi-modal tools.""" + reward_score: Optional[float] = None + """Reward score for the trajectory.""" + num_turns: int = 0 + """Number of chat turns, including user, assistant, tool.""" + metrics: AgentLoopMetrics + """Auxiliary performance metrics""" + extra_fields: dict[str, Any] = {} + """Extra fields for dynamic addition.""" + mm_processor_kwargs: Optional[dict[str, Any]] = None + """Processor/backend kwargs that must stay aligned across rollout and training paths.""" + + def as_dict(self) -> dict[str, Any]: + """Convert agent loop output to a dictionary.""" + output = self.model_dump(exclude_unset=True) + + output["prompts"] = torch.tensor(output.pop("prompt_ids"), dtype=torch.int64) + output["responses"] = torch.tensor(output.pop("response_ids"), dtype=torch.int64) + output["response_mask"] = torch.tensor(output.pop("response_mask"), dtype=torch.int64) + + response_logprobs = output.pop("response_logprobs", None) + if response_logprobs is not None: + output["rollout_log_probs"] = torch.tensor(response_logprobs, dtype=torch.float32) + + routed_experts = output.pop("routed_experts", None) + if routed_experts is not None: + output["routed_experts"] = torch.tensor(routed_experts, dtype=torch.int64) + + # rm_scores: reward score for each token + reward_score = output.pop("reward_score", None) + if reward_score is not None: + rm_scores = torch.zeros_like(output["response_mask"], dtype=torch.float32) + rm_scores[-1] = reward_score + output["rm_scores"] = rm_scores + + teacher_ids, teacher_logprobs = ( + output["extra_fields"].pop("teacher_ids", None), + output["extra_fields"].pop("teacher_logprobs", None), + ) + if teacher_ids is not None: + output["teacher_ids"] = teacher_ids + if teacher_logprobs is not None: + output["teacher_logprobs"] = teacher_logprobs + return output + + +class _InternalAgentLoopOutput(AgentLoopOutput): + """Internal agent loop output with padded sequences.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + prompt_ids: torch.Tensor + """Padded prompt token ids.""" + response_ids: torch.Tensor + """Padded response token ids.""" + input_ids: torch.Tensor + """Padded input ids(prompt_ids + response_ids).""" + position_ids: torch.Tensor + """Padded position ids.""" + response_mask: torch.Tensor + """Padded response mask.""" + attention_mask: torch.Tensor + """Padded attention mask.""" + response_logprobs: Optional[torch.Tensor] = None + """Padded log probabilities for the response tokens.""" + teacher_logprobs: Optional[torch.Tensor] = None + """Padded log probabilities from teacher model for prompt/response tokens.""" + teacher_ids: Optional[torch.Tensor] = None + """Padded token ids corresponding to the teacher log probabilities.""" + routed_experts: Optional[torch.Tensor] = None + """Padded routed experts for the total tokens.""" + multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None + """Multi-modal inputs for processors (e.g. pixel_values, image_grid_thw, video_grid_thw).""" + extra_fields: dict[str, Any] = {} + """Extra fields for dynamic addition.""" + + +class DictConfigWrap: + """Wrapper for DictConfig to avoid hydra.utils.instantiate recursive resolve.""" + + def __init__(self, config: DictConfig): + self.config = config + + +class ToolListWrap: + """Wraps a tool list so ``hydra.utils.instantiate`` doesn't recursively + resolve its elements (which would demote them to ``DictConfig``).""" + + def __init__(self, tools: list): + self.tools = tools + + +class AgentLoopBase(ABC): + """An agent loop takes an input message, chat with OpenAI compatible LLM server and interact with various + environments. + + Args: + trainer_config (DictConfig): whole config for main entrypoint. + server_manager (LLMServerClient): OpenAI compatible LLM server manager. + tokenizer (AutoTokenizer): Tokenizer for tokenize messages. + processor (AutoProcessor): Processor for process messages. + dataset_cls (type[Dataset]): Dataset class for creating dataset, Defaults to RLHFDataset. + data_config (DictConfigWrap): Dataset config. + """ + + def __init__( + self, + trainer_config: DictConfigWrap, + server_manager: LLMServerClient, + tokenizer: AutoTokenizer, + processor: AutoProcessor, + dataset_cls: type[RLHFDataset], + data_config: DictConfigWrap, + **kwargs, + ): + self.config = trainer_config.config + self.rollout_config = self.config.actor_rollout_ref.rollout + self.server_manager = server_manager + self.tokenizer = tokenizer + self.processor = processor + self.dataset_cls = dataset_cls + self.data_config = data_config.config + self.apply_chat_template_kwargs = self.data_config.get("apply_chat_template_kwargs", {}) + self.mm_processor_kwargs = self.data_config.get("mm_processor_kwargs", {}) + self.continuous_token_builder = None + self.enable_continuous_token = False + continuous_token_config = self.data_config.continuous_token + if continuous_token_config.enable and self.processor is None: + model_config = self.config.actor_rollout_ref.model + self.continuous_token_builder = create_continuous_token_builder( + self.tokenizer, + model_family=continuous_token_config.model_family, + model_path=model_config.path, + tokenizer_name_or_path=model_config.tokenizer_path, + chat_template_kwargs=self.apply_chat_template_kwargs, + ) + self.enable_continuous_token = True + # Continuous Token doesn't use the legacy removable system prompt. + self.system_prompt = None + else: + if continuous_token_config.enable and self.processor is not None: + logger.warning( + "Continuous Token is enabled but processor is set; falling back to legacy multimodal path." + ) + processing_class = self.processor if self.processor is not None else self.tokenizer + self.system_prompt = initialize_system_prompt(processing_class, **self.apply_chat_template_kwargs) + self.loop = get_event_loop() + + def _get_mm_processor_kwargs(self, audio_data: Optional[list[Any]] = None) -> dict[str, Any]: + mm_processor_kwargs = dict(self.mm_processor_kwargs or {}) + if audio_data is not None and "sampling_rate" not in mm_processor_kwargs: + sampling_rate = getattr(getattr(self.processor, "feature_extractor", None), "sampling_rate", None) + if sampling_rate is not None: + mm_processor_kwargs["sampling_rate"] = int(sampling_rate) + return mm_processor_kwargs + + async def process_vision_info(self, messages: list[dict]) -> dict: + """Backward-compatible wrapper for multi-modal extraction.""" + return await self.process_multi_modal_info(messages) + + async def process_multi_modal_info(self, messages: list[dict]) -> dict: + """Extract images, videos and audios from messages. + + Args: + messages (list[dict]): Input messages. + + Returns: + dict: Multi-modal data with keys like "images", "videos" and "audios". + """ + multi_modal_data = {} + if self.processor is not None: + image_patch_size = getattr(getattr(self.processor, "image_processor", None), "patch_size", 14) + if hasattr(self.dataset_cls, "process_multi_modal_info"): + images, videos, audios = await self.dataset_cls.process_multi_modal_info( + messages, image_patch_size=image_patch_size, config=self.data_config + ) + else: + images, videos = await self.dataset_cls.process_vision_info( + messages, image_patch_size=image_patch_size, config=self.data_config + ) + audios = None + if images is not None: + multi_modal_data["images"] = images + if videos is not None: + multi_modal_data["videos"] = videos + if audios is not None: + multi_modal_data["audios"] = audios + + return multi_modal_data + + async def ct_build_initial_tokens( + self, + messages: list[dict], + tools: list[dict] = None, + ) -> list[int]: + """Build the initial prompt token ids with Continuous Token.""" + prompt_ids = await self.loop.run_in_executor( + None, + lambda: self.continuous_token_builder.build_initial_tokens(messages, tools=tools), + ) + return self._cap_text_prompt_length(prompt_ids) + + async def ct_merge_non_assistant_msg( + self, + previous_messages: list[dict], + updated_messages: list[dict], + runtime_token_ids: list[int], + response_mask: list[int], + response_logprobs: Optional[list[float]] = None, + tools: list[dict] = None, + ): + """Merge appended non-assistant messages into runtime tokens and metadata.""" + merge_result = await self.loop.run_in_executor( + None, + lambda: self.continuous_token_builder.merge_non_assistant_tokens( + previous_messages, + updated_messages, + runtime_token_ids, + tools=tools, + ), + ) + aligned_response_mask, aligned_response_logprobs = self.continuous_token_builder.align_response_metadata( + merge_result, response_mask, response_logprobs + ) + return merge_result, aligned_response_mask, aligned_response_logprobs + + async def ct_merge_assistant_token( + self, + runtime_token_ids: list[int], + assistant_token_ids: list[int], + response_mask: list[int], + response_logprobs: Optional[list[float]] = None, + assistant_logprobs: Optional[list[float]] = None, + ): + """Merge assistant-generated tokens and align response metadata.""" + merge_result = await self.loop.run_in_executor( + None, + lambda: self.continuous_token_builder.merge_assistant_tokens( + runtime_token_ids, + assistant_token_ids, + ), + ) + aligned_response_mask, aligned_response_logprobs = self.continuous_token_builder.align_response_metadata( + merge_result, + response_mask, + response_logprobs, + assistant_logprobs=assistant_logprobs, + ) + return merge_result, aligned_response_mask, aligned_response_logprobs + + def _cap_text_prompt_length(self, prompt_ids: list[int]) -> list[int]: + prompt_length = self.rollout_config.prompt_length + if len(prompt_ids) > prompt_length: + logger.warning( + "Prompt of %d tokens exceeds rollout.prompt_length=%d; left-truncating.", + len(prompt_ids), + prompt_length, + ) + return prompt_ids[-prompt_length:] + return prompt_ids + + async def apply_chat_template( + self, + messages: list[dict], + tools: list[dict] = None, + images: list[Image.Image] = None, + videos: list[tuple[torch.Tensor, dict]] = None, + audios: list[Any] = None, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + remove_system_prompt: bool = False, + ): + """Apply chat template to messages with optional tools, images, and videos. + + Args: + messages (list[dict]): Input messages. + tools (list[dict], optional): Tools schemas. Defaults to None. + images (list[Image.Image], optional): Input images. Defaults to None. + videos (list[tuple[torch.Tensor, dict]], optional): Input videos. Defaults to None. + remove_system_prompt (bool, optional): Whether to remove system prompt. Defaults to False. + + Returns: + list[int]: Prompt token ids. + """ + if self.processor is not None: + raw_prompt = await self.loop.run_in_executor( + None, + lambda: apply_chat_template( + self.processor, + messages, + tools=tools, + add_generation_prompt=True, + tokenize=False, + **self.apply_chat_template_kwargs, + ), + ) + + model_inputs = build_multimodal_processor_inputs( + self.processor, + text=[raw_prompt], + images=images, + videos=videos, + audio=audios, + mm_processor_kwargs=mm_processor_kwargs + if mm_processor_kwargs is not None + else self._get_mm_processor_kwargs(audios), + ) + prompt_ids = normalize_token_ids(model_inputs.pop("input_ids")) + else: + tokenized_prompt = await self.loop.run_in_executor( + None, + lambda: apply_chat_template( + self.tokenizer, + messages, + tools=tools, + add_generation_prompt=True, + tokenize=True, + **self.apply_chat_template_kwargs, + ), + ) + prompt_ids = normalize_token_ids(tokenized_prompt) + + if remove_system_prompt: + prompt_ids = prompt_ids[len(self.system_prompt) :] + + # Mirror the response-side ``response_ids[:response_length]`` cap on the prompt side: + # every prompt produced by the agent loop must fit in ``rollout.prompt_length`` so that + # ``_pad_token_ids`` (and downstream ``torch.cat``) can rely on uniform shapes. + # Multimodal prompts cannot be sliced here because placeholder tokens must remain + # aligned 1:1 with ``multi_modal_inputs`` features, so we fail loudly instead. + prompt_length = self.rollout_config.prompt_length + if len(prompt_ids) > prompt_length: + if images or videos or audios: + raise ValueError( + f"Multimodal prompt produced {len(prompt_ids)} tokens, exceeding " + f"rollout.prompt_length={prompt_length}. Truncating multimodal token " + f"sequences corrupts vision/audio feature alignment, so this is treated " + f"as a configuration error. Reduce the multimodal input size " + f"(e.g. ``total_pixels`` / ``max_pixels`` / fps / number of frames) or " + f"increase ``rollout.prompt_length``." + ) + prompt_ids = self._cap_text_prompt_length(prompt_ids) + + return prompt_ids + + @abstractmethod + async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput: + """Run agent loop to interact with LLM server and environment. + + Args: + sampling_params (Dict[str, Any]): LLM sampling params. + **kwargs: dataset fields from `verl.utils.dataset.RLHFDataset`. + + Returns: + AgentLoopOutput: Agent loop output. + """ + raise NotImplementedError + + +"""Agent loop registry: key is agent_name, value is a dict of agent loop config +used by hydra.utils.instantiate to initialize agent loop instance. + +https://hydra.cc/docs/advanced/instantiate_objects/overview/ +""" +_agent_loop_registry: dict[str, dict] = {} + + +def register(agent_name: str): + """Register agent loop class.""" + + def decorator(subclass: type[AgentLoopBase]) -> type[AgentLoopBase]: + fqdn = f"{subclass.__module__}.{subclass.__qualname__}" + _agent_loop_registry[agent_name] = {"_target_": fqdn} + return subclass + + return decorator + + +class AgentLoopWorker: + """Agent loop worker takes a batch of messages and run each message in an agent loop. + + Args: + config (DictConfig): whole config for main entrypoint. + llm_client (LLMServerClient): Client for the LLM server. + teacher_client (dict[str, LLMServerClient]): Client for multiple teacher servers. + reward_loop_worker_handles (List[ray.actor.ActorHandle]): Actor handles for streaming reward computation. + """ + + def __init__( + self, + config: DictConfig, + llm_client: LLMServerClient, + teacher_client: dict[str, LLMServerClient] = None, + reward_loop_worker_handles: list[ray.actor.ActorHandle] = None, + ): + self.config = config + self.llm_client = llm_client + self.teacher_client = teacher_client + self.reward_loop_worker_handles = reward_loop_worker_handles + + rollout_config, model_config = config.actor_rollout_ref.rollout, config.actor_rollout_ref.model + self.rollout_config: RolloutConfig = omega_conf_to_dataclass(rollout_config) + self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config) + + self.dataset_cls = get_dataset_class(config.data) + self.tokenizer = self.model_config.tokenizer + self.processor = self.model_config.processor + self.mm_processor_kwargs = config.data.get("mm_processor_kwargs", {}) + + # Online policy distillation + self.distillation_enabled = is_distillation_enabled(config.distillation) + # 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 + self.teacher_server_manager = AsyncTeacherLLMServerManager( + config=config, + teacher_client=teacher_client, + ) + + # Load tools once per worker; each trajectory just reuses self.tools. + tool_config_path = self.rollout_config.multi_turn.tool_config_path + function_tool_path = self.rollout_config.multi_turn.function_tool_path + self.tools = load_all_tools( + tool_config_path=resolve_config_path(tool_config_path) if tool_config_path else None, + function_tool_path=resolve_config_path(function_tool_path) if function_tool_path else None, + ) + + # Load custom agent loop implementations from config path + agent_loop_config_path = self.rollout_config.agent.agent_loop_config_path + if agent_loop_config_path: + resolved_path = resolve_config_path(agent_loop_config_path) + agent_loop_configs = OmegaConf.load(resolved_path) + for agent_loop_config in agent_loop_configs: + _agent_loop_registry[agent_loop_config.name] = agent_loop_config + if self.model_config.get("custom_chat_template", None) is not None: + if self.model_config.processor is not None: + self.model_config.processor.chat_template = self.model_config.custom_chat_template + self.model_config.tokenizer.chat_template = self.model_config.custom_chat_template + + trace_config = self.rollout_config.trace + RolloutTraceConfig.init( + self.rollout_config.trace.project_name, + self.rollout_config.trace.experiment_name, + trace_config.get("backend"), + trace_config.get("token2text", False), + trace_config.get("max_samples_per_step_per_worker", None), + ) + + def _get_mm_processor_kwargs(self, audio_data: Optional[list[Any]] = None) -> dict[str, Any]: + """Return multimodal processor kwargs with audio sampling-rate defaults.""" + mm_processor_kwargs = dict(self.mm_processor_kwargs or {}) + if audio_data is not None and "sampling_rate" not in mm_processor_kwargs: + sampling_rate = getattr(getattr(self.processor, "feature_extractor", None), "sampling_rate", None) + if sampling_rate is not None: + mm_processor_kwargs["sampling_rate"] = int(sampling_rate) + return mm_processor_kwargs + + async def generate_sequences(self, batch: DataProto) -> DataProto: + """Generate sequences from agent loop. + + Args: + batch (DataProto): Input batch. + + Returns: + DataProto: Output batch. + - prompts: [bsz, prompt_length], prompt token ids from dataset. + - responses: [bsz, response_length], output token ids include response tokens + from LLM generation and observation tokens from tool_calls. + - response_mask: [bsz, response_length], 1 for LLM generated tokens, 0 for observation/padding tokens. + - input_ids: [bsz, prompt_length + response_length], whole sequence token ids, including prompt tokens + and response tokens. + - attention_mask: [bsz, prompt_length + response_length], 0 for padding tokens, 1 for other tokens. + - position_ids: [bsz, prompt_length + response_length], incremental position ids. + + For multi-turn conversations: + responses: |<- LLM generation ->|<- tool_calls ->|<- LLM generation ->|<- padding ->| + response_mask: | 1, 1, 1, ..., 1, 1 | 0, 0, .., 0, 0 | 1, 1, 1, ..., 1, 1 | 0, 0, ..., 0| + """ + config = self.rollout_config + validate = batch.meta_info.get("validate", False) + sampling_params = dict( + temperature=config.temperature, + top_p=config.top_p, + top_k=config.top_k, + repetition_penalty=1.0, + logprobs=config.calculate_log_probs, + ) + + def apply_greedy_sampling_params(params: dict[str, Any]) -> None: + params["top_p"] = 1.0 + params["top_k"] = -1 + params["temperature"] = 0 + + # override sampling params for validation + if validate: + sampling_params["top_p"] = config.val_kwargs.top_p + sampling_params["top_k"] = config.val_kwargs.top_k + sampling_params["temperature"] = config.val_kwargs.temperature + + # by default, we assume it's a single turn agent + if "agent_name" not in batch.non_tensor_batch: + default_agent_loop = config.agent.default_agent_loop + batch.non_tensor_batch["agent_name"] = np.array([default_agent_loop] * len(batch), dtype=object) + + if "index" in batch.non_tensor_batch: + index = batch.non_tensor_batch["index"] + else: + index = np.arange(len(batch)) + + max_samples_per_worker = RolloutTraceConfig.get_instance().max_samples_per_step_per_worker + + # For n rollouts per sample, we trace all n rollouts for selected samples + # Note: This sampling happens per-worker, so total traces = max_samples_per_worker * num_workers * n + if max_samples_per_worker is not None: + unique_sample_indices = np.unique(index) + if max_samples_per_worker < len(unique_sample_indices): + selected_samples = set( + np.random.choice(unique_sample_indices, max_samples_per_worker, replace=False).tolist() + ) + traced_indices = set(i for i in range(len(batch)) if index[i] in selected_samples) + else: + traced_indices = set(range(len(batch))) + else: + traced_indices = set(range(len(batch))) + + trajectory_info = await get_trajectory_info( + batch.meta_info.get("global_steps", -1), index.tolist(), batch.meta_info.get("validate", False) + ) + + # NOTE: __do_sample__ is an internal per-sample override used by REMAX combined rollout. + # Do not forward it to concrete agent loops, which may reject unknown kwargs. + per_sample_do_sample = batch.non_tensor_batch.get("__do_sample__") + tasks = [] + for i in range(len(batch)): + trace_this_sample = i in traced_indices + kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items() if k != "__do_sample__"} + sample_sampling_params = dict(sampling_params) + if not validate and per_sample_do_sample is not None and not bool(per_sample_do_sample[i]): + apply_greedy_sampling_params(sample_sampling_params) + tasks.append( + asyncio.create_task( + self._run_agent_loop(sample_sampling_params, trajectory_info[i], trace=trace_this_sample, **kwargs) + ) + ) + outputs = await asyncio.gather(*tasks) + + output = self._postprocess( + outputs, input_non_tensor_batch=batch.non_tensor_batch, validate=batch.meta_info.get("validate", False) + ) + return output + + async def _run_agent_loop( + self, + sampling_params: dict[str, Any], + trajectory: dict[str, Any], + *, + agent_name: str, + trace: bool = True, + **kwargs, + ) -> _InternalAgentLoopOutput: + with rollout_trace_attr( + step=trajectory["step"], + sample_index=trajectory["sample_index"], + rollout_n=trajectory["rollout_n"], + validate=trajectory["validate"], + name="agent_loop", + trace=trace, + ): + assert agent_name in _agent_loop_registry, ( + f"Agent loop {agent_name} not registered, registered agent loops: {_agent_loop_registry.keys()}" + ) + + agent_loop_config = _agent_loop_registry[agent_name] + agent_loop = hydra.utils.instantiate( + config=agent_loop_config, + trainer_config=DictConfigWrap(config=self.config), + server_manager=self.llm_client, + tokenizer=self.tokenizer, + processor=self.processor, + dataset_cls=self.dataset_cls, + data_config=DictConfigWrap(self.config.data), + tools=ToolListWrap(self.tools), + ) + output: AgentLoopOutput = await agent_loop.run(sampling_params, **kwargs) + return await self._agent_loop_postprocess(output, trajectory["validate"], **kwargs) + + def _pad_token_ids( + self, + tokens: list[int], + *, + max_length: int, + padding_side: str, + return_attention_mask: bool, + ) -> dict[str, torch.Tensor]: + """Right/left pad a flat list of token ids to a ``(1, max_length)`` tensor.""" + # tokenizer.pad() with empty input returns dict with list values + # instead of tensors, which breaks downstream .dim() calls. + if not tokens: + pad_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0 + result = {"input_ids": torch.full((1, max_length), pad_id, dtype=torch.long)} + if return_attention_mask: + result["attention_mask"] = torch.zeros((1, max_length), dtype=torch.long) + return result + self.tokenizer.padding_side = padding_side + padded = self.tokenizer.pad( + {"input_ids": tokens}, + padding="max_length", + max_length=max_length, + return_tensors="pt", + return_attention_mask=return_attention_mask, + ) + if padded["input_ids"].dim() == 1: + padded["input_ids"] = padded["input_ids"].unsqueeze(0) + if return_attention_mask: + padded["attention_mask"] = padded["attention_mask"].unsqueeze(0) + return padded + + async def _agent_loop_postprocess(self, output, validate, **kwargs) -> _InternalAgentLoopOutput: + """Perform post-processing operations on the output of each individual agent loop.""" + output.extra_fields["raw_prompt"] = kwargs["raw_prompt"] + + # Some AgentLoop may have already computed the reward score, e.g SWE-agent. + + # NOTE: consistent with the legacy batch version of generate_sequences that existed in the + # deprecated vLLM SPMD rollout implementation. + # prompt_ids: left padded with zeros (e.g., [0,0,0,0,1,2,3,4]) + # response_ids: right padded with zeros (e.g., [5,6,7,8,0,0,0,0]) + # input_ids: concatenation of prompt + response + # Mask: + # For example, if the prompt is [1,2,3,4] and the response is [5,6,7,(tool start)8,9(tool end),10,11,12] + # - prompt_attention_mask: 0s for padding, 1s for tokens + # e.g., [0,0,0,0,1,1,1,1] + # - response_attention_mask: 0s for padding, 1s for tokens + # e.g., [1,1,1,1,1,1,1,1,1,1,1,0,0,0,0] + # attention_mask: concatenation of prompt_attention_mask and response_attention_mask + # e.g., [0,0,0,0,1,1,1,1(prompt),1,1,1,1,1,1,1,1,1,1,1,0,0,0,0(response)] + # - response_mask: 1s for LLM generated tokens, 0 for tool response/padding tokens + # e.g., [1,1,1,1,1,1,1,(tool start),0,0(tool end),1,1,0,0,0,0] + # - position_ids: sequential positions for tokens, starting at 0 + # e.g., [0,0,0,0,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,0,0,0,0] + + # TODO(wuxibin): remove padding and use tensordict. + prompt_output = self._pad_token_ids( + output.prompt_ids, + max_length=self.rollout_config.prompt_length, + padding_side="left", + return_attention_mask=True, + ) + + response_output = self._pad_token_ids( + output.response_ids, + max_length=self.rollout_config.response_length, + padding_side="right", + return_attention_mask=True, + ) + + response_mask_output = self._pad_token_ids( + output.response_mask, + max_length=self.rollout_config.response_length, + padding_side="right", + return_attention_mask=False, + ) + + response_logprobs = None + if output.response_logprobs is not None: + pad_size = self.rollout_config.response_length - len(output.response_logprobs) + response_logprobs = torch.tensor(output.response_logprobs + [0.0] * pad_size).unsqueeze(0) + + response_mask = response_mask_output["input_ids"] * response_output["attention_mask"] + attention_mask = torch.cat([prompt_output["attention_mask"], response_output["attention_mask"]], dim=1) + input_ids = torch.cat([prompt_output["input_ids"], response_output["input_ids"]], dim=1) + + routed_experts = None + if output.routed_experts is not None: + total_length = input_ids.shape[1] + length, layer_num, topk_num = output.routed_experts.shape + if isinstance(output.routed_experts, np.ndarray): + routed_experts_array = output.routed_experts + if not routed_experts_array.flags.writeable: + routed_experts_array = routed_experts_array.copy() + experts_tensor = torch.from_numpy(routed_experts_array) + elif isinstance(output.routed_experts, torch.Tensor): + experts_tensor = output.routed_experts + else: + raise TypeError(f"Unsupported type for routed_experts: {type(output.routed_experts)}") + routed_experts = torch.zeros(1, total_length, layer_num, topk_num, dtype=experts_tensor.dtype) + + # Calculate start position: left padding means original prompt starts at the end + start_pos = prompt_output["input_ids"].shape[1] - len(output.prompt_ids) + end_pos = min(start_pos + length, total_length) + + # Add boundary checks for robustness + if start_pos < 0 or end_pos > total_length: + raise ValueError( + f"Invalid position range: start_pos={start_pos}, end_pos={end_pos}, total_length={total_length}" + ) + + routed_experts[:, start_pos:end_pos] = experts_tensor.unsqueeze(0) + + multi_modal_inputs = self._compute_multi_modal_inputs(output, input_ids) + position_ids = self._compute_position_ids( + input_ids, + attention_mask, + multi_modal_inputs, + output.mm_processor_kwargs + if output.mm_processor_kwargs is not None + else self._get_mm_processor_kwargs( + output.multi_modal_data.get("audios") if output.multi_modal_data else None + ), + ) + await self._compute_score([output], kwargs=kwargs) + await self._compute_teacher_logprobs( + output, + prompt_ids=output.prompt_ids, + response_ids=output.response_ids, + validate=validate, + sample_kwargs=kwargs, + ) + teacher_ids, teacher_logprobs = ( + output.extra_fields.pop("teacher_ids", None), + output.extra_fields.pop("teacher_logprobs", None), + ) + if teacher_ids is not None and teacher_logprobs is not None: + # TODO(wuxibin): remove padding and use tensordict. + from verl.experimental.teacher_loop.teacher_manager import _pad_teacher_outputs + + teacher_ids, teacher_logprobs = _pad_teacher_outputs( + teacher_ids, + teacher_logprobs, + prompt_width=prompt_output["input_ids"].shape[1], + response_width=response_output["input_ids"].shape[1], + prompt_length=len(output.prompt_ids), + response_length=len(output.response_ids), + pad_token_id=self.tokenizer.pad_token_id, + ) + + return _InternalAgentLoopOutput( + prompt_ids=prompt_output["input_ids"], + response_ids=response_output["input_ids"], + input_ids=input_ids, + position_ids=position_ids, + response_mask=response_mask, + attention_mask=attention_mask, + response_logprobs=response_logprobs, + routed_experts=routed_experts, + multi_modal_inputs=multi_modal_inputs, + multi_modal_data=output.multi_modal_data, + mm_processor_kwargs=output.mm_processor_kwargs, + teacher_logprobs=teacher_logprobs, + teacher_ids=teacher_ids, + reward_score=output.reward_score, + num_turns=output.num_turns, + metrics=output.metrics, + extra_fields=output.extra_fields, + ) + + def _compute_multi_modal_inputs(self, output, input_ids) -> dict[str, torch.Tensor]: + """Compute multi-modal inputs with image, video and audio.""" + multi_modal_inputs = {} + if self.processor is None: + return multi_modal_inputs + + multi_modal_data = output.multi_modal_data or {} + images = multi_modal_data.get("images") + videos = multi_modal_data.get("videos") + audios = multi_modal_data.get("audios") + current_text = self.tokenizer.decode(input_ids.squeeze(0), skip_special_tokens=True) + + multi_modal_inputs = build_multimodal_processor_inputs( + self.processor, + text=[current_text], + images=images, + videos=videos, + audio=audios, + mm_processor_kwargs=output.mm_processor_kwargs + if output.mm_processor_kwargs is not None + else self._get_mm_processor_kwargs(audios), + ) + multi_modal_inputs.pop("input_ids", None) + multi_modal_inputs.pop("attention_mask", None) + + # We must use dict(multi_modal_inputs) to convert BatchFeature values to a new dict + # because np.array() only keeps the keys for BatchFeature. + multi_modal_inputs = dict(multi_modal_inputs.convert_to_tensors("pt")) + image_grid_thw = multi_modal_inputs.get("image_grid_thw") + if image_grid_thw is not None: + images_seqlens = torch.repeat_interleave(image_grid_thw[:, 1] * image_grid_thw[:, 2], image_grid_thw[:, 0]) + multi_modal_inputs["images_seqlens"] = images_seqlens + return multi_modal_inputs + + def _compute_position_ids( + self, + input_ids, + attention_mask, + multi_modal_inputs, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + ) -> torch.Tensor: + """Compute position ids for multi-modal inputs.""" + # text-only OR non-M-RoPE multimodal (e.g. Gemma4) -> standard 1D positions + if self.processor is None or not hasattr(self.processor, "get_rope_index"): + return compute_position_id_with_mask(attention_mask) # (1, seq_len) + + multi_modal_kwargs = { + "image_grid_thw": multi_modal_inputs.get("image_grid_thw"), + "video_grid_thw": multi_modal_inputs.get("video_grid_thw"), + } + # For transformers>=5.3.0, mm_token_type_ids is only used to calculate position ids. + if multi_modal_inputs.pop("mm_token_type_ids", None) is not None: + mm_token_type_ids = torch.zeros_like(input_ids) + image_token_id = get_processor_token_id(self.processor, "image") + video_token_id = get_processor_token_id(self.processor, "video") + if image_token_id is not None: + mm_token_type_ids[0][input_ids[0] == image_token_id] = 1 + if video_token_id is not None: + mm_token_type_ids[0][input_ids[0] == video_token_id] = 2 + multi_modal_kwargs["mm_token_type_ids"] = mm_token_type_ids + + # Model's get_rope_index has been dynamically bind to the processor. + vision_position_ids, _ = self.processor.get_rope_index( + input_ids=input_ids, + attention_mask=attention_mask, + **multi_modal_kwargs, + ) + vision_position_ids = vision_position_ids.transpose(0, 1) # (3, 1, seq_len) => (1, 3, seq_len) + + valid_mask = attention_mask[0].bool() + text_position_ids = torch.ones((1, len(input_ids[0])), dtype=torch.long) + text_position_ids[0, valid_mask] = torch.arange(valid_mask.sum().item()) + text_position_ids = text_position_ids.unsqueeze(0) + position_ids = torch.cat((text_position_ids, vision_position_ids), dim=1) # (1, 4, seq_length) + return position_ids + + async def _compute_score(self, outputs: list[AgentLoopOutput], kwargs: dict) -> None: + """Compute reward score for all outputs in a trajectory; assigns result to outputs[-1].""" + enable_async_reward = self.reward_loop_worker_handles is not None + + final_output = outputs[-1] + if final_output.reward_score is None and enable_async_reward: + timing = {} + with simple_timer("compute_score", timing): + all_prompts, all_responses, all_input_ids, all_attention_mask, all_position_ids = [], [], [], [], [] + for output in outputs: + prompts = torch.tensor(output.prompt_ids, dtype=torch.int64) + responses = torch.tensor(output.response_ids, dtype=torch.int64) + input_ids = torch.cat([prompts, responses], dim=0) + attention_mask = torch.ones_like(input_ids, dtype=torch.int64) + multi_modal_inputs = self._compute_multi_modal_inputs(output, input_ids) + position_ids = self._compute_position_ids( + input_ids.unsqueeze(0), + attention_mask.unsqueeze(0), + multi_modal_inputs, + output.mm_processor_kwargs + if output.mm_processor_kwargs is not None + else self._get_mm_processor_kwargs( + output.multi_modal_data.get("audios") if output.multi_modal_data else None + ), + ).squeeze(0) + all_prompts.append(prompts) + all_responses.append(responses) + all_input_ids.append(input_ids) + all_attention_mask.append(attention_mask) + all_position_ids.append(position_ids) + + n = len(outputs) + batch = TensorDict( + { + "prompts": torch.nn.utils.rnn.pad_sequence(all_prompts, batch_first=True, padding_value=0), + "responses": torch.nn.utils.rnn.pad_sequence(all_responses, batch_first=True, padding_value=0), + "attention_mask": torch.nn.utils.rnn.pad_sequence( + all_attention_mask, batch_first=True, padding_value=0 + ), + "input_ids": torch.nn.utils.rnn.pad_sequence(all_input_ids, batch_first=True, padding_value=0), + "position_ids": torch.nn.utils.rnn.pad_sequence( + all_position_ids, batch_first=True, padding_value=0 + ), + }, + batch_size=n, + ) + non_tensor_batch = { + **{k: np.array([v] * n) for k, v in kwargs.items()}, + "__num_turns__": np.array([o.num_turns for o in outputs]), + "tool_extra_fields": np.array([o.extra_fields for o in outputs], dtype=object), + "prompt_len": np.array([len(o.prompt_ids) for o in outputs]), + "response_len": np.array([len(o.response_ids) for o in outputs]), + } + + data = DataProto( + batch=batch, + non_tensor_batch=non_tensor_batch, + ) + selected_reward_loop_worker_handle = random.choice(self.reward_loop_worker_handles) + result = await selected_reward_loop_worker_handle.compute_score.remote(data) + final_output.reward_score = result["reward_score"] + final_output.extra_fields["reward_extra_info"] = result["reward_extra_info"] + final_output.metrics.compute_score = timing["compute_score"] + + async def _compute_teacher_logprobs( + self, + output: AgentLoopOutput, + prompt_ids: list[int], + response_ids: list[int], + validate: bool, + sample_kwargs: Optional[dict[str, Any]] = None, + ) -> None: + """Compute teacher logprobs for single sample.""" + # 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) + if routing_value is not None: + # Non-tensor batch values arrive as 0-d numpy objects / arrays; normalize to Python. + routing_key = routing_value.item() if hasattr(routing_value, "item") else routing_value + teacher_ids, teacher_logprobs = await self.teacher_server_manager.compute_teacher_logprobs_single( + sequence_ids=prompt_ids + response_ids, + multi_modal_data=output.multi_modal_data, + mm_processor_kwargs=output.mm_processor_kwargs, + routing_key=routing_key, + ) + output.extra_fields["teacher_ids"] = teacher_ids + output.extra_fields["teacher_logprobs"] = teacher_logprobs + + def _postprocess( + self, + inputs: list[_InternalAgentLoopOutput], + input_non_tensor_batch: dict | None = None, + validate: bool = False, + ) -> DataProto: + """Process the padded outputs from _run_agent_loop and combine them into a batch.""" + # Convert lists back to tensors and stack them to create a batch. + prompt_ids = torch.cat([input.prompt_ids for input in inputs], dim=0) + response_ids = torch.cat([input.response_ids for input in inputs], dim=0) + response_mask = torch.cat([input.response_mask for input in inputs], dim=0) + attention_mask = torch.cat([input.attention_mask for input in inputs], dim=0) + input_ids = torch.cat([input.input_ids for input in inputs], dim=0) + position_ids = torch.cat([input.position_ids for input in inputs], dim=0) + optional_outputs = {} + if inputs[0].response_logprobs is not None: + optional_outputs["rollout_log_probs"] = torch.cat([input.response_logprobs for input in inputs], dim=0) + if inputs[0].routed_experts is not None: + optional_outputs["routed_experts"] = torch.cat([input.routed_experts for input in inputs], dim=0) + if inputs[0].teacher_logprobs is not None and inputs[0].teacher_ids is not None: + optional_outputs["teacher_logprobs"] = torch.cat([input.teacher_logprobs for input in inputs], dim=0) + optional_outputs["teacher_ids"] = torch.cat([input.teacher_ids for input in inputs], dim=0) + batch = TensorDict( + { + "prompts": prompt_ids, # [bsz, prompt_length] + "responses": response_ids, # [bsz, response_length] + "response_mask": response_mask, # [bsz, response_length] + "input_ids": input_ids, # [bsz, prompt_length + response_length] + "attention_mask": attention_mask, # [bsz, prompt_length + response_length] + # position_ids: [bsz, 3, prompt_length + response_length] or [bsz, prompt_length + response_length] + "position_ids": position_ids, + **optional_outputs, + }, + batch_size=len(inputs), + ) + + scores = [input.reward_score for input in inputs] + if all(score is not None for score in scores): + prompt_length = prompt_ids.size(1) + response_length = attention_mask[:, prompt_length:].sum(dim=1) - 1 + rm_scores = torch.zeros_like(response_mask, dtype=torch.float32) + rm_scores[torch.arange(response_mask.size(0)), response_length] = torch.tensor(scores, dtype=torch.float32) + batch["rm_scores"] = rm_scores + + non_tensor_batch = { + "__num_turns__": np.array([input.num_turns for input in inputs], dtype=np.int32), + } + if self.reward_loop_worker_handles is None and input_non_tensor_batch: + non_tensor_batch.update(input_non_tensor_batch) + + # add reward_extra_info to non_tensor_batch + reward_extra_infos = [input.extra_fields.get("reward_extra_info", {}) for input in inputs] + reward_extra_keys = list(reward_extra_infos[0].keys()) + for key in reward_extra_keys: + non_tensor_batch[key] = np.array([info[key] for info in reward_extra_infos]) + + # Add multi_modal_inputs to non_tensor_batch if any samples have them + multi_modal_inputs_list = [input.multi_modal_inputs for input in inputs] + if any(mmi is not None for mmi in multi_modal_inputs_list): + non_tensor_batch["multi_modal_inputs"] = np.array(multi_modal_inputs_list, dtype=object) + + metrics = [input.metrics.model_dump() for input in inputs] + # Collect extra fields from all inputs and convert them to np.ndarray + # Keep a stable set of keys so downstream batch concat stays consistent across agent loops. + extra_fields = {} + default_extra_keys = { + "turn_scores", + "tool_rewards", + "min_global_steps", + "max_global_steps", + "extras", + } + all_keys = set(key for input_item in inputs for key in input_item.extra_fields) | default_extra_keys + for key in all_keys: + temp_arr = np.empty(len(inputs), dtype=object) + temp_arr[:] = [input.extra_fields.get(key) for input in inputs] + extra_fields[key] = temp_arr + + non_tensor_batch.update(extra_fields) + + # Only include reward_extra_keys in meta_info if rm_scores is in batch + # This avoids conflicts when reward_tensor is merged later in ray_trainer.py + if "rm_scores" in batch.keys(): + meta_info = {"metrics": metrics, "reward_extra_keys": reward_extra_keys} + else: + meta_info = {"metrics": metrics} + + return DataProto( + batch=batch, + non_tensor_batch=non_tensor_batch, + meta_info=meta_info, + ) + + +async def get_trajectory_info(step, index, validate): + """Get trajectory info. + + Args: + step (int): global steps in the trainer. + index (list): form datastore extra_info.index column. + validate (bool): whether is a validate step. + + Returns: + list: trajectory. + """ + trajectory_info = [] + rollout_n = 0 + for i in range(len(index)): + if i > 0 and index[i - 1] == index[i]: + rollout_n += 1 + else: + rollout_n = 0 + trajectory_info.append({"step": step, "sample_index": index[i], "rollout_n": rollout_n, "validate": validate}) + return trajectory_info + + +class AgentLoopManager: + """Agent loop manager that manages a group of agent loop workers. + + Args: + config (DictConfig): whole config for main entrypoint. + llm_client (LLMServerClient): Client for the LLM server. + teacher_client (dict[str, LLMServerClient]): Client for multiple teacher servers. + reward_loop_worker_handles (List[ray.actor.ActorHandle]): Actor handles for streaming reward computation. + """ + + def __init__( + self, + config: DictConfig, + llm_client: LLMServerClient, + teacher_client: dict[str, LLMServerClient] = None, + reward_loop_worker_handles: list[ray.actor.ActorHandle] = None, + ): + self.config = config + self.rollout_config = config.actor_rollout_ref.rollout + self.model_config = config.actor_rollout_ref.model + self.llm_client = llm_client + self.teacher_client = teacher_client + self.reward_loop_worker_handles = reward_loop_worker_handles + + if not hasattr(self, "agent_loop_workers_class"): + self.agent_loop_workers_class = ray.remote(AgentLoopWorker) + + @classmethod + @auto_await + async def create(cls, *args, **kwargs): + """Create agent loop manager.""" + instance = cls(*args, **kwargs) + await instance._init_agent_loop_workers() + return instance + + async def _init_agent_loop_workers(self): + self.agent_loop_workers = [] + num_workers = self.rollout_config.agent.num_workers + + node_ids = [node["NodeID"] for node in ray.nodes() if node["Alive"] and node["Resources"].get("CPU", 0) > 0] + for i in range(num_workers): + # Round-robin scheduling over the all nodes + node_id = node_ids[i % len(node_ids)] + self.agent_loop_workers.append( + self.agent_loop_workers_class.options( + name=f"agent_loop_worker_{i}" + f"_{uuid4().hex[:8]}", + scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( + node_id=node_id, soft=True + ), + ).remote( + self.config, + self.llm_client, + self.teacher_client, + self.reward_loop_worker_handles, + ) + ) + + @auto_await + @SkipManager.annotate(role="rollout") + async def generate_sequences(self, prompts: DataProto) -> DataProto: + """Split input batch and dispatch to agent loop workers. + + Args: + prompts (DataProto): Input batch. + + Returns: + DataProto: Output batch. + """ + # Attach per-sample priority to the batch (like ``uid``) so each sample gets + # a globally-unique priority that flows to vLLM request scheduling. Assigned + # before chunking so chunks own disjoint ranges without per-worker offsets. + if "priority" not in prompts.non_tensor_batch: + prompts.non_tensor_batch["priority"] = np.arange(len(prompts), dtype=np.int64) + + chunkes = prompts.chunk(len(self.agent_loop_workers)) + outputs = await asyncio.gather( + *[ + worker.generate_sequences.remote(chunk) + for worker, chunk in zip(self.agent_loop_workers, chunkes, strict=True) + ] + ) + output = DataProto.concat(outputs) + + # calculate performance metrics + metrics = [output.meta_info.pop("metrics") for output in outputs] # List[List[Dict[str, str]]] + timing = self._performance_metrics(metrics, output) + + output.meta_info = {"timing": timing, **outputs[0].meta_info} + return output + + def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: DataProto) -> dict[str, float]: + timing = {} + t_generate_sequences = np.array([metric["generate_sequences"] for chunk in metrics for metric in chunk]) + t_tool_calls = np.array([metric["tool_calls"] for chunk in metrics for metric in chunk]) + t_compute_score = np.array([metric["compute_score"] for chunk in metrics for metric in chunk]) + num_preempted = np.array([metric["num_preempted"] for chunk in metrics for metric in chunk]) + timing["agent_loop/num_preempted/min"] = num_preempted.min() + timing["agent_loop/num_preempted/max"] = num_preempted.max() + timing["agent_loop/num_preempted/mean"] = num_preempted.mean() + timing["agent_loop/generate_sequences/min"] = t_generate_sequences.min() + timing["agent_loop/generate_sequences/max"] = t_generate_sequences.max() + timing["agent_loop/generate_sequences/mean"] = t_generate_sequences.mean() + timing["agent_loop/tool_calls/min"] = t_tool_calls.min() + timing["agent_loop/tool_calls/max"] = t_tool_calls.max() + timing["agent_loop/tool_calls/mean"] = t_tool_calls.mean() + timing["agent_loop/compute_score/min"] = t_compute_score.min() + timing["agent_loop/compute_score/max"] = t_compute_score.max() + timing["agent_loop/compute_score/mean"] = t_compute_score.mean() + + # batch sequence generation is bounded by the slowest sample + slowest = np.argmax(t_generate_sequences + t_tool_calls + t_compute_score) + prompt_length = output.batch["prompts"].shape[1] + timing["agent_loop/slowest/generate_sequences"] = t_generate_sequences[slowest] + timing["agent_loop/slowest/tool_calls"] = t_tool_calls[slowest] + timing["agent_loop/slowest/compute_score"] = t_compute_score[slowest] + timing["agent_loop/slowest/num_preempted"] = num_preempted[slowest] + + if "attention_mask" in output.batch: + attention_mask = output.batch["attention_mask"][slowest] + timing["agent_loop/slowest/prompt_length"] = attention_mask[:prompt_length].sum().item() + timing["agent_loop/slowest/response_length"] = attention_mask[prompt_length:].sum().item() + + return timing diff --git a/verl/experimental/teacher_loop/teacher_fsdp_manager.py b/verl/experimental/teacher_loop/teacher_fsdp_manager.py index 7c0f162158c..6ea6a0bb884 100644 --- a/verl/experimental/teacher_loop/teacher_fsdp_manager.py +++ b/verl/experimental/teacher_loop/teacher_fsdp_manager.py @@ -1,157 +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"] +# 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/distillation/fsdp/losses.py b/verl/trainer/distillation/fsdp/losses.py index ca7d4a8c80b..97d3ebf485a 100644 --- a/verl/trainer/distillation/fsdp/losses.py +++ b/verl/trainer/distillation/fsdp/losses.py @@ -1,272 +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 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 +# 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/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index a5f90e45f42..5fa6effe221 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -1,1935 +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 - - # 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() +# 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 0d9e9dd8ce8..4c4e7bf69a8 100644 --- a/verl/workers/config/distillation.py +++ b/verl/workers/config/distillation.py @@ -1,364 +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 .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 +# 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/teacher/fsdp_teacher_worker.py b/verl/workers/teacher/fsdp_teacher_worker.py index b55809998f1..bd03d07fbed 100644 --- a/verl/workers/teacher/fsdp_teacher_worker.py +++ b/verl/workers/teacher/fsdp_teacher_worker.py @@ -1,171 +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 +# 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 From 87cd055e872f9fdd7b976fcd797a0934b6dba265 Mon Sep 17 00:00:00 2001 From: YangxSong <1713458583@qq.com> Date: Mon, 6 Jul 2026 12:38:18 +0800 Subject: [PATCH 14/14] fix: add use_vllm_teacher attribute to agent_loop test mock --- .../agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py | 1 + 1 file changed, 1 insertion(+) 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()