diff --git a/examples/dpo_trainer/README.md b/examples/dpo_trainer/README.md index 3b1cf9681..9002b11ae 100644 --- a/examples/dpo_trainer/README.md +++ b/examples/dpo_trainer/README.md @@ -1,6 +1,6 @@ # DPO Training -Last updated: 08/14/2026 +Last updated: 09/03/2026 This directory contains examples for **direct-preference** training (DPO and related losses). Three workflows are supported: @@ -47,6 +47,14 @@ bash examples/dpo_trainer/qwen_image/run_qwen_image_online_dpo_lora.sh \ data.val_files=$WORKSPACE/data/ocr/qwen_image/test.parquet ``` +For CUDA V1 sync (TransferQueue + ReplayBuffer), use `examples/dpo_trainer/qwen_image/run_qwen_image_online_dpo_lora_v1.sh`. + +```bash +bash examples/dpo_trainer/qwen_image/run_qwen_image_online_dpo_lora_v1.sh \ + data.train_files=$WORKSPACE/data/ocr/qwen_image/train.parquet \ + data.val_files=$WORKSPACE/data/ocr/qwen_image/test.parquet +``` + #### NPU For Huawei Ascend NPUs, use the NPU-optimized script: diff --git a/examples/dpo_trainer/qwen_image/run_qwen_image_online_dpo_lora_v1.sh b/examples/dpo_trainer/qwen_image/run_qwen_image_online_dpo_lora_v1.sh new file mode 100644 index 000000000..e1ac03842 --- /dev/null +++ b/examples/dpo_trainer/qwen_image/run_qwen_image_online_dpo_lora_v1.sh @@ -0,0 +1,88 @@ +#!/usr/bin/env bash +# Qwen-Image online DPO LoRA (V1 trainer: TransferQueue + ReplayBuffer + sync mode). +# +# This is the v1 counterpart of run_qwen_image_online_dpo_lora.sh. It uses the +# new `verl_omni.trainer.main_diffusion_v1` entrypoint, which selects +# `PolicyGradientDiffusionTrainerV1Sync` via `trainer.v1.trainer_mode=sync` and +# wires verl's `AgentLoopManagerTQ` with `DiffusionAgentLoopWorkerTQ`. +# TransferQueue is force-enabled inside the runner, so it does not need to be +# set on the CLI. +# +# Reference (legacy v0 script): +# verl-omni/examples/dpo_trainer/qwen_image/run_qwen_image_online_dpo_lora.sh +set -x + +# Set WORKSPACE to any writable directory; defaults to $HOME. +WORKSPACE=${WORKSPACE:-$HOME} + +ocr_train_path=$WORKSPACE/data/ocr/qwen_image/train.parquet +ocr_test_path=$WORKSPACE/data/ocr/qwen_image/test.parquet + +model_name=Qwen/Qwen-Image +reward_model_name=Qwen/Qwen3-VL-8B-Instruct +reward_function_path=verl_omni/utils/reward_score/genrm_ocr.py + +NUM_GPUS_ACTOR_ROLLOUT_REWARD=4 +ROLLOUT_TP=1 +REWARD_TP=4 + +ENGINE=vllm_omni +REWARD_ENGINE=vllm + +python3 -m verl_omni.trainer.main_diffusion_v1 \ + algorithm.trainer_type=direct_preference \ + algorithm.sample_source=online \ + algorithm.paired_preference=true \ + data.train_files=$ocr_train_path \ + data.val_files=$ocr_test_path \ + data.train_batch_size=32 \ + data.max_prompt_length=256 \ + actor_rollout_ref.model.path=$model_name \ + actor_rollout_ref.model.algorithm=dpo \ + actor_rollout_ref.model.model_type=diffusion_dpo_model \ + actor_rollout_ref.model.external_lib=verl_omni.pipelines.qwen_image_dpo \ + actor_rollout_ref.model.lora_rank=64 \ + actor_rollout_ref.model.lora_alpha=128 \ + actor_rollout_ref.model.target_modules="['to_q','to_k','to_v','to_out.0','add_q_proj','add_k_proj','add_v_proj','to_add_out','img_mlp.net.0.proj','img_mlp.net.2','txt_mlp.net.0.proj','txt_mlp.net.2']" \ + actor_rollout_ref.actor.diffusion_loss.loss_mode=dpo \ + actor_rollout_ref.actor.diffusion_loss.dpo_beta=100.0 \ + actor_rollout_ref.actor.optim.lr=3e-4 \ + actor_rollout_ref.actor.optim.weight_decay=0.0001 \ + actor_rollout_ref.actor.ppo_mini_batch_size=16 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \ + actor_rollout_ref.actor.fsdp_config.param_offload=True \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.rollout.name=$ENGINE \ + actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP \ + actor_rollout_ref.rollout.n=16 \ + actor_rollout_ref.rollout.calculate_log_probs=false \ + actor_rollout_ref.rollout.agent.num_workers=$((NUM_GPUS_ACTOR_ROLLOUT_REWARD / ROLLOUT_TP)) \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=True \ + actor_rollout_ref.rollout.pipeline.num_inference_steps=35 \ + actor_rollout_ref.rollout.pipeline.true_cfg_scale=1.0 \ + actor_rollout_ref.rollout.pipeline.max_sequence_length=256 \ + actor_rollout_ref.rollout.val_kwargs.pipeline.num_inference_steps=50 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \ + reward.num_workers=$((NUM_GPUS_ACTOR_ROLLOUT_REWARD / REWARD_TP)) \ + reward.reward_model.enable=True \ + reward.reward_model.model_path=$reward_model_name \ + reward.reward_model.rollout.name=$REWARD_ENGINE \ + reward.reward_model.rollout.tensor_model_parallel_size=$REWARD_TP \ + reward.reward_model.rollout.enforce_eager=False \ + reward.custom_reward_function.path=$reward_function_path \ + reward.custom_reward_function.name=compute_score_ocr \ + trainer.logger='["console", "wandb"]' \ + trainer.project_name=online_dpo \ + trainer.experiment_name=qwen_image_online_dpo_lora_v1 \ + trainer.log_val_generations=8 \ + trainer.val_before_train=False \ + trainer.n_gpus_per_node=$NUM_GPUS_ACTOR_ROLLOUT_REWARD \ + trainer.nnodes=1 \ + trainer.save_freq=20 \ + trainer.test_freq=20 \ + trainer.total_epochs=15 \ + trainer.total_training_steps=300 \ + trainer.use_v1=true \ + trainer.v1.trainer_mode=sync "$@" diff --git a/tests/trainer/diffusion/test_v1_direct_preference_on_cpu.py b/tests/trainer/diffusion/test_v1_direct_preference_on_cpu.py new file mode 100644 index 000000000..8c09062e1 --- /dev/null +++ b/tests/trainer/diffusion/test_v1_direct_preference_on_cpu.py @@ -0,0 +1,199 @@ +# 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. +"""CPU tests for online DPO on the v1 diffusion trainer. + +Necessity: the v1 PG loop recomputes ``old_log_probs`` via ``infer_actor_batch``. +DPO engines return ``noise_pred`` (log_probs=None), which crashed +``run_qwen_image_online_dpo_lora_v1.sh``. These tests lock the direct-preference +branch that pairs rewards and uses ref noise preds instead. +""" + +import os +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch +from hydra import compose, initialize_config_dir +from transfer_queue import KVBatchMeta +from verl import DataProto +from verl.utils import tensordict_utils as tu + +import verl_omni + +CONFIG_DIR = os.path.join(os.path.dirname(os.path.abspath(verl_omni.__file__)), "trainer", "config") + +DPO_OVERRIDES = [ + "algorithm.trainer_type=direct_preference", + "algorithm.sample_source=online", + "algorithm.paired_preference=true", + "actor_rollout_ref.actor.diffusion_loss.loss_mode=dpo", +] + + +def compose_cfg(overrides): + with initialize_config_dir(config_dir=CONFIG_DIR, version_base=None): + return compose(config_name="diffusion_trainer", overrides=overrides) + + +def make_dpo_trainer(overrides=None): + from verl_omni.trainer.diffusion.v1.trainer_sync import PolicyGradientDiffusionTrainerV1Sync + + return PolicyGradientDiffusionTrainerV1Sync(compose_cfg(DPO_OVERRIDES + list(overrides or []))) + + +def test_offline_direct_preference_rejected_on_v1(): + from verl_omni.trainer.diffusion.v1.trainer_sync import PolicyGradientDiffusionTrainerV1Sync + + with pytest.raises(NotImplementedError, match="offline DPO stays on the v0 trainer"): + PolicyGradientDiffusionTrainerV1Sync( + compose_cfg( + [ + "algorithm.trainer_type=direct_preference", + "algorithm.sample_source=offline", + "actor_rollout_ref.actor.diffusion_loss.loss_mode=dpo", + ] + ) + ) + + +def test_dpo_enables_reference_policy_without_kl(): + trainer = make_dpo_trainer() + assert trainer._is_direct_preference + assert trainer.use_reference_policy + assert trainer._has_old_adapter is False + + +def test_dpo_train_step_skips_old_log_prob_and_pairs_batch(monkeypatch): + trainer = make_dpo_trainer() + trainer.tokenizer = SimpleNamespace(pad_token_id=0) + trainer.reward_loop_manager = SimpleNamespace(reward_loop_worker_handles=object()) + trainer.global_steps = 1 + + uid = np.array(["p0", "p0", "p1", "p1"], dtype=object) + data = DataProto.from_dict( + tensors={ + "rm_scores": torch.tensor([[1.0], [0.0], [0.2], [0.8]]), + "latents_clean": torch.zeros(4, 2, 4, 4), + }, + non_tensors={"uid": uid}, + ) + monkeypatch.setattr( + "verl_omni.trainer.diffusion.v1.trainer_base.diffusion_tq_batch_to_dataproto", + lambda meta, pad_token_id: data, + ) + monkeypatch.setattr( + "verl_omni.trainer.diffusion.v1.trainer_base.extract_reward", + lambda batch: (batch.batch["rm_scores"], {}), + ) + tq_writes: list[list[str]] = [] + + def capture_tq(batch_meta, batch, fields): + del batch_meta, batch + tq_writes.append(list(fields)) + + monkeypatch.setattr( + "verl_omni.trainer.diffusion.v1.trainer_base.put_dataproto_fields_to_tq", + capture_tq, + ) + + def fail_old_log_prob(_data): + raise AssertionError("DPO must not recompute old_log_probs") + + monkeypatch.setattr(trainer, "_compute_old_log_prob", fail_old_log_prob) + + def fail_balance(d, metrics): + raise AssertionError("DPO must not DP-pad before pairing") + + monkeypatch.setattr(trainer, "_balance_batch", fail_balance) + + captured = {} + + def capture_ref(batch): + captured["ref"] = batch + return DataProto.from_tensordict(tu.get_tensordict({"ref_noise_pred": torch.ones(len(batch), 2, 4, 4)})) + + def capture_update(batch): + captured["update"] = batch + return DataProto.from_single_dict(data={}, meta_info={"metrics": {"actor/dpo_loss": 0.1}}) + + monkeypatch.setattr(trainer, "_compute_ref_noise_pred", capture_ref) + monkeypatch.setattr(trainer, "_update_actor", capture_update) + + batch_meta = KVBatchMeta( + partition_id="train", + keys=["p0_0_0", "p0_1_0", "p1_0_0", "p1_1_0"], + tags=[{"is_padding": False}] * 4, + ) + result = trainer._train_sampled_batch({}, {}, batch_meta) + + assert result is batch_meta + assert "old_log_probs" not in tq_writes[0] + assert "sample_level_scores" in tq_writes[0] + paired = captured["update"] + assert len(paired) == 4 + assert list(paired.non_tensor_batch["uid"]) == ["p0", "p0", "p1", "p1"] + scores = paired.batch["sample_level_scores"].reshape(-1) + assert scores[0] >= scores[1] + assert scores[2] >= scores[3] + assert "ref_noise_pred" in captured["update"].batch + assert "old_log_probs" not in captured["update"].batch + + +def test_dpo_update_actor_uses_paired_mini_batch_size(): + from verl_omni.trainer.diffusion.v1.trainer_base import PolicyGradientDiffusionTrainerV1 + + actor = MagicMock() + actor.update_actor.return_value = tu.get_tensordict({}, non_tensor_dict={"metrics": {}}) + trainer = SimpleNamespace( + config=compose_cfg( + DPO_OVERRIDES + + [ + "actor_rollout_ref.actor.ppo_mini_batch_size=2", + "actor_rollout_ref.actor.ppo_epochs=1", + "actor_rollout_ref.actor.data_loader_seed=0", + "actor_rollout_ref.actor.shuffle=true", + "actor_rollout_ref.rollout.n=16", + ] + ), + _is_direct_preference=True, + actor_rollout_wg=actor, + ) + batch = DataProto.from_dict(tensors={"latents_clean": torch.zeros(4, 2, 2, 2)}) + + PolicyGradientDiffusionTrainerV1._update_actor(trainer, batch) + + sent = actor.update_actor.call_args.args[0] + assert tu.get_non_tensor_data(sent, "mini_batch_size", None) == 4 + assert tu.get_non_tensor_data(sent, "global_batch_size", None) == 4 + assert tu.get_non_tensor_data(sent, "dataloader_kwargs", {})["shuffle"] is False + + +def test_compute_old_log_prob_fails_closed_when_log_probs_missing(): + from verl_omni.trainer.diffusion.v1.trainer_base import PolicyGradientDiffusionTrainerV1 + + actor = MagicMock() + actor.infer_actor_batch.return_value = tu.get_tensordict( + {"noise_pred": torch.zeros(2, 1)}, + non_tensor_dict={"metrics": {}}, + ) + trainer = SimpleNamespace( + config=compose_cfg([]), + actor_rollout_wg=actor, + ) + batch = DataProto.from_dict(tensors={"all_latents": torch.zeros(2, 1, 2, 2, 2)}) + + with pytest.raises(RuntimeError, match="log_probs=None"): + PolicyGradientDiffusionTrainerV1._compute_old_log_prob(trainer, batch) diff --git a/tests/trainer/diffusion/test_worker_batch_projection_on_cpu.py b/tests/trainer/diffusion/test_worker_batch_projection_on_cpu.py index a12752386..b1ad4231b 100644 --- a/tests/trainer/diffusion/test_worker_batch_projection_on_cpu.py +++ b/tests/trainer/diffusion/test_worker_batch_projection_on_cpu.py @@ -152,6 +152,7 @@ def test_teacher_manager_hop_excludes_responses(): [ ("_compute_old_log_prob", "actor_rollout_wg", "infer_actor_batch"), ("_compute_ref_log_prob", "ref_policy_wg", "infer_ref_batch"), + ("_compute_ref_noise_pred", "ref_policy_wg", "infer_ref_batch"), ("_update_actor", "actor_rollout_wg", "update_actor"), ], ) diff --git a/verl_omni/trainer/diffusion/v1/trainer_base.py b/verl_omni/trainer/diffusion/v1/trainer_base.py index 3f86ad6f9..513342af0 100644 --- a/verl_omni/trainer/diffusion/v1/trainer_base.py +++ b/verl_omni/trainer/diffusion/v1/trainer_base.py @@ -58,13 +58,16 @@ from verl.utils.tracking import Tracking, ValidationGenerationsLogger from verl.workers.rollout.llm_server import LLMServerManager +from verl_omni.trainer.diffusion.diffusion_algos import get_diffusion_loss_fn from verl_omni.trainer.diffusion.diffusion_metric_utils import ( compute_data_metrics_diffusion, + compute_old_policy_metrics, compute_reward_extra_metrics_diffusion, compute_throughput_metrics_diffusion, compute_timing_metrics_diffusion, ) from verl_omni.trainer.diffusion.diffusion_trainer_utils import ( + old_policy_decay, validate_distillation_config, worker_group_port_ranges, ) @@ -78,6 +81,7 @@ from verl_omni.trainer.diffusion.teacher_manager import DiffusionTeacherManager from verl_omni.trainer.diffusion.v1.tq_utils import ( diffusion_tq_batch_to_dataproto, + put_dataproto_fields_to_tq, sort_diffusion_tq_keys, ) from verl_omni.workers.engine_workers import ActorRolloutRefWorker, resolve_teacher_infer_micro_batch_size @@ -122,7 +126,25 @@ def __init__(self, config): self.config = config self.trainer_mode = config.trainer.v1.trainer_mode self.parameter_sync_step = config.trainer.v1.get(self.trainer_mode, {}).get("parameter_sync_step", 1) - self.use_reference_policy = need_reference_policy(config) + loss_mode = config.actor_rollout_ref.actor.diffusion_loss.loss_mode + self._is_direct_preference = config.algorithm.get("trainer_type", "policy_gradient") == "direct_preference" + if self._is_direct_preference: + if config.algorithm.get("sample_source", "online") == "offline": + raise NotImplementedError( + "Diffusion offline DPO stays on the v0 trainer. Use " + "`python -m verl_omni.trainer.main_diffusion` with trainer.use_v1=false." + ) + self._loss_fn = get_diffusion_loss_fn(loss_mode) + self._has_old_adapter = "old" in tuple( + config.actor_rollout_ref.model.get("policy_state_adapters", ("default",)) + ) + if self._has_old_adapter: + self._validate_old_adapter_config() + else: + self._loss_fn = None + self._has_old_adapter = False + # DPO needs trainer-side ref noise preds even when KL is disabled. + self.use_reference_policy = need_reference_policy(config) or (loss_mode == "dpo") self.use_rm = need_reward_model(config) self.use_teacher_policy = is_distillation_enabled(config.get("distillation")) self.distillation_config = omega_conf_to_dataclass(config.distillation) if self.use_teacher_policy else None @@ -164,6 +186,8 @@ def _build_replay_buffer(self) -> ReplayBuffer: def init(self): """Initialize workers, rollout server, reward loop, checkpoint engine.""" self._setup() + if self._has_old_adapter: + self.actor_rollout_wg.copy_adapter(source="default", target="old") self.on_init_end() def fit(self, agent_loop_manager: AgentLoopManager): @@ -327,6 +351,9 @@ def _train_sampled_batch(self, metrics: dict, timing_raw: dict, batch_meta: KVBa data = data.union(self._compute_reward_colocate(data)) self.checkpoint_manager.update_weights(self.global_steps) + if self._is_direct_preference: + return self._train_direct_preference_batch(metrics, timing_raw, batch_meta, data) + data = self._balance_batch(data, metrics=metrics) # Bypass mode: skip old_log_prob recompute (2 policies). @@ -374,8 +401,6 @@ def _train_sampled_batch(self, metrics: dict, timing_raw: dict, batch_meta: KVBa # Persist computed fields back to TransferQueue so the sampled keys carry # the full trajectory for metrics/dumping (keys are cleared after step). # Slice to the original key count in case ``_balance_batch`` appended pad rows. - from verl_omni.trainer.diffusion.v1.tq_utils import put_dataproto_fields_to_tq - n_keys = len(batch_meta.keys) if len(data) > n_keys: data_for_tq = data.select_idxs(list(range(n_keys))) @@ -388,6 +413,121 @@ def _train_sampled_batch(self, metrics: dict, timing_raw: dict, batch_meta: KVBa ) return batch_meta + def _validate_old_adapter_config(self): + """Require the NFT old-adapter rollout/loss contract.""" + rollout_cfg = self.config.actor_rollout_ref.rollout + actor_loss_cfg = self.config.actor_rollout_ref.actor.diffusion_loss + if rollout_cfg.rollout_adapter != "old": + raise ValueError("Old-adapter algorithms require actor_rollout_ref.rollout.rollout_adapter=old.") + if actor_loss_cfg.loss_mode != "diffusion_nft": + raise ValueError( + "Old-adapter algorithms require actor_rollout_ref.actor.diffusion_loss.loss_mode=diffusion_nft." + ) + + def _prepare_actor_batch(self, batch: DataProto, reward_tensor: torch.Tensor) -> DataProto: + """Delegate algorithm-specific rollout-to-actor batch preparation.""" + reward_tensor = reward_tensor.squeeze(-1).float() if reward_tensor.ndim > 1 else reward_tensor.float() + return self._loss_fn.prepare_actor_batch(batch, reward_tensor, self.config) + + def _compute_ref_noise_pred(self, data: DataProto) -> DataProto | None: + """Reference transformer output and shared flow tensors for DPO.""" + batch_td = _to_diffusion_worker_tensordict(data) + batch_td = embeds_padding_2_no_padding(batch_td) + metadata = { + "compute_loss": False, + "height": self.config.actor_rollout_ref.model.pipeline.height, + "width": self.config.actor_rollout_ref.model.pipeline.width, + "vae_scale_factor": self.config.actor_rollout_ref.model.get("vae_scale_factor", 8), + } + 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.infer_actor_batch(batch_td) + else: + output = self.ref_policy_wg.infer_ref_batch(batch_td) + if output is None: + return None + + noise_pred = tu.get(output, "noise_pred") + if noise_pred is None: + raise RuntimeError( + "Reference infer returned noise_pred=None. Diffusion DPO requires " + "model_type=diffusion_dpo_model so infer_actor_batch / infer_ref_batch " + "emit noise_pred rather than SDE log_probs." + ) + if noise_pred.ndim >= 2 and noise_pred.shape[1] == 1: + noise_pred = noise_pred[:, 0] + noise = tu.get(output, "noise") + if noise.ndim >= 2 and noise.shape[1] == 1: + noise = noise[:, 0] + timesteps = tu.get(output, "timesteps") + if timesteps.ndim >= 2 and timesteps.shape[1] == 1: + timesteps = timesteps[:, 0] + ref_output = { + "ref_noise_pred": noise_pred.float(), + "noise": noise.float(), + "timesteps": timesteps.float(), + } + return DataProto.from_tensordict(tu.get_tensordict(ref_output)) + + def _update_old_policy(self) -> tuple[bool, float, str]: + """Refresh the NFT old-policy adapter (copy or EMA).""" + algo_cfg = self.config.algorithm + if self.global_steps % algo_cfg.old_policy_update_interval != 0: + return False, 0.0, "none" + + decay = algo_cfg.old_policy_decay + if decay is None: + decay = old_policy_decay(self.global_steps, algo_cfg.old_policy_decay_schedule) + + if decay == 0: + self.actor_rollout_wg.copy_adapter(source="default", target="old") + return True, float(decay), "copy" + self.actor_rollout_wg.ema_update_adapter(source="default", target="old", decay=decay) + return True, float(decay), "ema" + + def _train_direct_preference_batch( + self, metrics: dict, timing_raw: dict, batch_meta: KVBatchMeta, data: DataProto + ) -> KVBatchMeta: + """Online DPO / DiffusionNFT update: pair (or NFT-prep) then ref noise + actor. + + Unlike Flow-GRPO, DPO does not recompute SDE ``old_log_probs``. The DPO + engine returns ``noise_pred`` from ``infer_actor_batch``, so the PG + old-log-prob hop would crash with ``log_probs is None``. + """ + with marked_timer("prepare_actor_batch", timing_raw, color="brown"): + reward_tensor, reward_extra_infos_dict = extract_reward(data) + data.batch["sample_level_scores"] = reward_tensor + if reward_extra_infos_dict: + data.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()}) + # Persist unpaired scores for metrics before pairing shrinks the batch. + data.batch["sample_level_rewards"] = data.batch["sample_level_scores"] + n_keys = len(batch_meta.keys) + data_for_tq = data.select_idxs(list(range(n_keys))) if len(data) > n_keys else data + put_dataproto_fields_to_tq( + batch_meta, + data_for_tq, + fields=["sample_level_scores", "sample_level_rewards"], + ) + data = self._prepare_actor_batch(data, reward_tensor) + data.batch["sample_level_rewards"] = data.batch["sample_level_scores"] + + if self.use_reference_policy: + with marked_timer("ref", timing_raw, color="olive"): + ref_infer_res = self._compute_ref_noise_pred(data) + if ref_infer_res is not None: + data = data.union(ref_infer_res) + + with marked_timer("update_actor", timing_raw, color="red"): + actor_output = self._update_actor(data) + actor_metrics = reduce_metrics(actor_output.meta_info["metrics"]) + metrics.update(actor_metrics) + if self._has_old_adapter: + metrics.update(compute_old_policy_metrics(self._update_old_policy())) + + return batch_meta + def on_init_end(self): """Called after initialization ends.""" return @@ -954,6 +1094,12 @@ def _compute_old_log_prob(self, data: DataProto) -> DataProto: ) output = self.actor_rollout_wg.infer_actor_batch(batch_td) log_probs = tu.get(output, "log_probs") + if log_probs is None: + raise RuntimeError( + "Actor infer_actor_batch returned log_probs=None. Direct-preference " + "algorithms (DPO / DiffusionNFT) must set algorithm.trainer_type=" + "direct_preference so the trainer skips old-log-prob recomputation." + ) old_log_prob_dict = {"old_log_probs": log_probs.float()} prev_sample_mean = tu.get(output, "prev_sample_mean") if prev_sample_mean is not None: @@ -1011,14 +1157,30 @@ def _update_actor(self, data: DataProto) -> DataProto: batch_td = _to_diffusion_worker_tensordict(data) batch_td = embeds_padding_2_no_padding(batch_td) 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 + paired = bool(self.config.algorithm.get("paired_preference", False)) and getattr( + self, "_is_direct_preference", False + ) + if paired: + ppo_mini_batch_size = ppo_mini_batch_size * 2 + else: + 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 + if paired and shuffle: + logger.warning( + "Shuffle is not supported for direct preference during actor update." + "This is to prevent the chosen/rejected pairs from being split across different micro batches." + "Setting shuffle to False." + ) + shuffle = False tu.assign_non_tensor( batch_td, global_batch_size=ppo_mini_batch_size, mini_batch_size=ppo_mini_batch_size, - epochs=self.config.actor_rollout_ref.actor.ppo_epochs, - seed=self.config.actor_rollout_ref.actor.data_loader_seed, - dataloader_kwargs={"shuffle": self.config.actor_rollout_ref.actor.shuffle}, + epochs=ppo_epochs, + seed=seed, + dataloader_kwargs={"shuffle": shuffle}, height=self.config.actor_rollout_ref.model.pipeline.height, width=self.config.actor_rollout_ref.model.pipeline.width, vae_scale_factor=self.config.actor_rollout_ref.model.get("vae_scale_factor", 8),