From 924dc87a15cd65eb9323685f70218f1d0aef301e Mon Sep 17 00:00:00 2001 From: msz12345 Date: Thu, 2 Jul 2026 10:15:42 +0800 Subject: [PATCH] add fusenode on teacher forward and student train in fully_async_opd --- .../test_fused_teacher_config.py | 58 ++++ .../run_fully_async_opd_fusenode.sh | 252 ++++++++++++++++++ verl/experimental/agent_loop/agent_loop.py | 8 +- .../fully_async_policy/fully_async_main.py | 4 + .../fully_async_rollouter.py | 6 +- .../fully_async_policy/fully_async_trainer.py | 65 ++++- .../experimental/separation/engine_workers.py | 133 ++++++++- .../_generated_ppo_megatron_trainer.yaml | 1 + .../_generated_ppo_torchtitan_trainer.yaml | 1 + .../config/_generated_ppo_trainer.yaml | 1 + .../config/_generated_ppo_veomni_trainer.yaml | 1 + .../config/distillation/distillation.yaml | 7 +- verl/trainer/distillation/losses.py | 8 +- verl/workers/config/distillation.py | 58 +++- 14 files changed, 576 insertions(+), 27 deletions(-) create mode 100644 tests/experimental/fully_async_policy/test_fused_teacher_config.py create mode 100755 tests/special_e2e/run_fully_async_opd_fusenode.sh diff --git a/tests/experimental/fully_async_policy/test_fused_teacher_config.py b/tests/experimental/fully_async_policy/test_fused_teacher_config.py new file mode 100644 index 00000000000..aba33d6309d --- /dev/null +++ b/tests/experimental/fully_async_policy/test_fused_teacher_config.py @@ -0,0 +1,58 @@ +# 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. + +import pytest + +from verl.workers.config import ( + DistillationConfig, + DistillationLossConfig, + DistillationTeacherModelConfig, +) + + +def _teacher(): + return DistillationTeacherModelConfig( + key="business-data", + model_path="/tmp/teacher", + num_replicas=0, + ) + + +def test_trainer_teacher_does_not_require_a_resource_pool(): + config = DistillationConfig( + enabled=True, + teacher_execution="trainer", + n_gpus_per_node=0, + nnodes=0, + teacher_models={"teacher_model": _teacher()}, + distillation_loss=DistillationLossConfig(loss_mode="k1", use_policy_gradient=True), + ) + + assert list(config.teacher_models) == ["default"] + assert config.teacher_models["default"].model_path == "/tmp/teacher" + + +def test_trainer_teacher_rejects_topk_until_scorer_is_implemented(): + with pytest.raises(NotImplementedError, match="forward_kl_topk"): + DistillationConfig( + enabled=True, + teacher_execution="trainer", + n_gpus_per_node=0, + nnodes=0, + teacher_models={"teacher_model": _teacher()}, + distillation_loss=DistillationLossConfig( + loss_mode="forward_kl_topk", + use_policy_gradient=False, + ), + ) diff --git a/tests/special_e2e/run_fully_async_opd_fusenode.sh b/tests/special_e2e/run_fully_async_opd_fusenode.sh new file mode 100755 index 00000000000..8ee617f08f8 --- /dev/null +++ b/tests/special_e2e/run_fully_async_opd_fusenode.sh @@ -0,0 +1,252 @@ +#!/usr/bin/env bash +set -xeuo pipefail + +# Workaround for NVIDIA driver bug (r560-r575) causing SIGSEGV in ncclCuMemHostEnable() +# on PCIe machines without P2P access. See: https://github.com/NVIDIA/nccl/issues/1838 +export NCCL_CUMEM_ENABLE=0 +export NCCL_CUMEM_HOST_ENABLE=0 + +# Fully async OPD on GSM8K. +# Student rollout uses single-turn generation; teacher forward and student +# update share the fused Megatron trainer resource pool. + +############################ Quick Config ############################ + +ROLLOUT_NAME="sglang" + +# Keep MTP/speculative decoding disabled while validating fused-node OPD. +mtp_params=( + actor_rollout_ref.model.mtp.enable=False + actor_rollout_ref.model.mtp.enable_train=False + actor_rollout_ref.model.mtp.enable_rollout=False +) + +STUDENT_MODEL=${STUDENT_MODEL:-deepseek-ai/DeepSeek-R1-Distill-Qwen-7B} +GSM8K_TEACHER_MODEL=${TEACHER_MODEL:-Qwen/Qwen2.5-7B} + +DISTILLATION_LOSS_MODE="k1" +USE_POLICY_GRADIENT=True + +MAX_PROMPT=${MAX_PROMPT:-1600} +MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-32768} +MAX_NUM_TOKENS=$(( MAX_PROMPT + MAX_RESPONSE_LENGTH + 1 )) + +# Fully async specific +ROLLOUT_NNODES=2 +N_GPUS_ROLLOUT=8 +TRAINER_NNODES=2 +N_GPUS_TRAINING=8 +TOTAL_ROLLOUT_STEPS=${TOTAL_ROLLOUT_STEPS:-4096} + +# Megatron parallelism +GEN_TP=4 +TRAIN_TP=4 +TRAIN_PP=2 + +STALENESS_THRESHOLD=0.5 +TRIGGER_PARAMETER_SYNC_STEP=4 +SAVE_EVERY_TRAIN_STEPS=${SAVE_EVERY_TRAIN_STEPS:-16} +CHECKPOINT_DIR=${CHECKPOINT_DIR:-"/your_local_dir/"} + +if (( SAVE_EVERY_TRAIN_STEPS % TRIGGER_PARAMETER_SYNC_STEP != 0 )); then + echo "SAVE_EVERY_TRAIN_STEPS must be divisible by TRIGGER_PARAMETER_SYNC_STEP" >&2 + exit 1 +fi +SAVE_FREQ_PARAM_VERSIONS=$(( SAVE_EVERY_TRAIN_STEPS / TRIGGER_PARAMETER_SYNC_STEP )) + +############################ Data ############################ + +GSM8K_TRAIN="/your_local_dir/to_gsm8k/train.parquet" +GSM8K_TEST="/your_local_dir/to_gsm8k/test.parquet" + +TRAIN_FILES="['${GSM8K_TRAIN}']" +TEST_FILES="['${GSM8K_TEST}']" + +ACTOR_OFFLOAD=True + +############################ Parameter Groups ############################ + +DATA=( + data.train_files="$TRAIN_FILES" + data.val_files="$TEST_FILES" + data.prompt_key=prompt + data.truncation='left' + data.max_prompt_length=$MAX_PROMPT + data.max_response_length=$MAX_RESPONSE_LENGTH + data.train_batch_size=0 + data.gen_batch_size=1 + data.return_raw_chat=True + data.image_key=images +) + +MODEL=( + actor_rollout_ref.model.path="${STUDENT_MODEL}" + actor_rollout_ref.model.enable_gradient_checkpointing=True + actor_rollout_ref.model.use_remove_padding=True +) + +STUDENT=( + actor_rollout_ref.actor.strategy=megatron + actor_rollout_ref.actor.optim.lr=1e-6 + actor_rollout_ref.actor.optim.lr_warmup_steps=-1 + actor_rollout_ref.actor.optim.lr_decay_steps=10000000 + actor_rollout_ref.actor.optim.weight_decay=0.1 + actor_rollout_ref.actor.ppo_mini_batch_size=16 + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 + actor_rollout_ref.actor.entropy_coeff=0 + actor_rollout_ref.actor.loss_agg_mode="token-mean" + actor_rollout_ref.actor.clip_ratio_low=0.2 + actor_rollout_ref.actor.clip_ratio_high=0.28 + actor_rollout_ref.actor.clip_ratio_c=10.0 + actor_rollout_ref.actor.use_kl_loss=False + actor_rollout_ref.actor.kl_loss_coef=0.0 + actor_rollout_ref.actor.use_dynamic_bsz=True + actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$MAX_NUM_TOKENS + actor_rollout_ref.actor.megatron.param_offload=False + actor_rollout_ref.actor.megatron.optimizer_offload=False + actor_rollout_ref.actor.megatron.grad_offload=True + actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${TRAIN_PP} + actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TRAIN_TP} + actor_rollout_ref.actor.megatron.expert_model_parallel_size=1 + actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=1 + actor_rollout_ref.actor.megatron.context_parallel_size=1 + actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${TRAIN_PP} + actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TRAIN_TP} + actor_rollout_ref.ref.megatron.param_offload=True + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True + actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=$MAX_NUM_TOKENS +) + +ROLLOUT=( + actor_rollout_ref.rollout.name=$ROLLOUT_NAME + actor_rollout_ref.rollout.mode=async + actor_rollout_ref.rollout.n=4 + actor_rollout_ref.rollout.calculate_log_probs=True + actor_rollout_ref.rollout.prompt_length=$MAX_PROMPT + actor_rollout_ref.rollout.response_length=$MAX_RESPONSE_LENGTH + actor_rollout_ref.rollout.single_turn_response_length=$MAX_RESPONSE_LENGTH + actor_rollout_ref.rollout.gpu_memory_utilization=0.5 + actor_rollout_ref.rollout.temperature=1.0 + actor_rollout_ref.rollout.top_p=1.0 + actor_rollout_ref.rollout.top_k=-1 + actor_rollout_ref.rollout.disable_log_stats=False + actor_rollout_ref.rollout.max_model_len=$MAX_NUM_TOKENS + actor_rollout_ref.rollout.max_num_batched_tokens=$MAX_NUM_TOKENS + actor_rollout_ref.rollout.max_num_seqs=16 + actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True + actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=$MAX_NUM_TOKENS + actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP} + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 + actor_rollout_ref.rollout.val_kwargs.top_p=0.7 + actor_rollout_ref.rollout.val_kwargs.top_k=-1 + actor_rollout_ref.rollout.val_kwargs.do_sample=True + actor_rollout_ref.rollout.val_kwargs.n=1 + actor_rollout_ref.rollout.multi_turn.enable=False + actor_rollout_ref.rollout.agent.num_workers=1 + actor_rollout_ref.rollout.checkpoint_engine.backend='nccl' + actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=1024 + actor_rollout_ref.rollout.enforce_eager=False + +actor_rollout_ref.rollout.engine_kwargs.sglang.mamba_scheduler_strategy=no_buffer + +actor_rollout_ref.rollout.engine_kwargs.sglang.disable_radix_cache=True + +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_memory_saver=False + +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_weights_cpu_backup=False + +actor_rollout_ref.rollout.engine_kwargs.sglang.disable_overlap_schedule=True +) + +# Single-teacher OPD: teacher Megatron forward shares the trainer resource pool. +DISTILLATION=( + distillation.enabled=True + distillation.teacher_execution=trainer + distillation.teacher_key=data_source + distillation.n_gpus_per_node=0 + distillation.nnodes=0 + # --- single teacher --- + +distillation.teacher_models.gsm8k.key="openai/gsm8k" + +distillation.teacher_models.gsm8k.model_path="${GSM8K_TEACHER_MODEL}" + +distillation.teacher_models.gsm8k.num_replicas=0 + +distillation.teacher_models.gsm8k.inference.name=$ROLLOUT_NAME + # --- loss --- + distillation.distillation_loss.loss_mode=$DISTILLATION_LOSS_MODE + distillation.distillation_loss.topk=1 + distillation.distillation_loss.use_task_rewards=False + distillation.distillation_loss.use_policy_gradient=$USE_POLICY_GRADIENT + distillation.distillation_loss.loss_max_clamp=10.0 + distillation.distillation_loss.log_prob_min_clamp=-10.0 +) + +ALGORITHM=( + algorithm.adv_estimator=grpo + algorithm.use_kl_in_reward=False + algorithm.kl_ctrl.kl_coef=0.0 + algorithm.rollout_correction.bypass_mode=False +) + +REWARD=( + reward.reward_manager.name=dapo_judge + ++reward.custom_reward_function.path='verl/utils/reward_score/zero.py' + ++reward.custom_reward_function.name='compute_score' + +reward.reward_kwargs.overlong_buffer_cfg.enable=False + +reward.reward_kwargs.overlong_buffer_cfg.len=128 + +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=1.0 + +reward.reward_kwargs.overlong_buffer_cfg.log=False + +reward.reward_kwargs.max_resp_len=${MAX_RESPONSE_LENGTH} +) + +TRAINER=( + trainer.logger='["console","tensorboard"]' + trainer.project_name='verl-test-fully-async-opd' + trainer.experiment_name="qwen2.5-7b-fully-async-fused-node-opd" + trainer.val_before_train=False + trainer.save_freq=${SAVE_FREQ_PARAM_VERSIONS} + trainer.default_local_dir="${CHECKPOINT_DIR}" + trainer.resume_mode=disable + trainer.nnodes=${TRAINER_NNODES} + trainer.n_gpus_per_node=${N_GPUS_TRAINING} + trainer.log_val_generations=10 + +trainer.use_legacy_worker_impl=disable + trainer.total_epochs=2 + trainer.test_freq=-1 +) + +ASYNC_TRAINING=( + rollout.nnodes=${ROLLOUT_NNODES} + rollout.n_gpus_per_node=${N_GPUS_ROLLOUT} + rollout.total_rollout_steps=${TOTAL_ROLLOUT_STEPS} + async_training.staleness_threshold=${STALENESS_THRESHOLD} + async_training.partial_rollout=True + async_training.trigger_parameter_sync_step=${TRIGGER_PARAMETER_SYNC_STEP} + async_training.require_batches=1 + async_training.use_trainer_do_validate=False +) + +############################ Launch ############################ + +echo "Running fully_async_policy + Single-Teacher OPD" +echo "Student: ${STUDENT_MODEL}" +echo "Teacher: ${GSM8K_TEACHER_MODEL}" +echo "Dataset: ${GSM8K_TRAIN}" +echo "Single-turn: prompt=${MAX_PROMPT}, response=${MAX_RESPONSE_LENGTH}, total_tokens=${MAX_NUM_TOKENS}" +echo "MTP/speculative decoding: disabled" +echo "GPUs: ${N_GPUS_ROLLOUT}x${ROLLOUT_NNODES} rollout + ${N_GPUS_TRAINING}x${TRAINER_NNODES} fused teacher/training" +echo "Checkpoints: every ${SAVE_EVERY_TRAIN_STEPS} trainer steps -> ${CHECKPOINT_DIR}" + +python3 -m verl.experimental.fully_async_policy.fully_async_main \ + --config-path=config \ + --config-name='fully_async_ppo_megatron_trainer.yaml' \ + actor_rollout_ref.hybrid_engine=False \ + critic.strategy=megatron \ + "${DATA[@]}" \ + "${MODEL[@]}" \ + "${STUDENT[@]}" \ + "${ROLLOUT[@]}" \ + "${DISTILLATION[@]}" \ + "${ALGORITHM[@]}" \ + "${REWARD[@]}" \ + "${TRAINER[@]}" \ + "${ASYNC_TRAINING[@]}" \ + "${mtp_params[@]}" \ + "$@" + +echo "Fully async fused-node OPD completed successfully" diff --git a/verl/experimental/agent_loop/agent_loop.py b/verl/experimental/agent_loop/agent_loop.py index 9b618f89d2f..a969beaeae7 100644 --- a/verl/experimental/agent_loop/agent_loop.py +++ b/verl/experimental/agent_loop/agent_loop.py @@ -511,7 +511,11 @@ def __init__( # Online policy distillation self.distillation_enabled = is_distillation_enabled(config.distillation) - if self.distillation_enabled: + self.teacher_execution = ( + config.distillation.get("teacher_execution", "rollout") if self.distillation_enabled else "rollout" + ) + self.rollout_teacher_enabled = self.distillation_enabled and self.teacher_execution == "rollout" + if self.rollout_teacher_enabled: from verl.experimental.teacher_loop.teacher_manager import AsyncTeacherLLMServerManager self.teacher_key: str = config.distillation.teacher_key @@ -1006,7 +1010,7 @@ 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: + if self.rollout_teacher_enabled 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/fully_async_policy/fully_async_main.py b/verl/experimental/fully_async_policy/fully_async_main.py index 62a2ddf50a9..d0855a4d46b 100644 --- a/verl/experimental/fully_async_policy/fully_async_main.py +++ b/verl/experimental/fully_async_policy/fully_async_main.py @@ -112,6 +112,10 @@ def _initialize_components(self, config) -> None: if config.trainer.get("val_before_train", True): ray.get(self.components["trainer"]._fit_validate.remote(True)) + # Trainer-colocated OPD keeps the teacher resident while waiting for + # trajectories. Initial actor weight sync and validation must finish first. + ray.get(self.components["trainer"].prepare_fused_teacher.remote()) + print("[ASYNC MAIN] All components initialized successfully") def _create_rollouter(self, config) -> None: diff --git a/verl/experimental/fully_async_policy/fully_async_rollouter.py b/verl/experimental/fully_async_policy/fully_async_rollouter.py index ea578724908..b9987ded35d 100644 --- a/verl/experimental/fully_async_policy/fully_async_rollouter.py +++ b/verl/experimental/fully_async_policy/fully_async_rollouter.py @@ -633,7 +633,11 @@ async def _create_teacher_model_manager(self): from verl.trainer.ppo.utils import Role self.teacher_model_manager = None - if is_distillation_enabled(self.config.get("distillation")): + distillation_enabled = is_distillation_enabled(self.config.get("distillation")) + teacher_execution = ( + self.config.distillation.get("teacher_execution", "rollout") if distillation_enabled else "rollout" + ) + if distillation_enabled and teacher_execution == "rollout": from verl.experimental.teacher_loop import MultiTeacherModelManager resource_pool_spec = {} diff --git a/verl/experimental/fully_async_policy/fully_async_trainer.py b/verl/experimental/fully_async_policy/fully_async_trainer.py index 20ac0330ba9..9ae5409c066 100644 --- a/verl/experimental/fully_async_policy/fully_async_trainer.py +++ b/verl/experimental/fully_async_policy/fully_async_trainer.py @@ -35,10 +35,12 @@ from verl.trainer.ppo import core_algos from verl.trainer.ppo.ray_trainer import ResourcePoolManager from verl.trainer.ppo.utils import Role, WorkerType, need_critic, need_reference_policy, need_reward_model +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.tracking import Tracking +from verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding logger = logging.getLogger(__name__) @@ -87,6 +89,9 @@ def __init__( self.distillation_config = omega_conf_to_dataclass(self.config.distillation) else: self.distillation_config = None + self.fused_teacher_enabled = bool( + self.distillation_config is not None and self.distillation_config.teacher_execution == "trainer" + ) self.use_critic = need_critic(self.config) self.ray_worker_group_cls = ray_worker_group_cls @@ -404,6 +409,7 @@ async def fit(self): break self.progress_bar.close() + self._activate_fused_actor() if self.current_param_version % self.config.trainer.test_freq != 0 or self.local_trigger_step > 1: weights_updated = await self._fit_update_weights() if weights_updated: @@ -436,6 +442,7 @@ async def fit_step(self, batch_dict: dict = None): with marked_timer("step", self.timing_raw): batch = await self._fit_generate(None) + batch = self._fit_compute_teacher_log_prob(batch) batch = self._fit_compute_reward(batch) batch = self._fit_compute_log_prob(batch) batch = self._fit_compute_ref_log_prob(batch) @@ -454,6 +461,7 @@ async def fit_step(self, batch_dict: dict = None): if weights_updated: self._fit_log_aggregated_training_metrics() self._fit_postprocess_step() + self._activate_fused_teacher() async def _fit_generate(self, batch: DataProto = None) -> DataProto | None: metrics = self.metrics @@ -466,6 +474,44 @@ async def _fit_generate(self, batch: DataProto = None) -> DataProto | None: batch.meta_info["temperature"] = self.config.actor_rollout_ref.rollout.temperature return batch + def _activate_fused_actor(self): + if self.fused_teacher_enabled: + self.actor_rollout_wg.activate_actor() + + def _activate_fused_teacher(self): + if self.fused_teacher_enabled: + self.actor_rollout_wg.activate_teacher() + + def prepare_fused_teacher(self): + """Place the fused node in teacher-resident state before queue consumption.""" + self._activate_fused_teacher() + + def _fit_compute_teacher_log_prob(self, batch: DataProto) -> DataProto: + if not self.fused_teacher_enabled: + return batch + + timing_raw = self.timing_raw + try: + with marked_timer("teacher_log_prob", timing_raw, color="olive"): + batch_td = left_right_2_no_padding(batch.to_tensordict()) + tu.assign_non_tensor( + batch_td, + calculate_entropy=False, + compute_loss=False, + ) + output = self.actor_rollout_wg.compute_teacher_log_prob(batch_td) + teacher_log_probs = tu.get(output, "log_probs") + teacher_log_probs = no_padding_2_padding(teacher_log_probs, batch_td) + # Preserve the existing [batch, response, K] contract with K=1. + batch.batch["teacher_logprobs"] = teacher_log_probs.float().unsqueeze(-1) + teacher_metrics = tu.get(output, "metrics") + if teacher_metrics and "mfu" in teacher_metrics: + self.metrics["perf/mfu/teacher_infer"] = teacher_metrics["mfu"] + finally: + # Optimizer, checkpoint, and weight-sync APIs all expect student weights. + self._activate_fused_actor() + return batch + def _compute_old_log_prob(self, batch: DataProto): """ If algorithm.rollout_correction.bypass_mode is False, @@ -477,15 +523,22 @@ def _compute_old_log_prob(self, batch: DataProto): If local_trigger_step == 2, 3, ..., restore the parameters of version 1 to calculate the old_log_prob, then restore the parameters of the current version. """ + # State 1 is the teacher in fused mode; reserve 2 for the fixed old + # student and 3 for a temporarily displaced current student. + old_student_state = 2 if self.fused_teacher_enabled else 1 + current_student_state = 3 if self.fused_teacher_enabled else self.local_trigger_step + if self.local_trigger_step == 1: - self.actor_rollout_wg.save_model_to_cpu(1) + self.actor_rollout_wg.save_model_to_cpu(old_student_state) old_log_prob, old_log_prob_mfu = super()._compute_old_log_prob(batch) else: - self.actor_rollout_wg.save_model_to_cpu(self.local_trigger_step) - self.actor_rollout_wg.restore_model_from_cpu(1) - old_log_prob, old_log_prob_mfu = super()._compute_old_log_prob(batch) - self.actor_rollout_wg.restore_model_from_cpu(self.local_trigger_step) - self.actor_rollout_wg.clear_cpu_model(self.local_trigger_step) + self.actor_rollout_wg.save_model_to_cpu(current_student_state) + try: + self.actor_rollout_wg.restore_model_from_cpu(old_student_state) + old_log_prob, old_log_prob_mfu = super()._compute_old_log_prob(batch) + finally: + self.actor_rollout_wg.restore_model_from_cpu(current_student_state) + self.actor_rollout_wg.clear_cpu_model(current_student_state) return old_log_prob, old_log_prob_mfu def _fit_update_local_step(self): diff --git a/verl/experimental/separation/engine_workers.py b/verl/experimental/separation/engine_workers.py index e9571da51f0..89fb5e9128c 100644 --- a/verl/experimental/separation/engine_workers.py +++ b/verl/experimental/separation/engine_workers.py @@ -17,12 +17,15 @@ import os from typing import Callable, Optional -from omegaconf import DictConfig +from omegaconf import DictConfig, OmegaConf -from verl.single_controller.base.decorator import Dispatch, register +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.config import omega_conf_to_dataclass from verl.utils.device import ( get_device_name, ) +from verl.workers.config import HFModelConfig, MtpConfig from verl.workers.engine_workers import ActorRolloutRefWorker, DistillationConfig logger = logging.getLogger(__file__) @@ -42,6 +45,10 @@ class DetachActorWorker(ActorRolloutRefWorker): FSDP, FSDP2, VeOmni, and Megatron strategies. """ + FUSED_TEACHER_STATE = 1 + FUSED_OLD_STUDENT_STATE = 2 + FUSED_CURRENT_STUDENT_STATE = 3 + def __init__( self, config: DictConfig, role: str, distillation_config: Optional[DistillationConfig] = None, **kwargs ): @@ -56,6 +63,128 @@ def __init__( """ ActorRolloutRefWorker.__init__(self, config, role, distillation_config=distillation_config, **kwargs) self._strategy_handlers = None + self._fused_teacher_enabled = bool( + self.distillation_enabled + and distillation_config is not None + and distillation_config.get("teacher_execution", "rollout") == "trainer" + ) + self._active_model = "actor" + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def init_model(self): + super().init_model() + if self._fused_teacher_enabled: + self._init_fused_teacher() + + def _init_fused_teacher(self): + if self.config.actor.strategy != "megatron": + raise NotImplementedError("Trainer-colocated teachers currently require actor.strategy=megatron.") + if self.role not in {"actor", "actor_rollout", "actor_rollout_ref"}: + raise ValueError(f"Trainer-colocated teacher requires an actor role, got {self.role!r}.") + + distillation_config = omega_conf_to_dataclass(self.distillation_config) + if len(distillation_config.teacher_models) != 1: + raise NotImplementedError("Trainer-colocated distillation currently supports exactly one teacher.") + teacher_config = next(iter(distillation_config.teacher_models.values())) + + teacher_model_dict = { + "_target_": "verl.workers.config.HFModelConfig", + "path": teacher_config.model_path, + "hf_config_path": teacher_config.model_path, + # Token-level OPD requires a shared token-id space. Load the + # student's tokenizer while loading the teacher architecture. + "tokenizer_path": self.config.model.get("tokenizer_path") or self.config.model.path, + "load_tokenizer": True, + "use_shm": self.config.model.get("use_shm", False), + "trust_remote_code": self.config.model.get("trust_remote_code", False), + "external_lib": self.config.model.get("external_lib", None), + "enable_gradient_checkpointing": False, + "use_remove_padding": self.config.model.get("use_remove_padding", True), + "mtp": OmegaConf.to_container( + OmegaConf.structured(MtpConfig(enable=False, enable_train=False, enable_rollout=False)), + resolve=True, + ), + } + teacher_model_config: HFModelConfig = omega_conf_to_dataclass( + teacher_model_dict, dataclass_type=HFModelConfig + ) + + student_model_config = self.actor.model_config + if student_model_config.mtp.enable: + raise NotImplementedError("Single-engine fused OPD currently requires actor MTP to be disabled.") + if student_model_config.architectures != teacher_model_config.architectures: + raise ValueError( + "Single-engine fused OPD requires identical model architectures, " + f"got student={student_model_config.architectures} and " + f"teacher={teacher_model_config.architectures}." + ) + + structural_fields = ( + "model_type", + "vocab_size", + "hidden_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "intermediate_size", + "num_experts", + "num_experts_per_tok", + ) + student_hf_config = getattr(student_model_config.hf_config, "text_config", student_model_config.hf_config) + teacher_hf_config = getattr(teacher_model_config.hf_config, "text_config", teacher_model_config.hf_config) + mismatches = { + field: (getattr(student_hf_config, field, None), getattr(teacher_hf_config, field, None)) + for field in structural_fields + if getattr(student_hf_config, field, None) != getattr(teacher_hf_config, field, None) + } + if mismatches: + raise ValueError(f"Single-engine fused OPD requires isomorphic teacher/student models: {mismatches}") + + # Build the static teacher snapshot in the initialized student module. + # Optimizer/main-parameter state remains untouched and resident. + self.actor.to(device="device", model=True, optimizer=False, grad=False) + self.save_model_to_cpu(self.FUSED_CURRENT_STUDENT_STATE) + try: + engine = self.actor.engine + if engine.vanilla_bridge: + engine.bridge.load_weights(engine.module, teacher_model_config.local_path) + else: + engine.bridge.load_hf_weights(engine.module, teacher_model_config.local_path) + self.save_model_to_cpu(self.FUSED_TEACHER_STATE) + finally: + self.restore_model_from_cpu(self.FUSED_CURRENT_STUDENT_STATE) + self.clear_cpu_model(self.FUSED_CURRENT_STUDENT_STATE) + self._active_model = "actor" + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def activate_teacher(self): + if not self._fused_teacher_enabled: + return + if self._active_model == "teacher": + return + self.save_model_to_cpu(self.FUSED_CURRENT_STUDENT_STATE) + self.restore_model_from_cpu(self.FUSED_TEACHER_STATE) + self._active_model = "teacher" + + @register(dispatch_mode=Dispatch.ONE_TO_ALL) + def activate_actor(self): + if not self._fused_teacher_enabled: + return + if self._active_model == "actor": + return + self.restore_model_from_cpu(self.FUSED_CURRENT_STUDENT_STATE) + self.clear_cpu_model(self.FUSED_CURRENT_STUDENT_STATE) + self._active_model = "actor" + + @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor")) + def compute_teacher_log_prob(self, data): + if not self._fused_teacher_enabled: + raise RuntimeError("Trainer-colocated teacher is not initialized.") + if self._active_model != "teacher": + raise RuntimeError(f"Teacher scoring requires active_model='teacher', got {self._active_model!r}.") + tu.assign_non_tensor(data, disable_auto_offload=True, calculate_entropy=False, compute_loss=False) + output = self.actor.infer_batch(data) + return output.cpu() if output is not None else None def _get_strategy_handlers(self): """ diff --git a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml index fb2d87a9451..bc0a1b56ae7 100644 --- a/verl/trainer/config/_generated_ppo_megatron_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_megatron_trainer.yaml @@ -850,6 +850,7 @@ algorithm: distillation: _target_: verl.workers.config.DistillationConfig enabled: false + teacher_execution: rollout distillation_loss: _target_: verl.workers.config.DistillationLossConfig loss_mode: k3 diff --git a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml index 848ec0d6543..ad401652114 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -769,6 +769,7 @@ algorithm: distillation: _target_: verl.workers.config.DistillationConfig enabled: false + teacher_execution: rollout distillation_loss: _target_: verl.workers.config.DistillationLossConfig loss_mode: k3 diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index f91eb3a41f9..ce841c026d1 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -791,6 +791,7 @@ algorithm: distillation: _target_: verl.workers.config.DistillationConfig enabled: false + teacher_execution: rollout distillation_loss: _target_: verl.workers.config.DistillationLossConfig loss_mode: k3 diff --git a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml index dd93ec973f6..fac9298f5ee 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -789,6 +789,7 @@ algorithm: distillation: _target_: verl.workers.config.DistillationConfig enabled: false + teacher_execution: rollout distillation_loss: _target_: verl.workers.config.DistillationLossConfig loss_mode: k3 diff --git a/verl/trainer/config/distillation/distillation.yaml b/verl/trainer/config/distillation/distillation.yaml index 84ee7932ab6..6f4e339f586 100644 --- a/verl/trainer/config/distillation/distillation.yaml +++ b/verl/trainer/config/distillation/distillation.yaml @@ -13,6 +13,11 @@ defaults: # Whether to enable distillation. enabled: false +# Where teacher log probabilities are computed: +# - rollout: dedicated vLLM/SGLang teacher resource pool +# - trainer: isomorphic teacher/student snapshots share the trainer Megatron engine +teacher_execution: rollout + # distillation loss config distillation_loss: @@ -110,4 +115,4 @@ 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 diff --git a/verl/trainer/distillation/losses.py b/verl/trainer/distillation/losses.py index 3467e59ebaa..d15f7f0e984 100644 --- a/verl/trainer/distillation/losses.py +++ b/verl/trainer/distillation/losses.py @@ -381,7 +381,13 @@ def compute_distillation_loss_reverse_kl_estimator( - distillation_metrics: Dictionary of metrics. """ student_log_probs = no_padding_2_padding(model_output["log_probs"], data) - teacher_log_probs = no_padding_2_padding(data["teacher_logprobs"], data).squeeze(-1) + teacher_log_probs = data["teacher_logprobs"] + if teacher_log_probs.is_nested: + teacher_log_probs = no_padding_2_padding(teacher_log_probs, data) + # Dedicated rollout teachers provide full-sequence nested values, while + # trainer-colocated k1 scoring injects response-aligned padded values. + if teacher_log_probs.ndim == 3 and teacher_log_probs.shape[-1] == 1: + teacher_log_probs = teacher_log_probs.squeeze(-1) if data["response_mask"].is_nested: response_mask_bool = data["response_mask"].bool().to_padded_tensor(False) else: diff --git a/verl/workers/config/distillation.py b/verl/workers/config/distillation.py index b58f8bccffd..6c120d828bf 100644 --- a/verl/workers/config/distillation.py +++ b/verl/workers/config/distillation.py @@ -257,9 +257,18 @@ class DistillationConfig(BaseConfig): ``` """ - _mutable_fields = BaseConfig._mutable_fields | {"teacher_models", "n_gpus_per_node", "nnodes"} + _mutable_fields = BaseConfig._mutable_fields | { + "teacher_models", + "n_gpus_per_node", + "nnodes", + "teacher_execution", + } enabled: bool = False + # "rollout": score trajectories on dedicated vLLM/SGLang teacher resources. + # "trainer": score trajectories on the trainer's Megatron engine by swapping + # isomorphic teacher/student parameter snapshots between phases. + teacher_execution: str = "rollout" n_gpus_per_node: int = 0 nnodes: int = 0 teacher_models: dict[str, DistillationTeacherModelConfig] = field(default_factory=dict) @@ -270,7 +279,25 @@ def __post_init__(self): if not self.enabled: return + if self.teacher_execution not in {"rollout", "trainer"}: + raise ValueError( + "distillation.teacher_execution must be either 'rollout' or 'trainer', " + f"got {self.teacher_execution!r}." + ) + self.teacher_models = self._resolve_teacher_models() + if self.teacher_execution == "trainer": + if len(self.teacher_models) != 1: + raise NotImplementedError( + "Trainer-colocated distillation currently supports exactly one teacher model." + ) + if self.distillation_loss.loss_settings.use_topk: + raise NotImplementedError( + "Trainer-colocated distillation currently supports sampled-token losses " + "(for example k1), but not forward_kl_topk." + ) + return + teacher_world_size_sum = 0 for teacher_model in self.teacher_models.values(): teacher_model.validate_and_prepare_for_distillation( @@ -289,21 +316,24 @@ def __post_init__(self): def _resolve_teacher_models(self) -> dict[str, DistillationTeacherModelConfig]: assert "teacher_model" in self.teacher_models if len(self.teacher_models) == 1: - # Single teacher occupies the entire teacher resource pool. teacher_model = self.teacher_models["teacher_model"] - inference = teacher_model.inference - per_replica = ( - inference.tensor_model_parallel_size - * inference.data_parallel_size - * inference.pipeline_model_parallel_size - ) - pool_size = self.n_gpus_per_node * self.nnodes - if pool_size % per_replica != 0: - raise ValueError( - f"Single teacher's per_replica_world_size ({per_replica}) must divide the distillation " - f"resource pool size ({self.n_gpus_per_node=} * {self.nnodes=} = {pool_size})." + if self.teacher_execution == "rollout": + # Single teacher occupies the entire teacher resource pool. + 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 + else: + teacher_model.num_replicas = 0 teacher_model.key = "default" else: # Multiple teachers: remove default single teacher config